[TRTLLM-14903][fix] Free partially-allocated warmup dummy KV blocks and count spec extra tokens in warmup block estimates - #17162
Conversation
|
/bot run |
|
PR_Github #63215 [ run ] triggered by Bot. Commit: |
|
PR_Github #63215 [ run ] completed with state |
|
Validated beyond the original truncated-model repro: on a Mamba-hybrid MoE model with suffix-automaton speculative decoding and KV-cache estimation enabled (the previously hanging configuration), a full speculative-decoding logits-parity integration run now completes in ~17 minutes with parity statistics identical to the estimation-skipped and pre-regression baselines (52 prompts, 0 drift), with the estimation phase confirmed active in the logs and no warmup-overflow warnings. A GSM8K accuracy run on the full model at 16-GPU scale with speculative decoding and estimation enabled scores 96.66, matching the reference measured with estimation skipped. An audit of the V2 KV-cache manager's Marking the PR ready for review. The temporary estimation-skip workaround on |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe change updates warmup KV-cache capacity estimation, makes dummy-request setup transactional, and removes ChangesKV-cache management
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 (2)
tensorrt_llm/_torch/pyexecutor/resource_manager.py (2)
976-980: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnbind the unused loop variable.
token_numis not read in this loop. Ruff reports B007.♻️ Proposed fix
- for req_id, token_num, _ in batch_request_infos: + for req_id, _, _ in batch_request_infos:🤖 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 `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 976 - 980, Update the loop over batch_request_infos to bind the unused token_num element to an underscore, preserving req_id and the existing token-addition loops.Source: Linters/SAST tools
1039-1046: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNarrow the swallowed exception and log it.
The inner handler discards every exception type. A real failure in
remove_sequence(for example a binding signature change) then disappears, and the rollback silently stops for that request. The C++ binding raisesRuntimeErrorfor an unregistered sequence, so catch that type and log at debug level. Ruff reports S110 and BLE001 here.♻️ Proposed fix
for req in freeing_requests: try: freeing_impl.remove_sequence(req.py_request_id, req, False) - except Exception: + except RuntimeError as e: # The sequence may never have been registered (the # batched add itself failed); nothing to clean up. - pass + logger.debug( + "Dummy request rollback skipped for request " + f"{req.py_request_id}: {e}")As per coding guidelines "Catch the narrowest possible exceptions, keep duck-typing try blocks minimal, prefer
isinstance(), use built-in exception types".🤖 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 `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 1039 - 1046, Update the exception handler in the freeing_requests rollback loop around freeing_impl.remove_sequence to catch only RuntimeError, log the caught exception at debug level, and allow other exception types to propagate. Keep the try block limited to the remove_sequence call so genuine rollback failures are not silently swallowed and Ruff S110/BLE001 are resolved.Sources: Coding guidelines, Linters/SAST tools
🤖 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 `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 976-980: Update the loop over batch_request_infos to bind the
unused token_num element to an underscore, preserving req_id and the existing
token-addition loops.
- Around line 1039-1046: Update the exception handler in the freeing_requests
rollback loop around freeing_impl.remove_sequence to catch only RuntimeError,
log the caught exception at debug level, and allow other exception types to
propagate. Keep the try block limited to the remove_sequence call so genuine
rollback failures are not silently swallowed and Ruff S110/BLE001 are resolved.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 43777f60-714f-419b-8a39-2e0d0fcfe29b
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/resource_manager.py
|
/bot run |
|
PR_Github #63288 [ run ] triggered by Bot. Commit: |
|
PR_Github #63288 [ run ] completed with state
|
|
The Folded a harness fix into this PR (dc5a809): strip |
|
/bot run |
|
PR_Github #64015 [ run ] triggered by Bot. Commit: |
|
PR_Github #64010 [ run ] completed with state |
|
PR_Github #64015 [ run ] completed with state
|
…nd count spec extra tokens in warmup block estimates Two defects combined to hang LLM startup indefinitely during KV cache size estimation for Mamba-hybrid models with speculative decoding: 1. _create_warmup_request under-counted blocks_to_use: it ignored the per-sequence extra tokens (num_extra_kv_tokens, num_extra_decoding_steps, and the draft-token reserve for generation dummies) that add_dummy_requests actually allocates. With spec decoding, block-aligned multi-sequence warmup shapes (e.g. the Mamba hybrid multi-seq warmup) passed the estimate but overflowed the pool at allocation time. 2. add_dummy_requests leaked every already-registered sequence when a later add_token raised (e.g. "no free blocks left"). On the minimal KV pool built for cache-size estimation the leak left too few blocks for the estimation requests themselves, so the executor loop spun forever without scheduling them and LLM() never returned. Fix blocks_to_use to mirror the real allocation, and make add_dummy_requests remove already-registered sequences before re-raising, preserving callers' skip-on-failure semantics. Verified on a spec-decoding estimation-phase integration run on a Mamba-hybrid model (previously hung unboundedly; now completes with logits parity against a non-speculative baseline). Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
multi-round-qa.py imports its sibling utils.py through the implicit script-directory sys.path entry. When the test environment sets PYTHONSAFEPATH=1, that entry is disabled and the benchmark client exits immediately with ModuleNotFoundError: No module named 'utils', failing TestServePrefixAwareScheduling tests with 'Smoke warmup failed with rc=1' while the server is still healthy. Strip the variable from the client subprocess environment. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…ailures remove_sequence is already a no-op for request ids the failed batched add never registered, so the per-request exception handler only masked real releaseBlocks failures — which leave the KV pool poisoned, the same hang mechanism this cleanup exists to prevent. Attempt cleanup for every target and draft request, then re-raise the first cleanup failure instead of swallowing it. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
A 4-block pool admits both dummy sequences but runs out of blocks in the per-request draft add_token loop, exercising the partial-failure path: the exception must propagate and every already-allocated block must be freed. Fails against the pre-fix code (blocks leak), passes with the cleanup. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
45e7638 to
01196c0
Compare
|
PR_Github #64044 [ run ] triggered by Bot. Commit: |
|
PR_Github #64044 [ run ] completed with state
|
unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_moe.py and unittest/auto_deploy/singlegpu/transformations/library/test_moe_fusion.py fail in DGX_B200-AutoDeploy-1 for any PR rebased onto current main: identical failures on pipelines 51974 (this PR) and 51977 (PR NVIDIA#17225, zero file overlap). Suspect commit 89bba4c (NVIDIA#15297), which modifies the trtllm_moe custom op and Blackwell blockScaleMoe kernels. Waived pending an NVBug; the entries will be updated with the bug link once filed. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run --stage-list "DGX_B200-AutoDeploy-1" |
|
PR_Github #64078 [ run ] triggered by Bot. Commit: |
Cites nvbugs/6564714 for the two waives added in the previous commit. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run --stage-list "DGX_B200-AutoDeploy-1" |
|
PR_Github #64087 [ run ] triggered by Bot. Commit: |
|
PR_Github/17162-1fa12af #64078 was force-killed by a newer pipeline run. |
|
PR_Github #64087 [ run ] completed with state |
|
/bot skip --comment "Waive-only change on top of fully qualified head; targeted DGX_B200-AutoDeploy-1 run on this exact commit passed (pipeline 52014); waived tests are a pre-existing main-side breakage tracked in nvbugs/6564714." |
|
PR_Github #64107 [ skip ] triggered by Bot. Commit: |
|
PR_Github #64107 [ skip ] completed with state |
7 participants
Description
Fixes TRTLLM-14903:
LLM()startup hangs indefinitely during KV cache size estimation for Mamba-hybrid models with speculative decoding enabled.Two defects combined to produce the hang:
_create_warmup_requestunder-countedblocks_to_use: it ignored the per-sequence extra tokens (num_extra_kv_tokens,num_extra_decoding_steps, and the draft-token reserve for generation dummies) thatadd_dummy_requestsactually allocates. With spec decoding, block-aligned multi-sequence warmup shapes (e.g. the Mamba hybrid multi-seq warmup added in [None][perf] Close Mamba hybrid warmup gap in autotuner warmup #16177) passed the estimate but overflowed the pool at allocation time.add_dummy_requestsleaked every already-registered sequence when a lateradd_tokenraised (e.g. "no free blocks left"). On the minimal KV pool built for cache-size estimation, the leak left too few blocks for the estimation requests themselves, so the executor loop spun forever without scheduling them andLLM()never returned.The fix makes
blocks_to_usemirror the real allocation, and makesadd_dummy_requestsremove already-registered sequences before re-raising, preserving callers' skip-on-failure semantics.Note:
feat/kimi_k3currently carries a temporary workaround for this issue (estimation skipped whenever a speculative config is set, inpy_executor_creator.py, marked TRTLLM-14903); that workaround should be reverted once this fix merges.Test Coverage
Verified on a spec-decoding estimation-phase integration run on a Mamba-hybrid model: previously hung unboundedly during estimation warmup; with the fix the run completes and passes logits parity against a non-speculative baseline. The "Mamba hybrid warmup skipped" overflow path is now taken before any allocation, so the pool is no longer poisoned.
PR Checklist
Dev Engineer Review
_create_warmup_requestnow accounts for extra KV tokens, decoding steps, draft-loop reservations, and beam width.add_dummy_requestsremoves partially registered target and draft sequences after allocation failure.PYTHONSAFEPATHbefore launching the client.QA Engineer Review
tests/integration/defs/kv_cache/test_prefix_aware_scheduling.pywas modified.