[TRTLLM-13458][feat] Support Minimax M3 NVFP4 checkpoint

tensorrt_llm/_torch/model_config.py

@@ -45,6 +45,11 @@
 _DEEPSEEK_V4_ARCHITECTURES = {"DeepseekV4ForCausalLM"}
 _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT = "layers.0.ffn.experts.0.w1.weight"
 
+_MINIMAX_M3_ARCHITECTURES = {
+    "MiniMaxM3SparseForCausalLM",
+    "MiniMaxM3SparseForConditionalGeneration",
+}
+
 
 def _unified_kv_pool_includes_mamba(
         is_disagg: bool, spec_config: Optional['SpeculativeConfig']) -> bool:
@@ -674,6 +679,81 @@ def _set_deepseek_v4_routed_moe_quant_config(pretrained_config,
             experts_quant_config.quant_algo)
         return layer_quant_config
 
+    @staticmethod
+    def _set_minimax_m3_layer_quant_config(pretrained_config,
+                                           layer_quant_config):
+        """Normalize the Minimax M3 MIXED_PRECISION per-layer quant config.
+
+        Two fix-ups are applied:
+
+        1. Strip the ``language_model.`` prefix from every per-layer key.
+           The M3 VL checkpoint stores keys like
+           ``language_model.model.layers.0.self_attn.o_proj -> MXFP8`` in
+           ``hf_quant_config.json``, but the TRT-LLM module tree names the text
+           decoder ``model.layers.0.self_attn.o_proj`` (no ``language_model.``
+           prefix -- the loader strips it). ``apply_layerwise_quant_config``
+           matches *standalone* Linears (e.g. ``o_proj``, ``down_proj``) with an
+           **exact** ``name == key`` comparison, so the prefixed keys never match
+           and those layers silently fall back to the global ``MIXED_PRECISION``
+           config -> loaded unquantized (MXFP8 ``weight_scale`` dropped) -> the
+           attention/MLP output magnitude explodes. Stripping the prefix makes
+           the exact match succeed. (Fused qkv/gate_up Linears and the Attention
+           wrapper use substring matches and happened to work regardless.)
+
+        2. Inject a single coarse ``model.layers.N.block_sparse_moe.experts``
+           entry per MoE layer so ``MiniMaxM3MoE._get_experts_quant_config`` can
+           select the NVFP4 backend for the routed experts (the fine-grained
+           per-linear NVFP4 expert keys can't be used directly).
+
+        Does nothing when there is no per-layer config (e.g. the uniform MXFP8
+        or a BF16 checkpoint).
+        """
+        from tensorrt_llm.models.modeling_utils import QuantAlgo
+        if layer_quant_config is None:
+            return layer_quant_config
+
+        # (1) Strip the ``language_model.`` prefix so exact-match per-layer
+        # quant assignment works for standalone base Linears.
+        _LM_PREFIX = "language_model."
+        layer_quant_config = {
+            (k[len(_LM_PREFIX):] if k.startswith(_LM_PREFIX) else k): v
+            for k, v in layer_quant_config.items()
+        }
+
+        # (2) Inject coarse NVFP4 expert entries (only when routed experts are
+        # NVFP4).
+        has_nvfp4_experts = any(
+            "block_sparse_moe.experts" in k and isinstance(v, QuantConfig)
+            and v.quant_algo == QuantAlgo.NVFP4
+            for k, v in layer_quant_config.items())
+        if not has_nvfp4_experts:
+            return layer_quant_config
+
+        experts_quant_config = QuantConfig()
+        experts_quant_config.quant_algo = QuantAlgo.NVFP4
+        # TODO: remove the hardcoded group_size and read it from the per-linear
+        # NVFP4 expert entries in hf_quant_config.json instead. 16 is correct
+        # for standard NVFP4 today, but this is a latent bug if a checkpoint
+        # ever ships a different group size.
+        experts_quant_config.group_size = 16
+
+        text_config = getattr(pretrained_config, "text_config",
+                              pretrained_config)
+        if isinstance(text_config, dict):
+            moe_layer_freq = text_config.get("moe_layer_freq", [])
+        else:
+            moe_layer_freq = getattr(text_config, "moe_layer_freq", [])
+
+        for layer_idx, freq in enumerate(moe_layer_freq):
+            if int(freq) != 0:
+                layer_quant_config[
+                    f"model.layers.{layer_idx}.block_sparse_moe.experts"] = experts_quant_config
+
+        logger.info(
+            "Detected Minimax M3 NVFP4 routed MoE checkpoint; using NVFP4 "
+            "for routed experts.")
+        return layer_quant_config
+
     @staticmethod
     def load_quant_config_from_dtypes_json(dtypes_json_file, moe_backend: str):
         quant_config = QuantConfig()
@@ -1060,6 +1140,10 @@ def _recursive_update_config(config: transformers.PretrainedConfig,
                 kwargs.get('spec_config', None),
                 require_layout=require_deepseek_v4_routed_moe_layout)
 
+        if architecture in _MINIMAX_M3_ARCHITECTURES:
+            layer_quant_config = cls._set_minimax_m3_layer_quant_config(
+                pretrained_config, layer_quant_config)
+
         model_config = cls(pretrained_config=pretrained_config,
                            quant_config=quant_config,
                            quant_config_dict=layer_quant_config,

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -31,6 +31,7 @@
 
 from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType
 from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.models.modeling_utils import QuantConfig
 
 from ..attention_backend import AttentionMetadata
 from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams
@@ -352,6 +353,24 @@ class MiniMaxM3MoE(nn.Module):
     that the previous independent-reduction wiring incurred.
     """
 
+    @staticmethod
+    def _get_experts_quant_config(model_config: "ModelConfig", layer_idx: int) -> QuantConfig:
+        """Return the per-layer quant config for the routed experts.
+
+        For MIXED_PRECISION checkpoints (MXFP8 base + NVFP4 experts),
+        ``ModelConfig._set_minimax_m3_moe_quant_config`` pre-populates
+        ``quant_config_dict`` with coarse entries keyed by
+        ``model.layers.N.block_sparse_moe.experts``.  Falls back to the
+        global ``quant_config`` when no per-layer entry exists (e.g. BF16
+        or user-supplied global NVFP4 config).
+        """
+        if getattr(model_config, "quant_config_dict", None) is None:
+            return model_config.quant_config
+        return model_config.quant_config_dict.get(
+            f"model.layers.{layer_idx}.block_sparse_moe.experts",
+            model_config.quant_config,
+        )
+
     def __init__(
         self,
         model_config: "ModelConfig[PretrainedConfig]",
@@ -415,13 +434,15 @@ def __init__(
         # (matches DeepSeekV3). Under Attention DP the fused MoE already
         # skips its in-op all-reduce, so ``reduce_results=False`` is
         # also the correct flag there.
+        experts_quant_config = MiniMaxM3MoE._get_experts_quant_config(model_config, layer_idx)
         self.experts = create_moe(
             routing_method=self.gate.routing_method,
             num_experts=self.num_experts,
             aux_stream_dict=aux_stream_dict,
             reduce_results=False,
             model_config=model_config,
             layer_idx=layer_idx,
+            override_quant_config=experts_quant_config,
             swiglu_alpha=self.swiglu_alpha,
             swiglu_beta=self.swiglu_beta,
             swiglu_limit=self.swiglu_limit,

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

@@ -325,6 +325,10 @@ def _to_trtllm_gen_activation_type(self,
                                        activation_type: ActivationType) -> int:
         if activation_type == ActivationType.Swiglu:
             return 0
+        elif activation_type == ActivationType.SwigluBias:
+            # SwigluBias uses the same SwiGlu kernel path (ActType::SwiGlu == 0);
+            # the per-expert alpha/beta/clamp_limit are passed as separate tensors.
+            return 0
         elif activation_type == ActivationType.Relu2:
             return 1
         elif activation_type == ActivationType.Silu:

tests/integration/defs/accuracy/references/gsm8k.yaml

@@ -447,6 +447,9 @@ MiniMaxAI/MiniMax-M3:
 MiniMaxAI/MiniMax-M3-MXFP8:
   - quant_algo: MXFP8
     accuracy: 89
+nvidia/MiniMax-M3-NVFP4:
+  - quant_algo: MIXED_PRECISION
+    accuracy: 88
 nvidia/NVIDIA-Nemotron-Nano-9B-v2:
   - accuracy: 85.027
   - quant_algo: FP8

tests/integration/defs/accuracy/references/mmlu.yaml

@@ -278,6 +278,9 @@ MiniMaxAI/MiniMax-M3:
 MiniMaxAI/MiniMax-M3-MXFP8:
   - quant_algo: MXFP8
     accuracy: 85
+nvidia/MiniMax-M3-NVFP4:
+  - quant_algo: MIXED_PRECISION
+    accuracy: 83
 moonshotai/Kimi-K2-Instruct:
   - quant_algo: FP8_BLOCK_SCALES
     accuracy: 87.65

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -7233,10 +7233,11 @@ class TestMiniMaxM3(LlmapiAccuracyTestHarness):
     # config requires ``trust_remote_code`` and the runtime requires the
     # MiniMax-M3 sparse attention backend + matching KV-cache manager v2
     # (selected by ``sparse_attention_config``). Blackwell-only (SM100+).
-    # Two checkpoints are covered: the upstream BF16 checkpoint and the
+    # Three checkpoints are covered: the upstream BF16 checkpoint, the
     # NVIDIA MXFP8 checkpoint (weights in MXFP8, activations / KV cache
-    # remain BF16; the runtime path is identical aside from the quant
-    # config).
+    # remain BF16), and the NVIDIA NVFP4 checkpoint (MXFP8 base layers +
+    # NVFP4 routed experts); the runtime path is identical aside from the
+    # quant config.
     MODEL_NAME = "MiniMaxAI/MiniMax-M3"
     MODEL_PATH = f"{llm_models_root()}/MiniMax-M3"
 
@@ -7261,16 +7262,20 @@ def test_auto_dtype(self, tp_size, ep_size):
             task = GSM8K(self.MODEL_NAME)
             task.evaluate(llm)
 
-    @pytest.mark.skip_less_device(8)
+    @pytest.mark.skip_less_device(4)
     @pytest.mark.skip_less_device_memory(140000)
-    @parametrize_with_ids("tp_size,ep_size", [(8, 8)])
+    @parametrize_with_ids("tp_size,ep_size", [(4, 4)])
     def test_mxfp8(self, tp_size, ep_size):
         # MXFP8 checkpoint: weights are MXFP8 (e4m3 + UE8M0 1x32 block
         # scales) with MXFP8 dynamic activations; the KV cache stays in
         # BF16 and the sparse attention path is unchanged from BF16.
         model_name = "MiniMaxAI/MiniMax-M3-MXFP8"
         model_path = f"{llm_models_root()}/MiniMax-M3-MXFP8"
-        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6,
+        # Halving TP from the BF16 reference (TP=8) doubles per-rank
+        # model + KV footprint; cap KV cache at 0.4 of free memory and
+        # constrain batch / token budget so the runtime allocator stays
+        # under the PyTorch cap.
+        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4,
                                         enable_block_reuse=False)
         sparse_attention_config = MiniMaxM3SparseAttentionConfig()
         with LLM(model_path,
@@ -7279,13 +7284,42 @@ def test_mxfp8(self, tp_size, ep_size):
                  kv_cache_config=kv_cache_config,
                  sparse_attention_config=sparse_attention_config,
                  max_seq_len=4096,
+                 max_batch_size=32,
+                 max_num_tokens=4096,
                  trust_remote_code=True) as llm:
             assert llm.args.quant_config.quant_algo == QuantAlgo.MXFP8
             task = MMLU(model_name)
             task.evaluate(llm)
             task = GSM8K(model_name)
             task.evaluate(llm)
 
+    @pytest.mark.skip_less_device(4)
+    @pytest.mark.skip_less_device_memory(140000)
+    @parametrize_with_ids("tp_size,ep_size", [(4, 4)])
+    def test_nvfp4(self, tp_size, ep_size):
+        # NVFP4 checkpoint: MXFP8 base layers with NVFP4 routed experts
+        # (MIXED_PRECISION checkpoint); the KV cache stays in BF16 and the
+        # sparse attention path is unchanged from BF16.
+        model_name = "nvidia/MiniMax-M3-NVFP4"
+        model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4"
+        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6,
+                                        enable_block_reuse=False)
+        sparse_attention_config = MiniMaxM3SparseAttentionConfig()
+        moe_config = MoeConfig(backend="CUTLASS")
+        with LLM(model_path,
+                 tensor_parallel_size=tp_size,
+                 moe_expert_parallel_size=ep_size,
+                 kv_cache_config=kv_cache_config,
+                 sparse_attention_config=sparse_attention_config,
+                 moe_config=moe_config,
+                 max_seq_len=4096,
+                 trust_remote_code=True) as llm:
+            assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION
+            task = MMLU(model_name)
+            task.evaluate(llm)
+            task = GSM8K(model_name)
+            task.evaluate(llm)
+
 
 @skip_pre_blackwell
 class TestGLM5FP8(LlmapiAccuracyTestHarness):

tests/integration/test_lists/qa/llm_function_core.txt

@@ -664,7 +664,8 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch
 accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4]
 accuracy/test_llm_api_pytorch.py::TestMiniMaxM2_5::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4]
 accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] TIMEOUT (180)
-accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=8-ep_size=8] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] TIMEOUT (180)
 accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_auto_dtype
 accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_fp8
 accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_fp8[latency_moe_deepgemm]