feat(deepseek-v3.2 / glm-moe-dsa): implement the DeepSeek Sparse Attention (DSA) indexer to replace the full-attention fallback

[inureyes]

Summary

deepseek_v32 (DeepSeek V3.2) and glm_moe_dsa currently run with a full-attention fallback; the DeepSeek Sparse Attention (DSA) "lightning indexer" that these architectures are built around is deferred. This issue tracks implementing the indexer so these models attend to the top-k most relevant KV entries per query (the behavior the checkpoints were trained for), restoring long-context fidelity and the intended decode efficiency.

Current state in mlxcel

• src/models/deepseek_v32.rs:21-28 documents the deferral verbatim: "Indexer for sparse attention (deferred - using full attention fallback) ... Full attention is used instead", plus a TODO to add the single-token (L==1) fast path when the indexer lands.
• src/models/glm_moe_dsa.rs wraps DeepSeekV32Model (glm_moe_dsa.rs:223-228) and therefore inherits the same full-attention path. It already parses index_head_dim / index_n_heads / index_topk as config fields (glm_moe_dsa.rs:100-104) but leaves them unused (set to None, glm_moe_dsa.rs:315-316).
• deepseek_v32.rs parses q_lora_rank (deepseek_v32.rs:73-74) but has no indexer config (index_*, indexer_rope_interleave).
• No interleaved-vs-traditional RoPE handling for the indexer exists in either file (grep for interleave / indexer_rope / traditional returns nothing).

Consequence: at short context (when index_topk >= kv_len) the fallback is numerically identical to the trained model, but at long context the attended key set diverges from the reference, and the intended sparse-attention decode speedup is unrealized.

What needs to be built

1. Indexer module (shared, used by both models). Per decoder layer, alongside MLA:

   • Projections: wq_b: Linear(q_lora_rank -> index_n_heads * index_head_dim), wk: Linear(hidden -> index_head_dim), k_norm: LayerNorm(index_head_dim), weights_proj: Linear(hidden -> index_n_heads).
   • Apply RoPE to the indexer q and k over qk_rope_head_dim with the running cache offset (see interleave below).
   • Score: scores = relu(q @ k^T); head weights w = weights_proj(x) * (index_n_heads^-0.5 * head_dim^-0.5); combine scores = (scores * w).sum(over heads); apply the causal mask.
   • Select the top-index_topk keys per query. When kv_len <= index_topk, skip indexing (dense path). Use the selected indices to gather the KV subset for the main attention (take_along_axis) rather than constructing a sparse mask.

2. Per-model indexer_rope_interleave flag (load-bearing correctness). The indexer RoPE is non-interleaved (traditional = false) for deepseek_v32 and interleaved (traditional = true) for glm_moe_dsa. Add the config field with the correct per-model default and thread it into the indexer RoPE construction. Getting this wrong silently corrupts indexer scores, which corrupts key selection, with no crash.

3. Single-token decode fast path (L==1). Per the existing TODO at deepseek_v32.rs:26-28, gather the top-k KV entries directly via take_along_axis instead of materializing a sparse mask each step.

4. Config plumbing. Add index_topk / index_n_heads / index_head_dim to deepseek_v32 (glm_moe_dsa already has them); add indexer_rope_interleave to both with defaults deepseek_v32 = false, glm_moe_dsa = true; load the wq_b / wk / k_norm / weights_proj indexer weights.

Acceptance criteria

• Indexer implemented and wired for both deepseek_v32 and glm_moe_dsa; full-attention fallback removed (or kept behind an env flag for A/B only).
• indexer_rope_interleave correct per model (v3.2 non-interleaved, glm interleaved), covered by a unit test that asserts the RoPE path per model type.
• Short-context parity: with kv_len <= index_topk, decode is RMS-equivalent (ideally bit-for-bit) to the current full-attention fallback (the indexer must reduce to dense).
• Long-context behavior matches the reference key selection (validate the top-k index set on a constructed case).
• L==1 decode fast path uses gather (no per-step sparse-mask allocation).
• Validated on a real checkpoint (a DeepSeek-V3.2 or GLM-MoE-DSA mlx-community export): generation runs and stays coherent at context length greater than index_topk.

Files

• src/models/deepseek_v32.rs (indexer, config, attention wiring)
• src/models/glm_moe_dsa.rs (config defaults, interleave flag)
• a shared indexer helper if warranted (mirror the existing qwen3_next_helpers.rs split pattern)
• tests: new src/models/deepseek_v32_tests.rs and/or extend the glm_moe_dsa tests

Reference

The indexer algorithm and the per-model interleave defaults are specified upstream: https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/deepseek_v32.py and ml-explore/mlx-lm#1431 (the commit that parameterized indexer_rope_interleave: false for deepseek_v32, true for glm_moe_dsa).