feat(deepseek-v3.2 / glm-moe-dsa): implement the DSA lightning indexer#583
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replaces the full-attention fallback in
deepseek_v32(DeepSeek V3.2) andglm_moe_dsawith 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
src/models/deepseek_v32_indexer.rs): thewq_b/wk/k_norm/weights_projprojections,relu(q @ k^T)scoring combined with the per-headweights_projscale, causal masking, and top-index_topkkey selection via gather.indexer_rope_interleaveflag threaded into the indexer RoPE, with the correct per-model default: non-interleaved (false) fordeepseek_v32, interleaved (true) forglm_moe_dsa. Getting this wrong silently corrupts key selection, so it is covered by explicit tests.index_topk/index_n_heads/index_head_dim/indexer_rope_interleaveadded todeepseek_v32(glm already parsedindex_*), and thewq_b/wk/k_norm/weights_projindexer weights are loaded.glm_moe_dsa's previously-unusedindex_*fields are now wired.kv_len <= index_topkthe indexer reduces to the dense path (short context stays numerically equivalent to the previous fallback), and theL==1decode 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_interleavedandglm_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_fallbackandindexer_top_indices_none_at_short_context: reduces to dense when weights are absent or whenkv_len <= index_topk.indexer_top_indices_selects_highest_scoring_keysandindexer_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.index_topk. Those real-model parity tests are gatedignored(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