diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py index 7917f96c71dc..d6dcf26fd834 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -134,7 +134,12 @@ def run_msa_paged_gqa( head_dim = attn.head_dim kv_cache_manager = metadata.kv_cache_manager num_tokens = int(q.shape[0]) - if k is not None and v is not None: + # The fused per-layer scatter (msa_write_layer_caches) may have written + # this layer's K/V already; consume the marker so it never goes stale. + prewritten = getattr(metadata, "_msa_prewritten_layer", None) == layer_idx + if prewritten: + metadata._msa_prewritten_layer = None + if k is not None and v is not None and not prewritten: write_msa_main_kv( kv_cache_manager, layer_idx, metadata.msa_out_cache_loc[:num_tokens], k, v ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py index e80c76857c6a..e30a71b832ec 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -254,6 +254,10 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): msa_kv_indices: Optional[torch.Tensor] = None msa_max_score: Optional[torch.Tensor] = None msa_n_valid_blocks: Optional[torch.Tensor] = None + # Layer whose K/V/index-K caches were already written this step by the + # fused scatter (msa_write_layer_caches); run_msa_paged_gqa consumes and + # clears it so the legacy per-cache writes are skipped exactly once. + _msa_prewritten_layer: Optional[int] = None # _msa_buffers_ready gates the once-only device buffers; # _msa_fields_ready marks that the current step's buffers are populated. @@ -839,6 +843,9 @@ def _build_msa_fields(self) -> None: buffers. The transient builder tensors are discarded. """ self._msa_fields_ready = False + # Drop any prewritten marker a failed prior step left unconsumed, so + # it can never suppress a later step's cache write. + self._msa_prewritten_layer = None if not self._msa_buffers_ready: return request_ids = self.request_ids @@ -920,6 +927,53 @@ def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: layout="HND", ) + def msa_write_layer_caches( + self, + layer_idx: int, + k: torch.Tensor, + v: torch.Tensor, + idx_k: Optional[torch.Tensor] = None, + ) -> None: + """Write a layer's new-token K, V, and (sparse layers) index-K. + + One fused kernel launch when the source/cache layouts allow it, else + the legacy per-cache writes. Runs before the indexer's proxy pass + reads the index-K cache; the layer is recorded in + _msa_prewritten_layer so run_msa_paged_gqa skips its own K/V write. + Requires prepared metadata (msa_out_cache_loc filled), the same + contract as the writes it replaces. + """ + from .msa_scatter import fused_write_layer_caches + + buffers = self.kv_cache_manager.get_buffers(layer_idx, kv_layout="HND") + k_view, v_view = buffers[:, 0], buffers[:, 1] + idx_cache = self.msa_idx_k_cache(layer_idx) if idx_k is not None else None + num_tokens = int(k.shape[0]) + out_cache_loc = self.msa_out_cache_loc[:num_tokens] + if not fused_write_layer_caches(k_view, v_view, idx_cache, out_cache_loc, k, v, idx_k): + 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", + ) + self._msa_prewritten_layer = layer_idx + def msa_proxy_max_score_view( self, num_index_heads: int, plan_max_k_tiles: int, num_tokens: int ) -> torch.Tensor: @@ -1014,13 +1068,16 @@ def run_indexer( 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.topk_indices. Returns [total_q, num_kv_heads, topk]. Decode uses the prebuilt graph-safe proxy plan; prefill and mixed - batches use the prebuilt eager proxy plan. + batches use the prebuilt eager proxy plan. `idx_k_prewritten` marks + that the fused per-layer cache write (msa_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 @@ -1030,10 +1087,17 @@ def run_indexer( # scatter below both honor the source strides. idx_q_view = idx_q.reshape(num_tokens, config.num_index_heads, config.sparse_index_dim) idx_k_cache = metadata.msa_idx_k_cache(self.layer_idx) - if idx_k is not None: + # Index-K may already be in the cache by two routes: the fused per-layer + # write (msa_write_layer_caches, idx_k_prewritten=True) stored a live + # bf16 idx_k, or the FP8 producer inserted FP8 index-K and passed + # idx_k=None. Write here only when neither owns it — i.e. a live idx_k + # that was not pre-written. + if idx_k is not None and not idx_k_prewritten: idx_k_view = idx_k.reshape(num_tokens, 1, config.sparse_index_dim) metadata.msa_write_idx_k(self.layer_idx, idx_k_view) - elif idx_k_cache.dtype != torch.float8_e4m3fn or idx_q_view.dtype != torch.float8_e4m3fn: + elif idx_k is None and ( + idx_k_cache.dtype != torch.float8_e4m3fn or idx_q_view.dtype != torch.float8_e4m3fn + ): raise ValueError( "A missing live index-K is valid only when the fused MiniMax-M3 " "producer already emitted FP8 index-Q and inserted FP8 index-K." diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_scatter.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_scatter.py new file mode 100644 index 000000000000..94e333bcf653 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_scatter.py @@ -0,0 +1,163 @@ +# 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 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 head dim 1. + `k`/`v` are the layer's new-token values, [T, H*D] or [T, H, D] row views; + `idx_k` is [T, D] or [T, 1, D]. + """ + if not (k.is_cuda and k_cache.is_cuda): + return False + if k_cache.dim() != 4 or v_cache.dim() != 4: + 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 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 + + _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"] diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 60b7c70f9368..563471e0ab8e 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -1489,11 +1489,24 @@ def _msa_attention_core( """ if self.is_sparse_attention_layer: assert idx_q is not None + # One launch writes this layer's K/V and, on the bf16 path, index-K, + # ahead of the proxy pass that reads the index-K cache. On the FP8 + # indexer path idx_k is None because the fused producer already + # inserted FP8 index-K into the side cache; the fused write then + # stores K/V only. + attn_metadata.msa_write_layer_caches(self.attn.layer_idx, k, v, idx_k) # 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 marks that index-K is already in the cache (via + # the fused write above on bf16, or the FP8 producer when idx_k is + # None), so run_indexer must not write it again. + kv_block_indexes = self.attn.run_indexer( + idx_q, idx_k, attn_metadata, idx_k_prewritten=True + ) forward_args = AttentionForwardArgs(output=output, topk_indices=kv_block_indexes) else: assert idx_q is None and idx_k is None + # Dense layers get the same fused K/V write. + attn_metadata.msa_write_layer_caches(self.attn.layer_idx, k, v) # 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) diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index 3ce0bdd9033a..90639786c8cf 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -2,15 +2,20 @@ # SPDX-License-Identifier: Apache-2.0 """Structural tests for the MiniMax-M3 MSA sparse attention backend. -These validate backend selection and decode scratch-buffer sizing without -launching kernels. Numerical parity against the Triton reference is covered -by the SM100 integration accuracy test. +Most of these validate backend selection and decode scratch-buffer sizing +without launching kernels; the CUDA-gated tests cover the fused cache-scatter +parity and strided-cache kernel paths. Numerical parity against the Triton +reference is covered by the SM100 integration accuracy test. """ import pytest import torch from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import write_kv_slots +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_scatter import ( + fused_write_layer_caches, +) from tensorrt_llm._torch.attention_backend.sparse.utils import _resolve_minimax_m3_backend_cls from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig @@ -353,3 +358,60 @@ def test_per_token_valid_blocks_multi_token_decode(): n_valid = per_token_valid_blocks(qo, kv, off, causal=True, block_size=4) # Row 0: 9 positions -> 3 blocks. Row 1 tokens attend 4, 5, 6 -> 1, 2, 2. assert n_valid.tolist() == [3, 1, 2, 2] + + +def _reference_scatter_write(k_cache, v_cache, idx_cache, slots, k, v, idx_k): + num_tokens = int(slots.shape[0]) + num_heads, head_dim = int(k_cache.shape[1]), int(k_cache.shape[3]) + write_kv_slots(k_cache, slots, k.reshape(num_tokens, num_heads, head_dim), layout="HND") + write_kv_slots(v_cache, slots, v.reshape(num_tokens, num_heads, head_dim), layout="HND") + if idx_k is not None: + write_kv_slots(idx_cache, slots, idx_k.reshape(num_tokens, 1, head_dim), layout="HND") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("num_kv_heads", [1, 4]) +@pytest.mark.parametrize("with_idx", [True, False]) +def test_fused_scatter_matches_reference(cache_dtype, num_kv_heads, with_idx): + """The fused per-layer cache scatter must match the legacy write_kv_slots + path exactly on production-shaped inputs: non-contiguous HND cache views + carved from a pooled allocation and strided source rows sliced from a fused + projection, including the bf16 -> fp8 cache cast. Asserting on the whole + pool also catches stray writes outside the targeted slots.""" + torch.manual_seed(0) + device = "cuda" + num_pages, tokens_per_block, head_dim = 6, 32, 128 + num_tokens = 17 + inner = num_kv_heads * head_dim + + # Paged HND caches carved from a pool with a coalescing axis, so the + # views are non-contiguous like production get_buffers(...) output. + pool = torch.zeros( + num_pages, 2, num_kv_heads, tokens_per_block, head_dim, dtype=cache_dtype, device=device + ) + k_cache, v_cache = pool[:, 0], pool[:, 1] + idx_pool = torch.zeros( + num_pages, 2, 1, tokens_per_block, head_dim, dtype=torch.bfloat16, device=device + ) + idx_cache = idx_pool[:, 0] + + # Strided sources: rows sliced out of a wider fused-projection tensor. + qkv = torch.randn(num_tokens, 3 * inner + 64, dtype=torch.bfloat16, device=device) + k = qkv[:, :inner] + v = qkv[:, inner : 2 * inner] + idx_k = qkv[:, 2 * inner : 2 * inner + head_dim] if with_idx else None + + slots = torch.randperm(num_pages * tokens_per_block, device=device)[:num_tokens].to(torch.int32) + + ref_pool = pool.clone() + ref_idx_pool = idx_pool.clone() + _reference_scatter_write(ref_pool[:, 0], ref_pool[:, 1], ref_idx_pool[:, 0], slots, k, v, idx_k) + + wrote = fused_write_layer_caches( + k_cache, v_cache, idx_cache if with_idx else None, slots, k, v, idx_k + ) + assert wrote + + torch.testing.assert_close(pool.to(torch.float32), ref_pool.to(torch.float32)) + torch.testing.assert_close(idx_pool, ref_idx_pool)