diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 19099d0b1ac0..c5e39bf23afb 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -997,12 +997,71 @@ def __init__( mapping: Mapping, dtype: DataType = DataType.HALF, spec_config: Optional["DecodingBaseConfig"] = None, - layer_mask: Optional[List[bool]] = None, + layer_mask: Optional[ + List[bool]] = None, # this is the full attention layer mask is_estimating_kv_cache: bool = False, use_replay_state_update: bool = False, **kwargs, ) -> None: + # 3 kinds of layers: + # 1) Mamba layers (mamba_layer_mask is True) + # 2) Full attention layers (full_attention_layer_mask is True) + # 3) Not managed layers (both masks are False) + total_layers = len(mamba_layer_mask) + if layer_mask is None: + full_attention_layer_mask = [False] * total_layers + elif len(layer_mask) != total_layers: + raise ValueError( + f"layer_mask length ({len(layer_mask)}) must match " + f"mamba_layer_mask length ({total_layers})") + else: + full_attention_layer_mask = list(layer_mask) + layer_mask = [ + mamba_layer_mask[i] or full_attention_layer_mask[i] + for i in range(total_layers) + ] + # PP sharding is done across all layers. + # This is called again in the super().__init__, but we want it to run first + # to set up mtp states before the C++ backend is initialized. + self.pp_layers, _ = get_pp_layers( + mamba_num_layers + num_layers, + mapping, + spec_config=spec_config, + layer_mask=layer_mask, + ) + self.mamba_pp_layers = [ + layer_idx for layer_idx in self.pp_layers + if mamba_layer_mask[layer_idx] + ] + self.local_num_mamba_layers = len(self.mamba_pp_layers) + self.requests = [] + # Seed externally visible mamba fields before any early return so that + # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work + # on ranks with no local mamba layers. self._use_replay_state_update = use_replay_state_update + self.ssm_state_dtype = mamba_ssm_cache_dtype + + if self.local_num_mamba_layers == 0: + logger.info( + "No local mamba layers for this rank, skipping mamba cache initialization" + ) + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=full_attention_layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + ) + return + # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 d_inner = mamba_head_dim * mamba_num_heads @@ -1018,7 +1077,6 @@ def __init__( nheads = nheads // tp_size self.conv_state_shape = [conv_dim, mamba_d_conv - 1] self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] - self.ssm_state_dtype = mamba_ssm_cache_dtype self.conv_state_dtype = mamba_cache_dtype self.ssm_count = math.prod(self.ssm_state_shape) self.conv_count = math.prod(self.conv_state_shape) @@ -1049,18 +1107,7 @@ def __init__( ) kv_cache_config.enable_partial_reuse = False - full_attention_layer_mask = layer_mask.copy() - kv_cache_config.max_attention_window = [] - # 3 kinds of layers: - # 1) Mamba layers (mamba_layer_mask is True) - # 2) Full attention layers (full_attention_layer_mask is True) - # 3) Not managed layers (both masks are False) - total_layers = len(mamba_layer_mask) - layer_mask = [ - mamba_layer_mask[i] or full_attention_layer_mask[i] - for i in range(total_layers) - ] for i in range(len(layer_mask)): if layer_mask[i]: kv_cache_config.max_attention_window.append( @@ -1083,21 +1130,6 @@ def __init__( if is_mamba: per_layer_kv_heads[i] = 0 - # PP sharding is done across all layers. - # This is called again in the super().__init__, but we want it to run first - # to set up mtp states before the C++ backend is initialized. - self.pp_layers, _ = get_pp_layers( - mamba_num_layers + num_layers, - mapping, - spec_config=spec_config, - layer_mask=layer_mask, - ) - self.mamba_pp_layers = [ - layer_idx for layer_idx in self.pp_layers - if mamba_layer_mask[layer_idx] - ] - self.local_num_mamba_layers = len(self.mamba_pp_layers) - self._setup_mtp_intermediate_states(spec_config, max_batch_size) # pass remaining arguments to super class @@ -1128,7 +1160,6 @@ def __init__( ], dtype=torch.int32, device="cpu") - self.requests = [] self.recurrent_states_pool_index = self.kv_cache_pool_mapping[ self.layer_offsets[self.mamba_pp_layers[0]]][0] @@ -1287,6 +1318,8 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): def prepare_resources(self, scheduled_batch: ScheduledRequests): super().prepare_resources(scheduled_batch) + if self.local_num_mamba_layers == 0: + return self._prepare_resources(scheduled_batch) def is_speculative(self) -> bool: @@ -1296,6 +1329,8 @@ def update_mamba_states(self, attn_metadata: "AttentionMetadata", num_accepted_tokens: torch.Tensor, state_indices: Optional[torch.Tensor] = None): + if self.local_num_mamba_layers == 0: + return # Note: cannot use @torch.compile here because all_ssm_states and # all_conv_states are dtype-reinterpreted views of the C++ pool # (uint8 -> typed), and aot_autograd does not support mutations on @@ -1394,6 +1429,8 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): super().free_resources(request, pin_on_release) def _setup_state_indices(self) -> None: + if self.local_num_mamba_layers == 0: + return block_indices = [] for req in self.requests: if req.is_context_finished: diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 5ad4ce4b1748..d959cdfc2547 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Regression tests for MambaCacheManager padding-slot behavior.""" +"""Regression tests for MambaCacheManager padding-slot behavior and +CppMambaHybridCacheManager PP-sharding edge cases.""" from types import SimpleNamespace from unittest.mock import MagicMock @@ -11,8 +12,12 @@ from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( CppMambaCacheManager, + CppMambaHybridCacheManager, PythonMambaCacheManager, ) +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -251,3 +256,108 @@ def test_cpp_get_state_indices_resolves_sentinel_to_reserved_slot(): ) # Resolve again — reserved slot must be stable across calls. assert mgr.get_state_indices(request_ids, is_padding) == indices + + +# --------------------------------------------------------------------------- +# CppMambaHybridCacheManager: rank with zero local mamba layers +# +# Regression test for the early-exit path added when a rank ends up with no +# mamba layers (e.g. under PP sharding when all mamba layers fall on other +# ranks). On that path, the constructor must: +# - call the real parent KVCacheManager with the union layer_mask and +# num_layers=num_layers (not mamba_num_layers + num_layers), +# - skip allocating any mamba-only state, and +# - leave self.requests = [] so the guards on prepare_resources / +# update_mamba_states / _setup_state_indices can no-op without touching +# uninitialized state. +# +# We exercise the same Python branch with world_size=1 (so the real C++ +# KVCacheManager init doesn't need MPI) and a layer mask that contains zero +# mamba layers. +# --------------------------------------------------------------------------- + + +def _build_zero_mamba_hybrid(): + """Construct a real CppMambaHybridCacheManager whose this-rank slice has + no mamba layers. world_size=1 / pp_size=1 keeps the real parent + KVCacheManager off the MPI path.""" + # [other, other, full_attn, full_attn] + mamba_mask = [False, False, False, False] + attn_mask = [False, False, True, True] + mamba_num_layers = sum(mamba_mask) # 0 + num_layers = sum(attn_mask) # 4 + + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + # Cap KV pool size so the real C++ allocator only takes a tiny slice of + # GPU memory; we don't actually use the cache. + kv_cache_config = KvCacheConfig(max_tokens=128) + + mgr = CppMambaHybridCacheManager( + # mamba cache parameters — values are unused on the early-exit path + # but must be type-valid. + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=mamba_num_layers, + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.float16, + mamba_ssm_cache_dtype=torch.float16, + # kv cache parameters + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELF, + num_layers=num_layers, + num_kv_heads=4, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=2, + mapping=mapping, + spec_config=None, + layer_mask=attn_mask, + ) + return mgr + + +@skip_no_cuda +def test_cpp_hybrid_zero_local_mamba_layers(): + """End-to-end: real parent KVCacheManager + real early-exit. Verifies + early-exit invariants on the manager state AND that the three guarded + methods no-op without raising on uninitialized mamba-only state.""" + mgr = _build_zero_mamba_hybrid() + + # Early-exit indicators. + assert mgr.local_num_mamba_layers == 0 + assert mgr.mamba_pp_layers == [] + assert mgr.requests == [] + assert mgr.pp_layers == [2, 3] + + # Parent KVCacheManager was really initialized. self.impl is the C++ + # KVCacheManagerCpp object; blocks_per_window is set up by it. + assert hasattr(mgr, "impl") + assert hasattr(mgr, "blocks_per_window") + # Parent saw num_layers = num_layers (4), not mamba_num_layers + num_layers. + # On the early-exit branch, num_layers is forwarded as-is. + assert mgr.num_layers == 4 + assert mgr.num_local_layers == 2 + + # No mamba-only state was allocated. + for attr in ( + "ssm_state_shape", + "conv_state_shape", + "mamba_layer_offsets", + "cuda_state_indices", + "host_block_offsets", + "recurrent_states_pool_index", + ): + assert not hasattr(mgr, attr), f"{attr} must not be set on the zero-mamba early-exit path" + # Parent must not have been told to treat this as linear attention. + assert mgr.is_linear_attention is False + + # Guards on the three mamba-only methods must turn them into no-ops + # instead of crashing on the missing state above. + empty_batch = ScheduledRequests() + mgr.prepare_resources(empty_batch) # super() runs, then guard returns + mgr.update_mamba_states(attn_metadata=None, num_accepted_tokens=None, state_indices=None) + mgr._setup_state_indices()