diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 42acbda93146..efb02ded1f21 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -89,6 +89,11 @@ # Default: "0" (only rank 0 prints, matching existing behavior). PROFILE_LOG_RANKS_ENV_VAR_NAME = "TLLM_PROFILE_LOG_RANKS" +# C++ LlmRequest.pause() requires a prompt-length cap. Recompute pause should +# replay all generated tokens in PyTorch instead of inheriting TRT build-time +# max_input_len truncation. +_UNBOUNDED_PAUSE_MAX_INPUT_LEN = 0x7fffffff + class PPCommTag(IntEnum): """ @@ -1234,8 +1239,9 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests, stats.inflight_batching_stats.num_context_requests = scheduled_batch.num_context_requests stats.inflight_batching_stats.num_gen_requests = scheduled_batch.num_generation_requests stats.inflight_batching_stats.num_scheduled_requests = stats.inflight_batching_stats.num_context_requests + stats.inflight_batching_stats.num_gen_requests - stats.inflight_batching_stats.num_paused_requests = len( - scheduled_batch.paused_requests) + paused_requests = (scheduled_batch.paused_requests + + scheduled_batch.recompute_paused_requests) + stats.inflight_batching_stats.num_paused_requests = len(paused_requests) stats.inflight_batching_stats.avg_num_decoded_tokens_per_iter = 0 stats.inflight_batching_stats.micro_batch_id = micro_batch_id @@ -1307,7 +1313,7 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests, # RuntimeError on a mutated request. num_ctx_kv_tokens = 0 for req in scheduled_batch.context_requests: - if getattr(req, "is_attention_dp_dummy", False): + if req.is_attention_dp_dummy: continue last_chunk = getattr(req, "py_last_context_chunk", None) if last_chunk is not None and last_chunk[0] is not None: @@ -1324,7 +1330,7 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests, # summed across scheduled generation requests. num_gen_kv_tokens = 0 for req in scheduled_batch.generation_requests: - if getattr(req, "is_attention_dp_dummy", False): + if req.is_attention_dp_dummy: continue try: num_gen_kv_tokens += req.get_num_tokens(0) @@ -1375,8 +1381,8 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests, # requests — were decoding but got evicted back to the waiting # pool for this iteration. num_paused_kv_tokens = 0 - for req in scheduled_batch.paused_requests: - if getattr(req, "is_attention_dp_dummy", False): + for req in paused_requests: + if req.is_attention_dp_dummy: continue try: num_paused_kv_tokens += req.get_num_tokens(0) @@ -1629,6 +1635,9 @@ def _executor_loop_pp(self): self.scheduler.schedule_request(self.active_requests, self.inflight_req_ids) + self._terminate_recompute_paused_requests(scheduled_batch) + self._pause_recompute_paused_requests(scheduled_batch) + # For requests that are fitting disagg gen init, also prepare resources for KV cache manager if self.kv_cache_transceiver: self._prepare_disagg_gen_init( @@ -2428,6 +2437,9 @@ def _executor_loop(self): self._revert_gen_alloc(scheduled_batch) continue + self._terminate_recompute_paused_requests(scheduled_batch) + self._pause_recompute_paused_requests(scheduled_batch) + if not self._scheduler_manages_kv_suspend: self._terminate_requests(scheduled_batch.paused_requests) self._pause_requests(scheduled_batch.paused_requests) @@ -2720,6 +2732,8 @@ def _executor_loop_overlap(self): self._revert_gen_alloc(scheduled_batch) continue + self._terminate_recompute_paused_requests(scheduled_batch) + if not self._scheduler_manages_kv_suspend: self._terminate_requests(scheduled_batch.paused_requests) @@ -2849,6 +2863,7 @@ def _executor_loop_overlap(self): if not self._scheduler_manages_kv_suspend: self._pause_requests(scheduled_batch.paused_requests) + self._pause_recompute_paused_requests(scheduled_batch) if can_queue: guided_decoder_failed_requests = None @@ -3452,6 +3467,7 @@ def _schedule(self): scheduled_requests.reset_context_requests(scheduled_context_requests) scheduled_requests.generation_requests = scheduler_output.generation_requests scheduled_requests.paused_requests = scheduler_output.paused_requests + scheduled_requests.recompute_paused_requests = scheduler_output.recompute_paused_requests return scheduled_requests, scheduler_output.fitting_disagg_gen_init_requests, num_fitting @@ -4488,6 +4504,44 @@ def _pause_requests(self, requests_to_pause): for req in requests_to_pause: req.pause(self.max_input_len) + def _pause_recompute_request(self, req): + req.pause(_UNBOUNDED_PAUSE_MAX_INPUT_LEN) + req.py_batch_idx = None + req.py_seq_slot = None + req.py_prompt_len = req.prompt_len + req.py_orig_prompt_len = req.prompt_len + req.py_max_new_tokens = req.max_new_tokens + req.seqlen_this_rank_cp = req.prompt_len + req.total_input_len_cp = req.prompt_len + req.py_draft_pages_allocated = 0 + req.py_rewind_len = 0 + req.py_draft_tokens = [] + req.draft_tokens = [] + req.py_last_context_chunk = (None, None) + req.py_last_draft_tokens = None + req.py_num_accepted_draft_tokens = 0 + req.py_num_accepted_draft_tokens_indices = [] + req.py_rewind_draft_token_separate_adjustment = 0 + req.py_decoding_iter = 0 + req.py_ctx_pre_resize_cap = None + req._cached_tokens = 0 + req._cached_tokens_set = False + + def _terminate_recompute_paused_requests( + self, scheduled_batch: ScheduledRequests): + requests = scheduled_batch.recompute_paused_requests + if not requests: + return + self._terminate_requests(requests) + + def _pause_recompute_paused_requests(self, + scheduled_batch: ScheduledRequests): + requests = scheduled_batch.recompute_paused_requests + if not requests: + return + for req in requests: + self._pause_recompute_request(req) + def _add_inflight_ids(self, scheduled_requests: ScheduledRequests): """Add request IDs of current sampling requests to self.inflight_req_ids. diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 8047c81620c0..4a6b513a1518 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2102,6 +2102,8 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], logger.info( f"KV cache manager v2 host cache quota set to {host_quota / (1 << 30):.2f}GiB" ) + self.has_host_cache_tier = any( + isinstance(tier, HostCacheTierConfig) for tier in cache_tiers) self.vocab_size = vocab_size @@ -2133,6 +2135,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], cache_tiers=cache_tiers_gpu_only, ) cache_tiers = cache_tiers_gpu_only + self.has_host_cache_tier = False self.kv_cache_manager_py_config = config self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 3716f2397635..35fdd4633240 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -1,3 +1,18 @@ +# 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import dataclasses import inspect from abc import ABC, abstractmethod @@ -49,16 +64,46 @@ def _call_with_optional_summary( return fn(*args, cached_summary=cached_summary) -SchedulerOutput = namedtuple( - "SchedulerOutput", - [ - "context_requests", - "generation_requests", - "paused_requests", - "fitting_disagg_gen_init_requests", - "num_fitting_requests", - ], -) +class SchedulerOutput( + namedtuple( + "_SchedulerOutputBase", + [ + "context_requests", + "generation_requests", + "paused_requests", + "fitting_disagg_gen_init_requests", + "num_fitting_requests", + "recompute_paused_requests", + ], + ) +): + """Scheduler result. + + ``recompute_paused_requests`` is V2-only and defaults to an empty list so + existing V1 schedulers can keep constructing the original five-field + output. + """ + + __slots__ = () + + def __new__( + cls, + context_requests: RequestList, + generation_requests: RequestList, + paused_requests: RequestList, + fitting_disagg_gen_init_requests: RequestList, + num_fitting_requests: int, + recompute_paused_requests: RequestList | None = None, + ): + return super(SchedulerOutput, cls).__new__( + cls, + context_requests, + generation_requests, + paused_requests, + fitting_disagg_gen_init_requests, + num_fitting_requests, + [] if recompute_paused_requests is None else recompute_paused_requests, + ) class ScheduledRequests: @@ -77,13 +122,16 @@ class ScheduledRequests: generation_requests: RequestList """Requests that are in the generation phase.""" paused_requests: RequestList - """Requests that are paused.""" + """Requests whose KV cache was suspended without resetting request state.""" + recompute_paused_requests: RequestList + """Requests that must release resources and restart from context.""" def __init__(self): self.context_requests_chunking: RequestList = [] self.context_requests_last_chunk: RequestList = [] self.generation_requests: RequestList = [] self.paused_requests: RequestList = [] + self.recompute_paused_requests: RequestList = [] @property def is_generation_only(self) -> bool: @@ -173,6 +221,8 @@ class SerializableSchedulerOutput: int ] # request ids of fitting disaggregated generation initialization requests num_fitting_requests: int # number of fitting requests + recompute_paused_requests: list[int] = dataclasses.field(default_factory=list) + """Request ids of recompute-paused requests.""" @classmethod def from_scheduler_result( @@ -194,6 +244,9 @@ def from_scheduler_result( req.request_id for req in fitting_disagg_gen_init_requests ], num_fitting_requests=num_fitting_requests, + recompute_paused_requests=[ + req.request_id for req in scheduled_requests.recompute_paused_requests + ], ) def to_scheduler_result( @@ -213,6 +266,9 @@ def to_scheduler_result( scheduled_requests.paused_requests = [ id_to_request[req_id] for req_id in self.paused_requests ] + scheduled_requests.recompute_paused_requests = [ + id_to_request[req_id] for req_id in self.recompute_paused_requests + ] fitting_disagg_gen_init_requests = [ id_to_request[req_id] for req_id in self.fitting_disagg_gen_init_requests ] diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 191f4e51ab39..4d7800e5bf60 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -197,9 +197,14 @@ def schedule_request( self, active_requests: RequestList, inflight_request_ids: set[int] ) -> SchedulerOutput: # Main scheduling loop - (scheduled_ctx, scheduled_gen, evicted, disagg_candidates, has_chunking) = ( - self._schedule_loop(active_requests, inflight_request_ids) - ) + ( + scheduled_ctx, + scheduled_gen, + evicted, + recompute_paused, + disagg_candidates, + has_chunking, + ) = self._schedule_loop(active_requests, inflight_request_ids) # Sort by LoRA task ID self._sort_requests(scheduled_ctx, scheduled_gen, has_chunking) @@ -208,6 +213,7 @@ def schedule_request( context_requests=scheduled_ctx, generation_requests=scheduled_gen, paused_requests=evicted, + recompute_paused_requests=recompute_paused, fitting_disagg_gen_init_requests=disagg_candidates, num_fitting_requests=len(scheduled_ctx) + len(scheduled_gen), ) @@ -218,6 +224,7 @@ def _schedule_loop(self, active_requests, inflight_request_ids): scheduled_ctx: RequestList = [] scheduled_gen: RequestList = [] evicted: RequestList = [] + recompute_paused: RequestList = [] disagg_candidates: RequestList = [] scheduled_beam_width = 0 has_chunking = False @@ -338,6 +345,8 @@ def _schedule_loop(self, active_requests, inflight_request_ids): req_it, req_it_end, evicted, + recompute_paused, + inflight_request_ids, scheduled_beam_width, ) if action is ScheduleAction.STOP: @@ -387,17 +396,29 @@ def _schedule_loop(self, active_requests, inflight_request_ids): and not r.is_generation_to_complete_state and r.request_id not in inflight_request_ids ) - if num_gen_candidates > 0 and not evicted: + if ( + num_gen_candidates > 0 + and not evicted + and not recompute_paused + and not inflight_request_ids + ): raise RuntimeError( f"V2 scheduler deadlock: {num_gen_candidates} generation " f"request(s) active but none could be scheduled or " - f"evicted. KV cache pool is likely exhausted with no " + f"evicted or recompute-paused. KV cache pool is likely exhausted with no " f"host cache tier for suspend/resume offload. " f"Configure kv_cache_config.host_cache_size or increase " f"kv_cache_config.max_tokens." ) - return scheduled_ctx, scheduled_gen, evicted, disagg_candidates, has_chunking + return ( + scheduled_ctx, + scheduled_gen, + evicted, + recompute_paused, + disagg_candidates, + has_chunking, + ) # ---- Per-type scheduling methods ---- @@ -564,6 +585,8 @@ def _try_schedule_generation( req_it: int, req_it_end: int, evicted: RequestList, + recompute_paused: RequestList, + inflight_request_ids: set[int], scheduled_beam_width: int, ) -> tuple[ScheduleAction, int, int, int]: """Try to schedule a generation request. @@ -586,7 +609,18 @@ def _try_schedule_generation( if not success: req_it_end, success = self._try_evict_for_gen( - req, requests_list, req_it, req_it_end, evicted + req, requests_list, req_it, req_it_end, evicted, inflight_request_ids + ) + + if not success: + req_it_end, success = self._try_recompute_pause_for_gen( + req, + requests_list, + req_it, + req_it_end, + evicted, + recompute_paused, + inflight_request_ids, ) if success: @@ -634,17 +668,34 @@ def _suspend_request(self, req: LlmRequest) -> None: def _clear_request_runtime_state(self, req: LlmRequest) -> None: req.py_batch_idx = None - def _is_evictable(self, req: LlmRequest) -> bool: + def _is_evictable(self, req: LlmRequest, inflight_request_ids: set[int]) -> bool: """A started request whose KV cache is still active on GPU. Already-suspended requests are not useful eviction victims because suspending them again is a no-op that frees no pages. """ + if req.request_id in inflight_request_ids: + return False if not self._is_started_request(req): return False return self.kv_cache_manager.is_request_active(req.py_request_id) - def _try_evict_for_gen(self, req, requests_list, req_it, req_it_end, evicted): + def _is_recompute_pause_candidate( + self, req: LlmRequest, inflight_request_ids: set[int] + ) -> bool: + if req.request_id in inflight_request_ids: + return False + return self._is_started_request(req) + + def _recompute_pause_request(self, req: LlmRequest) -> None: + self._clear_request_runtime_state(req) + self.kv_cache_manager.free_resources(req) + if self.draft_kv_cache_manager is not None: + self.draft_kv_cache_manager.free_resources(req) + + def _try_evict_for_gen( + self, req, requests_list, req_it, req_it_end, evicted, inflight_request_ids + ): """Evict started requests from active_requests tail to make room. Search backwards from req_it_end @@ -662,7 +713,7 @@ def _try_evict_for_gen(self, req, requests_list, req_it, req_it_end, evicted): while req_it_end > req_it: victim_idx = None for i in range(req_it_end - 1, req_it, -1): - if self._is_evictable(requests_list[i]): + if self._is_evictable(requests_list[i], inflight_request_ids): victim_idx = i break @@ -683,6 +734,66 @@ def _try_evict_for_gen(self, req, requests_list, req_it, req_it_end, evicted): return req_it_end, False + def _try_recompute_pause_for_gen( + self, + req, + requests_list, + req_it, + req_it_end, + evicted, + recompute_paused, + inflight_request_ids, + ): + """Use destructive recompute pause when ordinary suspend is insufficient. + + Prefer requests already suspended during this allocation attempt. These + victims have no active GPU pages, so freeing them can only make room for + subsequent ordinary suspends to offload more active GPU pages. If + suspended victims are not enough, destructively pause active non-inflight + victims from the unprocessed tail. + """ + if self.kv_cache_manager.has_host_cache_tier: + while evicted: + victim = evicted.pop(0) + if not self._is_recompute_pause_candidate(victim, inflight_request_ids): + continue + logger.debug( + f"[V2Scheduler] Recompute-pausing suspended request {victim.py_request_id} " + f"before evicting more requests for request {req.py_request_id}" + ) + self._recompute_pause_request(victim) + recompute_paused.append(victim) + + req_it_end, success = self._try_evict_for_gen( + req, requests_list, req_it, req_it_end, evicted, inflight_request_ids + ) + if success: + return req_it_end, True + + while req_it_end > req_it: + victim_idx = None + for i in range(req_it_end - 1, req_it, -1): + if self._is_recompute_pause_candidate(requests_list[i], inflight_request_ids): + victim_idx = i + break + + if victim_idx is None: + break + + victim = requests_list[victim_idx] + logger.debug( + f"[V2Scheduler] Recompute-pausing active request {victim.py_request_id} " + f"to free pages for request {req.py_request_id}" + ) + self._recompute_pause_request(victim) + recompute_paused.append(victim) + req_it_end = victim_idx + + if self.kv_cache_manager.try_allocate_generation(req): + return req_it_end, True + + return req_it_end, False + # ---- Sorting ---- @staticmethod diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index 7028f5bdb7c7..1802dcb5fdc9 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -150,14 +150,21 @@ def make_kv_cache_manager( prepare_context_fn=None, resize_context_fn=None, try_allocate_generation_fn=None, + has_host_tier=False, ): mgr = Mock() mgr.tokens_per_block = tokens_per_block + mgr.has_host_cache_tier = has_host_tier mgr.kv_cache_map = _KVCacheMap() mgr.prepare_context.side_effect = prepare_context_fn or (lambda req: True) mgr.resize_context.side_effect = resize_context_fn or (lambda req, n, history_length=None: True) mgr.try_allocate_generation.side_effect = try_allocate_generation_fn or (lambda req: True) - mgr.suspend_request.return_value = None + + def suspend_request(req): + if has_host_tier: + mgr.kv_cache_map[req.py_request_id].is_active = False + + mgr.suspend_request.side_effect = suspend_request mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active return mgr @@ -471,6 +478,28 @@ def alloc_fn(req): # gen99 evicted as victim for gen1; gen1 self-evicts after assert set(ids(out.paused_requests)) == {1, 99} + def test_gen_alloc_fails_recompute_pauses_host_victim(self): + """host-tier victim is recompute-paused before gen self-evicts.""" + call_count = [0] + + def alloc_fn(req): + call_count[0] += 1 + # Only gen0 succeeds. A fourth call would mean the scheduler retried + # allocation immediately after recompute-pausing the suspended + # victim, which should not happen because it frees no GPU pages. + return call_count[0] in (1, 4) + + mgr = make_kv_cache_manager(try_allocate_generation_fn=alloc_fn, has_host_tier=True) + sched = make_scheduler(mgr, max_num_tokens=100) + victim = make_gen_request(99) + reqs = [make_gen_request(0), make_gen_request(1), victim] + out = sched.schedule_request(reqs, set()) + assert ids(out.generation_requests) == [0] + assert ids(out.paused_requests) == [1] + assert ids(out.recompute_paused_requests) == [99] + assert call_count[0] == 3 + mgr.free_resources.assert_called_once_with(victim) + def test_multiple_evictions_needed(self): """gen fails, 2 victims needed to free enough space.""" call_count = [0] @@ -1712,6 +1741,7 @@ def test_output_fields_correct(self): assert len(out.context_requests) == 1 assert len(out.generation_requests) == 1 assert len(out.fitting_disagg_gen_init_requests) == 1 + assert out.recompute_paused_requests == [] assert out.num_fitting_requests == 2 def test_num_fitting_requests(self): @@ -1842,6 +1872,22 @@ def selective_gen_alloc(req): assert len(out.generation_requests) == 0 assert set(ids(out.paused_requests)) == {0, 1, 2} + def test_multiple_gen_after_gen_fail_with_host_tier_recompute_pause(self): + """Host-tier evicted victims are recompute-paused before self-evict.""" + + def selective_gen_alloc(req): + return req.request_id != 0 + + mgr = make_kv_cache_manager( + try_allocate_generation_fn=selective_gen_alloc, has_host_tier=True + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + reqs = [make_gen_request(0), make_gen_request(1), make_gen_request(2)] + out = sched.schedule_request(reqs, set()) + assert len(out.generation_requests) == 0 + assert ids(out.paused_requests) == [0] + assert set(ids(out.recompute_paused_requests)) == {1, 2} + # =========================================================================== # can_schedule diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index bb7ae2448ac5..82987744c851 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -2292,6 +2292,7 @@ def test_full_pipeline_output_structure(self): assert hasattr(output, "paused_requests") assert hasattr(output, "fitting_disagg_gen_init_requests") assert hasattr(output, "num_fitting_requests") + assert len(output.recompute_paused_requests) == 0 assert len(output.context_requests) == 1 assert len(output.generation_requests) == 1 assert output.context_requests[0].request_id == 0 diff --git a/tests/unittest/_torch/executor/test_scheduler_serializable_output.py b/tests/unittest/_torch/executor/test_scheduler_serializable_output.py index fcce6d88ea9e..ccaabea7600e 100644 --- a/tests/unittest/_torch/executor/test_scheduler_serializable_output.py +++ b/tests/unittest/_torch/executor/test_scheduler_serializable_output.py @@ -27,6 +27,7 @@ def test_serializable_scheduler_output_round_trip(): scheduled_requests.context_requests_last_chunk = [request_pool[1], request_pool[2]] scheduled_requests.generation_requests = [request_pool[3]] scheduled_requests.paused_requests = [request_pool[4]] + scheduled_requests.recompute_paused_requests = [request_pool[7]] fitting_disagg_gen_init_requests = [request_pool[5], request_pool[6]] num_fitting_requests = 3 @@ -59,4 +60,7 @@ def test_serializable_scheduler_output_round_trip(): assert _request_ids(restored_schedule.paused_requests) == _request_ids( scheduled_requests.paused_requests ) + assert _request_ids(restored_schedule.recompute_paused_requests) == _request_ids( + scheduled_requests.recompute_paused_requests + ) assert _request_ids(restored_fitting) == _request_ids(fitting_disagg_gen_init_requests) diff --git a/tests/unittest/pyexecutor/test_iter_stats_populate.py b/tests/unittest/pyexecutor/test_iter_stats_populate.py index f5496361879c..355f5e0640fa 100644 --- a/tests/unittest/pyexecutor/test_iter_stats_populate.py +++ b/tests/unittest/pyexecutor/test_iter_stats_populate.py @@ -80,10 +80,17 @@ def get_num_tokens(self, beam: int = 0) -> int: class _StubScheduledBatch: - def __init__(self, context_reqs=None, gen_reqs=None, paused_reqs=None): + def __init__( + self, + context_reqs=None, + gen_reqs=None, + paused_reqs=None, + recompute_paused_reqs=None, + ): self.context_requests = list(context_reqs or []) self.generation_requests = list(gen_reqs or []) self.paused_requests = list(paused_reqs or []) + self.recompute_paused_requests = list(recompute_paused_reqs or []) @property def num_context_requests(self): @@ -364,10 +371,15 @@ def test_paused_decode_requests(): _StubRequest(num_tokens=300), _StubRequest(num_tokens=800), ] - stats = _invoke_update_iter_stats(_StubScheduledBatch(paused_reqs=paused), [], num_ctx_tokens=0) + recompute_paused = [_StubRequest(num_tokens=700)] + stats = _invoke_update_iter_stats( + _StubScheduledBatch(paused_reqs=paused, recompute_paused_reqs=recompute_paused), + [], + num_ctx_tokens=0, + ) ifb = stats.inflight_batching_stats - assert ifb.num_paused_requests == 2 - assert ifb.num_paused_kv_tokens == 1100 + assert ifb.num_paused_requests == 3 + assert ifb.num_paused_kv_tokens == 1800 def test_attention_dp_dummy_filtering_on_kv_token_fields(): diff --git a/tests/unittest/pyexecutor/test_recompute_pause.py b/tests/unittest/pyexecutor/test_recompute_pause.py new file mode 100644 index 000000000000..9aa14f9842f7 --- /dev/null +++ b/tests/unittest/pyexecutor/test_recompute_pause.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import types + +from tensorrt_llm._torch.pyexecutor.py_executor import _UNBOUNDED_PAUSE_MAX_INPUT_LEN, PyExecutor + + +class _StubRequest: + def __init__(self) -> None: + self.generated_tokens = 6 + self.max_new_tokens = 20 + self.prompt_len = 4 + self.pause_max_input_len = None + + def pause(self, max_input_len: int) -> None: + self.pause_max_input_len = max_input_len + new_prompt_len = min(max_input_len, self.prompt_len + self.generated_tokens) + self.max_new_tokens -= new_prompt_len - self.prompt_len + self.prompt_len = new_prompt_len + + +def test_recompute_pause_does_not_apply_executor_max_input_len() -> None: + executor = types.SimpleNamespace(max_input_len=5) + request = _StubRequest() + + PyExecutor._pause_recompute_request(executor, request) + + assert request.pause_max_input_len == _UNBOUNDED_PAUSE_MAX_INPUT_LEN + assert request.prompt_len == 10 + assert request.py_prompt_len == 10 + assert request.py_orig_prompt_len == 10 + assert request.py_max_new_tokens == 14 + assert request.py_draft_tokens == [] + assert request.draft_tokens == [] + assert request._cached_tokens == 0 + assert request._cached_tokens_set is False