[TRTLLM-14138][fix] Pre-allocate CUDA graph padding dummy during warmup - #16072
[TRTLLM-14138][fix] Pre-allocate CUDA graph padding dummy during warmup#16072kaiyux wants to merge 4 commits into
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: |
|
PR_Github #58578 [ run ] completed with state
|
|
[by Codex] @lowsfer Friendly reminder to review this PR. Thanks! |
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 on every subsequent step and padded batches silently fell back to eager mode for the rest of the process lifetime. Allocate the draft_len-0 padding dummy at the end of warmup instead, while the KV cache still has free blocks. Refactor the dummy creation into _get_or_create_padding_dummy() (keeping the encoder-decoder handling) and log a warning_once when padding falls back to eager so the failure mode is visible. Extend test_pytorch_model_engine.py to assert the dummy exists after warmup and that no KV blocks leak beyond it. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
…can use The warmup-time preallocation broke three spec-decode CI tests: - It always allocated the draft_len-0 dummy, which spec-decode runtime padding never requests (it pads with the runtime draft length), so the dummy permanently held KV blocks and hybrid-cache slots on top of the lazily created one. TestQwen3_5_4B::test_dflash failed with "No free blocks left" for the full-attention window pool. - It ran even when padding cannot occur at all (max_batch_size=1 with a batch-size-1 graph in test_eagle3_cdl_sampling), retaining resources the lazy path never consumes. - It ran during the KV cache capacity estimation phase, whose cache has zero headroom when the user-provided max_tokens clamps the estimate; the held block left the estimation request unschedulable and wedged executor initialization (test_eagle3_cdl_sampling, test_dflash_qwen3_8b). Fix: - Preallocate one dummy per draft length of the captured CUDA graphs - exactly the set runtime padding requests, which also covers dynamic draft-length schedules - and only when padding can actually occur for that draft length (_can_pad_any_batch mirrors _get_padded_batch's rounding and capacity guards). - Skip preallocation for estimation-phase KV cache managers. KvCacheCreator.build_managers now stamps is_estimating_kv_cache on the managers it builds. - Extend test_pytorch_model_engine.py with tests for the three gates. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
…-up dispatch Store the existing is_estimating_kv_cache constructor flag on KVCacheManager/KVCacheManagerV2 instead of stamping it onto manager instances in KvCacheCreator.build_managers; the constructors already receive the flag at every creation site, and MambaHybridCacheManager and the AutoDeploy shim already treat it as a constructor parameter. Drop the duplicated dynamic-draft-len enablement check from _get_padded_batch and _can_pad_any_batch: the dynamic mapping is only populated when the feature is active, so _round_up_batch_size_with_draft_len already reduces to plain batch-size rounding without it. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
… draft engines In two-model speculation the target and draft engines share one spec resource manager, and each engine's CUDAGraphRunner registers its padding dummy under the same draft-length-derived request ID (CUDA_GRAPH_DUMMY_REQUEST_ID - draft_len). SlotManager.add_slot only tolerates re-adding the exact sentinel ID, so the second registration died with an AssertionError. Warmup preallocation made this deterministic: the target engine warms up first and registers the dummy for draft_len == max_draft_len, and the draft engine's first-draft graphs use the same draft length, so its preallocation re-registered the same ID and killed the executor during init (test_eagle3_cuda_graph_padding CI failures on pipeline 47174). The lazy runtime path has the same latent collision, but the test never reached it: with 2048-token generations against the 4096-token KV cache the guaranteed-no-evict scheduler caps the batch at two requests, which exactly matches a CUDA graph batch size, so runtime padding never engages. Make dummy registration idempotent in Eagle3ResourceManager and MTPHiddenStatesManager (mirroring SuffixAutomatonManager): reserve the slot once and share it between the engines. Dummy outputs are discarded, so sharing the slot is safe. Shorten the generations in the non-overlap variant of test_eagle3_cuda_graph_padding so runtime padding actually engages and covers the shared-slot path end to end. The overlap variant keeps long generations for now: padded draft batches trip a pre-existing shape mismatch in Eagle3SpecMetadata.prepare on the incremental-update path, which is unrelated to this change and tracked separately. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
4b853e1 to
4d646f7
Compare
|
[by Codex] @yizhang-nv Friendly reminder to review this PR as the primary KV-cache-manager reviewer when you have a chance. Thanks! |
|
[by Codex] @yizhang-nv Could you please review this PR? Thank you! |
| # can distinguish the throwaway estimation-phase managers from the | ||
| # final ones: the estimation cache is sized with no headroom for | ||
| # retained dummy requests. | ||
| self.is_estimating_kv_cache = is_estimating_kv_cache |
There was a problem hiding this comment.
Where is is_estimating_kv_cache used?
4 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