通过消除重复计算共享提示，显著加速大回合长上下文强化学习训练。
DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts
- 设计双键值核，在一次内核调用中处理共享提示与序列响应，避免重复计算。
- 在32个序列、8K上下文下，策略更新速度提升1.63至2.09倍，显存利用率从36%升至76%。
- 适用于大规模模型和混合注意力机制，适合需要高效长序列训练的研究者。
现代强化学习后训练方法（如GRPO和DAPO）在包含P个标记的共享提示下采样N条长度为R的响应序列进行训练，但标准FlashAttention在前向和反向传播中将所有P个提示标记复制N次，造成计算与内存冗余。在大回合、长上下文训练中（N≥16，P≥8K），此冗余主导了策略更新开销。我们发现，在仅解码器模型中，因果掩码使每层的提示表示在各序列间不变，因此所有逐标记操作（归一化、投影、MLP）及注意力可仅处理一次提示——这一性质尚未在内核层面被利用。本文提出DualKV，首个针对强化学习训练的FlashAttention变体，通过（1）融合CUDA前向与反向内核，单次启动迭代两个不相交的键值区域——共享上下文与每序列响应；（2）在veRL中重构数据流水线，将N(P+R)个标记重排为每微批次P+NR个，使令牌减少因子ρ = N(P+R)/(P+NR)。DualKV在数学上等价于标准注意力，无近似误差。在Qwen3-8B GRPO训练中（8xH100，N=32，8K上下文），实现1.63–2.09倍策略更新加速，支持2倍更大微批次，显卡利用率从36%提升至76%。对DAPO同样取得2.47倍加速与77%显卡利用率。在30B MoE规模下（16xH100），相比需4路Ulysses序列并行以避免溢出的FlashAttention，DualKV实现3.82倍策略更新与3.38倍端到端步骤加速。该方法还扩展至头维度512的混合滑动/全局注意力（FA2不支持），并集成Ulysses序列并行，已在Gemma-4-31B GRPO 64K上下文中验证。
原文摘要 · Abstract (English)
Modern RL post-training methods such as GRPO and DAPO train on N response sequences of R tokens sampled from a shared prompt of P tokens, but standard FlashAttention replicates all P prompt tokens N times across both forward and backward passes -- duplicating compute and memory on identical hidden states. In large-rollout, long-context RL training (N>=16, P>=8K), this redundancy dominates the policy update cost. We observe that in decoder-only models, causal masking makes prompt representations invariant across sequences at every layer, so all per-token operations (norms, projections, MLP) and attention can process the prompt once -- a property not yet exploited at the kernel level for training. We propose DualKV, the first FlashAttention kernel variant that eliminates shared-prompt replication during RL training, via (1) fused CUDA forward and backward kernels that iterate over two disjoint KV regions -- shared context and per-sequence response -- in a single kernel launch, and (2) a data-pipeline redesign in veRL that repacks N(P+R) tokens into P+NR tokens per micro-batch, extending the token reduction from attention to the entire model by a factor rho = N(P+R)/(P+NR). DualKV is mathematically equivalent to standard attention and introduces no approximation. On Qwen3-8B GRPO training with 8xH100 GPUs (N=32, 8K-context), DualKV achieves 1.63--2.09x policy-update speedup, enables 2x larger micro-batches, and raises MFU from 36% to 76%. Similar gains hold for DAPO (2.47x speedup, 77% MFU). At 30B MoE scale on 16xH100, DualKV achieves 3.82x policy-update and 3.38x end-to-end step speedup over FlashAttention (which requires 4-way Ulysses sequence parallelism to avoid OOM). DualKV also extends to hybrid sliding/global attention with head dimension 512 (which FA2 does not support) and integrates with Ulysses sequence parallelism, demonstrated on Gemma-4-31B GRPO at 64K context.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。