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

2079-2091: ⚠️ Potential issue | 🟠 Major

Resource leak on resize() failure in is_gen block — inconsistent with other error paths.

When kv_cache.resize(new_capacity) fails (Line 2081) or draft_kv_cache.resize(new_capacity) fails (Line 2087), a ValueError is raised. Unlike every other failure path in this function, release_resources is not called before raising, leaving the already-allocated kv_cache (and draft_kv_cache, if applicable) in kv_cache_map without being freed.

🐛 Proposed fix

             if prepare_resource:
                 new_capacity = kv_cache.capacity + max_num_draft_tokens + 1
                 success = kv_cache.resize(new_capacity)
                 if not success:
-                    raise ValueError(
-                        f"Failed to resize capacity of KV cache for request {req.py_request_id} to {new_capacity} tokens for dummy request"
-                    )
+                    release_resources(req)
+                    return None
                 if draft_kv_cache is not None:
                     success = draft_kv_cache.resize(new_capacity)
                     if not success:
-                        raise ValueError(
-                            f"Failed to resize capacity of draft KV cache for request {req.py_request_id} to {new_capacity} tokens for dummy request"
-                        )
+                        release_resources(req, free_draft_resources=True)
+                        return None

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 2079 - 2091,
In the is_gen branch where prepare_resource is true, ensure you free
already-allocated resources before raising on resize failure: when
kv_cache.resize(new_capacity) or draft_kv_cache.resize(new_capacity) returns
False, call release_resources(req.py_request_id) (or the existing
release_resources function that cleans up entries in kv_cache_map for this
request) to remove and free kv_cache and draft_kv_cache, then raise the
ValueError; update the error paths in the block that contains prepare_resource,
kv_cache.resize, draft_kv_cache.resize and the existing raise statements so they
always call release_resources(req.py_request_id) prior to raising.

---

2264-2266: ⚠️ Potential issue | 🟡 Minor

get_max_resource_count() returning 1 is still a placeholder and may cause scheduling issues.

Returning 1 prevents division-by-zero but any scheduler that uses this to compute maximum capacity or admission thresholds will severely underestimate available resources for KVCacheManagerV2, potentially admitting far fewer requests than the hardware supports. The V1 equivalent returns self.impl.max_num_blocks.

Consider wiring through the actual page count, e.g.:

def get_max_resource_count(self) -> int:
    return max(
        self.impl.get_page_index_upper_bound(layer_id, Role.KEY)
        for layer_id in typed_range(LayerId(self.num_local_layers))
    ) // self.kv_factor

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 2264 - 2266,
get_max_resource_count currently returns a placeholder 1 which underestimates
capacity; replace it to compute the real max resource count by iterating local
layers (use typed_range(LayerId(self.num_local_layers))) and calling
self.impl.get_page_index_upper_bound(layer_id, Role.KEY) for each, take the
maximum of those values and divide by self.kv_factor (use integer division if
intended) to mirror the V1 behavior (self.impl.max_num_blocks) so schedulers get
an accurate capacity for KVCacheManagerV2.

---

2046-2067: ⚠️ Potential issue | 🔴 Critical

Replace ValueError with release_resources() pattern to prevent resource leak on resize failures.

At lines 2082–2091, when kv_cache.resize() or draft_kv_cache.resize() fails in the is_gen block, the code raises ValueError directly without freeing allocated resources. Unlike the error paths at lines 2051–2072 which correctly call release_resources(), these failures leak both the primary and draft KV cache objects. Use release_resources() followed by return None instead of raising.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 2046 - 2067,
In the is_gen branch, replace the direct raising of ValueError when
kv_cache.resize(...) or draft_kv_cache.resize(...) fails with the same cleanup
pattern used elsewhere: call release_resources(req) (and release_resources(req,
free_draft_resources=True) where the draft cache was allocated) and then return
None; update the failure paths around kw functions kv_cache.resize and
draft_kv_cache.resize so both primary and draft KV cache objects are freed via
release_resources() instead of raising an exception.