Back in June, while experimenting with Diffusion Gemma 4 on vLLM, I've noticed it won't start on my 4 Nvidia RTX A4000 graphic cards. I sent a patch to vLLM that makes DiffusionGemma servable on more than one GPU. It went in as PR #46177 against issue #45719, merged on June 26, and shipped in the v0.25.0 release on July 11. This post covers what was actually broken and why the fix ended up shaped the way it is.
The symptom
DiffusionGemma is a block diffusion model built on the Gemma 4 backbone: 25.2B total parameters, 3.8B active, with a mixture-of-experts FFN that runs 8 of 128 experts per token. At that size most people need several GPUs, which means tensor parallelism. Every launch with --tensor-parallel-size above 1 died during engine warmup:
RuntimeError: a and b must have same reduction dim ... [65536, 2816]A single GPU was completely fine. That combination, working on one card and crashing on four, is the whole clue.
Where the shapes stop agreeing
The sampler uses self-conditioning: at each denoising step it feeds the model a summary of its own previous prediction. For a discrete model the natural summary of a predicted distribution is the expected token embedding, which is just the distribution multiplied by the embedding matrix.
\mathbf{e}^{sc}_i = \mathbb{E}_{x\sim\mathbf{p}_i}\big[\mathrm{emb}(x)\big] = \sum_{v=1}^{V} p_{i,v}\,\mathbf{W}_{v,:} = \mathbf{p}_i\mathbf{W}
Here \mathbf{p}_i is a distribution over the full vocabulary, which for this model is V = 262{,}144 tokens, and \mathbf{W}\in\mathbb{R}^{V\times d} is the token embedding matrix with hidden size d = 2{,}816. The original code did exactly that:
soft_embeds = torch.matmul(
probs.to(embed_weight.dtype), embed_weight) * normalizerThe problem is that vLLM shards embeddings along the vocabulary axis. With T tensor-parallel ranks, rank r owns only the contiguous slice [s_r, e_r) of the vocabulary:
\mathbf{W}^{(r)} = \mathbf{W}[s_r{:}e_r,\,:] \;\in\; \mathbb{R}^{(V/T)\times d}
So at T = 4 each rank holds 65,536 rows instead of 262,144. Meanwhile probs is still full width, because the sampler runs on gathered logits. The contraction dimension is V on the left and V/T on the right, so PyTorch refuses.
At T = 1 the shard is the whole matrix, the two dimensions agree, and nothing looks wrong. The bug is only reachable in the configuration a 26B model actually needs.
The fix
A matrix product whose contracted axis is partitioned is a sum of partial products over the blocks of that partition. Since the vocabulary partition is disjoint and exhaustive:
\mathbf{P}\mathbf{W} \;=\; \sum_{v=1}^{V}\mathbf{P}_{:,v}\,\mathbf{W}_{v,:} \;=\; \sum_{r=0}^{T-1}\mathbf{P}_{:,\,s_r:e_r}\,\mathbf{W}^{(r)}
Every rank already holds its own \mathbf{W}^{(r)}, and it can slice the matching columns out of probs by itself. So each rank computes a local partial \mathbf{S}^{(r)} of shape [L_c \times d], and one all-reduce with summation puts the exact result on every rank:
\mathbf{E}^{sc} \;=\; \gamma \sum_{r=0}^{T-1}\mathbf{S}^{(r)}, \qquad \mathbf{S}^{(r)} = \mathbf{P}_{:,\,s_r:e_r}\,\mathbf{W}^{(r)}
where \gamma is the learned normalizer. No weight ever moves.
local_probs = probs[..., sc_vocab_start:sc_vocab_end] \
.to(embed_weight.dtype)
soft_embeds = torch.matmul(
local_probs,
embed_weight[: sc_vocab_end - sc_vocab_start])
if tp_size > 1:
soft_embeds = torch.ops.vllm.all_reduce(
soft_embeds, group_name=tp_group_name)
soft_embeds = soft_embeds * normalizerThis is the same communication pattern as a row-parallel linear layer in Megatron-LM: operands sharded on the contracted axis, local products giving partial sums, one all-reduce reconstructing the output. The only departure from bit-identity with the single GPU result is the order of additions, which is a benign floating point reassociation.
Two details matter here because the sampler is compiled.
The collective is torch.ops.vllm.all_reduce and not the eager tensor_model_parallel_all_reduce. The sampler step is wrapped in @torch.compile and captured into CUDA graphs, and only the functionalized custom op is traceable and graph safe. That constraint comes from the engine rather than from the model, and it decides what the fix is allowed to look like.
The tp_size > 1 guard is also load bearing. At T = 1 the slice is the whole vocabulary, the partial is already the full product, and the collective is skipped, so the single GPU path stays bitwise identical to what it was. The new code generalizes the old code instead of adding a second path that has to be kept in sync with it.
One smaller point worth mentioning: the shard bounds come from the embedding's own metadata, org_vocab_start_index and org_vocab_end_index, rather than from arithmetic on V and T. vLLM pads shard layouts to hardware friendly multiples, and those padding rows carry no probability mass, so reading the true bounds keeps them out of the slice.
Why not just all-gather the weights
There was a simpler fix available. All-gather \mathbf{W} to every rank once at startup, then leave the original matmul alone. It is equally correct. I did not do it because of what it costs.
At V = 262{,}144, d = 2{,}816 and BF16, the embedding is
V d b = 262{,}144 \times 2{,}816 \times 2 = 1{,}476{,}395{,}008 \text{ bytes} \approx 1.38 \text{ GiB}
All-gather puts that on every GPU and keeps it there. The hardware I was testing on is four RTX A4000s at 16 GiB each, already sitting near 15.4 GiB per card in service. An extra 1.38 GiB per rank is the difference between serving and an OOM. The reason anyone reaches for tensor parallelism in the first place is that the model does not fit on one card, so memory pressure is the normal case here, not an edge case.
All-gather also makes every rank redundantly compute the full matmul: 378 GFLOP per rank per step at T = 4, against 94.5 GFLOP for the sharded version.
What the all-reduce design costs instead is a collective on the hot path, one per denoising step. But the tensor being reduced is [L_c \times d], with no factor of V in it anywhere. At canvas length L_c = 256 that is 1.38 MiB, roughly 2.06 MiB per rank per request per step through a ring all-reduce.
One caveat on that comparison. If you count only wire volume, the one time all-gather wins over a long enough run. With K denoising steps and N active sequences, the crossover sits at
2 K N L_c d \;\lesssim\; V d \quad\Longleftrightarrow\quad K N L_c \;\lesssim\; V/2
which at these sizes means K N \lesssim 512. The T = 4 runs came in near 11 denoising steps per canvas, well under that for small batches. So the accurate description of the merged design is memory optimal inside vLLM's existing vocab-sharded layout, and communication aware rather than unconditionally communication optimal.
Where it sits in the sampler
The starred node is the only place in the loop where a full-vocabulary distribution meets the vocab-sharded embedding, which makes it the only place TP correctness for self-conditioning has to be enforced.
Running it
- Model: diffusiongemma-26B-A4B-it, INT8 dynamic
- Hardware: 4 × RTX A4000, 16 GiB each, Ampere SM86, about 15.4 GiB used per GPU in service
- Config:
--tensor-parallel-size 4,TRITON_ATTN, canvas length 256, max model length 131,072 - Before: warmup crash on the reduction dim
- After: warmup completes, CUDA graphs capture, four TP workers active, coherent generation
For a throughput sample I asked the service to write a complete offline pixel art editor as a single HTML file. 278 prompt tokens in, 8,814 completion tokens out, 50.56 seconds wall clock, 174.3 tokens/s measured client side. Server side logs showed generation windows up to 230.4 tokens/s. The output ran to 27,005 characters over 793 lines and was a complete, well-formed document.
That rate moves around quite a bit, and it should. Block diffusion commits a variable number of tokens after a variable number of denoising iterations. In that run the logs ranged from 13.3 to 24.2 steps per canvas, and the windows with fewer steps and more tokens committed per step read faster. That variation comes from the sampler's acceptance dynamics rather than from anything in this change.
None of this is a benchmark. There was no T > 1 baseline to be faster than, because T > 1 did not run at all. The claim here is availability, not speedup.
Keeping it fixed
The merged change adds a T = 2 GSM8K evaluation config on the FP8 checkpoint, 1,319 questions, 5-shot. The sharded sampler path now gets exercised on real hardware as part of vLLM's regression surface, rather than only by a synthetic shape test.
What it does not fix
Pipeline parallelism. The original issue reported a PP failure too, but PP needs diffusion canvas state propagated across pipeline stages, which is a different problem from making one projection layout aware. I scoped the change to TP and said so in the PR.
The general point
Autoregressive sampling propagates a sampled token id, so the full vocabulary distribution is a temporary that dies inside the sampler. Diffusion sampling propagates the distribution itself. Any operator that consumes a full-vocabulary distribution against a vocab-sharded parameter therefore has to be rewritten as a sharded reduction, and self-conditioning is simply the first one you hit in DiffusionGemma. For this class of model I doubt it will be the last.
The full derivation, the proof that the sharded version is exact, and the cost model behind that comparison are written up in a pre-print with my doctoral advisor, Sicong Shao.
