iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Saddss
2026-07-08 13:34:19 +01:00
committed by GitHub
co-authored by Yifan Qiao Saddss Nick Hill
parent 934eeaecfb
commit cd0de48d08
14 changed files with 549 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/multimodal/generation/test_voxtral_realtime.py
+1 -1
View File
@@ -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
tests/v1/core/test_kv_cache_utils.py
+3
View File
@@ -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,
tests/v1/core/test_prefix_caching.py
+4 -2
View File
@@ -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,
)
tests/v1/core/test_swa_inflight_window_free.py
+259
View File
@@ -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
tests/v1/core/utils.py
+10 -11
View File
@@ -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)
tests/v1/e2e/general/test_mamba_prefix_cache.py
+183 -61
View File
@@ -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)
vllm/config/vllm.py
+11
View File
@@ -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 (
vllm/v1/core/kv_cache_coordinator.py
+16 -16
View File
@@ -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,
vllm/v1/core/kv_cache_manager.py
+12 -9
View File
@@ -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]:
vllm/v1/core/sched/scheduler.py
+9 -4
View File
@@ -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,
)
vllm/v1/core/single_type_kv_cache_manager.py
+17 -21
View File
@@ -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,
)
)
vllm/v1/kv_cache_interface.py
+18 -15
View File
@@ -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
vllm/v1/request.py
+5
View File
@@ -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
vllm/v1/simple_kv_offload/manager.py
+1 -3
View File
@@ -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,