From 2a58057b5ce1a8aedca846292a69454b407b97f1 Mon Sep 17 00:00:00 2001
From: peihengh <259410613+peihu-nv@users.noreply.github.com>
Date: Mon, 19 Oct 2026 17:37:34 -0700
Subject: [PATCH 1/9] [None][feat] Enable MiniMax-M3 on Rubin (SM107): NVFP4 KV
 exemption and 128-token paged FMHA kernels

- trtllm-gen FMHA: on SM107, generate dense paged-KV context and
  KeepsMmaAb generation kernels with NVRTC for 128-token pages
  (matched BF16/E4M3 Q and KV, BF16 output); the cubin table has none.
- FallbackFmha: on SM107 (no TRTLLM-Gen FMHA), let the thop paged-context
  kernel serve the 128-token pages MiniMax-M3's KV managers opt in;
  every other SM keeps the refusal.
- MiniMax-M3 sparse decode: admit SM107 for the NVFP4 Triton path.
- validate_and_set_kv_cache_quant: keep NVFP4 KV for MiniMax-M3 MSA on
  SM107 instead of downgrading to FP8; its sparse layers are served by
  MSA/Triton and its dense and draft layers stay FP8.
- MiniMax-M3 KV cache manager: zero every KV pool view once at
  allocation (TRTLLM_M3_KV_POOL_ZERO_INIT=0 disables it), so a context
  step that reads a request's whole last page never sees uninitialized
  bytes in its unwritten tail.
- Tests for these changes, including an SM107 forward test of cached-prefix
  and chunked context at 128-token pages against the unchunked run (BF16
  and FP8 KV), and pin a non-SM107 device in the explicit KV-dtype
  override test.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
---
 .../trtllmGenKernels/fmha/fmhaKernels.h       |  17 +-
 .../attention/backends/fmha/fallback.py       |  15 ++
 .../sparse/minimax_m3/cache_manager.py        |  19 ++
 .../kernels/triton_sparse_decode.py           |   6 +-
 .../_torch/pyexecutor/model_loader.py         |  10 +-
 .../msa/test_minimax_m3_kv_pool_zero_init.py  |  93 +++++++
 .../msa/test_minimax_m3_sparse_attn_decode.py |  10 +-
 .../attention/sparse/msa/test_msa_backend.py  |  30 ++-
 .../test_context_fmha_kernel_presence.py      |  31 +++
 .../test_fallback_fmha_extra_page_sizes.py    | 251 ++++++++++++++++++
 tests/unittest/_torch/test_model_config.py    |  33 ++-
 11 files changed, 496 insertions(+), 19 deletions(-)
 create mode 100644 tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_pool_zero_init.py
 create mode 100644 tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py

diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h
index 79cc1eb5dd7e..a0c192d12a1c 100644
--- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h
+++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h
@@ -254,7 +254,7 @@ class TllmGenFmhaKernel
         }
     }
 
-    static bool shouldUseNvrtc(FmhaOptions const& options)
+    bool shouldUseNvrtc(FmhaOptions const& options) const
     {
 #ifdef TLLM_RUBIN_FEATURES
         // Spcompress kernels must use precompiled cubins.
@@ -263,7 +263,20 @@ class TllmGenFmhaKernel
             return false;
         }
 #endif // TLLM_RUBIN_FEATURES
-       // Sparse MQA/GQA uses NVRTC path for now because no model really uses it.
+       // The cubin table has no dense paged-KV kernels for 128-token pages, and the page size is
+       // part of the kernel hash. On SM107, generate those BF16/E4M3 context and KeepsMmaAb
+       // generation kernels for the actual cache layout instead.
+        constexpr int kPageSize128 = 128;
+        if (mSM == kSM_107
+            && (isContextKernel(options.mFmhaKernelType) || isKeepsMmaAbForGenerationKernel(options.mFmhaKernelType))
+            && !options.mIsMlaGen && !isTokenSparse(options.mSparseType) && isPagedKv(options.mQkvLayout)
+            && options.mNumTokensPerPage == kPageSize128 && options.mDtypeQ == options.mDtypeKv
+            && (options.mDtypeQ == tg::Dtype::Bfloat16 || options.mDtypeQ == tg::Dtype::E4m3)
+            && options.mDtypeOut == tg::Dtype::Bfloat16)
+        {
+            return true;
+        }
+        // Sparse MQA/GQA uses NVRTC path for now because no model really uses it.
         if (isStaticTokenSparse(options.mSparseType) && !options.mIsMlaGen)
         {
             return true;
diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py
index 448d64df9925..be974499b494 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py
@@ -53,6 +53,20 @@
     }
 )
 
+# SM107 has no TRTLLM-Gen FMHA, so the thop paged-context kernel serves every
+# context phase there, at the 128-token pages MiniMax-M3's KV managers opt in
+# (as before those managers listed 128 as a TRTLLM-Gen extra page size). The
+# SM107 trtllm-gen runner generates those kernels with NVRTC (fmhaKernels.h);
+# test_context_fmha_kernel_presence.py checks that the build reports them
+# present. Every other SM keeps the refusal.
+_THOP_PAGED_CONTEXT_EXTRA_PAGES_SMS: frozenset = frozenset({107})
+
+
+def _thop_serves_extra_page_sizes() -> bool:
+    """Whether the thop paged-context kernel serves a manager's TRTLLM-Gen extra page sizes on this device."""
+    return get_sm_version() in _THOP_PAGED_CONTEXT_EXTRA_PAGES_SMS
+
+
 # ``thop.attention`` kwargs hard-wired to a literal at the call site (no
 # rich object owns them). Sync test enforces both the kwarg name and the
 # literal value.
@@ -255,6 +269,7 @@ def _is_supported(
                 and phase != FmhaPhase.GENERATION
                 and metadata.use_paged_context_fmha
                 and metadata.tokens_per_block in extra_page_sizes
+                and not _thop_serves_extra_page_sizes()
             ):
                 raise RuntimeError(
                     f"FallbackFmha cannot serve paged prefill with "
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..ebd8432a5a78 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
@@ -28,6 +28,7 @@
 
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple
 
 import torch
@@ -227,6 +228,11 @@ def extend_attention_op_pools_for_shared_draft_layers(
     )
 
 
+# Zero every KV pool view (NVFP4 + FP8 K/V, NVFP4 block scales, paged index keys) once at allocation, so no reader
+# ever sees uninitialized bytes. Kill switch (diagnostics only): TRTLLM_M3_KV_POOL_ZERO_INIT=0.
+_KV_POOL_ZERO_INIT = os.environ.get("TRTLLM_M3_KV_POOL_ZERO_INIT", "1") == "1"
+
+
 class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2):
     """KVCacheManagerV2 subclass with a V2-managed paged index-K cache
     per sparse layer.
@@ -399,6 +405,19 @@ def __init__(
                 self._block_scale_buffers[layer_idx, kv_layout] = self.get_block_scale_buffers(
                     layer_idx, kv_layout
                 )
+        if _KV_POOL_ZERO_INIT:
+            # 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
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py
index 071f907067b4..0150347026a1 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py
+++ b/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)
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 9f2f4e12a855..e11d3f87d335 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -31,6 +31,7 @@
 from tensorrt_llm._utils import get_sm_version, str_dtype_to_torch
 from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig,
                                           ExecutorMemoryType,
+                                          MiniMaxM3SparseAttentionConfig,
                                           ModelExpressConfig,
                                           SparseAttentionConfig, TorchLlmArgs)
 from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto,
@@ -165,9 +166,16 @@ 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 = (isinstance(sparse_config,
+                                      MiniMaxM3SparseAttentionConfig)
+                           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.")
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_pool_zero_init.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_pool_zero_init.py
new file mode 100644
index 000000000000..988df58473f0
--- /dev/null
+++ b/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()
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py
index 88760a79452f..07376d86d2d7 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py
@@ -1158,7 +1158,7 @@ def test_resolve_num_topk_chunks_is_shape_only_power_of_two(total_q, num_kv_head
         (33, {}),
     ],
 )
-@pytest.mark.parametrize("capability", [(10, 0), (10, 3)])
+@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)])
 def test_sm100f_nvfp4_launch_options(local_batch, expected, capability):
     if capability != (10, 3):
         expected = {}
@@ -1192,7 +1192,7 @@ def test_sm100f_nvfp4_launch_options(local_batch, expected, capability):
         (15, None),
     ],
 )
-@pytest.mark.parametrize("capability", [(10, 0), (10, 3)])
+@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)])
 def test_sm100f_nvfp4_num_topk_chunks(local_batch, expected, capability):
     if capability != (10, 3):
         expected = None
@@ -1213,7 +1213,7 @@ def test_sm100f_nvfp4_num_topk_chunks(local_batch, expected, capability):
     ("local_batch", "expected"),
     [(10, 1), (11, 2), (12, 2), (14, 2), (15, 1)],
 )
-@pytest.mark.parametrize("capability", [(10, 0), (10, 3)])
+@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)])
 def test_sm100f_nvfp4_query_group_size_is_narrowly_scoped(local_batch, expected, capability):
     if capability != (10, 3):
         expected = 1
@@ -1259,7 +1259,7 @@ def test_sm100f_nvfp4_launch_options_rejects_unmeasured_shapes():
         (8, 14, False),
     ],
 )
-@pytest.mark.parametrize("capability", [(10, 0), (10, 3)])
+@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)])
 def test_sm100f_nvfp4_linear_softmax_policy(num_kv_heads, local_batch, expected, capability):
     if capability != (10, 3):
         expected = False
@@ -1290,7 +1290,7 @@ def test_sm100f_nvfp4_linear_softmax_policy(num_kv_heads, local_batch, expected,
         (12, {}),
     ],
 )
-@pytest.mark.parametrize("capability", [(10, 0), (10, 3)])
+@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)])
 def test_sm100f_nvfp4_merge_launch_options_is_narrowly_scoped(local_batch, expected, capability):
     if capability != (10, 3):
         expected = {}
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
index 8c83072141c8..eb93e3d2f135 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
@@ -258,7 +258,7 @@ def test_nvfp4_sparse_decode_rejects_unsupported_device(
     block_scale = torch.empty((1, 1, 128, 8), dtype=torch.uint8)
     global_scale = torch.ones(1)
 
-    with pytest.raises(NotImplementedError, match="NVFP4 sparse decode requires SM100/SM103"):
+    with pytest.raises(NotImplementedError, match="NVFP4 sparse decode requires SM100/SM103/SM107"):
         minimax_m3_sparse_attn_decode(
             q,
             kv,
@@ -277,6 +277,34 @@ def test_nvfp4_sparse_decode_rejects_unsupported_device(
     get_capability.assert_called_once_with(q.device)
 
 
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)])
+def test_nvfp4_sparse_decode_admits_the_sm100_family(
+    monkeypatch: pytest.MonkeyPatch, capability: tuple[int, int]
+) -> None:
+    """SM100-family devices, Rubin included, pass the device check and reach
+    input validation, which here rejects the missing scales."""
+    get_capability = Mock(spec_set=torch.cuda.get_device_capability, return_value=capability)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", get_capability)
+    q = torch.empty((1, 8, 128), dtype=torch.bfloat16)
+    kv = torch.empty((1, 1, 128, 64), dtype=torch.uint8)
+
+    with pytest.raises(ValueError, match="needs all four"):
+        minimax_m3_sparse_attn_decode(
+            q,
+            kv,
+            kv,
+            torch.zeros((1, 1, 1), dtype=torch.int32),
+            torch.zeros((1, 1), dtype=torch.int32),
+            torch.ones(1, dtype=torch.int32),
+            sm_scale=128**-0.5,
+            output=torch.empty_like(q),
+            decode_query_len=1,
+            k_block_scale=torch.empty((1, 1, 128, 8), dtype=torch.uint8),
+        )
+    get_capability.assert_called_once_with(q.device)
+
+
 @pytest.mark.cpu_only
 @pytest.mark.parametrize("indexer_kv_dtype", ["bf16", "fp8"])
 @pytest.mark.parametrize("implementation", ["msa", "triton"])
diff --git a/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py b/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py
index a9e1c62b1e0e..83ab912fdad6 100644
--- a/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py
+++ b/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py
@@ -35,6 +35,10 @@
 # A head width the SM100-family kernel table does carry for 16-bit paged-context
 # attention.
 _SUPPORTED_HEAD_SIZE = 64
+# The page size and head width of MiniMax-M3's dense and draft attention. The
+# kernel table carries no 128-token-page kernels; SM107 generates them at runtime.
+_NVRTC_TOKENS_PER_BLOCK = 128
+_NVRTC_HEAD_SIZE = 128
 
 _NO_GPU = not torch.cuda.is_available()
 _NO_GPU_REASON = "needs a GPU; missing bindings on a GPU are a test failure"
@@ -188,3 +192,30 @@ def test_lookup_is_page_size_sensitive():
     )
     assert present
     assert not odd_page
+
+
+@pytest.mark.parametrize("kv_cache_dtype", [DataType.BF16, DataType.FP8], ids=["bf16", "fp8"])
+def test_lookup_reports_present_for_sm107_128_token_pages(kv_cache_dtype):
+    """SM107 serves dense paged context over 128-token pages with BF16 output.
+
+    No precompiled kernel exists for that page size, so matched BF16/E4M3 Q
+    and KV with BF16 output must be generated at runtime rather than reported
+    absent, which would refuse paged-context attention for such models. The
+    runtime path is limited to that output precision: FP8 output at the same
+    page size has no kernel either way and must still report absent.
+    """
+    if get_sm_version() != 107:
+        pytest.skip("runtime-generated 128-token-page context kernels are specific to SM107")
+    assert thop.fused_context_fmha_kernel_exists(
+        head_size=_NVRTC_HEAD_SIZE,
+        kv_cache_dtype=kv_cache_dtype,
+        tokens_per_block=_NVRTC_TOKENS_PER_BLOCK,
+        output_dtype=DataType.BF16,
+    )
+    if kv_cache_dtype == DataType.FP8:
+        assert not thop.fused_context_fmha_kernel_exists(
+            head_size=_NVRTC_HEAD_SIZE,
+            kv_cache_dtype=kv_cache_dtype,
+            tokens_per_block=_NVRTC_TOKENS_PER_BLOCK,
+            output_dtype=DataType.FP8,
+        )
diff --git a/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py b/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py
new file mode 100644
index 000000000000..00a766aac4ee
--- /dev/null
+++ b/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py
@@ -0,0 +1,251 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""FallbackFmha's refusal of paged prefill at a KV manager's TRTLLM-Gen extra page sizes: kept on every SM but SM107,
+where no TRTLLM-Gen FMHA exists and the thop paged-context kernel serves those pages (its kernel presence is
+test_context_fmha_kernel_presence.py). The dispatch tests run on CPU: the check runs before any kernel. The SM107 tests
+run the admitted path's forward at 128-token pages: fresh, cached-prefix and chunked contexts must match the same
+sequences run unchunked, BF16 and FP8 KV."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import tensorrt_llm
+from tensorrt_llm._torch.attention.backends.fmha import fallback
+from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase
+from tensorrt_llm._torch.attention.backends.interface import (
+    AttentionRuntimeFeatures,
+    PredefinedAttentionMask,
+)
+from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata
+from tensorrt_llm._torch.metadata import KVCacheParams
+from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.bindings import DataType
+from tensorrt_llm.llmapi.llm_args import KvCacheConfig
+from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+
+class _PastRefusal(Exception):
+    """Raised by the stubbed paged-context validation that follows the refusal check."""
+
+
+def _fmha():
+    fmha = object.__new__(fallback.FallbackFmha)
+
+    def _validate(metadata):
+        raise _PastRefusal()
+
+    fmha._validate_paged_context_fmha = _validate
+    return fmha
+
+
+def _context_batch(tokens_per_block: int = 128, paged: bool = True):
+    return SimpleNamespace(
+        num_contexts=1,
+        kv_cache_manager=SimpleNamespace(trtllm_gen_extra_tokens_per_block=frozenset({128})),
+        use_paged_context_fmha=paged,
+        tokens_per_block=tokens_per_block,
+    )
+
+
+@pytest.mark.parametrize("sm", [90, 100, 103, 120])
+def test_extra_page_paged_prefill_is_refused_off_sm107(monkeypatch, sm) -> None:
+    monkeypatch.setattr(fallback, "get_sm_version", lambda: sm)
+    with pytest.raises(RuntimeError, match="cannot serve paged prefill with 128-token KV pages"):
+        _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT)
+
+
+def test_extra_page_paged_prefill_is_served_on_sm107(monkeypatch) -> None:
+    monkeypatch.setattr(fallback, "get_sm_version", lambda: 107)
+    with pytest.raises(_PastRefusal):
+        _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT)
+
+
+@pytest.mark.parametrize("sm", [100, 107])
+def test_other_page_sizes_and_unpaged_prefill_pass_the_check(monkeypatch, sm) -> None:
+    monkeypatch.setattr(fallback, "get_sm_version", lambda: sm)
+    for batch in (_context_batch(tokens_per_block=32), _context_batch(paged=False)):
+        with pytest.raises(_PastRefusal):
+            _fmha()._is_supported(None, None, None, batch, None, phase=FmhaPhase.CONTEXT)
+
+
+_ON_SM107 = torch.cuda.is_available() and get_sm_version() == 107
+_HEADS, _KV_HEADS, _HEAD_DIM, _PAGE = 8, 2, 128, 128
+_REF_RID = 100  # the unchunked run of request r is request _REF_RID + r
+
+# (case id, steps): each step is a batch of (request id, cached tokens, new tokens) context rows, run in order.
+_CACHED_CONTEXT_CASES = [
+    ("fresh_512", [[(0, 0, 512)]]),
+    ("prefix_1000_unaligned", [[(0, 0, 1000)], [(0, 1000, 700)]]),
+    ("prefix_1024_aligned", [[(0, 0, 1024)], [(0, 1024, 300)]]),
+    ("chunked_6096_at_2048", [[(0, 0, 2048)], [(0, 2048, 2048)], [(0, 4096, 2000)]]),
+    ("prefix_beside_fresh", [[(0, 0, 384)], [(0, 384, 257), (1, 0, 129)]]),
+]
+# Worst row's max |out - ref| / max |ref| per KV dtype: chunked vs unchunked (same cached K / V), and unchunked vs a
+# float reference (the FP8 KV kernel also quantizes Q and the softmax to FP8).
+_TOLERANCE = {
+    torch.bfloat16: {"unchunked": 1e-2, "float": 2e-2},
+    torch.float8_e4m3fn: {"unchunked": 3e-2, "float": 1.5e-1},
+}
+
+
+def _kv_manager(kv_dtype: torch.dtype, num_tokens: int) -> KVCacheManager:
+    manager = KVCacheManager(
+        KvCacheConfig(max_tokens=num_tokens, enable_block_reuse=False),
+        tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
+        num_layers=1,
+        num_kv_heads=_KV_HEADS,
+        head_dim=_HEAD_DIM,
+        tokens_per_block=_PAGE,
+        max_seq_len=8192,
+        max_batch_size=8,
+        mapping=Mapping(world_size=1, tp_size=1, rank=0),
+        dtype=DataType.FP8 if kv_dtype == torch.float8_e4m3fn else DataType.BF16,
+    )
+    # As MiniMax-M3's KV managers list it: this is what the admission in FallbackFmha keys on.
+    manager.trtllm_gen_extra_tokens_per_block = frozenset({_PAGE})
+    # Zero the pool: a context step reads a request's whole last page, and stale bytes in its unwritten tail that
+    # decode as FP8 NaN turn the masked P = 0 into NaN outputs.
+    manager.get_buffers(0, kv_layout="HND").view(torch.uint8).zero_()
+    return manager
+
+
+def _context_step(attn, manager, rows, qkv: torch.Tensor) -> torch.Tensor:
+    """One context-only step as the serving runner prepares it: a request's prompt length is its current chunk, with
+    the prefix in ``num_cached_tokens_per_seq``."""
+    metadata = TrtllmAttentionMetadata(
+        max_num_requests=8,
+        max_num_tokens=8192,
+        kv_cache_manager=manager,
+        runtime_features=AttentionRuntimeFeatures(chunked_prefill=True),
+        num_contexts=len(rows),
+        seq_lens=torch.tensor([n for _, _, n in rows], dtype=torch.int),
+        request_ids=[rid for rid, _, _ in rows],
+        prompt_lens=[n for _, _, n in rows],
+        kv_cache_params=KVCacheParams(
+            use_cache=True, num_cached_tokens_per_seq=[c for _, c, _ in rows]
+        ),
+    )
+    assert metadata.use_paged_context_fmha and metadata.tokens_per_block == _PAGE
+    metadata.prepare()
+    return attn.forward(qkv, None, None, metadata, attention_mask=PredefinedAttentionMask.CAUSAL)
+
+
+def _cached_kv(manager, rid: int, length: int) -> torch.Tensor:
+    pool = manager.get_buffers(0, kv_layout="HND")  # [pages, K/V, kv heads, page tokens, head dim]
+    pages = manager.get_batch_cache_indices([rid], 0)[0][: -(-length // _PAGE)]
+    kv = pool[torch.tensor(pages, device=pool.device)].permute(1, 2, 0, 3, 4)
+    return kv.reshape(2, _KV_HEADS, -1, _HEAD_DIM)[:, :, :length].view(torch.uint8)
+
+
+def _float_reference(
+    qkv: torch.Tensor, start: int, n: int, kv_from: int, fp8: bool
+) -> torch.Tensor:
+    """Causal attention of rows start..start+n over keys kv_from..start+n, in float32."""
+    q_size, kv_size = _HEADS * _HEAD_DIM, _KV_HEADS * _HEAD_DIM
+    rows = qkv[kv_from : start + n]
+    if fp8:
+        rows = rows.to(torch.float8_e4m3fn)
+    rows = rows.float()
+    q = rows[start - kv_from :, :q_size].view(n, _HEADS, _HEAD_DIM).transpose(0, 1)
+    k = rows[:, q_size : q_size + kv_size].view(-1, _KV_HEADS, _HEAD_DIM).transpose(0, 1)
+    v = rows[:, q_size + kv_size :].view(-1, _KV_HEADS, _HEAD_DIM).transpose(0, 1)
+    k = k.repeat_interleave(_HEADS // _KV_HEADS, dim=0)
+    v = v.repeat_interleave(_HEADS // _KV_HEADS, dim=0)
+    scores = q @ k.transpose(1, 2) / _HEAD_DIM**0.5
+    q_pos = torch.arange(start, start + n, device=qkv.device)[:, None]
+    k_pos = torch.arange(kv_from, start + n, device=qkv.device)[None, :]
+    scores = scores.masked_fill(k_pos > q_pos, float("-inf"))
+    return (scores.softmax(-1) @ v).transpose(0, 1).reshape(n, q_size)
+
+
+def _relative_error(out: torch.Tensor, ref: torch.Tensor) -> float:
+    out, ref = out.float(), ref.float()
+    scale = ref.abs().amax(dim=1).clamp_min(1e-6)
+    return ((out - ref).abs().amax(dim=1) / scale).max().item()
+
+
+@pytest.mark.skipif(not _ON_SM107, reason="the SM107 thop paged-context path at 128-token pages")
+@pytest.mark.parametrize(
+    "kv_dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16_kv", "fp8_kv"]
+)
+@pytest.mark.parametrize("case", _CACHED_CONTEXT_CASES, ids=[c[0] for c in _CACHED_CONTEXT_CASES])
+def test_sm107_cached_context_matches_unchunked(monkeypatch, case, kv_dtype) -> None:
+    """The admitted path attends to the cached prefix: each chunk's output matches the same rows of the request run
+    in one unchunked step, and the chunks leave the same K / V in the cache."""
+    name, steps = case
+    fp8 = kv_dtype == torch.float8_e4m3fn
+    tolerance = _TOLERANCE[kv_dtype]
+    lengths = {}
+    for step in steps:
+        for rid, c, n in step:
+            assert c == lengths.get(rid, 0), "each chunk starts at its request's cached length"
+            lengths[rid] = c + n
+    served = []
+    forward = fallback.FallbackFmha.forward
+    monkeypatch.setattr(
+        fallback.FallbackFmha,
+        "forward",
+        lambda self, q, k, v, metadata, args: served.append(metadata.num_contexts)
+        or forward(self, q, k, v, metadata, args),
+    )
+
+    gen = torch.Generator(device="cuda").manual_seed(len(name))
+    width = (_HEADS + 2 * _KV_HEADS) * _HEAD_DIM
+    qkv = {
+        rid: torch.randn(n, width, generator=gen, device="cuda").to(torch.bfloat16)
+        for rid, n in lengths.items()
+    }
+    num_tokens = 2 * sum(-(-n // _PAGE) * _PAGE for n in lengths.values())
+    manager = _kv_manager(kv_dtype, num_tokens)
+    try:
+        attn = TrtllmAttention(
+            layer_idx=0,
+            num_heads=_HEADS,
+            head_dim=_HEAD_DIM,
+            num_kv_heads=_KV_HEADS,
+            quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) if fp8 else None,
+        )
+        rids = sorted(lengths)
+        manager.add_dummy_requests(
+            rids + [_REF_RID + r for r in rids], [lengths[r] for r in rids] * 2
+        )
+        # The unchunked run: every request as one fresh context.
+        rows = [(_REF_RID + r, 0, lengths[r]) for r in rids]
+        out = _context_step(attn, manager, rows, torch.cat([qkv[r] for r in rids]))
+        unchunked, first = {}, 0
+        for r in rids:
+            unchunked[r] = out[first : first + lengths[r]]
+            first += lengths[r]
+            error = _relative_error(unchunked[r], _float_reference(qkv[r], 0, lengths[r], 0, fp8))
+            print(f"\n[{name}] request {r} unchunked vs float reference: {error:.4f}")
+            assert error <= tolerance["float"]
+
+        for step in steps:
+            out = _context_step(
+                attn, manager, step, torch.cat([qkv[r][c : c + n] for r, c, n in step])
+            )
+            first = 0
+            for r, c, n in step:
+                got = out[first : first + n]
+                first += n
+                assert torch.isfinite(got.float()).all()
+                error = _relative_error(got, unchunked[r][c : c + n])
+                print(f"[{name}] request {r} rows {c}..{c + n} vs unchunked: {error:.4f}")
+                assert error <= tolerance["unchunked"]
+                if c:
+                    # A kernel that dropped the cached prefix would match attention over the chunk alone.
+                    dropped = _float_reference(qkv[r], c, n, c, fp8)
+                    assert _relative_error(got, dropped) > 4 * tolerance["float"]
+        assert served and all(served), "the context steps ran through FallbackFmha"
+        torch.cuda.synchronize()
+        for r in rids:
+            assert torch.equal(
+                _cached_kv(manager, r, lengths[r]), _cached_kv(manager, _REF_RID + r, lengths[r])
+            ), f"request {r}: the chunks' K / V differ from the unchunked run's"
+    finally:
+        manager.shutdown()
diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py
index 0ed31615b378..98692ef294a9 100644
--- a/tests/unittest/_torch/test_model_config.py
+++ b/tests/unittest/_torch/test_model_config.py
@@ -15,6 +15,7 @@
     validate_and_set_kv_cache_quant,
     validate_encoder_decoder_kv_cache_config,
 )
+from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
 
@@ -130,7 +131,11 @@ def test_validate_and_set_kv_cache_quant_auto_uses_checkpoint():
     assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8
 
 
-def test_validate_and_set_kv_cache_quant_explicit_dtype_overrides():
+def test_validate_and_set_kv_cache_quant_explicit_dtype_overrides(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # Pin a non-SM107 device: on SM107 an explicit NVFP4 request falls back to FP8.
+    monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.model_loader.get_sm_version", lambda: 100)
     model_config = _make_model_config_with_kv_quant(QuantAlgo.FP8)
     validate_and_set_kv_cache_quant(model_config, "nvfp4")
     assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4
@@ -189,18 +194,32 @@ def test_validate_and_set_kv_cache_quant_auto_keeps_quant_config_dict():
         assert layer_quant_config.kv_cache_quant_algo is None
 
 
+@pytest.mark.parametrize("implementation", [None, "triton", "msa"])
+@pytest.mark.parametrize("requested_dtype", ["nvfp4", "auto"])
 def test_validate_and_set_kv_cache_quant_sm107_downgrade_propagates_to_quant_config_dict(
     monkeypatch: pytest.MonkeyPatch,
+    implementation: str | None,
+    requested_dtype: str,
 ) -> None:
-    """The SM107 NVFP4->FP8 downgrade must update per-layer QuantConfigs too,
-    otherwise a mixed-precision checkpoint keeps NVFP4 in quant_config_dict
-    while the KV pool (sized from the global config) uses FP8."""
+    """The SM107 NVFP4->FP8 decision must reach per-layer QuantConfigs too,
+    otherwise a mixed-precision checkpoint and the KV pool (sized from the
+    global config) disagree. MiniMax-M3 MSA keeps NVFP4, which its hybrid
+    cache serves on SM107; every other configuration falls back to FP8."""
     _mock_sm107(monkeypatch)
     model_config = _make_mixed_precision_model_config()
-    validate_and_set_kv_cache_quant(model_config, "nvfp4")
-    assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8
+    if implementation is not None:
+        model_config.sparse_attention_config = MiniMaxM3SparseAttentionConfig(
+            implementation=implementation
+        )
+    if requested_dtype == "auto":
+        model_config.quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4
+        for layer_quant_config in model_config.quant_config_dict.values():
+            layer_quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4
+    validate_and_set_kv_cache_quant(model_config, requested_dtype)
+    expected = QuantAlgo.NVFP4 if implementation == "msa" else QuantAlgo.FP8
+    assert model_config.quant_config.kv_cache_quant_algo == expected
     for layer_quant_config in model_config.quant_config_dict.values():
-        assert layer_quant_config.kv_cache_quant_algo == QuantAlgo.FP8
+        assert layer_quant_config.kv_cache_quant_algo == expected
 
 
 def _write_safetensors_header(checkpoint_dir, tensor_dtype, tensor_shape):

From e94d5a5e9ebf4e2db6534c0866d846b724f7880e Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Tue, 20 Oct 2026 09:35:09 +0000
Subject: [PATCH 2/9] [None][fix] Use FlashInfer TRTLLM-Gen FMHA on SM107 for
 extra KV page sizes (cherry-pick of #20676)

Cherry-pick of NVIDIA/TensorRT-LLM#20676 by @yuxianq (725b63af03; on main as
2dc625cf8e, https://nvbugs/6946881). The changed lines are identical to the
main commit, so a later merge of main resolves without conflicts.

Adaptations to this base:
- flashinfer_trtllm_gen.py: the SM gate lives in is_environment_available(),
  not _is_available(), so the SM107 admission lands there.
- test_small_context_fallback_preserves_mixed_batch_generation: this base
  carries @pytest.mark.cpu_only and the page_size/paged_context
  parametrization, so the monkeypatch parameter and the get_sm_version pin
  are inserted by hand.

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
---
 .../backends/fmha/flashinfer_trtllm_gen.py    |  16 ++-
 .../test_flashinfer_trtllm_gen_fmha.py        | 106 ++++++++++++++++++
 2 files changed, 119 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
index 6d53abd9e065..3f11871eb74b 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
@@ -407,9 +407,10 @@ 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
 
@@ -790,12 +791,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 get_sm_version() == 107 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, ""
 
diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py
index 2165c36576a5..8fb01d9f3d2f 100644
--- a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py
+++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py
@@ -32,6 +32,8 @@
 from tensorrt_llm.bindings import DataType
 from tensorrt_llm.quantization.mode import QuantMode
 
+_FLASHINFER_TRTLLM_GEN_MODULE = "tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen"
+
 
 def test_flashinfer_fp8_mode_remains_implementation_local() -> None:
     attn = FakeAttention()
@@ -62,6 +64,7 @@ def test_flashinfer_fp8_mode_remains_implementation_local() -> None:
 @pytest.mark.cpu_only
 @pytest.mark.parametrize("page_size,paged_context", [(32, False), (128, True)])
 def test_small_context_fallback_preserves_mixed_batch_generation(
+    monkeypatch: pytest.MonkeyPatch,
     num_contexts: int,
     num_generations: int,
     dtype: torch.dtype,
@@ -69,6 +72,7 @@ def test_small_context_fallback_preserves_mixed_batch_generation(
     page_size: int,
     paged_context: bool,
 ) -> None:
+    monkeypatch.setattr(f"{_FLASHINFER_TRTLLM_GEN_MODULE}.get_sm_version", lambda: 100)
     attn = FakeAttention()
     attn.head_dim = 256
     attn.sparse_params = None
@@ -191,3 +195,105 @@ def test_flashinfer_quantized_kv_context_avoids_fp16_bf16_fallback(
 
         assert supported, reason
         assert reason == ""
+
+
+@pytest.mark.parametrize(
+    "sm_version,expected",
+    [(90, False), (100, True), (103, True), (107, True), (120, False)],
+)
+def test_flashinfer_availability_by_sm(
+    monkeypatch: pytest.MonkeyPatch,
+    sm_version: int,
+    expected: bool,
+) -> None:
+    monkeypatch.setattr(f"{_FLASHINFER_TRTLLM_GEN_MODULE}.IS_FLASHINFER_AVAILABLE", True)
+    monkeypatch.setattr(
+        FlashInferTrtllmGenFmha, "_missing_fused_nanobind_ops", staticmethod(lambda: [])
+    )
+    monkeypatch.setattr(f"{_FLASHINFER_TRTLLM_GEN_MODULE}.get_sm_version", lambda: sm_version)
+    attn = FakeAttention()
+    attn.sparse_params = None
+
+    assert FlashInferTrtllmGenFmha._is_available(attn) is expected
+
+
+@pytest.mark.parametrize(
+    "input_type",
+    [AttentionInputType.context_only, AttentionInputType.generation_only],
+    ids=["context", "generation"],
+)
+@pytest.mark.parametrize(
+    "sm_version,tokens_per_block,extra_tokens_per_block,expected",
+    [
+        (107, 128, frozenset({128}), True),
+        (107, 32, frozenset({128}), False),
+        (107, 32, frozenset(), False),
+        (103, 32, frozenset(), True),
+        (103, 128, frozenset({128}), True),
+    ],
+    ids=[
+        "sm107_extra_page_size",
+        "sm107_default_page_size_with_extra",
+        "sm107_default_page_size",
+        "sm103_default_page_size",
+        "sm103_extra_page_size",
+    ],
+)
+def test_sm107_serves_only_extra_page_sizes(
+    monkeypatch: pytest.MonkeyPatch,
+    input_type: AttentionInputType,
+    sm_version: int,
+    tokens_per_block: int,
+    extra_tokens_per_block: frozenset,
+    expected: bool,
+) -> None:
+    """FP8-KV GQA 16:1 head-128 paged attention, the MiniMax-M3 Eagle3 draft layer shape."""
+    monkeypatch.setattr(f"{_FLASHINFER_TRTLLM_GEN_MODULE}.get_sm_version", lambda: sm_version)
+    fmha = object.__new__(FlashInferTrtllmGenFmha)
+    fmha.kv_factor = 2
+    monkeypatch.setattr(fmha, "_get_total_num_blocks", lambda _: 0)
+    attn = SimpleNamespace(
+        is_mla_enable=False,
+        sparse_params=None,
+        position_embedding_type=0,
+        head_dim=128,
+        num_heads=16,
+        num_kv_heads=1,
+    )
+    is_context = input_type == AttentionInputType.context_only
+    num_tokens = 4
+    q_hidden_size = attn.num_heads * attn.head_dim
+    qkv_hidden_size = q_hidden_size + 2 * attn.num_kv_heads * attn.head_dim
+    q = torch.empty((num_tokens, qkv_hidden_size), dtype=torch.bfloat16)
+    metadata = SimpleNamespace(
+        num_contexts=1 if is_context else 0,
+        num_generations=0 if is_context else 1,
+        helix_position_offsets=None,
+        num_sparse_topk=0,
+        use_spec_decoding=False,
+        is_spec_dec_tree=False,
+        kv_cache_block_offsets=object(),
+        kv_cache_manager=SimpleNamespace(
+            dtype=DataType.FP8,
+            trtllm_gen_extra_tokens_per_block=extra_tokens_per_block,
+        ),
+        is_cross=False,
+        is_spec_decoding_enabled=False,
+        tokens_per_block=tokens_per_block,
+        beam_width=1,
+    )
+    forward_args = AttentionForwardArgs(
+        output=torch.empty((num_tokens, q_hidden_size), dtype=torch.bfloat16),
+        attention_input_type=input_type,
+        is_fused_qkv=True,
+        update_kv_cache=True,
+    )
+
+    phase = FmhaPhase.CONTEXT if is_context else FmhaPhase.GENERATION
+    for requested_phase in (None, phase):
+        supported, reason = fmha._is_supported_with_reason(
+            q, None, None, attn, metadata, forward_args, phase=requested_phase
+        )
+        assert supported is expected, reason
+        if not expected:
+            assert "SM107" in reason

From e78055740a1056f3fff732741e1144d3837df350 Mon Sep 17 00:00:00 2001
From: peihengh <259410613+peihu-nv@users.noreply.github.com>
Date: Thu, 22 Oct 2026 00:03:33 -0700
Subject: [PATCH 3/9] rework #20167: route SM107 P128 FMHA to FlashInfer, drop
 the NVRTC path

- Remove the SM107 P128 NVRTC rule in fmhaKernels.h, the FallbackFmha SM107
  extra-page-size admission and the NVRTC kernel-presence test.
- MiniMaxM3KVCacheManagerV2: a separate (is_draft) manager at 128-token pages
  now lists 128 as a trtllm-gen extra page size, so every M3 manager that runs
  P128 native layers keeps the FlashInfer route.
- FlashInferTrtllmGenFmha.serves_only_extra_page_sizes(): single source for the
  SM107 extras-only rule; MiniMax-M3 dense layers keep SDPA on SM107.
- Re-point the SM107 P128 oracle at FlashInfer
  (test_flashinfer_trtllm_gen_sm107_p128.py): cached prefix / chunked context,
  generation q_len 1, verify 1+3, mixed, BF16 and FP8 KV, NaN tails, 16:1 hd128.
- M3 deployment guide: SM107 needs FlashInfer; pre-populate FLASHINFER_CUBIN_DIR.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
---
 .../trtllmGenKernels/fmha/fmhaKernels.h       |  17 +-
 ...ployment-guide-for-minimax-m3-on-trtllm.md |   1 +
 .../attention/backends/fmha/fallback.py       |  15 -
 .../backends/fmha/flashinfer_trtllm_gen.py    |  11 +-
 .../sparse/minimax_m3/cache_manager.py        |  13 +-
 .../_torch/models/modeling_minimaxm3.py       |   5 +
 .../test_minimax_m3_shared_draft_layers.py    |  39 ++
 .../test_context_fmha_kernel_presence.py      |  31 --
 .../test_fallback_fmha_extra_page_sizes.py    | 251 ------------
 .../test_flashinfer_trtllm_gen_sm107_p128.py  | 366 ++++++++++++++++++
 .../unittest/_torch/models/test_minimax_m3.py |  23 +-
 11 files changed, 450 insertions(+), 322 deletions(-)
 delete mode 100644 tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py
 create mode 100644 tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py

diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h
index a0c192d12a1c..79cc1eb5dd7e 100644
--- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h
+++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h
@@ -254,7 +254,7 @@ class TllmGenFmhaKernel
         }
     }
 
-    bool shouldUseNvrtc(FmhaOptions const& options) const
+    static bool shouldUseNvrtc(FmhaOptions const& options)
     {
 #ifdef TLLM_RUBIN_FEATURES
         // Spcompress kernels must use precompiled cubins.
@@ -263,20 +263,7 @@ class TllmGenFmhaKernel
             return false;
         }
 #endif // TLLM_RUBIN_FEATURES
-       // The cubin table has no dense paged-KV kernels for 128-token pages, and the page size is
-       // part of the kernel hash. On SM107, generate those BF16/E4M3 context and KeepsMmaAb
-       // generation kernels for the actual cache layout instead.
-        constexpr int kPageSize128 = 128;
-        if (mSM == kSM_107
-            && (isContextKernel(options.mFmhaKernelType) || isKeepsMmaAbForGenerationKernel(options.mFmhaKernelType))
-            && !options.mIsMlaGen && !isTokenSparse(options.mSparseType) && isPagedKv(options.mQkvLayout)
-            && options.mNumTokensPerPage == kPageSize128 && options.mDtypeQ == options.mDtypeKv
-            && (options.mDtypeQ == tg::Dtype::Bfloat16 || options.mDtypeQ == tg::Dtype::E4m3)
-            && options.mDtypeOut == tg::Dtype::Bfloat16)
-        {
-            return true;
-        }
-        // Sparse MQA/GQA uses NVRTC path for now because no model really uses it.
+       // Sparse MQA/GQA uses NVRTC path for now because no model really uses it.
         if (isStaticTokenSparse(options.mSparseType) && !options.mIsMlaGen)
         {
             return true;
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..f3413a6f8fae 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,6 +65,7 @@ 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, populate a shared `FLASHINFER_CUBIN_DIR` before the multi-rank launch (for example with a single-rank warm-up run that serves one Eagle3 request) and set `FLASHINFER_NO_DOWNLOAD=1`; see [FlashInfer JIT workspaces for MPI workers](../llm-api/index.md#flashinfer-jit-workspaces-for-mpi-workers).
 * **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.
diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py
index be974499b494..448d64df9925 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py
@@ -53,20 +53,6 @@
     }
 )
 
-# SM107 has no TRTLLM-Gen FMHA, so the thop paged-context kernel serves every
-# context phase there, at the 128-token pages MiniMax-M3's KV managers opt in
-# (as before those managers listed 128 as a TRTLLM-Gen extra page size). The
-# SM107 trtllm-gen runner generates those kernels with NVRTC (fmhaKernels.h);
-# test_context_fmha_kernel_presence.py checks that the build reports them
-# present. Every other SM keeps the refusal.
-_THOP_PAGED_CONTEXT_EXTRA_PAGES_SMS: frozenset = frozenset({107})
-
-
-def _thop_serves_extra_page_sizes() -> bool:
-    """Whether the thop paged-context kernel serves a manager's TRTLLM-Gen extra page sizes on this device."""
-    return get_sm_version() in _THOP_PAGED_CONTEXT_EXTRA_PAGES_SMS
-
-
 # ``thop.attention`` kwargs hard-wired to a literal at the call site (no
 # rich object owns them). Sync test enforces both the kwarg name and the
 # literal value.
@@ -269,7 +255,6 @@ def _is_supported(
                 and phase != FmhaPhase.GENERATION
                 and metadata.use_paged_context_fmha
                 and metadata.tokens_per_block in extra_page_sizes
-                and not _thop_serves_extra_page_sizes()
             ):
                 raise RuntimeError(
                     f"FallbackFmha cannot serve paged prefill with "
diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
index 3f11871eb74b..47652e5b9efa 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
+++ b/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 = {
@@ -416,6 +420,11 @@ def is_environment_available(cls) -> bool:
 
         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():
@@ -801,7 +810,7 @@ def _is_supported_with_reason(
         # 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 get_sm_version() == 107 and tokens_per_block not in extra_tokens_per_block:
+        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."
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 ebd8432a5a78..6dd95c6c1cf1 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
@@ -559,7 +559,15 @@ 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:
+                self.trtllm_gen_extra_tokens_per_block = frozenset({128})
+            return
+        if not native_layers:
             return
         if self.enable_swa_scratch_reuse:
             raise NotImplementedError(
@@ -571,7 +579,8 @@ 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:
             self.trtllm_gen_extra_tokens_per_block = frozenset({128})
         if self._use_per_layer_page_tables:
diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
index 9a878e561ebc..952647505aca 100644
--- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py
+++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
@@ -1233,6 +1233,11 @@ def _init_native_dense_attention(self, model_config: ModelConfig) -> None:
             return
         if not FlashInferTrtllmGenFmha.is_environment_available():
             return
+        if FlashInferTrtllmGenFmha.serves_only_extra_page_sizes():
+            # 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,
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..227b6c24d459 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
@@ -301,6 +301,45 @@ def fake_base_prepare(self, index_mapper_capacity):
     assert manager.trtllm_gen_extra_tokens_per_block == frozenset(expected_extra_pages)
 
 
+@pytest.mark.parametrize(
+    "dtype,tokens_per_block,expected",
+    [
+        (DataType.FP8, 128, {128}),
+        (DataType.BF16, 128, {128}),
+        (DataType.FP8, 32, set()),
+        (DataType.NVFP4, 128, set()),
+    ],
+)
+def test_separate_draft_manager_opts_into_p128(monkeypatch, dtype, tokens_per_block, expected):
+    """A separate (is_draft) manager holds only native draft layers: at 128-token pages
+    it lists 128 as a trtllm-gen extra page size (the only P128 route on SM107) and
+    builds no virtual pools."""
+    pointers = torch.tensor([[0x6000_0000, 0]])
+    mapping = torch.tensor([[0, 0]], dtype=torch.int32)
+
+    def fake_base_prepare(self, index_mapper_capacity):
+        self._use_per_layer_page_tables = False
+        self.kv_cache_pool_pointers = pointers.clone()
+        self.kv_cache_pool_mapping = mapping.clone()
+        self.num_attention_op_pools = 1
+
+    monkeypatch.setattr(KVCacheManagerV2, "_prepare_page_table_tensor", fake_base_prepare)
+    manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2)
+    manager.dtype = dtype
+    manager._shared_draft_layer_ids = []
+    manager.layer_offsets = {0: 0}
+    manager.is_draft = True
+    manager.enable_swa_scratch_reuse = False
+    manager.tokens_per_block = tokens_per_block
+
+    manager._prepare_page_table_tensor(8)
+
+    assert manager.trtllm_gen_extra_tokens_per_block == frozenset(expected)
+    assert manager._draft_op_pools == ()
+    assert torch.equal(manager.kv_cache_pool_pointers, pointers)
+    assert manager.num_attention_op_pools == (0 if dtype == DataType.NVFP4 else 1)
+
+
 def test_update_resources_refuses_tree_relocation(monkeypatch):
     """Linear acceptance rewinds through the base; tree acceptance is refused."""
     calls = []
diff --git a/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py b/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py
index 83ab912fdad6..a9e1c62b1e0e 100644
--- a/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py
+++ b/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py
@@ -35,10 +35,6 @@
 # A head width the SM100-family kernel table does carry for 16-bit paged-context
 # attention.
 _SUPPORTED_HEAD_SIZE = 64
-# The page size and head width of MiniMax-M3's dense and draft attention. The
-# kernel table carries no 128-token-page kernels; SM107 generates them at runtime.
-_NVRTC_TOKENS_PER_BLOCK = 128
-_NVRTC_HEAD_SIZE = 128
 
 _NO_GPU = not torch.cuda.is_available()
 _NO_GPU_REASON = "needs a GPU; missing bindings on a GPU are a test failure"
@@ -192,30 +188,3 @@ def test_lookup_is_page_size_sensitive():
     )
     assert present
     assert not odd_page
-
-
-@pytest.mark.parametrize("kv_cache_dtype", [DataType.BF16, DataType.FP8], ids=["bf16", "fp8"])
-def test_lookup_reports_present_for_sm107_128_token_pages(kv_cache_dtype):
-    """SM107 serves dense paged context over 128-token pages with BF16 output.
-
-    No precompiled kernel exists for that page size, so matched BF16/E4M3 Q
-    and KV with BF16 output must be generated at runtime rather than reported
-    absent, which would refuse paged-context attention for such models. The
-    runtime path is limited to that output precision: FP8 output at the same
-    page size has no kernel either way and must still report absent.
-    """
-    if get_sm_version() != 107:
-        pytest.skip("runtime-generated 128-token-page context kernels are specific to SM107")
-    assert thop.fused_context_fmha_kernel_exists(
-        head_size=_NVRTC_HEAD_SIZE,
-        kv_cache_dtype=kv_cache_dtype,
-        tokens_per_block=_NVRTC_TOKENS_PER_BLOCK,
-        output_dtype=DataType.BF16,
-    )
-    if kv_cache_dtype == DataType.FP8:
-        assert not thop.fused_context_fmha_kernel_exists(
-            head_size=_NVRTC_HEAD_SIZE,
-            kv_cache_dtype=kv_cache_dtype,
-            tokens_per_block=_NVRTC_TOKENS_PER_BLOCK,
-            output_dtype=DataType.FP8,
-        )
diff --git a/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py b/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py
deleted file mode 100644
index 00a766aac4ee..000000000000
--- a/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py
+++ /dev/null
@@ -1,251 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""FallbackFmha's refusal of paged prefill at a KV manager's TRTLLM-Gen extra page sizes: kept on every SM but SM107,
-where no TRTLLM-Gen FMHA exists and the thop paged-context kernel serves those pages (its kernel presence is
-test_context_fmha_kernel_presence.py). The dispatch tests run on CPU: the check runs before any kernel. The SM107 tests
-run the admitted path's forward at 128-token pages: fresh, cached-prefix and chunked contexts must match the same
-sequences run unchunked, BF16 and FP8 KV."""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-import tensorrt_llm
-from tensorrt_llm._torch.attention.backends.fmha import fallback
-from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase
-from tensorrt_llm._torch.attention.backends.interface import (
-    AttentionRuntimeFeatures,
-    PredefinedAttentionMask,
-)
-from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata
-from tensorrt_llm._torch.metadata import KVCacheParams
-from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
-from tensorrt_llm._utils import get_sm_version
-from tensorrt_llm.bindings import DataType
-from tensorrt_llm.llmapi.llm_args import KvCacheConfig
-from tensorrt_llm.mapping import Mapping
-from tensorrt_llm.models.modeling_utils import QuantConfig
-from tensorrt_llm.quantization.mode import QuantAlgo
-
-
-class _PastRefusal(Exception):
-    """Raised by the stubbed paged-context validation that follows the refusal check."""
-
-
-def _fmha():
-    fmha = object.__new__(fallback.FallbackFmha)
-
-    def _validate(metadata):
-        raise _PastRefusal()
-
-    fmha._validate_paged_context_fmha = _validate
-    return fmha
-
-
-def _context_batch(tokens_per_block: int = 128, paged: bool = True):
-    return SimpleNamespace(
-        num_contexts=1,
-        kv_cache_manager=SimpleNamespace(trtllm_gen_extra_tokens_per_block=frozenset({128})),
-        use_paged_context_fmha=paged,
-        tokens_per_block=tokens_per_block,
-    )
-
-
-@pytest.mark.parametrize("sm", [90, 100, 103, 120])
-def test_extra_page_paged_prefill_is_refused_off_sm107(monkeypatch, sm) -> None:
-    monkeypatch.setattr(fallback, "get_sm_version", lambda: sm)
-    with pytest.raises(RuntimeError, match="cannot serve paged prefill with 128-token KV pages"):
-        _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT)
-
-
-def test_extra_page_paged_prefill_is_served_on_sm107(monkeypatch) -> None:
-    monkeypatch.setattr(fallback, "get_sm_version", lambda: 107)
-    with pytest.raises(_PastRefusal):
-        _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT)
-
-
-@pytest.mark.parametrize("sm", [100, 107])
-def test_other_page_sizes_and_unpaged_prefill_pass_the_check(monkeypatch, sm) -> None:
-    monkeypatch.setattr(fallback, "get_sm_version", lambda: sm)
-    for batch in (_context_batch(tokens_per_block=32), _context_batch(paged=False)):
-        with pytest.raises(_PastRefusal):
-            _fmha()._is_supported(None, None, None, batch, None, phase=FmhaPhase.CONTEXT)
-
-
-_ON_SM107 = torch.cuda.is_available() and get_sm_version() == 107
-_HEADS, _KV_HEADS, _HEAD_DIM, _PAGE = 8, 2, 128, 128
-_REF_RID = 100  # the unchunked run of request r is request _REF_RID + r
-
-# (case id, steps): each step is a batch of (request id, cached tokens, new tokens) context rows, run in order.
-_CACHED_CONTEXT_CASES = [
-    ("fresh_512", [[(0, 0, 512)]]),
-    ("prefix_1000_unaligned", [[(0, 0, 1000)], [(0, 1000, 700)]]),
-    ("prefix_1024_aligned", [[(0, 0, 1024)], [(0, 1024, 300)]]),
-    ("chunked_6096_at_2048", [[(0, 0, 2048)], [(0, 2048, 2048)], [(0, 4096, 2000)]]),
-    ("prefix_beside_fresh", [[(0, 0, 384)], [(0, 384, 257), (1, 0, 129)]]),
-]
-# Worst row's max |out - ref| / max |ref| per KV dtype: chunked vs unchunked (same cached K / V), and unchunked vs a
-# float reference (the FP8 KV kernel also quantizes Q and the softmax to FP8).
-_TOLERANCE = {
-    torch.bfloat16: {"unchunked": 1e-2, "float": 2e-2},
-    torch.float8_e4m3fn: {"unchunked": 3e-2, "float": 1.5e-1},
-}
-
-
-def _kv_manager(kv_dtype: torch.dtype, num_tokens: int) -> KVCacheManager:
-    manager = KVCacheManager(
-        KvCacheConfig(max_tokens=num_tokens, enable_block_reuse=False),
-        tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
-        num_layers=1,
-        num_kv_heads=_KV_HEADS,
-        head_dim=_HEAD_DIM,
-        tokens_per_block=_PAGE,
-        max_seq_len=8192,
-        max_batch_size=8,
-        mapping=Mapping(world_size=1, tp_size=1, rank=0),
-        dtype=DataType.FP8 if kv_dtype == torch.float8_e4m3fn else DataType.BF16,
-    )
-    # As MiniMax-M3's KV managers list it: this is what the admission in FallbackFmha keys on.
-    manager.trtllm_gen_extra_tokens_per_block = frozenset({_PAGE})
-    # Zero the pool: a context step reads a request's whole last page, and stale bytes in its unwritten tail that
-    # decode as FP8 NaN turn the masked P = 0 into NaN outputs.
-    manager.get_buffers(0, kv_layout="HND").view(torch.uint8).zero_()
-    return manager
-
-
-def _context_step(attn, manager, rows, qkv: torch.Tensor) -> torch.Tensor:
-    """One context-only step as the serving runner prepares it: a request's prompt length is its current chunk, with
-    the prefix in ``num_cached_tokens_per_seq``."""
-    metadata = TrtllmAttentionMetadata(
-        max_num_requests=8,
-        max_num_tokens=8192,
-        kv_cache_manager=manager,
-        runtime_features=AttentionRuntimeFeatures(chunked_prefill=True),
-        num_contexts=len(rows),
-        seq_lens=torch.tensor([n for _, _, n in rows], dtype=torch.int),
-        request_ids=[rid for rid, _, _ in rows],
-        prompt_lens=[n for _, _, n in rows],
-        kv_cache_params=KVCacheParams(
-            use_cache=True, num_cached_tokens_per_seq=[c for _, c, _ in rows]
-        ),
-    )
-    assert metadata.use_paged_context_fmha and metadata.tokens_per_block == _PAGE
-    metadata.prepare()
-    return attn.forward(qkv, None, None, metadata, attention_mask=PredefinedAttentionMask.CAUSAL)
-
-
-def _cached_kv(manager, rid: int, length: int) -> torch.Tensor:
-    pool = manager.get_buffers(0, kv_layout="HND")  # [pages, K/V, kv heads, page tokens, head dim]
-    pages = manager.get_batch_cache_indices([rid], 0)[0][: -(-length // _PAGE)]
-    kv = pool[torch.tensor(pages, device=pool.device)].permute(1, 2, 0, 3, 4)
-    return kv.reshape(2, _KV_HEADS, -1, _HEAD_DIM)[:, :, :length].view(torch.uint8)
-
-
-def _float_reference(
-    qkv: torch.Tensor, start: int, n: int, kv_from: int, fp8: bool
-) -> torch.Tensor:
-    """Causal attention of rows start..start+n over keys kv_from..start+n, in float32."""
-    q_size, kv_size = _HEADS * _HEAD_DIM, _KV_HEADS * _HEAD_DIM
-    rows = qkv[kv_from : start + n]
-    if fp8:
-        rows = rows.to(torch.float8_e4m3fn)
-    rows = rows.float()
-    q = rows[start - kv_from :, :q_size].view(n, _HEADS, _HEAD_DIM).transpose(0, 1)
-    k = rows[:, q_size : q_size + kv_size].view(-1, _KV_HEADS, _HEAD_DIM).transpose(0, 1)
-    v = rows[:, q_size + kv_size :].view(-1, _KV_HEADS, _HEAD_DIM).transpose(0, 1)
-    k = k.repeat_interleave(_HEADS // _KV_HEADS, dim=0)
-    v = v.repeat_interleave(_HEADS // _KV_HEADS, dim=0)
-    scores = q @ k.transpose(1, 2) / _HEAD_DIM**0.5
-    q_pos = torch.arange(start, start + n, device=qkv.device)[:, None]
-    k_pos = torch.arange(kv_from, start + n, device=qkv.device)[None, :]
-    scores = scores.masked_fill(k_pos > q_pos, float("-inf"))
-    return (scores.softmax(-1) @ v).transpose(0, 1).reshape(n, q_size)
-
-
-def _relative_error(out: torch.Tensor, ref: torch.Tensor) -> float:
-    out, ref = out.float(), ref.float()
-    scale = ref.abs().amax(dim=1).clamp_min(1e-6)
-    return ((out - ref).abs().amax(dim=1) / scale).max().item()
-
-
-@pytest.mark.skipif(not _ON_SM107, reason="the SM107 thop paged-context path at 128-token pages")
-@pytest.mark.parametrize(
-    "kv_dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16_kv", "fp8_kv"]
-)
-@pytest.mark.parametrize("case", _CACHED_CONTEXT_CASES, ids=[c[0] for c in _CACHED_CONTEXT_CASES])
-def test_sm107_cached_context_matches_unchunked(monkeypatch, case, kv_dtype) -> None:
-    """The admitted path attends to the cached prefix: each chunk's output matches the same rows of the request run
-    in one unchunked step, and the chunks leave the same K / V in the cache."""
-    name, steps = case
-    fp8 = kv_dtype == torch.float8_e4m3fn
-    tolerance = _TOLERANCE[kv_dtype]
-    lengths = {}
-    for step in steps:
-        for rid, c, n in step:
-            assert c == lengths.get(rid, 0), "each chunk starts at its request's cached length"
-            lengths[rid] = c + n
-    served = []
-    forward = fallback.FallbackFmha.forward
-    monkeypatch.setattr(
-        fallback.FallbackFmha,
-        "forward",
-        lambda self, q, k, v, metadata, args: served.append(metadata.num_contexts)
-        or forward(self, q, k, v, metadata, args),
-    )
-
-    gen = torch.Generator(device="cuda").manual_seed(len(name))
-    width = (_HEADS + 2 * _KV_HEADS) * _HEAD_DIM
-    qkv = {
-        rid: torch.randn(n, width, generator=gen, device="cuda").to(torch.bfloat16)
-        for rid, n in lengths.items()
-    }
-    num_tokens = 2 * sum(-(-n // _PAGE) * _PAGE for n in lengths.values())
-    manager = _kv_manager(kv_dtype, num_tokens)
-    try:
-        attn = TrtllmAttention(
-            layer_idx=0,
-            num_heads=_HEADS,
-            head_dim=_HEAD_DIM,
-            num_kv_heads=_KV_HEADS,
-            quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) if fp8 else None,
-        )
-        rids = sorted(lengths)
-        manager.add_dummy_requests(
-            rids + [_REF_RID + r for r in rids], [lengths[r] for r in rids] * 2
-        )
-        # The unchunked run: every request as one fresh context.
-        rows = [(_REF_RID + r, 0, lengths[r]) for r in rids]
-        out = _context_step(attn, manager, rows, torch.cat([qkv[r] for r in rids]))
-        unchunked, first = {}, 0
-        for r in rids:
-            unchunked[r] = out[first : first + lengths[r]]
-            first += lengths[r]
-            error = _relative_error(unchunked[r], _float_reference(qkv[r], 0, lengths[r], 0, fp8))
-            print(f"\n[{name}] request {r} unchunked vs float reference: {error:.4f}")
-            assert error <= tolerance["float"]
-
-        for step in steps:
-            out = _context_step(
-                attn, manager, step, torch.cat([qkv[r][c : c + n] for r, c, n in step])
-            )
-            first = 0
-            for r, c, n in step:
-                got = out[first : first + n]
-                first += n
-                assert torch.isfinite(got.float()).all()
-                error = _relative_error(got, unchunked[r][c : c + n])
-                print(f"[{name}] request {r} rows {c}..{c + n} vs unchunked: {error:.4f}")
-                assert error <= tolerance["unchunked"]
-                if c:
-                    # A kernel that dropped the cached prefix would match attention over the chunk alone.
-                    dropped = _float_reference(qkv[r], c, n, c, fp8)
-                    assert _relative_error(got, dropped) > 4 * tolerance["float"]
-        assert served and all(served), "the context steps ran through FallbackFmha"
-        torch.cuda.synchronize()
-        for r in rids:
-            assert torch.equal(
-                _cached_kv(manager, r, lengths[r]), _cached_kv(manager, _REF_RID + r, lengths[r])
-            ), f"request {r}: the chunks' K / V differ from the unchunked run's"
-    finally:
-        manager.shutdown()
diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py
new file mode 100644
index 000000000000..db6106ba727e
--- /dev/null
+++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py
@@ -0,0 +1,366 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""FlashInfer TRTLLM-Gen FMHA at 128-token pages on SM107 (Rubin).
+
+On SM107 the FallbackFmha (thop) has no paged-context kernels for 128-token pages,
+so the MiniMax-M3 KV managers list 128 as a TRTLLM-Gen extra page size and
+FlashInferTrtllmGenFmha serves those steps (and only those). These tests run the
+real forward at the MiniMax-M3 Eagle3 draft layer shape (16 q : 1 kv head,
+head dim 128, BF16 Q / output) with BF16 and FP8 KV and check, for every step:
+
+- FlashInfer served it (run_context / run_generation) and FallbackFmha did not;
+- the output is within a tolerance of an FP32 reference built from the pool
+  contents (FlashInfer is precision-equivalent to the thop P32 kernels, not
+  bitwise, so there is no bit-exact oracle);
+
+over fresh, cached-prefix and chunked contexts (each chunk also matches the
+same rows run unchunked and leaves the same K / V), generation (q_len 1),
+linear Eagle3 verify (q_len 1 + 3) and mixed context + generation batches.
+
+The masked tail of a request's last page is read by the kernel: NaN bytes there
+turn the rows that attend to that page into NaN (measured on SM107 for both this
+route and the thop kernels). MiniMax-M3 guards against it by zeroing its KV pools
+at allocation (TRTLLM_M3_KV_POOL_ZERO_INIT, test_minimax_m3_kv_pool_zero_init.py);
+the "nan" tail cases check that rows of requests without a tail stay exact.
+"""
+
+import pytest
+import torch
+
+import tensorrt_llm
+from tensorrt_llm._torch.attention.backends.fmha import fallback
+from tensorrt_llm._torch.attention.backends.fmha import flashinfer_trtllm_gen as fi_mod
+from tensorrt_llm._torch.attention.backends.interface import (
+    AttentionRuntimeFeatures,
+    PredefinedAttentionMask,
+)
+from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata
+from tensorrt_llm._torch.metadata import KVCacheParams
+from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
+from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm.bindings import DataType
+from tensorrt_llm.llmapi.llm_args import KvCacheConfig
+from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+_ON_SM107 = torch.cuda.is_available() and get_sm_version() == 107
+pytestmark = pytest.mark.skipif(
+    not _ON_SM107, reason="FlashInfer serves 128-token pages only on SM107 in this route"
+)
+
+# The MiniMax-M3 Eagle3 draft attention shape on one rank.
+_HEADS, _KV_HEADS, _HEAD_DIM, _PAGE = 16, 1, 128, 128
+_REF_RID = 100  # the unchunked run of request r is request _REF_RID + r
+_KV_DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
+_KV_IDS = ["bf16_kv", "fp8_kv"]
+
+# Worst row's max |out - ref| / max |ref|. "float" is vs the FP32 reference (FP8 KV context
+# also quantizes Q to E4M3, which the reference mirrors; FP8 KV generation keeps BF16 Q);
+# "unchunked" is a chunk vs the same rows of the unchunked run (same cached K / V).
+# FP8 bounds sit above the SM107 measurements (Hecate 822960: context 0.175 with an
+# unquantized-Q reference, generation 0.077, verify 0.094, mixed generation 0.089).
+_TOLERANCE = {
+    torch.bfloat16: {"ctx": 2e-2, "gen": 2e-2, "unchunked": 1e-2},
+    torch.float8_e4m3fn: {"ctx": 1.5e-1, "gen": 1.2e-1, "unchunked": 3e-2},
+}
+
+# (case id, steps): each step is a batch of (request id, cached tokens, new tokens) context rows, run in order.
+_CACHED_CONTEXT_CASES = [
+    ("fresh_512", [[(0, 0, 512)]]),
+    ("prefix_1000_unaligned", [[(0, 0, 1000)], [(0, 1000, 700)]]),
+    ("prefix_1024_aligned", [[(0, 0, 1024)], [(0, 1024, 300)]]),
+    ("chunked_6096_at_2048", [[(0, 0, 2048)], [(0, 2048, 2048)], [(0, 4096, 2000)]]),
+    ("prefix_beside_fresh", [[(0, 0, 384)], [(0, 384, 257), (1, 0, 129)]]),
+]
+
+_PASTS = [2047, 16383, 165000, 640]
+# (case id, context rows [(cached, new)], generation rows [(past, q_len)]); contexts come first, as in the runtime.
+_GENERATION_CASES = [
+    ("gen_q1", [], [(p, 1) for p in _PASTS]),
+    ("verify_1p3", [], [(p, 4) for p in _PASTS]),
+    ("mixed_q1", [(0, 2048), (16384, 2048)], [(640, 1), (16383, 1), (165000, 1)]),
+    ("mixed_1p3", [(0, 2048), (16384, 2048)], [(640, 4), (16383, 4), (165000, 4)]),
+]
+
+
+@pytest.fixture
+def served(monkeypatch):
+    """Count the steps each FMHA library ran."""
+    counts = {"fi_ctx": 0, "fi_gen": 0, "fallback": 0}
+    for cls, name, key in (
+        (fi_mod.FlashInferTrtllmGenFmha, "run_context", "fi_ctx"),
+        (fi_mod.FlashInferTrtllmGenFmha, "run_generation", "fi_gen"),
+        (fallback.FallbackFmha, "forward", "fallback"),
+    ):
+        original = getattr(cls, name)
+
+        def wrapped(self, *args, _original=original, _key=key, **kwargs):
+            counts[_key] += 1
+            return _original(self, *args, **kwargs)
+
+        monkeypatch.setattr(cls, name, wrapped)
+    return counts
+
+
+def _attention(kv_dtype: torch.dtype) -> TrtllmAttention:
+    assert fi_mod.FlashInferTrtllmGenFmha.is_environment_available(), (
+        "SM107 serves 128-token pages only through FlashInfer TRTLLM-Gen FMHA"
+    )
+    fp8 = kv_dtype == torch.float8_e4m3fn
+    return TrtllmAttention(
+        layer_idx=0,
+        num_heads=_HEADS,
+        head_dim=_HEAD_DIM,
+        num_kv_heads=_KV_HEADS,
+        quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) if fp8 else None,
+    )
+
+
+def _kv_manager(kv_dtype: torch.dtype, num_tokens: int) -> KVCacheManager:
+    manager = KVCacheManager(
+        KvCacheConfig(max_tokens=num_tokens, enable_block_reuse=False),
+        tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
+        num_layers=1,
+        num_kv_heads=_KV_HEADS,
+        head_dim=_HEAD_DIM,
+        tokens_per_block=_PAGE,
+        max_seq_len=num_tokens,
+        max_batch_size=8,
+        mapping=Mapping(world_size=1, tp_size=1, rank=0),
+        dtype=DataType.FP8 if kv_dtype == torch.float8_e4m3fn else DataType.BF16,
+    )
+    # As MiniMax-M3's KV managers list it: this is what admits 128-token pages to FlashInfer on SM107.
+    manager.trtllm_gen_extra_tokens_per_block = frozenset({_PAGE})
+    # Zero the pool, as MiniMax-M3's managers do at allocation (see the module docstring).
+    _raw_pool(manager).zero_()
+    return manager
+
+
+def _raw_pool(manager) -> torch.Tensor:
+    """The layer's pool as bytes: [pages, K/V, kv heads, page tokens, head dim bytes]."""
+    return manager.get_buffers(0, kv_layout="HND").view(torch.uint8)
+
+
+def _token_slots(manager, rid: int, length: int):
+    """Pool (page, offset) of tokens 0..length-1 of request rid."""
+    pages = manager.get_batch_cache_indices([rid], 0)[0]
+    blocks = torch.tensor(pages, device="cuda", dtype=torch.long)
+    tokens = torch.arange(length, device="cuda")
+    return blocks[tokens // _PAGE], tokens % _PAGE
+
+
+def _step(attn, manager, rows, qkv: torch.Tensor, num_contexts: int) -> torch.Tensor:
+    """One step as the serving runner prepares it. rows = [(request id, cached, new)], contexts first.
+
+    A context row's prompt length is its current chunk, with the prefix in num_cached_tokens_per_seq;
+    a generation row's prompt length is its past.
+    """
+    metadata = TrtllmAttentionMetadata(
+        max_num_requests=8,
+        max_num_tokens=max(8192, qkv.shape[0] + 8),
+        kv_cache_manager=manager,
+        runtime_features=AttentionRuntimeFeatures(chunked_prefill=True),
+        num_contexts=num_contexts,
+        seq_lens=torch.tensor([n for _, _, n in rows], dtype=torch.int),
+        request_ids=[rid for rid, _, _ in rows],
+        prompt_lens=[n if i < num_contexts else c for i, (_, c, n) in enumerate(rows)],
+        kv_cache_params=KVCacheParams(
+            use_cache=True, num_cached_tokens_per_seq=[c for _, c, _ in rows]
+        ),
+    )
+    assert metadata.tokens_per_block == _PAGE
+    metadata.prepare()
+    out = attn.forward(qkv, None, None, metadata, attention_mask=PredefinedAttentionMask.CAUSAL)
+    return out[0] if isinstance(out, tuple) else out
+
+
+def _reference_from_pool(
+    manager, kv_dtype, rid: int, q: torch.Tensor, cached: int, quantize_q: bool
+) -> torch.Tensor:
+    """FP32 causal attention of rows cached..cached+n over the K / V the pool holds for tokens 0..cached+n."""
+    n = q.shape[0]
+    pages, offsets = _token_slots(manager, rid, cached + n)
+    kv = _raw_pool(manager)[pages, :, :, offsets, :].view(kv_dtype).float()  # [T, K/V, kv heads, D]
+    group = _HEADS // _KV_HEADS
+    k = kv[:, 0].repeat_interleave(group, dim=1)  # [T, heads, D]
+    v = kv[:, 1].repeat_interleave(group, dim=1)
+    if quantize_q:
+        q = q.to(torch.float8_e4m3fn)
+    q = q.float().view(n, _HEADS, _HEAD_DIM)
+    scores = torch.einsum("nhd,thd->nht", q, k) * _HEAD_DIM**-0.5
+    q_pos = cached + torch.arange(n, device="cuda")
+    k_pos = torch.arange(cached + n, device="cuda")
+    scores = scores.masked_fill(k_pos[None, None, :] > q_pos[:, None, None], float("-inf"))
+    return torch.einsum("nht,thd->nhd", scores.softmax(-1), v).reshape(n, _HEADS * _HEAD_DIM)
+
+
+def _float_reference(qkv: torch.Tensor, start: int, n: int, kv_from: int, fp8: bool):
+    """Causal attention of rows start..start+n over keys kv_from..start+n of one request's QKV, in FP32."""
+    q_size, kv_size = _HEADS * _HEAD_DIM, _KV_HEADS * _HEAD_DIM
+    rows = qkv[kv_from : start + n]
+    if fp8:
+        rows = rows.to(torch.float8_e4m3fn)
+    rows = rows.float()
+    q = rows[start - kv_from :, :q_size].view(n, _HEADS, _HEAD_DIM).transpose(0, 1)
+    k = rows[:, q_size : q_size + kv_size].view(-1, _KV_HEADS, _HEAD_DIM).transpose(0, 1)
+    v = rows[:, q_size + kv_size :].view(-1, _KV_HEADS, _HEAD_DIM).transpose(0, 1)
+    k = k.repeat_interleave(_HEADS // _KV_HEADS, dim=0)
+    v = v.repeat_interleave(_HEADS // _KV_HEADS, dim=0)
+    scores = q @ k.transpose(1, 2) / _HEAD_DIM**0.5
+    q_pos = torch.arange(start, start + n, device=qkv.device)[:, None]
+    k_pos = torch.arange(kv_from, start + n, device=qkv.device)[None, :]
+    scores = scores.masked_fill(k_pos > q_pos, float("-inf"))
+    return (scores.softmax(-1) @ v).transpose(0, 1).reshape(n, q_size)
+
+
+def _relative_error(out: torch.Tensor, ref: torch.Tensor) -> float:
+    out, ref = out.float(), ref.float()
+    scale = ref.abs().amax(dim=1).clamp_min(1e-6)
+    return ((out - ref).abs().amax(dim=1) / scale).max().item()
+
+
+def _cached_kv(manager, rid: int, length: int) -> torch.Tensor:
+    pages, offsets = _token_slots(manager, rid, length)
+    return _raw_pool(manager)[pages, :, :, offsets, :]
+
+
+@pytest.mark.parametrize("kv_dtype", _KV_DTYPES, ids=_KV_IDS)
+@pytest.mark.parametrize("case", _CACHED_CONTEXT_CASES, ids=[c[0] for c in _CACHED_CONTEXT_CASES])
+def test_cached_context_matches_unchunked(served, case, kv_dtype) -> None:
+    """FlashInfer attends to the cached prefix: each chunk's output matches the same rows of the request run
+    in one unchunked step and an FP32 reference, and the chunks leave the same K / V in the cache."""
+    name, steps = case
+    fp8 = kv_dtype == torch.float8_e4m3fn
+    tolerance = _TOLERANCE[kv_dtype]
+    lengths = {}
+    for step in steps:
+        for rid, c, n in step:
+            assert c == lengths.get(rid, 0), "each chunk starts at its request's cached length"
+            lengths[rid] = c + n
+
+    gen = torch.Generator(device="cuda").manual_seed(len(name))
+    width = (_HEADS + 2 * _KV_HEADS) * _HEAD_DIM
+    qkv = {
+        rid: torch.randn(n, width, generator=gen, device="cuda").to(torch.bfloat16)
+        for rid, n in lengths.items()
+    }
+    num_tokens = 2 * sum(-(-n // _PAGE) * _PAGE for n in lengths.values()) + 8 * _PAGE
+    manager = _kv_manager(kv_dtype, num_tokens)
+    try:
+        attn = _attention(kv_dtype)
+        rids = sorted(lengths)
+        manager.add_dummy_requests(
+            rids + [_REF_RID + r for r in rids], [lengths[r] for r in rids] * 2
+        )
+        # The unchunked run: every request as one fresh context.
+        rows = [(_REF_RID + r, 0, lengths[r]) for r in rids]
+        out = _step(attn, manager, rows, torch.cat([qkv[r] for r in rids]), len(rows))
+        unchunked, first = {}, 0
+        for r in rids:
+            unchunked[r] = out[first : first + lengths[r]]
+            first += lengths[r]
+            error = _relative_error(unchunked[r], _float_reference(qkv[r], 0, lengths[r], 0, fp8))
+            print(f"\n[{name}] request {r} unchunked vs FP32 reference: {error:.4f}")
+            assert error <= tolerance["ctx"]
+
+        for step in steps:
+            out = _step(
+                attn,
+                manager,
+                step,
+                torch.cat([qkv[r][c : c + n] for r, c, n in step]),
+                len(step),
+            )
+            first = 0
+            for r, c, n in step:
+                got = out[first : first + n]
+                first += n
+                assert torch.isfinite(got.float()).all()
+                error = _relative_error(got, unchunked[r][c : c + n])
+                print(f"[{name}] request {r} rows {c}..{c + n} vs unchunked: {error:.4f}")
+                assert error <= tolerance["unchunked"]
+                reference = _float_reference(qkv[r], c, n, 0, fp8)
+                assert _relative_error(got, reference) <= tolerance["ctx"]
+                if c:
+                    # A kernel that dropped the cached prefix would match attention over the chunk alone.
+                    dropped = _float_reference(qkv[r], c, n, c, fp8)
+                    assert _relative_error(got, dropped) > 4 * tolerance["ctx"]
+        assert served["fi_ctx"] > 0 and served["fallback"] == 0, served
+        torch.cuda.synchronize()
+        for r in rids:
+            assert torch.equal(
+                _cached_kv(manager, r, lengths[r]), _cached_kv(manager, _REF_RID + r, lengths[r])
+            ), f"request {r}: the chunks' K / V differ from the unchunked run's"
+    finally:
+        manager.shutdown()
+
+
+@pytest.mark.parametrize("tail", ["zero", "nan"])
+@pytest.mark.parametrize("kv_dtype", _KV_DTYPES, ids=_KV_IDS)
+@pytest.mark.parametrize("case", _GENERATION_CASES, ids=[c[0] for c in _GENERATION_CASES])
+def test_generation_verify_and_mixed_steps(served, case, kv_dtype, tail) -> None:
+    """Generation (q_len 1), linear verify (1 + 3) and mixed context + generation steps over long cached
+    contexts: FlashInfer serves every phase and each row is within tolerance of an FP32 reference from the
+    pool contents. With NaN bytes in a last page's unwritten tail, rows of requests whose last page is full
+    stay exact; the other requests' NaN rows are what the KV-pool zero-init prevents."""
+    name, ctx, gen_rows = case
+    fp8 = kv_dtype == torch.float8_e4m3fn
+    tolerance = _TOLERANCE[kv_dtype]
+    reqs = list(ctx) + list(gen_rows)  # (cached, new) per request
+    num_contexts = len(ctx)
+    totals = [c + n for c, n in reqs]
+    num_tokens = sum(-(-t // _PAGE) * _PAGE for t in totals) + 8 * _PAGE
+    manager = _kv_manager(kv_dtype, num_tokens)
+    try:
+        attn = _attention(kv_dtype)
+        rids = list(range(len(reqs)))
+        manager.add_dummy_requests(rids, totals)
+        raw = _raw_pool(manager)
+        gen = torch.Generator(device="cuda").manual_seed(5)
+        for rid, (c, n) in zip(rids, reqs):
+            if c:
+                # Cached K / V written straight into the pool.
+                kv = torch.randn(2, c, _KV_HEADS, _HEAD_DIM, generator=gen, device="cuda") * 2
+                pages, offsets = _token_slots(manager, rid, c)
+                for s in (0, 1):
+                    raw[pages, s, :, offsets, :] = kv[s].to(kv_dtype).view(torch.uint8)
+            if tail == "nan" and (c + n) % _PAGE:
+                last_page = manager.get_batch_cache_indices([rid], 0)[0][(c + n - 1) // _PAGE]
+                raw[last_page, :, :, (c + n) % _PAGE :, :] = 0xFF  # NaN in BF16 and E4M3
+        width = (_HEADS + 2 * _KV_HEADS) * _HEAD_DIM
+        qkv = torch.cat(
+            [
+                torch.randn(n, width, generator=gen, device="cuda").to(torch.bfloat16)
+                for _, n in reqs
+            ]
+        )
+        rows = [(rid, c, n) for rid, (c, n) in zip(rids, reqs)]
+        out = _step(attn, manager, rows, qkv, num_contexts)
+        torch.cuda.synchronize()
+
+        assert served["fallback"] == 0, served
+        assert (served["fi_ctx"] > 0) == (num_contexts > 0), served
+        assert served["fi_gen"] > 0, served
+
+        first = 0
+        for i, (rid, (c, n)) in enumerate(zip(rids, reqs)):
+            got = out[first : first + n].float()
+            q = qkv[first : first + n, : _HEADS * _HEAD_DIM]
+            first += n
+            is_context = i < num_contexts
+            has_nan_tail = tail == "nan" and (c + n) % _PAGE != 0
+            nonfinite = int((~torch.isfinite(got)).any(dim=1).sum())
+            if has_nan_tail:
+                print(f"[{name}] request {rid} (tail NaN): {nonfinite}/{n} non-finite rows")
+                continue
+            assert nonfinite == 0, f"request {rid}: {nonfinite}/{n} non-finite rows"
+            ref = _reference_from_pool(manager, kv_dtype, rid, q, c, quantize_q=fp8 and is_context)
+            error = _relative_error(got, ref)
+            phase = "ctx" if is_context else "gen"
+            print(
+                f"[{name}] request {rid} ({phase}, past {c}, q_len {n}) vs FP32 reference: {error:.4f}"
+            )
+            assert error <= tolerance[phase]
+    finally:
+        manager.shutdown()
diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py
index 45c2d427c2fb..43ff606c38b3 100644
--- a/tests/unittest/_torch/models/test_minimax_m3.py
+++ b/tests/unittest/_torch/models/test_minimax_m3.py
@@ -2716,15 +2716,19 @@ def test_minimax_m3_swiglu_oai_fused_matches_reference(dtype):
 
 @pytest.mark.cpu_only
 @pytest.mark.parametrize(
-    "libraries,available,use_native",
+    "libraries,available,extra_pages_only,use_native",
     [
-        ("flashinfer_trtllm_gen,fallback", True, True),
-        ("fallback,flashinfer_trtllm_gen", True, False),
-        ("fallback", True, False),
-        ("flashinfer_trtllm_gen,fallback", False, False),
+        ("flashinfer_trtllm_gen,fallback", True, False, True),
+        ("fallback,flashinfer_trtllm_gen", True, False, False),
+        ("fallback", True, False, False),
+        ("flashinfer_trtllm_gen,fallback", False, False, False),
+        # SM107: the library serves only extra page sizes; dense layers keep SDPA.
+        ("flashinfer_trtllm_gen,fallback", True, True, False),
     ],
 )
-def test_minimax_m3_dense_backend_library_order(monkeypatch, libraries, available, use_native):
+def test_minimax_m3_dense_backend_library_order(
+    monkeypatch, libraries, available, extra_pages_only, use_native
+):
     from types import SimpleNamespace
     from unittest.mock import Mock
 
@@ -2745,6 +2749,11 @@ def test_minimax_m3_dense_backend_library_order(monkeypatch, libraries, availabl
         "is_available",
         classmethod(lambda cls, attn: available),
     )
+    monkeypatch.setattr(
+        model_module.FlashInferTrtllmGenFmha,
+        "serves_only_extra_page_sizes",
+        classmethod(lambda cls: extra_pages_only),
+    )
     layer = SimpleNamespace(
         attn=legacy,
         layer_idx=0,
@@ -2762,7 +2771,7 @@ def test_minimax_m3_dense_backend_library_order(monkeypatch, libraries, availabl
     )
     assert layer.attn is (native if use_native else legacy)
     assert layer.support_fused_qkv is use_native
-    if libraries == "flashinfer_trtllm_gen,fallback" and available:
+    if libraries == "flashinfer_trtllm_gen,fallback" and available and not extra_pages_only:
         construct.assert_called_once()
         assert construct.call_args.kwargs["pos_embd_params"] is None
     else:

From f36816ad9e08def2749e45887f1b0f54b00461c6 Mon Sep 17 00:00:00 2001
From: peihengh <259410613+peihu-nv@users.noreply.github.com>
Date: Thu, 22 Oct 2026 07:49:04 -0700
Subject: [PATCH 4/9] fail fast when the SM107 P128 opt-in has no FlashInfer
 route

Check at KV cache manager setup (separate draft manager, target native
layers, NVFP4 draft view) that flashinfer_trtllm_gen is enabled and
available on SM107, and raise an error naming the fix instead of failing
at the first prefill. UT covers the OK and missing cases.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
---
 .../backends/fmha/flashinfer_trtllm_gen.py    | 26 ++++++++
 .../sparse/minimax_m3/cache_manager.py        | 13 ++++
 .../test_minimax_m3_shared_draft_layers.py    | 62 +++++++++++++++++++
 3 files changed, 101 insertions(+)

diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
index 47652e5b9efa..bca658689cec 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
@@ -425,6 +425,32 @@ 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 require_extra_page_size_route(cls, tokens_per_block: int) -> None:
+        """Raise at setup when a KV manager's extra page size has no FMHA on this device.
+
+        On SM107 this library is the only one with paged-context kernels for the
+        extra page sizes, so without it the first prefill would fail. Call this
+        where a KV cache manager opts in to ``tokens_per_block``.
+        """
+        if not cls.serves_only_extra_page_sizes():
+            return
+        from .registry import get_enabled_fmha_lib_classes
+
+        where = f"SM{get_sm_version()} with {tokens_per_block}-token KV pages"
+        if cls not in get_enabled_fmha_lib_classes():
+            raise RuntimeError(
+                f"{where} 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)."
+            )
+        if not cls.is_environment_available():
+            raise RuntimeError(
+                f"{where} needs the flashinfer_trtllm_gen FMHA library, but it is unavailable: "
+                "install flashinfer-python with its fused nanobind ops (see the debug log for "
+                "the missing piece)."
+            )
+
     @classmethod
     def _is_available(cls, attn: "TrtllmAttention") -> bool:
         if not cls.is_environment_available():
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 6dd95c6c1cf1..3e0dcef7286d 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
@@ -233,6 +233,15 @@ def extend_attention_op_pools_for_shared_draft_layers(
 _KV_POOL_ZERO_INIT = os.environ.get("TRTLLM_M3_KV_POOL_ZERO_INIT", "1") == "1"
 
 
+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)."""
+    from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import (
+        FlashInferTrtllmGenFmha,
+    )
+
+    FlashInferTrtllmGenFmha.require_extra_page_size_route(tokens_per_block)
+
+
 class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2):
     """KVCacheManagerV2 subclass with a V2-managed paged index-K cache
     per sparse layer.
@@ -565,6 +574,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
             # 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:
@@ -582,6 +592,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
         # 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:
+            _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.
@@ -1059,6 +1070,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])
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 227b6c24d459..c5506f059695 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
@@ -416,6 +416,68 @@ def test_draft_kv_cache_view_accessor_respects_local_layers(
     assert manager._shared_draft_layer_ids == [60, 61]
 
 
+@pytest.mark.cpu_only
+@pytest.mark.parametrize(
+    "sm,fmha_libs,env_ok,error",
+    [
+        (107, None, True, None),
+        (107, "-flashinfer_trtllm_gen", True, "TLLM_FMHA_LIBS excludes it"),
+        (107, "fallback", True, "TLLM_FMHA_LIBS excludes it"),
+        (107, None, False, "install flashinfer-python"),
+        # Elsewhere the fallback has P128 paged-context kernels, so nothing is required.
+        (103, "-flashinfer_trtllm_gen", False, None),
+    ],
+)
+@pytest.mark.parametrize("site", ["separate_draft_manager", "nvfp4_draft_view"])
+def test_p128_opt_in_fails_fast_without_the_flashinfer_route(
+    monkeypatch, sm, fmha_libs, env_ok, error, site
+):
+    """On SM107 the P128 opt-in is checked when the KV manager is set up, not at the
+    first prefill: a missing FlashInfer route raises an error that names the fix."""
+    monkeypatch.setattr(
+        "tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen.get_sm_version",
+        lambda: sm,
+    )
+    monkeypatch.setattr(
+        FlashInferTrtllmGenFmha, "is_environment_available", classmethod(lambda cls: env_ok)
+    )
+    if fmha_libs is None:
+        monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False)
+    else:
+        monkeypatch.setenv("TLLM_FMHA_LIBS", fmha_libs)
+
+    if site == "separate_draft_manager":
+        monkeypatch.setattr(
+            KVCacheManagerV2,
+            "_prepare_page_table_tensor",
+            lambda self, capacity: setattr(self, "_use_per_layer_page_tables", False),
+        )
+        manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2)
+        manager.dtype = DataType.FP8
+        manager._shared_draft_layer_ids = []
+        manager.layer_offsets = {0: 0}
+        manager.is_draft = True
+        manager.enable_swa_scratch_reuse = False
+        manager.tokens_per_block = 128
+        setup = lambda: manager._prepare_page_table_tensor(8)  # noqa: E731
+    else:
+        # Stop right after the check: the rest of the view needs real pools.
+        manager = SimpleNamespace(tokens_per_block=128, layer_offsets={})
+        setup = lambda: MiniMaxM3DraftKVCacheView(manager, [DRAFT_LOCAL_LAYER])  # noqa: E731
+
+    if error is not None:
+        with pytest.raises(RuntimeError, match=error) as excinfo:
+            setup()
+        assert "128-token KV pages" in str(excinfo.value)
+    elif site == "separate_draft_manager":
+        setup()
+        assert manager.trtllm_gen_extra_tokens_per_block == frozenset({128})
+    else:
+        # The check passed; the view then fails on the stub's missing layer.
+        with pytest.raises(KeyError):
+            setup()
+
+
 @pytest.mark.cpu_only
 @pytest.mark.parametrize("via_accessor", [False, True])
 def test_draft_kv_cache_view_rejects_multiple_local_layers(via_accessor: bool) -> None:

From c968f35daeeda598075eacae810e839f80d2a743 Mon Sep 17 00:00:00 2001
From: peihengh <259410613+peihu-nv@users.noreply.github.com>
Date: Thu, 22 Oct 2026 08:45:10 -0700
Subject: [PATCH 5/9] derive the SM107 P128 oracle tolerances from a seed sweep
 at 16:1

Set each bound to 1.25x the worst case over a many-seed sweep at the
MiniMax-M3 draft shape (32 seeds per context case, 16 seeds x zero/NaN
tails per generation case) and record the per-case maxima. Chunks now
match the unchunked run bitwise; the BF16 and FP8 context bounds tighten;
the FP8 generation bound rises from 0.12 to 0.18 (worst 0.143 at a
165000-token past). The dropped-prefix margin stays 4x.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
---
 .../test_flashinfer_trtllm_gen_sm107_p128.py  | 21 +++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py
index db6106ba727e..117bf9e71d97 100644
--- a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py
+++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py
@@ -57,13 +57,22 @@
 
 # Worst row's max |out - ref| / max |ref|. "float" is vs the FP32 reference (FP8 KV context
 # also quantizes Q to E4M3, which the reference mirrors; FP8 KV generation keeps BF16 Q);
-# "unchunked" is a chunk vs the same rows of the unchunked run (same cached K / V).
-# FP8 bounds sit above the SM107 measurements (Hecate 822960: context 0.175 with an
-# unquantized-Q reference, generation 0.077, verify 0.094, mixed generation 0.089).
+# "unchunked" is a chunk vs the same rows of the unchunked run (same cached K / V): 0 means
+# bitwise equal. Each bound is 1.25x the worst case of a seed sweep at this 16:1 shape
+# (Hecate 826693: 32 seeds per context case, 16 seeds x {zero, NaN} tails per generation
+# case), rounded up. Per-case worst (BF16 / FP8 KV):
+#   context rows: fresh 0.0048 / 0.045, cached prefixes 0.0053 / 0.047, chunked 0.0052 / 0.046,
+#     mixed batches 0.0061 / 0.049;
+#   generation rows: q_len 1 0.0054 / 0.094, verify 1+3 0.0057 / 0.116, mixed 0.0057 / 0.143.
+#     FP8 grows with the past (q_len 1: 0.083 at 640 -> 0.143 at 165000); BF16 does not;
+#   chunks vs unchunked: 0 for every seed and case.
 _TOLERANCE = {
-    torch.bfloat16: {"ctx": 2e-2, "gen": 2e-2, "unchunked": 1e-2},
-    torch.float8_e4m3fn: {"ctx": 1.5e-1, "gen": 1.2e-1, "unchunked": 3e-2},
+    torch.bfloat16: {"ctx": 7.7e-3, "gen": 7.2e-3, "unchunked": 0.0},
+    torch.float8_e4m3fn: {"ctx": 6.2e-2, "gen": 1.8e-1, "unchunked": 0.0},
 }
+# A kernel that dropped the cached prefix would match attention over the chunk alone; in the
+# sweep that error was >= 1.05, at least 4x above every "ctx" bound.
+_DROPPED_PREFIX_MARGIN = 4
 
 # (case id, steps): each step is a batch of (request id, cached tokens, new tokens) context rows, run in order.
 _CACHED_CONTEXT_CASES = [
@@ -285,7 +294,7 @@ def test_cached_context_matches_unchunked(served, case, kv_dtype) -> None:
                 if c:
                     # A kernel that dropped the cached prefix would match attention over the chunk alone.
                     dropped = _float_reference(qkv[r], c, n, c, fp8)
-                    assert _relative_error(got, dropped) > 4 * tolerance["ctx"]
+                    assert _relative_error(got, dropped) > _DROPPED_PREFIX_MARGIN * tolerance["ctx"]
         assert served["fi_ctx"] > 0 and served["fallback"] == 0, served
         torch.cuda.synchronize()
         for r in rids:

From 620402cc0a23cc569ef4659eb6bfec6562e5a920 Mon Sep 17 00:00:00 2001
From: peihengh <259410613+peihu-nv@users.noreply.github.com>
Date: Thu, 22 Oct 2026 08:45:40 -0700
Subject: [PATCH 6/9] document the SM107 P128 FlashInfer cubin set and offline
 seeding

Replace the single-rank warm-up hint with a warm-up at the production
TP size and KV dtype (the cubins depend on the per-rank head layout), and
list the 10 cubins the FP8 KV draft layer at TP2 needs.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
---
 .../deployment-guide-for-minimax-m3-on-trtllm.md             | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

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 f3413a6f8fae..02f68eb40184 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,7 +65,10 @@ 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, populate a shared `FLASHINFER_CUBIN_DIR` before the multi-rank launch (for example with a single-rank warm-up run that serves one Eagle3 request) and set `FLASHINFER_NO_DOWNLOAD=1`; see [FlashInfer JIT workspaces for MPI workers](../llm-api/index.md#flashinfer-jit-workspaces-for-mpi-workers).
+* **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.

From ed19a8a68fbe4898fb322a582ffb128c05080040 Mon Sep 17 00:00:00 2001
From: peihengh <259410613+peihu-nv@users.noreply.github.com>
Date: Thu, 22 Oct 2026 09:07:01 -0700
Subject: [PATCH 7/9] fail fast when TLLM_FMHA_LIBS tries fallback before the
 SM107 P128 FlashInfer route

Libraries are tried in TLLM_FMHA_LIBS order and FallbackFmha raises on a
paged prefill at an extra page size, so a fallback-first list passed the
setup check but failed at the first prefill. Reject that order at KV-manager
setup with an error that names the fix.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
---
 .../backends/fmha/flashinfer_trtllm_gen.py        | 15 ++++++++++++++-
 .../msa/test_minimax_m3_shared_draft_layers.py    |  3 +++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
index bca658689cec..026c1d6972ca 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
@@ -435,15 +435,28 @@ def require_extra_page_size_route(cls, tokens_per_block: int) -> None:
         """
         if not cls.serves_only_extra_page_sizes():
             return
+        from .fallback import FallbackFmha
         from .registry import get_enabled_fmha_lib_classes
 
         where = f"SM{get_sm_version()} with {tokens_per_block}-token KV pages"
-        if cls not in get_enabled_fmha_lib_classes():
+        enabled_libs = get_enabled_fmha_lib_classes()
+        if cls not in enabled_libs:
             raise RuntimeError(
                 f"{where} 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)."
             )
+        # Libraries are tried in TLLM_FMHA_LIBS order, and FallbackFmha raises on a paged
+        # prefill at an extra page size, so it must not come first.
+        if FallbackFmha in enabled_libs and enabled_libs.index(FallbackFmha) < enabled_libs.index(
+            cls
+        ):
+            raise RuntimeError(
+                f"{where} needs the flashinfer_trtllm_gen FMHA library to be tried before "
+                "fallback, which cannot serve paged prefill at this page size, but "
+                "TLLM_FMHA_LIBS lists fallback first. "
+                "List flashinfer_trtllm_gen before fallback in TLLM_FMHA_LIBS (or unset it)."
+            )
         if not cls.is_environment_available():
             raise RuntimeError(
                 f"{where} needs the flashinfer_trtllm_gen FMHA library, but it is unavailable: "
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 c5506f059695..22269a9bce8e 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
@@ -423,6 +423,9 @@ def test_draft_kv_cache_view_accessor_respects_local_layers(
         (107, None, True, None),
         (107, "-flashinfer_trtllm_gen", True, "TLLM_FMHA_LIBS excludes it"),
         (107, "fallback", True, "TLLM_FMHA_LIBS excludes it"),
+        # Libraries are tried in list order; fallback raises on P128 paged prefill.
+        (107, "fallback,flashinfer_trtllm_gen", True, "lists fallback first"),
+        (107, "flashinfer_trtllm_gen,fallback", True, None),
         (107, None, False, "install flashinfer-python"),
         # Elsewhere the fallback has P128 paged-context kernels, so nothing is required.
         (103, "-flashinfer_trtllm_gen", False, None),

From 8e7b3e3c34025b84451ff9da4d4141496e7f87b6 Mon Sep 17 00:00:00 2001
From: peihengh <259410613+peihu-nv@users.noreply.github.com>
Date: Thu, 22 Oct 2026 09:10:25 -0700
Subject: [PATCH 8/9] require the SM107 P128 FlashInfer route only for shared
 draft layers

On SM107 a target's dense layers keep SDPA, so a BF16 target without a shared
draft layer must not fail setup when FlashInfer is excluded or unavailable.
The P128 opt-in itself is unchanged.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
---
 .../sparse/minimax_m3/cache_manager.py        |  5 ++-
 .../test_minimax_m3_shared_draft_layers.py    | 35 +++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

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 3e0dcef7286d..852462d0e890 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
@@ -592,7 +592,10 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
         # 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:
-            _require_trtllm_gen_extra_page_route(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.
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 22269a9bce8e..d65c7372df3c 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
@@ -481,6 +481,41 @@ def test_p128_opt_in_fails_fast_without_the_flashinfer_route(
             setup()
 
 
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("shared_draft", [False, True])
+def test_p128_target_fails_fast_only_with_a_shared_draft_layer(monkeypatch, shared_draft):
+    """On SM107 a BF16 target's dense layers keep SDPA, so a missing FlashInfer route only
+    fails setup when a shared draft layer needs it."""
+    monkeypatch.setattr(
+        "tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen.get_sm_version",
+        lambda: 107,
+    )
+    monkeypatch.setattr(
+        FlashInferTrtllmGenFmha, "is_environment_available", classmethod(lambda cls: True)
+    )
+    monkeypatch.setenv("TLLM_FMHA_LIBS", "-flashinfer_trtllm_gen")
+    monkeypatch.setattr(
+        KVCacheManagerV2,
+        "_prepare_page_table_tensor",
+        lambda self, capacity: setattr(self, "_use_per_layer_page_tables", True),
+    )
+    manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2)
+    manager.dtype = DataType.HALF
+    manager.is_draft = False
+    manager.enable_swa_scratch_reuse = False
+    manager.tokens_per_block = 128
+    manager._dense_native_layer_ids = (0, 1, 2)
+    manager._shared_draft_layer_ids = [60] if shared_draft else []
+    manager.layer_offsets = {0: 0, 1: 1, 2: 2, **({60: 3} if shared_draft else {})}
+
+    if shared_draft:
+        with pytest.raises(RuntimeError, match="TLLM_FMHA_LIBS excludes it"):
+            manager._prepare_page_table_tensor(8)
+    else:
+        manager._prepare_page_table_tensor(8)
+        assert manager.trtllm_gen_extra_tokens_per_block == frozenset({128})
+
+
 @pytest.mark.cpu_only
 @pytest.mark.parametrize("via_accessor", [False, True])
 def test_draft_kv_cache_view_rejects_multiple_local_layers(via_accessor: bool) -> None:

From cbb1420d4a09fffab2b1c8b773bdadb3e09a3023 Mon Sep 17 00:00:00 2001
From: peihengh <259410613+peihu-nv@users.noreply.github.com>
Date: Thu, 22 Oct 2026 21:26:32 -0700
Subject: [PATCH 9/9] [None][fix] Simplify the MiniMax-M3 SM107 P128 route
 check and always zero the KV pools

Review follow-ups for the Rubin (SM107) MiniMax-M3 enablement:
- Zero the MiniMax-M3 KV pools at allocation unconditionally; drop the
  TRTLLM_M3_KV_POOL_ZERO_INIT switch.
- On SM107, the P128 setup check now only requires flashinfer_trtllm_gen to be
  enabled in TLLM_FMHA_LIBS: its 128-token-page cubins ship with it.
- The model keeps SDPA for its dense layers by checking SM107 directly.
- The NVFP4 KV exemption reuses is_minimax_m3().
- Run the SM107 P128 FlashInfer oracle in the GR100 post-merge stage, tighten
  a KeyError assertion and fix a list paragraph in the deployment guide.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
---
 ...ployment-guide-for-minimax-m3-on-trtllm.md |  1 +
 .../backends/fmha/flashinfer_trtllm_gen.py    | 39 ---------------
 .../sparse/minimax_m3/cache_manager.py        | 49 +++++++++++--------
 .../_torch/models/modeling_minimaxm3.py       |  3 +-
 .../_torch/pyexecutor/model_loader.py         |  6 +--
 .../test_lists/test-db/l0_gr100.yml           |  1 +
 .../test_minimax_m3_shared_draft_layers.py    | 36 +++++---------
 .../test_flashinfer_trtllm_gen_sm107_p128.py  |  2 +-
 .../unittest/_torch/models/test_minimax_m3.py |  6 +--
 9 files changed, 47 insertions(+), 96 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 02f68eb40184..bc23d6f39777 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
@@ -68,6 +68,7 @@ server (see [Launch the TensorRT LLM Server](#launch-the-tensorrt-llm-server)).
 * **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.
diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
index 62b4d93cea56..e465eda7d093 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
@@ -425,45 +425,6 @@ 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 require_extra_page_size_route(cls, tokens_per_block: int) -> None:
-        """Raise at setup when a KV manager's extra page size has no FMHA on this device.
-
-        On SM107 this library is the only one with paged-context kernels for the
-        extra page sizes, so without it the first prefill would fail. Call this
-        where a KV cache manager opts in to ``tokens_per_block``.
-        """
-        if not cls.serves_only_extra_page_sizes():
-            return
-        from .fallback import FallbackFmha
-        from .registry import get_enabled_fmha_lib_classes
-
-        where = f"SM{get_sm_version()} with {tokens_per_block}-token KV pages"
-        enabled_libs = get_enabled_fmha_lib_classes()
-        if cls not in enabled_libs:
-            raise RuntimeError(
-                f"{where} 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)."
-            )
-        # Libraries are tried in TLLM_FMHA_LIBS order, and FallbackFmha raises on a paged
-        # prefill at an extra page size, so it must not come first.
-        if FallbackFmha in enabled_libs and enabled_libs.index(FallbackFmha) < enabled_libs.index(
-            cls
-        ):
-            raise RuntimeError(
-                f"{where} needs the flashinfer_trtllm_gen FMHA library to be tried before "
-                "fallback, which cannot serve paged prefill at this page size, but "
-                "TLLM_FMHA_LIBS lists fallback first. "
-                "List flashinfer_trtllm_gen before fallback in TLLM_FMHA_LIBS (or unset it)."
-            )
-        if not cls.is_environment_available():
-            raise RuntimeError(
-                f"{where} needs the flashinfer_trtllm_gen FMHA library, but it is unavailable: "
-                "install flashinfer-python with its fused nanobind ops (see the debug log for "
-                "the missing piece)."
-            )
-
     @classmethod
     def _is_available(cls, attn: "TrtllmAttention") -> bool:
         if not cls.is_environment_available():
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 852462d0e890..c52c6d665126 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
@@ -28,7 +28,6 @@
 
 from __future__ import annotations
 
-import os
 from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple
 
 import torch
@@ -39,6 +38,7 @@
     TensorWrapper,
     binding_to_torch_dtype,
     convert_to_torch_tensor,
+    get_sm_version,
     prefer_pinned,
 )
 from tensorrt_llm.bindings import DataType
@@ -228,18 +228,26 @@ def extend_attention_op_pools_for_shared_draft_layers(
     )
 
 
-# Zero every KV pool view (NVFP4 + FP8 K/V, NVFP4 block scales, paged index keys) once at allocation, so no reader
-# ever sees uninitialized bytes. Kill switch (diagnostics only): TRTLLM_M3_KV_POOL_ZERO_INIT=0.
-_KV_POOL_ZERO_INIT = os.environ.get("TRTLLM_M3_KV_POOL_ZERO_INIT", "1") == "1"
-
-
 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)."""
+    """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
 
-    FlashInferTrtllmGenFmha.require_extra_page_size_route(tokens_per_block)
+    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):
@@ -414,19 +422,18 @@ def __init__(
                 self._block_scale_buffers[layer_idx, kv_layout] = self.get_block_scale_buffers(
                     layer_idx, kv_layout
                 )
-        if _KV_POOL_ZERO_INIT:
-            # 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_()
+        # 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
diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
index 0bc2dab4145d..adaaaf93682d 100644
--- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py
+++ b/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,7 +1300,7 @@ def _init_native_dense_attention(self, model_config: ModelConfig) -> None:
             return
         if not FlashInferTrtllmGenFmha.is_environment_available():
             return
-        if FlashInferTrtllmGenFmha.serves_only_extra_page_sizes():
+        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.
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index e11d3f87d335..5b6a0514cbd3 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -31,7 +31,6 @@
 from tensorrt_llm._utils import get_sm_version, str_dtype_to_torch
 from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig,
                                           ExecutorMemoryType,
-                                          MiniMaxM3SparseAttentionConfig,
                                           ModelExpressConfig,
                                           SparseAttentionConfig, TorchLlmArgs)
 from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto,
@@ -57,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)
@@ -169,8 +168,7 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig,
     # 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 = (isinstance(sparse_config,
-                                      MiniMaxM3SparseAttentionConfig)
+    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)
diff --git a/tests/integration/test_lists/test-db/l0_gr100.yml b/tests/integration/test_lists/test-db/l0_gr100.yml
index 45cd6e6a35a8..104093738356 100644
--- a/tests/integration/test_lists/test-db/l0_gr100.yml
+++ b/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
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 d65c7372df3c..0225dea33702 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
@@ -54,6 +54,7 @@
 from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode
 
 DRAFT_LOCAL_LAYER = 60
+_CACHE_MANAGER_MODULE = "tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager"
 SCALE = 179  # sub-pages per M3 mega-slot: 3 dense x 2 + 57 sparse x 3 + draft x 2
 DRAFT_K_ADDR = 0x7000_0000
 
@@ -418,32 +419,23 @@ def test_draft_kv_cache_view_accessor_respects_local_layers(
 
 @pytest.mark.cpu_only
 @pytest.mark.parametrize(
-    "sm,fmha_libs,env_ok,error",
+    "sm,fmha_libs,error",
     [
-        (107, None, True, None),
-        (107, "-flashinfer_trtllm_gen", True, "TLLM_FMHA_LIBS excludes it"),
-        (107, "fallback", True, "TLLM_FMHA_LIBS excludes it"),
-        # Libraries are tried in list order; fallback raises on P128 paged prefill.
-        (107, "fallback,flashinfer_trtllm_gen", True, "lists fallback first"),
-        (107, "flashinfer_trtllm_gen,fallback", True, None),
-        (107, None, False, "install flashinfer-python"),
+        (107, None, None),
+        (107, "-flashinfer_trtllm_gen", "TLLM_FMHA_LIBS excludes it"),
+        (107, "fallback", "TLLM_FMHA_LIBS excludes it"),
+        (107, "flashinfer_trtllm_gen,fallback", None),
         # Elsewhere the fallback has P128 paged-context kernels, so nothing is required.
-        (103, "-flashinfer_trtllm_gen", False, None),
+        (103, "-flashinfer_trtllm_gen", None),
     ],
 )
 @pytest.mark.parametrize("site", ["separate_draft_manager", "nvfp4_draft_view"])
 def test_p128_opt_in_fails_fast_without_the_flashinfer_route(
-    monkeypatch, sm, fmha_libs, env_ok, error, site
+    monkeypatch, sm, fmha_libs, error, site
 ):
     """On SM107 the P128 opt-in is checked when the KV manager is set up, not at the
     first prefill: a missing FlashInfer route raises an error that names the fix."""
-    monkeypatch.setattr(
-        "tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen.get_sm_version",
-        lambda: sm,
-    )
-    monkeypatch.setattr(
-        FlashInferTrtllmGenFmha, "is_environment_available", classmethod(lambda cls: env_ok)
-    )
+    monkeypatch.setattr(f"{_CACHE_MANAGER_MODULE}.get_sm_version", lambda: sm)
     if fmha_libs is None:
         monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False)
     else:
@@ -477,7 +469,7 @@ def test_p128_opt_in_fails_fast_without_the_flashinfer_route(
         assert manager.trtllm_gen_extra_tokens_per_block == frozenset({128})
     else:
         # The check passed; the view then fails on the stub's missing layer.
-        with pytest.raises(KeyError):
+        with pytest.raises(KeyError, match=str(DRAFT_LOCAL_LAYER)):
             setup()
 
 
@@ -486,13 +478,7 @@ def test_p128_opt_in_fails_fast_without_the_flashinfer_route(
 def test_p128_target_fails_fast_only_with_a_shared_draft_layer(monkeypatch, shared_draft):
     """On SM107 a BF16 target's dense layers keep SDPA, so a missing FlashInfer route only
     fails setup when a shared draft layer needs it."""
-    monkeypatch.setattr(
-        "tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen.get_sm_version",
-        lambda: 107,
-    )
-    monkeypatch.setattr(
-        FlashInferTrtllmGenFmha, "is_environment_available", classmethod(lambda cls: True)
-    )
+    monkeypatch.setattr(f"{_CACHE_MANAGER_MODULE}.get_sm_version", lambda: 107)
     monkeypatch.setenv("TLLM_FMHA_LIBS", "-flashinfer_trtllm_gen")
     monkeypatch.setattr(
         KVCacheManagerV2,
diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py
index 117bf9e71d97..f6b06575988b 100644
--- a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py
+++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py
@@ -20,7 +20,7 @@
 The masked tail of a request's last page is read by the kernel: NaN bytes there
 turn the rows that attend to that page into NaN (measured on SM107 for both this
 route and the thop kernels). MiniMax-M3 guards against it by zeroing its KV pools
-at allocation (TRTLLM_M3_KV_POOL_ZERO_INIT, test_minimax_m3_kv_pool_zero_init.py);
+at allocation (test_minimax_m3_kv_pool_zero_init.py);
 the "nan" tail cases check that rows of requests without a tail stay exact.
 """
 
diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py
index 8874ab13cfba..5e3ce1d1292b 100644
--- a/tests/unittest/_torch/models/test_minimax_m3.py
+++ b/tests/unittest/_torch/models/test_minimax_m3.py
@@ -2896,11 +2896,7 @@ def test_minimax_m3_dense_backend_library_order(
         "is_available",
         classmethod(lambda cls, attn: available),
     )
-    monkeypatch.setattr(
-        model_module.FlashInferTrtllmGenFmha,
-        "serves_only_extra_page_sizes",
-        classmethod(lambda cls: extra_pages_only),
-    )
+    monkeypatch.setattr(model_module, "get_sm_version", lambda: 107 if extra_pages_only else 100)
     layer = SimpleNamespace(
         attn=legacy,
         layer_idx=0,