[None][perf] Fuse MiniMax-M3 MSA per-layer KV-cache writes into one kernel

tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py

@@ -0,0 +1,173 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Fused paged-cache scatter for the MiniMax-M3 MSA backend.
+
+One Triton launch writes a layer's new-token main K, main V, and (sparse
+layers) index-K into their paged HND caches at the step's write slots.
+The legacy path costs three aten advanced-indexing writes per layer plus
+their index preprocessing; at 60 layers per forward step, all captured
+into decode CUDA graphs, the launch count dominates the cost. The kernel
+derives each token's (page, within-page) split from ``out_cache_loc``
+in-register, so it needs no precomputed index tensors at all.
+
+Sources may be strided row views (slices of the fused QKV projection);
+only the innermost [num_heads * head_dim] extent must be contiguous.
+Stores cast to the cache dtype, which folds the FP8 KV-cache cast in.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _fused_paged_scatter_kernel(
+    k_src,
+    v_src,
+    idx_src,
+    k_cache,
+    v_cache,
+    idx_cache,
+    out_cache_loc,
+    k_src_row_stride,
+    v_src_row_stride,
+    idx_src_row_stride,
+    kc_stride_page,
+    kc_stride_head,
+    kc_stride_tok,
+    vc_stride_page,
+    vc_stride_head,
+    vc_stride_tok,
+    ic_stride_page,
+    ic_stride_tok,
+    tokens_per_block,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    HAS_IDX: tl.constexpr,
+):
+    # int64 throughout: t * row_stride can exceed 2^31 elements on large
+    # eager prefill steps (num_tokens up to max_num_tokens times the fused
+    # QKV row stride), and the slot * page-stride products likewise.
+    t = tl.program_id(0).to(tl.int64)
+    slot = tl.load(out_cache_loc + t).to(tl.int64)
+    page = slot // tokens_per_block
+    within = slot % tokens_per_block
+    d = tl.arange(0, D)
+    for h in tl.static_range(H):
+        k_vals = tl.load(k_src + t * k_src_row_stride + h * D + d)
+        v_vals = tl.load(v_src + t * v_src_row_stride + h * D + d)
+        k_dst = k_cache + page * kc_stride_page + h * kc_stride_head + within * kc_stride_tok + d
+        v_dst = v_cache + page * vc_stride_page + h * vc_stride_head + within * vc_stride_tok + d
+        tl.store(k_dst, k_vals.to(k_cache.dtype.element_ty))
+        tl.store(v_dst, v_vals.to(v_cache.dtype.element_ty))
+    if HAS_IDX:
+        i_vals = tl.load(idx_src + t * idx_src_row_stride + d)
+        i_dst = idx_cache + page * ic_stride_page + within * ic_stride_tok + d
+        tl.store(i_dst, i_vals.to(idx_cache.dtype.element_ty))
+
+
+def _row_stride_if_fusable(src: torch.Tensor, inner: int) -> Optional[int]:
+    """Row stride (elements) if `src` is a [T, inner] row view with contiguous
+    rows (e.g. a column slice of the fused QKV projection); None otherwise."""
+    if src.dim() != 2 or src.shape[1] != inner or src.stride(1) != 1:
+        return None
+    return src.stride(0)
+
+
+def fused_write_layer_caches(
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    idx_cache: Optional[torch.Tensor],
+    out_cache_loc: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    idx_k: Optional[torch.Tensor],
+) -> bool:
+    """Fused single-launch write of new-token K/V (+index-K) into paged HND
+    caches. Returns False when a layout or device precondition fails, so the caller can
+    keep the legacy per-cache writes.
+
+    `k_cache`/`v_cache` are [num_pages, num_kv_heads, tokens_per_block,
+    head_dim] HND views; `idx_cache` is the MQA index-K view with one head.
+    `k`/`v` are the layer's new-token values as [T, H*D] row views;
+    `idx_k` is [T, D]. Their inner dimension must be contiguous.
+    """
+    if not k_cache.is_cuda or any(
+        tensor.device != k_cache.device for tensor in (k, v, v_cache, out_cache_loc)
+    ):
+        return False
+    if k_cache.dim() != 4 or v_cache.dim() != 4:
+        return False
+    if v_cache.shape != k_cache.shape:
+        return False
+    if k_cache.stride(-1) != 1 or v_cache.stride(-1) != 1:
+        return False
+    num_pages, num_heads, tokens_per_block, head_dim = k_cache.shape
+    if (head_dim & (head_dim - 1)) != 0:
+        return False
+    inner = num_heads * head_dim
+    k_stride = _row_stride_if_fusable(k, inner)
+    v_stride = _row_stride_if_fusable(v, inner)
+    if k_stride is None or v_stride is None:
+        return False
+
+    has_idx = idx_k is not None
+    idx_stride = 0
+    ic_stride_page = 0
+    ic_stride_tok = 0
+    if has_idx:
+        if idx_cache is None or idx_cache.dim() != 4 or idx_cache.stride(-1) != 1:
+            return False
+        if idx_k.device != k_cache.device or idx_cache.device != k_cache.device:
+            return False
+        if int(idx_cache.shape[1]) != 1 or int(idx_cache.shape[3]) != head_dim:
+            return False
+        if int(idx_cache.shape[2]) != tokens_per_block:
+            return False
+        idx_stride = _row_stride_if_fusable(idx_k, head_dim)
+        if idx_stride is None:
+            return False
+        ic_stride_page = idx_cache.stride(0)
+        ic_stride_tok = idx_cache.stride(2)
+
+    num_tokens = int(out_cache_loc.shape[0])
+    if num_tokens == 0:
+        return True
+    if k.shape[0] < num_tokens or v.shape[0] < num_tokens:
+        return False
+    if has_idx and idx_k.shape[0] < num_tokens:
+        return False
+
+    _fused_paged_scatter_kernel[(num_tokens,)](
+        k,
+        v,
+        idx_k if has_idx else k,  # unused when HAS_IDX=False
+        k_cache,
+        v_cache,
+        idx_cache if has_idx else k_cache,  # unused when HAS_IDX=False
+        out_cache_loc,
+        k_stride,
+        v_stride,
+        idx_stride,
+        k_cache.stride(0),
+        k_cache.stride(1),
+        k_cache.stride(2),
+        v_cache.stride(0),
+        v_cache.stride(1),
+        v_cache.stride(2),
+        ic_stride_page,
+        ic_stride_tok,
+        tokens_per_block,
+        H=num_heads,
+        D=head_dim,
+        HAS_IDX=has_idx,
+        num_warps=2,
+    )
+    return True
+
+
+__all__ = ["fused_write_layer_caches"]

tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py

@@ -140,7 +140,11 @@ def write_msa_phase_kv(
     phases cover the step between them and neither repeats the other's write.
 
     k and v are the phase's token slice, and token_offset its first token on
-    the step's token axis, which is what msa_out_cache_loc is indexed by.
+    the step's token axis, which is what msa_out_cache_loc is indexed by. A
+    phase handed no K/V (k and v None) has nothing to write: that is how the
+    MiniMax-M3 model layer, which stores the whole step's K/V itself through
+    MiniMaxM3MsaSparseAttention.write_layer_caches ahead of its indexer,
+    tells both libraries the cache is already resident.
     """
     if attention_input_type != AttentionInputType.mixed:
         raise NotImplementedError(

tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py

@@ -1208,20 +1208,74 @@ def support_fused_rope(cls) -> bool:
         # index branches explicitly.
         return False
 
+    def write_layer_caches(
+        self,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        idx_k: Optional[torch.Tensor],
+        metadata,
+    ) -> None:
+        """Write this layer's new-token K, V and (bf16 indexer) index-K.
+
+        One fused kernel launch when the source/cache layouts allow it, else
+        the legacy per-cache writes. The model layer calls this first, so the
+        index-K cache is populated before run_indexer's proxy pass reads it,
+        and then hands forward() k=v=None: write_msa_phase_kv writes nothing
+        for a phase without live K/V, so neither FMHA library repeats the
+        write. `idx_k` is None on the FP8 indexer path, where the fused
+        producer has already inserted E4M3 index-K into the side cache.
+        `metadata` only supplies the step's write slots (msa_out_cache_loc,
+        filled by prepare()) and the cache manager.
+        """
+        from .kernels.msa_scatter import fused_write_layer_caches
+
+        layer_idx = self.layer_idx
+        buffers = metadata.kv_cache_manager.get_buffers(layer_idx, kv_layout="HND")
+        k_view, v_view = buffers[:, 0], buffers[:, 1]
+        idx_cache = metadata.msa_idx_k_cache(layer_idx) if idx_k is not None else None
+        num_tokens = int(k.shape[0])
+        out_cache_loc = metadata.msa_out_cache_loc[:num_tokens]
+        if fused_write_layer_caches(k_view, v_view, idx_cache, out_cache_loc, k, v, idx_k):
+            return
+        num_kv_heads = int(k_view.shape[1])
+        head_dim = int(k_view.shape[3])
+        write_kv_slots(
+            k_view,
+            out_cache_loc,
+            k.reshape(num_tokens, num_kv_heads, head_dim),
+            layout="HND",
+        )
+        write_kv_slots(
+            v_view,
+            out_cache_loc,
+            v.reshape(num_tokens, num_kv_heads, head_dim),
+            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])),
+                layout="HND",
+            )
+
     def run_indexer(
         self,
         idx_q: torch.Tensor,
         idx_k: Optional[torch.Tensor],
         metadata,
         *,
         idx_sm_scale: Optional[float] = None,
+        idx_k_prewritten: bool = False,
     ) -> torch.Tensor:
         """Write the index-K cache and return the selected block indices.
 
         The model layer runs this before forward and threads the result through
         forward_args.sparse_backend_args. Returns [total_q, num_kv_heads, topk].
         The generation rows are scored by the CuTe DSL kernel and any context
         rows by the fmha_sm100 proxy pass, over the plan prepare() built.
+        `idx_k_prewritten` marks that the fused per-layer cache write
+        (write_layer_caches) already stored this layer's index-K.
         """
         config = self.m3_config
         idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5
@@ -1254,13 +1308,18 @@ def run_indexer(
                     "The MiniMax-M3 BF16 indexer requires BF16 index-Q and a live "
                     f"BF16 index-K tensor; got Q={idx_q_view.dtype}, K={live_k_dtype}."
                 )
-            idx_k_view = idx_k.view(num_tokens, 1, config.sparse_index_dim)
-            metadata.msa_write_idx_k(self.layer_idx, idx_k_view)
+            # The fused per-layer write (write_layer_caches, signalled by
+            # idx_k_prewritten) may already have stored this live bf16 index-K
+            # ahead of the proxy pass; write it here only when it did not.
+            if not idx_k_prewritten:
+                idx_k_view = idx_k.view(num_tokens, 1, config.sparse_index_dim)
+                metadata.msa_write_idx_k(self.layer_idx, idx_k_view)
         # The FP8 indexer mirrors vLLM's unscaled E4M3 contract: normalized
         # index Q/K are cast directly and the proxy accumulates their QK scores
         # in FP32. Block ordering is invariant to the omitted positive scale.
         # The fused production path arrives here with E4M3 Q and an already
-        # populated cache; the BF16 path writes its live K above.
+        # populated cache; the BF16 path writes its live K above unless the
+        # fused per-layer write already did.
 
         # Inputs for the CuTe DSL scorer, which takes this step's generation
         # span. Left None on a pure-prefill step, which has no span, so the

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -1398,20 +1398,34 @@ def _msa_attention_core(
         The backend runs the sparse GQA or dense paged GQA through its inherited
         FMHA forward; this layer selects the top-k blocks (sparse only) and
         builds the forward_args the FMHA reads.
+
+        This layer owns the cache write: write_layer_caches stores the
+        new-token K/V (and, on the bf16 indexer path, index-K) in one launch
+        before the indexer's proxy pass reads the index-K cache. forward()
+        then receives k=v=None, which is the backend's contract for "K/V are
+        already resident", so neither FMHA phase writes them again.
         """
         if self.is_sparse_attention_layer:
             assert idx_q is not None
+            # On the FP8 indexer path idx_k is None: the fused producer already
+            # inserted E4M3 index-K into the side cache, so only K/V are written.
+            self.attn.write_layer_caches(k, v, idx_k, attn_metadata)
             # Publish the selected blocks so the FMHA runs the sparse path.
-            kv_block_indexes = self.attn.run_indexer(idx_q, idx_k, attn_metadata)
+            # idx_k_prewritten: index-K is already in the cache (written above
+            # on bf16, or by the FP8 producer), so run_indexer must not write it.
+            kv_block_indexes = self.attn.run_indexer(
+                idx_q, idx_k, attn_metadata, idx_k_prewritten=True
+            )
             forward_args = AttentionForwardArgs(
                 output=output,
                 sparse_backend_args=SparseBackendForwardArgs(topk_indices=kv_block_indexes),
             )
         else:
             assert idx_q is None and idx_k is None
+            self.attn.write_layer_caches(k, v, None, attn_metadata)
             # No top-k selection means the FMHA attends the full page table.
             forward_args = AttentionForwardArgs(output=output)
-        self.attn.forward(q, k, v, attn_metadata, forward_args=forward_args)
+        self.attn.forward(q, None, None, attn_metadata, forward_args=forward_args)
         return output
 
     def _sparse_forward(