diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index 6ddac1bc05..d49769c843 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -39,11 +39,12 @@ from vllm.models.deepseek_v4.common.ops.fused_mtp_input_rmsnorm import ( _FUSED_MTP_INPUT_RMSNORM_KERNEL, _MTP_SHARED_HEAD_RMSNORM_KERNEL, ) +from vllm.models.deepseek_v4.quant_config import align_mtp_quant_config from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from .mega_moe import finalize_mega_moe_weights -from .model import DeepseekV4DecoderLayer +from .model import DeepseekV4DecoderLayer, DeepseekV4ForCausalLM logger = init_logger(__name__) @@ -289,10 +290,15 @@ class DeepSeekV4MultiTokenPredictor(nn.Module): class DeepSeekV4MTP(nn.Module): + # The draft head reuses DeepseekV4DecoderLayer, so its fused linears must + # unfold to the same checkpoint shards for per-layer quant config lookup. + packed_modules_mapping = DeepseekV4ForCausalLM.packed_modules_mapping + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config self.quant_config = vllm_config.quant_config + align_mtp_quant_config(self.quant_config, self.config) self.model = DeepSeekV4MultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) diff --git a/vllm/models/deepseek_v4/quant_config.py b/vllm/models/deepseek_v4/quant_config.py index f3aaf72d9b..30f5d369a3 100644 --- a/vllm/models/deepseek_v4/quant_config.py +++ b/vllm/models/deepseek_v4/quant_config.py @@ -21,11 +21,41 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( _DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8") if TYPE_CHECKING: + from transformers import PreTrainedConfig + + from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4Config, ) +def align_mtp_quant_config( + quant_config: QuantizationConfig | None, config: PreTrainedConfig +) -> None: + """Re-key the quant config's MTP entries onto the draft module prefixes. + + ``_make_deepseek_v4_weights_mapper`` renames the checkpoint's ``mtp.{i}.`` + layers to ``model.mtp.{i}.``, but the draft modules are built under + ``model.layers.{num_hidden_layers + i}.``. Without this remap no per-layer + entry matches, so every MTP linear falls back to the global scheme. On a + mixed checkpoint (FP8 attention and shared experts, NVFP4 routed experts) + that builds NVFP4 parameters for FP8 weights and the load then fails a + shape assertion. + """ + from vllm.model_executor.models.utils import WeightsMapper + + if quant_config is None: + return + quant_config.apply_vllm_mapper( + WeightsMapper( + orig_to_new_prefix={ + f"model.mtp.{i}.": f"model.layers.{config.num_hidden_layers + i}." + for i in range(config.num_nextn_predict_layers) + } + ) + ) + + class DeepseekV4FP8Config(Fp8Config): """FP8 config for DeepSeek V4 with expert-dtype-aware MoE dispatch.