[TRTLLM-14903][fix] Free partially-allocated warmup dummy KV blocks and count spec extra tokens in warmup block estimates

[brnguyen2]

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:

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 added in [None][perf] Close Mamba hybrid warmup gap in autotuner warmup #16177) 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.

The fix makes blocks_to_use mirror the real allocation, and makes add_dummy_requests remove already-registered sequences before re-raising, preserving callers' skip-on-failure semantics.

Note: feat/kimi_k3 currently carries a temporary workaround for this issue (estimation skipped whenever a speculative config is set, in py_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

• PR title and description follow the repo conventions
• Test coverage noted above

Dev Engineer Review

• Updated _create_warmup_request to include extra KV tokens, decoding steps, draft-loop reservations, and beam width when estimating KV-cache block usage.
• Updated add_dummy_requests to remove partially registered target and draft sequences when a later allocation fails.
• Preserved skip-on-failure behavior and prevented leaked KV blocks from exhausting estimation pools.
• No public API, configuration, or test-list changes were identified.
• Integration validation passed for Mamba-hybrid speculative decoding. Logits matched non-speculative and estimation-skipped baselines.
• A full-model GSM8K run at 16-GPU scale scored 96.66 and matched the reference run.
• The V2 KV-cache manager already releases partially allocated dummy requests. No equivalent change is required.
• CI was disabled because the top instance was under maintenance.

QA Engineer Review

No test changes.

[brnguyen2]

/bot run

[tensorrt-cicd]

PR_Github #63215 [ run ] triggered by Bot. Commit: e6b4c3d Link to invocation

[tensorrt-cicd]

PR_Github #63215 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance. For urgent request, contact Yiteng Niu

Link to invocation

[brnguyen2]

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 add_dummy_requests found it already frees partially-allocated dummy requests on capacity failure (success-flag based signaling with an explicit release_resources() on every capacity-failure path), so it does not need the equivalent change.

Marking the PR ready for review. The temporary estimation-skip workaround on feat/kimi_k3 has been reverted on that branch now that the fix is validated there.

[coderabbitai]

Walkthrough

The change improves warmup KV-cache capacity estimation and makes dummy-request allocation transactional. Partial target or draft sequence registration is rolled back before the original exception is re-raised.

Changes

KV-cache management

Layer / File(s)	Summary
Warmup capacity estimation tensorrt_llm/_torch/pyexecutor/model_engine.py	Warmup capacity estimation now includes extra KV and decode tokens, draft-loop reservations, block boundaries, and beam width.
Dummy-request rollback tensorrt_llm/_torch/pyexecutor/resource_manager.py	Dummy target and draft sequence setup removes partially registered sequences when allocation or generation-state initialization fails, then re-raises the original exception.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: schetlur-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the ticket, fix type, and both primary changes: warmup block estimation and cleanup of partially allocated KV blocks.
Description check	✅ Passed	The description explains the failure, root causes, fixes, and relevant integration validation, with the required sections present.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (2)

tensorrt_llm/_torch/pyexecutor/resource_manager.py (2)

976-980: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Unbind the unused loop variable.

token_num is 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 win

Narrow 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 raises RuntimeError for 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

📥 Commits

Reviewing files that changed from the base of the PR and between fdf7bd5 and e6b4c3d.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/model_engine.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py