mla: use the absorbed path for multi-token steps - #2150
Conversation
73807c3 to
8831903
Compare
8831903 to
3bb364d
Compare
This comment was marked as low quality.
This comment was marked as low quality.
Sorry, something went wrong.
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.
3bb364d to
64ff26c
Compare
|
Thanks for the review! Everything is addressed and pushed; replies are inline on each thread. Summary: Cold-cache regression. Reproduced it before changing anything, 1.28x to 1.95x on the attention op at Gate wiring. Both guards now read one boolean computed once per call, so the mismatched-pair failure you described cannot be expressed rather than merely being tested for. Tests. They run the real CI. Fixed, and it was
Two corrections. Your absorbed FLOP figure is 51.4M rather than 51.9M, which leaves the 1.75x ratio intact, and the One thing I could not confirm: I verified the first half of your quantized-KV note, the materialized branch does fail on a quantized latent with Test suite: unmodified main 3131 passed, this branch 3140, the difference being the nine tests added here, per-file otherwise identical. |
3 participants
The absorbed MLA path is gated on
L == 1. Any step with more than one query materialises K and V across the whole cached latent instead, and materialising costs O(T) in the cached length regardless of how many queries you have, while the absorbed path costs O(L). The crossover is therefore a fixed query count per model:For DeepSeek V3 that is 170 queries and for GLM-4 MoE Lite it is 398, against a gate that admits exactly one. So every speculative verify step, and every short continuation on a warm cache, has been taking the expensive branch.
This adds
max_absorbed_queries()tomodels/mla.pyand applies it 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_nextandyoutu_vl. The helper is byte-identical to the one in ml-explore/mlx-lm#1817, including the attribute name, so the two runtimes pick the same width for a shared model.It is the same reasoning as my mlx-lm PR ml-explore/mlx-lm#1817, extended here to four models that do not exist in mlx-lm.
Multi-turn chat, which needs no drafter
DEFAULT_PREFILL_STEP_SIZEis 2048, so a long prompt chunks above the crossover and is unaffected. A follow-up turn on a warm cache is not chunked though, it is one step of however many tokens the turn holds, and that is where most people will see this. Moonlight-16B-A3B-Instruct-4-bit at 32K context, full forward pass, M3 Max, medians of 5:Speculative verify steps
Same model and setup, at the draft lengths the in-tree MTP drafters use:
At 16384 the same three steps are 10.7x, 8.9x and 5.9x, and at 4096 they are 4.2x, 3.4x and 2.5x. L=1 is unchanged because the gate already picked the absorbed path there, which makes that row a control on the rest of the table.
The stock row is nearly flat across L because the cost is dominated by materialising K and V over the cache, which does not depend on L at all, so one extra token costs 21x more than the first one does at 32K.
Memory
Materialising also expands the latent into per-head K and V. Peak memory for a single attention call at 32768 context drops from 2478 MB to 462 MB at L=4, and the absorbed path stays lower at every L up to the threshold (2742 MB against 3602 MB at L=144).
The sparse indexer gate is left alone
deepseek_v32,longcat_flash_sparse,glm5_nextandglm_moe_dsahave a secondL == 1gate on the sparse top-k gather. That one selects withtopk_indices[:, :, 0, :], the first query's top-k, which is only correct for a single query, so widening it would silently apply one query's selection to all of them. It stays atL == 1and there is a test asserting it stays there.Numerics
In float32 the two forms agree to a relative 3.4e-06 across L = 1 to 170 at batch 1, 2 and 4, because they are the same linear algebra associated differently. At L = 1 the patch is a no-op and output is bit-identical.
I want to be clear about the quantised case. On 4-bit weights the two forms reorder quantised matmuls, so a wide step is not bit-identical to the other form. Teacher-forced over 1024 positions with token-by-token prefill as the reference, the stock wide chunk agrees on 98.54% of argmaxes and the patched wide chunk on 98.34%. Every disagreement sits at a near-tie, with top-2 margins between 0.00 and 0.30 against a median margin of 0.664, on a logit scale of
23.4. Drift of this size is already present without the patch, because the stock
wide chunk diverges from sequential prefill at the same rate, and in free-running greedy decode at the same token.
Testing
Unmodified main gives
3131 passedand this branch gives3137 passed, the difference being the six tests added here, with per-file results otherwise identical, so the change moves no existing test either way.test_smoke.pyfails to collect on this branch and on unmodified main alike, which is unrelated to this change.mlx_vlm/tests/test_absorbed_mla.pycovers the threshold formula, wide-step parity against the materialised form, and the preserved indexer gate.The glm5_next hunk overlaps #2127
glm5_nextis one of the ten models changed here, and #2127 adds that file at full length, so this is the one place the two PRs touch the same lines. That overlap is what superseded my earlier #2141, which only changedglm5_next. I am happy to wait on #2127 and re-apply this hunk afterwards, or to dropglm5_nextfrom this PR entirely if that is easier to review, since the other nine models stand on their own either way. Let me know which you prefer.What I did not measure
Moonlight-16B is the only one of the ten models I have local weights for, so the end-to-end numbers above are all from
deepseek_v3and the other nine are covered by the unit tests and by the block being identical. I also have no cached model that pairs a dense-MLA target with a working MTP drafter, so the speculative numbers are per-step costs rather than a full generation loop with measured acceptance.