[https://nvbugs/6550275][fix] Bound V2 residency for non-droppable state pools

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -1121,9 +1121,7 @@ def append_to_kv_heads_per_layer(
         # Pad max_blocks_per_seq to next multiple of 4 (copy_block_offsets kernel).
         # Account for max single-sequence capacity = seq_len + extra KV tokens +
         # _kv_reserve_draft_tokens (see __init__) + 1 base decode token.
-        max_seq_capacity = (
-            self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1
-        )
+        max_seq_capacity = self._max_sequence_capacity()
         self.max_blocks_per_seq = (max_seq_capacity + tokens_per_block - 1) // tokens_per_block
         if self.max_blocks_per_seq % 4 != 0:
             self.max_blocks_per_seq = ((self.max_blocks_per_seq + 3) // 4) * 4
@@ -2151,6 +2149,45 @@ 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_sequence_capacity(self) -> int:
+        """Return the largest capacity passed to the V2 cache."""
+        return self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1
+
+    def max_resident_sequences(self) -> Optional[int]:
+        """Upper bound on sequences that can co-reside at ``max_seq_len``.
+
+        MAX_UTILIZATION admits new sequences up to ``max_batch_size`` and
+        relies on suspend/resume to survive over-subscription. That recovery
+        only works while some resident sequence can still be evicted to free
+        the pages a suspended one needs to resume. A sequence whose state is
+        non-droppable (a hybrid Mamba recurrent state is fixed-size per
+        sequence and cannot be recomputed from tokens) contributes no
+        evictable pages, so once every sequence is suspended the pool can no
+        longer drain and the scheduler makes no further progress.
+
+        Returns None when no such non-droppable pool exists, keeping the
+        unbounded MAX_UTILIZATION behavior for plain attention models.
+        """
+        life_cycle_metadata = self._stats_life_cycle_metadata()
+        if not any(kind != "attention" for _, _, kind in life_cycle_metadata.values()):
+            return None
+
+        # A physical group can host multiple lifecycle variants. V2 allocates
+        # each variant separately from the shared group, so their costs add.
+        attention_slots = math.ceil(self._max_sequence_capacity() / self.tokens_per_block)
+        slots_per_pool_group: dict[int, int] = defaultdict(int)
+        for pool_group_id, _, kind in life_cycle_metadata.values():
+            slots_per_pool_group[pool_group_id] += attention_slots if kind == "attention" else 1
+
+        storage_stats = self._get_storage_statistics(GPU_LEVEL)
+        return max(
+            1,
+            min(
+                storage_stats[pool_group_id].total // slots_per_sequence
+                for pool_group_id, slots_per_sequence in slots_per_pool_group.items()
+            ),
+        )
+
     def _effective_draft_len(self, req: LlmRequest) -> int:
         """Draft token length to use for next-step KV capacity calculation.
 

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py

@@ -183,6 +183,9 @@ def __init__(
         self.chunk_unit_size = 0
         self.max_context_length = max_num_tokens
         self.tokens_per_block = kv_cache_manager.tokens_per_block
+        # Cap on concurrently-started sequences, None when unbounded. See
+        # KVCacheManagerV2.max_resident_sequences() for why it is needed.
+        self.max_resident_sequences = kv_cache_manager.max_resident_sequences()
         draft_mgr_name = (
             type(draft_kv_cache_manager).__name__ if draft_kv_cache_manager is not None else "None"
         )
@@ -192,6 +195,7 @@ def __init__(
         logger.info(
             f"KVCacheV2Scheduler: tokens_per_block={self.tokens_per_block}, "
             f"max_num_tokens={max_num_tokens}, max_batch_size={max_batch_size}, "
+            f"max_resident_sequences={self.max_resident_sequences}, "
             f"draft_mgr={draft_mgr_name}, cross_mgr={cross_mgr_name}, "
             f"enable_prefix_aware_scheduling={enable_prefix_aware_scheduling}"
         )
@@ -210,6 +214,12 @@ def __init__(
         self._context_init_state_value = LlmRequestState.CONTEXT_INIT.value
         self._encoder_init_state_value = LlmRequestState.ENCODER_INIT.value
         self._disagg_gen_init_state_value = LlmRequestState.DISAGG_GENERATION_INIT.value
+        self._disagg_gen_trans_in_progress_state_value = (
+            LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS.value
+        )
+        self._disagg_gen_trans_complete_state_value = (
+            LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE.value
+        )
         self._gen_to_complete_state_value = LlmRequestState.GENERATION_TO_COMPLETE.value
 
         # Opt-in (default off): on the disagg generation server, schedule
@@ -286,6 +296,13 @@ def _schedule_loop(self, active_requests, inflight_request_ids):
                 )
             )
 
+        residency_cap = self.max_resident_sequences
+        resident_request_ids = (
+            {req.request_id for req in requests_list if self._is_resident_request(req)}
+            if residency_cap is not None
+            else set()
+        )
+
         req_it_end = len(requests_list)
         req_it = 0
 
@@ -332,6 +349,12 @@ def _schedule_loop(self, active_requests, inflight_request_ids):
             # no free slots remain, so the request is skipped and retried next
             # iteration. PEFT budget is still checked and committed.
             if req_state_value == self._disagg_gen_init_state_value:
+                needs_residency = (
+                    residency_cap is not None and req.request_id not in resident_request_ids
+                )
+                if needs_residency and len(resident_request_ids) >= residency_cap:
+                    req_it += 1
+                    continue
                 peft_pages = budget.peft_pages_needed(req)
                 if peft_pages is None:
                     break
@@ -343,6 +366,8 @@ def _schedule_loop(self, active_requests, inflight_request_ids):
                     req_it += 1
                     continue
                 disagg_candidates.append(req)
+                if needs_residency:
+                    resident_request_ids.add(req.request_id)
                 # Disagg requests only commit PEFT (not num_requests/num_tokens)
                 # because they don't participate in the forward pass. Counting
                 # them toward num_requests would steal batch slots from gen/ctx
@@ -400,9 +425,20 @@ def _schedule_loop(self, active_requests, inflight_request_ids):
 
         # --- Phase 2: schedule deferred context / encoder requests ---
         # Generation PEFT pages are now fully committed in the budget.
+        #
+        # Starting a new sequence is what grows the resident set, so the
+        # residency cap is enforced here. Requests already started keep
+        # their slot; the rest wait in the queue until one drains.
         for req in pending_ctx:
             if budget.requests_full:
                 break
+            needs_residency = (
+                residency_cap is not None
+                and req.state_value == self._context_init_state_value
+                and req.request_id not in resident_request_ids
+            )
+            if needs_residency and len(resident_request_ids) >= residency_cap:
+                continue
             peft_pages = budget.peft_pages_needed(req)
             if peft_pages is None:
                 continue
@@ -418,6 +454,8 @@ def _schedule_loop(self, active_requests, inflight_request_ids):
                     break
                 if action is ScheduleAction.SKIP:
                     continue
+                if needs_residency:
+                    resident_request_ids.add(req.request_id)
                 has_chunking = has_chunking or chunking_flag
                 scheduled_ctx.append(req)
                 budget.commit(req, tokens, peft_pages)
@@ -985,6 +1023,17 @@ def _is_started_request(req: LlmRequest) -> bool:
             req.is_context_init_state and not req.is_first_context_chunk
         ) or req.is_generation_in_progress_state
 
+    def _is_resident_request(self, req: LlmRequest) -> bool:
+        """Return whether a request already owns or reserves V2 residency."""
+        if self._is_started_request(req):
+            return True
+        if req.state_value in (
+            self._disagg_gen_trans_in_progress_state_value,
+            self._disagg_gen_trans_complete_state_value,
+        ):
+            return True
+        return self.kv_cache_manager.is_request_active(req.py_request_id)
+
     def _suspend_request(self, req: LlmRequest) -> None:
         """Suspend a request's KV cache in both main and draft managers.
 

tests/unittest/_torch/executor/test_kv_cache_manager_v2.py

@@ -40,6 +40,69 @@
 MAX_SEQ_LEN = 16
 
 
+def _make_residency_manager(
+    *,
+    max_seq_len: int,
+    tokens_per_block: int,
+    num_extra_kv_tokens: int,
+    reserve_draft_tokens: int,
+    life_cycle_metadata: dict[int, tuple[int, int | None, str]],
+    pool_group_totals: list[int],
+) -> KVCacheManagerV2:
+    manager = object.__new__(KVCacheManagerV2)
+    manager.max_seq_len = max_seq_len
+    manager.tokens_per_block = tokens_per_block
+    manager.num_extra_kv_tokens = num_extra_kv_tokens
+    manager._kv_reserve_draft_tokens = reserve_draft_tokens
+    manager._stats_life_cycle_metadata = lambda: life_cycle_metadata
+    manager._get_storage_statistics = lambda _level: [
+        SimpleNamespace(total=total) for total in pool_group_totals
+    ]
+    return manager
+
+
+@pytest.mark.parametrize(
+    ("num_extra_kv_tokens", "reserve_draft_tokens", "expected"),
+    [(0, 0, 6), (1, 0, 4), (0, 32, 4)],
+)
+def test_max_resident_sequences_uses_full_rounded_capacity(
+    num_extra_kv_tokens: int,
+    reserve_draft_tokens: int,
+    expected: int,
+) -> None:
+    manager = _make_residency_manager(
+        max_seq_len=63,
+        tokens_per_block=32,
+        num_extra_kv_tokens=num_extra_kv_tokens,
+        reserve_draft_tokens=reserve_draft_tokens,
+        life_cycle_metadata={
+            0: (0, 63, "attention"),
+            1: (1, None, "ssm"),
+        },
+        pool_group_totals=[12, 100],
+    )
+
+    assert manager.max_resident_sequences() == expected
+
+
+def test_max_resident_sequences_sums_coalesced_life_cycles() -> None:
+    manager = _make_residency_manager(
+        max_seq_len=31,
+        tokens_per_block=16,
+        num_extra_kv_tokens=0,
+        reserve_draft_tokens=0,
+        life_cycle_metadata={
+            0: (0, 31, "attention"),
+            1: (0, None, "ssm"),
+            2: (0, 15, "attention"),
+        },
+        pool_group_totals=[51],
+    )
+
+    # Each sequence consumes 2 + 1 + 2 slots from the shared physical group.
+    assert manager.max_resident_sequences() == 10
+
+
 class _FakeKVCache:
     def __init__(self, num_committed_tokens: int) -> None:
         self.num_committed_tokens = num_committed_tokens