File size: 2,683 Bytes
19d9f3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | from __future__ import annotations
import struct
import torch
from wal_tat.runtime.cache import (
HARDWARE_BASE_HEADER, HARDWARE_BASE_MAGIC, build_hardware_cache,
load_hardware_base_cache,
)
def test_load_minimal_hardware_cache(tmp_path):
rows, columns, groups = 1, 128, 1
payloads = (
bytes(columns // 4), bytes(groups * 8), bytes(groups),
torch.ones((rows, groups), dtype=torch.float16).numpy().tobytes(),
torch.zeros((rows, groups), dtype=torch.float16).numpy().tobytes(),
)
header = HARDWARE_BASE_HEADER.pack(
HARDWARE_BASE_MAGIC, 1, rows, columns, groups,
*(len(value) for value in payloads),
)
path = tmp_path / "matrix.walhw"
path.write_bytes(header + b"".join(payloads))
loaded = load_hardware_base_cache(path, torch.device("cpu"))
assert loaded[-2:] == (rows, columns)
assert loaded[0].shape == (rows, columns // 4)
assert loaded[3].item() == 1.0
def test_rejects_truncated_hardware_cache(tmp_path):
path = tmp_path / "bad.walhw"
path.write_bytes(struct.pack("<Q", 1))
try:
load_hardware_base_cache(path, torch.device("cpu"))
except ValueError as error:
assert "short" in str(error)
else:
raise AssertionError("truncated cache accepted")
def test_build_and_attest_minimal_cache(tmp_path):
checkpoint = tmp_path / "checkpoint"
base = checkpoint / "base"
base.mkdir(parents=True)
(checkpoint / "manifest.json").write_text("{}")
source = base / "m.wal"
source.write_bytes(b"source")
class Endpoint:
@staticmethod
def _read_matrix_header(path):
return {"rows": 1, "columns": 128}
class Reference:
@staticmethod
def load_manifests(root):
return root, {}, base, {"matrices": [{
"name": "m.weight", "file": "m.wal", "sha256": "source-sha",
}]}
@staticmethod
def bundled_runtime(root):
return Endpoint(), object()
@staticmethod
def _iter_t3_sparse_rows(endpoint, path, row_chunk):
ternary = torch.zeros((1, 128), dtype=torch.int8)
sparse = torch.zeros((1, 1, 128), dtype=torch.int8)
sparse[0, 0, :8] = 1
yield 0, ternary, sparse, torch.ones((1, 1)), torch.ones((1, 1))
output = tmp_path / "cache"
report = build_hardware_cache(checkpoint, output, Reference())
assert report["matrix_count"] == 1
assert (output / "manifest.json").is_file()
assert (output / "attestation.json").is_file()
loaded = load_hardware_base_cache(output / "m.wal.walhw", torch.device("cpu"))
assert loaded[-2:] == (1, 128)
|