[None][fix] Remove context KV cache allocation rollback

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

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,17 +2812,13 @@ 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:
             if req.is_first_context_chunk:
                 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:

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),