"""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 collections.abc import Iterable 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", } OPTIONAL_FIELDS = {"tensor_manifest"} KNOWN_FIELDS = REQUIRED_FIELDS | OPTIONAL_FIELDS ALLOWED_DTYPES = {"f16", "bf16", "f32", "i8", "u8"} DTYPE_BYTES = {"f16": 2, "bf16": 2, "f32": 4, "i8": 1, "u8": 1} ALLOWED_LAYOUTS = {"layer-major-k-then-v"} MAX_METADATA_BYTES = 1_048_576 MANIFEST_FIELDS = {"version", "byte_order", "batch_size", "cache_class", "attention_type", "tensors"} TENSOR_FIELDS = {"name", "kind", "layer", "shape", "strides", "offset", "nbytes", "dtype"} 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 _exact_fields(value: Mapping[str, Any], expected: set[str], label: str) -> None: missing = expected.difference(value) extra = set(value).difference(expected) if missing or extra: raise CacheFormatError(f"{label} fields differ: missing={sorted(missing)}, extra={sorted(extra)}") def _positive_integer(value: Any, label: str) -> int: if not isinstance(value, int) or isinstance(value, bool) or value <= 0: raise CacheFormatError(f"{label} must be a positive integer") return value def _contiguous_byte_strides(shape: list[int], item_bytes: int) -> list[int]: strides = [item_bytes] * len(shape) for index in range(len(shape) - 2, -1, -1): strides[index] = strides[index + 1] * shape[index + 1] return strides def _validate_tensor_manifest(metadata: Mapping[str, Any], manifest: Any) -> int: if not isinstance(manifest, Mapping): raise CacheFormatError("tensor_manifest must be a mapping") manifest = dict(manifest) _exact_fields(manifest, MANIFEST_FIELDS, "tensor_manifest") if manifest["version"] != 1: raise CacheFormatError("tensor_manifest version must be 1") if manifest["byte_order"] not in {"little", "big"}: raise CacheFormatError("tensor_manifest byte_order must be little or big") batch_size = _positive_integer(manifest["batch_size"], "tensor_manifest batch_size") for field in ("cache_class", "attention_type"): if not isinstance(manifest[field], str) or not manifest[field].strip(): raise CacheFormatError(f"tensor_manifest {field} must be a non-empty string") tensors = manifest["tensors"] if not isinstance(tensors, list): raise CacheFormatError("tensor_manifest tensors must be a list") expected_count = metadata["layers"] * 2 if len(tensors) != expected_count: raise CacheFormatError(f"tensor_manifest must contain exactly {expected_count} key/value tensors") expected_shape = [batch_size, metadata["kv_heads"], metadata["sequence_length"], metadata["head_dim"]] expected_offset = 0 seen_names: set[str] = set() seen_slots: set[tuple[int, str]] = set() for index, raw_tensor in enumerate(tensors): if not isinstance(raw_tensor, Mapping): raise CacheFormatError(f"tensor_manifest tensor {index} must be a mapping") tensor = dict(raw_tensor) _exact_fields(tensor, TENSOR_FIELDS, f"tensor_manifest tensor {index}") name = tensor["name"] if not isinstance(name, str) or not name: raise CacheFormatError(f"tensor_manifest tensor {index} name must be non-empty") if name in seen_names: raise CacheFormatError(f"tensor_manifest tensor name is duplicated: {name}") seen_names.add(name) kind = tensor["kind"] if kind not in {"key", "value"}: raise CacheFormatError(f"tensor_manifest tensor {index} kind must be key or value") layer = tensor["layer"] if not isinstance(layer, int) or isinstance(layer, bool) or not 0 <= layer < metadata["layers"]: raise CacheFormatError(f"tensor_manifest tensor {index} layer is out of range") slot = (layer, kind) if slot in seen_slots: raise CacheFormatError(f"tensor_manifest duplicates layer {layer} {kind}") seen_slots.add(slot) shape = tensor["shape"] if shape != expected_shape: raise CacheFormatError(f"tensor_manifest tensor {index} shape must equal {expected_shape}") dtype = tensor["dtype"] if dtype != metadata["dtype"]: raise CacheFormatError(f"tensor_manifest tensor {index} dtype must match metadata dtype") expected_nbytes = DTYPE_BYTES[dtype] for dimension in shape: expected_nbytes *= _positive_integer(dimension, f"tensor_manifest tensor {index} shape dimension") if tensor["nbytes"] != expected_nbytes: raise CacheFormatError(f"tensor_manifest tensor {index} nbytes must equal {expected_nbytes}") expected_strides = _contiguous_byte_strides(shape, DTYPE_BYTES[dtype]) if tensor["strides"] != expected_strides: raise CacheFormatError( f"tensor_manifest tensor {index} must use C-contiguous byte strides {expected_strides}" ) if tensor["offset"] != expected_offset: raise CacheFormatError( f"tensor_manifest tensor {index} offsets must be contiguous; expected {expected_offset}" ) expected_offset += expected_nbytes expected_slots = {(layer, kind) for layer in range(metadata["layers"]) for kind in ("key", "value")} if seen_slots != expected_slots: raise CacheFormatError("tensor_manifest must contain one key and one value tensor for every layer") return expected_offset def _validate_payload_length(metadata: Mapping[str, Any], payload_length: int) -> None: manifest = metadata.get("tensor_manifest") if manifest is None: return expected = _validate_tensor_manifest(metadata, manifest) if payload_length != expected: raise CacheFormatError( f"tensor manifest declares {expected} payload bytes but container has {payload_length}" ) 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(KNOWN_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") if "tensor_manifest" in value: _validate_tensor_manifest(value, value["tensor_manifest"]) 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") return write_cache_stream(Path(path), metadata, (payload,)) def write_cache_stream(path: str | os.PathLike[str], metadata: Mapping[str, Any], chunks: Iterable[bytes | bytearray | memoryview]) -> None: destination = Path(path) metadata_dict = _plain_dict(metadata) metadata_bytes = _metadata_bytes(metadata_dict) destination.parent.mkdir(parents=True, exist_ok=True) 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 placeholder = PREFIX.pack(MAGIC, len(metadata_bytes), 0, b"\x00" * 32) temporary.write(placeholder) temporary.write(metadata_bytes) sha = hashlib.sha256() length = 0 for chunk in chunks: if not isinstance(chunk, (bytes, bytearray, memoryview)): raise CacheFormatError("chunk must be bytes, bytearray, or memoryview") chunk_bytes = bytes(chunk) temporary.write(chunk_bytes) sha.update(chunk_bytes) length += len(chunk_bytes) _validate_payload_length(metadata_dict, length) digest = sha.digest() temporary.seek(0) temporary.write(PREFIX.pack(MAGIC, len(metadata_bytes), length, digest)) 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") _validate_payload_length(metadata, len(payload)) return metadata, payload, actual_digest.hex() def read_cache(path: str | os.PathLike[str], expected_identity: Mapping[str, Any] | None = None) -> tuple[dict[str, Any], bytes]: metadata, payload, _ = _decode(path) if expected_identity is not None: if not isinstance(expected_identity, Mapping): raise CacheCompatibilityError("expected_identity must be a mapping") expected = dict(expected_identity) unknown = set(expected).difference(KNOWN_FIELDS) if unknown: raise CacheCompatibilityError(f"expected_identity has unknown fields: {sorted(unknown)}") try: _plain_dict({**metadata, **expected}) except CacheFormatError as error: raise CacheCompatibilityError(f"expected_identity is invalid: {error}") from error differences = [field for field in sorted(expected) 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) manifest = metadata.get("tensor_manifest") receipt = { "format": "KVC1", "metadata": metadata, "payload_bytes": len(payload), "payload_sha256": digest, } if manifest is not None: receipt["tensor_manifest_verified"] = True receipt["tensor_count"] = len(manifest["tensors"]) return receipt 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())