Follow-up to #111, whose stage 1 (SharpAI/mlx-swift-lm#48, submodule bump #138) loads glm_moe_dsa / deepseek_v3_2 with dense attention. Dense is numerically exact while the KV cache is at or under index_topk (2048 for GLM-5.2) because the reference indexer returns no selection until the cache outgrows it — beyond that, stage 1 runs dense where the reference runs top-2048 sparse. Stage 2 makes long context faithful (and fast: dense attention over ~1M positions is also the wrong cost profile, given the model advertises max_position_embeddings: 1048576).
Reference: mlx_lm/models/deepseek_v32.py (the Indexer class and DeepseekV32Attention.__call__).
What needs building
1. The Indexer module (~60 lines in the reference). Sub-modules wq_b (from q_lora_rank), wk, k_norm (LayerNorm), weights_proj; RoPE over qk_rope_head_dim with indexer_rope_interleave (GLM-5.2: true → traditional). Scores are ReLU'd q·k, weighted by weights_proj, summed over heads; selection is argPartition(scores, kth: -indexTopk). All primitives exist in mlx-swift: argPartition (Ops.swift:257), takeAlong/putAlong (already used in DeepseekV3.swift:286-291).
2. Paired KV caches per attention layer. The reference threads cache[0] (main MLA latent) and cache[1] (indexer keys) per layer, and couples them with mx.depends after the indexer runs so the lazy graph stays bounded. No Swift model here uses that exact shape — the hybrid models (Jamba, FalconH1, NemotronH) index cache[1] for layer state, which is precedent but not the same contract. newCache needs to produce the pair and everything downstream (canTrimPromptCache, trim, quantized-KV) needs to keep working.
3. The two sparse application paths. Decode (L == 1): gather the top-k latent/rope entries with takeAlong and attend over the gathered set. Prefill (L > 1): build a boolean mask with putAlong and AND it with the causal mask. Both are in the reference at DeepseekV32Attention.__call__.
4. Absorbed MLA (embed_q / unembed_out). The reference's decode path attends in latent space, absorbing the q/out projections. No model in this repo implements absorbed MLA — this is written-from-scratch, not ported-by-analogy, and is where reviewer attention should go.
5. Stop dropping weights. Remove the .indexer. filter from DeepseekV32Model.sanitize (and close the DeepseekV4Compressor/DeepseekV4Indexer TODO at DeepseekV4.swift:849 — same mechanism, three architectures).
The test that makes this verifiable without real weights
The inertness property gives stage 2 a strong self-check that needs no checkpoint: with the indexer implemented, output for any context ≤ index_topk must be identical to stage 1's dense output (the indexer returns nothing in that regime). A tiny random-init model with index_topk lowered to, say, 8 can then assert the sparse path engages past it and still produces finite, shape-correct logits — and that with index_topk large, sparse ≡ dense. The stage-1 tests in Tests/MLXLMTests/DeepseekV32Tests.swift are the scaffold.
Standing caveat
Nothing in stages 1 or 2 is validated against real weights — the smallest glm_moe_dsa checkpoint is 308 GB. First real-weights run (per #111, avlp12/GLM-5.2-Alis-MLX-Dynamic-3.5bpw with --stream-experts) should be treated as part of the verification, not an afterthought.
Follow-up to #111, whose stage 1 (SharpAI/mlx-swift-lm#48, submodule bump #138) loads
glm_moe_dsa/deepseek_v3_2with dense attention. Dense is numerically exact while the KV cache is at or underindex_topk(2048 for GLM-5.2) because the reference indexer returns no selection until the cache outgrows it — beyond that, stage 1 runs dense where the reference runs top-2048 sparse. Stage 2 makes long context faithful (and fast: dense attention over ~1M positions is also the wrong cost profile, given the model advertisesmax_position_embeddings: 1048576).Reference:
mlx_lm/models/deepseek_v32.py(theIndexerclass andDeepseekV32Attention.__call__).What needs building
1. The
Indexermodule (~60 lines in the reference). Sub-moduleswq_b(fromq_lora_rank),wk,k_norm(LayerNorm),weights_proj; RoPE overqk_rope_head_dimwithindexer_rope_interleave(GLM-5.2:true→ traditional). Scores are ReLU'd q·k, weighted byweights_proj, summed over heads; selection isargPartition(scores, kth: -indexTopk). All primitives exist in mlx-swift:argPartition(Ops.swift:257),takeAlong/putAlong(already used inDeepseekV3.swift:286-291).2. Paired KV caches per attention layer. The reference threads
cache[0](main MLA latent) andcache[1](indexer keys) per layer, and couples them withmx.dependsafter the indexer runs so the lazy graph stays bounded. No Swift model here uses that exact shape — the hybrid models (Jamba,FalconH1,NemotronH) indexcache[1]for layer state, which is precedent but not the same contract.newCacheneeds to produce the pair and everything downstream (canTrimPromptCache, trim, quantized-KV) needs to keep working.3. The two sparse application paths. Decode (
L == 1): gather the top-k latent/rope entries withtakeAlongand attend over the gathered set. Prefill (L > 1): build a boolean mask withputAlongand AND it with the causal mask. Both are in the reference atDeepseekV32Attention.__call__.4. Absorbed MLA (
embed_q/unembed_out). The reference's decode path attends in latent space, absorbing the q/out projections. No model in this repo implements absorbed MLA — this is written-from-scratch, not ported-by-analogy, and is where reviewer attention should go.5. Stop dropping weights. Remove the
.indexer.filter fromDeepseekV32Model.sanitize(and close theDeepseekV4Compressor/DeepseekV4IndexerTODO atDeepseekV4.swift:849— same mechanism, three architectures).The test that makes this verifiable without real weights
The inertness property gives stage 2 a strong self-check that needs no checkpoint: with the indexer implemented, output for any context ≤
index_topkmust be identical to stage 1's dense output (the indexer returns nothing in that regime). A tiny random-init model withindex_topklowered to, say, 8 can then assert the sparse path engages past it and still produces finite, shape-correct logits — and that withindex_topklarge, sparse ≡ dense. The stage-1 tests inTests/MLXLMTests/DeepseekV32Tests.swiftare the scaffold.Standing caveat
Nothing in stages 1 or 2 is validated against real weights — the smallest
glm_moe_dsacheckpoint is 308 GB. First real-weights run (per #111,avlp12/GLM-5.2-Alis-MLX-Dynamic-3.5bpwwith--stream-experts) should be treated as part of the verification, not an afterthought.