feat: DeepSeek V3.2 pure-Swift port — zero-fork overlay (parity 1e-4)#52
Merged
Conversation
First macMLX-owned model architecture (overlay, zero-fork). DeepSeek V3.2 is the foundation for a later V4 port; its DSA sparse attention (lightning indexer) is the shared novel piece. This session (S1) lands the leaf components: - DeepseekV32Configuration — Codable, matches config.json / the Python ModelArgs schema. 2 decode tests (real config + defaults), pass under swift test. - DeepseekMultiLinear — batched per-head linear, copied from mlx-swift-lm's internal MultiLinear (not public) for the absorbed-MLA embed_q/unembed_out projections. - DeepseekV32Indexer — the DSA lightning indexer: scores keys, returns top-index_topk indices (argPartition) or nil when the context fits in top-k. Prefill path (cache==nil); KV-cache decode integration lands with the attention module (S2). Translated from docs/reference/deepseek_v32_mlx_lm_reference.py. Feasibility audit + component checklist + parity strategy in docs/superpowers/plans/2026-05-11-deepseek-v32-port.md. Compiles clean (swift build). Indexer structural tests (shape, short-circuit) are Metal-gated; numerical parity vs the Python reference is the S1.5 follow-up (needs offline reference-value capture). Not yet registered in ModelOverlay — registration flips on after the full model passes parity.
Discovered the MLX-gated tests in this repo (PromptCache, LoRA-MLX, now DeepSeek) have NEVER actually run — the old requireMetalOrSkip checked for a bundled default.metallib, which is absent under both `swift test` (correct: MLX fatalErrors there) AND xcodebuild (wrong: MLX runs fine there via the Metal Toolchain). So everything skipped everywhere; MLX numerical correctness rested on compile-check + manual smoke only. New shared XCTestCase.requireMLXRuntimeOrSkip() discriminates on the test-bundle path: SPM's `swift test` bundle lives under `.build/` (skip — no metallib, MLX aborts); xcodebuild's lives under DerivedData (run — Metal Toolchain JITs the shaders). DeepseekV32Indexer tests now genuinely EXECUTE under `xcodebuild test -scheme MacMLXCore -skipPackagePluginValidation`: both the sparse (argPartition top-k) and short-circuit paths pass on Metal. This gives the DeepSeek port a real numerical-parity harness for S2 (attention) + S3 (MoE), instead of blind compile-only checks. Prereq: Metal Toolchain component (xcodebuild -downloadComponent MetalToolchain), installed locally.
The Swift DeepseekV32Indexer now selects the EXACT same top-k key indices as mlx-lm 0.31.3's reference Indexer given identical weights + inputs. The DSA lightning-indexer math (partial RoPE, score compute, weights_proj weighting, argPartition top-k) is numerically verified, not just compile-checked. Harness (reusable for S2 attention + S3 MoE): - docs/reference/capture_indexer.py — offline uv-venv capture from the Python reference (Python never enters macMLX). Fixed deterministic weights + inputs → runs reference Indexer → saves weights + inputs + expected sorted top-k to a safetensors fixture. - Tests/MacMLXCoreTests/Fixtures/indexer_parity_fixture.safetensors — the captured fixture (test resource; Package.swift copies it). - DeepseekV32IndexerParityTests — loads the fixture, sets the same weights on the Swift indexer, runs it, compares sorted top-k sets. Passes under xcodebuild (Metal); skips under swift test via the runtime-aware gate. Also fixes the indexer RoPE: mlx-lm 0.31.3 hardcodes traditional=true (a later upstream commit makes it configurable, default false — a behavior flip to resolve against real weights in S4). Match 0.31.3 so the fixture is valid. Includes the oMLX-parity roadmap doc (2026-05-11) for the broader plan.
DeepseekMultiLinear gains the transpose:Bool=true argument mirroring mlx-lm's mla.MultiLinear.__call__(x, transpose=True). The absorbed-MLA attention needs both directions: embed_q forward (transpose=true, inputDims→outputDims) and the prefill embed_q/unembed_out reverse (transpose=false, outputDims→inputDims). Compiles clean.
Mark components 1-3 done (Configuration, MultiLinear, Indexer — the last numerically verified vs Python). Record the S2 kickoff recipe so the next focused session starts fast: the absorbed-MLA two-branch forward (decode L==1 / prefill L>1), the CacheList double sub-cache, the parity-capture plan (both branches, uv venv, transformers<5.13), the weight key names, and the gotchas already learned.
Bring the branch base current before the rope-interleave hotfix: main gained MCPClientPool (PR #45) since this branch forked.
Stock mlx-lm 0.31.3 hardcodes traditional=True for the DSA indexer rope — a GLM5-era regression upstream reverted in PR #1431 (merged 2026-06-24, unreleased) because it silently degrades long-sequence quality. Correct value: config indexer_rope_interleave, default false. - Indexer rope: traditional: true → config.indexerRopeInterleave (the Configuration field already existed and decoded correctly) - Parity fixture regenerated against 0.31.3 + the #1431 one-line patch; top-k selection differs from the old fixture (the mode is numerically significant), and the parity test re-passes under xcodebuild against the corrected pair - Provenance headers added: the checked-in reference .py is a post-#1431 main snapshot, NOT PyPI 0.31.3 — the mismatch between those two is what made this bug hard to see in review - Port plan updated (component 3 notes + S2 gotchas)
LoRAAdapterConverterMLXTests (6 tests) and PromptCacheStoreTests (4 MLX-gated tests) still used the old metallib-probing helper, which skips under BOTH swift test and xcodebuild — those 10 tests had never actually executed anywhere. Migrated to MLXTestSupport's requireMLXRuntimeOrSkip; first real run: 10/10 pass.
CI previously ran only swift test (where every MLX-gated test self-skips) and an app build (no tests) — the numerical-parity fixtures protected nothing in CI. New 'metal' job runs the full MacMLXCore suite under xcodebuild with the Metal Toolchain installed. Known first-run risk: if the runner VM exposes no Metal device the job fails — then gate on a device probe rather than deleting the job.
Pool landed in PR #45 with two dead-server robustness fixes; remaining C2 tail (server-everything integration test, zombie reaping, stderr reader off the cooperative pool, connectAll signature) rolls into the hardening pass before C3.
DeepseekV32Attention: absorbed Multi-head Latent Attention with the DSA indexer's sparse mask folded in. Keeps K/V in the kv_lora_rank latent space (embed_q / unembed_out per-head projections) and feeds the RoPE contribution (pe_scores) to SDPA as an additive mask, so the full score is scale*(qNope*kT) + pe_scores in one softmax. All three branches implemented (prefill dense, sparse-prefill scatter, decode gather) — prefill is parity-verified now; sparse (S2.2) and decode+cache (S2.3) get their own fixtures next. S2.1 parity: attn_prefill_fixture (s=3<=index_topk=4 so the indexer short-circuits, isolating the absorbed-MLA math) matches the Python reference to atol/rtol 1e-4 under xcodebuild. Main-attention rope is traditional=true (distinct from the indexer's interleave=false).
Verifies the absorbed-MLA prefill path WITH sparsification: s=8 > index_topk=4 so the indexer returns a real top-k selection, driving the put_along_axis sparse-mask scatter branch (AND with the causal mask) that S2.1's short-circuit skipped. put_along_axis's scalar-True broadcast works as in Python. Matches the reference to atol/rtol 1e-4 under xcodebuild. The parity test grew a shared assertAttentionParity(fixture:) helper; prefill (S2.1) and sparse-prefill (S2.2) both call it.
Extends DeepseekV32Indexer for incremental decode: rope offset comes from the cache, keys accumulate via cache.update (values an empty head-dim-0 placeholder, mirroring Python's cache.update_and_fetch(k, zeros([b,1,s,0]))), and the short-circuit tests the total (cached) key length. Attention threads cache[1] to the indexer. The cache param defaults to nil so the prefill paths (S2.1 / S2.2) stay byte-for-byte equivalent. Compiles clean; full suite green under xcodebuild (47 tests, attention prefill + sparse + indexer parity all pass). Remaining for S2.3: primed-cache decode fixture + a CacheList two-stage test for the L==1 take_along_axis gather branch (the decode path is written but not yet numerically verified). Plan milestones updated.
…stage prefill+decode)
… (main already removed + gitignored them)
…itize, factory registration Completes the pure-Swift DeepSeek V3.2 port (components #1-12 all green, every stage parity-gated at 1e-4 against the Python mlx-lm reference): - S2.4: dense DeepseekV32DecoderLayer (attn + SwiGLU MLP + rmsNormEps norms + residuals), DeepseekV32MLP, deepseekV32NewCache/KVHeads helpers - S3: DeepseekV32MoEGate (noaux_tc routing with all five V3.2-vs-V3 divergences honored: bias-free gather from origScores, unconditional routedScalingFactor, nGroup==1 skip, no epsilon, float32 sigmoid) + DeepseekV32MoE (SwitchGLU routed experts + optional shared experts) - S4: MoE-vs-dense layer predicate (firstKDenseReplace/moeLayerFreq), DeepseekV32ModelInner/DeepseekV32Model (LLMModel + LoRAModel, newCache override → per-layer CacheList(main MLA, DSA indexer)), sanitize (drop MTP layers, drop fp8 scale keys, stack per-expert w1/w2/w3 into switch_mlp, split absorbed-MLA kv_b_proj → embed_q/unembed_out), ModelOverlay.registerAll() wires model_type "deepseek_v32" into LLMTypeRegistry.shared — zero fork of mlx-swift-lm Tests: full-model parity (2-layer dense+MoE fixture exercises the predicate), sanitize round-trip (Metal-free), 3 registration tests, plus per-component parity suites; capture scripts in docs/reference/.
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
…rlay # Conflicts: # docs/deepseek-v32-swift-port.md
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
This was referenced Jul 8, 2026
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
1 participant
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.
Completes the pure-Swift DeepSeek V3.2 architecture port started in S1: all 12 components green, every stage independently parity-gated at
1e-4against the Python mlx-lm reference (offline fixture capture → SwiftallCloseunder xcodebuild).What lands
CacheList), absorbed MLA, MoE (noaux_tc routing with all five V3.2-vs-V3 router divergences honored), MoE-vs-dense layer predicate, fullDeepseekV32Model(LLMModel/LoRAModel,newCacheoverride),sanitize(MTP drop, fp8-scale drop, expert stacking,kv_b_proj→embed_q/unembed_outsplit).ModelOverlay.registerAll()registersmodel_type: deepseek_v32into the stockLLMTypeRegistry.sharedfrom outside the package — same cross-module pattern Apple ships in MLXVLM. When upstream mlx-swift-lm lands DeepSeek V3.2, we delete our implementation and the registration; the overlay stays a thin, shrinking layer.Verification
Full-model parity on a 2-layer dense+MoE fixture (exercises the layer predicate end-to-end), sanitize round-trip (Metal-free), 3 registration tests, per-component parity suites (indexer / attention prefill+sparse+decode / decoder layer / MoE). Suite: 0 failures under
xcodebuild -skipPackagePluginValidation; adversarially code-reviewed (APPROVE — numerics faithful line-by-line, incl. mask sourcing fromcache[0][0]and the kv_b_proj split axes).Real-checkpoint smoke deferred (needs a downloaded model; the synthetic full-model parity gate stands in until then).
Also stops tracking
docs/superpowers/plans/on this branch (main already removed + gitignored them), avoiding a modify/delete conflict.