@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/attention/kvq4-decode.ts

@@ -0,0 +1,155 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+// One-row decode attention over a 4-bit, group-64 affine KV cache, head dim
+// 256: the shape of MLX's duplication-free GQA vector kernel
+// (sdpa_vector_2pass_1_gqa: each simdgroup reads a K/V row once into registers
+// and serves several query heads) with dequantize-on-load. Each lane owns 8 of
+// the 256 dims: one 4-byte load of packed K (and V) per row plus its group's
+// scale and bias. A threadgroup = one KV head × one block of BLOCK tokens
+// (the caller's `block`), split across SUB = 4 token sub-chunks × 2 head groups
+// (REP/2 heads each; threadgroup memory bounds SUB at 4 for D 256); the
+// sub-chunks combine in threadgroup memory and a second kernel merges blocks.
+// The cache length is an argument, not a specialization: one pipeline serves
+// every decode step. K/V bytes per token vs bf16: ~0.28× (with scales). Online
+// softmax in f32: agrees with the unfused path to bf16 precision, not bit for
+// bit; the block partitions the softmax, so it is part of the numbers.
+const SUB = 4;
+/** The tokens per threadgroup of the M4 Pro composition, which measured it. */
+export const KV4_DECODE_BLOCK = 512;
+const PARTIAL_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint hk = threadgroup_position_in_grid.y, block = threadgroup_position_in_grid.x;
+  const uint hg = sg % 2u, sc = sg / 2u;              // head group, token sub-chunk
+  constexpr uint HPG = REP / 2;                       // heads per group
+  const uint N = kp_shape[2];                         // cached positions: an argument, not a specialization
+  const uint t0 = block * (uint)BLOCK, tEnd = metal::min(N, t0 + (uint)BLOCK);
+  const uint sub = ((uint)BLOCK + (uint)SUB - 1u) / (uint)SUB;
+  const uint s0 = t0 + sc * sub, s1 = metal::min(tEnd, s0 + sub);
+  // Queries (pre-scaled), 8 dims per lane.
+  float qv[HPG][8];
+  #pragma clang loop unroll(full)
+  for (uint j = 0; j < HPG; ++j)
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) qv[j][i] = float(q[(hk * (uint)REP + hg * HPG + j) * (uint)D + lane * 8u + i]);
+  float m[HPG], l[HPG], o[HPG][8];
+  #pragma clang loop unroll(full)
+  for (uint j = 0; j < HPG; ++j) { m[j] = -INFINITY; l[j] = 0.0f; for (uint i = 0; i < 8u; ++i) o[j][i] = 0.0f; }
+  const uint g = lane / 8u;                           // this lane's 64-dim group
+  const ulong kph = (ulong)kp_strides[1], kpr = (ulong)kp_strides[2];
+  const ulong ksh = (ulong)ks_strides[1], ksr = (ulong)ks_strides[2];
+  const ulong vph = (ulong)vp_strides[1], vpr = (ulong)vp_strides[2];
+  const ulong vsh = (ulong)vs_strides[1], vsr = (ulong)vs_strides[2];
+  for (uint t = s0; t < s1; ++t) {
+    const uint kword = kp[(ulong)hk * kph + (ulong)t * kpr + lane];
+    const uint vword = vp[(ulong)hk * vph + (ulong)t * vpr + lane];
+    const float kscale = float(ks[(ulong)hk * ksh + (ulong)t * ksr + g]), kbias = float(kb[(ulong)hk * ksh + (ulong)t * ksr + g]);
+    const float vscale = float(vs[(ulong)hk * vsh + (ulong)t * vsr + g]), vbias = float(vb[(ulong)hk * vsh + (ulong)t * vsr + g]);
+    float kr[8], vr[8];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) {
+      kr[i] = metal::fma(kscale, float((kword >> (4u * i)) & 15u), kbias);
+      vr[i] = metal::fma(vscale, float((vword >> (4u * i)) & 15u), vbias);
+    }
+    #pragma clang loop unroll(full)
+    for (uint j = 0; j < HPG; ++j) {
+      float score = 0.0f;
+      #pragma clang loop unroll(full)
+      for (uint i = 0; i < 8u; ++i) score = metal::fma(qv[j][i], kr[i], score);
+      score = metal::simd_sum(score);
+      const float mn = metal::max(m[j], score);
+      const float factor = metal::exp(m[j] - mn), e = metal::exp(score - mn);
+      m[j] = mn;
+      l[j] = l[j] * factor + e;
+      #pragma clang loop unroll(full)
+      for (uint i = 0; i < 8u; ++i) o[j][i] = metal::fma(o[j][i], factor, e * vr[i]);
+    }
+  }
+  // Combine the SUB token sub-chunks of each head in threadgroup memory.
+  threadgroup float osh[SUB * REP * D];
+  threadgroup float msh[SUB * REP], lsh[SUB * REP];
+  #pragma clang loop unroll(full)
+  for (uint j = 0; j < HPG; ++j) {
+    const uint h = hg * HPG + j, slot = sc * (uint)REP + h;
+    for (uint i = 0; i < 8u; ++i) osh[slot * (uint)D + lane * 8u + i] = o[j][i];
+    if (lane == 0u) { msh[slot] = m[j]; lsh[slot] = l[j]; }
+  }
+  threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  for (uint e = thread_position_in_threadgroup.x; e < (uint)(REP * D); e += threads_per_threadgroup.x) {
+    const uint h = e / (uint)D, d = e % (uint)D;
+    float gmax = -INFINITY;
+    for (uint s = 0; s < (uint)SUB; ++s) gmax = metal::max(gmax, msh[s * (uint)REP + h]);
+    float acc = 0.0f, den = 0.0f;
+    for (uint s = 0; s < (uint)SUB; ++s) {
+      const float ms = msh[s * (uint)REP + h];
+      if (ms == -INFINITY) continue;
+      const float w = metal::exp(ms - gmax);
+      acc += w * osh[(s * (uint)REP + h) * (uint)D + d];
+      den += w * lsh[s * (uint)REP + h];
+    }
+    const ulong base = ((ulong)block * (uint)KVH + hk) * (uint)REP + h;
+    opart[base * (uint)D + d] = acc;
+    if (d == 0u) { mpart[base] = gmax; lpart[base] = den; }
+  }
+`;
+
+const MERGE_SOURCE = String.raw`
+  const uint d = thread_position_in_grid.x, qv = thread_position_in_grid.y;
+  if (d >= (uint)D || qv >= (uint)(KVH * REP)) return;
+  const uint BLOCKS = mpart_shape[0];
+  float mmax = -INFINITY;
+  for (uint s = 0; s < BLOCKS; ++s) mmax = metal::max(mmax, mpart[(ulong)s * (uint)(KVH * REP) + qv]);
+  float num = 0.0f, den = 0.0f;
+  for (uint s = 0; s < BLOCKS; ++s) {
+    const ulong b = (ulong)s * (uint)(KVH * REP) + qv;
+    const float ms = mpart[b];
+    if (ms == -INFINITY) continue;
+    const float w = metal::exp(ms - mmax);
+    num += w * opart[b * (uint)D + d];
+    den += w * lpart[b];
+  }
+  out[(ulong)qv * (uint)D + d] = T(num / den);
+`;
+
+let partial: MetalKernel | undefined, merge: MetalKernel | undefined;
+
+/** One query row, head dim 256, a 4-bit group-64 cache, an even GQA factor. */
+export function kvq4DecodeEligible(q: MlxArray, keys: ops.QuantizedTensor, groupSize: number, bits: number): boolean {
+  const [B, H, L, D] = q.shape as [number, number, number, number];
+  const KVH = keys.packed.shape[1]!;
+  return B === 1 && L === 1 && D === 256 && bits === 4 && groupSize === 64 && H % KVH === 0 && (H / KVH) % 2 === 0;
+}
+
+/** q [1, H, 1, 256]; keys/values: 4-bit group-64 affine [1, KVH, N, ·] (any strides) → [1, H, 1, 256].
+ *  `block`: tokens per threadgroup, a device's measured choice. */
+export function kvq4DecodeSdpa(q: MlxArray, keys: ops.QuantizedTensor, values: ops.QuantizedTensor, scale: number,
+  block = KV4_DECODE_BLOCK): MlxArray {
+  const [, H, , D] = q.shape as [number, number, number, number];
+  const KVH = keys.packed.shape[1]!, N = keys.packed.shape[2]!, REP = H / KVH;
+  const blocks = Math.ceil(N / block);
+  partial ??= new MetalKernel({ name: "mlx_bun_kvq4_decode_partial", inputNames: ["q", "kp", "ks", "kb", "vp", "vs", "vb"],
+    outputNames: ["opart", "mpart", "lpart"], source: PARTIAL_SOURCE, ensureRowContiguous: false });
+  merge ??= new MetalKernel({ name: "mlx_bun_kvq4_decode_merge", inputNames: ["opart", "mpart", "lpart"],
+    outputNames: ["out"], source: MERGE_SOURCE, ensureRowContiguous: true });
+  using qs = ops.mulScalar(q, scale);
+  using qc = ops.contiguous(qs);
+  const [opart, mpart, lpart] = partial.apply([qc, keys.packed, keys.scales, keys.biases, values.packed, values.scales, values.biases], {
+    outputs: [{ shape: [blocks, KVH, REP, D], dtype: Dtype.float32 }, { shape: [blocks, KVH, REP], dtype: Dtype.float32 },
+      { shape: [blocks, KVH, REP], dtype: Dtype.float32 }],
+    grid: [32 * 2 * SUB * blocks, KVH, 1], threadGroup: [32 * 2 * SUB, 1, 1],
+    templateDtypes: { T: q.dtype },
+    templateInts: { D, KVH, REP, BLOCK: block, SUB },
+  });
+  try {
+    const [out] = merge.apply([opart!, mpart!, lpart!], {
+      outputs: [{ shape: [1, H, 1, D], dtype: q.dtype }],
+      grid: [D, KVH * REP, 1], threadGroup: [Math.min(256, D), 1, 1],
+      templateDtypes: { T: q.dtype },
+      templateInts: { D, KVH, REP },
+    });
+    return out!;
+  } finally { opart!.dispose(); mpart!.dispose(); lpart!.dispose(); }
+}

package/src/kernels/attention/masks.ts

@@ -0,0 +1,40 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Mask } from "../../contracts/mlx/cache";
+
+
+/** Port of base.py create_causal_mask (bool, [N, offset+N]). */
+export function createCausalMask(N: number, offset: number, windowSize: number | null): MlxArray {
+  const rinds = ops.arange(0, offset + N, 1, Dtype.int32);
+  const lindsFlat = offset ? ops.arange(offset, offset + N, 1, Dtype.int32) : rinds;
+  const linds = ops.reshape(lindsFlat, [N, 1]);
+  const rindsB = ops.reshape(rinds, [1, offset + N]);
+  let mask = ops.greaterEqual(linds, rindsB);
+  if (windowSize !== null) {
+    const w = ops.fromInt32([windowSize], []);
+    const rPlusW = ops.add(rindsB, w);
+    const inWindow = ops.less(linds, rPlusW);
+    const combined = ops.logicalAnd(mask, inWindow);
+    for (const a of [w, rPlusW, inWindow, mask]) a.dispose();
+    mask = combined;
+  }
+  if (lindsFlat !== rinds) lindsFlat.dispose();
+  rinds.dispose();
+  linds.dispose();
+  rindsB.dispose();
+  return mask;
+}
+/** Causal(+window) mask OR'd with image×image bidirectional attention. */
+export function bidirMask(L: number, windowSize: number | null, bidir: MlxArray): Mask {
+  const causal = createCausalMask(L, 0, windowSize);
+  const col = ops.reshape(bidir, [L, 1]);
+  const row = ops.reshape(bidir, [1, L]);
+  const outer = ops.logicalAnd(col, row);
+  col.dispose();
+  row.dispose();
+  const allow = ops.logicalOr(causal, outer);
+  causal.dispose();
+  outer.dispose();
+  return { mode: "array", arr: allow };
+}

package/src/kernels/attention/multi-query.ts

@@ -0,0 +1,193 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+// Multi-query flash-decoding: causal attention for a few query rows (speculative verify, L <= 8) over a long KV
+// cache, reading every K and V row ONCE for all query vectors of a KV head.
+// MLX's vector kernel accepts at most 32 query vectors per KV head (5 rows at
+// GQA 6) and re-reads K/V per query row; above that it falls back to unfused
+// attention. Here one threadgroup owns one KV head and one key split; its
+// rep·L query vectors (padded to 8-row tiles, one simdgroup per tile) multiply
+// each 32-key block staged once in threadgroup memory on the simdgroup matrix
+// unit, with an online softmax per row. A second kernel merges the splits.
+// Row (r, i) — head kv·rep + r, query i — sees keys j <= N - L + i.
+// Not bit-identical to MLX's kernels (summation order); gate by KL.
+const PARTIAL_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint hk = threadgroup_position_in_grid.y;
+  const uint split = threadgroup_position_in_grid.x;
+  const uint fm = ((lane / 4u) & 4u) + ((lane / 2u) % 4u);
+  const uint fn = ((lane / 4u) & 2u) * 2u + (lane % 2u) * 2u;
+  threadgroup half kv[BK * D];
+  // Query tile: rows sg*8 .. sg*8+7 of the head group's rep·L vectors.
+  const bool tile = sg * 8u < (uint)NQ;               // simdgroups past the tiles only stage K/V
+  const uint row = sg * 8u + fm;                      // this lane's A-fragment row
+  const bool rowOk = row < (uint)NQ;
+  const uint rr = rowOk ? row / (uint)L : 0u, qi = rowOk ? row % (uint)L : 0u;
+  const device T* qp = q + ((ulong)(hk * (uint)REP + rr) * (uint)L + qi) * (uint)D;
+  metal::simdgroup_half8x8 qf[D / 8];
+  #pragma clang loop unroll(full)
+  for (uint f = 0; f < (uint)D / 8u; ++f) {
+    qf[f].thread_elements()[0] = rowOk ? half(qp[8u * f + fn]) : half(0.0h);
+    qf[f].thread_elements()[1] = rowOk ? half(qp[8u * f + fn + 1u]) : half(0.0h);
+  }
+  metal::simdgroup_float8x8 o[D / 8];
+  #pragma clang loop unroll(full)
+  for (uint f = 0; f < (uint)D / 8u; ++f) o[f] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+  float m = -INFINITY, l = 0.0f;
+  // Key range of this split; the row's causal limit.
+  const uint per = ((uint)N + (uint)SPLITS - 1u) / (uint)SPLITS;
+  const uint k0 = split * per, k1 = metal::min((uint)N, k0 + per);
+  const uint limit = (uint)N - (uint)L + qi;          // keys <= limit are visible
+  const ulong hs = (ulong)k_strides[1], rs = (ulong)k_strides[2];
+  const ulong vhs = (ulong)v_strides[1], vrs = (ulong)v_strides[2];
+  for (uint kb = k0; kb < k1; kb += (uint)BK) {
+    // Stage K block [BK, D] as half.
+    for (uint i = thread_position_in_threadgroup.x; i < (uint)(BK * D / 8); i += threads_per_threadgroup.x) {
+      const uint j = (8u * i) / (uint)D, d = (8u * i) % (uint)D, key = kb + j;
+      threadgroup half* dst = kv + j * (uint)D + d;
+      if (key < k1) {
+        const device metal::vec<T, 4>* src = (const device metal::vec<T, 4>*)(k + (ulong)hk * hs + (ulong)key * rs + d);
+        const float4 a = float4(src[0]), b = float4(src[1]);
+        *(threadgroup half4*)dst = half4(a); *(threadgroup half4*)(dst + 4) = half4(b);
+      } else { *(threadgroup half4*)dst = half4(0.0h); *(threadgroup half4*)(dst + 4) = half4(0.0h); }
+    }
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    // S = Q · Kᵀ for BK keys: BK/8 fragments of [8 rows × 8 keys].
+    metal::simdgroup_float8x8 s[BK / 8];
+    metal::simdgroup_half8x8 p[BK / 8];
+    if (tile) {
+    #pragma clang loop unroll(full)
+    for (uint c = 0; c < (uint)BK / 8u; ++c) {
+      s[c] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+      #pragma clang loop unroll(full)
+      for (uint f = 0; f < (uint)D / 8u; ++f) {
+        metal::simdgroup_half8x8 kt;
+        metal::simdgroup_load(kt, kv + (8u * c) * (uint)D + 8u * f, (ulong)D, ulong2(0, 0), true);
+        metal::simdgroup_multiply_accumulate(s[c], qf[f], kt, s[c]);
+      }
+    }
+    // Causal / range mask and online softmax for this lane's row fm.
+    float bmax = -INFINITY;
+    #pragma clang loop unroll(full)
+    for (uint c = 0; c < (uint)BK / 8u; ++c)
+      #pragma clang loop unroll(full)
+      for (uint e = 0; e < 2u; ++e) {
+        const uint key = kb + 8u * c + fn + e;
+        float v = s[c].thread_elements()[e];
+        if (!rowOk || key >= k1 || key > limit) v = -INFINITY;
+        s[c].thread_elements()[e] = v;
+        bmax = metal::max(bmax, v);
+      }
+    bmax = metal::max(bmax, metal::simd_shuffle_xor(bmax, 1u));
+    bmax = metal::max(bmax, metal::simd_shuffle_xor(bmax, 8u));
+    const float mnew = metal::max(m, bmax);
+    const float corr = mnew == -INFINITY ? 1.0f : metal::exp(m - mnew);
+    float bsum = 0.0f;
+    #pragma clang loop unroll(full)
+    for (uint c = 0; c < (uint)BK / 8u; ++c)
+      #pragma clang loop unroll(full)
+      for (uint e = 0; e < 2u; ++e) {
+        const float v = s[c].thread_elements()[e];
+        const float w = v == -INFINITY ? 0.0f : metal::exp(v - mnew);
+        bsum += w;
+        p[c].thread_elements()[e] = half(w);
+      }
+    bsum += metal::simd_shuffle_xor(bsum, 1u);
+    bsum += metal::simd_shuffle_xor(bsum, 8u);
+    l = l * corr + bsum;
+    m = mnew;
+    #pragma clang loop unroll(full)
+    for (uint f = 0; f < (uint)D / 8u; ++f) {
+      o[f].thread_elements()[0] *= corr;
+      o[f].thread_elements()[1] *= corr;
+    }
+    }
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    // Stage V block [BK, D] as half; O += P · V.
+    for (uint i = thread_position_in_threadgroup.x; i < (uint)(BK * D / 8); i += threads_per_threadgroup.x) {
+      const uint j = (8u * i) / (uint)D, d = (8u * i) % (uint)D, key = kb + j;
+      threadgroup half* dst = kv + j * (uint)D + d;
+      if (key < k1) {
+        const device metal::vec<T, 4>* src = (const device metal::vec<T, 4>*)(v + (ulong)hk * vhs + (ulong)key * vrs + d);
+        const float4 a = float4(src[0]), b = float4(src[1]);
+        *(threadgroup half4*)dst = half4(a); *(threadgroup half4*)(dst + 4) = half4(b);
+      } else { *(threadgroup half4*)dst = half4(0.0h); *(threadgroup half4*)(dst + 4) = half4(0.0h); }
+    }
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    if (tile)
+    #pragma clang loop unroll(full)
+    for (uint f = 0; f < (uint)D / 8u; ++f)
+      #pragma clang loop unroll(full)
+      for (uint c = 0; c < (uint)BK / 8u; ++c) {
+        metal::simdgroup_half8x8 vb;
+        metal::simdgroup_load(vb, kv + (8u * c) * (uint)D + 8u * f, (ulong)D);
+        metal::simdgroup_multiply_accumulate(o[f], p[c], vb, o[f]);
+      }
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  }
+  // Partial results: unnormalized O, running max and sum per row.
+  if (!tile || !rowOk) return;
+  const ulong base = ((ulong)split * (uint)KVH + hk) * (uint)NQ + row;
+  #pragma clang loop unroll(full)
+  for (uint f = 0; f < (uint)D / 8u; ++f) {
+    opart[base * (uint)D + 8u * f + fn] = o[f].thread_elements()[0];
+    opart[base * (uint)D + 8u * f + fn + 1u] = o[f].thread_elements()[1];
+  }
+  if (fn == 0u) { mpart[base] = m; lpart[base] = l; }
+`;
+
+const MERGE_SOURCE = String.raw`
+  const uint d = thread_position_in_grid.x;            // 0..D-1
+  const uint qv = thread_position_in_grid.y;           // 0..KVH·NQ-1
+  if (d >= (uint)D || qv >= (uint)(KVH * NQ)) return;
+  const uint hk = qv / (uint)NQ, row = qv % (uint)NQ;
+  float mmax = -INFINITY;
+  for (uint s = 0; s < (uint)SPLITS; ++s) mmax = metal::max(mmax, mpart[(ulong)s * (uint)(KVH * NQ) + qv]);
+  float num = 0.0f, den = 0.0f;
+  for (uint s = 0; s < (uint)SPLITS; ++s) {
+    const ulong b = (ulong)s * (uint)(KVH * NQ) + qv;
+    const float ms = mpart[b];
+    if (ms == -INFINITY) continue;
+    const float w = metal::exp(ms - mmax);
+    num += w * opart[b * (uint)D + d];
+    den += w * lpart[b];
+  }
+  const uint rr = row / (uint)L, qi = row % (uint)L;
+  out[((ulong)(hk * (uint)REP + rr) * (uint)L + qi) * (uint)D + d] = T(num / den);
+`;
+
+let partial: MetalKernel | undefined, merge: MetalKernel | undefined;
+const BK = 32;
+
+/** q [1, H, L, D], keys/values [1, KVH, N, D] (any head/row strides) → [1, H, L, D]. */
+export function multiQueryCausalSdpa(q: MlxArray, keys: MlxArray, values: MlxArray, scale: number): MlxArray {
+  const [, H, L, D] = q.shape as [number, number, number, number];
+  const KVH = keys.shape[1]!, N = keys.shape[2]!, REP = H / KVH, NQ = REP * L;
+  if (NQ > 64) throw new Error(`multiQueryCausalSdpa: ${NQ} query vectors per KV head exceed 64`);
+  const P = Math.max(1, Math.min(128, Math.ceil(N / 256)));   // key splits
+  partial ??= new MetalKernel({ name: "mlx_bun_attn_multi_query_partial", inputNames: ["q", "k", "v"],
+    outputNames: ["opart", "mpart", "lpart"], source: PARTIAL_SOURCE, ensureRowContiguous: false });
+  merge ??= new MetalKernel({ name: "mlx_bun_attn_multi_query_merge", inputNames: ["opart", "mpart", "lpart"],
+    outputNames: ["out"], source: MERGE_SOURCE, ensureRowContiguous: true });
+  using qs = ops.mulScalar(q, scale);                       // softmax scale folded into q
+  using qc = ops.contiguous(qs);
+  const [opart, mpart, lpart] = partial.apply([qc, keys, values], {
+    outputs: [{ shape: [P, KVH, NQ, D], dtype: Dtype.float32 }, { shape: [P, KVH, NQ], dtype: Dtype.float32 },
+      { shape: [P, KVH, NQ], dtype: Dtype.float32 }],
+    grid: [256 * P, KVH, 1], threadGroup: [256, 1, 1],
+    templateDtypes: { T: q.dtype },
+    templateInts: { L, D, N, KVH, REP, NQ, BK, SPLITS: P },
+  });
+  try {
+    const [out] = merge.apply([opart!, mpart!, lpart!], {
+      outputs: [{ shape: [1, H, L, D], dtype: q.dtype }],
+      grid: [D, KVH * NQ, 1], threadGroup: [Math.min(256, D), 1, 1],
+      templateDtypes: { T: q.dtype },
+      templateInts: { L, D, KVH, REP, NQ, SPLITS: P },
+    });
+    return out!;
+  } finally { opart!.dispose(); mpart!.dispose(); lpart!.dispose(); }
+}

package/src/kernels/attention/paged.ts

@@ -0,0 +1,130 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { KvAttentionView, Mask } from "../../contracts/mlx/cache";
+import { disposeResources } from "../../runtime/resources";
+/** Borrowed numerical storage. Allocation, appends, and eviction belong to state. */
+export interface PagedAttentionStorage {
+  readonly keys: MlxArray;
+  readonly values: MlxArray;
+  readonly blockSize: number;
+  readonly headDim: number;
+  readonly vHeadDim: number;
+  readonly dtype: Dtype;
+  readonly quantization?: { bits: number; groupSize: number } | null;
+  readonly keyScales?: MlxArray | null;
+  readonly keyBiases?: MlxArray | null;
+  readonly valueScales?: MlxArray | null;
+  readonly valueBiases?: MlxArray | null;
+}
+
+// Each SIMD group reduces one query/head over a bounded sequence partition.
+// Only partial output vectors and normalizers materialize, never gathered KV.
+const partialSpec = { name: "paged_attention_partial", ensureRowContiguous: true,
+  inputNames: ["q", "keys", "values", "table", "scale"], outputNames: ["out", "normal"], source: `
+  uint lane = thread_position_in_threadgroup.x;
+  uint hq = threadgroup_position_in_grid.y % HQ;
+  uint query = threadgroup_position_in_grid.y / HQ;
+  uint part = threadgroup_position_in_grid.z;
+  uint hk = hq / (HQ / HK);
+  float accum[(DV + 31) / 32];
+  for (uint j = 0; j < (DV + 31) / 32; ++j) accum[j] = 0;
+  float maximum = -INFINITY, total = 0;
+  uint length = uint(scale[1]);
+  uint count = (length + PART - 1) / PART;
+  uint end = min(uint(length - N + query + 1), (part + 1) * PART);
+  for (uint token = part * PART; token < end; ++token) {
+    uint physical = table[token / BS];
+    uint keyBase = ((physical * HK + hk) * BS + token % BS) * DK;
+    float score = 0;
+    for (uint d = lane; d < DK; d += 32)
+      score += float(q[(hq * N + query) * DK + d]) * float(keys[keyBase + d]);
+    score = simd_sum(score) * scale[0];
+    float next = max(maximum, score);
+    float previous = exp(maximum - next), weight = exp(score - next);
+    uint valueBase = ((physical * HK + hk) * BS + token % BS) * DV;
+    for (uint j = 0; j < (DV + 31) / 32; ++j) {
+      uint d = j * 32 + lane;
+      if (d < DV) accum[j] = accum[j] * previous + weight * float(values[valueBase + d]);
+    }
+    total = total * previous + weight; maximum = next;
+  }
+  uint base = ((hq * N + query) * count + part);
+  for (uint j = 0; j < (DV + 31) / 32; ++j) {
+    uint d = j * 32 + lane;
+    if (d < DV) out[base * DV + d] = accum[j];
+  }
+  if (lane == 0) { normal[base * 2] = maximum; normal[base * 2 + 1] = total; }
+` };
+const partial = new MetalKernel(partialSpec);
+const quantizedPartial = new MetalKernel({ ...partialSpec, name: "paged_attention_affine_partial",
+  inputNames: [...partialSpec.inputNames, "ks", "kb", "vs", "vb"],
+  source: partialSpec.source
+    .replace("float(keys[keyBase + d])", "float(InT(float(ks[(keyBase+d)/GROUP]) * float((keys[(keyBase+d)/PACK] >> (((keyBase+d)%PACK)*BITS)) & ((1u<<BITS)-1u)) + float(kb[(keyBase+d)/GROUP])))")
+    .replace("float(values[valueBase + d])", "float(InT(float(vs[(valueBase+d)/GROUP]) * float((values[(valueBase+d)/PACK] >> (((valueBase+d)%PACK)*BITS)) & ((1u<<BITS)-1u)) + float(vb[(valueBase+d)/GROUP])))"),
+});
+const merge = new MetalKernel({ name: "paged_attention_merge", inputNames: ["parts", "normal"], outputNames: ["out"], source: `
+  uint d = thread_position_in_grid.x;
+  uint row = thread_position_in_grid.y;
+  if (d >= DV) return;
+  uint count = normal_shape[2];
+  float maximum = -INFINITY;
+  for (uint part = 0; part < count; ++part) maximum = max(maximum, normal[(row * count + part) * 2]);
+  float sum = 0, value = 0;
+  for (uint part = 0; part < count; ++part) {
+    uint base = row * count + part;
+    float weight = exp(normal[base * 2] - maximum);
+    sum += normal[base * 2 + 1] * weight;
+    value += parts[base * DV + d] * weight;
+  }
+  out[row * DV + d] = OutT(value / sum);
+` });
+
+/** Independent immutable view: later appends or row retirement cannot alter it. */
+export function pagedAttentionView(pool: PagedAttentionStorage, blockTable: number[], length: number,
+  direct: boolean): KvAttentionView {
+  const keys = ops.contiguous(pool.keys), values = ops.contiguous(pool.values);
+  const table = MlxArray.fromInt32(Int32Array.from(blockTable), [blockTable.length]);
+  const bs = pool.blockSize, quantization = pool.quantization;
+  const extra = quantization ? [pool.keyScales!, pool.keyBiases!, pool.valueScales!, pool.valueBiases!].map(a => ops.contiguous(a)) : [];
+  const dkStored = pool.headDim, dvStored = pool.vHeadDim;
+  const dtype = pool.dtype;
+  return {
+    attend(q: MlxArray, scale: number, mask: Mask): MlxArray {
+      const [, hq, n, dk] = q.shape as [number, number, number, number];
+      const hk = keys.shape[1]!, dv = dvStored;
+      // Long prefill and explicit masks use the established attention oracle.
+      // This is shape dispatch within one attention interface, not scheduling.
+      if (!direct || n > 8 || mask.mode === "array") {
+        const gather = (tensor: MlxArray) => {
+          using picked = ops.takeAxis(tensor, table, 0);
+          using transposed = ops.transposeAxes(picked, [1, 0, 2, 3]);
+          using flat = ops.reshape(transposed, [1, hk, blockTable.length * bs, tensor.shape[3]!]);
+          return flat.slice([0, 0, 0, 0], [1, hk, length, tensor.shape[3]!]);
+        };
+        using k = gather(keys), v = gather(values);
+        if (!quantization) return ops.sdpa(q, k, v, scale, mask.mode, mask.arr);
+        using ks = gather(extra[0]!), kb = gather(extra[1]!), vs = gather(extra[2]!), vb = gather(extra[3]!);
+        using kd = ops.dequantize(k, ks, kb, { ...quantization, mode: "affine" });
+        using vd = ops.dequantize(v, vs, vb, { ...quantization, mode: "affine" });
+        return ops.sdpa(q, kd, vd, scale, mask.mode, mask.arr);
+      }
+      const part = 64, parts = Math.ceil(length / part);
+      using multiplier = MlxArray.fromFloat32(new Float32Array([scale, length]), [2]);
+      const [out, normal] = (quantization ? quantizedPartial : partial).apply([q, keys, values, table, multiplier, ...extra], {
+        outputs: [{ shape: [hq, n, parts, dv], dtype: Dtype.float32 }, { shape: [hq, n, parts, 2], dtype: Dtype.float32 }],
+        grid: [32, hq * n, parts], threadGroup: [32, 1, 1],
+        templateDtypes: quantization ? { InT: dtype } : {},
+        templateInts: { HQ: hq, HK: hk, N: n, DK: dkStored, DV: dv, BS: bs, PART: part,
+          ...(quantization ? { BITS: quantization.bits, GROUP: quantization.groupSize, PACK: 32 / quantization.bits } : {}) },
+      });
+      try {
+        return merge.apply([out!, normal!], { outputs: [{ shape: [1, hq, n, dv], dtype: q.dtype }],
+          grid: [Math.ceil(dv / 32) * 32, hq * n, 1], threadGroup: [32, 1, 1],
+          templateDtypes: { OutT: q.dtype }, templateInts: { DV: dv } })[0]!;
+      } finally { disposeResources([out!, normal!]); }
+    },
+    dispose() { disposeResources([keys, values, table, ...extra]); },
+  };
+}