wehub-resource-sync
248 lines
8.1 KiB
Python
248 lines
8.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Factory for creating cache instances from configuration.
|
|
|
|
This module provides a unified way to instantiate cache components
|
|
based on configuration settings.
|
|
|
|
Note: oMLX only supports paged SSD-based caching. Memory KV cache is managed
|
|
by mlx-lm's BatchGenerator. When paged SSD cache is disabled, no oMLX caching
|
|
is performed.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Optional, TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from .paged_cache import PagedCacheManager
|
|
from .prefix_cache import BlockAwarePrefixCache
|
|
from .paged_ssd_cache import PagedSSDCacheManager
|
|
from ..memory_monitor import MemoryMonitor
|
|
|
|
|
|
@dataclass
|
|
class CacheConfig:
|
|
"""
|
|
Configuration for cache system (paged SSD-only).
|
|
|
|
Note: oMLX only supports paged SSD-based caching via BlockAwarePrefixCache.
|
|
Memory KV cache is managed entirely by mlx-lm's BatchGenerator.
|
|
|
|
Attributes:
|
|
block_size: Number of tokens per cache block.
|
|
max_num_blocks: Maximum number of cache blocks.
|
|
initial_blocks: Initial number of blocks to allocate.
|
|
paged_ssd_cache_dir: Directory for paged SSD cache storage. If None, caching is disabled.
|
|
max_paged_ssd_cache_size: Maximum size of paged SSD cache in bytes.
|
|
max_kv_cache_memory: Maximum GPU memory for KV cache.
|
|
model_name: Model name for cache isolation.
|
|
"""
|
|
|
|
block_size: int = 64
|
|
max_num_blocks: int = 1024
|
|
initial_blocks: int = 256
|
|
paged_ssd_cache_dir: Optional[Path] = None
|
|
max_paged_ssd_cache_size: int = 100 * 1024 * 1024 * 1024 # 100GB
|
|
max_kv_cache_memory: Optional[int] = None
|
|
model_name: str = ""
|
|
|
|
|
|
class CacheFactory:
|
|
"""
|
|
Factory for creating cache instances (paged SSD-only).
|
|
|
|
This class provides static methods to create cache components
|
|
based on a unified CacheConfig. oMLX only supports paged SSD-based caching
|
|
via BlockAwarePrefixCache.
|
|
|
|
Example:
|
|
config = CacheConfig(
|
|
block_size=64,
|
|
paged_ssd_cache_dir=Path("/tmp/cache"),
|
|
)
|
|
|
|
paged_cache = CacheFactory.create_paged_cache(config)
|
|
paged_ssd_cache = CacheFactory.create_paged_ssd_cache(config, "llama-3b")
|
|
prefix_cache = CacheFactory.create_prefix_cache(
|
|
config, paged_cache, paged_ssd_cache
|
|
)
|
|
"""
|
|
|
|
@staticmethod
|
|
def create_paged_cache(
|
|
config: CacheConfig,
|
|
num_layers: Optional[int] = None,
|
|
) -> Optional["PagedCacheManager"]:
|
|
"""
|
|
Create a PagedCacheManager instance.
|
|
|
|
Args:
|
|
config: Cache configuration.
|
|
num_layers: Number of model layers (unused, kept for API compat).
|
|
|
|
Returns:
|
|
Configured PagedCacheManager instance, or None if paged SSD cache is disabled.
|
|
"""
|
|
if config.paged_ssd_cache_dir is None:
|
|
return None
|
|
|
|
from .paged_cache import PagedCacheManager
|
|
|
|
return PagedCacheManager(
|
|
block_size=config.block_size,
|
|
max_blocks=config.max_num_blocks,
|
|
enable_caching=True,
|
|
model_name=config.model_name,
|
|
initial_blocks=config.initial_blocks,
|
|
)
|
|
|
|
@staticmethod
|
|
def create_paged_ssd_cache(
|
|
config: CacheConfig,
|
|
model_name: Optional[str] = None,
|
|
) -> Optional["PagedSSDCacheManager"]:
|
|
"""
|
|
Create a PagedSSDCacheManager instance.
|
|
|
|
Args:
|
|
config: Cache configuration.
|
|
model_name: Override model name for cache isolation.
|
|
|
|
Returns:
|
|
Configured PagedSSDCacheManager instance, or None if disabled.
|
|
"""
|
|
if config.paged_ssd_cache_dir is None:
|
|
return None
|
|
|
|
from .paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
cache_dir = config.paged_ssd_cache_dir
|
|
if model_name:
|
|
cache_dir = cache_dir / model_name
|
|
|
|
return PagedSSDCacheManager(
|
|
cache_dir=cache_dir,
|
|
max_size_bytes=config.max_paged_ssd_cache_size,
|
|
expected_model_name=config.model_name or "",
|
|
expected_block_size=config.block_size,
|
|
expected_block_size_tokens=config.block_size,
|
|
# No memory monitor available here, so leave
|
|
# ``expected_kv_bytes_per_token`` at its 200 KB default.
|
|
# The Scheduler construction path is the one that benefits
|
|
# from the model-derived value; this factory is used mostly
|
|
# by direct/test callers.
|
|
)
|
|
|
|
@staticmethod
|
|
def create_prefix_cache(
|
|
config: CacheConfig,
|
|
model: Any = None,
|
|
paged_cache: Optional["PagedCacheManager"] = None,
|
|
paged_ssd_cache: Optional["PagedSSDCacheManager"] = None,
|
|
) -> Optional["BlockAwarePrefixCache"]:
|
|
"""
|
|
Create a BlockAwarePrefixCache instance.
|
|
|
|
Note: oMLX only supports paged SSD-based caching via BlockAwarePrefixCache.
|
|
Returns None if paged SSD cache is disabled.
|
|
|
|
Args:
|
|
config: Cache configuration.
|
|
model: Model instance for cache identification.
|
|
paged_cache: PagedCacheManager for block-based caching.
|
|
paged_ssd_cache: PagedSSDCacheManager for SSD storage.
|
|
|
|
Returns:
|
|
Configured BlockAwarePrefixCache instance, or None if disabled.
|
|
"""
|
|
if config.paged_ssd_cache_dir is None or paged_cache is None:
|
|
return None
|
|
|
|
from .prefix_cache import BlockAwarePrefixCache
|
|
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=paged_ssd_cache,
|
|
)
|
|
|
|
@staticmethod
|
|
def create_memory_monitor(
|
|
config: CacheConfig,
|
|
paged_cache: Optional["PagedCacheManager"] = None,
|
|
) -> "MemoryMonitor":
|
|
"""
|
|
Create a MemoryMonitor instance.
|
|
|
|
Args:
|
|
config: Cache configuration.
|
|
paged_cache: PagedCacheManager to monitor.
|
|
|
|
Returns:
|
|
Configured MemoryMonitor instance.
|
|
"""
|
|
from ..memory_monitor import MemoryMonitor
|
|
|
|
# Use configured max KV cache memory, or default to 4GB
|
|
max_kv_cache_memory = config.max_kv_cache_memory or (4 * 1024 * 1024 * 1024)
|
|
|
|
monitor = MemoryMonitor(
|
|
max_kv_cache_memory=max_kv_cache_memory,
|
|
)
|
|
|
|
if paged_cache is not None:
|
|
monitor.set_paged_cache_manager(paged_cache)
|
|
|
|
return monitor
|
|
|
|
@staticmethod
|
|
def create_full_cache_stack(
|
|
config: CacheConfig,
|
|
model: Any = None,
|
|
num_layers: Optional[int] = None,
|
|
) -> dict:
|
|
"""
|
|
Create a complete cache stack with all components.
|
|
|
|
This is a convenience method that creates all cache components
|
|
and wires them together. Returns all None values if paged SSD cache
|
|
is disabled (oMLX only supports paged SSD-based caching).
|
|
|
|
Args:
|
|
config: Cache configuration.
|
|
model: Model instance for cache identification.
|
|
num_layers: Number of model layers.
|
|
|
|
Returns:
|
|
Dictionary with keys: paged_cache, paged_ssd_cache, prefix_cache, memory_monitor
|
|
"""
|
|
paged_cache = None
|
|
paged_ssd_cache = None
|
|
prefix_cache = None
|
|
memory_monitor = None
|
|
|
|
# Only create cache components if paged SSD cache is enabled
|
|
if config.paged_ssd_cache_dir is not None:
|
|
paged_cache = CacheFactory.create_paged_cache(config, num_layers)
|
|
paged_ssd_cache = CacheFactory.create_paged_ssd_cache(
|
|
config, config.model_name
|
|
)
|
|
|
|
if paged_cache is not None and paged_ssd_cache is not None:
|
|
paged_cache.set_paged_ssd_cache_manager(paged_ssd_cache)
|
|
|
|
prefix_cache = CacheFactory.create_prefix_cache(
|
|
config, model, paged_cache, paged_ssd_cache
|
|
)
|
|
|
|
if paged_cache is not None:
|
|
memory_monitor = CacheFactory.create_memory_monitor(config, paged_cache)
|
|
|
|
return {
|
|
"paged_cache": paged_cache,
|
|
"paged_ssd_cache": paged_ssd_cache,
|
|
"prefix_cache": prefix_cache,
|
|
"memory_monitor": memory_monitor,
|
|
}
|