mlx-bun 0.2.0 → 0.4.0

package/src/model/trellis-shared-m.ts

@@ -0,0 +1,151 @@
+// Experimental small-M trellis kernels (variant 7).
+// Decode each weight once per SIMD group and reuse it across M input vectors.
+// Callers restrict M to 2..4 and pass variant 6's decoder/rounding constants.
+// Each vector keeps the original two-accumulator FMA order and simd_sum;
+// fused gate/up also keeps the original bf16 activation rounding boundaries.
+// Baseline implementation and the shared 1MAD header: trellis-linear.ts.
+// Full-model paired performance and quality gates are required before a default.
+
+export 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 += 128u) 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);
+  }
+`;
+
+export const GATEUP_SHARED_M_SOURCE = String.raw`
+  threadgroup float lutTG[4096];
+  if ((VARIANT) == 2) {
+    for (uint i = thread_position_in_threadgroup.x; i < 4096u; i += 128u) 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 ulong rowWords = (ulong)((C / BT) * wpb);
+  const device uint32_t* gRow = gcodes + (ulong)row * rowWords;
+  const device uint32_t* uRow = ucodes + (ulong)row * rowWords;
+  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 * wpb;
+    const device uint32_t* uw = uRow + blk * wpb;
+    const uint w0 = ((uint)BT - 32u - t0) * (uint)K / 32u;
+    uint g[K + 1];
+    uint u[K + 1];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)K; ++i) {
+      const uint wi = w0 + i;
+      const uint idx = wi >= wpb ? wi - wpb : wi;
+      g[i] = gw[idx];
+      u[i] = uw[idx];
+    }
+    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 b = (31u - j) * (uint)K;
+        const uint wi = b >> 5;
+        const uint off = b & 31u;
+        uint gwin = g[wi] >> off;
+        uint uwin = u[wi] >> off;
+        if (off + (uint)L > 32u) { gwin |= g[wi + 1] << (32u - off); uwin |= u[wi + 1] << (32u - off); }
+        gwin &= (1u << (uint)L) - 1u;
+        uwin &= (1u << (uint)L) - 1u;
+        const float gv = TRELLIS_ROUND(TRELLIS_DECODE(gwin, lutTG, lut) * gScale);
+        const float uv = TRELLIS_ROUND(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);
+      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));
+    }
+  }
+`;

package/src/model/trellis-shared-scatter.ts

@@ -0,0 +1,144 @@
+// Experimental variant 10: reuse each decoded scatter weight across M=2..4.
+// Both kernels retain each output's ordered row accumulation, split boundaries
+// and the caller's partial reduction/cast. M=1 keeps the incumbent kernel.
+// The balanced path requires k=3 / T=256 / L<=12; the generic path preserves
+// the existing packed-word lane mapping for other supported geometries.
+
+export const SCATTER_SHARED_M_SOURCE = String.raw`
+  threadgroup float lutTG[4096];
+  if ((VARIANT) == 2) {
+    for (uint i = thread_position_in_threadgroup.x; i < 4096u; i += 128u) 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 wpb = (uint)(BT * K / 32);
+  const uint nb = 32u / wpb;                       // blocks per SIMD group
+  const uint nBlocks = (uint)C / (uint)BT;
+  const uint groups = (nBlocks + nb - 1u) / nb;
+  const uint group = thread_position_in_grid.y * (uint)SG_TG + sg;
+  const uint zi = thread_position_in_grid.z;
+  const uint sample = 0u;
+  const uint split = zi;
+  const uint bi = lane / wpb;
+  const uint wi = lane - bi * wpb;
+  const uint blk = group * nb + bi;
+  const bool active = (group < groups) && (bi < nb) && (blk < nBlocks) && (sample < (uint)M);
+  const uint nbrLane = bi * wpb + ((wi + 1u == wpb) ? 0u : wi + 1u);
+  // This lane's positions: symbol offsets p = multiples of K inside [32wi, 32wi+32).
+  const uint pBase = 32u * wi;
+  const uint p0 = ((pBase + (uint)K - 1u) / (uint)K) * (uint)K;
+  uint offs[NP];
+  bool valid[NP];
+  bool spill[NP];
+  uint cols[NP];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < (uint)NP; ++i) {
+    const uint p = p0 + i * (uint)K;
+    valid[i] = active && (p < pBase + 32u);
+    offs[i] = p - pBase;
+    spill[i] = offs[i] + (uint)L > 32u;
+    const uint t = (uint)BT - 1u - p / (uint)K;
+    cols[i] = blk * (uint)BT + t;
+  }
+  const ulong rowWords = (ulong)(INTERLEAVE ? 2 : nBlocks) * wpb;
+  const uint rowsPer = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+
+  const device uint32_t* col0 = codes + (INTERLEAVE
+    ? (ulong)(blk / 2) * R * 2 * wpb + (blk % 2) * wpb : (ulong)blk * wpb) + wi;
+  const uint mask = (1u << (uint)L) - 1u;
+  float acc[M][NP];
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m)
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)NP; ++i) acc[m][i] = 0.0f;
+  for (uint r = r0; r < r1; ++r) {
+    const uint word = active ? col0[(ulong)r * rowWords] : 0u;
+    const uint nxt = metal::simd_shuffle(word, (ushort)nbrLane);
+    float xr[M];
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) xr[m] = float(x[(ulong)m * (ulong)R + r]);
+    const float sr = float(scales[r]);
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)NP; ++i) {
+      uint win = word >> offs[i];
+      if (spill[i]) win |= nxt << (32u - offs[i]);
+      win &= mask;
+      const float weight = TRELLIS_ROUND(TRELLIS_DECODE(win, lutTG, lut) * sr);
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m) acc[m][i] = metal::fma(weight, xr[m], acc[m][i]);
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    const ulong outBase = ((ulong)m * (ulong)SPLITS + split) * (ulong)C;
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)NP; ++i)
+      if (valid[i]) partial[outBase + cols[i]] = acc[m][i];
+  }
+`;
+
+export const BALANCED_SCATTER_SHARED_M_SOURCE = String.raw`
+  // Model-owned tuning selects this table for measured small-M shapes.
+  // Store integers so the reciprocal multiply retains its original rounding.
+  threadgroup short codebookY[CODEBOOK ? 4096 : 1];
+  if (CODEBOOK) {
+    for (uint i = thread_position_in_threadgroup.x; i < 4096u; i += threads_per_threadgroup.x)
+      codebookY[i] = (short)trellis_y(i);
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  }
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint block = thread_position_in_grid.y * (uint)SG_TG + sg;
+  const uint nBlocks = (uint)C / (uint)BT;
+  const uint zi = thread_position_in_grid.z;
+  const uint sample = 0u;
+  const uint split = zi;
+  if (block >= nBlocks || sample >= (uint)M) return;
+  const uint wpb = (uint)(BT * K / 32);
+  const ulong rowWords = (ulong)(INTERLEAVE ? 2 : nBlocks) * wpb;
+  const uint pBase = lane * 8u * (uint)K;
+  const uint wi = pBase >> 5u;
+  const uint off = pBase & 31u;
+  const uint next = (wi + 1u == wpb) ? 0u : wi + 1u;
+  const device uint32_t* words = codes + (INTERLEAVE
+    ? (ulong)(block / 2) * R * 2 * wpb + (block % 2) * wpb : (ulong)block * wpb);
+  const uint rowsPer = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+  float acc[M][8];
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m)
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) acc[m][i] = 0.0f;
+  for (uint r = r0; r < r1; ++r) {
+    const uint word = words[(ulong)r * rowWords + wi];
+    const uint nxt = words[(ulong)r * rowWords + next];
+    const uint lo = off == 0u ? word : (word >> off) | (nxt << (32u - off));
+    const uint hi = off == 0u ? nxt : nxt >> off;
+    float xr[M];
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) xr[m] = float(x[(ulong)m * (ulong)R + r]);
+    const float sr = float(scales[r]);
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) {
+      const uint shift = i * (uint)K;
+      uint win = lo >> shift;
+      if (shift + (uint)L > 32u) win |= hi << (32u - shift);
+      win &= (1u << (uint)L) - 1u;
+      const float value = (CODEBOOK ? (float)codebookY[win] : (float)trellis_y(win)) * (1.0f / 147.800537109375f);
+      const float weight = value * sr;
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m) acc[m][i] = metal::fma(weight, xr[m], acc[m][i]);
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    const ulong outBase = ((ulong)m * (ulong)SPLITS + split) * (ulong)C;
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i)
+      partial[outBase + block * (uint)BT + (uint)BT - 1u - lane * 8u - i] = acc[m][i];
+  }
+`;

package/src/model/trellis-splitk-prefill.ts

@@ -0,0 +1,158 @@
+// Packed axis-0 prefill with in-tile weight reconstruction and float partials.
+// Partitioning, SIMD fragment layout and ordered accumulation follow MLX 0.31.2
+// (MIT, Apple Inc.): mlx/backend/metal/matmul.cpp and steel/gemm/kernels/
+// steel_gemm_splitk.h. The shared operation preserves that reduction contract.
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import { MetalKernel } from "../mlx/metal-kernel";
+import type { TrellisGeometry } from "./trellis-linear";
+
+export function splitKTrellisPrefillEligible(g: TrellisGeometry, m: number, dtype: Dtype): boolean {
+  return dtype === Dtype.bfloat16 && m >= 5 && m <= 8 &&
+    g.axis === 0 && g.T === 256 && g.L === 12 && [2, 3, 4].includes(g.k) &&
+    g.inFeatures === 17408 && g.outFeatures === 5120;
+}
+
+let partialKernel: MetalKernel | undefined;
+let reductionKernel: MetalKernel | undefined;
+function kernels(): [MetalKernel, MetalKernel] {
+  partialKernel ??= new MetalKernel({ name: "mlx_bun_trellis_splitk_prefill",
+    inputNames: ["x", "codes", "scales"], outputNames: ["out"],
+    header: "#include <metal_simdgroup_matrix>\n", 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;
+    // The eligible M=5..8 path needs one row of 8x8 SIMD-group tiles.
+    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;
+    const uint split=threadgroup_position_in_grid.z;
+    const uint k0=split*PARTITION;
+    const uint k1=split+1==SPLITS?KC:k0+PARTITION;
+    for(uint kb=k0;kb<k1;kb+=BK) {
+      for(uint a=tid;a<BM*BK;a+=128) {
+        const uint m=mbase+a/BK, k=kb+a%BK;
+        As[(a/BK)*BK+a%BK]=(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=k;
+          const uint col=n;
+          const uint block=col/BT, t=col%BT;
+          const device uint* bw=codes+(INTERLEAVE
+            ? (ulong)(block/2)*KC*2*wpb+(ulong)row*2*wpb+(block%2)*wpb
+            : (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;
+          const uint z=win*34038481u+76625530u;
+          const uint p=(z&0x00FF00FFu)+((z>>8)&0x00FF00FFu);
+          const int code=int((p&0xFFFFu)+(p>>16))-510;
+          // Expansion consumes the precise host LUT even for variant 6.
+          // Refine the reciprocal product before scaling and bf16 rounding.
+          const float reciprocal=1.0f/147.800537109375f;
+          const float approx=float(code)*reciprocal;
+          const float residual=metal::fma(-approx,147.800537109375f,float(code));
+          const float value=metal::fma(residual,reciprocal,approx);
+          w=T(value*float(scales[row]));
+        }
+        Bs[(b/BN)*BN+b%BN]=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)split*M*N+(ulong)m*N+n]=acc[i][j].thread_elements()[0];
+        if(m<M && n+1<N)out[(ulong)split*M*N+(ulong)m*N+n+1]=acc[i][j].thread_elements()[1];
+      }
+    }
+  ` });
+  reductionKernel ??= new MetalKernel({ name: "mlx_bun_trellis_splitk_accumulate",
+    inputNames: ["partials"], outputNames: ["out"], source: String.raw`
+      const uint i = thread_position_in_grid.x;
+      if (i >= SIZE) return;
+      float total = 0.0f;
+      for (uint s = 0; s < SPLITS; s++) total += partials[(ulong)s * SIZE + i];
+      out[i] = T(total);
+    ` });
+  return [partialKernel, reductionKernel];
+}
+
+/** Borrow bf16 [M,K] and axis-0 1MAD weights; return owned lazy [M,N].
+ * M=5..8, T=256/L=12, k=2/3/4. Match MLX's non-NAX split-K matrix regime;
+ * other shapes and hardware require their own dispatch and full-model gates.
+ * Peak workspace is SPLITS*M*N float elements, with no dense weight buffer. */
+export function splitKTrellisPrefill(
+  x: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry,
+): MlxArray {
+  const [m, k] = x.shape, n = g.outFeatures;
+  if (x.ndim !== 2 || x.dtype !== Dtype.bfloat16 || !m || m < 5 || m > 8 ||
+      k !== g.inFeatures || g.axis !== 0 || g.T !== 256 || g.L !== 12 ||
+      ![2, 3, 4].includes(g.k) || g.cols % 256 !== 0 || n < 1 ||
+      k < 128 || k < n || Math.ceil(m / 16) * Math.ceil(n / 16) > 1024)
+    throw new Error("unsupported split-K trellis prefill geometry or dtype");
+  const tilesM = Math.ceil(m / 32), tilesN = Math.ceil(n / 32), tilesK = Math.floor(k / 16);
+  const ratio = Math.max(1, Math.floor(tilesK / (tilesM * tilesN)));
+  const splits = Math.min(32, Math.max(2, 2 ** Math.ceil(Math.log2(ratio))));
+  const partition = Math.floor(tilesK / splits) * 16;
+  const [partial, reduce] = kernels();
+  const [workspace] = partial.apply([x, codes, scales], {
+    outputs: [{ shape: [splits, m, n], dtype: Dtype.float32 }],
+    grid: [Math.ceil(n / 32) * 128, 1, splits], threadGroup: [128, 1, 1],
+    templateDtypes: { T: x.dtype }, templateInts: { M: m, N: n, KC: k, C: g.cols,
+      BITS: g.k, L: g.L, BT: g.T, BM: 8, BN: 32, BK: 16, SPLITS: splits,
+      PARTITION: partition, INTERLEAVE: g.blockInterleave ?? 0 },
+  });
+  try {
+    return reduce.apply([workspace!], {
+      outputs: [{ shape: [m, n], dtype: x.dtype }],
+      grid: [m * n, 1, 1], threadGroup: [256, 1, 1], templateDtypes: { T: x.dtype },
+      templateInts: { SIZE: m * n, SPLITS: splits },
+    })[0]!;
+  } finally { workspace!.dispose(); }
+}

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

@@ -0,0 +1,136 @@
+// 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/array";
+import { Dtype } from "../mlx/ffi";
+import { MetalKernel } from "../mlx/metal-kernel";
+import type { TrellisGeometry } from "./trellis-linear";
+import { nativeTrellisWidePrefill } from "./trellis-wide-prefill";
+
+/** Dispatch profile measured on Qwen's MLP geometry. No model-name check. */
+export function tiledTrellisPrefillEligible(g: TrellisGeometry, m: number, dtype: Dtype): boolean {
+  return dtype === Dtype.bfloat16 && m >= 5 && m <= 32 && !nativeTrellisWidePrefill(m) &&
+    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",
+    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;
+          const uint z=win*34038481u+76625530u;
+          const uint p=(z&0x00FF00FFu)+((z>>8)&0x00FF00FFu);
+          const int code=int((p&0xFFFFu)+(p>>16))-510;
+          // Expansion consumes the precise host LUT even for variant 6.
+          // Refine the reciprocal product before scaling and bf16 rounding.
+          const float reciprocal=1.0f/147.800537109375f;
+          const float approx=float(code)*reciprocal;
+          const float residual=metal::fma(-approx,147.800537109375f,float(code));
+          const float value=metal::fma(residual,reciprocal,approx);
+          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/model/trellis-vector-expand.ts

@@ -0,0 +1,62 @@
+// 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/array";
+import { Dtype } from "../mlx/ffi";
+import { MetalKernel } from "../mlx/metal-kernel";
+import type { TrellisGeometry } from "./trellis-linear";
+
+const kernel = new MetalKernel({
+  name: "mlx_bun_trellis_vector_expand",
+  inputNames: ["codes", "scales"], outputNames: ["out"],
+  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);
+      const uint z = s * 34038481u + 76625530u;
+      const uint p = (z & 0x00FF00FFu) + ((z >> 8) & 0x00FF00FFu);
+      const float y = float(int((p & 0xFFFFu) + (p >> 16)) - 510);
+      // Residual refinement reproduces the host f32 LUT before scaling.
+      const float reciprocal = 1.0f / 147.800537109375f;
+      const float q = y * reciprocal;
+      const float value = metal::fma(metal::fma(-q, 147.800537109375f, y), reciprocal, q);
+      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]!;
+}