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.
- package/README.md +112 -20
- package/docs/design/batching-v2-plan.md +515 -0
- package/docs/design/bucketing-stage.md +196 -0
- package/docs/design/compat-cli-surface-design.md +11 -11
- package/docs/design/diffusion-gemma-port.md +602 -0
- package/docs/design/docs-reorg-plan.md +1 -1
- package/docs/design/dreaming-nightly-pipeline.md +291 -0
- package/docs/design/dspark-speculative-decoding.md +106 -0
- package/docs/design/generic-model-support.md +381 -0
- package/docs/design/memory-inference-path.md +422 -0
- package/docs/design/memory-synthesis.md +287 -0
- package/docs/design/minicpm5-decode-megakernel.md +484 -0
- package/docs/design/mlx-lm-tool-parity-plan.md +535 -0
- package/docs/design/orpo-dynamic-lambda.md +41 -0
- package/docs/design/orpo-training.md +35 -5
- package/docs/design/parallel-slots.md +12 -1
- package/docs/design/parity-tier-dag.md +13 -5
- package/docs/design/spec-decode-larger-targets.md +229 -0
- package/docs/design/tauri-desktop-app.md +2 -2
- package/docs/design/the-dreaming-handoff.md +141 -0
- package/docs/design/the-dreaming-master-plan.md +446 -0
- package/docs/design/write-pipeline-entity.md +468 -0
- package/docs/investigations/curve-bisector-routing.md +1 -1
- package/docs/investigations/decode-roofline-lookagain.md +262 -0
- package/docs/investigations/dspark-handoff.md +58 -0
- package/docs/investigations/kernel-perf-review-2026-07.md +92 -0
- package/docs/planning/PRODUCT_ROADMAP.md +2 -2
- package/docs/planning/curve-sampler-research-plan.md +331 -0
- package/docs/planning/memory-docs-and-dag-plan.md +185 -0
- package/docs/planning/pi-integration-review.md +178 -0
- package/docs/planning/release-notes-v0.0.9.md +122 -0
- package/docs/planning/web-ui-pass-plan.md +213 -0
- package/docs/planning/website-readme-pass-plan.md +224 -0
- package/docs/reference/cli.md +255 -0
- package/docs/reference/embedding.md +7 -2
- package/docs/reference/library-api.md +22 -3
- package/docs/reference/memory.md +5 -5
- package/docs/reference/models.md +137 -0
- package/docs/reference/orpo-quickstart.md +6 -0
- package/docs/reference/server-api.md +187 -9
- package/docs/reference/server-config.md +40 -9
- package/docs/reference/training.md +25 -11
- package/package.json +1 -1
- package/src/anthropic.ts +22 -0
- package/src/assets/curve-designer.html +1 -1
- package/src/chat-template.ts +25 -7
- package/src/cli.ts +1029 -86
- package/src/config.ts +54 -3
- package/src/diffusion/diffusion-generate.ts +436 -0
- package/src/download.ts +99 -8
- package/src/embed.ts +77 -0
- package/src/eval/perplexity.ts +164 -0
- package/src/eval/runner.ts +19 -7
- package/src/eval/tasks/bfcl.ts +1 -1
- package/src/generate.ts +196 -19
- package/src/harness-pi.ts +43 -17
- package/src/index.ts +2 -0
- package/src/lora.ts +2 -2
- package/src/memory/article.ts +558 -0
- package/src/memory/chunk-validation.ts +150 -0
- package/src/memory/chunk.ts +305 -12
- package/src/memory/cluster.ts +344 -17
- package/src/memory/coverage.ts +134 -0
- package/src/memory/crosslink.ts +0 -0
- package/src/memory/db.ts +184 -45
- package/src/memory/entity.ts +203 -0
- package/src/memory/gate.ts +187 -0
- package/src/memory/history.ts +152 -0
- package/src/memory/ingest.ts +324 -11
- package/src/memory/model.ts +321 -0
- package/src/memory/normalize.ts +671 -0
- package/src/memory/parse.ts +22 -0
- package/src/memory/pipeline.ts +242 -30
- package/src/memory/prompts.ts +42 -0
- package/src/memory/query.ts +334 -0
- package/src/memory/reconcile.ts +734 -0
- package/src/memory/redirect.ts +108 -0
- package/src/memory/reindex.ts +325 -0
- package/src/memory/resolve.ts +307 -0
- package/src/memory/route.ts +272 -0
- package/src/memory/schedule.ts +6 -4
- package/src/memory/stages.ts +774 -0
- package/src/memory/synthesize.ts +1629 -20
- package/src/memory/tools.ts +446 -96
- package/src/memory/vault.ts +462 -35
- package/src/memory/wikify.ts +1057 -16
- package/src/mlx/ffi.ts +11 -0
- package/src/mlx/ops.ts +113 -0
- package/src/model/compiled-decode.ts +43 -8
- package/src/model/diffusion-gemma.ts +1026 -0
- package/src/model/factory.ts +35 -8
- package/src/model/flash-attention.ts +4 -2
- package/src/model/fused-decode-kernel.ts +14 -7
- package/src/model/fused-mlp-kernel.ts +169 -0
- package/src/model/fused-swiglu-kernel.ts +240 -0
- package/src/model/gemma4-base.ts +62 -6
- package/src/model/gemma4.ts +33 -2
- package/src/model/minicpm5.ts +106 -8
- package/src/model/qwen3.ts +276 -0
- package/src/model/steel-linear-kernel.ts +96 -0
- package/src/model/support.ts +45 -7
- package/src/model/universal/archs.ts +332 -0
- package/src/model/universal/dense.ts +576 -0
- package/src/model/universal/modules.ts +163 -0
- package/src/model/universal/rope.ts +366 -0
- package/src/pi-provider.ts +1 -1
- package/src/pi-terminal.ts +29 -8
- package/src/pi-web.ts +19 -2
- package/src/quantize/config-writer.ts +1 -1
- package/src/registry.ts +215 -7
- package/src/responses.ts +23 -0
- package/src/sampler.ts +220 -9
- package/src/serve/batch-scheduler.ts +272 -80
- package/src/serve/generation-gateway.ts +90 -18
- package/src/server.ts +643 -50
- package/src/spec/drafter.ts +81 -49
- package/src/spec/dspark/data-dflash.ts +156 -0
- package/src/spec/dspark/data.ts +201 -0
- package/src/spec/dspark/generate-dflash.ts +193 -0
- package/src/spec/dspark/generate.ts +290 -0
- package/src/spec/dspark/loss.ts +213 -0
- package/src/spec/dspark/module-dflash.ts +361 -0
- package/src/spec/dspark/module.ts +540 -0
- package/src/spec/dspark/sample.ts +112 -0
- package/src/spec/generate.ts +40 -19
- package/src/status-page.html +1 -1
- package/src/tokenizer.ts +7 -0
- package/src/train/diffusion-lora.ts +216 -0
- package/src/train/flash-cce.ts +70 -15
- package/src/train/fuse.ts +243 -0
- package/src/train/index.ts +3 -0
- package/src/train/job.ts +8 -0
- package/src/train/loss.ts +289 -30
- package/src/train/prefix-shared.ts +33 -5
- package/src/train/segmented.ts +67 -11
- package/src/train/trainer.ts +92 -28
- package/src/vision/diffusion-vision.ts +620 -0
- 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.
|