wehub-resource-sync
3003 lines
105 KiB
Python
3003 lines
105 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Tests for BlockAwarePrefixCache and related components.
|
|
|
|
This module tests the block-aware prefix caching system that uses
|
|
PagedCacheManager for block-based storage with SSD persistence.
|
|
"""
|
|
|
|
import sys
|
|
import time
|
|
import types
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from omlx.cache.paged_cache import (
|
|
BlockTable,
|
|
PagedCacheManager,
|
|
compute_block_hash,
|
|
)
|
|
from omlx.cache.prefix_cache import BlockAwarePrefixCache, BlockCacheEntry
|
|
from omlx.cache.stats import PrefixCacheStats
|
|
|
|
|
|
class MockModel:
|
|
"""Mock model for testing."""
|
|
|
|
def __init__(self, num_layers: int = 32):
|
|
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 TestBlockCacheEntry:
|
|
"""Tests for BlockCacheEntry dataclass."""
|
|
|
|
def test_creation(self):
|
|
"""Test creating a cache entry."""
|
|
block_table = BlockTable(request_id="req-001")
|
|
entry = BlockCacheEntry(
|
|
block_table=block_table,
|
|
last_access=time.time(),
|
|
)
|
|
|
|
assert entry.block_table is block_table
|
|
assert entry.last_access > 0
|
|
|
|
|
|
class TestBlockAwarePrefixCache:
|
|
"""Tests for BlockAwarePrefixCache."""
|
|
|
|
@pytest.fixture
|
|
def paged_cache(self):
|
|
"""Create a PagedCacheManager for testing."""
|
|
return PagedCacheManager(
|
|
block_size=4, # Small block size for testing
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def mock_model(self):
|
|
"""Create a mock model."""
|
|
return MockModel(num_layers=4)
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self, mock_model, paged_cache):
|
|
"""Create a BlockAwarePrefixCache for testing."""
|
|
return BlockAwarePrefixCache(
|
|
model=mock_model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
|
|
def test_initialization(self, prefix_cache, mock_model, paged_cache):
|
|
"""Test cache initialization."""
|
|
assert prefix_cache.model is mock_model
|
|
assert prefix_cache.paged_cache is paged_cache
|
|
assert prefix_cache.block_size == 4
|
|
assert prefix_cache.expected_num_layers == 4
|
|
|
|
def test_get_model_num_layers(self, paged_cache):
|
|
"""Test getting model num_layers from various sources."""
|
|
# Model with layers attribute
|
|
model1 = MockModel(num_layers=16)
|
|
cache1 = BlockAwarePrefixCache(model1, paged_cache)
|
|
assert cache1.expected_num_layers == 16
|
|
|
|
# Model with args.num_hidden_layers
|
|
model2 = MagicMock()
|
|
model2.layers = None
|
|
model2.args.num_hidden_layers = 24
|
|
delattr(model2, "layers")
|
|
cache2 = BlockAwarePrefixCache(model2, paged_cache)
|
|
assert cache2.expected_num_layers == 24
|
|
|
|
def test_get_model_num_layers_with_make_cache(self, paged_cache):
|
|
"""Test make_cache() takes priority over model.layers for hybrid models."""
|
|
# Hybrid model: 88 architectural layers but only 48 produce cache
|
|
model = MockModel(num_layers=88)
|
|
model.make_cache = lambda: [MagicMock() for _ in range(48)]
|
|
cache = BlockAwarePrefixCache(model, paged_cache)
|
|
assert cache.expected_num_layers == 48
|
|
|
|
# Non-hybrid model: make_cache() returns same count as model.layers
|
|
model2 = MockModel(num_layers=32)
|
|
model2.make_cache = lambda: [MagicMock() for _ in range(32)]
|
|
cache2 = BlockAwarePrefixCache(model2, paged_cache)
|
|
assert cache2.expected_num_layers == 32
|
|
|
|
def test_get_model_num_layers_make_cache_exception(self, paged_cache):
|
|
"""Test fallback to model.layers when make_cache() raises."""
|
|
model = MockModel(num_layers=88)
|
|
|
|
def bad_make_cache():
|
|
raise RuntimeError("model not initialized")
|
|
|
|
model.make_cache = bad_make_cache
|
|
cache = BlockAwarePrefixCache(model, paged_cache)
|
|
assert cache.expected_num_layers == 88
|
|
|
|
def test_get_model_num_layers_make_cache_empty(self, paged_cache):
|
|
"""Test fallback when make_cache() returns empty list."""
|
|
model = MockModel(num_layers=32)
|
|
model.make_cache = lambda: []
|
|
cache = BlockAwarePrefixCache(model, paged_cache)
|
|
assert cache.expected_num_layers == 32
|
|
|
|
def test_get_model_num_layers_make_cache_non_list(self, paged_cache):
|
|
"""Test fallback when make_cache() returns non-list."""
|
|
model = MockModel(num_layers=32)
|
|
model.make_cache = lambda: None
|
|
cache = BlockAwarePrefixCache(model, paged_cache)
|
|
assert cache.expected_num_layers == 32
|
|
|
|
def test_fetch_cache_miss(self, prefix_cache):
|
|
"""Test fetch_cache returns miss when no cache exists."""
|
|
tokens = [1, 2, 3, 4, 5, 6, 7, 8]
|
|
|
|
block_table, remaining = prefix_cache.fetch_cache("req-001", tokens)
|
|
|
|
assert block_table is None
|
|
assert remaining == tokens
|
|
assert prefix_cache._misses == 1
|
|
|
|
def test_fetch_cache_empty_tokens(self, prefix_cache):
|
|
"""Test fetch_cache with empty tokens."""
|
|
block_table, remaining = prefix_cache.fetch_cache("req-001", [])
|
|
|
|
assert block_table is None
|
|
assert remaining == []
|
|
|
|
def test_store_cache_empty_tokens(self, prefix_cache):
|
|
"""Test store_cache with empty tokens returns None."""
|
|
result = prefix_cache.store_cache("req-001", [], [])
|
|
assert result is None
|
|
|
|
def test_store_cache_creates_block_table(self, prefix_cache):
|
|
"""Test store_cache creates a block table."""
|
|
tokens = [1, 2, 3, 4] # One full block
|
|
# Mock cache data (tensor data format)
|
|
cache_data = [
|
|
{"state": (MagicMock(), MagicMock()), "cache_type": "KVCache"}
|
|
for _ in range(4)
|
|
]
|
|
|
|
result = prefix_cache.store_cache("req-001", tokens, cache_data)
|
|
|
|
assert result is not None
|
|
assert result.request_id == "req-001"
|
|
assert "req-001" in prefix_cache._request_tables
|
|
|
|
def test_release_cache(self, prefix_cache, paged_cache):
|
|
"""Test releasing cache for a request."""
|
|
# First create a block table with blocks
|
|
table = paged_cache.create_block_table("req-001")
|
|
block = paged_cache.allocate_block()
|
|
table.block_ids.append(block.block_id)
|
|
|
|
# Add to prefix cache tracking
|
|
prefix_cache._request_tables["req-001"] = BlockCacheEntry(
|
|
block_table=table,
|
|
last_access=time.time(),
|
|
)
|
|
|
|
initial_free = paged_cache.free_blocks
|
|
|
|
prefix_cache.release_cache("req-001")
|
|
|
|
assert "req-001" not in prefix_cache._request_tables
|
|
assert paged_cache.free_blocks == initial_free + 1
|
|
|
|
def test_clear_request_entry(self, prefix_cache, paged_cache):
|
|
"""Test clearing request entry without freeing blocks."""
|
|
table = paged_cache.create_block_table("req-001")
|
|
prefix_cache._request_tables["req-001"] = BlockCacheEntry(
|
|
block_table=table,
|
|
last_access=time.time(),
|
|
)
|
|
|
|
prefix_cache.clear_request_entry("req-001")
|
|
|
|
assert "req-001" not in prefix_cache._request_tables
|
|
# Blocks should still be tracked in paged_cache
|
|
assert "req-001" in paged_cache.request_tables
|
|
|
|
def test_fork_cache(self, prefix_cache, paged_cache):
|
|
"""Test forking cache from one request to another."""
|
|
# Create source with blocks
|
|
source_table = paged_cache.create_block_table("req-source")
|
|
block = paged_cache.allocate_block()
|
|
source_table.block_ids.append(block.block_id)
|
|
|
|
prefix_cache._request_tables["req-source"] = BlockCacheEntry(
|
|
block_table=source_table,
|
|
last_access=time.time(),
|
|
)
|
|
|
|
forked = prefix_cache.fork_cache("req-source", "req-forked")
|
|
|
|
assert forked is not None
|
|
assert forked.request_id == "req-forked"
|
|
assert "req-forked" in prefix_cache._request_tables
|
|
# Ref count should be incremented
|
|
assert block.ref_count == 2
|
|
|
|
def test_fork_cache_nonexistent(self, prefix_cache):
|
|
"""Test forking nonexistent cache returns None."""
|
|
result = prefix_cache.fork_cache("nonexistent", "new-req")
|
|
assert result is None
|
|
|
|
def test_get_stats(self, prefix_cache):
|
|
"""Test getting statistics."""
|
|
prefix_cache._hits = 5
|
|
prefix_cache._misses = 3
|
|
prefix_cache._tokens_saved = 100
|
|
|
|
stats = prefix_cache.get_stats()
|
|
|
|
assert isinstance(stats, PrefixCacheStats)
|
|
assert stats.hits == 5
|
|
assert stats.misses == 3
|
|
assert stats.tokens_saved == 100
|
|
assert stats.partial_block_skips == 0
|
|
assert stats.partial_tokens_skipped == 0
|
|
assert stats.block_size == prefix_cache.block_size
|
|
|
|
def test_get_stats_dict(self, prefix_cache):
|
|
"""Test getting statistics as dictionary."""
|
|
prefix_cache._hits = 10
|
|
prefix_cache._misses = 5
|
|
|
|
stats_dict = prefix_cache.get_stats_dict()
|
|
|
|
assert "hits" in stats_dict
|
|
assert "misses" in stats_dict
|
|
assert "hit_rate" in stats_dict
|
|
assert "partial_block_skips" in stats_dict
|
|
assert "partial_tokens_skipped" in stats_dict
|
|
assert "last_tokens_to_next_block" in stats_dict
|
|
assert stats_dict["hit_rate"] == pytest.approx(10 / 15)
|
|
|
|
def test_reset_stats(self, prefix_cache):
|
|
"""Test resetting statistics."""
|
|
prefix_cache._hits = 10
|
|
prefix_cache._misses = 5
|
|
prefix_cache._tokens_saved = 500
|
|
prefix_cache._partial_block_skips = 3
|
|
prefix_cache._partial_tokens_skipped = 42
|
|
prefix_cache._last_partial_tokens_skipped = 2
|
|
prefix_cache._last_tokens_to_next_block = 254
|
|
|
|
prefix_cache.reset_stats()
|
|
|
|
assert prefix_cache._hits == 0
|
|
assert prefix_cache._misses == 0
|
|
assert prefix_cache._tokens_saved == 0
|
|
assert prefix_cache._partial_block_skips == 0
|
|
assert prefix_cache._partial_tokens_skipped == 0
|
|
assert prefix_cache._last_partial_tokens_skipped == 0
|
|
assert prefix_cache._last_tokens_to_next_block == 0
|
|
|
|
def test_clear(self, prefix_cache, paged_cache):
|
|
"""Test clearing all cache data."""
|
|
# Add some data
|
|
table = paged_cache.create_block_table("req-001")
|
|
prefix_cache._request_tables["req-001"] = BlockCacheEntry(
|
|
block_table=table,
|
|
last_access=time.time(),
|
|
)
|
|
prefix_cache._prefix_index[b"test_hash"] = (10, [1, 2], 2)
|
|
prefix_cache._hits = 10
|
|
|
|
cleared = prefix_cache.clear()
|
|
|
|
assert cleared > 0
|
|
assert len(prefix_cache._request_tables) == 0
|
|
assert len(prefix_cache._prefix_index) == 0
|
|
assert prefix_cache._hits == 0
|
|
|
|
def test_len(self, prefix_cache):
|
|
"""Test __len__ returns number of request entries."""
|
|
assert len(prefix_cache) == 0
|
|
|
|
table = BlockTable(request_id="req-001")
|
|
prefix_cache._request_tables["req-001"] = BlockCacheEntry(
|
|
block_table=table,
|
|
last_access=time.time(),
|
|
)
|
|
prefix_cache._request_tables["req-002"] = BlockCacheEntry(
|
|
block_table=BlockTable(request_id="req-002"),
|
|
last_access=time.time(),
|
|
)
|
|
|
|
assert len(prefix_cache) == 2
|
|
|
|
def test_cache_manager_interface_fetch(self, prefix_cache):
|
|
"""Test CacheManager ABC fetch interface."""
|
|
# Invalid key format
|
|
value, hit = prefix_cache.fetch("invalid")
|
|
assert hit is False
|
|
assert value is None
|
|
|
|
# Valid key format but miss
|
|
value, hit = prefix_cache.fetch(("req-001", [1, 2, 3, 4]))
|
|
assert hit is False
|
|
|
|
def test_cache_manager_interface_store(self, prefix_cache):
|
|
"""Test CacheManager ABC store interface."""
|
|
# Invalid key format
|
|
result = prefix_cache.store("invalid", [])
|
|
assert result is False
|
|
|
|
def test_cache_manager_interface_evict(self, prefix_cache, paged_cache):
|
|
"""Test CacheManager ABC evict interface."""
|
|
# Evict nonexistent
|
|
result = prefix_cache.evict("req-nonexistent")
|
|
assert result is False
|
|
|
|
# Create entry then evict
|
|
table = paged_cache.create_block_table("req-001")
|
|
prefix_cache._request_tables["req-001"] = BlockCacheEntry(
|
|
block_table=table,
|
|
last_access=time.time(),
|
|
)
|
|
|
|
result = prefix_cache.evict("req-001")
|
|
assert result is True
|
|
assert "req-001" not in prefix_cache._request_tables
|
|
|
|
def test_cache_manager_interface_properties(self, prefix_cache):
|
|
"""Test CacheManager ABC property interface."""
|
|
assert prefix_cache.size == 0
|
|
assert prefix_cache.max_size == 100 # max_blocks
|
|
|
|
def test_set_paged_ssd_cache_manager(self, prefix_cache):
|
|
"""Test setting SSD cache manager."""
|
|
mock_ssd_cache = MagicMock()
|
|
|
|
prefix_cache.set_paged_ssd_cache_manager(mock_ssd_cache)
|
|
|
|
assert prefix_cache.paged_ssd_cache is mock_ssd_cache
|
|
|
|
def test_set_cold_restore_callback(self, prefix_cache):
|
|
"""Test setting cold restore callback."""
|
|
|
|
def restore_callback(block_id: int, block_hash: bytes) -> bool:
|
|
return True
|
|
|
|
prefix_cache.set_cold_restore_callback(restore_callback)
|
|
assert prefix_cache._cold_restore_callback is restore_callback
|
|
|
|
|
|
class TestBlockAwarePrefixCacheWithSSD:
|
|
"""Tests for BlockAwarePrefixCache with SSD cache manager."""
|
|
|
|
@pytest.fixture
|
|
def paged_cache(self):
|
|
"""Create a PagedCacheManager."""
|
|
return PagedCacheManager(
|
|
block_size=4,
|
|
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)
|
|
mock.has_block.return_value = False
|
|
return mock
|
|
|
|
@pytest.fixture
|
|
def prefix_cache_with_ssd(self, paged_cache, mock_ssd_cache):
|
|
"""Create a BlockAwarePrefixCache with SSD manager."""
|
|
model = MockModel(num_layers=4)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd_cache,
|
|
)
|
|
|
|
def test_get_cache_for_generation_no_entry(self, prefix_cache_with_ssd):
|
|
"""Test get_cache_for_generation with no entry."""
|
|
result, was_loaded = prefix_cache_with_ssd.get_cache_for_generation("req-001")
|
|
|
|
assert result is None
|
|
assert was_loaded is False
|
|
|
|
def test_reconstruct_cache_empty_table(self, prefix_cache_with_ssd):
|
|
"""Test reconstruct_cache with empty block table."""
|
|
table = BlockTable(request_id="req-001")
|
|
|
|
result = prefix_cache_with_ssd.reconstruct_cache(table)
|
|
|
|
assert result is None
|
|
|
|
def test_reconstruct_cache_no_ssd_manager(self, paged_cache):
|
|
"""Test reconstruct_cache without SSD manager."""
|
|
model = MockModel(num_layers=4)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
|
|
table = BlockTable(request_id="req-001", block_ids=[1, 2])
|
|
|
|
result = cache.reconstruct_cache(table)
|
|
|
|
assert result is None
|
|
|
|
def test_forget_incompatible_ssd_block_calls_forget(self, prefix_cache_with_ssd):
|
|
"""The helper clears local SSD indexes without deleting the file."""
|
|
block_hash = b"\xab" * 32
|
|
prefix_cache_with_ssd._forget_incompatible_ssd_block(block_hash)
|
|
|
|
prefix_cache_with_ssd.paged_ssd_cache.forget_block.assert_called_once_with(
|
|
block_hash
|
|
)
|
|
|
|
def test_forget_incompatible_ssd_block_noop_when_hash_missing(
|
|
self, prefix_cache_with_ssd
|
|
):
|
|
"""No hash means the block was never persisted to SSD — nothing to
|
|
forget."""
|
|
prefix_cache_with_ssd._forget_incompatible_ssd_block(None)
|
|
prefix_cache_with_ssd.paged_ssd_cache.forget_block.assert_not_called()
|
|
|
|
def test_forget_incompatible_ssd_block_noop_without_ssd(self, paged_cache):
|
|
"""Pure paged-cache deployments have no SSD manager attached — the
|
|
helper must not raise."""
|
|
model = MockModel(num_layers=4)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
# Should not raise.
|
|
cache._forget_incompatible_ssd_block(b"\xab" * 32)
|
|
|
|
def test_reconstruct_block_size_mismatch_forgets_without_delete(
|
|
self, prefix_cache_with_ssd
|
|
):
|
|
"""A block from another block_size is rejected without unlinking SSD."""
|
|
block = prefix_cache_with_ssd.paged_cache.allocate_block()
|
|
block.block_hash = b"\xcd" * 32
|
|
block.token_count = 4
|
|
table = BlockTable(
|
|
request_id="req-001",
|
|
block_ids=[block.block_id],
|
|
num_tokens=4,
|
|
)
|
|
prefix_cache_with_ssd.paged_ssd_cache.load_block_with_metadata.return_value = (
|
|
[(MagicMock(), MagicMock()) for _ in range(4)],
|
|
{
|
|
"model_name": "test-model",
|
|
"num_layers": 4,
|
|
"block_size": 2048,
|
|
"layer_cache_types": ["KVCache"] * 4,
|
|
},
|
|
)
|
|
|
|
result = prefix_cache_with_ssd.reconstruct_cache(table)
|
|
|
|
assert result is None
|
|
prefix_cache_with_ssd.paged_ssd_cache.forget_block.assert_called_once_with(
|
|
block.block_hash
|
|
)
|
|
prefix_cache_with_ssd.paged_ssd_cache.delete_block.assert_not_called()
|
|
|
|
def test_minimax_m3_sliceable_nstate_blocks_round_trip(self, monkeypatch):
|
|
"""MiniMax M3 single-cache blocks store K/V/index slices, not snapshots."""
|
|
mx = pytest.importorskip("mlx.core")
|
|
|
|
module = types.ModuleType("mlx_vlm.models.minimax_m3_vl.language")
|
|
|
|
class FakeInnerKVCache:
|
|
def __init__(self):
|
|
self.state = None
|
|
|
|
class MiniMaxM3KVCache:
|
|
def __init__(self):
|
|
self.kv_cache = FakeInnerKVCache()
|
|
self.index_keys = None
|
|
self.index_offset = 0
|
|
|
|
module.MiniMaxM3KVCache = MiniMaxM3KVCache
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"mlx_vlm.models.minimax_m3_vl.language",
|
|
module,
|
|
)
|
|
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
saved_by_hash = {}
|
|
|
|
def save_block(**kwargs):
|
|
saved_by_hash[kwargs["block_hash"]] = (
|
|
kwargs["cache_data"],
|
|
{
|
|
"model_name": "test-model",
|
|
"num_layers": 1,
|
|
"block_size": 4,
|
|
"layer_cache_types": kwargs["layer_cache_types"],
|
|
"layer_meta_states": kwargs["layer_meta_states"],
|
|
},
|
|
)
|
|
return True
|
|
|
|
mock_ssd.save_block.side_effect = save_block
|
|
cache = BlockAwarePrefixCache(
|
|
model=MockModel(num_layers=1),
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
keys = mx.arange(1 * 2 * 8 * 3, dtype=mx.float32).reshape(1, 2, 8, 3)
|
|
values = keys + 100
|
|
index_keys = mx.arange(1 * 1 * 8 * 3, dtype=mx.float32).reshape(1, 1, 8, 3)
|
|
mx.eval(keys, values, index_keys)
|
|
cache_data = [
|
|
{
|
|
"state": (keys, values, index_keys),
|
|
"cache_type": "MiniMaxM3KVCache",
|
|
"class_name": "MiniMaxM3KVCache",
|
|
"meta_state": (999,),
|
|
}
|
|
]
|
|
|
|
block_table = cache.store_cache("req-minimax", list(range(8)), cache_data)
|
|
|
|
assert block_table is not None
|
|
assert len(block_table.block_ids) == 2
|
|
assert mock_ssd.save_block.call_count == 2
|
|
for idx, call in enumerate(mock_ssd.save_block.call_args_list):
|
|
marker = call.kwargs["cache_data"][0]
|
|
assert marker[0] == "__nstate__"
|
|
assert marker[1] == "MiniMaxM3KVCache"
|
|
saved_keys, saved_values, saved_index = marker[2]
|
|
start = idx * 4
|
|
end = start + 4
|
|
assert saved_keys.tolist() == keys[:, :, start:end, :].tolist()
|
|
assert saved_values.tolist() == values[:, :, start:end, :].tolist()
|
|
assert saved_index.tolist() == index_keys[:, :, start:end, :].tolist()
|
|
|
|
def load_block_with_metadata(block_hash, **kwargs):
|
|
return saved_by_hash.get(block_hash, (None, None))
|
|
|
|
mock_ssd.load_block_with_metadata.side_effect = load_block_with_metadata
|
|
restored = cache.reconstruct_cache(block_table)
|
|
|
|
assert restored is not None
|
|
restored_cache = restored[0]
|
|
restored_keys, restored_values = restored_cache.kv_cache.state
|
|
assert restored_keys.tolist() == keys.tolist()
|
|
assert restored_values.tolist() == values.tolist()
|
|
assert restored_cache.index_keys.tolist() == index_keys.tolist()
|
|
assert restored_cache.index_offset == 8
|
|
|
|
|
|
class TestPrefixIndexOperations:
|
|
"""Tests for prefix index operations."""
|
|
|
|
@pytest.fixture
|
|
def paged_cache(self):
|
|
"""Create a PagedCacheManager."""
|
|
return PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self, paged_cache):
|
|
"""Create a BlockAwarePrefixCache."""
|
|
model = MockModel(num_layers=4)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
def test_update_prefix_index(self, prefix_cache, paged_cache):
|
|
"""Test _update_prefix_index method."""
|
|
tokens = [1, 2, 3, 4, 5, 6, 7, 8] # 2 blocks
|
|
block_ids = [1, 2]
|
|
|
|
# Allocate blocks and set them up
|
|
blocks = paged_cache.get_new_blocks(2)
|
|
block_ids = [b.block_id for b in blocks]
|
|
|
|
prefix_cache._update_prefix_index(tokens, block_ids)
|
|
|
|
# Prefix index should have entries
|
|
assert len(prefix_cache._prefix_index) > 0
|
|
|
|
def test_find_best_prefix_match_no_match(self, prefix_cache):
|
|
"""Test _find_best_prefix_match with no matching prefix."""
|
|
tokens = [1, 2, 3, 4]
|
|
|
|
result = prefix_cache._find_best_prefix_match(tokens)
|
|
|
|
assert result is None
|
|
|
|
def test_find_best_prefix_match_with_match(self, prefix_cache, paged_cache):
|
|
"""Test _find_best_prefix_match finding a matching prefix."""
|
|
tokens = [1, 2, 3, 4]
|
|
block_ids = [1, 2]
|
|
|
|
# Manually add to prefix index
|
|
block_hash = compute_block_hash(b"", tokens, model_name=paged_cache.model_name)
|
|
prefix_cache._prefix_index[block_hash] = (4, block_ids, 1)
|
|
|
|
result = prefix_cache._find_best_prefix_match(tokens)
|
|
|
|
assert result is not None
|
|
prefix_len, matched_ids, num_blocks = result
|
|
assert prefix_len == 4
|
|
|
|
def test_prefix_index_immutable_after_store(self, prefix_cache, paged_cache):
|
|
"""Test that _prefix_index entries are not affected by later mutations
|
|
of the original block_ids list (e.g., from CoW or block reallocation).
|
|
|
|
Regression test for: storing a mutable list reference in _prefix_index
|
|
allows CoW operations to silently corrupt the index.
|
|
"""
|
|
tokens = [1, 2, 3, 4, 5, 6, 7, 8] # 2 blocks (block_size=4)
|
|
|
|
# Allocate blocks
|
|
blocks = paged_cache.get_new_blocks(2)
|
|
block_ids = [b.block_id for b in blocks]
|
|
original_ids = list(block_ids) # snapshot for assertion
|
|
|
|
# Store into prefix index
|
|
prefix_cache._update_prefix_index(tokens, block_ids)
|
|
|
|
# Simulate CoW: mutate the original list in-place
|
|
block_ids[0] = 9999
|
|
|
|
# Verify: prefix_index must still contain the original block IDs
|
|
result = prefix_cache._find_best_prefix_match(tokens)
|
|
assert result is not None
|
|
_, matched_ids, num_blocks = result
|
|
assert list(matched_ids[:num_blocks]) == original_ids[:num_blocks]
|
|
|
|
|
|
class TestValidateBlockCacheData:
|
|
"""Tests for _validate_block_cache_data method."""
|
|
|
|
@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_validate_empty_data(self, prefix_cache):
|
|
"""Test validation of empty cache data."""
|
|
assert prefix_cache._validate_block_cache_data([]) is False
|
|
|
|
def test_validate_with_none_values(self, prefix_cache):
|
|
"""Test validation fails with None keys/values."""
|
|
cache_data = [(None, MagicMock())]
|
|
assert prefix_cache._validate_block_cache_data(cache_data) is False
|
|
|
|
cache_data = [(MagicMock(), None)]
|
|
assert prefix_cache._validate_block_cache_data(cache_data) is False
|
|
|
|
def test_validate_with_arrays_cache_skipped(self, prefix_cache):
|
|
"""Test validation skips ArraysCache layers for seq_len check."""
|
|
# Create mock tensors with shape
|
|
mock_keys1 = MagicMock()
|
|
mock_keys1.shape = (1, 8, 64, 64) # KVCache shape
|
|
mock_values1 = MagicMock()
|
|
|
|
mock_keys2 = MagicMock()
|
|
mock_keys2.shape = (1, 128) # ArraysCache has different shape
|
|
mock_values2 = MagicMock()
|
|
|
|
cache_data = [
|
|
(mock_keys1, mock_values1),
|
|
(mock_keys2, mock_values2),
|
|
]
|
|
layer_cache_types = ["KVCache", "ArraysCache"]
|
|
|
|
# Should pass because ArraysCache is skipped in validation
|
|
result = prefix_cache._validate_block_cache_data(cache_data, layer_cache_types)
|
|
assert result is True
|
|
|
|
def test_validate_seq_len_mismatch(self, prefix_cache):
|
|
"""Test validation fails on seq_len mismatch."""
|
|
mock_keys1 = MagicMock()
|
|
mock_keys1.shape = (1, 8, 64, 64)
|
|
mock_values1 = MagicMock()
|
|
|
|
mock_keys2 = MagicMock()
|
|
mock_keys2.shape = (1, 8, 32, 64) # Different seq_len
|
|
mock_values2 = MagicMock()
|
|
|
|
cache_data = [
|
|
(mock_keys1, mock_values1),
|
|
(mock_keys2, mock_values2),
|
|
]
|
|
layer_cache_types = ["KVCache", "KVCache"]
|
|
|
|
result = prefix_cache._validate_block_cache_data(cache_data, layer_cache_types)
|
|
assert result is False
|
|
|
|
|
|
class TestArraysCacheLastBlockOnly:
|
|
"""Tests for ArraysCache last-block-only storage and partial match rejection."""
|
|
|
|
@pytest.fixture
|
|
def mx(self):
|
|
"""Import MLX or skip."""
|
|
try:
|
|
import mlx.core as mx
|
|
|
|
return mx
|
|
except ImportError:
|
|
pytest.skip("MLX not available")
|
|
|
|
@pytest.fixture
|
|
def paged_cache(self):
|
|
"""Create a PagedCacheManager."""
|
|
return PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self, paged_cache):
|
|
"""Create a BlockAwarePrefixCache."""
|
|
model = MockModel(num_layers=2)
|
|
return BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
def test_extract_block_arrays_cache_last_block_stores_full_state(
|
|
self, prefix_cache, mx
|
|
):
|
|
"""Last block should store the full ArraysCache state."""
|
|
conv_state = mx.ones((1, 3, 64))
|
|
ssm_state = mx.ones((1, 32, 128, 128))
|
|
cache_data = [
|
|
{
|
|
"state": (conv_state, ssm_state),
|
|
"cache_type": "ArraysCache",
|
|
"class_name": "ArraysCache",
|
|
},
|
|
]
|
|
|
|
result = prefix_cache._extract_block_tensor_slice(
|
|
cache_data,
|
|
0,
|
|
4,
|
|
model_cache_config=None,
|
|
is_last_block=True,
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result) == 1
|
|
keys, values = result[0]
|
|
# Should be full state, not placeholder
|
|
assert keys.shape == (1, 3, 64)
|
|
assert values.shape == (1, 32, 128, 128)
|
|
|
|
def test_extract_block_arrays_cache_non_last_block_stores_placeholder(
|
|
self, prefix_cache, mx
|
|
):
|
|
"""Non-last block should store placeholder (1,) for ArraysCache."""
|
|
conv_state = mx.ones((1, 3, 64))
|
|
ssm_state = mx.ones((1, 32, 128, 128))
|
|
cache_data = [
|
|
{
|
|
"state": (conv_state, ssm_state),
|
|
"cache_type": "ArraysCache",
|
|
"class_name": "ArraysCache",
|
|
},
|
|
]
|
|
|
|
result = prefix_cache._extract_block_tensor_slice(
|
|
cache_data,
|
|
0,
|
|
4,
|
|
model_cache_config=None,
|
|
is_last_block=False,
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result) == 1
|
|
keys, values = result[0]
|
|
# Should be placeholder
|
|
assert keys.shape == (1,)
|
|
assert values.shape == (1,)
|
|
|
|
def test_extract_block_hybrid_model_arrays_cache_and_kvcache(
|
|
self, prefix_cache, mx
|
|
):
|
|
"""Hybrid model: KVCache sliced normally, ArraysCache last-block-only."""
|
|
from omlx.cache.hybrid_cache import ModelCacheConfig
|
|
|
|
kv_keys = mx.ones((1, 8, 8, 64))
|
|
kv_values = mx.ones((1, 8, 8, 64))
|
|
conv_state = mx.ones((1, 3, 64))
|
|
ssm_state = mx.ones((1, 32, 128, 128))
|
|
|
|
cache_data = [
|
|
{
|
|
"state": (kv_keys, kv_values),
|
|
"cache_type": "KVCache",
|
|
"class_name": "KVCache",
|
|
},
|
|
{
|
|
"state": (conv_state, ssm_state),
|
|
"cache_type": "ArraysCache",
|
|
"class_name": "ArraysCache",
|
|
},
|
|
]
|
|
|
|
config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "ArraysCache"], model_name="test"
|
|
)
|
|
|
|
# Non-last block
|
|
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 should be sliced normally
|
|
assert result[0][0].shape[2] == 4 # seq_len slice
|
|
# ArraysCache layer should be placeholder
|
|
assert result[1][0].shape == (1,)
|
|
|
|
# Last block
|
|
result = prefix_cache._extract_block_tensor_slice(
|
|
cache_data,
|
|
4,
|
|
8,
|
|
model_cache_config=config,
|
|
is_last_block=True,
|
|
)
|
|
assert result is not None
|
|
assert len(result) == 2
|
|
# KVCache layer should be sliced
|
|
assert result[0][0].shape[2] == 4
|
|
# ArraysCache layer should have full state
|
|
assert result[1][0].shape == (1, 3, 64)
|
|
|
|
def test_reconstruct_arrays_cache_partial_match_returns_none(
|
|
self, prefix_cache, mx
|
|
):
|
|
"""Partial match (placeholder in last block) should return None."""
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
# Create mock SSD cache
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
|
|
model = MockModel(num_layers=2)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
# Allocate blocks
|
|
block1 = paged_cache.allocate_block()
|
|
block1.block_hash = b"hash1"
|
|
block1.token_count = 4
|
|
block2 = paged_cache.allocate_block()
|
|
block2.block_hash = b"hash2"
|
|
block2.token_count = 4
|
|
|
|
block_table = BlockTable(
|
|
request_id="req-001",
|
|
block_ids=[block1.block_id, block2.block_id],
|
|
num_tokens=8,
|
|
)
|
|
|
|
# Block data: KVCache layer sliced, ArraysCache layer placeholder
|
|
kv_slice = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
|
|
block1_data = [kv_slice, placeholder] # ArraysCache = placeholder (non-last)
|
|
block2_data = [
|
|
kv_slice,
|
|
placeholder,
|
|
] # ArraysCache = placeholder (still non-last in original)
|
|
|
|
mock_ssd.load_block_with_metadata.side_effect = [
|
|
(
|
|
block1_data,
|
|
{
|
|
"model_name": "test-model",
|
|
"num_layers": 2,
|
|
"layer_cache_types": ["KVCache", "ArraysCache"],
|
|
"layer_meta_states": [(), ()],
|
|
},
|
|
),
|
|
(
|
|
block2_data,
|
|
{
|
|
"model_name": "test-model",
|
|
"num_layers": 2,
|
|
"layer_cache_types": ["KVCache", "ArraysCache"],
|
|
"layer_meta_states": [(), ()],
|
|
},
|
|
),
|
|
]
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
|
|
# Should return None because ArraysCache layer has placeholder
|
|
assert result is None
|
|
|
|
def test_reconstruct_arrays_cache_exact_match_succeeds(self, prefix_cache, mx):
|
|
"""Exact match (full state in last block) should reconstruct successfully."""
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
|
|
model = MockModel(num_layers=1)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
block1 = paged_cache.allocate_block()
|
|
block1.block_hash = b"hash1"
|
|
block1.token_count = 4
|
|
block2 = paged_cache.allocate_block()
|
|
block2.block_hash = b"hash2"
|
|
block2.token_count = 4
|
|
|
|
block_table = BlockTable(
|
|
request_id="req-001",
|
|
block_ids=[block1.block_id, block2.block_id],
|
|
num_tokens=8,
|
|
)
|
|
|
|
# First block: placeholder, Second block (last): full state
|
|
conv_state = mx.ones((1, 3, 64))
|
|
ssm_state = mx.ones((1, 32, 128, 128))
|
|
|
|
block1_data = [(mx.zeros((1,)), mx.zeros((1,)))] # placeholder
|
|
block2_data = [(conv_state, ssm_state)] # full state
|
|
|
|
mock_ssd.load_block_with_metadata.side_effect = [
|
|
(
|
|
block1_data,
|
|
{
|
|
"model_name": "test-model",
|
|
"num_layers": 1,
|
|
"layer_cache_types": ["ArraysCache"],
|
|
"layer_meta_states": [()],
|
|
},
|
|
),
|
|
(
|
|
block2_data,
|
|
{
|
|
"model_name": "test-model",
|
|
"num_layers": 1,
|
|
"layer_cache_types": ["ArraysCache"],
|
|
"layer_meta_states": [()],
|
|
},
|
|
),
|
|
]
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
|
|
# Should succeed because last block has full state
|
|
assert result is not None
|
|
assert len(result) == 1
|
|
|
|
def test_store_cache_skips_partial_blocks(self, mx):
|
|
"""store_cache should only create full blocks, skipping partial trailing tokens.
|
|
|
|
get_computed_blocks() matches full blocks only (floor division), so
|
|
partial blocks are never matched. Skipping them ensures is_last_block
|
|
points to the last full block, which is critical for ArraysCache/
|
|
RotatingKVCache last-block-only storage.
|
|
"""
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
# 10 tokens = 2 full blocks (8 tokens) + 2 partial tokens
|
|
tokens = list(range(10))
|
|
cache_data = [
|
|
{
|
|
"state": (mx.ones((1, 8, 10, 64)), mx.ones((1, 8, 10, 64))),
|
|
"cache_type": "KVCache",
|
|
"class_name": "KVCache",
|
|
}
|
|
]
|
|
|
|
result = cache.store_cache("req-001", tokens, cache_data)
|
|
|
|
assert result is not None
|
|
# Should create exactly 2 blocks (10 // 4 = 2), not 3
|
|
assert len(result.block_ids) == 2
|
|
# num_tokens should reflect only full blocks
|
|
assert result.num_tokens == 8 # 2 blocks * 4 tokens
|
|
|
|
stats = cache.get_stats()
|
|
assert stats.partial_block_skips == 1
|
|
assert stats.partial_tokens_skipped == 2
|
|
assert stats.last_partial_tokens_skipped == 2
|
|
assert stats.last_tokens_to_next_block == 2
|
|
|
|
def test_store_cache_arrayscache_partial_trailing_uses_last_full_block_state(
|
|
self, mx
|
|
):
|
|
"""ArraysCache with trailing partial tokens stores only full blocks safely."""
|
|
from omlx.cache.hybrid_cache import ModelCacheConfig
|
|
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
mock_ssd.save_block.return_value = True
|
|
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
# 7 tokens = 1 full block (4) + 3 trailing partial tokens
|
|
tokens = list(range(7))
|
|
conv_state = mx.ones((1, 3, 64))
|
|
ssm_state = mx.ones((1, 32, 128, 128))
|
|
cache_data = [
|
|
{
|
|
"state": (conv_state, ssm_state),
|
|
"cache_type": "ArraysCache",
|
|
"class_name": "ArraysCache",
|
|
}
|
|
]
|
|
model_cache_config = ModelCacheConfig.from_type_list(
|
|
["ArraysCache"], model_name="test-model"
|
|
)
|
|
|
|
result = cache.store_cache(
|
|
"req-001",
|
|
tokens,
|
|
cache_data,
|
|
model_cache_config=model_cache_config,
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result.block_ids) == 1
|
|
assert result.num_tokens == 4
|
|
mock_ssd.save_block.assert_called_once()
|
|
|
|
saved_data = mock_ssd.save_block.call_args.kwargs["cache_data"]
|
|
saved_conv_state, saved_ssm_state = saved_data[0]
|
|
assert saved_conv_state.shape == conv_state.shape
|
|
assert saved_ssm_state.shape == ssm_state.shape
|
|
|
|
def test_store_cache_all_partial_creates_no_blocks(self, mx):
|
|
"""Tokens fewer than block_size should create no blocks."""
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
# 3 tokens < block_size=4 -> 0 full blocks
|
|
tokens = [1, 2, 3]
|
|
cache_data = [
|
|
{
|
|
"state": (mx.ones((1, 8, 3, 64)), mx.ones((1, 8, 3, 64))),
|
|
"cache_type": "KVCache",
|
|
"class_name": "KVCache",
|
|
}
|
|
]
|
|
|
|
result = cache.store_cache("req-001", tokens, cache_data)
|
|
|
|
assert result is not None
|
|
assert len(result.block_ids) == 0
|
|
assert result.num_tokens == 0
|
|
|
|
stats = cache.get_stats()
|
|
assert stats.partial_block_skips == 1
|
|
assert stats.partial_tokens_skipped == 3
|
|
assert stats.last_partial_tokens_skipped == 3
|
|
assert stats.last_tokens_to_next_block == 1
|
|
|
|
def test_store_cache_exact_multiple_creates_all_blocks(self, mx):
|
|
"""Tokens exactly divisible by block_size should create all blocks."""
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
# 8 tokens = exactly 2 blocks
|
|
tokens = list(range(8))
|
|
cache_data = [
|
|
{
|
|
"state": (mx.ones((1, 8, 8, 64)), mx.ones((1, 8, 8, 64))),
|
|
"cache_type": "KVCache",
|
|
"class_name": "KVCache",
|
|
}
|
|
]
|
|
|
|
result = cache.store_cache("req-001", tokens, cache_data)
|
|
|
|
assert result is not None
|
|
assert len(result.block_ids) == 2
|
|
assert result.num_tokens == 8
|
|
|
|
def test_fetch_cache_with_segmented_extra_key_ranges(self):
|
|
"""Later image changes should preserve reuse before their boundary."""
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
)
|
|
|
|
tokens = list(range(12))
|
|
original_ranges = [
|
|
(5, ("image-1",)),
|
|
(9, ("image-1", "image-2")),
|
|
]
|
|
|
|
stored = cache.store_cache(
|
|
"req-store",
|
|
tokens,
|
|
[],
|
|
extra_key_ranges=original_ranges,
|
|
)
|
|
assert stored is not None
|
|
assert stored.num_tokens == 12
|
|
|
|
exact_table, exact_remaining = cache.fetch_cache(
|
|
"req-exact",
|
|
tokens,
|
|
extra_key_ranges=original_ranges,
|
|
)
|
|
assert exact_table is not None
|
|
assert exact_table.num_tokens == 12
|
|
assert exact_remaining == []
|
|
|
|
changed_later_image_table, changed_later_image_remaining = cache.fetch_cache(
|
|
"req-later-image",
|
|
tokens,
|
|
extra_key_ranges=[
|
|
(5, ("image-1",)),
|
|
(9, ("image-1", "image-3")),
|
|
],
|
|
)
|
|
assert changed_later_image_table is not None
|
|
assert changed_later_image_table.num_tokens == 8
|
|
assert changed_later_image_remaining == tokens[8:]
|
|
|
|
changed_earlier_image_table, changed_earlier_image_remaining = (
|
|
cache.fetch_cache(
|
|
"req-earlier-image",
|
|
tokens,
|
|
extra_key_ranges=[
|
|
(5, ("image-x",)),
|
|
(9, ("image-x", "image-2")),
|
|
],
|
|
)
|
|
)
|
|
assert changed_earlier_image_table is not None
|
|
assert changed_earlier_image_table.num_tokens == 4
|
|
assert changed_earlier_image_remaining == tokens[4:]
|
|
|
|
def test_store_cache_with_existing_prefix_uses_global_cache_indices(self, mx):
|
|
"""Store new blocks from full-sequence cache slices after cache hit.
|
|
|
|
When a request reuses prefix blocks, extracted cache data includes the
|
|
full sequence (prefix + newly processed suffix). New block slicing must
|
|
use global token indices, otherwise wrong KV ranges are persisted.
|
|
"""
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
mock_ssd.save_block.return_value = True
|
|
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
# Existing cached prefix (1 full block).
|
|
existing_tokens = [10, 11, 12, 13]
|
|
block_table = paged_cache.create_block_table("req-001")
|
|
existing_block = paged_cache.allocate_block()
|
|
assert existing_block is not None
|
|
existing_hash = compute_block_hash(
|
|
None, existing_tokens, model_name="test-model"
|
|
)
|
|
existing_block.block_hash = existing_hash
|
|
existing_block.token_count = block_size
|
|
block_table.block_ids.append(existing_block.block_id)
|
|
block_table.num_tokens = block_size
|
|
paged_cache.register_block_hash(existing_block, existing_tokens, None)
|
|
|
|
# Full sequence and full-sequence cache tensors.
|
|
full_tokens = existing_tokens + [20, 21, 22, 23]
|
|
keys = mx.arange(8, dtype=mx.float32).reshape(1, 1, 8, 1)
|
|
values = (mx.arange(8, dtype=mx.float32) + 100).reshape(1, 1, 8, 1)
|
|
cache_data = [
|
|
{"state": (keys, values), "cache_type": "KVCache", "class_name": "KVCache"}
|
|
]
|
|
|
|
result = cache.store_cache("req-001", full_tokens, cache_data)
|
|
|
|
assert result is not None
|
|
# Existing 1 block + newly stored 1 block
|
|
assert len(result.block_ids) == 2
|
|
assert mock_ssd.save_block.call_count == 1
|
|
|
|
saved_block_data = mock_ssd.save_block.call_args.kwargs["cache_data"]
|
|
saved_keys, saved_values = saved_block_data[0]
|
|
|
|
# New block must use global slice [4:8], not [0:4].
|
|
expected_keys = keys[:, :, 4:8, :]
|
|
expected_values = values[:, :, 4:8, :]
|
|
assert saved_keys.tolist() == expected_keys.tolist()
|
|
assert saved_values.tolist() == expected_values.tolist()
|
|
|
|
def test_store_cache_rolls_back_when_ssd_save_fails(self, mx):
|
|
"""Failed SSD save should not retain block metadata in paged cache."""
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
mock_ssd.save_block.return_value = False
|
|
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
tokens = [1, 2, 3, 4] # exactly one full block
|
|
keys = mx.ones((1, 8, 4, 64))
|
|
values = mx.ones((1, 8, 4, 64))
|
|
cache_data = [
|
|
{"state": (keys, values), "cache_type": "KVCache", "class_name": "KVCache"}
|
|
]
|
|
|
|
result = cache.store_cache("req-rollback", tokens, cache_data)
|
|
|
|
assert result is not None
|
|
# If persistence fails, block should be rolled back (not indexed/retained).
|
|
assert len(result.block_ids) == 0
|
|
assert result.num_tokens == 0
|
|
assert paged_cache.stats.allocated_blocks == 1 # null block only
|
|
|
|
failed_hash = compute_block_hash(None, tokens, model_name="test-model")
|
|
assert paged_cache.cached_block_hash_to_block.get_block(failed_hash) is None
|
|
|
|
def test_store_cache_keeps_valid_prefix_when_later_ssd_save_fails(self, mx):
|
|
"""A later SSD save failure should roll back only the failed block."""
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
mock_ssd.save_block.side_effect = [True, False]
|
|
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
# Seed request with one already-cached prefix block.
|
|
existing_tokens = [10, 11, 12, 13]
|
|
existing_block = paged_cache.allocate_block()
|
|
assert existing_block is not None
|
|
existing_block.token_count = block_size
|
|
block_table = paged_cache.create_block_table("req-partial-rollback")
|
|
block_table.block_ids.append(existing_block.block_id)
|
|
block_table.num_tokens = block_size
|
|
paged_cache.register_block_hash(existing_block, existing_tokens, None)
|
|
assert existing_block.block_hash is not None
|
|
|
|
# Add two more blocks; first save succeeds, second save fails.
|
|
tokens = existing_tokens + [20, 21, 22, 23, 30, 31, 32, 33]
|
|
keys = mx.arange(12, dtype=mx.float32).reshape(1, 1, 12, 1)
|
|
values = (mx.arange(12, dtype=mx.float32) + 100).reshape(1, 1, 12, 1)
|
|
cache_data = [
|
|
{"state": (keys, values), "cache_type": "KVCache", "class_name": "KVCache"}
|
|
]
|
|
|
|
result = cache.store_cache("req-partial-rollback", tokens, cache_data)
|
|
|
|
assert result is not None
|
|
|
|
first_new_hash = compute_block_hash(
|
|
existing_block.block_hash, tokens[4:8], model_name="test-model"
|
|
)
|
|
failed_hash = compute_block_hash(
|
|
first_new_hash, tokens[8:12], model_name="test-model"
|
|
)
|
|
|
|
# Keep existing block + first new block; drop only the failed second new block.
|
|
assert len(result.block_ids) == 2
|
|
assert result.num_tokens == 8
|
|
assert result.block_ids[0] == existing_block.block_id
|
|
|
|
first_new_block = paged_cache.cached_block_hash_to_block.get_block(
|
|
first_new_hash
|
|
)
|
|
assert first_new_block is not None
|
|
assert result.block_ids[1] == first_new_block.block_id
|
|
assert result.block_ids == [existing_block.block_id, first_new_block.block_id]
|
|
|
|
# save_block should attempt exactly two writes in this scenario.
|
|
calls = mock_ssd.save_block.call_args_list
|
|
assert len(calls) == 2
|
|
attempted_hashes = [call.kwargs["block_hash"] for call in calls]
|
|
assert attempted_hashes == [first_new_hash, failed_hash]
|
|
assert [call.kwargs["token_count"] for call in calls] == [
|
|
block_size,
|
|
block_size,
|
|
]
|
|
|
|
# Verify global-index slices were persisted for both attempted new blocks.
|
|
first_saved_keys, first_saved_values = calls[0].kwargs["cache_data"][0]
|
|
failed_saved_keys, failed_saved_values = calls[1].kwargs["cache_data"][0]
|
|
assert first_saved_keys.tolist() == keys[:, :, 4:8, :].tolist()
|
|
assert first_saved_values.tolist() == values[:, :, 4:8, :].tolist()
|
|
assert failed_saved_keys.tolist() == keys[:, :, 8:12, :].tolist()
|
|
assert failed_saved_values.tolist() == values[:, :, 8:12, :].tolist()
|
|
|
|
assert (
|
|
paged_cache.cached_block_hash_to_block.get_block(first_new_hash) is not None
|
|
)
|
|
assert paged_cache.cached_block_hash_to_block.get_block(failed_hash) is None
|
|
# Failed block should be freed, not just removed from hash index.
|
|
allocated_non_null_ids = {
|
|
block.block_id
|
|
for block in paged_cache.allocated_blocks.values()
|
|
if not block.is_null
|
|
}
|
|
assert allocated_non_null_ids == {
|
|
existing_block.block_id,
|
|
first_new_block.block_id,
|
|
}
|
|
assert all(
|
|
b.block_hash != failed_hash for b in paged_cache.allocated_blocks.values()
|
|
)
|
|
|
|
# Public contract after partial failure: only valid prefix should be reused.
|
|
expected_partial_ids = [existing_block.block_id, first_new_block.block_id]
|
|
fetched_partial, remaining_partial = cache.fetch_cache(
|
|
"req-partial-rollback-hit", tokens
|
|
)
|
|
assert fetched_partial is not None
|
|
assert fetched_partial.block_ids == expected_partial_ids
|
|
assert fetched_partial.num_tokens == 8
|
|
assert remaining_partial == tokens[8:12]
|
|
|
|
def test_store_cache_retry_after_partial_failure_saves_only_missing_tail(self, mx):
|
|
"""Retry should preserve valid prefix and only save the missing tail block."""
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
mock_ssd.save_block.side_effect = [True, False, True]
|
|
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
existing_tokens = [10, 11, 12, 13]
|
|
existing_block = paged_cache.allocate_block()
|
|
assert existing_block is not None
|
|
existing_block.token_count = block_size
|
|
block_table = paged_cache.create_block_table("req-retry")
|
|
block_table.block_ids.append(existing_block.block_id)
|
|
block_table.num_tokens = block_size
|
|
paged_cache.register_block_hash(existing_block, existing_tokens, None)
|
|
assert existing_block.block_hash is not None
|
|
|
|
tokens = existing_tokens + [20, 21, 22, 23, 30, 31, 32, 33]
|
|
keys = mx.arange(12, dtype=mx.float32).reshape(1, 1, 12, 1)
|
|
values = (mx.arange(12, dtype=mx.float32) + 100).reshape(1, 1, 12, 1)
|
|
cache_data = [
|
|
{"state": (keys, values), "cache_type": "KVCache", "class_name": "KVCache"}
|
|
]
|
|
|
|
first_result = cache.store_cache("req-retry", tokens, cache_data)
|
|
assert first_result is not None
|
|
|
|
first_new_hash = compute_block_hash(
|
|
existing_block.block_hash, tokens[4:8], model_name="test-model"
|
|
)
|
|
tail_hash = compute_block_hash(
|
|
first_new_hash, tokens[8:12], model_name="test-model"
|
|
)
|
|
first_new_block = paged_cache.cached_block_hash_to_block.get_block(
|
|
first_new_hash
|
|
)
|
|
assert first_new_block is not None
|
|
retained_prefix_ids = first_result.block_ids.copy()
|
|
assert retained_prefix_ids == [
|
|
existing_block.block_id,
|
|
first_new_block.block_id,
|
|
]
|
|
|
|
retry_result = cache.store_cache("req-retry", tokens, cache_data)
|
|
assert retry_result is not None
|
|
|
|
calls = mock_ssd.save_block.call_args_list
|
|
assert len(calls) == 3
|
|
attempted_hashes = [call.kwargs["block_hash"] for call in calls]
|
|
assert attempted_hashes == [first_new_hash, tail_hash, tail_hash]
|
|
assert attempted_hashes.count(first_new_hash) == 1
|
|
assert attempted_hashes.count(tail_hash) == 2
|
|
assert [call.kwargs["token_count"] for call in calls] == [
|
|
block_size,
|
|
block_size,
|
|
block_size,
|
|
]
|
|
retry_saved_keys, retry_saved_values = calls[2].kwargs["cache_data"][0]
|
|
assert retry_saved_keys.tolist() == keys[:, :, 8:12, :].tolist()
|
|
assert retry_saved_values.tolist() == values[:, :, 8:12, :].tolist()
|
|
|
|
assert len(retry_result.block_ids) == 3
|
|
assert retry_result.num_tokens == 12
|
|
assert retry_result.block_ids[:2] == retained_prefix_ids
|
|
assert len(set(retry_result.block_ids)) == 3
|
|
|
|
tail_block = paged_cache.cached_block_hash_to_block.get_block(tail_hash)
|
|
assert tail_block is not None
|
|
assert retry_result.block_ids[2] == tail_block.block_id
|
|
|
|
# Reconstruct full cache and verify tensor content for retry flow.
|
|
saved_by_hash = {
|
|
existing_block.block_hash: [
|
|
(keys[:, :, 0:4, :], values[:, :, 0:4, :]),
|
|
],
|
|
}
|
|
for call in calls:
|
|
saved_by_hash[call.kwargs["block_hash"]] = call.kwargs["cache_data"]
|
|
|
|
def load_block_with_metadata(block_hash):
|
|
block_data = saved_by_hash.get(block_hash)
|
|
if block_data is None:
|
|
return None, None
|
|
return (
|
|
block_data,
|
|
{
|
|
"model_name": "test-model",
|
|
"num_layers": 1,
|
|
"layer_cache_types": ["KVCache"],
|
|
"layer_meta_states": [()],
|
|
},
|
|
)
|
|
|
|
mock_ssd.load_block_with_metadata.side_effect = load_block_with_metadata
|
|
reconstructed = cache.reconstruct_cache(retry_result)
|
|
assert reconstructed is not None
|
|
assert len(reconstructed) == 1
|
|
layer_cache = reconstructed[0]
|
|
if hasattr(layer_cache, "state"):
|
|
reconstructed_keys, reconstructed_values = layer_cache.state
|
|
elif isinstance(layer_cache, (list, tuple)) and len(layer_cache) == 2:
|
|
reconstructed_keys, reconstructed_values = layer_cache
|
|
else:
|
|
reconstructed_keys, reconstructed_values = (
|
|
layer_cache.keys,
|
|
layer_cache.values,
|
|
)
|
|
|
|
assert reconstructed_keys.tolist() == keys.tolist()
|
|
assert reconstructed_values.tolist() == values.tolist()
|
|
|
|
# Force prefix-index fallback by removing chain-hash index entries.
|
|
for block_id in retry_result.block_ids:
|
|
block = paged_cache.allocated_blocks.get(block_id)
|
|
assert block is not None
|
|
assert block.block_hash is not None
|
|
paged_cache.cached_block_hash_to_block.pop(block.block_hash, block.block_id)
|
|
|
|
# Explicitly prove shared-hash path cannot succeed in this fixture.
|
|
assert paged_cache._paged_ssd_cache_manager is None
|
|
shared_block_ids, _ = paged_cache.find_shared_prefix(tokens)
|
|
assert shared_block_ids == []
|
|
|
|
expected_ids = retry_result.block_ids.copy()
|
|
# Public contract via prefix-index fallback: full prefix hit, no remaining tokens.
|
|
fetched_table, remaining = cache.fetch_cache(
|
|
"req-retry-prefix-index-hit", tokens
|
|
)
|
|
assert fetched_table is not None
|
|
assert fetched_table.block_ids == expected_ids
|
|
assert fetched_table.num_tokens == 12
|
|
assert remaining == []
|
|
|
|
|
|
class TestPrefixCacheCacheList:
|
|
"""Tests for CacheList support in BlockAwarePrefixCache."""
|
|
|
|
@pytest.fixture
|
|
def mx(self):
|
|
"""Import MLX or skip."""
|
|
try:
|
|
import mlx.core as mx
|
|
|
|
return mx
|
|
except ImportError:
|
|
pytest.skip("MLX not available")
|
|
|
|
@pytest.fixture
|
|
def paged_cache(self):
|
|
"""Create a PagedCacheManager for testing."""
|
|
return PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def mock_model(self):
|
|
"""Create a mock model."""
|
|
return MockModel(num_layers=4)
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self, mock_model, paged_cache):
|
|
"""Create a BlockAwarePrefixCache for testing."""
|
|
return BlockAwarePrefixCache(
|
|
model=mock_model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
|
|
def test_get_cache_seq_len_cache_list_only(self, prefix_cache, mx):
|
|
"""Test _get_cache_seq_len with all-CacheList model (e.g., deepseek_v32)."""
|
|
# CacheList sub-states with 4D tensors
|
|
sub_keys = mx.zeros((1, 8, 128, 64))
|
|
sub_values = mx.zeros((1, 8, 128, 64))
|
|
|
|
cache_data = [
|
|
{
|
|
"state": [(sub_keys, sub_values)], # CacheList sub-states
|
|
"cache_type": "CacheList",
|
|
"class_name": "CacheList",
|
|
},
|
|
{
|
|
"state": [(sub_keys, sub_values)],
|
|
"cache_type": "CacheList",
|
|
"class_name": "CacheList",
|
|
},
|
|
]
|
|
|
|
seq_len = prefix_cache._get_cache_seq_len(cache_data)
|
|
assert seq_len == 128
|
|
|
|
def test_get_cache_seq_len_mixed_kvcache_and_cache_list(self, prefix_cache, mx):
|
|
"""Test _get_cache_seq_len with mixed KVCache + CacheList model."""
|
|
kv_keys = mx.zeros((1, 8, 256, 64))
|
|
kv_values = mx.zeros((1, 8, 256, 64))
|
|
sub_keys = mx.zeros((1, 4, 128, 64))
|
|
|
|
cache_data = [
|
|
{
|
|
"state": (kv_keys, kv_values),
|
|
"cache_type": "KVCache",
|
|
"class_name": "KVCache",
|
|
},
|
|
{
|
|
"state": [(sub_keys, MagicMock())],
|
|
"cache_type": "CacheList",
|
|
"class_name": "CacheList",
|
|
},
|
|
]
|
|
|
|
seq_len = prefix_cache._get_cache_seq_len(cache_data)
|
|
# Should use KVCache layer (Step 1), not CacheList
|
|
assert seq_len == 256
|
|
|
|
def test_get_cache_seq_len_cache_list_skipped_in_step1(self, prefix_cache, mx):
|
|
"""Test CacheList is skipped in Step 1 and Step 2."""
|
|
sub_keys = mx.zeros((1, 8, 64, 64))
|
|
|
|
cache_data = [
|
|
{
|
|
"state": [(sub_keys, MagicMock())],
|
|
"cache_type": "CacheList",
|
|
"class_name": "CacheList",
|
|
},
|
|
]
|
|
|
|
# CacheList is in non_sliceable_types, so Step 1 & 2 skip it
|
|
# Step 3 finds the sub-cache seq_len
|
|
seq_len = prefix_cache._get_cache_seq_len(cache_data)
|
|
assert seq_len == 64
|
|
|
|
def test_get_cache_seq_len_pure_rotating_kvcache(self, prefix_cache, mx):
|
|
"""Test _get_cache_seq_len for pure RotatingKVCache model."""
|
|
rot_keys = mx.zeros((1, 8, 96, 64))
|
|
rot_values = mx.zeros((1, 8, 96, 64))
|
|
|
|
cache_data = [
|
|
{
|
|
"state": (rot_keys, rot_values),
|
|
"cache_type": "RotatingKVCache",
|
|
"class_name": "RotatingKVCache",
|
|
},
|
|
]
|
|
|
|
# Step 1 skips RotatingKVCache; Step 2 must still recover seq_len.
|
|
seq_len = prefix_cache._get_cache_seq_len(cache_data)
|
|
assert seq_len == 96
|
|
|
|
def test_extract_block_tensor_slice_cache_list_last_block(self, prefix_cache, mx):
|
|
"""Test _extract_block_tensor_slice for CacheList on last block."""
|
|
from omlx.cache.hybrid_cache import ModelCacheConfig
|
|
|
|
sub_keys = mx.zeros((1, 8, 32, 64))
|
|
sub_values = mx.ones((1, 8, 32, 64))
|
|
|
|
cache_data = [
|
|
{
|
|
"state": [(sub_keys, sub_values)],
|
|
"cache_type": "CacheList",
|
|
"class_name": "CacheList",
|
|
},
|
|
]
|
|
config = ModelCacheConfig.from_type_list(["CacheList"])
|
|
|
|
result = prefix_cache._extract_block_tensor_slice(
|
|
cache_data,
|
|
0,
|
|
32,
|
|
model_cache_config=config,
|
|
is_last_block=True,
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result) == 1
|
|
# CacheList marker format
|
|
assert result[0][0] == "__cache_list__"
|
|
assert len(result[0][1]) == 1 # One sub-cache
|
|
assert result[0][1][0][0].shape == (1, 8, 32, 64)
|
|
|
|
def test_extract_block_tensor_slice_cache_list_non_last_sliceable(
|
|
self, prefix_cache, mx
|
|
):
|
|
"""Test _extract_block_tensor_slice for CacheList with sliceable sub-caches on non-last block.
|
|
|
|
When all sub-caches are 4D KVCache tensors, they should be sliced
|
|
per-block instead of using last-block-only placeholder storage.
|
|
"""
|
|
from omlx.cache.hybrid_cache import ModelCacheConfig
|
|
|
|
sub_keys = mx.zeros((1, 8, 32, 64))
|
|
sub_values = mx.ones((1, 8, 32, 64))
|
|
|
|
cache_data = [
|
|
{
|
|
"state": [(sub_keys, sub_values)],
|
|
"cache_type": "CacheList",
|
|
"class_name": "CacheList",
|
|
},
|
|
]
|
|
config = ModelCacheConfig.from_type_list(["CacheList"])
|
|
|
|
result = prefix_cache._extract_block_tensor_slice(
|
|
cache_data,
|
|
0,
|
|
16,
|
|
model_cache_config=config,
|
|
is_last_block=False,
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result) == 1
|
|
# Sliceable sub-caches: per-block sliced data, not placeholder
|
|
assert result[0][0] == "__cache_list__"
|
|
assert len(result[0][1]) == 1
|
|
assert result[0][1][0][0].shape == (1, 8, 16, 64)
|
|
assert result[0][1][0][1].shape == (1, 8, 16, 64)
|
|
|
|
def test_extract_block_tensor_slice_cache_list_zero_dim_values(
|
|
self, prefix_cache, mx
|
|
):
|
|
"""Test per-block slicing for CacheList with zero-dim values (DSA indexer)."""
|
|
from omlx.cache.hybrid_cache import ModelCacheConfig
|
|
|
|
# GLM-5 style: main attention + indexer with zero head_dim
|
|
sub_keys1 = mx.zeros((1, 1, 64, 512))
|
|
sub_values1 = mx.zeros((1, 1, 64, 64))
|
|
sub_keys2 = mx.zeros((1, 1, 64, 128))
|
|
sub_values2 = mx.zeros((1, 1, 64, 0)) # zero head_dim
|
|
|
|
cache_data = [
|
|
{
|
|
"state": [(sub_keys1, sub_values1), (sub_keys2, sub_values2)],
|
|
"cache_type": "CacheList",
|
|
"class_name": "CacheList",
|
|
},
|
|
]
|
|
config = ModelCacheConfig.from_type_list(["CacheList"])
|
|
|
|
result = prefix_cache._extract_block_tensor_slice(
|
|
cache_data,
|
|
0,
|
|
32,
|
|
model_cache_config=config,
|
|
is_last_block=False,
|
|
)
|
|
|
|
assert result is not None
|
|
assert result[0][0] == "__cache_list__"
|
|
assert len(result[0][1]) == 2
|
|
# Sub-cache 0: sliced normally
|
|
assert result[0][1][0][0].shape == (1, 1, 32, 512)
|
|
assert result[0][1][0][1].shape == (1, 1, 32, 64)
|
|
# Sub-cache 1: sliced, values remain zero-dim
|
|
assert result[0][1][1][0].shape == (1, 1, 32, 128)
|
|
assert result[0][1][1][1].shape == (1, 1, 32, 0)
|
|
|
|
def test_validate_block_cache_data_cache_list(self, prefix_cache, mx):
|
|
"""Test _validate_block_cache_data with CacheList layers."""
|
|
# CacheList as list format (last block)
|
|
cache_data = [
|
|
[
|
|
(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64)))
|
|
], # CacheList sub-cache list
|
|
(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64))), # Standard KVCache
|
|
]
|
|
layer_cache_types = ["CacheList", "KVCache"]
|
|
|
|
result = prefix_cache._validate_block_cache_data(cache_data, layer_cache_types)
|
|
assert result is True
|
|
|
|
def test_validate_block_cache_data_cache_list_placeholder(self, prefix_cache, mx):
|
|
"""Test _validate_block_cache_data with CacheList placeholder."""
|
|
cache_data = [
|
|
(
|
|
mx.zeros((1,)),
|
|
mx.zeros((1,)),
|
|
), # CacheList placeholder (falls through to tuple check)
|
|
(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64))), # Standard KVCache
|
|
]
|
|
layer_cache_types = ["CacheList", "KVCache"]
|
|
|
|
result = prefix_cache._validate_block_cache_data(cache_data, layer_cache_types)
|
|
assert result is True
|
|
|
|
def test_find_kv_shape_ref_skips_cache_list(self, prefix_cache, mx):
|
|
"""Test _find_kv_shape_ref skips CacheList layers."""
|
|
all_block_data = [
|
|
[
|
|
[
|
|
(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64)))
|
|
], # CacheList: List[Tuple]
|
|
(mx.zeros((1, 4, 32, 128)), mx.zeros((1, 4, 32, 128))), # KVCache
|
|
]
|
|
]
|
|
layer_cache_types = ["CacheList", "KVCache"]
|
|
|
|
result = prefix_cache._find_kv_shape_ref(all_block_data, layer_cache_types)
|
|
assert result == (4, 128) # From KVCache layer, not CacheList
|
|
|
|
def test_reconstruct_cache_list_partial_match_reject(self, mx):
|
|
"""Test reconstruct_cache rejects CacheList with placeholder (partial match)."""
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
|
|
model = MockModel(num_layers=1)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
# Allocate a block
|
|
block1 = paged_cache.allocate_block()
|
|
block1.block_hash = b"hash_cl1"
|
|
block1.token_count = 4
|
|
|
|
block_table = BlockTable(
|
|
request_id="req-cl-partial",
|
|
block_ids=[block1.block_id],
|
|
num_tokens=4,
|
|
)
|
|
|
|
# Block data: CacheList placeholder (partial match)
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
block1_data = [placeholder]
|
|
|
|
mock_ssd.load_block_with_metadata.return_value = (
|
|
block1_data,
|
|
{
|
|
"model_name": "test-model",
|
|
"num_layers": 1,
|
|
"layer_cache_types": ["CacheList"],
|
|
"layer_meta_states": [()],
|
|
},
|
|
)
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
# Should return None because CacheList layer has placeholder
|
|
assert result is None
|
|
|
|
|
|
class TestWalkBackTruncation:
|
|
"""Tests for walk-back truncation of non-sliceable caches."""
|
|
|
|
@pytest.fixture
|
|
def mx(self):
|
|
"""Import MLX or skip."""
|
|
try:
|
|
import mlx.core as mx
|
|
|
|
return mx
|
|
except ImportError:
|
|
pytest.skip("MLX not available")
|
|
|
|
# ------------------------------------------------------------------
|
|
# _is_placeholder_state
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_is_placeholder_state_detects_placeholder(self, mx):
|
|
"""Placeholder tuple with shape (1,) should be detected."""
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
assert BlockAwarePrefixCache._is_placeholder_state(placeholder) is True
|
|
|
|
def test_is_placeholder_state_rejects_real_arrays_cache(self, mx):
|
|
"""Real ArraysCache state should not be flagged as placeholder."""
|
|
real_state = (mx.ones((1, 3, 64)), mx.ones((1, 32, 128, 128)))
|
|
assert BlockAwarePrefixCache._is_placeholder_state(real_state) is False
|
|
|
|
def test_is_placeholder_state_rejects_kv_cache(self, mx):
|
|
"""Standard 4D KVCache tensors should not be flagged."""
|
|
kv_state = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
assert BlockAwarePrefixCache._is_placeholder_state(kv_state) is False
|
|
|
|
def test_is_placeholder_state_rejects_list(self, mx):
|
|
"""CacheList real data (list format) should not be flagged."""
|
|
cache_list_data = [
|
|
(mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64))),
|
|
]
|
|
assert BlockAwarePrefixCache._is_placeholder_state(cache_list_data) is False
|
|
|
|
# ------------------------------------------------------------------
|
|
# _find_walk_back_truncation_point
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_walk_back_no_truncation_when_last_block_valid(self, mx):
|
|
"""No truncation when the last block has real state."""
|
|
model = MockModel(num_layers=2)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
real_state = (mx.ones((1, 3, 64)), mx.ones((1, 32, 128, 128)))
|
|
kv = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
|
|
all_block_data = [
|
|
[kv, placeholder], # block 0
|
|
[kv, real_state], # block 1 (last, valid)
|
|
]
|
|
layer_cache_types = ["KVCache", "ArraysCache"]
|
|
|
|
result = cache._find_walk_back_truncation_point(
|
|
all_block_data, layer_cache_types
|
|
)
|
|
assert result is None # No truncation needed
|
|
|
|
def test_walk_back_multi_turn_pattern(self, mx):
|
|
"""Walk-back finds the latest block with valid ArraysCache state.
|
|
|
|
Simulates multi-turn pattern:
|
|
A[p] B[p] C[real] D[p] E[real] F[p]
|
|
Should walk back to E (index 4).
|
|
"""
|
|
model = MockModel(num_layers=2)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
real_state = (mx.ones((1, 3, 64)), mx.ones((1, 32, 128, 128)))
|
|
kv = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
|
|
# A[p] B[p] C[real] D[p] E[real] F[p]
|
|
all_block_data = [
|
|
[kv, placeholder], # A
|
|
[kv, placeholder], # B
|
|
[kv, real_state], # C (turn 1 last block)
|
|
[kv, placeholder], # D
|
|
[kv, real_state], # E (turn 2 last block)
|
|
[kv, placeholder], # F (last loaded, placeholder)
|
|
]
|
|
layer_cache_types = ["KVCache", "ArraysCache"]
|
|
|
|
result = cache._find_walk_back_truncation_point(
|
|
all_block_data, layer_cache_types
|
|
)
|
|
assert result == 4 # Block E (index 4)
|
|
|
|
def test_walk_back_all_placeholders_returns_none(self, mx):
|
|
"""All blocks have placeholders -- no valid fallback exists."""
|
|
model = MockModel(num_layers=2)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
kv = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
|
|
all_block_data = [
|
|
[kv, placeholder],
|
|
[kv, placeholder],
|
|
[kv, placeholder],
|
|
]
|
|
layer_cache_types = ["KVCache", "ArraysCache"]
|
|
|
|
result = cache._find_walk_back_truncation_point(
|
|
all_block_data, layer_cache_types
|
|
)
|
|
assert result is None # No valid block found
|
|
|
|
def test_walk_back_includes_rotating_kv_cache(self, mx):
|
|
"""RotatingKVCache placeholders should walk back to latest valid block."""
|
|
model = MockModel(num_layers=2)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
rotating_real = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
kv = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
|
|
# Last block placeholder, previous block valid rotating state.
|
|
all_block_data = [
|
|
[kv, rotating_real],
|
|
[kv, placeholder],
|
|
]
|
|
layer_cache_types = ["KVCache", "RotatingKVCache"]
|
|
|
|
result = cache._find_walk_back_truncation_point(
|
|
all_block_data, layer_cache_types
|
|
)
|
|
assert result == 0
|
|
|
|
# ------------------------------------------------------------------
|
|
# Full reconstruct_cache integration with walk-back
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_reconstruct_arrays_cache_walks_back_to_valid_block(self, mx):
|
|
"""Partial match should walk back instead of rejecting entirely.
|
|
|
|
3 blocks: block0[p] block1[real] block2[p]
|
|
Should truncate to blocks 0-1, returning valid cache.
|
|
"""
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
|
|
model = MockModel(num_layers=2)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
# Allocate 3 blocks
|
|
blocks = []
|
|
for i in range(3):
|
|
b = paged_cache.allocate_block()
|
|
b.block_hash = f"hash{i}".encode()
|
|
b.token_count = 4
|
|
b.ref_count = 2 # Simulate fetch_cache having incremented ref
|
|
blocks.append(b)
|
|
|
|
block_table = BlockTable(
|
|
request_id="req-001",
|
|
block_ids=[b.block_id for b in blocks],
|
|
num_tokens=12,
|
|
)
|
|
|
|
kv_slice = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
conv_state = mx.ones((1, 3, 64))
|
|
ssm_state = mx.ones((1, 32, 128, 128))
|
|
real_state = (conv_state, ssm_state)
|
|
|
|
metadata = {
|
|
"model_name": "test-model",
|
|
"num_layers": 2,
|
|
"layer_cache_types": ["KVCache", "ArraysCache"],
|
|
"layer_meta_states": [(), ()],
|
|
}
|
|
|
|
mock_ssd.load_block_with_metadata.side_effect = [
|
|
([kv_slice, placeholder], metadata), # block 0: placeholder
|
|
([kv_slice, real_state], metadata), # block 1: real state (turn 1 last)
|
|
([kv_slice, placeholder], metadata), # block 2: placeholder
|
|
]
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
|
|
# Should NOT be None -- walk-back recovered blocks 0-1
|
|
assert result is not None
|
|
assert len(result) == 2 # 2 layers reconstructed
|
|
|
|
# block_table should be truncated to 2 blocks
|
|
assert len(block_table.block_ids) == 2
|
|
assert block_table.num_tokens == 8
|
|
|
|
# Block 2 ref_count should have been decremented (freed)
|
|
assert blocks[2].ref_count == 1
|
|
|
|
def test_reconstruct_rotating_cache_walks_back_to_valid_block(self, mx):
|
|
"""Rotating partial match should walk back to latest valid block."""
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
|
|
model = MockModel(num_layers=2)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
# Allocate 3 blocks
|
|
blocks = []
|
|
for i in range(3):
|
|
b = paged_cache.allocate_block()
|
|
b.block_hash = f"hash{i}".encode()
|
|
b.token_count = 4
|
|
b.ref_count = 2 # Simulate fetch_cache having incremented ref
|
|
blocks.append(b)
|
|
|
|
block_table = BlockTable(
|
|
request_id="req-001",
|
|
block_ids=[b.block_id for b in blocks],
|
|
num_tokens=12,
|
|
)
|
|
|
|
kv_slice = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
rotating_real = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
|
|
metadata = {
|
|
"model_name": "test-model",
|
|
"num_layers": 2,
|
|
"layer_cache_types": ["KVCache", "RotatingKVCache"],
|
|
"layer_meta_states": [(), (0, 4, 4, 4)],
|
|
}
|
|
|
|
mock_ssd.load_block_with_metadata.side_effect = [
|
|
([kv_slice, placeholder], metadata), # block 0: placeholder
|
|
([kv_slice, rotating_real], metadata), # block 1: real rotating state
|
|
([kv_slice, placeholder], metadata), # block 2: placeholder
|
|
]
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
|
|
# Should recover blocks 0-1 via walk-back.
|
|
assert result is not None
|
|
assert len(result) == 2
|
|
assert len(block_table.block_ids) == 2
|
|
assert block_table.num_tokens == 8
|
|
assert blocks[2].ref_count == 1
|
|
|
|
rotating_cache = result[1]
|
|
assert hasattr(rotating_cache, "max_size")
|
|
assert rotating_cache.max_size == 4
|
|
|
|
def test_reconstruct_all_placeholders_still_rejects(self, mx):
|
|
"""When no block has valid state, walk-back finds nothing and
|
|
the existing per-layer rejection returns None."""
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
|
|
model = MockModel(num_layers=2)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
blocks = []
|
|
for i in range(2):
|
|
b = paged_cache.allocate_block()
|
|
b.block_hash = f"hash{i}".encode()
|
|
b.token_count = 4
|
|
blocks.append(b)
|
|
|
|
block_table = BlockTable(
|
|
request_id="req-001",
|
|
block_ids=[b.block_id for b in blocks],
|
|
num_tokens=8,
|
|
)
|
|
|
|
kv_slice = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
placeholder = (mx.zeros((1,)), mx.zeros((1,)))
|
|
|
|
metadata = {
|
|
"model_name": "test-model",
|
|
"num_layers": 2,
|
|
"layer_cache_types": ["KVCache", "ArraysCache"],
|
|
"layer_meta_states": [(), ()],
|
|
}
|
|
|
|
mock_ssd.load_block_with_metadata.side_effect = [
|
|
([kv_slice, placeholder], metadata),
|
|
([kv_slice, placeholder], metadata),
|
|
]
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
|
|
# Should still return None -- all placeholders, no walk-back target
|
|
assert result is None
|
|
|
|
def test_partial_reconstruction_frees_dropped_blocks(self, mx):
|
|
"""Blocks dropped during partial reconstruction should have
|
|
their ref_counts decremented."""
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
|
|
model = MockModel(num_layers=1)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
block1 = paged_cache.allocate_block()
|
|
block1.block_hash = b"hash1"
|
|
block1.token_count = 4
|
|
block1.ref_count = 2 # Simulate fetch_cache increment
|
|
|
|
block2 = paged_cache.allocate_block()
|
|
block2.block_hash = b"hash2"
|
|
block2.token_count = 4
|
|
block2.ref_count = 2 # Simulate fetch_cache increment
|
|
|
|
block_table = BlockTable(
|
|
request_id="req-001",
|
|
block_ids=[block1.block_id, block2.block_id],
|
|
num_tokens=8,
|
|
)
|
|
|
|
kv_slice = (mx.ones((1, 8, 4, 64)), mx.ones((1, 8, 4, 64)))
|
|
|
|
# First block loads fine, second fails
|
|
mock_ssd.load_block_with_metadata.side_effect = [
|
|
(
|
|
[kv_slice],
|
|
{
|
|
"model_name": "test-model",
|
|
"num_layers": 1,
|
|
"layer_cache_types": ["KVCache"],
|
|
"layer_meta_states": [()],
|
|
},
|
|
),
|
|
(None, None), # Second block fails to load
|
|
]
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
|
|
# Should partially reconstruct with block1 only
|
|
assert result is not None
|
|
assert len(block_table.block_ids) == 1
|
|
assert block_table.num_tokens == 4
|
|
|
|
# block2 ref_count should have been decremented
|
|
assert block2.ref_count == 1
|
|
|
|
|
|
class TestTurboQuantFormatMismatchRecovery:
|
|
"""Regression tests for TurboQuant/fp16 prefix-chain format mismatches."""
|
|
|
|
@pytest.fixture
|
|
def mx(self):
|
|
"""Import MLX or skip."""
|
|
try:
|
|
import mlx.core as mx
|
|
|
|
return mx
|
|
except ImportError:
|
|
pytest.skip("MLX not available")
|
|
|
|
def test_reconstruct_truncates_turboquant_chain_at_fp16_tail(self, mx):
|
|
"""A pre-fix fp16 tail after TQ blocks should heal by truncation."""
|
|
from mlx_lm.models.cache import KVCache
|
|
from mlx_vlm.turboquant import TurboQuantKVCache
|
|
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
mock_ssd.forget_block.return_value = True
|
|
|
|
model = MockModel(num_layers=1)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
blocks = []
|
|
for i in range(2):
|
|
block = paged_cache.allocate_block()
|
|
block.block_hash = f"hash-tq-{i}".encode()
|
|
block.token_count = 4
|
|
block.ref_count = 2
|
|
paged_cache.cached_block_hash_to_block.insert(block.block_hash, block)
|
|
blocks.append(block)
|
|
|
|
block_table = BlockTable(
|
|
request_id="req-tq-mixed",
|
|
block_ids=[block.block_id for block in blocks],
|
|
num_tokens=8,
|
|
)
|
|
|
|
kv_cache = KVCache()
|
|
kv_cache.update_and_fetch(
|
|
mx.random.normal((1, 2, 4, 32)),
|
|
mx.random.normal((1, 2, 4, 32)),
|
|
)
|
|
tq_cache = TurboQuantKVCache.from_cache(kv_cache, bits=4.0)
|
|
tq_keys, tq_values = tq_cache.state
|
|
tq_block = [("__turboquant_v2__", (tq_keys, tq_values))]
|
|
|
|
fp16_tail = [
|
|
(
|
|
mx.random.normal((1, 2, 4, 32)),
|
|
mx.random.normal((1, 2, 4, 32)),
|
|
)
|
|
]
|
|
tq_metadata = {
|
|
"model_name": "test-model",
|
|
"num_layers": 1,
|
|
"block_size": 4,
|
|
"layer_cache_types": ["TurboQuantKVCache"],
|
|
"layer_meta_states": [tq_cache.meta_state],
|
|
}
|
|
fp16_metadata = {
|
|
"model_name": "test-model",
|
|
"num_layers": 1,
|
|
"block_size": 4,
|
|
"layer_cache_types": ["KVCache"],
|
|
"layer_meta_states": [(4,)],
|
|
}
|
|
mock_ssd.load_block_with_metadata.side_effect = [
|
|
(tq_block, tq_metadata),
|
|
(fp16_tail, fp16_metadata),
|
|
]
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
|
|
assert result is not None
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], TurboQuantKVCache)
|
|
assert block_table.block_ids == [blocks[0].block_id]
|
|
assert block_table.num_tokens == 4
|
|
assert blocks[1].ref_count == 1
|
|
mock_ssd.forget_block.assert_called_once_with(blocks[1].block_hash)
|
|
assert paged_cache.cached_block_hash_to_block.get_block(
|
|
blocks[1].block_hash
|
|
) is None
|
|
|
|
def test_reconstruct_rejects_stale_first_block_with_manager_signature(self, mx):
|
|
"""A live manager signature must make stale block 0 fail its own check."""
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
mock_ssd._expected_layer_cache_types = ["TurboQuantKVCache"]
|
|
mock_ssd.forget_block.return_value = True
|
|
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=MockModel(num_layers=1),
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
block = paged_cache.allocate_block()
|
|
block.block_hash = b"stale-first"
|
|
block.token_count = 4
|
|
block.ref_count = 2
|
|
paged_cache.cached_block_hash_to_block.insert(block.block_hash, block)
|
|
block_table = BlockTable(
|
|
request_id="req-stale-first",
|
|
block_ids=[block.block_id],
|
|
num_tokens=4,
|
|
)
|
|
|
|
stale_block = [
|
|
(
|
|
mx.random.normal((1, 2, 4, 32)),
|
|
mx.random.normal((1, 2, 4, 32)),
|
|
)
|
|
]
|
|
stale_metadata = {
|
|
"model_name": "test-model",
|
|
"num_layers": 1,
|
|
"block_size": 4,
|
|
"layer_cache_types": ["KVCache"],
|
|
"layer_meta_states": [(4,)],
|
|
}
|
|
mock_ssd.load_block_with_metadata.return_value = (
|
|
stale_block,
|
|
stale_metadata,
|
|
)
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
|
|
assert result is None
|
|
assert block_table.block_ids == []
|
|
assert block_table.num_tokens == 0
|
|
mock_ssd.forget_block.assert_called_once_with(block.block_hash)
|
|
|
|
def test_reconstruct_accepts_sized_arrays_metadata_with_turboquant(self, mx):
|
|
"""SizedArraysCache is a restored ArraysCache wrapper, not a mismatch."""
|
|
from mlx_lm.models.cache import KVCache
|
|
from mlx_vlm.turboquant import TurboQuantKVCache
|
|
|
|
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager
|
|
from omlx.cache.type_handlers import SizedArraysCache
|
|
|
|
mock_ssd = MagicMock(spec=PagedSSDCacheManager)
|
|
|
|
model = MockModel(num_layers=2)
|
|
paged_cache = PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
blocks = []
|
|
for i in range(2):
|
|
block = paged_cache.allocate_block()
|
|
block.block_hash = f"hash-arrays-tq-{i}".encode()
|
|
block.token_count = 4
|
|
block.ref_count = 2
|
|
paged_cache.cached_block_hash_to_block.insert(block.block_hash, block)
|
|
blocks.append(block)
|
|
|
|
block_table = BlockTable(
|
|
request_id="req-arrays-sized-tq",
|
|
block_ids=[block.block_id for block in blocks],
|
|
num_tokens=8,
|
|
)
|
|
|
|
arrays_state_1 = (
|
|
mx.ones((1, 3, 16)),
|
|
mx.ones((1, 2, 16, 16)),
|
|
)
|
|
arrays_state_2 = (
|
|
mx.ones((1, 3, 16)) * 2,
|
|
mx.ones((1, 2, 16, 16)) * 2,
|
|
)
|
|
|
|
kv_cache = KVCache()
|
|
kv_cache.update_and_fetch(
|
|
mx.random.normal((1, 2, 4, 32)),
|
|
mx.random.normal((1, 2, 4, 32)),
|
|
)
|
|
tq_cache = TurboQuantKVCache.from_cache(kv_cache, bits=4.0)
|
|
tq_keys, tq_values = tq_cache.state
|
|
tq_block = ("__turboquant_v2__", (tq_keys, tq_values))
|
|
|
|
arrays_metadata = {
|
|
"model_name": "test-model",
|
|
"num_layers": 2,
|
|
"block_size": 4,
|
|
"layer_cache_types": ["ArraysCache", "TurboQuantKVCache"],
|
|
"layer_meta_states": [(), tq_cache.meta_state],
|
|
}
|
|
sized_metadata = {
|
|
"model_name": "test-model",
|
|
"num_layers": 2,
|
|
"block_size": 4,
|
|
"layer_cache_types": ["SizedArraysCache", "TurboQuantKVCache"],
|
|
"layer_meta_states": [(), tq_cache.meta_state],
|
|
}
|
|
mock_ssd.load_block_with_metadata.side_effect = [
|
|
([arrays_state_1, tq_block], arrays_metadata),
|
|
([arrays_state_2, tq_block], sized_metadata),
|
|
]
|
|
|
|
result = cache.reconstruct_cache(block_table)
|
|
|
|
assert result is not None
|
|
assert len(result) == 2
|
|
assert isinstance(result[0], SizedArraysCache)
|
|
assert isinstance(result[1], TurboQuantKVCache)
|
|
assert block_table.block_ids == [blocks[0].block_id, blocks[1].block_id]
|
|
assert block_table.num_tokens == 8
|
|
mock_ssd.forget_block.assert_not_called()
|
|
|
|
|
|
class TestPerBlockMetaStates:
|
|
"""Tests for per-block meta_states in store_cache with boundary snapshots.
|
|
|
|
Verifies that blocks stored with boundary snapshots use the snapshot's
|
|
meta_state (correct per-boundary offset) rather than the shared final
|
|
meta_state from _extract_cache_states.
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def mx(self):
|
|
"""Import MLX or skip."""
|
|
try:
|
|
import mlx.core as mx
|
|
|
|
return mx
|
|
except ImportError:
|
|
pytest.skip("MLX not available")
|
|
|
|
def test_store_cache_uses_snapshot_meta_for_rotating_cache(self, mx):
|
|
"""Boundary snapshot meta_state should override shared meta for RotatingKVCache blocks."""
|
|
from omlx.cache.hybrid_cache import ModelCacheConfig
|
|
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
mock_ssd.save_block.return_value = True
|
|
|
|
model = MockModel(num_layers=2)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
# 8 tokens = 2 full blocks of 4
|
|
tokens = list(range(8))
|
|
|
|
# Simulate a hybrid model: layer 0 = KVCache, layer 1 = RotatingKVCache
|
|
# Final cache state has offset=8 (end of request)
|
|
cache_data = [
|
|
{
|
|
"state": (mx.ones((1, 4, 8, 64)), mx.ones((1, 4, 8, 64))),
|
|
"cache_type": "KVCache",
|
|
"class_name": "KVCache",
|
|
"meta_state": ("8",),
|
|
},
|
|
{
|
|
"state": (mx.ones((1, 1, 4, 256)), mx.ones((1, 1, 4, 256))),
|
|
"cache_type": "RotatingKVCache",
|
|
"class_name": "RotatingKVCache",
|
|
"meta_state": (
|
|
"0",
|
|
"4",
|
|
"8",
|
|
"4",
|
|
), # keep, max_size, offset=8 (final), _idx
|
|
},
|
|
]
|
|
model_cache_config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "RotatingKVCache"], model_name="test-model"
|
|
)
|
|
|
|
# Boundary snapshot at token 4 (end of block 1) with correct offset=4
|
|
boundary_snapshots = {
|
|
4: [
|
|
{
|
|
"state": (),
|
|
"meta_state": (),
|
|
"class_name": "KVCache",
|
|
"cache_type": "KVCache",
|
|
},
|
|
{
|
|
"state": (mx.ones((1, 1, 4, 256)), mx.ones((1, 1, 4, 256))),
|
|
"meta_state": ("0", "4", "4", "4"), # offset=4 at boundary
|
|
"class_name": "RotatingKVCache",
|
|
"cache_type": "RotatingKVCache",
|
|
},
|
|
],
|
|
}
|
|
|
|
result = cache.store_cache(
|
|
"req-001",
|
|
tokens,
|
|
cache_data,
|
|
model_cache_config=model_cache_config,
|
|
boundary_snapshots=boundary_snapshots,
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result.block_ids) == 2
|
|
|
|
# Verify save_block was called twice (one per block)
|
|
assert mock_ssd.save_block.call_count == 2
|
|
|
|
# Block 1 (has boundary snapshot): should use snapshot meta for
|
|
# RotatingKVCache layer (offset=4), not shared meta (offset=8)
|
|
block1_call = mock_ssd.save_block.call_args_list[0]
|
|
block1_meta = block1_call.kwargs["layer_meta_states"]
|
|
# RotatingKVCache meta (layer 1): offset should be 4 from snapshot
|
|
assert block1_meta[1] == ("0", "4", "4", "4"), (
|
|
f"Block 1 RotatingKVCache meta should use snapshot offset=4, "
|
|
f"got {block1_meta[1]}"
|
|
)
|
|
|
|
# Block 2 (last block, uses main state): should use shared meta
|
|
block2_call = mock_ssd.save_block.call_args_list[1]
|
|
block2_meta = block2_call.kwargs["layer_meta_states"]
|
|
# Last block has no separate boundary snapshot override (boundary at
|
|
# token 8 matches the request end), so it uses the shared meta
|
|
assert block2_meta[1] == ("0", "4", "8", "4"), (
|
|
f"Block 2 RotatingKVCache meta should use shared meta offset=8, "
|
|
f"got {block2_meta[1]}"
|
|
)
|
|
|
|
def test_store_cache_kvcache_meta_falls_back_to_shared(self, mx):
|
|
"""KVCache layers in boundary snapshots have empty meta, should fall back to shared."""
|
|
from omlx.cache.hybrid_cache import ModelCacheConfig
|
|
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
mock_ssd.save_block.return_value = True
|
|
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
tokens = list(range(8))
|
|
cache_data = [
|
|
{
|
|
"state": (mx.ones((1, 4, 8, 64)), mx.ones((1, 4, 8, 64))),
|
|
"cache_type": "KVCache",
|
|
"class_name": "KVCache",
|
|
"meta_state": ("8",),
|
|
},
|
|
]
|
|
model_cache_config = ModelCacheConfig.from_type_list(
|
|
["KVCache"], model_name="test-model"
|
|
)
|
|
|
|
# KVCache layers have empty meta in boundary snapshots
|
|
boundary_snapshots = {
|
|
4: [
|
|
{
|
|
"state": (),
|
|
"meta_state": (),
|
|
"class_name": "KVCache",
|
|
"cache_type": "KVCache",
|
|
},
|
|
],
|
|
}
|
|
|
|
result = cache.store_cache(
|
|
"req-001",
|
|
tokens,
|
|
cache_data,
|
|
model_cache_config=model_cache_config,
|
|
boundary_snapshots=boundary_snapshots,
|
|
)
|
|
|
|
assert result is not None
|
|
assert mock_ssd.save_block.call_count == 2
|
|
|
|
# Block 1: KVCache meta should fall back to shared meta (empty snapshot meta)
|
|
block1_call = mock_ssd.save_block.call_args_list[0]
|
|
block1_meta = block1_call.kwargs["layer_meta_states"]
|
|
assert block1_meta[0] == (
|
|
"8",
|
|
), f"KVCache should fall back to shared meta, got {block1_meta[0]}"
|
|
|
|
def test_store_cache_no_snapshot_uses_shared_meta(self, mx):
|
|
"""Blocks without boundary snapshots should use shared meta (existing behavior)."""
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
mock_ssd.save_block.return_value = True
|
|
|
|
model = MockModel(num_layers=1)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
tokens = list(range(4))
|
|
cache_data = [
|
|
{
|
|
"state": (mx.ones((1, 4, 4, 64)), mx.ones((1, 4, 4, 64))),
|
|
"cache_type": "KVCache",
|
|
"class_name": "KVCache",
|
|
"meta_state": ("4",),
|
|
},
|
|
]
|
|
|
|
# No boundary snapshots
|
|
result = cache.store_cache("req-001", tokens, cache_data)
|
|
|
|
assert result is not None
|
|
assert mock_ssd.save_block.call_count == 1
|
|
|
|
block_call = mock_ssd.save_block.call_args_list[0]
|
|
block_meta = block_call.kwargs["layer_meta_states"]
|
|
assert block_meta[0] == ("4",)
|
|
|
|
def test_store_cache_last_block_with_snapshot_uses_snapshot_meta(self, mx):
|
|
"""Last block should also prefer snapshot meta when a boundary snapshot exists."""
|
|
from omlx.cache.hybrid_cache import ModelCacheConfig
|
|
|
|
block_size = 4
|
|
paged_cache = PagedCacheManager(
|
|
block_size=block_size,
|
|
max_blocks=100,
|
|
model_name="test-model",
|
|
initial_blocks=100,
|
|
)
|
|
mock_ssd = MagicMock()
|
|
mock_ssd.save_block.return_value = True
|
|
|
|
# Hybrid model: KVCache + RotatingKVCache (need KVCache so
|
|
# _get_cache_seq_len can determine the full sequence length)
|
|
model = MockModel(num_layers=2)
|
|
cache = BlockAwarePrefixCache(
|
|
model=model,
|
|
paged_cache_manager=paged_cache,
|
|
paged_ssd_cache_manager=mock_ssd,
|
|
)
|
|
|
|
# 8 tokens = 2 blocks, but cache_data reflects processing of
|
|
# 11 total tokens (3 trailing partial not stored).
|
|
# Shared meta has offset=11 (final state) for RotatingKVCache.
|
|
tokens = list(range(8))
|
|
cache_data = [
|
|
{
|
|
"state": (mx.ones((1, 4, 11, 64)), mx.ones((1, 4, 11, 64))),
|
|
"cache_type": "KVCache",
|
|
"class_name": "KVCache",
|
|
"meta_state": ("11",),
|
|
},
|
|
{
|
|
"state": (mx.ones((1, 1, 4, 256)), mx.ones((1, 1, 4, 256))),
|
|
"cache_type": "RotatingKVCache",
|
|
"class_name": "RotatingKVCache",
|
|
"meta_state": ("0", "4", "11", "4"), # offset=11 (final)
|
|
},
|
|
]
|
|
model_cache_config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "RotatingKVCache"], model_name="test-model"
|
|
)
|
|
|
|
# Both blocks have boundary snapshots with correct per-boundary offsets
|
|
boundary_snapshots = {
|
|
4: [
|
|
{
|
|
"state": (),
|
|
"meta_state": (),
|
|
"class_name": "KVCache",
|
|
"cache_type": "KVCache",
|
|
},
|
|
{
|
|
"state": (mx.ones((1, 1, 4, 256)), mx.ones((1, 1, 4, 256))),
|
|
"meta_state": ("0", "4", "4", "4"),
|
|
"class_name": "RotatingKVCache",
|
|
"cache_type": "RotatingKVCache",
|
|
},
|
|
],
|
|
8: [
|
|
{
|
|
"state": (),
|
|
"meta_state": (),
|
|
"class_name": "KVCache",
|
|
"cache_type": "KVCache",
|
|
},
|
|
{
|
|
"state": (mx.ones((1, 1, 4, 256)), mx.ones((1, 1, 4, 256))),
|
|
"meta_state": ("0", "4", "8", "4"), # offset=8 at boundary
|
|
"class_name": "RotatingKVCache",
|
|
"cache_type": "RotatingKVCache",
|
|
},
|
|
],
|
|
}
|
|
|
|
result = cache.store_cache(
|
|
"req-001",
|
|
tokens,
|
|
cache_data,
|
|
model_cache_config=model_cache_config,
|
|
boundary_snapshots=boundary_snapshots,
|
|
)
|
|
|
|
assert result is not None
|
|
assert mock_ssd.save_block.call_count == 2
|
|
|
|
# Block 1: RotatingKVCache offset=4 from snapshot
|
|
b1_meta = mock_ssd.save_block.call_args_list[0].kwargs["layer_meta_states"]
|
|
assert b1_meta[1] == ("0", "4", "4", "4")
|
|
|
|
# Block 2 (last): RotatingKVCache offset=8 from snapshot,
|
|
# NOT offset=11 from shared meta
|
|
b2_meta = mock_ssd.save_block.call_args_list[1].kwargs["layer_meta_states"]
|
|
assert b2_meta[1] == ("0", "4", "8", "4"), (
|
|
f"Last block should use snapshot offset=8, not shared offset=11, "
|
|
f"got {b2_meta[1]}"
|
|
)
|
|
|
|
|
|
class TestSetPagedSSDCacheManagerTriggersSweep:
|
|
"""``set_paged_ssd_cache_manager`` must trigger the one-shot sweep so
|
|
stale-signature blocks left over from a previous cache-config run for
|
|
this model are evicted before the first prefix lookup runs."""
|
|
|
|
@pytest.fixture
|
|
def prefix_cache(self):
|
|
return BlockAwarePrefixCache(
|
|
model=MockModel(num_layers=4),
|
|
paged_cache_manager=PagedCacheManager(
|
|
block_size=4,
|
|
max_blocks=16,
|
|
model_name="test",
|
|
initial_blocks=16,
|
|
),
|
|
paged_ssd_cache_manager=None,
|
|
)
|
|
|
|
def test_attach_calls_invalidate(self, prefix_cache):
|
|
mock_mgr = MagicMock()
|
|
mock_mgr.invalidate_stale_layer_signature.return_value = 0
|
|
|
|
prefix_cache.set_paged_ssd_cache_manager(mock_mgr)
|
|
|
|
mock_mgr.invalidate_stale_layer_signature.assert_called_once_with()
|
|
assert prefix_cache.paged_ssd_cache is mock_mgr
|
|
|
|
def test_attach_survives_sweep_exception(self, prefix_cache):
|
|
mock_mgr = MagicMock()
|
|
mock_mgr.invalidate_stale_layer_signature.side_effect = RuntimeError("boom")
|
|
|
|
# Must not raise — sweep failure is logged but the manager
|
|
# connection must still complete.
|
|
prefix_cache.set_paged_ssd_cache_manager(mock_mgr)
|
|
|
|
assert prefix_cache.paged_ssd_cache is mock_mgr
|
|
|
|
def test_attach_none_no_call(self, prefix_cache):
|
|
# Detaching the manager must not invoke anything.
|
|
prefix_cache.set_paged_ssd_cache_manager(None)
|
|
assert prefix_cache.paged_ssd_cache is None
|
|
|
|
|
|
class TestCanonicalLayerCacheTypes:
|
|
"""The canonicalizer normalizes wrapper class names but must NOT
|
|
collapse types that change tensor representation (TurboQuantKVCache
|
|
stores 4-bit packed tensors; KVCache stores fp16) — collapsing those
|
|
would silently mix incompatible cache blocks."""
|
|
|
|
def test_none_passthrough(self):
|
|
assert BlockAwarePrefixCache._canonical_layer_cache_types(None) is None
|
|
|
|
def test_sized_arrays_normalized(self):
|
|
result = BlockAwarePrefixCache._canonical_layer_cache_types(
|
|
["SizedArraysCache", "SizedArraysCache", "KVCache"]
|
|
)
|
|
assert result == ["ArraysCache", "ArraysCache", "KVCache"]
|
|
|
|
def test_prefill_ready_rotating_normalized(self):
|
|
result = BlockAwarePrefixCache._canonical_layer_cache_types(
|
|
["KVCache", "PrefillReadyRotatingKVCache", "RotatingKVCache"]
|
|
)
|
|
assert result == ["KVCache", "RotatingKVCache", "RotatingKVCache"]
|
|
|
|
def test_turboquant_not_collapsed(self):
|
|
result = BlockAwarePrefixCache._canonical_layer_cache_types(
|
|
["ArraysCache", "TurboQuantKVCache", "KVCache"]
|
|
)
|
|
# TurboQuant must remain distinct from plain KVCache.
|
|
assert "TurboQuantKVCache" in result
|
|
assert result != ["ArraysCache", "KVCache", "KVCache"]
|