feat(distributed): orchestrate same-model prefill workers and KV memory tier

.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()

.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" \

.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: |

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