From 12b1dd65301cd5de3b8d10cc40340400a9956911 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Sun, 13 Sep 2026 21:17:59 -0700 Subject: [PATCH 1/3] [None][perf] Fuse MiniMax-M3 MSA per-layer KV-cache writes into one kernel Port of #16755 from feat/m3_with_msa to main. Each MSA layer wrote its new-token main K, main V and (sparse layers) index-K through three separate aten advanced-indexing writes, each with its own division / remainder / cast preprocessing: ~12 tiny launches per sparse layer, ~720 per decode step at 60 layers, all captured into the decode CUDA graphs. Replace them with one Triton launch per layer that derives (page, within-page) from out_cache_loc in-register and writes K, V and index-K together before the indexer's proxy pass. Layouts the kernel cannot take fall back to the legacy per-cache writes. Rebased onto the MsaPrefillFmha / MsaDecodeFmha split (#18611): the kernel lives in minimax_m3/kernels alongside the other cache writes, and the per-phase write_msa_phase_kv is what now skips a layer the fused scatter already wrote. run_indexer keeps main's strict indexer_kv_dtype validation and gates the bf16 index-K write on idx_k_prewritten. Signed-off-by: Zheyu Fu --- .../sparse/minimax_m3/kernels/msa_scatter.py | 173 ++++++++++++++++++ .../sparse/minimax_m3/kernels/msa_utils.py | 5 + .../backends/sparse/minimax_m3/msa_backend.py | 68 ++++++- .../_torch/models/modeling_minimaxm3.py | 15 +- .../attention/sparse/msa/test_msa_backend.py | 104 ++++++++++- 5 files changed, 359 insertions(+), 6 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py 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..b71c8c9f0b6c 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 @@ -150,6 +150,11 @@ def write_msa_phase_kv( ) if k is None or v is None: return + # The fused per-layer scatter (metadata.msa_write_layer_caches) may have + # written this whole step's K/V for the layer already; both phases of a + # mixed step then skip, and prepare() clears the marker for the next step. + if getattr(metadata, "_msa_prewritten_layer", None) == attn.layer_idx: + return num_tokens = int(k.shape[0]) if num_tokens == 0: return 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..1513b8f8e547 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 @@ -155,6 +155,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); write_msa_phase_kv skips the + # legacy per-cache writes for it. Reset when the next step is prepared. + _msa_prewritten_layer: Optional[int] = None # The same page table and lengths as msa_kv_indices / msa_kv_lens, in the # per-request 2-D form the decode kernels index directly # (block_table[request, block] and seq_lens[request]). fmha_sm100 instead @@ -962,6 +966,9 @@ def _build_msa_fields(self) -> None: decoding the inputs for on_update_kv_lens are staged as well. """ self._msa_fields_ready = False + # Drop the previous step's prewritten marker so it can never suppress + # this step's cache write. + self._msa_prewritten_layer = None if not self._msa_buffers_ready: return request_ids = self.request_ids @@ -1085,6 +1092,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 write_msa_phase_kv skips its own K/V write. + Requires prepared metadata (msa_out_cache_loc filled), the same + contract as the writes it replaces. + """ + from .kernels.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: @@ -1215,6 +1269,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 +1277,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 + (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 @@ -1254,13 +1311,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 (msa_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..0d9b4079bb7d 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -1401,14 +1401,27 @@ 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, sparse_backend_args=SparseBackendForwardArgs(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/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 210db3093c41..942b47e7c9f3 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,95 @@ 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 + + +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]) +@pytest.mark.parametrize( + "input_case", + [ + "supported", + "empty", + "strided", + "short_k", + "short_v", + "short_idx", + "v_shape", + "cpu_v", + "cpu_slots", + ], +) +def test_fused_scatter_matches_reference(cache_dtype, num_kv_heads, with_idx, input_case): + """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() + if input_case == "empty": + slots = slots[:0] + elif input_case == "strided": + k = qkv[:, : 2 * inner : 2] + elif input_case == "short_k": + k = k[:-1] + elif input_case == "short_v": + v = v[:-1] + elif input_case == "short_idx": + if idx_k is None: + pytest.skip("Requires index-K") + idx_k = idx_k[:-1] + elif input_case == "v_shape": + v_cache = v_cache[:, :, :-1, :] + elif input_case == "cpu_v": + v = v.cpu() + elif input_case == "cpu_slots": + slots = slots.cpu() + else: + _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 == (input_case in ("supported", "empty")) + + torch.testing.assert_close(pool.to(torch.float32), ref_pool.to(torch.float32)) + torch.testing.assert_close(idx_pool, ref_idx_pool) From 3ec3bb80da10ac233e79ffdb1d5c9328a0b693a5 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Sun, 13 Sep 2026 21:21:27 -0700 Subject: [PATCH 2/3] [None][refactor] Make the MiniMax-M3 model layer own the MSA cache write Review follow-up on the fused per-layer KV-cache write. The fused scatter left a per-step marker (_msa_prewritten_layer) on the attention metadata, set by a metadata method and consumed by the FMHA's K/V write. That is control flow, not a description of the step: the metadata is shared read-only by every layer, and "the model layer writes K/V for MSA layers" holds for every layer and every step, so it should not be per-step state at all. Drop the marker and the metadata method. The write moves to MiniMaxM3MsaSparseAttention.write_layer_caches, next to run_indexer, which the model layer already calls; metadata only supplies the write slots and the cache manager. After writing, the model layer hands forward() k=v=None, which is already the phase libraries' contract for "K/V are resident" (write_msa_phase_kv writes nothing without live K/V), so no cross-module state is needed to suppress the second write. Tests: cover the fp8 source into fp8 cache pairing the FP8-KV production path takes (the fused QK-norm+RoPE kernel emits E4M3 k/v), the phase libraries' no-K/V contract, and the model layer's call order (write, then indexer with idx_k_prewritten, then forward with k=v=None) on the bf16 indexer, FP8 indexer and dense layers. Signed-off-by: Zheyu Fu --- .../sparse/minimax_m3/kernels/msa_utils.py | 11 +- .../backends/sparse/minimax_m3/msa_backend.py | 109 ++++++++-------- .../_torch/models/modeling_minimaxm3.py | 25 ++-- .../attention/sparse/msa/test_msa_backend.py | 121 +++++++++++++++++- 4 files changed, 187 insertions(+), 79 deletions(-) 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 b71c8c9f0b6c..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( @@ -150,11 +154,6 @@ def write_msa_phase_kv( ) if k is None or v is None: return - # The fused per-layer scatter (metadata.msa_write_layer_caches) may have - # written this whole step's K/V for the layer already; both phases of a - # mixed step then skip, and prepare() clears the marker for the next step. - if getattr(metadata, "_msa_prewritten_layer", None) == attn.layer_idx: - return num_tokens = int(k.shape[0]) if num_tokens == 0: return 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 1513b8f8e547..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 @@ -155,10 +155,6 @@ 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); write_msa_phase_kv skips the - # legacy per-cache writes for it. Reset when the next step is prepared. - _msa_prewritten_layer: Optional[int] = None # The same page table and lengths as msa_kv_indices / msa_kv_lens, in the # per-request 2-D form the decode kernels index directly # (block_table[request, block] and seq_lens[request]). fmha_sm100 instead @@ -966,9 +962,6 @@ def _build_msa_fields(self) -> None: decoding the inputs for on_update_kv_lens are staged as well. """ self._msa_fields_ready = False - # Drop the previous step's prewritten marker so it can never suppress - # this step's cache write. - self._msa_prewritten_layer = None if not self._msa_buffers_ready: return request_ids = self.request_ids @@ -1092,53 +1085,6 @@ 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 write_msa_phase_kv skips its own K/V write. - Requires prepared metadata (msa_out_cache_loc filled), the same - contract as the writes it replaces. - """ - from .kernels.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: @@ -1262,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, @@ -1278,7 +1275,7 @@ def run_indexer( 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 - (msa_write_layer_caches) already stored this layer's index-K. + (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 @@ -1311,7 +1308,7 @@ 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}." ) - # The fused per-layer write (msa_write_layer_caches, signalled by + # 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: diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 0d9b4079bb7d..782c5e6a3e6c 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -1398,19 +1398,21 @@ 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 - # 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) + # 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. - # 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. + # 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 ) @@ -1420,11 +1422,10 @@ def _msa_attention_core( ) 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) + 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 942b47e7c9f3..369cfb55004c 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -1596,7 +1596,20 @@ def _reference_scatter_write(k_cache, v_cache, idx_cache, slots, k, v, idx_k): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@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("num_kv_heads", [1, 4]) @pytest.mark.parametrize("with_idx", [True, False]) @pytest.mark.parametrize( @@ -1613,12 +1626,15 @@ def _reference_scatter_write(k_cache, v_cache, idx_cache, slots, k, v, idx_k): "cpu_slots", ], ) -def test_fused_scatter_matches_reference(cache_dtype, num_kv_heads, with_idx, input_case): +def test_fused_scatter_matches_reference( + src_dtype, cache_dtype, num_kv_heads, with_idx, input_case +): """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.""" + 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, tokens_per_block, head_dim = 6, 32, 128 @@ -1637,10 +1653,15 @@ def test_fused_scatter_matches_reference(cache_dtype, num_kv_heads, with_idx, in 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] - 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) @@ -1676,3 +1697,93 @@ def test_fused_scatter_matches_reference(cache_dtype, num_kv_heads, with_idx, in torch.testing.assert_close(pool.to(torch.float32), ref_pool.to(torch.float32)) torch.testing.assert_close(idx_pool, ref_idx_pool) + + +def test_a_phase_handed_no_kv_writes_nothing(): + """k=v=None is the model layer's signal that it already wrote the step's + K/V through write_layer_caches; a phase must then leave the cache alone + rather than fail, or the fused write would be undone or duplicated.""" + from tensorrt_llm._torch.attention.backends.interface import AttentionInputType + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + write_msa_phase_kv, + ) + + buffers = torch.zeros(2, 2, 1, 8, 4) + get_buffers_calls = [] + + def get_buffers(layer_idx, kv_layout=None): + get_buffers_calls.append(layer_idx) + return buffers + + metadata = SimpleNamespace( + kv_cache_manager=SimpleNamespace(get_buffers=get_buffers), + msa_out_cache_loc=torch.tensor([1, 2, 11], dtype=torch.int32), + ) + attention = SimpleNamespace(layer_idx=0) + + write_msa_phase_kv(attention, None, None, metadata, AttentionInputType.mixed, token_offset=0) + write_msa_phase_kv(attention, None, None, metadata, AttentionInputType.mixed, token_offset=2) + + assert get_buffers_calls == [] + assert int((buffers != 0).sum()) == 0 + + +@pytest.mark.parametrize("layer_case", ["sparse_bf16_indexer", "sparse_fp8_indexer", "dense"]) +def test_msa_attention_core_owns_the_cache_write(layer_case): + """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. Checked on the bf16 indexer (live idx_k), the FP8 indexer (idx_k + None, cache populated by the fused producer) and the dense layers.""" + from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3Attention + + sparse = layer_case != "dense" + 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 layer_case == "sparse_bf16_indexer" 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 From 158db48516f5f4c1c451793e3c80577bec28cdad Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Mon, 14 Sep 2026 11:16:46 -0700 Subject: [PATCH 3/3] [None][test] Keep only the essential fused MSA cache-write tests The fused-scatter test fanned out to 108 cases; eight of its nine input cases exercised pure-Python layout preconditions that do not depend on dtype or head count. Keep the numerical check against the legacy write_kv_slots path per source/cache dtype pairing, with and without index-K, at one head count. Drop the no-K/V phase test, which covered a pre-existing early return, and the FP8-indexer variant of the call-order test, which only differed in passing idx_k=None through. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Zheyu Fu --- .../attention/sparse/msa/test_msa_backend.py | 102 +++--------------- 1 file changed, 16 insertions(+), 86 deletions(-) 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 369cfb55004c..488ef12a8459 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -1586,15 +1586,6 @@ def test_on_update_kv_lens_is_a_noop_without_speculative_decoding(monkeypatch): assert metadata.msa_n_valid_blocks.tolist() == [0] * 5 -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( ("src_dtype", "cache_dtype"), @@ -1610,25 +1601,8 @@ def _reference_scatter_write(k_cache, v_cache, idx_cache, slots, k, v, idx_k): (torch.float8_e4m3fn, torch.float8_e4m3fn), ], ) -@pytest.mark.parametrize("num_kv_heads", [1, 4]) @pytest.mark.parametrize("with_idx", [True, False]) -@pytest.mark.parametrize( - "input_case", - [ - "supported", - "empty", - "strided", - "short_k", - "short_v", - "short_idx", - "v_shape", - "cpu_v", - "cpu_slots", - ], -) -def test_fused_scatter_matches_reference( - src_dtype, cache_dtype, num_kv_heads, with_idx, input_case -): +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 @@ -1637,7 +1611,7 @@ def test_fused_scatter_matches_reference( slots.""" torch.manual_seed(0) device = "cuda" - num_pages, tokens_per_block, head_dim = 6, 32, 128 + num_pages, num_kv_heads, tokens_per_block, head_dim = 6, 4, 32, 128 num_tokens = 17 inner = num_kv_heads * head_dim @@ -1667,78 +1641,34 @@ def test_fused_scatter_matches_reference( ref_pool = pool.clone() ref_idx_pool = idx_pool.clone() - if input_case == "empty": - slots = slots[:0] - elif input_case == "strided": - k = qkv[:, : 2 * inner : 2] - elif input_case == "short_k": - k = k[:-1] - elif input_case == "short_v": - v = v[:-1] - elif input_case == "short_idx": - if idx_k is None: - pytest.skip("Requires index-K") - idx_k = idx_k[:-1] - elif input_case == "v_shape": - v_cache = v_cache[:, :, :-1, :] - elif input_case == "cpu_v": - v = v.cpu() - elif input_case == "cpu_slots": - slots = slots.cpu() - else: - _reference_scatter_write( - ref_pool[:, 0], ref_pool[:, 1], ref_idx_pool[:, 0], slots, k, v, idx_k + 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" ) - wrote = fused_write_layer_caches( + assert fused_write_layer_caches( k_cache, v_cache, idx_cache if with_idx else None, slots, k, v, idx_k ) - assert wrote == (input_case in ("supported", "empty")) torch.testing.assert_close(pool.to(torch.float32), ref_pool.to(torch.float32)) torch.testing.assert_close(idx_pool, ref_idx_pool) -def test_a_phase_handed_no_kv_writes_nothing(): - """k=v=None is the model layer's signal that it already wrote the step's - K/V through write_layer_caches; a phase must then leave the cache alone - rather than fail, or the fused write would be undone or duplicated.""" - from tensorrt_llm._torch.attention.backends.interface import AttentionInputType - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( - write_msa_phase_kv, - ) - - buffers = torch.zeros(2, 2, 1, 8, 4) - get_buffers_calls = [] - - def get_buffers(layer_idx, kv_layout=None): - get_buffers_calls.append(layer_idx) - return buffers - - metadata = SimpleNamespace( - kv_cache_manager=SimpleNamespace(get_buffers=get_buffers), - msa_out_cache_loc=torch.tensor([1, 2, 11], dtype=torch.int32), - ) - attention = SimpleNamespace(layer_idx=0) - - write_msa_phase_kv(attention, None, None, metadata, AttentionInputType.mixed, token_offset=0) - write_msa_phase_kv(attention, None, None, metadata, AttentionInputType.mixed, token_offset=2) - - assert get_buffers_calls == [] - assert int((buffers != 0).sum()) == 0 - - -@pytest.mark.parametrize("layer_case", ["sparse_bf16_indexer", "sparse_fp8_indexer", "dense"]) -def test_msa_attention_core_owns_the_cache_write(layer_case): +@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. Checked on the bf16 indexer (live idx_k), the FP8 indexer (idx_k - None, cache populated by the fused producer) and the dense layers.""" + again.""" from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3Attention - sparse = layer_case != "dense" num_tokens, width = 3, 128 topk_indices = torch.zeros(num_tokens, 1, 16, dtype=torch.int32) events = [] @@ -1759,7 +1689,7 @@ def forward(self, q, k, v, metadata, forward_args=None): 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 layer_case == "sparse_bf16_indexer" else None + idx_k = torch.zeros(num_tokens, width) if sparse else None metadata = object() output = torch.empty(num_tokens, width)