---
title: Context-Sharded Block Parallelism (CSBP)
url: https://www.emergentmind.com/topics/context-sharded-block-parallelism-csbp
type: topic
---

# Context-Sharded Block Parallelism (CSBP)

Context-Sharded Block Parallelism (CSBP) is a distributed training method for block diffusion language models (BDLMs) that combines block parallelism over corrupted target blocks with context parallelism over the shared clean sequence. It assigns each target-block computation, including its corrupted activations, loss, and corrupted gradients, to one rank while sharding reusable clean keys and values across ranks. This organization preserves BDLM training semantics, avoids communication of block-specific corrupted K/V tensors, reduces clean-prefix replication relative to pure block parallelism, and enables long-context training at context lengths reaching $1$ million tokens [2609.19242].

## 1. BDLM objective and dependency structure

BDLMs generate sequences block by block while denoising tokens within each block in parallel. For a sequence of length $L$ divided into $B$ contiguous blocks of $M=L/B$ tokens,

$$
x=\left(x^{(1)},\ldots,x^{(B)}\right),
\qquad
x^{(<b)}=\left(x^{(1)},\ldots,x^{(b-1)}\right).
$$

Their probabilistic factorization is

$$
p_\theta(x)=\prod_{b=1}^{B}p_\theta\!\left(x^{(b)}\mid x^{(<b)}\right).
$$

During training, each target block $b$ is sampled at a continuous diffusion time $t_b$ and corrupted, producing $x_{t_b}^{(b)}$. Preceding blocks remain clean. The BDLM objective decomposes as

$$
\mathcal{L}_{\mathrm{BDLM}}
=
\sum_{b=1}^{B}
\mathbb{E}
\left[
-w(t_b)
\log p_\theta
\left(
x^{(b)}
\mid
x_{t_b}^{(b)},x^{(<b)}
\right)
\right]
=
\sum_{b=1}^{B}\mathcal{L}_b.
$$

This additive decomposition is the basis of block parallelism. The corrupted computation associated with block $b$ contributes only to $\mathcal{L}_b$. For a corrupted query in block $b$, the valid attention inputs are

$$
O_b=
\operatorname{Attn}
\left(
Q_b,
[K_b;K_{<b}^{\mathrm{clean}}],
[V_b;V_{<b}^{\mathrm{clean}}]
\right).
$$

Thus, corrupted queries attend bidirectionally within their own target block and attend to clean preceding blocks. They do not attend to corrupted K/V from other target blocks.

This asymmetric dependency structure distinguishes BDLM training from ordinary autoregressive or uniformly position-sharded attention. Clean K/V are reused by many later target blocks, whereas corrupted K/V from block $b$ are needed only by queries in block $b$. Conventional context parallelism does not exploit this distinction because it shards the combined clean-plus-corrupted representation by token position.

Long-context BDLM training is constrained by quadratic attention computation, distributed K/V communication and gradient communication, and activation memory. The reference implementation commonly represents a length-$2L$ sequence consisting of one clean copy and one block-specific corrupted copy, with a block-causal mask defining valid dependencies.

## 2. Relationship to conventional and pure block parallelism

Conventional context parallelism (CP) partitions the attention input by token position across $P$ ranks. A rank owns a query shard, while distributed attention provides access to K/V values held by other ranks. For BDLMs, this causes both clean and corrupted K/V to participate in distributed communication. In the backward pass, both clean and corrupted $dK/dV$ may also be communicated.

This organization is inefficient for corrupted K/V because corrupted block $b$ is consumed only by the queries of block $b$. A rank may nevertheless transmit those tensors to other ranks solely because position-based sharding does not encode target-block ownership. Causal attention additionally produces load imbalance because later positions attend to longer prefixes; zigzag or early–late pairing can mitigate, but does not remove, the communication asymmetry.

Pure block parallelism (BP) assigns complete target blocks to ranks. If $\mathcal{B}_r$ denotes the target blocks assigned to rank $r$, then

$$
\mathcal{L}_r=\sum_{b\in\mathcal{B}_r}\mathcal{L}_b,
\qquad
\sum_{r=1}^{P}\nabla_\theta\mathcal{L}_r
=
\nabla_\theta\mathcal{L}_{\mathrm{BDLM}}.
$$

BP keeps corrupted $Q_b,K_b,V_b$, corrupted activations, $\mathcal{L}_b$, and corrupted gradients local to the rank owning $b$. Its limitation is memory: later target blocks require long clean prefixes, and pure BP can replicate or overlap almost the entire clean sequence across ranks. The supplied evaluation reports that pure BP is slower and uses more memory for NemotronDiffusion 14B at $64$K, runs out of memory at $128$K, and runs out of memory for DiffusionGemma 26B-A4B at both $64$K and $128$K under the matched native objective.

CSBP combines the two decompositions:

- BP assigns complete corrupted target-block computations to ranks.
- CP shards the shared clean sequence across the same ranks.
- Corrupted K/V and corrupted gradients remain local.
- Clean K/V are communicated through the context-parallel mechanism.
- Clean K/V gradients are accumulated on their owning ranks.

CSBP therefore avoids both the corrupted-K/V communication of conventional CP and the replicated clean-prefix memory of pure BP.

## 3. Forward execution

Each rank simultaneously owns a clean-sequence shard and a set of complete corrupted target blocks. For rank $r$, the local data include

- clean projections $Q_r^{\mathrm{clean}},K_r^{\mathrm{clean}},V_r^{\mathrm{clean}}$;
- corrupted projections $(Q_b,K_b,V_b)$ for $b\in\mathcal{B}_r$.

The forward workflow has several stages.

### Clean-context communication

Clean K/V shards are communicated among ranks. The implementation described in the paper uses an all-gather:

$$
(K_r^{\mathrm{clean}},V_r^{\mathrm{clean}})
\longrightarrow
(K^{\mathrm{clean}},V^{\mathrm{clean}}).
$$

Production execution starts this communication asynchronously and overlaps it with local corrupted-block attention. The broader CSBP principle does not require a particular transport primitive; clean K/V may be supplied through an all-gather or a ring/streaming distributed-attention mechanism.

### Local corrupted-block attention

The owner of each target block computes its corrupted-block attention locally:

$$
O_b^{\mathrm{corr}}
=
\operatorname{Attn}(Q_b,K_b,V_b),
\qquad b\in\mathcal{B}_r.
$$

No corrupted K/V communication is required. This preserves the block-locality implied by the BDLM objective.

### Clean-prefix attention

After clean K/V are available, the rank computes attention from the relevant corrupted queries to the clean prefix under the block-causal mask. A query in block $b$ can access only clean positions preceding block $b$.

The local corrupted-block key set and the clean-prefix key set are disjoint. Their outputs and row-wise log-sum-exp values can therefore be computed separately and merged exactly using an online-softmax operation. Given partial outputs $O_1,O_2$ with log-sum-exp values $z_1,z_2$, let

$$
m=\max(z_1,z_2),
\qquad
w_j=e^{z_j-m}.
$$

The merged result is

$$
O=\frac{w_1O_1+w_2O_2}{w_1+w_2},
\qquad
z=m+\log(w_1+w_2).
$$

The implementation stores the numerator, maximum, and normalizer in FP32 until final normalization. This avoids materializing the complete attention matrix while preserving the combined contribution of local corrupted and distributed clean keys.

### Layout adaptation

The transparent reshuffler hides CSBP’s block-aware layout from the remainder of the model. Tensor layouts are changed at the attention boundary and restored afterward, allowing FSDP, tensor parallelism, expert parallelism, sequence parallelism, positional embeddings, and feed-forward layers to retain their existing layouts.

## 4. Backward propagation and exactness

The backward pass follows the forward dependency graph.

### Local corrupted gradients

Because corrupted block $b$ is visible only to queries from block $b$,

$$
\frac{\partial O_b}{\partial Z_{b'}}=0
\qquad
\text{for }b'\neq b,
\qquad
Z\in\{Q,K,V\}.
$$

Consequently, the owner of block $b$ can compute $dQ_b,dK_b,dV_b$ locally. No other rank needs corrupted gradients from that block.

### Distributed clean gradients

A clean K/V row may be used by clean queries and by corrupted queries from many later target blocks. For the clean K/V rows owned by rank $r$,

$$
\frac{\partial\mathcal{L}}
{\partial Z_r^{\mathrm{clean}}}
=
\sum_q
\left(
\frac{\partial O_q}
{\partial Z_r^{\mathrm{clean}}}
\right)^{\!\top}
\frac{\partial\mathcal{L}}{\partial O_q},
\qquad
Z\in\{K,V\}.
$$

The owning rank accumulates contributions from all relevant queries. The implementation uses reduce-scatter for clean K/V gradients. The resulting communication pattern is asymmetric:

- corrupted $dQ,dK,dV$ terminate at the target-block owner;
- clean $dK,dV$ contributions are reduced to clean-shard owners;
- corrupted-gradient communication is eliminated.

The fused implementation packs local corrupted K/V together with gathered clean K/V, performs masked attention backward, separates local corrupted gradients from clean gradients, and reduce-scatters only the clean component.

### Objective preservation

CSBP preserves the sampled reference loss and parameter gradient up to floating-point reduction order. If $C_r$ is the clean-position shard owned by rank $r$, then

$$
\bigcup_{r=1}^{P}C_r=\{1,\ldots,L\},
\qquad
C_r\cap C_{r'}=\varnothing
\quad(r\neq r').
$$

For every query, the visible keys are partitioned among the local corrupted block and the relevant clean shards. Online-softmax merging reconstructs the same attention result as the reference length-$2L$ computation, subject to numerical reduction order:

$$
H_d^{\mathrm{CSBP}}=H_d^{\mathrm{ref}}.
$$

Consequently,

$$
\mathcal{L}_b^{\mathrm{CSBP}}
=
\mathcal{L}_b^{\mathrm{ref}},
\qquad
\mathcal{L}^{\mathrm{CSBP}}
=
\sum_{b=1}^{B}\mathcal{L}_b^{\mathrm{ref}},
$$

and

$$
\sum_{r=1}^{P}
\nabla_\theta\mathcal{L}_r^{\mathrm{CSBP}}
=
\nabla_\theta\mathcal{L}^{\mathrm{ref}}.
$$

The exactness claim assumes known target blocks and corruptions, no corrupted cross-block attention, correct clean-shard aggregation, correct clean-gradient accumulation, and identical loss normalization and model randomness.

At $256$K, the reported absolute and relative loss differences between CSBP and the best non-BP baseline were:

| Model | Absolute loss difference | Relative loss difference |
|---|---:|---:|
| NemotronDiffusion 3B | $2.83\times10^{-3}$ | $2.11\times10^{-4}$ |
| NemotronDiffusion 8B | $2.59\times10^{-3}$ | $1.97\times10^{-4}$ |
| NemotronDiffusion 14B | $9.54\times10^{-7}$ | $7.18\times10^{-8}$ |
| DiffusionGemma 26B-A4B | $6.51\times10^{-3}$ | $2.11\times10^{-4}$ |
| Qwen3.8-27B | $7.21\times10^{-6}$ | $5.11\times10^{-7}$ |
| Qwen3.5-27B | $9.50\times10^{-5}$ | $6.69\times10^{-6}$ |

These differences were attributed to floating-point reduction order rather than a change in the objective.

## 5. Load balancing and resource complexity

CSBP must balance both target-block computation and clean-prefix computation. For equal-sized blocks, the target-block work associated with block $b$ is

$$
W_b^{\mathrm{target}}=bM^2,
$$

because its $M$ queries attend to $bM$ visible keys. The clean-prefix work is

$$
W_b^{\mathrm{clean}}
=
(b-1)M^2+\frac{M(M+1)}{2}.
$$

Assigning early and late blocks to the same rank produces constant paired workloads:

$$
W_b^{\mathrm{target}}
+
W_{B+1-b}^{\mathrm{target}}
=
(B+1)M^2,
$$

and

$$
W_b^{\mathrm{clean}}
+
W_{B+1-b}^{\mathrm{clean}}
=
BM^2+M.
$$

This early–late or zigzag pairing balances causal attention more effectively than contiguous assignment. When $B$ is divisible by $2P$, the per-rank work is

$$
W_r^{\mathrm{target}}
=
\frac{B(B+1)M^2}{2P},
$$

and

$$
W_r^{\mathrm{clean}}
=
\frac{L(L+1)}{2P}.
$$

If divisibility is imperfect, complete pairs are distributed as evenly as possible and an unpaired middle block is assigned to the least-loaded rank.

At $256$K, replacing contiguous assignment with dual-end pairing improved throughput by $1.34\times$ for NemotronDiffusion 3B and $1.28\times$ for DiffusionGemma 26B-A4B. In the DiffusionGemma case, peak HBM was reduced by $6.5$ GiB.

CSBP does not asymptotically reduce attention FLOPs. Conventional CP and CSBP both perform

$$
\Theta(L^2)
$$

total attention work and approximately

$$
\Theta\left(\frac{L^2}{P}\right)
$$

work per rank under balanced partitioning. Its performance benefits arise from communication and memory organization:

- corrupted K/V are not communicated across ranks;
- corrupted gradients remain local;
- clean prefixes are sharded rather than replicated;
- clean K/V reuse is preserved;
- block ownership enables more favorable configurations of data, context, tensor, and expert parallelism.

Under balanced assignment, each CSBP rank stores approximately $L/P$ clean tokens and $L/P$ corrupted tokens, or approximately $2L/P$ token activations. Pure BP can require nearly $L$ clean tokens on a rank assigned late blocks, making the replicated clean-prefix component potentially up to $P$ times larger.

## 6. Communication, memory, and implementation trade-offs

Conventional CP communicates clean and corrupted K/V in the forward pass and clean and corrupted K/V gradients in the backward pass. CSBP communicates only shared clean context:

| Quantity | Conventional CP | CSBP |
|---|---|---|
| Forward distributed K/V | Clean and corrupted | Clean only |
| Backward distributed K/V gradients | Clean and corrupted | Clean only |
| Corrupted K/V ownership | Position-sharded | Target-block-local |
| Clean-prefix storage | Position-sharded | Position-sharded |
| Corrupted gradients | Potentially communicated | Local to block owner |

The paper reports detailed $256$K profiles:

| Measure | NemotronDiffusion 14B CP8 $\rightarrow$ CSBP | DiffusionGemma CP8 $\rightarrow$ CSBP |
|---|---:|---:|
| Step time | $66.50\to56.37$ s | $49.66\to34.13$ s |
| Forward GPU span | $17.57\to13.00$ s | $10.99\to5.83$ s |
| Backward GPU span | $49.46\to45.02$ s | $36.50\to26.62$ s |
| Attention traffic per GPU | $210.00\to105.00$ GiB | $433.13\to28.30$ GiB |
| Exposed attention communication | $3.034\to2.414$ s | $10.939\to0.723$ s |

For NemotronDiffusion 14B, logical attention traffic was halved, whereas exposed communication declined by $20.4\%$ because overlap already hid part of the communication. For DiffusionGemma, 25 of 30 attention layers use a $1{,}024$-token sliding window; CSBP exchanged only the needed clean rows in those layers, producing a $93.5\%$ logical and $93.4\%$ exposed attention-communication reduction.

The method remains dependent on clean-context communication. All-gathering full clean K/V can remain expensive at very large context lengths unless it is streamed or overlapped. The implementation uses asynchronous communication, multiple buffers, NCCL, fused attention kernels, and layout transformation at the attention boundary. The full-model experiments used BF16, PyTorch, FlashAttention-4, FlexAttention, NCCL, activation checkpointing, DeepSpeed ZeRO Stage 2, and fused or token-chunked cross entropy.

The strongest gains occur when context lengths are large, CP degree is high, clean prefixes are heavily reused, corrupted cross-block attention is absent, and communication can overlap with local block computation. Conventional CP may remain preferable for short contexts, small CP degrees, workloads without separable target blocks, or regimes in which data parallelism avoids sequence-parallel communication. In the reported DFlash2 $256$K experiment, for example, DP8 was faster than the CSBP configuration because full replicas fit and avoided sequence-parallel overhead.

## 7. Evaluation, applications, and limitations

Full-model training used two nodes with eight NVIDIA H200 GPUs per node. Evaluated models included NemotronDiffusion 3B, 8B, and 14B; DiffusionGemma 26B-A4B; and Qwen3.5-27B and Qwen3.8-27B for autoregressive-to-BDLM conversion.

At $256$K, reported throughput and peak HBM were:

| Workload and model | Baseline / CSBP throughput | Speedup | Baseline / CSBP peak HBM |
|---|---:|---:|---:|
| NemotronDiffusion 3B SFT | $13{,}434/16{,}027$ tok/s | $1.19\times$ | $57.1/57.0$ GiB |
| NemotronDiffusion 8B SFT | $9{,}668/11{,}379$ tok/s | $1.18\times$ | $97.5/97.5$ GiB |
| NemotronDiffusion 14B SFT | $7{,}884/9{,}301$ tok/s | $1.18\times$ | $91.1/90.6$ GiB |
| DiffusionGemma 26B-A4B SFT | $10{,}557/15{,}359$ tok/s | $1.45\times$ | $118.1/113.6$ GiB |
| Qwen3.8-27B conversion | $2{,}623/3{,}490$ tok/s | $1.33\times$ | $120.7/111.1$ GiB |
| Qwen3.5-27B conversion | $2{,}641/3{,}351$ tok/s | $1.27\times$ | $120.7/111.1$ GiB |

At $512$K, the strongest full-model speedups were $1.19\times$ for NemotronDiffusion 14B and $1.61\times$ for DiffusionGemma 26B-A4B. For DiffusionGemma at $512$K, CSBP achieved $9{,}815$ tok/s versus $6{,}079$ tok/s for the best baseline, while reducing peak HBM from $127.4$ to $106.5$ GiB.

The DFlash2 speculative-decoder experiments on eight NVIDIA H100 GPUs showed that the benefit can increase sharply when context length makes ordinary CP or replication infeasible. For Qwen3.8-27B, reported speedups were:

| Context | Baseline | CSBP | Speedup |
|---|---:|---:|---:|
| $256$K | $167{,}228$ tok/s | $140{,}350$ tok/s | $0.84\times$ |
| $512$K | $57{,}027$ tok/s | $141{,}222$ tok/s | $2.48\times$ |
| $1$M | $18{,}389$ tok/s | $139{,}541$ tok/s | $7.59\times$ |

The $256$K result favors ordinary data parallelism because full replicas fit. At $512$K, CSBP reduced exposed attention communication from $105.4$ to $7.5$ ms per GPU. At $1$M, CSBP enabled four data-parallel replicas instead of the baseline’s single CP8 replica.

The method’s limitations follow directly from its assumptions. It requires a block-decomposable training objective and known target blocks and corruptions. If corrupted K/V from one block can be consumed by queries in another block, corrupted gradients are no longer local and the principal communication advantage is lost. Load balancing is essential because equal numbers of target blocks do not imply equal causal attention work. The results are also hardware- and implementation-dependent, relying on BF16, fused attention kernels, NCCL, tuned communication, and specific H100/H200 configurations.

CSBP is therefore best characterized as a training-time context-sharded block-parallel execution strategy rather than a general replacement for context parallelism. Its defining principle is the alignment of distributed ownership with the BDLM dependency graph: target-block computation is sharded by block, while reusable clean context is sharded by position. This alignment preserves the BDLM objective while reducing unnecessary communication, limiting activation replication, and enabling long-context training at scales where conventional CP or pure BP become inefficient or infeasible.

Source: https://www.emergentmind.com/topics/context-sharded-block-parallelism-csbp