Faster prefill, tokenizer fix, draft-model speculation (1.3.0)

CHANGELOG.md

@@ -1,5 +1,19 @@
 # Changelog
 
+## 1.3.0
+
+Faster prefill, a tokenizer fix for Gemma, speculative decoding with a draft model, and cheaper 2-token verification.
+
+- **Prefill is about 1.45× faster.** The quantized prefill GEMM is one template over the per-simdgroup accumulator grid: the 32×32 tile stays for short batches, and a 64×64 tile (4×4 accumulators, two multiply-accumulates per fragment load instead of one, fully unrolled) runs long ones, reaching about 3.5 TFLOPS instead of 2.4. Both keep every output's arithmetic chain. Padded threadgroup rows and register prefetch measured slower on an M2 Pro and are not used. `flash_prefill` stages K/V as half (exact: the cache is half) with 16-byte loads of the next block during the current one's multiplies, rescales element-wise instead of through a diagonal matrix multiply, and schedules the longest query blocks first. Qwen3 0.6B prefills 512 tokens at 3,020 tok/s (was 2,040) and Mistral 7B at 247 (was 166). Decode is unchanged.
+- **Fixes wrong and unstable tokens on tokenizers with canonically equivalent entries.** Swift compares strings by canonical equivalence, but `tokenizer.json` keys are exact. Gemma lists both `;` and the Greek question mark U+037E, so they collided: which id a `;` got changed from run to run, and the losing id decoded to nothing. Vocabulary, merge, added-token and word-cache lookups now compare exact text.
+- **Draft-model speculative decoding** (`mlxr run --draft <model|auto>`, `mlxr serve --draft`, a "Draft model" picker in the app's inspector, `draft_tokens` / `accepted_draft_tokens` in the server's `usage`). A smaller model of the same family drafts tokens with the GPU sampler, queued back to back, and the target verifies them in one batched pass with the existing exact acceptance rule. Greedy output stays identical and sampled output keeps the target's distribution. The draft has its own KV cache, rolls back by position, and catches up on the prompt after the first token, so it does not delay the first token. Loading refuses a draft unless the tokenizers match, and the memory check sums both models. `auto` pairs the smallest installed model of the same type and tokenizer at most a quarter of the size. The app drops a saved pick that does not fit the model being loaded. **It is off by default and currently slower:** Qwen3 8B with a Qwen3 0.6B draft goes from 38.1 to 27.9 tok/s on a short prompt and from 35.6 to 21.3 at a 1,890-token context, because verifying 3–8 tokens costs about three decode steps and about 45% of drafts are accepted on the tested prompts.
+- **Verifying 2 tokens costs half of what it did** (Qwen3 8B: 78.6 → 40.7 ms; Qwen3 0.6B: 10.9 → 5.4 ms), through a 2-token matvec that reads each weight once. Batches of 3–8 tokens use a new simdgroup-matrix kernel (`qmvx4`) that dequantizes each weight once. On an M2 Pro it is ALU-bound at about 3× a single-token step (replacing the weight loads with constants costs the same), so 3–4 tokens are no faster than before, but the batched path now reaches 8. Batches of 2–8 tokens use the per-token decode attention instead of the prefill attention.
+- Keeps a chat message's Copy, Edit and Regenerate buttons visible while the pointer is on them, for a report of them hiding on hover and becoming unclickable.
+- `mlxr bench`: profile sweeps warm up first and take the best of three (the first sweep after loading read 2–3× slow), verification is reported at a short and a 1,900-token context, and `MLXR_BENCH_NOPROFILE=1` skips the profile.
+- Tests: GEMM tiles against the CPU reference up to 517 tokens, tiles agreeing bitwise across batch sizes, prefill vs incremental decode for a 600-token prompt, verification batches of 2–8 tokens vs one-at-a-time decode, exact speculative acceptance on synthetic logits, self-draft greedy identity, and tokenizer entries that are canonically equivalent.
+
+Tried in the same pass and left out: folding `qkv_post` into the decode attention, SiLU into the down projection and the residual add + RMSNorm into the next matvec (same output, but the step got slower on Qwen3 0.6B: 3.5 → 6.6 ms).
+
 ## 1.2.2
 
 - Fixes the chat window freezing during a long reply (any model) and then showing the rest at once, and the transcript going blank when a reply finished. Sampled on an M2 Pro, the main thread was 100% busy in SwiftUI/AppKit layout while the engine (which streams at 6–7 ms per token) had already finished the reply into the stream buffer. Five causes, all in the chat view:

README.md

@@ -103,8 +103,8 @@ Each checkpoint above passes end-to-end generation and the prefill-versus-decode
 - **Unsupported model types:** mixture-of-experts (Mixtral, Qwen3-MoE, Gemma 4 26B-A4B), Gemma 3n, and the vision and audio parts of multimodal models (Gemma 3 and Gemma 4 exports run as text-only). Discover marks unsupported types as not runnable.
 - **Gemma:** no tool calling and no Gemma 4 thinking mode (Gemma 4 runs with thinking off, its template's default). Sliding-window layers keep a ring of window + 512 K/V rows, so a chat whose shared prefix with the previous request ends more than 512 tokens before that request's end is processed from the start instead of reusing the cache.
 - **Chat templates:** the prompt format is detected (ChatML, Llama 3, Mistral, Phi-3, Gemma, Gemma 4) rather than executed as Jinja, so a model with an unusual template may get a near-match format.
-- **Prefill speed:** about 165 tok/s on a 7B model on an M2 Pro, because the prefill GEMM reaches about 35% of fp32 peak.
-- **Speculative decoding is opt-in.** Prompt-lookup drafting with exact verification is implemented and tested (greedy output is identical), but checking 4 drafted tokens on a 7B model still costs about 3 decode steps. So it only helps when most drafts are accepted (copying or editing text), and it is off by default (`SamplingParameters.speculative`, or `mlxr run --spec`).
+- **Prefill speed:** about 247 tok/s on a 7B model on an M2 Pro, against about 305 for mlx-lm on the same files. The 64×64 GEMM tile measures about 3.5 TFLOPS. Prompts shorter than about 128 tokens use the 32×32 tile and are further behind mlx-lm.
+- **Speculative decoding is opt-in and, for now, slower.** Prompt lookup (`mlxr run --spec`) and a draft model (`mlxr run --draft <model|auto>`, or the app's "Draft model" setting) both verify exactly, so greedy output is identical to a plain run. But verifying 3–8 drafted tokens costs about three decode steps on this GPU (the batched kernel is ALU-bound, see CHANGELOG 1.3.0), verifying 2 costs about 1.5, and Qwen3 0.6B's drafts for Qwen3 8B are accepted about 45% of the time on chat text: Qwen3 8B generates at 28 tok/s with the draft instead of 38. Both stay off by default until the verify kernel pays for them.
 - **Glass effects:** set `MLXRUNNER_NO_GLASS=1` to use plain materials instead of Liquid Glass.
 
 ## Performance
@@ -113,13 +113,29 @@ Measured with `mlxr bench` on an **M2 Pro (16 GB)**. Decode is measured with gre
 
 | Model (4-bit) | Weights | Prefill (512 tok) | Decode @16 | Decode @1K ctx |
 |---|---|---|---|---|
-| Qwen2.5 0.5B Instruct | 278 MB | 2,960 tok/s | 358 tok/s | 295 tok/s |
-| Qwen3 0.6B | 336 MB | 2,040 tok/s | 289 tok/s | 239 tok/s |
-| Llama 3.2 1B Instruct | 695 MB | 1,180 tok/s | 203 tok/s | 188 tok/s |
-| Phi-3.5 mini | 2.15 GB | 320 tok/s | 70 tok/s | 61 tok/s |
-| Mistral 7B v0.3 | 4.08 GB | 166 tok/s | 41.0 tok/s | 39.5 tok/s |
-| Gemma 3 1B (QAT) | 732 MB | 1,450 tok/s | 195 tok/s | 178 tok/s |
-| Gemma 3 4B | 2.56 GB | 345 tok/s | 67 tok/s | 62 tok/s |
+| Qwen2.5 0.5B Instruct | 278 MB | 4,270 tok/s | 343 tok/s | 278 tok/s |
+| Qwen3 0.6B | 336 MB | 3,020 tok/s | 286 tok/s | 234 tok/s |
+| Llama 3.2 1B Instruct | 695 MB | 1,720 tok/s | 202 tok/s | 185 tok/s |
+| Phi-3.5 mini | 2.15 GB | 466 tok/s | 70 tok/s | 61 tok/s |
+| Mistral 7B v0.3 | 4.08 GB | 247 tok/s | 41.0 tok/s | 39.3 tok/s |
+| Gemma 3 1B (QAT) | 732 MB | 2,070 tok/s | 189 tok/s | 173 tok/s |
+| Gemma 3 4B | 2.56 GB | 498 tok/s | 67 tok/s | 62 tok/s |
+
+**Compared with mlx-lm.** Same weight files on the same M2 Pro, mlx-lm 0.32.2 (mlx 0.32.4), one run each. Decode is 128 tokens at a short context for MLX Runner and `stream_generate` for mlx-lm, with greedy sampling. Prefill is a timed raw forward of 512 tokens from an empty cache.
+
+| Model (4-bit) | Decode: MLX Runner | mlx-lm | Prefill 512: MLX Runner | mlx-lm |
+|---|---|---|---|---|
+| Qwen3 0.6B | 286 tok/s | 260 | 3,020 tok/s | 2,340 |
+| Qwen2.5 0.5B | 343 | 296 | 4,270 | 4,030 |
+| Llama 3.2 1B | 202 | 205 | 1,720 | 1,720 |
+| Gemma 3 1B | 189 | 181 | 2,070 | 2,010 |
+| Gemma 2 2B | 94 | 113 | 782 | 791 |
+| Gemma 1.1 2B | 78 | 81 | 763 | 868 |
+| Gemma 4 E2B | 73 | 76 | 770 | 639 |
+| Phi-3.5 mini | 70 | 66 | 466 | 572 |
+| Mistral 7B | 41.0 | 41.2 | 247 | 305 |
+
+MLX Runner decodes faster on the smallest Qwen models and Phi-3.5, is close on most others, and is well behind on Gemma 2. It prefills 512 tokens faster on Qwen3 and Gemma 4 and is behind on Gemma 1.1, Phi-3.5 and Mistral 7B. At 128 tokens, where the 32×32 tile runs, mlx-lm prefills 4–36% faster on every model except Gemma 4, which is level.
 
 **The 7B ceiling.** Decoding one token reads every weight once, so speed is capped by memory bandwidth. This M2 Pro measures 197 GB/s with a plain read kernel, which caps a 4.08 GB model at about 48 tok/s. MLX Runner reaches 41.0 tok/s, about 85% of that cap. Reaching 50+ tok/s on a 4-bit 7B model therefore needs one of: a Mac with more bandwidth (M-series Max chips have about 400 GB/s), a 3-bit model (which would raise the ceiling to about 60 tok/s, but only 4-bit has an optimized decode kernel so far), or speculative decoding (see Limitations).
 
@@ -148,7 +164,8 @@ mlxr download mlx-community/Qwen3-4B-4bit   # download into the app's library
 mlxr list                                   # installed models
 mlxr run mlx-community/Qwen3-4B-4bit -p "Explain RoPE" --max 300 [--temp 0.7] [--no-think]
 mlxr serve mlx-community/Qwen3-4B-4bit --port 1234     # OpenAI-compatible server (--lan prints an API key)
-mlxr bench mlx-community/Qwen3-4B-4bit      # prefill/decode throughput + kernel micro-benchmarks
+mlxr run mlx-community/Qwen3-8B-4bit --draft mlx-community/Qwen3-0.6B-4bit   # speculative decoding (or --draft auto)
+mlxr bench mlx-community/Qwen3-4B-4bit      # prefill/decode throughput, verification cost at short and 1,900-token contexts, kernel profile (MLXR_BENCH_NOPROFILE=1 skips it)
 mlxr tokenize <model> "some text"
 ```