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

docs/source/_static/config_db.json

@@ -169,15 +169,15 @@
       "scenario": "Max Throughput"
     },
     {
-      "command": "trtllm-serve MiniMaxAI/MiniMax-M3 --config ${TRTLLM_DIR}/examples/configs/curated/minimax-m3-throughput.yaml",
+      "command": "trtllm-serve MiniMaxAI/MiniMax-M3-MXFP8 --config ${TRTLLM_DIR}/examples/configs/curated/minimax-m3-throughput.yaml",
       "config_filename": "minimax-m3-throughput.yaml",
       "config_github_url": "https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/configs/curated/minimax-m3-throughput.yaml",
       "config_path": "examples/configs/curated/minimax-m3-throughput.yaml",
       "config_raw_url": "https://raw.githubusercontent.com/NVIDIA/TensorRT-LLM/main/examples/configs/curated/minimax-m3-throughput.yaml",
       "gpu_compatibility": "GB200",
-      "model": "MiniMaxAI/MiniMax-M3",
-      "model_display_name": "MiniMax-M3 (BF16)",
-      "model_url": "https://huggingface.co/MiniMaxAI/MiniMax-M3",
+      "model": "MiniMaxAI/MiniMax-M3-MXFP8",
+      "model_display_name": "MiniMax-M3 (MXFP8)",
+      "model_url": "https://huggingface.co/MiniMaxAI/MiniMax-M3-MXFP8",
       "scenario": "Max Throughput"
     }
   ],
@@ -3516,9 +3516,9 @@
     }
   ],
   "models": {
-    "MiniMaxAI/MiniMax-M3": {
-      "display_name": "MiniMax-M3 (BF16)",
-      "url": "https://huggingface.co/MiniMaxAI/MiniMax-M3"
+    "MiniMaxAI/MiniMax-M3-MXFP8": {
+      "display_name": "MiniMax-M3 (MXFP8)",
+      "url": "https://huggingface.co/MiniMaxAI/MiniMax-M3-MXFP8"
     },
     "Qwen/Qwen3-30B-A3B": {
       "display_name": "Qwen3-30B-A3B",

docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md

@@ -6,15 +6,20 @@ This deployment guide provides step-by-step instructions for running the MiniMax
 
 MiniMax-M3 is a Mixture-of-Experts (MoE) model that uses MiniMax block-sparse attention. The first few layers use dense attention with a dense MLP, while the remaining layers combine a sparse attention path (an index-K block selector followed by sparse grouped-query attention) with MoE (top-4 of 128 routed experts plus one shared expert). In TensorRT LLM it is served through the `MiniMaxM3SparseForConditionalGeneration` architecture (text, image, and video) and the text-only `MiniMaxM3SparseForCausalLM` architecture.
 
-MiniMax-M3 is served in **BF16**; no FP8/NVFP4 serving path is supported at this time. The block-sparse attention path does **not** currently support KV cache reuse or Multi-Token Prediction (MTP) in this release.
+TensorRT LLM supports two precisions for MiniMax-M3:
+
+* **BF16** — the official upstream checkpoint from MiniMaxAI.
+* **MXFP8** — an NVIDIA-published checkpoint that quantizes the MoE/Linear weights to MXFP8 while keeping activations and the KV cache in BF16. The weights occupy ~half the memory of BF16, which is the recommended choice for throughput-oriented deployments and is the default for this guide.
+
+The block-sparse attention path does **not** currently support KV cache reuse or Multi-Token Prediction (MTP) in this release.
 
 This guide deploys MiniMax-M3 on **8x NVIDIA GB200 GPUs across 2 nodes** (4 GPUs per node) using Slurm and the `trtllm-llmapi-launch` multi-node launcher, with the MoE experts distributed via expert parallelism. The attention layers can run with either Tensor-Expert Parallelism (TEP) or Data-Expert Parallelism (DEP); see [Choosing the Parallelism Strategy](#choosing-the-parallelism-strategy-tep-vs-dep).
 
 The guide is intended for developers and practitioners seeking high-throughput or low-latency inference using NVIDIA's accelerated stack.
 
 ## Prerequisites
 
-* GPU: 8x NVIDIA GB200 GPUs across 2 nodes (4 GPUs per node). Tensor/expert parallelism of 8 (`--tp_size 8 --moe_expert_parallel_size 8`) spans all 8 GPUs, and the model is served in BF16, so plan for the corresponding memory footprint.
+* GPU: 8x NVIDIA GB200 GPUs across 2 nodes (4 GPUs per node). Tensor/expert parallelism of 8 (`tensor_parallel_size: 8`, `moe_expert_parallel_size: 8` in the curated YAML) spans all 8 GPUs. Plan for the corresponding memory footprint — MXFP8 weights occupy roughly half of BF16.
 * Multi-node launcher: Slurm with the pyxis/enroot container plugin (or an equivalent MPI launcher) to start one rank per GPU across both nodes.
 * High-speed inter-node interconnect (e.g., InfiniBand) for tensor/expert-parallel traffic.
 * Shared filesystem visible to both nodes for the model weights and the configuration file.
@@ -24,24 +29,36 @@ The guide is intended for developers and practitioners seeking high-throughput o
 
 ## Models
 
-The following checkpoint is available:
+Two checkpoints are supported. Both are loaded through the same `MiniMaxM3SparseForConditionalGeneration` / `MiniMaxM3SparseForCausalLM` architectures and share the same chat template and serving CLI; the only difference is the on-disk weight format.
+
+### MXFP8 (recommended for throughput)
+
+* [MiniMaxAI/MiniMax-M3-MXFP8](https://huggingface.co/MiniMaxAI/MiniMax-M3-MXFP8) — MiniMaxAI-published MXFP8-quantized checkpoint. Weights are stored in MXFP8 (block size 1×32); activations and the KV cache stay in BF16.
+
+```bash
+git lfs install
+git clone https://huggingface.co/MiniMaxAI/MiniMax-M3-MXFP8 /models/MiniMax-M3-MXFP8
+```
+
+### BF16 (upstream)
 
-* [MiniMaxAI/MiniMax-M3](https://huggingface.co/MiniMaxAI/MiniMax-M3) — Official BF16 checkpoint
+* [MiniMaxAI/MiniMax-M3](https://huggingface.co/MiniMaxAI/MiniMax-M3) — Official BF16 checkpoint from MiniMaxAI.
 
 ```bash
 git lfs install
 git clone https://huggingface.co/MiniMaxAI/MiniMax-M3 /models/MiniMax-M3
 ```
 
-The checkpoint ships its own chat template (`chat_template.jinja`), which is passed explicitly to the server (see [Launch the TensorRT LLM Server](#launch-the-tensorrt-llm-server)).
+Both checkpoints ship their own chat template (`chat_template.jinja`), which is passed explicitly to the server (see [Launch the TensorRT LLM Server](#launch-the-tensorrt-llm-server)).
 
 ## Feature Support Notes
 
 * **Block-sparse attention is required.** MiniMax-M3 runs on the block-sparse attention backend, which must be selected via `sparse_attention_config.algorithm: minimax_m3` in the YAML configuration. There is no dense fallback for the sparse layers.
-* **BF16 only.** MiniMax-M3 is served in BF16. No FP8/NVFP4 serving path is supported at this time. The default MoE backend is used.
+* **Supported precisions: BF16 and MXFP8.** No additional FP8/NVFP4 serving paths are supported at this time. MXFP8 quantizes only the weights; activations and the KV cache stay in BF16, so the curated YAML is identical for both checkpoints. The default MoE backend is used.
 * **KV cache reuse must be disabled.** KV cache reuse is not supported on the sparse-attention path, so set `kv_cache_config.enable_block_reuse: false`.
 * **MTP is not supported** on the sparse-attention path in this release.
-* **Parallelism**: MoE experts run with expert parallelism. The attention layers support both Tensor-Expert Parallelism (TEP) and Data-Expert Parallelism (DEP, via `--enable_attention_dp`). The overlap scheduler and CUDA graphs are also supported.
+* **`max_seq_len` must be capped for CUDA graphs.** The dense GQA expansion in the first few attention layers and the per-Q FP32 expansion in the sparse decode kernel allocate temporary tensors whose size grows linearly with the warmup decode's `max_k`. If `max_seq_len` is left at the checkpoint default, that `max_k` follows `max_position_embeddings` (1M for MXFP8, 512K for BF16) and the resulting gigabyte-scale single-allocation request exceeds the caching allocator's CUDA-graph-safe path, so capture fails with `cudaErrorStreamCaptureUnsupported` / OOM. The curated YAML therefore sets `max_seq_len` to a small value just above ISL+OSL (`2068` for the 1k/1k benchmark) so CUDA graphs can capture cleanly. Raise it for longer-context workloads, but expect a corresponding cut in `max_batch_size`.
+* **Parallelism.** MoE experts run with expert parallelism. The attention layers support both Tensor-Expert Parallelism (TEP) and Data-Expert Parallelism (DEP, via `enable_attention_dp: true`). The overlap scheduler is enabled by default.
 * **Multimodal.** `MiniMaxM3SparseForConditionalGeneration` supports text, image, and video inputs. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture.
 
 ## Deployment Steps
@@ -89,7 +106,7 @@ The configuration uses Data-Expert Parallelism (DEP): `enable_attention_dp: true
 MiniMax-M3 is launched through the `trtllm-llmapi-launch` wrapper, which sets up the multi-rank (MPI/Slurm) environment that the parallel server requires. The wrapper is run once per rank by Slurm (`srun`), with one task (rank) per GPU. The example below launches the server across 2 nodes (`-N 2`), 4 GPUs per node (`--ntasks-per-node 4`, 8 ranks total), using the curated YAML to drive parallelism, batching, and the MiniMax-M3 sparse-attention backend:
 
 ```bash
-export MODEL=/models/MiniMax-M3   # path on the shared filesystem; mounted into the container
+export MODEL=/models/MiniMax-M3-MXFP8   # MXFP8 (recommended); use /models/MiniMax-M3 for BF16
 
 srun -N 2 \
     --ntasks 8 --ntasks-per-node 4 \
@@ -107,7 +124,7 @@ srun -N 2 \
           --extra_llm_api_options $EXTRA_LLM_API_FILE"
 ```
 
-The parallelism, batch, KV-cache, sparse-attention, and CUDA-graph settings all live in the YAML; no CLI flags need to change to tune them.
+The parallelism, batch, KV-cache, sparse-attention, and CUDA-graph settings all live in the YAML; no CLI flags need to change to tune them. The same `$EXTRA_LLM_API_FILE` is used for both the MXFP8 and BF16 checkpoints — point `$MODEL` at the directory of whichever checkpoint you want to serve.
 
 > [!NOTE]
 > Adjust `-N`, `--ntasks`, `--ntasks-per-node`, and `--gres=gpu:` to match your cluster's GPUs-per-node. The total number of tasks (ranks) must equal `tensor_parallel_size` (`8`). Add the partition / account / node-list flags (`-p`, `-A`, `-w`) required by your Slurm setup, and ensure `/models`, `/workspace`, and the TensorRT LLM repository resolve to the same shared paths on both nodes.
@@ -156,7 +173,7 @@ After the TensorRT LLM server is set up and shows *Application startup complete*
 curl http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
-      "model": "MiniMaxAI/MiniMax-M3",
+      "model": "MiniMaxAI/MiniMax-M3-MXFP8",
       "messages": [
           {"role": "user", "content": "What is the capital of France?"}
       ],
@@ -171,7 +188,7 @@ Example response:
 {
   "id": "chatcmpl-...",
   "object": "chat.completion",
-  "model": "MiniMaxAI/MiniMax-M3",
+  "model": "MiniMaxAI/MiniMax-M3-MXFP8",
   "choices": [
     {
       "index": 0,
@@ -220,7 +237,7 @@ result_dir=/tmp/minimax_m3_output
 for concurrency in ${concurrency_list}; do
     num_prompts=$((concurrency * multi_round))
     python -m tensorrt_llm.serve.scripts.benchmark_serving \
-        --model MiniMaxAI/MiniMax-M3 \
+        --model MiniMaxAI/MiniMax-M3-MXFP8 \
         --backend openai \
         --dataset-name "random" \
         --random-input-len ${isl} \

examples/configs/curated/lookup.yaml

@@ -81,8 +81,8 @@
   scenario: Max Throughput
   gpu_compatibility: "B200, GB200"
   disagg: true
-- model: MiniMaxAI/MiniMax-M3
-  arch: MiniMaxM3SparseForCausalLM
+- model: MiniMaxAI/MiniMax-M3-MXFP8
+  arch: MiniMaxM3SparseForConditionalGeneration
   config_path: examples/configs/curated/minimax-m3-throughput.yaml
   scenario: Max Throughput
   gpu_compatibility: "GB200"

examples/configs/curated/minimax-m3-throughput.yaml

@@ -1,13 +1,22 @@
 max_batch_size: 256
 max_num_tokens: 8192
+# Cap max_seq_len explicitly. Without this cap the warmup decode used for
+# CUDA-graph capture inherits the checkpoint's max_position_embeddings
+# (1M for MXFP8, 512K for BF16), which makes both the dense GQA expansion
+# in `_dense_forward` and the per-Q FP32 slab in `_sparse_gqa_masked`
+# allocate gigabyte-scale temporaries that exceed the caching allocator's
+# graph-safe path and fail capture with cudaErrorStreamCaptureUnsupported.
+# 2068 leaves a small margin over ISL+OSL=2048 (1k/1k benchmark).
+max_seq_len: 2068
 tensor_parallel_size: 8
 moe_expert_parallel_size: 8
 enable_attention_dp: true
 trust_remote_code: true
 cuda_graph_config:
   enable_padding: true
+  max_batch_size: 256
 kv_cache_config:
-  free_gpu_memory_fraction: 0.3
+  free_gpu_memory_fraction: 0.7
   enable_block_reuse: false
 sparse_attention_config:
   algorithm: minimax_m3

tensorrt_llm/_torch/modules/fused_moe/routing.py

@@ -677,9 +677,16 @@ def __init__(
         )
         self.routed_scaling_factor = float(routed_scaling_factor)
 
-    def apply(self,
-              router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        topk_idx, top_k_weights = super().apply(router_logits)
+    def apply(
+        self,
+        router_logits: torch.Tensor,
+        input_ids: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # ``moe_scheduler.MoEScheduler`` calls ``routing_method.apply(router_logits,
+        # input_ids)`` positionally; mirror the parent's signature so the
+        # MiniMax-M3 routing accepts (but ignores) the optional ``input_ids``
+        # argument — M3's sigmoid+bias+top-k+renorm flow doesn't consume it.
+        topk_idx, top_k_weights = super().apply(router_logits, input_ids)
         if self.routed_scaling_factor != 1.0:
             top_k_weights = top_k_weights * self.routed_scaling_factor
         return topk_idx, top_k_weights

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

@@ -441,10 +441,12 @@ MiniMaxAI/MiniMax-M2.5:
   - quant_algo: FP8_BLOCK_SCALES
     accuracy: 93.75
 # MiniMax-M3 text decoder (block-sparse MoE) GSM8K, full 1319-sample split.
-# BF16 checkpoint only. Threshold is set ~1pp below the measured TRT-LLM
-# hardpath score (~88) to absorb normal eval noise.
+# Thresholds tracked per checkpoint / quantization scheme.
 MiniMaxAI/MiniMax-M3:
-  - accuracy: 87
+  - accuracy: 89
+MiniMaxAI/MiniMax-M3-MXFP8:
+  - quant_algo: MXFP8
+    accuracy: 89
 nvidia/NVIDIA-Nemotron-Nano-9B-v2:
   - accuracy: 85.027
   - quant_algo: FP8

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

@@ -274,7 +274,10 @@ Qwen/Qwen3.5-397B-A17B:
 MiniMaxAI/MiniMax-M2:
   - accuracy: 83.82
 MiniMaxAI/MiniMax-M3:
-  - accuracy: 80
+  - accuracy: 85
+MiniMaxAI/MiniMax-M3-MXFP8:
+  - quant_algo: MXFP8
+    accuracy: 85
 moonshotai/Kimi-K2-Instruct:
   - quant_algo: FP8_BLOCK_SCALES
     accuracy: 87.65

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -7232,8 +7232,11 @@ class TestMiniMaxM3(LlmapiAccuracyTestHarness):
     # text-only GSM8K / MMLU exercise the text-decoder path. The custom HF
     # 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+);
-    # BF16 checkpoint only.
+    # (selected by ``sparse_attention_config``). Blackwell-only (SM100+).
+    # Two checkpoints are covered: the upstream BF16 checkpoint and the
+    # NVIDIA MXFP8 checkpoint (weights in MXFP8, activations / KV cache
+    # remain BF16; the runtime path is identical aside from the quant
+    # config).
     MODEL_NAME = "MiniMaxAI/MiniMax-M3"
     MODEL_PATH = f"{llm_models_root()}/MiniMax-M3"
 
@@ -7258,6 +7261,31 @@ 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_memory(140000)
+    @parametrize_with_ids("tp_size,ep_size", [(8, 8)])
+    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,
+                                        enable_block_reuse=False)
+        sparse_attention_config = MiniMaxM3SparseAttentionConfig()
+        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,
+                 max_seq_len=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)
+
 
 @skip_pre_blackwell
 class TestGLM5FP8(LlmapiAccuracyTestHarness):

tests/integration/test_lists/qa/llm_function_core.txt

@@ -663,6 +663,7 @@ 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::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]