mlx-bun 0.0.13 → 0.3.0

package/src/quantize/trellis.ts

@@ -0,0 +1,586 @@
+// Bitshift trellis codec (QTIP TCQ) — the ENGINE's copy, with the packed
+// bit-stream the Q2b kernel reads. scripts/turboquant/tq-trellis.ts is the
+// frozen fake-quant codec that produced the Q3 artifact of record; this module
+// is a copy of it plus `encodeStates` / `fakeQuantRowsPacked` / the host pack
+// and unpack helpers, so nothing that reproduces the record changes.
+//
+// QTIP-style bitshift trellis-coded quantization (TCQ), fake-quant, on MLX.
+//
+// Reference (verified against source, not memory):
+//   Tseng et al., "QTIP: Quantization with Trellises and Incoherence Processing",
+//   arXiv:2406.11235v3 — Fig. 2 (bitshift trellis), §3.1.1 (computed codes),
+//   §3.2 + Algorithm 4 (tail-biting approximation), Table 2 (the (12,k,1)
+//   tail-biting distortion numbers this file's --validate reproduces).
+//   github.com/Cornell-RelaxML/qtip @ main, lib/codebook/bitshift.py.
+//
+// Bitshift trellis (L, k, V=1). State = L bits. Successor n of state p must
+// satisfy `n >> k == p & (2^(L-k) - 1)` (Fig. 2: "the 2^kV nodes that share
+// their top L-kV bits with its bottom L-kV bits"), i.e.
+//     n = ((p << k) | b) mod 2^L,   b in [0, 2^k)
+// so the predecessors of n are { (n >> k) + m·2^(L-k) : m in [0, 2^k) }.
+// The emitted value for state s is a computed pseudo-Gaussian f(s) — no
+// lookup table on the decode side (we materialize one anyway for the encoder).
+//
+// Encoder = Viterbi over 2^L states, UNWEIGHTED squared error (no Hessian /
+// LDLQ — v2 lever), batched over thousands of 256-weight blocks at a time.
+// Sign convention: we MAXIMIZE g = -(cost), and drop the per-step constant
+// x², so the emission term is a single addmm:  g_err = -lut² + 2·x·lut.
+// Dropping x² shifts every path in a block by the same amount, so argmax
+// decisions and the traceback are unchanged.
+
+import { MlxArray as Arr, gpuStream, type MlxArray } from "../mlx/array";
+import { Dtype, clearCache, activeMemory, cacheMemory, peakMemory } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+
+const S = gpuStream;
+
+// ------------------------------------------------------------- codebooks ---
+
+/** QTIP `decode_1mad` (lib/codebook/bitshift.py) evaluated over [0, 2^L).
+ *  x = (x·34038481 + 76625530) mod 2^32; y = Σ bytes(x) − 510; y / 147.800537109375.
+ *  (510 = 2·255 = E[Σ of 4 uniform bytes]; 147.8005… ≈ √(4·(2^16−1)/12) is the
+ *  std of that sum, so the code is unit-variance by construction.) */
+export function lut1mad(L: number): Float32Array {
+  const n = 1 << L;
+  const out = new Float32Array(n);
+  const M = (1n << 32n) - 1n;
+  for (let i = 0; i < n; i++) {
+    let x = BigInt(i) & M;
+    x = (x * 34038481n + 76625530n) & M;
+    const y =
+      Number((x & 255n) + ((x >> 8n) & 255n) + ((x >> 16n) & 255n) + ((x >> 24n) & 255n)) - 510;
+    out[i] = y / 147.800537109375;
+  }
+  return out;
+}
+
+/** QTIP `decode_3inst`: two bf16-shaped 16-bit lanes masked out of one
+ *  multiply-add, xor'd with an exponent mask, read as fp16 and summed. */
+export function lut3inst(L: number): Float32Array {
+  const n = 1 << L;
+  const out = new Float32Array(n);
+  const M = (1n << 32n) - 1n;
+  const a = 89226354n, b = 64248484n, fpmask = 996162400n;
+  const half = (1n << 15n) + ((1n << 12n) - 1n);
+  const mask = (half << 16n) + half;
+  const f16 = (bits: number): number => {
+    const sgn = bits & 0x8000 ? -1 : 1;
+    const exp = (bits >> 10) & 0x1f;
+    const man = bits & 0x3ff;
+    if (exp === 0) return sgn * man * 2 ** -24;
+    if (exp === 31) return sgn * (man ? NaN : Infinity);
+    return sgn * (1 + man / 1024) * 2 ** (exp - 15);
+  };
+  for (let i = 0; i < n; i++) {
+    let x = BigInt(i) & M;
+    x = (x * a + b) & M;
+    const res = (mask & x) ^ fpmask;
+    out[i] = f16(Number(res >> 16n) & 0xffff) + f16(Number(res & 0xffffn));
+  }
+  return out;
+}
+
+/** Pure-lookup random Gaussian trellis code (paper's RPTC) — the L-bit LUT is
+ *  2^L iid N(0,1) draws. Used only as a codec-validation reference. */
+export function lutRandom(L: number, seed = 0xc0ffee): Float32Array {
+  const n = 1 << L;
+  const out = new Float32Array(n);
+  let s = seed >>> 0;
+  const u = () => {
+    s = (s + 0x9e3779b9) >>> 0;
+    let z = s;
+    z = Math.imul(z ^ (z >>> 16), 0x21f0aaad) >>> 0;
+    z = Math.imul(z ^ (z >>> 15), 0x735a2d97) >>> 0;
+    return ((z ^ (z >>> 15)) >>> 0) / 4294967296 + 2 ** -33;
+  };
+  for (let i = 0; i < n; i += 2) {
+    const r = Math.sqrt(-2 * Math.log(u()));
+    const t = 2 * Math.PI * u();
+    out[i] = r * Math.cos(t);
+    if (i + 1 < n) out[i + 1] = r * Math.sin(t);
+  }
+  return out;
+}
+
+// ----------------------------------------------------------------- codec ---
+
+export type CodeName = "1mad" | "3inst" | "rand";
+
+export interface TrellisCfg {
+  L: number;          // state bits
+  K: number;          // bits per weight (V = 1)
+  T: number;          // block length (weights per tail-biting trellis)
+  code: CodeName;
+  tailBiting: boolean;
+}
+
+export class Trellis {
+  readonly cfg: TrellisCfg;
+  readonly nStates: number;
+  readonly G: number;      // 2^K branches
+  readonly J: number;      // 2^(L-K) predecessor groups
+  readonly lutRms: number;
+  /** [S] f32 code values, and the [1,S] row / -lut² row the encoder needs. */
+  private lutFlat: MlxArray;
+  private lutRow: MlxArray;
+  private negLut2: MlxArray;
+  private sHigh: MlxArray;   // [1,S] int32, state >> K
+  private sLow: MlxArray;    // [1,S] int32, state & (J-1)
+  private mrow: MlxArray;    // [1,G,1] uint8, branch index 0..G-1
+  private sIdx: MlxArray;    // [1,S] int32, arange(S)
+  private negOne: MlxArray;  // int32 scalar -1
+  private kConst: MlxArray;  // int32 scalar K
+  private jConst: MlxArray;  // int32 scalar J
+  private negBig: MlxArray;  // f32 scalar -1e30
+
+  constructor(cfg: TrellisCfg) {
+    this.cfg = cfg;
+    if (!Number.isInteger(cfg.K) || cfg.K < 1 || cfg.K >= cfg.L)
+      throw new Error(`Trellis: K=${cfg.K} must be an integer in [1, L=${cfg.L})`);
+    const n = (this.nStates = 1 << cfg.L);
+    this.G = 1 << cfg.K;
+    this.J = n >> cfg.K;
+    const host =
+      cfg.code === "1mad" ? lut1mad(cfg.L) : cfg.code === "3inst" ? lut3inst(cfg.L) : lutRandom(cfg.L);
+    let ss = 0;
+    for (const v of host) ss += v * v;
+    this.lutRms = Math.sqrt(ss / n);
+
+    this.lutFlat = Arr.fromFloat32(host, [n]);
+    this.lutRow = ops.reshape(this.lutFlat, [1, n], S);
+    const sq = ops.square(this.lutRow, S);
+    this.negLut2 = ops.neg(sq, S);
+    sq.dispose();
+
+    const idx = new Int32Array(n);
+    for (let i = 0; i < n; i++) idx[i] = i >> cfg.K;
+    this.sHigh = Arr.fromInt32(idx, [1, n]);
+    const idx2 = new Int32Array(n);
+    for (let i = 0; i < n; i++) idx2[i] = i & (this.J - 1);
+    this.sLow = Arr.fromInt32(idx2, [1, n]);
+
+    const idx3 = new Int32Array(n);
+    for (let i = 0; i < n; i++) idx3[i] = i;
+    this.sIdx = Arr.fromInt32(idx3, [1, n]);
+    this.negOne = Arr.fromInt32(new Int32Array([-1]), []);
+
+    const mh = new Int32Array(this.G);
+    for (let i = 0; i < this.G; i++) mh[i] = i;
+    const m32 = Arr.fromInt32(mh, [1, this.G, 1]);
+    this.mrow = m32.astype(Dtype.uint8, S);
+    m32.dispose();
+
+    this.kConst = Arr.fromInt32(new Int32Array([cfg.K]), []);
+    this.jConst = Arr.fromInt32(new Int32Array([this.J]), []);
+    this.negBig = Arr.fromFloat32(new Float32Array([-1e30]), []);
+    ops.evalAll([this.lutFlat, this.lutRow, this.negLut2, this.sHigh, this.sLow, this.mrow, this.sIdx]);
+  }
+
+  dispose(): void {
+    for (const a of [this.lutFlat, this.lutRow, this.negLut2, this.sHigh, this.sLow,
+      this.mrow, this.sIdx, this.negOne, this.kConst, this.jConst, this.negBig]) a.dispose();
+  }
+
+  /** One Viterbi pass over [B,T] (time on axis 1, already scaled to the code's
+   *  variance). `overlap` [B,1] int32 or null. Returns the decoded state at
+   *  every t >= `from` ([B,1] int32 each; earlier slots are null). */
+  private viterbi(X2T: MlxArray, B: number, overlap: MlxArray | null, from: number): (MlxArray | null)[] {
+    const { T } = this.cfg;
+    const n = this.nStates, G = this.G, J = this.J;
+    const col = (t: number): MlxArray => {
+      const row = X2T.slice([t, 0], [t + 1, B], S);   // [1,B]
+      const c = ops.reshape(row, [B, 1], S);
+      row.dispose();
+      return c;
+    };
+
+    // t = 0 (initial state free, or pinned to the tail-biting overlap).
+    let x = col(0);
+    let g = ops.addmm(this.negLut2, x, this.lutRow, S);
+    x.dispose();
+    if (overlap) {
+      const m = ops.equal(this.sHigh, overlap, S);
+      const masked = ops.where(m, g, this.negBig, S);
+      m.dispose(); g.dispose();
+      g = masked;
+    }
+
+    // NOTE: mlx's ArgReduce kernel is ~85x slower than Reduce on this shape
+    // (measured M1 Max: argmax over [4096,8,512] axis 1 = 42.9 ms vs max =
+    // 0.5 ms), so the branch index is recovered from the max with a
+    // compare + uint8 max instead of a real argmax.
+    const back: (MlxArray | null)[] = new Array(T).fill(null);
+    for (let t = 1; t < T; t++) {
+      const g3 = ops.reshape(g, [B, G, J], S);
+      const mx = ops.maxAxis(g3, 1, true, S);                    // [B,1,J]
+      const eq = ops.equal(g3, mx, S);                           // bool [B,G,J]
+      const e8 = eq.astype(Dtype.uint8, S);
+      const sel = ops.mul(e8, this.mrow, S);
+      back[t] = ops.maxAxis(sel, 1, false, S);                   // [B,J] uint8
+      const best = ops.reshape(mx, [B, J, 1], S);
+      x = col(t);
+      const err = ops.addmm(this.negLut2, x, this.lutRow, S);    // [B,S]
+      const err3 = ops.reshape(err, [B, J, G], S);
+      const gn3 = ops.add(err3, best, S);
+      const gn = ops.reshape(gn3, [B, n], S);
+      ops.evalAll([gn, back[t]!]);
+      for (const a of [g3, mx, eq, e8, sel, best, x, err, err3, gn3, g]) a.dispose();
+      g = gn;
+    }
+
+    if (overlap) {
+      const m = ops.equal(this.sLow, overlap, S);
+      const masked = ops.where(m, g, this.negBig, S);
+      m.dispose(); g.dispose();
+      g = masked;
+    }
+    // argmax over all 2^L states, again via max+compare (see note above).
+    const fmx = ops.maxAxis(g, 1, true, S);                       // [B,1]
+    const feq = ops.equal(g, fmx, S);
+    const fsel = ops.where(feq, this.sIdx, this.negOne, S);       // [B,S] int32
+    const finI = ops.maxAxis(fsel, 1, true, S);                   // [B,1] int32
+    for (const a of [fmx, feq, fsel, g]) a.dispose();
+
+    const states: (MlxArray | null)[] = new Array(T).fill(null);
+    let cur = finI;
+    let curKept = T - 1 >= from;
+    if (curKept) states[T - 1] = cur;
+    for (let t = T - 1; t >= 1; t--) {
+      const j = ops.rightShift(cur, this.kConst, S);              // [B,1] int32
+      const m8 = ops.takeAlongAxis(back[t]!, j, 1, S);            // [B,1] uint8
+      const mi = m8.astype(Dtype.int32, S);
+      const off = ops.mul(mi, this.jConst, S);
+      const prev = ops.add(j, off, S);
+      for (const a of [m8, mi, off, j]) a.dispose();
+      back[t]!.dispose();
+      back[t] = null;
+      if (!curKept) cur.dispose();
+      cur = prev;
+      curKept = t - 1 >= from;
+      if (curKept) states[t - 1] = cur;
+    }
+    if (!curKept) cur.dispose();
+    return states;
+  }
+
+  /** Trellis-encode + decode [B,T] rows (already variance-matched to the code).
+   *  Returns [B,T] f32 reconstruction. */
+  encodeDecode(X: MlxArray): MlxArray {
+    const [B, T] = X.shape as [number, number];
+    if (T !== this.cfg.T) throw new Error(`block length ${T} != cfg.T ${this.cfg.T}`);
+    // NOTE: every intermediate gets its own handle and an explicit dispose —
+    // an undisposed handle for a view (transpose/concat/slice) pins its PARENT
+    // buffer, which is how the first 27-tensor run leaked ~1.4 GB per tensor.
+    const two = ops.mulScalar(X, 2, S);
+    const twoT = ops.transposeAxes(two, [1, 0], S);
+    const X2T = ops.contiguous(twoT, S);                                // [T,B]
+    two.dispose(); twoT.dispose();
+    ops.evalAll([X2T]);
+
+    let overlap: MlxArray | null = null;
+    if (this.cfg.tailBiting) {
+      // Algorithm 4: rotate right by ⌊T/2⌋, unconstrained Viterbi, read the
+      // state that lands on original t=0 and use its top L-k bits as overlap.
+      const H = T >> 1;
+      const a = X2T.slice([H, 0], [T, B], S);
+      const b = X2T.slice([0, 0], [H, B], S);
+      const cat = ops.concatAxis([a, b], 0, S);
+      const rot = ops.contiguous(cat, S);
+      a.dispose(); b.dispose(); cat.dispose();
+      ops.evalAll([rot]);
+      const st = this.viterbi(rot, B, null, H);
+      rot.dispose();
+      overlap = ops.rightShift(st[H]!, this.kConst, S);
+      ops.evalAll([overlap]);
+      for (const s of st) s?.dispose();
+      clearCache();
+    }
+
+    const st = this.viterbi(X2T, B, overlap, 0);
+    X2T.dispose();
+    overlap?.dispose();
+    const idx = ops.concatAxis(st as MlxArray[], 1, S);                 // [B,T] int32
+    for (const s of st) s?.dispose();
+    const out = ops.takeAxis(this.lutFlat, idx, 0, S);
+    idx.dispose();
+    ops.evalAll([out]);
+    return out;
+  }
+
+  /** encodeDecode over [B,T] in chunks of at most `maxB` rows. The caller has
+   *  already variance-matched X to the code (BlockLDLQ path). */
+  encodeDecodeChunked(X: MlxArray, maxB: number): MlxArray {
+    const [B, T] = X.shape as [number, number];
+    if (B <= maxB) return this.encodeDecode(X);
+    const parts: MlxArray[] = [];
+    for (let b0 = 0; b0 < B; b0 += maxB) {
+      const b1 = Math.min(B, b0 + maxB);
+      const sl = X.slice([b0, 0], [b1, T], S);
+      const chunk = ops.contiguous(sl, S);
+      sl.dispose();
+      ops.evalAll([chunk]);
+      parts.push(this.encodeDecode(chunk));
+      chunk.dispose();
+      clearCache();
+    }
+    const out = ops.concatAxis(parts, 0, S);
+    for (const p of parts) p.dispose();
+    ops.evalAll([out]);
+    return out;
+  }
+
+  /** Fake-quant a 2-D tensor whose trellis axis is the LAST axis. Per-row
+   *  (per rotated vector) fp16 scale, QTIP's Wscale = RMS(W)/RMS(lut) applied
+   *  per row instead of per tensor (deviation: our incoherence processing is
+   *  one-sided, so rows do not share a norm). */
+  fakeQuantRows(W: MlxArray, batchBlocks: number, onProgress?: (done: number, total: number) => void): MlxArray {
+    const [N, C] = W.shape as [number, number];
+    const T = this.cfg.T;
+    if (C % T !== 0) throw new Error(`trellis axis ${C} not a multiple of block ${T}`);
+    const sq = ops.square(W, S);
+    const ms = ops.meanAxis(sq, 1, true, S);
+    sq.dispose();
+    const rms = ops.sqrt(ms, S);
+    ms.dispose();
+    const inv = ops.mulScalar(rms, 1 / this.lutRms, S);   // scale = rms/lutRms
+    rms.dispose();
+    // guard all-zero rows
+    const zero = Arr.fromFloat32(new Float32Array([0]), []);
+    const one = Arr.fromFloat32(new Float32Array([1]), []);
+    const isz = ops.equal(inv, zero, S);
+    const scale = ops.where(isz, one, inv, S);
+    for (const a of [zero, one, isz, inv]) a.dispose();
+    // fp16 round-trip: the scale we would actually store
+    const s16 = scale.astype(Dtype.float16, S);
+    const scaleQ = s16.astype(Dtype.float32, S);
+    scale.dispose(); s16.dispose();
+
+    const Xn = ops.div(W, scaleQ, S);
+    const flat = ops.reshape(Xn, [(N * C) / T, T], S);
+    Xn.dispose();
+    ops.evalAll([flat]);
+
+    const nBlocks = (N * C) / T;
+    const parts: MlxArray[] = [];
+    for (let b0 = 0; b0 < nBlocks; b0 += batchBlocks) {
+      const b1 = Math.min(nBlocks, b0 + batchBlocks);
+      const sl = flat.slice([b0, 0], [b1, T], S);
+      const chunk = ops.contiguous(sl, S);
+      sl.dispose();
+      ops.evalAll([chunk]);
+      const rec = this.encodeDecode(chunk);
+      chunk.dispose();
+      parts.push(rec);
+      clearCache();
+      onProgress?.(b1, nBlocks);
+    }
+    flat.dispose();
+    const all = parts.length === 1 ? parts[0]! : ops.concatAxis(parts, 0, S);
+    if (parts.length > 1) for (const p of parts) p.dispose();
+    const rec2 = ops.reshape(all, [N, C], S);
+    all.dispose();
+    const out = ops.mul(rec2, scaleQ, S);
+    rec2.dispose();
+    scaleQ.dispose();
+    ops.evalAll([out]);
+    return out;
+  }
+
+  /** Trellis-encode [B,T] rows (already variance-matched) and return the
+   *  decoded STATE at every t as [B,T] int32 — the k low bits of state t are
+   *  the emitted symbol, and (bitshift trellis) state t is exactly the last
+   *  L/k symbols, which is what makes random-access decode possible. */
+  encodeStates(X: MlxArray): MlxArray {
+    const [B, T] = X.shape as [number, number];
+    if (T !== this.cfg.T) throw new Error(`block length ${T} != cfg.T ${this.cfg.T}`);
+    const two = ops.mulScalar(X, 2, S);
+    const twoT = ops.transposeAxes(two, [1, 0], S);
+    const X2T = ops.contiguous(twoT, S);
+    two.dispose(); twoT.dispose();
+    ops.evalAll([X2T]);
+    let overlap: MlxArray | null = null;
+    if (this.cfg.tailBiting) {
+      const H = T >> 1;
+      const a = X2T.slice([H, 0], [T, B], S);
+      const b = X2T.slice([0, 0], [H, B], S);
+      const cat = ops.concatAxis([a, b], 0, S);
+      const rot = ops.contiguous(cat, S);
+      a.dispose(); b.dispose(); cat.dispose();
+      ops.evalAll([rot]);
+      const st = this.viterbi(rot, B, null, H);
+      rot.dispose();
+      overlap = ops.rightShift(st[H]!, this.kConst, S);
+      ops.evalAll([overlap]);
+      for (const s of st) s?.dispose();
+      clearCache();
+    }
+    const st = this.viterbi(X2T, B, overlap, 0);
+    X2T.dispose();
+    overlap?.dispose();
+    const idx = ops.concatAxis(st as MlxArray[], 1, S);
+    for (const s of st) s?.dispose();
+    ops.evalAll([idx]);
+    return idx;
+  }
+
+  /** States → decoded values ([B,T] f32), the exact map encodeDecode uses. */
+  decodeStates(idx: MlxArray): MlxArray {
+    const out = ops.takeAxis(this.lutFlat, idx, 0, S);
+    ops.evalAll([out]);
+    return out;
+  }
+
+  /** fakeQuantRows that ALSO returns the packed bit-stream and the fp16 row
+   *  scales — the on-disk trellis format (see packStates). `rec` is the same
+   *  f32 reconstruction fakeQuantRows returns, so the fake-quant and packed
+   *  artifacts are built from one Viterbi pass. */
+  fakeQuantRowsPacked(
+    W: MlxArray, batchBlocks: number, onProgress?: (done: number, total: number) => void,
+  ): { rec: MlxArray; codes: MlxArray; scales: MlxArray } {
+    const [N, C] = W.shape as [number, number];
+    const T = this.cfg.T;
+    if (C % T !== 0) throw new Error(`trellis axis ${C} not a multiple of block ${T}`);
+    const sq = ops.square(W, S);
+    const ms = ops.meanAxis(sq, 1, true, S);
+    sq.dispose();
+    const rms = ops.sqrt(ms, S);
+    ms.dispose();
+    const inv = ops.mulScalar(rms, 1 / this.lutRms, S);
+    rms.dispose();
+    const zero = Arr.fromFloat32(new Float32Array([0]), []);
+    const one = Arr.fromFloat32(new Float32Array([1]), []);
+    const isz = ops.equal(inv, zero, S);
+    const scale = ops.where(isz, one, inv, S);
+    for (const a of [zero, one, isz, inv]) a.dispose();
+    const s16 = scale.astype(Dtype.float16, S);          // [N,1] f16 — stored
+    const scaleQ = s16.astype(Dtype.float32, S);
+    scale.dispose();
+    const scales = ops.reshape(s16, [N], S);
+    s16.dispose();
+    ops.evalAll([scales]);
+
+    const Xn = ops.div(W, scaleQ, S);
+    const flat = ops.reshape(Xn, [(N * C) / T, T], S);
+    Xn.dispose();
+    ops.evalAll([flat]);
+
+    const nBlocks = (N * C) / T;
+    const wordsPerBlock = (T * this.cfg.K) / 32;
+    const packed = new Uint32Array(nBlocks * wordsPerBlock);
+    const parts: MlxArray[] = [];
+    for (let b0 = 0; b0 < nBlocks; b0 += batchBlocks) {
+      const b1 = Math.min(nBlocks, b0 + batchBlocks);
+      const sl = flat.slice([b0, 0], [b1, T], S);
+      const chunk = ops.contiguous(sl, S);
+      sl.dispose();
+      ops.evalAll([chunk]);
+      const idx = this.encodeStates(chunk);
+      chunk.dispose();
+      const rec = this.decodeStates(idx);
+      parts.push(rec);
+      const u32 = idx.astype(Dtype.uint32, S);           // rawBytes has no int32 path; states are >= 0
+      const raw = u32.rawBytes();                        // [B,T] row-major
+      u32.dispose();
+      const st = new Int32Array(raw.buffer, raw.byteOffset, raw.byteLength / 4);
+      packStates(st, b1 - b0, T, this.cfg.K, packed.subarray(b0 * wordsPerBlock, b1 * wordsPerBlock));
+      idx.dispose();
+      clearCache();
+      onProgress?.(b1, nBlocks);
+    }
+    flat.dispose();
+    const all = parts.length === 1 ? parts[0]! : ops.concatAxis(parts, 0, S);
+    if (parts.length > 1) for (const p of parts) p.dispose();
+    const rec2 = ops.reshape(all, [N, C], S);
+    all.dispose();
+    const out = ops.mul(rec2, scaleQ, S);
+    rec2.dispose();
+    scaleQ.dispose();
+    ops.evalAll([out]);
+    const codes = Arr.fromBytesCopy(
+      new Uint8Array(packed.buffer, packed.byteOffset, packed.byteLength),
+      [N, (C * this.cfg.K) / 32], Dtype.uint32,
+    );
+    return { rec: out, codes, scales };
+  }
+}
+
+// ------------------------------------------------------------ packed format ---
+//
+// One coded row = C/T blocks; one block = T·k bits = T·k/32 uint32 words
+// (little-endian words, bit i of the block = word i>>5, bit i&31). Symbol b_t
+// (= state_t & (2^k−1)) sits at bit offset (T−1−t)·k — REVERSED time — so the
+// L-bit window starting at that offset, read LSB-first and wrapping around the
+// block (tail-biting), is exactly state_t = b_t | b_{t−1}<<k | … . The kernel
+// therefore decodes any position in O(1): two word loads, a shift, a mask, and
+// the 1MAD hash (or a 4096-entry LUT). The fp16 per-row scale is stored as a
+// separate `.scales` [rows] tensor; the reconstructed weight is
+// bf16(f32(lut[state]) · f32(scale)), bit-identical to the fake-quant artifact.
+
+export const TRELLIS_L = 12;
+export const TRELLIS_T = 256;
+
+/** Words per T-block for k bits/weight (T·k must be a multiple of 32). */
+export function wordsPerBlock(T: number, k: number): number {
+  if ((T * k) % 32 !== 0) throw new Error(`T·k = ${T * k} not a multiple of 32`);
+  return (T * k) / 32;
+}
+
+/** Reorder k3/T256 codes into [two-block groups, rows, 48 words].
+ *  The caller selects an axis0/1MAD/L12 tensor; no code values change. */
+export function interleaveTrellisCodes(codes: MlxArray): MlxArray {
+  if (codes.dtype !== Dtype.uint32 || codes.ndim !== 2 ||
+      codes.shape[0]! < 1 || codes.shape[1]! < 48 || codes.shape[1]! % 48 !== 0)
+    throw new Error("trellis interleave: expected uint32 rows with whole two-block groups");
+  const [rows, words] = codes.shape as [number, number];
+  const view = ops.reshape(codes, [rows, words / 48, 48]);
+  const transposed = ops.transposeAxes(view, [1, 0, 2]);
+  view.dispose();
+  try { return ops.contiguous(transposed); }
+  finally { transposed.dispose(); }
+}
+
+/** Pack `nBlocks` blocks of T states ([nBlocks·T] int32, row-major) into the
+ *  reversed-time bit-stream described above, writing into `out`
+ *  (nBlocks · wordsPerBlock uint32). */
+export function packStates(states: Int32Array, nBlocks: number, T: number, k: number, out: Uint32Array): void {
+  const wpb = wordsPerBlock(T, k);
+  const mask = (1 << k) - 1;
+  for (let b = 0; b < nBlocks; b++) {
+    const base = b * wpb;
+    for (let w = 0; w < wpb; w++) out[base + w] = 0;
+    for (let t = 0; t < T; t++) {
+      const sym = states[b * T + t]! & mask;
+      const p = (T - 1 - t) * k;
+      const wi = p >>> 5, off = p & 31;
+      out[base + wi] = (out[base + wi]! | (sym << off)) >>> 0;
+      if (off + k > 32) out[base + wi + 1] = (out[base + wi + 1]! | (sym >>> (32 - off))) >>> 0;
+    }
+  }
+}
+
+/** Read state_t back out of one packed block (host reference for the kernel). */
+export function stateAt(words: Uint32Array, blockWordOffset: number, wpb: number, t: number, T: number, k: number, L: number): number {
+  const p = (T - 1 - t) * k;
+  const wi = p >>> 5, off = p & 31;
+  const w0 = words[blockWordOffset + wi]!;
+  const w1 = words[blockWordOffset + ((wi + 1) % wpb)]!;
+  const lo = off === 0 ? w0 : ((w0 >>> off) | (w1 << (32 - off))) >>> 0;
+  return lo & ((1 << L) - 1);
+}
+
+/** Host decode of a whole packed [N, C] tensor to f32 (reference only). */
+export function unpackDecodeHost(
+  words: Uint32Array, scales: Float32Array, N: number, C: number, k: number, T: number, lut: Float32Array,
+): Float32Array {
+  const wpb = wordsPerBlock(T, k);
+  const L = Math.log2(lut.length);
+  const out = new Float32Array(N * C);
+  const blocksPerRow = C / T;
+  for (let n = 0; n < N; n++) {
+    for (let b = 0; b < blocksPerRow; b++) {
+      const bw = (n * blocksPerRow + b) * wpb;
+      for (let t = 0; t < T; t++) out[n * C + b * T + t] = lut[stateAt(words, bw, wpb, t, T, k, L)]! * scales[n]!;
+    }
+  }
+  return out;
+}