Engram-protocol/engram

1"""2ENGRAM Protocol — Blob Parser Tests3Tests for llama.cpp state blob → structured tensors (D1).4Uses synthetic blobs matching the real llama_state_get_data() format.5"""6 7from __future__ import annotations8 9import struct10 11import numpy as np12import pytest13import torch14 15from kvcos.core.blob_parser import (16    GGML_TYPE_F16,17    BlobParseError,18    ParsedKVCache,19    parse_state_blob,20)21 22 23def _make_blob(24    n_cells: int,25    n_layers: int,26    n_kv_heads: int,27    head_dim: int,28    arch: str = "llama",29    v_trans: bool = True,30) -> bytes:31    """Build a synthetic blob matching llama_state_get_data() format."""32    parts: list[bytes] = []33 34    # 1. Architecture string header35    parts.append(struct.pack("<I", len(arch)))36    parts.append(arch.encode("ascii"))37 38    # 2. KV stream header39    parts.append(struct.pack("<I", 1))  # n_stream = 140    parts.append(struct.pack("<I", n_cells))  # cell_count41 42    # 3. Cell metadata: (pos:i32, n_seq:u32, seq_id:i32) per cell43    for i in range(n_cells):44        parts.append(struct.pack("<i", i))  # pos45        parts.append(struct.pack("<I", 1))  # n_seq_id = 146        parts.append(struct.pack("<i", 0))  # seq_id = 047 48    # 4. Data section header49    parts.append(struct.pack("<I", 1 if v_trans else 0))  # v_trans50    parts.append(struct.pack("<I", n_layers))51 52    n_embd_kv = n_kv_heads * head_dim53    row_size = n_embd_kv * 2  # fp1654 55    # 5. K layers56    for _ in range(n_layers):57        parts.append(struct.pack("<i", GGML_TYPE_F16))  # type_k58        parts.append(struct.pack("<Q", row_size))  # row_size_k59        data = np.random.randn(n_cells * n_embd_kv).astype(np.float16)60        parts.append(data.tobytes())61 62    # 6. V layers63    for _ in range(n_layers):64        parts.append(struct.pack("<i", GGML_TYPE_F16))  # type_v65        if v_trans:66            parts.append(struct.pack("<I", 2))  # el_size (fp16)67            parts.append(struct.pack("<I", n_embd_kv))  # n_embd_v_gqa68        else:69            parts.append(struct.pack("<Q", row_size))  # row_size_v70        data = np.random.randn(n_cells * n_embd_kv).astype(np.float16)71        parts.append(data.tobytes())72 73    return b"".join(parts)74 75 76class TestBlobParser:77    """Parse synthetic blobs in real llama_state_get_data format."""78 79    def test_parse_shape(self) -> None:80        blob = _make_blob(16, 32, 8, 128)81        result = parse_state_blob(blob, n_kv_heads=8, head_dim=128)82        assert result.keys.shape == (32, 8, 16, 128)83        assert result.values.shape == (32, 8, 16, 128)84 85    def test_parse_metadata(self) -> None:86        blob = _make_blob(8, 32, 8, 128)87        result = parse_state_blob(blob, n_kv_heads=8, head_dim=128)88        assert result.n_cells == 889        assert result.n_layers == 3290        assert result.arch == "llama"91        assert result.v_trans is True92        assert len(result.cells) == 893        assert result.cells[0].pos == 094        assert result.cells[7].pos == 795 96    def test_dtype_float16(self) -> None:97        blob = _make_blob(4, 28, 8, 128)98        result = parse_state_blob(blob, n_kv_heads=8, head_dim=128)99        assert result.keys.dtype == torch.float16100        assert result.values.dtype == torch.float16101 102    def test_non_transposed_v(self) -> None:103        blob = _make_blob(4, 28, 8, 128, v_trans=False)104        result = parse_state_blob(blob, n_kv_heads=8, head_dim=128)105        assert result.values.shape == (28, 8, 4, 128)106        assert result.v_trans is False107 108 109class TestBlobParserErrors:110    """Edge cases."""111 112    def test_zero_cells_raises(self) -> None:113        blob = struct.pack("<I", 5) + b"llama" + struct.pack("<II", 1, 0) + b"\x00" * 20114        with pytest.raises(BlobParseError, match="0 cells"):115            parse_state_blob(blob, n_kv_heads=8, head_dim=128)116 117    def test_truncated_blob_raises(self) -> None:118        blob = _make_blob(4, 28, 8, 128)119        with pytest.raises(BlobParseError):120            parse_state_blob(blob[:100], n_kv_heads=8, head_dim=128)121 122    def test_bad_arch_length_raises(self) -> None:123        blob = struct.pack("<I", 999) + b"x" * 100124        with pytest.raises(BlobParseError, match="too large"):125            parse_state_blob(blob, n_kv_heads=8, head_dim=128)126