Skip to content

[None][fix] Pre-subtract non-first-chunk context costs in reuse budge…#12806

Merged
liji-nv merged 1 commit into
NVIDIA:feat/bench_yfrom
liji-nv:fix/mnt-overflow-chunked-budget
Apr 7, 2026
Merged

[None][fix] Pre-subtract non-first-chunk context costs in reuse budge…#12806
liji-nv merged 1 commit into
NVIDIA:feat/bench_yfrom
liji-nv:fix/mnt-overflow-chunked-budget

Conversation

@liji-nv

@liji-nv liji-nv commented Apr 7, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

…t check

The prepare_resources budget re-validation for KV cache reuse was not accounting for non-first-chunk context requests before evaluating first-chunk requests. When a non-first-chunk request (with fixed compute cost) appeared after first-chunk requests in the batch, the budget check for first-chunk requests did not see the committed cost, allowing too many tokens to be scheduled and causing MNT overflow.

Fix by pre-subtracting all non-first-chunk context_chunk_size from remaining_budget before the allocation loop, so first-chunk budget checks see the correct available capacity.

@coderabbitai summary

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • 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)

  • 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.

@liji-nv
[None][fix] Pre-subtract non-first-chunk context costs in reuse budge…
ed9017e 
…t check

The prepare_resources budget re-validation for KV cache reuse was not
accounting for non-first-chunk context requests before evaluating
first-chunk requests.  When a non-first-chunk request (with fixed
compute cost) appeared after first-chunk requests in the batch, the
budget check for first-chunk requests did not see the committed cost,
allowing too many tokens to be scheduled and causing MNT overflow.

Fix by pre-subtracting all non-first-chunk context_chunk_size from
remaining_budget before the allocation loop, so first-chunk budget
checks see the correct available capacity.

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
@liji-nv
liji-nv requested a review from a team as a code owner April 7, 2026 12:02
@liji-nv
liji-nv merged commit 01c4947 into NVIDIA:feat/bench_y Apr 7, 2026
6 of 7 checks passed
dongfengy pushed a commit to dongfengy/TensorRT-LLM that referenced this pull request Apr 8, 2026
@liji-nv @dongfengy
[None][fix] Pre-subtract non-first-chunk context costs in reuse budge… (
cdf7717 
NVIDIA#12806)

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
dongfengy pushed a commit to dongfengy/TensorRT-LLM that referenced this pull request Apr 10, 2026
@liji-nv @dongfengy
[None][fix] Pre-subtract non-first-chunk context costs in reuse budge… (
5dc2445 
NVIDIA#12806)

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
dongfengy pushed a commit to dongfengy/TensorRT-LLM that referenced this pull request Apr 10, 2026
@liji-nv @dongfengy
[None][fix] Pre-subtract non-first-chunk context costs in reuse budge… (
3d59f85 
NVIDIA#12806)

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
@Saddss Saddss mentioned this pull request Apr 22, 2026
Open
5 tasks
This was referenced Apr 22, 2026
Closed
Closed
Draft
@nv-yna nv-yna mentioned this pull request Apr 29, 2026
Closed
nv-yna added a commit to nv-yna/TensorRT-LLM that referenced this pull request Apr 29, 2026
@nv-yna
[None][fix] Add Python-side remaining_budget guard for KV reuse budge…
f0a16e8 
…t overflow

Cherry-pick of Python portions from feat/bench_y PRs NVIDIA#12682 and NVIDIA#12806
that were not included in NVIDIA#12976 (which ported only the C++ fix to main).

Adds a remaining_budget re-validation guard in KVCacheManager.prepare_resources()
that re-probes the radix tree for actual reusable blocks after KV cache allocation
and skips requests whose forward cost exceeds the remaining budget. This catches
the estimation-vs-reality gap when cache eviction between scheduling and
prepare_resources() reduces actual reuse below the scheduler's estimate.

Original authors: Liao Lanyu (@lancelly), Jin Li (@liji-nv)

Signed-off-by: Yuewei Na <nv-yna@users.noreply.github.com>
chienchunhung added a commit to chienchunhung/TensorRT-LLM that referenced this pull request May 19, 2026
@chienchunhung
[None][docs] Add issue-13318 _prepare_tp_inputs wedge investigation
a3d622c 
Document the root cause of GitHub issue NVIDIA#13318 (V1 chunked-prefill
admission budget mismatch under enable_block_reuse + enable_partial_reuse
+ KV host offload) and draft a two-track fix plan: (A) port PR NVIDIA#12806's
pre-subtract semantic to the C++ V1 MicroBatchScheduler + audit the
Python _reuse_adjusted_compute callers; (B) wrap _event_loop_wrapper so
a single batch error fails the offending requests instead of killing the
executor thread.

Includes gap analysis against prior NVIDIA#12718 / NVIDIA#13119 hardening and
verifies that PR NVIDIA#12806 is merged only to feat/bench_y, not main.

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
chienchunhung added a commit to chienchunhung/TensorRT-LLM that referenced this pull request May 20, 2026
@chienchunhung
[None][doc] update issue-13318 fix plan with verified main state
c8b91b5 
Adds a "Verified State on main (2026-05-19)" section that pins down the
three claims the rest of the plan depends on, then folds those findings
back into the existing tracks.

- V1 (PR NVIDIA#12718 status): NVIDIA#12718 is merged; documents the primitives that
  landed (_error_queue, _handle_background_error, ManagedThread,
  signal.raise_signal(SIGINT) on CppExecutorError, LLM._check_health)
  with file:line, and the three reasons they do not catch this bug:
  (a) PyExecutor's event loop is a plain threading.Thread, not a
  ManagedThread, so its dying exception never reaches _error_queue;
  (b) LLM._check_health() observes Executor.doing_shutdown, not
  PyExecutor.is_shutdown, so /health stays 200 after thread death;
  (c) the escaping exception is AssertionError, not CppExecutorError,
  so the SIGINT self-terminate path never fires.

- V2 (response delivery on thread death): traces the 8-step path from
  _prepare_tp_inputs raising through _executor_loop_cleanup. Shows the
  in-flight batch DOES receive a 5xx today (the per-request channel via
  await_response_thread works), but every request submitted after the
  thread dies hangs forever because enqueue_requests has no liveness
  check, _await_single_response has no is_shutdown short-circuit, and
  /health still reports 200. Two corollaries: B3 only needs to use the
  existing channel correctly (not build one), and B4 must close both
  the wrong-flag hole in _check_health and the no-check hole in
  /v1/models.

- V3 (C++ analog for Track A1): cpp/tensorrt_llm/batch_manager/
  microBatchScheduler.cpp has the same reuse_adjusted_compute helper
  (L33-44) and three accumulator sites with the same PR NVIDIA#12806 defect
  (L66-68 in fitDraftTokens, L358-378/L385-386 in operator() chunked
  and non-chunked branches, L459-461 in the post-chunk-adjust loop).
  No C++ analog of the prep-time total_num_tokens > max_num_tokens
  assertion exists; the Python _prepare_tp_inputs assertion is the
  first place the overshoot is observed. Confirms B3's
  BatchAdmissionError has a single Python producer and needs no
  nanobind binding work.

Updates existing tracks:
- A1 now cites the exact C++ patch sites instead of "verify file/line
  during implementation," and adds an assumption note about
  pre-subtracting context_chunk_size for non-first-chunk requests.
- A2 corrects the _reuse_adjusted_compute call-site count from 6 to 9
  with refreshed line numbers.
- B-track gains an "Architectural note (per V2)" calling out the
  over-broad requests=None fan-out in today's _handle_errors call.
- B1 is reframed as "single edit sufficient to prevent the customer's
  exact wedge; ship independently if helpful," reflecting V2's finding
  that the assertion is the proximate cause.
- B2 keeps the _fail_all_inflight design but adds a preferred option:
  route the wrapper's exception into _error_queue to fold this failure
  mode into PR NVIDIA#12718's existing classification path. Adds a sync-
  safety note about response_cv / enqueue_lock re-entrancy.
- B3 is now concrete with a BatchAdmissionError patch shape for both
  model_engine.py and the forward-step except, explicitly fixing the
  over-broad fan-out today.
- B4 is split into the two distinct holes (_check_health wrong-flag,
  /v1/models no-check), each with its concrete patch.
- Sequencing updated: B1 can ship alone; B1+B2+B4 should ship together.
- Test Plan updated to validate the new findings.

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
chienchunhung added a commit to chienchunhung/TensorRT-LLM that referenced this pull request May 20, 2026
@chienchunhung
[None][doc] align issue-13318 README with verified main state
e1a932a 
Follow-up to c8b91b5 (fix-plan.md update). Brings the README into
consistency with the verified findings in fix-plan.md §V1/V2/V3.

- Executive summary: reframe "do not close this case" — both NVIDIA#12718 and
  NVIDIA#13119 are now on main; the design-gap argument still holds. Point at
  fix-plan.md §V1 for the verified primitive names.

- "Why the Mismatch Bricks the Server" failure-flow box: rewrite per
  fix-plan.md §V2. The in-flight batch DOES get a 5xx today via the
  existing per-request error channel (_handle_errors enqueues responses;
  await_response_thread drains them). The customer-visible wedge is
  steps 6-7: requests submitted after the thread dies hang forever
  (enqueue_requests has no liveness check; _await_single_response has
  no is_shutdown short-circuit), and both /health (wrong-flag bug) and
  /v1/models (no-check bug) continue to report 200. Updated line numbers
  to current main (2668, 1491/2444, 652, 3411).

- "Why Prior Hardening Did Not Catch It > Timing": note both PRs are now
  on main; the design-gap analysis below applies to a fresh rebuild.

- "Why Prior Hardening Did Not Catch It > Design gaps": rewrite to use
  the actual primitive names that landed in NVIDIA#12718 (_error_queue,
  _handle_background_error, ManagedThread, signal.raise_signal(SIGINT)
  on CppExecutorError) instead of the placeholder names from the
  original investigation (_fatal_error, ErrorBudget). Reorder so that
  point (2) — PyExecutor's event loop is a plain threading.Thread, not
  a ManagedThread — is named as the root architectural fact, with (1),
  (3), (4) called out as downstream consequences. Point (1) is
  additionally an independent observability hole (LLM._check_health
  reads Executor.doing_shutdown, not PyExecutor.is_shutdown).

- "Why NVIDIA#12806 Is Not on main": correct the Python call-site count for
  _reuse_adjusted_compute from 6 to 9 with refreshed line numbers
  (L441/L472/L545/L676/L677/L720/L734/L746/L779). Add explicit pointer
  to cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp as the C++
  surface the customer's failing path actually runs.

- "Affected Components": update all line numbers to current main; add
  two new rows (py_executor.py:698-700 for the plain-Thread root cause,
  llmapi/llm.py:990 for the wrong-flag bug); cite the exact C++ file
  path and the three accumulator sites named in fix-plan.md §V3.

- "Related PRs > Hardening adjacent": note both PRs are merged; tighten
  the NVIDIA#12718 entry to point at the verified primitives and the
  plain-Thread reason it doesn't catch this bug.

- "Fix Plan" summary: align with the updated fix-plan.md sequencing.
  B1 (remove bare assertions) is now called out as the smallest patch
  that closes the customer's exact wedge. B2 is reframed as routing
  the wrapper's exception into _error_queue (preferred option from
  the updated fix-plan.md). B3 calls out the over-broad requests=None
  fan-out as a separate fix in scope.

- "Cross-References": correct primitive names in the nvbug-6043291
  entry to match what actually landed in NVIDIA#12718.

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
@thorjohnsen thorjohnsen mentioned this pull request Jun 10, 2026
Open
5 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

@liji-nv liji-nv

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@liji-nv