diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index fbadd440d611..4c92c3b2c236 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -392,13 +392,16 @@ def __init__( # kv cache events self.kv_cache_manager = self.resource_manager.resource_managers.get( ResourceManagerType.KV_CACHE_MANAGER) - # V2 scheduler calls suspend_request() during scheduling, which - # offloads GPU pages while preserving the radix tree. The executor - # does not need to call _terminate_requests (GPU resources are already - # freed by suspend) or _pause_requests (V2's prepare_context handles - # resume internally, so resetting to CONTEXT_INIT is unnecessary). - self._scheduler_manages_kv_suspend = isinstance(self.kv_cache_manager, - KVCacheManagerV2) + # V2 manager owns KV alloc + suspend during scheduling: it + # eagerly grows ctx/gen capacity in the schedule loop and calls + # suspend_request() when needed (offloads GPU pages while + # preserving the radix tree). The executor therefore does not + # need to call _terminate_requests (GPU resources are already + # freed by suspend) or _pause_requests (V2's prepare_context + # handles resume internally, so resetting to CONTEXT_INIT is + # unnecessary). Several revert/skip paths gate on this flag. + self._is_kv_manager_v2 = isinstance(self.kv_cache_manager, + KVCacheManagerV2) self.enable_kv_cache_events = self.kv_cache_manager is not None and self.kv_cache_manager.event_buffer_max_size > 0 self.enable_kv_cache_reuse = self.kv_cache_manager is not None and self.kv_cache_manager.enable_block_reuse self.enable_partial_reuse_for_disagg = ( @@ -1959,10 +1962,26 @@ def _revert_gen_alloc(self, scheduled_batch): can_queue check. V1 allocates in prepare_resources() after the can_queue check, so no revert is needed. """ - if self._scheduler_manages_kv_suspend: + if self._is_kv_manager_v2: for req in scheduled_batch.generation_requests: self.kv_cache_manager.revert_allocate_generation(req) + def _revert_ctx_alloc(self, dropped_context_requests): + """Revert KV cache capacity growth for ctx requests deferred by + delay batching. + + With KV cache manager V2 + scheduler V2, ctx KV cache is grown + during scheduling (``resize_context``). When delay batching + (``_balance_adp_requests`` for ADP, or ``_waiting_requests`` + for non-ADP batch waiting) defers ctx requests, the + freshly-allocated pages would otherwise sit idle until the + request is re-scheduled, blocking pool space — particularly + painful for long-context workloads where each deferred ctx can + hold GBs of KV. + """ + for req in dropped_context_requests: + self.kv_cache_manager.revert_allocate_context(req) + def _prepare_and_schedule_batch(self): new_requests = self._fetch_and_activate_new_requests() if self.should_stop_processing: @@ -2204,7 +2223,7 @@ def _executor_loop(self): if should_retry: continue - if not self._scheduler_manages_kv_suspend: + if not self._is_kv_manager_v2: self._terminate_requests(scheduled_batch.paused_requests) self._pause_requests(scheduled_batch.paused_requests) @@ -2448,7 +2467,7 @@ def _executor_loop_overlap(self): if should_retry: continue - if not self._scheduler_manages_kv_suspend: + if not self._is_kv_manager_v2: self._terminate_requests(scheduled_batch.paused_requests) can_queue, can_queue_this_rank = self._can_queue( @@ -2573,7 +2592,7 @@ def _executor_loop_overlap(self): # Cleanup previous draft resources used in the draft model self.drafter.cleanup_previous_draft_resources() - if not self._scheduler_manages_kv_suspend: + if not self._is_kv_manager_v2: self._pause_requests(scheduled_batch.paused_requests) if can_queue: @@ -3120,7 +3139,8 @@ def _schedule(self): scheduler_output = self.scheduler.schedule_request( self.active_requests, self.inflight_req_ids) - scheduled_context_requests = scheduler_output.context_requests + original_ctx_requests = scheduler_output.context_requests + scheduled_context_requests = original_ctx_requests if self.enable_attention_dp and self.attention_dp_enable_balance: scheduled_context_requests = self._balance_adp_requests( scheduler_output.context_requests, @@ -3143,6 +3163,20 @@ def _schedule(self): scheduled_context_requests = self.kv_cache_manager.filter_ctx_requests_by_capacity( scheduled_context_requests) num_fitting = len(scheduled_context_requests) + + # V2 scheduler grew KV cache for ctx during scheduling; release + # those pages for any ctx that delay batching has dropped, so + # the wait window does not hold pool capacity hostage. V1 + # allocates after delay batching, so skip the dropped-set + # computation entirely on V1. + if (self._is_kv_manager_v2 and len(scheduled_context_requests) + < len(original_ctx_requests)): + kept = {r.py_request_id for r in scheduled_context_requests} + dropped = [ + r for r in original_ctx_requests if r.py_request_id not in kept + ] + self._revert_ctx_alloc(dropped) + scheduled_requests = ScheduledRequests() scheduled_requests.reset_context_requests(scheduled_context_requests) scheduled_requests.generation_requests = scheduler_output.generation_requests diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index fae72ccae7de..5843127cf0b0 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2246,6 +2246,37 @@ def revert_allocate_generation(self, req: LlmRequest) -> None: f"{req.py_request_id} from {kv_cache.capacity} to " f"{reverted_cap}") + def revert_allocate_context(self, req: LlmRequest) -> None: + """Undo the capacity growth from this iter's ``resize_context``. + + When delay batching (``_balance_adp_requests`` / + ``_waiting_requests``) defers a context request after V2 + scheduling, the forward pass is skipped for that request but the + scheduler already grew its KV cache capacity to cover the chunk. + This shrinks capacity back to the pre-resize value so the + freshly-allocated pages can be reused during the wait window — + important for long contexts where one deferred request can hold + GBs of KV. + """ + pre_cap = getattr(req, "py_ctx_pre_resize_cap", None) + if pre_cap is None: + return + # Mark as consumed even if the resize below is skipped, so a + # later iter does not see a stale snapshot. + req.py_ctx_pre_resize_cap = None + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None or not kv_cache.is_active: + return + if pre_cap >= kv_cache.capacity: + return + if not kv_cache.resize(pre_cap): + raise RuntimeError( + f"Failed to revert KV cache capacity for context " + f"request {req.py_request_id} from " + f"{kv_cache.capacity} to {pre_cap}") + if pre_cap > 0: + kv_cache.suspend() + def _restore_page_index_bufs(self, request_id: int, kv_cache) -> None: """Re-connect host page-index buffers after resume(). @@ -2326,6 +2357,10 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: overlaps with existing capacity are handled correctly. Returns True on success, False if resize failed (first chunk is suspended on failure). + + Snapshots the pre-resize capacity on ``req.py_ctx_pre_resize_cap`` + when growth happens so ``revert_allocate_context`` can undo it if + delay batching defers the request. """ kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is None: @@ -2333,12 +2368,16 @@ 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) + pre_cap = kv_cache.capacity if not kv_cache.resize(capacity): if req.is_first_context_chunk: kv_cache.suspend() return False + # None means "no growth this iter, nothing to revert"; this also + # invalidates a stale snapshot from a prior iter on the same req. + req.py_ctx_pre_resize_cap = pre_cap if capacity > pre_cap else None return True def extend_capacity_for_tokens(self, request: LlmRequest) -> None: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 9ab58dc71869..d42ffcdbdc27 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -2453,21 +2453,26 @@ def test_cute_dsl_nvfp4_4gpus( task.evaluate(llm) @skip_pre_blackwell + @parametrize_with_ids("v2_kv_cache", [False, True]) @parametrize_with_ids("torch_compile", [False, True]) @parametrize_with_ids("fp8kv,cuda_graph,overlap_scheduler", [(False, False, False), (True, True, True)]) @parametrize_with_ids("mtp_nextn", [0, 2]) @parametrize_with_ids( "batch_wait_timeout_iters,batch_wait_max_tokens_ratio", [(0, 0), - (10, 0.75), + (10, 1.0), (10, 0), (0, 0.75)]) def test_nvfp4_batch_waiting(self, torch_compile, fp8kv, cuda_graph, overlap_scheduler, mtp_nextn, batch_wait_timeout_iters, - batch_wait_max_tokens_ratio): + batch_wait_max_tokens_ratio, v2_kv_cache): moe_backend = "CUTLASS" - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) + # V2 KV cache manager auto-selects the V2 scheduler, which is the + # path that grows ctx KV during scheduling and needs the + # delay-batching ctx revert to release pages during the wait window. + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75, + use_kv_cache_manager_v2=v2_kv_cache) torch_compile_config = _get_default_torch_compile_config(torch_compile) pytorch_config = dict( disable_overlap_scheduler=not overlap_scheduler, diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 9285526830dd..80c2925f563e 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -426,8 +426,10 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backe accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=False-moe_backend=WIDEEP] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=TRTLLM] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=WIDEEP] -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=0.75-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=0.75-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=1.0-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-v2_kv_cache=False] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=1.0-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-v2_kv_cache=True] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=1.0-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True-v2_kv_cache=False] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=1.0-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True-v2_kv_cache=True] accuracy/test_llm_api_pytorch.py::TestEXAONE4::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype_vswa_chunked_prefill_reuse diff --git a/tests/integration/test_lists/qa/llm_function_rtx6k.txt b/tests/integration/test_lists/qa/llm_function_rtx6k.txt index dbf77fe1fbef..a5542794700e 100644 --- a/tests/integration/test_lists/qa/llm_function_rtx6k.txt +++ b/tests/integration/test_lists/qa/llm_function_rtx6k.txt @@ -59,7 +59,8 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUT accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=0.75-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=1.0-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-v2_kv_cache=False] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=1.0-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-v2_kv_cache=True] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_online_eplb[mtp_nextn=0-moe_backend=WIDEEP] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_online_eplb[mtp_nextn=2-moe_backend=WIDEEP] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=False-moe_backend=WIDEEP] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 4fdfa3eafad3..7772040885d3 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -50,6 +50,7 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=nvfp4-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_guided_decoding[xgrammar-mtp_nextn=2] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_dummy_load_format + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=1.0-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-v2_kv_cache=True] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-cutlass-auto] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-trtllm-fp8] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-triton-auto]