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
Source: [https://arxiv.org/html/2605.15422](https://arxiv.org/html/2605.15422)
Jiading Gai1,\*Shuai Zhang1,\*Xiang Song2Bernie Wang1George Karypis3 1Amazon Web Services2Google3University of Minnesota \{jiadingg, shuaizs, yuyawang\}@amazon\.comxiangsx@google\.comkarypis@umn\.edu \*Equal contribution\.

###### Abstract

Modern RL post\-training methods such as GRPO and DAPO train onNNresponse sequences ofRRtokens sampled from a shared prompt ofPPtokens, but standard FlashAttention replicates allPPprompt tokensNNtimes across both forward and backward passes — duplicating compute and memory on identical hidden states\. In large\-rollout, long\-context RL training \(N≥16N\{\\geq\}16,P≥8​KP\{\\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 proposeDualKV, 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 repacksN​\(P\+R\)N\(P\{\+\}R\)tokens intoP\+N​RP\{\+\}NRtokens per micro\-batch, extending the token reduction from attention to the entire model by a factorρ=N​\(P\+R\)/\(P\+N​R\)\\rho=N\(P\{\+\}R\)/\(P\{\+\}NR\)\. DualKV is mathematically equivalent to standard attention and introduces no approximation\. On Qwen3\-8B GRPO training with 8×\\timesH100 GPUs \(N=32N\{=\}32, 8K\-context\), DualKV achieves1\.631\.63–2\.09×2\.09\\timespolicy\-update speedup, enables2×2\\timeslarger micro\-batches, and raises MFU from36%36\\%to76%76\\%\. Similar gains hold for DAPO \(2\.47×2\.47\\timesspeedup,77%77\\%MFU\)\. At 30B MoE scale on 16×\\timesH100, DualKV achieves3\.82×3\.82\\timespolicy\-update and3\.38×3\.38\\timesend\-to\-end step speedup over FlashAttention \(which requires 4\-way Ulysses sequence parallelism to avoid OOM\)\.

## 1Introduction

Modern RL post\-training methods such as GRPO\(Shao et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib20)\)and DAPO\(Yu et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib24)\)compute log\-probabilities and gradients forNNresponse sequences that share a prompt ofPPtokens\. With standard FlashAttention\-2 \(FA2\)\(Dao,[2024](https://arxiv.org/html/2605.15422#bib.bib3)\), the training micro\-batch containsNNsequences of lengthSi=P\+RiS\_\{i\}=P\+R\_\{i\}, repeating thePP\-token promptNNtimes\. Each forward and backward pass recomputesK,VK,Vactivations and their gradients over\(N−1\)×P\(N\{\-\}1\)\\times Predundant prompt tokens per layer\.

Recent work demonstrates that model accuracy continues to improve with increasingly large rollout factors\.Lightman et al\. \([2023](https://arxiv.org/html/2605.15422#bib.bib12)\)show that best\-of\-NNperformance on the MATH dataset scales log\-linearly up toN=1860N\{=\}1860\. GRPO and related methods similarly benefit from largeNNfor more accurate reward signal estimation \(e\.g\., DAPO usesN=16N\{=\}16\)\. However, the compute cost of the standardNN\-copy packing grows asO​\(N⋅S2\)O\(N\\cdot S^\{2\}\): each of theNNsequences independently recomputes attention over the fullPPprompt tokens, duplicating\(N−1\)×P\(N\{\-\}1\)\\times Pprompt tokens in both compute and memory at every layer\. As we show in Section[4\.2](https://arxiv.org/html/2605.15422#S4.SS2), eliminating this replication cuts policy\-update time by2×2\\times\.

Prompt KV sharing for inference has been addressed from multiple angles: paged attention\(Kwon et al\.,[2023](https://arxiv.org/html/2605.15422#bib.bib11)\)and prefix caching avoid redundant prompt KV storage via copy\-on\-write block tables, while bifurcated attention\(Athiwaratkun et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib1)\)decomposes attention into shared\-prompt and per\-sequence phases for parallel decoding\. However, these mechanisms target inference only, where shared prompt KV is read\-only; they have no training backward, no autograd\-compatible gradient accumulation for the shared prompt, and cannot be dropped into a RL policy\-update pipeline \(Appendix[I](https://arxiv.org/html/2605.15422#A9)\)\. Training raises three new challenges: \(1\) the backward pass hasNNsequences concurrently accumulating gradients into the shared KV buffer — a pattern absent from inference that requires both atomic accumulation for race\-free writes and fp32 accumulators with a final cast to match FA2’s per\-element precision; \(2\) autograd must aggregate gradients correctly across the context self\-attention and decoded\-attention calls that both touch the shared KV; and \(3\) the RL training framework’s data pipeline must be restructured to group same\-prompt responses per micro\-batch \(Appendix[B](https://arxiv.org/html/2605.15422#A2)\)\. DualKV addresses all three at the kernel and system level\.

We observe that in decoder\-only models, prompt hidden states are*identical*across allNNsequences at every layer: causal masking ensures each prompt token attends only to preceding prompt tokens, never to the differing response tokens that follow, so the prompt representations are independent of which response is generated\. Prefix Grouper\(Liu et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib17)\)applies this at the framework level but provides no kernel support — it still passes the fullNN\-replicated KV through standard FA2, retaining theO​\(N⋅P⋅d\)O\(N\\cdot P\\cdot d\)memory bottleneck that causes OOM and forces sequence parallelism\. DualKV contributes the first FlashAttention kernel variant that eliminates this replication, with fused forward and backward CUDA kernels that read shared KV from a single physical buffer and accumulate gradients fromNNconcurrent sequences via fp32 atomic writes\. By packing the micro\-batch with a single prompt copy and decomposing attention into context self\-attention \(computed once\) and a fused DualKV kernel for decoded attention, DualKV eliminates redundant prompt computation throughout the entire model — not just attention, but also norms, projections, and MLP layers — yielding both memory and compute savings without approximation\.

## 2Eliminating Prompt Replication in RL Training: The DualKV Design

Each GRPO training step processesNNresponse sequences of lengthSi=P\+RiS\_\{i\}=P\+R\_\{i\}that share a prompt ofPPtokens\. The step runs three passes before the optimizer step: \(1\)old\_log\_prob: forward throughπθold\\pi\_\{\\theta\_\{\\text\{old\}\}\}for the importance ratio; \(2\)ref\_log\_prob: forward throughπref\\pi\_\{\\text\{ref\}\}for the KL penalty; and \(3\)policy update: forward\+backward throughπθ\\pi\_\{\\theta\}producing the policy gradient\. All three passes see the same packed micro\-batch with\(N−1\)⋅P\(N\{\-\}1\)\\cdot Predundant prompt tokens, so DualKV accelerates all three\. The same applies to any RL method that packsNNsame\-prompt responses, e\.g\., DAPO\.

### 2\.1Baseline: Standard Packing with Replicated Prompts

Standard implementations \(e\.g\., veRL with FA2\) pack theNNsequences into a singleflash\_attn\_varlen\_funccall using cumulative sequence lengths:

\[P,R1\]⏟seq 1,\[P,R2\]⏟seq 2,…,\[P,RN\]⏟seq​N— total​Tstd=N​\(P\+R\)​tokens\\underbrace\{\[P,R\_\{1\}\]\}\_\{\\text\{seq 1\}\},\\;\\underbrace\{\[P,R\_\{2\}\]\}\_\{\\text\{seq 2\}\},\\;\\ldots,\\;\\underbrace\{\[P,R\_\{N\}\]\}\_\{\\text\{seq \}N\}\\quad\\text\{\-\-\- total \}T\_\{\\text\{std\}\}=N\(P\+R\)\\text\{ tokens\}\(1\)The prompt tokens are*replicated*NNtimes: each sequence carries its own copy of the prompt’s hidden states, QKV projections, and attention computation\. Every per\-token operation \(Norm, QKV projection, RoPE, MLP, output projection\) processes allN​\(P\+R\)N\(P\+R\)tokens\. For attention, FA2 computesNNindependent causal self\-attention passes, one per sequenceii:

O\(i\)=softmax​\(Q\(i\)​\(K\(i\)\)⊤d\+Mcausal\)​V\(i\),i=1,…,N,O^\{\(i\)\}\\;=\\;\\text\{softmax\}\\\!\\left\(\\frac\{Q^\{\(i\)\}\(K^\{\(i\)\}\)^\{\\top\}\}\{\\sqrt\{d\}\}\+M\_\{\\text\{causal\}\}\\right\)V^\{\(i\)\},\\qquad i=1,\\ldots,N,\(2\)whereQ\(i\),K\(i\),V\(i\)∈ℝ\(P\+Ri\)×H×dQ^\{\(i\)\},K^\{\(i\)\},V^\{\(i\)\}\\in\\mathbb\{R\}^\{\(P\+R\_\{i\}\)\\times H\\times d\}are sequenceii’s packed projections andMcausalM\_\{\\text\{causal\}\}is the per\-sequence causal mask\. FA2 computes Eq\.[2](https://arxiv.org/html/2605.15422#S2.E2)via tile\-wise online softmax\. FLOPs:O​\(N⋅Si2⋅H⋅d\)O\(N\\cdot S\_\{i\}^\{2\}\\cdot H\\cdot d\)\.

### 2\.2DualKV: Single\-Prompt Packing

Prompt invariance and redundancy\.In Eq\.[2](https://arxiv.org/html/2605.15422#S2.E2), split the output rows ofO\(i\)O^\{\(i\)\}into prompt\-query rowsOP\(i\)O\_\{P\}^\{\(i\)\}\(the firstPP\) and response\-query rowsORi\(i\)O\_\{R\_\{i\}\}^\{\(i\)\}\(the remainingRiR\_\{i\}\)\. For the prompt\-query rows, causal masking blocks attention to response keys, so only the prompt keys participate:

OP=softmax​\(QP​\(KP\)⊤d\+Mcausal,P\)​VP,O\_\{P\}\\;=\\;\\text\{softmax\}\\\!\\left\(\\frac\{Q\_\{P\}\(K\_\{P\}\)^\{\\top\}\}\{\\sqrt\{d\}\}\+M\_\{\\text\{causal\},P\}\\right\)V\_\{P\},\(3\)whereQP,KP,VPQ\_\{P\},K\_\{P\},V\_\{P\}are the prompt\-row projections\. Because allNNsequences share the same prompt tokens,QP,KP,VPQ\_\{P\},K\_\{P\},V\_\{P\}are identical acrossii\(Appendix[A\.2](https://arxiv.org/html/2605.15422#A1.SS2)\), soOPO\_\{P\}is the same for every sequence\. TheN−1N\-1prompt\-on\-prompt computations in the baseline are therefore redundant\.

For the response\-query rows, each response attends causally to*both*the shared prompt keys and its own response keys:

ORi\(i\)=softmax​\(QRi\(i\)​\[KP;KRi\(i\)\]⊤d\+Mcausal\)​\[VP;VRi\(i\)\],O\_\{R\_\{i\}\}^\{\(i\)\}\\;=\\;\\text\{softmax\}\\\!\\left\(\\frac\{Q\_\{R\_\{i\}\}^\{\(i\)\}\[K\_\{P\};\\,K\_\{R\_\{i\}\}^\{\(i\)\}\]^\{\\top\}\}\{\\sqrt\{d\}\}\+M\_\{\\text\{causal\}\}\\right\)\[V\_\{P\};\\,V\_\{R\_\{i\}\}^\{\(i\)\}\],\(4\)where\[⋅;⋅\]\[\\cdot\\,;\\,\\cdot\]denotes concatenation along the token axis\. This block is not redundant — each response has distinct queries — but it reuses the sameKP,VPK\_\{P\},V\_\{P\}for allii\.

DualKV packing\.DualKV packs the micro\-batch as a single shared prompt followed byNNper\-sequence responses:

\[P\]⏟prompt \(once\),\[R1\]⏟resp 1,…,\[RN\]⏟resp​N— total​Tdk=P\+N​R​tokens\\underbrace\{\[P\]\}\_\{\\text\{prompt \(once\)\}\},\\;\\underbrace\{\[R\_\{1\}\]\}\_\{\\text\{resp 1\}\},\\;\\ldots,\\;\\underbrace\{\[R\_\{N\}\]\}\_\{\\text\{resp \}N\}\\quad\\text\{\-\-\- total \}T\_\{\\text\{dk\}\}=P\+NR\\text\{ tokens\}\(5\)All per\-token operations \(norms, projections, RoPE, MLP, output projection\) now processP\+N​RP\+NRtokens instead ofN​\(P\+R\)N\(P\+R\)— saving\(N−1\)​P\(N\{\-\}1\)Ptokens throughout the entire model\.

Two\-call decomposition\.Eq\.[3](https://arxiv.org/html/2605.15422#S2.E3)and Eq\.[4](https://arxiv.org/html/2605.15422#S2.E4)can be computed separately to eliminate the redundancy while preserving exact attention:

- •Call 1 — context self\-attention\(flash\_attn\_varlen\_func, standard FA2\): computes Eq\.[3](https://arxiv.org/html/2605.15422#S2.E3)over a single copy of the prompt,*once*\. FLOPs:O​\(P2⋅H⋅d\)O\(P^\{2\}\\cdot H\\cdot d\)\.
- •Call 2 — decoded attention\(flash\_attn\_dualkv\_varlen\_func, DualKV kernel\): computes Eq\.[4](https://arxiv.org/html/2605.15422#S2.E4)for allNNresponses, attending to the sharedKP,VPK\_\{P\},V\_\{P\}\(from Call 1\) plus per\-sequenceKRi,VRiK\_\{R\_\{i\}\},V\_\{R\_\{i\}\}\. FLOPs:O​\(N⋅R⋅S⋅H⋅d\)O\(N\\cdot R\\cdot S\\cdot H\\cdot d\)\.

Why a new kernel\.Eq\.[4](https://arxiv.org/html/2605.15422#S2.E4)’s attention pattern — response queries reading from the concatenation\[KP;KRi\]\[K\_\{P\};K\_\{R\_\{i\}\}\]— cannot be consumed by standard FA2 without materializingNNcopies ofKP,VPK\_\{P\},V\_\{P\}\(reintroducing the replication we just eliminated\)\. The DualKV kernel iterates over the two physically disjoint KV regions \(KP,VPK\_\{P\},V\_\{P\}shared across the batch;KRi,VRiK\_\{R\_\{i\}\},V\_\{R\_\{i\}\}per\-sequence\) within a single launch\. The kernel’s interface and algorithms are detailed in Section[3](https://arxiv.org/html/2605.15422#S3)\.

Data pipeline\.This packing requires that allNNresponses from the same prompt are co\-located on the same GPU within the same micro\-batch\. veRL’s rollout engine already produces prompt\-contiguous sequences; DualKV preserves this ordering by skipping veRL’sbalance\_batchstep \(which reorders by sequence length\) and disabling intra\-epoch shuffle in the mini\-batch iterator\. The resulting pipeline preserves the exact mini\-batch gradient estimator \(Appendix[B](https://arxiv.org/html/2605.15422#A2)\)\. When the micro\-batch contains multiple prompt groups, DualKV handles each group independently with its own context KV and per\-groupcu\_seqlens, supporting heterogeneous prompt lengths within a single forward pass\.

Gradient correctness\.The sharedKP,VPK\_\{P\},V\_\{P\}receive gradients from both calls — autograd sums the two contributions, while the DualKV kernel internally accumulates theNNper\-sequence terms within Call 2:

∂ℒ∂KP=\(∂ℒ∂KP\)Call 1\+∑i=1N\(∂ℒ∂KP\)Call 2\(i\)\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{P\}\}\\;=\\;\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{P\}\}\\right\)\_\{\\\!\\text\{Call 1\}\}\\;\+\\;\\sum\_\{i=1\}^\{N\}\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{P\}\}\\right\)^\{\(i\)\}\_\{\\\!\\text\{Call 2\}\}\(6\)\(same structure forVPV\_\{P\}\)\. The first term is Call 1’s prompt self\-attention gradient; the∑i=1N\\sum\_\{i=1\}^\{N\}collects Call 2’s per\-response contributions, creating anNN\-way concurrent write into sharedKP,VPK\_\{P\},V\_\{P\}— a pattern with no analog in FA2’s backward, resolved at the kernel level in Section[3\.2](https://arxiv.org/html/2605.15422#S3.SS2)\. The total is mathematically identical to the FA2 baseline \(Appendix[A\.5](https://arxiv.org/html/2605.15422#A1.SS5)\)\.

When DualKV helps\.The token reduction ratioρ=N​\(P\+R\)/\(P\+N​R\)\\rho=N\(P\+R\)/\(P\+NR\)governs the per\-token operation speedup \(norms, projections, RoPE, MLP, output projection\) throughout the entire model; attention has additional structural savings \(Appendix[C](https://arxiv.org/html/2605.15422#A3)\)\.ρ\\rhois largest in two regimes:large rollout\(N≥16N\\geq 16, common in GRPO and DAPO; atN=32,P=16​KN\{=\}32,P\{=\}16\\text\{K\}the reduction reaches7\.2×7\.2\\times\) andlong prompts\(P≥16​KP\\geq 16\\text\{K\}, prevalent in agentic tasks\(Jimenez et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib9)\)and repository\-level code generation\(Liu et al\.,[2023](https://arxiv.org/html/2605.15422#bib.bib16)\); atP=64​K,N=16P\{=\}64\\text\{K\},N\{=\}16the reduction reaches14\.3×14\.3\\times\)\. In practice,ρ\\rhois bounded by the per\-GPU micro\-batch size \(Section[4\.2](https://arxiv.org/html/2605.15422#S4.SS2)\)\. See Appendix[C\.3](https://arxiv.org/html/2605.15422#A3.SS3)for the full analysis\.

## 3New Computational Primitive: The DualKV Kernel

This section describes the DualKV CUDA kernel that executes Call 2 of the two\-call decomposition \(Section[2\.2](https://arxiv.org/html/2605.15422#S2.SS2)\)\. Implemented as an extension of FA2’s codebase, the kernel accepts five input tensors — a decoded queryqq, the shared contextKc,VcK\_\{c\},V\_\{c\}\(single copy, stored once\), and per\-sequence varlen\-packed decodedKd,VdK\_\{d\},V\_\{d\}, plus acontext\_seqlen=P\{\}=Pparameter that shifts the causal mask so decoded queries attend to allPPcontext keys and their own preceding decoded keys\. We writeKc,VcK\_\{c\},V\_\{c\}andKd,VdK\_\{d\},V\_\{d\}\(following FlashAttention’s naming\) forKP,VPK\_\{P\},V\_\{P\}andKRi,VRiK\_\{R\_\{i\}\},V\_\{R\_\{i\}\}in Section[2\.2](https://arxiv.org/html/2605.15422#S2.SS2)\. Full Python signature in Appendix[H](https://arxiv.org/html/2605.15422#A8); pseudocode in Appendix[G](https://arxiv.org/html/2605.15422#A7); theoretical analysis in Appendix[C](https://arxiv.org/html/2605.15422#A3)\.

### 3\.1Forward Pass

The DualKV forward kernel extends FA2’s tiled online\-softmax algorithm \(tile sizeBNB\_\{N\}\) to iterate over two physically disjoint KV regions within a single kernel launch\. For each decoded queryqr\(i\)q\_\{r\}^\{\(i\)\}\(positionrrin sequenceii\), the kernel:\(1\)Context phase: Iterate over⌈P/BN⌉\\lceil P/B\_\{N\}\\rceiltiles ofkctx,vctxk\_\{\\text\{ctx\}\},v\_\{\\text\{ctx\}\}\(tile\-level slices ofKc,VcK\_\{c\},V\_\{c\}; lowercase denotes per\-tile data following FA2’s convention\), shared across allNNsequences in the grid\. Accumulate running softmax statistics\(m,ℓ,𝐨\)\(m,\\ell,\\mathbf\{o\}\)using the online\-softmax algorithm; \(2\)Decoded phase: Continue iteration over⌈Ri/BN⌉\\lceil R\_\{i\}/B\_\{N\}\\rceiltiles of sequenceii’s ownkdec\(i\),vdec\(i\)k\_\{\\text\{dec\}\}^\{\(i\)\},v\_\{\\text\{dec\}\}^\{\(i\)\}\. The running statistics carry over from the context phase, so the two\-phase iteration produces the same output as if\[kctx;kdec\(i\)\]\[k\_\{\\text\{ctx\}\};k\_\{\\text\{dec\}\}^\{\(i\)\}\]were stored contiguously; \(3\)Causal masking: All context positions \(j<Pj<P\) are unmasked for every decoded query; decoded positions are masked causally using thePP\-shifted logical positionP\+rP\+r\(thecontext\_seqlenparameter\)\.

The kernel decouples the physical block index from the logical key position: physical blocknnhas logical startn⋅BNn\\cdot B\_\{N\}whenn<⌈P/BN⌉n<\\lceil P/B\_\{N\}\\rceil\(context\) andP\+\(n−⌈P/BN⌉\)⋅BNP\+\(n\-\\lceil P/B\_\{N\}\\rceil\)\\cdot B\_\{N\}otherwise \(decoded\), so the causal mask operator sees the two physically disjoint regions as a single contiguous causal sequence\.

Per\-tile out\-of\-bounds \(OOB\) masking\.Because the context and decoded regions are tiled toBNB\_\{N\}independently, up to two tiles within a single kernel instance can be partial — the last context tile \(whenPmodBN≠0P\\bmod B\_\{N\}\\neq 0\) and the sequence’s last decoded tile \(whenRimodBN≠0R\_\{i\}\\bmod B\_\{N\}\\neq 0\) — rather than FA2’s single\-tile boundary\. The kernel bounds\-checks every tile load and applies the per\-tile OOB mask on every iteration, forgoing FA2’s fast\-path unmasked\-steps loop\. This cost is negligible compared to the\(N−1\)⋅⌈P/BN⌉\(N\{\-\}1\)\\cdot\\lceil P/B\_\{N\}\\rceilcontext\-tile passes eliminated\.

The kernel supports grouped\-query attention \(GQA\) and variable decoded lengths \(viacu\_seqlens\_q\)\. Its output matches Eq\.[4](https://arxiv.org/html/2605.15422#S2.E4)exactly; full pseudocode is in Algorithm[3](https://arxiv.org/html/2605.15422#alg3)\(Appendix[G](https://arxiv.org/html/2605.15422#A7)\), and the full derivation is in Appendix[A](https://arxiv.org/html/2605.15422#A1), §A\.3\.

### 3\.2Backward Pass with fp32 Atomic Accumulation

The backward pass computes gradients for all five inputs \(d​QdQ,d​KcdK\_\{c\},d​VcdV\_\{c\},d​KddK\_\{d\},d​VddV\_\{d\}\) via the main DualKV backward kernel \(Algorithm[1](https://arxiv.org/html/2605.15422#alg1)\) and three auxiliary launches \(described below\)\. Two pieces distinguish it from FA2’s backward: \(1\)Two\-region grid iterationmirroring the forward \(Section[3\.1](https://arxiv.org/html/2605.15422#S3.SS1)\): per\-block pointer dispatch selectsKc/VcK\_\{c\}/V\_\{c\}versusKd/VdK\_\{d\}/V\_\{d\}, logical key positions feed the causal mask, and per\-tile OOB masking handles two independent partial tiles\. \(2\)fp32 atomic accumulationfor the shared context gradients, resolving a concurrent\-write race and the precision loss from accumulatingNNpartial gradients in bf16, both unique to DualKV’s shared context buffer\.

Context gradient structure\.The loss gradient with respect to the shared context KV is a sum over allNNsequences in the batch:

∂L∂Kc=∑i=1N\(∂L∂Kc\)\(i\),∂L∂Vc=∑i=1N\(∂L∂Vc\)\(i\),\\frac\{\\partial L\}\{\\partial K\_\{c\}\}=\\sum\_\{i=1\}^\{N\}\\left\(\\frac\{\\partial L\}\{\\partial K\_\{c\}\}\\right\)^\{\(i\)\},\\quad\\frac\{\\partial L\}\{\\partial V\_\{c\}\}=\\sum\_\{i=1\}^\{N\}\\left\(\\frac\{\\partial L\}\{\\partial V\_\{c\}\}\\right\)^\{\(i\)\},\(7\)where\(∂L/∂Kc\)\(i\)\(\\partial L/\\partial K\_\{c\}\)^\{\(i\)\}is the contribution from sequenceii’s decoded queries attending to the context \(closed form in Appendix[A](https://arxiv.org/html/2605.15422#A1), §A\.4\)\. This additive structure drives the backward kernel design: theNNsummands are computed byNNseparate thread blocks, creating the concurrent\-write problem Algorithm[1](https://arxiv.org/html/2605.15422#alg1)resolves\.

Algorithm 1DualKV backward main kernel \(launched on grid\(nctx\+ndec,N,H\)\(n\_\{\\text\{ctx\}\}\+n\_\{\\text\{dec\}\},\\;N,\\;H\)\)1:

d​OdO,

OO,

QQ, savedsoftmax\_lse,

Kc,Vc,Kd,VdK\_\{c\},V\_\{c\},K\_\{d\},V\_\{d\}
2:fp32 scratch

d​Kc,acc,d​Vc,accdK\_\{c,\\text\{acc\}\},dV\_\{c,\\text\{acc\}\}\(shared across batch; zero\-initialized\);

d​QaccdQ\_\{\\text\{acc\}\}\(per batch, fp32\)

3:

D=rowsum​\(d​O⊙O\)D=\\text\{rowsum\}\(dO\\odot O\)\(fromdot\_do\_o, reused from FA2\)

4:foreach

\(n,b,h\)\(n,b,h\)in paralleldo⊳\\trianglerightgrid: \(K\-col block, batch, head\)

5:

nctx←⌈P/BN⌉n\_\{\\text\{ctx\}\}\\leftarrow\\lceil P/B\_\{N\}\\rceil;is\_context

←\(n<nctx\)\\leftarrow\(n<n\_\{\\text\{ctx\}\}\)
6:ifis\_contextthen

7:

\(Kn,Vn\)←\(Kc,Vc\)\[nBN:\(n\+1\)BN\]\(K\_\{n\},V\_\{n\}\)\\leftarrow\(K\_\{c\},V\_\{c\}\)\[nB\_\{N\}:\(n\{\+\}1\)B\_\{N\}\];

jbase←n​BNj\_\{\\text\{base\}\}\\leftarrow nB\_\{N\}
8:else

9:

\(Kn,Vn\)←\(Kd\(b\),Vd\(b\)\)\[\(n−nctx\)BN:…\]\(K\_\{n\},V\_\{n\}\)\\leftarrow\(K\_\{d\}^\{\(b\)\},V\_\{d\}^\{\(b\)\}\)\[\(n\{\-\}n\_\{\\text\{ctx\}\}\)B\_\{N\}:\\ldots\];

jbase←P\+\(n−nctx\)​BNj\_\{\\text\{base\}\}\\leftarrow P\+\(n\{\-\}n\_\{\\text\{ctx\}\}\)B\_\{N\}
10:endif

11:

a​c​cd​K,a​c​cd​V←0acc\_\{dK\},acc\_\{dV\}\\leftarrow 0in fp32

12:foreach Q\-row block

QmQ\_\{m\}in causal support of

KnK\_\{n\}do

13:Load

Qm,d​Om,softmax\_lse​\[m\],DmQ\_\{m\},dO\_\{m\},\\texttt\{softmax\\\_lse\}\[m\],D\_\{m\}from HBM to on\-chip SRAM

14:On chip, compute

P=exp⁡\(Qm​Kn⊤/d−lsem\)P=\\exp\(Q\_\{m\}K\_\{n\}^\{\\top\}/\\sqrt\{d\}\-\\texttt\{lse\}\_\{m\}\)with causal mask⊳\\trianglerightper\-tile OOB mask if partial

15:

accd​V\+=P⊤dOmacc\_\{dV\}\\mathrel\{\+\}=P^\{\\top\}dO\_\{m\}
16:

d​P←d​Om​Vn⊤dP\\leftarrow dO\_\{m\}V\_\{n\}^\{\\top\};

d​S←P⊙\(d​P−Dm\)dS\\leftarrow P\\odot\(dP\-D\_\{m\}\)
17:

accd​K\+=dS⊤Qm/dacc\_\{dK\}\\mathrel\{\+\}=dS^\{\\top\}Q\_\{m\}/\\sqrt\{d\}
18:atomicAdd\(

d​Qacc\(b\)​\[m\]dQ\_\{\\text\{acc\}\}^\{\(b\)\}\[m\],

d​S​Kn/ddS\\,K\_\{n\}/\\sqrt\{d\}\)⊳\\trianglerightFA2’s existingd​QdQpattern

19:endfor

20:ifis\_contextthen⊳\\trianglerightNNblocks \(one perbb\) race on same address

21:atomicAdd\(

d​Kc,acc​\[n​BN\]dK\_\{c,\\text\{acc\}\}\[nB\_\{N\}\],

a​c​cd​Kacc\_\{dK\}\)

22:atomicAdd\(

d​Vc,acc​\[n​BN\]dV\_\{c,\\text\{acc\}\}\[nB\_\{N\}\],

a​c​cd​Vacc\_\{dV\}\)

23:else⊳\\trianglerightDecoded block: exclusive per\-bbregion, no race

24:

d​Kd\(b\)​\[…\]←bf16​\(a​c​cd​K\)dK\_\{d\}^\{\(b\)\}\[\\ldots\]\\leftarrow\\text\{bf16\}\(acc\_\{dK\}\);

d​Vd\(b\)​\[…\]←bf16​\(a​c​cd​V\)dV\_\{d\}^\{\(b\)\}\[\\ldots\]\\leftarrow\\text\{bf16\}\(acc\_\{dV\}\)
25:endif

26:endfor

Concurrent\-write problem \(race on line[21](https://arxiv.org/html/2605.15422#alg1.l21)\)\.In a standard FA2 backward, the grid is\(n,b,h\)\(n,b,h\)and each sequenceiiholds a privateKc\(i\),Vc\(i\)K\_\{c\}^\{\(i\)\},V\_\{c\}^\{\(i\)\}, so the batch index maps each block to a separated​Kc\(i\)dK\_\{c\}^\{\(i\)\}region — each block can cast its fp32 accumulator directly to bf16 and write with no contention\. DualKV collapses this batch dimension for the shared context: the singled​KcdK\_\{c\}andd​VcdV\_\{c\}buffer has no batch offset, yetNNthread blocks — one per decoded sequence — concurrently reach Algorithm[1](https://arxiv.org/html/2605.15422#alg1)’s context epilogue\. A standard direct\-write epilogue would produce a data race\. Resolving the race via bf16atomicAddwould compoundN−1N\{\-\}1bf16 rounding steps on top of the initial cast, turning the cross\-sequence reduction itself into a half\-precision accumulation\.

fp32 atomic accumulation\.We resolve this with a two\-stage design\. Stage 1 \(Algorithm[1](https://arxiv.org/html/2605.15422#alg1)\): context blocksatomicAddinto an fp32 scratch buffer; decoded blocks bypass the scratch and write bf16 directly \(they own disjointd​Kd/d​VddK\_\{d\}/dV\_\{d\}regions\)\. Stage 2 \(Algorithm[2](https://arxiv.org/html/2605.15422#alg2)\): after all Stage 1 blocks complete, a separate kernel casts the fp32 scratch to bf16\. This reproduces FA2’s single\-cast precision profile despite theNN\-way concurrent write: allNNcontext contributions are summed in fp32, and only the final value is cast to bf16\.

Algorithm 2convert\_dkv\_context: fp32\-to\-bf16 cast, launched on grid\(nctx,1,H\)\(n\_\{\\text\{ctx\}\},\\;1,\\;H\)after Algorithm[1](https://arxiv.org/html/2605.15422#alg1)completes \(batch dimension collapsed because context is shared\)1:fp32 scratch

d​Kc,acc,d​Vc,accdK\_\{c,\\text\{acc\}\},dV\_\{c,\\text\{acc\}\}\(aggregated over

NNsequences by Algorithm[1](https://arxiv.org/html/2605.15422#alg1)\)

2:foreach

\(n,h\)\(n,h\)in parallel, and each

\(token,dim\)\(\\text\{token\},\\text\{dim\}\)in tiledo

3:

d​Kc​\[n​BN\+token,dim\]←bf16​\(d​Kc,acc​\[n​BN\+token,dim\]\)dK\_\{c\}\[nB\_\{N\}\{\+\}\\text\{token\},\\text\{dim\}\]\\leftarrow\\text\{bf16\}\(dK\_\{c,\\text\{acc\}\}\[nB\_\{N\}\{\+\}\\text\{token\},\\text\{dim\}\]\)
4:

d​Vc​\[n​BN\+token,dim\]←bf16​\(d​Vc,acc​\[n​BN\+token,dim\]\)dV\_\{c\}\[nB\_\{N\}\{\+\}\\text\{token\},\\text\{dim\}\]\\leftarrow\\text\{bf16\}\(dV\_\{c,\\text\{acc\}\}\[nB\_\{N\}\{\+\}\\text\{token\},\\text\{dim\}\]\)
5:endfor

Analogy to FA2’sd​QdQpipeline\.FA2 uses an fp32\-then\-cast pipeline ford​QdQunder K\-block\-parallel scheduling: the main backward kernel atomic\-adds intod​QaccdQ\_\{\\text\{acc\}\}\(Algorithm[1](https://arxiv.org/html/2605.15422#alg1)’sd​QdQatomicAddline\), and a separateconvert\_dqkernel castsd​QaccdQ\_\{\\text\{acc\}\}to bf16\. DualKV replicates this pipeline ford​Kc,d​VcdK\_\{c\},dV\_\{c\}— now indexed by*batch element*\(not K\-block\) — with Algorithm[1](https://arxiv.org/html/2605.15422#alg1)’s context\-epilogue atomicAdd and the separateconvert\_dkv\_contextkernel \(Algorithm[2](https://arxiv.org/html/2605.15422#alg2)\) casting to bf16\. The fp32 accumulator keeps the context gradient within standard fp32 rounding of the FA2 baseline; formal equivalence is proven in Section[A\.5](https://arxiv.org/html/2605.15422#A1.SS5)\.

## 4Experiments

We first benchmark the DualKV attention kernel in isolation \(Section[4\.1](https://arxiv.org/html/2605.15422#S4.SS1)\), then evaluate end\-to\-end DAPO and GRPO training at H100 scale on LongReason \(Section[4\.2](https://arxiv.org/html/2605.15422#S4.SS2)\) and on multi\-node MoE \(Section[4\.3](https://arxiv.org/html/2605.15422#S4.SS3)\)\. We also study the memory scaling with context length \(Section[4\.4](https://arxiv.org/html/2605.15422#S4.SS4)\)\. Appendix[E\.1](https://arxiv.org/html/2605.15422#A5.SS1)confirms DualKV benefits extend to the short\-prompt regime \(GSM8K, A100\)\. Our primary baseline is FA2 in veRL’s FSDP worker\(LinkedIn AI,[2025](https://arxiv.org/html/2605.15422#bib.bib14)\); we also include FA3 in Section[4\.2](https://arxiv.org/html/2605.15422#S4.SS2)\.

### 4\.1Kernel\-Level Speedup: Forward, Backward, and Combined

Table[1](https://arxiv.org/html/2605.15422#S4.T1)reports separate forward\-only, backward\-only, and combined wall\-clock times on a single NVIDIA A100 GPU with Qwen3\-8B model dimensions \(H=32H\{=\}32,Hk=8H\_\{k\}\{=\}8,d=128d\{=\}128, GQA4:14\{:\}1\), fp16 with gradient computation enabled\. The backward pass has different compute characteristics than the forward — it computes gradients∂ℒ/∂Q\\partial\\mathcal\{L\}/\\partial Q,∂ℒ/∂K\\partial\\mathcal\{L\}/\\partial K,∂ℒ/∂V\\partial\\mathcal\{L\}/\\partial Vand involves the gradient accumulation for sharedKc,VcK\_\{c\},V\_\{c\}described in Section[3\.2](https://arxiv.org/html/2605.15422#S3.SS2)\(Appendix[A](https://arxiv.org/html/2605.15422#A1), Eqs\.[21](https://arxiv.org/html/2605.15422#A1.E21)–[22](https://arxiv.org/html/2605.15422#A1.E22)\)\. DualKV gradients match FA2 to half\-precision rounding \(verified viatorch\.allclosewithatol=10−3\\text\{atol\}\{=\}10^\{\-3\},rtol=10−3\\text\{rtol\}\{=\}10^\{\-3\}\)\.

Table 1:FA2 vs\. DualKV wall\-clock time \(ms\) and peak memory across varying rollout factorNNand prompt lengthPP; response lengthR=2048R\{=\}2048throughout\. OOM = out of memory\.FA2 \(ms\)DualKV \(ms\)SpeedupMemNNPPfwdbwdf\+bfwdbwdf\+bfwdbwdf\+b↓\\downarrow284K43\.6128\.6172\.227\.177\.1104\.21\.61×\\times1\.67×\\times1\.65×\\times63%2816K381\.31059\.61441\.098\.7273\.0371\.73\.86×\\times3\.88×\\times3\.88×\\times85%1632K771\.52111\.82883\.4139\.5386\.2525\.85\.53×\\times5\.47×\\times5\.48×\\times86%1664KOOM363\.3994\.51357\.8∞\\infty—Combined fwd\+bwd speedup ranges from1\.65×1\.65\\times\(P=4​KP\{=\}4\\text\{K\}\) to5\.48×5\.48\\times\(P=32​KP\{=\}32\\text\{K\}\), with up to86%86\\%peak memory reduction; atP≥32​KP\{\\geq\}32\\text\{K\}standard FA2 OOMs entirely\. Backward speedup tracks forward within3%3\\%as expected: the attention backward computes the sameQ​K⊤QK^\{\\top\}products as the forward\. The backward\-to\-forward ratio \(∼2\.7×\{\\sim\}2\.7\\times\) is stable across all configurations for both FA2 and DualKV, confirming that DualKV’s two\-region backward kernel adds no structural overhead beyond attention’s intrinsic backward cost\. Gradient accumulation for sharedKc,VcK\_\{c\},V\_\{c\}\(Section[3\.2](https://arxiv.org/html/2605.15422#S3.SS2)\) adds no measurable overhead — fp32 atomic writes and the follow\-up bf16 cast are bandwidth\-bound on context positions only and land inside the backward’s GEMM timing\.

Table 2:Single\-layer fwd\+bwd: DualKV vs\. PG vs\. FA2 \(Qwen3\-8B, H100\)\. Full sweep in Appendix[E\.4](https://arxiv.org/html/2605.15422#A5.SS4)\.Comparison with Prefix Grouper\.Table[2](https://arxiv.org/html/2605.15422#S4.T2)benchmarks against Prefix Grouper \(PG\)\(Liu et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib17)\), a framework\-level shared\-prompt optimization for training \(no kernel\-level deduplication\) — FA2 still receivesNN\-replicated KV on padded tensors, so projections, MLP, and norms process the fullNN\-copy prompt\. DualKV is22–4×4\\timesfaster and uses33–5×5\\timesless memory than PG across all configurations where both run\. At high micro\-batch \(P=5​KP\{=\}5\\text\{K\}, mb=32\), PG’s padded representation makes it even slower than FA2 varlen\. AtP≥32​KP\{\\geq\}32\\text\{K\}, both FA2 and PG OOM while DualKV continues to scale\. PG only optimizes attention \(∼30%\{\\sim\}30\\%of layer FLOPs\), while DualKV eliminates redundancy across all operations via packed varlen representation\.

### 4\.2End\-to\-End GRPO and DAPO Training on LongReason

To validate DualKV in a higher\-ρ\\rhoregime with longer prompts and larger rollout factors, we evaluate on LongReason\(Ling et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib13)\)— a long\-context mathematical reasoning dataset with prompts up to 8,192 tokens and response lengths up to 2,048 tokens — using both GRPO and DAPO\(Yu et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib24)\)\(which removes the KL reference pass\)\.

![Refer to caption](https://arxiv.org/html/2605.15422v1/x1.png)Figure 1:Per\-phase step breakdown\.Setup\.We train Qwen3\-8B on a singlep5\.48xlargeinstance \(8×\\timesH100\-SXM5\-80GB\) with FSDP2 in BF16 and gradient checkpointing\. Rollout generation uses vLLM with tensor parallelism=2=2andN=32N\{=\}32responses per prompt\. Training usestrain\_batch\_size=128\{\}=128andppo\_mini\_batch\_size=64\{\}=64\(both fixed across all configurations\), and we sweepppo\_micro\_batch\_size\_per\_gpu∈\{4,8\}\\in\\\{4,8\\\}\. Each configuration runs for 25 training steps with validation every 2 steps\. Although the rollout factor isN=32N\{=\}32, the token reduction ratioρ\\rho\(Eq\.[35](https://arxiv.org/html/2605.15422#A3.E35)\) is bounded by the*micro\-batch*size — only sequences co\-located in the same forward/backward pass share a prompt\. WithP=8192P\{=\}8192,R=2048R\{=\}2048:ρ=2\.5×\\rho=2\.5\\timesat mb=4 andρ=3\.3×\\rho=3\.3\\timesat mb=8, setting upper bounds on the measured policy\-update speedup\.

Comparisons\.We compare four configurations \(mb= per\-GPU micro\-batch size\): FA2 \(mb=4\), the largest micro\-batch FA2 can handle \(mb=8 OOMs\); FA3 \(mb=4\), evaluating the Hopper\-specific kernel; DualKV \(mb=4\), an apple\-to\-apple comparison; and DualKV \(mb=8\), a configuration FA2 cannot run\.

Results\.Table[3](https://arxiv.org/html/2605.15422#S4.T3)and Figure[2](https://arxiv.org/html/2605.15422#S4.F2)summarize both experiments\. All configs track identical training accuracy curves \(Figure[2](https://arxiv.org/html/2605.15422#S4.F2)A\), confirming DualKV introduces no convergence degradation for either method \(formal proof in Appendix[B](https://arxiv.org/html/2605.15422#A2)\)\.

![Refer to caption](https://arxiv.org/html/2605.15422v1/x2.png)![Refer to caption](https://arxiv.org/html/2605.15422v1/x3.png)

Figure 2:LongReason training \(Qwen3\-8B, 8×\\timesH100\), four configurations\.Top row: GRPO \(N=32N\{=\}32, steps 1–24\)\.Bottom row: DAPO \(N=32N\{=\}32, steps 1–25\)\. Panels:\(A\)Training reward — all configs track identically\.\(B\)Peak memory\.\(C\)Policy\-update latency\.\(D\)Policy\-update MFU\. GRPO: DualKV mb=8 achieves2\.09×2\.09\\timesspeedup,75\.8%75\.8\\%MFU\. DAPO:2\.47×2\.47\\timesspeedup,77\.4%77\.4\\%MFU\.Speedups, memory, and MFU\.At mb=4, DualKV cuts policy\-update time by1\.63×1\.63\\times\(GRPO\) and1\.91×1\.91\\times\(DAPO\), while reducing peak memory from106106GB to8181GB — a 25 GB reduction that brings the workload within the H100’s 80 GB physical HBM\. FA2’s 106 GB allocation exceeds HBM capacity via memory oversubscription, degrading throughput\. Doubling to mb=8 — which FA2 OOMs on — extends the speedup to2\.09×2\.09\\times\(GRPO\) and2\.47×2\.47\\times\(DAPO\), halving the gradient\-accumulation passes per mini\-batch \(2→12\\to 1\)\. Policy\-update MFU rises from36%36\\%/31%31\\%\(FA2\) to76%76\\%/77%77\\%\(DualKV mb=8\) for GRPO/DAPO respectively — a2\.12\.1–2\.5×2\.5\\timesutilization gain\. DAPO’s higher speedups reflect the absence of the ref phase: the policy update dominates a larger fraction of step time, so DualKV’s per\-phase compression translates more directly to wall\-clock gains\.

Table 3:End\-to\-end LongReason results\.Step= total training\-step time \(s\)\.MFU= policy\-update Model FLOPs Utilization\. Parentheses = speedup vs FA2 mb=4\. Per\-phase breakdown is in Figure[1](https://arxiv.org/html/2605.15422#S4.F1)\.†PyTorch allocator reservation; actual HBM use stays within 80 GB\.

End\-to\-end step speedup\.Total training step time drops1\.64×1\.64\\times/1\.82×1\.82\\times\(GRPO/DAPO\) at mb=8\. The remaining gap versus the per\-phase speedups is explained by rollout generation \(∼240\{\\sim\}240s/step, unaffected by DualKV\), whose share of step time grows as training phases accelerate — DualKV shifts the bottleneck from the policy update to rollout generation \(Figure[1](https://arxiv.org/html/2605.15422#S4.F1)\)\. Full per\-phase DAPO breakdowns are in Appendix[E\.2](https://arxiv.org/html/2605.15422#A5.SS2)\.

### 4\.3Multi\-Node RL Training with MoE Models

![Refer to caption](https://arxiv.org/html/2605.15422v1/x4.png)Figure 3:Per\-step training accuracy\.To validate DualKV at production scale with cross\-node communication and MoE expert routing, we train Qwen3\-30B\-A3B\(Yang et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib23)\)with GRPO on twop5\.48xlargenodes \(16×\\timesH100 GPUs, Elastic Fabric Adapter\): FSDP2 \+ BF16 \+ gradient checkpointing; rollout via vLLM \(TP=2\{=\}2,N=32N\{=\}32\)\. All runs:train\_batch\_size=128\{=\}128,ppo\_mini\_batch\_size=64\{=\}64,ppo\_micro\_batch\_size\_per\_gpu=8\{=\}8,Pmax=8192P\_\{\\max\}\{=\}8192,Rmax=2048R\_\{\\max\}\{=\}2048\. Timing is mean over 29 steps\. We compare three configurations: DualKV \(SP=1\), no sequence parallelism; FA2 \(SP=2\), which OOMs during the first policy update; and FA2 \(SP=4\), the minimum SP that avoids OOM\.

Results\.Table[4](https://arxiv.org/html/2605.15422#S4.T4)summarizes the three configurations\. DualKV delivers a3\.82×3\.82\\timespolicy\-update speedup and3\.38×3\.38\\timestotal step speedup over FA2 SP=4, without requiring sequence parallelism\. Both configs track identical training accuracy curves \(Figure[3](https://arxiv.org/html/2605.15422#S4.F3)\), confirming DualKV introduces no convergence degradation at multi\-node training scale \(formal proof in Appendix[B](https://arxiv.org/html/2605.15422#A2)\)\.

Table 4:Wall\-clock values are means over 29 steps; parenthesized values are speedups relative to FA2 SP=4\.†PyTorch allocator reservation; actual HBM use stays within 80 GB\.

Speedups at multi\-node MoE scale\.DualKV delivers a3\.82×3\.82\\timespolicy\-update speedup and3\.38×3\.38\\timestotal\-step speedup \(5284 s→\\to1564 s, throughput 293→\\to991 tokens/s\)\. The larger gains versus the 8B single\-node range \(1\.631\.63–2\.09×2\.09\\times\) arise because \(i\) MoE expert GEMMs amplify the per\-token savings from DualKV’s\(N−1\)⋅P\(N\{\-\}1\)\\cdot Preduction, and \(ii\) DualKV eliminates the all\-to\-all communication that FA2 with SP=4 pays on every attention call\. Over 29 steps, DualKV completes in 12\.6 hours vs\. 42\.6 hours for FA2 SP=4 on the same 16\-GPU cluster \(∼$​6\{\\sim\}\\mathdollar 6K saved at AWSp5\.48xlargeon\-demand list price\)\. Figure[9](https://arxiv.org/html/2605.15422#A6.F9)\(Appendix[F](https://arxiv.org/html/2605.15422#A6)\) visualizes the phase\-wise breakdown\.

Eliminating sequence parallelism\.DualKV’s memory reduction enables training at SP=1, whereas FA2 OOMs at SP=2 and requires SP=4\. This translates memory savings into additional speedup: DualKV avoids the per\-call all\-to\-all communication that SP imposes, which is why end\-to\-end gains \(3\.38×3\.38\\times\) exceed what token dedup alone would predict\.

### 4\.4Memory Scaling with Context Length

To characterize how DualKV’s advantage scales with prompt length, we sweepP∈\{8​K–​96​K\}P\\in\\\{8\\text\{K\}\\text\{\-\-\}96\\text\{K\}\\\}andmb∈\{4,8,16\}\\text\{mb\}\\in\\\{4,8,16\\\}on Llama\-3\.1\-8B \(16×\\timesH100\)\. This experiment also validates that DualKV generalizes across model families \(Llama here versus Qwen3 in Sections[4\.2](https://arxiv.org/html/2605.15422#S4.SS2)–[4\.3](https://arxiv.org/html/2605.15422#S4.SS3)\)\.

![Refer to caption](https://arxiv.org/html/2605.15422v1/x5.png)Figure 4:Memory scaling: DualKV vs\. FA2\.Memory decouples from micro\-batch size\.DualKV’s memory followsM0\+cP⋅P\+cmb​R⋅mb⋅RM\_\{0\}\+c\_\{P\}\\cdot P\+c\_\{\\text\{mb\}R\}\\cdot\\text\{mb\}\\cdot R: the prompt cost is paid once regardless of micro\-batch size, so the three mb curves cluster tightly \(Figure[4](https://arxiv.org/html/2605.15422#S4.F4)\)\. FA2’s memory scales asM0\+c⋅mb⋅\(P\+R\)M\_\{0\}\+c\\cdot\\text\{mb\}\\cdot\(P\{\+\}R\), diverging rapidly — atP=96​KP\{=\}96\\text\{K\}, increasing mb from 4 to 16 adds only∼4\{\\sim\}4GB under DualKV \(<5%<5\\%\), whereas FA2 would require225→775225\\to 775GB — nearly10×10\\timesthe physical capacity of an H100\. Practitioners can freely increase mb to maximize throughput without memory penalty in the long\-prompt regime\.

## 5Related Work

Efficient attention for long sequences\.FlashAttention\(Dao et al\.,[2022](https://arxiv.org/html/2605.15422#bib.bib4); Dao,[2024](https://arxiv.org/html/2605.15422#bib.bib3); Shah et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib19)\)reduces attention memory fromO​\(S2\)O\(S^\{2\}\)toO​\(S\)O\(S\)via tiling and online softmax; Ring attention\(Liu et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib15)\), Ulysses\(Jacobs et al\.,[2023](https://arxiv.org/html/2605.15422#bib.bib8)\), Megatron\-SP\(Korthikanti et al\.,[2023](https://arxiv.org/html/2605.15422#bib.bib10)\), and Context Parallelism\(Dubey et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib5)\)distribute sequences across GPUs via sequence parallelism\. All are orthogonal to DualKV’s within\-GPU shared\-prompt deduplication and compose with it\.

Shared\-prompt optimization\.Paged attention\(Kwon et al\.,[2023](https://arxiv.org/html/2605.15422#bib.bib11)\), RadixAttention\(Zheng et al\.,[2023](https://arxiv.org/html/2605.15422#bib.bib25)\), and bifurcated attention\(Athiwaratkun et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib1)\)share prompt KV for inference only\. Prefix Grouper\(Liu et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib17)\)is a framework\-level data reorganization that applies the shared\-prompt decomposition to RL policy updates but replicatesK,VK,VNNtimes before the attention call, retainingO​\(N⋅P⋅d\)O\(N\\cdot P\\cdot d\)memory at the kernel level\. DualKV contributes the kernel\-level primitive that eliminates this replication: a FlashAttention variant that reads shared KV from a single physical buffer and accumulates gradients fromNNconcurrent sequences via fp32 atomic writes, fully realizing the decomposition’s memory and compute savings for training \(Appendix[E\.4](https://arxiv.org/html/2605.15422#A5.SS4)shows DualKV is 2–4×\\timesfaster and uses 3–5×\\timesless memory than Prefix Grouper\)\.

RL training systems\.veRL\(Sheng et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib21)\), OpenRLHF\(Hu et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib6)\), and TRL\(von Werra et al\.,[2020](https://arxiv.org/html/2605.15422#bib.bib22)\)implement PPO\(Schulman et al\.,[2017](https://arxiv.org/html/2605.15422#bib.bib18)\)/GRPO\(Shao et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib20)\)/REINFORCE\+\+\(Hu et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib7)\)/DAPO\(Yu et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib24)\)policy updates by packingNNsequences with replicated prompts, without exploiting shared\-prompt structure\. DualKV is a drop\-in replacement for the attention layer in these systems, requiring only data pipeline changes to group same\-prompt responses\.

## 6Conclusion

DualKV eliminates the redundant prompt replication in RL policy updates by introducing a two\-region FlashAttention kernel that reads shared KV from a single physical buffer, paired with a data\-pipeline redesign that extends the token reduction to the entire model with provable gradient equivalence\. Across dense and MoE models, DualKV achieves22–3\.8×3\.8\\timespolicy\-update speedup, raises MFU from36%36\\%to76%76\\%, and eliminates the need for sequence parallelism at long contexts\. DualKV shifts the RL training bottleneck from the policy update to rollout generation\. The current design assumes a single shared prefix per micro\-batch group and benefits scale withP/RP/R\. Generalizing the kernel to arbitrary tree structures with shared sub\-graphs \(e\.g\., multi\-turn conversations with partially overlapping prefixes or tree\-of\-thought search\) would lift the single\-prefix constraint\. Future directions also include composing DualKV with Ulysses SP for ultra\-long contexts and porting to FA3/Hopper and FA4/Blackwell backends\.

## References

- Athiwaratkun et al\. \[2024\]Ben Athiwaratkun, Sujan Kumar Gonugondla, Sanjay Krishna Gouda, Haifeng Qian, Hantian Ding, Qing Sun, Jun Wang, Jiacheng Guo, Liangfu Chen, Parminder Bhatia, Ramesh Nallapati, Sudipta Sengupta, and Bing Xiang\.Bifurcated attention: Accelerating massively parallel decoding with shared prefixes in LLMs\.*arXiv preprint arXiv:2403\.08845*, 2024\.
- Cobbe et al\. \[2021\]Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman\.Training verifiers to solve math word problems\.*arXiv preprint arXiv:2110\.14168*, 2021\.
- Dao \[2024\]Tri Dao\.FlashAttention\-2: Faster attention with better parallelism and work partitioning\.In*International Conference on Learning Representations*, 2024\.
- Dao et al\. \[2022\]Tri Dao, Daniel Y\. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré\.FlashAttention: Fast and memory\-efficient exact attention with IO\-awareness\.In*Advances in Neural Information Processing Systems*, 2022\.
- Dubey et al\. \[2024\]Abhimanyu Dubey et al\.The Llama 3 herd of models\.*arXiv preprint arXiv:2407\.21783*, 2024\.
- Hu et al\. \[2024\]Jian Hu, Xibin Wu, Wei Shen, Jason Klein Liu, et al\.OpenRLHF: An easy\-to\-use, scalable and high\-performance RLHF framework\.*arXiv preprint arXiv:2405\.11143*, 2024\.
- Hu et al\. \[2025\]Jian Hu, Jason Klein Liu, Haotian Xu, and Wei Shen\.REINFORCE\+\+: Stabilizing critic\-free policy optimization with global advantage normalization\.*arXiv preprint arXiv:2501\.03262*, 2025\.
- Jacobs et al\. \[2023\]Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He\.DeepSpeed Ulysses: System optimizations for enabling training of extreme long sequence transformer models\.*arXiv preprint arXiv:2309\.14509*, 2023\.
- Jimenez et al\. \[2024\]Carlos E\. Jimenez, John Yang, Alexander Wettig, Shunyu Yao, Kexin Pei, Ofir Press, and Karthik Narasimhan\.SWE\-bench: Can language models resolve real\-world GitHub issues?In*International Conference on Learning Representations*, 2024\.
- Korthikanti et al\. \[2023\]Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro\.Reducing activation recomputation in large transformer models\.*Proceedings of Machine Learning and Systems \(MLSys\)*, 5, 2023\.
- Kwon et al\. \[2023\]Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica\.Efficient memory management for large language model serving with PagedAttention\.In*Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles*, pages 611–626, 2023\.
- Lightman et al\. \[2023\]Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe\.Let’s verify step by step\.*arXiv preprint arXiv:2305\.20050*, 2023\.
- Ling et al\. \[2025\]Zhan Ling, Kang Liu, Kai Yan, Yifan Yang, Weijian Lin, Ting\-Han Fan, Lingfeng Shen, Zhengyin Du, and Jiecao Chen\.LongReason: A synthetic long\-context reasoning benchmark via context expansion\.*arXiv preprint arXiv:2501\.15089*, 2025\.
- LinkedIn AI \[2025\]LinkedIn AI\.Post\-training gpt\-oss for agentic reasoning with reinforcement learning\.Hugging Face blog,[https://huggingface\.co/blog/LinkedIn/gpt\-oss\-agentic\-rl](https://huggingface.co/blog/LinkedIn/gpt-oss-agentic-rl), 2025\.
- Liu et al\. \[2024\]Hao Liu, Matei Zaharia, and Pieter Abbeel\.Ring attention with blockwise transformers for near\-infinite context\.In*International Conference on Learning Representations*, 2024\.
- Liu et al\. \[2023\]Tianyang Liu, Canwen Xu, and Julian McAuley\.RepoBench: Benchmarking repository\-level code auto\-completion systems\.*arXiv preprint arXiv:2306\.03091*, 2023\.
- Liu et al\. \[2025\]Zikang Liu, Tongtian Yue, Yepeng Tang, Longteng Guo, Junxian Cai, Qingbin Liu, Xi Chen, and Jing Liu\.Prefix grouper: Efficient GRPO training through shared\-prefix forward\.*arXiv preprint arXiv:2506\.05433*, 2025\.
- Schulman et al\. \[2017\]John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov\.Proximal policy optimization algorithms\.*arXiv preprint arXiv:1707\.06347*, 2017\.
- Shah et al\. \[2024\]Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao\.FlashAttention\-3: Fast and accurate attention with asynchrony and low\-precision\.In*Advances in Neural Information Processing Systems*, 2024\.
- Shao et al\. \[2024\]Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, YK Li, Y Wu, and Daya Guo\.DeepSeekMath: Pushing the limits of mathematical reasoning in open language models\.*arXiv preprint arXiv:2402\.03300*, 2024\.
- Sheng et al\. \[2024\]Guangming Sheng, Chi Cao, Zilingfeng Lin, Song Bian, Da Wei, Wenbo Xu, Caicai Yang, Jian Liu, and Tao Zhang\.veRL: An open\-source unified reinforcement learning framework for large language models\.*arXiv preprint arXiv:2409\.19951*, 2024\.
- von Werra et al\. \[2020\]Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, et al\.TRL: Transformers reinforcement learning, 2020\.URL[https://github\.com/huggingface/trl](https://github.com/huggingface/trl)\.
- Yang et al\. \[2025\]An Yang, Baosong Yang, Beichen Zhang, et al\.Qwen3 technical report\.*arXiv preprint arXiv:2505\.09388*, 2025\.
- Yu et al\. \[2025\]Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, Haibin Lin, Zhiqi Lin, Bole Ma, Guangming Sheng, Yuxuan Tong, Chi Zhang, Mofan Zhang, Wang Zhang, Hang Zhu, Jinhua Zhu, Jiaze Chen, Jiangjie Chen, Chengyi Wang, Hongli Yu, Yuxuan Song, Xiangpeng Wei, Hao Zhou, Jingjing Liu, Wei\-Ying Ma, Ya\-Qin Zhang, Lin Yan, Mu Qiao, Yonghui Wu, and Mingxuan Wang\.DAPO: An open\-source LLM reinforcement learning system at scale\.*arXiv preprint arXiv:2503\.14476*, 2025\.
- Zheng et al\. \[2023\]Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Shuo Cheng, Jeff Huang, Baris Kasikci, and Ion Stoica\.SGLang: Efficient execution of structured language model programs\.*arXiv preprint arXiv:2312\.07104*, 2023\.

## Appendix ADualKV Formal Derivation

This appendix provides a self\-contained mathematical derivation of the DualKV forward and backward passes and proves equivalence to the FA2 baseline\.

### A\.1Setup and Notation

ConsiderNNsequences sharing a prompt ofPPtokens\. SequenceiihasRiR\_\{i\}response tokens, with total lengthSi=P\+RiS\_\{i\}=P\+R\_\{i\}\. LetT=P\+∑i=1NRiT=P\+\\sum\_\{i=1\}^\{N\}R\_\{i\}denote the total packed token count\.

In a standard transformer layer, leth∈ℝT×Dh\\in\\mathbb\{R\}^\{T\\times D\}denote the hidden states after the input norm\. The QKV projection produces:

Q=h​WQ,K=h​WK,V=h​WVQ=hW\_\{Q\},\\quad K=hW\_\{K\},\\quad V=hW\_\{V\}\(8\)whereWQ∈ℝD×H​dW\_\{Q\}\\in\\mathbb\{R\}^\{D\\times Hd\},WK∈ℝD×Hk​dW\_\{K\}\\in\\mathbb\{R\}^\{D\\times H\_\{k\}d\},WV∈ℝD×Hk​dW\_\{V\}\\in\\mathbb\{R\}^\{D\\times H\_\{k\}d\}, withHHquery heads,HkH\_\{k\}KV heads, and head dimensiondd\.

### A\.2Prompt Hidden State Invariant

###### Claim 1\(Prompt Hidden State Invariant\)\.

In a decoder\-only model, for any prompt positionj<Pj<P, the hidden statehjh\_\{j\}is identical across allNNsequences at every layer\.

###### Proof\.

By induction on layers\. At the embedding layer, prompt tokens are shared, sohj\(i\)=hjh\_\{j\}^\{\(i\)\}=h\_\{j\}for allii\. Suppose the invariant holds at the input to layerℓ\\ell\. Every operation in the layer preserves it:

- •RMSNorm: per\-token, soNorm​\(hj\(i\)\)=Norm​\(hj\)\\text\{Norm\}\(h\_\{j\}^\{\(i\)\}\)=\\text\{Norm\}\(h\_\{j\}\)for allii\.
- •QKV projection: linear per\-token, soQj\(i\)=QjQ\_\{j\}^\{\(i\)\}=Q\_\{j\},Kj\(i\)=KjK\_\{j\}^\{\(i\)\}=K\_\{j\},Vj\(i\)=VjV\_\{j\}^\{\(i\)\}=V\_\{j\}\.
- •RoPE: depends only on position indexjj, applied per\-token\.
- •Causal attention: forj<Pj<P, the causal mask restricts attention to keys at positions\{0,…,j\}\\\{0,\\ldots,j\\\}, all prompt positions with identicalKKandVV\. Combined withQj\(i\)=QjQ\_\{j\}^\{\(i\)\}=Q\_\{j\}from the QKV\-projection step, the attention output satisfiesOj\(i\)=OjO\_\{j\}^\{\(i\)\}=O\_\{j\}for allii\.
- •Output projection, residual, MLP: all per\-token or element\-wise\.

The invariant holds at the output of layerℓ\\ell, completing the induction\. ∎

### A\.3Forward Pass

#### Packing\.

We represent the micro\-batch as a single packed sequence:

𝐱=\[p0,…,pP−1⏟prompt,r0\(1\),…,rR1−1\(1\)⏟response 1,…,r0\(N\),…,rRN−1\(N\)⏟response​N\]\\mathbf\{x\}=\[\\underbrace\{p\_\{0\},\\ldots,p\_\{P\-1\}\}\_\{\\text\{prompt\}\},\\;\\underbrace\{r\_\{0\}^\{\(1\)\},\\ldots,r\_\{R\_\{1\}\-1\}^\{\(1\)\}\}\_\{\\text\{response 1\}\},\\;\\ldots,\\;\\underbrace\{r\_\{0\}^\{\(N\)\},\\ldots,r\_\{R\_\{N\}\-1\}^\{\(N\)\}\}\_\{\\text\{response \}N\}\]\(9\)
All per\-token operations processT=P\+∑iRiT=P\+\\sum\_\{i\}R\_\{i\}tokens, saving\(N−1\)⋅P\(N\{\-\}1\)\\cdot Ptokens versus the standard∑iSi=N​P\+∑iRi\\sum\_\{i\}S\_\{i\}=NP\+\\sum\_\{i\}R\_\{i\}layout\.

#### QKV split\.

After the QKV projection onTTtokens, we split at the prompt boundaryPP\(zero\-copy tensor views\):

Qc\\displaystyle Q\_\{c\}=Q\[0:P\],Kc=K\[0:P\],Vc=V\[0:P\]\\displaystyle=Q\_\{\[0:P\]\},\\quad K\_\{c\}=K\_\{\[0:P\]\},\\quad V\_\{c\}=V\_\{\[0:P\]\}\(10\)Qd\\displaystyle Q\_\{d\}=Q\[P:T\],Kd=K\[P:T\],Vd=V\[P:T\]\\displaystyle=Q\_\{\[P:T\]\},\\quad K\_\{d\}=K\_\{\[P:T\]\},\\quad V\_\{d\}=V\_\{\[P:T\]\}\(11\)

#### Call 1: Context self\-attention\.

Standard causal attention over thePPprompt tokens, computed once:

Oc=CausalAttn​\(Qc,Kc,Vc\)∈ℝP×H×dO\_\{c\}=\\text\{CausalAttn\}\(Q\_\{c\},K\_\{c\},V\_\{c\}\)\\in\\mathbb\{R\}^\{P\\times H\\times d\}\(12\)

#### Call 2: Decoded attention \(DualKV kernel\)\.

For sequenceiiand decoded positionr∈\{0,…,Ri−1\}r\\in\\\{0,\\ldots,R\_\{i\}\-1\\\}, the query attends to the full key\-value sequence\[Kc;Kd\(i\)\]\[K\_\{c\};\\,K\_\{d\}^\{\(i\)\}\]with causal masking at logical positionP\+rP\+r:

Od\(i\)​\[r\]=∑j=0P\+rαj\(i,r\)⋅vjO\_\{d\}^\{\(i\)\}\[r\]=\\sum\_\{j=0\}^\{P\+r\}\\alpha\_\{j\}^\{\(i,r\)\}\\cdot v\_\{j\}\(13\)wherevjv\_\{j\}indexes into\[Vc;Vd\(i\)\]\[V\_\{c\};\\,V\_\{d\}^\{\(i\)\}\]and the attention weights are:

αj\(i,r\)=exp⁡\(qr\(i\)⋅kj/d\)∑j′=0P\+rexp⁡\(qr\(i\)⋅kj′/d\)\\alpha\_\{j\}^\{\(i,r\)\}=\\frac\{\\exp\\\!\\big\(q\_\{r\}^\{\(i\)\}\\cdot k\_\{j\}/\\sqrt\{d\}\\big\)\}\{\\sum\_\{j^\{\\prime\}=0\}^\{P\+r\}\\exp\\\!\\big\(q\_\{r\}^\{\(i\)\}\\cdot k\_\{j^\{\\prime\}\}/\\sqrt\{d\}\\big\)\}\(14\)
Thecontext\_seqlen=P\{\}=Pparameter shifts the kernel’s internal position counter so that decoded query at kernel indexrrmaps to logical positionP\+rP\+r, ensuring:

- •AllPPcontext positions are unmasked \(j<P≤P\+rj<P\\leq P\+r\)\.
- •Decoded positions are causally masked \(j≤P\+rj\\leq P\+r\)\.

#### Reassembly\.

The final attention output is:

O=\[Oc;Od\]∈ℝT×H×dO=\[O\_\{c\};\\,O\_\{d\}\]\\in\\mathbb\{R\}^\{T\\times H\\times d\}\(15\)which flows throughWOW\_\{O\}, residual addition, and MLP onTTtokens\.

### A\.4Backward Pass

Letℒ\\mathcal\{L\}denote the training loss \(computed only over response tokens\)\. We receive∂ℒ/∂O∈ℝT×H×d\\partial\\mathcal\{L\}/\\partial O\\in\\mathbb\{R\}^\{T\\times H\\times d\}and split it:

∂ℒ∂Oc=\(∂ℒ∂O\)\[0:P\],∂ℒ∂Od=\(∂ℒ∂O\)\[P:T\]\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\_\{c\}\}=\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\}\\right\)\_\{\\\!\[0:P\]\},\\quad\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\_\{d\}\}=\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\}\\right\)\_\{\\\!\[P:T\]\}\(16\)
Note that∂ℒ/∂Oc≠0\\partial\\mathcal\{L\}/\\partial O\_\{c\}\\neq 0in general: even though the loss is on response tokens only, gradients propagate back to prompt positions through subsequent layers’ residual connections\.

#### Call 2 backward \(DualKV kernel\)\.

Given∂ℒ/∂Od\\partial\\mathcal\{L\}/\\partial O\_\{d\}and the saved tensors from the forward pass, the backward produces:

∂ℒ∂Qd,∂ℒ∂Kd,∂ℒ∂Vd\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial Q\_\{d\}\},\\quad\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{d\}\},\\quad\\frac\{\\partial\\mathcal\{L\}\}\{\\partial V\_\{d\}\}\(17\)and gradients with respect to the shared context KV:

\(∂ℒ∂Kc\)d\\displaystyle\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}\\right\)\_\{\\\!d\}=∑i=1N∑r=0Ri−1∑h∂ℒ∂Od\(i\)​\[r,h\]⋅∂Od\(i\)​\[r,h\]∂Kc\\displaystyle=\\sum\_\{i=1\}^\{N\}\\sum\_\{r=0\}^\{R\_\{i\}\-1\}\\sum\_\{h\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\_\{d\}^\{\(i\)\}\[r,h\]\}\\cdot\\frac\{\\partial O\_\{d\}^\{\(i\)\}\[r,h\]\}\{\\partial K\_\{c\}\}\(18\)\(∂ℒ∂Vc\)d\\displaystyle\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial V\_\{c\}\}\\right\)\_\{\\\!d\}=∑i=1N∑r=0Ri−1∑h∂ℒ∂Od\(i\)​\[r,h\]⋅∂Od\(i\)​\[r,h\]∂Vc\\displaystyle=\\sum\_\{i=1\}^\{N\}\\sum\_\{r=0\}^\{R\_\{i\}\-1\}\\sum\_\{h\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\_\{d\}^\{\(i\)\}\[r,h\]\}\\cdot\\frac\{\\partial O\_\{d\}^\{\(i\)\}\[r,h\]\}\{\\partial V\_\{c\}\}\(19\)

#### Implementation: fp32 gradient accumulation\.

The summation overNNsequences in Eq\.[18](https://arxiv.org/html/2605.15422#A1.E18)–[19](https://arxiv.org/html/2605.15422#A1.E19)introduces a concurrent\-write pattern absent from standard FlashAttention: the single sharedd​KcdK\_\{c\}andd​VcdV\_\{c\}buffer has no batch offset, yetNNthread blocks \(one per decoded sequence\) concurrently compute partial contributions to the same context positions\. DualKV resolves this with fp32 atomic accumulation in a dedicated accumulator buffer, followed by a cast kernel that converts the accumulator to bf16\. The kernel\-level description and pseudocode are given in Section[3\.2](https://arxiv.org/html/2605.15422#S3.SS2)\(Algorithms[1](https://arxiv.org/html/2605.15422#alg1)–[2](https://arxiv.org/html/2605.15422#alg2)\)\. For the purposes of the mathematical derivation that follows, the relevant property is that the kernel produces\(∂ℒ/∂Kc\)d=∑i∑r∑h∂ℒ/∂Od\(i\)​\[r,h\]⋅∂Od\(i\)​\[r,h\]/∂Kc\(\\partial\\mathcal\{L\}/\\partial K\_\{c\}\)\_\{d\}=\\sum\_\{i\}\\sum\_\{r\}\\sum\_\{h\}\\partial\\mathcal\{L\}/\\partial O\_\{d\}^\{\(i\)\}\[r,h\]\\cdot\\partial O\_\{d\}^\{\(i\)\}\[r,h\]/\\partial K\_\{c\}exactly \(up to fp32 rounding\) as stated in Eq\.[18](https://arxiv.org/html/2605.15422#A1.E18)–[19](https://arxiv.org/html/2605.15422#A1.E19)\.

#### Call 1 backward \(standard FA2\)\.

Given∂ℒ/∂Oc\\partial\\mathcal\{L\}/\\partial O\_\{c\}, produces:

∂ℒ∂Qc,\(∂ℒ∂Kc\)c,\(∂ℒ∂Vc\)c\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial Q\_\{c\}\},\\quad\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}\\right\)\_\{\\\!c\},\\quad\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial V\_\{c\}\}\\right\)\_\{\\\!c\}\(20\)

#### Gradient accumulation\.

SinceKcK\_\{c\}andVcV\_\{c\}are the*same tensor objects*passed to both Call 1 and Call 2, PyTorch autograd accumulates their gradients:

∂ℒ∂Kc\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}=\(∂ℒ∂Kc\)c\+\(∂ℒ∂Kc\)d\\displaystyle=\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}\\right\)\_\{\\\!c\}\+\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}\\right\)\_\{\\\!d\}\(21\)∂ℒ∂Vc\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial V\_\{c\}\}=\(∂ℒ∂Vc\)c\+\(∂ℒ∂Vc\)d\\displaystyle=\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial V\_\{c\}\}\\right\)\_\{\\\!c\}\+\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial V\_\{c\}\}\\right\)\_\{\\\!d\}\(22\)
The full QKV gradients are then reassembled:

∂ℒ∂Q=\[∂ℒ∂Qc;∂ℒ∂Qd\],∂ℒ∂K=\[∂ℒ∂Kc;∂ℒ∂Kd\],∂ℒ∂V=\[∂ℒ∂Vc;∂ℒ∂Vd\]\\frac\{\\partial\\mathcal\{L\}\}\{\\partial Q\}=\\left\[\\frac\{\\partial\\mathcal\{L\}\}\{\\partial Q\_\{c\}\};\\,\\frac\{\\partial\\mathcal\{L\}\}\{\\partial Q\_\{d\}\}\\right\],\\quad\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\}=\\left\[\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\};\\,\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{d\}\}\\right\],\\quad\\frac\{\\partial\\mathcal\{L\}\}\{\\partial V\}=\\left\[\\frac\{\\partial\\mathcal\{L\}\}\{\\partial V\_\{c\}\};\\,\\frac\{\\partial\\mathcal\{L\}\}\{\\partial V\_\{d\}\}\\right\]\(23\)
These propagate through the QKV projection backward to produce∂ℒ/∂h\\partial\\mathcal\{L\}/\\partial hand∂ℒ/∂WQ\\partial\\mathcal\{L\}/\\partial W\_\{Q\},∂ℒ/∂WK\\partial\\mathcal\{L\}/\\partial W\_\{K\},∂ℒ/∂WV\\partial\\mathcal\{L\}/\\partial W\_\{V\}\.

### A\.5Equivalence to the FA2 Baseline

###### Theorem 1\(Gradient Equivalence\)\.

The parameter gradients produced by DualKV are identical to those produced by the FA2 baseline\.

###### Proof\.

In the standard approach, each sequencei∈\{1,…,N\}i\\in\\\{1,\\ldots,N\\\}has its own copy of the prompt KV:Kc\(i\),Vc\(i\)K\_\{c\}^\{\(i\)\},V\_\{c\}^\{\(i\)\}\. The gradient ofKc\(i\)K\_\{c\}^\{\(i\)\}has two components:

∂ℒ∂Kc\(i\)=gself\(i\)⏟prompt self\-attn\+gcross\(i\)⏟decoded→context\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}^\{\(i\)\}\}=\\underbrace\{g\_\{\\text\{self\}\}^\{\(i\)\}\}\_\{\\text\{prompt self\-attn\}\}\+\\underbrace\{g\_\{\\text\{cross\}\}^\{\(i\)\}\}\_\{\\text\{decoded\}\\to\\text\{context\}\}\(24\)
The parameter gradient from theWKW\_\{K\}projection at prompt positions is:

∂ℒ∂WK\|prompt=∑i=1N\(∂ℒ∂Kc\(i\)\)⊤​hc=\(∑i=1N∂ℒ∂Kc\(i\)\)⊤​hc\\frac\{\\partial\\mathcal\{L\}\}\{\\partial W\_\{K\}\}\\bigg\|\_\{\\text\{prompt\}\}=\\sum\_\{i=1\}^\{N\}\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}^\{\(i\)\}\}\\right\)^\{\\top\}h\_\{c\}=\\left\(\\sum\_\{i=1\}^\{N\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}^\{\(i\)\}\}\\right\)^\{\\top\}h\_\{c\}\(25\)wherehch\_\{c\}is the prompt hidden state \(identical acrossiiby the invariant, so it factors out of the sum\)\.

In DualKV, the singleKcK\_\{c\}receives gradient \(Eq\.[21](https://arxiv.org/html/2605.15422#A1.E21)\):

∂ℒ∂Kc=\(∂ℒ∂Kc\)c\+∑i=1Ngcross\(i\)\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}=\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}\\right\)\_\{\\\!c\}\+\\sum\_\{i=1\}^\{N\}g\_\{\\text\{cross\}\}^\{\(i\)\}\(26\)
The parameter gradient is:

∂ℒ∂WK\|promptDualKV=\(∂ℒ∂Kc\)⊤​hc\\frac\{\\partial\\mathcal\{L\}\}\{\\partial W\_\{K\}\}\\bigg\|\_\{\\text\{prompt\}\}^\{\\text\{DualKV\}\}=\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}\\right\)^\{\\top\}h\_\{c\}\(27\)
It remains to show that Eq\.[25](https://arxiv.org/html/2605.15422#A1.E25)equals Eq\.[27](https://arxiv.org/html/2605.15422#A1.E27), i\.e\.:

∑i=1N\(gself\(i\)\+gcross\(i\)\)=\(∂ℒ∂Kc\)c\+∑i=1Ngcross\(i\)\\sum\_\{i=1\}^\{N\}\\left\(g\_\{\\text\{self\}\}^\{\(i\)\}\+g\_\{\\text\{cross\}\}^\{\(i\)\}\\right\)=\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}\\right\)\_\{\\\!c\}\+\\sum\_\{i=1\}^\{N\}g\_\{\\text\{cross\}\}^\{\(i\)\}\(28\)
This reduces to showing:

∑i=1Ngself\(i\)=\(∂ℒ∂Kc\)c\\sum\_\{i=1\}^\{N\}g\_\{\\text\{self\}\}^\{\(i\)\}=\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}\\right\)\_\{\\\!c\}\(29\)
In the standard approach,gself\(i\)g\_\{\\text\{self\}\}^\{\(i\)\}is the gradient ofKc\(i\)K\_\{c\}^\{\(i\)\}from the context self\-attention of sequenceii\. This depends on∂ℒ/∂Oc\(i\)\\partial\\mathcal\{L\}/\\partial O\_\{c\}^\{\(i\)\}, which propagates back from subsequent layers\. While the forward values at prompt positions are identical acrossii, the backward gradients∂ℒ/∂Oc\(i\)\\partial\\mathcal\{L\}/\\partial O\_\{c\}^\{\(i\)\}generally differ \(each sequenceiihas different response tokens producing different loss contributions\)\.

In DualKV, the single prompt copy receives the*aggregated*upstream gradient:

∂ℒ∂Oc=∑i=1N∂ℒ∂Oc\(i\)\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\_\{c\}\}=\\sum\_\{i=1\}^\{N\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\_\{c\}^\{\(i\)\}\}\(30\)
This aggregation happens naturally through the residual stream: in the packed layout, prompt hidden states at layerℓ\+1\\ell\{\+\}1receive gradients from allNNresponse branches propagating backward through subsequent layers\. Since all intervening operations \(norms, projections, MLP, residual\) are per\-token and*linear in the backward pass*, the gradient at the prompt positions is exactly the sum over allNNsequences’ contributions\.

The context self\-attention backward is linear in∂ℒ/∂Oc\\partial\\mathcal\{L\}/\\partial O\_\{c\}:

\(∂ℒ∂Kc\)c=fattn\-bwd​\(∂ℒ∂Oc\)=fattn\-bwd​\(∑i=1N∂ℒ∂Oc\(i\)\)=∑i=1Nfattn\-bwd​\(∂ℒ∂Oc\(i\)\)=∑i=1Ngself\(i\)\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial K\_\{c\}\}\\right\)\_\{\\\!c\}=f\_\{\\text\{attn\-bwd\}\}\\\!\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\_\{c\}\}\\right\)=f\_\{\\text\{attn\-bwd\}\}\\\!\\left\(\\sum\_\{i=1\}^\{N\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\_\{c\}^\{\(i\)\}\}\\right\)=\\sum\_\{i=1\}^\{N\}f\_\{\\text\{attn\-bwd\}\}\\\!\\left\(\\frac\{\\partial\\mathcal\{L\}\}\{\\partial O\_\{c\}^\{\(i\)\}\}\\right\)=\\sum\_\{i=1\}^\{N\}g\_\{\\text\{self\}\}^\{\(i\)\}\(31\)
wherefattn\-bwdf\_\{\\text\{attn\-bwd\}\}denotes the attention backward function mappingd​OdOtod​KdK, which is linear ind​OdO\(since the forward valuesQc,Kc,VcQ\_\{c\},K\_\{c\},V\_\{c\}are fixed during backward\)\. This establishes Eq\.[29](https://arxiv.org/html/2605.15422#A1.E29)and completes the proof\. ∎

### A\.6Complexity Summary

This subsection summarizes the per\-component FLOPs and memory scaling of FA2 baseline vs\. DualKV in asymptotic notation, abstracting the concrete numbers of Section[2](https://arxiv.org/html/2605.15422#S2)\. The table shows that the savings are dominated by the token\-count reductionTstd→TdkT\_\{\\text\{std\}\}\\to T\_\{\\text\{dk\}\}: every per\-token cost \(norms, projections, MLP, KV storage, activation memory\) shrinks by the ratioρ=Tstd/Tdk\\rho=T\_\{\\text\{std\}\}/T\_\{\\text\{dk\}\}, and only attention has additional structural savings beyond that\.

Table 5:Complexity comparison\.Tstd=N​P\+∑iRiT\_\{\\text\{std\}\}=NP\+\\sum\_\{i\}R\_\{i\},Tdk=P\+∑iRiT\_\{\\text\{dk\}\}=P\+\\sum\_\{i\}R\_\{i\}\.WhenP≫RP\\gg RandRi=RR\_\{i\}=Rfor allii:

TstdTdk=N​P\+N​RP\+N​R→P≫RN\\frac\{T\_\{\\text\{std\}\}\}\{T\_\{\\text\{dk\}\}\}=\\frac\{NP\+NR\}\{P\+NR\}\\xrightarrow\{P\\gg R\}N\(32\)
The savings approach a factor ofNN\(the batch size\) across all components — norms, projections, MLP, attention, and activation memory — yielding near\-linear scaling\.

## Appendix BTraining\-Step Gradient Equivalence

Appendix[A](https://arxiv.org/html/2605.15422#A1)proves the DualKV kernel produces attention gradients identical to the FA2 baseline in exact arithmetic for any fixed packed input\. This appendix lifts that kernel\-level invariance to the training\-step level\. DualKV introduces two modifications to veRL’s policy\-update data pipeline, both motivated by the kernel’s same\-prompt\-per\-micro\-batch packing contract:

1. 1\.Skipbalance\_batch\.veRL’s default pipeline reorders the rollout list by per\-sample sequence length to equalize per\-GPU token counts before dispatch\. DualKV skips this step, since the seqlen permutation scatters rollouts from the same prompt across different GPUs, breaking the single\-prompt packing the kernel requires\. The global rollout multiset is unchanged; only its ordering and per\-GPU partition differ\.
2. 2\.Forceshuffle=Falsewithin the per\-rank mini\-batch iterator\. Under the default pipeline, samples within each rank are reshuffled between PPO epochs\. DualKV disables this to preserve contiguous prompt\-grouped ordering across epochs\.

For one training step, letBBbe the global multiset of\(p,r,A\)\(p,r,A\)triples produced by that step’s rollout, whereppis a prompt,rra sampled response, andA=A​\(p,r\)A=A\(p,r\)a scalar advantage\. The policy\-gradient estimator isg^​\(B,θ\)=∑\(p,r,A\)∈Bf​\(p,r,A,θ\)\\hat\{g\}\(B,\\theta\)=\\sum\_\{\(p,r,A\)\\in B\}f\(p,r,A,\\theta\), covering standard RL objectives \(PPO, GRPO, DAPO\)\. Let𝒫def\\mathcal\{P\}\_\{\\text\{def\}\}denote veRL’s default pipeline and𝒫dk\\mathcal\{P\}\_\{\\text\{dk\}\}the DualKV pipeline\.

###### Theorem 2\(No systematic bias at training\-step aggregate\)\.

For any parametersθ\\theta,

∑\(p,r,A\)∈Bf𝒫def​\(p,r,A,θ\)=∑\(p,r,A\)∈Bf𝒫dk​\(p,r,A,θ\)in exact arithmetic\.\\sum\_\{\(p,r,A\)\\in B\}f^\{\\mathcal\{P\}\_\{\\text\{def\}\}\}\(p,r,A,\\theta\)\\;=\\;\\sum\_\{\(p,r,A\)\\in B\}f^\{\\mathcal\{P\}\_\{\\text\{dk\}\}\}\(p,r,A,\\theta\)\\qquad\\text\{in exact arithmetic\.\}\(33\)

###### Proof\.

The per\-sample gradientf​\(p,r,A,θ\)f\(p,r,A,\\theta\)depends onθ\\theta, the advantageAA, and the per\-token log\-probabilitiesπθ​\(r∣p\)\\pi\_\{\\theta\}\(r\\mid p\)\. Advantages are computed upstream of the policy update and are pipeline\-independent\. The kernel\-level equivalence theorem of Appendix[A](https://arxiv.org/html/2605.15422#A1)establishes the primary lemma we rely on: log\-probabilities and their gradients at response positions are identical under𝒫def\\mathcal\{P\}\_\{\\text\{def\}\}and𝒫dk\\mathcal\{P\}\_\{\\text\{dk\}\}for any fixed sample\. Hencef𝒫def​\(p,r,A,θ\)=f𝒫dk​\(p,r,A,θ\)f^\{\\mathcal\{P\}\_\{\\text\{def\}\}\}\(p,r,A,\\theta\)=f^\{\\mathcal\{P\}\_\{\\text\{dk\}\}\}\(p,r,A,\\theta\)termwise\. The two pipeline modifications are both permutations of the rollout list, so the global multisetBBis the same under both pipelines\. The equation follows by termwise equality and order\-invariance of finite sums\. ∎

## Appendix CTheoretical Speedup and Memory Analysis

This appendix derives per\-layer kernel\-level memory savings and the closed\-form attention speedup ratio for DualKV\. These are theoretical ceilings; measured end\-to\-end numbers are in Sections[4\.2](https://arxiv.org/html/2605.15422#S4.SS2)–[4\.3](https://arxiv.org/html/2605.15422#S4.SS3)\(and Appendix[E\.1](https://arxiv.org/html/2605.15422#A5.SS1)\)\.

### C\.1Memory Savings

Per layer, the DualKV kernel avoids storingN−1N\{\-\}1redundant copies of three tensors that an equivalentNN\-copy FA2 call would materialize:

Table 6:Kernel\-level memory savings per layer relative to anNN\-copy FA2 call over the same packed batch\. Configuration: Qwen3\-8B \(N=8N\{=\}8,P=16384P\{=\}16384,R=2048R\{=\}2048,H=32H\{=\}32,Hk=8H\_\{k\}\{=\}8,d=128d\{=\}128\) in bf16\.These are*kernel\-level*savings: they count the per\-layer attention\-input tensors \(QQat prompt positions,KcK\_\{c\},VcV\_\{c\}\) and the prompt\-positionsoftmax\_lsethat an FA2 baseline would materializeNNtimes and DualKV materializes once\. DualKV’s single\-prompt packing also reduces framework\-level activation memory across every per\-token operation \(norms, QKV/output projections, MLP, residual stream\), since these now processP\+N​RP\+NRtokens instead ofN​\(P\+R\)N\(P\+R\)\.

### C\.2Speedup Scaling

The attention speedup ratio scales approximately as:

Speedupattn≈N⋅S2P2\+N⋅R⋅S=N​\(P\+R\)2P2\+N​R​\(P\+R\)\\text\{Speedup\}\_\{\\text\{attn\}\}\\approx\\frac\{N\\cdot S^\{2\}\}\{P^\{2\}\+N\\cdot R\\cdot S\}=\\frac\{N\(P\+R\)^\{2\}\}\{P^\{2\}\+NR\(P\+R\)\}\(34\)
ForP≫RP\\gg R, this simplifies to≈N\\approx N, meaning DualKV approaches a linear speedup in batch size as prompt length dominates\. ForP=32​KP\{=\}32\\text\{K\},R=2​KR\{=\}2\\text\{K\},N=16N\{=\}16: theoretical speedup is6\.1×6\.1\\times, measured fwd\+bwd is5\.48×5\.48\\times\(90% efficiency\)\. The consistent speedup across forward and backward passes confirms that the theoretical FLOPs analysis \(Section[2](https://arxiv.org/html/2605.15422#S2)\) applies equally to both directions\.

The end\-to\-end training speedup is lower than this kernel\-level ceiling, as attention is only one component of the training step \(alongside generation, projections, MLP, optimizer, and communication\)\. The token reduction ratioρ\\rho\(Section[2](https://arxiv.org/html/2605.15422#S2)\) provides an upper bound on the full\-model speedup\.

### C\.3Token Reduction Ratioρ\\rho: Scaling Analysis and Regimes

The token reduction ratio of DualKV over standard attention is:

ρ=TstdTdk=N​\(P\+R\)P\+N​R\\rho=\\frac\{T\_\{\\text\{std\}\}\}\{T\_\{\\text\{dk\}\}\}=\\frac\{N\(P\+R\)\}\{P\+NR\}\(35\)
This ratio governs the speedup across*all*model components \(norms, projections, MLP, activation memory\), not just attention\. We analyze its behavior along the two key axes: rollout factorNNand prompt\-to\-response ratioP/RP/R\.

#### Scaling withNN\(rollout factor\)\.

FixingPPandRR, taking the limit:

limN→∞ρ=limN→∞N​\(P\+R\)P\+N​R=P\+RR=1\+PR\\lim\_\{N\\to\\infty\}\\rho=\\lim\_\{N\\to\\infty\}\\frac\{N\(P\+R\)\}\{P\+NR\}=\\frac\{P\+R\}\{R\}=1\+\\frac\{P\}\{R\}\(36\)
For largeNN, the speedup saturates at1\+P/R1\+P/R\. WithP=16​KP\{=\}16\\text\{K\},R=2​KR\{=\}2\\text\{K\}: ceiling is9×9\\times\. IncreasingNNbeyond this point yields diminishing returns — the bottleneck shifts to decoded\-token computeO​\(N​R\)O\(NR\)which scales linearly in both approaches\.

#### Scaling withP/RP/R\(prompt length ratio\)\.

FixingNNandRR, asP→∞P\\to\\infty:

limP→∞ρ=limP→∞N​P\+N​RP\+N​R=N\\lim\_\{P\\to\\infty\}\\rho=\\lim\_\{P\\to\\infty\}\\frac\{NP\+NR\}\{P\+NR\}=N\(37\)
For long prompts, the speedup approachesNN— near\-linear in the rollout factor\. This is the regime where DualKV is most impactful: the\(N−1\)​P\(N\{\-\}1\)Predundant prompt tokens dominate total compute\.

#### Concrete scenarios\.

Table[7](https://arxiv.org/html/2605.15422#A3.T7)shows the token reduction ratioρ\\rhofor representative RL training configurations\.

Table 7:Token reduction ratioρ=Tstd/Tdk\\rho=T\_\{\\text\{std\}\}/T\_\{\\text\{dk\}\}for various\(N,P,R\)\(N,P,R\)configurations\.Two regimes emerge where DualKV provides substantial benefits:

1. 1\.Large rollout factor\(N≥16N\\geq 16\): Common in GRPO and DAPO, where generating many candidate responses per prompt improves reward signal quality\. AtN=32N\{=\}32withP=16​KP\{=\}16\\text\{K\}, the reduction reaches7\.2×7\.2\\times\.
2. 2\.Long prompts\(P≥16​KP\\geq 16\\text\{K\}\): Prevalent in agentic tasks, repository\-level code generation, and multi\-turn dialogue\. AtP=64​KP\{=\}64\\text\{K\}, even moderate rollout \(N=8N\{=\}8\) yields7\.6×7\.6\\timesreduction, andN=16N\{=\}16reaches14\.3×14\.3\\times\.

## Appendix DSoftware Environment

All kernel\-level benchmarks \(Section[4\.1](https://arxiv.org/html/2605.15422#S4.SS1)\) and end\-to\-end GRPO training runs \(Sections[4\.2](https://arxiv.org/html/2605.15422#S4.SS2)–[4\.3](https://arxiv.org/html/2605.15422#S4.SS3)and Appendix[E\.1](https://arxiv.org/html/2605.15422#A5.SS1)\) use an identical Python virtual environment\. Configuration flags \(use\_dualkv,ulysses\_sequence\_parallel\_size,attn\_implementation\) are the only variables that differ between runs\.

Table 8:Software stack used across all experiments\. “DualKV build” indicates our fork offlash\_attn2\.8\.4 with the DualKV kernel \(Section[3](https://arxiv.org/html/2605.15422#S3)\) added; base FA2 behavior is unchanged\. “DualKV integration” indicates our patch toverl0\.7\.0 that adds the data\-pipeline changes \(Section[2\.2](https://arxiv.org/html/2605.15422#S2.SS2)\) and actor/trainer hooks needed to use the DualKV kernel\.The veRL patch \(∼1200\{\\sim\}1200lines across 10 files\) modifies the actor’s forward pass, the trainer’s data pipeline \(skippingbalance\_batch, enforcing prompt\-grouping\), and a monkey\-patch hook that installs the DualKV attention call on supported model classes\. Both theflash\_attnandverlpatches will be released upon publication\.

## Appendix EAdditional Experiments

### E\.1GRPO Training on GSM8K \(Short Prompts\)

This appendix presents our short\-prompt validation experiment \(referenced in Section[4](https://arxiv.org/html/2605.15422#S4)introduction\)\. The experiment stress\-tests DualKV in the low\-ρ\\rhoregime \(ρ=1\.8\\rho=1\.8–3\.8×3\.8\\times\), where kernel\-level speedups are modest but memory savings still extend the feasible training frontier\.

We evaluate DualKV in end\-to\-end GRPO training using the veRL framework on GSM8K\[Cobbe et al\.,[2021](https://arxiv.org/html/2605.15422#bib.bib2)\], comparing against standard FA2 packing across a range of prompt lengths\.

#### Setup\.

We train Qwen3\-8B\[Yang et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib23)\]with GRPO on an 8\-node cluster of NVIDIA A100\-SXM4\-40GB GPUs \(64 GPUs total\)\. The model is sharded via FSDP across 32 data\-parallel ranks \(tensor parallelism=2=2\), with gradient checkpointing enabled and optimizer states offloaded to CPU\. Each training step consists of: \(1\) rollout generation via vLLM \(N=8N\{=\}8responses per prompt,Rmax=256R\_\{\\max\}\{=\}256\), \(2\) log\-probability computation, and \(3\) policy update with gradient accumulation \(micro\-batch size=8=8per GPU\)\. We useppo\_mini\_batch\_size=128\{\}=128andtrain\_batch\_size=512\{\}=512\.

#### Dataset\.

To control prompt length, we construct synthetic GSM8K variants by prepending few\-shot examples to each prompt until the target token count is reached\. We sweep prompt lengthsP∈\{1​K,1\.5​K,2​K,2\.5​K,3​K\}P\\in\\\{1\\text\{K\},1\.5\\text\{K\},2\\text\{K\},2\.5\\text\{K\},3\\text\{K\}\\\}tokens, all withR=256R\{=\}256andN=8N\{=\}8\. Each configuration runs 3 training steps; we report metrics from step 3 \(after 2 warmup steps\) to exclude one\-time initialization costs\.

#### Prompt grouping\.

DualKV requires allNNresponses from the same prompt to be co\-located on the same GPU within each micro\-batch\. We modify veRL’s data pipeline to \(1\) skip the defaultbalance\_batchreordering when DualKV is enabled, preserving prompt\-group contiguity, and \(2\) add runtime assertions verifying that every micro\-batch contains only same\-prompt rollouts\. The micro\-batch size \(=8=8\) equalsNN, ensuring each forward/backward pass processes exactly one prompt group\. Appendix[B](https://arxiv.org/html/2605.15422#A2)proves that these data\-pipeline modifications preserve the exact mini\-batch gradient estimatorg^​\(B,θ\)\\hat\{g\}\(B,\\theta\)under data\-parallel training, so they do not alter training dynamics at any parameter setting\.

#### Results\.

Table[9](https://arxiv.org/html/2605.15422#A5.T9)reports peak GPU memory \(reserved\) and wall\-clock time for both the full training step and the policy update phase \(where DualKV applies\)\. Rollout generation uses vLLM identically in both configurations and is excluded from the policy update timing\.

Table 9:End\-to\-end GRPO training: FA2 vs\. DualKV at micro\-batch size 8,N=8N\{=\}8,R=256R\{=\}256\. Qwen3\-8B on 64×\\timesA100\-40GB\. “Update” = policy update time\. OOM = out of memory during policy update\.ρ=N​\(P\+R\)/\(P\+N​R\)\\rho=N\(P\{\+\}R\)/\(P\{\+\}NR\)is the token reduction ratio\.‡Exceeds A100\-40GB physical memory; PyTorch uses unified memory with severe performance degradation\.

Memory scaling\.FA2 peak memory grows rapidly with prompt length \(37\.8→\\to44\.4→\\to46\.8 GB\), exceeding the A100’s 40 GB physical memory atP=1\.5​KP\{=\}1\.5\\text\{K\}and hitting OOM atP=2\.5​KP\{=\}2\.5\\text\{K\}\. DualKV memory grows slowly \(26\.5→\\to29\.8 GB\), remaining well within the 40 GB budget at all tested contexts\. The memory gap widens from 11\.3 GB \(−30%\-30\\%\) atP=1​KP\{=\}1\\text\{K\}to 18\.7 GB \(−40%\-40\\%\) atP=2​KP\{=\}2\\text\{K\}\. BeyondP=2​KP\{=\}2\\text\{K\}, FA2 cannot train at all at micro\-batch size 8, while DualKV continues to scale\.

Policy update speedup\.The policy update phase \(forward/backward passes with gradient accumulation\) shows increasing speedup:1\.04×1\.04\\timesatP=1​KP\{=\}1\\text\{K\},1\.05×1\.05\\timesatP=1\.5​KP\{=\}1\.5\\text\{K\}, and1\.19×1\.19\\timesatP=2​KP\{=\}2\\text\{K\}\. The update time for DualKV is nearly constant across prompt lengths \(168168–176176s\), confirming that the shared prompt KV is processed only once regardless of prompt length\.

End\-to\-end step time\.Total step time improvement is more modest \(1\.021\.02–1\.08×1\.08\\times\) because rollout generation \(vLLM inference\) and log\-probability computation — which are identical for both methods — dominate\. AtP=2​KP\{=\}2\\text\{K\}, the policy update accounts for only 36% of FA2’s step time but receives1\.19×1\.19\\timesspeedup, yielding an overall1\.08×1\.08\\timesimprovement\.

Extending the training frontier\.The most impactful result is that DualKV*enables training*at context lengths where FA2 cannot run\. AtP=3​KP\{=\}3\\text\{K\}withN=8N\{=\}8and micro\-batch size 8, FA2 requires∼8×\(3072\+256\)=26,624\{\\sim\}8\\times\(3072\{\+\}256\)=26,624tokens per micro\-batch forward pass, while DualKV requires only3072\+8×256=5,1203072\+8\\times 256=5,120tokens — a5\.2×5\.2\\timesreduction in activation memory\. This enables practitioners to train with longer prompts on existing hardware without reducing batch size or adding gradient checkpointing overhead\.

#### Numerical correctness\.

We verified that DualKV matches standard FA2 to within half\-precision rounding in two independent tests: \(1\) kernel\-level forward pass comparison usingtorch\.allclose\(atol=10−3\\text\{atol\}\{=\}10^\{\-3\},rtol=10−3\\text\{rtol\}\{=\}10^\{\-3\}in fp16\) across all supported head dimensions \(64, 96, 128, 192, 256\), GQA configurations, and context/decoded lengths; and \(2\) end\-to\-end inference with Qwen2\-1\.5B via vLLM, where DualKV produced*identical token sequences*to the standard attention baseline across 4 sequences sharing a common prompt \(greedy decoding, 128 tokens\)\. These results confirm that the two\-phase attention decomposition and gradient accumulation introduce no numerical divergence beyond the expected half\-precision rounding\.

#### Note onρ\\rhovalues\.

The token reduction ratios in this experiment \(ρ=1\.8\\rho=1\.8–3\.8×3\.8\\times\) are moderate becauseN=8N\{=\}8andR=256R\{=\}256\(short responses\)\. The kernel\-level benchmarks in Table[1](https://arxiv.org/html/2605.15422#S4.T1)useN=28N\{=\}28andR=2048R\{=\}2048, achievingρ\\rhoup to7×7\\timeswith correspondingly larger speedups\. In production RL training withN=16N\{=\}16–3232and longer responses, the end\-to\-end benefits will be substantially greater\.

### E\.2DAPO Training on LongReason

DualKV benefits any sampling\-based RL method that packsNNsame\-prompt responses into a micro\-batch\. To demonstrate this generality beyond GRPO, we repeat the LongReason experiment \(Section[4\.2](https://arxiv.org/html/2605.15422#S4.SS2)\) using DAPO\[Yu et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib24)\]— a sampling\-based RL method that removes the KL penalty \(no reference model forward pass\) and uses dynamic clipping with token\-level loss aggregation\.

#### Setup\.

Identical to Section[4\.2](https://arxiv.org/html/2605.15422#S4.SS2): Qwen3\-8B on a singlep5\.48xlargeinstance \(8×\\timesH100\), FSDP2 \+ BF16 \+ gradient checkpointing, vLLM rollout \(TP=2\{=\}2,N=32N\{=\}32\),train\_batch\_size=128\{=\}128,ppo\_mini\_batch\_size=64\{=\}64,Pmax=8192P\_\{\\max\}\{=\}8192,Rmax=2048R\_\{\\max\}\{=\}2048, 25 training steps\. DAPO\-specific:clip\_ratio\_low=0\.2\{=\}0\.2,clip\_ratio\_high=0\.28\{=\}0\.28,clip\_ratio\_c=10\.0\{=\}10\.0,loss\_agg\_mode=\{=\}token\-mean, no KL loss\. Because DAPO does not use a reference model, theref\_log\_probphase is absent; the training step consists of rollout,old\_log\_prob, and the policy update only\.

#### Comparisons\.

Same four configurations as Section[4\.2](https://arxiv.org/html/2605.15422#S4.SS2): FA2 \(mb=4\), FA3 \(mb=4\), DualKV \(mb=4\), and DualKV \(mb=8\)\.

#### Results\.

Table[10](https://arxiv.org/html/2605.15422#A5.T10)summarizes the results\. DualKV delivers1\.91×1\.91\\times\(mb=4\) and2\.47×2\.47\\times\(mb=8\) policy\-update speedup over FA2, comparable to the GRPO speedups \(1\.63×1\.63\\times/2\.09×2\.09\\times\)\. The slightly higher speedups reflect the absence of the ref phase: the policy update’s share of step time is larger in DAPO, so DualKV’s per\-phase compression translates more directly to end\-to\-end gains\. Step speedup is1\.63×1\.63\\times\(mb=4\) and1\.82×1\.82\\times\(mb=8\)\. MFU rises from30\.9%30\.9\\%\(FA2\) to77\.4%77\.4\\%\(DualKV mb=8\) — a2\.5×2\.5\\timesutilization gain\. Peak memory drops from106106GB to8080GB \(mb=4\) and9393GB \(mb=8\), confirming the same memory\-headroom story\. All four configs track identical training accuracy curves \(final reward≈0\.85\\approx 0\.85–0\.870\.87\), confirming no convergence degradation\.

Table 10:DAPO training on LongReason \(Qwen3\-8B, 8×\\timesH100\)\. Same setup as Table[3](https://arxiv.org/html/2605.15422#S4.T3)but with DAPO instead of GRPO\. Noref\_log\_probphase \(DAPO removes the KL penalty\)\. Speedups relative to FA2 mb=4\.∗Rollout\(vLLM inference\) is not targeted by DualKV\.

#### Comparison with GRPO\.

The per\-phase speedups are consistent with Section[4\.2](https://arxiv.org/html/2605.15422#S4.SS2): DualKV compressesold\_log\_probat∼2\.4×\{\\sim\}2\.4\\times\(vs∼2\.1×\{\\sim\}2\.1\\timesin GRPO\) and the policy update at1\.911\.91–2\.47×2\.47\\times\(vs1\.631\.63–2\.09×2\.09\\times\)\. The slightly larger DAPO speedups are attributable to \(i\) absence of the ref phase reducing the denominator, and \(ii\) DAPO’s token\-level loss aggregation producing marginally different gradient\-accumulation patterns that favor larger effective batch sizes\. The key takeaway is that DualKV uniformly accelerates any sampling\-based RL method that packsNNsame\-prompt responses\.

![Refer to caption](https://arxiv.org/html/2605.15422v1/x6.png)Figure 5:DAPO training on LongReason \(Qwen3\-8B, 8×\\timesH100\)\.\(A\)All four configurations track identical reward curves, confirming no convergence degradation\.\(B\)Peak memory: FA2/FA3 exceed the H100 80 GB physical limit \(∼106\{\\sim\}106GB allocated\), while DualKV mb=4 fits at∼80\{\\sim\}80GB\.\(C\)Policy\-update latency: DualKV mb=8 achieves2\.47×2\.47\\timesspeedup over FA2\.\(D\)MFU rises from31%31\\%\(FA2\) to77%77\\%\(DualKV mb=8\)\.

### E\.3Context\-Length Scaling Sweep \(Llama\-3\.1\-8B\)

To validate that DualKV’s benefits generalize beyond Qwen and to empirically characterize memory scaling, we conduct a context\-length sweep on Llama\-3\.1\-8B\-Instruct\[Dubey et al\.,[2024](https://arxiv.org/html/2605.15422#bib.bib5)\]across prompt lengthsP∈\{8​K,16​K,32​K,48​K,64​K,96​K\}P\\in\\\{8\\text\{K\},16\\text\{K\},32\\text\{K\},48\\text\{K\},64\\text\{K\},96\\text\{K\}\\\}and micro\-batch sizesmb∈\{4,8,16\}\\text\{mb\}\\in\\\{4,8,16\\\}\.

#### Setup\.

Hardware: 2×\\timesp5\.48xlarge\(16×\\timesH100\-80GB\)\. GRPO withN=32N\{=\}32,Rmax=2048R\_\{\\max\}\{=\}2048,train\_batch\_size=128\{=\}128, vLLM rollout \(TP=2\{=\}2\), DualKV \+ FusedLinearForPPO enabled,gpu\_memory\_utilization=0\.4\{=\}0\.4\.

#### Memory scaling law\.

Figure[6](https://arxiv.org/html/2605.15422#A5.F6)plots peak GPU memory \(max\_memory\_allocated\) against prompt length for all 18 configurations\. DualKV’s memory follows:

memDualKV=M0⏟model \+ optim \+framework\+cP⋅P⏟prompt activations\(paid once\)\+cmb​R⋅mb⋅R⏟response activations\(scales with mb\)\\text\{mem\}\_\{\\text\{DualKV\}\}=\\underbrace\{M\_\{0\}\}\_\{\\begin\{subarray\}\{c\}\\text\{model \+ optim \+\}\\\\ \\text\{framework\}\\end\{subarray\}\}\+\\underbrace\{c\_\{P\}\\cdot P\}\_\{\\begin\{subarray\}\{c\}\\text\{prompt activations\}\\\\ \\text\{\(paid once\)\}\\end\{subarray\}\}\+\\underbrace\{c\_\{\\text\{mb\}R\}\\cdot\\text\{mb\}\\cdot R\}\_\{\\begin\{subarray\}\{c\}\\text\{response activations\}\\\\ \\text\{\(scales with mb\)\}\\end\{subarray\}\}\(38\)where the coefficients are estimated by least\-squares regression across all 18 data points \(66prompt lengths×\\times33micro\-batch sizes\):M0≈41\.1M\_\{0\}\\approx 41\.1GB,cP≈0\.475c\_\{P\}\\approx 0\.475GB/Ktoken,cmb​R≈0\.122c\_\{\\text\{mb\}R\}\\approx 0\.122GB/Ktoken \(r2=0\.989r^\{2\}=0\.989\)\. The asymmetrycP≫cmb​Rc\_\{P\}\\gg c\_\{\\text\{mb\}R\}arises because prompt tokens incur higher per\-token activation cost under gradient checkpointing: block\-boundary activations are retained for allPPpositions across all layers, and the context self\-attention backward \(Call 1\) savessoftmax\_lseof sizeO​\(P×H\)O\(P\\times H\)\. Response tokens in Call 2 attend only to their own shorter decoded subsequence, requiring less stored state\. The key insight is structural: the prompt termcP⋅Pc\_\{P\}\\cdot Pdoes not multiply by mb — DualKV processes the shared prompt once regardless of how many responses are in the micro\-batch\. Under standard FA2 packing, each of the mb sequences carries its own full copy of the prompt, so memory scales as:

memFA2=M0⏟same\+cP⋅mb⋅\(P\+R\)⏟prompt replicated mb times\\text\{mem\}\_\{\\text\{FA2\}\}=\\underbrace\{M\_\{0\}\}\_\{\\text\{same\}\}\+\\underbrace\{c\_\{P\}\\cdot\\text\{mb\}\\cdot\(P\+R\)\}\_\{\\text\{prompt replicated mb times\}\}\(39\)making memory*linear in mb*\. FA2 values in Figure[6](https://arxiv.org/html/2605.15422#A5.F6)are projections from Eq\.[39](https://arxiv.org/html/2605.15422#A5.E39)\(using the same fittedM0M\_\{0\}andcPc\_\{P\}\) rather than measurements, because FA2 triggers CUDA OOM at these context lengths — the standardmb×\(P\+R\)\\text\{mb\}\\times\(P\+R\)packing exceeds the H100’s 80 GB HBM for most configurations in this sweep\.

#### Practical implication\.

WhenP≫mb⋅RP\\gg\\text\{mb\}\\cdot R\(the long\-prompt regime common in agentic and reasoning tasks\), memory is dominated by the prompt term and becomes effectively independent of micro\-batch size\. Practitioners can freely increase mb to maximize GPU utilization without memory penalty: atP=96​KP\{=\}96\\text\{K\}, increasing mb from 4 to 16 adds only∼4\{\\sim\}4GB \(<5%<5\\%\), whereas FA2 would require225→775225\\to 775GB — nearly10×10\\timesthe physical capacity of an H100\.

![Refer to caption](https://arxiv.org/html/2605.15422v1/x7.png)Figure 6:Memory scaling law: DualKV \(measured, solid\) vs\. FA2 \(projected, dashed\)\. Llama\-3\.1\-8B,N=32N\{=\}32,R=2048R\{=\}2048, 16×\\timesH100\. DualKV’s three micro\-batch curves cluster tightly \(mem∝P\+mb⋅R\\text\{mem\}\\propto P\+\\text\{mb\}\\cdot R\), while FA2’s diverge \(mem∝mb⋅\(P\+R\)\\text\{mem\}\\propto\\text\{mb\}\\cdot\(P\+R\)\)\. FA2 curves are projected from the fitted model since FA2 OOMs at these configurations\.

### E\.4Comparison with Prefix Grouper: Single\-Layer Benchmark

To empirically validate DualKV’s advantage over framework\-level prefix sharing, we benchmark against Prefix Grouper\[Liu et al\.,[2025](https://arxiv.org/html/2605.15422#bib.bib17)\]on a single Qwen3\-8B decoder layer \(forward \+ backward\) on one H100\-80GB GPU\. Prefix Grouper decomposes attention into prefix self\-attention and suffix cross\-attention \(2×\\timesFA2 calls\), but operates on*padded*2D tensors: allNNsequences are padded to the longest response length and the full padded tensor passes through QKV projections, output projection, LayerNorm, and MLP\. We give Prefix Grouper the best possible backend \(paddedflash\_attn\_func, not the default HuggingFace wrapper\) and compare against both standard FA2 varlen \(the baseline\) and DualKV\.

#### Setup\.

Qwen3\-8B dimensions \(H=32H\{=\}32,Hk=8H\_\{k\}\{=\}8,d=128d\{=\}128\), bf16, response lengthsR∼U​\(131,2048\)R\\sim U\(131,2048\)\. We sweep prompt lengthsP∈\{5059,8192,16384,32768,65536,131072\}P\\in\\\{5059,8192,16384,32768,65536,131072\\\}and micro\-batch sizesmb∈\{1,2,4,8,16,32\}\\text\{mb\}\\in\\\{1,2,4,8,16,32\\\}\. Timing: 2 warmup \+ 5 measured iterations\. Correctness verified viatorch\.allcloseagainst FA2 at smaller configs \(both DualKV and PG pass; PG has zero error since it uses FA2 internally\)\.

Table 11:Single\-layer forward\+backward: DualKV vs\. Prefix Grouper \(PG\) vs\. FA2\. Time in ms, memory in GB\. “DK/FA2” and “DK/PG” are DualKV’s speedup ratios\. OOM = out of memory on H100\-80GB\.PG≈\\approxFA2 in latency\.Prefix Grouper’s attention optimization \(avoiding redundant prefix\-on\-prefix QK⊤\) is offset by the overhead of padding, KV expansion \(expand \+ contiguous\), and concatenation\. Across all configs where both run, PG is within±\\pm20% of FA2 — sometimes slightly faster \(attention savings dominate at largePP\), sometimes slightly slower \(padding overhead dominates at large mb\)\.

DualKV is 2–4×\\timesfaster than both\.The speedup comes from processing fewer tokens through*every*operation in the layer\. At the paper’s primary config \(P=8​KP\{=\}8\\text\{K\}, mb=8\), DualKV processes 18K tokens vs\. FA2’s 75K and PG’s 81K — a4\.2×4\.2\\timesreduction that translates directly to a3\.5×3\.5\\timeswall\-clock speedup\.

Memory gap is catastrophic for PG\.PG allocates a padded\[mb,P\+Rmax,H\]\[\\text\{mb\},P\+R\_\{\\max\},H\]tensor, consuming33–5×5\\timesthe memory of DualKV’s packed representation\. AtP=8​KP\{=\}8\\text\{K\}mb=16, PG uses 50\.7 GB vs\. DualKV’s 9\.9 GB \(5\.1×5\.1\\timesreduction\); at mb=32, PG OOMs while DualKV fits in 14\.6 GB\.

Scalability\.Of 54 total configurations tested \(9 prompt lengths×\\times6 micro\-batch sizes\), DualKV runs 36 configs, while both FA2 and PG run only 21\. BeyondP=131​KP\{=\}131\\text\{K\}, only DualKV can scale past mb=1 on a single H100\.

Root cause: PG only optimizes attention\.Attention accounts for∼30%\{\\sim\}30\\%of a decoder layer’s FLOPs; the remaining 70% \(QKV projections, output projection, MLP, norms\) still operates on the full padded tensor withNNduplicated prompt copies\. DualKV eliminates redundancy across*all*operations by packing the prompt once in a varlen representation\.

## Appendix FAdditional Experimental Figures

![Refer to caption](https://arxiv.org/html/2605.15422v1/x8.png)Figure 7:Per\-phase breakdown of an average GRPO training step on LongReason \(Section[4\.2](https://arxiv.org/html/2605.15422#S4.SS2)\)\. DualKV compresses the three training\-side phases \(old\_log\_prob,ref\_log\_prob,policy update\) by∼2×\{\\sim\}2\\timeseach; rollout latency is identical across all four configs\.\(i\)the policy\-update slice shrinks from520520s \(FA2, FA3\) to247247s \(2\.1×2\.1\\timesover FA2\);\(ii\)the rollout slice is invariant \(∼247\{\\sim\}247s across all four configs\);\(iii\)rollout’s share of step time grows from23%23\\%\(FA2 mb=4\) to40%40\\%\(DualKV mb=8\) — DualKV shifts the bottleneck from the policy update to rollout generation\.![Refer to caption](https://arxiv.org/html/2605.15422v1/x9.png)Figure 8:Multi\-node MoE GRPO training on LongReason \(Section[4\.3](https://arxiv.org/html/2605.15422#S4.SS3); Qwen3\-30B\-A3B, 16×\\timesH100, 2 nodes\), comparing DualKV mb=8 SP=1 against FA2 mb=8 SP=4 over the common 1–29\-step window\.\(A\)Per\-step training reward — both configs track closely\.\(B\)Per\-step policy\-update latency — DualKV averages 1078 s vs FA2’s 4115 s \(3\.82×3\.82\\times\)\.\(C\)Per\-step total training step time — DualKV averages 1564 s vs FA2’s 5284 s \(3\.38×3\.38\\times\)\.![Refer to caption](https://arxiv.org/html/2605.15422v1/x10.png)Figure 9:Per\-phase breakdown of an average GRPO training step for the multi\-node MoE experiment \(Section[4\.3](https://arxiv.org/html/2605.15422#S4.SS3)\)\. DualKV compresses the three training\-side phases by3\.43\.4–3\.6×3\.6\\timeseach; rollout latency is identical across both configs\. The FA2 bar uses 4\-way Ulysses SP — the smallest SP that avoids OOM at this scale — while DualKV runs at SP=1\.
## Appendix GDualKV Kernel Algorithms

This appendix gives full pseudocode for the DualKV forward and backward kernels \(referenced from Sections[3\.1](https://arxiv.org/html/2605.15422#S3.SS1)–[3\.2](https://arxiv.org/html/2605.15422#S3.SS2)\)\. The listings elide K/V parallels: gradients, loads, and writes forVVmirror those forKKand are omitted for brevity\.

Algorithm 3DualKV forward kernel \(one thread block, assigned\(m,b,h\)\(m,b,h\)= query\-row block, batch, head\)1:

QmQ\_\{m\}for batch

bb, head

hh; shared

Kc,Vc∈ℝP×Hk×dK\_\{c\},V\_\{c\}\\in\\mathbb\{R\}^\{P\\times H\_\{k\}\\times d\}; per\-sequence

Kd\(b\),Vd\(b\)∈ℝRb×Hk×dK\_\{d\}^\{\(b\)\},V\_\{d\}^\{\(b\)\}\\in\\mathbb\{R\}^\{R\_\{b\}\\times H\_\{k\}\\times d\}; tile size

BNB\_\{N\};

PP\(=context\_seqlen\)

2:

nctx←⌈P/BN⌉n\_\{\\text\{ctx\}\}\\leftarrow\\lceil P/B\_\{N\}\\rceil
3:

ndec←⌈Rb/BN⌉n\_\{\\text\{dec\}\}\\leftarrow\\lceil R\_\{b\}/B\_\{N\}\\rceil
4:

nmax←nctx\+ndecn\_\{\\max\}\\leftarrow n\_\{\\text\{ctx\}\}\+n\_\{\\text\{dec\}\}
5:Apply causal pruning: reduce

nmaxn\_\{\\max\}so no tile exceeds the causal horizon of

QmQ\_\{m\}
6:Init

mi←−∞m\_\{i\}\\leftarrow\-\\infty,

ℓi←0\\ell\_\{i\}\\leftarrow 0,

𝐨i←0\\mathbf\{o\}\_\{i\}\\leftarrow 0⊳\\trianglerightonline\-softmax state per row inQmQ\_\{m\}

7:for

n=nmax−1n=n\_\{\\max\}\{\-\}1down to

0do⊳\\trianglerightiterate K\-tiles in reverse \(FA2 convention\)

8:if

n<nctxn<n\_\{\\text\{ctx\}\}then⊳\\trianglerightcontext tile: shared buffer

9:

Kn,Vn←Kc\[nBN:nBN\+BN\],Vc\[…\]K\_\{n\},V\_\{n\}\\leftarrow K\_\{c\}\[nB\_\{N\}:nB\_\{N\}\+B\_\{N\}\],V\_\{c\}\[\\ldots\]
10:

jbase←n⋅BNj\_\{\\text\{base\}\}\\leftarrow n\\cdot B\_\{N\}⊳\\trianglerightlogical K position

11:else⊳\\trianglerightdecoded tile: per\-sequence buffer

12:

Kn,Vn←Kd\(b\)\[\(n−nctx\)BN:…\]K\_\{n\},V\_\{n\}\\leftarrow K\_\{d\}^\{\(b\)\}\[\(n\{\-\}n\_\{\\text\{ctx\}\}\)B\_\{N\}:\\ldots\]
13:

jbase←P\+\(n−nctx\)⋅BNj\_\{\\text\{base\}\}\\leftarrow P\+\(n\-n\_\{\\text\{ctx\}\}\)\\cdot B\_\{N\}
14:endif

15:Mask out\-of\-bounds lanes \(partial tile at last context / last decoded tile\)

16:

S←Qm​Kn⊤/dS\\leftarrow Q\_\{m\}K\_\{n\}^\{\\top\}/\\sqrt\{d\}
17:Apply causal mask using

jbasej\_\{\\text\{base\}\}\(logical K position\) vs\.

QmQ\_\{m\}’s logical row position

18:

minew←max⁡\(mi,rowmax​\(S\)\)m\_\{i\}^\{\\text\{new\}\}\\leftarrow\\max\(m\_\{i\},\\;\\text\{rowmax\}\(S\)\)
19:

P←exp⁡\(S−minew\)P\\leftarrow\\exp\(S\-m\_\{i\}^\{\\text\{new\}\}\)
20:

α←emi−minew\\alpha\\leftarrow e^\{m\_\{i\}\-m\_\{i\}^\{\\text\{new\}\}\}⊳\\trianglerightonline\-softmax rescale factor

21:

ℓi←α​ℓi\+rowsum​\(P\)\\ell\_\{i\}\\leftarrow\\alpha\\,\\ell\_\{i\}\+\\text\{rowsum\}\(P\)
22:

𝐨i←α​𝐨i\+P⋅Vn\\mathbf\{o\}\_\{i\}\\leftarrow\\alpha\\,\\mathbf\{o\}\_\{i\}\+P\\cdot V\_\{n\}
23:

mi←minewm\_\{i\}\\leftarrow m\_\{i\}^\{\\text\{new\}\}
24:endfor

25:

𝐨i←𝐨i/ℓi\\mathbf\{o\}\_\{i\}\\leftarrow\\mathbf\{o\}\_\{i\}/\\ell\_\{i\}⊳\\trianglerightepilogue: normalize

26:Write

𝐨i\\mathbf\{o\}\_\{i\}to

o​\[m\]o\[m\]
27:Save

L=mi\+log⁡ℓiL=m\_\{i\}\+\\log\\ell\_\{i\}tosoftmax\_lse\[m\]⊳\\trianglerightused by backward

The backward main kernel \(Algorithm[1](https://arxiv.org/html/2605.15422#alg1)\) and the post\-kernel context\-gradient cast \(Algorithm[2](https://arxiv.org/html/2605.15422#alg2)\) are in Section[3\.2](https://arxiv.org/html/2605.15422#S3.SS2)\.

## Appendix HKernel Interface: Python Signature and Tensor Shapes

This appendix gives the full Python signature of the DualKV kernel \(referenced from Section[3](https://arxiv.org/html/2605.15422#S3)\) and the per\-tensor shape contract\.

out=flash\_attn\_dualkv\_varlen\_func\(

q,

k\_context,v\_context,

k\_decoded,v\_decoded,

cu\_seqlens\_q,

cu\_seqlens\_k\_decoded,

max\_seqlen\_q,

context\_seqlen,

max\_seqlen\_k\_decoded,

softmax\_scale=None,

causal=True,

\)

Listing 1:Signature of DualKV kernel interfaceflash\_attn\_dualkv\_varlen\_funcThe key contract is thatkctxk\_\{\\text\{ctx\}\}andvctxv\_\{\\text\{ctx\}\}\(k\_context,v\_contextin the listing\) are stored once \(not replicatedNNtimes\), whileqq,kdeck\_\{\\text\{dec\}\}, andvdecv\_\{\\text\{dec\}\}are varlen\-packed over theNNdecoded sequences\. The kernel is implemented as a PyTorchtorch\.autograd\.Function\. Internally, the forward also computes a per\-query softmax log\-sum\-exp tensorsoftmax\_lseof shape\(H,∑iRi\)\(H,\\sum\_\{i\}R\_\{i\}\)that is saved onto the autograd context for the backward pass \(Section[3\.2](https://arxiv.org/html/2605.15422#S3.SS2)\); it is not part of the user\-facing return value\.

## Appendix IComparison with Other Attention Variants

Readers familiar with attention\-optimization literature may ask whether existing variants — paged attention, prefix caching, bifurcated attention, sequence\-parallel methods, or sparse/linear attention — already address the shared\-prompt redundancy we target\. This appendix systematically positions DualKV against these variants along four capabilities relevant to RL policy\-update training:

1. 1\.Shared\-prompt KV dedup: the method stores the shared prompt’sK,VK,Vonce rather thanNNtimes\.
2. 2\.Shared\-prompt compute dedup: the method also processes the shared prompt through per\-token operations \(RMSNorm, QKV projection, RoPE, MLP, output projection\)*once*rather thanNNtimes — i\.e\., eliminates the\(N−1\)⋅P\(N\{\-\}1\)\\cdot Ptokens of compute outside attention\.
3. 3\.Training backward \+ autograd integration: the method provides a differentiable backward pass, with gradients for all inputs compatible with PyTorch autograd \(including gradient accumulation into the sharedKc,VcK\_\{c\},V\_\{c\}across theNNconcurrent sequences\)\.
4. 4\.Mathematically exact: the method produces outputs identical to standard attention up to floating\-point rounding, not an approximation\.

Table[12](https://arxiv.org/html/2605.15422#A9.T12)summarizes the comparison\. DualKV is the only variant that satisfies all four criteria simultaneously\.

Table 12:Comparison of attention variants along four capabilities relevant to RL training\.✓\\checkmark= satisfies;×\\times= does not\. DualKV is the only variant that is✓\\checkmarkon all four\.#### Inference\-only variants \(paged, RadixAttention, bifurcated\)\.

These methods share the shared\-prompt KV at*storage*level \(via block tables or memory\-layout tricks\) during inference\. None defines a backward pass for training\. Naively invoking a paged or bifurcated kernel in a training loop would fail because \(i\) PyTorch autograd cannot differentiate through block\-table pointer indirection — the shared KV must be a proper tensor with a computation graph; \(ii\) withNNconcurrent sequences backward\-propagating into the*same*Kc,VcK\_\{c\},V\_\{c\}, the gradient accumulation requires race\-free concurrent writes with precision preservation \(DualKV’s fp32 atomic accumulation, Section[3\.2](https://arxiv.org/html/2605.15422#S3.SS2)\); and \(iii\) they address attention only, leaving the\(N−1\)⋅P\(N\{\-\}1\)\\cdot Pper\-token cost in MLP, norms, and projections unaffected\.

#### Sequence\-parallel variants \(Ring, Ulysses\)\.

These are orthogonal techniques for long\-sequence training: they shard a*single*long sequence across multiple ranks via collective communication\. They do not deduplicate the*shared prompt across multiple sequences*— that is a different axis of redundancy\. DualKV composes with sequence parallelism; combining the two is an interesting future direction\.

#### Approximate variants \(linear, sparse, compressed\)\.

These trade accuracy for efficiency by replacing attention withO​\(S\)O\(S\)approximations or by pruning KV entries\. They violate DualKV’s “mathematically equivalent to standard attention” guarantee, and are orthogonal to shared\-prompt dedup — they can be applied per\-sequence independently of whether sequences share a prompt\.

#### Summary\.

DualKV is the first attention variant to jointly provide \(1\) shared\-prompt KV dedup, \(2\) shared\-prompt compute dedup through the full transformer layer, \(3\) a mathematically exact training backward with gradient accumulation into the sharedKc,VcK\_\{c\},V\_\{c\}, and \(4\) autograd\-compatible integration with PyTorch\. No existing technique can be dropped into a RL policy\-update pipeline to replicate DualKV’s savings\.