Engram-protocol/engram

1"""2ENGRAM Protocol — State Extractor Tests3Tests for all 3 EGR extraction modes (D3).4"""5 6from __future__ import annotations7 8import torch9 10from kvcos.core.cache_spec import LLAMA_3_1_8B, PHI_3_MINI11from kvcos.core.types import StateExtractionMode12from kvcos.core.state_extractor import MARStateExtractor13from tests.conftest import make_synthetic_kv14 15 16class TestMeanPool:17    """mean_pool: mean over layers, heads, context → [head_dim]."""18 19    def test_output_dim(self) -> None:20        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)21        ext = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)22        result = ext.extract(keys, LLAMA_3_1_8B)23        assert result.state_vec.shape == (128,)24 25    def test_output_dim_api(self) -> None:26        ext = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)27        assert ext.output_dim(LLAMA_3_1_8B) == 12828 29    def test_l2_norm_positive(self) -> None:30        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)31        ext = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)32        result = ext.extract(keys, LLAMA_3_1_8B)33        assert result.l2_norm > 034 35    def test_deterministic(self) -> None:36        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)37        ext = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)38        r1 = ext.extract(keys, LLAMA_3_1_8B)39        r2 = ext.extract(keys, LLAMA_3_1_8B)40        assert torch.equal(r1.state_vec, r2.state_vec)41 42 43class TestSVDProject:44    """svd_project: truncated SVD, rank-160 → [rank]."""45 46    def test_output_dim(self) -> None:47        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)48        ext = MARStateExtractor(mode=StateExtractionMode.SVD_PROJECT, rank=160)49        result = ext.extract(keys, LLAMA_3_1_8B)50        assert result.state_vec.shape == (128,)  # clamped to head_dim51 52    def test_projection_stored(self) -> None:53        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)54        ext = MARStateExtractor(mode=StateExtractionMode.SVD_PROJECT, rank=160)55        ext.extract(keys, LLAMA_3_1_8B)56        proj = ext.last_projection57        assert proj is not None58        assert 0.0 < proj.explained_variance_ratio <= 1.059 60    def test_n_layers_used(self) -> None:61        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)62        ext = MARStateExtractor(mode=StateExtractionMode.SVD_PROJECT)63        result = ext.extract(keys, LLAMA_3_1_8B)64        assert result.n_layers_used == 24  # layers 8-3165 66 67class TestXKVProject:68    """xkv_project: grouped cross-layer SVD."""69 70    def test_output_dim(self) -> None:71        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)72        ext = MARStateExtractor(mode=StateExtractionMode.XKV_PROJECT, rank=160)73        result = ext.extract(keys, LLAMA_3_1_8B)74        expected_dim = ext.output_dim(LLAMA_3_1_8B)75        assert result.state_vec.shape == (expected_dim,)76 77    def test_different_from_mean_pool(self) -> None:78        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)79        ext_mp = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)80        ext_xkv = MARStateExtractor(mode=StateExtractionMode.XKV_PROJECT)81        r_mp = ext_mp.extract(keys, LLAMA_3_1_8B)82        r_xkv = ext_xkv.extract(keys, LLAMA_3_1_8B)83        assert r_mp.state_vec.shape != r_xkv.state_vec.shape84 85    def test_phi3_works(self) -> None:86        keys, _ = make_synthetic_kv(PHI_3_MINI, ctx_len=64)87        ext = MARStateExtractor(mode=StateExtractionMode.XKV_PROJECT, rank=96)88        result = ext.extract(keys, PHI_3_MINI)89        assert result.state_vec.dim() == 190        assert result.state_vec.shape[0] > 091