Detect MTP routed-expert layout separately (dense NVFP4 vs MTP MXFP4)

tensorrt_llm/_torch/model_config.py

@@ -681,15 +681,47 @@ def _set_deepseek_v4_routed_moe_quant_config(pretrained_config,
         else:
             layer_quant_config = dict(layer_quant_config)
 
-        num_moe_layers = pretrained_config.num_hidden_layers
-        if (spec_config is not None
-                and spec_config.spec_dec_mode.is_mtp_one_model()):
-            num_moe_layers += spec_config.num_nextn_predict_layers
-
-        for layer_idx in range(num_moe_layers):
+        num_hidden = pretrained_config.num_hidden_layers
+        for layer_idx in range(num_hidden):
             layer_quant_config[
                 f"model.layers.{layer_idx}.mlp.experts"] = experts_quant_config
 
+        # The MTP layers' routed experts can carry a DIFFERENT layout than the
+        # dense layers: the ModelOpt experts-only repack (e.g.
+        # nvidia/DeepSeek-V4-Pro-NVFP4) re-quantizes only the dense experts to
+        # NVFP4 (U8) and leaves the MTP routed experts at the base model's MXFP4
+        # (I8). Detect the MTP expert layout separately so the MTP layers do not
+        # inherit the dense NVFP4 config and crash in fused_moe load_quant_scales.
+        num_mtp = 0
+        if (spec_config is not None
+                and spec_config.spec_dec_mode.is_mtp_one_model()):
+            num_mtp = spec_config.num_nextn_predict_layers or 0
+        if num_mtp:
+            mtp_info = ModelConfig._get_safetensors_header_for_tensor(
+                checkpoint_dir, "mtp.0.ffn.experts.0.w1.weight")
+            mtp_dtype = mtp_info.get("dtype") if mtp_info else None
+            mtp_layout = {"I8": "mxfp4", "U8": "nvfp4"}.get(mtp_dtype)
+            if mtp_layout is not None and mtp_layout != layout:
+                mtp_experts_quant_config = QuantConfig()
+                if mtp_layout == "mxfp4":
+                    mtp_experts_quant_config.quant_algo = (
+                        ModelConfig.get_mxfp4_quant_algo(moe_backend))
+                    mtp_experts_quant_config.group_size = 32
+                else:
+                    mtp_experts_quant_config.quant_algo = QuantAlgo.NVFP4
+                    mtp_experts_quant_config.group_size = 16
+                mtp_experts_quant_config.exclude_modules = (
+                    experts_quant_config.exclude_modules)
+                logger.info(
+                    "DeepSeek-V4 MTP routed experts use a different layout (%s) "
+                    "than the dense experts (%s).", mtp_layout, layout)
+            else:
+                mtp_experts_quant_config = experts_quant_config
+            for i in range(num_mtp):
+                layer_quant_config[
+                    f"model.layers.{num_hidden + i}.mlp.experts"] = (
+                        mtp_experts_quant_config)
+
         logger.info(
             "Detected DeepSeek-V4 routed MoE %s checkpoint layout; using "
             "%s for routed experts.", layout.upper(),