Skip to content

Commit 9bbbc09

Browse files
cursoragentFluffyAIcode
cursoragent
and
FluffyAIcode
committed
feat(mac): scripts/run_kakeya_mac.sh — one-command local launcher for the FULL Kakeya engine
Wraps the full fused-engine chat (gemma-4 verifier + DFlash proposer + f_θ + S5 bounded KV) into a single command. f_θ runs by default; --fast = all-MLX proposer path (f_θ bypassed, faster). Resolves model/drafter/f_θ from the KAKEYA_MAC_* env vars with fallbacks; preflights MLX + model dir; --dry-run prints the command. Adds bridge preset mlx-kakeya-launcher-smoke to verify the launcher end-to-end on the Mac. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
1 parent 6911092 commit 9bbbc09

3 files changed

Lines changed: 136 additions & 0 deletions

File tree

‎inference_engine/bridge/manifest.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -749,6 +749,28 @@ def _harness_preset(
749749
},
750750
validate_reports=False,
751751
),
752+
Preset(
753+
name="mlx-kakeya-launcher-smoke",
754+
description="Verify the one-command local launcher "
755+
"scripts/run_kakeya_mac.sh runs the engine end-to-end on "
756+
"the Mac: invokes it in --fast scripted mode (all-MLX "
757+
"proposer path) with a fixed prompt and writes a "
758+
"transcript. Proves launcher → harness → engine wiring + "
759+
"env resolution + preflight on the real machine.",
760+
command_templates=(
761+
(
762+
"bash", "scripts/run_kakeya_mac.sh", "--fast",
763+
"--max-new-tokens", "{max_new_tokens}",
764+
"--chat-scripted",
765+
"What is the capital of France? Answer in one short sentence.",
766+
"--output",
767+
"results/research/k3_mac_bridge_launcher_smoke.json",
768+
),
769+
),
770+
timeout_minutes=45,
771+
params={"max_new_tokens": ("int:max_new_tokens", "64")},
772+
validate_reports=False,
773+
),
752774
)
753775
}
754776

‎scripts/run_kakeya_mac.sh‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
#!/usr/bin/env bash
2+
# Run the FULL Kakeya Inference Engine locally on a Mac (Apple Silicon).
3+
#
4+
# Launches an interactive chat on the complete engine:
5+
# gemma-4 verifier (MLX) + DFlash proposer (fused spec-decode)
6+
# + f_θ K/V restoration + S5 bounded KV.
7+
# f_θ runs by DEFAULT (the full verifier/proposer/f_θ pipeline). Use --fast for
8+
# the all-MLX proposer path (f_θ bypassed via S5 native prefill — much faster on
9+
# Mac, but the f_θ projection does not execute).
10+
#
11+
# Model facts come from env vars (set on the kakeya-mac-m4 runner), with sane
12+
# fallbacks; override on the CLI if needed:
13+
# KAKEYA_MAC_VERIFIER_PATH local MLX gemma-4 dir
14+
# KAKEYA_MAC_DRAFTER_ID DFlash drafter repo/dir
15+
# KAKEYA_MAC_FTHETA_DIR trained f_θ projection dir
16+
#
17+
# Usage:
18+
# bash scripts/run_kakeya_mac.sh # full engine (f_θ on), interactive
19+
# bash scripts/run_kakeya_mac.sh --fast # proposer-only (f_θ bypassed), faster
20+
# bash scripts/run_kakeya_mac.sh --max-new-tokens 2048 --window 128
21+
# bash scripts/run_kakeya_mac.sh --dry-run # print the command, run nothing
22+
# echo 'Explain proof-of-work.' | bash scripts/run_kakeya_mac.sh # one-shot via stdin
23+
set -euo pipefail
24+
25+
repo_root="$(cd "$(dirname "$0")/.." && pwd)"
26+
cd "$repo_root"
27+
28+
VERIFIER="${KAKEYA_MAC_VERIFIER_PATH:-$HOME/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit}"
29+
DRAFTER="${KAKEYA_MAC_DRAFTER_ID:-z-lab/gemma-4-26B-A4B-it-DFlash}"
30+
FTHETA="${KAKEYA_MAC_FTHETA_DIR:-results/research/f_theta_v5_s5_sliding}"
31+
SINK="${KAKEYA_SINK:-4}"
32+
WINDOW="${KAKEYA_WINDOW:-64}"
33+
BLOCK="${KAKEYA_BLOCK_SIZE:-4}"
34+
MAX_NEW="${KAKEYA_MAX_NEW_TOKENS:-1024}"
35+
36+
FAST=0
37+
DRY_RUN=0
38+
EXTRA=()
39+
while [[ $# -gt 0 ]]; do
40+
case "$1" in
41+
--fast) FAST=1 ;;
42+
--dry-run) DRY_RUN=1 ;;
43+
--verifier-path) shift; VERIFIER="${1:?}" ;;
44+
--drafter-id) shift; DRAFTER="${1:?}" ;;
45+
--f-theta-dir) shift; FTHETA="${1:?}" ;;
46+
--max-new-tokens) shift; MAX_NEW="${1:?}" ;;
47+
--window) shift; WINDOW="${1:?}" ;;
48+
--sink) shift; SINK="${1:?}" ;;
49+
--block-size) shift; BLOCK="${1:?}" ;;
50+
-h|--help) sed -n '2,28p' "$0"; exit 0 ;;
51+
*) EXTRA+=("$1") ;; # pass-through (e.g. --chat-scripted ...)
52+
esac
53+
shift
54+
done
55+
56+
log() { echo "[run-kakeya-mac] $*" >&2; }
57+
58+
# ---- argv for the full-engine harness chat ----
59+
args=(
60+
--verifier-path "$VERIFIER"
61+
--drafter-id "$DRAFTER"
62+
--f-theta-dir "$FTHETA"
63+
--s5-exact-full-attn --fused-specdecode
64+
--sink-size "$SINK" --window-size "$WINDOW" --block-size "$BLOCK"
65+
--max-new-tokens "$MAX_NEW" --chat
66+
)
67+
if [[ "$FAST" == "1" ]]; then
68+
# all-MLX proposer + bounded trim: faster, but f_θ is bypassed (S5 free lunch).
69+
args+=( --all-mlx-drafter --cuda-trim )
70+
MODE="FAST (verifier + proposer + S5 bounded KV; f_θ BYPASSED)"
71+
else
72+
# torch drafter + f_θ: the harness auto-enables --force-f-theta in --chat, so
73+
# f_θ projection ACTUALLY RUNS each turn (the full pipeline).
74+
MODE="FULL (verifier + proposer + f_θ + S5 bounded KV; f_θ runs)"
75+
fi
76+
77+
log "mode : $MODE"
78+
log "verifier: $VERIFIER"
79+
log "drafter : $DRAFTER"
80+
log "f_theta : $FTHETA"
81+
log "params : sink=$SINK window=$WINDOW block=$BLOCK max_new=$MAX_NEW"
82+
83+
cmd=( python3 scripts/research/k3_integrated_niah_eval_mac.py "${args[@]}" "${EXTRA[@]}" )
84+
85+
if [[ "$DRY_RUN" == "1" ]]; then
86+
echo "PYTHONPATH=.:sdks/python ${cmd[*]}"
87+
exit 0
88+
fi
89+
90+
# ---- preflight (Apple Silicon + MLX + model) ----
91+
command -v python3 >/dev/null || { log "python3 not found"; exit 1; }
92+
python3 -c "import mlx.core" 2>/dev/null \
93+
|| { log "MLX not importable — this needs Apple Silicon + 'pip install mlx mlx-lm'"; exit 2; }
94+
[[ -d "$VERIFIER" ]] \
95+
|| { log "verifier model dir not found: $VERIFIER (set KAKEYA_MAC_VERIFIER_PATH)"; exit 3; }
96+
if [[ "$FAST" != "1" && ! -e "$FTHETA" ]]; then
97+
log "f_θ dir not found: $FTHETA — set KAKEYA_MAC_FTHETA_DIR, or use --fast (f_θ bypassed)"
98+
exit 4
99+
fi
100+
101+
log "starting... (type a message, blank line / Ctrl-D to quit)"
102+
PYTHONPATH=".:sdks/python" exec "${cmd[@]}"

‎tests/inference_engine/bridge/test_manifest.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ def test_allowlist_contains_exactly_the_documented_presets():
8383
"mlx-kakeya-chat-smoke",
8484
"mlx-kakeya-fused-chat-ftheta",
8585
"mlx-kakeya-fused-chat-smoke",
86+
"mlx-kakeya-launcher-smoke",
8687
"mlx-multitenant-pressure",
8788
"mlx-upgrade",
8889
"mlx-upstream-batch-probe",
@@ -150,6 +151,17 @@ def test_mlx_kakeya_chat_smoke_preset_resolves():
150151
assert not [t for t in argv if t.startswith("{") and t.endswith("}")]
151152

152153

154+
def test_mlx_kakeya_launcher_smoke_preset_invokes_launcher():
155+
request = parse_manifest(_manifest(
156+
preset="mlx-kakeya-launcher-smoke", params={"max_new_tokens": "64"}))
157+
(argv,) = build_commands(request, {})
158+
assert argv[0] == "bash"
159+
assert argv[1].endswith("run_kakeya_mac.sh")
160+
assert "--fast" in argv
161+
assert "--chat-scripted" in argv
162+
assert argv[argv.index("--max-new-tokens") + 1] == "64"
163+
164+
153165
def test_mlx_kakeya_fused_chat_ftheta_preset_runs_f_theta_path():
154166
request = parse_manifest(_manifest(
155167
preset="mlx-kakeya-fused-chat-ftheta",

0 commit comments

Comments
 (0)