jacklarmer/laguna-s-2.1-mlx-optimized

Laguna S 2.1 — Apple Silicon MLX Optimization

Maximize Laguna S 2.1 (117B MoE,3B active) decode speed on Apple Silicon via pure-MLX, holding per-token logit correlation ≥0.95 vs the BF16 reference.

Results

Metric	Value
Decode speed	77.9 tok/s peak (77.6 avg) on M5 Max 128GB
Quality	0.9555 Pearson logit correlation vs BF16 reference (≥0.95 gate)
Baseline	65.2 tok/s stock mlx-lm, → +19% with our optimizations
Custom baseline	7.5 tok/s (naive Python MoE loop) → +937%

What's in this repo

• bench_speed.py — the optimized decode benchmark (one-big-function @mx.compile + zmlx sigmoid gating fix)
• test_quality.py — quality verification (Pearson logit correlation vs BF16 reference)
• test_generation.py — end-to-end generation smoke test
• kernels/fused_sdpa_decode.py — custom Metal SDPA kernel (correctness-verified, see kernels/test_fused_sdpa.py)
• laguna_pipenetwork.py — model registration script (copies laguna.py into mlx-lm's models directory)
• serve.py — OpenAI-compatible API server
• EXHAUSTION.md — full documentation of every optimization lever tried and why the remaining ones can't help
• probes/ — scratch/profiling scripts from the optimization process

Quick start

Prerequisites

• Apple Silicon Mac (M-series, unified memory). Tested on M5 Max 128GB. The 4-bit model needs ~66GB free RAM. Smaller-memory Macs may swap.
• Python 3.11+
• ~66GB disk for model download

Install

git clone https://github.com/jacklarmer/laguna-apple.git
cd laguna-apple
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Register the Laguna model architecture with mlx-lm:

python laguna_pipenetwork.py

Run the benchmark

python bench_speed.py

Expected output:

Fused SDPA kernel verified on decode path (kernels/test_fused_sdpa.py: 4/4 PASS)
Decode: 77.6 tok/s for 20 tokens

Verify quality

# First time: generate the BF16 reference (needs poolside/Laguna-S-2.1-NVFP4-mlx)
python make_reference.py

# Then: run quality test (threshold 0.95)
LAGUNA_QUALITY_THRESHOLD=0.95 python test_quality.py

Interactive generation

python run.py

OpenAI-compatible server

python serve.py
# POST to http://127.0.0.1:8090/v1/chat/completions

How it works

The speed comes from five stacked optimizations, all preserving the 0.955 quality gate:

1. One-big-function @mx.compile — all 48 layers compiled into a single MLX function. Model parameters are closure constants; KV cache arrays are explicit arguments. This eliminates ~96 Python→Metal boundary crossings per decode step.

2. zmlx moe_mlp patch — applies optimized MoE combine logic from ZMLX. Gives ~3% speed on the 47 MoE layers.

3. Sigmoid gating fix — the key quality breakthrough. ZMLX's default gating uses softmax, but Laguna uses sigmoid routing with e_score_correction_bias. We monkey-patch ZMLX's _gating function to use mx.sigmoid() instead of softmax. Without this fix, zmlx drops quality from 0.955 to 0.810.

4. Precomputed attention masks — masks for each decode offset are computed outside the compiled function and passed as arguments, eliminating ~8 mask dispatches per layer per step.

5. bfloat16 attention gating — the softplus gate (logaddexp(g_proj(x), 0)) is computed in bfloat16 instead of upcasting to float32, saving 2 dispatches per layer.

A custom Metal SDPA kernel (kernels/fused_sdpa_decode.py) was also written and correctness-verified (4/4 layer types pass), but stock MLX's scaled_dot_product_attention is faster because it uses optimized matrix-multiply units. The custom kernel uses scalar float loops and can't compete. See EXHAUSTION.md for the full analysis.

⚠️ LM Studio / Ollama compatibility

This model will NOT load in LM Studio or Ollama right now. Here's why:

LM Studio (since v0.3.4) and Ollama both support MLX engines, but neither supports trust_remote_code — they only load model architectures already built into their bundled mlx-lm/llama.cpp. Laguna is a custom architecture — the laguna.py model file is not in upstream mlx-lm (we install it manually via laguna_pipenetwork.py). LM Studio's mlx-engine does not execute custom repo code, so it cannot load Laguna.

What you can do instead (works today):

• Use the run.py / mlx-lm path in this repo (3-command setup, see Quick Start above)
• Use serve.py for an OpenAI-compatible API server on localhost:8090 (any client that speaks OpenAI API can connect to it)

Path to LM Studio support (future):

1. A PR to mlx-lm main merging the Laguna architecture code
2. Wait for an mlx-lm release that includes it
3. Wait for LM Studio to update its bundled mlx-lm

This is external to this repo — it depends on Apple's mlx-lm maintainers accepting the PR. We can't control that timeline.

GGUF alternative: A separate GGUF conversion (for llama.cpp / LM Studio / Ollama) is possible but would lose all our MLX-specific optimizations (zmlx, compile, the sigmoid gating fix). That's a different project — not in scope here.

Model

• Base model: poolside/Laguna-S-2.1 (BF16, 117B MoE)
• 4-bit MLX build used: pipenetwork/Laguna-S-2.1-MLX-4bit (affine 4-bit, group_size=64, ~66GB)
• BF16 reference model: poolside/Laguna-S-2.1-NVFP4-mlx (BF16 attention + NVFP4 experts, ~70GB)

We did not requantize or modify the model weights. All optimizations are runtime (Python/Metal code), not weight changes.

Hardware

Tested on Apple M5 Max with 128GB unified memory (~600GB/s peak bandwidth). The model is bandwidth-bound: ~2.5GB of expert weights are read per decode step (top-10 of 256 experts × 47 MoE layers) at ~400GB/s effective = ~6.25ms. The ~77.9 tok/s = 12.8ms/step is within 1.5ms of the raw bandwidth floor.

See EXHAUSTION.md for the full bottleneck analysis and why ~78 tok/s is the ceiling on this hardware holding the 0.95 quality gate.

License

The optimization code in this repo is MIT-licensed. The model weights belong to their respective creators (see model links above). The Laguna model architecture code (laguna.py registered via laguna_pipenetwork.py) is Copyright PipeNetwork, Apache 2.0.