Fix deepseek_v32 Indexer evicting attention sinks from sparse top-k

mlx_lm/models/deepseek_v32.py

@@ -109,6 +109,37 @@ 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 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]
+        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[:, None],
+            mx.array(float("inf"), scores.dtype),
+            scores,
+        )
+
         return mx.argpartition(scores, kth=-self.index_topk, axis=-1)[
             ..., -self.index_topk :
         ]

tests/test_models.py

@@ -1453,6 +1453,76 @@ 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_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