Add accepted KVC1 reference implementation

reference/kvcache.py

@@ -0,0 +1,194 @@
+"""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())