[None][fix] Guard MiniMax-M3 FP8 indexer against padded -1 cache slots

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
@@ -103,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.
 
@@ -116,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,
     )
 
 
@@ -163,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,
     )
 
 
@@ -268,6 +295,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",

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

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}; "

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])
 

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,24 +1292,30 @@ 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(
                 idx_cache,
                 out_cache_loc,
                 idx_k.reshape(num_tokens, 1, int(idx_cache.shape[-1])),
                 layout="HND",
+                num_live_tokens=num_tokens,
             )
 
     def run_indexer(

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(

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(