65ca13102d
Signed-off-by: Prajwal A <prajwalanagani@gmail.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Samit <285365963@qq.com>
101 lines
3.2 KiB
Python
101 lines
3.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
Cache adapter registry for diffusion models.
|
|
|
|
This module provides a registry pattern for cache adapters, allowing dynamic
|
|
registration and instantiation of different cache types (TeaCache, DeepCache, etc.).
|
|
"""
|
|
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.cache.base import CacheAdapter
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class CacheType(Enum):
|
|
"""Supported cache adapter types."""
|
|
|
|
NONE = "none"
|
|
TEA_CACHE = "tea_cache"
|
|
# Future cache types can be added here:
|
|
# DEEP_CACHE = "deep_cache"
|
|
# DISTRI_FUSION = "distri_fusion"
|
|
|
|
|
|
# Global registry mapping cache types to adapter classes
|
|
CACHE_ADAPTER_REGISTRY: dict[CacheType, type[CacheAdapter]] = {}
|
|
|
|
|
|
def register_cache_adapter(cache_type: CacheType, adapter_class: type[CacheAdapter]) -> None:
|
|
"""
|
|
Register a cache adapter class for a given cache type.
|
|
|
|
This allows extending the system with new cache types without modifying
|
|
the core cache infrastructure.
|
|
|
|
Args:
|
|
cache_type: CacheType enum value
|
|
adapter_class: CacheAdapter subclass to register
|
|
|
|
Example:
|
|
>>> register_cache_adapter(CacheType.TEA_CACHE, TeaCacheAdapter)
|
|
"""
|
|
if not issubclass(adapter_class, CacheAdapter):
|
|
raise TypeError(f"{adapter_class} must be a subclass of CacheAdapter")
|
|
|
|
CACHE_ADAPTER_REGISTRY[cache_type] = adapter_class
|
|
logger.debug(f"Registered cache adapter: {cache_type.value} -> {adapter_class.__name__}")
|
|
|
|
|
|
def get_cache_adapter(cache_type: str, config: dict[str, Any]) -> CacheAdapter:
|
|
"""
|
|
Factory function to get cache adapter instance.
|
|
|
|
Converts string cache type to enum, looks up in registry, and instantiates
|
|
the appropriate adapter class with the provided configuration.
|
|
|
|
Args:
|
|
cache_type: String name of cache type ("tea_cache", "deep_cache", etc.)
|
|
config: Configuration dictionary to pass to adapter constructor
|
|
|
|
Returns:
|
|
Instantiated CacheAdapter subclass
|
|
|
|
Raises:
|
|
ValueError: If cache_type is unknown or not registered
|
|
|
|
Example:
|
|
>>> adapter = get_cache_adapter("tea_cache", {"rel_l1_thresh": 0.2})
|
|
>>> adapter.apply(transformer)
|
|
"""
|
|
# Normalize cache type string
|
|
cache_type_str = cache_type.lower().strip()
|
|
|
|
# Convert string to enum
|
|
try:
|
|
cache_enum = CacheType(cache_type_str)
|
|
except ValueError:
|
|
available = [ct.value for ct in CacheType if ct != CacheType.NONE]
|
|
raise ValueError(f"Unknown cache type: '{cache_type}'. Available types: {available}")
|
|
|
|
# Check if it's the special "none" case
|
|
if cache_enum == CacheType.NONE:
|
|
raise ValueError("Cannot instantiate adapter for cache_type='none'. Use setup_cache() which handles this case.")
|
|
|
|
# Lookup in registry
|
|
if cache_enum not in CACHE_ADAPTER_REGISTRY:
|
|
raise ValueError(
|
|
f"Cache type '{cache_type}' is not registered. Registered types: {list(CACHE_ADAPTER_REGISTRY.keys())}"
|
|
)
|
|
|
|
adapter_class = CACHE_ADAPTER_REGISTRY[cache_enum]
|
|
|
|
# Instantiate and return
|
|
return adapter_class(config)
|