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

Paper

Two Designs That Make Speculative Decoding Practical for Large-Scale, Long-Context RL Post-Training

TL;DR : Rollout generation dominates the wall-clock time of RL post-training. This work accelerates it with speculative decoding, and adds online draft co-training so the draft does not lag behind as the policy evolves. However, existing systems fail to handle the branch attention used by advanced drafts (unsupported by context parallelism) and target features scattered across stages (a pipeline parallelism problem). This paper proposes (1) a CP technique that decomposes branch attention into a causal main-sequence component and a rank-local branch and merges them, and (2) TapChannel, which delivers target features outside the pipeline schedule, achieving 1.16–1.88× end-to-end speedup from 8B to 122B (source: §3.3, Tab. 1).

Core Idea

In RL post-training, speculative decoding (SD) accelerates autoregressive generation without changing the output distribution: a small draft first proposes several tokens, and the larger policy model verifies them in parallel (source: §1). If the draft is kept fixed, acceptance length degrades as the policy evolves, so co-training the draft online on the policy’s rollout tokens yields larger speedups (source: §1). The problem is that state-of-the-art drafts such as EAGLE-3, DFlash, and DSpark (a) use branch-structured attention and (b) consume the target model’s intermediate hidden states (source: §1, §2.2). These collide head-on with the limitations of standard causal context parallelism (CP) and pipeline parallelism (PP), respectively. The authors’ core claim is simple.

“If we keep the target model’s existing CP·PP topology intact and adapt the system around it, draft co-training can be made practical at large scale and long context.”

In other words, unlike prior approaches that reshape the parallelization scheme to fit the draft, this paper’s originality lies in preserving the target’s layout and adding two mechanisms on top of it (source: §2).

Background: The Problem They Solved

RL post-training (GRPO and others) has become a standard paradigm exemplified by DeepSeekMath, DAPO, GLM5, and most of its wall-clock time is spent on rollout generation (source: §1). SD mitigates this with parallel verification, but prior work focused mostly on draft sourcing, serving setups, and verification optimization (source: §2, Related Work). The system-level challenges of actually “keeping the draft trained to follow an evolving policy” are two.

1. CP does not support branch attention. EAGLE-3 uses TTT (Train-Time Test) to re-input the previously predicted hidden at each draft position, creating a separate branch, and DFlash/DFlash-family drafts (DSpark) also require in-block positions to attend to a branch-local context (source: §2.2). Standard causal CP cannot express such branch structures. SpecForge handles EAGLE-3 TTT, but its sequential ring sharding leaves the causal workload imbalanced and constrains the Ulysses dimension (source: §2.2).

2. Under PP, target features live on stages far from the draft. The draft is trained as a submodule residing on the policy’s last pipeline stage, but the target hidden states it consumes (so-called taps) are produced on several earlier stages. Standard pipeline communication only connects adjacent stages, so these non-adjacent features cannot be delivered (source: §2.3).

The authors integrated an end-to-end system that resolves both obstacles into NeMo-RL (source: §1).

The New Approach: Branch Attention under CP + TapChannel under PP

1) Decomposing Branch Attention under CP

Each branch query must attend to two sets of keys: (a) the causal prefix of the main sequence (sharded across CP ranks) and (b) branch-local keys (present only on the rank owning the branch’s anchor). The authors compute the two independently and merge them via an online-softmax (LSE) reduction (source: §2.2, Fig. cp-branch-attention). The main-sequence component follows the same packed zigzag-ring attention as standard CP, while the branch-local component is computed locally on the anchor-owning rank.

Here $(O_{\mathrm{m}}, \ell_{\mathrm{m}})$ are the output and log-sum-exp of the main-sequence component, and $(O_{\mathrm{b}}, \ell_{\mathrm{b}})$ those of the branch-local component (source: §2.2, Eq. 3). The strength of this design is that this single merge rule covers EAGLE-3 (TTT), DFlash (block parallel), and DSpark (Markov head) simultaneously (source: §2.2).

2) TapChannel under PP

The draft lives only on the last stage, but the layers that produce taps span multiple stages. Since taps need no return path, the authors designed TapChannel, which delivers features through a side pass outside the pipeline schedule (source: §2.3, Fig. tap-channel). Per-source mailboxes (pre-allocated buffer slots) live on the draft stage: once a source finishes that micro-batch’s policy forward, it writes the tap into its slot, and the draft reads it right before its own forward for the same micro-batch. Ordering is handled by incrementing a sequence stamp per slot; same-node sources use CUDA IPC, while cross-node sources use a dedicated NCCL communicator backed by GPUDirect RDMA (source: §2.3).

How It Works: A Concrete Walkthrough

Let’s compress CP branch attention to a 3×3 toy case. Suppose the sequence is split across CP=3 ranks as [A B C | D E F | G H I] (source: §2.2). In EAGLE-3 TTT, assume that after position C, the draft generates branch tokens c1, c2 from its own predicted hidden states (source: §2.2).

• Main-sequence component: each rank’s queries must see all main-sequence tokens that precede them. In the zigzag ring, K/V circulate among ranks while queries stay in place, so a rank holding D E F receives the K/V of A B C and attends to it. Call the resulting output $(O_{\mathrm{m}}, \ell_{\mathrm{m}})$.
• Branch-local component: the anchor of branch tokens c1, c2 is rank 0 (owning A B C), so the keys and values of c1, c2 exist only on rank 0. Only rank 0 computes branch attention locally to obtain $(O_{\mathrm{b}}, \ell_{\mathrm{b}})$.
• Merge: combine the two results with the equation above. If $\ell_{\mathrm{m}}=3.1$, $\ell_{\mathrm{b}}=1.2$, then $\ell=\log(e^{3.1}+e^{1.2})\approx3.13$, and the exponentially larger component (main sequence) dominates the final output. Because this is exactly the same form as the online-softmax reduction used between ring attention steps, it layers naturally onto existing kernels (source: §2.2).

From a communication cost perspective, if the CP degree is $C$ and the $N$ main-sequence tokens are evenly split across ranks, the per-rank send volume in the forward ring is

where $d_{\mathrm{kv}}$ is the per-token K/V width and $b$ is the bytes per element. Since branch-local K/V stays on the anchor-owning rank, this cost is independent of the number or depth of branches (source: §2.2, Eq. 4). Because attention grows quadratically with context length while communication grows linearly, longer contexts leave more overlap headroom to hide communication behind compute (source: §2.2).

TapChannel’s cost is equally intuitive. Let $n$ be the number of token rows in a micro-batch, $\mathcal{S}$ the set of source stages sending taps, and $d_s$ the per-token feature dimension of source $s$. The payload per micro-batch is

and each feature is transferred directly exactly once, independent of the number of PP hops (source: §2.3, Eq. 6). Because the pipeline schedule provides natural slack between a source’s tap production and the draft’s forward, the transfer has zero latency if it completes within this slack; only the part exceeding the slack becomes visible overhead (source: §2.3, Eq. 7).

Evaluation: Key Results

1) Training Trajectories Are Preserved

The first thing verified was “does SD/co-training not break RL training?” Comparing the baseline (no SD, no co-training) against EAGLE-3/DFlash/DSpark co-training with Qwen3-8B as the target, reward, verification accuracy, and training-inference KL divergence all closely tracked the baseline trajectory (source: §3.2, Fig. learning-stability). In particular, a near-zero KL divergence means the training and inference backends are numerically identical at the same weights—a key condition for honoring GRPO’s on-policy assumption (source: §3.1).

2) Three Draft Families across 8B–122B Scales

Co-trained drafts achieved acceptance lengths of 2.28–4.78, rollout speedups of 1.19–2.23×, and end-to-end training speedups of 1.16–1.88× (source: §3.3, Tab. 1).

Two patterns stand out (source: §3.3). First, DFlash/DSpark have longer acceptance lengths than EAGLE-3 and are consistently faster. Second, the large MoE targets (Qwen3.5-122B, GPT-OSS-120B) have high acceptance lengths yet lower end-to-end speedups—because sparse routing triggers more expert compute per verification forward (source: §3.3).

3) Limits on Multi-Turn Workloads

In NeMo Gym Workplace Assistant (a multi-turn, agentic, tool-using environment), end-to-end speedup is 1.25–1.43×, well below the 1.75–2.23× rollout-stage speedup. The reason is that rollout accounts for only 55.8% of step time—the remaining tool execution and environment latency cannot be touched by decoding acceleration (source: §3.3, Fig. workplace-learning).

4) CP Attention Performance: Advantage over USP and 256K Scaling

On the same workload, packed zigzag attention was compared against SpecForge’s USP. USP pads the batch to 2.25× the actual token count, while the packed implementation avoids this overhead. At CP=2, 4, 8, packed zigzag showed 2.9×, 2.3×, 1.5× lower latency than the best USP and 2.7× lower peak memory per GPU (source: §3.4, Fig. pack-dist). In long-context scaling, TTT attention latency dropped from 17.7s at CP=1 to 2.35s at CP=8, a 7.5× reduction (94% parallel efficiency), and per-GPU memory decreased almost linearly from 53.2 GB to 7.5 GB (source: §3.4, Fig. cp-scaling).

5) PP Overhead Is Overlappable

TapChannel’s raw transfer hits 27–39 GB/s via one-sided writes, 4.5–8.5× faster than host staging, with noise-level impact on source stages and only 1.6% HBM contention on the receiving draft stage. Host staging, by contrast, slows every rank by 80% or more (source: §3.5, Fig. tap-transport). In full runs, draft co-training adds under 15% update-time overhead for block-style drafts and 34% for EAGLE-3 (due to its TTT passes). Yet it shortens rollout by 28–52%, yielding a net speedup of 1.31–1.85×, and tap wait time is only 0.4–0.6s per policy update (1.5–2.2% of optimization time)—almost entirely hidden in the pipeline schedule (source: §3.5, Tab. 2).

Our Take: Strengths, Limitations, and Why It Matters

The strengths are clear. First, the design principle of target topology invariance stands out. Where prior work reshaped the parallel layout to fit the draft, this paper adapts the system without touching the policy’s CP·PP configuration, making it possible to co-train the draft while the policy keeps training (source: §2, Related Work). Second is the generality: a single LSE merge rule simultaneously supports three structurally different drafts—EAGLE-3, DFlash, and DSpark (source: §2.2). Third, the broad validation from 8B to 122B and from single-turn to multi-turn, together with quantifying trajectory preservation via KL divergence, boosts practical confidence (source: §3.2).

The limitations are equally clear. First, on MoE targets and linear-attention models the benefit of speculative decoding is structurally limited: MoE pays expert-routing cost per verification, and linear attention makes verification itself already cheap, leaving little relative speedup headroom (source: §3.3, §4). Second, EAGLE-3’s multiple TTT passes incur 34% update overhead, putting it at a disadvantage relative to block-style drafts (source: §3.5, Tab. 2). Third, when rollout occupies only about half of total step time, as in multi-turn, tool-using environments, the end-to-end transfer of decoding acceleration drops sharply—suggesting this approach’s value is concentrated on rollout-dominated workloads (source: §3.3). Also, the paper leans heavily on “building the system,” so its essence is making existing drafts practical for large-scale distributed training rather than contributing new draft architectures.

Even so, the research matters because it attacks head-on, at the system level, the fact that the real bottleneck of RL post-training is not “model training” but “rollout generation.” Without a framework that co-trains the draft while preserving the training trajectory, SD-based RL acceleration would have remained a one-off trick viable only at small scale and short context.

What’s Next?: The Road Ahead

The authors list adapting speculative decoding to sparse MoE and linear-attention models as future work (source: §4). Building on that, reasonable next steps would be:

• Cut MoE verification cost: routing caching that trims the expert compute in verification forwards, or expert-level speculative verification, to convert acceptance-length gains into end-to-end speedup even on large MoEs.
• Ease the non-decoding bottleneck in multi-turn rollouts: overlap the tool/environment-latency windows with decoding, or use session-level draft conditioning to raise acceptance in multi-turn settings.
• Shrink EAGLE-3’s TTT overhead: reduce the number of TTT passes, or graft on a low-cost DSpark-style Markov head, to converge to block-style update costs.
• Ease exposed communication under strong CP scaling: branch-aware K/V compression or shuffle-schedule optimization that improves the regime where communication becomes the bottleneck as the per-rank context shrinks (§2.2).

In the end, the paper’s significance is not “a specific draft” but a pattern: making drafts that live and move with an evolving policy practical in large-scale distributed training.

Tables from the paper

Tables converted mechanically from the arXiv e-print LaTeX source. The numbers are the paper’s own and did not pass through a model.

Table 1. End-to-end performance across target scales and draft families.

Table 2. Pipeline-parallel overhead for Qwen3-8B, averaged over first 10 policy update steps. DSpark achieves the largest end-to-end speedup (1.85$\times$) with 14.6% update-time overhead.

Figures in this post are taken from the original arXiv:2609.07108 (CC BY 4.0). Only size and format were changed.