[None][perf] Custom decode kernels for MinimaxM3 - rebased

tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py

@@ -124,9 +124,20 @@ def run_msa_paged_gqa(
     with the sparse plan) and the dense layers (kv_block_indexes None, with the
     dense plan, attending the full page table). fmha_sm100 reads the paged cache
     directly, so the new-token K/V must be resident before the run.
+
+    This is also where a step splits by request phase. TensorRT-LLM orders a
+    batch context-first, so the generation requests are its token suffix: a
+    ported decode kernel takes that suffix and fmha_sm100 keeps the context
+    prefix, running under the plan prepare() built over those rows alone (see
+    _msa_fmha_plan_rows). The prefix is empty on a pure-decode step, so one code
+    path covers both. The split lives here rather than in a PhasedFmha subclass
+    because MiniMaxM3MsaSparseAttention.forward_prepopulated_kv also calls this
+    helper directly, bypassing TrtllmAttention.forward.
     """
     from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import (
+        msa_decode_span_bounds,
         msa_paged_kv,
+        msa_ported_decode_active,
         write_msa_main_kv,
     )
 
@@ -155,29 +166,140 @@ def run_msa_paged_gqa(
     k_paged, v_paged = msa_paged_kv(kv_cache_manager, layer_idx)
     sm_scale = (head_dim**-0.5) / float(attn.q_scaling)
 
+    # Query tokens and batch rows the ported kernels own; (num_tokens, batch)
+    # of them on a pure-decode step, the trailing generation slice on a mixed
+    # one. gen_tok0 is 0 whenever nothing is ported.
+    gen_tok0, gen_row0, gen_row1, decode_query_len = msa_decode_span_bounds(metadata, num_tokens)
+    # Leading query tokens fmha_sm100 must still run: the whole batch until a
+    # ported kernel takes the generation slice, then the context prefix alone.
+    fmha_tokens = num_tokens
+    ported = msa_ported_decode_active(metadata)
+
+    if kv_block_indexes is not None and ported:
+        from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.triton_sparse_decode import (
+            minimax_m3_sparse_attn_decode,
+        )
+
+        # The Triton kernel dequantizes the FP8 cache into q's dtype, so q stays
+        # wide here. An already-E4M3 q from the fused producer widens exactly,
+        # leaving the same values the fmha_sm100 path would have used.
+        gen_q = q_view[gen_tok0:]
+        if gen_q.dtype != out_view.dtype:
+            gen_q = gen_q.to(out_view.dtype)
+        minimax_m3_sparse_attn_decode(
+            gen_q,
+            k_paged,
+            v_paged,
+            # [total_q, num_kv_heads, topk] -> head-major, contiguous when the
+            # indexer emitted a head-major table and the slice is the whole
+            # batch (see msa_ported_decode_active, which both sites read). The
+            # kernel takes every stride, so a mixed step's strided suffix is fine.
+            kv_block_indexes[gen_tok0:].permute(1, 0, 2),
+            metadata.msa_block_table[gen_row0:gen_row1],
+            metadata.msa_seq_lens_cuda[gen_row0:gen_row1],
+            sm_scale=sm_scale,
+            output=out_view[gen_tok0:],
+            decode_query_len=decode_query_len,
+        )
+        fmha_tokens = gen_tok0
+
+    elif kv_block_indexes is None and ported:
+        from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.trtllm_gen_dense_decode import (
+            dense_decode_unsupported_reason,
+            minimax_m3_trtllm_gen_dense_decode,
+        )
+
+        unsupported = dense_decode_unsupported_reason(kv_cache_manager, head_dim)
+        if unsupported is not None:
+            raise RuntimeError(
+                "MiniMax-M3 resolved a generation span for this step's dense "
+                f"layers and skipped the fmha_sm100 dense plan, but {unsupported} "
+                "The two must agree, and there is no plan left to run the span; "
+                "see _resolve_decode_kernels."
+            )
+        # The sub-page block table prepare() staged, if it could; the kernel
+        # expands its own when the factor does not match this layer's.
+        staged_subpage_rows = getattr(metadata, "msa_subpage_rows", None)
+        staged_table, staged_factor = (
+            staged_subpage_rows(gen_row0, gen_row1)
+            if staged_subpage_rows is not None
+            else (None, 0)
+        )
+        minimax_m3_trtllm_gen_dense_decode(
+            q_view[gen_tok0:],
+            kv_cache_manager,
+            layer_idx,
+            metadata.msa_block_table[gen_row0:gen_row1],
+            metadata.msa_seq_lens_cuda[gen_row0:gen_row1],
+            sm_scale=sm_scale,
+            output=out_view[gen_tok0:],
+            decode_query_len=decode_query_len,
+            # Bounded by the span's own rows, so a long context request
+            # cannot inflate the kernel's scheduling hint.
+            max_seq_len=int(metadata.msa_max_kv_len),
+            max_num_requests=int(metadata.max_num_requests),
+            staged_subpage_table=staged_table,
+            staged_subpages_per_slot=staged_factor,
+        )
+        fmha_tokens = gen_tok0
+
+    if fmha_tokens == 0:
+        return
+
+    if fmha_tokens == num_tokens:
+        # Reaching fmha_sm100 for the whole batch when a span was resolved means
+        # neither branch above took it, after prepare had already skipped this
+        # layer's plan and, on a pure-decode step, the flattened page table the
+        # call below reads. Fail loudly: running on with a stale msa_kv_indices
+        # would silently attend the wrong pages.
+        if ported:
+            raise RuntimeError(
+                "MiniMax-M3 paged GQA reached fmha_sm100 with no plan for a "
+                f"{'sparse' if kv_block_indexes is not None else 'dense'} layer. "
+                "The step resolved a generation span, which the ported decode "
+                "kernels own; see _resolve_decode_kernels."
+            )
+        fmha_rows = None
+    else:
+        # The context prefix, matching the rows `plan` was built over. The
+        # flattened page table needs no slice: context pages are its prefix and
+        # the plan implies how many of them to read.
+        fmha_rows = gen_row0
+
     # The fmha_sm100 variant is chosen from q.dtype and shares one dtype across
     # q/k/v, so q must be FP8 to match an FP8 paged K/V. MiniMax-M3 has no
     # KV-cache scales, so the scale is 1.0 and this is a plain E4M3 cast. When the
     # model's fused QK-norm+RoPE already emitted FP8 q/k/v (the FP8-KV fast path),
     # this .to() is a no-op; it stays as a safety net for callers that pass bf16 q.
+    fmha_q = q_view[:fmha_tokens]
     use_fp8 = k_paged.dtype == torch.float8_e4m3fn
-    if use_fp8 and q_view.dtype != torch.float8_e4m3fn:
-        q_view = q_view.to(torch.float8_e4m3fn)
+    if use_fp8 and fmha_q.dtype != torch.float8_e4m3fn:
+        fmha_q = fmha_q.to(torch.float8_e4m3fn)
+
+    def rows_of(lens: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        """Narrow a per-request host length tensor to the rows fmha_sm100 runs.
+
+        Slicing a pinned tensor keeps the pinned backing, so the inline planner
+        still stages these with non-blocking copies.
+        """
+        if lens is None or fmha_rows is None:
+            return lens
+        return lens[:fmha_rows]
 
     run_msa_sparse_gqa(
-        q_view,
+        fmha_q,
         k_paged,
         v_paged,
-        kv_block_indexes,
+        None if kv_block_indexes is None else kv_block_indexes[:fmha_tokens],
         kv_indices=metadata.msa_kv_indices,
         sm_scale=sm_scale,
-        qo_lens_cpu=metadata.msa_qo_lens_cpu,
-        kv_lens_cpu=metadata.msa_kv_lens_cpu,
-        qo_offset_cpu=metadata.msa_qo_offset_cpu,
+        qo_lens_cpu=rows_of(metadata.msa_qo_lens_cpu),
+        kv_lens_cpu=rows_of(metadata.msa_kv_lens_cpu),
+        qo_offset_cpu=rows_of(metadata.msa_qo_offset_cpu),
         causal=True,
         head_dim=head_dim,
         plan=plan,
-        out=out_view,
+        out=out_view[:fmha_tokens],
         use_fp8=use_fp8,
     )
 
@@ -190,11 +312,13 @@ class MsaSparseGqaFmha(Fmha):
     and attend those blocks; dense layers leave the indices None and attend the
     full page table.
 
-        Inherits Fmha rather than PhasedFmha: fmha_sm100 takes a single plan and
-        the selected block indices span the whole batch, so it handles a mixed
-        context and generation batch in one call and there is no
-        context/generation split from PhasedFmha to reuse. Requires head_dim 128
-        and 4-D HND paged K/V.
+        Inherits Fmha rather than PhasedFmha even though a mixed batch is split
+        by phase, because that split has to happen in run_msa_paged_gqa rather
+        than in forward: MiniMaxM3MsaSparseAttention.forward_prepopulated_kv
+        calls that helper directly, so a split placed in PhasedFmha.forward
+        would miss it. PhasedFmha also cannot reach the third ported kernel, the
+        indexer scorer, which runs before forward from run_indexer. Requires
+        head_dim 128 and 4-D HND paged K/V.
     """
 
     @classmethod

tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py

@@ -25,7 +25,7 @@
 
 from __future__ import annotations
 
-from typing import List, Optional, Sequence
+from typing import List, Optional, Sequence, Tuple
 
 import torch
 
@@ -324,33 +324,25 @@ def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]:
     def has_index_value(self, layer_idx: int) -> bool:
         return layer_idx in self._index_v_buffers
 
-    def get_buffers(
-        self, layer_idx: int, kv_layout: Optional[str] = None
-    ) -> Optional[torch.Tensor]:
-        """Return a paged K+V view with strides spanning the coalesced pool.
-
-        The base :meth:`KVCacheManagerV2.get_buffers` produces a
-        ``[num_pages, kv_factor, ...]`` view with contiguous strides
-        that assume the slot holds exactly one layer's K+V. In M3's
-        pool the slot packs K+V for *all* layers of the group
-        (``scale >= 2 * num_layers_in_group``), so the base view's
-        dim-0 stride does not reach the next slot's K for this layer.
-        (When INDEX_KEY's per-block size coincides with K/V's, it is
-        coalesced into the same pool and contributes to ``scale`` too.)
-
-        The override builds a ``[num_slots, scale, ...]`` view rooted
-        at K's base, then slices ``[:, :2]`` to extract K+V. The slice
-        preserves the dim-0 stride (``scale * page_stride``), so
-        ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``.
-        When omitted, ``kv_layout`` follows the selected sparse backend.
+    def _kv_slot_geometry(
+        self, layer_idx: int, kv_layout: Optional[str]
+    ) -> Tuple[int, torch.dtype, int, int, List[int]]:
+        """Resolve one layer's position in the coalesced K/V pool.
+
+        Returns ``(addr_key, torch_dtype, num_slots, scale, page_shape)``,
+        where ``scale`` is the number of equal-sized sub-pages a slot packs
+        and ``page_shape`` is one sub-page's shape in ``kv_layout``. This
+        layer's K is sub-page 0 and its V sub-page 1, counting from
+        ``addr_key``. When ``kv_layout`` is None it follows the selected
+        sparse backend.
         """
         if kv_layout is None:
             kv_layout = self._main_kv_layout
         if kv_layout not in ("NHD", "HND"):
             raise ValueError(f"Unsupported kv_layout: {kv_layout}")
         if self.kv_cache_type == CacheTypeCpp.SELFKONLY:
             raise NotImplementedError(
-                "MiniMaxM3KVCacheManagerV2.get_buffers does not support SELFKONLY cache type"
+                "MiniMaxM3KVCacheManagerV2 does not support the SELFKONLY cache type"
             )
 
         layer_offset = self.layer_offsets[layer_idx]
@@ -361,13 +353,13 @@ def get_buffers(
         # V2 always lays V immediately after K within the per-layer
         # contribution to a slot. The slice ``[:, :2]`` depends on this.
         assert addr_key + page_stride_value == addr_value, (
-            f"MiniMaxM3 get_buffers requires addr_K + page_stride "
+            f"MiniMaxM3 requires addr_K + page_stride "
             f"== addr_V (V immediately after K in slot); got "
             f"addr_K={addr_key} page_stride_V={page_stride_value} "
             f"addr_V={addr_value} for layer {layer_idx}."
         )
         assert page_stride_key == page_stride_value, (
-            f"MiniMaxM3 get_buffers requires equal K and V page "
+            f"MiniMaxM3 requires equal K and V page "
             f"strides; got K={page_stride_key} V="
             f"{page_stride_value}."
         )
@@ -394,27 +386,67 @@ def get_buffers(
 
         layer_head_dim = self.head_dim_per_layer[layer_offset]
         num_kv_heads = self.num_kv_heads_per_layer[layer_offset]
+        containers = layer_head_dim // element_per_container
 
         if kv_layout == "NHD":
-            full_slot_shape = [
-                num_slots,
-                scale,
-                self.tokens_per_block,
-                num_kv_heads,
-                layer_head_dim // element_per_container,
-            ]
+            page_shape = [self.tokens_per_block, num_kv_heads, containers]
         else:
-            full_slot_shape = [
-                num_slots,
-                scale,
-                num_kv_heads,
-                self.tokens_per_block,
-                layer_head_dim // element_per_container,
-            ]
+            page_shape = [num_kv_heads, self.tokens_per_block, containers]
+        return addr_key, torch_dtype, num_slots, scale, page_shape
+
+    def get_buffers(
+        self, layer_idx: int, kv_layout: Optional[str] = None
+    ) -> Optional[torch.Tensor]:
+        """Return a paged K+V view with strides spanning the coalesced pool.
+
+        The base :meth:`KVCacheManagerV2.get_buffers` produces a
+        ``[num_pages, kv_factor, ...]`` view with contiguous strides
+        that assume the slot holds exactly one layer's K+V. In M3's
+        pool the slot packs K+V for *all* layers of the group
+        (``scale >= 2 * num_layers_in_group``), so the base view's
+        dim-0 stride does not reach the next slot's K for this layer.
+        (When INDEX_KEY's per-block size coincides with K/V's, it is
+        coalesced into the same pool and contributes to ``scale`` too.)
 
+        The override builds a ``[num_slots, scale, ...]`` view rooted
+        at K's base, then slices ``[:, :2]`` to extract K+V. The slice
+        preserves the dim-0 stride (``scale * page_stride``), so
+        ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``.
+        When omitted, ``kv_layout`` follows the selected sparse backend.
+        """
+        addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry(
+            layer_idx, kv_layout
+        )
+        full_slot_shape = [num_slots, scale, *page_shape]
         full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape))
         return full_view[:, :2]
 
+    def get_kv_subpage_pool(
+        self, layer_idx: int, kv_layout: str = "HND"
+    ) -> Tuple[torch.Tensor, int]:
+        """Return ``(flat_pool, subpages_per_slot)`` for flat-block consumers.
+
+        trtllm-gen addresses K and V pages independently, through a
+        ``[batch, 2, max_blocks]`` block table into one flat
+        ``[num_subpages, *page_shape]`` pool. That is expressible here even
+        though the per-layer stride is not uniform: a slot packs ``scale``
+        equal-sized sub-pages, of which this layer owns two adjacent ones, so
+        rooting the flat pool at this layer's K puts slot ``s``'s K at
+        ``s * scale`` and its V at ``s * scale + 1``.
+
+        The view stops two sub-pages past the last slot's K rather than
+        spanning ``num_slots * scale``, which would run off the pool by
+        whatever this layer's K offset is inside a slot.
+        """
+        addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry(
+            layer_idx, kv_layout
+        )
+        num_subpages = (num_slots - 1) * scale + 2
+        flat = convert_to_torch_tensor(
+            TensorWrapper(addr_key, torch_dtype, [num_subpages, *page_shape])
+        )
+        return flat, scale
+
     def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> int:
         """Pool-mapping offset from the layer's physical position in its pool.