mlx-bun 0.0.8 → 0.0.9

package/src/spec/generate.ts

@@ -1,11 +1,17 @@
 // Speculative decoding loop — port of optiq runtime/spec/runtime.py
 // (γ-draft / verify-in-one-pass / accept-longest-prefix / rollback).
 //
-// Correctness contract: greedy spec output is token-for-token identical
-// to greedy non-spec decode — the verify step re-derives every emitted
-// token with the SAME pick function as the stock path (logprobs argmax
-// over [1, V]); drafts only ever skip target forwards, never change
-// what the target would emit.
+// Correctness oracle: OPTIQ's spec_generate, NOT stock decode. Assistant-
+// drafter speculation is an optiq feature (mlx-lm has only generic
+// two-model spec, and can't drive a KV-borrowing drafter). Both optiq AND
+// mlx-lm BATCH the verify lm-head (one matmul over the whole γ+1 window,
+// then argmax per position) — so neither is bit-exact to stock token-at-
+// a-time decode; they diverge from it at bf16 knife-edges. We do the same
+// (picksBatched): greedy spec is bit-exact to optiq's spec, and agrees
+// with stock only on tie-free prompts (where batched == per-position).
+// An earlier version verified per-position to stay bit-exact to STOCK —
+// that matched no real oracle and cost an extra γ× read of the lm-head
+// weight per step; superseded. See docs/design/spec-decode-larger-targets.md.
 //
 // Limitation (same as the reference): partial-accept rollback requires
 // trimmable caches; RotatingKVCache loses trimability once its ring
@@ -14,7 +20,6 @@
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { Gemma4Model, type Cache } from "../model/gemma4";
-import { toLogprobs } from "../sampler";
 import type { GemmaAssistantDrafter } from "./drafter";
 
 export interface SpecOptions {
@@ -39,12 +44,8 @@ export interface SpecResult {
   stats: SpecStats;
 }
 
-/** The stock path's greedy pick, bit-identical shapes: slice ONE hidden
- *  position [1,1,H], lm-head on it (exactly what generate() does for a
- *  decode step), then logprobs argmax over [1, V]. Computing the lm-head
- *  per position (instead of batched over the verify window) keeps every
- *  pick in the same kernel shapes as stock decode — batched lm-head
- *  rounds differently and flips bf16 knife-edges. */
+/** Single greedy pick from one hidden position [1,1,H] (lm-head + argmax).
+ *  Used for the NON-speculative prefill token (matches stock there). */
 function pickFromHidden(model: Gemma4Model, hidden: MlxArray, pos: number): number {
   const H = hidden.shape[2]!;
   const hSl = hidden.slice([0, pos, 0], [1, pos + 1, H]);
@@ -53,15 +54,36 @@ function pickFromHidden(model: Gemma4Model, hidden: MlxArray, pos: number): numb
   const V = logits.shape[2]!;
   const flat = ops.reshape(logits, [1, V]);
   logits.dispose();
-  const lp = toLogprobs(flat);
+  const am = ops.argmaxAxis(flat, -1);
   flat.dispose();
-  const am = ops.argmaxAxis(lp, -1);
-  lp.dispose();
   const t = ops.itemUint32(am);
   am.dispose();
   return t;
 }
 
+/** optiq-faithful batched verify: ONE lm-head over the whole γ+1 window
+ *  (reads the lm-head weight ONCE, as optiq and mlx-lm both do), then
+ *  argmax per position. This is what makes greedy spec bit-exact to
+ *  optiq's spec_generate. argmax is invariant to the logsumexp shift and
+ *  to the monotone final-logit softcap, so raw-logit argmax here picks the
+ *  same token optiq does. */
+function picksBatched(model: Gemma4Model, hidden: MlxArray, count: number): number[] {
+  const logits = model.logitsFromHidden(hidden); // [1, L, V] — one matmul
+  const V = logits.shape[2]!;
+  const out: number[] = [];
+  for (let k = 0; k < count; k++) {
+    const sl = logits.slice([0, k, 0], [1, k + 1, V]);
+    const flat = ops.reshape(sl, [1, V]);
+    sl.dispose();
+    const am = ops.argmaxAxis(flat, -1);
+    flat.dispose();
+    out.push(ops.itemUint32(am));
+    am.dispose();
+  }
+  logits.dispose();
+  return out;
+}
+
 /** Donor caches: the target's LAST sliding and LAST full cache owners
  *  (port of optiq kv_view.find_donor_layers). */
 function donorIndices(model: Gemma4Model): { sliding: number; full: number } {
@@ -163,10 +185,9 @@ export function specGenerate(
       vIds.dispose();
       stats.targetCalls++;
 
-      // 2c. accept the longest matching prefix (per-position lm-head —
-      // see pickFromHidden)
-      const gt: number[] = [];
-      for (let k = 0; k <= gamma; k++) gt.push(pickFromHidden(model, vHidden, k));
+      // 2c. accept the longest matching prefix. Batched verify lm-head
+      // (optiq-faithful — see picksBatched).
+      const gt = picksBatched(model, vHidden, gamma + 1);
       let kAccept = 0;
       while (kAccept < gamma && drafts[kAccept] === gt[kAccept]) kAccept++;
       stats.accepted += kAccept;

package/src/status-page.html

@@ -285,7 +285,7 @@
           <thead><tr><th>Command</th><th style="text-align:left">What it does</th></tr></thead>
           <tbody>
             <tr><td><code>mlx-bun pi</code></td><td style="text-align:left">Agent session on a local model — picks/downloads the right model, serves, drops into pi. Flags: <code>--query</code>, <code>--port</code>; everything else passes to pi (<code>-p</code>, <code>--continue</code>, <code>@file</code>…)</td></tr>
-            <tr><td><code>mlx-bun serve &lt;query&gt;</code></td><td style="text-align:left">Persistent server + this page. <code>--port 8090</code>, <code>--memory-budget GB</code></td></tr>
+            <tr><td><code>mlx-bun serve &lt;query&gt;</code></td><td style="text-align:left">Persistent server + this page. <code>--port 8080</code>, <code>--memory-budget GB</code></td></tr>
             <tr><td><code>mlx-bun get &lt;org/repo&gt;</code></td><td style="text-align:left">Resumable, sha256-verified model download</td></tr>
             <tr><td><code>mlx-bun ls</code> / <code>fit</code></td><td style="text-align:left">List downloaded models / will-it-fit assessment (<code>--ctx</code>, <code>--skus</code>)</td></tr>
             <tr><td><code>mlx-bun harness pi</code></td><td style="text-align:left">Register this server as a provider in your own pi install (<code>--remove</code> undoes)</td></tr>

package/src/tokenizer.ts

@@ -11,6 +11,10 @@ export interface LoadedTokenizer {
    *  you get a corrupting double-BOS `<s><s>…`). */
   encode(text: string, addSpecialTokens?: boolean): number[];
   decode(ids: number[], skipSpecialTokens?: boolean): string;
+  /** Token id → raw vocab token string (mlx-lm's convert_ids_to_tokens —
+   *  the undecoded piece, e.g. "▁Hello", used in logprobs responses).
+   *  Falls back to decode() for ids outside the base vocab (added tokens). */
+  idToToken(id: number): string;
   readonly bosTokenId: number | null;
   readonly eosTokenId: number | null;
 }
@@ -36,6 +40,9 @@ export async function loadTokenizer(modelDir: string): Promise<LoadedTokenizer>
     decode: (ids, skipSpecialTokens = false) =>
       // python's decode([]) === ""; the JS lib throws on empty input
       ids.length === 0 ? "" : tok.decode(ids, { skip_special_tokens: skipSpecialTokens }),
+    idToToken: (id) =>
+      (tok as unknown as { id_to_token(id: number): string | undefined })
+        .id_to_token(id) ?? tok.decode([id], { skip_special_tokens: false }),
     bosTokenId: idOf("bos_token"),
     eosTokenId: idOf("eos_token"),
   };

package/src/train/flash-cce.ts

@@ -723,12 +723,16 @@ const BWD_SG_SOURCE = String.raw`
   }
 `;
 
-// Apple-CCE coeff filter epsilon (phase-2 per-tile skip). DEFAULT OFF ("0") — the
-// Apple-production shortcut skips vocab tiles whose whole coeff block < eps, which gave
-// the OLD (slow, phase-2-dominated) SG kernel ~3.5×. On the steel kernel — whose phase-2
-// is now QuantizedBlockLoader-fast — it is a BAD trade: measured e4b M=512 REAL it cut dh
-// accuracy 0.66%→2.7% for only ~7% speedup (the win was already captured by the fast
-// dequant). Opt in with MLX_BUN_CCE_BWD_FILTER_EPS=1e-5 only on genuinely peaked data.
+// Apple-CCE coeff filter epsilon (phase-2 per-tile skip). DEFAULT OFF ("0") today —
+// no production path sets the env — but the VERDICT is measured and positive: the old
+// 0.66%→2.7% dh-accuracy scare was on RANDOM targets (flat softmax, the filter's
+// worst case — see PLAN.md "filter-on-real-data"); the real-data gate
+// (scripts/experiments/flash-cce-filter-realdata.ts, commit 3d259ef, 2026-06-22)
+// measured eps 1e-5 = 0.16% dh (under the bf16 floor) for a 1.35× faster backward,
+// and PLAN.md's recorded decision is "enable the filter at eps ~1e-5". (The ~3.5×
+// figure floating around older notes was the retired SG kernel, not steel.) Flipping
+// the default is the open backlog item; until then opt in with
+// MLX_BUN_CCE_BWD_FILTER_EPS=1e-5 on real (peaked) training data.
 const BWD_FILTER_EPS = process.env.MLX_BUN_CCE_BWD_FILTER_EPS ?? "0";
 // blockMax vocab-block early-exit eps (skip whole cold (token-block,vocab-block) programs
 // — attacks BOTH phases). DEFAULT OFF ("0"). It is LOSSLESS when it fires (only skips
@@ -1005,6 +1009,22 @@ function fwdSteelKernel(gs: number, bits: number): MetalKernel {
   return k;
 }
 
+// One-time dispatch-route warnings. A failed tiling guard silently reroutes the
+// whole flash head onto the legacy kernels (~5× slower class: e4b M=512 backward
+// 754 → 3687+ ms) with no observable signal — the exact "forward passes are
+// flag-gated, verify in code" hazard. Name the failed guard once per route so a
+// new model port that misses steel is visible in the training log.
+const _steelFallbackWarned = new Set<string>();
+function warnSteelFallback(which: "forward" | "backward", detail: string, fallback: string): void {
+  if (_steelFallbackWarned.has(which)) return;
+  _steelFallbackWarned.add(which);
+  console.warn(
+    `[flash-cce] ${which}: steel kernel guard failed (${detail}) — dispatching the ` +
+      `${fallback} instead (~5x-slower class). If this model should be steel-capable, ` +
+      `check the tiling guards in flash-cce.ts.`,
+  );
+}
+
 export function flashCceForward(
   hResp: MlxArray, head: FlashCceHead, targets: number[],
 ): { logp: MlxArray; lse: MlxArray; blockMax: MlxArray } {
@@ -1039,6 +1059,15 @@ export function flashCceForward(
   const blockV0 = Math.ceil(V / NBLK);
   const useSteel = process.env.MLX_BUN_CCE_NOSTEEL !== "1" && H % 32 === 0 && V % 32 === 0 && blockV0 % 32 === 0;
   const useSimd = !useSteel && USE_SIMDGROUP_FWD && H % SG_BD === 0;
+  if (!useSteel && process.env.MLX_BUN_CCE_NOSTEEL !== "1") {
+    // NOTE blockV0 = ceil(V/NBLK) can violate %32 even when V%32==0 (e.g. V=50304)
+    const fails = [
+      H % 32 !== 0 ? `H%32!=0 (H=${H})` : null,
+      V % 32 !== 0 ? `V%32!=0 (V=${V})` : null,
+      blockV0 % 32 !== 0 ? `blockV0%32!=0 (blockV0=${blockV0}, V=${V}, NBLK=${NBLK})` : null,
+    ].filter(Boolean).join(", ");
+    warnSteelFallback("forward", fails, useSimd ? "simdgroup kernel" : "scalar/lane kernel");
+  }
   const useLane = !useSteel && !useSimd && process.env.MLX_BUN_CCE_LANE === "1" && WPR % 32 === 0;
   const kernel = useSteel ? fwdSteelKernel(GS, head.bits) : useSimd ? fwdSgKernel() : useLane ? fwdLaneKernel() : fwdKernel();
   const tokTile = useSteel ? STEEL_BM : useSimd ? SG_BB : BLOCK_B;
@@ -1103,10 +1132,11 @@ export function flashCceForward(
  *  (no atomics). Caller owns the result. */
 export function flashCceBackward(
   hResp: MlxArray, head: FlashCceHead, targets: number[], lse: MlxArray, cot: number[],
-  // Apple-CCE coeff filter eps. Default ON (BWD_FILTER_EPS) for the ORPO training
-  // path — a pretrained model's softmax is peaked, so skipping near-zero-coeff rows
-  // is near-lossless and ~3.5× faster. Pass "0" for exact gradients (the parity gate
-  // does this, since synthetic-flat softmax has no negligible rows to skip).
+  // Apple-CCE coeff filter eps. Defaults to BWD_FILTER_EPS, which is OFF ("0")
+  // unless MLX_BUN_CCE_BWD_FILTER_EPS is set — exact gradients by default. The
+  // measured real-data case FOR enabling it (eps 1e-5 = 0.16% dh, 1.35×; PLAN.md
+  // "filter-on-real-data") is a pending default flip — see BWD_FILTER_EPS's note.
+  // Pass "0" explicitly for guaranteed-exact (the parity gate does).
   filterEps: string = BWD_FILTER_EPS,
   // The forward's `blockMax [M, NBLK]` enables the vocab-block early exit (skips
   // whole cold programs — attacks phase 1). Omit (or set blockEps "0") for exact.
@@ -1119,7 +1149,10 @@ export function flashCceBackward(
   const GS = head.groupSize;
   const WPR = head.w.shape[1]!;
   if (H > H_MAX) throw new Error(`flashCceBackward: hidden ${H} > H_MAX ${H_MAX}`);
-  if (H % TG !== 0) throw new Error(`flashCceBackward: hidden ${H} not divisible by TG ${TG}`);
+  // NOTE: H % TG is a LEGACY-SCALAR-kernel requirement only (it owns DPER = H/TG
+  // dims per thread — floor division, so a remainder would leave dims silently
+  // zero). The steel/simdgroup kernels have their own (weaker) guards below; the
+  // throw is applied only when dispatch actually lands on the scalar kernel.
 
   const hF = hResp.dtype === Dtype.float32 ? hResp : hResp.astype(Dtype.float32);
   const scF = head.scales.dtype === Dtype.float32 ? head.scales : head.scales.astype(Dtype.float32);
@@ -1133,10 +1166,12 @@ export function flashCceBackward(
 
   // STEEL is the DEFAULT (verbatim MLX BlockMMA GEMM for both phases, [M,V]-free) when
   // the tiling is clean (H,V,blockV all %32 AND H % HTw for the H-tiled accumulator).
-  // MLX_BUN_CCE_BWD_NOSTEEL=1 falls back to the old kernels. The steel backward now
-  // carries the SAME Apple-CCE skips as the SG kernel — coeff filter (phase-2 per-tile)
-  // + blockMax vocab-block early exit (whole-program) — both compiled out at eps=0
-  // (exact, the parity gate) and default-on for training (peaked softmax → ~all vocab cold).
+  // MLX_BUN_CCE_BWD_NOSTEEL=1 falls back to the old kernels. The steel backward carries
+  // the SAME Apple-CCE skips as the SG kernel — coeff filter (phase-2 per-tile) +
+  // blockMax vocab-block early exit (whole-program) — both compiled out at eps=0.
+  // BOTH skips are OFF by default (the envs default "0", nothing in production sets
+  // them): exact gradients unless explicitly opted in. Enabling the coeff filter at
+  // ~1e-5 is the measured-and-decided pending flip (PLAN.md "filter-on-real-data").
   const blockV0 = Math.ceil(V / NBLK);
   // HTw == group_size for the phase-2 QuantizedBlockLoader (one group per H-tile); needs
   // H % GS == 0 (always, GR*GS=H) and GS % 32 == 0 (TNv = GS/32 frags per H-tile integral).
@@ -1144,6 +1179,26 @@ export function flashCceBackward(
   // simdgroup_matrix phase-2 GEMM is the next tier (5512→3913 ms, 1.41×) when the H
   // tiling is clean; MLX_BUN_CCE_BWD_SCALAR=1 forces the float4 scalar kernel.
   const useSg = !useSteel && process.env.MLX_BUN_CCE_BWD_SCALAR !== "1" && H % (SGN * 8) === 0 && H % SG_BD === 0;
+  if (!useSteel && process.env.MLX_BUN_CCE_BWD_NOSTEEL !== "1") {
+    // NOTE blockV0 = ceil(V/NBLK) can violate %32 even when V%32==0 (e.g. V=50304)
+    const fails = [
+      H % 32 !== 0 ? `H%32!=0 (H=${H})` : null,
+      V % 32 !== 0 ? `V%32!=0 (V=${V})` : null,
+      blockV0 % 32 !== 0 ? `blockV0%32!=0 (blockV0=${blockV0}, V=${V}, NBLK=${NBLK})` : null,
+      H % GS !== 0 ? `H%GS!=0 (H=${H}, GS=${GS})` : null,
+      GS % 32 !== 0 ? `GS%32!=0 (GS=${GS})` : null,
+    ].filter(Boolean).join(", ");
+    warnSteelFallback("backward", fails, useSg ? "simdgroup kernel" : "scalar kernel");
+  }
+  // Legacy scalar kernel: each thread owns DPER = H/TG dims (floor) — a remainder
+  // would leave the last H%TG dims silently zero, so it hard-requires H % TG == 0.
+  // Only enforced when dispatch actually lands there (steel/SG have weaker guards).
+  if (!useSteel && !useSg && H % TG !== 0)
+    throw new Error(
+      `flashCceBackward: hidden ${H} not divisible by TG ${TG} — required by the legacy ` +
+        `scalar kernel, which is the selected route because the steel/simdgroup tiling ` +
+        `guards also failed for this shape (see the [flash-cce] warning above).`,
+    );
   // Block-skip is on (either kernel) only when blockMax is supplied AND blockEps>0 (the
   // kernel reads it as an extra input). The coeff filter needs no extra input (just eps).
   const useBlockSkip = (useSteel || useSg) && blockMax != null && Number.parseFloat(blockEps) > 0;

package/src/train/fuse.ts

@@ -0,0 +1,243 @@
+// Fuse a LoRA adapter into base model weights → a standalone model snapshot
+// (mlx_lm.fuse counterpart, minus GGUF/de-quantize export).
+//
+// Math (matches mlx-lm's LoRALinear.fuse and our QuantizedLinear.forward):
+//   forward:  y = x @ W_dequantᵀ + scale · ((x @ A) @ B)
+//   fused:    W'[out,in] = W + (scale · B)ᵀ.astype(bf16) @ Aᵀ.astype(bf16)
+// mlx-lm casts the (scaled) LoRA factors to the base dtype BEFORE the matmul,
+// then re-quantizes with the module's OWN source spec (bits/group unchanged) —
+// we do the same, so the output config.json is the source config verbatim.
+//
+// Everything runs on the CPU stream (same discipline as src/quantize/quantizer:
+// the safetensors loader lives there, and one module is materialized, fused,
+// and disposed at a time). Modules the adapter does not touch pass through
+// bit-identical (no dequant→requant round trip).
+
+import { existsSync, mkdirSync } from "node:fs";
+import { join } from "node:path";
+import { clearCache, Dtype } from "../mlx/ffi";
+import { cpuStream, type MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { Weights } from "../weights";
+import { loadModelConfig, quantFor } from "../config";
+import { loadAdapterTensors, readAdapterScale, adapterWeightsFile } from "../lora";
+import { writeShardedSafetensors, type NamedTensor } from "../quantize/safetensors-writer";
+import { copyAuxFiles } from "../quantize/config-writer";
+
+/** Quantize the allocator cache this often (in modules processed). */
+const CACHE_CLEAR_EVERY = 16;
+
+export interface FuseStats {
+  outDir: string;
+  /** Modules whose weights had an adapter delta folded in. */
+  fusedModules: number;
+  /** Adapter tensors that matched no base weight (skipped). */
+  skippedAdapterTensors: number;
+  totalTensors: number;
+}
+
+export interface FuseProgress {
+  stage: string;
+  message: string;
+  /** 0..1 fraction of modules processed. */
+  progress: number;
+}
+
+/** Resolve the adapter safetensors file: `<dir>/adapters.safetensors`,
+ *  `<dir>/best/adapters.safetensors`, `adapter_model.safetensors`, or an
+ *  explicit file path. */
+function resolveAdapterFile(dir: string): string {
+  const best = join(dir, "best");
+  if (existsSync(join(best, "adapters.safetensors"))) return join(best, "adapters.safetensors");
+  if (dir.endsWith(".safetensors") && existsSync(dir)) return dir;
+  return adapterWeightsFile(dir);
+}
+
+/** Group adapter tensors into (modulePath → {a, b}); disposes non-LoRA
+ *  tensors and returns how many were skipped. */
+function groupPairs(tensors: Map<string, MlxArray>): {
+  pairs: Map<string, { a?: MlxArray; b?: MlxArray }>;
+  skipped: number;
+} {
+  const pairs = new Map<string, { a?: MlxArray; b?: MlxArray }>();
+  let skipped = 0;
+  for (const [name, arr] of tensors) {
+    const m = name.match(/^(.*)\.(lora_a|lora_A|lora_b|lora_B)(\.weight)?$/);
+    if (!m) { skipped++; arr.dispose(); continue; }
+    const slot = pairs.get(m[1]!) ?? {};
+    if (m[2]!.toLowerCase() === "lora_a") slot.a = arr; else slot.b = arr;
+    pairs.set(m[1]!, slot);
+  }
+  return { pairs, skipped };
+}
+
+/**
+ * Fuse the adapter at `adapterDir` into the base model at `modelDir` and
+ * write a complete, loadable snapshot to `outDir` (weights + config.json +
+ * tokenizer aux files). Quantization layout is preserved verbatim: fused
+ * modules are dequantized, updated, and re-quantized with their own source
+ * spec; untouched modules pass through bit-identical.
+ */
+export async function fuseAdapter(
+  modelDir: string,
+  adapterDir: string,
+  outDir: string,
+  onProgress?: (e: FuseProgress) => void,
+): Promise<FuseStats> {
+  const progress = (stage: string, message: string, frac: number) =>
+    onProgress?.({ stage, message, progress: frac });
+
+  progress("loading", `Reading ${modelDir}`, 0);
+  const config = await loadModelConfig(modelDir);
+  const weights = await Weights.open(modelDir);
+
+  const adapterFile = resolveAdapterFile(adapterDir);
+  const { scale, rsLora } = await readAdapterScale(adapterDir);
+  const { pairs, skipped } = groupPairs(loadAdapterTensors(adapterFile));
+
+  // Map each adapter module path to the base weight-file module path,
+  // probing the VLM-wrapped/unwrapped `language_model.` prefixes both ways
+  // (same normalization as AdapterManager.mount).
+  const weightNames = new Set(weights.tensorNames);
+  const byModule = new Map<string, { a: MlxArray; b: MlxArray }>();
+  let unmatched = 0;
+  const disposeAll = () => {
+    for (const { a, b } of pairs.values()) { a?.dispose(); b?.dispose(); }
+  };
+  try {
+    for (const [path, { a, b }] of pairs) {
+      if (!a || !b) throw new Error(`${path}: adapter has only one of lora_a/lora_b`);
+      const candidates = [
+        path,
+        `language_model.${path}`,
+        path.startsWith("language_model.") ? path.slice("language_model.".length) : null,
+      ].filter((c): c is string => c !== null);
+      const hit = candidates.find((c) => weightNames.has(`${c}.weight`));
+      if (!hit) { unmatched++; continue; }
+      byModule.set(hit, { a, b });
+    }
+    if (byModule.size === 0)
+      throw new Error(
+        `no adapter tensors match the base model's weights — was ` +
+        `${adapterFile} trained for a different base?`,
+      );
+  } catch (e) {
+    disposeAll();
+    weights.dispose();
+    throw e;
+  }
+
+  mkdirSync(outDir, { recursive: true });
+
+  // Walk tensors in source order (diff-friendly like the quantizer): group
+  // names into module bases so a fused module's weight/scales/biases are
+  // regenerated together.
+  const names = weights.tensorNames;
+  const moduleSuffixes = new Map<string, Set<string>>();
+  const passthroughNames: string[] = [];
+  for (const name of names) {
+    const m = /^(.*)\.(weight|scales|biases)$/.exec(name);
+    if (!m) { passthroughNames.push(name); continue; }
+    let set = moduleSuffixes.get(m[1]!);
+    if (!set) { set = new Set(); moduleSuffixes.set(m[1]!, set); }
+    set.add(m[2]!);
+  }
+
+  const out: NamedTensor[] = [];
+  let fusedModules = 0;
+  let processed = 0;
+  const totalModules = moduleSuffixes.size;
+
+  try {
+    for (const [base, suffixes] of moduleSuffixes) {
+      processed++;
+      if (processed % CACHE_CLEAR_EVERY === 0) clearCache();
+      const pair = suffixes.has("weight") ? byModule.get(base) : undefined;
+      if (!pair) {
+        for (const suf of suffixes) {
+          const n = `${base}.${suf}`;
+          out.push({ name: n, array: weights.tensor(n) });
+        }
+        continue;
+      }
+      progress("fusing", `Module ${processed}/${totalModules}: ${base}`, processed / totalModules);
+
+      // Materialize the full-precision weight [out, in].
+      const srcSpec = suffixes.has("scales") ? quantFor(config.quantization, base) : null;
+      let w: MlxArray;
+      if (srcSpec) {
+        const packed = weights.tensor(`${base}.weight`);
+        const scales = weights.tensor(`${base}.scales`);
+        const biases = suffixes.has("biases") ? weights.tensor(`${base}.biases`) : null;
+        w = ops.dequantize(packed, scales, biases, srcSpec, cpuStream);
+      } else {
+        w = weights.tensor(`${base}.weight`);
+      }
+      const [outF, inF] = w.shape as [number, number];
+      const { a, b } = pair;
+      const [aIn, rank] = a.shape as [number, number];
+      const [bRank, bOut] = b.shape as [number, number];
+      if (aIn !== inF || bRank !== rank || bOut !== outF)
+        throw new Error(
+          `${base}: shape mismatch — lora_a [${a.shape}] / lora_b [${b.shape}] ` +
+          `vs base [out ${outF}, in ${inF}]`,
+        );
+
+      // delta[out,in] = (scale·B)ᵀ.astype(bf16) @ Aᵀ.astype(bf16)
+      // (rsLoRA: effective per-layer scale is α/√rank, matching mount/training.)
+      const effScale = rsLora ? scale / Math.sqrt(rank) : scale;
+      const bScaled = ops.mulScalar(b, effScale, cpuStream);
+      const bT = ops.transposeAxes(bScaled, [1, 0], cpuStream);
+      bScaled.dispose();
+      const bT16 = bT.astype(Dtype.bfloat16, cpuStream);
+      bT.dispose();
+      const aT = ops.transposeAxes(a, [1, 0], cpuStream);
+      const aT16 = aT.astype(Dtype.bfloat16, cpuStream);
+      aT.dispose();
+      const delta = ops.matmul(bT16, aT16, cpuStream);
+      bT16.dispose();
+      aT16.dispose();
+
+      const w16 = w.dtype === Dtype.bfloat16 ? w : w.astype(Dtype.bfloat16, cpuStream);
+      if (w16 !== w) w.dispose();
+      const fusedW = ops.add(w16, delta, cpuStream);
+      w16.dispose();
+      delta.dispose();
+
+      if (srcSpec) {
+        const q = ops.quantize(fusedW, srcSpec.groupSize, srcSpec.bits, srcSpec.mode, cpuStream);
+        fusedW.dispose();
+        out.push({ name: `${base}.weight`, array: q.packed });
+        out.push({ name: `${base}.scales`, array: q.scales });
+        out.push({ name: `${base}.biases`, array: q.biases });
+      } else {
+        out.push({ name: `${base}.weight`, array: fusedW });
+      }
+      fusedModules++;
+    }
+    for (const name of passthroughNames) out.push({ name, array: weights.tensor(name) });
+
+    progress("writing", `Writing ${out.length} tensors to ${outDir}`, 1);
+    writeShardedSafetensors(outDir, out);
+
+    // Config is unchanged (same quantization layout) — copy it and the aux
+    // files (tokenizer, chat template, …) verbatim for a loadable snapshot.
+    await Bun.write(join(outDir, "config.json"), Bun.file(join(modelDir, "config.json")));
+    await copyAuxFiles(modelDir, outDir);
+
+    progress("done", `Fused ${fusedModules} module(s) → ${outDir}`, 1);
+    return {
+      outDir,
+      fusedModules,
+      skippedAdapterTensors: skipped + unmatched * 2,
+      totalTensors: out.length,
+    };
+  } finally {
+    for (const t of out) {
+      try { t.array.dispose(); } catch {}
+    }
+    disposeAll();
+    weights.dispose();
+    clearCache();
+  }
+}

package/src/train/index.ts

@@ -14,6 +14,7 @@ export { trainForward, buildBatchedPadMask } from "./forward";
 
 export {
   sftLoss, dpoLoss, dpoRefLogps, dpoMetrics,
+  orpoLoss, orpoMetrics, type SftScope,
 } from "./loss";
 
 export {
@@ -38,6 +39,8 @@ export {
 
 export { mergeAdapters, type MergeStats } from "./merge";
 
+export { fuseAdapter, type FuseStats, type FuseProgress } from "./fuse";
+
 export { exportAdapter, type ExportManifest } from "./export";
 
 export { finetuneRunner, inspectDataset } from "./job";

package/src/train/job.ts

@@ -51,6 +51,7 @@ interface FinetuneSubmit {
   orpo_fused_ce?: boolean;
   orpo_flash_ce?: boolean;
   orpo_prefix_shared?: boolean;
+  sft_scope?: "full" | "response";
   warm_start_adapter?: string;
 }
 
@@ -68,6 +69,12 @@ function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir:
   // quantized-head pieces if the base turns out to be unquantized.
   const isOrpo = (c.method ?? "sft") === "orpo";
 
+  // ORPO L_SFT scope: "full" (default; paper/TRL-faithful) | "response"
+  // (pre-2026-07 behavior). Validate here so a typo fails at submit, not mid-run.
+  const sftScope = c.sft_scope ?? DEFAULT_TRAIN_CONFIG.sftScope;
+  if (sftScope !== "full" && sftScope !== "response")
+    throw new Error(`finetune job: sft_scope must be "full" or "response" (got ${JSON.stringify(c.sft_scope)})`);
+
   const cfg: TrainConfig = {
     ...DEFAULT_TRAIN_CONFIG,
     method: c.method ?? "sft",
@@ -105,6 +112,7 @@ function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir:
     orpoFusedCe: c.orpo_fused_ce ?? DEFAULT_TRAIN_CONFIG.orpoFusedCe,
     orpoFlashCe: c.orpo_flash_ce ?? (isOrpo ? true : DEFAULT_TRAIN_CONFIG.orpoFlashCe),
     orpoPrefixShared: c.orpo_prefix_shared ?? (isOrpo ? true : DEFAULT_TRAIN_CONFIG.orpoPrefixShared),
+    sftScope,
     warmStartAdapter: c.warm_start_adapter ?? DEFAULT_TRAIN_CONFIG.warmStartAdapter,
     adapterPath: c.adapter_path,
     baseModel: c.model_dir,