[TRTLLM-14903][fix] Free partially-allocated warmup dummy KV blocks and count spec extra tokens in warmup block estimates - #17162
Conversation
…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>
|
/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 |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
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
|
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>
|
The Folded a harness fix into this PR (dc5a809): strip |
|
/bot run |
|
PR_Github #63453 [ run ] triggered by Bot. Commit: |
|
PR_Github #63453 [ run ] completed with state
|
|
Looked into the two failures in the latest CI run; neither reproduces with this PR's change:
Rerunning CI. |
|
/bot run |
|
PR_Github #63499 [ run ] triggered by Bot. Commit: |
|
PR_Github #63499 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63510 [ run ] triggered by Bot. Commit: |
|
PR_Github #63510 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63521 [ run ] triggered by Bot. Commit: |
| blocks_to_use = (num_full_seqs * math.ceil( | ||
| (max_seq_len + extra_ctx_tokens) / tokens_per_block) + (math.ceil( | ||
| (num_left_over_tokens + extra_ctx_tokens) / | ||
| tokens_per_block) if num_left_over_tokens > 0 else 0) + | ||
| num_gen_requests * self.max_beam_width * math.ceil( | ||
| (1 + extra_gen_tokens) / tokens_per_block)) |
There was a problem hiding this comment.
This expression is not very readable. Can we clean it up
There was a problem hiding this comment.
Cleaned up in b57fba9: pulled out a blocks_for_seq helper and split the sum into one statement per term (full context sequences, leftover-token sequence, gen requests).
| tokens_per_block = kv_cache_manager.tokens_per_block | ||
| extra_ctx_tokens = (getattr(kv_cache_manager, "num_extra_kv_tokens", 0) | ||
| or 0) + num_extra_decoding_steps | ||
| extra_gen_tokens = extra_ctx_tokens + self.max_draft_loop_tokens |
There was a problem hiding this comment.
This is going to be 2 * max_draft_len - 1 in most cases, why?
There was a problem hiding this comment.
It mirrors what add_dummy_requests actually allocates per gen dummy: num_extra_kv_tokens (max_draft_len - 1 in one-engine modes) + num_extra_decoding_steps + max_draft_loop_tokens (max_draft_len for a linear draft loop) add_token calls on top of the single prompt token. The old estimate ignored these, so warmup could pass this check and then fail midway through the actual allocation, which is the leak this PR fixes. Added a comment spelling out the arithmetic in 5bdb732.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
PR_Github #63558 [ run ] triggered by Bot. Commit: |
|
PR_Github #63521 [ run ] completed with state |
|
PR_Github #63558 [ run ] completed with state
|
4 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.