[TRTLLM-12714][fix] Free CUDA-graph padding dummy KV caches before pool rebalance adjust() - #16157
Conversation
…ol rebalance adjust() CudaGraphRunner._pad_batch creates persistent padding dummy requests whose KVCacheManagerV2 caches stay ACTIVE across iterations but never appear in PyExecutor.active_requests. The rebalance hook therefore never suspends them, and the first live adjust() fails its all-caches-suspended precondition assert, terminating the executor event loop with all in-flight requests. CUDA-graph padding is enabled by default, so any deployment of enable_kv_pool_rebalance on a multi-pool-group model (e.g. gemma-3 with VSWA) hits this within the tuner's first adjustment. Close the padding dummies before suspending and clear padding_dummy_requests so _pad_batch lazily recreates them with the post-adjust pool layout, mirroring the existing reset idiom in ModelEngine.warmup and CudaGraphRunner.clear. Verified live on 1x H100 with gemma-3-1b-it (VSWA, 2 pool groups): 9 rebalances across three benchmark runs, including one run under TLLM_KV_CACHE_MANAGER_V2_DEBUG=1, with no assertion failures. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
|
/bot run |
📝 WalkthroughWalkthroughThe KV pool rebalance hook in PyExecutor now frees CUDA-graph padding dummy requests' KV resources and clears the padding dummy map before suspending active requests and calling adjust. A test fixture default and a new regression test validate this cleanup order and behavior. ChangesKV pool rebalance padding dummy cleanup
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant CudaGraphRunner
participant KvCacheManager
PyExecutor->>CudaGraphRunner: check padding_dummy_requests
loop each dummy request
PyExecutor->>KvCacheManager: free_resources(dummy)
end
PyExecutor->>CudaGraphRunner: clear padding_dummy_requests
PyExecutor->>KvCacheManager: impl.adjust()
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_kv_pool_rebalance.py (1)
206-235: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueGood regression coverage for the free-before-adjust ordering.
The test correctly exercises the fix's core contract:
free_resourcesis called exactly once with the dummy,padding_dummy_requestsends up empty, and the free happens beforeadjust(). One gap: there's no test for therunner is None(nocuda_graph_runnerattribute) or multi-dummy-map case; the current fixture always provides aMagicMockrunner with a dict, so that guard branch (if runner is not None and runner.padding_dummy_requests:) is untested. Given it's a trivialgetattr/dict-truthiness guard, this is a minor coverage gap rather than a blocker — consider adding it if you want full branch coverage, but not required for this PR.As per path instructions, "Act as a QA engineer reviewing test changes and coverage... suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR": coverage here is sufficient for the fix's primary contract; the
runner is Nonebranch is an optional follow-up.🤖 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 `@tests/unittest/_torch/executor/test_kv_pool_rebalance.py` around lines 206 - 235, Add a small follow-up test around PyExecutor._maybe_rebalance_kv_pools to cover the guard path where cuda_graph_runner is None or where padding_dummy_requests is empty, since the current test only exercises the non-empty MagicMock runner case. Reuse the existing _make_executor and _make_request helpers, but set model_engine.cuda_graph_runner to None (or an empty dummy map) and assert the rebalance path still runs without trying to free padding dummies. This will cover the runner-null/dict-truthiness branch in the same area as test_frees_cuda_graph_padding_dummies_before_adjust.Source: Path instructions
🤖 Prompt for all review comments with 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.
Nitpick comments:
In `@tests/unittest/_torch/executor/test_kv_pool_rebalance.py`:
- Around line 206-235: Add a small follow-up test around
PyExecutor._maybe_rebalance_kv_pools to cover the guard path where
cuda_graph_runner is None or where padding_dummy_requests is empty, since the
current test only exercises the non-empty MagicMock runner case. Reuse the
existing _make_executor and _make_request helpers, but set
model_engine.cuda_graph_runner to None (or an empty dummy map) and assert the
rebalance path still runs without trying to free padding dummies. This will
cover the runner-null/dict-truthiness branch in the same area as
test_frees_cuda_graph_padding_dummies_before_adjust.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 34f50a4b-a97f-4d3f-b709-faf252a59d14
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_kv_pool_rebalance.py
|
PR_Github #58329 [ run ] triggered by Bot. Commit: |
|
PR_Github #58329 [ run ] completed with state
|
|
/bot |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run |
|
PR_Github #58504 [ run ] triggered by Bot. Commit: |
|
PR_Github #58504 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58989 [ run ] triggered by Bot. Commit: |
|
PR_Github #58989 [ run ] completed with state
|
chienchunhung
left a comment
chienchunhung
left a comment
There was a problem hiding this comment.
Thanks for the PR!
| runner = getattr(self.model_engine, "cuda_graph_runner", None) | ||
| if runner is not None and runner.padding_dummy_requests: | ||
| for dummy in runner.padding_dummy_requests.values(): | ||
| mgr.free_resources(dummy) |
There was a problem hiding this comment.
This only releases the primary KV cache, shall we also clear the other places? For example, CUDAGraphRunner can also register the same dummy ID with other places, e.g., one-model draft KV manager, speculative resource manager, and encoder-decoder cross-KV manager.
I am also wondering if there is a way that we can centralize the cleanup.
|
Should address the other manager's free resource as well. |
| # living cache to be suspended. Close them and let _pad_batch | ||
| # recreate them on the next padded batch, with the post-adjust | ||
| # pool layout. | ||
| runner = getattr(self.model_engine, "cuda_graph_runner", None) |
|
/bot run --disable-fail-fast |
|
PR_Github #61357 [ run ] triggered by Bot. Commit: |
|
PR_Github #61357 [ run ] completed with state
|
5 participants
Background
KvCacheConfig.enable_kv_pool_rebalance(the KVCacheManagerV2 auto-tuner hook, #14578) suspends every request inPyExecutor.active_requestsbefore callingadjust(), which asserts that all living KV caches are suspended.CudaGraphRunner._pad_batchcreates persistent padding dummy requests whose V2 KV caches stay ACTIVE across iterations but never appear inactive_requests. The hook therefore never suspends them, and the first liveadjust()fails its precondition assert, terminating the executor event loop with all in-flight requests errored. CUDA-graph padding is enabled by default, so any deployment of the flag on a multi-pool-group model (e.g. gemma-3 with VSWAmax_attention_window) hits this within the tuner's first adjustment (~2 minutes of serving).Summary
_maybe_rebalance_kv_pools: close the padding dummies (mgr.free_resources) and clearrunner.padding_dummy_requestsbefore the suspend loop;_pad_batchlazily recreates them on the next padded batch with the post-adjust pool layout. Mirrors the existing reset idiom inModelEngine.warmupandCudaGraphRunner.clear.test_frees_cuda_graph_padding_dummies_before_adjust(verifies the dummies are freed beforeadjust()runs and the map is cleared).Impact
enable_kv_pool_rebalance=True, default off); no behavior change otherwise.TLLM_KV_CACHE_MANAGER_V2_DEBUG=1— with no assertion failures.tests/unittest/_torch/executor/test_kv_pool_rebalance.py: 16/16 pass.Test Coverage
tests/unittest/_torch/executor/test_kv_pool_rebalance.py::TestMaybeRebalanceKvPools::test_frees_cuda_graph_padding_dummies_before_adjust🤖 Generated with Claude Code
Summary by CodeRabbit