DualKV: Computing Shared Prompts Once Instead of N Times in RL Training of LLMs

Forum topic · 小凯 · 2026-05-19

Summary

A zhichai.net forum post discusses DualKV, a technique by Gai, Zhang, Song, Wang, and Karypis that eliminates redundant computation in RL post-training of large language models such as GRPO and DAPO. In these methods, N sampled responses share the same prompt, yet standard FlashAttention recomputes the prompt's attention N times. DualKV introduces a fused CUDA kernel that iterates over two disjoint KV regions—the shared prompt context and per-sequence response tokens—computing prompt KV only once, plus a data pipeline that reshapes a micro-batch from N(P+R) tokens to P+NR tokens, with compression factor rho = N(P+R)/(P+NR) applied to the entire model. The method is mathematically equivalent to standard attention with no approximation. Benchmarks: Qwen3-8B GRPO training on 8x H100 (N=32, 8K context) achieves 1.63-2.09x policy-update speedup, 2x larger micro-batches, and MFU rising from 36% to 76%; DAPO gains 2.47x speedup with 77% MFU; a 30B MoE model on 16x H100 sees 3.82x policy-update and 3.38x end-to-end speedups. Open questions include portability beyond H100, benefits at small N, and effectiveness for long prompts with short responses.

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

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

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

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

Benchmark results

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

Open questions

Does DualKV's implementation rely on H100-specific features, or can it be ported to other GPU architectures?
N is typically large in GRPO (32+); if N is small (e.g., 2–4), the benefit of sharing the prompt shrinks considerably.
KV sharing adds kernel complexity — how well does the optimization perform for scenarios 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.