[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,7 @@ 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 +132,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),
+        num_live_tokens,
+        layout="HND",
     )
     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),
+        num_live_tokens,
+        layout="HND",
     )
 
 
@@ -163,6 +187,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_tokens,
     )
 
 
@@ -268,6 +294,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

@@ -13,6 +13,7 @@ def write_kv_slots(
     cache: torch.Tensor,
     out_cache_loc: torch.Tensor,
     values: torch.Tensor,
+    num_live_tokens: int,
     *,
     layout: Literal["NHD", "HND"] = "NHD",
 ) -> None:
@@ -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 carries no default 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

@@ -45,6 +45,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
 
@@ -1048,11 +1050,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:
@@ -258,6 +260,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

@@ -1044,6 +1044,15 @@ def _build_msa_fields(self) -> None:
                 f"MSA out_cache_loc buffer ({self.msa_out_cache_loc.shape[0]}) is "
                 f"smaller than the step's new-token count ({total_new_tokens})."
             )
+        # The cache writers trim to num_tokens, so that count and the mapping
+        # have to describe the same rows. The mapping emits one slot per new
+        # token, so they agree unless a caller staged lengths this metadata's
+        # seq_lens does not match.
+        if total_new_tokens != int(self.num_tokens):
+            raise ValueError(
+                f"MSA slot mapping covers {total_new_tokens} new tokens, but the "
+                f"step's token count is {int(self.num_tokens)}."
+            )
         if kv_indices is not None and int(kv_indices.shape[0]) > self.msa_kv_indices.shape[0]:
             raise ValueError(
                 f"MSA kv_indices buffer ({self.msa_kv_indices.shape[0]}) is "
@@ -1058,8 +1067,11 @@ def _build_msa_fields(self) -> None:
             )
 
         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.
+        # Captured producers also execute padded rows: the fused index producer
+        # sits inside the captured region, so trimming it to a host-side count
+        # would make its shape dynamic, and a negative slot is what makes those
+        # rows cache-write no-ops instead. Invalidate the unwritten tail so they
+        # cannot reuse the previous step's live slots, which address real pages.
         if total_new_tokens < self.msa_out_cache_loc.shape[0]:
             self.msa_out_cache_loc[total_new_tokens:].fill_(-1)
         if kv_indices is not None:
@@ -1168,6 +1180,9 @@ def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None:
             cache,
             self.msa_out_cache_loc[:num_tokens],
             idx_k.reshape(num_tokens, 1, sparse_index_dim),
+            # idx_k arrives over the padded token extent; the live prefix is
+            # where msa_out_cache_loc stops holding real slots.
+            int(self.num_tokens),
             layout="HND",
         )
 
@@ -1349,23 +1364,29 @@ 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),
+            num_tokens,
             layout="HND",
         )
         write_kv_slots(
             v_view,
             out_cache_loc,
             v.reshape(num_tokens, num_kv_heads, head_dim),
+            num_tokens,
             layout="HND",
         )
         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])),
+                num_tokens,
                 layout="HND",
             )
 

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, 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, int(values.shape[0]))
 
 
 def _scatter_topk_to_block_mask(

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -952,6 +952,40 @@ 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.
+
+    The pad rows of output are left as the buffer supplied them. Nothing reads
+    them for its own result, and clearing them would put a device launch behind
+    a host-side count on every layer of every step.
+    """
+    num_tokens = int(attn_metadata.num_tokens)
+    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],
@@ -989,15 +1023,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)
@@ -1959,7 +1987,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(