Optimize Metal prefill and decode on all Apple Silicon chips #555
Optimize Metal prefill and decode on all Apple Silicon chips #555ivanfioravanti wants to merge 3 commits into
Conversation
52e77a7 to
de0858c
Compare
|
Waiting for ds4-eval to complete before marking this ready for review. So far, so good. |
|
New: ds4-eval: 77/92 passed, 15 failed, runtime 04h:51m 100% correct. |
- Stage argsort top-k scores in threadgroup memory instead of gathered device loads (same comparison network, bit-identical permutation). - Cache the graph dump-prefix getenv and memoize flash-attn vec/reduce pipeline lookups on the per-token hot path. - Default prefill routed MoE to a fused Q4_K gate+up+SwiGLU grouped matmul sharing one staged B tile; generalizes the IQ2 fused kernel into a quant-generic template. Removes one GEMM dispatch and the F32 gate/up intermediate round-trip. All paths are bit-identical to previous logits, verified with frontier logit dumps (2k-32k ctx), local golden vectors, IQ2 fused-vs-unfused, and make test (tensor equivalence max_abs=0).
|
Great! Let me check :) |
|
Also thanks for submitting a PR that contains a human-readable description ;D |
|
M5 Max results (community datapoint) Thanks for this! Since the PR is titled "Pre-M5" and was benchmarked on an M3 Ultra, I ran it on an M5 Max (128 GB unified) to check whether the gains carry over to the newer chip — they largely do. Setup
Results (M5 Max)
So on M5 the low-context decode is essentially flat, but high-context decode (+7%) and prefill (+15%) both improve meaningfully — the prefill gain is even slightly larger than the M3 numbers reported here. Nice work. (Note: I measured throughput only and did not re-verify byte-identical logits; I'm trusting the deterministic-trace check in the PR description for correctness.) |
|
WOW @Flor1an-B this is a side effect! A very positive one! I’m working aside on M5 but was struggling a bit there, I’ll let agent point to this PR to get inspiration then! |
|
@ivanfioravanti re: "working aside on M5 but struggling a bit" — here's what we've measured on M5, in case it saves you time. M5 Max 128GB, macOS 27.0 (26A5378n), running ds4 with GLM-5.2 IQ2_XXS (197GB, SSD-streamed) and V4-Flash. Apple already ships a routed-MoE quantized GEMM, and it's MLX's
The metallib reflection gives away the design — shader family So it's a fused affine expert-loader on the same TensorOps matmul2d engine you'd target yourself on M5 — not new silicon. Its ceiling is TensorOps. That was the useful part for us: it told us where the roof is. It's affine-only, by hard assert. Hand it a LUT/codebook descriptor: So it can't take ds4's IQ2_XXS (codebook magnitudes + separate sign channel + per-K-block scales). MLX-style affine (q·scale + bias) only. Its affine scale blocks also run along N, not K — K-grouped layouts get Measured at ds4's GLM gate shape (6144→2048, top-8-of-256), affine-u8: Note the efficiency collapse as fan-out saturates all 256 experts. Caveat: the one gather config that encodes cleanly mis-computes on this build (the same descriptor is exact to 8.8e-8 on the non-gather QMM), so treat those as timings only. We killed it for our format. One more for your M5 work: Metal 27 adds Happy to share the probe sources if any of that's worth reproducing. |
M3 Ultra (Q4): +5 to 6% prefill, +4% decode, no regressionsI reproduced this on a second M3 Ultra and it holds up. Every measured context frontier improved, prefill and decode, with byte-identical KV cache sizes. Setup
Base vs base + this PR
Average across the sweep: prefill +5.8%, decode +4.1%. No frontier regressed. Output is unchanged. It composes cleanly. My daily driver is One honest note on magnitude. My sustained prefill gain (~+5 to 6%) lands a little below the +10.7% headline, though decode matches (+9.3% at 2K on my prod build). Likely a different base GGUF and single-run variance rather than a real disagreement. Nice work @ivanfioravanti . I would ❤️ to see this merged. |
|
@gilbert-barajas thanks for the hint! I'm throwing all possible context and hints to the agent, I'm finally starting to see positive results! I'll add this info too. Keep you posted. |
|
@ivanfioravanti glad it helped. If it's useful to your agent, I can gist the probe sources — a self-contained Obj-C tool (17 subcommands: the support matrix, numeric gates, weight-map extraction, and a perf ladder at real MoE shapes) that dumps the MPS quant runtime surface and drives the gather kernel directly on-box. Just say the word. Rooting for the M5 push. |
|
@antirez please wait on this one, I'm finalizing last tests on M5 speed-up, they should finish by tomorrow morning, if everything ok I'll merge changes in this branch so that we can land optimizations for the whole Apple Ecosystem 💪 |
Dense direct-RHS matmul (Q8_0/F16), attention-output low-rank, and indexer scores NAX kernels on Metal 4 / M5: - dense + attn-out: double-buffered weight tile (2->1 barrier per k-step), flattened staging with vector half4 stores, dead bounds checks removed under the host M%64/K%32 dispatch guarantee, and dequantize_q8_0_pairs producing identical halves from eight aligned 16-bit loads instead of sixteen byte loads. - indexer scores: double-buffered q tile (10->5 barriers per head), vectorized float4/half4 q/k staging, and head-packing (2 heads per matmul with a TQ=32 descriptor, halving mm.run count). Kernel-level: dense NAX sweep ~+3.6%, attn-out op ~+3%, indexer ~2.7x. End-to-end prefill ~+5% (7189-token chunked prompt) and up to ~+10% in cool same-binary A/B; decode unchanged (NAX is prefill-only). Logits verified bit-identical to the pre-change tree: full 129280-float logits after prefill and after each of 64 greedy decode steps, on both an unchunked 3840-token prompt (65 frames) and a chunked 7189-token prompt (33 frames, exercising the pos0>0 indexer path). Project tests: --metal-kernels, --local-golden-vectors, --logprob-vectors all OK.
ivanfioravanti
changed the title
|
Ready to go @antirez a very small bump ~1.5% on M5 using NAX, but it gets benefits from previous changes too. I updated the name of the PR. |
5 participants
Summary
This PR improves GGUF inference performance on the Apple M3 Metal backend while preserving the existing mathematical behavior and byte-identical logits.
The work was developed as a sequence of small experiments. Each optimization was retained only after:
Main changes
No model semantics, attention selection behavior, or approximation settings were changed.
Performance
Tested on an Apple M3 Ultra with 512 GB unified memory using the Q4 GGUF model,
promessi_sposi.txt, warmed model residency, and 128 generated tokens.Baseline:
80ebbc3Optimized branch:
52e77a7The measurements used order-balanced runs:
Correctness
A deterministic 64-token trace was compared directly with the original baseline. The complete logprob files were byte-identical:
6657ac9df2d53f9b41f329db323af89506894449277a3c8805cc06624c44975eThe focused Metal kernel suite also passes: