[Bugfix] Fix DeepSeek V4 mHC broadcast buffer for dummy load

tests/models/test_deepseek_v4_mega_moe.py

@@ -9,10 +9,13 @@
 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
     bind_routed_experts_capturer,
 )
+from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM
 from vllm.models.deepseek_v4.nvidia.model import (
+    DeepseekV4ForCausalLM,
     DeepseekV4MegaMoEExperts,
     make_deepseek_v4_expert_params_mapping,
 )
+from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP
 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
 from vllm.platforms import current_platform
 
@@ -243,6 +246,37 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
     )
 
 
+def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast():
+    """The loader invokes the model-level PWAL hook for every load format,
+    so it must finalize megamoe + mhc broadcast weights to cover dummy
+    load, which skips load_weights()."""
+    calls = []
+    stub = SimpleNamespace(
+        model=SimpleNamespace(
+            finalize_mega_moe_weights=lambda: calls.append("mega_moe"),
+            finalize_mhc_broadcast_weights=lambda: calls.append("mhc"),
+        )
+    )
+
+    DeepseekV4ForCausalLM.process_weights_after_loading(stub)
+
+    assert calls == ["mega_moe", "mhc"]
+
+
+def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe():
+    """MTP/DSpark drafters load as their own top-level models, so each needs
+    its own PWAL hook now that the megamoe forward no longer finalizes
+    weights lazily on first use."""
+    calls = []
+    mtp = SimpleNamespace(finalize_mega_moe_weights=lambda: calls.append("mtp"))
+    DeepSeekV4MTP.process_weights_after_loading(mtp)
+
+    dspark = SimpleNamespace(_finalize_moe=lambda: calls.append("dspark"))
+    DSparkDeepseekV4ForCausalLM.process_weights_after_loading(dspark)
+
+    assert calls == ["mtp", "dspark"]
+
+
 @pytest.mark.skipif(
     not torch.cuda.is_available(),
     reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",

vllm/models/deepseek_v4/nvidia/dspark.py

@@ -504,16 +504,19 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 weight_loader(param, loaded_weight)
                 loaded_params.add(name)
 
-        self._finalize_moe()
         if self.model.confidence_head is not None and not loaded_confidence_head:
             self.model.confidence_head = None
+        self.process_weights_after_loading()
         logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
         return loaded_params
 
     def _finalize_moe(self) -> None:
         for layer in self.model.layers:
             layer.ffn.finalize_mega_moe_weights()
 
+    def process_weights_after_loading(self) -> None:
+        self._finalize_moe()
+
     def _remap_dspark_name(self, name: str) -> str | None:
         """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
 

vllm/models/deepseek_v4/nvidia/model.py

@@ -504,10 +504,6 @@ def forward(
             is_padding=is_padding,
         )
 
-        # This method must have been already called during the weight loading phase.
-        # We call it again here to cover the dummy weight loading case.
-        self.finalize_weights()
-
         assert self._transformed_l1_weights is not None
         assert self._transformed_l2_weights is not None
         deep_gemm.fp8_fp4_mega_moe(
@@ -1543,9 +1539,12 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
         loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        self.process_weights_after_loading()
+        return loaded_params
+
+    def process_weights_after_loading(self) -> None:
         self.model.finalize_mega_moe_weights()
         self.model.finalize_mhc_broadcast_weights()
-        return loaded_params
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         return self.model.get_expert_mapping()

vllm/models/deepseek_v4/nvidia/mtp.py

@@ -502,14 +502,17 @@ def _find_mtp_layer_idx(name: str) -> int:
                     f"Use a checkpoint that includes MTP layer weights, "
                     f"or disable speculative decoding."
                 )
-        self.finalize_mega_moe_weights()
+        self.process_weights_after_loading()
         logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
         return loaded_params
 
     def finalize_mega_moe_weights(self) -> None:
         for layer in self.model.layers.values():
             layer.mtp_block.ffn.finalize_mega_moe_weights()
 
+    def process_weights_after_loading(self) -> None:
+        self.finalize_mega_moe_weights()
+
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
         """
         Rewrite the weight name to match the format of the original model.