Split-KV: Efficient Context Parallelism for Training Diffusion Drafters

Geng Zhang; Fangzhao Zhang; Cong Xie; Tiantian Fan; Yang You

cs.AINov 4, 2026

Split-KV: Efficient Context Parallelism for Training Diffusion Drafters

Authors: Geng Zhang, Fangzhao Zhang, Cong Xie, Tiantian Fan, Yang You

Organizations: National University of Singapore · Stanford University · ByteDance Seed

Abstract

Diffusion drafters accelerate speculative decoding by denoising multiple draft tokens in parallel. Their training uses anchor sampling to construct compact draft queries and performs cross-attention with a block-sparse mask over a much longer context. We empirically show that acceptance length remains stable as context scales with a fixed query budget, while increasing the query budget yields diminishing quality gains, jointly establishing a practical regime with queries far shorter than the context. No existing context parallelism (CP) method jointly exploits this query--context asymmetry and meets the nonlinear, sample-dependent load-balancing requirement created by anchor sampling. We present Split-KV, an exact CP method that keeps key/value shards local and limits per-layer attention communication to query-sized tensors, making this communication independent of context length. Its global query view further enables a cost-aware longest-processing-time-first (LPT) strategy that balances context key/value chunks by their realized attention cost. At a CP size of 64, Split-KV retains 85--91% weak-scaling efficiency and achieves up to $3.67\times$ speedup over the fastest available CP baseline. LPT achieves $<$ 1% imbalance and plans three orders of magnitude faster than the best load-balancing baseline. Code is available at https://github.com/zxgx/verl.

Figures & tables

	communication	peak memory	load balancing
no CP ( $P{=}1$ )	—	$\big(2(S{+}Q)+2Q\big)H$	N/A
Ulysses	$\frac{P-1}{P}\cdot\frac{4H(\hbox{\pagecolor{baselineRed}$ \displaystyle S $}+2Q)}{P}$	$\big(4(S{+}Q)+2Q\big)\frac{H}{P}$	Head sharding
Ring Attention	$\frac{2(3P-2)}{P}\cdot(\hbox{\pagecolor{baselineRed}$ \displaystyle S $}{+}Q)H$	$\big(4(S+Q)+2Q\big)\frac{H}{P}$	Zigzag
RingX	$\frac{P-1}{P}\cdot 2H(2\hbox{\pagecolor{baselineRed}$ \displaystyle S $}{+}3Q)$	$\big(4(S+Q)+2Q\big)\frac{H}{P}$	None
LV-XAttn	$\frac{2(3P-2)}{P}\cdot QH$	$\big(2(S+Q)+8Q\big)\frac{H}{P}$	None
Split-KV	$\frac{P-1}{P}\cdot 4QH$	$\big(2(S+Q)+2PQ\big)\frac{H}{P}$	Cost-aware LPT

Table 1: Per-rank training cost for a diffusion drafter layer. Multi-head attention is used for brevity.

H

is hidden size,

P

is CP size. Communication includes both forward and backward passes. Peak memory counts input and output tensors and communication buffers. Red highlights mark baseline limitations jointly addressed by Split-KV.

Panel of Figure 1 — Figure 1: The DFlash block-sparse mask for one micro-batch with two draft blocks ( $B{=}3$ ). One cell is one query–key pair. Each block’s anchor row $a_{i}$ and its $B{-}1$ masked rows ([M]) attend the context window $[s_{i},a_{i})$ of their own sample and their own block.

Figure 2: Mean acceptance length over five downstream tasks for two controlled SFT studies. (a) Fix

Q

–Scale

S

varies the maximum context length at a fixed query budget. (b) Fix

S

–Scale

Q

varies the query budget at a fixed maximum context length.

Figure 3: Number of queries attending key

x

in a 16-token sample. Causal self-attention has a linear profile whose mirrored positions have constant total load. Diffusion drafter anchor sampling (

B=2

, anchors

\{6,10,11,12\}

) produces a piecewise profile.

Panel of Figure 4 — Figure 4: Split-KV exploits $Q\ll S$ by all-gathering the global query while keeping the long key/value shards local. The global query view decouples key/value-chunk work from query ownership and enables the cost-aware balancing design.

Panel of Figure 5 — Figure 5: Load balancing of the context keys in Split-KV for the same two packed samples and sampled anchors. A cell shows its position and realized context cost $c_{\mathrm{anchor}}(x)$ ; chunk cells show the chunk label above its summed context cost. Zigzag independently cuts each sample into $2P$ chunks and pairs mirrored positions, yielding $(C_{0},\ldots,C_{3})=[24,24,30,34]$ . LPT cuts the global packed context into $PJ$ chunks and assigns them by realized chunk cost, yielding $(C_{0},\ldots,C_{3})=[28,28,28,28]$ . Adding the identical $D=2B^{2}=8$ gives $(W_{0},\ldots,W_{3})=[32,32,38,42]$ and $(W_{0},\ldots,W_{3})=[36,36,36,36]$ , respectively.

Target	Layers	$d_{\mathrm{model}}$	$d_{\mathrm{ffn}}$	$n_{q}$	$n_{kv}$	$d_{h}$
Qwen3.5-35B-A3B	8	2048	6144	32	4	128
Qwen3.5-27B	5	5120	17408	32	8	128

Table 2: Architectures of the evaluated DFlash drafters, where

n_{q}

n_{kv}

, and

d_{h}

denote the numbers of query heads, key/value heads, and the head dimension, respectively.

Figure 6: Drafter training step time under weak scaling. Each bar is labeled with absolute time in seconds and weak-scaling efficiency

E_{\mathrm{weak}}(P)=T_{8}/T_{P}

, separated by a slash. Lower time and higher efficiency are better; 100% efficiency denotes unchanged drafter training step time as the context and

P

grow proportionally. N/A denotes an unsupported configuration.

Figure 7: Drafter training step time under strong scaling. Each bar is labeled with absolute time in seconds and strong-scaling efficiency

E_{\mathrm{strong}}(P)=(T_{8}/T_{P})/(P/8)

, separated by a slash. Lower time and higher efficiency are better; the gap from 100% quantifies the scale-out loss. N/A denotes an unsupported configuration.

Figure 8: Drafter training step time as the query-to-context ratio decreases. Lower is better; labels report absolute time in seconds.

Figure 9: Average work imbalance under varying context lengths and CP sizes. Each label reports

\mathrm{avg}\;(\mathrm{std})

, with both values measured in percentage points. The lower panel magnifies the cost-aware strategies.

Figure 10: Planning latency across context-length and CP-size scaling for Qwen3.5-35B-A3B. The shared logarithmic axis labels every decade explicitly so that all three planners remain visible. No balance has zero planning cost and is omitted.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Fix $Q$ –Scale $S$
Target model	$S$	$Q$	GSM8K	MATH-500	HumanEval	MBPP	MT-Bench	Avg.
Qwen3.5-35B-A3B	$4\mathrm{k}$	$8\mathrm{k}$	5.98	6.99	7.00	5.96	4.48	6.08
	$16\mathrm{k}$	$8\mathrm{k}$	6.09	7.05	7.03	5.98	4.52	6.13
	$64\mathrm{k}$	$8\mathrm{k}$	6.09	6.95	7.05	5.97	4.54	6.12
Qwen3.5-9B	$4\mathrm{k}$	$8\mathrm{k}$	5.71	6.59	6.30	5.50	4.40	5.70
	$16\mathrm{k}$	$8\mathrm{k}$	5.70	6.51	6.44	5.56	4.47	5.74

Appendix

Table 1: Acceptance length on five downstream tasks. Fix

Q

–Scale

S

holds 512 anchors per sequence and varies the maximum context length; Fix

S

–Scale

Q

holds the maximum context length at

64\mathrm{k}

and varies the anchor budget. With block size

B=16

, the anchor budgets

\{128,256,512,1024\}

correspond to query lengths

\{2\mathrm{k},4\mathrm{k},8\mathrm{k},16\mathrm{k}\}

. Avg. is the unweighted mean of the five displayed task values.

Figure 1: Speedup in drafter training step time for each load-balancing strategy over No balance with per-sample context and query lengths of

256\mathrm{k}

and

8\mathrm{k}

, respectively. The dashed line denotes performance equal to No balance.

Context Parallelism

See all 11 papers

Nov 4, 2026

D3D: Decoupled Diffusion Drafting for Asynchronous Speculative Decoding

Target-conditioned diffusion drafters predict candidate tokens in parallel, yet each new block must wait for updated target-model hidden states. This forces drafting and verification to run sequentially, limiting speedup even though each block is generated in one pass. We introduce D3D, a standalone diffusion drafter that removes this inter-round dependency. D3D combines blockwise denoising with auxiliary autoregressive supervision to achieve long accepted prefixes without target hidden states. This independence lets D3D draft successor blocks during verification in one packed pass, enabling asynchronous speculative decoding. We progressively prune layers according to their effect on acceptance length to reduce the cost of packed drafting. Removing five of 36 layers retains over 99% of mean acceptance length across three benchmarks with Qwen3-8B as the target. Across seven benchmarks, D3D averages 5.00

\times

and 5.71

\times

decoding speedup with Qwen3-8B and Qwen3-32B targets, compared with 4.03

\times

and 3.67

\times

for DFlash. On GSM8K with Qwen3-32B, speedup reaches 7.48

\times

. Code is publicly available at https://github.com/fudan-generative-vision/D3D.

Baoyou Chen, Ruiqiao Mei, Guancong Lin +4

Shanghai Innovation Institute · Shanghai Academy of AI for Science · Fudan University +2

Speculative Decoding Speculative Decoding for Diffusion Language Models Diffusion Language Model Inference Blockwise Diffusion Language Models

Ask PDF

May 29, 2026

BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding

Speculative decoding speeds up autoregressive decoding by using a drafter to propose multiple tokens that a verifier validates in parallel. In resource-constrained deployments, the drafter uses a sparse KV cache to limit peak GPU memory and end-to-end latency under a fixed KV budget, while the verifier keeps a full KV cache. Mid-to-long context inference (4K--16K context length) is common in real applications. However, naive sparse/full speculative decoding suffers from the sparse/full mismatch as context length grows, causing the acceptance rate to drop quickly. We propose BudgetDraft, a multi-view sparse training method for sparse drafting in mid-to-long inference. The drafter is exposed to multiple sampled KV budgets during training and learns to align each sparse view with one shared full-cache teacher target. BudgetDraft combines an acceptance-aware loss on a full-cache branch with a multi-view loss on a sparse-cache branch, producing a single budget-robust drafter that recovers acceptance across sparsity levels without extra inference-time components. Experimental results on PG-19, LongBench, and LWM show that BudgetDraft achieves up to 6.55x, 4.46x, and 2.10x end-to-end speedup vs AR at 4K, 8K, and 16K context lengths, while keeping the inference pipeline memory-friendly.

Liang He, Jingbo Wen, Qishi Zhan +4

Shanghai Institute of Optics and Fine Mechanics · The University of Sydney · Marquette University +4

KV Caching Long-Context Language Model Inference KV-Cache Management Memory-Efficient Inference Speculative Decoding

Ask PDF

Oct 4, 2026

SharpDraft: Accelerating Long-Context Speculative Decoding with Cardinality-Aware Query Scaling

Long-form reasoning makes inference expensive, and speculative decoding mitigates this cost by verifying multiple draft tokens in parallel. Its speedup, however, can fade as context grows and draft acceptance declines. We focus on attention-mass dilution: as softmax normalizes over more visible Keys, the mass concentrated on the highest-scoring Keys can decrease. We introduce SharpDraft, a training-free method that counteracts this effect through cardinality-aware Query scaling, without the computational overhead of online adaptation. Under explicit assumptions, we derive an exact top-

k

mass correction and deploy a closed-form fixed-slope approximation. Across AIME-26, GPQA-Diamond, and LongGenBench Diary, SharpDraft achieves

2.59

3.19\times

geometric-mean end-to-end speedups over target-only autoregressive decoding when applied to DFlash, PARD, and EAGLE 3.1. With DFlash, it improves decoding speed and outperforms full-parameter and LoRA-based online adaptation in end-to-end speedup, while matching the unmodified drafter's reported peak allocated GPU memory.

Jeonghoon Park, Seongwoon Jo, Jongwon Lee +1

UNIST

Long-Context Language Model Inference Speculative Decoding LLM Inference Acceleration Parallel Speculative Decoding

Ask PDF

Abstract

Context Parallelism

Explore similar work

D3D: Decoupled Diffusion Drafting for Asynchronous Speculative Decoding

BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding

SharpDraft: Accelerating Long-Context Speculative Decoding with Cardinality-Aware Query Scaling