diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 25ea9952c94c..ee1e765db6cf 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/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.", diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index b702d7ec9cbd..1a3c2c780d28 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -504,9 +504,9 @@ 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 @@ -514,6 +514,9 @@ 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. diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 17f29f4a1ae1..bd4e7614563b 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/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() diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index eaf105205c1c..7dd32fe97267 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -502,7 +502,7 @@ 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 @@ -510,6 +510,9 @@ 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.