[#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

tensorrt-cicd · 2026-07-31T17:11:10Z

PR_Github #63075 [ kill ] triggered by Bot. Commit: db2bd31 Link to invocation

tensorrt-cicd · 2026-07-31T17:14:52Z

PR_Github #63073 [ run ] completed with state ABORTED. Commit: db2bd31

Link to invocation

tensorrt-cicd · 2026-07-31T17:15:03Z

PR_Github #63075 [ kill ] completed with state SUCCESS. Commit: db2bd31
Successfully killed previous jobs for commit db2bd31

Link to invocation

…atch Driving the token-budget fallback from _prepare_and_schedule_batch gave that method a new dependency on self.resource_manager, but unit tests exercise it on partially-constructed executors built with object.__new__(PyExecutor) that set only the attributes under test. Six cases in tests/unittest/_torch/executor/test_benchmark_disagg.py raised AttributeError: 'PyExecutor' object has no attribute 'resource_manager'. Guard the lookup the same way model_engine is guarded a few lines above, and for the same reason. A real executor always has a resource manager -- __init__ assigns it unconditionally -- so this only affects the partially-constructed executors the tests build. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

thorjohnsen · 2026-07-31T17:19:59Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-31T17:25:54Z

PR_Github #63079 [ run ] triggered by Bot. Commit: 1a2942c Link to invocation

tensorrt-cicd · 2026-07-31T22:03:23Z

PR_Github #63079 [ run ] completed with state FAILURE. Commit: 1a2942c
/LLM/main/L0_MergeRequest_PR pipeline #51173 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-01T00:46:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-01T00:53:16Z

PR_Github #63161 [ run ] triggered by Bot. Commit: 1a2942c Link to invocation

tensorrt-cicd · 2026-08-01T04:02:25Z

PR_Github #63161 [ run ] completed with state FAILURE. Commit: 1a2942c
/LLM/main/L0_MergeRequest_PR pipeline #51245 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-03T15:47:11Z

/bot run --disable-fail-fast

thorjohnsen · 2026-08-04T15:22:18Z

/bot run --disable-fail-fast

thorjohnsen · 2026-08-04T15:25:18Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-04T15:28:03Z

PR_Github #63781 [ run ] triggered by Bot. Commit: 1e3f516 Link to invocation

tensorrt-cicd · 2026-08-04T15:31:44Z

PR_Github #63782 [ run ] triggered by Bot. Commit: 1e3f516 Link to invocation

tensorrt-cicd · 2026-08-04T15:31:55Z

PR_Github #63781 [ run ] completed with state ABORTED. Commit: 1e3f516

Link to invocation

tensorrt-cicd · 2026-08-04T19:52:44Z

PR_Github #63782 [ run ] completed with state SUCCESS. Commit: 1e3f516
/LLM/main/L0_MergeRequest_PR pipeline #51731 completed with status: 'FAILURE'

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

CI Agent Failure Analysis

Link to invocation

The token-budget guard ran before ResourceManager.prepare_resources, where context_chunk_size still spans the reusable KV prefix and is not a forward-pass token count. Reading it as one charges a request for tokens the forward pass never computes: measured on H100, a request with chunk=19212 and estimated_reusable_tokens=19200 has a true cost of 12 tokens but was charged 19212 against an 8192 budget, so the guard deferred requests that already fit. That divergence is not knowable any earlier. The micro-batch scheduler admits a batch on estimated_reusable_tokens, a radix-tree guess made during capacity scheduling; the real figure is prepopulated_prompt_len, computed later inside addSequence. This issue is precisely the case where the two disagree. Move the trim to the end of ResourceManager.prepare_resources, after every manager has prepared, where context_current_position and context_chunk_size are final and the cost model is exact. The trim is shrink-only: blocks are allocated for the whole prompt rather than for the chunk, so trimming a chunk changes no block accounting and the tokens simply move to the next iteration. Nothing leaves the batch, so the attention-DP _can_queue vote, the pipeline-parallel inflight set and every earlier manager's per-request state stay consistent. Validated on H100 at this base: - test_token_budget_fallback.py: 23 passed - tests/unittest/_torch/executor: 1357 passed - shared-prefix reuse repro: keep=15, rechunk=0, defer=0 - defaults repro: keep=49 then keep=97, defer=0 - fault-injection A/B: with the trim, max forward tokens 2048 == budget and all requests complete; without it, total_num_tokens (3606) > max_num_tokens (2048) trips the _prepare_tp_inputs assert and kills the executor loop. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

coderabbitai · 2026-08-06T18:40:46Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

thorjohnsen · 2026-08-06T18:41:05Z

/bot run

coderabbitai

🧹 Nitpick comments (2)

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

397-402: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Update the stale comment for the renamed method and the new behavior.

The comment names _fit_token_budget and states the fallback is "defer instead of re-chunk". The method is now fit_token_budget (Line 835) and it only shrinks chunks; it never defers.

📝 Proposed comment fix

-        # Whether chunked prefill is enabled for this engine. Gates the re-chunk
-        # path in _fit_token_budget: a context request may only be shrunk into a
-        # partial chunk when the attention backend is set up for chunked context.
-        # Defaults to False (safe: defer instead of re-chunk) and is set to the
-        # finalized value by _create_kv_cache_manager.
+        # Whether chunked prefill is enabled for this engine. Gates the shrink
+        # path in fit_token_budget: a context request may only be shrunk into a
+        # partial chunk when the attention backend is set up for chunked context.
+        # Defaults to False (safe: skip the shrink) and is set to the finalized
+        # value by _create_kv_cache_manager.

🤖 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 397 - 402,
Update the comment above enable_chunked_prefill to reference fit_token_budget
instead of _fit_token_budget and describe that this flag gates shrinking a
context request into a partial chunk. Remove the inaccurate “defer instead of
re-chunk” fallback wording while preserving the explanation that the finalized
value is set by _create_kv_cache_manager.

tensorrt_llm/_torch/pyexecutor/_util.py (1)

2438-2447: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Pass enable_chunked_prefill through the V1 manager constructor.

Add it to manager_extra_kwargs for KVCacheManager subclasses. Use False when model_engine is None for the one-model draft path. PyTorchModelEngine initializes attn_runtime_features for target and draft engines, so the guarded access is safe during estimation and draft-manager creation.
🤖 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/_util.py` around lines 2438 - 2447, Update the
KVCacheManager construction flow to pass enable_chunked_prefill through
manager_extra_kwargs for KVCacheManager subclasses, using
model_engine.attn_runtime_features.chunked_prefill when available and False when
model_engine is None. Remove the separate post-construction assignment and
preserve the finalized target/draft engine value during manager creation.

🤖 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/_util.py`:
- Around line 2438-2447: Update the KVCacheManager construction flow to pass
enable_chunked_prefill through manager_extra_kwargs for KVCacheManager
subclasses, using model_engine.attn_runtime_features.chunked_prefill when
available and False when model_engine is None. Remove the separate
post-construction assignment and preserve the finalized target/draft engine
value during manager creation.

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 397-402: Update the comment above enable_chunked_prefill to
reference fit_token_budget instead of _fit_token_budget and describe that this
flag gates shrinking a context request into a partial chunk. Remove the
inaccurate “defer instead of re-chunk” fallback wording while preserving the
explanation that the finalized value is set by _create_kv_cache_manager.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6ad4df2a-3dca-49c4-9d9e-1e13093eb8aa

📥 Commits

Reviewing files that changed from the base of the PR and between 0448fe9 and dfd0b6e.

📒 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

The comment still described the pre-rebase behaviour: it named the method _fit_token_budget (now fit_token_budget) and said the flag defaults to deferring instead of re-chunking. The trim no longer defers anything -- it shrinks context chunks and never changes batch membership -- so with chunked prefill disabled the chunk is simply left at its scheduled size. Comment only; no behaviour change. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

thorjohnsen · 2026-08-06T19:16:50Z

/bot run

coderabbitai · 2026-08-06T19:56:18Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

thorjohnsen · 2026-08-06T19:57:00Z

/bot run --disable-fail-fast

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

coderabbitai · 2026-08-06T20:01:11Z

+    def test_disagg_gen_init_requests_are_left_alone(self):
+        # They only allocate/transfer KV cache and contribute no compute tokens.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        disagg = _FakeRequest(context_chunk_size=4096, is_disagg_generation_init_state=True)
+        ctx = _FakeRequest(context_chunk_size=16, prompt_len=16)
+        batch = _make_batch([disagg, ctx], [])
+
+        mgr.fit_token_budget(batch)
+
+        self.assertEqual(disagg.context_chunk_size, 4096)
+        self.assertEqual(ctx.context_chunk_size, 16)


🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Strengthen the disagg test so it discriminates the exclusion logic.

This test passes through the early return, not through the disagg guards. The totals are 16 tokens against a 128-token budget, so fit_token_budget returns at the excess <= 0 check and never reaches the loop.

Trace the two guards it is meant to cover:

If the not req.is_disagg_generation_init_state filter in the cost sum were removed, total would become 4112 and excess 3984.

The loop would then visit ctx first, which cannot shrink (a 16-token chunk at position 0 is already one block), and skip disagg at the in-loop guard.

Both assertions would still hold. The test cannot fail for either defect.

Add a case with a real overshoot so the disagg request is reachable in the loop and a shrinkable peer exists.

As per path instructions, tests/**: "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM" and "Keep feedback actionable".

💚 Proposed additional test

def test_disagg_gen_init_requests_are_left_alone(self): # They only allocate/transfer KV cache and contribute no compute tokens. mgr = _make_manager(max_num_tokens=128, tokens_per_block=16) disagg = _FakeRequest(context_chunk_size=4096, is_disagg_generation_init_state=True) ctx = _FakeRequest(context_chunk_size=16, prompt_len=16) batch = _make_batch([disagg, ctx], []) mgr.fit_token_budget(batch) self.assertEqual(disagg.context_chunk_size, 4096) self.assertEqual(ctx.context_chunk_size, 16) + + def test_disagg_gen_init_is_neither_costed_nor_shrunk_under_overshoot(self): + # A real overshoot, so the loop is reached. The disagg chunk must stay + # out of the cost sum and must never be shrunk, even though it is the + # largest shrinkable-looking chunk in the batch. + mgr = _make_manager(max_num_tokens=128, tokens_per_block=16) + disagg = _FakeRequest( + context_chunk_size=4096, prompt_len=4096, is_disagg_generation_init_state=True + ) + ctx = _FakeRequest(context_chunk_size=64, prompt_len=64) + gen = _FakeRequest(py_beam_width=100) # leaves a 28-token budget + batch = _make_batch([disagg, ctx], [gen]) + + # The disagg chunk contributes nothing: 100 gen + 64 ctx, not 4260. + self.assertEqual(_forward_tokens(mgr, batch), 164) + + mgr.fit_token_budget(batch) + + self.assertEqual(disagg.context_chunk_size, 4096, "disagg must never be shrunk") + self.assertEqual(ctx.context_chunk_size, 16, "the real context chunk absorbs the excess") + self.assertEqual(batch.num_context_requests, 2) + self.assertIn(disagg, batch.context_requests)

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

Source: Path instructions

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

thorjohnsen requested a review from a team as a code owner June 10, 2026 00:44

github-actions Bot assigned thorjohnsen Jun 10, 2026

thorjohnsen marked this pull request as draft June 10, 2026 00:45

coderabbitai Bot reviewed Jun 10, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/resource_manager.py Outdated

Comment thread tensorrt_llm/_torch/pyexecutor/resource_manager.py Outdated

Comment thread tests/unittest/_torch/executor/test_token_budget_fallback.py Outdated

Merge branch 'main' into fix/token-budget-prep-fallback

1e3f516

coderabbitai Bot reviewed Aug 6, 2026

View reviewed changes

Merge branch 'main' into fix/token-budget-prep-fallback

0cf9cdc

coderabbitai Bot reviewed Aug 6, 2026

View reviewed changes

Uh oh!

Conversation

thorjohnsen commented Jun 10, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

What goes wrong

What this PR adds

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

The KV connector is told the trimmed batch

How this behaves under each parallelism mode

What this PR deliberately does not do

Test Coverage

Validation

Dev Engineer Review

QA Engineer Review

Uh oh!

thorjohnsen commented Jun 10, 2026

Uh oh!

coderabbitai Bot commented Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (2 warnings)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

thorjohnsen commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

thorjohnsen commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

thorjohnsen commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

thorjohnsen commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 12, 2026

Uh oh!

thorjohnsen commented Jun 12, 2026

Uh oh!

tensorrt-cicd commented Jun 12, 2026

Uh oh!

tensorrt-cicd commented Jun 13, 2026

Uh oh!

thorjohnsen commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jul 31, 2026

Uh oh!

tensorrt-cicd commented Jul 31, 2026

Uh oh!

tensorrt-cicd commented Jul 31, 2026

Uh oh!

thorjohnsen commented Jul 31, 2026

Uh oh!

tensorrt-cicd commented Jul 31, 2026

Uh oh!

thorjohnsen commented Jun 10, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jun 10, 2026 •

edited

Loading