[DeepSeek-V4][fix] Detect MTP routed-expert layout separately

tensorrt_llm/_torch/model_config.py

@@ -60,6 +60,7 @@
 
 _DEEPSEEK_V4_ARCHITECTURES = {"DeepseekV4ForCausalLM"}
 _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT = "layers.0.ffn.experts.0.w1.weight"
+_DEEPSEEK_V4_MTP_ROUTED_EXPERT_WEIGHT = "mtp.0.ffn.experts.0.w1.weight"
 
 _MINIMAX_M3_ARCHITECTURES = {
     "MiniMaxM3SparseForCausalLM",
@@ -608,9 +609,11 @@ def _get_safetensors_header_for_tensor(checkpoint_dir: str,
 
     @staticmethod
     def _detect_deepseek_v4_routed_moe_layout(
-            checkpoint_dir: str) -> Optional[str]:
+        checkpoint_dir: str,
+        tensor_name: str = _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT,
+    ) -> Optional[str]:
         tensor_info = ModelConfig._get_safetensors_header_for_tensor(
-            checkpoint_dir, _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT)
+            checkpoint_dir, tensor_name)
         if tensor_info is None:
             return None
 
@@ -622,6 +625,24 @@ def _detect_deepseek_v4_routed_moe_layout(
             return "nvfp4"
         return None
 
+    @staticmethod
+    def _make_routed_experts_quant_config(layout: str,
+                                          moe_backend: str) -> QuantConfig:
+        """Build the routed-experts QuantConfig for a detected MoE layout."""
+        quant_config = QuantConfig()
+        if layout == "mxfp4":
+            quant_config.quant_algo = ModelConfig.get_mxfp4_quant_algo(
+                moe_backend)
+            quant_config.group_size = 32
+        else:
+            quant_config.quant_algo = QuantAlgo.NVFP4
+            quant_config.group_size = 16
+        quant_config.exclude_modules = [
+            'block.*.attn.out', 'block.*.mlp.gate', 'block.*.attn.qkv',
+            'embedding', 'unembedding'
+        ]
+        return quant_config
+
     @staticmethod
     def _is_deepseek_v4_base_checkpoint(checkpoint_dir: str) -> bool:
         tensor_info = ModelConfig._get_safetensors_header_for_tensor(
@@ -661,33 +682,59 @@ def _set_deepseek_v4_routed_moe_quant_config(pretrained_config,
                     "for MXFP4 or U8 for NVFP4.")
             return layer_quant_config
 
-        experts_quant_config = QuantConfig()
-        if layout == "mxfp4":
-            experts_quant_config.quant_algo = ModelConfig.get_mxfp4_quant_algo(
-                moe_backend)
-            experts_quant_config.group_size = 32
-        else:
-            experts_quant_config.quant_algo = QuantAlgo.NVFP4
-            experts_quant_config.group_size = 16
-        experts_quant_config.exclude_modules = [
-            'block.*.attn.out', 'block.*.mlp.gate', 'block.*.attn.qkv',
-            'embedding', 'unembedding'
-        ]
+        experts_quant_config = ModelConfig._make_routed_experts_quant_config(
+            layout, moe_backend)
 
         if layer_quant_config is None:
             layer_quant_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_layout = ModelConfig._detect_deepseek_v4_routed_moe_layout(
+                checkpoint_dir, _DEEPSEEK_V4_MTP_ROUTED_EXPERT_WEIGHT)
+            if mtp_layout is None:
+                # Probe missed (unknown naming convention, MTP shard absent from
+                # the index, or an unexpected dtype/rank). Falling back to the
+                # dense config is exactly the failure this detection exists to
+                # prevent, so say so instead of failing later inside
+                # fused_moe load_quant_scales.
+                logger.warning(
+                    "DeepSeek-V4 MTP routed-expert layout could not be detected "
+                    f"from {_DEEPSEEK_V4_MTP_ROUTED_EXPERT_WEIGHT}; assuming the "
+                    f"dense layout ({layout}). If this checkpoint "
+                    "stores its MTP routed experts in a different format, "
+                    "loading will fail in fused_moe load_quant_scales.")
+                mtp_experts_quant_config = experts_quant_config
+            elif mtp_layout != layout:
+                mtp_experts_quant_config = (
+                    ModelConfig._make_routed_experts_quant_config(
+                        mtp_layout, moe_backend))
+                logger.info(
+                    f"DeepSeek-V4 MTP routed experts use a different layout "
+                    f"({mtp_layout}) than the dense experts ({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(),

tests/unittest/_torch/modeling/test_modeling_deepseekv4.py

@@ -555,6 +555,110 @@ class MTPConfig:
         assert layer_quant_config[f"model.layers.{layer_idx}.mlp.experts"].quant_algo == quant_algo
 
 
+def _write_deepseek_v4_mtp_checkpoint(tmp_path, dense_dtype, mtp_dtype):
+    """Two-shard checkpoint whose dense and MTP routed experts may differ.
+
+    Mirrors the ModelOpt experts-only repacks (e.g. nvidia/DeepSeek-V4-Pro-NVFP4),
+    which re-quantize the dense experts but leave the MTP shard at the base
+    model's layout. Passing mtp_dtype=None omits the MTP entry entirely.
+    """
+    dense_tensor = "layers.0.ffn.experts.0.w1.weight"
+    mtp_tensor = "mtp.0.ffn.experts.0.w1.weight"
+    dense_shard = "model-00001-of-00002.safetensors"
+    mtp_shard = "model-00002-of-00002.safetensors"
+
+    _write_safetensors_header(tmp_path / dense_shard, dense_tensor, dense_dtype, [2, 2])
+    weight_map = {dense_tensor: dense_shard}
+    if mtp_dtype is not None:
+        _write_safetensors_header(tmp_path / mtp_shard, mtp_tensor, mtp_dtype, [2, 2])
+        weight_map[mtp_tensor] = mtp_shard
+
+    (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map}))
+
+
+def _deepseek_v4_mtp_spec_config(num_nextn_predict_layers):
+    class MTPMode:
+        @staticmethod
+        def is_mtp_one_model():
+            return True
+
+    class MTPConfig:
+        spec_dec_mode = MTPMode()
+
+    MTPConfig.num_nextn_predict_layers = num_nextn_predict_layers
+    return MTPConfig()
+
+
+@pytest.mark.parametrize(
+    "dense_dtype,mtp_dtype,dense_algo,dense_group,mtp_algo,mtp_group",
+    [
+        # ModelOpt experts-only NVFP4 repack: dense re-quantized to NVFP4 (U8),
+        # MTP left at the base model's MXFP4 (I8). This is the layout that
+        # crashed fused_moe load_quant_scales before the split detection.
+        ("U8", "I8", QuantAlgo.NVFP4, 16, QuantAlgo.W4A8_MXFP4_MXFP8, 32),
+        # Inverse direction, to pin that neither layout is hard-coded.
+        ("I8", "U8", QuantAlgo.W4A8_MXFP4_MXFP8, 32, QuantAlgo.NVFP4, 16),
+    ],
+)
+def test_deepseek_v4_mtp_routed_experts_detected_separately(
+    tmp_path, monkeypatch, dense_dtype, mtp_dtype, dense_algo, dense_group, mtp_algo, mtp_group
+):
+    monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 100)
+    _write_deepseek_v4_mtp_checkpoint(tmp_path, dense_dtype, mtp_dtype)
+
+    num_hidden_layers = 2
+    num_mtp = 3
+    layer_quant_config = ModelConfig._set_deepseek_v4_routed_moe_quant_config(
+        DeepseekV4Config(num_hidden_layers=num_hidden_layers),
+        str(tmp_path),
+        "TRTLLM",
+        None,
+        _deepseek_v4_mtp_spec_config(num_mtp),
+    )
+
+    for layer_idx in range(num_hidden_layers):
+        dense_config = layer_quant_config[f"model.layers.{layer_idx}.mlp.experts"]
+        assert dense_config.quant_algo == dense_algo
+        assert dense_config.group_size == dense_group
+
+    for i in range(num_mtp):
+        mtp_config = layer_quant_config[f"model.layers.{num_hidden_layers + i}.mlp.experts"]
+        assert mtp_config.quant_algo == mtp_algo
+        assert mtp_config.group_size == mtp_group
+
+
+def test_deepseek_v4_mtp_routed_experts_warn_when_probe_missing(tmp_path, monkeypatch):
+    """An undetectable MTP layout must fall back loudly, not silently.
+
+    The warning is captured by patching logger.warning rather than via caplog:
+    tensorrt_llm's logger sets propagate=False, so root-handler capture would
+    not see it.
+    """
+    monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 100)
+    _write_deepseek_v4_mtp_checkpoint(tmp_path, "U8", None)
+
+    warnings = []
+    monkeypatch.setattr(
+        "tensorrt_llm._torch.model_config.logger.warning",
+        warnings.append,
+    )
+
+    num_hidden_layers = 2
+    num_mtp = 1
+    layer_quant_config = ModelConfig._set_deepseek_v4_routed_moe_quant_config(
+        DeepseekV4Config(num_hidden_layers=num_hidden_layers),
+        str(tmp_path),
+        "TRTLLM",
+        None,
+        _deepseek_v4_mtp_spec_config(num_mtp),
+    )
+
+    # Falls back to the dense layout, and says which probe missed.
+    mtp_config = layer_quant_config[f"model.layers.{num_hidden_layers}.mlp.experts"]
+    assert mtp_config.quant_algo == QuantAlgo.NVFP4
+    assert any("mtp.0.ffn.experts.0.w1.weight" in w for w in warnings)
+
+
 def test_deepseek_v4_mtp_projection_uses_fp8_quant_config(monkeypatch):
     def fake_decoder_layer_init(self, model_config, *_args, **_kwargs):
         torch.nn.Module.__init__(self)