[Bugfix][DSv4] Remap Quark MTP layer_quant_config onto draft prefixes (#56054)

Signed-off-by: jimmy-adams <41593649+jimmy-adams@users.noreply.github.com>

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")
         )

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.