From d467a2a7f2f088dd360c7bef2f3cf5c59a1ffde8 Mon Sep 17 00:00:00 2001 From: llx <54896441+llx-08@users.noreply.github.com> Date: Tue, 16 Jun 2026 05:36:09 +0800 Subject: [PATCH] [Bugfix] Defer block freeing until in-flight steps finish under async scheduling + PD KV consumer (#45357) Signed-off-by: llx-08 <2596671364@qq.com> Signed-off-by: Nick Hill Co-authored-by: Nick Hill Co-authored-by: Jiangyun Zhu --- tests/v1/core/test_async_scheduler.py | 1 + tests/v1/core/test_deferred_block_free.py | 414 ++++++++++++++++++ tests/v1/core/test_scheduler.py | 1 + .../config_sweep_accuracy_test.sh | 5 +- vllm/v1/core/kv_cache_coordinator.py | 19 + vllm/v1/core/kv_cache_manager.py | 13 + vllm/v1/core/sched/scheduler.py | 73 ++- vllm/v1/core/single_type_kv_cache_manager.py | 33 +- vllm/v1/request.py | 4 + 9 files changed, 543 insertions(+), 20 deletions(-) create mode 100644 tests/v1/core/test_deferred_block_free.py diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index a77a50173f3..5e9c9280dbe 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -284,6 +284,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.vllm_config.model_config.enable_return_routed_experts = False scheduler.enable_return_routed_experts = False scheduler.recompute_kv_load_failures = False + scheduler.defer_block_free = False scheduler.make_stats = Mock(return_value=None) scheduler.max_model_len = 128 diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py new file mode 100644 index 00000000000..8cab620f0e3 --- /dev/null +++ b/tests/v1/core/test_deferred_block_free.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for deferred block freeing under async scheduling. + +With async scheduling, a finished/preempted request's blocks may still be +written by a speculatively over-scheduled in-flight GPU step (mamba/GDN +layers rewrite the whole state block every step). If such a block is +reallocated to a request arriving via PD disaggregation, the NIC/RDMA write +of the received state races with the in-flight stale write. The scheduler +closes the race by deferring the return of blocks to the block pool until +the newest scheduled step's output has been processed. +""" + +import os +import time +from unittest.mock import PropertyMock, patch + +import pytest + +from vllm.config import VllmConfig +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.request import RequestStatus + +from .utils import create_requests, create_scheduler, mock_kv + +pytestmark = pytest.mark.cpu_test + +# Allow overriding the model with a local path for offline environments. +MODEL = os.environ.get("VLLM_TEST_DEFER_FREE_MODEL", "facebook/opt-125m") +STOP_TOKEN_ID = 42 +NUM_PROMPT_TOKENS = 33 # 3 blocks with block_size=16 + + +def _make_model_runner_output( + scheduler_output: SchedulerOutput, + token_id: int = 0, +) -> ModelRunnerOutput: + req_ids = list(scheduler_output.num_scheduled_tokens.keys()) + return ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=[[token_id] for _ in req_ids], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + +def _create_deferring_scheduler(): + """Async scheduler with deferred block freeing forced on. + + The production gate additionally requires a PD KV-consumer connector; + the mechanism itself is independent of it. + """ + scheduler = create_scheduler(model=MODEL, async_scheduling=True) + scheduler.defer_block_free = True + return scheduler + + +def _setup_request_with_inflight_step(scheduler, max_tokens: int = 5): + """Schedule a request's prefill (step 1) and one speculatively + over-scheduled decode (step 2), mimicking async scheduling depth 1. + + Returns (request, out0, out1). + """ + request = create_requests( + num_requests=1, + num_tokens=NUM_PROMPT_TOKENS, + max_tokens=max_tokens, + stop_token_ids=[STOP_TOKEN_ID], + )[0] + scheduler.add_request(request) + out0 = scheduler.schedule() + assert out0.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS + out1 = scheduler.schedule() + assert out1.num_scheduled_tokens[request.request_id] == 1 + return request, out0, out1 + + +def test_gate_enabled_for_async_consumer(): + # Overlapping batches + consumer-side connector enables the gate. Async + # scheduling (which would give >1 concurrent batches) is force-disabled on + # CPU, where this test runs, and PP can't be built without GPUs, so force + # max_concurrent_batches to exercise the enabled path on any platform. + with patch.object( + VllmConfig, + "max_concurrent_batches", + new_callable=PropertyMock, + return_value=2, + ): + scheduler = create_scheduler( + model=MODEL, + async_scheduling=True, + use_kv_connector=mock_kv(matched_tokens=0, is_async=False), + ) + assert scheduler.defer_block_free + + +def test_gate_disabled_without_connector(): + # Async scheduling alone (no PD connector): the gate must stay off + # and freeing must remain immediate. + scheduler = create_scheduler(model=MODEL, async_scheduling=True) + assert not scheduler.defer_block_free + + pool = scheduler.kv_cache_manager.block_pool + num_free_initially = pool.get_num_free_blocks() + + request, out0, out1 = _setup_request_with_inflight_step(scheduler) + assert pool.get_num_free_blocks() < num_free_initially + + # Request stops early while step 2 is in flight: blocks are freed + # immediately because deferral is disabled. + scheduler.update_from_output( + out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID) + ) + assert request.is_finished() + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_initially + + +def test_finish_defers_free_until_inflight_step_done(): + scheduler = _create_deferring_scheduler() + pool = scheduler.kv_cache_manager.block_pool + num_free_initially = pool.get_num_free_blocks() + + request, out0, out1 = _setup_request_with_inflight_step(scheduler) + num_free_running = pool.get_num_free_blocks() + assert num_free_running < num_free_initially + + # The request stops early (stop token) while the over-scheduled step 2 + # is still in flight: its blocks must NOT return to the pool yet. + scheduler.update_from_output( + out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID) + ) + assert request.is_finished() + assert len(scheduler.deferred_frees) == 1 + assert pool.get_num_free_blocks() == num_free_running + + # Step 2's output is processed: every GPU write of step 2 has + # completed, so the blocks can now be returned to the pool. + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_initially + + +def test_finish_frees_immediately_when_no_inflight_step(): + scheduler = _create_deferring_scheduler() + pool = scheduler.kv_cache_manager.block_pool + num_free_initially = pool.get_num_free_blocks() + + request = create_requests( + num_requests=1, + num_tokens=NUM_PROMPT_TOKENS, + max_tokens=5, + stop_token_ids=[STOP_TOKEN_ID], + )[0] + scheduler.add_request(request) + out0 = scheduler.schedule() + + # Synchronous-like flow: out0 is the newest scheduled step and its + # output is being processed, so no other step can still write the + # blocks and the free happens immediately. + scheduler.update_from_output( + out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID) + ) + assert request.is_finished() + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_initially + + +def test_abort_defers_free(): + scheduler = _create_deferring_scheduler() + pool = scheduler.kv_cache_manager.block_pool + num_free_initially = pool.get_num_free_blocks() + + request, out0, out1 = _setup_request_with_inflight_step(scheduler) + num_free_running = pool.get_num_free_blocks() + + # External abort arrives while steps 1 and 2 are both in flight. + scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) + assert len(scheduler.deferred_frees) == 1 + assert pool.get_num_free_blocks() == num_free_running + + # Step 1's output: step 2 is still in flight, keep holding the blocks. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert len(scheduler.deferred_frees) == 1 + assert pool.get_num_free_blocks() == num_free_running + + # Step 2's output: now the blocks can be freed. + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_initially + + +def test_preempt_defers_free_and_clears_bookkeeping(): + scheduler = _create_deferring_scheduler() + pool = scheduler.kv_cache_manager.block_pool + num_free_initially = pool.get_num_free_blocks() + + request, out0, out1 = _setup_request_with_inflight_step(scheduler) + num_free_running = pool.get_num_free_blocks() + + # Preempt the request while steps are in flight (mirrors the + # preemption path inside schedule()). + scheduler.running.remove(request) + scheduler._preempt_request(request, time.monotonic()) + assert request.status == RequestStatus.PREEMPTED + + # Blocks are withheld from the pool, but the manager bookkeeping is + # cleared immediately so the request can be rescheduled safely. + assert len(scheduler.deferred_frees) == 1 + assert pool.get_num_free_blocks() == num_free_running + for manager in scheduler.kv_cache_manager.coordinator.single_type_managers: + assert request.request_id not in manager.req_to_blocks + + # Outputs of both in-flight steps are processed: blocks return to the + # pool only after the newest one. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert len(scheduler.deferred_frees) == 1 + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_initially + + +def test_multiple_deferred_frees_drain_in_order(): + scheduler = _create_deferring_scheduler() + pool = scheduler.kv_cache_manager.block_pool + num_free_initially = pool.get_num_free_blocks() + + requests = create_requests( + num_requests=2, + num_tokens=NUM_PROMPT_TOKENS, + max_tokens=5, + stop_token_ids=[STOP_TOKEN_ID], + ) + for request in requests: + scheduler.add_request(request) + out0 = scheduler.schedule() + out1 = scheduler.schedule() + + # Both requests stop early at step 1's output while step 2 is in + # flight: two deferred entries with the same fence. + scheduler.update_from_output( + out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID) + ) + assert len(scheduler.deferred_frees) == 2 + assert pool.get_num_free_blocks() < num_free_initially + + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_initially + + +def test_fence_held_across_multiple_inflight_steps(): + """Pipeline-parallel / deep async: with several steps scheduled ahead, + a freed request's blocks must stay held until the *newest* in-flight + step's output is processed, not the first. + + Depth-1 tests only check a single intervening update; with PP the + scheduler can dispatch up to pp_size steps ahead, so the fence must + survive multiple intervening update_from_output calls. + """ + scheduler = _create_deferring_scheduler() + pool = scheduler.kv_cache_manager.block_pool + num_free_initially = pool.get_num_free_blocks() + + request = create_requests( + num_requests=1, + num_tokens=NUM_PROMPT_TOKENS, + max_tokens=10, + )[0] + scheduler.add_request(request) + + # Schedule three steps ahead without processing any output: a prefill + # plus two speculatively over-scheduled decodes, all in flight at once. + outs = [scheduler.schedule() for _ in range(3)] + assert outs[0].num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS + assert outs[1].num_scheduled_tokens[request.request_id] == 1 + assert outs[2].num_scheduled_tokens[request.request_id] == 1 + assert scheduler.sched_step_seq == 3 + num_free_running = pool.get_num_free_blocks() + assert num_free_running < num_free_initially + + # Abort while all three steps are in flight: the fence is the newest + # scheduled step (3), since any of them may still write the blocks. + scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) + assert len(scheduler.deferred_frees) == 1 + assert scheduler.deferred_frees[0][0] == 3 + assert pool.get_num_free_blocks() == num_free_running + + # Draining the two earlier in-flight steps must NOT release the blocks: + # their outputs don't fence the still-pending newest write. + for out in (outs[0], outs[1]): + scheduler.update_from_output(out, _make_model_runner_output(out)) + assert len(scheduler.deferred_frees) == 1 + assert pool.get_num_free_blocks() == num_free_running + + # Only once the newest scheduled step's output is processed do the + # blocks return to the pool. + scheduler.update_from_output(outs[2], _make_model_runner_output(outs[2])) + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_initially + + +def test_max_tokens_finish_frees_immediately_with_other_inflight(): + """A request finishing by reaching max_tokens is never over-scheduled past + its final-token step, so no in-flight step writes its blocks: it is freed + immediately even while another request's step is still in flight. + """ + scheduler = _create_deferring_scheduler() + pool = scheduler.kv_cache_manager.block_pool + + # Short request finishes at max_tokens=1; long request keeps running. + short = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=1, req_ids=["short"] + )[0] + long = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=100, req_ids=["long"] + )[0] + scheduler.add_request(short) + scheduler.add_request(long) + + out0 = scheduler.schedule() # prefill both + out1 = scheduler.schedule() # short is skipped (at max_tokens); long decodes + assert "short" not in out1.num_scheduled_tokens + assert "long" in out1.num_scheduled_tokens + + free_before = pool.get_num_free_blocks() + # Process step 0: `short` reaches max_tokens and finishes while step 1 + # (which scheduled `long`, not `short`) is still in flight. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + + assert short.is_finished() + # A step IS globally in flight (the old global fence would have deferred), + # but the per-request gate frees `short` immediately since nothing writes + # its blocks anymore. + assert scheduler.sched_step_seq > scheduler.processed_step_seq + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() > free_before # short's blocks returned + + +def test_abort_mid_prefill_defers_free(): + """Intermediate prefill chunks don't allocate output placeholders, so the + deferral must key off is_prefill_chunk: aborting a request whose prefill + chunk is still in flight must withhold its blocks. + """ + scheduler = create_scheduler( + model=MODEL, async_scheduling=True, long_prefill_token_threshold=16 + ) + scheduler.defer_block_free = True + pool = scheduler.kv_cache_manager.block_pool + num_free_initially = pool.get_num_free_blocks() + + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=5 + )[0] + scheduler.add_request(request) + + out0 = scheduler.schedule() + # Partial prefill: a chunk is in flight, with no output placeholders yet. + assert out0.num_scheduled_tokens[request.request_id] == 16 + assert request.num_output_placeholders == 0 + assert request.is_prefill_chunk + num_free_running = pool.get_num_free_blocks() + assert num_free_running < num_free_initially + + # Abort while the prefill chunk is in flight: blocks must be withheld + # (keyed off is_prefill_chunk, since there are no placeholders). + scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) + assert len(scheduler.deferred_frees) == 1 + assert pool.get_num_free_blocks() == num_free_running + + # Once the in-flight prefill step's output is processed, blocks return. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_initially + + +def test_non_async_abort_defers_via_last_sched_seq(): + """Without async (e.g. PP filling the pipeline) there are no placeholders + and a full prefill isn't a partial chunk, yet an abort with a step in flight + must defer. Only the last-scheduled-step fence catches this. + + PP=2 can't be built on a single-GPU host, so force the flag and exercise the + mechanism; the gate itself is covered by test_gate_enabled_for_async_consumer. + """ + scheduler = create_scheduler(model=MODEL, async_scheduling=False) + scheduler.defer_block_free = True + pool = scheduler.kv_cache_manager.block_pool + num_free_initially = pool.get_num_free_blocks() + + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=5 + )[0] + scheduler.add_request(request) + + out0 = scheduler.schedule() + # Neither async-only signal marks this request as in flight. + assert request.num_output_placeholders == 0 + assert not request.is_prefill_chunk + # Only the last-scheduled-step fence does. + assert request.last_sched_seq > scheduler.processed_step_seq + num_free_running = pool.get_num_free_blocks() + assert num_free_running < num_free_initially + + # Abort while out0 is in flight: blocks must be withheld. + scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) + assert len(scheduler.deferred_frees) == 1 + assert pool.get_num_free_blocks() == num_free_running + + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_initially diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 6b446fbc952..9ffd6f4cc0e 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -2571,6 +2571,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.vllm_config.model_config.enable_return_routed_experts = False scheduler.enable_return_routed_experts = False scheduler.recompute_kv_load_failures = False + scheduler.defer_block_free = False scheduler.make_stats = Mock(return_value=None) scheduler.max_model_len = 128 diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh index 432e7de3e99..bf9b15e7c78 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh @@ -24,11 +24,10 @@ dp_ep_configs=( # We assume HMA enabled by default. hybrid_ssm_configs=( "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" - # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models. - "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling" + "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" # GDN (Qwen3.5) "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B" - "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B VLLM_SERVE_EXTRA_ARGS=--no-async-scheduling" + "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B" ) sw_attn_configs=( # NOTE: gemma3 does not work with FlashInfer diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index bd528c66a00..376f65f6697 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -291,6 +291,25 @@ class KVCacheCoordinator(ABC): for manager in self.single_type_managers: manager.free(request_id) + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + """ + Pop the request's bookkeeping from all single-type managers and + return its blocks without returning them to the block pool. The + caller must eventually pass the returned blocks to + `block_pool.free_blocks`, freeing them in reverse order (so that + tail blocks are evicted first). + + Args: + request_id: The request ID. + + Returns: + The request's blocks in allocation order. + """ + blocks: list[KVCacheBlock] = [] + for manager in self.single_type_managers: + blocks.extend(manager.pop_blocks_for_free(request_id)) + return blocks + def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: """ Get the number of common prefix blocks for all requests with allocated diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 9af54e0a249..b0f6655bf95 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -480,6 +480,19 @@ class KVCacheManager: """ self.coordinator.remove_skipped_blocks(request_id, total_computed_tokens) + def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]: + """Pop the request's bookkeeping and return its blocks without + returning them to the block pool. The caller must eventually free + them in reverse order (so that tail blocks are evicted first). + + Args: + request: The request to pop the blocks for. + + Returns: + The request's blocks in allocation order. + """ + return self.coordinator.pop_blocks_for_free(request.request_id) + def evict_blocks(self, block_ids: set[int]) -> None: """evict blocks from the prefix cache by their block IDs. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 3b63ba32100..4d94d149050 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -37,6 +37,7 @@ from vllm.v1.core.encoder_cache_manager import ( from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector +from vllm.v1.core.kv_cache_utils import KVCacheBlock from vllm.v1.core.sched.interface import PauseState, SchedulerInterface from vllm.v1.core.sched.output import ( CachedRequestData, @@ -125,7 +126,9 @@ class Scheduler(SchedulerInterface): self.connector = None self.connector_prefix_cache_stats: PrefixCacheStats | None = None self.recompute_kv_load_failures = True - if self.vllm_config.kv_transfer_config is not None: + self.defer_block_free = False + kv_transfer_config = self.vllm_config.kv_transfer_config + if kv_transfer_config is not None: assert not self.is_encoder_decoder, ( "Encoder-decoder models are not currently supported with KV connectors" ) @@ -136,11 +139,17 @@ class Scheduler(SchedulerInterface): ) if self.log_stats: self.connector_prefix_cache_stats = PrefixCacheStats() - kv_load_failure_policy = ( - self.vllm_config.kv_transfer_config.kv_load_failure_policy - ) + kv_load_failure_policy = kv_transfer_config.kv_load_failure_policy self.recompute_kv_load_failures = kv_load_failure_policy == "recompute" + # With overlapping batches (async scheduling or PP), a step may + # still be writing a freed request's KV blocks. A consumer KV + # Connector can reallocate and fill those blocks via a load that + # isn't ordered against that write, so defer freeing them. + multiple_inflight_batches = self.vllm_config.max_concurrent_batches > 1 + if multiple_inflight_batches and kv_transfer_config.is_kv_consumer: + self.defer_block_free = True + self.kv_event_publisher = EventPublisherFactory.create( self.kv_events_config, self.parallel_config.data_parallel_index, @@ -275,6 +284,15 @@ class Scheduler(SchedulerInterface): self.need_mamba_block_aligned_split = ( self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" ) + + # Counts of non-empty steps scheduled / processed. update_from_output + # is called once per scheduled step in FIFO order, so these stay in sync. + self.sched_step_seq = 0 + self.processed_step_seq = 0 + # FIFO of (fence_seq, blocks): blocks become safe to free once + # processed_step_seq >= fence_seq. + self.deferred_frees: deque[tuple[int, list[KVCacheBlock]]] = deque() + self.perf_metrics: ModelMetrics | None = None if self.log_stats and vllm_config.observability_config.enable_mfu_metrics: self.perf_metrics = ModelMetrics(vllm_config) @@ -1044,6 +1062,11 @@ class Scheduler(SchedulerInterface): ) scheduler_output.ec_connector_metadata = ec_meta + # Advance the fence only for non-empty steps (those that actually + # write KV and have their output processed later in update_from_output). + if self.defer_block_free and total_num_scheduled_tokens > 0: + self.sched_step_seq += 1 + with record_function_or_nullcontext("schedule: update_after_schedule"): self._update_after_schedule(scheduler_output) return scheduler_output @@ -1062,7 +1085,7 @@ class Scheduler(SchedulerInterface): assert request.status == RequestStatus.RUNNING, ( "Only running requests can be preempted" ) - self.kv_cache_manager.free(request) + self._free_request_blocks(request) self.encoder_cache_manager.free(request) self._inflight_prefills.discard(request) request.status = RequestStatus.PREEMPTED @@ -1090,6 +1113,9 @@ class Scheduler(SchedulerInterface): for req_id, num_scheduled_token in num_scheduled_tokens.items(): request = self.requests[req_id] request.num_computed_tokens += num_scheduled_token + if self.defer_block_free: + # Record the in-flight step, to fence deferred block freeing. + request.last_sched_seq = self.sched_step_seq request.is_prefill_chunk = request.num_computed_tokens < ( request.num_tokens + request.num_output_placeholders ) @@ -1422,6 +1448,12 @@ class Scheduler(SchedulerInterface): kv_connector_output = model_runner_output.kv_connector_output cudagraph_stats = model_runner_output.cudagraph_stats + # Every GPU write enqueued by this and earlier steps has completed, so it is + # safe to return deferred-free blocks to the pool. + if self.defer_block_free and scheduler_output.total_num_scheduled_tokens > 0: + self.processed_step_seq += 1 + self._drain_deferred_frees() + perf_stats: PerfStats | None = None if self.perf_metrics and self.perf_metrics.is_enabled(): perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output) @@ -2006,7 +2038,7 @@ class Scheduler(SchedulerInterface): def _free_blocks(self, request: Request): assert request.is_finished() - self.kv_cache_manager.free(request) + self._free_request_blocks(request) del self.requests[request.request_id] @property @@ -2016,6 +2048,35 @@ class Scheduler(SchedulerInterface): def set_pause_state(self, pause_state: PauseState) -> None: self._pause_state = pause_state + def _free_request_blocks(self, request: Request): + """Free the request's KV blocks, deferring the return to the block + pool when an in-flight GPU step may still write them. + """ + if not self.defer_block_free or ( + # Last scheduled step already processed: no in-flight write remains + # (always the case for a normal finish), so free now. + request.last_sched_seq <= self.processed_step_seq + ): + self.kv_cache_manager.free(request) + return + blocks = self.kv_cache_manager.pop_blocks_for_free(request) + if blocks: + self.deferred_frees.append((self.sched_step_seq, blocks)) + + def _drain_deferred_frees(self): + """Return deferred blocks whose fence step has completed. + + Entries are appended with monotonically non-decreasing fences, so + stop at the first one that is still pending. + """ + while self.deferred_frees: + fence, _ = self.deferred_frees[0] + if fence > self.processed_step_seq: + break + _, blocks = self.deferred_frees.popleft() + # Free in reverse order so that the tail blocks are evicted first. + self.kv_cache_manager.block_pool.free_blocks(reversed(blocks)) + def get_num_unfinished_requests(self) -> int: if self._pause_state == PauseState.PAUSED_ALL: return 0 diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index bfc396c23c3..ad47e321e16 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -378,6 +378,24 @@ class SingleTypeKVCacheManager(ABC): """ return None + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + """ + Pop the request's bookkeeping and return its blocks without yet + returning them to the block pool. The caller is responsible for + eventually passing the returned blocks to `block_pool.free_blocks`, + freeing them in reverse order (so that tail blocks are evicted first). + + Args: + request_id: The request ID. + + Returns: + The request's blocks in allocation order. + """ + # Default to [] in case a request is freed (aborted) before alloc. + req_blocks = self.req_to_blocks.pop(request_id, []) + self.num_cached_block.pop(request_id, None) + return req_blocks + def free(self, request_id: str) -> None: """ Free the blocks for the request. @@ -385,15 +403,8 @@ class SingleTypeKVCacheManager(ABC): Args: request_id: The request ID. """ - # Default to [] in case a request is freed (aborted) before alloc. - req_blocks = self.req_to_blocks.pop(request_id, []) - - # Free blocks in reverse order so that the tail blocks are - # freed first. - ordered_blocks = reversed(req_blocks) - - self.block_pool.free_blocks(ordered_blocks) - self.num_cached_block.pop(request_id, None) + # Free blocks in reverse order so that the tail blocks are freed first. + self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id))) @abstractmethod def get_num_common_prefix_blocks(self, running_request_id: str) -> int: @@ -1212,11 +1223,11 @@ class MambaManager(SingleTypeKVCacheManager): self._allocated_block_reqs.add(request_id) return req_blocks[prev_block_len:] - def free(self, request_id: str) -> None: + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: if self.mamba_cache_mode == "align": self._allocated_block_reqs.discard(request_id) self.last_state_block_idx.pop(request_id, None) - super().free(request_id) + return super().pop_blocks_for_free(request_id) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: """ diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 44246e70a8b..0e8d4ee006f 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -145,6 +145,10 @@ class Request: # so the worker's broadcast slot ring stays consistent. self.next_decode_eligible_step = 0 + # Seq of the most recent step this request was scheduled in; fences + # deferred block freeing (see Scheduler._free_request_blocks). + self.last_sched_seq = 0 + self.spec_token_ids: list[int] = [] self.num_computed_tokens = 0 self.cache_salt: str | None = cache_salt