diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 234817ccb3fb..9d1cb753af6c 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2376,13 +2376,14 @@ def prepare_context(self, req: LlmRequest) -> bool: ) return self._resume_and_restore(req.py_request_id, kv_cache) - def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: + def resize_context(self, + req: LlmRequest, + num_tokens: int, + history_length: int | None = None) -> bool: """Resize KV cache to cover context_current_position + num_tokens. - num_tokens is the number of tokens to be processed (i.e., - context_remaining_length or a chunk thereof). The target capacity is - computed as context_current_position + num_tokens so that block reuse - overlaps with existing capacity are handled correctly. + history_length, when set, lets SWA life cycles compute their stale + range at allocation time so pre-window blocks are never allocated. Returns True on success, False if resize failed (first chunk is suspended on failure). """ @@ -2393,7 +2394,7 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: target = req.context_current_position + num_tokens + self.num_extra_kv_tokens capacity = max(kv_cache.capacity, target) - if not kv_cache.resize(capacity): + if not kv_cache.resize(capacity, history_length): if req.is_first_context_chunk: kv_cache.suspend() return False diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index e8490bacab3e..4ab36f10dd6c 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -368,8 +368,12 @@ def _try_schedule_disagg_gen_init( logger.debug("prepare_context failed for disagg gen init request %s", req.py_request_id) return ScheduleAction.SKIP, 0 + # Prompt is already historic (processed by context server) — pass + # history_length so SWA layers skip pre-window block allocation. if not self.kv_cache_manager.resize_context( - req, req.context_remaining_length + get_draft_token_length(req) + req, + req.context_remaining_length + get_draft_token_length(req), + history_length=req.context_remaining_length, ): return ScheduleAction.SKIP, 0 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 9aaa35492ea2..dbfd4182e918 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -155,7 +155,7 @@ def make_kv_cache_manager( mgr.tokens_per_block = tokens_per_block 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: 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 mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active @@ -1156,7 +1156,7 @@ def test_disagg_resize_context_fails_skips(self): """resize_context failure skips the request, loop continues.""" call_count = [0] - def resize_fn(req, n): + def resize_fn(req, n, history_length=None): call_count[0] += 1 return call_count[0] > 1 # first fails, second succeeds @@ -1166,6 +1166,24 @@ def resize_fn(req, n): out = sched.schedule_request(reqs, set()) assert ids(out.fitting_disagg_gen_init_requests) == [1] + def test_disagg_resize_context_passes_history_length(self): + """Disagg init forwards history_length=prompt_len to resize_context.""" + captured = {} + + def resize_fn(req, n, history_length=None): + captured[req.py_request_id] = history_length + return True + + mgr = make_kv_cache_manager(resize_context_fn=resize_fn) + sched = make_scheduler(mgr, max_num_tokens=100) + reqs = [ + make_disagg_request(0, context_remaining_length=200), + make_disagg_request(1, context_remaining_length=512), + ] + sched.schedule_request(reqs, set()) + assert captured[0] == 200 + assert captured[1] == 512 + def test_disagg_exceeds_max_batch_size(self): """Disagg count can exceed max_batch_size since it bypasses budget.""" mgr = make_kv_cache_manager()