[Bugfix][Model] Fix DiffusionGemma silently freezing attention mask under CUDA graph replay (#51994)
Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Fabian Joswig committed
Sep 30, 2026 at 07:34 UTC
1117140edbaf3d1e1f73b01b09cde3cf98ec24ff
3 files changed
+15
-8
vllm/model_executor/models/diffusion_gemma.py
+9
-6
@@ -909,9 +909,12 @@ class DiffusionGemmaModelState(ModelState):
909
self._req_id_to_index: dict[str, int] = {}
910
911
# Persistent buffer for per-request causal flags, updated in-place
912
- # so FULL CUDA graph replay sees the latest values.
912
+ # so FULL CUDA graph replay sees the latest values. Must be int32:
913
+ # FlashAttentionMetadataBuilder.build() rejects other dtypes, since
914
+ # an out-of-place cast there would detach the captured graph from
915
+ # this buffer and freeze replay at the capture-time snapshot.
916
self._causal_buf = torch.zeros(
914
- self.max_num_reqs, dtype=torch.bool, device=device
917
+ self.max_num_reqs, dtype=torch.int32, device=device
918
)
919
920
# Persistent inputs_embeds buffer — required so FULL CUDA graph
@@ -1116,11 +1119,11 @@ class DiffusionGemmaModelState(ModelState):
1119
# Invariant: the sampler flips is_encoder_phase to False only after a
1120
# request's FINAL prompt chunk, so a prompt spanning multiple chunks
1121
# (longer than the token budget) stays causal for every chunk.
1119
- self._causal_buf[:actual_num_reqs] = self.diffusion_states.is_encoder_phase[
1120
- slots
1121
- ]
1122
+ self._causal_buf[:actual_num_reqs].copy_(
1123
+ self.diffusion_states.is_encoder_phase[slots]
1124
+ )
1125
if actual_num_reqs < num_reqs:
1123
- self._causal_buf[actual_num_reqs:num_reqs] = False
1126
+ self._causal_buf[actual_num_reqs:num_reqs] = 0
1127
causal: bool | torch.Tensor = self._causal_buf[:num_reqs]
1128
1129
return build_attn_metadata(
vllm/v1/attention/backends/flash_attn.py
+5
-1
@@ -987,7 +987,11 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
987
scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)
988
989
if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32:
990
- causal = causal.to(torch.int32)
990
+ raise ValueError(
991
+ f"Per-request causal tensor must be int32, got {causal.dtype}. "
992
+ "Casting here would allocate a fresh tensor each build() and "
993
+ "break FULL CUDA graph replay; allocate the buffer as int32."
994
+ )
995
996
attn_metadata = FlashAttentionMetadata(
997
num_actual_tokens=num_actual_tokens,
vllm/v1/attention/ops/triton_unified_attention.py
+1
-1
@@ -217,7 +217,7 @@ def kernel_unified_attention(
217
SLIDING_WINDOW: tl.constexpr, # int
218
USE_CAUSAL: tl.constexpr, # bool
219
USE_PER_SEQ_CAUSAL: tl.constexpr, # bool
220
- per_seq_causal_ptr, # [num_seqs] bool, or None
220
+ per_seq_causal_ptr, # [num_seqs] int32, or None
221
USE_MM_PREFIX: tl.constexpr, # bool
222
MAX_MM_RANGES: tl.constexpr, # int
223
mm_prefix_range_ptr,