Split-KV：面向 Diffusion Drafter 训练的高效上下文并行方法
Split-KV: Efficient Context Parallelism for Training Diffusion Drafters
AI 导读
arXiv 论文提出 Split-KV，一种用于训练 Diffusion drafter 的精确上下文并行（CP）方法，将 key/value 分片保留在本地，使每层 attention 通信限制在 query 大小的张量、与上下文长度无关。
正文 · 原文
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