[None][fix] Derive MiniMax-M3 index-K cache layers from the model config

tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py

@@ -53,7 +53,11 @@
     PageIndexMode,
 )
 
+from .common import get_sparse_layer_ids
+
 if TYPE_CHECKING:
+    from transformers import PretrainedConfig
+
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
 
 
@@ -184,6 +188,39 @@ def derive_shared_draft_layout(
     return list(range(num_target, total)), num_target
 
 
+def default_sparse_layer_ids(
+    pretrained_config: Optional[PretrainedConfig], num_target_layers: Optional[int]
+) -> List[int]:
+    """Sparse target layers for a manager built without ``sparse_layer_ids``.
+
+    The model picks dense or sparse attention per layer from the checkpoint's
+    ``sparse_attention_config`` (:func:`get_sparse_layer_ids`), so the index-K
+    cache follows the same split whenever the config describes exactly the
+    target layers. Without such a config, layers 0..2 are dense and the rest
+    sparse, as in the MiniMax-M3 checkpoint.
+    """
+    if num_target_layers is None:
+        return []
+    num_hidden_layers = getattr(pretrained_config, "num_hidden_layers", None)
+    if num_hidden_layers is not None and int(num_hidden_layers) == num_target_layers:
+        return get_sparse_layer_ids(pretrained_config)[1]
+    if num_hidden_layers is not None:
+        # Expected for a manager that does not hold the target decoder layers
+        # (a separate draft manager). Say so when the config's own split is not
+        # the default, since the cache then follows a different split than the model.
+        try:
+            config_sparse = get_sparse_layer_ids(pretrained_config)[1]
+        except ValueError:
+            config_sparse = None
+        if config_sparse is not None and config_sparse != list(range(3, int(num_hidden_layers))):
+            logger.info(
+                f"MiniMax-M3 KV cache: the model config has {int(num_hidden_layers)} layers but this "
+                f"manager holds {num_target_layers} target layers, so its sparse layers default to "
+                f"layers 3..{num_target_layers - 1} instead of the config's split."
+            )
+    return list(range(3, num_target_layers))
+
+
 def extend_attention_op_pools_for_shared_draft_layers(
     pool_pointers: torch.Tensor,
     pool_mapping: torch.Tensor,
@@ -242,7 +279,8 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2):
     stays empty in production.
 
     Constructor extras (forwarded kwargs go to :class:`KVCacheManagerV2`):
-      * ``sparse_layer_ids`` — layer indices using sparse attention.
+      * ``sparse_layer_ids`` — layer indices using sparse attention;
+        defaults to :func:`default_sparse_layer_ids`.
       * ``disable_index_value_layer_ids`` — subset whose index-V is
         omitted.
       * ``sparse_index_dim`` — width of the index-K/V vectors.
@@ -280,11 +318,12 @@ def __init__(
         **kwargs,
     ):
         # Resolve M3 sparse-layer metadata from explicit kwargs first, then
-        # from the executor's ``sparse_attention_config`` keyword, then from
-        # the M3 checkpoint convention (layers 0..2 dense, 3..N-1 sparse,
+        # from the executor's ``sparse_attention_config`` and
+        # ``pretrained_config`` keywords, then from the M3 checkpoint
+        # convention (layers 0..2 dense, 3..N-1 sparse,
         # disable_index_value=True, sparse_index_dim=128). Honoring the
-        # executor keyword also makes non-default sparse_index_dim values
-        # authoritative for the cache layout instead of falling back to 128.
+        # executor keywords also makes a non-default sparse layer split or
+        # sparse_index_dim authoritative for the cache layout.
         # Peeked (not popped) so the base __init__ still receives them.
         sparse_attention_config = kwargs.get("sparse_attention_config")
         num_layers = kwargs.get("num_layers")
@@ -311,10 +350,9 @@ def __init__(
             shared_draft_layer_count(kwargs.get("spec_config"), kwargs.get("layer_mask")),
         )
         if sparse_layer_ids is None:
-            if num_target_layers is not None:
-                sparse_layer_ids = list(range(3, num_target_layers))
-            else:
-                sparse_layer_ids = []
+            sparse_layer_ids = default_sparse_layer_ids(
+                kwargs.get("pretrained_config"), num_target_layers
+            )
         if disable_index_value_layer_ids is None:
             disable_index_value_layer_ids = list(sparse_layer_ids)
 
@@ -1117,6 +1155,7 @@ def get_minimax_m3_kv_cache_manager_cls():
     "MiniMaxM3DraftKVCacheView",
     "MiniMaxM3KVCacheManagerV2",
     "MiniMaxM3SparseIndexCache",
+    "default_sparse_layer_ids",
     "derive_shared_draft_layout",
     "extend_attention_op_pools_for_shared_draft_layers",
     "get_minimax_m3_kv_cache_manager_cls",

tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py

@@ -4,14 +4,15 @@
 
 Both the Triton reference and the MSA (fmha_sm100) path share these
 backend-neutral pieces: the lowered parameter and per-rank kernel config
-bundles, block-priority sentinels, and the paged-cache slot mapping builder.
+bundles, block-priority sentinels, the paged-cache slot mapping builder, and
+the dense/sparse layer split the model and the KV cache manager both read.
 MSA-only helpers live in :mod:`.kernels.msa_utils`.
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Tuple
+from typing import TYPE_CHECKING, List, Literal, NamedTuple, Optional, Tuple
 
 import torch
 
@@ -24,6 +25,8 @@
 from .kernels.paged_cache import write_kv_slots
 
 if TYPE_CHECKING:
+    from transformers import PretrainedConfig
+
     from tensorrt_llm.mapping import Mapping
 
 # Sentinel scores that force init and local blocks into the top-k regardless
@@ -56,6 +59,26 @@ def index_head_range(
     return rank * count, (rank + 1) * count
 
 
+def get_sparse_layer_ids(text_config: "PretrainedConfig") -> Tuple[List[int], List[int]]:
+    """Return ``(dense_layer_ids, sparse_layer_ids)`` for the M3 text model."""
+    sparse_cfg = getattr(text_config, "sparse_attention_config", None)
+    n_layers = int(text_config.num_hidden_layers)
+    if sparse_cfg is None:
+        return list(range(n_layers)), []
+    if not sparse_cfg.get("use_sparse_attention", True):
+        return list(range(n_layers)), []
+    freq = sparse_cfg.get("sparse_attention_freq")
+    if freq is None:
+        return list(range(n_layers)), []
+    if len(freq) != n_layers:
+        raise ValueError(
+            f"sparse_attention_freq length {len(freq)} does not match num_hidden_layers {n_layers}"
+        )
+    dense = [i for i, f in enumerate(freq) if int(f) == 0]
+    sparse = [i for i, f in enumerate(freq) if int(f) != 0]
+    return dense, sparse
+
+
 @dataclass(frozen=True)
 class MiniMaxM3SparseParams(SparseParams):
     """Lowered runtime parameters for the MiniMax-M3 sparse backend."""
@@ -291,5 +314,6 @@ def build_paged_kv_slot_mapping(
     "MiniMaxM3SparseParams",
     "PagedKvSlotMapping",
     "build_paged_kv_slot_mapping",
+    "get_sparse_layer_ids",
     "write_kv_slots",
 ]

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -54,7 +54,7 @@
     _gather_paged_batched,
     _write_main_kv_slots_to_pool,
 )
-from ..attention.backends.sparse.minimax_m3.common import index_head_range
+from ..attention.backends.sparse.minimax_m3.common import get_sparse_layer_ids, index_head_range
 from ..attention.backends.sparse.params import SparseBackendForwardArgs
 from ..attention.backends.trtllm import TrtllmAttention
 from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams, MiniMaxAllReduceRMS
@@ -431,26 +431,6 @@ def _validate_sparse_attention_runtime_config(
             )
 
 
-def get_sparse_layer_ids(text_config: PretrainedConfig) -> Tuple[List[int], List[int]]:
-    """Return ``(dense_layer_ids, sparse_layer_ids)`` for the M3 text model."""
-    sparse_cfg = getattr(text_config, "sparse_attention_config", None)
-    n_layers = int(text_config.num_hidden_layers)
-    if sparse_cfg is None:
-        return list(range(n_layers)), []
-    if not sparse_cfg.get("use_sparse_attention", True):
-        return list(range(n_layers)), []
-    freq = sparse_cfg.get("sparse_attention_freq")
-    if freq is None:
-        return list(range(n_layers)), []
-    if len(freq) != n_layers:
-        raise ValueError(
-            f"sparse_attention_freq length {len(freq)} does not match num_hidden_layers {n_layers}"
-        )
-    dense = [i for i, f in enumerate(freq) if int(f) == 0]
-    sparse = [i for i, f in enumerate(freq) if int(f) != 0]
-    return dense, sparse
-
-
 def get_sparse_disable_index_value_layer_ids(
     text_config: PretrainedConfig,
 ) -> List[int]:

tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py

@@ -33,6 +33,7 @@
 from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import (
     MiniMaxM3DraftKVCacheView,
     MiniMaxM3KVCacheManagerV2,
+    default_sparse_layer_ids,
     derive_shared_draft_layout,
     extend_attention_op_pools_for_shared_draft_layers,
     shared_draft_layer_count,
@@ -343,6 +344,88 @@ def test_shared_draft_layer_count_follows_the_base_manager():
     assert shared_draft_layer_count(None, None) == 0
 
 
+def _text_config(sparse_attention_freq):
+    return SimpleNamespace(
+        num_hidden_layers=len(sparse_attention_freq),
+        sparse_attention_config={"sparse_attention_freq": list(sparse_attention_freq)},
+    )
+
+
+def test_default_sparse_layers_follow_the_text_config():
+    # The model reads its dense/sparse split from the config; so does the cache.
+    assert default_sparse_layer_ids(_text_config([1, 0, 1, 1, 0]), 5) == [0, 2, 3]
+    assert default_sparse_layer_ids(_text_config([0, 0, 0, 1, 1]), 5) == [3, 4]
+    # Without a config spanning the target layers, layers 0..2 stay dense.
+    assert default_sparse_layer_ids(None, 5) == [3, 4]
+    assert default_sparse_layer_ids(_text_config([1, 0, 1, 1, 0]), 6) == [3, 4, 5]
+    assert default_sparse_layer_ids(_text_config([1, 0, 1, 1, 0]), None) == []
+
+
+def test_default_sparse_layers_log_a_fallback_that_ignores_the_config(monkeypatch):
+    import tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager as cache_manager
+
+    log = Mock()
+    monkeypatch.setattr(cache_manager, "logger", log)
+    # A config whose split is not the default, on a manager with a different layer count.
+    assert default_sparse_layer_ids(_text_config([1, 0, 1, 1, 0]), 6) == [3, 4, 5]
+    assert log.info.call_count == 1
+    # The default split, a matching layer count, or no config: nothing to report.
+    default_sparse_layer_ids(_text_config([0, 0, 0, 1, 1]), 6)
+    default_sparse_layer_ids(_text_config([1, 0, 1, 1, 0]), 5)
+    default_sparse_layer_ids(None, 6)
+    assert log.info.call_count == 1
+
+
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("implementation", ["msa", "triton"])
+def test_index_k_cache_follows_the_config_sparse_layers(
+    monkeypatch: pytest.MonkeyPatch, implementation: str
+) -> None:
+    """Every layer the config makes sparse gets an index-K view, and only those."""
+    resolved = []
+
+    def fake_base_init(self, *args, **kwargs):
+        self.is_disagg = False
+        self.dtype = DataType.BF16
+        # Five target layers and the shared Eagle3 layer appended after them.
+        self.layer_offsets = {layer: layer for layer in range(6)}
+
+    def fake_get_index_k_buffer(self, layer_idx, **kwargs):
+        resolved.append(layer_idx)
+        return torch.empty(0)
+
+    monkeypatch.setattr(KVCacheManagerV2, "__init__", fake_base_init)
+    monkeypatch.setattr(KVCacheManagerV2, "get_index_k_buffer", fake_get_index_k_buffer)
+    monkeypatch.setattr(MiniMaxM3KVCacheManagerV2, "_compute_num_total_slots", lambda self: 0)
+    monkeypatch.setattr(MiniMaxM3KVCacheManagerV2, "get_buffers", lambda self, *args: None)
+    mode = SimpleNamespace(
+        is_mtp_eagle_one_model=lambda: False,
+        is_mtp_vanilla=lambda: False,
+        is_eagle3_one_model=lambda: True,
+    )
+
+    manager = MiniMaxM3KVCacheManagerV2(
+        num_layers=5,
+        spec_config=SimpleNamespace(spec_dec_mode=mode, _num_draft_hidden_layers=None),
+        sparse_attention_config=SimpleNamespace(
+            implementation=implementation, indexer_kv_dtype="bf16"
+        ),
+        pretrained_config=_text_config([1, 0, 1, 1, 0]),
+    )
+
+    assert manager._shared_draft_layer_ids == [5]
+    assert manager.sparse_layer_ids == [0, 2, 3]
+    assert manager.disable_index_value_layer_ids == {0, 2, 3}
+    assert resolved == [0, 2, 3]
+    for layer in (0, 2, 3):
+        assert manager.get_index_k_buffer(layer) is not None
+    for layer in (1, 4, 5):
+        assert manager.get_index_k_buffer(layer) is None
+    # The Triton path runs the dense target layers on the native attention op.
+    expected_dense_native = (1, 4) if implementation == "triton" else ()
+    assert manager._dense_native_layer_ids == expected_dense_native
+
+
 @pytest.mark.cpu_only
 @pytest.mark.parametrize("local_draft_layers", [[], [60], [61]])
 @pytest.mark.parametrize("swa_scratch_reuse", [False, True])