feat(prefill): add bit-packed KakeyaLattice snapshots

deploy/install_prefill_worker_launchd.sh

@@ -17,6 +17,8 @@ CACHE_MODEL_ID="${KAKEYA_CACHE_MODEL_ID:-$KAKEYA_WORKER_MODEL}"
 MODEL_REVISION="${KAKEYA_MODEL_REVISION:-}"
 TOKENIZER_REVISION="${KAKEYA_TOKENIZER_REVISION:-}"
 QUANTIZATION="${KAKEYA_CACHE_QUANTIZATION:-4bit-mlx}"
+CACHE_FORMAT_VERSION="${KAKEYA_CACHE_FORMAT_VERSION:-kakeya-prefill-v2-zlib}"
+CACHE_COMPRESSION="${KAKEYA_CACHE_COMPRESSION:-zlib}"
 ROPE_HASH="${KAKEYA_ROPE_HASH:-}"
 SINK="${KAKEYA_WORKER_SINK:-4}"
 WINDOW="${KAKEYA_WORKER_WINDOW:-64}"
@@ -59,6 +61,8 @@ cat > "$PLIST" <<EOF
     <string>--model-revision</string><string>$MODEL_REVISION</string>
     <string>--tokenizer-revision</string><string>$TOKENIZER_REVISION</string>
     <string>--quantization</string><string>$QUANTIZATION</string>
+    <string>--cache-format-version</string><string>$CACHE_FORMAT_VERSION</string>
+    <string>--cache-compression</string><string>$CACHE_COMPRESSION</string>
     <string>--rope-hash</string><string>$ROPE_HASH</string>
     <string>--layer-geometry-hash</string><string>$KAKEYA_LAYER_GEOMETRY_HASH</string>
     <string>--tenant-id</string><string>$TENANT</string>

deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist

@@ -22,6 +22,8 @@
     <string>--cache-model-id</string><string>gemma-4-26B-A4B-it-mlx-4bit</string>
     <string>--model-revision</string><string>local-4bit-v1</string>
     <string>--tokenizer-revision</string><string>gemma4-v1</string>
+    <string>--cache-format-version</string><string>kakeya-prefill-v3-kl-d4-q38</string>
+    <string>--cache-compression</string><string>kakeyalattice-d4</string>
     <string>--cache-quantization</string><string>4bit-mlx</string>
     <string>--cache-kv-dtype</string><string>bfloat16</string>
     <string>--cache-block-tokens</string><string>64</string>

docs/ops/distributed-prefill-kv-network.md

@@ -119,6 +119,7 @@ chmod 600 ~/.kakeya/fleet.psk
 Install the worker:
 
 ```bash
+python -m pip install -r requirements-kakeyalattice.txt
 export KAKEYA_WORKER_REPO="$HOME/Kakeya-LLM-Inference-engine"
 export KAKEYA_WORKER_PYTHON="$HOME/kakeya-venv/bin/python"
 export KAKEYA_WORKER_MODEL="$HOME/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit"
@@ -132,6 +133,8 @@ export KAKEYA_LAYER_GEOMETRY_HASH="<same-value-as-primary>"
 export KAKEYA_WORKER_SINK="4"
 export KAKEYA_WORKER_WINDOW="2048"
 export KAKEYA_CACHE_BLOCK_TOKENS="64"
+export KAKEYA_CACHE_FORMAT_VERSION="kakeya-prefill-v3-kl-d4-q38"
+export KAKEYA_CACHE_COMPRESSION="kakeyalattice-d4"
 export KAKEYA_WORKER_NETWORK="thunderbolt"
 export KAKEYA_WORKER_PRIORITY="100"
 export KAKEYA_WORKER_RTT_MS="0.55"
@@ -153,6 +156,12 @@ The primary must use the same compatibility and auth values:
 --cache-replication-factor 1
 ```
 
+For bit-packed KakeyaLattice snapshots, replace the primary compression line
+with `--cache-compression kakeyalattice-d4` and set
+`--cache-format-version kakeya-prefill-v3-kl-d4-q38`. Both nodes must install
+the pinned optional dependency from `requirements-kakeyalattice.txt`. D4 Q=38
+is lossy; live acceptance must validate output quality as well as byte savings.
+
 `--cache-peer` remains an emergency static override. Normal worker/cache
 selection is derived from compatible live capability cards and their TTL/load
 metrics.

inference_engine/distributed/capability.py

@@ -56,6 +56,7 @@ class CompressionCodec(enum.IntEnum):
     UNSPECIFIED = 0
     NONE = 1
     ZLIB = 2
+    KAKEYA_LATTICE_D4 = 3
 
 
 @dataclass(frozen=True)

inference_engine/distributed/prefill_compression.py

@@ -20,16 +20,25 @@ def compress_payload(
     raw = bytes(payload)
     if codec in (CompressionCodec.UNSPECIFIED, CompressionCodec.NONE):
         return raw
-    if codec != CompressionCodec.ZLIB:
+    if codec not in (
+        CompressionCodec.ZLIB,
+        CompressionCodec.KAKEYA_LATTICE_D4,
+    ):
         raise ValueError(f"unsupported compression codec {codec!r}")
     if not (0 <= level <= 9):
         raise ValueError("zlib level must be in [0, 9]")
-    compressed = zlib.compress(raw, level)
+    framed = raw
+    if codec == CompressionCodec.KAKEYA_LATTICE_D4:
+        from inference_engine.distributed.prefill_kakeyalattice import (
+            encode_snapshot,
+        )
+        framed = encode_snapshot(raw)
+    compressed = zlib.compress(framed, level)
     return _HEADER.pack(
         _MAGIC,
         int(codec),
         len(raw),
-        hashlib.sha256(raw).digest(),
+        hashlib.sha256(framed).digest(),
     ) + compressed
 
 
@@ -54,7 +63,10 @@ def decompress_payload(
         codec = CompressionCodec(raw_codec)
     except ValueError as exc:
         raise ValueError(f"unsupported compression codec {raw_codec}") from exc
-    if codec != CompressionCodec.ZLIB:
+    if codec not in (
+        CompressionCodec.ZLIB,
+        CompressionCodec.KAKEYA_LATTICE_D4,
+    ):
         raise ValueError(f"unsupported framed compression codec {codec.name}")
     decompressor = zlib.decompressobj()
     raw = decompressor.decompress(
@@ -69,12 +81,17 @@ def decompress_payload(
         raise ValueError("decompressed payload exceeds import budget")
     if decompressor.unused_data:
         raise ValueError("compressed prefill payload has trailing data")
+    if hashlib.sha256(raw).digest() != expected_sha:
+        raise ValueError("decompressed prefill payload checksum mismatch")
+    if codec == CompressionCodec.KAKEYA_LATTICE_D4:
+        from inference_engine.distributed.prefill_kakeyalattice import (
+            decode_snapshot,
+        )
+        raw = decode_snapshot(raw)
     if len(raw) != expected_size:
         raise ValueError(
             f"decompressed payload size {len(raw)} != expected {expected_size}",
         )
-    if hashlib.sha256(raw).digest() != expected_sha:
-        raise ValueError("decompressed prefill payload checksum mismatch")
     return raw
 
 

inference_engine/distributed/prefill_kakeyalattice.py

@@ -0,0 +1,213 @@
+"""Bit-packed KakeyaLattice D4 codec for portable MLX prefill snapshots."""
+from __future__ import annotations
+
+import json
+import struct
+from typing import Any
+
+import torch
+
+from inference_engine.backends.mlx.prefill_snapshot import _pack, _unpack
+from inference_engine.distributed.tensor_codec import (
+    from_proto_fields,
+    to_proto_fields,
+    torch_to_wire,
+    wire_to_torch,
+)
+
+_MAGIC = b"KPKL1"
+_HEADER_LEN = struct.Struct("<I")
+_TORCH_DTYPES = {
+    "uint8": torch.uint8,
+    "int8": torch.int8,
+    "int32": torch.int32,
+    "float16": torch.float16,
+}
+_WIRE_DTYPES = {
+    "bfloat16": torch.bfloat16,
+    "float16": torch.float16,
+    "float32": torch.float32,
+}
+
+
+def _imports():
+    try:
+        from kakeyalattice import V14KakeyaZamirLatticeGPU
+        from kakeyalattice.hf.bitpack import (
+            pack_lattice_codes,
+            unpack_lattice_codes,
+        )
+        from kakeyalattice.hf.quantized_cache import (
+            decode_from_indices,
+            encode_to_indices,
+        )
+    except (ImportError, ModuleNotFoundError) as exc:
+        raise RuntimeError(
+            "KakeyaLattice bit-packed prefill requires kakeyalattice>=1.6.1",
+        ) from exc
+    return (
+        V14KakeyaZamirLatticeGPU,
+        pack_lattice_codes,
+        unpack_lattice_codes,
+        encode_to_indices,
+        decode_from_indices,
+    )
+
+
+def _device() -> torch.device:
+    return torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+
+
+def encode_snapshot(payload: bytes, *, q_range: int = 38) -> bytes:
+    """Quantize every K/V tensor and serialize losslessly packed lattice codes."""
+    (
+        codec_cls,
+        pack_lattice_codes,
+        _unpack_lattice_codes,
+        encode_to_indices,
+        _decode_from_indices,
+    ) = _imports()
+    snapshot_metadata, tensors = _unpack(payload)
+    device = _device()
+    codecs: dict[int, Any] = {}
+    parts: list[bytes] = []
+    records: list[dict[str, Any]] = []
+
+    def add_part(tensor: torch.Tensor) -> dict[str, Any]:
+        value = tensor.detach().contiguous().cpu()
+        dtype = str(value.dtype).removeprefix("torch.")
+        raw = value.numpy().tobytes()
+        record = {
+            "offset": sum(len(part) for part in parts),
+            "length": len(raw),
+            "shape": list(value.shape),
+            "dtype": dtype,
+        }
+        parts.append(raw)
+        return record
+
+    for name, (wire, framework) in tensors.items():
+        if not name.startswith("layer."):
+            dtype, shape, data = to_proto_fields(wire)
+            records.append({
+                "name": name,
+                "kind": "raw",
+                "framework": framework,
+                "wire_dtype": dtype,
+                "wire_shape": shape,
+                "data": add_part(torch.frombuffer(bytearray(data), dtype=torch.uint8)),
+            })
+            continue
+        original = wire_to_torch(wire).to(device)
+        head_dim = int(original.shape[-1])
+        codec = codecs.get(head_dim)
+        if codec is None:
+            codec = codec_cls(D=head_dim, q_range=q_range, device=str(device))
+            codecs[head_dim] = codec
+        codes, norms, qmax = encode_to_indices(codec, original)
+        packed = pack_lattice_codes(codes, "d4", q_range)
+        records.append({
+            "name": name,
+            "kind": "kakeyalattice-d4",
+            "framework": framework,
+            "wire_dtype": wire.dtype,
+            "head_dim": head_dim,
+            "q_range": q_range,
+            "packed": {
+                "width": packed["width"],
+                "mode": packed["mode"],
+                "n_blocks": packed["n_blocks"],
+                "shape": list(packed["shape"]),
+                "bd": packed["bd"],
+                "variant": packed["variant"],
+                "buf": add_part(packed["buf"]),
+                "exc_idx": add_part(packed["exc_idx"]),
+                "exc_vals": add_part(packed["exc_vals"]),
+            },
+            "norms": add_part(norms),
+            "qmax": add_part(qmax),
+        })
+    header = json.dumps(
+        {"snapshot": snapshot_metadata, "tensors": records},
+        sort_keys=True,
+        separators=(",", ":"),
+    ).encode()
+    return _MAGIC + _HEADER_LEN.pack(len(header)) + header + b"".join(parts)
+
+
+def decode_snapshot(payload: bytes) -> bytes:
+    """Restore a standard KPKV1 snapshot from bit-packed lattice tensors."""
+    (
+        codec_cls,
+        _pack_lattice_codes,
+        unpack_lattice_codes,
+        _encode_to_indices,
+        decode_from_indices,
+    ) = _imports()
+    if not payload.startswith(_MAGIC) or len(payload) < len(_MAGIC) + 4:
+        raise ValueError("invalid KakeyaLattice prefill snapshot magic")
+    header_len = _HEADER_LEN.unpack(payload[len(_MAGIC):len(_MAGIC) + 4])[0]
+    start = len(_MAGIC) + 4
+    end = start + header_len
+    if end > len(payload):
+        raise ValueError("truncated KakeyaLattice prefill header")
+    metadata = json.loads(payload[start:end])
+    raw = memoryview(payload)[end:]
+    device = _device()
+    codecs: dict[int, Any] = {}
+
+    def read_part(record: dict[str, Any], *, device_: torch.device | None = None):
+        offset = int(record["offset"])
+        part_end = offset + int(record["length"])
+        if offset < 0 or part_end > len(raw):
+            raise ValueError("truncated KakeyaLattice tensor component")
+        dtype = _TORCH_DTYPES[record["dtype"]]
+        if part_end == offset:
+            tensor = torch.empty(record["shape"], dtype=dtype)
+        else:
+            tensor = torch.frombuffer(
+                bytearray(raw[offset:part_end]),
+                dtype=dtype,
+            ).reshape(record["shape"])
+        return tensor.to(device_) if device_ is not None else tensor
+
+    tensors = []
+    for record in metadata["tensors"]:
+        if record["kind"] == "raw":
+            data = read_part(record["data"]).numpy().tobytes()
+            wire = from_proto_fields(
+                record["wire_dtype"],
+                record["wire_shape"],
+                data,
+            )
+        else:
+            head_dim = int(record["head_dim"])
+            q_range = int(record["q_range"])
+            codec = codecs.get(head_dim)
+            if codec is None:
+                codec = codec_cls(D=head_dim, q_range=q_range, device=str(device))
+                codecs[head_dim] = codec
+            packed_record = record["packed"]
+            packed = {
+                "width": int(packed_record["width"]),
+                "mode": packed_record["mode"],
+                "n_blocks": int(packed_record["n_blocks"]),
+                "shape": tuple(int(v) for v in packed_record["shape"]),
+                "bd": int(packed_record["bd"]),
+                "q_range": q_range,
+                "variant": packed_record["variant"],
+                "buf": read_part(packed_record["buf"], device_=device),
+                "exc_idx": read_part(packed_record["exc_idx"], device_=device),
+                "exc_vals": read_part(packed_record["exc_vals"], device_=device),
+            }
+            codes = unpack_lattice_codes(packed)
+            restored = decode_from_indices(
+                codec,
+                codes,
+                read_part(record["norms"], device_=device),
+                read_part(record["qmax"], device_=device),
+                out_dtype=_WIRE_DTYPES[record["wire_dtype"]],
+            )
+            wire = torch_to_wire(restored.cpu())
+        tensors.append((record["name"], wire, record["framework"]))
+    return _pack(tensors, metadata=metadata["snapshot"])