File size: 7,780 Bytes
2c76aad | 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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """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())
|