[None][fix] Avoid pre-window SWA over-allocation on disagg gen server

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

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
 

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