132 lines
2.7 KiB
Python
132 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Cache module - unified cache management for oMLX.
|
|
|
|
This package contains cache implementations including:
|
|
- Prefix caching for KV state reuse
|
|
- Paged cache for memory-efficient KV storage
|
|
- VLM cache for vision-language model optimization
|
|
- SSD cache for disk-based persistence
|
|
"""
|
|
|
|
# Stats
|
|
from .stats import (
|
|
BaseCacheStats,
|
|
PagedCacheStats,
|
|
VLMCacheStats,
|
|
PagedSSDCacheStats,
|
|
)
|
|
|
|
# Interfaces
|
|
from .interface import CacheManager
|
|
|
|
# Paged cache implementations
|
|
from .paged_cache import (
|
|
PagedCacheManager,
|
|
CacheBlock,
|
|
BlockTable,
|
|
FreeKVCacheBlockQueue,
|
|
BlockHashToBlockMap,
|
|
BlockHash,
|
|
compute_block_hash,
|
|
)
|
|
|
|
# Prefix cache implementations (SSD-only)
|
|
from .prefix_cache import (
|
|
BlockAwarePrefixCache,
|
|
BlockCacheEntry,
|
|
)
|
|
|
|
# Paged SSD cache implementations
|
|
from .paged_ssd_cache import (
|
|
PagedSSDCacheManager,
|
|
PagedSSDBlockMetadata,
|
|
PagedSSDCacheIndex,
|
|
parse_size,
|
|
)
|
|
|
|
# Vision feature cache
|
|
from .vision_feature_cache import (
|
|
VisionFeatureSSDCache,
|
|
VisionFeatureSSDEntry,
|
|
)
|
|
|
|
# Managers
|
|
from .recovery import CacheRecoveryManager
|
|
|
|
# Factory
|
|
from .factory import CacheConfig, CacheFactory
|
|
|
|
# Type handlers
|
|
from .type_handlers import (
|
|
CacheType,
|
|
CacheTypeHandler,
|
|
CacheStateInfo,
|
|
KVCacheHandler,
|
|
RotatingKVCacheHandler,
|
|
ArraysCacheHandler,
|
|
CacheListHandler,
|
|
DefaultCacheHandler,
|
|
SizedArraysCache,
|
|
)
|
|
|
|
# Type registry
|
|
from .type_registry import CacheTypeRegistry
|
|
|
|
# Hybrid cache config
|
|
from .hybrid_cache import (
|
|
LayerCacheConfig,
|
|
ModelCacheConfig,
|
|
create_default_kvcache_config,
|
|
)
|
|
|
|
__all__ = [
|
|
# Stats
|
|
"BaseCacheStats",
|
|
"PagedCacheStats",
|
|
"VLMCacheStats",
|
|
"PagedSSDCacheStats",
|
|
# Interfaces
|
|
"CacheManager",
|
|
# Paged cache
|
|
"PagedCacheManager",
|
|
"CacheBlock",
|
|
"BlockTable",
|
|
"FreeKVCacheBlockQueue",
|
|
"BlockHashToBlockMap",
|
|
"BlockHash",
|
|
"compute_block_hash",
|
|
# Prefix cache (SSD-only)
|
|
"BlockAwarePrefixCache",
|
|
"BlockCacheEntry",
|
|
# Paged SSD cache
|
|
"PagedSSDCacheManager",
|
|
"PagedSSDBlockMetadata",
|
|
"PagedSSDCacheIndex",
|
|
"parse_size",
|
|
# Vision feature cache
|
|
"VisionFeatureSSDCache",
|
|
"VisionFeatureSSDEntry",
|
|
# Managers
|
|
"CacheRecoveryManager",
|
|
# Factory
|
|
"CacheConfig",
|
|
"CacheFactory",
|
|
# Type handlers
|
|
"CacheType",
|
|
"CacheTypeHandler",
|
|
"CacheStateInfo",
|
|
"KVCacheHandler",
|
|
"RotatingKVCacheHandler",
|
|
"ArraysCacheHandler",
|
|
"CacheListHandler",
|
|
"DefaultCacheHandler",
|
|
"SizedArraysCache",
|
|
# Type registry
|
|
"CacheTypeRegistry",
|
|
# Hybrid cache config
|
|
"LayerCacheConfig",
|
|
"ModelCacheConfig",
|
|
"create_default_kvcache_config",
|
|
]
|