Engram-protocol/engram

1"""2ENGRAM Protocol — Block Pool Tests3Tests for 256-token block segmentation/assembly/extend.4"""5 6from __future__ import annotations7 8import pytest9import torch10 11from kvcos.core.block_pool import BlockPool, KVBlock12from kvcos.core.types import BLOCK_SIZE_TOKENS13 14 15def _kv(n_layers: int, n_heads: int, ctx: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]:16    k = torch.randn(n_layers, n_heads, ctx, dim, dtype=torch.float16)17    return k, k.clone()18 19 20class TestSegment:21    """Segment full KV cache into 256-token blocks."""22 23    def test_exact_blocks(self) -> None:24        keys, vals = _kv(32, 8, 512, 128)25        pool = BlockPool(agent_id="a", model_id="m")26        blocks = pool.segment(keys, vals)27        assert len(blocks) == 228        assert all(b.is_full for b in blocks)29 30    def test_partial_last_block(self) -> None:31        keys, vals = _kv(32, 8, 300, 128)32        pool = BlockPool(agent_id="a", model_id="m")33        blocks = pool.segment(keys, vals)34        assert len(blocks) == 235        assert blocks[0].is_full36        assert not blocks[1].is_full37        assert blocks[1].block_len == 4438 39    def test_total_tokens(self) -> None:40        keys, vals = _kv(32, 8, 700, 128)41        pool = BlockPool(agent_id="a", model_id="m")42        pool.segment(keys, vals)43        assert pool.total_tokens == 70044 45 46class TestAssemble:47    """Assemble blocks back into full KV cache."""48 49    def test_round_trip(self) -> None:50        keys, vals = _kv(4, 2, 512, 64)51        pool = BlockPool(agent_id="a", model_id="m")52        pool.segment(keys, vals)53        k_out, v_out = pool.assemble()54        assert torch.equal(k_out, keys)55 56    def test_subset_assembly(self) -> None:57        keys, vals = _kv(4, 2, 768, 64)58        pool = BlockPool(agent_id="a", model_id="m")59        pool.segment(keys, vals)60        k_out, _ = pool.assemble(block_indices=[0, 2])61        assert k_out.shape[2] == BLOCK_SIZE_TOKENS * 262 63    def test_empty_raises(self) -> None:64        pool = BlockPool(agent_id="a", model_id="m")65        with pytest.raises(ValueError, match="No blocks"):66            pool.assemble()67 68 69class TestExtend:70    """Extend pool with new tokens."""71 72    def test_fills_partial_block(self) -> None:73        keys, vals = _kv(4, 2, 200, 64)74        pool = BlockPool(agent_id="a", model_id="m")75        pool.segment(keys, vals)76        assert not pool.blocks[-1].is_full77 78        new_k, new_v = _kv(4, 2, 56, 64)79        pool.extend(new_k, new_v)80        assert pool.blocks[-1].is_full81        assert pool.total_tokens == 25682 83    def test_extend_creates_new_blocks(self) -> None:84        keys, vals = _kv(4, 2, 256, 64)85        pool = BlockPool(agent_id="a", model_id="m")86        pool.segment(keys, vals)87        assert pool.n_blocks == 188 89        new_k, new_v = _kv(4, 2, 300, 64)90        pool.extend(new_k, new_v)91        assert pool.n_blocks == 392        assert pool.total_tokens == 55693