From 8d38f0fa92a6156d16ff43bf8a885b5e3171032c Mon Sep 17 00:00:00 2001 From: d3nb Date: Sat, 11 Jul 2026 17:06:13 +0000 Subject: [PATCH 1/3] [fix] Detect MTP routed-expert layout separately (dense NVFP4 vs MTP MXFP4) The ModelOpt experts-only NVFP4 repacks (e.g. nvidia/DeepSeek-V4-Pro-NVFP4) re-quantize only the dense routed experts to NVFP4 (U8) and leave the MTP routed experts at the base model's MXFP4 (I8). _set_deepseek_v4_routed_moe_ quant_config detected a single layout from layers.0 and applied it to every MoE layer including the MTP layer, so the MTP experts got NVFP4 and crashed in fused_moe load_quant_scales. Detect the MTP expert dtype separately and assign the MTP layer indices the correct (MXFP4) config. Validated end-to-end on real weights: nvidia/DeepSeek-V4-Pro-NVFP4 on 4x B300, TP4, moe_backend=TRTLLM, rc15.post1 (equivalent change), with the construction fix from 3972f5e: loads 100% + serves + generates, MTP=1 and MTP=3 (mtp_eagle_one_model builds one shared MTP layer). MTP accept_len ~2.86. Signed-off-by: d3nb --- tensorrt_llm/_torch/model_config.py | 44 +++++++++++++++++++++++++---- 1 file changed, 38 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index dbab6f07b75a..fa04e9109bb4 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -679,15 +679,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(), From 627d9833ed9a60f262d4126b6f8570403f8798d3 Mon Sep 17 00:00:00 2001 From: d3nb Date: Fri, 14 Aug 2026 09:55:14 +0800 Subject: [PATCH 2/3] [fix] Cover DeepSeek-V4 MTP routed-expert layouts Signed-off-by: WEI CHENG CHIU Signed-off-by: d3nb --- tensorrt_llm/_torch/model_config.py | 73 +++++++----- .../modeling/test_modeling_deepseekv4.py | 104 ++++++++++++++++++ 2 files changed, 148 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index fa04e9109bb4..c1e141c91b70 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/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,18 +682,8 @@ 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 = {} @@ -695,21 +706,25 @@ def _set_deepseek_v4_routed_moe_quant_config(pretrained_config, 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) + 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 " + "from %s; assuming the dense layout (%s). If this checkpoint " + "stores its MTP routed experts in a different format, " + "loading will fail in fused_moe load_quant_scales.", + _DEEPSEEK_V4_MTP_ROUTED_EXPERT_WEIGHT, layout) + 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( "DeepSeek-V4 MTP routed experts use a different layout (%s) " "than the dense experts (%s).", mtp_layout, layout) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index b46ed3e477b9..608a2e5a3d5c 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/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", + lambda msg, *args: warnings.append(msg % args if args else msg), + ) + + 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) From d906f2f50ff80ae25d6e6bcc67141969d1c1ddd4 Mon Sep 17 00:00:00 2001 From: WEI CHENG CHIU Date: Fri, 14 Aug 2026 10:04:40 +0800 Subject: [PATCH 3/3] [fix] Format DeepSeek-V4 MTP layout diagnostics Signed-off-by: WEI CHENG CHIU --- tensorrt_llm/_torch/model_config.py | 10 +++++----- .../_torch/modeling/test_modeling_deepseekv4.py | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index c1e141c91b70..867594ea1baa 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -716,18 +716,18 @@ def _set_deepseek_v4_routed_moe_quant_config(pretrained_config, # fused_moe load_quant_scales. logger.warning( "DeepSeek-V4 MTP routed-expert layout could not be detected " - "from %s; assuming the dense layout (%s). If this checkpoint " + 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.", - _DEEPSEEK_V4_MTP_ROUTED_EXPERT_WEIGHT, layout) + "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( - "DeepSeek-V4 MTP routed experts use a different layout (%s) " - "than the dense experts (%s).", mtp_layout, layout) + 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): diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 608a2e5a3d5c..cfd443cac756 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -640,7 +640,7 @@ def test_deepseek_v4_mtp_routed_experts_warn_when_probe_missing(tmp_path, monkey warnings = [] monkeypatch.setattr( "tensorrt_llm._torch.model_config.logger.warning", - lambda msg, *args: warnings.append(msg % args if args else msg), + warnings.append, ) num_hidden_layers = 2