From 9b986b902ee100cf9760b64491c93b20abac24a5 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:13:59 -0700 Subject: [PATCH] [None][perf] Avoid Index-K cache materialization for MSA Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../sparse/minimax_m3/cache_manager.py | 8 +- .../sparse/minimax_m3/msa_backend.py | 5 +- .../sparse/minimax_m3/msa_indexer.py | 4 +- .../sparse/minimax_m3/msa_utils.py | 19 --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 35 ++-- .../sparse/test_minimax_m3_msa_backend.py | 149 ++++++++++++++++++ .../test_kv_cache_v2_extra_buffers.py | 89 ++++++++++- 7 files changed, 260 insertions(+), 49 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py index 65c89b170c5c..4609414a8731 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py +++ b/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]: 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 86cd67fa7272..3b0bbb0d2b77 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 @@ -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( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py index a6a2e10e0dc1..b4a836d03ffa 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py +++ b/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( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py index 83380bc730c0..262f5e10ec0d 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py +++ b/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", diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 24808e8b55b5..b9dea96c6174 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/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] 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 58d651631608..2e44c27e9bd3 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 @@ -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) diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py index d23a677dad47..073190eacf71 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py @@ -75,6 +75,15 @@ def _extra_buffers_per_layer(self, *, tokens_per_block): } +class _CoalescedIndexKeyV2(_IndexKeyOnSparseLayersV2): + """Mirror M3's pool mapping when INDEX_KEY coalesces with K/V.""" + + def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr): + del key_base_addr + layers_in_group = list(self.impl.layer_grouping[int(layer_group_id)]) + return layers_in_group.index(int(layer_id)) + + class _DuplicateRoleV2(KVCacheManagerV2): """Negative-control subclass: register Role.KEY as an "extra" so the standard buffer + extra duplicate. Must raise.""" @@ -238,11 +247,10 @@ def test_duplicate_role_against_standard_buffer_asserts(self): class TestIndexKeyBufferAccessor(unittest.TestCase): """CUDA/GPU regressions for :meth:`KVCacheManagerV2.get_index_k_buffer`. - The accessor returns a paged torch view shaped - ``[num_pages, tokens_per_block, num_heads, head_dim]`` over the - managed ``Role.INDEX_KEY`` pool, returns ``None`` for dense layers, - rejects wiring mismatches against the V2-reported page stride, and - is zero-copy (writes propagate; ``data_ptr`` stable across calls). + The accessor returns NHD or HND paged torch views over the managed + ``Role.INDEX_KEY`` pool, returns ``None`` for dense layers, rejects + wiring mismatches against the V2-reported page stride, and is zero-copy + (writes propagate; ``data_ptr`` stable across calls). """ NUM_HEADS = 1 @@ -353,6 +361,62 @@ def test_accessor_view_is_zero_copy_over_pool(self): mgr.shutdown() del mgr + def test_accessor_hnd_and_nhd_alias_coalesced_pool(self): + # Match MiniMax-M3's one-head Index-K byte size to one main K/V + # head so V2 coalesces all roles. Both layout views must retain the + # larger physical page stride while addressing the same bytes. + mgr = _CoalescedIndexKeyV2( + sparse_layer_ids=[0, 1, 2, 3], + **_make_kwargs( + num_layers=4, + num_kv_heads=1, + dtype=DataType.BF16, + ), + ) + try: + layer_idx = 3 + nhd = mgr.get_index_k_buffer( + layer_idx, + num_heads=self.NUM_HEADS, + head_dim=self.HEAD_DIM, + dtype=torch.bfloat16, + ) + hnd = mgr.get_index_k_buffer( + layer_idx, + num_heads=self.NUM_HEADS, + head_dim=self.HEAD_DIM, + dtype=torch.bfloat16, + kv_layout="HND", + ) + self.assertIsNotNone(nhd) + self.assertIsNotNone(hnd) + + converter = mgr.impl.get_page_index_converter( + mgr.layer_offsets[layer_idx], Role.INDEX_KEY + ) + self.assertGreater(int(converter.scale), 1) + self.assertEqual( + nhd.shape, + (nhd.shape[0], mgr.tokens_per_block, self.NUM_HEADS, self.HEAD_DIM), + ) + self.assertEqual( + hnd.shape, + (hnd.shape[0], self.NUM_HEADS, mgr.tokens_per_block, self.HEAD_DIM), + ) + self.assertEqual(nhd.data_ptr(), hnd.data_ptr()) + self.assertEqual(nhd.stride(0), hnd.stride(0)) + self.assertEqual(hnd.stride(2), self.HEAD_DIM) + self.assertEqual(hnd.stride(3), 1) + self.assertFalse(hnd.is_contiguous()) + + sentinel = torch.tensor(37.0, dtype=torch.bfloat16, device="cuda") + hnd[0, 0, 3, 7] = sentinel + torch.cuda.synchronize() + self.assertEqual(nhd[0, 3, 0, 7].item(), float(sentinel)) + finally: + mgr.shutdown() + del mgr + def test_accessor_pointer_matches_v2_pool_base(self): # Verify the wrapper points at the exact V2-managed pool base # for INDEX_KEY, so the view participates in the same lifecycle @@ -416,6 +480,21 @@ def test_accessor_rejects_wrong_dtype(self): mgr.shutdown() del mgr + def test_accessor_rejects_unknown_layout(self): + mgr = self._make_sparse_mgr(sparse_layer_ids=(3,)) + try: + with self.assertRaisesRegex(ValueError, "Unsupported kv_layout"): + mgr.get_index_k_buffer( + 3, + num_heads=self.NUM_HEADS, + head_dim=self.HEAD_DIM, + dtype=torch.bfloat16, + kv_layout="HDN", + ) + finally: + mgr.shutdown() + del mgr + def test_accessor_returns_none_on_base_v2_default_manager(self): # The base KVCacheManagerV2 (no extra buffers) must report None # for every local layer when queried for INDEX_KEY, because the