[None][perf] Avoid Index-K cache materialization for MSA

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

@@ -258,18 +258,18 @@ def _torch_dtype_for_index_cache(self) -> torch.dtype:
             return torch.float32
         return torch.bfloat16
 
-    def get_index_k_buffer(self, layer_idx: int) -> Optional[torch.Tensor]:
+    def get_index_k_buffer(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]:
         """Return the V2-managed paged index-K view for ``layer_idx``.
 
-        Shape: ``[num_pages, tokens_per_block, 1, sparse_index_dim]``.
-        Reads/writes decompose ``slot = (page, within)`` and use
-        multi-dim fancy indexing; writes propagate to pool storage.
+        NHD shape is ``[num_pages, tokens_per_block, 1, sparse_index_dim]``;
+        HND shape is ``[num_pages, 1, tokens_per_block, sparse_index_dim]``.
         """
         return super().get_index_k_buffer(
             layer_idx,
             num_heads=1,
             head_dim=self.sparse_index_dim,
             dtype=self._torch_dtype_for_index_cache(),
+            kv_layout=kv_layout,
         )
 
     def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]:

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

@@ -609,8 +609,8 @@ def _build_msa_fields(self) -> None:
         self._msa_fields_ready = True
 
     def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor:
-        """Paged index-K view for the indexer; HND conversion is done there."""
-        return self.kv_cache_manager.get_index_k_buffer(layer_idx)
+        """Return the paged index-K cache in the HND layout MSA consumes."""
+        return self.kv_cache_manager.get_index_k_buffer(layer_idx, kv_layout="HND")
 
     def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None:
         """Write the new-token index-K into the side cache at out_cache_loc."""
@@ -621,6 +621,7 @@ def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None:
             cache,
             self.msa_out_cache_loc[:num_tokens],
             idx_k.reshape(num_tokens, 1, sparse_index_dim),
+            layout="HND",
         )
 
     def msa_proxy_max_score_view(

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

@@ -19,7 +19,6 @@
 
 from .msa_utils import (
     MSA_REQUIRED_TOPK,
-    cache_view_to_msa_paged,
     per_token_valid_blocks,
     require_msa_module,
     select_blocks_from_maxscore,
@@ -122,7 +121,7 @@ def __init__(self, config: "MiniMaxM3SparseConfig"):
     def select_blocks(
         self,
         idx_q: torch.Tensor,
-        idx_k_cache: torch.Tensor,
+        idx_k_paged: torch.Tensor,
         *,
         idx_sm_scale: float,
         kv_indices: torch.Tensor,
@@ -144,7 +143,6 @@ def select_blocks(
         both, and generation is the one-query-token-per-request special case.
         """
         config = self.config
-        idx_k_paged = cache_view_to_msa_paged(idx_k_cache)
 
         if proxy_plan is None:
             max_score = _proxy_max_score(

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

@@ -72,24 +72,6 @@ def require_msa_module():
     return fmha_sm100
 
 
-def cache_view_to_msa_paged(cache_view: torch.Tensor) -> torch.Tensor:
-    """Convert a KV cache view to the fmha_sm100 HND paged layout.
-
-    A 4-D paged view [num_pages, page_size, num_heads, head_dim] permutes to
-    [num_pages, num_heads, page_size, head_dim]. A 3-D flat-slot cache
-    [num_slots, num_heads, head_dim] is treated as one virtual page, giving
-    [1, num_heads, num_slots, head_dim].
-    """
-    if cache_view.dim() == 4:
-        return cache_view.permute(0, 2, 1, 3).contiguous()
-    if cache_view.dim() == 3:
-        return cache_view.permute(1, 0, 2).unsqueeze(0).contiguous()
-    raise ValueError(
-        f"Unsupported cache view rank {cache_view.dim()} for MSA paged conversion; "
-        "expected 3 (flat-slot) or 4 (paged)."
-    )
-
-
 def msa_paged_kv(kv_cache_manager, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
     """Return per-layer paged K and V in fmha_sm100 HND layout, zero-copy.
 
@@ -247,7 +229,6 @@ def select_blocks_from_maxscore(
     "MSA_REQUIRED_HEAD_DIM",
     "MSA_REQUIRED_TOPK",
     "build_kv_page_indices",
-    "cache_view_to_msa_paged",
     "msa_package_available",
     "msa_paged_kv",
     "per_token_valid_blocks",

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -1891,18 +1891,19 @@ def get_index_k_buffer(
         num_heads: int = 1,
         head_dim: int,
         dtype: Union[torch.dtype, "DataType", str] = torch.bfloat16,
+        kv_layout: str = "NHD",
     ) -> Optional[torch.Tensor]:
         """Return a torch view over the V2-managed paged ``Role.INDEX_KEY``
         buffer for ``layer_idx``, or ``None`` when the layer has no
         INDEX_KEY buffer registered (e.g. dense layers in a sparse model,
         or non-local layers on the current PP rank).
 
-        The view has shape ``[num_pages, tokens_per_block, num_heads,
-        head_dim]`` where ``num_pages == impl.get_page_index_upper_bound(
-        layer_idx, Role.INDEX_KEY)``. Sparse modeling code addresses
-        entries by ``(page, within_page, head, dim)`` after decomposing
-        the per-token slot id used by the main paged K/V cache into
-        ``(page, within_page)``.
+        For ``kv_layout="NHD"``, the view has shape ``[num_pages,
+        tokens_per_block, num_heads, head_dim]``. For ``kv_layout="HND"``,
+        it has shape ``[num_pages, num_heads, tokens_per_block, head_dim]``.
+        Sparse modeling code decomposes the per-token slot id used by the
+        main paged K/V cache into ``(page, within_page)`` and indexes the
+        token axis selected by ``kv_layout``.
 
         Because :class:`BufferConfig` only carries an opaque byte ``size``
         per block, the dtype and head shape are caller-side contracts.
@@ -1921,6 +1922,8 @@ def get_index_k_buffer(
         propagate to the pool, and successive calls return views over
         the same backing storage.
         """
+        if kv_layout not in ("NHD", "HND"):
+            raise ValueError(f"Unsupported kv_layout: {kv_layout}")
         if layer_idx not in self.layer_offsets:
             return None
         layer_offset = self.layer_offsets[layer_idx]
@@ -1975,21 +1978,21 @@ def get_index_k_buffer(
             f"{num_slots_total} is not divisible by scale = {scale}."
         )
         num_slots = num_slots_total // scale
+        if kv_layout == "NHD":
+            page_shape = [self.tokens_per_block, num_heads, head_dim]
+        else:
+            page_shape = [num_heads, self.tokens_per_block, head_dim]
 
         if scale == 1:
             # Non-coalesced INDEX_KEY pool: the per-buffer stride is
-            # the entire page, so ``[page_upper, tokens_per_block,
-            # num_heads, head_dim]`` is the correct contiguous view.
-            shape = [page_upper, self.tokens_per_block, num_heads, head_dim]
+            # the entire page, so either layout is a contiguous view.
+            shape = [page_upper, *page_shape]
             return convert_to_torch_tensor(TensorWrapper(addr, torch_dtype, shape))
 
-        # Coalesced pool: build a ``[num_slots, scale, tokens_per_block,
-        # num_heads, head_dim]`` view at INDEX_KEY's base, then slice
-        # ``[:, 0]`` to extract this layer's INDEX_KEY data. The slice
-        # preserves dim-0 stride = ``scale * page_stride`` bytes, so
-        # ``view[s, w, h, d]`` lands on the correct byte for any
-        # ``s`` in [0, num_slots).
-        full_slot_shape = [num_slots, scale, self.tokens_per_block, num_heads, head_dim]
+        # Coalesced pool: include the buffers-per-slot dimension, then
+        # slice ``[:, 0]`` to extract this layer's INDEX_KEY data while
+        # preserving dim-0 stride = ``scale * page_stride`` bytes.
+        full_slot_shape = [num_slots, scale, *page_shape]
         full_view = convert_to_torch_tensor(TensorWrapper(addr, torch_dtype, full_slot_shape))
         return full_view[:, 0]
 

tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py

@@ -74,3 +74,152 @@ def test_msa_proxy_max_score_view_is_contiguous_over_stable_store():
     # Oversized requests are rejected rather than silently corrupting memory.
     with pytest.raises(ValueError, match=r"msa_max_score backing store"):
         metadata.msa_proxy_max_score_view(num_index_heads, worst_k, max_batch + 1)
+
+
+def test_msa_index_k_uses_hnd_cache_view_and_writer():
+    metadata_cls = MiniMaxM3MsaSparseAttention.Metadata
+    metadata = metadata_cls.__new__(metadata_cls)
+    num_pages, coalescing_scale, page_size, head_dim = 2, 7, 8, 16
+    pool = torch.zeros(
+        num_pages,
+        coalescing_scale,
+        1,
+        page_size,
+        head_dim,
+        dtype=torch.bfloat16,
+    )
+    hnd_cache = pool[:, 0]
+
+    class FakeCacheManager:
+        def __init__(self):
+            self.calls = []
+
+        def get_index_k_buffer(self, layer_idx, kv_layout="NHD"):
+            self.calls.append((layer_idx, kv_layout))
+            return hnd_cache
+
+    manager = FakeCacheManager()
+    metadata.kv_cache_manager = manager
+    metadata.msa_out_cache_loc = torch.tensor([2, page_size + 5], dtype=torch.int32)
+    values = torch.arange(2 * head_dim, dtype=torch.float32).reshape(2, 1, head_dim)
+
+    returned = metadata.msa_idx_k_cache(3)
+    metadata.msa_write_idx_k(3, values)
+
+    assert returned.data_ptr() == hnd_cache.data_ptr()
+    assert not returned.is_contiguous()
+    assert manager.calls == [(3, "HND"), (3, "HND")]
+    torch.testing.assert_close(hnd_cache[0, 0, 2], values[0, 0].to(torch.bfloat16))
+    torch.testing.assert_close(hnd_cache[1, 0, 5], values[1, 0].to(torch.bfloat16))
+
+
+def test_msa_indexer_preserves_strided_hnd_index_k(monkeypatch):
+    import tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer as indexer_module
+    from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import MiniMaxM3SparseConfig
+
+    config = MiniMaxM3SparseConfig(
+        num_q_heads=4,
+        num_kv_heads=1,
+        head_dim=128,
+        num_index_heads=4,
+        sparse_index_dim=128,
+        block_size=128,
+        topk=16,
+    )
+    indexer = indexer_module.MsaIndexer(config)
+    pool = torch.randn(2, 7, 1, 128, 128, dtype=torch.bfloat16)
+    idx_k_paged = pool[:, 0]
+    captured = {}
+
+    def fake_proxy_max_score(idx_q, passed_idx_k, **kwargs):
+        del kwargs
+        captured["idx_k"] = passed_idx_k
+        return torch.zeros(4, 2, idx_q.shape[0])
+
+    expected = torch.zeros(1, 1, 16, dtype=torch.int32)
+
+    def fake_select_blocks_from_maxscore(*args, **kwargs):
+        del args, kwargs
+        return expected
+
+    monkeypatch.setattr(indexer_module, "_proxy_max_score", fake_proxy_max_score)
+    monkeypatch.setattr(
+        indexer_module,
+        "select_blocks_from_maxscore",
+        fake_select_blocks_from_maxscore,
+    )
+
+    result = indexer.select_blocks(
+        torch.zeros(1, 4, 128, dtype=torch.bfloat16),
+        idx_k_paged,
+        idx_sm_scale=128**-0.5,
+        kv_indices=torch.arange(2, dtype=torch.int32),
+        qo_lens_cpu=torch.tensor([1], dtype=torch.int32),
+        kv_lens_cpu=torch.tensor([256], dtype=torch.int32),
+        qo_offset_cpu=torch.tensor([255], dtype=torch.int32),
+    )
+
+    assert captured["idx_k"] is idx_k_paged
+    assert captured["idx_k"].data_ptr() == idx_k_paged.data_ptr()
+    assert not captured["idx_k"].is_contiguous()
+    assert result is expected
+
+
+def test_msa_proxy_max_score_strided_index_k_matches_packed():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    if torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("SM100 (Blackwell) required")
+
+    from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer import _proxy_max_score
+    from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import (
+        msa_package_available,
+    )
+
+    if not msa_package_available():
+        pytest.skip("fmha_sm100 (MSA) not importable")
+
+    page_size = head_dim = 128
+    num_index_heads = 4
+    coalescing_scale = 57
+    kv_lens_cpu = torch.tensor([1, 130, 257, 128, 511, 1024, 33, 900], dtype=torch.int32)
+    pages_per_sequence = (kv_lens_cpu + page_size - 1) // page_size
+    num_pages = int(pages_per_sequence.sum().item())
+
+    generator = torch.Generator(device="cuda").manual_seed(0)
+    index_k_pool = torch.randn(
+        num_pages,
+        coalescing_scale,
+        1,
+        page_size,
+        head_dim,
+        generator=generator,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    index_k_strided = index_k_pool[:, 0]
+    index_k_packed = index_k_strided.contiguous()
+    index_q = torch.randn(
+        kv_lens_cpu.numel(),
+        num_index_heads,
+        head_dim,
+        generator=generator,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    kwargs = {
+        "qo_lens_cpu": torch.ones_like(kv_lens_cpu),
+        "kv_lens_cpu": kv_lens_cpu,
+        "qo_offset_cpu": kv_lens_cpu - 1,
+        "kv_indices": torch.arange(num_pages, device="cuda", dtype=torch.int32),
+        "sm_scale": head_dim**-0.5,
+        "causal": True,
+    }
+
+    strided_scores = _proxy_max_score(index_q, index_k_strided, **kwargs)
+    packed_scores = _proxy_max_score(index_q, index_k_packed, **kwargs)
+    torch.cuda.synchronize()
+
+    assert not index_k_strided.is_contiguous()
+    assert index_k_strided.stride(0) == coalescing_scale * page_size * head_dim
+    assert torch.equal(strided_scores, packed_scores)