wehub-resource-sync
339 lines
11 KiB
Python
339 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Hybrid cache configuration for models with mixed cache types.
|
|
|
|
This module provides configuration classes for models that use different
|
|
cache types across layers (e.g., Qwen3-Next with ArraysCache + KVCache).
|
|
"""
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, List, Optional, Tuple
|
|
import logging
|
|
|
|
from .type_handlers import CacheType, CacheTypeHandler
|
|
from .type_registry import CacheTypeRegistry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class LayerCacheConfig:
|
|
"""Configuration for a single layer's cache.
|
|
|
|
Attributes:
|
|
layer_idx: Index of the layer
|
|
cache_type: Type of cache used by this layer
|
|
supports_block_slicing: Whether this layer's cache can be block-sliced
|
|
class_name: Original mlx-lm class name
|
|
"""
|
|
|
|
layer_idx: int
|
|
cache_type: CacheType
|
|
supports_block_slicing: bool
|
|
class_name: str = ""
|
|
|
|
@property
|
|
def handler(self) -> CacheTypeHandler:
|
|
"""Get the handler for this layer's cache type."""
|
|
return CacheTypeRegistry.get_handler(self.cache_type)
|
|
|
|
|
|
@dataclass
|
|
class ModelCacheConfig:
|
|
"""Cache configuration for an entire model.
|
|
|
|
Contains per-layer cache type information and model-level metadata.
|
|
This is used to properly handle hybrid models that mix different
|
|
cache types (e.g., KVCache + ArraysCache).
|
|
|
|
Attributes:
|
|
model_name: Name of the model
|
|
num_layers: Total number of layers
|
|
layer_configs: Per-layer cache configurations
|
|
is_hybrid: True if model uses multiple cache types
|
|
sliceable_layer_count: Number of layers that support block slicing
|
|
"""
|
|
|
|
model_name: str = ""
|
|
num_layers: int = 0
|
|
layer_configs: List[LayerCacheConfig] = field(default_factory=list)
|
|
is_hybrid: bool = False
|
|
sliceable_layer_count: int = 0
|
|
|
|
@classmethod
|
|
def from_cache_list(
|
|
cls,
|
|
cache_list: List[Any],
|
|
model_name: str = "",
|
|
) -> "ModelCacheConfig":
|
|
"""Create configuration from mlx-lm cache list.
|
|
|
|
Args:
|
|
cache_list: List of cache objects from model.make_cache()
|
|
model_name: Optional model name for identification
|
|
|
|
Returns:
|
|
ModelCacheConfig with per-layer type information
|
|
"""
|
|
if not cache_list:
|
|
return cls(model_name=model_name)
|
|
|
|
layer_configs = []
|
|
cache_types_seen = set()
|
|
sliceable_count = 0
|
|
max_window_size = 0
|
|
|
|
for idx, cache_obj in enumerate(cache_list):
|
|
cache_type = CacheTypeRegistry.detect_cache_type(cache_obj)
|
|
handler = CacheTypeRegistry.get_handler(cache_type)
|
|
class_name = type(cache_obj).__name__
|
|
|
|
cache_types_seen.add(cache_type)
|
|
if handler.supports_block_slicing:
|
|
sliceable_count += 1
|
|
|
|
# Extract window_size from RotatingKVCache layers
|
|
if cache_type == CacheType.ROTATING_KVCACHE:
|
|
window_size = getattr(cache_obj, "max_size", 0)
|
|
if window_size > max_window_size:
|
|
max_window_size = window_size
|
|
|
|
# Extract window_size from CacheList sub-caches (e.g., RotatingKVCache inside)
|
|
if cache_type == CacheType.CACHE_LIST:
|
|
sub_caches = getattr(cache_obj, "caches", ())
|
|
for sub_c in sub_caches:
|
|
sub_type = CacheTypeRegistry.detect_cache_type(sub_c)
|
|
if sub_type == CacheType.ROTATING_KVCACHE:
|
|
ws = getattr(sub_c, "max_size", 0)
|
|
if ws > max_window_size:
|
|
max_window_size = ws
|
|
|
|
layer_configs.append(
|
|
LayerCacheConfig(
|
|
layer_idx=idx,
|
|
cache_type=cache_type,
|
|
supports_block_slicing=handler.supports_block_slicing,
|
|
class_name=class_name,
|
|
)
|
|
)
|
|
|
|
config = cls(
|
|
model_name=model_name,
|
|
num_layers=len(cache_list),
|
|
layer_configs=layer_configs,
|
|
is_hybrid=len(cache_types_seen) > 1,
|
|
sliceable_layer_count=sliceable_count,
|
|
)
|
|
config._max_window_size = max_window_size
|
|
return config
|
|
|
|
@classmethod
|
|
def from_type_list(
|
|
cls,
|
|
cache_types: List[str],
|
|
model_name: str = "",
|
|
) -> "ModelCacheConfig":
|
|
"""Create configuration from list of type names.
|
|
|
|
Useful for reconstructing config from serialized metadata.
|
|
|
|
Args:
|
|
cache_types: List of cache type names (e.g., ["KVCache", "ArraysCache"])
|
|
model_name: Optional model name
|
|
|
|
Returns:
|
|
ModelCacheConfig
|
|
"""
|
|
if not cache_types:
|
|
return cls(model_name=model_name)
|
|
|
|
layer_configs = []
|
|
cache_types_seen = set()
|
|
sliceable_count = 0
|
|
|
|
for idx, type_name in enumerate(cache_types):
|
|
handler = CacheTypeRegistry.get_handler_by_class_name(type_name)
|
|
cache_type = handler.cache_type
|
|
|
|
cache_types_seen.add(cache_type)
|
|
if handler.supports_block_slicing:
|
|
sliceable_count += 1
|
|
|
|
layer_configs.append(
|
|
LayerCacheConfig(
|
|
layer_idx=idx,
|
|
cache_type=cache_type,
|
|
supports_block_slicing=handler.supports_block_slicing,
|
|
class_name=type_name,
|
|
)
|
|
)
|
|
|
|
return cls(
|
|
model_name=model_name,
|
|
num_layers=len(cache_types),
|
|
layer_configs=layer_configs,
|
|
is_hybrid=len(cache_types_seen) > 1,
|
|
sliceable_layer_count=sliceable_count,
|
|
)
|
|
|
|
def get_sliceable_layers(self) -> List[int]:
|
|
"""Get indices of layers that support block slicing.
|
|
|
|
Returns:
|
|
List of layer indices
|
|
"""
|
|
return [cfg.layer_idx for cfg in self.layer_configs if cfg.supports_block_slicing]
|
|
|
|
def get_non_sliceable_layers(self) -> List[int]:
|
|
"""Get indices of layers that don't support block slicing.
|
|
|
|
Returns:
|
|
List of layer indices (e.g., ArraysCache, RotatingKVCache layers)
|
|
"""
|
|
return [
|
|
cfg.layer_idx for cfg in self.layer_configs if not cfg.supports_block_slicing
|
|
]
|
|
|
|
def get_layer_type(self, layer_idx: int) -> CacheType:
|
|
"""Get cache type for a specific layer.
|
|
|
|
Args:
|
|
layer_idx: Layer index
|
|
|
|
Returns:
|
|
CacheType enum
|
|
"""
|
|
if 0 <= layer_idx < len(self.layer_configs):
|
|
return self.layer_configs[layer_idx].cache_type
|
|
return CacheType.KVCACHE # Default
|
|
|
|
def get_layer_handler(self, layer_idx: int) -> CacheTypeHandler:
|
|
"""Get handler for a specific layer.
|
|
|
|
Args:
|
|
layer_idx: Layer index
|
|
|
|
Returns:
|
|
CacheTypeHandler instance
|
|
"""
|
|
if 0 <= layer_idx < len(self.layer_configs):
|
|
return self.layer_configs[layer_idx].handler
|
|
return CacheTypeRegistry.get_handler(CacheType.KVCACHE)
|
|
|
|
def get_type_names(self) -> List[str]:
|
|
"""Get list of cache type names for serialization.
|
|
|
|
Returns:
|
|
List of class name strings
|
|
"""
|
|
return [cfg.class_name for cfg in self.layer_configs]
|
|
|
|
def get_meta_states(self, cache_list: List[Any]) -> List[Tuple]:
|
|
"""Extract meta_states from cache objects.
|
|
|
|
For CacheList layers, the meta_state is a composite:
|
|
([class_names], [sub_meta_states]).
|
|
|
|
Args:
|
|
cache_list: List of cache objects
|
|
|
|
Returns:
|
|
List of meta_state tuples
|
|
"""
|
|
meta_states = []
|
|
for idx, cache_obj in enumerate(cache_list):
|
|
if idx < len(self.layer_configs):
|
|
handler = self.layer_configs[idx].handler
|
|
|
|
if self.layer_configs[idx].cache_type == CacheType.CACHE_LIST:
|
|
# CacheList: extract composite meta_state from sub-caches
|
|
state = handler.extract_state(cache_obj)
|
|
meta_state = (
|
|
state.get("sub_class_names", []),
|
|
state.get("sub_meta_states", []),
|
|
)
|
|
meta_states.append(meta_state)
|
|
else:
|
|
state = handler.extract_state(cache_obj)
|
|
meta_state = state.get("meta_state", ())
|
|
if not meta_state:
|
|
# Try to extract offset at minimum
|
|
offset = state.get("offset", 0)
|
|
meta_state = (offset,)
|
|
meta_states.append(meta_state)
|
|
else:
|
|
meta_states.append(())
|
|
return meta_states
|
|
|
|
def supports_full_block_slicing(self) -> bool:
|
|
"""Check if all layers support block slicing.
|
|
|
|
Returns:
|
|
True if all layers can be block-sliced
|
|
"""
|
|
return self.sliceable_layer_count == self.num_layers
|
|
|
|
def has_rotating_layers(self) -> bool:
|
|
"""Check if model has any RotatingKVCache layers.
|
|
|
|
Returns:
|
|
True if any layer uses RotatingKVCache
|
|
"""
|
|
return any(
|
|
cfg.cache_type == CacheType.ROTATING_KVCACHE for cfg in self.layer_configs
|
|
)
|
|
|
|
def get_max_window_size(self) -> int:
|
|
"""Get maximum window size from RotatingKVCache layers.
|
|
|
|
The window size is extracted from RotatingKVCache meta_state during
|
|
from_cache_list() construction, or can be set manually via
|
|
_max_window_size attribute.
|
|
|
|
Returns:
|
|
Maximum window_size, or 0 if no RotatingKVCache layers
|
|
"""
|
|
return getattr(self, "_max_window_size", 0)
|
|
|
|
def __repr__(self) -> str:
|
|
types = [cfg.class_name for cfg in self.layer_configs[:5]]
|
|
if len(self.layer_configs) > 5:
|
|
types.append(f"... +{len(self.layer_configs) - 5} more")
|
|
return (
|
|
f"ModelCacheConfig(model='{self.model_name}', "
|
|
f"layers={self.num_layers}, "
|
|
f"hybrid={self.is_hybrid}, "
|
|
f"types={types})"
|
|
)
|
|
|
|
|
|
def create_default_kvcache_config(num_layers: int, model_name: str = "") -> ModelCacheConfig:
|
|
"""Create a default KVCache-only configuration.
|
|
|
|
Convenience function for models that use only KVCache.
|
|
|
|
Args:
|
|
num_layers: Number of model layers
|
|
model_name: Optional model name
|
|
|
|
Returns:
|
|
ModelCacheConfig with all KVCache layers
|
|
"""
|
|
layer_configs = [
|
|
LayerCacheConfig(
|
|
layer_idx=idx,
|
|
cache_type=CacheType.KVCACHE,
|
|
supports_block_slicing=True,
|
|
class_name="KVCache",
|
|
)
|
|
for idx in range(num_layers)
|
|
]
|
|
|
|
return ModelCacheConfig(
|
|
model_name=model_name,
|
|
num_layers=num_layers,
|
|
layer_configs=layer_configs,
|
|
is_hybrid=False,
|
|
sliceable_layer_count=num_layers,
|
|
)
|