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 37d9564664a2..b6fbb0729246 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 @@ -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", 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 97c2f91805f4..f2aff89ec673 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/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]: 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 e9abf2039073..7e9838625985 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, @@ -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])