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

[Bugfix][Prefix caching] Fix inconsistent hybrid Mamba prefix cache boundaries (#60533)

Browse Source 
Signed-off-by: Wei Zhao <weizha@aws-cmh-slurm-1-login-01.cm.cluster>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Wei Zhao <weizha@aws-cmh-slurm-1-login-01.cm.cluster>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Wei Zhao
2026-10-08 18:33:48 +00:00
committed by GitHub
co-authored by Wei Zhao Codex zjy0516 Claude Opus 5.5 mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 1664c4f444
commit d8fa463e2b
17 changed files with 546 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/kimi_k3/test_amd_kda_checkpoint.py
+6 -1
View File
@@ -75,7 +75,7 @@ def _builder(
vllm_config.cache_config.prefix_match_unit = prefix_match_unit
if num_spec:
vllm_config.speculative_config = _StubSpeculativeConfig(num_spec)
return KimiK3ROCmKDAMetadataBuilder(
builder = KimiK3ROCmKDAMetadataBuilder(
kv_cache_spec=MambaSpec(
block_size=block_size,
shapes=((16, 64),),
@@ -91,6 +91,11 @@ def _builder(
vllm_config=vllm_config,
device=torch.device("cuda"),
)
vllm_config.cache_config.hash_block_size = prefix_match_unit or block_size
vllm_config.cache_config.cache_hit_alignment_tokens = (
prefix_match_unit or block_size
)
return builder
def _build(batch: BatchSpec, block_size: int = BLOCK_SIZE, **kw):
tests/models/kimi_k3/test_kda_metadata.py
+52 -8
View File
@@ -100,6 +100,8 @@ def _make_builder(
prefix_match_unit: int | None = None,
use_eagle: bool = False,
disable_eagle_block_drop: bool = False,
cache_hit_alignment_tokens: int | None = None,
hash_block_size: int | None = None,
) -> AttentionMetadataBuilder:
vllm_config = create_vllm_config(
model_name="Qwen/Qwen3.5-0.8B",
@@ -122,6 +124,11 @@ def _make_builder(
vllm_config.cache_config.use_replayssm = use_recoverssm
vllm_config.cache_config.use_kda_recoverssm = use_recoverssm
vllm_config.cache_config.prefix_match_unit = prefix_match_unit
hash_block_size = hash_block_size or prefix_match_unit or mamba_block_size
vllm_config.cache_config.hash_block_size = hash_block_size
vllm_config.cache_config.cache_hit_alignment_tokens = (
cache_hit_alignment_tokens or hash_block_size
)
builder = builder_cls(
kv_cache_spec=MambaSpec(
block_size=mamba_block_size,
@@ -220,6 +227,8 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
vllm_config=builder_config,
device=DEVICE,
)
builder_config.cache_config.hash_block_size = BLOCK_SIZE
builder_config.cache_config.cache_hit_alignment_tokens = BLOCK_SIZE
# An all-prefill speculative batch used to leave an all-false spec mask
# alive, then access active_non_spec_mask_cpu before it was initialized.
@@ -287,18 +296,45 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
@pytest.mark.parametrize(
"batch,prefix_match_unit,hit_alignment,expected_offsets,expected_states",
[
(
BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]),
None,
16,
[48, 0],
[2, NULL_BLOCK_ID],
),
(
BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]),
16,
16,
[48, 0],
[2, NULL_BLOCK_ID],
),
(BatchSpec(seq_lens=[530], query_lens=[530]), None, 128, [512], [32]),
(BatchSpec(seq_lens=[530], query_lens=[50]), None, 128, [32], [32]),
(BatchSpec(seq_lens=[530], query_lens=[530]), 16, 16, [528], [32]),
(BatchSpec(seq_lens=[530], query_lens=[50]), 16, 16, [48], [32]),
],
)
def test_internal_checkpoint_metadata_targets_last_aligned_boundary(
batch, prefix_match_unit, hit_alignment, expected_offsets, expected_states
):
device = torch.device("cuda")
batch = BatchSpec(seq_lens=[50, 32], query_lens=[50, 16])
common_attn_metadata = create_common_attn_metadata(
batch, BLOCK_SIZE, device, arange_block_indices=True
).replace(is_prefilling=torch.tensor([True, True]))
).replace(is_prefilling=torch.tensor([True] * len(batch.seq_lens)))
builder = _make_builder(
KimiK3KDAMetadataBuilder,
num_speculative_tokens=0,
full_cuda_graph=False,
mamba_cache_mode="align",
num_prefill_checkpoint_blocks=1,
prefix_match_unit=prefix_match_unit,
cache_hit_alignment_tokens=hit_alignment,
hash_block_size=16,
device=device,
)
assert isinstance(builder, KimiK3KDAMetadataBuilder)
@@ -313,26 +349,33 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
assert actual.checkpoint is not None
torch.testing.assert_close(
actual.checkpoint.state_indices,
torch.tensor([2, NULL_BLOCK_ID], dtype=torch.int32, device=device),
torch.tensor(expected_states, dtype=torch.int32, device=device),
)
torch.testing.assert_close(
actual.checkpoint.checkpoint_offsets,
torch.tensor([48, 0], dtype=torch.int32, device=device),
torch.tensor(expected_offsets, dtype=torch.int32, device=device),
)
@pytest.mark.parametrize(
("disable_eagle_block_drop", "prefix_match_unit", "expected_offset"),
[(False, 16, 80), (True, 16, 96), (False, 8, None)],
("seq_len", "disable_eagle_block_drop", "prefix_match_unit", "expected_offset"),
[
(100, False, 16, 80),
(100, True, 16, 96),
(100, False, 8, None),
(96, False, 16, 80),
(96, True, 16, 80),
],
)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
def test_spec_internal_checkpoint_metadata_targets_replay_boundary(
seq_len: int,
disable_eagle_block_drop: bool,
prefix_match_unit: int,
expected_offset: int | None,
) -> None:
device = torch.device("cuda")
batch = BatchSpec(seq_lens=[100], query_lens=[100])
batch = BatchSpec(seq_lens=[seq_len], query_lens=[seq_len])
common_attn_metadata = create_common_attn_metadata(
batch, BLOCK_SIZE, device, arange_block_indices=True
)
@@ -390,6 +433,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset():
full_cuda_graph=False,
mamba_cache_mode="align",
num_prefill_checkpoint_blocks=1,
prefix_match_unit=16,
device=device,
)
assert isinstance(builder, KimiK3KDAMetadataBuilder)
tests/models/kimi_k3/test_prefix_cache.py
+40
View File
@@ -8,6 +8,7 @@ import random
from collections.abc import Iterator
from dataclasses import dataclass
from math import lcm
from uuid import uuid4
import pytest
import requests
@@ -353,6 +354,45 @@ def _format(turns: list[Turn]) -> str:
return "\n".join(rows)
@pytest.mark.parametrize(
"prompt_len",
[
# Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned:
# ensure lookup reads the cached attention proof at 7040 despite
# the final hit limit of 7039, then applies EAGLE's 128-token drop.
pytest.param(7040, id="eagle-proof-scan"),
# Blocks=6144, PMU=128. P=7296 is an exact PMU multiple:
# ensure the Mamba checkpoint is 7168, not 7040 from applying both
# the P - 1 cap and EAGLE's 128-token drop.
pytest.param(7296, id="eagle-double-cap"),
],
)
def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None:
"""Resends reuse the Mamba checkpoint and its companion EAGLE attention proof."""
mode = Mode(128, True)
salt = str(uuid4())
rng = random.Random(0)
prompt = [rng.randint(1000, 150000) for _ in range(prompt_len)]
deployment = DEPLOYMENTS["plain"]
with _serve(deployment, mode) as servers:
url = servers[0].url_root
cold = _complete(url, prompt, salt, max_tokens=1)
resend = _complete(url, prompt, salt, max_tokens=1)
recompute = _complete(url, prompt, f"{salt}-recompute", max_tokens=1)
# EAGLE's dropped hash unit already leaves tokens to recompute on resend.
expected = (prompt_len // 128 - 1) * 128
actual = _cached(resend)
assert _cached(cold) == _cached(recompute) == 0
assert actual == expected, f"{prompt_len=}: {actual=}, {expected=}"
check_logprobs_close(
outputs_0_lst=[_tokens_text_logprobs(recompute)],
outputs_1_lst=[_tokens_text_logprobs(resend)],
name_0="recompute",
name_1="resend",
)
@pytest.mark.parametrize(
"name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES]
)
tests/v1/attention/test_gdn_metadata_builder.py
+6 -2
View File
@@ -385,7 +385,7 @@ def test_cudagraph_capture_batch_stays_decode_only():
def _create_checkpoint_builder_and_batch(
num_spec: int, prefix_match_unit: int = 16
num_spec: int, prefix_match_unit: int | None = 16
) -> tuple[GDNAttentionMetadataBuilder, CommonAttentionMetadata, dict]:
vllm_config = create_vllm_config(
model_name="Qwen/Qwen3.5-0.8B",
@@ -413,6 +413,8 @@ def _create_checkpoint_builder_and_batch(
vllm_config=vllm_config,
device=DEVICE,
)
vllm_config.cache_config.hash_block_size = prefix_match_unit or 64
vllm_config.cache_config.cache_hit_alignment_tokens = prefix_match_unit or 64
batch = BatchSpec(
seq_lens=[65, 100, 100], query_lens=[3 if num_spec else 1, 100, 99]
)
@@ -431,7 +433,9 @@ def _create_checkpoint_builder_and_batch(
@pytest.mark.parametrize("num_spec", [0, 3])
@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)])
@pytest.mark.parametrize(
"prefix_match_unit, expected_offset", [(None, 64), (16, 96), (8, 96)]
)
def test_checkpoint_metadata_preserves_non_spec_order(
num_spec, prefix_match_unit, expected_offset
):
tests/v1/attention/test_mamba2_prefill_checkpoints.py
+7 -2
View File
@@ -77,7 +77,9 @@ def _create_mamba2_builder(
num_prefill_checkpoint_blocks: int = 1,
prefill_checkpoint_alignment: int | None = 1,
) -> Mamba2AttentionMetadataBuilder:
vllm_config = create_vllm_config(block_size=MAMBA_BLOCK_SIZE)
vllm_config = create_vllm_config(
model_name="Qwen/Qwen3.5-0.8B", block_size=MAMBA_BLOCK_SIZE
)
vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
# get_mamba_chunk_size() reads `mamba_chunk_size` then `chunk_size`.
vllm_config.model_config.hf_text_config.mamba_chunk_size = CHUNK_SIZE
@@ -89,12 +91,15 @@ def _create_mamba2_builder(
num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
prefill_checkpoint_alignment=prefill_checkpoint_alignment,
)
return Mamba2AttentionMetadataBuilder(
builder = Mamba2AttentionMetadataBuilder(
kv_cache_spec=spec,
layer_names=["layer.0"],
vllm_config=vllm_config,
device=DEVICE,
)
vllm_config.cache_config.hash_block_size = MAMBA_BLOCK_SIZE
vllm_config.cache_config.cache_hit_alignment_tokens = MAMBA_BLOCK_SIZE
return builder
def _build(builder, seq_lens, query_lens):
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+201 -52
View File
@@ -22,6 +22,7 @@ from vllm.v1.core.kv_cache_utils import (
get_block_hash,
get_group_id,
init_none_hash,
resolve_cache_hit_alignment_tokens,
)
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.kv_cache_interface import (
@@ -148,29 +149,84 @@ def make_full_mamba_manager(
)
def test_dcp_fine_hit_retention_uses_hash_alignment_without_eagle():
"""DCP must not discard a reusable Mamba state at hash alignment."""
hash_block_size = 2
@pytest.mark.parametrize(
"dcp_world_size,hash_block_size,block_size,prompt_tokens,checkpoint_blocks",
[
(4, 2, 2, 7, 0),
(8, 16, 16, 512, 1),
(8, 16, 16, 513, 1),
*[(8, 128, 896, p, 1) for p in (127, 128, 129, 7040, 7168, 7296, 7449)],
],
)
@pytest.mark.parametrize("use_eagle", [False, True])
@pytest.mark.parametrize("extension_tokens", [0, 32])
def test_dcp_fine_hit_retention_uses_hash_alignment(
dcp_world_size,
hash_block_size,
block_size,
prompt_tokens,
checkpoint_blocks,
use_eagle,
extension_tokens,
):
"""Resends and extensions reuse the retained hash-aligned checkpoint."""
manager = make_full_mamba_manager(
dcp_world_size=4,
dcp_world_size=dcp_world_size,
hash_block_size=hash_block_size,
full_block_size=hash_block_size,
mamba_block_size=hash_block_size,
num_blocks=64,
full_block_size=block_size,
mamba_block_size=block_size,
num_blocks=128,
use_eagle=use_eagle,
num_speculative_blocks=4 if use_eagle else 0,
num_prefill_checkpoint_blocks=checkpoint_blocks,
)
manager.coordinator.retention_interval = 0
token_ids = list(range(7))
producer = make_request("producer", token_ids, hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(producer)
assert num_computed == 0
assert manager.allocate_slots(producer, 6, 0, computed_blocks) is not None
assert manager.coordinator.enable_partial_hash_hits
scheduler = SimpleNamespace(
block_size=block_size * dcp_world_size,
cache_config=SimpleNamespace(block_size=block_size),
hash_block_size=hash_block_size,
max_num_scheduled_tokens=8192,
scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
use_eagle_block_drop=use_eagle,
mamba_partial_cache_hit=True,
mamba_shared_prefix_checkpoint=False,
mamba_has_prefill_checkpoint_blocks=bool(checkpoint_blocks),
mamba_prefill_checkpoint_alignment=16,
)
producer = make_request(
"producer", list(range(prompt_tokens)), hash_block_size, sha256
)
while producer.num_computed_tokens < prompt_tokens:
scheduled = Scheduler._mamba_block_aligned_split(
scheduler, producer, prompt_tokens - producer.num_computed_tokens
)
assert scheduled > 0
manager.new_step_starts()
assert manager.allocate_slots(producer, scheduled) is not None
producer.num_computed_tokens += scheduled
manager.free(producer)
manager.new_step_starts()
consumer = make_request("consumer", token_ids, hash_block_size, sha256)
_, num_computed, _ = manager.get_computed_blocks(consumer)
assert num_computed == 6
replay = make_request(
"replay", list(range(prompt_tokens + extension_tokens)), hash_block_size, sha256
)
_, hit, _ = manager.get_computed_blocks(replay)
expected = (
max((prompt_tokens // hash_block_size - 1) * hash_block_size, 0)
if use_eagle
else ((prompt_tokens - 1) // hash_block_size * hash_block_size)
)
assert hit == expected
cached_states = {
end
for end in range(hash_block_size, prompt_tokens + 1, hash_block_size)
if manager.block_pool.get_cached_block(
producer.block_hashes[end // hash_block_size - 1], [1]
)
is not None
}
assert cached_states == ({expected} if expected else set())
@pytest.mark.parametrize("dcp_world_size", [1, 4])
@@ -339,7 +395,7 @@ def test_hybrid_mamba_align_partial_hash_hit():
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert num_computed == 0
blocks = manager.allocate_slots(req0, 6, num_computed, computed_blocks)
@@ -488,7 +544,7 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert num_computed == 0
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
@@ -538,7 +594,7 @@ def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks():
num_prefill_checkpoint_blocks=1,
)
owner = make_request("owner", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
owner = make_request("owner", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(owner)
assert manager.allocate_slots(owner, 6, num_computed, computed_blocks) is not None
manager.free(owner)
@@ -619,44 +675,70 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle():
assert num_computed == 112
@pytest.mark.parametrize(
"retention_interval,transient_published",
[(None, True), (0, False)],
)
@pytest.mark.parametrize("retention_interval", [None, 0, 32, 64, 256])
@pytest.mark.parametrize("hash_block_size", [16, 32])
@pytest.mark.parametrize("use_eagle", [False, True])
def test_internal_checkpoint_publication_respects_retention(
retention_interval, transient_published
retention_interval, hash_block_size, use_eagle
):
"""With retention_interval=0, a mid-prompt internal checkpoint is
request-local and must not enter the prefix cache; the prompt-end
checkpoint keeps the existing publication behavior. Dense retention
(None) is unchanged.
"""
hash_block_size = 16
"""Publish only retained block boundaries and the final prompt checkpoint."""
manager = make_full_mamba_manager(
dcp_world_size=1,
hash_block_size=hash_block_size,
full_block_size=hash_block_size,
mamba_block_size=32,
num_prefill_checkpoint_blocks=1,
use_eagle=use_eagle,
)
manager.coordinator.retention_interval = retention_interval
request = make_request("producer", list(range(240)), hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
# Mid-prompt chunk: exports a transient checkpoint at 112.
# The EAGLE adjustment can make a chunk checkpoint block-aligned or partial.
assert manager.allocate_slots(request, 128, num_computed, computed_blocks)
transient_hash = request.block_hashes[112 // hash_block_size - 1]
checkpoint = (128 if use_eagle else 127) // hash_block_size * hash_block_size - (
hash_block_size if use_eagle else 0
)
retained_checkpoint = checkpoint % 32 == 0 and (
retention_interval is None
or (retention_interval > 0 and checkpoint % retention_interval == 0)
)
transient_hash = request.block_hashes[checkpoint // hash_block_size - 1]
transient_hit = manager.block_pool.get_cached_block(transient_hash, [1])
assert (transient_hit is not None) == transient_published
assert (transient_hit is not None) == retained_checkpoint
offloads = drain_boundary_state_offloads(manager).get(request.request_id, [])
assert any(position == checkpoint for _, _, position in offloads) == (
retained_checkpoint
)
assert all(
position % 32 == 0
and (
retention_interval is None
or (retention_interval > 0 and position % retention_interval == 0)
)
for _, _, position in offloads
)
if retention_interval in (None, 32, 64):
retained_hash = request.block_hashes[128 // hash_block_size - 1]
assert manager.block_pool.get_cached_block(retained_hash, [1]) is not None
request.num_computed_tokens = 128
manager.new_step_starts()
# Final chunk: the prompt-end checkpoint at 224 stays published.
# The final checkpoint is reusable regardless of the retention interval.
assert manager.allocate_slots(request, 112) is not None
end_hash = request.block_hashes[224 // hash_block_size - 1]
prompt_checkpoint = (
240 if use_eagle else 239
) // hash_block_size * hash_block_size - (hash_block_size if use_eagle else 0)
end_hash = request.block_hashes[prompt_checkpoint // hash_block_size - 1]
assert manager.block_pool.get_cached_block(end_hash, [1]) is not None
assert any(
position == prompt_checkpoint
for _, _, position in drain_boundary_state_offloads(manager).get(
request.request_id, []
)
)
def test_transient_checkpoint_evicts_retained_boundary_hash():
@@ -802,9 +884,9 @@ def test_hash_aligned_query_end_uses_regular_partial_tail():
mamba_block_size=mamba_block_size,
num_prefill_checkpoint_blocks=1,
)
request = make_request("producer", list(range(14)), hash_block_size, sha256)
request = make_request("producer", list(range(15)), hash_block_size, sha256)
new_blocks = manager.allocate_slots(request, request.num_tokens)
new_blocks = manager.allocate_slots(request, 14)
assert new_blocks is not None
mamba_manager = manager.coordinator.single_type_managers[1]
@@ -928,7 +1010,7 @@ def test_boundary_state_offloads_returns_cow_target():
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
@@ -1004,7 +1086,7 @@ def test_finished_partial_tail_uses_table_source_once():
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
source_block_id = manager.get_blocks("0").get_block_ids()[1][1]
@@ -1024,7 +1106,7 @@ def test_finished_partial_tail_uses_table_source_once():
enable_caching=True,
hash_block_size=hash_block_size,
)
req1 = make_request("1", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = advanced_manager.get_computed_blocks(req1)
assert (
advanced_manager.allocate_slots(req1, 6, num_computed, computed_blocks)
@@ -1212,7 +1294,7 @@ def test_truncate_computed_blocks_preserves_sparse_prefix_positions():
enable_caching=True,
hash_block_size=hash_block_size,
)
producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
producer = make_request("producer", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
blocks, num_computed, _ = manager.get_computed_blocks(producer)
assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None
manager.free(producer)
@@ -1267,7 +1349,7 @@ def test_truncate_computed_blocks_allows_short_mamba_group_only():
enable_caching=True,
hash_block_size=hash_block_size,
)
producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
producer = make_request("producer", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
blocks, num_computed, _ = manager.get_computed_blocks(producer)
assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None
manager.free(producer)
@@ -1331,7 +1413,7 @@ def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit():
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert num_computed == 0
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
@@ -1411,7 +1493,7 @@ def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit():
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert num_computed == 0
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
@@ -1470,7 +1552,7 @@ def test_hybrid_full_attention_partial_hash_hit_uses_cow():
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert num_computed == 0
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
@@ -1544,7 +1626,7 @@ def test_hybrid_partial_hit_cow_target_starts_uncached():
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert num_computed == 0
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
@@ -1700,7 +1782,7 @@ def test_cow_retained_blocks_returned_for_release():
enable_caching=True,
hash_block_size=hash_block_size,
)
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
@@ -2133,6 +2215,72 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits():
assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed
@pytest.mark.parametrize("with_sliding_window", [False, True])
def test_resolved_cache_hit_alignment_matches_coordinator(with_sliding_window):
"""Workers place Mamba checkpoints on the engine-resolved alignment, so it
must be where the coordinator lands hits, including when a group without
fine-grained lookups disables partial hits."""
hash_block_size = 2
mamba_block_size = 4 * hash_block_size
kv_cache_groups = [
KVCacheGroupSpec(
["full"],
FullAttentionSpec(
block_size=hash_block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
),
KVCacheGroupSpec(
["mamba"],
MambaSpec(
block_size=mamba_block_size,
shapes=(1, 1),
dtypes=(torch.float32,),
mamba_cache_mode="align",
),
),
]
if with_sliding_window:
kv_cache_groups.append(
KVCacheGroupSpec(
["swa"],
SlidingWindowSpec(
block_size=2 * hash_block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
sliding_window=2 * hash_block_size,
),
)
)
kv_cache_config = KVCacheConfig(
num_blocks=64, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
)
coordinator = make_kv_cache_manager(
kv_cache_config=kv_cache_config,
max_model_len=8192,
enable_caching=True,
hash_block_size=hash_block_size,
).coordinator
vllm_config = SimpleNamespace(
cache_config=SimpleNamespace(enable_prefix_caching=True),
parallel_config=SimpleNamespace(decode_context_parallel_size=1),
)
assert coordinator.enable_partial_hash_hits != with_sliding_window
assert (
resolve_cache_hit_alignment_tokens(
kv_cache_config,
vllm_config,
coordinator.scheduler_block_size,
hash_block_size,
)
== coordinator._cache_hit_alignment_tokens
)
def test_opted_out_scratch_group_keeps_partial_hash_hits():
hash_block_size = 2
mamba_block_size = 2 * hash_block_size
@@ -2179,7 +2327,7 @@ def test_opted_out_scratch_group_keeps_partial_hash_hits():
assert manager.coordinator.enable_partial_hash_hits
req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
manager.free(req0)
@@ -2240,7 +2388,7 @@ def test_kpool_tail_supports_128_token_partial_hash_hits():
assert manager.coordinator.enable_partial_hash_hits
shared_prefix = [10] * (12 * hash_block_size)
req0 = make_request("0", shared_prefix, hash_block_size, sha256)
req0 = make_request("0", shared_prefix + [11], hash_block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert (
manager.allocate_slots(req0, len(shared_prefix), num_computed, computed_blocks)
@@ -2266,7 +2414,7 @@ def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int):
)
assert manager.coordinator.enable_partial_hash_hits
req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2], 2, sha256)
req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2, 3], 2, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
manager.free(req0)
@@ -2318,7 +2466,7 @@ def test_dcp_partial_hit_resumes_on_replicated_mamba_snapshot(
assert manager.coordinator.single_type_managers[1].block_size == block_size
prefix = list(range(12))
req0 = make_request("snapshot-owner", prefix, block_size, sha256)
req0 = make_request("snapshot-owner", prefix + [12], block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert manager.allocate_slots(req0, 12, num_computed, computed_blocks) is not None
manager.free(req0)
@@ -2355,7 +2503,7 @@ def test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots():
mamba_block_size=block_size,
)
prefix = list(range(12))
req0 = make_request("joint-owner", prefix, block_size, sha256)
req0 = make_request("joint-owner", prefix + [12], block_size, sha256)
computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None
manager.new_step_starts()
@@ -2702,6 +2850,7 @@ def test_checkpoint_reservation_matches_worker(
mamba_block_size=mamba_block_size,
checkpoint_alignment=16,
drop_eagle_block=use_eagle,
cache_hit_alignment_tokens=mamba_manager.cache_hit_alignment_tokens,
)
reserved = mamba_manager._checkpoints.get(request.request_id)
if reserved is not None:
tests/v1/core/test_prefix_caching.py
+10 -26
View File
@@ -5047,19 +5047,10 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary():
def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt():
"""An identical resend and a longer sibling resume at DIFFERENT positions.
"""EAGLE resend and extension reuse the same retained checkpoint.
How far a lookup matches depends on who is asking. A resend of the same
prompt caps its lookup at ``num_tokens - 1`` (the last token is recomputed
for logits), while a sibling whose prompt merely starts with this one caps
above the prompt. The two coincide unless the prompt length is an exact
multiple of the alignment -- there they differ by one alignment unit, and
under the EAGLE drop BOTH are reachable.
Retaining only the higher one leaves the resend with every retained state
above every candidate its lookup can produce, and the reconciled hit
collapses to 0 -- the same zero-hit failure sparse retention already fixes
at unaligned prompt lengths.
Attention may inspect the prompt's final block before dropping it, so an
aligned resend must not lose another block to the P - 1 lookup limit.
"""
block_size = 32
num_spec = 3
@@ -5097,16 +5088,12 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt():
use_eagle=True,
)
# 128 tokens, an exact multiple of the 32-token alignment. A longer sibling
# matches 128 and drops to 96; this prompt's own resend caps at 127, matches
# 96 and drops to 64. Both states must survive retention.
# Both requests prove 128 attention tokens, drop one block, and reuse 96.
token_ids = [i for i in range(4) for _ in range(block_size)]
req0 = make_request("0", token_ids, block_size, sha256)
computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
assert num_computed_tokens == 0
# Prefill in block-aligned chunks the way the align-mode scheduler does: a
# state only materializes as a chunk's running-state block, so a
# single-shot prefill could not retain the lower one.
# Materialize the checkpoint through normal block-aligned prefill chunks.
for chunk_end in (32, 64, 96, 128):
blocks = manager.allocate_slots(
req0,
@@ -5118,8 +5105,7 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt():
assert blocks is not None
req0.num_computed_tokens = chunk_end
# Block ``i`` ends at token ``(i + 1) * 32``, so positions 64 and 96 are
# mamba blocks 1 and 2.
# Both materialized replay checkpoints survive; resend and extension use 96.
pool = manager.block_pool
expected_mamba_cached = {1, 2}
for i in range(4):
@@ -5130,15 +5116,13 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt():
assert cached is None, f"mamba hash {i} should not be cached"
manager.free(req0)
# The identical resend: full attention matches blocks 0-2 (96 tokens, capped
# by num_tokens - 1) and the EAGLE drop caps the candidate at 64. Without
# the lower state retained the reconciled hit would be 0.
# The identical resend must reuse the retained checkpoint, not miss to zero.
req1 = make_request("1", token_ids, block_size, sha256)
computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
assert num_computed_tokens == 2 * block_size
assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 2]
assert num_computed_tokens == 3 * block_size
assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 3]
# The longer sibling resumes one alignment unit higher, off the same prompt.
# The longer sibling reuses that same checkpoint.
longer = make_request("2", token_ids + [9] * block_size, block_size, sha256)
computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(longer)
assert num_computed_tokens == 3 * block_size
vllm/config/cache.py
+11
View File
@@ -230,6 +230,15 @@ class CacheConfig:
for hybrid models where requests occupy multiple KV cache groups."""
kv_cache_max_concurrency: float | None = field(default=None, init=False)
"""Per-DP-engine maximum concurrency at max_model_len tokens."""
# The two differ when hashes are finer than the scheduler block but partial
# hash hits are off (e.g. a sliding-window group cannot look up sub-block
# hits), so hits only land on the scheduler block.
# TODO(zjy0516): Merge into one value
hash_block_size: int | None = field(default=None, init=False)
"""Tokens per prefix-cache block hash."""
cache_hit_alignment_tokens: int | None = field(default=None, init=False)
"""Token granularity at which prefix-cache hits land: `hash_block_size`
with partial hash hits, otherwise the scheduler block size."""
kv_sharing_fast_prefill: bool = False
"""In some KV sharing setups, e.g. YOCO (https://arxiv.org/abs/2405.05254),
@@ -298,6 +307,8 @@ class CacheConfig:
"effective_attention_block_size",
"kv_cache_size_tokens",
"kv_cache_max_concurrency",
"hash_block_size",
"cache_hit_alignment_tokens",
# Feature toggles not impacting compiled graph shape
"kv_sharing_fast_prefill",
"swa_bounded_replay",
vllm/model_executor/layers/mamba/checkpoint.py
+9 -2
View File
@@ -25,6 +25,7 @@ def compute_mamba_prefill_checkpoints(
mamba_block_size: int,
checkpoint_alignment: int | None,
drop_eagle_block: bool,
cache_hit_alignment_tokens: int,
) -> tuple[list[int], list[int]]:
"""Per-row internal prefill checkpoint offsets and cache block columns.
@@ -43,7 +44,7 @@ def compute_mamba_prefill_checkpoints(
for seq_len, query_len in zip(seq_lens, query_lens):
query_start = seq_len - query_len
position = get_mamba_prefill_checkpoint_position(
seq_len, hash_block_size, drop_eagle_block=drop_eagle_block
seq_len, cache_hit_alignment_tokens, drop_eagle_block=drop_eagle_block
)
valid = is_mamba_prefill_checkpoint_valid(
query_start=query_start,
@@ -106,12 +107,17 @@ class MambaPrefillCheckpointBuilder:
return None
if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0:
return None
cache_config = self.vllm_config.cache_config
hash_block_size = cache_config.hash_block_size
cache_hit_alignment_tokens = cache_config.cache_hit_alignment_tokens
if hash_block_size is None or cache_hit_alignment_tokens is None:
# Unresolved only for the minimal KV cache built for profiling.
return None
assert m.seq_lens_cpu_upper_bound is not None
all_query_lens = m.query_start_loc_cpu.diff().tolist()
query_lens = [all_query_lens[row] for row in request_rows]
seq_lens = m.seq_lens_cpu_upper_bound.tolist()
block_size = self.kv_cache_spec.block_size
hash_block_size = self.vllm_config.cache_config.prefix_match_unit or block_size
speculative_config = self.vllm_config.speculative_config
drop_eagle_block = (
speculative_config is not None and speculative_config.use_eagle_block_drop()
@@ -123,6 +129,7 @@ class MambaPrefillCheckpointBuilder:
mamba_block_size=block_size,
checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment,
drop_eagle_block=drop_eagle_block,
cache_hit_alignment_tokens=cache_hit_alignment_tokens,
)
if not any(checkpoint_offsets):
return None
vllm/v1/core/block_pool.py
+2 -5
View File
@@ -479,8 +479,8 @@ class BlockPool:
kv_cache_group_id: KV cache group that owns the partial entry.
block_size: Cache block size for the owning group. The partial
entry hash itself is always the prefix-chain hash at
``num_tokens``; ``block_size`` is used to assert that the
entry is partial within the owning cache block.
``num_tokens``; ``block_size`` must be a multiple of the hash
block size.
replace_existing_hashes: Whether the block contents were replaced
and all existing cache entries must be removed before the new
entry is registered.
@@ -494,9 +494,6 @@ class BlockPool:
return None
assert block_size % self.hash_block_size == 0
assert replace_existing_hashes or (
block_size > self.hash_block_size and num_tokens % block_size != 0
)
block_hash = self._get_partial_block_hash(request, num_tokens)
num_hash_blocks = num_tokens // self.hash_block_size
block_hash_with_group_id = make_block_hash_with_group_id(
vllm/v1/core/kv_cache_coordinator.py
+29 -40
View File
@@ -4,13 +4,14 @@ from abc import ABC, abstractmethod
from collections.abc import Sequence
from typing import NamedTuple
from vllm.logger import init_logger
from vllm.utils.math_utils import cdiv, round_down
from vllm.v1.core.block_pool import BlockPool
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
from vllm.v1.core.kv_cache_utils import (
BlockHash,
KVCacheBlock,
eagle_proof_margin,
partial_hash_hits_enabled,
)
from vllm.v1.core.single_type_kv_cache_manager import (
CrossAttentionManager,
@@ -28,8 +29,6 @@ from vllm.v1.kv_cache_interface import (
)
from vllm.v1.request import Request
logger = init_logger(__name__)
def _validate_prefix_cache_retention_interval(
retention_interval: int | None,
@@ -695,37 +694,15 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
# managers in every group. TP needs hashing finer than the Mamba block;
# DCP accepts equality because it scales the effective full-attention
# block instead.
has_partial_mamba_group = any(
isinstance(g.kv_cache_spec, MambaSpec)
and g.kv_cache_spec.mamba_cache_mode == "align"
and (
(dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size)
or (
dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size
)
)
for g in kv_cache_config.kv_cache_groups
)
self.enable_partial_hash_hits = (
allow_partial_hash_hits and has_partial_mamba_group
allow_partial_hash_hits
and partial_hash_hits_enabled(
kv_cache_config.kv_cache_groups,
hash_block_size,
dcp_world_size,
[type(manager) for manager in self.single_type_managers],
)
)
if self.enable_partial_hash_hits:
unsupported_partial_hit_managers = {
type(manager).__name__
for manager, group in zip(
self.single_type_managers, kv_cache_config.kv_cache_groups
)
if group.kv_cache_spec.prefix_cacheable
and not manager.supports_fine_grained_hash_lookup
and manager.block_size != hash_block_size
}
if unsupported_partial_hit_managers:
self.enable_partial_hash_hits = False
logger.warning_once(
"Disabling fine-grained prefix-cache hits because these KV "
"cache managers require block-aligned lookups: %s.",
", ".join(sorted(unsupported_partial_hit_managers)),
)
cache_hit_alignment_tokens = self._cache_hit_alignment_tokens
for manager in self.single_type_managers:
manager.cache_hit_alignment_tokens = cache_hit_alignment_tokens
@@ -899,15 +876,15 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
# mamba: its finder never drops (draft models have no mamba
# layers), so the hit would grow past the candidate.
if drop_eagle_block and not isinstance(spec, MambaSpec):
eagle_margin = (
self.hash_block_size
if self.enable_partial_hash_hits
and manager_cls.supports_fine_grained_hash_lookup
and group_block_size > self.hash_block_size
else group_block_size
eagle_margin = eagle_proof_margin(
group_block_size,
self.hash_block_size,
self.enable_partial_hash_hits
and manager_cls.supports_fine_grained_hash_lookup,
)
_max_length = min(
curr_hit_length + eagle_margin, max_cache_hit_length
curr_hit_length + eagle_margin,
len(block_hashes) * self.hash_block_size,
)
hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit(
block_hashes=block_hashes,
@@ -983,9 +960,21 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
for spec, group_ids, manager_cls, use_eagle in self.attention_groups:
manager = self.single_type_managers[group_ids[0]]
lookup_length = max_cache_hit_length
if use_eagle and not isinstance(spec, MambaSpec):
eagle_margin = eagle_proof_margin(
manager.block_size,
self.hash_block_size,
self.enable_partial_hash_hits
and manager_cls.supports_fine_grained_hash_lookup,
)
lookup_length = min(
max_cache_hit_length + eagle_margin,
len(block_hashes) * self.hash_block_size,
)
blocks, group_hit = manager_cls.find_longest_cache_hit(
block_hashes=block_hashes,
max_length=max_cache_hit_length,
max_length=lookup_length,
kv_cache_group_ids=group_ids,
block_pool=manager.block_pool,
kv_cache_spec=spec,
vllm/v1/core/kv_cache_utils.py
+93 -15
View File
@@ -53,6 +53,7 @@ from vllm.v1.utils import tensor_data
if TYPE_CHECKING:
from vllm.v1.core.block_pool import BlockPool
from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager
# BlockHash represents the hash of a single KV-cache block used for
@@ -747,6 +748,60 @@ def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCache
return replace(spec, block_size=block_size)
def partial_hash_hits_enabled(
kv_cache_groups: Sequence[KVCacheGroupSpec],
hash_block_size: int,
dcp_world_size: int = 1,
manager_classes: Sequence[type["SingleTypeKVCacheManager"]] | None = None,
) -> bool:
"""Whether aligned Mamba states support sub-block prefix-cache hits.
If ``manager_classes`` (one per group) is given, also require every other
prefix-cacheable group to support fine-grained lookups.
"""
if not any(
isinstance(spec, MambaSpec)
and spec.mamba_cache_mode == "align"
and (
(dcp_world_size == 1 and spec.block_size > hash_block_size)
or (dcp_world_size > 1 and spec.block_size >= hash_block_size)
)
for group in kv_cache_groups
for spec in iter_layer_specs(group.kv_cache_spec)
):
return False
if manager_classes is None:
return True
unsupported_managers = {
manager_cls.__name__
for group, manager_cls in zip(kv_cache_groups, manager_classes)
if group.kv_cache_spec.prefix_cacheable
and not manager_cls.supports_fine_grained_hash_lookup
and resolve_dcp_kv_block_size(group.kv_cache_spec, dcp_world_size)
!= hash_block_size
}
if unsupported_managers:
logger.warning_once(
"Disabling fine-grained prefix-cache hits because these KV "
"cache managers require block-aligned lookups: %s.",
", ".join(sorted(unsupported_managers)),
)
return False
return True
def eagle_proof_margin(
block_size: int, hash_block_size: int, fine_grained_lookup: bool
) -> int:
"""Tokens an EAGLE group matches past a cache hit before dropping them.
Fine-grained lookups drop one hash unit; others drop one cache block.
"""
if fine_grained_lookup and block_size > hash_block_size:
return hash_block_size
return block_size
def resolve_kv_cache_block_sizes(
kv_cache_config: KVCacheConfig,
vllm_config: VllmConfig,
@@ -775,8 +830,7 @@ def resolve_kv_cache_block_sizes(
if groups and not groups[0].kv_cache_spec.dcp_sharded:
dcp = 1
bs = cache_config.block_size * dcp
# The Mamba prefill checkpoint builder reads prefix_match_unit directly,
# so a value dropped here puts its checkpoints off the scheduler's grid.
# Reject a prefix_match_unit that would otherwise be silently ignored.
if (
cache_config.prefix_match_unit not in (None, bs)
and len(groups) == 1
@@ -797,8 +851,8 @@ def resolve_kv_cache_block_sizes(
]
scheduler_block_size = math.lcm(*group_block_sizes)
logger.info("kv cache group sizes %s", group_block_sizes)
logger.info("kv lcm block sizes %s", scheduler_block_size)
logger.info_once("kv cache group sizes %s", tuple(group_block_sizes))
logger.info_once("kv lcm block sizes %s", scheduler_block_size)
# Block hashes are only consumed by prefix caching and KV connectors
# (P/D, offloading); when neither is active, keep hash_block_size equal
@@ -839,18 +893,10 @@ def resolve_kv_cache_block_sizes(
and isinstance(spec.tokens_per_state, int)
and spec.tokens_per_state > 1
}
has_partial_mamba_group = any(
isinstance(spec, MambaSpec)
and spec.mamba_cache_mode == "align"
and (
(dcp == 1 and block_size > hash_block_size)
or (dcp > 1 and block_size >= hash_block_size)
)
for group, block_size in zip(groups, group_block_sizes)
for spec in iter_layer_specs(group.kv_cache_spec)
)
cache_hit_alignment = (
hash_block_size if has_partial_mamba_group else scheduler_block_size
hash_block_size
if partial_hash_hits_enabled(groups, hash_block_size, dcp)
else scheduler_block_size
)
if any(cache_hit_alignment % alignment for alignment in prefix_alignments):
raise ValueError(
@@ -861,6 +907,38 @@ def resolve_kv_cache_block_sizes(
return scheduler_block_size, hash_block_size
def resolve_cache_hit_alignment_tokens(
kv_cache_config: KVCacheConfig,
vllm_config: VllmConfig,
scheduler_block_size: int,
hash_block_size: int,
) -> int:
"""Token granularity at which prefix-cache hits land.
Mirrors the alignment the hybrid KV cache coordinator pushes to its
managers: fine-grained hits land on ``hash_block_size``, all others on
``scheduler_block_size``.
"""
groups = kv_cache_config.kv_cache_groups
if not vllm_config.cache_config.enable_prefix_caching or len(groups) <= 1:
return scheduler_block_size
manager_classes = []
for group in groups:
manager_cls = KVCacheSpecRegistry.get_manager_class(
group.kv_cache_spec, group.role
)
assert manager_cls is not None
manager_classes.append(manager_cls)
if partial_hash_hits_enabled(
groups,
hash_block_size,
vllm_config.parallel_config.decode_context_parallel_size,
manager_classes,
):
return hash_block_size
return scheduler_block_size
def get_request_block_hasher(
hash_block_size: int,
caching_hash_fn: Callable[[Any], bytes],
vllm/v1/core/sched/scheduler.py
+17 -18
View File
@@ -381,9 +381,8 @@ class Scheduler(SchedulerInterface):
or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0
for group in kv_cache_config.kv_cache_groups
)
# A finer prefix_match_unit is configured: a mamba partial tail entry
# can only be registered by a step ending exactly at the prompt's last
# hash boundary, so the split adds that stop.
# Fine-grained hits reuse the prompt's hash-aligned Mamba checkpoint.
# Split there only if the backend cannot save it within the final chunk.
self.mamba_partial_cache_hit = (
self.need_mamba_block_aligned_split
and self.hash_block_size < self.block_size
@@ -460,14 +459,19 @@ class Scheduler(SchedulerInterface):
# The last block-aligned position whose state can be cached. With
# Eagle, FullAttn prunes the last matching block, so back off one
# block to avoid a Mamba cache miss.
last_cache_position = request.num_tokens - request.num_tokens % block_size
if self.use_eagle_block_drop:
last_cache_position = max(last_cache_position - block_size, 0)
ckpt_block_size = (
self.hash_block_size if self.mamba_partial_cache_hit else self.block_size
)
last_cache_position = get_mamba_prefill_checkpoint_position(
request.num_tokens,
block_size,
self.use_eagle_block_drop,
)
end = start + num_new_tokens
checkpoint_position = get_mamba_prefill_checkpoint_position(
prefill_end,
self.hash_block_size,
ckpt_block_size,
drop_eagle_block=self.use_eagle_block_drop,
)
use_internal_checkpoint = (
@@ -500,18 +504,14 @@ class Scheduler(SchedulerInterface):
next_block_boundary = (start // block_size + 1) * block_size
tail_boundary = (
request.num_prompt_tokens // self.hash_block_size * self.hash_block_size
get_mamba_prefill_checkpoint_position(
request.num_prompt_tokens,
self.hash_block_size,
self.use_eagle_block_drop,
)
if self.mamba_partial_cache_hit and not use_internal_checkpoint
else 0
)
if tail_boundary and self.use_eagle_block_drop:
# Eagle matches one hash unit past the candidate and drops it, so
# nothing proves the prompt's own last hash boundary. Materialize
# the state one unit lower, where the hit can actually land. Keyed on
# the block-drop bit, not plain use_eagle: this shift exists only to
# compensate for the drop, and the Mamba manager's matching gate
# reads the same bit (the coordinator is handed use_eagle_block_drop).
tail_boundary = max(tail_boundary - self.hash_block_size, 0)
junction = request.shared_prefix_boundary
# Block-floored: a sub-block junction's state is not separately cacheable.
block_floored = start + (junction - start) // block_size * block_size
@@ -532,8 +532,7 @@ class Scheduler(SchedulerInterface):
else 0,
# Never run past the last cacheable block boundary mid-chunk.
last_cache_position,
# Fine-grained hits: the prompt's partial-tail entry can only be
# registered by a chunk ending exactly at its last hash boundary.
# Stop at the prompt checkpoint unless the backend saves it internally.
tail_boundary
if last_cache_position < tail_boundary < request.num_prompt_tokens
else 0,
vllm/v1/core/single_type_kv_cache_manager.py
+34 -26
View File
@@ -876,7 +876,14 @@ class FullAttentionManager(SingleTypeKVCacheManager):
block are intentionally skipped.
"""
hash_block_size = self.block_pool.hash_block_size
boundary_tokens = request.num_prompt_tokens // hash_block_size * hash_block_size
# A resend matches at most `prompt_len - 1` tokens, but EAGLE lookups
# read the whole prompt before dropping a hash unit.
token_limit = (
request.num_prompt_tokens
if self.use_eagle
else request.num_prompt_tokens - 1
)
boundary_tokens = token_limit // hash_block_size * hash_block_size
if boundary_tokens == 0 or boundary_tokens > num_tokens:
return
if boundary_tokens % self.block_size == 0:
@@ -1788,7 +1795,7 @@ class MambaManager(SingleTypeKVCacheManager):
# is kept, so `prefill_end` skips reserving blocks for the others.
checkpoint_position = get_mamba_prefill_checkpoint_position(
num_tokens,
self.block_pool.hash_block_size,
self.cache_hit_alignment_tokens,
self.drop_eagle_checkpoint_block,
)
if not self._needs_internal_checkpoint(
@@ -2073,20 +2080,29 @@ class MambaManager(SingleTypeKVCacheManager):
checkpoint_position
== get_mamba_prefill_checkpoint_position(
request.num_prompt_tokens,
hash_block_size,
self.cache_hit_alignment_tokens,
self.drop_eagle_checkpoint_block,
)
)
is_retained_checkpoint = False
if checkpoint_position % self.block_size == 0:
block_idx = checkpoint_position // self.block_size - 1
retained = self.reachable_block_mask(
block_idx,
block_idx + 1,
self.cache_hit_alignment_tokens,
self.kv_cache_spec,
self.drop_eagle_checkpoint_block,
retention_interval=retention_interval,
)
is_retained_checkpoint = retained is None or retained[0]
if (
retention_interval == 0
and num_tokens < request.num_prompt_tokens
and not is_prompt_checkpoint
not is_prompt_checkpoint
and not is_retained_checkpoint
and checkpoint_position != request.shared_prefix_boundary
):
# retention_interval == 0 keeps this transient checkpoint
# request-local. The slot may carry a hash from this step's
# full-block pass; that must go too, since the checkpoint
# state is about to overwrite the block.
# The transient state overwrites this slot, so neither it nor
# the full-block hash registered earlier this step is reusable.
self.block_pool._maybe_evict_cached_block(checkpoint_block)
return None
if checkpoint_block.block_hash_num_tokens == checkpoint_position:
@@ -2099,21 +2115,13 @@ class MambaManager(SingleTypeKVCacheManager):
block_size=self.block_size,
replace_existing_hashes=True,
)
if self.block_size == hash_block_size:
if num_tokens <= 0 or num_tokens % hash_block_size != 0:
return None
if num_tokens % self.block_size == 0:
return None
if num_tokens % hash_block_size != 0:
return None
latest_prompt_hash_boundary = (
request.num_prompt_tokens // hash_block_size
) * hash_block_size
if self.drop_eagle_checkpoint_block:
# Eagle groups match one hash unit past the candidate and drop it,
# so register the tail one unit lower.
latest_prompt_hash_boundary = max(
latest_prompt_hash_boundary - hash_block_size, 0
)
latest_prompt_hash_boundary = get_mamba_prefill_checkpoint_position(
request.num_prompt_tokens,
hash_block_size,
self.drop_eagle_checkpoint_block,
)
# The junction is the other position a sibling resumes at: where one was
# observed to stop, and where the scheduler already ends a chunk. Bounded
# to the prompt chunk being computed -- during decode the target is the
@@ -2126,7 +2134,7 @@ class MambaManager(SingleTypeKVCacheManager):
):
return None
block_idx = num_tokens // self.block_size
block_idx = cdiv(num_tokens, self.block_size) - 1
blocks = self.req_to_blocks[request.request_id]
if block_idx >= len(blocks):
return None
@@ -2141,7 +2149,7 @@ class MambaManager(SingleTypeKVCacheManager):
kv_cache_group_id=self.kv_cache_group_id,
block_size=self.block_size,
)
if partial_hash is not None:
if partial_hash is not None and num_tokens % self.block_size:
self._partial_hit_reqs[request.request_id] = (block_idx, source_block)
self.num_cached_block[request.request_id] = block_idx
# Producer of this partial tail: the boundary state currently lives
vllm/v1/engine/core.py
+14
View File
@@ -53,6 +53,7 @@ from vllm.v1.core.kv_cache_utils import (
get_kv_cache_configs,
get_request_block_hasher,
init_none_hash,
resolve_cache_hit_alignment_tokens,
resolve_kv_cache_block_sizes,
update_kv_cache_capacity,
)
@@ -352,6 +353,19 @@ class EngineCore:
vllm_config.validate_block_size()
scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
scheduler_kv_cache_config, vllm_config
)
cache_hit_alignment_tokens = resolve_cache_hit_alignment_tokens(
scheduler_kv_cache_config,
vllm_config,
scheduler_block_size,
hash_block_size,
)
for kv_cache_config in kv_cache_configs:
kv_cache_config.hash_block_size = hash_block_size
kv_cache_config.cache_hit_alignment_tokens = cache_hit_alignment_tokens
self.model_executor.initialize_from_config(kv_cache_configs)
if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
self.model_executor.compile_or_warm_up_model()
vllm/v1/kv_cache_interface.py
+8 -1
View File
@@ -1127,7 +1127,10 @@ def get_mamba_prefill_checkpoint_position(
drop_eagle_block: bool,
) -> int:
"""Return the reusable Mamba checkpoint boundary for a prefill."""
checkpoint_position = (num_tokens - 1) // hash_block_size * hash_block_size
# Without EAGLE, leave one token to recompute on resend.
# EAGLE's block drop already leaves tokens to recompute.
proof_limit = num_tokens if drop_eagle_block else num_tokens - 1
checkpoint_position = proof_limit // hash_block_size * hash_block_size
if drop_eagle_block:
checkpoint_position -= hash_block_size
return max(checkpoint_position, 0)
@@ -1487,6 +1490,10 @@ class KVCacheConfig:
"""Resolved retention policy for local prefix-cache checkpoints."""
kv_cache_layout: str | None = None
"""The KV cache layout resolved by the engine core, adopted by all workers."""
hash_block_size: int | None = None
"""Tokens per prefix-cache block hash, resolved by the engine core."""
cache_hit_alignment_tokens: int | None = None
"""Token granularity of prefix-cache hits, resolved by the engine core."""
hisparse_host_num_blocks: int | None = None
"""Capacity of the dedicated HiSparse host-block manager, when enabled."""
vllm/v1/worker/gpu_worker.py
+7 -2
View File
@@ -821,10 +821,15 @@ class Worker(WorkerBase):
# so that it's available to the warmup stage.
self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks
# Adopt the engine core's layout; workers spawned after resolution
# (e.g. elastic EP scale-up) only see it through the config.
# Adopt the engine core's layout and prefix-cache granularity; workers
# spawned after resolution (e.g. elastic EP scale-up) only see them
# through the config.
if kv_cache_config.kv_cache_layout is not None:
record_kv_cache_layout(self.cache_config, kv_cache_config.kv_cache_layout)
self.cache_config.hash_block_size = kv_cache_config.hash_block_size
self.cache_config.cache_hit_alignment_tokens = (
kv_cache_config.cache_hit_alignment_tokens
)
# Init kv cache connector here, because it requires
# `kv_cache_config`.