@tduncan / Nvidia-Optimization-Stack-VLLM / commits / cd0de48d08

[Bugfix][V1] Free out-of-window blocks on the processed-token basis under async scheduling (#47728)

Signed-off-by: Saddss <28726669061@qq.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Saddss <28726669061@qq.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Saddss committed Jul 8, 2026 at 20:34 UTC cd0de48d0883ecb8e1ef350a99baa0c158f58e82
14 files changed +549 -143
tests/models/multimodal/generation/test_voxtral_realtime.py
+1 -1
@@ -81,7 +81,7 @@ def assert_encoder_kv_cache_spec(engine: LLM) -> None:
81 assert spec.sliding_window == cdiv(750, 4) + 1 == 189
82 assert (
83 spec.max_admission_blocks_per_request(
84 - max_num_batched_tokens=1,
84 + max_in_flight_tokens=1,
85 max_model_len=vllm_config.model_config.max_model_len,
86 )
87 == 13
tests/v1/core/test_kv_cache_utils.py
+3
@@ -1404,12 +1404,15 @@ def test_get_max_concurrency_for_kv_cache_config():
1404 enable_chunked_prefill=True,
1405 max_model_len=model_config.max_model_len,
1406 is_encoder_decoder=model_config.is_encoder_decoder,
1407 + # Pin to sync: SWA per-request bounds grow with overlapping batches.
1408 + async_scheduling=False,
1409 )
1410
1411 vllm_config = VllmConfig(
1412 model_config=model_config,
1413 scheduler_config=scheduler_config,
1414 )
1415 + assert vllm_config.max_concurrent_batches == 1
1416
1417 full_attention_spec = FullAttentionSpec(
1418 block_size=16,
tests/v1/core/test_prefix_caching.py
+4 -2
@@ -3453,7 +3453,8 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
3453 manager = make_kv_cache_manager(
3454 config,
3455 max_model_len=max_model_len,
3456 - max_num_batched_tokens=max_num_batched_tokens,
3456 + # Single (sync) batch in flight, so in-flight tokens == batched tokens.
3457 + max_in_flight_tokens=max_num_batched_tokens,
3458 enable_caching=True,
3459 hash_block_size=block_size,
3460 )
@@ -3509,7 +3510,8 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized():
3510 manager = make_kv_cache_manager(
3511 config,
3512 max_model_len=max_model_len,
3512 - max_num_batched_tokens=max_num_batched_tokens,
3513 + # Single (sync) batch in flight, so in-flight tokens == batched tokens.
3514 + max_in_flight_tokens=max_num_batched_tokens,
3515 enable_caching=True,
3516 hash_block_size=block_size,
3517 )
tests/v1/core/test_swa_inflight_window_free.py new
+259
@@ -0,0 +1,259 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Out-of-window block frees vs in-flight GPU steps.
4 +
5 +With async scheduling / PP, `num_computed_tokens` optimistically includes
6 +tokens of unprocessed steps whose attention windows still read the blocks
7 +just below the optimistic boundary (and rejected spec tokens can roll it
8 +back), so `allocate_slots` frees on the processed-token basis:
9 +`num_computed_tokens - num_in_flight_tokens`.
10 +"""
11 +
12 +import torch
13 +
14 +from vllm.v1.core.sched.output import SchedulerOutput
15 +from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec
16 +from vllm.v1.outputs import ModelRunnerOutput
17 +
18 +from .utils import create_requests, create_scheduler, mock_kv
19 +
20 +NUM_PROMPT_TOKENS = 100
21 +BLOCK_SIZE = 16
22 +SLIDING_WINDOW = 16
23 +# Tokens 0..84 are outside the window of the next token to compute
24 +# (100 - 16 + 1 = 85), i.e. 5 full blocks.
25 +NUM_OUT_OF_WINDOW_BLOCKS = 85 // BLOCK_SIZE
26 +# Chunked-local skips whole chunks left of the current one:
27 +# (100 // 32) * 32 = 96 settled tokens -> 6 full blocks.
28 +CHUNK_SIZE = 32
29 +NUM_OUT_OF_CHUNK_BLOCKS = (NUM_PROMPT_TOKENS // CHUNK_SIZE) * CHUNK_SIZE // BLOCK_SIZE
30 +
31 +
32 +def _make_model_runner_output(
33 + scheduler_output: SchedulerOutput,
34 + token_id: int = 0,
35 +) -> ModelRunnerOutput:
36 + req_ids = list(scheduler_output.num_scheduled_tokens.keys())
37 + return ModelRunnerOutput(
38 + req_ids=req_ids,
39 + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},
40 + sampled_token_ids=[[token_id] for _ in req_ids],
41 + logprobs=None,
42 + prompt_logprobs_dict={},
43 + pooler_output=[],
44 + )
45 +
46 +
47 +def _create_swa_scheduler(async_scheduling: bool):
48 + return create_scheduler(
49 + block_size=BLOCK_SIZE,
50 + async_scheduling=async_scheduling,
51 + kv_cache_spec=SlidingWindowSpec(
52 + block_size=BLOCK_SIZE,
53 + num_kv_heads=1,
54 + head_size=1,
55 + dtype=torch.float32,
56 + sliding_window=SLIDING_WINDOW,
57 + ),
58 + )
59 +
60 +
61 +def _create_chunked_scheduler(async_scheduling: bool):
62 + return create_scheduler(
63 + block_size=BLOCK_SIZE,
64 + async_scheduling=async_scheduling,
65 + kv_cache_spec=ChunkedLocalAttentionSpec(
66 + block_size=BLOCK_SIZE,
67 + num_kv_heads=1,
68 + head_size=1,
69 + dtype=torch.float32,
70 + attention_chunk_size=CHUNK_SIZE,
71 + ),
72 + )
73 +
74 +
75 +def _num_null_blocks(scheduler, request_id: str) -> int:
76 + manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0]
77 + null_block = manager._null_block
78 + return sum(1 for b in manager.req_to_blocks[request_id] if b is null_block)
79 +
80 +
81 +def test_num_in_flight_tokens_accounting():
82 + scheduler = create_scheduler(async_scheduling=True)
83 + request = create_requests(num_requests=1, num_tokens=NUM_PROMPT_TOKENS)[0]
84 + scheduler.add_request(request)
85 +
86 + out0 = scheduler.schedule()
87 + assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS
88 + # Async: decode scheduled before the prefill output is processed.
89 + out1 = scheduler.schedule()
90 + assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS + 1
91 +
92 + scheduler.update_from_output(out0, _make_model_runner_output(out0))
93 + assert request.num_in_flight_tokens == 1
94 + scheduler.update_from_output(out1, _make_model_runner_output(out1))
95 + assert request.num_in_flight_tokens == 0
96 +
97 +
98 +def test_swa_free_waits_for_in_flight_step():
99 + """Async: out-of-window blocks stay allocated until the step that still
100 + reads them has been processed."""
101 + scheduler = _create_swa_scheduler(async_scheduling=True)
102 + request = create_requests(
103 + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
104 + )[0]
105 + scheduler.add_request(request)
106 + req_id = request.request_id
107 + block_pool = scheduler.kv_cache_manager.block_pool
108 +
109 + out0 = scheduler.schedule() # prefill, in flight from here on
110 + free_after_prefill = block_pool.get_num_free_blocks()
111 +
112 + # Decode scheduled while the prefill still reads the out-of-window blocks:
113 + # they must not be freed yet.
114 + out1 = scheduler.schedule()
115 + assert _num_null_blocks(scheduler, req_id) == 0
116 + assert block_pool.get_num_free_blocks() == free_after_prefill
117 +
118 + # Prefill output processed; the next allocate frees the out-of-window
119 + # blocks.
120 + scheduler.update_from_output(out0, _make_model_runner_output(out0))
121 + scheduler.schedule()
122 + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
123 + assert (
124 + block_pool.get_num_free_blocks()
125 + == free_after_prefill + NUM_OUT_OF_WINDOW_BLOCKS
126 + )
127 + # Not double-freed on the following steps.
128 + scheduler.update_from_output(out1, _make_model_runner_output(out1))
129 + scheduler.schedule()
130 + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
131 +
132 +
133 +def test_swa_free_immediate_when_sync():
134 + """Sync: no in-flight step at schedule time, frees happen at the first
135 + decode allocation as before."""
136 + scheduler = _create_swa_scheduler(async_scheduling=False)
137 + request = create_requests(
138 + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
139 + )[0]
140 + scheduler.add_request(request)
141 + req_id = request.request_id
142 +
143 + out0 = scheduler.schedule()
144 + scheduler.update_from_output(out0, _make_model_runner_output(out0))
145 + assert request.num_in_flight_tokens == 0
146 +
147 + scheduler.schedule()
148 + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
149 +
150 +
151 +def test_swa_admission_cap_accounts_for_overlapping_batches():
152 + spec = SlidingWindowSpec(
153 + block_size=16,
154 + num_kv_heads=1,
155 + head_size=1,
156 + dtype=torch.float32,
157 + sliding_window=1024,
158 + )
159 + base = spec.max_admission_blocks_per_request(
160 + max_in_flight_tokens=1024, max_model_len=16384
161 + )
162 + # (1024 - 1 + 1024) tokens -> 128 blocks, +1 for window misalignment.
163 + assert base == 129
164 + overlapped = spec.max_admission_blocks_per_request(
165 + max_in_flight_tokens=2 * 1024, max_model_len=16384
166 + )
167 + # One extra in-flight chunk is held back: (1024 - 1 + 2 * 1024) tokens.
168 + assert overlapped == 193
169 +
170 +
171 +def test_chunked_local_free_waits_for_in_flight_step():
172 + """Chunked-local attention frees whole chunks left of the current one, and
173 + is exposed to the same load-WAR: with async scheduling those chunks must
174 + stay allocated until the in-flight step that still reads them settles."""
175 + scheduler = _create_chunked_scheduler(async_scheduling=True)
176 + request = create_requests(
177 + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
178 + )[0]
179 + scheduler.add_request(request)
180 + req_id = request.request_id
181 + block_pool = scheduler.kv_cache_manager.block_pool
182 +
183 + out0 = scheduler.schedule() # prefill, in flight from here on
184 + free_after_prefill = block_pool.get_num_free_blocks()
185 +
186 + # Decode scheduled while the prefill still reads the out-of-chunk blocks.
187 + out1 = scheduler.schedule()
188 + assert _num_null_blocks(scheduler, req_id) == 0
189 + assert block_pool.get_num_free_blocks() == free_after_prefill
190 +
191 + # Prefill output processed; the next allocate frees the out-of-chunk blocks.
192 + scheduler.update_from_output(out0, _make_model_runner_output(out0))
193 + scheduler.schedule()
194 + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS
195 + assert (
196 + block_pool.get_num_free_blocks() == free_after_prefill + NUM_OUT_OF_CHUNK_BLOCKS
197 + )
198 + # Not double-freed on the following steps.
199 + scheduler.update_from_output(out1, _make_model_runner_output(out1))
200 + scheduler.schedule()
201 + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS
202 +
203 +
204 +def test_chunked_local_admission_cap_accounts_for_overlapping_batches():
205 + spec = ChunkedLocalAttentionSpec(
206 + block_size=16,
207 + num_kv_heads=1,
208 + head_size=1,
209 + dtype=torch.float32,
210 + attention_chunk_size=1024,
211 + )
212 + base = spec.max_admission_blocks_per_request(
213 + max_in_flight_tokens=1024, max_model_len=16384
214 + )
215 + # (1024 + 1024) tokens -> 128 blocks.
216 + assert base == 128
217 + overlapped = spec.max_admission_blocks_per_request(
218 + max_in_flight_tokens=2 * 1024, max_model_len=16384
219 + )
220 + # One extra in-flight chunk is held back: (1024 + 2 * 1024) tokens.
221 + assert overlapped == 192
222 +
223 +
224 +def test_connector_finish_frees_on_settled_basis():
225 + """The out-of-window prune done at request finish, before the block table
226 + is handed to a KV connector (simple CPU offload / NIXL store), must use the
227 + same processed-token basis. Otherwise the connector reads/hands off a block
228 + the still-in-flight step is optimistically counted as done with, which is
229 + the load-path WAR this fix closes."""
230 + scheduler = create_scheduler(
231 + block_size=BLOCK_SIZE,
232 + async_scheduling=True,
233 + use_kv_connector=mock_kv(matched_tokens=0, is_async=False),
234 + kv_cache_spec=SlidingWindowSpec(
235 + block_size=BLOCK_SIZE,
236 + num_kv_heads=1,
237 + head_size=1,
238 + dtype=torch.float32,
239 + sliding_window=SLIDING_WINDOW,
240 + ),
241 + )
242 + request = create_requests(
243 + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
244 + )[0]
245 + scheduler.add_request(request)
246 + req_id = request.request_id
247 +
248 + out0 = scheduler.schedule() # prefill, in flight from here on
249 + scheduler.schedule() # decode over-scheduled: num_computed_tokens optimistic
250 +
251 + # Finishing now (connector store) must NOT prune out-of-window blocks the
252 + # in-flight prefill still reads.
253 + scheduler._connector_finished(request)
254 + assert _num_null_blocks(scheduler, req_id) == 0
255 +
256 + # Once the in-flight step settles, the same prune releases them.
257 + scheduler.update_from_output(out0, _make_model_runner_output(out0))
258 + scheduler._connector_finished(request)
259 + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
tests/v1/core/utils.py
+10 -11
@@ -30,6 +30,7 @@ from vllm.v1.kv_cache_interface import (
30 FullAttentionSpec,
31 KVCacheConfig,
32 KVCacheGroupSpec,
33 + KVCacheSpec,
34 )
35 from vllm.v1.request import Request
36 from vllm.v1.structured_output import StructuredOutputManager
@@ -63,6 +64,7 @@ def create_scheduler(
64 use_ec_connector: bool = False,
65 ec_role: str | None = None,
66 use_v2_model_runner: bool | None = None,
67 + kv_cache_spec: KVCacheSpec | None = None,
68 ) -> Scheduler | AsyncScheduler:
69 """Create scheduler under test.
70
@@ -164,20 +166,17 @@ def create_scheduler(
166 speculative_config=speculative_config,
167 ec_transfer_config=ec_transfer_config,
168 )
169 + if kv_cache_spec is None:
170 + kv_cache_spec = FullAttentionSpec(
171 + block_size=block_size,
172 + num_kv_heads=1,
173 + head_size=1,
174 + dtype=torch.float32,
175 + )
176 kv_cache_config = KVCacheConfig(
177 num_blocks=num_blocks, # A large number of blocks to hold all requests
178 kv_cache_tensors=[],
170 - kv_cache_groups=[
171 - KVCacheGroupSpec(
172 - ["layer"],
173 - FullAttentionSpec(
174 - block_size=block_size,
175 - num_kv_heads=1,
176 - head_size=1,
177 - dtype=torch.float32,
178 - ),
179 - )
180 - ],
179 + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
180 )
181 cache_config.num_gpu_blocks = num_blocks
182 register_all_kvcache_specs(vllm_config)
tests/v1/e2e/general/test_mamba_prefix_cache.py
+183 -61
@@ -39,13 +39,18 @@ from vllm.v1.worker.mamba_utils import get_mamba_groups
39 class StepAction:
40 num_computed_tokens_start: int
41 num_scheduled_tokens: int
42 - kv_cache_block_ids: list[int] # [] to follow last step
42 + kv_cache_block_ids: list[int] # per-block mask: 1=held, 0=freed/nulled
43 preprocess_copy_idx: tuple[int, int] # -1, -1 for no copy
44 postprocess_copy_idx: tuple[int, int] # -1, -1 for no copy
45
46
47 num_speculative_tokens = 3
48
49 +# Whether the run under test uses async scheduling. Set by each test entrypoint
50 +# before generation; consulted where the scheduler's optimistic token count must
51 +# be corrected for in-flight (possibly-rejected) speculative tokens.
52 +async_scheduling_mode = False
53 +
54 num_accepted_tokens = 1
55 prompt_token_ids: list[int] = []
56 MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
@@ -249,7 +254,8 @@ def get_fake_execute_model_fn(original_execute_model_fn: Callable):
254 scheduler_output.scheduled_cached_reqs.num_computed_tokens[0]
255 )
256 if (
252 - self.num_spec_tokens
257 + async_scheduling_mode
258 + and self.num_spec_tokens
259 and num_prompt_tokens is not None
260 and num_computed_tokens > num_prompt_tokens
261 ):
@@ -496,7 +502,10 @@ def apply_patch(monkeypatch: pytest.MonkeyPatch):
502 monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn)
503
504
499 -def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
505 +def get_mamba_prefix_cache_step_configs(
506 + async_scheduling: bool = False,
507 +) -> dict[str, TestConfig]:
508 + a = async_scheduling
509 tests = {
510 "accept_1": TestConfig(
511 num_prompt_tokens=554,
@@ -504,14 +513,24 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
513 num_accepted_tokens=1,
514 step_actions=[
515 StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
507 - StepAction(554, 4, [], (-1, -1), (-1, -1)),
508 - StepAction(555, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
509 - StepAction(556, 4, [], (-1, -1), (-1, -1)),
510 - StepAction(557, 4, [], (0, 1), (-1, -1)),
511 - StepAction(558, 4, [], (-1, -1), (-1, -1)),
512 - StepAction(559, 4, [], (-1, -1), (1, 0)),
513 - StepAction(560, 4, [], (-1, -1), (-1, -1)),
514 - StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
516 + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
517 + StepAction(
518 + 555, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
519 + ),
520 + StepAction(
521 + 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
522 + ),
523 + StepAction(557, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
524 + StepAction(558, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
525 + StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
526 + StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
527 + StepAction(
528 + 561,
529 + 4,
530 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
531 + (-1, -1),
532 + (-1, -1),
533 + ),
534 ],
535 ),
536 # test case 2.1: no hit, accept 2 tokens
@@ -521,11 +540,19 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
540 num_accepted_tokens=2,
541 step_actions=[
542 StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
524 - StepAction(554, 4, [], (-1, -1), (-1, -1)),
525 - StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
526 - StepAction(558, 4, [], (1, 1), (2, 0)),
527 - StepAction(560, 4, [], (-1, -1), (-1, -1)),
528 - StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
543 + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
544 + StepAction(
545 + 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
546 + ),
547 + StepAction(558, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)),
548 + StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
549 + StepAction(
550 + 562,
551 + 4,
552 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
553 + (-1, -1),
554 + (-1, -1),
555 + ),
556 ],
557 ),
558 # test case 2.2: no hit, accept 2 tokens
@@ -535,10 +562,16 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
562 num_accepted_tokens=2,
563 step_actions=[
564 StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
538 - StepAction(555, 4, [], (-1, -1), (-1, -1)),
565 + StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
566 StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
540 - StepAction(559, 4, [], (-1, -1), (1, 0)),
541 - StepAction(561, 4, [], (-1, -1), (-1, -1)),
567 + StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
568 + StepAction(
569 + 561,
570 + 4,
571 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
572 + (-1, -1),
573 + (-1, -1),
574 + ),
575 StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
576 ],
577 ),
@@ -548,10 +581,18 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
581 num_accepted_tokens=3,
582 step_actions=[
583 StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
551 - StepAction(553, 4, [], (-1, -1), (-1, -1)),
552 - StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
553 - StepAction(559, 4, [], (2, 1), (1, 0)),
554 - StepAction(562, 4, [], (-1, -1), (-1, -1)),
584 + StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
585 + StepAction(
586 + 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
587 + ),
588 + StepAction(559, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)),
589 + StepAction(
590 + 562,
591 + 4,
592 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
593 + (-1, -1),
594 + (-1, -1),
595 + ),
596 StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
597 ],
598 ),
@@ -561,10 +602,16 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
602 num_accepted_tokens=3,
603 step_actions=[
604 StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
564 - StepAction(554, 4, [], (-1, -1), (-1, -1)),
605 + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
606 StepAction(557, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)),
566 - StepAction(560, 4, [], (-1, -1), (-1, -1)),
567 - StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
607 + StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
608 + StepAction(
609 + 563,
610 + 4,
611 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
612 + (-1, -1),
613 + (-1, -1),
614 + ),
615 ],
616 ),
617 "accept_3_3": TestConfig(
@@ -573,9 +620,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
620 num_accepted_tokens=3,
621 step_actions=[
622 StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
576 - StepAction(555, 4, [], (-1, -1), (-1, -1)),
623 + StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
624 StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
578 - StepAction(561, 4, [], (-1, -1), (-1, -1)),
625 + StepAction(
626 + 561,
627 + 4,
628 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
629 + (-1, -1),
630 + (-1, -1),
631 + ),
632 StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
633 ],
634 ),
@@ -585,9 +638,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
638 num_accepted_tokens=4,
639 step_actions=[
640 StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
588 - StepAction(553, 4, [], (-1, -1), (-1, -1)),
641 + StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
642 StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
590 - StepAction(561, 4, [], (-1, -1), (-1, -1)),
643 + StepAction(
644 + 561,
645 + 4,
646 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
647 + (-1, -1),
648 + (-1, -1),
649 + ),
650 StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
651 ],
652 ),
@@ -597,9 +656,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
656 num_accepted_tokens=4,
657 step_actions=[
658 StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
600 - StepAction(554, 4, [], (-1, -1), (-1, -1)),
659 + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
660 StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
602 - StepAction(562, 4, [], (-1, -1), (-1, -1)),
661 + StepAction(
662 + 562,
663 + 4,
664 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
665 + (-1, -1),
666 + (-1, -1),
667 + ),
668 StepAction(566, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
669 ],
670 ),
@@ -609,9 +674,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
674 num_accepted_tokens=4,
675 step_actions=[
676 StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
612 - StepAction(555, 4, [], (-1, -1), (-1, -1)),
677 + StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
678 StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
614 - StepAction(563, 4, [], (-1, -1), (-1, -1)),
679 + StepAction(
680 + 563,
681 + 4,
682 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
683 + (-1, -1),
684 + (-1, -1),
685 + ),
686 StepAction(567, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
687 ],
688 ),
@@ -621,9 +692,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
692 num_accepted_tokens=4,
693 step_actions=[
694 StepAction(0, 556, [1, 1, 1, 1], (-1, -1), (-1, -1)),
624 - StepAction(556, 4, [], (-1, -1), (3, 0)),
695 + StepAction(556, 4, [1, 1, 1, 1], (-1, -1), (3, 0)),
696 StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
626 - StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
697 + StepAction(
698 + 564,
699 + 4,
700 + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
701 + (-1, -1),
702 + (-1, -1),
703 + ),
704 ],
705 ),
706 "prompt_block_size": TestConfig(
@@ -642,7 +719,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
719 step_actions=[
720 StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
721 StepAction(560, 560, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
645 - StepAction(560 * 2, 4, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
722 + StepAction(
723 + 560 * 2,
724 + 4,
725 + [1, 1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1, 1],
726 + (1, 2),
727 + (-1, -1),
728 + ),
729 ],
730 ),
731 "prompt_2_block_size_10": TestConfig(
@@ -652,7 +735,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
735 step_actions=[
736 StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
737 StepAction(560, 570, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
655 - StepAction(560 * 2 + 10, 4, [0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
738 + StepAction(
739 + 560 * 2 + 10,
740 + 4,
741 + [1, 0, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1],
742 + (-1, -1),
743 + (-1, -1),
744 + ),
745 ],
746 ),
747 "prompt_3_block_size": TestConfig(
@@ -662,7 +751,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
751 step_actions=[
752 StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
753 StepAction(560 * 2, 560, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
665 - StepAction(560 * 3, 4, [0, 0, 1, 1, 1, 1, 1], (2, 3), (-1, -1)),
754 + StepAction(
755 + 560 * 3,
756 + 4,
757 + [0, 1, 1, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1, 1],
758 + (2, 3),
759 + (-1, -1),
760 + ),
761 ],
762 ),
763 "prompt_3_block_size_10": TestConfig(
@@ -672,7 +767,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
767 step_actions=[
768 StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
769 StepAction(560 * 2, 570, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
675 - StepAction(560 * 3 + 10, 4, [0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
770 + StepAction(
771 + 560 * 3 + 10,
772 + 4,
773 + [0, 1, 0, 1, 1, 1, 1] if a else [0, 0, 0, 1, 1, 1, 1],
774 + (-1, -1),
775 + (-1, -1),
776 + ),
777 ],
778 ),
779 "prompt_10_block_size": TestConfig(
@@ -691,14 +792,18 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
792 StepAction(
793 560 * 9,
794 560,
694 - [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
795 + [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1]
796 + if a
797 + else [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
798 (8, 9),
799 (-1, -1),
800 ),
801 StepAction(
802 560 * 10,
803 4,
701 - [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
804 + [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1]
805 + if a
806 + else [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
807 (9, 10),
808 (-1, -1),
809 ),
@@ -720,34 +825,28 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
825 StepAction(
826 560 * 9,
827 560 + 10,
723 - [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1],
828 + [0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1]
829 + if a
830 + else [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1],
831 (8, 10),
832 (-1, -1),
833 ),
834 ],
835 ),
836 }
730 -
837 return tests
838
839
734 -def fill_following_kv_cache_block_ids(test_config: TestConfig) -> None:
735 - for step_action_prev, step_action_next in zip(
736 - test_config.step_actions[:-1], test_config.step_actions[1:]
737 - ):
738 - if len(step_action_next.kv_cache_block_ids) == 0:
739 - step_action_next.kv_cache_block_ids = (
740 - step_action_prev.kv_cache_block_ids.copy()
741 - )
742 -
743 -
744 -@create_new_process_for_each_test()
745 -def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
840 +def _run_mamba_prefix_cache_mrv1(
841 + monkeypatch: pytest.MonkeyPatch, async_scheduling: bool
842 +):
843 + global async_scheduling_mode
844 + async_scheduling_mode = async_scheduling
845 run_ref_mamba_state_in_subprocess()
846 apply_patch(monkeypatch)
847 prompt_dataset = datasets.load_dataset("heheda/a_long_article")
848 full_prompt = prompt_dataset["train"][0]["text"]
750 - tests = get_mamba_prefix_cache_step_configs()
849 + tests = get_mamba_prefix_cache_step_configs(async_scheduling)
850
851 engine = LLM(
852 model=MODEL,
@@ -761,6 +860,7 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
860 },
861 max_num_batched_tokens=3072,
862 hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
863 + async_scheduling=async_scheduling,
864 seed=42,
865 )
866 global prompt_token_ids
@@ -777,7 +877,6 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
877 )
878 global cur_step_action_idx
879 cur_step_action_idx = 0
780 - fill_following_kv_cache_block_ids(test_config)
880 global step_actions
881 step_actions = test_config.step_actions
882 _ = engine.generate(
@@ -800,7 +899,20 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
899
900
901 @create_new_process_for_each_test()
803 -def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch):
902 +def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
903 + _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=False)
904 +
905 +
906 +@create_new_process_for_each_test()
907 +def test_mamba_prefix_cache_mrv1_async(monkeypatch: pytest.MonkeyPatch):
908 + _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=True)
909 +
910 +
911 +def _run_mamba_prefix_cache_mrv2(
912 + monkeypatch: pytest.MonkeyPatch, async_scheduling: bool
913 +):
914 + global async_scheduling_mode
915 + async_scheduling_mode = async_scheduling
916 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
917 monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
918 envs.disable_envs_cache()
@@ -998,11 +1110,12 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch):
1110 max_num_batched_tokens=3072,
1111 max_model_len=BLOCK_SIZE * 12,
1112 hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
1113 + async_scheduling=async_scheduling,
1114 seed=42,
1115 )
1116
1117 try:
1005 - tests = get_mamba_prefix_cache_step_configs()
1118 + tests = get_mamba_prefix_cache_step_configs(async_scheduling)
1119
1120 global step_actions
1121 global cur_step_action_idx
@@ -1010,7 +1123,6 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch):
1123 for test_name, test_config in tests.items():
1124 num_accepted_tokens = test_config.num_accepted_tokens
1125 cur_step_action_idx = 0
1013 - fill_following_kv_cache_block_ids(test_config)
1126 step_actions = test_config.step_actions
1127 sampling_params = SamplingParams(
1128 temperature=0.0,
@@ -1053,3 +1165,13 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch):
1165 del engine
1166 torch.accelerator.empty_cache()
1167 cleanup_dist_env_and_memory()
1168 +
1169 +
1170 +@create_new_process_for_each_test()
1171 +def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch):
1172 + _run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=False)
1173 +
1174 +
1175 +@create_new_process_for_each_test()
1176 +def test_mamba_prefix_cache_mrv2_async(monkeypatch: pytest.MonkeyPatch):
1177 + _run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=True)
vllm/config/vllm.py
+11
@@ -501,6 +501,17 @@ class VllmConfig:
501 return 2
502 return pp_size
503
504 + @property
505 + def max_in_flight_tokens(self) -> int:
506 + # Upper bound on tokens that are scheduled but not yet settled (freed):
507 + # every concurrent batch may hold up to a full `max_num_batched_tokens`.
508 + # Recycling-aware KV cache specs (sliding-window, chunked-local) reserve
509 + # for this because out-of-window blocks are freed on the processed-token
510 + # basis, so in-flight steps transiently keep their blocks.
511 + return (
512 + self.max_concurrent_batches * self.scheduler_config.max_num_batched_tokens
513 + )
514 +
515 @property
516 def num_speculative_tokens(self) -> int:
517 if (
vllm/v1/core/kv_cache_coordinator.py
+16 -16
@@ -67,7 +67,7 @@ class KVCacheCoordinator(ABC):
67 self,
68 kv_cache_config: KVCacheConfig,
69 max_model_len: int,
70 - max_num_batched_tokens: int,
70 + max_in_flight_tokens: int,
71 use_eagle: bool,
72 enable_caching: bool,
73 enable_kv_cache_events: bool,
@@ -107,7 +107,7 @@ class KVCacheCoordinator(ABC):
107 self.single_type_managers = tuple(
108 get_manager_for_kv_cache_spec(
109 kv_cache_spec=kv_cache_group.kv_cache_spec,
110 - max_num_batched_tokens=max_num_batched_tokens,
110 + max_in_flight_tokens=max_in_flight_tokens,
111 max_model_len=max_model_len,
112 block_pool=self.block_pool,
113 enable_caching=enable_caching,
@@ -332,7 +332,7 @@ class KVCacheCoordinator(ABC):
332 def remove_skipped_blocks(
333 self,
334 request_id: str,
335 - total_computed_tokens: int,
335 + processed_computed_tokens: int,
336 num_prompt_tokens: int | None = None,
337 ) -> None:
338 """
@@ -341,15 +341,15 @@ class KVCacheCoordinator(ABC):
341
342 Args:
343 request_id: The request ID.
344 - total_computed_tokens: The total number of computed tokens, including
345 - local computed tokens and external computed tokens.
344 + processed_computed_tokens: Computed-token prefix length covering
345 + fully processed and committed tokens only (safe to free).
346 num_prompt_tokens: Optional prompt length. R-SWA managers use this to
347 free gap blocks between the prefill tail and decode window; other
348 manager types ignore it.
349 """
350 for manager in self.single_type_managers:
351 manager.remove_skipped_blocks(
352 - request_id, total_computed_tokens, num_prompt_tokens
352 + request_id, processed_computed_tokens, num_prompt_tokens
353 )
354
355 def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]:
@@ -387,7 +387,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
387 self,
388 kv_cache_config: KVCacheConfig,
389 max_model_len: int,
390 - max_num_batched_tokens: int,
390 + max_in_flight_tokens: int,
391 use_eagle: bool,
392 enable_kv_cache_events: bool,
393 dcp_world_size: int,
@@ -399,7 +399,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
399 super().__init__(
400 kv_cache_config,
401 max_model_len,
402 - max_num_batched_tokens,
402 + max_in_flight_tokens,
403 use_eagle,
404 False,
405 enable_kv_cache_events,
@@ -436,7 +436,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
436 self,
437 kv_cache_config: KVCacheConfig,
438 max_model_len: int,
439 - max_num_batched_tokens: int,
439 + max_in_flight_tokens: int,
440 use_eagle: bool,
441 enable_caching: bool,
442 enable_kv_cache_events: bool,
@@ -449,7 +449,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
449 super().__init__(
450 kv_cache_config,
451 max_model_len,
452 - max_num_batched_tokens,
452 + max_in_flight_tokens,
453 use_eagle,
454 enable_caching,
455 enable_kv_cache_events,
@@ -522,7 +522,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
522 self,
523 kv_cache_config: KVCacheConfig,
524 max_model_len: int,
525 - max_num_batched_tokens: int,
525 + max_in_flight_tokens: int,
526 use_eagle: bool,
527 enable_caching: bool,
528 enable_kv_cache_events: bool,
@@ -535,7 +535,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
535 super().__init__(
536 kv_cache_config,
537 max_model_len,
538 - max_num_batched_tokens,
538 + max_in_flight_tokens,
539 use_eagle,
540 enable_caching,
541 enable_kv_cache_events,
@@ -783,7 +783,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
783 def get_kv_cache_coordinator(
784 kv_cache_config: KVCacheConfig,
785 max_model_len: int,
786 - max_num_batched_tokens: int,
786 + max_in_flight_tokens: int,
787 use_eagle: bool,
788 enable_caching: bool,
789 enable_kv_cache_events: bool,
@@ -797,7 +797,7 @@ def get_kv_cache_coordinator(
797 return KVCacheCoordinatorNoPrefixCache(
798 kv_cache_config,
799 max_model_len,
800 - max_num_batched_tokens,
800 + max_in_flight_tokens,
801 use_eagle,
802 enable_kv_cache_events,
803 dcp_world_size=dcp_world_size,
@@ -810,7 +810,7 @@ def get_kv_cache_coordinator(
810 return UnitaryKVCacheCoordinator(
811 kv_cache_config,
812 max_model_len,
813 - max_num_batched_tokens,
813 + max_in_flight_tokens,
814 use_eagle,
815 enable_caching,
816 enable_kv_cache_events,
@@ -823,7 +823,7 @@ def get_kv_cache_coordinator(
823 return HybridKVCacheCoordinator(
824 kv_cache_config,
825 max_model_len,
826 - max_num_batched_tokens,
826 + max_in_flight_tokens,
827 use_eagle,
828 enable_caching,
829 enable_kv_cache_events,
vllm/v1/core/kv_cache_manager.py
+12 -9
@@ -114,7 +114,7 @@ class KVCacheManager:
114 max_model_len: int,
115 scheduler_block_size: int,
116 hash_block_size: int,
117 - max_num_batched_tokens: int | None = None,
117 + max_in_flight_tokens: int | None = None,
118 enable_caching: bool = True,
119 use_eagle: bool = False,
120 log_stats: bool = False,
@@ -128,8 +128,8 @@ class KVCacheManager:
128 # When unset, fall back to `max_model_len` so the recycling-aware cap
129 # collapses to the prior (uncapped) admission behavior. The scheduler
130 # always supplies the real value at runtime.
131 - if max_num_batched_tokens is None:
132 - max_num_batched_tokens = max_model_len
131 + if max_in_flight_tokens is None:
132 + max_in_flight_tokens = max_model_len
133
134 self.enable_caching = enable_caching
135 self.use_eagle = use_eagle
@@ -143,7 +143,7 @@ class KVCacheManager:
143 self.coordinator = get_kv_cache_coordinator(
144 kv_cache_config=kv_cache_config,
145 max_model_len=self.max_model_len,
146 - max_num_batched_tokens=max_num_batched_tokens,
146 + max_in_flight_tokens=max_in_flight_tokens,
147 use_eagle=self.use_eagle,
148 enable_caching=self.enable_caching,
149 enable_kv_cache_events=enable_kv_cache_events,
@@ -397,9 +397,12 @@ class KVCacheManager:
397 # insufficient free blocks.
398 # Should call this function before allocating new blocks to reduce
399 # the number of evicted blocks.
400 + # Free on the processed-token basis: in-flight steps' attention windows
401 + # still read blocks below the optimistic boundary, and rejected spec
402 + # tokens can roll it back.
403 self.coordinator.remove_skipped_blocks(
404 request.request_id,
402 - total_computed_tokens,
405 + max(0, total_computed_tokens - request.num_in_flight_tokens),
406 num_prompt_tokens=request.num_prompt_tokens,
407 )
408
@@ -472,7 +475,7 @@ class KVCacheManager:
475 def remove_skipped_blocks(
476 self,
477 request_id: str,
475 - total_computed_tokens: int,
478 + processed_computed_tokens: int,
479 num_prompt_tokens: int | None = None,
480 ) -> None:
481 """Remove the blocks that are no longer needed from `blocks` and replace
@@ -480,12 +483,12 @@ class KVCacheManager:
483
484 Args:
485 request_id: The request ID.
483 - total_computed_tokens: The total number of computed tokens, including
484 - local computed tokens and external computed tokens.
486 + processed_computed_tokens: Computed-token prefix length covering
487 + fully processed and committed tokens only (safe to free).
488 num_prompt_tokens: Optional prompt length for R-SWA gap eviction.
489 """
490 self.coordinator.remove_skipped_blocks(
488 - request_id, total_computed_tokens, num_prompt_tokens
491 + request_id, processed_computed_tokens, num_prompt_tokens
492 )
493
494 def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]:
vllm/v1/core/sched/scheduler.py
+9 -4
@@ -262,7 +262,7 @@ class Scheduler(SchedulerInterface):
262 self.kv_cache_manager = KVCacheManager(
263 kv_cache_config=kv_cache_config,
264 max_model_len=self.max_model_len,
265 - max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
265 + max_in_flight_tokens=vllm_config.max_in_flight_tokens,
266 enable_caching=self.cache_config.enable_prefix_caching,
267 use_eagle=self.use_eagle,
268 log_stats=self.log_stats,
@@ -1180,6 +1180,7 @@ class Scheduler(SchedulerInterface):
1180 for req_id, num_scheduled_token in num_scheduled_tokens.items():
1181 request = self.requests[req_id]
1182 request.num_computed_tokens += num_scheduled_token
1183 + request.num_in_flight_tokens += num_scheduled_token
1184 if self.defer_block_free:
1185 # Record the in-flight step, to fence deferred block freeing.
1186 request.last_sched_seq = self.sched_step_seq
@@ -1566,10 +1567,12 @@ class Scheduler(SchedulerInterface):
1567 stopped_preempted_reqs: set[Request] = set()
1568 for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
1569 assert num_tokens_scheduled > 0
1570 + request = self.requests.get(req_id)
1571 + if request is not None:
1572 + request.num_in_flight_tokens -= num_tokens_scheduled
1573 if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
1574 # skip failed or rescheduled requests from KV load failure
1575 continue
1572 - request = self.requests.get(req_id)
1576 if request is None or request.is_finished():
1577 # The request is already finished. This can happen if the
1578 # request is aborted while the model is executing it (e.g.,
@@ -2367,10 +2370,12 @@ class Scheduler(SchedulerInterface):
2370 return False, None
2371
2372 # Free any out-of-window prefix blocks before we hand the block table to
2370 - # the connector.
2373 + # the connector, on the processed-token basis (see `allocate_slots`).
2374 self.kv_cache_manager.remove_skipped_blocks(
2375 request_id=request.request_id,
2373 - total_computed_tokens=request.num_computed_tokens,
2376 + processed_computed_tokens=max(
2377 + 0, request.num_computed_tokens - request.num_in_flight_tokens
2378 + ),
2379 num_prompt_tokens=request.num_prompt_tokens,
2380 )
2381
vllm/v1/core/single_type_kv_cache_manager.py
+17 -21
@@ -508,7 +508,7 @@ class SingleTypeKVCacheManager(ABC):
508 def remove_skipped_blocks(
509 self,
510 request_id: str,
511 - total_computed_tokens: int,
511 + processed_computed_tokens: int,
512 num_prompt_tokens: int | None = None,
513 ) -> None:
514 """
@@ -520,15 +520,15 @@ class SingleTypeKVCacheManager(ABC):
520
521 Args:
522 request_id: The request ID.
523 - total_computed_tokens: The total number of computed tokens, including
524 - local computed tokens and external computed tokens.
523 + processed_computed_tokens: Computed-token prefix length covering
524 + fully processed and committed tokens only (safe to free).
525 num_prompt_tokens: Optional prompt length for attention types (e.g.
526 R-SWA) that evict a middle gap rather than a head prefix. Ignored
527 by the default implementation.
528 """
529 del num_prompt_tokens
530 # Remove the blocks that will be skipped during attention computation.
531 - num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens)
531 + num_skipped_tokens = self.get_num_skipped_tokens(processed_computed_tokens)
532 if num_skipped_tokens <= 0:
533 # This indicates that ALL tokens are inside attention window.
534 # Thus we do not need to free any blocks outside attention window.
@@ -639,14 +639,14 @@ class RSWAManager(FullAttentionManager):
639 def remove_skipped_blocks(
640 self,
641 request_id: str,
642 - total_computed_tokens: int,
642 + processed_computed_tokens: int,
643 num_prompt_tokens: int | None = None,
644 ) -> None:
645 """Free gap blocks that are no longer needed for attention.
646
647 Gap = blocks entirely within
648 [ceil(prefix_len / block_size) * block_size,
649 - max(prefix_len, total_computed_tokens - rswa_window))
649 + max(prefix_len, processed_computed_tokens - rswa_window))
650
651 Freed blocks are replaced with null_block in req_to_blocks so the
652 block_table passed to FA4 is valid (null_block KV is all-zero;
@@ -654,7 +654,7 @@ class RSWAManager(FullAttentionManager):
654 """
655 if num_prompt_tokens is None:
656 super().remove_skipped_blocks(
657 - request_id, total_computed_tokens, num_prompt_tokens
657 + request_id, processed_computed_tokens, num_prompt_tokens
658 )
659 return
660
@@ -662,7 +662,9 @@ class RSWAManager(FullAttentionManager):
662 # First block fully after the prefill boundary.
663 first_gap_block = cdiv(num_prompt_tokens, bs)
664 # Decode window start position; blocks before this are evictable.
665 - window_start = max(num_prompt_tokens, total_computed_tokens - self.rswa_window)
665 + window_start = max(
666 + num_prompt_tokens, processed_computed_tokens - self.rswa_window
667 + )
668 last_gap_block = window_start // bs # exclusive upper bound
669 self._remove_blocks_in_range(request_id, first_gap_block, last_gap_block)
670
@@ -1144,20 +1146,13 @@ class MambaManager(SingleTypeKVCacheManager):
1146 def remove_skipped_blocks(
1147 self,
1148 request_id: str,
1147 - num_computed_tokens: int,
1149 + processed_computed_tokens: int,
1150 num_prompt_tokens: int | None = None,
1151 ) -> None:
1152 assert isinstance(self.kv_cache_spec, MambaSpec)
1153
1152 - # NOTE (tdoublep) with async scheduling, the num_computed_tokens can contain
1153 - # draft tokens from the previous step that may or may not be rejected later.
1154 - # This can make us think we are further ahead in the sequence than we actually
1155 - # are, so let's assume that all tokens are rejected so we don't free blocks
1156 - # that we might actually need.
1157 - num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks)
1158 -
1154 super().remove_skipped_blocks(
1160 - request_id, num_computed_tokens, num_prompt_tokens
1155 + request_id, processed_computed_tokens, num_prompt_tokens
1156 )
1157 if self.mamba_cache_mode == "align":
1158 # `last_state_block_idx` refers to the block index allocated two steps ago.
@@ -1171,7 +1166,7 @@ class MambaManager(SingleTypeKVCacheManager):
1166 if (
1167 last_state_block_idx is not None
1168 and last_state_block_idx
1174 - < cdiv(num_computed_tokens, self.block_size) - 1
1169 + < cdiv(processed_computed_tokens, self.block_size) - 1
1170 ):
1171 blocks = self.req_to_blocks[request_id]
1172 if blocks[last_state_block_idx] != self._null_block:
@@ -1455,7 +1450,7 @@ class SinkFullAttentionManager(FullAttentionManager):
1450
1451 def get_manager_for_kv_cache_spec(
1452 kv_cache_spec: KVCacheSpec,
1458 - max_num_batched_tokens: int,
1453 + max_in_flight_tokens: int,
1454 max_model_len: int,
1455 **kwargs,
1456 ) -> SingleTypeKVCacheManager:
@@ -1468,7 +1463,8 @@ def get_manager_for_kv_cache_spec(
1463
1464 Args:
1465 kv_cache_spec: The KVCacheSpec instance
1471 - max_num_batched_tokens: The maximum number of tokens in a batch
1466 + max_in_flight_tokens: The max tokens scheduled but not yet settled
1467 + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`
1468 max_model_len: The maximum context length the model could serve
1469 Returns:
1470 An instance of the appropriate SingleTypeKVCacheManager subclass
@@ -1489,7 +1485,7 @@ def get_manager_for_kv_cache_spec(
1485 ):
1486 kwargs["max_admission_blocks_per_request"] = (
1487 kv_cache_spec.max_admission_blocks_per_request(
1492 - max_num_batched_tokens=max_num_batched_tokens,
1488 + max_in_flight_tokens=max_in_flight_tokens,
1489 max_model_len=max_model_len,
1490 )
1491 )
vllm/v1/kv_cache_interface.py
+18 -15
@@ -482,25 +482,28 @@ class ChunkedLocalAttentionSpec(AttentionSpec):
482 attention_chunk_size: int
483
484 def max_admission_blocks_per_request(
485 - self, max_num_batched_tokens: int, max_model_len: int
485 + self, max_in_flight_tokens: int, max_model_len: int
486 ) -> int:
487 """Per-request admission cap, in blocks.
488
489 Single source of truth for both startup pool sizing
490 (`max_memory_usage_bytes`) and the runtime admission gate, so requests
491 admitted by startup can also be admitted at runtime.
492 +
493 + `max_in_flight_tokens` is the max tokens scheduled but not yet settled
494 + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`.
495 """
493 - # During chunked prefill, we hold KV for at most one chunk window.
496 + # During chunked prefill, we hold KV for at most one chunk window plus
497 + # the in-flight tokens, since frees happen on the processed-token basis.
498 num_tokens = min(
495 - self.attention_chunk_size + max_num_batched_tokens, max_model_len
499 + self.attention_chunk_size + max_in_flight_tokens, max_model_len
500 )
501 return cdiv(num_tokens, self.block_size)
502
503 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
500 - max_model_len = vllm_config.model_config.max_model_len
501 - max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
504 max_blocks = self.max_admission_blocks_per_request(
503 - max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
505 + max_in_flight_tokens=vllm_config.max_in_flight_tokens,
506 + max_model_len=vllm_config.model_config.max_model_len,
507 )
508 return max_blocks * self.page_size_bytes
509
@@ -544,7 +547,7 @@ class SlidingWindowSpec(AttentionSpec):
547 )
548
549 def max_admission_blocks_per_request(
547 - self, max_num_batched_tokens: int, max_model_len: int
550 + self, max_in_flight_tokens: int, max_model_len: int
551 ) -> int:
552 """Per-request admission cap, in blocks.
553
@@ -553,13 +556,14 @@ class SlidingWindowSpec(AttentionSpec):
556 real-held blocks plateau at this bound because
557 `SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots`
558 before each chunk's `get_num_blocks_to_allocate`.
559 +
560 + `max_in_flight_tokens` is the max tokens scheduled but not yet settled
561 + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`.
562 """
563 # During chunked prefill, we hold KV for the last `sliding_window-1`
558 - # computed tokens plus the newly scheduled tokens, and never more
559 - # than `max_model_len`.
560 - num_tokens = min(
561 - self.sliding_window - 1 + max_num_batched_tokens, max_model_len
562 - )
564 + # computed tokens plus the in-flight tokens (frees happen on the
565 + # processed-token basis); never more than `max_model_len`.
566 + num_tokens = min(self.sliding_window - 1 + max_in_flight_tokens, max_model_len)
567 # +1 because the sliding window may not start from the beginning of
568 # the block. E.g. block size 4 and num_token 4 needs two blocks
569 # [XXCD][EF] to store the 6-token window [CDEF].
@@ -569,10 +573,9 @@ class SlidingWindowSpec(AttentionSpec):
573 assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
574 "DCP not support sliding window."
575 )
572 - max_model_len = vllm_config.model_config.max_model_len
573 - max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
576 max_blocks = self.max_admission_blocks_per_request(
575 - max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
577 + max_in_flight_tokens=vllm_config.max_in_flight_tokens,
578 + max_model_len=vllm_config.model_config.max_model_len,
579 )
580 return max_blocks * self.page_size_bytes
581
vllm/v1/request.py
+5
@@ -141,6 +141,11 @@ class Request:
141 self.num_output_placeholders = 0
142 self.async_tokens_to_discard = 0
143
144 + # Tokens of steps whose output is not yet processed (async scheduling
145 + # and PP run ahead of the GPU); `num_computed_tokens` counts them
146 + # optimistically.
147 + self.num_in_flight_tokens = 0
148 +
149 # V2+PP+async: Enforces `pp_size` cadence between same-request decode steps
150 # so the worker's broadcast slot ring stays consistent.
151 self.next_decode_eligible_step = 0
vllm/v1/simple_kv_offload/manager.py
+1 -3
@@ -122,9 +122,7 @@ class SimpleCPUOffloadScheduler:
122 self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
123 kv_cache_config=self.cpu_kv_cache_config,
124 max_model_len=vllm_config.model_config.max_model_len,
125 - max_num_batched_tokens=(
126 - vllm_config.scheduler_config.max_num_batched_tokens
127 - ),
125 + max_in_flight_tokens=vllm_config.max_in_flight_tokens,
126 use_eagle=False,
127 enable_caching=True,
128 enable_kv_cache_events=self.enable_kv_cache_events,