When to hand-write a GPU kernel on Apple Silicon (and when the compiler already won)

Jun 30, 2026

I wrote five GPU kernels from scratch on a 16 GB M3 to learn how LLM inference actually works at the metal. The most useful thing I learned isn’t a kernel — it’s a decision rule:

• Elementwise ops: don’t bother. My hand-written fused SwiGLU kernel only matched mx.compile. The compiler already fuses elementwise work; you can’t beat it by hand.
• Reductions: hand-writing wins. My RMSNorm kernel beat mx.compile by 1.7–3.3×, because the compiler can’t fuse across the cross-thread communication a reduction needs.
• The famous trick is rarely the hard part. Online softmax — the keystone of FlashAttention — was the easy, correct part; the performance is all tiling and matmul hardware I haven’t built. My correct version is 13–33× slower than the expert kernel.

Everything below is measured on one passively-cooled M3 (16 GB). I write Metal, not CUDA, on purpose: it’s the hardware I own, and almost nobody is doing LLM kernels there — so the gaps are mine to find.

Rung 1 — SwiGLU: the compiler already won

SwiGLU (silu(gate) * up) is the MLP activation in Llama/Qwen/Gemma. The idiomatic form is two ops that round-trip an intermediate through memory; my kernel fused them into one pass. Against the naive form it looked like a 1.6× win. Then I benchmarked against mx.compile — and the win vanished: parity (0.98–1.17×). The compiler fuses elementwise ops for free. For memory-bound elementwise work the optimum is “one pass,” and the compiler already gets there.

Here’s the whole thing — MLX hands you the buffers and a thread index; you write the body:

kernel = mx.fast.metal_kernel(
    name="fused_swiglu",
    input_names=["gate", "up", "size"], output_names=["out"],
    source=SOURCE,                      # just the body; MLX wraps the [[kernel]] signature
)

uint i = thread_position_in_grid.x;
if (i >= size[0]) return;
float g = float(gate[i]);                        // upcast: do the activation in fp32
out[i] = T(g / (1.0f + exp(-g)) * float(up[i]));  // silu(gate) * up, one pass

(I also nearly published a 2.8× number — that was against gate*sigmoid(gate), a strawman; the real baseline nn.silu is faster. Always benchmark the strongest baseline.)

Rung 2 — RMSNorm: where hand-writing pays

RMSNorm has a reduction (mean of squares across the row), which needs threads to share partial sums — something elementwise fusion can’t express. So here the hand kernel decisively beat mx.compile: 1.7–3.3×, every trial. That’s the rule, proven both ways: elementwise → trust the compiler; reductions → reach for a kernel.

The reduction is the part the compiler can’t generate — threads have to share partial sums:

threadgroup float partial[256];                  // one threadgroup per row
float acc = 0.0f;
for (uint c = lane; c < d; c += tg) { float v = float(x[base + c]); acc += v * v; }
partial[lane] = acc;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint s = tg / 2; s > 0; s >>= 1) {          // tree-reduce in shared memory
    if (lane < s) partial[lane] += partial[lane + s];
    threadgroup_barrier(mem_flags::mem_threadgroup);
}
float inv = rsqrt(partial[0] / float(d) + eps[0]);

Rung 3 — the measurement was lying to me

Comparing my RMSNorm kernel to MLX’s expert mx.fast.rms_norm, the numbers were unrankable — the same op timed 3× apart across runs. The reflex is “I need GPU timestamps.” But the real bug was the experiment: I timed 50 runs of mine, then 50 of theirs, and a passively-cooled M3 throttles in the seconds between those blocks. A better clock wouldn’t help — a throttled GPU shows slow in the timestamp too. The fix was interleaving: alternate the two kernels call-by-call so each pair shares one thermal window, and report the ratio. That made them rankable: fp32 parity, fp16 ~15% behind — and a reason why.

Rung 4 — closing the fp16 gap

The fp16 gap was load efficiency: scalar half loads waste the memory bus. Switching to vectorized half4 loads plus a hardware simd_sum reduction closed it — fp16 reached parity with Apple’s expert kernel (within measurement noise, at the shapes I tested), a ~20–28% gain over my own scalar version. In the bandwidth-bound regime, load width is the lever.

float4 v = float4(x4[base4 + c]);  acc += dot(v, v);  // half4 loads: 4 elements/instruction
float s = simd_sum(acc);                              // reduce within a 32-lane SIMD group (hardware)
if (lane == 0) sg[wid] = s;                           // then combine 8 partials — threadgroup[8], not [256]

Rung 5 — online softmax, correct and honestly slow

Online softmax lets attention skip materializing the full score matrix by streaming over keys with a running-max rescale. The whole trick is six lines — rescale the running state whenever a bigger score shows up, so you never hold the full score matrix:

float s = dot(q_row, k_j) * scale;
float m_new = max(m, s);
float corr  = exp(m - m_new);                    // rescale prior state when the max grows
float p     = exp(s - m_new);
l = l * corr + p;
for (uint t = 0; t < d; t++) acc[t] = acc[t] * corr + p * float(v_j[t]);
m = m_new;                                        // out = acc / l  →  exact softmax(QKᵀ)V

I implemented it as a kernel; it’s numerically exact (matches both a reference and MLX’s scaled_dot_product_attention) — and 13–33× slower. That gap is the point: the softmax trick is famous, but FlashAttention’s speed is tiling and simdgroup-matrix matmul, none of which this does. I built the correct easy part and measured the distance to the hard part.

What I’m not claiming

These are textbook kernels — this is a learning journey, not research. “Parity with Apple’s kernel” means within measurement noise at the shapes I tested (hidden 4096, batch 8–16k), on wall-clock, not hardware counters. And I haven’t yet shown any of this moves real tokens/sec end-to-end in a model — that’s the next thing I owe this work, and I’ll report it honestly whether it’s a win or noise.

The takeaway

You don’t need a datacenter or CUDA to learn this. A laptop, MLX’s metal_kernel, and the discipline to benchmark the strongest baseline and name your noise floor will teach you where the real costs live in LLM inference.

All five kernels, the benchmark harness, and the honest-ranking timer are on GitHub: github.com/robertlangdonn/mlx-metal-kernels — run them on your own M3.