fix: Speed up multi-token decode by computing the absorbed MLA threshold

[stooit]

Speculative decoding on MLA models is currently slower than plain decoding: the absorbed path is gated on L == 1, so every verify step expands the entire cached latent instead. The cost grows with cache length, so it gets worse the longer you generate. Any step with L >= 2 pays it -- e.g. mlx-lm's draft-model speculation (and any multi-token forward pass).

The crossover between the two paths is a query count set by head geometry and the cache length: L < r*d / (r*d/T + 2r - d), with d = qk_nope_head_dim + v_head_dim. For a long cache that tends to r*d / (2r - d), 170 for DeepSeek V3 and 398 for GLM-4.7-Flash; on a cold cache, where T == L, it reverses and materialising always wins. This PR adds max_absorbed_queries and latent_length to models/mla.py so each model computes it from its own config and the current cache length. deepseek_v32's third guard, the sparse top-k gather, is deliberately left at L == 1 -- that branch is about correctness, not speed.

Single-token decode is within 2% across all models, because L=1 takes the same branch either way.

Example: on GLM-4.7-Flash with an 8K prompt, mlx_lm.generate --draft-model runs at 2.8 tok/s against 55.8 tok/s autoregressive. This change takes that to 14.7 tok/s (5.2×). Whether speculation then wins depends on draft cost (e.g. with a cheap draft head it does, with a same-architecture draft model it still loses).

Prior art: SGLang exposes --speculative-attention-mode {prefill,decode} to pick which path target-verify uses (server_args.py, deepseek_v2.py); vLLM tracks the same problem class in 21984; the formula is from Simon Dong's write-up.

Results

M3 Max 128GB. Stock and patched are two checkouts of the same commit selected via sys.path, so nothing but the two guards differs -- no runtime patching.

Verify-step cost (ms), stock → patched

Model	Module	Threshold	cache 512	cache 2048	cache 8192
GLM-4.7-Flash	glm4_moe_lite	398	29.4 → 2.4	93.8 → 2.9	363.7 → 3.3
Moonlight-16B-A3B	deepseek_v3	170	6.6 → 2.3	21.7 → 2.4	83.0 → 2.5
Kimi-Linear-48B	kimi_linear	170	4.4 → 1.9	12.8 → 1.7	45.9 → 2.0

Marginal cost of the second token in a decode step -- what a verify pays. Stock grows with the cache, patched is flat. Kimi-Linear is per-step only: mlx-lm refuses speculative decoding on MLA hybrids (ValueError: Speculative decoding requires a trimmable prompt cache (got {'ArraysCache'})), since linear-attention state cannot be rewound on rejection.

Where the paths cross

One GLM-4.7-Flash layer at cache 8192 with each branch forced:

L	absorbed	materialised
1	0.5	7.3
398	16.0	16.9
512	19.9	18.6
2048	89.8	59.6

The empirical crossover is between 398 and 512 and the formula computes 379 at this cache length (398 in the limit), so it lands on the correct side with a small margin. Prefill steps are 512 or 2048, where materialised correctly wins -- prefill is unaffected.

Downstream: MTPLX

MTPLX depends on mlx-lm and gets its MLA implementation from these files. Its tuner benchmarks autoregressive decode against each speculative depth and picks a winner.

Context	Arm	AR	MTP D1	vs AR	Verify	Acceptance
~700 tok	stock	60.7	47.6	0.78×	41.6 ms	—
~700 tok	patched	61.1	87.0	1.42×	20.1 ms	90.3%
2K–16K	stock	47.8	7.1	0.15×	413.9 ms	72.1%
2K–16K	patched	47.7	64.4	1.35×	27.7 ms	73.8%

Acceptance barely moves while verify cost falls 14.9×: the drafting was always fine but the verification was not. Every draft depth flips from losing to winning -- 0.83× / 0.82× / 0.64× at depths 1/2/3 becomes 1.39× / 1.15× / 1.06×.

Tests

test_models covers all seven touched architectures: 85 OK on both arms. test_generate 30, test_gated_delta 8, test_utils 11, test_losses 4 -- all OK on both. test_prompt_cache (4 errors) and test_sample_utils (1 failure) are identical on both arms and reproduce on unmodified main.

tests/test_absorbed_mla.py is new: it covers the large-cache and cold-cache limits, the quantised latent form, that each of the seven models makes one gate decision used by both guards, and that deepseek_v32's indexer gate is untouched.

Scope

Measured on deepseek_v3, glm4_moe_lite and kimi_linear. kimi_k3, longcat_flash and bailing_moe_v3 share the same 512/128/128 head geometry and threshold of 170 but had no checkpoint small enough to run here. deepseek_v32 is sparse, so its benefit is limited to caches below the top-k budget and neutral above (3.4× at 512, 0.99-1.02× at 2048/8192/32768 on the equivalent mlx-vlm model).

Outputs are not bit-identical: at 4-bit the two branches differ by 4.883e-04 max on values of scale ~0.1, against a bf16 ULP of ~2.44e-04 -- one to two units in the last place, and flat across L, so it is quantised-matmul ordering.

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Claude Code (Opus 5) wrote the patch and the benchmark scripts. I reviewed the diff, ran every measurement myself on an M3 Max, and take responsibility for every line. This description is mine.

[pierre427]

Worked on something similar recently, dropping a couple of tips for you:

consider exact integer arithmetic for the crossover calculation to avoid float-boundary rounding; our implementation also applies the same gate consistently across more MLA model families.