test: upload 220 tests

tests/conftest.py

@@ -0,0 +1,201 @@
+"""
+ENGRAM Protocol — Test Fixtures
+
+
+Shared pytest fixtures for all test modules.
+Provides synthetic KV cache tensors at correct shapes,
+temp directories, and model specs.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+
+from kvcos.core.cache_spec import GEMMA_4_26B_A4B, LLAMA_3_1_8B, PHI_3_MINI
+from kvcos.core.types import AttentionType, CacheSection, ModelCacheSpec
+
+
+@pytest.fixture
+def llama_spec() -> ModelCacheSpec:
+    """Llama 3.1 8B model spec."""
+    return LLAMA_3_1_8B
+
+
+@pytest.fixture
+def phi3_spec() -> ModelCacheSpec:
+    """Phi-3-Mini model spec."""
+    return PHI_3_MINI
+
+
+@pytest.fixture
+def gemma4_spec() -> ModelCacheSpec:
+    """Gemma 4 26B-A4B ISWA model spec."""
+    return GEMMA_4_26B_A4B
+
+
+@pytest.fixture
+def tmp_data_dir(tmp_path: Path) -> Path:
+    """Temporary data directory for storage tests."""
+    data_dir = tmp_path / "engram_data"
+    data_dir.mkdir()
+    return data_dir
+
+
+@pytest.fixture
+def tmp_index_dir(tmp_path: Path) -> Path:
+    """Temporary directory for FAISS index persistence tests."""
+    index_dir = tmp_path / "engram_index"
+    index_dir.mkdir()
+    return index_dir
+
+
+def make_synthetic_kv(
+    spec: ModelCacheSpec,
+    ctx_len: int = 256,
+    seed: int = 42,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Create synthetic KV cache tensors with correct shapes.
+
+    Returns (keys, values) each [n_layers, n_kv_heads, ctx_len, head_dim].
+    Values are random but reproducible via seed.
+    """
+    torch.manual_seed(seed)
+    shape = (spec["n_layers"], spec["n_kv_heads"], ctx_len, spec["head_dim"])
+    keys = torch.randn(shape, dtype=torch.float16)
+    values = torch.randn(shape, dtype=torch.float16)
+    return keys, values
+
+
+@pytest.fixture
+def llama_kv_256(llama_spec: ModelCacheSpec) -> tuple[torch.Tensor, torch.Tensor]:
+    """Synthetic Llama 3.1 8B KV cache, 256 tokens.
+
+    Shape: [32, 8, 256, 128] for both keys and values.
+    """
+    return make_synthetic_kv(llama_spec, ctx_len=256)
+
+
+@pytest.fixture
+def llama_kv_1024(llama_spec: ModelCacheSpec) -> tuple[torch.Tensor, torch.Tensor]:
+    """Synthetic Llama 3.1 8B KV cache, 1024 tokens."""
+    return make_synthetic_kv(llama_spec, ctx_len=1024, seed=123)
+
+
+@pytest.fixture
+def phi3_kv_256(phi3_spec: ModelCacheSpec) -> tuple[torch.Tensor, torch.Tensor]:
+    """Synthetic Phi-3-Mini KV cache, 256 tokens.
+
+    Shape: [32, 32, 256, 96] for both keys and values.
+    """
+    return make_synthetic_kv(phi3_spec, ctx_len=256, seed=99)
+
+
+# ── ISWA Fixtures ────────────────────────────────────────────────────────────
+
+
+def make_synthetic_iswa_blob(
+    sections: tuple[CacheSection, ...],
+    n_cells: int = 4,
+    arch: str = "gemma4",
+    v_trans: bool = True,
+    seed: int = 42,
+) -> bytes:
+    """Build a synthetic ISWA blob with multiple KV cache sections.
+
+    Matches llama.cpp state blob format for ISWA models:
+      1. Architecture string header
+      2. n_stream = len(sections)
+      3. Per stream: cell metadata + K/V data per layer
+
+    Args:
+        sections: Cache sections (e.g., global + SWA for Gemma 4).
+        n_cells: Number of KV cells per section.
+        arch: Architecture string in blob header.
+        v_trans: Whether V tensors are stored transposed.
+        seed: Random seed for reproducible data.
+    """
+    import struct
+
+    import numpy as np
+
+    from kvcos.core.blob_parser import GGML_TYPE_F16
+
+    rng = np.random.RandomState(seed)
+    parts: list[bytes] = []
+
+    # 1. Architecture string header
+    parts.append(struct.pack("<I", len(arch)))
+    parts.append(arch.encode("ascii"))
+
+    # 2. Stream count = number of cache sections
+    parts.append(struct.pack("<I", len(sections)))
+
+    # 3. Per-stream data
+    for section in sections:
+        n_embd_kv = section.n_kv_heads * section.head_dim
+        row_size = n_embd_kv * 2  # fp16
+
+        # Cell metadata
+        parts.append(struct.pack("<I", n_cells))
+        for i in range(n_cells):
+            parts.append(struct.pack("<i", i))    # pos
+            parts.append(struct.pack("<I", 1))    # n_seq_id = 1
+            parts.append(struct.pack("<i", 0))    # seq_id = 0
+
+        # Data section header
+        parts.append(struct.pack("<I", 1 if v_trans else 0))
+        parts.append(struct.pack("<I", section.n_layers))
+
+        # K layers
+        for _ in range(section.n_layers):
+            parts.append(struct.pack("<i", GGML_TYPE_F16))
+            parts.append(struct.pack("<Q", row_size))
+            data = rng.randn(n_cells * n_embd_kv).astype(np.float16)
+            parts.append(data.tobytes())
+
+        # V layers
+        for _ in range(section.n_layers):
+            parts.append(struct.pack("<i", GGML_TYPE_F16))
+            if v_trans:
+                parts.append(struct.pack("<I", 2))         # el_size (fp16)
+                parts.append(struct.pack("<I", n_embd_kv)) # n_embd_v_gqa
+            else:
+                parts.append(struct.pack("<Q", row_size))
+            data = rng.randn(n_cells * n_embd_kv).astype(np.float16)
+            parts.append(data.tobytes())
+
+    return b"".join(parts)
+
+
+# Gemma 4 ISWA section constants (reverse-engineered)
+GEMMA4_GLOBAL_SECTION = CacheSection(
+    attention_type=AttentionType.FULL,
+    n_layers=5,
+    n_kv_heads=2,
+    head_dim=512,
+)
+
+GEMMA4_SWA_SECTION = CacheSection(
+    attention_type=AttentionType.SLIDING,
+    n_layers=25,
+    n_kv_heads=8,
+    head_dim=256,
+    window_size=1024,
+)
+
+GEMMA4_SECTIONS = (GEMMA4_GLOBAL_SECTION, GEMMA4_SWA_SECTION)
+
+
+@pytest.fixture
+def gemma4_iswa_blob() -> bytes:
+    """Synthetic Gemma 4 ISWA blob with 2 sections, 4 cells."""
+    return make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=4)
+
+
+@pytest.fixture
+def gemma4_iswa_blob_8cells() -> bytes:
+    """Synthetic Gemma 4 ISWA blob with 2 sections, 8 cells."""
+    return make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=8, seed=99)

tests/test_blob_parser.py

@@ -0,0 +1,125 @@
+"""
+ENGRAM Protocol — Blob Parser Tests
+Tests for llama.cpp state blob → structured tensors (D1).
+Uses synthetic blobs matching the real llama_state_get_data() format.
+"""
+
+from __future__ import annotations
+
+import struct
+
+import numpy as np
+import pytest
+import torch
+
+from kvcos.core.blob_parser import (
+    GGML_TYPE_F16,
+    BlobParseError,
+    ParsedKVCache,
+    parse_state_blob,
+)
+
+
+def _make_blob(
+    n_cells: int,
+    n_layers: int,
+    n_kv_heads: int,
+    head_dim: int,
+    arch: str = "llama",
+    v_trans: bool = True,
+) -> bytes:
+    """Build a synthetic blob matching llama_state_get_data() format."""
+    parts: list[bytes] = []
+
+    # 1. Architecture string header
+    parts.append(struct.pack("<I", len(arch)))
+    parts.append(arch.encode("ascii"))
+
+    # 2. KV stream header
+    parts.append(struct.pack("<I", 1))  # n_stream = 1
+    parts.append(struct.pack("<I", n_cells))  # cell_count
+
+    # 3. Cell metadata: (pos:i32, n_seq:u32, seq_id:i32) per cell
+    for i in range(n_cells):
+        parts.append(struct.pack("<i", i))  # pos
+        parts.append(struct.pack("<I", 1))  # n_seq_id = 1
+        parts.append(struct.pack("<i", 0))  # seq_id = 0
+
+    # 4. Data section header
+    parts.append(struct.pack("<I", 1 if v_trans else 0))  # v_trans
+    parts.append(struct.pack("<I", n_layers))
+
+    n_embd_kv = n_kv_heads * head_dim
+    row_size = n_embd_kv * 2  # fp16
+
+    # 5. K layers
+    for _ in range(n_layers):
+        parts.append(struct.pack("<i", GGML_TYPE_F16))  # type_k
+        parts.append(struct.pack("<Q", row_size))  # row_size_k
+        data = np.random.randn(n_cells * n_embd_kv).astype(np.float16)
+        parts.append(data.tobytes())
+
+    # 6. V layers
+    for _ in range(n_layers):
+        parts.append(struct.pack("<i", GGML_TYPE_F16))  # type_v
+        if v_trans:
+            parts.append(struct.pack("<I", 2))  # el_size (fp16)
+            parts.append(struct.pack("<I", n_embd_kv))  # n_embd_v_gqa
+        else:
+            parts.append(struct.pack("<Q", row_size))  # row_size_v
+        data = np.random.randn(n_cells * n_embd_kv).astype(np.float16)
+        parts.append(data.tobytes())
+
+    return b"".join(parts)
+
+
+class TestBlobParser:
+    """Parse synthetic blobs in real llama_state_get_data format."""
+
+    def test_parse_shape(self) -> None:
+        blob = _make_blob(16, 32, 8, 128)
+        result = parse_state_blob(blob, n_kv_heads=8, head_dim=128)
+        assert result.keys.shape == (32, 8, 16, 128)
+        assert result.values.shape == (32, 8, 16, 128)
+
+    def test_parse_metadata(self) -> None:
+        blob = _make_blob(8, 32, 8, 128)
+        result = parse_state_blob(blob, n_kv_heads=8, head_dim=128)
+        assert result.n_cells == 8
+        assert result.n_layers == 32
+        assert result.arch == "llama"
+        assert result.v_trans is True
+        assert len(result.cells) == 8
+        assert result.cells[0].pos == 0
+        assert result.cells[7].pos == 7
+
+    def test_dtype_float16(self) -> None:
+        blob = _make_blob(4, 28, 8, 128)
+        result = parse_state_blob(blob, n_kv_heads=8, head_dim=128)
+        assert result.keys.dtype == torch.float16
+        assert result.values.dtype == torch.float16
+
+    def test_non_transposed_v(self) -> None:
+        blob = _make_blob(4, 28, 8, 128, v_trans=False)
+        result = parse_state_blob(blob, n_kv_heads=8, head_dim=128)
+        assert result.values.shape == (28, 8, 4, 128)
+        assert result.v_trans is False
+
+
+class TestBlobParserErrors:
+    """Edge cases."""
+
+    def test_zero_cells_raises(self) -> None:
+        blob = struct.pack("<I", 5) + b"llama" + struct.pack("<II", 1, 0) + b"\x00" * 20
+        with pytest.raises(BlobParseError, match="0 cells"):
+            parse_state_blob(blob, n_kv_heads=8, head_dim=128)
+
+    def test_truncated_blob_raises(self) -> None:
+        blob = _make_blob(4, 28, 8, 128)
+        with pytest.raises(BlobParseError):
+            parse_state_blob(blob[:100], n_kv_heads=8, head_dim=128)
+
+    def test_bad_arch_length_raises(self) -> None:
+        blob = struct.pack("<I", 999) + b"x" * 100
+        with pytest.raises(BlobParseError, match="too large"):
+            parse_state_blob(blob, n_kv_heads=8, head_dim=128)

tests/test_block_pool.py

@@ -0,0 +1,92 @@
+"""
+ENGRAM Protocol — Block Pool Tests
+Tests for 256-token block segmentation/assembly/extend.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from kvcos.core.block_pool import BlockPool, KVBlock
+from kvcos.core.types import BLOCK_SIZE_TOKENS
+
+
+def _kv(n_layers: int, n_heads: int, ctx: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]:
+    k = torch.randn(n_layers, n_heads, ctx, dim, dtype=torch.float16)
+    return k, k.clone()
+
+
+class TestSegment:
+    """Segment full KV cache into 256-token blocks."""
+
+    def test_exact_blocks(self) -> None:
+        keys, vals = _kv(32, 8, 512, 128)
+        pool = BlockPool(agent_id="a", model_id="m")
+        blocks = pool.segment(keys, vals)
+        assert len(blocks) == 2
+        assert all(b.is_full for b in blocks)
+
+    def test_partial_last_block(self) -> None:
+        keys, vals = _kv(32, 8, 300, 128)
+        pool = BlockPool(agent_id="a", model_id="m")
+        blocks = pool.segment(keys, vals)
+        assert len(blocks) == 2
+        assert blocks[0].is_full
+        assert not blocks[1].is_full
+        assert blocks[1].block_len == 44
+
+    def test_total_tokens(self) -> None:
+        keys, vals = _kv(32, 8, 700, 128)
+        pool = BlockPool(agent_id="a", model_id="m")
+        pool.segment(keys, vals)
+        assert pool.total_tokens == 700
+
+
+class TestAssemble:
+    """Assemble blocks back into full KV cache."""
+
+    def test_round_trip(self) -> None:
+        keys, vals = _kv(4, 2, 512, 64)
+        pool = BlockPool(agent_id="a", model_id="m")
+        pool.segment(keys, vals)
+        k_out, v_out = pool.assemble()
+        assert torch.equal(k_out, keys)
+
+    def test_subset_assembly(self) -> None:
+        keys, vals = _kv(4, 2, 768, 64)
+        pool = BlockPool(agent_id="a", model_id="m")
+        pool.segment(keys, vals)
+        k_out, _ = pool.assemble(block_indices=[0, 2])
+        assert k_out.shape[2] == BLOCK_SIZE_TOKENS * 2
+
+    def test_empty_raises(self) -> None:
+        pool = BlockPool(agent_id="a", model_id="m")
+        with pytest.raises(ValueError, match="No blocks"):
+            pool.assemble()
+
+
+class TestExtend:
+    """Extend pool with new tokens."""
+
+    def test_fills_partial_block(self) -> None:
+        keys, vals = _kv(4, 2, 200, 64)
+        pool = BlockPool(agent_id="a", model_id="m")
+        pool.segment(keys, vals)
+        assert not pool.blocks[-1].is_full
+
+        new_k, new_v = _kv(4, 2, 56, 64)
+        pool.extend(new_k, new_v)
+        assert pool.blocks[-1].is_full
+        assert pool.total_tokens == 256
+
+    def test_extend_creates_new_blocks(self) -> None:
+        keys, vals = _kv(4, 2, 256, 64)
+        pool = BlockPool(agent_id="a", model_id="m")
+        pool.segment(keys, vals)
+        assert pool.n_blocks == 1
+
+        new_k, new_v = _kv(4, 2, 300, 64)
+        pool.extend(new_k, new_v)
+        assert pool.n_blocks == 3
+        assert pool.total_tokens == 556

tests/test_chunker.py

@@ -0,0 +1,99 @@
+"""Tests for kvcos.engram.chunker — markdown-aware semantic chunker."""
+
+import pytest
+
+from kvcos.engram.chunker import Chunk, chunk_markdown, eng_filename, slug_from_path
+
+
+class TestChunkMarkdown:
+    def test_empty_content(self):
+        assert chunk_markdown("") == []
+        assert chunk_markdown("   ") == []
+
+    def test_small_file_single_chunk(self):
+        content = "# Title\n\nSome short content."
+        chunks = chunk_markdown(content, max_chars=2000)
+        assert len(chunks) == 1
+        assert chunks[0].index == 0
+        assert chunks[0].char_start == 0
+        assert chunks[0].char_end == len(content)
+
+    def test_large_file_splits(self):
+        # Create content that exceeds max_chars
+        content = "# Section 1\n\n" + "A" * 1500 + "\n\n# Section 2\n\n" + "B" * 1500
+        chunks = chunk_markdown(content, max_chars=2000)
+        assert len(chunks) >= 2
+
+    def test_chunks_cover_full_content(self):
+        content = "# A\n\nText A.\n\n# B\n\nText B.\n\n# C\n\nText C."
+        chunks = chunk_markdown(content, max_chars=15)
+        # All original content should be present across chunks
+        combined = " ".join(c.raw_text for c in chunks)
+        for word in ["Text A", "Text B", "Text C"]:
+            assert word in combined
+
+    def test_context_prefix(self):
+        content = "Hello world"
+        chunks = chunk_markdown(content, context_prefix="Source: test.md")
+        assert len(chunks) == 1
+        assert chunks[0].text.startswith("Source: test.md")
+
+    def test_indices_sequential(self):
+        content = "# A\n\n" + "X" * 3000 + "\n\n# B\n\n" + "Y" * 3000
+        chunks = chunk_markdown(content, max_chars=2000)
+        for i, chunk in enumerate(chunks):
+            assert chunk.index == i
+
+    def test_merge_small_sections(self):
+        """Small consecutive sections should merge into one chunk."""
+        content = "# A\n\nShort.\n\n# B\n\nAlso short.\n\n# C\n\nStill short."
+        chunks = chunk_markdown(content, max_chars=2000, min_chars=100)
+        # All three small sections should merge into 1 chunk
+        assert len(chunks) == 1
+
+    def test_paragraph_split_fallback(self):
+        """Content without headers should split on paragraphs."""
+        paragraphs = ["Paragraph " + str(i) + ". " + "X" * 500
+                       for i in range(6)]
+        content = "\n\n".join(paragraphs)
+        chunks = chunk_markdown(content, max_chars=1500)
+        assert len(chunks) >= 2
+
+
+class TestSlugFromPath:
+    def test_simple_filename(self):
+        assert slug_from_path("readme.md") == "readme"
+
+    def test_uppercase_underscores(self):
+        assert slug_from_path("EIGENGRAM_SPEC.md") == "eigengram-spec"
+
+    def test_already_kebab(self):
+        assert slug_from_path("coding-style.md") == "coding-style"
+
+    def test_full_path(self):
+        assert slug_from_path("/Users/test/docs/my_doc.md") == "my-doc"
+
+    def test_special_chars(self):
+        assert slug_from_path("file (copy).md") == "file-copy"
+
+
+class TestEngFilename:
+    def test_single_chunk(self):
+        name = eng_filename("engram", "readme", "2026-04-02")
+        assert name == "readme_2026-04-02.eng"
+
+    def test_multi_chunk(self):
+        name = eng_filename("engram", "geodesic3", "2026-04-02",
+                           chunk_index=0, chunk_total=5)
+        assert name == "geodesic3_001_2026-04-02.eng"
+
+    def test_with_time(self):
+        name = eng_filename("engram", "session", "2026-04-02",
+                           time_str="1430")
+        assert name == "session_2026-04-02_1430.eng"
+
+    def test_single_chunk_no_index(self):
+        """Single-chunk files should not have chunk number."""
+        name = eng_filename("engram", "small", "2026-04-02",
+                           chunk_index=0, chunk_total=1)
+        assert name == "small_2026-04-02.eng"

tests/test_compression.py

@@ -0,0 +1,356 @@
+"""
+ENGRAM Protocol — Compression Tests
+
+
+Tests for kvcos.core.compression:
+  - FP16 passthrough
+  - Q8_0 round-trip accuracy & shape preservation
+  - PolarQuant round-trip accuracy & rotation invariants
+  - Dispatcher routing and Q4_0 fallback warning
+  - Edge cases: padding, single-element groups
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import pytest
+import torch
+
+from kvcos.core.compression import (
+    Q8_GROUP_SIZE,
+    CompressionResult,
+    compress,
+    compress_fp16,
+    compress_polarquant,
+    compress_q8_0,
+    decompress,
+    decompress_fp16,
+    decompress_polarquant,
+    decompress_q8_0,
+)
+from kvcos.core.types import CompressionMethod
+
+
+# ── FP16 Passthrough ──────────────────────────────────────────────────────────
+
+
+class TestFP16:
+    """FP16 passthrough: no quantization, just dtype normalization."""
+
+    def test_fp16_passthrough_shape(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_fp16(keys)
+        assert result.data.shape == keys.shape
+
+    def test_fp16_passthrough_dtype(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_fp16(keys)
+        assert result.data.dtype == torch.float16
+
+    def test_fp16_passthrough_exact(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_fp16(keys)
+        assert torch.equal(result.data, keys.to(torch.float16))
+
+    def test_fp16_compression_ratio_one(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_fp16(keys)
+        assert result.compression_ratio == 1.0
+
+    def test_fp16_method_tag(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_fp16(keys)
+        assert result.method == CompressionMethod.FP16
+
+    def test_fp16_from_fp32(self) -> None:
+        """FP32 input is cast to FP16."""
+        t = torch.randn(4, 8, 32, 128, dtype=torch.float32)
+        result = compress_fp16(t)
+        assert result.data.dtype == torch.float16
+        assert result.original_dtype == torch.float32
+
+    def test_fp16_decompress_identity(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_fp16(keys)
+        out = decompress_fp16(result.data)
+        assert torch.equal(out, result.data)
+
+
+# ── Q8_0 Quantization ────────────────────────────────────────────────────────
+
+
+class TestQ8_0:
+    """Q8_0: group quantization matching llama.cpp GGML_TYPE_Q8_0."""
+
+    def test_q8_0_shape_preserved(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_q8_0(keys)
+        assert result.data.shape == keys.shape
+
+    def test_q8_0_output_dtype(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """Q8_0 stores dequantized bfloat16 for safetensors compat."""
+        keys, _ = llama_kv_256
+        result = compress_q8_0(keys)
+        assert result.data.dtype == torch.bfloat16
+
+    def test_q8_0_method_tag(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_q8_0(keys)
+        assert result.method == CompressionMethod.Q8_0
+
+    def test_q8_0_metadata_group_size(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_q8_0(keys)
+        assert result.metadata["q8_group_size"] == str(Q8_GROUP_SIZE)
+
+    def test_q8_0_round_trip_low_error(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """Q8_0 quantization error should be < 1% relative MSE."""
+        keys, _ = llama_kv_256
+        result = compress_q8_0(keys)
+        decompressed = decompress_q8_0(result.data)
+
+        original = keys.float()
+        restored = decompressed.float()
+
+        mse = ((original - restored) ** 2).mean()
+        signal_power = (original**2).mean()
+        relative_mse = (mse / signal_power).item()
+        assert relative_mse < 0.01, f"Q8_0 relative MSE {relative_mse:.6f} > 1%"
+
+    def test_q8_0_round_trip_values(
+        self, phi3_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """Q8_0 round-trip on Phi-3 (head_dim=96, needs padding)."""
+        keys, values = phi3_kv_256
+        for tensor in (keys, values):
+            result = compress_q8_0(tensor)
+            assert result.data.shape == tensor.shape
+
+    def test_q8_0_compression_ratio_fp32(self) -> None:
+        """FP32 input → bfloat16 output gives 2x compression ratio."""
+        t = torch.randn(2, 4, 64, 128, dtype=torch.float32)
+        result = compress_q8_0(t)
+        assert abs(result.compression_ratio - 2.0) < 0.01
+
+    def test_q8_0_compression_ratio_fp16(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """FP16 input → bfloat16 output gives 1x ratio (same byte width)."""
+        keys, _ = llama_kv_256
+        result = compress_q8_0(keys)
+        assert abs(result.compression_ratio - 1.0) < 0.01
+
+    def test_q8_0_preserves_original_dtype(self) -> None:
+        t = torch.randn(4, 8, 32, 128, dtype=torch.float32)
+        result = compress_q8_0(t)
+        assert result.original_dtype == torch.float32
+
+    def test_q8_0_padding_dim_not_divisible(self) -> None:
+        """Head dims not divisible by 32 get padded then unpadded."""
+        t = torch.randn(2, 4, 16, 96, dtype=torch.float16)  # 96 = 3*32, exact
+        result = compress_q8_0(t)
+        assert result.data.shape == t.shape
+
+        t2 = torch.randn(2, 4, 16, 100, dtype=torch.float16)  # 100 not div by 32
+        result2 = compress_q8_0(t2)
+        assert result2.data.shape == t2.shape
+
+    def test_q8_0_zero_tensor(self) -> None:
+        """All-zero tensor should round-trip exactly."""
+        t = torch.zeros(2, 4, 16, 128, dtype=torch.float16)
+        result = compress_q8_0(t)
+        decompressed = decompress_q8_0(result.data)
+        assert torch.allclose(decompressed, t.to(torch.float16), atol=1e-6)
+
+
+# ── PolarQuant ───────────────────────────────────────────────────────────────
+
+
+class TestPolarQuant:
+    """PolarQuant: MSE-optimal random rotation + Lloyd-Max at 3 bits.
+    QJL intentionally absent (D5).
+    """
+
+    def test_polarquant_shape_preserved(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_polarquant(keys)
+        assert result.data.shape == keys.shape
+
+    def test_polarquant_output_dtype(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_polarquant(keys)
+        assert result.data.dtype == torch.bfloat16
+
+    def test_polarquant_method_tag(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        keys, _ = llama_kv_256
+        result = compress_polarquant(keys)
+        assert result.method == CompressionMethod.POLARQUANT
+
+    def test_polarquant_metadata_qjl_disabled(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """D5: QJL must be marked disabled in metadata."""
+        keys, _ = llama_kv_256
+        result = compress_polarquant(keys)
+        assert result.metadata["qjl_enabled"] == "false"
+        assert result.metadata["polarquant_bits"] == "3"
+
+    def test_polarquant_round_trip_bounded_error(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """PolarQuant 3-bit error should be < 15% relative MSE.
+
+        3-bit Lloyd-Max on rotated Gaussian: theoretical ~10% for 8 centroids.
+        Allow margin for rotation + dtype casting.
+        """
+        keys, _ = llama_kv_256
+        result = compress_polarquant(keys)
+        decompressed = decompress_polarquant(result.data)
+
+        original = keys.float()
+        restored = decompressed.float()
+
+        mse = ((original - restored) ** 2).mean()
+        signal_power = (original**2).mean()
+        relative_mse = (mse / signal_power).item()
+        assert relative_mse < 0.15, f"PolarQuant relative MSE {relative_mse:.4f} > 15%"
+
+    def test_polarquant_worse_than_q8_0(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """3-bit PolarQuant should have higher error than 8-bit Q8_0."""
+        keys, _ = llama_kv_256
+        original = keys.float()
+
+        q8_result = compress_q8_0(keys)
+        pq_result = compress_polarquant(keys)
+
+        q8_mse = ((original - decompress_q8_0(q8_result.data).float()) ** 2).mean()
+        pq_mse = (
+            (original - decompress_polarquant(pq_result.data).float()) ** 2
+        ).mean()
+
+        assert pq_mse > q8_mse, "PolarQuant 3-bit should be less accurate than Q8_0"
+
+    def test_polarquant_deterministic(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """Same input → same output (fixed seed rotation matrix)."""
+        keys, _ = llama_kv_256
+        r1 = compress_polarquant(keys)
+        r2 = compress_polarquant(keys)
+        assert torch.equal(r1.data, r2.data)
+
+    def test_polarquant_phi3_shape(
+        self, phi3_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """Phi-3 head_dim=96 works with PolarQuant."""
+        keys, _ = phi3_kv_256
+        result = compress_polarquant(keys)
+        assert result.data.shape == keys.shape
+
+
+# ── Dispatcher ───────────────────────────────────────────────────────────────
+
+
+class TestDispatcher:
+    """compress() and decompress() dispatch to correct implementations."""
+
+    @pytest.mark.parametrize(
+        "method",
+        [CompressionMethod.FP16, CompressionMethod.Q8_0, CompressionMethod.POLARQUANT],
+    )
+    def test_compress_dispatches(self, method: CompressionMethod) -> None:
+        t = torch.randn(2, 4, 16, 128, dtype=torch.float16)
+        result = compress(t, method)
+        assert isinstance(result, CompressionResult)
+        assert result.method == method
+
+    @pytest.mark.parametrize(
+        "method",
+        [CompressionMethod.FP16, CompressionMethod.Q8_0, CompressionMethod.POLARQUANT],
+    )
+    def test_decompress_returns_fp16(self, method: CompressionMethod) -> None:
+        t = torch.randn(2, 4, 16, 128, dtype=torch.float16)
+        result = compress(t, method)
+        out = decompress(result.data, method)
+        assert out.dtype == torch.float16
+
+    def test_q4_0_warns_and_falls_back(self) -> None:
+        """D5: Q4_0 emits warning and uses Q8_0 instead."""
+        t = torch.randn(2, 4, 16, 128, dtype=torch.float16)
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            result = compress(t, CompressionMethod.Q4_0)
+            assert len(w) == 1
+            assert "Q4_0" in str(w[0].message)
+            assert "92%" in str(w[0].message)
+        assert result.method == CompressionMethod.Q8_0
+
+    def test_unknown_method_raises(self) -> None:
+        t = torch.randn(2, 4, 16, 128, dtype=torch.float16)
+        with pytest.raises(ValueError, match="Unknown compression method"):
+            compress(t, "invalid_method")  # type: ignore[arg-type]
+
+    def test_decompress_unknown_raises(self) -> None:
+        t = torch.randn(2, 4, 16, 128, dtype=torch.float16)
+        with pytest.raises(ValueError, match="Unknown compression method"):
+            decompress(t, "invalid_method")  # type: ignore[arg-type]
+
+
+# ── Round-trip Integration ───────────────────────────────────────────────────
+
+
+class TestRoundTrip:
+    """Full compress → decompress round-trip through dispatcher."""
+
+    @pytest.mark.parametrize(
+        "method",
+        [CompressionMethod.FP16, CompressionMethod.Q8_0, CompressionMethod.POLARQUANT],
+    )
+    def test_round_trip_shape_preserved(self, method: CompressionMethod) -> None:
+        t = torch.randn(4, 8, 64, 128, dtype=torch.float16)
+        result = compress(t, method)
+        out = decompress(result.data, method)
+        assert out.shape == t.shape
+
+    def test_round_trip_both_kv(
+        self, llama_kv_256: tuple[torch.Tensor, torch.Tensor]
+    ) -> None:
+        """Compress and decompress both keys and values."""
+        keys, values = llama_kv_256
+        for tensor in (keys, values):
+            for method in (CompressionMethod.FP16, CompressionMethod.Q8_0):
+                result = compress(tensor, method)
+                out = decompress(result.data, method)
+                assert out.shape == tensor.shape
+                assert out.dtype == torch.float16

tests/test_eigengram.py

@@ -0,0 +1,152 @@
+"""
+EIGENGRAM test suite — no model calls, pure format verification.
+"""
+
+from __future__ import annotations
+
+import os
+import struct
+
+import pytest
+import torch
+
+from kvcos.engram.format import (
+    EigramDecoder,
+    EigramEncoder,
+    EIGENGRAM_MAGIC,
+    EIGENGRAM_VERSION,
+)
+
+BASIS_PATH = "results/corpus_basis_fcdb_v2.pt"
+
+
+@pytest.fixture(scope="module")
+def basis():
+    if not os.path.exists(BASIS_PATH):
+        pytest.skip("FCDB v2 basis not built yet")
+    return torch.load(BASIS_PATH, weights_only=False)
+
+
+@pytest.fixture(scope="module")
+def sample_cert(basis):
+    enc = EigramEncoder()
+    R = basis["basis"].shape[0]
+    return enc.encode(
+        vec_perdoc=torch.randn(R),
+        vec_fcdb=torch.randn(R),
+        joint_center=basis["joint_center"],
+        corpus_hash="a" * 32,
+        model_id="Llama-3.1-8B",
+        basis_rank=R,
+        n_corpus=200,
+        layer_range=(8, 24),
+        context_len=512,
+        l2_norm=1.234,
+        scs=0.42,
+        margin_proof=0.013,
+        task_description="Test document for transformer attention.",
+        cache_id="test-doc-001",
+    )
+
+
+class TestFormat:
+    def test_magic_present(self, sample_cert: bytes) -> None:
+        assert sample_cert[:4] == EIGENGRAM_MAGIC
+
+    def test_version_byte(self, sample_cert: bytes) -> None:
+        assert struct.unpack_from("<B", sample_cert, 4)[0] == EIGENGRAM_VERSION
+
+    def test_minimum_size(self, sample_cert: bytes, basis) -> None:
+        R = basis["basis"].shape[0]
+        min_size = 99 + R * 2 + R * 2 + 128 * 2
+        assert len(sample_cert) >= min_size
+
+    def test_file_size_reasonable(self, sample_cert: bytes) -> None:
+        assert len(sample_cert) < 2048
+
+
+class TestRoundTrip:
+    def test_model_id(self, sample_cert: bytes) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert rec["model_id"] == "Llama-3.1-8B"
+
+    def test_basis_rank(self, sample_cert: bytes, basis) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert rec["basis_rank"] == basis["basis"].shape[0]
+
+    def test_vec_perdoc_shape(self, sample_cert: bytes, basis) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert rec["vec_perdoc"].shape == (basis["basis"].shape[0],)
+
+    def test_vec_fcdb_shape(self, sample_cert: bytes, basis) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert rec["vec_fcdb"].shape == (basis["basis"].shape[0],)
+
+    def test_joint_center_shape(self, sample_cert: bytes) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert rec["joint_center"].shape == (128,)
+
+    def test_scs(self, sample_cert: bytes) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert abs(rec["scs"] - 0.42) < 0.01
+
+    def test_margin_proof(self, sample_cert: bytes) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert abs(rec["margin_proof"] - 0.013) < 0.001
+
+    def test_task_description(self, sample_cert: bytes) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert "transformer" in rec["task_description"]
+
+    def test_cache_id(self, sample_cert: bytes) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert rec["cache_id"] == "test-doc-001"
+
+    def test_layer_range(self, sample_cert: bytes) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert rec["layer_range"] == (8, 24)
+
+    def test_n_corpus(self, sample_cert: bytes) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert rec["n_corpus"] == 200
+
+    def test_context_len(self, sample_cert: bytes) -> None:
+        rec = EigramDecoder().decode(sample_cert)
+        assert rec["context_len"] == 512
+
+    def test_float16_cosine_preserved(self, basis) -> None:
+        enc = EigramEncoder()
+        R = basis["basis"].shape[0]
+        v = torch.randn(R)
+        v = v / v.norm()
+        cert = enc.encode(
+            vec_perdoc=v, vec_fcdb=v,
+            joint_center=basis["joint_center"],
+            corpus_hash="a" * 32, model_id="test",
+            basis_rank=R, n_corpus=200,
+            layer_range=(8, 24), context_len=0,
+            l2_norm=1.0, scs=0.5, margin_proof=0.0,
+            task_description="cosine test", cache_id="cos",
+        )
+        rec = EigramDecoder().decode(cert)
+        cos = torch.nn.functional.cosine_similarity(
+            v.unsqueeze(0), rec["vec_perdoc"].unsqueeze(0)
+        ).item()
+        assert cos > 0.999, f"Cosine after round-trip: {cos:.5f}"
+
+
+class TestErrorHandling:
+    def test_bad_magic_raises(self) -> None:
+        bad = b"XXXX" + b"\x00" * 200
+        with pytest.raises(ValueError, match="magic"):
+            EigramDecoder().decode(bad)
+
+    def test_wrong_version_raises(self, sample_cert: bytes) -> None:
+        data = bytearray(sample_cert)
+        data[4] = 99
+        with pytest.raises(ValueError, match="version"):
+            EigramDecoder().decode(bytes(data))
+
+    def test_truncated_raises(self, sample_cert: bytes) -> None:
+        with pytest.raises(Exception):
+            EigramDecoder().decode(sample_cert[:20])

tests/test_embedder.py

@@ -0,0 +1,106 @@
+"""Tests for kvcos.engram.embedder — unified fingerprint embedding."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from kvcos.engram.embedder import (
+    HashEmbedder,
+    get_embedder,
+    get_fingerprint,
+    reset_embedder,
+)
+
+
+class TestHashEmbedder:
+    def test_deterministic(self):
+        emb = HashEmbedder(dim=128)
+        fp1 = emb.embed("hello")
+        fp2 = emb.embed("hello")
+        assert torch.allclose(fp1, fp2)
+
+    def test_different_text(self):
+        emb = HashEmbedder(dim=128)
+        fp1 = emb.embed("hello")
+        fp2 = emb.embed("world")
+        assert not torch.allclose(fp1, fp2)
+
+    def test_normalized(self):
+        emb = HashEmbedder(dim=128)
+        fp = emb.embed("test")
+        norm = torch.norm(fp).item()
+        assert abs(norm - 1.0) < 0.01
+
+    def test_dimension(self):
+        emb = HashEmbedder(dim=256)
+        fp = emb.embed("test")
+        assert fp.shape == (256,)
+        assert emb.dim == 256
+
+    def test_source_tag(self):
+        emb = HashEmbedder()
+        assert emb.source == "hash-fallback"
+
+
+class TestGetFingerprint:
+    def test_returns_tensor_and_source(self):
+        fp, source = get_fingerprint("test text")
+        assert isinstance(fp, torch.Tensor)
+        assert isinstance(source, str)
+        assert source in ("llama_cpp", "sbert", "hash-fallback")
+
+    def test_deterministic(self):
+        fp1, _ = get_fingerprint("same text")
+        fp2, _ = get_fingerprint("same text")
+        assert torch.allclose(fp1, fp2)
+
+
+class TestSBertEmbedder:
+    """Test sbert if available (installed in this venv)."""
+
+    def test_sbert_available(self):
+        """Verify sentence-transformers is usable."""
+        try:
+            from kvcos.engram.embedder import SBertEmbedder
+            emb = SBertEmbedder()
+            assert emb.source == "sbert"
+            assert emb.dim == 384
+        except ImportError:
+            pytest.skip("sentence-transformers not installed")
+
+    def test_semantic_discrimination(self):
+        """Related texts should be more similar than unrelated."""
+        try:
+            from kvcos.engram.embedder import SBertEmbedder
+            emb = SBertEmbedder()
+        except ImportError:
+            pytest.skip("sentence-transformers not installed")
+
+        fp_a = emb.embed("machine learning neural network training")
+        fp_b = emb.embed("deep learning model optimization")
+        fp_c = emb.embed("chocolate cake baking recipe")
+
+        sim_ab = F.cosine_similarity(fp_a.unsqueeze(0), fp_b.unsqueeze(0)).item()
+        sim_ac = F.cosine_similarity(fp_a.unsqueeze(0), fp_c.unsqueeze(0)).item()
+
+        assert sim_ab > sim_ac, (
+            f"Related topics ({sim_ab:.4f}) should be more similar "
+            f"than unrelated ({sim_ac:.4f})"
+        )
+
+
+class TestGetEmbedder:
+    def test_singleton(self):
+        reset_embedder()
+        e1 = get_embedder()
+        e2 = get_embedder()
+        assert e1 is e2
+
+    def test_reset(self):
+        reset_embedder()
+        e1 = get_embedder()
+        reset_embedder()
+        e2 = get_embedder()
+        # After reset, a new instance is created
+        # (may or may not be same object depending on strategy)
+        assert e2 is not None

tests/test_hnsw_index.py

@@ -0,0 +1,108 @@
+"""Tests for kvcos.engram.hnsw_index — HNSW nearest-neighbor index."""
+
+import torch
+import pytest
+
+from kvcos.engram.hnsw_index import EngramIndex, HNSWResult
+
+
+@pytest.fixture
+def small_index():
+    """Build a small 8-dim HNSW index with 5 documents."""
+    idx = EngramIndex(dim=8)
+    ids = [f"doc_{i}" for i in range(5)]
+    # deterministic orthogonal-ish vectors
+    vecs = torch.eye(5, 8)
+    idx.add_batch(ids, vecs)
+    return idx
+
+
+class TestEngramIndexBuild:
+    def test_add_batch_len(self, small_index):
+        assert len(small_index) == 5
+
+    def test_add_batch_ids_stored(self, small_index):
+        assert small_index._ids == [f"doc_{i}" for i in range(5)]
+
+    def test_repr(self, small_index):
+        r = repr(small_index)
+        assert "n=5" in r
+        assert "dim=8" in r
+
+
+class TestEngramIndexSearch:
+    def test_search_returns_results(self, small_index):
+        query = torch.eye(5, 8)[0]  # matches doc_0
+        results = small_index.search(query, top_k=3)
+        assert len(results) == 3
+        assert results[0].doc_id == "doc_0"
+
+    def test_search_scores_descending(self, small_index):
+        query = torch.eye(5, 8)[2]
+        results = small_index.search(query, top_k=5)
+        scores = [r.score for r in results]
+        assert scores == sorted(scores, reverse=True)
+
+    def test_search_margin(self, small_index):
+        query = torch.eye(5, 8)[0]
+        results = small_index.search(query, top_k=3)
+        assert results[0].margin >= 0
+
+    def test_search_raises_before_build(self):
+        idx = EngramIndex(dim=8)
+        with pytest.raises(RuntimeError, match="not built"):
+            idx.search(torch.randn(8), top_k=1)
+
+
+class TestEngramIndexGetVector:
+    def test_get_vector_returns_tensor(self, small_index):
+        vec = small_index.get_vector("doc_0")
+        assert vec is not None
+        assert isinstance(vec, torch.Tensor)
+        assert vec.shape == (8,)
+
+    def test_get_vector_none_for_missing(self, small_index):
+        vec = small_index.get_vector("nonexistent")
+        assert vec is None
+
+    def test_get_vector_reconstructs_normalized(self, small_index):
+        """Vectors are L2-normalized on add, so reconstruction should be unit-length."""
+        vec = small_index.get_vector("doc_0")
+        norm = torch.norm(vec).item()
+        assert abs(norm - 1.0) < 0.01
+
+    def test_get_vector_matches_original_direction(self, small_index):
+        """Reconstructed vector should point in the same direction as the original."""
+        original = torch.nn.functional.normalize(torch.eye(5, 8)[3:4], dim=-1)[0]
+        reconstructed = small_index.get_vector("doc_3")
+        cosine = torch.dot(original, reconstructed).item()
+        assert cosine > 0.99
+
+
+class TestEngramIndexPersistence:
+    def test_save_and_load(self, small_index, tmp_path):
+        path = str(tmp_path / "test_hnsw")
+        small_index.save(path)
+
+        loaded = EngramIndex.load(path)
+        assert len(loaded) == 5
+        assert loaded._ids == small_index._ids
+
+    def test_loaded_search_matches_original(self, small_index, tmp_path):
+        path = str(tmp_path / "test_hnsw")
+        small_index.save(path)
+        loaded = EngramIndex.load(path)
+
+        query = torch.eye(5, 8)[1]
+        orig_results = small_index.search(query, top_k=3)
+        load_results = loaded.search(query, top_k=3)
+        assert [r.doc_id for r in orig_results] == [r.doc_id for r in load_results]
+
+    def test_loaded_get_vector(self, small_index, tmp_path):
+        path = str(tmp_path / "test_hnsw")
+        small_index.save(path)
+        loaded = EngramIndex.load(path)
+
+        vec = loaded.get_vector("doc_2")
+        assert vec is not None
+        assert vec.shape == (8,)

tests/test_integration_synthetic.py

@@ -0,0 +1,160 @@
+"""
+ENGRAM Protocol — Synthetic Integration Test
+Full pipeline E2E with synthetic tensors — no real model needed.
+
+Pipeline: create KV → extract state → serialize .eng → load → index → query → retrieve
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import load_file
+
+from kvcos.core.cache_spec import LLAMA_3_1_8B
+from kvcos.core.serializer import EngramSerializer
+from kvcos.core.types import CompressionMethod, StateExtractionMode
+from kvcos.core.manifold_index import ManifoldIndex
+from kvcos.core.retriever import EGRRetriever
+from kvcos.core.state_extractor import MARStateExtractor
+from kvcos.storage.local import LocalStorageBackend
+from tests.conftest import make_synthetic_kv
+
+
+class TestFullPipeline:
+    """End-to-end: store → index → query → retrieve using synthetic data."""
+
+    def test_serialize_round_trip(self, tmp_data_dir: Path) -> None:
+        """Step 1-4: Create → serialize → load → verify shape."""
+        keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=256)
+        assert keys.shape == (32, 8, 256, 128)
+
+        serializer = EngramSerializer()
+        eng_path = tmp_data_dir / "roundtrip.eng"
+
+        serializer.serialize(
+            keys=keys, values=values,
+            agent_id="integration-test", task_description="round-trip test",
+            model_id=LLAMA_3_1_8B["model_id"], output_path=eng_path,
+            compression=CompressionMethod.FP16,
+        )
+        assert eng_path.exists()
+
+        # Verify valid safetensors
+        tensors = load_file(str(eng_path))
+        assert "layer_0_keys" in tensors
+
+        k_out, v_out, meta = serializer.deserialize(eng_path)
+        assert k_out.shape == keys.shape
+        assert v_out.shape == values.shape
+
+    @pytest.mark.parametrize("mode", list(StateExtractionMode))
+    def test_extraction_all_modes(self, mode: StateExtractionMode) -> None:
+        """Step 2: Extract state vector in all 3 modes."""
+        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=256)
+        extractor = MARStateExtractor(mode=mode, rank=128)
+        result = extractor.extract(keys, LLAMA_3_1_8B)
+
+        assert result.state_vec.dim() == 1
+        assert result.state_vec.shape[0] > 0
+        assert result.l2_norm > 0
+        assert result.mode == mode
+
+    def test_index_and_query(self, tmp_data_dir: Path) -> None:
+        """Step 5-6: Index state vector → query with different tensor → get result."""
+        keys_a, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=256, seed=42)
+        keys_b, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=256, seed=99)
+
+        extractor = MARStateExtractor(
+            mode=StateExtractionMode.MEAN_POOL,
+        )
+        dim = extractor.output_dim(LLAMA_3_1_8B)
+        index = ManifoldIndex(dim=dim)
+
+        # Extract and index first tensor
+        from kvcos.core.manifold_index import IndexEntry
+
+        result_a = extractor.extract(keys_a, LLAMA_3_1_8B)
+        index.add(
+            result_a.state_vec,
+            IndexEntry(
+                cache_id="test-cache-a",
+                task_description="indexed engram",
+                model_id=LLAMA_3_1_8B["model_id"],
+                created_at="2026-01-01T00:00:00Z",
+                context_len=256,
+                l2_norm=result_a.l2_norm,
+            ),
+        )
+
+        # Query with second tensor
+        result_b = extractor.extract(keys_b, LLAMA_3_1_8B)
+        results = index.search(result_b.state_vec, top_k=1)
+
+        assert len(results) >= 1
+        assert results[0]["cache_id"] == "test-cache-a"
+
+    def test_full_egr_pipeline(self, tmp_data_dir: Path) -> None:
+        """Step 7: Full EGR retrieval — store → index → query → retrieve."""
+        keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=256, seed=42)
+        query_keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=256, seed=99)
+
+        extractor = MARStateExtractor(
+            mode=StateExtractionMode.MEAN_POOL,
+        )
+        dim = extractor.output_dim(LLAMA_3_1_8B)
+        index = ManifoldIndex(dim=dim)
+        storage = LocalStorageBackend(data_dir=tmp_data_dir)
+        retriever = EGRRetriever(extractor, index, storage)
+
+        # Store
+        cache_id = retriever.index_engram(
+            keys=keys, values=values, spec=LLAMA_3_1_8B,
+            agent_id="integration-test",
+            task_description="full pipeline test",
+            model_id=LLAMA_3_1_8B["model_id"],
+            output_dir=tmp_data_dir,
+        )
+        assert isinstance(cache_id, str)
+        assert index.n_entries == 1
+
+        # Retrieve
+        response = retriever.retrieve(query_keys, LLAMA_3_1_8B, top_k=1)
+        assert len(response.results) >= 1
+
+        result = response.results[0]
+        assert result.cache_id == cache_id
+        assert result.keys.shape == keys.shape
+        assert result.values.shape == values.shape
+        assert result.similarity != 0.0
+
+    def test_multi_engram_ranking(self, tmp_data_dir: Path) -> None:
+        """Store 3 engrams, query, verify results are ranked by similarity."""
+        extractor = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)
+        dim = extractor.output_dim(LLAMA_3_1_8B)
+        index = ManifoldIndex(dim=dim)
+        storage = LocalStorageBackend(data_dir=tmp_data_dir)
+        retriever = EGRRetriever(extractor, index, storage)
+
+        for seed in (10, 20, 30):
+            keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64, seed=seed)
+            retriever.index_engram(
+                keys=keys, values=values, spec=LLAMA_3_1_8B,
+                agent_id="test", task_description=f"seed-{seed}",
+                model_id=LLAMA_3_1_8B["model_id"],
+                output_dir=tmp_data_dir,
+            )
+
+        assert index.n_entries == 3
+
+        query_keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64, seed=10)
+        response = retriever.retrieve(query_keys, LLAMA_3_1_8B, top_k=3)
+
+        assert len(response.results) == 3
+        # Results should be sorted by descending similarity
+        sims = [r.similarity for r in response.results]
+        assert sims == sorted(sims, reverse=True)
+        # Closest match should be seed=10 (same as query)
+        assert response.results[0].task_description == "seed-10"

tests/test_iswa_blob_parser.py

@@ -0,0 +1,136 @@
+"""
+ENGRAM Protocol — ISWA Blob Parser Tests
+Tests for multi-section KV cache parsing (Gemma 4 ISWA format).
+
+Uses synthetic ISWA blobs from conftest.make_synthetic_iswa_blob().
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from kvcos.core.blob_parser import (
+    BlobParseError,
+    ParsedKVCache,
+    ParsedMultiSectionCache,
+    parse_multi_section_blob,
+    parse_state_blob,
+)
+from kvcos.core.types import AttentionType, CacheSection, ModelCacheSpec
+from tests.conftest import (
+    GEMMA4_GLOBAL_SECTION,
+    GEMMA4_SECTIONS,
+    GEMMA4_SWA_SECTION,
+    make_synthetic_iswa_blob,
+)
+
+
+class TestParseMultiSectionBlob:
+    """Parse ISWA blobs with multiple KV cache sections."""
+
+    def test_parse_gemma4_shape(self) -> None:
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=4)
+        result = parse_multi_section_blob(blob, GEMMA4_SECTIONS)
+
+        assert len(result.sections) == 2
+
+        # Global section: [5, 2, 4, 512]
+        s0 = result.sections[0]
+        assert s0.keys.shape == (5, 2, 4, 512)
+        assert s0.values.shape == (5, 2, 4, 512)
+
+        # SWA section: [25, 8, 4, 256]
+        s1 = result.sections[1]
+        assert s1.keys.shape == (25, 8, 4, 256)
+        assert s1.values.shape == (25, 8, 4, 256)
+
+    def test_parse_metadata(self) -> None:
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=4)
+        result = parse_multi_section_blob(blob, GEMMA4_SECTIONS)
+
+        assert result.arch == "gemma4"
+        assert result.n_sections == 2
+        assert result.total_layers == 30
+
+        assert result.sections[0].n_layers == 5
+        assert result.sections[0].arch == "gemma4"
+        assert result.sections[1].n_layers == 25
+
+    def test_parse_cells(self) -> None:
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=4)
+        result = parse_multi_section_blob(blob, GEMMA4_SECTIONS)
+
+        for sec in result.sections:
+            assert sec.n_cells == 4
+            assert len(sec.cells) == 4
+            assert sec.cells[0].pos == 0
+            assert sec.cells[3].pos == 3
+
+    def test_dtype_float16(self) -> None:
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=2)
+        result = parse_multi_section_blob(blob, GEMMA4_SECTIONS)
+
+        for sec in result.sections:
+            assert sec.keys.dtype == torch.float16
+            assert sec.values.dtype == torch.float16
+
+    def test_different_cell_counts(self) -> None:
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=8)
+        result = parse_multi_section_blob(blob, GEMMA4_SECTIONS)
+
+        assert result.sections[0].n_cells == 8
+        assert result.sections[1].n_cells == 8
+
+    def test_non_transposed_v(self) -> None:
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=2, v_trans=False)
+        result = parse_multi_section_blob(blob, GEMMA4_SECTIONS)
+
+        for sec in result.sections:
+            assert sec.v_trans is False
+
+    def test_single_section_works(self) -> None:
+        """Single-section ISWA parse should work identically to standard."""
+        single = (GEMMA4_GLOBAL_SECTION,)
+        blob = make_synthetic_iswa_blob(single, n_cells=4)
+        result = parse_multi_section_blob(blob, single)
+
+        assert len(result.sections) == 1
+        assert result.sections[0].keys.shape == (5, 2, 4, 512)
+
+
+class TestParseMultiSectionErrors:
+    """Error handling for ISWA blob parsing."""
+
+    def test_section_mismatch_raises(self) -> None:
+        """Blob has 2 sections but we pass specs for 3."""
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=4)
+        three_sections = GEMMA4_SECTIONS + (GEMMA4_GLOBAL_SECTION,)
+        with pytest.raises(BlobParseError, match="Expected 3.*got 2"):
+            parse_multi_section_blob(blob, three_sections)
+
+    def test_truncated_blob_raises(self) -> None:
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=4)
+        with pytest.raises(BlobParseError):
+            parse_multi_section_blob(blob[:100], GEMMA4_SECTIONS)
+
+    def test_wrong_dimensions_raises(self) -> None:
+        """Pass wrong KV head count for a section."""
+        wrong_sections = (
+            CacheSection(AttentionType.FULL, 5, 4, 512),  # wrong: 4 heads not 2
+            GEMMA4_SWA_SECTION,
+        )
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=4)
+        with pytest.raises(BlobParseError):
+            parse_multi_section_blob(blob, wrong_sections)
+
+
+class TestStandardBlobBackwardCompat:
+    """Ensure parse_state_blob still works for single-stream blobs."""
+
+    def test_single_stream_still_works(self) -> None:
+        from tests.test_blob_parser import _make_blob
+
+        blob = _make_blob(16, 32, 8, 128)
+        result = parse_state_blob(blob, n_kv_heads=8, head_dim=128)
+        assert result.keys.shape == (32, 8, 16, 128)

tests/test_iswa_bridge.py

@@ -0,0 +1,92 @@
+"""
+ENGRAM Protocol — ISWA Bridge Tests
+Tests for multi-architecture metadata detection and ISWA cache extraction.
+Does NOT require a real GGUF model — tests the metadata helpers and spec logic.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from integrations.llama_cpp_bridge import _meta_get
+from kvcos.core.cache_spec import (
+    GEMMA_4_26B_A4B,
+    LLAMA_3_1_8B,
+    get_model_spec,
+    is_iswa_spec,
+)
+
+
+class TestMetaGet:
+    """Metadata key fallback chain across architecture prefixes."""
+
+    def test_llama_prefix(self) -> None:
+        meta = {"llama.block_count": "32"}
+        assert _meta_get(meta, "block_count") == "32"
+
+    def test_gemma4_prefix(self) -> None:
+        meta = {"gemma4.block_count": "30"}
+        assert _meta_get(meta, "block_count") == "30"
+
+    def test_gemma_prefix(self) -> None:
+        meta = {"gemma.attention.head_count": "8"}
+        assert _meta_get(meta, "attention.head_count") == "8"
+
+    def test_general_fallback(self) -> None:
+        meta = {"general.block_count": "28"}
+        assert _meta_get(meta, "block_count") == "28"
+
+    def test_default_when_missing(self) -> None:
+        meta = {}
+        assert _meta_get(meta, "block_count", "32") == "32"
+
+    def test_llama_takes_priority(self) -> None:
+        meta = {
+            "llama.block_count": "32",
+            "gemma4.block_count": "30",
+            "general.block_count": "28",
+        }
+        assert _meta_get(meta, "block_count") == "32"
+
+    def test_gemma4_before_general(self) -> None:
+        meta = {
+            "gemma4.embedding_length": "3072",
+            "general.embedding_length": "4096",
+        }
+        assert _meta_get(meta, "embedding_length") == "3072"
+
+
+class TestISWASpecDetection:
+    """Registry and ISWA detection."""
+
+    def test_gemma4_in_registry(self) -> None:
+        spec = get_model_spec("google/gemma-4-26b-a4b-it")
+        assert spec is not None
+        assert spec["model_family"] == "gemma"
+
+    def test_gemma4_is_iswa(self) -> None:
+        assert is_iswa_spec(GEMMA_4_26B_A4B) is True
+
+    def test_llama_not_iswa(self) -> None:
+        assert is_iswa_spec(LLAMA_3_1_8B) is False
+
+    def test_gemma4_sections_correct(self) -> None:
+        sections = GEMMA_4_26B_A4B["cache_sections"]
+        assert len(sections) == 2
+
+        # Global section
+        assert sections[0].n_layers == 5
+        assert sections[0].n_kv_heads == 2
+        assert sections[0].head_dim == 512
+        assert sections[0].n_embd_kv == 1024
+
+        # SWA section
+        assert sections[1].n_layers == 25
+        assert sections[1].n_kv_heads == 8
+        assert sections[1].head_dim == 256
+        assert sections[1].window_size == 1024
+
+    def test_gemma4_total_layers(self) -> None:
+        sections = GEMMA_4_26B_A4B["cache_sections"]
+        total = sum(s.n_layers for s in sections)
+        assert total == GEMMA_4_26B_A4B["n_layers"]

tests/test_iswa_fingerprint.py

@@ -0,0 +1,94 @@
+"""
+ENGRAM Protocol — ISWA Fingerprint Tests
+Tests for per-section Fourier fingerprint computation and concatenation.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from kvcos.core.blob_parser import ParsedKVCache, ParsedMultiSectionCache, parse_multi_section_blob
+from kvcos.core.fingerprint import (
+    compute_fourier_fingerprint_v2,
+    compute_iswa_fingerprint,
+)
+from kvcos.core.types import AttentionType, CacheSection
+from tests.conftest import GEMMA4_SECTIONS, make_synthetic_iswa_blob
+
+
+class TestISWAFingerprint:
+    """Per-section fingerprint computation for ISWA models."""
+
+    def _make_parsed(self, n_cells: int = 4) -> ParsedMultiSectionCache:
+        blob = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=n_cells)
+        return parse_multi_section_blob(blob, GEMMA4_SECTIONS)
+
+    def test_fingerprint_shape(self) -> None:
+        parsed = self._make_parsed()
+        fp = compute_iswa_fingerprint(parsed, freqs=[0, 1])
+
+        # Global: 2 * 512 * 2 = 2048
+        # SWA:    8 * 256 * 2 = 4096
+        # Total:  6144
+        assert fp.shape == (6144,)
+
+    def test_fingerprint_dtype(self) -> None:
+        parsed = self._make_parsed()
+        fp = compute_iswa_fingerprint(parsed)
+        assert fp.dtype == torch.float32
+
+    def test_fingerprint_normalized(self) -> None:
+        """Each section's sub-FP is concat of per-freq L2-normalized vectors."""
+        parsed = self._make_parsed()
+        fp = compute_iswa_fingerprint(parsed, freqs=[0, 1])
+
+        # Global section FP: first 2048 dims (1024 per freq, 2 freqs)
+        global_fp = fp[:2048]
+        # SWA section FP: next 4096 dims (2048 per freq, 2 freqs)
+        swa_fp = fp[2048:]
+
+        # Each sub-section is 2 concatenated unit vectors → norm = sqrt(2)
+        import math
+        expected_norm = math.sqrt(2)
+        assert abs(global_fp.norm().item() - expected_norm) < 0.05
+        assert abs(swa_fp.norm().item() - expected_norm) < 0.05
+
+    def test_deterministic(self) -> None:
+        parsed = self._make_parsed()
+        fp1 = compute_iswa_fingerprint(parsed)
+        fp2 = compute_iswa_fingerprint(parsed)
+        assert torch.allclose(fp1, fp2)
+
+    def test_different_inputs_differ(self) -> None:
+        p1 = self._make_parsed(n_cells=4)
+        blob2 = make_synthetic_iswa_blob(GEMMA4_SECTIONS, n_cells=4, seed=999)
+        p2 = parse_multi_section_blob(blob2, GEMMA4_SECTIONS)
+
+        fp1 = compute_iswa_fingerprint(p1)
+        fp2 = compute_iswa_fingerprint(p2)
+        cos = torch.nn.functional.cosine_similarity(fp1.unsqueeze(0), fp2.unsqueeze(0))
+        assert cos.item() < 0.99  # different inputs → different FPs
+
+    def test_single_section_matches_standard(self) -> None:
+        """Single-section ISWA FP should match standard FP."""
+        section = CacheSection(AttentionType.FULL, 5, 2, 512)
+        blob = make_synthetic_iswa_blob((section,), n_cells=4)
+        parsed = parse_multi_section_blob(blob, (section,))
+
+        iswa_fp = compute_iswa_fingerprint(parsed, freqs=[0, 1])
+
+        # Compare with standard FP on same data
+        layer_keys = parsed.sections[0].keys.float().mean(dim=2)
+        standard_fp = compute_fourier_fingerprint_v2(layer_keys, freqs=[0, 1])
+
+        assert torch.allclose(iswa_fp, standard_fp, atol=1e-5)
+
+    def test_custom_freqs(self) -> None:
+        parsed = self._make_parsed()
+        fp_f0 = compute_iswa_fingerprint(parsed, freqs=[0])
+        fp_f01 = compute_iswa_fingerprint(parsed, freqs=[0, 1])
+
+        # f0 only: Global(1024) + SWA(2048) = 3072
+        assert fp_f0.shape == (3072,)
+        # f0+f1: double
+        assert fp_f01.shape == (6144,)

tests/test_iswa_types.py

@@ -0,0 +1,132 @@
+"""
+ENGRAM Protocol — ISWA Type System Tests
+Tests for CacheSection, AttentionType, and ISWA-aware ModelCacheSpec.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from kvcos.core.types import AttentionType, CacheSection, ModelCacheSpec
+
+
+class TestAttentionType:
+    """AttentionType enum values."""
+
+    def test_full_value(self) -> None:
+        assert AttentionType.FULL == "full"
+
+    def test_sliding_value(self) -> None:
+        assert AttentionType.SLIDING == "sliding"
+
+    def test_is_str(self) -> None:
+        assert isinstance(AttentionType.FULL, str)
+
+
+class TestCacheSection:
+    """CacheSection frozen dataclass."""
+
+    def test_global_section(self) -> None:
+        sec = CacheSection(
+            attention_type=AttentionType.FULL,
+            n_layers=5,
+            n_kv_heads=2,
+            head_dim=512,
+        )
+        assert sec.n_layers == 5
+        assert sec.n_kv_heads == 2
+        assert sec.head_dim == 512
+        assert sec.window_size is None
+
+    def test_sliding_section(self) -> None:
+        sec = CacheSection(
+            attention_type=AttentionType.SLIDING,
+            n_layers=25,
+            n_kv_heads=8,
+            head_dim=256,
+            window_size=1024,
+        )
+        assert sec.attention_type == AttentionType.SLIDING
+        assert sec.window_size == 1024
+
+    def test_n_embd_kv(self) -> None:
+        sec = CacheSection(
+            attention_type=AttentionType.FULL,
+            n_layers=5,
+            n_kv_heads=2,
+            head_dim=512,
+        )
+        assert sec.n_embd_kv == 1024  # 2 * 512
+
+    def test_frozen(self) -> None:
+        sec = CacheSection(
+            attention_type=AttentionType.FULL,
+            n_layers=5,
+            n_kv_heads=2,
+            head_dim=512,
+        )
+        with pytest.raises(AttributeError):
+            sec.n_layers = 10  # type: ignore[misc]
+
+    def test_equality(self) -> None:
+        a = CacheSection(AttentionType.FULL, 5, 2, 512)
+        b = CacheSection(AttentionType.FULL, 5, 2, 512)
+        assert a == b
+
+    def test_inequality(self) -> None:
+        a = CacheSection(AttentionType.FULL, 5, 2, 512)
+        b = CacheSection(AttentionType.SLIDING, 25, 8, 256, 1024)
+        assert a != b
+
+
+class TestModelCacheSpecISWA:
+    """ModelCacheSpec with optional cache_sections."""
+
+    def test_standard_spec_no_sections(self) -> None:
+        spec = ModelCacheSpec(
+            model_id="meta-llama/Llama-3.1-8B-Instruct",
+            model_family="llama",
+            n_layers=32,
+            n_heads=32,
+            n_kv_heads=8,
+            head_dim=128,
+            rope_enabled=True,
+            extraction_layers=tuple(range(8, 32)),
+        )
+        assert "cache_sections" not in spec
+        assert spec["n_kv_heads"] == 8
+
+    def test_iswa_spec_with_sections(self) -> None:
+        sections = (
+            CacheSection(AttentionType.FULL, 5, 2, 512),
+            CacheSection(AttentionType.SLIDING, 25, 8, 256, 1024),
+        )
+        spec = ModelCacheSpec(
+            model_id="google/gemma-4-26b-a4b-it",
+            model_family="gemma",
+            n_layers=30,
+            n_heads=32,
+            n_kv_heads=8,
+            head_dim=256,
+            rope_enabled=True,
+            extraction_layers=tuple(range(8, 30)),
+            cache_sections=sections,
+        )
+        assert "cache_sections" in spec
+        assert len(spec["cache_sections"]) == 2
+        assert spec["cache_sections"][0].n_embd_kv == 1024
+        assert spec["cache_sections"][1].n_embd_kv == 2048
+
+    def test_iswa_total_layers_match(self) -> None:
+        sections = (
+            CacheSection(AttentionType.FULL, 5, 2, 512),
+            CacheSection(AttentionType.SLIDING, 25, 8, 256, 1024),
+        )
+        total = sum(s.n_layers for s in sections)
+        assert total == 30
+
+    def test_backward_compat_existing_specs(self) -> None:
+        """Existing specs without cache_sections still work."""
+        from kvcos.core.cache_spec import LLAMA_3_1_8B
+        assert LLAMA_3_1_8B["n_kv_heads"] == 8
+        assert "cache_sections" not in LLAMA_3_1_8B

tests/test_knowledge_index.py

@@ -0,0 +1,138 @@
+"""Tests for kvcos.engram.knowledge_index — HNSW knowledge search."""
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+
+from kvcos.engram.embedder import get_fingerprint
+from kvcos.engram.format import EigramEncoder
+from kvcos.engram.knowledge_index import KnowledgeIndex
+
+
+@pytest.fixture
+def knowledge_dir(tmp_path):
+    """Create a temporary knowledge directory with test .eng files."""
+    encoder = EigramEncoder()
+    project_dir = tmp_path / "test_project"
+    project_dir.mkdir()
+
+    docs = [
+        ("doc_ml", "Machine learning model training and optimization"),
+        ("doc_db", "PostgreSQL database schema migration tools"),
+        ("doc_api", "REST API endpoint authentication and authorization"),
+        ("doc_test", "Unit testing with pytest fixtures and mocking"),
+        ("doc_deploy", "Docker container deployment to Kubernetes cluster"),
+    ]
+
+    for doc_id, text in docs:
+        fp, source = get_fingerprint(text)
+        dim = fp.shape[0]
+
+        blob = encoder.encode(
+            vec_perdoc=torch.zeros(116),
+            vec_fcdb=torch.zeros(116),
+            joint_center=torch.zeros(128),
+            corpus_hash="test" * 8,
+            model_id=source[:16],
+            basis_rank=116,
+            n_corpus=0,
+            layer_range=(0, 0),
+            context_len=len(text),
+            l2_norm=float(torch.norm(fp).item()),
+            scs=0.0,
+            margin_proof=0.0,
+            task_description=text[:256],
+            cache_id=doc_id,
+            vec_fourier=fp if dim == 2048 else None,
+            vec_fourier_v2=fp,
+            confusion_flag=False,
+        )
+
+        eng_path = project_dir / f"{doc_id}.eng"
+        eng_path.write_bytes(blob)
+
+        meta = {
+            "cache_id": doc_id,
+            "task_description": text,
+            "source_path": f"/test/{doc_id}.md",
+            "project": "test_project",
+            "fp_source": source,
+            "chunk_index": 0,
+            "chunk_total": 1,
+            "headers": [],
+        }
+        meta_path = Path(str(eng_path) + ".meta.json")
+        meta_path.write_text(json.dumps(meta))
+
+    return tmp_path
+
+
+class TestKnowledgeIndexBuild:
+    def test_build_from_directory(self, knowledge_dir):
+        kidx = KnowledgeIndex.build_from_knowledge_dir(
+            knowledge_dir, verbose=False
+        )
+        assert len(kidx) == 5
+
+    def test_build_empty_directory(self, tmp_path):
+        with pytest.raises(ValueError, match="No .eng files"):
+            KnowledgeIndex.build_from_knowledge_dir(tmp_path, verbose=False)
+
+
+class TestKnowledgeIndexSearch:
+    def test_search_returns_results(self, knowledge_dir):
+        kidx = KnowledgeIndex.build_from_knowledge_dir(
+            knowledge_dir, verbose=False
+        )
+        results = kidx.search("database query optimization", k=3)
+        assert len(results) == 3
+        assert all(r.score > 0 for r in results)
+
+    def test_search_result_fields(self, knowledge_dir):
+        kidx = KnowledgeIndex.build_from_knowledge_dir(
+            knowledge_dir, verbose=False
+        )
+        results = kidx.search("testing", k=1)
+        r = results[0]
+        assert r.doc_id
+        assert isinstance(r.score, float)
+        assert r.rank == 0
+        assert r.project == "test_project"
+
+    def test_search_with_tensor(self, knowledge_dir):
+        kidx = KnowledgeIndex.build_from_knowledge_dir(
+            knowledge_dir, verbose=False
+        )
+        query_fp, _ = get_fingerprint("unit tests")
+        results = kidx.search(query_fp, k=2)
+        assert len(results) == 2
+
+    def test_search_margin(self, knowledge_dir):
+        kidx = KnowledgeIndex.build_from_knowledge_dir(
+            knowledge_dir, verbose=False
+        )
+        results = kidx.search("testing", k=3)
+        # Top result should have a margin
+        assert results[0].margin >= 0
+
+
+class TestKnowledgeIndexPersistence:
+    def test_save_and_load(self, knowledge_dir, tmp_path):
+        kidx = KnowledgeIndex.build_from_knowledge_dir(
+            knowledge_dir, verbose=False
+        )
+        index_dir = tmp_path / "index"
+        kidx.save(index_dir)
+
+        loaded = KnowledgeIndex.load(index_dir)
+        assert len(loaded) == len(kidx)
+
+        # Search should work on loaded index
+        results = loaded.search("database", k=2)
+        assert len(results) == 2
+
+    def test_load_nonexistent(self, tmp_path):
+        with pytest.raises(FileNotFoundError):
+            KnowledgeIndex.load(tmp_path / "nonexistent")

tests/test_manifest.py

@@ -0,0 +1,167 @@
+"""Tests for kvcos.engram.manifest — knowledge index registry."""
+
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from kvcos.engram.manifest import ChunkRecord, Manifest, SourceRecord, _content_hash
+
+
+@pytest.fixture
+def tmp_manifest(tmp_path):
+    """Create a Manifest with a temporary path."""
+    return Manifest.load(tmp_path / "manifest.json")
+
+
+class TestContentHash:
+    def test_deterministic(self):
+        assert _content_hash("hello") == _content_hash("hello")
+
+    def test_different_content(self):
+        assert _content_hash("hello") != _content_hash("world")
+
+
+class TestManifestLoad:
+    def test_load_nonexistent_creates_empty(self, tmp_path):
+        m = Manifest.load(tmp_path / "does_not_exist.json")
+        assert m.total_sources == 0
+        assert m.total_chunks == 0
+
+    def test_load_existing(self, tmp_path):
+        # Write a manifest, then load it
+        m = Manifest.load(tmp_path / "manifest.json")
+        m = m.register(
+            source_path="/test/file.md",
+            content_hash="abc123",
+            project="test",
+            file_size=100,
+            chunks=[ChunkRecord(
+                eng_path="/test/file.eng",
+                chunk_index=0,
+                chunk_total=1,
+                char_start=0,
+                char_end=100,
+                indexed_at=1000.0,
+            )],
+        )
+
+        # Load again from disk
+        m2 = Manifest.load(tmp_path / "manifest.json")
+        assert m2.total_sources == 1
+        assert m2.total_chunks == 1
+
+
+class TestManifestRegister:
+    def test_register_new(self, tmp_manifest):
+        chunks = [ChunkRecord(
+            eng_path="/out/test.eng",
+            chunk_index=0,
+            chunk_total=1,
+            char_start=0,
+            char_end=50,
+            indexed_at=1000.0,
+        )]
+        m = tmp_manifest.register(
+            source_path="/src/test.md",
+            content_hash="hash1",
+            project="myproject",
+            file_size=50,
+            chunks=chunks,
+        )
+        assert m.total_sources == 1
+        assert m.total_chunks == 1
+        assert "myproject" in m.projects
+
+    def test_register_overwrites_existing(self, tmp_manifest):
+        chunks1 = [ChunkRecord(
+            eng_path="/out/v1.eng", chunk_index=0, chunk_total=1,
+            char_start=0, char_end=50, indexed_at=1000.0,
+        )]
+        m = tmp_manifest.register(
+            "/src/test.md", "hash1", "proj", 50, chunks1,
+        )
+        assert m.total_chunks == 1
+
+        chunks2 = [
+            ChunkRecord("/out/v2_1.eng", 0, 2, 0, 25, 2000.0),
+            ChunkRecord("/out/v2_2.eng", 1, 2, 25, 50, 2000.0),
+        ]
+        m = m.register("/src/test.md", "hash2", "proj", 50, chunks2)
+        assert m.total_sources == 1  # still 1 source
+        assert m.total_chunks == 2   # now 2 chunks
+
+    def test_register_returns_new_manifest(self, tmp_manifest):
+        """Register returns a new Manifest (immutability)."""
+        m1 = tmp_manifest
+        m2 = m1.register("/src/a.md", "h", "p", 10, [])
+        assert m1.total_sources == 0  # original unchanged
+        assert m2.total_sources == 1
+
+
+class TestManifestNeedsReindex:
+    def test_unknown_file_needs_index(self, tmp_manifest):
+        assert tmp_manifest.needs_reindex("/new/file.md", "any_hash")
+
+    def test_same_hash_no_reindex(self, tmp_manifest):
+        m = tmp_manifest.register("/src/a.md", "hash1", "p", 10, [])
+        assert not m.needs_reindex("/src/a.md", "hash1")
+
+    def test_different_hash_needs_reindex(self, tmp_manifest):
+        m = tmp_manifest.register("/src/a.md", "hash1", "p", 10, [])
+        assert m.needs_reindex("/src/a.md", "hash2")
+
+
+class TestManifestUnregister:
+    def test_unregister_existing(self, tmp_manifest):
+        m = tmp_manifest.register("/src/a.md", "h", "p", 10, [])
+        m = m.unregister("/src/a.md")
+        assert m.total_sources == 0
+
+    def test_unregister_nonexistent(self, tmp_manifest):
+        m = tmp_manifest.unregister("/not/here.md")
+        assert m.total_sources == 0
+
+
+class TestManifestQueries:
+    def test_get_project_records(self, tmp_manifest):
+        m = tmp_manifest
+        m = m.register("/a.md", "h1", "proj_a", 10, [])
+        m = m.register("/b.md", "h2", "proj_b", 20, [])
+        m = m.register("/c.md", "h3", "proj_a", 30, [])
+
+        a_recs = m.get_project_records("proj_a")
+        assert len(a_recs) == 2
+
+    def test_summary(self, tmp_manifest):
+        m = tmp_manifest.register("/a.md", "h", "p", 10, [
+            ChunkRecord("/a.eng", 0, 1, 0, 10, 1000.0),
+        ])
+        s = m.summary()
+        assert s["total_sources"] == 1
+        assert s["total_chunks"] == 1
+        assert "p" in s["projects"]
+
+    def test_contains(self, tmp_manifest):
+        m = tmp_manifest.register("/a.md", "h", "p", 10, [])
+        assert "/a.md" in m
+        assert "/b.md" not in m
+
+    def test_len(self, tmp_manifest):
+        m = tmp_manifest.register("/a.md", "h", "p", 10, [])
+        assert len(m) == 1
+
+
+class TestManifestPersistence:
+    def test_atomic_write(self, tmp_path):
+        m = Manifest.load(tmp_path / "manifest.json")
+        m = m.register("/a.md", "h", "p", 10, [])
+
+        # File should exist
+        assert (tmp_path / "manifest.json").exists()
+
+        # Content should be valid JSON
+        data = json.loads((tmp_path / "manifest.json").read_text())
+        assert data["version"] == 1
+        assert len(data["sources"]) == 1

tests/test_manifold_index.py

@@ -0,0 +1,97 @@
+"""
+ENGRAM Protocol — Manifold Index Tests
+Tests for FAISS IndexFlatIP add/search/remove/persist (D2, D4).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from kvcos.core.manifold_index import IndexEntry, ManifoldIndex
+
+
+def _entry(cid: str = "c1", model: str = "llama") -> IndexEntry:
+    return IndexEntry(
+        cache_id=cid, task_description="test",
+        model_id=model, created_at="2026-01-01T00:00:00Z",
+        context_len=256, l2_norm=1.0,
+    )
+
+
+class TestAddAndSearch:
+    """Add vectors, search via MIPS."""
+
+    def test_add_increments(self) -> None:
+        idx = ManifoldIndex(dim=8)
+        idx.add(torch.randn(8), _entry("a"))
+        idx.add(torch.randn(8), _entry("b"))
+        assert idx.n_entries == 2
+
+    def test_search_returns_correct_order(self) -> None:
+        idx = ManifoldIndex(dim=4)
+        v1 = torch.tensor([1.0, 0.0, 0.0, 0.0])
+        v2 = torch.tensor([0.0, 1.0, 0.0, 0.0])
+        idx.add(v1, _entry("close"))
+        idx.add(v2, _entry("far"))
+
+        query = torch.tensor([1.0, 0.0, 0.0, 0.0])
+        results = idx.search(query, top_k=2)
+        assert results[0]["cache_id"] == "close"
+        assert results[0]["similarity"] > results[1]["similarity"]
+
+    def test_search_empty_returns_empty(self) -> None:
+        idx = ManifoldIndex(dim=4)
+        results = idx.search(torch.randn(4), top_k=5)
+        assert results == []
+
+    def test_model_filter(self) -> None:
+        idx = ManifoldIndex(dim=4)
+        idx.add(torch.randn(4), _entry("a", model="llama"))
+        idx.add(torch.randn(4), _entry("b", model="phi"))
+        results = idx.search(torch.randn(4), top_k=10, model_id="phi")
+        assert all(r["model_id"] == "phi" for r in results)
+
+
+class TestRemoveAndRebuild:
+    """Remove entries and rebuild index."""
+
+    def test_remove_hides_from_search(self) -> None:
+        idx = ManifoldIndex(dim=4)
+        v = torch.tensor([1.0, 0.0, 0.0, 0.0])
+        idx.add(v, _entry("target"))
+        assert idx.remove("target")
+        results = idx.search(v, top_k=1)
+        assert len(results) == 0
+
+    def test_rebuild_compacts(self) -> None:
+        idx = ManifoldIndex(dim=4)
+        for i in range(5):
+            idx.add(torch.randn(4), _entry(f"c{i}"))
+        idx.remove("c1")
+        idx.remove("c3")
+        active = idx.rebuild()
+        assert active == 3
+
+
+class TestPersistence:
+    """Save/load round-trip (D2: serialize_index/deserialize_index)."""
+
+    def test_save_load_round_trip(self, tmp_index_dir: Path) -> None:
+        idx = ManifoldIndex(dim=4)
+        v1 = torch.tensor([1.0, 0.0, 0.0, 0.0])
+        idx.add(v1, _entry("persisted"))
+        idx.save(tmp_index_dir / "test.faiss")
+
+        idx2 = ManifoldIndex(dim=4, index_path=tmp_index_dir / "test.faiss")
+        assert idx2.n_entries == 1
+        results = idx2.search(v1, top_k=1)
+        assert results[0]["cache_id"] == "persisted"
+
+    def test_dim_mismatch_raises(self) -> None:
+        idx = ManifoldIndex(dim=4)
+        with pytest.raises(ValueError, match="dim"):
+            idx.add(torch.randn(8), _entry("wrong"))

tests/test_retriever.py

@@ -0,0 +1,86 @@
+"""
+ENGRAM Protocol — Retriever Tests
+Tests for EGRRetriever: store → index → query → retrieve pipeline.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from kvcos.core.cache_spec import LLAMA_3_1_8B
+from kvcos.core.serializer import EngramSerializer
+from kvcos.core.types import CompressionMethod, StateExtractionMode
+from kvcos.core.manifold_index import ManifoldIndex
+from kvcos.core.retriever import EGRRetriever, RetrievalResponse
+from kvcos.core.state_extractor import MARStateExtractor
+from kvcos.storage.local import LocalStorageBackend
+from tests.conftest import make_synthetic_kv
+
+
+def _build_retriever(
+    data_dir: Path, mode: StateExtractionMode = StateExtractionMode.MEAN_POOL,
+) -> EGRRetriever:
+    ext = MARStateExtractor(mode=mode, rank=128)
+    dim = ext.output_dim(LLAMA_3_1_8B)
+    idx = ManifoldIndex(dim=dim)
+    storage = LocalStorageBackend(data_dir=data_dir)
+    return EGRRetriever(ext, idx, storage)
+
+
+class TestIndexAndRetrieve:
+    """Full store → search → load pipeline."""
+
+    def test_index_returns_cache_id(self, tmp_data_dir: Path) -> None:
+        keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        retriever = _build_retriever(tmp_data_dir)
+
+        cid = retriever.index_engram(
+            keys=keys, values=values, spec=LLAMA_3_1_8B,
+            agent_id="test", task_description="test engram",
+            model_id=LLAMA_3_1_8B["model_id"],
+            output_dir=tmp_data_dir,
+        )
+        assert isinstance(cid, str)
+        assert len(cid) > 0
+
+    def test_retrieve_finds_stored(self, tmp_data_dir: Path) -> None:
+        keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        retriever = _build_retriever(tmp_data_dir)
+
+        retriever.index_engram(
+            keys=keys, values=values, spec=LLAMA_3_1_8B,
+            agent_id="test", task_description="findable engram",
+            model_id=LLAMA_3_1_8B["model_id"],
+            output_dir=tmp_data_dir,
+        )
+
+        query_keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64, seed=99)
+        response = retriever.retrieve(query_keys, LLAMA_3_1_8B, top_k=1)
+
+        assert isinstance(response, RetrievalResponse)
+        assert len(response.results) == 1
+        assert response.results[0].keys.shape == keys.shape
+
+    def test_retrieve_empty_index(self, tmp_data_dir: Path) -> None:
+        retriever = _build_retriever(tmp_data_dir)
+        query_keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        response = retriever.retrieve(query_keys, LLAMA_3_1_8B, top_k=5)
+        assert len(response.results) == 0
+
+    def test_delete_removes(self, tmp_data_dir: Path) -> None:
+        keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        retriever = _build_retriever(tmp_data_dir)
+
+        cid = retriever.index_engram(
+            keys=keys, values=values, spec=LLAMA_3_1_8B,
+            agent_id="test", task_description="deletable",
+            model_id=LLAMA_3_1_8B["model_id"],
+            output_dir=tmp_data_dir,
+        )
+        assert retriever.delete_engram(cid)
+
+        query_keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        response = retriever.retrieve(query_keys, LLAMA_3_1_8B, top_k=5)
+        assert len(response.results) == 0

tests/test_serializer.py

@@ -0,0 +1,118 @@
+"""
+ENGRAM Protocol — Serializer Tests
+Tests for .eng safetensors serialize/deserialize round-trip (D7).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import load_file
+
+from kvcos.core.serializer import EngramSerializer, SerializationError
+from kvcos.core.types import CompressionMethod
+from tests.conftest import make_synthetic_kv
+from kvcos.core.cache_spec import LLAMA_3_1_8B
+
+
+class TestSerializeRoundTrip:
+    """Serialize → deserialize preserves shape, dtype, metadata."""
+
+    def test_round_trip_shape(self, tmp_data_dir: Path) -> None:
+        keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=256)
+        s = EngramSerializer()
+        eng = tmp_data_dir / "test.eng"
+
+        s.serialize(
+            keys=keys, values=values,
+            agent_id="test-agent", task_description="unit test",
+            model_id=LLAMA_3_1_8B["model_id"], output_path=eng,
+            compression=CompressionMethod.FP16,
+        )
+        k_out, v_out, meta = s.deserialize(eng)
+
+        assert k_out.shape == keys.shape
+        assert v_out.shape == values.shape
+
+    def test_metadata_fields(self, tmp_data_dir: Path) -> None:
+        keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        s = EngramSerializer()
+        eng = tmp_data_dir / "meta.eng"
+
+        s.serialize(
+            keys=keys, values=values,
+            agent_id="agent-42", task_description="metadata check",
+            model_id=LLAMA_3_1_8B["model_id"], output_path=eng,
+            compression=CompressionMethod.Q8_0,
+        )
+        _, _, meta = s.deserialize(eng)
+
+        assert meta["agent_id"] == "agent-42"
+        assert meta["task_description"] == "metadata check"
+        assert meta["compression"] == "q8_0"
+        assert meta["n_layers"] == "32"
+        assert meta["model_family"] == "llama"
+
+    def test_safetensors_loadable(self, tmp_data_dir: Path) -> None:
+        """D7: File must be valid safetensors."""
+        keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        s = EngramSerializer()
+        eng = tmp_data_dir / "valid.eng"
+
+        s.serialize(
+            keys=keys, values=values,
+            agent_id="test", task_description="safetensors check",
+            model_id=LLAMA_3_1_8B["model_id"], output_path=eng,
+            compression=CompressionMethod.FP16,
+        )
+        tensors = load_file(str(eng))
+        assert "layer_0_keys" in tensors
+        assert "layer_0_values" in tensors
+
+    def test_result_dict(self, tmp_data_dir: Path) -> None:
+        keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        s = EngramSerializer()
+        eng = tmp_data_dir / "result.eng"
+
+        result = s.serialize(
+            keys=keys, values=values,
+            agent_id="test", task_description="result check",
+            model_id=LLAMA_3_1_8B["model_id"], output_path=eng,
+        )
+        assert "cache_id" in result
+        assert result["size_bytes"] > 0
+        assert result["n_layers"] == 32
+
+
+class TestSerializerErrors:
+    """Edge cases and error handling."""
+
+    def test_shape_mismatch_raises(self, tmp_data_dir: Path) -> None:
+        keys = torch.randn(32, 8, 64, 128, dtype=torch.float16)
+        values = torch.randn(32, 8, 32, 128, dtype=torch.float16)
+        s = EngramSerializer()
+
+        with pytest.raises(SerializationError, match="mismatch"):
+            s.serialize(
+                keys=keys, values=values,
+                agent_id="t", task_description="t",
+                model_id="test", output_path=tmp_data_dir / "bad.eng",
+            )
+
+    def test_3d_tensor_raises(self, tmp_data_dir: Path) -> None:
+        keys = torch.randn(8, 64, 128, dtype=torch.float16)
+        s = EngramSerializer()
+
+        with pytest.raises(SerializationError, match="4D"):
+            s.serialize(
+                keys=keys, values=keys,
+                agent_id="t", task_description="t",
+                model_id="test", output_path=tmp_data_dir / "bad.eng",
+            )
+
+    def test_missing_file_raises(self, tmp_data_dir: Path) -> None:
+        s = EngramSerializer()
+        with pytest.raises(SerializationError, match="not found"):
+            s.deserialize(tmp_data_dir / "nonexistent.eng")

tests/test_state_extractor.py

@@ -0,0 +1,90 @@
+"""
+ENGRAM Protocol — State Extractor Tests
+Tests for all 3 EGR extraction modes (D3).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from kvcos.core.cache_spec import LLAMA_3_1_8B, PHI_3_MINI
+from kvcos.core.types import StateExtractionMode
+from kvcos.core.state_extractor import MARStateExtractor
+from tests.conftest import make_synthetic_kv
+
+
+class TestMeanPool:
+    """mean_pool: mean over layers, heads, context → [head_dim]."""
+
+    def test_output_dim(self) -> None:
+        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        ext = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)
+        result = ext.extract(keys, LLAMA_3_1_8B)
+        assert result.state_vec.shape == (128,)
+
+    def test_output_dim_api(self) -> None:
+        ext = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)
+        assert ext.output_dim(LLAMA_3_1_8B) == 128
+
+    def test_l2_norm_positive(self) -> None:
+        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        ext = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)
+        result = ext.extract(keys, LLAMA_3_1_8B)
+        assert result.l2_norm > 0
+
+    def test_deterministic(self) -> None:
+        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        ext = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)
+        r1 = ext.extract(keys, LLAMA_3_1_8B)
+        r2 = ext.extract(keys, LLAMA_3_1_8B)
+        assert torch.equal(r1.state_vec, r2.state_vec)
+
+
+class TestSVDProject:
+    """svd_project: truncated SVD, rank-160 → [rank]."""
+
+    def test_output_dim(self) -> None:
+        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        ext = MARStateExtractor(mode=StateExtractionMode.SVD_PROJECT, rank=160)
+        result = ext.extract(keys, LLAMA_3_1_8B)
+        assert result.state_vec.shape == (128,)  # clamped to head_dim
+
+    def test_projection_stored(self) -> None:
+        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        ext = MARStateExtractor(mode=StateExtractionMode.SVD_PROJECT, rank=160)
+        ext.extract(keys, LLAMA_3_1_8B)
+        proj = ext.last_projection
+        assert proj is not None
+        assert 0.0 < proj.explained_variance_ratio <= 1.0
+
+    def test_n_layers_used(self) -> None:
+        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        ext = MARStateExtractor(mode=StateExtractionMode.SVD_PROJECT)
+        result = ext.extract(keys, LLAMA_3_1_8B)
+        assert result.n_layers_used == 24  # layers 8-31
+
+
+class TestXKVProject:
+    """xkv_project: grouped cross-layer SVD."""
+
+    def test_output_dim(self) -> None:
+        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        ext = MARStateExtractor(mode=StateExtractionMode.XKV_PROJECT, rank=160)
+        result = ext.extract(keys, LLAMA_3_1_8B)
+        expected_dim = ext.output_dim(LLAMA_3_1_8B)
+        assert result.state_vec.shape == (expected_dim,)
+
+    def test_different_from_mean_pool(self) -> None:
+        keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64)
+        ext_mp = MARStateExtractor(mode=StateExtractionMode.MEAN_POOL)
+        ext_xkv = MARStateExtractor(mode=StateExtractionMode.XKV_PROJECT)
+        r_mp = ext_mp.extract(keys, LLAMA_3_1_8B)
+        r_xkv = ext_xkv.extract(keys, LLAMA_3_1_8B)
+        assert r_mp.state_vec.shape != r_xkv.state_vec.shape
+
+    def test_phi3_works(self) -> None:
+        keys, _ = make_synthetic_kv(PHI_3_MINI, ctx_len=64)
+        ext = MARStateExtractor(mode=StateExtractionMode.XKV_PROJECT, rank=96)
+        result = ext.extract(keys, PHI_3_MINI)
+        assert result.state_vec.dim() == 1
+        assert result.state_vec.shape[0] > 0