arXiv cs.AI· Geng Zhang, Fangzhao Zhang, Cong Xie, Tiantian Fan, Yang You·· 11 小时前AI 评分24
Split-KV：面向扩散草稿模型训练的高效上下文并行方法
Split-KV: Efficient Context Parallelism for Training Diffusion Drafters
AI 导读
Split-KV 是一种精确的上下文并行（CP）方法，将 key/value 分片保留在本地，把每层注意力通信限制在 query 规模的张量上，使通信量与上下文长度无关。在 CP 规模为 64 时，它保持 85–91% 的弱扩展效率，相比最快的 CP 基线最高实现 3.67 倍加速。其基于全局 query 视图的 LPT 负载均衡策略将不均衡度控制在 1% 以下，规划速度比最佳基线快三个数量级。
正文 · 原文
Abstract: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 this https URL.
| Subjects: | Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2611.05926 [cs.AI] |
| (or arXiv:2611.05926v1 [cs.AI] for this version) | |
| https://doi.org/10.48550/arXiv.2611.05926 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Geng Zhang [view email]
[v1]
Wed, 4 Nov 2026 04:16:46 UTC (272 KB)
来源：arXiv cs.AI · arxiv.org