From ded67ab7105c0221e9502dae4887e097cb09eac0 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 3 Aug 2026 06:44:14 -0700 Subject: [PATCH 1/3] [nvbugs/6550275][fix] Bound V2 scheduler residency for non-droppable state pools MAX_UTILIZATION admits new sequences up to max_batch_size and relies on suspend/resume to survive over-subscription. That recovery needs some resident sequence to still be evictable so its pages can be freed for a suspended one to resume. A hybrid Mamba recurrent state is fixed-size per sequence and cannot be recomputed from tokens, so such a sequence yields no evictable pages; once every sequence is suspended the pool can no longer drain and the scheduler stops making progress. On L40S, qwen3.5_9b at 500-in/2000-out admitted 485 sequences while the attention pool holds only ~118 at max_seq_len, then spun 8676 iterations scheduling nothing before raising 'V2 scheduler deadlock'. Add KVCacheManagerV2.max_resident_sequences(), derived from per-pool-group page counts, and gate new-sequence admission on it. Returns None (unbounded, unchanged behavior) for models without a non-droppable state pool. Pin the return value on the scheduler test's manager double: a bare Mock() is not None, so the new gate would otherwise compare an int against a child Mock and raise TypeError in every test that schedules a first context chunk. Add direct coverage for the cap, which had none. Two config levers suggested by triage were measured and do not fix this: avg_seq_len=2500 reproduces the failure identically, and max_util_for_resume=1.0 replaces the raise with an unbounded livelock. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 37 ++++++++++++++ .../pyexecutor/scheduler/scheduler_v2.py | 22 ++++++++ .../executor/test_kv_cache_v2_scheduler.py | 51 +++++++++++++++++++ 3 files changed, 110 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 56c6c00e0218..d991ca2648e4 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2151,6 +2151,43 @@ 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]: + """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. + """ + state_pool_groups = { + pool_group_id + for pool_group_id, _, kind in self._stats_life_cycle_metadata().values() + if kind != "attention" + } + if not state_pool_groups: + return None + + # Charging every attention pool the full max_seq_len is deliberately + # worst-case: a genuinely sliding-window pool needs fewer pages per + # sequence, so the bound errs low rather than over-admitting. + pages_per_seq = math.ceil(self.max_seq_len / self.tokens_per_block) + # A state pool holds one fixed slot per sequence; an attention pool + # must hold every page of each resident sequence. + return max( + 1, + min( + stat.total // (1 if pool_group_id in state_pool_groups else pages_per_seq) + for pool_group_id, stat in enumerate(self._get_storage_statistics(GPU_LEVEL)) + ), + ) + def _effective_draft_len(self, req: LlmRequest) -> int: """Draft token length to use for next-step KV capacity calculation. diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 958afc59ac04..1d762a28129a 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/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}" ) @@ -400,9 +404,25 @@ 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. + residency_cap = self.max_resident_sequences + num_started = ( + sum(1 for r in requests_list if self._is_started_request(r)) + if residency_cap is not None + else 0 + ) + for req in pending_ctx: if budget.requests_full: break + # Read before scheduling: _try_schedule_context advances the + # request past its first chunk. + is_new_sequence = residency_cap is not None and req.is_first_context_chunk + if is_new_sequence and num_started >= residency_cap: + break peft_pages = budget.peft_pages_needed(req) if peft_pages is None: continue @@ -418,6 +438,8 @@ def _schedule_loop(self, active_requests, inflight_request_ids): break if action is ScheduleAction.SKIP: continue + if is_new_sequence: + num_started += 1 has_chunking = has_chunking or chunking_flag scheduled_ctx.append(req) budget.commit(req, tokens, peft_pages) diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index 0a7c67e9e702..031b003bbf0f 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/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 + # Must be pinned: a bare Mock() is not None, so the scheduler's residency + # gate would compare an int against a child Mock and raise TypeError. + mgr.max_resident_sequences.return_value = max_resident_sequences return mgr @@ -2217,6 +2221,53 @@ def test_interleaved_states(self): assert ids(out.context_requests) == [4] +# =========================================================================== +# Residency Cap (non-droppable state pools) +# =========================================================================== + + +class TestResidencyCap: + """A manager reporting max_resident_sequences bounds newly started + sequences, so a pool whose state cannot be evicted never over-subscribes. + """ + + def test_none_leaves_admission_unbounded(self): + """Plain attention models report None and keep MAX_UTILIZATION behavior.""" + mgr = make_kv_cache_manager(max_resident_sequences=None) + sched = make_scheduler(mgr) + reqs = [make_ctx_request(i, context_remaining_length=10) for i in range(20)] + out = sched.schedule_request(reqs, set()) + assert len(out.context_requests) == 20 + + def test_caps_newly_started_sequences(self): + mgr = make_kv_cache_manager(max_resident_sequences=3) + sched = make_scheduler(mgr) + reqs = [make_ctx_request(i, context_remaining_length=10) for i in range(20)] + out = sched.schedule_request(reqs, set()) + assert ids(out.context_requests) == [0, 1, 2] + + def test_already_started_sequences_consume_the_cap(self): + """In-progress generation holds its slot, so no new context is admitted.""" + mgr = make_kv_cache_manager(max_resident_sequences=2) + sched = make_scheduler(mgr) + reqs = [ + make_gen_request(0), + make_gen_request(1), + make_ctx_request(2, context_remaining_length=10), + ] + out = sched.schedule_request(reqs, set()) + assert ids(out.generation_requests) == [0, 1] + assert ids(out.context_requests) == [] + + def test_continuing_chunk_is_not_charged_again(self): + """Only a first chunk starts a sequence; later chunks keep their slot.""" + mgr = make_kv_cache_manager(max_resident_sequences=1) + sched = make_scheduler(mgr) + 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] + + # =========================================================================== # Block Reuse Boundary Alignment (Chunked + Partial Reuse) # =========================================================================== From e0c11f9aa529d5f0d0215e505f95a893b83aa0a9 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:27:44 +0800 Subject: [PATCH 2/3] [nvbugs/6550275][fix] Address residency review feedback (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 34 ++++----- .../pyexecutor/scheduler/scheduler_v2.py | 55 +++++++++++---- .../executor/test_kv_cache_manager_v2.py | 63 +++++++++++++++++ .../executor/test_kv_cache_v2_scheduler.py | 70 ++++++++++++------- 4 files changed, 165 insertions(+), 57 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index d991ca2648e4..deb28f0624ab 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/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,10 @@ 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``. @@ -2166,25 +2168,23 @@ def max_resident_sequences(self) -> Optional[int]: Returns None when no such non-droppable pool exists, keeping the unbounded MAX_UTILIZATION behavior for plain attention models. """ - state_pool_groups = { - pool_group_id - for pool_group_id, _, kind in self._stats_life_cycle_metadata().values() - if kind != "attention" - } - if not state_pool_groups: + life_cycle_metadata = self._stats_life_cycle_metadata() + if not any(kind != "attention" for _, _, kind in life_cycle_metadata.values()): return None - # Charging every attention pool the full max_seq_len is deliberately - # worst-case: a genuinely sliding-window pool needs fewer pages per - # sequence, so the bound errs low rather than over-admitting. - pages_per_seq = math.ceil(self.max_seq_len / self.tokens_per_block) - # A state pool holds one fixed slot per sequence; an attention pool - # must hold every page of each resident sequence. + # 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( - stat.total // (1 if pool_group_id in state_pool_groups else pages_per_seq) - for pool_group_id, stat in enumerate(self._get_storage_statistics(GPU_LEVEL)) + storage_stats[pool_group_id].total // slots_per_sequence + for pool_group_id, slots_per_sequence in slots_per_pool_group.items() ), ) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 1d762a28129a..37aa63084472 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -214,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 @@ -290,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 @@ -336,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 @@ -347,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 @@ -408,21 +429,16 @@ def _schedule_loop(self, active_requests, inflight_request_ids): # 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. - residency_cap = self.max_resident_sequences - num_started = ( - sum(1 for r in requests_list if self._is_started_request(r)) - if residency_cap is not None - else 0 - ) - for req in pending_ctx: if budget.requests_full: break - # Read before scheduling: _try_schedule_context advances the - # request past its first chunk. - is_new_sequence = residency_cap is not None and req.is_first_context_chunk - if is_new_sequence and num_started >= residency_cap: - 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 @@ -438,8 +454,8 @@ def _schedule_loop(self, active_requests, inflight_request_ids): break if action is ScheduleAction.SKIP: continue - if is_new_sequence: - num_started += 1 + 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) @@ -1007,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. diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index b625cc974dee..06940e0e6c5c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/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 diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index 031b003bbf0f..b42b8f26a650 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -170,7 +170,12 @@ def make_kv_cache_manager( mgr.prepare_disagg_gen_init.side_effect = prepare_disagg_gen_init_fn or (lambda req: True) 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 + if max_resident_sequences is None: + mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active + else: + mgr.is_request_active.side_effect = ( + lambda req_id: req_id in mgr.kv_cache_map and mgr.kv_cache_map[req_id].is_active + ) # Must be pinned: a bare Mock() is not None, so the scheduler's residency # gate would compare an int against a child Mock and raise TypeError. mgr.max_resident_sequences.return_value = max_resident_sequences @@ -1363,30 +1368,19 @@ def selective_prepare(req): out = sched.schedule_request(reqs, set()) assert ids(out.fitting_disagg_gen_init_requests) == [1, 3] - def test_disagg_cross_iteration_slot_overflow(self): - """Reproduce IndexMapper slot overflow across iterations. - - Simulates the real crash scenario: - Iter 1: disagg_gen_init scheduled → slots consumed → state becomes - DISAGG_GENERATION_TRANS_IN_PROGRESS (value=9). - Iter 2: TRANS_IN_PROGRESS (9) is invisible to both the disagg branch - (checks ==8) and state gating ([10,14)), so budget resets to - 0. New disagg_gen_init passes budget → prepare_disagg_gen_init - called again → IndexMapper would crash in production. - - This test counts prepare_disagg_gen_init calls across two iterations. - With scheduler_capacity=2 (simulating IndexMapper=3 slots, 1 dummy), a - correct implementation should cap total prepare_disagg_gen_init calls to - 2 (the IndexMapper capacity). The bug allows 4. - """ + def test_disagg_transfer_counts_against_residency_across_iterations(self): + """In-progress transfers retain residency and defer new init requests.""" prepare_count = [0] def counting_prepare(req): prepare_count[0] += 1 return True - mgr = make_kv_cache_manager(prepare_disagg_gen_init_fn=counting_prepare) - sched = make_scheduler(mgr, max_num_tokens=200, scheduler_capacity=2) + mgr = make_kv_cache_manager( + prepare_disagg_gen_init_fn=counting_prepare, + max_resident_sequences=2, + ) + sched = make_scheduler(mgr, max_num_tokens=200) # Iteration 1: two disagg_gen_init requests reqs_iter1 = [make_disagg_request(0), make_disagg_request(1)] @@ -1399,18 +1393,14 @@ def counting_prepare(req): for req in reqs_iter1: req.state_value = DISAGG_GEN_TRANS_IN_PROGRESS - # Iteration 2: old requests (now TRANS_IN_PROGRESS) + 2 new disagg. - # TRANS_IN_PROGRESS (9) is invisible to scheduler: not ==8, not in [10,14). + # Iteration 2: old transfers still occupy both resident slots. new_reqs = [make_disagg_request(2), make_disagg_request(3)] all_active = reqs_iter1 + new_reqs out2 = sched.schedule_request(all_active, set()) - # BUG: budget resets to 0, TRANS_IN_PROGRESS not counted, so - # both new disagg pass → prepare_disagg_gen_init called 4 times total. - # In production, this would crash IndexMapper (only 3 slots = 2+1 dummy). - assert ids(out2.fitting_disagg_gen_init_requests) == [2, 3] - assert prepare_count[0] == 4 # <-- proves the overflow + assert ids(out2.fitting_disagg_gen_init_requests) == [] + assert prepare_count[0] == 2 # =========================================================================== @@ -2267,6 +2257,34 @@ def test_continuing_chunk_is_not_charged_again(self): out = sched.schedule_request([req], set()) assert ids(out.context_requests) == [0] + def test_caps_successful_disagg_init_in_phase_one(self): + mgr = make_kv_cache_manager(max_resident_sequences=2) + sched = make_scheduler(mgr) + + out = sched.schedule_request([make_disagg_request(i) for i in range(3)], set()) + + assert ids(out.fitting_disagg_gen_init_requests) == [0, 1] + assert mgr.prepare_disagg_gen_init.call_count == 2 + + def test_disagg_init_does_not_block_continuing_context(self): + mgr = make_kv_cache_manager(max_resident_sequences=2) + sched = make_scheduler(mgr) + continuing = make_ctx_request( + 2, + context_remaining_length=10, + is_first_context_chunk=False, + ) + reqs = [ + make_disagg_request(0), + make_ctx_request(1, context_remaining_length=10), + continuing, + ] + + out = sched.schedule_request(reqs, set()) + + assert ids(out.fitting_disagg_gen_init_requests) == [0] + assert ids(out.context_requests) == [2] + # =========================================================================== # Block Reuse Boundary Alignment (Chunked + Partial Reuse) From ce537324a34220d9f7ecbc7b5146586ea685d89d Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Sat, 8 Aug 2026 16:08:57 +0800 Subject: [PATCH 3/3] [nvbugs/6550275][fix] Synchronize V2 residency cap (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 36 ++++++--- .../executor/test_kv_cache_manager_v2.py | 74 ++++++++++++++++++- 2 files changed, 100 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index deb28f0624ab..5fcc8775dc0a 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2165,11 +2165,23 @@ def max_resident_sequences(self) -> Optional[int]: 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. + Returns None when no rank has such a non-droppable pool, keeping the + unbounded MAX_UTILIZATION behavior for plain attention models. When + one exists, every rank contributes its local full-sequence capacity + so pipeline-parallel ranks admit the same request set even when some + ranks contain only attention layers or no local layers. """ life_cycle_metadata = self._stats_life_cycle_metadata() - if not any(kind != "attention" for _, _, kind in life_cycle_metadata.values()): + has_non_droppable_pool = any( + kind != "attention" for _, _, kind in life_cycle_metadata.values() + ) + dist = None + if self.mapping.world_size > 1: + dist = Distributed.get(self.mapping) + has_non_droppable_pool = bool( + dist.allreduce(int(has_non_droppable_pool), op=ReduceOp.MAX) + ) + if not has_non_droppable_pool: return None # A physical group can host multiple lifecycle variants. V2 allocates @@ -2179,14 +2191,20 @@ def max_resident_sequences(self) -> Optional[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( + if slots_per_pool_group: + storage_stats = self._get_storage_statistics(GPU_LEVEL) + local_capacity = min( storage_stats[pool_group_id].total // slots_per_sequence for pool_group_id, slots_per_sequence in slots_per_pool_group.items() - ), - ) + ) + else: + # A PP rank with no local cache layers must not constrain ranks + # that do own physical cache pools. + local_capacity = sys.maxsize + + if dist is not None: + local_capacity = dist.allreduce(local_capacity, op=ReduceOp.MIN) + return max(1, int(local_capacity)) def _effective_draft_len(self, req: LlmRequest) -> int: """Draft token length to use for next-step KV capacity calculation. diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index 06940e0e6c5c..fe0d956a4f2f 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -13,13 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +import sys from dataclasses import dataclass, field from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import call, patch import pytest import torch +from tensorrt_llm._torch.distributed.communicator import ReduceOp from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType @@ -48,8 +50,10 @@ def _make_residency_manager( reserve_draft_tokens: int, life_cycle_metadata: dict[int, tuple[int, int | None, str]], pool_group_totals: list[int], + world_size: int = 1, ) -> KVCacheManagerV2: manager = object.__new__(KVCacheManagerV2) + manager.mapping = SimpleNamespace(world_size=world_size) manager.max_seq_len = max_seq_len manager.tokens_per_block = tokens_per_block manager.num_extra_kv_tokens = num_extra_kv_tokens @@ -103,6 +107,74 @@ def test_max_resident_sequences_sums_coalesced_life_cycles() -> None: assert manager.max_resident_sequences() == 10 +def test_max_resident_sequences_returns_none_for_attention_only() -> 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")}, + pool_group_totals=[4], + ) + + assert manager.max_resident_sequences() is None + + +def test_max_resident_sequences_floors_zero_capacity_at_one() -> 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, None, "ssm")}, + pool_group_totals=[0], + ) + + assert manager.max_resident_sequences() == 1 + + +def test_max_resident_sequences_syncs_attention_only_pp_rank() -> 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")}, + pool_group_totals=[8], + world_size=2, + ) + + with patch("tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2.Distributed.get") as get_dist: + get_dist.return_value.allreduce.side_effect = [1, 3] + + assert manager.max_resident_sequences() == 3 + + get_dist.return_value.allreduce.assert_has_calls( + [call(0, op=ReduceOp.MAX), call(4, op=ReduceOp.MIN)] + ) + + +def test_max_resident_sequences_ignores_pp_rank_without_cache_layers() -> 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={}, + pool_group_totals=[], + world_size=2, + ) + + with patch("tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2.Distributed.get") as get_dist: + get_dist.return_value.allreduce.side_effect = [1, 7] + + assert manager.max_resident_sequences() == 7 + + get_dist.return_value.allreduce.assert_has_calls( + [call(0, op=ReduceOp.MAX), call(sys.maxsize, op=ReduceOp.MIN)] + ) + + class _FakeKVCache: def __init__(self, num_committed_tokens: int) -> None: self.num_committed_tokens = num_committed_tokens