Repository navigation
[Bugfix][DiffusionGemma] Read the causal mask as int32 and compile the sample step - #59992
Merged
Merged
Conversation
Signed-off-by: ojeda-e <eojeda@league.com>
ojeda-e
requested review from
AndreasKaratzas,
WoosukKwon,
bigPYJ1151,
mgoin,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
October 5, 2026 00:28
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ojeda-e
commented
Oct 5, 2026
| token_num = sum(query_lens) | ||
| dynamic_causal_tensor = ( | ||
| torch.tensor(dynamic_causal, dtype=torch.bool) | ||
| torch.tensor(dynamic_causal, dtype=torch.int32) |
Contributor
Author
There was a problem hiding this comment.
Leaving this note as a heads up.
Since tying the test to the model's buffer felt out of scope here, I only updated this dtype.
However, note that the kernel test still builds its own mask instead of reading DiffusionGemma's _causal_buf. As a consequence, a future dtype change in the model would not fail this test. It would fail at the kernel instead, because data_ptr<int32_t>() checks the dtype.
Signed-off-by: ojeda-e <eojeda@league.com>
ojeda-e
changed the title
ojeda-e
changed the title
Chessing234
reviewed
Oct 6, 2026
Chessing234
left a comment
Chessing234
left a comment
There was a problem hiding this comment.
everything seems fine, ready to merge.
bigPYJ1151
added
the
verified
Member
|
/ci run |
|
✅ Triggered Buildkite CI #93067 for commit |
Contributor
CI selector (shadow): 11 test steps (25 jobs) instead of 68 (112 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (11)
Would skip (today's rules run them) (63)
Would add (today's rules do not run them) (6)
AMD mirrors: would skip (60)
AMD mirrors: would add (0)none 5 changed files · base |
6 tasks done
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
4 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.
Overview
Fixes #59991.
This PR has two fixes for DiffusionGemma. The first unblocks the CPU backend. The second applies to every platform.
DiffusionGemma fails on every CPU run. [Bugfix][Model] Fix DiffusionGemma silently freezing attention mask under CUDA graph replay #51994 made the per-request causal mask
int32and updated FlashAttention and Triton. CPU attention still reads it asbool. This PR makes the CPU kernel readint32.[Perf] DiffusionGemma: one-pass sampler statistics kernel #58226 left
_compiled_sample_stepuncompiled on every platform. This PR moves itstorch.compiledecorators back from the_denoise_temperaturehelper.No open issue or PR covers either fix. For the causal mask I searched for the error message,
dynamic_causalon CPU and DiffusionGemma on CPU. For the sample step I searched open PRs for_compiled_sample_step,_denoise_temperatureand58226. The only hits, #45954, #47278 and #56210, do not touch the decorators.Claims
CPU causal mask (7b1cee8):
RuntimeError: expected scalar type Bool but found Inton the CPU backend. It is raised at engine startup. No DiffusionGemma request can run on CPU onmain.int32contract from [Bugfix][Model] Fix DiffusionGemma silently freezing attention mask under CUDA graph replay #51994. The mask is passed through without a cast, like FlashAttention, Triton and the XPU support in [XPU][Attention] Enable per-sequence causal masking with FA2 #57747.test_cpu_attn.pybuilds the mask asint32, like the model does. The existing per-request causal test reproduces the bug onmain.Sample step compile (c4a07b1):
_compiled_sample_stepis compiled again, as it was before [Perf] DiffusionGemma: one-pass sampler statistics kernel #58226._denoise_temperaturestays compiled and drops the raised recompile limit, which it does not need.test_sample_step_is_compiledfails iftorch.compilemoves off the step again.Validation
Hardware: 22 vCPU Intel Xeon Platinum 8481C (AMX), Ubuntu 22.04, PyTorch 2.13.0+cpu. vLLM built from source for CPU at 7867d6c.
CPU causal mask
maintest_varlen_with_paged_kv_dynamic_causalexpected scalar type Bool but found Inttests/kernels/attention/test_cpu_attn.pygoogle/diffusiongemma-26B-A4B-itin BF16Any DiffusionGemma request on a CPU build reproduces it. I used
LLM.generatewith a structured read (32-token canvas, read-only, one step).Sample step compile
test_sample_step_is_compiledfails before the change and passes after. Its module skips without CUDA and I have no GPU, so I called the test function directly on CPU.fallback_randomon. Sampled tokens and state match. No graph breaks.Sample step alone, vocab 262144, hidden 2816, bf16 embeddings, median of 10 calls. Eager is the same build under
torch.compiler.set_stance("force_eager").End to end with
LLM.generate:google/diffusiongemma-26B-A4B-itin BF16,enforce_eager=True(it does not turn off the step's owntorch.compile), canvas 64, 8 chat prompts per call,max_tokens=128. 3 timed calls after a warmup, arms run in the order shown.The difference is within run-to-run noise. On CPU the sample step is a small share of each call. The compiled arm words answers differently because the renoise tokens now come from inductor's RNG.
I couldn't test in GPU (running short today).
Details
CPU causal mask
#51994 fixed a FULL CUDA graph bug. FlashAttention cast the
boolmask toint32on everybuild(), and the captured graph kept reading the tensor from capture time. The fix allocates the mask asint32in DiffusionGemma, and FlashAttention now rejects other dtypes. CPU attention passes the same tensor to its C++ ops, which read it withdata_ptr<bool>(). The CPU tests built their ownboolmask. CI stayed green.The kernel reads
int32instead ofCPUAttentionMetadataBuilder.build()casting the mask. A cast there would bring back the per-step allocation #51994 removed. It is harmless on CPU today, since CPU has no graph capture, but it would break the contract every other backend follows.No measurable tradeoffs. The kernel reads the flag once per request and once per work group. The mask grows from 1 to 4 bytes per request.
Limitations:
boolmasks. The only caller passes the model'sint32mask.test_varlen_with_paged_kv_dynamic_causalruns only on CPUs with AMX. Other ISAs share the changed lines, but no test covers their per-request causal path, before or after this PR.Sample step compile
#58226 added
_denoise_temperaturedirectly below the@torch._dynamo.config.patch(recompile_limit=64)and@torch.compiledecorators of_compiled_sample_step, and Python applied them to the helper. Since then the step has run eager on every platform. The step body has not changed since #58226. The fix restores the earlier layout.Tradeoffs, each the same as before #58226:
torch.randintuses inductor's RNG. At a given global seed, outputs differ from the eager step. Diffusion models reject per-request seeds.Limitations:
test_sample_step_is_compiledchecks_torchdynamo_orig_callable, whichtest_batched_count_greater_than.pyalso reads. Compiled and eager differ only in random draws, so no output assertion can tell them apart.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.