[None][fix] CppMambaHybridCacheManager: handle ranks with zero local mamba layers

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:

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()