diff --git a/deploy/install_prefill_worker_launchd.sh b/deploy/install_prefill_worker_launchd.sh
index f0bf1bcd..8c62a5d3 100755
--- a/deploy/install_prefill_worker_launchd.sh
+++ b/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>
diff --git a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist
index 7a5a3a58..2751dc04 100644
--- a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist
+++ b/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>
diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md
index baa090f8..1c7df1bb 100644
--- a/docs/ops/distributed-prefill-kv-network.md
+++ b/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.
diff --git a/inference_engine/distributed/capability.py b/inference_engine/distributed/capability.py
index b4a4ef83..7ee9d026 100644
--- a/inference_engine/distributed/capability.py
+++ b/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)
diff --git a/inference_engine/distributed/prefill_compression.py b/inference_engine/distributed/prefill_compression.py
index 23f7ef2c..8248440e 100644
--- a/inference_engine/distributed/prefill_compression.py
+++ b/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
 
 
diff --git a/inference_engine/distributed/prefill_kakeyalattice.py b/inference_engine/distributed/prefill_kakeyalattice.py
new file mode 100644
index 00000000..5abac6b1
--- /dev/null
+++ b/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"])
diff --git a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py
index 7fef240f..6d700b9d 100644
--- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py
+++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py
@@ -24,7 +24,7 @@
 
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1bkakeya/v1/distributed.proto\x12\tkakeya.v1\"}\n\x0fModelCapability\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\'\n\x04role\x18\x02 \x01(\x0e\x32\x19.kakeya.v1.CapabilityRole\x12\x14\n\x0cquantization\x18\x03 \x01(\t\x12\x19\n\x11tokens_per_second\x18\x04 \x01(\x01\"\x83\x03\n\x0eNodeCapability\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12\x14\n\x0cgrpc_address\x18\x02 \x01(\t\x12\x10\n\x08platform\x18\x03 \x01(\t\x12\x1c\n\x14unified_memory_bytes\x18\x04 \x01(\x04\x12\x13\n\x0bmlx_version\x18\x05 \x01(\t\x12*\n\x06models\x18\x06 \x03(\x0b\x32\x1a.kakeya.v1.ModelCapability\x12\x19\n\x11\x61nnounced_at_unix\x18\x07 \x01(\x01\x12\x13\n\x0bttl_seconds\x18\x08 \x01(\x01\x12\x14\n\x0cring_address\x18\t \x01(\t\x12*\n\x06\x63\x61\x63hes\x18\n \x03(\x0b\x32\x1a.kakeya.v1.CacheCapability\x12*\n\tendpoints\x18\x0b \x03(\x0b\x32\x17.kakeya.v1.NodeEndpoint\x12;\n\x0fprefill_workers\x18\x0c \x03(\x0b\x32\".kakeya.v1.PrefillWorkerCapability\"[\n\x0cNodeEndpoint\x12\x0f\n\x07\x61\x64\x64ress\x18\x01 \x01(\t\x12\x0f\n\x07network\x18\x02 \x01(\t\x12\x10\n\x08priority\x18\x03 \x01(\r\x12\x17\n\x0fmeasured_rtt_ms\x18\x04 \x01(\x01\"\xad\x02\n\x12\x43\x61\x63heCompatibility\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x16\n\x0emodel_revision\x18\x02 \x01(\t\x12\x1a\n\x12tokenizer_revision\x18\x03 \x01(\t\x12\x1c\n\x14\x63\x61\x63he_format_version\x18\x04 \x01(\t\x12\x14\n\x0cquantization\x18\x05 \x01(\t\x12\x11\n\trope_hash\x18\x06 \x01(\t\x12\x1b\n\x13layer_geometry_hash\x18\x07 \x01(\t\x12\x10\n\x08kv_dtype\x18\x08 \x01(\t\x12\x19\n\x11\x62lock_size_tokens\x18\t \x01(\r\x12\x18\n\x10tenant_namespace\x18\n \x01(\t\x12\x11\n\tsink_size\x18\x0b \x01(\r\x12\x13\n\x0bwindow_size\x18\x0c \x01(\r\"\xcd\x02\n\x0f\x43\x61\x63heCapability\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x15\n\rcache_address\x18\x02 \x01(\t\x12\x18\n\x10\x63\x61\x63he_bytes_used\x18\x03 \x01(\x04\x12\x18\n\x10\x63\x61\x63he_bytes_free\x18\x04 \x01(\x04\x12\x13\n\x0b\x65ntry_count\x18\x05 \x01(\x04\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x06 \x01(\x04\x12\x0c\n\x04load\x18\x07 \x01(\x01\x12\x15\n\rtokens_served\x18\x08 \x01(\x04\x12\x14\n\x0c\x62loom_filter\x18\t \x01(\x0c\x12\x38\n\x13\x64\x65\x66\x61ult_compression\x18\n \x01(\x0e\x32\x1b.kakeya.v1.CompressionCodec\x12\x1a\n\x12replication_factor\x18\x0b \x01(\r\"\xae\x02\n\x17PrefillWorkerCapability\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x16\n\x0eworker_address\x18\x02 \x01(\t\x12\x1b\n\x13max_concurrent_jobs\x18\x03 \x01(\r\x12\x15\n\rinflight_jobs\x18\x04 \x01(\r\x12\x13\n\x0bqueued_jobs\x18\x05 \x01(\r\x12\x0c\n\x04load\x18\x06 \x01(\x01\x12!\n\x19tokens_per_second_prefill\x18\x07 \x01(\x01\x12\x16\n\x0eram_bytes_free\x18\x08 \x01(\x04\x12\x1c\n\x14\x61\x63\x63\x65pts_compute_jobs\x18\t \x01(\x08\x12\x15\n\rqueued_tokens\x18\n \x01(\x04\"M\n\x1b\x45xchangeCapabilitiesRequest\x12.\n\x0bknown_nodes\x18\x01 \x03(\x0b\x32\x19.kakeya.v1.NodeCapability\"N\n\x1c\x45xchangeCapabilitiesResponse\x12.\n\x0bknown_nodes\x18\x01 \x03(\x0b\x32\x19.kakeya.v1.NodeCapability\"\x1a\n\x18GetNodeCapabilityRequest\"D\n\x19GetNodeCapabilityResponse\x12\'\n\x04node\x18\x01 \x01(\x0b\x32\x19.kakeya.v1.NodeCapability\"N\n\x16GetCacheSummaryRequest\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\"V\n\x17GetCacheSummaryResponse\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12*\n\x06\x63\x61\x63hes\x18\x02 \x03(\x0b\x32\x1a.kakeya.v1.CacheCapability\"a\n\x13LookupPrefixRequest\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x14\n\x0c\x62lock_hashes\x18\x02 \x03(\x0c\"\xcf\x01\n\x14LookupPrefixResponse\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12\x17\n\x0fhit_block_count\x18\x02 \x01(\r\x12\x17\n\x0fhit_token_count\x18\x03 \x01(\x04\x12\x16\n\x0etransfer_bytes\x18\x04 \x01(\x04\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x05 \x01(\x04\x12\x10\n\x08lease_id\x18\x06 \x01(\t\x12\x1d\n\x15lease_expires_at_unix\x18\x07 \x01(\x01\x12\x16\n\x0epayload_sha256\x18\x08 \x01(\x0c\"&\n\x12\x46\x65tchBlocksRequest\x12\x10\n\x08lease_id\x18\x01 \x01(\t\"\xb7\x01\n\x13\x46\x65tchBlocksResponse\x12\x12\n\nblock_hash\x18\x01 \x01(\x0c\x12\x13\n\x0b\x62lock_index\x18\x02 \x01(\r\x12\x13\n\x0btoken_count\x18\x03 \x01(\r\x12\x13\n\x0b\x63hunk_index\x18\x04 \x01(\r\x12\x14\n\x0ctotal_chunks\x18\x05 \x01(\r\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x14\n\x0c\x62lock_sha256\x18\x07 \x01(\x0c\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x08 \x01(\x04\"\xed\x01\n\x13PublishBlockRequest\x12\x12\n\nblock_hash\x18\x01 \x01(\x0c\x12\x13\n\x0b\x62lock_index\x18\x02 \x01(\r\x12\x13\n\x0btoken_count\x18\x03 \x01(\r\x12\x13\n\x0b\x63hunk_index\x18\x04 \x01(\r\x12\x14\n\x0ctotal_chunks\x18\x05 \x01(\r\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x14\n\x0c\x62lock_sha256\x18\x07 \x01(\x0c\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x08 \x01(\x04\x12\x34\n\rcompatibility\x18\t \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\";\n\x14PublishBlockResponse\x12\x0e\n\x06stored\x18\x01 \x01(\x08\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x02 \x01(\x04\"\xf0\x01\n\x17SubmitPrefillJobRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\x12\x34\n\rcompatibility\x18\x03 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x11\n\ttoken_ids\x18\x04 \x03(\r\x12\x14\n\x0c\x62lock_hashes\x18\x05 \x03(\x0c\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x06 \x01(\r\x12:\n\x15preferred_compression\x18\x07 \x01(\x0e\x32\x1b.kakeya.v1.CompressionCodec\"\x85\x01\n\x18SubmitPrefillJobResponse\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12+\n\x06status\x18\x02 \x01(\x0e\x32\x1b.kakeya.v1.PrefillJobStatus\x12\x16\n\x0eworker_node_id\x18\x03 \x01(\t\x12\x14\n\x0cqueue_eta_ms\x18\x04 \x01(\x01\"?\n\x1aGetPrefillJobStatusRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\"\x8c\x02\n\x1bGetPrefillJobStatusResponse\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12+\n\x06status\x18\x02 \x01(\x0e\x32\x1b.kakeya.v1.PrefillJobStatus\x12\x17\n\x0ftokens_computed\x18\x03 \x01(\r\x12\x10\n\x08lease_id\x18\x04 \x01(\t\x12\x12\n\nblock_hash\x18\x05 \x01(\x0c\x12\x16\n\x0epayload_sha256\x18\x06 \x01(\x0c\x12\x16\n\x0etransfer_bytes\x18\x07 \x01(\x04\x12\x16\n\x0e\x66\x61ilure_reason\x18\x08 \x01(\t\x12\x12\n\ncompute_ms\x18\t \x01(\x01\x12\x15\n\rcache_address\x18\n \x01(\t\"<\n\x17\x43\x61ncelPrefillJobRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\"-\n\x18\x43\x61ncelPrefillJobResponse\x12\x11\n\tcancelled\x18\x01 \x01(\x08\"k\n\x13ProposeBlockRequest\x12\x1b\n\x13\x63ommitted_token_ids\x18\x01 \x03(\r\x12\x12\n\nblock_size\x18\x02 \x01(\r\x12\x11\n\tnum_steps\x18\x03 \x01(\r\x12\x10\n\x08model_id\x18\x04 \x01(\t\"y\n\x14ProposeBlockResponse\x12\x11\n\ttoken_ids\x18\x01 \x03(\r\x12\x17\n\x0f\x64iffusion_steps\x18\x02 \x01(\r\x12\x16\n\x0e\x66orward_passes\x18\x03 \x01(\r\x12\x1d\n\x15peak_activation_bytes\x18\x04 \x01(\x04\"4\n\x06Tensor\x12\r\n\x05\x64type\x18\x01 \x01(\t\x12\r\n\x05shape\x18\x02 \x03(\x03\x12\x0c\n\x04\x64\x61ta\x18\x03 \x01(\x0c\"T\n\x07LayerKV\x12\r\n\x05layer\x18\x01 \x01(\x05\x12\x1c\n\x01k\x18\x02 \x01(\x0b\x32\x11.kakeya.v1.Tensor\x12\x1c\n\x01v\x18\x03 \x01(\x0b\x32\x11.kakeya.v1.Tensor\"\x84\x01\n\x0eRestoreRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x12\n\nprompt_ids\x18\x02 \x03(\r\x12\x0c\n\x04sink\x18\x03 \x01(\r\x12\x0e\n\x06window\x18\x04 \x01(\r\x12\x1a\n\x12s5_exact_full_attn\x18\x05 \x01(\x08\x12\x10\n\x08model_id\x18\x06 \x01(\t\"f\n\x0fRestoreResponse\x12$\n\x08restored\x18\x01 \x03(\x0b\x32\x12.kakeya.v1.LayerKV\x12\x19\n\x11\x65victed_positions\x18\x02 \x03(\x05\x12\x12\n\nprompt_len\x18\x03 \x01(\r\"[\n\x12SeedContextRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x1e\n\x03\x61ux\x18\x02 \x03(\x0b\x32\x11.kakeya.v1.Tensor\x12\x11\n\tpositions\x18\x03 \x03(\x05\"*\n\x13SeedContextResponse\x12\x13\n\x0b\x63ontext_len\x18\x01 \x01(\r\"h\n\x11\x44raftBlockRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x16\n\x0e\x62onus_token_id\x18\x02 \x01(\r\x12\x13\n\x0b\x63ontext_len\x18\x03 \x01(\r\x12\x12\n\nblock_size\x18\x04 \x01(\r\"d\n\x12\x44raftBlockResponse\x12\x17\n\x0f\x64raft_token_ids\x18\x01 \x03(\r\x12\x16\n\x0e\x66orward_passes\x18\x02 \x01(\r\x12\x1d\n\x15peak_activation_bytes\x18\x03 \x01(\x04\"]\n\x14\x45xtendContextRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x1e\n\x03\x61ux\x18\x02 \x03(\x0b\x32\x11.kakeya.v1.Tensor\x12\x11\n\tpositions\x18\x03 \x03(\x05\",\n\x15\x45xtendContextResponse\x12\x13\n\x0b\x63ontext_len\x18\x01 \x01(\r\">\n(DFlashProposerServiceCloseSessionRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\"+\n)DFlashProposerServiceCloseSessionResponse*\xed\x01\n\x0e\x43\x61pabilityRole\x12\x1f\n\x1b\x43\x41PABILITY_ROLE_UNSPECIFIED\x10\x00\x12\x1c\n\x18\x43\x41PABILITY_ROLE_VERIFIER\x10\x01\x12\x1c\n\x18\x43\x41PABILITY_ROLE_PROPOSER\x10\x02\x12\x1c\n\x18\x43\x41PABILITY_ROLE_EMBEDDER\x10\x03\x12\x18\n\x14\x43\x41PABILITY_ROLE_TOOL\x10\x04\x12!\n\x1d\x43\x41PABILITY_ROLE_PREFILL_CACHE\x10\x05\x12#\n\x1f\x43\x41PABILITY_ROLE_PREFILL_COMPUTE\x10\x06*m\n\x10\x43ompressionCodec\x12!\n\x1d\x43OMPRESSION_CODEC_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x43OMPRESSION_CODEC_NONE\x10\x01\x12\x1a\n\x16\x43OMPRESSION_CODEC_ZLIB\x10\x02*\xd8\x01\n\x10PrefillJobStatus\x12\"\n\x1ePREFILL_JOB_STATUS_UNSPECIFIED\x10\x00\x12\x1d\n\x19PREFILL_JOB_STATUS_QUEUED\x10\x01\x12\x1e\n\x1aPREFILL_JOB_STATUS_RUNNING\x10\x02\x12 \n\x1cPREFILL_JOB_STATUS_COMPLETED\x10\x03\x12\x1d\n\x19PREFILL_JOB_STATUS_FAILED\x10\x04\x12 \n\x1cPREFILL_JOB_STATUS_CANCELLED\x10\x05\x32\xdc\x01\n\x11\x43\x61pabilityService\x12g\n\x14\x45xchangeCapabilities\x12&.kakeya.v1.ExchangeCapabilitiesRequest\x1a\'.kakeya.v1.ExchangeCapabilitiesResponse\x12^\n\x11GetNodeCapability\x12#.kakeya.v1.GetNodeCapabilityRequest\x1a$.kakeya.v1.GetNodeCapabilityResponse2b\n\x0fProposerService\x12O\n\x0cProposeBlock\x12\x1e.kakeya.v1.ProposeBlockRequest\x1a\x1f.kakeya.v1.ProposeBlockResponse2\xe3\x02\n\x13PrefillCacheService\x12X\n\x0fGetCacheSummary\x12!.kakeya.v1.GetCacheSummaryRequest\x1a\".kakeya.v1.GetCacheSummaryResponse\x12O\n\x0cLookupPrefix\x12\x1e.kakeya.v1.LookupPrefixRequest\x1a\x1f.kakeya.v1.LookupPrefixResponse\x12N\n\x0b\x46\x65tchBlocks\x12\x1d.kakeya.v1.FetchBlocksRequest\x1a\x1e.kakeya.v1.FetchBlocksResponse0\x01\x12Q\n\x0cPublishBlock\x12\x1e.kakeya.v1.PublishBlockRequest\x1a\x1f.kakeya.v1.PublishBlockResponse(\x01\x32\xb6\x02\n\x14PrefillWorkerService\x12[\n\x10SubmitPrefillJob\x12\".kakeya.v1.SubmitPrefillJobRequest\x1a#.kakeya.v1.SubmitPrefillJobResponse\x12\x64\n\x13GetPrefillJobStatus\x12%.kakeya.v1.GetPrefillJobStatusRequest\x1a&.kakeya.v1.GetPrefillJobStatusResponse\x12[\n\x10\x43\x61ncelPrefillJob\x12\".kakeya.v1.CancelPrefillJobRequest\x1a#.kakeya.v1.CancelPrefillJobResponse2\xc1\x03\n\x15\x44\x46lashProposerService\x12@\n\x07Restore\x12\x19.kakeya.v1.RestoreRequest\x1a\x1a.kakeya.v1.RestoreResponse\x12L\n\x0bSeedContext\x12\x1d.kakeya.v1.SeedContextRequest\x1a\x1e.kakeya.v1.SeedContextResponse\x12I\n\nDraftBlock\x12\x1c.kakeya.v1.DraftBlockRequest\x1a\x1d.kakeya.v1.DraftBlockResponse\x12R\n\rExtendContext\x12\x1f.kakeya.v1.ExtendContextRequest\x1a .kakeya.v1.ExtendContextResponse\x12y\n\x0c\x43loseSession\x12\x33.kakeya.v1.DFlashProposerServiceCloseSessionRequest\x1a\x34.kakeya.v1.DFlashProposerServiceCloseSessionResponseb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1bkakeya/v1/distributed.proto\x12\tkakeya.v1\"}\n\x0fModelCapability\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\'\n\x04role\x18\x02 \x01(\x0e\x32\x19.kakeya.v1.CapabilityRole\x12\x14\n\x0cquantization\x18\x03 \x01(\t\x12\x19\n\x11tokens_per_second\x18\x04 \x01(\x01\"\x83\x03\n\x0eNodeCapability\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12\x14\n\x0cgrpc_address\x18\x02 \x01(\t\x12\x10\n\x08platform\x18\x03 \x01(\t\x12\x1c\n\x14unified_memory_bytes\x18\x04 \x01(\x04\x12\x13\n\x0bmlx_version\x18\x05 \x01(\t\x12*\n\x06models\x18\x06 \x03(\x0b\x32\x1a.kakeya.v1.ModelCapability\x12\x19\n\x11\x61nnounced_at_unix\x18\x07 \x01(\x01\x12\x13\n\x0bttl_seconds\x18\x08 \x01(\x01\x12\x14\n\x0cring_address\x18\t \x01(\t\x12*\n\x06\x63\x61\x63hes\x18\n \x03(\x0b\x32\x1a.kakeya.v1.CacheCapability\x12*\n\tendpoints\x18\x0b \x03(\x0b\x32\x17.kakeya.v1.NodeEndpoint\x12;\n\x0fprefill_workers\x18\x0c \x03(\x0b\x32\".kakeya.v1.PrefillWorkerCapability\"[\n\x0cNodeEndpoint\x12\x0f\n\x07\x61\x64\x64ress\x18\x01 \x01(\t\x12\x0f\n\x07network\x18\x02 \x01(\t\x12\x10\n\x08priority\x18\x03 \x01(\r\x12\x17\n\x0fmeasured_rtt_ms\x18\x04 \x01(\x01\"\xad\x02\n\x12\x43\x61\x63heCompatibility\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x16\n\x0emodel_revision\x18\x02 \x01(\t\x12\x1a\n\x12tokenizer_revision\x18\x03 \x01(\t\x12\x1c\n\x14\x63\x61\x63he_format_version\x18\x04 \x01(\t\x12\x14\n\x0cquantization\x18\x05 \x01(\t\x12\x11\n\trope_hash\x18\x06 \x01(\t\x12\x1b\n\x13layer_geometry_hash\x18\x07 \x01(\t\x12\x10\n\x08kv_dtype\x18\x08 \x01(\t\x12\x19\n\x11\x62lock_size_tokens\x18\t \x01(\r\x12\x18\n\x10tenant_namespace\x18\n \x01(\t\x12\x11\n\tsink_size\x18\x0b \x01(\r\x12\x13\n\x0bwindow_size\x18\x0c \x01(\r\"\xcd\x02\n\x0f\x43\x61\x63heCapability\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x15\n\rcache_address\x18\x02 \x01(\t\x12\x18\n\x10\x63\x61\x63he_bytes_used\x18\x03 \x01(\x04\x12\x18\n\x10\x63\x61\x63he_bytes_free\x18\x04 \x01(\x04\x12\x13\n\x0b\x65ntry_count\x18\x05 \x01(\x04\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x06 \x01(\x04\x12\x0c\n\x04load\x18\x07 \x01(\x01\x12\x15\n\rtokens_served\x18\x08 \x01(\x04\x12\x14\n\x0c\x62loom_filter\x18\t \x01(\x0c\x12\x38\n\x13\x64\x65\x66\x61ult_compression\x18\n \x01(\x0e\x32\x1b.kakeya.v1.CompressionCodec\x12\x1a\n\x12replication_factor\x18\x0b \x01(\r\"\xae\x02\n\x17PrefillWorkerCapability\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x16\n\x0eworker_address\x18\x02 \x01(\t\x12\x1b\n\x13max_concurrent_jobs\x18\x03 \x01(\r\x12\x15\n\rinflight_jobs\x18\x04 \x01(\r\x12\x13\n\x0bqueued_jobs\x18\x05 \x01(\r\x12\x0c\n\x04load\x18\x06 \x01(\x01\x12!\n\x19tokens_per_second_prefill\x18\x07 \x01(\x01\x12\x16\n\x0eram_bytes_free\x18\x08 \x01(\x04\x12\x1c\n\x14\x61\x63\x63\x65pts_compute_jobs\x18\t \x01(\x08\x12\x15\n\rqueued_tokens\x18\n \x01(\x04\"M\n\x1b\x45xchangeCapabilitiesRequest\x12.\n\x0bknown_nodes\x18\x01 \x03(\x0b\x32\x19.kakeya.v1.NodeCapability\"N\n\x1c\x45xchangeCapabilitiesResponse\x12.\n\x0bknown_nodes\x18\x01 \x03(\x0b\x32\x19.kakeya.v1.NodeCapability\"\x1a\n\x18GetNodeCapabilityRequest\"D\n\x19GetNodeCapabilityResponse\x12\'\n\x04node\x18\x01 \x01(\x0b\x32\x19.kakeya.v1.NodeCapability\"N\n\x16GetCacheSummaryRequest\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\"V\n\x17GetCacheSummaryResponse\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12*\n\x06\x63\x61\x63hes\x18\x02 \x03(\x0b\x32\x1a.kakeya.v1.CacheCapability\"a\n\x13LookupPrefixRequest\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x14\n\x0c\x62lock_hashes\x18\x02 \x03(\x0c\"\xcf\x01\n\x14LookupPrefixResponse\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12\x17\n\x0fhit_block_count\x18\x02 \x01(\r\x12\x17\n\x0fhit_token_count\x18\x03 \x01(\x04\x12\x16\n\x0etransfer_bytes\x18\x04 \x01(\x04\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x05 \x01(\x04\x12\x10\n\x08lease_id\x18\x06 \x01(\t\x12\x1d\n\x15lease_expires_at_unix\x18\x07 \x01(\x01\x12\x16\n\x0epayload_sha256\x18\x08 \x01(\x0c\"&\n\x12\x46\x65tchBlocksRequest\x12\x10\n\x08lease_id\x18\x01 \x01(\t\"\xb7\x01\n\x13\x46\x65tchBlocksResponse\x12\x12\n\nblock_hash\x18\x01 \x01(\x0c\x12\x13\n\x0b\x62lock_index\x18\x02 \x01(\r\x12\x13\n\x0btoken_count\x18\x03 \x01(\r\x12\x13\n\x0b\x63hunk_index\x18\x04 \x01(\r\x12\x14\n\x0ctotal_chunks\x18\x05 \x01(\r\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x14\n\x0c\x62lock_sha256\x18\x07 \x01(\x0c\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x08 \x01(\x04\"\xed\x01\n\x13PublishBlockRequest\x12\x12\n\nblock_hash\x18\x01 \x01(\x0c\x12\x13\n\x0b\x62lock_index\x18\x02 \x01(\r\x12\x13\n\x0btoken_count\x18\x03 \x01(\r\x12\x13\n\x0b\x63hunk_index\x18\x04 \x01(\r\x12\x14\n\x0ctotal_chunks\x18\x05 \x01(\r\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x14\n\x0c\x62lock_sha256\x18\x07 \x01(\x0c\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x08 \x01(\x04\x12\x34\n\rcompatibility\x18\t \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\";\n\x14PublishBlockResponse\x12\x0e\n\x06stored\x18\x01 \x01(\x08\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x02 \x01(\x04\"\xf0\x01\n\x17SubmitPrefillJobRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\x12\x34\n\rcompatibility\x18\x03 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x11\n\ttoken_ids\x18\x04 \x03(\r\x12\x14\n\x0c\x62lock_hashes\x18\x05 \x03(\x0c\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x06 \x01(\r\x12:\n\x15preferred_compression\x18\x07 \x01(\x0e\x32\x1b.kakeya.v1.CompressionCodec\"\x85\x01\n\x18SubmitPrefillJobResponse\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12+\n\x06status\x18\x02 \x01(\x0e\x32\x1b.kakeya.v1.PrefillJobStatus\x12\x16\n\x0eworker_node_id\x18\x03 \x01(\t\x12\x14\n\x0cqueue_eta_ms\x18\x04 \x01(\x01\"?\n\x1aGetPrefillJobStatusRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\"\x8c\x02\n\x1bGetPrefillJobStatusResponse\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12+\n\x06status\x18\x02 \x01(\x0e\x32\x1b.kakeya.v1.PrefillJobStatus\x12\x17\n\x0ftokens_computed\x18\x03 \x01(\r\x12\x10\n\x08lease_id\x18\x04 \x01(\t\x12\x12\n\nblock_hash\x18\x05 \x01(\x0c\x12\x16\n\x0epayload_sha256\x18\x06 \x01(\x0c\x12\x16\n\x0etransfer_bytes\x18\x07 \x01(\x04\x12\x16\n\x0e\x66\x61ilure_reason\x18\x08 \x01(\t\x12\x12\n\ncompute_ms\x18\t \x01(\x01\x12\x15\n\rcache_address\x18\n \x01(\t\"<\n\x17\x43\x61ncelPrefillJobRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\"-\n\x18\x43\x61ncelPrefillJobResponse\x12\x11\n\tcancelled\x18\x01 \x01(\x08\"k\n\x13ProposeBlockRequest\x12\x1b\n\x13\x63ommitted_token_ids\x18\x01 \x03(\r\x12\x12\n\nblock_size\x18\x02 \x01(\r\x12\x11\n\tnum_steps\x18\x03 \x01(\r\x12\x10\n\x08model_id\x18\x04 \x01(\t\"y\n\x14ProposeBlockResponse\x12\x11\n\ttoken_ids\x18\x01 \x03(\r\x12\x17\n\x0f\x64iffusion_steps\x18\x02 \x01(\r\x12\x16\n\x0e\x66orward_passes\x18\x03 \x01(\r\x12\x1d\n\x15peak_activation_bytes\x18\x04 \x01(\x04\"4\n\x06Tensor\x12\r\n\x05\x64type\x18\x01 \x01(\t\x12\r\n\x05shape\x18\x02 \x03(\x03\x12\x0c\n\x04\x64\x61ta\x18\x03 \x01(\x0c\"T\n\x07LayerKV\x12\r\n\x05layer\x18\x01 \x01(\x05\x12\x1c\n\x01k\x18\x02 \x01(\x0b\x32\x11.kakeya.v1.Tensor\x12\x1c\n\x01v\x18\x03 \x01(\x0b\x32\x11.kakeya.v1.Tensor\"\x84\x01\n\x0eRestoreRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x12\n\nprompt_ids\x18\x02 \x03(\r\x12\x0c\n\x04sink\x18\x03 \x01(\r\x12\x0e\n\x06window\x18\x04 \x01(\r\x12\x1a\n\x12s5_exact_full_attn\x18\x05 \x01(\x08\x12\x10\n\x08model_id\x18\x06 \x01(\t\"f\n\x0fRestoreResponse\x12$\n\x08restored\x18\x01 \x03(\x0b\x32\x12.kakeya.v1.LayerKV\x12\x19\n\x11\x65victed_positions\x18\x02 \x03(\x05\x12\x12\n\nprompt_len\x18\x03 \x01(\r\"[\n\x12SeedContextRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x1e\n\x03\x61ux\x18\x02 \x03(\x0b\x32\x11.kakeya.v1.Tensor\x12\x11\n\tpositions\x18\x03 \x03(\x05\"*\n\x13SeedContextResponse\x12\x13\n\x0b\x63ontext_len\x18\x01 \x01(\r\"h\n\x11\x44raftBlockRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x16\n\x0e\x62onus_token_id\x18\x02 \x01(\r\x12\x13\n\x0b\x63ontext_len\x18\x03 \x01(\r\x12\x12\n\nblock_size\x18\x04 \x01(\r\"d\n\x12\x44raftBlockResponse\x12\x17\n\x0f\x64raft_token_ids\x18\x01 \x03(\r\x12\x16\n\x0e\x66orward_passes\x18\x02 \x01(\r\x12\x1d\n\x15peak_activation_bytes\x18\x03 \x01(\x04\"]\n\x14\x45xtendContextRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x1e\n\x03\x61ux\x18\x02 \x03(\x0b\x32\x11.kakeya.v1.Tensor\x12\x11\n\tpositions\x18\x03 \x03(\x05\",\n\x15\x45xtendContextResponse\x12\x13\n\x0b\x63ontext_len\x18\x01 \x01(\r\">\n(DFlashProposerServiceCloseSessionRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\"+\n)DFlashProposerServiceCloseSessionResponse*\xed\x01\n\x0e\x43\x61pabilityRole\x12\x1f\n\x1b\x43\x41PABILITY_ROLE_UNSPECIFIED\x10\x00\x12\x1c\n\x18\x43\x41PABILITY_ROLE_VERIFIER\x10\x01\x12\x1c\n\x18\x43\x41PABILITY_ROLE_PROPOSER\x10\x02\x12\x1c\n\x18\x43\x41PABILITY_ROLE_EMBEDDER\x10\x03\x12\x18\n\x14\x43\x41PABILITY_ROLE_TOOL\x10\x04\x12!\n\x1d\x43\x41PABILITY_ROLE_PREFILL_CACHE\x10\x05\x12#\n\x1f\x43\x41PABILITY_ROLE_PREFILL_COMPUTE\x10\x06*\x96\x01\n\x10\x43ompressionCodec\x12!\n\x1d\x43OMPRESSION_CODEC_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x43OMPRESSION_CODEC_NONE\x10\x01\x12\x1a\n\x16\x43OMPRESSION_CODEC_ZLIB\x10\x02\x12\'\n#COMPRESSION_CODEC_KAKEYA_LATTICE_D4\x10\x03*\xd8\x01\n\x10PrefillJobStatus\x12\"\n\x1ePREFILL_JOB_STATUS_UNSPECIFIED\x10\x00\x12\x1d\n\x19PREFILL_JOB_STATUS_QUEUED\x10\x01\x12\x1e\n\x1aPREFILL_JOB_STATUS_RUNNING\x10\x02\x12 \n\x1cPREFILL_JOB_STATUS_COMPLETED\x10\x03\x12\x1d\n\x19PREFILL_JOB_STATUS_FAILED\x10\x04\x12 \n\x1cPREFILL_JOB_STATUS_CANCELLED\x10\x05\x32\xdc\x01\n\x11\x43\x61pabilityService\x12g\n\x14\x45xchangeCapabilities\x12&.kakeya.v1.ExchangeCapabilitiesRequest\x1a\'.kakeya.v1.ExchangeCapabilitiesResponse\x12^\n\x11GetNodeCapability\x12#.kakeya.v1.GetNodeCapabilityRequest\x1a$.kakeya.v1.GetNodeCapabilityResponse2b\n\x0fProposerService\x12O\n\x0cProposeBlock\x12\x1e.kakeya.v1.ProposeBlockRequest\x1a\x1f.kakeya.v1.ProposeBlockResponse2\xe3\x02\n\x13PrefillCacheService\x12X\n\x0fGetCacheSummary\x12!.kakeya.v1.GetCacheSummaryRequest\x1a\".kakeya.v1.GetCacheSummaryResponse\x12O\n\x0cLookupPrefix\x12\x1e.kakeya.v1.LookupPrefixRequest\x1a\x1f.kakeya.v1.LookupPrefixResponse\x12N\n\x0b\x46\x65tchBlocks\x12\x1d.kakeya.v1.FetchBlocksRequest\x1a\x1e.kakeya.v1.FetchBlocksResponse0\x01\x12Q\n\x0cPublishBlock\x12\x1e.kakeya.v1.PublishBlockRequest\x1a\x1f.kakeya.v1.PublishBlockResponse(\x01\x32\xb6\x02\n\x14PrefillWorkerService\x12[\n\x10SubmitPrefillJob\x12\".kakeya.v1.SubmitPrefillJobRequest\x1a#.kakeya.v1.SubmitPrefillJobResponse\x12\x64\n\x13GetPrefillJobStatus\x12%.kakeya.v1.GetPrefillJobStatusRequest\x1a&.kakeya.v1.GetPrefillJobStatusResponse\x12[\n\x10\x43\x61ncelPrefillJob\x12\".kakeya.v1.CancelPrefillJobRequest\x1a#.kakeya.v1.CancelPrefillJobResponse2\xc1\x03\n\x15\x44\x46lashProposerService\x12@\n\x07Restore\x12\x19.kakeya.v1.RestoreRequest\x1a\x1a.kakeya.v1.RestoreResponse\x12L\n\x0bSeedContext\x12\x1d.kakeya.v1.SeedContextRequest\x1a\x1e.kakeya.v1.SeedContextResponse\x12I\n\nDraftBlock\x12\x1c.kakeya.v1.DraftBlockRequest\x1a\x1d.kakeya.v1.DraftBlockResponse\x12R\n\rExtendContext\x12\x1f.kakeya.v1.ExtendContextRequest\x1a .kakeya.v1.ExtendContextResponse\x12y\n\x0c\x43loseSession\x12\x33.kakeya.v1.DFlashProposerServiceCloseSessionRequest\x1a\x34.kakeya.v1.DFlashProposerServiceCloseSessionResponseb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -33,10 +33,10 @@
   DESCRIPTOR._loaded_options = None
   _globals['_CAPABILITYROLE']._serialized_start=4889
   _globals['_CAPABILITYROLE']._serialized_end=5126
-  _globals['_COMPRESSIONCODEC']._serialized_start=5128
-  _globals['_COMPRESSIONCODEC']._serialized_end=5237
-  _globals['_PREFILLJOBSTATUS']._serialized_start=5240
-  _globals['_PREFILLJOBSTATUS']._serialized_end=5456
+  _globals['_COMPRESSIONCODEC']._serialized_start=5129
+  _globals['_COMPRESSIONCODEC']._serialized_end=5279
+  _globals['_PREFILLJOBSTATUS']._serialized_start=5282
+  _globals['_PREFILLJOBSTATUS']._serialized_end=5498
   _globals['_MODELCAPABILITY']._serialized_start=42
   _globals['_MODELCAPABILITY']._serialized_end=167
   _globals['_NODECAPABILITY']._serialized_start=170
@@ -113,14 +113,14 @@
   _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_end=4841
   _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_start=4843
   _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_end=4886
-  _globals['_CAPABILITYSERVICE']._serialized_start=5459
-  _globals['_CAPABILITYSERVICE']._serialized_end=5679
-  _globals['_PROPOSERSERVICE']._serialized_start=5681
-  _globals['_PROPOSERSERVICE']._serialized_end=5779
-  _globals['_PREFILLCACHESERVICE']._serialized_start=5782
-  _globals['_PREFILLCACHESERVICE']._serialized_end=6137
-  _globals['_PREFILLWORKERSERVICE']._serialized_start=6140
-  _globals['_PREFILLWORKERSERVICE']._serialized_end=6450
-  _globals['_DFLASHPROPOSERSERVICE']._serialized_start=6453
-  _globals['_DFLASHPROPOSERSERVICE']._serialized_end=6902
+  _globals['_CAPABILITYSERVICE']._serialized_start=5501
+  _globals['_CAPABILITYSERVICE']._serialized_end=5721
+  _globals['_PROPOSERSERVICE']._serialized_start=5723
+  _globals['_PROPOSERSERVICE']._serialized_end=5821
+  _globals['_PREFILLCACHESERVICE']._serialized_start=5824
+  _globals['_PREFILLCACHESERVICE']._serialized_end=6179
+  _globals['_PREFILLWORKERSERVICE']._serialized_start=6182
+  _globals['_PREFILLWORKERSERVICE']._serialized_end=6492
+  _globals['_DFLASHPROPOSERSERVICE']._serialized_start=6495
+  _globals['_DFLASHPROPOSERSERVICE']._serialized_end=6944
 # @@protoc_insertion_point(module_scope)
diff --git a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi
index 02da80a0..9f2f8cc1 100644
--- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi
+++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi
@@ -22,6 +22,7 @@ class CompressionCodec(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     COMPRESSION_CODEC_UNSPECIFIED: _ClassVar[CompressionCodec]
     COMPRESSION_CODEC_NONE: _ClassVar[CompressionCodec]
     COMPRESSION_CODEC_ZLIB: _ClassVar[CompressionCodec]
+    COMPRESSION_CODEC_KAKEYA_LATTICE_D4: _ClassVar[CompressionCodec]
 
 class PrefillJobStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -41,6 +42,7 @@ CAPABILITY_ROLE_PREFILL_COMPUTE: CapabilityRole
 COMPRESSION_CODEC_UNSPECIFIED: CompressionCodec
 COMPRESSION_CODEC_NONE: CompressionCodec
 COMPRESSION_CODEC_ZLIB: CompressionCodec
+COMPRESSION_CODEC_KAKEYA_LATTICE_D4: CompressionCodec
 PREFILL_JOB_STATUS_UNSPECIFIED: PrefillJobStatus
 PREFILL_JOB_STATUS_QUEUED: PrefillJobStatus
 PREFILL_JOB_STATUS_RUNNING: PrefillJobStatus
diff --git a/proto/kakeya/v1/distributed.proto b/proto/kakeya/v1/distributed.proto
index cf9b679d..702d1e72 100644
--- a/proto/kakeya/v1/distributed.proto
+++ b/proto/kakeya/v1/distributed.proto
@@ -114,6 +114,8 @@ enum CompressionCodec {
   COMPRESSION_CODEC_UNSPECIFIED = 0;
   COMPRESSION_CODEC_NONE = 1;
   COMPRESSION_CODEC_ZLIB = 2;
+  // Lossy D4 Q=38 lattice quantization with bit-packed integer codes.
+  COMPRESSION_CODEC_KAKEYA_LATTICE_D4 = 3;
 }
 
 enum PrefillJobStatus {
diff --git a/requirements-kakeyalattice.txt b/requirements-kakeyalattice.txt
new file mode 100644
index 00000000..7022b0f6
--- /dev/null
+++ b/requirements-kakeyalattice.txt
@@ -0,0 +1,2 @@
+# Bit-packed snapshot support is unreleased on PyPI as of 2026-07-12.
+kakeyalattice @ git+https://github.com/FluffyAIcode/LLM-KV--Cache-compress.git@61ad37f8cd833dff6346cf138bcd37dd12bdf1d9#subdirectory=kakeyalattice
diff --git a/scripts/start_grpc_runtime_server.py b/scripts/start_grpc_runtime_server.py
index 2059e050..00192170 100755
--- a/scripts/start_grpc_runtime_server.py
+++ b/scripts/start_grpc_runtime_server.py
@@ -56,6 +56,15 @@
 _LOG = logging.getLogger("kakeya.grpc-runtime")
 
 
+def _compression_codec(name: str):
+    from inference_engine.distributed.capability import CompressionCodec
+    return {
+        "none": CompressionCodec.NONE,
+        "zlib": CompressionCodec.ZLIB,
+        "kakeyalattice-d4": CompressionCodec.KAKEYA_LATTICE_D4,
+    }[name]
+
+
 def _resolve_kv_dims(verifier) -> Tuple[int, int, int]:
     """Derive (num_layers, num_kv_heads, head_dim) from a loaded
     HF / MLX verifier.
@@ -213,9 +222,7 @@ def _build_capability_registry(
             cache_capability(
                 cache_store,
                 cache_address=args.cache_advertise or args.advertise or args.bind,
-                default_compression=(
-                    2 if args.cache_compression == "zlib" else 1
-                ),
+                default_compression=_compression_codec(args.cache_compression),
                 replication_factor=args.cache_replication_factor,
             ),
         )
@@ -385,7 +392,6 @@ async def _serve(args: argparse.Namespace) -> int:
         import hashlib
         from inference_engine.distributed.capability import (
             CacheCompatibility,
-            CompressionCodec,
         )
         from inference_engine.distributed.prefill_auth import FleetAuthConfig
         from inference_engine.distributed.prefill_cache import PrefixCacheStore
@@ -436,11 +442,7 @@ async def _serve(args: argparse.Namespace) -> int:
             remote_compute_min_tokens=args.remote_prefill_min_tokens,
             max_import_bytes=int(args.cache_max_import_gb * (1 << 30)),
             estimated_snapshot_bytes_per_token=args.cache_estimated_bytes_per_token,
-            compression=(
-                CompressionCodec.ZLIB
-                if args.cache_compression == "zlib"
-                else CompressionCodec.NONE
-            ),
+            compression=_compression_codec(args.cache_compression),
             replication_factor=args.cache_replication_factor,
             cost_config=PrefillCostConfig(
                 local_prefill_tps=args.local_prefill_tps,
@@ -587,9 +589,7 @@ async def _serve(args: argparse.Namespace) -> int:
                 args.exchange_interval_s,
                 cache_store=prefill_store,
                 cache_address=args.cache_advertise or args.advertise or args.bind,
-                cache_compression=(
-                    2 if args.cache_compression == "zlib" else 1
-                ),
+                cache_compression=int(_compression_codec(args.cache_compression)),
                 cache_replication_factor=args.cache_replication_factor,
             ),
         )
@@ -729,7 +729,8 @@ def main() -> int:
     ap.add_argument("--fleet-psk-file", default="",
                     help="Optional fleet PSK file for authenticated prefill RPCs "
                          "and tenant-HMAC prefix hashes.")
-    ap.add_argument("--cache-compression", choices=["none", "zlib"],
+    ap.add_argument("--cache-compression",
+                    choices=["none", "zlib", "kakeyalattice-d4"],
                     default="zlib")
     ap.add_argument("--cache-replication-factor", type=int, default=1)
     ap.add_argument("--cache-max-import-gb", type=float, default=1.0,
diff --git a/scripts/start_prefill_cache_node.py b/scripts/start_prefill_cache_node.py
index 6778e5d3..df061337 100644
--- a/scripts/start_prefill_cache_node.py
+++ b/scripts/start_prefill_cache_node.py
@@ -45,6 +45,14 @@
 from inference_engine.network.state import NetworkState
 
 
+def compression_codec(name: str) -> CompressionCodec:
+    return {
+        "none": CompressionCodec.NONE,
+        "zlib": CompressionCodec.ZLIB,
+        "kakeyalattice-d4": CompressionCodec.KAKEYA_LATTICE_D4,
+    }[name]
+
+
 def physical_memory_bytes() -> int:
     try:
         return int(__import__("os").sysconf("SC_PAGE_SIZE")
@@ -98,11 +106,7 @@ async def serve(args) -> None:
         caches=(cache_capability(
             store,
             cache_address=args.advertise,
-            default_compression=(
-                CompressionCodec.ZLIB
-                if args.cache_compression == "zlib"
-                else CompressionCodec.NONE
-            ),
+            default_compression=compression_codec(args.cache_compression),
             replication_factor=args.replication_factor,
         ),),
         endpoints=(
@@ -202,7 +206,8 @@ def main() -> None:
     ap.add_argument("--sink", type=int, default=4)
     ap.add_argument("--window", type=int, default=64)
     ap.add_argument("--fleet-psk-file", default="")
-    ap.add_argument("--cache-compression", choices=["none", "zlib"],
+    ap.add_argument("--cache-compression",
+                    choices=["none", "zlib", "kakeyalattice-d4"],
                     default="zlib")
     ap.add_argument("--replication-factor", type=int, default=1)
     ap.add_argument("--cache-gb", type=float, default=4)
diff --git a/scripts/start_prefill_worker_node.py b/scripts/start_prefill_worker_node.py
index 98272f66..b5ef96ea 100644
--- a/scripts/start_prefill_worker_node.py
+++ b/scripts/start_prefill_worker_node.py
@@ -52,6 +52,14 @@
 _LOG = logging.getLogger("kakeya.prefill-worker")
 
 
+def compression_codec(name: str) -> CompressionCodec:
+    return {
+        "none": CompressionCodec.NONE,
+        "zlib": CompressionCodec.ZLIB,
+        "kakeyalattice-d4": CompressionCodec.KAKEYA_LATTICE_D4,
+    }[name]
+
+
 def physical_memory_bytes() -> int:
     try:
         import os
@@ -150,11 +158,7 @@ def card() -> NodeCapability:
                 store,
                 cache_address=args.advertise,
                 load=load,
-                default_compression=(
-                    CompressionCodec.ZLIB
-                    if args.cache_compression == "zlib"
-                    else CompressionCodec.NONE
-                ),
+                default_compression=compression_codec(args.cache_compression),
                 replication_factor=args.replication_factor,
             ),),
             endpoints=(
@@ -236,7 +240,8 @@ def main() -> None:
     parser.add_argument("--sink", type=int, default=4)
     parser.add_argument("--window", type=int, default=64)
     parser.add_argument("--cache-gb", type=float, default=4.0)
-    parser.add_argument("--cache-compression", choices=["none", "zlib"],
+    parser.add_argument("--cache-compression",
+                        choices=["none", "zlib", "kakeyalattice-d4"],
                         default="zlib")
     parser.add_argument("--replication-factor", type=int, default=1)
     parser.add_argument("--max-concurrent-jobs", type=int, default=1)
diff --git a/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts b/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts
index 868f3d3b..df8d60ca 100644
--- a/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts
+++ b/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts
@@ -108,6 +108,8 @@ export enum CompressionCodec {
   UNSPECIFIED = 0,
   NONE = 1,
   ZLIB = 2,
+  /** KAKEYA_LATTICE_D4 - Lossy D4 Q=38 lattice quantization with bit-packed integer codes. */
+  KAKEYA_LATTICE_D4 = 3,
   UNRECOGNIZED = -1,
 }
 
@@ -122,6 +124,9 @@ export function compressionCodecFromJSON(object: any): CompressionCodec {
     case 2:
     case "COMPRESSION_CODEC_ZLIB":
       return CompressionCodec.ZLIB;
+    case 3:
+    case "COMPRESSION_CODEC_KAKEYA_LATTICE_D4":
+      return CompressionCodec.KAKEYA_LATTICE_D4;
     case -1:
     case "UNRECOGNIZED":
     default:
@@ -137,6 +142,8 @@ export function compressionCodecToJSON(object: CompressionCodec): string {
       return "COMPRESSION_CODEC_NONE";
     case CompressionCodec.ZLIB:
       return "COMPRESSION_CODEC_ZLIB";
+    case CompressionCodec.KAKEYA_LATTICE_D4:
+      return "COMPRESSION_CODEC_KAKEYA_LATTICE_D4";
     case CompressionCodec.UNRECOGNIZED:
     default:
       return "UNRECOGNIZED";
diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py
index 542f54fa..e32d0ced 100644
--- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py
+++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py
@@ -12,6 +12,8 @@ def test_worker_installer_emits_full_cache_compatibility_contract():
         "--sink",
         "--window",
         "--block-size-tokens",
+        "--cache-format-version",
+        "--cache-compression",
         "--prefill-tps",
         "--network",
         "--priority",
@@ -44,3 +46,8 @@ def test_head_runtime_discovers_and_uses_worker_cache_port():
         in plist
     )
     assert "<string>--fleet-psk-file</string>" in plist
+    assert (
+        "<string>--cache-compression</string>"
+        "<string>kakeyalattice-d4</string>"
+        in plist
+    )
diff --git a/tests/inference_engine/distributed/test_prefill_compression.py b/tests/inference_engine/distributed/test_prefill_compression.py
index 027b8e92..c3539edd 100644
--- a/tests/inference_engine/distributed/test_prefill_compression.py
+++ b/tests/inference_engine/distributed/test_prefill_compression.py
@@ -28,6 +28,27 @@ def test_none_is_backward_compatible_raw_payload():
     assert payload_sizes(raw) == (len(raw), len(raw))
 
 
+def test_kakeyalattice_framing_uses_optional_snapshot_codec(monkeypatch):
+    import sys
+    from types import SimpleNamespace
+
+    monkeypatch.setitem(
+        sys.modules,
+        "inference_engine.distributed.prefill_kakeyalattice",
+        SimpleNamespace(
+            encode_snapshot=lambda raw: raw[::-1],
+            decode_snapshot=lambda packed: packed[::-1],
+        ),
+    )
+    raw = b"portable-kpkv" * 100
+    framed = compress_payload(raw, CompressionCodec.KAKEYA_LATTICE_D4)
+    assert payload_sizes(framed) == (len(framed), len(raw))
+    assert decompress_payload(
+        framed,
+        max_uncompressed_bytes=len(raw),
+    ) == raw
+
+
 def test_compression_and_import_limits_validate():
     with pytest.raises(ValueError, match="zlib level"):
         compress_payload(b"x", CompressionCodec.ZLIB, level=10)
diff --git a/tests/integration/test_prefill_kakeyalattice_snapshot.py b/tests/integration/test_prefill_kakeyalattice_snapshot.py
new file mode 100644
index 00000000..529a48fb
--- /dev/null
+++ b/tests/integration/test_prefill_kakeyalattice_snapshot.py
@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytest.importorskip("kakeyalattice.hf.bitpack")
+
+from inference_engine.backends.mlx.prefill_snapshot import _pack, _unpack
+from inference_engine.distributed.prefill_kakeyalattice import (
+    decode_snapshot,
+    encode_snapshot,
+)
+from inference_engine.distributed.tensor_codec import torch_to_wire, wire_to_torch
+
+
+def test_bitpacked_snapshot_handles_heterogeneous_head_dimensions():
+    torch.manual_seed(7)
+    original = {
+        "layer.0.k": torch.randn(1, 2, 16, 256, dtype=torch.bfloat16),
+        "layer.0.v": torch.randn(1, 2, 16, 256, dtype=torch.bfloat16),
+        "layer.1.k": torch.randn(1, 1, 16, 512, dtype=torch.bfloat16),
+        "layer.1.v": torch.randn(1, 1, 16, 512, dtype=torch.bfloat16),
+    }
+    raw = _pack(
+        [(name, torch_to_wire(value), "torch") for name, value in original.items()],
+        metadata={"layer_count": 2, "token_count": 16},
+    )
+    restored_raw = decode_snapshot(encode_snapshot(raw))
+    metadata, restored = _unpack(restored_raw)
+    assert metadata["token_count"] == 16
+    for name, expected in original.items():
+        actual = wire_to_torch(restored[name][0])
+        assert actual.shape == expected.shape
+        relative_mse = (
+            (actual.float() - expected.float()).square().sum()
+            / expected.float().square().sum()
+        )
+        assert float(relative_mse) < 0.002