Split-KV: Breaking the Context Parallelism Bottleneck in Diffusion Drafter Training

Diffusion Drafters can denoise multiple draft tokens in parallel in a single forward pass, making them an important approach to accelerating speculative decoding. However, their training process has an easily overlooked structural characteristic: short, compact draft queries are constructed through anchor sampling, followed by block-sparse cross-attention over a context far longer than the query. This extremely short query, extremely long context asymmetry makes traditional context parallelism (CP) methods difficult to apply directly. Split-KV is an exact CP solution proposed precisely for this contradiction, maintaining 85–91% weak scaling efficiency even at a CP scale of 64, and achieving up to 3.67x speedup over the fastest existing CP baseline.

The Query-Context Asymmetry: An Overlooked Training Regime

The research team first presents a set of key empirical observations: when the query budget is fixed, the acceptance length remains stable as the context size grows; while simply expanding the query budget yields diminishing quality returns. The two curves intersect at a practical training regime—query length far smaller than context length.

The significance of this finding is that it redefines the problem of "context parallelism." If the query itself is short, then the communication volume of attention should in theory not grow linearly with context length. But most existing CP methods assume queries and context are of comparable scale, or do not exploit this asymmetry to compress communication.

More challenging is load balancing. Anchor sampling causes each sample to actually participate in a different range of context, so the attention computation cost exhibits nonlinear, sample-dependent fluctuation. Static context partitioning is easily unbalanced in the face of such fluctuation, and Split-KV's design revolves around these two pain points.

Speculative Decoding is a technique for accelerating large model generation during inference: a lightweight "draft model" first quickly generates multiple candidate tokens, which are then verified in parallel by the target large model, thereby reducing the number of serial forward passes of the target model while guaranteeing that the output distribution remains unchanged. Acceptance length refers to the average number of draft tokens actually accepted by the target model in each round of speculation, which directly determines the speedup—the longer the acceptance length, the more pronounced the acceleration.

Diffusion Drafters model draft generation as a diffusion denoising process, and can therefore produce a batch of draft tokens in parallel in a single forward pass, breaking the one-by-one generation limitation of autoregressive draft models. But parallel denoising also means the organization of draft queries needs to be redesigned, which is precisely why anchor sampling and block-sparse cross-attention emerged.

Split-KV's Core Mechanism: Decoupling Communication from Context Length

Split-KV is an exact CP method whose core idea is: keep key/value shards locally on each device, and have each attention layer communicate only tensors comparable in size to the query. Since the query is far shorter than the context, this communication volume becomes decoupled from context length.

Specifically, each rank holds its own KV shard and, during attention computation, aggregates the local attention results of each shard through a global query view. What is communicated is the query-side tensor rather than the full KV, and the savings this brings in long-context scenarios are on the order of magnitude.

This design also brings an additional benefit: the global query view enables the system to perceive the actual attention cost of each KV shard, thereby providing an actionable basis for load balancing.

Context Parallelism (CP) is a parallelism dimension commonly used when training long-context models: the context along the sequence dimension is partitioned across multiple devices, with each GPU responsible for only one segment of KV, but attention computation needs to span all shards, so information must be exchanged between devices. Traditional CP methods (such as Ring Attention) typically need to transmit the full KV or its chunks at each layer, with communication volume growing linearly with context length.

Split-KV's name derives precisely from this partitioning approach—KV is split and stored across devices, while attention results are aggregated through query-side tensors, avoiding full KV communication. Here exact is relative to approximate methods: Split-KV does not sacrifice numerical correctness, and its computation results are strictly consistent with the unsharded case, which is crucial for gradient correctness during training.

LPT Scheduling: Load Balancing with Real Attention Costs

With the global query view, Split-KV adopts a cost-aware Longest Processing Time first (LPT) strategy to allocate context KV blocks. Unlike the traditional practice of roughly partitioning by token count, LPT balances the load across ranks based on the actual attention cost each KV block incurs.

The effect is quite direct: the load imbalance rate is below 1%, while planning speed is three orders of magnitude faster than the best load-balancing baseline. For anchor-sampling training that requires frequent re-planning, the planning overhead itself cannot be ignored, and this speedup has real engineering value.

LPT (Longest Processing Time first) is a classic multi-machine scheduling heuristic: tasks are sorted by processing time from longest to shortest and assigned in turn to the currently least-loaded machine, yielding a makespan close to optimal. It has a theoretical approximation ratio guarantee of 4/3.

Split-KV's contribution lies in replacing the "processing time" in LPT with the real attention computation cost of each KV block, rather than a rough token count. Because anchor sampling causes each sample to participate in a different range of context, costs vary greatly between blocks, and this cost-aware ordering can significantly improve balancing quality. Weak scaling efficiency measures the ratio of actual throughput to ideal linear growth when keeping per-device workload constant and increasing the number of devices proportionally; it is a core metric for evaluating the scalability of a parallel scheme.

Scalability Performance and Open Code

At a CP scale of 64, Split-KV maintains 85–91% weak scaling efficiency, achieving up to 3.67x speedup over the fastest existing CP baseline. Considering the real-world demand for continuously growing context lengths in diffusion drafter training, this scaling efficiency means that adding devices still brings near-linear returns.

The code is open-sourced in the verl repository on GitHub (https://github.com/zxgx/verl). For teams training diffusion drafters and struggling with long-context parallel efficiency, this is an implementation worth evaluating directly.

From a broader perspective, Split-KV's value lies not just in a specific parallel strategy, but in demonstrating an approach: first identify the structural asymmetry of the workload, then design communication and scheduling around that asymmetry. As contexts continue to lengthen and query budgets continue to tighten, this "query-centric" parallel perspective may become the default option for diffusion drafter training infrastructure.