Fugu-MT 論文翻訳(概要): DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts

論文の概要: DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts

arxiv url: http://arxiv.org/abs/2605.15422v1
Date: Thu, 14 May 2026 21:11:32 GMT
ステータス: 翻訳完了
システム内更新日: 2026-05-18 21:22:26.107758
Title: DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts
Title（参考訳）: DualKV: 大きなロールアウトと長期コンテキストによる効率的なRLトレーニングのための共有プロンプトフラッシュアテンション
Authors: Jiading Gai, Shuai Zhang, Xiang Song, Bernie Wang, George Karypis,
Abstract要約: GRPOやDAPOのような現代のRLポストトレーニングメソッドは、共有プロンプトからサンプリングされた$R$トークンの$N$レスポンスシーケンスでトレーニングする。標準FlashAttentionは、すべての$P$プロンプトトークンを前方と後方の両方で$N回複製する。トレーニング中に共有プロンプトレプリケーションを排除した最初のFlashカーネルである textDualAttention を提案する。
参考スコア（独自算出の注目度）: 32.57334288518084
License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
Abstract: 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{\geq}16$, $P{\geq}8\text{K}$), 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 \textbf{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 $ρ= N(P{+}R)/(P{+}NR)$. DualKV is mathematically equivalent to standard attention and introduces no approximation. On Qwen3-8B GRPO training with 8$\times$H100 GPUs ($N{=}32$, 8K-context), DualKV achieves $1.63$--$2.09\times$ policy-update speedup, enables $2\times$ larger micro-batches, and raises MFU from $36\%$ to $76\%$. Similar gains hold for DAPO ($2.47\times$ speedup, $77\%$ MFU). At 30B MoE scale on 16$\times$H100, DualKV achieves $3.82\times$ policy-update and $3.38\times$ end-to-end step speedup over FlashAttention (which requires 4-way Ulysses sequence parallelism to avoid OOM).
Abstract（参考訳）: GRPOやDAPOといったモダンなRLポストトレーニングメソッドは、$R$トークンの共有プロンプトからサンプリングされた$R$トークンのレスポンスシーケンスをN$でトレーニングするが、標準のFlashAttentionは、すべての$P$プロンプトトークンを前方と後方の両方で複製する。大きなロールアウトで長期のRLトレーニング(N{\geq}16$, $P{\geq}8\text{K}$)では、この冗長性がポリシー更新コストを支配している。デコーダのみのモデルでは、因果マスキングは各レイヤのシーケンス間で即時表現を不変にするため、すべてのトーケン操作(ノーム、プロジェクション、MLP)と注意が、カーネルレベルでまだ利用されていないプロパティであるプロンプトを一度に処理できる。 In a single kernel launch, (2) a data-pipeline redesign in veRL that repacks $N(P{+}R)$ tokens into $P{+}NR$ tokens per micro-batch, extended the token reduction from attention to the model by a factor $ρ=N(P{+}R)/(P{+}NR)$$$$$。 DualKVは数学的には標準的な注意と等価であり、近似を導入しない。 Qwen3-8B 8$\times$H100 GPUs$N{=}32$, 8K-context)によるGRPOトレーニングでは、DualKVが1.63$-2.09\times$ Policy-update speedupを達成した。 DAPO(2.47\times$ speedup, 7.7\%$ MFU)も同様に上昇している。 16$\times$H100での30B MoEスケールでは、DualKVはポリシー更新で3.82ドル、FlashAttention上でのエンドツーエンドのステップスピードアップで3.38ドルを達成する(OOMを避けるには4ウェイのユリシーズシーケンス並列性が必要となる)。

論文の概要: DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts

関連論文リスト