diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py new file mode 100644 index 000000000000..7636f5660b17 --- /dev/null +++ b/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"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py index 28716a6dece5..b3aa264d893e 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py +++ b/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( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index 5a3f652fbe8c..78944e3c0dab 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -1208,6 +1208,57 @@ 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, @@ -1215,6 +1266,7 @@ 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. @@ -1222,6 +1274,8 @@ def run_indexer( 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 diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 2495e8dd446b..782c5e6a3e6c 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/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( diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 210db3093c41..488ef12a8459 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -4,8 +4,10 @@ These validate MiniMax-M3 backend selection, indexer/cache integration, decode scratch-buffer sizing, and the paged HND contract passed to the packaged MSA -kernel. Generic block-sparse MQA/GQA numerical coverage lives in the parent -``test_sparse_mqa_gqa.py`` module. +kernel; the CUDA-gated fused cache-scatter test checks the single-launch +K/V/index-K write against the legacy per-cache path. Generic block-sparse +MQA/GQA numerical coverage lives in the parent ``test_sparse_mqa_gqa.py`` +module. """ import sys @@ -20,10 +22,16 @@ MiniMaxM3KVCacheManagerV2, MiniMaxM3MsaSparseAttention, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( MSA_REQUIRED_TOPK, msa_paged_kv, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import MsaDecodeSpan from tensorrt_llm._torch.attention.backends.sparse.registry import _resolve_minimax_m3_backend_cls from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -1576,3 +1584,136 @@ def test_on_update_kv_lens_is_a_noop_without_speculative_decoding(monkeypatch): assert metadata.msa_seq_lens_cuda.tolist() == [4, 6] assert metadata.msa_out_cache_loc.tolist() == [-1] * 5 assert metadata.msa_n_valid_blocks.tolist() == [0] * 5 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + ("src_dtype", "cache_dtype"), + [ + # bf16 K/V into a bf16 cache: the plain path. + (torch.bfloat16, torch.bfloat16), + # bf16 K/V into an fp8 cache: the kernel folds in the E4M3 cast. + (torch.bfloat16, torch.float8_e4m3fn), + # fp8 K/V into an fp8 cache: production with an FP8 KV cache, where + # the fused QK-norm+RoPE kernel already emits E4M3 k/v + # (MiniMaxM3Attention._emit_fp8_main_qkv), so the kernel stores + # without a cast. + (torch.float8_e4m3fn, torch.float8_e4m3fn), + ], +) +@pytest.mark.parametrize("with_idx", [True, False]) +def test_fused_scatter_matches_reference(src_dtype, cache_dtype, 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, for every source/cache dtype pairing the model produces. + Asserting on the whole pool also catches stray writes outside the targeted + slots.""" + torch.manual_seed(0) + device = "cuda" + num_pages, num_kv_heads, tokens_per_block, head_dim = 6, 4, 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. + # randn has no fp8 variant, so generate bf16 and cast the whole buffer, as + # the fused producer would, before slicing the column views. + qkv = torch.randn(num_tokens, 3 * inner + 64, dtype=torch.bfloat16, device=device) + # The index branch stays bf16 on the bf16 indexer path even when the main + # K/V are fp8, so carve index-K out before casting. + idx_k = qkv[:, 2 * inner : 2 * inner + head_dim] if with_idx else None + qkv = qkv.to(src_dtype) + k = qkv[:, :inner] + v = qkv[:, inner : 2 * inner] + + 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() + write_kv_slots( + ref_pool[:, 0], slots, k.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ) + write_kv_slots( + ref_pool[:, 1], slots, v.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ) + if with_idx: + write_kv_slots( + ref_idx_pool[:, 0], slots, idx_k.reshape(num_tokens, 1, head_dim), layout="HND" + ) + + assert fused_write_layer_caches( + k_cache, v_cache, idx_cache if with_idx else None, slots, k, v, idx_k + ) + + torch.testing.assert_close(pool.to(torch.float32), ref_pool.to(torch.float32)) + torch.testing.assert_close(idx_pool, ref_idx_pool) + + +@pytest.mark.parametrize("sparse", [True, False]) +def test_msa_attention_core_owns_the_cache_write(sparse): + """The model layer's MSA core must write the caches exactly once and in + the right place: write_layer_caches runs before run_indexer (whose proxy + pass reads the index-K cache), run_indexer is told index-K is already + resident, and forward() receives k=v=None so no FMHA phase writes K/V + again.""" + from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3Attention + + num_tokens, width = 3, 128 + topk_indices = torch.zeros(num_tokens, 1, 16, dtype=torch.int32) + events = [] + + class FakeBackend: + layer_idx = 7 + + def write_layer_caches(self, k, v, idx_k, metadata): + events.append(("write", k, v, idx_k, metadata)) + + def run_indexer(self, idx_q, idx_k, metadata, *, idx_k_prewritten=False): + events.append(("indexer", idx_q, idx_k, metadata, idx_k_prewritten)) + return topk_indices + + def forward(self, q, k, v, metadata, forward_args=None): + events.append(("forward", q, k, v, metadata, forward_args)) + + layer = SimpleNamespace(is_sparse_attention_layer=sparse, attn=FakeBackend()) + q, k, v = (torch.zeros(num_tokens, width) for _ in range(3)) + idx_q = torch.zeros(num_tokens, width) if sparse else None + idx_k = torch.zeros(num_tokens, width) if sparse else None + metadata = object() + output = torch.empty(num_tokens, width) + + result = MiniMaxM3Attention._msa_attention_core(layer, q, k, v, idx_q, idx_k, metadata, output) + + assert result is output + names = [event[0] for event in events] + if sparse: + assert names == ["write", "indexer", "forward"] + _, indexer_q, indexer_k, indexer_metadata, prewritten = events[1] + assert indexer_q is idx_q and indexer_k is idx_k and indexer_metadata is metadata + assert prewritten is True + else: + assert names == ["write", "forward"] + + _, written_k, written_v, written_idx_k, write_metadata = events[0] + assert written_k is k and written_v is v and write_metadata is metadata + assert written_idx_k is idx_k + + _, forward_q, forward_k, forward_v, forward_metadata, forward_args = events[-1] + assert forward_q is q and forward_metadata is metadata + assert forward_k is None and forward_v is None + assert forward_args.output is output + if sparse: + assert forward_args.sparse_backend_args.topk_indices is topk_indices + else: + assert forward_args.sparse_backend_args is None