[None][feat] Enable MiniMax-M3 on Rubin (SM107): NVFP4 KV exemption and 128-token paged FMHA kernels

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

@@ -65,6 +65,11 @@ server (see [Launch the TensorRT LLM Server](#launch-the-tensorrt-llm-server)).
 * **Weight and KV precision.** BF16, MXFP8, and mixed MXFP8/NVFP4 checkpoints are supported. On SM100/SM103, the MSA backend supports FP8 and NVFP4 KV cache. NVFP4 storage is used in sparse target layers; dense target layers and the shared Eagle3 draft layer use FP8. The curated BF16-KV YAML remains a separate starting configuration.
 * **KV cache reuse must be disabled.** KV cache reuse is not supported on the sparse-attention path, so set `kv_cache_config.enable_block_reuse: false`.
 * **Eagle3.** The MSA path supports one shared GQA draft layer with linear Eagle3 drafting. Other MTP and tree-drafting configurations are unsupported. Keep NVFP4 KV and piecewise CUDA graph qualification separate: the existing PCG accuracy cases use FP8 KV.
+* **Rubin (SM107) Eagle3 draft layer.** On SM107 the shared Eagle3 draft layer runs at 128-token pages through FlashInfer's TRTLLM-Gen FMHA, the only FMHA with kernels for that page size there: keep `flashinfer_trtllm_gen` enabled in `TLLM_FMHA_LIBS`. `flashinfer-python` downloads those kernels on first use. On nodes without internet access, or to keep the download off the serving path, pre-populate the cubin directory and serve offline:
+  1. On one node with internet access, start the server with the production configuration (same tensor-parallel size and KV cache dtype: the kernels depend on the per-rank head layout and KV dtype), `FLASHINFER_CUBIN_DIR` set to an empty directory, and serve one Eagle3 request with a long prompt (a context and a generation step at 128-token pages).
+  2. Copy that directory to each serving node, set `FLASHINFER_CUBIN_DIR` to it and `FLASHINFER_NO_DOWNLOAD=1`; see [FlashInfer JIT workspaces for MPI workers](../llm-api/index.md#flashinfer-jit-workspaces-for-mpi-workers).
+
+  For the FP8 KV draft layer at TP2 (16 query heads and 1 KV head per rank, head size 128) the directory holds, under `<artifact hash>/fmha/trtllm-gen/`, `include/flashInferMetaInfo.h`, `checksums.txt` and 10 cubins: the `fmhaSm107aKernel_QkvE4m3OBfloat16H128PagedKvCausalPaged` context kernel `VarSeqQ256Kv128PersistentContext` and three generation kernels (`MultiCtasKvCgaVarSeqQ64Kv256StaticGroupedKeepsAbForGen`, `MultiCtasKvVarSeqQ64Kv256StaticGroupedKeepsAbForGen`, `VarSeqQ64Kv256PersistentGroupedKeepsAbForGen`), plus six `fmhaSm100fKernel_QkvE4m3OBfloat16H128PagedKvCausalPaged` generation kernels (`Q16` and `Q32` variants).
 * **Use matching builds for disaggregated context and generation workers.** The Triton unquantized main KV cache uses HND physical storage. Builds using the previous NHD storage are not compatible for KV transfer; upgrade both workers together.
 * **Size `max_seq_len` for the active attention paths.** On SM100/SM103, supported dense layers with unquantized KV cache in the Triton configuration use native paged attention and do not gather historical KV or expand GQA heads into temporary tensors. The SDPA dense fallback still allocates those buffers when TRTLLM-Gen is unavailable, disabled, ordered after the C++ fallback in `TLLM_FMHA_LIBS`, or the Triton configuration uses FP8/NVFP4 KV. FP8 KV retains SDPA to preserve its existing Q precision. The Triton sparse decode path separately has per-query FP32 temporary buffers that grow with the attended KV length. Keep `max_seq_len` bounded for paths with these allocations: checkpoint defaults (1M for MXFP8, 512K for BF16) can cause OOM or CUDA graph capture failures. The curated YAML sets `max_seq_len` just above ISL+OSL (`2068` for the 1k/1k benchmark); longer-context workloads may require a lower `max_batch_size`. MSA retains its own dense/sparse kernels and cache layout.
 * **Parallelism.** MoE experts run with expert parallelism. The attention layers support both Tensor-Expert Parallelism (TEP) and Data-Expert Parallelism (DEP, via `enable_attention_dp: true`). The overlap scheduler is enabled by default.

tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py

@@ -369,6 +369,10 @@ class FlashInferTrtllmGenFmha(PhasedFmha):
     MAX_HEADS_RATIO_GENERATION = 32
     MIN_TOKENS_PER_BLOCK = 8
     SUPPORTED_TOKENS_PER_BLOCK = {16, 32, 64}
+    # SMs on which this library serves only a KV cache manager's extra page
+    # sizes (``trtllm_gen_extra_tokens_per_block``); every other request there
+    # stays on the fallback FMHA.
+    EXTRA_PAGE_SIZES_ONLY_SMS = frozenset({107})
     # FlashInfer narrows key_cache.size(0) to int before constructing its TMA descriptor.
     MAX_NUM_PAGES_IN_MEM_POOL = (1 << 31) - 1
     SUPPORTED_MLA_GENERATION_HEAD_DIMS = {
@@ -407,14 +411,20 @@ def is_environment_available(cls) -> bool:
             return False
 
         sm = get_sm_version()
-        if sm not in (100, 103):
+        if sm not in (100, 103, 107):
             logger.debug(
-                f"FlashInfer TRTLLM-Gen FMHA is unavailable: requires SM100 or SM103, got SM{sm}."
+                "FlashInfer TRTLLM-Gen FMHA is unavailable: "
+                f"requires SM100, SM103 or SM107, got SM{sm}."
             )
             return False
 
         return True
 
+    @classmethod
+    def serves_only_extra_page_sizes(cls) -> bool:
+        """Whether this device uses the library only for a KV cache manager's extra page sizes."""
+        return get_sm_version() in cls.EXTRA_PAGE_SIZES_ONLY_SMS
+
     @classmethod
     def _is_available(cls, attn: "TrtllmAttention") -> bool:
         if not cls.is_environment_available():
@@ -790,12 +800,21 @@ def _is_supported_with_reason(
         if tokens_per_block & (tokens_per_block - 1) != 0:
             return False, f"tokens_per_block ({tokens_per_block}) that is not a power of 2."
         # A KV cache manager may allow extra page sizes, e.g. MiniMax-M3 adds 128.
-        supported_tokens_per_block = self.SUPPORTED_TOKENS_PER_BLOCK | set(
+        extra_tokens_per_block = frozenset(
             getattr(meta.kv_cache_manager, "trtllm_gen_extra_tokens_per_block", ())
         )
+        supported_tokens_per_block = self.SUPPORTED_TOKENS_PER_BLOCK | extra_tokens_per_block
         if tokens_per_block not in supported_tokens_per_block:
             supported = sorted(supported_tokens_per_block)
             return False, f"tokens_per_block ({tokens_per_block}). Supported: {supported}."
+        # On SM107 this library serves only the extra page sizes, which the
+        # fallback thop FMHA has no paged-context kernels for. Every other SM107
+        # request stays on the fallback.
+        if self.serves_only_extra_page_sizes() and tokens_per_block not in extra_tokens_per_block:
+            return False, (
+                f"tokens_per_block ({tokens_per_block}) on SM107, which uses this library "
+                "only for the KV cache manager's extra page sizes."
+            )
 
         return True, ""
 

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

@@ -38,6 +38,7 @@
     TensorWrapper,
     binding_to_torch_dtype,
     convert_to_torch_tensor,
+    get_sm_version,
     prefer_pinned,
 )
 from tensorrt_llm.bindings import DataType
@@ -227,6 +228,28 @@ def extend_attention_op_pools_for_shared_draft_layers(
     )
 
 
+def _require_trtllm_gen_extra_page_route(tokens_per_block: int) -> None:
+    """Fail at setup, not at the first prefill, when the P128 opt-in has no FMHA (SM107).
+
+    On SM107 only FlashInfer TRTLLM-Gen has paged-context kernels for the extra
+    page sizes (its cubins ship with it), so the route exists iff it is enabled.
+    """
+    if get_sm_version() != 107:
+        return
+    # Imported here: the FMHA backends import the attention classes this module serves.
+    from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import (
+        FlashInferTrtllmGenFmha,
+    )
+    from tensorrt_llm._torch.attention.backends.fmha.registry import get_enabled_fmha_lib_classes
+
+    if FlashInferTrtllmGenFmha not in get_enabled_fmha_lib_classes():
+        raise RuntimeError(
+            f"SM107 with {tokens_per_block}-token KV pages needs the flashinfer_trtllm_gen FMHA "
+            "library, the only one with paged-context kernels for this page size here, but "
+            "TLLM_FMHA_LIBS excludes it. Keep flashinfer_trtllm_gen in TLLM_FMHA_LIBS (or unset it)."
+        )
+
+
 class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2):
     """KVCacheManagerV2 subclass with a V2-managed paged index-K cache
     per sparse layer.
@@ -399,6 +422,18 @@ def __init__(
                 self._block_scale_buffers[layer_idx, kv_layout] = self.get_block_scale_buffers(
                     layer_idx, kv_layout
                 )
+        # The KV pools are written only at live slots, and readers that consume a whole block or page (the CSR
+        # sparse prefill's V dequant, the full-context decode on a last page's tail) would turn uninitialized
+        # bytes into 0 * NaN = NaN. Zero every layer's K/V bytes (NVFP4 sparse, FP8 dense, the drafter's in its
+        # own manager), the NVFP4 block scales and the paged index keys once, before any request runs: every
+        # byte a kernel writes is unchanged.
+        for view in (
+            *self._kv_buffers.values(),
+            *self._block_scale_buffers.values(),
+            *self._index_k_buffers.values(),
+        ):
+            if view is not None:
+                view.zero_()
 
     def _validate_speculative_config(self, spec_config: DecodingBaseConfig | None) -> None:
         # Dynamic-tree relocation moves packed K/V bytes without their NVFP4
@@ -540,7 +575,16 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
             for layer_idx in (*self._dense_native_layer_ids, *self._shared_draft_layer_ids)
             if layer_idx in self.layer_offsets
         ]
-        if self.is_draft or not native_layers:
+        if self.is_draft:
+            # A separate draft manager holds only draft layers, which all run
+            # native attention and use their own pool tables; it needs the P128
+            # opt-in but no virtual pools. (MiniMax-M3 Eagle3 shares the target
+            # manager today; this keeps a separate one on the trtllm-gen route.)
+            if self.tokens_per_block == 128 and self.dtype != DataType.NVFP4:
+                _require_trtllm_gen_extra_page_route(128)
+                self.trtllm_gen_extra_tokens_per_block = frozenset({128})
+            return
+        if not native_layers:
             return
         if self.enable_swa_scratch_reuse:
             raise NotImplementedError(
@@ -552,8 +596,13 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
             return
         # Native layers run at the target's 128-token pages. trtllm-gen has P128
         # kernels for their dense-GQA shapes but not for every shape, so opt in
-        # here rather than in the global allowlist.
+        # here rather than in the global allowlist. On SM107 this opt-in is the
+        # only route for these layers: FallbackFmha has no P128 paged-context kernels.
         if self.tokens_per_block == 128:
+            # On SM107 the dense layers keep SDPA (see the model's native dense attention
+            # setup), so only the shared draft layers need the route there.
+            if any(i in self.layer_offsets for i in self._shared_draft_layer_ids):
+                _require_trtllm_gen_extra_page_route(128)
             self.trtllm_gen_extra_tokens_per_block = frozenset({128})
         if self._use_per_layer_page_tables:
             # Per-layer page tables already give every layer its own pool.
@@ -1031,6 +1080,8 @@ def __init__(self, manager: MiniMaxM3KVCacheManagerV2, draft_layer_ids: Sequence
                 "MiniMax-M3 draft KV cache views support exactly one local shared draft layer; "
                 f"got {len(draft_layer_ids)}. Multiple layers require separate pool roots."
             )
+        if manager.tokens_per_block in self.trtllm_gen_extra_tokens_per_block:
+            _require_trtllm_gen_extra_page_route(manager.tokens_per_block)
         self._manager = manager
         layer_id = draft_layer_ids[0]
         local = int(manager.layer_offsets[layer_id])

tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py

@@ -799,8 +799,8 @@ def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> i
     return 1 << (target.bit_length() - 1)
 
 
-# NVFP4 instructions are supported on both datacenter Blackwell targets.
-_SM100F_CAPABILITIES = ((10, 0), (10, 3))
+# NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets.
+_SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7))
 # Shape-specific tuning was measured on GB300 only. SM100 uses the defaults.
 _NVFP4_TUNED_CAPABILITIES = ((10, 3),)
 
@@ -1063,7 +1063,7 @@ def minimax_m3_sparse_attn_decode(
         capability = torch.cuda.get_device_capability(q.device)
         if capability not in _SM100F_CAPABILITIES:
             raise NotImplementedError(
-                "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103; "
+                "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103/SM107; "
                 f"got SM{capability[0] * 10 + capability[1]}."
             )
         _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols)

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -30,6 +30,7 @@
 from torch.nn.attention import SDPBackend, sdpa_kernel
 from transformers import PretrainedConfig
 
+from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.bindings import DataType
 from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType
 from tensorrt_llm.logger import logger
@@ -1299,6 +1300,11 @@ def _init_native_dense_attention(self, model_config: ModelConfig) -> None:
             return
         if not FlashInferTrtllmGenFmha.is_environment_available():
             return
+        if get_sm_version() == 107:
+            # SM107 admits the library only for the KV manager's extra page sizes
+            # (the Eagle3 draft layers at 128-token pages). Dense layers keep SDPA
+            # there, as before SM107 was admitted.
+            return
         # Q/K norm and partial RoPE run in the model before backend dispatch.
         native_attn = TrtllmAttention(
             self.layer_idx,

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -56,7 +56,7 @@
                                                maybe_create_moe_load_balancer)
 from ..virtual_memory import RestoreMode
 from ..virtual_memory import scope as virtual_memory_scope
-from .config_utils import (is_hybrid_linear, is_mla,
+from .config_utils import (is_hybrid_linear, is_minimax_m3, is_mla,
                            resolve_auto_ssm_cache_dtype,
                            supports_fp4_mla_attention, uses_fp4_mla_attention,
                            validate_kimi_kda_state_dtype)
@@ -165,9 +165,15 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig,
 
     effective_kv_cache_quant = (kv_cache_quant if pyt_kv_cache_dtype == "auto"
                                 else mapped_pyt_quant)
+    # MiniMax-M3 MSA keeps NVFP4 only on its sparse layers, which MSA and
+    # Triton serve; its dense and draft layers stay FP8 for trtllm-gen.
+    sparse_config = model_config.sparse_attention_config
+    uses_minimax_m3_msa = (is_minimax_m3(sparse_config)
+                           and sparse_config.implementation == "msa")
     if (effective_kv_cache_quant in (QuantAlgo.NVFP4, QuantAlgo.NVFP4.value)
             and not supports_fp4_mla_attention(model_config)
-            and torch.cuda.is_available() and get_sm_version() == 107):
+            and not uses_minimax_m3_msa and torch.cuda.is_available()
+            and get_sm_version() == 107):
         logger.warning(
             "NVFP4 KV cache is not supported by trtllm-gen on SM107; "
             "using FP8 KV cache instead.")

tests/integration/test_lists/test-db/l0_gr100.yml

@@ -14,3 +14,4 @@ l0_gr100:
       backend: pytorch
   tests:
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True]
+  - unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py

tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_pool_zero_init.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MiniMax-M3 KV pools start zeroed.
+
+Readers that consume a whole block or page (the CSR sparse prefill's V dequant,
+the full-context decode on a last page's tail) also read bytes no request has
+written, so every pool view must hold zeros, not uninitialized memory, before
+the first request.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import (
+    MiniMaxM3KVCacheManagerV2,
+)
+from tensorrt_llm.bindings import DataType
+from tensorrt_llm.bindings.internal.batch_manager import CacheType
+from tensorrt_llm.llmapi.llm_args import (
+    Eagle3DecodingConfig,
+    KvCacheConfig,
+    MiniMaxM3SparseAttentionConfig,
+)
+from tensorrt_llm.mapping import Mapping
+
+
+def _nonzero_bytes(view: torch.Tensor) -> int:
+    """Count nonzero bytes; a strided view that cannot be reinterpreted in place is copied first."""
+    try:
+        raw = view.view(torch.uint8)
+    except RuntimeError:
+        raw = view.contiguous().view(torch.uint8)
+    return int(torch.count_nonzero(raw).item())
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
+@pytest.mark.parametrize("with_draft", [False, True], ids=["target_only", "shared_eagle3"])
+def test_every_pool_view_is_zero_after_construction(with_draft: bool) -> None:
+    torch.cuda.init()
+    manager = MiniMaxM3KVCacheManagerV2(
+        KvCacheConfig(
+            max_gpu_total_bytes=16 << 20,
+            host_cache_size=0,
+            enable_block_reuse=False,
+            enable_swa_scratch_reuse=False,
+        ),
+        CacheType.SELF,
+        num_layers=4,
+        num_kv_heads=[2, 2, 2, 2, 1] if with_draft else [2, 2, 2, 2],
+        head_dim=128,
+        tokens_per_block=128,
+        max_seq_len=256,
+        max_batch_size=2,
+        max_num_tokens=256,
+        mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1),
+        dtype=DataType.NVFP4,
+        vocab_size=16,
+        spec_config=(
+            Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused")
+            if with_draft
+            else None
+        ),
+        sparse_attention_config=MiniMaxM3SparseAttentionConfig(
+            implementation="msa", indexer_kv_dtype="fp8"
+        ),
+    )
+    try:
+        views = {
+            **{("kv",) + key: view for key, view in manager._kv_buffers.items()},
+            **{("scale",) + key: view for key, view in manager._block_scale_buffers.items()},
+            **{("index_k", key): view for key, view in manager._index_k_buffers.items()},
+        }
+        assert any(name[0] == "kv" for name in views)
+        assert any(name[0] == "scale" for name in views)
+        assert any(name[0] == "index_k" and view is not None for name, view in views.items())
+        torch.cuda.synchronize()
+        counts = {name: _nonzero_bytes(view) for name, view in views.items() if view is not None}
+        nonzero = {name: count for name, count in counts.items() if count}
+        assert not nonzero, f"pool views with nonzero bytes after construction: {nonzero}"
+    finally:
+        manager.shutdown()