@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/glm52/streamed-experts.ts

@@ -0,0 +1,360 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import { positiveInteger } from "../../runtime/integers";
+import { GLM52_EXPERT_SLOT_ALIGNMENT, validateRange, type Glm52CanonicalMetalLayout } from "./layout";
+export { GLM52_EXPERT_SLOT_ALIGNMENT } from "./layout";
+export type { Glm52CanonicalMetalLayout, Glm52CanonicalQ4MetalLayout, Glm52CanonicalQ8MetalLayout } from "./layout";
+
+const SIMD_WIDTH = 32;
+const SIMD_GROUPS = 4;
+const THREADS = SIMD_WIDTH * SIMD_GROUPS;
+const ROWS_PER_THREADGROUP = SIMD_GROUPS;
+
+/** Wrap a live, aligned ExpertIOSlabStore slot without copying it.
+ *
+ * The residency lease remains owned by the caller and must not be released
+ * until the returned expert output has been evaluated and the GPU stream has
+ * completed.
+ */
+export function glm52CanonicalQ4SlotView(
+  pointer: number,
+  layout: Glm52CanonicalMetalLayout,
+): MlxArray {
+  if (!Number.isSafeInteger(pointer) || pointer <= 0)
+    throw new Error("expert slot pointer must be a positive safe integer");
+  if (pointer % GLM52_EXPERT_SLOT_ALIGNMENT !== 0) {
+    throw new Error(
+      "expert slot pointer must be page-aligned for zero-copy Metal access",
+    );
+  }
+  return MlxArray.fromPointer(pointer, [layout.slotBytes], Dtype.uint8);
+}
+
+export const glm52CanonicalQ8SlotView = glm52CanonicalQ4SlotView;
+
+// One simdgroup owns one output row. The packed Q4 byte stream and F32 row
+// scales are read straight from the canonical residency slot. Every product is
+// explicitly dequantized to float and accumulated with an F32 FMA.
+const GATE_UP_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint row = thread_position_in_grid.y * (uint)ROWS_TG + sg;
+  const uint sample = thread_position_in_grid.z;
+  if (row >= (uint)I || sample >= (uint)M) return;
+  const device T* sampleX = x + (ulong)sample * (ulong)H;
+
+  const device uint32_t* gateRow =
+    (const device uint32_t*)(slot + GATE_W) + (ulong)row * (ulong)(H / 8);
+  const device uint32_t* upRow =
+    (const device uint32_t*)(slot + UP_W) + (ulong)row * (ulong)(H / 8);
+  const device float* gateScales = (const device float*)(slot + GATE_S);
+  const device float* upScales = (const device float*)(slot + UP_S);
+  const float gateScale = gateScales[row];
+  const float upScale = upScales[row];
+
+  float gateAcc = 0.0f;
+  float upAcc = 0.0f;
+  for (uint word = lane; word < (uint)(H / 8); word += 32u) {
+    const uint32_t gatePacked = gateRow[word];
+    const uint32_t upPacked = upRow[word];
+    const uint k0 = word * 8u;
+    for (uint place = 0; place < 8u; ++place) {
+      const float xv = float(sampleX[k0 + place]);
+      const float gateWeight =
+        float((gatePacked >> (place * 4u)) & 0xFu) - 8.0f;
+      const float upWeight =
+        float((upPacked >> (place * 4u)) & 0xFu) - 8.0f;
+      gateAcc = metal::fma(gateWeight, xv, gateAcc);
+      upAcc = metal::fma(upWeight, xv, upAcc);
+    }
+  }
+
+  const float gate = metal::simd_sum(gateAcc) * gateScale;
+  const float up = metal::simd_sum(upAcc) * upScale;
+  if (lane == 0u) {
+    // Match the production BF16 graph's materialization boundaries: gate and
+    // up are rounded to T, then SiLU and the product are rounded to T.
+    const T gateT = T(gate);
+    const T upT = T(up);
+    const T sigmoidT = T(1.0f / (1.0f + metal::precise::exp(-float(gateT))));
+    const T siluT = T(float(gateT) * float(sigmoidT));
+    mid[(ulong)sample * (ulong)I + row] =
+      T(float(siluT) * float(upT));
+  }
+`;
+
+const DOWN_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint row = thread_position_in_grid.y * (uint)ROWS_TG + sg;
+  const uint sample = thread_position_in_grid.z;
+  if (row >= (uint)H || sample >= (uint)M) return;
+  const device T* sampleMid = mid + (ulong)sample * (ulong)I;
+
+  const device uint32_t* downRow =
+    (const device uint32_t*)(slot + DOWN_W) + (ulong)row * (ulong)(I / 8);
+  const device float* downScales = (const device float*)(slot + DOWN_S);
+  const float scale = downScales[row];
+  float acc = 0.0f;
+  for (uint word = lane; word < (uint)(I / 8); word += 32u) {
+    const uint32_t packed = downRow[word];
+    const uint k0 = word * 8u;
+    for (uint place = 0; place < 8u; ++place) {
+      const float xv = float(sampleMid[k0 + place]);
+      const float weight =
+        float((packed >> (place * 4u)) & 0xFu) - 8.0f;
+      acc = metal::fma(weight, xv, acc);
+    }
+  }
+  const float value = metal::simd_sum(acc) * scale;
+  if (lane == 0u)
+    out[(ulong)sample * (ulong)H + row] = T(value);
+`;
+
+const Q8_GATE_UP_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint row = thread_position_in_grid.y * (uint)ROWS_TG + sg;
+  const uint sample = thread_position_in_grid.z;
+  if (row >= (uint)I || sample >= (uint)M) return;
+
+  const device int8_t* gateRow =
+    (const device int8_t*)(slot + GATE_W) + (ulong)row * (ulong)H;
+  const device int8_t* upRow =
+    (const device int8_t*)(slot + UP_W) + (ulong)row * (ulong)H;
+  const device float* gateScales = (const device float*)(slot + GATE_S);
+  const device float* upScales = (const device float*)(slot + UP_S);
+  const float gateScale = gateScales[row];
+  const float upScale = upScales[row];
+  float gateAcc = 0.0f;
+  float upAcc = 0.0f;
+  for (uint k = lane; k < (uint)H; k += 32u) {
+    const float xv = float(x[(ulong)sample * (ulong)H + k]);
+    gateAcc = metal::fma(float(gateRow[k]), xv, gateAcc);
+    upAcc = metal::fma(float(upRow[k]), xv, upAcc);
+  }
+  const float gate = metal::simd_sum(gateAcc) * gateScale;
+  const float up = metal::simd_sum(upAcc) * upScale;
+  if (lane == 0u) {
+    const T gateT = T(gate);
+    const T upT = T(up);
+    const T sigmoidT = T(1.0f / (1.0f + metal::precise::exp(-float(gateT))));
+    const T siluT = T(float(gateT) * float(sigmoidT));
+    mid[(ulong)sample * (ulong)I + row] = T(float(siluT) * float(upT));
+  }
+`;
+
+const Q8_DOWN_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint row = thread_position_in_grid.y * (uint)ROWS_TG + sg;
+  const uint sample = thread_position_in_grid.z;
+  if (row >= (uint)H || sample >= (uint)M) return;
+
+  const device int8_t* downRow =
+    (const device int8_t*)(slot + DOWN_W) + (ulong)row * (ulong)I;
+  const device float* downScales = (const device float*)(slot + DOWN_S);
+  const float scale = downScales[row];
+  float acc = 0.0f;
+  for (uint k = lane; k < (uint)I; k += 32u) {
+    const float xv = float(mid[(ulong)sample * (ulong)I + k]);
+    acc = metal::fma(float(downRow[k]), xv, acc);
+  }
+  const float value = metal::simd_sum(acc) * scale;
+  if (lane == 0u) out[(ulong)sample * (ulong)H + row] = T(value);
+`;
+
+/** What differs between the Q4 and Q8 kernel families: the Metal sources, the
+ *  weight packing the validators check, and the error wording. */
+interface CanonicalExecutorSpec {
+  readonly label: "Q4" | "Q8";
+  readonly kernelPrefix: string;
+  readonly gateUpSource: string;
+  readonly downSource: string;
+  /** Weights stored per byte (Q4 packs two, Q8 one). */
+  readonly weightsPerByte: 1 | 2;
+  /** Output dimensions must divide by this many packed weights (Q4: one uint32 word). */
+  readonly dimensionMultiple: number;
+  readonly subject: string;
+}
+
+function validateCall(
+  input: MlxArray,
+  slot: MlxArray,
+  layout: Glm52CanonicalMetalLayout,
+  spec: CanonicalExecutorSpec,
+): void {
+  const hidden = positiveInteger(layout.hiddenSize, "hidden size");
+  const intermediate = positiveInteger(
+    layout.intermediateSize,
+    "intermediate size",
+  );
+  positiveInteger(layout.slotBytes, "expert slot bytes");
+  if (hidden % spec.dimensionMultiple !== 0 || intermediate % spec.dimensionMultiple !== 0) {
+    throw new Error(
+      `${spec.label} expert dimensions must be divisible by ${spec.dimensionMultiple}`,
+    );
+  }
+  if (hidden % ROWS_PER_THREADGROUP !== 0 ||
+      intermediate % ROWS_PER_THREADGROUP !== 0) {
+    throw new Error(
+      `${spec.label} expert output dimensions must be divisible by ${ROWS_PER_THREADGROUP}`,
+    );
+  }
+  if (
+    input.shape.length !== 2 ||
+    !Number.isSafeInteger(input.shape[0]) ||
+    input.shape[0]! < 1 ||
+    input.shape[1] !== hidden
+  ) {
+    throw new Error(`${spec.subject} requires input [M,${hidden}]`);
+  }
+  if (input.dtype !== Dtype.bfloat16 && input.dtype !== Dtype.float32)
+    throw new Error(`${spec.subject} requires bfloat16 or float32 input`);
+  if (slot.dtype !== Dtype.uint8 || slot.shape.length !== 1)
+    throw new Error("canonical expert slot must be a flat uint8 array");
+  if (slot.size < layout.slotBytes)
+    throw new Error("canonical expert slot array is shorter than its layout");
+
+  const perByte = spec.weightsPerByte;
+  const ranges = [
+    [layout.downWeightOffset, hidden * intermediate / perByte, "down weights"],
+    [layout.gateWeightOffset, intermediate * hidden / perByte, "gate weights"],
+    [layout.upWeightOffset, intermediate * hidden / perByte, "up weights"],
+    [layout.downScaleOffset, hidden * 4, "down scales"],
+    [layout.gateScaleOffset, intermediate * 4, "gate scales"],
+    [layout.upScaleOffset, intermediate * 4, "up scales"],
+  ] as const;
+  for (const [offset, length, label] of ranges)
+    validateRange(offset, length, layout.slotBytes, label);
+}
+
+/**
+ * Row-independent routed-SwiGLU kernel pair (gate/up, then down) that reads a
+ * canonical residency slot directly. M=1 decode and a pinned speculative
+ * verify batch use the identical source and dispatch geometry; only grid.z
+ * changes. The Q4 and Q8 families below differ only in their spec.
+ */
+class Glm52CanonicalMetalExecutor {
+  readonly #gateUp: MetalKernel;
+  readonly #down: MetalKernel;
+  readonly #spec: CanonicalExecutorSpec;
+  #disposed = false;
+
+  constructor(spec: CanonicalExecutorSpec) {
+    this.#spec = spec;
+    this.#gateUp = new MetalKernel({
+      name: `${spec.kernelPrefix}_gate_up`,
+      inputNames: ["x", "slot"],
+      outputNames: ["mid"],
+      source: spec.gateUpSource,
+      ensureRowContiguous: true,
+    });
+    this.#down = new MetalKernel({
+      name: `${spec.kernelPrefix}_down`,
+      inputNames: ["mid", "slot"],
+      outputNames: ["out"],
+      source: spec.downSource,
+      ensureRowContiguous: true,
+    });
+  }
+
+  execute(
+    input: MlxArray,
+    slot: MlxArray,
+    layout: Glm52CanonicalMetalLayout,
+  ): MlxArray {
+    if (this.#disposed)
+      throw new Error(`${this.#spec.subject} executor used after dispose`);
+    validateCall(input, slot, layout, this.#spec);
+    const samples = input.shape[0]!;
+    const templateInts = {
+      M: samples,
+      H: layout.hiddenSize,
+      I: layout.intermediateSize,
+      DOWN_W: layout.downWeightOffset,
+      GATE_W: layout.gateWeightOffset,
+      UP_W: layout.upWeightOffset,
+      DOWN_S: layout.downScaleOffset,
+      GATE_S: layout.gateScaleOffset,
+      UP_S: layout.upScaleOffset,
+      ROWS_TG: ROWS_PER_THREADGROUP,
+    };
+    const [mid] = this.#gateUp.apply([input, slot], {
+      outputs: [{
+        shape: [samples, layout.intermediateSize],
+        dtype: input.dtype,
+      }],
+      grid: [
+        THREADS,
+        layout.intermediateSize / ROWS_PER_THREADGROUP,
+        samples,
+      ],
+      threadGroup: [THREADS, 1, 1],
+      templateDtypes: { T: input.dtype },
+      templateInts,
+    });
+    if (!mid) throw new Error(`${this.#spec.subject} gate/up kernel returned no output`);
+    try {
+      const [output] = this.#down.apply([mid, slot], {
+        outputs: [{
+          shape: [samples, layout.hiddenSize],
+          dtype: input.dtype,
+        }],
+        grid: [
+          THREADS,
+          layout.hiddenSize / ROWS_PER_THREADGROUP,
+          samples,
+        ],
+        threadGroup: [THREADS, 1, 1],
+        templateDtypes: { T: input.dtype },
+        templateInts,
+      });
+      if (!output) throw new Error(`${this.#spec.subject} down kernel returned no output`);
+      return output;
+    } finally {
+      mid.dispose();
+    }
+  }
+
+  dispose(): void {
+    if (this.#disposed) return;
+    this.#disposed = true;
+    this.#gateUp.dispose();
+    this.#down.dispose();
+  }
+}
+
+/** Q4 (packed nibbles, affine zero-point 8) kernel family for the GLM-5.2
+ *  production geometry. */
+export class Glm52CanonicalQ4MetalExecutor extends Glm52CanonicalMetalExecutor {
+  constructor() {
+    super({
+      label: "Q4",
+      kernelPrefix: "mlx_bun_glm52_q4_slot",
+      gateUpSource: GATE_UP_SOURCE,
+      downSource: DOWN_SOURCE,
+      weightsPerByte: 2,
+      dimensionMultiple: 8,
+      subject: "GLM Metal streamed expert decode",
+    });
+  }
+}
+
+/** Fixed M=1..gamma signed-Q8 kernel family for native MTP draft and accepted
+ *  token absorption. The row dot-product and materialization boundaries are
+ *  identical for every batch width; only grid.z changes. */
+export class Glm52CanonicalQ8MetalExecutor extends Glm52CanonicalMetalExecutor {
+  constructor() {
+    super({
+      label: "Q8",
+      kernelPrefix: "mlx_bun_glm52_q8_slot",
+      gateUpSource: Q8_GATE_UP_SOURCE,
+      downSource: Q8_DOWN_SOURCE,
+      weightsPerByte: 1,
+      dimensionMultiple: 1,
+      subject: "GLM MTP Metal streamed expert",
+    });
+  }
+}

package/src/kernels/logits.ts

@@ -0,0 +1,38 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+
+
+/** tanh(x / cap) * cap with weak-scalar semantics, UNFUSED (standalone div, tanh,
+ *  mul), as plain-Python oracle sites compose it. The mx.compile'd one
+ *  (gemma4_text.py) is `compiledLogitSoftcap` in softcap.ts; the two are
+ *  separate on purpose so each site dispatches the same kernels as its oracle. */
+export function logitSoftcap(logits: MlxArray, cap: number): MlxArray {
+  const capArr = ops.scalarLike(cap, logits);
+  const scaled = ops.div(logits, capArr);
+  const t = ops.tanh(scaled);
+  scaled.dispose();
+  const out = ops.mul(t, capArr);
+  t.dispose();
+  capArr.dispose();
+  return out;
+}
+
+/** argmax over vocab at the last sequence position of [1, L, V]. */
+export function argmaxLastPosition(logits: MlxArray): number {
+  const [, L, V] = logits.shape as [number, number, number];
+  const last = logits.slice([0, L - 1, 0], [1, L, V]);
+  const am = ops.argmaxAxis(last, -1);
+  const id = ops.itemUint32(am);
+  last.dispose();
+  am.dispose();
+  return id;
+}
+
+/** Last-position logits as f32 (for the parity harness). */
+export function lastPositionLogits(logits: MlxArray): Float32Array {
+  const [, L, V] = logits.shape as [number, number, number];
+  const last = logits.slice([0, L - 1, 0], [1, L, V]);
+  const f32 = last.toFloat32();
+  last.dispose();
+  return f32;
+}

package/src/kernels/quantization/affine-row.ts

@@ -0,0 +1,154 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+// One-row matmul over MLX affine 3- or 4-bit, group-64 weights: y = x · Wᵀ for
+// a single activation row, bit for bit MLX 0.32.2's affine_qmv_fast (the
+// kernel MLX's quantized matmul runs at one row). Every output keeps MLX's
+// arithmetic: each lane owns 16 values of every 512-value K step (lane·16 on),
+// computes MLX's load_vector (x_thread, its bf16 sum) and qdot over them,
+// adds each step's qdot into its partial from 0 in K order, and the 32
+// partials meet in simd_sum. Only what does not touch a bit changes:
+// - load_vector runs once per threadgroup, every 16-value chunk of x into
+//   threadgroup memory, where MLX recomputes it per simdgroup and K step (its
+//   bf16 sum chain is software-rounded on the M1 Max);
+// - each masked weight field's float is formed as (q | 2^23 bits) − 2^23, which
+//   is float(q) exactly (q < 2^23), in place of an integer convert;
+// - R = 4 rows per simdgroup and S = 16 simdgroups per threadgroup (MLX: 4
+//   and 2), so 64 rows share one threadgroup's x.
+// Measured on the M1 Max, each layer's weights back to back, against MLX's
+// quantized matmul at one row: 3-bit 5120 to 10240 outputs 0.082 ms against
+// 0.100, 5120 to 6144 0.053 against 0.062, 6144 to 5120 0.054 against 0.062,
+// 4-bit 5120 to 248320 (the head) 1.96 against 2.26; 4-bit projections of 1024
+// to 12288 outputs level or up to 3% slower (left on MLX's).
+const R = 4, S = 16;
+
+const HEADER = String.raw`
+// MLX 0.32.2 load_vector (quantized.h) at values_per_thread 16, bits 3 and 4.
+template <typename T, int bits>
+inline float mlxq_load_vector(const device T* x, thread float* x_thread) {
+  float sum = 0;
+  if (bits == 3) {
+    for (int i = 0; i < 16; i += 8) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
+          x[i + 6] + x[i + 7];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 8.0f;
+      x_thread[i + 2] = x[i + 2] / 64.0f;
+      x_thread[i + 3] = x[i + 3] / 2.0f;
+      x_thread[i + 4] = x[i + 4] / 16.0f;
+      x_thread[i + 5] = x[i + 5] / 128.0f;
+      x_thread[i + 6] = x[i + 6] / 4.0f;
+      x_thread[i + 7] = x[i + 7] / 32.0f;
+    }
+  } else {
+    for (int i = 0; i < 16; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 16.0f;
+      x_thread[i + 2] = x[i + 2] / 256.0f;
+      x_thread[i + 3] = x[i + 3] / 4096.0f;
+    }
+  }
+  return sum;
+}
+
+// float(q) for a masked field q < 2^23: q in the mantissa of 2^23, minus 2^23.
+#define MLXQ_F(q) (as_type<float>((uint)(q) | 0x4B000000u) - 8388608.0f)
+
+// MLX 0.32.2 qdot (quantized.h) at values_per_thread 16, bits 3 and 4: the
+// same products and additions in the same order.
+template <int bits>
+inline float mlxq_qdot(const device uint8_t* w, const threadgroup float* x_thread, float scale, float bias, float sum) {
+  float accum = 0;
+  if (bits == 3) {
+    for (int i = 0; i < 2; i++) {
+      x_thread += 8 * i;
+      w += 3 * i;
+      accum += MLXQ_F(w[0] & 0x07) * x_thread[0];
+      accum += MLXQ_F(w[0] & 0x38) * x_thread[1];
+      accum += MLXQ_F(w[0] & 0xc0) * x_thread[2];
+      accum += MLXQ_F(w[1] & 0x01) * (x_thread[2] * 256.0f);
+      accum += MLXQ_F(w[1] & 0x0e) * x_thread[3];
+      accum += MLXQ_F(w[1] & 0x70) * x_thread[4];
+      accum += MLXQ_F(w[1] & 0x80) * x_thread[5];
+      accum += MLXQ_F(w[2] & 0x03) * (x_thread[5] * 256.0f);
+      accum += MLXQ_F(w[2] & 0x1c) * x_thread[6];
+      accum += MLXQ_F(w[2] & 0xe0) * x_thread[7];
+    }
+  } else {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < 4; i++) {
+      accum +=
+          (x_thread[4 * i] * MLXQ_F(ws[i] & 0x000f) +
+           x_thread[4 * i + 1] * MLXQ_F(ws[i] & 0x00f0) +
+           x_thread[4 * i + 2] * MLXQ_F(ws[i] & 0x0f00) +
+           x_thread[4 * i + 3] * MLXQ_F(ws[i] & 0xf000));
+    }
+  }
+  return scale * accum + sum * bias;
+}
+`;
+
+const SOURCE = String.raw`
+  constexpr int BPL = 2 * BITS;                   // weight bytes per lane per K step
+  constexpr int ROWB = K * BITS / 8;
+  constexpr int GROUPS = K / 64;
+  const uint sg = simdgroup_index_in_threadgroup, lane = thread_index_in_simdgroup;
+  // MLX's x_thread for every 16-value chunk of x, and each chunk's sum.
+  threadgroup float xs[K];
+  threadgroup float ss[K / 16];
+  for (int c = (int)(sg * 32 + lane); c < K / 16; c += S * 32) {
+    thread float t[16];
+    ss[c] = mlxq_load_vector<T, BITS>(x + c * 16, t);
+    for (int j = 0; j < 16; j++) xs[c * 16 + j] = t[j];
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  const int row0 = (threadgroup_position_in_grid.y * S + sg) * R;
+  const device uint8_t* ws = (const device uint8_t*)w + row0 * ROWB + lane * BPL;
+  const device T* sc = scales + row0 * GROUPS + lane / 4;
+  const device T* bi = biases + row0 * GROUPS + lane / 4;
+  float result[R] = {0};
+  #pragma clang loop unroll_count(1)
+  for (int k = 0; k < K; k += 512) {
+    const float sum = ss[(k + lane * 16) / 16];
+    const threadgroup float* x_thread = xs + k + lane * 16;
+    for (int r = 0; r < R; r++) {
+      const float s = sc[r * GROUPS];
+      const float b = bi[r * GROUPS];
+      result[r] += mlxq_qdot<BITS>(ws + r * ROWB, x_thread, s, b, sum);
+    }
+    ws += 32 * BPL;
+    sc += 8;
+    bi += 8;
+  }
+  for (int r = 0; r < R; r++) {
+    const float v = simd_sum(result[r]);
+    if (lane == 0) y[row0 + r] = static_cast<T>(v);
+  }
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** Eligible: affine 3- or 4-bit, group 64, one row, K a multiple of 512 whose x_thread fits threadgroup memory
+ *  (K <= 7680), N a multiple of 64. */
+export function affineRowEligible(rows: number, w: MlxArray, scales: MlxArray, spec: { bits: number; groupSize: number; mode: string }): boolean {
+  if (spec.mode !== "affine" || spec.groupSize !== 64 || (spec.bits !== 3 && spec.bits !== 4)) return false;
+  const N = scales.shape[0]!, K = scales.shape[1]! * 64;
+  return rows === 1 && K % 512 === 0 && K <= 7680 && N % (R * S) === 0 && w.shape.length === 2;
+}
+
+/** x [..., K] (one row) · Wᵀ → [..., N] in x's dtype, for `bits`-bit weights: MLX's quantized matmul's bits.
+ *  Borrows inputs. */
+export function affineRow(x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray, bits: 3 | 4): MlxArray {
+  const K = x.shape.at(-1)!, M = x.size / K, N = scales.shape[0]!;
+  if (M !== 1) throw new Error(`affineRow: M=${M}, one row only`);
+  kernel ??= new MetalKernel({ name: "mlx_bun_affine_row", inputNames: ["x", "w", "scales", "biases"],
+    outputNames: ["y"], source: SOURCE, header: HEADER, ensureRowContiguous: true });
+  const [y] = kernel.apply([x, w, scales, biases], {
+    outputs: [{ shape: [...x.shape.slice(0, -1), N], dtype: x.dtype }],
+    grid: [32, N / R, 1], threadGroup: [32, S, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { N, K, BITS: bits, R, S },
+  });
+  return y!;
+}

package/src/kernels/quantization/affine-rows.ts

@@ -0,0 +1,36 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/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();
+  }
+}