Thirty tiny kernels per layer: a 16% decode win for GLM-5.3-Flash on Apple Silicon

Carmen, my assistant, runs on GLM-5.3-Flash on a Mac Studio M3 Ultra with 512 GB of memory. The build is a mixed quantisation I published yesterday (frederic-ai/GLM-5.3-Flash-MLX-mixed-4_6bit): experts at 4-bit, everything else at 6-bit. It decodes at about 30 tokens per second on a single request.

I wanted to know whether there was any speed left in the runtime itself, without touching the model or the weights. There was, a little. This post is about where it came from, what it cost, and why the same change is worth less than half as much in the upstream project.

Where the time goes

Generating one token on Apple Silicon means streaming every active weight through the GPU once. On this build that is roughly 12 GB per token: 5.5 GB for the eight routed experts the token selects, and the rest for attention, the shared experts and the embeddings. The M3 Ultra moves about 800 GB/s, so if the GPU did nothing but read weights, the ceiling would be about 64 tokens per second.

We see 30. Half the time, the memory bus is idle.

Some of that gap is the model's structure. GLM-5.3-Flash has 45 layers. Eleven use sparse attention with a lightning indexer; 34 use Kimi Delta Attention, a form of linear attention with a small recurrent state instead of a growing key-value cache. A decode step through one of those 34 layers looks like this:

• three input projections (q, k, v)
• a short causal convolution over the last four tokens, then SiLU
• an L2 normalisation of q and of k, with scaling
• a forget-gate projection and a "safe gate" transform
• a sigmoid for the beta gate
• the delta-rule update of the recurrent state
• a gated RMS norm on the output
• the output projection

In the eager path, each of those is its own GPU kernel launch, and several of them are two or three (a cast to float32, the operation, a cast back). Call it 25 to 30 launches per layer, per token. Each one does microseconds of arithmetic. The launch itself is not free, and the GPU idles between them. Across 34 layers that adds up to a few milliseconds of a ~33 ms token.

Why mx.compile fits here and not elsewhere

MLX's mx.compile traces a function once, fuses chains of elementwise operations into single kernels, and replays the fused graph on later calls. The condition is that the input shapes stay fixed; a new shape means a new trace.

That is why nobody compiles the decode step of a normal attention layer: the key-value cache grows by one token every step, so every step would be a retrace. Linear attention is different. The state is a fixed [heads, d, d] matrix plus a three-token convolution window. Every decode step has exactly the same shapes as the last one. Compile once, reuse forever.

The other obstacle is that model code updates the cache in place, and a compiled function has to be pure. So I wrote the step as a function that takes the states in and returns them:

def _decode_step_pure(self, x, conv_state, state):
    """(x[1,1,D], conv_state, state) -> (out[1,1,D], new_conv_state, new_state)"""
    mixed = mx.concatenate([self.q_proj(x), self.k_proj(x), self.v_proj(x)], axis=-1)
    conv_input = mx.concatenate([conv_state, mixed], axis=1)
    new_conv = mx.contiguous(conv_input[:, -(self.conv_kernel_size - 1):, :])
    q, k, v = self._split_heads(nn.silu(self.conv1d(conv_input)))
    q = self._l2norm(q) * self.head_dim ** -0.5
    k = self._l2norm(k)
    out, new_state = gated_delta_update(q, k, v, self._gate(x), self._beta(x),
                                        self.A_log, self.dt_bias, state=state)
    out = self.o_norm(out, self._out_gate(x))
    return self.o_proj(out), new_conv, new_state


def _fused_decode(self, x, cache):
    step = self._compiled or mx.compile(self._decode_step_pure)
    out, cache[0], cache[1] = step(x, cache[0], cache[1])
    return out

It is the same operations in the same order as the eager path, mirrored into a pure function. The dispatch at the top of the layer sends single-token, single-request decode through it and leaves prefill and batches on the eager path.

One small thing made the measurement honest: the switch is a file in /tmp, checked on every call. That meant I could flip it between requests on the production server, with no restart, and interleave off/on for every prompt so both settings saw the same background load.

What it bought

Four prompts, four rounds each, off and on alternating per request, with the other model on the box idle:

An earlier sweep on the previous 4/8-bit build showed the gain holds as context grows: +4.2% at a short prompt, +4.0% at 29k tokens, +3.4% at 100k. At the layer level it is 0.673 ms → 0.595 ms per KDA layer, about 2.7 ms saved per token.

Five percent. Not a headline, but Carmen produces a lot of tokens, and this cost nothing at runtime.

The part that looked like it wasn't free

I expected the outputs to be bit-identical. Same operations, same order, same weights. They were not. At temperature 0 the text matched for somewhere between 30 and 140 tokens, then forked at a near-tie between two candidate tokens and carried on, equally well but differently. The recurrent state stayed exact over 32 steps; the layer outputs differed in the last bit.

"Different but fine" is a property you cannot write a test for, so I chased it. I compiled each op of the chain on its own, a million random values each, and compared bits with eager. exp, log, tanh, rsqrt, division, multiply-add, where, the fp32 norm chain: all identical. Exactly one op was not: mx.sigmoid. Inside a compiled kernel its exp is lowered to Metal's fast approximation, while the prebuilt sigmoid kernel uses the precise one; 17% of fp32 sigmoid outputs differ by one ulp. mx.exp itself stays precise under compilation, so spelling out the prebuilt kernel's formula, y = 1/(1+exp(|x|)) and then y or 1-y by sign, reproduces it bit for bit. There are three sigmoids per layer. Replaced.

That made the upstream version exact. My own runtime still differed, and the second culprit was a constant. My fork computes the q/k norms by hand and then multiplies by head_dim ** -0.5, which is 0.08838834764831845. MLX prints scalar constants captured by a compiled function into the kernel source with 7 significant digits, one short of what a float32 needs to round-trip, so the compiled kernel multiplied by 0.08838835 instead. Constants like 0.3 or 1e-6 survive 7 digits and hide the bug; 1/3, 1/√2 and π do not. Passing the scale into the compiled function as an input rather than capturing it fixes it, and the bug report is filed upstream as ml-explore/mlx#4503.

With both changes the fused step is bit-identical to eager: outputs, conv window and recurrent state, over 64 consecutive steps, and on the real model every generated token matches. Exactness cost nothing. The upstream port still measures +2%, and my runtime still measures 1.13× per layer.

Taking it upstream, and finding out why it's worth less there

Carmen's server runs a patched fork of mlx-vlm's GLM implementation. To be useful to anyone else, the change had to go into mlx-vlm itself, and a single-request-only fast path would rightly be rejected by a project whose server batches requests by default.

So the upstream version handles batches: a left-padded batch passes its mask into the compiled graph, a ragged continuous batch passes per-row lengths, and each variant gets its own compiled function, traced once per batch size. It switches itself off during speculative decoding, where the cache records intermediate states the pure function cannot produce. It comes with a unit test that checks the compiled path actually ran. The project's existing GLM, speculative-decoding and cache suites (1,067 tests) pass with the flag off and on.

Then I measured it on the real model, and the gain was 2%, not 5%. Every generated token was identical to the eager path.

The reason is a rewrite that had landed five days earlier. Pull request #2127, merged on September 9, rebuilt the GLM-5 layer on shared components: one fused projection instead of three, mx.fast.rms_norm instead of a hand-written norm with casts around it, a shared delta-rule kernel. That removed roughly half the small launches my compile was saving. My fork predates it. So on my runtime the compile recovers 5%, and on main there is only 2% left to recover.

I also found I was not the first. PR #2105, by another contributor, had folded the entire post-projection chain into one hand-written Metal kernel per layer: bit-exact, and +7.8% end-to-end. It was closed on September 5 because it would have conflicted with #2127's shared kernels, and the maintainer asked for an issue with a spec instead.

That is what I filed: issue #2254 with the measurements, the two rounding traps and the design options, and PR #2255 with the opt-in, bit-exact implementation, offering to close it if folding the work into the shared kernel is the preferred route. Two percent is a small number next to a 7.8% kernel that was already turned down; what the PR has going for it is that it is exact, ninety lines, off by default, and comes with the data.

The kernel that was already written

The contributor behind PR #2105, avlp12, had written something better than my compiled step: a single hand-written Metal kernel that fuses the entire post-projection chain of the KDA layer — convolution, norms, gates, delta-rule update, output norm. Bit-exact against eager, and +7.8% end-to-end on their setup. It was closed not because it was wrong but because it collided with the #2127 refactor's shared kernels; the maintainer asked for a spec issue instead of a conflicting implementation. That kernel, plus PipeNetwork's follow-up port in glm53-flash-mlx#2, already existed.

My runtime is a fork of the pre-#2127 layout — exactly the layout the kernel was written for. So I ported it unchanged, wired it behind a flag with the same file-based kill switch as the compiled step, and measured at both the layer level and the server level.

Layer test first: 64 decode tokens, random 4-bit weights at real shapes, bf16 activations. Zero output elements and zero state elements differ from eager. Per layer, per decode step: 0.662 ms eager, 0.581 ms compiled, 0.509 ms kernel.

Then the production server, six runs per prompt — two eager, two compiled, two kernel, interleaved so all three settings saw the same background load:

Stacked on the compiled path the kernel adds another 5%, and the two together take the runtime from the 29.0 tok/s this journey started at to 33.6 — +16% decode in total. Since decode is roughly 70% of an assistant turn (the rest is prefill and tool overhead), that is about +11% on a full turn. I enabled it in production at 02:11 by creating one file, with no restart.

Raw data: prodkernel_live.json, stored next to this draft.

How much a verify step really costs

The obvious next lever is speculative decoding, and whether it pays depends on one number nobody publishes for MoE models: what a cached forward over a multi-token block costs, relative to a one-token decode step. I measured it on the real model at both 1.5k and 11.7k context, with no memory spike:

Each extra token costs about 0.36 of a step. The arithmetic that follows is unforgiving: a verify step over S tokens pays off only if the expected number of accepted tokens exceeds S × 0.36. At a 50–60% draft acceptance rate, a 5-token block needs roughly 2 accepted tokens to break even — and typical acceptance gains fall short of that. A 5-token block is a net loss; blocks of 2–3 are the only viable operating point on this MoE, and they cap the speedup at something modest.

There is also a blocker: DFlash2 currently cannot run at all on mlx-vlm main for this model, because the #2127 refactor left the speculative-cache hook unimplemented. So speculative decoding on GLM-5.3-Flash stays unmeasured — the numbers above say it would need small blocks even once it runs.

So why hadn't MLX done this?

Because it is only easy in one place, and only recently. It works on linear-attention layers because their state has a fixed shape; models built that way have existed for about a year. It needs a duplicate, pure copy of the layer's forward pass, which maintainers avoid because every future fix has to land twice. It needs care with batching or it retraces constantly. Until you find the two rounding traps above, it changes greedy outputs, which makes regressions harder to spot. And the expensive part, the delta-rule update, was already a custom kernel; what was left was glue. Five percent of glue, per model, is easy to deprioritise in a project where correctness of the port comes first.

What's left

The attention chain is done. The remaining lever is the batch-size-1 efficiency of the expert matmuls themselves, so I microbenchedmarked them at the real expert shapes (288 experts, top-8 routing, 4-bit): at batch 1, MLX's gather_qmm moves the 113 MB of expert weights per layer at 258 GB/s in isolation — against a 514 GB/s plain read of the same bytes, and 509 GB/s once the batch reaches 4. That gap is the rest of the distance to the bandwidth ceiling.

A custom expert kernel might recover +10–15% end-to-end, at the price of weeks of Metal work. The first move is an upstream report, not a kernel: the batch-1 case deserves a baseline better than 258 GB/s before anyone writes Metal for it.

What it cost

One thing went wrong, and it deserves its own paragraph. While testing, I loaded a second copy of the model next to production, and macOS's jetsam killed it — the event was vm-compressor-space-shortage: production GLM had grown to 188 GB, the test copy added 172 GB, gemma sat at 60 GB. Production restarted in 4 seconds and nothing was lost beyond a RAM cache, but the lesson is now a rule on this box: never load a second big model next to a serving process. Two copies do not share weight pages; "it fit last time" is not a memory plan.

The bigger levers for Carmen are now elsewhere anyway: generating fewer tokens per turn (compact tool-call JSON alone is −23% on those turns), a finer quantisation mix, and routing routine turns to a smaller model. None of those is a miracle either. They stack.