| """ |
| 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 |
|
|
