From 75710b1ac14db44352fbc50bcc7612af09e4eb86 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Thu, 30 Jul 2026 07:41:45 -0700 Subject: [PATCH] [nvbugs/6525898][fix] Constrain V2 Mamba SSM pool to its live slot floor MambaHybridCacheManagerV2.__init__ requires the SSM pool to hold one slot per resident lineage plus every reserved dummy slot. When avg_seq_len is unset, _build_base_config emits no constraints, so the storage manager received no min_slots floor for the SSM pool group; only the fallback typical_step steered the pool ratio. A ratio is not a floor, so the grain-based split rounded the SSM pool below the required count and initialization failed with 'The V2 Mamba state pool has only 31 slots but needs at least 34 live/dummy slots'. Emit the live/dummy slot count as an explicit constraint. The descriptors carry no capacity, so the attention pool floor is unchanged. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../kv_cache/mamba_cache_manager.py | 26 +++++----- .../kv_cache/test_mamba_cache_manager.py | 48 +++++++++++++++++++ 2 files changed, 59 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 7287940d5eef..4fb9f55c6bd6 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3240,8 +3240,7 @@ def __init__( LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // self._ssm_page_index_scale) - required_live_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) + required_live_slots = self._num_required_state_slots() if num_ssm_slots < required_live_slots: KVCacheManagerV2.shutdown(self) raise ValueError( @@ -3785,6 +3784,11 @@ def _get_pool_roles(self, def _max_resident_sequences(self) -> int: return self.max_batch_size * self.mapping.pp_size + def _num_required_state_slots(self) -> int: + """Return the SSM slots that must always be live: one per resident + request lineage plus every reserved dummy slot.""" + return self._max_resident_sequences() + self._num_reserved_dummy_slots + def _mamba_state_bytes_per_slot(self) -> int: base_bytes = self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes) @@ -3924,9 +3928,8 @@ def _minimum_live_gpu_quota(self) -> int: """Return the minimum quota for live states and one attention page.""" attention_block_quota = (self._attention_cache_bytes_per_token() * self.tokens_per_block) - num_state_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - state_quota = num_state_slots * self._mamba_state_bytes_per_slot() + state_quota = (self._num_required_state_slots() * + self._mamba_state_bytes_per_slot()) return max( self._get_quota_from_max_tokens(0), state_quota + attention_block_quota, @@ -3962,10 +3965,8 @@ def _build_cache_config( buffers=buffers, ) - dummy_requests = [ - KVCacheDesc(capacity=0, history_length=0) - for _ in range(self._num_reserved_dummy_slots) - ] + empty_desc = KVCacheDesc(capacity=0, history_length=0) + dummy_requests = [empty_desc] * self._num_reserved_dummy_slots constraints = [ replace( batch, @@ -3992,14 +3993,9 @@ def _build_cache_config( # / __init__). Add a min-slots constraint of zero-capacity requests: # these cost no attention pages but reserve one SSM slot each. if any(isinstance(layer, SsmLayerConfig) for layer in layers): - ssm_floor_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) constraints = [ *constraints, - BatchDesc([ - KVCacheDesc(capacity=0, history_length=0) - for _ in range(ssm_floor_slots) - ]), + BatchDesc([empty_desc] * self._num_required_state_slots()), ] return replace( config, diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 9174796cd075..808a27a405bb 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2121,6 +2121,54 @@ def _slot_sizes(stat): assert high_mamba_allocation[1] < low_mamba_allocation[1] +def test_v2_hybrid_constrains_ssm_pool_to_live_slot_floor(): + """The SSM pool floor must be a constraint, not just a typical_step ratio. + + With avg_seq_len unset the base config emits no constraints, so only the + fallback typical_step steers the pool ratio. A ratio is not a floor, so the + grain split could round the SSM pool below the live/dummy slot count that + __init__ validates (https://nvbugs/6525898). Assert the floor reaches the + storage manager as a constraint carrying no attention capacity. + """ + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.kv_cache_type = CacheTypeCpp.SELF + mgr.head_dim_per_layer = [64, 64] + mgr.pp_layers = [0, 1] + mgr._mamba_layer_mask = [True, False] + mgr.ssm_bytes = 64 + mgr.conv_bytes = 32 + mgr.max_attention_window_vec = [128, 128] + mgr.max_batch_size = 32 + mgr.mapping = Mapping(world_size=2, rank=0, tp_size=1, pp_size=2) + mgr.max_seq_len = 128 + mgr.max_num_tokens = 128 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.local_num_mamba_layers = 1 + mgr._num_reserved_dummy_slots = 2 + mgr.dtype = DataType.HALF + mgr.enable_swa_scratch_reuse = False + mgr.enable_stats = False + mgr.num_extra_kv_tokens = 0 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + mgr._minimum_live_gpu_quota = lambda: 0 + mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) + + base_config = mgr._build_base_config( + mgr.kv_cache_config, + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + ) + # avg_seq_len is unset, so nothing pins the pool sizes yet. + assert base_config.typical_step is None + assert not base_config.constraints + + config = mgr._build_cache_config(base_config) + + # 32 max_batch_size * pp_size 2 resident lineages + 2 reserved dummy slots. + assert config.constraints == [BatchDesc([KVCacheDesc(capacity=0, history_length=0)] * 66)] + + # --------------------------------------------------------------------------- # Cpp/V2 Mamba hybrid managers: recurrent-state allocation and reuse #