@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0c6631f02a
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 2, 01:00 UTC

[KVCache] Support Pluggable KVCacheSpec (#37505)

Signed-off-by: MengqingCao <cmq0113@163.com> Signed-off-by: Mengqing Cao <cmq0113@163.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Mengqing Cao committed Jun 4, 2026 at 00:05 UTC 0c6631f02a05eec5bd3209d3cadaa2c7dcac0793
11 files changed +698 -59
tests/v1/core/utils.py
+2
@@ -24,6 +24,7 @@ from vllm.utils.hashing import sha256
24 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
25 from vllm.v1.core.sched.async_scheduler import AsyncScheduler
26 from vllm.v1.core.sched.scheduler import Scheduler
27 +from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
28 from vllm.v1.kv_cache_interface import (
29 FullAttentionSpec,
30 KVCacheConfig,
@@ -160,6 +161,7 @@ def create_scheduler(
161 ],
162 )
163 cache_config.num_gpu_blocks = num_blocks
164 + register_all_kvcache_specs(vllm_config)
165 scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
166 return scheduler_cls(
167 vllm_config=vllm_config,
tests/v1/simple_kv_offload/test_scheduler.py
+6
@@ -30,6 +30,9 @@ from vllm.v1.core.sched.output import (
30 NewRequestData,
31 SchedulerOutput,
32 )
33 +from vllm.v1.core.single_type_kv_cache_manager import (
34 + register_all_kvcache_specs,
35 +)
36 from vllm.v1.kv_cache_interface import (
37 FullAttentionSpec,
38 KVCacheConfig,
@@ -68,6 +71,9 @@ def _make_kv_cache_config(
71 """Build a KVCacheConfig with non-empty kv_cache_tensors."""
72 groups = []
73 tensors = []
74 + register_all_kvcache_specs(
75 + vllm_config=None
76 + ) # Ensure specs are registered for tests
77 for g in range(num_groups):
78 layer_names = [f"layer_{g}"]
79 groups.append(
tests/v1/test_kv_cache_spec_registry.py new
+322
@@ -0,0 +1,322 @@
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 + MambaManager,
20 + SingleTypeKVCacheManager,
21 + SinkFullAttentionManager,
22 + SlidingWindowManager,
23 + register_all_kvcache_specs,
24 +)
25 +from vllm.v1.kv_cache_interface import (
26 + ChunkedLocalAttentionSpec,
27 + CrossAttentionSpec,
28 + FullAttentionSpec,
29 + HiddenStateCacheSpec,
30 + KVCacheSpec,
31 + MambaSpec,
32 + MLAAttentionSpec,
33 + SinkFullAttentionSpec,
34 + SlidingWindowMLASpec,
35 + SlidingWindowSpec,
36 + TQFullAttentionSpec,
37 + UniformTypeKVCacheSpecs,
38 +)
39 +from vllm.v1.kv_cache_spec_registry import (
40 + _REGISTRY_KVCACHESPEC_LIST,
41 + KVCacheSpecRegistry,
42 + register_kv_cache_spec,
43 +)
44 +
45 +
46 +def make_vllm_config() -> VllmConfig:
47 + return VllmConfig(
48 + cache_config=CacheConfig(
49 + block_size=64,
50 + cache_dtype="bfloat16",
51 + ),
52 + device_config=DeviceConfig(device="cpu"),
53 + )
54 +
55 +
56 +vllm_config = make_vllm_config()
57 +register_all_kvcache_specs(vllm_config)
58 +
59 +
60 +@pytest.fixture(autouse=True)
61 +def restore_kv_cache_spec_registry():
62 + registry = _REGISTRY_KVCACHESPEC_LIST.copy()
63 + yield
64 + _REGISTRY_KVCACHESPEC_LIST.clear()
65 + _REGISTRY_KVCACHESPEC_LIST.update(registry)
66 +
67 +
68 +@dataclass(frozen=True)
69 +class _TrulyUnregisteredSpec(KVCacheSpec):
70 + """
71 + A spec that inherits directly from KVCacheSpec with no registered
72 + ancestor in the MRO. Used to test that the registry correctly raises
73 + when no entry can be found.
74 + """
75 +
76 + @property
77 + def page_size_bytes(self) -> int:
78 + return self.block_size * 128
79 +
80 + def max_memory_usage_bytes(self, _) -> int:
81 + return 0
82 +
83 +
84 +spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
85 + FullAttentionSpec: FullAttentionManager,
86 + TQFullAttentionSpec: FullAttentionManager,
87 + MLAAttentionSpec: FullAttentionManager,
88 + HiddenStateCacheSpec: FullAttentionManager,
89 + SlidingWindowSpec: SlidingWindowManager,
90 + SlidingWindowMLASpec: SlidingWindowManager,
91 + ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,
92 + MambaSpec: MambaManager,
93 + CrossAttentionSpec: CrossAttentionManager,
94 + SinkFullAttentionSpec: SinkFullAttentionManager,
95 +}
96 +
97 +spec_uniform_base_map: dict[type[KVCacheSpec], type[KVCacheSpec]] = {
98 + FullAttentionSpec: FullAttentionSpec,
99 + TQFullAttentionSpec: 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 + TQFullAttentionSpec: dict(
115 + block_size=64,
116 + num_kv_heads=8,
117 + head_size=128,
118 + dtype=torch.bfloat16,
119 + tq_slot_size=256,
120 + ),
121 + MLAAttentionSpec: dict(
122 + block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16
123 + ),
124 + HiddenStateCacheSpec: dict(
125 + block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16
126 + ),
127 + SlidingWindowSpec: dict(
128 + block_size=64,
129 + num_kv_heads=8,
130 + head_size=128,
131 + dtype=torch.bfloat16,
132 + sliding_window=128,
133 + ),
134 + SlidingWindowMLASpec: dict(
135 + block_size=64,
136 + num_kv_heads=1,
137 + head_size=128,
138 + dtype=torch.bfloat16,
139 + sliding_window=128,
140 + ),
141 + ChunkedLocalAttentionSpec: dict(
142 + block_size=64,
143 + num_kv_heads=8,
144 + head_size=128,
145 + dtype=torch.bfloat16,
146 + attention_chunk_size=4,
147 + ),
148 + MambaSpec: dict(
149 + block_size=64,
150 + shapes=((2, 512), (3, 32, 32)),
151 + dtypes=(torch.float32, torch.float32),
152 + mamba_cache_mode="align",
153 + num_speculative_blocks=2,
154 + ),
155 + CrossAttentionSpec: dict(
156 + block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16
157 + ),
158 + SinkFullAttentionSpec: dict(
159 + block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16, sink_len=16
160 + ),
161 +}
162 +
163 +
164 +def make_spec(spec_cls: type[KVCacheSpec]) -> KVCacheSpec:
165 + return spec_cls(**spec_args_map[spec_cls])
166 +
167 +
168 +def are_uniform_specs(*specs: KVCacheSpec) -> bool:
169 + return UniformTypeKVCacheSpecs.is_uniform_type(
170 + {f"layer_{i}": spec for i, spec in enumerate(specs)}
171 + )
172 +
173 +
174 +class TestKVCacheSpecRegistry:
175 + """Test the core registry functionality."""
176 +
177 + def test_builtin_kvcache_specs_registered(self):
178 + assert set(spec_manager_map) <= set(_REGISTRY_KVCACHESPEC_LIST)
179 + for spec_cls, manager in spec_manager_map.items():
180 + spec = make_spec(spec_cls)
181 + assert KVCacheSpecRegistry.get_manager_class(spec) is manager
182 + assert (
183 + KVCacheSpecRegistry.get_uniform_type_base_spec(spec)
184 + is spec_uniform_base_map[spec_cls]
185 + )
186 +
187 + @pytest.mark.parametrize("spec_cls", list(spec_manager_map))
188 + def test_custom_spec_register(self, spec_cls):
189 + """A decorated custom spec resolves to the declared manager."""
190 + manager = spec_manager_map[spec_cls]
191 + uniform_base_spec = spec_uniform_base_map[spec_cls]
192 +
193 + @register_kv_cache_spec(
194 + manager_class=manager,
195 + uniform_type_base_spec=uniform_base_spec,
196 + )
197 + @dataclass(frozen=True, kw_only=True)
198 + class _CustomSpec(spec_cls): # type: ignore[valid-type,misc]
199 + custom_param: int = 16
200 +
201 + spec = _CustomSpec(**spec_args_map[spec_cls], custom_param=100)
202 +
203 + assert KVCacheSpecRegistry.get_manager_class(spec) is manager
204 + assert KVCacheSpecRegistry.get_uniform_type_base_spec(spec) is uniform_base_spec
205 +
206 + def test_custom_spec_register_requires_manager(self):
207 + """Invalid register decorator arguments fail early."""
208 +
209 + with pytest.raises(AssertionError, match="manager_class is required"):
210 +
211 + @register_kv_cache_spec(
212 + uniform_type_base_spec=FullAttentionSpec,
213 + )
214 + @dataclass(frozen=True, kw_only=True)
215 + class _CustomFullSpecWithoutManager(FullAttentionSpec):
216 + custom_param: int = 16
217 +
218 + def test_unregistered_spec_no_registered_parent_raises(self):
219 + """
220 + A spec whose entire MRO contains no registered class resolves to None.
221 + Runtime callers should use check_kv_cache_spec_registry to fail early.
222 + Subclasses of registered specs intentionally do not fail — they inherit
223 + their parent's manager via MRO walking.
224 + """
225 + spec = _TrulyUnregisteredSpec(block_size=16)
226 +
227 + assert KVCacheSpecRegistry.get_manager_class(spec) is None
228 + assert KVCacheSpecRegistry.get_uniform_type_base_spec(spec) is None
229 +
230 + with pytest.raises(
231 + ValueError, match="Unsupported KV cache spec type for layer layer_0"
232 + ):
233 + KVCacheSpecRegistry.check_kv_cache_spec_registry({"layer_0": spec})
234 +
235 + with pytest.raises(AssertionError, match="Unsupported KV cache spec type"):
236 + UniformTypeKVCacheSpecs.is_uniform_type({"layer_0": spec})
237 +
238 + def test_unregistered_subclass_inherits_parent_manager(self):
239 + """
240 + An unregistered subclass of a registered spec resolves via MRO
241 + to its parent's manager — this is intentional registry behaviour.
242 + """
243 +
244 + @dataclass(frozen=True, kw_only=True)
245 + class _ImplicitlyInheritedSpec(FullAttentionSpec):
246 + pass
247 +
248 + spec = _ImplicitlyInheritedSpec(
249 + block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16
250 + )
251 +
252 + # MRO walk finds FullAttentionSpec → FullAttentionManager
253 + assert KVCacheSpecRegistry.get_manager_class(spec) is FullAttentionManager
254 +
255 + @pytest.mark.parametrize("spec_cls", list(spec_manager_map))
256 + def test_builtin_specs_are_uniform_with_same_spec_type(self, spec_cls):
257 + spec = make_spec(spec_cls)
258 + assert are_uniform_specs(spec, replace(spec))
259 +
260 + def test_full_attention_family_specs_are_uniform(self):
261 + specs = [
262 + make_spec(FullAttentionSpec),
263 + make_spec(TQFullAttentionSpec),
264 + make_spec(MLAAttentionSpec),
265 + make_spec(HiddenStateCacheSpec),
266 + make_spec(SinkFullAttentionSpec),
267 + ]
268 +
269 + assert are_uniform_specs(*specs)
270 +
271 + @pytest.mark.parametrize(
272 + ("spec_cls", "field", "value"),
273 + [
274 + (SlidingWindowSpec, "sliding_window", 256),
275 + (SlidingWindowMLASpec, "sliding_window", 256),
276 + (ChunkedLocalAttentionSpec, "attention_chunk_size", 8),
277 + (MambaSpec, "num_speculative_blocks", 4),
278 + ],
279 + )
280 + def test_specs_with_type_specific_uniform_fields(self, spec_cls, field, value):
281 + spec = make_spec(spec_cls)
282 + changed_spec = replace(spec, **{field: value})
283 +
284 + assert not are_uniform_specs(spec, changed_spec)
285 +
286 + @pytest.mark.parametrize(
287 + ("left_cls", "right_cls"),
288 + [
289 + (FullAttentionSpec, CrossAttentionSpec),
290 + (FullAttentionSpec, SlidingWindowSpec),
291 + (FullAttentionSpec, ChunkedLocalAttentionSpec),
292 + (FullAttentionSpec, MambaSpec),
293 + (SlidingWindowMLASpec, SlidingWindowSpec),
294 + (ChunkedLocalAttentionSpec, SlidingWindowSpec),
295 + (MambaSpec, CrossAttentionSpec),
296 + ],
297 + )
298 + def test_different_uniform_groups_are_not_uniform(self, left_cls, right_cls):
299 + assert not are_uniform_specs(make_spec(left_cls), make_spec(right_cls))
300 +
301 + def test_different_block_sizes_are_not_uniform(self):
302 + spec = make_spec(FullAttentionSpec)
303 +
304 + assert not are_uniform_specs(spec, replace(spec, block_size=32))
305 +
306 + def test_registered_custom_spec_uses_base_uniform_rule(self):
307 + @register_kv_cache_spec(
308 + manager_class=FullAttentionManager,
309 + uniform_type_base_spec=FullAttentionSpec,
310 + )
311 + @dataclass(frozen=True, kw_only=True)
312 + class _CustomFullSpec(FullAttentionSpec):
313 + custom_param: int = 16
314 +
315 + custom_spec = _CustomFullSpec(
316 + block_size=64,
317 + num_kv_heads=8,
318 + head_size=128,
319 + dtype=torch.bfloat16,
320 + )
321 +
322 + assert are_uniform_specs(custom_spec, make_spec(FullAttentionSpec))
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+5 -2
@@ -13,7 +13,6 @@ from vllm.v1.core.kv_cache_utils import (
13 )
14 from vllm.v1.core.single_type_kv_cache_manager import (
15 SingleTypeKVCacheManager,
16 - spec_manager_map,
16 )
17 from vllm.v1.kv_cache_interface import (
18 FullAttentionSpec,
@@ -21,6 +20,7 @@ from vllm.v1.kv_cache_interface import (
20 KVCacheSpec,
21 UniformTypeKVCacheSpecs,
22 )
23 +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
24
25 # Dummy placeholder hash for store_mask's template computation.
26 _DUMMY_BLOCK_HASH = BlockHash(b"\x00" * 32)
@@ -89,7 +89,10 @@ class MooncakeStoreCoordinator:
89 ] = []
90 for i, g in enumerate(self.kv_cache_groups):
91 spec = _unwrap_spec(g.kv_cache_spec)
92 - manager_cls = spec_manager_map[type(spec)]
92 + manager_cls = KVCacheSpecRegistry.get_manager_class(spec)
93 + assert manager_cls is not None, (
94 + f"No manager registered for KVCacheSpec {spec}"
95 + )
96 for existing_spec, group_ids, existing_cls in attention_groups:
97 if existing_spec == spec:
98 assert manager_cls is existing_cls
vllm/platforms/interface.py
+7
@@ -698,6 +698,13 @@ class Platform:
698 mamba_padding_pct,
699 )
700
701 + @classmethod
702 + def register_custom_kv_cache_specs(cls, vllm_config: "VllmConfig") -> None:
703 + """
704 + Register custom KVCacheSpec class on current platform.
705 + """
706 + pass
707 +
708 @classmethod
709 def verify_model_arch(cls, model_arch: str) -> None:
710 """
vllm/v1/core/kv_cache_utils.py
+4
@@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import (
33 SlidingWindowSpec,
34 UniformTypeKVCacheSpecs,
35 )
36 +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
37 from vllm.v1.request import Request
38 from vllm.v1.utils import tensor_data
39
@@ -1991,6 +1992,9 @@ def get_kv_cache_configs(
1992 "across workers. This is not supported yet."
1993 )
1994
1995 + # Check if the KV cache specs are registered correctly.
1996 + # This is to prevent that some layers are initialized with unregistered specs.
1997 + KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs)
1998 # Get global KV cache groups. This also handles spec unification for
1999 # hybrid models when disable_hybrid_kv_cache_manager is enabled.
2000 # After this call, merged_kv_cache_specs may be modified in-place.
vllm/v1/core/single_type_kv_cache_manager.py
+80 -15
@@ -25,6 +25,7 @@ from vllm.v1.kv_cache_interface import (
25 SlidingWindowSpec,
26 TQFullAttentionSpec,
27 )
28 +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
29 from vllm.v1.request import Request
30
31
@@ -1247,27 +1248,30 @@ class SinkFullAttentionManager(FullAttentionManager):
1248 self.sink_blocks = self.block_pool.free_block_queue.popleft_n(num_sink_block)
1249
1250
1250 -spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
1251 - FullAttentionSpec: FullAttentionManager,
1252 - TQFullAttentionSpec: FullAttentionManager,
1253 - MLAAttentionSpec: FullAttentionManager,
1254 - HiddenStateCacheSpec: FullAttentionManager,
1255 - SlidingWindowSpec: SlidingWindowManager,
1256 - SlidingWindowMLASpec: SlidingWindowManager,
1257 - ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,
1258 - MambaSpec: MambaManager,
1259 - CrossAttentionSpec: CrossAttentionManager,
1260 - SinkFullAttentionSpec: SinkFullAttentionManager,
1261 -}
1262 -
1263 -
1251 def get_manager_for_kv_cache_spec(
1252 kv_cache_spec: KVCacheSpec,
1253 max_num_batched_tokens: int,
1254 max_model_len: int,
1255 **kwargs,
1256 ) -> SingleTypeKVCacheManager:
1270 - manager_class = spec_manager_map[type(kv_cache_spec)]
1257 + """
1258 + Get the appropriate manager for a given KVCacheSpec.
1259 +
1260 + Uses the KVCacheSpecRegistry to look up the manager class, supporting
1261 + both built-in and custom specs registered via @register_kv_cache_spec
1262 + and KVCacheSpecRegistry.register.
1263 +
1264 + Args:
1265 + kv_cache_spec: The KVCacheSpec instance
1266 + max_num_batched_tokens: The maximum number of tokens in a batch
1267 + max_model_len: The maximum context length the model could serve
1268 + Returns:
1269 + An instance of the appropriate SingleTypeKVCacheManager subclass
1270 + """
1271 + manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec)
1272 + assert manager_class is not None, (
1273 + f"No manager registered for KVCacheSpec {type(kv_cache_spec)}"
1274 + )
1275 # SlidingWindow / ChunkedLocalAttention managers recycle blocks across
1276 # chunks; the runtime admission cap must match the recycling-aware bound
1277 # the startup pool sizer uses (single source of truth: the spec method).
@@ -1280,3 +1284,64 @@ def get_manager_for_kv_cache_spec(
1284 )
1285 manager = manager_class(kv_cache_spec, **kwargs)
1286 return manager
1287 +
1288 +
1289 +def register_all_kvcache_specs(vllm_config):
1290 + """Built-in spec registration"""
1291 + KVCacheSpecRegistry.register(
1292 + FullAttentionSpec,
1293 + FullAttentionManager,
1294 + uniform_type_base_spec=FullAttentionSpec,
1295 + )
1296 +
1297 + KVCacheSpecRegistry.register(
1298 + SlidingWindowSpec,
1299 + SlidingWindowManager,
1300 + uniform_type_base_spec=SlidingWindowSpec,
1301 + )
1302 + KVCacheSpecRegistry.register(
1303 + SlidingWindowMLASpec,
1304 + SlidingWindowManager,
1305 + uniform_type_base_spec=SlidingWindowMLASpec,
1306 + )
1307 +
1308 + KVCacheSpecRegistry.register(
1309 + MambaSpec, MambaManager, uniform_type_base_spec=MambaSpec
1310 + )
1311 + KVCacheSpecRegistry.register(
1312 + ChunkedLocalAttentionSpec,
1313 + ChunkedLocalAttentionManager,
1314 + uniform_type_base_spec=ChunkedLocalAttentionSpec,
1315 + )
1316 + KVCacheSpecRegistry.register(
1317 + CrossAttentionSpec,
1318 + CrossAttentionManager,
1319 + uniform_type_base_spec=CrossAttentionSpec,
1320 + )
1321 +
1322 + # FullAttentionSpec subclasses — grouped with FullAttentionSpec
1323 + KVCacheSpecRegistry.register(
1324 + TQFullAttentionSpec,
1325 + FullAttentionManager,
1326 + uniform_type_base_spec=FullAttentionSpec,
1327 + )
1328 + KVCacheSpecRegistry.register(
1329 + MLAAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec
1330 + )
1331 + # NOTE(Mengqing): HiddenStateCacheSpec won't take part in
1332 + # grouping, thus the uniform_type_base_spec is just a
1333 + # placeholder.
1334 + KVCacheSpecRegistry.register(
1335 + HiddenStateCacheSpec,
1336 + FullAttentionManager,
1337 + uniform_type_base_spec=FullAttentionSpec,
1338 + )
1339 + KVCacheSpecRegistry.register(
1340 + SinkFullAttentionSpec,
1341 + SinkFullAttentionManager,
1342 + uniform_type_base_spec=FullAttentionSpec,
1343 + )
1344 +
1345 + from vllm.platforms import current_platform
1346 +
1347 + current_platform.register_custom_kv_cache_specs(vllm_config)
vllm/v1/engine/core.py
+4
@@ -52,6 +52,7 @@ from vllm.v1.core.kv_cache_utils import (
52 )
53 from vllm.v1.core.sched.interface import PauseState, SchedulerInterface
54 from vllm.v1.core.sched.output import SchedulerOutput
55 +from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
56 from vllm.v1.engine import (
57 EEP_NOTIFICATION_CALL_ID,
58 EEPNotificationType,
@@ -235,6 +236,9 @@ class EngineCore:
236 def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig:
237 start = time.time()
238
239 + # register all kvcache specs in enginecore process.
240 + register_all_kvcache_specs(vllm_config)
241 +
242 # Get all kv cache needed by the model
243 kv_cache_specs = self.model_executor.get_kv_cache_specs()
244
vllm/v1/kv_cache_interface.py
+57 -42
@@ -17,6 +17,7 @@ from vllm.logger import init_logger
17 from vllm.utils.math_utils import cdiv, round_up
18 from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
19 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
20 +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
21
22 if TYPE_CHECKING:
23 from vllm.config import VllmConfig
@@ -139,6 +140,21 @@ class KVCacheSpec:
140 )
141 return copy.deepcopy(specs[0])
142
143 + def is_uniform_with_collection(
144 + self, kv_cache_specs: dict[str, KVCacheSpec]
145 + ) -> bool:
146 + """
147 + Whether this KVCacheSpec is uniform with all specs of all layers.
148 + """
149 + uniform_type_base_spec = KVCacheSpecRegistry.get_uniform_type_base_spec(self)
150 + assert uniform_type_base_spec is not None, (
151 + f"Unsupported KV cache spec type: {type(self)}. "
152 + "Please register it using @register_kv_cache_spec decorator."
153 + )
154 + return all(
155 + isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values()
156 + )
157 +
158
159 @dataclass(frozen=True, kw_only=True)
160 class AttentionSpec(KVCacheSpec):
@@ -430,6 +446,15 @@ class ChunkedLocalAttentionSpec(AttentionSpec):
446 )
447 return max_blocks * self.page_size_bytes
448
449 + def is_uniform_with_collection(
450 + self, kv_cache_specs: dict[str, KVCacheSpec]
451 + ) -> bool:
452 + return all(
453 + isinstance(spec, ChunkedLocalAttentionSpec)
454 + and spec.attention_chunk_size == self.attention_chunk_size
455 + for spec in kv_cache_specs.values()
456 + )
457 +
458
459 @dataclass(frozen=True, kw_only=True)
460 class SlidingWindowSpec(AttentionSpec):
@@ -493,6 +518,15 @@ class SlidingWindowSpec(AttentionSpec):
518 )
519 return max_blocks * self.page_size_bytes
520
521 + def is_uniform_with_collection(
522 + self, kv_cache_specs: dict[str, KVCacheSpec]
523 + ) -> bool:
524 + return all(
525 + isinstance(spec, SlidingWindowSpec)
526 + and spec.sliding_window == self.sliding_window
527 + for spec in kv_cache_specs.values()
528 + )
529 +
530
531 @dataclass(frozen=True, kw_only=True)
532 class SlidingWindowMLASpec(SlidingWindowSpec):
@@ -558,6 +592,15 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
592 model_version=model_version_set.pop(),
593 )
594
595 + def is_uniform_with_collection(
596 + self, kv_cache_specs: dict[str, KVCacheSpec]
597 + ) -> bool:
598 + return all(
599 + isinstance(spec, SlidingWindowMLASpec)
600 + and spec.sliding_window == self.sliding_window
601 + for spec in kv_cache_specs.values()
602 + )
603 +
604
605 @dataclass(frozen=True)
606 class MambaSpec(KVCacheSpec):
@@ -590,6 +633,15 @@ class MambaSpec(KVCacheSpec):
633 else:
634 return self.page_size_bytes * (1 + self.num_speculative_blocks)
635
636 + def is_uniform_with_collection(
637 + self, kv_cache_specs: dict[str, KVCacheSpec]
638 + ) -> bool:
639 + return all(
640 + isinstance(spec, MambaSpec)
641 + and spec.num_speculative_blocks == self.num_speculative_blocks
642 + for spec in kv_cache_specs.values()
643 + )
644 +
645
646 @dataclass(frozen=True)
647 class EncoderOnlyAttentionSpec(AttentionSpec):
@@ -689,53 +741,16 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
741 def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool:
742 """
743 Whether all layers have the same type of KV cache spec.
744 +
745 + Uses the registry to determine grouping base classes, so custom specs
746 + that inherit from FullAttentionSpec are treated as full attention.
747 """
748 block_sizes = set(spec.block_size for spec in kv_cache_specs.values())
749 if len(block_sizes) > 1:
750 # Different block sizes, not uniform.
751 return False
697 - one_spec = next(iter(kv_cache_specs.values()))
698 - # NOTE: Check subclasses before parent classes since isinstance()
699 - # returns True for subclasses.
700 - if isinstance(one_spec, SlidingWindowMLASpec):
701 - # SlidingWindowMLASpec is uniform if all specs are SlidingWindowMLASpec
702 - # with the same sliding_window size.
703 - return all(
704 - isinstance(spec, SlidingWindowMLASpec)
705 - and spec.sliding_window == one_spec.sliding_window
706 - for spec in kv_cache_specs.values()
707 - )
708 - elif isinstance(one_spec, FullAttentionSpec):
709 - return all(
710 - isinstance(spec, FullAttentionSpec) for spec in kv_cache_specs.values()
711 - )
712 - elif isinstance(one_spec, CrossAttentionSpec):
713 - return all(
714 - isinstance(spec, CrossAttentionSpec) for spec in kv_cache_specs.values()
715 - )
716 - elif isinstance(one_spec, SlidingWindowSpec):
717 - return all(
718 - isinstance(spec, SlidingWindowSpec)
719 - and spec.sliding_window == one_spec.sliding_window
720 - for spec in kv_cache_specs.values()
721 - )
722 - elif isinstance(one_spec, ChunkedLocalAttentionSpec):
723 - return all(
724 - isinstance(spec, ChunkedLocalAttentionSpec)
725 - and spec.attention_chunk_size == one_spec.attention_chunk_size
726 - for spec in kv_cache_specs.values()
727 - )
728 - elif isinstance(one_spec, MambaSpec):
729 - return all(
730 - isinstance(spec, MambaSpec)
731 - and spec.num_speculative_blocks == one_spec.num_speculative_blocks
732 - for spec in kv_cache_specs.values()
733 - )
734 - else:
735 - # NOTE(Chen): Please add new branches for new KV cache spec types.
736 - raise NotImplementedError(
737 - f"Unsupported KV cache spec type: {type(one_spec)}"
738 - )
752 + first_spec = next(iter(kv_cache_specs.values()))
753 + return first_spec.is_uniform_with_collection(kv_cache_specs)
754
755 @classmethod
756 def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None:
vllm/v1/kv_cache_spec_registry.py new
+209
@@ -0,0 +1,209 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Registry for KVCacheSpec types and their associated managers.
6 +
7 +This module provides a pluggable architecture for registering custom KVCacheSpec
8 +subclasses without modifying vLLM core code. Out-of-tree platforms can define
9 +custom specs and managers by using the @register_kv_cache_spec decorator.
10 +"""
11 +
12 +from dataclasses import dataclass
13 +from typing import TYPE_CHECKING
14 +
15 +from vllm.logger import init_logger
16 +
17 +logger = init_logger(__name__)
18 +
19 +if TYPE_CHECKING:
20 + from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager
21 + from vllm.v1.kv_cache_interface import KVCacheSpec
22 +
23 +
24 +@dataclass(frozen=True)
25 +class KVCacheSpecMetadata:
26 + """Metadata for a registered KVCacheSpec."""
27 +
28 + kvcache_spec_cls: type["KVCacheSpec"]
29 + manager_class: type["SingleTypeKVCacheManager"]
30 + # The base spec class for grouping compatibility checks.
31 + # KVCacheSpecs with the same uniform_type_base_spec will be
32 + # grouped into one kvcache group
33 + uniform_type_base_spec: type["KVCacheSpec"]
34 +
35 +
36 +_REGISTRY_KVCACHESPEC_LIST: dict[type["KVCacheSpec"], KVCacheSpecMetadata] = {}
37 +
38 +
39 +class KVCacheSpecRegistry:
40 + """Global registry for KVCacheSpec types and their associated managers."""
41 +
42 + @classmethod
43 + def _ensure_registered(cls, vllm_config=None) -> None:
44 + """
45 + Run full KVCacheSpec registration if the registration is not done.
46 + """
47 + if _REGISTRY_KVCACHESPEC_LIST:
48 + return
49 +
50 + if vllm_config is None:
51 + from vllm.config import get_current_vllm_config_or_none
52 +
53 + vllm_config = get_current_vllm_config_or_none()
54 +
55 + # lazy import to avoid circular dependency
56 + from vllm.v1.core.single_type_kv_cache_manager import (
57 + register_all_kvcache_specs,
58 + )
59 +
60 + register_all_kvcache_specs(vllm_config)
61 +
62 + @classmethod
63 + def register(
64 + cls,
65 + kvcache_spec_cls: type["KVCacheSpec"],
66 + manager_class: type["SingleTypeKVCacheManager"] | None = None,
67 + uniform_type_base_spec: type["KVCacheSpec"] | None = None,
68 + ) -> None:
69 + """
70 + Register a KVCacheSpec class with its manager and base spec.
71 +
72 + Args:
73 + kvcache_spec_cls: The KVCacheSpec subclass to register
74 + manager_class: The SingleTypeKVCacheManager to use for this spec
75 + uniform_type_base_spec: The base spec class for grouping compatibility.
76 + instead of being grouped to different kvcache group, `kvcache_spec_cls`
77 + and `uniform_type_base_spec` will be trated as uniform type.
78 + If None, defaults to kvcache_spec_cls itself (for built-in base specs).
79 + """
80 + assert manager_class is not None, "manager_class is required"
81 + if uniform_type_base_spec is None:
82 + uniform_type_base_spec = kvcache_spec_cls
83 + assert issubclass(kvcache_spec_cls, uniform_type_base_spec), (
84 + f"{kvcache_spec_cls.__name__} must inherit from its declared "
85 + f"uniform_type_base_spec {uniform_type_base_spec.__name__}."
86 + )
87 +
88 + if kvcache_spec_cls in _REGISTRY_KVCACHESPEC_LIST:
89 + registered_spec = _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls]
90 + is_same_registration = (
91 + manager_class == registered_spec.manager_class
92 + and uniform_type_base_spec == registered_spec.uniform_type_base_spec
93 + )
94 + assert is_same_registration, (
95 + f"Conflicting registration for KVCacheSpec "
96 + f": {kvcache_spec_cls.__name__}"
97 + )
98 +
99 + _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls] = KVCacheSpecMetadata(
100 + kvcache_spec_cls=kvcache_spec_cls,
101 + manager_class=manager_class,
102 + uniform_type_base_spec=uniform_type_base_spec,
103 + )
104 +
105 + @classmethod
106 + def get_manager_class(
107 + cls, kvcache_spec: "KVCacheSpec"
108 + ) -> type["SingleTypeKVCacheManager"] | None:
109 + """
110 + Get the single type kvcache manager class for a given kvcache spec instance.
111 +
112 + Args:
113 + kvcache_spec: A KVCacheSpec instance
114 +
115 + Returns:
116 + The SingleTypeKVCacheManager class to use for this kvcache_spec
117 + """
118 + cls._ensure_registered()
119 + kvcache_spec_cls = type(kvcache_spec)
120 +
121 + # Walk up the MRO to find a registered base class
122 + for base in kvcache_spec_cls.__mro__:
123 + if base in _REGISTRY_KVCACHESPEC_LIST:
124 + return _REGISTRY_KVCACHESPEC_LIST[base].manager_class
125 +
126 + return None
127 +
128 + @classmethod
129 + def get_uniform_type_base_spec(
130 + cls, kvcache_spec: "KVCacheSpec"
131 + ) -> type["KVCacheSpec"] | None:
132 + """
133 + Get the base kvcache spec class for grouping compatibility checks.
134 + KVCacheSpecs with uniform_type_base_spec will be trated as one group.
135 +
136 + Args:
137 + kvcache_spec: A KVCacheSpec instance
138 +
139 + Returns:
140 + The base KVCacheSpec class for checking uniform type kvcache specs
141 + """
142 + cls._ensure_registered()
143 + kvcache_spec_cls = type(kvcache_spec)
144 +
145 + # Walk up the MRO to find a registered base spec
146 + for base in kvcache_spec_cls.__mro__:
147 + if base in _REGISTRY_KVCACHESPEC_LIST:
148 + return _REGISTRY_KVCACHESPEC_LIST[base].uniform_type_base_spec
149 +
150 + return None
151 +
152 + @classmethod
153 + def check_kv_cache_spec_registry(
154 + cls, kv_cache_spec: dict[str, "KVCacheSpec"]
155 + ) -> None:
156 + """
157 + Check if the KVCacheSpecs of each layer are registered as expected.
158 + """
159 + cls._ensure_registered()
160 + for layer_name, spec in kv_cache_spec.items():
161 + # use raise instead of assert to make it effective in production environment
162 + if cls.get_uniform_type_base_spec(spec) is None:
163 + raise ValueError(
164 + f"Unsupported KV cache spec type for layer {layer_name}: "
165 + f"{type(spec)}. Please register it using "
166 + f"@register_kv_cache_spec decorator."
167 + )
168 + if cls.get_manager_class(spec) is None:
169 + raise ValueError(
170 + f"No manager found for KV cache spec type for layer "
171 + f"{layer_name}: {type(spec)}. Please register it using "
172 + f"@register_kv_cache_spec decorator."
173 + )
174 +
175 +
176 +def register_kv_cache_spec(
177 + manager_class: type["SingleTypeKVCacheManager"] | None = None,
178 + uniform_type_base_spec: type["KVCacheSpec"] | None = None,
179 +):
180 + """
181 + Decorator to register a custom KVCacheSpec class.
182 +
183 + Args:
184 + manager_class: The SingleTypeKVCacheManager to use for this spec.
185 + Required for all registered specs.
186 + uniform_type_base_spec: The base spec class for uniform type kv cache specs
187 + compatibility. If None, the spec is treated as a new base
188 + type.
189 +
190 + Examples:
191 + - Register a new specs:
192 + @register_kv_cache_spec(
193 + manager_class=FullAttentionManager,
194 + uniform_type_base_spec=FullAttentionSpec
195 + )
196 + @dataclass(frozen=True, kw_only=True)
197 + class CustomFullAttentionSpec(FullAttentionSpec):
198 + pass
199 + """
200 +
201 + def decorator(kvcache_spec_cls: type["KVCacheSpec"]) -> type["KVCacheSpec"]:
202 + KVCacheSpecRegistry.register(
203 + kvcache_spec_cls=kvcache_spec_cls,
204 + manager_class=manager_class,
205 + uniform_type_base_spec=uniform_type_base_spec,
206 + )
207 + return kvcache_spec_cls
208 +
209 + return decorator
vllm/v1/worker/gpu_model_runner.py
+2
@@ -152,6 +152,7 @@ from vllm.v1.kv_cache_interface import (
152 SlidingWindowSpec,
153 UniformTypeKVCacheSpecs,
154 )
155 +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
156 from vllm.v1.outputs import (
157 EMPTY_MODEL_RUNNER_OUTPUT,
158 AsyncModelRunnerOutput,
@@ -6235,6 +6236,7 @@ class GPUModelRunner(
6236 )
6237
6238 kv_cache_spec = self.get_kv_cache_spec()
6239 + KVCacheSpecRegistry.check_kv_cache_spec_registry(kv_cache_spec)
6240 kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec)
6241 min_blocks = self.compilation_config.max_cudagraph_capture_size or 1
6242