test(distributed): add run_distributed_demo.sh + mlx-distributed-spec-decode-demo bridge preset

Two-process (proposer + Qwen3-0.6B verifier) distributed spec-decode over real
gRPC sockets, asserting byte-identical-to-greedy. Lets the Mac bridge validate
the ADR 0009 distributed engine on-device (mirrors the GPU-host run).

Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

Author: cursoragent

inference_engine/bridge/manifest.py

@@ -102,6 +102,25 @@ def _harness_preset(
 PRESETS: Dict[str, Preset] = {
     p.name: p
     for p in (
+        Preset(
+            name="mlx-distributed-spec-decode-demo",
+            description="ADR 0009 distributed spec-decode, on-device: two local "
+                        "processes (n-gram ProposerService + Qwen3-0.6B verifier) "
+                        "over real gRPC sockets — capability gossip, placement, "
+                        "remote drafts + LOCAL greedy verify, asserting the output "
+                        "is byte-identical to local greedy. Validates the "
+                        "distributed engine on the Mac.",
+            command_templates=(
+                (
+                    "bash", "scripts/run_distributed_demo.sh",
+                    "--verifier-id", "Qwen/Qwen3-0.6B",
+                    "--max-new-tokens", "{max_new_tokens}",
+                ),
+            ),
+            timeout_minutes=45,
+            params={"max_new_tokens": ("int:max_new_tokens", "48")},
+            validate_reports=False,
+        ),
         Preset(
             name="mlx-env-probe",
             description="Probe Metal/MLX + mlx.distributed availability.",

scripts/run_distributed_demo.sh

@@ -0,0 +1,59 @@
+#!/usr/bin/env bash
+# Run the ADR 0009 distributed speculative-decode demo as TWO local processes
+# (proposer node + verifier node over real gRPC sockets) and assert the
+# distributed output is byte-identical to local greedy. Used for on-device
+# validation on the Mac bridge (and runnable anywhere with the deps).
+#
+# Picks an mlx_lm-free but torch+transformers+grpcio-capable Python: prefers
+# KAKEYA_MAC_PYTHON, then the repo venv, then python3.
+#
+# Usage:
+#   bash scripts/run_distributed_demo.sh [--verifier-id ID] [--max-new-tokens N]
+set -euo pipefail
+
+repo_root="$(cd "$(dirname "$0")/.." && pwd)"
+cd "$repo_root"
+
+VERIFIER_ID="Qwen/Qwen3-0.6B"
+MAXNEW="48"
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --verifier-id) shift; VERIFIER_ID="${1:?}" ;;
+    --max-new-tokens) shift; MAXNEW="${1:?}" ;;
+    *) echo "[dist-demo] ignoring arg: $1" >&2 ;;
+  esac
+  shift
+done
+
+_can() { [ -n "${1:-}" ] && "$1" -c 'import grpc, torch, transformers' >/dev/null 2>&1; }
+PYBIN=""
+for c in "${KAKEYA_MAC_PYTHON:-}" "$repo_root/.venv-mac/bin/python3.13" \
+         "$repo_root/.venv-mac/bin/python" "$HOME/kakeya-venv/bin/python" \
+         "$(command -v python3 2>/dev/null || true)"; do
+  if _can "$c"; then PYBIN="$c"; break; fi
+done
+if [[ -z "$PYBIN" ]]; then
+  echo "[dist-demo] no Python with grpc+torch+transformers found; set KAKEYA_MAC_PYTHON" >&2
+  exit 2
+fi
+echo "[dist-demo] python=$PYBIN verifier=$VERIFIER_ID max_new=$MAXNEW" >&2
+
+export PYTHONPATH="$repo_root:$repo_root/sdks/python"
+export HF_HUB_DISABLE_PROGRESS_BARS=1
+
+# Clean any stale demo procs, start the proposer node, run the verifier node.
+for p in $(pgrep -f demo_distributed_spec_decode 2>/dev/null || true); do kill "$p" 2>/dev/null || true; done
+sleep 1
+"$PYBIN" scripts/demo_distributed_spec_decode.py \
+    --role proposer-node --bind 127.0.0.1:50061 --node-id node-b \
+    > /tmp/kakeya_dist_proposer.log 2>&1 &
+PP=$!
+trap 'kill "$PP" 2>/dev/null || true' EXIT
+sleep 6
+"$PYBIN" scripts/demo_distributed_spec_decode.py \
+    --role verifier-node --bind 127.0.0.1:50060 --node-id node-a \
+    --peer 127.0.0.1:50061 --verifier-id "$VERIFIER_ID" \
+    --max-new-tokens "$MAXNEW"
+RC=$?
+echo "[dist-demo] VERIFIER_RC=$RC" >&2
+exit $RC

tests/inference_engine/bridge/test_manifest.py

@@ -79,6 +79,7 @@ def test_allowlist_contains_exactly_the_documented_presets():
         "mlx-batched-manual-sdpa",
         "mlx-batched-multitenant",
         "mlx-batched-pad-decode",
+        "mlx-distributed-spec-decode-demo",
         "mlx-env-probe",
         "mlx-kakeya-chat-smoke",
         "mlx-kakeya-chat-stream-probe",