dflash-mlx-bonsai2

DFlash 2 speculative decoding for PrismML's Ternary-Bonsai-2-27B: the 2-bit 27B model, faster on code and math, same greedy output. On a Mac, an NVIDIA GPU or a browser tab.

Apple Silicon 24 GB, NVIDIA 24 GB, or Chrome with WebGPU. Runs on your own hardware. No account, no telemetry.

Install

Step-by-step for every row, with downloads, flags and a chat UI for the Mac server: docs/USAGE.md. The rest of this README is the Mac path. On a Mac, the script makes a venv, installs this repo, and downloads the model pack (8.6 GB) and the drafter (3.85 GB) into the Hugging Face cache; it is idempotent and --dry-run shows the plan. Then start the server and point any OpenAI-compatible client at it:

bash scripts/serve-bonsai2.sh                      # dflash serve on http://127.0.0.1:8790/v1
curl http://127.0.0.1:8790/v1/chat/completions -H "Content-Type: application/json" \
  -d '{"model":"Ternary-Bonsai-2-27B-mlx-2bit","messages":[{"role":"user","content":"Write a Python LRU cache."}],"temperature":0,"max_tokens":512,"stream":true}'

Use temperature: 0: speculation only engages on greedy requests. No config file, no account. Try it without installing anything by reading the measured tables in docs/BONSAI2.md.

Why

You have a 27B model packed to 2 bits, and on a Mac it answers at 21 tokens a second because every token re-reads all of its weights. Speculative decoding fixes that in principle: a small drafter guesses a block of tokens and the big model checks them in one pass. In practice, on the 2-bit Bonsai 2 target, two things were broken, and this fork fixes both.

The drafter (z-lab's DFlash 2) was trained on bf16 Qwen3.8 and guesses what the bf16 model would say, not what the ternary one does. And MLX's 2-bit matmul is tuned for one row, so checking 8 tokens cost as much as decoding 8 tokens one by one. This fork of bstnxbt/dflash-mlx adds a loader for the Hadamard-rotated pack, a Metal 8-row 2-bit GEMM for the check, and a drafter re-fitted on the ternary model's own output.

Use something else if you want the base Qwen3.8-27B at 4 bits, where upstream dflash-mlx already works and has more memory headroom; you have a CUDA box, where PrismML's llama.cpp fork with this same drafter is the right path (docs/USAGE.md; our Metal port of it is in lab/leg9/patches, and it is not fast); or your traffic is chat and thinking, where this fork is break-even and plain mlx-lm is simpler.

What it does to the numbers

Apple M4 Pro, 24 GB, quiet machine, thinking off, greedy. Plain decode through the same pack: ~21.5 tok/s.

The same drafter on one NVIDIA L4 through llama.cpp is faster still, measured by bonsai2-run: 2.2× on GSM8K, MBPP and MATH-500 (31 → 68 tok/s). On code edits, llama.cpp's prompt lookup (--spec-type ngram-mod) runs first and reuses text from your prompt, and the drafter covers the rest: 3.15× together, against 2.46× for the drafter alone and 1.57× for prompt lookup alone. The MLX server has its own prompt lookup from upstream (--copyspec-mode, on by default); its gain on code edits on a Mac is not measured yet.

Output equals plain greedy decoding up to fp16 ties (the 8-row and 1-row kernels accumulate in different orders; on 1000 real tokens the v7 path made 0 argmax flips against the fp32 reference). The gains and the caveats behind each row are in docs/BONSAI2.md.

What this fork adds

The pack loader (dflash_mlx/runtime/prism_pack.py) builds a stock mlx_lm Qwen3.5 text model and installs the pack's Hadamard + 2-bit modules into it, so upstream's hidden capture, tape-replay rollback and prefix cache work unchanged (parity with the pack's own loader: max logit delta 3e-5). The verify kernel (dflash_mlx/runtime/prism_qmm.py, DFLASH_PRISM_VERIFY=v7) is a simdgroup 8×8 MMA GEMM that dequantizes 2-bit weights straight into matrix-tile registers, with a fused sign, Hadamard, transpose and row-sum prep; 2–8-row verify calls are padded to 8 because the real DFlash 2 verify is 4–5 rows wide. It cut the 8-token check from 438 ms to 105–115 ms against a 46 ms decode step.

The drafter (naklitechie/Qwen3.8-27B-DFlash2-ternary-bonsai2) is z-lab's DFlash 2 fine-tuned on 1.5 M tokens of the ternary model's own greedy generations, captured with a llama.cpp hidden-state dumper and trained on a rented GPU (lab/aws/). Against the shipped drafter it accepts more on every prompt measured (drafter A/B, its own runs): chat 2.46 to 2.68, code 3.07 to 3.28, raw code 4.41 to 4.57 tokens per cycle. DFLASH_TOOL_PARSER=off lets clients that parse <tool_call> text themselves stream it through the server.

The browser port (lab/webgpu/) runs the same loop inside the WebGPU engine that LocalMind uses for Bonsai 2: the drafter in WGSL with Q4_K/Q6_K weights kept packed on the GPU, an M ≤ 8 ternary GEMM for the verify block, and a recurrence-only rewind after a rollback (the target is a Gated-DeltaNet hybrid). runner/spec-runner.js is the class LocalMind vendors; in-app on the same M4 Pro it is 1.18× on a 1,486-token code answer (34.4 vs 40.7 ms/token), output identical.

Commands

bash scripts/setup-bonsai2.sh [--dry-run]                       # venv + install + both model downloads
bash scripts/serve-bonsai2.sh                                   # OpenAI-compatible server on :8790
DFLASH_PRISM_VERIFY=v7 dflash generate  --model <pack-dir> --draft <drafter-dir> --prompt "..."
DFLASH_PRISM_VERIFY=v7 dflash benchmark --model <pack-dir> --draft <drafter-dir> --only-dflash --block-tokens 8 --verify-mode dflash --max-tokens 512 --no-eos --prompt "..."
DFLASH_PRISM_VERIFY=v4b|fp16|off ...                            # older verify paths, for A/B
DFLASH_PRISM_VERIFY_STATS=1 ...                                 # rows-per-call histogram at exit
DFLASH_TOOL_PARSER=off dflash serve ...                         # stream tool-call text unparsed

--model must be the pack's local directory (a bare repo id falls through to mlx_lm.load, which cannot build it); dflash benchmark needs --only-dflash for the same reason.

Verify it yourself

python lab/padded_parity.py      # pack from $BONSAI2_PACK or the HF cache; 3–8-row verify path vs the stock fp32 path: argmax identical, fp16-level deltas
python lab/tie_flip_audit.py     # 1000 real tokens: argmax flips per verify path vs the fp32 reference
python lab/verify_timing_modes.py v4b v7   # isolated verify(8) vs decode(1), per kernel

The parity and tie-flip audits are what "greedy-identical up to fp16 ties" points at; the timing script is what the 105–115 ms figure points at. Verified end to end on one M4 Pro: dflash serve with this drafter driving a tool-calling chat client at 70–73 % acceptance. Not verified: other Apple chips, other memory sizes.

License

Apache-2.0, as upstream; see NOTICE for z-lab, PrismML and Qwen attribution. Founding document and every measurement: docs/BONSAI2.md · Mac, NVIDIA and browser setup: docs/USAGE.md · runtime flags: docs/runtime-flags.md · research trail: lab/.