Research map / Architectures and efficiency / Attention and transformers

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

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes DualKV, a FlashAttention variant plus veRL data-pipeline change that processes the shared RL prompt once across N GRPO/DAPO rollouts, giving exact attention with 1.6-3.8x policy-update speedups.

Problem

Standard FlashAttention in RL post-training (GRPO, DAPO) replicates a shared P-token prompt across N response sequences in both forward and backward passes, wasting compute and memory at large rollout factors and long contexts.

Approach

Pack the micro-batch as one shared prompt plus N per-sequence responses, and use a fused FlashAttention forward/backward kernel that iterates over two disjoint KV regions (shared context and per-sequence response) with fp32 atomic accumulation of shared-prompt gradients, yielding exact attention.

Key results

• On Qwen3-8B GRPO with 8xH100 (N=32, 8K context), DualKV gives 2.09x policy-update speedup at mb=8 and raises MFU from 36% to 76%; DAPO reaches 2.47x.
• On Qwen3-30B-A3B across 16xH100, DualKV (SP=1) achieves 3.82x policy-update and 3.38x end-to-end step speedup over FA2 with SP=4.
• Kernel-level fwd+bwd speedup over FA2 ranges from 1.65x (P=4K) to 5.48x (P=32K, N=16), with up to 86% peak memory reduction.

Surprising findings

• Kernel speedup reaches 5.48x at P=32K, with FA2 OOM at 64K while DualKV runs.
• Gradients match FA2 only to half-precision rounding in the kernel test, though the training step is claimed exact in arithmetic.

Limitations

• Assumes a single shared prefix per micro-batch group; arbitrary tree-structured prefixes not supported.
• End-to-end gains are bounded by rollout generation time, which DualKV does not accelerate.
• Evaluations are limited to a few models and hardware setups (H100/H200/A100) and short training runs of 24-29 steps.

Experimental setup

Datasets: LongReason GSM8K

Benchmarks: LongReason GSM8K

Metrics: policy-update time end-to-end step time MFU peak memory validation accuracy kernel fwd/bwd wall-clock time

Baselines: FlashAttention-2 FlashAttention-3 Prefix Grouper Ulysses sequence parallelism

Builds on: FlashAttention-2 veRL GRPO DAPO Prefix Grouper

Closest papers by meaning

• Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models (ICLR 2025 · similarity 0.87)
• FlashMask: Efficient and Rich Mask Extension of FlashAttention (ICLR 2025 · similarity 0.86)
• DUET: Optimize Token-Budget Allocation for Reinforcement Learning with Verifiable Rewards (NeurIPS 2026 · similarity 0.86)
• MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention (NeurIPS 2024 · similarity 0.86)
• Adam on Local Time: Addressing Nonstationarity in RL with Relative Adam Timesteps (NeurIPS 2024 · similarity 0.86)
• DPC: Dual-Prompt Collaboration for Tuning Vision-Language Models (CVPR 2025 · similarity 0.86)