Unfusing a Kernel Made It 1.63x Faster

Porting a Metal prefill engine from a laptop chip to a workstation chip. The single biggest win was deleting a fusion — and two of our claims were wrong.

Summary

We ported an open-source Metal prefill engine, tuned for an M4 laptop chip, to an M3 Ultra workstation. Per-layer time went from 94.15 ms to 45.21 ms — 2.08×, taking the engine from 47% to 98.3% of the reference framework's throughput. The largest single win came from removing a fused kernel and dispatching two existing kernels instead. Two of the findings we published upstream turned out to be wrong and had to be publicly retracted, which is the part of this most worth reading.

The headline result

Fidelity improved at every step, incidentally — the original scalar kernels summed partial results in half precision, and the matrix units accumulate in float.

Finding 1: M4-era fusion is a net loss on a wide chip

The upstream engine fuses the gate and up projections of the SwiGLU MLP into a single kernel. That's textbook optimization: one memory pass instead of two.

We unfused it. No new kernel code — just dispatch the existing GEMM twice and then the existing activation.

1.63×.

The arithmetic is the whole lesson. The fusion saves one memory pass, worth about 0.4 ms. It costs 64 accumulator registers per thread, worth about 17 ms in lost occupancy.

On a narrow chip with fewer cores, memory passes dominate and the fusion is correct. On a wide chip with many cores, occupancy dominates and the same fusion is catastrophic. Nothing about the kernel changed — the hardware it runs on changed, and the optimization inverted.

When porting tuned kernels across chip widths, re-derive the fusion decisions rather than carrying them over. They encode a resource tradeoff specific to the original hardware, and they're some of the hardest choices to spot as assumptions because they look like straightforward wins.

Finding 2: a quantized KV cache buys nothing while you're occupancy-bound

Upstream, Q8_0 attention measured 17.27 ms against FP16's 17.24 ms. Halving the KV cache bytes changed nothing at all.

After the port: 1.354 ms vs 1.576 ms. Now it's a real 14% win.

The quantized cache was always going to help. It just couldn't, while the kernel was bottlenecked on something else entirely. Measured before the port, it looks like a dead feature.

Fix occupancy first, then re-evaluate every memory optimization you'd previously rejected. A bandwidth optimization measured on a compute-bound kernel reads as zero, and you'll file it as refuted.

Finding 3: threadgroup memory, not barriers, is the binding constraint

Across four independent changes, every variant that raised threadgroup memory lost and every variant that lowered it won — monotonically, no exceptions.

Doubling the K-block halves the barrier count and still loses. That's the cleanest evidence available: it isolates barrier overhead from threadgroup memory pressure and shows barriers aren't what matter here.

This also let us reject double-buffering without implementing it. Double-buffering trades threadgroup memory for latency hiding, and we had four data points saying that trade always loses on this chip.

A monotonic relationship across independent experiments is strong enough to prune the search space with. It's the difference between a tuning campaign and a search.

Negative results, so nobody repeats them: dropping the M-block to 64 on the fused MLP kernel is 27× slower (32 KB of epilogue staging leaves one threadgroup per core). Raising the K-block loses once the epilogue drains per row-block. Attention with a 32-wide C-block is 1.3% slower for 44% more threadgroup memory.

Two published claims we got wrong

We filed these findings upstream. A verification pass caught two errors, and we posted corrections publicly. This is the most useful section of this post.

Wrong claim 1: "694 GB/s bandwidth ceiling."

That number came from a read+write elementwise operation in the reference framework. The real ceilings are 776 GB/s read-only (95% of theoretical) and 695 GB/s read+write (85%).

The measurement was correct. It reproduces exactly. It was simply the wrong reference for weight streaming, which is read-dominated — so every "percent of ceiling" figure derived from it was understated.

A number can be reproducible, correct, and still the wrong reference for the claim you're making.

Wrong claim 2: "it spills registers."

We attributed the 27× slowdown on the small-M-block variant to register spilling. An occupancy probe shows maxTotalThreadsPerThreadgroup = 1024 for all sixteen kernels, including the slow one. No spilling anywhere.

The 32 KB of threadgroup memory — one threadgroup per core — explains the entire slowdown on its own.

The correct diagnosis was already in hand. We added a mechanism on top of it that sounded right, and that mechanism was fiction.

Don't attribute a slowdown to a mechanism you haven't measured. "Register spilling" is a plausible-sounding cause that requires an actual measurement to claim, and we had the tool to check — we just didn't run it before publishing.

Worth being explicit: the occupancy claims that survive in our writeup are inferred from throughput behavior, not verified with hardware counters. That's exactly the gap that produced error 2, and it's stated as a limitation rather than quietly omitted.

The methodological gap we couldn't close

GPU counters are unreachable on this chip. The device exposes only a timestamp counter set — no stage utilization, no statistics, no dispatch-boundary sampling. That needs full Xcode, and the workstation has Command Line Tools only.

So the entire diagnosis was achieved-throughput-versus-reference plus source reading. That's the biggest methodological weakness in the work, and it's the direct cause of one published error.

Get a roofline from a mature framework on the same machine before touching anything. Two minutes of running the reference framework's quantized matmul and attention primitives at our real shapes gave 24.8 TFLOP/s and the bandwidth numbers. That reference is what made "we're at 8% of ceiling" an actionable statement instead of a vague dissatisfaction.

Traps worth stealing

Probe language feature support by compiling snippets at runtime rather than guessing. We needed to know whether a particular half-precision accumulate-to-float operation was supported. Compile a five-line shader and find out, rather than reading spec documents.

Patch remote files with copied scripts, not inline SSH commands. Inline ssh host '...python...' breaks on quoting in ways that produce confusing failures far from the actual problem.

zsh does not word-split unquoted variables. This cost us one bogus benchmark row before we noticed — a flags variable that expanded as a single argument on one shell and multiple on another.

A shallow clone cannot be pushed. git clone --depth 1 in a setup step means an --unshallow fetch later, at exactly the wrong moment.

git add -A in a "docs only" commit swept a kernel change into a rebrand commit. The branch shipped correct code with a stale README, and the history now misattributes a real change to a cosmetic commit. Stage deliberately when a commit is meant to be docs-only.

The finding that keeps us honest about microbenchmarks

We repacked the quantized weight layout. The original stored weights so that adjacent thread-pairs read blocks 2,304 bytes apart, and each 18-byte block was read twice — once for the low nibbles and once for the high nibbles of the same bytes.

The isolated access-pattern microbenchmark showed 3.09× at matched thread counts: 784.9 GB/s against 254.3 GB/s.

The end-to-end gain was 2.0% per layer.

Not because the microbenchmark was wrong — the weight load overlaps compute rather than serializing, so fixing it recovers only the part that wasn't already hidden.

Don't quote microbenchmark ratios as promises. A 3× improvement on an isolated operation that's already 90% overlapped is a 2% improvement in the thing you care about. We shipped it — 2% is 2% — but a plan built on the 3.09× would have been wrong by two orders of magnitude.

Was it worth it?

Honestly: as an engineering exercise, yes. As a product, no, and it's worth saying so.

The engine is at 98.3% of the mature framework we already use in production. It only implements dense multi-head attention with a fixed head dimension. The model we actually serve has a different head dimension, grouped-query attention, a hybrid attention/state-space architecture, partial rotary embeddings, and multimodal inputs. This engine cannot run it, and closing that gap would be a rewrite whose ceiling is roughly break-even against what we already have.

The reusable output is the three findings at the top, and they transfer to any Apple Silicon kernel work.

Key takeaways

1. Re-derive fusion decisions when porting across chip widths. Ours saved 0.4 ms of memory traffic and cost 17 ms of occupancy on the wider chip.
2. Fix occupancy before evaluating memory optimizations. A quantized KV cache read as exactly zero benefit until the kernel stopped being compute-bound.
3. A monotonic relationship across independent experiments prunes the search space. Four changes all pointing at threadgroup memory let us reject double-buffering without building it.
4. Get a roofline from a mature framework on the same hardware first. Two minutes, and it converts vague dissatisfaction into a percentage.
5. A correct, reproducible measurement can still be the wrong reference. Ours was read+write bandwidth used as a weight-streaming ceiling.
6. Don't publish a mechanism you haven't measured. "Register spilling" sounded right, was checkable, and was false.
7. Microbenchmark ratios are not end-to-end promises. 3.09× isolated became 2.0% in situ, because the operation was already overlapped.