diff --git a/docs/MAC.md b/docs/MAC.md index d63b1f8..319d953 100644 --- a/docs/MAC.md +++ b/docs/MAC.md @@ -118,3 +118,66 @@ Set `MODEL_FILE` in `.env` to a different GGUF in `models/` and restart. The bun `templates/chat_template.jinja` is **Qwen3.6-specific**; for a non-Qwen GGUF set `CHAT_TEMPLATE_FILE` to that model's own Jinja template (a wrong template silently corrupts tool-calling). + +## Alternative backend: MLX (faster prefill) + +`scripts/mac-run.sh` can serve `:8080` with Apple's **mlx-lm** instead of llama.cpp, behind a +`BACKEND` knob. It's opt-in; the default stays llama.cpp. + +```bash +BACKEND=mlx bash scripts/mac-run.sh # or set BACKEND=mlx in .env +``` + +On an M3 Pro / 36 GB serving the 4-bit MLX model, prefill is ~530 tok/s vs llama.cpp's ~90 (**~6x**), +with decode at or above the MTP path (~39 vs ~37 tok/s). Prefill is what the compaction loop is bound +by, so this is the win that matters for long jobs. Greedy output matches llama.cpp (9/10 exact, the +lone diff a synonym). + +| Backend | Engine | Model | Prefill | Context ceiling | +| --- | --- | --- | --- | --- | +| `llamacpp` *(default)* | llama.cpp (Metal, GGUF) | `models/mtp/...Q4_K_XL.gguf` | ~90 tok/s | >128K (stable) | +| `mlx` | mlx-lm (`mlx_lm.server`) | `models/mlx/Qwen3.6-35B-A3B-UD-MLX-4bit` | ~530 tok/s | ~75K (auto-capped) | + +Use `llamacpp` for jobs that need >75K context; use `mlx` for throughput on the compaction-heavy +loop and jobs under that cap. + +### MLX setup (one-time) + +mlx-lm must live in its **own** venv - installing it into the app `.venv` bumps `transformers` and +breaks the embedder. + +```bash +uv venv .venv-mlx +VIRTUAL_ENV=$PWD/.venv-mlx uv pip install mlx-lm +# Download or convert a 4-bit MLX build into models/mlx/Qwen3.6-35B-A3B-UD-MLX-4bit +# (e.g. mlx_lm.convert, or pull a pre-quantized 4-bit MLX repo). +``` + +`mac-run.sh` checks both the venv and the model exist before starting and errors with the fix if not. + +### MLX env knobs + +| Var | Default | Meaning | +| --- | --- | --- | +| `BACKEND` | `llamacpp` | `llamacpp` or `mlx`. | +| `MLX_MODEL` | `models/mlx/Qwen3.6-35B-A3B-UD-MLX-4bit` | Path to the MLX model dir. | +| `MLX_CTX_CAP` | `75000` | Max context for MLX; `CTX_SIZE` is auto-capped to this. | +| `MODEL_ID` | (auto = `MLX_MODEL`) | Pinned so Pi's requests match the loaded model (see below). | + +### MLX caveats + +- **Context auto-cap.** Stock `mlx_lm.server` runs an **fp16 KV cache** (no `--kv-bits` flag - + upstream [ml-explore/mlx-lm#1043](https://github.com/ml-explore/mlx-lm/issues/1043)). fp16 KV + OOMs (`kIOGPUCommandBufferCallbackErrorOutOfMemory`) at ~78-92K actual tokens on 36 GB, so the + script caps context to `MLX_CTX_CAP` (75K). Quantized KV is greedy-lossless with 15/15 fact + recall at 83.5K locally; once the upstream server gains `--kv-bits`, switch it on and raise the + cap toward ~85K (kv4 ceiling ~92-113K). +- **`MODEL_ID` is pinned to the model path.** `mlx_lm.server` resolves the request's `model` field + against the loaded model and otherwise tries to fetch it from HuggingFace (a request for the + friendly label `qwen3.6` -> 404). llama.cpp ignores the label; MLX needs the match, so the script + exports `MODEL_ID=$MLX_MODEL` for this backend. +- **Dashboard tok/s + KV gauge are llama.cpp-only.** They read llama.cpp's `/metrics` and `/slots`, + which `mlx_lm.server` doesn't expose, so those live widgets stay blank under `mlx` (the job still + runs and packages normally). + +Stop the MLX model with `pkill -f mlx_lm.server`. diff --git a/scripts/mac-run.sh b/scripts/mac-run.sh index 4e55da9..ba596c5 100755 --- a/scripts/mac-run.sh +++ b/scripts/mac-run.sh @@ -7,6 +7,10 @@ # The Pi agent and the v5-nano embedder run inside the app process tree (embedder on CPU). # See docs/MAC.md for the full setup. This script only LAUNCHES; install steps live there. # +# BACKEND selects the :8080 server: 'llamacpp' (default, GGUF via llama.cpp) or 'mlx' (mlx-lm, +# Apple-native, ~6x faster prefill). mlx uses .venv-mlx + models/mlx/... and auto-caps context to +# MLX_CTX_CAP (fp16-KV OOM headroom). See docs/MAC.md. +# # It honours the same env knobs as the docker-compose path (MODEL_FILE, CTX_SIZE, SPEC_ARGS, # CHAT_TEMPLATE_FILE), just with Mac-appropriate defaults: # - SPEC_ARGS defaults to '--spec-type draft-mtp --spec-draft-n-max 2' (measured ~1.23x @@ -26,9 +30,18 @@ if [ -z "$JINA_API_KEY" ] || [ "$JINA_API_KEY" = "jina_xxxx" ]; then exit 1 fi -command -v llama-server >/dev/null || { echo "ERROR: llama-server not found. Install: brew install llama.cpp" >&2; exit 1; } +# Inference backend: llamacpp (GGUF via llama.cpp, default) | mlx (mlx-lm, Apple-native). +BACKEND="${BACKEND:-llamacpp}" + [ -x "$ROOT/.venv/bin/python" ] || { echo "ERROR: .venv missing. See docs/MAC.md (uv venv + uv pip install)." >&2; exit 1; } command -v pi >/dev/null || { echo "ERROR: pi not found. Install: npm install -g @earendil-works/pi-coding-agent@0.78.0" >&2; exit 1; } +case "$BACKEND" in + llamacpp) + command -v llama-server >/dev/null || { echo "ERROR: llama-server not found. Install: brew install llama.cpp" >&2; exit 1; } ;; + mlx) + [ -x "$ROOT/.venv-mlx/bin/mlx_lm.server" ] || { echo "ERROR: .venv-mlx is missing mlx-lm. Create it: uv venv .venv-mlx && VIRTUAL_ENV=\$PWD/.venv-mlx uv pip install mlx-lm (see docs/MAC.md)" >&2; exit 1; } ;; + *) echo "ERROR: BACKEND must be 'llamacpp' or 'mlx' (got '$BACKEND')" >&2; exit 1 ;; +esac MODEL_FILE="${MODEL_FILE:-mtp/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf}" CTX_SIZE="${CTX_SIZE:-65536}" @@ -36,7 +49,39 @@ NGL="${NGL:-999}" SPEC_ARGS="${SPEC_ARGS:---spec-type draft-mtp --spec-draft-n-max 2}" # ~1.23x on Metal; set SPEC_ARGS= to disable CHAT_TEMPLATE_FILE="${CHAT_TEMPLATE_FILE:-$ROOT/templates/chat_template.jinja}" MODEL_PATH="$ROOT/models/$MODEL_FILE" -[ -f "$MODEL_PATH" ] || { echo "ERROR: model not found: $MODEL_PATH (see docs/MAC.md to download the GGUF)" >&2; exit 1; } + +# --- MLX backend knobs (only used when BACKEND=mlx) --- +MLX_MODEL="${MLX_MODEL:-$ROOT/models/mlx/Qwen3.6-35B-A3B-UD-MLX-4bit}" +# KV cache quantization. Stock mlx_lm.server ran fp16 KV (OOM ~78-92K actual tokens on a 36GB +# Mac). The --kv-bits flag (ml-explore/mlx-lm#1353) pushes the ceiling to ~92-113K with 4-bit KV +# (greedy-lossless, long-ctx recall verified). We auto-detect the flag below: when present we +# default to 4-bit KV and the higher context cap; otherwise fp16 KV with the conservative cap. +MLX_KV_BITS="${MLX_KV_BITS:-4}" +MLX_KV_GROUP_SIZE="${MLX_KV_GROUP_SIZE:-64}" +MLX_KV_ARGS=() + +if [ "$BACKEND" = "mlx" ]; then + [ -d "$MLX_MODEL" ] || { echo "ERROR: MLX model not found: $MLX_MODEL (convert with mlx_lm.convert; see docs/MAC.md)" >&2; exit 1; } + # mlx_lm.server resolves the request's `model` field against the loaded model (and otherwise + # tries to fetch it from HuggingFace). llama.cpp ignores the label, but MLX needs it to match + # the loaded path, so pin Pi's MODEL_ID to it. (llama.cpp keeps the friendly default.) + export MODEL_ID="${MODEL_ID:-$MLX_MODEL}" + # Use quantized KV only if this mlx_lm.server build supports it; raise the cap accordingly. + if [ -n "$MLX_KV_BITS" ] && "$ROOT/.venv-mlx/bin/mlx_lm.server" --help 2>/dev/null | grep -q -- '--kv-bits'; then + MLX_KV_ARGS=(--kv-bits "$MLX_KV_BITS" --kv-group-size "$MLX_KV_GROUP_SIZE") + MLX_CTX_CAP="${MLX_CTX_CAP:-85000}" + echo "NOTE: BACKEND=mlx using ${MLX_KV_BITS}-bit KV (group $MLX_KV_GROUP_SIZE); ctx cap $MLX_CTX_CAP" + else + MLX_CTX_CAP="${MLX_CTX_CAP:-75000}" + [ -n "$MLX_KV_BITS" ] && echo "NOTE: this mlx_lm.server lacks --kv-bits (need mlx-lm#1353); fp16 KV, ctx cap $MLX_CTX_CAP" + fi + if [ "$CTX_SIZE" -gt "$MLX_CTX_CAP" ]; then + echo "NOTE: BACKEND=mlx caps CTX_SIZE $CTX_SIZE -> $MLX_CTX_CAP" + CTX_SIZE="$MLX_CTX_CAP" + fi +else + [ -f "$MODEL_PATH" ] || { echo "ERROR: model not found: $MODEL_PATH (see docs/MAC.md to download the GGUF)" >&2; exit 1; } +fi mkdir -p logs "${JOBS_DIR:-./data/jobs}" @@ -45,9 +90,30 @@ export PATH="$ROOT/.venv/bin:$(dirname "$(command -v pi)"):$PATH" export PI_BIN="$(command -v pi)" export PI_SKIP_VERSION_CHECK=1 -# --- 1. llama-server (Metal) -------------------------------------------------- +# Wait for the :8080 server to answer /health, or tail its log and bail. +wait_for_server() { + local label="$1" logf="$2" + echo -n "waiting for $label" + for i in $(seq 1 120); do + if curl -fsS "http://127.0.0.1:8080/health" >/dev/null 2>&1; then echo " ready"; return 0; fi + echo -n "."; sleep 2 + [ "$i" = 120 ] && { echo " TIMEOUT"; tail -30 "$logf"; exit 1; } + done +} + +# --- 1. inference server (:8080) ---------------------------------------------- if curl -fsS "http://127.0.0.1:8080/health" >/dev/null 2>&1; then - echo "llama-server already up on :8080" + echo "inference server already up on :8080" +elif [ "$BACKEND" = "mlx" ]; then + echo "=== starting mlx_lm.server (Metal) - loads ~19GB, first run ~30-60s ===" + nohup "$ROOT/.venv-mlx/bin/mlx_lm.server" \ + --model "$MLX_MODEL" \ + --host 127.0.0.1 --port 8080 \ + --max-tokens 8192 \ + ${MLX_KV_ARGS[@]+"${MLX_KV_ARGS[@]}"} \ + > "$ROOT/logs/mlx.log" 2>&1 & + echo "mlx_lm.server PID: $! (logs: logs/mlx.log)" + wait_for_server "mlx_lm.server" "$ROOT/logs/mlx.log" else echo "=== starting llama-server (Metal) - loads ~22GB, first run ~30-60s ===" # shellcheck disable=SC2086 @@ -67,12 +133,7 @@ else $SPEC_ARGS \ > "$ROOT/logs/llama.log" 2>&1 & echo "llama-server PID: $! (logs: logs/llama.log)" - echo -n "waiting for llama-server" - for i in $(seq 1 120); do - if curl -fsS "http://127.0.0.1:8080/health" >/dev/null 2>&1; then echo " ready"; break; fi - echo -n "."; sleep 2 - [ "$i" = 120 ] && { echo " TIMEOUT"; tail -30 "$ROOT/logs/llama.log"; exit 1; } - done + wait_for_server "llama-server" "$ROOT/logs/llama.log" fi # --- 2. FastAPI app ----------------------------------------------------------- @@ -85,6 +146,16 @@ export PORT="${PORT:-8000}" echo "=== starting Dataroom app on :$PORT ===" echo " web UI: http://localhost:$PORT/" -echo " LLAMA_URL: $LLAMA_URL ctx=$CTX_SIZE ngl=$NGL embedder=$EMBED_DEVICE" -[ -z "$SPEC_ARGS" ] && echo " spec: (disabled)" || echo " spec: $SPEC_ARGS" +echo " backend: $BACKEND" +if [ "$BACKEND" = "mlx" ]; then + echo " LLAMA_URL: $LLAMA_URL ctx=$CTX_SIZE (cap $MLX_CTX_CAP) model=$MLX_MODEL embedder=$EMBED_DEVICE" + if [ ${#MLX_KV_ARGS[@]} -gt 0 ]; then + echo " note: ${MLX_KV_BITS}-bit KV (group $MLX_KV_GROUP_SIZE, mlx-lm#1353); dashboard tok/s + KV gauge are llama.cpp-only" + else + echo " note: fp16 KV (this mlx_lm.server lacks --kv-bits, see mlx-lm#1353); dashboard tok/s + KV gauge are llama.cpp-only" + fi +else + echo " LLAMA_URL: $LLAMA_URL ctx=$CTX_SIZE ngl=$NGL embedder=$EMBED_DEVICE" + [ -z "$SPEC_ARGS" ] && echo " spec: (disabled)" || echo " spec: $SPEC_ARGS" +fi exec "$ROOT/.venv/bin/python" -m server.app