Split-KV：用于训练扩散草稿模型的高效上下文并行方法

扩散草稿模型通过并行对多个草稿词元进行去噪来加速推测解码。其训练采用锚点采样来构建紧凑的草稿查询，并在长得多的上下文上使用块稀疏掩码执行交叉注意力。我们的实证结果表明，在查询预算固定时，接受长度随着上下文规模增大而保持稳定，而增加查询预算带来的质量收益则逐渐递减。这两点共同确立了一种查询远短于上下文的实用运行区间。现有上下文并行（CP）方法均无法同时利用这种查询与上下文之间的不对称性，并满足锚点采样所带来的非线性、依赖样本的负载均衡要求。我们提出了 Split-KV，这是一种精确的 CP 方法，将键/值分片保留在本地，并将每层的注意力通信限制为查询规模的张量，从而使该通信量与上下文长度无关。其全局查询视图进一步支持一种成本感知的最长处理时间优先（LPT）策略，根据实际产生的注意力计算成本，对上下文键/值块进行负载均衡。在 CP 规模为 64 时，Split-KV 保持了 85–91% 的弱扩展效率，相比当前最快的 CP 基线实现了最高 3.67 倍的加速。LPT 将负载不均衡率降至 1% 以下，且规划速度比最佳负载均衡基线快了三个数量级。代码已公开：https://github.com/zxgx/verl。