Skip to content

vllm.v1.kv_cache_spec_registry

Registry for KVCacheSpec types and their associated managers.

This module provides a pluggable architecture for registering custom KVCacheSpec subclasses without modifying vLLM core code. Out-of-tree platforms can define custom specs and managers by using the @register_kv_cache_spec decorator.

Classes:

Functions:

KVCacheSpecMetadata dataclass

Metadata for a registered KVCacheSpec.

Source code in vllm/v1/kv_cache_spec_registry.py 
23
24
25
26
27
28
29
30
31
32
@dataclass(frozen=True)
class KVCacheSpecMetadata:
    """Metadata for a registered KVCacheSpec."""

    kvcache_spec_cls: type["KVCacheSpec"]
    manager_class: type["SingleTypeKVCacheManager"]
    # The base spec class for grouping compatibility checks.
    # KVCacheSpecs with the same uniform_type_base_spec will be
    # grouped into one kvcache group
    uniform_type_base_spec: type["KVCacheSpec"]

KVCacheSpecRegistry

Global registry for KVCacheSpec types and their associated managers.

Methods:

Source code in vllm/v1/kv_cache_spec_registry.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
class KVCacheSpecRegistry:
    """Global registry for KVCacheSpec types and their associated managers."""

    @classmethod
    def _ensure_registered(cls, vllm_config=None) -> None:
        """Run full KVCacheSpec registration if the registration is not done."""
        if _REGISTRY_KVCACHESPEC_LIST:
            return

        if vllm_config is None:
            from vllm.config import get_current_vllm_config_or_none

            vllm_config = get_current_vllm_config_or_none()

        # lazy import to avoid circular dependency
        from vllm.v1.core.single_type_kv_cache_manager import (
            register_all_kvcache_specs,
        )

        register_all_kvcache_specs(vllm_config)

    @classmethod
    def register(
        cls,
        kvcache_spec_cls: type["KVCacheSpec"],
        manager_class: type["SingleTypeKVCacheManager"] | None = None,
        uniform_type_base_spec: type["KVCacheSpec"] | None = None,
    ) -> None:
        """Register a KVCacheSpec class with its manager and base spec.

        Args:
            kvcache_spec_cls: The KVCacheSpec subclass to register
            manager_class: The SingleTypeKVCacheManager to use for this spec
            uniform_type_base_spec: The base spec class for grouping compatibility.
                instead of being grouped to different kvcache group, `kvcache_spec_cls`
                and `uniform_type_base_spec` will be trated as uniform type.
                If None, defaults to kvcache_spec_cls itself (for built-in base specs).

        """
        assert manager_class is not None, "manager_class is required"
        if uniform_type_base_spec is None:
            uniform_type_base_spec = kvcache_spec_cls
        assert issubclass(kvcache_spec_cls, uniform_type_base_spec), (
            f"{kvcache_spec_cls.__name__} must inherit from its declared "
            f"uniform_type_base_spec {uniform_type_base_spec.__name__}."
        )

        if kvcache_spec_cls in _REGISTRY_KVCACHESPEC_LIST:
            registered_spec = _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls]
            is_same_registration = (
                manager_class == registered_spec.manager_class
                and uniform_type_base_spec == registered_spec.uniform_type_base_spec
            )
            assert is_same_registration, (
                f"Conflicting registration for KVCacheSpec "
                f": {kvcache_spec_cls.__name__}"
            )

        _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls] = KVCacheSpecMetadata(
            kvcache_spec_cls=kvcache_spec_cls,
            manager_class=manager_class,
            uniform_type_base_spec=uniform_type_base_spec,
        )

    @classmethod
    def register_role_manager(
        cls, role: str, manager_class: type["SingleTypeKVCacheManager"]
    ) -> None:
        """Select ``manager_class`` for every group with the given role."""
        registered = _REGISTRY_ROLE_MANAGERS.get(role)
        assert registered is None or registered is manager_class, (
            f"Conflicting manager registration for group role {role}"
        )
        _REGISTRY_ROLE_MANAGERS[role] = manager_class

    @classmethod
    def get_manager_class(
        cls, kvcache_spec: "KVCacheSpec", role: str | None = None
    ) -> type["SingleTypeKVCacheManager"] | None:
        """Get the single type kvcache manager class for a given kvcache spec instance.

        Args:
            kvcache_spec: A KVCacheSpec instance
            role: The group's role; a role-registered manager takes precedence.

        Returns:
            The SingleTypeKVCacheManager class to use for this kvcache_spec

        """
        cls._ensure_registered()
        if role is not None and (manager_class := _REGISTRY_ROLE_MANAGERS.get(role)):
            return manager_class
        kvcache_spec_cls = type(kvcache_spec)

        # Walk up the MRO to find a registered base class
        for base in kvcache_spec_cls.__mro__:
            if base in _REGISTRY_KVCACHESPEC_LIST:
                return _REGISTRY_KVCACHESPEC_LIST[base].manager_class

        return None

    @classmethod
    def get_uniform_type_base_spec(
        cls, kvcache_spec: "KVCacheSpec"
    ) -> type["KVCacheSpec"] | None:
        """Get the base kvcache spec class for grouping compatibility checks.
        KVCacheSpecs with uniform_type_base_spec will be trated as one group.

        Args:
            kvcache_spec: A KVCacheSpec instance

        Returns:
            The base KVCacheSpec class for checking uniform type kvcache specs

        """
        cls._ensure_registered()
        kvcache_spec_cls = type(kvcache_spec)

        # Walk up the MRO to find a registered base spec
        for base in kvcache_spec_cls.__mro__:
            if base in _REGISTRY_KVCACHESPEC_LIST:
                return _REGISTRY_KVCACHESPEC_LIST[base].uniform_type_base_spec

        return None

    @classmethod
    def check_kv_cache_spec_registry(
        cls, kv_cache_spec: dict[str, "KVCacheSpec"]
    ) -> None:
        """Check if the KVCacheSpecs of each layer are registered as expected."""
        cls._ensure_registered()
        for layer_name, spec in kv_cache_spec.items():
            # use raise instead of assert to make it effective in production environment
            if cls.get_uniform_type_base_spec(spec) is None:
                raise ValueError(
                    f"Unsupported KV cache spec type for layer {layer_name}: "
                    f"{type(spec)}. Please register it using "
                    f"@register_kv_cache_spec decorator."
                )
            if cls.get_manager_class(spec) is None:
                raise ValueError(
                    f"No manager found for KV cache spec type for layer "
                    f"{layer_name}: {type(spec)}. Please register it using "
                    f"@register_kv_cache_spec decorator."
                )

_ensure_registered(vllm_config=None) classmethod

Run full KVCacheSpec registration if the registration is not done.

Source code in vllm/v1/kv_cache_spec_registry.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@classmethod
def _ensure_registered(cls, vllm_config=None) -> None:
    """Run full KVCacheSpec registration if the registration is not done."""
    if _REGISTRY_KVCACHESPEC_LIST:
        return

    if vllm_config is None:
        from vllm.config import get_current_vllm_config_or_none

        vllm_config = get_current_vllm_config_or_none()

    # lazy import to avoid circular dependency
    from vllm.v1.core.single_type_kv_cache_manager import (
        register_all_kvcache_specs,
    )

    register_all_kvcache_specs(vllm_config)

check_kv_cache_spec_registry(kv_cache_spec) classmethod

Check if the KVCacheSpecs of each layer are registered as expected.

Source code in vllm/v1/kv_cache_spec_registry.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
@classmethod
def check_kv_cache_spec_registry(
    cls, kv_cache_spec: dict[str, "KVCacheSpec"]
) -> None:
    """Check if the KVCacheSpecs of each layer are registered as expected."""
    cls._ensure_registered()
    for layer_name, spec in kv_cache_spec.items():
        # use raise instead of assert to make it effective in production environment
        if cls.get_uniform_type_base_spec(spec) is None:
            raise ValueError(
                f"Unsupported KV cache spec type for layer {layer_name}: "
                f"{type(spec)}. Please register it using "
                f"@register_kv_cache_spec decorator."
            )
        if cls.get_manager_class(spec) is None:
            raise ValueError(
                f"No manager found for KV cache spec type for layer "
                f"{layer_name}: {type(spec)}. Please register it using "
                f"@register_kv_cache_spec decorator."
            )

get_manager_class(kvcache_spec, role=None) classmethod

Get the single type kvcache manager class for a given kvcache spec instance.

Parameters:

  • kvcache_spec

    (KVCacheSpec) –

    A KVCacheSpec instance

  • role

    (str | None, default: None ) –

    The group's role; a role-registered manager takes precedence.

Returns:

Source code in vllm/v1/kv_cache_spec_registry.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
@classmethod
def get_manager_class(
    cls, kvcache_spec: "KVCacheSpec", role: str | None = None
) -> type["SingleTypeKVCacheManager"] | None:
    """Get the single type kvcache manager class for a given kvcache spec instance.

    Args:
        kvcache_spec: A KVCacheSpec instance
        role: The group's role; a role-registered manager takes precedence.

    Returns:
        The SingleTypeKVCacheManager class to use for this kvcache_spec

    """
    cls._ensure_registered()
    if role is not None and (manager_class := _REGISTRY_ROLE_MANAGERS.get(role)):
        return manager_class
    kvcache_spec_cls = type(kvcache_spec)

    # Walk up the MRO to find a registered base class
    for base in kvcache_spec_cls.__mro__:
        if base in _REGISTRY_KVCACHESPEC_LIST:
            return _REGISTRY_KVCACHESPEC_LIST[base].manager_class

    return None

get_uniform_type_base_spec(kvcache_spec) classmethod

Get the base kvcache spec class for grouping compatibility checks. KVCacheSpecs with uniform_type_base_spec will be trated as one group.

Parameters:

Returns:

  • type[KVCacheSpec] | None –

    The base KVCacheSpec class for checking uniform type kvcache specs

Source code in vllm/v1/kv_cache_spec_registry.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
@classmethod
def get_uniform_type_base_spec(
    cls, kvcache_spec: "KVCacheSpec"
) -> type["KVCacheSpec"] | None:
    """Get the base kvcache spec class for grouping compatibility checks.
    KVCacheSpecs with uniform_type_base_spec will be trated as one group.

    Args:
        kvcache_spec: A KVCacheSpec instance

    Returns:
        The base KVCacheSpec class for checking uniform type kvcache specs

    """
    cls._ensure_registered()
    kvcache_spec_cls = type(kvcache_spec)

    # Walk up the MRO to find a registered base spec
    for base in kvcache_spec_cls.__mro__:
        if base in _REGISTRY_KVCACHESPEC_LIST:
            return _REGISTRY_KVCACHESPEC_LIST[base].uniform_type_base_spec

    return None

register(kvcache_spec_cls, manager_class=None, uniform_type_base_spec=None) classmethod

Register a KVCacheSpec class with its manager and base spec.

Parameters:

  • kvcache_spec_cls

    (type[KVCacheSpec]) –

    The KVCacheSpec subclass to register

  • manager_class

    (type[SingleTypeKVCacheManager] | None, default: None ) –

    The SingleTypeKVCacheManager to use for this spec

  • uniform_type_base_spec

    (type[KVCacheSpec] | None, default: None ) –

    The base spec class for grouping compatibility. instead of being grouped to different kvcache group, kvcache_spec_cls and uniform_type_base_spec will be trated as uniform type. If None, defaults to kvcache_spec_cls itself (for built-in base specs).

Source code in vllm/v1/kv_cache_spec_registry.py 
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
@classmethod
def register(
    cls,
    kvcache_spec_cls: type["KVCacheSpec"],
    manager_class: type["SingleTypeKVCacheManager"] | None = None,
    uniform_type_base_spec: type["KVCacheSpec"] | None = None,
) -> None:
    """Register a KVCacheSpec class with its manager and base spec.

    Args:
        kvcache_spec_cls: The KVCacheSpec subclass to register
        manager_class: The SingleTypeKVCacheManager to use for this spec
        uniform_type_base_spec: The base spec class for grouping compatibility.
            instead of being grouped to different kvcache group, `kvcache_spec_cls`
            and `uniform_type_base_spec` will be trated as uniform type.
            If None, defaults to kvcache_spec_cls itself (for built-in base specs).

    """
    assert manager_class is not None, "manager_class is required"
    if uniform_type_base_spec is None:
        uniform_type_base_spec = kvcache_spec_cls
    assert issubclass(kvcache_spec_cls, uniform_type_base_spec), (
        f"{kvcache_spec_cls.__name__} must inherit from its declared "
        f"uniform_type_base_spec {uniform_type_base_spec.__name__}."
    )

    if kvcache_spec_cls in _REGISTRY_KVCACHESPEC_LIST:
        registered_spec = _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls]
        is_same_registration = (
            manager_class == registered_spec.manager_class
            and uniform_type_base_spec == registered_spec.uniform_type_base_spec
        )
        assert is_same_registration, (
            f"Conflicting registration for KVCacheSpec "
            f": {kvcache_spec_cls.__name__}"
        )

    _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls] = KVCacheSpecMetadata(
        kvcache_spec_cls=kvcache_spec_cls,
        manager_class=manager_class,
        uniform_type_base_spec=uniform_type_base_spec,
    )

register_role_manager(role, manager_class) classmethod

Select manager_class for every group with the given role.

Source code in vllm/v1/kv_cache_spec_registry.py 
103
104
105
106
107
108
109
110
111
112
@classmethod
def register_role_manager(
    cls, role: str, manager_class: type["SingleTypeKVCacheManager"]
) -> None:
    """Select ``manager_class`` for every group with the given role."""
    registered = _REGISTRY_ROLE_MANAGERS.get(role)
    assert registered is None or registered is manager_class, (
        f"Conflicting manager registration for group role {role}"
    )
    _REGISTRY_ROLE_MANAGERS[role] = manager_class

register_kv_cache_spec(manager_class=None, uniform_type_base_spec=None)

Decorator to register a custom KVCacheSpec class.

Parameters:

  • manager_class

    (type[SingleTypeKVCacheManager] | None, default: None ) –

    The SingleTypeKVCacheManager to use for this spec. Required for all registered specs.

  • uniform_type_base_spec

    (type[KVCacheSpec] | None, default: None ) –

    The base spec class for uniform type kv cache specs compatibility. If None, the spec is treated as a new base type.

Examples:

  • Register a new specs: @register_kv_cache_spec( manager_class=FullAttentionManager, uniform_type_base_spec=FullAttentionSpec ) @dataclass(frozen=True, kw_only=True) class CustomFullAttentionSpec(FullAttentionSpec): pass
Source code in vllm/v1/kv_cache_spec_registry.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
def register_kv_cache_spec(
    manager_class: type["SingleTypeKVCacheManager"] | None = None,
    uniform_type_base_spec: type["KVCacheSpec"] | None = None,
):
    """Decorator to register a custom KVCacheSpec class.

    Args:
        manager_class: The SingleTypeKVCacheManager to use for this spec.
            Required for all registered specs.
        uniform_type_base_spec: The base spec class for uniform type kv cache specs
            compatibility. If None, the spec is treated as a new base
            type.

    Examples:
    - Register a new specs:
        @register_kv_cache_spec(
            manager_class=FullAttentionManager,
            uniform_type_base_spec=FullAttentionSpec
        )
        @dataclass(frozen=True, kw_only=True)
        class CustomFullAttentionSpec(FullAttentionSpec):
            pass

    """

    def decorator(kvcache_spec_cls: type["KVCacheSpec"]) -> type["KVCacheSpec"]:
        KVCacheSpecRegistry.register(
            kvcache_spec_cls=kvcache_spec_cls,
            manager_class=manager_class,
            uniform_type_base_spec=uniform_type_base_spec,
        )
        return kvcache_spec_cls

    return decorator