Skip to content

feat(mac): add BACKEND={llamacpp|mlx} knob for native MLX inference #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
soobrosa wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(mac): add BACKEND={llamacpp|mlx} knob for native MLX inference #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c22fcd4
feat(mac): add BACKEND={llamacpp|mlx} knob to mac-run.sh
soobrosa Jun 5, 2026
23e71bb
mac: auto-enable 4-bit KV on mlx backend when --kv-bits is supported
soobrosa Jun 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 63 additions & 0 deletions docs/MAC.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
95 changes: 83 additions & 12 deletions scripts/mac-run.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -26,17 +30,58 @@ 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}"
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}"

Expand All @@ -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
Expand All @@ -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 -----------------------------------------------------------
Expand All @@ -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