DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts

• URL: http://arxiv.org/abs/2605.15422v1
• Date: Thu, 14 May 2026 21:11:32 GMT
• Title: DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts
• Abstract summary: Modern RL post-training methods such as GRPO and DAPO train on $N$ response sequences of $R$ tokens sampled from a shared prompt.<n>Standard FlashAttention replicates all $P$ prompt tokens $N times across both forward and backward passes.<n>We propose textDualAttention, the first Flash kernel variant eliminates shared-prompt replication during training.
• Score: 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).

Related papers

• Concurrent Image Understanding and Generation: Self-Correcting Coupled Markov Jump Processes [70.61868608402723]
  We introduce $textbfSelf-Correcting Coupled Markov Jump Processes (SC-CMJP)<n>In conjunction with SC-CMJP, we introduce $textttCO_texttt2textttJump$, a novel training-free single-pass sampler for joint multimodal geneneration.<n>For training and evaluation purposes, we have created and will release three large-scale joint multimodal generation corpora.
  arXiv  Detail & Related papers  (2026-07-14T18:39:29Z)
• Pruning Deep Neural Networks via the Marchenko--Pastur Distribution [0.0]
  We study a Marchenko--Pastur (MP) random-matrix approach to pruning deep neural networks with very small post-pruning fine-tuning budgets.<n>The main practical contribution is accuracy retention under short calibration and fine-tuning schedules.
  arXiv  Detail & Related papers  (2026-05-23T19:44:02Z)
• When Quantization Is Free: An int4 KV Cache That Outruns fp16 on Apple Silicon [0.0]
  KV-cache quantization is framed as a quality--latency trade-off.<n>We show it is emphinverted on Apple Silicon's unified memory.
  arXiv  Detail & Related papers  (2026-05-07T05:44:39Z)
• How Much Is One Recurrence Worth? Iso-Depth Scaling Laws for Looped Language Models [33.509836193149795]
  We measure how much one extra recurrence is worth to a looped (depth-recurrent) language model.<n>Our method applies to any looped LM and separates true loop improvements from token-budget gains.
  arXiv  Detail & Related papers  (2026-04-22T21:51:11Z)
• Copy-as-Decode: Grammar-Constrained Parallel Prefill for LLM Editing [2.6382975801439836]
  LLMs edit text and code by autoregressively regenerating the full output, even when most tokens appear verbatim in the input.<n>We study Copy-as-Decode, a decoding-layer mechanism that recasts edit generation as structured decoding over a two-primitive grammar.
  arXiv  Detail & Related papers  (2026-04-20T12:29:53Z)
• RACER: Retrieval-Augmented Contextual Rapid Speculative Decoding [80.12789199134511]
  Autoregressive decoding in Large Language Models (LLMs) generates one token per step, causing high inference latency.<n>We propose $textbfRACER, a lightweight and training-free method that integrates retrieved exact patterns with logit-driven future cues.<n>Experiments on Spec-Bench, HumanEval, and MGSM-ZH demonstrate that RACER consistently accelerates inference, achieving more than $2times$ speedup over autoregressive decoding.
  arXiv  Detail & Related papers  (2026-04-16T11:23:55Z)
• One Pass Streaming Algorithm for Super Long Token Attention Approximation in Sublinear Space [11.735802740426294]
  Attention computation takes both the time complexity of $O(n2)$ and the space complexity of $O(n2)$ simultaneously. We introduce a new algorithm that only reads one pass of data in a streaming fashion. Notably, our algorithm exhibits exceptional memory-efficient performance with super-long tokens.
  arXiv  Detail & Related papers  (2023-11-24T18:35:00Z)
• H$_2$O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models [110.06476624089679]
  We introduce a novel approach for implementing the KV cache which significantly reduces its memory footprint. Our approach is based on the observation that a small portion of tokens contributes most of the value when computing attention scores. We propose Heavy Hitter (H$$O), a KV cache eviction policy that dynamically retains a balance of recent and H$$ tokens.
  arXiv  Detail & Related papers  (2023-06-24T20:11:14Z)
• Near-Optimal Regret Bounds for Multi-batch Reinforcement Learning [54.806166861456035]
  We study the episodic reinforcement learning (RL) problem modeled by finite-horizon Markov Decision Processes (MDPs) with constraint on the number of batches. We design a computational efficient algorithm to achieve near-optimal regret of $tildeO(sqrtSAH3Kln (1/delta))$tildeO(cdot) hides logarithmic terms of $(S,A,H,K)$ in $K$ episodes. Our technical contribution are two-fold: 1) a near-optimal design scheme to explore
  arXiv  Detail & Related papers  (2022-10-15T09:22:22Z)
• Reward-Mixing MDPs with a Few Latent Contexts are Learnable [75.17357040707347]
  We consider episodic reinforcement learning in reward-mixing Markov decision processes (RMMDPs) Our goal is to learn a near-optimal policy that nearly maximizes the $H$ time-step cumulative rewards in such a model.
  arXiv  Detail & Related papers  (2022-10-05T22:52:00Z)
• On the Power of Multitask Representation Learning in Linear MDP [61.58929164172968]
  This paper presents analyses for the statistical benefit of multitask representation learning in linear Markov Decision Process (MDP) We first discover a emphLeast-Activated-Feature-Abundance (LAFA) criterion, denoted as $kappa$, with which we prove that a straightforward least-square algorithm learns a policy which is $tildeO(H2sqrtfrackappa mathcalC(Phi)2 kappa dNT+frackappa dn)
  arXiv  Detail & Related papers  (2021-06-15T11:21:06Z)
• Nearly Minimax Optimal Reinforcement Learning for Linear Mixture Markov Decision Processes [91.38793800392108]
  We study reinforcement learning with linear function approximation where the underlying transition probability kernel of the Markov decision process (MDP) is a linear mixture model. We propose a new, computationally efficient algorithm with linear function approximation named $textUCRL-VTR+$ for the aforementioned linear mixture MDPs. To the best of our knowledge, these are the first computationally efficient, nearly minimax optimal algorithms for RL with linear function approximation.
  arXiv  Detail & Related papers  (2020-12-15T18:56:46Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.