@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/sampling/normalized-argmax.ts

@@ -0,0 +1,116 @@
+// Greedy selection preserves dtype-rounded normalized scores, including ties.
+// Native logsumexp retains the reference reduction. Tiled selection avoids
+// writing the vocabulary-sized normalized array. Inputs remain borrowed.
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { CompiledFunction, isShapelessTracing } from "@mlx-bun/mlx/compile";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+const header = `
+inline bool better(float x, uint xi, float y, uint yi) {
+  return x > y || (x == y && xi < yi);
+}`;
+
+// Four simdgroups reduce score/index pairs. Lowest token ID wins every tie.
+// The caller finishes the simd==0 block by writing its output from lane zero.
+const reduction = `
+  for (uint distance = 16; distance > 0; distance >>= 1) {
+    float other = simd_shuffle_down(best, distance);
+    uint oi = simd_shuffle_down(bestId, distance);
+    if (lane + distance < 32 && better(other, oi, best, bestId)) {
+      best = other; bestId = oi;
+    }
+  }
+  threadgroup float scores[4];
+  threadgroup uint ids[4];
+  if (lane == 0) { scores[simd] = best; ids[simd] = bestId; }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd == 0) {
+    best = lane < 4 ? scores[lane] : -INFINITY;
+    bestId = lane < 4 ? ids[lane] : 0xffffffffu;
+    for (uint distance = 2; distance > 0; distance >>= 1) {
+      float other = simd_shuffle_down(best, distance);
+      uint oi = simd_shuffle_down(bestId, distance);
+      if (lane + distance < 4 && better(other, oi, best, bestId)) {
+        best = other; bestId = oi;
+      }
+    }
+`;
+
+const partialSource = `
+  const uint tid = thread_position_in_threadgroup.x;
+  const uint lane = tid % 32, simd = tid / 32;
+  const uint row = threadgroup_position_in_grid.y;
+  const uint tile = threadgroup_position_in_grid.x;
+  float best = -INFINITY;
+  uint bestId = 0xffffffffu;
+  for (uint i = tile * 1024 + tid; i < min(uint(V), (tile + 1) * 1024); i += 128) {
+    float value = float(T(logits[row * V + i] - lse[row]));
+    // A nonfinite logsumexp produces only NaN/-inf normalized scores.
+    // Native argmax selects token zero for this degenerate row.
+    value = isnan(value) ? -INFINITY : value;
+    if (better(value, i, best, bestId)) { best = value; bestId = i; }
+  }
+  ${reduction}
+    if (lane == 0) { values[row * P + tile] = best; indices[row * P + tile] = bestId; }
+  }
+`;
+
+const finishSource = `
+  const uint tid = thread_position_in_threadgroup.x;
+  const uint lane = tid % 32, simd = tid / 32;
+  const uint row = threadgroup_position_in_grid.y;
+  float best = -INFINITY;
+  uint bestId = 0xffffffffu;
+  for (uint i = tid; i < P; i += 128) {
+    float value = values[row * P + i];
+    uint id = indices[row * P + i];
+    if (better(value, id, best, bestId)) { best = value; bestId = id; }
+  }
+  ${reduction}
+    if (lane == 0) out[row] = bestId;
+  }
+`;
+
+let partial: MetalKernel | undefined;
+let finish: MetalKernel | undefined;
+let compiled: CompiledFunction | undefined;
+
+function kernelGraph(logits: MlxArray): MlxArray {
+  const shape = logits.shape;
+  const vocab = shape.at(-1)!, rows = logits.size / vocab, tiles = Math.ceil(vocab / 1024);
+  partial ??= new MetalKernel({ name: "normalized_argmax_partial", header,
+    inputNames: ["logits", "lse"], outputNames: ["values", "indices"], source: partialSource });
+  finish ??= new MetalKernel({ name: "normalized_argmax_finish", header,
+    inputNames: ["values", "indices"], outputNames: ["out"], source: finishSource });
+  using lse = ops.logsumexpAxis(logits, -1, true);
+  const [values, indices] = partial.apply([logits, lse], {
+    outputs: [{ shape: [rows, tiles], dtype: Dtype.float32 }, { shape: [rows, tiles], dtype: Dtype.uint32 }],
+    grid: [tiles * 128, rows, 1], threadGroup: [128, 1, 1],
+    templateInts: { V: vocab, P: tiles }, templateDtypes: { T: logits.dtype },
+  });
+  try {
+    using tokens = finish.apply([values!, indices!], {
+      outputs: [{ shape: [rows], dtype: Dtype.uint32 }],
+      grid: [128, rows, 1], threadGroup: [128, 1, 1], templateInts: { P: tiles },
+    })[0]!;
+    return ops.reshape(tokens, shape.slice(0, -1));
+  } finally { values!.dispose(); indices!.dispose(); }
+}
+
+/** Return owned greedy IDs with the same shape/rounding as argmax(logprobs).
+ * Fixed-shape compilation caches graph/FFI work for repeated row geometries.
+ * Shapeless enclosing graphs retain native operations, which support symbolic
+ * sizes; integer/empty inputs also retain the original operation semantics. */
+export function normalizedArgmax(logits: MlxArray): MlxArray {
+  const dtype = logits.dtype;
+  if (isShapelessTracing() || logits.size === 0 ||
+      (dtype !== Dtype.float32 && dtype !== Dtype.float16 && dtype !== Dtype.bfloat16)) {
+    using lse = ops.logsumexpAxis(logits, -1, true);
+    using scores = ops.sub(logits, lse);
+    return ops.argmaxAxis(scores, -1);
+  }
+  compiled ??= new CompiledFunction(([input]) => [kernelGraph(input!)], false);
+  return compiled.apply([logits])[0]!;
+}

package/src/kernels/sampling/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 "@mlx-bun/mlx/array";
+import { CompiledFunction } from "@mlx-bun/mlx/compile";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { MetalKernel } from "@mlx-bun/mlx/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/kernels/softcap.ts

@@ -0,0 +1,35 @@
+// Compiled (fused) Gemma logit softcap, matching mlx-lm's gemma4_text.py
+// `@partial(mx.compile) logit_softcap`. The unfused composition lives in
+// logits.ts; the two are separate on purpose (see compiledLogitSoftcap).
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { CompiledFunction } from "@mlx-bun/mlx/compile";
+
+/** gemma4_text.py `@partial(mx.compile) def logit_softcap(softcap, x):
+ *    return mx.tanh(x / softcap) * softcap`
+ *  The oracle fuses divide+tanh+multiply into ONE kernel; our default
+ *  `logitSoftcap` (logits.ts) runs them UNFUSED (standalone div + tanh + mul →
+ *  a separate `vn_Tanh` + `vsn_Divide`). Keyed by cap so the softcap is baked
+ *  as a graph constant (matching the oracle's weak-scalar → the
+ *  `…Divide…Tanh…Multiply…_VC_` kernel). The two implementations are
+ *  intentionally separate: this one reproduces the oracle's compiled kernel
+ *  (Gemma4's final softcap, outside a compiled trace); `logitSoftcap` is the
+ *  plain composition used elsewhere (Universal attention/final softcap, and
+ *  inside a compiled trace). Neither may replace the other. */
+const _softcapClosures = new Map<number, CompiledFunction>();
+export function compiledLogitSoftcap(x: MlxArray, cap: number): MlxArray {
+  let fn = _softcapClosures.get(cap);
+  if (!fn) {
+    fn = new CompiledFunction((inputs) => {
+      const xx = inputs[0]!;                    // tanh(x / softcap) * softcap
+      const capArr = ops.scalarLike(cap, xx);
+      const scaled = ops.div(xx, capArr);
+      const t = ops.tanh(scaled); scaled.dispose();
+      const out = ops.mul(t, capArr); t.dispose(); capArr.dispose();
+      return [out];
+    });
+    _softcapClosures.set(cap, fn);
+  }
+  return fn.apply([x])[0]!;
+}

package/src/kernels/trellis/codebook.ts

@@ -0,0 +1,123 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+
+/** The only definition of the 1MAD decode arithmetic; every kernel header includes it. */
+export const HEADER = String.raw`
+static inline float trellis_bf16(float v) {
+  uint u = as_type<uint>(v);
+  u = (u + 0x7FFFu + ((u >> 16) & 1u)) & 0xFFFF0000u;
+  return as_type<float>(u);
+}
+// QTIP 1MAD computed code for state s: x = s*34038481 + 76625530 (mod 2^32),
+// y = sum of the four bytes of x - 510, value = y / 147.800537109375. Two
+// byte-pair adds replace the four extracts; precise divide matches the host
+// LUT (lut1mad below) bit for bit.
+static inline int trellis_y(uint s) {
+  const uint x = s * 34038481u + 76625530u;
+  const uint p = (x & 0x00FF00FFu) + ((x >> 8) & 0x00FF00FFu);
+  return (int)((p & 0xFFFFu) + (p >> 16)) - 510;
+}
+static inline float trellis_val(uint s) {
+  return metal::precise::divide((float)trellis_y(s), 147.800537109375f);
+}
+// y × (1/d) refined by one residual step: q' = fma(fma(-q, d, y), 1/d, q).
+// Correctly rounded for every y in [-510, 510] (checked host-side against
+// the LUT), 3 FMA-class ops instead of a precise divide. Every expansion and
+// prefill kernel that reproduces the host LUT decodes through this.
+static inline float trellis_refined_y(int y) {
+  const float yf = (float)y;
+  const float r = 1.0f / 147.800537109375f;
+  const float q = yf * r;
+  const float e = metal::fma(-q, 147.800537109375f, yf);
+  return metal::fma(e, r, q);
+}
+static inline float trellis_val_rcp(uint s) { return trellis_refined_y(trellis_y(s)); }
+// Unrefined y × (1/d) (≤1 ulp f32): the packed matvec/scatter weight is this
+// f32 code value, a different numerical contract from the LUT-exact decode above.
+static inline float trellis_unrefined_y(int y) {
+  return (float)y * (1.0f / 147.800537109375f);
+}
+// The code value y as f32 without an integer-to-float conversion: y + 510
+// (0..1020) ORed into the mantissa of 2^23, whose ULP is one, minus 2^23 + 510.
+// Equal to (float)trellis_y(s). Factored kernels accumulate y·x and apply
+// scale·TRELLIS_RCP once per output (or fold it into each input row).
+static inline float trellis_y_f32(uint s) {
+  return as_type<float>(0x4b000000u | uint(trellis_y(s) + 510)) - 8389118.0f;
+}
+#define TRELLIS_RCP (1.0f / 147.800537109375f)
+// For y in [-510, 510], 2^23 has unit f32 spacing, so setting its low
+// mantissa bits constructs 2^23+y+510 exactly. Subtracting 2^23+510
+// recovers the same f32 y as the integer conversion; the reciprocal is unchanged.
+static inline float trellis_unrefined_bits_y(int y) {
+  const float lifted = as_type<float>(0x4b000000u | uint(y + 510));
+  return (lifted - 8389118.0f) * (1.0f / 147.800537109375f);
+}
+// VARIANT: 0 = inline 1MAD + precise divide; 1 = inline 1MAD × reciprocal
+// (1 ulp risk vs the host LUT); 2 = 4096-entry f32 LUT in threadgroup memory;
+// 3 = the same LUT gathered from device memory.
+// 4 = NO decode (timing floor only, wrong numerics); 5 = rcp decode, no bf16
+// rounding; 6 = unrefined y×(1/d) (≤1 ulp f32), no bf16 rounding — the
+// weight is f32 code×scale, more accurate than the artifact's stored bf16.
+#define TRELLIS_DECODE(win, lutTG, lut) \
+  ((VARIANT) == 0 ? trellis_val(win) : \
+   (VARIANT) == 1 || (VARIANT) == 5 ? trellis_val_rcp(win) : \
+   (VARIANT) == 6 ? trellis_unrefined_y(trellis_y(win)) : \
+   (VARIANT) == 2 ? lutTG[win] : \
+   (VARIANT) == 4 ? (float)(win) : lut[win])
+#define TRELLIS_ROUND(v) ((VARIANT) >= 4 ? (v) : trellis_bf16(v))
+// state_t of one packed block: L bits at offset (BT-1-t)*K, wrapping (32-bit ops).
+static inline uint trellis_state(const device uint32_t* blk, uint wpb, uint t, uint bt, uint k, uint l) {
+  const uint p = (bt - 1u - t) * k;
+  const uint wi = p >> 5;
+  const uint off = p & 31u;
+  uint win = blk[wi] >> off;
+  if (off + l > 32u) win |= blk[(wi + 1u == wpb) ? 0u : (wi + 1u)] << (32u - off);
+  return win & ((1u << l) - 1u);
+}
+`;
+
+/** The one-row threadgroup-codebook kernels' table read (gate-up-codebook, down-codebook): the code value of the
+ *  L-bit window at bit `off` of a lane's words w, read as the L + 2 bits [off - 2, off + L), the window's byte offset
+ *  in a float table plus two low bits that the M1 Max (applegpu_g13s) threadgroup load ignores for a 4-byte element,
+ *  so one bit-field extract serves each window; the kernels' bit-identity tests hold them to the reference bits. A
+ *  field that crosses a word, or starts below bit 0, reads the virtual words v (v[i]: the high half of w[i - 1] and
+ *  the low half of w[i], bits [32i - 16, 32i + 16)). `off` and `l` are compile-time constants at every call. */
+export const CODEBOOK_WINDOW = String.raw`
+static inline float trellis_codebook_window(const threadgroup uchar* table, thread const uint* w, thread const uint* v,
+                                            uint off, uint l) {
+  const int f = int(off) - 2;
+  const uint mask = (1u << (l + 2u)) - 1u;
+  uint bits;
+  if (f >= 0 && uint(f & 31) + l + 2u <= 32u) bits = (w[f >> 5] >> uint(f & 31)) & mask;
+  else { const uint vi = uint(f + 16) >> 5; bits = (v[vi] >> uint(f + 16 - 32 * int(vi))) & mask; }
+  return *(const threadgroup float*)(table + bits);
+}
+`;
+
+export function lut1mad(L: number): Float32Array {
+  const n = 1 << L;
+  const out = new Float32Array(n);
+  const M = (1n << 32n) - 1n;
+  for (let i = 0; i < n; i++) {
+    let x = BigInt(i) & M;
+    x = (x * 34038481n + 76625530n) & M;
+    const y =
+      Number((x & 255n) + ((x >> 8n) & 255n) + ((x >> 16n) & 255n) + ((x >> 24n) & 255n)) - 510;
+    out[i] = y / 147.800537109375;
+  }
+  return out;
+}
+
+export function wordsPerBlock(T: number, k: number): number {
+  if ((T * k) % 32 !== 0) throw new Error(`T·k = ${T * k} not a multiple of 32`);
+  return (T * k) / 32;
+}
+
+const luts = new Map<number, MlxArray>();
+export function lutFor(L: number): MlxArray {
+  let a = luts.get(L);
+  if (!a) { a = MlxArray.fromFloat32(lut1mad(L), [1 << L]); a.eval(); luts.set(L, a); }
+  return a;
+}
+
+// Variants 7..13 retain variant 6 decoded values.
+export function decoderVariant(selected: number): number { return selected >= 7 && selected <= 13 ? 6 : selected; }

package/src/kernels/trellis/down-codebook.ts

@@ -0,0 +1,184 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import type { TrellisGeometry } from "./geometry";
+import { CODEBOOK_WINDOW, HEADER } from "./codebook";
+import { TRELLIS_CODEBOOK_DOWN_THREADS } from "./launch";
+
+// Axis-0 (down) projection of one activation row over T 256 Trellis codes,
+// decoding through a threadgroup codebook (every code value y·(1/147.8) of the
+// 2^L states, generated once per threadgroup). Each lane owns P = 16
+// consecutive positions of one block: it loads the words they span (plus the
+// next lane's first word for the last windows) for U input rows ahead, aligns
+// them to its first bit once per row where P·K is not a multiple of 32 (the
+// 3-bit codes, whose odd lanes start at bit 16), then reads every window with
+// one bit-field extract at a compile-time shift (CODEBOOK_WINDOW, codebook.ts).
+// No lane exchanges data, so lanes past the last block return at once.
+// Split-K over input rows into SPLITS f32 partials, summed as MLX sums them
+// (a threadgroup adds the four splits of the first stage, FOLD the rest) and
+// cast to the activation dtype. Every position accumulates the same f32 code·scale weights over the same
+// rows in the same order as trellisScatter's single-row kernels at variant 13
+// (the balanced one for the 3-bit block interleave), so the output is theirs
+// bit for bit.
+const SPLITS = 128, P = 16, U = 2;
+
+const SOURCE = String.raw`
+  threadgroup float codebook[4096];  // 2^L codes, then the eight SIMD groups' partials
+  for (uint i = thread_position_in_threadgroup.x; i < (1u << (uint)L); i += (uint)THREADS)
+    codebook[i] = trellis_unrefined_bits_y(trellis_y(i));
+  threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  constexpr uint BITS = (uint)(P * K);               // a lane's bits
+  constexpr uint SHIFT = BITS % 32u;                 // 0, or 16: odd lanes start at bit 16 of a word
+  constexpr bool ALIGNED = SHIFT == 0u;
+  constexpr uint NA = (BITS + (uint)L - (uint)K + 31u) / 32u;  // words its windows span from its first bit
+  constexpr uint NL = ALIGNED ? NA : NA + 1u;        // words it loads
+  constexpr uint LB = 256u / (uint)P;                // lanes per block
+  constexpr uint WPB = 8u * (uint)K;                 // words per block
+  constexpr uint NB = (uint)(C / 256);               // blocks
+  constexpr uint SGS = (NB + 32u / LB - 1u) / (32u / LB); // SIMD groups per split
+  constexpr ulong ROWW = INTERLEAVE ? 2ul * WPB : (ulong)NB * WPB;
+  constexpr uint ROWS_PER = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
+  constexpr uint SG_COLS = 32u * (uint)P;            // a SIMD group's columns
+  // Threadgroup t: row group y = t % 32 and SIMD groups (k, c) = (s % 4, s / 4) take split y + 32k of column group
+  // 2·(t / 32) + c, so the threadgroup holds the four splits MLX's column sum adds first for 2·SG_COLS columns.
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint ty = threadgroup_position_in_grid.y;
+  const uint y = ty % 32u;
+  const uint split = y + 32u * (sg % 4u);
+  const uint colGroup = 2u * (ty / 32u) + sg / 4u;
+  const uint li = lane % LB;
+  const uint blk = colGroup * (32u / LB) + lane / LB;
+  const bool active = colGroup < SGS && blk < NB;
+  const device uint32_t* blockWords = codes + (!active ? 0ul : INTERLEAVE
+    ? (ulong)(blk / 2u) * (ulong)R * 2ul * WPB + (blk % 2u) * WPB : (ulong)blk * WPB);
+  const uint first = (li * BITS) >> 5;
+  const bool shifted = !ALIGNED && (li & 1u);
+  uint idx[NL];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < NL; ++i) idx[i] = (first + i) % WPB;
+  const uint r0 = split * ROWS_PER;
+  const uint rows = !active ? 0u : (uint)R % (uint)SPLITS == 0u ? ROWS_PER : r0 >= (uint)R ? 0u : metal::min(ROWS_PER, (uint)R - r0);
+  const threadgroup uchar* table = (const threadgroup uchar*)codebook;
+  float acc[P];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < (uint)P; ++i) acc[i] = 0.0f;
+  uint r = 0;
+  for (; r + (uint)U <= rows; r += (uint)U) {
+    uint ld[U][NL];
+    float xr[U], sr[U];
+    #pragma clang loop unroll(full)
+    for (uint u = 0; u < (uint)U; ++u) {
+      const device uint32_t* rowWords = blockWords + (ulong)(r0 + r + u) * ROWW;
+      #pragma clang loop unroll(full)
+      for (uint i = 0; i < NL; ++i) ld[u][i] = rowWords[idx[i]];
+      xr[u] = float(x[r0 + r + u]);
+      sr[u] = float(scales[r0 + r + u]);
+    }
+    #pragma clang loop unroll(full)
+    for (uint u = 0; u < (uint)U; ++u) {
+      uint w[NA], v[NA];
+      #pragma clang loop unroll(full)
+      for (uint i = 0; i < NA; ++i) w[i] = shifted ? (ld[u][i] >> SHIFT) | (ld[u][i + 1] << (32u - SHIFT)) : ld[u][i];
+      v[0] = w[0] << 16;
+      #pragma clang loop unroll(full)
+      for (uint i = 1; i < NA; ++i) v[i] = (w[i - 1] >> 16) | (w[i] << 16);
+      #pragma clang loop unroll(full)
+      for (uint i = 0; i < (uint)P; ++i) {
+        const float code = trellis_codebook_window(table, w, v, i * (uint)K, (uint)L);
+        acc[i] = metal::fma(code * sr[u], xr[u], acc[i]);
+      }
+    }
+  }
+  for (; r < rows; ++r) {
+    const device uint32_t* rowWords = blockWords + (ulong)(r0 + r) * ROWW;
+    uint ld[NL];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < NL; ++i) ld[i] = rowWords[idx[i]];
+    const float xr = float(x[r0 + r]);
+    const float sr = float(scales[r0 + r]);
+    uint w[NA], v[NA];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < NA; ++i) w[i] = shifted ? (ld[i] >> SHIFT) | (ld[i + 1] << (32u - SHIFT)) : ld[i];
+    v[0] = w[0] << 16;
+    #pragma clang loop unroll(full)
+    for (uint i = 1; i < NA; ++i) v[i] = (w[i - 1] >> 16) | (w[i] << 16);
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)P; ++i) {
+      const float code = trellis_codebook_window(table, w, v, i * (uint)K, (uint)L);
+      acc[i] = metal::fma(code * sr, xr, acc[i]);
+    }
+  }
+  // The four splits' partials meet in the codebook's memory; each column's first sum stage, ((0 + p[y]) + p[y + 32])
+  // + ..., in MLX's order, goes out as row y of the 32 the fold sums.
+  threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  threadgroup float* sums = codebook;
+  const uint local = (lane / LB) * 256u + 255u - li * (uint)P;
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < (uint)P; ++i) sums[sg * SG_COLS + local - i] = acc[i];
+  threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  for (uint c = thread_position_in_threadgroup.x; c < 2u * SG_COLS; c += (uint)THREADS) {
+    const uint group2 = c / SG_COLS, col = (2u * (ty / 32u) + group2) * SG_COLS + c % SG_COLS;
+    if (col >= (uint)C) continue;
+    float t = 0.0f;
+    #pragma clang loop unroll(full)
+    for (uint k = 0; k < 4u; ++k) t = sums[(group2 * 4u + k) * SG_COLS + c % SG_COLS] + t;
+    partial[(ulong)y * (ulong)C + col] = t;
+  }
+`;
+
+// The second stage of MLX's column sum over the SPLITS partials (col_reduce_looped, 32x32 tiles, the reduction it runs
+// at this shape) and the cast to T, bit for bit: SIMD group s of a threadgroup's 32 columns sums each of its four
+// columns over the 32 first-stage rows (lane y holds row y) with simd_sum.
+const FOLD = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint col = threadgroup_position_in_grid.x * 32u + simdgroup_index_in_threadgroup * 4u;
+  const float4 t = *(const device float4*)(partial + (ulong)lane * (ulong)C + col);
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < 4u; ++i) {
+    const float v = metal::simd_sum(t[i]);
+    if (lane == 0u) out[col + i] = T(v);
+  }
+`;
+
+let kernel: MetalKernel | undefined;
+let fold: MetalKernel | undefined;
+
+function launch(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry, interleave: boolean): MlxArray {
+  const groupsPerSplit = Math.ceil((g.cols / 256) / (P / 8));
+  if (x2.shape[0] !== 1) throw new Error("downCodebookRow: expects one activation row");
+  if (g.axis !== 0 || g.T !== 256 || g.L > 12 || g.cols % 256) throw new Error("downCodebookRow: needs axis-0 codes at T 256, L <= 12");
+  kernel ??= new MetalKernel({ name: "mlx_bun_trellis_down_codebook_row", inputNames: ["x", "codes", "scales"],
+    outputNames: ["partial"], source: SOURCE, header: HEADER + CODEBOOK_WINDOW, ensureRowContiguous: true });
+  fold ??= new MetalKernel({ name: "mlx_bun_trellis_down_codebook_fold", inputNames: ["partial"], outputNames: ["out"],
+    source: FOLD, ensureRowContiguous: true });
+  // 32 row groups, each threadgroup two SIMD groups' columns of four splits.
+  const [partial] = kernel.apply([x2, codes, scales], {
+    outputs: [{ shape: [1, 32, g.cols], dtype: Dtype.float32 }],
+    grid: [TRELLIS_CODEBOOK_DOWN_THREADS, 32 * Math.ceil(groupsPerSplit / 2), 1],
+    threadGroup: [TRELLIS_CODEBOOK_DOWN_THREADS, 1, 1],
+    templateDtypes: { T: x2.dtype },
+    templateInts: { R: g.rows, C: g.cols, K: g.k, L: g.L, P, U, SPLITS, THREADS: TRELLIS_CODEBOOK_DOWN_THREADS,
+      INTERLEAVE: Number(interleave) },
+  });
+  const [out] = fold.apply([partial!], {
+    outputs: [{ shape: [1, g.cols], dtype: x2.dtype }],
+    grid: [8 * 32 * g.cols / 32, 1, 1], threadGroup: [8 * 32, 1, 1],
+    templateDtypes: { T: x2.dtype }, templateInts: { C: g.cols },
+  });
+  partial!.dispose();
+  return out!;
+}
+
+/** x2 [1, inFeatures] · W → [1, outFeatures] in x2's dtype over row-major 2- or
+ *  4-bit codes [rows, cols·K/32]. Borrows inputs. */
+export function downCodebookRow(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  if (g.blockInterleave || codes.ndim !== 2 || (g.k !== 2 && g.k !== 4)) throw new Error("downCodebookRow: needs row-major 2- or 4-bit codes");
+  return launch(x2, codes, scales, g, false);
+}
+
+/** The same over 3-bit block-interleaved codes [cols/512, rows, 48]. */
+export function downK3InterleavedCodebookRow(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  if (g.blockInterleave !== 2 || g.k !== 3) throw new Error("downK3InterleavedCodebookRow: needs 3-bit block-interleaved codes");
+  return launch(x2, codes, scales, g, true);
+}