[None][fix] Cherry-pick Python-side remaining_budget guard for KV reuse budget overflow

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -2678,6 +2678,24 @@ 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:
+            ctx_details = []
+            for r in scheduled_requests.context_requests:
+                pos = r.context_current_position
+                csz = r.context_chunk_size
+                full = len(r.get_tokens(0))
+                tokens = min(csz, max(0, full - pos))
+                ctx_details.append(
+                    f"rid={r.py_request_id} pos={pos} chunk={csz} "
+                    f"full={full} tokens={tokens}")
+            gen_count = len(scheduled_requests.generation_requests)
+            from tensorrt_llm.logger import logger as _mnt_logger
+            _mnt_logger.error(
+                f"MNT overflow: total={total_num_tokens} "
+                f"max={self.max_num_tokens} "
+                f"ctx_reqs={len(scheduled_requests.context_requests)} "
+                f"gen_reqs={gen_count} "
+                f"ctx_breakdown=[{'; '.join(ctx_details)}]")
         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})"
         )

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -3094,6 +3094,78 @@ def _compute_scheduled_tokens(context_requests, generation_requests):
                                        for gen_req in generation_requests)
         return num_scheduled_ctx_tokens + num_scheduled_gen_tokens
 
+    def _maybe_log_batch_wait_decision(
+        self,
+        context_requests: list[LlmRequest],
+        generation_requests: list[LlmRequest],
+        num_scheduled_tokens: int,
+        wait_threshold: float,
+        should_waiting: bool,
+    ) -> None:
+        """Emit per-iteration batch-wait diagnostics on rank 0.
+
+        Logs the scheduling-formula token count next to the actual chunk-token
+        sum so that mismatches (e.g., from KV cache reuse estimation drift) are
+        visible in the rank-0 stream. Other ranks early-return.
+        """
+        if self.dist.rank != 0:
+            return
+
+        num_scheduled_gen_tokens = sum(1 + gen_req.num_draft_tokens
+                                       for gen_req in generation_requests)
+        num_scheduled_ctx_formula = num_scheduled_tokens - num_scheduled_gen_tokens
+
+        chunk_ctx_sum = 0
+        ctx_summaries: List[str] = []
+        max_detail = 4
+        for i, ctx_req in enumerate(context_requests):
+            full_len = len(ctx_req.get_tokens(0))
+            begin = ctx_req.context_current_position
+            chunk_sz = ctx_req.context_chunk_size
+            this_chunk = min(chunk_sz, max(0, full_len - begin))
+            chunk_ctx_sum += this_chunk
+            reusable = (ctx_req.estimated_reusable_tokens
+                        if ctx_req.is_first_context_chunk else 0)
+            reusable_in_chunk = max(0, reusable - begin)
+            remaining = ctx_req.context_remaining_length
+            if (reusable_in_chunk > 0
+                    and reusable_in_chunk + chunk_sz < remaining):
+                formula_contrib = chunk_sz
+            else:
+                formula_contrib = max(1, chunk_sz - reusable_in_chunk)
+            if i < max_detail:
+                ctx_summaries.append(
+                    f"rid={ctx_req.py_request_id} full={full_len} pos={begin} "
+                    f"chunk_sz={chunk_sz} this_chunk={this_chunk} "
+                    f"reusable={reusable} formula_contrib={formula_contrib}")
+        n_ctx = len(context_requests)
+        if n_ctx > max_detail:
+            ctx_summaries.append(f"... +{n_ctx - max_detail} more ctx req(s)")
+
+        logger.info(
+            "batch_wait: formula_total=",
+            num_scheduled_tokens,
+            " formula_ctx=",
+            num_scheduled_ctx_formula,
+            " formula_gen=",
+            num_scheduled_gen_tokens,
+            " chunk_ctx_sum=",
+            chunk_ctx_sum,
+            " threshold=",
+            wait_threshold,
+            " wait_iter=",
+            self.batch_wait_iters_count,
+            "/",
+            self.batch_wait_timeout_iters,
+            " should_defer_ctx=",
+            should_waiting,
+            " num_gen=",
+            len(generation_requests),
+            " ctx_detail=[",
+            "; ".join(ctx_summaries),
+            "]",
+        )
+
     def _waiting_requests(self, context_requests: list[LlmRequest],
                           generation_requests: list[LlmRequest]):
         """
@@ -3105,8 +3177,14 @@ def _waiting_requests(self, context_requests: list[LlmRequest],
 
         num_scheduled_tokens = self._compute_scheduled_tokens(
             context_requests, generation_requests)
-
-        should_waiting = self.batch_wait_iters_count < self.batch_wait_timeout_iters and num_scheduled_tokens < self.batch_wait_max_tokens_ratio * self.max_num_tokens
+        wait_threshold = (self.batch_wait_max_tokens_ratio *
+                          self.max_num_tokens)
+
+        should_waiting = self.batch_wait_iters_count < self.batch_wait_timeout_iters and num_scheduled_tokens < wait_threshold
+        self._maybe_log_batch_wait_decision(context_requests,
+                                            generation_requests,
+                                            num_scheduled_tokens,
+                                            wait_threshold, should_waiting)
         if should_waiting:
             self.batch_wait_iters_count += 1
             return []

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -703,6 +703,30 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
             # wait for all pending work to finish before launching offload/onboarding/partial copy
             self.impl.sync_transfer_manager_with_buffer_manager()
 
+            # Pre-addSequence budget re-validation.  The C++ scheduler
+            # should already account for the chunk-shift cost, but under
+            # heavy KV-cache eviction the actual reuse may be lower than
+            # estimated.  We re-probe the radix tree and estimate the
+            # true forward cost; if it exceeds the remaining budget the
+            # request is skipped (re-scheduled next iteration).
+            remaining_budget = None
+            if self.enable_block_reuse and not self.is_draft:
+                gen_tokens = sum(
+                    req.get_beam_width_by_iter(for_next_iteration=False) +
+                    get_draft_token_length(req)
+                    for req in scheduled_batch.generation_requests)
+                remaining_budget = self.max_num_tokens - gen_tokens
+
+                # Pre-subtract the fixed cost of non-first-chunk context
+                # requests.  These have no reuse to re-validate and their
+                # compute cost is committed, so first-chunk budget checks
+                # must see the budget with these costs already removed.
+                for req in scheduled_batch.context_requests:
+                    if not req.is_first_context_chunk:
+                        remaining_budget -= req.context_chunk_size
+
+            accepted_ctx_requests = []
+
             # Collect first-chunk requests eligible for add_sequence_batch.
             # When block reuse is enabled, addSequenceBatch uses a two-phase
             # claim-then-onboard strategy that prevents host offloading from
@@ -731,11 +755,42 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                 else:
                     if req.is_first_context_chunk and self._kv_connector_should_add_sequence(
                             req):
+                        if remaining_budget is not None:
+                            unique_tokens = req.get_unique_tokens(0)
+                            # reusable_blocks_all counts every block that
+                            # matches the prefix, including blocks not yet
+                            # allocated to a request — the right basis for
+                            # the budget re-probe since any of them could be
+                            # claimed before add_sequence_batch runs.
+                            reusable_blocks = self.impl.analyze_prefix_reuse(
+                                unique_tokens, req).reusable_blocks_all
+                            actual_reuse = (reusable_blocks *
+                                            self.tokens_per_block)
+                            req_compute = self._estimate_post_reuse_compute(
+                                actual_reuse, req.context_chunk_size,
+                                req.prompt_len)
+                            if req_compute > remaining_budget:
+                                logger.warning(
+                                    f"Reuse budget: skip req "
+                                    f"{req.py_request_id} "
+                                    f"(compute={req_compute}, "
+                                    f"chunk={req.context_chunk_size}, "
+                                    f"reuse={actual_reuse}, "
+                                    f"remaining={remaining_budget})")
+                                continue
+                            remaining_budget -= req_compute
+
                         # Batch path: two-phase claim-then-onboard
                         batch_request_infos.append(
                             (req.py_request_id, req.prompt_len, req_beam_width))
                         batch_llm_requests.append(req)
                         batch_ctx_requests.append(req)
+                    elif remaining_budget is not None and req.is_first_context_chunk:
+                        reusable = req.estimated_reusable_tokens
+                        remaining_budget -= self._estimate_post_reuse_compute(
+                            reusable, req.context_chunk_size, req.prompt_len)
+
+                accepted_ctx_requests.append(req)
 
             if batch_request_infos:
                 self.impl.add_sequence_batch(batch_request_infos,
@@ -753,7 +808,7 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
 
             # A request may change from `context_requests_chunking` to `context_requests_last_chunk` in
             # `add_sequence_batch` due to KV cache reuse, so we rebuild the context request lists here.
-            scheduled_batch.reset_context_requests()
+            scheduled_batch.reset_context_requests(accepted_ctx_requests)
 
             for req in scheduled_batch.generation_requests:
                 if self.mapping.has_cp_helix():
@@ -779,6 +834,26 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
             self.kv_connector_manager.build_scheduler_output(
                 scheduled_batch, self)
 
+    def _estimate_post_reuse_compute(self, reuse_tokens: int, chunk_size: int,
+                                     prompt_len: int) -> int:
+        """Estimate forward compute tokens for a context chunk given how many
+        prefix tokens will be served from KV cache reuse.
+
+        For non-last chunks the chunk window shifts right by the reused amount
+        and the forward cost is approximately ``chunk_size`` (rounded down to
+        a ``tokens_per_block`` boundary).  For last chunks the cost is
+        ``prompt_len - reuse``.  Returns ``chunk_size`` verbatim when reuse is
+        out of range (``reuse <= 0`` or ``reuse >= prompt_len``).
+        """
+        P = reuse_tokens
+        if P > 0 and P < prompt_len:
+            if P + chunk_size < prompt_len:
+                aligned_end = ((P + chunk_size) // self.tokens_per_block *
+                               self.tokens_per_block)
+                return max(1, aligned_end - P)
+            return max(1, prompt_len - P)
+        return chunk_size
+
     def extend_capacity_for_tokens(self, request: LlmRequest) -> None:
         """No-op for V1; interface kept consistent with V2."""
 

tests/unittest/_torch/executor/test_batch_wait_log.py

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+"""Smoke test for ``PyExecutor._maybe_log_batch_wait_decision`` rank gating.
+
+The method early-returns when ``self.dist.rank != 0`` and otherwise reads
+``self.batch_wait_iters_count`` and ``self.batch_wait_timeout_iters`` to
+format diagnostic output. If a refactor removes the rank gate, the function
+would crash on the wait counters since this test never sets them on the fake.
+"""
+
+from unittest.mock import Mock
+
+from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+
+
+def test_log_decision_early_returns_on_nonzero_rank():
+    """Non-rank-0 path early-returns without touching wait counters."""
+    fake = PyExecutor.__new__(PyExecutor)
+    fake.dist = Mock()
+    fake.dist.rank = 1  # non-zero rank → early return
+
+    # No exception expected. If the gate is broken, the function would proceed
+    # to read attributes (``batch_wait_iters_count``, ``batch_wait_timeout_iters``)
+    # that we never set on the fake, raising AttributeError.
+    PyExecutor._maybe_log_batch_wait_decision(
+        fake,
+        context_requests=[],
+        generation_requests=[],
+        num_scheduled_tokens=0,
+        wait_threshold=0.5,
+        should_waiting=False,
+    )