Split-KV：面向 Diffusion Drafter 训练的高效上下文并行方法
Split-KV: Efficient Context Parallelism for Training Diffusion Drafters
训练 diffusion drafter 的人可以看看，Split-KV 在 CP 64 下最高加速 3.67 倍，负载均衡还快三个数量级，代码已开源在 GitHub。
Diffusion drafter 通过并行去噪多个草稿 token 加速投机解码，其训练中草稿查询远短于上下文，但现有上下文并行（CP）方法无法利用这种不对称性。论文提出 Split-KV，一种精确 CP 方法，将 key/value 分片保留在本地，把每层注意力通信限制在查询大小的张量，使通信量与上下文长度无关。在 CP size 为 64 时，Split-KV 保持 85–91% 的弱扩展效率，相比最快的 CP 基线实现最高 3.67 倍加速。其基于 LPT 的负载均衡策略不平衡度低于 1%，规划速度比最佳基线快三个数量级。
Split-KV: Efficient Context Parallelism for Training Diffusion Drafters
Diffusion drafters accelerate speculative decoding by denoising multiple draft tokens in parallel. Their training uses anchor sampling to construct compact draft queries and performs cross-attention with a block-sparse mask over a much longer context. We empirically show that acceptance length remains stable as context scales with a fixed query budget, while increasing the query budget yields diminishing quality gains, jointly establishing a practical regime with queries far shorter than the context. No existing context parallelism (CP) method jointly exploits this query--context asymmetry and meets the nonlinear, sample-dependent load-balancing requirement created by anchor sampling. We present Split-KV, an exact CP method that keeps key/value shards local and limits per-layer attention communication to query-sized tensors, making this communication independent of context length. Its global query view further enables a cost-aware longest-processing-time-first (LPT) strategy that balances context key/value chunks by their realized attention cost. At a CP size of 64, Split-KV retains 85--91\% weak-scaling efficiency and achieves up to $3.67\times$ speedup over the fastest available CP baseline. LPT achieves $<$1\% imbalance and plans three orders of magnitude faster than the best load-balancing baseline. Code is available at https://github.com/zxgx/verl.