fix(models): make Gemma 4 suppress-token masking graph safe

Use a device buffer and index_fill_ for suppress-token masking during CUDA graph capture.

Fixes #48503

Co-authored-by: Devin AI <noreply@devin.ai>

Signed-off-by: Sidharth <dumko.raj@gmail.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Author: dumko2001

vllm/model_executor/models/gemma4_mm.py

@@ -1131,7 +1131,20 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.num_redundant_experts = self.language_model.num_redundant_experts
 
         gen_cfg = vllm_config.model_config.try_get_generation_config()
-        self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None
+        suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None
+        self.register_buffer(
+            "_suppress_token_ids",
+            (
+                torch.tensor(
+                    suppress_token_ids,
+                    dtype=torch.long,
+                    device=vllm_config.device_config.device,
+                )
+                if suppress_token_ids
+                else None
+            ),
+            persistent=False,
+        )
 
     # ------------------------------------------------------------------ #
     # Input parsing
@@ -1612,8 +1625,8 @@ def compute_logits(
         hidden_states: torch.Tensor,
     ) -> torch.Tensor | None:
         logits = self.language_model.compute_logits(hidden_states)
-        if logits is not None and self._suppress_token_ids:
-            logits[:, self._suppress_token_ids] = -float("inf")
+        if logits is not None and self._suppress_token_ids is not None:
+            logits.index_fill_(1, self._suppress_token_ids, -float("inf"))
         return logits
 
     # ------------------------------------------------------------------ #

vllm/model_executor/models/gemma4_mtp.py

@@ -529,7 +529,20 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         draft_cfg = vllm_config.speculative_config.draft_model_config
         gen_cfg = draft_cfg.try_get_generation_config()
-        self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None
+        suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None
+        self.register_buffer(
+            "_suppress_token_ids",
+            (
+                torch.tensor(
+                    suppress_token_ids,
+                    dtype=torch.long,
+                    device=vllm_config.device_config.device,
+                )
+                if suppress_token_ids
+                else None
+            ),
+            persistent=False,
+        )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)
@@ -581,8 +594,8 @@ def compute_logits(
             )
         else:
             logits = self.logits_processor(self.lm_head, hidden_states)
-        if logits is not None and self._suppress_token_ids:
-            logits[:, self._suppress_token_ids] = -float("inf")
+        if logits is not None and self._suppress_token_ids is not None:
+            logits.index_fill_(1, self._suppress_token_ids, -float("inf"))
         return logits
 
     def get_top_tokens(

vllm/model_executor/models/gemma4_unified.py

@@ -346,9 +346,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.num_redundant_experts = self.language_model.num_redundant_experts
         self.set_eplb_state = self.language_model.set_eplb_state
 
-        gen_cfg = vllm_config.model_config.try_get_generation_config()
-        self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None
-
     # ------------------------------------------------------------------ #
     # Multimodal processing (encoder-free overrides)
     # ------------------------------------------------------------------ #