From db521b8091b3704ab4067f66383ffae3e2e8a18e Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Wed, 10 Jun 2026 00:09:21 +0000
Subject: [PATCH 01/16] [None][fix] Gracefully fit token budget at prep
 boundary (#13318)

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>
---
 .../_torch/pyexecutor/resource_manager.py     | 107 ++++++++++++
 .../executor/test_token_budget_fallback.py    | 158 ++++++++++++++++++
 2 files changed, 265 insertions(+)
 create mode 100644 tests/unittest/_torch/executor/test_token_budget_fallback.py

diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index f4bc390060e3..f10682bdbbdd 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -655,6 +655,9 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
         # Import here to avoid circular imports
         from ..speculative import get_num_extra_kv_tokens
         self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config)
+        # Kept so prepare_resources can re-validate the per-step token budget
+        # (the forward-pass scratch size enforced in _prepare_tp_inputs).
+        self.max_num_tokens = max_num_tokens
         self.event_buffer_max_size = kv_cache_config.event_buffer_max_size
         self.attention_dp_events_gather_period_ms = kv_cache_config.attention_dp_events_gather_period_ms
         self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0
@@ -984,7 +987,111 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int:
             remaining_tokens / self.tokens_per_block)
         return need_blocks
 
+    @staticmethod
+    def _has_mm_bidirectional_block(req: LlmRequest) -> bool:
+        # Mirror the gate in scheduler_v2._align_chunk_to_mm_block: re-chunking
+        # a request whose boundary would split a bidirectional multimodal block
+        # silently breaks attention, so such requests are deferred whole rather
+        # than re-chunked.
+        mm = getattr(req, "py_multimodal_data", None)
+        return isinstance(mm, dict) and mm.get("mm_bidirectional_blocks", False)
+
+    def _request_forward_tokens(self, req: LlmRequest, *,
+                                is_context: bool) -> int:
+        """Upper bound on the number of position ids ``req`` contributes to a
+        forward pass in ``_prepare_tp_inputs``.
+
+        This MUST over-estimate. Under-counting would reintroduce the
+        ``total_num_tokens <= max_num_tokens`` assert in ``_prepare_tp_inputs``
+        that this guard exists to prevent.
+        """
+        draft_len = get_draft_token_length(req)
+        if is_context:
+            # Context contributes ``context_chunk_size`` positions; draft tokens
+            # are appended only on the last chunk.
+            return req.context_chunk_size + (draft_len if
+                                             req.is_last_context_chunk else 0)
+        # Generation: one position per beam for the new token, plus draft tokens
+        # (speculative verification) per beam.
+        return req.py_beam_width * (1 + draft_len)
+
+    def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
+        """Defer or re-chunk context requests so the scheduled batch cannot
+        exceed ``max_num_tokens`` in the forward pass.
+
+        The micro-batch scheduler's token-budget estimate can diverge from the
+        tokens actually materialized by ``_prepare_tp_inputs`` -- for example
+        when a reuse-discounted last context chunk lands next to a near-full
+        generation batch (see GitHub issue #13318). Rather than letting that
+        divergence trip a hard assert and wedge the executor loop, re-validate
+        the budget here -- before any KV cache is allocated -- and gracefully
+        shed only the deferrable work (context chunks), leaving in-flight
+        generation requests untouched.
+
+        Deferred context requests are simply dropped from this iteration's
+        ``scheduled_batch``; they remain in the active pool and are rescheduled
+        on a later iteration with a fresh budget.
+        """
+        budget = self.max_num_tokens
+
+        # Generation requests are in-flight and cannot be deferred. If they
+        # alone exceed the budget something is genuinely misconfigured -- fail
+        # this batch loudly rather than overshoot silently.
+        gen_tokens = sum(
+            self._request_forward_tokens(req, is_context=False)
+            for req in scheduled_batch.generation_requests)
+        if gen_tokens > budget:
+            raise RuntimeError(
+                f"In-flight generation requests need {gen_tokens} tokens, "
+                f"exceeding max_num_tokens ({budget}); cannot schedule.")
+
+        remaining = budget - gen_tokens
+        kept: RequestList = []
+        deferring = False
+        for req in scheduled_batch.context_requests:
+            # Disagg generation-init requests only allocate/transfer KV cache
+            # and contribute no compute tokens, so never shed them.
+            if deferring and not req.is_disagg_generation_init_state:
+                continue
+            cost = self._request_forward_tokens(req, is_context=True)
+            if cost <= remaining:
+                kept.append(req)
+                remaining -= cost
+                continue
+
+            # Doesn't fit. Try re-chunking the compute (fewer tokens this step)
+            # before deferring. Re-chunking only reduces compute tokens -- KV is
+            # allocated for the full prompt regardless -- so block accounting is
+            # unaffected. Only safe when the request can be chunked further, the
+            # shrunk chunk still holds at least one block (aligned to block
+            # size), and the boundary won't split a bidirectional multimodal
+            # block.
+            new_chunk = (remaining //
+                         self.tokens_per_block) * self.tokens_per_block
+            if (new_chunk >= self.tokens_per_block
+                    and new_chunk < req.context_chunk_size
+                    and not self._has_mm_bidirectional_block(req)):
+                req.context_chunk_size = new_chunk  # now a non-last chunk
+                kept.append(req)
+                remaining -= new_chunk
+            # remaining budget is now < one block, so no further context
+            # request can fit this iteration.
+            deferring = True
+
+        if len(kept) != scheduled_batch.num_context_requests:
+            logger.debug(
+                f"_fit_token_budget: kept {len(kept)}/"
+                f"{scheduled_batch.num_context_requests} context requests to "
+                f"stay within max_num_tokens={budget}")
+            # reset_context_requests re-derives chunking vs last-chunk from each
+            # request's (possibly updated) is_last_context_chunk.
+            scheduled_batch.reset_context_requests(kept)
+
     def prepare_resources(self, scheduled_batch: ScheduledRequests):
+        if not self.is_draft:
+            # The draft-model engine builds inputs with a different token shape;
+            # its budget is handled separately.
+            self._fit_token_budget(scheduled_batch)
         with request_context(self.is_draft, scheduled_batch):
             # wait for all pending work to finish before launching offload/onboarding/partial copy
             self.impl.sync_transfer_manager_with_buffer_manager()
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
new file mode 100644
index 000000000000..429d95f6e5cd
--- /dev/null
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Unit tests for KVCacheManager._fit_token_budget.
+
+These exercise the prep-boundary token-budget fallback that defers or
+re-chunks context requests so a scheduled batch cannot overshoot
+``max_num_tokens`` in the forward pass (GitHub issue #13318). The fallback is
+pure scheduling logic and does not touch the GPU, so the tests build a bare
+KVCacheManager via ``__new__`` and drive the method with lightweight fake
+requests.
+"""
+
+import unittest
+
+from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
+from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
+
+
+class _FakeRequest:
+    """Minimal stand-in exposing only the attributes _fit_token_budget reads."""
+
+    def __init__(
+        self,
+        *,
+        context_chunk_size=0,
+        is_last_context_chunk=True,
+        py_beam_width=1,
+        py_draft_tokens=None,
+        is_disagg_generation_init_state=False,
+        mm_bidirectional=False,
+    ):
+        self.context_chunk_size = context_chunk_size
+        self.is_last_context_chunk = is_last_context_chunk
+        self.py_beam_width = py_beam_width
+        self.py_draft_tokens = py_draft_tokens
+        self.is_disagg_generation_init_state = is_disagg_generation_init_state
+        self.py_multimodal_data = {"mm_bidirectional_blocks": True} if mm_bidirectional else None
+
+
+def _make_manager(max_num_tokens, tokens_per_block):
+    # Skip the heavy (GPU-allocating) __init__; the method under test only
+    # needs these two attributes plus its own (bound) helper methods.
+    mgr = KVCacheManager.__new__(KVCacheManager)
+    mgr.max_num_tokens = max_num_tokens
+    mgr.tokens_per_block = tokens_per_block
+    return mgr
+
+
+def _make_batch(context_requests=(), generation_requests=()):
+    batch = ScheduledRequests()
+    for req in context_requests:
+        batch.append_context_request(req)
+    batch.generation_requests = list(generation_requests)
+    return batch
+
+
+class TestFitTokenBudget(unittest.TestCase):
+    def test_request_forward_tokens_upper_bound(self):
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+
+        # Context: chunk size, plus draft tokens only on the last chunk.
+        last = _FakeRequest(
+            context_chunk_size=10, is_last_context_chunk=True, py_draft_tokens=[1, 2]
+        )
+        self.assertEqual(mgr._request_forward_tokens(last, is_context=True), 12)
+        mid = _FakeRequest(
+            context_chunk_size=10, is_last_context_chunk=False, py_draft_tokens=[1, 2]
+        )
+        self.assertEqual(mgr._request_forward_tokens(mid, is_context=True), 10)
+
+        # Generation: (1 + draft) per beam.
+        gen = _FakeRequest(py_beam_width=2, py_draft_tokens=[1, 2, 3])
+        self.assertEqual(mgr._request_forward_tokens(gen, is_context=False), 8)
+
+    def test_within_budget_is_noop(self):
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        ctx = _FakeRequest(context_chunk_size=16)
+        gen = _FakeRequest(py_beam_width=100)  # 100 gen tokens
+        batch = _make_batch([ctx], [gen])
+
+        mgr._fit_token_budget(batch)
+
+        self.assertEqual(batch.num_context_requests, 1)
+        self.assertEqual(ctx.context_chunk_size, 16)  # untouched
+
+    def test_overshoot_rechunks_context(self):
+        # 100 gen tokens leave a 28-token budget; a 64-token last chunk does not
+        # fit but can be re-chunked down to a block-aligned 16.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        ctx = _FakeRequest(context_chunk_size=64, is_last_context_chunk=True)
+        gen = _FakeRequest(py_beam_width=100)
+        batch = _make_batch([ctx], [gen])
+
+        mgr._fit_token_budget(batch)
+
+        self.assertEqual(batch.num_context_requests, 1)
+        self.assertEqual(ctx.context_chunk_size, 16)  # (28 // 16) * 16
+        total = mgr._request_forward_tokens(ctx, is_context=True) + mgr._request_forward_tokens(
+            gen, is_context=False
+        )
+        self.assertLessEqual(total, mgr.max_num_tokens)
+
+    def test_overshoot_defers_when_cannot_rechunk(self):
+        # Only an 8-token budget remains -- smaller than one block -- so the
+        # context request cannot be re-chunked and must be deferred entirely.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        ctx = _FakeRequest(context_chunk_size=64)
+        gen = _FakeRequest(py_beam_width=120)  # remaining = 8 < tokens_per_block
+        batch = _make_batch([ctx], [gen])
+
+        mgr._fit_token_budget(batch)
+
+        self.assertEqual(batch.num_context_requests, 0)
+        self.assertEqual(ctx.context_chunk_size, 64)  # not re-chunked
+
+    def test_mm_bidirectional_is_deferred_not_rechunked(self):
+        # A re-chunkable budget exists, but splitting a bidirectional MM block
+        # would corrupt attention, so the request is deferred whole.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        ctx = _FakeRequest(context_chunk_size=64, mm_bidirectional=True)
+        gen = _FakeRequest(py_beam_width=100)  # remaining = 28
+        batch = _make_batch([ctx], [gen])
+
+        mgr._fit_token_budget(batch)
+
+        self.assertEqual(batch.num_context_requests, 0)
+        self.assertEqual(ctx.context_chunk_size, 64)
+
+    def test_defers_all_subsequent_context_requests(self):
+        # ctx1 fits; ctx2 overshoots and cannot re-chunk; ctx3 (small) must
+        # still be deferred to preserve context-progress ordering.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        ctx1 = _FakeRequest(context_chunk_size=96)
+        ctx2 = _FakeRequest(context_chunk_size=64)
+        ctx3 = _FakeRequest(context_chunk_size=16)
+        gen = _FakeRequest(py_beam_width=16)  # remaining = 112
+        batch = _make_batch([ctx1, ctx2, ctx3], [gen])
+
+        mgr._fit_token_budget(batch)
+
+        # ctx1 (96) fits into 112; remaining 16. ctx2 (64) doesn't fit and
+        # (16 // 16) * 16 == 16 but 16 < 64 so it *could* re-chunk to 16...
+        # remaining after is 0, so ctx3 is deferred.
+        kept = batch.context_requests
+        self.assertIn(ctx1, kept)
+        self.assertNotIn(ctx3, kept)
+
+    def test_generation_alone_over_budget_raises(self):
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        gen = _FakeRequest(py_beam_width=200)
+        batch = _make_batch([], [gen])
+
+        with self.assertRaises(RuntimeError):
+            mgr._fit_token_budget(batch)
+
+
+if __name__ == "__main__":
+    unittest.main()

From 389dffb86b655ebf44271cc176ebfee447f60deb Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Wed, 10 Jun 2026 18:50:15 +0000
Subject: [PATCH 02/16] [#13318][fix] Re-bin re-chunked context requests in
 token-budget 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>
---
 .../_torch/pyexecutor/resource_manager.py     | 37 +++++++++++++----
 .../executor/test_token_budget_fallback.py    | 41 ++++++++++++++++++-
 2 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index f10682bdbbdd..140d492792a8 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -989,10 +989,14 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int:
 
     @staticmethod
     def _has_mm_bidirectional_block(req: LlmRequest) -> bool:
-        # Mirror the gate in scheduler_v2._align_chunk_to_mm_block: re-chunking
-        # a request whose boundary would split a bidirectional multimodal block
-        # silently breaks attention, so such requests are deferred whole rather
-        # than re-chunked.
+        """Whether ``req`` carries a bidirectional multimodal block that makes
+        re-chunking unsafe.
+
+        Mirrors the gate in ``scheduler_v2._align_chunk_to_mm_block``:
+        re-chunking a request whose boundary would split a bidirectional
+        multimodal block silently breaks attention, so such requests are
+        deferred whole rather than re-chunked.
+        """
         mm = getattr(req, "py_multimodal_data", None)
         return isinstance(mm, dict) and mm.get("mm_bidirectional_blocks", False)
 
@@ -1048,10 +1052,16 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
         remaining = budget - gen_tokens
         kept: RequestList = []
         deferring = False
+        # Tracks whether we changed the batch at all -- either by dropping a
+        # context request (deferral) or by shrinking one's chunk (re-chunk).
+        # Re-chunking does not change len(kept), so the count alone is not a
+        # sufficient signal that the batch's last-chunk/chunking bins are stale.
+        modified = False
         for req in scheduled_batch.context_requests:
             # Disagg generation-init requests only allocate/transfer KV cache
             # and contribute no compute tokens, so never shed them.
             if deferring and not req.is_disagg_generation_init_state:
+                modified = True
                 continue
             cost = self._request_forward_tokens(req, is_context=True)
             if cost <= remaining:
@@ -1071,20 +1081,31 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
             if (new_chunk >= self.tokens_per_block
                     and new_chunk < req.context_chunk_size
                     and not self._has_mm_bidirectional_block(req)):
-                req.context_chunk_size = new_chunk  # now a non-last chunk
+                # Shrinking context_chunk_size flips is_last_context_chunk (a
+                # computed property: context_current_position + chunk_size ==
+                # prompt_len) to False, so this is now a non-last chunk and must
+                # be re-binned into the chunking list below -- otherwise
+                # downstream treats it as a final chunk and appends generation /
+                # draft tokens to it, corrupting the forward pass.
+                req.context_chunk_size = new_chunk
                 kept.append(req)
                 remaining -= new_chunk
+                modified = True
+            else:
+                # Cannot re-chunk: defer this request entirely.
+                modified = True
             # remaining budget is now < one block, so no further context
             # request can fit this iteration.
             deferring = True
 
-        if len(kept) != scheduled_batch.num_context_requests:
+        if modified:
             logger.debug(
                 f"_fit_token_budget: kept {len(kept)}/"
                 f"{scheduled_batch.num_context_requests} context requests to "
                 f"stay within max_num_tokens={budget}")
-            # reset_context_requests re-derives chunking vs last-chunk from each
-            # request's (possibly updated) is_last_context_chunk.
+            # Re-bin kept requests into chunking vs last-chunk from each
+            # request's (possibly updated) is_last_context_chunk, and drop any
+            # deferred requests from this iteration's batch.
             scheduled_batch.reset_context_requests(kept)
 
     def prepare_resources(self, scheduled_batch: ScheduledRequests):
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
index 429d95f6e5cd..a12c41df663f 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fallback.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -24,18 +24,33 @@ def __init__(
         *,
         context_chunk_size=0,
         is_last_context_chunk=True,
+        prompt_len=None,
+        context_current_position=0,
         py_beam_width=1,
         py_draft_tokens=None,
         is_disagg_generation_init_state=False,
         mm_bidirectional=False,
     ):
         self.context_chunk_size = context_chunk_size
-        self.is_last_context_chunk = is_last_context_chunk
+        self.context_current_position = context_current_position
+        # Mirrors the C++ semantics: is_last_context_chunk is a *computed*
+        # property (context_current_position + context_chunk_size == prompt_len),
+        # so shrinking the chunk during re-chunk flips it to False. When
+        # prompt_len is None the flag is a fixed override (for tests that don't
+        # exercise re-chunk re-binning).
+        self._prompt_len = prompt_len
+        self._is_last_override = is_last_context_chunk
         self.py_beam_width = py_beam_width
         self.py_draft_tokens = py_draft_tokens
         self.is_disagg_generation_init_state = is_disagg_generation_init_state
         self.py_multimodal_data = {"mm_bidirectional_blocks": True} if mm_bidirectional else None
 
+    @property
+    def is_last_context_chunk(self):
+        if self._prompt_len is None:
+            return self._is_last_override
+        return self.context_current_position + self.context_chunk_size == self._prompt_len
+
 
 def _make_manager(max_num_tokens, tokens_per_block):
     # Skip the heavy (GPU-allocating) __init__; the method under test only
@@ -100,6 +115,30 @@ def test_overshoot_rechunks_context(self):
         )
         self.assertLessEqual(total, mgr.max_num_tokens)
 
+    def test_rechunk_only_rebins_to_chunking(self):
+        # Regression for the prep-boundary corruption (issue #13318 follow-up):
+        # when the overshoot is absorbed purely by re-chunking the *last*
+        # context request (no deferral), len(kept) is unchanged, but the request
+        # has flipped from last-chunk to non-last and MUST be moved out of the
+        # last-chunk bin. Otherwise downstream treats it as a final chunk and
+        # appends generation/draft tokens, corrupting the forward pass.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        # Full prompt is 64 tokens, processed in one (last) chunk.
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        self.assertTrue(ctx.is_last_context_chunk)
+        gen = _FakeRequest(py_beam_width=100)  # remaining = 28
+        batch = _make_batch([ctx], [gen])
+
+        mgr._fit_token_budget(batch)
+
+        # Re-chunked to (28 // 16) * 16 == 16, now a non-last chunk.
+        self.assertEqual(ctx.context_chunk_size, 16)
+        self.assertFalse(ctx.is_last_context_chunk)
+        # Count is unchanged, but it must have been re-binned into chunking.
+        self.assertEqual(batch.num_context_requests, 1)
+        self.assertIn(ctx, batch.context_requests_chunking)
+        self.assertNotIn(ctx, batch.context_requests_last_chunk)
+
     def test_overshoot_defers_when_cannot_rechunk(self):
         # Only an 8-token budget remains -- smaller than one block -- so the
         # context request cannot be re-chunked and must be deferred entirely.

From 69c378673607abfad2890b22841eb7e5289d6037 Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Fri, 12 Jun 2026 23:17:07 +0000
Subject: [PATCH 03/16] [#13318][fix] Only re-chunk token-budget fallback when
 chunked 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 #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>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       | 10 +++++++
 .../_torch/pyexecutor/resource_manager.py     | 28 ++++++++++++++----
 .../executor/test_token_budget_fallback.py    | 29 +++++++++++++++++--
 3 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 31b0c8d28328..effa288949fe 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1452,6 +1452,16 @@ def _create_kv_cache_manager(
     # via cache_layer_idx — shared layers use target layer's index for
     # get_buffers(). No layer_offsets remapping needed here.
 
+    # Propagate the finalized chunked-prefill flag so KVCacheManager._fit_token_budget
+    # only re-chunks context requests when the attention backend can consume a
+    # partial context chunk; otherwise it defers them instead. The flag is read
+    # from attn_runtime_features, which py_executor_creator finalizes (including
+    # the SM-version / attention-backend overrides) before build_managers runs.
+    if isinstance(kv_cache_manager,
+                  KVCacheManager) and model_engine is not None:
+        kv_cache_manager.enable_chunked_prefill = bool(
+            model_engine.attn_runtime_features.chunked_prefill)
+
     return kv_cache_manager
 
 
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 140d492792a8..ea9fac0916ff 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -571,6 +571,7 @@ def __init__(
         # attention types (e.g. Gemma4 SWA head_dim=256 + full-attention
         # head_dim=512).
         pool_configurations: Optional[List[PoolConfiguration]] = None,
+        enable_chunked_prefill: bool = False,
         **kwargs,
     ) -> None:
         self.mapping = mapping
@@ -658,6 +659,12 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
         # Kept so prepare_resources can re-validate the per-step token budget
         # (the forward-pass scratch size enforced in _prepare_tp_inputs).
         self.max_num_tokens = max_num_tokens
+        # 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.
+        self.enable_chunked_prefill = enable_chunked_prefill
         self.event_buffer_max_size = kv_cache_config.event_buffer_max_size
         self.attention_dp_events_gather_period_ms = kv_cache_config.attention_dp_events_gather_period_ms
         self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0
@@ -1035,6 +1042,13 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
         Deferred context requests are simply dropped from this iteration's
         ``scheduled_batch``; they remain in the active pool and are rescheduled
         on a later iteration with a fresh budget.
+
+        An over-budget context request is shrunk (re-chunked) in place only when
+        chunked prefill is enabled; otherwise the attention backend is not set
+        up to consume a partial context chunk and the request is deferred whole
+        instead. Re-chunking with chunked prefill disabled produces an invalid
+        forward pass (empty-query asserts / missing quantized KV buffers /
+        cudaErrorInvalidValue) -- see the regression covered by PR #15187.
         """
         budget = self.max_num_tokens
 
@@ -1072,13 +1086,17 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
             # Doesn't fit. Try re-chunking the compute (fewer tokens this step)
             # before deferring. Re-chunking only reduces compute tokens -- KV is
             # allocated for the full prompt regardless -- so block accounting is
-            # unaffected. Only safe when the request can be chunked further, the
-            # shrunk chunk still holds at least one block (aligned to block
-            # size), and the boundary won't split a bidirectional multimodal
-            # block.
+            # unaffected. Only safe when chunked prefill is enabled (otherwise
+            # the attention backend is not set up to consume a partial context
+            # chunk -- shrinking the chunk produces an invalid forward pass,
+            # e.g. cudaErrorInvalidValue / empty-query asserts), the request can
+            # be chunked further, the shrunk chunk still holds at least one block
+            # (aligned to block size), and the boundary won't split a
+            # bidirectional multimodal block.
             new_chunk = (remaining //
                          self.tokens_per_block) * self.tokens_per_block
-            if (new_chunk >= self.tokens_per_block
+            if (self.enable_chunked_prefill
+                    and new_chunk >= self.tokens_per_block
                     and new_chunk < req.context_chunk_size
                     and not self._has_mm_bidirectional_block(req)):
                 # Shrinking context_chunk_size flips is_last_context_chunk (a
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
index a12c41df663f..cfa8e3e718be 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fallback.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -52,12 +52,17 @@ def is_last_context_chunk(self):
         return self.context_current_position + self.context_chunk_size == self._prompt_len
 
 
-def _make_manager(max_num_tokens, tokens_per_block):
+def _make_manager(max_num_tokens, tokens_per_block, enable_chunked_prefill=True):
     # Skip the heavy (GPU-allocating) __init__; the method under test only
-    # needs these two attributes plus its own (bound) helper methods.
+    # needs these attributes plus its own (bound) helper methods.
     mgr = KVCacheManager.__new__(KVCacheManager)
     mgr.max_num_tokens = max_num_tokens
     mgr.tokens_per_block = tokens_per_block
+    # Re-chunking is only valid when chunked prefill is enabled; otherwise the
+    # attention backend cannot consume a partial context chunk and the fallback
+    # must defer instead. Default to enabled so the re-chunk tests exercise that
+    # path; the disabled case is covered explicitly below.
+    mgr.enable_chunked_prefill = enable_chunked_prefill
     return mgr
 
 
@@ -139,6 +144,26 @@ def test_rechunk_only_rebins_to_chunking(self):
         self.assertIn(ctx, batch.context_requests_chunking)
         self.assertNotIn(ctx, batch.context_requests_last_chunk)
 
+    def test_overshoot_defers_when_chunked_prefill_disabled(self):
+        # Regression for the CI failures (q.numel()==0 / "Separate quantized
+        # buffer is not provided" / cudaErrorInvalidValue) seen in PR #15187:
+        # when chunked prefill is disabled the attention backend cannot consume
+        # a partial context chunk, so an over-budget request that *would* be
+        # re-chunkable must instead be deferred whole -- never re-chunked.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False)
+        # Same shape as test_overshoot_rechunks_context (a 28-token budget and a
+        # 64-token last chunk that is block-aligned re-chunkable to 16), but with
+        # chunked prefill off the request must be deferred, not shrunk.
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        gen = _FakeRequest(py_beam_width=100)  # remaining = 28
+        batch = _make_batch([ctx], [gen])
+
+        mgr._fit_token_budget(batch)
+
+        self.assertEqual(batch.num_context_requests, 0)
+        self.assertEqual(ctx.context_chunk_size, 64)  # not re-chunked
+        self.assertTrue(ctx.is_last_context_chunk)  # still a whole last chunk
+
     def test_overshoot_defers_when_cannot_rechunk(self):
         # Only an 8-token budget remains -- smaller than one block -- so the
         # context request cannot be re-chunked and must be deferred entirely.

From 9266f94023c499b02356bc50b807bebddbd7bad7 Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Sat, 13 Jun 2026 00:05:23 +0000
Subject: [PATCH 04/16] [#13318][feat] Make the token-budget fallback an
 opt-out flag

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_llm/_torch/pyexecutor/_util.py       | 10 +++++-
 .../_torch/pyexecutor/resource_manager.py     | 10 ++++--
 tensorrt_llm/llmapi/llm_args.py               | 10 ++++++
 .../executor/test_token_budget_fallback.py    | 31 +++++++++++++++++++
 .../api_stability/references/llm.yaml         |  4 +++
 5 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index effa288949fe..2d081cb97bdc 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -784,6 +784,8 @@ def _create_kv_cache_manager(
             execution_stream=self._execution_stream,
             layer_mask=spec_dec_layer_mask,
             is_disagg=self._is_disagg,
+            enable_token_budget_fallback=getattr(
+                self._llm_args, "enable_token_budget_fallback", True),
         )
 
         if not self._skip_est:
@@ -1126,7 +1128,8 @@ def _create_kv_cache_manager(
         is_draft: Optional[bool] = None,
         layer_mask: Optional[List[bool]] = None,
         num_layers: Optional[int] = None,
-        is_disagg: bool = False) -> KVCacheManager:
+        is_disagg: bool = False,
+        enable_token_budget_fallback: bool = True) -> KVCacheManager:
     """
     Returns:
         A KVCacheManager instance for the given model engine or model config
@@ -1461,6 +1464,11 @@ def _create_kv_cache_manager(
                   KVCacheManager) and model_engine is not None:
         kv_cache_manager.enable_chunked_prefill = bool(
             model_engine.attn_runtime_features.chunked_prefill)
+    # Opt-out switch (TorchLlmArgs.enable_token_budget_fallback) for the
+    # prep-boundary token-budget fallback in _fit_token_budget.
+    if isinstance(kv_cache_manager, KVCacheManager):
+        kv_cache_manager.enable_token_budget_fallback = (
+            enable_token_budget_fallback)
 
     return kv_cache_manager
 
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index ea9fac0916ff..c5edd058ca23 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -572,6 +572,7 @@ def __init__(
         # head_dim=512).
         pool_configurations: Optional[List[PoolConfiguration]] = None,
         enable_chunked_prefill: bool = False,
+        enable_token_budget_fallback: bool = True,
         **kwargs,
     ) -> None:
         self.mapping = mapping
@@ -665,6 +666,10 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
         # Defaults to False (safe: defer instead of re-chunk) and is set to the
         # finalized value by _create_kv_cache_manager.
         self.enable_chunked_prefill = enable_chunked_prefill
+        # Opt-out switch for the prep-boundary token-budget fallback
+        # (_fit_token_budget). Enabled by default; set False to restore the
+        # pre-fallback behavior. Wired from TorchLlmArgs.enable_token_budget_fallback.
+        self.enable_token_budget_fallback = enable_token_budget_fallback
         self.event_buffer_max_size = kv_cache_config.event_buffer_max_size
         self.attention_dp_events_gather_period_ms = kv_cache_config.attention_dp_events_gather_period_ms
         self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0
@@ -1127,9 +1132,10 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
             scheduled_batch.reset_context_requests(kept)
 
     def prepare_resources(self, scheduled_batch: ScheduledRequests):
-        if not self.is_draft:
+        if not self.is_draft and self.enable_token_budget_fallback:
             # The draft-model engine builds inputs with a different token shape;
-            # its budget is handled separately.
+            # its budget is handled separately. Gated by an opt-out flag so the
+            # fallback can be disabled to restore pre-fallback behavior.
             self._fit_token_budget(scheduled_batch)
         with request_context(self.is_draft, scheduled_batch):
             # wait for all pending work to finish before launching offload/onboarding/partial copy
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index 7219d59cc76c..64eb00199d60 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -4054,6 +4054,16 @@ class TorchLlmArgs(BaseLlmArgs):
         "Lower values trigger more frequent garbage collection.",
         status="beta")
 
+    enable_token_budget_fallback: bool = Field(
+        default=True,
+        description=
+        "Re-validate the per-step token budget at the prep boundary and "
+        "gracefully defer or re-chunk over-budget context requests instead of "
+        "letting a scheduler/materialization divergence trip the forward-pass "
+        "token assert and wedge the executor loop (GitHub issue #13318). "
+        "Disable to restore the pre-fallback behavior.",
+        status="beta")
+
     cuda_graph_config: Optional[CudaGraphConfigType] = Field(
         default_factory=CudaGraphConfig,
         description="CUDA graph config. If true, use CUDA graphs for decoding. \
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
index cfa8e3e718be..a354f162d415 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fallback.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -209,6 +209,37 @@ def test_defers_all_subsequent_context_requests(self):
         self.assertIn(ctx1, kept)
         self.assertNotIn(ctx3, kept)
 
+    def test_fallback_can_be_disabled_via_flag(self):
+        # The fallback is opt-out (TorchLlmArgs.enable_token_budget_fallback,
+        # default True). When disabled, prepare_resources must NOT invoke
+        # _fit_token_budget, leaving the scheduled batch untouched.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        mgr.is_draft = False
+        mgr.enable_token_budget_fallback = False
+
+        called = []
+        mgr._fit_token_budget = lambda batch: called.append(batch)
+
+        # Reproduce the gate from KVCacheManager.prepare_resources without the
+        # surrounding GPU work.
+        if not mgr.is_draft and mgr.enable_token_budget_fallback:
+            mgr._fit_token_budget(object())
+
+        self.assertEqual(called, [])
+
+        # Sanity: enabling it does call through.
+        mgr.enable_token_budget_fallback = True
+        if not mgr.is_draft and mgr.enable_token_budget_fallback:
+            mgr._fit_token_budget(object())
+        self.assertEqual(len(called), 1)
+
+    def test_torch_llm_args_flag_default_is_opt_out(self):
+        # The user-facing flag must default to enabled (opt-out semantics).
+        from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
+
+        field = TorchLlmArgs.model_fields["enable_token_budget_fallback"]
+        self.assertEqual(field.default, True)
+
     def test_generation_alone_over_budget_raises(self):
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
         gen = _FakeRequest(py_beam_width=200)
diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml
index a0cccbfa287e..a855c3e6f526 100644
--- a/tests/unittest/api_stability/references/llm.yaml
+++ b/tests/unittest/api_stability/references/llm.yaml
@@ -74,6 +74,10 @@ methods:
         annotation: int
         default: 20000
         status: beta
+      enable_token_budget_fallback:
+        annotation: bool
+        default: True
+        status: beta
       # Misc
       backend:
         annotation: Literal["pytorch"]

From 1b0276cc7c751a2a87abada5c6df67fdc42cc96b Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Mon, 15 Jun 2026 14:52:39 +0000
Subject: [PATCH 05/16] [#13318][fix] Run token-budget fallback before all
 resource managers

The prep-boundary token-budget fallback (#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>
---
 .../_torch/pyexecutor/resource_manager.py     | 33 ++++++-
 .../executor/test_token_budget_fallback.py    | 89 ++++++++++++++++++-
 2 files changed, 120 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index c5edd058ca23..807dc78ff4a0 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -1131,12 +1131,29 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
             # deferred requests from this iteration's batch.
             scheduled_batch.reset_context_requests(kept)
 
-    def prepare_resources(self, scheduled_batch: ScheduledRequests):
+    def maybe_fit_token_budget(self,
+                               scheduled_batch: ScheduledRequests) -> None:
+        """Apply the prep-boundary token-budget fallback to ``scheduled_batch``.
+
+        This is a *batch-level* scheduling decision (defer/re-chunk context
+        requests so the forward pass cannot exceed ``max_num_tokens``) and MUST
+        run before any resource manager allocates KV cache for the batch. It is
+        therefore driven once by ``ResourceManager.prepare_resources`` rather
+        than from this manager's own ``prepare_resources``: the target KV cache
+        manager is deliberately invoked *last* (see ``_util.py``'s
+        ``move_to_end(KV_CACHE_MANAGER)``), so running the fallback here would
+        let an earlier manager -- e.g. a separate draft KV cache manager under
+        MTP -- add sequences for context requests the fallback then defers,
+        orphaning those sequences and tripping a double-add (``emplaceDone``,
+        kvCacheManager.cpp) when the deferred requests reschedule.
+        """
         if not self.is_draft and self.enable_token_budget_fallback:
             # The draft-model engine builds inputs with a different token shape;
             # its budget is handled separately. Gated by an opt-out flag so the
             # fallback can be disabled to restore pre-fallback behavior.
             self._fit_token_budget(scheduled_batch)
+
+    def prepare_resources(self, scheduled_batch: ScheduledRequests):
         with request_context(self.is_draft, scheduled_batch):
             # wait for all pending work to finish before launching offload/onboarding/partial copy
             self.impl.sync_transfer_manager_with_buffer_manager()
@@ -3964,6 +3981,20 @@ def get_resource_manager(
 
     @nvtx_range("prepare_resources")
     def prepare_resources(self, scheduled_batch: ScheduledRequests):
+        # Apply the prep-boundary token-budget fallback (#13318) once, before
+        # any manager allocates resources. It defers/re-chunks context requests
+        # so the forward pass cannot exceed max_num_tokens, and mutates
+        # scheduled_batch in place. It must run up front so every manager --
+        # including a separate draft KV cache manager (MTP) that is invoked
+        # before the target KV cache manager -- observes the same deferred
+        # batch; otherwise an earlier manager adds sequences for context
+        # requests the fallback later defers, orphaning them and tripping a
+        # double-add (emplaceDone) when those requests reschedule.
+        kv_cache_manager = self.resource_managers.get(
+            ResourceManagerType.KV_CACHE_MANAGER)
+        if kv_cache_manager is not None and hasattr(kv_cache_manager,
+                                                    "maybe_fit_token_budget"):
+            kv_cache_manager.maybe_fit_token_budget(scheduled_batch)
         for _, resource_manager in self.resource_managers.items():
             if hasattr(resource_manager, "prepare_resources"):
                 resource_manager.prepare_resources(scheduled_batch)
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
index a354f162d415..6e57b880d0e1 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fallback.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -11,14 +11,21 @@
 """
 
 import unittest
+from collections import OrderedDict
 
-from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
+from tensorrt_llm._torch.pyexecutor.resource_manager import (
+    KVCacheManager,
+    ResourceManager,
+    ResourceManagerType,
+)
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 
 
 class _FakeRequest:
     """Minimal stand-in exposing only the attributes _fit_token_budget reads."""
 
+    _next_id = 0
+
     def __init__(
         self,
         *,
@@ -31,6 +38,8 @@ def __init__(
         is_disagg_generation_init_state=False,
         mm_bidirectional=False,
     ):
+        _FakeRequest._next_id += 1
+        self.py_request_id = _FakeRequest._next_id
         self.context_chunk_size = context_chunk_size
         self.context_current_position = context_current_position
         # Mirrors the C++ semantics: is_last_context_chunk is a *computed*
@@ -240,6 +249,84 @@ def test_torch_llm_args_flag_default_is_opt_out(self):
         field = TorchLlmArgs.model_fields["enable_token_budget_fallback"]
         self.assertEqual(field.default, True)
 
+    def test_maybe_fit_token_budget_honors_flag_and_draft(self):
+        # maybe_fit_token_budget is the single entry point driven by the
+        # aggregate ResourceManager. It must apply the fallback only for the
+        # non-draft manager and only when the opt-out flag is enabled.
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        gen = _FakeRequest(py_beam_width=120)  # remaining = 8 -> defer ctx
+
+        # Non-draft + enabled -> defers.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False)
+        mgr.is_draft = False
+        mgr.enable_token_budget_fallback = True
+        batch = _make_batch([ctx], [gen])
+        mgr.maybe_fit_token_budget(batch)
+        self.assertEqual(batch.num_context_requests, 0)
+
+        # Draft manager -> never fits (handled separately).
+        mgr.is_draft = True
+        batch = _make_batch([ctx], [gen])
+        mgr.maybe_fit_token_budget(batch)
+        self.assertEqual(batch.num_context_requests, 1)
+
+        # Disabled flag -> no-op.
+        mgr.is_draft = False
+        mgr.enable_token_budget_fallback = False
+        batch = _make_batch([ctx], [gen])
+        mgr.maybe_fit_token_budget(batch)
+        self.assertEqual(batch.num_context_requests, 1)
+
+    def test_fallback_runs_before_other_managers(self):
+        # Regression for the emplaceDone double-add (PR #15187): the token-budget
+        # fallback must mutate scheduled_batch BEFORE any resource manager
+        # allocates sequences. A separate draft KV cache manager (MTP) is
+        # invoked before the target KV cache manager (the target is moved to the
+        # end of the manager dict on purpose), so if the fallback ran inside the
+        # target's own prepare_resources the draft manager would already have
+        # added sequences for context requests the fallback then defers --
+        # orphaning them and causing a double-add when they reschedule.
+        #
+        # Build the aggregate ResourceManager with the same ordering as
+        # production (draft-like manager first, KV cache manager last) and assert
+        # the earlier manager observes the *already-deferred* batch.
+        target = _make_manager(
+            max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False
+        )
+        target.is_draft = False
+        target.enable_token_budget_fallback = True
+        # Don't touch the GPU: only the budget fallback matters for ordering.
+        target.prepare_resources = lambda batch: None
+
+        observed = []
+
+        class _RecordingManager:
+            def prepare_resources(self, batch):
+                observed.append([r.py_request_id for r in batch.context_requests])
+
+        ctx_keep = _FakeRequest(context_chunk_size=96)
+        ctx_defer = _FakeRequest(context_chunk_size=64)
+        gen = _FakeRequest(py_beam_width=16)  # remaining = 112
+        batch = _make_batch([ctx_keep, ctx_defer], [gen])
+
+        # Draft-like manager registered FIRST, KV cache manager LAST (mirrors
+        # _util.py's move_to_end(KV_CACHE_MANAGER)).
+        rm = ResourceManager(
+            OrderedDict(
+                [
+                    (ResourceManagerType.DRAFT_KV_CACHE_MANAGER, _RecordingManager()),
+                    (ResourceManagerType.KV_CACHE_MANAGER, target),
+                ]
+            )
+        )
+        rm.prepare_resources(batch)
+
+        # ctx_keep (96) fits into 112; ctx_defer (64) does not and is deferred.
+        # The draft-like manager, though invoked first, must have seen only the
+        # kept request -- proving the fallback ran up front.
+        self.assertEqual(observed, [[ctx_keep.py_request_id]])
+        self.assertEqual(batch.num_context_requests, 1)
+
     def test_generation_alone_over_budget_raises(self):
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
         gen = _FakeRequest(py_beam_width=200)

From 84f5094a0abacc822aed31121d23c7ce14292a6c Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Mon, 15 Jun 2026 16:52:58 +0000
Subject: [PATCH 06/16] [#13318][feat] Terminate server on token-budget
 overshoot when fallback disabled

When the prep-boundary token-budget fallback is disabled
(enable_token_budget_fallback=False), an over-budget batch reaches the forward
pass and _prepare_tp_inputs asserted. That AssertionError only killed the
executor's worker thread, leaving the server process up but with every request
hanging forever -- the server looks alive but makes no progress.

Convert the assert into a typed TokenBudgetExceededError and route it, in the
overlap and non-overlap executor loops, through a fatal shutdown
(_handle_errors with a new immediate_fatal flag): all in-flight and queued
requests fail with the error message and the executor enqueues a shutdown, so
the server terminates cleanly with a clear diagnostic instead of zombie-hanging.

- error_classification.py: new TokenBudgetExceededError.
- model_engine._prepare_tp_inputs: raise it instead of asserting.
- py_executor: _handle_errors gains immediate_fatal (bypasses the error budget);
  new _handle_token_budget_error helper; _forward_step re-raises the typed error
  (instead of funneling it through the generic per-request handler, which would
  only charge the budget and leave the server running); the overlap/non-overlap
  loops catch it and route to the fatal shutdown. The pipeline-parallel loop
  intentionally keeps the prior behavior -- a mid-iteration bail would skip PP
  send/recv handle bookkeeping and risk a collective hang.
- Unit tests for the fatal-shutdown routing and the immediate_fatal flag.

Verified on H100 (DeepSeek-V3-Lite + MTP, fallback disabled): a forced overshoot
now surfaces the message to the client (RequestError) and terminates instead of
hanging.

Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
---
 .../_torch/pyexecutor/error_classification.py |  15 +++
 .../_torch/pyexecutor/model_engine.py         |  18 ++-
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  62 ++++++++--
 .../test_token_budget_fatal_shutdown.py       | 113 ++++++++++++++++++
 4 files changed, 197 insertions(+), 11 deletions(-)
 create mode 100644 tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py

diff --git a/tensorrt_llm/_torch/pyexecutor/error_classification.py b/tensorrt_llm/_torch/pyexecutor/error_classification.py
index c356e6ca3a62..ea723988e51d 100644
--- a/tensorrt_llm/_torch/pyexecutor/error_classification.py
+++ b/tensorrt_llm/_torch/pyexecutor/error_classification.py
@@ -21,6 +21,21 @@
 import dataclasses
 import time
 
+
+class TokenBudgetExceededError(RuntimeError):
+    """A scheduled batch overshot ``max_num_tokens`` in the forward pass.
+
+    Raised by ``_prepare_tp_inputs`` when the materialized token count exceeds
+    ``max_num_tokens``. This is normally prevented by the prep-boundary
+    token-budget fallback (``KVCacheManager._fit_token_budget``); when that
+    fallback is disabled (``enable_token_budget_fallback=False``) an over-budget
+    batch reaches the forward pass and trips this error. The executor routes it
+    to a fatal, server-terminating shutdown so every in-flight and queued
+    request fails with this message -- rather than the bare assert killing only
+    the executor loop thread and leaving the server up but hanging.
+    """
+
+
 # Patterns that corrupt the CUDA context beyond recovery.
 # Matched case-insensitively against the error message.
 IMMEDIATE_FATAL_PATTERNS: list[str] = [
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 789e2155da97..36648da85fda 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -65,6 +65,7 @@
 from .cuda_graph_runner import (CUDAGraphRunner, CUDAGraphRunnerConfig,
                                 EncoderCUDAGraphRunner,
                                 EncoderCUDAGraphRunnerConfig)
+from .error_classification import TokenBudgetExceededError
 from .guided_decoder import CapturableGuidedDecoder
 from .layerwise_nvtx_marker import LayerwiseNvtxMarker
 from .llm_request import LlmRequest, get_draft_token_length
@@ -3066,9 +3067,20 @@ def previous_seq_slots_device():
         num_tokens = len(input_ids)
         num_draft_tokens = len(draft_tokens)
         total_num_tokens = len(position_ids)
-        assert total_num_tokens <= self.max_num_tokens, (
-            f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self.max_num_tokens})"
-        )
+        if total_num_tokens > self.max_num_tokens:
+            # The scheduled batch overshot the token budget. This is normally
+            # prevented by the prep-boundary token-budget fallback
+            # (KVCacheManager._fit_token_budget); when that fallback is disabled
+            # the overshoot reaches here. Raise a typed error (instead of a bare
+            # assert) so the executor can route it to a clean, server-terminating
+            # shutdown rather than silently killing only the loop thread and
+            # leaving the server up but hanging.
+            raise TokenBudgetExceededError(
+                f"total_num_tokens ({total_num_tokens}) exceeds max_num_tokens "
+                f"({self.max_num_tokens}); the scheduled batch overshot the "
+                "token budget. This is normally avoided by the prep-boundary "
+                "token-budget fallback -- it is disabled "
+                "(enable_token_budget_fallback=False).")
         # if exist requests that do not have previous batch, copy input_ids and draft_tokens
         if num_tokens > 0:
             input_ids = torch.tensor(input_ids,
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index b438a31e2719..1b478c60bfe4 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -54,7 +54,7 @@
 from .adp_iter_stats import ADPIterStatsBuffer
 from .connectors.kv_cache_connector import KvCacheConnectorManager
 from .dwdp import DwdpManager
-from .error_classification import ErrorBudget
+from .error_classification import ErrorBudget, TokenBudgetExceededError
 from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem
 from .guided_decoder import GuidedDecoder
 from .handle_additional_outputs import HandleAdditionalOutputs
@@ -2013,6 +2013,13 @@ def _executor_loop_pp(self):
                                 self.guided_decoder.add_batch(scheduled_batch)
                                 self.guided_decoder.init_disagg_gen_requests()
 
+                            # NOTE: the pipeline-parallel loop intentionally does
+                            # not translate TokenBudgetExceededError into a fatal
+                            # shutdown here -- a mid-iteration bail would skip the
+                            # PP send/recv handle bookkeeping and risk a collective
+                            # hang. PP retains the pre-existing assert behavior for
+                            # this path; the single-node (overlap / non-overlap)
+                            # loops handle it (see _handle_token_budget_error).
                             batch_outputs = self._forward_step(scheduled_batch)
 
                             guided_decoder_failed_requests = None
@@ -2792,7 +2799,11 @@ def _executor_loop(self):
                             gpu_forward_start, gpu_forward_end) as fwd_timing:
                         if self.dwdp_manager is not None:
                             self.dwdp_manager.prefetch_first_layers()
-                        batch_outputs = self._forward_step(scheduled_batch)
+                        try:
+                            batch_outputs = self._forward_step(scheduled_batch)
+                        except TokenBudgetExceededError as e:
+                            self._handle_token_budget_error(e)
+                            continue
 
                     guided_decoder_failed_requests = None
                     if self.guided_decoder is not None:
@@ -3190,9 +3201,13 @@ def _executor_loop_overlap(self):
 
                     with self.perf_manager.record_perf_events(
                             gpu_forward_start, gpu_forward_end) as fwd_timing:
-                        batch_outputs = self._forward_step(
-                            scheduled_batch, previous_tensors_device,
-                            num_accepted_tokens_device)
+                        try:
+                            batch_outputs = self._forward_step(
+                                scheduled_batch, previous_tensors_device,
+                                num_accepted_tokens_device)
+                        except TokenBudgetExceededError as e:
+                            self._handle_token_budget_error(e)
+                            continue
 
                 if self.previous_batch is not None and should_process_previous_batch:
                     self._update_requests(self.previous_batch.sample_state)
@@ -4281,6 +4296,14 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
             self._kv_connector_wait_for_save()
 
             return outputs
+        except TokenBudgetExceededError:
+            # A token-budget overshoot (fallback disabled) is fatal for the whole
+            # engine, not a per-request failure. Re-raise so the executor loop
+            # routes it to a clean, server-terminating shutdown
+            # (_handle_token_budget_error) instead of the generic per-request
+            # handler below, which would only charge the error budget and leave
+            # the server running.
+            raise
         except Exception as e:
             traceback.print_exc()
             error_msg = str(e)
@@ -4416,13 +4439,36 @@ def _update_requests(self,
             logger.error(f"Encountered an error in sampling: {error_msg}")
             self._handle_errors(error_msg)
 
+    def _handle_token_budget_error(self,
+                                   error: TokenBudgetExceededError) -> None:
+        """Turn an over-budget batch into a clean, server-terminating shutdown.
+
+        ``_prepare_tp_inputs`` raises ``TokenBudgetExceededError`` when a
+        scheduled batch overshoots ``max_num_tokens`` -- only reachable with the
+        prep-boundary token-budget fallback disabled
+        (``enable_token_budget_fallback=False``). Route it through the fatal
+        path so every in-flight and queued request fails with the message and
+        the executor shuts down, instead of the exception killing only this loop
+        thread and leaving the server up but hanging.
+        """
+        logger.error(f"Token budget exceeded; terminating server: {error}")
+        self._handle_errors(str(error), immediate_fatal=True)
+
     def _handle_errors(self,
                        error_msg: Optional[str] = None,
                        *,
                        requests: Optional[List[LlmRequest]] = None,
-                       charge_budget: bool = True) -> None:
+                       charge_budget: bool = True,
+                       immediate_fatal: bool = False) -> None:
         """Fail requests and optionally initiate shutdown on fatal errors.
 
+        When ``immediate_fatal`` is True, the error is treated as fatal
+        unconditionally (bypassing the error budget): **all** active and queued
+        requests are failed and a shutdown is enqueued. Use this for conditions
+        that are known to be unrecoverable for the whole engine (e.g. a token
+        budget overshoot with the fallback disabled), so the server terminates
+        with the error message instead of leaving a half-dead, hanging process.
+
         When ``charge_budget`` is True (the default), classifies the error
         via the error budget.  If deemed fatal (immediate-fatal pattern or
         budget exhausted), **all** active requests are failed and a shutdown
@@ -4456,8 +4502,8 @@ def _handle_errors(self,
         error_responses: Dict[int, LlmResponse] = {}
         error_msg = error_msg or "error"
 
-        is_fatal = (self._error_budget.consume(error_msg)
-                    if charge_budget else False)
+        is_fatal = immediate_fatal or (self._error_budget.consume(error_msg)
+                                       if charge_budget else False)
         if is_fatal and self._error_budget.budget < 1e-9:
             logger.error(f"Error budget exhausted "
                          f"(budget={self._error_budget.budget:.3f}), "
diff --git a/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py b/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py
new file mode 100644
index 000000000000..3b96cccab488
--- /dev/null
+++ b/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py
@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Unit tests for the fatal token-budget shutdown path.
+
+When the prep-boundary token-budget fallback is *disabled*
+(``enable_token_budget_fallback=False``), an over-budget batch reaches the
+forward pass and ``_prepare_tp_inputs`` raises ``TokenBudgetExceededError``.
+The executor must convert that into a fatal, server-terminating shutdown --
+failing every active/queued request with the message and enqueuing a shutdown --
+rather than letting the exception kill only the loop thread and leave the server
+up but hanging. These tests drive ``PyExecutor._handle_errors`` /
+``_handle_token_budget_error`` on a bare instance (``__new__``) with the minimal
+collaborators stubbed; no GPU is touched.
+"""
+
+import unittest
+
+from tensorrt_llm._torch.pyexecutor.error_classification import (
+    ErrorBudget,
+    TokenBudgetExceededError,
+)
+from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+
+class _FakeReq:
+    def __init__(self, rid):
+        self.py_request_id = rid
+        self.py_client_id = rid
+        self.state = None
+
+
+class _EmptyRawQueue:
+    def empty(self):
+        return True
+
+
+class _FakeExecQueue:
+    def __init__(self):
+        self.shutdown_enqueued = False
+
+    def get_request_queue(self):
+        return _EmptyRawQueue()
+
+    def enqueue_shutdown_request(self):
+        self.shutdown_enqueued = True
+
+
+def _make_executor():
+    ex = PyExecutor.__new__(PyExecutor)
+    ex._error_budget = ErrorBudget()
+    ex._fatal_error = None
+    ex.is_shutdown = False
+    ex.waiting_queue = []
+    ex.executor_request_queue = _FakeExecQueue()
+    ex.active_requests = []
+    ex.gather_all_responses = False
+    enqueued = []
+    terminated = []
+    ex._enqueue_responses = lambda items: enqueued.extend(items)
+    ex._terminate_request = lambda r: terminated.append(r.py_request_id)
+    return ex, enqueued, terminated
+
+
+class TestTokenBudgetFatalShutdown(unittest.TestCase):
+    def test_token_budget_error_terminates_server(self):
+        # _handle_token_budget_error must fail ALL active requests with the
+        # message, mark shutdown, and enqueue a shutdown request.
+        ex, enqueued, terminated = _make_executor()
+        ex.active_requests = [_FakeReq(1), _FakeReq(2)]
+
+        ex._handle_token_budget_error(TokenBudgetExceededError("overshot by 100 tokens"))
+
+        self.assertIsNotNone(ex._fatal_error)
+        self.assertTrue(ex.is_shutdown)
+        self.assertTrue(ex.executor_request_queue.shutdown_enqueued)
+        self.assertEqual({rid for rid, _ in enqueued}, {1, 2})
+        for _, resp in enqueued:
+            self.assertIn("overshot by 100 tokens", resp.error_msg)
+        self.assertEqual(set(terminated), {1, 2})
+        self.assertEqual(ex.active_requests, [])
+
+    def test_immediate_fatal_bypasses_error_budget(self):
+        # immediate_fatal forces a fatal shutdown even with a pristine budget
+        # and charge_budget=False (the budget is never consulted).
+        ex, _, _ = _make_executor()
+        ex.active_requests = [_FakeReq(1)]
+
+        ex._handle_errors("boom", charge_budget=False, immediate_fatal=True)
+
+        self.assertIsNotNone(ex._fatal_error)
+        self.assertTrue(ex.is_shutdown)
+        self.assertTrue(ex.executor_request_queue.shutdown_enqueued)
+
+    def test_request_scoped_error_does_not_shutdown(self):
+        # Guard against regressing the per-request path: a non-fatal,
+        # budget-free request error must NOT trigger shutdown.
+        ex, enqueued, terminated = _make_executor()
+        req = _FakeReq(1)
+        other = _FakeReq(2)
+        ex.active_requests = [req, other]
+
+        ex._handle_errors("bad input", requests=[req], charge_budget=False)
+
+        self.assertIsNone(ex._fatal_error)
+        self.assertFalse(ex.is_shutdown)
+        self.assertFalse(ex.executor_request_queue.shutdown_enqueued)
+        # Only the named request was failed; the other stays active.
+        self.assertEqual([rid for rid, _ in enqueued], [1])
+        self.assertEqual(ex.active_requests, [other])
+
+
+if __name__ == "__main__":
+    unittest.main()

From acdf6d8a2edae49f3055ed1076860dd2a0c06aea Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Tue, 16 Jun 2026 18:05:48 +0000
Subject: [PATCH 07/16] [#13318][fix] Keep disagg gen-init requests cost-free
 in token-budget fit

The token-budget fit loop in KVCacheManager._fit_token_budget guarded disagg
generation-init requests only in the deferring fast-path
(not req.is_disagg_generation_init_state), which still let them fall through to
cost-accounting, re-chunking, and deferral -- contradicting the comment that
they are never shed. Hoist the handling to the top of the loop so any such
request is kept unconditionally and cost-free. In practice the capacity
scheduler already partitions these into fitting_disagg_gen_init_requests
(handled by _prepare_disagg_gen_init), so they do not reach this loop; the
change makes the guard defensive and its intent explicit.

Also add a re-chunk regression test with draft tokens, asserting a re-chunked
request flips off the last-chunk path and drops its last-chunk draft tokens
from the forward-token cost.

Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
---
 .../_torch/pyexecutor/resource_manager.py     | 12 ++++++--
 .../executor/test_token_budget_fallback.py    | 30 +++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 874cb3d05ab3..0d5561218231 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -762,8 +762,16 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
         modified = False
         for req in scheduled_batch.context_requests:
             # Disagg generation-init requests only allocate/transfer KV cache
-            # and contribute no compute tokens, so never shed them.
-            if deferring and not req.is_disagg_generation_init_state:
+            # and contribute no compute tokens. The capacity scheduler already
+            # partitions them into a separate fitting_disagg_gen_init_requests
+            # list (capacityScheduler.cpp) handled by _prepare_disagg_gen_init,
+            # so they should never appear in context_requests here -- but if one
+            # ever does, keep it unconditionally and cost-free rather than
+            # accounting, re-chunking, or deferring (shedding) it.
+            if req.is_disagg_generation_init_state:
+                kept.append(req)
+                continue
+            if deferring:
                 modified = True
                 continue
             cost = self._request_forward_tokens(req, is_context=True)
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
index 6e57b880d0e1..e04a42a00c6a 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fallback.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -153,6 +153,36 @@ def test_rechunk_only_rebins_to_chunking(self):
         self.assertIn(ctx, batch.context_requests_chunking)
         self.assertNotIn(ctx, batch.context_requests_last_chunk)
 
+    def test_rechunk_drops_last_chunk_draft_tokens(self):
+        # Same re-chunk regression as above, but with draft tokens, which are
+        # appended only on the *last* chunk (see _request_forward_tokens). If a
+        # re-chunked request were left on the last-chunk path, its draft tokens
+        # would still be counted/materialized and re-introduce the overshoot
+        # this guard prevents. After re-chunking, the request must be a non-last
+        # chunk and its forward-token cost must no longer include the draft.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        # 64-token last chunk + 2 draft tokens; a 28-token budget cannot fit
+        # 64 (+2), but the chunk re-chunks to a block-aligned 16.
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64, py_draft_tokens=[1, 2])
+        self.assertTrue(ctx.is_last_context_chunk)
+        gen = _FakeRequest(py_beam_width=100)  # remaining = 28
+        batch = _make_batch([ctx], [gen])
+
+        mgr._fit_token_budget(batch)
+
+        # Re-chunked to (28 // 16) * 16 == 16 and flipped to a non-last chunk.
+        self.assertEqual(ctx.context_chunk_size, 16)
+        self.assertFalse(ctx.is_last_context_chunk)
+        self.assertIn(ctx, batch.context_requests_chunking)
+        self.assertNotIn(ctx, batch.context_requests_last_chunk)
+        # Cost is now the chunk size alone -- the 2 draft tokens are dropped
+        # because the request is no longer the last chunk.
+        self.assertEqual(mgr._request_forward_tokens(ctx, is_context=True), 16)
+        total = mgr._request_forward_tokens(ctx, is_context=True) + mgr._request_forward_tokens(
+            gen, is_context=False
+        )
+        self.assertLessEqual(total, mgr.max_num_tokens)
+
     def test_overshoot_defers_when_chunked_prefill_disabled(self):
         # Regression for the CI failures (q.numel()==0 / "Separate quantized
         # buffer is not provided" / cudaErrorInvalidValue) seen in PR #15187:

From 0f8bdbccb8e29f17a7d2a31cddd843f84c69cf95 Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Thu, 18 Jun 2026 13:34:34 +0000
Subject: [PATCH 08/16] [#13318][fix] Use correct sparse_attention_config kwarg
 in cross KV pool

The cross-attention KV cache manager creation for encoder-decoder
models passed sparse_attn_config=None to _create_kv_cache_manager,
but the function parameter is named sparse_attention_config. This
raised TypeError: _create_kv_cache_manager() got an unexpected
keyword argument 'sparse_attn_config' on the T5/BART encoder-decoder
paths. The mismatch was introduced while merging main's cross-pool
support with this branch.

Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 76c9f10e8a5a..39b14e1619a5 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1303,7 +1303,7 @@ def _create_cross_kv_cache_manager(
             max_seq_len=max_seq_len,
             max_batch_size=self._max_batch_size,
             spec_config=None,
-            sparse_attn_config=None,
+            sparse_attention_config=None,
             max_num_tokens=self._max_num_tokens,
             max_beam_width=1,
             kv_connector_manager=None,

From 90e3f5d13468a795fe9a24f0575d33d55cb0fb1b Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Fri, 17 Jul 2026 20:31:07 +0000
Subject: [PATCH 09/16] [#13318][test] Stub enable_attention_dp in
 fatal-shutdown unit test

After merging main, PyExecutor._handle_errors' fatal drain path gates an
attention-DP collective on self.enable_attention_dp. The bare-__new__ test
instance did not set it, so both fatal-path cases raised AttributeError.
Stub it False (single-rank, no collective) to match the other stubbed flags.

Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
---
 .../_torch/executor/test_token_budget_fatal_shutdown.py        | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py b/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py
index 3b96cccab488..a4e71d0a32f1 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py
@@ -54,6 +54,9 @@ def _make_executor():
     ex.executor_request_queue = _FakeExecQueue()
     ex.active_requests = []
     ex.gather_all_responses = False
+    # The fatal-shutdown drain path gates an attention-DP collective on this
+    # flag (py_executor._handle_errors); single-rank test path never collects.
+    ex.enable_attention_dp = False
     enqueued = []
     terminated = []
     ex._enqueue_responses = lambda items: enqueued.extend(items)

From 65f1ff020b8f1e134717f1dff046bc5ad5c32892 Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Fri, 17 Jul 2026 20:31:07 +0000
Subject: [PATCH 10/16] [#13318][chore] Add enable_token_budget_fallback to
 telemetry golden manifest

The new TorchLlmArgs.enable_token_budget_fallback bool is captured by the LLM
args telemetry manifest gate (test_llmapi_config_telemetry_docs). Regenerate
llm_args_golden_manifest.json to include it. Requires telemetry/privacy
CODEOWNER approval.

Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
---
 tensorrt_llm/usage/llm_args_golden_manifest.json | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json
index ca6e762490cf..2483f1a502bf 100644
--- a/tensorrt_llm/usage/llm_args_golden_manifest.json
+++ b/tensorrt_llm/usage/llm_args_golden_manifest.json
@@ -467,6 +467,13 @@
       "kind": "value",
       "path": "enable_speculative_beam_history_d2h"
     },
+    {
+      "allowed_values": [],
+      "annotation": "<class 'bool'>",
+      "converter": "",
+      "kind": "value",
+      "path": "enable_token_budget_fallback"
+    },
     {
       "allowed_values": [],
       "annotation": "<class 'bool'>",

From 84292d5ba03f0269067157b27864e3bd0f59402b Mon Sep 17 00:00:00 2001
From: thorjohnsen <tjohnsen@nvidia.com>
Date: Wed, 29 Jul 2026 17:25:01 +0000
Subject: [PATCH 11/16] [#13318][fix] Drop the token-budget opt-out flag and
 its fatal-shutdown path

Address review feedback (#15187): make the prep-boundary token-budget
fallback unconditional and revert the executor error handling it gated.

The TorchLlmArgs.enable_token_budget_fallback flag did not restore
pre-fallback behavior when disabled -- it selected a third behavior that
terminated the server -- so it was useless as an operator kill switch while
still costing permanent public API surface.

The machinery it gated is also unnecessary. The assert in
_prepare_tp_inputs does not escape the executor loop: it fires inside
ModelEngine.forward, is caught by the generic handler in _forward_step, and
is routed to _handle_errors, which fails the batch's requests and charges
the error budget. Separately, 3ec9e9b0 ("Propagate event loop errors to
await_response", nvbugs/6038228) already fails every pending request when
the loop thread does die. TokenBudgetExceededError therefore did not
prevent a hang; it converted a survivable per-batch error into an
unconditional server shutdown. It also called _handle_errors from inside
the loop, which performs collective gathers in attention-DP / gather-all
modes -- the deadlock hazard _event_loop_wrapper documents and avoids --
for a condition that can be rank-local under attention-DP.

Removed: the TorchLlmArgs field and its golden-manifest / api-stability
entries, TokenBudgetExceededError, the model_engine assert-to-raise
conversion, the two executor-loop except branches, the forward() re-raise,
_handle_token_budget_error, the _handle_errors immediate_fatal parameter,
and test_token_budget_fatal_shutdown.py.

Retained: the fallback itself in KVCacheManager._fit_token_budget, the
enable_chunked_prefill propagation in _util.py, and the unit tests for the
scheduling logic. This PR now touches 3 files with no API change.

Signed-off-by: thorjohnsen <tjohnsen@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       |  10 +-
 .../_torch/pyexecutor/error_classification.py |  15 ---
 .../_torch/pyexecutor/model_engine.py         |  18 +--
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  61 ++-------
 .../_torch/pyexecutor/resource_manager.py     |  20 ++-
 tensorrt_llm/llmapi/llm_args.py               |  10 --
 .../usage/llm_args_golden_manifest.json       |   7 --
 .../executor/test_token_budget_fallback.py    |  49 +-------
 .../test_token_budget_fatal_shutdown.py       | 116 ------------------
 .../api_stability/references/llm.yaml         |   4 -
 10 files changed, 24 insertions(+), 286 deletions(-)
 delete mode 100644 tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 256e69e6828d..872e4ccf0121 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1011,8 +1011,6 @@ def _create_kv_cache_manager(
             execution_stream=self._execution_stream,
             layer_mask=spec_dec_layer_mask,
             is_disagg=self._is_disagg,
-            enable_token_budget_fallback=getattr(
-                self._llm_args, "enable_token_budget_fallback", True),
         )
 
         if not self._skip_est:
@@ -1722,8 +1720,7 @@ def _create_kv_cache_manager(
         num_kv_heads: Optional[Union[int, List[int]]] = None,
         head_dim: Optional[int] = None,
         kv_cache_type=None,
-        is_disagg: bool = False,
-        enable_token_budget_fallback: bool = True) -> KVCacheManager:
+        is_disagg: bool = False) -> KVCacheManager:
     """
     Returns:
         A KVCacheManager instance for the given model engine or model config
@@ -2076,11 +2073,6 @@ def _create_kv_cache_manager(
                   KVCacheManager) and model_engine is not None:
         kv_cache_manager.enable_chunked_prefill = bool(
             model_engine.attn_runtime_features.chunked_prefill)
-    # Opt-out switch (TorchLlmArgs.enable_token_budget_fallback) for the
-    # prep-boundary token-budget fallback in _fit_token_budget.
-    if isinstance(kv_cache_manager, KVCacheManager):
-        kv_cache_manager.enable_token_budget_fallback = (
-            enable_token_budget_fallback)
 
     return kv_cache_manager
 
diff --git a/tensorrt_llm/_torch/pyexecutor/error_classification.py b/tensorrt_llm/_torch/pyexecutor/error_classification.py
index 75a69ed7f2fd..be10331f6a9e 100644
--- a/tensorrt_llm/_torch/pyexecutor/error_classification.py
+++ b/tensorrt_llm/_torch/pyexecutor/error_classification.py
@@ -21,21 +21,6 @@
 import dataclasses
 import time
 
-
-class TokenBudgetExceededError(RuntimeError):
-    """A scheduled batch overshot ``max_num_tokens`` in the forward pass.
-
-    Raised by ``_prepare_tp_inputs`` when the materialized token count exceeds
-    ``max_num_tokens``. This is normally prevented by the prep-boundary
-    token-budget fallback (``KVCacheManager._fit_token_budget``); when that
-    fallback is disabled (``enable_token_budget_fallback=False``) an over-budget
-    batch reaches the forward pass and trips this error. The executor routes it
-    to a fatal, server-terminating shutdown so every in-flight and queued
-    request fails with this message -- rather than the bare assert killing only
-    the executor loop thread and leaving the server up but hanging.
-    """
-
-
 # Patterns that corrupt the CUDA context beyond recovery.
 # Matched case-insensitively against the error message.
 IMMEDIATE_FATAL_PATTERNS: list[str] = [
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 5adca74cba73..23f679155940 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -76,7 +76,6 @@
                                 CUDAGraphRunner, CUDAGraphRunnerConfig,
                                 EncoderCUDAGraphRunner,
                                 EncoderCUDAGraphRunnerConfig)
-from .error_classification import TokenBudgetExceededError
 from .guided_decoder import CapturableGuidedDecoder
 from .kv_cache_manager_v2 import KVCacheManagerV2
 from .layerwise_nvtx_marker import LayerwiseNvtxMarker
@@ -4149,20 +4148,9 @@ def previous_seq_slots_device():
         num_tokens = len(input_ids)
         num_draft_tokens = len(draft_tokens)
         total_num_tokens = len(position_ids)
-        if total_num_tokens > self.max_num_tokens:
-            # The scheduled batch overshot the token budget. This is normally
-            # prevented by the prep-boundary token-budget fallback
-            # (KVCacheManager._fit_token_budget); when that fallback is disabled
-            # the overshoot reaches here. Raise a typed error (instead of a bare
-            # assert) so the executor can route it to a clean, server-terminating
-            # shutdown rather than silently killing only the loop thread and
-            # leaving the server up but hanging.
-            raise TokenBudgetExceededError(
-                f"total_num_tokens ({total_num_tokens}) exceeds max_num_tokens "
-                f"({self.max_num_tokens}); the scheduled batch overshot the "
-                "token budget. This is normally avoided by the prep-boundary "
-                "token-budget fallback -- it is disabled "
-                "(enable_token_budget_fallback=False).")
+        assert total_num_tokens <= self.max_num_tokens, (
+            f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self.max_num_tokens})"
+        )
         # if exist requests that do not have previous batch, copy input_ids and draft_tokens
         if num_tokens > 0:
             input_ids = torch.tensor(input_ids,
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index a16c44a482a7..3ab3b8d4f6d9 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -61,7 +61,7 @@
 from .adp_iter_stats import ADPIterStatsBuffer
 from .connectors.kv_cache_connector import KvCacheConnectorManager
 from .dwdp import DwdpManager
-from .error_classification import ErrorBudget, TokenBudgetExceededError
+from .error_classification import ErrorBudget
 from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem
 from .guided_decoder import GuidedDecoder
 from .handle_additional_outputs import HandleAdditionalOutputs
@@ -2648,13 +2648,6 @@ def _executor_loop_pp(self):
                                 self.guided_decoder.add_batch(scheduled_batch)
                                 self.guided_decoder.init_disagg_gen_requests()
 
-                            # NOTE: the pipeline-parallel loop intentionally does
-                            # not translate TokenBudgetExceededError into a fatal
-                            # shutdown here -- a mid-iteration bail would skip the
-                            # PP send/recv handle bookkeeping and risk a collective
-                            # hang. PP retains the pre-existing assert behavior for
-                            # this path; the single-node (overlap / non-overlap)
-                            # loops handle it (see _handle_token_budget_error).
                             with self.perf_manager.record_perf_events(
                                     gpu_forward_start, gpu_forward_end):
                                 batch_outputs = self._forward_step(
@@ -4081,11 +4074,7 @@ def _executor_loop(self):
                             gpu_forward_start, gpu_forward_end) as fwd_timing:
                         if self.dwdp_manager is not None:
                             self.dwdp_manager.prefetch_first_layers()
-                        try:
-                            batch_outputs = self._forward_step(scheduled_batch)
-                        except TokenBudgetExceededError as e:
-                            self._handle_token_budget_error(e)
-                            continue
+                        batch_outputs = self._forward_step(scheduled_batch)
 
                     self._maybe_prefetch_next_iter_mm_encoders(scheduled_batch)
 
@@ -4580,13 +4569,9 @@ def _executor_loop_overlap(self):
 
                     with self.perf_manager.record_perf_events(
                             gpu_forward_start, gpu_forward_end) as fwd_timing:
-                        try:
-                            batch_outputs = self._forward_step(
-                                scheduled_batch, previous_tensors_device,
-                                num_accepted_tokens_device)
-                        except TokenBudgetExceededError as e:
-                            self._handle_token_budget_error(e)
-                            continue
+                        batch_outputs = self._forward_step(
+                            scheduled_batch, previous_tensors_device,
+                            num_accepted_tokens_device)
 
                     self._maybe_prefetch_next_iter_mm_encoders(scheduled_batch)
 
@@ -6276,14 +6261,6 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
             self._kv_connector_wait_for_save()
 
             return outputs
-        except TokenBudgetExceededError:
-            # A token-budget overshoot (fallback disabled) is fatal for the whole
-            # engine, not a per-request failure. Re-raise so the executor loop
-            # routes it to a clean, server-terminating shutdown
-            # (_handle_token_budget_error) instead of the generic per-request
-            # handler below, which would only charge the error budget and leave
-            # the server running.
-            raise
         except Exception as e:
             traceback.print_exc()
             error_msg = str(e)
@@ -6419,36 +6396,13 @@ def _update_requests(self,
             logger.error(f"Encountered an error in sampling: {error_msg}")
             self._handle_errors(error_msg)
 
-    def _handle_token_budget_error(self,
-                                   error: TokenBudgetExceededError) -> None:
-        """Turn an over-budget batch into a clean, server-terminating shutdown.
-
-        ``_prepare_tp_inputs`` raises ``TokenBudgetExceededError`` when a
-        scheduled batch overshoots ``max_num_tokens`` -- only reachable with the
-        prep-boundary token-budget fallback disabled
-        (``enable_token_budget_fallback=False``). Route it through the fatal
-        path so every in-flight and queued request fails with the message and
-        the executor shuts down, instead of the exception killing only this loop
-        thread and leaving the server up but hanging.
-        """
-        logger.error(f"Token budget exceeded; terminating server: {error}")
-        self._handle_errors(str(error), immediate_fatal=True)
-
     def _handle_errors(self,
                        error_msg: Optional[str] = None,
                        *,
                        requests: Optional[List[LlmRequest]] = None,
-                       charge_budget: bool = True,
-                       immediate_fatal: bool = False) -> None:
+                       charge_budget: bool = True) -> None:
         """Fail requests and optionally initiate shutdown on fatal errors.
 
-        When ``immediate_fatal`` is True, the error is treated as fatal
-        unconditionally (bypassing the error budget): **all** active and queued
-        requests are failed and a shutdown is enqueued. Use this for conditions
-        that are known to be unrecoverable for the whole engine (e.g. a token
-        budget overshoot with the fallback disabled), so the server terminates
-        with the error message instead of leaving a half-dead, hanging process.
-
         When ``charge_budget`` is True (the default), classifies the error
         via the error budget.  If deemed fatal (immediate-fatal pattern or
         budget exhausted), **all** active requests are failed and a shutdown
@@ -6485,8 +6439,7 @@ def _handle_errors(self,
 
         budget_fatal = (self._error_budget.consume(error_msg)
                         if charge_budget else False)
-        is_fatal = (immediate_fatal or self._fatal_error is not None
-                    or budget_fatal)
+        is_fatal = self._fatal_error is not None or budget_fatal
         if budget_fatal and self._error_budget.budget < 1e-9:
             logger.error(f"Error budget exhausted "
                          f"(budget={self._error_budget.budget:.3f}), "
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 54f3b1b6581f..5e83f71c5a8e 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -303,7 +303,6 @@ def __init__(
         # head_dim=512).
         pool_configurations: Optional[List[PoolConfiguration]] = None,
         enable_chunked_prefill: bool = False,
-        enable_token_budget_fallback: bool = True,
         **kwargs,
     ) -> None:
         self.mapping = mapping
@@ -397,10 +396,6 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
         # Defaults to False (safe: defer instead of re-chunk) and is set to the
         # finalized value by _create_kv_cache_manager.
         self.enable_chunked_prefill = enable_chunked_prefill
-        # Opt-out switch for the prep-boundary token-budget fallback
-        # (_fit_token_budget). Enabled by default; set False to restore the
-        # pre-fallback behavior. Wired from TorchLlmArgs.enable_token_budget_fallback.
-        self.enable_token_budget_fallback = enable_token_budget_fallback
         self.event_buffer_max_size = kv_cache_config.event_buffer_max_size
         self.attention_dp_events_gather_period_ms = kv_cache_config.attention_dp_events_gather_period_ms
         self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0
@@ -804,10 +799,12 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
         tokens actually materialized by ``_prepare_tp_inputs`` -- for example
         when a reuse-discounted last context chunk lands next to a near-full
         generation batch (see GitHub issue #13318). Rather than letting that
-        divergence trip a hard assert and wedge the executor loop, re-validate
-        the budget here -- before any KV cache is allocated -- and gracefully
-        shed only the deferrable work (context chunks), leaving in-flight
-        generation requests untouched.
+        divergence trip the ``total_num_tokens <= max_num_tokens`` assert in
+        ``_prepare_tp_inputs`` -- which fails every request in the batch and
+        charges the executor's error budget -- re-validate the budget here,
+        before any KV cache is allocated, and gracefully shed only the
+        deferrable work (context chunks), leaving in-flight generation requests
+        untouched.
 
         Deferred context requests are simply dropped from this iteration's
         ``scheduled_batch``; they remain in the active pool and are rescheduled
@@ -920,10 +917,9 @@ def maybe_fit_token_budget(self,
         orphaning those sequences and tripping a double-add (``emplaceDone``,
         kvCacheManager.cpp) when the deferred requests reschedule.
         """
-        if not self.is_draft and self.enable_token_budget_fallback:
+        if not self.is_draft:
             # The draft-model engine builds inputs with a different token shape;
-            # its budget is handled separately. Gated by an opt-out flag so the
-            # fallback can be disabled to restore pre-fallback behavior.
+            # its budget is handled separately.
             self._fit_token_budget(scheduled_batch)
 
     def _context_seq_len(self, req: LlmRequest, is_cross: bool,
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index 672095447ffb..57c3de25dcfe 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -4586,16 +4586,6 @@ class TorchLlmArgs(BaseLlmArgs):
         "Lower values trigger more frequent garbage collection.",
         status="beta")
 
-    enable_token_budget_fallback: bool = Field(
-        default=True,
-        description=
-        "Re-validate the per-step token budget at the prep boundary and "
-        "gracefully defer or re-chunk over-budget context requests instead of "
-        "letting a scheduler/materialization divergence trip the forward-pass "
-        "token assert and wedge the executor loop (GitHub issue #13318). "
-        "Disable to restore the pre-fallback behavior.",
-        status="beta")
-
     cuda_graph_config: Optional[CudaGraphConfigType] = Field(
         default_factory=CudaGraphConfig,
         description="CUDA graph config. If true, use CUDA graphs for decoding. \
diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json
index 2483f1a502bf..ca6e762490cf 100644
--- a/tensorrt_llm/usage/llm_args_golden_manifest.json
+++ b/tensorrt_llm/usage/llm_args_golden_manifest.json
@@ -467,13 +467,6 @@
       "kind": "value",
       "path": "enable_speculative_beam_history_d2h"
     },
-    {
-      "allowed_values": [],
-      "annotation": "<class 'bool'>",
-      "converter": "",
-      "kind": "value",
-      "path": "enable_token_budget_fallback"
-    },
     {
       "allowed_values": [],
       "annotation": "<class 'bool'>",
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
index e04a42a00c6a..ea2a28ec3eaa 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fallback.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -248,48 +248,17 @@ def test_defers_all_subsequent_context_requests(self):
         self.assertIn(ctx1, kept)
         self.assertNotIn(ctx3, kept)
 
-    def test_fallback_can_be_disabled_via_flag(self):
-        # The fallback is opt-out (TorchLlmArgs.enable_token_budget_fallback,
-        # default True). When disabled, prepare_resources must NOT invoke
-        # _fit_token_budget, leaving the scheduled batch untouched.
-        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
-        mgr.is_draft = False
-        mgr.enable_token_budget_fallback = False
-
-        called = []
-        mgr._fit_token_budget = lambda batch: called.append(batch)
-
-        # Reproduce the gate from KVCacheManager.prepare_resources without the
-        # surrounding GPU work.
-        if not mgr.is_draft and mgr.enable_token_budget_fallback:
-            mgr._fit_token_budget(object())
-
-        self.assertEqual(called, [])
-
-        # Sanity: enabling it does call through.
-        mgr.enable_token_budget_fallback = True
-        if not mgr.is_draft and mgr.enable_token_budget_fallback:
-            mgr._fit_token_budget(object())
-        self.assertEqual(len(called), 1)
-
-    def test_torch_llm_args_flag_default_is_opt_out(self):
-        # The user-facing flag must default to enabled (opt-out semantics).
-        from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
-
-        field = TorchLlmArgs.model_fields["enable_token_budget_fallback"]
-        self.assertEqual(field.default, True)
-
-    def test_maybe_fit_token_budget_honors_flag_and_draft(self):
+    def test_maybe_fit_token_budget_skips_draft_manager(self):
         # maybe_fit_token_budget is the single entry point driven by the
-        # aggregate ResourceManager. It must apply the fallback only for the
-        # non-draft manager and only when the opt-out flag is enabled.
+        # aggregate ResourceManager. It must apply the fallback for the target
+        # manager only -- the draft-model engine builds inputs with a different
+        # token shape and its budget is handled separately.
         ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
         gen = _FakeRequest(py_beam_width=120)  # remaining = 8 -> defer ctx
 
-        # Non-draft + enabled -> defers.
+        # Non-draft -> defers.
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False)
         mgr.is_draft = False
-        mgr.enable_token_budget_fallback = True
         batch = _make_batch([ctx], [gen])
         mgr.maybe_fit_token_budget(batch)
         self.assertEqual(batch.num_context_requests, 0)
@@ -300,13 +269,6 @@ def test_maybe_fit_token_budget_honors_flag_and_draft(self):
         mgr.maybe_fit_token_budget(batch)
         self.assertEqual(batch.num_context_requests, 1)
 
-        # Disabled flag -> no-op.
-        mgr.is_draft = False
-        mgr.enable_token_budget_fallback = False
-        batch = _make_batch([ctx], [gen])
-        mgr.maybe_fit_token_budget(batch)
-        self.assertEqual(batch.num_context_requests, 1)
-
     def test_fallback_runs_before_other_managers(self):
         # Regression for the emplaceDone double-add (PR #15187): the token-budget
         # fallback must mutate scheduled_batch BEFORE any resource manager
@@ -324,7 +286,6 @@ def test_fallback_runs_before_other_managers(self):
             max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False
         )
         target.is_draft = False
-        target.enable_token_budget_fallback = True
         # Don't touch the GPU: only the budget fallback matters for ordering.
         target.prepare_resources = lambda batch: None
 
diff --git a/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py b/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py
deleted file mode 100644
index a4e71d0a32f1..000000000000
--- a/tests/unittest/_torch/executor/test_token_budget_fatal_shutdown.py
+++ /dev/null
@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-"""Unit tests for the fatal token-budget shutdown path.
-
-When the prep-boundary token-budget fallback is *disabled*
-(``enable_token_budget_fallback=False``), an over-budget batch reaches the
-forward pass and ``_prepare_tp_inputs`` raises ``TokenBudgetExceededError``.
-The executor must convert that into a fatal, server-terminating shutdown --
-failing every active/queued request with the message and enqueuing a shutdown --
-rather than letting the exception kill only the loop thread and leave the server
-up but hanging. These tests drive ``PyExecutor._handle_errors`` /
-``_handle_token_budget_error`` on a bare instance (``__new__``) with the minimal
-collaborators stubbed; no GPU is touched.
-"""
-
-import unittest
-
-from tensorrt_llm._torch.pyexecutor.error_classification import (
-    ErrorBudget,
-    TokenBudgetExceededError,
-)
-from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
-
-
-class _FakeReq:
-    def __init__(self, rid):
-        self.py_request_id = rid
-        self.py_client_id = rid
-        self.state = None
-
-
-class _EmptyRawQueue:
-    def empty(self):
-        return True
-
-
-class _FakeExecQueue:
-    def __init__(self):
-        self.shutdown_enqueued = False
-
-    def get_request_queue(self):
-        return _EmptyRawQueue()
-
-    def enqueue_shutdown_request(self):
-        self.shutdown_enqueued = True
-
-
-def _make_executor():
-    ex = PyExecutor.__new__(PyExecutor)
-    ex._error_budget = ErrorBudget()
-    ex._fatal_error = None
-    ex.is_shutdown = False
-    ex.waiting_queue = []
-    ex.executor_request_queue = _FakeExecQueue()
-    ex.active_requests = []
-    ex.gather_all_responses = False
-    # The fatal-shutdown drain path gates an attention-DP collective on this
-    # flag (py_executor._handle_errors); single-rank test path never collects.
-    ex.enable_attention_dp = False
-    enqueued = []
-    terminated = []
-    ex._enqueue_responses = lambda items: enqueued.extend(items)
-    ex._terminate_request = lambda r: terminated.append(r.py_request_id)
-    return ex, enqueued, terminated
-
-
-class TestTokenBudgetFatalShutdown(unittest.TestCase):
-    def test_token_budget_error_terminates_server(self):
-        # _handle_token_budget_error must fail ALL active requests with the
-        # message, mark shutdown, and enqueue a shutdown request.
-        ex, enqueued, terminated = _make_executor()
-        ex.active_requests = [_FakeReq(1), _FakeReq(2)]
-
-        ex._handle_token_budget_error(TokenBudgetExceededError("overshot by 100 tokens"))
-
-        self.assertIsNotNone(ex._fatal_error)
-        self.assertTrue(ex.is_shutdown)
-        self.assertTrue(ex.executor_request_queue.shutdown_enqueued)
-        self.assertEqual({rid for rid, _ in enqueued}, {1, 2})
-        for _, resp in enqueued:
-            self.assertIn("overshot by 100 tokens", resp.error_msg)
-        self.assertEqual(set(terminated), {1, 2})
-        self.assertEqual(ex.active_requests, [])
-
-    def test_immediate_fatal_bypasses_error_budget(self):
-        # immediate_fatal forces a fatal shutdown even with a pristine budget
-        # and charge_budget=False (the budget is never consulted).
-        ex, _, _ = _make_executor()
-        ex.active_requests = [_FakeReq(1)]
-
-        ex._handle_errors("boom", charge_budget=False, immediate_fatal=True)
-
-        self.assertIsNotNone(ex._fatal_error)
-        self.assertTrue(ex.is_shutdown)
-        self.assertTrue(ex.executor_request_queue.shutdown_enqueued)
-
-    def test_request_scoped_error_does_not_shutdown(self):
-        # Guard against regressing the per-request path: a non-fatal,
-        # budget-free request error must NOT trigger shutdown.
-        ex, enqueued, terminated = _make_executor()
-        req = _FakeReq(1)
-        other = _FakeReq(2)
-        ex.active_requests = [req, other]
-
-        ex._handle_errors("bad input", requests=[req], charge_budget=False)
-
-        self.assertIsNone(ex._fatal_error)
-        self.assertFalse(ex.is_shutdown)
-        self.assertFalse(ex.executor_request_queue.shutdown_enqueued)
-        # Only the named request was failed; the other stays active.
-        self.assertEqual([rid for rid, _ in enqueued], [1])
-        self.assertEqual(ex.active_requests, [other])
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml
index 00db34c8c85d..7ff208c41ce3 100644
--- a/tests/unittest/api_stability/references/llm.yaml
+++ b/tests/unittest/api_stability/references/llm.yaml
@@ -78,10 +78,6 @@ methods:
         annotation: int
         default: 20000
         status: beta
-      enable_token_budget_fallback:
-        annotation: bool
-        default: True
-        status: beta
       # Misc
       backend:
         annotation: Literal["pytorch"]

From 9b974fb966dae2e5554b0fc8bfa517648f200b69 Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Fri, 31 Jul 2026 16:43:18 +0000
Subject: [PATCH 12/16] [#13318][fix] Drive token-budget fallback from the
 executor loop

The fallback ran inside ResourceManager.prepare_resources, which is past two
points that must observe the trimmed batch:

- _can_queue all-gathers scheduled_batch.batch_size under attention DP and
  gates on no rank being empty. Trimming after that vote let a rank shed its
  way to an empty batch once every rank had already agreed to run -- the exact
  state the gate exists to prevent.
- _add_inflight_ids (PP) registers the batch's last-chunk context requests in
  the inflight set, and _remove_inflight_ids re-derived that set from the batch
  afterwards. A deferred request (dropped from the batch) or a re-chunked one
  (moved out of context_requests_last_chunk) was registered but never erased.
  The scheduler skips inflight ids, so such a request was never scheduled again
  while still holding its KV blocks and sequence slot.

Move the driver to ResourceManager.maybe_fit_token_budget, called from
_prepare_and_schedule_batch and _executor_loop_pp right after scheduling, so
the batch every later step sees is the one that will run. It still precedes
every resource manager, so the emplaceDone double-add ordering is preserved.
Record the ids _add_inflight_ids actually inserted on ScheduledRequests so the
paired removal erases exactly those, independent of any later mutation.

Return early from _fit_token_budget when the batch has no context requests: the
fallback can only shed context work, so the generation-token scan was dead work
on the gen-only path (~50us at 512 generation requests, ~100us at 1024). This
also stops the gen-only over-budget RuntimeError, which killed the event loop on
a condition that is rank-local under attention DP; that case falls back to the
pre-existing _prepare_tp_inputs assert.

Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  46 ++++--
 .../_torch/pyexecutor/resource_manager.py     |  71 ++++++---
 .../_torch/pyexecutor/scheduler/scheduler.py  |  10 ++
 .../executor/test_token_budget_fallback.py    | 145 +++++++++++++++++-
 4 files changed, 238 insertions(+), 34 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 2efc017e0e4c..e3d1ce3887be 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -2592,6 +2592,12 @@ def _executor_loop_pp(self):
                         num_fitting_reqs, fitting_disagg_gen_init_requests,
                         wait_for_disagg_gen_transfer_progress, all_gen_first)
 
+                # Re-validate the per-step token budget and shed deferrable
+                # work here, before the batch is measured, registered in the
+                # inflight set, voted on by _can_queue, or allocated against.
+                # See ResourceManager.maybe_fit_token_budget.
+                self.resource_manager.maybe_fit_token_budget(scheduled_batch)
+
                 self.num_scheduled_requests = scheduled_batch.batch_size
 
                 logger.debug(
@@ -3722,6 +3728,11 @@ def _prepare_and_schedule_batch(self):
                 self._handle_errors(error_msg, requests=self.active_requests)
                 return None, None
 
+        # Re-validate the per-step token budget and shed deferrable work here,
+        # before the batch is measured, voted on by _can_queue, or allocated
+        # against. See ResourceManager.maybe_fit_token_budget.
+        self.resource_manager.maybe_fit_token_budget(scheduled_batch)
+
         self.num_scheduled_requests = scheduled_batch.batch_size
         logger.debug(
             f'has {len(self.active_requests)} active_requests, '
@@ -6979,30 +6990,47 @@ def _add_inflight_ids(self, scheduled_requests: ScheduledRequests):
         Only requests that sample new tokens should be added to the inflight set since their next iteration depends
         on these new tokens, so they should be skipped in the scheduler until the new tokens are generated.
         This includes context requests that finish context phase and generation requests.
+
+        The inserted ids are recorded on ``scheduled_requests`` so the paired
+        ``_remove_inflight_ids`` erases exactly what was added, rather than
+        re-deriving the set from a batch that may have changed in between. An id
+        left behind is not recoverable: the scheduler skips inflight ids, so the
+        request is never scheduled again while still holding its KV blocks and
+        sequence slot.
+
+        Callers currently trim the batch before this point (see
+        ``ResourceManager.maybe_fit_token_budget``, which can defer a context
+        request out of the batch entirely or re-chunk one out of
+        ``context_requests_last_chunk``), so the two views agree today. Snapshot
+        anyway: the pairing must not silently depend on that ordering.
         """
+        added: List[int] = []
         for req in scheduled_requests.context_requests_last_chunk:
             logger.debug(
                 f"Context request with ID {req.request_id} added to DECODER model inflight set"
             )
             self.inflight_req_ids.insert(req.request_id)
+            added.append(req.request_id)
         for req in scheduled_requests.generation_requests:
             logger.debug(
                 f"Generation request with ID {req.request_id} added to DECODER model inflight set"
             )
             self.inflight_req_ids.insert(req.request_id)
+            added.append(req.request_id)
+        scheduled_requests.added_inflight_req_ids = added
 
     def _remove_inflight_ids(self, scheduled_requests: ScheduledRequests):
-        """Remove request IDs of current sampling requests from self.inflight_req_ids."""
-        for req in scheduled_requests.context_requests_last_chunk:
-            logger.debug(
-                f"Context request with ID {req.request_id} removed from DECODER model inflight set"
-            )
-            self.inflight_req_ids.erase(req.request_id)
-        for req in scheduled_requests.generation_requests:
+        """Remove the request IDs this batch added to self.inflight_req_ids.
+
+        Erases the ids recorded by ``_add_inflight_ids`` rather than re-deriving
+        them from the batch, which may have been trimmed since (see there).
+        """
+        for req_id in scheduled_requests.added_inflight_req_ids:
             logger.debug(
-                f"Generation request with ID {req.request_id} removed from DECODER model inflight set"
+                f"Request with ID {req_id} removed from DECODER model inflight set"
             )
-            self.inflight_req_ids.erase(req.request_id)
+            self.inflight_req_ids.erase(req_id)
+        scheduled_requests.added_inflight_req_ids = []
 
     def _handle_speculative_decoding(
         self, scheduled_batch, previous_tensors, target_inputs
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 2d842855b2df..0984bb8ee857 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -822,7 +822,21 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
         instead. Re-chunking with chunked prefill disabled produces an invalid
         forward pass (empty-query asserts / missing quantized KV buffers /
         cudaErrorInvalidValue) -- see the regression covered by PR #15187.
+
+        Context requests are the only work this fallback can shed, so a batch
+        with none scheduled returns immediately -- including before the
+        generation-token accounting below, which cannot change the outcome when
+        there is nothing to defer or re-chunk. That keeps the gen-only batch
+        (the executor loop's hottest, most latency-sensitive path) free of an
+        O(num_generation_requests) scan that costs ~50us at 512 generation
+        requests and ~100us at 1024.
         """
+        # Tested against the two lists rather than the ``context_requests``
+        # property, which concatenates them into a fresh list on every access.
+        if (not scheduled_batch.context_requests_chunking
+                and not scheduled_batch.context_requests_last_chunk):
+            return
+
         budget = self.max_num_tokens
 
         # Generation requests are in-flight and cannot be deferred. If they
@@ -912,16 +926,17 @@ def maybe_fit_token_budget(self,
         """Apply the prep-boundary token-budget fallback to ``scheduled_batch``.
 
         This is a *batch-level* scheduling decision (defer/re-chunk context
-        requests so the forward pass cannot exceed ``max_num_tokens``) and MUST
-        run before any resource manager allocates KV cache for the batch. It is
-        therefore driven once by ``ResourceManager.prepare_resources`` rather
-        than from this manager's own ``prepare_resources``: the target KV cache
-        manager is deliberately invoked *last* (see ``_util.py``'s
-        ``move_to_end(KV_CACHE_MANAGER)``), so running the fallback here would
-        let an earlier manager -- e.g. a separate draft KV cache manager under
-        MTP -- add sequences for context requests the fallback then defers,
-        orphaning those sequences and tripping a double-add (``emplaceDone``,
-        kvCacheManager.cpp) when the deferred requests reschedule.
+        requests so the forward pass cannot exceed ``max_num_tokens``), so it is
+        driven once by ``ResourceManager.maybe_fit_token_budget`` from the
+        executor loop -- see there for why it has to run that early -- rather
+        than from this manager's own ``prepare_resources``. In particular the
+        target KV cache manager is deliberately invoked *last* (see
+        ``_util.py``'s ``move_to_end(KV_CACHE_MANAGER)``), so running the
+        fallback here would let an earlier manager -- e.g. a separate draft KV
+        cache manager under MTP -- add sequences for context requests the
+        fallback then defers, orphaning those sequences and tripping a
+        double-add (``emplaceDone``, kvCacheManager.cpp) when the deferred
+        requests reschedule.
         """
         if not self.is_draft:
             # The draft-model engine builds inputs with a different token shape;
@@ -2771,22 +2786,36 @@ def get_resource_manager(
             self, type: ResourceManagerType) -> Optional[BaseResourceManager]:
         return self.resource_managers.get(type)
 
-    @nvtx_range("prepare_resources")
-    def prepare_resources(self, scheduled_batch: ScheduledRequests):
-        # Apply the prep-boundary token-budget fallback (#13318) once, before
-        # any manager allocates resources. It defers/re-chunks context requests
-        # so the forward pass cannot exceed max_num_tokens, and mutates
-        # scheduled_batch in place. It must run up front so every manager --
-        # including a separate draft KV cache manager (MTP) that is invoked
-        # before the target KV cache manager -- observes the same deferred
-        # batch; otherwise an earlier manager adds sequences for context
-        # requests the fallback later defers, orphaning them and tripping a
-        # double-add (emplaceDone) when those requests reschedule.
+    @nvtx_range("maybe_fit_token_budget")
+    def maybe_fit_token_budget(self, scheduled_batch: ScheduledRequests):
+        """Apply the prep-boundary token-budget fallback (#13318) to the batch.
+
+        Defers or re-chunks context requests so the forward pass cannot exceed
+        max_num_tokens, mutating scheduled_batch in place. Driven by the
+        executor loop right after scheduling -- not from prepare_resources --
+        because the batch it produces must be the one every later step sees:
+
+        - _can_queue all-gathers scheduled_batch.batch_size under attention DP
+          and gates on no rank being empty. Trimming after that vote lets a rank
+          shed its way to an empty batch once the ranks have already agreed to
+          run, which is the exact state that gate exists to prevent.
+        - _add_inflight_ids (PP) registers the batch's last-chunk context
+          requests in the inflight set. Trimming after it registers requests
+          that are no longer in the batch.
+        - Every resource manager keys off the batch's contents, including a
+          separate draft KV cache manager (MTP) invoked before the target one.
+          A manager that adds sequences for a context request the fallback then
+          defers orphans them, tripping a double-add (emplaceDone) when those
+          requests reschedule.
+        """
         kv_cache_manager = self.resource_managers.get(
             ResourceManagerType.KV_CACHE_MANAGER)
         if kv_cache_manager is not None and hasattr(kv_cache_manager,
                                                     "maybe_fit_token_budget"):
             kv_cache_manager.maybe_fit_token_budget(scheduled_batch)
+
+    @nvtx_range("prepare_resources")
+    def prepare_resources(self, scheduled_batch: ScheduledRequests):
         for _, resource_manager in self.resource_managers.items():
             if hasattr(resource_manager, "prepare_resources"):
                 resource_manager.prepare_resources(scheduled_batch)
diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
index caa8e3cb3de1..bdcb4e718cb1 100644
--- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
+++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
@@ -152,6 +152,15 @@ class ScheduledRequests:
     """Requests that are in the generation phase."""
     paused_requests: RequestList
     """Requests that are paused."""
+    added_inflight_req_ids: list[int]
+    """Request ids this batch inserted into the executor's inflight set.
+
+    Recorded by ``PyExecutor._add_inflight_ids`` so the paired
+    ``_remove_inflight_ids`` erases exactly what was added. The batch can be
+    trimmed between the two calls -- ``ResourceManager.prepare_resources``
+    defers over-budget context requests and re-chunks others -- so the ids are
+    no longer derivable from the request lists at removal time.
+    """
 
     def __init__(self):
         self.encoder_requests: RequestList = []
@@ -159,6 +168,7 @@ def __init__(self):
         self.context_requests_last_chunk: RequestList = []
         self.generation_requests: RequestList = []
         self.paused_requests: RequestList = []
+        self.added_inflight_req_ids: list[int] = []
 
     @property
     def is_generation_only(self) -> bool:
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
index ea2a28ec3eaa..dcc6364df639 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fallback.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -40,6 +40,10 @@ def __init__(
     ):
         _FakeRequest._next_id += 1
         self.py_request_id = _FakeRequest._next_id
+        # PyExecutor's inflight-set bookkeeping reads ``request_id`` (the C++
+        # binding's name) rather than ``py_request_id``; keep them in sync so the
+        # same fake drives both the fallback and TestInflightIdsSurviveTrim.
+        self.request_id = self.py_request_id
         self.context_chunk_size = context_chunk_size
         self.context_current_position = context_current_position
         # Mirrors the C++ semantics: is_last_context_chunk is a *computed*
@@ -279,9 +283,12 @@ def test_fallback_runs_before_other_managers(self):
         # added sequences for context requests the fallback then defers --
         # orphaning them and causing a double-add when they reschedule.
         #
-        # Build the aggregate ResourceManager with the same ordering as
-        # production (draft-like manager first, KV cache manager last) and assert
-        # the earlier manager observes the *already-deferred* batch.
+        # The executor loop drives the fallback via
+        # ResourceManager.maybe_fit_token_budget before it calls
+        # prepare_resources; build the aggregate ResourceManager with the same
+        # ordering as production (draft-like manager first, KV cache manager
+        # last) and assert the earlier manager observes the *already-deferred*
+        # batch.
         target = _make_manager(
             max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False
         )
@@ -310,6 +317,7 @@ def prepare_resources(self, batch):
                 ]
             )
         )
+        rm.maybe_fit_token_budget(batch)
         rm.prepare_resources(batch)
 
         # ctx_keep (96) fits into 112; ctx_defer (64) does not and is deferred.
@@ -318,14 +326,143 @@ def prepare_resources(self, batch):
         self.assertEqual(observed, [[ctx_keep.py_request_id]])
         self.assertEqual(batch.num_context_requests, 1)
 
+    def test_prepare_resources_does_not_trim(self):
+        # The fallback must NOT be reachable from prepare_resources: the
+        # executor loop drives it earlier so that _can_queue's attention-DP
+        # tp_allgather(batch_size) and the PP inflight-set registration both see
+        # the trimmed batch. A second trim here would be redundant at best, and
+        # leaving it as the *only* trim point is what let a rank shed its way to
+        # an empty batch after the ranks had already voted to run.
+        target = _make_manager(
+            max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False
+        )
+        target.is_draft = False
+        target.prepare_resources = lambda batch: None
+
+        ctx_over_budget = _FakeRequest(context_chunk_size=64)
+        gen = _FakeRequest(py_beam_width=100)  # remaining = 28, so ctx cannot fit
+        batch = _make_batch([ctx_over_budget], [gen])
+
+        rm = ResourceManager(OrderedDict([(ResourceManagerType.KV_CACHE_MANAGER, target)]))
+        rm.prepare_resources(batch)
+
+        self.assertEqual(batch.num_context_requests, 1)
+
+        # ...and the batch is trimmed only once maybe_fit_token_budget is called.
+        rm.maybe_fit_token_budget(batch)
+        self.assertEqual(batch.num_context_requests, 0)
+
     def test_generation_alone_over_budget_raises(self):
+        # A context request must be present for the fallback to engage at all
+        # (see test_gen_only_batch_is_left_alone); generation requests that
+        # exceed the budget by themselves leave nothing to shed.
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
         gen = _FakeRequest(py_beam_width=200)
-        batch = _make_batch([], [gen])
+        ctx = _FakeRequest(context_chunk_size=16)
+        batch = _make_batch([ctx], [gen])
 
         with self.assertRaises(RuntimeError):
             mgr._fit_token_budget(batch)
 
+    def test_gen_only_batch_is_left_alone(self):
+        # Context requests are the only thing the fallback can shed, so a
+        # gen-only batch returns before the generation-token accounting: it
+        # keeps that scan off the executor loop's hottest path, and it avoids
+        # raising on a condition nothing here can fix. An over-budget gen-only
+        # batch stays the concern of the _prepare_tp_inputs assert, which fails
+        # one batch rather than killing the (possibly only rank-local) event
+        # loop.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        gen = _FakeRequest(py_beam_width=200)  # 200 > max_num_tokens
+        batch = _make_batch([], [gen])
+
+        mgr._fit_token_budget(batch)  # must not raise
+
+        self.assertEqual(batch.num_context_requests, 0)
+        self.assertEqual(batch.generation_requests, [gen])
+
+
+class TestInflightIdsSurviveTrim(unittest.TestCase):
+    """The fallback must not strand ids in PyExecutor's inflight set.
+
+    ``_executor_loop_pp`` calls ``_add_inflight_ids`` before
+    ``ResourceManager.prepare_resources`` and ``_remove_inflight_ids`` after, so
+    the fallback runs between them and mutates the batch: a deferred context
+    request is dropped from it, and a re-chunked one moves out of
+    ``context_requests_last_chunk``. Removal must therefore erase the ids that
+    were actually inserted, not re-derive them from the trimmed batch -- an id
+    left behind makes the scheduler skip that request forever (scheduler.py's
+    ``if req.request_id in inflight_request_ids: continue``).
+    """
+
+    @staticmethod
+    def _bare_executor():
+        # Same trick as _make_manager: PyExecutor.__init__ builds an engine, so
+        # instantiate bare and supply only the inflight set the methods touch.
+        from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+        from tensorrt_llm.bindings.internal.batch_manager import ReqIdsSet
+
+        executor = PyExecutor.__new__(PyExecutor)
+        executor.inflight_req_ids = ReqIdsSet()
+        return executor
+
+    def test_trimmed_context_requests_leave_no_inflight_ids(self):
+        executor = self._bare_executor()
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+
+        gen = _FakeRequest(py_beam_width=100)  # leaves a 28-token budget
+        # Both start as last-chunk requests, so both are registered inflight.
+        # ctx_rechunked shrinks to a block-aligned 16 and flips to a non-last
+        # chunk; ctx_deferred no longer fits and is dropped from the batch.
+        ctx_rechunked = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        ctx_deferred = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        batch = _make_batch([ctx_rechunked, ctx_deferred], [gen])
+
+        executor._add_inflight_ids(batch)
+        self.assertEqual(
+            sorted(batch.added_inflight_req_ids),
+            sorted([ctx_rechunked.request_id, ctx_deferred.request_id, gen.request_id]),
+        )
+
+        mgr._fit_token_budget(batch)
+
+        # Preconditions for the regression: the batch really did change shape.
+        self.assertEqual(ctx_rechunked.context_chunk_size, 16)
+        self.assertIn(ctx_rechunked, batch.context_requests_chunking)
+        self.assertEqual(batch.context_requests_last_chunk, [])
+        self.assertNotIn(ctx_deferred, batch.context_requests)
+
+        executor._remove_inflight_ids(batch)
+
+        for req in (ctx_rechunked, ctx_deferred, gen):
+            self.assertNotIn(
+                req.request_id,
+                executor.inflight_req_ids,
+                f"request {req.request_id} left in the inflight set; the scheduler "
+                "would never schedule it again",
+            )
+        self.assertEqual(batch.added_inflight_req_ids, [])
+
+    def test_untrimmed_batch_round_trips(self):
+        # The batch the fallback leaves alone must behave exactly as before.
+        executor = self._bare_executor()
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+
+        gen = _FakeRequest(py_beam_width=1)
+        ctx = _FakeRequest(context_chunk_size=16)
+        batch = _make_batch([ctx], [gen])
+
+        executor._add_inflight_ids(batch)
+        for req in (ctx, gen):
+            self.assertIn(req.request_id, executor.inflight_req_ids)
+
+        mgr._fit_token_budget(batch)
+        self.assertEqual(batch.num_context_requests, 1)
+
+        executor._remove_inflight_ids(batch)
+        for req in (ctx, gen):
+            self.assertNotIn(req.request_id, executor.inflight_req_ids)
+
 
 if __name__ == "__main__":
     unittest.main()

From 1a2942ce82410305b029bbf0a8737f10e91f0400 Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Fri, 31 Jul 2026 17:15:59 +0000
Subject: [PATCH 13/16] [#13318][fix] Guard resource_manager in
 _prepare_and_schedule_batch

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>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 0c2ff9287b81..c48b22dfaaa6 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3734,7 +3734,14 @@ def _prepare_and_schedule_batch(self):
         # Re-validate the per-step token budget and shed deferrable work here,
         # before the batch is measured, voted on by _can_queue, or allocated
         # against. See ResourceManager.maybe_fit_token_budget.
-        self.resource_manager.maybe_fit_token_budget(scheduled_batch)
+        #
+        # resource_manager is guarded for the same reason model_engine is above:
+        # unit tests drive this method on partially-constructed executors
+        # (object.__new__) that set only the attributes under test. A real
+        # executor always has one -- __init__ assigns it unconditionally.
+        resource_manager = getattr(self, "resource_manager", None)
+        if resource_manager is not None:
+            resource_manager.maybe_fit_token_budget(scheduled_batch)
 
         self.num_scheduled_requests = scheduled_batch.batch_size
         logger.debug(

From dfd0b6ed202941a5f4b369c3bc2139d61dffdd2d Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Thu, 6 Aug 2026 18:08:27 +0000
Subject: [PATCH 14/16] [#13318][fix] Move token-budget trim after resource
 preparation

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>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       |  10 +-
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  32 +-
 .../_torch/pyexecutor/resource_manager.py     | 305 ++++++-----
 .../_torch/pyexecutor/scheduler/scheduler.py  |   9 +-
 .../executor/test_token_budget_fallback.py    | 516 ++++++++++--------
 5 files changed, 450 insertions(+), 422 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 30cfd2737fcb..8bbdbdbef2fc 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -2435,11 +2435,11 @@ def _create_kv_cache_manager(
     # via cache_layer_idx — shared layers use target layer's index for
     # get_buffers(). No layer_offsets remapping needed here.
 
-    # Propagate the finalized chunked-prefill flag so KVCacheManager._fit_token_budget
-    # only re-chunks context requests when the attention backend can consume a
-    # partial context chunk; otherwise it defers them instead. The flag is read
-    # from attn_runtime_features, which py_executor_creator finalizes (including
-    # the SM-version / attention-backend overrides) before build_managers runs.
+    # Propagate the finalized chunked-prefill flag so KVCacheManager.fit_token_budget
+    # only shrinks context chunks when the attention backend can consume a
+    # partial context chunk. The flag is read from attn_runtime_features, which
+    # py_executor_creator finalizes (including the SM-version /
+    # attention-backend overrides) before build_managers runs.
     if isinstance(kv_cache_manager,
                   KVCacheManager) and model_engine is not None:
         kv_cache_manager.enable_chunked_prefill = bool(
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 933ce9f6b1fd..7f72ae95457e 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -2597,12 +2597,6 @@ def _executor_loop_pp(self):
                         num_fitting_reqs, fitting_disagg_gen_init_requests,
                         wait_for_disagg_gen_transfer_progress, all_gen_first)
 
-                # Re-validate the per-step token budget and shed deferrable
-                # work here, before the batch is measured, registered in the
-                # inflight set, voted on by _can_queue, or allocated against.
-                # See ResourceManager.maybe_fit_token_budget.
-                self.resource_manager.maybe_fit_token_budget(scheduled_batch)
-
                 self.num_scheduled_requests = scheduled_batch.batch_size
 
                 logger.debug(
@@ -3733,18 +3727,6 @@ def _prepare_and_schedule_batch(self):
                 self._handle_errors(error_msg, requests=self.active_requests)
                 return None, None
 
-        # Re-validate the per-step token budget and shed deferrable work here,
-        # before the batch is measured, voted on by _can_queue, or allocated
-        # against. See ResourceManager.maybe_fit_token_budget.
-        #
-        # resource_manager is guarded for the same reason model_engine is above:
-        # unit tests drive this method on partially-constructed executors
-        # (object.__new__) that set only the attributes under test. A real
-        # executor always has one -- __init__ assigns it unconditionally.
-        resource_manager = getattr(self, "resource_manager", None)
-        if resource_manager is not None:
-            resource_manager.maybe_fit_token_budget(scheduled_batch)
-
         self.num_scheduled_requests = scheduled_batch.batch_size
         logger.debug(
             f'has {len(self.active_requests)} active_requests, '
@@ -7111,16 +7093,16 @@ def _add_inflight_ids(self, scheduled_requests: ScheduledRequests):
 
         The inserted ids are recorded on ``scheduled_requests`` so the paired
         ``_remove_inflight_ids`` erases exactly what was added, rather than
-        re-deriving the set from a batch that may have changed in between. An id
-        left behind is not recoverable: the scheduler skips inflight ids, so the
+        re-deriving the set from a batch that has changed in between. An id left
+        behind is not recoverable: the scheduler skips inflight ids, so the
         request is never scheduled again while still holding its KV blocks and
         sequence slot.
 
-        Callers currently trim the batch before this point (see
-        ``ResourceManager.maybe_fit_token_budget``, which can defer a context
-        request out of the batch entirely or re-chunk one out of
-        ``context_requests_last_chunk``), so the two views agree today. Snapshot
-        anyway: the pairing must not silently depend on that ordering.
+        This is load-bearing, not defensive. In this loop the next step is
+        ``resource_manager.prepare_resources``, which runs
+        ``ResourceManager.maybe_fit_token_budget`` at its end -- and that can
+        shrink a context request out of ``context_requests_last_chunk``. By
+        removal time the batch no longer agrees with what was inserted here.
         """
         added: List[int] = []
         for req in scheduled_requests.context_requests_last_chunk:
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index bde57311f681..b955e4c20400 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -773,64 +773,99 @@ def _has_mm_bidirectional_block(req: LlmRequest) -> bool:
 
         Mirrors the gate in ``scheduler_v2._align_chunk_to_mm_block``:
         re-chunking a request whose boundary would split a bidirectional
-        multimodal block silently breaks attention, so such requests are
-        deferred whole rather than re-chunked.
+        multimodal block silently breaks attention, so such requests are left
+        at their scheduled chunk size rather than shrunk.
         """
         mm = getattr(req, "py_multimodal_data", None)
         return isinstance(mm, dict) and mm.get("mm_bidirectional_blocks", False)
 
     def _request_forward_tokens(self, req: LlmRequest, *,
                                 is_context: bool) -> int:
-        """Upper bound on the number of position ids ``req`` contributes to a
-        forward pass in ``_prepare_tp_inputs``.
+        """Number of position ids ``req`` contributes to the forward pass.
 
-        This MUST over-estimate. Under-counting would reintroduce the
-        ``total_num_tokens <= max_num_tokens`` assert in ``_prepare_tp_inputs``
-        that this guard exists to prevent.
+        Exact for context requests, but **only once every resource manager has
+        prepared** -- see ``fit_token_budget``. Before that,
+        ``context_chunk_size`` still spans the reusable KV prefix and this
+        over-counts by up to the whole cached prefix.
+
+        Mirrors ``_prepare_tp_inputs``: a context request contributes
+        ``all_prompt_tokens[pos : pos + chunk]``, which Python slicing clamps to
+        what is left of the prompt, plus draft tokens on the last chunk only.
         """
         draft_len = get_draft_token_length(req)
         if is_context:
-            # Context contributes ``context_chunk_size`` positions; draft tokens
-            # are appended only on the last chunk.
-            return req.context_chunk_size + (draft_len if
-                                             req.is_last_context_chunk else 0)
+            materialized = min(req.context_chunk_size,
+                               req.context_remaining_length)
+            return materialized + (draft_len
+                                   if req.is_last_context_chunk else 0)
         # Generation: one position per beam for the new token, plus draft tokens
         # (speculative verification) per beam.
         return req.py_beam_width * (1 + draft_len)
 
-    def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
-        """Defer or re-chunk context requests so the scheduled batch cannot
-        exceed ``max_num_tokens`` in the forward pass.
-
-        The micro-batch scheduler's token-budget estimate can diverge from the
-        tokens actually materialized by ``_prepare_tp_inputs`` -- for example
-        when a reuse-discounted last context chunk lands next to a near-full
-        generation batch (see GitHub issue #13318). Rather than letting that
-        divergence trip the ``total_num_tokens <= max_num_tokens`` assert in
-        ``_prepare_tp_inputs`` -- which fails every request in the batch and
-        charges the executor's error budget -- re-validate the budget here,
-        before any KV cache is allocated, and gracefully shed only the
-        deferrable work (context chunks), leaving in-flight generation requests
-        untouched.
-
-        Deferred context requests are simply dropped from this iteration's
-        ``scheduled_batch``; they remain in the active pool and are rescheduled
-        on a later iteration with a fresh budget.
-
-        An over-budget context request is shrunk (re-chunked) in place only when
-        chunked prefill is enabled; otherwise the attention backend is not set
-        up to consume a partial context chunk and the request is deferred whole
-        instead. Re-chunking with chunked prefill disabled produces an invalid
-        forward pass (empty-query asserts / missing quantized KV buffers /
-        cudaErrorInvalidValue) -- see the regression covered by PR #15187.
-
-        Context requests are the only work this fallback can shed, so a batch
-        with none scheduled returns immediately -- including before the
-        generation-token accounting below, which cannot change the outcome when
-        there is nothing to defer or re-chunk. That keeps the gen-only batch
-        (the executor loop's hottest, most latency-sensitive path) free of an
-        O(num_generation_requests) scan that costs ~50us at 512 generation
-        requests and ~100us at 1024.
+    def _shrink_context_chunk(self, req: LlmRequest, excess: int) -> int:
+        """Shrink ``req``'s chunk by up to ``excess`` tokens. Returns the number
+        of forward-pass tokens actually shed (0 if the request cannot shrink).
+
+        The new chunk must keep ``context_current_position + context_chunk_size``
+        on a block boundary: ``setPrepopulatedPromptLen`` (llmRequest.h) asserts
+        that for every non-last chunk, to keep the KV cache unfragmented. So the
+        chunk is trimmed to the largest block-aligned end that sheds at least
+        ``excess`` tokens, and never below one block of forward progress -- a
+        zero-token chunk would leave the request scheduled but computing
+        nothing, which never terminates.
+        """
+        pos = req.context_current_position
+        current = min(req.context_chunk_size, req.context_remaining_length)
+
+        # Smallest chunk that still lands on a block boundary and makes progress.
+        floor_chunk = self.tokens_per_block - (pos % self.tokens_per_block)
+        target_end = ((pos + current - excess) //
+                      self.tokens_per_block) * self.tokens_per_block
+        new_chunk = max(floor_chunk, target_end - pos)
+        if new_chunk >= current:
+            return 0
+
+        # Shrinking flips is_last_context_chunk (a computed property:
+        # context_current_position + chunk_size == prompt_len) to False, so the
+        # caller must re-bin the batch afterwards -- otherwise downstream treats
+        # this as a final chunk and appends generation / draft tokens to it.
+        req.context_chunk_size = new_chunk
+        return current - new_chunk
+
+    def fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
+        """Shrink over-budget context chunks so the forward pass cannot exceed
+        ``max_num_tokens``.
+
+        Driven by ``ResourceManager.prepare_resources`` **after** every manager
+        has prepared -- see there for why that is the only correct point. The
+        micro-batch scheduler admits a batch on an estimate: it charges
+        ``reuse_adjusted_compute(chunk, estimated_reusable_tokens, remaining)``
+        (microBatchScheduler.cpp), where ``estimated_reusable_tokens`` is a
+        radix-tree guess made during capacity scheduling. The real figure is
+        ``prepopulated_prompt_len``, computed later inside ``addSequence``. When
+        the real reuse comes in lower than the estimate, the forward pass
+        materializes more tokens than were charged and trips the
+        ``total_num_tokens <= max_num_tokens`` assert in ``_prepare_tp_inputs``
+        -- which fails every request in the batch and charges the executor's
+        error budget (GitHub issue #13318).
+
+        That divergence is not knowable before ``prepare_resources``: the whole
+        point is that the estimate was wrong, and only ``addSequence`` knows by
+        how much. Running here, ``context_current_position`` and
+        ``context_chunk_size`` are final and ``_request_forward_tokens`` is
+        exact.
+
+        Shrink only, never defer. KV cache is already allocated and sequences
+        are already added, so a request cannot be dropped from the batch at this
+        point -- but it does not need to be. Blocks are allocated for the full
+        chunk; the tokens trimmed here are simply computed on the next
+        iteration. Because nothing leaves the batch, this cannot desynchronize
+        the attention-DP ``_can_queue`` vote, the inflight set, or any earlier
+        manager's per-request state.
+
+        Context requests are the only work this can shed, so a batch with none
+        returns immediately -- keeping the gen-only batch (the executor loop's
+        hottest path) free of an O(num_generation_requests) scan.
         """
         # Tested against the two lists rather than the ``context_requests``
         # property, which concatenates them into a fresh list on every access.
@@ -839,110 +874,75 @@ def _fit_token_budget(self, scheduled_batch: ScheduledRequests) -> None:
             return
 
         budget = self.max_num_tokens
-
-        # Generation requests are in-flight and cannot be deferred. If they
-        # alone exceed the budget something is genuinely misconfigured -- fail
-        # this batch loudly rather than overshoot silently.
-        gen_tokens = sum(
+        total = sum(
             self._request_forward_tokens(req, is_context=False)
             for req in scheduled_batch.generation_requests)
-        if gen_tokens > budget:
-            raise RuntimeError(
-                f"In-flight generation requests need {gen_tokens} tokens, "
-                f"exceeding max_num_tokens ({budget}); cannot schedule.")
-
-        remaining = budget - gen_tokens
-        kept: RequestList = []
-        deferring = False
-        # Tracks whether we changed the batch at all -- either by dropping a
-        # context request (deferral) or by shrinking one's chunk (re-chunk).
-        # Re-chunking does not change len(kept), so the count alone is not a
-        # sufficient signal that the batch's last-chunk/chunking bins are stale.
-        modified = False
-        for req in scheduled_batch.context_requests:
+        context_requests = scheduled_batch.context_requests
+        total += sum(
+            self._request_forward_tokens(req, is_context=True)
+            for req in context_requests
+            if not req.is_disagg_generation_init_state)
+
+        excess = total - budget
+        if excess <= 0:
+            return
+
+        if not self.enable_chunked_prefill:
+            # 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 (empty-query asserts /
+            # missing quantized KV buffers / cudaErrorInvalidValue). Nothing
+            # safe is left to do, so let the assert fire as it does on main.
+            logger.warning(
+                f"Scheduled batch needs {total} forward-pass tokens, exceeding "
+                f"max_num_tokens ({budget}) by {excess}. Cannot trim: chunked "
+                "prefill is disabled. See GitHub issue #13318.")
+            return
+
+        # Shed from the back. ``context_requests`` is
+        # ``context_requests_chunking + context_requests_last_chunk``, so this
+        # trims last-chunk requests first -- which is where the overshoot comes
+        # from: only a last chunk carries a reuse discount
+        # (reuse_adjusted_compute's last-chunk branch) and draft tokens, so it is
+        # the request whose cost the scheduler can have under-charged. Trimming
+        # it converts it back into a chunking request, which is exactly the
+        # repair. Mid-prefill chunks are touched only if that is not enough.
+        shed = 0
+        for req in reversed(context_requests):
+            if excess - shed <= 0:
+                break
             # Disagg generation-init requests only allocate/transfer KV cache
-            # and contribute no compute tokens. The capacity scheduler already
-            # partitions them into a separate fitting_disagg_gen_init_requests
-            # list (capacityScheduler.cpp) handled by _prepare_disagg_gen_init,
-            # so they should never appear in context_requests here -- but if one
-            # ever does, keep it unconditionally and cost-free rather than
-            # accounting, re-chunking, or deferring (shedding) it.
+            # and contribute no compute tokens, so there is nothing to shed.
             if req.is_disagg_generation_init_state:
-                kept.append(req)
                 continue
-            if deferring:
-                modified = True
-                continue
-            cost = self._request_forward_tokens(req, is_context=True)
-            if cost <= remaining:
-                kept.append(req)
-                remaining -= cost
+            # Re-chunking a request whose boundary would split a bidirectional
+            # multimodal block silently breaks attention.
+            if self._has_mm_bidirectional_block(req):
                 continue
+            shed += self._shrink_context_chunk(req, excess - shed)
 
-            # Doesn't fit. Try re-chunking the compute (fewer tokens this step)
-            # before deferring. Re-chunking only reduces compute tokens -- KV is
-            # allocated for the full prompt regardless -- so block accounting is
-            # unaffected. Only safe when chunked prefill is enabled (otherwise
-            # the attention backend is not set up to consume a partial context
-            # chunk -- shrinking the chunk produces an invalid forward pass,
-            # e.g. cudaErrorInvalidValue / empty-query asserts), the request can
-            # be chunked further, the shrunk chunk still holds at least one block
-            # (aligned to block size), and the boundary won't split a
-            # bidirectional multimodal block.
-            new_chunk = (remaining //
-                         self.tokens_per_block) * self.tokens_per_block
-            if (self.enable_chunked_prefill
-                    and new_chunk >= self.tokens_per_block
-                    and new_chunk < req.context_chunk_size
-                    and not self._has_mm_bidirectional_block(req)):
-                # Shrinking context_chunk_size flips is_last_context_chunk (a
-                # computed property: context_current_position + chunk_size ==
-                # prompt_len) to False, so this is now a non-last chunk and must
-                # be re-binned into the chunking list below -- otherwise
-                # downstream treats it as a final chunk and appends generation /
-                # draft tokens to it, corrupting the forward pass.
-                req.context_chunk_size = new_chunk
-                kept.append(req)
-                remaining -= new_chunk
-                modified = True
-            else:
-                # Cannot re-chunk: defer this request entirely.
-                modified = True
-            # remaining budget is now < one block, so no further context
-            # request can fit this iteration.
-            deferring = True
-
-        if modified:
+        if shed:
             logger.debug(
-                f"_fit_token_budget: kept {len(kept)}/"
+                f"fit_token_budget: trimmed {shed} context tokens from "
                 f"{scheduled_batch.num_context_requests} context requests to "
                 f"stay within max_num_tokens={budget}")
-            # Re-bin kept requests into chunking vs last-chunk from each
-            # request's (possibly updated) is_last_context_chunk, and drop any
-            # deferred requests from this iteration's batch.
-            scheduled_batch.reset_context_requests(kept)
+            # Re-bin from each request's (possibly updated)
+            # is_last_context_chunk: a shrunk request is no longer a last chunk.
+            scheduled_batch.reset_context_requests(context_requests)
+
+        if shed < excess:
+            logger.warning(
+                f"Scheduled batch needs {total} forward-pass tokens, exceeding "
+                f"max_num_tokens ({budget}); could only trim {shed}. See "
+                "GitHub issue #13318.")
 
     def maybe_fit_token_budget(self,
                                scheduled_batch: ScheduledRequests) -> None:
-        """Apply the prep-boundary token-budget fallback to ``scheduled_batch``.
-
-        This is a *batch-level* scheduling decision (defer/re-chunk context
-        requests so the forward pass cannot exceed ``max_num_tokens``), so it is
-        driven once by ``ResourceManager.maybe_fit_token_budget`` from the
-        executor loop -- see there for why it has to run that early -- rather
-        than from this manager's own ``prepare_resources``. In particular the
-        target KV cache manager is deliberately invoked *last* (see
-        ``_util.py``'s ``move_to_end(KV_CACHE_MANAGER)``), so running the
-        fallback here would let an earlier manager -- e.g. a separate draft KV
-        cache manager under MTP -- add sequences for context requests the
-        fallback then defers, orphaning those sequences and tripping a
-        double-add (``emplaceDone``, kvCacheManager.cpp) when the deferred
-        requests reschedule.
-        """
+        """Apply the post-allocation token-budget trim to ``scheduled_batch``."""
         if not self.is_draft:
             # The draft-model engine builds inputs with a different token shape;
             # its budget is handled separately.
-            self._fit_token_budget(scheduled_batch)
+            self.fit_token_budget(scheduled_batch)
 
     def _context_seq_len(self, req: LlmRequest, is_cross: bool,
                          is_star_cp: bool) -> Optional[int]:
@@ -2781,25 +2781,25 @@ def get_resource_manager(
 
     @nvtx_range("maybe_fit_token_budget")
     def maybe_fit_token_budget(self, scheduled_batch: ScheduledRequests):
-        """Apply the prep-boundary token-budget fallback (#13318) to the batch.
-
-        Defers or re-chunks context requests so the forward pass cannot exceed
-        max_num_tokens, mutating scheduled_batch in place. Driven by the
-        executor loop right after scheduling -- not from prepare_resources --
-        because the batch it produces must be the one every later step sees:
-
-        - _can_queue all-gathers scheduled_batch.batch_size under attention DP
-          and gates on no rank being empty. Trimming after that vote lets a rank
-          shed its way to an empty batch once the ranks have already agreed to
-          run, which is the exact state that gate exists to prevent.
-        - _add_inflight_ids (PP) registers the batch's last-chunk context
-          requests in the inflight set. Trimming after it registers requests
-          that are no longer in the batch.
-        - Every resource manager keys off the batch's contents, including a
-          separate draft KV cache manager (MTP) invoked before the target one.
-          A manager that adds sequences for a context request the fallback then
-          defers orphans them, tripping a double-add (emplaceDone) when those
-          requests reschedule.
+        """Apply the post-allocation token-budget trim (#13318) to the batch.
+
+        Shrinks over-budget context chunks so the forward pass cannot exceed
+        max_num_tokens, mutating scheduled_batch in place. Driven from
+        prepare_resources, after every manager has run, because that is the
+        first point at which the batch's forward-pass token count is knowable:
+
+        - The micro-batch scheduler admits the batch on estimated_reusable_tokens
+          (a radix-tree guess). The real reuse is prepopulated_prompt_len,
+          computed inside addSequence. #13318 is the case where the two differ,
+          so no check running before addSequence can see the divergence.
+        - Until setPrepopulatedPromptLen runs, context_chunk_size still spans the
+          reusable prefix -- it is not a token count, and reading it as one
+          over-charges a reuse hit by the whole cached prefix.
+
+        Trimming this late is only safe because it shrinks rather than defers;
+        nothing leaves the batch, so the attention-DP _can_queue vote, the PP
+        inflight set and every earlier manager's per-request state all stay
+        consistent with the batch they were computed from.
         """
         kv_cache_manager = self.resource_managers.get(
             ResourceManagerType.KV_CACHE_MANAGER)
@@ -2812,6 +2812,9 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
         for _, resource_manager in self.resource_managers.items():
             if hasattr(resource_manager, "prepare_resources"):
                 resource_manager.prepare_resources(scheduled_batch)
+        # After every manager, so context_current_position / context_chunk_size
+        # are final. See maybe_fit_token_budget.
+        self.maybe_fit_token_budget(scheduled_batch)
 
     @nvtx_range("update_resources")
     def update_resources(
diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
index bdcb4e718cb1..8f75e9cfc001 100644
--- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
+++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
@@ -156,10 +156,11 @@ class ScheduledRequests:
     """Request ids this batch inserted into the executor's inflight set.
 
     Recorded by ``PyExecutor._add_inflight_ids`` so the paired
-    ``_remove_inflight_ids`` erases exactly what was added. The batch can be
-    trimmed between the two calls -- ``ResourceManager.prepare_resources``
-    defers over-budget context requests and re-chunks others -- so the ids are
-    no longer derivable from the request lists at removal time.
+    ``_remove_inflight_ids`` erases exactly what was added. The batch is trimmed
+    between the two calls -- ``ResourceManager.prepare_resources`` shrinks
+    over-budget context chunks at its end, which can move a request out of
+    ``context_requests_last_chunk`` -- so the ids are no longer derivable from
+    the request lists at removal time.
     """
 
     def __init__(self):
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
index dcc6364df639..e68a52dedef0 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fallback.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -1,13 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-"""Unit tests for KVCacheManager._fit_token_budget.
-
-These exercise the prep-boundary token-budget fallback that defers or
-re-chunks context requests so a scheduled batch cannot overshoot
-``max_num_tokens`` in the forward pass (GitHub issue #13318). The fallback is
-pure scheduling logic and does not touch the GPU, so the tests build a bare
-KVCacheManager via ``__new__`` and drive the method with lightweight fake
-requests.
+"""Unit tests for KVCacheManager.fit_token_budget.
+
+These exercise the post-allocation token-budget trim that shrinks over-budget
+context chunks so a scheduled batch cannot overshoot ``max_num_tokens`` in the
+forward pass (GitHub issue #13318). The trim is pure scheduling logic and does
+not touch the GPU, so the tests build a bare KVCacheManager via ``__new__`` and
+drive the method with lightweight fake requests.
+
+The trim runs at the end of ``ResourceManager.prepare_resources``, which is the
+first point where ``context_current_position`` and ``context_chunk_size`` mean
+"forward-pass tokens" -- before ``setPrepopulatedPromptLen`` the chunk still
+spans the reusable KV prefix. ``TestReuseDiscountedChunk`` is the regression
+test for reading it too early.
 """
 
 import unittest
@@ -22,7 +27,7 @@
 
 
 class _FakeRequest:
-    """Minimal stand-in exposing only the attributes _fit_token_budget reads."""
+    """Minimal stand-in exposing only the attributes fit_token_budget reads."""
 
     _next_id = 0
 
@@ -42,15 +47,14 @@ def __init__(
         self.py_request_id = _FakeRequest._next_id
         # PyExecutor's inflight-set bookkeeping reads ``request_id`` (the C++
         # binding's name) rather than ``py_request_id``; keep them in sync so the
-        # same fake drives both the fallback and TestInflightIdsSurviveTrim.
+        # same fake drives both the trim and TestInflightIdsSurviveTrim.
         self.request_id = self.py_request_id
         self.context_chunk_size = context_chunk_size
         self.context_current_position = context_current_position
         # Mirrors the C++ semantics: is_last_context_chunk is a *computed*
         # property (context_current_position + context_chunk_size == prompt_len),
-        # so shrinking the chunk during re-chunk flips it to False. When
-        # prompt_len is None the flag is a fixed override (for tests that don't
-        # exercise re-chunk re-binning).
+        # so shrinking the chunk flips it to False. When prompt_len is None the
+        # flag is a fixed override (for tests that don't exercise re-binning).
         self._prompt_len = prompt_len
         self._is_last_override = is_last_context_chunk
         self.py_beam_width = py_beam_width
@@ -64,6 +68,14 @@ def is_last_context_chunk(self):
             return self._is_last_override
         return self.context_current_position + self.context_chunk_size == self._prompt_len
 
+    @property
+    def context_remaining_length(self):
+        # C++: mPromptLen - getContextCurrentPosition(). With no prompt_len the
+        # chunk is by definition all that is left.
+        if self._prompt_len is None:
+            return self.context_chunk_size
+        return self._prompt_len - self.context_current_position
+
 
 def _make_manager(max_num_tokens, tokens_per_block, enable_chunked_prefill=True):
     # Skip the heavy (GPU-allocating) __init__; the method under test only
@@ -71,11 +83,11 @@ def _make_manager(max_num_tokens, tokens_per_block, enable_chunked_prefill=True)
     mgr = KVCacheManager.__new__(KVCacheManager)
     mgr.max_num_tokens = max_num_tokens
     mgr.tokens_per_block = tokens_per_block
-    # Re-chunking is only valid when chunked prefill is enabled; otherwise the
-    # attention backend cannot consume a partial context chunk and the fallback
-    # must defer instead. Default to enabled so the re-chunk tests exercise that
-    # path; the disabled case is covered explicitly below.
+    # Shrinking produces a partial context chunk, which only chunked prefill's
+    # attention path can consume. Default to enabled; the disabled case is
+    # covered explicitly below.
     mgr.enable_chunked_prefill = enable_chunked_prefill
+    mgr.is_draft = False
     return mgr
 
 
@@ -87,11 +99,69 @@ def _make_batch(context_requests=(), generation_requests=()):
     return batch
 
 
+def _forward_tokens(mgr, batch):
+    """What _prepare_tp_inputs will materialize for this batch."""
+    return sum(
+        mgr._request_forward_tokens(r, is_context=False) for r in batch.generation_requests
+    ) + sum(
+        mgr._request_forward_tokens(r, is_context=True)
+        for r in batch.context_requests
+        if not r.is_disagg_generation_init_state
+    )
+
+
+class TestReuseDiscountedChunk(unittest.TestCase):
+    """Regression for the defect this trim shipped with.
+
+    Read before ``prepare_resources``, ``context_chunk_size`` spans the reusable
+    KV prefix: a 19212-token prompt with a 19200-token cache hit still reports
+    ``context_chunk_size == 19212`` while the forward pass will compute 12
+    tokens. Costing it at 19212 makes every reuse hit look 1600x too expensive,
+    so it is endlessly re-chunked (chunked prefill on) or deferred (off).
+
+    Read after ``prepare_resources`` -- where the trim now runs --
+    ``setPrepopulatedPromptLen`` has advanced ``context_current_position`` past
+    the prefix and the same request costs 12.
+    """
+
+    def test_reuse_hit_costs_only_the_uncached_tail(self):
+        mgr = _make_manager(max_num_tokens=8192, tokens_per_block=32)
+        # Post-setPrepopulatedPromptLen state for a 19212-token prompt with a
+        # 19200-token cache hit.
+        req = _FakeRequest(context_chunk_size=12, context_current_position=19200, prompt_len=19212)
+        self.assertEqual(mgr._request_forward_tokens(req, is_context=True), 12)
+
+    def test_reuse_hit_is_not_trimmed(self):
+        mgr = _make_manager(max_num_tokens=8192, tokens_per_block=32)
+        reqs = [
+            _FakeRequest(context_chunk_size=12, context_current_position=19200, prompt_len=19212)
+            for _ in range(4)
+        ]
+        batch = _make_batch(reqs, [_FakeRequest(py_beam_width=1) for _ in range(3)])
+
+        mgr.fit_token_budget(batch)
+
+        self.assertEqual(batch.num_context_requests, 4, "no request may be dropped")
+        for req in reqs:
+            self.assertEqual(req.context_chunk_size, 12, "chunk must be untouched")
+
+    def test_chunk_beyond_prompt_end_is_clamped(self):
+        # _prepare_tp_inputs slices all_prompt_tokens[pos:pos + chunk], and
+        # Python clamps that to the end of the list. A chunk that overhangs the
+        # prompt must be costed at what is actually left, not at its nominal
+        # size.
+        mgr = _make_manager(max_num_tokens=8192, tokens_per_block=32)
+        req = _FakeRequest(
+            context_chunk_size=8160, context_current_position=19200, prompt_len=19212
+        )
+        self.assertEqual(mgr._request_forward_tokens(req, is_context=True), 12)
+
+
 class TestFitTokenBudget(unittest.TestCase):
-    def test_request_forward_tokens_upper_bound(self):
+    def test_request_forward_tokens(self):
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
 
-        # Context: chunk size, plus draft tokens only on the last chunk.
+        # Context: materialized chunk, plus draft tokens only on the last chunk.
         last = _FakeRequest(
             context_chunk_size=10, is_last_context_chunk=True, py_draft_tokens=[1, 2]
         )
@@ -111,204 +181,224 @@ def test_within_budget_is_noop(self):
         gen = _FakeRequest(py_beam_width=100)  # 100 gen tokens
         batch = _make_batch([ctx], [gen])
 
-        mgr._fit_token_budget(batch)
+        mgr.fit_token_budget(batch)
 
         self.assertEqual(batch.num_context_requests, 1)
         self.assertEqual(ctx.context_chunk_size, 16)  # untouched
 
-    def test_overshoot_rechunks_context(self):
+    def test_overshoot_shrinks_context_to_fit(self):
         # 100 gen tokens leave a 28-token budget; a 64-token last chunk does not
-        # fit but can be re-chunked down to a block-aligned 16.
+        # fit and must be shrunk to the largest block-aligned chunk that does.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        gen = _FakeRequest(py_beam_width=100)
+        batch = _make_batch([ctx], [gen])
+
+        mgr.fit_token_budget(batch)
+
+        self.assertEqual(ctx.context_chunk_size, 16)
+        self.assertLessEqual(_forward_tokens(mgr, batch), 128)
+
+    def test_nothing_is_ever_dropped(self):
+        # The defining property of the post-allocation trim: KV is allocated and
+        # sequences are added, so a request may be shrunk but never removed.
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
-        ctx = _FakeRequest(context_chunk_size=64, is_last_context_chunk=True)
+        ctxs = [_FakeRequest(context_chunk_size=64, prompt_len=64) for _ in range(3)]
         gen = _FakeRequest(py_beam_width=100)
+        batch = _make_batch(ctxs, [gen])
+
+        mgr.fit_token_budget(batch)
+
+        self.assertEqual(batch.num_context_requests, 3)
+        for ctx in ctxs:
+            self.assertIn(ctx, batch.context_requests)
+
+    def test_shrink_keeps_chunk_end_block_aligned(self):
+        # setPrepopulatedPromptLen asserts (pos + chunk) % tokens_per_block == 0
+        # for every non-last chunk, to keep the KV cache unfragmented.
+        mgr = _make_manager(max_num_tokens=100, tokens_per_block=16)
+        ctx = _FakeRequest(context_chunk_size=96, context_current_position=32, prompt_len=128)
+        gen = _FakeRequest(py_beam_width=53)
         batch = _make_batch([ctx], [gen])
 
-        mgr._fit_token_budget(batch)
+        mgr.fit_token_budget(batch)
 
-        self.assertEqual(batch.num_context_requests, 1)
-        self.assertEqual(ctx.context_chunk_size, 16)  # (28 // 16) * 16
-        total = mgr._request_forward_tokens(ctx, is_context=True) + mgr._request_forward_tokens(
-            gen, is_context=False
+        self.assertLess(ctx.context_chunk_size, 96)
+        self.assertEqual(
+            (ctx.context_current_position + ctx.context_chunk_size) % 16,
+            0,
+            "chunk end must land on a block boundary",
         )
-        self.assertLessEqual(total, mgr.max_num_tokens)
-
-    def test_rechunk_only_rebins_to_chunking(self):
-        # Regression for the prep-boundary corruption (issue #13318 follow-up):
-        # when the overshoot is absorbed purely by re-chunking the *last*
-        # context request (no deferral), len(kept) is unchanged, but the request
-        # has flipped from last-chunk to non-last and MUST be moved out of the
-        # last-chunk bin. Otherwise downstream treats it as a final chunk and
-        # appends generation/draft tokens, corrupting the forward pass.
+
+    def test_shrink_never_produces_a_zero_token_chunk(self):
+        # A zero-token chunk leaves the request scheduled but computing nothing,
+        # which never terminates. One block of progress is the floor even when
+        # that overshoots the budget.
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
-        # Full prompt is 64 tokens, processed in one (last) chunk.
         ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
-        self.assertTrue(ctx.is_last_context_chunk)
-        gen = _FakeRequest(py_beam_width=100)  # remaining = 28
+        gen = _FakeRequest(py_beam_width=127)  # leaves a 1-token budget
         batch = _make_batch([ctx], [gen])
 
-        mgr._fit_token_budget(batch)
+        mgr.fit_token_budget(batch)
 
-        # Re-chunked to (28 // 16) * 16 == 16, now a non-last chunk.
         self.assertEqual(ctx.context_chunk_size, 16)
-        self.assertFalse(ctx.is_last_context_chunk)
-        # Count is unchanged, but it must have been re-binned into chunking.
-        self.assertEqual(batch.num_context_requests, 1)
-        self.assertIn(ctx, batch.context_requests_chunking)
-        self.assertNotIn(ctx, batch.context_requests_last_chunk)
-
-    def test_rechunk_drops_last_chunk_draft_tokens(self):
-        # Same re-chunk regression as above, but with draft tokens, which are
-        # appended only on the *last* chunk (see _request_forward_tokens). If a
-        # re-chunked request were left on the last-chunk path, its draft tokens
-        # would still be counted/materialized and re-introduce the overshoot
-        # this guard prevents. After re-chunking, the request must be a non-last
-        # chunk and its forward-token cost must no longer include the draft.
+
+    def test_shrink_rebins_to_chunking(self):
+        # Shrinking flips is_last_context_chunk to False, so the request must
+        # move out of context_requests_last_chunk -- otherwise downstream treats
+        # it as a final chunk and appends generation / draft tokens to it.
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
-        # 64-token last chunk + 2 draft tokens; a 28-token budget cannot fit
-        # 64 (+2), but the chunk re-chunks to a block-aligned 16.
-        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64, py_draft_tokens=[1, 2])
-        self.assertTrue(ctx.is_last_context_chunk)
-        gen = _FakeRequest(py_beam_width=100)  # remaining = 28
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        gen = _FakeRequest(py_beam_width=100)
         batch = _make_batch([ctx], [gen])
+        self.assertEqual(batch.context_requests_last_chunk, [ctx])
 
-        mgr._fit_token_budget(batch)
+        mgr.fit_token_budget(batch)
+
+        self.assertEqual(batch.context_requests_last_chunk, [])
+        self.assertEqual(batch.context_requests_chunking, [ctx])
+
+    def test_shrink_drops_last_chunk_draft_tokens(self):
+        # Draft tokens ride only on the last chunk, so a shrunk request stops
+        # contributing them and the budget must account for that.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64, py_draft_tokens=[1, 2, 3, 4])
+        gen = _FakeRequest(py_beam_width=100)
+        batch = _make_batch([ctx], [gen])
+        self.assertEqual(mgr._request_forward_tokens(ctx, is_context=True), 68)
+
+        mgr.fit_token_budget(batch)
 
-        # Re-chunked to (28 // 16) * 16 == 16 and flipped to a non-last chunk.
-        self.assertEqual(ctx.context_chunk_size, 16)
         self.assertFalse(ctx.is_last_context_chunk)
-        self.assertIn(ctx, batch.context_requests_chunking)
-        self.assertNotIn(ctx, batch.context_requests_last_chunk)
-        # Cost is now the chunk size alone -- the 2 draft tokens are dropped
-        # because the request is no longer the last chunk.
         self.assertEqual(mgr._request_forward_tokens(ctx, is_context=True), 16)
-        total = mgr._request_forward_tokens(ctx, is_context=True) + mgr._request_forward_tokens(
-            gen, is_context=False
-        )
-        self.assertLessEqual(total, mgr.max_num_tokens)
-
-    def test_overshoot_defers_when_chunked_prefill_disabled(self):
-        # Regression for the CI failures (q.numel()==0 / "Separate quantized
-        # buffer is not provided" / cudaErrorInvalidValue) seen in PR #15187:
-        # when chunked prefill is disabled the attention backend cannot consume
-        # a partial context chunk, so an over-budget request that *would* be
-        # re-chunkable must instead be deferred whole -- never re-chunked.
+
+    def test_sheds_the_last_chunk_first(self):
+        # context_requests is chunking + last_chunk, so the trim walks last-chunk
+        # requests first. That is the request whose cost the scheduler can have
+        # under-charged (only a last chunk carries a reuse discount), so it is
+        # the right one to repair; mid-prefill chunks are touched only if
+        # trimming it is not enough.
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        chunking = _FakeRequest(
+            context_chunk_size=96, prompt_len=256
+        )  # pos 0 + 96 != 256 -> chunking
+        last = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        batch = _make_batch([chunking, last], [])
+        self.assertEqual(batch.context_requests_chunking, [chunking])
+        self.assertEqual(batch.context_requests_last_chunk, [last])
+
+        mgr.fit_token_budget(batch)  # 160 tokens vs a 128 budget
+
+        self.assertEqual(chunking.context_chunk_size, 96, "mid-prefill untouched")
+        self.assertEqual(last.context_chunk_size, 32, "last chunk absorbed the excess")
+        self.assertLessEqual(_forward_tokens(mgr, batch), 128)
+
+    def test_shrinks_multiple_requests_when_one_is_not_enough(self):
+        mgr = _make_manager(max_num_tokens=64, tokens_per_block=16)
+        ctxs = [_FakeRequest(context_chunk_size=64, prompt_len=64) for _ in range(3)]
+        batch = _make_batch(ctxs, [])
+
+        mgr.fit_token_budget(batch)
+
+        self.assertLessEqual(_forward_tokens(mgr, batch), 64)
+        self.assertEqual(batch.num_context_requests, 3)
+
+    def test_no_shrink_when_chunked_prefill_disabled(self):
+        # A partial context chunk is only valid under chunked prefill; forcing
+        # one produces an invalid forward pass. Nothing safe is left to do.
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False)
-        # Same shape as test_overshoot_rechunks_context (a 28-token budget and a
-        # 64-token last chunk that is block-aligned re-chunkable to 16), but with
-        # chunked prefill off the request must be deferred, not shrunk.
         ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
-        gen = _FakeRequest(py_beam_width=100)  # remaining = 28
+        gen = _FakeRequest(py_beam_width=100)
         batch = _make_batch([ctx], [gen])
 
-        mgr._fit_token_budget(batch)
+        mgr.fit_token_budget(batch)
 
-        self.assertEqual(batch.num_context_requests, 0)
-        self.assertEqual(ctx.context_chunk_size, 64)  # not re-chunked
-        self.assertTrue(ctx.is_last_context_chunk)  # still a whole last chunk
+        self.assertEqual(ctx.context_chunk_size, 64)
+        self.assertEqual(batch.num_context_requests, 1)
 
-    def test_overshoot_defers_when_cannot_rechunk(self):
-        # Only an 8-token budget remains -- smaller than one block -- so the
-        # context request cannot be re-chunked and must be deferred entirely.
+    def test_mm_bidirectional_is_not_shrunk(self):
+        # Splitting a bidirectional multimodal block silently breaks attention.
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
-        ctx = _FakeRequest(context_chunk_size=64)
-        gen = _FakeRequest(py_beam_width=120)  # remaining = 8 < tokens_per_block
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64, mm_bidirectional=True)
+        gen = _FakeRequest(py_beam_width=100)
         batch = _make_batch([ctx], [gen])
 
-        mgr._fit_token_budget(batch)
+        mgr.fit_token_budget(batch)
 
-        self.assertEqual(batch.num_context_requests, 0)
-        self.assertEqual(ctx.context_chunk_size, 64)  # not re-chunked
+        self.assertEqual(ctx.context_chunk_size, 64)
 
-    def test_mm_bidirectional_is_deferred_not_rechunked(self):
-        # A re-chunkable budget exists, but splitting a bidirectional MM block
-        # would corrupt attention, so the request is deferred whole.
+    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)
-        ctx = _FakeRequest(context_chunk_size=64, mm_bidirectional=True)
-        gen = _FakeRequest(py_beam_width=100)  # remaining = 28
-        batch = _make_batch([ctx], [gen])
+        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)
+        mgr.fit_token_budget(batch)
 
-        self.assertEqual(batch.num_context_requests, 0)
-        self.assertEqual(ctx.context_chunk_size, 64)
+        self.assertEqual(disagg.context_chunk_size, 4096)
+        self.assertEqual(ctx.context_chunk_size, 16)
 
-    def test_defers_all_subsequent_context_requests(self):
-        # ctx1 fits; ctx2 overshoots and cannot re-chunk; ctx3 (small) must
-        # still be deferred to preserve context-progress ordering.
-        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
-        ctx1 = _FakeRequest(context_chunk_size=96)
-        ctx2 = _FakeRequest(context_chunk_size=64)
-        ctx3 = _FakeRequest(context_chunk_size=16)
-        gen = _FakeRequest(py_beam_width=16)  # remaining = 112
-        batch = _make_batch([ctx1, ctx2, ctx3], [gen])
+    def test_gen_only_batch_is_left_alone(self):
+        # Generation cannot be shed, so a batch with no context requests returns
+        # immediately -- keeping the executor loop's hottest path free of an
+        # O(num_generation_requests) scan.
+        mgr = _make_manager(max_num_tokens=8, tokens_per_block=16)
+        batch = _make_batch([], [_FakeRequest(py_beam_width=64)])
+
+        mgr.fit_token_budget(batch)  # must not raise
+
+        self.assertEqual(len(batch.generation_requests), 1)
+
+    def test_generation_alone_over_budget_does_not_raise(self):
+        # There is nothing to shed, but raising here would be rank-local and
+        # would deadlock the surviving ranks under attention DP. Warn and let
+        # the forward pass report it.
+        mgr = _make_manager(max_num_tokens=64, tokens_per_block=16)
+        ctx = _FakeRequest(context_chunk_size=16, prompt_len=16)
+        gen = _FakeRequest(py_beam_width=100)
+        batch = _make_batch([ctx], [gen])
 
-        mgr._fit_token_budget(batch)
+        mgr.fit_token_budget(batch)  # must not raise
 
-        # ctx1 (96) fits into 112; remaining 16. ctx2 (64) doesn't fit and
-        # (16 // 16) * 16 == 16 but 16 < 64 so it *could* re-chunk to 16...
-        # remaining after is 0, so ctx3 is deferred.
-        kept = batch.context_requests
-        self.assertIn(ctx1, kept)
-        self.assertNotIn(ctx3, kept)
+        self.assertEqual(batch.num_context_requests, 1)
 
     def test_maybe_fit_token_budget_skips_draft_manager(self):
-        # maybe_fit_token_budget is the single entry point driven by the
-        # aggregate ResourceManager. It must apply the fallback for the target
-        # manager only -- the draft-model engine builds inputs with a different
-        # token shape and its budget is handled separately.
-        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
-        gen = _FakeRequest(py_beam_width=120)  # remaining = 8 -> defer ctx
+        # The draft-model engine builds inputs with a different token shape and
+        # its budget is handled separately.
+        gen = _FakeRequest(py_beam_width=100)
 
-        # Non-draft -> defers.
-        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False)
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
         mgr.is_draft = False
-        batch = _make_batch([ctx], [gen])
-        mgr.maybe_fit_token_budget(batch)
-        self.assertEqual(batch.num_context_requests, 0)
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        mgr.maybe_fit_token_budget(_make_batch([ctx], [gen]))
+        self.assertEqual(ctx.context_chunk_size, 16)
 
-        # Draft manager -> never fits (handled separately).
         mgr.is_draft = True
-        batch = _make_batch([ctx], [gen])
-        mgr.maybe_fit_token_budget(batch)
-        self.assertEqual(batch.num_context_requests, 1)
-
-    def test_fallback_runs_before_other_managers(self):
-        # Regression for the emplaceDone double-add (PR #15187): the token-budget
-        # fallback must mutate scheduled_batch BEFORE any resource manager
-        # allocates sequences. A separate draft KV cache manager (MTP) is
-        # invoked before the target KV cache manager (the target is moved to the
-        # end of the manager dict on purpose), so if the fallback ran inside the
-        # target's own prepare_resources the draft manager would already have
-        # added sequences for context requests the fallback then defers --
-        # orphaning them and causing a double-add when they reschedule.
-        #
-        # The executor loop drives the fallback via
-        # ResourceManager.maybe_fit_token_budget before it calls
-        # prepare_resources; build the aggregate ResourceManager with the same
-        # ordering as production (draft-like manager first, KV cache manager
-        # last) and assert the earlier manager observes the *already-deferred*
-        # batch.
-        target = _make_manager(
-            max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False
-        )
-        target.is_draft = False
-        # Don't touch the GPU: only the budget fallback matters for ordering.
-        target.prepare_resources = lambda batch: None
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        mgr.maybe_fit_token_budget(_make_batch([ctx], [gen]))
+        self.assertEqual(ctx.context_chunk_size, 64)
 
+    def test_trim_runs_after_every_manager(self):
+        # The trim must observe the batch as prepare_resources leaves it: only
+        # after addSequence has run does context_chunk_size mean forward-pass
+        # tokens (setPrepopulatedPromptLen advances context_current_position
+        # past the reusable prefix). Registering the KV cache manager last
+        # mirrors _util.py's move_to_end(KV_CACHE_MANAGER); the trim must still
+        # run after that.
         observed = []
 
+        target = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        target.prepare_resources = lambda batch: observed.append("kv_cache_manager")
+
         class _RecordingManager:
             def prepare_resources(self, batch):
-                observed.append([r.py_request_id for r in batch.context_requests])
+                observed.append("draft_manager")
 
-        ctx_keep = _FakeRequest(context_chunk_size=96)
-        ctx_defer = _FakeRequest(context_chunk_size=64)
-        gen = _FakeRequest(py_beam_width=16)  # remaining = 112
-        batch = _make_batch([ctx_keep, ctx_defer], [gen])
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        gen = _FakeRequest(py_beam_width=100)
+        batch = _make_batch([ctx], [gen])
 
-        # Draft-like manager registered FIRST, KV cache manager LAST (mirrors
-        # _util.py's move_to_end(KV_CACHE_MANAGER)).
         rm = ResourceManager(
             OrderedDict(
                 [
@@ -317,78 +407,33 @@ def prepare_resources(self, batch):
                 ]
             )
         )
-        rm.maybe_fit_token_budget(batch)
         rm.prepare_resources(batch)
 
-        # ctx_keep (96) fits into 112; ctx_defer (64) does not and is deferred.
-        # The draft-like manager, though invoked first, must have seen only the
-        # kept request -- proving the fallback ran up front.
-        self.assertEqual(observed, [[ctx_keep.py_request_id]])
-        self.assertEqual(batch.num_context_requests, 1)
+        self.assertEqual(observed, ["draft_manager", "kv_cache_manager"])
+        self.assertEqual(ctx.context_chunk_size, 16, "trim ran after both managers")
 
-    def test_prepare_resources_does_not_trim(self):
-        # The fallback must NOT be reachable from prepare_resources: the
-        # executor loop drives it earlier so that _can_queue's attention-DP
-        # tp_allgather(batch_size) and the PP inflight-set registration both see
-        # the trimmed batch. A second trim here would be redundant at best, and
-        # leaving it as the *only* trim point is what let a rank shed its way to
-        # an empty batch after the ranks had already voted to run.
-        target = _make_manager(
-            max_num_tokens=128, tokens_per_block=16, enable_chunked_prefill=False
-        )
-        target.is_draft = False
+    def test_prepare_resources_trims(self):
+        # prepare_resources is the only entry point; the executor loops must not
+        # need their own call.
+        target = _make_manager(max_num_tokens=128, tokens_per_block=16)
         target.prepare_resources = lambda batch: None
 
-        ctx_over_budget = _FakeRequest(context_chunk_size=64)
-        gen = _FakeRequest(py_beam_width=100)  # remaining = 28, so ctx cannot fit
-        batch = _make_batch([ctx_over_budget], [gen])
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        gen = _FakeRequest(py_beam_width=100)
+        batch = _make_batch([ctx], [gen])
 
         rm = ResourceManager(OrderedDict([(ResourceManagerType.KV_CACHE_MANAGER, target)]))
         rm.prepare_resources(batch)
 
-        self.assertEqual(batch.num_context_requests, 1)
-
-        # ...and the batch is trimmed only once maybe_fit_token_budget is called.
-        rm.maybe_fit_token_budget(batch)
-        self.assertEqual(batch.num_context_requests, 0)
-
-    def test_generation_alone_over_budget_raises(self):
-        # A context request must be present for the fallback to engage at all
-        # (see test_gen_only_batch_is_left_alone); generation requests that
-        # exceed the budget by themselves leave nothing to shed.
-        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
-        gen = _FakeRequest(py_beam_width=200)
-        ctx = _FakeRequest(context_chunk_size=16)
-        batch = _make_batch([ctx], [gen])
-
-        with self.assertRaises(RuntimeError):
-            mgr._fit_token_budget(batch)
-
-    def test_gen_only_batch_is_left_alone(self):
-        # Context requests are the only thing the fallback can shed, so a
-        # gen-only batch returns before the generation-token accounting: it
-        # keeps that scan off the executor loop's hottest path, and it avoids
-        # raising on a condition nothing here can fix. An over-budget gen-only
-        # batch stays the concern of the _prepare_tp_inputs assert, which fails
-        # one batch rather than killing the (possibly only rank-local) event
-        # loop.
-        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
-        gen = _FakeRequest(py_beam_width=200)  # 200 > max_num_tokens
-        batch = _make_batch([], [gen])
-
-        mgr._fit_token_budget(batch)  # must not raise
-
-        self.assertEqual(batch.num_context_requests, 0)
-        self.assertEqual(batch.generation_requests, [gen])
+        self.assertEqual(ctx.context_chunk_size, 16)
 
 
 class TestInflightIdsSurviveTrim(unittest.TestCase):
-    """The fallback must not strand ids in PyExecutor's inflight set.
+    """The trim must not strand ids in PyExecutor's inflight set.
 
     ``_executor_loop_pp`` calls ``_add_inflight_ids`` before
     ``ResourceManager.prepare_resources`` and ``_remove_inflight_ids`` after, so
-    the fallback runs between them and mutates the batch: a deferred context
-    request is dropped from it, and a re-chunked one moves out of
+    the trim runs between them and moves shrunk requests out of
     ``context_requests_last_chunk``. Removal must therefore erase the ids that
     were actually inserted, not re-derive them from the trimmed batch -- an id
     left behind makes the scheduler skip that request forever (scheduler.py's
@@ -406,35 +451,32 @@ def _bare_executor():
         executor.inflight_req_ids = ReqIdsSet()
         return executor
 
-    def test_trimmed_context_requests_leave_no_inflight_ids(self):
+    def test_shrunk_context_requests_leave_no_inflight_ids(self):
         executor = self._bare_executor()
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
 
         gen = _FakeRequest(py_beam_width=100)  # leaves a 28-token budget
-        # Both start as last-chunk requests, so both are registered inflight.
-        # ctx_rechunked shrinks to a block-aligned 16 and flips to a non-last
-        # chunk; ctx_deferred no longer fits and is dropped from the batch.
-        ctx_rechunked = _FakeRequest(context_chunk_size=64, prompt_len=64)
-        ctx_deferred = _FakeRequest(context_chunk_size=64, prompt_len=64)
-        batch = _make_batch([ctx_rechunked, ctx_deferred], [gen])
+        # Starts as a last-chunk request, so it is registered inflight; the trim
+        # then shrinks it to a block-aligned 16 and it stops being a last chunk.
+        ctx = _FakeRequest(context_chunk_size=64, prompt_len=64)
+        batch = _make_batch([ctx], [gen])
 
         executor._add_inflight_ids(batch)
         self.assertEqual(
             sorted(batch.added_inflight_req_ids),
-            sorted([ctx_rechunked.request_id, ctx_deferred.request_id, gen.request_id]),
+            sorted([ctx.request_id, gen.request_id]),
         )
 
-        mgr._fit_token_budget(batch)
+        mgr.fit_token_budget(batch)
 
-        # Preconditions for the regression: the batch really did change shape.
-        self.assertEqual(ctx_rechunked.context_chunk_size, 16)
-        self.assertIn(ctx_rechunked, batch.context_requests_chunking)
+        # Precondition for the regression: the batch really did change shape.
+        self.assertEqual(ctx.context_chunk_size, 16)
+        self.assertIn(ctx, batch.context_requests_chunking)
         self.assertEqual(batch.context_requests_last_chunk, [])
-        self.assertNotIn(ctx_deferred, batch.context_requests)
 
         executor._remove_inflight_ids(batch)
 
-        for req in (ctx_rechunked, ctx_deferred, gen):
+        for req in (ctx, gen):
             self.assertNotIn(
                 req.request_id,
                 executor.inflight_req_ids,
@@ -444,7 +486,7 @@ def test_trimmed_context_requests_leave_no_inflight_ids(self):
         self.assertEqual(batch.added_inflight_req_ids, [])
 
     def test_untrimmed_batch_round_trips(self):
-        # The batch the fallback leaves alone must behave exactly as before.
+        # The batch the trim leaves alone must behave exactly as before.
         executor = self._bare_executor()
         mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
 
@@ -456,7 +498,7 @@ def test_untrimmed_batch_round_trips(self):
         for req in (ctx, gen):
             self.assertIn(req.request_id, executor.inflight_req_ids)
 
-        mgr._fit_token_budget(batch)
+        mgr.fit_token_budget(batch)
         self.assertEqual(batch.num_context_requests, 1)
 
         executor._remove_inflight_ids(batch)

From 5c63043690e093bdaaf40edd0cc8912a4c154357 Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Thu, 6 Aug 2026 19:14:55 +0000
Subject: [PATCH 15/16] [#13318][chore] Correct stale comment on
 enable_chunked_prefill

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>
---
 tensorrt_llm/_torch/pyexecutor/resource_manager.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index b955e4c20400..ab00e7c89dea 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -394,11 +394,11 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
         # Kept so prepare_resources can re-validate the per-step token budget
         # (the forward-pass scratch size enforced in _prepare_tp_inputs).
         self.max_num_tokens = max_num_tokens
-        # 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
+        # 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: defer instead of re-chunk) and is set to the
-        # finalized value by _create_kv_cache_manager.
+        # Defaults to False (safe: leave the chunk at its scheduled size) and is
+        # set to the finalized value by _create_kv_cache_manager.
         self.enable_chunked_prefill = enable_chunked_prefill
         self.event_buffer_max_size = kv_cache_config.event_buffer_max_size
         self.attention_dp_events_gather_period_ms = kv_cache_config.attention_dp_events_gather_period_ms

From b026f501a70e65317a7ae7ab619b929596b14637 Mon Sep 17 00:00:00 2001
From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Date: Thu, 6 Aug 2026 20:23:25 +0000
Subject: [PATCH 16/16] [#13318][fix] Report the trimmed batch to the KV
 connector

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_llm/_torch/pyexecutor/py_executor.py | 10 +++
 .../_torch/pyexecutor/resource_manager.py     | 20 +++++
 .../executor/test_token_budget_fallback.py    | 84 +++++++++++++++++++
 3 files changed, 114 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 327ded2adfdd..ab73e4245641 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -5982,6 +5982,16 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests):
                         resource_mgr_type].prepare_resources(
                             disagg_gen_init_to_prepare)
 
+            # Reporting this mini-batch to the KV connector used to happen
+            # inside KVCacheManager.prepare_resources; it now runs after the
+            # token-budget trim, which this path does not go through. Kept here
+            # so the connector's per-request state advances exactly as before.
+            kv_cache_manager = self.resource_manager.resource_managers.get(
+                ResourceManagerType.KV_CACHE_MANAGER)
+            if hasattr(kv_cache_manager, "publish_connector_scheduler_output"):
+                kv_cache_manager.publish_connector_scheduler_output(
+                    disagg_gen_init_to_prepare)
+
             # Trigger KV cache exchange for new disagg_gen_init_requests
             self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests)
 
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index db31eaeb0396..d86560800387 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -1046,6 +1046,19 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
         # reuse, so we rebuild the context request lists here.
         scheduled_batch.reset_context_requests()
 
+    def publish_connector_scheduler_output(
+            self, scheduled_batch: ScheduledRequests) -> None:
+        """Report the batch to the KV connector.
+
+        Driven by ``ResourceManager.prepare_resources`` *after* the token-budget
+        trim rather than from the end of ``prepare_resources`` above, because
+        ``RequestData.num_scheduled_tokens`` is documented as "the number of
+        scheduled tokens for the upcoming forward pass" and is built from
+        ``context_chunk_size`` (``connectors/kv_cache_connector.py``). Reporting
+        before the trim over-states it for every chunk the trim shrinks, and a
+        connector that decides what to save or offload from that count would
+        publish KV for tokens the forward pass never computed.
+        """
         if self.kv_connector_manager is not None:
             self.kv_connector_manager.build_scheduler_output(
                 scheduled_batch, self)
@@ -2867,6 +2880,13 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
         # After every manager, so context_current_position / context_chunk_size
         # are final. See maybe_fit_token_budget.
         self.maybe_fit_token_budget(scheduled_batch)
+        # Strictly after the trim: the connector is told how many tokens the
+        # forward pass will compute, which is only settled once the trim has
+        # run. See KVCacheManager.publish_connector_scheduler_output.
+        kv_cache_manager = self.resource_managers.get(
+            ResourceManagerType.KV_CACHE_MANAGER)
+        if hasattr(kv_cache_manager, "publish_connector_scheduler_output"):
+            kv_cache_manager.publish_connector_scheduler_output(scheduled_batch)
 
     @nvtx_range("update_resources")
     def update_resources(
diff --git a/tests/unittest/_torch/executor/test_token_budget_fallback.py b/tests/unittest/_torch/executor/test_token_budget_fallback.py
index e68a52dedef0..e45827e238da 100644
--- a/tests/unittest/_torch/executor/test_token_budget_fallback.py
+++ b/tests/unittest/_torch/executor/test_token_budget_fallback.py
@@ -88,6 +88,9 @@ def _make_manager(max_num_tokens, tokens_per_block, enable_chunked_prefill=True)
     # covered explicitly below.
     mgr.enable_chunked_prefill = enable_chunked_prefill
     mgr.is_draft = False
+    # Read by publish_connector_scheduler_output; most tests run without a
+    # connector attached.
+    mgr.kv_connector_manager = None
     return mgr
 
 
@@ -506,5 +509,86 @@ def test_untrimmed_batch_round_trips(self):
             self.assertNotIn(req.request_id, executor.inflight_req_ids)
 
 
+class TestConnectorSeesTheTrimmedBatch(unittest.TestCase):
+    """The KV connector must see the batch that will actually run.
+
+    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 reporting the batch before the trim over-states it for every chunk the
+    trim shrinks -- and a connector that decides what to save or offload from
+    that count would publish KV for tokens the forward pass never computed.
+    """
+
+    class _RecordingKvManager:
+        """Stands in for KVCacheManager: records when it is asked to publish."""
+
+        def __init__(self, log, batch, shrink_to):
+            self._log = log
+            self._batch = batch
+            self._shrink_to = shrink_to
+
+        def prepare_resources(self, scheduled_batch):
+            self._log.append(("prepare", self._chunks()))
+
+        def maybe_fit_token_budget(self, scheduled_batch):
+            for req in scheduled_batch.context_requests:
+                req.context_chunk_size = self._shrink_to
+            self._log.append(("trim", self._chunks()))
+
+        def publish_connector_scheduler_output(self, scheduled_batch):
+            self._log.append(("publish", self._chunks()))
+
+        def _chunks(self):
+            return [r.context_chunk_size for r in self._batch.context_requests]
+
+    def _resource_manager(self, log, batch, shrink_to):
+        return ResourceManager(
+            OrderedDict(
+                [
+                    (
+                        ResourceManagerType.KV_CACHE_MANAGER,
+                        self._RecordingKvManager(log, batch, shrink_to),
+                    )
+                ]
+            )
+        )
+
+    def test_the_connector_is_told_after_the_trim(self):
+        log = []
+        req = _FakeRequest(context_chunk_size=4096, prompt_len=4096)
+        batch = _make_batch([req])
+
+        self._resource_manager(log, batch, shrink_to=64).prepare_resources(batch)
+
+        self.assertEqual([step for step, _ in log], ["prepare", "trim", "publish"])
+        # The count the connector sees is the trimmed one, not the scheduled one.
+        self.assertEqual(dict(log)["publish"], [64])
+
+    def test_managers_without_a_connector_hook_are_skipped(self):
+        rm = ResourceManager(OrderedDict([(ResourceManagerType.KV_CACHE_MANAGER, object())]))
+        rm.prepare_resources(_make_batch())  # must not raise
+
+    def test_publishing_is_a_no_op_without_a_connector(self):
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        mgr.kv_connector_manager = None
+        mgr.publish_connector_scheduler_output(_make_batch())  # must not raise
+
+    def test_publishing_forwards_the_batch_to_the_connector(self):
+        class _FakeConnector:
+            def __init__(self):
+                self.calls = []
+
+            def build_scheduler_output(self, scheduled_batch, kv_cache_manager):
+                self.calls.append((scheduled_batch, kv_cache_manager))
+
+        mgr = _make_manager(max_num_tokens=128, tokens_per_block=16)
+        mgr.kv_connector_manager = _FakeConnector()
+        batch = _make_batch([_FakeRequest(context_chunk_size=16, prompt_len=16)])
+
+        mgr.publish_connector_scheduler_output(batch)
+
+        self.assertEqual(mgr.kv_connector_manager.calls, [(batch, mgr)])
+
+
 if __name__ == "__main__":
     unittest.main()