diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 3c6c9dda..ce5a63d6 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -105,11 +105,9 @@ jobs:
-v
coverage report \
--include='inference_engine/server/auth.py,inference_engine/server/config.py,inference_engine/server/errors.py,inference_engine/server/grpc_app.py,inference_engine/server/metrics.py,inference_engine/server/schemas.py,inference_engine/server/proto_gen/**/*.py,inference_engine/memory/*,inference_engine/bridge/*,inference_engine/distributed/*,inference_engine/network/*,inference_engine/scheduler/config.py,inference_engine/scheduler/session.py,inference_engine/pipeline/*,inference_engine/session/store.py,inference_engine/setup/*,sdks/python/kakeya/__init__.py,sdks/python/kakeya/errors.py,training/repr_align/*' \
- --omit='inference_engine/distributed/prefill_cache_runtime.py' \
--fail-under=100
coverage xml -o coverage.xml \
- --include='inference_engine/server/auth.py,inference_engine/server/config.py,inference_engine/server/errors.py,inference_engine/server/grpc_app.py,inference_engine/server/metrics.py,inference_engine/server/schemas.py,inference_engine/server/proto_gen/**/*.py,inference_engine/memory/*,inference_engine/bridge/*,inference_engine/distributed/*,inference_engine/network/*,inference_engine/scheduler/config.py,inference_engine/scheduler/session.py,inference_engine/pipeline/*,inference_engine/session/store.py,inference_engine/setup/*,sdks/python/kakeya/__init__.py,sdks/python/kakeya/errors.py,training/repr_align/*' \
- --omit='inference_engine/distributed/prefill_cache_runtime.py'
+ --include='inference_engine/server/auth.py,inference_engine/server/config.py,inference_engine/server/errors.py,inference_engine/server/grpc_app.py,inference_engine/server/metrics.py,inference_engine/server/schemas.py,inference_engine/server/proto_gen/**/*.py,inference_engine/memory/*,inference_engine/bridge/*,inference_engine/distributed/*,inference_engine/network/*,inference_engine/scheduler/config.py,inference_engine/scheduler/session.py,inference_engine/pipeline/*,inference_engine/session/store.py,inference_engine/setup/*,sdks/python/kakeya/__init__.py,sdks/python/kakeya/errors.py,training/repr_align/*'
- name: Upload coverage artifact
if: always()
diff --git a/.github/workflows/integration.yaml b/.github/workflows/integration.yaml
index e4e2cebd..193407d4 100644
--- a/.github/workflows/integration.yaml
+++ b/.github/workflows/integration.yaml
@@ -4,12 +4,10 @@ name: Integration (Mac M4)
# tests/integration/ against real Qwen3-0.6B on Apple Silicon.
#
# Trigger model:
-# - Pull-request events. Only fires when the PR carries the
-# ``needs-mac-m4`` label (auto-applied by .github/workflows/
-# auto-label-mac.yaml when a PR touches inference_engine/,
-# sdks/, proto/, or tests/integration/). PRs that don't touch
-# verifier-dependent code skip this gate entirely so the runner
-# pool isn't burned on doc-only or CI-only PRs.
+# - Pull-request events touching runtime/model/proto/integration paths.
+# Path filtering is used directly instead of depending on an auto-label:
+# workflows triggered with GITHUB_TOKEN do not recursively trigger the
+# ``labeled`` event, which previously made the first Mac gate skip.
# - Manual workflow_dispatch for re-runs from the Actions UI.
#
# Runner requirements (self-hosted):
@@ -27,6 +25,14 @@ on:
# Only run on PR events for branches targeting main.
types: [opened, synchronize, reopened, labeled]
branches: [main]
+ paths:
+ - "inference_engine/**"
+ - "kv_cache_proposer/**"
+ - "proto/**"
+ - "sdks/**"
+ - "tests/integration/**"
+ - "tests/backends/mlx/**"
+ - ".github/workflows/integration.yaml"
workflow_dispatch: {}
# Cancel superseded runs on the same PR — saves runner time when
@@ -39,16 +45,33 @@ concurrency:
jobs:
integration:
name: pytest -m integration on Mac M4
- # Only fire on labeled PRs (this saves the runner pool from
- # doc-only / CI-only PRs that don't touch verifier-dependent
- # code). The auto-label workflow adds 'needs-mac-m4' on file
- # paths that warrant the GA gate.
- if: |
- github.event_name == 'workflow_dispatch' ||
- contains(github.event.pull_request.labels.*.name, 'needs-mac-m4')
runs-on: [self-hosted, macOS, ARM64, kakeya-mac-m4]
timeout-minutes: 90
steps:
+ - name: Bootstrap git-lfs before checkout
+ # Reused worktrees can retain a git-lfs post-checkout hook. The hook
+ # executes inside checkout@v4, so repair PATH before checkout.
+ run: |
+ set -euo pipefail
+ if [ -x /opt/homebrew/bin/git-lfs ]; then
+ echo "/opt/homebrew/bin" >> "$GITHUB_PATH"
+ exit 0
+ fi
+ if [ -x /usr/local/bin/git-lfs ]; then
+ echo "/usr/local/bin" >> "$GITHUB_PATH"
+ exit 0
+ fi
+ brew_bin=""
+ for candidate in /opt/homebrew/bin/brew /usr/local/bin/brew; do
+ if [ -x "$candidate" ]; then brew_bin="$candidate"; break; fi
+ done
+ if [ -z "$brew_bin" ]; then
+ echo "::error::git-lfs is missing and Homebrew is unavailable."
+ exit 1
+ fi
+ "$brew_bin" install git-lfs
+ echo "$(dirname "$brew_bin")" >> "$GITHUB_PATH"
+
- uses: actions/checkout@v4
with:
# Full history so the runner can compare against base for
@@ -80,23 +103,81 @@ jobs:
fi
echo "Found $MODEL_DIR"
+ - name: Gate real MLX distributed-prefill continuation equivalence
+ env:
+ KAKEYA_MAC_VERIFIER_PATH_VAR: ${{ vars.KAKEYA_MAC_VERIFIER_PATH || '' }}
+ run: |
+ set -euo pipefail
+ default_verifier="$HOME/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit"
+ if [ ! -d "$default_verifier" ]; then
+ default_verifier="models/gemma-4-26B-A4B-it-mlx-4bit"
+ fi
+ export KAKEYA_MAC_VERIFIER_PATH="${KAKEYA_MAC_VERIFIER_PATH_VAR:-$default_verifier}"
+ if [ ! -d "$KAKEYA_MAC_VERIFIER_PATH" ]; then
+ echo "::error::KAKEYA_MAC_VERIFIER_PATH must point to the pre-warmed MLX verifier."
+ exit 1
+ fi
+ PYBIN="$(
+ python3 - <<'PY'
+ import os, shutil, subprocess
+ candidates = [
+ os.environ.get("KAKEYA_MAC_PYTHON"),
+ os.path.expanduser("~/kakeya-venv/bin/python"),
+ os.path.expanduser("~/.venv/bin/python"),
+ os.path.expanduser(
+ "~/Documents/Kakeya-LLM-Inference-engine-pr109/"
+ ".venv-mac/bin/python3.13"
+ ),
+ os.path.expanduser(
+ "~/Documents/Kakeya-LLM-Inference-engine-pr109/"
+ ".venv-mac/bin/python"
+ ),
+ shutil.which("python3.13"),
+ shutil.which("python3"),
+ ]
+ for candidate in candidates:
+ if not candidate or not os.path.isfile(candidate):
+ continue
+ if subprocess.run(
+ [candidate, "-c", "import mlx_lm, torch, pytest"],
+ stdout=subprocess.DEVNULL,
+ stderr=subprocess.DEVNULL,
+ ).returncode == 0:
+ print(candidate)
+ break
+ PY
+ )"
+ test -n "$PYBIN"
+ PYTHONPATH=.:sdks/python "$PYBIN" -m pytest \
+ -m integration \
+ tests/integration/test_prefill_snapshot_mlx_equivalence.py \
+ -q
+
- name: Install Python dependencies
run: |
- # The runner is expected to have a long-lived venv.
- # If a per-run venv is preferred, swap to ``python3 -m venv .venv``.
- python3 -m pip install --upgrade pip
+ # Keep the legacy Qwen/Transformers-4 suite isolated from both
+ # Homebrew's PEP-668-managed system Python and the Transformers-5 MLX
+ # production venv used by the preceding real-model gate.
+ LEGACY_VENV="$HOME/.kakeya/integration-legacy-venv"
+ if [ ! -x "$LEGACY_VENV/bin/python" ]; then
+ mkdir -p "$(dirname "$LEGACY_VENV")"
+ python3 -m venv "$LEGACY_VENV"
+ fi
+ LEGACY_PY="$LEGACY_VENV/bin/python"
+ "$LEGACY_PY" -m pip install --upgrade pip
# The repo runs via PYTHONPATH (see ci.yaml) — it is NOT a pip package
# (no setup.py/pyproject.toml), so install runtime deps from
# requirements.txt rather than an editable `-e .` (which errors with
# "does not appear to be a Python project").
- python3 -m pip install -r requirements.txt
+ "$LEGACY_PY" -m pip install -r requirements.txt
# The integration suite exercises the legacy dllm-hub Qwen proposer,
# whose remote modeling file depends on the Transformers 4.x
# decoder_layer.attention_type API. Keep this runner in the dedicated
# legacy range; K3/Gemma production paths use requirements.txt's
# unbounded Transformers 5.x-compatible environment.
- python3 -m pip install 'transformers>=4.45,<5.0'
- python3 -m pip install pytest pytest-asyncio pytest-timeout coverage
+ "$LEGACY_PY" -m pip install 'transformers>=4.45,<5.0'
+ "$LEGACY_PY" -m pip install pytest pytest-asyncio pytest-timeout coverage
+ echo "KAKEYA_INTEGRATION_PY=$LEGACY_PY" >> "$GITHUB_ENV"
- name: Run integration suite
env:
@@ -107,7 +188,7 @@ jobs:
run: |
mkdir -p results/platform-tests
stamp=$(date +%s)
- python3 -m pytest \
+ "$KAKEYA_INTEGRATION_PY" -m pytest \
-m integration \
tests/integration/ \
--junitxml="results/platform-tests/integration-mac-m4-${stamp}.junit.xml" \
diff --git a/.github/workflows/mac-bridge.yaml b/.github/workflows/mac-bridge.yaml
index 1eb13c64..07e63fb3 100644
--- a/.github/workflows/mac-bridge.yaml
+++ b/.github/workflows/mac-bridge.yaml
@@ -51,17 +51,35 @@ jobs:
runs-on: [self-hosted, macOS, ARM64, kakeya-mac-m4]
timeout-minutes: 150
steps:
+ - name: Bootstrap git-lfs before checkout
+ # A previous `git lfs install --local` leaves a post-checkout hook in
+ # the reused worktree. That hook runs inside actions/checkout itself,
+ # even with lfs:false, so git-lfs must be available BEFORE checkout.
+ run: |
+ set -euo pipefail
+ if [ -x /opt/homebrew/bin/git-lfs ]; then
+ echo "/opt/homebrew/bin" >> "$GITHUB_PATH"
+ exit 0
+ fi
+ if [ -x /usr/local/bin/git-lfs ]; then
+ echo "/usr/local/bin" >> "$GITHUB_PATH"
+ exit 0
+ fi
+ brew_bin=""
+ for candidate in /opt/homebrew/bin/brew /usr/local/bin/brew; do
+ if [ -x "$candidate" ]; then brew_bin="$candidate"; break; fi
+ done
+ if [ -z "$brew_bin" ]; then
+ echo "::error::git-lfs is missing and Homebrew is unavailable."
+ exit 1
+ fi
+ "$brew_bin" install git-lfs
+ echo "$(dirname "$brew_bin")" >> "$GITHUB_PATH"
+
- uses: actions/checkout@v4
with:
# Push results back to the request branch.
persist-credentials: true
- # k3-* presets load LFS-tracked checkpoints from the repo
- # (e.g. results/research/f_theta_v5_s5_sliding/
- # f_theta_weights.pt). Without lfs:true the workspace holds
- # pointer files and torch.load fails with the cryptic
- # "Unsupported operand 118" (ASCII 'v' = the first byte of
- # an LFS pointer).
- lfs: true
- name: Show request
run: |
diff --git a/README.md b/README.md
index 3b46af79..c105ef8e 100644
--- a/README.md
+++ b/README.md
@@ -9,17 +9,19 @@
Kakeya is a **memory-bounded local agent runtime**: a long-running inference
server that holds session state on the server side, exposes a gRPC
`RuntimeService`, and bounds **KV memory + per-turn latency** on long
-conversations — its KV footprint **does not grow with the conversation** (see
-[Kakeya Attention](#how-this-differs--kakeya-attention-vs-pagedattention--radixattention)).
-
-**v0.4** pairs a frozen **AR verifier** (Gemma-4 26B-A4B) with a **dLLM proposer**
-(DFlash) and a trained projection **f_θ**: a sliding-window-bounded KV cache plus
-**K/V restoration** reconstructs evicted context on demand, so memory is bounded
-**without trading away recall, throughput, or context length**. It ships per
-platform — **`v0.4-mac`** (Apple-Silicon MLX) and **`v0.4-cuda`** (NVIDIA) — atop
-the session-bound gRPC runtime whose foundation landed in June 2026
-([ADR 0008](docs/adr/0008-session-bound-runtime-and-grpc-protocol.md): 9 ms
-latency drift over a 4-hour, 480-turn Mac M4 run; bounded memory).
+conversations — its active decode KV footprint **does not grow with the
+conversation**.
+
+**Current architecture (ADR 0017):** one primary Mac owns RuntimeService,
+sessions and every decode step. Compatible peer Macs are either
+`PREFILL_COMPUTE` workers (same model, prefill only) or `PREFILL_CACHE` nodes
+(RAM-only). The primary imports one immutable longest-prefix KV snapshot,
+computes a missing suffix locally, and decodes entirely locally.
+
+> **Legacy notice:** DFlash, f_θ and fused speculative-decode experiments are
+> retained for reproducibility, but they are no longer the product architecture
+> or a completeness criterion. Historical v0.4 sections below are labeled as
+> such.
```
┌────────────────────┐ gRPC bidi-stream ┌────────────────────────┐
@@ -33,27 +35,26 @@ latency drift over a 4-hour, 480-turn Mac M4 run; bounded memory).
│ └────────┬─────────┘ │
│ ▼ │
│ ┌──────────────────┐ │
- │ │ Restored verifier│ │
+ │ │ Primary verifier│ │
│ │ Gemma-4 26B (AR) │ │
- │ │ + DFlash proposer│ │
- │ │ + f_θ / S5 │ │
│ │ bounded sink+win │ │
+ │ │ peer-prefill KV │ │
│ │ per-session bind │ │
│ └──────────────────┘ │
└────────────────────────┘
```
-> The verifier slot is pluggable: a small **Qwen3 (CPU/MLX)** sink+window
-> verifier for lightweight serving, or the **restored Gemma-4 26B** path
-> (proposer + f_θ/S5) for the memory-bounded, recall-preserving engine below.
+> The verifier slot remains pluggable. All nodes in one prefill group must use
+> the exact same model/tokenizer/quantization/RoPE/cache geometry.
## Distributed Prefill KV Cache Network
-Kakeya can use trusted peer Mac minis as an **immutable prefill-cache tier**.
-Every node advertises model/cache compatibility through the existing P2P
-`CapabilityService`; a cold inference node queries local and remote caches in
-parallel, imports the longest valid token-prefix snapshot once, computes only
-the missing suffix, and keeps autoregressive decode entirely local.
+Kakeya uses trusted peer Mac minis as an **immutable prefill-cache and compute
+tier**. Every node advertises exact compatibility through P2P
+`CapabilityService`. A cold primary queries compatible snapshots; on a miss it
+can submit the prompt to a load-aware `PREFILL_COMPUTE` worker. The worker
+prefills with the same model and retains the snapshot in RAM. The primary
+imports it once and keeps autoregressive decode entirely local.
This is not remote attention and not coherent shared RAM:
@@ -73,6 +74,9 @@ Key properties:
- exact model/tokenizer/quantization/RoPE/cache-format compatibility;
- longest **contiguous** prefix reuse — arbitrary holes are never reused;
- memory-bounded LRU storage with leases and cache epochs;
+- load/cost-aware prefill workers, idempotent jobs and local fallback;
+- gossip-driven worker/cache discovery and deterministic bounded replication;
+- optional zlib framing, fleet-PSK authentication and tenant-HMAC prefix hashes;
- point-to-point chunked gRPC publish/fetch with SHA-256 validation;
- failure-safe fallback to local prefill;
- Thunderbolt/LAN/Tailscale endpoint priority;
@@ -91,6 +95,7 @@ observed speedup ≈97×
Architecture and operations:
- [ADR 0016 — Distributed Prefill KV Cache Network](docs/adr/0016-distributed-prefill-kv-cache-network.md)
+- [ADR 0017 — Primary-decode / distributed-prefill orchestration](docs/adr/0017-prefill-compute-worker-orchestration.md)
- [Two-Mac live report](docs/reports/distributed-prefill-kv-mac-thunderbolt.md)
- [Operator runbook](docs/ops/distributed-prefill-kv-network.md)
@@ -136,7 +141,7 @@ For a full 10-minute walkthrough — Mac vs Linux setup, troubleshooting,
HuggingFace cache pre-warm, mainland-China mirror routing, gRPC SDK
patterns — see [`docs/quickstart.md`](docs/quickstart.md).
-## What's in the v0.4 architecture
+## Historical: v0.4 DFlash/f_θ architecture (no longer the product path)
| Component | What it does | Where |
| --- | --- | --- |
@@ -451,7 +456,7 @@ H200, gemma-4-26B-A4B, recall **1.0**:
> models — the large ~6× memory differentiator, where a bounded window *without*
> restoration would destroy recall — is the **v0.6** roadmap item.
-## v0.4 for Mac — MLX speculative-decode port (the journey to parity)
+## Historical: v0.4 for Mac — MLX speculative-decode research
After the **CUDA** path (f_θ + S5 K/V-restoration verifier, **fused DFlash
spec-decode at 1.79–2.06× AR, recall 1.0 on Gemma-4-26B-A4B / H200**), the engine
diff --git a/deploy/install_prefill_worker_launchd.sh b/deploy/install_prefill_worker_launchd.sh
new file mode 100755
index 00000000..f55c0bf0
--- /dev/null
+++ b/deploy/install_prefill_worker_launchd.sh
@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+# Install the ADR 0017 prefill-compute worker as a per-user LaunchAgent.
+set -euo pipefail
+
+: "${KAKEYA_WORKER_REPO:?set KAKEYA_WORKER_REPO}"
+: "${KAKEYA_WORKER_PYTHON:?set KAKEYA_WORKER_PYTHON}"
+: "${KAKEYA_WORKER_MODEL:?set KAKEYA_WORKER_MODEL}"
+: "${KAKEYA_WORKER_NODE_ID:?set KAKEYA_WORKER_NODE_ID}"
+: "${KAKEYA_WORKER_ADVERTISE:?set KAKEYA_WORKER_ADVERTISE (host:port)}"
+: "${KAKEYA_LAYER_GEOMETRY_HASH:?set KAKEYA_LAYER_GEOMETRY_HASH}"
+
+BIND="${KAKEYA_WORKER_BIND:-0.0.0.0:53051}"
+TENANT="${KAKEYA_TENANT_ID:-default}"
+CACHE_GB="${KAKEYA_WORKER_CACHE_GB:-4}"
+PSK_FILE="${KAKEYA_FLEET_PSK_FILE:-}"
+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}"
+ROPE_HASH="${KAKEYA_ROPE_HASH:-}"
+LABEL="ai.kakeya.prefill-worker"
+PLIST="$HOME/Library/LaunchAgents/$LABEL.plist"
+LOG_DIR="$HOME/.kakeya"
+mkdir -p "$(dirname "$PLIST")" "$LOG_DIR"
+
+psk_xml=""
+if [[ -n "$PSK_FILE" ]]; then
+ psk_xml="--fleet-psk-file$PSK_FILE"
+fi
+
+cat > "$PLIST" <
+
+
+ Label$LABEL
+ ProgramArguments
+ $KAKEYA_WORKER_PYTHON
+ $KAKEYA_WORKER_REPO/scripts/start_prefill_worker_node.py
+ --node-id$KAKEYA_WORKER_NODE_ID
+ --bind$BIND
+ --advertise$KAKEYA_WORKER_ADVERTISE
+ --model-id$KAKEYA_WORKER_MODEL
+ --cache-model-id$CACHE_MODEL_ID
+ --model-revision$MODEL_REVISION
+ --tokenizer-revision$TOKENIZER_REVISION
+ --quantization$QUANTIZATION
+ --rope-hash$ROPE_HASH
+ --layer-geometry-hash$KAKEYA_LAYER_GEOMETRY_HASH
+ --tenant-id$TENANT
+ --cache-gb$CACHE_GB
+ $psk_xml
+
+ WorkingDirectory$KAKEYA_WORKER_REPO
+ EnvironmentVariables
+ PYTHONPATH$KAKEYA_WORKER_REPO:$KAKEYA_WORKER_REPO/sdks/python
+
+ RunAtLoad
+ KeepAlive
+ ProcessTypeInteractive
+ StandardOutPath$LOG_DIR/prefill-worker.log
+ StandardErrorPath$LOG_DIR/prefill-worker.log
+
+EOF
+
+launchctl bootout "gui/$(id -u)/$LABEL" 2>/dev/null || true
+launchctl bootstrap "gui/$(id -u)" "$PLIST"
+echo "installed $LABEL -> $PLIST"
+
diff --git a/docs/adr/0017-prefill-compute-worker-orchestration.md b/docs/adr/0017-prefill-compute-worker-orchestration.md
new file mode 100644
index 00000000..662c24e4
--- /dev/null
+++ b/docs/adr/0017-prefill-compute-worker-orchestration.md
@@ -0,0 +1,131 @@
+# ADR 0017 — Primary-decode / distributed-prefill worker orchestration
+
+- **Status:** Accepted / implementation
+- **Date:** 2026-07-11
+- **Supersedes for product architecture:** DFlash/f_θ/fused speculative-decode
+ serving paths
+- **Extends:** ADR 0009 (capability gossip), ADR 0016 (distributed prefill KV)
+
+## Context
+
+The product topology is one primary Mac mini that owns RuntimeService, session
+state and every autoregressive decode step. Peer Mac minis run the exact same
+model but are restricted to prefill work. They retain immutable prefill K/V
+snapshots in unified memory and answer lookup/fetch requests from the primary.
+Cache-only peers remain useful when maximizing snapshot capacity matters more
+than compute.
+
+ADR 0016 implemented the cache half of this design, but its deployed peer is
+cache-only: the primary computes every miss and replicates snapshots. It does
+not let a peer compute a cold prefix. Peer addresses are also static CLI flags,
+and a fetch/import failure can escape instead of falling back to local prefill.
+
+## Decision
+
+Kakeya adopts three explicit fleet roles:
+
+1. **PRIMARY_DECODE** (`VERIFIER` on the wire): serves users and performs all
+ decode. No remote RPC is allowed in the token loop.
+2. **PREFILL_COMPUTE:** loads the same model, accepts bounded/idempotent prefill
+ jobs, writes snapshots to its local `PrefixCacheStore`, and never serves
+ user decode.
+3. **PREFILL_CACHE:** loads no model; spends RAM on immutable snapshots only.
+
+### Request flow
+
+For a cold append, the primary:
+
+1. computes tenant-HMAC chained prefix hashes;
+2. queries its local cache and compatible live cache/worker cards from gossip;
+3. imports the best cache hit when transfer cost beats local recomputation;
+4. on a miss, submits an idempotent job to the least-loaded compatible
+ `PREFILL_COMPUTE` worker when remote compute+transfer is cheaper;
+5. fetches and imports the completed snapshot once;
+6. computes any missing suffix locally;
+7. decodes entirely locally;
+8. publishes snapshots to a deterministic subset of cache peers.
+
+Every remote error (lookup, job, lease, fetch, checksum, decompress, import)
+resets the verifier and falls back to full local prefill. Cache availability
+must never determine request correctness.
+
+### Discovery and placement
+
+Capability gossip is the only membership source. Static `--cache-peer` and
+`--prefill-worker` flags remain emergency/operator overrides. Cards advertise
+exact cache compatibility, cache address, compute address, queue depth,
+inflight jobs, measured prefill throughput, free RAM and endpoint RTT.
+
+The scheduler is deterministic and cost-aware:
+
+```text
+local_ms = missing_tokens / local_prefill_tps * 1000
+import_ms = endpoint_rtt_ms + transfer_bytes / link_bytes_per_ms
+remote_ms = queue_eta_ms + prompt_tokens / worker_tps * 1000 + import_ms
+```
+
+The longest safe hit is preferred only when `import_ms < local_ms`. A worker is
+used only when `remote_ms < local_ms`. Unknown metrics use conservative
+operator-configured defaults.
+
+### Memory/storage policy
+
+- Decode KV remains local to the primary.
+- Peer memory is a pre-decode snapshot tier, not coherent remote attention RAM.
+- Snapshot payloads support zlib framing and retain SHA-256 of the uncompressed
+ bytes.
+- Replication uses rendezvous hashing and a bounded replication factor instead
+ of publishing every snapshot to every peer.
+- Block size controls checkpoint sparsity. Cache-only nodes should use larger
+ blocks for long prompts; delta snapshots remain a future format revision.
+- Import checks advertised transfer size against a configurable byte budget
+ before allocating/reassembling.
+
+### Security and tenant isolation
+
+Trusted-LAN mode remains available, but production mode uses a fleet PSK:
+
+- request metadata is HMAC-SHA256 signed with timestamp and node/tenant id;
+- replay window is bounded;
+- prefix hashes are HMACed per tenant, preventing cross-tenant prefix probing;
+- tenant namespace is part of `CacheCompatibility`;
+- cache/worker services reject unauthenticated requests before allocation.
+
+mTLS can replace PSK transport later without changing the application protocol.
+
+## Correctness gates
+
+Blocking tests cover:
+
+- scheduler decisions, worker job lifecycle and idempotency;
+- auth/replay/tenant isolation;
+- compression/checksum and replication placement;
+- every remote failure falling back to local prefill;
+- real MLX local-prefill vs remote-prefill-import continuation logits and
+ argmax equivalence;
+- zero prefill/cache RPCs during `Generate`.
+
+## Consequences
+
+Positive:
+
+- the primary dedicates compute to decode;
+- peer compute and RAM scale independently;
+- repeated system/RAG prefixes survive primary restarts;
+- peers join/leave through existing gossip/TTL;
+- local privacy is preserved inside the authenticated fleet.
+
+Costs:
+
+- every compute worker duplicates model weights, reducing RAM available for KV;
+- first-use remote prefill only wins when worker compute plus transfer beats the
+ primary;
+- snapshots still duplicate bounded state at checkpoint boundaries;
+- MLX worker execution is serialized per loaded verifier in the MVP.
+
+## Legacy paths
+
+DFlash, f_θ and fused speculative-decode modules remain temporarily under
+`research/legacy` compatibility surfaces for reproducibility. They are not
+product completeness criteria and must not be wired into RuntimeService.
+
diff --git a/docs/adr/README.md b/docs/adr/README.md
index bd7a7410..cf487188 100644
--- a/docs/adr/README.md
+++ b/docs/adr/README.md
@@ -1,7 +1,7 @@
# Architecture Decision Records
This directory contains Architecture Decision Records (ADRs) for the
-DLM-proposer + AR-verifier project. Each ADR captures a single architectural
+Kakeya distributed inference runtime. Each ADR captures a single architectural
decision, the context that led to it, the alternatives considered, and the
consequences of choosing one path over another.
@@ -40,12 +40,13 @@ reader what was *not* chosen.
| 0006 | [Project positioning as local agent infrastructure](0006-local-agent-infrastructure-positioning.md) | Accepted |
| 0007 | [Cross-request KV cache reuse for long sessions](0007-cross-request-kv-reuse.md) | Superseded by 0008 |
| 0008 | [Session-bound runtime + gRPC protocol](0008-session-bound-runtime-and-grpc-protocol.md) | Accepted |
-| 0009 | [Multi-host milestone: AR-verifier / dLM-proposer on mlx.distributed + agent capability exchange](0009-mlx-distributed-spec-decode-and-capability-exchange.md) | Accepted |
+| 0009 | [Multi-host milestone: AR-verifier / dLM-proposer on mlx.distributed + agent capability exchange](0009-mlx-distributed-spec-decode-and-capability-exchange.md) | Capability gossip retained; proposer path legacy |
| 0012 | [Proposer/verifier value proposition: bounded-memory + recall, platform-forked throughput](0012-proposer-verifier-value-proposition.md) | Accepted |
| 0013 | [Distributed inference topology: what AR sequentiality allows](0013-distributed-inference-topology.md) | Accepted |
| 0014 | [Agent-connection capacity & cross-host proposer/verifier topology: test plan & results](0014-agent-connection-capacity-and-cross-host-topology-tests.md) | Accepted |
-| 0015 | [Kakeya Inference Engine: a product-grade vLLM replacement, Kakeya Attention native](0015-kakeya-attention-and-engine-substrate.md) | Accepted |
-| 0016 | [Distributed Prefill KV Cache Network](0016-distributed-prefill-kv-cache-network.md) | Proposed / MVP |
+| 0015 | [Kakeya Inference Engine: a product-grade vLLM replacement, Kakeya Attention native](0015-kakeya-attention-and-engine-substrate.md) | Superseded by 0017 for product architecture |
+| 0016 | [Distributed Prefill KV Cache Network](0016-distributed-prefill-kv-cache-network.md) | Accepted foundation |
+| 0017 | [Primary-decode / distributed-prefill worker orchestration](0017-prefill-compute-worker-orchestration.md) | Accepted / implementation |
Note: ADR numbering is monotonically increasing; in-flight or
planned numbers (0005) appear in the index so readers can
diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md
index 50c69535..93f26b7d 100644
--- a/docs/ops/distributed-prefill-kv-network.md
+++ b/docs/ops/distributed-prefill-kv-network.md
@@ -1,8 +1,8 @@
# Distributed Prefill KV Cache Network — Operator Runbook
-This runbook deploys one inference head and one cache peer over a private
-Thunderbolt Bridge. The same services work over LAN/Tailscale with lower
-endpoint priority.
+This runbook deploys one primary decode node, prefill-compute workers and
+optional cache-only peers over a private Thunderbolt Bridge. The same services
+work over LAN/Tailscale with lower endpoint priority.
## Production surfaces
@@ -17,13 +17,15 @@ require `X-API-Key`.
## Components
-| Component | Head | Cache peer |
-|---|---|---|
-| Kakeya RuntimeService | `127.0.0.1:51051` | optional |
-| PrefillCacheService | runtime + `:52051` control node | `169.254.27.104:52051` |
-| CapabilityService gossip | enabled | enabled / pull-only if macOS blocks outbound Python sockets |
-| Dashboard/API | `127.0.0.1:8090` | no |
-| Cloudflare public edge | `kakeya.ai/*` Worker | no |
+| Component | Primary | Prefill worker | Cache-only peer |
+|---|---|---|---|
+| Kakeya RuntimeService / decode | `127.0.0.1:51051` | no | no |
+| Same MLX model loaded | yes | yes | no |
+| PrefillWorkerService | no | `:53051` | no |
+| PrefillCacheService | runtime/local LRU | co-located | `:52051` |
+| CapabilityService gossip | enabled | enabled | enabled / pull-only |
+| Dashboard/API | `127.0.0.1:8090` | no | no |
+| Cloudflare public edge | `kakeya.ai/*` Worker | no | no |
## Compatibility lock
@@ -101,6 +103,52 @@ If `nc` works but Python/gRPC outbound calls return `Errno 65`, grant Local
Network access to that Python executable in macOS Privacy & Security. Head→peer
lookup/publish/fetch remains usable while reverse gossip is disabled.
+## Prefill-compute worker
+
+The worker loads the exact same MLX model as the primary, accepts queued
+prefill-only jobs, writes immutable snapshots into its co-located RAM cache and
+never serves user decode.
+
+Create a fleet PSK once and copy it to every trusted node:
+
+```bash
+openssl rand -hex 32 > ~/.kakeya/fleet.psk
+chmod 600 ~/.kakeya/fleet.psk
+```
+
+Install the worker:
+
+```bash
+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"
+export KAKEYA_CACHE_MODEL_ID="gemma-4-26B-A4B-it-mlx-4bit"
+export KAKEYA_MODEL_REVISION="local-4bit-v1"
+export KAKEYA_TOKENIZER_REVISION="gemma4-v1"
+export KAKEYA_WORKER_NODE_ID="prefill-mini-1"
+export KAKEYA_WORKER_ADVERTISE="169.254.27.104:53051"
+export KAKEYA_LAYER_GEOMETRY_HASH=""
+export KAKEYA_FLEET_PSK_FILE="$HOME/.kakeya/fleet.psk"
+export KAKEYA_TENANT_ID="private-fleet"
+bash deploy/install_prefill_worker_launchd.sh
+```
+
+The primary must use the same compatibility and auth values:
+
+```text
+--enable-prefill-cache
+--enable-capability-exchange
+--peer 169.254.27.104:53051
+--cache-tenant-id private-fleet
+--fleet-psk-file ~/.kakeya/fleet.psk
+--cache-compression zlib
+--cache-replication-factor 1
+```
+
+`--cache-peer` remains an emergency static override. Normal worker/cache
+selection is derived from compatible live capability cards and their TTL/load
+metrics.
+
## Node registration and groups
Create a registration:
@@ -132,6 +180,10 @@ Expected invariants:
- imported snapshot checksum and compatibility fingerprint match;
- remote failure falls back to local prefill;
- no remote RPC occurs in autoregressive decode;
+- a cache miss is assigned to a compatible `PREFILL_COMPUTE` worker when its
+ queue+compute+transfer estimate beats blocking the primary;
+- worker failure/timeout/import rejection resets the verifier and performs full
+ local prefill;
- completed and KV-assisted token counters increase after live calls.
Minimal acceptance:
diff --git a/docs/ops/mac-m4-runner-setup.md b/docs/ops/mac-m4-runner-setup.md
index 2cb233c2..ac759646 100644
--- a/docs/ops/mac-m4-runner-setup.md
+++ b/docs/ops/mac-m4-runner-setup.md
@@ -1,11 +1,8 @@
# Mac M4 self-hosted runner setup
-This runner backs the **Integration (Mac M4)** GitHub Actions workflow
-(`.github/workflows/integration.yaml`). It runs `pytest -m integration`
-against real Qwen3-0.6B on every PR labelled `needs-mac-m4`
-(auto-applied by `.github/workflows/auto-label-mac.yaml` when a PR
-touches `inference_engine/`, `sdks/`, `proto/`, `tests/integration/`,
-or `kv_cache_proposer/`).
+This runner backs the **Integration (Mac M4)** and **Mac bridge** workflows.
+Integration is triggered directly by runtime/model/proto/integration path
+changes; it does not depend on a label being applied by another workflow.
## Hardware requirements
@@ -66,13 +63,11 @@ pyenv global 3.12.7
Confirm `python3 --version` returns 3.12.x and `python3 -c 'import platform; print(platform.machine())'` returns `arm64`.
-### 4. (Optional) long-lived venv
+### 4. Pin the MLX workload venv
-The workflow currently does `pip install -e .` per run, which is
-~30 s on a warm pip cache. If you want to skip even that, create a
-venv at `~/kakeya-runner-venv` and add a step to the workflow that
-activates it before `pytest`. v0.3 keeps the per-run install for
-simplicity.
+Set `KAKEYA_MAC_PYTHON` in the runner service environment to the long-lived
+venv that imports `mlx_lm`, `torch` and `pytest`. Real MLX gates resolve this
+interpreter before the legacy Qwen integration environment is installed.
## Runtime expectations
@@ -123,12 +118,14 @@ cd ~/actions-runner
Workflow failures are visible at `Actions → Integration (Mac M4)`. The "Surface failure summary" step inlines the test names + first-line error messages so triage doesn't require downloading the JUnit XML.
-If the runner itself is offline (queue depth grows, no jobs pick up), check on the Mac:
+If the runner itself is offline (multiple differently-labelled jobs remain
+`queued` and no first step starts), check on the Mac:
```bash
cd ~/actions-runner
sudo ./svc.sh status
tail -200 ~/Library/Logs/actions-runner/Runner_*.log
+bash /path/to/repo/scripts/mac_bridge/runner_healthcheck.sh
```
Common causes:
@@ -136,6 +133,24 @@ Common causes:
- HF cache was purged; the verify step fails. Re-warm.
- Disk full from accumulated pip downloads; clear cache.
+### Headless reboot recovery
+
+A user LaunchAgent runs only after GUI login. For an unattended Mac mini,
+install the system watchdog once so the runner recovers before login:
+
+```bash
+cd /path/to/Kakeya-LLM-Inference-engine
+sudo "$HOME/actions-runner/svc.sh" install 2>/dev/null || true
+sudo "$HOME/actions-runner/svc.sh" start 2>/dev/null || true
+bash scripts/mac_bridge/install_autorecover_launchagent.sh --system
+bash scripts/mac_bridge/recover_runner_after_reboot.sh
+```
+
+For jobs that are already queued, the last command is the immediate recovery
+action. GitHub assigns queued jobs automatically once `Runner.Listener`
+reconnects. The system LaunchDaemon retries every 60 seconds and does not
+require a logged-in desktop session.
+
## Mac bridge (cloud-agent access)
The same runner also serves the **Mac bridge**
diff --git a/inference_engine/backends/mlx/prefill_snapshot.py b/inference_engine/backends/mlx/prefill_snapshot.py
index 969b9afd..1fa5abdf 100644
--- a/inference_engine/backends/mlx/prefill_snapshot.py
+++ b/inference_engine/backends/mlx/prefill_snapshot.py
@@ -35,6 +35,7 @@ class ImportedPrefillSnapshot:
token_count: int
cached_token_ids: tuple[int, ...]
next_token_logits: Any | None
+ block_hash: bytes = b""
def export_mlx_prefill_snapshot(
@@ -44,6 +45,7 @@ def export_mlx_prefill_snapshot(
cached_token_ids: Sequence[int],
compatibility: CacheCompatibility,
next_token_logits: Any | None = None,
+ block_hash: bytes = b"",
) -> bytes:
"""Serialize current MLX cache state at one prefix boundary."""
if token_count <= 0:
@@ -68,6 +70,7 @@ def export_mlx_prefill_snapshot(
"token_count": int(token_count),
"cached_token_ids": [int(token) for token in cached_token_ids],
"layer_count": len(cache),
+ "block_hash": bytes(block_hash).hex(),
},
)
@@ -108,6 +111,7 @@ def import_mlx_prefill_snapshot(
token_count=token_count,
cached_token_ids=tuple(int(t) for t in metadata["cached_token_ids"]),
next_token_logits=next_logits,
+ block_hash=bytes.fromhex(metadata.get("block_hash", "")),
)
diff --git a/inference_engine/backends/mlx/prefill_worker.py b/inference_engine/backends/mlx/prefill_worker.py
new file mode 100644
index 00000000..6b8f1942
--- /dev/null
+++ b/inference_engine/backends/mlx/prefill_worker.py
@@ -0,0 +1,95 @@
+"""MLX prefill-only compute engine for ADR 0017 worker nodes."""
+from __future__ import annotations
+
+import threading
+from typing import Sequence
+
+from inference_engine.backends.mlx.prefill_snapshot import (
+ export_mlx_prefill_snapshot,
+)
+from inference_engine.distributed.capability import (
+ CacheCompatibility,
+ CompressionCodec,
+)
+from inference_engine.distributed.prefill_cache import CacheBlock
+from inference_engine.distributed.prefill_compression import compress_payload
+
+
+class MLXPrefillComputeEngine:
+ """Serially runs prefill with a loaded MLX verifier and exports one snapshot."""
+
+ def __init__(self, verifier, compatibility: CacheCompatibility) -> None:
+ self.verifier = verifier
+ self.compatibility = compatibility
+ self._lock = threading.Lock()
+
+ def compute_prefill(
+ self,
+ token_ids: Sequence[int],
+ block_hashes: Sequence[bytes],
+ *,
+ compression: CompressionCodec,
+ cancelled: threading.Event,
+ ) -> Sequence[CacheBlock]:
+ tokens = [int(token) for token in token_ids]
+ if not tokens or not block_hashes:
+ raise ValueError("token_ids and block_hashes must be non-empty")
+ size = self.compatibility.block_size_tokens
+ expected_blocks = (len(tokens) + size - 1) // size
+ if len(block_hashes) != expected_blocks:
+ raise ValueError(
+ f"expected {expected_blocks} block hashes, got {len(block_hashes)}",
+ )
+ with self._lock:
+ if cancelled.is_set():
+ raise InterruptedError("prefill job cancelled")
+ first_end = min(size, len(tokens))
+ self.verifier.prefill(tokens[:first_end])
+ blocks: list[CacheBlock] = [
+ self._snapshot(
+ token_count=first_end,
+ block_hash=block_hashes[0],
+ compression=compression,
+ ),
+ ]
+ for block_index, start in enumerate(
+ range(first_end, len(tokens), size),
+ start=1,
+ ):
+ if cancelled.is_set():
+ raise InterruptedError("prefill job cancelled")
+ block = tokens[start:start + size]
+ logits = self.verifier.forward_block(block)
+ self.verifier.commit_or_truncate(
+ forwarded=len(block),
+ accepted=len(block),
+ )
+ self.verifier.next_token_logits = logits[-1].clone()
+ blocks.append(self._snapshot(
+ token_count=min(start + size, len(tokens)),
+ block_hash=block_hashes[block_index],
+ compression=compression,
+ ))
+ return blocks
+
+ def _snapshot(
+ self,
+ *,
+ token_count: int,
+ block_hash: bytes,
+ compression: CompressionCodec,
+ ) -> CacheBlock:
+ raw = export_mlx_prefill_snapshot(
+ self.verifier.cache,
+ token_count=token_count,
+ cached_token_ids=self.verifier.cached_token_sequence,
+ compatibility=self.compatibility,
+ next_token_logits=self.verifier.next_token_logits,
+ block_hash=block_hash,
+ )
+ return CacheBlock.create(
+ bytes(block_hash),
+ token_count,
+ compress_payload(raw, compression),
+ )
+
diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py
index 483d63db..73535794 100644
--- a/inference_engine/bridge/manifest.py
+++ b/inference_engine/bridge/manifest.py
@@ -104,6 +104,22 @@ def _harness_preset(
PRESETS: Dict[str, Preset] = {
p.name: p
for p in (
+ Preset(
+ name="mlx-prefill-snapshot-equivalence",
+ description="ADR 0017 real-model gate: local MLX prefill vs "
+ "snapshot export/import must preserve continuation "
+ "logits, argmax, retained KV tokens and one decode step.",
+ command_templates=(
+ (
+ "python3", "-m", "pytest",
+ "-m", "integration",
+ "tests/integration/test_prefill_snapshot_mlx_equivalence.py",
+ "-q",
+ ),
+ ),
+ timeout_minutes=45,
+ validate_reports=False,
+ ),
Preset(
name="mlx-distributed-dflash-e2e-inproc",
description="Real-model distributed DFlash+f_θ E2E (in-process): loads "
diff --git a/inference_engine/bridge/runner_python.py b/inference_engine/bridge/runner_python.py
index 1b0c27b4..84eb848e 100644
--- a/inference_engine/bridge/runner_python.py
+++ b/inference_engine/bridge/runner_python.py
@@ -55,6 +55,14 @@ def workload_python_candidates(
environ.get("KAKEYA_MAC_PYTHON"),
expanduser("~/kakeya-venv/bin/python"),
expanduser("~/.venv/bin/python"),
+ expanduser(
+ "~/Documents/Kakeya-LLM-Inference-engine-pr109/"
+ ".venv-mac/bin/python3.13",
+ ),
+ expanduser(
+ "~/Documents/Kakeya-LLM-Inference-engine-pr109/"
+ ".venv-mac/bin/python",
+ ),
which("python3.13"),
which("python3"),
]
diff --git a/inference_engine/distributed/__init__.py b/inference_engine/distributed/__init__.py
index 110bcbe5..d0832c0b 100644
--- a/inference_engine/distributed/__init__.py
+++ b/inference_engine/distributed/__init__.py
@@ -1,20 +1,19 @@
-"""Multi-host plane for Kakeya (ADR 0009, v0.5-M1).
+"""Multi-host plane for Kakeya (ADR 0009 / 0016 / 0017).
Subpackage layout:
- :mod:`capability` — capability cards + the converging gossip
registry (``NodeCapability`` / ``ModelCapability`` /
``CapabilityRegistry``).
-- :mod:`placement` — deterministic spec-decode placement over a
- fleet snapshot.
+- :mod:`prefill_worker` — queued prefill-only compute workers.
+- :mod:`prefill_cache` / :mod:`prefill_cache_service` — immutable RAM cache.
+- :mod:`prefill_cache_runtime` — primary-side orchestration and fallback.
+- :mod:`prefill_scheduler` — load/cost-aware worker and replica placement.
+- :mod:`prefill_auth` — fleet-PSK authentication and tenant hash isolation.
- :mod:`exchange` — ``CapabilityService`` gRPC servicer + the
``exchange_once`` gossip client.
-- :mod:`ngram` — model-free prompt-lookup proposer (the always-
- available proposer capability every node can advertise).
-- :mod:`proposer_service` — ``ProposerService`` gRPC servicer +
- ``RemoteProposer`` client (drop-in ``DLMProposer`` substitute).
-- :mod:`spec_decode` — pure greedy accept rule +
- ``DistributedSpeculativeDecoder``.
+- :mod:`ngram`, :mod:`proposer_service`, :mod:`spec_decode` — legacy research
+ proposer paths retained for reproducibility; not the product architecture.
- :mod:`mlx_ring` — optional ``mlx.distributed`` ring probe
(bulk-tensor data plane advertisement, ADR 0009 §4 item 4).
@@ -26,8 +25,11 @@
from inference_engine.distributed.capability import (
CapabilityRegistry,
CapabilityRole,
+ CacheCompatibility,
+ CompressionCodec,
ModelCapability,
NodeCapability,
+ PrefillWorkerCapability,
)
from inference_engine.distributed.placement import (
PlacementError,
@@ -38,8 +40,11 @@
__all__ = [
"CapabilityRegistry",
"CapabilityRole",
+ "CacheCompatibility",
+ "CompressionCodec",
"ModelCapability",
"NodeCapability",
+ "PrefillWorkerCapability",
"PlacementError",
"SpecDecodePlacement",
"plan_spec_decode_placement",
diff --git a/inference_engine/distributed/capability.py b/inference_engine/distributed/capability.py
index 627396ea..b4a4ef83 100644
--- a/inference_engine/distributed/capability.py
+++ b/inference_engine/distributed/capability.py
@@ -49,6 +49,13 @@ class CapabilityRole(enum.IntEnum):
EMBEDDER = 3
TOOL = 4
PREFILL_CACHE = 5
+ PREFILL_COMPUTE = 6
+
+
+class CompressionCodec(enum.IntEnum):
+ UNSPECIFIED = 0
+ NONE = 1
+ ZLIB = 2
@dataclass(frozen=True)
@@ -112,12 +119,15 @@ class CacheCompatibility:
model_id: str
model_revision: str = ""
tokenizer_revision: str = ""
- cache_format_version: str = "kv-v1"
+ cache_format_version: str = "kakeya-prefill-v2-zlib"
quantization: str = ""
rope_hash: str = ""
layer_geometry_hash: str = ""
kv_dtype: str = ""
block_size_tokens: int = 64
+ tenant_namespace: str = ""
+ sink_size: int = 4
+ window_size: int = 64
def to_proto(self) -> distributed_pb2.CacheCompatibility:
return distributed_pb2.CacheCompatibility(
@@ -130,6 +140,9 @@ def to_proto(self) -> distributed_pb2.CacheCompatibility:
layer_geometry_hash=self.layer_geometry_hash,
kv_dtype=self.kv_dtype,
block_size_tokens=self.block_size_tokens,
+ tenant_namespace=self.tenant_namespace,
+ sink_size=self.sink_size,
+ window_size=self.window_size,
)
@classmethod
@@ -146,6 +159,9 @@ def from_proto(
layer_geometry_hash=msg.layer_geometry_hash,
kv_dtype=msg.kv_dtype,
block_size_tokens=msg.block_size_tokens,
+ tenant_namespace=msg.tenant_namespace,
+ sink_size=msg.sink_size,
+ window_size=msg.window_size,
)
@@ -162,6 +178,8 @@ class CacheCapability:
load: float = 0.0
tokens_served: int = 0
bloom_filter: bytes = b""
+ default_compression: CompressionCodec = CompressionCodec.NONE
+ replication_factor: int = 1
def to_proto(self) -> distributed_pb2.CacheCapability:
return distributed_pb2.CacheCapability(
@@ -174,6 +192,8 @@ def to_proto(self) -> distributed_pb2.CacheCapability:
load=self.load,
tokens_served=self.tokens_served,
bloom_filter=self.bloom_filter,
+ default_compression=int(self.default_compression),
+ replication_factor=self.replication_factor,
)
@classmethod
@@ -188,6 +208,55 @@ def from_proto(cls, msg: distributed_pb2.CacheCapability) -> "CacheCapability":
load=msg.load,
tokens_served=msg.tokens_served,
bloom_filter=msg.bloom_filter,
+ default_compression=CompressionCodec(msg.default_compression),
+ replication_factor=msg.replication_factor,
+ )
+
+
+@dataclass(frozen=True)
+class PrefillWorkerCapability:
+ """One prefill-only compute offering on a node."""
+
+ compatibility: CacheCompatibility
+ worker_address: str = ""
+ max_concurrent_jobs: int = 1
+ inflight_jobs: int = 0
+ queued_jobs: int = 0
+ load: float = 0.0
+ tokens_per_second_prefill: float = 0.0
+ ram_bytes_free: int = 0
+ accepts_compute_jobs: bool = True
+ queued_tokens: int = 0
+
+ def to_proto(self) -> distributed_pb2.PrefillWorkerCapability:
+ return distributed_pb2.PrefillWorkerCapability(
+ compatibility=self.compatibility.to_proto(),
+ worker_address=self.worker_address,
+ max_concurrent_jobs=self.max_concurrent_jobs,
+ inflight_jobs=self.inflight_jobs,
+ queued_jobs=self.queued_jobs,
+ load=self.load,
+ tokens_per_second_prefill=self.tokens_per_second_prefill,
+ ram_bytes_free=self.ram_bytes_free,
+ accepts_compute_jobs=self.accepts_compute_jobs,
+ queued_tokens=self.queued_tokens,
+ )
+
+ @classmethod
+ def from_proto(
+ cls, msg: distributed_pb2.PrefillWorkerCapability,
+ ) -> "PrefillWorkerCapability":
+ return cls(
+ compatibility=CacheCompatibility.from_proto(msg.compatibility),
+ worker_address=msg.worker_address,
+ max_concurrent_jobs=msg.max_concurrent_jobs,
+ inflight_jobs=msg.inflight_jobs,
+ queued_jobs=msg.queued_jobs,
+ load=msg.load,
+ tokens_per_second_prefill=msg.tokens_per_second_prefill,
+ ram_bytes_free=msg.ram_bytes_free,
+ accepts_compute_jobs=msg.accepts_compute_jobs,
+ queued_tokens=msg.queued_tokens,
)
@@ -206,6 +275,7 @@ class NodeCapability:
ring_address: str = ""
caches: Tuple[CacheCapability, ...] = ()
endpoints: Tuple[NodeEndpoint, ...] = ()
+ prefill_workers: Tuple[PrefillWorkerCapability, ...] = ()
def __post_init__(self) -> None:
if not self.node_id:
@@ -239,6 +309,7 @@ def to_proto(self) -> distributed_pb2.NodeCapability:
ring_address=self.ring_address,
caches=[c.to_proto() for c in self.caches],
endpoints=[e.to_proto() for e in self.endpoints],
+ prefill_workers=[w.to_proto() for w in self.prefill_workers],
)
@classmethod
@@ -255,6 +326,10 @@ def from_proto(cls, msg: distributed_pb2.NodeCapability) -> "NodeCapability":
ring_address=msg.ring_address,
caches=tuple(CacheCapability.from_proto(c) for c in msg.caches),
endpoints=tuple(NodeEndpoint.from_proto(e) for e in msg.endpoints),
+ prefill_workers=tuple(
+ PrefillWorkerCapability.from_proto(w)
+ for w in msg.prefill_workers
+ ),
)
diff --git a/inference_engine/distributed/prefill_auth.py b/inference_engine/distributed/prefill_auth.py
new file mode 100644
index 00000000..db30e720
--- /dev/null
+++ b/inference_engine/distributed/prefill_auth.py
@@ -0,0 +1,143 @@
+"""Fleet-PSK authentication and tenant isolation for prefill RPCs.
+
+The transport can still be a private-network insecure gRPC channel, but every
+request is authenticated before allocation/compute. The signature covers the
+deterministic protobuf bytes plus caller/tenant/timestamp metadata. Prefix
+hashes use a tenant-derived HMAC key so one tenant cannot probe another
+tenant's prompt prefixes.
+"""
+from __future__ import annotations
+
+import hashlib
+import hmac
+import time
+from dataclasses import dataclass
+from typing import Iterable, Sequence, Tuple
+
+AUTH_TENANT = "x-kakeya-tenant-id"
+AUTH_NODE = "x-kakeya-node-id"
+AUTH_TS = "x-kakeya-auth-ts"
+AUTH_MAC = "x-kakeya-auth-mac"
+
+
+class PrefillAuthError(ValueError):
+ """Authentication or replay-window failure."""
+
+
+@dataclass(frozen=True)
+class FleetAuthConfig:
+ psk: bytes
+ tenant_id: str
+ node_id: str
+ max_clock_skew_s: float = 60.0
+
+ def __post_init__(self) -> None:
+ if len(self.psk) < 16:
+ raise ValueError("fleet PSK must be at least 16 bytes")
+ if not self.tenant_id:
+ raise ValueError("tenant_id must be non-empty")
+ if not self.node_id:
+ raise ValueError("node_id must be non-empty")
+ if self.max_clock_skew_s <= 0:
+ raise ValueError("max_clock_skew_s must be > 0")
+
+ @classmethod
+ def from_file(
+ cls, path: str, *, tenant_id: str, node_id: str,
+ max_clock_skew_s: float = 60.0,
+ ) -> "FleetAuthConfig":
+ with open(path, "rb") as fh:
+ secret = fh.read().strip()
+ return cls(secret, tenant_id, node_id, max_clock_skew_s)
+
+ def tenant_hash_key(self) -> bytes:
+ return hmac.new(
+ self.psk,
+ b"kakeya-prefill-tenant\0" + self.tenant_id.encode(),
+ hashlib.sha256,
+ ).digest()
+
+
+def _request_bytes(request) -> bytes:
+ serialize = getattr(request, "SerializeToString", None)
+ if serialize is None:
+ raise TypeError("authenticated request must be a protobuf message")
+ return serialize(deterministic=True)
+
+
+def _mac(
+ request, *, psk: bytes, tenant_id: str, node_id: str, timestamp: str,
+) -> str:
+ body_hash = hashlib.sha256(_request_bytes(request)).digest()
+ payload = b"\0".join((
+ tenant_id.encode(),
+ node_id.encode(),
+ timestamp.encode(),
+ body_hash,
+ ))
+ return hmac.new(psk, payload, hashlib.sha256).hexdigest()
+
+
+def signed_metadata(
+ request, config: FleetAuthConfig, *, now: float | None = None,
+) -> Tuple[Tuple[str, str], ...]:
+ timestamp = str(int(time.time() if now is None else now))
+ return (
+ (AUTH_TENANT, config.tenant_id),
+ (AUTH_NODE, config.node_id),
+ (AUTH_TS, timestamp),
+ (AUTH_MAC, _mac(
+ request,
+ psk=config.psk,
+ tenant_id=config.tenant_id,
+ node_id=config.node_id,
+ timestamp=timestamp,
+ )),
+ )
+
+
+def verify_metadata(
+ metadata: Iterable[Tuple[str, str]],
+ request,
+ config: FleetAuthConfig,
+ *,
+ now: float | None = None,
+) -> Tuple[str, str]:
+ values = {key.lower(): value for key, value in metadata}
+ tenant = values.get(AUTH_TENANT, "")
+ node = values.get(AUTH_NODE, "")
+ timestamp = values.get(AUTH_TS, "")
+ supplied = values.get(AUTH_MAC, "")
+ if tenant != config.tenant_id:
+ raise PrefillAuthError("tenant mismatch")
+ if not node or not timestamp or not supplied:
+ raise PrefillAuthError("missing prefill authentication metadata")
+ try:
+ ts = int(timestamp)
+ except ValueError as exc:
+ raise PrefillAuthError("invalid authentication timestamp") from exc
+ current = time.time() if now is None else now
+ if abs(current - ts) > config.max_clock_skew_s:
+ raise PrefillAuthError("authentication timestamp outside replay window")
+ expected = _mac(
+ request,
+ psk=config.psk,
+ tenant_id=tenant,
+ node_id=node,
+ timestamp=timestamp,
+ )
+ if not hmac.compare_digest(supplied, expected):
+ raise PrefillAuthError("invalid prefill authentication MAC")
+ return tenant, node
+
+
+def metadata_pairs(metadata: Sequence) -> Tuple[Tuple[str, str], ...]:
+ """Normalize grpc metadata objects or plain pairs for verification."""
+ result = []
+ for item in metadata:
+ if hasattr(item, "key") and hasattr(item, "value"):
+ result.append((item.key, item.value))
+ else:
+ result.append((item[0], item[1]))
+ return tuple(result)
+
diff --git a/inference_engine/distributed/prefill_cache.py b/inference_engine/distributed/prefill_cache.py
index d931d924..2386f7b2 100644
--- a/inference_engine/distributed/prefill_cache.py
+++ b/inference_engine/distributed/prefill_cache.py
@@ -13,6 +13,7 @@
from __future__ import annotations
import hashlib
+import hmac
import json
import secrets
import threading
@@ -37,7 +38,10 @@ def compatibility_fingerprint(compatibility: CacheCompatibility) -> bytes:
"model_revision": compatibility.model_revision,
"quantization": compatibility.quantization,
"rope_hash": compatibility.rope_hash,
+ "sink_size": compatibility.sink_size,
+ "tenant_namespace": compatibility.tenant_namespace,
"tokenizer_revision": compatibility.tokenizer_revision,
+ "window_size": compatibility.window_size,
}
return hashlib.sha256(
json.dumps(payload, sort_keys=True, separators=(",", ":")).encode(),
@@ -47,6 +51,8 @@ def compatibility_fingerprint(compatibility: CacheCompatibility) -> bytes:
def chained_block_hashes(
token_ids: Sequence[int],
compatibility: CacheCompatibility,
+ *,
+ hmac_key: bytes = b"",
) -> list[bytes]:
"""Hash fixed-size token blocks, chaining each hash to its predecessor.
@@ -62,7 +68,12 @@ def chained_block_hashes(
for start in range(0, len(token_ids), size):
block = token_ids[start:start + size]
encoded = b"".join(int(t).to_bytes(4, "little", signed=False) for t in block)
- previous = hashlib.sha256(namespace + previous + encoded).digest()
+ material = namespace + previous + encoded
+ previous = (
+ hmac.new(hmac_key, material, hashlib.sha256).digest()
+ if hmac_key
+ else hashlib.sha256(material).digest()
+ )
hashes.append(previous)
return hashes
@@ -149,6 +160,7 @@ def put(self, block: CacheBlock) -> bool:
if block.nbytes > self.max_bytes:
raise ValueError("block payload exceeds cache capacity")
with self._lock:
+ self._expire_leases(time.time())
existing = self._blocks.get(block.block_hash)
if existing is not None:
if existing.payload_sha256 != block.payload_sha256:
@@ -159,6 +171,10 @@ def put(self, block: CacheBlock) -> bool:
self._bytes_used += block.nbytes
self._epoch += 1
self._evict_to_budget()
+ if block.block_hash not in self._blocks:
+ raise ValueError(
+ "cache capacity is pinned by active leases",
+ )
return True
def put_prefix(
@@ -249,6 +265,19 @@ def block_hashes(self) -> tuple[bytes, ...]:
with self._lock:
return tuple(self._blocks)
+ def invalidate(self, block_hash: bytes) -> bool:
+ """Drop a corrupt/rejected snapshot so local recompute can replace it."""
+ with self._lock:
+ block = self._blocks.pop(bytes(block_hash), None)
+ if block is None:
+ return False
+ self._bytes_used -= block.nbytes
+ self._epoch += 1
+ for lease_id, lease in list(self._leases.items()):
+ if bytes(block_hash) in lease.block_hashes:
+ del self._leases[lease_id]
+ return True
+
def _expire_leases(self, now: float) -> None:
for lease_id, lease in list(self._leases.items()):
if now > lease.expires_at_unix:
diff --git a/inference_engine/distributed/prefill_cache_runtime.py b/inference_engine/distributed/prefill_cache_runtime.py
index 0b2c5484..e4b6b439 100644
--- a/inference_engine/distributed/prefill_cache_runtime.py
+++ b/inference_engine/distributed/prefill_cache_runtime.py
@@ -9,9 +9,11 @@
from __future__ import annotations
import hashlib
+import time
+import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
-from typing import Any, Sequence
+from typing import Any, Callable, Sequence
import grpc
@@ -19,12 +21,39 @@
export_mlx_prefill_snapshot,
import_mlx_prefill_snapshot,
)
-from inference_engine.distributed.capability import CacheCompatibility
+from inference_engine.distributed.capability import (
+ CacheCompatibility,
+ CompressionCodec,
+ NodeCapability,
+)
+from inference_engine.distributed.prefill_auth import (
+ FleetAuthConfig,
+ signed_metadata,
+)
from inference_engine.distributed.prefill_cache import (
CacheBlock,
PrefixCacheStore,
chained_block_hashes,
)
+from inference_engine.distributed.prefill_compression import (
+ compress_payload,
+ decompress_payload,
+)
+from inference_engine.distributed.prefill_scheduler import (
+ PrefillCostConfig,
+ choose_prefill_worker,
+ remote_import_wins,
+ select_cache_replicas,
+)
+from inference_engine.distributed.prefill_worker import (
+ PrefillJobState,
+ cancel_prefill_job_sync,
+ get_prefill_job_sync,
+ submit_prefill_job_sync,
+)
+from inference_engine.distributed.prefill_cache_service import (
+ compatible_cache_peers,
+)
from inference_engine.server.proto_gen.kakeya.v1 import (
distributed_pb2,
distributed_pb2_grpc,
@@ -39,6 +68,10 @@ class PrefillReuseStats:
tokens_reused: int = 0
tokens_computed: int = 0
bytes_received: int = 0
+ remote_jobs: int = 0
+ remote_job_failures: int = 0
+ fallbacks: int = 0
+ last_fallback_reason: str = ""
@dataclass(frozen=True)
@@ -49,6 +82,9 @@ class _Hit:
hit_tokens: int
transfer_bytes: int
payload: bytes | None = None
+ rtt_ms: float = 0.0
+ block_hash: bytes = b""
+ payload_sha256: bytes = b""
class DistributedPrefillCacheHook:
@@ -59,19 +95,53 @@ def __init__(
local_store: PrefixCacheStore,
*,
peers: Sequence[str] = (),
+ registry_provider: Callable[[], Sequence[NodeCapability]] | None = None,
lookup_timeout_s: float = 2.0,
fetch_timeout_s: float = 30.0,
+ worker_timeout_s: float = 120.0,
+ worker_poll_interval_s: float = 0.05,
+ remote_compute_min_tokens: int = 128,
+ max_import_bytes: int = 1 << 30,
+ estimated_snapshot_bytes_per_token: int = 400_000,
+ compression: CompressionCodec = CompressionCodec.ZLIB,
+ replication_factor: int = 1,
+ cost_config: PrefillCostConfig | None = None,
+ auth: FleetAuthConfig | None = None,
on_reuse=None,
) -> None:
+ if min(
+ lookup_timeout_s,
+ fetch_timeout_s,
+ worker_timeout_s,
+ worker_poll_interval_s,
+ max_import_bytes,
+ estimated_snapshot_bytes_per_token,
+ ) <= 0:
+ raise ValueError("prefill runtime limits must be > 0")
+ if remote_compute_min_tokens < 0 or replication_factor < 0:
+ raise ValueError("prefill thresholds must be >= 0")
self.local_store = local_store
self.compatibility = local_store.compatibility
self.peers = tuple(dict.fromkeys(peer for peer in peers if peer))
+ self.registry_provider = registry_provider
self.lookup_timeout_s = float(lookup_timeout_s)
self.fetch_timeout_s = float(fetch_timeout_s)
+ self.worker_timeout_s = float(worker_timeout_s)
+ self.worker_poll_interval_s = float(worker_poll_interval_s)
+ self.remote_compute_min_tokens = int(remote_compute_min_tokens)
+ self.max_import_bytes = int(max_import_bytes)
+ self.estimated_snapshot_bytes_per_token = int(
+ estimated_snapshot_bytes_per_token,
+ )
+ self.compression = CompressionCodec(compression)
+ self.replication_factor = int(replication_factor)
+ self.cost_config = cost_config or PrefillCostConfig()
+ self.auth = auth
+ self._hash_key = auth.tenant_hash_key() if auth is not None else b""
self.stats = PrefillReuseStats()
self._on_reuse = on_reuse
self._publisher = ThreadPoolExecutor(
- max_workers=max(1, min(4, len(self.peers))),
+ max_workers=4,
thread_name_prefix="prefill-kv-publish",
)
@@ -83,18 +153,82 @@ def prepare(self, verifier: Any, token_ids: Sequence[int]) -> int:
tokens = [int(token) for token in token_ids]
if not tokens:
return 0
- hashes = chained_block_hashes(tokens, self.compatibility)
+ hashes = chained_block_hashes(
+ tokens,
+ self.compatibility,
+ hmac_key=self._hash_key,
+ )
hit = self._best_hit(hashes)
reused = 0
if hit is not None:
- payload = hit.payload if hit.payload is not None else self._fetch_remote(hit)
+ reused = self._try_import(verifier, tokens, hit)
+ elif len(tokens) >= self.remote_compute_min_tokens:
+ remote_hit = self._compute_remote(tokens, hashes)
+ if remote_hit is not None:
+ reused = self._try_import(verifier, tokens, remote_hit)
+ if reused == 0:
+ self.stats.misses += 1
+
+ self._compute_and_publish(verifier, tokens, hashes, reused)
+ return reused
+
+ def _try_import(
+ self,
+ verifier: Any,
+ tokens: list[int],
+ hit: _Hit,
+ ) -> int:
+ try:
+ expected_hit_tokens = min(
+ hit.hit_blocks * self.compatibility.block_size_tokens,
+ len(tokens),
+ )
+ if (
+ hit.hit_blocks <= 0
+ or hit.hit_tokens != expected_hit_tokens
+ ):
+ raise ValueError("prefill hit does not match a token boundary")
+ if (
+ hit.payload is None
+ and hit.transfer_bytes > self.max_import_bytes
+ ):
+ raise ValueError("prefill snapshot exceeds wire import budget")
+ payload = (
+ hit.payload
+ if hit.payload is not None
+ else self._fetch_remote(hit)
+ )
+ if len(payload) > self.max_import_bytes:
+ raise ValueError("prefill snapshot exceeds wire import budget")
+ raw = decompress_payload(
+ payload,
+ max_uncompressed_bytes=self.max_import_bytes,
+ )
verifier.reset()
imported = import_mlx_prefill_snapshot(
- payload,
+ raw,
verifier.cache,
compatibility=self.compatibility,
)
+ if (
+ not (0 < imported.token_count <= len(tokens))
+ or imported.token_count != hit.hit_tokens
+ ):
+ raise ValueError("prefill snapshot token_count is invalid")
+ if hit.block_hash and imported.block_hash != hit.block_hash:
+ raise ValueError("prefill snapshot block hash mismatch")
+ if imported.next_token_logits is None:
+ raise ValueError("prefill snapshot is missing continuation logits")
reused = min(imported.token_count, len(tokens))
+ expected_prefix = tokens[:reused]
+ sink_window = getattr(verifier, "_sink_window_slice", None)
+ expected_cached = (
+ list(sink_window(expected_prefix))
+ if callable(sink_window)
+ else expected_prefix
+ )
+ if list(imported.cached_token_ids) != expected_cached:
+ raise ValueError("prefill snapshot cached token sequence mismatch")
verifier.cached_token_sequence = list(imported.cached_token_ids)
verifier.next_global_position = reused
if imported.next_token_logits is not None:
@@ -106,11 +240,106 @@ def prepare(self, verifier: Any, token_ids: Sequence[int]) -> int:
self.stats.local_hits += 1
else:
self.stats.remote_hits += 1
- else:
- self.stats.misses += 1
+ return reused
+ except Exception as exc:
+ # Cache is an optimization. A corrupt/expired/unreachable hit must
+ # never determine request correctness.
+ self.stats.fallbacks += 1
+ self.stats.last_fallback_reason = f"{type(exc).__name__}: {exc}"
+ if hit.source == "local" and hit.block_hash:
+ self.local_store.invalidate(hit.block_hash)
+ verifier.reset()
+ return 0
- self._compute_and_publish(verifier, tokens, hashes, reused)
- return reused
+ def _compute_remote(
+ self,
+ tokens: list[int],
+ hashes: list[bytes],
+ ) -> _Hit | None:
+ cards = tuple(
+ card for card in self._cards()
+ if card.node_id != self.local_store.node_id
+ )
+ target = choose_prefill_worker(
+ cards,
+ self.compatibility,
+ prompt_tokens=len(tokens),
+ estimated_snapshot_bytes=(
+ len(tokens) * self.estimated_snapshot_bytes_per_token
+ ),
+ config=self.cost_config,
+ )
+ if target is None:
+ return None
+ request = distributed_pb2.SubmitPrefillJobRequest(
+ request_id=uuid.uuid4().hex,
+ tenant_id=self.compatibility.tenant_namespace or "default",
+ compatibility=self.compatibility.to_proto(),
+ token_ids=tokens,
+ block_hashes=hashes,
+ deadline_ms=int(self.worker_timeout_s * 1000),
+ preferred_compression=int(self.compression),
+ )
+ try:
+ response = None
+ response = submit_prefill_job_sync(
+ target.address,
+ request,
+ timeout_s=self.lookup_timeout_s,
+ auth=self.auth,
+ )
+ self.stats.remote_jobs += 1
+ deadline = time.monotonic() + self.worker_timeout_s
+ while time.monotonic() < deadline:
+ status_request = distributed_pb2.GetPrefillJobStatusRequest(
+ job_id=response.job_id,
+ tenant_id=request.tenant_id,
+ )
+ status = get_prefill_job_sync(
+ target.address,
+ status_request,
+ timeout_s=self.lookup_timeout_s,
+ auth=self.auth,
+ )
+ if status.status == int(PrefillJobState.COMPLETED):
+ return _Hit(
+ source=status.cache_address or target.address,
+ lease_id=status.lease_id,
+ hit_blocks=len(hashes),
+ hit_tokens=status.tokens_computed,
+ transfer_bytes=status.transfer_bytes,
+ rtt_ms=target.rtt_ms,
+ block_hash=hashes[-1],
+ payload_sha256=bytes(status.payload_sha256),
+ )
+ if status.status in (
+ int(PrefillJobState.FAILED),
+ int(PrefillJobState.CANCELLED),
+ ):
+ raise RuntimeError(
+ status.failure_reason or "remote prefill job failed",
+ )
+ time.sleep(self.worker_poll_interval_s)
+ raise TimeoutError("remote prefill job timed out")
+ except Exception as exc:
+ if response is not None:
+ try:
+ cancel_request = distributed_pb2.CancelPrefillJobRequest(
+ job_id=response.job_id,
+ tenant_id=request.tenant_id,
+ )
+ cancel_prefill_job_sync(
+ target.address,
+ cancel_request,
+ timeout_s=self.lookup_timeout_s,
+ auth=self.auth,
+ )
+ except Exception:
+ pass
+ self.stats.remote_job_failures += 1
+ self.stats.fallbacks += 1
+ self.stats.last_fallback_reason = f"{type(exc).__name__}: {exc}"
+ return None
def _compute_and_publish(
self,
@@ -120,21 +349,17 @@ def _compute_and_publish(
reused: int,
) -> None:
size = self.compatibility.block_size_tokens
- start_block = reused // size
if reused == 0:
first_end = min(size, len(tokens))
verifier.prefill(tokens[:first_end])
self.stats.tokens_computed += first_end
self._publish_boundary(verifier, tokens, hashes, 0, first_end)
- start_block = 1
- for block_index in range(start_block, len(hashes)):
- start = block_index * size
- if start < reused:
- continue
- end = min(start + size, len(tokens))
- block_tokens = tokens[start:end]
- if not block_tokens:
- continue
+ reused = first_end
+ cursor = reused
+ while cursor < len(tokens):
+ block_index = cursor // size
+ end = min((block_index + 1) * size, len(tokens))
+ block_tokens = tokens[cursor:end]
logits = verifier.forward_block(block_tokens)
verifier.commit_or_truncate(
forwarded=len(block_tokens),
@@ -143,6 +368,7 @@ def _compute_and_publish(
verifier.next_token_logits = logits[-1].clone()
self.stats.tokens_computed += len(block_tokens)
self._publish_boundary(verifier, tokens, hashes, block_index, end)
+ cursor = end
def _publish_boundary(
self,
@@ -152,31 +378,69 @@ def _publish_boundary(
block_index: int,
prefix_end: int,
) -> None:
- payload = export_mlx_prefill_snapshot(
+ raw_payload = export_mlx_prefill_snapshot(
verifier.cache,
token_count=prefix_end,
cached_token_ids=verifier.cached_token_sequence,
compatibility=self.compatibility,
next_token_logits=verifier.next_token_logits,
+ block_hash=hashes[block_index],
)
+ payload = compress_payload(raw_payload, self.compression)
block = CacheBlock.create(hashes[block_index], prefix_end, payload)
self.local_store.put(block)
- if self.peers:
+ peers = self._publish_peers(block.block_hash)
+ if peers:
from inference_engine.distributed.prefill_cache_service import (
publish_block_sync,
)
- for peer in self.peers:
+ for peer in peers:
self._publisher.submit(
publish_block_sync,
peer,
self.compatibility,
block,
timeout_s=self.fetch_timeout_s,
+ auth=self.auth,
)
def close(self) -> None:
self._publisher.shutdown(wait=False, cancel_futures=True)
+ def _cards(self) -> tuple[NodeCapability, ...]:
+ if self.registry_provider is None:
+ return ()
+ try:
+ return tuple(self.registry_provider())
+ except Exception:
+ return ()
+
+ def _cache_peers(self) -> tuple[str, ...]:
+ dynamic = compatible_cache_peers(
+ tuple(
+ card for card in self._cards()
+ if card.node_id != self.local_store.node_id
+ ),
+ self.compatibility,
+ )
+ return tuple(dict.fromkeys((*self.peers, *dynamic)))
+
+ def _publish_peers(self, block_hash: bytes) -> tuple[str, ...]:
+ cards = tuple(
+ card for card in self._cards()
+ if card.node_id != self.local_store.node_id
+ )
+ if cards:
+ selected = select_cache_replicas(
+ cards,
+ self.compatibility,
+ block_hash=block_hash,
+ replication_factor=self.replication_factor,
+ )
+ if selected:
+ return tuple(selected)
+ return self._cache_peers()[:self.replication_factor]
+
def _best_hit(self, hashes: Sequence[bytes]) -> _Hit | None:
candidates: list[_Hit] = []
local = self.local_store.lookup(hashes)
@@ -189,12 +453,15 @@ def _best_hit(self, hashes: Sequence[bytes]) -> _Hit | None:
hit_tokens=local.hit_token_count,
transfer_bytes=local.transfer_bytes,
payload=blocks[-1].payload,
+ block_hash=blocks[-1].block_hash,
+ payload_sha256=blocks[-1].payload_sha256,
))
- if self.peers:
- with ThreadPoolExecutor(max_workers=min(8, len(self.peers))) as pool:
+ peers = self._cache_peers()
+ if peers:
+ with ThreadPoolExecutor(max_workers=min(8, len(peers))) as pool:
futures = {
pool.submit(self._lookup_peer, peer, hashes): peer
- for peer in self.peers
+ for peer in peers
}
for future in as_completed(futures):
hit = future.result()
@@ -202,7 +469,7 @@ def _best_hit(self, hashes: Sequence[bytes]) -> _Hit | None:
candidates.append(hit)
if not candidates:
return None
- return max(
+ best = max(
candidates,
key=lambda hit: (
hit.hit_tokens,
@@ -210,48 +477,119 @@ def _best_hit(self, hashes: Sequence[bytes]) -> _Hit | None:
-hit.transfer_bytes,
),
)
+ if (
+ best.source != "local"
+ and not remote_import_wins(
+ hit_tokens=best.hit_tokens,
+ transfer_bytes=best.transfer_bytes,
+ rtt_ms=best.rtt_ms,
+ config=self.cost_config,
+ )
+ ):
+ local = next(
+ (candidate for candidate in candidates if candidate.source == "local"),
+ None,
+ )
+ return local
+ return best
def _lookup_peer(self, peer: str, hashes: Sequence[bytes]) -> _Hit | None:
try:
with grpc.insecure_channel(peer) as channel:
stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel)
+ request = distributed_pb2.LookupPrefixRequest(
+ compatibility=self.compatibility.to_proto(),
+ block_hashes=hashes,
+ )
response = stub.LookupPrefix(
- distributed_pb2.LookupPrefixRequest(
- compatibility=self.compatibility.to_proto(),
- block_hashes=hashes,
- ),
+ request,
timeout=self.lookup_timeout_s,
+ metadata=(
+ signed_metadata(request, self.auth)
+ if self.auth is not None else None
+ ),
)
except grpc.RpcError:
return None
if not response.lease_id or response.hit_block_count == 0:
return None
+ if response.hit_block_count > len(hashes):
+ return None
return _Hit(
source=peer,
lease_id=response.lease_id,
hit_blocks=response.hit_block_count,
hit_tokens=response.hit_token_count,
transfer_bytes=response.transfer_bytes,
+ rtt_ms=self._peer_rtt(peer),
+ block_hash=bytes(hashes[response.hit_block_count - 1]),
+ payload_sha256=bytes(response.payload_sha256),
)
+ def _peer_rtt(self, address: str) -> float:
+ for card in self._cards():
+ for endpoint in card.endpoints:
+ if endpoint.address == address:
+ return endpoint.measured_rtt_ms
+ return 0.0
+
def _fetch_remote(self, hit: _Hit) -> bytes:
+ if hit.transfer_bytes > self.max_import_bytes:
+ raise RuntimeError("remote prefill payload exceeds import budget")
parts: dict[int, bytes] = {}
expected_chunks = 0
expected_sha = b""
+ expected_block_hash = b""
+ received = 0
try:
with grpc.insecure_channel(hit.source) as channel:
stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel)
+ request = distributed_pb2.FetchBlocksRequest(
+ lease_id=hit.lease_id,
+ )
for chunk in stub.FetchBlocks(
- distributed_pb2.FetchBlocksRequest(lease_id=hit.lease_id),
+ request,
timeout=self.fetch_timeout_s,
+ metadata=(
+ signed_metadata(request, self.auth)
+ if self.auth is not None else None
+ ),
):
- parts[chunk.chunk_index] = bytes(chunk.data)
- expected_chunks = chunk.total_chunks
- expected_sha = bytes(chunk.block_sha256)
+ if expected_chunks == 0:
+ expected_chunks = chunk.total_chunks
+ if expected_chunks > 65_536:
+ raise RuntimeError(
+ "remote prefill cache chunk count exceeds limit",
+ )
+ expected_sha = bytes(chunk.block_sha256)
+ expected_block_hash = bytes(chunk.block_hash)
+ elif (
+ chunk.total_chunks != expected_chunks
+ or bytes(chunk.block_sha256) != expected_sha
+ or bytes(chunk.block_hash) != expected_block_hash
+ ):
+ raise RuntimeError(
+ "remote prefill cache chunk metadata changed",
+ )
+ if (
+ chunk.chunk_index < 0
+ or chunk.chunk_index >= expected_chunks
+ or chunk.chunk_index in parts
+ ):
+ raise RuntimeError("invalid or duplicate prefill chunk")
+ data = bytes(chunk.data)
+ received += len(data)
+ if received > self.max_import_bytes:
+ raise RuntimeError(
+ "remote prefill payload exceeds import budget",
+ )
+ parts[chunk.chunk_index] = data
except grpc.RpcError as exc:
raise RuntimeError(f"remote prefill cache fetch failed: {exc}") from exc
if expected_chunks <= 0 or len(parts) != expected_chunks:
raise RuntimeError("remote prefill cache stream was incomplete")
+ if hit.payload_sha256 and expected_sha != hit.payload_sha256:
+ raise RuntimeError("remote prefill cache lease checksum changed")
payload = b"".join(parts[index] for index in range(expected_chunks))
if hashlib.sha256(payload).digest() != expected_sha:
raise RuntimeError("remote prefill cache checksum mismatch")
diff --git a/inference_engine/distributed/prefill_cache_service.py b/inference_engine/distributed/prefill_cache_service.py
index ec70056c..5168d180 100644
--- a/inference_engine/distributed/prefill_cache_service.py
+++ b/inference_engine/distributed/prefill_cache_service.py
@@ -14,6 +14,12 @@
CacheCompatibility,
NodeCapability,
)
+from inference_engine.distributed.prefill_auth import (
+ FleetAuthConfig,
+ PrefillAuthError,
+ signed_metadata,
+ verify_metadata,
+)
from inference_engine.distributed.prefill_cache import (
CacheBlock,
PrefixCacheStore,
@@ -32,7 +38,10 @@ def cache_capability(
*,
cache_address: str,
load: float = 0.0,
+ default_compression=None,
+ replication_factor: int = 1,
) -> CacheCapability:
+ from inference_engine.distributed.capability import CompressionCodec
stats = store.stats()
return CacheCapability(
compatibility=store.compatibility,
@@ -43,6 +52,12 @@ def cache_capability(
cache_epoch=stats.cache_epoch,
load=load,
tokens_served=stats.tokens_served,
+ default_compression=(
+ CompressionCodec.NONE
+ if default_compression is None
+ else CompressionCodec(default_compression)
+ ),
+ replication_factor=replication_factor,
)
@@ -55,19 +70,48 @@ def __init__(
*,
cache_address: str,
chunk_bytes: int = DEFAULT_CHUNK_BYTES,
+ auth: FleetAuthConfig | None = None,
+ max_payload_bytes: int | None = None,
) -> None:
if chunk_bytes <= 0:
raise ValueError("chunk_bytes must be > 0")
self.store = store
self.cache_address = cache_address
self.chunk_bytes = int(chunk_bytes)
+ self.auth = auth
+ self.max_payload_bytes = (
+ int(max_payload_bytes)
+ if max_payload_bytes is not None else store.max_bytes
+ )
+ if self.max_payload_bytes <= 0:
+ raise ValueError("max_payload_bytes must be > 0")
+
+ async def _authenticate(self, request, context) -> None:
+ if self.auth is None:
+ return
+ try:
+ verify_metadata(context.invocation_metadata(), request, self.auth)
+ except PrefillAuthError as exc:
+ await context.abort(grpc.StatusCode.UNAUTHENTICATED, str(exc))
+
+ async def _authorize_compatibility(self, compatibility, context) -> None:
+ if (
+ self.auth is not None
+ and compatibility.tenant_namespace != self.auth.tenant_id
+ ):
+ await context.abort(
+ grpc.StatusCode.PERMISSION_DENIED,
+ "cache tenant namespace mismatch",
+ )
async def GetCacheSummary( # noqa: N802
self,
request: distributed_pb2.GetCacheSummaryRequest,
context: grpc.aio.ServicerContext,
) -> distributed_pb2.GetCacheSummaryResponse:
+ await self._authenticate(request, context)
requested = CacheCompatibility.from_proto(request.compatibility)
+ await self._authorize_compatibility(requested, context)
caches = []
if requested == self.store.compatibility:
caches.append(
@@ -86,7 +130,9 @@ async def LookupPrefix( # noqa: N802
request: distributed_pb2.LookupPrefixRequest,
context: grpc.aio.ServicerContext,
) -> distributed_pb2.LookupPrefixResponse:
+ await self._authenticate(request, context)
requested = CacheCompatibility.from_proto(request.compatibility)
+ await self._authorize_compatibility(requested, context)
if requested != self.store.compatibility:
return distributed_pb2.LookupPrefixResponse(
node_id=self.store.node_id,
@@ -100,6 +146,7 @@ async def FetchBlocks( # noqa: N802
request: distributed_pb2.FetchBlocksRequest,
context: grpc.aio.ServicerContext,
):
+ await self._authenticate(request, context)
try:
blocks = self.store.fetch(request.lease_id)
except KeyError as exc:
@@ -130,9 +177,23 @@ async def PublishBlock( # noqa: N802
) -> distributed_pb2.PublishBlockResponse:
parts: dict[int, bytes] = {}
first = None
+ received = 0
async for chunk in request_iterator:
if first is None:
first = chunk
+ await self._authenticate(first, context)
+ requested = CacheCompatibility.from_proto(first.compatibility)
+ await self._authorize_compatibility(requested, context)
+ if (
+ first.total_chunks <= 0
+ or first.total_chunks
+ > (self.max_payload_bytes + self.chunk_bytes - 1)
+ // self.chunk_bytes
+ ):
+ await context.abort(
+ grpc.StatusCode.RESOURCE_EXHAUSTED,
+ "publish stream exceeds payload budget",
+ )
elif (
chunk.block_hash != first.block_hash
or chunk.total_chunks != first.total_chunks
@@ -142,7 +203,23 @@ async def PublishBlock( # noqa: N802
grpc.StatusCode.INVALID_ARGUMENT,
"inconsistent publish chunk metadata",
)
- parts[chunk.chunk_index] = bytes(chunk.data)
+ if (
+ chunk.chunk_index >= chunk.total_chunks
+ or chunk.chunk_index in parts
+ or len(chunk.data) > self.chunk_bytes
+ ):
+ await context.abort(
+ grpc.StatusCode.INVALID_ARGUMENT,
+ "invalid or duplicate publish chunk",
+ )
+ data = bytes(chunk.data)
+ received += len(data)
+ if received > self.max_payload_bytes:
+ await context.abort(
+ grpc.StatusCode.RESOURCE_EXHAUSTED,
+ "publish payload exceeds cache allocation budget",
+ )
+ parts[chunk.chunk_index] = data
if first is None:
await context.abort(
grpc.StatusCode.INVALID_ARGUMENT,
@@ -185,11 +262,15 @@ def add_prefill_cache_service(
*,
cache_address: str,
chunk_bytes: int = DEFAULT_CHUNK_BYTES,
+ auth: FleetAuthConfig | None = None,
+ max_payload_bytes: int | None = None,
) -> PrefillCacheServiceServicer:
servicer = PrefillCacheServiceServicer(
store,
cache_address=cache_address,
chunk_bytes=chunk_bytes,
+ auth=auth,
+ max_payload_bytes=max_payload_bytes,
)
distributed_pb2_grpc.add_PrefillCacheServiceServicer_to_server(
servicer, server,
@@ -220,16 +301,19 @@ async def lookup_peer(
block_hashes: Sequence[bytes],
*,
timeout_s: float = 3.0,
+ auth: FleetAuthConfig | None = None,
) -> RemotePrefixHit:
try:
async with grpc.aio.insecure_channel(address) as channel:
stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel)
+ request = distributed_pb2.LookupPrefixRequest(
+ compatibility=compatibility.to_proto(),
+ block_hashes=block_hashes,
+ )
response = await stub.LookupPrefix(
- distributed_pb2.LookupPrefixRequest(
- compatibility=compatibility.to_proto(),
- block_hashes=block_hashes,
- ),
+ request,
timeout=timeout_s,
+ metadata=signed_metadata(request, auth) if auth else None,
)
except grpc.aio.AioRpcError:
return RemotePrefixHit(address, "", "", 0, 0, 0, 0, 0.0, b"")
@@ -252,6 +336,7 @@ async def lookup_best_peer(
block_hashes: Sequence[bytes],
*,
timeout_s: float = 3.0,
+ auth: FleetAuthConfig | None = None,
) -> RemotePrefixHit | None:
"""Fan out concurrently and choose longest hit, then smallest transfer."""
if not peers:
@@ -262,6 +347,7 @@ async def lookup_best_peer(
compatibility,
block_hashes,
timeout_s=timeout_s,
+ auth=auth,
)
for peer in peers
))
@@ -282,12 +368,15 @@ async def fetch_remote_blocks(
hit: RemotePrefixHit,
*,
timeout_s: float = 30.0,
+ auth: FleetAuthConfig | None = None,
) -> list[distributed_pb2.FetchBlocksResponse]:
async with grpc.aio.insecure_channel(hit.address) as channel:
stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel)
+ request = distributed_pb2.FetchBlocksRequest(lease_id=hit.lease_id)
stream = stub.FetchBlocks(
- distributed_pb2.FetchBlocksRequest(lease_id=hit.lease_id),
+ request,
timeout=timeout_s,
+ metadata=signed_metadata(request, auth) if auth else None,
)
return [chunk async for chunk in stream]
@@ -299,6 +388,7 @@ def publish_block_sync(
*,
timeout_s: float = 30.0,
chunk_bytes: int = DEFAULT_CHUNK_BYTES,
+ auth: FleetAuthConfig | None = None,
) -> bool:
"""Publish one immutable snapshot to a peer (used by background workers)."""
try:
@@ -306,24 +396,31 @@ def publish_block_sync(
1, (block.nbytes + chunk_bytes - 1) // chunk_bytes,
)
+ chunk_messages = []
+ for chunk_index in range(total_chunks):
+ start = chunk_index * chunk_bytes
+ chunk_messages.append(distributed_pb2.PublishBlockRequest(
+ block_hash=block.block_hash,
+ token_count=block.token_count,
+ chunk_index=chunk_index,
+ total_chunks=total_chunks,
+ data=block.payload[start:start + chunk_bytes],
+ block_sha256=block.payload_sha256,
+ compatibility=compatibility.to_proto(),
+ ))
+
def chunks():
- for chunk_index in range(total_chunks):
- start = chunk_index * chunk_bytes
- yield distributed_pb2.PublishBlockRequest(
- block_hash=block.block_hash,
- token_count=block.token_count,
- chunk_index=chunk_index,
- total_chunks=total_chunks,
- data=block.payload[start:start + chunk_bytes],
- block_sha256=block.payload_sha256,
- compatibility=compatibility.to_proto(),
- )
+ yield from chunk_messages
with grpc.insecure_channel(address) as channel:
stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel)
response = stub.PublishBlock(
chunks(),
timeout=timeout_s,
+ metadata=(
+ signed_metadata(chunk_messages[0], auth)
+ if auth is not None else None
+ ),
)
except grpc.RpcError:
return False
diff --git a/inference_engine/distributed/prefill_compression.py b/inference_engine/distributed/prefill_compression.py
new file mode 100644
index 00000000..23f7ef2c
--- /dev/null
+++ b/inference_engine/distributed/prefill_compression.py
@@ -0,0 +1,87 @@
+"""Snapshot payload framing and compression for distributed prefill K/V."""
+from __future__ import annotations
+
+import hashlib
+import struct
+import zlib
+
+from inference_engine.distributed.capability import CompressionCodec
+
+_MAGIC = b"KPC1"
+_HEADER = struct.Struct("<4sBQ32s")
+
+
+def compress_payload(
+ payload: bytes,
+ codec: CompressionCodec,
+ *,
+ level: int = 3,
+) -> bytes:
+ raw = bytes(payload)
+ if codec in (CompressionCodec.UNSPECIFIED, CompressionCodec.NONE):
+ return raw
+ if codec != CompressionCodec.ZLIB:
+ 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)
+ return _HEADER.pack(
+ _MAGIC,
+ int(codec),
+ len(raw),
+ hashlib.sha256(raw).digest(),
+ ) + compressed
+
+
+def decompress_payload(
+ payload: bytes,
+ *,
+ max_uncompressed_bytes: int,
+) -> bytes:
+ data = bytes(payload)
+ if not data.startswith(_MAGIC):
+ if len(data) > max_uncompressed_bytes:
+ raise ValueError("uncompressed prefill payload exceeds import budget")
+ return data
+ if len(data) < _HEADER.size:
+ raise ValueError("truncated compressed prefill payload header")
+ _magic, raw_codec, expected_size, expected_sha = _HEADER.unpack(
+ data[:_HEADER.size],
+ )
+ if expected_size > max_uncompressed_bytes:
+ raise ValueError("prefill payload exceeds uncompressed import budget")
+ try:
+ codec = CompressionCodec(raw_codec)
+ except ValueError as exc:
+ raise ValueError(f"unsupported compression codec {raw_codec}") from exc
+ if codec != CompressionCodec.ZLIB:
+ raise ValueError(f"unsupported framed compression codec {codec.name}")
+ decompressor = zlib.decompressobj()
+ raw = decompressor.decompress(
+ data[_HEADER.size:],
+ max_uncompressed_bytes + 1,
+ )
+ if decompressor.unconsumed_tail or len(raw) > max_uncompressed_bytes:
+ raise ValueError("decompressed payload exceeds import budget")
+ remaining = max_uncompressed_bytes - len(raw)
+ raw += decompressor.flush(remaining + 1)
+ if len(raw) > max_uncompressed_bytes: # pragma: no cover - zlib flush guard
+ raise ValueError("decompressed payload exceeds import budget")
+ if decompressor.unused_data:
+ raise ValueError("compressed prefill payload has trailing data")
+ 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
+
+
+def payload_sizes(payload: bytes) -> tuple[int, int]:
+ """Return ``(wire_bytes, uncompressed_bytes)`` without decompressing."""
+ data = bytes(payload)
+ if data.startswith(_MAGIC) and len(data) >= _HEADER.size:
+ return len(data), int(_HEADER.unpack(data[:_HEADER.size])[2])
+ return len(data), len(data)
+
diff --git a/inference_engine/distributed/prefill_scheduler.py b/inference_engine/distributed/prefill_scheduler.py
new file mode 100644
index 00000000..c8cd8475
--- /dev/null
+++ b/inference_engine/distributed/prefill_scheduler.py
@@ -0,0 +1,190 @@
+"""Pure placement/cost functions for distributed prefill compute and storage."""
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass
+from enum import Enum
+from typing import Sequence
+
+from inference_engine.distributed.capability import (
+ CacheCompatibility,
+ NodeCapability,
+ PrefillWorkerCapability,
+)
+
+
+class PrefillAction(str, Enum):
+ IMPORT = "import"
+ REMOTE_COMPUTE = "remote_compute"
+ LOCAL_COMPUTE = "local_compute"
+
+
+@dataclass(frozen=True)
+class PrefillCostConfig:
+ local_prefill_tps: float = 20.0
+ default_worker_tps: float = 20.0
+ link_mbps: float = 1000.0
+ default_rtt_ms: float = 2.0
+ minimum_savings_ratio: float = 0.10
+ primary_compute_penalty_ms: float = 0.0
+
+ def __post_init__(self) -> None:
+ if min(
+ self.local_prefill_tps,
+ self.default_worker_tps,
+ self.link_mbps,
+ self.default_rtt_ms,
+ ) <= 0:
+ raise ValueError("prefill cost metrics must be > 0")
+ if not (0 <= self.minimum_savings_ratio < 1):
+ raise ValueError("minimum_savings_ratio must be in [0, 1)")
+ if self.primary_compute_penalty_ms < 0:
+ raise ValueError("primary_compute_penalty_ms must be >= 0")
+
+
+@dataclass(frozen=True)
+class WorkerTarget:
+ node_id: str
+ address: str
+ capability: PrefillWorkerCapability
+ rtt_ms: float
+
+ @property
+ def queue_eta_ms(self) -> float:
+ tps = self.capability.tokens_per_second_prefill
+ if tps <= 0:
+ return 0.0
+ return self.capability.queued_tokens * 1000.0 / tps
+
+
+def estimate_local_prefill_ms(tokens: int, config: PrefillCostConfig) -> float:
+ return (
+ max(0, tokens) / config.local_prefill_tps * 1000.0
+ + config.primary_compute_penalty_ms
+ )
+
+
+def estimate_import_ms(
+ transfer_bytes: int,
+ *,
+ rtt_ms: float,
+ config: PrefillCostConfig,
+) -> float:
+ bytes_per_ms = config.link_mbps * 1_000_000.0 / 8.0 / 1000.0
+ return max(rtt_ms, 0.0) + max(transfer_bytes, 0) / bytes_per_ms
+
+
+def remote_import_wins(
+ *,
+ hit_tokens: int,
+ transfer_bytes: int,
+ rtt_ms: float,
+ config: PrefillCostConfig,
+) -> bool:
+ local = estimate_local_prefill_ms(hit_tokens, config)
+ remote = estimate_import_ms(
+ transfer_bytes,
+ rtt_ms=rtt_ms or config.default_rtt_ms,
+ config=config,
+ )
+ return remote <= local * (1.0 - config.minimum_savings_ratio)
+
+
+def compatible_prefill_workers(
+ cards: Sequence[NodeCapability],
+ compatibility: CacheCompatibility,
+) -> list[WorkerTarget]:
+ targets: list[WorkerTarget] = []
+ for card in cards:
+ rtt = min(
+ (
+ endpoint.measured_rtt_ms
+ for endpoint in card.endpoints
+ if endpoint.measured_rtt_ms > 0
+ ),
+ default=0.0,
+ )
+ for worker in card.prefill_workers:
+ if (
+ worker.accepts_compute_jobs
+ and worker.compatibility == compatibility
+ and (worker.worker_address or card.grpc_address)
+ ):
+ targets.append(WorkerTarget(
+ card.node_id,
+ worker.worker_address or card.grpc_address,
+ worker,
+ rtt,
+ ))
+ return targets
+
+
+def choose_prefill_worker(
+ cards: Sequence[NodeCapability],
+ compatibility: CacheCompatibility,
+ *,
+ prompt_tokens: int,
+ estimated_snapshot_bytes: int,
+ config: PrefillCostConfig,
+) -> WorkerTarget | None:
+ candidates = compatible_prefill_workers(cards, compatibility)
+ if not candidates:
+ return None
+ local_ms = estimate_local_prefill_ms(prompt_tokens, config)
+
+ def cost(target: WorkerTarget) -> float:
+ tps = (
+ target.capability.tokens_per_second_prefill
+ or config.default_worker_tps
+ )
+ compute_ms = prompt_tokens / tps * 1000.0
+ import_ms = estimate_import_ms(
+ estimated_snapshot_bytes,
+ rtt_ms=target.rtt_ms or config.default_rtt_ms,
+ config=config,
+ )
+ load_penalty = max(0.0, target.capability.load) * compute_ms
+ return target.queue_eta_ms + compute_ms + import_ms + load_penalty
+
+ best = min(candidates, key=lambda target: (
+ cost(target),
+ -target.capability.ram_bytes_free,
+ target.node_id,
+ ))
+ if cost(best) > local_ms * (1.0 - config.minimum_savings_ratio):
+ return None
+ return best
+
+
+def select_cache_replicas(
+ cards: Sequence[NodeCapability],
+ compatibility: CacheCompatibility,
+ *,
+ block_hash: bytes,
+ replication_factor: int,
+) -> list[str]:
+ """Deterministic rendezvous placement avoids publishing to every peer."""
+ if replication_factor <= 0:
+ return []
+ candidates: list[tuple[int, int, str]] = []
+ for card in cards:
+ for cache in card.caches:
+ if cache.compatibility != compatibility or not cache.cache_address:
+ continue
+ score = int.from_bytes(hashlib.sha256(
+ bytes(block_hash) + card.node_id.encode(),
+ ).digest(), "big")
+ candidates.append((
+ score,
+ cache.cache_bytes_free,
+ cache.cache_address,
+ ))
+ candidates.sort(reverse=True)
+ result: list[str] = []
+ for _score, _free, address in candidates:
+ if address not in result:
+ result.append(address)
+ if len(result) >= replication_factor:
+ break
+ return result
+
diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py
new file mode 100644
index 00000000..e4b41e36
--- /dev/null
+++ b/inference_engine/distributed/prefill_worker.py
@@ -0,0 +1,475 @@
+"""Queued prefill-only worker jobs and gRPC service (ADR 0017)."""
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass, field
+from enum import IntEnum
+from typing import Protocol, Sequence
+
+import grpc
+
+from inference_engine.distributed.capability import (
+ CacheCompatibility,
+ CompressionCodec,
+)
+from inference_engine.distributed.prefill_auth import (
+ FleetAuthConfig,
+ PrefillAuthError,
+ signed_metadata,
+ verify_metadata,
+)
+from inference_engine.distributed.prefill_cache import (
+ CacheBlock,
+ PrefixCacheStore,
+ compatibility_fingerprint,
+)
+from inference_engine.server.proto_gen.kakeya.v1 import (
+ distributed_pb2,
+ distributed_pb2_grpc,
+)
+
+
+class PrefillJobState(IntEnum):
+ UNSPECIFIED = 0
+ QUEUED = 1
+ RUNNING = 2
+ COMPLETED = 3
+ FAILED = 4
+ CANCELLED = 5
+
+
+class PrefillComputeEngine(Protocol):
+ """Model-specific worker that produces the final restorable snapshot."""
+
+ def compute_prefill(
+ self,
+ token_ids: Sequence[int],
+ block_hashes: Sequence[bytes],
+ *,
+ compression: CompressionCodec,
+ cancelled: threading.Event,
+ ) -> Sequence[CacheBlock]: ...
+
+
+@dataclass
+class PrefillJob:
+ job_id: str
+ request_id: str
+ tenant_id: str
+ token_ids: tuple[int, ...]
+ block_hashes: tuple[bytes, ...]
+ compression: CompressionCodec
+ state: PrefillJobState = PrefillJobState.QUEUED
+ tokens_computed: int = 0
+ lease_id: str = ""
+ block_hash: bytes = b""
+ payload_sha256: bytes = b""
+ transfer_bytes: int = 0
+ failure_reason: str = ""
+ compute_ms: float = 0.0
+ created_at: float = field(default_factory=time.time)
+ finished_at: float = 0.0
+ cancelled: threading.Event = field(default_factory=threading.Event)
+ future: Future | None = field(default=None, repr=False)
+ request_digest: bytes = b""
+ deadline_at: float = 0.0
+
+
+class PrefillJobStore:
+ """Bounded, idempotent job queue around one or more prefill engines."""
+
+ def __init__(
+ self,
+ engine: PrefillComputeEngine,
+ cache_store: PrefixCacheStore,
+ *,
+ max_concurrent_jobs: int = 1,
+ max_jobs: int = 128,
+ completed_ttl_s: float = 600.0,
+ max_prompt_tokens: int = 131_072,
+ ) -> None:
+ if min(
+ max_concurrent_jobs,
+ max_jobs,
+ completed_ttl_s,
+ max_prompt_tokens,
+ ) <= 0:
+ raise ValueError("worker limits must be > 0")
+ self.engine = engine
+ self.cache_store = cache_store
+ self.max_concurrent_jobs = int(max_concurrent_jobs)
+ self.max_jobs = int(max_jobs)
+ self.completed_ttl_s = float(completed_ttl_s)
+ self.max_prompt_tokens = int(max_prompt_tokens)
+ self._jobs: dict[str, PrefillJob] = {}
+ self._requests: dict[tuple[str, str], str] = {}
+ self._lock = threading.RLock()
+ self._executor = ThreadPoolExecutor(
+ max_workers=self.max_concurrent_jobs,
+ thread_name_prefix="kakeya-prefill-worker",
+ )
+
+ def submit(
+ self,
+ *,
+ request_id: str,
+ tenant_id: str,
+ token_ids: Sequence[int],
+ block_hashes: Sequence[bytes],
+ compatibility: CacheCompatibility,
+ compression: CompressionCodec,
+ deadline_ms: int = 0,
+ ) -> PrefillJob:
+ if not request_id or not tenant_id:
+ raise ValueError("request_id and tenant_id must be non-empty")
+ if deadline_ms < 0:
+ raise ValueError("deadline_ms must be >= 0")
+ if compatibility != self.cache_store.compatibility:
+ raise ValueError("prefill worker compatibility mismatch")
+ if not token_ids or not block_hashes:
+ raise ValueError("token_ids and block_hashes must be non-empty")
+ if len(token_ids) > self.max_prompt_tokens:
+ raise ValueError("prefill prompt exceeds worker token limit")
+ expected_blocks = (
+ len(token_ids) + compatibility.block_size_tokens - 1
+ ) // compatibility.block_size_tokens
+ if len(block_hashes) != expected_blocks:
+ raise ValueError(
+ f"expected {expected_blocks} block hashes, got {len(block_hashes)}",
+ )
+ if any(len(bytes(block_hash)) != 32 for block_hash in block_hashes):
+ raise ValueError("every prefill block hash must be SHA-256 (32 bytes)")
+ digest = hashlib.sha256(
+ compatibility_fingerprint(compatibility)
+ + b"".join(int(token).to_bytes(4, "little") for token in token_ids)
+ + b"".join(bytes(h) for h in block_hashes)
+ + int(compression).to_bytes(2, "little")
+ + int(deadline_ms).to_bytes(8, "little", signed=False)
+ ).digest()
+ with self._lock:
+ self._gc_locked()
+ request_key = (tenant_id, request_id)
+ existing_id = self._requests.get(request_key)
+ if existing_id is not None:
+ existing = self._jobs[existing_id]
+ if existing.request_digest != digest:
+ raise ValueError(
+ "idempotency key reused with different prefill request",
+ )
+ return existing
+ active = sum(
+ job.state in (PrefillJobState.QUEUED, PrefillJobState.RUNNING)
+ for job in self._jobs.values()
+ )
+ if active >= self.max_jobs:
+ raise ValueError("prefill worker job queue is full")
+ job = PrefillJob(
+ job_id=uuid.uuid4().hex,
+ request_id=request_id,
+ tenant_id=tenant_id,
+ token_ids=tuple(int(token) for token in token_ids),
+ block_hashes=tuple(bytes(h) for h in block_hashes),
+ compression=compression,
+ request_digest=digest,
+ deadline_at=(
+ time.time() + deadline_ms / 1000.0
+ if deadline_ms > 0 else 0.0
+ ),
+ )
+ self._jobs[job.job_id] = job
+ self._requests[request_key] = job.job_id
+ job.future = self._executor.submit(self._run, job.job_id)
+ return job
+
+ def get(self, job_id: str, tenant_id: str) -> PrefillJob:
+ with self._lock:
+ self._gc_locked()
+ job = self._jobs.get(job_id)
+ if job is None or job.tenant_id != tenant_id:
+ raise KeyError(job_id)
+ return job
+
+ def cancel(self, job_id: str, tenant_id: str) -> bool:
+ with self._lock:
+ job = self.get(job_id, tenant_id)
+ if job.state in (
+ PrefillJobState.COMPLETED,
+ PrefillJobState.FAILED,
+ PrefillJobState.CANCELLED,
+ ):
+ return False
+ job.cancelled.set()
+ if job.future is not None and job.future.cancel():
+ job.state = PrefillJobState.CANCELLED
+ job.finished_at = time.time()
+ return True
+
+ def stats(self) -> tuple[int, int, float, int]:
+ with self._lock:
+ running = sum(j.state == PrefillJobState.RUNNING for j in self._jobs.values())
+ queued = sum(j.state == PrefillJobState.QUEUED for j in self._jobs.values())
+ queued_tokens = sum(
+ len(j.token_ids)
+ for j in self._jobs.values()
+ if j.state in (PrefillJobState.QUEUED, PrefillJobState.RUNNING)
+ )
+ load = min(1.0, (running + queued) / self.max_concurrent_jobs)
+ return running, queued, load, queued_tokens
+
+ def close(self) -> None:
+ self._executor.shutdown(wait=False, cancel_futures=True)
+
+ def _run(self, job_id: str) -> None:
+ with self._lock:
+ job = self._jobs[job_id]
+ if job.cancelled.is_set():
+ job.state = PrefillJobState.CANCELLED
+ job.finished_at = time.time()
+ return
+ job.state = PrefillJobState.RUNNING
+ started = time.perf_counter()
+ timer = None
+ if job.deadline_at:
+ remaining = job.deadline_at - time.time()
+ if remaining <= 0:
+ job.cancelled.set()
+ else:
+ timer = threading.Timer(remaining, job.cancelled.set)
+ timer.daemon = True
+ timer.start()
+ try:
+ blocks = tuple(self.engine.compute_prefill(
+ job.token_ids,
+ job.block_hashes,
+ compression=job.compression,
+ cancelled=job.cancelled,
+ ))
+ if job.cancelled.is_set():
+ raise InterruptedError("prefill job cancelled")
+ if len(blocks) != len(job.block_hashes):
+ raise RuntimeError(
+ "prefill engine must return one snapshot per block hash",
+ )
+ for block in blocks:
+ if job.cancelled.is_set():
+ raise InterruptedError("prefill job cancelled")
+ self.cache_store.put(block)
+ lease = self.cache_store.lookup(job.block_hashes)
+ if not lease.lease_id:
+ raise RuntimeError("computed snapshot was not discoverable")
+ if job.cancelled.is_set():
+ raise InterruptedError("prefill job cancelled")
+ with self._lock:
+ if job.cancelled.is_set():
+ raise InterruptedError("prefill job cancelled")
+ job.state = PrefillJobState.COMPLETED
+ job.tokens_computed = len(job.token_ids)
+ job.lease_id = lease.lease_id
+ final = blocks[-1]
+ job.block_hash = final.block_hash
+ job.payload_sha256 = final.payload_sha256
+ job.transfer_bytes = final.nbytes
+ except InterruptedError as exc:
+ with self._lock:
+ job.state = PrefillJobState.CANCELLED
+ job.failure_reason = str(exc)
+ except Exception as exc:
+ with self._lock:
+ job.state = PrefillJobState.FAILED
+ job.failure_reason = f"{type(exc).__name__}: {exc}"
+ finally:
+ if timer is not None:
+ timer.cancel()
+ with self._lock:
+ job.compute_ms = (time.perf_counter() - started) * 1000.0
+ job.finished_at = time.time()
+
+ def _gc_locked(self) -> None:
+ cutoff = time.time() - self.completed_ttl_s
+ for job_id, job in list(self._jobs.items()):
+ if job.finished_at and job.finished_at < cutoff:
+ self._jobs.pop(job_id, None)
+ self._requests.pop((job.tenant_id, job.request_id), None)
+ overflow = len(self._jobs) - self.max_jobs * 2
+ if overflow > 0:
+ finished = sorted(
+ (
+ job for job in self._jobs.values()
+ if job.finished_at
+ ),
+ key=lambda job: job.finished_at,
+ )
+ for job in finished[:overflow]:
+ self._jobs.pop(job.job_id, None)
+ self._requests.pop((job.tenant_id, job.request_id), None)
+
+
+class PrefillWorkerServiceServicer(
+ distributed_pb2_grpc.PrefillWorkerServiceServicer,
+):
+ def __init__(
+ self,
+ jobs: PrefillJobStore,
+ *,
+ node_id: str,
+ cache_address: str,
+ auth: FleetAuthConfig | None = None,
+ tokens_per_second_prefill: float = 0.0,
+ ) -> None:
+ self.jobs = jobs
+ self.node_id = node_id
+ self.cache_address = cache_address
+ self.auth = auth
+ self.tokens_per_second_prefill = float(tokens_per_second_prefill)
+
+ async def _authenticate(self, request, context) -> None:
+ if self.auth is None:
+ return
+ try:
+ verify_metadata(
+ context.invocation_metadata(),
+ request,
+ self.auth,
+ )
+ except PrefillAuthError as exc:
+ await context.abort(grpc.StatusCode.UNAUTHENTICATED, str(exc))
+
+ async def SubmitPrefillJob(self, request, context): # noqa: N802
+ await self._authenticate(request, context)
+ if self.auth is not None and request.tenant_id != self.auth.tenant_id:
+ await context.abort(
+ grpc.StatusCode.PERMISSION_DENIED,
+ "prefill job tenant mismatch",
+ )
+ try:
+ job = await asyncio.to_thread(
+ self.jobs.submit,
+ request_id=request.request_id,
+ tenant_id=request.tenant_id,
+ token_ids=list(request.token_ids),
+ block_hashes=list(request.block_hashes),
+ compatibility=CacheCompatibility.from_proto(request.compatibility),
+ compression=CompressionCodec(request.preferred_compression),
+ deadline_ms=request.deadline_ms,
+ )
+ except ValueError as exc:
+ await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(exc))
+ return distributed_pb2.SubmitPrefillJobResponse(
+ job_id=job.job_id,
+ status=int(job.state),
+ worker_node_id=self.node_id,
+ queue_eta_ms=(
+ self.jobs.stats()[3] / self.tokens_per_second_prefill * 1000.0
+ if self.tokens_per_second_prefill > 0 else 0.0
+ ),
+ )
+
+ async def GetPrefillJobStatus(self, request, context): # noqa: N802
+ await self._authenticate(request, context)
+ if self.auth is not None and request.tenant_id != self.auth.tenant_id:
+ await context.abort(grpc.StatusCode.PERMISSION_DENIED, "tenant mismatch")
+ try:
+ job = self.jobs.get(request.job_id, request.tenant_id)
+ except KeyError as exc:
+ await context.abort(grpc.StatusCode.NOT_FOUND, str(exc))
+ return _job_status_proto(job, self.cache_address)
+
+ async def CancelPrefillJob(self, request, context): # noqa: N802
+ await self._authenticate(request, context)
+ if self.auth is not None and request.tenant_id != self.auth.tenant_id:
+ await context.abort(grpc.StatusCode.PERMISSION_DENIED, "tenant mismatch")
+ try:
+ cancelled = self.jobs.cancel(request.job_id, request.tenant_id)
+ except KeyError as exc:
+ await context.abort(grpc.StatusCode.NOT_FOUND, str(exc))
+ return distributed_pb2.CancelPrefillJobResponse(cancelled=cancelled)
+
+
+def add_prefill_worker_service(
+ server: grpc.aio.Server,
+ jobs: PrefillJobStore,
+ *,
+ node_id: str,
+ cache_address: str,
+ auth: FleetAuthConfig | None = None,
+ tokens_per_second_prefill: float = 0.0,
+) -> PrefillWorkerServiceServicer:
+ servicer = PrefillWorkerServiceServicer(
+ jobs,
+ node_id=node_id,
+ cache_address=cache_address,
+ auth=auth,
+ tokens_per_second_prefill=tokens_per_second_prefill,
+ )
+ distributed_pb2_grpc.add_PrefillWorkerServiceServicer_to_server(
+ servicer,
+ server,
+ )
+ return servicer
+
+
+def submit_prefill_job_sync(
+ address: str,
+ request: distributed_pb2.SubmitPrefillJobRequest,
+ *,
+ timeout_s: float,
+ auth: FleetAuthConfig | None = None,
+) -> distributed_pb2.SubmitPrefillJobResponse:
+ metadata = signed_metadata(request, auth) if auth is not None else None
+ with grpc.insecure_channel(address) as channel:
+ stub = distributed_pb2_grpc.PrefillWorkerServiceStub(channel)
+ return stub.SubmitPrefillJob(request, timeout=timeout_s, metadata=metadata)
+
+
+def get_prefill_job_sync(
+ address: str,
+ request: distributed_pb2.GetPrefillJobStatusRequest,
+ *,
+ timeout_s: float,
+ auth: FleetAuthConfig | None = None,
+) -> distributed_pb2.GetPrefillJobStatusResponse:
+ metadata = signed_metadata(request, auth) if auth is not None else None
+ with grpc.insecure_channel(address) as channel:
+ stub = distributed_pb2_grpc.PrefillWorkerServiceStub(channel)
+ return stub.GetPrefillJobStatus(request, timeout=timeout_s, metadata=metadata)
+
+
+def cancel_prefill_job_sync(
+ address: str,
+ request: distributed_pb2.CancelPrefillJobRequest,
+ *,
+ timeout_s: float,
+ auth: FleetAuthConfig | None = None,
+) -> distributed_pb2.CancelPrefillJobResponse:
+ metadata = signed_metadata(request, auth) if auth is not None else None
+ with grpc.insecure_channel(address) as channel:
+ stub = distributed_pb2_grpc.PrefillWorkerServiceStub(channel)
+ return stub.CancelPrefillJob(
+ request,
+ timeout=timeout_s,
+ metadata=metadata,
+ )
+
+
+def _job_status_proto(
+ job: PrefillJob,
+ cache_address: str,
+) -> distributed_pb2.GetPrefillJobStatusResponse:
+ return distributed_pb2.GetPrefillJobStatusResponse(
+ job_id=job.job_id,
+ status=int(job.state),
+ tokens_computed=job.tokens_computed,
+ lease_id=job.lease_id,
+ block_hash=job.block_hash,
+ payload_sha256=job.payload_sha256,
+ transfer_bytes=job.transfer_bytes,
+ failure_reason=job.failure_reason,
+ compute_ms=job.compute_ms,
+ cache_address=cache_address,
+ )
+
diff --git a/inference_engine/network/state.py b/inference_engine/network/state.py
index 82a1be2f..0e194804 100644
--- a/inference_engine/network/state.py
+++ b/inference_engine/network/state.py
@@ -95,6 +95,7 @@ def nodes(self) -> list[dict[str, Any]]:
for card in self.registry.snapshot():
registration = registrations.get(card.node_id, {})
cache = card.caches[0] if card.caches else None
+ worker = card.prefill_workers[0] if card.prefill_workers else None
endpoint = sorted(
card.endpoints,
key=lambda item: item.priority,
@@ -132,6 +133,19 @@ def nodes(self) -> list[dict[str, Any]]:
}
if cache else None
),
+ "prefill_worker": (
+ {
+ "address": worker.worker_address,
+ "max_concurrent_jobs": worker.max_concurrent_jobs,
+ "inflight_jobs": worker.inflight_jobs,
+ "queued_jobs": worker.queued_jobs,
+ "queued_tokens": worker.queued_tokens,
+ "load": worker.load,
+ "tokens_per_second": worker.tokens_per_second_prefill,
+ "ram_bytes_free": worker.ram_bytes_free,
+ }
+ if worker else None
+ ),
"endpoint": (
{
"address": endpoint[0].address,
@@ -160,6 +174,7 @@ def nodes(self) -> list[dict[str, Any]]:
"memory_bytes": 0,
"models": [],
"cache": None,
+ "prefill_worker": None,
"endpoint": {
"address": registration["address"],
"network": "pending",
diff --git a/inference_engine/server/grpc_app.py b/inference_engine/server/grpc_app.py
index 02a093ad..3a7f45af 100644
--- a/inference_engine/server/grpc_app.py
+++ b/inference_engine/server/grpc_app.py
@@ -382,6 +382,7 @@ def create_grpc_server(
default_proposer_model_id: str = "",
prefill_cache_store: Optional[object] = None,
prefill_cache_address: str = "",
+ prefill_auth: Optional[object] = None,
) -> grpc.aio.Server:
"""Build, but do not start, a configured gRPC asyncio server.
@@ -455,6 +456,7 @@ def create_grpc_server(
server,
prefill_cache_store,
cache_address=prefill_cache_address or config.bind_address,
+ auth=prefill_auth,
)
_logger.info(
"gRPC PrefillCacheService enabled at %s",
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 699e21ec..7fef240f 100644
--- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py
+++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py
@@ -24,83 +24,103 @@
-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\"\xc6\x02\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\"[\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\"\xeb\x01\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\"\xf7\x01\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\"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\"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*\xc8\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\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\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*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')
_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'kakeya.v1.distributed_pb2', _globals)
if not _descriptor._USE_C_DESCRIPTORS:
DESCRIPTOR._loaded_options = None
- _globals['_CAPABILITYROLE']._serialized_start=3547
- _globals['_CAPABILITYROLE']._serialized_end=3747
+ _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['_MODELCAPABILITY']._serialized_start=42
_globals['_MODELCAPABILITY']._serialized_end=167
_globals['_NODECAPABILITY']._serialized_start=170
- _globals['_NODECAPABILITY']._serialized_end=496
- _globals['_NODEENDPOINT']._serialized_start=498
- _globals['_NODEENDPOINT']._serialized_end=589
- _globals['_CACHECOMPATIBILITY']._serialized_start=592
- _globals['_CACHECOMPATIBILITY']._serialized_end=827
- _globals['_CACHECAPABILITY']._serialized_start=830
- _globals['_CACHECAPABILITY']._serialized_end=1077
- _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_start=1079
- _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_end=1156
- _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_start=1158
- _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_end=1236
- _globals['_GETNODECAPABILITYREQUEST']._serialized_start=1238
- _globals['_GETNODECAPABILITYREQUEST']._serialized_end=1264
- _globals['_GETNODECAPABILITYRESPONSE']._serialized_start=1266
- _globals['_GETNODECAPABILITYRESPONSE']._serialized_end=1334
- _globals['_GETCACHESUMMARYREQUEST']._serialized_start=1336
- _globals['_GETCACHESUMMARYREQUEST']._serialized_end=1414
- _globals['_GETCACHESUMMARYRESPONSE']._serialized_start=1416
- _globals['_GETCACHESUMMARYRESPONSE']._serialized_end=1502
- _globals['_LOOKUPPREFIXREQUEST']._serialized_start=1504
- _globals['_LOOKUPPREFIXREQUEST']._serialized_end=1601
- _globals['_LOOKUPPREFIXRESPONSE']._serialized_start=1604
- _globals['_LOOKUPPREFIXRESPONSE']._serialized_end=1811
- _globals['_FETCHBLOCKSREQUEST']._serialized_start=1813
- _globals['_FETCHBLOCKSREQUEST']._serialized_end=1851
- _globals['_FETCHBLOCKSRESPONSE']._serialized_start=1854
- _globals['_FETCHBLOCKSRESPONSE']._serialized_end=2037
- _globals['_PUBLISHBLOCKREQUEST']._serialized_start=2040
- _globals['_PUBLISHBLOCKREQUEST']._serialized_end=2277
- _globals['_PUBLISHBLOCKRESPONSE']._serialized_start=2279
- _globals['_PUBLISHBLOCKRESPONSE']._serialized_end=2338
- _globals['_PROPOSEBLOCKREQUEST']._serialized_start=2340
- _globals['_PROPOSEBLOCKREQUEST']._serialized_end=2447
- _globals['_PROPOSEBLOCKRESPONSE']._serialized_start=2449
- _globals['_PROPOSEBLOCKRESPONSE']._serialized_end=2570
- _globals['_TENSOR']._serialized_start=2572
- _globals['_TENSOR']._serialized_end=2624
- _globals['_LAYERKV']._serialized_start=2626
- _globals['_LAYERKV']._serialized_end=2710
- _globals['_RESTOREREQUEST']._serialized_start=2713
- _globals['_RESTOREREQUEST']._serialized_end=2845
- _globals['_RESTORERESPONSE']._serialized_start=2847
- _globals['_RESTORERESPONSE']._serialized_end=2949
- _globals['_SEEDCONTEXTREQUEST']._serialized_start=2951
- _globals['_SEEDCONTEXTREQUEST']._serialized_end=3042
- _globals['_SEEDCONTEXTRESPONSE']._serialized_start=3044
- _globals['_SEEDCONTEXTRESPONSE']._serialized_end=3086
- _globals['_DRAFTBLOCKREQUEST']._serialized_start=3088
- _globals['_DRAFTBLOCKREQUEST']._serialized_end=3192
- _globals['_DRAFTBLOCKRESPONSE']._serialized_start=3194
- _globals['_DRAFTBLOCKRESPONSE']._serialized_end=3294
- _globals['_EXTENDCONTEXTREQUEST']._serialized_start=3296
- _globals['_EXTENDCONTEXTREQUEST']._serialized_end=3389
- _globals['_EXTENDCONTEXTRESPONSE']._serialized_start=3391
- _globals['_EXTENDCONTEXTRESPONSE']._serialized_end=3435
- _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_start=3437
- _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_end=3499
- _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_start=3501
- _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_end=3544
- _globals['_CAPABILITYSERVICE']._serialized_start=3750
- _globals['_CAPABILITYSERVICE']._serialized_end=3970
- _globals['_PROPOSERSERVICE']._serialized_start=3972
- _globals['_PROPOSERSERVICE']._serialized_end=4070
- _globals['_PREFILLCACHESERVICE']._serialized_start=4073
- _globals['_PREFILLCACHESERVICE']._serialized_end=4428
- _globals['_DFLASHPROPOSERSERVICE']._serialized_start=4431
- _globals['_DFLASHPROPOSERSERVICE']._serialized_end=4880
+ _globals['_NODECAPABILITY']._serialized_end=557
+ _globals['_NODEENDPOINT']._serialized_start=559
+ _globals['_NODEENDPOINT']._serialized_end=650
+ _globals['_CACHECOMPATIBILITY']._serialized_start=653
+ _globals['_CACHECOMPATIBILITY']._serialized_end=954
+ _globals['_CACHECAPABILITY']._serialized_start=957
+ _globals['_CACHECAPABILITY']._serialized_end=1290
+ _globals['_PREFILLWORKERCAPABILITY']._serialized_start=1293
+ _globals['_PREFILLWORKERCAPABILITY']._serialized_end=1595
+ _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_start=1597
+ _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_end=1674
+ _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_start=1676
+ _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_end=1754
+ _globals['_GETNODECAPABILITYREQUEST']._serialized_start=1756
+ _globals['_GETNODECAPABILITYREQUEST']._serialized_end=1782
+ _globals['_GETNODECAPABILITYRESPONSE']._serialized_start=1784
+ _globals['_GETNODECAPABILITYRESPONSE']._serialized_end=1852
+ _globals['_GETCACHESUMMARYREQUEST']._serialized_start=1854
+ _globals['_GETCACHESUMMARYREQUEST']._serialized_end=1932
+ _globals['_GETCACHESUMMARYRESPONSE']._serialized_start=1934
+ _globals['_GETCACHESUMMARYRESPONSE']._serialized_end=2020
+ _globals['_LOOKUPPREFIXREQUEST']._serialized_start=2022
+ _globals['_LOOKUPPREFIXREQUEST']._serialized_end=2119
+ _globals['_LOOKUPPREFIXRESPONSE']._serialized_start=2122
+ _globals['_LOOKUPPREFIXRESPONSE']._serialized_end=2329
+ _globals['_FETCHBLOCKSREQUEST']._serialized_start=2331
+ _globals['_FETCHBLOCKSREQUEST']._serialized_end=2369
+ _globals['_FETCHBLOCKSRESPONSE']._serialized_start=2372
+ _globals['_FETCHBLOCKSRESPONSE']._serialized_end=2555
+ _globals['_PUBLISHBLOCKREQUEST']._serialized_start=2558
+ _globals['_PUBLISHBLOCKREQUEST']._serialized_end=2795
+ _globals['_PUBLISHBLOCKRESPONSE']._serialized_start=2797
+ _globals['_PUBLISHBLOCKRESPONSE']._serialized_end=2856
+ _globals['_SUBMITPREFILLJOBREQUEST']._serialized_start=2859
+ _globals['_SUBMITPREFILLJOBREQUEST']._serialized_end=3099
+ _globals['_SUBMITPREFILLJOBRESPONSE']._serialized_start=3102
+ _globals['_SUBMITPREFILLJOBRESPONSE']._serialized_end=3235
+ _globals['_GETPREFILLJOBSTATUSREQUEST']._serialized_start=3237
+ _globals['_GETPREFILLJOBSTATUSREQUEST']._serialized_end=3300
+ _globals['_GETPREFILLJOBSTATUSRESPONSE']._serialized_start=3303
+ _globals['_GETPREFILLJOBSTATUSRESPONSE']._serialized_end=3571
+ _globals['_CANCELPREFILLJOBREQUEST']._serialized_start=3573
+ _globals['_CANCELPREFILLJOBREQUEST']._serialized_end=3633
+ _globals['_CANCELPREFILLJOBRESPONSE']._serialized_start=3635
+ _globals['_CANCELPREFILLJOBRESPONSE']._serialized_end=3680
+ _globals['_PROPOSEBLOCKREQUEST']._serialized_start=3682
+ _globals['_PROPOSEBLOCKREQUEST']._serialized_end=3789
+ _globals['_PROPOSEBLOCKRESPONSE']._serialized_start=3791
+ _globals['_PROPOSEBLOCKRESPONSE']._serialized_end=3912
+ _globals['_TENSOR']._serialized_start=3914
+ _globals['_TENSOR']._serialized_end=3966
+ _globals['_LAYERKV']._serialized_start=3968
+ _globals['_LAYERKV']._serialized_end=4052
+ _globals['_RESTOREREQUEST']._serialized_start=4055
+ _globals['_RESTOREREQUEST']._serialized_end=4187
+ _globals['_RESTORERESPONSE']._serialized_start=4189
+ _globals['_RESTORERESPONSE']._serialized_end=4291
+ _globals['_SEEDCONTEXTREQUEST']._serialized_start=4293
+ _globals['_SEEDCONTEXTREQUEST']._serialized_end=4384
+ _globals['_SEEDCONTEXTRESPONSE']._serialized_start=4386
+ _globals['_SEEDCONTEXTRESPONSE']._serialized_end=4428
+ _globals['_DRAFTBLOCKREQUEST']._serialized_start=4430
+ _globals['_DRAFTBLOCKREQUEST']._serialized_end=4534
+ _globals['_DRAFTBLOCKRESPONSE']._serialized_start=4536
+ _globals['_DRAFTBLOCKRESPONSE']._serialized_end=4636
+ _globals['_EXTENDCONTEXTREQUEST']._serialized_start=4638
+ _globals['_EXTENDCONTEXTREQUEST']._serialized_end=4731
+ _globals['_EXTENDCONTEXTRESPONSE']._serialized_start=4733
+ _globals['_EXTENDCONTEXTRESPONSE']._serialized_end=4777
+ _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_start=4779
+ _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
# @@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 8141a68c..02da80a0 100644
--- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi
+++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi
@@ -15,12 +15,38 @@ class CapabilityRole(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
CAPABILITY_ROLE_EMBEDDER: _ClassVar[CapabilityRole]
CAPABILITY_ROLE_TOOL: _ClassVar[CapabilityRole]
CAPABILITY_ROLE_PREFILL_CACHE: _ClassVar[CapabilityRole]
+ CAPABILITY_ROLE_PREFILL_COMPUTE: _ClassVar[CapabilityRole]
+
+class CompressionCodec(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+ __slots__ = ()
+ COMPRESSION_CODEC_UNSPECIFIED: _ClassVar[CompressionCodec]
+ COMPRESSION_CODEC_NONE: _ClassVar[CompressionCodec]
+ COMPRESSION_CODEC_ZLIB: _ClassVar[CompressionCodec]
+
+class PrefillJobStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+ __slots__ = ()
+ PREFILL_JOB_STATUS_UNSPECIFIED: _ClassVar[PrefillJobStatus]
+ PREFILL_JOB_STATUS_QUEUED: _ClassVar[PrefillJobStatus]
+ PREFILL_JOB_STATUS_RUNNING: _ClassVar[PrefillJobStatus]
+ PREFILL_JOB_STATUS_COMPLETED: _ClassVar[PrefillJobStatus]
+ PREFILL_JOB_STATUS_FAILED: _ClassVar[PrefillJobStatus]
+ PREFILL_JOB_STATUS_CANCELLED: _ClassVar[PrefillJobStatus]
CAPABILITY_ROLE_UNSPECIFIED: CapabilityRole
CAPABILITY_ROLE_VERIFIER: CapabilityRole
CAPABILITY_ROLE_PROPOSER: CapabilityRole
CAPABILITY_ROLE_EMBEDDER: CapabilityRole
CAPABILITY_ROLE_TOOL: CapabilityRole
CAPABILITY_ROLE_PREFILL_CACHE: CapabilityRole
+CAPABILITY_ROLE_PREFILL_COMPUTE: CapabilityRole
+COMPRESSION_CODEC_UNSPECIFIED: CompressionCodec
+COMPRESSION_CODEC_NONE: CompressionCodec
+COMPRESSION_CODEC_ZLIB: CompressionCodec
+PREFILL_JOB_STATUS_UNSPECIFIED: PrefillJobStatus
+PREFILL_JOB_STATUS_QUEUED: PrefillJobStatus
+PREFILL_JOB_STATUS_RUNNING: PrefillJobStatus
+PREFILL_JOB_STATUS_COMPLETED: PrefillJobStatus
+PREFILL_JOB_STATUS_FAILED: PrefillJobStatus
+PREFILL_JOB_STATUS_CANCELLED: PrefillJobStatus
class ModelCapability(_message.Message):
__slots__ = ("model_id", "role", "quantization", "tokens_per_second")
@@ -35,7 +61,7 @@ class ModelCapability(_message.Message):
def __init__(self, model_id: _Optional[str] = ..., role: _Optional[_Union[CapabilityRole, str]] = ..., quantization: _Optional[str] = ..., tokens_per_second: _Optional[float] = ...) -> None: ...
class NodeCapability(_message.Message):
- __slots__ = ("node_id", "grpc_address", "platform", "unified_memory_bytes", "mlx_version", "models", "announced_at_unix", "ttl_seconds", "ring_address", "caches", "endpoints")
+ __slots__ = ("node_id", "grpc_address", "platform", "unified_memory_bytes", "mlx_version", "models", "announced_at_unix", "ttl_seconds", "ring_address", "caches", "endpoints", "prefill_workers")
NODE_ID_FIELD_NUMBER: _ClassVar[int]
GRPC_ADDRESS_FIELD_NUMBER: _ClassVar[int]
PLATFORM_FIELD_NUMBER: _ClassVar[int]
@@ -47,6 +73,7 @@ class NodeCapability(_message.Message):
RING_ADDRESS_FIELD_NUMBER: _ClassVar[int]
CACHES_FIELD_NUMBER: _ClassVar[int]
ENDPOINTS_FIELD_NUMBER: _ClassVar[int]
+ PREFILL_WORKERS_FIELD_NUMBER: _ClassVar[int]
node_id: str
grpc_address: str
platform: str
@@ -58,7 +85,8 @@ class NodeCapability(_message.Message):
ring_address: str
caches: _containers.RepeatedCompositeFieldContainer[CacheCapability]
endpoints: _containers.RepeatedCompositeFieldContainer[NodeEndpoint]
- def __init__(self, node_id: _Optional[str] = ..., grpc_address: _Optional[str] = ..., platform: _Optional[str] = ..., unified_memory_bytes: _Optional[int] = ..., mlx_version: _Optional[str] = ..., models: _Optional[_Iterable[_Union[ModelCapability, _Mapping]]] = ..., announced_at_unix: _Optional[float] = ..., ttl_seconds: _Optional[float] = ..., ring_address: _Optional[str] = ..., caches: _Optional[_Iterable[_Union[CacheCapability, _Mapping]]] = ..., endpoints: _Optional[_Iterable[_Union[NodeEndpoint, _Mapping]]] = ...) -> None: ...
+ prefill_workers: _containers.RepeatedCompositeFieldContainer[PrefillWorkerCapability]
+ def __init__(self, node_id: _Optional[str] = ..., grpc_address: _Optional[str] = ..., platform: _Optional[str] = ..., unified_memory_bytes: _Optional[int] = ..., mlx_version: _Optional[str] = ..., models: _Optional[_Iterable[_Union[ModelCapability, _Mapping]]] = ..., announced_at_unix: _Optional[float] = ..., ttl_seconds: _Optional[float] = ..., ring_address: _Optional[str] = ..., caches: _Optional[_Iterable[_Union[CacheCapability, _Mapping]]] = ..., endpoints: _Optional[_Iterable[_Union[NodeEndpoint, _Mapping]]] = ..., prefill_workers: _Optional[_Iterable[_Union[PrefillWorkerCapability, _Mapping]]] = ...) -> None: ...
class NodeEndpoint(_message.Message):
__slots__ = ("address", "network", "priority", "measured_rtt_ms")
@@ -73,7 +101,7 @@ class NodeEndpoint(_message.Message):
def __init__(self, address: _Optional[str] = ..., network: _Optional[str] = ..., priority: _Optional[int] = ..., measured_rtt_ms: _Optional[float] = ...) -> None: ...
class CacheCompatibility(_message.Message):
- __slots__ = ("model_id", "model_revision", "tokenizer_revision", "cache_format_version", "quantization", "rope_hash", "layer_geometry_hash", "kv_dtype", "block_size_tokens")
+ __slots__ = ("model_id", "model_revision", "tokenizer_revision", "cache_format_version", "quantization", "rope_hash", "layer_geometry_hash", "kv_dtype", "block_size_tokens", "tenant_namespace", "sink_size", "window_size")
MODEL_ID_FIELD_NUMBER: _ClassVar[int]
MODEL_REVISION_FIELD_NUMBER: _ClassVar[int]
TOKENIZER_REVISION_FIELD_NUMBER: _ClassVar[int]
@@ -83,6 +111,9 @@ class CacheCompatibility(_message.Message):
LAYER_GEOMETRY_HASH_FIELD_NUMBER: _ClassVar[int]
KV_DTYPE_FIELD_NUMBER: _ClassVar[int]
BLOCK_SIZE_TOKENS_FIELD_NUMBER: _ClassVar[int]
+ TENANT_NAMESPACE_FIELD_NUMBER: _ClassVar[int]
+ SINK_SIZE_FIELD_NUMBER: _ClassVar[int]
+ WINDOW_SIZE_FIELD_NUMBER: _ClassVar[int]
model_id: str
model_revision: str
tokenizer_revision: str
@@ -92,10 +123,13 @@ class CacheCompatibility(_message.Message):
layer_geometry_hash: str
kv_dtype: str
block_size_tokens: int
- def __init__(self, model_id: _Optional[str] = ..., model_revision: _Optional[str] = ..., tokenizer_revision: _Optional[str] = ..., cache_format_version: _Optional[str] = ..., quantization: _Optional[str] = ..., rope_hash: _Optional[str] = ..., layer_geometry_hash: _Optional[str] = ..., kv_dtype: _Optional[str] = ..., block_size_tokens: _Optional[int] = ...) -> None: ...
+ tenant_namespace: str
+ sink_size: int
+ window_size: int
+ def __init__(self, model_id: _Optional[str] = ..., model_revision: _Optional[str] = ..., tokenizer_revision: _Optional[str] = ..., cache_format_version: _Optional[str] = ..., quantization: _Optional[str] = ..., rope_hash: _Optional[str] = ..., layer_geometry_hash: _Optional[str] = ..., kv_dtype: _Optional[str] = ..., block_size_tokens: _Optional[int] = ..., tenant_namespace: _Optional[str] = ..., sink_size: _Optional[int] = ..., window_size: _Optional[int] = ...) -> None: ...
class CacheCapability(_message.Message):
- __slots__ = ("compatibility", "cache_address", "cache_bytes_used", "cache_bytes_free", "entry_count", "cache_epoch", "load", "tokens_served", "bloom_filter")
+ __slots__ = ("compatibility", "cache_address", "cache_bytes_used", "cache_bytes_free", "entry_count", "cache_epoch", "load", "tokens_served", "bloom_filter", "default_compression", "replication_factor")
COMPATIBILITY_FIELD_NUMBER: _ClassVar[int]
CACHE_ADDRESS_FIELD_NUMBER: _ClassVar[int]
CACHE_BYTES_USED_FIELD_NUMBER: _ClassVar[int]
@@ -105,6 +139,8 @@ class CacheCapability(_message.Message):
LOAD_FIELD_NUMBER: _ClassVar[int]
TOKENS_SERVED_FIELD_NUMBER: _ClassVar[int]
BLOOM_FILTER_FIELD_NUMBER: _ClassVar[int]
+ DEFAULT_COMPRESSION_FIELD_NUMBER: _ClassVar[int]
+ REPLICATION_FACTOR_FIELD_NUMBER: _ClassVar[int]
compatibility: CacheCompatibility
cache_address: str
cache_bytes_used: int
@@ -114,7 +150,33 @@ class CacheCapability(_message.Message):
load: float
tokens_served: int
bloom_filter: bytes
- def __init__(self, compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ..., cache_address: _Optional[str] = ..., cache_bytes_used: _Optional[int] = ..., cache_bytes_free: _Optional[int] = ..., entry_count: _Optional[int] = ..., cache_epoch: _Optional[int] = ..., load: _Optional[float] = ..., tokens_served: _Optional[int] = ..., bloom_filter: _Optional[bytes] = ...) -> None: ...
+ default_compression: CompressionCodec
+ replication_factor: int
+ def __init__(self, compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ..., cache_address: _Optional[str] = ..., cache_bytes_used: _Optional[int] = ..., cache_bytes_free: _Optional[int] = ..., entry_count: _Optional[int] = ..., cache_epoch: _Optional[int] = ..., load: _Optional[float] = ..., tokens_served: _Optional[int] = ..., bloom_filter: _Optional[bytes] = ..., default_compression: _Optional[_Union[CompressionCodec, str]] = ..., replication_factor: _Optional[int] = ...) -> None: ...
+
+class PrefillWorkerCapability(_message.Message):
+ __slots__ = ("compatibility", "worker_address", "max_concurrent_jobs", "inflight_jobs", "queued_jobs", "load", "tokens_per_second_prefill", "ram_bytes_free", "accepts_compute_jobs", "queued_tokens")
+ COMPATIBILITY_FIELD_NUMBER: _ClassVar[int]
+ WORKER_ADDRESS_FIELD_NUMBER: _ClassVar[int]
+ MAX_CONCURRENT_JOBS_FIELD_NUMBER: _ClassVar[int]
+ INFLIGHT_JOBS_FIELD_NUMBER: _ClassVar[int]
+ QUEUED_JOBS_FIELD_NUMBER: _ClassVar[int]
+ LOAD_FIELD_NUMBER: _ClassVar[int]
+ TOKENS_PER_SECOND_PREFILL_FIELD_NUMBER: _ClassVar[int]
+ RAM_BYTES_FREE_FIELD_NUMBER: _ClassVar[int]
+ ACCEPTS_COMPUTE_JOBS_FIELD_NUMBER: _ClassVar[int]
+ QUEUED_TOKENS_FIELD_NUMBER: _ClassVar[int]
+ compatibility: CacheCompatibility
+ worker_address: str
+ max_concurrent_jobs: int
+ inflight_jobs: int
+ queued_jobs: int
+ load: float
+ tokens_per_second_prefill: float
+ ram_bytes_free: int
+ accepts_compute_jobs: bool
+ queued_tokens: int
+ def __init__(self, compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ..., worker_address: _Optional[str] = ..., max_concurrent_jobs: _Optional[int] = ..., inflight_jobs: _Optional[int] = ..., queued_jobs: _Optional[int] = ..., load: _Optional[float] = ..., tokens_per_second_prefill: _Optional[float] = ..., ram_bytes_free: _Optional[int] = ..., accepts_compute_jobs: _Optional[bool] = ..., queued_tokens: _Optional[int] = ...) -> None: ...
class ExchangeCapabilitiesRequest(_message.Message):
__slots__ = ("known_nodes",)
@@ -236,6 +298,82 @@ class PublishBlockResponse(_message.Message):
cache_epoch: int
def __init__(self, stored: _Optional[bool] = ..., cache_epoch: _Optional[int] = ...) -> None: ...
+class SubmitPrefillJobRequest(_message.Message):
+ __slots__ = ("request_id", "tenant_id", "compatibility", "token_ids", "block_hashes", "deadline_ms", "preferred_compression")
+ REQUEST_ID_FIELD_NUMBER: _ClassVar[int]
+ TENANT_ID_FIELD_NUMBER: _ClassVar[int]
+ COMPATIBILITY_FIELD_NUMBER: _ClassVar[int]
+ TOKEN_IDS_FIELD_NUMBER: _ClassVar[int]
+ BLOCK_HASHES_FIELD_NUMBER: _ClassVar[int]
+ DEADLINE_MS_FIELD_NUMBER: _ClassVar[int]
+ PREFERRED_COMPRESSION_FIELD_NUMBER: _ClassVar[int]
+ request_id: str
+ tenant_id: str
+ compatibility: CacheCompatibility
+ token_ids: _containers.RepeatedScalarFieldContainer[int]
+ block_hashes: _containers.RepeatedScalarFieldContainer[bytes]
+ deadline_ms: int
+ preferred_compression: CompressionCodec
+ def __init__(self, request_id: _Optional[str] = ..., tenant_id: _Optional[str] = ..., compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ..., token_ids: _Optional[_Iterable[int]] = ..., block_hashes: _Optional[_Iterable[bytes]] = ..., deadline_ms: _Optional[int] = ..., preferred_compression: _Optional[_Union[CompressionCodec, str]] = ...) -> None: ...
+
+class SubmitPrefillJobResponse(_message.Message):
+ __slots__ = ("job_id", "status", "worker_node_id", "queue_eta_ms")
+ JOB_ID_FIELD_NUMBER: _ClassVar[int]
+ STATUS_FIELD_NUMBER: _ClassVar[int]
+ WORKER_NODE_ID_FIELD_NUMBER: _ClassVar[int]
+ QUEUE_ETA_MS_FIELD_NUMBER: _ClassVar[int]
+ job_id: str
+ status: PrefillJobStatus
+ worker_node_id: str
+ queue_eta_ms: float
+ def __init__(self, job_id: _Optional[str] = ..., status: _Optional[_Union[PrefillJobStatus, str]] = ..., worker_node_id: _Optional[str] = ..., queue_eta_ms: _Optional[float] = ...) -> None: ...
+
+class GetPrefillJobStatusRequest(_message.Message):
+ __slots__ = ("job_id", "tenant_id")
+ JOB_ID_FIELD_NUMBER: _ClassVar[int]
+ TENANT_ID_FIELD_NUMBER: _ClassVar[int]
+ job_id: str
+ tenant_id: str
+ def __init__(self, job_id: _Optional[str] = ..., tenant_id: _Optional[str] = ...) -> None: ...
+
+class GetPrefillJobStatusResponse(_message.Message):
+ __slots__ = ("job_id", "status", "tokens_computed", "lease_id", "block_hash", "payload_sha256", "transfer_bytes", "failure_reason", "compute_ms", "cache_address")
+ JOB_ID_FIELD_NUMBER: _ClassVar[int]
+ STATUS_FIELD_NUMBER: _ClassVar[int]
+ TOKENS_COMPUTED_FIELD_NUMBER: _ClassVar[int]
+ LEASE_ID_FIELD_NUMBER: _ClassVar[int]
+ BLOCK_HASH_FIELD_NUMBER: _ClassVar[int]
+ PAYLOAD_SHA256_FIELD_NUMBER: _ClassVar[int]
+ TRANSFER_BYTES_FIELD_NUMBER: _ClassVar[int]
+ FAILURE_REASON_FIELD_NUMBER: _ClassVar[int]
+ COMPUTE_MS_FIELD_NUMBER: _ClassVar[int]
+ CACHE_ADDRESS_FIELD_NUMBER: _ClassVar[int]
+ job_id: str
+ status: PrefillJobStatus
+ tokens_computed: int
+ lease_id: str
+ block_hash: bytes
+ payload_sha256: bytes
+ transfer_bytes: int
+ failure_reason: str
+ compute_ms: float
+ cache_address: str
+ def __init__(self, job_id: _Optional[str] = ..., status: _Optional[_Union[PrefillJobStatus, str]] = ..., tokens_computed: _Optional[int] = ..., lease_id: _Optional[str] = ..., block_hash: _Optional[bytes] = ..., payload_sha256: _Optional[bytes] = ..., transfer_bytes: _Optional[int] = ..., failure_reason: _Optional[str] = ..., compute_ms: _Optional[float] = ..., cache_address: _Optional[str] = ...) -> None: ...
+
+class CancelPrefillJobRequest(_message.Message):
+ __slots__ = ("job_id", "tenant_id")
+ JOB_ID_FIELD_NUMBER: _ClassVar[int]
+ TENANT_ID_FIELD_NUMBER: _ClassVar[int]
+ job_id: str
+ tenant_id: str
+ def __init__(self, job_id: _Optional[str] = ..., tenant_id: _Optional[str] = ...) -> None: ...
+
+class CancelPrefillJobResponse(_message.Message):
+ __slots__ = ("cancelled",)
+ CANCELLED_FIELD_NUMBER: _ClassVar[int]
+ cancelled: bool
+ def __init__(self, cancelled: _Optional[bool] = ...) -> None: ...
+
class ProposeBlockRequest(_message.Message):
__slots__ = ("committed_token_ids", "block_size", "num_steps", "model_id")
COMMITTED_TOKEN_IDS_FIELD_NUMBER: _ClassVar[int]
diff --git a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2_grpc.py b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2_grpc.py
index 8885f712..dcfd367b 100644
--- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2_grpc.py
+++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2_grpc.py
@@ -468,6 +468,173 @@ def PublishBlock(request_iterator,
_registered_method=True)
+class PrefillWorkerServiceStub:
+ """PrefillWorkerService executes model prefill only. Workers load the same model
+ as the primary, materialize immutable snapshots into their co-located
+ PrefillCacheService, and never participate in autoregressive decode.
+ """
+
+ def __init__(self, channel):
+ """Constructor.
+
+ Args:
+ channel: A grpc.Channel.
+ """
+ self.SubmitPrefillJob = channel.unary_unary(
+ '/kakeya.v1.PrefillWorkerService/SubmitPrefillJob',
+ request_serializer=kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobRequest.SerializeToString,
+ response_deserializer=kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobResponse.FromString,
+ _registered_method=True)
+ self.GetPrefillJobStatus = channel.unary_unary(
+ '/kakeya.v1.PrefillWorkerService/GetPrefillJobStatus',
+ request_serializer=kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusRequest.SerializeToString,
+ response_deserializer=kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusResponse.FromString,
+ _registered_method=True)
+ self.CancelPrefillJob = channel.unary_unary(
+ '/kakeya.v1.PrefillWorkerService/CancelPrefillJob',
+ request_serializer=kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobRequest.SerializeToString,
+ response_deserializer=kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobResponse.FromString,
+ _registered_method=True)
+
+
+class PrefillWorkerServiceServicer:
+ """PrefillWorkerService executes model prefill only. Workers load the same model
+ as the primary, materialize immutable snapshots into their co-located
+ PrefillCacheService, and never participate in autoregressive decode.
+ """
+
+ def SubmitPrefillJob(self, request, context):
+ """Missing associated documentation comment in .proto file."""
+ context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+ context.set_details('Method not implemented!')
+ raise NotImplementedError('Method not implemented!')
+
+ def GetPrefillJobStatus(self, request, context):
+ """Missing associated documentation comment in .proto file."""
+ context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+ context.set_details('Method not implemented!')
+ raise NotImplementedError('Method not implemented!')
+
+ def CancelPrefillJob(self, request, context):
+ """Missing associated documentation comment in .proto file."""
+ context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+ context.set_details('Method not implemented!')
+ raise NotImplementedError('Method not implemented!')
+
+
+def add_PrefillWorkerServiceServicer_to_server(servicer, server):
+ rpc_method_handlers = {
+ 'SubmitPrefillJob': grpc.unary_unary_rpc_method_handler(
+ servicer.SubmitPrefillJob,
+ request_deserializer=kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobRequest.FromString,
+ response_serializer=kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobResponse.SerializeToString,
+ ),
+ 'GetPrefillJobStatus': grpc.unary_unary_rpc_method_handler(
+ servicer.GetPrefillJobStatus,
+ request_deserializer=kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusRequest.FromString,
+ response_serializer=kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusResponse.SerializeToString,
+ ),
+ 'CancelPrefillJob': grpc.unary_unary_rpc_method_handler(
+ servicer.CancelPrefillJob,
+ request_deserializer=kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobRequest.FromString,
+ response_serializer=kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobResponse.SerializeToString,
+ ),
+ }
+ generic_handler = grpc.method_handlers_generic_handler(
+ 'kakeya.v1.PrefillWorkerService', rpc_method_handlers)
+ server.add_generic_rpc_handlers((generic_handler,))
+ server.add_registered_method_handlers('kakeya.v1.PrefillWorkerService', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class PrefillWorkerService:
+ """PrefillWorkerService executes model prefill only. Workers load the same model
+ as the primary, materialize immutable snapshots into their co-located
+ PrefillCacheService, and never participate in autoregressive decode.
+ """
+
+ @staticmethod
+ def SubmitPrefillJob(request,
+ target,
+ options=(),
+ channel_credentials=None,
+ call_credentials=None,
+ insecure=False,
+ compression=None,
+ wait_for_ready=None,
+ timeout=None,
+ metadata=None):
+ return grpc.experimental.unary_unary(
+ request,
+ target,
+ '/kakeya.v1.PrefillWorkerService/SubmitPrefillJob',
+ kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobRequest.SerializeToString,
+ kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobResponse.FromString,
+ options,
+ channel_credentials,
+ insecure,
+ call_credentials,
+ compression,
+ wait_for_ready,
+ timeout,
+ metadata,
+ _registered_method=True)
+
+ @staticmethod
+ def GetPrefillJobStatus(request,
+ target,
+ options=(),
+ channel_credentials=None,
+ call_credentials=None,
+ insecure=False,
+ compression=None,
+ wait_for_ready=None,
+ timeout=None,
+ metadata=None):
+ return grpc.experimental.unary_unary(
+ request,
+ target,
+ '/kakeya.v1.PrefillWorkerService/GetPrefillJobStatus',
+ kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusRequest.SerializeToString,
+ kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusResponse.FromString,
+ options,
+ channel_credentials,
+ insecure,
+ call_credentials,
+ compression,
+ wait_for_ready,
+ timeout,
+ metadata,
+ _registered_method=True)
+
+ @staticmethod
+ def CancelPrefillJob(request,
+ target,
+ options=(),
+ channel_credentials=None,
+ call_credentials=None,
+ insecure=False,
+ compression=None,
+ wait_for_ready=None,
+ timeout=None,
+ metadata=None):
+ return grpc.experimental.unary_unary(
+ request,
+ target,
+ '/kakeya.v1.PrefillWorkerService/CancelPrefillJob',
+ kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobRequest.SerializeToString,
+ kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobResponse.FromString,
+ options,
+ channel_credentials,
+ insecure,
+ call_credentials,
+ compression,
+ wait_for_ready,
+ timeout,
+ metadata,
+ _registered_method=True)
+
+
class DFlashProposerServiceStub:
"""DFlashProposerService: stateful remote DFlash drafter + f_θ restoration.
Per turn: Restore (prompt -> f_θ-projected verifier K/V) then SeedContext
diff --git a/proto/kakeya/v1/distributed.proto b/proto/kakeya/v1/distributed.proto
index 8e51fa6b..cf9b679d 100644
--- a/proto/kakeya/v1/distributed.proto
+++ b/proto/kakeya/v1/distributed.proto
@@ -72,6 +72,15 @@ service PrefillCacheService {
rpc PublishBlock(stream PublishBlockRequest) returns (PublishBlockResponse);
}
+// PrefillWorkerService executes model prefill only. Workers load the same model
+// as the primary, materialize immutable snapshots into their co-located
+// PrefillCacheService, and never participate in autoregressive decode.
+service PrefillWorkerService {
+ rpc SubmitPrefillJob(SubmitPrefillJobRequest) returns (SubmitPrefillJobResponse);
+ rpc GetPrefillJobStatus(GetPrefillJobStatusRequest) returns (GetPrefillJobStatusResponse);
+ rpc CancelPrefillJob(CancelPrefillJobRequest) returns (CancelPrefillJobResponse);
+}
+
// -----------------------------------------------------------------------------
// Capability messages
// -----------------------------------------------------------------------------
@@ -96,6 +105,24 @@ enum CapabilityRole {
// The node can answer PrefillCacheService lookups and stream compatible
// immutable prefill K/V blocks.
CAPABILITY_ROLE_PREFILL_CACHE = 5;
+
+ // The node has the compatible model loaded and accepts prefill-only jobs.
+ CAPABILITY_ROLE_PREFILL_COMPUTE = 6;
+}
+
+enum CompressionCodec {
+ COMPRESSION_CODEC_UNSPECIFIED = 0;
+ COMPRESSION_CODEC_NONE = 1;
+ COMPRESSION_CODEC_ZLIB = 2;
+}
+
+enum PrefillJobStatus {
+ PREFILL_JOB_STATUS_UNSPECIFIED = 0;
+ PREFILL_JOB_STATUS_QUEUED = 1;
+ PREFILL_JOB_STATUS_RUNNING = 2;
+ PREFILL_JOB_STATUS_COMPLETED = 3;
+ PREFILL_JOB_STATUS_FAILED = 4;
+ PREFILL_JOB_STATUS_CANCELLED = 5;
}
// ModelCapability is one (model, role) a node offers.
@@ -164,6 +191,9 @@ message NodeCapability {
// Reachable interfaces ordered by operator preference (Thunderbolt before
// LAN/Tailscale). grpc_address remains the compatibility default.
repeated NodeEndpoint endpoints = 11;
+
+ // Prefill-only compute offerings. Cache-only nodes leave this empty.
+ repeated PrefillWorkerCapability prefill_workers = 12;
}
message NodeEndpoint {
@@ -183,6 +213,12 @@ message CacheCompatibility {
string layer_geometry_hash = 7;
string kv_dtype = 8;
uint32 block_size_tokens = 9;
+ // Tenant namespace is part of the compatibility fingerprint. It prevents a
+ // snapshot from being reused across isolation domains even with identical
+ // model/tokenizer geometry.
+ string tenant_namespace = 10;
+ uint32 sink_size = 11;
+ uint32 window_size = 12;
}
message CacheCapability {
@@ -196,6 +232,21 @@ message CacheCapability {
uint64 tokens_served = 8;
// Compact probabilistic summary. Empty means "query me directly".
bytes bloom_filter = 9;
+ CompressionCodec default_compression = 10;
+ uint32 replication_factor = 11;
+}
+
+message PrefillWorkerCapability {
+ CacheCompatibility compatibility = 1;
+ string worker_address = 2;
+ uint32 max_concurrent_jobs = 3;
+ uint32 inflight_jobs = 4;
+ uint32 queued_jobs = 5;
+ double load = 6;
+ double tokens_per_second_prefill = 7;
+ uint64 ram_bytes_free = 8;
+ bool accepts_compute_jobs = 9;
+ uint64 queued_tokens = 10;
}
message ExchangeCapabilitiesRequest {
@@ -277,6 +328,57 @@ message PublishBlockResponse {
uint64 cache_epoch = 2;
}
+// -----------------------------------------------------------------------------
+// Prefill-only compute jobs
+// -----------------------------------------------------------------------------
+
+message SubmitPrefillJobRequest {
+ // Caller-generated idempotency key. Repeating it returns the existing job.
+ string request_id = 1;
+ string tenant_id = 2;
+ CacheCompatibility compatibility = 3;
+ repeated uint32 token_ids = 4;
+ // Chained hashes computed by the primary (tenant-HMACed in authenticated
+ // mode). One hash per token block.
+ repeated bytes block_hashes = 5;
+ uint32 deadline_ms = 6;
+ CompressionCodec preferred_compression = 7;
+}
+
+message SubmitPrefillJobResponse {
+ string job_id = 1;
+ PrefillJobStatus status = 2;
+ string worker_node_id = 3;
+ double queue_eta_ms = 4;
+}
+
+message GetPrefillJobStatusRequest {
+ string job_id = 1;
+ string tenant_id = 2;
+}
+
+message GetPrefillJobStatusResponse {
+ string job_id = 1;
+ PrefillJobStatus status = 2;
+ uint32 tokens_computed = 3;
+ string lease_id = 4;
+ bytes block_hash = 5;
+ bytes payload_sha256 = 6;
+ uint64 transfer_bytes = 7;
+ string failure_reason = 8;
+ double compute_ms = 9;
+ string cache_address = 10;
+}
+
+message CancelPrefillJobRequest {
+ string job_id = 1;
+ string tenant_id = 2;
+}
+
+message CancelPrefillJobResponse {
+ bool cancelled = 1;
+}
+
// -----------------------------------------------------------------------------
// Remote proposal messages
// -----------------------------------------------------------------------------
diff --git a/scripts/mac_bridge/install_autorecover_launchagent.sh b/scripts/mac_bridge/install_autorecover_launchagent.sh
index a102c53b..14ac198e 100755
--- a/scripts/mac_bridge/install_autorecover_launchagent.sh
+++ b/scripts/mac_bridge/install_autorecover_launchagent.sh
@@ -1,20 +1,29 @@
#!/usr/bin/env bash
-# Install a user LaunchAgent that re-checks the mac-bridge runner
-# after reboot and periodically self-heals it.
+# Install a runner watchdog. `--system` installs a headless LaunchDaemon that
+# runs before user login (recommended for remote Mac minis). Without it, a
+# per-user LaunchAgent is installed as a fallback.
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
RECOVER_SCRIPT="${REPO_ROOT}/scripts/mac_bridge/recover_runner_after_reboot.sh"
-PLIST_DIR="${HOME}/Library/LaunchAgents"
-PLIST_PATH="${PLIST_DIR}/com.kakeya.mac-bridge-runner-autorecover.plist"
LABEL="com.kakeya.mac-bridge-runner-autorecover"
UID_NUM="$(id -u)"
+USER_NAME="$(id -un)"
+MODE="user"
+if [ "${1:-}" = "--system" ]; then
+ MODE="system"
+elif [ -n "${1:-}" ]; then
+ echo "usage: $0 [--system]" >&2
+ exit 2
+fi
-mkdir -p "$PLIST_DIR"
[ -x "$RECOVER_SCRIPT" ] || chmod +x "$RECOVER_SCRIPT"
+mkdir -p "${HOME}/actions-runner/_diag"
+TMP_PLIST="$(mktemp)"
+trap 'rm -f "$TMP_PLIST"' EXIT
-cat >"$PLIST_PATH" <"$TMP_PLIST" <
@@ -31,8 +40,18 @@ cat >"$PLIST_PATH" <WorkingDirectory
${REPO_ROOT}
+ EnvironmentVariables
+
+ HOME${HOME}
+ RUNNER_DIR${HOME}/actions-runner
+
+
RunAtLoad
+ KeepAlive
+
+ SuccessfulExit
+
StartInterval
60
@@ -44,10 +63,28 @@ cat >"$PLIST_PATH" <
EOF
-launchctl bootout "gui/${UID_NUM}" "${PLIST_PATH}" >/dev/null 2>&1 || true
-launchctl bootstrap "gui/${UID_NUM}" "${PLIST_PATH}"
-launchctl enable "gui/${UID_NUM}/${LABEL}" || true
-launchctl kickstart -k "gui/${UID_NUM}/${LABEL}" || true
+if [ "$MODE" = "system" ]; then
+ PLIST_PATH="/Library/LaunchDaemons/${LABEL}.plist"
+ # UserName lets the runner retain access to its registration, models and
+ # workspace while the daemon itself starts before GUI login.
+ /usr/libexec/PlistBuddy -c "Add :UserName string ${USER_NAME}" "$TMP_PLIST"
+ sudo install -o root -g wheel -m 644 "$TMP_PLIST" "$PLIST_PATH"
+ sudo launchctl bootout system "$PLIST_PATH" >/dev/null 2>&1 || true
+ sudo launchctl bootstrap system "$PLIST_PATH"
+ sudo launchctl enable "system/${LABEL}" || true
+ sudo launchctl kickstart -k "system/${LABEL}" || true
+ DOMAIN="system"
+else
+ PLIST_DIR="${HOME}/Library/LaunchAgents"
+ PLIST_PATH="${PLIST_DIR}/${LABEL}.plist"
+ mkdir -p "$PLIST_DIR"
+ install -m 644 "$TMP_PLIST" "$PLIST_PATH"
+ launchctl bootout "gui/${UID_NUM}" "$PLIST_PATH" >/dev/null 2>&1 || true
+ launchctl bootstrap "gui/${UID_NUM}" "$PLIST_PATH"
+ launchctl enable "gui/${UID_NUM}/${LABEL}" || true
+ launchctl kickstart -k "gui/${UID_NUM}/${LABEL}" || true
+ DOMAIN="gui/${UID_NUM}"
+fi
echo "[mac-bridge-autorecover] installed: ${PLIST_PATH}"
-echo "[mac-bridge-autorecover] label: ${LABEL}"
+echo "[mac-bridge-autorecover] label: ${DOMAIN}/${LABEL}"
diff --git a/scripts/mac_bridge/recover_runner_after_reboot.sh b/scripts/mac_bridge/recover_runner_after_reboot.sh
index c80e9204..6497d00e 100755
--- a/scripts/mac_bridge/recover_runner_after_reboot.sh
+++ b/scripts/mac_bridge/recover_runner_after_reboot.sh
@@ -12,37 +12,85 @@ set -euo pipefail
RUNNER_DIR="${RUNNER_DIR:-$HOME/actions-runner}"
LOG_DIR="${RUNNER_LOG_DIR:-$HOME/actions-runner/_diag}"
mkdir -p "$LOG_DIR"
+LOCK_DIR="${TMPDIR:-/tmp}/kakeya-runner-recover.lock"
log() { echo "[mac-bridge-recover] $*" >&2; }
-if pgrep -f "Runner.Listener.*${RUNNER_DIR}" >/dev/null 2>&1; then
+listener_running() {
+ pgrep -f "${RUNNER_DIR}/bin/Runner.Listener" >/dev/null 2>&1 \
+ || pgrep -f "Runner.Listener run" >/dev/null 2>&1
+}
+
+if ! mkdir "$LOCK_DIR" 2>/dev/null; then
+ lock_pid="$(cat "$LOCK_DIR/pid" 2>/dev/null || true)"
+ if [ -n "$lock_pid" ] && kill -0 "$lock_pid" 2>/dev/null; then
+ log "another recovery attempt is active (pid=$lock_pid)."
+ exit 0
+ fi
+ rm -rf "$LOCK_DIR"
+ mkdir "$LOCK_DIR"
+fi
+echo "$$" >"$LOCK_DIR/pid"
+trap 'rm -rf "$LOCK_DIR" 2>/dev/null || true' EXIT
+
+if listener_running; then
log "Runner.Listener already running."
exit 0
fi
if [ ! -x "${RUNNER_DIR}/run.sh" ]; then
- log "runner not found at ${RUNNER_DIR}; skipping."
- exit 0
+ log "runner not found at ${RUNNER_DIR}."
+ exit 1
+fi
+if [ ! -s "${RUNNER_DIR}/.runner" ]; then
+ log "runner exists but is not registered (.runner missing)."
+ exit 1
fi
-if [ -x "${RUNNER_DIR}/svc.sh" ] && "${RUNNER_DIR}/svc.sh" status 2>/dev/null | grep -q "installed"; then
- log "service installed; attempting svc.sh start"
- if "${RUNNER_DIR}/svc.sh" start >/dev/null 2>&1; then
- sleep 2
- if pgrep -f "Runner.Listener.*${RUNNER_DIR}" >/dev/null 2>&1; then
- log "runner started via service."
- exit 0
+if ! curl -fsSI --max-time 10 https://github.com/ >/dev/null 2>&1; then
+ log "github.com is unreachable; leaving runner stopped for the next retry."
+ exit 1
+fi
+
+if [ -x "${RUNNER_DIR}/svc.sh" ]; then
+ status="$("${RUNNER_DIR}/svc.sh" status 2>&1 || true)"
+ # Do not match the phrase "not installed".
+ if echo "$status" | grep -qi "installed" \
+ && ! echo "$status" | grep -qi "not installed"; then
+ log "official service is installed; attempting svc.sh start"
+ if [ "$(id -u)" -eq 0 ]; then
+ "${RUNNER_DIR}/svc.sh" start >/dev/null 2>&1 || true
+ elif sudo -n true >/dev/null 2>&1; then
+ sudo -n "${RUNNER_DIR}/svc.sh" start >/dev/null 2>&1 || true
fi
+ for _ in 1 2 3 4 5; do
+ sleep 2
+ if listener_running; then
+ log "runner started via official service."
+ exit 0
+ fi
+ done
+ log "official service did not bring up listener; using direct fallback."
fi
- log "service start did not bring up listener, falling back to run.sh"
fi
ts="$(date +%Y%m%d_%H%M%S)"
-nohup "${RUNNER_DIR}/run.sh" >"${LOG_DIR}/runner-nohup-${ts}.log" 2>&1 &
-sleep 2
-if pgrep -f "Runner.Listener.*${RUNNER_DIR}" >/dev/null 2>&1; then
- log "runner started via nohup run.sh fallback."
- exit 0
+(
+ cd "$RUNNER_DIR"
+ nohup ./run.sh >"${LOG_DIR}/runner-nohup-${ts}.log" 2>&1 &
+)
+for _ in 1 2 3 4 5 6 7 8; do
+ sleep 2
+ if listener_running; then
+ log "runner started via direct run.sh fallback."
+ exit 0
+ fi
+done
+
+latest_log="$(ls -t "$LOG_DIR"/runner-nohup-*.log 2>/dev/null | head -1 || true)"
+if [ -n "$latest_log" ]; then
+ log "last runner output:"
+ tail -20 "$latest_log" >&2 || true
fi
log "failed to start runner listener."
diff --git a/scripts/mac_bridge/runner_healthcheck.sh b/scripts/mac_bridge/runner_healthcheck.sh
new file mode 100755
index 00000000..e49967ce
--- /dev/null
+++ b/scripts/mac_bridge/runner_healthcheck.sh
@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+# Diagnose why GitHub Actions jobs are queued for kakeya-mac-m4.
+set -uo pipefail
+
+RUNNER_DIR="${RUNNER_DIR:-$HOME/actions-runner}"
+LABEL="com.kakeya.mac-bridge-runner-autorecover"
+failed=0
+
+check() {
+ local name="$1"; shift
+ if "$@" >/dev/null 2>&1; then
+ printf 'PASS %s\n' "$name"
+ else
+ printf 'FAIL %s\n' "$name"
+ failed=1
+ fi
+}
+
+echo "runner_dir=$RUNNER_DIR"
+check "runner registration (.runner)" test -s "$RUNNER_DIR/.runner"
+check "runner executable" test -x "$RUNNER_DIR/run.sh"
+check "github.com connectivity" curl -fsSI --max-time 10 https://github.com/
+check "Runner.Listener process" sh -c \
+ "pgrep -f '$RUNNER_DIR/bin/Runner.Listener' >/dev/null || pgrep -f 'Runner.Listener run' >/dev/null"
+check "disk has >=10 GiB free" sh -c \
+ "[ \$(df -k '$RUNNER_DIR' | awk 'NR==2 {print \$4}') -ge 10485760 ]"
+
+echo
+echo "official service:"
+"$RUNNER_DIR/svc.sh" status 2>&1 || true
+echo
+echo "user watchdog:"
+launchctl print "gui/$(id -u)/$LABEL" 2>&1 | head -30 || true
+echo
+echo "system watchdog:"
+sudo -n launchctl print "system/$LABEL" 2>&1 | head -30 || true
+echo
+echo "recent runner diagnostics:"
+ls -t "$RUNNER_DIR"/_diag/Runner_*.log \
+ "$RUNNER_DIR"/_diag/runner-nohup-*.log 2>/dev/null \
+ | head -2 | while read -r log; do
+ echo "=== $log ==="
+ tail -20 "$log"
+ done
+
+exit "$failed"
+
diff --git a/scripts/start_grpc_runtime_server.py b/scripts/start_grpc_runtime_server.py
index d5529ade..acda21bd 100755
--- a/scripts/start_grpc_runtime_server.py
+++ b/scripts/start_grpc_runtime_server.py
@@ -213,6 +213,10 @@ 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
+ ),
+ replication_factor=args.cache_replication_factor,
),
)
@@ -253,6 +257,8 @@ async def _exchange_loop(
*,
cache_store=None,
cache_address: str = "",
+ cache_compression: int = 1,
+ cache_replication_factor: int = 1,
) -> None:
"""Periodic gossip with seed peers until the task is cancelled."""
from inference_engine.distributed.exchange import exchange_once
@@ -269,6 +275,8 @@ async def _exchange_loop(
cache_capability(
cache_store,
cache_address=cache_address,
+ default_compression=cache_compression,
+ replication_factor=cache_replication_factor,
),
),
)
@@ -369,15 +377,24 @@ async def _serve(args: argparse.Namespace) -> int:
prefill_store = None
prefill_hook = None
+ prefill_auth = None
+ capability_registry_holder = [None]
if args.enable_prefill_cache:
if args.backend != "mlx":
raise SystemExit("--enable-prefill-cache currently requires --backend mlx")
import hashlib
- from inference_engine.distributed.capability import CacheCompatibility
+ from inference_engine.distributed.capability import (
+ CacheCompatibility,
+ CompressionCodec,
+ )
+ from inference_engine.distributed.prefill_auth import FleetAuthConfig
from inference_engine.distributed.prefill_cache import PrefixCacheStore
from inference_engine.distributed.prefill_cache_runtime import (
DistributedPrefillCacheHook,
)
+ from inference_engine.distributed.prefill_scheduler import (
+ PrefillCostConfig,
+ )
geometry = f"{num_layers}:{num_kv_heads}:{head_dim}"
compatibility = CacheCompatibility(
@@ -390,6 +407,9 @@ async def _serve(args: argparse.Namespace) -> int:
layer_geometry_hash=hashlib.sha256(geometry.encode()).hexdigest(),
kv_dtype=args.cache_kv_dtype,
block_size_tokens=args.cache_block_tokens,
+ tenant_namespace=args.cache_tenant_id,
+ sink_size=args.sink,
+ window_size=args.window,
)
prefill_store = PrefixCacheStore(
compatibility,
@@ -397,11 +417,40 @@ async def _serve(args: argparse.Namespace) -> int:
node_id=args.node_id or (__import__("platform").node() or "localhost"),
)
telemetry_callback = _build_token_telemetry_callback(args)
+ if args.fleet_psk_file:
+ prefill_auth = FleetAuthConfig.from_file(
+ args.fleet_psk_file,
+ tenant_id=args.cache_tenant_id,
+ node_id=args.node_id or "primary",
+ )
prefill_hook = DistributedPrefillCacheHook(
prefill_store,
peers=args.cache_peer,
+ registry_provider=lambda: (
+ capability_registry_holder[0].snapshot()
+ if capability_registry_holder[0] is not None else ()
+ ),
lookup_timeout_s=args.cache_lookup_timeout_s,
fetch_timeout_s=args.cache_fetch_timeout_s,
+ worker_timeout_s=args.prefill_worker_timeout_s,
+ 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
+ ),
+ replication_factor=args.cache_replication_factor,
+ cost_config=PrefillCostConfig(
+ local_prefill_tps=args.local_prefill_tps,
+ default_worker_tps=args.worker_prefill_tps,
+ link_mbps=args.cache_link_mbps,
+ default_rtt_ms=args.cache_default_rtt_ms,
+ minimum_savings_ratio=args.prefill_min_savings_ratio,
+ primary_compute_penalty_ms=args.primary_prefill_penalty_ms,
+ ),
+ auth=prefill_auth,
on_reuse=(
(lambda count: telemetry_callback(count, count))
if telemetry_callback is not None else None
@@ -467,6 +516,7 @@ async def _serve(args: argparse.Namespace) -> int:
backend=args.backend,
cache_store=prefill_store,
)
+ capability_registry_holder[0] = registry
if args.serve_ngram_proposer:
from inference_engine.distributed.capability import NGRAM_MODEL_ID
from inference_engine.distributed.ngram import NGramProposer
@@ -486,6 +536,7 @@ async def _serve(args: argparse.Namespace) -> int:
proposers=proposers,
prefill_cache_store=prefill_store,
prefill_cache_address=args.cache_advertise or args.advertise or args.bind,
+ prefill_auth=prefill_auth,
)
await server.start()
@@ -532,6 +583,10 @@ 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_replication_factor=args.cache_replication_factor,
),
)
_LOG.info(
@@ -563,6 +618,8 @@ def _on_signal(sig: int) -> None:
http_server.should_exit = True
if http_task is not None:
await http_task
+ if prefill_hook is not None:
+ prefill_hook.close()
await server.stop(grace=args.shutdown_grace_s)
_LOG.info("kakeya gRPC RuntimeService stopped cleanly")
return 0
@@ -648,7 +705,7 @@ def main() -> int:
"Defaults to --advertise/--bind.")
ap.add_argument("--cache-block-tokens", type=int, default=64,
help="Token boundary interval for restorable snapshots.")
- ap.add_argument("--cache-format-version", default="kakeya-prefill-v1")
+ ap.add_argument("--cache-format-version", default="kakeya-prefill-v2-zlib")
ap.add_argument("--cache-model-id", default="",
help="Logical cache model id; defaults to --verifier-id. "
"Use this when verifier-id is a host-specific path.")
@@ -663,6 +720,28 @@ def main() -> int:
help="RoPE/position configuration fingerprint.")
ap.add_argument("--cache-lookup-timeout-s", type=float, default=2.0)
ap.add_argument("--cache-fetch-timeout-s", type=float, default=30.0)
+ ap.add_argument("--cache-tenant-id", default="default",
+ help="Tenant namespace included in cache compatibility.")
+ 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"],
+ default="zlib")
+ ap.add_argument("--cache-replication-factor", type=int, default=1)
+ ap.add_argument("--cache-max-import-gb", type=float, default=1.0,
+ help="Reject remote snapshots whose wire or expanded size "
+ "would exceed this import budget.")
+ ap.add_argument("--cache-estimated-bytes-per-token", type=int, default=400000)
+ ap.add_argument("--cache-link-mbps", type=float, default=1000.0)
+ ap.add_argument("--cache-default-rtt-ms", type=float, default=2.0)
+ ap.add_argument("--local-prefill-tps", type=float, default=20.0)
+ ap.add_argument("--worker-prefill-tps", type=float, default=20.0)
+ ap.add_argument("--prefill-min-savings-ratio", type=float, default=0.10)
+ ap.add_argument("--primary-prefill-penalty-ms", type=float, default=5000.0,
+ help="Opportunity cost assigned to blocking the primary "
+ "with prefill; drives work to compute peers.")
+ ap.add_argument("--remote-prefill-min-tokens", type=int, default=128)
+ ap.add_argument("--prefill-worker-timeout-s", type=float, default=120.0)
ap.add_argument("--network-label", default="lan",
help="Advertised interface: thunderbolt|lan|tailscale|public.")
ap.add_argument("--network-priority", type=int, default=50)
diff --git a/scripts/start_prefill_cache_node.py b/scripts/start_prefill_cache_node.py
index 4a55f3ae..6778e5d3 100644
--- a/scripts/start_prefill_cache_node.py
+++ b/scripts/start_prefill_cache_node.py
@@ -24,6 +24,7 @@
from inference_engine.distributed.capability import (
CacheCompatibility,
+ CompressionCodec,
CapabilityRegistry,
CapabilityRole,
ModelCapability,
@@ -35,6 +36,7 @@
exchange_once,
)
from inference_engine.distributed.prefill_cache import PrefixCacheStore
+from inference_engine.distributed.prefill_auth import FleetAuthConfig
from inference_engine.distributed.prefill_cache_service import (
add_prefill_cache_service,
cache_capability,
@@ -62,12 +64,23 @@ async def serve(args) -> None:
layer_geometry_hash=args.layer_geometry_hash,
kv_dtype=args.kv_dtype,
block_size_tokens=args.block_size_tokens,
+ tenant_namespace=args.tenant_id,
+ sink_size=args.sink,
+ window_size=args.window,
)
store = PrefixCacheStore(
compatibility,
max_bytes=int(args.cache_gb * (1 << 30)),
node_id=args.node_id,
)
+ auth = (
+ FleetAuthConfig.from_file(
+ args.fleet_psk_file,
+ tenant_id=args.tenant_id,
+ node_id=args.node_id,
+ )
+ if args.fleet_psk_file else None
+ )
card = NodeCapability(
node_id=args.node_id,
grpc_address=args.advertise,
@@ -82,7 +95,16 @@ async def serve(args) -> None:
),
announced_at_unix=time.time(),
ttl_seconds=args.ttl_seconds,
- caches=(cache_capability(store, cache_address=args.advertise),),
+ caches=(cache_capability(
+ store,
+ cache_address=args.advertise,
+ default_compression=(
+ CompressionCodec.ZLIB
+ if args.cache_compression == "zlib"
+ else CompressionCodec.NONE
+ ),
+ replication_factor=args.replication_factor,
+ ),),
endpoints=(
NodeEndpoint(
args.advertise,
@@ -99,6 +121,7 @@ async def serve(args) -> None:
grpc_server,
store,
cache_address=args.advertise,
+ auth=auth,
)
grpc_server.add_insecure_port(args.bind)
await grpc_server.start()
@@ -169,12 +192,19 @@ def main() -> None:
ap.add_argument("--model-id", required=True)
ap.add_argument("--model-revision", default="")
ap.add_argument("--tokenizer-revision", default="")
- ap.add_argument("--cache-format-version", default="kakeya-prefill-v1")
+ ap.add_argument("--cache-format-version", default="kakeya-prefill-v2-zlib")
ap.add_argument("--quantization", default="")
ap.add_argument("--rope-hash", default="")
ap.add_argument("--layer-geometry-hash", default="")
ap.add_argument("--kv-dtype", default="bfloat16")
ap.add_argument("--block-size-tokens", type=int, default=64)
+ ap.add_argument("--tenant-id", default="default")
+ 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"],
+ default="zlib")
+ ap.add_argument("--replication-factor", type=int, default=1)
ap.add_argument("--cache-gb", type=float, default=4)
ap.add_argument("--platform", default="")
ap.add_argument("--memory-bytes", type=int, default=0)
diff --git a/scripts/start_prefill_worker_node.py b/scripts/start_prefill_worker_node.py
new file mode 100644
index 00000000..dda9138c
--- /dev/null
+++ b/scripts/start_prefill_worker_node.py
@@ -0,0 +1,252 @@
+"""Start a prefill-only MLX compute worker with a co-located RAM cache.
+
+The worker loads the same model as the primary, accepts PrefillWorkerService
+jobs, stores immutable snapshots, and never serves user decode.
+"""
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import platform
+import signal
+import sys
+import time
+from dataclasses import replace
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+import grpc
+import torch
+
+from inference_engine.backends.mlx.prefill_worker import MLXPrefillComputeEngine
+from inference_engine.backends.mlx.verifier import MLXSinkWindowVerifier
+from inference_engine.distributed.capability import (
+ CacheCompatibility,
+ CapabilityRegistry,
+ CapabilityRole,
+ CompressionCodec,
+ ModelCapability,
+ NodeCapability,
+ NodeEndpoint,
+ PrefillWorkerCapability,
+)
+from inference_engine.distributed.exchange import (
+ add_capability_service,
+ exchange_once,
+)
+from inference_engine.distributed.prefill_auth import FleetAuthConfig
+from inference_engine.distributed.prefill_cache import PrefixCacheStore
+from inference_engine.distributed.prefill_cache_service import (
+ add_prefill_cache_service,
+ cache_capability,
+)
+from inference_engine.distributed.prefill_worker import (
+ PrefillJobStore,
+ add_prefill_worker_service,
+)
+from kv_cache_proposer.verifier import VerifierConfig
+
+_LOG = logging.getLogger("kakeya.prefill-worker")
+
+
+def physical_memory_bytes() -> int:
+ try:
+ import os
+ return int(os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES"))
+ except (ValueError, OSError, AttributeError):
+ return 0
+
+
+async def serve(args) -> None:
+ compatibility = CacheCompatibility(
+ model_id=args.cache_model_id or args.model_id,
+ model_revision=args.model_revision,
+ tokenizer_revision=args.tokenizer_revision,
+ cache_format_version=args.cache_format_version,
+ quantization=args.quantization,
+ rope_hash=args.rope_hash,
+ layer_geometry_hash=args.layer_geometry_hash,
+ kv_dtype=args.kv_dtype,
+ block_size_tokens=args.block_size_tokens,
+ tenant_namespace=args.tenant_id,
+ sink_size=args.sink,
+ window_size=args.window,
+ )
+ auth = (
+ FleetAuthConfig.from_file(
+ args.fleet_psk_file,
+ tenant_id=args.tenant_id,
+ node_id=args.node_id,
+ )
+ if args.fleet_psk_file else None
+ )
+ verifier = MLXSinkWindowVerifier(VerifierConfig(
+ model_id=args.model_id,
+ sink_size=args.sink,
+ window_size=args.window,
+ dtype=torch.bfloat16,
+ device="cpu",
+ ))
+ store = PrefixCacheStore(
+ compatibility,
+ max_bytes=int(args.cache_gb * (1 << 30)),
+ node_id=args.node_id,
+ )
+ engine = MLXPrefillComputeEngine(verifier, compatibility)
+ jobs = PrefillJobStore(
+ engine,
+ store,
+ max_concurrent_jobs=args.max_concurrent_jobs,
+ max_jobs=args.max_jobs,
+ completed_ttl_s=args.job_ttl_s,
+ max_prompt_tokens=args.max_prompt_tokens,
+ )
+
+ def card() -> NodeCapability:
+ inflight, queued, load, queued_tokens = jobs.stats()
+ worker = PrefillWorkerCapability(
+ compatibility=compatibility,
+ worker_address=args.advertise,
+ max_concurrent_jobs=args.max_concurrent_jobs,
+ inflight_jobs=inflight,
+ queued_jobs=queued,
+ load=load,
+ tokens_per_second_prefill=args.prefill_tps,
+ ram_bytes_free=max(
+ 0,
+ physical_memory_bytes() - store.stats().bytes_used,
+ ),
+ queued_tokens=queued_tokens,
+ )
+ return NodeCapability(
+ node_id=args.node_id,
+ grpc_address=args.advertise,
+ platform=f"{platform.system()}-{platform.machine()}",
+ unified_memory_bytes=physical_memory_bytes(),
+ models=(
+ ModelCapability(
+ args.cache_model_id or args.model_id,
+ CapabilityRole.PREFILL_COMPUTE,
+ args.quantization,
+ args.prefill_tps,
+ ),
+ ),
+ announced_at_unix=time.time(),
+ ttl_seconds=args.ttl_seconds,
+ caches=(cache_capability(
+ store,
+ cache_address=args.advertise,
+ load=load,
+ default_compression=(
+ CompressionCodec.ZLIB
+ if args.cache_compression == "zlib"
+ else CompressionCodec.NONE
+ ),
+ replication_factor=args.replication_factor,
+ ),),
+ endpoints=(
+ NodeEndpoint(
+ args.advertise,
+ args.network,
+ args.priority,
+ args.rtt_ms,
+ ),
+ ),
+ prefill_workers=(worker,),
+ )
+
+ registry = CapabilityRegistry(card())
+ server = grpc.aio.server(
+ maximum_concurrent_rpcs=args.max_concurrent_rpcs,
+ )
+ add_capability_service(server, registry)
+ add_prefill_cache_service(
+ server,
+ store,
+ cache_address=args.advertise,
+ auth=auth,
+ )
+ add_prefill_worker_service(
+ server,
+ jobs,
+ node_id=args.node_id,
+ cache_address=args.advertise,
+ auth=auth,
+ tokens_per_second_prefill=args.prefill_tps,
+ )
+ server.add_insecure_port(args.bind)
+ await server.start()
+ _LOG.info("prefill worker ready on %s for %s", args.bind, args.model_id)
+
+ stop = asyncio.Event()
+ loop = asyncio.get_running_loop()
+ for sig in (signal.SIGINT, signal.SIGTERM):
+ loop.add_signal_handler(sig, stop.set)
+
+ async def gossip() -> None:
+ while not stop.is_set():
+ registry.self_card = replace(card(), announced_at_unix=time.time())
+ if args.peer:
+ await exchange_once(registry, args.peer, timeout_s=args.gossip_timeout_s)
+ try:
+ await asyncio.wait_for(stop.wait(), timeout=args.gossip_interval_s)
+ except asyncio.TimeoutError:
+ pass
+
+ gossip_task = asyncio.create_task(gossip())
+ await stop.wait()
+ gossip_task.cancel()
+ jobs.close()
+ await server.stop(grace=2.0)
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--node-id", required=True)
+ parser.add_argument("--bind", default="127.0.0.1:53051")
+ parser.add_argument("--advertise", default="127.0.0.1:53051")
+ parser.add_argument("--peer", action="append", default=[])
+ parser.add_argument("--model-id", required=True)
+ parser.add_argument("--model-revision", default="")
+ parser.add_argument("--tokenizer-revision", default="")
+ parser.add_argument("--cache-format-version", default="kakeya-prefill-v2-zlib")
+ parser.add_argument("--cache-model-id", default="",
+ help="Logical model id used for compatibility; defaults "
+ "to --model-id (which may be a host-specific path).")
+ parser.add_argument("--quantization", default="4bit-mlx")
+ parser.add_argument("--rope-hash", default="")
+ parser.add_argument("--layer-geometry-hash", required=True)
+ parser.add_argument("--kv-dtype", default="bfloat16")
+ parser.add_argument("--block-size-tokens", type=int, default=64)
+ parser.add_argument("--tenant-id", default="default")
+ parser.add_argument("--fleet-psk-file")
+ 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"],
+ default="zlib")
+ parser.add_argument("--replication-factor", type=int, default=1)
+ parser.add_argument("--max-concurrent-jobs", type=int, default=1)
+ parser.add_argument("--max-jobs", type=int, default=128)
+ parser.add_argument("--max-prompt-tokens", type=int, default=131072)
+ parser.add_argument("--job-ttl-s", type=float, default=600.0)
+ parser.add_argument("--prefill-tps", type=float, default=20.0)
+ parser.add_argument("--max-concurrent-rpcs", type=int, default=32)
+ parser.add_argument("--ttl-seconds", type=float, default=120.0)
+ parser.add_argument("--gossip-interval-s", type=float, default=10.0)
+ parser.add_argument("--gossip-timeout-s", type=float, default=3.0)
+ parser.add_argument("--network", default="lan")
+ parser.add_argument("--priority", type=int, default=50)
+ parser.add_argument("--rtt-ms", type=float, default=1.0)
+ parser.add_argument("--log-level", default="INFO")
+ args = parser.parse_args()
+ logging.basicConfig(level=getattr(logging, args.log_level.upper()))
+ asyncio.run(serve(args))
+
+
+if __name__ == "__main__":
+ main()
+
diff --git a/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts b/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts
index e35e7130..868f3d3b 100644
--- a/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts
+++ b/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts
@@ -47,6 +47,8 @@ export enum CapabilityRole {
* immutable prefill K/V blocks.
*/
PREFILL_CACHE = 5,
+ /** PREFILL_COMPUTE - The node has the compatible model loaded and accepts prefill-only jobs. */
+ PREFILL_COMPUTE = 6,
UNRECOGNIZED = -1,
}
@@ -70,6 +72,9 @@ export function capabilityRoleFromJSON(object: any): CapabilityRole {
case 5:
case "CAPABILITY_ROLE_PREFILL_CACHE":
return CapabilityRole.PREFILL_CACHE;
+ case 6:
+ case "CAPABILITY_ROLE_PREFILL_COMPUTE":
+ return CapabilityRole.PREFILL_COMPUTE;
case -1:
case "UNRECOGNIZED":
default:
@@ -91,12 +96,110 @@ export function capabilityRoleToJSON(object: CapabilityRole): string {
return "CAPABILITY_ROLE_TOOL";
case CapabilityRole.PREFILL_CACHE:
return "CAPABILITY_ROLE_PREFILL_CACHE";
+ case CapabilityRole.PREFILL_COMPUTE:
+ return "CAPABILITY_ROLE_PREFILL_COMPUTE";
case CapabilityRole.UNRECOGNIZED:
default:
return "UNRECOGNIZED";
}
}
+export enum CompressionCodec {
+ UNSPECIFIED = 0,
+ NONE = 1,
+ ZLIB = 2,
+ UNRECOGNIZED = -1,
+}
+
+export function compressionCodecFromJSON(object: any): CompressionCodec {
+ switch (object) {
+ case 0:
+ case "COMPRESSION_CODEC_UNSPECIFIED":
+ return CompressionCodec.UNSPECIFIED;
+ case 1:
+ case "COMPRESSION_CODEC_NONE":
+ return CompressionCodec.NONE;
+ case 2:
+ case "COMPRESSION_CODEC_ZLIB":
+ return CompressionCodec.ZLIB;
+ case -1:
+ case "UNRECOGNIZED":
+ default:
+ return CompressionCodec.UNRECOGNIZED;
+ }
+}
+
+export function compressionCodecToJSON(object: CompressionCodec): string {
+ switch (object) {
+ case CompressionCodec.UNSPECIFIED:
+ return "COMPRESSION_CODEC_UNSPECIFIED";
+ case CompressionCodec.NONE:
+ return "COMPRESSION_CODEC_NONE";
+ case CompressionCodec.ZLIB:
+ return "COMPRESSION_CODEC_ZLIB";
+ case CompressionCodec.UNRECOGNIZED:
+ default:
+ return "UNRECOGNIZED";
+ }
+}
+
+export enum PrefillJobStatus {
+ UNSPECIFIED = 0,
+ QUEUED = 1,
+ RUNNING = 2,
+ COMPLETED = 3,
+ FAILED = 4,
+ CANCELLED = 5,
+ UNRECOGNIZED = -1,
+}
+
+export function prefillJobStatusFromJSON(object: any): PrefillJobStatus {
+ switch (object) {
+ case 0:
+ case "PREFILL_JOB_STATUS_UNSPECIFIED":
+ return PrefillJobStatus.UNSPECIFIED;
+ case 1:
+ case "PREFILL_JOB_STATUS_QUEUED":
+ return PrefillJobStatus.QUEUED;
+ case 2:
+ case "PREFILL_JOB_STATUS_RUNNING":
+ return PrefillJobStatus.RUNNING;
+ case 3:
+ case "PREFILL_JOB_STATUS_COMPLETED":
+ return PrefillJobStatus.COMPLETED;
+ case 4:
+ case "PREFILL_JOB_STATUS_FAILED":
+ return PrefillJobStatus.FAILED;
+ case 5:
+ case "PREFILL_JOB_STATUS_CANCELLED":
+ return PrefillJobStatus.CANCELLED;
+ case -1:
+ case "UNRECOGNIZED":
+ default:
+ return PrefillJobStatus.UNRECOGNIZED;
+ }
+}
+
+export function prefillJobStatusToJSON(object: PrefillJobStatus): string {
+ switch (object) {
+ case PrefillJobStatus.UNSPECIFIED:
+ return "PREFILL_JOB_STATUS_UNSPECIFIED";
+ case PrefillJobStatus.QUEUED:
+ return "PREFILL_JOB_STATUS_QUEUED";
+ case PrefillJobStatus.RUNNING:
+ return "PREFILL_JOB_STATUS_RUNNING";
+ case PrefillJobStatus.COMPLETED:
+ return "PREFILL_JOB_STATUS_COMPLETED";
+ case PrefillJobStatus.FAILED:
+ return "PREFILL_JOB_STATUS_FAILED";
+ case PrefillJobStatus.CANCELLED:
+ return "PREFILL_JOB_STATUS_CANCELLED";
+ case PrefillJobStatus.UNRECOGNIZED:
+ default:
+ return "UNRECOGNIZED";
+ }
+}
+
/** ModelCapability is one (model, role) a node offers. */
export interface ModelCapability {
/**
@@ -172,6 +275,8 @@ export interface NodeCapability {
* LAN/Tailscale). grpc_address remains the compatibility default.
*/
endpoints: NodeEndpoint[];
+ /** Prefill-only compute offerings. Cache-only nodes leave this empty. */
+ prefillWorkers: PrefillWorkerCapability[];
}
export interface NodeEndpoint {
@@ -192,6 +297,14 @@ export interface CacheCompatibility {
layerGeometryHash: string;
kvDtype: string;
blockSizeTokens: number;
+ /**
+ * Tenant namespace is part of the compatibility fingerprint. It prevents a
+ * snapshot from being reused across isolation domains even with identical
+ * model/tokenizer geometry.
+ */
+ tenantNamespace: string;
+ sinkSize: number;
+ windowSize: number;
}
export interface CacheCapability {
@@ -205,6 +318,21 @@ export interface CacheCapability {
tokensServed: string;
/** Compact probabilistic summary. Empty means "query me directly". */
bloomFilter: Uint8Array;
+ defaultCompression: CompressionCodec;
+ replicationFactor: number;
+}
+
+export interface PrefillWorkerCapability {
+ compatibility?: CacheCompatibility | undefined;
+ workerAddress: string;
+ maxConcurrentJobs: number;
+ inflightJobs: number;
+ queuedJobs: number;
+ load: number;
+ tokensPerSecondPrefill: number;
+ ramBytesFree: string;
+ acceptsComputeJobs: boolean;
+ queuedTokens: string;
}
export interface ExchangeCapabilitiesRequest {
@@ -285,6 +413,55 @@ export interface PublishBlockResponse {
cacheEpoch: string;
}
+export interface SubmitPrefillJobRequest {
+ /** Caller-generated idempotency key. Repeating it returns the existing job. */
+ requestId: string;
+ tenantId: string;
+ compatibility?: CacheCompatibility | undefined;
+ tokenIds: number[];
+ /**
+ * Chained hashes computed by the primary (tenant-HMACed in authenticated
+ * mode). One hash per token block.
+ */
+ blockHashes: Uint8Array[];
+ deadlineMs: number;
+ preferredCompression: CompressionCodec;
+}
+
+export interface SubmitPrefillJobResponse {
+ jobId: string;
+ status: PrefillJobStatus;
+ workerNodeId: string;
+ queueEtaMs: number;
+}
+
+export interface GetPrefillJobStatusRequest {
+ jobId: string;
+ tenantId: string;
+}
+
+export interface GetPrefillJobStatusResponse {
+ jobId: string;
+ status: PrefillJobStatus;
+ tokensComputed: number;
+ leaseId: string;
+ blockHash: Uint8Array;
+ payloadSha256: Uint8Array;
+ transferBytes: string;
+ failureReason: string;
+ computeMs: number;
+ cacheAddress: string;
+}
+
+export interface CancelPrefillJobRequest {
+ jobId: string;
+ tenantId: string;
+}
+
+export interface CancelPrefillJobResponse {
+ cancelled: boolean;
+}
+
export interface ProposeBlockRequest {
/**
* The committed prefix (prompt + accepted tokens), raw token ids in
@@ -542,6 +719,7 @@ function createBaseNodeCapability(): NodeCapability {
ringAddress: "",
caches: [],
endpoints: [],
+ prefillWorkers: [],
};
}
@@ -580,6 +758,9 @@ export const NodeCapability: MessageFns = {
for (const v of message.endpoints) {
NodeEndpoint.encode(v!, writer.uint32(90).fork()).join();
}
+ for (const v of message.prefillWorkers) {
+ PrefillWorkerCapability.encode(v!, writer.uint32(98).fork()).join();
+ }
return writer;
},
@@ -678,6 +859,14 @@ export const NodeCapability: MessageFns = {
message.endpoints.push(NodeEndpoint.decode(reader, reader.uint32()));
continue;
}
+ case 12: {
+ if (tag !== 98) {
+ break;
+ }
+
+ message.prefillWorkers.push(PrefillWorkerCapability.decode(reader, reader.uint32()));
+ continue;
+ }
}
if ((tag & 7) === 4 || tag === 0) {
break;
@@ -734,6 +923,11 @@ export const NodeCapability: MessageFns = {
endpoints: globalThis.Array.isArray(object?.endpoints)
? object.endpoints.map((e: any) => NodeEndpoint.fromJSON(e))
: [],
+ prefillWorkers: globalThis.Array.isArray(object?.prefillWorkers)
+ ? object.prefillWorkers.map((e: any) => PrefillWorkerCapability.fromJSON(e))
+ : globalThis.Array.isArray(object?.prefill_workers)
+ ? object.prefill_workers.map((e: any) => PrefillWorkerCapability.fromJSON(e))
+ : [],
};
},
@@ -772,6 +966,9 @@ export const NodeCapability: MessageFns = {
if (message.endpoints?.length) {
obj.endpoints = message.endpoints.map((e) => NodeEndpoint.toJSON(e));
}
+ if (message.prefillWorkers?.length) {
+ obj.prefillWorkers = message.prefillWorkers.map((e) => PrefillWorkerCapability.toJSON(e));
+ }
return obj;
},
@@ -791,6 +988,7 @@ export const NodeCapability: MessageFns = {
message.ringAddress = object.ringAddress ?? "";
message.caches = object.caches?.map((e) => CacheCapability.fromPartial(e)) || [];
message.endpoints = object.endpoints?.map((e) => NodeEndpoint.fromPartial(e)) || [];
+ message.prefillWorkers = object.prefillWorkers?.map((e) => PrefillWorkerCapability.fromPartial(e)) || [];
return message;
},
};
@@ -918,6 +1116,9 @@ function createBaseCacheCompatibility(): CacheCompatibility {
layerGeometryHash: "",
kvDtype: "",
blockSizeTokens: 0,
+ tenantNamespace: "",
+ sinkSize: 0,
+ windowSize: 0,
};
}
@@ -950,6 +1151,15 @@ export const CacheCompatibility: MessageFns = {
if (message.blockSizeTokens !== 0) {
writer.uint32(72).uint32(message.blockSizeTokens);
}
+ if (message.tenantNamespace !== "") {
+ writer.uint32(82).string(message.tenantNamespace);
+ }
+ if (message.sinkSize !== 0) {
+ writer.uint32(88).uint32(message.sinkSize);
+ }
+ if (message.windowSize !== 0) {
+ writer.uint32(96).uint32(message.windowSize);
+ }
return writer;
},
@@ -1032,6 +1242,30 @@ export const CacheCompatibility: MessageFns = {
message.blockSizeTokens = reader.uint32();
continue;
}
+ case 10: {
+ if (tag !== 82) {
+ break;
+ }
+
+ message.tenantNamespace = reader.string();
+ continue;
+ }
+ case 11: {
+ if (tag !== 88) {
+ break;
+ }
+
+ message.sinkSize = reader.uint32();
+ continue;
+ }
+ case 12: {
+ if (tag !== 96) {
+ break;
+ }
+
+ message.windowSize = reader.uint32();
+ continue;
+ }
}
if ((tag & 7) === 4 || tag === 0) {
break;
@@ -1084,6 +1318,21 @@ export const CacheCompatibility: MessageFns = {
: isSet(object.block_size_tokens)
? globalThis.Number(object.block_size_tokens)
: 0,
+ tenantNamespace: isSet(object.tenantNamespace)
+ ? globalThis.String(object.tenantNamespace)
+ : isSet(object.tenant_namespace)
+ ? globalThis.String(object.tenant_namespace)
+ : "",
+ sinkSize: isSet(object.sinkSize)
+ ? globalThis.Number(object.sinkSize)
+ : isSet(object.sink_size)
+ ? globalThis.Number(object.sink_size)
+ : 0,
+ windowSize: isSet(object.windowSize)
+ ? globalThis.Number(object.windowSize)
+ : isSet(object.window_size)
+ ? globalThis.Number(object.window_size)
+ : 0,
};
},
@@ -1116,6 +1365,15 @@ export const CacheCompatibility: MessageFns = {
if (message.blockSizeTokens !== 0) {
obj.blockSizeTokens = Math.round(message.blockSizeTokens);
}
+ if (message.tenantNamespace !== "") {
+ obj.tenantNamespace = message.tenantNamespace;
+ }
+ if (message.sinkSize !== 0) {
+ obj.sinkSize = Math.round(message.sinkSize);
+ }
+ if (message.windowSize !== 0) {
+ obj.windowSize = Math.round(message.windowSize);
+ }
return obj;
},
@@ -1133,6 +1391,9 @@ export const CacheCompatibility: MessageFns = {
message.layerGeometryHash = object.layerGeometryHash ?? "";
message.kvDtype = object.kvDtype ?? "";
message.blockSizeTokens = object.blockSizeTokens ?? 0;
+ message.tenantNamespace = object.tenantNamespace ?? "";
+ message.sinkSize = object.sinkSize ?? 0;
+ message.windowSize = object.windowSize ?? 0;
return message;
},
};
@@ -1148,6 +1409,8 @@ function createBaseCacheCapability(): CacheCapability {
load: 0,
tokensServed: "0",
bloomFilter: new Uint8Array(0),
+ defaultCompression: 0,
+ replicationFactor: 0,
};
}
@@ -1180,6 +1443,12 @@ export const CacheCapability: MessageFns = {
if (message.bloomFilter.length !== 0) {
writer.uint32(74).bytes(message.bloomFilter);
}
+ if (message.defaultCompression !== 0) {
+ writer.uint32(80).int32(message.defaultCompression);
+ }
+ if (message.replicationFactor !== 0) {
+ writer.uint32(88).uint32(message.replicationFactor);
+ }
return writer;
},
@@ -1262,6 +1531,22 @@ export const CacheCapability: MessageFns = {
message.bloomFilter = reader.bytes();
continue;
}
+ case 10: {
+ if (tag !== 80) {
+ break;
+ }
+
+ message.defaultCompression = reader.int32() as any;
+ continue;
+ }
+ case 11: {
+ if (tag !== 88) {
+ break;
+ }
+
+ message.replicationFactor = reader.uint32();
+ continue;
+ }
}
if ((tag & 7) === 4 || tag === 0) {
break;
@@ -1310,6 +1595,16 @@ export const CacheCapability: MessageFns = {
: isSet(object.bloom_filter)
? bytesFromBase64(object.bloom_filter)
: new Uint8Array(0),
+ defaultCompression: isSet(object.defaultCompression)
+ ? compressionCodecFromJSON(object.defaultCompression)
+ : isSet(object.default_compression)
+ ? compressionCodecFromJSON(object.default_compression)
+ : 0,
+ replicationFactor: isSet(object.replicationFactor)
+ ? globalThis.Number(object.replicationFactor)
+ : isSet(object.replication_factor)
+ ? globalThis.Number(object.replication_factor)
+ : 0,
};
},
@@ -1342,6 +1637,12 @@ export const CacheCapability: MessageFns = {
if (message.bloomFilter.length !== 0) {
obj.bloomFilter = base64FromBytes(message.bloomFilter);
}
+ if (message.defaultCompression !== 0) {
+ obj.defaultCompression = compressionCodecToJSON(message.defaultCompression);
+ }
+ if (message.replicationFactor !== 0) {
+ obj.replicationFactor = Math.round(message.replicationFactor);
+ }
return obj;
},
@@ -1361,26 +1662,66 @@ export const CacheCapability: MessageFns = {
message.load = object.load ?? 0;
message.tokensServed = object.tokensServed ?? "0";
message.bloomFilter = object.bloomFilter ?? new Uint8Array(0);
+ message.defaultCompression = object.defaultCompression ?? 0;
+ message.replicationFactor = object.replicationFactor ?? 0;
return message;
},
};
-function createBaseExchangeCapabilitiesRequest(): ExchangeCapabilitiesRequest {
- return { knownNodes: [] };
+function createBasePrefillWorkerCapability(): PrefillWorkerCapability {
+ return {
+ compatibility: undefined,
+ workerAddress: "",
+ maxConcurrentJobs: 0,
+ inflightJobs: 0,
+ queuedJobs: 0,
+ load: 0,
+ tokensPerSecondPrefill: 0,
+ ramBytesFree: "0",
+ acceptsComputeJobs: false,
+ queuedTokens: "0",
+ };
}
-export const ExchangeCapabilitiesRequest: MessageFns = {
- encode(message: ExchangeCapabilitiesRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
- for (const v of message.knownNodes) {
- NodeCapability.encode(v!, writer.uint32(10).fork()).join();
+export const PrefillWorkerCapability: MessageFns = {
+ encode(message: PrefillWorkerCapability, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
+ if (message.compatibility !== undefined) {
+ CacheCompatibility.encode(message.compatibility, writer.uint32(10).fork()).join();
+ }
+ if (message.workerAddress !== "") {
+ writer.uint32(18).string(message.workerAddress);
+ }
+ if (message.maxConcurrentJobs !== 0) {
+ writer.uint32(24).uint32(message.maxConcurrentJobs);
+ }
+ if (message.inflightJobs !== 0) {
+ writer.uint32(32).uint32(message.inflightJobs);
+ }
+ if (message.queuedJobs !== 0) {
+ writer.uint32(40).uint32(message.queuedJobs);
+ }
+ if (message.load !== 0) {
+ writer.uint32(49).double(message.load);
+ }
+ if (message.tokensPerSecondPrefill !== 0) {
+ writer.uint32(57).double(message.tokensPerSecondPrefill);
+ }
+ if (message.ramBytesFree !== "0") {
+ writer.uint32(64).uint64(message.ramBytesFree);
+ }
+ if (message.acceptsComputeJobs !== false) {
+ writer.uint32(72).bool(message.acceptsComputeJobs);
+ }
+ if (message.queuedTokens !== "0") {
+ writer.uint32(80).uint64(message.queuedTokens);
}
return writer;
},
- decode(input: BinaryReader | Uint8Array, length?: number): ExchangeCapabilitiesRequest {
+ decode(input: BinaryReader | Uint8Array, length?: number): PrefillWorkerCapability {
const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
const end = length === undefined ? reader.len : reader.pos + length;
- const message = createBaseExchangeCapabilitiesRequest();
+ const message = createBasePrefillWorkerCapability();
while (reader.pos < end) {
const tag = reader.uint32();
switch (tag >>> 3) {
@@ -1389,7 +1730,79 @@ export const ExchangeCapabilitiesRequest: MessageFns NodeCapability.fromJSON(e))
- : globalThis.Array.isArray(object?.known_nodes)
- ? object.known_nodes.map((e: any) => NodeCapability.fromJSON(e))
- : [],
+ compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined,
+ workerAddress: isSet(object.workerAddress)
+ ? globalThis.String(object.workerAddress)
+ : isSet(object.worker_address)
+ ? globalThis.String(object.worker_address)
+ : "",
+ maxConcurrentJobs: isSet(object.maxConcurrentJobs)
+ ? globalThis.Number(object.maxConcurrentJobs)
+ : isSet(object.max_concurrent_jobs)
+ ? globalThis.Number(object.max_concurrent_jobs)
+ : 0,
+ inflightJobs: isSet(object.inflightJobs)
+ ? globalThis.Number(object.inflightJobs)
+ : isSet(object.inflight_jobs)
+ ? globalThis.Number(object.inflight_jobs)
+ : 0,
+ queuedJobs: isSet(object.queuedJobs)
+ ? globalThis.Number(object.queuedJobs)
+ : isSet(object.queued_jobs)
+ ? globalThis.Number(object.queued_jobs)
+ : 0,
+ load: isSet(object.load) ? globalThis.Number(object.load) : 0,
+ tokensPerSecondPrefill: isSet(object.tokensPerSecondPrefill)
+ ? globalThis.Number(object.tokensPerSecondPrefill)
+ : isSet(object.tokens_per_second_prefill)
+ ? globalThis.Number(object.tokens_per_second_prefill)
+ : 0,
+ ramBytesFree: isSet(object.ramBytesFree)
+ ? globalThis.String(object.ramBytesFree)
+ : isSet(object.ram_bytes_free)
+ ? globalThis.String(object.ram_bytes_free)
+ : "0",
+ acceptsComputeJobs: isSet(object.acceptsComputeJobs)
+ ? globalThis.Boolean(object.acceptsComputeJobs)
+ : isSet(object.accepts_compute_jobs)
+ ? globalThis.Boolean(object.accepts_compute_jobs)
+ : false,
+ queuedTokens: isSet(object.queuedTokens)
+ ? globalThis.String(object.queuedTokens)
+ : isSet(object.queued_tokens)
+ ? globalThis.String(object.queued_tokens)
+ : "0",
};
},
- toJSON(message: ExchangeCapabilitiesRequest): unknown {
+ toJSON(message: PrefillWorkerCapability): unknown {
const obj: any = {};
- if (message.knownNodes?.length) {
- obj.knownNodes = message.knownNodes.map((e) => NodeCapability.toJSON(e));
+ if (message.compatibility !== undefined) {
+ obj.compatibility = CacheCompatibility.toJSON(message.compatibility);
+ }
+ if (message.workerAddress !== "") {
+ obj.workerAddress = message.workerAddress;
+ }
+ if (message.maxConcurrentJobs !== 0) {
+ obj.maxConcurrentJobs = Math.round(message.maxConcurrentJobs);
+ }
+ if (message.inflightJobs !== 0) {
+ obj.inflightJobs = Math.round(message.inflightJobs);
+ }
+ if (message.queuedJobs !== 0) {
+ obj.queuedJobs = Math.round(message.queuedJobs);
+ }
+ if (message.load !== 0) {
+ obj.load = message.load;
+ }
+ if (message.tokensPerSecondPrefill !== 0) {
+ obj.tokensPerSecondPrefill = message.tokensPerSecondPrefill;
+ }
+ if (message.ramBytesFree !== "0") {
+ obj.ramBytesFree = message.ramBytesFree;
+ }
+ if (message.acceptsComputeJobs !== false) {
+ obj.acceptsComputeJobs = message.acceptsComputeJobs;
+ }
+ if (message.queuedTokens !== "0") {
+ obj.queuedTokens = message.queuedTokens;
}
return obj;
},
- create, I>>(base?: I): ExchangeCapabilitiesRequest {
- return ExchangeCapabilitiesRequest.fromPartial(base ?? ({} as any));
+ create, I>>(base?: I): PrefillWorkerCapability {
+ return PrefillWorkerCapability.fromPartial(base ?? ({} as any));
},
- fromPartial, I>>(object: I): ExchangeCapabilitiesRequest {
- const message = createBaseExchangeCapabilitiesRequest();
- message.knownNodes = object.knownNodes?.map((e) => NodeCapability.fromPartial(e)) || [];
+ fromPartial, I>>(object: I): PrefillWorkerCapability {
+ const message = createBasePrefillWorkerCapability();
+ message.compatibility = (object.compatibility !== undefined && object.compatibility !== null)
+ ? CacheCompatibility.fromPartial(object.compatibility)
+ : undefined;
+ message.workerAddress = object.workerAddress ?? "";
+ message.maxConcurrentJobs = object.maxConcurrentJobs ?? 0;
+ message.inflightJobs = object.inflightJobs ?? 0;
+ message.queuedJobs = object.queuedJobs ?? 0;
+ message.load = object.load ?? 0;
+ message.tokensPerSecondPrefill = object.tokensPerSecondPrefill ?? 0;
+ message.ramBytesFree = object.ramBytesFree ?? "0";
+ message.acceptsComputeJobs = object.acceptsComputeJobs ?? false;
+ message.queuedTokens = object.queuedTokens ?? "0";
return message;
},
};
-function createBaseExchangeCapabilitiesResponse(): ExchangeCapabilitiesResponse {
+function createBaseExchangeCapabilitiesRequest(): ExchangeCapabilitiesRequest {
return { knownNodes: [] };
}
-export const ExchangeCapabilitiesResponse: MessageFns = {
+export const ExchangeCapabilitiesRequest: MessageFns = {
+ encode(message: ExchangeCapabilitiesRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
+ for (const v of message.knownNodes) {
+ NodeCapability.encode(v!, writer.uint32(10).fork()).join();
+ }
+ return writer;
+ },
+
+ decode(input: BinaryReader | Uint8Array, length?: number): ExchangeCapabilitiesRequest {
+ const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
+ const end = length === undefined ? reader.len : reader.pos + length;
+ const message = createBaseExchangeCapabilitiesRequest();
+ while (reader.pos < end) {
+ const tag = reader.uint32();
+ switch (tag >>> 3) {
+ case 1: {
+ if (tag !== 10) {
+ break;
+ }
+
+ message.knownNodes.push(NodeCapability.decode(reader, reader.uint32()));
+ continue;
+ }
+ }
+ if ((tag & 7) === 4 || tag === 0) {
+ break;
+ }
+ reader.skip(tag & 7);
+ }
+ return message;
+ },
+
+ fromJSON(object: any): ExchangeCapabilitiesRequest {
+ return {
+ knownNodes: globalThis.Array.isArray(object?.knownNodes)
+ ? object.knownNodes.map((e: any) => NodeCapability.fromJSON(e))
+ : globalThis.Array.isArray(object?.known_nodes)
+ ? object.known_nodes.map((e: any) => NodeCapability.fromJSON(e))
+ : [],
+ };
+ },
+
+ toJSON(message: ExchangeCapabilitiesRequest): unknown {
+ const obj: any = {};
+ if (message.knownNodes?.length) {
+ obj.knownNodes = message.knownNodes.map((e) => NodeCapability.toJSON(e));
+ }
+ return obj;
+ },
+
+ create, I>>(base?: I): ExchangeCapabilitiesRequest {
+ return ExchangeCapabilitiesRequest.fromPartial(base ?? ({} as any));
+ },
+ fromPartial, I>>(object: I): ExchangeCapabilitiesRequest {
+ const message = createBaseExchangeCapabilitiesRequest();
+ message.knownNodes = object.knownNodes?.map((e) => NodeCapability.fromPartial(e)) || [];
+ return message;
+ },
+};
+
+function createBaseExchangeCapabilitiesResponse(): ExchangeCapabilitiesResponse {
+ return { knownNodes: [] };
+}
+
+export const ExchangeCapabilitiesResponse: MessageFns = {
encode(message: ExchangeCapabilitiesResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
for (const v of message.knownNodes) {
NodeCapability.encode(v!, writer.uint32(10).fork()).join();
@@ -2339,25 +2891,743 @@ export const PublishBlockRequest: MessageFns = {
if (message.totalChunks !== 0) {
writer.uint32(40).uint32(message.totalChunks);
}
- if (message.data.length !== 0) {
- writer.uint32(50).bytes(message.data);
+ if (message.data.length !== 0) {
+ writer.uint32(50).bytes(message.data);
+ }
+ if (message.blockSha256.length !== 0) {
+ writer.uint32(58).bytes(message.blockSha256);
+ }
+ if (message.cacheEpoch !== "0") {
+ writer.uint32(64).uint64(message.cacheEpoch);
+ }
+ if (message.compatibility !== undefined) {
+ CacheCompatibility.encode(message.compatibility, writer.uint32(74).fork()).join();
+ }
+ return writer;
+ },
+
+ decode(input: BinaryReader | Uint8Array, length?: number): PublishBlockRequest {
+ const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
+ const end = length === undefined ? reader.len : reader.pos + length;
+ const message = createBasePublishBlockRequest();
+ while (reader.pos < end) {
+ const tag = reader.uint32();
+ switch (tag >>> 3) {
+ case 1: {
+ if (tag !== 10) {
+ break;
+ }
+
+ message.blockHash = reader.bytes();
+ continue;
+ }
+ case 2: {
+ if (tag !== 16) {
+ break;
+ }
+
+ message.blockIndex = reader.uint32();
+ continue;
+ }
+ case 3: {
+ if (tag !== 24) {
+ break;
+ }
+
+ message.tokenCount = reader.uint32();
+ continue;
+ }
+ case 4: {
+ if (tag !== 32) {
+ break;
+ }
+
+ message.chunkIndex = reader.uint32();
+ continue;
+ }
+ case 5: {
+ if (tag !== 40) {
+ break;
+ }
+
+ message.totalChunks = reader.uint32();
+ continue;
+ }
+ case 6: {
+ if (tag !== 50) {
+ break;
+ }
+
+ message.data = reader.bytes();
+ continue;
+ }
+ case 7: {
+ if (tag !== 58) {
+ break;
+ }
+
+ message.blockSha256 = reader.bytes();
+ continue;
+ }
+ case 8: {
+ if (tag !== 64) {
+ break;
+ }
+
+ message.cacheEpoch = reader.uint64().toString();
+ continue;
+ }
+ case 9: {
+ if (tag !== 74) {
+ break;
+ }
+
+ message.compatibility = CacheCompatibility.decode(reader, reader.uint32());
+ continue;
+ }
+ }
+ if ((tag & 7) === 4 || tag === 0) {
+ break;
+ }
+ reader.skip(tag & 7);
+ }
+ return message;
+ },
+
+ fromJSON(object: any): PublishBlockRequest {
+ return {
+ blockHash: isSet(object.blockHash)
+ ? bytesFromBase64(object.blockHash)
+ : isSet(object.block_hash)
+ ? bytesFromBase64(object.block_hash)
+ : new Uint8Array(0),
+ blockIndex: isSet(object.blockIndex)
+ ? globalThis.Number(object.blockIndex)
+ : isSet(object.block_index)
+ ? globalThis.Number(object.block_index)
+ : 0,
+ tokenCount: isSet(object.tokenCount)
+ ? globalThis.Number(object.tokenCount)
+ : isSet(object.token_count)
+ ? globalThis.Number(object.token_count)
+ : 0,
+ chunkIndex: isSet(object.chunkIndex)
+ ? globalThis.Number(object.chunkIndex)
+ : isSet(object.chunk_index)
+ ? globalThis.Number(object.chunk_index)
+ : 0,
+ totalChunks: isSet(object.totalChunks)
+ ? globalThis.Number(object.totalChunks)
+ : isSet(object.total_chunks)
+ ? globalThis.Number(object.total_chunks)
+ : 0,
+ data: isSet(object.data) ? bytesFromBase64(object.data) : new Uint8Array(0),
+ blockSha256: isSet(object.blockSha256)
+ ? bytesFromBase64(object.blockSha256)
+ : isSet(object.block_sha256)
+ ? bytesFromBase64(object.block_sha256)
+ : new Uint8Array(0),
+ cacheEpoch: isSet(object.cacheEpoch)
+ ? globalThis.String(object.cacheEpoch)
+ : isSet(object.cache_epoch)
+ ? globalThis.String(object.cache_epoch)
+ : "0",
+ compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined,
+ };
+ },
+
+ toJSON(message: PublishBlockRequest): unknown {
+ const obj: any = {};
+ if (message.blockHash.length !== 0) {
+ obj.blockHash = base64FromBytes(message.blockHash);
+ }
+ if (message.blockIndex !== 0) {
+ obj.blockIndex = Math.round(message.blockIndex);
+ }
+ if (message.tokenCount !== 0) {
+ obj.tokenCount = Math.round(message.tokenCount);
+ }
+ if (message.chunkIndex !== 0) {
+ obj.chunkIndex = Math.round(message.chunkIndex);
+ }
+ if (message.totalChunks !== 0) {
+ obj.totalChunks = Math.round(message.totalChunks);
+ }
+ if (message.data.length !== 0) {
+ obj.data = base64FromBytes(message.data);
+ }
+ if (message.blockSha256.length !== 0) {
+ obj.blockSha256 = base64FromBytes(message.blockSha256);
+ }
+ if (message.cacheEpoch !== "0") {
+ obj.cacheEpoch = message.cacheEpoch;
+ }
+ if (message.compatibility !== undefined) {
+ obj.compatibility = CacheCompatibility.toJSON(message.compatibility);
+ }
+ return obj;
+ },
+
+ create, I>>(base?: I): PublishBlockRequest {
+ return PublishBlockRequest.fromPartial(base ?? ({} as any));
+ },
+ fromPartial, I>>(object: I): PublishBlockRequest {
+ const message = createBasePublishBlockRequest();
+ message.blockHash = object.blockHash ?? new Uint8Array(0);
+ message.blockIndex = object.blockIndex ?? 0;
+ message.tokenCount = object.tokenCount ?? 0;
+ message.chunkIndex = object.chunkIndex ?? 0;
+ message.totalChunks = object.totalChunks ?? 0;
+ message.data = object.data ?? new Uint8Array(0);
+ message.blockSha256 = object.blockSha256 ?? new Uint8Array(0);
+ message.cacheEpoch = object.cacheEpoch ?? "0";
+ message.compatibility = (object.compatibility !== undefined && object.compatibility !== null)
+ ? CacheCompatibility.fromPartial(object.compatibility)
+ : undefined;
+ return message;
+ },
+};
+
+function createBasePublishBlockResponse(): PublishBlockResponse {
+ return { stored: false, cacheEpoch: "0" };
+}
+
+export const PublishBlockResponse: MessageFns = {
+ encode(message: PublishBlockResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
+ if (message.stored !== false) {
+ writer.uint32(8).bool(message.stored);
+ }
+ if (message.cacheEpoch !== "0") {
+ writer.uint32(16).uint64(message.cacheEpoch);
+ }
+ return writer;
+ },
+
+ decode(input: BinaryReader | Uint8Array, length?: number): PublishBlockResponse {
+ const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
+ const end = length === undefined ? reader.len : reader.pos + length;
+ const message = createBasePublishBlockResponse();
+ while (reader.pos < end) {
+ const tag = reader.uint32();
+ switch (tag >>> 3) {
+ case 1: {
+ if (tag !== 8) {
+ break;
+ }
+
+ message.stored = reader.bool();
+ continue;
+ }
+ case 2: {
+ if (tag !== 16) {
+ break;
+ }
+
+ message.cacheEpoch = reader.uint64().toString();
+ continue;
+ }
+ }
+ if ((tag & 7) === 4 || tag === 0) {
+ break;
+ }
+ reader.skip(tag & 7);
+ }
+ return message;
+ },
+
+ fromJSON(object: any): PublishBlockResponse {
+ return {
+ stored: isSet(object.stored) ? globalThis.Boolean(object.stored) : false,
+ cacheEpoch: isSet(object.cacheEpoch)
+ ? globalThis.String(object.cacheEpoch)
+ : isSet(object.cache_epoch)
+ ? globalThis.String(object.cache_epoch)
+ : "0",
+ };
+ },
+
+ toJSON(message: PublishBlockResponse): unknown {
+ const obj: any = {};
+ if (message.stored !== false) {
+ obj.stored = message.stored;
+ }
+ if (message.cacheEpoch !== "0") {
+ obj.cacheEpoch = message.cacheEpoch;
+ }
+ return obj;
+ },
+
+ create, I>>(base?: I): PublishBlockResponse {
+ return PublishBlockResponse.fromPartial(base ?? ({} as any));
+ },
+ fromPartial, I>>(object: I): PublishBlockResponse {
+ const message = createBasePublishBlockResponse();
+ message.stored = object.stored ?? false;
+ message.cacheEpoch = object.cacheEpoch ?? "0";
+ return message;
+ },
+};
+
+function createBaseSubmitPrefillJobRequest(): SubmitPrefillJobRequest {
+ return {
+ requestId: "",
+ tenantId: "",
+ compatibility: undefined,
+ tokenIds: [],
+ blockHashes: [],
+ deadlineMs: 0,
+ preferredCompression: 0,
+ };
+}
+
+export const SubmitPrefillJobRequest: MessageFns = {
+ encode(message: SubmitPrefillJobRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
+ if (message.requestId !== "") {
+ writer.uint32(10).string(message.requestId);
+ }
+ if (message.tenantId !== "") {
+ writer.uint32(18).string(message.tenantId);
+ }
+ if (message.compatibility !== undefined) {
+ CacheCompatibility.encode(message.compatibility, writer.uint32(26).fork()).join();
+ }
+ writer.uint32(34).fork();
+ for (const v of message.tokenIds) {
+ writer.uint32(v);
+ }
+ writer.join();
+ for (const v of message.blockHashes) {
+ writer.uint32(42).bytes(v!);
+ }
+ if (message.deadlineMs !== 0) {
+ writer.uint32(48).uint32(message.deadlineMs);
+ }
+ if (message.preferredCompression !== 0) {
+ writer.uint32(56).int32(message.preferredCompression);
+ }
+ return writer;
+ },
+
+ decode(input: BinaryReader | Uint8Array, length?: number): SubmitPrefillJobRequest {
+ const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
+ const end = length === undefined ? reader.len : reader.pos + length;
+ const message = createBaseSubmitPrefillJobRequest();
+ while (reader.pos < end) {
+ const tag = reader.uint32();
+ switch (tag >>> 3) {
+ case 1: {
+ if (tag !== 10) {
+ break;
+ }
+
+ message.requestId = reader.string();
+ continue;
+ }
+ case 2: {
+ if (tag !== 18) {
+ break;
+ }
+
+ message.tenantId = reader.string();
+ continue;
+ }
+ case 3: {
+ if (tag !== 26) {
+ break;
+ }
+
+ message.compatibility = CacheCompatibility.decode(reader, reader.uint32());
+ continue;
+ }
+ case 4: {
+ if (tag === 32) {
+ message.tokenIds.push(reader.uint32());
+
+ continue;
+ }
+
+ if (tag === 34) {
+ const end2 = reader.uint32() + reader.pos;
+ while (reader.pos < end2) {
+ message.tokenIds.push(reader.uint32());
+ }
+
+ continue;
+ }
+
+ break;
+ }
+ case 5: {
+ if (tag !== 42) {
+ break;
+ }
+
+ message.blockHashes.push(reader.bytes());
+ continue;
+ }
+ case 6: {
+ if (tag !== 48) {
+ break;
+ }
+
+ message.deadlineMs = reader.uint32();
+ continue;
+ }
+ case 7: {
+ if (tag !== 56) {
+ break;
+ }
+
+ message.preferredCompression = reader.int32() as any;
+ continue;
+ }
+ }
+ if ((tag & 7) === 4 || tag === 0) {
+ break;
+ }
+ reader.skip(tag & 7);
+ }
+ return message;
+ },
+
+ fromJSON(object: any): SubmitPrefillJobRequest {
+ return {
+ requestId: isSet(object.requestId)
+ ? globalThis.String(object.requestId)
+ : isSet(object.request_id)
+ ? globalThis.String(object.request_id)
+ : "",
+ tenantId: isSet(object.tenantId)
+ ? globalThis.String(object.tenantId)
+ : isSet(object.tenant_id)
+ ? globalThis.String(object.tenant_id)
+ : "",
+ compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined,
+ tokenIds: globalThis.Array.isArray(object?.tokenIds)
+ ? object.tokenIds.map((e: any) => globalThis.Number(e))
+ : globalThis.Array.isArray(object?.token_ids)
+ ? object.token_ids.map((e: any) => globalThis.Number(e))
+ : [],
+ blockHashes: globalThis.Array.isArray(object?.blockHashes)
+ ? object.blockHashes.map((e: any) => bytesFromBase64(e))
+ : globalThis.Array.isArray(object?.block_hashes)
+ ? object.block_hashes.map((e: any) => bytesFromBase64(e))
+ : [],
+ deadlineMs: isSet(object.deadlineMs)
+ ? globalThis.Number(object.deadlineMs)
+ : isSet(object.deadline_ms)
+ ? globalThis.Number(object.deadline_ms)
+ : 0,
+ preferredCompression: isSet(object.preferredCompression)
+ ? compressionCodecFromJSON(object.preferredCompression)
+ : isSet(object.preferred_compression)
+ ? compressionCodecFromJSON(object.preferred_compression)
+ : 0,
+ };
+ },
+
+ toJSON(message: SubmitPrefillJobRequest): unknown {
+ const obj: any = {};
+ if (message.requestId !== "") {
+ obj.requestId = message.requestId;
+ }
+ if (message.tenantId !== "") {
+ obj.tenantId = message.tenantId;
+ }
+ if (message.compatibility !== undefined) {
+ obj.compatibility = CacheCompatibility.toJSON(message.compatibility);
+ }
+ if (message.tokenIds?.length) {
+ obj.tokenIds = message.tokenIds.map((e) => Math.round(e));
+ }
+ if (message.blockHashes?.length) {
+ obj.blockHashes = message.blockHashes.map((e) => base64FromBytes(e));
+ }
+ if (message.deadlineMs !== 0) {
+ obj.deadlineMs = Math.round(message.deadlineMs);
+ }
+ if (message.preferredCompression !== 0) {
+ obj.preferredCompression = compressionCodecToJSON(message.preferredCompression);
+ }
+ return obj;
+ },
+
+ create, I>>(base?: I): SubmitPrefillJobRequest {
+ return SubmitPrefillJobRequest.fromPartial(base ?? ({} as any));
+ },
+ fromPartial, I>>(object: I): SubmitPrefillJobRequest {
+ const message = createBaseSubmitPrefillJobRequest();
+ message.requestId = object.requestId ?? "";
+ message.tenantId = object.tenantId ?? "";
+ message.compatibility = (object.compatibility !== undefined && object.compatibility !== null)
+ ? CacheCompatibility.fromPartial(object.compatibility)
+ : undefined;
+ message.tokenIds = object.tokenIds?.map((e) => e) || [];
+ message.blockHashes = object.blockHashes?.map((e) => e) || [];
+ message.deadlineMs = object.deadlineMs ?? 0;
+ message.preferredCompression = object.preferredCompression ?? 0;
+ return message;
+ },
+};
+
+function createBaseSubmitPrefillJobResponse(): SubmitPrefillJobResponse {
+ return { jobId: "", status: 0, workerNodeId: "", queueEtaMs: 0 };
+}
+
+export const SubmitPrefillJobResponse: MessageFns = {
+ encode(message: SubmitPrefillJobResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
+ if (message.jobId !== "") {
+ writer.uint32(10).string(message.jobId);
+ }
+ if (message.status !== 0) {
+ writer.uint32(16).int32(message.status);
+ }
+ if (message.workerNodeId !== "") {
+ writer.uint32(26).string(message.workerNodeId);
+ }
+ if (message.queueEtaMs !== 0) {
+ writer.uint32(33).double(message.queueEtaMs);
+ }
+ return writer;
+ },
+
+ decode(input: BinaryReader | Uint8Array, length?: number): SubmitPrefillJobResponse {
+ const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
+ const end = length === undefined ? reader.len : reader.pos + length;
+ const message = createBaseSubmitPrefillJobResponse();
+ while (reader.pos < end) {
+ const tag = reader.uint32();
+ switch (tag >>> 3) {
+ case 1: {
+ if (tag !== 10) {
+ break;
+ }
+
+ message.jobId = reader.string();
+ continue;
+ }
+ case 2: {
+ if (tag !== 16) {
+ break;
+ }
+
+ message.status = reader.int32() as any;
+ continue;
+ }
+ case 3: {
+ if (tag !== 26) {
+ break;
+ }
+
+ message.workerNodeId = reader.string();
+ continue;
+ }
+ case 4: {
+ if (tag !== 33) {
+ break;
+ }
+
+ message.queueEtaMs = reader.double();
+ continue;
+ }
+ }
+ if ((tag & 7) === 4 || tag === 0) {
+ break;
+ }
+ reader.skip(tag & 7);
+ }
+ return message;
+ },
+
+ fromJSON(object: any): SubmitPrefillJobResponse {
+ return {
+ jobId: isSet(object.jobId)
+ ? globalThis.String(object.jobId)
+ : isSet(object.job_id)
+ ? globalThis.String(object.job_id)
+ : "",
+ status: isSet(object.status) ? prefillJobStatusFromJSON(object.status) : 0,
+ workerNodeId: isSet(object.workerNodeId)
+ ? globalThis.String(object.workerNodeId)
+ : isSet(object.worker_node_id)
+ ? globalThis.String(object.worker_node_id)
+ : "",
+ queueEtaMs: isSet(object.queueEtaMs)
+ ? globalThis.Number(object.queueEtaMs)
+ : isSet(object.queue_eta_ms)
+ ? globalThis.Number(object.queue_eta_ms)
+ : 0,
+ };
+ },
+
+ toJSON(message: SubmitPrefillJobResponse): unknown {
+ const obj: any = {};
+ if (message.jobId !== "") {
+ obj.jobId = message.jobId;
+ }
+ if (message.status !== 0) {
+ obj.status = prefillJobStatusToJSON(message.status);
+ }
+ if (message.workerNodeId !== "") {
+ obj.workerNodeId = message.workerNodeId;
+ }
+ if (message.queueEtaMs !== 0) {
+ obj.queueEtaMs = message.queueEtaMs;
+ }
+ return obj;
+ },
+
+ create, I>>(base?: I): SubmitPrefillJobResponse {
+ return SubmitPrefillJobResponse.fromPartial(base ?? ({} as any));
+ },
+ fromPartial, I>>(object: I): SubmitPrefillJobResponse {
+ const message = createBaseSubmitPrefillJobResponse();
+ message.jobId = object.jobId ?? "";
+ message.status = object.status ?? 0;
+ message.workerNodeId = object.workerNodeId ?? "";
+ message.queueEtaMs = object.queueEtaMs ?? 0;
+ return message;
+ },
+};
+
+function createBaseGetPrefillJobStatusRequest(): GetPrefillJobStatusRequest {
+ return { jobId: "", tenantId: "" };
+}
+
+export const GetPrefillJobStatusRequest: MessageFns = {
+ encode(message: GetPrefillJobStatusRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
+ if (message.jobId !== "") {
+ writer.uint32(10).string(message.jobId);
+ }
+ if (message.tenantId !== "") {
+ writer.uint32(18).string(message.tenantId);
+ }
+ return writer;
+ },
+
+ decode(input: BinaryReader | Uint8Array, length?: number): GetPrefillJobStatusRequest {
+ const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
+ const end = length === undefined ? reader.len : reader.pos + length;
+ const message = createBaseGetPrefillJobStatusRequest();
+ while (reader.pos < end) {
+ const tag = reader.uint32();
+ switch (tag >>> 3) {
+ case 1: {
+ if (tag !== 10) {
+ break;
+ }
+
+ message.jobId = reader.string();
+ continue;
+ }
+ case 2: {
+ if (tag !== 18) {
+ break;
+ }
+
+ message.tenantId = reader.string();
+ continue;
+ }
+ }
+ if ((tag & 7) === 4 || tag === 0) {
+ break;
+ }
+ reader.skip(tag & 7);
+ }
+ return message;
+ },
+
+ fromJSON(object: any): GetPrefillJobStatusRequest {
+ return {
+ jobId: isSet(object.jobId)
+ ? globalThis.String(object.jobId)
+ : isSet(object.job_id)
+ ? globalThis.String(object.job_id)
+ : "",
+ tenantId: isSet(object.tenantId)
+ ? globalThis.String(object.tenantId)
+ : isSet(object.tenant_id)
+ ? globalThis.String(object.tenant_id)
+ : "",
+ };
+ },
+
+ toJSON(message: GetPrefillJobStatusRequest): unknown {
+ const obj: any = {};
+ if (message.jobId !== "") {
+ obj.jobId = message.jobId;
+ }
+ if (message.tenantId !== "") {
+ obj.tenantId = message.tenantId;
+ }
+ return obj;
+ },
+
+ create, I>>(base?: I): GetPrefillJobStatusRequest {
+ return GetPrefillJobStatusRequest.fromPartial(base ?? ({} as any));
+ },
+ fromPartial, I>>(object: I): GetPrefillJobStatusRequest {
+ const message = createBaseGetPrefillJobStatusRequest();
+ message.jobId = object.jobId ?? "";
+ message.tenantId = object.tenantId ?? "";
+ return message;
+ },
+};
+
+function createBaseGetPrefillJobStatusResponse(): GetPrefillJobStatusResponse {
+ return {
+ jobId: "",
+ status: 0,
+ tokensComputed: 0,
+ leaseId: "",
+ blockHash: new Uint8Array(0),
+ payloadSha256: new Uint8Array(0),
+ transferBytes: "0",
+ failureReason: "",
+ computeMs: 0,
+ cacheAddress: "",
+ };
+}
+
+export const GetPrefillJobStatusResponse: MessageFns = {
+ encode(message: GetPrefillJobStatusResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
+ if (message.jobId !== "") {
+ writer.uint32(10).string(message.jobId);
+ }
+ if (message.status !== 0) {
+ writer.uint32(16).int32(message.status);
+ }
+ if (message.tokensComputed !== 0) {
+ writer.uint32(24).uint32(message.tokensComputed);
+ }
+ if (message.leaseId !== "") {
+ writer.uint32(34).string(message.leaseId);
+ }
+ if (message.blockHash.length !== 0) {
+ writer.uint32(42).bytes(message.blockHash);
+ }
+ if (message.payloadSha256.length !== 0) {
+ writer.uint32(50).bytes(message.payloadSha256);
}
- if (message.blockSha256.length !== 0) {
- writer.uint32(58).bytes(message.blockSha256);
+ if (message.transferBytes !== "0") {
+ writer.uint32(56).uint64(message.transferBytes);
}
- if (message.cacheEpoch !== "0") {
- writer.uint32(64).uint64(message.cacheEpoch);
+ if (message.failureReason !== "") {
+ writer.uint32(66).string(message.failureReason);
}
- if (message.compatibility !== undefined) {
- CacheCompatibility.encode(message.compatibility, writer.uint32(74).fork()).join();
+ if (message.computeMs !== 0) {
+ writer.uint32(73).double(message.computeMs);
+ }
+ if (message.cacheAddress !== "") {
+ writer.uint32(82).string(message.cacheAddress);
}
return writer;
},
- decode(input: BinaryReader | Uint8Array, length?: number): PublishBlockRequest {
+ decode(input: BinaryReader | Uint8Array, length?: number): GetPrefillJobStatusResponse {
const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
const end = length === undefined ? reader.len : reader.pos + length;
- const message = createBasePublishBlockRequest();
+ const message = createBaseGetPrefillJobStatusResponse();
while (reader.pos < end) {
const tag = reader.uint32();
switch (tag >>> 3) {
@@ -2366,7 +3636,7 @@ export const PublishBlockRequest: MessageFns = {
break;
}
- message.blockHash = reader.bytes();
+ message.jobId = reader.string();
continue;
}
case 2: {
@@ -2374,7 +3644,7 @@ export const PublishBlockRequest: MessageFns = {
break;
}
- message.blockIndex = reader.uint32();
+ message.status = reader.int32() as any;
continue;
}
case 3: {
@@ -2382,23 +3652,23 @@ export const PublishBlockRequest: MessageFns = {
break;
}
- message.tokenCount = reader.uint32();
+ message.tokensComputed = reader.uint32();
continue;
}
case 4: {
- if (tag !== 32) {
+ if (tag !== 34) {
break;
}
- message.chunkIndex = reader.uint32();
+ message.leaseId = reader.string();
continue;
}
case 5: {
- if (tag !== 40) {
+ if (tag !== 42) {
break;
}
- message.totalChunks = reader.uint32();
+ message.blockHash = reader.bytes();
continue;
}
case 6: {
@@ -2406,31 +3676,39 @@ export const PublishBlockRequest: MessageFns = {
break;
}
- message.data = reader.bytes();
+ message.payloadSha256 = reader.bytes();
continue;
}
case 7: {
- if (tag !== 58) {
+ if (tag !== 56) {
break;
}
- message.blockSha256 = reader.bytes();
+ message.transferBytes = reader.uint64().toString();
continue;
}
case 8: {
- if (tag !== 64) {
+ if (tag !== 66) {
break;
}
- message.cacheEpoch = reader.uint64().toString();
+ message.failureReason = reader.string();
continue;
}
case 9: {
- if (tag !== 74) {
+ if (tag !== 73) {
break;
}
- message.compatibility = CacheCompatibility.decode(reader, reader.uint32());
+ message.computeMs = reader.double();
+ continue;
+ }
+ case 10: {
+ if (tag !== 82) {
+ break;
+ }
+
+ message.cacheAddress = reader.string();
continue;
}
}
@@ -2442,136 +3720,147 @@ export const PublishBlockRequest: MessageFns = {
return message;
},
- fromJSON(object: any): PublishBlockRequest {
+ fromJSON(object: any): GetPrefillJobStatusResponse {
return {
+ jobId: isSet(object.jobId)
+ ? globalThis.String(object.jobId)
+ : isSet(object.job_id)
+ ? globalThis.String(object.job_id)
+ : "",
+ status: isSet(object.status) ? prefillJobStatusFromJSON(object.status) : 0,
+ tokensComputed: isSet(object.tokensComputed)
+ ? globalThis.Number(object.tokensComputed)
+ : isSet(object.tokens_computed)
+ ? globalThis.Number(object.tokens_computed)
+ : 0,
+ leaseId: isSet(object.leaseId)
+ ? globalThis.String(object.leaseId)
+ : isSet(object.lease_id)
+ ? globalThis.String(object.lease_id)
+ : "",
blockHash: isSet(object.blockHash)
? bytesFromBase64(object.blockHash)
: isSet(object.block_hash)
? bytesFromBase64(object.block_hash)
: new Uint8Array(0),
- blockIndex: isSet(object.blockIndex)
- ? globalThis.Number(object.blockIndex)
- : isSet(object.block_index)
- ? globalThis.Number(object.block_index)
- : 0,
- tokenCount: isSet(object.tokenCount)
- ? globalThis.Number(object.tokenCount)
- : isSet(object.token_count)
- ? globalThis.Number(object.token_count)
- : 0,
- chunkIndex: isSet(object.chunkIndex)
- ? globalThis.Number(object.chunkIndex)
- : isSet(object.chunk_index)
- ? globalThis.Number(object.chunk_index)
- : 0,
- totalChunks: isSet(object.totalChunks)
- ? globalThis.Number(object.totalChunks)
- : isSet(object.total_chunks)
- ? globalThis.Number(object.total_chunks)
- : 0,
- data: isSet(object.data) ? bytesFromBase64(object.data) : new Uint8Array(0),
- blockSha256: isSet(object.blockSha256)
- ? bytesFromBase64(object.blockSha256)
- : isSet(object.block_sha256)
- ? bytesFromBase64(object.block_sha256)
+ payloadSha256: isSet(object.payloadSha256)
+ ? bytesFromBase64(object.payloadSha256)
+ : isSet(object.payload_sha256)
+ ? bytesFromBase64(object.payload_sha256)
: new Uint8Array(0),
- cacheEpoch: isSet(object.cacheEpoch)
- ? globalThis.String(object.cacheEpoch)
- : isSet(object.cache_epoch)
- ? globalThis.String(object.cache_epoch)
+ transferBytes: isSet(object.transferBytes)
+ ? globalThis.String(object.transferBytes)
+ : isSet(object.transfer_bytes)
+ ? globalThis.String(object.transfer_bytes)
: "0",
- compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined,
+ failureReason: isSet(object.failureReason)
+ ? globalThis.String(object.failureReason)
+ : isSet(object.failure_reason)
+ ? globalThis.String(object.failure_reason)
+ : "",
+ computeMs: isSet(object.computeMs)
+ ? globalThis.Number(object.computeMs)
+ : isSet(object.compute_ms)
+ ? globalThis.Number(object.compute_ms)
+ : 0,
+ cacheAddress: isSet(object.cacheAddress)
+ ? globalThis.String(object.cacheAddress)
+ : isSet(object.cache_address)
+ ? globalThis.String(object.cache_address)
+ : "",
};
},
- toJSON(message: PublishBlockRequest): unknown {
+ toJSON(message: GetPrefillJobStatusResponse): unknown {
const obj: any = {};
- if (message.blockHash.length !== 0) {
- obj.blockHash = base64FromBytes(message.blockHash);
+ if (message.jobId !== "") {
+ obj.jobId = message.jobId;
}
- if (message.blockIndex !== 0) {
- obj.blockIndex = Math.round(message.blockIndex);
+ if (message.status !== 0) {
+ obj.status = prefillJobStatusToJSON(message.status);
}
- if (message.tokenCount !== 0) {
- obj.tokenCount = Math.round(message.tokenCount);
+ if (message.tokensComputed !== 0) {
+ obj.tokensComputed = Math.round(message.tokensComputed);
}
- if (message.chunkIndex !== 0) {
- obj.chunkIndex = Math.round(message.chunkIndex);
+ if (message.leaseId !== "") {
+ obj.leaseId = message.leaseId;
}
- if (message.totalChunks !== 0) {
- obj.totalChunks = Math.round(message.totalChunks);
+ if (message.blockHash.length !== 0) {
+ obj.blockHash = base64FromBytes(message.blockHash);
}
- if (message.data.length !== 0) {
- obj.data = base64FromBytes(message.data);
+ if (message.payloadSha256.length !== 0) {
+ obj.payloadSha256 = base64FromBytes(message.payloadSha256);
}
- if (message.blockSha256.length !== 0) {
- obj.blockSha256 = base64FromBytes(message.blockSha256);
+ if (message.transferBytes !== "0") {
+ obj.transferBytes = message.transferBytes;
}
- if (message.cacheEpoch !== "0") {
- obj.cacheEpoch = message.cacheEpoch;
+ if (message.failureReason !== "") {
+ obj.failureReason = message.failureReason;
}
- if (message.compatibility !== undefined) {
- obj.compatibility = CacheCompatibility.toJSON(message.compatibility);
+ if (message.computeMs !== 0) {
+ obj.computeMs = message.computeMs;
+ }
+ if (message.cacheAddress !== "") {
+ obj.cacheAddress = message.cacheAddress;
}
return obj;
},
- create, I>>(base?: I): PublishBlockRequest {
- return PublishBlockRequest.fromPartial(base ?? ({} as any));
+ create, I>>(base?: I): GetPrefillJobStatusResponse {
+ return GetPrefillJobStatusResponse.fromPartial(base ?? ({} as any));
},
- fromPartial, I>>(object: I): PublishBlockRequest {
- const message = createBasePublishBlockRequest();
+ fromPartial, I>>(object: I): GetPrefillJobStatusResponse {
+ const message = createBaseGetPrefillJobStatusResponse();
+ message.jobId = object.jobId ?? "";
+ message.status = object.status ?? 0;
+ message.tokensComputed = object.tokensComputed ?? 0;
+ message.leaseId = object.leaseId ?? "";
message.blockHash = object.blockHash ?? new Uint8Array(0);
- message.blockIndex = object.blockIndex ?? 0;
- message.tokenCount = object.tokenCount ?? 0;
- message.chunkIndex = object.chunkIndex ?? 0;
- message.totalChunks = object.totalChunks ?? 0;
- message.data = object.data ?? new Uint8Array(0);
- message.blockSha256 = object.blockSha256 ?? new Uint8Array(0);
- message.cacheEpoch = object.cacheEpoch ?? "0";
- message.compatibility = (object.compatibility !== undefined && object.compatibility !== null)
- ? CacheCompatibility.fromPartial(object.compatibility)
- : undefined;
+ message.payloadSha256 = object.payloadSha256 ?? new Uint8Array(0);
+ message.transferBytes = object.transferBytes ?? "0";
+ message.failureReason = object.failureReason ?? "";
+ message.computeMs = object.computeMs ?? 0;
+ message.cacheAddress = object.cacheAddress ?? "";
return message;
},
};
-function createBasePublishBlockResponse(): PublishBlockResponse {
- return { stored: false, cacheEpoch: "0" };
+function createBaseCancelPrefillJobRequest(): CancelPrefillJobRequest {
+ return { jobId: "", tenantId: "" };
}
-export const PublishBlockResponse: MessageFns = {
- encode(message: PublishBlockResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
- if (message.stored !== false) {
- writer.uint32(8).bool(message.stored);
+export const CancelPrefillJobRequest: MessageFns = {
+ encode(message: CancelPrefillJobRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
+ if (message.jobId !== "") {
+ writer.uint32(10).string(message.jobId);
}
- if (message.cacheEpoch !== "0") {
- writer.uint32(16).uint64(message.cacheEpoch);
+ if (message.tenantId !== "") {
+ writer.uint32(18).string(message.tenantId);
}
return writer;
},
- decode(input: BinaryReader | Uint8Array, length?: number): PublishBlockResponse {
+ decode(input: BinaryReader | Uint8Array, length?: number): CancelPrefillJobRequest {
const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
const end = length === undefined ? reader.len : reader.pos + length;
- const message = createBasePublishBlockResponse();
+ const message = createBaseCancelPrefillJobRequest();
while (reader.pos < end) {
const tag = reader.uint32();
switch (tag >>> 3) {
case 1: {
- if (tag !== 8) {
+ if (tag !== 10) {
break;
}
- message.stored = reader.bool();
+ message.jobId = reader.string();
continue;
}
case 2: {
- if (tag !== 16) {
+ if (tag !== 18) {
break;
}
- message.cacheEpoch = reader.uint64().toString();
+ message.tenantId = reader.string();
continue;
}
}
@@ -2583,35 +3872,97 @@ export const PublishBlockResponse: MessageFns = {
return message;
},
- fromJSON(object: any): PublishBlockResponse {
+ fromJSON(object: any): CancelPrefillJobRequest {
return {
- stored: isSet(object.stored) ? globalThis.Boolean(object.stored) : false,
- cacheEpoch: isSet(object.cacheEpoch)
- ? globalThis.String(object.cacheEpoch)
- : isSet(object.cache_epoch)
- ? globalThis.String(object.cache_epoch)
- : "0",
+ jobId: isSet(object.jobId)
+ ? globalThis.String(object.jobId)
+ : isSet(object.job_id)
+ ? globalThis.String(object.job_id)
+ : "",
+ tenantId: isSet(object.tenantId)
+ ? globalThis.String(object.tenantId)
+ : isSet(object.tenant_id)
+ ? globalThis.String(object.tenant_id)
+ : "",
};
},
- toJSON(message: PublishBlockResponse): unknown {
+ toJSON(message: CancelPrefillJobRequest): unknown {
const obj: any = {};
- if (message.stored !== false) {
- obj.stored = message.stored;
+ if (message.jobId !== "") {
+ obj.jobId = message.jobId;
}
- if (message.cacheEpoch !== "0") {
- obj.cacheEpoch = message.cacheEpoch;
+ if (message.tenantId !== "") {
+ obj.tenantId = message.tenantId;
}
return obj;
},
- create, I>>(base?: I): PublishBlockResponse {
- return PublishBlockResponse.fromPartial(base ?? ({} as any));
+ create, I>>(base?: I): CancelPrefillJobRequest {
+ return CancelPrefillJobRequest.fromPartial(base ?? ({} as any));
},
- fromPartial, I>>(object: I): PublishBlockResponse {
- const message = createBasePublishBlockResponse();
- message.stored = object.stored ?? false;
- message.cacheEpoch = object.cacheEpoch ?? "0";
+ fromPartial, I>>(object: I): CancelPrefillJobRequest {
+ const message = createBaseCancelPrefillJobRequest();
+ message.jobId = object.jobId ?? "";
+ message.tenantId = object.tenantId ?? "";
+ return message;
+ },
+};
+
+function createBaseCancelPrefillJobResponse(): CancelPrefillJobResponse {
+ return { cancelled: false };
+}
+
+export const CancelPrefillJobResponse: MessageFns = {
+ encode(message: CancelPrefillJobResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter {
+ if (message.cancelled !== false) {
+ writer.uint32(8).bool(message.cancelled);
+ }
+ return writer;
+ },
+
+ decode(input: BinaryReader | Uint8Array, length?: number): CancelPrefillJobResponse {
+ const reader = input instanceof BinaryReader ? input : new BinaryReader(input);
+ const end = length === undefined ? reader.len : reader.pos + length;
+ const message = createBaseCancelPrefillJobResponse();
+ while (reader.pos < end) {
+ const tag = reader.uint32();
+ switch (tag >>> 3) {
+ case 1: {
+ if (tag !== 8) {
+ break;
+ }
+
+ message.cancelled = reader.bool();
+ continue;
+ }
+ }
+ if ((tag & 7) === 4 || tag === 0) {
+ break;
+ }
+ reader.skip(tag & 7);
+ }
+ return message;
+ },
+
+ fromJSON(object: any): CancelPrefillJobResponse {
+ return { cancelled: isSet(object.cancelled) ? globalThis.Boolean(object.cancelled) : false };
+ },
+
+ toJSON(message: CancelPrefillJobResponse): unknown {
+ const obj: any = {};
+ if (message.cancelled !== false) {
+ obj.cancelled = message.cancelled;
+ }
+ return obj;
+ },
+
+ create, I>>(base?: I): CancelPrefillJobResponse {
+ return CancelPrefillJobResponse.fromPartial(base ?? ({} as any));
+ },
+ fromPartial, I>>(object: I): CancelPrefillJobResponse {
+ const message = createBaseCancelPrefillJobResponse();
+ message.cancelled = object.cancelled ?? false;
return message;
},
};
@@ -4373,6 +5724,111 @@ export const PrefillCacheServiceClient = makeGenericClientConstructor(
serviceName: string;
};
+/**
+ * PrefillWorkerService executes model prefill only. Workers load the same model
+ * as the primary, materialize immutable snapshots into their co-located
+ * PrefillCacheService, and never participate in autoregressive decode.
+ */
+export type PrefillWorkerServiceService = typeof PrefillWorkerServiceService;
+export const PrefillWorkerServiceService = {
+ submitPrefillJob: {
+ path: "/kakeya.v1.PrefillWorkerService/SubmitPrefillJob" as const,
+ requestStream: false as const,
+ responseStream: false as const,
+ requestSerialize: (value: SubmitPrefillJobRequest): Buffer =>
+ Buffer.from(SubmitPrefillJobRequest.encode(value).finish()),
+ requestDeserialize: (value: Buffer): SubmitPrefillJobRequest => SubmitPrefillJobRequest.decode(value),
+ responseSerialize: (value: SubmitPrefillJobResponse): Buffer =>
+ Buffer.from(SubmitPrefillJobResponse.encode(value).finish()),
+ responseDeserialize: (value: Buffer): SubmitPrefillJobResponse => SubmitPrefillJobResponse.decode(value),
+ },
+ getPrefillJobStatus: {
+ path: "/kakeya.v1.PrefillWorkerService/GetPrefillJobStatus" as const,
+ requestStream: false as const,
+ responseStream: false as const,
+ requestSerialize: (value: GetPrefillJobStatusRequest): Buffer =>
+ Buffer.from(GetPrefillJobStatusRequest.encode(value).finish()),
+ requestDeserialize: (value: Buffer): GetPrefillJobStatusRequest => GetPrefillJobStatusRequest.decode(value),
+ responseSerialize: (value: GetPrefillJobStatusResponse): Buffer =>
+ Buffer.from(GetPrefillJobStatusResponse.encode(value).finish()),
+ responseDeserialize: (value: Buffer): GetPrefillJobStatusResponse => GetPrefillJobStatusResponse.decode(value),
+ },
+ cancelPrefillJob: {
+ path: "/kakeya.v1.PrefillWorkerService/CancelPrefillJob" as const,
+ requestStream: false as const,
+ responseStream: false as const,
+ requestSerialize: (value: CancelPrefillJobRequest): Buffer =>
+ Buffer.from(CancelPrefillJobRequest.encode(value).finish()),
+ requestDeserialize: (value: Buffer): CancelPrefillJobRequest => CancelPrefillJobRequest.decode(value),
+ responseSerialize: (value: CancelPrefillJobResponse): Buffer =>
+ Buffer.from(CancelPrefillJobResponse.encode(value).finish()),
+ responseDeserialize: (value: Buffer): CancelPrefillJobResponse => CancelPrefillJobResponse.decode(value),
+ },
+} as const;
+
+export interface PrefillWorkerServiceServer extends UntypedServiceImplementation {
+ submitPrefillJob: handleUnaryCall;
+ getPrefillJobStatus: handleUnaryCall;
+ cancelPrefillJob: handleUnaryCall;
+}
+
+export interface PrefillWorkerServiceClient extends Client {
+ submitPrefillJob(
+ request: SubmitPrefillJobRequest,
+ callback: (error: ServiceError | null, response: SubmitPrefillJobResponse) => void,
+ ): ClientUnaryCall;
+ submitPrefillJob(
+ request: SubmitPrefillJobRequest,
+ metadata: Metadata,
+ callback: (error: ServiceError | null, response: SubmitPrefillJobResponse) => void,
+ ): ClientUnaryCall;
+ submitPrefillJob(
+ request: SubmitPrefillJobRequest,
+ metadata: Metadata,
+ options: Partial,
+ callback: (error: ServiceError | null, response: SubmitPrefillJobResponse) => void,
+ ): ClientUnaryCall;
+ getPrefillJobStatus(
+ request: GetPrefillJobStatusRequest,
+ callback: (error: ServiceError | null, response: GetPrefillJobStatusResponse) => void,
+ ): ClientUnaryCall;
+ getPrefillJobStatus(
+ request: GetPrefillJobStatusRequest,
+ metadata: Metadata,
+ callback: (error: ServiceError | null, response: GetPrefillJobStatusResponse) => void,
+ ): ClientUnaryCall;
+ getPrefillJobStatus(
+ request: GetPrefillJobStatusRequest,
+ metadata: Metadata,
+ options: Partial,
+ callback: (error: ServiceError | null, response: GetPrefillJobStatusResponse) => void,
+ ): ClientUnaryCall;
+ cancelPrefillJob(
+ request: CancelPrefillJobRequest,
+ callback: (error: ServiceError | null, response: CancelPrefillJobResponse) => void,
+ ): ClientUnaryCall;
+ cancelPrefillJob(
+ request: CancelPrefillJobRequest,
+ metadata: Metadata,
+ callback: (error: ServiceError | null, response: CancelPrefillJobResponse) => void,
+ ): ClientUnaryCall;
+ cancelPrefillJob(
+ request: CancelPrefillJobRequest,
+ metadata: Metadata,
+ options: Partial,
+ callback: (error: ServiceError | null, response: CancelPrefillJobResponse) => void,
+ ): ClientUnaryCall;
+}
+
+export const PrefillWorkerServiceClient = makeGenericClientConstructor(
+ PrefillWorkerServiceService,
+ "kakeya.v1.PrefillWorkerService",
+) as unknown as {
+ new (address: string, credentials: ChannelCredentials, options?: Partial): PrefillWorkerServiceClient;
+ service: typeof PrefillWorkerServiceService;
+ serviceName: string;
+};
+
/**
* DFlashProposerService: stateful remote DFlash drafter + f_θ restoration.
* Per turn: Restore (prompt -> f_θ-projected verifier K/V) then SeedContext
diff --git a/tests/inference_engine/bridge/test_manifest.py b/tests/inference_engine/bridge/test_manifest.py
index ce623c9e..02b51f59 100644
--- a/tests/inference_engine/bridge/test_manifest.py
+++ b/tests/inference_engine/bridge/test_manifest.py
@@ -96,6 +96,7 @@ def test_allowlist_contains_exactly_the_documented_presets():
"mlx-kakeya-launcher-full",
"mlx-kakeya-launcher-smoke",
"mlx-multitenant-pressure",
+ "mlx-prefill-snapshot-equivalence",
"mlx-upgrade",
"mlx-upstream-batch-probe",
"pytest-path",
diff --git a/tests/inference_engine/bridge/test_runner_python.py b/tests/inference_engine/bridge/test_runner_python.py
index 8a0a6bb5..acf3c5a4 100644
--- a/tests/inference_engine/bridge/test_runner_python.py
+++ b/tests/inference_engine/bridge/test_runner_python.py
@@ -30,6 +30,8 @@ def test_candidates_prioritise_pin_then_venvs_then_path():
"/pin/bin/python",
"/home/me/kakeya-venv/bin/python",
"/home/me/.venv/bin/python",
+ "/home/me/Documents/Kakeya-LLM-Inference-engine-pr109/.venv-mac/bin/python3.13",
+ "/home/me/Documents/Kakeya-LLM-Inference-engine-pr109/.venv-mac/bin/python",
"/usr/bin/python3.13",
"/usr/bin/python3",
]
@@ -44,6 +46,8 @@ def test_candidates_drop_empty_and_dedupe():
assert cands == [
"/home/me/kakeya-venv/bin/python",
"/home/me/.venv/bin/python",
+ "/home/me/Documents/Kakeya-LLM-Inference-engine-pr109/.venv-mac/bin/python3.13",
+ "/home/me/Documents/Kakeya-LLM-Inference-engine-pr109/.venv-mac/bin/python",
]
assert None not in cands
diff --git a/tests/inference_engine/distributed/test_capability.py b/tests/inference_engine/distributed/test_capability.py
index c8978ef8..e730eed3 100644
--- a/tests/inference_engine/distributed/test_capability.py
+++ b/tests/inference_engine/distributed/test_capability.py
@@ -18,9 +18,11 @@
CapabilityRole,
CacheCapability,
CacheCompatibility,
+ CompressionCodec,
ModelCapability,
NodeCapability,
NodeEndpoint,
+ PrefillWorkerCapability,
)
T0 = 1_000_000.0
@@ -115,6 +117,7 @@ def test_cache_capability_and_endpoints_proto_round_trip():
layer_geometry_hash="geometry",
kv_dtype="bfloat16",
block_size_tokens=64,
+ tenant_namespace="tenant",
)
card = NodeCapability(
node_id="cache-peer",
@@ -130,6 +133,8 @@ def test_cache_capability_and_endpoints_proto_round_trip():
load=0.5,
tokens_served=100,
bloom_filter=b"filter",
+ default_compression=CompressionCodec.ZLIB,
+ replication_factor=2,
),
),
endpoints=(
@@ -140,6 +145,19 @@ def test_cache_capability_and_endpoints_proto_round_trip():
0.45,
),
),
+ prefill_workers=(
+ PrefillWorkerCapability(
+ compatibility,
+ worker_address="169.254.27.104:53051",
+ max_concurrent_jobs=1,
+ inflight_jobs=1,
+ queued_jobs=2,
+ load=0.5,
+ tokens_per_second_prefill=33.0,
+ ram_bytes_free=1234,
+ queued_tokens=456,
+ ),
+ ),
)
assert NodeCapability.from_proto(card.to_proto()) == card
assert CapabilityRole.PREFILL_CACHE.value == 5
diff --git a/tests/inference_engine/distributed/test_prefill_auth.py b/tests/inference_engine/distributed/test_prefill_auth.py
new file mode 100644
index 00000000..bea8fee6
--- /dev/null
+++ b/tests/inference_engine/distributed/test_prefill_auth.py
@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import pytest
+
+from inference_engine.distributed.prefill_auth import (
+ FleetAuthConfig,
+ PrefillAuthError,
+ metadata_pairs,
+ signed_metadata,
+ verify_metadata,
+)
+from inference_engine.server.proto_gen.kakeya.v1 import distributed_pb2
+
+
+def _request():
+ return distributed_pb2.FetchBlocksRequest(lease_id="lease")
+
+
+def _config():
+ return FleetAuthConfig(b"x" * 32, "tenant-a", "node-a", 30)
+
+
+def test_sign_and_verify_round_trip():
+ request = _request()
+ metadata = signed_metadata(request, _config(), now=100)
+ assert verify_metadata(metadata, request, _config(), now=110) == (
+ "tenant-a", "node-a",
+ )
+
+
+def test_tenant_hash_keys_are_isolated():
+ a = _config().tenant_hash_key()
+ b = FleetAuthConfig(b"x" * 32, "tenant-b", "node-a").tenant_hash_key()
+ assert a != b
+
+
+@pytest.mark.parametrize("mutator,match", [
+ (lambda md: tuple((k, "tenant-b" if k.endswith("tenant-id") else v)
+ for k, v in md), "tenant mismatch"),
+ (lambda md: tuple((k, "bad" if k.endswith("auth-mac") else v)
+ for k, v in md), "invalid prefill authentication MAC"),
+ (lambda md: tuple((k, v) for k, v in md if not k.endswith("auth-mac")),
+ "missing prefill authentication metadata"),
+])
+def test_auth_rejects_bad_metadata(mutator, match):
+ request = _request()
+ with pytest.raises(PrefillAuthError, match=match):
+ verify_metadata(
+ mutator(signed_metadata(request, _config(), now=100)),
+ request,
+ _config(),
+ now=100,
+ )
+
+
+def test_auth_rejects_replay_and_bad_timestamp():
+ request = _request()
+ metadata = signed_metadata(request, _config(), now=100)
+ with pytest.raises(PrefillAuthError, match="replay window"):
+ verify_metadata(metadata, request, _config(), now=200)
+ bad = tuple(("x-kakeya-auth-ts", "nan") if k == "x-kakeya-auth-ts"
+ else (k, v) for k, v in metadata)
+ with pytest.raises(PrefillAuthError, match="invalid authentication timestamp"):
+ verify_metadata(bad, request, _config(), now=100)
+
+
+def test_auth_rejects_tampered_request():
+ request = _request()
+ metadata = signed_metadata(request, _config(), now=100)
+ with pytest.raises(PrefillAuthError, match="invalid prefill authentication MAC"):
+ verify_metadata(
+ metadata,
+ distributed_pb2.FetchBlocksRequest(lease_id="other"),
+ _config(),
+ now=100,
+ )
+
+
+def test_config_validation_and_file(tmp_path):
+ with pytest.raises(ValueError):
+ FleetAuthConfig(b"short", "t", "n")
+ with pytest.raises(ValueError):
+ FleetAuthConfig(b"x" * 32, "", "n")
+ with pytest.raises(ValueError):
+ FleetAuthConfig(b"x" * 32, "t", "")
+ with pytest.raises(ValueError):
+ FleetAuthConfig(b"x" * 32, "t", "n", 0)
+ path = tmp_path / "psk"
+ path.write_bytes(b"y" * 32 + b"\n")
+ assert FleetAuthConfig.from_file(
+ str(path), tenant_id="t", node_id="n",
+ ).psk == b"y" * 32
+
+
+def test_metadata_pairs_accepts_tuples():
+ assert metadata_pairs((("a", "b"),)) == (("a", "b"),)
+
+ class Item:
+ key = "c"
+ value = "d"
+
+ assert metadata_pairs((Item(),)) == (("c", "d"),)
+
+
+def test_sign_rejects_non_protobuf():
+ with pytest.raises(TypeError):
+ signed_metadata(object(), _config())
+
diff --git a/tests/inference_engine/distributed/test_prefill_cache.py b/tests/inference_engine/distributed/test_prefill_cache.py
index 5295080b..2aacbae8 100644
--- a/tests/inference_engine/distributed/test_prefill_cache.py
+++ b/tests/inference_engine/distributed/test_prefill_cache.py
@@ -104,3 +104,16 @@ def test_pinned_eviction_and_missing_leased_block_guards():
store._blocks.pop(block.block_hash)
with pytest.raises(KeyError, match="evicted"):
store.fetch(lease.lease_id, now=1)
+
+
+def test_put_rejects_when_active_lease_pins_capacity():
+ import time
+
+ store = PrefixCacheStore(_compat(), max_bytes=5, node_id="x")
+ first = CacheBlock.create(bytes(32), 1, b"12345")
+ store.put(first)
+ store.lookup([first.block_hash], now=time.time())
+ second = CacheBlock.create(bytes.fromhex("01" * 32), 1, b"abc")
+ with pytest.raises(ValueError, match="pinned"):
+ store.put(second)
+ assert store.block_hashes() == (first.block_hash,)
diff --git a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py
new file mode 100644
index 00000000..5dd696d1
--- /dev/null
+++ b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py
@@ -0,0 +1,593 @@
+from __future__ import annotations
+
+from inference_engine.distributed.capability import (
+ CacheCapability,
+ CacheCompatibility,
+ CompressionCodec,
+ NodeCapability,
+ NodeEndpoint,
+ PrefillWorkerCapability,
+)
+from inference_engine.distributed.prefill_cache import (
+ CacheBlock,
+ PrefixCacheStore,
+ chained_block_hashes,
+)
+from inference_engine.distributed.prefill_cache_runtime import (
+ DistributedPrefillCacheHook,
+ _Hit,
+)
+from inference_engine.distributed.prefill_scheduler import PrefillCostConfig
+from inference_engine.server.proto_gen.kakeya.v1 import distributed_pb2
+
+
+class _Verifier:
+ def __init__(self):
+ self.cache = []
+ self.cached_token_sequence = []
+ self.next_global_position = 0
+ self.next_token_logits = None
+ self.prefill_calls = 0
+
+ def reset(self):
+ self.cache = []
+ self.cached_token_sequence = []
+ self.next_global_position = 0
+
+ def prefill(self, tokens):
+ self.reset()
+ self.prefill_calls += 1
+ self.cached_token_sequence = list(tokens)
+ self.next_global_position = len(tokens)
+ self.next_token_logits = _Row(tokens[-1])
+
+ def forward_block(self, tokens):
+ self.cached_token_sequence.extend(tokens)
+ self.next_global_position += len(tokens)
+ return [_Row(token) for token in tokens]
+
+ def commit_or_truncate(self, *, forwarded, accepted):
+ assert forwarded == accepted
+
+
+class _Row:
+ def __init__(self, value):
+ self.value = value
+
+ def clone(self):
+ return _Row(self.value)
+
+
+def test_import_failure_always_falls_back_to_full_local_prefill(monkeypatch):
+ compatibility = CacheCompatibility(model_id="m", block_size_tokens=2)
+ store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head")
+ hashes = chained_block_hashes([1, 2], compatibility)
+ store.put(CacheBlock.create(hashes[0], 2, b"corrupt"))
+ hook = DistributedPrefillCacheHook(
+ store,
+ compression=CompressionCodec.NONE,
+ )
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "import_mlx_prefill_snapshot",
+ lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("bad snapshot")),
+ )
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "export_mlx_prefill_snapshot",
+ lambda *args, **kwargs: b"fresh",
+ )
+ verifier = _Verifier()
+ assert hook.prepare(verifier, [1, 2]) == 0
+ assert verifier.prefill_calls == 1
+ assert verifier.cached_token_sequence == [1, 2]
+ assert hook.stats.fallbacks == 1
+ assert "bad snapshot" in hook.stats.last_fallback_reason
+ assert not store.invalidate(b"missing")
+ hook.close()
+
+
+def test_tenant_hmac_changes_hash_namespace(monkeypatch, tmp_path):
+ from inference_engine.distributed.prefill_auth import FleetAuthConfig
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "export_mlx_prefill_snapshot",
+ lambda *args, **kwargs: b"snapshot",
+ )
+ a = CacheCompatibility(
+ model_id="m", block_size_tokens=2, tenant_namespace="a",
+ )
+ b = CacheCompatibility(
+ model_id="m", block_size_tokens=2, tenant_namespace="b",
+ )
+ store_a = PrefixCacheStore(a, max_bytes=1024, node_id="a")
+ store_b = PrefixCacheStore(b, max_bytes=1024, node_id="b")
+ auth_a = FleetAuthConfig(b"x" * 32, "a", "node")
+ auth_b = FleetAuthConfig(b"x" * 32, "b", "node")
+ hook_a = DistributedPrefillCacheHook(
+ store_a, auth=auth_a, compression=CompressionCodec.NONE,
+ )
+ hook_b = DistributedPrefillCacheHook(
+ store_b, auth=auth_b, compression=CompressionCodec.NONE,
+ )
+ hook_a.prepare(_Verifier(), [1, 2])
+ hook_b.prepare(_Verifier(), [1, 2])
+ assert store_a.block_hashes() != store_b.block_hashes()
+ hook_a.close()
+ hook_b.close()
+
+
+def test_runtime_validation_empty_and_provider_failure():
+ store = PrefixCacheStore(
+ CacheCompatibility(model_id="m"),
+ max_bytes=1024,
+ node_id="head",
+ )
+ with __import__("pytest").raises(ValueError):
+ DistributedPrefillCacheHook(store, lookup_timeout_s=0)
+ with __import__("pytest").raises(ValueError):
+ DistributedPrefillCacheHook(store, replication_factor=-1)
+ hook = DistributedPrefillCacheHook(
+ store,
+ registry_provider=lambda: (_ for _ in ()).throw(RuntimeError("boom")),
+ )
+ assert hook.prepare(_Verifier(), []) == 0
+ assert hook._cards() == ()
+ hook.close()
+
+
+def test_successful_local_import_suffix_and_on_reuse(monkeypatch):
+ compatibility = CacheCompatibility(model_id="m", block_size_tokens=2)
+ store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head")
+ hashes = chained_block_hashes([1, 2, 3, 4], compatibility)
+ store.put(CacheBlock.create(hashes[0], 2, b"snapshot"))
+ imported = type("Imported", (), {
+ "token_count": 2,
+ "cached_token_ids": (1, 2),
+ "next_token_logits": _Row(2),
+ "block_hash": hashes[0],
+ })()
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "import_mlx_prefill_snapshot",
+ lambda *args, **kwargs: imported,
+ )
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "export_mlx_prefill_snapshot",
+ lambda *args, **kwargs: b"fresh",
+ )
+ reused = []
+ hook = DistributedPrefillCacheHook(
+ store,
+ compression=CompressionCodec.NONE,
+ on_reuse=reused.append,
+ )
+ verifier = _Verifier()
+ assert hook.prepare(verifier, [1, 2, 3, 4]) == 2
+ assert verifier.cached_token_sequence == [1, 2, 3, 4]
+ assert hook.stats.local_hits == 1
+ assert reused == [2]
+ hook.close()
+
+
+def test_import_budget_and_remote_worker_failures(monkeypatch):
+ compatibility = CacheCompatibility(model_id="m", block_size_tokens=2)
+ store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head")
+ hook = DistributedPrefillCacheHook(
+ store,
+ max_import_bytes=2,
+ remote_compute_min_tokens=1,
+ worker_timeout_s=0.001,
+ worker_poll_interval_s=0.001,
+ )
+ verifier = _Verifier()
+ assert hook._try_import(
+ verifier, [1, 2], _Hit("local", "l", 1, 1, 1, b"x"),
+ ) == 0
+ assert hook._try_import(
+ verifier, [1, 2], _Hit("local", "l", 1, 2, 3, b"big"),
+ ) == 0
+ assert hook._try_import(
+ verifier, [1, 2], _Hit("peer", "l", 1, 2, 3),
+ ) == 0
+
+ invalid = [
+ type("Imported", (), {
+ "token_count": 0,
+ "cached_token_ids": (),
+ "next_token_logits": _Row(1),
+ "block_hash": b"h" * 32,
+ })(),
+ type("Imported", (), {
+ "token_count": 2,
+ "cached_token_ids": (1, 2),
+ "next_token_logits": _Row(1),
+ "block_hash": b"x" * 32,
+ })(),
+ type("Imported", (), {
+ "token_count": 2,
+ "cached_token_ids": (1, 2),
+ "next_token_logits": None,
+ "block_hash": b"h" * 32,
+ })(),
+ type("Imported", (), {
+ "token_count": 2,
+ "cached_token_ids": (9, 9),
+ "next_token_logits": _Row(1),
+ "block_hash": b"h" * 32,
+ })(),
+ ]
+ monkeypatch.setattr(
+ hook,
+ "_fetch_remote",
+ lambda hit: b"ok",
+ )
+ for imported in invalid:
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "import_mlx_prefill_snapshot",
+ lambda *args, _imported=imported, **kwargs: _imported,
+ )
+ assert hook._try_import(
+ verifier,
+ [1, 2],
+ _Hit("peer", "l", 1, 2, 2, block_hash=b"h" * 32),
+ ) == 0
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "choose_prefill_worker",
+ lambda *args, **kwargs: None,
+ )
+ assert hook._compute_remote([1, 2], [b"a" * 32]) is None
+
+ target = type("Target", (), {
+ "address": "worker:1",
+ "rtt_ms": 1.0,
+ })()
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "choose_prefill_worker",
+ lambda *args, **kwargs: target,
+ )
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "submit_prefill_job_sync",
+ lambda *args, **kwargs: type("Response", (), {"job_id": "j"})(),
+ )
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "get_prefill_job_sync",
+ lambda *args, **kwargs: type("Status", (), {
+ "status": 4,
+ "failure_reason": "failed",
+ })(),
+ )
+ assert hook._compute_remote([1, 2], [b"a" * 32]) is None
+ assert hook.stats.remote_job_failures == 1
+
+ # A perpetually queued job reaches the bounded deadline and falls back.
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "get_prefill_job_sync",
+ lambda *args, **kwargs: type("Status", (), {
+ "status": 1,
+ "failure_reason": "",
+ })(),
+ )
+ assert hook._compute_remote([1, 2], [b"a" * 32]) is None
+ assert hook.stats.remote_job_failures == 2
+ hook.close()
+
+
+def test_dynamic_replica_selection_and_cost_reject(monkeypatch):
+ compatibility = CacheCompatibility(model_id="m", block_size_tokens=2)
+ card = NodeCapability(
+ "peer",
+ "peer:1",
+ caches=(CacheCapability(compatibility, "peer:2"),),
+ endpoints=(NodeEndpoint("peer:2", "lan", 1, 5),),
+ )
+ store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head")
+ hook = DistributedPrefillCacheHook(
+ store,
+ registry_provider=lambda: (card,),
+ replication_factor=1,
+ cost_config=PrefillCostConfig(
+ local_prefill_tps=10000,
+ default_worker_tps=1,
+ link_mbps=1,
+ default_rtt_ms=10,
+ minimum_savings_ratio=0.5,
+ ),
+ )
+ assert hook._publish_peers(b"h" * 32) == ("peer:2",)
+ monkeypatch.setattr(
+ hook,
+ "_lookup_peer",
+ lambda peer, hashes: _Hit(peer, "l", 1, 1, 10_000_000),
+ )
+ assert hook._best_hit([b"h" * 32]) is None
+ assert hook._peer_rtt("peer:2") == 5
+ assert hook._peer_rtt("unknown") == 0
+ hook.close()
+
+
+def test_publish_boundary_dispatches_selected_replica(monkeypatch):
+ compatibility = CacheCompatibility(model_id="m", block_size_tokens=2)
+ store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head")
+ hook = DistributedPrefillCacheHook(store, peers=("peer:1",))
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "export_mlx_prefill_snapshot",
+ lambda *args, **kwargs: b"snapshot",
+ )
+ calls = []
+
+ class Publisher:
+ def submit(self, fn, *args, **kwargs):
+ calls.append((fn, args, kwargs))
+
+ def shutdown(self, **kwargs):
+ pass
+
+ hook._publisher = Publisher()
+ verifier = _Verifier()
+ verifier.prefill([1, 2])
+ hashes = chained_block_hashes([1, 2], compatibility)
+ hook._publish_boundary(verifier, [1, 2], hashes, 0, 2)
+ assert calls and calls[0][1][0] == "peer:1"
+ hook.close()
+
+
+def test_unaligned_reused_prefix_computes_boundary_remainder(monkeypatch):
+ compatibility = CacheCompatibility(model_id="m", block_size_tokens=2)
+ store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head")
+ hook = DistributedPrefillCacheHook(store)
+ monkeypatch.setattr(hook, "_publish_boundary", lambda *args, **kwargs: None)
+ verifier = _Verifier()
+ verifier.prefill([1, 2, 3])
+ hashes = chained_block_hashes([1, 2, 3, 4, 5, 6], compatibility)
+ hook._compute_and_publish(verifier, [1, 2, 3, 4, 5, 6], hashes, 3)
+ # Snapshot ends at token 3; token 4 completes that block before [5,6].
+ assert verifier.cached_token_sequence[-3:] == [4, 5, 6]
+ hook.close()
+
+
+class _Context:
+ def __init__(self, stub):
+ self.stub = stub
+
+ def __enter__(self):
+ return type("Channel", (), {
+ "unary_unary": lambda *args, **kwargs: None,
+ })()
+
+ def __exit__(self, *args):
+ return False
+
+
+def test_fetch_remote_validates_stream(monkeypatch):
+ import hashlib
+
+ compatibility = CacheCompatibility(model_id="m")
+ hook = DistributedPrefillCacheHook(
+ PrefixCacheStore(compatibility, max_bytes=1024, node_id="head"),
+ )
+ payload = b"payload"
+
+ class Stub:
+ def __init__(self, channel):
+ pass
+
+ def FetchBlocks(self, request, **kwargs):
+ return [distributed_pb2.FetchBlocksResponse(
+ chunk_index=0,
+ total_chunks=1,
+ data=payload,
+ block_sha256=hashlib.sha256(payload).digest(),
+ )]
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime.grpc.insecure_channel",
+ lambda address: _Context(Stub),
+ )
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ Stub,
+ )
+ assert hook._fetch_remote(_Hit("peer", "lease", 1, 1, 7)) == payload
+
+ class EmptyStub(Stub):
+ def FetchBlocks(self, request, **kwargs):
+ return []
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ EmptyStub,
+ )
+ with __import__("pytest").raises(RuntimeError, match="incomplete"):
+ hook._fetch_remote(_Hit("peer", "lease", 1, 1, 7))
+
+ class ErrorStub(Stub):
+ def FetchBlocks(self, request, **kwargs):
+ import grpc
+ raise grpc.RpcError("down")
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ ErrorStub,
+ )
+ with __import__("pytest").raises(RuntimeError, match="fetch failed"):
+ hook._fetch_remote(_Hit("peer", "lease", 1, 1, 7))
+
+ class CorruptStub(Stub):
+ def FetchBlocks(self, request, **kwargs):
+ return [distributed_pb2.FetchBlocksResponse(
+ chunk_index=0,
+ total_chunks=1,
+ data=payload,
+ block_sha256=b"x" * 32,
+ )]
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ CorruptStub,
+ )
+ with __import__("pytest").raises(RuntimeError, match="checksum"):
+ hook._fetch_remote(_Hit("peer", "lease", 1, 1, 7))
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ Stub,
+ )
+ with __import__("pytest").raises(RuntimeError, match="lease checksum"):
+ hook._fetch_remote(_Hit(
+ "peer",
+ "lease",
+ 1,
+ 1,
+ 7,
+ payload_sha256=b"x" * 32,
+ ))
+
+ with __import__("pytest").raises(RuntimeError, match="import budget"):
+ hook._fetch_remote(
+ _Hit(
+ "peer", "lease", 1, 1, hook.max_import_bytes + 1,
+ ),
+ )
+
+ class ChangedMetadataStub(Stub):
+ def FetchBlocks(self, request, **kwargs):
+ digest = hashlib.sha256(b"ab").digest()
+ return [
+ distributed_pb2.FetchBlocksResponse(
+ chunk_index=0, total_chunks=2, data=b"a",
+ block_hash=b"h", block_sha256=digest,
+ ),
+ distributed_pb2.FetchBlocksResponse(
+ chunk_index=1, total_chunks=3, data=b"b",
+ block_hash=b"h", block_sha256=digest,
+ ),
+ ]
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ ChangedMetadataStub,
+ )
+ with __import__("pytest").raises(RuntimeError, match="metadata changed"):
+ hook._fetch_remote(_Hit("peer", "lease", 1, 1, 2))
+
+ class TooManyChunksStub(Stub):
+ def FetchBlocks(self, request, **kwargs):
+ return [distributed_pb2.FetchBlocksResponse(
+ chunk_index=0,
+ total_chunks=65_537,
+ data=b"",
+ block_hash=b"h",
+ block_sha256=hashlib.sha256(b"").digest(),
+ )]
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ TooManyChunksStub,
+ )
+ with __import__("pytest").raises(RuntimeError, match="chunk count"):
+ hook._fetch_remote(_Hit("peer", "lease", 1, 1, 0))
+
+ class DuplicateStub(Stub):
+ def FetchBlocks(self, request, **kwargs):
+ digest = hashlib.sha256(b"aa").digest()
+ return [
+ distributed_pb2.FetchBlocksResponse(
+ chunk_index=0, total_chunks=2, data=b"a",
+ block_hash=b"h", block_sha256=digest,
+ ),
+ distributed_pb2.FetchBlocksResponse(
+ chunk_index=0, total_chunks=2, data=b"a",
+ block_hash=b"h", block_sha256=digest,
+ ),
+ ]
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ DuplicateStub,
+ )
+ with __import__("pytest").raises(RuntimeError, match="duplicate"):
+ hook._fetch_remote(_Hit("peer", "lease", 1, 1, 2))
+
+ hook.max_import_bytes = 1
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ Stub,
+ )
+ with __import__("pytest").raises(RuntimeError, match="import budget"):
+ hook._fetch_remote(_Hit("peer", "lease", 1, 1, 1))
+ hook.close()
+
+
+def test_lookup_peer_success_miss_and_rpc_error(monkeypatch):
+ import grpc
+
+ compatibility = CacheCompatibility(model_id="m")
+ hook = DistributedPrefillCacheHook(
+ PrefixCacheStore(compatibility, max_bytes=1024, node_id="head"),
+ )
+
+ class LookupStub:
+ response = distributed_pb2.LookupPrefixResponse(
+ lease_id="lease",
+ hit_block_count=1,
+ hit_token_count=5,
+ transfer_bytes=7,
+ )
+
+ def __init__(self, channel):
+ pass
+
+ def LookupPrefix(self, request, **kwargs):
+ return self.response
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime.grpc.insecure_channel",
+ lambda address: _Context(LookupStub),
+ )
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ LookupStub,
+ )
+ hit = hook._lookup_peer("peer", [b"h" * 32])
+ assert hit is not None and hit.hit_tokens == 5
+ LookupStub.response = distributed_pb2.LookupPrefixResponse()
+ assert hook._lookup_peer("peer", [b"h" * 32]) is None
+ LookupStub.response = distributed_pb2.LookupPrefixResponse(
+ lease_id="lease",
+ hit_block_count=2,
+ )
+ assert hook._lookup_peer("peer", [b"h" * 32]) is None
+
+ class ErrorStub(LookupStub):
+ def LookupPrefix(self, request, **kwargs):
+ raise grpc.RpcError("down")
+
+ monkeypatch.setattr(
+ "inference_engine.distributed.prefill_cache_runtime."
+ "distributed_pb2_grpc.PrefillCacheServiceStub",
+ ErrorStub,
+ )
+ assert hook._lookup_peer("peer", [b"h" * 32]) is None
+ hook.close()
+
diff --git a/tests/inference_engine/distributed/test_prefill_cache_service.py b/tests/inference_engine/distributed/test_prefill_cache_service.py
index e4f4cabd..ac093ab5 100644
--- a/tests/inference_engine/distributed/test_prefill_cache_service.py
+++ b/tests/inference_engine/distributed/test_prefill_cache_service.py
@@ -13,6 +13,10 @@
)
from inference_engine.distributed.prefill_cache import PrefixCacheStore
from inference_engine.distributed.prefill_cache import CacheBlock
+from inference_engine.distributed.prefill_auth import (
+ FleetAuthConfig,
+ signed_metadata,
+)
from inference_engine.distributed.prefill_cache_service import (
PrefillCacheServiceServicer,
add_prefill_cache_service,
@@ -67,6 +71,54 @@ async def test_lookup_and_fetch_over_real_grpc():
await server.stop(0)
+@pytest.mark.asyncio
+async def test_authenticated_cache_service_rejects_unsigned_requests():
+ compatibility = CacheCompatibility(
+ model_id="m",
+ block_size_tokens=2,
+ tenant_namespace="tenant",
+ )
+ store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="peer")
+ hashes = store.put_prefix([1, 2], [b"payload"])
+ auth = FleetAuthConfig(b"k" * 32, "tenant", "client")
+ server = grpc.aio.server()
+ port = server.add_insecure_port("127.0.0.1:0")
+ address = f"127.0.0.1:{port}"
+ add_prefill_cache_service(
+ server, store, cache_address=address, auth=auth,
+ )
+ await server.start()
+ try:
+ request = distributed_pb2.LookupPrefixRequest(
+ compatibility=_compat().to_proto(),
+ # overwritten below to use the authenticated tenant namespace
+ block_hashes=hashes,
+ )
+ request.compatibility.CopyFrom(compatibility.to_proto())
+ async with grpc.aio.insecure_channel(address) as channel:
+ stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel)
+ with pytest.raises(grpc.aio.AioRpcError) as error:
+ await stub.LookupPrefix(request)
+ assert error.value.code() == grpc.StatusCode.UNAUTHENTICATED
+ response = await stub.LookupPrefix(
+ request,
+ metadata=signed_metadata(request, auth),
+ )
+ assert response.hit_block_count == 1
+ wrong = distributed_pb2.LookupPrefixRequest(
+ compatibility=_compat().to_proto(),
+ block_hashes=hashes,
+ )
+ with pytest.raises(grpc.aio.AioRpcError) as error:
+ await stub.LookupPrefix(
+ wrong,
+ metadata=signed_metadata(wrong, auth),
+ )
+ assert error.value.code() == grpc.StatusCode.PERMISSION_DENIED
+ finally:
+ await server.stop(0)
+
+
@pytest.mark.asyncio
async def test_incompatible_and_dead_peers_are_misses():
store = PrefixCacheStore(_compat(), max_bytes=1024, node_id="peer")
@@ -224,5 +276,73 @@ def test_service_validation_and_dead_publish():
cache_address="peer:1",
chunk_bytes=0,
)
+ with pytest.raises(ValueError, match="max_payload_bytes"):
+ PrefillCacheServiceServicer(
+ store,
+ cache_address="peer:1",
+ max_payload_bytes=0,
+ )
block = CacheBlock.create(bytes(32), 1, b"x")
assert not publish_block_sync("127.0.0.1:1", _compat(), block, timeout_s=0.1)
+
+
+@pytest.mark.asyncio
+async def test_publish_rejects_duplicate_over_budget_and_incomplete_streams():
+ store = PrefixCacheStore(_compat(), max_bytes=1024, node_id="peer")
+ server = grpc.aio.server()
+ port = server.add_insecure_port("127.0.0.1:0")
+ address = f"127.0.0.1:{port}"
+ add_prefill_cache_service(
+ server,
+ store,
+ cache_address=address,
+ chunk_bytes=3,
+ max_payload_bytes=4,
+ )
+ await server.start()
+
+ async def call(chunks):
+ async with grpc.aio.insecure_channel(address) as channel:
+ stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel)
+ return await stub.PublishBlock(iter(chunks))
+
+ import hashlib
+ base = dict(
+ block_hash=bytes(32),
+ token_count=2,
+ total_chunks=2,
+ block_sha256=hashlib.sha256(b"abcdef").digest(),
+ compatibility=_compat().to_proto(),
+ )
+ try:
+ with pytest.raises(grpc.aio.AioRpcError) as exc:
+ await call([
+ distributed_pb2.PublishBlockRequest(
+ **base, chunk_index=0, data=b"abc",
+ ),
+ distributed_pb2.PublishBlockRequest(
+ **base, chunk_index=0, data=b"abc",
+ ),
+ ])
+ assert exc.value.code() == grpc.StatusCode.INVALID_ARGUMENT
+
+ with pytest.raises(grpc.aio.AioRpcError) as exc:
+ await call([
+ distributed_pb2.PublishBlockRequest(
+ **base, chunk_index=0, data=b"abc",
+ ),
+ distributed_pb2.PublishBlockRequest(
+ **base, chunk_index=1, data=b"def",
+ ),
+ ])
+ assert exc.value.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
+
+ with pytest.raises(grpc.aio.AioRpcError) as exc:
+ await call([
+ distributed_pb2.PublishBlockRequest(
+ **base, chunk_index=0, data=b"abc",
+ ),
+ ])
+ assert exc.value.code() == grpc.StatusCode.DATA_LOSS
+ finally:
+ await server.stop(0)
diff --git a/tests/inference_engine/distributed/test_prefill_compression.py b/tests/inference_engine/distributed/test_prefill_compression.py
new file mode 100644
index 00000000..027b8e92
--- /dev/null
+++ b/tests/inference_engine/distributed/test_prefill_compression.py
@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import pytest
+
+from inference_engine.distributed.capability import CompressionCodec
+from inference_engine.distributed.prefill_compression import (
+ compress_payload,
+ decompress_payload,
+ payload_sizes,
+)
+
+
+def test_zlib_round_trip_and_sizes():
+ raw = (b"prefill-kv-" * 1000)
+ compressed = compress_payload(raw, CompressionCodec.ZLIB)
+ assert len(compressed) < len(raw)
+ assert payload_sizes(compressed) == (len(compressed), len(raw))
+ assert decompress_payload(
+ compressed,
+ max_uncompressed_bytes=len(raw),
+ ) == raw
+
+
+def test_none_is_backward_compatible_raw_payload():
+ raw = b"legacy-kpkv"
+ assert compress_payload(raw, CompressionCodec.NONE) == raw
+ assert decompress_payload(raw, max_uncompressed_bytes=100) == raw
+ assert payload_sizes(raw) == (len(raw), len(raw))
+
+
+def test_compression_and_import_limits_validate():
+ with pytest.raises(ValueError, match="zlib level"):
+ compress_payload(b"x", CompressionCodec.ZLIB, level=10)
+ with pytest.raises(ValueError, match="unsupported compression"):
+ compress_payload(b"x", 99) # type: ignore[arg-type]
+ framed = compress_payload(b"x" * 100, CompressionCodec.ZLIB)
+ with pytest.raises(ValueError, match="import budget"):
+ decompress_payload(framed, max_uncompressed_bytes=10)
+ with pytest.raises(ValueError, match="import budget"):
+ decompress_payload(b"x" * 11, max_uncompressed_bytes=10)
+
+
+def test_truncated_and_corrupt_payloads_fail():
+ import hashlib
+ import struct
+
+ framed = compress_payload(b"x" * 100, CompressionCodec.ZLIB)
+ with pytest.raises(ValueError):
+ decompress_payload(framed[:10], max_uncompressed_bytes=1000)
+ damaged = framed[:-1] + bytes([framed[-1] ^ 1])
+ with pytest.raises((ValueError, __import__("zlib").error)):
+ decompress_payload(damaged, max_uncompressed_bytes=1000)
+ magic = b"KPC1"
+ with pytest.raises(ValueError, match="unsupported compression codec"):
+ decompress_payload(
+ struct.pack("<4sBQ32s", magic, 99, 1, hashlib.sha256(b"x").digest())
+ + b"x",
+ max_uncompressed_bytes=100,
+ )
+ with pytest.raises(ValueError, match="unsupported framed"):
+ decompress_payload(
+ struct.pack(
+ "<4sBQ32s",
+ magic,
+ int(CompressionCodec.NONE),
+ 1,
+ hashlib.sha256(b"x").digest(),
+ ) + b"x",
+ max_uncompressed_bytes=100,
+ )
+ with pytest.raises(ValueError, match="decompressed payload size"):
+ decompress_payload(
+ framed[:5] + struct.pack("= 5
+
+ def set(self):
+ self.calls = 5
+
+ final_store = PrefixCacheStore(COMPAT, max_bytes=4096, node_id="final")
+ final_jobs = PrefillJobStore(_Engine(), final_store)
+ final_job = PrefillJob(
+ "final", "final", "tenant", (1, 2), (b"a" * 32,),
+ CompressionCodec.NONE,
+ )
+ final_job.cancelled = SequenceEvent()
+ final_jobs._jobs[final_job.job_id] = final_job
+ final_jobs._run(final_job.job_id)
+ assert final_job.state == PrefillJobState.CANCELLED
+ final_jobs.close()
+
+
+@pytest.mark.asyncio
+async def test_worker_service_tenant_and_not_found_errors(worker):
+ address, _, _, _ = worker
+ wrong_tenant = _submit("wrong-tenant")
+ wrong_tenant.tenant_id = "other"
+ with pytest.raises(grpc.aio.AioRpcError) as exc:
+ await _rpc(address, "SubmitPrefillJob", wrong_tenant)
+ assert exc.value.code() == grpc.StatusCode.PERMISSION_DENIED
+ wrong_get = distributed_pb2.GetPrefillJobStatusRequest(
+ job_id="anything", tenant_id="other",
+ )
+ with pytest.raises(grpc.aio.AioRpcError) as exc:
+ await _rpc(address, "GetPrefillJobStatus", wrong_get)
+ assert exc.value.code() == grpc.StatusCode.PERMISSION_DENIED
+ wrong_cancel = distributed_pb2.CancelPrefillJobRequest(
+ job_id="anything", tenant_id="other",
+ )
+ with pytest.raises(grpc.aio.AioRpcError) as exc:
+ await _rpc(address, "CancelPrefillJob", wrong_cancel)
+ assert exc.value.code() == grpc.StatusCode.PERMISSION_DENIED
+ missing = distributed_pb2.GetPrefillJobStatusRequest(
+ job_id="missing",
+ tenant_id="tenant",
+ )
+ with pytest.raises(grpc.aio.AioRpcError) as exc:
+ await _rpc(address, "GetPrefillJobStatus", missing)
+ assert exc.value.code() == grpc.StatusCode.NOT_FOUND
+ cancel = distributed_pb2.CancelPrefillJobRequest(
+ job_id="missing",
+ tenant_id="tenant",
+ )
+ with pytest.raises(grpc.aio.AioRpcError) as exc:
+ await _rpc(address, "CancelPrefillJob", cancel)
+ assert exc.value.code() == grpc.StatusCode.NOT_FOUND
+
+
+@pytest.mark.asyncio
+async def test_sync_clients_work_without_auth():
+ cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="worker")
+ jobs = PrefillJobStore(_Engine(), cache)
+ server = grpc.aio.server()
+ add_prefill_worker_service(
+ server,
+ jobs,
+ node_id="worker",
+ cache_address="cache:1",
+ )
+ port = server.add_insecure_port("127.0.0.1:0")
+ await server.start()
+ address = f"127.0.0.1:{port}"
+ try:
+ submitted = await asyncio.to_thread(
+ submit_prefill_job_sync,
+ address,
+ _submit("no-auth"),
+ timeout_s=2,
+ )
+ request = distributed_pb2.GetPrefillJobStatusRequest(
+ job_id=submitted.job_id,
+ tenant_id="tenant",
+ )
+ status = await asyncio.to_thread(
+ get_prefill_job_sync,
+ address,
+ request,
+ timeout_s=2,
+ )
+ assert status.job_id == submitted.job_id
+ finally:
+ jobs.close()
+ await server.stop(0)
+
diff --git a/tests/inference_engine/network/test_network_state.py b/tests/inference_engine/network/test_network_state.py
index c92cc96d..418671eb 100644
--- a/tests/inference_engine/network/test_network_state.py
+++ b/tests/inference_engine/network/test_network_state.py
@@ -6,6 +6,7 @@
CapabilityRegistry,
NodeCapability,
NodeEndpoint,
+ PrefillWorkerCapability,
)
from inference_engine.distributed.prefill_cache import PrefixCacheStore
from inference_engine.network.state import NetworkState
@@ -27,6 +28,12 @@ def _state(tmp_path):
),
),
endpoints=(NodeEndpoint("head:2", "thunderbolt", 100, 0.4),),
+ prefill_workers=(PrefillWorkerCapability(
+ compatibility,
+ worker_address="head:3",
+ queued_tokens=128,
+ tokens_per_second_prefill=32,
+ ),),
)
return NetworkState(
CapabilityRegistry(self_card=card),
diff --git a/tests/integration/test_prefill_snapshot_mlx_equivalence.py b/tests/integration/test_prefill_snapshot_mlx_equivalence.py
new file mode 100644
index 00000000..d499731f
--- /dev/null
+++ b/tests/integration/test_prefill_snapshot_mlx_equivalence.py
@@ -0,0 +1,83 @@
+"""Real-MLX gate: imported prefill state must preserve continuation logits."""
+from __future__ import annotations
+
+import os
+
+import pytest
+
+
+def test_real_mlx_prefill_snapshot_preserves_continuation_logits():
+ pytest.importorskip("mlx.core")
+ torch = pytest.importorskip("torch")
+
+ from inference_engine.backends.mlx.verifier import MLXSinkWindowVerifier
+ from inference_engine.distributed.capability import (
+ CacheCompatibility,
+ CompressionCodec,
+ )
+ from inference_engine.distributed.prefill_cache import PrefixCacheStore
+ from inference_engine.distributed.prefill_cache_runtime import (
+ DistributedPrefillCacheHook,
+ )
+ from kv_cache_proposer.verifier import VerifierConfig
+
+ model_path = os.environ["KAKEYA_MAC_VERIFIER_PATH"]
+ verifier = MLXSinkWindowVerifier(VerifierConfig(
+ model_id=model_path,
+ dtype=torch.bfloat16,
+ device="cpu",
+ sink_size=4,
+ window_size=64,
+ ))
+ prompt = (
+ "Kakeya distributed prefill equivalence test. "
+ "The imported cache must produce exactly the same continuation logits. "
+ ) * 8
+ token_ids = verifier.tokenizer.encode(prompt)
+ compatibility = CacheCompatibility(
+ model_id=os.environ.get("KAKEYA_CACHE_MODEL_ID", model_path),
+ model_revision=os.environ.get("KAKEYA_MODEL_REVISION", ""),
+ tokenizer_revision=os.environ.get("KAKEYA_TOKENIZER_REVISION", ""),
+ cache_format_version="kakeya-prefill-v2-zlib",
+ quantization="4bit-mlx",
+ # This is a single-verifier round-trip gate; production head/workers
+ # derive and compare the real geometry hash separately. Keep the
+ # namespace stable without requiring another machine-local env var.
+ layer_geometry_hash=os.environ.get(
+ "KAKEYA_LAYER_GEOMETRY_HASH",
+ "integration-single-verifier",
+ ),
+ kv_dtype="bfloat16",
+ block_size_tokens=64,
+ tenant_namespace="integration",
+ )
+ store = PrefixCacheStore(
+ compatibility,
+ max_bytes=1 << 30,
+ node_id="integration-head",
+ )
+ hook = DistributedPrefillCacheHook(
+ store,
+ compression=CompressionCodec.ZLIB,
+ )
+ try:
+ assert hook.prepare(verifier, token_ids) == 0
+ baseline_logits = verifier.next_token_logits.clone()
+ baseline_argmax = int(torch.argmax(baseline_logits).item())
+ baseline_tokens = list(verifier.cached_token_sequence)
+
+ reused = hook.prepare(verifier, token_ids)
+ assert reused == len(token_ids)
+ assert int(torch.argmax(verifier.next_token_logits).item()) == baseline_argmax
+ assert torch.equal(verifier.next_token_logits, baseline_logits)
+ assert verifier.cached_token_sequence == baseline_tokens
+
+ # One real decode step must remain bit-identical after re-import.
+ next_token = baseline_argmax
+ imported_row = verifier.forward_block([next_token])[-1].clone()
+ hook.prepare(verifier, token_ids)
+ local_row = verifier.forward_block([next_token])[-1].clone()
+ assert torch.equal(imported_row, local_row)
+ finally:
+ hook.close()
+