[TRTLLM-14138][fix] Pre-allocate CUDA graph padding dummy during warmup - #16072
Conversation
📝 WalkthroughWalkthroughThe CUDA graph padding dummy request allocation in ChangesPadding Dummy Allocation Refactor
Estimated code review effort: 3 (Moderate) | ~25 minutes 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.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_pytorch_model_engine.py (1)
340-349: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for the allocation-failure/fallback path.
test_warmupnow verifies the happy path (preallocation succeeds and frees cleanly), but there's no test exercising the failure path this PR also touches:_get_padded_batch's newwarning_once+ return-0 fallback when_get_or_create_padding_dummyreturnsNone(e.g., KV cache saturated by the time a padded batch is first requested), norpreallocate_padding_dummy's own warning-and-continue branch when preallocation itself fails at warmup.Coverage is currently insufficient for the fallback branch introduced in
cuda_graph_runner.py(_get_padded_batchlines 527-533,preallocate_padding_dummylines 617-621). Suggest adding a test in this file that exhausts KV cache blocks (e.g., via dummy requests) before callingpad_batch/preallocate_padding_dummy, then asserts a0return / warning and that no exception is raised.Based on path instructions for
tests/**: "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM. Keep feedback actionable: suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR."🤖 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_pytorch_model_engine.py` around lines 340 - 349, Coverage is insufficient for the new allocation-failure fallback in cuda_graph_runner.py: add a test in test_pytorch_model_engine.py that forces KV cache exhaustion before the first padded batch/preallocation attempt, then calls pad_batch and preallocate_padding_dummy to verify the warning-once path, a 0 fallback from _get_padded_batch, and that preallocate_padding_dummy logs a warning and continues without raising. Use the existing test_warmup setup plus padding_dummy_requests and kv_cache_manager to locate the relevant flow.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.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 605-625: preallocate_padding_dummy() only allocates the padding
dummy for draft_len=0, but _get_padded_batch() can later request separate
dummies for each runtime draft length when dynamic draft-length scheduling is
enabled. Update preallocate_padding_dummy() to iterate over
self._dynamic_draft_len_mapping.values() and call
_get_or_create_padding_dummy(resource_manager, draft_len) for every reachable
draft length, while keeping the existing enabled/padding_enabled guard and
warning/logging behavior in CudaGraphRunner.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 340-349: Coverage is insufficient for the new allocation-failure
fallback in cuda_graph_runner.py: add a test in test_pytorch_model_engine.py
that forces KV cache exhaustion before the first padded batch/preallocation
attempt, then calls pad_batch and preallocate_padding_dummy to verify the
warning-once path, a 0 fallback from _get_padded_batch, and that
preallocate_padding_dummy logs a warning and continues without raising. Use the
existing test_warmup setup plus padding_dummy_requests and kv_cache_manager to
locate the relevant flow.
🪄 Autofix (Beta)
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: ab97c5a3-44a7-4595-846c-13afd6e6a70d
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/executor/test_pytorch_model_engine.py
kaiyux
changed the title
|
/bot run |
|
PR_Github #58124 [ run ] triggered by Bot. Commit: |
|
PR_Github #58124 [ run ] completed with state
|
4e49bae to
c2fc7f3
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58164 [ run ] triggered by Bot. Commit: |
|
PR_Github #58164 [ run ] completed with state
|
18142d5 to
81ba27b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58448 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #58451 [ run ] triggered by Bot. Commit: |
|
PR_Github #58448 [ run ] completed with state |
|
PR_Github #58451 [ run ] completed with state
|
d10363b to
4b853e1
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58578 [ run ] triggered by Bot. Commit: |
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #62733 [ run ] triggered by Bot. Commit: |
|
PR_Github #62733 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62904 [ run ] triggered by Bot. Commit: |
|
PR_Github #62904 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63339 [ run ] triggered by Bot. Commit: |
|
PR_Github #63339 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63373 [ run ] triggered by Bot. Commit: |
|
PR_Github #63373 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63668 [ run ] triggered by Bot. Commit: |
|
PR_Github #63668 [ run ] completed with state |
|
The preallocation has no ctx/gen role gate, and the feasibility check doesn't mirror the no-ctx-requests condition. A |
@Shixiaowei02 Would CTX disable cuda graph? If CTX disable the gen only cudagraph, the code just does not take effect by checking whether cuda graph is enabled at the beginning of preallocate_padding_dummies. |
…balance adjust() CUDAGraphRunner retains one padding dummy request per captured draft length, whose KVCacheManagerV2 cache stays ACTIVE across iterations but never appears in PyExecutor.active_requests. The rebalance hook therefore never suspends them, and the first live adjust() fails its all-caches-suspended precondition, terminating the executor event loop along with all in-flight requests. Suspend the dummies alongside the active requests and resume them after adjust(). 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 and are never refreshed when a page later migrates, so a cache left ACTIVE across adjust() can end up addressing slots that now belong to other sequences. The dummies are suspended rather than freed so that the warmup pre-allocation added in NVIDIA#16072 survives a rebalance. Freeing them would return to lazy re-creation against a KV cache that is under load -- rebalance only fires after 2000 sampled caches and a 120s cooldown -- and possibly just shrunk, which is exactly the case where allocation fails and padded batches silently fall back to eager mode for the rest of the process lifetime. A dummy that cannot be resumed is released and dropped from the runner instead, because nothing reschedules a padding dummy and _get_or_create_padding_dummy returns a cached dummy without checking that its cache is live. Reproduced on main with gemma-3-1b-it (VSWA, 2 pool groups) and CudaGraphConfig(enable_padding=True) under TLLM_KV_CACHE_MANAGER_V2_BACKEND=python, where the precondition is a plain assert rather than the C++ backend's debug-gated TLLM_CHECK_DEBUG: adjust() raised AssertionError and 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. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
7 participants
Description
The CUDA graph padding dummy request was allocated lazily at the first padded step. If the KV cache was already saturated by that point, the allocation failed — and kept failing on every subsequent step — so padded batches silently fell back to eager mode for the rest of the process lifetime, losing CUDA graph coverage exactly when the server is under load.
This PR:
CUDAGraphRunner.preallocate_padding_dummies()called fromPyTorchModelEngine.warmup). One dummy is created per draft length of the captured graphs — exactly the set runtime padding can request (speculative engines pad with non-zero draft lengths, so unconditionally preallocatingdraft_len=0would hold KV blocks and spec/hybrid slots the lazy path never uses)._can_pad_any_batch, e.g.max_batch_size=1with a captured graph for batch size 1), so infeasible configs don't retain KV blocks for nothing.KVCacheManager/KVCacheManagerV2now retain theis_estimating_kv_cacheflag they already receive as a constructor parameter so the runner can detect estimation-phase managers (MambaHybridCacheManageralready did)._get_or_create_padding_dummy(), keeping the existing encoder-decoder handling, and drops the duplicated dynamic-draft-len enablement check at the round-up call sites: the dynamic mapping is only populated when the feature is active, so_round_up_batch_size_with_draft_lenalready reduces to plain batch-size rounding without it.warning_oncewhen CUDA graph padding falls back to eager, so the failure mode is visible instead of silent, and logs the preallocation outcome per draft length at warmup.Trade-off: the final executor now permanently holds a few KV blocks per captured draft length — the same blocks the lazy path would have held from the first successfully padded step onward. The feasibility and estimation gates keep this from affecting configs that never pad.
Split out of #16033 (1 of 3, together with #16073 and #16074); this PR is independent of the other two.
Test Coverage
tests/unittest/_torch/executor/test_pytorch_model_engine.py:test_warmupasserts the padding dummy exists after warmup and that no KV cache blocks leak beyond it.test_warmup_skips_padding_dummy_when_padding_impossibleasserts no dummy is retained when every reachable batch size already matches a captured graph.test_warmup_skips_padding_dummy_during_estimationasserts no dummy is retained on estimation-phase KV cache managers.test_preallocate_padding_dummies_uses_captured_draft_lensasserts dummies are created for the captured draft lengths, notdraft_len=0unconditionally.045705139d) on 1x B200 (Gemma-4-31B-IT-NVFP4, ISL 1024 / OSL 128, 1024 prompts, concurrency 1024,trtllm-serve+benchmark_serving): 1319.0 vs 1331.8 tok/s output throughput (-1.0%, within run-to-run noise). Throughput-neutral is expected for this particular workload: at concurrency 1024 the steady-state decode batch size (~230) exceeds the CUDA graph max batch size, so the padded-batch path is not exercised at all. The change is a robustness fix for deployments where CUDA graph padding is active — there, KV-cache saturation before the first padded step previously caused a silent, permanent fallback to eager for every padded batch, and the pre-allocation preserves the CUDA graph coverage.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.🤖 Generated with Claude Code