mlx-bun 0.0.7 → 0.0.9

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (138) hide show
  1. package/README.md +112 -20
  2. package/docs/design/batching-v2-plan.md +515 -0
  3. package/docs/design/bucketing-stage.md +196 -0
  4. package/docs/design/compat-cli-surface-design.md +11 -11
  5. package/docs/design/diffusion-gemma-port.md +602 -0
  6. package/docs/design/docs-reorg-plan.md +1 -1
  7. package/docs/design/dreaming-nightly-pipeline.md +291 -0
  8. package/docs/design/dspark-speculative-decoding.md +106 -0
  9. package/docs/design/generic-model-support.md +381 -0
  10. package/docs/design/memory-inference-path.md +422 -0
  11. package/docs/design/memory-synthesis.md +287 -0
  12. package/docs/design/minicpm5-decode-megakernel.md +484 -0
  13. package/docs/design/mlx-lm-tool-parity-plan.md +535 -0
  14. package/docs/design/orpo-dynamic-lambda.md +41 -0
  15. package/docs/design/orpo-training.md +35 -5
  16. package/docs/design/parallel-slots.md +12 -1
  17. package/docs/design/parity-tier-dag.md +13 -5
  18. package/docs/design/spec-decode-larger-targets.md +229 -0
  19. package/docs/design/tauri-desktop-app.md +2 -2
  20. package/docs/design/the-dreaming-handoff.md +141 -0
  21. package/docs/design/the-dreaming-master-plan.md +446 -0
  22. package/docs/design/write-pipeline-entity.md +468 -0
  23. package/docs/investigations/curve-bisector-routing.md +1 -1
  24. package/docs/investigations/decode-roofline-lookagain.md +262 -0
  25. package/docs/investigations/dspark-handoff.md +58 -0
  26. package/docs/investigations/kernel-perf-review-2026-07.md +92 -0
  27. package/docs/planning/PRODUCT_ROADMAP.md +2 -2
  28. package/docs/planning/curve-sampler-research-plan.md +331 -0
  29. package/docs/planning/memory-docs-and-dag-plan.md +185 -0
  30. package/docs/planning/pi-integration-review.md +178 -0
  31. package/docs/planning/release-notes-v0.0.9.md +122 -0
  32. package/docs/planning/web-ui-pass-plan.md +213 -0
  33. package/docs/planning/website-readme-pass-plan.md +224 -0
  34. package/docs/reference/cli.md +255 -0
  35. package/docs/reference/embedding.md +7 -2
  36. package/docs/reference/library-api.md +22 -3
  37. package/docs/reference/memory.md +5 -5
  38. package/docs/reference/models.md +137 -0
  39. package/docs/reference/orpo-quickstart.md +6 -0
  40. package/docs/reference/server-api.md +187 -9
  41. package/docs/reference/server-config.md +40 -9
  42. package/docs/reference/training.md +25 -11
  43. package/package.json +1 -1
  44. package/src/anthropic.ts +22 -0
  45. package/src/assets/curve-designer.html +1 -1
  46. package/src/chat-template.ts +25 -7
  47. package/src/cli.ts +1029 -86
  48. package/src/config.ts +54 -3
  49. package/src/diffusion/diffusion-generate.ts +436 -0
  50. package/src/download.ts +99 -8
  51. package/src/embed.ts +77 -0
  52. package/src/eval/perplexity.ts +164 -0
  53. package/src/eval/runner.ts +19 -7
  54. package/src/eval/tasks/bfcl.ts +1 -1
  55. package/src/generate.ts +196 -19
  56. package/src/harness-pi.ts +43 -17
  57. package/src/index.ts +2 -0
  58. package/src/lora.ts +2 -2
  59. package/src/memory/article.ts +558 -0
  60. package/src/memory/chunk-validation.ts +150 -0
  61. package/src/memory/chunk.ts +305 -12
  62. package/src/memory/cluster.ts +344 -17
  63. package/src/memory/coverage.ts +134 -0
  64. package/src/memory/crosslink.ts +0 -0
  65. package/src/memory/db.ts +184 -45
  66. package/src/memory/entity.ts +203 -0
  67. package/src/memory/gate.ts +187 -0
  68. package/src/memory/history.ts +152 -0
  69. package/src/memory/ingest.ts +324 -11
  70. package/src/memory/model.ts +321 -0
  71. package/src/memory/normalize.ts +671 -0
  72. package/src/memory/parse.ts +22 -0
  73. package/src/memory/pipeline.ts +242 -30
  74. package/src/memory/prompts.ts +42 -0
  75. package/src/memory/query.ts +334 -0
  76. package/src/memory/reconcile.ts +734 -0
  77. package/src/memory/redirect.ts +108 -0
  78. package/src/memory/reindex.ts +325 -0
  79. package/src/memory/resolve.ts +307 -0
  80. package/src/memory/route.ts +272 -0
  81. package/src/memory/schedule.ts +6 -4
  82. package/src/memory/stages.ts +774 -0
  83. package/src/memory/synthesize.ts +1629 -20
  84. package/src/memory/tools.ts +446 -96
  85. package/src/memory/vault.ts +462 -35
  86. package/src/memory/wikify.ts +1057 -16
  87. package/src/mlx/ffi.ts +11 -0
  88. package/src/mlx/ops.ts +113 -0
  89. package/src/model/compiled-decode.ts +43 -8
  90. package/src/model/diffusion-gemma.ts +1026 -0
  91. package/src/model/factory.ts +35 -8
  92. package/src/model/flash-attention.ts +4 -2
  93. package/src/model/fused-decode-kernel.ts +14 -7
  94. package/src/model/fused-mlp-kernel.ts +169 -0
  95. package/src/model/fused-swiglu-kernel.ts +240 -0
  96. package/src/model/gemma4-base.ts +62 -6
  97. package/src/model/gemma4.ts +33 -2
  98. package/src/model/minicpm5.ts +106 -8
  99. package/src/model/qwen3.ts +276 -0
  100. package/src/model/steel-linear-kernel.ts +96 -0
  101. package/src/model/support.ts +45 -7
  102. package/src/model/universal/archs.ts +332 -0
  103. package/src/model/universal/dense.ts +576 -0
  104. package/src/model/universal/modules.ts +163 -0
  105. package/src/model/universal/rope.ts +366 -0
  106. package/src/pi-provider.ts +1 -1
  107. package/src/pi-terminal.ts +29 -8
  108. package/src/pi-web.ts +19 -2
  109. package/src/quantize/config-writer.ts +1 -1
  110. package/src/registry.ts +215 -7
  111. package/src/responses.ts +23 -0
  112. package/src/sampler.ts +220 -9
  113. package/src/serve/batch-scheduler.ts +272 -80
  114. package/src/serve/generation-gateway.ts +90 -18
  115. package/src/server.ts +643 -50
  116. package/src/spec/drafter.ts +81 -49
  117. package/src/spec/dspark/data-dflash.ts +156 -0
  118. package/src/spec/dspark/data.ts +201 -0
  119. package/src/spec/dspark/generate-dflash.ts +193 -0
  120. package/src/spec/dspark/generate.ts +290 -0
  121. package/src/spec/dspark/loss.ts +213 -0
  122. package/src/spec/dspark/module-dflash.ts +361 -0
  123. package/src/spec/dspark/module.ts +540 -0
  124. package/src/spec/dspark/sample.ts +112 -0
  125. package/src/spec/generate.ts +40 -19
  126. package/src/status-page.html +1 -1
  127. package/src/tokenizer.ts +7 -0
  128. package/src/train/diffusion-lora.ts +216 -0
  129. package/src/train/flash-cce.ts +70 -15
  130. package/src/train/fuse.ts +243 -0
  131. package/src/train/index.ts +3 -0
  132. package/src/train/job.ts +8 -0
  133. package/src/train/loss.ts +289 -30
  134. package/src/train/prefix-shared.ts +33 -5
  135. package/src/train/segmented.ts +67 -11
  136. package/src/train/trainer.ts +92 -28
  137. package/src/vision/diffusion-vision.ts +620 -0
  138. package/src/web/skills/memory/SKILL.md +40 -17
@@ -0,0 +1,381 @@
1
+ # Generic model support — run anything mlx-lm runs
2
+
3
+ **Status: Phase 1 + 1.5 BUILT (2026-07-01, commit 9bd9f1b)** — `src/model/universal/`
4
+ (UniversalDense + descriptor table for 11 archs, rope factory bit-exact vs oracle
5
+ fixtures, DenseLinear/bf16 checkpoints, weight audit, dispatch ladder in
6
+ factory/support, manifest-driven parity harness). Parity gates light per-model
7
+ once the small checkpoints are downloaded (`scripts/regen-universal-goldens.ts` →
8
+ `MLX_BUN_TEST_UNIVERSAL=1`). Phases 2 (UniversalMoE), 3 (SSM/MLA one-offs),
9
+ 4 (`mlx-bun port` scaffold) remain design.
10
+
11
+ Direction (Josh): *"we should be able to run any model that mlx-lm can run.
12
+ Run any model generically; for the specific models we target, optimize them
13
+ further — push the quality/performance border outwards. Like the gemma
14
+ generator methodology."*
15
+
16
+ The architecture already embodies the two-tier idea for one family:
17
+ `src/model/gemma4-base.ts` (shared machinery) + `gemma4.ts` (config-driven
18
+ monolith) + `src/model/generated/*.ts` (codegen'd specializations via
19
+ `scripts/gen-model.ts`, gated bit-exact by `tests/generated-parity.test.ts`).
20
+ This doc generalizes that shape to the whole mlx-lm zoo:
21
+
22
+ - **Tier-0 GENERIC** — one config-driven "universal" module per architecture
23
+ *family* (dense llama-family first, MoE second). Runs any covered checkpoint
24
+ at **L1 parity** (bit-exact vs mlx-lm on this machine's GPU). Monolith path
25
+ only: no perf kernels, no compiled decode, no kv-quant. Slow, never broken.
26
+ - **Tier-1 TARGETED** — a model *graduates*: dedicated/generated forward,
27
+ optiq-quant **L2** parity, **L3** original kernels, benchmark entries.
28
+ Exactly what gemma4 / minicpm5 / qwen3_5 already are.
29
+
30
+ The fidelity contract is unchanged (three-level tree): **L1 = mlx-lm
31
+ bit-exact is the gate for ANY new model**, generic or targeted. A model the
32
+ generic tier cannot hold to L1 parity is not "supported with caveats" — it is
33
+ unsupported until fixed.
34
+
35
+ ---
36
+
37
+ ## 1. The mlx-lm model zoo, inventoried
38
+
39
+ Source: `mlx-lm 0.31.3` (the oracle venv),
40
+ `site-packages/mlx_lm/models/*.py`. **108 model files** after excluding
41
+ infrastructure (`base`, `cache`, `rope_utils`, `activations`,
42
+ `switch_layers`, `ssm`, `gated_delta`, `mla`, `bitlinear_layers`,
43
+ `pipeline`). `utils.py MODEL_REMAPPING` adds free aliases on top
44
+ (`mistral→llama`, `kimi_k2→deepseek_v3`, `qwen2_5_vl→qwen2_vl`,
45
+ `minimax_m2→minimax`, `iquestcoder→llama`, `falcon_mamba→mamba`, …).
46
+
47
+ ### Buckets (by reading llama.py in full + feature-grep + sampling ~20 files)
48
+
49
+ **A. Dense llama-family — ~36 files, one parameterized module covers them.**
50
+ All are: embed → N × (norm → GQA-attention(+rope) → residual → norm →
51
+ MLP → residual) → norm → head. Every file is a ~170–240-line permutation of
52
+ the same dozen knobs (verified deltas listed in §3.1):
53
+
54
+ > llama (+mistral, iquestcoder via remap), qwen2, qwen3, gemma, gemma2,
55
+ > helium, olmo2, olmo3, smollm3, starcoder2, stablelm, phi3, internlm2,
56
+ > internlm3, granite, exaone, exaone4, glm, glm4, nemotron, minicpm, mimo,
57
+ > ministral3, seed_oss, hunyuan_v1_dense, ernie4_5, cohere, cohere2,
58
+ > apertus, nanochat, lille-130m, telechat3, solar_open, olmo, phi, plamo
59
+
60
+ (`smollm3` is literally a 75-line subclass of `llama` that swaps some
61
+ layers' rope for a no-op; `solar_open` is 38 lines. mlx-lm itself treats
62
+ this family as one module with per-file ModelArgs.)
63
+
64
+ **B. MoE llama-family (SwitchGLU experts) — ~17 files, a second module.**
65
+ Same block, MLP replaced by router + `SwitchGLU` (+ optional shared expert):
66
+
67
+ > mixtral, qwen2_moe, qwen3_moe, olmoe, phimoe, dots1, granitemoe,
68
+ > ernie4_5_moe, glm4_moe, hunyuan, Klear, bailing_moe, exaone_moe, deepseek
69
+ > (v1: MoE without MLA), step3p5, mimo_v2_flash, afmoe
70
+ > (llama4_text / gpt_oss are near-family but carry extra deltas —
71
+ > chunked-attention/NoPE, attention sinks.)
72
+
73
+ **C. MLA lineage (DeepSeek latent attention, usually + MoE + yarn) — ~11:**
74
+ > deepseek_v2, deepseek_v3 (+kimi_k2 remap), deepseek_v32, glm4_moe_lite,
75
+ > glm_moe_dsa, kimi_k25, kimi_vl, longcat_flash(+ngram), minicpm3, youtu_llm
76
+
77
+ **D. SSM / hybrid / linear-attention — ~17, individually shaped:**
78
+ > mamba, mamba2 (+falcon_mamba remap), falcon_h1, jamba, granitemoehybrid,
79
+ > nemotron_h, plamo2, lfm2, lfm2_moe, qwen3_next, kimi_linear, minimax,
80
+ > recurrent_gemma, rwkv7, baichuan_m1, **qwen3_5 (already ported)**,
81
+ > qwen3_5_moe
82
+
83
+ **E. Legacy / exotic one-offs — ~13, deprioritized:**
84
+ > gpt2, gpt_bigcode, gpt_neox, qwen (v1), openelm (per-layer varying dims),
85
+ > phi3small (blocksparse attn), phixtral, dbrx, nemotron-nas (variable
86
+ > blocks), bitnet (bitlinear), afm7, baichuan_m1 (also in D), lille-130m-ish
87
+ > oddities
88
+
89
+ **F. Vision/text wrappers — ~10, ~free once the text base exists:**
90
+ > gemma3, gemma4 (ours already), qwen2_vl (+qwen2_5_vl), qwen3_vl,
91
+ > qwen3_vl_moe, kimi_vl, lfm2-vl, mistral3 (+llava remap), pixtral, llama4
92
+ > — mlx-lm runs only their *text* stack; each wrapper is ~50–120 lines of
93
+ > config unwrapping + weight-prefix stripping.
94
+
95
+ ### Coverage math
96
+
97
+ | Bucket | Files | Unlocked by |
98
+ |---|---|---|
99
+ | A. dense llama-family | ~36 (33%) | **one** UniversalDense module (Phase 1) |
100
+ | B. MoE llama-family | ~17 (16%) | **one** UniversalMoE module (Phase 2) |
101
+ | F. vision wrappers | ~10 (9%) | thin config/prefix shims over A/B (Phase 1–2) |
102
+ | C. MLA lineage | ~11 (10%) | one MLA attention module (Phase 3, demand-driven — most are >24 GB even at 4-bit) |
103
+ | D. SSM/hybrid | ~17 (16%) | individual ports, qwen3_5 pattern (Phase 3) |
104
+ | E. legacy/exotic | ~13 (12%) | individual ports, likely never |
105
+ | already in mlx-bun | 5 | gemma4(+text), qwen3, qwen3_5, minicpm5(llama-typed), diffusion_gemma |
106
+
107
+ **Two generic modules + the wrapper shims ≈ 60% of the zoo by file count,
108
+ and >90% of what actually gets downloaded** (Llama 3.x, Mistral/Ministral,
109
+ Qwen 2.5/3 dense+MoE, Gemma 1–3, Phi-3/4, SmolLM3, OLMo-2/3, GLM-4,
110
+ Granite, Mixtral, OLMoE are all in A/B/F).
111
+
112
+ ---
113
+
114
+ ## 2. What mlx-bun already has to build on
115
+
116
+ Audited 2026-07-01; file references are load-bearing.
117
+
118
+ - **Config parsing** — `src/config.ts loadModelConfig`: already a *normalized*
119
+ config (`TextConfig`) with per-model_type quirk handling inline (gemma
120
+ nested `text_config`, qwen3_5 flat rope dict, llama defaults, diffusion
121
+ backfill). Already carries: layerTypes/slidingWindow, headDim fallback,
122
+ tie, softcap, partialRotaryFactor, `ropeParameters` map, per-layer quant
123
+ map (`quantFor`), kv_config. **This is 70% of the normalization layer;**
124
+ it needs the delta fields from §3.1 and a table-driven (not if-chain)
125
+ quirk mechanism before it drowns.
126
+ - **Primitives** — `src/model/gemma4-base.ts`: `QuantizedLinear` (+LoRA
127
+ residual), `QuantizedEmbedding` (+tied `asLinear`), `RMSNorm`,
128
+ `QuantizedSwitchLinear` (MoE experts — already used by 26B),
129
+ `KVCache`/`RotatingKVCache`(+quantized variants), mask machinery,
130
+ `logitSoftcap`, `disposing`. `src/model/qwen3.ts` (276 lines) is the
131
+ proof-of-shape: a complete dense port using only these primitives.
132
+ `minicpm5.ts`/`qwen3_5.ts` show the "dedicated module conforming to
133
+ RuntimeModel" pattern.
134
+ - **Gating/dispatch** — `src/model/support.ts` (allowlist by model_type +
135
+ config predicates), `src/model/factory.ts` (specific classes →
136
+ fingerprint-matched generated class → monolith; *rejects* unknown types),
137
+ `src/model/fingerprint.ts` (structural hash for generated dispatch).
138
+ - **Generator** — `scripts/gen-model.ts`: emits unrolled, branch-resolved
139
+ `forwardLayers` overrides *of Gemma4Model specifically* (reads the shard
140
+ index for layout facts, folds kv-quant constants, guards cache signature,
141
+ falls back to monolith). Gated by `tests/generated-parity.test.ts`
142
+ (bit-exact vs monolith, fingerprint asserted, fast-path-use asserted).
143
+ - **Parity harness** — the `scripts/regen-*-goldens.ts` +
144
+ `tests/*-parity.test.ts` pattern (e.g. `regen-qwen-parity-goldens.ts`):
145
+ oracle venv emits `<prefix>-parity.json` (explicit prompt ids + greedy
146
+ ids) + per-step raw logit `.bin`s; the test replays ids (no cross-stack
147
+ tokenizer dependency) and compares bit-exact on the same GPU. Plus
148
+ `scripts/check-minicpm5-port.ts` (hard, no-skip port check) and
149
+ `scripts/inspect-model.ts` (tensor name/shape/dtype dump — the seed of
150
+ the weight audit).
151
+ - **Tokenizer** — `src/tokenizer.ts`: `@huggingface/tokenizers` (pure JS)
152
+ reads stock `tokenizer.json` + `tokenizer_config.json`; already
153
+ generic across BPE/SentencePiece-as-tokenizer.json. Known landmine class:
154
+ BOS conventions (the double-BOS bug cost a debug session — see memory);
155
+ chat-template rendering rides Pi, with `/v1` as the boundary.
156
+
157
+ ### Gaps the generic tier must close (found in this audit)
158
+
159
+ 1. **No unquantized Linear.** `QuantizedLinear.load` throws without a
160
+ `.scales` tensor (`gemma4-base.ts:94`). Tier-0 v1 can inherit this
161
+ constraint (target mlx-community 4/8-bit checkpoints — what users pull
162
+ anyway), but a plain `DenseLinear` (bf16/fp16 `ops.matmul`) is a required
163
+ fast-follow: many small models ship only fp16.
164
+ 2. **No additive linear bias.** `QuantizedLinear`'s `biases` are quant
165
+ zero-points, not a bias *term*. qwen2 (attention qkv bias), stablelm,
166
+ starcoder2, phi et al. need an optional post-matmul `+ bias` in the
167
+ linear primitive.
168
+ 3. **Rope variants.** Only default/proportional exist in the runtime.
169
+ mlx-lm `rope_utils.py` also ships: `linear`, `llama3` (Llama3RoPE),
170
+ `yarn` (3 spellings), `longrope`/`su` (SuScaledRoPE), `mrope`. llama3 +
171
+ yarn + longrope are hard requirements for bucket A (Llama 3.x, Phi-3,
172
+ DeepSeek-lineage, various long-context finetunes).
173
+ 4. **Attention-score softcap** (gemma2): needs the manual
174
+ (q·kᵀ → tanh-cap → softmax) path — `ops.sdpa` can't express it.
175
+ 5. **factory.ts rejects unknown types** rather than attempting a generic
176
+ fallback (by design today; changes in §3.3).
177
+
178
+ ---
179
+
180
+ ## 3. Design
181
+
182
+ ### 3.1 Tier-0: `UniversalDense` (src/model/universal/)
183
+
184
+ One module, shaped exactly like `qwen3.ts` but every delta a config field.
185
+ **The normalized config is the union of mlx-lm's per-arch ModelArgs**, and a
186
+ declarative **arch descriptor table** supplies per-model_type defaults +
187
+ config-key aliases — mirroring mlx-lm's structure (shared base module,
188
+ per-arch args) without one file per arch:
189
+
190
+ ```ts
191
+ // src/model/universal/archs.ts — ~15–25 lines per arch vs mlx-lm's ~200
192
+ export interface ArchDescriptor {
193
+ // block structure
194
+ block: "sequential" | "sandwich" | "parallel"; // llama | gemma2/glm4 | cohere/phi
195
+ norm: { kind: "rmsnorm" | "layernorm"; plusOne?: boolean; bias?: boolean };
196
+ qkNorm: "none" | "head-rms" | "layernorm-2d"; // qwen3/olmo2 | cohere
197
+ mlp: { kind: "gated" | "plain"; act: "silu" | "gelu" | "gelu_tanh" | "relu2" };
198
+ fusedQkv?: boolean; fusedGateUp?: boolean; // phi3 — split in sanitize
199
+ // scalars (all read from config.json when present; descriptor = default)
200
+ attnLogitSoftcap?: number; // gemma2
201
+ logitScale?: ConfigPath; // cohere (multiply), granite logits_scaling (divide)
202
+ embedScale?: "sqrt-hidden" | "mup" | null; // gemma family, minicpm
203
+ residualMultiplier?: ConfigPath; // granite
204
+ bias: { qkv?: boolean; o?: boolean; mlp?: boolean };
205
+ nope?: "interval" | "list"; // smollm3, llama4-style no-rope layers
206
+ // weights
207
+ weightMap?: Partial<Record<CanonicalName, string>>; // starcoder2 c_fc/c_proj, gpt2 h.N.*
208
+ sanitize?: (names: string[]) => Rename[]; // fused-proj splits, inv_freq drops
209
+ configAliases?: Record<string, string>; // norm_epsilon→rms_norm_eps etc.
210
+ }
211
+ export const ARCHS: Record<string /* model_type */, ArchDescriptor> = { ... };
212
+ ```
213
+
214
+ Key choices:
215
+
216
+ - **Descriptor keyed by `model_type` + mlx-lm's MODEL_REMAPPING table
217
+ vendored in.** No config sniffing/fingerprinting to *identify* an arch —
218
+ fingerprints stay what they are today: a dispatch guard for generated
219
+ specializations. Explicit table = explicit support surface = honest
220
+ `mlx-bun models` output.
221
+ - **Transcribe mlx-lm op-for-op** inside the universal forward, per
222
+ descriptor branch — same discipline as every existing port (q/k-norm
223
+ before vs after transpose, rope order, cast points). The descriptor
224
+ selects branches; it never approximates.
225
+ - **Rope factory** `src/model/universal/rope.ts` ports
226
+ `rope_utils.initialize_rope` verbatim: default/linear (native
227
+ `ops.rope` with scale), llama3 (precomputed freqs), yarn, longrope/su.
228
+ This is the single riskiest chunk of Phase 1 — rope bugs are
229
+ off-by-a-little and only caught by logit goldens, which is why the parity
230
+ gate feeds explicit ids and compares raw logits per step.
231
+ - **KV cache**: reuse `KVCache`/`RotatingKVCache` unchanged;
232
+ `layer_types`/`sliding_window` already normalized. No kv-quant at Tier-0
233
+ (that's an L2/graduation concern).
234
+ - **Quantization**: v1 requires MLX-quantized checkpoints (`.scales`
235
+ present), same as the rest of the runtime; `DenseLinear` for
236
+ fp16/bf16 checkpoints is Phase 1.5. Add the additive-bias term to both.
237
+ - **RuntimeModel surface**: implement the same minimal contract
238
+ `Qwen3Model` does (config, weightsBytes, loraState, makeCache,
239
+ forwardHidden, logitsFromHidden, forward, generate, loraTargets) so
240
+ serving/eval/LoRA hot-swap work day one. `forwardEmbeddings` throws
241
+ (vision is per-arch work).
242
+
243
+ ### 3.2 Tier-0.5: `UniversalMoE`
244
+
245
+ Same block with the MLP swapped for router + `QuantizedSwitchLinear`
246
+ (already exists and is serving 26B-A4B). Deltas to parameterize: top-k,
247
+ `norm_topk_probs`, routing score fn (softmax/sigmoid), routing scale,
248
+ shared-expert branch (+ optional shared-expert gate — qwen2_moe),
249
+ first-k-dense-replace (deepseek/glm4_moe layouts), grouped routing
250
+ (deepseek_v3-style, defer). gpt_oss (attention sinks) and llama4 (chunked
251
+ attention) are *not* claimed by the generic MoE — they graduate to
252
+ one-offs if wanted.
253
+
254
+ Memory reality check on the reference machines (24–32 GB): Mixtral-8x7B
255
+ 4-bit ≈ 25 GB — no. OLMoE-1B-7B, Qwen3-30B-A3B (≈17 GB @4bit),
256
+ GLM-4 lite variants, granitemoe small — yes. The generic MoE tier is
257
+ worth building, but expectations belong in `mlx-bun models` output
258
+ (weightsBytes vs RAM), not in marketing.
259
+
260
+ ### 3.3 Dispatch changes (support.ts / factory.ts)
261
+
262
+ Order matters — **generic never shadows a dedicated port**:
263
+
264
+ ```
265
+ createModel(weights, config):
266
+ 1. dedicated classes (diffusion, minicpm5, qwen3_5, qwen3) — unchanged
267
+ 2. gemma4*: GENERATED.get(fingerprint) ?? Gemma4Model — unchanged
268
+ 3. ARCHS[remap(config.modelType)] → new UniversalDense/MoE(...) — NEW
269
+ 4. throw "unsupported model_type X (nearest: …)" — unchanged shape
270
+ ```
271
+
272
+ - `support.ts isSupportedModelRecord`: add `GENERIC_MODEL_TYPES` (the
273
+ descriptor table's keys, post-remap). Keep the drafter exclusion.
274
+ `mlx-bun models` labels the tier: `supported (targeted)` vs
275
+ `supported (generic)` — sets perf expectations and doubles as the
276
+ graduation worklist.
277
+ - The current hard throw for non-MiniCPM5 `llama` configs is *replaced* by
278
+ the generic path (that throw exists only because nothing generic existed).
279
+ `isMiniCPM5Config`'s exact-shape predicate keeps MiniCPM5 on its
280
+ dedicated class ahead of the generic.
281
+
282
+ ### 3.4 Weight-name mapping + load audit
283
+
284
+ Bucket A is already ~uniform in HF land
285
+ (`model.layers.N.self_attn.{q,k,v,o}_proj`, `mlp.{gate,up,down}_proj`,
286
+ `input_layernorm`, `post_attention_layernorm`, `model.norm`,
287
+ `model.embed_tokens`, optional `lm_head`). The descriptor handles the
288
+ exceptions: rename maps (starcoder2 `c_fc/c_proj`), fused-proj split rules
289
+ (phi3 `qkv_proj`, `gate_up_proj` — split at load into views/copies),
290
+ prefix strips (vision wrappers `language_model.model.*` — gemma4 already
291
+ does this via `prefixBase`), and drop rules (`rotary_emb.inv_freq`,
292
+ mlx-lm `sanitize`).
293
+
294
+ **Every generic load runs a weight audit** (grow `scripts/inspect-model.ts`
295
+ into `src/model/universal/audit.ts`): after module construction, diff
296
+ {tensors in shard index} vs {tensors consumed}. Unconsumed or missing
297
+ tensors ⇒ load error naming them, not a silent wrong-answer model. This is
298
+ the cheap tripwire that catches descriptor mistakes before parity does
299
+ (the "verify the foundation" rule, mechanized).
300
+
301
+ ### 3.5 The parity gate for a NEW arch (cost: ~2–4 h/arch, mostly waiting)
302
+
303
+ Pattern already proven three times (gemma/minicpm5/qwen). Parameterize it
304
+ once:
305
+
306
+ 1. `scripts/regen-generic-goldens.ts <snapshot> <prefix>` — oracle venv
307
+ (`/Users/joshrossi/Code/mlx-lm/.venv/bin/python`, mlx-lm `load()` +
308
+ `make_prompt_cache`) emits `<prefix>-parity.json` (explicit prompt ids,
309
+ greedy ids) + 12 step-logit `.bin`s. Same-machine rule applies (logit
310
+ goldens are GPU-specific — label goldens with host).
311
+ 2. `tests/generic-parity.test.ts` — iterates a manifest
312
+ (`tests/generic-manifest.ts`: prefix → snapshot path → model_type),
313
+ skipIf-gated on snapshot presence like every parity test. Asserts:
314
+ fingerprint-free construction, weight audit clean, greedy trajectory
315
+ match, per-step logits bit-exact.
316
+ 3. Checkpoint choice: **smallest mlx-community 4-bit of the arch**
317
+ (0.5B–3B exists for nearly all of bucket A) — keeps the gate runnable
318
+ and the download small. Downloads stay a user action (ground rules).
319
+
320
+ Definition of supported-generic = manifest entry green on the current
321
+ machine. No entry, no support claim.
322
+
323
+ ### 3.6 Graduation to Tier-1 (targeted)
324
+
325
+ **Criteria** (all three): (a) L1 generic parity green and stable across a
326
+ few sessions of real use; (b) it's a daily driver / benchmark / product
327
+ target (e.g. the Lucien pipeline models); (c) measured perf or memory gap
328
+ vs a dedicated path that's worth the maintenance (number in
329
+ `benchmarks/RESULTS.md`, per the every-perf-claim-gets-a-number rule).
330
+
331
+ **Process** (the gemma methodology, in order — stop at any step):
332
+ 1. **L2**: optiq-quant goldens (`regen-*-kv-goldens` pattern, mixed-KV vs
333
+ the optiq oracle) if an optiq checkpoint/kv_config exists for it.
334
+ 2. **Dedicated or generated forward**: if the arch is gemma4-base-shaped,
335
+ extend `gen-model.ts`; otherwise hand-port a dedicated module
336
+ (`qwen3_5.ts` pattern). Note: today's `gen-model.ts` subclasses
337
+ `Gemma4Model` specifically — generalizing it to emit from
338
+ `UniversalDense` is Phase 4 work, and *optional*: hand ports are proven
339
+ and cheap ("faithful port took ~1 hr, not the huge effort claimed").
340
+ 3. **L3 kernels**: fused decode / compiled decode / kv-quant / flash paths,
341
+ each behind its existing flag, each frozen-oracle-gated.
342
+ 4. RESULTS.md + docs entries; `mlx-bun models` label flips to `targeted`.
343
+
344
+ Nothing about graduation removes the generic path — it stays as the
345
+ always-correct fallback, exactly like the monolith does for generated
346
+ gemma classes today.
347
+
348
+ ---
349
+
350
+ ## 4. Phased rollout
351
+
352
+ | Phase | Scope | Unlocks | Effort | Key risks |
353
+ |---|---|---|---|---|
354
+ | **1** | `UniversalDense` + descriptor table + rope factory (llama3/yarn/longrope) + bias term + goldens harness. Launch archs: **llama (⇒ Llama 3.x + Mistral + Ministral), qwen2, qwen3(fold-in), gemma, gemma2, smollm3, olmo2, olmo3, phi3, glm4, granite, starcoder2** | ~20 files at launch; the rest of bucket A (~36 total) as S-sized descriptor+golden follow-ons; F-wrappers for covered bases | **L** (module M, rope M, harness S, 12×goldens S-each) | rope variants (llama3/yarn/longrope math), gemma2 softcap manual-attn path, phi3 fused-proj split, BOS/tokenizer conventions per family (the double-BOS class of bug), qwen2 qkv bias |
355
+ | **1.5** | `DenseLinear` (bf16/fp16 checkpoints, no `.scales`) | fp16-only small models; removes the "MLX-quantized only" constraint | **S–M** | dtype/cast points vs mlx-lm (bf16 matmul accumulation — see e4b length-sensitivity memory) |
356
+ | **2** | `UniversalMoE` (router + SwitchGLU + shared experts + first-k-dense) : **mixtral, qwen2_moe, qwen3_moe, olmoe, glm4_moe, dots1, granitemoe, ernie4_5_moe, phimoe** | ~12–15 of bucket B (gpt_oss/llama4 excluded) | **M** (SwitchGLU exists; routing variants are config) | routing-score parity (softmax vs sigmoid + norm_topk), RAM honesty on 24 GB |
357
+ | **3** | Demand-driven one-offs, qwen3_5 pattern: SSM/hybrid (**lfm2, qwen3_next, mamba2, nemotron_h**…), MLA module if a runnable target appears (**minicpm3**; deepseek-lineage mostly exceeds local RAM) | bucket C/D picks, ~1 model per effort unit | **M–L each** (cumulatively XL; never a blocker for A/B) | new cache types (conv/ssm state) touch the Cache interface + serving slots; each is a real port with its own goldens |
358
+ | **4** | Tooling: **`mlx-bun port <hf-repo>`** — download smallest snapshot → descriptor coverage check (config keys + weight-audit dry run) → auto-regen goldens vs oracle venv → run parity gate → print verdict + manifest entry. Optionally: generalize `gen-model.ts` to emit from UniversalDense for graduation | turns "new arch" into a ~1-command, ~1-hour operation; makes the coverage claim self-verifying | **M** (scaffold) + **L** (generator generalization, optional) | oracle-venv coupling (mlx-lm version drift changes goldens — pin + record versions in parity.json, already the pattern) |
359
+
360
+ Sequencing note: Phase 1's descriptor table should be *born* with only the
361
+ 12 launch archs actually gated green — adding an arch without its manifest
362
+ goldens is the one anti-pattern this design forbids.
363
+
364
+ ## 5. Top design decisions (summary)
365
+
366
+ 1. **Explicit descriptor table keyed by model_type (+ vendored
367
+ MODEL_REMAPPING), not config sniffing** — support is a declared,
368
+ testable surface; fingerprints remain a generated-dispatch guard only.
369
+ 2. **One normalized config** — extend the existing `TextConfig` (it's
370
+ already 70% there) with the bucket-A delta union; move per-arch quirks
371
+ out of `loadModelConfig`'s if-chain into descriptor `configAliases`.
372
+ 3. **L1 bit-exact is the gate for generic too** — manifest-driven golden
373
+ harness (explicit ids, raw per-step logits, same-GPU), one regen script
374
+ + one test file for all generic archs.
375
+ 4. **Weight audit at load** — unconsumed/missing tensors are a load error,
376
+ never a silently-wrong model.
377
+ 5. **Generic never shadows targeted; graduation never deletes generic** —
378
+ the dispatch ladder is dedicated → generated → generic → reject, and the
379
+ generic path is the permanent correctness fallback.
380
+ 6. **Honest tiering in UX** — `mlx-bun models` shows `targeted` vs
381
+ `generic` (and RAM fit), which doubles as the graduation worklist.