diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py index 46c7064b0dfd..fa81ff0bc672 100644 --- a/tests/v1/spec_decode/test_mtp.py +++ b/tests/v1/spec_decode/test_mtp.py @@ -57,6 +57,40 @@ def test_glm_mtp_defers_lm_head(default_vllm_config): assert layer.shared_head.head is None +@pytest.mark.parametrize( + "module_name", + ["vllm.models.deepseek_v32.nvidia.mtp", "vllm.models.deepseek_v32.amd.mtp"], +) +def test_deepseek_v32_mtp_defers_lm_head(default_vllm_config, module_name): + import importlib + + from vllm.model_executor.models import deepseek_mtp + + mtp = importlib.import_module(module_name) + + config = mock.MagicMock( + hidden_size=16, + rms_norm_eps=1e-5, + index_topk=8, + model_type="deepseek_mtp", + ) + vllm_config = mock.MagicMock() + vllm_config.speculative_config.draft_model_config.hf_config = config + vllm_config.speculative_config.num_speculative_tokens = 1 + vllm_config.scheduler_config.max_num_batched_tokens = 4 + vllm_config.scheduler_config.max_num_seqs = 1 + + with ( + mock.patch.object(mtp, "DeepseekV32DecoderLayer", return_value=nn.Identity()), + mock.patch.object(deepseek_mtp, "ParallelLMHead") as parallel_lm_head, + mock.patch.object(mtp.current_platform, "device_type", "cpu"), + ): + layer = mtp.DeepseekV32MultiTokenPredictorLayer(vllm_config, "model.layers.1") + + parallel_lm_head.assert_not_called() + assert layer.shared_head.head is None + + def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer: """Create an MTP proposer with unified model configuration.""" model_config = ModelConfig( diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py index 70ddeae374c8..6c7e7ff07654 100644 --- a/vllm/models/deepseek_v32/amd/mtp.py +++ b/vllm/models/deepseek_v32/amd/mtp.py @@ -51,7 +51,6 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: assert vllm_config.speculative_config is not None config = vllm_config.speculative_config.draft_model_config.hf_config self.config = config - quant_config = vllm_config.quant_config self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -64,7 +63,9 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: device=current_platform.device_type, ) self.shared_head = SharedHead( - config=config, prefix=prefix, quant_config=quant_config + config=config, + prefix=prefix, + defer_lm_head=True, ) self.mtp_block = DeepseekV32DecoderLayer( vllm_config, @@ -295,6 +296,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: is_fusion_moe_shared_experts_layer = ( self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) ) + if ".shared_head.head." in name: + continue name = self._rewrite_spec_layer_name(spec_layer, name) if _try_load_fp8_indexer_wk( diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index bdc6b95c15e4..c2a6b2e035b3 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -67,7 +67,6 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: assert vllm_config.speculative_config is not None config = vllm_config.speculative_config.draft_model_config.hf_config self.config = config - quant_config = vllm_config.quant_config self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -89,7 +88,9 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: get_sparse_mla_index_group_max_rows(vllm_config), ) self.shared_head = SharedHead( - config=config, prefix=prefix, quant_config=quant_config + config=config, + prefix=prefix, + defer_lm_head=True, ) self.mtp_block = DeepseekV32DecoderLayer( vllm_config, @@ -392,6 +393,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, loaded_weight) loaded_params.add(name) continue + if ".shared_head.head." in name: + continue name = self._rewrite_spec_layer_name(spec_layer, name) if _try_load_fp8_indexer_wk(