@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/models/gemma4/serving-blocks.ts

@@ -0,0 +1,72 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache } from "../../contracts/mlx/cache";
+import { CompiledGegluMLP, PlainGegluMLP } from "../../layers/geglu";
+import { disposing } from "../../layers/helpers";
+import { QuantizedLinearPlain } from "../../layers/quantized-linear";
+import { DecoderLayer, FixedExperts, type DonorRead, type Gemma4Phase, type Router } from "./blocks";
+import { isExpertTracing, recordRouting } from "../../runtime/expert-trace";
+
+/** Serving operations borrow the same projections that training and LoRA mount. */
+export class ServingRouter {
+  readonly projection: QuantizedLinearPlain;
+  constructor(readonly base: Router) { this.projection = new QuantizedLinearPlain(base.proj); }
+  forward(x: MlxArray): { indices: MlxArray; weights: MlxArray } {
+    using normed = ops.rmsNorm(x, this.base.normWeight, this.base.eps);
+    using scores = this.projection.forward(normed);
+    using part = ops.argpartitionAxis(scores, this.base.numExperts - this.base.topK, -1);
+    const [B, L] = scores.shape;
+    const indices = part.slice([0, 0, this.base.numExperts - this.base.topK], [B!, L!, this.base.numExperts]);
+    let weights = ops.takeAlongAxis(scores, indices, -1);
+    weights = disposing(weights, ops.softmaxAxis(weights, -1, false));
+    using gathered = ops.takeAxis(this.base.perExpertScale, indices, 0);
+    weights = disposing(weights, ops.mul(weights, gathered));
+    if (isExpertTracing()) recordRouting(this.base.layerIdx, indices);
+    return { indices, weights };
+  }
+}
+
+export class ServingDenseLayer {
+  readonly finish: (hidden: MlxArray, input: MlxArray | null) => MlxArray;
+  constructor(readonly base: DecoderLayer, readonly mlp: CompiledGegluMLP | PlainGegluMLP) {
+    // Artifact features are composed once; the request executes these operations.
+    const input = base.perLayerGate ? (h: MlxArray, p: MlxArray | null) => base.perLayerGated(h, p!) : (h: MlxArray) => h;
+    const scale = base.layerScalar ? (h: MlxArray) => base.scaled(h) : (h: MlxArray) => h;
+    this.finish = (h, p) => scale(input(h, p));
+  }
+  forwardAttn(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, donor: DonorRead | null) { return this.base.forwardAttn(x, cache, phase, donor); }
+  forwardMlp(h: MlxArray, input: MlxArray | null): MlxArray {
+    let f = this.base.preFfNorm.forward(h);
+    f = disposing(f, this.mlp.forward(f));
+    f = disposing(f, this.base.postFfNorm.forward(f));
+    const out = ops.add(h, f); f.dispose();
+    return this.finish(out, input);
+  }
+  forward(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, donor: DonorRead | null, input: MlxArray | null) {
+    const result = this.forwardAttn(x, cache, phase, donor);
+    try { return {h:this.forwardMlp(result.h,input),donor:result.donor}; }
+    catch(error) { result.donor?.dispose(); throw error; }
+    finally { result.h.dispose(); }
+  }
+}
+
+export class ServingMoeLayer extends ServingDenseLayer {
+  readonly router: ServingRouter;
+  constructor(base: DecoderLayer, mlp: CompiledGegluMLP, readonly experts: FixedExperts) {
+    super(base, mlp); this.router = new ServingRouter(base.router!);
+  }
+  override forwardMlp(h: MlxArray, input: MlxArray | null): MlxArray {
+    let dense = this.base.preFfNorm.forward(h);
+    dense = disposing(dense, this.mlp.forward(dense));
+    dense = disposing(dense, this.base.postFfNorm1!.forward(dense));
+    const {indices,weights} = this.router.forward(h);
+    let expert = this.base.preFfNorm2!.forward(h);
+    expert = disposing(expert, this.experts.forward(expert,indices,weights));
+    expert = disposing(expert, this.base.postFfNorm2!.forward(expert));
+    indices.dispose(); weights.dispose();
+    let f = ops.add(dense,expert); dense.dispose(); expert.dispose();
+    f = disposing(f,this.base.postFfNorm.forward(f));
+    const out = ops.add(h,f); f.dispose();
+    return this.finish(out,input);
+  }
+}

package/src/models/gemma4/trainable.ts

@@ -0,0 +1,186 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { Checkpoint } from "@mlx-bun/mlx/checkpoint";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache, SharedKv } from "../../contracts/mlx/cache";
+import type {
+  GradCheckpointRun, HeadQuant, LoraLeaves, PrefixLayout, SegmentedPass, TrainableGraph,
+} from "../../contracts/mlx/trainable";
+import { AttentionMasks } from "../../state/attention-read";
+import { prefixSharedCaches, selectedTrainingAttention, TrainingCache } from "../../state/training-cache";
+import type { LoraWeights } from "../../layers/lora";
+import type { RMSNorm } from "../../layers/normalization";
+import type { QuantizedEmbedding } from "../../layers/quantized-embedding";
+
+/** Optional gradient-checkpointing context for the training forward. When set,
+ *  the training forward wraps each layer in a Checkpoint so its interior activations
+ *  are recomputed in the backward pass. `byLayer` maps a layer index to that
+ *  layer's trainable LoRA weights, partitioned into the attention sub-block
+ *  (q/k/v/o_proj) and the MLP sub-block (gate/up/down_proj + per-layer gate/
+ *  projection) — re-swapped inside the closure so they are explicit checkpoint
+ *  inputs (required because the autograd primals are disposed before the
+ *  recompute). When `splitMlp` is set, the layer is wrapped in TWO checkpoints
+ *  (attn sub-block + MLP sub-block) with the post-attn residual `hMid` as the
+ *  boundary, so the backward recompute holds `max(attn,MLP)+hMid` instead of the
+ *  whole layer (the intra-layer MLP split, lever 4). `keepAlive` collects the
+ *  Checkpoint objects; the trainer disposes them after value_and_grad. */
+export interface LayerLoras {
+  attn: LoraWeights[];
+  mlp: LoraWeights[];
+}
+export interface GradCheckpointCtx {
+  byLayer: Map<number, LayerLoras>;
+  splitMlp: boolean;
+  keepAlive: Checkpoint[];
+}
+
+/** What the training declaration reads from a Gemma 4 model (structural, so this
+ *  module does not depend on the model file). */
+export interface Gemma4Training {
+  readonly config: { readonly text: { readonly finalLogitSoftcapping: number | null } };
+  readonly embed: QuantizedEmbedding;
+  readonly finalNorm: RMSNorm;
+  readonly layers: readonly { readonly layerType: string }[];
+  /** The sliding layers' window. */
+  readonly windowSize: number;
+  /** Donor count: layers [0, numDonors) own caches; the rest share. */
+  readonly numDonors: number;
+  readonly previousKvs: readonly number[];
+  readonly cacheIndex: readonly number[];
+  readonly reusedDonors: Set<number>;
+  gradCkpt: GradCheckpointCtx | null;
+  /** The training forward over one training cache per donor layer, which owns
+   * its read; donors hand their sharers plain keys and values. */
+  forwardTrainingHidden(ids: MlxArray, cache: Cache[]): MlxArray;
+  embedForSegmented(ids: MlxArray): { hScaled: MlxArray; perLayer: MlxArray | null };
+  runLayerRange(
+    h: MlxArray, aIdx: number, bIdx: number, cache: Cache[],
+    perLayer: MlxArray | null, donorKvIn: Map<number, SharedKv>,
+  ): { h: MlxArray; donorKvOut: Map<number, SharedKv> };
+}
+
+/** One stateless training cache per donor layer, with that layer's sliding
+ *  window, sharing one mask memo and attending with the run's training
+ *  attention: prefix-shared caches under a layout, plain ones otherwise. */
+function trainingCaches(model: Gemma4Training, prefix?: PrefixLayout): Cache[] {
+  const attention = selectedTrainingAttention();
+  const windows = model.layers.slice(0, model.numDonors)
+    .map(layer => layer.layerType === "sliding_attention" ? model.windowSize : null);
+  if (prefix) return prefixSharedCaches(prefix, windows, attention);
+  const masks = new AttentionMasks();
+  return windows.map(window => new TrainingCache(window, masks, attention));
+}
+
+/** Copy `a` into a graph-free leaf (an evaluated, row-major copy), disposing `a`. */
+function detached(a: MlxArray): MlxArray {
+  const c = ops.contiguous(a);
+  const leaf = c.detachCopy();
+  c.dispose();
+  a.dispose();
+  return leaf;
+}
+
+/** One segmented pass over a Gemma 4 sequence: the per-layer-input tensor (a
+ *  constant boundary sliced per layer) and one stateless training cache per
+ *  donor layer (block-sparse with a logical-position window when
+ *  prefix-shared), which builds its masks once for the pass; sharers read the
+ *  donors' K/V through the donors' caches. */
+class Gemma4SegmentedPass implements SegmentedPass {
+  readonly input: MlxArray;
+  private readonly caches: Cache[];
+  private readonly perLayer: MlxArray | null = null;
+
+  constructor(private readonly model: Gemma4Training, ids: MlxArray, prefix?: PrefixLayout) {
+    this.caches = trainingCaches(model, prefix);
+    let input: MlxArray | null = null;
+    let perLayer: MlxArray | null = null;
+    try {
+      const embedded = model.embedForSegmented(ids);
+      input = embedded.hScaled;
+      perLayer = embedded.perLayer ? detached(embedded.perLayer) : null;
+    } catch (error) {
+      input?.dispose();
+      perLayer?.dispose();
+      this.releaseCaches();
+      throw error;
+    }
+    this.input = input;
+    this.perLayer = perLayer;
+  }
+
+  private releaseCaches(): void {
+    for (const c of this.caches) c.dispose();
+  }
+
+  runRange(h: MlxArray, lo: number, hi: number, donorKvIn: Map<number, SharedKv>) {
+    return this.model.runLayerRange(h, lo, hi, this.caches, this.perLayer, donorKvIn);
+  }
+
+  dispose(): void {
+    this.releaseCaches();
+    this.perLayer?.dispose();
+  }
+}
+
+function checkpointRun(model: Gemma4Training, targets: readonly { modulePath: string; lw: LoraLeaves }[], splitMlp: boolean): GradCheckpointRun {
+  // Partition each layer's LoRA into the attention sub-block (self_attn.*) and
+  // the MLP sub-block (everything else in the layer: gate/up/down_proj +
+  // per_layer_*) so the split-MLP checkpoint can wrap them independently.
+  const byLayer: GradCheckpointCtx["byLayer"] = new Map();
+  for (const t of targets) {
+    const m = t.modulePath.match(/\.layers\.(\d+)\./);
+    if (!m) continue;
+    const li = Number(m[1]);
+    const ll = byLayer.get(li) ?? byLayer.set(li, { attn: [], mlp: [] }).get(li)!;
+    (t.modulePath.includes(".self_attn.") ? ll.attn : ll.mlp).push(t.lw);
+  }
+  // Checkpoints are created per forward and disposed after each step's backward.
+  const ctx: GradCheckpointCtx = { byLayer, splitMlp, keepAlive: [] };
+  model.gradCkpt = ctx;
+  return {
+    releaseStep() {
+      for (const ck of ctx.keepAlive) ck.dispose();
+      ctx.keepAlive.length = 0;
+    },
+    suspend() { model.gradCkpt = null; },
+    resume() { model.gradCkpt = ctx; },
+    end() {
+      model.gradCkpt = null;
+      for (const ck of ctx.keepAlive) ck.dispose();
+    },
+  };
+}
+
+/** Gemma 4's training declaration. */
+export function gemma4Trainable(model: Gemma4Training): TrainableGraph {
+  const donorOf = model.layers.map((_, i) => (model.cacheIndex[i] === -1 ? model.previousKvs[i]! : null));
+  return {
+    lmHead(): HeadQuant {
+      // Tied: the quantized embedding is the head; the softcap comes from the config.
+      const e = model.embed;
+      return { w: e.w, scales: e.scales, biases: e.biases, spec: e.spec, softcap: model.config.text.finalLogitSoftcapping };
+    },
+    flashAttention: {
+      supported: false,
+      reason: "e4b SIGTRAPs on this path at seq >= 2048 and it has not been re-validated at that scale since the kernel fixes",
+    },
+    segmented: {
+      layerCount: model.layers.length,
+      sharedKv: { reusedDonors: model.reusedDonors, donorOf },
+      finalNorm: (h) => model.finalNorm.forward(h),
+      begin: (ids, prefix) => new Gemma4SegmentedPass(model, ids, prefix),
+    },
+    prefixShared: {
+      prefixSharedHidden(ids: MlxArray, layout: PrefixLayout): MlxArray {
+        // One prefix cache per DONOR layer (sharers read donors' K/V through
+        // the donors' caches); the caches read under the block-sparse masks.
+        const caches = trainingCaches(model, layout);
+        try {
+          return model.forwardTrainingHidden(ids, caches); // [1, T, hidden], post-finalNorm
+        } finally {
+          for (const c of caches) c.dispose();
+        }
+      },
+    },
+    gradCheckpoint: { enable: (targets, splitMlp) => checkpointRun(model, targets, splitMlp) },
+  };
+}

package/src/models/gemma4/written-blocks.ts

@@ -0,0 +1,232 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { ModelConfig } from "../../artifacts/config";
+import type { Weights } from "../../artifacts/weights";
+import type { Cache } from "../../contracts/mlx/cache";
+import { RMSNorm } from "../../layers/normalization";
+import { QuantizedLinear, QuantizedLinearPlain } from "../../layers/quantized-linear";
+import { GegluMLP, PlainGegluMLP, CompiledGegluMLP, CompiledGeglu } from "../../layers/geglu";
+import { Attention, Router, SwitchGLU, SortedSwitchGLU, UnsortedSwitchGLU, FixedExperts, type DonorRead, type Gemma4Phase } from "./blocks";
+import { disposing } from "../../layers/helpers";
+
+export class SlidingAttention extends Attention {
+  constructor(weights: Weights, config: ModelConfig, prefix: string) { super(weights, config, prefix, "sliding_attention", true, undefined, false); }
+  override forward(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, _donor: DonorRead | null) { return this.forwardSeparate(x, cache!, phase); }
+}
+export class FullAttention extends Attention {
+  constructor(weights: Weights, config: ModelConfig, prefix: string) { super(weights, config, prefix, "full_attention", true, undefined, false); }
+  override forward(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, _donor: DonorRead | null) { return this.forwardSeparate(x, cache!, phase); }
+}
+export class FullKEqVAttention extends Attention {
+  constructor(weights: Weights, config: ModelConfig, prefix: string) { super(weights, config, prefix, "full_attention", true, undefined, true); }
+  override forward(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, _donor: DonorRead | null) { return this.forwardKEqV(x, cache!, phase); }
+}
+export class SharedSlidingAttention extends Attention {
+  constructor(weights: Weights, config: ModelConfig, prefix: string) { super(weights, config, prefix, "sliding_attention", false, undefined, false); }
+  override forward(x: MlxArray, _cache: Cache | null, _phase: Gemma4Phase, donor: DonorRead | null) { return this.forwardShared(x, donor!); }
+}
+export class SharedFullAttention extends Attention {
+  constructor(weights: Weights, config: ModelConfig, prefix: string) { super(weights, config, prefix, "full_attention", false, undefined, false); }
+  override forward(x: MlxArray, _cache: Cache | null, _phase: Gemma4Phase, donor: DonorRead | null) { return this.forwardShared(x, donor!); }
+}
+
+/** Stored MLP projections. Each operation view names its activation literally. */
+export class GemmaMlpWeights extends GegluMLP {
+  constructor(weights: Weights, config: ModelConfig, prefix: string) { super(weights, config, prefix, true); }
+}
+
+/** Norms and scale owned by one written layer; model features are supplied as
+ * concrete attention and feed-forward operations by the file. */
+export class GemmaLayerWeights {
+  readonly inputNorm: RMSNorm;
+  readonly postAttnNorm: RMSNorm;
+  readonly preFfNorm: RMSNorm;
+  readonly postFfNorm: RMSNorm;
+  readonly layerScalar: MlxArray;
+  readonly layerType: string;
+  readonly mlp: GemmaMlpWeights;
+  constructor(weights: Weights, config: ModelConfig, readonly prefix: string, readonly attn: Attention) {
+    const norm = (name: string) => new RMSNorm(weights.tensor(`${prefix}.${name}.weight`), config.text.rmsNormEps);
+    this.inputNorm = norm("input_layernorm"); this.postAttnNorm = norm("post_attention_layernorm");
+    this.preFfNorm = norm("pre_feedforward_layernorm"); this.postFfNorm = norm("post_feedforward_layernorm");
+    this.layerScalar = weights.tensor(`${prefix}.layer_scalar`);
+    this.layerType = attn.isSliding ? "sliding_attention" : "full_attention";
+    this.mlp = new GemmaMlpWeights(weights, config, `${prefix}.mlp`);
+  }
+  forwardAttn(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, donor: DonorRead | null): { h: MlxArray; donor: DonorRead | null } {
+    const normed = this.inputNorm.forward(x);
+    let result: { out: MlxArray; donor: DonorRead | null };
+    try { result = this.attn.forward(normed, cache, phase, donor); } finally { normed.dispose(); }
+    let hidden = result.out;
+    try {
+      hidden = disposing(hidden, this.postAttnNorm.forward(hidden));
+      hidden = disposing(hidden, ops.add(x, hidden));
+      return { h: hidden, donor: result.donor };
+    } catch (error) { hidden.dispose(); result.donor?.dispose(); throw error; }
+  }
+  scaled(x: MlxArray): MlxArray { return disposing(x, ops.mul(x, this.layerScalar)); }
+}
+
+export class GemmaInputLayerWeights extends GemmaLayerWeights {
+  readonly perLayerGate: QuantizedLinear;
+  readonly perLayerProjection: QuantizedLinear;
+  readonly postPerLayerNorm: RMSNorm;
+  readonly #gate: QuantizedLinearPlain;
+  readonly #projection: QuantizedLinearPlain;
+  constructor(weights: Weights, config: ModelConfig, prefix: string, attn: Attention) {
+    super(weights, config, prefix, attn);
+    this.perLayerGate = QuantizedLinear.load(weights, `${prefix}.per_layer_input_gate`, config);
+    this.perLayerProjection = QuantizedLinear.load(weights, `${prefix}.per_layer_projection`, config);
+    this.postPerLayerNorm = new RMSNorm(weights.tensor(`${prefix}.post_per_layer_input_norm.weight`), config.text.rmsNormEps);
+    this.#gate = new QuantizedLinearPlain(this.perLayerGate); this.#projection = new QuantizedLinearPlain(this.perLayerProjection);
+  }
+  perLayerGated(h: MlxArray, input: MlxArray): MlxArray {
+    let gate = this.#gate.forward(h);
+    gate = disposing(gate, ops.geluApprox(gate));
+    gate = disposing(gate, ops.mul(gate, input));
+    gate = disposing(gate, this.#projection.forward(gate));
+    gate = disposing(gate, this.postPerLayerNorm.forward(gate));
+    const out = ops.add(h, gate); h.dispose(); gate.dispose(); return out;
+  }
+}
+
+export class GemmaMoeLayerWeights extends GemmaLayerWeights {
+  readonly router: Router;
+  readonly experts: SwitchGLU;
+  readonly postFfNorm1: RMSNorm;
+  readonly postFfNorm2: RMSNorm;
+  readonly preFfNorm2: RMSNorm;
+  constructor(weights: Weights, config: ModelConfig, prefix: string, attn: Attention) {
+    super(weights, config, prefix, attn);
+    this.router = new PlainRouter(weights, config, `${prefix}.router`);
+    this.experts = new SwitchGLU(weights, config, `${prefix}.experts.switch_glu`, true);
+    const norm = (name: string) => new RMSNorm(weights.tensor(`${prefix}.${name}.weight`), config.text.rmsNormEps);
+    this.postFfNorm1 = norm("post_feedforward_layernorm_1"); this.postFfNorm2 = norm("post_feedforward_layernorm_2");
+    this.preFfNorm2 = norm("pre_feedforward_layernorm_2");
+  }
+}
+
+export class GemmaDenseLayer {
+  constructor(readonly weights: GemmaLayerWeights, readonly mlp: PlainGegluMLP | CompiledGegluMLP) {}
+  get attn(): Attention { return this.weights.attn; }
+  get layerType(): string { return this.weights.layerType; }
+  forwardAttn(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, donor: DonorRead | null) { return this.weights.forwardAttn(x, cache, phase, donor); }
+  feedforward(h: MlxArray): MlxArray {
+    let feedforward = this.weights.preFfNorm.forward(h);
+    feedforward = disposing(feedforward, this.mlp.forward(feedforward));
+    feedforward = disposing(feedforward, this.weights.postFfNorm.forward(feedforward));
+    const output = ops.add(h, feedforward); feedforward.dispose();
+    return this.weights.scaled(output);
+  }
+  forward(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, donor: DonorRead | null, _input: MlxArray | null = null) {
+    const { h, donor: made } = this.forwardAttn(x, cache, phase, donor);
+    try { return { h: this.feedforward(h), donor: made }; }
+    catch (error) { made?.dispose(); throw error; } finally { h.dispose(); }
+  }
+}
+
+export class GemmaInputLayer {
+  constructor(readonly weights: GemmaInputLayerWeights, readonly mlp: PlainGegluMLP | CompiledGegluMLP) {}
+  get attn(): Attention { return this.weights.attn; }
+  get layerType(): string { return this.weights.layerType; }
+  forwardAttn(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, donor: DonorRead | null) { return this.weights.forwardAttn(x, cache, phase, donor); }
+  feedforward(h: MlxArray, input: MlxArray): MlxArray {
+    let feedforward = this.weights.preFfNorm.forward(h);
+    feedforward = disposing(feedforward, this.mlp.forward(feedforward));
+    feedforward = disposing(feedforward, this.weights.postFfNorm.forward(feedforward));
+    const output = ops.add(h, feedforward); feedforward.dispose();
+    return this.weights.scaled(this.weights.perLayerGated(output, input));
+  }
+  forward(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, donor: DonorRead | null, input: MlxArray | null) {
+    const { h, donor: made } = this.forwardAttn(x, cache, phase, donor);
+    try { return { h: this.feedforward(h, input!), donor: made }; }
+    catch (error) { made?.dispose(); throw error; } finally { h.dispose(); }
+  }
+}
+
+export class GemmaMoeLayer {
+  constructor(readonly weights: GemmaMoeLayerWeights, readonly mlp: CompiledGegluMLP, readonly experts: FixedExperts) {}
+  get attn(): Attention { return this.weights.attn; }
+  get layerType(): string { return this.weights.layerType; }
+  forwardAttn(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, donor: DonorRead | null) { return this.weights.forwardAttn(x, cache, phase, donor); }
+  feedforward(h: MlxArray): MlxArray {
+    let dense = this.weights.preFfNorm.forward(h);
+    dense = disposing(dense, this.mlp.forward(dense));
+    dense = disposing(dense, this.weights.postFfNorm1.forward(dense));
+    const { indices, weights } = this.weights.router.forward(h);
+    let expert = this.weights.preFfNorm2.forward(h);
+    expert = disposing(expert, this.experts.forward(expert, indices, weights));
+    expert = disposing(expert, this.weights.postFfNorm2.forward(expert));
+    indices.dispose(); weights.dispose();
+    let feedforward = ops.add(dense, expert); dense.dispose(); expert.dispose();
+    feedforward = disposing(feedforward, this.weights.postFfNorm.forward(feedforward));
+    const out = ops.add(h, feedforward); feedforward.dispose(); return this.weights.scaled(out);
+  }
+  forward(x: MlxArray, cache: Cache | null, phase: Gemma4Phase, donor: DonorRead | null, _input: MlxArray | null = null) {
+    const { h, donor: made } = this.forwardAttn(x, cache, phase, donor);
+    try { return { h: this.feedforward(h), donor: made }; }
+    catch (error) { made?.dispose(); throw error; } finally { h.dispose(); }
+  }
+}
+
+export { PlainGegluMLP, CompiledGegluMLP, CompiledGeglu, SortedSwitchGLU, UnsortedSwitchGLU, FixedExperts };
+
+/** Ordinary layer-input slices; trace slices have their own named layer. */
+export class FixedPerLayerInputs extends PerLayerInputs {
+  readonly #projection: QuantizedLinearPlain;
+  readonly #layerCount: number;
+  readonly #hiddenSize: number;
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    super(weights, config, prefix);
+    this.#projection = new QuantizedLinearPlain(this.modelProjection);
+    this.#layerCount = config.text.numHiddenLayers; this.#hiddenSize = config.text.hiddenSize;
+  }
+  override compute(ids: MlxArray, hidden: MlxArray): MlxArray {
+    const [B, L] = ids.shape;
+    let inputs = this.embed.encode(ids);
+    inputs = disposing(inputs, ops.mulScalar(inputs, Math.sqrt(this.width)));
+    inputs = disposing(inputs, ops.reshape(inputs, [B!, L!, this.#layerCount, this.width]));
+    let projection = this.#projection.forward(hidden);
+    projection = disposing(projection, ops.mulScalar(projection, 1 / Math.sqrt(this.#hiddenSize)));
+    projection = disposing(projection, ops.reshape(projection, [B!, L!, this.#layerCount, this.width]));
+    projection = disposing(projection, this.projectionNorm.forward(projection));
+    const combined = ops.add(projection, inputs); projection.dispose(); inputs.dispose();
+    return disposing(combined, ops.mulScalar(combined, Math.SQRT1_2));
+  }
+  override layer(inputs: MlxArray, index: number, batch: number, length: number): MlxArray {
+    using slice = inputs.slice([0, 0, index, 0], [batch, length, index + 1, this.width]);
+    return ops.reshape(slice, [batch, length, this.width]);
+  }
+}
+export class TracePerLayerInputs extends FixedPerLayerInputs {
+  override layer(inputs: MlxArray, index: number, batch: number, length: number): MlxArray {
+    using start = ops.fromInt32([index], [1]);
+    using slice = ops.sliceDynamic(inputs, start, [2], [batch, length, 1, this.width]);
+    return ops.reshape(slice, [batch, length, this.width]);
+  }
+}
+
+import { PerLayerInputs } from "./blocks";
+
+/** Ordinary router slice, with no compiled-trace form selection. */
+export class PlainRouter extends Router {
+  readonly #projection: QuantizedLinearPlain;
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    super(weights, config, prefix); this.#projection = new QuantizedLinearPlain(this.proj);
+  }
+  override forward(x: MlxArray): { indices: MlxArray; weights: MlxArray } {
+    using normed = ops.rmsNorm(x, this.normWeight, this.eps);
+    using scores = this.#projection.forward(normed);
+    using part = ops.argpartitionAxis(scores, this.numExperts - this.topK, -1);
+    const [B, L] = scores.shape;
+    const indices = part.slice([0, 0, this.numExperts - this.topK], [B!, L!, this.numExperts]);
+    let weights = ops.takeAlongAxis(scores, indices, -1);
+    weights = disposing(weights, ops.softmaxAxis(weights, -1, false));
+    using gathered = ops.takeAxis(this.perExpertScale, indices, 0);
+    weights = disposing(weights, ops.mul(weights, gathered));
+    if (isExpertTracing()) recordRouting(this.layerIdx, indices);
+    return { indices, weights };
+  }
+}
+
+import { isExpertTracing, recordRouting } from "../../runtime/expert-trace";

package/src/models/gemma4/written-forward.ts

@@ -0,0 +1,159 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { MlxLayerCapture } from "../../contracts/mlx/graph";
+import type { Cache } from "../../contracts/mlx/cache";
+import { GraphForwardScope } from "../../runtime/graph-forward";
+import type { DonorRead, PerLayerInputs, Gemma4Phase } from "./blocks";
+import type { GemmaLayerWeights, GemmaInputLayerWeights, GemmaMoeLayerWeights } from "./written-blocks";
+import type { QuantizedLinear } from "../../layers/quantized-linear";
+
+export interface GemmaLayerResult { h: MlxArray; donor: DonorRead | null }
+
+/** Carries reads between explicitly written layer calls; owns every retained
+ * donor read and the current hidden handle during failures as well as success. */
+export class GemmaForwardScope extends GraphForwardScope {
+  readonly reads: (DonorRead | null)[];
+  constructor(hidden: MlxArray, capture: MlxLayerCapture | null, caches: number) {
+    super(hidden, capture); this.reads = Array(caches).fill(null);
+  }
+  own(index: number, result: GemmaLayerResult): void {
+    result.donor!.dispose(); this.accept(index, result.h);
+  }
+  donor(index: number, result: GemmaLayerResult): void {
+    this.reads[index] = result.donor; this.accept(index, result.h);
+  }
+  shared(index: number, result: GemmaLayerResult): void { this.accept(index, result.h); }
+  override [Symbol.dispose](): void {
+    super[Symbol.dispose](); for (const read of this.reads) read?.dispose();
+  }
+}
+
+/** Inputs are sliced one layer at a time in the same order as the graph calls. */
+export class GemmaInputForwardScope extends GemmaForwardScope {
+  readonly #inputs: MlxArray;
+  #current: MlxArray | null = null;
+  readonly #batch: number;
+  readonly #length: number;
+  constructor(hidden: MlxArray, ids: MlxArray, capture: MlxLayerCapture | null, caches: number,
+    readonly perLayer: PerLayerInputs) {
+    super(hidden, capture, caches);
+    [this.#batch, this.#length] = ids.shape as [number, number];
+    try { this.#inputs = perLayer.compute(ids, hidden); }
+    catch (error) { super[Symbol.dispose](); throw error; }
+  }
+  input(index: number): MlxArray {
+    this.#current?.dispose();
+    this.#current = this.perLayer.layer(this.#inputs, index, this.#batch, this.#length);
+    return this.#current;
+  }
+  override accept(index: number, next: MlxArray): void {
+    this.#current?.dispose(); this.#current = null; super.accept(index, next);
+  }
+  override [Symbol.dispose](): void {
+    this.#current?.dispose(); this.#inputs.dispose(); super[Symbol.dispose]();
+  }
+}
+
+/** Record only mounting targets; it never chooses a projection operation. */
+export function denseLoraTargets(parts: readonly GemmaLayerWeights[]): Map<string, QuantizedLinear> {
+  const targets = new Map<string, QuantizedLinear>();
+  for (const layer of parts) {
+    const add = (suffix: string, projection: QuantizedLinear | null) => { if (projection) targets.set(`${layer.prefix}.${suffix}`, projection); };
+    add("self_attn.q_proj", layer.attn.qProj); add("self_attn.k_proj", layer.attn.kProj);
+    add("self_attn.v_proj", layer.attn.vProj); add("self_attn.o_proj", layer.attn.oProj);
+    add("mlp.gate_proj", layer.mlp.gate); add("mlp.up_proj", layer.mlp.up); add("mlp.down_proj", layer.mlp.down);
+  }
+  return targets;
+}
+export function inputLoraTargets(parts: readonly GemmaInputLayerWeights[]): Map<string, QuantizedLinear> {
+  const targets = denseLoraTargets(parts);
+  for (const layer of parts) {
+    targets.set(`${layer.prefix}.per_layer_input_gate`, layer.perLayerGate);
+    targets.set(`${layer.prefix}.per_layer_projection`, layer.perLayerProjection);
+  }
+  return targets;
+}
+export interface GemmaMixedGroup {
+  readonly ids: MlxArray;
+  readonly cache: Cache[];
+  readonly phase: Gemma4Phase;
+  readonly reads: (DonorRead | null)[];
+  h: MlxArray;
+  inputs: MlxArray | null;
+}
+
+/** Request-data iteration for a mixed forward. The file calls each layer in
+ * order and supplies literal attention and feed-forward operations; this scope
+ * only packs that layer's token rows and carries each group's owned state. */
+export class GemmaMixedScope {
+  readonly groups: GemmaMixedGroup[] = [];
+  readonly #results: MlxArray[] = [];
+  constructor(readonly work: readonly { ids: MlxArray; state: Cache[]; phase: "decode" | "window" | "verify"; capture: MlxLayerCapture | null }[],
+    embed: (ids: MlxArray) => MlxArray, readonly caches: number, readonly perLayer: PerLayerInputs | null = null) {
+    try {
+      for (const row of work) {
+        const group: GemmaMixedGroup = { ids: row.ids, cache: row.state, phase: row.phase === "decode" ? "decode" : "window",
+          reads: Array(caches).fill(null), h: embed(row.ids), inputs: null };
+        this.groups.push(group);
+        if (perLayer) group.inputs = perLayer.compute(row.ids, group.h);
+      }
+    } catch (error) { this[Symbol.dispose](); throw error; }
+  }
+  layer(index: number, attend: (x: MlxArray, group: GemmaMixedGroup) => MlxArray,
+    feedforward: (x: MlxArray, input: MlxArray | null) => MlxArray): void {
+    const mids: MlxArray[] = [], inputs: MlxArray[] = [];
+    try {
+      for (const group of this.groups) {
+        mids.push(attend(group.h, group));
+        if (this.perLayer) inputs.push(this.perLayer.layer(group.inputs!, index, group.ids.shape[0]!, group.ids.shape[1]!));
+      }
+      let packedInput: MlxArray | null = null;
+      if (inputs.length) {
+        const flat = inputs.map(input => ops.reshape(input, [1, -1, this.perLayer!.width]));
+        try { packedInput = ops.concatAxis(flat, 1); } finally { for (const input of flat) input.dispose(); }
+      }
+      using held = packedInput;
+      const outputs = mapPackedTokens(mids, packed => feedforward(packed, held));
+      for (const [row, group] of this.groups.entries()) { group.h.dispose(); group.h = outputs[row]!; }
+      for (const [row, group] of this.groups.entries()) this.work[row]!.capture?.(index, group.h);
+    } finally { for (const mid of mids) mid.dispose(); for (const input of inputs) input.dispose(); }
+  }
+  finish(layerCount: number, norm: { forward(x: MlxArray): MlxArray }): MlxArray[] {
+    try {
+      for (const [index, group] of this.groups.entries()) {
+        const hidden = norm.forward(group.h); this.#results.push(hidden);
+        this.work[index]!.capture?.(layerCount, hidden);
+      }
+      return this.#results.splice(0);
+    } catch (error) { for (const result of this.#results) result.dispose(); this.#results.length = 0; throw error; }
+  }
+  [Symbol.dispose](): void {
+    for (const group of this.groups) { group.h.dispose(); group.inputs?.dispose(); for (const read of group.reads) read?.dispose(); }
+  }
+}
+
+export function ownMixed(result: GemmaLayerResult): MlxArray { result.donor!.dispose(); return result.h; }
+export function donorMixed(group: GemmaMixedGroup, cache: number, result: GemmaLayerResult): MlxArray { group.reads[cache] = result.donor; return result.h; }
+export function sharedMixed(result: GemmaLayerResult): MlxArray { return result.h; }
+
+import { mapPackedTokens } from "../../input/token-groups";
+
+export function weightBytes(weights: { shards: { files: Map<string, { mmap: { size: number } }> } }): number {
+  return [...weights.shards.files.values()].reduce((total, file) => total + file.mmap.size, 0);
+}
+
+/** Construction ownership until the written graph transfers its constants. */
+export class GemmaBuildScope {
+  readonly #constants: MlxArray[] = [];
+  hold<T extends GemmaLayerWeights>(layer: T): T {
+    if (layer.attn.ropeFreqs) this.#constants.push(layer.attn.ropeFreqs);
+    return layer;
+  }
+  holdMoe(layer: GemmaMoeLayerWeights): GemmaMoeLayerWeights {
+    this.hold(layer); this.#constants.push(layer.router.normWeight); return layer;
+  }
+  finish() { return ownResource(this.#constants.splice(0), disposeResources); }
+  [Symbol.dispose](): void { disposeResources(this.#constants.splice(0)); }
+}
+
+import { ownResource, disposeResources } from "../../runtime/resources";