A high-throughput and memory-efficient inference and serving engine for LLMs
[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>
llx committed
Jun 16, 2026 at 05:36 UTC
d467a2a7f2f088dd360c7bef2f3cf5c59a1ffde8
9 files changed
+541
-18
tests/v1/core/test_async_scheduler.py
+1
@@ -284,6 +284,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
284
scheduler.vllm_config.model_config.enable_return_routed_experts = False
285
scheduler.enable_return_routed_experts = False
286
scheduler.recompute_kv_load_failures = False
287
+ scheduler.defer_block_free = False
288
scheduler.make_stats = Mock(return_value=None)
289
scheduler.max_model_len = 128
290
tests/v1/core/test_deferred_block_free.py
new
+414
@@ -0,0 +1,414 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for deferred block freeing under async scheduling.
4
+
5
+With async scheduling, a finished/preempted request's blocks may still be
6
+written by a speculatively over-scheduled in-flight GPU step (mamba/GDN
7
+layers rewrite the whole state block every step). If such a block is
8
+reallocated to a request arriving via PD disaggregation, the NIC/RDMA write
9
+of the received state races with the in-flight stale write. The scheduler
10
+closes the race by deferring the return of blocks to the block pool until
11
+the newest scheduled step's output has been processed.
12
+"""
13
+
14
+import os
15
+import time
16
+from unittest.mock import PropertyMock, patch
17
+
18
+import pytest
19
+
20
+from vllm.config import VllmConfig
21
+from vllm.v1.core.sched.output import SchedulerOutput
22
+from vllm.v1.outputs import ModelRunnerOutput
23
+from vllm.v1.request import RequestStatus
24
+
25
+from .utils import create_requests, create_scheduler, mock_kv
26
+
27
+pytestmark = pytest.mark.cpu_test
28
+
29
+# Allow overriding the model with a local path for offline environments.
30
+MODEL = os.environ.get("VLLM_TEST_DEFER_FREE_MODEL", "facebook/opt-125m")
31
+STOP_TOKEN_ID = 42
32
+NUM_PROMPT_TOKENS = 33 # 3 blocks with block_size=16
33
+
34
+
35
+def _make_model_runner_output(
36
+ scheduler_output: SchedulerOutput,
37
+ token_id: int = 0,
38
+) -> ModelRunnerOutput:
39
+ req_ids = list(scheduler_output.num_scheduled_tokens.keys())
40
+ return ModelRunnerOutput(
41
+ req_ids=req_ids,
42
+ req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},
43
+ sampled_token_ids=[[token_id] for _ in req_ids],
44
+ logprobs=None,
45
+ prompt_logprobs_dict={},
46
+ pooler_output=[],
47
+ )
48
+
49
+
50
+def _create_deferring_scheduler():
51
+ """Async scheduler with deferred block freeing forced on.
52
+
53
+ The production gate additionally requires a PD KV-consumer connector;
54
+ the mechanism itself is independent of it.
55
+ """
56
+ scheduler = create_scheduler(model=MODEL, async_scheduling=True)
57
+ scheduler.defer_block_free = True
58
+ return scheduler
59
+
60
+
61
+def _setup_request_with_inflight_step(scheduler, max_tokens: int = 5):
62
+ """Schedule a request's prefill (step 1) and one speculatively
63
+ over-scheduled decode (step 2), mimicking async scheduling depth 1.
64
+
65
+ Returns (request, out0, out1).
66
+ """
67
+ request = create_requests(
68
+ num_requests=1,
69
+ num_tokens=NUM_PROMPT_TOKENS,
70
+ max_tokens=max_tokens,
71
+ stop_token_ids=[STOP_TOKEN_ID],
72
+ )[0]
73
+ scheduler.add_request(request)
74
+ out0 = scheduler.schedule()
75
+ assert out0.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS
76
+ out1 = scheduler.schedule()
77
+ assert out1.num_scheduled_tokens[request.request_id] == 1
78
+ return request, out0, out1
79
+
80
+
81
+def test_gate_enabled_for_async_consumer():
82
+ # Overlapping batches + consumer-side connector enables the gate. Async
83
+ # scheduling (which would give >1 concurrent batches) is force-disabled on
84
+ # CPU, where this test runs, and PP can't be built without GPUs, so force
85
+ # max_concurrent_batches to exercise the enabled path on any platform.
86
+ with patch.object(
87
+ VllmConfig,
88
+ "max_concurrent_batches",
89
+ new_callable=PropertyMock,
90
+ return_value=2,
91
+ ):
92
+ scheduler = create_scheduler(
93
+ model=MODEL,
94
+ async_scheduling=True,
95
+ use_kv_connector=mock_kv(matched_tokens=0, is_async=False),
96
+ )
97
+ assert scheduler.defer_block_free
98
+
99
+
100
+def test_gate_disabled_without_connector():
101
+ # Async scheduling alone (no PD connector): the gate must stay off
102
+ # and freeing must remain immediate.
103
+ scheduler = create_scheduler(model=MODEL, async_scheduling=True)
104
+ assert not scheduler.defer_block_free
105
+
106
+ pool = scheduler.kv_cache_manager.block_pool
107
+ num_free_initially = pool.get_num_free_blocks()
108
+
109
+ request, out0, out1 = _setup_request_with_inflight_step(scheduler)
110
+ assert pool.get_num_free_blocks() < num_free_initially
111
+
112
+ # Request stops early while step 2 is in flight: blocks are freed
113
+ # immediately because deferral is disabled.
114
+ scheduler.update_from_output(
115
+ out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID)
116
+ )
117
+ assert request.is_finished()
118
+ assert not scheduler.deferred_frees
119
+ assert pool.get_num_free_blocks() == num_free_initially
120
+
121
+
122
+def test_finish_defers_free_until_inflight_step_done():
123
+ scheduler = _create_deferring_scheduler()
124
+ pool = scheduler.kv_cache_manager.block_pool
125
+ num_free_initially = pool.get_num_free_blocks()
126
+
127
+ request, out0, out1 = _setup_request_with_inflight_step(scheduler)
128
+ num_free_running = pool.get_num_free_blocks()
129
+ assert num_free_running < num_free_initially
130
+
131
+ # The request stops early (stop token) while the over-scheduled step 2
132
+ # is still in flight: its blocks must NOT return to the pool yet.
133
+ scheduler.update_from_output(
134
+ out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID)
135
+ )
136
+ assert request.is_finished()
137
+ assert len(scheduler.deferred_frees) == 1
138
+ assert pool.get_num_free_blocks() == num_free_running
139
+
140
+ # Step 2's output is processed: every GPU write of step 2 has
141
+ # completed, so the blocks can now be returned to the pool.
142
+ scheduler.update_from_output(out1, _make_model_runner_output(out1))
143
+ assert not scheduler.deferred_frees
144
+ assert pool.get_num_free_blocks() == num_free_initially
145
+
146
+
147
+def test_finish_frees_immediately_when_no_inflight_step():
148
+ scheduler = _create_deferring_scheduler()
149
+ pool = scheduler.kv_cache_manager.block_pool
150
+ num_free_initially = pool.get_num_free_blocks()
151
+
152
+ request = create_requests(
153
+ num_requests=1,
154
+ num_tokens=NUM_PROMPT_TOKENS,
155
+ max_tokens=5,
156
+ stop_token_ids=[STOP_TOKEN_ID],
157
+ )[0]
158
+ scheduler.add_request(request)
159
+ out0 = scheduler.schedule()
160
+
161
+ # Synchronous-like flow: out0 is the newest scheduled step and its
162
+ # output is being processed, so no other step can still write the
163
+ # blocks and the free happens immediately.
164
+ scheduler.update_from_output(
165
+ out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID)
166
+ )
167
+ assert request.is_finished()
168
+ assert not scheduler.deferred_frees
169
+ assert pool.get_num_free_blocks() == num_free_initially
170
+
171
+
172
+def test_abort_defers_free():
173
+ scheduler = _create_deferring_scheduler()
174
+ pool = scheduler.kv_cache_manager.block_pool
175
+ num_free_initially = pool.get_num_free_blocks()
176
+
177
+ request, out0, out1 = _setup_request_with_inflight_step(scheduler)
178
+ num_free_running = pool.get_num_free_blocks()
179
+
180
+ # External abort arrives while steps 1 and 2 are both in flight.
181
+ scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)
182
+ assert len(scheduler.deferred_frees) == 1
183
+ assert pool.get_num_free_blocks() == num_free_running
184
+
185
+ # Step 1's output: step 2 is still in flight, keep holding the blocks.
186
+ scheduler.update_from_output(out0, _make_model_runner_output(out0))
187
+ assert len(scheduler.deferred_frees) == 1
188
+ assert pool.get_num_free_blocks() == num_free_running
189
+
190
+ # Step 2's output: now the blocks can be freed.
191
+ scheduler.update_from_output(out1, _make_model_runner_output(out1))
192
+ assert not scheduler.deferred_frees
193
+ assert pool.get_num_free_blocks() == num_free_initially
194
+
195
+
196
+def test_preempt_defers_free_and_clears_bookkeeping():
197
+ scheduler = _create_deferring_scheduler()
198
+ pool = scheduler.kv_cache_manager.block_pool
199
+ num_free_initially = pool.get_num_free_blocks()
200
+
201
+ request, out0, out1 = _setup_request_with_inflight_step(scheduler)
202
+ num_free_running = pool.get_num_free_blocks()
203
+
204
+ # Preempt the request while steps are in flight (mirrors the
205
+ # preemption path inside schedule()).
206
+ scheduler.running.remove(request)
207
+ scheduler._preempt_request(request, time.monotonic())
208
+ assert request.status == RequestStatus.PREEMPTED
209
+
210
+ # Blocks are withheld from the pool, but the manager bookkeeping is
211
+ # cleared immediately so the request can be rescheduled safely.
212
+ assert len(scheduler.deferred_frees) == 1
213
+ assert pool.get_num_free_blocks() == num_free_running
214
+ for manager in scheduler.kv_cache_manager.coordinator.single_type_managers:
215
+ assert request.request_id not in manager.req_to_blocks
216
+
217
+ # Outputs of both in-flight steps are processed: blocks return to the
218
+ # pool only after the newest one.
219
+ scheduler.update_from_output(out0, _make_model_runner_output(out0))
220
+ assert len(scheduler.deferred_frees) == 1
221
+ scheduler.update_from_output(out1, _make_model_runner_output(out1))
222
+ assert not scheduler.deferred_frees
223
+ assert pool.get_num_free_blocks() == num_free_initially
224
+
225
+
226
+def test_multiple_deferred_frees_drain_in_order():
227
+ scheduler = _create_deferring_scheduler()
228
+ pool = scheduler.kv_cache_manager.block_pool
229
+ num_free_initially = pool.get_num_free_blocks()
230
+
231
+ requests = create_requests(
232
+ num_requests=2,
233
+ num_tokens=NUM_PROMPT_TOKENS,
234
+ max_tokens=5,
235
+ stop_token_ids=[STOP_TOKEN_ID],
236
+ )
237
+ for request in requests:
238
+ scheduler.add_request(request)
239
+ out0 = scheduler.schedule()
240
+ out1 = scheduler.schedule()
241
+
242
+ # Both requests stop early at step 1's output while step 2 is in
243
+ # flight: two deferred entries with the same fence.
244
+ scheduler.update_from_output(
245
+ out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID)
246
+ )
247
+ assert len(scheduler.deferred_frees) == 2
248
+ assert pool.get_num_free_blocks() < num_free_initially
249
+
250
+ scheduler.update_from_output(out1, _make_model_runner_output(out1))
251
+ assert not scheduler.deferred_frees
252
+ assert pool.get_num_free_blocks() == num_free_initially
253
+
254
+
255
+def test_fence_held_across_multiple_inflight_steps():
256
+ """Pipeline-parallel / deep async: with several steps scheduled ahead,
257
+ a freed request's blocks must stay held until the *newest* in-flight
258
+ step's output is processed, not the first.
259
+
260
+ Depth-1 tests only check a single intervening update; with PP the
261
+ scheduler can dispatch up to pp_size steps ahead, so the fence must
262
+ survive multiple intervening update_from_output calls.
263
+ """
264
+ scheduler = _create_deferring_scheduler()
265
+ pool = scheduler.kv_cache_manager.block_pool
266
+ num_free_initially = pool.get_num_free_blocks()
267
+
268
+ request = create_requests(
269
+ num_requests=1,
270
+ num_tokens=NUM_PROMPT_TOKENS,
271
+ max_tokens=10,
272
+ )[0]
273
+ scheduler.add_request(request)
274
+
275
+ # Schedule three steps ahead without processing any output: a prefill
276
+ # plus two speculatively over-scheduled decodes, all in flight at once.
277
+ outs = [scheduler.schedule() for _ in range(3)]
278
+ assert outs[0].num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS
279
+ assert outs[1].num_scheduled_tokens[request.request_id] == 1
280
+ assert outs[2].num_scheduled_tokens[request.request_id] == 1
281
+ assert scheduler.sched_step_seq == 3
282
+ num_free_running = pool.get_num_free_blocks()
283
+ assert num_free_running < num_free_initially
284
+
285
+ # Abort while all three steps are in flight: the fence is the newest
286
+ # scheduled step (3), since any of them may still write the blocks.
287
+ scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)
288
+ assert len(scheduler.deferred_frees) == 1
289
+ assert scheduler.deferred_frees[0][0] == 3
290
+ assert pool.get_num_free_blocks() == num_free_running
291
+
292
+ # Draining the two earlier in-flight steps must NOT release the blocks:
293
+ # their outputs don't fence the still-pending newest write.
294
+ for out in (outs[0], outs[1]):
295
+ scheduler.update_from_output(out, _make_model_runner_output(out))
296
+ assert len(scheduler.deferred_frees) == 1
297
+ assert pool.get_num_free_blocks() == num_free_running
298
+
299
+ # Only once the newest scheduled step's output is processed do the
300
+ # blocks return to the pool.
301
+ scheduler.update_from_output(outs[2], _make_model_runner_output(outs[2]))
302
+ assert not scheduler.deferred_frees
303
+ assert pool.get_num_free_blocks() == num_free_initially
304
+
305
+
306
+def test_max_tokens_finish_frees_immediately_with_other_inflight():
307
+ """A request finishing by reaching max_tokens is never over-scheduled past
308
+ its final-token step, so no in-flight step writes its blocks: it is freed
309
+ immediately even while another request's step is still in flight.
310
+ """
311
+ scheduler = _create_deferring_scheduler()
312
+ pool = scheduler.kv_cache_manager.block_pool
313
+
314
+ # Short request finishes at max_tokens=1; long request keeps running.
315
+ short = create_requests(
316
+ num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=1, req_ids=["short"]
317
+ )[0]
318
+ long = create_requests(
319
+ num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=100, req_ids=["long"]
320
+ )[0]
321
+ scheduler.add_request(short)
322
+ scheduler.add_request(long)
323
+
324
+ out0 = scheduler.schedule() # prefill both
325
+ out1 = scheduler.schedule() # short is skipped (at max_tokens); long decodes
326
+ assert "short" not in out1.num_scheduled_tokens
327
+ assert "long" in out1.num_scheduled_tokens
328
+
329
+ free_before = pool.get_num_free_blocks()
330
+ # Process step 0: `short` reaches max_tokens and finishes while step 1
331
+ # (which scheduled `long`, not `short`) is still in flight.
332
+ scheduler.update_from_output(out0, _make_model_runner_output(out0))
333
+
334
+ assert short.is_finished()
335
+ # A step IS globally in flight (the old global fence would have deferred),
336
+ # but the per-request gate frees `short` immediately since nothing writes
337
+ # its blocks anymore.
338
+ assert scheduler.sched_step_seq > scheduler.processed_step_seq
339
+ assert not scheduler.deferred_frees
340
+ assert pool.get_num_free_blocks() > free_before # short's blocks returned
341
+
342
+
343
+def test_abort_mid_prefill_defers_free():
344
+ """Intermediate prefill chunks don't allocate output placeholders, so the
345
+ deferral must key off is_prefill_chunk: aborting a request whose prefill
346
+ chunk is still in flight must withhold its blocks.
347
+ """
348
+ scheduler = create_scheduler(
349
+ model=MODEL, async_scheduling=True, long_prefill_token_threshold=16
350
+ )
351
+ scheduler.defer_block_free = True
352
+ pool = scheduler.kv_cache_manager.block_pool
353
+ num_free_initially = pool.get_num_free_blocks()
354
+
355
+ request = create_requests(
356
+ num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=5
357
+ )[0]
358
+ scheduler.add_request(request)
359
+
360
+ out0 = scheduler.schedule()
361
+ # Partial prefill: a chunk is in flight, with no output placeholders yet.
362
+ assert out0.num_scheduled_tokens[request.request_id] == 16
363
+ assert request.num_output_placeholders == 0
364
+ assert request.is_prefill_chunk
365
+ num_free_running = pool.get_num_free_blocks()
366
+ assert num_free_running < num_free_initially
367
+
368
+ # Abort while the prefill chunk is in flight: blocks must be withheld
369
+ # (keyed off is_prefill_chunk, since there are no placeholders).
370
+ scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)
371
+ assert len(scheduler.deferred_frees) == 1
372
+ assert pool.get_num_free_blocks() == num_free_running
373
+
374
+ # Once the in-flight prefill step's output is processed, blocks return.
375
+ scheduler.update_from_output(out0, _make_model_runner_output(out0))
376
+ assert not scheduler.deferred_frees
377
+ assert pool.get_num_free_blocks() == num_free_initially
378
+
379
+
380
+def test_non_async_abort_defers_via_last_sched_seq():
381
+ """Without async (e.g. PP filling the pipeline) there are no placeholders
382
+ and a full prefill isn't a partial chunk, yet an abort with a step in flight
383
+ must defer. Only the last-scheduled-step fence catches this.
384
+
385
+ PP=2 can't be built on a single-GPU host, so force the flag and exercise the
386
+ mechanism; the gate itself is covered by test_gate_enabled_for_async_consumer.
387
+ """
388
+ scheduler = create_scheduler(model=MODEL, async_scheduling=False)
389
+ scheduler.defer_block_free = True
390
+ pool = scheduler.kv_cache_manager.block_pool
391
+ num_free_initially = pool.get_num_free_blocks()
392
+
393
+ request = create_requests(
394
+ num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=5
395
+ )[0]
396
+ scheduler.add_request(request)
397
+
398
+ out0 = scheduler.schedule()
399
+ # Neither async-only signal marks this request as in flight.
400
+ assert request.num_output_placeholders == 0
401
+ assert not request.is_prefill_chunk
402
+ # Only the last-scheduled-step fence does.
403
+ assert request.last_sched_seq > scheduler.processed_step_seq
404
+ num_free_running = pool.get_num_free_blocks()
405
+ assert num_free_running < num_free_initially
406
+
407
+ # Abort while out0 is in flight: blocks must be withheld.
408
+ scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)
409
+ assert len(scheduler.deferred_frees) == 1
410
+ assert pool.get_num_free_blocks() == num_free_running
411
+
412
+ scheduler.update_from_output(out0, _make_model_runner_output(out0))
413
+ assert not scheduler.deferred_frees
414
+ assert pool.get_num_free_blocks() == num_free_initially
tests/v1/core/test_scheduler.py
+1
@@ -2571,6 +2571,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
2571
scheduler.vllm_config.model_config.enable_return_routed_experts = False
2572
scheduler.enable_return_routed_experts = False
2573
scheduler.recompute_kv_load_failures = False
2574
+ scheduler.defer_block_free = False
2575
scheduler.make_stats = Mock(return_value=None)
2576
scheduler.max_model_len = 128
2577
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+2
-3
@@ -24,11 +24,10 @@ dp_ep_configs=(
24
# We assume HMA enabled by default.
25
hybrid_ssm_configs=(
26
"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"
27
- # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
28
- "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"
27
+ "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"
28
# GDN (Qwen3.5)
29
"VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B"
31
- "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"
30
+ "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"
31
)
32
sw_attn_configs=(
33
# NOTE: gemma3 does not work with FlashInfer
vllm/v1/core/kv_cache_coordinator.py
+19
@@ -291,6 +291,25 @@ class KVCacheCoordinator(ABC):
291
for manager in self.single_type_managers:
292
manager.free(request_id)
293
294
+ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
295
+ """
296
+ Pop the request's bookkeeping from all single-type managers and
297
+ return its blocks without returning them to the block pool. The
298
+ caller must eventually pass the returned blocks to
299
+ `block_pool.free_blocks`, freeing them in reverse order (so that
300
+ tail blocks are evicted first).
301
+
302
+ Args:
303
+ request_id: The request ID.
304
+
305
+ Returns:
306
+ The request's blocks in allocation order.
307
+ """
308
+ blocks: list[KVCacheBlock] = []
309
+ for manager in self.single_type_managers:
310
+ blocks.extend(manager.pop_blocks_for_free(request_id))
311
+ return blocks
312
+
313
def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
314
"""
315
Get the number of common prefix blocks for all requests with allocated
vllm/v1/core/kv_cache_manager.py
+13
@@ -480,6 +480,19 @@ class KVCacheManager:
480
"""
481
self.coordinator.remove_skipped_blocks(request_id, total_computed_tokens)
482
483
+ def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]:
484
+ """Pop the request's bookkeeping and return its blocks without
485
+ returning them to the block pool. The caller must eventually free
486
+ them in reverse order (so that tail blocks are evicted first).
487
+
488
+ Args:
489
+ request: The request to pop the blocks for.
490
+
491
+ Returns:
492
+ The request's blocks in allocation order.
493
+ """
494
+ return self.coordinator.pop_blocks_for_free(request.request_id)
495
+
496
def evict_blocks(self, block_ids: set[int]) -> None:
497
"""evict blocks from the prefix cache by their block IDs.
498
vllm/v1/core/sched/scheduler.py
+67
-6
@@ -37,6 +37,7 @@ from vllm.v1.core.encoder_cache_manager import (
37
from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
38
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
39
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
40
+from vllm.v1.core.kv_cache_utils import KVCacheBlock
41
from vllm.v1.core.sched.interface import PauseState, SchedulerInterface
42
from vllm.v1.core.sched.output import (
43
CachedRequestData,
@@ -125,7 +126,9 @@ class Scheduler(SchedulerInterface):
126
self.connector = None
127
self.connector_prefix_cache_stats: PrefixCacheStats | None = None
128
self.recompute_kv_load_failures = True
128
- if self.vllm_config.kv_transfer_config is not None:
129
+ self.defer_block_free = False
130
+ kv_transfer_config = self.vllm_config.kv_transfer_config
131
+ if kv_transfer_config is not None:
132
assert not self.is_encoder_decoder, (
133
"Encoder-decoder models are not currently supported with KV connectors"
134
)
@@ -136,11 +139,17 @@ class Scheduler(SchedulerInterface):
139
)
140
if self.log_stats:
141
self.connector_prefix_cache_stats = PrefixCacheStats()
139
- kv_load_failure_policy = (
140
- self.vllm_config.kv_transfer_config.kv_load_failure_policy
141
- )
142
+ kv_load_failure_policy = kv_transfer_config.kv_load_failure_policy
143
self.recompute_kv_load_failures = kv_load_failure_policy == "recompute"
144
145
+ # With overlapping batches (async scheduling or PP), a step may
146
+ # still be writing a freed request's KV blocks. A consumer KV
147
+ # Connector can reallocate and fill those blocks via a load that
148
+ # isn't ordered against that write, so defer freeing them.
149
+ multiple_inflight_batches = self.vllm_config.max_concurrent_batches > 1
150
+ if multiple_inflight_batches and kv_transfer_config.is_kv_consumer:
151
+ self.defer_block_free = True
152
+
153
self.kv_event_publisher = EventPublisherFactory.create(
154
self.kv_events_config,
155
self.parallel_config.data_parallel_index,
@@ -275,6 +284,15 @@ class Scheduler(SchedulerInterface):
284
self.need_mamba_block_aligned_split = (
285
self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align"
286
)
287
+
288
+ # Counts of non-empty steps scheduled / processed. update_from_output
289
+ # is called once per scheduled step in FIFO order, so these stay in sync.
290
+ self.sched_step_seq = 0
291
+ self.processed_step_seq = 0
292
+ # FIFO of (fence_seq, blocks): blocks become safe to free once
293
+ # processed_step_seq >= fence_seq.
294
+ self.deferred_frees: deque[tuple[int, list[KVCacheBlock]]] = deque()
295
+
296
self.perf_metrics: ModelMetrics | None = None
297
if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:
298
self.perf_metrics = ModelMetrics(vllm_config)
@@ -1044,6 +1062,11 @@ class Scheduler(SchedulerInterface):
1062
)
1063
scheduler_output.ec_connector_metadata = ec_meta
1064
1065
+ # Advance the fence only for non-empty steps (those that actually
1066
+ # write KV and have their output processed later in update_from_output).
1067
+ if self.defer_block_free and total_num_scheduled_tokens > 0:
1068
+ self.sched_step_seq += 1
1069
+
1070
with record_function_or_nullcontext("schedule: update_after_schedule"):
1071
self._update_after_schedule(scheduler_output)
1072
return scheduler_output
@@ -1062,7 +1085,7 @@ class Scheduler(SchedulerInterface):
1085
assert request.status == RequestStatus.RUNNING, (
1086
"Only running requests can be preempted"
1087
)
1065
- self.kv_cache_manager.free(request)
1088
+ self._free_request_blocks(request)
1089
self.encoder_cache_manager.free(request)
1090
self._inflight_prefills.discard(request)
1091
request.status = RequestStatus.PREEMPTED
@@ -1090,6 +1113,9 @@ class Scheduler(SchedulerInterface):
1113
for req_id, num_scheduled_token in num_scheduled_tokens.items():
1114
request = self.requests[req_id]
1115
request.num_computed_tokens += num_scheduled_token
1116
+ if self.defer_block_free:
1117
+ # Record the in-flight step, to fence deferred block freeing.
1118
+ request.last_sched_seq = self.sched_step_seq
1119
request.is_prefill_chunk = request.num_computed_tokens < (
1120
request.num_tokens + request.num_output_placeholders
1121
)
@@ -1422,6 +1448,12 @@ class Scheduler(SchedulerInterface):
1448
kv_connector_output = model_runner_output.kv_connector_output
1449
cudagraph_stats = model_runner_output.cudagraph_stats
1450
1451
+ # Every GPU write enqueued by this and earlier steps has completed, so it is
1452
+ # safe to return deferred-free blocks to the pool.
1453
+ if self.defer_block_free and scheduler_output.total_num_scheduled_tokens > 0:
1454
+ self.processed_step_seq += 1
1455
+ self._drain_deferred_frees()
1456
+
1457
perf_stats: PerfStats | None = None
1458
if self.perf_metrics and self.perf_metrics.is_enabled():
1459
perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
@@ -2006,7 +2038,7 @@ class Scheduler(SchedulerInterface):
2038
2039
def _free_blocks(self, request: Request):
2040
assert request.is_finished()
2009
- self.kv_cache_manager.free(request)
2041
+ self._free_request_blocks(request)
2042
del self.requests[request.request_id]
2043
2044
@property
@@ -2016,6 +2048,35 @@ class Scheduler(SchedulerInterface):
2048
def set_pause_state(self, pause_state: PauseState) -> None:
2049
self._pause_state = pause_state
2050
2051
+ def _free_request_blocks(self, request: Request):
2052
+ """Free the request's KV blocks, deferring the return to the block
2053
+ pool when an in-flight GPU step may still write them.
2054
+ """
2055
+ if not self.defer_block_free or (
2056
+ # Last scheduled step already processed: no in-flight write remains
2057
+ # (always the case for a normal finish), so free now.
2058
+ request.last_sched_seq <= self.processed_step_seq
2059
+ ):
2060
+ self.kv_cache_manager.free(request)
2061
+ return
2062
+ blocks = self.kv_cache_manager.pop_blocks_for_free(request)
2063
+ if blocks:
2064
+ self.deferred_frees.append((self.sched_step_seq, blocks))
2065
+
2066
+ def _drain_deferred_frees(self):
2067
+ """Return deferred blocks whose fence step has completed.
2068
+
2069
+ Entries are appended with monotonically non-decreasing fences, so
2070
+ stop at the first one that is still pending.
2071
+ """
2072
+ while self.deferred_frees:
2073
+ fence, _ = self.deferred_frees[0]
2074
+ if fence > self.processed_step_seq:
2075
+ break
2076
+ _, blocks = self.deferred_frees.popleft()
2077
+ # Free in reverse order so that the tail blocks are evicted first.
2078
+ self.kv_cache_manager.block_pool.free_blocks(reversed(blocks))
2079
+
2080
def get_num_unfinished_requests(self) -> int:
2081
if self._pause_state == PauseState.PAUSED_ALL:
2082
return 0
vllm/v1/core/single_type_kv_cache_manager.py
+20
-9
@@ -378,22 +378,33 @@ class SingleTypeKVCacheManager(ABC):
378
"""
379
return None
380
381
- def free(self, request_id: str) -> None:
381
+ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
382
"""
383
- Free the blocks for the request.
383
+ Pop the request's bookkeeping and return its blocks without yet
384
+ returning them to the block pool. The caller is responsible for
385
+ eventually passing the returned blocks to `block_pool.free_blocks`,
386
+ freeing them in reverse order (so that tail blocks are evicted first).
387
388
Args:
389
request_id: The request ID.
390
+
391
+ Returns:
392
+ The request's blocks in allocation order.
393
"""
394
# Default to [] in case a request is freed (aborted) before alloc.
395
req_blocks = self.req_to_blocks.pop(request_id, [])
396
+ self.num_cached_block.pop(request_id, None)
397
+ return req_blocks
398
391
- # Free blocks in reverse order so that the tail blocks are
392
- # freed first.
393
- ordered_blocks = reversed(req_blocks)
399
+ def free(self, request_id: str) -> None:
400
+ """
401
+ Free the blocks for the request.
402
395
- self.block_pool.free_blocks(ordered_blocks)
396
- self.num_cached_block.pop(request_id, None)
403
+ Args:
404
+ request_id: The request ID.
405
+ """
406
+ # Free blocks in reverse order so that the tail blocks are freed first.
407
+ self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id)))
408
409
@abstractmethod
410
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
@@ -1212,11 +1223,11 @@ class MambaManager(SingleTypeKVCacheManager):
1223
self._allocated_block_reqs.add(request_id)
1224
return req_blocks[prev_block_len:]
1225
1215
- def free(self, request_id: str) -> None:
1226
+ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
1227
if self.mamba_cache_mode == "align":
1228
self._allocated_block_reqs.discard(request_id)
1229
self.last_state_block_idx.pop(request_id, None)
1219
- super().free(request_id)
1230
+ return super().pop_blocks_for_free(request_id)
1231
1232
def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
1233
"""
vllm/v1/request.py
+4
@@ -145,6 +145,10 @@ class Request:
145
# so the worker's broadcast slot ring stays consistent.
146
self.next_decode_eligible_step = 0
147
148
+ # Seq of the most recent step this request was scheduled in; fences
149
+ # deferred block freeing (see Scheduler._free_request_blocks).
150
+ self.last_sched_seq = 0
151
+
152
self.spec_token_ids: list[int] = []
153
self.num_computed_tokens = 0
154
self.cache_salt: str | None = cache_salt