Split-KV：破解扩散草稿模型训练中的上下文并行瓶颈

扩散草稿模型（Diffusion Drafter）能在单次前向中并行去噪多个草稿 token，是加速投机解码的重要路线。但其训练过程有一个容易被忽视的结构性特点：通过锚点采样构造短小精悍的草稿查询，再对远长于查询的上下文做块稀疏交叉注意力。这种查询极短、上下文极长的非对称性，让传统的上下文并行（CP）方法难以直接套用。Split-KV 正是针对这一矛盾提出的精确 CP 方案，在 CP 规模为 64 时仍保持 85–91% 的弱扩展效率，相比最快的现有 CP 基线最高实现 3.67 倍加速。

查询与上下文的非对称性：一个被忽略的训练区间

研究团队首先通过实证给出了一组关键观察：当查询预算固定时，接受长度（acceptance length）随着上下文规模的增长保持稳定；而单纯扩大查询预算，带来的质量收益却呈现递减。两条曲线交汇出一个实用的训练区间——查询长度远小于上下文长度。

这一发现的意义在于，它把「上下文并行」的问题重新定义了一遍。如果查询本身很短，那么注意力的通信量理论上就不应随上下文长度线性膨胀。但现有 CP 方法大多默认查询与上下文同量级，或者没有利用这种非对称性来压缩通信。

更棘手的是负载均衡。锚点采样使得每个样本实际参与的上下文范围不同，注意力计算代价呈现出非线性、样本依赖的波动。静态切分上下文的方式在这种波动面前很容易失衡，而 Split-KV 的设计正是围绕这两个痛点展开。

投机解码（Speculative Decoding） 是一种在推理阶段加速大模型生成的技术：用一个轻量级的「草稿模型」先快速生成多个候选 token，再由目标大模型并行验证，从而在保证输出分布不变的前提下减少目标模型的串行前向次数。接受长度（acceptance length） 指的是每一轮投机过程中，草稿 token 被目标模型实际接受的平均数量，它直接决定了加速比的高低——接受长度越长，加速越明显。

扩散草稿模型 将草稿生成建模为扩散去噪过程，因而可以在单次前向中并行产出一批草稿 token，突破了自回归草稿模型逐个生成的限制。但并行去噪也意味着草稿查询的组织方式需要重新设计，这正是锚点采样与块稀疏交叉注意力出现的原因。

Split-KV 的核心机制：让通信量与上下文长度解耦

Split-KV 是一种精确（exact）的 CP 方法，其核心思路是：把 key/value 分片保留在各设备本地，每层注意力只通信与查询规模相当的张量。由于查询远短于上下文，这一通信量就与上下文长度脱钩了。

具体来说，各 rank 持有自己的 KV 分片，在注意力计算中通过全局查询视图汇总各分片的局部注意力结果。通信的是查询侧的张量而非完整的 KV，这在长上下文场景下带来的节省是数量级层面的。

这一设计还带来一个额外好处：全局查询视图让系统能够感知每个 KV 分片的实际注意力代价，从而为负载均衡提供了可操作的依据。

上下文并行（Context Parallelism, CP） 是训练长上下文模型时常用的一种并行维度：把序列维度上的上下文切分到多个设备，每张卡只负责一段 KV，但注意力计算需要跨越所有分片，因此设备间必须交换信息。传统 CP 方法（如环形注意力 Ring Attention）通常需要在每层传输完整的 KV 或其分块，通信量随上下文长度线性增长。

Split-KV 的名称正源于这种切分方式——KV 在设备间被拆分（split）保存，而注意力结果通过查询侧张量进行汇总，避免了 KV 的全量通信。这里的 exact（精确） 是相对近似方法而言：Split-KV 不牺牲数值正确性，计算结果与不分片时严格一致，这对训练阶段的梯度正确性至关重要。

LPT 调度：用真实注意力代价做负载均衡

有了全局查询视图，Split-KV 采用了一种代价感知的最长处理时间优先（LPT）策略来分配上下文 KV 块。与按 token 数量粗略切分的传统做法不同，LPT 依据每个 KV 块实际产生的注意力代价来平衡各 rank 的负载。

效果相当直接：负载不均衡率低于 1%，而规划速度比最好的负载均衡基线快三个数量级。对于需要频繁重新规划的锚点采样训练来说，规划开销本身就不可忽视，这一提速具有实际工程价值。

LPT（Longest Processing Time first，最长处理时间优先） 是经典的多机调度启发式算法：将任务按处理时间从长到短排序，依次分配给当前负载最轻的机器，从而获得接近最优的 makespan。它在理论上具有 4/3 的近似比保证。

Split-KV 的贡献在于把 LPT 中的「处理时间」替换为每个 KV 块真实的注意力计算代价，而非粗略的 token 计数。由于锚点采样使每个样本参与的上下文范围不同，块与块之间的代价差异很大，这种代价感知的排序能显著提升均衡质量。弱扩展效率（weak scaling efficiency） 衡量的是在保持单设备工作量不变、成比例增加设备数时，实际吞吐相对理想线性增长的比值，是评估并行方案可扩展性的核心指标。

扩展性表现与开放代码

在 CP 规模为 64 的配置下，Split-KV 保持了 85–91% 的弱扩展效率，相比现有最快的 CP 基线最高取得 3.67 倍加速。考虑到扩散草稿模型训练中上下文长度持续增长的现实需求，这种扩展效率意味着增加设备仍能带来接近线性的收益。

代码已在 GitHub 的 verl 仓库开源（https://github.com/zxgx/verl），对于正在训练扩散草稿模型、又受困于长上下文并行效率的团队，这是一个值得直接评估的实现。

从更广的视角看，Split-KV 的价值不只在于一个具体的并行策略，而在于它示范了一种思路：先识别工作负载的结构性非对称，再让通信和调度围绕这种非对称来设计。当上下文继续变长、查询预算继续收紧时，这种「以查询为中心」的并行视角可能会成为扩散草稿模型训练基础设施的默认选项。