1"""Serialization of caches (safetensors) with a compatibility manifest; the loader rejects incompatible manifests2before any model execution. Independent of mlx_lm's save_prompt_cache so that the schema is under WP30 control."""3from __future__ import annotations4import json, os, gzip, tempfile5import mlx.core as mx6from mlx_lm.models.cache import KVCache, QuantizedKVCache7from .sparse import SparseKVCache, SparseQuantizedKVCache, _arrays_of8from ..util import sha256_text, utc_now910SCHEMA_VERSION = "wp30-cache-2"111213class IncompatibleCache(Exception):14 pass151617def _layer_desc(i, c):18 k, v = c.keys_and_values()19 if hasattr(c, "bits"):20 return {"layer": i, "cls": type(c).__name__, "k_shapes": [list(x.shape) for x in k], "v_shapes": [list(x.shape) for x in v],21 "dtypes": [str(x.dtype) for x in k], "bits": c.bits, "group_size": c.group_size,22 "stored_positions": int(k[0].shape[2]), "logical_offset": int(c.offset)}23 return {"layer": i, "cls": type(c).__name__, "k_shape": list(k.shape), "v_shape": list(v.shape), "dtype": str(k.dtype),24 "stored_positions": int(k.shape[2]), "logical_offset": int(c.offset)}252627def build_manifest(model_info, prefix_token_ids: list[int], cache: list, document_id: str, method_id: str,28 compression_config: dict, extra: dict | None = None) -> dict:29 m = {30 "schema_version": SCHEMA_VERSION, "created_utc": utc_now(),31 "model_repo_id": model_info.repo_id, "model_revision": model_info.revision, "config_hash": model_info.config_hash,32 "weight_files_sha256": model_info.weight_files_sha256, "tokenizer_hash": model_info.tokenizer_hash,33 "chat_template_hash": model_info.chat_template_hash,34 "prefix_token_hash": sha256_text(",".join(map(str, prefix_token_ids))), "prefix_token_count": len(prefix_token_ids),35 "num_layers": len(cache), "num_kv_heads": model_info.num_kv_heads, "head_dim": model_info.head_dim,36 "rope_theta": model_info.rope_theta, "rope_scaling": model_info.rope_scaling, "sliding_window": model_info.sliding_window,37 "positions": "keys stored post-RoPE at original positions; logical_offset = original processed-token count",38 "causal_mask_semantics": "retained positions fully visible to new tokens; new tokens causal among themselves",39 "mlx_version": model_info.mlx_version, "mlx_lm_version": model_info.mlx_lm_version,40 "document_id": document_id, "method_id": method_id, "compression_config": compression_config,41 "layers": [_layer_desc(i, c) for i, c in enumerate(cache)],42 }43 if extra:44 m.update(extra)45 return m464748def _export(cache: list) -> tuple[dict, list]:49 arrays, metas = {}, []50 for i, c in enumerate(cache):51 k, v = c.keys_and_values()52 if hasattr(c, "bits"):53 for j, x in enumerate(k): arrays[f"L{i}.k{j}"] = x54 for j, x in enumerate(v): arrays[f"L{i}.v{j}"] = x55 metas.append({"cls": type(c).__name__, "offset": int(c.offset), "group_size": int(c.group_size), "bits": int(c.bits),56 "n_kept": int(k[0].shape[2])})57 else:58 arrays[f"L{i}.k"] = k; arrays[f"L{i}.v"] = v59 metas.append({"cls": type(c).__name__, "offset": int(c.offset), "n_kept": int(k.shape[2])})60 return arrays, metas616263def serialize_cache(cache: list, manifest: dict) -> bytes:64 arrays, metas = _export(cache)65 md = {"wp30_manifest": json.dumps(manifest, separators=(",", ":")), "wp30_layers": json.dumps(metas)}66 with tempfile.TemporaryDirectory() as d:67 p = os.path.join(d, "c.safetensors")68 mx.save_safetensors(p, arrays, md)69 with open(p, "rb") as f:70 return f.read()717273def check_manifest(manifest: dict, expected_model_info=None, expected_prefix_hash: str | None = None):74 if manifest.get("schema_version") != SCHEMA_VERSION:75 raise IncompatibleCache(f"schema {manifest.get('schema_version')} != {SCHEMA_VERSION}")76 if expected_model_info is not None:77 pairs = [("config_hash", "config_hash"), ("tokenizer_hash", "tokenizer_hash"), ("chat_template_hash", "chat_template_hash"), ("model_revision", "revision")]78 for mk, ak in pairs:79 if manifest[mk] != getattr(expected_model_info, ak):80 raise IncompatibleCache(f"{mk} mismatch: cache from {manifest['model_repo_id']}@{manifest['model_revision']}, receiver {expected_model_info.repo_id}@{expected_model_info.revision}")81 if (manifest["num_layers"], manifest["num_kv_heads"], manifest["head_dim"]) != (expected_model_info.num_layers, expected_model_info.num_kv_heads, expected_model_info.head_dim):82 raise IncompatibleCache("layer/head geometry mismatch")83 if expected_prefix_hash is not None and manifest["prefix_token_hash"] != expected_prefix_hash:84 raise IncompatibleCache("prefix token hash mismatch")858687def deserialize_cache(payload: bytes, expected_model_info=None, expected_prefix_hash: str | None = None, check: bool = True) -> tuple[list, dict]:88 with tempfile.TemporaryDirectory() as d:89 p = os.path.join(d, "c.safetensors")90 with open(p, "wb") as f:91 f.write(payload)92 arrays, md = mx.load(p, return_metadata=True)93 mx.eval(list(arrays.values()))94 manifest = json.loads(md["wp30_manifest"]); metas = json.loads(md["wp30_layers"])95 if check:96 check_manifest(manifest, expected_model_info, expected_prefix_hash)97 cache = []98 for i, m in enumerate(metas):99 cls = m["cls"]100 if cls in ("SparseQuantizedKVCache", "QuantizedKVCache"):101 c = SparseQuantizedKVCache(group_size=m["group_size"], bits=m["bits"])102 c.keys = tuple(arrays[f"L{i}.k{j}"] for j in range(3)); c.values = tuple(arrays[f"L{i}.v{j}"] for j in range(3))103 else:104 c = SparseKVCache(); c.keys, c.values = arrays[f"L{i}.k"], arrays[f"L{i}.v"]105 c.offset, c.n_kept = m["offset"], m["n_kept"]106 cache.append(c)107 return cache, manifest108109110def codec_compress(b: bytes, codec: str, level: int | None = None) -> tuple[bytes, float]:111 import time as _t112 t0 = _t.perf_counter()113 if codec == "none":114 out = b115 elif codec == "gzip":116 out = gzip.compress(b, compresslevel=6 if level is None else level)117 elif codec == "zstd":118 import zstandard as zstd119 out = zstd.ZstdCompressor(level=3 if level is None else level).compress(b)120 else:121 raise ValueError(codec)122 return out, (_t.perf_counter() - t0) * 1000.0123124125def codec_decompress(b: bytes, codec: str) -> tuple[bytes, float]:126 import time as _t127 t0 = _t.perf_counter()128 if codec == "none":129 out = b130 elif codec == "gzip":131 out = gzip.decompress(b)132 elif codec == "zstd":133 import zstandard as zstd134 out = zstd.ZstdDecompressor().decompress(b)135 else:136 raise ValueError(codec)137 return out, (_t.perf_counter() - t0) * 1000.0138