[TRTLLM-12714][fix] Suspend CUDA-graph padding dummies before pool rebalance adjust() - #16157
Conversation
|
/bot run |
WalkthroughKV pool rebalance now suspends active CUDA-graph padding dummies before adjustment. It resumes dummies afterward without freeing them. Unresumable dummies release resources and are removed for later recreation. Tests cover these paths. ChangesKV pool rebalance padding dummy lifecycle
Estimated code review effort: 4 (Complex) | ~40 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant CUDAGraphRunner
participant KvCacheManager
participant ResourceManagers
PyExecutor->>CUDAGraphRunner: suspend padding dummies
PyExecutor->>KvCacheManager: adjust KV pool
PyExecutor->>CUDAGraphRunner: resume padding dummies
CUDAGraphRunner->>ResourceManagers: release unresumable dummy resources
CUDAGraphRunner->>CUDAGraphRunner: remove cached dummies
Possibly related PRs
Suggested reviewers: 🚥 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!
|
Should address the other manager's free resource as well. |
|
/bot run --disable-fail-fast |
|
PR_Github #65075 [ run ] triggered by Bot. Commit: |
…s managers _get_or_create_padding_dummy spreads one dummy request ID across up to four managers: the main KV cache manager, the one-model draft KV cache manager, the speculative resource manager slot and, for encoder-decoder, the cross-KV cache manager. Releasing only the main one leaves the others holding the ID, and re-creation reuses the same CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len. Add CUDAGraphRunner.release_padding_dummy(), the release path symmetric with _get_or_create_padding_dummy: it frees the dummy from every manager that allocated part of it and drops it from the runner so a later padded step re-creates it. The manager list lives in _padding_dummy_managers() next to the creation path so the two stay in step, and is deduplicated by identity since double-freeing one manager is not safe in general. The rebalance hook's resume-failure branch now goes through it instead of calling free_resources on the main KV cache manager alone. That branch is the only place the dummies are released at all: the normal path suspends and resumes them. Addresses review feedback from chienchunhung and yizhang-nv on NVIDIA#16157. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
|
@yizhang-nv — done in Short version: the PR has been reworked to suspend the padding dummies for The one remaining release, on the rare resume-failure branch, now goes through a new |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 1480-1483: Extend the test around the non-encoder-decoder
assertion to cover the encoder-decoder branch: set runner.is_encoder_decoder to
True, create a padding dummy through the existing setup, and verify
cross_manager.free_resources is called exactly once with that dummy. Preserve
the existing assertion that non-encoder-decoder runners do not invoke the
cross-KV manager, and exercise the conditional manager from
_padding_dummy_managers.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5516db92-6777-4f16-9b2a-d823308db8a8
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_kv_pool_rebalance.pytests/unittest/_torch/executor/test_pytorch_model_engine.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- tests/unittest/_torch/executor/test_kv_pool_rebalance.py
|
/bot run --disable-fail-fast |
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
|
PR_Github #65077 [ run ] triggered by Bot. Commit: |
|
PR_Github #65079 [ run ] triggered by Bot. Commit: |
|
PR_Github #65077 [ run ] completed with state |
|
PR_Github #65075 [ run ] completed with state |
chienchunhung
left a comment
chienchunhung
left a comment
There was a problem hiding this comment.
Thanks for addressing the comments; LGTM!
The Known limitation paragraph in the PR description is stale at the current head: release_padding_dummy() now covers the primary, one-model draft, speculative-resource, and encoder-decoder cross-KV managers. Could you remove or update that paragraph so the description matches the implementation?
|
PR_Github #65079 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65134 [ run ] triggered by Bot. Commit: |
|
@asfiyab-nvidia @JunyiXu-nv — could one of you take a look when you get a chance? The PR has two approvals but is missing one from |
|
PR_Github #65134 [ run ] completed with state
|
|
@mikeiovine @litaotju @QiJune @pcastonguay @schetlur-nv @juney-nvidia — could one of you add PR #16157 already has two code-review approvals but needs an approval from |
|
/bot run --disable-fail-fast |
|
PR_Github #65390 [ run ] triggered by Bot. Commit: |
|
PR_Github #65390 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65459 [ run ] triggered by Bot. Commit: |
|
PR_Github #65459 [ run ] completed with state |
schetlur-nv
left a comment
schetlur-nv
left a comment
There was a problem hiding this comment.
Rubber stamping for runtime devs based on others' reviews
8 participants
Description
CUDAGraphRunnerretains one padding dummy request per captured draft length, whoseKVCacheManagerV2cache staysACTIVEacross iterations but never appears inPyExecutor.active_requests. The rebalance hook therefore never suspends them, and the first liveadjust()fails its all-caches-suspended precondition, terminating the executor event loop along with all in-flight requests.This PR suspends the dummies alongside the active requests and resumes them after
adjust().Why suspend rather than free. Suspension is what the precondition actually asks for: it tears down the cache's base-page-index buffers and releases its page locks, which is what makes the pages safe to migrate. Those buffers are written only when a page lock is taken (
_page.py:449/page.cpp:380) and cleared on unlock (_page.py:467/page.cpp:391) — nothing refreshes them when a page later migrates. A cache leftACTIVEacrossadjust()can therefore end up addressing slots that now belong to other sequences, and the shrink path migrates held pages by design (_gather_persistent_pages→shrink_pool_group→_batchedMigrate(defrag=True)).Interaction with #16072. An earlier revision of this PR freed the dummies and let
_pad_batchre-create them. That is no longer safe: #16072 pre-allocates them at the end of warmup precisely because lazy allocation against a saturated KV cache fails and silently drops padded batches to eager mode for the rest of the process lifetime. Freeing them mid-run returns to that lazy path at the worst moment — rebalance only fires after 2000 sampled caches and a 120 s cooldown, i.e. under sustained load, andadjust()may have just shrunk the pool group the dummy maps to.preallocate_padding_dummiesonly runs fromModelEngine.warmup, so the pre-allocation would not be re-established. Suspending holds the reservation across the whole rebalance instead.Resume failure.
resume_requestcan legitimately returnFalse(GPU pressure abovemax_util_for_resume, or out of pages). A real request is then left suspended for the scheduler to reactivate, but nothing reschedules a padding dummy, and_get_or_create_padding_dummyreturns a cached dummy without checking that its cache is live — so a dummy left suspended would be padded into a batch with a torn-down block table. One that cannot be resumed is released and dropped from the runner instead, falling back to lazy re-creation.release_padding_dummy()releases the dummy from every manager that allocated part of it: the main KV cache manager, the one-model draft KV cache manager, the speculative resource manager slot, and (for encoder-decoder models) the cross-KV cache manager. Duplicates are dropped by identity.Test Coverage
tests/unittest/_torch/executor/test_kv_pool_rebalance.py, newTestPaddingDummiesclass (21 passed, was 15):test_padding_dummy_is_suspended_before_adjust— the dummy is suspended beforeadjust()runs.test_padding_dummy_is_resumed_and_never_freed— pins [TRTLLM-14138][fix] Pre-allocate CUDA graph padding dummy during warmup #16072's invariant: the dummy survives the rebalance and stays in the runner's map.test_every_captured_draft_length_is_covered— [TRTLLM-14138][fix] Pre-allocate CUDA graph padding dummy during warmup #16072 retains one dummy per captured draft length, not just one.test_unresumable_padding_dummy_is_released_and_dropped— a dummy that cannot be resumed is released and removed from the runner.test_already_suspended_padding_dummy_is_left_alonetest_missing_cuda_graph_runner_is_toleratedtests/unittest/_torch/executor/test_pytorch_model_engine.py, newtest_release_padding_dummy_covers_every_manager:release_padding_dummy()drops the dummy from the runner map, frees the spec resource manager slot, does not invoke the cross-KV manager for a non-encoder-decoder engine, and is idempotent on a second call.Mutation-tested in both directions:
test_padding_dummy_is_resumed_and_never_freedThe second confirms the suite would catch a regression against #16072.
Live reproduction. On
mainwith gemma-3-1b-it (VSWA, 2 pool groups) andCudaGraphConfig(enable_padding=True), underTLLM_KV_CACHE_MANAGER_V2_BACKEND=python— where the precondition is a plainassertrather than the C++ backend's debug-gatedTLLM_CHECK_DEBUG—adjust()raisedAssertionErrorat_kv_cache_manager.py:872and killed the executor loop. With this change the same run completes and the GPU pool ratio moves from 0.500/0.500 to 0.667/0.333.Also run:
kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py— 112 passed, 13 skipped (unaffected; the change touches onlypy_executor.py).PR Checklist
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.