[#13318][fix] Gracefully fit token budget at prep boundary#15187
[#13318][fix] Gracefully fit token budget at prep boundary#15187thorjohnsen wants to merge 24 commits into
Conversation
The micro-batch scheduler's per-step token-budget estimate can diverge from the tokens actually materialized by _prepare_tp_inputs -- e.g. when a reuse-discounted last context chunk lands next to a near-full generation batch. That over-admission tripped the `total_num_tokens <= max_num_tokens` assert in _prepare_tp_inputs, which killed the background executor loop and wedged the server (health checks kept returning 200). Re-validate the budget in KVCacheManager.prepare_resources, before any KV cache is allocated: keep in-flight generation requests, and defer or re-chunk context requests so the batch can never overshoot. Re-chunking only reduces compute tokens (KV is allocated for the full prompt regardless) and is skipped for bidirectional-multimodal requests. A generation-only batch that still overflows raises a clear error instead of corrupting state. Adds GPU-free unit tests covering the upper-bound cost math, re-chunk, defer, multimodal safety, defer-the-rest ordering, and the generation-overflow error path. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughKVCacheManager now enforces per-step token budgets to keep forward-pass overhead within ChangesToken budget constraint and scheduling
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 3
🤖 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/resource_manager.py`:
- Around line 1052-1059: The disaggregated generation-init requests are being
passed into token accounting; change the logic so that when deferring and
req.is_disagg_generation_init_state is true you treat the request as cost-free
and keep it unconditionally (append to kept and continue) instead of calling
self._request_forward_tokens and decrementing remaining; apply this same
early-return/keep pattern to the other similar block that currently computes
cost (the block around the other occurrence of self._request_forward_tokens /
remaining / kept).
- Around line 1071-1088: In _fit_token_budget, when you reassign
req.context_chunk_size (inside the loop that builds kept), mark that a re-chunk
occurred (e.g., set a local rechunked flag) and then call
scheduled_batch.reset_context_requests(kept) whenever rechunked is true (in
addition to the existing len(kept) mismatch check); this ensures
ScheduledRequests' chunk/last-chunk partition is rebuilt after any re-chunking
even if no requests were deferred. Use the existing symbols
req.context_chunk_size, kept, scheduled_batch.reset_context_requests, and the
_fit_token_budget function to locate and implement the change.
In `@tests/unittest/_torch/executor/test_token_budget_fallback.py`:
- Around line 86-101: Add a test variant that includes draft tokens so the
re-chunk path must update last-chunk bookkeeping: create a second scenario in
test_overshoot_rechunks_context where the context request (FakeRequest with
is_last_context_chunk=True and context_chunk_size=64) is accompanied by a
generator request that has non-zero draft tokens (set the FakeRequest attribute
draft_tokens > 0) before calling mgr._fit_token_budget(batch); after calling
_fit_token_budget assert the request is no longer treated as the last-chunk path
(check ctx.is_last_context_chunk is False and/or batch.num_context_requests
unchanged) in addition to the existing assertions on ctx.context_chunk_size and
total token requests computed via mgr._request_forward_tokens.
🪄 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: b853fbd4-3c6e-4218-8c50-07ee5d795e03
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/resource_manager.pytests/unittest/_torch/executor/test_token_budget_fallback.py
|
PR_Github #53182 [ run ] triggered by Bot. Commit: |
|
PR_Github #53182 [ run ] completed with state
|
|
/bot run |
|
PR_Github #53360 [ run ] triggered by Bot. Commit: |
…t fallback When _fit_token_budget absorbs a token-budget overshoot by re-chunking the last context request (rather than deferring one), len(kept) is unchanged, so the previous code skipped reset_context_requests and left the request in the last-chunk bin. Because is_last_context_chunk is a computed property that flips to False once context_chunk_size shrinks, downstream then treated a non-last chunk as final and appended generation/draft tokens to it, producing empty query tensors (q.numel()==0) and invalid attention-kernel arguments. Track whether the batch was modified at all (re-chunk or defer) and re-bin in every modified case. Add a regression test and a docstring for _has_mm_bidirectional_block. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
|
/bot run |
|
PR_Github #53365 [ run ] triggered by Bot. Commit: |
|
PR_Github #53360 [ run ] completed with state |
|
PR_Github #53365 [ run ] completed with state
|
|
/bot run |
|
PR_Github #53367 [ run ] triggered by Bot. Commit: |
|
PR_Github #53367 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53691 [ run ] triggered by Bot. Commit: |
|
PR_Github #53691 [ run ] completed with state
|
…prefill is enabled KVCacheManager._fit_token_budget re-chunked an over-budget context request even when chunked prefill was disabled. The non-chunked attention backend is not set up to consume a partial context chunk, so shrinking context_chunk_size produced an invalid forward pass -- manifesting across models/backends as q.numel()>0 asserts, "Separate quantized buffer is not provided", or cudaErrorInvalidValue. Because the fallback runs on every non-draft prepare_resources call, this broke a broad set of accuracy tests (DeepSeekV3Lite, Llama3 fp8, Qwen3, GPT-OSS) once the scheduler's reuse-discounted token estimate diverged from the materialized token count (the NVIDIA#13318 condition this guard targets) on a batch whose requests were not chunkable. Gate the re-chunk branch on chunked prefill being enabled; otherwise defer the request whole (deferral is always safe -- it drops the request from this iteration's batch and reschedules it later). The flag is threaded into KVCacheManager (default False, the safe defer-only behavior) and set from the finalized attn_runtime_features.chunked_prefill in _create_kv_cache_manager, which runs after py_executor_creator applies its SM-version / attention-backend overrides. Add a regression unit test covering the chunked-prefill-disabled deferral. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
|
/bot run |
|
PR_Github #53972 [ run ] triggered by Bot. Commit: |
Add TorchLlmArgs.enable_token_budget_fallback (default True) so the prep-boundary token-budget fallback (KVCacheManager._fit_token_budget) can be disabled to restore the pre-fallback behavior. The flag is threaded through _create_kv_cache_manager onto the KVCacheManager and gates the call site in prepare_resources. Update the api_stability reference (references/llm.yaml) for the new beta field and add unit tests for the disabled gate and the opt-out default. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
|
PR_Github #53972 [ run ] completed with state
|
…agers The prep-boundary token-budget fallback (NVIDIA#13318) defers/re-chunks context requests in `_fit_token_budget`, mutating `scheduled_batch` in place. It was invoked from `KVCacheManager.prepare_resources`, but the target KV cache manager is deliberately moved to the END of the resource-manager dict (`_util.py` `move_to_end(KV_CACHE_MANAGER)`). Under MTP with a separate draft KV cache manager, that draft manager's `prepare_resources` runs FIRST and adds C++ KV sequences for every context request in the batch -- including ones the fallback then defers. The deferred requests never complete, so their draft-side sequences are never freed; when those requests reschedule on a later iteration the draft manager adds them again, tripping `Assertion failed: emplaceDone (kvCacheManager.cpp)`. Token-budget fitting is a batch-level scheduling decision, not a per-pool one. Hoist it into `ResourceManager.prepare_resources` so it runs once, up front, before any manager allocates -- every manager (draft KV cache, MTP slot manager, etc.) then observes the same deferred batch. Reproduced on H100 with DeepSeek-V3-Lite + MTP(2) + chunked-prefill-off and verified the crash is gone. Adds a regression test asserting a manager registered before the target KV cache manager observes the already-deferred batch. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
|
/bot run |
|
PR_Github #54315 [ run ] triggered by Bot. Commit: |
CI failure triage (builds 44708 / 44814 / 44962 / 45229 / 45308)I went through the failure analyses for the last five Notably, the failures never repeat identically across runs (Triton-cache Every "real-looking" failure on the latest run (#45308) maps to a sibling parametrization already waived on
Plus two clearly-unrelated ones:
Mechanically this is expected: Re-triggering CI to get a clean run. |
|
/bot run --disable-fail-fast |
|
PR_Github #56984 [ run ] triggered by Bot. Commit: |
|
PR_Github #56984 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57210 [ run ] triggered by Bot. Commit: |
|
PR_Github #57210 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57256 [ run ] triggered by Bot. Commit: |
|
/bot run |
|
PR_Github #57256 [ run ] completed with state
|
|
[by Codex] @eopXD Friendly reminder: could you please review this PR? Thanks! |
|
/bot run |
|
PR_Github #57498 [ run ] triggered by Bot. Commit: |
|
PR_Github #57498 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57770 [ run ] triggered by Bot. Commit: |
|
PR_Github #57770 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58031 [ run ] triggered by Bot. Commit: |
|
PR_Github #58031 [ run ] completed with state
|
3 participants
Description
Fixes #13318.
The micro-batch scheduler's per-step token-budget estimate can diverge from the number of tokens actually materialized by
_prepare_tp_inputs. The known trigger: with chunked prefill + KV-cache reuse, a reuse-discounted last context chunk (whose compute budget is reserved asreuse_adjusted_compute = contextRemaining - reusableinmicroBatchScheduler.cpp) lands in the same batch as a near-full generation batch, but_prepare_tp_inputslays out position ids straight fromcontext_chunk_size. The two disagree, the batch overshootsmax_num_tokens, and the assertfires inside the background executor loop. That kills the loop and wedges the server — health-check endpoints keep returning 200 while all subsequent requests hang, so the failure is hard to detect. The reporter observed it under multi-turn chat replay at ~7 QPS (overshoots of +1 to +2890 tokens).
Fix
Re-validate the per-step token budget in
KVCacheManager.prepare_resources, before any KV cache is allocated, and shed only the deferrable work instead of asserting:max_num_tokens, raise a clearRuntimeError(genuine misconfiguration) rather than overshoot silently.(remaining // tokens_per_block) * tokens_per_block). Re-chunking only reduces compute tokens — KV is allocated for the full prompt regardless — so block accounting is unaffected. Skipped for bidirectional-multimodal requests, where splitting a block would corrupt attention (mirrors the gate inscheduler_v2._align_chunk_to_mm_block); those are deferred whole.Because the trim runs before allocation, deferred first-chunk requests never get
add_sequence/add_token, and in-progress chunks keep their existing KV. The cost estimator (_request_forward_tokens) is a deliberate upper bound on each request'sposition_idscontribution, so the assert can no longer fire; the only cost of a loose bound is a marginally smaller batch.This is scoped to the V1
KVCacheManagerpath (where the bug reproduces).KVCacheManagerV2already computeschunk_size = min(remaining_budget, context_remaining)and setscontext_chunk_sizeto that same value inline, so it is structurally immune and is left untouched.Opting out of deferment
The prep-boundary token-budget fallback (the re-chunk / defer behavior described above) is an opt-out feature, controlled by
TorchLlmArgs.enable_token_budget_fallback(defaultTrue). To opt out, set it toFalse— either via the LLM API:or in a YAML config passed to
trtllm-serve:With the fallback disabled, context requests that don't fit are no longer re-chunked or deferred. If the scheduler then hands
_prepare_tp_inputsa batch that exceedsmax_num_tokens, the code raisesTokenBudgetExceededErrorand routes it to_handle_errors(...)with theimmediate_fatalflag set toTrue. Withimmediate_fatal=True, the error bypasses the error budget and is treated as unconditionally fatal: all active and queued requests are failed with the error message and a shutdown is enqueued, so the server terminates cleanly with a clear, actionable error.This replaces the previous behavior, where the overshoot assert killed only the background executor loop thread while the process stayed alive — health-check endpoints kept returning 200 while every subsequent request hung indefinitely (the "wedged server" failure mode described above). Opting out therefore trades the default graceful degradation for fail-fast termination, surfacing the misconfiguration immediately instead of hanging.
Test Coverage
tests/unittest/_torch/executor/test_token_budget_fallback.py(new, GPU-free) drivesKVCacheManager._fit_token_budgetdirectly:test_request_forward_tokens_upper_bound— cost math: contextchunk_size± draft (last chunk only); genbeam × (1 + draft)test_within_budget_is_noop— batch within budget is unchangedtest_overshoot_rechunks_context— the [Bug]: Scheduler deadlock on main + #12976 + #13029: AssertionError total_num_tokens > max_num_tokens in _prepare_tp_inputs under KV offload + chunked prefill permanently hangs the event loop #13318 scenario: near-full gen batch + oversized last chunk → chunk shrunk to a block-aligned fit, total ≤max_num_tokenstest_overshoot_defers_when_cannot_rechunk— sub-block remaining budget → request deferredtest_mm_bidirectional_is_deferred_not_rechunked— multimodal safetytest_defers_all_subsequent_context_requests— defer-the-rest orderingtest_generation_alone_over_budget_raises— non-deferrable overflow →RuntimeErrorFollow-up validation recommended before/with merge:
max_num_tokenspassed toKVCacheManageris identical to themax_num_tokensthe_prepare_tp_inputsassert checks.PR Checklist
TorchLlmArgs.enable_token_budget_fallbackflag — backward-compatible, defaultTrue; carries theapi-compatiblelabel.)Summary by CodeRabbit
Refactor
Tests