@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bcb9c133ba

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 20:15 UTC

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

Signed-off-by: PeaBrane <yanrpei@gmail.com>

Yan Ru Pei committed May 12, 2026 at 08:58 UTC bcb9c133bafed161b3f4d0ea391c985465429680

8 files changed +381 -21

tests/distributed/test_kv_cache_events.py

+14 -2

@@ -9,7 +9,10 @@ from vllm.distributed.kv_events import BlockRemoved, BlockStored

9 _FAKE_HASH: bytes = b"\xab" * 32

12 -def _make_block_stored(group_idx: int | None = None) -> BlockStored:

12 +def _make_block_stored(

13 + group_idx: int | None = None,

14 + kv_cache_spec_sliding_window: int | None = None,

15 +) -> BlockStored:

16 return BlockStored(

17 block_hashes=[_FAKE_HASH],

18 parent_block_hash=None,

@@ -19,10 +22,13 @@ def _make_block_stored(group_idx: int | None = None) -> BlockStored:

22 medium="GPU",

23 lora_name=None,

24 group_idx=group_idx,

25 + kv_cache_spec_sliding_window=kv_cache_spec_sliding_window,

26 )

25 -def _make_block_removed(group_idx: int | None = None) -> BlockRemoved:

29 +def _make_block_removed(

30 + group_idx: int | None = None,

31 +) -> BlockRemoved:

32 return BlockRemoved(

33 block_hashes=[_FAKE_HASH],

34 medium="GPU",

@@ -72,3 +78,9 @@ def test_block_removed_hash_same_for_equal_group_idx():

78 event_a = _make_block_removed(group_idx=1)

79 event_b = _make_block_removed(group_idx=1)

80 assert hash(event_a) == hash(event_b)

81 +

82 +

83 +def test_block_stored_hash_differs_by_sliding_window():

84 + event_a = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=128)

85 + event_b = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=256)

86 + assert hash(event_a) != hash(event_b)

tests/v1/core/test_kv_cache_utils.py

+148

@@ -43,11 +43,16 @@ from vllm.v1.kv_cache_interface import (

43 KVCacheConfig,

44 KVCacheGroupSpec,

45 KVCacheSpec,

46 + KVCacheSpecKind,

47 KVCacheTensor,

48 MambaSpec,

49 MLAAttentionSpec,

50 + SinkFullAttentionSpec,

51 + SlidingWindowMLASpec,

52 SlidingWindowSpec,

53 UniformTypeKVCacheSpecs,

54 + get_kv_cache_spec_kind,

55 + get_kv_cache_spec_sliding_window,

56 )

57 from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats

58 from vllm.v1.request import Request

@@ -1865,6 +1870,149 @@ def new_mla_spec(cache_dtype_str=None):

1870 )

1871

1872

1873 +def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses():

1874 + assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION

1875 +

1876 + sliding_window_mla_spec = SlidingWindowMLASpec(

1877 + block_size=16,

1878 + num_kv_heads=1,

1879 + head_size=576,

1880 + dtype=torch.float32,

1881 + sliding_window=128,

1882 + )

1883 + assert (

1884 + get_kv_cache_spec_kind(sliding_window_mla_spec)

1885 + == KVCacheSpecKind.SLIDING_WINDOW_MLA

1886 + )

1887 +

1888 + sink_full_attention_spec = SinkFullAttentionSpec(

1889 + block_size=16,

1890 + num_kv_heads=1,

1891 + head_size=64,

1892 + dtype=torch.float32,

1893 + sink_len=4,

1894 + )

1895 + assert (

1896 + get_kv_cache_spec_kind(sink_full_attention_spec)

1897 + == KVCacheSpecKind.SINK_FULL_ATTENTION

1898 + )

1899 +

1900 +

1901 +def test_get_kv_cache_spec_kind_unwraps_uniform_type_specs():

1902 + uniform_mla_spec = UniformTypeKVCacheSpecs(

1903 + block_size=16,

1904 + kv_cache_specs={

1905 + "layer_1": new_mla_spec(),

1906 + "layer_2": new_mla_spec(cache_dtype_str="fp8"),

1907 + },

1908 + )

1909 + assert get_kv_cache_spec_kind(uniform_mla_spec) == KVCacheSpecKind.MLA_ATTENTION

1910 +

1911 + uniform_swa_mla_spec = UniformTypeKVCacheSpecs(

1912 + block_size=16,

1913 + kv_cache_specs={

1914 + "layer_1": SlidingWindowMLASpec(

1915 + block_size=16,

1916 + num_kv_heads=1,

1917 + head_size=576,

1918 + dtype=torch.float32,

1919 + sliding_window=128,

1920 + ),

1921 + "layer_2": SlidingWindowMLASpec(

1922 + block_size=16,

1923 + num_kv_heads=1,

1924 + head_size=1024,

1925 + dtype=torch.float32,

1926 + sliding_window=128,

1927 + ),

1928 + },

1929 + )

1930 + assert (

1931 + get_kv_cache_spec_kind(uniform_swa_mla_spec)

1932 + == KVCacheSpecKind.SLIDING_WINDOW_MLA

1933 + )

1934 +

1935 +

1936 +def test_get_kv_cache_spec_kind_unknown_for_mixed_uniform_type_specs():

1937 + uniform_mixed_spec = UniformTypeKVCacheSpecs(

1938 + block_size=16,

1939 + kv_cache_specs={

1940 + "layer_1": new_mla_spec(),

1941 + "layer_2": SlidingWindowMLASpec(

1942 + block_size=16,

1943 + num_kv_heads=1,

1944 + head_size=576,

1945 + dtype=torch.float32,

1946 + sliding_window=128,

1947 + ),

1948 + },

1949 + )

1950 + assert get_kv_cache_spec_kind(uniform_mixed_spec) == KVCacheSpecKind.UNKNOWN

1951 +

1952 +

1953 +def test_get_kv_cache_spec_sliding_window_reads_windowed_specs():

1954 + full_attention_spec = FullAttentionSpec(

1955 + block_size=16,

1956 + num_kv_heads=1,

1957 + head_size=64,

1958 + dtype=torch.float32,

1959 + )

1960 + sliding_window_spec = SlidingWindowSpec(

1961 + block_size=16,

1962 + num_kv_heads=1,

1963 + head_size=64,

1964 + dtype=torch.float32,

1965 + sliding_window=128,

1966 + )

1967 +

1968 + assert get_kv_cache_spec_sliding_window(full_attention_spec) is None

1969 + assert get_kv_cache_spec_sliding_window(sliding_window_spec) == 128

1970 +

1971 +

1972 +def test_get_kv_cache_spec_sliding_window_unwraps_uniform_type_specs():

1973 + uniform_window_spec = UniformTypeKVCacheSpecs(

1974 + block_size=16,

1975 + kv_cache_specs={

1976 + "layer_1": SlidingWindowSpec(

1977 + block_size=16,

1978 + num_kv_heads=1,

1979 + head_size=64,

1980 + dtype=torch.float32,

1981 + sliding_window=128,

1982 + ),

1983 + "layer_2": SlidingWindowSpec(

1984 + block_size=16,

1985 + num_kv_heads=2,

1986 + head_size=64,

1987 + dtype=torch.float32,

1988 + sliding_window=128,

1989 + ),

1990 + },

1991 + )

1992 + mixed_window_spec = UniformTypeKVCacheSpecs(

1993 + block_size=16,

1994 + kv_cache_specs={

1995 + "layer_1": SlidingWindowSpec(

1996 + block_size=16,

1997 + num_kv_heads=1,

1998 + head_size=64,

1999 + dtype=torch.float32,

2000 + sliding_window=128,

2001 + ),

2002 + "layer_2": SlidingWindowSpec(

2003 + block_size=16,

2004 + num_kv_heads=1,

2005 + head_size=64,

2006 + dtype=torch.float32,

2007 + sliding_window=256,

2008 + ),

2009 + },

2010 + )

2011 +

2012 + assert get_kv_cache_spec_sliding_window(uniform_window_spec) == 128

2013 + assert get_kv_cache_spec_sliding_window(mixed_window_spec) is None

2014 +

2015 +

2016 def test_merge_mla_spec():

2017 kv_cache_specs = [

2018 new_mla_spec(),

tests/v1/core/test_prefix_caching.py

+98 -9

@@ -8,6 +8,7 @@ from collections.abc import Callable

8 import pytest

9 import torch

11 +import vllm.v1.core.kv_cache_manager as kv_cache_manager

12 import vllm.v1.core.kv_cache_utils as kv_cache_utils

13 from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved, BlockStored

14 from vllm.lora.request import LoRARequest

@@ -35,6 +36,7 @@ from vllm.v1.kv_cache_interface import (

36 FullAttentionSpec,

37 KVCacheConfig,

38 KVCacheGroupSpec,

39 + KVCacheSpecKind,

40 MambaSpec,

41 SlidingWindowSpec,

42 )

@@ -1933,6 +1935,7 @@ def test_kv_cache_events(blocks_to_cache: int):

1935 == len(manager.block_pool.cached_block_hash_to_block)

1936 )

1937 assert len(block.token_ids) == block.block_size * len(block.block_hashes)

1938 + assert block.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value

1939 assert len(manager.block_pool.kv_event_queue) == 0

1940

1941 stored_block_hash = block.block_hashes

@@ -1946,6 +1949,7 @@ def test_kv_cache_events(blocks_to_cache: int):

1949 events = manager.take_events()

1950

1951 for blocks in events[:-1]:

1952 + assert isinstance(blocks, BlockRemoved)

1953 assert blocks.block_hashes[0] in stored_block_hash

1954 assert len(events) == blocks_to_cache + 1

1955 assert isinstance(events[-2], BlockRemoved)

@@ -2022,6 +2026,8 @@ def test_null_parent_block_hash():

2026 ]

2027 assert event.block_hashes == expected_new_hashes

2028 assert event.group_idx == kv_cache_group_id

2029 + assert event.kv_cache_spec_kind is None

2030 + assert event.kv_cache_spec_sliding_window is None

2031

2032 # Ensure we didn't accidentally assign a hash to the null block.

2033 assert pool.null_block.block_hash is None

@@ -2095,12 +2101,14 @@ def test_block_stored_event_group_idx(group_id: int):

2101 block_size = 4

2102 num_tokens = block_size * 2

2103

2098 - pool = BlockPool(

2099 - num_gpu_blocks=5,

2104 + manager = KVCacheManager(

2105 + make_kv_cache_config_three_types(block_size, num_blocks=5),

2106 + max_model_len=8192,

2107 enable_caching=True,

2101 - hash_block_size=block_size,

2108 enable_kv_cache_events=True,

2109 + hash_block_size=block_size,

2110 )

2111 + pool = manager.block_pool

2112

2113 req = make_request(

2114 "req_grp_idx",

@@ -2119,10 +2127,26 @@ def test_block_stored_event_group_idx(group_id: int):

2127 kv_cache_group_id=group_id,

2128 )

2129

2122 - events = pool.take_events()

2130 + events = manager.take_events()

2131 assert len(events) == 1

2132 assert isinstance(events[0], BlockStored)

2133 assert events[0].group_idx == group_id

2134 + assert (

2135 + events[0].kv_cache_spec_kind

2136 + == [

2137 + KVCacheSpecKind.FULL_ATTENTION.value,

2138 + KVCacheSpecKind.SLIDING_WINDOW.value,

2139 + KVCacheSpecKind.MAMBA.value,

2140 + ][group_id]

2141 + )

2142 + assert (

2143 + events[0].kv_cache_spec_sliding_window

2144 + == [

2145 + None,

2146 + 2 * block_size,

2147 + None,

2148 + ][group_id]

2149 + )

2150

2151

2152 def test_block_stored_event_group_idx_multiple_groups():

@@ -2137,13 +2161,38 @@ def test_block_stored_event_group_idx_multiple_groups():

2161 block_size = 4

2162 num_tokens = block_size * 2

2163

2140 - # null block + 4 usable (2 per group)

2141 - pool = BlockPool(

2142 - num_gpu_blocks=5,

2164 + manager = KVCacheManager(

2165 + KVCacheConfig(

2166 + num_blocks=5,

2167 + kv_cache_tensors=[],

2168 + kv_cache_groups=[

2169 + KVCacheGroupSpec(

2170 + ["layer1"],

2171 + FullAttentionSpec(

2172 + block_size=block_size,

2173 + num_kv_heads=1,

2174 + head_size=1,

2175 + dtype=torch.float32,

2176 + ),

2177 + ),

2178 + KVCacheGroupSpec(

2179 + ["layer2"],

2180 + SlidingWindowSpec(

2181 + block_size=block_size,

2182 + num_kv_heads=1,

2183 + head_size=1,

2184 + dtype=torch.float32,

2185 + sliding_window=128,

2186 + ),

2187 + ),

2188 + ],

2189 + ),

2190 + max_model_len=8192,

2191 enable_caching=True,

2144 - hash_block_size=block_size,

2192 enable_kv_cache_events=True,

2193 + hash_block_size=block_size,

2194 )

2195 + pool = manager.block_pool

2196

2197 req = make_request(

2198 "req_multi_grp",

@@ -2174,12 +2223,52 @@ def test_block_stored_event_group_idx_multiple_groups():

2223 kv_cache_group_id=1,

2224 )

2225

2177 - events = pool.take_events()

2226 + events = manager.take_events()

2227 assert len(events) == 2

2228 assert isinstance(events[0], BlockStored)

2229 assert events[0].group_idx == 0

2230 + assert events[0].kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value

2231 + assert events[0].kv_cache_spec_sliding_window is None

2232 assert isinstance(events[1], BlockStored)

2233 assert events[1].group_idx == 1

2234 + assert events[1].kv_cache_spec_kind == KVCacheSpecKind.SLIDING_WINDOW.value

2235 + assert events[1].kv_cache_spec_sliding_window == 128

2236 +

2237 +

2238 +def test_block_stored_event_group_idx_out_of_bounds(monkeypatch):

2239 + """Out-of-range group_idx events are returned without metadata annotation."""

2240 + block_size = 4

2241 + manager = KVCacheManager(

2242 + make_kv_cache_config(block_size, num_blocks=5),

2243 + max_model_len=8192,

2244 + enable_caching=True,

2245 + enable_kv_cache_events=True,

2246 + hash_block_size=block_size,

2247 + )

2248 + event = BlockStored(

2249 + block_hashes=[1],

2250 + parent_block_hash=None,

2251 + token_ids=list(range(block_size)),

2252 + block_size=block_size,

2253 + lora_id=None,

2254 + medium=None,

2255 + lora_name=None,

2256 + group_idx=1,

2257 + )

2258 + manager.block_pool.kv_event_queue.append(event)

2259 + warnings = []

2260 +

2261 + def collect_warning(message, *args, **kwargs):

2262 + del kwargs

2263 + warnings.append(message % args if args else message)

2264 +

2265 + monkeypatch.setattr(kv_cache_manager.logger, "warning", collect_warning)

2266 + events = manager.take_events()

2267 +

2268 + assert events == [event]

2269 + assert event.kv_cache_spec_kind is None

2270 + assert event.kv_cache_spec_sliding_window is None

2271 + assert warnings == ["Group index `1` not in KV cache metadata"]

2272

2273

2274 @pytest.mark.parametrize("group_id", [0, 1, 2])

vllm/distributed/kv_events.py

@@ -68,6 +68,10 @@ class BlockStored(KVCacheEvent):

68 """

70 group_idx: int | None = None

71 + # Store events carry cache-spec metadata so consumers can classify and

72 + # filter groups as they are learned. Remove events only need group_idx+hash.

73 + kv_cache_spec_kind: str | None = None

74 + kv_cache_spec_sliding_window: int | None = None

76 def __hash__(self) -> int:

77 return hash(

@@ -80,6 +84,8 @@ class BlockStored(KVCacheEvent):

84 self.medium,

85 tuple(self.extra_keys) if self.extra_keys else None,

86 self.group_idx,

87 + self.kv_cache_spec_kind,

88 + self.kv_cache_spec_sliding_window,

89 )

90 )

vllm/v1/core/kv_cache_coordinator.py

+5 -5

@@ -48,11 +48,11 @@ class KVCacheCoordinator(ABC):

48 self.enable_caching = enable_caching

50 self.block_pool = BlockPool(

51 - kv_cache_config.num_blocks,

52 - enable_caching,

53 - hash_block_size,

54 - enable_kv_cache_events,

55 - metrics_collector,

51 + num_gpu_blocks=kv_cache_config.num_blocks,

52 + enable_caching=enable_caching,

53 + hash_block_size=hash_block_size,

54 + enable_kv_cache_events=enable_kv_cache_events,

55 + metrics_collector=metrics_collector,

56 )

58 # KV cache group indices that get the EAGLE last-block drop.

vllm/v1/core/kv_cache_manager.py

+32 -3

@@ -6,12 +6,16 @@ from collections.abc import Sequence

6 from dataclasses import dataclass

7 from typing import Literal, overload

9 -from vllm.distributed.kv_events import KVCacheEvent

9 +from vllm.distributed.kv_events import BlockStored, KVCacheEvent

10 from vllm.logger import init_logger

11 from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator

12 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector

13 from vllm.v1.core.kv_cache_utils import KVCacheBlock

14 -from vllm.v1.kv_cache_interface import KVCacheConfig

14 +from vllm.v1.kv_cache_interface import (

15 + KVCacheConfig,

16 + get_kv_cache_spec_kind,

17 + get_kv_cache_spec_sliding_window,

18 +)

19 from vllm.v1.metrics.stats import PrefixCacheStats

20 from vllm.v1.request import Request

@@ -149,6 +153,13 @@ class KVCacheManager:

153 self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)

154 self.block_pool = self.coordinator.block_pool

155 self.kv_cache_config = kv_cache_config

156 + self.kv_cache_event_metadata = tuple(

157 + (

158 + get_kv_cache_spec_kind(group.kv_cache_spec).value,

159 + get_kv_cache_spec_sliding_window(group.kv_cache_spec),

160 + )

161 + for group in kv_cache_config.kv_cache_groups

162 + )

163

164 # Pre-constructed KVCacheBlocks with no blocks, callers should use this

165 # via create_kv_cache_blocks instead of creating new ones to avoid GC

@@ -502,7 +513,25 @@ class KVCacheManager:

513 Returns:

514 A list of KV cache events.

515 """

505 - return self.block_pool.take_events()

516 + events = self.block_pool.take_events()

517 + for event in events:

518 + if not isinstance(event, BlockStored):

519 + continue

520 + if event.group_idx is None:

521 + continue

522 + if event.group_idx < 0 or event.group_idx >= len(

523 + self.kv_cache_event_metadata

524 + ):

525 + logger.warning(

526 + "Group index `%s` not in KV cache metadata", event.group_idx

527 + )

528 + continue

529 + # Annotate here so BlockPool can keep emitting structural cache

530 + # events without owning semantic KV cache spec metadata.

531 + kind, sliding_window = self.kv_cache_event_metadata[event.group_idx]

532 + event.kv_cache_spec_kind = kind

533 + event.kv_cache_spec_sliding_window = sliding_window

534 + return events

535

536 def get_blocks(self, request_id: str) -> KVCacheBlocks:

537 """Get the blocks of a request."""

vllm/v1/engine/core.py

+20 -1

@@ -72,7 +72,7 @@ from vllm.v1.engine.utils import (

72 get_device_indices,

73 )

74 from vllm.v1.executor import Executor

75 -from vllm.v1.kv_cache_interface import KVCacheConfig

75 +from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind

76 from vllm.v1.metrics.stats import SchedulerStats

77 from vllm.v1.outputs import ModelRunnerOutput

78 from vllm.v1.request import Request, RequestStatus

@@ -312,6 +312,25 @@ class EngineCore:

312 def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

313 return self.model_executor.supported_tasks

314

315 + def get_kv_cache_group_metadata(self) -> list[dict[str, int | str | None]]:

316 + """Return msgspec-serializable metadata for scheduler KV cache groups."""

317 + kv_cache_config = getattr(self.scheduler, "kv_cache_config", None)

318 + if kv_cache_config is None:

319 + return []

320 +

321 + metadata: list[dict[str, int | str | None]] = []

322 + for group_idx, group in enumerate(kv_cache_config.kv_cache_groups):

323 + spec = group.kv_cache_spec

324 + metadata.append(

325 + {

326 + "group_idx": group_idx,

327 + "kind": get_kv_cache_spec_kind(spec).value,

328 + "block_size": spec.block_size,

329 + "sliding_window": getattr(spec, "sliding_window", None),

330 + }

331 + )

332 + return metadata

333 +

334 def add_request(self, request: Request, request_wave: int = 0):

335 """Add request to the scheduler.

336

vllm/v1/kv_cache_interface.py

+58 -1

@@ -6,7 +6,7 @@ from __future__ import annotations

6 import copy

7 from collections import Counter

8 from dataclasses import dataclass, fields, replace

9 -from enum import IntEnum

9 +from enum import Enum, IntEnum

10 from math import prod

11 from typing import TYPE_CHECKING

@@ -78,6 +78,19 @@ def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:

78 return get_kv_quant_mode(kv_cache_dtype).is_per_token_head

81 +class KVCacheSpecKind(str, Enum):

82 + FULL_ATTENTION = "full_attention"

83 + MLA_ATTENTION = "mla_attention"

84 + SLIDING_WINDOW = "sliding_window"

85 + SLIDING_WINDOW_MLA = "sliding_window_mla"

86 + MAMBA = "mamba"

87 + CHUNKED_LOCAL_ATTENTION = "chunked_local_attention"

88 + SINK_FULL_ATTENTION = "sink_full_attention"

89 + ENCODER_ONLY_ATTENTION = "encoder_only_attention"

90 + CROSS_ATTENTION = "cross_attention"

91 + UNKNOWN = "unknown"

92 +

93 +

94 @dataclass(frozen=True)

95 class KVCacheSpec:

96 """

@@ -732,6 +745,50 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):

745 )

746

747

748 +def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind:

749 + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):

750 + inner_kinds = {

751 + get_kv_cache_spec_kind(spec)

752 + for spec in kv_cache_spec.kv_cache_specs.values()

753 + }

754 + if len(inner_kinds) == 1:

755 + return next(iter(inner_kinds))

756 + return KVCacheSpecKind.UNKNOWN

757 + # Keep subclass checks before base classes so specialized specs keep their

758 + # more precise kind.

759 + if isinstance(kv_cache_spec, SlidingWindowMLASpec):

760 + return KVCacheSpecKind.SLIDING_WINDOW_MLA

761 + if isinstance(kv_cache_spec, MLAAttentionSpec):

762 + return KVCacheSpecKind.MLA_ATTENTION

763 + if isinstance(kv_cache_spec, SinkFullAttentionSpec):

764 + return KVCacheSpecKind.SINK_FULL_ATTENTION

765 + if isinstance(kv_cache_spec, FullAttentionSpec):

766 + return KVCacheSpecKind.FULL_ATTENTION

767 + if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec):

768 + return KVCacheSpecKind.CHUNKED_LOCAL_ATTENTION

769 + if isinstance(kv_cache_spec, SlidingWindowSpec):

770 + return KVCacheSpecKind.SLIDING_WINDOW

771 + if isinstance(kv_cache_spec, MambaSpec):

772 + return KVCacheSpecKind.MAMBA

773 + if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):

774 + return KVCacheSpecKind.ENCODER_ONLY_ATTENTION

775 + if isinstance(kv_cache_spec, CrossAttentionSpec):

776 + return KVCacheSpecKind.CROSS_ATTENTION

777 + return KVCacheSpecKind.UNKNOWN

778 +

779 +

780 +def get_kv_cache_spec_sliding_window(kv_cache_spec: KVCacheSpec) -> int | None:

781 + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):

782 + inner_windows = {

783 + get_kv_cache_spec_sliding_window(spec)

784 + for spec in kv_cache_spec.kv_cache_specs.values()

785 + }

786 + return next(iter(inner_windows)) if len(inner_windows) == 1 else None

787 + if isinstance(kv_cache_spec, SlidingWindowSpec):

788 + return kv_cache_spec.sliding_window

789 + return None

790 +

791 +

792 @dataclass

793 class KVCacheTensor:

794 """