[https://nvbugs/6550276][fix] Clamp residency to what the quota affords (allreduce-MIN across ranks) behind…

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -2156,6 +2156,16 @@ def is_request_active(self, request_id: int) -> bool:
         kv_cache = self.kv_cache_map.get(request_id)
         return kv_cache is not None and kv_cache.is_active
 
+    def max_resident_sequences(self) -> Optional[int]:
+        """Cap on concurrently resident sequences, or ``None`` if unbounded.
+
+        Attention pages are droppable, so pure-attention models let
+        suspend/resume absorb over-admission and need no cap. Managers that own
+        non-droppable per-sequence state (e.g. Mamba recurrent state) override
+        this so the scheduler stops admitting sequences the pool cannot hold.
+        """
+        return None
+
     def _effective_draft_len(self, req: LlmRequest) -> int:
         """Draft token length to use for next-step KV capacity calculation.
 

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -28,6 +28,7 @@
     from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
 
+from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import (
     BlockReusePolicy, KVCacheManagerV2, Role)
 from tensorrt_llm._torch.pyexecutor.llm_request import (
@@ -2477,6 +2478,11 @@ class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager):
 
     _supports_additional_snapshot_offsets = True
 
+    # Bound on concurrently resident sequences, set in _build_cache_config when
+    # the GPU cache quota cannot hold the requested max_batch_size. ``None``
+    # means no bound applies and max_batch_size is used as requested.
+    _resident_sequence_cap: Optional[int] = None
+
     def __init__(
         self,
         # mamba cache parameters
@@ -2740,12 +2746,62 @@ def _get_pool_roles(self,
             return MambaRole.SSM_STATE, None
         return super()._get_pool_roles(pool_id)
 
-    def _max_resident_sequences(self) -> int:
+    def _requested_resident_sequences(self) -> int:
+        """Sequences the configured ``max_batch_size`` asks to keep resident."""
         return self.max_batch_size * self.mapping.pp_size
 
+    def _max_resident_sequences(self) -> int:
+        if self._resident_sequence_cap is not None:
+            return self._resident_sequence_cap
+        return self._requested_resident_sequences()
+
+    def max_resident_sequences(self) -> Optional[int]:
+        """Number of sequences whose recurrent state can be resident at once."""
+        if self.local_num_mamba_layers == 0:
+            return None
+        return self._max_resident_sequences()
+
+    def _resident_sequences_for_quota(self, gpu_quota: int,
+                                      typical_descs: List[KVCacheDesc]) -> int:
+        """Return how many sequences the GPU quota can keep resident.
+
+        A recurrent state is fixed-size per sequence and cannot be recomputed
+        from tokens, so every resident sequence permanently occupies one state
+        slot. ``max_batch_size`` is a ceiling, not an allocation contract, so a
+        request for more sequences than the quota can hold is bounded here
+        instead of becoming a hard allocation floor.
+
+        The per-sequence cost reuses ``typical_step``'s request model
+        (``typical_descs``: one descriptor per live state, together spanning the
+        typical request capacity), so sizing here cannot drift from the pool
+        sizing it bounds. Charging the state cost alone would admit sequences
+        whose states consume the whole quota and leave nothing for attention.
+        """
+        state_bytes = self._mamba_state_bytes_per_slot()
+        if state_bytes <= 0:
+            # Attention-only rank: nothing non-droppable to bound. Return the
+            # request itself so this rank stays neutral in the allreduce(MIN).
+            return self._requested_resident_sequences()
+
+        attention_block_bytes = self._attention_block_bytes()
+        typical_capacity = sum(desc.capacity for desc in typical_descs)
+        attention_blocks = math.ceil(typical_capacity / self.tokens_per_block)
+        per_sequence = (len(typical_descs) * state_bytes +
+                        attention_blocks * attention_block_bytes)
+
+        # Reserved dummy slots and one attention page are charged up front:
+        # they exist regardless of how many real sequences are resident.
+        reserved = (self._num_reserved_dummy_slots * state_bytes +
+                    attention_block_bytes)
+        return (gpu_quota - reserved) // per_sequence
+
     def _mamba_state_bytes_per_slot(self) -> int:
         return self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes)
 
+    def _attention_block_bytes(self) -> int:
+        """Bytes of one attention page across all local attention layers."""
+        return self._attention_cache_bytes_per_token() * self.tokens_per_block
+
     def _num_ssm_snapshots_for_capacity(
         self,
         capacity: int,
@@ -2827,8 +2883,7 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int:
         # attention page per request lineage. This remains conservative when
         # the plan contains fewer than one non-live slot per lineage.
         extra_attention_quota = (num_request_lineages *
-                                 self._attention_cache_bytes_per_token() *
-                                 self.tokens_per_block
+                                 self._attention_block_bytes()
                                  if snapshot_slots > 0 else 0)
         return attention_quota + state_quota + extra_attention_quota
 
@@ -2854,21 +2909,44 @@ def _get_max_tokens_from_quota(self, quota: int) -> float:
 
     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()
         return max(
             self._get_quota_from_max_tokens(0),
-            state_quota + attention_block_quota,
+            state_quota + self._attention_block_bytes(),
         )
 
     def _build_cache_config(
             self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy:
         kv_cache_config = self.kv_cache_config
         cache_tiers = config.cache_tiers
         gpu_quota = cache_tiers[0].quota
+        # One descriptor per live SSM state of a typical request. Shared with
+        # the typical_step sizing below so both agree on the per-sequence cost
+        # (and so the avg_seq_len fallback warning is emitted only once).
+        request_descs = self._typical_request_descs(
+            self._get_typical_request_capacity(kv_cache_config),
+            kv_cache_config)
+        requested_resident = self._requested_resident_sequences()
+        affordable_resident = self._resident_sequences_for_quota(
+            gpu_quota, request_descs)
+        if self.mapping.world_size > 1:
+            # Quotas and per-rank state costs differ across ranks (uneven mamba
+            # layer splits, attention-only PP ranks). Every rank must apply the
+            # same bound or the schedulers admit different batches and desync.
+            affordable_resident = Distributed.get(self.mapping).allreduce(
+                affordable_resident, op=ReduceOp.MIN)
+        if affordable_resident < requested_resident:
+            self._resident_sequence_cap = max(1, affordable_resident)
+            logger.warning(
+                f"The V2 Mamba GPU cache quota ({gpu_quota} bytes) cannot keep "
+                f"{requested_resident} sequences resident: each one holds a "
+                f"fixed {self._mamba_state_bytes_per_slot()} bytes of recurrent "
+                "state that cannot be evicted. Limiting concurrently resident "
+                f"sequences to {self._resident_sequence_cap}. Reduce "
+                "max_batch_size, or raise free_gpu_memory_fraction / "
+                "max_gpu_total_bytes, to run the requested batch size.")
         minimum_live_quota = self._minimum_live_gpu_quota()
         if minimum_live_quota > gpu_quota:
             raise ValueError(
@@ -2892,25 +2970,26 @@ def _build_cache_config(
                     ],
                 )
 
+        max_resident = self._max_resident_sequences()
         dummy_requests = [
             KVCacheDesc(capacity=0, history_length=0)
             for _ in range(self._num_reserved_dummy_slots)
         ]
+        # The base class sizes its warmup constraints from the requested
+        # max_batch_size. Every entry of a constraint costs one SSM slot (the
+        # planner never shares recurrent state between requests), so a
+        # constraint wider than the clamped residency would restore the very
+        # floor the clamp removed. Truncate to what can actually be resident.
         constraints = [
             replace(
                 batch,
-                kv_caches=[*batch.kv_caches, *dummy_requests],
+                kv_caches=[*batch.kv_caches[:max_resident], *dummy_requests],
             ) for batch in config.constraints
         ]
 
         typical_step = config.typical_step
         if config.initial_pool_ratio is None:
-            typical_capacity = self._get_typical_request_capacity(
-                kv_cache_config)
-            request_descs = self._typical_request_descs(typical_capacity,
-                                                        kv_cache_config)
-            typical_step = BatchDesc(request_descs *
-                                     self._max_resident_sequences() +
+            typical_step = BatchDesc(request_descs * max_resident +
                                      dummy_requests)
         # The recurrent (SSM) state pool must hold one slot per resident
         # sequence plus every reserved dummy slot. Unlike attention pages, a
@@ -2922,8 +3001,7 @@ 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)
+            ssm_floor_slots = max_resident + self._num_reserved_dummy_slots
             constraints = [
                 *constraints,
                 BatchDesc([

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py

@@ -176,6 +176,11 @@ def __init__(
             scheduler_policy = CapacitySchedulerPolicy.MAX_UTILIZATION
         self.policy = scheduler_policy
         self.peft_cache_manager = peft_cache_manager
+        # Non-droppable per-sequence state (Mamba recurrent state) yields no
+        # evictable pages, so MAX_UTILIZATION's suspend/resume cannot recover
+        # from over-admission: once every resident sequence is suspended the
+        # pool can never drain. Bound admission instead.
+        self.max_resident_sequences = kv_cache_manager.max_resident_sequences()
 
         # Chunking config.
         self.chunking_enabled = False
@@ -193,7 +198,8 @@ def __init__(
             f"KVCacheV2Scheduler: tokens_per_block={self.tokens_per_block}, "
             f"max_num_tokens={max_num_tokens}, max_batch_size={max_batch_size}, "
             f"draft_mgr={draft_mgr_name}, cross_mgr={cross_mgr_name}, "
-            f"enable_prefix_aware_scheduling={enable_prefix_aware_scheduling}"
+            f"enable_prefix_aware_scheduling={enable_prefix_aware_scheduling}, "
+            f"max_resident_sequences={self.max_resident_sequences}"
         )
         if ctx_chunk_config is not None:
             self.chunking_enabled = True
@@ -309,6 +315,16 @@ def _schedule_loop(self, active_requests, inflight_request_ids):
             if req.state_value == self._gen_to_complete_state_value:
                 budget.pre_claim_peft(req)
 
+        # Sequences already holding a non-droppable state slot. Counted over all
+        # active requests (not just the ones scheduled this iteration) because a
+        # suspended sequence keeps its slot.
+        max_resident = self.max_resident_sequences
+        num_resident = (
+            sum(1 for req in requests_list if self._is_started_request(req))
+            if max_resident is not None
+            else 0
+        )
+
         # --- Phase 1: generation / disagg only ---
         while req_it < req_it_end:
             req = requests_list[req_it]
@@ -403,6 +419,11 @@ def _schedule_loop(self, active_requests, inflight_request_ids):
         for req in pending_ctx:
             if budget.requests_full:
                 break
+            # A first context chunk starts a new sequence and therefore claims a
+            # state slot for the rest of its lifetime.
+            starts_new_sequence = max_resident is not None and req.is_first_context_chunk
+            if starts_new_sequence and num_resident >= max_resident:
+                break
             peft_pages = budget.peft_pages_needed(req)
             if peft_pages is None:
                 continue
@@ -421,6 +442,8 @@ def _schedule_loop(self, active_requests, inflight_request_ids):
                 has_chunking = has_chunking or chunking_flag
                 scheduled_ctx.append(req)
                 budget.commit(req, tokens, peft_pages)
+                if starts_new_sequence:
+                    num_resident += 1
 
         # Deadlock detection: if generation requests exist but none were
         # scheduled and none were evicted, no forward pass will run and no

tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py

@@ -160,6 +160,7 @@ def make_kv_cache_manager(
     resize_context_fn=None,
     prepare_disagg_gen_init_fn=None,
     try_allocate_generation_fn=None,
+    max_resident_sequences=None,
 ):
     mgr = Mock()
     mgr.tokens_per_block = tokens_per_block
@@ -170,6 +171,9 @@ def make_kv_cache_manager(
     mgr.try_allocate_generation.side_effect = try_allocate_generation_fn or (lambda req: True)
     mgr.suspend_request.return_value = None
     mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active
+    # Pin explicitly: a bare Mock() would auto-vivify a truthy Mock here, which
+    # the residency gate compares against an int.
+    mgr.max_resident_sequences.return_value = max_resident_sequences
     return mgr
 
 
@@ -2143,6 +2147,59 @@ def track_resize(req: Mock, num_tokens: int) -> bool:
         assert ids(out.context_requests) == []
 
 
+# ===========================================================================
+# Residency cap (non-droppable per-sequence state, e.g. Mamba)
+# ===========================================================================
+
+
+class TestResidencyCap:
+    """A manager owning non-droppable per-sequence state bounds admission.
+
+    Mamba recurrent state yields no evictable pages, so suspend/resume cannot
+    recover from over-admission (nvbugs/6550276).
+    """
+
+    def test_new_sequences_capped_at_max_resident(self):
+        mgr = make_kv_cache_manager(max_resident_sequences=2)
+        sched = make_scheduler(mgr, max_num_tokens=100000, max_batch_size=100)
+        reqs = [make_ctx_request(i, context_remaining_length=10) for i in range(5)]
+
+        out = sched.schedule_request(reqs, set())
+
+        assert ids(out.context_requests) == [0, 1]
+
+    def test_started_sequences_count_against_the_cap(self):
+        """A resident sequence keeps its state slot, so it consumes the cap."""
+        mgr = make_kv_cache_manager(max_resident_sequences=2)
+        sched = make_scheduler(mgr, max_num_tokens=100000, max_batch_size=100)
+        gens = [make_gen_request(i) for i in range(2)]
+        ctxs = [make_ctx_request(10 + i, context_remaining_length=10) for i in range(2)]
+
+        out = sched.schedule_request(gens + ctxs, set())
+
+        assert ids(out.generation_requests) == [0, 1]
+        assert ids(out.context_requests) == []
+
+    def test_non_first_chunk_is_not_charged_again(self):
+        """Only the first chunk starts a sequence; later chunks already hold a slot."""
+        mgr = make_kv_cache_manager(max_resident_sequences=1)
+        sched = make_scheduler(mgr, max_num_tokens=100000, max_batch_size=100)
+        req = make_ctx_request(0, context_remaining_length=10, is_first_context_chunk=False)
+
+        out = sched.schedule_request([req], set())
+
+        assert ids(out.context_requests) == [0]
+
+    def test_unbounded_manager_admits_every_request(self):
+        mgr = make_kv_cache_manager(max_resident_sequences=None)
+        sched = make_scheduler(mgr, max_num_tokens=100000, max_batch_size=100)
+        reqs = [make_ctx_request(i, context_remaining_length=10) for i in range(5)]
+
+        out = sched.schedule_request(reqs, set())
+
+        assert ids(out.context_requests) == [0, 1, 2, 3, 4]
+
+
 # ===========================================================================
 # Edge Cases
 # ===========================================================================