"""Portable, fail-closed raw KV-cache container. KVC1 stores already-materialized key/value tensor bytes plus the exact model, tokenizer, tensor geometry, RoPE, dtype, layout, and sequence-position identity needed to decide whether a runtime may safely reuse them. Cross-architecture translation is deliberately out of scope for this byte container. """ from __future__ import annotations import argparse import hashlib import hmac import json import os from pathlib import Path import struct import tempfile from typing import Any, Mapping MAGIC = b"KVC1" PREFIX = struct.Struct(">4sIQ32s") REQUIRED_FIELDS = { "model_revision", "tokenizer_sha256", "rope_theta", "layers", "kv_heads", "head_dim", "dtype", "layout", "sequence_start", "sequence_length", } ALLOWED_DTYPES = {"f16", "bf16", "f32", "i8", "u8"} ALLOWED_LAYOUTS = {"layer-major-k-then-v"} MAX_METADATA_BYTES = 1_048_576 class CacheFormatError(ValueError): """The container is malformed, incomplete, corrupt, or unsupported.""" class CacheCompatibilityError(ValueError): """The container is valid but does not match the requested runtime.""" def _plain_dict(metadata: Mapping[str, Any]) -> dict[str, Any]: if not isinstance(metadata, Mapping): raise CacheFormatError("metadata must be a mapping") value = dict(metadata) missing = REQUIRED_FIELDS.difference(value) extra = set(value).difference(REQUIRED_FIELDS) if missing or extra: raise CacheFormatError(f"metadata fields differ: missing={sorted(missing)}, extra={sorted(extra)}") if not isinstance(value["model_revision"], str) or "@" not in value["model_revision"]: raise CacheFormatError("model_revision must identify an immutable revision") tokenizer_hash = value["tokenizer_sha256"] if not isinstance(tokenizer_hash, str) or len(tokenizer_hash) != 64: raise CacheFormatError("tokenizer_sha256 must contain 64 hexadecimal characters") try: int(tokenizer_hash, 16) except ValueError as error: raise CacheFormatError("tokenizer_sha256 is not hexadecimal") from error if not isinstance(value["rope_theta"], (int, float)) or isinstance(value["rope_theta"], bool) or value["rope_theta"] <= 0: raise CacheFormatError("rope_theta must be positive") for field in ("layers", "kv_heads", "head_dim", "sequence_length"): if not isinstance(value[field], int) or isinstance(value[field], bool) or value[field] <= 0: raise CacheFormatError(f"{field} must be a positive integer") if not isinstance(value["sequence_start"], int) or isinstance(value["sequence_start"], bool) or value["sequence_start"] < 0: raise CacheFormatError("sequence_start must be a non-negative integer") if value["dtype"] not in ALLOWED_DTYPES: raise CacheFormatError("unsupported dtype") if value["layout"] not in ALLOWED_LAYOUTS: raise CacheFormatError("unsupported layout") return value def _metadata_bytes(metadata: Mapping[str, Any]) -> bytes: try: encoded = json.dumps(_plain_dict(metadata), sort_keys=True, separators=(",", ":"), allow_nan=False).encode("utf-8") except (TypeError, ValueError) as error: if isinstance(error, CacheFormatError): raise raise CacheFormatError("metadata is not canonical JSON") from error if len(encoded) > MAX_METADATA_BYTES: raise CacheFormatError("metadata is too large") return encoded def write_cache(path: str | os.PathLike[str], metadata: Mapping[str, Any], payload: bytes) -> None: """Atomically publish one KVC1 generation.""" if type(payload) is not bytes: raise CacheFormatError("payload must be raw bytes") destination = Path(path) destination.parent.mkdir(parents=True, exist_ok=True) metadata_bytes = _metadata_bytes(metadata) digest = hashlib.sha256(payload).digest() prefix = PREFIX.pack(MAGIC, len(metadata_bytes), len(payload), digest) temporary_name: str | None = None try: with tempfile.NamedTemporaryFile( mode="wb", prefix=f".{destination.name}.", suffix=".tmp", dir=destination.parent, delete=False, ) as temporary: temporary_name = temporary.name temporary.write(prefix) temporary.write(metadata_bytes) temporary.write(payload) temporary.flush() os.fsync(temporary.fileno()) os.replace(temporary_name, destination) temporary_name = None finally: if temporary_name is not None: try: os.unlink(temporary_name) except FileNotFoundError: pass def _decode(path: str | os.PathLike[str]) -> tuple[dict[str, Any], bytes, str]: try: raw = Path(path).read_bytes() except OSError as error: raise CacheFormatError(f"cache could not be read: {error}") from error if len(raw) < PREFIX.size: raise CacheFormatError("container is truncated") try: magic, metadata_length, payload_length, expected_digest = PREFIX.unpack_from(raw) except struct.error as error: raise CacheFormatError("container prefix is malformed") from error if magic != MAGIC: raise CacheFormatError("unsupported container magic or version") if metadata_length == 0 or metadata_length > MAX_METADATA_BYTES: raise CacheFormatError("metadata length is invalid") expected_length = PREFIX.size + metadata_length + payload_length if len(raw) != expected_length: raise CacheFormatError("container length does not match its header") metadata_raw = raw[PREFIX.size:PREFIX.size + metadata_length] payload = raw[PREFIX.size + metadata_length:] try: decoded = json.loads(metadata_raw.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as error: raise CacheFormatError("metadata is not valid UTF-8 JSON") from error metadata = _plain_dict(decoded) if _metadata_bytes(metadata) != metadata_raw: raise CacheFormatError("metadata is not in canonical form") actual_digest = hashlib.sha256(payload).digest() if not hmac.compare_digest(actual_digest, expected_digest): raise CacheFormatError("payload checksum mismatch") return metadata, payload, actual_digest.hex() def read_cache(path: str | os.PathLike[str], expected_metadata: Mapping[str, Any] | None = None) -> tuple[dict[str, Any], bytes]: metadata, payload, _ = _decode(path) if expected_metadata is not None: expected = _plain_dict(expected_metadata) differences = [field for field in sorted(REQUIRED_FIELDS) if metadata[field] != expected[field]] if differences: raise CacheCompatibilityError(f"cache is incompatible: {', '.join(differences)}") return metadata, payload def inspect_cache(path: str | os.PathLike[str]) -> dict[str, Any]: metadata, payload, digest = _decode(path) return { "format": "KVC1", "metadata": metadata, "payload_bytes": len(payload), "payload_sha256": digest, } def main() -> int: parser = argparse.ArgumentParser(description="Inspect a portable raw KV-cache container.") subparsers = parser.add_subparsers(dest="command", required=True) inspect_parser = subparsers.add_parser("inspect") inspect_parser.add_argument("path") args = parser.parse_args() if args.command == "inspect": try: print(json.dumps(inspect_cache(args.path), sort_keys=True)) except (CacheFormatError, CacheCompatibilityError) as error: parser.exit(1, f"kvcache: {error}\n") return 0 if __name__ == "__main__": raise SystemExit(main())