[#13318][fix] Gracefully fit token budget at prep boundary

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

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