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 4609414a8731..67c4c7a3e4b5 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 @@ -150,6 +150,8 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): * ``sparse_index_dim`` — width of the index-K/V vectors. """ + _main_kv_layout = "NHD" + def __init__( self, *args, @@ -164,6 +166,8 @@ def __init__( # disable_index_value=True, sparse_index_dim=128). sparse_attn_config = kwargs.get("sparse_attn_config") num_layers = kwargs.get("num_layers") + implementation = getattr(sparse_attn_config, "implementation", "triton") + self._main_kv_layout = "HND" if implementation == "msa" else "NHD" if sparse_index_dim is None: sparse_index_dim = int(getattr(sparse_attn_config, "sparse_index_dim", 0) or 0) or 128 @@ -230,9 +234,10 @@ def _extra_buffers_per_layer(self, *, tokens_per_block): } def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: - """Declare MiniMax M3's token-major K/V and replicated index-K.""" + """Declare the backend's main K/V layout and replicated index-K.""" + main_kv_mapper = MapperKind.INDEXED if self._main_kv_layout == "HND" else MapperKind.NHD return { - Role.ALL: MapperKind.NHD, + Role.ALL: main_kv_mapper, Role.INDEX_KEY: MapperKind.REPLICATED, } @@ -258,12 +263,17 @@ def _torch_dtype_for_index_cache(self) -> torch.dtype: return torch.float32 return torch.bfloat16 - def get_index_k_buffer(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: + def get_index_k_buffer( + self, layer_idx: int, kv_layout: Optional[str] = None + ) -> Optional[torch.Tensor]: """Return the V2-managed paged index-K view for ``layer_idx``. NHD shape is ``[num_pages, tokens_per_block, 1, sparse_index_dim]``; HND shape is ``[num_pages, 1, tokens_per_block, sparse_index_dim]``. + When omitted, ``kv_layout`` follows the selected sparse backend. """ + if kv_layout is None: + kv_layout = self._main_kv_layout return super().get_index_k_buffer( layer_idx, num_heads=1, @@ -279,7 +289,9 @@ def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]: def has_index_value(self, layer_idx: int) -> bool: return layer_idx in self._index_v_buffers - def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: + def get_buffers( + self, layer_idx: int, kv_layout: Optional[str] = None + ) -> Optional[torch.Tensor]: """Return a paged K+V view with strides spanning the coalesced pool. The base :meth:`KVCacheManagerV2.get_buffers` produces a @@ -295,7 +307,10 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. at K's base, then slices ``[:, :2]`` to extract K+V. The slice preserves the dim-0 stride (``scale * page_stride``), so ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``. + When omitted, ``kv_layout`` follows the selected sparse backend. """ + if kv_layout is None: + kv_layout = self._main_kv_layout if kv_layout not in ("NHD", "HND"): raise ValueError(f"Unsupported kv_layout: {kv_layout}") if self.kv_cache_type == CacheTypeCpp.SELFKONLY: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 2c567824a8cb..ccd9553a66b9 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1917,7 +1917,8 @@ def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: ``_extra_buffers_per_layer``. Model-specific managers may declare logical layouts without requiring the shared extractor to inspect private attributes or role names. MiniMax M3, for example, maps - ordinary K/V to ``NHD`` and keeps index-key ``REPLICATED``. + ordinary K/V to ``INDEXED`` for its MSA backend and ``NHD`` for its + Triton backend, while keeping index-key ``REPLICATED``. This declaration does not influence storage pooling: V2 storage coalesces buffers purely by ``(life_cycle, buffer size)``, so roles diff --git a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py index a05b49a0781d..77b8983214bb 100644 --- a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py +++ b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py @@ -26,6 +26,8 @@ """ from collections.abc import Sequence +from types import SimpleNamespace +from typing import Any import kv_transfer_harness as transfer_harness import pytest @@ -93,13 +95,35 @@ def test_v2_disagg_role_mapper_kind_defaults() -> None: } -def test_minimax_disagg_role_mapper_kinds() -> None: - manager = object.__new__(MiniMaxM3KVCacheManagerV2) +@pytest.mark.parametrize( + "implementation,expected_main_mapper", + [("triton", MapperKind.NHD), ("msa", MapperKind.INDEXED)], +) +def test_minimax_disagg_role_mapper_kinds( + monkeypatch: pytest.MonkeyPatch, + implementation: str, + expected_main_mapper: MapperKind, +) -> None: + def fake_base_init(self: KVCacheManagerV2, *args: Any, **kwargs: Any) -> None: + self.is_disagg = kwargs.get("is_disagg", False) + self.dtype = kwargs.get("dtype", DataType.BF16) + self.layer_offsets = {} + self.max_batch_size = 1 + self.max_seq_len = 1 - role_mapper_kinds = manager.get_disagg_role_mapper_kinds() + monkeypatch.setattr(KVCacheManagerV2, "__init__", fake_base_init) + manager = MiniMaxM3KVCacheManagerV2( + num_layers=0, + sparse_layer_ids=[], + disable_index_value_layer_ids=[], + sparse_attn_config=SimpleNamespace( + implementation=implementation, + indexer_kv_dtype="bf16", + ), + ) - assert role_mapper_kinds == { - Role.ALL: MapperKind.NHD, + assert manager.get_disagg_role_mapper_kinds() == { + Role.ALL: expected_main_mapper, Role.INDEX_KEY: MapperKind.REPLICATED, } @@ -152,7 +176,10 @@ def test_minimax_kv_pool_mapping_offset_ignores_layer_grouping_order() -> None: def _create_manager( - mapping: Mapping, dtype: DataType, sparse_layers: list[int] | None = None + mapping: Mapping, + dtype: DataType, + sparse_layers: list[int] | None = None, + implementation: str = "triton", ) -> MiniMaxM3KVCacheManagerV2: max_num_tokens = 2048 kv_cache_dtype = { @@ -180,6 +207,10 @@ def _create_manager( sparse_layer_ids=sparse_layers if sparse_layers is not None else SPARSE_LAYERS, disable_index_value_layer_ids=sparse_layers if sparse_layers is not None else SPARSE_LAYERS, sparse_index_dim=INDEX_DIM, + sparse_attn_config=SimpleNamespace( + implementation=implementation, + indexer_kv_dtype="bf16", + ), ) @@ -243,6 +274,7 @@ def _create_managers( enable_dp: bool, dtype: DataType = DataType.BF16, sparse_layers: list[int] | None = None, + implementation: str = "triton", ) -> list[MiniMaxM3KVCacheManagerV2]: return [ _create_manager( @@ -255,6 +287,7 @@ def _create_managers( ), dtype, sparse_layers, + implementation, ) for rank in range(tp * pp) ] @@ -314,14 +347,21 @@ def _fill_position_dependent( *, layer_idx: int, first_global_head: int, + layout: str, ) -> None: - """Fill ``[block, role, token, head, dim]`` with exact small integers.""" + """Fill an NHD or HND cache view with exact position-dependent integers.""" block = torch.arange(tensor.shape[0], device=tensor.device)[:, None, None, None, None] role = torch.arange(tensor.shape[1], device=tensor.device)[None, :, None, None, None] - token = torch.arange(tensor.shape[2], device=tensor.device)[None, None, :, None, None] - head = (first_global_head + torch.arange(tensor.shape[3], device=tensor.device))[ - None, None, None, :, None - ] + if layout == "HND": + head = (first_global_head + torch.arange(tensor.shape[2], device=tensor.device))[ + None, None, :, None, None + ] + token = torch.arange(tensor.shape[3], device=tensor.device)[None, None, None, :, None] + else: + token = torch.arange(tensor.shape[2], device=tensor.device)[None, None, :, None, None] + head = (first_global_head + torch.arange(tensor.shape[3], device=tensor.device))[ + None, None, None, :, None + ] dim = torch.arange(tensor.shape[4], device=tensor.device)[None, None, None, None, :] values = (layer_idx * 17 + block * 11 + role * 13 + token * 3 + head * 19 + dim) % 97 tensor.copy_(values.to(tensor.dtype)) @@ -334,6 +374,8 @@ def _as_nvfp4_scale_tensor( scale_view = _get_nvfp4_scale_view(manager, layer_idx) if scale_view is None: return None + if manager._main_kv_layout != "NHD": + raise AssertionError("MSA uses an FP8 KV cache; NVFP4 scale pools are NHD-only") local_layer_id = manager.layer_offsets[layer_idx] local_heads = manager.num_kv_heads_per_layer[local_layer_id] bytes_per_token_head = HEAD_DIM // 16 @@ -405,12 +447,13 @@ def _initialize_cache( continue for layer_idx in manager.pp_layers: - kv = manager.get_buffers(layer_idx, kv_layout="NHD") + kv = manager.get_buffers(layer_idx) first_global_head = _first_global_head(manager) _fill_position_dependent( kv, layer_idx=layer_idx, first_global_head=first_global_head, + layout=manager._main_kv_layout, ) index_key = manager.get_index_k_buffer(layer_idx) @@ -420,6 +463,7 @@ def _initialize_cache( index_tensor, layer_idx=layer_idx, first_global_head=0, + layout=manager._main_kv_layout, ) scale_tensor = _as_nvfp4_scale_tensor(manager, layer_idx) @@ -428,6 +472,7 @@ def _initialize_cache( scale_tensor, layer_idx=layer_idx, first_global_head=first_global_head, + layout=manager._main_kv_layout, ) @@ -457,8 +502,9 @@ def _verify_cache( gen_indices = _valid_indices(manager, gen_request_id, layer_idx) assert gen_indices - gen_kv = manager.get_buffers(layer_idx, kv_layout="NHD")[gen_indices] - local_heads = gen_kv.shape[3] + gen_kv = manager.get_buffers(layer_idx)[gen_indices] + head_axis = 2 if manager._main_kv_layout == "HND" else 3 + local_heads = gen_kv.shape[head_axis] first_global_head = _first_global_head(manager) for kv_idx in range(2): for local_head in range(local_heads): @@ -474,13 +520,14 @@ def _verify_cache( ctx_indices = _valid_indices( ctx_manager, ctx_request_ids[req_idx], layer_idx ) - ctx_kv = ctx_manager.get_buffers(layer_idx, kv_layout="NHD") - torch.testing.assert_close( - gen_kv[:, kv_idx, :, local_head, :], - ctx_kv[ctx_indices, kv_idx, :, ctx_local_head, :], - rtol=0, - atol=0, - ) + ctx_kv = ctx_manager.get_buffers(layer_idx) + if manager._main_kv_layout == "HND": + actual = gen_kv[:, kv_idx, local_head, :, :] + expected = ctx_kv[ctx_indices, kv_idx, ctx_local_head, :, :] + else: + actual = gen_kv[:, kv_idx, :, local_head, :] + expected = ctx_kv[ctx_indices, kv_idx, :, ctx_local_head, :] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) index_key = manager.get_index_k_buffer(layer_idx) if index_key is not None: @@ -519,12 +566,13 @@ def _verify_cache( ) ctx_scales = _as_nvfp4_scale_tensor(ctx_manager, layer_idx) assert ctx_scales is not None - torch.testing.assert_close( - gen_scales[gen_indices, :, :, local_head, :], - ctx_scales[ctx_indices, :, :, ctx_local_head, :], - rtol=0, - atol=0, - ) + if manager._main_kv_layout == "HND": + actual = gen_scales[gen_indices, :, local_head, :, :] + expected = ctx_scales[ctx_indices, :, ctx_local_head, :, :] + else: + actual = gen_scales[gen_indices, :, :, local_head, :] + expected = ctx_scales[ctx_indices, :, :, ctx_local_head, :] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) # Production is expected to use TEP/DEP context and DEP generation. Bias the @@ -605,6 +653,36 @@ def test_minimax_m3_kv_transfer( ) +@pytest.mark.cuda +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "update_before_transfer", + [True, False], + ids=["update_before", "update_after"], +) +def test_minimax_m3_msa_hnd_head_mismatch_transfer( + update_before_transfer: bool, +) -> None: + transfer_harness.run_kv_transfer_test( + ctx_tp=1, + ctx_pp=1, + gen_tp=4, + gen_pp=1, + ctx_enable_dp=False, + gen_enable_dp=False, + update_before_transfer=update_before_transfer, + manager_factory=lambda tp, pp, enable_dp: _create_managers( + tp, + pp, + enable_dp, + DataType.FP8, + implementation="msa", + ), + init_fn=_initialize_cache, + verify_fn=_verify_cache, + ) + + # Multiple sparse layers spread across PP stages: the replicated index-key # pool overlaps only partially between peers, exercising the layer-strided # ReplicatedMapper offsets (a single-sparse-layer model always fully diff --git a/tests/unittest/disaggregated/test_pool_matching.py b/tests/unittest/disaggregated/test_pool_matching.py index 8a32a6821444..bcaa8225e03a 100644 --- a/tests/unittest/disaggregated/test_pool_matching.py +++ b/tests/unittest/disaggregated/test_pool_matching.py @@ -267,7 +267,7 @@ def test_kv_and_indexer_in_same_lg(): def test_minimax_split_kv_and_replicated_index_pools_match(): - """MiniMax M3 shape: NHD KV pool + REPLICATED index-key pool. + """MiniMax M3 Triton shape: NHD KV pool + REPLICATED index-key pool. coalescing_group derivation keeps INDEX_KEY in its own pool on every topology, so both sides always present the same two role sets and