[Issue-13318][fix] Trim context chunks in _prepare_tp_inputs to prevent over-admission crash

[yifjiang]

Summary

Fixes #13318. Adds a Python-side post-scheduler guard in _prepare_tp_inputs that pre-computes the predicted position_ids growth and trims the trailing context chunk(s) if the batch would overshoot max_num_tokens. Mitigates the event-loop-killing assert total_num_tokens <= self.max_num_tokens crash observed in production with enable_chunked_prefill: true + enable_block_reuse: true.

Root cause (quoting the issue)

• MicroBatchScheduler (V1, C++) admits batches using a reuse-discounted token budget. For a last-chunk context request with reusable > 0 it charges only max(0, context_remaining - reusable) tokens — see _reuse_adjusted_compute in tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py:321 and its C++ counterpart.
• _prepare_tp_inputs (this file, around line 2276) extends position_ids by the full context_chunk_size per context request regardless of reuse: the reusable prefix's position embeddings still need to be computed in the forward pass.
• When max_num_tokens is tight, the admission-vs-forward mismatch lets len(position_ids) overshoot max_num_tokens by up to min(chunk_size, reusable) per offending request. The assertion at model_engine.py:2678 fires, the event-loop thread dies, and the server hangs until docker kill + restart (issue [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 has the full reproduction).

The admission math in _reuse_adjusted_compute (and its C++ twin) is still wrong after this fix — it subtracts reusable tokens that the forward pass has to process anyway. A proper root-cause fix should align the admission budget with the actual forward-pass cost. Until that lands, this guard keeps the event loop alive and lets the server degrade gracefully instead of crashing.

What the fix does

At the top of _prepare_tp_inputs (before input_ids = []), after the _apply_incremental_update early-return:

1. Compute predicted_ctx_tokens = sum(req.context_chunk_size for req in context_requests) — exact match for what the context loop will extend onto position_ids.
2. Compute predicted_gen_tokens = sum(1 + max(get_draft_token_length(req), self.runtime_draft_len) for req in generation_requests) — conservative upper bound covering both extend-request and plain-gen / first-draft paths.
3. If predicted_total > self.max_num_tokens, iterate the context requests in reverse and shrink req.context_chunk_size until the overshoot is cleared (leaving each chunk at least 1 token).
4. If the overshoot cannot be fully cleared (pathological config), log an error and let the original assertion fire so the condition stays observable.

Why mutating context_chunk_size mid-pipeline is safe:

• isLastContextChunk() (cpp/include/tensorrt_llm/batch_manager/llmRequest.h:1636) is live-derived from getContextCurrentPosition() + getContextChunkSize() == mPromptLen, so shrinking the chunk automatically de-promotes a "last" chunk to a "middle" chunk for this iteration. The next scheduler tick picks up the remaining tokens through the normal chunked-prefill progression.
• context_chunk_size is a read/write property in the nanobind binding (cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp:128), mapped to setContextChunkSize which clamps to getContextRemainingLength() and rejects negatives.
• KV cache already reserved by prepare_resources for the original chunk size stays in place and is reclaimed when the request finishes.
• The cost is O(N_requests) per iteration and only fires when overshoot is detected.

Relationship to PR #12806

PR #12806 fixes the analogous bug on feat/bench_y, which carries a Python-side remaining_budget re-check in prepare_resources that doesn't exist on main. The fix here applies the same intent ("pre-account committed forward-pass cost; shrink admissions to fit") but at the layer that exists on main — the pre-loop guard in _prepare_tp_inputs. Once the C++ admission math is aligned with forward-pass cost (or a Python-side pre-scheduler budget lands on main), this guard becomes a redundant backstop and can be removed.

Test plan

• Run the reproduction from issue [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 (Mistral-Nemo-12B replay at qps=7.0 with max_num_tokens=4096, enable_chunked_prefill=true, enable_partial_reuse=true, enable_block_reuse=true): verify no AssertionError, no event-loop hang, [issue-13318] trimming context chunk warnings appear in the server log, and traffic continues to be served.
• Verify max_num_tokens=8192 baseline (clean across 10k+ requests per issue's matrix) is unchanged.
• Unit test: construct a ScheduledRequests batch whose summed context_chunk_size + gen_tokens exceeds a mocked max_num_tokens; assert the trailing context_chunk_size is reduced to fit and a warning is logged.
• No-regression check on non-reuse / non-chunked configs: the guard only activates when predicted_total > max_num_tokens, which should not happen under a correct scheduler, so steady-state behavior should be unchanged.

Summary by CodeRabbit

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds a pre-forward admission guard in _prepare_tp_inputs to prevent position_ids length from exceeding self.max_num_tokens. Computes predicted token total and conditionally trims trailing context requests' chunk sizes while logging warnings or errors.

Changes

Cohort / File(s)	Summary
Admission Guard Logic tensorrt_llm/_torch/pyexecutor/model_engine.py	Adds pre-forward validation in _prepare_tp_inputs that predicts total token count using context chunk sizes and generation draft lengths. Trims trailing context request chunk sizes if prediction exceeds max_num_tokens, logs per-request warnings, and emits error if trimming cannot fully resolve overshoot.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: adding a guard to prevent over-admission crashes by trimming context chunks. It references the issue, specifies the fix type, and clearly conveys the primary objective.
Description check	✅ Passed	The description provides comprehensive context covering the root cause, fix rationale, safety justifications, and test plan, though it lacks explicit checkboxes marked and uses narrative format instead of strictly following the template structure.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2270-2279: The current predicted_gen_tokens sum undercounts cases
where a request will append more position_ids than 1 +
max(get_draft_token_length(req), self.runtime_draft_len) (specifically: plain
generation with beam_width > 1 and py_is_first_draft requests that use
original_max_draft_len). Update the aggregation that computes
predicted_gen_tokens (over scheduled_requests.generation_requests) to use the
actual generation-side accounting per request: for each req compute
tokens_to_account = 1 + max(get_draft_token_length(req), self.runtime_draft_len,
(self.original_max_draft_len if req.py_is_first_draft else 0), (req.beam_width
if getattr(req, "beam_width", 1) > 1 else 0)) and sum those values; reference
get_draft_token_length, self.runtime_draft_len, req.py_is_first_draft,
self.original_max_draft_len, and req.beam_width when making the change so
trimming matches the generation-side logic.

🪄 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: Pro Plus

Run ID: 34c01a53-7e34-42b4-aef1-f32ba0562539

📥 Commits

Reviewing files that changed from the base of the PR and between 7a8bd87 and 9731386.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/pyexecutor/model_engine.py

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Mirror the actual generation-side token accounting here.

The formula on Line 2278 still undercounts batches that later append more position_ids than 1 + max(...): plain generation with beam_width > 1 adds one token per beam at Lines 2641-2694, and py_is_first_draft requests add 1 + self.original_max_draft_len at Lines 2574-2616. In those cases this guard can skip trimming and the assertion at Lines 2746-2748 still remains reachable.

Suggested direction

-            predicted_gen_tokens = sum(
-                1 + max(get_draft_token_length(req), self.runtime_draft_len)
-                for req in scheduled_requests.generation_requests)
+            predicted_gen_tokens = 0
+            for req in scheduled_requests.generation_requests:
+                if get_draft_token_length(req) > 0 or next_draft_tokens_device is not None:
+                    predicted_gen_tokens += 1 + max(
+                        get_draft_token_length(req), self.runtime_draft_len)
+                elif req.py_is_first_draft:
+                    predicted_gen_tokens += 1 + self.original_max_draft_len
+                else:
+                    predicted_gen_tokens += req.sampling_config.beam_width

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 2270 - 2279, The
current predicted_gen_tokens sum undercounts cases where a request will append
more position_ids than 1 + max(get_draft_token_length(req),
self.runtime_draft_len) (specifically: plain generation with beam_width > 1 and
py_is_first_draft requests that use original_max_draft_len). Update the
aggregation that computes predicted_gen_tokens (over
scheduled_requests.generation_requests) to use the actual generation-side
accounting per request: for each req compute tokens_to_account = 1 +
max(get_draft_token_length(req), self.runtime_draft_len,
(self.original_max_draft_len if req.py_is_first_draft else 0), (req.beam_width
if getattr(req, "beam_width", 1) > 1 else 0)) and sum those values; reference
get_draft_token_length, self.runtime_draft_len, req.py_is_first_draft,
self.original_max_draft_len, and req.beam_width when making the change so
trimming matches the generation-side logic.

[yifjiang]

Closing in favor of porting the prepare_resources remaining_budget re-validation from feat/bench_y to main, then applying #12806's accounting fix on top. That keeps the fix at the same architectural layer the upstream solution was designed for, rather than introducing a parallel post-scheduler guard in _prepare_tp_inputs. Production unblock for the immediate qwen-coder deployment will be handled via a downstream-only image patch, not via this PR.