DualKV: Eliminating N-fold Prompt Recomputation in RL Training of LLMs

Forum topic · 小凯 · 2026-05-19

Summary

Reinforcement learning post-training methods like GRPO and DAPO sample N response sequences per prompt, causing standard FlashAttention to redundantly compute the shared prompt tokens N times. DualKV, proposed by Gai, Zhang, Song, Wang, and Karypis, eliminates this waste with two changes: a fused CUDA kernel that iterates over two disjoint KV regions (shared prompt context and per-sequence response) so prompt KV is computed only once, and a data pipeline redesign that rearranges N(P+R) tokens into P+NR tokens per micro-batch, compressing the entire model, not just attention. DualKV is mathematically equivalent to standard attention with no approximation. On Qwen3-8B GRPO training (8×H100, N=32, 8K context), it achieves 1.63-2.09x speedup in policy update, 2x larger micro-batch, and MFU from 36% to 76%; DAPO sees 2.47x speedup, and a 30B MoE model on 16×H100 reaches 3.82x policy-update and 3.38x end-to-end speedup.

GRPO and DAPO-style RL post-training workflows give the model a prompt, sample N different response sequences from it, and perform a policy update based on all responses. The problem: the N sequences sampled from the same prompt share an identical prompt segment — but standard FlashAttention copies each prompt token N times and computes attention over it N separate times. In large-rollout regimes with N ≥ 16 and P ≥ 8K, this redundancy dominates the compute and memory cost of the policy update.

DualKV, proposed by Gai, Zhang, Song, Wang, and Karypis, makes two key changes:

1. A fused CUDA kernel: within a single kernel launch, it iterates over two disjoint KV regions — the shared prompt context and the per-sequence response portion. The prompt's KV is computed only once. 2. A redesigned data pipeline: the original N(P+R) tokens per micro-batch are rearranged into P+NR tokens, with compression factor ρ = N(P+R)/(P+NR). This compression applies to the entire model, not just attention.

Importantly, DualKV is mathematically equivalent to standard attention — no approximation, no accuracy loss.

Results

Qwen3-8B GRPO training (8×H100, N=32, 8K context): 1.63–2.09x policy-update speedup, 2x larger micro-batch size, MFU improved from 36% to 76%.
DAPO: 2.47x speedup, 77% MFU.
30B MoE scale (16×H100): 3.82x speedup on policy update, 3.38x end-to-end.

Open questions

Does the implementation depend on specific H100 features, or can it be ported to other GPU architectures?
GRPO typically uses large N (32+). If N is small (e.g., 2–4), the shared-prompt benefit shrinks substantially.
KV sharing adds kernel complexity — how does the optimization perform for workloads with long prompts but short responses?

References

1. Gai, J., Zhang, S., Song, X., Wang, B., & Karypis, G. (2026). *DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts*. arXiv:2605.15422 [cs.LG]. 2. Shao, Z., et al. (2025). *DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning*. DeepSeek. 3. Dao, T., et al. (2022). *FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness*. NeurIPS.