From 83d2117fa3fee54cdeca380459da8a48a45e8cca Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:57:54 -0700 Subject: [PATCH 1/5] [None][feat] MiniMax-M3: load the compressed-tensors NVFP4 MoE packaging MXFP8 non-expert projections with NVFP4 routed experts on the existing model: model_config expands the packaging's expert recipe into the mixed per-layer config, the weight mapper renames the packed expert tensors and inverts their global scales, MXFP8 index projections load their declared quantization, and the sparse cache manager takes the layer schedule from the checkpoint and counts index-K bytes. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- ...ployment-guide-for-minimax-m3-on-trtllm.md | 20 + .../sparse/minimax_m3/cache_manager.py | 128 +++- .../backends/sparse/minimax_m3/common.py | 39 +- tensorrt_llm/_torch/model_config.py | 158 +++- .../checkpoints/hf/minimaxm3_weight_mapper.py | 96 ++- .../_torch/models/modeling_minimaxm3.py | 90 ++- .../test_lists/test-db/l0_b200.yml | 4 + .../integration/test_lists/test-db/l0_cpu.yml | 11 + .../test_minimax_m3_shared_draft_layers.py | 107 +++ .../hf/test_minimaxm3_weight_mapper.py | 213 ++++++ .../unittest/_torch/models/test_minimax_m3.py | 703 +++++++++++++++++- 11 files changed, 1460 insertions(+), 109 deletions(-) diff --git a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md index ef58d5ed3b83..fa661a326da4 100644 --- a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md @@ -56,8 +56,28 @@ 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 is zero, negative or not finite, or an expert whose `w1` +and `w3` weight global scales differ, fails the load with an error naming the +checkpoint tensors. 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 diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 37d9564664a2..b14b9af7b900 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -53,6 +53,8 @@ PageIndexMode, ) +from .common import get_sparse_disable_index_value_layer_ids, get_sparse_layer_ids + if TYPE_CHECKING: from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig @@ -159,6 +161,47 @@ def shared_draft_layer_count(spec_config, layer_mask) -> int: return int(get_num_spec_layers(spec_config)) +def resolve_sparse_layer_schedule( + sparse_layer_ids: Optional[Sequence[int]], + disable_index_value_layer_ids: Optional[Sequence[int]], + pretrained_config, + num_target_layers: Optional[int], +) -> Tuple[List[int], List[int]]: + """Return ``(sparse_layer_ids, disable_index_value_layer_ids)`` for the cache. + + Explicit ids win. Otherwise the checkpoint's attention schedule + (``pretrained_config.sparse_attention_config``) decides, which is the + schedule the model builds its sparse layers from; it may make every layer + sparse, independently of the dense-MLP layers. Without either, the public + M3 convention applies: layers 0..2 dense, 3..N-1 sparse and K-only. + """ + if ( + sparse_layer_ids is None + and getattr(pretrained_config, "sparse_attention_config", None) is not None + ): + _, sparse_layer_ids = get_sparse_layer_ids(pretrained_config) + if disable_index_value_layer_ids is None: + disable_index_value_layer_ids = get_sparse_disable_index_value_layer_ids( + pretrained_config + ) + if num_target_layers is not None and any( + not 0 <= int(i) < num_target_layers + for i in (*sparse_layer_ids, *disable_index_value_layer_ids) + ): + raise ValueError( + "MiniMax M3 checkpoint attention schedule names layers outside the " + f"{num_target_layers} target layers: sparse={sparse_layer_ids}, " + f"disable_index_value={disable_index_value_layer_ids}." + ) + if sparse_layer_ids is None: + sparse_layer_ids = ( + list(range(3, num_target_layers)) if num_target_layers is not None else [] + ) + if disable_index_value_layer_ids is None: + disable_index_value_layer_ids = list(sparse_layer_ids) + return list(sparse_layer_ids), list(disable_index_value_layer_ids) + + def derive_shared_draft_layout( num_layers: Optional[int], num_kv_heads, @@ -280,13 +323,16 @@ def __init__( **kwargs, ): # Resolve M3 sparse-layer metadata from explicit kwargs first, then - # from the executor's ``sparse_attention_config`` keyword, then from - # the M3 checkpoint convention (layers 0..2 dense, 3..N-1 sparse, + # from the executor's ``sparse_attention_config`` keyword and the + # checkpoint's attention schedule (``pretrained_config``, the same + # schedule the model builds its sparse layers from), then from the + # public M3 convention (layers 0..2 dense, 3..N-1 sparse, # disable_index_value=True, sparse_index_dim=128). Honoring the # executor keyword also makes non-default sparse_index_dim values # authoritative for the cache layout instead of falling back to 128. # Peeked (not popped) so the base __init__ still receives them. sparse_attention_config = kwargs.get("sparse_attention_config") + pretrained_config = kwargs.get("pretrained_config") num_layers = kwargs.get("num_layers") implementation = getattr(sparse_attention_config, "implementation", "triton") self._main_kv_view_layout = "HND" if implementation == "msa" else "NHD" @@ -310,13 +356,9 @@ def __init__( kwargs.get("num_kv_heads"), shared_draft_layer_count(kwargs.get("spec_config"), kwargs.get("layer_mask")), ) - if sparse_layer_ids is None: - if num_target_layers is not None: - sparse_layer_ids = list(range(3, num_target_layers)) - else: - sparse_layer_ids = [] - if disable_index_value_layer_ids is None: - disable_index_value_layer_ids = list(sparse_layer_ids) + sparse_layer_ids, disable_index_value_layer_ids = resolve_sparse_layer_schedule( + sparse_layer_ids, disable_index_value_layer_ids, pretrained_config, num_target_layers + ) # Must be set BEFORE super().__init__ — the base # ``_build_base_config`` invokes ``_extra_buffers_per_layer`` @@ -437,25 +479,36 @@ def _build_cache_config(self, config): def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): """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: + if not has_index_k: + return 0 + index_dtype = self._torch_dtype_for_index_cache() + index_elem_bytes = torch.empty(0, dtype=index_dtype, device="cpu").element_size() + return self.sparse_index_dim * index_elem_bytes + + 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.""" @@ -621,20 +674,17 @@ def update_resources(self, scheduled_batch, attn_metadata=None, kv_cache_dtype_b super().update_resources(scheduled_batch, attn_metadata, kv_cache_dtype_byte_size) def _extra_buffers_per_layer(self, *, tokens_per_block): - """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 - 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 + """Register the paged index-K buffer for each sparse local layer.""" return { - self.layer_offsets[layer_id]: [BufferConfig(role=Role.INDEX_KEY, size=size_per_block)] + self.layer_offsets[layer_id]: [ + BufferConfig( + role=Role.INDEX_KEY, + size=self.get_layer_bytes_per_token( + self.layer_offsets[layer_id], Role.INDEX_KEY + ) + * tokens_per_block, + ) + ] for layer_id in self.sparse_layer_ids if layer_id in self.layer_offsets } diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py index f288495eeff1..93df670a4137 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py @@ -11,7 +11,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Tuple +from typing import TYPE_CHECKING, Any, List, Literal, NamedTuple, Optional, Tuple import torch @@ -32,6 +32,43 @@ _LOCAL_SCORE = 1e29 +def get_sparse_layer_ids(text_config: Any) -> Tuple[List[int], List[int]]: + """Return ``(dense_layer_ids, sparse_layer_ids)`` for the M3 text model. + + The checkpoint's ``sparse_attention_config.sparse_attention_freq`` is the + attention schedule; it is independent of the dense-MLP/MoE schedule. + Shared by the model, which builds sparse attention on these layers, and + the KV-cache manager, which allocates their index-K cache. + """ + sparse_cfg = getattr(text_config, "sparse_attention_config", None) + n_layers = int(text_config.num_hidden_layers) + if sparse_cfg is None: + return list(range(n_layers)), [] + if not sparse_cfg.get("use_sparse_attention", True): + return list(range(n_layers)), [] + freq = sparse_cfg.get("sparse_attention_freq") + if freq is None: + return list(range(n_layers)), [] + if len(freq) != n_layers: + raise ValueError( + f"sparse_attention_freq length {len(freq)} does not match num_hidden_layers {n_layers}" + ) + dense = [i for i, f in enumerate(freq) if int(f) == 0] + sparse = [i for i, f in enumerate(freq) if int(f) != 0] + return dense, sparse + + +def get_sparse_disable_index_value_layer_ids(text_config: Any) -> List[int]: + """Return layer ids where the sparse index-value branch is disabled.""" + sparse_cfg = getattr(text_config, "sparse_attention_config", None) + if sparse_cfg is None: + return [] + flags = sparse_cfg.get("sparse_disable_index_value") + if flags is None: + return [] + return [i for i, f in enumerate(flags) if int(f) != 0] + + def index_head_range( num_index_heads: int, num_kv_heads: int, mapping: Optional["Mapping"] = None ) -> Tuple[int, int]: diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 84edf64cfd06..32af9e56c144 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -21,6 +21,7 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path +from types import SimpleNamespace from typing import (TYPE_CHECKING, Any, Dict, Generic, Iterator, List, Optional, Sequence, Tuple, TypeVar) @@ -974,11 +975,147 @@ def _set_deepseek_v4_routed_moe_quant_config(pretrained_config, return layer_quant_config @staticmethod - def _set_minimax_m3_layer_quant_config(pretrained_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) + + @staticmethod + def _build_minimax_m3_mxfp8_moe_quant_config(pretrained_config, + quant_config): + """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 + + 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") + }) + excluded = mixed_quant_config.is_module_excluded_from_quantization + 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 _set_minimax_m3_layer_quant_config(pretrained_config, quant_config, layer_quant_config): - """Normalize the Minimax M3 MIXED_PRECISION per-layer quant config. + """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. - Two fix-ups are applied: + 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 @@ -998,13 +1135,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. @@ -1021,7 +1156,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 @@ -1046,7 +1181,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): @@ -1540,9 +1675,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, diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py index bb6d58957fd3..261559f186d4 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py @@ -13,8 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import re +from collections.abc import Mapping + +import torch from torch import Tensor, nn +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper @@ -22,16 +27,105 @@ r"^(.*\.block_sparse_moe)\.e_score_correction_bias$": r"\1.gate.e_score_correction_bias", } +# Routed-expert tensors of the compressed-tensors NVFP4 MoE packaging of the +# MiniMax M3 architecture. Per projection: ``weight_packed`` (uint8, two E2M1 +# values per byte), ``weight_scale`` (E4M3 per 16 elements) and the FP32 +# scalars ``weight_global_scale`` / ``input_global_scale``, stored as the +# multiplier 2688 / amax (2688 = E4M3 max 448 * E2M1 max 6). +_PACKED_NVFP4_EXPERT_KEY = re.compile( + r"^(?P.*\.block_sparse_moe\.experts\.\d+\.w[123])\." + r"(?Pweight_packed|weight_global_scale|input_global_scale)$" +) +_NVFP4_GLOBAL_SCALE_RANGE = 448.0 * 6.0 + + +def nvfp4_weight_scale_2(weight_global_scale: Tensor) -> Tensor: + """Return the NVFP4 MoE ``weight_scale_2`` of a stored global scale. + + compressed-tensors stores ``weight_global_scale = 2688 / amax`` (E4M3 max + 448 times E2M1 max 6); the dequantization scale is its reciprocal, + ``amax / 2688``, computed in FP32. + """ + return torch.reciprocal(weight_global_scale.float()) + + +def _read_global_scale(name: str, value) -> Tensor: + """Materialize one global scale, rejecting values that would disable an expert.""" + scale = (value if isinstance(value, Tensor) else value[...]).to(torch.float32) + if scale.numel() != 1 or not bool(torch.isfinite(scale).all()) or not bool((scale > 0).all()): + raise ValueError( + f"Checkpoint tensor {name} must hold one finite positive global scale, " + f"got {scale.flatten()[:4].tolist()}." + ) + return scale + + +def translate_packed_nvfp4_experts(weights: Mapping[str, Tensor]) -> Mapping[str, Tensor]: + """Rename packed NVFP4 routed-expert tensors to the NVFP4 MoE contract. + + ``weight_packed`` becomes ``weight`` with the same bytes, and the global + scales become the ``weight_scale_2`` / ``input_scale`` divisors + (amax / 2688) the NVFP4 MoE loader reads; ``weight_scale`` keeps its name + and bytes. Only the scalar scales are materialized. Every other key passes + through, including already canonical ``weight_scale_2`` / ``input_scale``, + so a second call is a no-op. + + Raises: + ValueError: naming the checkpoint tensor, when a global scale is not + finite and positive, or when an expert's gate (w1) and up (w3) + weight scales differ: the fused gate/up GEMM applies one weight + scale per expert. + """ + matches = {} + for name in weights.keys(): + match = _PACKED_NVFP4_EXPERT_KEY.match(name) + if match is not None: + matches[name] = match + if not matches: + return weights + + renamed = {} + gate_up_scales: dict[str, dict[str, tuple[str, Tensor]]] = {} + for name in weights.keys(): + value = weights[name] + match = matches.get(name) + if match is None: + renamed[name] = value + continue + proj, kind = match["proj"], match["kind"] + if kind == "weight_packed": + renamed[f"{proj}.weight"] = value + elif kind == "weight_global_scale": + scale_2 = nvfp4_weight_scale_2(_read_global_scale(name, value)) + renamed[f"{proj}.weight_scale_2"] = scale_2 + expert, w = proj.rsplit(".", 1) + if w != "w2": + gate_up_scales.setdefault(expert, {})[w] = (name, scale_2) + else: + renamed[f"{proj}.input_scale"] = torch.reciprocal(_read_global_scale(name, value)) + + for scales in gate_up_scales.values(): + if len(scales) == 2 and not torch.equal(scales["w1"][1], scales["w3"][1]): + raise ValueError( + f"Checkpoint tensors {scales['w1'][0]} and {scales['w3'][0]} decode to " + "different weight scales; the fused gate/up expert GEMM needs one." + ) + return ConsumableWeightsDict.take_ownership(weights, renamed) + @register_mapper("HF", "MiniMaxM3SparseForCausalLM") @register_mapper("HF", "MiniMaxM3SparseForConditionalGeneration") class MiniMaxM3HfWeightMapper(HfWeightMapper): - """Handle M3 gate naming and MXFP8 GQA duplication for loader v2.""" + """Handle M3 gate naming, MXFP8 GQA duplication and packed NVFP4 experts.""" def __init__(self) -> None: super().__init__() self.params_map = MINIMAX_M3_PARAMS_MAP + def preprocess_weights(self, weights: Mapping[str, Tensor]) -> Mapping[str, Tensor]: + """See :func:`translate_packed_nvfp4_experts`.""" + return translate_packed_nvfp4_experts(weights) + def _duplicate_kv_weights( self, module: nn.Module, new_name: str, weights: dict[str, Tensor] ) -> dict[str, Tensor]: diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 4359a6e4549b..b1c651e6a98b 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -54,7 +54,11 @@ _gather_paged_batched, _write_main_kv_slots_to_pool, ) -from ..attention.backends.sparse.minimax_m3.common import index_head_range +from ..attention.backends.sparse.minimax_m3.common import ( + get_sparse_disable_index_value_layer_ids, + get_sparse_layer_ids, + index_head_range, +) from ..attention.backends.sparse.params import SparseBackendForwardArgs from ..attention.backends.trtllm import TrtllmAttention from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams, MiniMaxAllReduceRMS @@ -76,7 +80,11 @@ from ..speculative import SpecMetadata from ..utils import AuxStreamType, EventType, get_model_extra_attrs, is_torch_compiling from .checkpoints.base_weight_mapper import BaseWeightMapper -from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper +from .checkpoints.hf.minimaxm3_weight_mapper import ( + MINIMAX_M3_PARAMS_MAP, + MiniMaxM3HfWeightMapper, + translate_packed_nvfp4_experts, +) from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, ModelConfig, filter_weights, register_auto_model @@ -431,39 +439,6 @@ def _validate_sparse_attention_runtime_config( ) -def get_sparse_layer_ids(text_config: PretrainedConfig) -> Tuple[List[int], List[int]]: - """Return ``(dense_layer_ids, sparse_layer_ids)`` for the M3 text model.""" - sparse_cfg = getattr(text_config, "sparse_attention_config", None) - n_layers = int(text_config.num_hidden_layers) - if sparse_cfg is None: - return list(range(n_layers)), [] - if not sparse_cfg.get("use_sparse_attention", True): - return list(range(n_layers)), [] - freq = sparse_cfg.get("sparse_attention_freq") - if freq is None: - return list(range(n_layers)), [] - if len(freq) != n_layers: - raise ValueError( - f"sparse_attention_freq length {len(freq)} does not match num_hidden_layers {n_layers}" - ) - dense = [i for i, f in enumerate(freq) if int(f) == 0] - sparse = [i for i, f in enumerate(freq) if int(f) != 0] - return dense, sparse - - -def get_sparse_disable_index_value_layer_ids( - text_config: PretrainedConfig, -) -> List[int]: - """Return layer ids where the sparse index-value branch is disabled.""" - sparse_cfg = getattr(text_config, "sparse_attention_config", None) - if sparse_cfg is None: - return [] - flags = sparse_cfg.get("sparse_disable_index_value") - if flags is None: - return [] - return [i for i, f in enumerate(flags) if int(f) != 0] - - def get_moe_layer_ids(text_config: PretrainedConfig) -> Tuple[List[int], List[int]]: """Return ``(dense_mlp_layer_ids, moe_layer_ids)`` for the M3 text model.""" freq = getattr(text_config, "moe_layer_freq", None) @@ -479,6 +454,36 @@ def get_moe_layer_ids(text_config: PretrainedConfig) -> Tuple[List[int], List[in return dense, moe +def _get_index_projection_quant_config( + model_config: "ModelConfig[PretrainedConfig]", layer_idx: Optional[int] +) -> Optional[QuantConfig]: + """Return the quant config the checkpoint declares for the index projections. + + ``index_q_proj`` and ``index_k_proj`` load into one fused ``index_qk_proj`` + that ``apply_layerwise_quant_config`` cannot match by name, so their + per-layer entries are read here. Without such entries the projection + stays unquantized, as before. + """ + quant_config_dict = getattr(model_config, "quant_config_dict", None) + if quant_config_dict is None or layer_idx is None: + return None + prefix = f"model.layers.{layer_idx}.self_attn" + q_config = quant_config_dict.get(f"{prefix}.index_q_proj") + k_config = quant_config_dict.get(f"{prefix}.index_k_proj") + if q_config is None and k_config is None: + return None + if ( + q_config is None + or k_config is None + or (q_config.quant_algo, q_config.group_size) != (k_config.quant_algo, k_config.group_size) + ): + raise ValueError( + f"{prefix}.index_q_proj and {prefix}.index_k_proj load into one fused " + f"projection and need the same quantization, got {q_config} and {k_config}." + ) + return q_config + + # Outside ``language_model.*``: deferred multimodal branches. _VL_PREFIXES = ("vision_tower.", "multi_modal_projector.", "patch_merge_mlp.") @@ -1251,7 +1256,9 @@ def __init__( dtype=config.torch_dtype, mapping=self.qkv_proj.mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, - quant_config=None, + # Quantized when the checkpoint declares per-layer entries + # for the index projections; None keeps them unquantized. + quant_config=_get_index_projection_quant_config(model_config, layer_idx), weights_loading_config=WeightsLoadingConfig( weight_mode=WeightMode.FUSED_GATE_UP_LINEAR ), @@ -2975,6 +2982,9 @@ def load_weights( params_map: Optional[Dict[str, str]] = None, allow_partial_loading: bool = False, ) -> None: + # Rename packed NVFP4 routed-expert tensors to the NVFP4 MoE contract + # (a no-op for other packagings and for already rewritten weights). + weights = translate_packed_nvfp4_experts(weights) # Pack the opt-in five-way projection with MiniMax-specific TP # sharding before the generic mapper handles the compatibility path. packed_projection_modules = _load_qkv_index_proj_weights(self, weights) @@ -3175,6 +3185,8 @@ def load_weights( # set of unknown top-level keys it surfaces as errors. from .modeling_minimaxm3_vl import load_minimax_m3_vl_state_dict, split_multimodal_weights + # Rewrite before the prefix strip so errors name the checkpoint tensor. + weights = translate_packed_nvfp4_experts(weights) text_or_other_weights, vision_weights = split_multimodal_weights(weights) text_weights, ignored = _strip_language_model_prefix(text_or_other_weights) @@ -3186,6 +3198,12 @@ def load_weights( ) self.last_loaded_vision_keys = loaded self.last_missing_vision_keys = missing + logger.info( + f"MiniMax-M3 VL weights: {len(self.last_loaded_vision_keys)} vision tensors " + f"loaded, {len(self.last_missing_vision_keys)} vision parameters missing, " + f"{len(ignored)} checkpoint tensors skipped " + f"({sorted(set(ignored.values()))})." + ) super().load_weights( weights=text_weights, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index f6f854d4296c..3c99beabc76f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -45,6 +45,10 @@ l0_b200: - unittest/_torch/models/test_minimax_m3.py::test_nvfp4_reload_preserves_attention_scale_storage[cuda] - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_gate_stays_fp32_when_tf32_is_allowed - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_gate_bf16_split_matches_fp32_router + # MiniMax-M3 compressed-tensors NVFP4 MoE packaging, one GPU: expert loading and + # eager/graph numerics on CUTLASS; MXFP8 index-projection scales per simulated TP4 rank. + - unittest/_torch/models/test_minimax_m3.py::test_packaged_nvfp4_experts_load_and_run_on_cutlass + - unittest/_torch/models/test_minimax_m3.py::test_packaged_index_projection_keeps_mxfp8_scales_per_tp4_rank - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_transfer_stress - unittest/others/test_lora_manager.py - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index e62770617cc3..9a3d5f50f0ad 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -101,6 +101,17 @@ l0_cpu: - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_projection_shard_geometry - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_loader_returns_exact_generic_skip - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_split_router_weight_reconstructs_fp32 + # MiniMax-M3 compressed-tensors NVFP4 MoE packaging: config expansion, + # index-projection quantization and the checkpoint attention schedule. + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fold_gemma_boundary_weights_accepts_partial_lazy_load + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_index_projection_quantization_follows_checkpoint_entries + - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_expands_to_mixed_precision + - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_honors_ignored_layers + - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_rejects_unsupported_expert_recipes + - unittest/_torch/models/test_minimax_m3.py::test_uniform_mxfp8_and_bf16_configs_are_unchanged + - unittest/_torch/models/test_minimax_m3.py::test_modelopt_mixed_precision_map_is_normalized_as_before + - unittest/_torch/models/test_minimax_m3.py::test_index_projection_quant_config_follows_declared_entries + - unittest/_torch/models/test_minimax_m3.py::test_kv_cache_schedule_follows_the_checkpoint_attention_schedule - unittest/_torch/models/checkpoints - unittest/_torch/modules - unittest/_torch/moe diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index e9abf2039073..9ba890ab50a0 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -525,3 +525,110 @@ def test_sparse_view_writes_native_hnd_storage_without_copy(monkeypatch): assert index_nhd.data_ptr() == index_hnd.data_ptr() torch.testing.assert_close(manager.get_index_k_buffer(0, "HND"), index_hnd) torch.testing.assert_close(index_hnd, before[:, 2, :1]) + + +# Bytes per token for local layers 0/1/2 (dense global 11, sparse global 27, +# shared Eagle global 60), each with 4 KV heads x head_dim 128 (K + V), plus a +# 128-wide index-K on the sparse layer: 256 B in BF16, 128 B in FP8. +# BF16 K/V is 2048 B per layer and FP8 K/V is 1024 B. With NVFP4 only the +# sparse layer is NVFP4 (512 B data + 64 B FP8 block scales); the dense and +# Eagle layers keep FP8 K/V (1024 B), the Eagle layer on P32 physical pages. +_M3_LAYER_BYTES_PER_TOKEN = { + (DataType.BF16, "bf16"): [2048, 2048 + 256, 2048], + (DataType.BF16, "fp8"): [2048, 2048 + 128, 2048], + (DataType.FP8, "bf16"): [1024, 1024 + 256, 1024], + (DataType.FP8, "fp8"): [1024, 1024 + 128, 1024], + (DataType.NVFP4, "bf16"): [1024, 576 + 256, 1024], + (DataType.NVFP4, "fp8"): [1024, 576 + 128, 1024], +} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("main_dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) +@pytest.mark.parametrize("index_dtype", ["bf16", "fp8"]) +def test_m3_budget_estimate_covers_actual_paged_buffers(main_dtype, index_dtype): + manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + manager.dtype = main_dtype + manager.kv_cache_type = CacheType.SELF + manager.kv_factor = 2 + manager.num_kv_heads_per_layer = [4, 4, 4] + manager.head_dim_per_layer = [128, 128, 128] + manager.num_local_layers = 3 + # Nonidentity PP mapping: local layers 0/1/2 are global layers 11/27/60. + manager.pp_layers = [11, 27, 60] + manager.layer_offsets = {11: 0, 27: 1, 60: 2} + manager.sparse_layer_ids = {27} + manager.disable_index_value_layer_ids = {27} + manager._shared_draft_layer_ids = [60] + manager.sparse_index_dim = 128 + manager.indexer_kv_dtype = index_dtype + manager.tokens_per_block = 128 + manager._ledger_tokens_per_block = 128 + manager._generation_kv_capacity_headroom = 1 + manager.max_num_tokens = 8192 + manager.max_seq_len = 8192 + manager.max_batch_size = 1 + manager.max_beam_width = 1 + manager.enable_partial_reuse = False + manager.max_draft_len = 0 + manager.max_attention_window_vec = [None, None, None] + manager.max_cuda_graph_batch_size = None + manager.num_extra_kv_tokens = 0 + manager.enable_swa_scratch_reuse = False + manager.enable_block_reuse = False + manager.block_reuse_policy = kv_cache_manager_v2.BlockReusePolicy.ALL_REUSABLE + manager.reuse_match_backoff = 0 + manager.enable_stats = False + manager._has_cp_helix = False + + config = manager._build_base_config( + KvCacheConfig( + pool_ratio=[1.0], + enable_block_reuse=False, + max_util_for_resume=0.9, + enable_partial_reuse=False, + ), + tokens_per_block=manager.tokens_per_block, + cache_tiers=[kv_cache_manager_v2.GpuCacheTierConfig(quota=1 << 20)], + ) + config = manager._build_cache_config(config) + config = manager._remove_zero_size_buffers(config) + + # Independently total final allocator descriptors. Eagle P32 buffers use + # their physical-page override; all other buffers use the ledger P128. + actual_bytes_per_token = [] + for layer in config.layers: + layer_bytes = 0 + for buffer in layer.buffers: + tokens_per_page = buffer.tokens_per_block_override or manager.tokens_per_block + assert buffer.size % tokens_per_page == 0 + layer_bytes += buffer.size // tokens_per_page + actual_bytes_per_token.append(layer_bytes) + expected_layer_bytes = _M3_LAYER_BYTES_PER_TOKEN[(main_dtype, index_dtype)] + assert actual_bytes_per_token == expected_layer_bytes + estimated_layer_bytes, _ = manager._get_runtime_cache_size_layer_components() + assert estimated_layer_bytes == expected_layer_bytes + assert [ + sum(buffer.role == Role.INDEX_KEY for buffer in layer.buffers) for layer in config.layers + ] == [0, 1, 0] + + expected_index_bytes = 128 * (2 if index_dtype == "bf16" else 1) + assert manager.get_layer_bytes_per_token(1, Role.INDEX_KEY) == expected_index_bytes + assert manager.get_layer_bytes_per_token(0, Role.INDEX_KEY) == 0 + assert manager.get_layer_bytes_per_token(2, Role.INDEX_KEY) == 0 + for role in (Role.KEY, Role.VALUE): + assert manager.get_layer_bytes_per_token(1, role) == ( + KVCacheManagerV2.get_layer_bytes_per_token(manager, 1, role) + ) + if main_dtype == DataType.NVFP4: + for role in (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE): + assert manager.get_layer_bytes_per_token(1, role) == ( + KVCacheManagerV2.get_layer_bytes_per_token(manager, 1, role) + ) + scale_roles = {Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE} + assert all(buffer.role not in scale_roles for buffer in config.layers[0].buffers) + assert all(buffer.role not in scale_roles for buffer in config.layers[2].buffers) + + configured_footprint = sum(actual_bytes_per_token) * manager.max_num_tokens + estimated_quota = manager._get_quota_from_max_tokens(manager.max_num_tokens) + assert estimated_quota >= configured_footprint diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py index 160bb485d5f8..d4ac9bee2469 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py @@ -14,6 +14,9 @@ # limitations under the License. import inspect +import math +import re +import struct from types import SimpleNamespace from unittest.mock import patch @@ -22,8 +25,11 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.auto_mapper import AutoCheckpointMapper +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( + MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper, + nvfp4_weight_scale_2, ) from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper from tensorrt_llm._torch.models.modeling_minimaxm3 import ( @@ -201,3 +207,210 @@ def test_load_weights_accepts_base_mapper_without_params_map() -> None: assert call_kwargs["weight_mapper"] is mapper renamed = mapper.rename_by_params_map(call_kwargs["params_map"], {source_name: bias}) assert renamed[target_name] is bias + + +# --------------------------------------------------------------------------- +# Compressed-tensors NVFP4 MoE packaging of the MiniMax M3 architecture +# --------------------------------------------------------------------------- + +# Checkpoint naming of one routed expert in the multimodal checkpoint layout. +_EXPERT = "language_model.model.layers.3.block_sparse_moe.experts.7" + +# weight_scale_2 = fp32(fp32(bf16(2688 / gs) * 2688) * fp32(1 / 2688**2)) as +# float32 bit patterns. Computed outside PyTorch: Python float arithmetic +# rounded to float32 after every operation (exact for these operations), BF16 +# by round-to-nearest-even on the float32 bits. +# FP32 reciprocals of the stored weight_global_scale values. +_WEIGHT_SCALE_2_GOLDENS = { + 2688.0: 0x39C30C31, + 3000.0: 0x39AEC33E, + 2913.75: 0x39B3EF92, + 16999.5: 0x3876BB24, + 10000.0: 0x38D1B717, +} +_RECIPROCAL_GOLDENS = {1.0: 0x3F800000, 0.5: 0x40000000, 3.0: 0x3EAAAAAB} + + +def _f32_bits(value) -> int: + return struct.unpack(" dict[str, torch.Tensor]: + """Synthetic tensors of one packed NVFP4 expert, shaped like the checkpoint's.""" + generator = torch.Generator().manual_seed(0) + weights = {} + # (rows, packed bytes per row) for hidden 64, intermediate 32: w1/w3 are + # [inter, hidden], w2 is [hidden, inter]; two E2M1 values per byte. + for proj, (rows, packed_cols) in {"w1": (32, 32), "w2": (64, 16), "w3": (32, 32)}.items(): + prefix = f"{_EXPERT}.{proj}" + weights[f"{prefix}.weight_packed"] = torch.randint( + 0, 256, (rows, packed_cols), dtype=torch.uint8, generator=generator + ) + weights[f"{prefix}.weight_scale"] = ( + torch.rand(rows, packed_cols * 2 // 16, generator=generator) * 4 + ).to(torch.float8_e4m3fn) + weights[f"{prefix}.weight_global_scale"] = torch.tensor( + down_global_scale if proj == "w2" else gate_up_global_scale + ) + weights[f"{prefix}.input_global_scale"] = torch.tensor(input_global_scale) + return weights + + +@pytest.mark.parametrize("global_scale", sorted(_WEIGHT_SCALE_2_GOLDENS)) +def test_nvfp4_weight_scale_2_matches_independent_goldens(global_scale: float) -> None: + scale_2 = nvfp4_weight_scale_2(torch.tensor(global_scale)) + assert scale_2.dtype == torch.float32 + assert _f32_bits(scale_2) == _WEIGHT_SCALE_2_GOLDENS[global_scale] + + +def test_packed_nvfp4_expert_translates_to_nvfp4_moe_contract() -> None: + source = _packed_expert() + unrelated = { + "language_model.model.layers.3.self_attn.q_proj.weight": torch.ones(4, 4), + "language_model.model.layers.3.block_sparse_moe.gate.weight": torch.ones(4, 4), + } + out = MiniMaxM3HfWeightMapper().preprocess_weights({**source, **unrelated}) + + for name, value in unrelated.items(): + assert out[name] is value + for proj in ("w1", "w2", "w3"): + prefix = f"{_EXPERT}.{proj}" + for stale in ("weight_packed", "weight_global_scale", "input_global_scale"): + assert f"{prefix}.{stale}" not in out + # Packed FP4 bytes and E4M3 group scales are passed through, not copied or cast. + assert out[f"{prefix}.weight"] is source[f"{prefix}.weight_packed"] + assert out[f"{prefix}.weight"].dtype == torch.uint8 + assert out[f"{prefix}.weight_scale"] is source[f"{prefix}.weight_scale"] + expected = _WEIGHT_SCALE_2_GOLDENS[2913.75 if proj == "w2" else 3000.0] + assert _f32_bits(out[f"{prefix}.weight_scale_2"]) == expected + assert out[f"{prefix}.weight_scale_2"].shape == () + assert _f32_bits(out[f"{prefix}.input_scale"]) == _RECIPROCAL_GOLDENS[1.0] + assert len(out) == len(source) + len(unrelated) + + +@pytest.mark.parametrize("input_global_scale", sorted(_RECIPROCAL_GOLDENS)) +def test_packed_nvfp4_input_scale_is_the_reciprocal(input_global_scale: float) -> None: + out = MiniMaxM3HfWeightMapper().preprocess_weights( + _packed_expert(input_global_scale=input_global_scale) + ) + for proj in ("w1", "w2", "w3"): + assert ( + _f32_bits(out[f"{_EXPERT}.{proj}.input_scale"]) + == _RECIPROCAL_GOLDENS[input_global_scale] + ) + + +@pytest.mark.parametrize("kind", ["weight_global_scale", "input_global_scale"]) +@pytest.mark.parametrize("bad", [0.0, -1.0, math.nan, math.inf, -math.inf]) +def test_packed_nvfp4_rejects_invalid_global_scale(kind: str, bad: float) -> None: + weights = _packed_expert() + name = f"{_EXPERT}.w2.{kind}" + weights[name] = torch.tensor(bad) + with pytest.raises(ValueError, match=re.escape(name)): + MiniMaxM3HfWeightMapper().preprocess_weights(weights) + + +def test_packed_nvfp4_rejects_different_gate_up_weight_scales() -> None: + weights = _packed_expert() + weights[f"{_EXPERT}.w3.weight_global_scale"] = torch.tensor(3100.0) + with pytest.raises(ValueError, match=re.escape(f"{_EXPERT}.w3.weight_global_scale")): + MiniMaxM3HfWeightMapper().preprocess_weights(weights) + + +def test_canonical_nvfp4_expert_tensors_are_not_inverted_again() -> None: + canonical = {} + for proj in ("w1", "w2", "w3"): + prefix = f"model.layers.3.block_sparse_moe.experts.0.{proj}" + canonical[f"{prefix}.weight"] = torch.zeros(8, 32, dtype=torch.uint8) + canonical[f"{prefix}.weight_scale"] = torch.ones(8, 4).to(torch.float8_e4m3fn) + canonical[f"{prefix}.weight_scale_2"] = torch.tensor(1.0 / 3000.0) + canonical[f"{prefix}.input_scale"] = torch.tensor(0.25) + out = MiniMaxM3HfWeightMapper().preprocess_weights(canonical) + assert out is canonical + + +def test_packed_nvfp4_translation_is_idempotent() -> None: + mapper = MiniMaxM3HfWeightMapper() + once = mapper.preprocess_weights(_packed_expert()) + assert mapper.preprocess_weights(once) is once + + +def test_packed_nvfp4_translation_takes_over_consumable_weights() -> None: + source = _packed_expert() + packed = source[f"{_EXPERT}.w1.weight_packed"] + consumable = ConsumableWeightsDict(dict(source)) + out = MiniMaxM3HfWeightMapper().preprocess_weights(consumable) + assert isinstance(out, ConsumableWeightsDict) + # The translated dict now owns the tensors, so consuming them frees memory. + assert len(consumable) == 0 + assert out[f"{_EXPERT}.w1.weight"] is packed + assert out.mark_consumed(_EXPERT) == len(source) + + +def _bare_model(model_class: type) -> torch.nn.Module: + config = SimpleNamespace(num_key_value_heads=_NUM_KV_HEADS, num_attention_heads=8) + model = object.__new__(model_class) + torch.nn.Module.__init__(model) + model.model_config = ModelConfig(pretrained_config=config, mapping=Mapping()) + model.model = torch.nn.Module() + model.model.layers = torch.nn.ModuleList() + model.spec_config = None + if model_class is MiniMaxM3VLForConditionalGeneration: + model.vision_tower = None + model.last_loaded_vision_keys = [] + model.last_missing_vision_keys = [] + return model + + +def test_vl_load_weights_translates_packed_experts_before_loading() -> None: + model = _bare_model(MiniMaxM3VLForConditionalGeneration) + with patch.object(DecoderModelForCausalLM, "load_weights") as base_load_weights: + model.load_weights(_packed_expert(), weight_mapper=MiniMaxM3HfWeightMapper()) + + loaded = base_load_weights.call_args.kwargs["weights"] + expert = _EXPERT[len("language_model.") :] + assert sorted(loaded.keys()) == sorted( + f"{expert}.{proj}.{kind}" + for proj in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + ) + assert _f32_bits(loaded[f"{expert}.w1.weight_scale_2"]) == _WEIGHT_SCALE_2_GOLDENS[3000.0] + assert _f32_bits(loaded[f"{expert}.w2.weight_scale_2"]) == _WEIGHT_SCALE_2_GOLDENS[2913.75] + + +def test_vl_load_weights_names_the_checkpoint_tensor_of_an_invalid_scale() -> None: + model = _bare_model(MiniMaxM3VLForConditionalGeneration) + weights = _packed_expert() + name = f"{_EXPERT}.w2.input_global_scale" + weights[name] = torch.tensor(0.0) + with patch.object(DecoderModelForCausalLM, "load_weights") as base_load_weights: + with pytest.raises(ValueError, match=re.escape(name)): + model.load_weights(weights, weight_mapper=HfWeightMapper()) + base_load_weights.assert_not_called() + + +def test_text_load_weights_translates_packed_experts_with_any_mapper() -> None: + # The translation does not depend on the mapper driving the load, and the + # M3 params map still reaches the base loader. + model = _bare_model(MiniMaxM3ForCausalLM) + expert = _EXPERT[len("language_model.") :] + weights = {k[len("language_model.") :]: v for k, v in _packed_expert().items()} + mapper = HfWeightMapper() + with patch.object(DecoderModelForCausalLM, "load_weights") as base_load_weights: + model.load_weights(weights, weight_mapper=mapper) + + call_kwargs = base_load_weights.call_args.kwargs + assert call_kwargs["weight_mapper"] is mapper + assert call_kwargs["params_map"] == MINIMAX_M3_PARAMS_MAP + loaded = call_kwargs["weights"] + assert sorted(loaded.keys()) == sorted( + f"{expert}.{proj}.{kind}" + for proj in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + ) + assert loaded[f"{expert}.w3.weight"] is weights[f"{expert}.w3.weight_packed"] + assert _f32_bits(loaded[f"{expert}.w3.weight_scale_2"]) == _WEIGHT_SCALE_2_GOLDENS[3000.0] diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 25fede204ec8..60426567f5f8 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -46,6 +46,7 @@ ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( MiniMaxM3KVCacheManagerV2, + resolve_sparse_layer_schedule, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.common import ( MiniMaxM3SparseConfig, @@ -57,6 +58,7 @@ MiniMaxM3MsaSparseAttentionMetadata, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import _group_max_reduce +from tensorrt_llm._torch.autotuner import autotune from tensorrt_llm._torch.compilation.remove_copy_pass import remove_copy_for_mutates_args from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( @@ -71,6 +73,9 @@ MiniMaxM3QKVIndexerLinear, _build_swiglu_oai_dense_mlp, _dispatch_attention_over_live_tokens, + _fold_gemma_boundary_norm_weights, + _get_index_projection_quant_config, + _load_index_qk_proj_weights, _load_qkv_index_proj_weights, _minimax_m3_swiglu_oai, _moe_routed_output_is_global, @@ -297,21 +302,15 @@ def test_minimax_m3_moe_reduces_only_local_terms( assert moe.allreduce.params == [None] -@pytest.mark.cpu_only -@pytest.mark.parametrize( - "scheduler_kind, expected_post_fusion", - [ - (MoESchedulerKind.EXTERNAL_COMM, True), - (MoESchedulerKind.FUSED_COMM, False), - ], -) -def test_minimax_m3_decoder_layer_sets_post_fusion_from_moe_scheduler( +def _build_decoder_layer_with_stubs( monkeypatch: pytest.MonkeyPatch, scheduler_kind: MoESchedulerKind, - expected_post_fusion: bool, -) -> None: - """A fused-communication MoE must not schedule a second boundary reduction.""" - + use_gemma_norm: bool, + *, + dtype: torch.dtype = torch.bfloat16, + hidden_size: int = 6144, +) -> MiniMaxM3DecoderLayer: + """A TP2 MoE decoder layer with real norms and stub attention, MoE and AllReduce.""" monkeypatch.setenv("TRTLLM_MINIMAX_M3_EAGER_FUSION_DISABLED", "0") class _DecoderComponent(nn.Module): @@ -334,20 +333,19 @@ def __init__(self, *args: object, **kwargs: object) -> None: "tensorrt_llm._torch.models.modeling_minimaxm3.MiniMaxM3MoE", _DecoderMoE, ) - monkeypatch.setattr( - "tensorrt_llm._torch.models.modeling_minimaxm3.RMSNorm", - _DecoderComponent, - ) monkeypatch.setattr( "tensorrt_llm._torch.models.modeling_minimaxm3.AllReduce", _DecoderComponent, ) + config = _make_text_config() + config.use_gemma_norm = use_gemma_norm + config.torch_dtype = dtype + config.hidden_size = hidden_size model_config = ModelConfig( - pretrained_config=_make_text_config(), + pretrained_config=config, mapping=Mapping(world_size=2, rank=0, tp_size=2), ) - - layer = MiniMaxM3DecoderLayer( + return MiniMaxM3DecoderLayer( model_config=model_config, layer_idx=3, aux_stream_dict={ @@ -356,6 +354,25 @@ def __init__(self, *args: object, **kwargs: object) -> None: }, ) + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "scheduler_kind, expected_post_fusion", + [ + (MoESchedulerKind.EXTERNAL_COMM, True), + (MoESchedulerKind.FUSED_COMM, False), + ], +) +@pytest.mark.parametrize("use_gemma_norm", [False, True]) +def test_minimax_m3_decoder_layer_sets_post_fusion_from_moe_scheduler( + monkeypatch: pytest.MonkeyPatch, + scheduler_kind: MoESchedulerKind, + expected_post_fusion: bool, + use_gemma_norm: bool, +) -> None: + """Plain and folded-Gemma norms preserve fusion without a second routed reduction.""" + layer = _build_decoder_layer_with_stubs(monkeypatch, scheduler_kind, use_gemma_norm) + assert layer.enable_fusion assert layer.pre_feed_forward_fusion assert layer.post_feed_forward_fusion is expected_post_fusion @@ -1326,6 +1343,206 @@ def test_get_moe_layer_ids_length_mismatch_raises(): get_moe_layer_ids(cfg) +# --------------------------------------------------------------------------- +# Compressed-tensors NVFP4 MoE packaging of the MiniMax M3 architecture +# --------------------------------------------------------------------------- + + +def _make_packaged_text_config(): + """The trimmed text config as the compressed-tensors package declares it. + + Every attention layer is sparse and K-only while layers 0-2 keep a dense + MLP; rope_theta 1e7, 1M positions, no MTP module and no text model_type. + """ + cfg = _make_text_config() + cfg.sparse_attention_config["sparse_attention_freq"] = [1] * _NUM_HIDDEN_LAYERS + cfg.sparse_attention_config["sparse_disable_index_value"] = [1] * _NUM_HIDDEN_LAYERS + cfg.rope_theta = 10000000 + cfg.max_position_embeddings = 1048576 + cfg.num_mtp_modules = 0 + del cfg.model_type + return cfg + + +def _packaged_vl_config(text_config_as_dict: bool = False, **quant_overrides): + quantization_config = { + "quant_method": "mxfp8", + "activation_scheme": "dynamic", + "weight_block_size": [1, 32], + "group_size": 16, + "moe_quant_algo": "NVFP4", + "moe_quant_format": "nvfp4-pack-quantized", + "ignored_layers": [ + "lm_head", + "model.embed_tokens", + "vision_tower", + "multi_modal_projector", + "patch_merge_mlp", + ] + + [ + f"language_model.model.layers.{i}.block_sparse_moe.gate" + for i in range(3, _NUM_HIDDEN_LAYERS) + ], + } + quantization_config.update(quant_overrides) + text_config = _make_packaged_text_config() + cfg = _make_vl_config() + cfg.text_config = dict(vars(text_config)) if text_config_as_dict else text_config + cfg.quantization_config = quantization_config + return cfg + + +def _expand_packaged_quant_config(pretrained_config): + """Parse the inline config as from_pretrained does, then apply the M3 hook.""" + quant_config, layer_quant_config = ModelConfig.load_hf_quant_config( + pretrained_config.quantization_config, "CUTLASS" + ) + assert quant_config.quant_algo == QuantAlgo.MXFP8 + assert layer_quant_config is None + return ModelConfig._set_minimax_m3_layer_quant_config( + pretrained_config, quant_config, layer_quant_config + ) + + +@pytest.mark.parametrize("text_config_as_dict", [False, True]) +def test_packaged_moe_config_expands_to_mixed_precision(text_config_as_dict: bool) -> None: + quant_config, layer_quant_config = _expand_packaged_quant_config( + _packaged_vl_config(text_config_as_dict) + ) + + assert quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert "model.layers.3.block_sparse_moe.gate" in quant_config.exclude_modules + assert not any(name.startswith("language_model.") for name in quant_config.exclude_modules) + assert quant_config.is_module_excluded_from_quantization("lm_head") + + expected = {} + for layer in range(_NUM_HIDDEN_LAYERS): + prefix = f"model.layers.{layer}" + names = [ + f"{prefix}.self_attn.{proj}" + for proj in ("q_proj", "k_proj", "v_proj", "o_proj", "index_q_proj", "index_k_proj") + ] + if _MOE_LAYER_FREQ[layer]: + names += [ + f"{prefix}.block_sparse_moe.shared_experts.{proj}" + for proj in ("gate_proj", "up_proj", "down_proj") + ] + expected[f"{prefix}.block_sparse_moe.experts"] = (QuantAlgo.NVFP4, 16) + else: + names += [f"{prefix}.mlp.{proj}" for proj in ("gate_proj", "up_proj", "down_proj")] + expected.update({name: (QuantAlgo.MXFP8, 32) for name in names}) + assert {k: (v.quant_algo, v.group_size) for k, v in layer_quant_config.items()} == expected + + +def test_packaged_moe_config_honors_ignored_layers() -> None: + cfg = _packaged_vl_config() + cfg.quantization_config["ignored_layers"].append( + "language_model.model.layers.4.self_attn.o_proj" + ) + quant_config, layer_quant_config = _expand_packaged_quant_config(cfg) + assert "model.layers.4.self_attn.o_proj" not in layer_quant_config + assert quant_config.is_module_excluded_from_quantization("model.layers.4.self_attn.o_proj") + assert "model.layers.5.self_attn.o_proj" in layer_quant_config + + +@pytest.mark.parametrize( + "overrides", + [ + {"moe_quant_algo": "MXFP4"}, + {"moe_quant_format": "float-quantized"}, + {"group_size": 32}, + {"group_size": None}, + {"moe_quant_format": None}, + {"moe_quant_algo": None}, + ], +) +def test_packaged_moe_config_rejects_unsupported_expert_recipes(overrides) -> None: + with pytest.raises(ValueError, match="routed-expert packaging"): + _expand_packaged_quant_config(_packaged_vl_config(**overrides)) + + +def test_uniform_mxfp8_and_bf16_configs_are_unchanged() -> None: + cfg = _packaged_vl_config() + for key in ("moe_quant_algo", "moe_quant_format", "group_size"): + del cfg.quantization_config[key] + quant_config, _ = ModelConfig.load_hf_quant_config(cfg.quantization_config, "CUTLASS") + result = ModelConfig._set_minimax_m3_layer_quant_config(cfg, quant_config, None) + assert result[0] is quant_config and result[1] is None + + no_quant = QuantConfig() + result = ModelConfig._set_minimax_m3_layer_quant_config(_make_vl_config(), no_quant, None) + assert result[0] is no_quant and result[1] is None + + +def test_modelopt_mixed_precision_map_is_normalized_as_before() -> None: + base_config = QuantConfig(quant_algo=QuantAlgo.MXFP8, group_size=32) + expert_config = QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16) + global_config = QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION) + layer_map = { + "language_model.model.layers.0.self_attn.o_proj": base_config, + "language_model.model.layers.3.block_sparse_moe.experts.0.w1": expert_config, + } + quant_config, out = ModelConfig._set_minimax_m3_layer_quant_config( + _make_vl_config(), global_config, layer_map + ) + assert quant_config is global_config + assert out["model.layers.0.self_attn.o_proj"] is base_config + assert out["model.layers.3.block_sparse_moe.experts.0.w1"] is expert_config + assert not any(key.startswith("language_model.") for key in out) + for layer in range(_NUM_HIDDEN_LAYERS): + key = f"model.layers.{layer}.block_sparse_moe.experts" + if _MOE_LAYER_FREQ[layer]: + assert (out[key].quant_algo, out[key].group_size) == (QuantAlgo.NVFP4, 16) + else: + assert key not in out + + +def test_index_projection_quant_config_follows_declared_entries() -> None: + _, layer_quant_config = _expand_packaged_quant_config(_packaged_vl_config()) + declared = _get_index_projection_quant_config( + SimpleNamespace(quant_config_dict=layer_quant_config), 4 + ) + assert (declared.quant_algo, declared.group_size) == (QuantAlgo.MXFP8, 32) + # Checkpoints without per-layer entries keep the unquantized index projection. + for quant_config_dict in (None, {}): + assert ( + _get_index_projection_quant_config( + SimpleNamespace(quant_config_dict=quant_config_dict), 4 + ) + is None + ) + + mismatched = dict(layer_quant_config) + mismatched["model.layers.4.self_attn.index_k_proj"] = QuantConfig( + quant_algo=QuantAlgo.NVFP4, group_size=16 + ) + with pytest.raises(ValueError, match="same quantization"): + _get_index_projection_quant_config(SimpleNamespace(quant_config_dict=mismatched), 4) + del mismatched["model.layers.4.self_attn.index_k_proj"] + with pytest.raises(ValueError, match="same quantization"): + _get_index_projection_quant_config(SimpleNamespace(quant_config_dict=mismatched), 4) + + +def test_kv_cache_schedule_follows_the_checkpoint_attention_schedule() -> None: + all_layers = list(range(_NUM_HIDDEN_LAYERS)) + assert resolve_sparse_layer_schedule( + None, None, _make_packaged_text_config(), _NUM_HIDDEN_LAYERS + ) == (all_layers, all_layers) + assert resolve_sparse_layer_schedule(None, None, _make_text_config(), _NUM_HIDDEN_LAYERS) == ( + [3, 4, 5, 6], + [3, 4, 5, 6], + ) + # Explicit ids win; without a checkpoint schedule the public convention holds. + assert resolve_sparse_layer_schedule([5], None, _make_packaged_text_config(), 7) == ([5], [5]) + for config in (None, SimpleNamespace(num_hidden_layers=_NUM_HIDDEN_LAYERS)): + assert resolve_sparse_layer_schedule(None, None, config, _NUM_HIDDEN_LAYERS) == ( + [3, 4, 5, 6], + [3, 4, 5, 6], + ) + with pytest.raises(ValueError, match="outside the 4 target layers"): + resolve_sparse_layer_schedule(None, None, _make_packaged_text_config(), 4) + + # --------------------------------------------------------------------------- # attention module transforms # --------------------------------------------------------------------------- @@ -3004,3 +3221,449 @@ def check_forward(packed, unused_k, unused_v, meta, args): ) kernel.forward.assert_called_once() torch.testing.assert_close(output, q) + + +# --------------------------------------------------------------------------- +# Compressed-tensors NVFP4 MoE packaging: CUDA loading and forward checks +# --------------------------------------------------------------------------- +# +# The references below are written from the checkpoint format (E4M3 / UE8M0 / +# E2M1 definitions) without calling TensorRT-LLM quantizers. Each forward is +# checked eagerly and again through CUDA graph capture and replay. + +_E2M1_GRID = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _relative_l2(out: torch.Tensor, ref: torch.Tensor) -> float: + return float((out.float() - ref.float()).norm() / ref.float().norm().clamp_min(1e-6)) + + +def _cosine(out: torch.Tensor, ref: torch.Tensor) -> float: + return float( + torch.nn.functional.cosine_similarity(out.float().flatten(), ref.float().flatten(), dim=0) + ) + + +def _graph_replay(fn, *inputs: torch.Tensor): + """Capture ``fn(*inputs)`` in a CUDA graph; return a replay(new_inputs) helper.""" + static_inputs = [t.clone() for t in inputs] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + fn(*static_inputs) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_output = fn(*static_inputs) + + def replay(*new_inputs: torch.Tensor) -> torch.Tensor: + for static, new in zip(static_inputs, new_inputs): + static.copy_(new) + graph.replay() + torch.cuda.synchronize() + return static_output.clone() + + return replay + + +def _mxfp8_dequant(weight: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tensor: + """E4M3 weight times 2**(UE8M0 - 127) per 32 input channels.""" + exponent = scale_inv.float().repeat_interleave(32, dim=1) - 127.0 + return weight.float() * torch.exp2(exponent) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MXFP8 index projection needs CUDA") +@pytest.mark.parametrize("use_cuda_graph", [False, True]) +def test_packaged_index_projection_keeps_mxfp8_scales_per_tp4_rank( + use_cuda_graph: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """index_q/k_proj load as MXFP8 with each rank's scale rows next to its weight rows. + + Every index-Q head gets its own scale exponent, so a scale row that lands + on another head's weights changes the output by powers of two; index-K is + replicated on every rank with all of its scale rows. + """ + import tensorrt_llm._torch.distributed as distributed + + # Simulate each rank on one GPU; COLUMN projections need no collective. + monkeypatch.setattr(distributed, "AllReduce", lambda **_kwargs: nn.Identity()) + torch.manual_seed(0) + hidden, index_dim, index_heads, tp_size = 256, 32, 4, 4 + head_exponents = torch.tensor([127, 123, 130, 125], dtype=torch.int32) + q_weight = (torch.randn(index_heads * index_dim, hidden) * 2).to(torch.float8_e4m3fn) + q_scale = ( + head_exponents.repeat_interleave(index_dim)[:, None] + + torch.randint(-1, 2, (index_heads * index_dim, hidden // 32)) + ).to(torch.uint8) + k_weight = (torch.randn(index_dim, hidden) * 2).to(torch.float8_e4m3fn) + k_scale = torch.randint(121, 129, (index_dim, hidden // 32)).to(torch.uint8) + x = torch.randn(48, hidden, device="cuda", dtype=torch.bfloat16) + x_next = torch.randn_like(x) + ref_q = x.float() @ _mxfp8_dequant(q_weight, q_scale).cuda().T + ref_k = x.float() @ _mxfp8_dequant(k_weight, k_scale).cuda().T + mxfp8 = QuantConfig(quant_algo=QuantAlgo.MXFP8, group_size=32) + + for rank in range(tp_size): + text_cfg, model_cfg = _make_attention_test_config() + text_cfg.hidden_size = hidden + text_cfg.num_attention_heads = 8 + text_cfg.num_key_value_heads = 4 + text_cfg.sparse_attention_config["sparse_num_index_heads"] = index_heads + model_cfg.mapping = Mapping(world_size=tp_size, rank=rank, tp_size=tp_size) + model_cfg.quant_config_dict = { + "model.layers.3.self_attn.index_q_proj": mxfp8, + "model.layers.3.self_attn.index_k_proj": mxfp8, + } + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=3, + is_sparse_attention_layer=True, + disable_index_value=True, + ) + proj = attn.index_qk_proj + assert proj.quant_config.quant_algo == QuantAlgo.MXFP8 + proj.create_weights() + assert type(proj.quant_method).__name__ == "MXFP8LinearMethod" + proj.cuda() + + layer = nn.Module() + layer.self_attn = attn + _load_index_qk_proj_weights( + layer, + { + "self_attn.index_q_proj.weight": q_weight, + "self_attn.index_q_proj.weight_scale_inv": q_scale, + "self_attn.index_k_proj.weight": k_weight, + "self_attn.index_k_proj.weight_scale_inv": k_scale, + }, + ) + rows = slice(rank * index_dim, (rank + 1) * index_dim) + expected_weight = torch.cat([q_weight[rows], k_weight]).view(torch.uint8).cuda() + assert torch.equal(proj.weight.view(torch.uint8), expected_weight) + expected_scale = torch.cat([q_scale[rows], k_scale]).cuda() + if proj.quant_method.use_cutlass: + expected_scale = torch.ops.trtllm.block_scale_interleave(expected_scale) + assert torch.equal(proj.weight_scale.view(torch.uint8).flatten(), expected_scale.flatten()) + + def forward(inp: torch.Tensor) -> torch.Tensor: + return proj(inp) + + if use_cuda_graph: + replay = _graph_replay(forward, x) + out, out_next = replay(x), replay(x_next) + assert torch.equal(out_next, forward(x_next)) + else: + out = forward(x) + ref = torch.cat([ref_q[:, rows], ref_k], dim=1) + # Dynamic per-token MXFP8 activations cost ~1% relative error; a scale + # row on the wrong head is off by at least a factor of two. + assert _relative_l2(out, ref) < 0.05, rank + assert _cosine(out, ref) > 0.998, rank + + +def _e2m1_values(packed: torch.Tensor) -> torch.Tensor: + """Unpack [rows, K/2] bytes (even element in the low nibble) to float32 [rows, K].""" + codes = torch.stack((packed & 0xF, packed >> 4), dim=-1).reshape(packed.shape[0], -1).long() + magnitude = torch.tensor(_E2M1_GRID, device=packed.device)[codes & 0x7] + return torch.where(codes >= 8, -magnitude, magnitude) + + +def _round_to_e2m1(x: torch.Tensor) -> torch.Tensor: + """Round to the E2M1 grid, ties to the even code, saturating at 6.""" + grid = torch.tensor(_E2M1_GRID, device=x.device) + bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], device=x.device) + magnitude = x.abs() + index = torch.bucketize(magnitude, bounds) + index = index + torch.isin(magnitude, bounds[[1, 3, 5]]).long() + return torch.copysign(grid[index], x) + + +def _nvfp4_activation(x: torch.Tensor, global_scale: float = 1.0) -> torch.Tensor: + """Quantize-dequantize with 16-element E4M3 block scales and a static global scale.""" + blocks = x.float().reshape(*x.shape[:-1], -1, 16) + scale = blocks.abs().amax(-1, keepdim=True) * (1.0 / 6.0) * global_scale + scale = scale.to(torch.float8_e4m3fn).float() + scaled = torch.where(scale > 0, blocks * (global_scale / scale), torch.zeros_like(blocks)) + return (_round_to_e2m1(scaled) * scale / global_scale).reshape(x.shape) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="CUTLASS NVFP4 MoE needs CUDA") +@pytest.mark.parametrize("use_cuda_graph", [False, True]) +@pytest.mark.parametrize("tune", [False, True]) +def test_packaged_nvfp4_experts_load_and_run_on_cutlass(use_cuda_graph: bool, tune: bool) -> None: + """Packed experts reach the CUTLASS NVFP4 path intact and compute the static recipe. + + Each expert has its own global scales, so a dropped scale or a wrong + expert assignment is visible in the output. The FC1 input is quantized + with the static input scale, the FC1 results round to BF16, the FC2 input + is quantized from that BF16 activation and the routing weights apply in + the finalize. + """ + from tensorrt_llm._torch.moe.fused_moe import SwigluBiasActivation, create_moe + from tensorrt_llm.quantization.utils.fp4_utils import float4_sf_dtype + + torch.manual_seed(0) + num_experts, top_k, hidden, inter, tokens = 8, 2, 256, 256, 24 + prefix = "language_model.model.layers.3.block_sparse_moe.experts" + generator = torch.Generator().manual_seed(1) + checkpoint = {} + for e in range(num_experts): + for proj, rows, cols, global_scale in ( + ("w1", inter, hidden, 50.0 + 7.0 * e), + ("w3", inter, hidden, 50.0 + 7.0 * e), + ("w2", hidden, inter, 60.0 + 5.0 * e), + ): + name = f"{prefix}.{e}.{proj}" + checkpoint[f"{name}.weight_packed"] = torch.randint( + 0, 256, (rows, cols // 2), dtype=torch.uint8, generator=generator + ) + checkpoint[f"{name}.weight_scale"] = ( + torch.rand(rows, cols // 16, generator=generator) * 2 + 0.25 + ).to(torch.float8_e4m3fn) + checkpoint[f"{name}.weight_global_scale"] = torch.tensor(global_scale) + checkpoint[f"{name}.input_global_scale"] = torch.tensor(1.0) + + _, layer_quant_config = _expand_packaged_quant_config(_packaged_vl_config()) + pretrained = _make_packaged_text_config() + pretrained.num_local_experts = num_experts + pretrained.hidden_size = hidden + pretrained.intermediate_size = inter + pretrained.torch_dtype = torch.bfloat16 + model_config = ModelConfig( + pretrained_config=pretrained, + quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION), + quant_config_dict=layer_quant_config, + moe_backend="CUTLASS", + max_num_tokens=256, + ) + bias = torch.randn(num_experts, device="cuda") * 0.1 + routing = MiniMaxM3MoeRoutingMethod( + top_k=top_k, + num_experts=num_experts, + callable_e_score_correction_bias=lambda: bias, + routed_scaling_factor=2.0, + ) + moe = create_moe( + routing_method=routing, + num_experts=num_experts, + model_config=model_config, + layer_idx=3, + override_quant_config=MiniMaxM3MoE._get_experts_quant_config(model_config, 3), + activation=SwigluBiasActivation(gate_sigmoid_scale=1.702, linear_offset=1.0, clamp=7.0), + allow_backend_degradation=False, + ) + backend = moe.backend + assert type(backend).__name__ == "CutlassFusedMoE" + translated = MiniMaxM3HfWeightMapper().preprocess_weights(dict(checkpoint)) + module_weights = { + key[len(prefix) + 1 :]: value + for key, value in translated.items() + if key.startswith(prefix + ".") + } + with torch.device("cuda"): + moe.create_weights() + assert type(backend.quant_method).__name__ == "NVFP4CutlassFusedMoEMethod" + moe.load_weights([module_weights]) + moe.post_load_weights() + + # Static input scales of 1.0 stay 1.0; alphas are the reciprocal weight global scales. + assert backend.fc31_input_scale.item() == 1.0 and backend.fc2_input_scale.item() == 1.0 + experts = [] + for e in range(num_experts): + source = { + proj: ( + checkpoint[f"{prefix}.{e}.{proj}.weight_packed"], + checkpoint[f"{prefix}.{e}.{proj}.weight_scale"], + torch.reciprocal(checkpoint[f"{prefix}.{e}.{proj}.weight_global_scale"]).item(), + ) + for proj in ("w1", "w2", "w3") + } + assert torch.equal( + backend.w3_w1_weight[e].view(torch.uint8), + torch.cat([source["w3"][0], source["w1"][0]]).cuda(), + ) + assert torch.equal(backend.w2_weight[e].view(torch.uint8), source["w2"][0].cuda()) + expected_fc31_scale = torch.ops.trtllm.block_scale_interleave( + torch.cat([source["w3"][1], source["w1"][1]]).cuda().view(float4_sf_dtype) + ) + assert torch.equal( + backend.w3_w1_weight_scale[e].view(torch.uint8).flatten(), + expected_fc31_scale.view(torch.uint8).flatten(), + ) + expected_fc2_scale = torch.ops.trtllm.block_scale_interleave( + source["w2"][1].cuda().view(float4_sf_dtype) + ) + assert torch.equal( + backend.w2_weight_scale[e].view(torch.uint8).flatten(), + expected_fc2_scale.view(torch.uint8).flatten(), + ) + alphas = (source["w1"][2], source["w3"][2], source["w2"][2]) + # The alpha is 1 / (input_scale * (1 / weight_scale_2)): within one ulp. + for alpha, golden in ( + (backend.fc31_alpha[e].item(), source["w1"][2]), + (backend.fc2_alpha[e].item(), source["w2"][2]), + ): + assert abs(alpha - golden) <= 2**-23 * golden, (e, alpha, golden) + # Dequantized FP4 weights without the global scale (exact in FP32), and that scale. + experts.append( + ( + tuple( + _e2m1_values(w).cuda() * s.float().cuda().repeat_interleave(16, dim=1) + for w, s, _ in (source["w1"], source["w3"], source["w2"]) + ), + alphas, + ) + ) + + def gemm(a: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """FP32 rounding of the FP64 sum of exactly representable products.""" + return (a.double() @ w.double().T).float() + + def reference(x: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: + """Static unit FC1 scale; BF16 FC1 results; FC2 input quantized from its + BF16 cast; routing weights in the finalize.""" + scores = torch.sigmoid(logits.float()) + ids = torch.topk(scores + bias, top_k, dim=-1).indices + weights = scores.gather(1, ids) + weights = weights / (weights.sum(-1, keepdim=True) + 1e-20) * 2.0 + x_q = _nvfp4_activation(x) + slots = torch.zeros(x.shape[0], top_k, hidden, device="cuda") + for e, ((w1, w3, w2), (a1, a3, a2)) in enumerate(experts): + rows, slot = (ids == e).nonzero(as_tuple=True) + if rows.numel() == 0: + continue + acc1, acc3 = gemm(x_q[rows], w1), gemm(x_q[rows], w3) + w = weights[rows, slot, None] + gate = (acc1 * a1).bfloat16().float().clamp(max=7.0) + up = (acc3 * a3).bfloat16().float().clamp(-7.0, 7.0) + act = (gate * torch.sigmoid(1.702 * gate) * (up + 1.0)).bfloat16().float() + y = (gemm(_nvfp4_activation(act), w2) * a2).bfloat16().float() + slots[rows, slot] = y * w + out = torch.zeros(x.shape[0], hidden, device="cuda") + for j in range(top_k): + out = out + slots[:, j] + return out + + def forward(x: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: + return moe(x, logits, all_rank_num_tokens=[tokens], use_dp_padding=False) + + x = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + logits = torch.randn(tokens, num_experts, device="cuda") + x_next, logits_next = torch.randn_like(x), torch.randn_like(logits) + with torch.inference_mode(): + # Tune the actual shape, then verify the selected tactic eagerly and in a graph. + if tune: + with autotune(tune_mode=True, skip_dynamic_tuning_buckets=True): + forward(x, logits) + if use_cuda_graph: + replay = _graph_replay(forward, x, logits) + outputs = [ + (replay(x, logits), x, logits), + (replay(x_next, logits_next), x_next, logits_next), + ] + # The fused top-k combine may accumulate in a different order per launch. + assert _relative_l2(outputs[1][0], forward(x_next, logits_next)) < 1e-3 + else: + outputs = [(forward(x, logits), x, logits)] + for out, inp, lg in outputs: + ref = reference(inp, lg) + assert torch.isfinite(out).all() + # Same recipe on both sides; the remainder is accumulation order + # and the BF16 rounding of the GEMM outputs. + assert _relative_l2(out, ref) < 0.02 + assert _cosine(out, ref) > 0.999 + + +@pytest.mark.cpu_only +def test_minimax_m3_fold_gemma_boundary_weights_accepts_partial_lazy_load(): + stored = torch.tensor([0.25, -0.25], dtype=torch.bfloat16) + + class LazyWeight: + def __getitem__(self, index): + assert index == slice(None) + return stored + + untouched = LazyWeight() + weights = { + "model.layers.2.input_layernorm.weight": LazyWeight(), + "model.layers.2.self_attn.q_norm.weight": untouched, + } + result = _fold_gemma_boundary_norm_weights(weights) + assert result is weights + torch.testing.assert_close( + result["model.layers.2.input_layernorm.weight"], stored + 1.0, rtol=0, atol=0 + ) + assert result["model.layers.2.self_attn.q_norm.weight"] is untouched + + +def _stub_minimax_m3_attention_construction(monkeypatch, *, is_neox=True): + """Keep the M3 constructor real while removing generic GPU setup and projection allocation.""" + + def initialize_attention(self, **kwargs): + nn.Module.__init__(self) + config = kwargs["config"] + self.pos_embd_params = SimpleNamespace(rope=kwargs["pos_embd_params"].rope, is_neox=is_neox) + self.rotary_emb = object() + self.head_dim = config.pretrained_config.head_dim + self.qkv_proj = nn.Module() + self.qkv_proj.mapping = config.mapping + + def make_linear(*args, **kwargs): + projection = nn.Module() + projection.quant_config = kwargs["quant_config"] + return projection + + monkeypatch.setattr(modeling_minimaxm3.Attention, "__init__", initialize_attention) + monkeypatch.setattr(torch.cuda, "Stream", lambda: None) + monkeypatch.setattr(torch.cuda, "Event", lambda: None) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + projection_factory = Mock(side_effect=make_linear) + monkeypatch.setattr(modeling_minimaxm3, "Linear", projection_factory) + return projection_factory + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("entries", ["matching", "absent", "mismatched", "missing_k"]) +def test_minimax_m3_index_projection_quantization_follows_checkpoint_entries(monkeypatch, entries): + """index_qk_proj is quantized exactly when the checkpoint declares matching per-layer entries.""" + _, model_config = _make_fused_qk_norm_rope_test_config() + # Other-layer entries must not affect the projection being constructed. + declared = QuantConfig(quant_algo=QuantAlgo.MXFP8, group_size=32) + other = QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16) + layer_quant_config = { + "model.layers.2.self_attn.index_q_proj": other, + "model.layers.2.self_attn.index_k_proj": other, + } + if entries != "absent": + layer_quant_config["model.layers.3.self_attn.index_q_proj"] = declared + if entries in ("matching", "mismatched"): + layer_quant_config["model.layers.3.self_attn.index_k_proj"] = ( + declared if entries == "matching" else other + ) + model_config.quant_config_dict = layer_quant_config + projection_factory = _stub_minimax_m3_attention_construction(monkeypatch) + lookup = Mock(wraps=_get_index_projection_quant_config) + monkeypatch.setattr(modeling_minimaxm3, "_get_index_projection_quant_config", lookup) + if entries in ("mismatched", "missing_k"): + with pytest.raises(ValueError, match="same quantization"): + MiniMaxM3Attention( + model_config=model_config, + layer_idx=3, + is_sparse_attention_layer=True, + disable_index_value=True, + ) + projection_factory.assert_not_called() + else: + layer = MiniMaxM3Attention( + model_config=model_config, + layer_idx=3, + is_sparse_attention_layer=True, + disable_index_value=True, + ) + projection_factory.assert_called_once() + expected = declared if entries == "matching" else None + assert layer.index_qk_proj.quant_config is expected + assert projection_factory.call_args.kwargs["quant_config"] is expected + lookup.assert_called_once_with(model_config, 3) From a949410a581c213b266c38bf417757baa2825671 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 13:31:44 -0700 Subject: [PATCH 2/5] [None][feat] MiniMax-M3: drop packaging pieces main now covers Packed NVFP4 experts now load through main's convert_compressed_tensors_nvfp4_weights, so translate_packed_nvfp4_experts, its mapper hook and tests go. Index-K buffer and byte estimate share _index_k_bytes_per_token; the deployment guide states how main treats uninvertible global scales; tests follow main's MoE class names and KVCM V2. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- ...ployment-guide-for-minimax-m3-on-trtllm.md | 13 +- .../sparse/minimax_m3/cache_manager.py | 36 +-- tensorrt_llm/_torch/model_config.py | 11 +- .../checkpoints/hf/minimaxm3_weight_mapper.py | 96 +------- .../_torch/models/modeling_minimaxm3.py | 14 +- .../test_minimax_m3_shared_draft_layers.py | 5 +- .../hf/test_minimaxm3_weight_mapper.py | 213 ------------------ .../unittest/_torch/models/test_minimax_m3.py | 38 +++- 8 files changed, 65 insertions(+), 361 deletions(-) diff --git a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md index aa445ef0dc48..58efce12646f 100644 --- a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md @@ -65,12 +65,13 @@ routed expert projection as `weight_packed`, `weight_scale`, 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 is zero, negative or not finite, or an expert whose `w1` -and `w3` weight global scales differ, fails the load with an error naming the -checkpoint tensors. 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. +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)). diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 91253d7c8b18..f741fb16a537 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -513,16 +513,12 @@ 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.""" global_layer_idx = int(self.pp_layers[int(local_layer_idx)]) has_index_k = global_layer_idx in self.sparse_layer_ids if data_role == Role.INDEX_KEY: - if not has_index_k: - return 0 - index_dtype = self._torch_dtype_for_index_cache() - index_elem_bytes = torch.empty(0, dtype=index_dtype, device="cpu").element_size() - return self.sparse_index_dim * index_elem_bytes + 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) @@ -741,18 +737,19 @@ 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): - """Register the paged index-K buffer for each sparse local layer.""" + 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**: :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. + """ + size_per_block = self._index_k_bytes_per_token() * tokens_per_block return { - self.layer_offsets[layer_id]: [ - BufferConfig( - role=Role.INDEX_KEY, - size=self.get_layer_bytes_per_token( - self.layer_offsets[layer_id], Role.INDEX_KEY - ) - * tokens_per_block, - ) - ] + self.layer_offsets[layer_id]: [BufferConfig(role=Role.INDEX_KEY, size=size_per_block)] for layer_id in self.sparse_layer_ids if layer_id in self.layer_offsets } @@ -782,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": diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 85e752efcaf6..882572875d2d 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -1024,8 +1024,9 @@ def _get_mxfp8_moe_quant_config( group_size=NVFP4_SF_VEC_SIZE) @staticmethod - def _build_minimax_m3_mxfp8_moe_quant_config(pretrained_config, - quant_config): + 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 @@ -1125,8 +1126,10 @@ def _build_minimax_m3_mxfp8_moe_quant_config(pretrained_config, return mixed_quant_config, layer_quant_config @staticmethod - def _set_minimax_m3_layer_quant_config(pretrained_config, quant_config, - layer_quant_config): + 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 diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py index 261559f186d4..bb6d58957fd3 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py @@ -13,13 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -import re -from collections.abc import Mapping - -import torch from torch import Tensor, nn -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper @@ -27,105 +22,16 @@ r"^(.*\.block_sparse_moe)\.e_score_correction_bias$": r"\1.gate.e_score_correction_bias", } -# Routed-expert tensors of the compressed-tensors NVFP4 MoE packaging of the -# MiniMax M3 architecture. Per projection: ``weight_packed`` (uint8, two E2M1 -# values per byte), ``weight_scale`` (E4M3 per 16 elements) and the FP32 -# scalars ``weight_global_scale`` / ``input_global_scale``, stored as the -# multiplier 2688 / amax (2688 = E4M3 max 448 * E2M1 max 6). -_PACKED_NVFP4_EXPERT_KEY = re.compile( - r"^(?P.*\.block_sparse_moe\.experts\.\d+\.w[123])\." - r"(?Pweight_packed|weight_global_scale|input_global_scale)$" -) -_NVFP4_GLOBAL_SCALE_RANGE = 448.0 * 6.0 - - -def nvfp4_weight_scale_2(weight_global_scale: Tensor) -> Tensor: - """Return the NVFP4 MoE ``weight_scale_2`` of a stored global scale. - - compressed-tensors stores ``weight_global_scale = 2688 / amax`` (E4M3 max - 448 times E2M1 max 6); the dequantization scale is its reciprocal, - ``amax / 2688``, computed in FP32. - """ - return torch.reciprocal(weight_global_scale.float()) - - -def _read_global_scale(name: str, value) -> Tensor: - """Materialize one global scale, rejecting values that would disable an expert.""" - scale = (value if isinstance(value, Tensor) else value[...]).to(torch.float32) - if scale.numel() != 1 or not bool(torch.isfinite(scale).all()) or not bool((scale > 0).all()): - raise ValueError( - f"Checkpoint tensor {name} must hold one finite positive global scale, " - f"got {scale.flatten()[:4].tolist()}." - ) - return scale - - -def translate_packed_nvfp4_experts(weights: Mapping[str, Tensor]) -> Mapping[str, Tensor]: - """Rename packed NVFP4 routed-expert tensors to the NVFP4 MoE contract. - - ``weight_packed`` becomes ``weight`` with the same bytes, and the global - scales become the ``weight_scale_2`` / ``input_scale`` divisors - (amax / 2688) the NVFP4 MoE loader reads; ``weight_scale`` keeps its name - and bytes. Only the scalar scales are materialized. Every other key passes - through, including already canonical ``weight_scale_2`` / ``input_scale``, - so a second call is a no-op. - - Raises: - ValueError: naming the checkpoint tensor, when a global scale is not - finite and positive, or when an expert's gate (w1) and up (w3) - weight scales differ: the fused gate/up GEMM applies one weight - scale per expert. - """ - matches = {} - for name in weights.keys(): - match = _PACKED_NVFP4_EXPERT_KEY.match(name) - if match is not None: - matches[name] = match - if not matches: - return weights - - renamed = {} - gate_up_scales: dict[str, dict[str, tuple[str, Tensor]]] = {} - for name in weights.keys(): - value = weights[name] - match = matches.get(name) - if match is None: - renamed[name] = value - continue - proj, kind = match["proj"], match["kind"] - if kind == "weight_packed": - renamed[f"{proj}.weight"] = value - elif kind == "weight_global_scale": - scale_2 = nvfp4_weight_scale_2(_read_global_scale(name, value)) - renamed[f"{proj}.weight_scale_2"] = scale_2 - expert, w = proj.rsplit(".", 1) - if w != "w2": - gate_up_scales.setdefault(expert, {})[w] = (name, scale_2) - else: - renamed[f"{proj}.input_scale"] = torch.reciprocal(_read_global_scale(name, value)) - - for scales in gate_up_scales.values(): - if len(scales) == 2 and not torch.equal(scales["w1"][1], scales["w3"][1]): - raise ValueError( - f"Checkpoint tensors {scales['w1'][0]} and {scales['w3'][0]} decode to " - "different weight scales; the fused gate/up expert GEMM needs one." - ) - return ConsumableWeightsDict.take_ownership(weights, renamed) - @register_mapper("HF", "MiniMaxM3SparseForCausalLM") @register_mapper("HF", "MiniMaxM3SparseForConditionalGeneration") class MiniMaxM3HfWeightMapper(HfWeightMapper): - """Handle M3 gate naming, MXFP8 GQA duplication and packed NVFP4 experts.""" + """Handle M3 gate naming and MXFP8 GQA duplication for loader v2.""" def __init__(self) -> None: super().__init__() self.params_map = MINIMAX_M3_PARAMS_MAP - def preprocess_weights(self, weights: Mapping[str, Tensor]) -> Mapping[str, Tensor]: - """See :func:`translate_packed_nvfp4_experts`.""" - return translate_packed_nvfp4_experts(weights) - def _duplicate_kv_weights( self, module: nn.Module, new_name: str, weights: dict[str, Tensor] ) -> dict[str, Tensor]: diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index e5d75dcddaf0..fdc8a7d45051 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -78,11 +78,7 @@ from ..speculative import SpecMetadata from ..utils import AuxStreamType, EventType, get_model_extra_attrs, is_torch_compiling from .checkpoints.base_weight_mapper import BaseWeightMapper -from .checkpoints.hf.minimaxm3_weight_mapper import ( - MINIMAX_M3_PARAMS_MAP, - MiniMaxM3HfWeightMapper, - translate_packed_nvfp4_experts, -) +from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper from .checkpoints.hf.weight_mapper import convert_compressed_tensors_nvfp4_weights from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import ( @@ -3218,8 +3214,6 @@ def load_weights( # set of unknown top-level keys it surfaces as errors. from .modeling_minimaxm3_vl import load_minimax_m3_vl_state_dict, split_multimodal_weights - # Rewrite before the prefix strip so errors name the checkpoint tensor. - weights = translate_packed_nvfp4_experts(weights) text_or_other_weights, vision_weights = split_multimodal_weights(weights) text_weights, ignored = _strip_language_model_prefix(text_or_other_weights) @@ -3231,12 +3225,6 @@ def load_weights( ) self.last_loaded_vision_keys = loaded self.last_missing_vision_keys = missing - logger.info( - f"MiniMax-M3 VL weights: {len(self.last_loaded_vision_keys)} vision tensors " - f"loaded, {len(self.last_missing_vision_keys)} vision parameters missing, " - f"{len(ignored)} checkpoint tensors skipped " - f"({sorted(set(ignored.values()))})." - ) super().load_weights( weights=text_weights, diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index befb9a45c4e3..03d56a617474 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -849,7 +849,9 @@ def test_sparse_view_writes_native_hnd_storage_without_copy(monkeypatch): @pytest.mark.cpu_only @pytest.mark.parametrize("main_dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) @pytest.mark.parametrize("index_dtype", ["bf16", "fp8"]) -def test_m3_budget_estimate_covers_actual_paged_buffers(main_dtype, index_dtype): +def test_m3_budget_estimate_covers_actual_paged_buffers( + main_dtype: DataType, index_dtype: str +) -> None: manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) manager.dtype = main_dtype manager.kv_cache_type = CacheType.SELF @@ -883,6 +885,7 @@ def test_m3_budget_estimate_covers_actual_paged_buffers(main_dtype, index_dtype) manager.reuse_match_backoff = 0 manager.enable_stats = False manager._has_cp_helix = False + manager.layer_specs = None config = manager._build_base_config( KvCacheConfig( diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py index 685254d40be6..0484c139fbee 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py @@ -14,9 +14,6 @@ # limitations under the License. import inspect -import math -import re -import struct from types import SimpleNamespace from unittest.mock import patch @@ -25,11 +22,8 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.auto_mapper import AutoCheckpointMapper -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( - MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper, - nvfp4_weight_scale_2, ) from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper from tensorrt_llm._torch.models.modeling_minimaxm3 import ( @@ -216,213 +210,6 @@ def test_load_weights_accepts_base_mapper_without_params_map() -> None: assert renamed[target_name] is bias -# --------------------------------------------------------------------------- -# Compressed-tensors NVFP4 MoE packaging of the MiniMax M3 architecture -# --------------------------------------------------------------------------- - -# Checkpoint naming of one routed expert in the multimodal checkpoint layout. -_EXPERT = "language_model.model.layers.3.block_sparse_moe.experts.7" - -# weight_scale_2 = fp32(fp32(bf16(2688 / gs) * 2688) * fp32(1 / 2688**2)) as -# float32 bit patterns. Computed outside PyTorch: Python float arithmetic -# rounded to float32 after every operation (exact for these operations), BF16 -# by round-to-nearest-even on the float32 bits. -# FP32 reciprocals of the stored weight_global_scale values. -_WEIGHT_SCALE_2_GOLDENS = { - 2688.0: 0x39C30C31, - 3000.0: 0x39AEC33E, - 2913.75: 0x39B3EF92, - 16999.5: 0x3876BB24, - 10000.0: 0x38D1B717, -} -_RECIPROCAL_GOLDENS = {1.0: 0x3F800000, 0.5: 0x40000000, 3.0: 0x3EAAAAAB} - - -def _f32_bits(value) -> int: - return struct.unpack(" dict[str, torch.Tensor]: - """Synthetic tensors of one packed NVFP4 expert, shaped like the checkpoint's.""" - generator = torch.Generator().manual_seed(0) - weights = {} - # (rows, packed bytes per row) for hidden 64, intermediate 32: w1/w3 are - # [inter, hidden], w2 is [hidden, inter]; two E2M1 values per byte. - for proj, (rows, packed_cols) in {"w1": (32, 32), "w2": (64, 16), "w3": (32, 32)}.items(): - prefix = f"{_EXPERT}.{proj}" - weights[f"{prefix}.weight_packed"] = torch.randint( - 0, 256, (rows, packed_cols), dtype=torch.uint8, generator=generator - ) - weights[f"{prefix}.weight_scale"] = ( - torch.rand(rows, packed_cols * 2 // 16, generator=generator) * 4 - ).to(torch.float8_e4m3fn) - weights[f"{prefix}.weight_global_scale"] = torch.tensor( - down_global_scale if proj == "w2" else gate_up_global_scale - ) - weights[f"{prefix}.input_global_scale"] = torch.tensor(input_global_scale) - return weights - - -@pytest.mark.parametrize("global_scale", sorted(_WEIGHT_SCALE_2_GOLDENS)) -def test_nvfp4_weight_scale_2_matches_independent_goldens(global_scale: float) -> None: - scale_2 = nvfp4_weight_scale_2(torch.tensor(global_scale)) - assert scale_2.dtype == torch.float32 - assert _f32_bits(scale_2) == _WEIGHT_SCALE_2_GOLDENS[global_scale] - - -def test_packed_nvfp4_expert_translates_to_nvfp4_moe_contract() -> None: - source = _packed_expert() - unrelated = { - "language_model.model.layers.3.self_attn.q_proj.weight": torch.ones(4, 4), - "language_model.model.layers.3.block_sparse_moe.gate.weight": torch.ones(4, 4), - } - out = MiniMaxM3HfWeightMapper().preprocess_weights({**source, **unrelated}) - - for name, value in unrelated.items(): - assert out[name] is value - for proj in ("w1", "w2", "w3"): - prefix = f"{_EXPERT}.{proj}" - for stale in ("weight_packed", "weight_global_scale", "input_global_scale"): - assert f"{prefix}.{stale}" not in out - # Packed FP4 bytes and E4M3 group scales are passed through, not copied or cast. - assert out[f"{prefix}.weight"] is source[f"{prefix}.weight_packed"] - assert out[f"{prefix}.weight"].dtype == torch.uint8 - assert out[f"{prefix}.weight_scale"] is source[f"{prefix}.weight_scale"] - expected = _WEIGHT_SCALE_2_GOLDENS[2913.75 if proj == "w2" else 3000.0] - assert _f32_bits(out[f"{prefix}.weight_scale_2"]) == expected - assert out[f"{prefix}.weight_scale_2"].shape == () - assert _f32_bits(out[f"{prefix}.input_scale"]) == _RECIPROCAL_GOLDENS[1.0] - assert len(out) == len(source) + len(unrelated) - - -@pytest.mark.parametrize("input_global_scale", sorted(_RECIPROCAL_GOLDENS)) -def test_packed_nvfp4_input_scale_is_the_reciprocal(input_global_scale: float) -> None: - out = MiniMaxM3HfWeightMapper().preprocess_weights( - _packed_expert(input_global_scale=input_global_scale) - ) - for proj in ("w1", "w2", "w3"): - assert ( - _f32_bits(out[f"{_EXPERT}.{proj}.input_scale"]) - == _RECIPROCAL_GOLDENS[input_global_scale] - ) - - -@pytest.mark.parametrize("kind", ["weight_global_scale", "input_global_scale"]) -@pytest.mark.parametrize("bad", [0.0, -1.0, math.nan, math.inf, -math.inf]) -def test_packed_nvfp4_rejects_invalid_global_scale(kind: str, bad: float) -> None: - weights = _packed_expert() - name = f"{_EXPERT}.w2.{kind}" - weights[name] = torch.tensor(bad) - with pytest.raises(ValueError, match=re.escape(name)): - MiniMaxM3HfWeightMapper().preprocess_weights(weights) - - -def test_packed_nvfp4_rejects_different_gate_up_weight_scales() -> None: - weights = _packed_expert() - weights[f"{_EXPERT}.w3.weight_global_scale"] = torch.tensor(3100.0) - with pytest.raises(ValueError, match=re.escape(f"{_EXPERT}.w3.weight_global_scale")): - MiniMaxM3HfWeightMapper().preprocess_weights(weights) - - -def test_canonical_nvfp4_expert_tensors_are_not_inverted_again() -> None: - canonical = {} - for proj in ("w1", "w2", "w3"): - prefix = f"model.layers.3.block_sparse_moe.experts.0.{proj}" - canonical[f"{prefix}.weight"] = torch.zeros(8, 32, dtype=torch.uint8) - canonical[f"{prefix}.weight_scale"] = torch.ones(8, 4).to(torch.float8_e4m3fn) - canonical[f"{prefix}.weight_scale_2"] = torch.tensor(1.0 / 3000.0) - canonical[f"{prefix}.input_scale"] = torch.tensor(0.25) - out = MiniMaxM3HfWeightMapper().preprocess_weights(canonical) - assert out is canonical - - -def test_packed_nvfp4_translation_is_idempotent() -> None: - mapper = MiniMaxM3HfWeightMapper() - once = mapper.preprocess_weights(_packed_expert()) - assert mapper.preprocess_weights(once) is once - - -def test_packed_nvfp4_translation_takes_over_consumable_weights() -> None: - source = _packed_expert() - packed = source[f"{_EXPERT}.w1.weight_packed"] - consumable = ConsumableWeightsDict(dict(source)) - out = MiniMaxM3HfWeightMapper().preprocess_weights(consumable) - assert isinstance(out, ConsumableWeightsDict) - # The translated dict now owns the tensors, so consuming them frees memory. - assert len(consumable) == 0 - assert out[f"{_EXPERT}.w1.weight"] is packed - assert out.mark_consumed(_EXPERT) == len(source) - - -def _bare_model(model_class: type) -> torch.nn.Module: - config = SimpleNamespace(num_key_value_heads=_NUM_KV_HEADS, num_attention_heads=8) - model = object.__new__(model_class) - torch.nn.Module.__init__(model) - model.model_config = ModelConfig(pretrained_config=config, mapping=Mapping()) - model.model = torch.nn.Module() - model.model.layers = torch.nn.ModuleList() - model.spec_config = None - if model_class is MiniMaxM3VLForConditionalGeneration: - model.vision_tower = None - model.last_loaded_vision_keys = [] - model.last_missing_vision_keys = [] - return model - - -def test_vl_load_weights_translates_packed_experts_before_loading() -> None: - model = _bare_model(MiniMaxM3VLForConditionalGeneration) - with patch.object(DecoderModelForCausalLM, "load_weights") as base_load_weights: - model.load_weights(_packed_expert(), weight_mapper=MiniMaxM3HfWeightMapper()) - - loaded = base_load_weights.call_args.kwargs["weights"] - expert = _EXPERT[len("language_model.") :] - assert sorted(loaded.keys()) == sorted( - f"{expert}.{proj}.{kind}" - for proj in ("w1", "w2", "w3") - for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") - ) - assert _f32_bits(loaded[f"{expert}.w1.weight_scale_2"]) == _WEIGHT_SCALE_2_GOLDENS[3000.0] - assert _f32_bits(loaded[f"{expert}.w2.weight_scale_2"]) == _WEIGHT_SCALE_2_GOLDENS[2913.75] - - -def test_vl_load_weights_names_the_checkpoint_tensor_of_an_invalid_scale() -> None: - model = _bare_model(MiniMaxM3VLForConditionalGeneration) - weights = _packed_expert() - name = f"{_EXPERT}.w2.input_global_scale" - weights[name] = torch.tensor(0.0) - with patch.object(DecoderModelForCausalLM, "load_weights") as base_load_weights: - with pytest.raises(ValueError, match=re.escape(name)): - model.load_weights(weights, weight_mapper=HfWeightMapper()) - base_load_weights.assert_not_called() - - -def test_text_load_weights_translates_packed_experts_with_any_mapper() -> None: - # The translation does not depend on the mapper driving the load, and the - # M3 params map still reaches the base loader. - model = _bare_model(MiniMaxM3ForCausalLM) - expert = _EXPERT[len("language_model.") :] - weights = {k[len("language_model.") :]: v for k, v in _packed_expert().items()} - mapper = HfWeightMapper() - with patch.object(DecoderModelForCausalLM, "load_weights") as base_load_weights: - model.load_weights(weights, weight_mapper=mapper) - - call_kwargs = base_load_weights.call_args.kwargs - assert call_kwargs["weight_mapper"] is mapper - assert call_kwargs["params_map"] == MINIMAX_M3_PARAMS_MAP - loaded = call_kwargs["weights"] - assert sorted(loaded.keys()) == sorted( - f"{expert}.{proj}.{kind}" - for proj in ("w1", "w2", "w3") - for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") - ) - assert loaded[f"{expert}.w3.weight"] is weights[f"{expert}.w3.weight_packed"] - assert _f32_bits(loaded[f"{expert}.w3.weight_scale_2"]) == _WEIGHT_SCALE_2_GOLDENS[3000.0] - - @pytest.mark.parametrize( ("model_class", "checkpoint_prefix"), [ diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 0774d3b986cd..dfacdd05707e 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -23,6 +23,7 @@ import json import os from types import SimpleNamespace +from typing import Any, Callable, Optional, Tuple from unittest.mock import Mock, create_autospec import pytest @@ -63,6 +64,9 @@ from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( MiniMaxM3HfWeightMapper, ) +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import ( + convert_compressed_tensors_nvfp4_weights, +) from tensorrt_llm._torch.models.modeling_minimaxm3 import ( MiniMaxM3Attention, MiniMaxM3DecoderLayer, @@ -1347,7 +1351,7 @@ def test_get_moe_layer_ids_length_mismatch_raises(): # --------------------------------------------------------------------------- -def _make_packaged_text_config(): +def _make_packaged_text_config() -> SimpleNamespace: """The trimmed text config as the compressed-tensors package declares it. Every attention layer is sparse and K-only while layers 0-2 keep a dense @@ -1363,7 +1367,9 @@ def _make_packaged_text_config(): return cfg -def _packaged_vl_config(text_config_as_dict: bool = False, **quant_overrides): +def _packaged_vl_config( + text_config_as_dict: bool = False, **quant_overrides: Any +) -> SimpleNamespace: quantization_config = { "quant_method": "mxfp8", "activation_scheme": "dynamic", @@ -1391,7 +1397,9 @@ def _packaged_vl_config(text_config_as_dict: bool = False, **quant_overrides): return cfg -def _expand_packaged_quant_config(pretrained_config): +def _expand_packaged_quant_config( + pretrained_config: SimpleNamespace, +) -> Tuple[QuantConfig, Optional[dict[str, QuantConfig]]]: """Parse the inline config as from_pretrained does, then apply the M3 hook.""" quant_config, layer_quant_config = ModelConfig.load_hf_quant_config( pretrained_config.quantization_config, "CUTLASS" @@ -3228,7 +3236,9 @@ def _cosine(out: torch.Tensor, ref: torch.Tensor) -> float: ) -def _graph_replay(fn, *inputs: torch.Tensor): +def _graph_replay( + fn: Callable[..., torch.Tensor], *inputs: torch.Tensor +) -> Callable[..., torch.Tensor]: """Capture ``fn(*inputs)`` in a CUDA graph; return a replay(new_inputs) helper.""" static_inputs = [t.clone() for t in inputs] stream = torch.cuda.Stream() @@ -3439,8 +3449,8 @@ def test_packaged_nvfp4_experts_load_and_run_on_cutlass(use_cuda_graph: bool, tu allow_backend_degradation=False, ) backend = moe.backend - assert type(backend).__name__ == "CutlassFusedMoE" - translated = MiniMaxM3HfWeightMapper().preprocess_weights(dict(checkpoint)) + assert type(backend).__name__ == "TrtllmCutlassNvfp4Impl" + translated = convert_compressed_tensors_nvfp4_weights(dict(checkpoint)) module_weights = { key[len(prefix) + 1 :]: value for key, value in translated.items() @@ -3561,11 +3571,11 @@ def forward(x: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: @pytest.mark.cpu_only -def test_minimax_m3_fold_gemma_boundary_weights_accepts_partial_lazy_load(): +def test_minimax_m3_fold_gemma_boundary_weights_accepts_partial_lazy_load() -> None: stored = torch.tensor([0.25, -0.25], dtype=torch.bfloat16) class LazyWeight: - def __getitem__(self, index): + def __getitem__(self, index: slice) -> torch.Tensor: assert index == slice(None) return stored @@ -3582,10 +3592,12 @@ def __getitem__(self, index): assert result["model.layers.2.self_attn.q_norm.weight"] is untouched -def _stub_minimax_m3_attention_construction(monkeypatch, *, is_neox=True): +def _stub_minimax_m3_attention_construction( + monkeypatch: pytest.MonkeyPatch, *, is_neox: bool = True +) -> Mock: """Keep the M3 constructor real while removing generic GPU setup and projection allocation.""" - def initialize_attention(self, **kwargs): + def initialize_attention(self: nn.Module, **kwargs: Any) -> None: nn.Module.__init__(self) config = kwargs["config"] self.pos_embd_params = SimpleNamespace(rope=kwargs["pos_embd_params"].rope, is_neox=is_neox) @@ -3594,7 +3606,7 @@ def initialize_attention(self, **kwargs): self.qkv_proj = nn.Module() self.qkv_proj.mapping = config.mapping - def make_linear(*args, **kwargs): + def make_linear(*args: Any, **kwargs: Any) -> nn.Module: projection = nn.Module() projection.quant_config = kwargs["quant_config"] return projection @@ -3610,7 +3622,9 @@ def make_linear(*args, **kwargs): @pytest.mark.cpu_only @pytest.mark.parametrize("entries", ["matching", "absent", "mismatched", "missing_k"]) -def test_minimax_m3_index_projection_quantization_follows_checkpoint_entries(monkeypatch, entries): +def test_minimax_m3_index_projection_quantization_follows_checkpoint_entries( + monkeypatch: pytest.MonkeyPatch, entries: str +) -> None: """index_qk_proj is quantized exactly when the checkpoint declares matching per-layer entries.""" _, model_config = _make_fused_qk_norm_rope_test_config() # Other-layer entries must not affect the projection being constructed. From 799153ca917abd299980647445707823ab5be64b Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 15:49:14 -0700 Subject: [PATCH 3/5] [None][test] MiniMax-M3: keep the decoder and Gemma-fold tests as on main The packaging change does not touch the decoder norms or the Gemma boundary fold, so their tests stay as main has them. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../integration/test_lists/test-db/l0_cpu.yml | 1 - .../unittest/_torch/models/test_minimax_m3.py | 75 +++++-------------- 2 files changed, 20 insertions(+), 56 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index d28ff03e8c3f..af9cefd82dd1 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -126,7 +126,6 @@ l0_cpu: - unittest/_torch/misc/test_xgrammar_structural_tag.py # MiniMax-M3 compressed-tensors NVFP4 MoE packaging: config expansion and # index-projection quantization. - - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fold_gemma_boundary_weights_accepts_partial_lazy_load - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_index_projection_quantization_follows_checkpoint_entries - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_expands_to_mixed_precision - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_honors_ignored_layers diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index dfacdd05707e..a34c262949be 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -76,7 +76,6 @@ MiniMaxM3QKVIndexerLinear, _build_swiglu_oai_dense_mlp, _dispatch_attention_over_live_tokens, - _fold_gemma_boundary_norm_weights, _get_index_projection_quant_config, _load_index_qk_proj_weights, _load_qkv_index_proj_weights, @@ -305,15 +304,21 @@ def test_minimax_m3_moe_reduces_only_local_terms( assert moe.allreduce.params == [None] -def _build_decoder_layer_with_stubs( +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "scheduler_kind, expected_post_fusion", + [ + (MoESchedulerKind.EXTERNAL_COMM, True), + (MoESchedulerKind.FUSED_COMM, False), + ], +) +def test_minimax_m3_decoder_layer_sets_post_fusion_from_moe_scheduler( monkeypatch: pytest.MonkeyPatch, scheduler_kind: MoESchedulerKind, - use_gemma_norm: bool, - *, - dtype: torch.dtype = torch.bfloat16, - hidden_size: int = 6144, -) -> MiniMaxM3DecoderLayer: - """A TP2 MoE decoder layer with real norms and stub attention, MoE and AllReduce.""" + expected_post_fusion: bool, +) -> None: + """A fused-communication MoE must not schedule a second boundary reduction.""" + monkeypatch.setenv("TRTLLM_MINIMAX_M3_EAGER_FUSION_DISABLED", "0") class _DecoderComponent(nn.Module): @@ -336,19 +341,20 @@ def __init__(self, *args: object, **kwargs: object) -> None: "tensorrt_llm._torch.models.modeling_minimaxm3.MiniMaxM3MoE", _DecoderMoE, ) + monkeypatch.setattr( + "tensorrt_llm._torch.models.modeling_minimaxm3.RMSNorm", + _DecoderComponent, + ) monkeypatch.setattr( "tensorrt_llm._torch.models.modeling_minimaxm3.AllReduce", _DecoderComponent, ) - config = _make_text_config() - config.use_gemma_norm = use_gemma_norm - config.torch_dtype = dtype - config.hidden_size = hidden_size model_config = ModelConfig( - pretrained_config=config, + pretrained_config=_make_text_config(), mapping=Mapping(world_size=2, rank=0, tp_size=2), ) - return MiniMaxM3DecoderLayer( + + layer = MiniMaxM3DecoderLayer( model_config=model_config, layer_idx=3, aux_stream_dict={ @@ -357,25 +363,6 @@ def __init__(self, *args: object, **kwargs: object) -> None: }, ) - -@pytest.mark.cpu_only -@pytest.mark.parametrize( - "scheduler_kind, expected_post_fusion", - [ - (MoESchedulerKind.EXTERNAL_COMM, True), - (MoESchedulerKind.FUSED_COMM, False), - ], -) -@pytest.mark.parametrize("use_gemma_norm", [False, True]) -def test_minimax_m3_decoder_layer_sets_post_fusion_from_moe_scheduler( - monkeypatch: pytest.MonkeyPatch, - scheduler_kind: MoESchedulerKind, - expected_post_fusion: bool, - use_gemma_norm: bool, -) -> None: - """Plain and folded-Gemma norms preserve fusion without a second routed reduction.""" - layer = _build_decoder_layer_with_stubs(monkeypatch, scheduler_kind, use_gemma_norm) - assert layer.enable_fusion assert layer.pre_feed_forward_fusion assert layer.post_feed_forward_fusion is expected_post_fusion @@ -3570,28 +3557,6 @@ def forward(x: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: assert _cosine(out, ref) > 0.999 -@pytest.mark.cpu_only -def test_minimax_m3_fold_gemma_boundary_weights_accepts_partial_lazy_load() -> None: - stored = torch.tensor([0.25, -0.25], dtype=torch.bfloat16) - - class LazyWeight: - def __getitem__(self, index: slice) -> torch.Tensor: - assert index == slice(None) - return stored - - untouched = LazyWeight() - weights = { - "model.layers.2.input_layernorm.weight": LazyWeight(), - "model.layers.2.self_attn.q_norm.weight": untouched, - } - result = _fold_gemma_boundary_norm_weights(weights) - assert result is weights - torch.testing.assert_close( - result["model.layers.2.input_layernorm.weight"], stored + 1.0, rtol=0, atol=0 - ) - assert result["model.layers.2.self_attn.q_norm.weight"] is untouched - - def _stub_minimax_m3_attention_construction( monkeypatch: pytest.MonkeyPatch, *, is_neox: bool = True ) -> Mock: From b3abe3a2fbad3b540221e0ca849e3a46a98a040f Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 16:03:52 -0700 Subject: [PATCH 4/5] [None][feat] MiniMax-M3: reject packaged recipes that half-quantize fused or head modules The packaged expansion rejects an untied lm_head left quantized and an ignore list that excludes only part of fused QKV or gate/up projections; the config tests run in the CPU stage under the cpu_only marker. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 36 +++++++++++++++++-- .../integration/test_lists/test-db/l0_cpu.yml | 2 ++ .../unittest/_torch/models/test_minimax_m3.py | 33 ++++++++++++++++- 3 files changed, 67 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 882572875d2d..9f15191f5cd2 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -22,8 +22,8 @@ from dataclasses import dataclass, field from pathlib import Path from types import SimpleNamespace -from typing import (TYPE_CHECKING, Any, Dict, Generic, Iterator, List, Optional, - Sequence, Tuple, TypeVar) +from typing import (TYPE_CHECKING, Any, Callable, Dict, Generic, Iterator, List, + Optional, Sequence, Tuple, TypeVar) import filelock import torch @@ -1086,6 +1086,14 @@ def _build_minimax_m3_mxfp8_moe_quant_config( 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): @@ -1113,7 +1121,19 @@ def _build_minimax_m3_mxfp8_moe_quant_config( f"{prefix}.mlp.{proj}": base_quant_config for proj in ("gate_proj", "up_proj", "down_proj") }) - excluded = mixed_quant_config.is_module_excluded_from_quantization + # 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) @@ -1125,6 +1145,16 @@ def _build_minimax_m3_mxfp8_moe_quant_config( 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, diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index af9cefd82dd1..85a3e7abae7d 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -129,6 +129,8 @@ l0_cpu: - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_index_projection_quantization_follows_checkpoint_entries - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_expands_to_mixed_precision - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_honors_ignored_layers + - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_rejects_a_quantized_head + - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_rejects_partly_excluded_fused_projections - unittest/_torch/models/test_minimax_m3.py::test_packaged_moe_config_rejects_unsupported_expert_recipes - unittest/_torch/models/test_minimax_m3.py::test_uniform_mxfp8_and_bf16_configs_are_unchanged - unittest/_torch/models/test_minimax_m3.py::test_modelopt_mixed_precision_map_is_normalized_as_before diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index a34c262949be..c6b28e838636 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -1398,6 +1398,7 @@ def _expand_packaged_quant_config( ) +@pytest.mark.cpu_only @pytest.mark.parametrize("text_config_as_dict", [False, True]) def test_packaged_moe_config_expands_to_mixed_precision(text_config_as_dict: bool) -> None: quant_config, layer_quant_config = _expand_packaged_quant_config( @@ -1428,6 +1429,7 @@ def test_packaged_moe_config_expands_to_mixed_precision(text_config_as_dict: boo assert {k: (v.quant_algo, v.group_size) for k, v in layer_quant_config.items()} == expected +@pytest.mark.cpu_only def test_packaged_moe_config_honors_ignored_layers() -> None: cfg = _packaged_vl_config() cfg.quantization_config["ignored_layers"].append( @@ -1439,6 +1441,7 @@ def test_packaged_moe_config_honors_ignored_layers() -> None: assert "model.layers.5.self_attn.o_proj" in layer_quant_config +@pytest.mark.cpu_only @pytest.mark.parametrize( "overrides", [ @@ -1450,11 +1453,37 @@ def test_packaged_moe_config_honors_ignored_layers() -> None: {"moe_quant_algo": None}, ], ) -def test_packaged_moe_config_rejects_unsupported_expert_recipes(overrides) -> None: +def test_packaged_moe_config_rejects_unsupported_expert_recipes( + overrides: dict[str, str | int | None], +) -> None: with pytest.raises(ValueError, match="routed-expert packaging"): _expand_packaged_quant_config(_packaged_vl_config(**overrides)) +@pytest.mark.cpu_only +def test_packaged_moe_config_rejects_a_quantized_head() -> None: + cfg = _packaged_vl_config() + cfg.quantization_config["ignored_layers"].remove("lm_head") + with pytest.raises(ValueError, match="must exclude lm_head"): + _expand_packaged_quant_config(cfg) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "module", + ["self_attn.q_proj", "mlp.gate_proj", "block_sparse_moe.shared_experts.up_proj"], +) +def test_packaged_moe_config_rejects_partly_excluded_fused_projections(module: str) -> None: + layer = 1 if module.startswith("mlp.") else 4 + cfg = _packaged_vl_config() + cfg.quantization_config["ignored_layers"].append( + f"language_model.model.layers.{layer}.{module}" + ) + with pytest.raises(ValueError, match="all or none of the fused projections"): + _expand_packaged_quant_config(cfg) + + +@pytest.mark.cpu_only def test_uniform_mxfp8_and_bf16_configs_are_unchanged() -> None: cfg = _packaged_vl_config() for key in ("moe_quant_algo", "moe_quant_format", "group_size"): @@ -1468,6 +1497,7 @@ def test_uniform_mxfp8_and_bf16_configs_are_unchanged() -> None: assert result[0] is no_quant and result[1] is None +@pytest.mark.cpu_only def test_modelopt_mixed_precision_map_is_normalized_as_before() -> None: base_config = QuantConfig(quant_algo=QuantAlgo.MXFP8, group_size=32) expert_config = QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16) @@ -1491,6 +1521,7 @@ def test_modelopt_mixed_precision_map_is_normalized_as_before() -> None: assert key not in out +@pytest.mark.cpu_only def test_index_projection_quant_config_follows_declared_entries() -> None: _, layer_quant_config = _expand_packaged_quant_config(_packaged_vl_config()) declared = _get_index_projection_quant_config( From ad05f4cc1f29fde02e70b15343f5132141d5017d Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 16:16:36 -0700 Subject: [PATCH 5/5] [None][test] MiniMax-M3: gate the packaged GPU tests by architecture The MXFP8 index-projection test runs on SM100/SM103, the CUTLASS NVFP4 MoE test on Blackwell and later. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tests/unittest/_torch/models/test_minimax_m3.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index c6b28e838636..eef025eb83fe 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -37,6 +37,7 @@ from torch.fx.experimental.symbolic_shapes import ShapeEnv from transformers import AutoConfig from utils.llm_data import llm_models_root +from utils.util import skip_blackwell_geforce, skip_pre_blackwell import tensorrt_llm._torch.models.modeling_minimaxm3 as modeling_minimaxm3 from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales @@ -3285,7 +3286,8 @@ def _mxfp8_dequant(weight: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tenso @pytest.mark.gpu -@pytest.mark.skipif(not _has_cuda(), reason="MXFP8 index projection needs CUDA") +@skip_pre_blackwell +@skip_blackwell_geforce @pytest.mark.parametrize("use_cuda_graph", [False, True]) def test_packaged_index_projection_keeps_mxfp8_scales_per_tp4_rank( use_cuda_graph: bool, monkeypatch: pytest.MonkeyPatch @@ -3401,7 +3403,7 @@ def _nvfp4_activation(x: torch.Tensor, global_scale: float = 1.0) -> torch.Tenso @pytest.mark.gpu -@pytest.mark.skipif(not _has_cuda(), reason="CUTLASS NVFP4 MoE needs CUDA") +@skip_pre_blackwell @pytest.mark.parametrize("use_cuda_graph", [False, True]) @pytest.mark.parametrize("tune", [False, True]) def test_packaged_nvfp4_experts_load_and_run_on_cutlass(use_cuda_graph: bool, tune: bool) -> None: