Skip to content

feat(deepseek-v3.2 / glm-moe-dsa): implement the DSA lightning indexer#583

Merged
inureyes merged 1 commit into
mainfrom
feature/issue-509-dsa-indexer
Jul 1, 2026
Merged

feat(deepseek-v3.2 / glm-moe-dsa): implement the DSA lightning indexer#583
inureyes merged 1 commit into
mainfrom
feature/issue-509-dsa-indexer

Conversation

@inureyes

@inureyes inureyes commented Jul 1, 2026

Copy link
Copy Markdown
Member

Summary

Replaces the full-attention fallback in deepseek_v32 (DeepSeek V3.2) and glm_moe_dsa with the DeepSeek Sparse Attention (DSA) "lightning indexer" these architectures were trained for, so the models attend to the top-k most relevant KV entries per query and restore long-context fidelity.

Changes

  • Shared indexer module (src/models/deepseek_v32_indexer.rs): the wq_b / wk / k_norm / weights_proj projections, relu(q @ k^T) scoring combined with the per-head weights_proj scale, causal masking, and top-index_topk key selection via gather.
  • Per-model indexer_rope_interleave flag threaded into the indexer RoPE, with the correct per-model default: non-interleaved (false) for deepseek_v32, interleaved (true) for glm_moe_dsa. Getting this wrong silently corrupts key selection, so it is covered by explicit tests.
  • Config plumbing: index_topk / index_n_heads / index_head_dim / indexer_rope_interleave added to deepseek_v32 (glm already parsed index_*), and the wq_b / wk / k_norm / weights_proj indexer weights are loaded. glm_moe_dsa's previously-unused index_* fields are now wired.
  • When kv_len <= index_topk the indexer reduces to the dense path (short context stays numerically equivalent to the previous fallback), and the L==1 decode path gathers the selected keys instead of materializing a per-step sparse mask.

Correctness (covered by unit tests)

  • deepseek_v32_indexer_rope_defaults_to_non_interleaved and glm_moe_dsa_indexer_rope_is_interleaved_by_default: the per-model RoPE interleave default is correct.
  • indexer_load_threads_rope_flag_per_model, glm_moe_dsa_falls_back_to_direct_rope_fields, glm_moe_dsa_prefers_rope_parameters_over_direct_fields, ..._config_overrides_are_respected, ..._config_defaults_thread_through.
  • indexer_load_absent_weights_is_dense_fallback and indexer_top_indices_none_at_short_context: reduces to dense when weights are absent or when kv_len <= index_topk.
  • indexer_top_indices_selects_highest_scoring_keys and indexer_top_indices_respects_causal_mask: top-k selection is correct and causal.

Validation

  • cargo check --lib --tests --features metal,accelerate: compiles clean.
  • cargo test --release --features metal,accelerate deepseek_v32: 7 pass. ... glm_moe_dsa: 5 pass. Total 12 indexer unit tests pass on Apple Silicon Metal.
  • Not run here: generation on a real DeepSeek-V3.2 or GLM-MoE-DSA checkpoint at context length greater than index_topk. Those real-model parity tests are gated ignored (they require local model weights and, for the pipeline variants, TCP-bound remote stages), and no such checkpoint is present in this environment. Long-context behavior is validated at the indexer scoring / top-k selection / dense-reduction level, not real-model generation, and should be confirmed on a real checkpoint before relying on it in production.

Closes #509

Replace the full-attention fallback in deepseek_v32 and glm_moe_dsa with the DeepSeek Sparse Attention lightning indexer. Adds a shared indexer module (wq_b, wk, k_norm, weights_proj projections; relu(q @ k^T) scoring combined with per-head weights_proj scaling; causal masking; top-index_topk key selection via gather) and threads the per-model indexer_rope_interleave flag (non-interleaved for deepseek_v32, interleaved for glm_moe_dsa) into the indexer RoPE. When kv_len is at most index_topk the indexer reduces to the dense path so short context stays numerically equivalent to the previous fallback, and the single-token decode path gathers the selected keys instead of materializing a per-step sparse mask. Wires the previously-unused index_* config on glm_moe_dsa and adds index_topk, index_n_heads, index_head_dim, and indexer_rope_interleave to deepseek_v32.
@inureyes inureyes added the status:done Completed label Jul 1, 2026
@inureyes
inureyes merged commit a4b0ed0 into main Jul 1, 2026
5 checks passed
@inureyes
inureyes deleted the feature/issue-509-dsa-indexer branch July 1, 2026 11:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

status:done Completed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant