Skip to content

Home

Jump to bottom
Jiading Gai edited this page May 27, 2026 · 6 revisions

DualKV: Shared-Prompt Flash-Attention for RL Training

Overview

DualKV eliminates redundant prompt computation in GRPO/DAPO training. Instead of computing attention over N*(P+R) tokens (where N responses share the same prompt P), DualKV computes over P + N*R — yielding up to 6x kernel speedup and 2x end-to-end throughput on long-context RL workloads.

Paper

DualKV: Shared-Prompt Flash-Attention Kernels for Efficient Policy Updates in RL Training

How It Works

Modern sampling-based RL training (GRPO, DAPO etc) generates N responses per prompt, then updates the policy on all of them. With flash-attention, each response independently attends to its own copy of the prompt KV — wasting compute and memory.

DualKV introduces a dual-KV kernel that factorizes flash attention into:

  1. Shared prompt self-attention — computed once for the group
  2. Per-response cross-attention — each response attends to both the shared prompt KV and its own decoded KV simultaneously

This is implemented as a custom CUDA kernel extending FlashAttention-2, with full forward + backward support for training.

Architecture

┌─────────────────────────────────────────────────┐
│  dp_actor.py: _dualkv_repack()                  │
│  [P+R1, P+R2, ..., P+RN] → [P, R1, R2, ..RN]  │
└──────────────────────┬──────────────────────────┘
                       │
┌──────────────────────▼──────────────────────────┐
│  monkey_patch.py: _dualkv_flash_forward()       │
│  Routes to DualKV kernel when context present   │
│  Supports Ulysses SP (all-to-all before/after)  │
└──────────────────────┬──────────────────────────┘
                       │
┌──────────────────────▼──────────────────────────┐
│  flash_attn_dualkv_varlen_func()                │
│  CUDA kernel: fwd + bwd with dual KV streams   │
└─────────────────────────────────────────────────┘

Key Results (H100-80GB)

Config Kernel Speedup E2E Throughput
P=4K, N=28 1.6x 1.4x
P=16K, N=28 3.7x 1.8x
P=32K, N=28 6.0x 2.1x
P=64K, N=16 ∞ (FA2 OOM) —

Kernel Design

The DualKV kernel extends FlashAttention-2's compute_attn_1rowblock to handle two KV streams simultaneously:

Core idea: For each Q row block, the kernel iterates over K/V tile blocks drawn from two separate memory regions — context (shared) and decoded (per-sequence) — while maintaining a single running softmax accumulator.

Block layout in the inner loop:

blocks 0 .. n_ctx-1          → load from kcontext_ptr (shared, bs=1)
blocks n_ctx .. n_ctx+n_dec-1 → load from kdecoded_ptr (varlen, per-sequence)

Key modifications from FA2:

Aspect FA2 DualKV
KV source Single contiguous buffer Two pointers: kcontext_ptr + kdecoded_ptr
KV offsets cu_seqlens_k[bidb] Context: no batch offset; Decoded: cu_seqlens_k_decoded[bidb]
Causal mask col <= row col <= context_seqlen + decoded_row (shifted for dual layout)
Block iteration Single range [0, n_blocks) Two ranges stitched: context blocks then decoded blocks
Memory savings N copies of prompt KV 1 copy of prompt KV (N× reduction)

What was stripped: Dropout, ALiBi, softcap, local/sliding window, paged KV — these complicate the dual-pointer logic and aren't needed for RL training. This keeps the kernel lean and the tile scheduling simple.

Backward kernel: Mirrors the forward structure — computes dQ by iterating over both context and decoded KV blocks, and computes dK/dV separately for context (accumulated across all N sequences) and decoded (per-sequence).

Ulysses SP Composition

DualKV composes with Ulysses Sequence Parallelism for additional memory savings. The integration:

  • Repacks tokens before the Ulysses sequence slice
  • Performs all-to-all (gather-seq, scatter-heads) inside the attention wrapper
  • Runs the DualKV kernel on the full token sequence with partial heads
  • Scatters back after attention

This enables training with both longer contexts and larger micro-batches simultaneously.

Future Work

  • Hopper/Blackwell backend — Port DualKV to use TMA (Tensor Memory Accelerator) and warp-specialized persistent kernels for H100/B200.
  • FP8 support — Extend the kernel to FP8 E4M3/E5M2 for next-gen hardware (H100/B200) where FP8 attention is the default
  • Multi-node DualKV + SP — Scale to SP=4/8 across nodes with 128K+ prompts, combining DualKV deduplication with ring/Ulysses parallelism
  • DPO/RLHF integration — Extend beyond GRPO/DAPO to DPO (chosen/rejected pairs share prompts) and RLHF (critic also shares prompts with actor)
  • Online/iterative RL — Support streaming prompt groups where N varies dynamically across micro-batches
  • Context parallelism composition — Explore combining DualKV with CP (ring attention) for ultra-long prompts that exceed single-device memory
  • Prefix caching synergy — Combine DualKV training kernels with vLLM-style prefix caching during rollout for end-to-end shared-prompt optimization
  • Variable group sizes — Support mixed group sizes within a batch (some prompts with N=32, others with N=8) for heterogeneous workloads
  • Backward kernel optimization — DualKV-pointer logic introduces extra branching in the inner loop; leveraging CUDA coding agents to hoist branches out of the critical path.
  • MCTS tree search — In tree-based RL (e.g., AlphaCode, LATS, rStar), nodes at the same depth share a common prefix. DualKV can deduplicate KV across sibling branches, turning O(BxD) attention into O(D + Bxleaf) where D is shared depth — natural fit for beam/tree rollouts
  • Upstream to veRL/flash attention — Contribute DualKV as a first-class feature in the veRL framework with auto-detection of shared prompts

Quick Start

See the README for setup and experiment reproduction instructions.

Clone this wiki locally