From 5442601c673db4cb5b49106fac6fab3ae121e16d Mon Sep 17 00:00:00 2001 From: d3nb Date: Sat, 11 Jul 2026 17:06:13 +0000 Subject: [PATCH] [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 dd92337543c6..6487412d5643 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/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(),