mlx-bun 0.0.4 → 0.0.5

package/docs/design/docs-reorg-plan.md

@@ -4,6 +4,19 @@
 (src/ and scripts/ moves — they change import paths / script references
 and need a test run, so deferred to their own pass). Authored 2026-06-14.
 
+**Execution state (verified 2026-06-17):**
+- P0 #1 — **DONE**: `AGENTS.md` is a symlink → `CLAUDE.md`; drift risk resolved.
+- P0 #2 — **DONE**: `benchmarks/RESULTS.md` created; old committed benchmark artifacts `git rm`'d. Root untracked spew (`benchmarks-h2h-2026-06-13-*` etc.) is now cleaned — no untracked report files found at root.
+- P0 #3 — **DONE**: root HTML/md spew cleared (no `benchmarks-h2h-*.md`, `benchmark-report-*.html`, `kernel-fusion-report-*.html`, or `project-review.html` at root).
+- P0 #4 — **DONE**: hardcoded test count removed from CLAUDE.md (no "All 72 tests" present).
+- P1 #5 — **DONE**: `STATUS.md` exists at root.
+- P1 #6 — **DONE**: `docs/planning/` exists with `IDEAS.md`, `PRODUCT_ROADMAP.md`, `ResearchTopics.md`, `journal.md`.
+- P1 #7 — open (optimization_plan.md still lives in docs/design/ as a standalone).
+- P1 #8 — status unknown (archive/ not verified in this pass).
+- P2 #9 — **NOT DONE**: `harness-pi.ts`, `pi-provider.ts`, `pi-terminal.ts`, `pi-web.ts` still live flat in `src/` root (no `src/pi/` subdirectory).
+- P2 #10 — **PARTIALLY DONE**: ~57 one-off research/debug scripts moved to `scripts/experiments/`; production + tooling + bench-* + regen-* + parity-* + gen-model + eval/serve remain at `scripts/` root (the full bench/regen/build split was not applied).
+- P2 #11 — **DONE**: CLI vocabulary reconciled.
+
 **Scope & ethos.** Private repo, single author + agents, MLX-by-identity.
 No OSS/community scaffolding (CONTRIBUTING, issue templates, badges) —
 out of scope by design. **Prefer move / consolidate over delete.** The
@@ -127,7 +140,7 @@ files, not benchmark output. Leave entirely alone.
 | 5 | Add `STATUS.md` (or a pinned top section in PLAN.md) holding only live handoff: current state + next action + blockers. Move PLAN.md's `NEXT UP` / `NEXT SESSION PICKUP` / `SESSION SWEEP` blocks into it. | Separates transient "do this next" from 2,475 lines of durable phase history; gives agents one entry point. |
 | 6 | Create `docs/planning/` and move `PRODUCT_ROADMAP.md`, `IDEAS.md`, `ResearchTopics.md`, `journal.md` under it. Leave only README + PLAN.md (+ STATUS.md) at root. Add a ~6-line "doc map" to README. | Root is currently a 6-file planning thicket; extends the taxonomy that already works in `docs/`. |
 | 7 | Pick one home for the optimization plan: fold `docs/design/optimization_plan.md` into PLAN.md Phase A–E, or have PLAN.md link out to it as canonical. | Two homes for the same plan = drift. |
-| 8 | Resolve the archive overlap: `archive/pi-builtin-investigation.html` duplicates [docs/investigations/pi-builtin-investigation.md](../investigations/pi-builtin-investigation.md). Decide whether `archive/` (tracked; 8 generated HTML files incl. a 374 KB pi-session transcript) should be tracked at all, or moved to an ignored location. Keep the `.md`, drop the duplicate. | Stops committing heavyweight generated HTML; one source per investigation. |
+| 8 | Resolve the archive overlap: `archive/pi-builtin-investigation.html` duplicates [docs/investigations/pi-builtin-investigation.md](../investigations/pi-builtin-investigation.md). Decide whether `archive/` (tracked; 2 HTML files: `mlx-bun-lab-report.html` + `wwdc-mlx-bun.html`) should be tracked at all, or moved to an ignored location. Keep the `.md`, drop the duplicate. | Stops committing heavyweight generated HTML; one source per investigation. |
 
 ## P2 — Polish (optional)
 

package/docs/design/hlg-sampling.md

@@ -1,6 +1,6 @@
 # HLG sampling — a piecewise tone curve on the logits
 
-Status: **design + phasing** (no code yet; this doc is the contract)
+Status: **implemented, found neutral, superseded by Curve Designer** — `applyHlg` shipped in `src/sampler.ts`; investigation found HLG a wash vs plain temperature; the Curve Designer (`src/curve-sampler.ts`) is its v2 replacement. See [docs/investigations/hlg-sampling-investigation.md](../investigations/hlg-sampling-investigation.md) for findings. This doc remains as the design rationale.
 Owner: sampling layer (`src/sampler.ts`)
 Default: **off** (`--hlg-sampling off` = today's temperature sampler, untouched)
 
@@ -10,6 +10,12 @@ Default: **off** (`--hlg-sampling off` = today's temperature sampler, untouched)
 > [docs/reference/server-config.md](../reference/server-config.md) and
 > [docs/reference/server-api.md](../reference/server-api.md) once shipped.
 
+## Outcome / Retrospective
+
+HLG was implemented through **Piece 2** (pure curve + neutrality tests + serial wiring, `applyHlg` live in `src/sampler.ts`). Evaluation found it **neutral vs plain temperature** — no measurable gain in diversity or quality at matched output entropy. As a result, HLG is **superseded by the Curve Designer** (`src/curve-sampler.ts`), which generalizes the same region-shaping idea with an editable control-point interface and cleaner UX. Full empirical findings are in [docs/investigations/hlg-sampling-investigation.md](../investigations/hlg-sampling-investigation.md).
+
+The sections below (Goal through Correctness) remain as design rationale. Pieces 3–6 of Phasing and the Open Questions were never executed; see the [Phasing](#phasing-the-piece-n-increments) and [Open Questions](#open-questions--risks) sections for historical context.
+
 ## Goal
 
 Add a sampling transform — gated entirely behind `--hlg-sampling on` and a
@@ -290,6 +296,8 @@ quality + diversity rather than parity — except for the degeneracy gates, whic
 
 ## Phasing (the "piece N" increments)
 
+> **Historical / superseded.** Pieces 1–2 shipped. Pieces 3–6 were never executed — HLG was evaluated after Piece 2 and found neutral; the feature was superseded by the Curve Designer before Piece 3 began.
+
 1. **Pure curve + neutrality tests** — `applyHlg(lp, params)` in `sampler.ts`,
    top-anchored pivot only; degeneracy + monotonicity tests green across all
    four models. (No flag wired yet — pure, fully-tested function.)
@@ -305,6 +313,8 @@ quality + diversity rather than parity — except for the degeneracy gates, whic
 
 ## Open questions / risks
 
+> **Historical / superseded.** These questions were open at the time of design. HLG evaluation concluded with a neutral verdict before they were answered; the Curve Designer supersedes this path. Questions are preserved as design context.
+
 - **Default knob values** are calibration-pending — set empirically by the
   Piece 5 sweep, not guessed. Initial exploration points: `c ≈ 6`, `w ≈ 4`,
   `β_h ≈ 4`, `β_t ≈ 6` nats; confirm or replace with measured values.
@@ -337,4 +347,4 @@ for it. That a transfer function is standard practice in another domain says
 nothing about whether it helps token sampling — that is decided only by
 measurement (the KL characterization, the capability guardrail, and the
 diversity lens in [Correctness](#correctness--evaluation)). The full empirical
-write-up goes in `docs/investigations/hlg-sampling-investigation.md` once run.
+write-up is in [docs/investigations/hlg-sampling-investigation.md](../investigations/hlg-sampling-investigation.md).

package/docs/design/memory-system.md

@@ -1,6 +1,6 @@
 # Memory system — design sketch
 
-Status: **design** (nothing built yet; capture prerequisite already landed)
+Status: **aspirational — not scheduled; no PLAN.md phase** (nothing built yet; capture prerequisite already landed). Note: `src/jobs/` is the generic background-job runner for quantize/finetune/fit — it is NOT this memory system.
 Owner: assistant layer (pi-web / pi-terminal / web-tools / jobs)
 Default: **off** until the user opts in (vault + schedule are explicit, consented actions)
 

package/docs/design/optimization_plan.md

@@ -1,8 +1,10 @@
 # mlx-bun: gemma4.ts decompose & optimize — action plan
 
 > This is the canonical **design/rationale** for the optimization work.
-> Live **status/tracking** (checkboxes, findings) lives in PLAN.md under
-> "Optimization plan Phase A–E".
+> Live **status/tracking** (checkboxes, findings) for Phases A–D lives in
+> PLAN-archive.md; Phase E remains in PLAN.md.
+
+**Status (verified 2026-06-17):** Phases A–D are **DONE** (`[x]` in PLAN-archive.md, all landed 2026-06-10/11). Phase E — per-config fused decode-SDPA kernel — is **in progress** (`[~]` in PLAN.md, started 2026-06-11; fused kernel exists and is quality-gated per PLAN.md Step 4). The body of this doc remains the durable design rationale; completed phase checkboxes and findings live in PLAN-archive.md; Phase E tracking lives in PLAN.md.
 
 **Premise (confirmed from the source read):** the current engine sits at
 parity-to-ahead vs mlx-lm on the 12B *while still interpreting the decode

package/docs/design/orpo-future-enhancements.md

@@ -0,0 +1,298 @@
+# ORPO future enhancements
+
+Performance-first follow-on ideas for mlx-bun's ORPO trainer, plus a short
+appendix of objective variants that may be useful once the runtime path is
+comfortable.
+
+> **LANDED since this was written (2026-06-19):** the top items here are now in
+> production — **Priority 1 (Cut Cross Entropy / fused linear CE)** and **§1
+> (vocab-axis chunked CE)** shipped as the `[M,vocab]`-free **flash-CCE Metal head**
+> (verbatim MLX steel GEMM, fwd+bwd; `orpo_flash_ce`), and **§3 (prefix sharing in the
+> trainer)** shipped as `orpo_prefix_shared`, composed with the flash head AND the
+> segmented backward for MiniCPM5 **and** e4b. See [training.md](../reference/training.md),
+> [orpo-quickstart.md](../reference/orpo-quickstart.md), and the live record
+> [steel-flash-cce-handoff.md](../investigations/steel-flash-cce-handoff.md). The
+> remaining sections (§2, §4–§9, the objective variants) are still genuinely future.
+
+Product relevance: ORPO is not only a researcher feature. It is the path to
+small, downloadable task adapters that make the default local assistant better
+without forcing every user onto a larger base model. The performance work below
+matters when it turns adapter training from a fragile lab run into a repeatable
+part of the mlx-bun product loop.
+
+## Priority 1: Cut Cross Entropy / Liger Fused Linear CE
+
+The top ORPO performance priority is a Cut Cross Entropy / Liger-style Fused
+Linear Cross Entropy path for the LM head. This should be treated as the next
+major kernel project before lower-impact objective variants.
+
+Why it is first:
+
+- Gemma's vocab is huge, so materializing `[tokens, vocab]` logits dominates
+  peak memory in ORPO.
+- ORPO needs chosen and rejected response log-probs, so the head cost is paid
+  twice unless the trainer can share or fuse more work.
+- Token-axis chunking already helps fit the run, but vocab-axis chunking cuts
+  the largest remaining logits allocation at the source.
+- The same fused engine can support SFT, DPO, ORPO, SimPO, IPO, and hinge-style
+  preference losses once it exposes mean log-prob with a custom upstream
+  cotangent.
+- A Metal fused version can avoid extra MLX dispatches and intermediate memory
+  traffic in the hot training path.
+
+The implementation target is not just "chunk logits"; it is a fused log-prob
+head that computes the linear projection, online vocab logsumexp, target-logit
+gather, loss reduction, and backward replay without ever storing full logits.
+
+## Performance focus
+
+### 1. Vocab-axis chunked cross entropy
+
+The current ORPO head optimization is **token-axis chunking**: split the response
+positions into chunks, run the LM head on `[chunk, vocab]`, and rematerialize
+those chunks in the backward. That bounds peak logits by `chunk_tokens * vocab`.
+
+The complementary trick is **vocab-axis chunking**: split the vocabulary instead,
+run the head on `[tokens, vocab_chunk]`, and combine chunks with an
+online-logsumexp recurrence. This is the idea behind **Cut Cross Entropy** and
+Liger-style **Fused Linear Cross Entropy**. This is priority 1 for the ORPO
+performance roadmap.
+
+Exact forward sketch:
+
+1. For each vocab chunk, compute partial logits `h @ W_chunk`.
+2. Gather the target logit only if the target id falls inside this chunk.
+3. Maintain per-token online `logsumexp` across vocab chunks.
+4. At the end, `logp = target_logit - lse`.
+
+Backward reconstructs the softmax chunk-by-chunk from the saved per-token LSE
+and never materializes full `[tokens, vocab]` logits. For Gemma's 262k vocab,
+this is likely the next real moat after segmented backward.
+
+Two implementation levels:
+
+- Pure MLX chunked VJP first: exact, easier to validate, probably enough to fit
+  longer contexts.
+- Metal fused kernel later: fuse quantized head matmul + online LSE + target
+  gather + backward. This is the high-performance version.
+
+### 2. Fold ORPO gradients into the fused CE path
+
+ORPO needs two mean log-probs, `lw` and `lr`, then applies the odds-ratio loss.
+For a fused CE/log-prob kernel, the clean design is to feed the upstream scalar
+gradient for each branch back into the chunked head backward. In other words,
+the fused head should expose "mean logp with custom upstream cotangent", not
+only cross-entropy.
+
+This avoids materializing logits while still letting ORPO, SimPO, and future
+preference losses share the same head engine.
+
+### 3. Prefix sharing in the trainer
+
+`src/train/prefix-shared.ts` already proves the MiniCPM5 shared-prefix idea:
+one forward over `[prompt; chosen; rejected]` with block-sparse masking and
+block-wise RoPE. It is not yet wired into `orpoLoop`.
+
+Next useful path:
+
+- Add an opt-in flag, for example `orpo_prefix_sharing`.
+- Gate it by prompt/response ratio so short-prompt data stays on the simple
+  two-forward path.
+- Wire MiniCPM5 first.
+- Port e4b later, where sliding-window masks, donor-KV reuse, and per-layer
+  input make it trickier.
+
+### 4. Sequence-chunk the MLP sub-block
+
+The MLP is per-token independent. Inside segmented backward, run MLP blocks over
+sequence chunks so the GeGLU/SwiGLU intermediate is `chunk_tokens * d_ff` instead
+of `seq_len * d_ff`.
+
+This stacks with depth segmentation and vocab-axis CE:
+
+`layer segment -> layer -> MLP seq chunk -> vocab chunk`
+
+Attention cannot be chunked this naively, but the MLP can.
+
+### 5. Rotor / budget-aware segment planner
+
+Replace fixed `segment_size` with a planner that chooses segment boundaries from
+estimated activation bytes and a target memory budget. The planner should isolate
+expensive full-attention layers and larger MLP regions instead of treating every
+layer equally.
+
+This is the practical version of rotor/checkpoint scheduling for this codebase:
+less hand tuning, better fit at 8K+.
+
+### 6. Faster flash-attention backward
+
+The custom flash backward is now correctness-useful but too slow. A serious
+performance version should mimic MLX's tuned SDPA style: simdgroup tiling,
+larger threadgroups, and FA2-style backward that stores only per-row LSE and
+recomputes attention probabilities.
+
+This matters most when `ops.sdpa` backward's O(L^2) memory is the blocker.
+
+### 7. Fused LoRA epilogue
+
+Today LoRA adds multiple ops after the quantized base matmul:
+
+`base(x) + scale * ((x @ A) @ B)`
+
+A fused dequant-GEMM + LoRA epilogue would reduce dispatches and intermediate
+traffic. Backward can recompute the cheap rank projection:
+
+- `h = x @ A`
+- `dB = scale * h.T @ dy`
+- `dA = scale * x.T @ (dy @ B.T)`
+- `dx = dy @ W.T + scale * (dy @ B.T) @ A.T`
+
+This is more kernel work, but it benefits SFT, DPO, ORPO, and serving with
+adapters.
+
+### 8. Shape-specialized compile for training graphs
+
+Shapeless compile is mainly required when a dimension changes every replay, such
+as autoregressive decode with a growing KV cache:
+
+`active_kv_len = prompt_len + generated_tokens`
+
+ORPO training is different. Each dataset row already has concrete lengths:
+
+- `P`: prompt length
+- `Rc`: chosen response length
+- `Rr`: rejected response length
+- `T`: prefix-shared length, `P + Rc + Rr`
+
+That means ORPO does not need shapeless compile in the same way decode does. A
+training graph can be compiled for an exact `(P, Rc, Rr)` shape, or for a padded
+length bucket, and cached under a key that includes the model, flags, segment
+plan, token/vocab chunk sizes, and objective mode.
+
+This is especially relevant for e4b prefix-sharing: the proof path should stay
+eager first, but a later L3 speed tier can use shape-specialized compiled graphs
+for known training shapes instead of requiring every primitive to support
+shapeless output-shape inference.
+
+The deeper L3 framing is **model-checksum-specific specialization**. For a fixed
+model snapshot, we know the architecture, layer count, hidden sizes, quantization
+layout, RoPE geometry, vocab, layer types, KV-sharing pattern, and supported
+context budget in advance. It is acceptable to spend much more work up front if
+that buys a faster hot path later.
+
+Practical version:
+
+- Treat the model snapshot checksum as part of the compilation key.
+- Generate/compile/cache a table of exact graph variants for the shapes and
+  modes we care about.
+- Bucket or pad examples into those graph variants when helpful.
+- Fall back to eager only for shapes/modes not in the table.
+
+This is effectively a lookup table of specialized execution plans: more compile
+and analysis cost up front, less branching, shape handling, and dispatch overhead
+at runtime.
+
+### 9. Fixed-shape decode cache as a compile route
+
+Decode is harder because the active KV length grows by one token per step. The
+current shapeless-compile path exists to avoid recompiling for:
+
+`L = 17, 18, 19, 20, ...`
+
+A separate route is to make decode shapes fixed:
+
+1. Preallocate K/V buffers to a known maximum context length.
+2. Write the new token into a dynamic position.
+3. Carry `active_length` / `write_position` as scalar array inputs.
+4. Mask out inactive cache positions, or preferably use an
+   active-length-aware attention kernel that simply does not scan them.
+
+The correctness version is "fixed cache + mask". The performance version is
+"fixed cache + active-length-aware attention"; otherwise early decode steps may
+pay for the full max context even when only a short prefix is active.
+
+There is also a more aggressive specialization route: prebuild a graph table for
+specific active lengths or length buckets. Even if that means many compiled graph
+variants, the trade can be valid for a fixed model/checksum if compile cost is
+paid once and reused. In that regime, the important design question is not "can
+one shapeless graph infer every shape?" but "which finite set of exact shapes and
+runtime states should we specialize for?"
+
+If every op in a chosen route is compile-safe, the whole decode graph can become
+a stable compiled graph. Custom Metal kernels still need to stay outside compiled
+regions unless the route is intentionally eager/custom-kernel-only or MLX compile
+can accept the primitive's output contract.
+
+## Quality and robustness ideas
+
+These are lower priority if the current goal is performance, but they are useful
+once ORPO can run comfortably.
+
+### SimPO
+
+SimPO is a reference-free preference loss that uses the **average log-probability
+of each response** as the reward and adds a target reward margin. It reuses the
+same `lw` / `lr` values ORPO already computes.
+
+Rough shape:
+
+`loss = softplus(-beta * ((lw - lr) - gamma))`
+
+where `gamma` is the desired chosen-vs-rejected margin. It should be cheap to add
+as a sibling objective because the expensive part, response mean log-probs, is
+already implemented.
+
+### IPO
+
+IPO means **Identity Preference Optimization**. It was introduced as a more
+overfit-resistant alternative to the original DPO logistic objective. Instead of
+driving preference probabilities toward hard 0/1 labels indefinitely, IPO uses a
+smoother identity-style objective over preference gaps. The practical point:
+when preference labels are noisy, IPO is often less eager to push margins to
+extremes.
+
+For mlx-bun, IPO is interesting mainly as a loss-family option after the shared
+logp engine exists.
+
+### Hinge / SLiC-style loss
+
+The hinge loss is a margin-ranking objective:
+
+`loss = max(0, margin - reward_gap)`
+
+For preference training, `reward_gap` is usually the chosen score minus rejected
+score. Once the gap exceeds the margin, that pair contributes no gradient. This
+can be useful when you want to stop over-optimizing already-solved pairs and
+focus compute on pairs whose preference margin is still too small.
+
+### Label smoothing and robust weighting
+
+Preference pairs are noisy. Cheap robustness knobs:
+
+- `label_smoothing`: treat a pair as, for example, 95% chosen / 5% rejected
+  rather than a hard label.
+- Per-example `weight`: let curated data carry confidence or score gap.
+- Loss clipping / winsorization: cap extreme high-loss pairs so noisy examples
+  do not dominate gradients.
+- Ambiguity filtering: skip pairs with tiny initial margin or near-identical
+  chosen/rejected responses.
+
+### Better curation
+
+The current UltraFeedback curation mostly extracts rows and filters by length.
+Quality can likely improve by filtering for:
+
+- Large enough chosen-vs-rejected score gap when available.
+- Similar response length, or at least bounded length ratio.
+- High semantic contrast but not unrelated answers.
+- Instruction-following categories if IFEval is the target metric.
+
+Data quality may beat objective cleverness on small ORPO runs.
+
+## Short answer to "token vs vocab"
+
+The thing to look for is **vocab-axis chunked / fused cross entropy**, especially
+under the names **Cut Cross Entropy** or **Liger Fused Linear Cross Entropy**.
+Token chunking slices `[tokens, vocab]` by rows. Vocab chunking slices it by
+columns and combines chunks with online logsumexp so the full logits matrix is
+never live.