mlx-bun 0.2.0 → 0.4.0

package/src/model/qwen3_5-checkpoint.ts

@@ -0,0 +1,164 @@
+// Qwen3.5 checkpoint-generation normalization — the load-seam port of
+// mlx_lm.models.qwen3_5's `Model.sanitize` + `TextModel.sanitize`.
+//
+// Two checkpoint generations ship the same graph under different conventions:
+//
+//   pre-5.8 (mlx-lm-converted)    5.8-family (transformers 5.8.0 export)
+//   ──────────────────────────    ───────────────────────────────────────
+//   language_model.model.*        model.language_model.*
+//   language_model.lm_head.*      lm_head.*                (top level)
+//   vision_tower.*                model.visual.*
+//   (no mtp tensors in-repo)      mtp.*                    (in-repo)
+//   conv1d.weight [C, K, 1]       conv1d.weight [C, 1, K]  (HF layout)
+//   RMSNorm gains stored as γ     RMSNorm gains stored as γ−1
+//
+// Our graph speaks the FIRST column (`PREFIX = "language_model"` in
+// qwen3_5.ts) — which is exactly mlx-lm's post-sanitize name space — so this
+// module maps the artifact onto it instead of teaching every consumer two
+// namings. The rules below are a line-for-line copy of the oracle; see
+// mlx_lm/models/qwen3_5.py `Model.sanitize` (naming/drop) and
+// `TextModel.sanitize` (γ shift, conv layout).
+//
+// The γ shift is the dangerous one: a spurious +1.0 silently corrupts an old
+// artifact. The oracle's discriminator is NOT the naming — it is
+//   should_shift = any("mtp." in k) or any conv1d.weight with shape[-1] != 1
+// evaluated over the whole (post-rename, pre-drop) weight map. Reproduced
+// verbatim. Measured on the two artifacts in play (2026-09-01, M1 Max):
+//   • ~/.cache/…/models--mjriii--Qwen3.8-27B/snapshots/staged — no mtp.*,
+//     conv1d [10240,4,1] → no shift; layer-0 input_layernorm ≡ 1.0.
+//   • /Volumes/MLX-Models/models/mjriii/Qwen3.8-27B-{tqalloc-norot,rtn4-g64}
+//     — 31 mtp.* tensors, conv1d [10240,1,4] → shift; layer-0
+//     input_layernorm mean −0.0334 (= γ−1), byte-identical to the raw
+//     Qwen/Qwen3.8-27B 5.8 checkpoint (the quantizer copies norms verbatim).
+
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { Weights, WeightsView } from "../weights";
+
+/** Header-level facts about one stored tensor (no bytes are read). */
+export interface CheckpointTensor {
+  readonly name: string;
+  readonly shape: readonly number[];
+}
+
+export interface Qwen35Sanitized {
+  /** canonical (graph) name → the name the artifact actually stores. */
+  readonly names: Map<string, string>;
+  /** Canonical names stored as γ−1: a +1.0 is due at load. */
+  readonly normShift: Set<string>;
+  /** Canonical names in HF conv layout: a moveaxis(2, 1) is due at load. */
+  readonly convMoveaxis: Set<string>;
+  /** True when the artifact already IS our canonical layout — no view needed
+   *  and the pre-existing load path runs untouched. */
+  readonly identity: boolean;
+}
+
+/** mlx-lm `TextModel.sanitize` norm_keys, verbatim and in order. */
+const NORM_KEYS = [
+  ".input_layernorm.weight",
+  ".post_attention_layernorm.weight",
+  "model.norm.weight",
+  ".q_norm.weight",
+  ".k_norm.weight",
+] as const;
+
+/** mlx-lm `Model.sanitize`: drop vision, fold every naming onto
+ *  `language_model.…`. Returns null for a dropped tensor. */
+function canonicalName(key: string): string | null {
+  if (key.startsWith("vision_tower") || key.startsWith("model.visual")) return null;
+  if (key.startsWith("model.language_model"))
+    return "language_model.model" + key.slice("model.language_model".length);
+  if (key.startsWith("language_model.")) return key;
+  return "language_model." + key;
+}
+
+/**
+ * Port of mlx_lm.models.qwen3_5 `Model.sanitize` ∘ `TextModel.sanitize`,
+ * name/shape half. Pure: takes header facts, returns the load plan.
+ */
+export function sanitizeQwen35Checkpoint(
+  tensors: Iterable<CheckpointTensor>,
+): Qwen35Sanitized {
+  // Pass 1 — rename (vision dropped), exactly as Model.sanitize.
+  const renamed: { canonical: string; source: string; shape: readonly number[] }[] = [];
+  let sawForeignName = false;
+  for (const t of tensors) {
+    const canonical = canonicalName(t.name);
+    if (canonical === null) {
+      sawForeignName ||= t.name.startsWith("model.visual");
+      continue;
+    }
+    if (canonical !== t.name) sawForeignName = true;
+    renamed.push({ canonical, source: t.name, shape: t.shape });
+  }
+
+  // Pass 2 — TextModel.sanitize's two flags, over the renamed map BEFORE the
+  // mtp drop (the oracle computes them first, then filters).
+  const hasMtpWeights = renamed.some((t) => t.canonical.includes("mtp."));
+  const hasUnsanitizedConv1d = renamed.some(
+    (t) => t.canonical.includes("conv1d.weight") && t.shape[t.shape.length - 1] !== 1,
+  );
+  const shouldShiftNormWeights = hasMtpWeights || hasUnsanitizedConv1d;
+
+  const names = new Map<string, string>();
+  const normShift = new Set<string>();
+  const convMoveaxis = new Set<string>();
+  for (const t of renamed) {
+    if (t.canonical.includes("mtp.")) continue; // oracle drops these outright
+    names.set(t.canonical, t.source);
+    if (t.canonical.includes("conv1d.weight") && t.shape[t.shape.length - 1] !== 1)
+      convMoveaxis.add(t.canonical);
+    if (
+      shouldShiftNormWeights &&
+      NORM_KEYS.some((sfx) => t.canonical.endsWith(sfx)) &&
+      t.shape.length === 1
+    )
+      normShift.add(t.canonical);
+  }
+
+  return {
+    names,
+    normShift,
+    convMoveaxis,
+    identity: !sawForeignName && !hasMtpWeights && !shouldShiftNormWeights &&
+      convMoveaxis.size === 0,
+  };
+}
+
+/**
+ * The {@link WeightsView} for a qwen3.5 artifact, or null when the artifact is
+ * already canonical (then nothing is installed and the load path is bit-for-bit
+ * what it was before this seam existed).
+ */
+export function qwen35WeightsView(weights: Weights): WeightsView | null {
+  const plan = sanitizeQwen35Checkpoint(
+    weights.tensorNames.map((name) => ({ name, shape: weights.info(name).shape })),
+  );
+  if (plan.identity) return null;
+  return {
+    names: plan.names,
+    fixup(canonical: string, arr: MlxArray): MlxArray | null {
+      // Oracle order: moveaxis first, then the γ shift (disjoint sets — a
+      // conv1d weight is 3-D and never a norm gain).
+      if (plan.convMoveaxis.has(canonical)) {
+        // `v.moveaxis(2, 1)` on [C, 1, K] → [C, K, 1]. Materialized: the
+        // conv1d kernel consumes it every forward, and the transposed view
+        // would otherwise pin the source shard buffer.
+        const moved = ops.transposeAxes(arr, [0, 2, 1]);
+        const out = ops.contiguous(moved);
+        moved.dispose();
+        return out;
+      }
+      if (plan.normShift.has(canonical)) {
+        // `v + 1.0` — a weak python scalar, so the add happens at the gain's
+        // own dtype (bf16 here); scalarLike builds the constant host-side at
+        // that dtype for the same single-kernel add the reference dispatches.
+        const one = ops.scalarLike(1, arr);
+        const out = ops.add(arr, one);
+        one.dispose();
+        return out;
+      }
+      return null;
+    },
+  };
+}

package/src/model/qwen3_5.ts

@@ -10,18 +10,20 @@
 // shared maybeQuantizeKv path).
 
 import type { ModelConfig } from "../config";
+import { gatedDeltaState } from "./qwen3-delta-state";
 import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
+import { Dtype, deviceArchitecture } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { CompiledFunction } from "../mlx/compile";
+import { qwenAppendChunkSize } from "./qwen-append";
+import { TrellisLinear, fusedGateUpEligible, fusedGateUpSwiglu, TRELLIS_MATVEC_MAX_M } from "./trellis-linear";
 import {
   argmaxLastPosition,
   disposeTriple,
   disposing,
   KVCache,
   LoraState,
-  QuantizedKVCache,
   QuantizedEmbedding,
   QuantizedLinear,
   quantizedSdpa,
@@ -29,7 +31,10 @@ import {
   type Cache,
   type Mask,
 } from "./gemma4-base";
+import { qwen35WeightsView } from "./qwen3_5-checkpoint";
+import { materializeCopy } from "../mlx/materialize";
 import { gatedDeltaUpdate, SSMCache } from "./qwen3-delta";
+import type { QwenConvolution } from "./qwen-conv";
 import {
   activeMrope, applyInterleavedRope, buildMropePositions, mropeInvFreq,
   setActiveMrope, type MropeRequestState,
@@ -107,6 +112,10 @@ export function compiledSilu(x: MlxArray): MlxArray {
 
 /** Gated-DeltaNet linear-attention layer (mlx-lm GatedDeltaNet). */
 export class GatedDeltaNet {
+  /** Per-model experiment seam; null retains the oracle graph. The same
+   * implementation advances speculative rollback prefixes. Borrowed inputs,
+   * owned activation and independent state tail; no global runtime mutation. */
+  convolution: QwenConvolution | null = null;
   readonly inProjQkv: QuantizedLinear;
   readonly inProjZ: QuantizedLinear;
   readonly inProjB: QuantizedLinear;
@@ -146,7 +155,40 @@ export class GatedDeltaNet {
     this.normWeight = weights.tensor(`${prefix}.norm.weight`);
   }
 
-  forward(x: MlxArray, cache: SSMCache): MlxArray {
+  #convolve(qkv: MlxArray, state: MlxArray, rowLengths?: readonly number[]): [MlxArray, MlxArray] {
+    if (this.convolution) {
+      const result = this.convolution(qkv, state, this.convWeight);
+      if (!rowLengths) return result;
+      using input = ops.concatAxis([state, qkv], 1);
+      const tail = this.#convTail(input, qkv.shape[1]!, rowLengths);
+      result[1].dispose();
+      return [result[0], tail];
+    }
+    const [, S, D] = qkv.shape as [number, number, number];
+    const input = ops.concatAxis([state, qkv], 1);
+    // MLX copy and contiguous can both alias the whole prefill buffer.
+    // Materialize only the tail so cache residency follows its logical size.
+    const tail = this.#convTail(input, S, rowLengths);
+    const conv = ops.conv1d(input, this.convWeight, 1, 0, 1, D);
+    input.dispose();
+    const out = compiledSilu(conv);
+    conv.dispose();
+    return [out, tail];
+  }
+
+  #convTail(input: MlxArray, processed: number, rowLengths?: readonly number[]): MlxArray {
+    const [B, , D] = input.shape as [number, number, number];
+    const nKeep = this.convKernel - 1;
+    if (rowLengths) {
+      using indices = ops.fromInt32(rowLengths.flatMap(length =>
+        Array.from({ length: nKeep }, (_, position) => length + position)), [B, nKeep, 1]);
+      return ops.takeAlongAxis(input, indices, 1);
+    }
+    using view = input.slice([0, processed, 0], [B, processed + nKeep, D]);
+    return materializeCopy(view);
+  }
+
+  forward(x: MlxArray, cache: SSMCache, independentRows = false, mask?: MlxArray | null): MlxArray {
     const [B, S] = x.shape as [number, number, number];
     const convDim = this.keyDim * 2 + this.valueDim;
     const nKeep = this.convKernel - 1;
@@ -168,38 +210,34 @@ export class GatedDeltaNet {
     }
 
     const t0 = prof ? performance.now() : 0;
-    const qkv = this.inProjQkv.forward(x); // [B,S,convDim]
-    let z = this.inProjZ.forward(x);
+    let qkv = this.inProjQkv.forward(x, independentRows); // [B,S,convDim]
+    let z = this.inProjZ.forward(x, independentRows);
     z = disposing(z, ops.reshape(z, [B, S, this.numVHeads, this.headVDim]));
-    const b = this.inProjB.forward(x); // [B,S,numVHeads]
-    const a = this.inProjA.forward(x);
+    const b = this.inProjB.forward(x, independentRows); // [B,S,numVHeads]
+    const a = this.inProjA.forward(x, independentRows);
     if (prof) { ops.evalAll([qkv, z, b, a]); prof.proj = (prof.proj ?? 0) + performance.now() - t0; }
     const tc = prof ? performance.now() : 0; // [B,S,numVHeads]
 
-    // Causal depthwise conv with the conv-state prefix (B=1: no ssm mask).
+    // The model shares one padding mask across recurrent layers. Direct layer
+    // callers can request the same cache-owned mask without a model wrapper.
+    using localMask = mask === undefined ? cache.prefillPadding?.makeMask(S) ?? null : null;
+    const ssmMask = mask ?? localMask;
+    if (ssmMask) {
+      using expanded = ops.reshape(ssmMask, [B, S, 1]);
+      using zero = MlxArray.fromBytesCopy(new Uint8Array(qkv.dtype === Dtype.float32 ? 4 : 2), [], qkv.dtype);
+      qkv = disposing(qkv, ops.where(expanded, qkv, zero));
+    }
+    // Convolution keeps each row's last real-token tail, including empty rows.
     const convState =
       cache.conv ?? ops.zeros([B, nKeep, convDim], x.dtype);
-    const convInput = ops.concatAxis([convState, qkv], 1); // [B,S+nKeep,convDim]
+    const [convOut, newConv] = this.#convolve(qkv, convState, cache.prefillPadding?.convolutionLengths(S));
     if (!cache.conv) convState.dispose();
     if (spec) spec.qkv = qkv;
     else qkv.dispose();
-    // New conv state = last nKeep rows. MUST be a TRUE copy (mx.copy):
-    // this tail slice is already row-contiguous at B=1, so ops.contiguous
-    // is a no-op VIEW sharing convInput's buffer — caching it pinned the
-    // whole [B, S+nKeep, convDim] chunk buffer (~42 MB/layer/chunk, 48
-    // layers = the measured 2.0 GB-per-2048-chunk prefill leak; 46 GB
-    // active at 32k, async-GPU-OOM on 24 GB — M4 2026-08-20).
-    const newConvView = convInput.slice([0, S, 0], [B, S + nKeep, convDim]);
-    const newConv = ops.copyOf(newConvView);
-    newConvView.dispose();
     if (spec) spec.prevConv = cache.conv;
     else cache.conv?.dispose();
     cache.conv = newConv;
 
-    const conv = ops.conv1d(convInput, this.convWeight, 1, 0, 1, convDim);
-    convInput.dispose();
-    const convOut = compiledSilu(conv); // nn.silu (mx.compile), matches mlx-lm
-    conv.dispose();
     if (prof) { ops.evalAll([convOut]); prof.conv = (prof.conv ?? 0) + performance.now() - tc; }
     const tn = prof ? performance.now() : 0;
 
@@ -223,7 +261,7 @@ export class GatedDeltaNet {
     const tk = prof ? performance.now() : 0;
 
     const [out, newState] = gatedDeltaUpdate(
-      q, k, v, a, b, this.aLog, this.dtBias, cache.recurrent,
+      q, k, v, a, b, this.aLog, this.dtBias, cache.recurrent, ssmMask,
     );
     if (prof) { ops.evalAll([out]); prof.kernel = (prof.kernel ?? 0) + performance.now() - tk; }
     const to = prof ? performance.now() : 0;
@@ -248,7 +286,7 @@ export class GatedDeltaNet {
     z.dispose();
     const merged = ops.reshape(gated, [B, S, this.valueDim]);
     gated.dispose();
-    const result = this.outProj.forward(merged);
+    const result = this.outProj.forward(merged, independentRows);
     merged.dispose();
     if (prof) { ops.evalAll([result]); prof.out = (prof.out ?? 0) + performance.now() - to; }
     return result;
@@ -262,66 +300,69 @@ export class GatedDeltaNet {
    *  recurrent state is BIT-EXACTLY what a forward over only the accepted
    *  prefix would have produced. Output projections (z / out_proj) are state-
    *  free and skipped: only the states matter here. */
-  #replaySpecPrefix(cache: SSMCache, keep: number): void {
+  #replaySpecPrefix(cache: SSMCache, keep: number | readonly number[]): void {
     const r = cache.specRound;
-    if (!r || !r.qkv || !r.a || !r.b)      throw new Error("GatedDeltaNet replay without recorded round inputs");
+    if (!r || !r.qkv || !r.a || !r.b)
+      throw new Error("GatedDeltaNet replay without recorded round inputs");
     const B = r.qkv.shape[0]!;
+    const count = typeof keep === "number" ? keep : Math.max(...keep);
+    const rowLengths = typeof keep === "number" ? undefined : keep;
     const convDim = this.keyDim * 2 + this.valueDim;
     const nKeep = this.convKernel - 1;
 
-    const qkvPfxView = r.qkv.slice([0, 0, 0], [B, keep, convDim]);
+    const qkvPfxView = r.qkv.slice([0, 0, 0], [B, count, convDim]);
     const qkvPfx = ops.contiguous(qkvPfxView);
     qkvPfxView.dispose();
     const convState = cache.conv ?? ops.zeros([B, nKeep, convDim], qkvPfx.dtype);
-    const convInput = ops.concatAxis([convState, qkvPfx], 1);
+    const [convOut, newConv] = this.#convolve(qkvPfx, convState, rowLengths);
     if (!cache.conv) convState.dispose();
     qkvPfx.dispose();
-    // TRUE copy, same reason as forward(): the contiguous tail slice would
-    // otherwise be a view pinning the whole replay convInput buffer.
-    const rConvView = convInput.slice([0, keep, 0], [B, keep + nKeep, convDim]);
-    const newConv = ops.copyOf(rConvView);
-    rConvView.dispose();
     cache.conv?.dispose();
     cache.conv = newConv;
 
-    const conv = ops.conv1d(convInput, this.convWeight, 1, 0, 1, convDim);
-    convInput.dispose();
-    const convOut = compiledSilu(conv);
-    conv.dispose();
     const [qFlat, kFlat, vFlat] = ops.split(
       convOut, [this.keyDim, 2 * this.keyDim], -1,
     ) as [MlxArray, MlxArray, MlxArray];
     convOut.dispose();
-    let q = ops.reshape(qFlat, [B, keep, this.numKHeads, this.headKDim]);
+    let q = rowLengths ? null : ops.reshape(qFlat, [B, count, this.numKHeads, this.headKDim]);
     qFlat.dispose();
-    let k = ops.reshape(kFlat, [B, keep, this.numKHeads, this.headKDim]);
+    let k = ops.reshape(kFlat, [B, count, this.numKHeads, this.headKDim]);
     kFlat.dispose();
-    const v = disposing(vFlat, ops.reshape(vFlat, [B, keep, this.numVHeads, this.headVDim]));
+    const v = disposing(vFlat, ops.reshape(vFlat, [B, count, this.numVHeads, this.headVDim]));
     const invScale = Math.pow(this.headKDim, -0.5);
-    q = disposing(q, ops.rmsNorm(q, null, 1e-6));
-    q = disposing(q, ops.mulScalar(q, invScale * invScale));
+    if (q) {
+      q = disposing(q, ops.rmsNorm(q, null, 1e-6));
+      q = disposing(q, ops.mulScalar(q, invScale * invScale));
+    }
     k = disposing(k, ops.rmsNorm(k, null, 1e-6));
     k = disposing(k, ops.mulScalar(k, invScale));
 
-    const aPfxView = r.a.slice([0, 0, 0], [B, keep, this.numVHeads]);
+    const aPfxView = r.a.slice([0, 0, 0], [B, count, this.numVHeads]);
     const aPfx = ops.contiguous(aPfxView);
     aPfxView.dispose();
-    const bPfxView = r.b.slice([0, 0, 0], [B, keep, this.numVHeads]);
+    const bPfxView = r.b.slice([0, 0, 0], [B, count, this.numVHeads]);
     const bPfx = ops.contiguous(bPfxView);
     bPfxView.dispose();
 
-    const [y, newState] = gatedDeltaUpdate(
-      q, k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent,
-    );
-    q.dispose();
+    let newState: MlxArray;
+    if (rowLengths) {
+      // Unequal prefixes need a row-length-aware recurrence. The state-only
+      // kernel omits query/output work. Uniform replay retains its measured
+      // kernel until the separate optimization demonstrates a serving win.
+      newState = gatedDeltaState(k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent, rowLengths);
+    } else {
+      const [output, state] = gatedDeltaUpdate(q!, k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent);
+      output.dispose(); newState = state;
+    }
+    q?.dispose();
     k.dispose();
     v.dispose();
     aPfx.dispose();
     bPfx.dispose();
-    y.dispose(); // only the state advance matters on the rollback path
     cache.recurrent?.dispose();
     cache.recurrent = newState;
-    cache.advance(keep);
+    if (typeof keep === "number") cache.advance(keep);
+    else cache.advanceRows(keep);
   }
 
   private rmsNormGated(hidden: MlxArray, gate: MlxArray): MlxArray {
@@ -365,18 +406,18 @@ export class Qwen3Attention {
     this.kNorm = new RMSNorm(weights.tensor(`${prefix}.k_norm.weight`), t.rmsNormEps);
   }
 
-  forward(x: MlxArray, mask: Mask, cache: Cache): MlxArray {
+  forward(x: MlxArray, mask: Mask, cache: Cache, independentRows = false): MlxArray {
     const [B, L] = x.shape as [number, number, number];
 
     // q_proj emits 2× head_dim per head → split into queries + gate.
-    const qp = this.qProj.forward(x);
+    const qp = this.qProj.forward(x, independentRows);
     const qpr = disposing(qp, ops.reshape(qp, [B, L, this.nHeads, this.headDim * 2]));
     const [qHeads, gateHeads] = ops.split(qpr, [this.headDim], -1) as [MlxArray, MlxArray];
     qpr.dispose();
     const gate = disposing(gateHeads, ops.reshape(gateHeads, [B, L, this.nHeads * this.headDim]));
 
-    let k = this.kProj.forward(x);
-    let v = this.vProj.forward(x);
+    let k = this.kProj.forward(x, independentRows);
+    let v = this.vProj.forward(x, independentRows);
 
     // q/k norm over head_dim BEFORE transpose (reference order).
     let q = this.qNorm.forward(qHeads);
@@ -391,7 +432,7 @@ export class Qwen3Attention {
     // Batched decode: the scheduler's mask wrapper exposes each row's REAL
     // position as ropeOffsetArr (rows have different prompt lengths); the
     // dynamic-offset kernel is the same fast::rope, bit-exact vs the static
-    // form (tests/compile.test.ts). Serial lane: scalar offset, unchanged.
+    // form (tests/unit/compile.test.ts). Serial lane: scalar offset, unchanged.
     // Vision requests (serial lane only) install activeMrope for the current
     // forward: 3D interleaved positions via the manual apply — the reference
     // never uses the fused fast-rope kernel when position_ids are supplied,
@@ -411,11 +452,16 @@ export class Qwen3Attention {
     }
 
     let attn: MlxArray;
-    if (cache instanceof QuantizedKVCache) {
-      const [keys, values] = cache.updateAndFetchQuantized(k, v);
+    const quantized = cache.quantizedAttention;
+    if (cache.attentionState) {
+      const view = cache.attentionState.appendAndFetch(k, v);
+      k.dispose(); v.dispose();
+      try { attn = view.attend(q, this.scale, mask); } finally { view.dispose(); }
+    } else if (quantized) {
+      const [keys, values] = quantized.updateAndFetchQuantized(k, v);
       k.dispose();
       v.dispose();
-      attn = quantizedSdpa(q, keys, values, this.scale, mask, cache.groupSize, cache.bits);
+      attn = quantizedSdpa(q, keys, values, this.scale, mask, quantized.groupSize, quantized.bits);
       disposeTriple(keys);
       disposeTriple(values);
     } else {
@@ -439,32 +485,56 @@ export class Qwen3Attention {
     const gated = ops.mul(merged, sig);
     merged.dispose();
     sig.dispose();
-    const out = this.oProj.forward(gated);
+    const out = this.oProj.forward(gated, independentRows);
     gated.dispose();
     return out;
   }
 }
 
+/** MLP projections are affine-quantized (QuantizedLinear) or, in a Q2b
+ *  packed-trellis artifact, TrellisLinear — chosen per module by the config's
+ *  quantization entry (`mode: "trellis"`). */
+export type MlpLinear = QuantizedLinear | TrellisLinear;
+function loadMlpLinear(weights: Weights, path: string, config: ModelConfig): MlpLinear {
+  if (!TrellisLinear.isTrellis(config, path)) return QuantizedLinear.load(weights, path, config);
+  // M4 Pro measurements support the 27B down projection at M3/4. The
+  // operation also checks variant, dtype and packed-code layout eligibility.
+  const useSharedScatterCodebook = config.text.hiddenSize === 5120 &&
+    config.text.intermediateSize === 17408 && path.endsWith(".down_proj");
+  return TrellisLinear.load(weights, path, config, useSharedScatterCodebook);
+}
+
 export class Qwen3MLP {
-  readonly gate: QuantizedLinear;
-  readonly up: QuantizedLinear;
-  readonly down: QuantizedLinear;
+  readonly gate: MlpLinear;
+  readonly up: MlpLinear;
+  readonly down: MlpLinear;
 
   constructor(weights: Weights, config: ModelConfig, prefix: string) {
-    this.gate = QuantizedLinear.load(weights, `${prefix}.gate_proj`, config);
-    this.up = QuantizedLinear.load(weights, `${prefix}.up_proj`, config);
-    this.down = QuantizedLinear.load(weights, `${prefix}.down_proj`, config);
+    this.gate = loadMlpLinear(weights, `${prefix}.gate_proj`, config);
+    this.up = loadMlpLinear(weights, `${prefix}.up_proj`, config);
+    this.down = loadMlpLinear(weights, `${prefix}.down_proj`, config);
   }
 
-  forward(x: MlxArray): MlxArray {
-    const g = this.gate.forward(x);
-    const u = this.up.forward(x);
+  forward(x: MlxArray, inputRowContiguous = false, independentRows = false): MlxArray {
+    // Packed-trellis decode (M ≤ 4): gate, up and the swiglu in ONE kernel —
+    // x read once, no gate/up vectors materialized. This Lab path retains
+    // the packed activation arithmetic documented in turboquant.md.
+    if (this.gate instanceof TrellisLinear && this.up instanceof TrellisLinear &&
+        fusedGateUpEligible(this.gate, this.up) &&
+        x.shape.slice(0, -1).reduce((a, b) => a * b, 1) <= TRELLIS_MATVEC_MAX_M) {
+      const hidden = fusedGateUpSwiglu(x, this.gate, this.up);
+      const out = this.down instanceof QuantizedLinear ? this.down.forward(hidden, independentRows) : this.down.forward(hidden);
+      hidden.dispose();
+      return out;
+    }
+    const g = this.gate instanceof TrellisLinear ? this.gate.forward(x, inputRowContiguous) : this.gate.forward(x, independentRows);
+    const u = this.up instanceof TrellisLinear ? this.up.forward(x, inputRowContiguous) : this.up.forward(x, independentRows);
     // Oracle: down_proj(swiglu(gate_proj(x), up_proj(x))) — swiglu ALWAYS compiled
     // (mlx-lm's @mx.compile swiglu; no unfused silu+mul, matches its kernel set).
     const hidden = compiledSwiglu(g, u);
     g.dispose();
     u.dispose();
-    const out = this.down.forward(hidden);
+    const out = this.down instanceof QuantizedLinear ? this.down.forward(hidden, independentRows) : this.down.forward(hidden);
     hidden.dispose();
     return out;
   }
@@ -489,16 +559,18 @@ export class Qwen3Layer {
     this.postAttnNorm = new RMSNorm(weights.tensor(`${prefix}.post_attention_layernorm.weight`), config.text.rmsNormEps);
   }
 
-  forward(x: MlxArray, faMask: Mask, cache: Cache): MlxArray {
+  forward(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false, ssmMask?: MlxArray | null): MlxArray {
     const xn = this.inputNorm.forward(x);
     const r = this.isLinear
-      ? this.linearAttn!.forward(xn, cache as SSMCache)
-      : this.selfAttn!.forward(xn, faMask, cache);
+      ? this.linearAttn!.forward(xn, cache as SSMCache, independentRows, ssmMask)
+      : this.selfAttn!.forward(xn, faMask, cache, independentRows);
     xn.dispose();
     const h = ops.add(x, r);
     r.dispose();
     const hn = this.postAttnNorm.forward(h);
-    const m = this.mlp.forward(hn);
+    // RMSNorm allocates an aligned row-contiguous output. Pass that layout
+    // proof so Trellis can match native small-prefill matmul without an eval.
+    const m = this.mlp.forward(hn, true, independentRows);
     hn.dispose();
     const out = ops.add(h, m);
     h.dispose();
@@ -522,6 +594,14 @@ export class Qwen35Model {
   readonly faIdx: number;
 
   constructor(weights: Weights, config: ModelConfig) {
+    // Checkpoint-generation normalization BEFORE the first tensor() call: the
+    // 5.8-family export names the trunk `model.language_model.*` with a
+    // top-level `lm_head`, ships `mtp.*` in-repo, and stores RMSNorm gains as
+    // γ−1 with HF-layout conv1d. The view maps all of that onto the names and
+    // values this graph reads (mlx-lm's post-sanitize space). No-op — nothing
+    // installed at all — for artifacts already in that space.
+    const view = qwen35WeightsView(weights);
+    if (view) weights.setView(view);
     this.config = config;
     this.tied = config.text.tieWordEmbeddings;
     this.weightsBytes = [...weights.shards.files.values()]
@@ -555,9 +635,10 @@ export class Qwen35Model {
         out.set(`${p}.linear_attn.in_proj_a`, l.linearAttn.inProjA);
         out.set(`${p}.linear_attn.out_proj`, l.linearAttn.outProj);
       }
-      out.set(`${p}.mlp.gate_proj`, l.mlp.gate);
-      out.set(`${p}.mlp.up_proj`, l.mlp.up);
-      out.set(`${p}.mlp.down_proj`, l.mlp.down);
+      // Trellis-coded MLP tensors are not LoRA targets (no adapter seam yet).
+      if (l.mlp.gate instanceof QuantizedLinear) out.set(`${p}.mlp.gate_proj`, l.mlp.gate);
+      if (l.mlp.up instanceof QuantizedLinear) out.set(`${p}.mlp.up_proj`, l.mlp.up);
+      if (l.mlp.down instanceof QuantizedLinear) out.set(`${p}.mlp.down_proj`, l.mlp.down);
     }
     return out;
   }
@@ -566,6 +647,38 @@ export class Qwen35Model {
     return this.layers.map((l) => (l.isLinear ? new SSMCache() : new KVCache()));
   }
 
+  /** The 27B append path keeps the native M1 affine arithmetic while Trellis
+   * shares up to four rows. Other shapes retain single-token execution. */
+  createAppend(policy: { hasAdapters: boolean; pagedKv: boolean }) {
+    const t = this.config.text;
+    if (policy.hasAdapters || policy.pagedKv || this.loraState.active.length || this.mrope ||
+        t.hiddenSize !== 5120 || t.intermediateSize !== 17408 || this.layers.length !== 64 ||
+        t.headDim !== 256 || t.numAttentionHeads !== 24 || t.numKeyValueHeads !== 4)
+      return null;
+    if (deviceArchitecture() !== "applegpu_g16s") return null;
+    const qualified = (linear: QuantizedLinear | TrellisLinear) => {
+      if (linear instanceof TrellisLinear) return true;
+      const { mode, bits, groupSize } = linear.spec;
+      return mode === "affine" && [2, 3, 4, 6, 8].includes(bits) && [32, 64, 128].includes(groupSize);
+    };
+    for (const layer of this.layers) {
+      const a = layer.linearAttn, s = layer.selfAttn;
+      const attention = a ? [a.inProjQkv, a.inProjZ, a.inProjB, a.inProjA, a.outProj]
+        : [s!.qProj, s!.kProj, s!.vProj, s!.oProj];
+      if (!attention.every(qualified) || ![layer.mlp.gate, layer.mlp.up, layer.mlp.down].every(qualified))
+        return null;
+    }
+    return {
+      maxChunkSize: (state: readonly Cache[]) => qwenAppendChunkSize(state[0]!.offset),
+      forwardHidden: (ids: MlxArray, cache: Cache[]): MlxArray => {
+        if (ids.shape.length !== 2 || ids.shape[0] !== 1 || ids.shape[1]! > 4)
+          throw new Error("Qwen committed-token append supports one row and at most four positions");
+        const hidden = this.embed.encode(ids);
+        return this.forwardLayers(hidden, cache, true);
+      },
+    };
+  }
+
   forwardHidden(ids: MlxArray, cache: Cache[]): MlxArray {
     const h = this.embed.encode(ids);
     return this.forwardLayers(h, cache);
@@ -609,11 +722,12 @@ export class Qwen35Model {
     tap.captured.set(i, copy);
   }
 
-  protected forwardLayers(h0: MlxArray, cache: Cache[]): MlxArray {
+  protected forwardLayers(h0: MlxArray, cache: Cache[], independentRows = false): MlxArray {
     const L = h0.shape[1]!;
     // One full-attention mask shared by all full layers (same offset); linear
     // layers see no ssm mask at B=1.
     const faMask = cache[this.faIdx]!.makeMask(L, null);
+    using ssmMask = (cache[0] as SSMCache).prefillPadding?.makeMask(L) ?? null;
     // Vision requests: one interleaved-mRoPE cos/sin table per forward,
     // shared by all 12 full-attention layers (positions are layer-invariant).
     let mropeFwd: ReturnType<typeof buildMropePositions> | null = null;
@@ -634,9 +748,19 @@ export class Qwen35Model {
     try {
       for (let i = 0; i < this.layers.length; i++) {
         const tl = prof ? performance.now() : 0;
-        const next = this.layers[i]!.forward(h, faMask, cache[i]!);
+        const next = this.layers[i]!.forward(h, faMask, cache[i]!, independentRows, ssmMask);
         h.dispose();
         h = next;
+        if (L > TRELLIS_MATVEC_MAX_M) {
+          // Affine prefills also need bounded evaluation: a retained prefix
+          // history can leave too little room for the entire deferred graph.
+          // Materialize the cache outputs too. In particular, the tiny
+          // copied conv tail is not a dependency of `h`; leaving it lazy
+          // pins the whole prefill conv buffer until the chunk ends.
+          const state = cache[i]!.state();
+          try { ops.evalAll([h, ...state]); }
+          finally { if (cache[i]!.stateNeedsDispose) for (const a of state) a.dispose(); }
+        }
         if (prof) {
           ops.evalAll([h]);
           const key = this.layers[i]!.isLinear ? "layerLin" : "layerFull";