mlx-bun 0.0.10 → 0.0.11
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 +66 -21
- package/THIRD_PARTY_LICENSES.md +1 -0
- package/docs/archive/mlx-bun-lab-report.html +275 -0
- package/docs/archive/wwdc-mlx-bun.html +848 -0
- package/docs/design/audio-input-plan.md +356 -0
- package/docs/design/batching-perf-path.md +58 -35
- package/docs/design/batching-v2-plan.md +23 -9
- package/docs/design/bucketing-stage.md +2 -1
- package/docs/design/compat-cli-surface-design.md +4 -3
- package/docs/design/decode-speed-program.md +172 -0
- package/docs/design/docs-reorg-plan.md +1 -1
- package/docs/design/dspark-serving-program.md +371 -0
- package/docs/design/dspark-speculative-decoding.md +167 -5
- package/docs/design/faithful-l1-consolidation.md +310 -0
- package/docs/design/generic-model-support.md +8 -5
- package/docs/design/grammar-spec-batching-integration.md +301 -0
- package/docs/design/minicpm5-decode-megakernel.md +16 -11
- package/docs/design/mlx-lm-tool-parity-plan.md +11 -1
- package/docs/design/omlx-adoption-map.md +14 -6
- package/docs/design/orpo-training.md +34 -30
- package/docs/design/paged-kv-cache.md +214 -0
- package/docs/design/parallel-slots.md +82 -49
- package/docs/design/parity-tier-dag.md +43 -33
- package/docs/design/repo-cleanup-plan.md +189 -0
- package/docs/design/runtime-isolation.md +19 -3
- package/docs/design/segmented-backward-training.md +3 -2
- package/docs/design/ssd-kv-cold-tier.md +130 -0
- package/docs/design/structured-output.md +484 -0
- package/docs/design/the-dreaming-master-plan.md +2 -3
- package/docs/design/turboquant-kv.md +235 -0
- package/docs/design/unified-engine-frontier-plan.md +774 -0
- package/docs/design/web-chat-beat-matrix.md +315 -0
- package/docs/design/web-chat-redesign.md +1112 -0
- package/docs/investigations/dspark-handoff.md +182 -16
- package/docs/investigations/lab-build-journal.md +2 -2
- package/docs/investigations/openwiki-evaluation.md +249 -0
- package/docs/planning/IDEAS.md +14 -0
- package/docs/planning/PRODUCT_ROADMAP.md +2 -3
- package/docs/planning/chat-ui-vision-memory-first.md +199 -0
- package/docs/planning/chat-ui-vision.md +50 -0
- package/docs/planning/memory-docs-and-dag-plan.md +15 -6
- package/docs/planning/release-notes-v0.0.11.md +239 -0
- package/docs/reference/cli.md +35 -16
- package/docs/reference/distribution.md +13 -6
- package/docs/reference/embedding.md +35 -14
- package/docs/reference/features-matrix.md +97 -0
- package/docs/reference/library-api.md +30 -9
- package/docs/reference/memory.md +26 -6
- package/docs/reference/models.md +11 -1
- package/docs/reference/orpo-quickstart.md +3 -4
- package/docs/reference/server-api.md +634 -22
- package/docs/reference/server-config.md +250 -121
- package/docs/reference/training.md +14 -15
- package/package.json +5 -3
- package/src/anthropic.ts +13 -0
- package/src/audio/conformer.ts +647 -0
- package/src/audio/decode.ts +133 -0
- package/src/audio/features.ts +202 -0
- package/src/audio/transcode.ts +67 -0
- package/src/cli.ts +356 -105
- package/src/config.ts +49 -2
- package/src/download.ts +40 -0
- package/src/eval/README.md +7 -6
- package/src/eval/kl.ts +59 -3
- package/src/eval/quality-db.ts +2 -2
- package/src/eval/runner.ts +3 -3
- package/src/expert-offload.ts +15 -2
- package/src/flags.ts +13 -0
- package/src/generate.ts +330 -24
- package/src/grammar.ts +585 -0
- package/src/hub-rest.ts +242 -0
- package/src/kv-store.ts +442 -88
- package/src/lora.ts +8 -0
- package/src/md-modules.d.ts +42 -0
- package/src/media-fetch.ts +284 -0
- package/src/memory/rest.ts +245 -0
- package/src/memory/vault.ts +74 -0
- package/src/mlx/array.ts +106 -19
- package/src/mlx/autograd.ts +2 -2
- package/src/mlx/ffi.ts +27 -3
- package/src/mlx/metal-kernel.ts +7 -4
- package/src/mlx/ops.ts +153 -23
- package/src/mlx/turboquant-ops.ts +831 -0
- package/src/mlx/turboquant-tables.ts +237 -0
- package/src/model/batched-mask.ts +62 -1
- package/src/model/batched-quant.ts +206 -0
- package/src/model/batched-rotating-quant.ts +346 -0
- package/src/model/batched-rotating.ts +46 -1
- package/src/model/compiled-decode.ts +4 -15
- package/src/model/factory.ts +10 -4
- package/src/model/gemma4-base.ts +363 -39
- package/src/model/gemma4.ts +122 -33
- package/src/model/generated/gemma4-12b.ts +9 -13
- package/src/model/generated/gemma4-26b.ts +10 -5
- package/src/model/generated/gemma4-e4b.ts +14 -9
- package/src/model/minicpm5.ts +58 -102
- package/src/model/paged-kv.ts +253 -0
- package/src/model/qwen3-delta.ts +76 -15
- package/src/model/qwen3-moe.ts +404 -0
- package/src/model/qwen3.ts +9 -8
- package/src/model/qwen3_5.ts +88 -23
- package/src/model/support.ts +12 -1
- package/src/model/universal/dense.ts +40 -5
- package/src/model/universal/rope.ts +36 -21
- package/src/pi-web.ts +1299 -82
- package/src/preflight.ts +27 -8
- package/src/prompt-cache.ts +293 -33
- package/src/quantize/quantizer.ts +16 -1
- package/src/registry.ts +57 -2
- package/src/serve/batch-scheduler.ts +965 -72
- package/src/serve/generation-gateway.ts +214 -11
- package/src/serve/isolate.ts +417 -0
- package/src/serve/lane-registry.ts +44 -0
- package/src/serve/session-search.ts +326 -0
- package/src/server.ts +1380 -135
- package/src/spec/assistant-source.ts +150 -0
- package/src/spec/deepspec-source.ts +143 -0
- package/src/spec/dflash-source.ts +124 -0
- package/src/spec/dspark/ab-stats.ts +136 -0
- package/src/spec/dspark/calibration.ts +112 -0
- package/src/spec/dspark/deepspec-module.ts +801 -0
- package/src/spec/dspark/generate-dflash.ts +72 -20
- package/src/spec/dspark/loader.ts +23 -0
- package/src/spec/dspark/module-dflash.ts +278 -36
- package/src/spec/dspark/quantize-drafter.ts +91 -0
- package/src/spec/generate.ts +59 -25
- package/src/spec/ngram-source.ts +132 -0
- package/src/spec/serve-loop.ts +536 -0
- package/src/spec/source.ts +93 -0
- package/src/spec/two-model.ts +171 -0
- package/src/ssd-cache.ts +107 -35
- package/src/tokenizer.ts +288 -2
- package/src/tool-approvals.ts +131 -0
- package/src/tool-call.ts +408 -20
- package/src/train/flash-cce.ts +2 -2
- package/src/train/job.ts +0 -2
- package/src/train/trainer.ts +1 -4
- package/src/vision/prompt.ts +217 -48
- package/src/web/app.html +1184 -1789
- package/src/web/app.js +5667 -0
- package/src/web/icon.svg +13 -0
- package/src/web/manifest.webmanifest +24 -0
- package/src/web/src/adapters-panel.ts +288 -0
- package/src/web/src/api.ts +57 -0
- package/src/web/src/assistant.ts +485 -0
- package/src/web/src/chat.ts +1050 -0
- package/src/web/src/composer.ts +959 -0
- package/src/web/src/dataset.ts +132 -0
- package/src/web/src/finetune.ts +267 -0
- package/src/web/src/hub.ts +428 -0
- package/src/web/src/main.ts +98 -0
- package/src/web/src/markdown.ts +0 -0
- package/src/web/src/memory-panel.ts +572 -0
- package/src/web/src/model-picker.ts +159 -0
- package/src/web/src/palette.ts +346 -0
- package/src/web/src/protocol.ts +62 -0
- package/src/web/src/quantize.ts +167 -0
- package/src/web/src/rag.ts +274 -0
- package/src/web/src/sessions.ts +469 -0
- package/src/web/src/shell.ts +804 -0
- package/src/web/src/status.ts +252 -0
- package/src/web/src/ui-catalog.ts +92 -0
- package/src/web/sw.js +71 -0
- package/src/web/vendor/README.md +61 -0
- package/src/web/vendor/hljs-LICENSE +29 -0
- package/src/web/vendor/hljs-theme.css +59 -0
- package/src/web/vendor/hljs.js +26 -0
- package/src/model/fused-decode-kernel.ts +0 -312
- package/src/model/fused-geglu-kernel.ts +0 -198
- package/src/model/fused-mlp-kernel.ts +0 -169
- package/src/model/fused-swiglu-kernel.ts +0 -240
- package/src/model/steel-linear-kernel.ts +0 -96
- package/src/status-page.html +0 -497
|
@@ -0,0 +1,214 @@
|
|
|
1
|
+
# Optional paged KV cache (vLLM-style) — v1 design
|
|
2
|
+
|
|
3
|
+
Status: **v1 in progress** (this PR)
|
|
4
|
+
Default: **off** (`--paged-kv` unset ⇒ today's `KVCache`/`RotatingKVCache`
|
|
5
|
+
path, byte-identical — the flag-off proof is that zero new code executes).
|
|
6
|
+
|
|
7
|
+
Plan produced by a multi-agent design pass (4 subsystem readers +
|
|
8
|
+
prior-art scout → 3 competing designs → judged synthesis), then vetted
|
|
9
|
+
against the code by hand. Judge's pick: minimal-seam (8.25) over
|
|
10
|
+
serving-first (6.5) and future-proof (7.25); build-order discipline
|
|
11
|
+
grafted from future-proof.
|
|
12
|
+
|
|
13
|
+
## Motivation — what paged buys here, honestly
|
|
14
|
+
|
|
15
|
+
vLLM's PagedAttention (Kwon et al., SOSP 2023, arXiv 2309.06180) solves
|
|
16
|
+
contiguous worst-case per-sequence KV reservation + allocator
|
|
17
|
+
fragmentation. mlx-bun is a single process on unified memory with
|
|
18
|
+
byte-budget admission already shipped (parallel-slots.md rung 2), so the
|
|
19
|
+
fragmentation story mostly doesn't transfer. What does:
|
|
20
|
+
|
|
21
|
+
- **Padded-batch waste removal** — today's batched buffer width tracks
|
|
22
|
+
the LONGEST live row (`mergeKVRows`/`extendKVRows`,
|
|
23
|
+
src/model/batched-mask.ts): a 50-token row batched with a 4000-token
|
|
24
|
+
row pays ~4000 tokens of KV while they cohabit. Block-granular
|
|
25
|
+
allocation removes this. **A batching feature; v1 does NOT implement
|
|
26
|
+
it** (non-goal — follow-up PR).
|
|
27
|
+
- **Block-level prefix sharing / copy-on-write** — PromptCache clones at
|
|
28
|
+
entry granularity; block CoW is vLLM's headline win for concurrent
|
|
29
|
+
agents sharing a system prompt. **Not in v1.**
|
|
30
|
+
- **At batch=1, paging buys nothing** — gather-before-attend is pure
|
|
31
|
+
bandwidth tax on a bandwidth-bound decode. Measured and reported
|
|
32
|
+
plainly, never spun.
|
|
33
|
+
|
|
34
|
+
v1 therefore ships the correctness-proven block-manager abstraction at
|
|
35
|
+
serial (batch=1) scope, expecting a small decode regression and zero
|
|
36
|
+
density win — scaffolding for the follow-ups, not a standalone perf win.
|
|
37
|
+
parallel-slots.md already names paged KV "rung 3 / S3+ deferred";
|
|
38
|
+
unified-engine-frontier-plan.md names block-paged KV the long-term
|
|
39
|
+
unifier. This PR is the first, narrowest slice.
|
|
40
|
+
|
|
41
|
+
Prior-rejection reconciliation (the record already says "no paging"
|
|
42
|
+
twice, for different axes, and both stand): the 2026-07-07 vLLM-mining
|
|
43
|
+
pass rejected paged KV as a PROMPT-CACHE substitute (take()'s zero-copy
|
|
44
|
+
entry clones already share physically), and ssd-kv-cold-tier.md D1
|
|
45
|
+
rejected paged BLOCKS as the SSD-spill granularity (per-block hashing on
|
|
46
|
+
the token loop = oMLX's ~20% tax). This PR is neither: it's the rung-3
|
|
47
|
+
ALLOCATION abstraction (parallel-slots.md), explicitly requested as an
|
|
48
|
+
optional flag, with the padded-batch/CoW payoff deferred to follow-ups.
|
|
49
|
+
|
|
50
|
+
**No external oracle exists**: mlx-lm's cache.py has no paged/block
|
|
51
|
+
cache (verified in the oracle venv). The correctness gate is mlx-bun's
|
|
52
|
+
own plain `KVCache` — valid because v1's claim is storage-layout
|
|
53
|
+
equivalence (same bytes, different physical arrangement), not new
|
|
54
|
+
arithmetic. Bit-exactness (tol 0) is the bar.
|
|
55
|
+
|
|
56
|
+
## v1 scope
|
|
57
|
+
|
|
58
|
+
- Serial (non-batched) single-sequence decode; Gemma4-family models
|
|
59
|
+
(modelType `gemma4*`), plain full-attention layers only. Sliding
|
|
60
|
+
(`RotatingKVCache`) layers keep today's scheme untouched — mixed
|
|
61
|
+
per-layer paged-full + rotating-sliding is the supported shape.
|
|
62
|
+
- bf16 only. No `--kv-quant`/TurboQuant composition.
|
|
63
|
+
- Strategy: **block pool + gather-to-contiguous, then the unchanged
|
|
64
|
+
`ops.sdpa()`** (gemma4.ts Attention.forward's plain-cache `else`
|
|
65
|
+
branch consumes any `Cache` whose updateAndFetch returns contiguous
|
|
66
|
+
[1,H,S,D] views). No new attention math, no new kernel, no new branch
|
|
67
|
+
in the forward.
|
|
68
|
+
- Fixed pool per request, sized up front from prompt length + maxTokens
|
|
69
|
+
(capacity is known at generate() time) — exhaustion is a typed
|
|
70
|
+
`PagedPoolExhausted` error, never silent truncation or OOM.
|
|
71
|
+
|
|
72
|
+
## Non-goals (v1)
|
|
73
|
+
|
|
74
|
+
- No `--batch N` integration (`--paged-kv` + `--batch N>1` refused at
|
|
75
|
+
startup). No `LayerInner` union change in batch-scheduler.ts.
|
|
76
|
+
- No sliding-window paging; no quantized blocks; no CoW/prefix sharing;
|
|
77
|
+
no MoE/SSM/universal/CPM/Diffusion models; no vision/audio; no
|
|
78
|
+
LoRA/adapters; no speculative decode (spec eligibility excludes paged).
|
|
79
|
+
- No fused paged-attention Metal kernel (follow-up once the gather tax
|
|
80
|
+
is measured to matter; metal-kernel.ts infra + optiq's
|
|
81
|
+
sdpa_2pass_paged are the port sources).
|
|
82
|
+
- No compiled-decode composability: a data-dependent block-list length
|
|
83
|
+
is the exact shape shapeless replay already broke on (tiled quantized
|
|
84
|
+
SDPA precedent). `CompiledDecode.supports()` excludes PagedKVCache
|
|
85
|
+
automatically (it's not one of the four supported classes); the
|
|
86
|
+
serve/CLI layer also refuses the explicit combination.
|
|
87
|
+
- No PromptCache integration: paged requests bypass take/put (the
|
|
88
|
+
vision precedent in server.ts runGeneration) and dispose caches on
|
|
89
|
+
completion. `state()` exists only for evalCacheState.
|
|
90
|
+
|
|
91
|
+
## Design
|
|
92
|
+
|
|
93
|
+
New file `src/model/paged-kv.ts`:
|
|
94
|
+
|
|
95
|
+
- **`BlockPool`** — per-layer arena. K and V pool tensors shaped
|
|
96
|
+
`[numBlocks, H_kv, blockSize, headDim]`, allocated via `ops.zeros`
|
|
97
|
+
(mlx-owned end to end — no host-pointer alignment/dtor hazards).
|
|
98
|
+
`alloc()` pops the free list or throws `PagedPoolExhausted`;
|
|
99
|
+
`free(idx)` pushes back. Incoming K/V write via the existing
|
|
100
|
+
`ops.sliceUpdate` at `[block, 0, within, 0]` — the incoming
|
|
101
|
+
`[1,H,l,D]` piece matches the destination slice directly, no
|
|
102
|
+
transpose.
|
|
103
|
+
- **`PagedKVCache implements Cache`** — standalone implementation
|
|
104
|
+
(deliberately NOT a `KVCache` subclass, same reasoning as
|
|
105
|
+
`TurboQuantKVCache`: every instanceof gate — CompiledDecode.supports,
|
|
106
|
+
gateway #modelCachesBatchable, generated `#matches()` — then excludes
|
|
107
|
+
it automatically and routes to the monolith forward).
|
|
108
|
+
- `updateAndFetch(k, v)`: split the incoming L along block boundaries,
|
|
109
|
+
sliceUpdate each piece into pool blocks (allocating from the free
|
|
110
|
+
list as the tail block fills), then fetch: `ops.takeAxis(pool,
|
|
111
|
+
blockTable, 0)` → `[nb,H,bs,D]`, transpose `[1,0,2,3]` →
|
|
112
|
+
`[H,nb,bs,D]`, reshape `[1,H,nb*bs,D]` (the gather copy), slice to
|
|
113
|
+
offset. **Uses only existing bound ops** — the plan's proposed
|
|
114
|
+
`mlx_gather` FFI binding is unnecessary (`takeAxis` gathers whole
|
|
115
|
+
blocks along axis 0), which deletes the bun:ffi stack-arg ABI risk
|
|
116
|
+
outright.
|
|
117
|
+
- `makeMask`: same logic as `KVCache.makeMask` (reads only offset).
|
|
118
|
+
- `trim(n)`: rewind offset, free now-unoccupied tail blocks. Stale
|
|
119
|
+
bytes past offset are never read (fetch slices to offset; writes
|
|
120
|
+
cover exact ranges) — same invariant as KVCache's padding.
|
|
121
|
+
- Block table is host-side `number[]` — no device indirection.
|
|
122
|
+
- All intermediates bound to locals and disposed
|
|
123
|
+
(mlx-inline-slice-leak-pattern discipline).
|
|
124
|
+
|
|
125
|
+
**Wiring** — mirrors the verified `maybeQuantizeKv` precedent
|
|
126
|
+
(post-construction in-place mutation of the cache array,
|
|
127
|
+
generate.ts:153; no makeCache() signature change): `maybePageKv(cache,
|
|
128
|
+
options, capacityTokens)` called ONCE at generation setup (unlike
|
|
129
|
+
quantization there is no "convert when populated" trigger — paged
|
|
130
|
+
replaces EMPTY plain `KVCache` entries, offset 0 only, before prefill).
|
|
131
|
+
`GenerateOptions.pagedKv?: { blockSize?: number }`.
|
|
132
|
+
|
|
133
|
+
Gates:
|
|
134
|
+
- server startup (explicit refuse, not silent downgrade): `--paged-kv` +
|
|
135
|
+
any of `--batch N>1`, `--kv-quant`≠off, TurboQuant, `--draft-model`,
|
|
136
|
+
non-gemma4 model ⇒ exit with a clear message. With no explicit
|
|
137
|
+
`--batch`, the CLI pins `--batch 1` (default is 8) and says so.
|
|
138
|
+
- request scope (documented v1 non-goal CELLS, scoped per request rather
|
|
139
|
+
than 400'd): media (vision/audio) and LoRA-adapter requests strip the
|
|
140
|
+
flag and run the plain cache path — one effective `pagedKv` value in
|
|
141
|
+
runGeneration keeps the strip and the prompt-cache bypass coherent.
|
|
142
|
+
- serve serial lane: paged requests skip promptCache.take/put and the
|
|
143
|
+
boundary snapshot; caches dispose on completion. Spec-decode
|
|
144
|
+
eligibility gains `!options.pagedKv` (belt — the startup refusal
|
|
145
|
+
already excludes the combination).
|
|
146
|
+
- batch lane: unreachable by construction (startup refusal pins serial),
|
|
147
|
+
so no RequestShape change — noted here so the next increment knows the
|
|
148
|
+
gate to add when batched paging lands.
|
|
149
|
+
|
|
150
|
+
Flag surface:
|
|
151
|
+
- CLI: `mlx-bun serve --paged-kv [--paged-kv-block-size N]` (default
|
|
152
|
+
block size 256 = `KVCache.STEP`, so v1's growth granularity is a
|
|
153
|
+
permutation of today's into reusable slots, not a new tuning axis).
|
|
154
|
+
- Env: `MLX_BUN_PAGED_KV=1` via `flagOn()` (explicit env wins).
|
|
155
|
+
- ServerOptions: `pagedKv?: { blockSize?: number }`, doc-commented like
|
|
156
|
+
the neighboring kvQuant/turboQuant fields.
|
|
157
|
+
|
|
158
|
+
## Implementation order
|
|
159
|
+
|
|
160
|
+
1. **BlockPool + PagedKVCache + model-free unit tests**
|
|
161
|
+
(tests/paged-kv.test.ts, FAST tier): block-boundary crossing,
|
|
162
|
+
free-list reuse after trim, pool-exhaustion typed error, gathered
|
|
163
|
+
fetch vs a plain-KVCache reference fed identical synthetic k/v
|
|
164
|
+
(bit-equal every step), dispose-doesn't-leak. Exit: green with no
|
|
165
|
+
weights.
|
|
166
|
+
2. **Wiring**: `maybePageKv` in generate.ts; GenerateOptions/
|
|
167
|
+
ServerOptions/CLI flags; startup refusals; gateway shape exclusion;
|
|
168
|
+
prompt-cache/spec bypass. Exit: flag parses, default-off inert;
|
|
169
|
+
flag-on serves a completion end-to-end (gemma4); incompatible combos
|
|
170
|
+
fail fast.
|
|
171
|
+
3. **Parity gate** (tests/paged-kv-parity.test.ts, weights-gated,
|
|
172
|
+
skipIf pattern): bit-exact (tol 0) full greedy trajectory paged-on
|
|
173
|
+
vs paged-off, incl. a ≥2-block prompt and a trim/rollback case.
|
|
174
|
+
Exit: bit-exact; any divergence documented, never a loosened
|
|
175
|
+
tolerance.
|
|
176
|
+
4. **Perf disclosure**: paged-on vs paged-off decode/prefill on a quiet
|
|
177
|
+
machine via the preflight-gated harness (expected: small decode
|
|
178
|
+
regression at M=1; report plainly). Numbers land in
|
|
179
|
+
benchmarks/RESULTS.md when run on the reference machine.
|
|
180
|
+
5. **Docs, same commit**: server-config.md, cli.md flag rows;
|
|
181
|
+
parallel-slots.md S3+ status flip (v1 shipped: serial, bf16,
|
|
182
|
+
full-attention-only, gather-then-sdpa; batched + CoW + kernel still
|
|
183
|
+
deferred); STATUS.md handoff; this doc.
|
|
184
|
+
|
|
185
|
+
## Risks
|
|
186
|
+
|
|
187
|
+
- Gather pays a full K/V copy per decode step once >1 block is occupied
|
|
188
|
+
— pure tax at M=1 (disclosed, measured; the fused-kernel follow-up is
|
|
189
|
+
the cure if it matters).
|
|
190
|
+
- Fixed pool sizing from prompt+maxTokens: a trim past a block boundary
|
|
191
|
+
followed by regrowth must re-alloc from the free list correctly —
|
|
192
|
+
dedicated test.
|
|
193
|
+
- trim/rollback against a block table is more failure-prone than
|
|
194
|
+
KVCache's offset rewind — dedicated test, not incidental coverage.
|
|
195
|
+
- Every caller that could receive a PagedKVCache it can't handle
|
|
196
|
+
(cloneKvCaches, prompt-cache put, compiled decode, batch scheduler)
|
|
197
|
+
is excluded by gates above; the ones reached only through excluded
|
|
198
|
+
paths throw typed errors if ever reached anyway.
|
|
199
|
+
|
|
200
|
+
## Follow-ups (out of scope, in dependency order)
|
|
201
|
+
|
|
202
|
+
1. **Batched integration** — PagedKVCache into the scheduler's
|
|
203
|
+
LayerInner union; block-count admission replacing byte-projection;
|
|
204
|
+
#mergeJoiner/#applyFilter/#extractAndPut allocate/free blocks
|
|
205
|
+
instead of pad+concat. This is where the padded-batch-waste win
|
|
206
|
+
actually lives.
|
|
207
|
+
2. **Block-level CoW prefix sharing** — refcounted block table
|
|
208
|
+
(retain/release), fork-on-divergent-write; extends PromptCache's
|
|
209
|
+
entry-level ref-counting down to blocks.
|
|
210
|
+
3. **Quantized paged blocks** — dtype-parametric block descriptor
|
|
211
|
+
(packed/scales/biases pools sharing one block-index space).
|
|
212
|
+
4. **Fused paged-attention Metal kernel** — port optiq's
|
|
213
|
+
sdpa_2pass_paged via the copy-verbatim methodology; kills the
|
|
214
|
+
per-step gather copy.
|
|
@@ -1,6 +1,8 @@
|
|
|
1
1
|
# Parallel slots — configurable batched serving
|
|
2
2
|
|
|
3
|
-
Status: **S0 + S1b SHIPPED** (S2
|
|
3
|
+
Status: **S0 + S1b + S2 SHIPPED** (S2's continuous inject/evict shipped
|
|
4
|
+
with the scheduler; its byte-budget admission half landed 2026-07-03 as
|
|
5
|
+
`--kv-budget`; S3+ deferred)
|
|
4
6
|
Owner: serving layer
|
|
5
7
|
Default: **off** (`--batch 1` = today's serialized path, untouched)
|
|
6
8
|
|
|
@@ -27,6 +29,16 @@ Direct/embedding mode stays batch=1.
|
|
|
27
29
|
> onToken failure containment, drain-on-serial-waiter (no starvation), and
|
|
28
30
|
> the engine-hygiene transplant (pipelined decode via asyncEval, clearCache
|
|
29
31
|
> 1/256 steps, chunked admit prefill interleaved with decode steps).
|
|
32
|
+
> Landed since: **SSM batched path** (Qwen3.5 batches, token-exact vs
|
|
33
|
+
> mlx-lm B=2 — 2026-07-02, batching-perf-path.md P5); **per-row logits
|
|
34
|
+
> processors** (repetition/presence/frequency penalties + logit_bias +
|
|
35
|
+
> min_p/XTC batch — 2026-07-02, same doc); **grammar under batching** (B1
|
|
36
|
+
> per-row matchers — structured-output.md); **KV-budget admission**
|
|
37
|
+
> (`--kv-budget`, queue-don't-OOM, `/stats.batch` fields — 2026-07-03,
|
|
38
|
+
> grammar-spec-batching-integration.md Phase D); **Tier-0 universal
|
|
39
|
+
> coverage** (plain full-attention archs, e.g. Llama, batch via per-row
|
|
40
|
+
> RoPE, token-exact vs mlx-lm B=2 — 2026-07-03, same doc's execution log;
|
|
41
|
+
> maskArray + sliding universal archs stay serial).
|
|
30
42
|
|
|
31
43
|
## STATUS (2026-06-14) — LIVE for full-attention models
|
|
32
44
|
|
|
@@ -62,8 +74,8 @@ default stays mixed-precision (optiq parity); an explicit `--kv-quant` under
|
|
|
62
74
|
novel extension: a memory-density win (batching × 4-bit KV compound), gated by
|
|
63
75
|
KL + 6-task quality (per the tree framing), deferred. Other follow-ups (all
|
|
64
76
|
optional polish): the `extend` join op (today joins re-merge), KV-budget
|
|
65
|
-
admission, prompt-cache reuse under
|
|
66
|
-
allocation).
|
|
77
|
+
admission (**LANDED 2026-07-03** — `--kv-budget`), prompt-cache reuse under
|
|
78
|
+
batching, and **paged KV** (zero-padding-waste allocation).
|
|
67
79
|
|
|
68
80
|
## Why it helps (and when it doesn't)
|
|
69
81
|
|
|
@@ -222,54 +234,58 @@ base library; **mlx-optiq and mlx-bun are both optimizations on top of it**
|
|
|
222
234
|
(optiq = mixed-precision per-layer KV quant, "drop-in for mlx_lm.server", no
|
|
223
235
|
batching; mlx-bun = per-model fused kernels + batched serving). Bit-parity with
|
|
224
236
|
an ancestor is just an **easy correctness test where we overlap one** — not a
|
|
225
|
-
porting goal. The "L1/L2
|
|
237
|
+
porting goal. The "L1/L2" labels below are **parity anchors / correctness
|
|
226
238
|
checkpoints**, not a strict stack: L1 = "matches mlx-lm (bf16)", L2 = "matches
|
|
227
|
-
optiq (mixed-precision quant)"
|
|
228
|
-
|
|
239
|
+
optiq (mixed-precision quant)"; beyond them is the no-oracle **Lab** class ("our
|
|
240
|
+
own optimizations, no ancestor does them → no oracle" — formerly labeled "L3";
|
|
241
|
+
the `--l3` product tier was deleted 2026-07-05,
|
|
242
|
+
docs/design/unified-engine-frontier-plan.md §6). A combo no ancestor does (e.g. **batched mixed-precision
|
|
229
243
|
serving** — mlx-lm can't do mixed-precision, optiq can't batch) is genuinely new
|
|
230
244
|
territory: gate it by KL + the 6-task quality + benchmarks, never bit-exactness.
|
|
231
245
|
The anchors are still useful — we can always dial back to a bit-exact-vs-ancestor
|
|
232
246
|
config as a correctness checkpoint and build outward.
|
|
233
247
|
|
|
234
248
|
- **L1 — mlx-lm parity (standard KV).** bf16 KV, **bit-exact vs mlx-lm**. The
|
|
235
|
-
foundation oracle
|
|
236
|
-
|
|
249
|
+
foundation oracle (and, since 2026-07-05, the naked default). Path: monolith
|
|
250
|
+
`Gemma4Model` / `MiniCPM5Model`, `ops.sdpa`.
|
|
237
251
|
- **L2 — mlx-optiq parity (mixed-precision KV quant).** quant/mixed KV,
|
|
238
|
-
**bit-exact vs optiq**. Path: `QuantizedKVCache` + `quantizedSdpa`
|
|
239
|
-
|
|
240
|
-
|
|
241
|
-
|
|
242
|
-
design
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
|
|
250
|
-
Degradation
|
|
251
|
-
|
|
252
|
+
**bit-exact vs optiq**. Path: `QuantizedKVCache` + `quantizedSdpa` (fused-sdpa
|
|
253
|
+
follows `--kv-quant`: on for `config` — the optiq golden composition; off →
|
|
254
|
+
*unfused* for uniform/bf16).
|
|
255
|
+
- **L3 — DELETED as a product mode 2026-07-05**
|
|
256
|
+
(docs/design/unified-engine-frontier-plan.md §6): no output-changing kernel
|
|
257
|
+
(perf-kernel, fused-decode, fused-gelu) beat the L1 baseline in a paired A/B,
|
|
258
|
+
so the kernels were deleted and `--l3` hard-errors. Its no-oracle territory is
|
|
259
|
+
now the **Lab**: output-changing experiments are env-flag items with a bench +
|
|
260
|
+
expiry, gated by **KL vs the compatible upstream + the 6 eval-task mean
|
|
261
|
+
(`src/eval/tasks/`) + perf numbers** — never bit-exactness — until one earns
|
|
262
|
+
promotion.
|
|
263
|
+
|
|
264
|
+
Degradation lever: bf16 KV drops L2→L1. The surviving kill switches
|
|
265
|
+
(`--compiled-decode`, `--compiled-activations`, `--fused-sdpa`) are per-kernel
|
|
266
|
+
escapes. Each layer must stand alone and stay reachable (don't-delete-
|
|
252
267
|
optionality as a fallback ladder).
|
|
253
268
|
|
|
254
269
|
So a feature like batched decode (`--batch`) is "done" only when green at
|
|
255
270
|
**every (model × layer)** cell — and each cell may need its own fix:
|
|
256
271
|
|
|
257
|
-
| Model · path | L1 (bf16 / mlx-lm) | L2 (quant / optiq) | L3 (
|
|
272
|
+
| Model · path | L1 (bf16 / mlx-lm) | L2 (quant / optiq) | ~~L3~~ (tier deleted 2026-07-05) |
|
|
258
273
|
| --- | --- | --- | --- |
|
|
259
|
-
| CPM (MiniCPM5) | ✅ **oracle-verified** (= mlx-lm B=2) | ☐ quant + unfused-mask |
|
|
260
|
-
| Gemma 12B | ✅ **oracle-verified** (= mlx-lm B=2, short-ctx) | ☐ generated + sliding-window |
|
|
261
|
-
| Gemma e4b | ✅ **oracle-verified** (= mlx-lm B=2, short-ctx) | ☐ generated + same |
|
|
262
|
-
| Gemma 26B | ✅ **oracle-verified** (= mlx-lm B=2, short-ctx) | ☐ generated + MoE |
|
|
263
|
-
| *(Qwen, future)* | *new path* | *new path* |
|
|
274
|
+
| CPM (MiniCPM5) | ✅ **oracle-verified** (= mlx-lm B=2) | ☐ quant + unfused-mask | — |
|
|
275
|
+
| Gemma 12B | ✅ **oracle-verified** (= mlx-lm B=2, short-ctx) | ☐ generated + sliding-window | — |
|
|
276
|
+
| Gemma e4b | ✅ **oracle-verified** (= mlx-lm B=2, short-ctx) | ☐ generated + same | — |
|
|
277
|
+
| Gemma 26B | ✅ **oracle-verified** (= mlx-lm B=2, short-ctx) | ☐ generated + MoE | — |
|
|
278
|
+
| *(Qwen, future)* | *new path* | *new path* | — |
|
|
264
279
|
|
|
265
280
|
**🎯 L1 batched decode COMPLETE (2026-06-14d): the entire L1 row is green** —
|
|
266
281
|
all 4 model paths bit-parity with mlx-lm B=2 (CPM full-attn, Gemma 12B
|
|
267
282
|
sliding+full, e4b per-layer-input+KV-share, 26B MoE). Only one per-path fix was
|
|
268
283
|
needed (e4b's `[1,L,…]` per-layer-input hardcode); sliding (RotatingKVCache),
|
|
269
284
|
KV-sharing, and MoE were all already B-generic. Caveat: short-context (pre-wrap)
|
|
270
|
-
— ring-wrap (>window) is the one remaining L1 follow-up. Next
|
|
271
|
-
(quantized KV → `quantizedSdpaUnfused` + 4-D mask, vs optiq)
|
|
272
|
-
kernels under batching
|
|
285
|
+
— ring-wrap (>window) is the one remaining L1 follow-up. Next row: **L2**
|
|
286
|
+
(quantized KV → `quantizedSdpaUnfused` + 4-D mask, vs optiq). (The former
|
|
287
|
+
"L3: perf kernels under batching" row is void — those kernels were deleted
|
|
288
|
+
2026-07-05, docs/design/unified-engine-frontier-plan.md §6.)
|
|
273
289
|
|
|
274
290
|
Harness: `tests/batched-decode-parity.test.ts` has two paths — (a)
|
|
275
291
|
`runBatchedDecodeParity` = the internal-consistency KL check (B=N vs our B=1,
|
|
@@ -329,15 +345,15 @@ Batched-decode landmines found by the modes analysis (2026-06-14):
|
|
|
329
345
|
breaks; plus KV-sharing (sharer `SharedKv`).
|
|
330
346
|
- **Gemmas**: `RotatingKVCache` sliding-window ring — per-row decode mask + ring
|
|
331
347
|
write (the hard one). 26B adds MoE (`GatherQMM`/`SwitchGLU`) batched routing.
|
|
332
|
-
- **Generated handlers (
|
|
348
|
+
- **Generated handlers (all 3 Gemmas)** repeat the K/Q rope timing trap —
|
|
333
349
|
they read `cache.ropeOffsetArr` twice around `updateAndFetch`
|
|
334
350
|
(generated/gemma4-12b.ts:63 & :74). The fix goes in the GENERATOR
|
|
335
351
|
`scripts/gen-model.ts`, then regenerate (the files are DO-NOT-EDIT).
|
|
336
|
-
- **Quant + batched**:
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
|
|
340
|
-
is
|
|
352
|
+
- **Quant + batched**: quant batched decode goes through `quantizedSdpaUnfused`
|
|
353
|
+
— correct. (The old "perf kernel bypassed under batching" debt is resolved by
|
|
354
|
+
deletion: the perf/fused-decode kernels were deleted 2026-07-05,
|
|
355
|
+
docs/design/unified-engine-frontier-plan.md §6 — no batched-aware fused decode
|
|
356
|
+
kernel is planned.) `quantizedSdpaUnfused` must be validated with a 4-D bool mask.
|
|
341
357
|
- **All cells**: compiled-decode forced off under `--batch>1`; LoRA batches only
|
|
342
358
|
within one adapter group (else drain).
|
|
343
359
|
|
|
@@ -363,7 +379,7 @@ most" fidelity and cost:
|
|
|
363
379
|
concurrent rows per GB. This is what "dynamically allocate based on
|
|
364
380
|
who needs it most" actually means. Cost: a custom attention kernel +
|
|
365
381
|
a block manager. Feasible — the project already ships custom Metal
|
|
366
|
-
|
|
382
|
+
kernel infrastructure (`src/mlx/metal-kernel.ts`).
|
|
367
383
|
|
|
368
384
|
Plan: ship rung 2 (S1–S2) for the throughput win + dynamic admission
|
|
369
385
|
with no new kernel; build rung 3 (S3+) as the memory-density upgrade.
|
|
@@ -376,8 +392,16 @@ growing in 256-token steps, width = longest live sequence, left-pad
|
|
|
376
392
|
`BatchScheduler` + `GenerationGateway` (S1b SHIPPED). But the
|
|
377
393
|
**budget/admission half of rung 2 is NOT built** (no `B × S_max`
|
|
378
394
|
projection, no KV budget enforcement) — that's scheduler work (S2, NOT
|
|
379
|
-
STARTED).
|
|
380
|
-
|
|
395
|
+
STARTED). *(Since LANDED 2026-07-03: `--kv-budget` aggregate projection
|
|
396
|
+
gate in the scheduler — grammar-spec-batching-integration.md Phase D.)*
|
|
397
|
+
And **rung 3 (paged): the v1 CACHE CLASS shipped** —
|
|
398
|
+
`PagedKVCache`/`BlockPool` behind default-off `--paged-kv`
|
|
399
|
+
(docs/design/paged-kv-cache.md): serial batch=1, Gemma4, bf16,
|
|
400
|
+
gather-then-stock-sdpa (no custom kernel yet), gated bit-exact vs the
|
|
401
|
+
plain path. The rung-3 PAYOFF (batched block allocation replacing the
|
|
402
|
+
pad+concat merge/extend, block-level CoW prefix sharing, the fused paged
|
|
403
|
+
kernel) is still the deferred follow-up — the shipped slice is the block
|
|
404
|
+
manager + correctness proof, not the density win.
|
|
381
405
|
|
|
382
406
|
KV quantization is a force multiplier: 4-bit KV (`generate.ts:67`,
|
|
383
407
|
already supported) is ~4× smaller, so the same budget buys ~4× the slots
|
|
@@ -456,7 +480,10 @@ server, concurrent + streaming + drain). No serial regression (server 17/17,
|
|
|
456
480
|
tools 13/13).
|
|
457
481
|
(3) **DONE 2026-06-14** — `_is_batchable` gate is the `GenerationGateway.willBatch`
|
|
458
482
|
predicate (full-attention model + no vision/adapter/repetition-penalty/user-seed
|
|
459
|
-
→ batch; else serial).
|
|
483
|
+
→ batch; else serial). *(The predicate has since widened: repetition
|
|
484
|
+
penalty + logits extras now batch via per-row processors, 2026-07-02.)*
|
|
485
|
+
`B×S_max` memory admission is still TODO — **DONE 2026-07-03**
|
|
486
|
+
(`--kv-budget`, see the v2-plan pointer above).
|
|
460
487
|
|
|
461
488
|
## Fallbacks (must degrade, never break)
|
|
462
489
|
|
|
@@ -478,9 +505,9 @@ predicate (full-attention model + no vision/adapter/repetition-penalty/user-seed
|
|
|
478
505
|
|
|
479
506
|
## Config surface
|
|
480
507
|
|
|
481
|
-
- `ServerOptions.batch` (default 1);
|
|
482
|
-
`promptConcurrency`.
|
|
483
|
-
- CLI: `--batch N` (shipped; `--slots N` was the original design name, `--decode-concurrency` accepted as mlx_lm.server alias);
|
|
508
|
+
- `ServerOptions.batch` (default 1); `kvBudgetBytes` (shipped 2026-07-03);
|
|
509
|
+
later `promptConcurrency`.
|
|
510
|
+
- CLI: `--batch N` (shipped; `--slots N` was the original design name, `--decode-concurrency` accepted as mlx_lm.server alias); `--kv-budget <GB>` (shipped 2026-07-03).
|
|
484
511
|
- `batch === 1` → today's serialized promise chain, untouched.
|
|
485
512
|
`batch > 1` → scheduler.
|
|
486
513
|
|
|
@@ -502,11 +529,17 @@ predicate (full-attention model + no vision/adapter/repetition-penalty/user-seed
|
|
|
502
529
|
**Teacher-forced gate**: a 2-row batch must produce per-row logits matching
|
|
503
530
|
two solo runs within bf16 tolerance (the train-batch-e2e methodology, now
|
|
504
531
|
applied to decode).
|
|
505
|
-
- **S2 — N-wide + continuous injection/eviction. (
|
|
506
|
-
queued requests prefill into freed rows
|
|
507
|
-
|
|
508
|
-
|
|
509
|
-
|
|
532
|
+
- **S2 — N-wide + continuous injection/eviction. (DONE)** Rows retire on EOS;
|
|
533
|
+
queued requests prefill into freed rows (shipped with the scheduler).
|
|
534
|
+
Dynamic byte-budget admission landed 2026-07-03 as `--kv-budget`
|
|
535
|
+
(aggregate projection gate, queue-don't-OOM —
|
|
536
|
+
grammar-spec-batching-integration.md Phase D).
|
|
537
|
+
- **S3+ — paged KV (v1 cache class SHIPPED; density payoff DEFERRED)** —
|
|
538
|
+
the block manager (`src/model/paged-kv.ts`) shipped behind default-off
|
|
539
|
+
`--paged-kv` at serial scope, bit-exact-gated (paged-kv-cache.md); the
|
|
540
|
+
batched block allocation, block-level CoW, and the custom
|
|
541
|
+
paged-attention kernel remain the follow-ups. KV-quant under batch
|
|
542
|
+
landed separately (Phase 3.1 + milestone 2); LoRA-group batching open.
|
|
510
543
|
|
|
511
544
|
Every phase ships default-off behind `batch=1`; the serialized path is
|
|
512
545
|
never removed (only bypassed).
|
|
@@ -13,15 +13,19 @@ matches bit-for-bit**:
|
|
|
13
13
|
|
|
14
14
|
- **L1** — drop-in for **mlx-lm** (its op is the oracle; bit-exact or it's a bug).
|
|
15
15
|
- **L2** — drop-in for **mlx-optiq** (optiq's behavior is the oracle).
|
|
16
|
-
- **
|
|
17
|
-
|
|
18
|
-
|
|
16
|
+
- **the Lab** (formerly the `--l3` tier — deleted as a product mode 2026-07-05,
|
|
17
|
+
[unified-engine-frontier-plan.md §6](unified-engine-frontier-plan.md)) — **our
|
|
18
|
+
originals** (flash-CCE head, ORPO loss, prefix-share, segmented backward, HLG
|
|
19
|
+
sampler, expert offload). **No oracle ships**, so correctness is gated by
|
|
20
|
+
finite-difference / KL / quality, *not* parity. Output-changing inference
|
|
21
|
+
experiments live here as env-flag items with a bench + expiry, until one beats
|
|
22
|
+
the L1 baseline in a paired A/B.
|
|
19
23
|
|
|
20
24
|
**The gate is the oracle.** A node can only live in a lower tier if that oracle
|
|
21
25
|
exists AND it matches bit-exactly; otherwise it bubbles up. Note the subtlety the
|
|
22
26
|
map already shows: an *optimization* that still matches the oracle stays low —
|
|
23
27
|
`compiled-decode` is ours but is L1 because it replays the same ops bit-for-bit.
|
|
24
|
-
Only nodes with **no** oracle (flash-CCE) are *forced*
|
|
28
|
+
Only nodes with **no** oracle (flash-CCE) are *forced* into the Lab. (Background +
|
|
25
29
|
post-mortem: [orpo-flash-cce-pin-leak.md](../investigations/orpo-flash-cce-pin-leak.md).)
|
|
26
30
|
|
|
27
31
|
## Two axes: parity gates performance
|
|
@@ -31,7 +35,7 @@ made the flag surface unreadable:
|
|
|
31
35
|
|
|
32
36
|
1. **Parity** — *which reference does it reproduce bit-for-bit?* This sets the
|
|
33
37
|
**lowest level** it may live in: matches mlx-lm → L1; matches mlx-optiq but not
|
|
34
|
-
mlx-lm → L2; matches neither →
|
|
38
|
+
mlx-lm → L2; matches neither → no-oracle, Lab-gated (KL + tests).
|
|
35
39
|
2. **Performance** — *is it the fastest correct way?* This sets whether it's the
|
|
36
40
|
**default (on)** within its level, or a kept-but-**off** opt-in.
|
|
37
41
|
|
|
@@ -49,11 +53,12 @@ your floor of trust; the default within it is the best way to honor that floor.
|
|
|
49
53
|
- **slower but still correct** → **keep it as a documented default-off flag**
|
|
50
54
|
(optionality for A/B; never the default, never deleted).
|
|
51
55
|
- **breaks the guarantee but offers a real tradeoff** (e.g. 2× tok/s at KL
|
|
52
|
-
0.0015) →
|
|
56
|
+
0.0015) → a **Lab item** (env-flag, bench + expiry); it earns a place only
|
|
57
|
+
by beating the L1 baseline in a paired A/B.
|
|
53
58
|
|
|
54
59
|
Worked examples:
|
|
55
60
|
- *5% less memory, still bit-for-bit with mlx-lm* → **default-on in L1.**
|
|
56
|
-
- *2× tok/s, but KL 0.0015 to mlx-lm* → **
|
|
61
|
+
- *2× tok/s, but KL 0.0015 to mlx-lm* → **Lab item** (proves itself or expires).
|
|
57
62
|
- *another way, 3% slower, still correct* → **kept, off in every level.**
|
|
58
63
|
|
|
59
64
|
### Our decode levers under this rule (today)
|
|
@@ -61,23 +66,28 @@ Worked examples:
|
|
|
61
66
|
| lever | parity | perf | placement |
|
|
62
67
|
|---|---|---|---|
|
|
63
68
|
| `compiled-decode` | bit-exact w/ uncompiled (`compiled-decode.test`) | faster | **default-on, every level** |
|
|
64
|
-
| `perf-kernel` |
|
|
65
|
-
| `fused-sdpa` | matches optiq bit-for-bit, not mlx-lm (`fused-sdpa.test` tier-a goldens) | faster | **default-on
|
|
66
|
-
| `fused-decode` | (
|
|
67
|
-
|
|
68
|
-
So `--l1/--l2
|
|
69
|
-
each tier are the **performance** axis (the fastest kernel that still holds it).
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
mixed-precision —
|
|
69
|
+
| `perf-kernel` | envelope-gated, no bit-exact oracle (was gated ≥56/64 teacher-forced argmax vs our own frozen compat trajectory) | regressed e4b 0.62–0.93× | **DELETED 2026-07-05, [unified-engine-frontier-plan.md §6](unified-engine-frontier-plan.md)** (kernel + frozen-oracle tests removed) |
|
|
70
|
+
| `fused-sdpa` | matches optiq bit-for-bit, not mlx-lm (`fused-sdpa.test` tier-a goldens) | faster | **default follows `--kv-quant`** (on for `config` — the L2 bridge; off for uniform/bf16) |
|
|
71
|
+
| `fused-decode` | (no oracle) | ~1.00× (`optimization_plan.md`) | **DELETED 2026-07-05, [unified-engine-frontier-plan.md §6](unified-engine-frontier-plan.md)** |
|
|
72
|
+
|
|
73
|
+
So `--l1/--l2` is the **parity** axis (the guarantee), and the defaults within
|
|
74
|
+
each tier are the **performance** axis (the fastest kernel that still holds it).
|
|
75
|
+
2026-07-05 outcome ([unified-engine-frontier-plan.md §6](unified-engine-frontier-plan.md)):
|
|
76
|
+
no output-changing decode kernel beat the L1 baseline in a paired A/B, so `--l3`
|
|
77
|
+
was deleted as a product mode (the flag now hard-errors) along with its kernels
|
|
78
|
+
(perf-kernel, fused-decode). The no-oracle territory — HLG sampler, expert
|
|
79
|
+
offload, batched mixed-precision — is now the **Lab**: env-flag experiments with
|
|
80
|
+
a bench + expiry, gated by KL + quality, off the decode axis.
|
|
75
81
|
|
|
76
82
|
> History: commit `f1bf5cc` put the perf kernel in the L2 preset claiming
|
|
77
83
|
> "perf-kernel-oracle.test asserts the fused kernel == the frozen optiq goldens".
|
|
78
84
|
> Both halves were wrong — that golden is frozen from mlx-bun's own compat engine,
|
|
79
85
|
> and the gate is argmax-agreement, not equality — so 2026-07-01 restored the tier
|
|
80
86
|
> rule: the bare tier is the guarantee; perf kernels are opt-ins within it.
|
|
87
|
+
> Deletion outcome: the perf kernel never beat the L1 baseline in a paired A/B
|
|
88
|
+
> (regressed e4b 0.62–0.93×; only win 12B@16k +6% with a KL WARN), so it was
|
|
89
|
+
> DELETED 2026-07-05 with its frozen-oracle tests
|
|
90
|
+
> ([unified-engine-frontier-plan.md §6](unified-engine-frontier-plan.md)).
|
|
81
91
|
|
|
82
92
|
## Why this matters for CLI flags
|
|
83
93
|
|
|
@@ -94,7 +104,7 @@ A flag falls into exactly one of three buckets:
|
|
|
94
104
|
|---|---|---|---|
|
|
95
105
|
| the only viable route (no alternative node) | — | not a choice | **always-on default; drop the flag** |
|
|
96
106
|
| between two routes at the SAME tier / same oracle | none | cosmetic / safety guard | **default it; expose only if it ever fails** |
|
|
97
|
-
| between a **lower-tier** route and
|
|
107
|
+
| between a **lower-tier** route and a **no-oracle (Lab)** route | **L1/L2 ↔ Lab** | a real **parity ⇄ optimization** knob | **keep, document as such, group by tier** |
|
|
98
108
|
|
|
99
109
|
So the DAG lets us answer "**what flags push down THAT route**" for any path, and
|
|
100
110
|
from that:
|
|
@@ -102,7 +112,7 @@ from that:
|
|
|
102
112
|
1. **Prune always-on flags to defaults** — anything whose alternative route is
|
|
103
113
|
never selected, or whose toggle doesn't change the tier, shouldn't be a flag.
|
|
104
114
|
2. **Surface the real axis.** The flags that genuinely matter all toggle
|
|
105
|
-
*bit-exact-lower-tier ⇄ our-
|
|
115
|
+
*bit-exact-lower-tier ⇄ our-no-oracle-optimization*. Those could collapse into a small
|
|
106
116
|
number of intent-level switches (e.g. a single `--exact` that forces every node
|
|
107
117
|
down to its lowest tier, vs the optimized default) instead of N independent
|
|
108
118
|
kernel toggles the user has to understand individually.
|
|
@@ -117,11 +127,11 @@ from that:
|
|
|
117
127
|
knob, **not** a parity choice. Should read as "memory budget," not "exactness."
|
|
118
128
|
- `--grad-accum N` → changes the accumulation route (L1; smoother grads) → a real
|
|
119
129
|
*training-quality* knob.
|
|
120
|
-
- `--no-flash` → **
|
|
121
|
-
fallback (the one place the head leaves
|
|
122
|
-
- `--no-prefix` → **
|
|
123
|
-
numerical cost.
|
|
124
|
-
- `--lambda` → ORPO hyperparameter (pure
|
|
130
|
+
- `--no-flash` → **no-oracle flash-CCE → L2 MLX fused head** → a genuine *parity ⇄
|
|
131
|
+
perf* fallback (the one place the head leaves the no-oracle class).
|
|
132
|
+
- `--no-prefix` → **no-oracle prefix-share → bit-exact two-forward** → a *perf* knob
|
|
133
|
+
with no numerical cost.
|
|
134
|
+
- `--lambda` → ORPO hyperparameter (pure no-oracle) → a *training* knob.
|
|
125
135
|
|
|
126
136
|
The pattern: most flags are **memory/training knobs or always-on guards**; only a
|
|
127
137
|
couple are true **parity ⇄ optimization** toggles. Today they're a flat list the
|
|
@@ -133,15 +143,15 @@ actually are.
|
|
|
133
143
|
The map is hand-authored today. To make the framework *enforce* itself:
|
|
134
144
|
|
|
135
145
|
1. **Derive the graph from code** — extract the call graph + per-node op/oracle tags
|
|
136
|
-
so it stays in sync and is **queryable** ("every
|
|
137
|
-
|
|
146
|
+
so it stays in sync and is **queryable** ("every no-oracle node", "paths that
|
|
147
|
+
cross into the Lab", "the pure-L1 subgraph", "flags that select route X").
|
|
138
148
|
2. **CI-enforce the gate** — a node tagged L1 must have a passing bit-exact parity
|
|
139
|
-
test vs mlx-lm; if it can't, it's a bug or must be re-tagged
|
|
140
|
-
the gate" becomes a check, not a principle.
|
|
141
|
-
3. **Shrink the
|
|
142
|
-
bit-exact-able and just unproven (e.g. the prefix-shared *forward* is
|
|
143
|
-
bit-exact — `prefix-shared-parity.ts` — so only the encode-once *fold*
|
|
144
|
-
|
|
149
|
+
test vs mlx-lm; if it can't, it's a bug or must be re-tagged no-oracle. "The
|
|
150
|
+
oracle is the gate" becomes a check, not a principle.
|
|
151
|
+
3. **Shrink the no-oracle surface** — audit whether anything in the Lab island is
|
|
152
|
+
actually bit-exact-able and just unproven (e.g. the prefix-shared *forward* is
|
|
153
|
+
already bit-exact — `prefix-shared-parity.ts` — so only the encode-once *fold*
|
|
154
|
+
is truly no-oracle). Minimizing it minimizes what needs KL/quality gates.
|
|
145
155
|
4. **Generate the flag surface from the graph** — once flags map to routes, the CLI
|
|
146
156
|
help + defaults can be derived: always-on routes have no flag, parity⇄opt routes
|
|
147
157
|
collapse to intent switches, knobs group by what they trade.
|