mlx-bun 0.0.10 → 0.0.11

package/src/model/universal/dense.ts

@@ -1,7 +1,7 @@
 // UniversalDense — the Tier-0 config-driven dense llama-family module
 // (docs/design/generic-model-support.md §3.1). One module, shaped exactly
 // like qwen3.ts, with every arch delta selected by the UniversalArgs
-// descriptor (archs.ts). Monolith path only: no perf kernels, no compiled
+// descriptor (archs.ts). Monolith path only: no compiled
 // decode, no kv-quant — slow, never broken. The bar is L1 bit-exactness
 // vs mlx-lm on this machine's GPU (tests/universal-parity.test.ts).
 //
@@ -28,6 +28,7 @@ import {
   type Mask,
 } from "../gemma4-base";
 import { genericArgsFor, type UniversalArgs } from "./archs";
+import { compiledSwiglu } from "../minicpm5";
 import { initializeRope, NoRope, phi3SuRope, UniversalRope } from "./rope";
 import {
   loadEmbedding,
@@ -140,8 +141,18 @@ class UniversalAttention {
     k = disposing(k, ops.transposeAxes(k, [0, 2, 1, 3]));
     v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));
 
-    q = disposing(q, this.rope.apply(q, cache.offset));
-    k = disposing(k, this.rope.apply(k, cache.offset));
+    // Batched decode: the scheduler's mask wrapper exposes each row's REAL
+    // position as ropeOffsetArr (rows have different prompt lengths under
+    // left-padding; the scalar offset would mis-position every padded row by
+    // its pad amount — the 2026-07-03 uneven-batch bug). Serial lane: scalar
+    // offset, byte-identical to before (the S1b.1 pattern, cf. qwen3_5.ts).
+    const offArr = (cache as { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
+    q = disposing(q, offArr
+      ? this.rope.applyDynamic(q, offArr)
+      : this.rope.apply(q, cache.offset));
+    k = disposing(k, offArr
+      ? this.rope.applyDynamic(k, offArr)
+      : this.rope.apply(k, cache.offset));
 
     const [keys, values] = cache.updateAndFetch(k, v);
     k.dispose();
@@ -215,11 +226,23 @@ class UniversalAttention {
   }
 }
 
-class UniversalMLP {
+/** swiglu(gate, up) = nn.silu(gate) * up as ONE call — the whole activation,
+ *  not just silu. Defaults to the `@mx.compile` version (compiledSwiglu) — the
+ *  same libmlx kernel mlx-lm dispatches (L1-faithful), bit-identical to the old
+ *  uncompiled composition but one dispatch instead of three. Set to null to force
+ *  the uncompiled path (the slower same-parity opt-in). */
+export type SwigluFn = (gate: MlxArray, up: MlxArray) => MlxArray;
+
+export class UniversalMLP {
   readonly gate: AnyLinear | null;
   readonly up: AnyLinear | null;
   readonly gateUp: AnyLinear | null;
   readonly down: AnyLinear;
+  /** Fused swiglu (silu(gate)*up), default = the oracle's `@mx.compile` closure.
+   *  Only consumed for the swiglu / fused_swiglu kinds; geglu/geglu_approx/gelu_mlp
+   *  are unaffected (the oracle does NOT @mx.compile those — gemma/gemma2/starcoder2
+   *  use a plain nn.gelu composition). Set to null for the uncompiled composition. */
+  swigluFn: SwigluFn | null = compiledSwiglu;
 
   constructor(
     weights: Weights,
@@ -268,6 +291,18 @@ class UniversalMLP {
       up = this.up!.forward(x);
     }
 
+    // swiglu / fused_swiglu run silu(gate)*up as ONE fused call (mlx-lm's
+    // `@mx.compile swiglu`), matching the oracle kernel set exactly. This is the
+    // default (swigluFn = compiledSwiglu); null forces the uncompiled composition.
+    if (this.swigluFn && (this.kind === "swiglu" || this.kind === "fused_swiglu")) {
+      const hidden = this.swigluFn(gate, up);
+      gate.dispose();
+      up.dispose();
+      const out = this.down.forward(hidden);
+      hidden.dispose();
+      return out;
+    }
+
     let act: MlxArray;
     if (this.kind === "geglu") {
       act = ops.geluPrecise(gate); // gemma-1: nn.gelu (erf)
@@ -289,7 +324,7 @@ class UniversalMLP {
   }
 }
 
-class UniversalLayer {
+export class UniversalLayer {
   readonly attn: UniversalAttention;
   readonly mlp: UniversalMLP;
   /** pre/gemma2/glm4: input_layernorm. post (olmo2): unused (null). */

package/src/model/universal/rope.ts

@@ -49,32 +49,44 @@ export class UniversalRope {
   ) {}
 
   apply(x: MlxArray, offset: number): MlxArray {
-    let input = x;
-    let scaled: MlxArray | null = null;
-    if (this.preScaleAlways || this.preScale !== 1.0) {
-      // x[..., :dims] = preScale * x[..., :dims] (weak-scalar mul)
-      const D = x.shape[x.shape.length - 1]!;
-      if (this.dims === D) {
-        scaled = ops.mulScalar(x, this.preScale);
-      } else {
-        const start = x.shape.map(() => 0);
-        const stop = [...x.shape];
-        stop[stop.length - 1] = this.dims;
-        const head = x.slice(start, stop);
-        const headScaled = ops.mulScalar(head, this.preScale);
-        head.dispose();
-        scaled = ops.sliceUpdate(x, headScaled, start, stop);
-        headScaled.dispose();
-      }
-      input = scaled;
-    }
+    const input = this.#preScaled(x);
     const out = ops.ropeScaled(
-      input, this.dims, this.traditional, this.base, this.scale, offset, this.freqs,
+      input ?? x, this.dims, this.traditional, this.base, this.scale, offset, this.freqs,
     );
-    scaled?.dispose();
+    input?.dispose();
     return out;
   }
 
+  /** Per-row array-offset variant (batched decode under left-padding: each
+   *  row's REAL position rides an int32 [B] array — the scalar offset would
+   *  mis-position every padded row by its pad amount). Same math as
+   *  `apply`; only the offset spelling differs. */
+  applyDynamic(x: MlxArray, offsetArr: MlxArray): MlxArray {
+    const input = this.#preScaled(x);
+    const out = ops.ropeScaledDynamic(
+      input ?? x, this.dims, this.traditional, this.base, this.scale, offsetArr, this.freqs,
+    );
+    input?.dispose();
+    return out;
+  }
+
+  /** x[..., :dims] = preScale * x[..., :dims] (weak-scalar mul); null when
+   *  no pre-scale applies (caller uses x unchanged). */
+  #preScaled(x: MlxArray): MlxArray | null {
+    if (!(this.preScaleAlways || this.preScale !== 1.0)) return null;
+    const D = x.shape[x.shape.length - 1]!;
+    if (this.dims === D) return ops.mulScalar(x, this.preScale);
+    const start = x.shape.map(() => 0);
+    const stop = [...x.shape];
+    stop[stop.length - 1] = this.dims;
+    const head = x.slice(start, stop);
+    const headScaled = ops.mulScalar(head, this.preScale);
+    head.dispose();
+    const scaled = ops.sliceUpdate(x, headScaled, start, stop);
+    headScaled.dispose();
+    return scaled;
+  }
+
   dispose(): void {
     this.freqs?.dispose();
   }
@@ -89,6 +101,9 @@ export class NoRope extends UniversalRope {
     // Return an owned copy-view so callers can dispose uniformly.
     return x.slice(x.shape.map(() => 0), [...x.shape]);
   }
+  override applyDynamic(x: MlxArray, _offsetArr: MlxArray): MlxArray {
+    return x.slice(x.shape.map(() => 0), [...x.shape]);
+  }
 }
 
 /** `base ** (arange(0, dims, 2) / dims)` with the oracle's dtypes: