wehub-resource-sync
1134 lines
42 KiB
Python
1134 lines
42 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Tests for hybrid caching with RotatingKVCache support.
|
|
|
|
This module tests the last-block-only storage strategy for RotatingKVCache
|
|
layers and strict partial-prefix rejection for boundary-safe restore.
|
|
"""
|
|
|
|
import math
|
|
import time
|
|
from typing import Any, Dict, List, Optional
|
|
from unittest.mock import MagicMock, PropertyMock, patch
|
|
|
|
import pytest
|
|
|
|
try:
|
|
import mlx.core as mx
|
|
|
|
HAS_MLX = True
|
|
except ImportError:
|
|
HAS_MLX = False
|
|
mx = None
|
|
|
|
from omlx.cache.hybrid_cache import (
|
|
LayerCacheConfig,
|
|
ModelCacheConfig,
|
|
create_default_kvcache_config,
|
|
)
|
|
from omlx.cache.paged_cache import (
|
|
BlockTable,
|
|
PagedCacheManager,
|
|
compute_block_hash,
|
|
)
|
|
from omlx.cache.prefix_cache import BlockAwarePrefixCache, BlockCacheEntry
|
|
from omlx.cache.type_handlers import CacheType
|
|
|
|
|
|
class MockModel:
|
|
"""Mock model for testing."""
|
|
|
|
def __init__(self, num_layers: int = 4):
|
|
self._num_layers = num_layers
|
|
self.layers = [MagicMock() for _ in range(num_layers)]
|
|
|
|
@property
|
|
def args(self):
|
|
mock_args = MagicMock()
|
|
mock_args.num_hidden_layers = self._num_layers
|
|
return mock_args
|
|
|
|
|
|
class TestModelCacheConfigWindowSize:
|
|
"""Tests for ModelCacheConfig window size functionality."""
|
|
|
|
def test_get_max_window_size_default(self):
|
|
"""Test get_max_window_size returns 0 when not set."""
|
|
config = create_default_kvcache_config(4)
|
|
assert config.get_max_window_size() == 0
|
|
|
|
def test_get_max_window_size_set(self):
|
|
"""Test get_max_window_size returns stored value."""
|
|
config = create_default_kvcache_config(4)
|
|
config._max_window_size = 1024
|
|
assert config.get_max_window_size() == 1024
|
|
|
|
def test_from_cache_list_extracts_window_size(self):
|
|
"""Test from_cache_list extracts window_size from RotatingKVCache."""
|
|
# Create mock cache objects with spec to avoid spurious attributes
|
|
mock_kv = MagicMock(spec=[])
|
|
mock_kv.__class__ = type("KVCache", (), {})
|
|
mock_kv.state = (MagicMock(), MagicMock())
|
|
mock_kv.offset = 64
|
|
mock_kv.keys = MagicMock()
|
|
mock_kv.values = MagicMock()
|
|
|
|
mock_rotating = MagicMock(spec=[])
|
|
mock_rotating.__class__ = type("RotatingKVCache", (), {})
|
|
mock_rotating.state = (MagicMock(), MagicMock())
|
|
mock_rotating.max_size = 1024
|
|
mock_rotating._idx = 0
|
|
mock_rotating.offset = 0
|
|
|
|
cache_list = [mock_kv, mock_rotating, mock_rotating, mock_kv]
|
|
config = ModelCacheConfig.from_cache_list(cache_list)
|
|
|
|
assert config.is_hybrid is True
|
|
assert config.has_rotating_layers() is True
|
|
assert config.get_max_window_size() == 1024
|
|
|
|
def test_from_cache_list_no_rotating(self):
|
|
"""Test from_cache_list with KVCache-only model."""
|
|
mock_kv = MagicMock(spec=[])
|
|
mock_kv.__class__ = type("KVCache", (), {})
|
|
mock_kv.state = (MagicMock(), MagicMock())
|
|
mock_kv.offset = 64
|
|
mock_kv.keys = MagicMock()
|
|
mock_kv.values = MagicMock()
|
|
|
|
cache_list = [mock_kv, mock_kv, mock_kv, mock_kv]
|
|
config = ModelCacheConfig.from_cache_list(cache_list)
|
|
|
|
assert config.is_hybrid is False
|
|
assert config.has_rotating_layers() is False
|
|
assert config.get_max_window_size() == 0
|
|
|
|
def test_from_cache_list_max_window_size_selected(self):
|
|
"""Test from_cache_list picks the maximum window_size."""
|
|
mock_rotating_small = MagicMock()
|
|
mock_rotating_small.__class__ = type("RotatingKVCache", (), {})
|
|
mock_rotating_small.__class__.__name__ = "RotatingKVCache"
|
|
mock_rotating_small.max_size = 512
|
|
mock_rotating_small._idx = 0
|
|
|
|
mock_rotating_large = MagicMock()
|
|
mock_rotating_large.__class__ = type("RotatingKVCache", (), {})
|
|
mock_rotating_large.__class__.__name__ = "RotatingKVCache"
|
|
mock_rotating_large.max_size = 2048
|
|
mock_rotating_large._idx = 0
|
|
|
|
cache_list = [mock_rotating_small, mock_rotating_large]
|
|
config = ModelCacheConfig.from_cache_list(cache_list)
|
|
|
|
assert config.get_max_window_size() == 2048
|
|
|
|
|
|
class TestApplyWindowPadding:
|
|
"""Tests for _apply_window_padding method."""
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self):
|
|
"""Create a BlockAwarePrefixCache with block_size=256."""
|
|
paged_cache = PagedCacheManager(
|
|
block_size=256,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=4)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
def test_no_config(self, prefix_cache):
|
|
"""Test with no model_cache_config returns all blocks."""
|
|
result = prefix_cache._apply_window_padding(16, None)
|
|
assert result == 16
|
|
|
|
def test_kvcache_only_model(self, prefix_cache):
|
|
"""Test KVCache-only model returns all blocks."""
|
|
config = create_default_kvcache_config(4)
|
|
result = prefix_cache._apply_window_padding(16, config)
|
|
assert result == 16
|
|
|
|
def test_hybrid_model_with_rotating(self, prefix_cache):
|
|
"""Test hybrid model subtracts padding blocks."""
|
|
# Gemma3-like: window_size=1024, block_size=256
|
|
config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "RotatingKVCache", "RotatingKVCache", "KVCache"]
|
|
)
|
|
config._max_window_size = 1024
|
|
|
|
# 16 blocks matched -> restore 12 (16 - ceil(1024/256) = 16 - 4 = 12)
|
|
result = prefix_cache._apply_window_padding(16, config)
|
|
assert result == 12
|
|
|
|
def test_padding_exceeds_matched(self, prefix_cache):
|
|
"""Test when padding blocks exceed matched blocks, returns 0."""
|
|
config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "RotatingKVCache"]
|
|
)
|
|
config._max_window_size = 1024
|
|
|
|
# 3 blocks matched, need 4 for padding -> 0
|
|
result = prefix_cache._apply_window_padding(3, config)
|
|
assert result == 0
|
|
|
|
def test_exact_padding(self, prefix_cache):
|
|
"""Test when matched blocks exactly equal padding blocks."""
|
|
config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "RotatingKVCache"]
|
|
)
|
|
config._max_window_size = 1024
|
|
|
|
# 4 blocks matched, need 4 for padding -> 0
|
|
result = prefix_cache._apply_window_padding(4, config)
|
|
assert result == 0
|
|
|
|
def test_window_size_not_multiple_of_block_size(self, prefix_cache):
|
|
"""Test with window_size not a multiple of block_size."""
|
|
config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "RotatingKVCache"]
|
|
)
|
|
config._max_window_size = 1000 # Not a multiple of 256
|
|
|
|
# ceil(1000/256) = 4 padding blocks
|
|
result = prefix_cache._apply_window_padding(16, config)
|
|
assert result == 12
|
|
|
|
def test_window_size_zero(self, prefix_cache):
|
|
"""Test with window_size=0 returns all blocks."""
|
|
config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "RotatingKVCache"]
|
|
)
|
|
config._max_window_size = 0
|
|
|
|
result = prefix_cache._apply_window_padding(16, config)
|
|
assert result == 16
|
|
|
|
def test_small_block_size(self):
|
|
"""Test with small block_size (more padding blocks needed)."""
|
|
paged_cache = PagedCacheManager(
|
|
block_size=64,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=4)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "RotatingKVCache"]
|
|
)
|
|
config._max_window_size = 1024
|
|
|
|
# ceil(1024/64) = 16 padding blocks
|
|
result = cache._apply_window_padding(20, config)
|
|
assert result == 4
|
|
|
|
|
|
@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
|
|
class TestExtractBlockTensorSliceLastBlock:
|
|
"""Tests for _extract_block_tensor_slice with is_last_block parameter."""
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self):
|
|
"""Create a BlockAwarePrefixCache."""
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=4)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
def _make_kvcache_layer(self, seq_len=64):
|
|
"""Create a mock KVCache layer state."""
|
|
return {
|
|
'state': (mx.zeros((1, 8, seq_len, 64)), mx.zeros((1, 8, seq_len, 64))),
|
|
'cache_type': 'KVCache',
|
|
}
|
|
|
|
def _make_rotating_layer(self, max_size=256):
|
|
"""Create a mock RotatingKVCache layer state."""
|
|
return {
|
|
'state': (mx.zeros((1, 8, max_size, 64)), mx.zeros((1, 8, max_size, 64))),
|
|
'cache_type': 'RotatingKVCache',
|
|
}
|
|
|
|
def test_kvcache_only_slicing_unchanged(self, prefix_cache):
|
|
"""Test KVCache-only model slicing is unchanged regardless of is_last_block."""
|
|
cache_data = [self._make_kvcache_layer(64) for _ in range(4)]
|
|
|
|
# Non-last block
|
|
result_non_last = prefix_cache._extract_block_tensor_slice(
|
|
cache_data, 0, 4, is_last_block=False
|
|
)
|
|
assert result_non_last is not None
|
|
assert len(result_non_last) == 4
|
|
for keys, values in result_non_last:
|
|
assert keys.shape == (1, 8, 4, 64)
|
|
|
|
# Last block
|
|
result_last = prefix_cache._extract_block_tensor_slice(
|
|
cache_data, 0, 4, is_last_block=True
|
|
)
|
|
assert result_last is not None
|
|
assert len(result_last) == 4
|
|
for keys, values in result_last:
|
|
assert keys.shape == (1, 8, 4, 64)
|
|
|
|
def test_rotating_non_last_block_placeholder(self, prefix_cache):
|
|
"""Test RotatingKVCache non-last block stores placeholder."""
|
|
cache_data = [
|
|
self._make_kvcache_layer(64),
|
|
self._make_rotating_layer(256),
|
|
]
|
|
config = ModelCacheConfig.from_type_list(["KVCache", "RotatingKVCache"])
|
|
|
|
result = prefix_cache._extract_block_tensor_slice(
|
|
cache_data, 0, 4, model_cache_config=config, is_last_block=False
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result) == 2
|
|
|
|
# KVCache layer: normal slice
|
|
keys0, values0 = result[0]
|
|
assert keys0.shape == (1, 8, 4, 64)
|
|
|
|
# RotatingKVCache layer: placeholder
|
|
keys1, values1 = result[1]
|
|
assert keys1.shape == (1,)
|
|
assert values1.shape == (1,)
|
|
|
|
def test_rotating_last_block_full_state(self, prefix_cache):
|
|
"""Test RotatingKVCache last block stores full state."""
|
|
cache_data = [
|
|
self._make_kvcache_layer(64),
|
|
self._make_rotating_layer(256),
|
|
]
|
|
config = ModelCacheConfig.from_type_list(["KVCache", "RotatingKVCache"])
|
|
|
|
result = prefix_cache._extract_block_tensor_slice(
|
|
cache_data, 0, 4, model_cache_config=config, is_last_block=True
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result) == 2
|
|
|
|
# KVCache layer: normal slice
|
|
keys0, values0 = result[0]
|
|
assert keys0.shape == (1, 8, 4, 64)
|
|
|
|
# RotatingKVCache layer: full state
|
|
keys1, values1 = result[1]
|
|
assert keys1.shape == (1, 8, 256, 64) # Full RotatingKVCache state
|
|
assert values1.shape == (1, 8, 256, 64)
|
|
|
|
def test_hybrid_model_multiple_blocks(self, prefix_cache):
|
|
"""Test storing multiple blocks for a hybrid model."""
|
|
# Simulate Gemma3-like model: KVCache + RotatingKVCache
|
|
cache_data = [
|
|
self._make_kvcache_layer(16), # KVCache with 16 tokens
|
|
self._make_rotating_layer(256), # RotatingKVCache with window=256
|
|
]
|
|
config = ModelCacheConfig.from_type_list(["KVCache", "RotatingKVCache"])
|
|
|
|
# Block 0 (non-last): KVCache sliced, RotatingKVCache placeholder
|
|
block0 = prefix_cache._extract_block_tensor_slice(
|
|
cache_data, 0, 4, model_cache_config=config, is_last_block=False
|
|
)
|
|
assert block0 is not None
|
|
assert block0[0][0].shape == (1, 8, 4, 64) # KVCache slice
|
|
assert block0[1][0].shape == (1,) # RotatingKVCache placeholder
|
|
|
|
# Block 1 (non-last): KVCache sliced, RotatingKVCache placeholder
|
|
block1 = prefix_cache._extract_block_tensor_slice(
|
|
cache_data, 4, 8, model_cache_config=config, is_last_block=False
|
|
)
|
|
assert block1 is not None
|
|
assert block1[0][0].shape == (1, 8, 4, 64)
|
|
assert block1[1][0].shape == (1,)
|
|
|
|
# Block 3 (last): KVCache sliced, RotatingKVCache full state
|
|
block3 = prefix_cache._extract_block_tensor_slice(
|
|
cache_data, 12, 16, model_cache_config=config, is_last_block=True
|
|
)
|
|
assert block3 is not None
|
|
assert block3[0][0].shape == (1, 8, 4, 64)
|
|
assert block3[1][0].shape == (1, 8, 256, 64) # Full state
|
|
|
|
|
|
@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
|
|
class TestValidateBlockCacheDataPlaceholder:
|
|
"""Tests for _validate_block_cache_data with RotatingKVCache placeholders."""
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self):
|
|
"""Create a BlockAwarePrefixCache."""
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=4)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
def test_placeholder_passes_validation(self, prefix_cache):
|
|
"""Test placeholder (1D tensor) passes validation for RotatingKVCache."""
|
|
# KVCache: normal 4D tensor, RotatingKVCache: placeholder
|
|
cache_data = [
|
|
(mx.zeros((1, 8, 4, 64)), mx.zeros((1, 8, 4, 64))), # KVCache
|
|
(mx.zeros((1,)), mx.zeros((1,))), # RotatingKVCache placeholder
|
|
]
|
|
layer_cache_types = ["KVCache", "RotatingKVCache"]
|
|
|
|
result = prefix_cache._validate_block_cache_data(cache_data, layer_cache_types)
|
|
assert result is True
|
|
|
|
def test_full_rotating_state_passes_validation(self, prefix_cache):
|
|
"""Test full RotatingKVCache state passes validation."""
|
|
cache_data = [
|
|
(mx.zeros((1, 8, 4, 64)), mx.zeros((1, 8, 4, 64))), # KVCache
|
|
(mx.zeros((1, 8, 256, 64)), mx.zeros((1, 8, 256, 64))), # RotatingKVCache full
|
|
]
|
|
layer_cache_types = ["KVCache", "RotatingKVCache"]
|
|
|
|
result = prefix_cache._validate_block_cache_data(cache_data, layer_cache_types)
|
|
assert result is True
|
|
|
|
|
|
class TestDetectWindowPaddingFromBlocks:
|
|
"""Tests for _detect_window_padding_from_blocks method."""
|
|
|
|
@pytest.fixture
|
|
def paged_cache(self):
|
|
"""Create a PagedCacheManager."""
|
|
return PagedCacheManager(
|
|
block_size=256,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def mock_ssd_cache(self):
|
|
"""Create a mock SSD cache manager."""
|
|
mock = MagicMock()
|
|
mock.save_block.return_value = True
|
|
mock.load_block.return_value = None
|
|
mock.load_block_with_metadata.return_value = (None, None)
|
|
return mock
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self, paged_cache, mock_ssd_cache):
|
|
"""Create a BlockAwarePrefixCache with SSD."""
|
|
model = MockModel(num_layers=4)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd_cache,
|
|
)
|
|
|
|
def test_no_block_ids(self, prefix_cache):
|
|
"""Test with empty block_ids returns None."""
|
|
result = prefix_cache._detect_window_padding_from_blocks([])
|
|
assert result is None
|
|
|
|
def test_no_ssd_cache(self, paged_cache):
|
|
"""Test without SSD cache returns None."""
|
|
model = MockModel(num_layers=4)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
result = cache._detect_window_padding_from_blocks([1, 2, 3])
|
|
assert result is None
|
|
|
|
def test_no_rotating_kvcache_in_metadata(self, prefix_cache, mock_ssd_cache, paged_cache):
|
|
"""Test with KVCache-only metadata returns None."""
|
|
# Allocate a block
|
|
block = paged_cache.allocate_block()
|
|
block.block_hash = b"test_hash"
|
|
|
|
mock_ssd_cache.load_block_with_metadata.return_value = (
|
|
[(MagicMock(), MagicMock())],
|
|
{
|
|
'layer_cache_types': ['KVCache', 'KVCache'],
|
|
'layer_meta_states': [(64,), (64,)],
|
|
}
|
|
)
|
|
|
|
result = prefix_cache._detect_window_padding_from_blocks([block.block_id])
|
|
assert result is None
|
|
|
|
def test_rotating_kvcache_detected(self, prefix_cache, mock_ssd_cache, paged_cache):
|
|
"""Test RotatingKVCache is detected and config built correctly."""
|
|
block = paged_cache.allocate_block()
|
|
block.block_hash = b"test_hash"
|
|
|
|
mock_ssd_cache.load_block_with_metadata.return_value = (
|
|
[(MagicMock(), MagicMock()), (MagicMock(), MagicMock())],
|
|
{
|
|
'layer_cache_types': ['KVCache', 'RotatingKVCache'],
|
|
'layer_meta_states': [
|
|
(64,), # KVCache: (offset,)
|
|
(4, 1024, 500, 100), # RotatingKVCache: (keep, max_size, offset, _idx)
|
|
],
|
|
}
|
|
)
|
|
|
|
result = prefix_cache._detect_window_padding_from_blocks([block.block_id])
|
|
|
|
assert result is not None
|
|
assert result.has_rotating_layers() is True
|
|
assert result.get_max_window_size() == 1024
|
|
|
|
def test_max_window_size_from_multiple_rotating_layers(
|
|
self, prefix_cache, mock_ssd_cache, paged_cache
|
|
):
|
|
"""Test max window_size is selected from multiple RotatingKVCache layers."""
|
|
block = paged_cache.allocate_block()
|
|
block.block_hash = b"test_hash"
|
|
|
|
mock_ssd_cache.load_block_with_metadata.return_value = (
|
|
[(MagicMock(), MagicMock())] * 3,
|
|
{
|
|
'layer_cache_types': ['RotatingKVCache', 'KVCache', 'RotatingKVCache'],
|
|
'layer_meta_states': [
|
|
(4, 512, 500, 100), # RotatingKVCache: max_size=512
|
|
(64,), # KVCache
|
|
(4, 2048, 500, 100), # RotatingKVCache: max_size=2048
|
|
],
|
|
}
|
|
)
|
|
|
|
result = prefix_cache._detect_window_padding_from_blocks([block.block_id])
|
|
|
|
assert result is not None
|
|
assert result.get_max_window_size() == 2048
|
|
|
|
|
|
class TestFetchCachePrefixMatching:
|
|
"""Tests for fetch_cache prefix matching behavior."""
|
|
|
|
@pytest.fixture
|
|
def paged_cache(self):
|
|
"""Create a PagedCacheManager."""
|
|
return PagedCacheManager(
|
|
block_size=256,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def mock_ssd_cache(self):
|
|
"""Create a mock SSD cache manager."""
|
|
mock = MagicMock()
|
|
mock.save_block.return_value = True
|
|
mock.load_block.return_value = None
|
|
mock.load_block_with_metadata.return_value = (None, None)
|
|
return mock
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self, paged_cache, mock_ssd_cache):
|
|
"""Create a BlockAwarePrefixCache with SSD."""
|
|
model = MockModel(num_layers=4)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd_cache,
|
|
)
|
|
|
|
def test_fetch_cache_miss_no_padding(self, prefix_cache):
|
|
"""Test cache miss doesn't trigger window padding."""
|
|
tokens = list(range(1024))
|
|
block_table, remaining = prefix_cache.fetch_cache("req-001", tokens)
|
|
|
|
assert block_table is None
|
|
assert remaining == tokens
|
|
|
|
def test_fetch_cache_hit_no_rotating(self, prefix_cache, paged_cache, mock_ssd_cache):
|
|
"""Test cache hit without RotatingKVCache doesn't apply padding."""
|
|
# Set up blocks in paged cache for shared prefix
|
|
tokens = list(range(256))
|
|
block = paged_cache.allocate_block()
|
|
block.token_count = 256
|
|
parent_hash = compute_block_hash(None, tokens, model_name="test-model")
|
|
block.block_hash = parent_hash
|
|
paged_cache.register_block_hash(block, tokens, None)
|
|
|
|
# SSD returns KVCache-only metadata
|
|
mock_ssd_cache.load_block_with_metadata.return_value = (
|
|
[(MagicMock(), MagicMock())],
|
|
{
|
|
'layer_cache_types': ['KVCache'],
|
|
'layer_meta_states': [(64,)],
|
|
}
|
|
)
|
|
|
|
# Extend tokens so there are remaining
|
|
extended_tokens = tokens + list(range(256, 512))
|
|
block_table, remaining = prefix_cache.fetch_cache("req-001", extended_tokens)
|
|
|
|
assert block_table is not None
|
|
assert len(remaining) == 256
|
|
|
|
def test_fetch_cache_hit_with_rotating_keeps_matched_blocks(
|
|
self, prefix_cache, paged_cache, mock_ssd_cache
|
|
):
|
|
"""Rotating models no longer prune matched blocks via window padding."""
|
|
tokens = list(range(256))
|
|
block = paged_cache.allocate_block()
|
|
block.token_count = 256
|
|
parent_hash = compute_block_hash(None, tokens, model_name="test-model")
|
|
block.block_hash = parent_hash
|
|
paged_cache.register_block_hash(block, tokens, None)
|
|
|
|
# SSD returns RotatingKVCache metadata with window_size=1024
|
|
mock_ssd_cache.load_block_with_metadata.return_value = (
|
|
[(MagicMock(), MagicMock()), (MagicMock(), MagicMock())],
|
|
{
|
|
'layer_cache_types': ['KVCache', 'RotatingKVCache'],
|
|
'layer_meta_states': [
|
|
(64,),
|
|
(4, 1024, 500, 100),
|
|
],
|
|
}
|
|
)
|
|
|
|
extended_tokens = tokens + list(range(256, 512))
|
|
block_table, remaining = prefix_cache.fetch_cache("req-001", extended_tokens)
|
|
|
|
assert block_table is not None
|
|
assert block_table.num_tokens == 256
|
|
assert remaining == extended_tokens[256:]
|
|
|
|
|
|
@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
|
|
class TestCreateEmptyRotatingCache:
|
|
"""Tests for _create_empty_rotating_cache method."""
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self):
|
|
"""Create a BlockAwarePrefixCache."""
|
|
paged_cache = PagedCacheManager(
|
|
block_size=256,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=4)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
def test_create_with_valid_meta_state_and_shape(self, prefix_cache):
|
|
"""Test creating empty RotatingKVCache with zero-length keys."""
|
|
meta_state = (4, 1024, 500, 100) # (keep, max_size, offset, _idx)
|
|
kv_shape_ref = (8, 64) # (kv_heads, head_dim)
|
|
cache = prefix_cache._create_empty_rotating_cache(
|
|
meta_state, kvcache_offset=768, kv_shape_ref=kv_shape_ref,
|
|
)
|
|
|
|
assert cache is not None
|
|
assert cache.max_size == 1024
|
|
assert cache.keep == 4
|
|
assert cache.offset == 768
|
|
# Zero-length keys so empty() returns False
|
|
assert cache.keys is not None
|
|
assert cache.keys.shape == (1, 8, 0, 64)
|
|
assert cache.values is not None
|
|
assert cache.values.shape == (1, 8, 0, 64)
|
|
assert cache._idx == 0
|
|
assert not cache.empty()
|
|
|
|
def test_create_without_shape_ref(self, prefix_cache):
|
|
"""Test creating with no shape ref falls back to keys=None."""
|
|
meta_state = (0, 512)
|
|
cache = prefix_cache._create_empty_rotating_cache(meta_state)
|
|
|
|
assert cache is not None
|
|
assert cache.max_size == 512
|
|
assert cache.keep == 0
|
|
# Without shape ref, keys remain None
|
|
assert cache.keys is None
|
|
|
|
def test_create_with_none_meta_state(self, prefix_cache):
|
|
"""Test creating with None meta_state returns None."""
|
|
cache = prefix_cache._create_empty_rotating_cache(None)
|
|
assert cache is None
|
|
|
|
def test_create_with_empty_meta_state(self, prefix_cache):
|
|
"""Test creating with empty tuple returns None."""
|
|
cache = prefix_cache._create_empty_rotating_cache(())
|
|
assert cache is None
|
|
|
|
def test_create_with_short_meta_state(self, prefix_cache):
|
|
"""Test creating with meta_state shorter than 2 returns None."""
|
|
cache = prefix_cache._create_empty_rotating_cache((64,))
|
|
assert cache is None
|
|
|
|
def test_empty_cache_reports_size_zero(self, prefix_cache):
|
|
"""Test that empty RotatingKVCache reports size=0, not min(offset, max_size).
|
|
|
|
This is critical for BatchRotatingKVCache.merge(): standard
|
|
RotatingKVCache.size() returns min(offset, max_size) which incorrectly
|
|
claims data exists when keys are zero-length, causing merge() to
|
|
create unmasked zero-filled buffers that dilute attention scores.
|
|
"""
|
|
meta_state = (0, 128, 500, 64) # (keep, max_size, offset, _idx)
|
|
kv_shape_ref = (8, 64)
|
|
cache = prefix_cache._create_empty_rotating_cache(
|
|
meta_state, kvcache_offset=512, kv_shape_ref=kv_shape_ref,
|
|
)
|
|
|
|
assert cache is not None
|
|
# Standard RotatingKVCache.size() would return min(512, 128) = 128
|
|
# Our subclass must return 0 for zero-length keys
|
|
assert cache.size() == 0
|
|
# offset is still correct for RoPE alignment
|
|
assert cache.offset == 512
|
|
|
|
def test_prefill_ready_subclass_type(self, prefix_cache):
|
|
"""Test _create_empty_rotating_cache returns the correct subclass type."""
|
|
from mlx_lm.models.cache import RotatingKVCache
|
|
|
|
meta_state = (0, 128)
|
|
kv_shape_ref = (8, 64)
|
|
cache = prefix_cache._create_empty_rotating_cache(
|
|
meta_state, kvcache_offset=256, kv_shape_ref=kv_shape_ref,
|
|
)
|
|
|
|
assert cache is not None
|
|
# Must be a RotatingKVCache (for merge/empty compatibility)
|
|
assert isinstance(cache, RotatingKVCache)
|
|
# But with overridden size()
|
|
assert cache.size() == 0
|
|
|
|
def test_prefill_ready_size_normal_after_data(self, prefix_cache):
|
|
"""Test that size() returns normal value after keys have data."""
|
|
meta_state = (0, 128)
|
|
kv_shape_ref = (8, 64)
|
|
cache = prefix_cache._create_empty_rotating_cache(
|
|
meta_state, kvcache_offset=256, kv_shape_ref=kv_shape_ref,
|
|
)
|
|
|
|
assert cache.size() == 0
|
|
|
|
# Simulate adding data (as happens during prefill)
|
|
cache.keys = mx.zeros((1, 8, 64, 64))
|
|
cache.values = mx.zeros((1, 8, 64, 64))
|
|
cache.offset = 64
|
|
# Now size() should return normal value: min(64, 128) = 64
|
|
assert cache.size() == 64
|
|
|
|
|
|
@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
|
|
class TestPrefillReadyMergeBehavior:
|
|
"""Tests for correct merge behavior with _PrefillReadyRotatingKVCache.
|
|
|
|
Verifies that BatchRotatingKVCache.merge() creates zero-length buffers
|
|
(not zero-filled) when merging empty RotatingKVCaches from SSD restore.
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self):
|
|
paged_cache = PagedCacheManager(
|
|
block_size=256,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=4)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
def test_merge_creates_zero_length_buffer(self, prefix_cache):
|
|
"""Test merge creates zero-length buffer, not zero-filled.
|
|
|
|
Without the fix, merge creates a (1, H, 128, D) zero-filled buffer
|
|
with left_padding=0, causing 128 unmasked zero positions in attention.
|
|
With the fix, merge creates a (1, H, 0, D) buffer.
|
|
"""
|
|
from mlx_lm.models.cache import BatchRotatingKVCache
|
|
|
|
meta_state = (0, 128, 500, 64)
|
|
kv_shape_ref = (8, 64)
|
|
cache = prefix_cache._create_empty_rotating_cache(
|
|
meta_state, kvcache_offset=512, kv_shape_ref=kv_shape_ref,
|
|
)
|
|
|
|
# Merge single cache (as happens with single request in batch)
|
|
batch_cache = BatchRotatingKVCache.merge([cache])
|
|
|
|
# With _PrefillReadyRotatingKVCache.size() == 0, merge skips
|
|
# the copy (keys is None after init), preserving zero-data state.
|
|
# In new mlx-lm, keys may be None when size() is 0.
|
|
if batch_cache.keys is not None:
|
|
assert batch_cache.keys.shape[2] == 0
|
|
assert batch_cache.values.shape[2] == 0
|
|
# After merge, the batch cache represents an empty state.
|
|
# The original per-request offset (512) is consumed by
|
|
# the merge → left_padding → offset arithmetic.
|
|
|
|
def test_merge_old_behavior_would_create_zero_filled(self, prefix_cache):
|
|
"""Demonstrate what standard RotatingKVCache.size() reports.
|
|
|
|
With standard RotatingKVCache.size() = min(512, 128) = 128,
|
|
merge would try to copy 128 elements from a 0-length buffer.
|
|
New mlx-lm (4469ad4+) slices with [..., -l:, :] which fails on
|
|
empty keys, so merge now raises ValueError for this case.
|
|
This confirms the _PrefillReadyRotatingKVCache fix is still needed.
|
|
"""
|
|
from mlx_lm.models.cache import RotatingKVCache, BatchRotatingKVCache
|
|
|
|
# Create using standard RotatingKVCache (NOT our subclass)
|
|
old_cache = RotatingKVCache(max_size=128, keep=0)
|
|
old_cache.offset = 512
|
|
old_cache.keys = mx.zeros((1, 8, 0, 64))
|
|
old_cache.values = mx.zeros((1, 8, 0, 64))
|
|
old_cache._idx = 0
|
|
|
|
# Standard size() reports 128 (the bug)
|
|
assert old_cache.size() == 128
|
|
|
|
# New mlx-lm merge raises ValueError due to shape mismatch
|
|
# (tries to broadcast (1,8,0,64) slice onto (1,8,128,64) slot)
|
|
with pytest.raises((ValueError, IndexError)):
|
|
BatchRotatingKVCache.merge([old_cache])
|
|
|
|
|
|
@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
|
|
class TestFindKVShapeRef:
|
|
"""Tests for _find_kv_shape_ref helper."""
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self):
|
|
paged_cache = PagedCacheManager(
|
|
block_size=256, max_blocks=100,
|
|
model_name="test-model", initial_blocks=100,
|
|
)
|
|
return BlockAwarePrefixCache(
|
|
model=MockModel(num_layers=4),
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
def test_finds_shape_from_kvcache_layer(self, prefix_cache):
|
|
"""Find kv_heads and head_dim from a KVCache layer."""
|
|
block_data = [
|
|
[(mx.zeros((1, 8, 256, 64)), mx.zeros((1, 8, 256, 64))), # KVCache
|
|
(mx.zeros((1,)), mx.zeros((1,)))], # RotatingKVCache placeholder
|
|
]
|
|
result = prefix_cache._find_kv_shape_ref(
|
|
block_data, ['KVCache', 'RotatingKVCache']
|
|
)
|
|
assert result == (8, 64)
|
|
|
|
def test_skips_non_kvcache_layers(self, prefix_cache):
|
|
"""Skip RotatingKVCache layers when finding shape ref."""
|
|
block_data = [
|
|
[(mx.zeros((1, 4, 256, 128)), mx.zeros((1, 4, 256, 128))), # Rot placeholder
|
|
(mx.zeros((1, 8, 256, 64)), mx.zeros((1, 8, 256, 64)))], # KVCache
|
|
]
|
|
result = prefix_cache._find_kv_shape_ref(
|
|
block_data, ['RotatingKVCache', 'KVCache']
|
|
)
|
|
assert result == (8, 64)
|
|
|
|
def test_returns_none_for_empty_data(self, prefix_cache):
|
|
"""Return None when no block data available."""
|
|
assert prefix_cache._find_kv_shape_ref([], None) is None
|
|
|
|
|
|
@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
|
|
class TestReconstructCachePartialRestore:
|
|
"""Tests for reconstruct_cache with partial restore (placeholder handling)."""
|
|
|
|
@pytest.fixture
|
|
def paged_cache(self):
|
|
"""Create a PagedCacheManager."""
|
|
return PagedCacheManager(
|
|
block_size=256,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def mock_ssd_cache(self):
|
|
"""Create a mock SSD cache manager."""
|
|
mock = MagicMock()
|
|
mock.save_block.return_value = True
|
|
return mock
|
|
|
|
def test_partial_restore_rejects_rotating_placeholder(
|
|
self, paged_cache, mock_ssd_cache
|
|
):
|
|
"""Partial prefix match with Rotating placeholder must be rejected."""
|
|
model = MockModel(num_layers=2)
|
|
prefix_cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd_cache,
|
|
)
|
|
|
|
# Allocate 2 blocks (both non-last, so RotatingKVCache has placeholders)
|
|
block1 = paged_cache.allocate_block()
|
|
block1.block_hash = b"block1_hash"
|
|
block1.token_count = 256
|
|
|
|
block2 = paged_cache.allocate_block()
|
|
block2.block_hash = b"block2_hash"
|
|
block2.token_count = 256
|
|
|
|
block_table = paged_cache.create_block_table("req-001")
|
|
block_table.block_ids = [block1.block_id, block2.block_id]
|
|
block_table.num_tokens = 512
|
|
|
|
# Mock SSD: 2 layers - KVCache (4D slice) + RotatingKVCache (placeholder)
|
|
kv_keys = mx.zeros((1, 8, 256, 64))
|
|
kv_values = mx.zeros((1, 8, 256, 64))
|
|
rot_placeholder = mx.zeros((1,))
|
|
|
|
block_data = [
|
|
(kv_keys, kv_values), # Layer 0: KVCache slice
|
|
(rot_placeholder, rot_placeholder), # Layer 1: RotatingKVCache placeholder
|
|
]
|
|
block_metadata = {
|
|
'layer_cache_types': ['KVCache', 'RotatingKVCache'],
|
|
'layer_meta_states': [
|
|
(256,), # KVCache: (offset,)
|
|
(4, 1024, 500, 100), # RotatingKVCache: (keep, max_size, offset, _idx)
|
|
],
|
|
'model_name': 'test-model',
|
|
'num_layers': 2,
|
|
}
|
|
|
|
mock_ssd_cache.load_block_with_metadata.return_value = (block_data, block_metadata)
|
|
|
|
result = prefix_cache.reconstruct_cache(block_table)
|
|
|
|
assert result is None
|
|
|
|
def test_full_restore_reconstructs_rotating_state(
|
|
self, paged_cache, mock_ssd_cache
|
|
):
|
|
"""Exact block match should restore full RotatingKVCache state."""
|
|
model = MockModel(num_layers=2)
|
|
prefix_cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd_cache,
|
|
)
|
|
|
|
# Allocate 1 block (the last block with full RotatingKVCache data)
|
|
block = paged_cache.allocate_block()
|
|
block.block_hash = b"block_hash"
|
|
block.token_count = 256
|
|
|
|
block_table = paged_cache.create_block_table("req-001")
|
|
block_table.block_ids = [block.block_id]
|
|
block_table.num_tokens = 256
|
|
|
|
# Mock SSD: last block has full RotatingKVCache data
|
|
kv_keys = mx.zeros((1, 8, 256, 64))
|
|
kv_values = mx.zeros((1, 8, 256, 64))
|
|
rot_keys = mx.ones((1, 8, 1024, 64)) # Full RotatingKVCache state
|
|
rot_values = mx.ones((1, 8, 1024, 64))
|
|
|
|
block_data = [
|
|
(kv_keys, kv_values), # Layer 0: KVCache
|
|
(rot_keys, rot_values), # Layer 1: RotatingKVCache full state
|
|
]
|
|
block_metadata = {
|
|
'layer_cache_types': ['KVCache', 'RotatingKVCache'],
|
|
'layer_meta_states': [
|
|
(256,),
|
|
(4, 1024, 500, 100),
|
|
],
|
|
'model_name': 'test-model',
|
|
'num_layers': 2,
|
|
}
|
|
|
|
mock_ssd_cache.load_block_with_metadata.return_value = (block_data, block_metadata)
|
|
|
|
result = prefix_cache.reconstruct_cache(block_table)
|
|
|
|
assert result is not None
|
|
assert len(result) == 2
|
|
|
|
# Layer 0 (KVCache): normal reconstruct
|
|
kv_cache = result[0]
|
|
assert kv_cache.keys.shape[2] == 256
|
|
|
|
# Layer 1 (RotatingKVCache): reconstructed from stored state
|
|
rot_cache = result[1]
|
|
assert rot_cache.max_size == 1024
|
|
assert rot_cache.keep == 4
|
|
assert rot_cache.offset == 500
|
|
assert rot_cache.keys is not None
|
|
assert rot_cache.keys.shape == (1, 8, 1024, 64)
|
|
assert rot_cache.values.shape == (1, 8, 1024, 64)
|
|
# mlx-lm 0.31.3 contract: _idx must equal keys.shape[2] so the
|
|
# buffer is read in temporal order (case 1 of _temporal_order).
|
|
# The legacy meta_state's _idx=100 hint is intentionally ignored.
|
|
assert rot_cache._idx == 1024
|
|
|
|
def test_undersized_rotating_cache_preserves_buffer_length(
|
|
self, paged_cache, mock_ssd_cache
|
|
):
|
|
"""Undersized RotatingKVCache from BatchRotatingKVCache.extract()
|
|
keeps its actual buffer length on reconstruct.
|
|
|
|
Pre-fix omlx zero-padded the buffer to max_size. That broke the
|
|
merge contract by leaking zero positions into attention (#934,
|
|
#903). Post-fix, the buffer length is preserved and
|
|
PrefillReadyRotatingKVCache.size() clamps the merge slice so
|
|
BatchRotatingKVCache.merge() never reads beyond the real data.
|
|
"""
|
|
model = MockModel(num_layers=2)
|
|
prefix_cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd_cache,
|
|
)
|
|
|
|
block = paged_cache.allocate_block()
|
|
block.block_hash = b"block_hash"
|
|
block.token_count = 256
|
|
|
|
block_table = paged_cache.create_block_table("req-undersized")
|
|
block_table.block_ids = [block.block_id]
|
|
block_table.num_tokens = 256
|
|
|
|
# Simulate undersized buffer from extract(): 845 entries, not 1024
|
|
kv_keys = mx.zeros((1, 8, 256, 64))
|
|
kv_values = mx.zeros((1, 8, 256, 64))
|
|
rot_keys = mx.ones((1, 8, 845, 64)) # Undersized!
|
|
rot_values = mx.ones((1, 8, 845, 64))
|
|
|
|
block_data = [
|
|
(kv_keys, kv_values),
|
|
(rot_keys, rot_values),
|
|
]
|
|
block_metadata = {
|
|
'layer_cache_types': ['KVCache', 'RotatingKVCache'],
|
|
'layer_meta_states': [
|
|
(256,),
|
|
(0, 1024, 44225, 845), # keep=0, max_size=1024, offset=44225, _idx=845
|
|
],
|
|
'model_name': 'test-model',
|
|
'num_layers': 2,
|
|
}
|
|
|
|
mock_ssd_cache.load_block_with_metadata.return_value = (block_data, block_metadata)
|
|
|
|
result = prefix_cache.reconstruct_cache(block_table)
|
|
|
|
assert result is not None
|
|
assert len(result) == 2
|
|
|
|
rot_cache = result[1]
|
|
assert rot_cache.max_size == 1024
|
|
# Buffer length is preserved — no zero padding to max_size.
|
|
assert rot_cache.keys.shape[2] == 845
|
|
assert rot_cache.values.shape[2] == 845
|
|
# Offset preserved from meta_state.
|
|
assert rot_cache.offset == 44225
|
|
# _idx matches the buffer so _temporal_order returns as-is.
|
|
assert rot_cache._idx == 845
|
|
# PrefillReadyRotatingKVCache clamps size() to the buffer length,
|
|
# protecting BatchRotatingKVCache.merge() from over-reading.
|
|
assert rot_cache.size() == 845
|
|
|
|
|
|
class TestModelCacheConfigCacheList:
|
|
"""Tests for CacheList support in ModelCacheConfig."""
|
|
|
|
def test_from_cache_list_with_cache_list(self):
|
|
"""Test from_cache_list detects CacheList."""
|
|
from omlx.cache.type_handlers import CacheType
|
|
from omlx.cache.type_registry import CacheTypeRegistry
|
|
|
|
# CacheList mock: has .caches attribute
|
|
mock_cache_list = MagicMock(spec=[])
|
|
mock_sub_kv = MagicMock(spec=[])
|
|
mock_sub_kv.__class__ = type("KVCache", (), {})
|
|
mock_cache_list.caches = (mock_sub_kv,)
|
|
mock_cache_list.__class__ = type("CacheList", (), {})
|
|
|
|
mock_kv = MagicMock(spec=[])
|
|
mock_kv.__class__ = type("KVCache", (), {})
|
|
mock_kv.keys = MagicMock()
|
|
mock_kv.values = MagicMock()
|
|
|
|
config = ModelCacheConfig.from_cache_list(
|
|
[mock_cache_list, mock_kv], model_name="test"
|
|
)
|
|
|
|
assert config.num_layers == 2
|
|
assert config.is_hybrid is True
|
|
assert config.layer_configs[0].cache_type == CacheType.CACHE_LIST
|
|
assert config.layer_configs[1].cache_type == CacheType.KVCACHE
|
|
|
|
def test_from_cache_list_cache_list_with_rotating_sub(self):
|
|
"""Test from_cache_list extracts window_size from CacheList sub-caches."""
|
|
mock_rotating = MagicMock(spec=[])
|
|
mock_rotating.__class__ = type("RotatingKVCache", (), {})
|
|
mock_rotating.max_size = 256
|
|
mock_rotating._idx = 0
|
|
|
|
mock_cache_list = MagicMock(spec=[])
|
|
mock_cache_list.__class__ = type("CacheList", (), {})
|
|
mock_cache_list.caches = (mock_rotating,)
|
|
|
|
config = ModelCacheConfig.from_cache_list([mock_cache_list])
|
|
|
|
assert config.get_max_window_size() == 256
|
|
|
|
def test_get_meta_states_cache_list(self):
|
|
"""Test get_meta_states extracts composite meta_state for CacheList."""
|
|
from omlx.cache.type_handlers import CacheType
|
|
|
|
# Use real thin classes so type(obj).__name__ returns the correct name
|
|
# (MagicMock(spec=[]) with __class__ override does not affect type().__name__)
|
|
KVCacheStub = type("KVCache", (), {
|
|
"state": (MagicMock(), MagicMock()),
|
|
"meta_state": (32,),
|
|
})
|
|
mock_sub = KVCacheStub()
|
|
|
|
CacheListStub = type("CacheList", (), {
|
|
"caches": (mock_sub,),
|
|
})
|
|
mock_cache_list = CacheListStub()
|
|
|
|
config = ModelCacheConfig.from_cache_list([mock_cache_list])
|
|
meta_states = config.get_meta_states([mock_cache_list])
|
|
|
|
assert len(meta_states) == 1
|
|
# CacheList meta_state is (class_names, sub_meta_states)
|
|
assert isinstance(meta_states[0], tuple)
|
|
assert len(meta_states[0]) == 2
|
|
class_names, sub_metas = meta_states[0]
|
|
assert class_names == ["KVCache"]
|