skhavin/proactive-cache

1"""Unit tests for proactive_cache."""2 3import pytest4import numpy as np5import torch6 7from proactive_cache.eviction import score_tokens, select_indices, prune_kv_cache, evict8from proactive_cache.prototypes import build_prototypes, save_prototypes, load_prototypes9from proactive_cache.utils import to_tuple_kv, to_dynamic_cache10 11 12# ── Fixtures ──────────────────────────────────────────────────────────────────13 14def make_dummy_patterns(num_docs=5, num_layers=2, num_heads=4, seq_len=64):15    """Create synthetic attention patterns for testing."""16    patterns = []17    for _ in range(num_docs):18        doc = {}19        for layer in range(num_layers):20            for head in range(num_heads):21                arr = np.random.rand(seq_len).astype(np.float32)22                arr /= arr.sum()23                doc[(layer, head)] = arr24        patterns.append(doc)25    return patterns26 27 28def make_dummy_kv_cache(num_layers=2, num_heads=4, seq_len=64, head_dim=32, device="cpu"):29    """Create a synthetic KV cache tuple."""30    return tuple(31        (torch.randn(1, num_heads, seq_len, head_dim, device=device),32         torch.randn(1, num_heads, seq_len, head_dim, device=device))33        for _ in range(num_layers)34    )35 36 37# ── Eviction tests ────────────────────────────────────────────────────────────38 39class TestScoreTokens:40    def test_returns_correct_shape(self):41        scores = score_tokens(None, seq_len=128, budget=64)42        assert scores.shape == (128,)43 44    def test_token_zero_has_highest_score(self):45        scores = score_tokens(None, seq_len=128, budget=64)46        # Sink boost means position 0 is always kept47        top_k = np.argsort(scores)[-64:]48        assert 0 in top_k, "Token 0 (attention sink) must always be selected"49 50    def test_recency_tokens_kept(self):51        seq_len, budget = 128, 6452        scores = score_tokens(None, seq_len=seq_len, budget=budget)53        top_k = np.argsort(scores)[-budget:]54        # Last few tokens should be in top-k55        assert (seq_len - 1) in top_k, "Most recent token must always be kept"56 57    def test_with_prototypes(self):58        patterns = make_dummy_patterns(seq_len=64)59        protos = build_prototypes(patterns, n_clusters=2, max_seq_len=64)60        scores = score_tokens(protos, seq_len=64, budget=32)61        assert scores.shape == (64,)62        assert np.all(np.isfinite(scores)), "Scores must be finite"63 64    def test_budget_proportional_recency(self):65        # Larger budget → larger recency window (proportional)66        s128 = score_tokens(None, seq_len=512, budget=128)67        s256 = score_tokens(None, seq_len=512, budget=256)68        # More positions should be elevated in s25669        # (just check both run without error)70        assert s128.shape == s256.shape == (512,)71 72 73class TestSelectIndices:74    def test_returns_sorted(self):75        scores = np.random.rand(100)76        idx = select_indices(scores, budget=20)77        assert idx == sorted(idx), "Indices must be in ascending order"78 79    def test_correct_count(self):80        scores = np.random.rand(100)81        idx = select_indices(scores, budget=30)82        assert len(idx) == 3083 84    def test_budget_larger_than_seq(self):85        scores = np.random.rand(10)86        idx = select_indices(scores, budget=50)87        assert len(idx) == 10  # clipped to seq_len88 89 90class TestPruneKVCache:91    def test_prunes_to_budget(self):92        kv = make_dummy_kv_cache(num_layers=3, num_heads=4, seq_len=128)93        indices = list(range(0, 64, 2))  # 32 indices94        pruned = prune_kv_cache(kv, indices, device=torch.device("cpu"))95        pruned_tuple = to_tuple_kv(pruned)96        assert pruned_tuple[0][0].shape[2] == 32, "Pruned KV must have budget tokens"97 98    def test_all_layers_pruned(self):99        num_layers = 4100        kv = make_dummy_kv_cache(num_layers=num_layers, seq_len=100)101        indices = list(range(50))102        pruned_tuple = to_tuple_kv(prune_kv_cache(kv, indices, torch.device("cpu")))103        assert len(pruned_tuple) == num_layers104 105    def test_no_prune_when_under_budget(self):106        kv = make_dummy_kv_cache(seq_len=32)107        result = evict(kv, budget=64, prototypes=None, seq_len=32, device=torch.device("cpu"))108        # Should return unchanged (seq_len <= budget)109        assert to_tuple_kv(result)[0][0].shape[2] == 32110 111 112# ── Prototype tests ───────────────────────────────────────────────────────────113 114class TestPrototypes:115    def test_build_returns_dict(self):116        patterns = make_dummy_patterns()117        protos = build_prototypes(patterns, n_clusters=2, max_seq_len=64)118        assert isinstance(protos, dict)119        assert len(protos) > 0120 121    def test_centroid_shapes(self):122        patterns = make_dummy_patterns(num_layers=2, num_heads=4, seq_len=64)123        protos = build_prototypes(patterns, n_clusters=3, max_seq_len=64)124        for key, val in protos.items():125            centroids = val["centroids"]126            assert centroids.shape == (3, 64), f"Wrong centroid shape: {centroids.shape}"127 128    def test_save_load_roundtrip(self, tmp_path):129        patterns = make_dummy_patterns()130        protos = build_prototypes(patterns, n_clusters=2, max_seq_len=64)131        path = str(tmp_path / "test_protos.pkl")132        save_prototypes(protos, path)133        loaded = load_prototypes(path)134        assert set(loaded.keys()) == set(protos.keys())135 136    def test_load_missing_raises(self, tmp_path):137        with pytest.raises(FileNotFoundError):138            load_prototypes(str(tmp_path / "does_not_exist.pkl"))139 140    def test_empty_patterns_raises(self):141        with pytest.raises(ValueError):142            build_prototypes([], n_clusters=2)143 144 145# ── Utils tests ───────────────────────────────────────────────────────────────146 147class TestUtils:148    def test_to_tuple_kv_from_tuple(self):149        kv = make_dummy_kv_cache(num_layers=2)150        result = to_tuple_kv(kv)151        assert len(result) == 2152        assert isinstance(result[0], tuple)153 154    def test_to_dynamic_cache_roundtrip(self):155        kv = make_dummy_kv_cache(num_layers=2, seq_len=32)156        kv_tuple = to_tuple_kv(kv)157        dynamic = to_dynamic_cache(kv_tuple)158        back = to_tuple_kv(dynamic)159        # Shapes should be preserved160        assert back[0][0].shape == kv_tuple[0][0].shape161