From 7d7630de6ec43309ba317dfab19789a2f4203872 Mon Sep 17 00:00:00 2001 From: robertlangdonn Date: Sun, 12 Jul 2026 01:35:19 +0530 Subject: [PATCH 1/2] Fix deepseek_v32 Indexer evicting attention sinks from sparse top-k 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. --- mlx_lm/models/deepseek_v32.py | 20 ++++++++++++++++++++ tests/test_models.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index d696c02c5..dcbb0f4ae 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/mlx_lm/models/deepseek_v32.py @@ -109,6 +109,26 @@ def __call__( scores = scores.sum(axis=1, keepdims=True) if mask is not None: scores = mx.where(mask, scores, -float("inf")) + + # The learned top-k does not reliably keep the first few key positions + # (attention sinks). Once sparse attention starts (past index_topk), + # losing a sink collapses generation into repetition (StreamingLLM + # effect, Xiao et al. 2023). Force sinks + a small recency window into + # the selection; non-causal picks for early rows are harmless since the + # caller ANDs this selection with the real causal mask before use. + n_sinks, local_window = 4, 128 + num_keys = scores.shape[-1] + query_pos = (mx.arange(scores.shape[2]) + offset).reshape(-1, 1) + key_pos = mx.arange(num_keys).reshape(1, num_keys) + force_keep = (key_pos < n_sinks) | ( + (key_pos <= query_pos) & (key_pos > query_pos - local_window) + ) + scores = mx.where( + force_keep.reshape(1, 1, scores.shape[2], num_keys), + mx.array(float("inf"), scores.dtype), + scores, + ) + return mx.argpartition(scores, kth=-self.index_topk, axis=-1)[ ..., -self.index_topk : ] diff --git a/tests/test_models.py b/tests/test_models.py index 6b0e00173..da9e12615 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1453,6 +1453,39 @@ def test_deepseek_v3(self): model, args.model_type, args.vocab_size, args.num_hidden_layers ) + def test_deepseek_v32_indexer_keeps_attention_sinks(self): + # Regression test for #1443: once a sequence exceeds `index_topk`, the + # Indexer's learned top-k selection must never let the first few key + # positions (attention sinks) drop out, or sparse decoding collapses. + from mlx_lm.models import deepseek_v32 + + mx.random.seed(0) + + args = deepseek_v32.ModelArgs( + model_type="deepseek_v32", + hidden_size=32, + index_head_dim=8, + index_n_heads=2, + index_topk=200, + q_lora_rank=16, + qk_rope_head_dim=4, + ) + indexer = deepseek_v32.Indexer(args) + + b, s_total = 1, 400 + x = mx.random.normal((b, s_total, args.hidden_size)) * 0.05 + # A "distractor" block outside both the sink zone (0..3) and the local + # recency window competes for the ranking on raw score alone. + x = mx.concatenate([x[:, :50], x[:, 50:150] * 20.0, x[:, 150:]], axis=1) + qr = mx.random.normal((b, s_total, args.q_lora_rank)) + mask = mx.tril(mx.ones((s_total, s_total), dtype=mx.bool_)) + + topk_indices = indexer(x, qr, mask, cache=None) + self.assertIsNotNone(topk_indices) + + last_row = set(topk_indices[0, 0, -1].tolist()) + self.assertEqual(last_row & set(range(4)), set(range(4))) + def test_gemma2(self): from mlx_lm.models import gemma2 From a37627b8671edbe6a8eff2fd146e7b69d0ed4bea Mon Sep 17 00:00:00 2001 From: robertlangdonn Date: Sun, 12 Jul 2026 01:45:47 +0530 Subject: [PATCH 2/2] Handle batched/left-padded offsets in the sink-forcing mask 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. --- mlx_lm/models/deepseek_v32.py | 31 +++++++++++++++++++---------- tests/test_models.py | 37 +++++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 10 deletions(-) diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index dcbb0f4ae..9e5236f28 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/mlx_lm/models/deepseek_v32.py @@ -110,21 +110,32 @@ def __call__( if mask is not None: scores = mx.where(mask, scores, -float("inf")) - # The learned top-k does not reliably keep the first few key positions - # (attention sinks). Once sparse attention starts (past index_topk), - # losing a sink collapses generation into repetition (StreamingLLM - # effect, Xiao et al. 2023). Force sinks + a small recency window into - # the selection; non-causal picks for early rows are harmless since the - # caller ANDs this selection with the real causal mask before use. + # The learned top-k does not reliably keep the first few real key + # positions of each sequence (attention sinks). Once sparse attention + # starts (past index_topk), losing a sink collapses generation into + # repetition (StreamingLLM effect, Xiao et al. 2023). Force sinks + a + # small recency window into the selection; non-causal picks for early + # rows are harmless since the caller ANDs this selection with the real + # causal mask before use. + # + # offset/left_padding can be a per-sequence array under BatchKVCache + # (not a python scalar), and "sink" means the first real tokens of + # each sequence, not buffer column 0 once left-padding is involved — + # so both are folded into the position math with an explicit batch + # axis instead of assuming a shared scalar offset. n_sinks, local_window = 4, 128 num_keys = scores.shape[-1] - query_pos = (mx.arange(scores.shape[2]) + offset).reshape(-1, 1) - key_pos = mx.arange(num_keys).reshape(1, num_keys) - force_keep = (key_pos < n_sinks) | ( + left_padding = mx.array(getattr(cache, "left_padding", 0)) + query_pos = mx.arange(scores.shape[2]).reshape(1, -1, 1) + ( + mx.array(offset) + left_padding + ).reshape(-1, 1, 1) + key_pos = mx.arange(num_keys).reshape(1, 1, num_keys) + sink_start = left_padding.reshape(-1, 1, 1) + force_keep = ((key_pos >= sink_start) & (key_pos < sink_start + n_sinks)) | ( (key_pos <= query_pos) & (key_pos > query_pos - local_window) ) scores = mx.where( - force_keep.reshape(1, 1, scores.shape[2], num_keys), + force_keep[:, None], mx.array(float("inf"), scores.dtype), scores, ) diff --git a/tests/test_models.py b/tests/test_models.py index da9e12615..8b2ce7063 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1486,6 +1486,43 @@ def test_deepseek_v32_indexer_keeps_attention_sinks(self): last_row = set(topk_indices[0, 0, -1].tolist()) self.assertEqual(last_row & set(range(4)), set(range(4))) + def test_deepseek_v32_indexer_keeps_padded_batch_sinks(self): + # Under BatchKVCache, offset/left_padding are per-sequence arrays, not + # a python scalar, and "sink" means the first real tokens of each + # sequence rather than buffer column 0 once left-padding shifts them. + from mlx_lm.models import deepseek_v32 + from mlx_lm.models.base import create_causal_mask + from mlx_lm.models.cache import BatchKVCache + + mx.random.seed(0) + + args = deepseek_v32.ModelArgs( + model_type="deepseek_v32", + hidden_size=32, + index_head_dim=8, + index_n_heads=2, + index_topk=200, + q_lora_rank=16, + qk_rope_head_dim=4, + ) + indexer = deepseek_v32.Indexer(args) + + b, s_total = 2, 400 + left_padding = [0, 30] + x = mx.random.normal((b, s_total, args.hidden_size)) * 0.05 + x = mx.concatenate([x[:, :50], x[:, 50:150] * 20.0, x[:, 150:]], axis=1) + qr = mx.random.normal((b, s_total, args.q_lora_rank)) + mask = create_causal_mask(s_total, left_padding=mx.array(left_padding)) + + cache = BatchKVCache(left_padding) + topk_indices = indexer(x, qr, mask, cache=cache) + self.assertEqual(topk_indices.shape[0], b) + + for i, pad in enumerate(left_padding): + real_sinks = set(range(pad, pad + 4)) + last_row = set(topk_indices[i, 0, -1].tolist()) + self.assertEqual(last_row & real_sinks, real_sinks) + def test_gemma2(self): from mlx_lm.models import gemma2