@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/trellis/tiled-prefill.ts

@@ -0,0 +1,131 @@
+// Packed axis-1 trellis prefill without a full expanded weight allocation.
+// SIMD-group fragment layout follows MLX steel/gemm/mma.h (MIT, Apple Inc.):
+// https://github.com/ml-explore/mlx/blob/v0.31.2/mlx/backend/metal/kernels/steel/gemm/mma.h
+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 { HEADER } from "./codebook";
+import { nativeTrellisWidePrefill } from "./wide-prefill";
+
+/** Dispatch profile measured on Qwen's MLP geometry. No model-name check.
+ *  `widePrefill` as for `wideTrellisPrefillEligible`: tiled steps aside where wide applies. */
+export function tiledTrellisPrefillEligible(g: TrellisGeometry, m: number, dtype: Dtype, widePrefill: boolean): boolean {
+  return dtype === Dtype.bfloat16 && m >= 5 && m <= 32 && !nativeTrellisWidePrefill(m, widePrefill) &&
+    g.axis === 1 && g.T === 256 && g.L === 12 && [2, 3, 4].includes(g.k) &&
+    g.inFeatures === 5120 && g.outFeatures === 17408;
+}
+
+let kernel: MetalKernel | undefined;
+function tiledKernel(): MetalKernel {
+  return kernel ??= new MetalKernel({ name: "mlx_bun_trellis_prefill_tile",
+    inputNames: ["x", "codes", "scales"], outputNames: ["out"],
+    header: "#include <metal_simdgroup_matrix>\n" + HEADER,
+    source: String.raw`
+    const uint tid=thread_position_in_threadgroup.x;
+    const uint lane=thread_index_in_simdgroup;
+    const uint sg=simdgroup_index_in_threadgroup;
+    const uint mbase=threadgroup_position_in_grid.y*BM;
+    const uint nbase=threadgroup_position_in_grid.x*BN;
+    const uint fm=((lane/4)&4)+((lane/2)%4);
+    const uint fn=((lane/4)&2)*2+(lane%2)*2;
+    // An eight-row tile assigns all four SIMD groups to output columns.
+    // Each output keeps the same 8x8 products and K accumulation order.
+    constexpr int WM=BM==8?1:2, WN=4/WM, SM=WM*8, SN=WN*8;
+    const uint sm=(sg/WN)*8+fm, sn=(sg%WN)*8+fn;
+    constexpr int TM=BM/SM, TN=BN/SN;
+    threadgroup T As[BM*BK];
+    threadgroup T Bs[BK*BN];
+    metal::simdgroup_float8x8 acc[TM][TN];
+    #pragma clang loop unroll(full)
+    for(uint i=0;i<TM;i++) {
+      #pragma clang loop unroll(full)
+      for(uint j=0;j<TN;j++) {
+        acc[i][j].thread_elements()[0]=0.0f;
+        acc[i][j].thread_elements()[1]=0.0f;
+      }
+    }
+    constexpr uint wpb=BT*BITS/32;
+    for(uint kb=0;kb<KC;kb+=BK) {
+      for(uint a=tid;a<BM*BK;a+=128) {
+        const uint m=mbase+a/BK, k=kb+a%BK;
+        As[a]=(m<M && k<KC)?x[(ulong)m*KC+k]:T(0);
+      }
+      for(uint b=tid;b<BK*BN;b+=128) {
+        const uint k=kb+b/BN, n=nbase+b%BN;
+        T w=T(0);
+        if(k<KC && n<N) {
+          const uint row=AXIS==1?n:k;
+          const uint col=AXIS==1?k:n;
+          const uint block=col/BT, t=col%BT;
+          const device uint* bw=codes+(ulong)row*(C/BT)*wpb+block*wpb;
+          const uint pos=(BT-1-t)*BITS, wi=pos>>5, off=pos&31u;
+          uint win=bw[wi]>>off;
+          if(off+L>32)win|=bw[wi+1==wpb?0:wi+1]<<(32-off);
+          win&=(1u<<L)-1;
+          // Expansion consumes the precise host LUT even for variant 6.
+          const float value=trellis_val_rcp(win);
+          w=T(value*float(scales[row]));
+        }
+        Bs[b]=w;
+      }
+      threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+      #pragma clang loop unroll(full)
+      for(uint kk=0;kk<BK;kk+=8) {
+        metal::simdgroup_float8x8 a[TM],b[TN];
+        #pragma clang loop unroll(full)
+        for(uint i=0;i<TM;i++) {
+          a[i].thread_elements()[0]=float(As[(sm+i*SM)*BK+kk+fn]);
+          a[i].thread_elements()[1]=float(As[(sm+i*SM)*BK+kk+fn+1]);
+        }
+        #pragma clang loop unroll(full)
+        for(uint j=0;j<TN;j++) {
+          b[j].thread_elements()[0]=float(Bs[(kk+fm)*BN+sn+j*SN]);
+          b[j].thread_elements()[1]=float(Bs[(kk+fm)*BN+sn+j*SN+1]);
+        }
+        #pragma clang loop unroll(full)
+        for(uint i=0;i<TM;i++) {
+          #pragma clang loop unroll(full)
+          for(uint j=0;j<TN;j++) {
+            const uint nj=(i%2)?TN-1-j:j;
+            metal::simdgroup_multiply_accumulate(acc[i][nj],a[i],b[nj],acc[i][nj]);
+          }
+        }
+      }
+      threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    }
+    #pragma clang loop unroll(full)
+    for(uint i=0;i<TM;i++) {
+      #pragma clang loop unroll(full)
+      for(uint j=0;j<TN;j++) {
+        const uint m=mbase+sm+i*SM,n=nbase+sn+j*SN;
+        if(m<M && n<N)out[(ulong)m*N+n]=T(acc[i][j].thread_elements()[0]);
+        if(m<M && n+1<N)out[(ulong)m*N+n+1]=T(acc[i][j].thread_elements()[1]);
+      }
+    }
+  ` });
+}
+
+/** Shared operation for bf16 [M, K], axis-1 1MAD, T=256/L=12, k=2/3/4.
+ * M=5..32. Borrow inputs and return owned [M, N]; caller controls evaluation.
+ * Reconstruct the precise LUT value before bf16 weight rounding. The default
+ * packed matvec's unrefined f32 code value is a different numerical contract.
+ * Other geometries need a measured dispatch profile and model-level gates. */
+export function tiledTrellisPrefill(
+  x: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry,
+): MlxArray {
+  const [m, k] = x.shape;
+  if (x.ndim !== 2 || x.dtype !== Dtype.bfloat16 || !m || m < 5 || m > 32 ||
+      k !== g.inFeatures || g.axis !== 1 || g.T !== 256 || g.L !== 12 ||
+      ![2, 3, 4].includes(g.k) || g.cols % 256 !== 0 || g.rows < 1)
+    throw new Error("unsupported tiled trellis prefill geometry or dtype");
+  const bm = m <= 8 ? 8 : m <= 16 ? 16 : 32, bn = 64, bk = 16;
+  const [out] = tiledKernel().apply([x, codes, scales], {
+    outputs: [{ shape: [m, g.outFeatures], dtype: x.dtype }],
+    grid: [Math.ceil(g.outFeatures / bn) * 128, Math.ceil(m / bm), 1],
+    threadGroup: [128, 1, 1], templateDtypes: { T: x.dtype },
+    templateInts: { M: m, N: g.outFeatures, KC: g.inFeatures, C: g.cols,
+      BITS: g.k, L: g.L, BT: g.T, AXIS: g.axis, BM: bm, BN: bn, BK: bk },
+  });
+  return out!;
+}

package/src/kernels/trellis/vector-expand.ts

@@ -0,0 +1,58 @@
+// Four adjacent weights share one circular packed-word window. Keep the
+// precise expansion values and MLX's matmul/evaluation boundary unchanged.
+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 { HEADER } from "./codebook";
+
+const kernel = new MetalKernel({
+  name: "mlx_bun_trellis_vector_expand",
+  inputNames: ["codes", "scales"], outputNames: ["out"], header: HEADER,
+  source: String.raw`
+    uint c = thread_position_in_grid.x * 4;
+    uint r = thread_position_in_grid.y;
+    if constexpr (ROW_TILE > 0) {
+      constexpr uint groups = C / (BT * 2);
+      const uint tg = threadgroup_position_in_grid.x;
+      const uint tile = tg / (groups * ROW_TILE), within = tg % (groups * ROW_TILE);
+      r = tile * ROW_TILE + within % ROW_TILE;
+      c = (within / ROW_TILE) * BT * 2 + thread_position_in_threadgroup.x * 4;
+    }
+    if (c >= C || r >= R) return;
+    constexpr uint wpb = BT * BITS / 32;
+    const uint block = c / BT, t = c % BT;
+    const device uint* bw = codes + (INTERLEAVE
+      ? (ulong)(block / 2) * R * 2 * wpb + (ulong)r * 2 * wpb + (block % 2) * wpb
+      : (ulong)r * (C / BT) * wpb + block * wpb);
+    const float scale = float(scales[r]);
+    // The last of four weights starts the window. Earlier weights occupy
+    // its higher bits; a second word also handles the circular block edge.
+    const uint pos = (BT - 1 - (t + 3)) * BITS, wi = pos >> 5, off = pos & 31u;
+    const ulong window = (ulong(bw[wi]) | (ulong(bw[wi + 1 == wpb ? 0 : wi + 1]) << 32)) >> off;
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 4; i++) {
+      const uint s = uint(window >> ((3 - i) * BITS)) & ((1u << L) - 1);
+      // Residual refinement reproduces the host f32 LUT before scaling.
+      const float value = trellis_val_rcp(s);
+      out[(ulong)r * C + c + i] = T(value * scale);
+    }
+  `,
+});
+
+export function vectorTrellisExpandEligible(g: TrellisGeometry, dtype: Dtype): boolean {
+  return dtype === Dtype.bfloat16 && g.L === 12 && g.T === 256 &&
+    (g.k === 2 || g.k === 3 || g.k === 4) && g.rows > 0 && g.cols > 0 && g.cols % 256 === 0;
+}
+
+/** Decode the stored matrix; the caller owns the projection's evaluation boundary. */
+export function vectorTrellisExpand(codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  const rowTile = g.blockInterleave && g.rows % 128 === 0 ? 128 : 0;
+  return kernel.apply([codes, scales], {
+    outputs: [{ shape: [g.rows, g.cols], dtype: Dtype.bfloat16 }],
+    grid: rowTile ? [g.cols / 4 * g.rows, 1, 1] : [g.cols / 4, g.rows, 1], threadGroup: [128, 1, 1],
+    templateDtypes: { T: Dtype.bfloat16 },
+    templateInts: { R: g.rows, C: g.cols, BT: g.T, BITS: g.k, L: g.L,
+      INTERLEAVE: g.blockInterleave ?? 0, ROW_TILE: rowTile },
+  })[0]!;
+}

package/src/kernels/trellis/wide-prefill.ts

@@ -0,0 +1,89 @@
+// Packed Trellis prefill with MLX 0.32.2 GemvWide arithmetic.
+// Ported from ml-explore/mlx v0.32.2, mlx/backend/metal/kernels/gemv.h
+// (MIT). Decode weights to bf16 inside the native dot/unroll/shuffle order;
+// share each decoded row across every input vector in this narrow regime.
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import { type MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { type TrellisGeometry } from "./geometry";
+import { HEADER } from "./codebook";
+
+/** Matches the bundled runtime's native reduction where the GPU family has it
+ *  (`widePrefill` from `qualifiedRepresentations`, resolved by the caller). */
+export function nativeTrellisWidePrefill(m: number, widePrefill: boolean): boolean {
+  return m >= 5 && m <= 15 && widePrefill;
+}
+
+export function wideTrellisPrefillEligible(g: TrellisGeometry, m: number, dtype: Dtype, widePrefill: boolean): boolean {
+  return dtype === Dtype.bfloat16 && nativeTrellisWidePrefill(m, widePrefill) &&
+    g.axis === 1 && g.T === 256 && g.L === 12 && [2, 3, 4].includes(g.k) &&
+    g.inFeatures === 5120 && g.outFeatures === 17408 && !g.blockInterleave;
+}
+
+let kernel: MetalKernel | undefined;
+function wideKernel() {
+  return kernel ??= new MetalKernel({ name: 'mlx_bun_trellis_wide_prefill',
+    inputNames: ['x', 'codes', 'scales'], outputNames: ['out'], header: HEADER, source: String.raw`
+    constexpr uint unroll=8, rows_per_simd=32/KL, simdgroups=KL/8;
+    const uint lane=thread_index_in_simdgroup, sg=simdgroup_index_in_threadgroup;
+    const uint kl=lane%KL, row=threadgroup_position_in_grid.y*4+sg*rows_per_simd+lane/KL;
+    const uint vec0=threadgroup_position_in_grid.x*SM;
+    constexpr uint wpb=BT*BITS/32, row_words=KC*BITS/32;
+    const float scale=float(scales[row]);
+    constexpr uint n_v4=KC/4, n_main=n_v4-n_v4%(KL*unroll);
+    float result[SM]={0};
+    for(uint base=0;base<n_main;base+=KL*unroll) {
+      float4 wf[unroll];
+      #pragma clang loop unroll(full)
+      for(uint i=0;i<unroll;i++) {
+        #pragma clang loop unroll(full)
+        for(uint c=0;c<4;c++) {
+          const uint k=(base+i*KL+kl)*4+c, block=k/BT, t=k%BT;
+          const device uint* bw=codes+(ulong)row*row_words+block*wpb;
+          const uint pos=(BT-1-t)*BITS, wi=pos>>5, off=pos&31u;
+          uint win=bw[wi]>>off;
+          if(off+L>32)win|=bw[wi+1==wpb?0:wi+1]<<(32-off);
+          win&=(1u<<L)-1;
+          const float value=trellis_val_rcp(win);
+          wf[i][c]=float(T(value*scale));
+        }
+      }
+      #pragma clang loop unroll(full)
+      for(uint v=0;v<SM;v++) {
+        const uint m=min(vec0+v,uint(M-1));
+        float acc=0;
+        #pragma clang loop unroll(full)
+        for(uint i=0;i<unroll;i++) {
+          const ulong offset=(ulong)m*KC+(base+i*KL+kl)*4;
+          const float4 xv=float4(float(x[offset]),float(x[offset+1]),float(x[offset+2]),float(x[offset+3]));
+          acc+=metal::dot(wf[i],xv);
+        }
+        result[v]+=acc;
+      }
+    }
+    #pragma clang loop unroll(full)
+    for(uint v=0;v<SM;v++) {
+      #pragma clang loop unroll(full)
+      for(ushort off=KL/2;off>=1;off>>=1) result[v]+=simd_shuffle_down(result[v],off);
+      if(kl==0 && vec0+v<M) out[(ulong)(vec0+v)*N+row]=T(result[v]);
+    }
+  ` });
+}
+/** Requires the caller to prove aligned, row-contiguous x, as Qwen does
+ * with its RMSNorm output. Preserve native fallback when that proof is absent. */
+export function wideTrellisPrefill(x: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  const [m, k] = x.shape;
+  if (x.ndim !== 2 || x.dtype !== Dtype.bfloat16 || !m || m < 5 || m > 15 ||
+      k !== 5120 || g.inFeatures !== k || g.axis !== 1 || g.T !== 256 || g.L !== 12 ||
+      ![2, 3, 4].includes(g.k) || g.blockInterleave || g.rows < 4 || g.rows % 4 ||
+      g.outFeatures !== g.rows || g.cols !== k || codes.dtype !== Dtype.uint32 ||
+      codes.ndim !== 2 || codes.shape[0] !== g.rows || codes.shape[1] !== k * g.k / 32 || scales.size !== g.rows)
+    throw new Error("unsupported wide trellis prefill geometry, dtype or layout");
+  const kl = m <= 5 ? 32 : 16;
+  return wideKernel().apply([x, codes, scales], {
+    outputs: [{ shape: [m, g.outFeatures], dtype: x.dtype }],
+    grid: [32, g.outFeatures / 4 * (kl / 8), 1], threadGroup: [32, kl / 8, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M: m, SM: m, KL: kl, KC: k, N: g.outFeatures, BT: g.T, L: g.L, BITS: g.k },
+  })[0]!;
+}

package/src/kernels/turboquant/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 '@mlx-bun/mlx/array';
+import { Dtype, type MlxHandle } from '@mlx-bun/mlx/ffi';
+import { isShapelessTracing } from '@mlx-bun/mlx/compile';
+import { MetalKernel } from '@mlx-bun/mlx/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/kernels/turboquant/index.ts

@@ -0,0 +1,5 @@
+export * from "./ops";
+export * from "./decode";
+export * from "./inverse";
+export * from "./tables";
+export * from "./signs";

package/src/kernels/turboquant/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 "@mlx-bun/mlx/array";
+import { Dtype, type MlxHandle } from "@mlx-bun/mlx/ffi";
+import { isShapelessTracing } from "@mlx-bun/mlx/compile";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import { type PackedKvArrays } from "./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];
+}