Paper Walkthrough: Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training — training the draft alongside a 122B policy at 256K tokens

Rollout generation dominates RL post-training time. NVIDIA speeds it up with speculative decoding and keeps training the draft model alongside the policy — and shows how to get past the two walls that context parallelism and pipeline parallelism put in the way. A from-scratch walkthrough grounded in the paper text.

Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training

What the paper solves

The paper is "Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training" (arXiv:2609.07108, NVIDIA, Technical Report), by Zili Wang, Zhaopeng Qiu, Yuekai Zhang, Shuang Yu, and Junjie Lai, published 2026-09-07.

Here is the abstract in plain terms. In reinforcement learning (RL) post-training, rollout generation — having the policy model write its answers — dominates the cost. Speculative decoding speeds that up, and if you keep training the draft model online alongside the policy, the draft gets more accurate and the speedup grows. But doing this on large models with long contexts runs into two obstacles: (1) the "branch attention" that draft training needs is not supported by standard causal context-parallel (CP) implementations, and (2) the intermediate policy features the draft consumes are spread across pipeline-parallel (PP) stages. The paper handles both: it extends packed, load-balanced zigzag ring attention to merge rank-local branch attention with causal main-sequence attention, and it introduces TapChannel, which moves intermediate features across stages on a separate path without touching the pipeline schedule. Experimentally, co-trained drafts track the policy baseline closely while delivering substantial rollout and end-to-end speedups up to 122B, the CP design scales at 256K tokens with large memory savings over prior work, and the PP transport adds modest overhead.

An analogy: the stenographer and the proofreader

Start with speculative decoding (SD) itself. A large language model writing one token at a time is like a slow proofreader looking up every word in a dictionary. Each word requires running the entire model once, so the GPU sits mostly idle.

So you put a fast stenographer (the draft model) next to the proofreader. The stenographer jots down the next few words; the proofreader checks them all at once, keeps the correct ones, and rewrites from the first mistake onward. Because the check uses rejection sampling, the output distribution is identical to what the proofreader would have written alone (§1). Quality is untouched; you only get faster as long as the stenographer keeps guessing right. The basics are in Speculative decoding from scratch.

Why bring this to RL? Because, in the paper's words, "the wall-clock time of RL post-training is often dominated by rollout generation" (§1). Methods like GRPO have the policy generate many answers, score them, and update. That generation is the slow part.

Why train them together

The catch is that in RL, the proofreader's writing style changes every step. A fixed stenographer drifts out of sync and the number of accepted words per check (the accepted length) drops. The paper puts it as "a fixed draft loses alignment with an evolving policy" (§2).

The fix is to keep training the stenographer too — online co-training. The objective is Eq. (1) in §3.1:

Read it as "the policy's RL loss (parameters $\theta$) plus $\lambda$ times the draft's loss (parameters $\phi$)." $x$ is the rollout token sequence, $H_\theta(x)$ is the set of features pulled from the policy's intermediate layers, and $\operatorname{sg}$ is stop-gradient — "don't send gradients back through these features." The draft learns by looking inside the policy, but the policy is never pulled around by the draft. The draft lives as a submodule on the policy's last pipeline stage and is updated jointly with it (§3.1).

The two walls

Everything so far exists in prior work (FastGRPO, ReSpec, MTP-RL, etc., §2). This paper's question is what breaks when you do it at scale with long contexts.

Wall 1: branch attention. Strong modern drafts — EAGLE-3, DFlash, DSpark — use branched attention during training. EAGLE-3 runs Training-Time Test (TTT): it feeds its own predictions back as input and rehearses several steps ahead, so a branch grows from every position. DFlash generates a whole block of tokens in one forward pass, and DSpark adds a light Markov head on top (§3.2). The shared requirement: every branch must attend to both the causal prefix of the main sequence and its own branch-local KV. Standard CP, which shards a long context across GPUs, assumes plain causal attention.

Wall 2: features straddle stages. The draft takes features from the policy's intermediate layers, but under PP those layers sit on several stages and the draft is only on the last one. Standard pipeline communication only connects adjacent stages (§3.3).

If the parallelism vocabulary is fuzzy, read Distributed training from scratch first. The paper's stance is explicit: prior work reshapes the parallel layout to accommodate draft training; "we keep the target's topology unchanged and instead adapt the system around it" (§2).