[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 <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>

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
 

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

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
 

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

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

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.
 

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

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:
         """

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