One training step of our latent-reasoning agent took 36 minutes on a single H100. That is the median over 46 fresh steps of a 92-step production run, which needed 43 hours before an upstream data source timed out and ended it. The model is Qwen 3.8 27B, trained with latent GRPO: it thinks in continuous vectors, then calls tools in a multi-turn gym, and every rollout is scored again with gradients.
Today the same recipe runs a fresh step in a median of 2.6 minutes on one 8× NVIDIA B300 node, and a reuse step in 16 seconds. The first five steps went from 2 hours 01 minutes to 8.7 minutes: 14× faster, with the same learning rate, group size, latent limits and correctness gates.
No single flag did this. It took a dozen fixes, some of them bugs that would have silently corrupted training. This article walks through each one, with the measurement that justified it.
First 5 Steps14×2 h 01 min → 8.7 min
vLLM Answer Decode62.8 → 5–9 ms/tokNF4 → BF16 + dense LoRA + MTP
1. Anatomy of a 36-Minute Step #
A fresh step has two halves. Generation runs four rollouts of a multi-turn episode: for each model turn, the HF trainer prefills the prompt and samples up to 16 latent thoughts, vLLM samples the text answer from those thoughts, and the gym executes the tool call. Replay then scores every turn again with gradients and back-propagates. A reuse step skips generation and replays the same rollouts once more.
| Median seconds per phase, fresh steps (H100: 46 steps of one production run; B300: current stack) | |||
|---|---|---|---|
| Phase | 1× H100 | 8× B300 | Speedup |
| --- | --- | --- | --- |
| HF prefill | 280 | 16 | 17× |
| HF latent steps | 72 | 26 | 2.8× |
| vLLM answers | 423 | 48 | 8.8× |
| Gym tools | 16 | 7 | — |
| Replay (score + backward) | 1,296 | 24 | 54× |
| Whole fresh step | 2,175 | 156 | 14× |
| Whole reuse step | 882 | 16 | 55× |
Rollouts differ from run to run, so these are not paired measurements. Normalizing by work tells the same story: a fresh step without tool waiting cost 59 s per model turn on the H100 and 4–6 s on the B300 node.
2. The Rule: Gates Never Move #
Every optimization below had to pass the same checks the slow run passed. We never loosened a threshold to make a change land. The checks that mattered most:
- Rollout logprob parity: for every real answer token, the trainer’s log-probability versus the one vLLM reported while sampling: mean gap ≤ 0.1, max ≤ 1.0.
- Teacher-forced argmax agreement between the trainer and vLLM after each adapter update.
- Bitwise thoughts: any change to the latent loop must reproduce the same thought vectors bit for bit from the same seed.
- Gradient equivalence against an FP32 reference on a saved trajectory: relative L2 ≤ 0.002, zero tensors outside tolerance.
Parity is not bookkeeping. GRPO with vLLM sampling corrects for the engine mismatch with an importance ratio, so a logprob error is multiplied into the gradient:
A one-nat error on one token means that token’s gradient is weighted almost three times too heavily. Section 6 is the story of exactly that error.
3. Replay: 1,296 s → 24 s #
Replay was 60% of the fresh step and all of the reuse step. Three changes removed it as a bottleneck.
3a. Stop saving memory you no longer need to save
On an 80 GB H100, replay offloaded activations to CPU, processed MLP layers in 512-token blocks and recomputed aggressively. A B300 has 288 GB. On one saved trajectory:
| Replay of one saved trajectory (microbatch 2) on one B300 | |||
|---|---|---|---|
| Memory policy | Time | Peak | Gradient gate |
| --- | --- | --- | --- |
| BF16 LoRA, H100 policy (offload, MLP 512) | 177.1 s | 38 GB | reference |
| BF16 LoRA, no offload, unchunked MLP | 45.8 s | 92 GB | failed: 2.2e-3 |
| FP32 LoRA masters, no offload, unchunked MLP | 45.9 s | 92 GB | passed: 1.1e-6 |
3b. Keep LoRA weights in FP32
The second row is instructive: the faster policy failed the gradient gate only because BF16 weight gradients are noisy at that size. There was a deeper problem too. The trainer stored the LoRA matrices in BF16, and at a learning rate of 1e-5 many updates are smaller than one BF16 step: in a CPU simulation, 48–72% of lora_A elements did not change at all after an optimizer step. We now keep FP32 master weights, run forward and input gradients in BF16, and compute weight gradients with FP32 accumulation. It costs nothing measurable and makes the updates real.
3c. Split replay by turn across 8 GPUs
The recipe has four rollouts per step, which does not divide well across 8 GPUs by trajectory. Turns do: each turn carries its complete prefix and its saved thoughts, so any GPU holding the same weights can score it independently.
Rank 0 keeps generation and the unmodified TRL loss. It evaluates the loss on the per-token log-probabilities returned by the helpers, and the derivative of that loss becomes one weight per token. Helpers back-propagate the weighted sum and all-reduce FP32 gradients. No loss formula is duplicated, so CISPO, truncated importance sampling and normalization stay exactly as they were. In real training, reuse steps went from 66.5 s to 13.8 s and from 103.8 s to 20.2 s. The limit is the longest single turn: up to 34k of 221k tokens in one step.
4. Blackwell Bring-Up: Kernels That Were Wrong, Not Slow #
B300 is sm_103, and our stack had only been qualified on H100. The dangerous findings were not crashes but silent numerical errors:
Gated DeltaNet output kernel
- The Triton kernel that produces the GDN chunk output returned a wrong forward in about 1 of 7 calls (1.5% error against a reference, versus the usual 0.5%).
- Through the density of latent thoughts, this produced logprob errors in the hundreds, and every replay gate failed.
- Fix: compute that one step in PyTorch with the kernel’s own arithmetic (BF16 products, FP32 sums). 0 mismatches in 50 forward+backward repeats; cost about 1.7 ms per layer at 9,274 tokens.
vLLM graphs under speculation
- With the default mix of FULL and PIECEWISE CUDA graphs, speculative decoding at batch 1 differed from eager execution by up to 0.03 logprob.
- Each mode alone was bitwise equal to eager.
- Fix: PIECEWISE graphs whenever speculation is on. 0.0 difference in 42 of 42 cases, same speed (5.23 s vs. 5.27 s).
Two more were plain failures: the fused recurrent GDN kernel hits an illegal memory access on sm_103 for sequences longer than one token (the model only uses it for single-token decode, where it works), and vLLM picked FlashInfer attention, whose JIT build needed CUDA headers our image did not ship. We pinned FlashAttention, as on H100.
5. vLLM Answers: BF16 Weights, Dense LoRA, MTP #
The text answers come from a vLLM engine that lives in the trainer process and receives the adapter after every update. On the H100 it decoded at a median of 62.8 ms per token. Three changes, each qualified on its own:
| vLLM answer decode, step by step | ||
|---|---|---|
| Change | Why it was slow | Measured effect |
| --- | --- | --- |
| BF16 base weights | vLLM served the NF4 checkpoint through bitsandbytes, dequantizing on every step. | Decode step 36.9 → 12.3 ms. The weights are the trainer’s NF4 base dequantized once, so values are identical; exported once per node (52 GB). |
| Dense single-adapter LoRA | The multi-adapter punica kernels run with split-K = 1 for determinism, leaving almost no parallelism. The adapter cost more than the model: 32.8 ms per step with it, 14.3 ms without. | With exactly one adapter, shrink and expand become two cuBLAS BF16 GEMMs with the same arithmetic. Step 33.0 → 24.0 ms, graphs still bitwise. |
| MTP speculation | N-gram drafts are rarely right on free-form reasoning text. | The checkpoint’s own multi-token-prediction head: 27.1 → 12.0 ms per token in the trainer engine, about 3 tokens accepted per verify step. |
In real training, answers went from 12.1–20.7 ms per token (NF4 on B300) to 7.0–8.9 ms with BF16 and dense LoRA, and to 5.0–9.4 ms with MTP. But MTP did not pass on the first try.
6. The One-Token Bug: MTP Meets the Hybrid Prefix Cache #
MTP passed our synthetic speculative-decoding qualification 14 of 14. The first real training step failed the rollout parity gate. The mean gap was fine (0.004–0.01), but the maximum over about 2,000 answer tokens was 1.15, above the 1.0 limit. We reran it four more times; every run failed at step one with a maximum between 1.02 and 2.70. N-gram runs had passed the same gate on every step.
One token repeated across runs with identical numbers: a newline after a closing tool call, trainer logprob −4.186, vLLM −3.163. A deterministic error is a gift, because it can be bisected. We added a diagnostic that runs only when the gate has already failed: for the worst tokens it rebuilds the exact turn (prompt, sampled thoughts, adapter) and scores the target token on every path we can isolate.
Every isolated path agreed, including MTP verification itself, alone and in a batch of eight. The diagnostic clears vLLM’s prefix cache before each probe; the live rollout does not. Qwen 3.8’s GDN layers keep a recurrent state, and vLLM caches those states at block boundaries in a mode it labels experimental. With MTP, almost every step verifies four draft tokens and rolls back the rejected ones, and a later turn that reuses a cached prefix appears to start from a state that does not match its tokens. With prefix caching off, all five steps passed with a maximum gap of 0.26–0.31, and answers still decoded at 5.0–9.4 ms per token.
Rule we adopted: MTP with prefix caching on a hybrid model is rejected at configuration time, before a model is loaded. This matters for RL, where the sampling logprob enters the gradient. For pure serving, where only the text matters, see our
7. Prefill Reuse for a Model Without a Plain KV Cache #
Four rollouts share a prompt, and each new turn extends the previous one, so most prefill work is repeated. The trainer had a prefill cache, but it only knew attention KV tensors. On a hybrid model it never produced a hit, so every turn re-read its full history.
The new cache stores, per rollout row, both the attention K/V and the GDN convolution and recurrent states, keyed by token prefix. Each row is prefilled alone from its longest stored prefix, without padding, which a recurrent state cannot tolerate. The group’s shared prompt is computed once, and afterwards only each turn’s new tokens. Rows are then stacked into a left-padded batch cache for decoding. In real training, prefill went from 28–51 s to 12–18 s per fresh step, computing only 10–20% of the input tokens.
8. Latent Steps: A 27B Model Waiting for Python #
After everything above, the latent loop stood out: about 300 ms per forward pass at batch 4, in real training and in an isolated benchmark alike. A profile explained it. For a one-token, batch-4 forward the GPU is busy about 15 ms. The rest is the host walking PEFT DoRA modules, FP32 LoRA casts and NF4 dispatch, launching thousands of small kernels. A faster GPU would not help at all.
The decoder already ran on a static cache with fixed input, position and mask buffers, and transformers 5.13 keeps the cache cursor on the device and updates GDN states in place. That is exactly what CUDA graphs need. The first step runs eagerly as warm-up, the second is captured, and every later step is a single replay:
| 16 latent steps, batch 4, 3k-token prompts, one B300 | ||
|---|---|---|
| Decoder | Time | Thoughts |
| --- | --- | --- |
| Dynamic cache | 4.81 s | reference |
| Static cache, eager | 5.24 s | bitwise equal |
| Static cache + CUDA graph replay | 1.65 s | bitwise equal |
In real training, latent time per fresh step dropped from 31–59 s to 11–30 s. What remains is the eager warm-up and capture paid once per rollout call, and sampling thoughts on the host.
9. The Plumbing Bugs #
Three small bugs cost more time than some kernels:
- Inherited distributed environment. Replay helpers were started from the vLLM process and inherited its
RANKandMASTER_PORT, so they tried to join the wrong process group and failed with an address in use. Helpers now start with a clean environment. - Inherited Triton cache. vLLM points
TRITON_CACHE_DIRat its compile cache, which is keyed by engine configuration. Every new configuration made helpers re-tune their kernels: up to 93 s inside a fresh step. With the normal Triton cache and persisted autotuning, 5 of 7 helpers tune nothing. - Dead helpers must fail loudly. A crashed helper now closes its channel and rank 0 stops at once, instead of waiting forever on a GPU that will never answer.
10. Results and What We Would Do Again #
| First five steps of the same recipe (different rollouts, so not paired) | ||||
|---|---|---|---|---|
| Step | Type | 1× H100 | 8× B300 | Speedup |
| --- | --- | --- | --- | --- |
| 1 | fresh | 34:42 | 2:11 | 16× |
| 2 | reuse | 10:30 | 0:13 | 49× |
| 3 | fresh | 43:52 | 2:37 | 17× |
| 4 | reuse | 8:19 | 0:18 | 27× |
| 5 | fresh | 23:32 | 3:24 | 7× |
| Total | 2:00:54 | 8:42 | 14× |
The run finished with every gate green: non-zero gradients on all five steps, rollout parity within limits on every step, teacher-forced agreement and a verified checkpoint. By a rough estimate that accounts for later steps being longer, the 43-hour, 92-step H100 run would take about 3 hours on this node. We will confirm it with a full run.
What paid off
- Measure the phase, not the step. Per-phase timers and work counters in every step showed where time went.
- Fix precision before speed. FP32 LoRA masters made a faster memory policy pass a gate it had failed.
- Parallelize what is independent. Turns, not trajectories, were the unit that fits 8 GPUs.
- Profile before you guess. The latent loop needed CUDA graphs, not a faster kernel.
What protected us
- Fixed thresholds. The MTP bug surfaced because a 1.0-nat limit stayed at 1.0.
- Diagnostics on failure. Rebuilding the failing token on every path turned a vague mismatch into one cache setting.
- Bitwise checks where possible. Thoughts, CUDA graphs and turn-parallel logprobs are compared bit for bit, so any drift is a bug.
- New hardware is new numerics. Two Blackwell kernel issues were silent until a gate compared against a reference.
For background on the method itself, see our latent GRPO deep dive and how asynchronous rollouts keep GPUs busy while gyms run tools.