@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 219 lines 8.13 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 """Registry for KVCacheSpec types and their associated managers.
5
6 This module provides a pluggable architecture for registering custom KVCacheSpec
7 subclasses without modifying vLLM core code. Out-of-tree platforms can define
8 custom specs and managers by using the @register_kv_cache_spec decorator.
9 """
10
11 from dataclasses import dataclass
12 from typing import TYPE_CHECKING
13
14 from vllm.logger import init_logger
15
16 logger = init_logger(__name__)
17
18 if TYPE_CHECKING:
19 from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager
20 from vllm.v1.kv_cache_interface import KVCacheSpec
21
22
23 @dataclass(frozen=True)
24 class KVCacheSpecMetadata:
25 """Metadata for a registered KVCacheSpec."""
26
27 kvcache_spec_cls: type["KVCacheSpec"]
28 manager_class: type["SingleTypeKVCacheManager"]
29 # The base spec class for grouping compatibility checks.
30 # KVCacheSpecs with the same uniform_type_base_spec will be
31 # grouped into one kvcache group
32 uniform_type_base_spec: type["KVCacheSpec"]
33
34
35 _REGISTRY_KVCACHESPEC_LIST: dict[type["KVCacheSpec"], KVCacheSpecMetadata] = {}
36 _REGISTRY_ROLE_MANAGERS: dict[str, type["SingleTypeKVCacheManager"]] = {}
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 """Run full KVCacheSpec registration if the registration is not done."""
45 if _REGISTRY_KVCACHESPEC_LIST:
46 return
47
48 if vllm_config is None:
49 from vllm.config import get_current_vllm_config_or_none
50
51 vllm_config = get_current_vllm_config_or_none()
52
53 # lazy import to avoid circular dependency
54 from vllm.v1.core.single_type_kv_cache_manager import (
55 register_all_kvcache_specs,
56 )
57
58 register_all_kvcache_specs(vllm_config)
59
60 @classmethod
61 def register(
62 cls,
63 kvcache_spec_cls: type["KVCacheSpec"],
64 manager_class: type["SingleTypeKVCacheManager"] | None = None,
65 uniform_type_base_spec: type["KVCacheSpec"] | None = None,
66 ) -> None:
67 """Register a KVCacheSpec class with its manager and base spec.
68
69 Args:
70 kvcache_spec_cls: The KVCacheSpec subclass to register
71 manager_class: The SingleTypeKVCacheManager to use for this spec
72 uniform_type_base_spec: The base spec class for grouping compatibility.
73 instead of being grouped to different kvcache group, `kvcache_spec_cls`
74 and `uniform_type_base_spec` will be trated as uniform type.
75 If None, defaults to kvcache_spec_cls itself (for built-in base specs).
76
77 """
78 assert manager_class is not None, "manager_class is required"
79 if uniform_type_base_spec is None:
80 uniform_type_base_spec = kvcache_spec_cls
81 assert issubclass(kvcache_spec_cls, uniform_type_base_spec), (
82 f"{kvcache_spec_cls.__name__} must inherit from its declared "
83 f"uniform_type_base_spec {uniform_type_base_spec.__name__}."
84 )
85
86 if kvcache_spec_cls in _REGISTRY_KVCACHESPEC_LIST:
87 registered_spec = _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls]
88 is_same_registration = (
89 manager_class == registered_spec.manager_class
90 and uniform_type_base_spec == registered_spec.uniform_type_base_spec
91 )
92 assert is_same_registration, (
93 f"Conflicting registration for KVCacheSpec "
94 f": {kvcache_spec_cls.__name__}"
95 )
96
97 _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls] = KVCacheSpecMetadata(
98 kvcache_spec_cls=kvcache_spec_cls,
99 manager_class=manager_class,
100 uniform_type_base_spec=uniform_type_base_spec,
101 )
102
103 @classmethod
104 def register_role_manager(
105 cls, role: str, manager_class: type["SingleTypeKVCacheManager"]
106 ) -> None:
107 """Select ``manager_class`` for every group with the given role."""
108 registered = _REGISTRY_ROLE_MANAGERS.get(role)
109 assert registered is None or registered is manager_class, (
110 f"Conflicting manager registration for group role {role}"
111 )
112 _REGISTRY_ROLE_MANAGERS[role] = manager_class
113
114 @classmethod
115 def get_manager_class(
116 cls, kvcache_spec: "KVCacheSpec", role: str | None = None
117 ) -> type["SingleTypeKVCacheManager"] | None:
118 """Get the single type kvcache manager class for a given kvcache spec instance.
119
120 Args:
121 kvcache_spec: A KVCacheSpec instance
122 role: The group's role; a role-registered manager takes precedence.
123
124 Returns:
125 The SingleTypeKVCacheManager class to use for this kvcache_spec
126
127 """
128 cls._ensure_registered()
129 if role is not None and (manager_class := _REGISTRY_ROLE_MANAGERS.get(role)):
130 return manager_class
131 kvcache_spec_cls = type(kvcache_spec)
132
133 # Walk up the MRO to find a registered base class
134 for base in kvcache_spec_cls.__mro__:
135 if base in _REGISTRY_KVCACHESPEC_LIST:
136 return _REGISTRY_KVCACHESPEC_LIST[base].manager_class
137
138 return None
139
140 @classmethod
141 def get_uniform_type_base_spec(
142 cls, kvcache_spec: "KVCacheSpec"
143 ) -> type["KVCacheSpec"] | None:
144 """Get the base kvcache spec class for grouping compatibility checks.
145 KVCacheSpecs with uniform_type_base_spec will be trated as one group.
146
147 Args:
148 kvcache_spec: A KVCacheSpec instance
149
150 Returns:
151 The base KVCacheSpec class for checking uniform type kvcache specs
152
153 """
154 cls._ensure_registered()
155 kvcache_spec_cls = type(kvcache_spec)
156
157 # Walk up the MRO to find a registered base spec
158 for base in kvcache_spec_cls.__mro__:
159 if base in _REGISTRY_KVCACHESPEC_LIST:
160 return _REGISTRY_KVCACHESPEC_LIST[base].uniform_type_base_spec
161
162 return None
163
164 @classmethod
165 def check_kv_cache_spec_registry(
166 cls, kv_cache_spec: dict[str, "KVCacheSpec"]
167 ) -> None:
168 """Check if the KVCacheSpecs of each layer are registered as expected."""
169 cls._ensure_registered()
170 for layer_name, spec in kv_cache_spec.items():
171 # use raise instead of assert to make it effective in production environment
172 if cls.get_uniform_type_base_spec(spec) is None:
173 raise ValueError(
174 f"Unsupported KV cache spec type for layer {layer_name}: "
175 f"{type(spec)}. Please register it using "
176 f"@register_kv_cache_spec decorator."
177 )
178 if cls.get_manager_class(spec) is None:
179 raise ValueError(
180 f"No manager found for KV cache spec type for layer "
181 f"{layer_name}: {type(spec)}. Please register it using "
182 f"@register_kv_cache_spec decorator."
183 )
184
185
186 def register_kv_cache_spec(
187 manager_class: type["SingleTypeKVCacheManager"] | None = None,
188 uniform_type_base_spec: type["KVCacheSpec"] | None = None,
189 ):
190 """Decorator to register a custom KVCacheSpec class.
191
192 Args:
193 manager_class: The SingleTypeKVCacheManager to use for this spec.
194 Required for all registered specs.
195 uniform_type_base_spec: The base spec class for uniform type kv cache specs
196 compatibility. If None, the spec is treated as a new base
197 type.
198
199 Examples:
200 - Register a new specs:
201 @register_kv_cache_spec(
202 manager_class=FullAttentionManager,
203 uniform_type_base_spec=FullAttentionSpec
204 )
205 @dataclass(frozen=True, kw_only=True)
206 class CustomFullAttentionSpec(FullAttentionSpec):
207 pass
208
209 """
210
211 def decorator(kvcache_spec_cls: type["KVCacheSpec"]) -> type["KVCacheSpec"]:
212 KVCacheSpecRegistry.register(
213 kvcache_spec_cls=kvcache_spec_cls,
214 manager_class=manager_class,
215 uniform_type_base_spec=uniform_type_base_spec,
216 )
217 return kvcache_spec_cls
218
219 return decorator