@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/trellis/mixed-gate-up.ts

@@ -0,0 +1,146 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import type { TrellisWeights } from "./geometry";
+import { HEADER, lutFor, decoderVariant } from "./codebook";
+import { TRELLIS_THREADS, TRELLIS_SG_PER_TG } from "./launch";
+
+import { TRELLIS_MATVEC_MAX_M as MATVEC_MAX_M } from "./reduce";
+
+const MIXED_GATEUP_SOURCE = String.raw`
+  threadgroup float lutTG[4096];
+  if ((VARIANT) == 2) {
+    for (uint i = thread_position_in_threadgroup.x; i < 4096u; i += ${TRELLIS_THREADS}u) lutTG[i] = lut[i];
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  }
+  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;
+  if (row >= (uint)R) return;
+  const uint gWpb = (uint)(BT * KG / 32);
+  const uint uWpb = (uint)(BT * KU / 32);
+  const uint runsPerBlock = (uint)BT / 32u;
+  const uint nRuns = (uint)C / 32u;
+  const device uint32_t* gRow = gcodes + (ulong)row * (ulong)((C / BT) * gWpb);
+  const device uint32_t* uRow = ucodes + (ulong)row * (ulong)((C / BT) * uWpb);
+  const float gScale = float(gscales[row]);
+  const float uScale = float(uscales[row]);
+  float gAcc[M], gAcc2[M], uAcc[M], uAcc2[M];
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    gAcc[m] = 0.0f; gAcc2[m] = 0.0f; uAcc[m] = 0.0f; uAcc2[m] = 0.0f;
+  }
+  for (uint run = lane; run < nRuns; run += 32u) {
+    const uint blk = run / runsPerBlock;
+    const uint t0 = (run - blk * runsPerBlock) * 32u;
+    const device uint32_t* gw = gRow + blk * gWpb;
+    const device uint32_t* uw = uRow + blk * uWpb;
+    const uint gW0 = ((uint)BT - 32u - t0) * (uint)KG / 32u;
+    const uint uW0 = ((uint)BT - 32u - t0) * (uint)KU / 32u;
+    uint g[KG + 1];
+    uint u[KU + 1];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)KG; ++i) {
+      const uint wi = gW0 + i;
+      g[i] = gw[wi >= gWpb ? wi - gWpb : wi];
+    }
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)KU; ++i) {
+      const uint wi = uW0 + i;
+      u[i] = uw[wi >= uWpb ? wi - uWpb : wi];
+    }
+    const uint c0 = blk * (uint)BT + t0;
+    #pragma clang loop unroll(full)
+    for (uint j4 = 0; j4 < 8u; ++j4) {
+      float4 xv4[M];
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m)
+        xv4[m] = float4(((const device metal::vec<T, 4>*)(x + (ulong)m * (ulong)C + c0))[j4]);
+      #pragma clang loop unroll(full)
+      for (uint q = 0; q < 4u; ++q) {
+        const uint j = j4 * 4u + q;
+        const uint gb = (31u - j) * (uint)KG;
+        const uint ub = (31u - j) * (uint)KU;
+        const uint gwi = gb >> 5, goff = gb & 31u;
+        const uint uwi = ub >> 5, uoff = ub & 31u;
+        uint gwin = g[gwi] >> goff;
+        uint uwin = u[uwi] >> uoff;
+        if (goff + (uint)L > 32u) gwin |= g[gwi + 1] << (32u - goff);
+        if (uoff + (uint)L > 32u) uwin |= u[uwi + 1] << (32u - uoff);
+        gwin &= (1u << (uint)L) - 1u;
+        uwin &= (1u << (uint)L) - 1u;
+        const float gv = TRELLIS_ROUND((BITS_Y ? trellis_unrefined_bits_y(trellis_y(gwin)) : TRELLIS_DECODE(gwin, lutTG, lut)) * gScale);
+        const float uv = TRELLIS_ROUND((BITS_Y ? trellis_unrefined_bits_y(trellis_y(uwin)) : TRELLIS_DECODE(uwin, lutTG, lut)) * uScale);
+        #pragma clang loop unroll(full)
+        for (uint m = 0; m < (uint)M; ++m) {
+          const float xv = xv4[m][q];
+          if (q & 1u) { gAcc2[m] = metal::fma(gv, xv, gAcc2[m]); uAcc2[m] = metal::fma(uv, xv, uAcc2[m]); }
+          else { gAcc[m] = metal::fma(gv, xv, gAcc[m]); uAcc[m] = metal::fma(uv, xv, uAcc[m]); }
+        }
+      }
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    const float gate = metal::simd_sum(gAcc[m] + gAcc2[m]);
+    const float up = metal::simd_sum(uAcc[m] + uAcc2[m]);
+    if (lane == 0u) {
+      const T gateT = T(gate);
+      const T upT = T(up);
+      if ((TAIL) == 1) {
+        const T sigY = 1 / (1 + metal::exp(metal::abs(gateT)));
+        const T sigmoidT = (gateT < 0) ? sigY : 1 - sigY;
+        const T siluT = gateT * sigmoidT;
+        mid[(ulong)m * (ulong)R + row] = siluT * upT;
+      } else {
+        const T sigmoidT = T(1.0f / (1.0f + metal::precise::exp(-float(gateT))));
+        const T siluT = T(float(gateT) * float(sigmoidT));
+        mid[(ulong)m * (ulong)R + row] = T(float(siluT) * float(upT));
+      }
+    }
+  }
+`;
+
+let mixedGateUp: MetalKernel | null = null;
+function mixedGateUpKernel(): MetalKernel {
+  return mixedGateUp ??= new MetalKernel({
+    name: "mlx_bun_trellis_gateup_mixed_k",
+    inputNames: ["x", "gcodes", "gscales", "ucodes", "uscales", "lut"],
+    outputNames: ["mid"], source: MIXED_GATEUP_SOURCE, header: HEADER, ensureRowContiguous: true,
+  });
+}
+/** Which SwiGLU arithmetic the mixed kernel reproduces: `"split"` is MLX's compiled
+ *  swiglu over the two bf16 projections (bit-identical to gate.forward + up.forward +
+ *  compiledSwiglu); `"fused"` is fusedGateUpSwiglu's float32 sigmoid. */
+export type MixedGateUpTail = "split" | "fused";
+
+/** silu(gate(x)) * up(x) for M <= 4 rows in one kernel when gate and up are coded
+ *  at different k. Same lanes, accumulation order and bf16 SwiGLU boundaries as
+ *  fusedGateUpSwiglu; each projection keeps its own word layout. `exactBits`
+ *  applies at the qualified single-row bf16 geometry, where the decoder builds
+ *  its value from exact float32 bits: layers pass the choice they resolved at
+ *  construction (`qualifiedRepresentations`), benches and tests the one they compare. */
+export function fusedGateUpSwigluMixed(x: MlxArray, gate: TrellisWeights, up: TrellisWeights, selected: number,
+  tail: MixedGateUpTail, exactBits: boolean): MlxArray {
+  const g = gate.geometry, u = up.geometry;
+  const lead = x.shape.slice(0, -1);
+  const M = lead.reduce((a, b) => a * b, 1);
+  if (M > MATVEC_MAX_M) throw new Error(`fusedGateUpSwigluMixed: M=${M} > ${MATVEC_MAX_M}`);
+  const bitsY = selected === 13 && M === 1 && x.dtype === Dtype.bfloat16 &&
+    g.T === 256 && g.L === 12 && g.rows === 17408 && g.cols === 5120 &&
+    [2, 3, 4].includes(g.k) && [2, 3, 4].includes(u.k) && exactBits;
+  const x2 = ops.reshape(x, [M, g.inFeatures]);
+  const [mid] = mixedGateUpKernel().apply([x2, gate.codes, gate.scales, up.codes, up.scales, lutFor(g.L)], {
+    outputs: [{ shape: [M, g.rows], dtype: x.dtype }],
+    grid: [TRELLIS_THREADS, Math.ceil(g.rows / TRELLIS_SG_PER_TG), 1],
+    threadGroup: [TRELLIS_THREADS, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M, R: g.rows, C: g.cols, BT: g.T, KG: g.k, KU: u.k, L: g.L, ROWS_TG: TRELLIS_SG_PER_TG,
+      VARIANT: decoderVariant(selected), TAIL: tail === "split" ? 1 : 0, BITS_Y: Number(bitsY) },
+  });
+  x2.dispose();
+  const out = ops.reshape(mid!, [...lead, g.rows]);
+  mid!.dispose();
+  return out;
+}

package/src/kernels/trellis/qualified.ts

@@ -0,0 +1,34 @@
+/** Single-row down representations at the qualified geometry (`scatterBitsEligible`). */
+export interface ScatterRepresentation {
+  /** Balanced 3-bit decoder value from exact float32 bits. */
+  bits: boolean;
+  /** The same for 2/4-bit codes. */
+  genericBits: boolean;
+  /** Generated float32 threadgroup codebook. */
+  floatCodebook: boolean;
+}
+
+/** What the Trellis kernels take from the GPU family, as explicit arguments.
+ *  The single-row representations were measured on the Qwen MLP geometry
+ *  (bf16, variant 13); each keeps output bits and is a speed choice. */
+export interface TrellisRepresentations {
+  /** Threadgroup codebook in the same-width gate/up decode. */
+  gateUpCodebook: boolean;
+  /** Exact float32 bits in the mixed-width gate/up decoder. */
+  mixedExactBits: boolean;
+  scatter: ScatterRepresentation;
+  /** The bundled runtime reduces M=5..15 matmuls with GemvWide on this GPU, so
+   *  wide prefill matches it there and tiled prefill steps aside. */
+  widePrefill: boolean;
+}
+
+/** The choices qualified on `architecture`: the single-row representations on
+ *  `applegpu_g13s` only, wide prefill on GPU family 15 (M3) and newer. A layer
+ *  resolves this once, when it is constructed, and passes the stored result to
+ *  the kernels on every call. */
+export function qualifiedRepresentations(architecture: string): TrellisRepresentations {
+  const qualified = architecture === "applegpu_g13s";
+  return { gateUpCodebook: qualified, mixedExactBits: qualified,
+    scatter: { bits: qualified, genericBits: qualified, floatCodebook: qualified },
+    widePrefill: Number(/^applegpu_[a-z](\d{2})/.exec(architecture)?.[1] ?? 0) >= 15 };
+}

package/src/kernels/trellis/reduce.ts

@@ -0,0 +1,150 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import type { TrellisGeometry } from "./geometry";
+import { HEADER, lutFor, decoderVariant } from "./codebook";
+import { TRELLIS_THREADS, TRELLIS_SG_PER_TG } from "./launch";
+
+export const TRELLIS_MATVEC_MAX_M = 4;
+
+const REDUCE_SOURCE = String.raw`
+  threadgroup float lutTG[4096];
+  if ((VARIANT) == 2) {
+    for (uint i = thread_position_in_threadgroup.x; i < 4096u; i += ${TRELLIS_THREADS}u) lutTG[i] = lut[i];
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  }
+  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)R || sample >= (uint)M) return;
+  const uint wpb = (uint)(BT * K / 32);
+  const uint runsPerBlock = (uint)BT / 32u;
+  const uint nRuns = (uint)C / 32u;
+  const device uint32_t* codeRow = codes + (ulong)row * (ulong)((C / BT) * wpb);
+  const device T* xs = x + (ulong)sample * (ulong)C;
+  const float scale = float(scales[row]);
+  float acc = 0.0f, acc2 = 0.0f;
+  for (uint run = lane; run < nRuns; run += 32u) {
+    const uint blk = run / runsPerBlock;
+    const uint t0 = (run - blk * runsPerBlock) * 32u;
+    const device uint32_t* bw = codeRow + blk * wpb;
+    const uint w0 = ((uint)BT - 32u - t0) * (uint)K / 32u;
+    uint w[K + 1];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)K; ++i) {
+      const uint wi = w0 + i;
+      w[i] = bw[wi >= wpb ? wi - wpb : wi];
+    }
+    const uint c0 = blk * (uint)BT + t0;
+    const device metal::vec<T, 4>* x4 = (const device metal::vec<T, 4>*)(xs + c0);
+    #pragma clang loop unroll(full)
+    for (uint j4 = 0; j4 < 8u; ++j4) {
+      const float4 xv = float4(x4[j4]);
+      #pragma clang loop unroll(full)
+      for (uint q = 0; q < 4u; ++q) {
+        const uint j = j4 * 4u + q;
+        const uint b = (31u - j) * (uint)K;
+        const uint wi = b >> 5;
+        const uint off = b & 31u;
+        uint win = w[wi] >> off;
+        if (off + (uint)L > 32u) win |= w[wi + 1] << (32u - off);
+        win &= (1u << (uint)L) - 1u;
+        const float wv = TRELLIS_ROUND(TRELLIS_DECODE(win, lutTG, lut) * scale);
+        if (q & 1u) acc2 = metal::fma(wv, xv[q], acc2);
+        else acc = metal::fma(wv, xv[q], acc);
+      }
+    }
+  }
+  const float v = metal::simd_sum(acc + acc2);
+  if (lane == 0u) out[(ulong)sample * (ulong)R + row] = T(v);
+`;
+
+const REDUCE_SHARED_M_SOURCE = String.raw`
+  threadgroup float lutTG[4096];
+  if ((VARIANT) == 2) {
+    for (uint i = thread_position_in_threadgroup.x; i < 4096u; i += ${TRELLIS_THREADS}u) lutTG[i] = lut[i];
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  }
+  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 = 0u;
+  if (row >= (uint)R || sample >= (uint)M) return;
+  const uint wpb = (uint)(BT * K / 32);
+  const uint runsPerBlock = (uint)BT / 32u;
+  const uint nRuns = (uint)C / 32u;
+  const device uint32_t* codeRow = codes + (ulong)row * (ulong)((C / BT) * wpb);
+  const float scale = float(scales[row]);
+  float acc[M], acc2[M];
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) { acc[m] = 0.0f; acc2[m] = 0.0f; }
+  for (uint run = lane; run < nRuns; run += 32u) {
+    const uint blk = run / runsPerBlock;
+    const uint t0 = (run - blk * runsPerBlock) * 32u;
+    const device uint32_t* bw = codeRow + blk * wpb;
+    const uint w0 = ((uint)BT - 32u - t0) * (uint)K / 32u;
+    uint w[K + 1];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)K; ++i) {
+      const uint wi = w0 + i;
+      w[i] = bw[wi >= wpb ? wi - wpb : wi];
+    }
+    const uint c0 = blk * (uint)BT + t0;
+    #pragma clang loop unroll(full)
+    for (uint j4 = 0; j4 < 8u; ++j4) {
+      float4 xv[M];
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m)
+        xv[m] = float4(((const device metal::vec<T, 4>*)(x + (ulong)m * (ulong)C + c0))[j4]);
+      #pragma clang loop unroll(full)
+      for (uint q = 0; q < 4u; ++q) {
+        const uint j = j4 * 4u + q;
+        const uint b = (31u - j) * (uint)K;
+        const uint wi = b >> 5;
+        const uint off = b & 31u;
+        uint win = w[wi] >> off;
+        if (off + (uint)L > 32u) win |= w[wi + 1] << (32u - off);
+        win &= (1u << (uint)L) - 1u;
+        const float wv = TRELLIS_ROUND(TRELLIS_DECODE(win, lutTG, lut) * scale);
+        #pragma clang loop unroll(full)
+        for (uint m = 0; m < (uint)M; ++m) {
+          if (q & 1u) acc2[m] = metal::fma(wv, xv[m][q], acc2[m]);
+          else acc[m] = metal::fma(wv, xv[m][q], acc[m]);
+        }
+      }
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    const float v = metal::simd_sum(acc[m] + acc2[m]);
+    if (lane == 0u) out[(ulong)m * (ulong)R + row] = T(v);
+  }
+`;
+
+let base: MetalKernel | undefined, shared: MetalKernel | undefined;
+function reduceKernel(): MetalKernel {
+  return base ??= new MetalKernel({ name: "mlx_bun_trellis_reduce",
+    inputNames: ["x", "codes", "scales", "lut"], outputNames: ["out"],
+    source: REDUCE_SOURCE, header: HEADER, ensureRowContiguous: true });
+}
+function sharedKernel(): MetalKernel {
+  return shared ??= new MetalKernel({ name: "mlx_bun_trellis_reduce_shared_m",
+    inputNames: ["x", "codes", "scales", "lut"], outputNames: ["out"],
+    source: REDUCE_SHARED_M_SOURCE, header: HEADER, ensureRowContiguous: true });
+}
+
+/** Borrow [M, inFeatures] and axis-1 weights; return owned lazy [M, outFeatures].
+ * The caller selects a supported variant and limits M to 1..4. */
+export function trellisReduce(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry, selected: number): MlxArray {
+  const M = x2.shape[0]!;
+  const shared = selected >= 7 && selected <= 13 && M > 1;
+  const kernel = shared ? sharedKernel() : reduceKernel();
+  const [out] = kernel.apply([x2, codes, scales, lutFor(g.L)], {
+    outputs: [{ shape: [M, g.rows], dtype: x2.dtype }],
+    grid: [TRELLIS_THREADS, Math.ceil(g.rows / TRELLIS_SG_PER_TG), shared ? 1 : M],
+    threadGroup: [TRELLIS_THREADS, 1, 1],
+    templateDtypes: { T: x2.dtype },
+    templateInts: { M, R: g.rows, C: g.cols, BT: g.T, K: g.k, L: g.L, ROWS_TG: TRELLIS_SG_PER_TG, VARIANT: decoderVariant(selected) },
+  });
+  return out!;
+}