Block Parallelism for Efficient Distributed Long-Context Diffusion Language Model Training

We introduce Context-Sharded Block Parallelism (CSBP), a new distributed parallelism strategy that unlocks significant long-context training efficiency.

For diffusion LLMs and speculative decoding drafters, and open-sourced in Turbo-dLLM, our new highly optimized distributed training library.

Tarun Suresh^* Pranshu Chaturvedi^* Hangoo Kang^* Parth Shroff
Ishan S. Khare Hermann Kumbong Azalia Mirhoseini Stanford University^*Equal contribution

Read the Paper Code

Faster long-context training, across the board

Training time per 1B tokens versus context length for three workloads. Autoregressive to block diffusion training of Qwen3.8-27B: CSBP is 1.26×, 1.30× and 1.33× faster at 64K, 128K and 256K. Block diffusion fine-tuning of DiffusionGemma 26B-A4B: 1.25×, 1.44×, 1.45× and 1.61× faster from 64K to 512K. DFlash2 speculative decoding drafter training for Qwen3.8-27B: 2.48× faster at 512K and 7.59× faster at 1M.

CSBP cuts the wall-clock cost of training on long sequences for every workload we tried, and the gap widens as contexts grow. On 16 H200 GPUs it converts Qwen3.8-27B into a block diffusion model 1.33× faster at 256K and fine-tunes DiffusionGemma 26B-A4B 1.61× faster at 512K. On a single 8 H100 node it trains a DFlash2 speculative decoding drafter 7.59× faster at 1M tokens of context, while computing exactly the same objective and gradients.

Background

Block diffusion language models

Generate text block by block, denoising every token inside a block in parallel.

Training input

Every token appears twice, so a 256K example is 512K positions of attention input and must be split across GPUs. Conventional context parallelism splits both copies and sends both copies' K/V between GPUs.

Block diffusion speculative decoding

A small diffusion drafter writes a whole block in one pass from the target model's features; the target then verifies it.

Training input

Drafters for long-context targets train on sequences of up to a million tokens. Conventional context parallelism repeats every block-local drafter computation on each context GPU.

In both, training is one long shared context plus many small block-local computations. Conventional context parallelism splits everything by token position and ignores that structure. CSBP is built around it.

Method

How does Context-Sharded Block Parallelism (CSBP) work?

Split the shared context across GPUs, but never split a block.

Corrupted keys and values are only ever read by their own block, yet conventional context parallelism sends them, and their gradients, between GPUs at every layer.
Block parallelism (BP) runs each block's whole forward pass, loss, and backward pass on one GPU, so corrupted K/V never leave it.
Context-sharded block parallelism (CSBP) also shards the shared clean context across those GPUs. Only clean K/V and their gradients cross GPUs, and the loss and gradients stay exact.

CSBP in motion

Load balancing

Later blocks attend to longer prefixes, so they cost more. Handing out blocks in order leaves the last GPU with the most work while the others sit idle waiting for it. CSBP pairs early and late blocks so every GPU does the same work, and the same pairing balances the sharded clean context, so the whole step finishes sooner.

Library

Turbo-dLLM: Optimized for Long-Context Spec Decoding and Diffusion Training End to End

CSBP decides where work runs. Turbo-dLLM also rewrites how each step runs, replacing the generic building blocks that stock trainers use for both workloads.

Speculative Decoding Optimizations

Layer	Turbo-dLLM	SGLang SpecForge	Effect
Target context	Compacts the union of the required sliding windows	Projects the full target sequence in every draft layer	Much less K/V compute and memory
Attention	Specialized FlashAttention-4 interval attention with cached sparse plans	Generic FlexAttention / SDPA / eager masks	Less mask construction and workspace
Projections	Packed QKV, packed context K/V, packed MLP	Separate projections and concatenations	Fewer kernel launches and intermediates
Vocabulary loss	Streams rows through the frozen LM head without materializing full logits	Materializes logits in checkpointed chunks and recomputes them in backward	Lower memory traffic, less recomputation
Parallelism	Context and block parallelism on top of data parallelism	FSDP sharding modes; trainer tensor parallelism unsupported offline	Fits contexts a single GPU cannot hold

Block Diffusion Optimizations

Layer	Turbo-dLLM	NVIDIA AutoModel	Effect
Attention	Represents the whole structured mask in one compiled FlexAttention call per layer	Evaluates each shard block by block, with a separate attention launch for the growing clean prefix and each noisy fragment	Far fewer launches and no per-block mask overhead
Vocabulary loss	Token-chunked projection and cross-entropy that never materializes full-sequence logits	Materializes logits for both streams, and runs out of memory at the matched topology	Fits where the baseline cannot
Parallelism	Dual-end zigzag layout pairs early and late regions across ranks	Assigns contiguous blocks, leaving later ranks with much longer prefixes	Balanced attention work instead of stragglers
Collectives	Fused K/V collectives, exact fused backward, and owner reduce-scatter	Autograd all-gather followed by a blockwise attention decomposition	Less exposed communication per layer

Results

CSBP Accelerates Long-Context Block Diffusion Training

Model	Best baseline		CSBP		Speedup	Peak memory (GiB) baseline → CSBP
Model	Topology	Tok/s	Topology	Tok/s	Speedup	Peak memory (GiB) baseline → CSBP
Block diffusion fine-tuning
NemotronDiffusion 3B	DP4/CP4	13,434	DP4/CP4/BP4	16,027	1.19×	57.1 → 57.0
NemotronDiffusion 8B	DP4/CP4	9,668	DP4/CP4/BP4	11,379	1.18×	97.5 → 97.5
NemotronDiffusion 14B	DP2/CP8	7,884	DP2/CP8/BP8	9,301	1.18×	91.1 → 90.6
DiffusionGemma 26B-A4B	DP1/EP2/CP8	10,557	DP1/EP2/CP8/BP8	15,359	1.45×	118.1 → 113.6
Autoregressive → block diffusion conversion
Qwen3.8-27B	DP1/TP2/CP8	2,623	DP1/TP2/CP8/BP8	3,490	1.33×	120.7 → 111.1
Qwen3.5-27B	DP1/TP2/CP8	2,641	DP1/TP2/CP8/BP8	3,351	1.27×	120.7 → 111.1

16× H200, 256K context. The baseline is the fastest feasible combination of data, tensor, expert, context and sequence parallelism we found for each model, tuned with the same kernels and optimizations as CSBP. Topology degrees: DP data, TP tensor (omitted when 1), EP expert, CP context, BP block parallelism.

CSBP is faster for every model and both objectives at 256K context: 1.18–1.45× for block diffusion fine-tuning and 1.27–1.33× for converting autoregressive models, while matching or reducing peak memory.

Scaling

Speedup Scales with Context Length

Training time per 1B tokens versus context length. Autoregressive to block diffusion training of Qwen3.8-27B from 64K to 256K: CSBP is 1.26×, 1.30× and 1.33× faster. Block diffusion fine-tuning of DiffusionGemma 26B-A4B from 64K to 512K: 1.25×, 1.44×, 1.45× and 1.61× faster.

Longer contexts make CSBP's advantage grow. Converting Qwen3.8-27B into a block diffusion model speeds up from 1.26× at 64K to 1.33× at 256K, and fine-tuning DiffusionGemma 26B-A4B from 1.25× at 64K to 1.61× at 512K: the longer the sequence, the more it pays to keep corrupted K/V off the wire.

Speculative decoding

Up to 7.59× faster drafter training at 1M context

DFlash2 drafter training time per 1B tokens on 8× H100. Qwen3.8-27B drafter: CSBP is 2.48× faster at 512K and 7.59× faster at 1M. Muse-Glimmer-30B drafter: 1.56× faster at 256K, 3.35× at 512K and 4.81× at 1M.

DFlash2 drafters see the largest gains. Conventional CP copies the block-local draft decoder, vocabulary loss and selector onto every context rank, while CSBP gives each draft block to exactly one GPU. Freeing that memory lets more replicas train in parallel: 7.59× faster for the Qwen3.8-27B drafter and 4.81× faster for the Muse-Glimmer-30B drafter at 1M tokens, on a single 8 H100 node.

Benchmark

Turbo-dLLM Accelerates Long-Context Training Over Stock Libraries

Speculative Decoding Drafter Training

Context	Turbo-dLLM		SGLang SpecForge		Speedup
Context	Topology	Tok/s	Topology	Tok/s	Speedup
64K	DP8	167,960	DP8	84,002	2.00×
128K	DP8	167,900	DP8	73,199	2.29×
256K	DP8	166,732	DP8	OOM	only Turbo-dLLM fits
512K	DP4/CP2/BP2	141,222	—	OOM	only Turbo-dLLM fits
1M	DP4/CP2/BP2	139,541	—	OOM	only Turbo-dLLM fits

DFlash2 drafter training for Qwen3.8-27B on one 8×H100 node, best topology per library and context.

Turbo-dLLM trains the same drafter 2.00× faster at 64K and 2.29× faster at 128K, and keeps going to 1M context where SGLang SpecForge has no configuration that fits.

Block Diffusion Fine-Tuning

Context	Turbo-dLLM			NVIDIA AutoModel			Advantage
Context	Topology	Tok/s	Peak HBM	Topology	Tok/s	Peak HBM	Advantage
64K	DP8/CP2/BP2	22,152	89.8 GiB	DP4/TP2/CP2/SP	2,934	114.1 GiB	7.55×
128K	DP4/CP4/BP4	15,136	90.1 GiB	DP2/TP2/CP4/SP	2,155	114.1 GiB	7.02×
256K	DP2/CP8/BP8	9,301	90.6 GiB	DP1/TP2/CP8/SP	633	130.1 GiB	14.69×

NemotronDiffusion 14B on 16× H200, aggregate input tokens/s for each library's fastest feasible setup. AutoModel uses two GPUs per replica for tensor parallelism and therefore runs half the global batch. The Turbo-dLLM rows are the paper's CSBP topologies.

Downstream

Same GPUs, same hours, better coding agents

Pass rate versus training time for DiffusionGemma 26B-A4B. SWE-bench Verified: CSBP reaches 33.8, 35.4, 36.0 and 36.4 percent at 3, 6, 9 and 12 hours, versus 32.2, 33.6, 34.6 and 35.4 for the best baseline. Terminal-Bench Lite: CSBP reaches 27, 29, 30 and 30 percent versus 26, 27, 28 and 29.

Higher throughput means more learning per GPU-hour. We fine-tuned DiffusionGemma 26B-A4B for 12 hours on eight H100s with each method, using identical data, optimizer and evaluation, and tested checkpoints every three hours. CSBP leads at every trained checkpoint on both SWE-bench Verified and Terminal-Bench Lite, by up to +1.8 and +2 points.

Details

Under the hood

Where does the time go? Attention traffic at 256K

Keeping corrupted K/V local halves attention traffic for NemotronDiffusion 14B, which uses full attention. DiffusionGemma 26B-A4B saves more because 25 of its 30 layers use a 1,024-token sliding window: there, CSBP sends only the clean rows each GPU's queries can see, while the CP8 baseline exchanges full shards in every layer.

Measure	NemotronDiffusion 14B CP8 → CSBP	DiffusionGemma 26B-A4B CP8 → CSBP
Step time (s)	66.50 → 56.37 1.18×	49.66 → 34.13 1.45×
Forward GPU span (s)	17.57 → 13.00 −26.0%	10.99 → 5.83 −47.0%
Backward GPU span (s)	49.46 → 45.02 −9.0%	36.50 → 26.62 −27.1%
Attention traffic (GiB/GPU)	210.00 → 105.00 −50.0%	433.13 → 28.30 −93.5%
Exposed attention comm. (s/GPU)	3.034 → 2.414 −20.4%	10.939 → 0.723 −93.4%

Why shard the context? Block parallelism on its own runs out of memory

Without context sharding, ranks that own later blocks must hold nearly the whole clean prefix. Pure BP is slower and uses more memory for NemotronDiffusion 14B at 64K, then runs out of memory at 128K. DiffusionGemma runs out of memory at both lengths. CSBP fits and is fastest in every case.

Model	Context	Method	Topology	Tok/s	Peak memory (GiB)
NemotronDiffusion 14B	64K	Best baseline	DP8/CP2	20,092	89.5
		Pure BP	DP8/BP2	15,223	110.7
		CSBP	DP8/CP2/BP2	22,152	89.8
NemotronDiffusion 14B	128K	Best baseline	DP4/CP4	13,114	90.0
		Pure BP	DP4/BP4	OOM	OOM
		CSBP	DP4/CP4/BP4	15,136	90.1
DiffusionGemma 26B-A4B	64K	Best baseline	DP4/EP2/CP2	22,154	109.4
		Pure BP	DP4/EP2/BP2	OOM	OOM
		CSBP	DP4/EP2/CP2/BP2	27,653	113.5
DiffusionGemma 26B-A4B	128K	Best baseline	DP2/EP2/CP4	14,584	112.0
		Pure BP	DP2/EP2/BP4	OOM	OOM
		CSBP	DP2/EP2/CP4/BP4	21,049	113.6

Load balancing: pairing early and late blocks

Later blocks attend to longer prefixes, so handing out blocks in order leaves the last rank doing the most work. CSBP gives block b and block B+1−b to the same rank and uses the same early–late pairing for zigzag clean-sequence sharding. Every pair then does identical work, which balances both target-block and clean causal attention across ranks exactly, with no extra memory.

Model (256K)	Topology	Contiguous tok/s	Dual-end tok/s	Speedup	Peak memory (GiB)
NemotronDiffusion 3B	DP4/CP4/BP4	11,788	15,768	1.34×	57.0 → 57.0
DiffusionGemma 26B-A4B	EP2/CP8/BP8	11,991	15,359	1.28×	120.1 → 113.6

Citation

BibTeX

@misc{suresh2026blockparallelismefficientdistributed,
      title={Block Parallelism For Efficient Distributed Long-Context Diffusion Language Model Training}, 
      author={Tarun Suresh and Pranshu Chaturvedi and Hangoo Kang and Parth Shroff and Ishan S. Khare and Hermann Kumbong and Azalia Mirhoseini},
      year={2026},
      eprint={2609.19242},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2609.19242}, 
}