[#13318][fix] Gracefully fit token budget at prep boundary by thorjohnsen · Pull Request #15187 · NVIDIA/TensorRT-LLM

thorjohnsen · 2026-06-10T00:44:12Z

Description

What goes wrong

The micro-batch scheduler admits a batch against an estimate of how much KV
cache each context request will reuse. microBatchScheduler.cpp charges
reuse_adjusted_compute(chunk, estimated_reusable_tokens, remaining), where
estimated_reusable_tokens is a radix-tree guess made during capacity
scheduling. The number that actually governs the forward pass is
prepopulated_prompt_len, and that is not computed until addSequence runs
inside resource preparation — after the batch has been admitted.

When the real reuse comes in lower than the guess, the forward pass materializes
more tokens than were ever charged, and _prepare_tp_inputs trips

total_num_tokens (6986) should be less than or equal to max_num_tokens (4096)

On main that assert is caught by the generic handler in _forward_step and
routed to _handle_errors, which fails every request in the batch and
charges the executor's error budget. The user-visible symptom is intermittent
request failures under load and, if the pattern repeats enough to exhaust the
error budget, a fatal shutdown. The reporter saw it under multi-turn chat replay
at ~7 QPS, with overshoots from +1 to +2890 tokens.

What this PR adds

Graceful handling of the case where the token budget is genuinely exceeded
because the scheduler mis-estimated reuse. This is not a tightening of the
estimate and it does not try to predict the divergence — the whole point is that
the estimate was already wrong and only addSequence knows by how much. Instead
the batch is re-measured once the true numbers exist, and any overshoot is
absorbed by shrinking context chunks so the forward pass cannot exceed
max_num_tokens.

Concretely, KVCacheManager.fit_token_budget runs at the end of
ResourceManager.prepare_resources, after every resource manager has
prepared. At that point context_current_position and context_chunk_size are
final, so _request_forward_tokens is an exact count of the position ids each
request will contribute, mirroring _prepare_tp_inputs:

context: min(context_chunk_size, context_remaining_length), plus draft
tokens on the last chunk only;
generation: beam_width * (1 + draft_len);
disagg generation-init requests contribute no compute tokens and are excluded.

If the total exceeds max_num_tokens, context chunks are shrunk from the back
until it fits. Trimming runs back-to-front because that is where the overshoot
comes from: only a last chunk carries a reuse discount and draft tokens, so it
is the request whose cost the scheduler can have under-charged. Shrinking it
converts it back into a chunking request, which is exactly the repair;
mid-prefill chunks are touched only if that is not enough.

Why running here is the only correct point. An earlier revision of this PR
re-validated the budget before allocation, from the executor loop. That cannot
work, and measurably did not: before setPrepopulatedPromptLen runs,
context_chunk_size still spans the reusable prefix and is not a token count at
all. Reading it as one charges a request for tokens the forward pass never
computes. Measured on an H100 with the earlier revision applied: a request with
context_chunk_size = 19212 and estimated_reusable_tokens = 19200 has a true
forward cost of 12 tokens but was charged 19212 against an 8192 budget —
so the guard deferred requests that already fit, and did so repeatedly.

Why shrinking is safe, and why nothing needs to be deferred

Shrinking is close to free and — importantly — it never changes the membership
of the batch:

No block accounting changes. KV blocks are allocated for the whole prompt,
not for the chunk (_collect_context_sequences sizes add_sequence_batch
from prompt_len), so trimming a chunk moves compute to the next iteration
and touches nothing the KV cache manager has already done.
Chunk ends stay block-aligned. setPrepopulatedPromptLen asserts that
position + chunk lands on a block boundary for every non-last chunk, so the
new chunk is the largest block-aligned size that sheds the excess.
A chunk never shrinks to zero. A zero-token chunk would leave the request
scheduled but computing nothing, which never terminates; shrinking floors at
one block of forward progress.
Re-binning is explicit. Shrinking flips the computed
is_last_context_chunk to False, so the batch is re-binned via
reset_context_requests; otherwise downstream would still treat the request
as a final chunk and append generation/draft tokens to it.
Multimodal and disagg requests are skipped. Re-chunking across a
bidirectional multimodal block silently breaks attention (mirrors the gate in
scheduler_v2._align_chunk_to_mm_block), and disagg generation-init requests
have no compute tokens to shed.

Because no request leaves the batch, none of the invariants that a deferring
trim would break are even reachable: no manager's per-request state is orphaned,
no sequence is added without a matching batch entry, and no rank can shed its
way to an empty batch.

The KV connector is told the trimmed batch

build_scheduler_output ran at the end of KVCacheManager.prepare_resources,
i.e. before the trim, and handle_metadata() consumes its output afterwards.
So the connector was handed a SchedulerOutput describing the untrimmed batch.
RequestData.num_scheduled_tokens is documented as "the number of scheduled
tokens for the upcoming forward pass" and is built from context_chunk_size, so
every chunk the trim shrinks was over-reported — and a connector that decides
what to save or offload from that count would publish KV for tokens the forward
pass never computed.

The call moved into KVCacheManager.publish_connector_scheduler_output, driven
by ResourceManager.prepare_resources after maybe_fit_token_budget. The
hasattr gate keeps KVCacheManagerV2 (which defines neither hook) on its
existing path. The disagg generation-init path calls the KV cache manager's
prepare_resources directly on its own mini-batch and does not go through the
trim, so it publishes explicitly and its behaviour is unchanged.

Measured on H100 with a recording connector that moves no KV, over 12 compared
context requests: publishing after the trim gives 0 mismatches between what
the connector was told and what the forward pass computed; publishing before it
(the previous ordering) gives 1 mismatch — the connector told 1442 tokens for
a request the forward pass computed 512 on.

How this behaves under each parallelism mode

Single GPU (TP=1, no attention DP). The trim is a local, deterministic
function of the batch and the budget. No collective is involved. _can_queue
reduces to batch_size > 0 on the one rank.

Tensor parallel (attention DP off). Every TP rank schedules the same request
set and therefore holds the same batch, and fit_token_budget is a pure
function of (batch, max_num_tokens) with no rank-dependent input. All ranks
shrink identically, so the batches stay bit-identical going into the forward
pass. _can_queue is batch_size > 0 evaluated locally and identically —
there is no cross-rank vote to invalidate.

Attention DP. Each rank schedules its own batch, so the trim decision is
rank-local by construction. This is safe because the trim only changes
chunk sizes, never batch membership. The _can_queue vote — can_queue = 0 not in tp_allgather(scheduled_batch.batch_size) — is taken before
prepare_resources and gates on no rank being empty; since batch_size is
unchanged by the trim, that vote remains valid afterwards and no re-vote or
extra collective is needed. A rank that shrinks more than its peers simply
computes fewer tokens that iteration.

Pipeline parallel. _pp_schedule_and_propagate gives every PP rank the same
batch, so all ranks reach the same trim decision. PP needs one extra
consideration, which this PR fixes: _add_inflight_ids registers the batch's
last-chunk context requests, and _remove_inflight_ids previously re-derived
that set from the batch at removal time. The trim can move a request out of
context_requests_last_chunk in between, so the two views no longer agree.
_add_inflight_ids now records exactly what it inserted on
ScheduledRequests.added_inflight_req_ids, and _remove_inflight_ids erases
that snapshot. This is load-bearing, not defensive: the scheduler skips inflight
ids, so an id left behind is unrecoverable — the request would never be
scheduled again while still holding its KV blocks and sequence slot.

What this PR deliberately does not do

Error handling is unchanged from main. The _prepare_tp_inputs assert
stays a bare assert and no executor loop grows a new except branch. An earlier
revision converted it into a typed error routed to a server-terminating
shutdown behind an opt-out TorchLlmArgs flag; both were removed after review.
The flag did not restore pre-fallback behaviour when disabled (it selected a
third behaviour), and the fatal path called _handle_errors from inside the
loop, which performs collective gathers under attention DP — the deadlock
hazard _event_loop_wrapper explicitly documents and avoids, for a condition
that can be rank-local.
Nothing raises on the generation side any more. Two related cases:
a batch with no context requests returns immediately — context work is the
only thing this can shed, so scanning generation requests was dead work on the
executor loop's hottest path; and a batch where the generation requests alone
already exceed the budget is warned about rather than raised on, even when
context requests are present. The earlier revision raised a RuntimeError
there. That raise was rank-local under attention DP and unwound into
_event_loop_wrapper, killing one rank's loop thread while its peers waited
in a collective. Generation-only overshoot is a configuration property —
max_batch_size x beam_width x (1 + max_draft_len) is known at startup — and
is better validated there than discovered in the scheduling path; it currently
falls back to the existing assert, which fails one batch instead of the
server.
With chunked prefill disabled nothing is trimmed. A shrunk chunk is a
partial context chunk, which the attention backend is only set up to consume
under chunked prefill; forcing one produces an invalid forward pass. The
overshoot is logged with a clear warning and the pre-existing assert fires as
it does on main.
When shrinking cannot absorb the whole overshoot (every chunk already at
its one-block floor), the remainder is logged and the existing assert fires.
Handling that case requires removing work from the batch, which is a larger
change and is left to a follow-up.
KVCacheManagerV2 is untouched. The V2 scheduler already sizes each chunk
as min(remaining_budget, context_remaining) and sets context_chunk_size to
that same value inline, so it is structurally immune;
ResourceManager.maybe_fit_token_budget gates on
hasattr(kv_cache_manager, "maybe_fit_token_budget"), which V2 does not
define. The draft-model manager is skipped as well — it builds inputs with a
different token shape and its budget is handled separately.

There is no API change. This PR touches 5 files.

Note: a prior fix for this issue (PR #12806) was authored against the
feat/bench_y branch, which had a remaining_budget re-validation block in
prepare_resources that was never upstreamed to main. That block does not
exist on main, so #12806 cannot be cherry-picked — this PR reimplements the
intent against the current code.

Test Coverage

tests/unittest/_torch/executor/test_token_budget_fallback.py (new, GPU-free)
drives KVCacheManager.fit_token_budget directly with lightweight fakes —
23 tests. Highlights:

test_trim_runs_after_every_manager / test_prepare_resources_trims — the
central design claim: the trim is driven from the end of
ResourceManager.prepare_resources, after every manager has prepared.
TestReuseDiscountedChunk — the regression for reading the chunk too early: a
19212-token prompt with a 19200-token cache hit costs 12 forward tokens
(test_reuse_hit_costs_only_the_uncached_tail) and must not be trimmed
(test_reuse_hit_is_not_trimmed). This is the defect the earlier revision
shipped with.
test_within_budget_is_noop / test_untrimmed_batch_round_trips — a batch
within budget is left untouched.
test_overshoot_shrinks_context_to_fit — 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
generation batch plus an oversized last chunk, shrunk to a block-aligned fit.
test_shrink_keeps_chunk_end_block_aligned /
test_shrink_never_produces_a_zero_token_chunk — the two invariants above.
test_shrink_rebins_to_chunking / test_shrink_drops_last_chunk_draft_tokens
— a shrunk request stops being a last chunk and stops contributing drafts.
test_nothing_is_ever_dropped — shrink-only: membership never changes. This
is the property every parallelism argument above rests on.
test_sheds_the_last_chunk_first /
test_shrinks_multiple_requests_when_one_is_not_enough — back-to-front
ordering, and spilling onto earlier requests only when needed.
test_gen_only_batch_is_left_alone /
test_generation_alone_over_budget_does_not_raise — the hot-path early-out,
and that generation-only overshoot does not raise.
test_no_shrink_when_chunked_prefill_disabled /
test_mm_bidirectional_is_not_shrunk /
test_disagg_gen_init_requests_are_left_alone — the three cases that must not
be re-chunked.
TestInflightIdsSurviveTrim
(test_shrunk_context_requests_leave_no_inflight_ids) — a shrunk context
request leaves no id stranded in the PP inflight set.
test_maybe_fit_token_budget_skips_draft_manager — the draft manager's batch
is not trimmed.

Validation

Measured on H100 at this base:

check	result
`test_token_budget_fallback.py`	27 passed
`tests/unittest/_torch/executor/`	1369 passed, 0 failures
shared-prefix reuse workload	`keep=15, rechunk=0, defer=0` — no request is deferred or re-chunked when reuse is estimated correctly
defaults workload	`keep=49` then `keep=97`, `defer=0`
fault injection, trim enabled	max forward tokens 2048 == `max_num_tokens`, all requests complete
connector ordering A/B	0 mismatches; control (publish before the trim) 1 mismatch, told 1442 / computed 512
fault injection, trim disabled	`total_num_tokens (3606) > max_num_tokens (2048)` → assert → `Sampling failed` → executor loop dies (i.e. #13318 reproduced)

The fault-injection A/B is the direct evidence that the trim converts the
reported failure into a survivable, correctly-sized batch.

Dev Engineer Review

Added V1 token-budget revalidation in KVCacheManager.
Added block-aligned context trimming and request re-binning.
Preserved generation requests and handled multimodal and disaggregated requests.
Added ResourceManager.maybe_fit_token_budget.
Preserved inflight request IDs across batch trimming.
Propagated the finalized chunked-prefill setting to V1 managers.
V2 and draft-model managers remain unchanged.
No configuration or test-list files changed.
No API compatibility or dependency concerns were identified.

QA Engineer Review

Added tests/unittest/_torch/executor/test_token_budget_fallback.py.
Added TestReuseDiscountedChunk for cache reuse and chunk bounds.
Added TestFitTokenBudget for token accounting, budget fitting, request preservation, block-aligned shrinking, re-binning, draft-token handling, trimming order, disabled chunked prefill, multimodal and disaggregated requests, generation overflow, manager ordering, and resource preparation.
Added TestInflightIdsSurviveTrim for inflight-ID handling.
Added TestConnectorSeesTheTrimmedBatch for connector visibility after trimming.
These tests are not listed in tests/integration/test_lists/test-db/ or tests/integration/test_lists/.
The existing token-budget test-list entry covers a separate V2 scheduler test.
Verdict: needs follow-up.

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>

thorjohnsen · 2026-06-10T00:48:12Z

/bot run --disable-fail-fast

coderabbitai · 2026-06-10T00:51:02Z

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

Walkthrough

KVCacheManager now enforces post-allocation token budgets. It accounts for cache reuse, generation, draft tokens, multimodal boundaries, and block alignment. Executor scheduling preserves inflight request IDs and publishes connector output after trimming.

Changes

Token budget constraint and scheduling

Layer / File(s)	Summary
Token budget fitting and context re-chunking `tensorrt_llm/_torch/pyexecutor/resource_manager.py`	`KVCacheManager` stores chunked-prefill settings, computes forward-token usage, shrinks eligible context chunks, and publishes connector scheduling output after fitting. `ResourceManager.prepare_resources` runs fitting after all managers prepare resources.
Runtime configuration and inflight ID preservation `tensorrt_llm/_torch/pyexecutor/_util.py`, `tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py`, `tensorrt_llm/_torch/pyexecutor/py_executor.py`	V1 managers receive the finalized chunked-prefill setting. `ScheduledRequests` records exact inflight IDs. `PyExecutor` removes those IDs after resource preparation and publishes disaggregated generation-init output.
Token budget behavior and integration tests `tests/unittest/_torch/executor/test_token_budget_fallback.py`	Tests cover token accounting, cache reuse, block-aligned shrinking, request preservation, unsupported cases, generation-only batches, draft-manager bypassing, manager ordering, inflight ID cleanup, and connector publication.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PyExecutor
  participant ResourceManager
  participant KVCacheManager
  participant ScheduledRequests
  participant KVConnector
  PyExecutor->>ScheduledRequests: record added inflight request IDs
  PyExecutor->>ResourceManager: prepare scheduled batch
  ResourceManager->>KVCacheManager: prepare resources
  ResourceManager->>KVCacheManager: fit token budget
  KVCacheManager->>ScheduledRequests: compute and trim eligible context chunks
  ResourceManager->>KVConnector: publish trimmed scheduler output
  PyExecutor->>ScheduledRequests: remove recorded inflight IDs

Suggested labels: ci: full pre-merge approved

Suggested reviewers: qijune, allisonlim-nv

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Linked Issues check	⚠️ Warning	The PR fixes the token-accounting mismatch, but it leaves assertion-driven event-loop failures for disabled chunked prefill and unabsorbable overshoots required by `#13318`.	Add graceful handling for all remaining overshoot cases so assertions cannot terminate the event loop or make the server unresponsive under overload.
Docstring Coverage	⚠️ Warning	Docstring coverage is 9.43% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Out of Scope Changes check	✅ Passed	The implementation, connector ordering, inflight-ID tracking, configuration propagation, and tests directly support the token-budget fix and its documented invariants.
Title check	✅ Passed	The title clearly identifies the fix and its token-budget preparation behavior.
Description check	✅ Passed	The description explains the problem, solution, design constraints, tests, and validation, but omits the template's PR Checklist section.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

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

📥 Commits

Reviewing files that changed from the base of the PR and between 9a7f76f and db521b8.

📒 Files selected for processing (2)

tensorrt_llm/_torch/pyexecutor/resource_manager.py
tests/unittest/_torch/executor/test_token_budget_fallback.py

tensorrt-cicd · 2026-06-10T00:54:16Z

PR_Github #53182 [ run ] triggered by Bot. Commit: db521b8 Link to invocation

tensorrt-cicd · 2026-06-10T05:16:59Z

PR_Github #53182 [ run ] completed with state FAILURE. Commit: db521b8
/LLM/main/L0_MergeRequest_PR pipeline #42382 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

thorjohnsen · 2026-06-10T18:31:02Z

/bot run

tensorrt-cicd · 2026-06-10T18:37:38Z

PR_Github #53360 [ run ] triggered by Bot. Commit: db521b8 Link to invocation

…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>

thorjohnsen · 2026-06-10T18:53:03Z

/bot run

tensorrt-cicd · 2026-06-10T18:58:52Z

PR_Github #53365 [ run ] triggered by Bot. Commit: 389dffb Link to invocation

tensorrt-cicd · 2026-06-10T19:02:43Z

PR_Github #53360 [ run ] completed with state ABORTED. Commit: db521b8

Link to invocation

tensorrt-cicd · 2026-06-10T19:23:38Z

PR_Github #53365 [ run ] completed with state FAILURE. Commit: 389dffb
/LLM/main/L0_MergeRequest_PR pipeline #42544 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

thorjohnsen · 2026-06-10T19:27:59Z

/bot run

tensorrt-cicd · 2026-06-10T19:35:03Z

PR_Github #53367 [ run ] triggered by Bot. Commit: 389dffb Link to invocation

tensorrt-cicd · 2026-06-10T21:39:03Z

PR_Github #53367 [ run ] completed with state FAILURE. Commit: 389dffb
/LLM/main/L0_MergeRequest_PR pipeline #42545 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

thorjohnsen · 2026-06-11T21:43:34Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-11T21:50:02Z

PR_Github #53691 [ run ] triggered by Bot. Commit: 389dffb Link to invocation

tensorrt-cicd · 2026-06-12T01:49:02Z

PR_Github #53691 [ run ] completed with state SUCCESS. Commit: 389dffb
/LLM/main/L0_MergeRequest_PR pipeline #42826 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…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>

thorjohnsen · 2026-06-12T23:19:40Z

/bot run

tensorrt-cicd · 2026-06-12T23:25:06Z

PR_Github #53972 [ run ] triggered by Bot. Commit: 69c3786 Link to invocation

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>

tensorrt-cicd · 2026-06-13T01:57:20Z

PR_Github #53972 [ run ] completed with state FAILURE. Commit: 69c3786
/LLM/main/L0_MergeRequest_PR pipeline #43061 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…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>

thorjohnsen · 2026-06-15T14:58:57Z

/bot run

tensorrt-cicd · 2026-06-15T15:05:02Z

PR_Github #54315 [ run ] triggered by Bot. Commit: 1b0276c Link to invocation

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/_torch/executor/test_token_budget_fallback.py (1)
1-16: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a residual-overshoot assertion to test_nothing_is_ever_dropped.

Test coverage summary

23 test functions were added. None were modified or removed.

tests/unittest/_torch/executor/test_token_budget_fallback.py is covered by directory entries in l0_cpu.yml, l0_b300.yml, l0_h100.yml, l0_gb300_multi_gpus.yml, and l0_dgx_b300.yml.

Coverage is insufficient. After trimming, test_nothing_is_ever_dropped leaves 148 forward-pass tokens against a budget of 128. Assert this residual overshoot to protect the one-block floor behavior.
🤖 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 `@tests/unittest/_torch/executor/test_token_budget_fallback.py` around lines 1
- 16, Update test_nothing_is_ever_dropped to assert that the post-trim
forward-pass token total remains 148 against the 128-token budget, preserving
coverage of the one-block floor behavior and residual overshoot.
Source: Path instructions

🤖 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 `@tests/unittest/_torch/executor/test_token_budget_fallback.py`:
- Around line 330-340: Strengthen test_disagg_gen_init_requests_are_left_alone
by creating a genuine token-budget overshoot with a shrinkable non-disagg
request, ensuring fit_token_budget reaches its reduction loop. Keep the disagg
request large enough to expose accidental inclusion in the cost sum and verify
its context_chunk_size remains unchanged while the peer is reduced, covering
both the cost-sum filter and in-loop disagg guard.

---

Nitpick comments:
In `@tests/unittest/_torch/executor/test_token_budget_fallback.py`:
- Around line 1-16: Update test_nothing_is_ever_dropped to assert that the
post-trim forward-pass token total remains 148 against the 128-token budget,
preserving coverage of the one-block floor behavior and residual overshoot.

🪄 Autofix

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: 5f50d6c7-5cd1-4405-942c-a2b24ca4c0d4

📥 Commits

Reviewing files that changed from the base of the PR and between 2224cb7 and 0cf9cdc.

📒 Files selected for processing (5)

tensorrt_llm/_torch/pyexecutor/_util.py
tensorrt_llm/_torch/pyexecutor/py_executor.py
tensorrt_llm/_torch/pyexecutor/resource_manager.py
tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
tests/unittest/_torch/executor/test_token_budget_fallback.py

🚧 Files skipped from review as they are similar to previous changes (3)

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
tensorrt_llm/_torch/pyexecutor/_util.py
tensorrt_llm/_torch/pyexecutor/py_executor.py

build_scheduler_output ran at the end of KVCacheManager.prepare_resources, i.e. before the token-budget trim, and handle_metadata() consumes its output afterwards. So the connector was handed a SchedulerOutput describing the untrimmed batch. RequestData.num_scheduled_tokens is documented as "the number of scheduled tokens for the upcoming forward pass" and is built from context_chunk_size, so every chunk the trim shrinks was over-reported. A connector that decides what to save or offload from that count would publish KV for tokens the forward pass never computed. Move the call into KVCacheManager.publish_connector_scheduler_output, driven by ResourceManager.prepare_resources after maybe_fit_token_budget. The hasattr gate keeps KVCacheManagerV2 (which defines neither hook) on its existing path. The disagg generation-init path calls the KV cache manager's prepare_resources directly on its own mini-batch and does not go through the trim, so it publishes explicitly and its behaviour is unchanged. Measured on H100 with a recording connector that moves no KV, over 12 compared context requests: - publishing after the trim: 0 mismatches between what the connector was told and what the forward pass computed; - publishing before it (the previous ordering): 1 mismatch -- the connector was told 1442 tokens for a request the forward pass computed 512 on. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

tensorrt-cicd · 2026-08-06T22:41:53Z

PR_Github #64394 [ run ] triggered by Bot. Commit: b026f50 Link to invocation

tensorrt-cicd · 2026-08-07T03:14:39Z

PR_Github #64394 [ run ] completed with state FAILURE. Commit: b026f50
/LLM/main/L0_MergeRequest_PR pipeline #52282 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

thorjohnsen · 2026-08-07T13:51:11Z

/bot --disable-fail-fast

github-actions · 2026-08-07T13:51:21Z

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental) --high-priority]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Supports wildcard * for pattern matching (e.g., "*PerfSanity*" matches all stages containing PerfSanity). Examples: "A10-PyTorch-1, xxx", "PerfSanity". The patterns "*", "*Post-Merge*", and "*PerfSanity*", including equivalent escaped or repeated-star forms and their use in comma-separated lists, require the ci: post-merge approved PR label. Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Requires the ci: full pre-merge approved label on the PR (ask a member of NVIDIA/trt-llm-ci-approvers). Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline. Requires the ci: full pre-merge approved label on the PR (ask a member of NVIDIA/trt-llm-ci-approvers).

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline. Requires the ci: post-merge approved PR label applied by an active member of NVIDIA/trt-llm-ci-approvers. The approval label remains in place when new commits are pushed.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Supports wildcard * for pattern matching. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx", --extra-stage "Post-Merge". The patterns "*", "*Post-Merge*", and "*PerfSanity*", including equivalent escaped or repeated-star forms and their use in comma-separated lists, require the ci: post-merge approved PR label.

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

--high-priority (OPTIONAL) : Run the pipeline with high priority. This option is restricted to authorized users only and will route the job to a high-priority queue.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

thorjohnsen · 2026-08-07T14:23:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-07T14:29:23Z

PR_Github #64620 [ run ] triggered by Bot. Commit: b026f50 Link to invocation

tensorrt-cicd · 2026-08-07T15:14:54Z

PR_Github #64620 [ run ] completed with state SUCCESS. Commit: b026f50
/LLM/main/L0_MergeRequest_PR pipeline #52480 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

thorjohnsen · 2026-08-07T16:00:11Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-07T16:07:34Z

PR_Github #64641 [ run ] triggered by Bot. Commit: b026f50 Link to invocation

tensorrt-cicd · 2026-08-07T16:35:06Z

PR_Github #64641 [ run ] completed with state FAILURE. Commit: b026f50
/LLM/main/L0_MergeRequest_PR pipeline #52502 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

thorjohnsen · 2026-08-07T17:23:15Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-07T17:32:21Z

PR_Github #64661 [ run ] triggered by Bot. Commit: b026f50 Link to invocation

tensorrt-cicd · 2026-08-07T17:53:41Z

PR_Github #64661 [ run ] completed with state ABORTED. Commit: b026f50

Link to invocation

thorjohnsen · 2026-08-10T14:20:56Z

/bot run --disable-fail-fast

github-actions · 2026-08-10T14:21:11Z