mlx-bun 0.0.9 → 0.0.11

package/src/mlx/ops.ts

@@ -43,10 +43,39 @@ export function sigmoid(a: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(outArray("sigmoid", (o) => C.mlx_sigmoid(o, a.handle, s)));
 }
 
-/** Scalar constant with the same dtype as `like` (python weak-scalar semantics). */
+/** f32 → bf16 truncation with round-to-nearest-even, returning the 2
+ *  little-endian bytes of the bf16 value. Matches mlx's
+ *  `static_cast<bfloat16_t>(float)` bit-for-bit so a host-built bf16 scalar is
+ *  numerically identical to `f32_scalar.astype(bf16)` — just without the GPU
+ *  Copy (astype) kernel. */
+function bf16Bytes(value: number): Uint8Array {
+  const fb = new Float32Array([value]);
+  const bits = new Uint32Array(fb.buffer)[0]!;
+  let hi16: number;
+  if (((bits >>> 23) & 0xff) === 0xff) {
+    // inf / nan: keep the high 16 bits, and keep a NaN a (quiet) NaN.
+    hi16 = (bits >>> 16) & 0xffff;
+    if ((bits & 0x7fffff) !== 0) hi16 |= 0x40;
+  } else {
+    const roundingBias = 0x7fff + ((bits >>> 16) & 1); // round to nearest even
+    hi16 = ((bits + roundingBias) >>> 16) & 0xffff;
+  }
+  return new Uint8Array([hi16 & 0xff, (hi16 >>> 8) & 0xff]);
+}
+
+/** Scalar constant with the same dtype as `like` (python weak-scalar semantics).
+ *  Built HOST-SIDE at the target dtype — mlx's own scalar→array path constructs
+ *  a bf16/f32 constant directly, so `scalar * bf16_array` is a single multiply.
+ *  Going through an f32 leaf + GPU `.astype(bf16)` (the old path) emits a
+ *  separate Copy kernel (`v_copy…`) that the reference never dispatches. */
 export function scalarLike(value: number, like: MlxArray): MlxArray {
+  if (like.dtype === Dtype.float32)
+    return MlxArray.fromFloat32(new Float32Array([value]), []);
+  if (like.dtype === Dtype.bfloat16)
+    return MlxArray.fromBytesCopy(bf16Bytes(value), [], Dtype.bfloat16);
+  // f16 / other float dtypes: keep the astype path (rare here; no bit-exact JS
+  // converter written for them yet).
   const f = MlxArray.fromFloat32(new Float32Array([value]), []);
-  if (like.dtype === Dtype.float32) return f;
   const cast = f.astype(like.dtype);
   f.dispose();
   return cast;
@@ -147,6 +176,25 @@ export function ropeDynamic(
   );
 }
 
+/** Array-offset fast::rope with every knob exposed (traditional + scale) —
+ *  the dynamic twin of `ropeScaled`, for the universal rope factory's
+ *  batched-decode path (per-row [B] offsets under left-padding). Same
+ *  kernel as `ropeDynamic`; static/dynamic bit-exactness is asserted in
+ *  tests/compile.test.ts. */
+export function ropeScaledDynamic(
+  x: MlxArray, dims: number, traditional: boolean, base: number | null,
+  scale: number, offset: MlxArray, freqs: MlxArray | null, s: S = gpuStream,
+): MlxArray {
+  return new MlxArray(
+    outArray("fast_rope_dynamic", (o) =>
+      C.mlx_fast_rope_dynamic(
+        o, x.handle, dims, traditional, optFloat(base), scale, offset.handle,
+        freqs?.handle ?? 0n, s,
+      ),
+    ),
+  );
+}
+
 export function sdpa(
   q: MlxArray, k: MlxArray, v: MlxArray, scale: number,
   maskMode: "" | "causal" | "array", maskArr: MlxArray | null = null,
@@ -172,6 +220,32 @@ export function quantizedMatmul(
   x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray | null,
   spec: QuantSpec, transpose = true, s: S = gpuStream,
 ): MlxArray {
+  // UPSTREAM MLX BUG WORKAROUND (found 2026-07-02, reproduced on stock mlx
+  // 0.31.2): quantized_matmul with transpose=false returns WRONG VALUES when
+  // the row count is 2 or 3 (relnorm ~0.3-1.1 vs the dequantized matmul);
+  // M=1 and M>=4 are correct, transpose=true is correct at every M. Pad the
+  // row axis to 4 with zeros and slice the result back. NOTE this wrapper
+  // cannot protect mlx's INTERNAL vjp of a transpose=true qmm (C++ calls the
+  // primitive directly) — training heads pad the span through the head
+  // instead (loss.ts logitsFromHiddenPadM).
+  const rowAxis = x.shape.length - 2;
+  const rows = x.shape[rowAxis]!;
+  if (!transpose && (rows === 2 || rows === 3)) {
+    const padShape = [...x.shape];
+    padShape[rowAxis] = 4 - rows;
+    const pad = zeros(padShape, x.dtype, s);
+    const xPad = concatAxis([x, pad], rowAxis, s);
+    pad.dispose();
+    const yPad = quantizedMatmul(xPad, w, scales, biases, spec, transpose, s);
+    xPad.dispose();
+    const start = MlxArray.fromInt32(new Int32Array([0]), [1]);
+    const outShape = [...yPad.shape];
+    outShape[rowAxis] = rows;
+    const y = sliceDynamic(yPad, start, [rowAxis], outShape, s);
+    start.dispose();
+    yPad.dispose();
+    return y;
+  }
   return new MlxArray(
     outArray("quantized_matmul", (o) =>
       C.mlx_quantized_matmul(
@@ -266,21 +340,35 @@ export function pow(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(outArray("power", (o) => C.mlx_power(o, a.handle, b.handle, s)));
 }
 
-/** nn.gelu_approx: 0.5 x (1 + tanh(√(2/π) (x + 0.044715 x³))) — composed
- *  exactly as mlx's python source: x**3 is mx.power (NOT x·x·x — they
- *  round differently in bf16), scalars promote weakly to x's dtype. */
+/** nn.gelu_approx: 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³))) — composed
+ *  exactly as mlx's python source, INCLUDING operand order: x**3 is mx.power
+ *  (NOT x·x·x — they round differently in bf16), and every scalar·array product
+ *  puts the SCALAR FIRST (`0.5 * x`, not `x * 0.5`). Multiply is commutative so
+ *  this is bit-identical, but it makes the compiled geglu graph's operand string
+ *  (`Multiply(C, V)`) match the oracle's fused kernel exactly. Scalars promote
+ *  weakly to x's dtype. */
 export function geluApprox(x: MlxArray, s: S = gpuStream): MlxArray {
-  const three = scalarLike(3, x);
-  const x3 = pow(x, three, s);
-  const cx3 = mulScalar(x3, 0.044715, s);
-  const inner = add(x, cx3, s);
-  const scaled = mulScalar(inner, Math.sqrt(2 / Math.PI), s);
+  // EXACT copy of mlx's graph for
+  //   0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x**3)))
+  // as dumped via mx.export_to_dot. mlx's `__rmul__`/`__radd__` canonicalize
+  // every `scalar OP array` to `array OP scalar`, so every commutative product/
+  // sum below is ARRAY-FIRST, constant-second (NOT scalar-first). Build order
+  // follows Python left-to-right (`0.5 * x` first). Same libmlx compiles this to
+  // the identical fused kernel.
+  const cHalf = scalarLike(0.5, x);
+  const halfx = mul(x, cHalf, s);                           // Multiply(x, 0.5)
+  const cThree = scalarLike(3, x);
+  const x3 = pow(x, cThree, s);                             // Power(x, 3)
+  const cCoef = scalarLike(0.044715, x);
+  const cx3 = mul(x3, cCoef, s);                            // Multiply(x**3, 0.044715)
+  const inner = add(x, cx3, s);                             // Add(x, cx3)
+  const cSqrt = scalarLike(Math.sqrt(2 / Math.PI), x);
+  const scaled = mul(inner, cSqrt, s);                      // Multiply(inner, √(2/π))
   const t = tanh(scaled, s);
-  const one = scalarLike(1, x);
-  const t1 = add(one, t, s);
-  const halfx = mulScalar(x, 0.5, s);
-  const out = mul(halfx, t1, s);
-  for (const a of [three, x3, cx3, inner, scaled, t, one, t1, halfx]) a.dispose();
+  const cOne = scalarLike(1, x);
+  const t1 = add(t, cOne, s);                               // Add(tanh, 1)
+  const out = mul(halfx, t1, s);                            // Multiply(halfx, 1+tanh)
+  for (const a of [cHalf, halfx, cThree, x3, cCoef, cx3, inner, cSqrt, scaled, t, cOne, t1]) a.dispose();
   return out;
 }
 
@@ -376,7 +464,7 @@ export function where(cond: MlxArray, x: MlxArray, y: MlxArray, s: S = gpuStream
 
 // arange is built host-side and uploaded. The original mlx_arange binding
 // broke once the calling path got JIT-optimized; a standalone repro
-// (repro/bun-ffi-f64/) traced the root cause to a Bun bug: after DFG
+// (lab/repro/bun-ffi-f64/) traced the root cause to a Bun bug: after DFG
 // tier-up, typed-array reads following a bun:ffi call return stale values
 // (the JIT eliminates the load across the native call). Not f64 marshaling
 // — args reach C intact. See PLAN.md Phase 4 findings. Large constant
@@ -392,13 +480,18 @@ export function arange(start: number, stop: number, step: number, dtype: Dtype,
   const cached = arangeCache.get(key);
   if (cached) return cached.slice([0], [n], s);
 
-  const data = new Int32Array(n);
-  for (let i = 0; i < n; i++) data[i] = start + i * step;
-  let arr = MlxArray.fromInt32(data, [n]);
-  if (dtype !== Dtype.int32) {
-    const cast = arr.astype(dtype, s);
-    arr.dispose();
-    arr = cast;
+  let arr: MlxArray;
+  if (dtype === Dtype.int32) {
+    // host-built int32 leaf (the Bun-f64-FFI-safe path — see note above).
+    const data = new Int32Array(n);
+    for (let i = 0; i < n; i++) data[i] = start + i * step;
+    arr = MlxArray.fromInt32(data, [n]);
+  } else {
+    // Native mlx Arange in the target dtype — matches mlx's `Arange` op exactly.
+    // The old `fromInt32 → .astype(dtype)` emitted a phantom int32→float Copy
+    // (an AsType the oracle never dispatches — e.g. the proportional-RoPE freq
+    // exponents). Integer start/step passed as doubles are exact.
+    arr = new MlxArray(outArray("arange", (o) => C.mlx_arange(o, start, stop, step, dtype, s)));
   }
   if (n >= ARANGE_CACHE_MIN) {
     arangeCache.set(key, arr);
@@ -785,6 +878,26 @@ export function conv1d(
   );
 }
 
+/** nn.Conv2d: input [B, H, W, C_in] (NHWC), weight [C_out, kH, kW, C_in/groups]
+ *  → [B, H', W', C_out]. mlx weight layout — the oracle transposes PyTorch's
+ *  [C_out, C_in, kH, kW] at load. For the gemma-4 audio SSCP subsampler.
+ *  dilation_1 and groups ride in one packed u64 — see the mlx_conv2d hazard
+ *  note in ffi.ts (bun:ffi stack-arg ABI bug). */
+export function conv2d(
+  input: MlxArray, weight: MlxArray, stride: [number, number] = [1, 1],
+  padding: [number, number] = [0, 0], dilation: [number, number] = [1, 1],
+  groups = 1, s: S = gpuStream,
+): MlxArray {
+  const dilation1AndGroups =
+    (BigInt(dilation[1] >>> 0) & 0xffffffffn) | (BigInt(groups >>> 0) << 32n);
+  return new MlxArray(
+    outArray("conv2d", (o) =>
+      C.mlx_conv2d(o, input.handle, weight.handle, stride[0], stride[1],
+        padding[0], padding[1], dilation[0], dilation1AndGroups, s),
+    ),
+  );
+}
+
 /** mx.split(a, indices, axis): split at the given boundary offsets along `axis`
  *  into N+1 contiguous slices (mirrors numpy/mlx split-by-indices). */
 export function split(a: MlxArray, indices: number[], axis: number): MlxArray[] {
@@ -818,3 +931,46 @@ export function silu(a: MlxArray, s: S = gpuStream): MlxArray {
   sig.dispose();
   return r;
 }
+
+/** fast::hadamard_transform over the last axis. `scale` defaults to
+ *  1/sqrt(dim) (mlx's own default when unset) — pass it explicitly to
+ *  match a reference that always states its normalization, per
+ *  turboquant-ops.ts's fwht(). */
+export function hadamardTransform(x: MlxArray, scale: number | null = null, s: S = gpuStream): MlxArray {
+  return new MlxArray(
+    outArray("hadamard_transform", (o) =>
+      C.mlx_hadamard_transform(o, x.handle, optFloat(scale), s),
+    ),
+  );
+}
+
+export function minAxis(a: MlxArray, axis: number, keepdims: boolean, s: S = gpuStream): MlxArray {
+  return new MlxArray(
+    outArray("min_axis", (o) => C.mlx_min_axis(o, a.handle, axis, keepdims, s)),
+  );
+}
+
+/** mx.round(x, decimals=0) — round-half-to-even at the ULP grid mlx uses. */
+export function round(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("round", (o) => C.mlx_round(o, a.handle, 0, s)));
+}
+
+export function greater(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("greater", (o) => C.mlx_greater(o, a.handle, b.handle, s)));
+}
+
+export function bitwiseAnd(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("bitwise_and", (o) => C.mlx_bitwise_and(o, a.handle, b.handle, s)));
+}
+
+export function bitwiseOr(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("bitwise_or", (o) => C.mlx_bitwise_or(o, a.handle, b.handle, s)));
+}
+
+export function leftShift(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("left_shift", (o) => C.mlx_left_shift(o, a.handle, b.handle, s)));
+}
+
+export function rightShift(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("right_shift", (o) => C.mlx_right_shift(o, a.handle, b.handle, s)));
+}