mlx-bun 0.4.0 → 1.0.0

package/src/mlx/quantized-rows.ts

@@ -1,36 +0,0 @@
-import type { MlxArray } from "./array";
-import * as ops from "./ops";
-
-/** Run independent affine M=1 matrix-vector products in one native batched call.
- * Weight/scales/bias views share their original buffers. MLX 0.32.2's
- * flattened M>=2 path uses different arithmetic, so appending known tokens
- * must explicitly select this layout when preserving decode state.
- */
-export function quantizedMatmulRows(
-  x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray | null,
-  spec: ops.QuantSpec,
-): MlxArray {
-  const shape = x.shape, width = shape[shape.length - 1]!;
-  const rows = shape.reduce((a, b) => a * b, 1) / width;
-  if (rows === 1) return ops.quantizedMatmul(x, w, scales, biases, spec);
-  if (spec.mode !== "affine")
-    throw new Error("Independent-row batched matmul requires affine quantization");
-  if (w.shape.length !== 2 || scales.shape.length !== 2 || (biases && biases.shape.length !== 2))
-    throw new Error("Independent-row matmul requires shared two-dimensional weights");
-  const views: MlxArray[] = [];
-  const view = (value: MlxArray, newShape: number[]) => {
-    const result = ops.reshape(value, newShape); views.push(result); return result;
-  };
-  let output: MlxArray | undefined;
-  try {
-    const input = view(x, [rows, 1, width]);
-    const weight = view(w, [1, ...w.shape]);
-    const scale = view(scales, [1, ...scales.shape]);
-    const bias = biases ? view(biases, [1, ...biases.shape]) : null;
-    output = ops.quantizedMatmul(input, weight, scale, bias, spec);
-    return ops.reshape(output, [...shape.slice(0, -1), output.shape.at(-1)!]);
-  } finally {
-    output?.dispose();
-    for (const value of views) value.dispose();
-  }
-}

package/src/mlx/token-bitmask.ts

@@ -1,52 +0,0 @@
-// XGrammar's additive 0/-Infinity mask contract, decoded from packed bits.
-// Reference: mlc-ai/xgrammar b5048784, apply_token_bitmask_mlx.py (Apache-2.0).
-import { MlxArray } from "./array";
-import { CompiledFunction } from "./compile";
-import { Dtype } from "./ffi";
-import { MetalKernel } from "./metal-kernel";
-
-let compiled: CompiledFunction | undefined;
-let kernel: MetalKernel | undefined;
-
-function maskGraph(logits: MlxArray, bits: MlxArray): MlxArray {
-  kernel ??= new MetalKernel({
-    name: "mlx_bun_token_bitmask",
-    inputNames: ["logits", "bits"],
-    outputNames: ["out"],
-    source: String.raw`
-      const uint i = thread_position_in_grid.x;
-      if (i >= SIZE) return;
-      const uint id = i % V;
-      const uint word = uint(bits[id >> 5]);
-      const bool valid = ((word >> (id & 31)) & 1u) != 0;
-      out[i] = logits[i] + T(valid ? 0.0f : -INFINITY);
-    `,
-  });
-  return kernel.apply([logits, bits], {
-    outputs: [{ shape: logits.shape, dtype: logits.dtype }],
-    grid: [logits.size, 1, 1],
-    threadGroup: [128, 1, 1],
-    templateDtypes: { T: logits.dtype },
-    templateInts: { V: logits.shape.at(-1)!, SIZE: logits.size },
-  })[0]!;
-}
-
-/** Apply one little-endian token bitmask to every row along the final axis.
- * Borrow logits; return an owned lazy result. Copy the mutable host bits now,
- * so a matcher can advance before GPU evaluation. Missing words reject tokens.
- * Preserve the additive-mask expression, including NaNs and signed zero.
- * Shape-specialized compilation handles bf16/f16/f32 and strided inputs. */
-export function applyTokenBitmask(logits: MlxArray, bits: Int32Array): MlxArray {
-  const vocab = logits.shape.at(-1);
-  if (!vocab || logits.size === 0 || logits.size > 0x7fffffff)
-    throw new Error("token bitmask requires nonempty logits with at most 2^31-1 elements");
-  if (![Dtype.bfloat16, Dtype.float16, Dtype.float32].includes(logits.dtype))
-    throw new Error("token bitmask requires bf16, f16 or f32 logits");
-  const copy = new Int32Array(Math.ceil(vocab / 32));
-  copy.set(bits.subarray(0, copy.length));
-  const deviceBits = MlxArray.fromInt32(copy, [copy.length]);
-  try {
-    compiled ??= new CompiledFunction(([x, mask]) => [maskGraph(x!, mask!)], false);
-    return compiled.apply([logits, deviceBits])[0]!;
-  } finally { deviceBits.dispose(); }
-}

package/src/mlx/turboquant-kv-decode.ts

@@ -1,82 +0,0 @@
-// Exact byte-packed TurboQuant K/V decoding. The caller owns output arrays.
-// Rotation and cache policy belong to the existing codec/cache layers.
-import { type MlxArray, gpuStream } from './array';
-import { Dtype, type MlxHandle } from './ffi';
-import { isShapelessTracing } from './compile';
-import { MetalKernel } from './metal-kernel';
-
-export type PackedKvArrays = readonly [
-  keyIndices: MlxArray, keyScales: MlxArray, keyZeros: MlxArray,
-  valueIndices: MlxArray, valueScales: MlxArray,
-];
-
-const source = `
-  uint col = thread_position_in_grid.x, pos = thread_position_in_grid.y;
-  uint bh = thread_position_in_grid.z;
-  if (col >= D || pos >= vp_shape[2] || bh >= vp_shape[0] * vp_shape[1]) return;
-  uint b = bh / vp_shape[1], h = bh % vp_shape[1];
-  uint ko = b * kp_strides[0] + h * kp_strides[1] + pos * kp_strides[2];
-  uint vpo = b * vp_strides[0] + h * vp_strides[1] + pos * vp_strides[2];
-  uint kso = b * ks_strides[0] + h * ks_strides[1] + pos * ks_strides[2] + (col / KG) * ks_strides[3];
-  uint kzo = b * kz_strides[0] + h * kz_strides[1] + pos * kz_strides[2] + (col / KG) * kz_strides[3];
-  uint vso = b * vs_strides[0] + h * vs_strides[1] + pos * vs_strides[2] + (col / VG) * vs_strides[3];
-  int key_code;
-  if (KBITS == 8) key_code = int(kp[ko + col * kp_strides[3]]);
-  else {
-    uint bit = col * KBITS, byte = bit / 8u, shift = bit % 8u;
-    uint word = uint(kp[ko + byte * kp_strides[3]]);
-    if (shift + KBITS > 8u) word |= uint(kp[ko + (byte + 1u) * kp_strides[3]]) << 8u;
-    key_code = int((word >> shift) & ((1u << KBITS) - 1u));
-  }
-  // Keep the incumbent f32 addition before multiplication.
-  float shifted = float(key_code) + float(kz[kzo]);
-  uint i = (bh * vp_shape[2] + pos) * D + col;
-  outK[i] = OutK(shifted * float(ks[kso]));
-  uint bit = col * VBITS, byte = bit / 8u, shift = bit % 8u;
-  uint word = uint(vp[vpo + byte * vp_strides[3]]);
-  if (shift + VBITS > 8u) word |= uint(vp[vpo + (byte + 1u) * vp_strides[3]]) << 8u;
-  uint code = (word >> shift) & ((1u << VBITS) - 1u);
-  outV[i] = OutV(centroids[code] * float(vs[vso]));
-`;
-
-let kernel: MetalKernel | undefined;
-function jointKernel(): MetalKernel {
-  return kernel ??= new MetalKernel({
-    name: 'turboquant_kv_decode',
-    inputNames: ['kp', 'ks', 'kz', 'vp', 'vs', 'centroids'],
-    outputNames: ['outK', 'outV'], source, ensureRowContiguous: false,
-  });
-}
-
-/** Return bf16 K and rotated V (f32 for eager inverse rotation, bf16 for
- * deferred rotation). Null means the ordinary codec must handle this input.
- * Input arrays and the shared centroid table remain borrowed. */
-export function tryJointDecodePackedKv(
-  inputs: PackedKvArrays, centroids: MlxArray, kBits: number, vBits: number,
-  headDim: number, deferV: boolean, stream: MlxHandle = gpuStream,
-): [MlxArray, MlxArray] | null {
-  if (stream !== gpuStream || isShapelessTracing() ||
-      ![64, 128, 256, 512].includes(headDim) ||
-      ![2, 4, 5, 8].includes(kBits) || ![2, 3, 4, 5, 8].includes(vBits)) return null;
-  const [kp, ks, kz, vp, vs] = inputs;
-  if (kp.dtype !== (kBits === 8 ? Dtype.int8 : Dtype.uint8) ||
-      vp.dtype !== Dtype.uint8 || centroids.dtype !== Dtype.float32 ||
-      centroids.ndim !== 1 || centroids.size !== (1 << vBits)) return null;
-  const shapes = inputs.map(a => a.shape);
-  if (shapes.some(shape => shape.length !== 4)) return null;
-  const [b, h, n] = shapes[3]!;
-  if (!b || !h || !n || b * h * n * headDim >= 2 ** 32 ||
-      shapes.some(shape => shape[0] !== b || shape[1] !== h || shape[2] !== n) ||
-      shapes[0]![3] !== headDim * kBits / 8 || shapes[3]![3] !== headDim * vBits / 8 ||
-      shapes[1]![3] !== shapes[2]![3]) return null;
-  const keyGroup = headDim / shapes[1]![3]!, valueGroup = headDim / shapes[4]![3]!;
-  if (![32, 64].includes(keyGroup) || ![32, 64].includes(valueGroup) ||
-      [ks, kz, vs].some(a => ![Dtype.float16, Dtype.bfloat16, Dtype.float32].includes(a.dtype))) return null;
-  const shape = [b, h, n, headDim], valueDtype = deferV ? Dtype.bfloat16 : Dtype.float32;
-  return jointKernel().apply([...inputs, centroids], {
-    templateDtypes: { OutK: Dtype.bfloat16, OutV: valueDtype },
-    templateInts: { D: headDim, KBITS: kBits, VBITS: vBits, KG: keyGroup, VG: valueGroup },
-    grid: [headDim, n, b * h], threadGroup: [Math.min(headDim, 256), Math.max(1, 256 / headDim), 1],
-    outputs: [{ shape, dtype: Dtype.bfloat16 }, { shape, dtype: valueDtype }], stream,
-  }) as [MlxArray, MlxArray];
-}

package/src/mlx/turboquant-kv-inverse.ts

@@ -1,122 +0,0 @@
-// Joint packed K/V decode with eager inverse rotation in the measured D256 regime.
-// Hadamard butterfly/layout follows MLX 0.31.2,
-// mlx/backend/metal/kernels/hadamard.h, Copyright © 2024 Apple Inc., MIT.
-import { type MlxArray, gpuStream } from "./array";
-import { Dtype, type MlxHandle } from "./ffi";
-import { isShapelessTracing } from "./compile";
-import { MetalKernel } from "./metal-kernel";
-import { type PackedKvArrays } from "./turboquant-kv-decode";
-
-const header = `
-template <short R> inline void tq_inverse_radix(thread float* x) {
-  constexpr short logR = __builtin_ctz(R);
-  short h = 1;
-  #pragma unroll
-  for (short s = 0; s < logR; s++) {
-    #pragma unroll
-    for (short i = 0; i < R / 2; i++) {
-      short k = i & (h - 1), j = ((i - k) << 1) + k;
-      float a = x[j], b = x[j + h];
-      x[j] = a + b; x[j + h] = a - b;
-    }
-    h <<= 1;
-  }
-}
-`;
-const source = `
-  constexpr short D = 256, R = 16, NT = 16, LOG_R = 4;
-  uint pos = thread_position_in_grid.y, bh = thread_position_in_grid.z;
-  short i = thread_position_in_grid.x;
-  bool valid = pos < vp_shape[2];
-  uint b = bh / vp_shape[1], hidx = bh % vp_shape[1];
-  uint ko = b * kp_strides[0] + hidx * kp_strides[1] + pos * kp_strides[2];
-  uint vo = b * vp_strides[0] + hidx * vp_strides[1] + pos * vp_strides[2];
-  uint kso = b * ks_strides[0] + hidx * ks_strides[1] + pos * ks_strides[2];
-  uint kzo = b * kz_strides[0] + hidx * kz_strides[1] + pos * kz_strides[2];
-  uint vso = b * vs_strides[0] + hidx * vs_strides[1] + pos * vs_strides[2];
-  uint row = (bh * vp_shape[2] + pos) * D;
-  threadgroup float storage[D * 16];
-  threadgroup float* buf = storage + thread_position_in_threadgroup.y * D;
-  #pragma unroll
-  for (short j = 0; j < R / 4; j++) {
-    short index = j * 4 * NT + i * 4;
-    #pragma unroll
-    for (short r = 0; r < 4; r++) {
-      uint col = index + r;
-      if (!valid) { buf[col] = 0.0f; continue; }
-      int key_code = int(kp[ko + col * kp_strides[3]]);
-      float shifted = float(key_code) + float(kz[kzo + (col / 32) * kz_strides[3]]);
-      outK[row + col] = OutK(shifted * float(ks[kso + (col / 32) * ks_strides[3]]));
-      uint bit = col * 3u, byte = bit / 8u, shift = bit % 8u;
-      uint word = uint(vp[vo + byte * vp_strides[3]]);
-      if (shift + 3u > 8u) word |= uint(vp[vo + (byte + 1u) * vp_strides[3]]) << 8u;
-      uint code = (word >> shift) & 7u;
-      buf[col] = centroids[code * centroids_strides[0]] * float(vs[vso + (col / 32) * vs_strides[3]]);
-    }
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  float x[R]; short h = 1;
-  #pragma unroll
-  for (short s = 0; s < 2; s++) {
-    short k = i & (h - 1), j = ((i - k) << LOG_R) + k;
-    #pragma unroll
-    for (short r = 0; r < R; r++) x[r] = buf[j + h * r];
-    tq_inverse_radix<R>(x);
-    #pragma unroll
-    for (short r = 0; r < R; r++) buf[j + h * r] = x[r];
-    h <<= LOG_R;
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
-  #pragma unroll
-  for (short j = 0; j < R / 4; j++) {
-    short index = j * 4 * NT + i * 4;
-    #pragma unroll
-    for (short r = 0; r < 4; r++) {
-      uint col = index + r;
-      float normalized = buf[col] * 0.0625f;
-      if (valid) outV[row + col] = OutV(normalized * signs[col * signs_strides[0]]);
-    }
-  }
-`;
-
-let kernel: MetalKernel | undefined;
-function inverseKernel(): MetalKernel {
-  return kernel ??= new MetalKernel({
-    name: "turboquant_kv_inverse256",
-    inputNames: ["kp", "ks", "kz", "vp", "vs", "centroids", "signs"],
-    outputNames: ["outK", "outV"], header, source, ensureRowContiguous: false,
-  });
-}
-
-/** Borrow packed inputs and codec tables; return owned bf16 K/V or null.
- * Only the measured eager D256/k8v3/B1/H4/fp16/group32 regime is eligible.
- * Other regimes continue through the existing joint decoder or original codec.
- * N affects dispatch size but is never a shader template parameter. */
-export function tryDecodePackedKvInverse256(
-  inputs: PackedKvArrays, centroids: MlxArray, signs: MlxArray,
-  kBits: number, vBits: number, headDim: number, s: MlxHandle = gpuStream,
-): [MlxArray, MlxArray] | null {
-  if (s !== gpuStream || isShapelessTracing() || headDim !== 256 || kBits !== 8 || vBits !== 3)
-    return null;
-  const [kp, ks, kz, vp, vs] = inputs;
-  const shape = kp.shape;
-  if (shape.length !== 4 || shape[0] !== 1 || shape[1] !== 4 || shape[3] !== 256)
-    return null;
-  const n = shape[2]!;
-  if (n < 8192 || n * 1024 >= 2 ** 32) return null;
-  const samePrefix = (a: MlxArray, width: number) => {
-    const x = a.shape;
-    return x.length === 4 && x[0] === 1 && x[1] === 4 && x[2] === n && x[3] === width;
-  };
-  if (kp.dtype !== Dtype.int8 || vp.dtype !== Dtype.uint8 || !samePrefix(vp, 96)) return null;
-  if (![ks, kz, vs].every(a => a.dtype === Dtype.float16 && samePrefix(a, 8))) return null;
-  if (centroids.dtype !== Dtype.float32 || centroids.shape.length !== 1 || centroids.shape[0] !== 8)
-    return null;
-  if (signs.dtype !== Dtype.float32 || signs.shape.length !== 1 || signs.shape[0] !== 256)
-    return null;
-  return inverseKernel().apply([...inputs, centroids, signs], {
-    templateDtypes: { OutK: Dtype.bfloat16, OutV: Dtype.bfloat16 },
-    grid: [16, Math.ceil(n / 16) * 16, 4], threadGroup: [16, 16, 1],
-    outputs: [{ shape, dtype: Dtype.bfloat16 }, { shape, dtype: Dtype.bfloat16 }], stream: s,
-  }) as [MlxArray, MlxArray];
-}