[Bugfix] Make Gemma4 MTP suppress_tokens masking CUDA-graph-safe (#53884)

Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>

vllm/model_executor/models/gemma4_mtp.py

@@ -524,6 +524,10 @@ class Gemma4MTP(nn.Module):
         draft_cfg = 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
+        # Materialized on-device in load_weights: compute_logits runs under CUDA
+        # graph capture in the V2 speculator, where indexing with a Python list
+        # would issue an unpinned H2D copy (illegal during capture).
+        self._suppress_idx: torch.Tensor | None = None
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)
@@ -576,8 +580,8 @@ class Gemma4MTP(nn.Module):
             )
         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_idx is not None:
+            logits.index_fill_(1, self._suppress_idx, -float("inf"))
         return logits
 
     def get_top_tokens(
@@ -594,4 +598,11 @@ class Gemma4MTP(nn.Module):
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         self._stable_full_lm_head_weight = None
         loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        if self._suppress_token_ids:
+            self._suppress_idx = torch.tensor(
+                self._suppress_token_ids,
+                dtype=torch.long,
+                device=next(self.parameters()).device,
+            )
+        return loaded