obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-24 06:30:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(kv-events): emit KV cache metadata (#40984)

Browse Source 
Signed-off-by: PeaBrane <[email protected]>
This commit is contained in:
Yan Ru Pei
2026-05-12 15:58:48 +00:00
committed by GitHub
parent c8a6e272e0
commit bcb9c133ba
8 changed files with 381 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/distributed/test_kv_cache_events.py
+14 -2
View File
@@ -9,7 +9,10 @@ from vllm.distributed.kv_events import BlockRemoved, BlockStored
_FAKE_HASH: bytes = b"\xab" * 32
def _make_block_stored(group_idx: int | None = None) -> BlockStored:
def _make_block_stored(
group_idx: int | None = None,
kv_cache_spec_sliding_window: int | None = None,
) -> BlockStored:
return BlockStored(
block_hashes=[_FAKE_HASH],
parent_block_hash=None,
@@ -19,10 +22,13 @@ def _make_block_stored(group_idx: int | None = None) -> BlockStored:
medium="GPU",
lora_name=None,
group_idx=group_idx,
kv_cache_spec_sliding_window=kv_cache_spec_sliding_window,
)
def _make_block_removed(group_idx: int | None = None) -> BlockRemoved:
def _make_block_removed(
group_idx: int | None = None,
) -> BlockRemoved:
return BlockRemoved(
block_hashes=[_FAKE_HASH],
medium="GPU",
@@ -72,3 +78,9 @@ def test_block_removed_hash_same_for_equal_group_idx():
event_a = _make_block_removed(group_idx=1)
event_b = _make_block_removed(group_idx=1)
assert hash(event_a) == hash(event_b)
def test_block_stored_hash_differs_by_sliding_window():
event_a = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=128)
event_b = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=256)
assert hash(event_a) != hash(event_b)
tests/v1/core/test_kv_cache_utils.py
+148
View File
@@ -43,11 +43,16 @@ from vllm.v1.kv_cache_interface import (
KVCacheConfig,
KVCacheGroupSpec,
KVCacheSpec,
KVCacheSpecKind,
KVCacheTensor,
MambaSpec,
MLAAttentionSpec,
SinkFullAttentionSpec,
SlidingWindowMLASpec,
SlidingWindowSpec,
UniformTypeKVCacheSpecs,
get_kv_cache_spec_kind,
get_kv_cache_spec_sliding_window,
)
from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
from vllm.v1.request import Request
@@ -1865,6 +1870,149 @@ def new_mla_spec(cache_dtype_str=None):
)
def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses():
assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION
sliding_window_mla_spec = SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
sliding_window=128,
)
assert (
get_kv_cache_spec_kind(sliding_window_mla_spec)
== KVCacheSpecKind.SLIDING_WINDOW_MLA
)
sink_full_attention_spec = SinkFullAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sink_len=4,
)
assert (
get_kv_cache_spec_kind(sink_full_attention_spec)
== KVCacheSpecKind.SINK_FULL_ATTENTION
)
def test_get_kv_cache_spec_kind_unwraps_uniform_type_specs():
uniform_mla_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": new_mla_spec(),
"layer_2": new_mla_spec(cache_dtype_str="fp8"),
},
)
assert get_kv_cache_spec_kind(uniform_mla_spec) == KVCacheSpecKind.MLA_ATTENTION
uniform_swa_mla_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
sliding_window=128,
),
"layer_2": SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=1024,
dtype=torch.float32,
sliding_window=128,
),
},
)
assert (
get_kv_cache_spec_kind(uniform_swa_mla_spec)
== KVCacheSpecKind.SLIDING_WINDOW_MLA
)
def test_get_kv_cache_spec_kind_unknown_for_mixed_uniform_type_specs():
uniform_mixed_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": new_mla_spec(),
"layer_2": SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
sliding_window=128,
),
},
)
assert get_kv_cache_spec_kind(uniform_mixed_spec) == KVCacheSpecKind.UNKNOWN
def test_get_kv_cache_spec_sliding_window_reads_windowed_specs():
full_attention_spec = FullAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
)
sliding_window_spec = SlidingWindowSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sliding_window=128,
)
assert get_kv_cache_spec_sliding_window(full_attention_spec) is None
assert get_kv_cache_spec_sliding_window(sliding_window_spec) == 128
def test_get_kv_cache_spec_sliding_window_unwraps_uniform_type_specs():
uniform_window_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": SlidingWindowSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sliding_window=128,
),
"layer_2": SlidingWindowSpec(
block_size=16,
num_kv_heads=2,
head_size=64,
dtype=torch.float32,
sliding_window=128,
),
},
)
mixed_window_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": SlidingWindowSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sliding_window=128,
),
"layer_2": SlidingWindowSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sliding_window=256,
),
},
)
assert get_kv_cache_spec_sliding_window(uniform_window_spec) == 128
assert get_kv_cache_spec_sliding_window(mixed_window_spec) is None
def test_merge_mla_spec():
kv_cache_specs = [
new_mla_spec(),
tests/v1/core/test_prefix_caching.py
+98 -9
View File
@@ -8,6 +8,7 @@ from collections.abc import Callable
import pytest
import torch
import vllm.v1.core.kv_cache_manager as kv_cache_manager
import vllm.v1.core.kv_cache_utils as kv_cache_utils
from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved, BlockStored
from vllm.lora.request import LoRARequest
@@ -35,6 +36,7 @@ from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheSpecKind,
MambaSpec,
SlidingWindowSpec,
)
@@ -1933,6 +1935,7 @@ def test_kv_cache_events(blocks_to_cache: int):
== len(manager.block_pool.cached_block_hash_to_block)
)
assert len(block.token_ids) == block.block_size * len(block.block_hashes)
assert block.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
assert len(manager.block_pool.kv_event_queue) == 0
stored_block_hash = block.block_hashes
@@ -1946,6 +1949,7 @@ def test_kv_cache_events(blocks_to_cache: int):
events = manager.take_events()
for blocks in events[:-1]:
assert isinstance(blocks, BlockRemoved)
assert blocks.block_hashes[0] in stored_block_hash
assert len(events) == blocks_to_cache + 1
assert isinstance(events[-2], BlockRemoved)
@@ -2022,6 +2026,8 @@ def test_null_parent_block_hash():
]
assert event.block_hashes == expected_new_hashes
assert event.group_idx == kv_cache_group_id
assert event.kv_cache_spec_kind is None
assert event.kv_cache_spec_sliding_window is None
# Ensure we didn't accidentally assign a hash to the null block.
assert pool.null_block.block_hash is None
@@ -2095,12 +2101,14 @@ def test_block_stored_event_group_idx(group_id: int):
block_size = 4
num_tokens = block_size * 2
pool = BlockPool(
num_gpu_blocks=5,
manager = KVCacheManager(
make_kv_cache_config_three_types(block_size, num_blocks=5),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
enable_kv_cache_events=True,
hash_block_size=block_size,
)
pool = manager.block_pool
req = make_request(
"req_grp_idx",
@@ -2119,10 +2127,26 @@ def test_block_stored_event_group_idx(group_id: int):
kv_cache_group_id=group_id,
)
events = pool.take_events()
events = manager.take_events()
assert len(events) == 1
assert isinstance(events[0], BlockStored)
assert events[0].group_idx == group_id
assert (
events[0].kv_cache_spec_kind
== [
KVCacheSpecKind.FULL_ATTENTION.value,
KVCacheSpecKind.SLIDING_WINDOW.value,
KVCacheSpecKind.MAMBA.value,
][group_id]
)
assert (
events[0].kv_cache_spec_sliding_window
== [
None,
2 * block_size,
None,
][group_id]
)
def test_block_stored_event_group_idx_multiple_groups():
@@ -2137,13 +2161,38 @@ def test_block_stored_event_group_idx_multiple_groups():
block_size = 4
num_tokens = block_size * 2
# null block + 4 usable (2 per group)
pool = BlockPool(
num_gpu_blocks=5,
manager = KVCacheManager(
KVCacheConfig(
num_blocks=5,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
["layer1"],
FullAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
),
KVCacheGroupSpec(
["layer2"],
SlidingWindowSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
sliding_window=128,
),
),
],
),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
enable_kv_cache_events=True,
hash_block_size=block_size,
)
pool = manager.block_pool
req = make_request(
"req_multi_grp",
@@ -2174,12 +2223,52 @@ def test_block_stored_event_group_idx_multiple_groups():
kv_cache_group_id=1,
)
events = pool.take_events()
events = manager.take_events()
assert len(events) == 2
assert isinstance(events[0], BlockStored)
assert events[0].group_idx == 0
assert events[0].kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
assert events[0].kv_cache_spec_sliding_window is None
assert isinstance(events[1], BlockStored)
assert events[1].group_idx == 1
assert events[1].kv_cache_spec_kind == KVCacheSpecKind.SLIDING_WINDOW.value
assert events[1].kv_cache_spec_sliding_window == 128
def test_block_stored_event_group_idx_out_of_bounds(monkeypatch):
"""Out-of-range group_idx events are returned without metadata annotation."""
block_size = 4
manager = KVCacheManager(
make_kv_cache_config(block_size, num_blocks=5),
max_model_len=8192,
enable_caching=True,
enable_kv_cache_events=True,
hash_block_size=block_size,
)
event = BlockStored(
block_hashes=[1],
parent_block_hash=None,
token_ids=list(range(block_size)),
block_size=block_size,
lora_id=None,
medium=None,
lora_name=None,
group_idx=1,
)
manager.block_pool.kv_event_queue.append(event)
warnings = []
def collect_warning(message, *args, **kwargs):
del kwargs
warnings.append(message % args if args else message)
monkeypatch.setattr(kv_cache_manager.logger, "warning", collect_warning)
events = manager.take_events()
assert events == [event]
assert event.kv_cache_spec_kind is None
assert event.kv_cache_spec_sliding_window is None
assert warnings == ["Group index `1` not in KV cache metadata"]
@pytest.mark.parametrize("group_id", [0, 1, 2])
vllm/distributed/kv_events.py
+6
View File
@@ -68,6 +68,10 @@ class BlockStored(KVCacheEvent):
"""
group_idx: int | None = None
# Store events carry cache-spec metadata so consumers can classify and
# filter groups as they are learned. Remove events only need group_idx+hash.
kv_cache_spec_kind: str | None = None
kv_cache_spec_sliding_window: int | None = None
def __hash__(self) -> int:
return hash(
@@ -80,6 +84,8 @@ class BlockStored(KVCacheEvent):
self.medium,
tuple(self.extra_keys) if self.extra_keys else None,
self.group_idx,
self.kv_cache_spec_kind,
self.kv_cache_spec_sliding_window,
)
)
vllm/v1/core/kv_cache_coordinator.py
+5 -5
View File
@@ -48,11 +48,11 @@ class KVCacheCoordinator(ABC):
self.enable_caching = enable_caching
self.block_pool = BlockPool(
kv_cache_config.num_blocks,
enable_caching,
hash_block_size,
enable_kv_cache_events,
metrics_collector,
num_gpu_blocks=kv_cache_config.num_blocks,
enable_caching=enable_caching,
hash_block_size=hash_block_size,
enable_kv_cache_events=enable_kv_cache_events,
metrics_collector=metrics_collector,
)
# KV cache group indices that get the EAGLE last-block drop.
vllm/v1/core/kv_cache_manager.py
+32 -3
View File
@@ -6,12 +6,16 @@ from collections.abc import Sequence
from dataclasses import dataclass
from typing import Literal, overload
from vllm.distributed.kv_events import KVCacheEvent
from vllm.distributed.kv_events import BlockStored, KVCacheEvent
from vllm.logger import init_logger
from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
from vllm.v1.core.kv_cache_utils import KVCacheBlock
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.kv_cache_interface import (
KVCacheConfig,
get_kv_cache_spec_kind,
get_kv_cache_spec_sliding_window,
)
from vllm.v1.metrics.stats import PrefixCacheStats
from vllm.v1.request import Request
@@ -149,6 +153,13 @@ class KVCacheManager:
self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
self.block_pool = self.coordinator.block_pool
self.kv_cache_config = kv_cache_config
self.kv_cache_event_metadata = tuple(
(
get_kv_cache_spec_kind(group.kv_cache_spec).value,
get_kv_cache_spec_sliding_window(group.kv_cache_spec),
)
for group in kv_cache_config.kv_cache_groups
)
# Pre-constructed KVCacheBlocks with no blocks, callers should use this
# via create_kv_cache_blocks instead of creating new ones to avoid GC
@@ -502,7 +513,25 @@ class KVCacheManager:
Returns:
A list of KV cache events.
"""
return self.block_pool.take_events()
events = self.block_pool.take_events()
for event in events:
if not isinstance(event, BlockStored):
continue
if event.group_idx is None:
continue
if event.group_idx < 0 or event.group_idx >= len(
self.kv_cache_event_metadata
):
logger.warning(
"Group index `%s` not in KV cache metadata", event.group_idx
)
continue
# Annotate here so BlockPool can keep emitting structural cache
# events without owning semantic KV cache spec metadata.
kind, sliding_window = self.kv_cache_event_metadata[event.group_idx]
event.kv_cache_spec_kind = kind
event.kv_cache_spec_sliding_window = sliding_window
return events
def get_blocks(self, request_id: str) -> KVCacheBlocks:
"""Get the blocks of a request."""
vllm/v1/engine/core.py
+20 -1
View File
@@ -72,7 +72,7 @@ from vllm.v1.engine.utils import (
get_device_indices,
)
from vllm.v1.executor import Executor
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind
from vllm.v1.metrics.stats import SchedulerStats
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
@@ -312,6 +312,25 @@ class EngineCore:
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
return self.model_executor.supported_tasks
def get_kv_cache_group_metadata(self) -> list[dict[str, int | str | None]]:
"""Return msgspec-serializable metadata for scheduler KV cache groups."""
kv_cache_config = getattr(self.scheduler, "kv_cache_config", None)
if kv_cache_config is None:
return []
metadata: list[dict[str, int | str | None]] = []
for group_idx, group in enumerate(kv_cache_config.kv_cache_groups):
spec = group.kv_cache_spec
metadata.append(
{
"group_idx": group_idx,
"kind": get_kv_cache_spec_kind(spec).value,
"block_size": spec.block_size,
"sliding_window": getattr(spec, "sliding_window", None),
}
)
return metadata
def add_request(self, request: Request, request_wave: int = 0):
"""Add request to the scheduler.
vllm/v1/kv_cache_interface.py
+58 -1
View File
@@ -6,7 +6,7 @@ from __future__ import annotations
import copy
from collections import Counter
from dataclasses import dataclass, fields, replace
from enum import IntEnum
from enum import Enum, IntEnum
from math import prod
from typing import TYPE_CHECKING
@@ -78,6 +78,19 @@ def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
return get_kv_quant_mode(kv_cache_dtype).is_per_token_head
class KVCacheSpecKind(str, Enum):
FULL_ATTENTION = "full_attention"
MLA_ATTENTION = "mla_attention"
SLIDING_WINDOW = "sliding_window"
SLIDING_WINDOW_MLA = "sliding_window_mla"
MAMBA = "mamba"
CHUNKED_LOCAL_ATTENTION = "chunked_local_attention"
SINK_FULL_ATTENTION = "sink_full_attention"
ENCODER_ONLY_ATTENTION = "encoder_only_attention"
CROSS_ATTENTION = "cross_attention"
UNKNOWN = "unknown"
@dataclass(frozen=True)
class KVCacheSpec:
"""
@@ -732,6 +745,50 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
)
def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind:
if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
inner_kinds = {
get_kv_cache_spec_kind(spec)
for spec in kv_cache_spec.kv_cache_specs.values()
}
if len(inner_kinds) == 1:
return next(iter(inner_kinds))
return KVCacheSpecKind.UNKNOWN
# Keep subclass checks before base classes so specialized specs keep their
# more precise kind.
if isinstance(kv_cache_spec, SlidingWindowMLASpec):
return KVCacheSpecKind.SLIDING_WINDOW_MLA
if isinstance(kv_cache_spec, MLAAttentionSpec):
return KVCacheSpecKind.MLA_ATTENTION
if isinstance(kv_cache_spec, SinkFullAttentionSpec):
return KVCacheSpecKind.SINK_FULL_ATTENTION
if isinstance(kv_cache_spec, FullAttentionSpec):
return KVCacheSpecKind.FULL_ATTENTION
if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec):
return KVCacheSpecKind.CHUNKED_LOCAL_ATTENTION
if isinstance(kv_cache_spec, SlidingWindowSpec):
return KVCacheSpecKind.SLIDING_WINDOW
if isinstance(kv_cache_spec, MambaSpec):
return KVCacheSpecKind.MAMBA
if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
return KVCacheSpecKind.ENCODER_ONLY_ATTENTION
if isinstance(kv_cache_spec, CrossAttentionSpec):
return KVCacheSpecKind.CROSS_ATTENTION
return KVCacheSpecKind.UNKNOWN
def get_kv_cache_spec_sliding_window(kv_cache_spec: KVCacheSpec) -> int | None:
if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
inner_windows = {
get_kv_cache_spec_sliding_window(spec)
for spec in kv_cache_spec.kv_cache_specs.values()
}
return next(iter(inner_windows)) if len(inner_windows) == 1 else None
if isinstance(kv_cache_spec, SlidingWindowSpec):
return kv_cache_spec.sliding_window
return None
@dataclass
class KVCacheTensor:
"""