mlx-bun 0.2.0 → 0.4.0

package/src/mlx/ffi.ts

@@ -10,10 +10,10 @@
 // - mlx_optional_int is `{ int32 value; bool has_value }` = 8 bytes, by
 //   value in one register: value in bits 0..31, has_value at bit 32.
 
-import { dlopen, FFIType, ptr, read } from "bun:ffi";
+import { dlopen, FFIType, ptr, read, type Pointer } from "bun:ffi";
 import { existsSync } from "node:fs";
 import { dirname, join } from "node:path";
-import { nativePackDir } from "../native-pack";
+import { MLX_CORE_VERSION, nativePackDir } from "../native-pack";
 import { runtimeValue } from "../runtime-config";
 
 const { ptr: P, i32, u64, f32, f64, cstring } = FFIType;
@@ -21,7 +21,7 @@ const { ptr: P, i32, u64, f32, f64, cstring } = FFIType;
 /** libmlxc resolution, in order: explicit env override → next to the
  *  executable (the `bun build --compile` sidecar layout, where
  *  libmlxc.dylib + libmlx.dylib + mlx.metallib ship beside the binary —
- *  see docs/reference/embedding.md) → the downloaded native-pack cache
+ *  see docs/reference/distribution.md) → the downloaded native-pack cache
  *  (src/native-pack.ts; populated on first run by the CLI) → homebrew
  *  (arm64, then Intel prefix). Keep in sync with nativeRuntimeDir(). */
 function resolveLibmlxc(): string {
@@ -38,8 +38,12 @@ function resolveLibmlxc(): string {
 }
 
 export const LIBMLXC_PATH = resolveLibmlxc();
-
 export const C = dlopen(LIBMLXC_PATH, {
+  // linked MLX core version, independent of the C API family
+  mlx_version: { args: [P], returns: i32 },
+  mlx_string_new: { args: [], returns: u64 },
+  mlx_string_data: { args: [u64], returns: cstring },
+  mlx_string_free: { args: [u64], returns: i32 },
   // array lifecycle
   mlx_array_new: { args: [], returns: u64 },
   mlx_array_new_data: { args: [P, P, i32, i32], returns: u64 },
@@ -78,6 +82,7 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_device_info_get: { args: [P, u64], returns: i32 },
   mlx_device_info_free: { args: [u64], returns: i32 },
   mlx_device_info_get_size: { args: [P, u64, cstring], returns: i32 },
+  mlx_device_info_get_string: { args: [P, u64, cstring], returns: i32 },
   mlx_get_default_device: { args: [P], returns: i32 },
   mlx_device_free: { args: [u64], returns: i32 },
   mlx_synchronize: { args: [u64], returns: i32 },
@@ -98,6 +103,10 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_fast_rms_norm: { args: [P, u64, u64, f32, u64], returns: i32 },
   // (res, x, weight may-be-null, bias may-be-null, eps, stream)
   mlx_fast_layer_norm: { args: [P, u64, u64, u64, f32, u64], returns: i32 },
+  // MLX 0.32.2: mask, sinks, force_fused, stream.
+  mlx_fast_scaled_dot_product_attention: {
+    args: [P, u64, u64, u64, f32, cstring, u64, u64, FFIType.bool, u64], returns: i32,
+  },
   mlx_matmul: { args: [P, u64, u64, u64], returns: i32 },
   mlx_contiguous: { args: [P, u64, FFIType.bool, u64], returns: i32 },
   mlx_copy: { args: [P, u64, u64], returns: i32 },
@@ -108,15 +117,14 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_softmax_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
   mlx_max_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
   mlx_sum_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
+  // Optional dtype uses the eight-byte value/has_value layout of optional int.
+  mlx_cumsum_axis: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64, u64], returns: i32 },
+  mlx_cummax_axis: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64], returns: i32 },
   mlx_maximum: { args: [P, u64, u64, u64], returns: i32 },
   mlx_expand_dims: { args: [P, u64, i32, u64], returns: i32 },
   mlx_logical_or: { args: [P, u64, u64, u64], returns: i32 },
   // (res, x, dims, traditional, opt base, scale, offset, freqs may-be-null, stream)
   mlx_fast_rope: { args: [P, u64, i32, FFIType.bool, u64, f32, i32, u64, u64], returns: i32 },
-  // (res, q, k, v, scale, mask_mode, mask may-be-null, sinks may-be-null, stream)
-  mlx_fast_scaled_dot_product_attention: {
-    args: [P, u64, u64, u64, f32, cstring, u64, u64, u64], returns: i32,
-  },
   // ops (grown as phases need them)
   mlx_add: { args: [P, u64, u64, u64], returns: i32 },
   mlx_subtract: { args: [P, u64, u64, u64], returns: i32 },
@@ -150,8 +158,6 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_arange: { args: [P, f64, f64, f64, i32, u64], returns: i32 },
   mlx_take_along_axis: { args: [P, u64, u64, i32, u64], returns: i32 },
   mlx_put_along_axis: { args: [P, u64, u64, u64, i32, u64], returns: i32 },
-  mlx_cumsum: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64], returns: i32 },
-  mlx_cummax: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64], returns: i32 },
   mlx_where: { args: [P, u64, u64, u64, u64], returns: i32 },
   mlx_logsumexp_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
   mlx_greater_equal: { args: [P, u64, u64, u64], returns: i32 },
@@ -220,7 +226,7 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_save_safetensors: { args: [cstring, u64, u64], returns: i32 },
   mlx_map_string_to_array_insert: { args: [u64, cstring, u64], returns: i32 },
   mlx_map_string_to_string_insert: { args: [u64, cstring, cstring], returns: i32 },
-  // compile via closures (Phase A compiled decode — docs/design/optimization_plan.md).
+  // compile via closures (Phase A compiled decode — docs/archive/investigations/optimization_plan.md).
   // mlx_closure is the usual one-pointer struct; the func-payload variant
   // carries an id we use to route the trace callback back to JS.
   mlx_closure_new: { args: [], returns: u64 },
@@ -314,6 +320,22 @@ export function takeMlxError(): string | null {
   return e;
 }
 
+function readRuntimeVersion(): string {
+  const slot = new BigUint64Array([C.mlx_string_new()]);
+  const p = ptr(slot);
+  try {
+    if (C.mlx_version(p) !== 0)
+      throw new Error(`mlx_version failed: ${takeMlxError() ?? "(no mlx message)"}`);
+    return C.mlx_string_data(read.u64(p, 0)).toString();
+  } finally { C.mlx_string_free(read.u64(p, 0)); }
+}
+
+/** Read once at startup; kernel eligibility must follow the linked core. */
+export const MLX_VERSION = readRuntimeVersion();
+if (MLX_VERSION !== MLX_CORE_VERSION) {
+  throw new Error(`mlx-bun requires its bundled MLX ${MLX_CORE_VERSION}; ${LIBMLXC_PATH} loaded MLX ${MLX_VERSION}`);
+}
+
 /** mlx_dtype enum values (array.h). */
 export const enum Dtype {
   bool = 0, uint8 = 1, uint16 = 2, uint32 = 3, uint64 = 4,
@@ -405,6 +427,31 @@ export function clearCache(): void {
   C.mlx_clear_cache();
 }
 
+let cachedDeviceArchitecture: string | undefined;
+
+/** Native architecture identifier, used to match MLX's hardware dispatch. */
+export function deviceArchitecture(): string {
+  if (cachedDeviceArchitecture !== undefined) return cachedDeviceArchitecture;
+  const devSlot = new BigUint64Array(1), devPtr = ptr(devSlot);
+  if (C.mlx_get_default_device(devPtr) !== 0)
+    throw new Error(`mlx_get_default_device failed: ${takeMlxError() ?? ""}`);
+  const dev = read.u64(devPtr, 0);
+  const infoSlot = new BigUint64Array([C.mlx_device_info_new()]), infoPtr = ptr(infoSlot);
+  try {
+    if (C.mlx_device_info_get(infoPtr, dev) !== 0)
+      throw new Error(`mlx_device_info_get failed: ${takeMlxError() ?? ""}`);
+    const out = new BigUint64Array(1), outPtr = ptr(out);
+    const key = Buffer.from("architecture\0", "utf8");
+    if (C.mlx_device_info_get_string(outPtr, read.u64(infoPtr, 0), ptr(key)) !== 0)
+      throw new Error("architecture not in device info");
+    // Copy before the device-info object releases the native string.
+    return cachedDeviceArchitecture = new CString(Number(read.u64(outPtr, 0)) as Pointer).toString();
+  } finally {
+    C.mlx_device_info_free(read.u64(infoPtr, 0));
+    C.mlx_device_free(dev);
+  }
+}
+
 /** Metal's recommended max working-set size for the default device. */
 export function maxRecommendedWorkingSetSize(): number {
   const devSlot = new BigUint64Array(1);

package/src/mlx/materialize.ts

@@ -0,0 +1,23 @@
+// Explicit compact storage for a value that must release a larger backing
+// buffer. MLX 0.31.2 Copy::eval shares storage; contiguous() may also alias.
+// MLX's dynamic slice allocates the logical output size and has a native VJP.
+import { MlxArray } from "./array";
+import * as ops from "./ops";
+
+let zero: MlxArray | undefined;
+
+/** Borrow an array; return an owned, row-contiguous copy on the GPU.
+ * Values and dtype are unchanged. Evaluation stays lazy and never reads data
+ * back to the host. Use only when independent compact storage is required;
+ * immutable array aliases normally avoid this allocation and dispatch.
+ * Slice sizes are concrete, so use shape-specialized compiled graphs. */
+export function materializeCopy(input: MlxArray): MlxArray {
+  zero ??= MlxArray.fromInt32(new Int32Array([0]), []);
+  if (input.ndim > 0) return ops.sliceDynamic(input, zero, [0], input.shape);
+  const vector = ops.reshape(input, [1]);
+  try {
+    const copy = ops.sliceDynamic(vector, zero, [0], [1]);
+    try { return ops.reshape(copy, []); }
+    finally { copy.dispose(); }
+  } finally { vector.dispose(); }
+}

package/src/mlx/metal-kernel.ts

@@ -1,4 +1,4 @@
-// mx.fast.metal_kernel from Bun (docs/design/optimization_plan.md Phase E): the
+// mx.fast.metal_kernel from Bun (docs/archive/investigations/optimization_plan.md Phase E): the
 // intended MLX extension point for custom Metal kernels — no fork. The
 // kernel body is Metal Shading Language; mlx generates the signature
 // from input/output names (inputs by name, outputs by name, plus
@@ -29,15 +29,15 @@ export interface MetalKernelSpec {
 }
 
 export interface MetalKernelCall {
-  /** Concrete output shapes/dtypes. Provide this OR `outputShapeFn` (the latter
-   *  makes the kernel derive its outputs from the input shapes, like a real MLX
-   *  primitive — required to compose inside an mx.compile'd / shapeless closure,
-   *  where the wrapper must answer "given these input shapes, what comes out?"). */
+  /** Concrete output shapes/dtypes. Provide this OR `outputShapeFn`.
+   *  Both are captured when apply() builds the custom primitive. */
   outputs?: { shape: number[]; dtype: Dtype }[];
-  /** Derive output shapes/dtypes from the (current/traced) input shapes. */
+  /** Derive outputs during eager construction or a shape-specific compile trace.
+   *  This JS callback is not native shape inference: MLX's CustomKernel lacks
+   *  output_shapes(), so it cannot participate in shapeless replay. */
   outputShapeFn?: (inputs: MlxArray[]) => { shape: number[]; dtype: Dtype }[];
-  /** Concrete launch grid, OR derive it from inputs (so it tracks shape under
-   *  compile replay the same way `outputShapeFn` tracks output shapes). */
+  /** Concrete launch grid, OR derive it during eager construction or tracing.
+   *  Native replay reuses the captured grid without invoking this callback. */
   grid: [number, number, number] | ((inputs: MlxArray[]) => [number, number, number]);
   threadGroup: [number, number, number];
   templateInts?: Record<string, number>;
@@ -73,9 +73,8 @@ export class MetalKernel {
     // flash-CCE pin bug). Pure-additive: holds references, changes no logic.
     const keepAlive: unknown[] = [];
     try {
-      // Derive outputs + grid from the (current) input shapes when a function is
-      // given — so the SAME kernel object answers correctly for any input shape
-      // during a compile/replay trace, instead of baking in one call's shapes.
+      // Resolve callbacks once for this eager call or shape-specific trace.
+      // A compiled graph captures these values; replay does not call JS again.
       const outputs = call.outputShapeFn ? call.outputShapeFn(inputs) : call.outputs;
       if (!outputs) throw new Error("metal_kernel: provide outputs or outputShapeFn");
       const grid = typeof call.grid === "function" ? call.grid(inputs) : call.grid;

package/src/mlx/ops.ts

@@ -170,7 +170,7 @@ export function ropeScaled(
  *  of a baked int — required inside compiled decode graphs, where the
  *  offset changes every step but the graph must not. Same kernel as
  *  `rope`; bit-exactness vs the static form is asserted in
- *  tests/compile.test.ts. */
+ *  tests/unit/compile.test.ts. */
 export function ropeDynamic(
   x: MlxArray, dims: number, base: number | null, offset: MlxArray,
   freqs: MlxArray | null, s: S = gpuStream,
@@ -189,7 +189,7 @@ export function ropeDynamic(
  *  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. */
+ *  tests/unit/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,
@@ -213,7 +213,7 @@ export function sdpa(
     outArray("fast_sdpa", (o) =>
       C.mlx_fast_scaled_dot_product_attention(
         o, q.handle, k.handle, v.handle, scale, ptr(cstr(maskMode)),
-        maskArr?.handle ?? 0n, 0n, s,
+        maskArr?.handle ?? 0n, 0n, false, s,
       ),
     ),
   );
@@ -314,7 +314,8 @@ export function takeAxis(a: MlxArray, indices: MlxArray, axis: number, s: S = gp
 }
 
 export function reshape(a: MlxArray, shape: number[], s: S = gpuStream): MlxArray {
-  const buf = new Int32Array(shape);
+  // bun:ffi cannot take a pointer to an empty view; ndim=0 ignores the dummy.
+  const buf = new Int32Array(shape.length ? shape : [0]);
   return new MlxArray(
     outArray("reshape", (o) => C.mlx_reshape(o, a.handle, ptr(buf), BigInt(shape.length), s)),
   );
@@ -454,14 +455,14 @@ export function putAlongAxis(
 
 export function cumsum(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
   return new MlxArray(
-    outArray("cumsum", (o) => C.mlx_cumsum(o, a.handle, axis, false, true, s)),
+    outArray("cumsum", (o) => C.mlx_cumsum_axis(o, a.handle, axis, false, true, 0n, s)),
   );
 }
 
 /** Cumulative maximum along `axis` (inclusive, forward) — mlx.core.cummax. */
 export function cummax(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
   return new MlxArray(
-    outArray("cummax", (o) => C.mlx_cummax(o, a.handle, axis, false, true, s)),
+    outArray("cummax", (o) => C.mlx_cummax_axis(o, a.handle, axis, false, true, s)),
   );
 }
 
@@ -476,10 +477,13 @@ export function where(cond: MlxArray, x: MlxArray, y: MlxArray, s: S = gpuStream
 // (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
-// ranges (topP's vocab arange) are cached for the process lifetime.
-const arangeCache = new Map<string, MlxArray>();
+// — args reach C intact. See PLAN.md Phase 4 findings. Reuse large ranges
+// across attention layers and sampling, but bound their retained storage:
+// causal masks request a new range on every round once context exceeds 64K.
+const arangeCache = new Map<string, { array: MlxArray; bytes: number }>();
 const ARANGE_CACHE_MIN = 65536;
+const ARANGE_CACHE_BYTES = 8 * 1024 * 1024;
+let arangeCacheBytes = 0;
 
 export function arange(start: number, stop: number, step: number, dtype: Dtype, s: S = gpuStream): MlxArray {
   if (!Number.isInteger(start) || !Number.isInteger(step) || step === 0)
@@ -487,7 +491,11 @@ export function arange(start: number, stop: number, step: number, dtype: Dtype,
   const n = Math.max(0, Math.ceil((stop - start) / step));
   const key = `${start}|${step}|${n}|${dtype}`;
   const cached = arangeCache.get(key);
-  if (cached) return cached.slice([0], [n], s);
+  if (cached) {
+    arangeCache.delete(key);
+    arangeCache.set(key, cached);
+    return cached.array.slice([0], [n], s);
+  }
 
   let arr: MlxArray;
   if (dtype === Dtype.int32) {
@@ -503,8 +511,20 @@ export function arange(start: number, stop: number, step: number, dtype: Dtype,
     arr = new MlxArray(outArray("arange", (o) => C.mlx_arange(o, start, stop, step, dtype, s)));
   }
   if (n >= ARANGE_CACHE_MIN) {
-    arangeCache.set(key, arr);
-    return arr.slice([0], [n], s);
+    const bytes = arr.nbytes;
+    if (bytes <= ARANGE_CACHE_BYTES) {
+      while (arangeCacheBytes + bytes > ARANGE_CACHE_BYTES) {
+        const [oldKey, old] = arangeCache.entries().next().value!;
+        arangeCache.delete(oldKey);
+        arangeCacheBytes -= old.bytes;
+        // Returned slices retain their own native reference, including lazy
+        // views that have not been evaluated when this owner is evicted.
+        old.array.dispose();
+      }
+      arangeCache.set(key, { array: arr, bytes });
+      arangeCacheBytes += bytes;
+      return arr.slice([0], [n], s);
+    }
   }
   return arr;
 }
@@ -611,7 +631,7 @@ export function sliceDynamic(
 /** slice_update with the start index as an ARRAY (one int32 per entry in
  *  `axes`) — the compiled-decode form of the per-step cache write. Same
  *  write as `sliceUpdate`; bit-exactness vs the static form is asserted
- *  in tests/compile.test.ts. */
+ *  in tests/unit/compile.test.ts. */
 export function sliceUpdateDynamic(
   src: MlxArray, update: MlxArray, start: MlxArray, axes: number[], s: S = gpuStream,
 ): MlxArray {

package/src/mlx/quantized-rows.ts

@@ -0,0 +1,36 @@
+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

@@ -0,0 +1,52 @@
+// 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

@@ -0,0 +1,82 @@
+// 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

@@ -0,0 +1,122 @@
+// 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];
+}