[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>
Luciano Martins committed
Aug 26, 2026 at 12:39 UTC
2ab187430b6be67abdecf116e1cd13e14968ddf5
1 file changed
+14
-3
vllm/model_executor/models/gemma4_mtp.py
+14
-3
@@ -524,6 +524,10 @@ class Gemma4MTP(nn.Module):
524
draft_cfg = speculative_config.draft_model_config
525
gen_cfg = draft_cfg.try_get_generation_config()
526
self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None
527
+ # Materialized on-device in load_weights: compute_logits runs under CUDA
528
+ # graph capture in the V2 speculator, where indexing with a Python list
529
+ # would issue an unpinned H2D copy (illegal during capture).
530
+ self._suppress_idx: torch.Tensor | None = None
531
532
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
533
return self.model.embed_input_ids(input_ids)
@@ -576,8 +580,8 @@ class Gemma4MTP(nn.Module):
580
)
581
else:
582
logits = self.logits_processor(self.lm_head, hidden_states)
579
- if logits is not None and self._suppress_token_ids:
580
- logits[:, self._suppress_token_ids] = -float("inf")
583
+ if logits is not None and self._suppress_idx is not None:
584
+ logits.index_fill_(1, self._suppress_idx, -float("inf"))
585
return logits
586
587
def get_top_tokens(
@@ -594,4 +598,11 @@ class Gemma4MTP(nn.Module):
598
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
599
self._stable_full_lm_head_weight = None
600
loader = AutoWeightsLoader(self)
597
- return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
601
+ loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
602
+ if self._suppress_token_ids:
603
+ self._suppress_idx = torch.tensor(
604
+ self._suppress_token_ids,
605
+ dtype=torch.long,
606
+ device=next(self.parameters()).device,
607
+ )
608
+ return loaded