面向大规模长上下文RL后训练的投机解码在线草稿协同训练

分支注意力CP分解与TapChannel旁路传输，让草稿在线协同训练实现最高1.88倍加速

前置知识

投机解码（Speculative Decoding）

一种无损加速自回归生成的技术：轻量草稿模型一次并行提出多个候选token，大目标模型对这批候选做一次并行验证，用拒绝采样决定接受哪些。由于拒绝采样保证最终输出分布与目标模型完全一致，加速不改变生成质量。实际加速比取决于平均接受长度（acceptance length）——每次验证平均被接受的token数。

本文整个系统就是为了在RL后训练的rollout生成中利用投机解码，并用在线协同训练维持草稿与演化策略的接受长度；不理解提议/验证/接受机制就无法读懂加速与接受长度指标。

上下文并行（Context Parallelism, CP）

把超长序列切分到多张GPU上训练注意力的方法。RingAttention让主序列的K/V沿环形拓扑在rank间循环传递、查询驻留本地，通信量随序列长度线性增长；zigzag（锯齿）分片让每个rank同时持有序列的头部和尾部chunk，使因果注意力的计算负载均衡；packed指变长序列打包消除padding浪费。

本文CP侧贡献是在标准因果zigzag-ring基础上扩展出分支注意力，必须理解K/V环传、查询驻留、zigzag负载均衡，才能读懂Figure 1和Eq.(2)(3)的合并与通信量分析。

流水线并行（Pipeline Parallelism, PP）

把模型按层切分为多个stage分布到不同GPU，microbatch像流水线一样依次流过各stage。标准PP通信只发生在相邻stage之间传递激活，且严格由调度器（如1F1B）编排，任何新增的跨stage数据流都可能干扰调度、产生气泡。

草稿模型只放在最后一个PP stage，而它条件化所需的目标中间特征（tap）分散在多个stage上；理解PP的相邻stage通信限制才能理解TapChannel为什么必须设计成调度外的旁路通道。

在线草稿协同训练（Online Draft Co-Training）

在RL训练过程中让草稿模型与策略同步演化的技术：草稿在策略自己的rollout tokens上以辅助损失训练，输入还包含从策略中间层抽取的目标特征（tap），并用stop-gradient阻断梯度回流策略。随策略更新，草稿持续适配，接受长度不衰减甚至上升。

这正是本文要做的事本身：联合目标 $\mathcal{L}_{RL}+\lambda\mathcal{L}_{draft}$、stop-gradient和tap接口是全文出发点，实验中接受长度曲线直接反映协同训练效果。

EAGLE-3 / DFlash / DSpark 三种草稿家族

三种先进草稿架构。EAGLE-3用目标多层融合特征做条件草稿器，训练时以TTT（Training-Time Test）模拟多步生成、把上一步预测的隐状态喂回输入，从而在每个草稿位置产生一个注意力不断增长的分支；DFlash用块扩散语言模型一次生成整块候选；DSpark在DFlash上叠加低秩、以前一token为条件的Markov头，从左到右恢复块内因果依赖。三者注意力掩码与分支结构各异。

本文CP机制的统一性正体现在同时支持这三种家族；理解各自分支形态（TTT多步分支、块分支、Markov头分支）才能评价分支注意力分解的通用性主张。

RL后训练与GRPO

在预训练/指令微调后的LLM上用强化学习继续优化推理与智能体能力。GRPO对同一prompt采样一组响应，用组内相对奖励估计优势，无需critic。rollout（用当前策略生成响应）占据RL训练的大部分墙钟时间，长上下文任务尤甚；GRPO还要求训练与推理后端数值一致以保持on-policy假设。

本文加速对象正是RL后训练的rollout阶段；GRPO对数值一致性的要求解释了论文为何用训练-推理KL散度（近零）作为系统正确性的核心验证指标。

研究动机

RL后训练（如GRPO）的墙钟时间往往由rollout生成主导，投机解码通过小草稿模型并行提议、大目标模型并行验证来加速生成且不改变输出分布，已被NeMo-RL、veRL等rollout引擎集成。近期工作发现随RL策略不断演化，固定草稿会逐渐失配，在线协同训练草稿能让接受长度随策略提升，进一步放大加速。但把在线协同训练扩展到大模型、长上下文场景存在两个系统级障碍：其一，EAGLE-3、DFlash、DSpark等先进草稿架构在训练时需要分支注意力——每个分支查询既要看主序列的因果前缀、又要看分支局部的KV——而标准因果上下文并行不支持这种结构，SpecForge虽支持EAGLE-3的TTT但依赖顺序环分片，因果负载不均衡且对Ulysses维度限制苛刻；其二，草稿条件化所需的目标中间特征分布在多个流水线并行stage上，而标准PP通信只连接相邻stage，草稿无法获取远处stage的特征。论文实验目标横跨Qwen3-8B到Qwen3.5-122B-A10B、GPT-OSS-120B，序列长度最高256K，这些障碍在实际规模下无法回避。

本文的目标是本文目标是构建一套端到端系统，使草稿模型能在已配置上下文并行与流水线并行的目标策略训练系统内在线协同训练，且完全不改变目标模型既有的并行拓扑。具体分解为四个可验证子目标：(1) CP侧设计在packed、负载均衡zigzag-ring执行下的分支注意力机制，用同一套机制统一支持EAGLE-3、DFlash、DSpark三种分支形态各异的草稿家族；(2) PP侧设计一条不进入、也不改变流水线调度的旁路通道TapChannel，把分散在各stage的目标特征按时送达草稿所在的最后一个stage；(3) 将两者集成进NeMo-RL，在8B到122B目标规模、单轮数学与多轮智能体任务上验证端到端可行性与加速（最终实测1.16-1.88倍E2E加速）；(4) 量化验证系统正确性（训练-推理KL一致性）、CP长序列强扩展性（至256K tokens）以及PP侧协同训练的增量开销（实测13.6%-34.3%更新时间开销）。

与已有工作不同的是，已有工作从三个方向接近该问题，但都未覆盖本文场景。目标函数层面：Fast-GRPO做并发感知的在线草稿学习，ReSpec动态选择投机参数并蒸馏策略，MTP-RL用优势感知系数优化共享MTP模块，OCC自适应平衡辅助损失与策略更新，Bebop直接以接受率为目标训练——它们都不解决CP/PP系统问题。不学习草稿的路线：SPEC-RL复用上一轮策略的响应片段做投机前缀，EfficientRollout构造量化自草稿并按接受率启停投机。并行布局路线：P-EAGLE用结构化掩码并行化EAGLE训练，LongSpec面向长上下文推理，SpecForge为目标-草稿解耦设计顺序环分片的分支注意力——它们都要修改并行布局。本文的独特切入是反其道而行：保持目标策略的CP+PP拓扑不动，把适配工作全部做在草稿侧——CP上将分支注意力分解为可复用标准zigzag-ring的因果分量加无需通信的rank本地分量，PP上用调度外的mailbox旁路传输特征，与目标函数层面的改进正交互补。

核心方法

直觉：草稿的训练样本就是策略自己产生的rollout tokens $x$，因此可与策略联合优化：$$\mathcal{L}(\theta,\phi)=\mathcal{L}_{RL}(\theta)+\lambda\,\mathcal{L}_{draft}(\phi;\,x,\,\mathrm{sg}(H_\theta(x)))$$ 其中 $H_\theta(x)$ 是策略中间层收集的目标特征（tap），$\mathrm{sg}$ 阻断梯度回流策略；草稿作为策略最后一个PP stage的子模块与 $\theta$ 联合更新。技术路线分两支。CP侧：每个分支查询面对两组key——主序列因果前缀（packed成zigzag chunks、K/V沿环循环）与分支局部keys（留在锚点rank上）；分别用标准ring注意力和本地注意力计算，再按在线softmax合并：$\ell=\log(e^{\ell_m}+e^{\ell_b})$，$O=(e^{\ell_m-\ell}O_m+e^{\ell_b-\ell}O_b)/(e^{\ell_m-\ell}+e^{\ell_b-\ell})$。PP侧：每个产tap的stage在草稿stage显存中有预分配mailbox slot，源stage完成前向后写入、草稿在同一microbatch前向前读取；slot带序列stamp握手，同节点走CUDA IPC、跨节点走GPUDirect RDMA，完全旁路流水线调度。

核心思想与所有修改并行布局的方案相反：目标策略的CP+PP拓扑一概不动，让系统围绕它适配。第一，分支注意力的CP分解：把分支查询的注意力拆成主序列因果分量与rank本地分支分量，因果分量完全复用标准packed zigzag-ring注意力，分支分量在锚点rank本地完成、零额外通信，再用在线softmax合并恢复精确输出。由于分支KV不参与环传，每rank前向通信量 $V^{fwd}_{CP}=\frac{2(C-1)N}{C}d_{kv}b$ 与分支数量和深度完全无关，反向只重放同样的环并本地累积梯度。第二，TapChannel的调度外传输：tap不需要返回路径，可绕开相邻stage的P2P通道，用mailbox加序列stamp握手，不理会流水线调度；调度天然提供了源stage产出tap与草稿同一microbatch前向之间的slack窗口 $\Delta_{s,m}$，只有超过slack的部分才变成可见开销 $E^{(m)}_{tap}=\max_{s\in S}(\tau_{s,m}-\Delta_{s,m})$，且每个特征只直接传一次、与PP跳数无关。

方法步骤详情

五步流程。第一步实例化：官方checkpoint的草稿（EAGLE-3/DFlash/DSpark）挂为策略最后一个PP stage的子模块，tap施加stop-gradient。第二步CP侧前向：主序列packed成zigzag chunks，$C$ 个rank各持 $N/C$ token；每个ring step把本地K/V发给其余 $C-1$ 个rank、查询驻留，注意力与下一轮通信重叠；锚点rank本地算分支KV注意力，两分量按Eq.(2)合并。第三步tap传输：各stage在策略前向产出中间特征，经TapChannel写入草稿stage预分配slot——同节点CUDA IPC、跨节点GPUDirect RDMA，序列stamp同步，实测27-39 GB/s。第四步草稿前向与损失：草稿读slot中的tap，EAGLE-3做多步TTT模拟、DFlash一次生成整块、DSpark用Markov头恢复块内因果，在rollout tokens上算草稿损失。第五步反向与更新：反向重放K/V环并本地累积梯度，与策略梯度一起更新 $\theta$ 与 $\phi$。

技术新颖性

技术新颖性有三层。第一，分支注意力的CP分解是首个在packed、负载均衡zigzag-ring执行下统一覆盖三种分支式草稿架构的机制：EAGLE-3的TTT在每个草稿位置产生注意力不断增长的分支，DFlash的块扩散整块并行生成，DSpark再叠加低秩Markov头——三者形态迥异，但共享'分支既要看主序列因果前缀、又要看分支局部KV'的需求，本文把它抽象为两组key集的分解加在线softmax合并。相比SpecForge的顺序ring分片（因果负载不均衡、Ulysses维度受限、需padding到2.25倍token），实测CP=2/4/8延迟提升2.9×/2.3×/1.5×、峰值显存低2.7×。第二，TapChannel是首个面向PP下目标特征fan-in的调度外one-sided传输通道：host staging基线让每个rank慢83-88%，TapChannel对源stage零干扰、接收stage仅1.6%争用。第三，系统级首次在NeMo-RL中打通8B-122B、TP/PP/CP/EP混合并行下的端到端在线草稿协同训练，并以训练-推理KL验证on-policy性。

**Figure 1**: Illustration of branch attention under context parallelism. (a) The causal prefix is zigzag-sharded across CP ranks; branch-local keys stay local. (b) Each rank merges the ring attention result with local branch attention via Eq. (2).

**Figure 2**: Illustration of TapChannel feature fan-in under PP=4. Tap-producing stages deliver target features to the draft stage over an out-of-schedule path, bypassing the standard adjacent-stage P2P communication (gray).

实验结果

核心发现分五组。(1) 正确性：Qwen3-8B在DAPOMath-17K（输入4096/响应16384，AIME2024评估）上，三种草稿协同训练的reward与验证准确率与基线重合，训练-推理KL约$6$-$7\times10^{-4}$，未破坏GRPO的on-policy假设。(2) 单轮：Qwen3-8B上EAGLE-3/DFlash/DSpark接受长度2.28/3.45/3.63，rollout加速1.63×/2.23×/2.18×，E2E 1.50×/1.88×/1.83×；更大模型上Qwen3.5-122B-A10B+DFlash接受长度4.78全场最高但E2E仅1.35×，因MoE验证前向触发稀疏路由更多专家计算，吞掉部分收益（各配置完整数字见表1）。(3) 多轮：Workplace Assistant（32K序列）上E2E 1.25-1.43×，低于rollout加速1.75-2.23×，工具执行与环境延迟不可被更快解码压缩。(4) CP效率：对USP延迟提升2.9×/2.3×/1.5×（CP=2/4/8）、显存低2.7×；256K时TTT注意力17.7s→2.35s（7.5×，94%并行效率），显存53.2→7.5GB。(5) PP开销：EAGLE-3更新时间+34.3%（TTT多pass）但E2E仍1.31×，DFlash +13.6%/1.37×，DSpark +14.6%/1.85×；tap等待0.37-0.57s，仅占优化时间1.5-2.2%。

**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×) with 14.6% update-time overhead.

**Figure 3**: Experimental results of Qwen3-8B online co-training with EAGLE-3 (yellow), DFlash (red) and DSpark (blue) on DAPOMath-17K, including reward, accuracy evaluated on AIME2024, KL divergence, accepted length, throughput and end-to-end step time.

**Figure 4**: Qwen3-8B online co-training with EAGLE-3 (yellow), DFlash (red), and DSpark (blue) on NeMo Gym Workplace Assistant.

**Figure 5**: Comparison of packed zigzag-ring attention (ours) vs. padded USP attention under the same workload. Top: latency of the forward+backward pass; bottom: per-GPU peak memory.

**Figure 6**: Scaling of draft-training attention over CP=1, 2, 4, 8. Each group fixes the global sequence length and reports forward+backward latency (top) and peak HBM (bottom).

**Figure 7**: TapChannel vs. host staging on a PP=4 fan-in. Left: fan-in latency; right: slowdown during transfers, measured as the time increased when transfers run concurrently.

查看结构化数据

任务	指标	本文	基线	提升
Qwen3-8B 单轮数学 RL（DAPOMath-17K）+ EAGLE-3	接受长度 / rollout加速 / E2E加速	2.28 / 1.63× / 1.50×	无投机解码基线 1.00×	E2E 1.50×
Qwen3-8B 单轮数学 RL + DFlash	接受长度 / rollout加速 / E2E加速	3.45 / 2.23× / 1.88×	无投机解码基线 1.00×	E2E 1.88×
Qwen3-8B 单轮数学 RL + DSpark	接受长度 / rollout加速 / E2E加速	3.63 / 2.18× / 1.83×	无投机解码基线 1.00×	E2E 1.83×
Qwen3.5-35B-A3B + DFlash	接受长度 / rollout加速 / E2E加速	4.58 / 1.50× / 1.46×	无投机解码基线 1.00×	E2E 1.46×
Nemotron-3.5-Lightning-30B-A3B + DSpark	接受长度 / rollout加速 / E2E加速	2.65 / 1.19× / 1.16×	无投机解码基线 1.00×	E2E 1.16×
Qwen3.5-122B-A10B + DFlash	接受长度 / rollout加速 / E2E加速	4.78 / 1.72× / 1.35×	无投机解码基线 1.00×	E2E 1.35×
GPT-OSS-120B + DFlash	接受长度 / rollout加速 / E2E加速	3.80 / 1.48× / 1.19×	无投机解码基线 1.00×	E2E 1.19×
CP分支注意力 vs SpecForge USP（EAGLE-3 TTT，最长20,480 tokens）	前向+反向延迟 / 每GPU峰值显存	packed zigzag，CP=2/4/8延迟提升2.9×/2.3×/1.5×，显存低2.7×	USP padded（batch padding至2.25×）	延迟最高2.9×
长上下文CP扩展（256K tokens，EAGLE-3 TTT注意力）	注意力延迟 / 并行效率 / 峰值显存	17.7s（CP=1）→ 2.35s（CP=8），94%并行效率，53.2→7.5GB	CP=1	7.5×延迟下降
TapChannel传输微基准（PP=4 fan-in，真实tap布局）	fan-in延迟 / 源stage干扰 / 接收stage争用	one-sided写27-39 GB/s，快4.5-8.5×，源stage噪声级（-0.1%），草稿stage+1.6%	host staging（所有rank慢83-88%）	4.5-8.5×
Qwen3-8B PP协同训练开销（前10个更新步平均）	更新时间开销 / tap等待 / E2E加速	EAGLE-3 +34.3%/0.37s/1.31×；DFlash +13.6%/0.57s/1.37×；DSpark +14.6%/0.54s/1.85×	无草稿训练（rollout 329.2s/步）	tap等待仅占优化时间1.5-2.2%
Qwen3-8B 多轮智能体（NeMo Gym Workplace Assistant，32K序列）	E2E训练加速	1.25-1.43×（接受长度单调升至约2.0-3.5）	无草稿基线	最高1.43×

局限与改进

作者明确承认的局限：一是稀疏MoE上收益缩水，Qwen3.5-122B-A10B接受长度4.78为全场最高但E2E加速只有1.35×，因为每次验证前向都会因投机token触发更多专家的稀疏路由计算；二是线性注意力模型上验证本身便宜，投机解码的相对收益天然更小（被列为未来工作）。EAGLE-3的TTT需要多个pass，带来34.3%的更新时间开销，远高于DFlash/DSpark的13.6%/14.6%。我的补充观察：其一，任务覆盖窄，只有数学（DAPOMath-17K/AIME2024）与一个工具调用环境（Workplace Assistant），代码、检索等rollout结构不同的任务未验证；其二，多轮场景rollout仅占步时55.8%，环境延迟构成加速上限；其三，草稿均从官方checkpoint起步，未探索从零协同训练或对分布外策略的适应速度；其四，TapChannel每源stage预分配slot，源数或特征维度大时草稿stage显存线性增长（未给规模上限）；其五，缺少与其他RL系统投机方案的横向对比。

独立分析的弱点

独立分析的弱点与改进方向：(1) MoE上加速受限——投机token的验证会激活更多专家，改进方向是草稿感知的路由（限制验证时专家激活数）或只在活跃专家上做投机；(2) EAGLE-3的34.3%更新开销来自TTT多pass模拟，可把TTT蒸馏为单pass近似，或按观测接受率动态调整pass数；(3) 多轮任务中环境延迟主导（E2E 1.25-1.43× vs rollout 1.75-2.23×），可把投机解码与异步工具执行、环境侧流水线正交组合；(4) 分支注意力合并依赖在线softmax数值稳定性，256K以上长度或低精度（bf16/fp8）下需监控精度退化，可引入分块重缩放；(5) TapChannel的slack依赖流水线气泡，在气泡极小的紧密调度（如zero-bubble）下暴露开销会放大，可自适应降级为错峰传输；(6) 每源stage预分配slot的显存随源数与特征宽度线性增长，可改为两级缓冲（显存+pinned内存）或按tap重要性选择性传输。

未来方向

作者提出的未来方向是把投机解码定制到稀疏MoE与线性注意力模型，因为当前收益在这两类模型受限。基于本文成果还可延伸：其一，接受率感知的联合训练——让 $\lambda$ 或草稿更新频率随观测接受率自适应，避免草稿欠拟合或拖慢策略；其二，TapChannel的mailbox机制可推广为PP下的通用跨stage旁路，服务于MTP模块共享、虚拟流水线等其他需要非相邻stage数据的场景；其三，与异步RL结合，草稿在稍旧策略的rollout上训练并容忍一步滞后，进一步隐藏tap等待；其四，分支注意力分解可推广到MLA/GQA等草稿头结构与fp8训练；其五，把slack模型 $E^{(m)}_{tap}$ 与1F1B/interleaved/zero-bubble等调度形式化结合，自动选择传输时机；其六，多轮场景中结合环境/工具延迟预测，把环境执行时间也纳入投机窗口；其七，在千卡、>256K上下文上验证zigzag-ring分支注意力能否维持94%并行效率。

复现评估

复现评估：论文声明代码已开源（'Code can be found here'），系统集成于开源的NVIDIA NeMo-RL框架。数据全部公开：DAPOMath-17K、AIME 2024、WorkBench/Workplace Assistant；草稿模型来自EAGLE-3、DFlash、DSpark官方checkpoint。算力门槛较高：Qwen3-8B实验用H100（TP/PP/CP=2/2/2，8卡），Qwen3.5-35B-A3B为8卡（TP/PP/CP/EP=2/2/2/8），30B/120B/122B级MoE在GB200上（配置至TP/PP/CP/EP=4/4/2/16）；个人或小团队复现8B级配置可行，全量复现需要数据中心级集群。难度评估：中等偏上——方法描述（Eq.1-6、通信量与slack公式、mailbox/stamp协议）相当完整，但涉及NCCL communicator、GPUDirect RDMA、CUDA IPC等底层工程，需熟悉NeMo-RL与Megatron式并行栈；建议先在CP=2小模型上验证分支注意力合并的数值一致性，再逐步加入PP与规模。

面向大规模长上下文RL后训练的投机解码在线草稿协同训练 Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training