One guard, 110× off the verify step. Speculative decoding on MLA models in MLX.

I run large language models on a Mac Studio and spend a lot of time trying to make them faster. Last week I worked out why speculative decoding was a net loss on every DeepSeek-lineage model I tried in mlx-lm. It came down to one comparison operator that had been copied into eight files across three codebases.

The mlx-lm change is open as ml-explore/mlx-lm#1817 and the mlx-vlm change as Blaizzy/mlx-vlm#2141. This is the longer version with all the measurements, since a PR description is not the place for them.

The defect

Multi-head Latent Attention (MLA) keeps its key-value cache as a compressed latent, not as full keys and values. At attention time there are two mathematically equivalent ways to use that cache:

• Absorbed. Fold the up-projection into the query and attend directly against the compressed latent. Cost scales with the number of new queries.
• Materialised. Expand the entire cached latent back into full K and V, then attend normally. Cost scales with the cached length.

mlx-lm picks the absorbed path with if L == 1:, where L is the number of query positions in this forward call. Plain autoregressive decode is always L == 1, so it gets the fast path. Speculative decoding is never L == 1: a verify step scores the draft tokens plus one bonus token in one call, so L is 2 or more. So every verify step re-expanded the whole cache, and a step that is supposed to be cheap grew linearly with context instead.

The same guard, character for character, sits in seven mlx-lm model files, in mlx-vlm’s glm5_next, and in the vendored copy of that file inside oMLX. I found the oMLX one telling, because that file had been hand-optimised around custom sparse-MLA Metal kernels and the L == 1 boundary still survived the rewrite.

The threshold is not a constant

Because the two costs scale on different axes, the crossover between them is a fixed number of queries that does not move with context length. Per head, the absorbed path does roughly 2·L·Kv·r work on attention plus L·r·(n+v) on the fold. The materialised path does Kv·r·(n+v) to expand and L·Kv·(n+v) to attend. Here r is the latent rank, n the no-position-encoding head dimension, v the value head dimension and Kv the cached length. Set them equal, drop the fold term because the cache is much longer than either r or L, and Kv cancels:

max_absorbed_queries = r · (n + v) / (2r − n − v)

I derived this myself and then went looking to see whether anyone else had. Simon Dong published the identical formula in Understanding MLA Decoding (Materialize vs Absorb) in July 2026, with “about 171 tokens” for DeepSeek-V3. The integer floor gives 170.

Plugging in real checkpoint configs:

geometry (r / n / v)	threshold	models
512 / 128 / 128	170	DeepSeek-V3, Moonlight-16B, Kimi-Linear, Ling-3.0-flash, LongCat
512 / 192 / 256	398	GLM-4.7-Flash
512 / 256 / 256	512	GLM-5.3-Flash

A hardcoded constant would be wrong for one family or the other, which is why the patch computes it from the config at construction time rather than picking a number.

An empirical sweep on GLM-4.7-Flash’s attention layer, cache primed to 8192 tokens, agrees with the formula. Absorbed wins at every width up to and including 398, and materialised wins from 512 up. So the formula lands on the right side of the real crossover, with a small margin.

Each branch forced on GLM-4.7-Flash's first attention layer. The lines cross between 398 and 512. Prefill chunks of 512 and 2048 sit to the right, where materialised is still the cheaper path.

The change

The patch is small. mlx_lm/models/mla.py gets one helper:

def max_absorbed_queries(kv_lora_rank, qk_nope_head_dim, v_head_dim) -> int:
    denom = 2 * kv_lora_rank - qk_nope_head_dim - v_head_dim
    if denom <= 0:
        return 1
    return max(1, int(kv_lora_rank * (qk_nope_head_dim + v_head_dim) / denom))

Each MLA model computes self._max_absorbed in its constructor, after the three head dimensions exist, and its two guards become if L <= self._max_absorbed:. In mlx-lm that is seven model files plus the helper, +56/−21 lines. In mlx-vlm it is two files, +20/−3.

I left two L == 1 guards alone on purpose. deepseek_v32 has a third one that gathers the top-k sparse indices for a single query, and that branch is there for correctness rather than speed, since the wider case already builds a mask. rwkv7 has one in its token-shift code that has nothing to do with attention.

What it buys

Everything below was measured on an M3 Max with 128 GB. Both arms are real clones of mlx-lm at the same upstream commit, selected by sys.path, so there is no monkey-patching involved. The number that matters for a verify step is the marginal cost of the second query, which I take as whole-model time at L=2 minus L=1, median of nine runs.

The harness and the patch were written with an AI coding agent driving the terminal, with me directing and checking, and every number here was run by me.

Log scale. Every patched dot sits in the same 1.7 to 3.3 ms band regardless of cache length.

Stock’s marginal grows with the cache, while patched stays between 1.7 and 3.3 ms at every length I tried. That flat line is the whole point of the change.

Here is the same thing from the whole model’s point of view, as cost per decode step at cache 8192 while the step widens from one query to four:

L=1 is identical in both arms. The entire penalty is the step from one query to two. Past that, stock's marginals are only 2 to 3 ms because it is already materialising.

Plain decode at L=1 takes the same branch in both arms and is unchanged, within 2 % on every model. I actually use that as a validity check. If the L=1 timings differ between arms, something other than the patch is different and I throw the run away. That caught one genuinely bad run, which I describe further down.

End to end

Marginal attention cost is not the same thing as throughput, so to see whether the fix changes what a runtime actually does, I ran MTPLX’s tune command, which benchmarks plain decode against each speculative depth and picks a winner, on a pack forged from the original zai-org/GLM-4.7-Flash checkpoint with its multi-token-prediction head. The only difference between arms is the two guards.

arm	context	AR tok/s	MTP depth 1 tok/s	ratio	verify ms	acceptance	verdict
stock	~700 tokens	60.7	47.6	0.78×	41.6	—	rejected
patched	~700 tokens	61.1	87.0	1.42×	20.1	90.3 %	—
stock	2K–16K tokens	47.8	7.1	0.15×	413.9	72.1 %	rejected
patched	2K–16K tokens	47.7	64.4	1.35×	27.7	73.8 %	selected

The long-context row is the one to look at. Verify cost falls 14.9× and throughput rises 9.1×, while acceptance barely moves. The draft head was producing good drafts the whole time, and it was the cost of checking them that made speculation a loss.

Plain decode is 61 and 48 tok/s in the two rows, so the patched arm beats it in both and the stock arm beats it in neither.

Every draft depth moves the same way. This is one run, both arms, depths one to three, with plain decode as the control at 62.6 and 62.9 tok/s:

Without the change there is no depth at which speculation is worth enabling on this model. With it, every depth beats plain decode. Depth 1 wins because this model's position-2 acceptance is only about 11 percent.

The drop-off with context is even clearer in MTPLX’s depth sweep (greedy, 256 tokens generated, same prompts in both arms). Ratio of speculative decode throughput to plain decode:

Greedy, 256 tokens generated per case, same prompts in both arms. Above 1.0 speculation wins. This harness uses a diagnostic history policy with lower acceptance than tune, so read the shape, not the absolute level.

At 16K, turning on speculative decoding without the patch made the model about 25× slower than leaving it off. The two MTPLX harnesses report different absolute acceptance because they use different history policies, so I quote tune for the headline numbers and use the sweep only for the shape of the curve. They agree on the direction.

The sparse model is different

GLM-5.3-Flash has sparse attention. An indexer selects the top 2048 cached positions and attention only runs over those. Below that budget the MLA fix gives 3.4× on the verify step, but above it the work is already capped and there is nothing left for the absorbed path to save. At 2048, 8192 and 32768 tokens of cache, patched is within 2 % of stock, so there is no regression and no gain either.

That means the attention fix stops helping this model at exactly the context lengths where speculation would be most useful. The remaining cost turned out to be in a different place.

The same bug, one layer over

The indexer that picks those 2048 positions pools the cached keys into windows of four before scoring them. Pooling the whole cache every step is O(T), so the mlx-vlm implementation caches completed windows and recomputes only the tail. That path was gated on S == 1, where S is the number of new tokens in the step. It is the same mistake as the attention guard, with an optimisation that is valid for any small S only switched on for one value, so every speculative verify fell back to pooling the entire cache.

Widening it is safe because windows start at multiples of four and tokens only append, so a window completed before the tail cannot change. The fix drops the S == 1 condition and keeps the guards that do the real work: the cache has a pool, there is no padding, and the batch shape matches. A review pass added a fourth one, that the cache is exactly as long as the pool remembers. That covers speculative rejection, where the cache gets trimmed and a stale pool would otherwise carry the rejected tokens into the next step.

S is the number of new tokens in the step. S=1 is unchanged at about 0.7 ms at 8K and 1.0 ms at 128K in both arms. Across the model's eleven sparse layers, a verify step at 128K spends about 63 ms in the indexer before the change and about 12 ms after.

A wide step now pools identically to the same tokens fed one at a time, checked for two, three, four and eight new tokens at aligned and misaligned cache lengths. Window indices, cache length and the chosen top-k match exactly. The pooled keys agree to a few float32 ULP because the two paths sum the same windows in a different order.

Neither fix does anything for drafting itself. GLM-5.3-Flash’s MTP head does work. It is stored under a prefix that mlx-vlm’s sanitize() drops on load, which is why no public runtime uses it. Read from the raw safetensors, with a properly aligned cache for the draft head, conditional acceptance is 0.84 at position one and decays gently: 0.73, 0.63, 0.62. Combined with the measured verify costs, the best depth is 2 to 3 for about 1.6×. Getting a runtime to actually do that is a separate job.

Things I have to disclose

Outputs of the attention change are not bit-identical. Forcing the absorbed and materialised paths on identical inputs through one 4-bit GLM-4.7-Flash attention layer gives a maximum absolute difference of 4.883e-04 on values near 0.1. bf16 spacing at 0.1 is 2⁻¹¹, or 4.88e-04, so the two paths differ by exactly one bf16 ULP. The delta is the same at L=2 and L=8, so it comes from quantised matmuls running in a different order rather than from anything that grows with query count. Over hundreds of greedy tokens that is enough for trajectories to diverge, and it showed up as MTP acceptance differing between arms on otherwise identical runs. vLLM and SGLang accept the same property for the same reason.

Three of the seven patched modules are unmeasured. kimi_k3, longcat_flash and bailing_moe_v3 share the 170 geometry but every checkpoint I could find is 255 GB or larger, or has no mlx-lm module at all. The argument carries by construction, but I want to be clear that for those three it is an argument and not a measurement.

One run was thrown out. A patched 16K sampled run reported 118 tok/s, three times plain decode at the same context. It had drafted 12 of 256 tokens, so the number was an artefact of a run that barely speculated. The greedy re-run gave a sane 31 tok/s and that is what appears above. A separate oMLX run was invalidated by the L=1 check: the second arm inherited a saturated MLX buffer cache on a 103 GB model and every point sat near 1550 ms. On a machine that size you have to interleave the arms sample by sample and check that L=1 matches, or the second arm quietly measures the allocator instead of the code.

Prefill is essentially untouched. Prefill runs in chunks of 512 or 2048, both above every computed threshold. The only change is that a short final partial chunk can now take the absorbed path, and the sweep above shows that is faster, not slower.

If you maintain an MLA runtime, it is worth grepping for L == 1 near your embed_q. It was in all three of the codebases I looked at.