Repository navigation
[Bugfix][Model] Fix DiffusionGemma silently freezing attention mask under CUDA graph replay - #51994
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
LucasWilkinson
left a comment
LucasWilkinson
left a comment
There was a problem hiding this comment.
Good catch! thanks; can we remove the .to(...) call too inorder to avoid this footgun?
|
Thanks @LucasWilkinson , I just pushed a new commit that removes the |
|
Given that this is a bug on the standard path in the official recipe, would you consider merging this @LucasWilkinson @DarkLight1337 @ywang96 @AndreasKaratzas @MatthewBonanni ? |
710b110 to
a3b03f5
Compare
|
Hi @mmastrac, @LucasWilkinson, @njhill, @NickLucche I saw that you merged #57250 to add support for Jev-like structured outputs. Quick reminder that this PR fixes an issue where DiffusionGemma can use an incorrect attention mask during FULL CUDA graph replay, switching between causal and bidirectional attention incorrectly. To my understanding, the example in #57250 uses The actual impact of a wrong attention mask is actually not super clear to me but the bug can impose a training-inference mismatch. I rebased the branch onto main and would recommend merging this simple fix. |
…uffer under CUDA graphs DiffusionGemmaModelState._causal_buf is bool. FlashAttention's metadata builder casts non-int32 causal tensors via `causal.to(torch.int32)`, which allocates a new tensor on every call instead of mutating in place. Under CUDAGraphMode.FULL, the address returned at capture time is what the replayed kernel reads for the graph's entire lifetime, so a bool _causal_buf means later updates never reach an already-captured graph. Verified on H100 (SM90, FlashAttention v4): constructing a real DiffusionGemmaModelState and driving its _causal_buf through an actual torch.cuda.CUDAGraph capture/replay cycle shows the replayed output frozen at the capture-time value, not the live one. test_causal_buf_survives_cudagraph_replay fails on this commit, demonstrating the bug. Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com>
…DA graph replay _causal_buf was bool, but FlashAttentionMetadataBuilder.build() casts non-int32 causal tensors via causal.to(torch.int32), an out-of-place op that allocates a new tensor every call. Under CUDAGraphMode.FULL, the captured graph binds to the address returned at capture time, so later in-place updates to a bool _causal_buf never reach an already-captured graph -- replay keeps using whatever causal/bidirectional flags existed at capture. Making the buffer int32 up front means build() takes the dtype-matched fast path and never re-allocates, so in-place updates (.copy_() / direct assignment) stay visible to the graph across every replay. Fixes test_causal_buf_survives_cudagraph_replay added in the prior commit. Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com>
Trim redundant docstrings/comments and inline the single-use helper; no behavior change. The FlashAttention-cast still runs once before graph capture, matching how build() behaves in production, so the test still fails against a bool _causal_buf and passes against int32. Verified on H100. Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com>
Per review: drop the silent causal.to(torch.int32) fallback. The cast is out-of-place by construction (the builder has no persistent buffer to cast into), so any non-int32 causal tensor silently detaches FULL CUDA graph replay from the caller's buffer. Raise instead so the next caller hits a loud error at the first build() rather than stale attention modes. DiffusionGemma is the only tensor-causal producer and already allocates int32; comment and regression test updated to state the contract instead of mirroring the removed cast. Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com>
a3b03f5 to
8c67d68
Compare
|
@fjosw LGTM. If you want to fix a related nit: @LucasWilkinson FWIW I think this should land as-is. |
LucasWilkinson
left a comment
LucasWilkinson
left a comment
There was a problem hiding this comment.
LGTM thanks for the fix! sorry this slipped under my radar for so long! thanks @mmastrac for flagging
There was a problem hiding this comment.
nit: I think this test is a bit overly specific, not sure how much value it brings. I think we should probably just remove it
LucasWilkinson
added
the
ready
|
✅ @fjosw, CI is now available for this PR.
|
|
Hi @fjosw, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
Update the Triton per-sequence causal buffer documentation to match the int32 buffer used by DiffusionGemma. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com>
|
Thanks for your comments @mmastrac and @LucasWilkinson, I followed your suggestions and removed the test and fixed the comment in |
|
/ci run |
|
✅ Triggered Buildkite CI #92028 for commit |
3 participants
Purpose
While looking into whether FULL CUDA graphs are worth the added complexity for DiffusionGemma, I found that
DiffusionGemmaModelState._causal_buf(the per-request flag that switches each request between causal/encoder and bidirectional/denoise attention) stops being updated onceCUDAGraphMode.FULLgraphs are captured. It's bool, but FlashAttentionMetadataBuilder.build() casts non-int32 causal tensors viacausal.to(torch.int32), allocating a new tensor every call. A FULL-mode graph binds to the address it saw at capture time, so once captured it keeps reading that orphaned tensor forever. Later updates to_causal_bufnever reach it. Any request served through an already-captured FULL graph gets whatever causal/bidirectional pattern existed at capture (warmup) time for its slot, for the graph's whole lifetime, silently.This isn't opt-in: vLLM's default optimization level (O2) resolves cudagraph_mode to FULL_AND_PIECEWISE, whose decode_mode() is FULL, so
vllm serve google/diffusiongemma-26B-A4B-itwith no flags already runs every denoise/decode step under FULL cudagraphs.I couldn't pin observed output corruption on this specifically: an E2E FULL-vs-PIECEWISE comparison was inconclusive, since DiffusionGemma's denoising has its own run-to-run variance of similar magnitude. The capture/replay test below reproduces the mechanism directly instead. So overall, this is a silent bug that likely has some unpredictable impact on output quality.
Fix: make
_causal_bufint32 from construction and update it in place (.copy_()/ direct assignment) instead of = from a bool source, so build() takes its dtype-matched fast path and never reallocates.Tagging @LucasWilkinson, who authored the original FULL-cudagraph path in #45163.
AI-assisted, every line reviewed and tested by the submitter.
Test Plan
pytest tests/models/test_diffusion_gemma_cudagraph.py -vTest Result
New regression test constructs a real DiffusionGemmaModelState and drives _causal_buf through an actual torch.cuda.CUDAGraph capture/replay cycle mirroring FlashAttention's cast. Fails on main: replay returns the stale
[0, 0, 0, 0]instead of the updated[1, 1, 1, 1]. Passes with this patch. Run on H100 (SM90).