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

D3D: Decoupled Diffusion Drafting for Asynchronous Speculative Decoding
Nov 4, 2026Baoyou Chen, Ruiqiao Mei, Guancong Lin +4Speculative Decoding Speculative Decoding for Diffusion Language Models
Shanghai Innovation Institute · Shanghai Academy of AI for Science · Fudan University +2
BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding
May 29, 2026Liang He, Jingbo Wen, Qishi Zhan +4KV Caching Long-Context Language Model Inference
Shanghai Institute of Optics and Fine Mechanics · The University of Sydney · Marquette University +4
SharpDraft: Accelerating Long-Context Speculative Decoding with Cardinality-Aware Query Scaling
Oct 4, 2026Jeonghoon Park, Seongwoon Jo, Jongwon Lee +1Long-Context Language Model Inference Speculative Decoding
UNIST

Nov 4, 2026J/K move · Enter open · S save

D3D: Decoupled Diffusion Drafting for Asynchronous Speculative Decoding

Baoyou Chen, Ruiqiao Mei, Guancong Lin, Tao Song+3

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

Open paper AskPDF arXiv

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.

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