From cd0de48d0883ecb8e1ef350a99baa0c158f58e82 Mon Sep 17 00:00:00 2001 From: Saddss <108515797+Saddss@users.noreply.github.com> Date: Wed, 8 Jul 2026 20:34:19 +0800 Subject: [PATCH] [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 Signed-off-by: Yifan Qiao Co-authored-by: Yifan Qiao Co-authored-by: Saddss <28726669061@qq.com> Co-authored-by: Nick Hill --- .../generation/test_voxtral_realtime.py | 2 +- tests/v1/core/test_kv_cache_utils.py | 3 + tests/v1/core/test_prefix_caching.py | 6 +- .../v1/core/test_swa_inflight_window_free.py | 259 ++++++++++++++++++ tests/v1/core/utils.py | 21 +- .../v1/e2e/general/test_mamba_prefix_cache.py | 244 ++++++++++++----- vllm/config/vllm.py | 11 + vllm/v1/core/kv_cache_coordinator.py | 32 +-- vllm/v1/core/kv_cache_manager.py | 21 +- vllm/v1/core/sched/scheduler.py | 13 +- vllm/v1/core/single_type_kv_cache_manager.py | 38 ++- vllm/v1/kv_cache_interface.py | 33 ++- vllm/v1/request.py | 5 + vllm/v1/simple_kv_offload/manager.py | 4 +- 14 files changed, 549 insertions(+), 143 deletions(-) create mode 100644 tests/v1/core/test_swa_inflight_window_free.py diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 4d37b36c36..df2d63ad84 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -81,7 +81,7 @@ def assert_encoder_kv_cache_spec(engine: LLM) -> None: assert spec.sliding_window == cdiv(750, 4) + 1 == 189 assert ( spec.max_admission_blocks_per_request( - max_num_batched_tokens=1, + max_in_flight_tokens=1, max_model_len=vllm_config.model_config.max_model_len, ) == 13 diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 5b8b43701d..947672c48a 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1404,12 +1404,15 @@ def test_get_max_concurrency_for_kv_cache_config(): enable_chunked_prefill=True, max_model_len=model_config.max_model_len, is_encoder_decoder=model_config.is_encoder_decoder, + # Pin to sync: SWA per-request bounds grow with overlapping batches. + async_scheduling=False, ) vllm_config = VllmConfig( model_config=model_config, scheduler_config=scheduler_config, ) + assert vllm_config.max_concurrent_batches == 1 full_attention_spec = FullAttentionSpec( block_size=16, diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 6246a23329..59260a499e 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -3453,7 +3453,8 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt(): manager = make_kv_cache_manager( config, max_model_len=max_model_len, - max_num_batched_tokens=max_num_batched_tokens, + # Single (sync) batch in flight, so in-flight tokens == batched tokens. + max_in_flight_tokens=max_num_batched_tokens, enable_caching=True, hash_block_size=block_size, ) @@ -3509,7 +3510,8 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized(): manager = make_kv_cache_manager( config, max_model_len=max_model_len, - max_num_batched_tokens=max_num_batched_tokens, + # Single (sync) batch in flight, so in-flight tokens == batched tokens. + max_in_flight_tokens=max_num_batched_tokens, enable_caching=True, hash_block_size=block_size, ) diff --git a/tests/v1/core/test_swa_inflight_window_free.py b/tests/v1/core/test_swa_inflight_window_free.py new file mode 100644 index 0000000000..737796a4bf --- /dev/null +++ b/tests/v1/core/test_swa_inflight_window_free.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Out-of-window block frees vs in-flight GPU steps. + +With async scheduling / PP, `num_computed_tokens` optimistically includes +tokens of unprocessed steps whose attention windows still read the blocks +just below the optimistic boundary (and rejected spec tokens can roll it +back), so `allocate_slots` frees on the processed-token basis: +`num_computed_tokens - num_in_flight_tokens`. +""" + +import torch + +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec +from vllm.v1.outputs import ModelRunnerOutput + +from .utils import create_requests, create_scheduler, mock_kv + +NUM_PROMPT_TOKENS = 100 +BLOCK_SIZE = 16 +SLIDING_WINDOW = 16 +# Tokens 0..84 are outside the window of the next token to compute +# (100 - 16 + 1 = 85), i.e. 5 full blocks. +NUM_OUT_OF_WINDOW_BLOCKS = 85 // BLOCK_SIZE +# Chunked-local skips whole chunks left of the current one: +# (100 // 32) * 32 = 96 settled tokens -> 6 full blocks. +CHUNK_SIZE = 32 +NUM_OUT_OF_CHUNK_BLOCKS = (NUM_PROMPT_TOKENS // CHUNK_SIZE) * CHUNK_SIZE // BLOCK_SIZE + + +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_swa_scheduler(async_scheduling: bool): + return create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=async_scheduling, + kv_cache_spec=SlidingWindowSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=SLIDING_WINDOW, + ), + ) + + +def _create_chunked_scheduler(async_scheduling: bool): + return create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=async_scheduling, + kv_cache_spec=ChunkedLocalAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + attention_chunk_size=CHUNK_SIZE, + ), + ) + + +def _num_null_blocks(scheduler, request_id: str) -> int: + manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0] + null_block = manager._null_block + return sum(1 for b in manager.req_to_blocks[request_id] if b is null_block) + + +def test_num_in_flight_tokens_accounting(): + scheduler = create_scheduler(async_scheduling=True) + request = create_requests(num_requests=1, num_tokens=NUM_PROMPT_TOKENS)[0] + scheduler.add_request(request) + + out0 = scheduler.schedule() + assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS + # Async: decode scheduled before the prefill output is processed. + out1 = scheduler.schedule() + assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS + 1 + + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert request.num_in_flight_tokens == 1 + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + assert request.num_in_flight_tokens == 0 + + +def test_swa_free_waits_for_in_flight_step(): + """Async: out-of-window blocks stay allocated until the step that still + reads them has been processed.""" + scheduler = _create_swa_scheduler(async_scheduling=True) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + block_pool = scheduler.kv_cache_manager.block_pool + + out0 = scheduler.schedule() # prefill, in flight from here on + free_after_prefill = block_pool.get_num_free_blocks() + + # Decode scheduled while the prefill still reads the out-of-window blocks: + # they must not be freed yet. + out1 = scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == 0 + assert block_pool.get_num_free_blocks() == free_after_prefill + + # Prefill output processed; the next allocate frees the out-of-window + # blocks. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS + assert ( + block_pool.get_num_free_blocks() + == free_after_prefill + NUM_OUT_OF_WINDOW_BLOCKS + ) + # Not double-freed on the following steps. + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS + + +def test_swa_free_immediate_when_sync(): + """Sync: no in-flight step at schedule time, frees happen at the first + decode allocation as before.""" + scheduler = _create_swa_scheduler(async_scheduling=False) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + + out0 = scheduler.schedule() + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert request.num_in_flight_tokens == 0 + + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS + + +def test_swa_admission_cap_accounts_for_overlapping_batches(): + spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=1024, + ) + base = spec.max_admission_blocks_per_request( + max_in_flight_tokens=1024, max_model_len=16384 + ) + # (1024 - 1 + 1024) tokens -> 128 blocks, +1 for window misalignment. + assert base == 129 + overlapped = spec.max_admission_blocks_per_request( + max_in_flight_tokens=2 * 1024, max_model_len=16384 + ) + # One extra in-flight chunk is held back: (1024 - 1 + 2 * 1024) tokens. + assert overlapped == 193 + + +def test_chunked_local_free_waits_for_in_flight_step(): + """Chunked-local attention frees whole chunks left of the current one, and + is exposed to the same load-WAR: with async scheduling those chunks must + stay allocated until the in-flight step that still reads them settles.""" + scheduler = _create_chunked_scheduler(async_scheduling=True) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + block_pool = scheduler.kv_cache_manager.block_pool + + out0 = scheduler.schedule() # prefill, in flight from here on + free_after_prefill = block_pool.get_num_free_blocks() + + # Decode scheduled while the prefill still reads the out-of-chunk blocks. + out1 = scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == 0 + assert block_pool.get_num_free_blocks() == free_after_prefill + + # Prefill output processed; the next allocate frees the out-of-chunk blocks. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS + assert ( + block_pool.get_num_free_blocks() == free_after_prefill + NUM_OUT_OF_CHUNK_BLOCKS + ) + # Not double-freed on the following steps. + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS + + +def test_chunked_local_admission_cap_accounts_for_overlapping_batches(): + spec = ChunkedLocalAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + attention_chunk_size=1024, + ) + base = spec.max_admission_blocks_per_request( + max_in_flight_tokens=1024, max_model_len=16384 + ) + # (1024 + 1024) tokens -> 128 blocks. + assert base == 128 + overlapped = spec.max_admission_blocks_per_request( + max_in_flight_tokens=2 * 1024, max_model_len=16384 + ) + # One extra in-flight chunk is held back: (1024 + 2 * 1024) tokens. + assert overlapped == 192 + + +def test_connector_finish_frees_on_settled_basis(): + """The out-of-window prune done at request finish, before the block table + is handed to a KV connector (simple CPU offload / NIXL store), must use the + same processed-token basis. Otherwise the connector reads/hands off a block + the still-in-flight step is optimistically counted as done with, which is + the load-path WAR this fix closes.""" + scheduler = create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=True, + use_kv_connector=mock_kv(matched_tokens=0, is_async=False), + kv_cache_spec=SlidingWindowSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=SLIDING_WINDOW, + ), + ) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + + out0 = scheduler.schedule() # prefill, in flight from here on + scheduler.schedule() # decode over-scheduled: num_computed_tokens optimistic + + # Finishing now (connector store) must NOT prune out-of-window blocks the + # in-flight prefill still reads. + scheduler._connector_finished(request) + assert _num_null_blocks(scheduler, req_id) == 0 + + # Once the in-flight step settles, the same prune releases them. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler._connector_finished(request) + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index d347a5e2d4..19beba1a53 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -30,6 +30,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheSpec, ) from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -63,6 +64,7 @@ def create_scheduler( use_ec_connector: bool = False, ec_role: str | None = None, use_v2_model_runner: bool | None = None, + kv_cache_spec: KVCacheSpec | None = None, ) -> Scheduler | AsyncScheduler: """Create scheduler under test. @@ -164,20 +166,17 @@ def create_scheduler( speculative_config=speculative_config, ec_transfer_config=ec_transfer_config, ) + if kv_cache_spec is None: + kv_cache_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ) kv_cache_config = KVCacheConfig( num_blocks=num_blocks, # A large number of blocks to hold all requests kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer"], - FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ) - ], + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], ) cache_config.num_gpu_blocks = num_blocks register_all_kvcache_specs(vllm_config) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index dd9efc6696..5a7af6f22c 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -39,13 +39,18 @@ from vllm.v1.worker.mamba_utils import get_mamba_groups class StepAction: num_computed_tokens_start: int num_scheduled_tokens: int - kv_cache_block_ids: list[int] # [] to follow last step + kv_cache_block_ids: list[int] # per-block mask: 1=held, 0=freed/nulled preprocess_copy_idx: tuple[int, int] # -1, -1 for no copy postprocess_copy_idx: tuple[int, int] # -1, -1 for no copy num_speculative_tokens = 3 +# Whether the run under test uses async scheduling. Set by each test entrypoint +# before generation; consulted where the scheduler's optimistic token count must +# be corrected for in-flight (possibly-rejected) speculative tokens. +async_scheduling_mode = False + num_accepted_tokens = 1 prompt_token_ids: list[int] = [] MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8" @@ -249,7 +254,8 @@ def get_fake_execute_model_fn(original_execute_model_fn: Callable): scheduler_output.scheduled_cached_reqs.num_computed_tokens[0] ) if ( - self.num_spec_tokens + async_scheduling_mode + and self.num_spec_tokens and num_prompt_tokens is not None and num_computed_tokens > num_prompt_tokens ): @@ -496,7 +502,10 @@ def apply_patch(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn) -def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: +def get_mamba_prefix_cache_step_configs( + async_scheduling: bool = False, +) -> dict[str, TestConfig]: + a = async_scheduling tests = { "accept_1": TestConfig( num_prompt_tokens=554, @@ -504,14 +513,24 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=1, step_actions=[ StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [], (-1, -1), (-1, -1)), - StepAction(555, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(556, 4, [], (-1, -1), (-1, -1)), - StepAction(557, 4, [], (0, 1), (-1, -1)), - StepAction(558, 4, [], (-1, -1), (-1, -1)), - StepAction(559, 4, [], (-1, -1), (1, 0)), - StepAction(560, 4, [], (-1, -1), (-1, -1)), - StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 555, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + ), + StepAction( + 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + ), + StepAction(557, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), + StepAction(558, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)), + StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 561, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), # test case 2.1: no hit, accept 2 tokens @@ -521,11 +540,19 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=2, step_actions=[ StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [], (-1, -1), (-1, -1)), - StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(558, 4, [], (1, 1), (2, 0)), - StepAction(560, 4, [], (-1, -1), (-1, -1)), - StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + ), + StepAction(558, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)), + StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 562, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), # test case 2.2: no hit, accept 2 tokens @@ -535,10 +562,16 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=2, step_actions=[ StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(555, 4, [], (-1, -1), (-1, -1)), + StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)), - StepAction(559, 4, [], (-1, -1), (1, 0)), - StepAction(561, 4, [], (-1, -1), (-1, -1)), + StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)), + StepAction( + 561, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -548,10 +581,18 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=3, step_actions=[ StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(553, 4, [], (-1, -1), (-1, -1)), - StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(559, 4, [], (2, 1), (1, 0)), - StepAction(562, 4, [], (-1, -1), (-1, -1)), + StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + ), + StepAction(559, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)), + StepAction( + 562, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -561,10 +602,16 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=3, step_actions=[ StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [], (-1, -1), (-1, -1)), + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(557, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)), - StepAction(560, 4, [], (-1, -1), (-1, -1)), - StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 563, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), "accept_3_3": TestConfig( @@ -573,9 +620,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=3, step_actions=[ StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(555, 4, [], (-1, -1), (-1, -1)), + StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)), - StepAction(561, 4, [], (-1, -1), (-1, -1)), + StepAction( + 561, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -585,9 +638,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=4, step_actions=[ StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(553, 4, [], (-1, -1), (-1, -1)), + StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)), - StepAction(561, 4, [], (-1, -1), (-1, -1)), + StepAction( + 561, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -597,9 +656,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=4, step_actions=[ StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [], (-1, -1), (-1, -1)), + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)), - StepAction(562, 4, [], (-1, -1), (-1, -1)), + StepAction( + 562, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(566, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -609,9 +674,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=4, step_actions=[ StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(555, 4, [], (-1, -1), (-1, -1)), + StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)), - StepAction(563, 4, [], (-1, -1), (-1, -1)), + StepAction( + 563, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(567, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -621,9 +692,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=4, step_actions=[ StepAction(0, 556, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(556, 4, [], (-1, -1), (3, 0)), + StepAction(556, 4, [1, 1, 1, 1], (-1, -1), (3, 0)), StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), - StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 564, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), "prompt_block_size": TestConfig( @@ -642,7 +719,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: step_actions=[ StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(560, 560, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), - StepAction(560 * 2, 4, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)), + StepAction( + 560 * 2, + 4, + [1, 1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1, 1], + (1, 2), + (-1, -1), + ), ], ), "prompt_2_block_size_10": TestConfig( @@ -652,7 +735,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: step_actions=[ StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(560, 570, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)), - StepAction(560 * 2 + 10, 4, [0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 560 * 2 + 10, + 4, + [1, 0, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), "prompt_3_block_size": TestConfig( @@ -662,7 +751,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: step_actions=[ StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(560 * 2, 560, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)), - StepAction(560 * 3, 4, [0, 0, 1, 1, 1, 1, 1], (2, 3), (-1, -1)), + StepAction( + 560 * 3, + 4, + [0, 1, 1, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1, 1], + (2, 3), + (-1, -1), + ), ], ), "prompt_3_block_size_10": TestConfig( @@ -672,7 +767,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: step_actions=[ StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(560 * 2, 570, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)), - StepAction(560 * 3 + 10, 4, [0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 560 * 3 + 10, + 4, + [0, 1, 0, 1, 1, 1, 1] if a else [0, 0, 0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), "prompt_10_block_size": TestConfig( @@ -691,14 +792,18 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: StepAction( 560 * 9, 560, - [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1] + if a + else [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], (8, 9), (-1, -1), ), StepAction( 560 * 10, 4, - [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1] + if a + else [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], (9, 10), (-1, -1), ), @@ -720,34 +825,28 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: StepAction( 560 * 9, 560 + 10, - [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1], + [0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1] + if a + else [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1], (8, 10), (-1, -1), ), ], ), } - return tests -def fill_following_kv_cache_block_ids(test_config: TestConfig) -> None: - for step_action_prev, step_action_next in zip( - test_config.step_actions[:-1], test_config.step_actions[1:] - ): - if len(step_action_next.kv_cache_block_ids) == 0: - step_action_next.kv_cache_block_ids = ( - step_action_prev.kv_cache_block_ids.copy() - ) - - -@create_new_process_for_each_test() -def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): +def _run_mamba_prefix_cache_mrv1( + monkeypatch: pytest.MonkeyPatch, async_scheduling: bool +): + global async_scheduling_mode + async_scheduling_mode = async_scheduling run_ref_mamba_state_in_subprocess() apply_patch(monkeypatch) prompt_dataset = datasets.load_dataset("heheda/a_long_article") full_prompt = prompt_dataset["train"][0]["text"] - tests = get_mamba_prefix_cache_step_configs() + tests = get_mamba_prefix_cache_step_configs(async_scheduling) engine = LLM( model=MODEL, @@ -761,6 +860,7 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): }, max_num_batched_tokens=3072, hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS}, + async_scheduling=async_scheduling, seed=42, ) global prompt_token_ids @@ -777,7 +877,6 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): ) global cur_step_action_idx cur_step_action_idx = 0 - fill_following_kv_cache_block_ids(test_config) global step_actions step_actions = test_config.step_actions _ = engine.generate( @@ -800,7 +899,20 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): @create_new_process_for_each_test() -def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): +def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): + _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=False) + + +@create_new_process_for_each_test() +def test_mamba_prefix_cache_mrv1_async(monkeypatch: pytest.MonkeyPatch): + _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=True) + + +def _run_mamba_prefix_cache_mrv2( + monkeypatch: pytest.MonkeyPatch, async_scheduling: bool +): + global async_scheduling_mode + async_scheduling_mode = async_scheduling monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") envs.disable_envs_cache() @@ -998,11 +1110,12 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): max_num_batched_tokens=3072, max_model_len=BLOCK_SIZE * 12, hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS}, + async_scheduling=async_scheduling, seed=42, ) try: - tests = get_mamba_prefix_cache_step_configs() + tests = get_mamba_prefix_cache_step_configs(async_scheduling) global step_actions global cur_step_action_idx @@ -1010,7 +1123,6 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): for test_name, test_config in tests.items(): num_accepted_tokens = test_config.num_accepted_tokens cur_step_action_idx = 0 - fill_following_kv_cache_block_ids(test_config) step_actions = test_config.step_actions sampling_params = SamplingParams( temperature=0.0, @@ -1053,3 +1165,13 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): del engine torch.accelerator.empty_cache() cleanup_dist_env_and_memory() + + +@create_new_process_for_each_test() +def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): + _run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=False) + + +@create_new_process_for_each_test() +def test_mamba_prefix_cache_mrv2_async(monkeypatch: pytest.MonkeyPatch): + _run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=True) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a1fe109289..dadf7e35f1 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -501,6 +501,17 @@ class VllmConfig: return 2 return pp_size + @property + def max_in_flight_tokens(self) -> int: + # Upper bound on tokens that are scheduled but not yet settled (freed): + # every concurrent batch may hold up to a full `max_num_batched_tokens`. + # Recycling-aware KV cache specs (sliding-window, chunked-local) reserve + # for this because out-of-window blocks are freed on the processed-token + # basis, so in-flight steps transiently keep their blocks. + return ( + self.max_concurrent_batches * self.scheduler_config.max_num_batched_tokens + ) + @property def num_speculative_tokens(self) -> int: if ( diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 5829e553f9..4756136f03 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -67,7 +67,7 @@ class KVCacheCoordinator(ABC): self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -107,7 +107,7 @@ class KVCacheCoordinator(ABC): self.single_type_managers = tuple( get_manager_for_kv_cache_spec( kv_cache_spec=kv_cache_group.kv_cache_spec, - max_num_batched_tokens=max_num_batched_tokens, + max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len, block_pool=self.block_pool, enable_caching=enable_caching, @@ -332,7 +332,7 @@ class KVCacheCoordinator(ABC): def remove_skipped_blocks( self, request_id: str, - total_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """ @@ -341,15 +341,15 @@ class KVCacheCoordinator(ABC): Args: request_id: The request ID. - total_computed_tokens: The total number of computed tokens, including - local computed tokens and external computed tokens. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). num_prompt_tokens: Optional prompt length. R-SWA managers use this to free gap blocks between the prefill tail and decode window; other manager types ignore it. """ for manager in self.single_type_managers: manager.remove_skipped_blocks( - request_id, total_computed_tokens, num_prompt_tokens + request_id, processed_computed_tokens, num_prompt_tokens ) def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]: @@ -387,7 +387,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator): self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_kv_cache_events: bool, dcp_world_size: int, @@ -399,7 +399,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator): super().__init__( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, False, enable_kv_cache_events, @@ -436,7 +436,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -449,7 +449,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): super().__init__( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -522,7 +522,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -535,7 +535,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): super().__init__( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -783,7 +783,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): def get_kv_cache_coordinator( kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -797,7 +797,7 @@ def get_kv_cache_coordinator( return KVCacheCoordinatorNoPrefixCache( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_kv_cache_events, dcp_world_size=dcp_world_size, @@ -810,7 +810,7 @@ def get_kv_cache_coordinator( return UnitaryKVCacheCoordinator( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -823,7 +823,7 @@ def get_kv_cache_coordinator( return HybridKVCacheCoordinator( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 57cd1490e8..89061c6bc1 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -114,7 +114,7 @@ class KVCacheManager: max_model_len: int, scheduler_block_size: int, hash_block_size: int, - max_num_batched_tokens: int | None = None, + max_in_flight_tokens: int | None = None, enable_caching: bool = True, use_eagle: bool = False, log_stats: bool = False, @@ -128,8 +128,8 @@ class KVCacheManager: # When unset, fall back to `max_model_len` so the recycling-aware cap # collapses to the prior (uncapped) admission behavior. The scheduler # always supplies the real value at runtime. - if max_num_batched_tokens is None: - max_num_batched_tokens = max_model_len + if max_in_flight_tokens is None: + max_in_flight_tokens = max_model_len self.enable_caching = enable_caching self.use_eagle = use_eagle @@ -143,7 +143,7 @@ class KVCacheManager: self.coordinator = get_kv_cache_coordinator( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, - max_num_batched_tokens=max_num_batched_tokens, + max_in_flight_tokens=max_in_flight_tokens, use_eagle=self.use_eagle, enable_caching=self.enable_caching, enable_kv_cache_events=enable_kv_cache_events, @@ -397,9 +397,12 @@ class KVCacheManager: # insufficient free blocks. # Should call this function before allocating new blocks to reduce # the number of evicted blocks. + # Free on the processed-token basis: in-flight steps' attention windows + # still read blocks below the optimistic boundary, and rejected spec + # tokens can roll it back. self.coordinator.remove_skipped_blocks( request.request_id, - total_computed_tokens, + max(0, total_computed_tokens - request.num_in_flight_tokens), num_prompt_tokens=request.num_prompt_tokens, ) @@ -472,7 +475,7 @@ class KVCacheManager: def remove_skipped_blocks( self, request_id: str, - total_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """Remove the blocks that are no longer needed from `blocks` and replace @@ -480,12 +483,12 @@ class KVCacheManager: Args: request_id: The request ID. - total_computed_tokens: The total number of computed tokens, including - local computed tokens and external computed tokens. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). num_prompt_tokens: Optional prompt length for R-SWA gap eviction. """ self.coordinator.remove_skipped_blocks( - request_id, total_computed_tokens, num_prompt_tokens + request_id, processed_computed_tokens, num_prompt_tokens ) def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 018f9c34d4..483cc2f543 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -262,7 +262,7 @@ class Scheduler(SchedulerInterface): self.kv_cache_manager = KVCacheManager( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, - max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, + max_in_flight_tokens=vllm_config.max_in_flight_tokens, enable_caching=self.cache_config.enable_prefix_caching, use_eagle=self.use_eagle, log_stats=self.log_stats, @@ -1180,6 +1180,7 @@ 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 + request.num_in_flight_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 @@ -1566,10 +1567,12 @@ class Scheduler(SchedulerInterface): stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 + request = self.requests.get(req_id) + if request is not None: + request.num_in_flight_tokens -= num_tokens_scheduled if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: # skip failed or rescheduled requests from KV load failure continue - request = self.requests.get(req_id) if request is None or request.is_finished(): # The request is already finished. This can happen if the # request is aborted while the model is executing it (e.g., @@ -2367,10 +2370,12 @@ class Scheduler(SchedulerInterface): return False, None # Free any out-of-window prefix blocks before we hand the block table to - # the connector. + # the connector, on the processed-token basis (see `allocate_slots`). self.kv_cache_manager.remove_skipped_blocks( request_id=request.request_id, - total_computed_tokens=request.num_computed_tokens, + processed_computed_tokens=max( + 0, request.num_computed_tokens - request.num_in_flight_tokens + ), num_prompt_tokens=request.num_prompt_tokens, ) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 5d554b5925..87c3f8feb7 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -508,7 +508,7 @@ class SingleTypeKVCacheManager(ABC): def remove_skipped_blocks( self, request_id: str, - total_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """ @@ -520,15 +520,15 @@ class SingleTypeKVCacheManager(ABC): Args: request_id: The request ID. - total_computed_tokens: The total number of computed tokens, including - local computed tokens and external computed tokens. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). num_prompt_tokens: Optional prompt length for attention types (e.g. R-SWA) that evict a middle gap rather than a head prefix. Ignored by the default implementation. """ del num_prompt_tokens # Remove the blocks that will be skipped during attention computation. - num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens) + num_skipped_tokens = self.get_num_skipped_tokens(processed_computed_tokens) if num_skipped_tokens <= 0: # This indicates that ALL tokens are inside attention window. # Thus we do not need to free any blocks outside attention window. @@ -639,14 +639,14 @@ class RSWAManager(FullAttentionManager): def remove_skipped_blocks( self, request_id: str, - total_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """Free gap blocks that are no longer needed for attention. Gap = blocks entirely within [ceil(prefix_len / block_size) * block_size, - max(prefix_len, total_computed_tokens - rswa_window)) + max(prefix_len, processed_computed_tokens - rswa_window)) Freed blocks are replaced with null_block in req_to_blocks so the block_table passed to FA4 is valid (null_block KV is all-zero; @@ -654,7 +654,7 @@ class RSWAManager(FullAttentionManager): """ if num_prompt_tokens is None: super().remove_skipped_blocks( - request_id, total_computed_tokens, num_prompt_tokens + request_id, processed_computed_tokens, num_prompt_tokens ) return @@ -662,7 +662,9 @@ class RSWAManager(FullAttentionManager): # First block fully after the prefill boundary. first_gap_block = cdiv(num_prompt_tokens, bs) # Decode window start position; blocks before this are evictable. - window_start = max(num_prompt_tokens, total_computed_tokens - self.rswa_window) + window_start = max( + num_prompt_tokens, processed_computed_tokens - self.rswa_window + ) last_gap_block = window_start // bs # exclusive upper bound self._remove_blocks_in_range(request_id, first_gap_block, last_gap_block) @@ -1144,20 +1146,13 @@ class MambaManager(SingleTypeKVCacheManager): def remove_skipped_blocks( self, request_id: str, - num_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: assert isinstance(self.kv_cache_spec, MambaSpec) - # NOTE (tdoublep) with async scheduling, the num_computed_tokens can contain - # draft tokens from the previous step that may or may not be rejected later. - # This can make us think we are further ahead in the sequence than we actually - # are, so let's assume that all tokens are rejected so we don't free blocks - # that we might actually need. - num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks) - super().remove_skipped_blocks( - request_id, num_computed_tokens, num_prompt_tokens + request_id, processed_computed_tokens, num_prompt_tokens ) if self.mamba_cache_mode == "align": # `last_state_block_idx` refers to the block index allocated two steps ago. @@ -1171,7 +1166,7 @@ class MambaManager(SingleTypeKVCacheManager): if ( last_state_block_idx is not None and last_state_block_idx - < cdiv(num_computed_tokens, self.block_size) - 1 + < cdiv(processed_computed_tokens, self.block_size) - 1 ): blocks = self.req_to_blocks[request_id] if blocks[last_state_block_idx] != self._null_block: @@ -1455,7 +1450,7 @@ class SinkFullAttentionManager(FullAttentionManager): def get_manager_for_kv_cache_spec( kv_cache_spec: KVCacheSpec, - max_num_batched_tokens: int, + max_in_flight_tokens: int, max_model_len: int, **kwargs, ) -> SingleTypeKVCacheManager: @@ -1468,7 +1463,8 @@ def get_manager_for_kv_cache_spec( Args: kv_cache_spec: The KVCacheSpec instance - max_num_batched_tokens: The maximum number of tokens in a batch + max_in_flight_tokens: The max tokens scheduled but not yet settled + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens` max_model_len: The maximum context length the model could serve Returns: An instance of the appropriate SingleTypeKVCacheManager subclass @@ -1489,7 +1485,7 @@ def get_manager_for_kv_cache_spec( ): kwargs["max_admission_blocks_per_request"] = ( kv_cache_spec.max_admission_blocks_per_request( - max_num_batched_tokens=max_num_batched_tokens, + max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len, ) ) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 323b1e763a..4204c31be5 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -482,25 +482,28 @@ class ChunkedLocalAttentionSpec(AttentionSpec): attention_chunk_size: int def max_admission_blocks_per_request( - self, max_num_batched_tokens: int, max_model_len: int + self, max_in_flight_tokens: int, max_model_len: int ) -> int: """Per-request admission cap, in blocks. Single source of truth for both startup pool sizing (`max_memory_usage_bytes`) and the runtime admission gate, so requests admitted by startup can also be admitted at runtime. + + `max_in_flight_tokens` is the max tokens scheduled but not yet settled + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`. """ - # During chunked prefill, we hold KV for at most one chunk window. + # During chunked prefill, we hold KV for at most one chunk window plus + # the in-flight tokens, since frees happen on the processed-token basis. num_tokens = min( - self.attention_chunk_size + max_num_batched_tokens, max_model_len + self.attention_chunk_size + max_in_flight_tokens, max_model_len ) return cdiv(num_tokens, self.block_size) def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: - max_model_len = vllm_config.model_config.max_model_len - max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens max_blocks = self.max_admission_blocks_per_request( - max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len + max_in_flight_tokens=vllm_config.max_in_flight_tokens, + max_model_len=vllm_config.model_config.max_model_len, ) return max_blocks * self.page_size_bytes @@ -544,7 +547,7 @@ class SlidingWindowSpec(AttentionSpec): ) def max_admission_blocks_per_request( - self, max_num_batched_tokens: int, max_model_len: int + self, max_in_flight_tokens: int, max_model_len: int ) -> int: """Per-request admission cap, in blocks. @@ -553,13 +556,14 @@ class SlidingWindowSpec(AttentionSpec): real-held blocks plateau at this bound because `SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots` before each chunk's `get_num_blocks_to_allocate`. + + `max_in_flight_tokens` is the max tokens scheduled but not yet settled + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`. """ # During chunked prefill, we hold KV for the last `sliding_window-1` - # computed tokens plus the newly scheduled tokens, and never more - # than `max_model_len`. - num_tokens = min( - self.sliding_window - 1 + max_num_batched_tokens, max_model_len - ) + # computed tokens plus the in-flight tokens (frees happen on the + # processed-token basis); never more than `max_model_len`. + num_tokens = min(self.sliding_window - 1 + max_in_flight_tokens, max_model_len) # +1 because the sliding window may not start from the beginning of # the block. E.g. block size 4 and num_token 4 needs two blocks # [XXCD][EF] to store the 6-token window [CDEF]. @@ -569,10 +573,9 @@ class SlidingWindowSpec(AttentionSpec): assert vllm_config.parallel_config.decode_context_parallel_size == 1, ( "DCP not support sliding window." ) - max_model_len = vllm_config.model_config.max_model_len - max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens max_blocks = self.max_admission_blocks_per_request( - max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len + max_in_flight_tokens=vllm_config.max_in_flight_tokens, + max_model_len=vllm_config.model_config.max_model_len, ) return max_blocks * self.page_size_bytes diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 0e8d4ee006..e9946a7f76 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -141,6 +141,11 @@ class Request: self.num_output_placeholders = 0 self.async_tokens_to_discard = 0 + # Tokens of steps whose output is not yet processed (async scheduling + # and PP run ahead of the GPU); `num_computed_tokens` counts them + # optimistically. + self.num_in_flight_tokens = 0 + # V2+PP+async: Enforces `pp_size` cadence between same-request decode steps # so the worker's broadcast slot ring stays consistent. self.next_decode_eligible_step = 0 diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index dfaa2234eb..44ba902f45 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -122,9 +122,7 @@ class SimpleCPUOffloadScheduler: self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator( kv_cache_config=self.cpu_kv_cache_config, max_model_len=vllm_config.model_config.max_model_len, - max_num_batched_tokens=( - vllm_config.scheduler_config.max_num_batched_tokens - ), + max_in_flight_tokens=vllm_config.max_in_flight_tokens, use_eagle=False, enable_caching=True, enable_kv_cache_events=self.enable_kv_cache_events,