From 2ab187430b6be67abdecf116e1cd13e14968ddf5 Mon Sep 17 00:00:00 2001 From: Luciano Martins <22145370+lucianommartins@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:39:19 -0300 Subject: [PATCH] [Bugfix] Make Gemma4 MTP suppress_tokens masking CUDA-graph-safe (#53884) Signed-off-by: Luciano Martins Co-authored-by: Luciano Martins --- vllm/model_executor/models/gemma4_mtp.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index 8e135d84c0..5e25a80d6b 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/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