From b03ba08604d3f026854b9b3d1b8cab40a233636b Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Sun, 12 Jul 2026 20:25:53 +0800 Subject: [PATCH] feat(prefill): add bit-packed KakeyaLattice snapshots Introduce an explicit D4 Q38 wire codec for heterogeneous Gemma KV tensors, pin the unreleased 1.6.1 dependency, and require matching cache format metadata so lossy packed snapshots are only exchanged by compatible nodes. Co-authored-by: Cursor --- deploy/install_prefill_worker_launchd.sh | 4 + .../ai.kakeya.grpc-runtime-prefill.plist | 2 + docs/ops/distributed-prefill-kv-network.md | 9 + inference_engine/distributed/capability.py | 1 + .../distributed/prefill_compression.py | 29 ++- .../distributed/prefill_kakeyalattice.py | 213 ++++++++++++++++++ .../proto_gen/kakeya/v1/distributed_pb2.py | 30 +-- .../proto_gen/kakeya/v1/distributed_pb2.pyi | 2 + proto/kakeya/v1/distributed.proto | 2 + requirements-kakeyalattice.txt | 2 + scripts/start_grpc_runtime_server.py | 27 +-- scripts/start_prefill_cache_node.py | 17 +- scripts/start_prefill_worker_node.py | 17 +- .../src/proto_gen/kakeya/v1/distributed.ts | 7 + .../bridge/test_prefill_worker_launchd.py | 7 + .../distributed/test_prefill_compression.py | 21 ++ .../test_prefill_kakeyalattice_snapshot.py | 38 ++++ 17 files changed, 382 insertions(+), 46 deletions(-) create mode 100644 inference_engine/distributed/prefill_kakeyalattice.py create mode 100644 requirements-kakeyalattice.txt create mode 100644 tests/integration/test_prefill_kakeyalattice_snapshot.py 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" <--model-revision$MODEL_REVISION --tokenizer-revision$TOKENIZER_REVISION --quantization$QUANTIZATION + --cache-format-version$CACHE_FORMAT_VERSION + --cache-compression$CACHE_COMPRESSION --rope-hash$ROPE_HASH --layer-geometry-hash$KAKEYA_LAYER_GEOMETRY_HASH --tenant-id$TENANT 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 @@ --cache-model-idgemma-4-26B-A4B-it-mlx-4bit --model-revisionlocal-4bit-v1 --tokenizer-revisiongemma4-v1 + --cache-format-versionkakeya-prefill-v3-kl-d4-q38 + --cache-compressionkakeyalattice-d4 --cache-quantization4bit-mlx --cache-kv-dtypebfloat16 --cache-block-tokens64 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="" 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("=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 "--fleet-psk-file" in plist + assert ( + "--cache-compression" + "kakeyalattice-d4" + 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