From f874c5d747428f27120de23d71bc3d0d98df74d3 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Mon, 1 Jun 2026 20:23:06 -0700 Subject: [PATCH] [None][fix] Remove context KV cache allocation rollback Signed-off-by: Jiagan Cheng --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 34 ++-------------- .../_torch/pyexecutor/resource_manager.py | 39 ------------------- .../test_kv_cache_stats_behavior.py | 19 +++++---- 3 files changed, 12 insertions(+), 80 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 98b919e4fd8b..88de275eff4e 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2146,22 +2146,6 @@ def _revert_gen_alloc(self, scheduled_batch): for req in scheduled_batch.generation_requests: self.kv_cache_manager.revert_allocate_generation(req) - def _revert_ctx_alloc(self, dropped_context_requests): - """Revert KV cache capacity growth for ctx requests deferred by - delay batching. - - With KV cache manager V2 + scheduler V2, ctx KV cache is grown - during scheduling (``resize_context``). When delay batching - (``_balance_adp_requests`` for ADP, or ``_waiting_requests`` - for non-ADP batch waiting) defers ctx requests, the - freshly-allocated pages would otherwise sit idle until the - request is re-scheduled, blocking pool space — particularly - painful for long-context workloads where each deferred ctx can - hold GBs of KV. - """ - for req in dropped_context_requests: - self.kv_cache_manager.revert_allocate_context(req) - def _commit_kv_cache_stats(self, scheduled_batch: ScheduledRequests) -> None: if self._scheduler_manages_kv_suspend and isinstance( @@ -3482,8 +3466,7 @@ def _schedule(self): scheduler_output = self.scheduler.schedule_request( self.active_requests, self.inflight_req_ids) - original_ctx_requests = scheduler_output.context_requests - scheduled_context_requests = original_ctx_requests + scheduled_context_requests = scheduler_output.context_requests if self.enable_attention_dp and self.attention_dp_enable_balance: scheduled_context_requests = self._balance_adp_requests( scheduler_output.context_requests, @@ -3494,6 +3477,8 @@ def _schedule(self): scheduler_output.context_requests) > 0 and len( scheduler_output.generation_requests) > 0 if should_check_waiting: + # With KV cache manager V2, scheduling has already grown context request KV cache capacity. Requests dropped + # for batch waiting still occupy KV cache and may reduce the batch size available for generation requests. scheduled_context_requests = self._waiting_requests( scheduler_output.context_requests, scheduler_output.generation_requests) @@ -3507,19 +3492,6 @@ def _schedule(self): scheduled_context_requests) num_fitting = len(scheduled_context_requests) - # V2 scheduler grew KV cache for ctx during scheduling; release - # those pages for any ctx that delay batching has dropped, so - # the wait window does not hold pool capacity hostage. V1 - # allocates after delay batching, so skip the dropped-set - # computation entirely on V1. - if (self._scheduler_manages_kv_suspend and - len(scheduled_context_requests) < len(original_ctx_requests)): - kept = {r.py_request_id for r in scheduled_context_requests} - dropped = [ - r for r in original_ctx_requests if r.py_request_id not in kept - ] - self._revert_ctx_alloc(dropped) - scheduled_requests = ScheduledRequests() scheduled_requests.reset_context_requests(scheduled_context_requests) scheduled_requests.generation_requests = scheduler_output.generation_requests diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index f88c0043f21b..f4b07dc94b40 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2711,37 +2711,6 @@ def revert_allocate_generation(self, req: LlmRequest) -> None: f"{req.py_request_id} from {kv_cache.capacity} to " f"{reverted_cap}") - def revert_allocate_context(self, req: LlmRequest) -> None: - """Undo the capacity growth from this iter's ``resize_context``. - - When delay batching (``_balance_adp_requests`` / - ``_waiting_requests``) defers a context request after V2 - scheduling, the forward pass is skipped for that request but the - scheduler already grew its KV cache capacity to cover the chunk. - This shrinks capacity back to the pre-resize value so the - freshly-allocated pages can be reused during the wait window — - important for long contexts where one deferred request can hold - GBs of KV. - """ - pre_cap = getattr(req, "py_ctx_pre_resize_cap", None) - if pre_cap is None: - return - # Mark as consumed even if the resize below is skipped, so a - # later iter does not see a stale snapshot. - req.py_ctx_pre_resize_cap = None - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None or not kv_cache.is_active: - return - if pre_cap >= kv_cache.capacity: - return - if not kv_cache.resize(pre_cap): - raise RuntimeError( - f"Failed to revert KV cache capacity for context " - f"request {req.py_request_id} from " - f"{kv_cache.capacity} to {pre_cap}") - if pre_cap > 0: - kv_cache.suspend() - def _restore_page_index_bufs(self, request_id: int, kv_cache) -> None: """Re-connect host page-index buffers after resume(). @@ -2834,10 +2803,6 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: Returns True on success, False if resize failed (first chunk is suspended on failure). - - Snapshots the pre-resize capacity on ``req.py_ctx_pre_resize_cap`` - when growth happens so ``revert_allocate_context`` can undo it if - delay batching defers the request. """ assert not req.is_disagg_generation_init_state, ( f"req {req.py_request_id}: use prepare_disagg_gen_init") @@ -2847,7 +2812,6 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: target = req.context_current_position + num_tokens + self.num_extra_kv_tokens capacity = max(kv_cache.capacity, target) - pre_cap = kv_cache.capacity success = kv_cache.resize(capacity) if not success: @@ -2855,9 +2819,6 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: kv_cache.suspend() return False - # None means "no growth this iter, nothing to revert"; this also - # invalidates a stale snapshot from a prior iter on the same req. - req.py_ctx_pre_resize_cap = pre_cap if capacity > pre_cap else None return True def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index efd778227aef..8bf24d288ad9 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py @@ -50,6 +50,7 @@ class _StatsRequest: is_dummy_request: bool = False is_attention_dp_dummy: bool = False is_cuda_graph_dummy: bool = False + is_disagg_generation_init_state: bool = False is_disagg_generation_transmission_complete: bool = False context_phase_params: None = None py_draft_tokens: list[int] = field(default_factory=list) @@ -359,7 +360,7 @@ def test_reverted_generation_allocation_does_not_report_stats(resource_guard) -> _assert_iteration_delta(stats_report.by_window_size[manager.max_seq_len]) -def test_reverted_context_allocation_does_not_report_pending_stats(resource_guard) -> None: +def test_waited_context_allocation_reports_pending_stats_when_scheduled(resource_guard) -> None: request = _StatsRequest(1, list(range(8)), context_remaining_length=8) manager = resource_guard(_create_manager(gpu_bytes=8 << 20), request) @@ -377,25 +378,23 @@ def test_reverted_context_allocation_does_not_report_pending_stats(resource_guar request.is_first_context_chunk = False assert manager.prepare_context(request) assert manager.resize_context(request, num_tokens=4) - manager.revert_allocate_context(request) - manager.commit_scheduled_kv_cache_stats(_context_batch(request)) - reverted_stats_report = manager.get_iteration_stats() - assert reverted_stats_report is not None - _assert_iteration_delta(reverted_stats_report.by_window_size[manager.max_seq_len]) + second_chunk_stats = _commit_and_get_stats(manager, _context_batch(request)) + _assert_iteration_delta(second_chunk_stats, alloc_total=1, alloc_new=1, missed=1) assert request.kv_cache_perf_metric_calls == [ _metric_call(alloc_total=1, alloc_new=1, missed=1), + _metric_call(alloc_total=1, alloc_new=1, missed=1), ] kv_stats = manager.get_kv_cache_stats() - assert kv_stats.alloc_total_blocks == 1 - assert kv_stats.alloc_new_blocks == 1 - assert kv_stats.missed_blocks == 1 + assert kv_stats.alloc_total_blocks == 2 + assert kv_stats.alloc_new_blocks == 2 + assert kv_stats.missed_blocks == 2 assert manager.prepare_context(request) assert manager.resize_context(request, num_tokens=4) _finish_context(manager, request) second_chunk_stats = _commit_and_get_stats(manager, _context_batch(request)) - _assert_iteration_delta(second_chunk_stats, alloc_total=1, alloc_new=1, missed=1) + _assert_iteration_delta(second_chunk_stats) assert request.kv_cache_perf_metric_calls == [ _metric_call(alloc_total=1, alloc_new=1, missed=1), _metric_call(alloc_total=1, alloc_new=1, missed=1),