tensorrt_llm/_torch/pyexecutor/resource_manager.py (1)

2351-2381: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Clear stale py_ctx_pre_resize_cap before attempting resize.

If kv_cache.resize(...) fails (Line 2373), req.py_ctx_pre_resize_cap is left untouched. A prior iteration’s snapshot can then be incorrectly consumed by revert_allocate_context, causing an outdated shrink target.

💡 Suggested fix

 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
+        # Clear stale snapshot from prior iterations before resize attempt.
+        req.py_ctx_pre_resize_cap = None
 
         if not kv_cache.resize(capacity):
             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
+        if capacity > pre_cap:
+            req.py_ctx_pre_resize_cap = pre_cap
         return True

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 2351 - 2381,
In resize_context, clear any stale snapshot before attempting the resize by
setting req.py_ctx_pre_resize_cap = None prior to calling kv_cache.resize(...);
then compute pre_cap and on successful growth (capacity > pre_cap) set
req.py_ctx_pre_resize_cap = pre_cap, otherwise leave it None; keep the existing
kv_cache.suspend() behavior when kv_cache.resize(...) fails for
req.is_first_context_chunk and return False as before so revert_allocate_context
cannot consume an outdated snapshot.

tensorrt_llm/_torch/pyexecutor/py_executor.py (1)

1566-1566: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Also revert context growth when can_queue skips the batch.

_schedule() now unwinds ctx growth for requests dropped by delay batching, but the later skip paths still only call _revert_gen_alloc(). On KV-manager V2, resize_context has already grown scheduled_batch.context_requests, so an attention-DP skip / post-connector recheck / PP queue skip can still leave those ctx pages pinned until a future iteration. That recreates the same pool-pressure leak for skipped batches.

Suggested direction

-    def _revert_gen_alloc(self, scheduled_batch):
-        """Revert KV cache capacity growth when the batch is skipped.
+    def _revert_skipped_batch_alloc(self, scheduled_batch):
+        """Revert KV cache capacity growth when a scheduled batch is skipped.
 
-        With attention DP, can_queue=False means another rank has an empty
-        batch so no forward pass will run.  The V2 scheduler already grew
-        each generation request's KV cache capacity during scheduling;
-        revert that growth so it does not accumulate across skipped
-        iterations and overflow the host page-index buffer.
+        With KV cache manager V2, scheduling may already have grown both
+        context and generation capacity before we discover that the batch
+        cannot run on this iteration.
         """
         if self._is_kv_manager_v2:
+            for req in scheduled_batch.context_requests:
+                self.kv_cache_manager.revert_allocate_context(req)
             for req in scheduled_batch.generation_requests:
                 self.kv_cache_manager.revert_allocate_generation(req)

Then switch the not can_queue branches in PP / overlap / non-overlap to call the new helper.

Also applies to: 1951-1983, 2257-2258, 2513-2514

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/pyexecutor/py_executor.py` at line 1566, When a batch is
skipped by a can_queue check, also undo any context growth (resize_context)
instead of only calling _revert_gen_alloc; replace the direct calls to
self._revert_gen_alloc(scheduled_batch) in _schedule and in the PP / overlap /
non-overlap "not can_queue" branches with the new helper that unrolls context
growth (the helper introduced alongside _revert_gen_alloc that also reverts
scheduled_batch.context_requests), so skipped batches do not leave ctx pages
pinned.