[TRTLLM-17388][feat] MiniMax-M3: load the compressed-tensors NVFP4 MoE packaging

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

@@ -56,8 +56,29 @@ The `MiniMax-M3-NVFP4` checkpoint uses MXFP8 base layers and NVFP4 routed expert
 (`MIXED_PRECISION`). Its name describes weight precision; set
 `kv_cache_config.dtype` explicitly to select FP8 or NVFP4 KV cache.
 
+The compressed-tensors NVFP4 MoE packaging of the MiniMax M3 architecture is
+also loaded as mixed MXFP8/NVFP4. Such a checkpoint declares
+`quant_method: mxfp8` together with `moe_quant_algo: NVFP4`,
+`moe_quant_format: nvfp4-pack-quantized` and `group_size: 16`, and stores each
+routed expert projection as `weight_packed`, `weight_scale`,
+`weight_global_scale` and `input_global_scale`. TensorRT LLM keeps the packed
+FP4 bytes and E4M3 block scales as stored and converts the global scales at
+load time: the stored `2688 / amax` multipliers become the reciprocal
+`weight_scale_2` and `input_scale` divisors of the NVFP4 MoE weight contract.
+A global scale that cannot be inverted (zero, as stored for an expert without
+calibration data, or negative or not finite) is set to 0 with a warning; a
+weight global scale of 0 zeroes that expert's output. The routed experts then
+run the standard static NVFP4 MoE recipe (static input scale, BF16 rounding
+after FC1, routing weights applied in the finalize) on any NVFP4-capable MoE
+backend; `moe_config.backend: CUTLASS` and `TRTLLM` are both qualified on this
+packaging.
+
 Use the checkpoint's chat template (`chat_template.jinja`) when launching the
 server (see [Launch the TensorRT LLM Server](#launch-the-tensorrt-llm-server)).
+Evaluate a checkpoint in this packaging through the chat template as well, for
+example `trtllm-eval ... gsm8k --apply_chat_template --chat_template_kwargs
+'{"thinking_mode": "disabled"}'`; the raw eight-shot prompts of the `gsm8k`
+default bypass the template the model was trained with.
 
 ## Feature Support Notes
 

tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py

@@ -513,27 +513,34 @@ def _build_cache_config(self, config):
             ]
         return super()._build_cache_config(config)
 
-    def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role):
+    def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role) -> int:
         """Report the hybrid sparse-NVFP4 / dense-FP8 storage footprint."""
-        if self.dtype != DataType.NVFP4:
-            return super().get_layer_bytes_per_token(local_layer_idx, data_role)
         global_layer_idx = int(self.pp_layers[int(local_layer_idx)])
-        if global_layer_idx in self.sparse_layer_ids:
-            return super().get_layer_bytes_per_token(local_layer_idx, data_role)
-
-        if data_role in (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE):
-            return 0
-        if data_role == Role.ALL:
-            kv_factor = self.kv_factor
-        elif data_role in (Role.KEY, Role.VALUE):
-            kv_factor = 1
+        has_index_k = global_layer_idx in self.sparse_layer_ids
+        if data_role == Role.INDEX_KEY:
+            return self._index_k_bytes_per_token() if has_index_k else 0
+
+        if self.dtype != DataType.NVFP4 or has_index_k:
+            cache_bytes = super().get_layer_bytes_per_token(local_layer_idx, data_role)
         else:
-            return super().get_layer_bytes_per_token(local_layer_idx, data_role)
-        return (
-            kv_factor
-            * self.num_kv_heads_per_layer[int(local_layer_idx)]
-            * self.head_dim_per_layer[int(local_layer_idx)]
-        )
+            # Existing NVFP4 dense-target / shared-Eagle FP8 cache calculation.
+            if data_role in (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE):
+                return 0
+            if data_role == Role.ALL:
+                kv_factor = self.kv_factor
+            elif data_role in (Role.KEY, Role.VALUE):
+                kv_factor = 1
+            else:
+                return super().get_layer_bytes_per_token(local_layer_idx, data_role)
+            cache_bytes = (
+                kv_factor
+                * self.num_kv_heads_per_layer[int(local_layer_idx)]
+                * self.head_dim_per_layer[int(local_layer_idx)]
+            )
+
+        if data_role == Role.ALL and has_index_k:
+            cache_bytes += self.get_layer_bytes_per_token(local_layer_idx, Role.INDEX_KEY)
+        return cache_bytes
 
     def is_nvfp4_layer(self, layer_idx: int) -> bool:
         """Whether ``layer_idx`` stores packed NVFP4 data and block scales."""
@@ -730,19 +737,17 @@ def update_resources(self, scheduled_batch, attn_metadata=None, kv_cache_dtype_b
             raise NotImplementedError("MiniMax-M3 does not relocate accepted draft tokens.")
         super().update_resources(scheduled_batch, attn_metadata, kv_cache_dtype_byte_size)
 
-    def _extra_buffers_per_layer(self, *, tokens_per_block):
+    def _extra_buffers_per_layer(self, *, tokens_per_block: int) -> dict[int, List[BufferConfig]]:
         """Register a per-sparse-layer ``Role.INDEX_KEY`` :class:`BufferConfig`.
 
-        ``size`` is bytes per **block**: ``1 * sparse_index_dim *
-        elem_bytes * tokens_per_block``. Keyed by **local** layer id —
-        the base ``_build_base_config`` iterates local ids, so keying
+        ``size`` is bytes per **block**: :meth:`_index_k_bytes_per_token`
+        times ``tokens_per_block``, the same per-token size
+        :meth:`get_layer_bytes_per_token` reports. Keyed by **local** layer
+        id — the base ``_build_base_config`` iterates local ids, so keying
         by global ids would silently skip registration on non-trivial
         PP ranks.
         """
-        torch_dtype = self._torch_dtype_for_index_cache()
-        elem_bytes = torch.tensor([], dtype=torch_dtype).element_size()
-        bytes_per_token = 1 * self.sparse_index_dim * elem_bytes
-        size_per_block = bytes_per_token * tokens_per_block
+        size_per_block = self._index_k_bytes_per_token() * tokens_per_block
         return {
             self.layer_offsets[layer_id]: [BufferConfig(role=Role.INDEX_KEY, size=size_per_block)]
             for layer_id in self.sparse_layer_ids
@@ -774,6 +779,11 @@ def _compute_num_total_slots(self) -> int:
         kv_factor = 1 if self.kv_cache_type == CacheTypeCpp.SELFKONLY else 2
         return int((page_upper // kv_factor) * self.tokens_per_block)
 
+    def _index_k_bytes_per_token(self) -> int:
+        """Bytes of one token's index-K row on a sparse layer."""
+        index_dtype = self._torch_dtype_for_index_cache()
+        return self.sparse_index_dim * torch.empty(0, dtype=index_dtype).element_size()
+
     def _torch_dtype_for_index_cache(self) -> torch.dtype:
         """Return the independently configured index-cache storage dtype."""
         if self.indexer_kv_dtype == "fp8":

tensorrt_llm/_torch/model_config.py

@@ -21,8 +21,9 @@
 import tempfile
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import (TYPE_CHECKING, Any, Dict, Generic, Iterator, List, Optional,
-                    Sequence, Tuple, TypeVar)
+from types import SimpleNamespace
+from typing import (TYPE_CHECKING, Any, Callable, Dict, Generic, Iterator, List,
+                    Optional, Sequence, Tuple, TypeVar)
 
 import filelock
 import torch
@@ -994,11 +995,180 @@ def _set_deepseek_v4_routed_moe_quant_config(pretrained_config,
         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.
+    def _get_mxfp8_moe_quant_config(
+            hf_quant_config: Dict[str, Any]) -> Optional[QuantConfig]:
+        """Return the routed-expert recipe declared next to an MXFP8 recipe.
+
+        The compressed-tensors NVFP4 MoE packaging of the MiniMax M3
+        architecture keeps ``quant_method: mxfp8`` for the base layers and
+        declares the routed experts beside it (``moe_quant_algo: NVFP4``,
+        ``moe_quant_format: nvfp4-pack-quantized``, ``group_size: 16``).
+        Returns None when no separate expert recipe is declared; raises for a
+        declared combination TensorRT-LLM cannot load.
+        """
+        moe_quant_algo = hf_quant_config.get("moe_quant_algo")
+        moe_quant_format = hf_quant_config.get("moe_quant_format")
+        if moe_quant_algo is None and moe_quant_format is None:
+            return None
+        group_size = hf_quant_config.get("group_size")
+        if (moe_quant_algo, moe_quant_format,
+                group_size) != ("NVFP4", "nvfp4-pack-quantized",
+                                NVFP4_SF_VEC_SIZE):
+            raise ValueError(
+                "Unsupported routed-expert packaging in an MXFP8 checkpoint: "
+                f"moe_quant_algo={moe_quant_algo!r}, moe_quant_format="
+                f"{moe_quant_format!r}, group_size={group_size!r}. Supported: "
+                "moe_quant_algo='NVFP4', moe_quant_format="
+                f"'nvfp4-pack-quantized', group_size={NVFP4_SF_VEC_SIZE}.")
+        return QuantConfig(quant_algo=QuantAlgo.NVFP4,
+                           group_size=NVFP4_SF_VEC_SIZE)
 
-        Two fix-ups are applied:
+    @staticmethod
+    def _build_minimax_m3_mxfp8_moe_quant_config(
+        pretrained_config: Any, quant_config: QuantConfig
+    ) -> Tuple[QuantConfig, Optional[Dict[str, QuantConfig]]]:
+        """Expand MXFP8 base + NVFP4 routed experts into MIXED_PRECISION.
+
+        Used for the compressed-tensors NVFP4 MoE packaging, which has no
+        per-layer map. Every quantized text linear gets an explicit MXFP8
+        entry, because an unmatched module would inherit the global
+        MIXED_PRECISION config and load unquantized; the routed experts of
+        each MoE layer get the coarse NVFP4 entry
+        ``MiniMaxM3MoE._get_experts_quant_config`` reads. Modules listed in
+        ``ignored_layers`` get no entry, and those names lose their
+        ``language_model.`` prefix so they match the TensorRT-LLM module tree.
+
+        Returns ``(quant_config, layer_quant_config)``; a checkpoint without the
+        expert recipe is returned unchanged with no per-layer map.
+        """
+        from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.common import \
+            get_sparse_layer_ids
+        hf_quant_config = getattr(pretrained_config, "quantization_config",
+                                  None)
+        if (quant_config is None or quant_config.quant_algo != QuantAlgo.MXFP8
+                or not isinstance(hf_quant_config, dict)):
+            return quant_config, None
+        experts_quant_config = ModelConfig._get_mxfp8_moe_quant_config(
+            hf_quant_config)
+        if experts_quant_config is None:
+            return quant_config, None
+
+        text_config = getattr(pretrained_config, "text_config",
+                              None) or pretrained_config
+        if isinstance(text_config, dict):
+            text_config = SimpleNamespace(**text_config)
+        num_layers = int(text_config.num_hidden_layers)
+        moe_layer_freq = getattr(text_config, "moe_layer_freq", None)
+        if moe_layer_freq is None:
+            moe_layer_freq = [1] * num_layers
+        if len(moe_layer_freq) != num_layers:
+            raise ValueError(
+                f"moe_layer_freq length {len(moe_layer_freq)} does not match "
+                f"num_hidden_layers {num_layers}")
+        _, sparse_layer_ids = get_sparse_layer_ids(text_config)
+        has_shared_experts = int(
+            getattr(text_config, "n_shared_experts", 0) or 0) > 0
+
+        lm_prefix = "language_model."
+        exclude_modules = None
+        if quant_config.exclude_modules is not None:
+            exclude_modules = list(
+                dict.fromkeys(
+                    m[len(lm_prefix):] if m.startswith(lm_prefix) else m
+                    for m in quant_config.exclude_modules))
+        mixed_quant_config = QuantConfig(
+            quant_algo=QuantAlgo.MIXED_PRECISION,
+            kv_cache_quant_algo=quant_config.kv_cache_quant_algo,
+            group_size=quant_config.group_size,
+            exclude_modules=exclude_modules)
+        base_quant_config = QuantConfig(
+            quant_algo=QuantAlgo.MXFP8,
+            kv_cache_quant_algo=quant_config.kv_cache_quant_algo,
+            group_size=quant_config.group_size)
+        experts_quant_config.kv_cache_quant_algo = quant_config.kv_cache_quant_algo
+        excluded = mixed_quant_config.is_module_excluded_from_quantization
+        # The per-layer map covers only the decoder layers, so an untied head
+        # must be excluded: a quantized head would load without its scales.
+        if not getattr(text_config, "tie_word_embeddings",
+                       False) and not excluded("lm_head"):
+            raise ValueError(
+                "The MiniMax M3 NVFP4 routed-expert packaging must exclude "
+                "lm_head from quantization (list it in ignored_layers).")
+
+        layer_quant_config = {}
+        for layer_idx in range(num_layers):
+            prefix = f"model.layers.{layer_idx}"
+            entries = {
+                f"{prefix}.self_attn.{proj}": base_quant_config
+                for proj in ("q_proj", "k_proj", "v_proj", "o_proj")
+            }
+            if layer_idx in sparse_layer_ids:
+                entries.update({
+                    f"{prefix}.self_attn.{proj}": base_quant_config
+                    for proj in ("index_q_proj", "index_k_proj")
+                })
+            if int(moe_layer_freq[layer_idx]) != 0:
+                if has_shared_experts:
+                    entries.update({
+                        f"{prefix}.block_sparse_moe.shared_experts.{proj}":
+                        base_quant_config
+                        for proj in ("gate_proj", "up_proj", "down_proj")
+                    })
+                entries[
+                    f"{prefix}.block_sparse_moe.experts"] = experts_quant_config
+            else:
+                entries.update({
+                    f"{prefix}.mlp.{proj}": base_quant_config
+                    for proj in ("gate_proj", "up_proj", "down_proj")
+                })
+            # Fused projections load as one module, so their members must
+            # agree on whether they are quantized.
+            ModelConfig._check_fused_exclusions([
+                f"{prefix}.self_attn.{proj}"
+                for proj in ("q_proj", "k_proj", "v_proj")
+            ], excluded)
+            for parent in ("block_sparse_moe.shared_experts", "mlp"):
+                names = [
+                    f"{prefix}.{parent}.{proj}"
+                    for proj in ("gate_proj", "up_proj")
+                ]
+                if names[0] in entries:
+                    ModelConfig._check_fused_exclusions(names, excluded)
+            layer_quant_config.update({
+                name: config
+                for name, config in entries.items() if not excluded(name)
+            })
+
+        logger.info(
+            "Detected the compressed-tensors NVFP4 MoE packaging of MiniMax M3: "
+            "MXFP8 base linears with NVFP4 routed experts "
+            f"({len(layer_quant_config)} per-layer quant entries).")
+        return mixed_quant_config, layer_quant_config
+
+    @staticmethod
+    def _check_fused_exclusions(names: List[str],
+                                excluded: Callable[[str], bool]) -> None:
+        """Reject an ignore list that excludes only part of a fused module."""
+        flags = [excluded(name) for name in names]
+        if any(flags) and not all(flags):
+            raise ValueError(
+                "The MiniMax M3 NVFP4 routed-expert packaging must exclude all "
+                f"or none of the fused projections {names}.")
+
+    @staticmethod
+    def _set_minimax_m3_layer_quant_config(
+        pretrained_config: Any, quant_config: QuantConfig,
+        layer_quant_config: Optional[Dict[str, QuantConfig]]
+    ) -> Tuple[QuantConfig, Optional[Dict[str, QuantConfig]]]:
+        """Normalize the Minimax M3 MIXED_PRECISION quant configs.
+
+        Returns ``(quant_config, layer_quant_config)``. A checkpoint without a
+        per-layer map goes through
+        ``_build_minimax_m3_mxfp8_moe_quant_config``, which builds one for the
+        compressed-tensors NVFP4 MoE packaging and leaves uniform MXFP8 or BF16
+        checkpoints unchanged.
+
+        For a per-layer map (ModelOpt MIXED_PRECISION), two fix-ups are applied:
 
         1. Strip the ``language_model.`` prefix from every per-layer key.
            The M3 VL checkpoint stores keys like
@@ -1018,13 +1188,11 @@ def _set_minimax_m3_layer_quant_config(pretrained_config,
            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
+            return ModelConfig._build_minimax_m3_mxfp8_moe_quant_config(
+                pretrained_config, quant_config)
 
         # (1) Strip the ``language_model.`` prefix so exact-match per-layer
         # quant assignment works for standalone base Linears.
@@ -1041,7 +1209,7 @@ def _set_minimax_m3_layer_quant_config(pretrained_config,
             and v.quant_algo == QuantAlgo.NVFP4
             for k, v in layer_quant_config.items())
         if not has_nvfp4_experts:
-            return layer_quant_config
+            return quant_config, layer_quant_config
 
         experts_quant_config = QuantConfig()
         experts_quant_config.quant_algo = QuantAlgo.NVFP4
@@ -1066,7 +1234,7 @@ def _set_minimax_m3_layer_quant_config(pretrained_config,
         logger.info(
             "Detected Minimax M3 NVFP4 routed MoE checkpoint; using NVFP4 "
             "for routed experts.")
-        return layer_quant_config
+        return quant_config, layer_quant_config
 
     @staticmethod
     def load_quant_config_from_dtypes_json(dtypes_json_file, moe_backend: str):
@@ -1561,9 +1729,8 @@ def _recursive_update_config(config: transformers.PretrainedConfig,
                 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)
-
+            quant_config, layer_quant_config = cls._set_minimax_m3_layer_quant_config(
+                pretrained_config, quant_config, layer_quant_config)
         model_config = cls(pretrained_config=pretrained_config,
                            quant_config=quant_config,
                            quant_config_dict=layer_quant_config,