mlx-bun 0.4.0 → 1.0.0

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

@@ -1,151 +0,0 @@
-// 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

@@ -1,144 +0,0 @@
-// 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

@@ -1,158 +0,0 @@
-// 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

@@ -1,136 +0,0 @@
-// 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

@@ -1,62 +0,0 @@
-// 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]!;
-}