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(),