[Weight Reload][Gemma4] Refresh MTP gathered LM head in stable storage

tests/model_executor/test_utils.py

@@ -161,6 +161,36 @@ def scale_for(group_size: int) -> int:
     assert layer.weight.scale_for(3) == 6
 
 
+def test_gemma4_mtp_gather_cache_keeps_pointer_and_new_values(monkeypatch):
+    from types import SimpleNamespace
+
+    from vllm.model_executor.models import gemma4_mtp
+    from vllm.model_executor.utils import register_derived_buffer
+
+    model = object.__new__(gemma4_mtp.Gemma4MTP)
+    torch.nn.Module.__init__(model)
+    model.lm_head = torch.nn.Linear(2, 2, bias=False)
+    model.masked_embedding = SimpleNamespace(vocab_size=3)
+    register_derived_buffer(model, "_stable_full_lm_head_weight")
+    monkeypatch.setattr(gemma4_mtp, "get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(
+        gemma4_mtp,
+        "tensor_model_parallel_all_gather",
+        lambda weight, dim: torch.cat((weight, weight), dim=dim),
+    )
+    first = model._get_full_lm_head_weight()
+    pointer = first.data_ptr()
+    for value in (1.0, 3.0, 1.0):
+        with torch.no_grad():
+            model.lm_head.weight.fill_(value)
+        model.post_weights_reload()
+        result = model._get_full_lm_head_weight()
+        assert result is first
+        assert result.data_ptr() == pointer
+        torch.testing.assert_close(result, torch.full((3, 2), value))
+    assert "_stable_full_lm_head_weight" not in model.state_dict()
+
+
 def test_replace_parameter_preserves_bound_method_attribute() -> None:
     """A callable already bound to another object must keep pointing at that
     object after replacement.

vllm/model_executor/models/gemma4_mtp.py

@@ -44,6 +44,7 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.utils import register_derived_buffer, set_derived_buffer
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
 
@@ -468,7 +469,7 @@
         text_config = _get_text_config(config)
         self.quant_config = get_draft_quant_config(vllm_config)
         self.config = config
-        self._stable_full_lm_head_weight: torch.Tensor | None = None
+        register_derived_buffer(self, "_stable_full_lm_head_weight")
 
         self.model = Gemma4MultiTokenPredictor(
             vllm_config=vllm_config,
@@ -520,7 +521,7 @@
         # 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
+        self.register_buffer("_suppress_idx", None, persistent=False)
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)
@@ -547,6 +548,9 @@
     def _get_full_lm_head_weight(self) -> torch.Tensor:
         if self._stable_full_lm_head_weight is not None:
             return self._stable_full_lm_head_weight
+        return self._refresh_full_lm_head_weight()
+
+    def _refresh_full_lm_head_weight(self) -> torch.Tensor:
         assert self.masked_embedding is not None
         lm_head_weight = self.lm_head.weight
         tp_size = get_tensor_model_parallel_world_size()
@@ -558,9 +562,15 @@
         lm_head_weight = lm_head_weight[: self.masked_embedding.vocab_size]
         if tp_size > 1:
             lm_head_weight = lm_head_weight.contiguous()
-            self._stable_full_lm_head_weight = lm_head_weight
+            lm_head_weight = set_derived_buffer(
+                self, "_stable_full_lm_head_weight", lm_head_weight
+            )
         return lm_head_weight
 
+    def post_weights_reload(self) -> None:
+        if self.masked_embedding is not None:
+            self._refresh_full_lm_head_weight()
+
     def compute_logits(
         self,
         hidden_states: torch.Tensor,
@@ -573,8 +583,8 @@
             )
         else:
             logits = self.logits_processor(self.lm_head, hidden_states)
         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(
@@ -589,10 +599,9 @@
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        self._stable_full_lm_head_weight = None
         loader = AutoWeightsLoader(self)
         loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-        if self._suppress_token_ids:
+        if self._suppress_token_ids and self._suppress_idx is None:
             self._suppress_idx = torch.tensor(
                 self._suppress_token_ids,
                 dtype=torch.long,