[None][fix] Use HND mapping for MiniMax-M3 MSA KV cache

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:

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

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

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