arXiv cs.AI· Geng Zhang, Fangzhao Zhang, Cong Xie, Tiantian Fan, Yang You·· 10 小时前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% 以下，规划速度比最佳基线快三个数量级。
正文 · AI 翻译
摘要：扩散草稿器通过并行去噪多个草稿 token 来加速推测解码。其训练使用锚点采样来构建紧凑的草稿查询，并在长得多的上下文上使用块稀疏掩码执行交叉注意力。我们通过实验表明，在固定查询预算下，接受长度随上下文扩展保持稳定，而增加查询预算带来的质量提升递减，这共同确立了一个查询远短于上下文的实用区间。现有上下文并行（CP）方法均未同时利用这种查询—上下文不对称性，也未满足锚点采样所产生、依赖于样本的非线性负载均衡需求。我们提出 Split-KV，一种精确的 CP 方法，它保持键/值分片本地化，并将每层注意力通信限制为查询大小的张量，使该通信与上下文长度无关。其全局查询视图进一步支持一种成本感知的最长处理时间优先（LPT）策略，该策略按实际注意力成本平衡上下文键/值块。在 CP 规模为 64 时，Split-KV 保持 85--91\% 的弱扩展效率，并相比最快的可用 CP 基线实现最高 $3.67\times$ 的加速。LPT 实现 $<$1\% 的不平衡度，且规划速度比最佳负载均衡基线快三个数量级。代码见 此 https URL。
| 学科： | 人工智能（cs.AI） |
| 引用为： | arXiv:2611.05926 [cs.AI] |
| （或 arXiv:2611.05926v1 [cs.AI] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2611.05926 由 DataCite 颁发的 arXiv DOI（待注册） |
提交历史
来自：Geng Zhang [查看邮箱]
[v1]
2026年11月4日星期三 04:16:46 UTC（272 KB）
来源：arXiv cs.AI · arxiv.org