| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass, replace |
| 5 | from typing import Any |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import ( |
| 11 | CacheConfig, |
| 12 | DeviceConfig, |
| 13 | VllmConfig, |
| 14 | ) |
| 15 | from vllm.v1.core.single_type_kv_cache_manager import ( |
| 16 | ChunkedLocalAttentionManager, |
| 17 | CrossAttentionManager, |
| 18 | FullAttentionManager, |
| 19 | HiSparseSourceManager, |
| 20 | MambaManager, |
| 21 | SingleTypeKVCacheManager, |
| 22 | SinkFullAttentionManager, |
| 23 | SlidingWindowManager, |
| 24 | register_all_kvcache_specs, |
| 25 | ) |
| 26 | from vllm.v1.kv_cache_interface import ( |
| 27 | ChunkedLocalAttentionSpec, |
| 28 | CrossAttentionSpec, |
| 29 | FullAttentionSpec, |
| 30 | HiddenStateCacheSpec, |
| 31 | KVCacheGroupRole, |
| 32 | KVCacheSpec, |
| 33 | KVCacheSpecKind, |
| 34 | MambaSpec, |
| 35 | MLAAttentionSpec, |
| 36 | SinkFullAttentionSpec, |
| 37 | SlidingWindowMLASpec, |
| 38 | SlidingWindowSpec, |
| 39 | UniformTypeKVCacheSpecs, |
| 40 | get_kv_cache_spec_kind, |
| 41 | ) |
| 42 | from vllm.v1.kv_cache_spec_registry import ( |
| 43 | _REGISTRY_KVCACHESPEC_LIST, |
| 44 | KVCacheSpecRegistry, |
| 45 | register_kv_cache_spec, |
| 46 | ) |
| 47 | |
| 48 | |
| 49 | def make_vllm_config() -> VllmConfig: |
| 50 | return VllmConfig( |
| 51 | cache_config=CacheConfig( |
| 52 | block_size=64, |
| 53 | cache_dtype="bfloat16", |
| 54 | ), |
| 55 | device_config=DeviceConfig(device="cpu"), |
| 56 | ) |
| 57 | |
| 58 | |
| 59 | vllm_config = make_vllm_config() |
| 60 | register_all_kvcache_specs(vllm_config) |
| 61 | |
| 62 | |
| 63 | @pytest.fixture(autouse=True) |
| 64 | def restore_kv_cache_spec_registry(): |
| 65 | registry = _REGISTRY_KVCACHESPEC_LIST.copy() |
| 66 | yield |
| 67 | _REGISTRY_KVCACHESPEC_LIST.clear() |
| 68 | _REGISTRY_KVCACHESPEC_LIST.update(registry) |
| 69 | |
| 70 | |
| 71 | @dataclass(frozen=True) |
| 72 | class _TrulyUnregisteredSpec(KVCacheSpec): |
| 73 | """A spec that inherits directly from KVCacheSpec with no registered |
| 74 | ancestor in the MRO. Used to test that the registry correctly raises |
| 75 | when no entry can be found. |
| 76 | """ |
| 77 | |
| 78 | @property |
| 79 | def page_size_bytes(self) -> int: |
| 80 | return self.block_size * 128 |
| 81 | |
| 82 | def max_memory_usage_bytes(self, _) -> int: |
| 83 | return 0 |
| 84 | |
| 85 | |
| 86 | spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = { |
| 87 | FullAttentionSpec: FullAttentionManager, |
| 88 | MLAAttentionSpec: FullAttentionManager, |
| 89 | HiddenStateCacheSpec: FullAttentionManager, |
| 90 | SlidingWindowSpec: SlidingWindowManager, |
| 91 | SlidingWindowMLASpec: SlidingWindowManager, |
| 92 | ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager, |
| 93 | MambaSpec: MambaManager, |
| 94 | CrossAttentionSpec: CrossAttentionManager, |
| 95 | SinkFullAttentionSpec: SinkFullAttentionManager, |
| 96 | } |
| 97 | |
| 98 | spec_uniform_base_map: dict[type[KVCacheSpec], type[KVCacheSpec]] = { |
| 99 | FullAttentionSpec: FullAttentionSpec, |
| 100 | MLAAttentionSpec: FullAttentionSpec, |
| 101 | HiddenStateCacheSpec: FullAttentionSpec, |
| 102 | SlidingWindowSpec: SlidingWindowSpec, |
| 103 | SlidingWindowMLASpec: SlidingWindowMLASpec, |
| 104 | ChunkedLocalAttentionSpec: ChunkedLocalAttentionSpec, |
| 105 | MambaSpec: MambaSpec, |
| 106 | CrossAttentionSpec: CrossAttentionSpec, |
| 107 | SinkFullAttentionSpec: FullAttentionSpec, |
| 108 | } |
| 109 | |
| 110 | spec_args_map: dict[type[KVCacheSpec], dict[str, Any]] = { |
| 111 | FullAttentionSpec: dict( |
| 112 | block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 |
| 113 | ), |
| 114 | MLAAttentionSpec: dict( |
| 115 | block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 |
| 116 | ), |
| 117 | HiddenStateCacheSpec: dict( |
| 118 | block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 |
| 119 | ), |
| 120 | SlidingWindowSpec: dict( |
| 121 | block_size=64, |
| 122 | num_kv_heads=8, |
| 123 | head_size=128, |
| 124 | dtype=torch.bfloat16, |
| 125 | sliding_window=128, |
| 126 | ), |
| 127 | SlidingWindowMLASpec: dict( |
| 128 | block_size=64, |
| 129 | num_kv_heads=1, |
| 130 | head_size=128, |
| 131 | dtype=torch.bfloat16, |
| 132 | sliding_window=128, |
| 133 | ), |
| 134 | ChunkedLocalAttentionSpec: dict( |
| 135 | block_size=64, |
| 136 | num_kv_heads=8, |
| 137 | head_size=128, |
| 138 | dtype=torch.bfloat16, |
| 139 | attention_chunk_size=4, |
| 140 | ), |
| 141 | MambaSpec: dict( |
| 142 | block_size=64, |
| 143 | shapes=((2, 512), (3, 32, 32)), |
| 144 | dtypes=(torch.float32, torch.float32), |
| 145 | mamba_cache_mode="align", |
| 146 | num_speculative_blocks=2, |
| 147 | ), |
| 148 | CrossAttentionSpec: dict( |
| 149 | block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 |
| 150 | ), |
| 151 | SinkFullAttentionSpec: dict( |
| 152 | block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16, sink_len=16 |
| 153 | ), |
| 154 | } |
| 155 | |
| 156 | |
| 157 | def make_spec(spec_cls: type[KVCacheSpec]) -> KVCacheSpec: |
| 158 | return spec_cls(**spec_args_map[spec_cls]) |
| 159 | |
| 160 | |
| 161 | def are_uniform_specs(*specs: KVCacheSpec) -> bool: |
| 162 | return UniformTypeKVCacheSpecs.is_uniform_type( |
| 163 | {f"layer_{i}": spec for i, spec in enumerate(specs)} |
| 164 | ) |
| 165 | |
| 166 | |
| 167 | class TestKVCacheSpecRegistry: |
| 168 | """Test the core registry functionality.""" |
| 169 | |
| 170 | def test_builtin_kvcache_specs_registered(self): |
| 171 | assert set(spec_manager_map) <= set(_REGISTRY_KVCACHESPEC_LIST) |
| 172 | for spec_cls, manager in spec_manager_map.items(): |
| 173 | spec = make_spec(spec_cls) |
| 174 | assert KVCacheSpecRegistry.get_manager_class(spec) is manager |
| 175 | assert ( |
| 176 | KVCacheSpecRegistry.get_uniform_type_base_spec(spec) |
| 177 | is spec_uniform_base_map[spec_cls] |
| 178 | ) |
| 179 | |
| 180 | @pytest.mark.parametrize("role", list(KVCacheGroupRole)) |
| 181 | def test_mla_manager_selection_by_role(self, role): |
| 182 | """Only the source role overrides ordinary MLA manager selection.""" |
| 183 | expected = ( |
| 184 | HiSparseSourceManager |
| 185 | if role == KVCacheGroupRole.HISPARSE_SOURCE |
| 186 | else FullAttentionManager |
| 187 | ) |
| 188 | assert ( |
| 189 | KVCacheSpecRegistry.get_manager_class(make_spec(MLAAttentionSpec), role) |
| 190 | is expected |
| 191 | ) |
| 192 | |
| 193 | @pytest.mark.parametrize("spec_cls", list(spec_manager_map)) |
| 194 | def test_custom_spec_register(self, spec_cls): |
| 195 | """A decorated custom spec resolves to the declared manager.""" |
| 196 | manager = spec_manager_map[spec_cls] |
| 197 | uniform_base_spec = spec_uniform_base_map[spec_cls] |
| 198 | |
| 199 | @register_kv_cache_spec( |
| 200 | manager_class=manager, |
| 201 | uniform_type_base_spec=uniform_base_spec, |
| 202 | ) |
| 203 | @dataclass(frozen=True, kw_only=True) |
| 204 | class _CustomSpec(spec_cls): # type: ignore[valid-type,misc] |
| 205 | custom_param: int = 16 |
| 206 | |
| 207 | spec = _CustomSpec(**spec_args_map[spec_cls], custom_param=100) |
| 208 | |
| 209 | assert KVCacheSpecRegistry.get_manager_class(spec) is manager |
| 210 | assert KVCacheSpecRegistry.get_uniform_type_base_spec(spec) is uniform_base_spec |
| 211 | |
| 212 | def test_custom_spec_register_requires_manager(self): |
| 213 | """Invalid register decorator arguments fail early.""" |
| 214 | with pytest.raises(AssertionError, match="manager_class is required"): |
| 215 | |
| 216 | @register_kv_cache_spec( |
| 217 | uniform_type_base_spec=FullAttentionSpec, |
| 218 | ) |
| 219 | @dataclass(frozen=True, kw_only=True) |
| 220 | class _CustomFullSpecWithoutManager(FullAttentionSpec): |
| 221 | custom_param: int = 16 |
| 222 | |
| 223 | def test_unregistered_spec_no_registered_parent_raises(self): |
| 224 | """A spec whose entire MRO contains no registered class resolves to None. |
| 225 | Runtime callers should use check_kv_cache_spec_registry to fail early. |
| 226 | Subclasses of registered specs intentionally do not fail — they inherit |
| 227 | their parent's manager via MRO walking. |
| 228 | """ |
| 229 | spec = _TrulyUnregisteredSpec(block_size=16) |
| 230 | |
| 231 | assert KVCacheSpecRegistry.get_manager_class(spec) is None |
| 232 | assert KVCacheSpecRegistry.get_uniform_type_base_spec(spec) is None |
| 233 | |
| 234 | with pytest.raises( |
| 235 | ValueError, match="Unsupported KV cache spec type for layer layer_0" |
| 236 | ): |
| 237 | KVCacheSpecRegistry.check_kv_cache_spec_registry({"layer_0": spec}) |
| 238 | |
| 239 | with pytest.raises(AssertionError, match="Unsupported KV cache spec type"): |
| 240 | UniformTypeKVCacheSpecs.is_uniform_type({"layer_0": spec}) |
| 241 | |
| 242 | def test_unregistered_subclass_inherits_parent_manager(self): |
| 243 | """An unregistered subclass of a registered spec resolves via MRO |
| 244 | to its parent's manager — this is intentional registry behaviour. |
| 245 | """ |
| 246 | |
| 247 | @dataclass(frozen=True, kw_only=True) |
| 248 | class _ImplicitlyInheritedSpec(FullAttentionSpec): |
| 249 | pass |
| 250 | |
| 251 | spec = _ImplicitlyInheritedSpec( |
| 252 | block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 |
| 253 | ) |
| 254 | |
| 255 | # MRO walk finds FullAttentionSpec → FullAttentionManager |
| 256 | assert KVCacheSpecRegistry.get_manager_class(spec) is FullAttentionManager |
| 257 | |
| 258 | @pytest.mark.parametrize("spec_cls", list(spec_manager_map)) |
| 259 | def test_builtin_specs_are_uniform_with_same_spec_type(self, spec_cls): |
| 260 | spec = make_spec(spec_cls) |
| 261 | assert are_uniform_specs(spec, replace(spec)) |
| 262 | |
| 263 | def test_full_attention_family_specs_are_uniform(self): |
| 264 | specs = [ |
| 265 | make_spec(FullAttentionSpec), |
| 266 | make_spec(MLAAttentionSpec), |
| 267 | make_spec(HiddenStateCacheSpec), |
| 268 | make_spec(SinkFullAttentionSpec), |
| 269 | ] |
| 270 | |
| 271 | assert are_uniform_specs(*specs) |
| 272 | |
| 273 | @pytest.mark.parametrize( |
| 274 | ("spec_cls", "field", "value"), |
| 275 | [ |
| 276 | (SlidingWindowSpec, "sliding_window", 256), |
| 277 | (SlidingWindowMLASpec, "sliding_window", 256), |
| 278 | (ChunkedLocalAttentionSpec, "attention_chunk_size", 8), |
| 279 | (MambaSpec, "num_speculative_blocks", 4), |
| 280 | ], |
| 281 | ) |
| 282 | def test_specs_with_type_specific_uniform_fields(self, spec_cls, field, value): |
| 283 | spec = make_spec(spec_cls) |
| 284 | changed_spec = replace(spec, **{field: value}) |
| 285 | |
| 286 | assert not are_uniform_specs(spec, changed_spec) |
| 287 | |
| 288 | @pytest.mark.parametrize( |
| 289 | ("left_cls", "right_cls"), |
| 290 | [ |
| 291 | (FullAttentionSpec, CrossAttentionSpec), |
| 292 | (FullAttentionSpec, SlidingWindowSpec), |
| 293 | (FullAttentionSpec, ChunkedLocalAttentionSpec), |
| 294 | (FullAttentionSpec, MambaSpec), |
| 295 | (SlidingWindowMLASpec, SlidingWindowSpec), |
| 296 | (ChunkedLocalAttentionSpec, SlidingWindowSpec), |
| 297 | (MambaSpec, CrossAttentionSpec), |
| 298 | ], |
| 299 | ) |
| 300 | def test_different_uniform_groups_are_not_uniform(self, left_cls, right_cls): |
| 301 | assert not are_uniform_specs(make_spec(left_cls), make_spec(right_cls)) |
| 302 | |
| 303 | def test_different_block_sizes_are_not_uniform(self): |
| 304 | spec = make_spec(FullAttentionSpec) |
| 305 | |
| 306 | assert not are_uniform_specs(spec, replace(spec, block_size=32)) |
| 307 | |
| 308 | def test_registered_custom_spec_uses_base_uniform_rule(self): |
| 309 | @register_kv_cache_spec( |
| 310 | manager_class=FullAttentionManager, |
| 311 | uniform_type_base_spec=FullAttentionSpec, |
| 312 | ) |
| 313 | @dataclass(frozen=True, kw_only=True) |
| 314 | class _CustomFullSpec(FullAttentionSpec): |
| 315 | custom_param: int = 16 |
| 316 | |
| 317 | custom_spec = _CustomFullSpec( |
| 318 | block_size=64, |
| 319 | num_kv_heads=8, |
| 320 | head_size=128, |
| 321 | dtype=torch.bfloat16, |
| 322 | ) |
| 323 | |
| 324 | assert are_uniform_specs(custom_spec, make_spec(FullAttentionSpec)) |
| 325 | |
| 326 | |
| 327 | class TestGetKVCacheSpecKind: |
| 328 | """get_kv_cache_spec_kind() for merged UniformTypeKVCacheSpecs groups.""" |
| 329 | |
| 330 | @pytest.mark.parametrize( |
| 331 | "spec_cls,expected_kind", |
| 332 | [ |
| 333 | (FullAttentionSpec, KVCacheSpecKind.FULL_ATTENTION), |
| 334 | (MLAAttentionSpec, KVCacheSpecKind.MLA_ATTENTION), |
| 335 | (SlidingWindowSpec, KVCacheSpecKind.SLIDING_WINDOW), |
| 336 | ], |
| 337 | ) |
| 338 | def test_homogeneous_group_keeps_member_kind(self, spec_cls, expected_kind): |
| 339 | """A group whose members all share one kind reports that kind.""" |
| 340 | spec = make_spec(spec_cls) |
| 341 | group = UniformTypeKVCacheSpecs( |
| 342 | block_size=64, kv_cache_specs={"layer_0": spec, "layer_1": spec} |
| 343 | ) |
| 344 | |
| 345 | assert get_kv_cache_spec_kind(group) == expected_kind |
| 346 | |
| 347 | @pytest.mark.parametrize("other_cls", [MLAAttentionSpec, SinkFullAttentionSpec]) |
| 348 | def test_mixed_kinds_sharing_full_attention_base(self, other_cls): |
| 349 | """Members with differing kinds but a single registered |
| 350 | uniform_type_base_spec of FullAttentionSpec report FULL_ATTENTION. |
| 351 | |
| 352 | This is the shape produced by a model that pairs a main K/V cache with a |
| 353 | differently shaped auxiliary cache (e.g. a K-only index cache). The two |
| 354 | merge into one group precisely because they share a base spec, so |
| 355 | falling back to UNKNOWN would discard what the merge established. |
| 356 | """ |
| 357 | group = UniformTypeKVCacheSpecs( |
| 358 | block_size=64, |
| 359 | kv_cache_specs={ |
| 360 | "layer_0": make_spec(FullAttentionSpec), |
| 361 | "layer_1": make_spec(other_cls), |
| 362 | }, |
| 363 | ) |
| 364 | |
| 365 | # Precondition: the members really do differ in kind but share a base. |
| 366 | assert ( |
| 367 | len({get_kv_cache_spec_kind(s) for s in group.kv_cache_specs.values()}) == 2 |
| 368 | ) |
| 369 | assert { |
| 370 | KVCacheSpecRegistry.get_uniform_type_base_spec(s) |
| 371 | for s in group.kv_cache_specs.values() |
| 372 | } == {FullAttentionSpec} |
| 373 | |
| 374 | assert get_kv_cache_spec_kind(group) == KVCacheSpecKind.FULL_ATTENTION |
| 375 | |
| 376 | def test_mixed_base_specs_report_unknown(self): |
| 377 | """Members resolving to different base specs remain UNKNOWN.""" |
| 378 | group = UniformTypeKVCacheSpecs( |
| 379 | block_size=64, |
| 380 | kv_cache_specs={ |
| 381 | "layer_0": make_spec(FullAttentionSpec), |
| 382 | "layer_1": make_spec(SlidingWindowSpec), |
| 383 | }, |
| 384 | ) |
| 385 | |
| 386 | assert ( |
| 387 | len( |
| 388 | { |
| 389 | KVCacheSpecRegistry.get_uniform_type_base_spec(s) |
| 390 | for s in group.kv_cache_specs.values() |
| 391 | } |
| 392 | ) |
| 393 | == 2 |
| 394 | ) |
| 395 | assert get_kv_cache_spec_kind(group) == KVCacheSpecKind.UNKNOWN |
| 396 | |
| 397 | def test_uniform_non_full_attention_base_reports_unknown(self): |
| 398 | """The FULL_ATTENTION result is gated on the shared base spec being |
| 399 | FullAttentionSpec; a uniform base of any other type stays UNKNOWN. |
| 400 | """ |
| 401 | |
| 402 | # SlidingWindowMLASpec subclasses SlidingWindowSpec but carries a |
| 403 | # distinct kind, so declaring SlidingWindowSpec as the base gives a |
| 404 | # group with differing kinds yet a single, non-full-attention base. |
| 405 | @register_kv_cache_spec( |
| 406 | manager_class=SlidingWindowManager, |
| 407 | uniform_type_base_spec=SlidingWindowSpec, |
| 408 | ) |
| 409 | @dataclass(frozen=True, kw_only=True) |
| 410 | class _SlidingBasedMLASpec(SlidingWindowMLASpec): |
| 411 | custom_param: int = 16 |
| 412 | |
| 413 | custom_spec = _SlidingBasedMLASpec(**spec_args_map[SlidingWindowMLASpec]) |
| 414 | group = UniformTypeKVCacheSpecs( |
| 415 | block_size=64, |
| 416 | kv_cache_specs={ |
| 417 | "layer_0": custom_spec, |
| 418 | "layer_1": make_spec(SlidingWindowSpec), |
| 419 | }, |
| 420 | ) |
| 421 | |
| 422 | # Two differing kinds, but a single shared base that is not |
| 423 | # FullAttentionSpec -- so the group must not be upgraded. |
| 424 | assert ( |
| 425 | len({get_kv_cache_spec_kind(s) for s in group.kv_cache_specs.values()}) == 2 |
| 426 | ) |
| 427 | assert { |
| 428 | KVCacheSpecRegistry.get_uniform_type_base_spec(s) |
| 429 | for s in group.kv_cache_specs.values() |
| 430 | } == {SlidingWindowSpec} |
| 431 | |
| 432 | assert get_kv_cache_spec_kind(group) == KVCacheSpecKind.UNKNOWN |