Split-KV speeds diffusion drafter training by keeping long-context data local

Diffusion drafters propose several tokens in parallel for a larger language model to verify, helping accelerate speculative decoding. But training those drafters presents its own systems problem: relatively few attention queries may need access to a much longer context distributed across accelerators.

In the preprint Split-KV: Efficient Context Parallelism for Training Diffusion Drafters, Geng Zhang, Fangzhao Zhang, Cong Xie, Tiantian Fan and Yang You propose a method tailored to this imbalance. Their reported gains concern drafter training step time, not a measured increase in language-model generation speed.

Why short queries change the problem

The paper uses DFlash as its concrete diffusion drafter architecture. During training, sampled positions called anchors open blocks containing tokens to be denoised. Each block attends to the preceding context within its own sample and to all tokens inside that block.

The anchor budget determines the number of query tokens, while the training samples determine the context length. Those quantities can therefore change independently. The authors first investigate whether a small query budget remains useful as the maximum training context grows.

In controlled studies with drafters for Qwen3.5-9B and Qwen3.5-35B-A3B, they hold the query budget at 8k and vary the maximum context length from 4k to 64k. Mean acceptance length across five tasks changes by less than 1% for either model. Increasing the query budget at a fixed maximum context length improves acceptance length, but with diminishing returns.

Acceptance length measures how many tokens are accepted per draft step, including the target model’s bonus token; it is not downstream answer accuracy. Each training variant was evaluated once, and the authors explicitly treat the sub-1% differences as equivalent rather than statistically significant rankings. Within the tested range, these results support using queries substantially shorter than the context.

Exchange the smaller side

Context parallelism divides sequence-related work across devices. Several existing approaches still communicate tensors whose size grows with the context. That is potentially wasteful when the query budget stays fixed.

Split-KV instead keeps each device’s attention keys and values local throughout the forward and backward passes. It gathers the projected queries so every device can evaluate them against its own key/value shard, then combines the partial results using normalization statistics and collective reductions.

The authors derive this combination to reconstruct the original attention exactly, rather than approximate it. Its per-layer attention communication depends on query size, not context length. That does not make the entire training step context-independent: context-side computation remains, and the balancing scheme performs a separate context redistribution once per training step.

Another baseline, LV-XAttn, also avoids context-length-dependent attention communication. Split-KV’s distinction includes its collective communication schedule and its integration with workload balancing, rather than context-independent traffic alone.

Balance actual attention work

Equal amounts of context data do not necessarily mean equal computation. Early context positions can be read by many sampled draft blocks, while later positions may be read by fewer. The resulting workload depends on the sampled anchors.

Split-KV gives every device the complete query view, allowing the scheduler to assign each context chunk a cost based on its actual attention usage. A longest-processing-time-first strategy, or LPT, sorts chunks by cost and assigns expensive chunks to the least-loaded eligible device, while keeping chunk counts equal.

Across the two load-balancing scaling studies, LPT achieves average work imbalance of 0.38% and 0.21%. Its planning latency ranges from 0.38 to 1.24 milliseconds for Qwen3.5-35B-A3B. The competing DCP planner reaches more than a second at the largest tested parallel size. The authors report that LPT improves drafter training step time by 1.15×–1.42× over leaving chunks unbalanced.

What the speedups cover

The systems benchmarks use DFlash drafters for Qwen3.5-35B-A3B and Qwen3.5-27B on up to 64 accelerators. With workload growing alongside device count, Split-KV retains 91% and 85% weak-scaling efficiency, respectively. With total workload fixed, its largest reported advantage over the fastest available baseline is 3.67× at 64 devices.

Scaling is not perfect. At that size, strong-scaling efficiency is 55% and 43% for the two drafters. The paper attributes the loss to shrinking local attention work alongside a fixed global query payload and increasing collective latency and synchronization.

These timing experiments use randomly generated fixed-length samples and precomputed target-model hidden states, excluding target-model feature extraction. Timings average three measured steps after five warmup steps. The results therefore demonstrate faster offline drafter training under the reported conditions, not an equivalent speedup for every stage of a complete training or inference pipeline.