From 7c7560ecbee8e1378235eb3ab9673afc8b841ab2 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:58:14 -0700 Subject: [PATCH 1/3] [None][fix] Derive MiniMax-M3 index-K cache layers from the model config MiniMaxM3KVCacheManagerV2 allocated index-K caches for layers 3..N-1 whenever the caller passed no sparse_layer_ids, while the model selects sparse attention per layer from sparse_attention_config.sparse_attention_freq. For a config whose sparse layers differ from that fixed split, a sparse layer found no index-K cache (get_index_k_buffer() returned None) and dense layers were laid out as sparse ones. Build the default from the pretrained_config keyword the executor already passes to the manager, with the same get_sparse_layer_ids() the model uses; the helper moves to minimax_m3/common.py so the cache manager does not import the model module. The fixed split stays the fallback when no config spanning the target layers is given, so configs that match it see no change. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- .../sparse/minimax_m3/cache_manager.py | 43 +++++++++--- .../backends/sparse/minimax_m3/common.py | 28 +++++++- .../_torch/models/modeling_minimaxm3.py | 22 +----- .../test_minimax_m3_shared_draft_layers.py | 68 +++++++++++++++++++ 4 files changed, 129 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 5b450c3742c9..6aa2c6d3aa6d 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -54,7 +54,11 @@ PageIndexMode, ) +from .common import get_sparse_layer_ids + if TYPE_CHECKING: + from transformers import PretrainedConfig + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig @@ -185,6 +189,25 @@ 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] + return list(range(3, num_target_layers)) + + def extend_attention_op_pools_for_shared_draft_layers( pool_pointers: torch.Tensor, pool_mapping: torch.Tensor, @@ -265,7 +288,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. @@ -303,11 +327,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") @@ -334,10 +359,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) @@ -1171,6 +1195,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", diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py index f288495eeff1..b527fbbdced7 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py +++ b/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", ] diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 311600ab67ca..2352f2dc7dd8 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -55,7 +55,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 @@ -439,26 +439,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]: diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index 19f9bfdba7ab..250b8690fd14 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/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, @@ -385,6 +386,73 @@ 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) == [] + + +@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]) From 4ac332980849cb4772df1151e444a401c93925a3 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Sat, 24 Oct 2026 18:16:08 -0700 Subject: [PATCH 2/3] [None][fix] Log when the MiniMax-M3 index-K default ignores the config split Review follow-up. When a manager's target layer count differs from the model config's num_hidden_layers, default_sparse_layer_ids keeps the range(3, N) fallback. That is right for a separate draft manager, but if a target manager is ever built with a different layer count, a checkpoint with a non-default split would quietly get the legacy layout. Log one info line in that case, when the config's own split is not the default. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- .../backends/sparse/minimax_m3/cache_manager.py | 14 ++++++++++++++ .../msa/test_minimax_m3_shared_draft_layers.py | 15 +++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 6aa2c6d3aa6d..c10b486f39d7 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -205,6 +205,20 @@ def default_sparse_layer_ids( 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)) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index 250b8690fd14..7265d656cdc1 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -403,6 +403,21 @@ def test_default_sparse_layers_follow_the_text_config(): 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( From bb58f422f21e00202a3f77f24cd30490dbf853ed Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:20:09 -0700 Subject: [PATCH 3/3] [None][fix] Assert the MiniMax-M3 index-K pool buffers follow the config split test_index_k_cache_follows_the_config_sparse_layers stubs the base KVCacheManagerV2.__init__, so _build_base_config never ran and the test checked only the resolved index-K views, not the Role.INDEX_KEY buffers the V2 pool registers. Call _extra_buffers_per_layer from the stubbed base __init__, where _build_base_config calls it, and assert it registers Role.INDEX_KEY for the config's sparse layers 0, 2 and 3 only, not for the dense layers 1 and 4 or the shared Eagle3 layer 5. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../sparse/msa/test_minimax_m3_shared_draft_layers.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index 7265d656cdc1..b479d6bc0faf 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -423,14 +423,17 @@ def test_default_sparse_layers_log_a_fallback_that_ignores_the_config(monkeypatc 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.""" + """Every layer the config makes sparse gets an index-K buffer and view, and only those.""" resolved = [] + registered = {} 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)} + # The base _build_base_config registers the extra pool buffers here. + registered.update(self._extra_buffers_per_layer(tokens_per_block=32)) def fake_get_index_k_buffer(self, layer_idx, **kwargs): resolved.append(layer_idx) @@ -458,6 +461,9 @@ def fake_get_index_k_buffer(self, layer_idx, **kwargs): 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 sorted(registered) == [0, 2, 3] + for buffers in registered.values(): + assert [buffer.role for buffer in buffers] == [Role.INDEX_KEY] assert resolved == [0, 2, 3] for layer in (0, 2, 3): assert manager.get_index_k_buffer(layer) is not None