From e8a045e5641ab4e4559dcffd0b724e3e4e1a6a09 Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Thu, 10 Sep 2026 10:59:36 -0700
Subject: [PATCH 1/6] [None][test] Cover the FP8 indexer's padded-slot guards

minimaxM3Fp8IndexerQKNormRopeKernel drops a row whose out_cache_loc entry is
negative and one whose page is past the end of the cache. Both guards are
exercised today by a single case with slots [0, -1, 128] over three tokens,
which leaves three gaps.

It asserts only that the pages flanking the cache view stay zero. A stray
store that lands on a valid page at an offset no live row claimed clears both
guard regions and goes unnoticed, so pin every unwritten slot to zero too.

It covers one out-of-range slot, the first page past the end at offset zero.
Add an out-of-range slot with a nonzero page offset and one a further page
out, since page and offset come from separate truncation arithmetic.

Its three slots interleave a sentinel between two valid ones, whereas the
padding this defends against presents a live prefix followed by a sentinel
tail. Run a 17-row extent over a 4-page cache with a 4-row live prefix, and
again with no live rows at all, which is the case where nothing may be
written anywhere.

Slot tails are chosen to land inside the guard region rather than outside the
allocation, so a regression trips an assert instead of faulting: -1 truncates
to page 0 at offset -1, one token below the base, and the largest tail slot
reaches two pages past the last, well inside the two guard pages above.

The file already runs in l0_h100 via the unittest/_torch/thop/
parallel_hw_agnostic entry, so no test list changes.

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 .../test_minimax_m3_fp8_indexer.py            | 80 +++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py
index 042ccd858e6a..bbe8cad52171 100644
--- a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py
+++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py
@@ -78,6 +78,37 @@ def _strided_cache(num_pages: int, page_size: int = 128, stride_scale: int = 7)
     return backing[::stride_scale]
 
 
+def _guarded_cache(
+    num_pages: int, page_size: int = 128, stride_scale: int = 7, guard_pages: int = 2
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Allocate a strided cache flanked by zeroed guard regions.
+
+    Returns the cache view plus the backing below and above it. A store from a
+    negative slot lands below the base and one from an out-of-range slot lands
+    above the last page, so both stay inspectable rather than depending on a
+    fault that an out-of-bounds FP8 store may never raise.
+    """
+    total_pages = num_pages + 2 * guard_pages
+    backing = torch.zeros(
+        total_pages * stride_scale,
+        1,
+        page_size,
+        128,
+        dtype=torch.float8_e4m3fn,
+        device="cuda",
+    )
+    first = guard_pages * stride_scale
+    last = (guard_pages + num_pages) * stride_scale
+    view = backing[first:last:stride_scale]
+    assert view.shape[0] == num_pages
+    return view, backing[:first], backing[last:]
+
+
+def _assert_all_zero(*regions: torch.Tensor) -> None:
+    for region in regions:
+        assert torch.count_nonzero(region.reshape(-1).view(torch.uint8)).item() == 0
+
+
 def _run(
     qk: torch.Tensor,
     cache: torch.Tensor,
@@ -163,6 +194,55 @@ def test_minimax_m3_fp8_indexer_defensively_skips_invalid_direct_op_slots() -> N
     assert torch.count_nonzero(backing[2].view(torch.uint8)).item() == 0
 
 
+@pytest.mark.parametrize("num_live_tokens", [0, 4])
+@pytest.mark.parametrize("tail_slot", [-1, 4 * 128, 4 * 128 + 61, 5 * 128 + 127])
+def test_minimax_m3_fp8_indexer_ignores_a_padded_slot_tail(
+    tail_slot: int, num_live_tokens: int
+) -> None:
+    """Rows past the live prefix must leave every cache byte untouched.
+
+    A direct caller can hand the kernel a padded token height whose trailing
+    slots are the -1 sentinel or a stale out-of-range id, which the two guards
+    have to drop. The failure modes differ: -1 truncates to page 0 at offset
+    -1, one token below the cache base, while an out-of-range page scatters
+    above the pool. The tail slots start at the first page past the end and
+    stay inside the guard region, so a regression trips an assert rather than
+    faulting.
+    """
+    torch.manual_seed(2468)
+    num_heads_q = 4
+    page_size = 128
+    num_pages = 4
+    padded_tokens = 17
+
+    qk = torch.randn(padded_tokens, (num_heads_q + 1) * 128, dtype=torch.bfloat16, device="cuda")
+    q_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda")
+    k_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda")
+    position_ids = torch.arange(padded_tokens, dtype=torch.int32, device="cuda") + 1024
+
+    cache, below, above = _guarded_cache(num_pages, page_size)
+    slots = torch.full((padded_tokens,), tail_slot, dtype=torch.int32, device="cuda")
+    # One page per live row, so a stray tail store cannot be mistaken for one.
+    pages = torch.arange(num_live_tokens, dtype=torch.int32, device="cuda")
+    within = (pages * 37) % page_size
+    slots[:num_live_tokens] = pages * page_size + within
+
+    q_out = _run(qk, cache, slots, q_weight, k_weight, position_ids, num_heads_q)
+    q_ref, k_ref = _reference(qk, num_heads_q, q_weight, k_weight, position_ids)
+
+    # Only the cache store is slot-gated; index-Q is produced for every row.
+    _assert_fp8_close(q_out, q_ref)
+    _assert_all_zero(below, above)
+    if num_live_tokens:
+        _assert_fp8_close(cache[pages.long(), 0, within.long()], k_ref[:num_live_tokens])
+
+    # A tail store landing on a valid page at the wrong offset would clear both
+    # guard regions, so require every unwritten slot to stay zero as well.
+    written = torch.zeros(num_pages, page_size, dtype=torch.bool, device="cuda")
+    written[pages.long(), within.long()] = True
+    _assert_all_zero(cache[:, 0][~written])
+
+
 def test_minimax_m3_fp8_indexer_cuda_graph_replay_updates_outputs() -> None:
     torch.manual_seed(5678)
     num_tokens = 16

From 9f4afa8e2ff05d00e2330b285e634ecaa49c2f89 Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Sat, 12 Sep 2026 14:54:33 -0700
Subject: [PATCH 2/6] [None][fix] Limit MiniMax-M3 decode kernels to the last
 real token

A piecewise CUDA graph pads token-shaped inputs up to its capture bucket
without adding requests to go with them, so q can outrun the rows the batch
has. The decode kernels derive a request as token // decode_query_len, so a
padded q reads page table rows and sequence lengths past the batch's last one.

minimax_m3_attn_custom_op_inplace already clipped to the live token count, but
two paths were left uncovered.

_forward_attention_core's non-compiled branch dispatched the full-height q. It
is padded all the same, since the bucket is agreed across ranks. Pull the clip
into _dispatch_attention_over_live_tokens so both branches share it rather than
each kernel repeating it.

The clip also left output's pad rows as new_empty supplied them. No kernel
writes them and they flow on into the rest of the network, where nobody can
tell leftover bytes from a real NaN, so zero them.

Promote the Triton sparse decode's shape check to check_decode_span_shape and
apply it to trtllm-gen dense decode too, which had none. Its multi-CTA KV
counters are sized against max_num_requests, so a batch read out of a longer q
undersizes them. The two kernels take the same num_tokens, block table and
seq_lens slices from MsaDecodeFmha.run, so this enforces on the dense path an
invariant its sparse sibling already required. Naming the kernel is the point:
the alternative is an assert several frames inside one of them.

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 .../sparse/minimax_m3/kernels/msa_utils.py    | 16 +++++
 .../kernels/triton_sparse_decode.py           | 10 ++--
 .../kernels/trtllm_gen_dense_decode.py        | 11 ++++
 .../_torch/models/modeling_minimaxm3.py       | 51 ++++++++++++----
 .../attention/sparse/msa/test_msa_backend.py  | 12 ++++
 .../unittest/_torch/models/test_minimax_m3.py | 59 +++++++++++++++++++
 6 files changed, 144 insertions(+), 15 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py
index b3aa264d893e..cf65d95e5951 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py
@@ -23,6 +23,21 @@
 MSA_REQUIRED_HEAD_DIM = 128
 
 
+def check_decode_span_shape(kernel: str, total_q: int, batch: int, query_len: int) -> None:
+    """Reject a q that does not cover exactly the batch it was handed.
+
+    The decode kernels derive the request id as token // query_len, so a longer
+    q reads page table rows and lengths past the batch's last one. The caller
+    names itself so the error does too, rather than surfacing as an assert
+    several frames inside a kernel.
+    """
+    if total_q != batch * query_len:
+        raise ValueError(
+            f"{kernel}: total_q ({total_q}) must be batch ({batch}) * "
+            f"decode_query_len ({query_len})."
+        )
+
+
 def is_msa_layer(attn) -> bool:
     """Whether this layer's attention is served by the MiniMax-M3 MSA kernels."""
     sparse_params = attn.sparse_params
@@ -268,6 +283,7 @@ def select_blocks_from_maxscore(
     "MSA_REQUIRED_HEAD_DIM",
     "MSA_REQUIRED_TOPK",
     "build_kv_page_indices",
+    "check_decode_span_shape",
     "msa_package_available",
     "msa_paged_kv",
     "per_token_valid_blocks",
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 4e00b384b857..664c131f7b9c 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
@@ -36,6 +36,8 @@
 
 from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers
 
+from .msa_utils import check_decode_span_shape
+
 # One sparse block is exactly one KV page.
 SPARSE_BLOCK_SIZE = 128
 
@@ -315,11 +317,9 @@ def minimax_m3_sparse_attn_decode(
     """
     total_q, num_heads, head_dim = q.shape
     num_kv_heads = int(k_paged.shape[1])
-    if total_q != int(seq_lens.shape[0]) * decode_query_len:
-        raise ValueError(
-            f"total_q ({total_q}) must be batch ({int(seq_lens.shape[0])}) * "
-            f"decode_query_len ({decode_query_len})."
-        )
+    check_decode_span_shape(
+        "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len
+    )
     if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE:
         raise ValueError(
             f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; "
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py
index b477dc06a8d7..c94991b836f6 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py
@@ -27,6 +27,8 @@
 
 from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers
 
+from .msa_utils import check_decode_span_shape
+
 
 @functools.lru_cache(maxsize=None)
 def _counter_size(num_heads: int, max_num_requests: int, device_index: int) -> int:
@@ -252,6 +254,15 @@ def minimax_m3_trtllm_gen_dense_decode(
     """
     import flashinfer
 
+    # The multi-CTA KV counters are sized against max_num_requests, so a batch
+    # read out of a longer q would undersize them.
+    check_decode_span_shape(
+        "MiniMax-M3 trtllm-gen dense decode",
+        int(q.shape[0]),
+        int(seq_lens.shape[0]),
+        decode_query_len,
+    )
+
     kv_pool, subpages_per_slot = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND")
     num_heads = int(q.shape[1])
 
diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
index a7722bccfc1e..d8223e46410e 100644
--- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py
+++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
@@ -886,6 +886,41 @@ def _minimax_m3_fused_sparse_qkv_producer_fake(
     )
 
 
+def _dispatch_attention_over_live_tokens(
+    attn_layer: "MiniMaxM3Attention",
+    q: torch.Tensor,
+    k: Optional[torch.Tensor],
+    v: Optional[torch.Tensor],
+    idx_q: Optional[torch.Tensor],
+    idx_k: Optional[torch.Tensor],
+    attn_metadata: AttentionMetadata,
+    output: torch.Tensor,
+) -> None:
+    """Run the attention core over the step's live tokens alone.
+
+    A piecewise CUDA graph pads token-shaped inputs up to its capture bucket
+    without adding requests to go with them (see _get_padding_params in
+    model_engine), so q can outrun the rows the batch has. The kernels below
+    read a request out of a token index, so the pad comes off here, once for
+    both dispatch paths.
+
+    No kernel writes the pad rows of output, so they are zeroed rather than left
+    holding stale values that nothing can distinguish from a real NaN.
+    """
+    num_tokens = int(attn_metadata.num_tokens)
+    if num_tokens < int(output.shape[0]):
+        output[num_tokens:].zero_()
+    attn_layer._dispatch_attention_backend(
+        q[:num_tokens],
+        k[:num_tokens] if k is not None else None,
+        v[:num_tokens] if v is not None else None,
+        idx_q[:num_tokens] if idx_q is not None else None,
+        idx_k[:num_tokens] if idx_k is not None else None,
+        attn_metadata,
+        output[:num_tokens],
+    )
+
+
 @torch.library.custom_op("trtllm::minimax_m3_attn_custom_op_inplace", mutates_args=("output",))
 def minimax_m3_attn_custom_op_inplace(
     q: Optional[torch.Tensor],
@@ -923,15 +958,9 @@ def minimax_m3_attn_custom_op_inplace(
         k = v = idx_k = None
     if q is None:
         raise RuntimeError(f"MiniMax-M3 attention layer {layer_idx} received no query tensor.")
-    attn_layer._dispatch_attention_backend(
-        q[:num_tokens],
-        k[:num_tokens] if k is not None else None,
-        v[:num_tokens] if v is not None else None,
-        idx_q[:num_tokens] if idx_q is not None else None,
-        idx_k[:num_tokens] if idx_k is not None else None,
-        attn_metadata,
-        output[:num_tokens],
-    )
+    # The live token count is a host value, so the compiled graph above must
+    # not see it: it would guard on it and recapture per count.
+    _dispatch_attention_over_live_tokens(attn_layer, q, k, v, idx_q, idx_k, attn_metadata, output)
 
 
 maybe_bcg_minimax_m3_attn_custom_op_inplace = eager_on_graph(minimax_m3_attn_custom_op_inplace)
@@ -1843,7 +1872,9 @@ def _forward_attention_core(
                 output,
             )
         else:
-            self._dispatch_attention_backend(q, k, v, idx_q, idx_k, attn_metadata, output)
+            # A step that runs here rather than through compile is padded all
+            # the same, since the bucket is agreed across ranks.
+            _dispatch_attention_over_live_tokens(self, q, k, v, idx_q, idx_k, attn_metadata, output)
         return output
 
     def _dispatch_attention_backend(
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 4fd1c5535b13..fa85bdbbf5a7 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
@@ -27,6 +27,7 @@
 )
 from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import (
     MSA_REQUIRED_TOPK,
+    check_decode_span_shape,
     msa_paged_kv,
 )
 from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import (
@@ -1186,6 +1187,17 @@ def test_the_decode_span_of_a_mixed_step_is_its_generation_suffix():
     assert metadata.msa_max_kv_len == 40
 
 
+def test_decode_span_shape_check_names_the_kernel_that_rejected_the_q():
+    """The guard both decode kernels share."""
+    # Eleven speculative decode requests of 4 query tokens each.
+    check_decode_span_shape("kernel", 44, 11, 4)
+
+    # A piecewise CUDA graph's pad folded into the batch: the kernels would read
+    # 128 page table rows out of a batch that only has 11.
+    with pytest.raises(ValueError, match=r"kernel: total_q \(512\) must be batch \(11\)"):
+        check_decode_span_shape("kernel", 512, 11, 4)
+
+
 def test_a_pure_prefill_step_has_no_decode_span():
     """A step with no generation row has nothing for the decode kernels, and
     fmha_sm100 keeps every plan and the page table they read."""
diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py
index 32705c5708e8..8b5821cb3b2b 100644
--- a/tests/unittest/_torch/models/test_minimax_m3.py
+++ b/tests/unittest/_torch/models/test_minimax_m3.py
@@ -66,6 +66,7 @@
     MiniMaxM3MoE,
     MiniMaxM3QKVIndexerLinear,
     _build_swiglu_oai_dense_mlp,
+    _dispatch_attention_over_live_tokens,
     _load_qkv_index_proj_weights,
     _minimax_m3_swiglu_oai,
     _moe_routed_output_is_global,
@@ -996,6 +997,64 @@ def _has_cuda() -> bool:
 # ---------------------------------------------------------------------------
 
 
+def test_attention_dispatch_clips_the_piecewise_token_pad():
+    """Only the live tokens reach the attention core, and the output's pad rows
+    come back zeroed rather than as the buffer supplied them."""
+    seen = {}
+
+    def capture(q, k, v, idx_q, idx_k, attn_metadata, output):
+        del attn_metadata
+        seen["rows"] = [None if t is None else int(t.shape[0]) for t in (q, k, v, idx_q, idx_k)]
+        seen["out_rows"] = int(output.shape[0])
+        output.fill_(7.0)
+
+    attn_layer = SimpleNamespace(_dispatch_attention_backend=capture)
+    # Eleven speculative decode requests of 4 query tokens, padded to 64.
+    padded, live, hidden = 64, 44, 8
+    q = torch.ones((padded, hidden))
+    output = torch.full((padded, hidden), float("nan"))
+
+    _dispatch_attention_over_live_tokens(
+        attn_layer,
+        q,
+        q,
+        q,
+        None,
+        None,
+        SimpleNamespace(num_tokens=live),
+        output,
+    )
+
+    assert seen["rows"] == [live, live, live, None, None]
+    assert seen["out_rows"] == live
+    assert torch.equal(output[:live], torch.full((live, hidden), 7.0))
+    assert torch.equal(output[live:], torch.zeros(padded - live, hidden))
+
+
+def test_attention_dispatch_leaves_an_unpadded_step_alone():
+    """No pad, so nothing to clip and nothing to zero."""
+    seen = {}
+
+    def capture(q, k, v, idx_q, idx_k, attn_metadata, output):
+        del q, k, v, idx_q, idx_k, attn_metadata
+        seen["out"] = output
+
+    output = torch.full((5, 8), float("nan"))
+    _dispatch_attention_over_live_tokens(
+        SimpleNamespace(_dispatch_attention_backend=capture),
+        torch.ones((5, 8)),
+        None,
+        None,
+        None,
+        None,
+        SimpleNamespace(num_tokens=5),
+        output,
+    )
+
+    assert seen["out"].shape == (5, 8)
+    assert output.isnan().all()
+
+
 def test_is_minimax_m3_vl_config_detects_vl():
     assert is_minimax_m3_vl_config(_make_vl_config()) is True
 

From 55d6e781094abad56d0b04c9d079bbdf0b7a92f4 Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Sat, 12 Sep 2026 15:17:49 -0700
Subject: [PATCH 3/6] [None][fix] Reset the MSA cache-slot tail a piecewise
 step pads into

The fused FP8 index producer runs on hidden_states, ahead of the clip the
attention dispatch applies, so under a piecewise CUDA graph it sees the capture
bucket's padded height rather than the step's live token count. It cannot be
trimmed there: it sits inside the captured region, and a host-side count would
make its shape dynamic and defeat the capture. So it passes
msa_out_cache_loc[:padded] and the tail reaches
minimax_m3_fp8_indexer_qk_norm_rope.

That kernel already declines a negative slot and bounds the store from above by
the cache's page count. What it has no way to recognize is a tail that was never
reset: msa_out_cache_loc is allocated empty and written only up to
total_new_tokens, so a step following a larger one finds that step's valid slot
ids sitting past its own live count. Those address real pages, pass both guards,
and are written - silently, into whichever request owns them. Fill the buffer
with -1 before staging the mapping, which is the sentinel the guards were
written against.

write_kv_slots now takes the live token count and writes only that prefix. The
sentinel tail is contiguous by construction, so this is a host-side slice and
costs no device sync, unlike masking on out_cache_loc >= 0. The argument is
required, because there is no default that is safe to assume and the padded
height is exactly what a caller would pass by accident; making every call site
state its live count is the point. Its failure mode is the quiet one: torch
wraps negative indices, so a -1 slot writes into the last page rather than
raising. write_msa_main_kv threads the count through from the phase slice it is
handed, and the two triton_backend helpers pass their own row count, which is
already sentinel-free - build_paged_kv_slot_mapping emits one real slot per new
token.

The tests cover the shrinking-step case the fill exists for, the eager writer
against a sentinel-padded input, and the refusal to write against an unprepared
step's slots.

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 .../sparse/minimax_m3/kernels/msa_utils.py    |  16 ++-
 .../sparse/minimax_m3/kernels/paged_cache.py  |  21 +++
 .../backends/sparse/minimax_m3/msa_backend.py |  41 ++++++
 .../sparse/minimax_m3/triton_backend.py       |   7 +-
 .../attention/sparse/msa/test_msa_backend.py  | 121 +++++++++++++++++-
 5 files changed, 199 insertions(+), 7 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py
index cf65d95e5951..db0312237d8d 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py
@@ -118,6 +118,8 @@ def write_msa_main_kv(
     out_cache_loc: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
+    *,
+    num_live_tokens: int,
 ) -> None:
     """Write new-token K and V into the paged main cache at out_cache_loc.
 
@@ -131,10 +133,18 @@ def write_msa_main_kv(
     head_dim = int(k_view.shape[3])
     num_tokens = int(k.shape[0])
     write_kv_slots(
-        k_view, out_cache_loc, k.reshape(num_tokens, num_kv_heads, head_dim), layout="HND"
+        k_view,
+        out_cache_loc,
+        k.reshape(num_tokens, num_kv_heads, head_dim),
+        layout="HND",
+        num_live_tokens=num_live_tokens,
     )
     write_kv_slots(
-        v_view, out_cache_loc, v.reshape(num_tokens, num_kv_heads, head_dim), layout="HND"
+        v_view,
+        out_cache_loc,
+        v.reshape(num_tokens, num_kv_heads, head_dim),
+        layout="HND",
+        num_live_tokens=num_live_tokens,
     )
 
 
@@ -178,6 +188,8 @@ def write_msa_phase_kv(
         metadata.msa_out_cache_loc[token_offset : token_offset + num_tokens],
         k,
         v,
+        # k and v are this phase's own token slice, so every row owns a slot.
+        num_live_tokens=num_tokens,
     )
 
 
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py
index 3a20070d56c3..44df6e9ee08c 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py
@@ -15,6 +15,7 @@ def write_kv_slots(
     values: torch.Tensor,
     *,
     layout: Literal["NHD", "HND"] = "NHD",
+    num_live_tokens: int,
 ) -> None:
     """Write per-token values into a K, V, or index-K cache at given slots.
 
@@ -28,7 +29,27 @@ def write_kv_slots(
     mapping satisfies this contract: ``get_block_ids_per_seq`` canonicalizes
     padded ``BAD_PAGE_INDEX`` entries before ``build_paged_kv_slot_mapping``
     selects only the allocated live-token positions.
+
+    `num_live_tokens` is how many leading rows own a real cache slot; the rest
+    are dropped. It is required rather than defaulted because a caller passing
+    the padded token extent of a piecewise CUDA graph corrupts the cache
+    silently: torch wraps a negative index, so the -1 sentinel past the live
+    count lands in the last page instead of raising.
     """
+    # Trimming by count keeps this sync-free, the sentinel tail being contiguous
+    # by construction, where masking on the slot values would not.
+    if num_live_tokens < 0:
+        raise ValueError(f"num_live_tokens must be non-negative, got {num_live_tokens}")
+    if out_cache_loc.shape[0] < num_live_tokens or values.shape[0] < num_live_tokens:
+        raise ValueError(
+            f"num_live_tokens={num_live_tokens} exceeds the rows supplied "
+            f"(out_cache_loc={out_cache_loc.shape[0]}, values={values.shape[0]})"
+        )
+    if num_live_tokens == 0:
+        return
+    out_cache_loc = out_cache_loc[:num_live_tokens]
+    values = values[:num_live_tokens]
+
     with torch.no_grad():
         if cache.ndim >= 4:
             token_axis = 2 if layout == "HND" else 1
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
index 5abbedc884aa..87baca8f4f04 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
@@ -182,6 +182,8 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata):
     # _msa_fields_ready marks that the current step's buffers are populated.
     _msa_buffers_ready: bool = False
     _msa_fields_ready: bool = False
+    # Read through msa_live_token_count, which states the contract.
+    _msa_live_total_q: int = 0
     # Sparse geometry the plans need.
     _msa_params: Optional[MiniMaxM3SparseMetadataParams] = None
     # This step's fmha_sm100 plans, plain tuples with no graph-stable buffers
@@ -974,6 +976,10 @@ def _build_msa_fields(self) -> None:
         decoding the inputs for on_update_kv_lens are staged as well.
         """
         self._msa_fields_ready = False
+        # The live count describes the slot mapping staged at the end of this
+        # method, so clear it with the ready flag. An early return below would
+        # otherwise leave a count describing a batch that is no longer scheduled.
+        self._msa_live_total_q = 0
         if not self._msa_buffers_ready:
             return
         request_ids = self.request_ids
@@ -1034,6 +1040,13 @@ def _build_msa_fields(self) -> None:
                 f"smaller than the step's per-request page count ({block_table_cols})."
             )
 
+        # Piecewise graphs pad token-shaped inputs to the capture bucket, and the
+        # fused index producer runs over that padded extent: it sits inside the
+        # captured region, so trimming it to a host-side count would make its
+        # shape dynamic. A negative slot is what makes those rows cache-write
+        # no-ops. Fill before the copy, because a prior larger step leaves valid
+        # slot ids in the tail and those address real pages.
+        self.msa_out_cache_loc.fill_(-1)
         self.msa_out_cache_loc[:total_new_tokens].copy_(out_cache_loc, non_blocking=True)
         # Captured producers also execute padded rows. Invalidate only the
         # unwritten tail so they cannot reuse the previous step's live slots.
@@ -1060,6 +1073,12 @@ def _build_msa_fields(self) -> None:
                 self.msa_subpage_block_table[:batch_size],
             )
 
+        # Live, unpadded new-token count for this step, staged before either
+        # exit so the non-dynamic early return below carries it too. Cache
+        # writers read it through msa_live_token_count to find where
+        # msa_out_cache_loc stops holding real slots.
+        self._msa_live_total_q = total_new_tokens
+
         self._msa_kv_lens_dynamic = self._msa_kv_lens_may_change()
         if not self._msa_kv_lens_dynamic:
             self._msa_fields_ready = True
@@ -1087,6 +1106,21 @@ def _build_msa_fields(self) -> None:
         self.msa_kv_lens_staged[:batch_size].copy_(kv_lens_cpu, non_blocking=True)
         self._msa_fields_ready = True
 
+    def msa_live_token_count(self) -> int:
+        """This step's live, unpadded new-token count.
+
+        Cache writers take the padded token extent and need this to find where
+        msa_out_cache_loc stops holding real slots. Raising when no mapping is
+        staged keeps an unprepared step from writing against another step's
+        slots, which a stale count would otherwise allow.
+        """
+        if not self._msa_fields_ready:
+            raise RuntimeError(
+                "MiniMax-M3 MSA cache write requires prepared metadata, but "
+                "prepare() did not stage a slot mapping for this step."
+            )
+        return self._msa_live_total_q
+
     def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor:
         """Return the paged index-K cache in the HND layout MSA consumes."""
         return self.kv_cache_manager.get_index_k_buffer(layer_idx)
@@ -1101,6 +1135,7 @@ def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None:
             self.msa_out_cache_loc[:num_tokens],
             idx_k.reshape(num_tokens, 1, sparse_index_dim),
             layout="HND",
+            num_live_tokens=self.msa_live_token_count(),
         )
 
     def msa_proxy_max_score_view(
@@ -1257,17 +1292,22 @@ def write_layer_caches(
             return
         num_kv_heads = int(k_view.shape[1])
         head_dim = int(k_view.shape[3])
+        # The dispatch clips k/v/idx_k to the step's live tokens before this
+        # runs, so every supplied row owns a real slot: num_tokens is the live
+        # count write_kv_slots requires.
         write_kv_slots(
             k_view,
             out_cache_loc,
             k.reshape(num_tokens, num_kv_heads, head_dim),
             layout="HND",
+            num_live_tokens=num_tokens,
         )
         write_kv_slots(
             v_view,
             out_cache_loc,
             v.reshape(num_tokens, num_kv_heads, head_dim),
             layout="HND",
+            num_live_tokens=num_tokens,
         )
         if idx_k is not None:
             write_kv_slots(
@@ -1275,6 +1315,7 @@ def write_layer_caches(
                 out_cache_loc,
                 idx_k.reshape(num_tokens, 1, int(idx_cache.shape[-1])),
                 layout="HND",
+                num_live_tokens=num_tokens,
             )
 
     def run_indexer(
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py
index 844fd979469c..a6ce0e0b4608 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py
@@ -158,8 +158,11 @@ def _write_main_kv_slots_to_pool(
     ``out_cache_loc`` is the 1-D ``[num_new_tokens]`` int tensor of flat slot
     ids to update. ``pool[:, kv_index]`` is a storage-sharing view, so the
     shared :func:`common.write_kv_slots` propagates the write to the pool.
+
+    Every row owns a slot, so the row count is the live count: the slot mapping
+    emits one real slot per new token and no sentinel.
     """
-    write_kv_slots(pool[:, kv_index], out_cache_loc, values)
+    write_kv_slots(pool[:, kv_index], out_cache_loc, values, num_live_tokens=int(values.shape[0]))
 
 
 def _write_main_kv_slots(
@@ -173,7 +176,7 @@ def _write_main_kv_slots(
     flat-slot layout used by focused unit tests and the 4-D paged view of
     ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``.
     """
-    write_kv_slots(cache, out_cache_loc, values)
+    write_kv_slots(cache, out_cache_loc, values, num_live_tokens=int(values.shape[0]))
 
 
 def _scatter_topk_to_block_mask(
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 fa85bdbbf5a7..1ddf333e576b 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
@@ -646,6 +646,8 @@ def get_index_k_buffer(self, layer_idx):
     metadata.kv_cache_manager = manager
     metadata.msa_out_cache_loc = torch.tensor([2, page_size + 5], dtype=torch.int32)
     values = torch.arange(2 * head_dim, dtype=torch.float32).reshape(2, 1, head_dim)
+    metadata._msa_fields_ready = True
+    metadata._msa_live_total_q = 2
 
     returned = metadata.msa_idx_k_cache(3)
     metadata.msa_write_idx_k(3, values)
@@ -771,6 +773,7 @@ def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None:
                 self.msa_out_cache_loc,
                 idx_k,
                 layout="HND",
+                num_live_tokens=int(idx_k.shape[0]),
             )
 
         def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor:
@@ -1198,6 +1201,106 @@ def test_decode_span_shape_check_names_the_kernel_that_rejected_the_q():
         check_decode_span_shape("kernel", 512, 11, 4)
 
 
+def test_a_shrinking_step_leaves_no_live_slot_in_the_padded_tail():
+    """The slot-guarded consumers recognize a negative slot and nothing else, so
+    every row a step does not own has to hold one."""
+    block_ids = torch.arange(3 * 4, dtype=torch.int32).reshape(3, 4)
+
+    class FakeCacheManager:
+        tokens_per_block = 4
+
+        def get_block_ids_per_seq(self, request_ids):
+            return block_ids[: len(request_ids)]
+
+        def get_buffers(self, layer_idx, kv_layout="NHD"):
+            # Only its device is read, which keeps this test off the GPU.
+            return torch.zeros(1, dtype=torch.bfloat16)
+
+    metadata_cls = MiniMaxM3MsaSparseAttention.Metadata
+    metadata = metadata_cls.__new__(metadata_cls)
+    metadata.kv_cache_manager = FakeCacheManager()
+    metadata.mapping = None
+    metadata._msa_buffers_ready = True
+    metadata._msa_params = None
+    metadata._msa_decode_span = None
+    metadata.max_num_sequences = 3
+    metadata.msa_subpage_block_table = None
+    metadata._msa_subpages_per_slot = 0
+    metadata.msa_out_cache_loc = torch.zeros(16, dtype=torch.int32)
+    metadata.msa_kv_indices = torch.zeros(3 * 4, dtype=torch.int32)
+    metadata.msa_block_table = torch.zeros(3, 4, dtype=torch.int32)
+    metadata.msa_seq_lens_cuda = torch.zeros(3, dtype=torch.int32)
+
+    def build(request_ids, qo_lens, kv_lens):
+        # The host lengths are read-only properties over these.
+        metadata.request_ids = request_ids
+        metadata._msa_qo_lens_cpu = torch.tensor(qo_lens, dtype=torch.int32)
+        metadata._msa_kv_lens_cpu = torch.tensor(kv_lens, dtype=torch.int32)
+        metadata._msa_qo_offset_cpu = metadata._msa_kv_lens_cpu - metadata._msa_qo_lens_cpu
+        metadata._build_msa_fields()
+
+    # A three-request step of six new tokens, then a smaller one of two.
+    build([0, 1, 2], (4, 1, 1), (9, 3, 5))
+    assert metadata.msa_live_token_count() == 6
+    wide_tail = metadata.msa_out_cache_loc[2:6].tolist()
+    assert all(slot >= 0 for slot in wide_tail)
+
+    build([0, 1], (1, 1), (3, 5))
+
+    assert metadata.msa_live_token_count() == 2
+    # The rows the shrunk step does not own, including those the wider one did.
+    assert metadata.msa_out_cache_loc[2:].tolist() == [-1] * 14
+    assert (metadata.msa_out_cache_loc[:2] >= 0).all()
+
+
+def test_the_eager_writer_drops_the_sentinel_tail_it_is_handed():
+    """The padded rows of a step must reach no page at all, the wrap target of a
+    surviving -1 included."""
+    num_pages, page_size, head_dim = 4, 8, 16
+    cache = torch.zeros(num_pages, 1, page_size, head_dim, dtype=torch.bfloat16)
+    # Two live tokens, then the -1 tail a capture bucket pads the step out to.
+    out_cache_loc = torch.tensor([2, page_size + 5, -1, -1], dtype=torch.int32)
+    values = torch.arange(4 * head_dim, dtype=torch.float32).reshape(4, 1, head_dim)
+
+    write_kv_slots(cache, out_cache_loc, values, layout="HND", num_live_tokens=2)
+
+    torch.testing.assert_close(cache[0, 0, 2], values[0, 0].to(torch.bfloat16))
+    torch.testing.assert_close(cache[1, 0, 5], values[1, 0].to(torch.bfloat16))
+    # The page a wrapped -1 would have hit.
+    assert not cache[num_pages - 1].any()
+
+
+def test_the_eager_writer_refuses_a_live_count_it_has_no_rows_for():
+    """A count past the rows supplied is a caller bug, not a short write."""
+    cache = torch.zeros(4, 1, 8, 16, dtype=torch.bfloat16)
+    out_cache_loc = torch.tensor([2, 5], dtype=torch.int32)
+    values = torch.zeros(2, 1, 16)
+
+    with pytest.raises(ValueError, match=r"num_live_tokens=3 exceeds the rows supplied"):
+        write_kv_slots(cache, out_cache_loc, values, layout="HND", num_live_tokens=3)
+
+    with pytest.raises(ValueError, match="must be non-negative"):
+        write_kv_slots(cache, out_cache_loc, values, layout="HND", num_live_tokens=-1)
+
+    # A step that scheduled nothing writes nothing rather than erroring.
+    write_kv_slots(cache, out_cache_loc, values, layout="HND", num_live_tokens=0)
+    assert not cache.any()
+
+
+def test_an_unprepared_step_has_no_live_count_to_write_against():
+    """A stale count would let a write land on another step's slots."""
+    metadata_cls = MiniMaxM3MsaSparseAttention.Metadata
+    metadata = metadata_cls.__new__(metadata_cls)
+    metadata._msa_fields_ready = False
+    metadata._msa_live_total_q = 7
+
+    with pytest.raises(RuntimeError, match="did not stage a slot mapping"):
+        metadata.msa_live_token_count()
+
+    metadata._msa_fields_ready = True
+    assert metadata.msa_live_token_count() == 7
+
+
 def test_a_pure_prefill_step_has_no_decode_span():
     """A step with no generation row has nothing for the decode kernels, and
     fmha_sm100 keeps every plan and the page table they read."""
@@ -1820,14 +1923,26 @@ def test_fused_scatter_matches_reference(src_dtype, cache_dtype, with_idx):
     ref_pool = pool.clone()
     ref_idx_pool = idx_pool.clone()
     write_kv_slots(
-        ref_pool[:, 0], slots, k.reshape(num_tokens, num_kv_heads, head_dim), layout="HND"
+        ref_pool[:, 0],
+        slots,
+        k.reshape(num_tokens, num_kv_heads, head_dim),
+        layout="HND",
+        num_live_tokens=num_tokens,
     )
     write_kv_slots(
-        ref_pool[:, 1], slots, v.reshape(num_tokens, num_kv_heads, head_dim), layout="HND"
+        ref_pool[:, 1],
+        slots,
+        v.reshape(num_tokens, num_kv_heads, head_dim),
+        layout="HND",
+        num_live_tokens=num_tokens,
     )
     if with_idx:
         write_kv_slots(
-            ref_idx_pool[:, 0], slots, idx_k.reshape(num_tokens, 1, head_dim), layout="HND"
+            ref_idx_pool[:, 0],
+            slots,
+            idx_k.reshape(num_tokens, 1, head_dim),
+            layout="HND",
+            num_live_tokens=num_tokens,
         )
 
     assert fused_write_layer_caches(

From 0f63888fe9cdf650c3c0a5f5db9e0df448d96333 Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Sat, 12 Sep 2026 15:44:52 -0700
Subject: [PATCH 4/6] [None][test] Cover the wiring of the MiniMax-M3 padding
 guards

Both guards against a padded step were covered only in isolation. The span
check was exercised by calling it with integers, and the live-token clip by
calling it with a stand-in backend, so neither test says anything about whether
the guard is reached. Deleting the call from either decode kernel, or
dispatching to the backend without going through the clip, left the suite
green, which is the shape of regression that put the guards there.

The decode kernels are now driven through their own entry points with a q that
does not cover the batch. Both read their shapes before touching a device, so
this costs no GPU, and the dense kernel is handed no cache manager to show it
declines before consulting one.

Which callers reach the attention backend is itself the invariant, so the clip
is covered by reading it off the module: the backend dispatch must be reached
from the clipping helper alone. A second path that skipped the clip would name
itself in the failure.

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 .../attention/sparse/msa/test_msa_backend.py  | 57 +++++++++++++++++++
 .../unittest/_torch/models/test_minimax_m3.py | 34 +++++++++++
 2 files changed, 91 insertions(+)

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 1ddf333e576b..1b9219342ec2 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
@@ -1201,6 +1201,63 @@ def test_decode_span_shape_check_names_the_kernel_that_rejected_the_q():
         check_decode_span_shape("kernel", 512, 11, 4)
 
 
+def test_both_decode_kernels_reject_a_q_that_outruns_the_batch():
+    """The guard has to be reached, not merely available.
+
+    Both kernels read their shapes before touching a device, so the refusal
+    happens at the call and needs no GPU. The dense kernel is handed no cache
+    manager for the same reason: it must decline before consulting one.
+    """
+    batch, query_len, total_q = 11, 4, 512
+    num_heads, head_dim, page_size = 4, 128, 128
+    q = torch.empty(total_q, num_heads, head_dim)
+    output = torch.empty(total_q, num_heads, head_dim)
+    block_table = torch.zeros(batch, 4, dtype=torch.int32)
+    seq_lens = torch.zeros(batch, dtype=torch.int32)
+
+    pytest.importorskip("triton")
+    from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import (
+        minimax_m3_sparse_attn_decode,
+    )
+
+    paged = torch.empty(1, 1, page_size, head_dim)
+    with pytest.raises(
+        ValueError, match=r"Triton sparse decode: total_q \(512\) must be batch \(11\)"
+    ):
+        minimax_m3_sparse_attn_decode(
+            q,
+            paged,
+            paged,
+            torch.zeros(1, total_q, 64, dtype=torch.int64),
+            block_table,
+            seq_lens,
+            sm_scale=head_dim**-0.5,
+            output=output,
+            decode_query_len=query_len,
+        )
+
+    pytest.importorskip("flashinfer")
+    from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import (
+        minimax_m3_trtllm_gen_dense_decode,
+    )
+
+    with pytest.raises(
+        ValueError, match=r"trtllm-gen dense decode: total_q \(512\) must be batch \(11\)"
+    ):
+        minimax_m3_trtllm_gen_dense_decode(
+            q,
+            None,
+            0,
+            block_table,
+            seq_lens,
+            sm_scale=head_dim**-0.5,
+            output=output,
+            decode_query_len=query_len,
+            max_seq_len=1024,
+            max_num_requests=batch,
+        )
+
+
 def test_a_shrinking_step_leaves_no_live_slot_in_the_padded_tail():
     """The slot-guarded consumers recognize a negative slot and nothing else, so
     every row a step does not own has to hold one."""
diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py
index 8b5821cb3b2b..2571cd4a4540 100644
--- a/tests/unittest/_torch/models/test_minimax_m3.py
+++ b/tests/unittest/_torch/models/test_minimax_m3.py
@@ -1055,6 +1055,40 @@ def capture(q, k, v, idx_q, idx_k, attn_metadata, output):
     assert output.isnan().all()
 
 
+@pytest.mark.cpu_only
+def test_every_attention_dispatch_goes_through_the_live_token_clip():
+    """A dispatch path that bypassed the clip would hand the kernels the pad.
+
+    The clipping helper is the only thing standing between a padded step and
+    kernels that read a request out of a token index, so which callers reach
+    the backend is itself the invariant.
+    """
+    import ast
+    import inspect
+
+    from tensorrt_llm._torch.models import modeling_minimaxm3
+
+    tree = ast.parse(inspect.getsource(modeling_minimaxm3))
+    parents = {child: parent for parent in ast.walk(tree) for child in ast.iter_child_nodes(parent)}
+
+    def enclosing_function(node):
+        while node is not None:
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                return node.name
+            node = parents.get(node)
+        return None
+
+    callers = {
+        enclosing_function(node)
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and node.func.attr == "_dispatch_attention_backend"
+    }
+
+    assert callers == {"_dispatch_attention_over_live_tokens"}
+
+
 def test_is_minimax_m3_vl_config_detects_vl():
     assert is_minimax_m3_vl_config(_make_vl_config()) is True
 

From 4795e086df2709e248a53338e49f5f95436cd2e8 Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Thu, 24 Sep 2026 16:00:01 -0700
Subject: [PATCH 5/6] fix test

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 .../unittest/_torch/attention/sparse/msa/test_msa_backend.py  | 4 ++++
 1 file changed, 4 insertions(+)

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 1b9219342ec2..eea1190ce693 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
@@ -1287,6 +1287,10 @@ def get_buffers(self, layer_idx, kv_layout="NHD"):
     metadata.msa_kv_indices = torch.zeros(3 * 4, dtype=torch.int32)
     metadata.msa_block_table = torch.zeros(3, 4, dtype=torch.int32)
     metadata.msa_seq_lens_cuda = torch.zeros(3, dtype=torch.int32)
+    # Non-speculative, so no on_update_kv_lens staging: the slot tail this test
+    # is about is the same either way, and the staging buffers are not part of
+    # the fixture.
+    metadata._msa_kv_lens_may_change = lambda: False
 
     def build(request_ids, qo_lens, kv_lens):
         # The host lengths are read-only properties over these.

From c522dff139d074014b0908314c2cc6e5668004e4 Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Fri, 25 Sep 2026 05:45:00 +0000
Subject: [PATCH 6/6] fix another test

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 tests/unittest/_torch/models/test_minimax_m3.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py
index 2571cd4a4540..59b250952953 100644
--- a/tests/unittest/_torch/models/test_minimax_m3.py
+++ b/tests/unittest/_torch/models/test_minimax_m3.py
@@ -458,7 +458,7 @@ def _dispatch_attention_backend(self, q, k, v, idx_q, idx_k, attn_metadata, outp
 
     assert layer.producer_shapes == ((2, 5), (1, 2), 2)
     torch.testing.assert_close(output[:2], packed[:2, :3])
-    torch.testing.assert_close(output[2:], torch.full((2, 3), -1.0))
+    torch.testing.assert_close(output[2:], torch.zeros((2, 3)))
 
 
 @pytest.mark.cpu_only