RL 训练 LLM 时 prompt 被算了 N 次——DualKV 用一次计算共享全部

GRPO 和 DAPO 这类 RL 后训练方法的工作流程是：给模型一个 prompt，从中采样 N 条不同的响应序列，然后基于所有响应做策略更新。问题在于，从同一个 prompt 采样的 N 条序列共享完全相同的 prompt 部分——但标准 FlashAttention 把 prompt 的每个 token 复制了 N 份，分别计算 N 次 attention。在 N≥16、P≥8K 的大 rollout 场景中，这个冗余占据了策略更新的绝大部分计算和内存成本。

Gai、Zhang、Song、Wang 和 Karypis 提出的 DualKV 做了两个关键改动。第一个是融合 CUDA 内核：在单个内核启动中，迭代两个不相交的 KV 区域——共享的 prompt 上下文和每条序列独有的响应部分。Prompt 的 KV 只需要计算一次。第二个是数据流水线重设计：把一个微批次原来的 N(P+R) 个 token 重排为 P+NR 个 token——压缩因子 ρ = N(P+R)/(P+NR)。这个压缩不仅作用于 attention，而是作用于整个模型。

重要的是：DualKV 在数学上等价于标准 attention，没有任何近似。不损失精度。

在 Qwen3-8B 的 GRPO 训练中（8×H100，N=32，8K 上下文），DualKV 实现了 1.63-2.09 倍的策略更新加速，微批次大小提升 2 倍，MFU 从 36% 提升到 76%。DAPO 下加速 2.47 倍，MFU 77%。在 30B MoE 规模上（16×H100），加速达到了 3.82 倍（策略更新）、3.38 倍（端到端）。

不清楚的地方：DualKV 的实现是否依赖于特定的 H100 特性，还是可以移植到其他 GPU 架构？GRPO 中 N 值通常较大（32+），如果 N 很小（如 2-4），共享 prompt 的收益会大幅降低。KV 共享增加了内核的复杂性——对于长 prompt 但短 response 的场景，优化效果如何？

参考文献

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.