SIGN IN SIGN UP

mla: use the absorbed path for multi-token steps (#2150)

The absorbed MLA path was gated on `L == 1`, so any step with more than
one query materialized K and V across the whole cached latent instead.
Materializing costs O(T) in the cached length and does not depend on L,
while the absorbed path costs O(L), so on a warm cache the crossover is a
fixed query count per model, 170 for DeepSeek V3 and 398 for GLM-4 MoE
Lite, against a gate that admitted exactly one.

`models/mla.py` gains `max_absorbed_queries()` and `latent_length()`,
applied in the ten models with an absorbed path: deepseek_v3,
deepseek_v32, kimi_k3, kimi_linear, longcat_flash, longcat_flash_sparse,
glm4_moe_lite, glm_moe_dsa, glm5_next and youtu_vl. The same change for
mlx-lm is ml-explore/mlx-lm#1817.

The threshold takes the cache length, so it is right at both ends. Per
head, with r = kv_lora_rank and d = nope + v, the absorbed path costs
L*r*d + 2*L*T*r and materializing costs T*r*d + L*T*d, giving

    L < r*d / (r*d/T + 2*r - d)

As T grows this tends to r*d / (2*r - d), the fixed count above. On a
cold cache, where T == L, it reduces to 2*r < d, false for all ten
models, so the absorbed path is correctly rejected. That case is
reachable: `generate/ar.py` chunks at min(prefill_step_size, N - 1), so
a prompt shorter than the chunk size is one step with T == L. Taking the
absorbed path there costs about 1.95x the attention FLOPs at L = T = 169,
which is 136.5 ms against 133.7 ms on a whole-model forward.

Both gates are now driven by one boolean computed once per call. With
two independent comparisons a mechanical edit across ten files can widen
one and leave the other, and the shape of the code does nothing to stop
it; with one decision that failure cannot be expressed.

The cache length comes from the latent rather than from `offset`, since
`offset` is not in scope at two of the ten sites, and `latent_length()`
tolerates the quantized 3-tuple form.

The sparse-indexer gate stays at `L == 1` in deepseek_v32,
longcat_flash_sparse, glm5_next and glm_moe_dsa. It selects with
`topk_indices[:, :, 0, :]`, the first query's top-k, which is correct for
a single query only.

Full forward pass, Moonlight-16B-A3B-Instruct-4-bit (DeepSeek V3
architecture, 27 layers), M3 Max, medians of 5, cache 32768:

| L | 1 | 2 | 4 | 8 | 16 | 32 | 64 |
|--:|--:|--:|--:|--:|---:|---:|---:|
| stock | 17.4 | 385.4 | 394.3 | 419.3 | 489.1 | 512.9 | 546.4 |
| patched | 17.3 | 22.0 | 28.5 | 48.7 | 84.2 | 152.7 | 281.0 |
| speedup | 1.00x | 17.5x | 13.8x | 8.6x | 5.8x | 3.4x | 1.9x |

L = 1 is unchanged. Cold-cache prompts are unchanged as well, at 32, 64,
128, 169 and 400 tokens.

Peak memory for one attention call at 32768 context falls from 2478 MB to
462 MB at L = 4, and stays lower at every L up to the threshold.

Numerics: in float32 the two forms agree to a relative 3.4e-06 across
L = 1 to 170 at batch 1, 2 and 4, and at L = 1 output is bit-identical.
On 4-bit weights they reorder quantized matmuls, so a wide step is not
bit-identical to the other form; teacher-forced over 1024 positions
against token-by-token prefill, stock agrees on 98.54% of argmaxes and
patched on 98.34%, with every disagreement at a near-tie.

Tests: unmodified main gives 3131 passed, this branch 3140, the
difference being the nine tests added here, with per-file results
otherwise identical. The new tests run the real `__call__` and force each
branch, and were checked against four mutations: unpaired gates, a full
revert of one file, a widened indexer gate, and dropping the cache-length
term. All four fail the suite.
S
Stuart Rowlands committed
5c9b5f52adfeab35b5ece0bb2d6e4d44541d9e32
Parent: 89baff4
Committed by GitHub <noreply@github.com> on 9/3/2026, 8:53:36 PM