diff --git a/tests/distributed/test_kv_cache_events.py b/tests/distributed/test_kv_cache_events.py index 57d1c9b546..aa39ab17b3 100644 --- a/tests/distributed/test_kv_cache_events.py +++ b/tests/distributed/test_kv_cache_events.py @@ -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) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 985b97c69c..94e9f6f4c1 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -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(), diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index c35c38911a..0a58fa2f4f 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -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]) diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index d3e304f8b6..ee21185969 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -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, ) ) diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 65993e8041..b5332ea0fd 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -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. diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 431776870c..9359d8843a 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -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.""" diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 07d3cca051..d19d928c70 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -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. diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index cf50dbff17..fa58297c42 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -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: """