Fix deepseek_v32 Indexer evicting attention sinks from sparse top-k - #1552
robertlangdonn wants to merge 2 commits into
Conversation
Once a sequence exceeds index_topk, the DSA Indexer's learned top-k selection does not reliably keep the first few key positions. Losing these attention sinks (Xiao et al. 2023) causes the softmax to redistribute onto irrelevant keys, and decode collapses into repetition/garbage exactly at the index_topk boundary. Affects deepseek_v32 and glm_moe_dsa (which reuses this Indexer). Force sink columns and a small recency window into the selection before argpartition. Non-causal picks for early prefill rows are harmless since the caller ANDs this selection with the real causal mask before use.
The previous commit assumed offset was a scalar. Under BatchKVCache (batched generate()), offset and left_padding are per-sequence arrays, and "sink" means the first real tokens of each sequence, not buffer column 0 once left-padding shifts them. Fold both into the position math with an explicit batch axis instead. Adds a regression test for the batched + left-padded case alongside the single-sequence one.
4ee25f2 to
a37627b
Compare
snapkv_keep_indices already staged sinks + recent window into the keep set before the score-ranked fill, but capped the sink floor at max(1, budget//8), so a tight budget silently kept fewer than sink_tokens sinks. Dropping an attention sink collapses decode into repetition -- the same failure mlx-lm ml-explore#1552 fixed for DeepSeek's DSA indexer top-k selection. Add a default-on guarantee_sinks guard that clamps the sink floor only by budget/seq_len, so the first N sinks (and the recent window) can never be evicted no matter how high a middle row scores. guarantee_sinks=False restores the old capped behaviour. Existing callers use budget>=32 (cap>=4) so behaviour is unchanged there; the guard only bites when budget//8 < sink_tokens. Adds regression tests: N sinks + recency survive budget=16 with every middle row scored 1e9, and the flag toggles the floor. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Independent verification on
|
| total tokens | stock 0.31.3 | with this PR |
|---|---|---|
| ≤ ~1000 | 0% | 0% |
| ~1176–1926 | 3–12% | 0–1% |
| ~2176 (crosses 2048) | 13% | 1% |
| ~2426 | 26% | 4% |
| ~2676 | 22% | 9% |
| ~2926 | 67% | 0% |
| ~3176 | 78% | 0% |
Stock degenerates into full repetition right past the threshold (tail: the model arguing with itself in a two-phrase loop); with the patch the essay stays coherent to the end.
Dense-attention control: the same generation through llama.cpp (GGUF build of the same model family, dense attention, same sampling) is 0% in every block — which rules out the checkpoint and sampling, and matches the sink-eviction mechanism: the failure belongs to the sparse path exactly at its activation boundary.
One methodological note that may help reviewers: needle-in-haystack tests at 6–8K prompt tokens pass on stock — prefill is not where this bites, which is presumably why it survived so long. A repetition profile of a long generation across the threshold catches it in one run.
Would be great to see this merged — the same failure signature is currently being re-discovered model by model (cf. the DeepSeek-V4 thread in #1189).
3 participants
Fix deepseek_v32 Indexer evicting attention sinks from sparse top-k
Summary
index_topk, the DSAIndexer's learned top-k selection does not reliably keep the first few key positions (attention sinks), and sparse decode collapses into repetition/garbage right at that boundary.deepseek_v32andglm_moe_dsa(which reuses thisIndexer/Model).sinks_keptinstrumentation on a real GLM-5.2 checkpoint) and an initial fix sketch — this PR verifies it independently, generalizes it to the batched/left-padded path, and adds regression tests.Problem
Indexer.__call__ranks all key positions by a learned score and keeps only the topindex_topkviaargpartition. Nothing guarantees the first few positions (attention sinks, Xiao et al. 2023, StreamingLLM) survive that ranking — if the router's raw score for a sink is low relative to other keys, it drops out. Once even one sink is evicted across enough of the ~60+ layers in a real model, the softmax redistributes onto irrelevant keys and generation collapses.Reproduced directly against
Indexer(no checkpoint needed) with a synthetic sequence containing a "distractor" block that outscores the sink columns on raw dot-product alone — pre-fix, the last query's top-k drops 3 of 4 sink columns.Fix
Force sink columns + a small recency window into the selection before
argpartition, by setting their score to+inf:Non-causal picks for early prefill rows are harmless — the caller ANDs this selection with the real causal mask (
sparse_mask & mask) before it reaches attention.offset/left_paddingare folded in with an explicit batch axis rather than assumed scalar, because underBatchKVCache(batchedgenerate()) both are per-sequence arrays, and "sink" means the first real tokens of each sequence — not buffer column 0 once left-padding shifts them.Total change: +31 in
mlx_lm/models/deepseek_v32.py.Test plan
test_deepseek_v32_indexer_keeps_attention_sinks— synthetic single-sequence case with a distractor block, asserts all 4 sinks survive top-k (fails pre-fix, passes post-fix).test_deepseek_v32_indexer_keeps_padded_batch_sinks— batchedBatchKVCachewith per-sequence left-padding, asserts each sequence's real sinks (offset by its own padding) survive.python -m unittest discover tests/— 188 tests, only pre-existing unrelated import errors (datasets/lm_eval/requests, optional deps not installed locally).deepseek_v32.Model(tiny synthetic config), prefill pastindex_topkthen 5 decode steps — runs cleanly, exercises theL==1decode path and growing cache offset.pre-commit run --files mlx_lm/models/deepseek_v32.py tests/test_models.py— clean.