mlx-bun 0.0.9 → 0.0.11

package/src/model/qwen3-moe.ts

@@ -0,0 +1,404 @@
+// Concrete model graph for Qwen3-MoE (model_type `qwen3_moe`,
+// Qwen3MoeForCausalLM). A STRAIGHT op-for-op transcription of the oracle
+// (`/Users/joshrossi/Code/mlx-lm/.venv/.../mlx_lm/models/qwen3_moe.py` +
+// switch_layers.py SwitchGLU + activations.py swiglu). Each op below carries
+// the exact oracle line it copies — no abstraction, no reorder, no
+// "optimization". The ONLY intentional deviation from a naive transcription is
+// that the expert swiglu is mx.compile-COMPILED (like the oracle's
+// activations.py `@partial(mx.compile) def swiglu`), matching both the greedy
+// token stream AND the dispatched Metal kernel set.
+//
+// Attention is identical to plain qwen3 (per-head q/k RMSNorm over head_dim
+// before RoPE, full-head RoPE, GQA, SDPA). The MoE block is
+// Qwen3MoeSparseMoeBlock: gate Linear -> precise softmax over ALL experts ->
+// argpartition top-k -> take_along_axis -> renormalize (norm_topk_prob) ->
+// SwitchGLU(swiglu) experts -> weighted sum. Per decoder_sparse_step /
+// mlp_only_layers a layer may instead be a dense MLP (MLP class); both are
+// handled faithfully. Weights are untied (separate lm_head). Config: 48 layers,
+// hidden 2048, 128 experts, top-8, moe_intermediate 768, rope_theta 1e7.
+
+import type { ModelConfig } from "../config";
+import type { Weights } from "../weights";
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { CompiledFunction } from "../mlx/compile";
+import {
+  argmaxLastPosition,
+  disposing,
+  isCompiledTrace,
+  KVCache,
+  LoraState,
+  QuantizedEmbedding,
+  QuantizedLinear,
+  QuantizedSwitchLinear,
+  RMSNorm,
+  type Cache,
+  type Mask,
+} from "./gemma4-base";
+
+// ── activations.py swiglu (mx.compile) ───────────────────────────────────────
+// `@partial(mx.compile, shapeless=True) def swiglu(gate, x): return nn.silu(gate) * x`
+// nn.silu(g) == g * sigmoid(g); mx.compile fuses sigmoid + mul + mul into ONE
+// kernel. Traced once (shapeless), replayed thereafter. Autograd-safe (mx.compile
+// threads the VJP through the traced graph). Exported so the parity test can
+// assert the closure exists and the expert path actually uses it.
+let _swigluClosure: CompiledFunction | null = null;
+export function compiledSwiglu(gate: MlxArray, up: MlxArray): MlxArray {
+  if (!_swigluClosure) {
+    _swigluClosure = new CompiledFunction((inputs) => {
+      const g = inputs[0]!, u = inputs[1]!;              // nn.silu(gate) * x
+      const sig = ops.sigmoid(g);
+      const silu = ops.mul(g, sig); sig.dispose();
+      const out = ops.mul(silu, u); silu.dispose();
+      return [out];
+    });
+  }
+  return _swigluClosure.apply([gate, up])[0]!;
+}
+
+// ── qwen3_moe.py Attention (== plain qwen3) ──────────────────────────────────
+class Attention {
+  readonly qProj: QuantizedLinear;
+  readonly kProj: QuantizedLinear;
+  readonly vProj: QuantizedLinear;
+  readonly oProj: QuantizedLinear;
+  readonly qNorm: RMSNorm;
+  readonly kNorm: RMSNorm;
+  readonly nHeads: number;
+  readonly nKvHeads: number;
+  readonly headDim: number;
+  readonly scale: number;
+  readonly ropeBase: number;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    const t = config.text;
+    this.nHeads = t.numAttentionHeads;                                  // self.n_heads
+    this.nKvHeads = t.numKeyValueHeads;                                 // self.n_kv_heads
+    this.headDim = t.headDim;                                           // head_dim
+    this.scale = Math.pow(this.headDim, -0.5);                          // self.scale = head_dim**-0.5
+    this.ropeBase = t.ropeParameters.full_attention?.ropeTheta ?? 10000000; // args.rope_theta
+    this.qProj = QuantizedLinear.load(weights, `${prefix}.q_proj`, config); // nn.Linear(dim, n_heads*head_dim, bias=False)
+    this.kProj = QuantizedLinear.load(weights, `${prefix}.k_proj`, config);
+    this.vProj = QuantizedLinear.load(weights, `${prefix}.v_proj`, config);
+    this.oProj = QuantizedLinear.load(weights, `${prefix}.o_proj`, config);
+    this.qNorm = new RMSNorm(weights.tensor(`${prefix}.q_norm.weight`), t.rmsNormEps); // nn.RMSNorm(head_dim)
+    this.kNorm = new RMSNorm(weights.tensor(`${prefix}.k_norm.weight`), t.rmsNormEps);
+  }
+
+  forward(x: MlxArray, mask: Mask, cache: Cache): MlxArray {
+    const [B, L] = x.shape as [number, number, number];                // B, L, D = x.shape
+    let q = this.qProj.forward(x);                                     // queries = self.q_proj(x)
+    let k = this.kProj.forward(x);                                     // keys = self.k_proj(x)
+    let v = this.vProj.forward(x);                                     // values = self.v_proj(x)
+    // queries = self.q_norm(queries.reshape(B,L,n_heads,-1)).transpose(0,2,1,3)
+    q = disposing(q, ops.reshape(q, [B, L, this.nHeads, this.headDim]));
+    k = disposing(k, ops.reshape(k, [B, L, this.nKvHeads, this.headDim]));
+    v = disposing(v, ops.reshape(v, [B, L, this.nKvHeads, this.headDim]));
+    q = disposing(q, this.qNorm.forward(q));                           // self.q_norm(...)
+    k = disposing(k, this.kNorm.forward(k));                           // self.k_norm(...)
+    q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));              //   .transpose(0,2,1,3)
+    k = disposing(k, ops.transposeAxes(k, [0, 2, 1, 3]));
+    v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));              // values.reshape(...).transpose(0,2,1,3)
+    q = disposing(q, ops.rope(q, this.headDim, this.ropeBase, cache.offset, null)); // self.rope(queries, offset=cache.offset)
+    k = disposing(k, ops.rope(k, this.headDim, this.ropeBase, cache.offset, null)); // self.rope(keys, offset=cache.offset)
+    const [keys, values] = cache.updateAndFetch(k, v);                 // cache.update_and_fetch(keys, values)
+    k.dispose(); v.dispose();
+    let out = ops.sdpa(q, keys, values, this.scale, mask.mode, mask.arr); // scaled_dot_product_attention(...)
+    q.dispose(); keys.dispose(); values.dispose();
+    out = disposing(out, ops.transposeAxes(out, [0, 2, 1, 3]));        // output.transpose(0,2,1,3)
+    out = disposing(out, ops.reshape(out, [B, L, -1]));                //   .reshape(B, L, -1)
+    const y = this.oProj.forward(out); out.dispose();                 // self.o_proj(output)
+    return y;
+  }
+}
+
+// ── qwen3_moe.py MLP (dense-layer FFN) ───────────────────────────────────────
+// def __call__(self, x): return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+class MLP {
+  readonly gateProj: QuantizedLinear;
+  readonly upProj: QuantizedLinear;
+  readonly downProj: QuantizedLinear;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    this.gateProj = QuantizedLinear.load(weights, `${prefix}.gate_proj`, config); // nn.Linear(dim, hidden_dim, bias=False)
+    this.upProj = QuantizedLinear.load(weights, `${prefix}.up_proj`, config);
+    this.downProj = QuantizedLinear.load(weights, `${prefix}.down_proj`, config);
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const g = this.gateProj.forward(x);                                // self.gate_proj(x)
+    const u = this.upProj.forward(x);                                  // self.up_proj(x)
+    const s = compiledSwiglu(g, u);                                    // swiglu(gate, x)  (mx.compile)
+    g.dispose(); u.dispose();
+    const y = this.downProj.forward(s); s.dispose();                  // self.down_proj(...)
+    return y;
+  }
+}
+
+// ── switch_layers.py SwitchGLU (swiglu activation, quantized experts) ─────────
+// The three per-expert projections are the stacked switch_mlp.{gate,up,down}_proj
+// tensors; the activation is the SAME @mx.compile swiglu as activations.py.
+class SwitchGLU {
+  readonly gateProj: QuantizedSwitchLinear;
+  readonly upProj: QuantizedSwitchLinear;
+  readonly downProj: QuantizedSwitchLinear;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    this.gateProj = QuantizedSwitchLinear.load(weights, `${prefix}.gate_proj`, config);
+    this.upProj = QuantizedSwitchLinear.load(weights, `${prefix}.up_proj`, config);
+    this.downProj = QuantizedSwitchLinear.load(weights, `${prefix}.down_proj`, config);
+  }
+
+  /** x [B, L, H], indices [B, L, k] → [B, L, k, H]. */
+  forward(x: MlxArray, indices: MlxArray): MlxArray {
+    // x = mx.expand_dims(x, (-2, -3))  → [B, L, 1, 1, H]
+    let h = ops.expandDims(x, -2);
+    h = disposing(h, ops.expandDims(h, -3));
+
+    // do_sort = indices.size >= 64
+    const doSort = indices.size >= 64;
+    let idx = indices;
+    let invOrder: MlxArray | null = null;
+    let order: MlxArray | null = null;
+    if (doSort) {
+      // _gather_sort(x, indices): flatten indices, argsort → order, argsort(order) →
+      // inv_order, gather rows by order // M (M = k = last dim of indices).
+      const M = indices.shape[indices.ndim - 1]!;                      // *_, M = indices.shape
+      const idxFlat = ops.reshape(indices, [indices.size]);           // indices.flatten()
+      order = ops.argsortAxis(idxFlat, 0);                            // order = mx.argsort(indices)
+      invOrder = ops.argsortAxis(order, 0);                          // inv_order = mx.argsort(order)
+      const mScalar = ops.scalarLike(M, order);
+      const rowIdx = ops.floorDivide(order, mScalar);                // order // M
+      mScalar.dispose();
+      const [, , , , H] = h.shape as number[];
+      const flat = ops.reshape(h, [-1, 1, H!]);                      // x.flatten(0, -3)
+      h.dispose();
+      h = ops.takeAxis(flat, rowIdx, 0);                             // x.flatten(0,-3)[order // M]
+      flat.dispose(); rowIdx.dispose();
+      idx = ops.takeAxis(idxFlat, order, 0);                         // indices[order]
+      idxFlat.dispose();
+    }
+    // (self.training path — mx.stop_gradient(idx) — is inference-irrelevant here)
+    const xUp = this.upProj.forward(h, idx, doSort);                  // x_up = self.up_proj(x, idx, sorted_indices=do_sort)
+    const xGate = this.gateProj.forward(h, idx, doSort);             // x_gate = self.gate_proj(x, idx, sorted_indices=do_sort)
+    h.dispose();
+    // self.activation(x_up, x_gate) == SwiGLU()(x_up, x_gate) == swiglu(x_gate, x_up)
+    // (SwiGLU.__call__(self, x, gate) returns swiglu(gate, x); called (x_up, x_gate)).
+    const mid = compiledSwiglu(xGate, xUp);                           // swiglu(gate=x_gate, x=x_up)  (mx.compile)
+    xGate.dispose(); xUp.dispose();
+    let y = this.downProj.forward(mid, idx, doSort);                 // x = self.down_proj(activation, idx, sorted_indices=do_sort)
+    mid.dispose();
+    if (idx !== indices) idx.dispose();
+
+    if (doSort) {
+      // _scatter_unsort(x, inv_order, indices.shape): x[inv_order], unflatten to
+      // indices.shape (+ the trailing [1, H] the squeeze(-2) below removes).
+      y = disposing(y, ops.takeAxis(y, invOrder!, 0));               // x = x[inv_order]
+      invOrder!.dispose(); order!.dispose();
+      const [B, L, k] = indices.shape as number[];
+      const H = y.shape[y.ndim - 1]!;
+      y = disposing(y, ops.reshape(y, [B!, L!, k!, 1, H]));          // mx.unflatten(x, 0, shape)
+    }
+    // return x.squeeze(-2)
+    const shape = y.shape;
+    shape.splice(shape.length - 2, 1);
+    return disposing(y, ops.reshape(y, shape));
+  }
+}
+
+// ── qwen3_moe.py Qwen3MoeSparseMoeBlock ──────────────────────────────────────
+class Qwen3MoeSparseMoeBlock {
+  readonly gate: QuantizedLinear;
+  readonly switchMlp: SwitchGLU;
+  readonly numExperts: number;
+  readonly topK: number;
+  readonly normTopkProb: boolean;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    const t = config.text;
+    this.numExperts = t.numExperts;                                   // self.num_experts
+    this.topK = t.topKExperts;                                        // self.top_k
+    this.normTopkProb = t.normTopkProb;                              // self.norm_topk_prob
+    this.gate = QuantizedLinear.load(weights, `${prefix}.gate`, config); // nn.Linear(dim, num_experts, bias=False)
+    this.switchMlp = new SwitchGLU(weights, config, `${prefix}.switch_mlp`); // SwitchGLU(dim, moe_inter, num_experts)
+  }
+
+  forward(x: MlxArray): MlxArray {
+    let gates = this.gate.forward(x);                                 // gates = self.gate(x)
+    gates = disposing(gates, ops.softmaxAxis(gates, -1, true));      // gates = mx.softmax(gates, axis=-1, precise=True)
+
+    const k = this.topK;
+    const [B, L] = gates.shape as [number, number, number];
+    // inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
+    const part = ops.argpartitionAxis(gates, this.numExperts - k, -1); // argpartition(..., kth=-k)
+    let inds: MlxArray;
+    if (isCompiledTrace()) {
+      const start = ops.fromInt32([this.numExperts - k], [1]);
+      inds = ops.sliceDynamic(part, start, [2], [B, L, k]);
+      start.dispose();
+    } else {
+      inds = part.slice([0, 0, this.numExperts - k], [B, L, this.numExperts]); // [..., -k:]
+    }
+    part.dispose();
+
+    let scores = ops.takeAlongAxis(gates, inds, -1);                 // scores = mx.take_along_axis(gates, inds, axis=-1)
+    gates.dispose();
+    if (this.normTopkProb) {
+      const denom = ops.sumAxis(scores, -1, true);                  // mx.sum(scores, axis=-1, keepdims=True)
+      scores = disposing(scores, ops.div(scores, denom));           // scores /= ...
+      denom.dispose();
+    }
+
+    const y = this.switchMlp.forward(x, inds);                       // y = self.switch_mlp(x, inds)
+    inds.dispose();
+    // y = (y * scores[..., None]).sum(axis=-2)
+    const sExp = ops.expandDims(scores, -1);                        // scores[..., None]
+    scores.dispose();
+    const wy = ops.mul(y, sExp);                                    // y * scores[..., None]
+    y.dispose(); sExp.dispose();
+    return disposing(wy, ops.sumAxis(wy, -2, false));               // .sum(axis=-2)
+  }
+}
+
+// ── qwen3_moe.py Qwen3MoeDecoderLayer ────────────────────────────────────────
+class Qwen3MoeDecoderLayer {
+  readonly selfAttn: Attention;
+  readonly mlp: Qwen3MoeSparseMoeBlock | MLP;
+  readonly inputLayernorm: RMSNorm;
+  readonly postAttentionLayernorm: RMSNorm;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string, layerIdx: number) {
+    const t = config.text;
+    this.selfAttn = new Attention(weights, config, `${prefix}.self_attn`);
+    this.inputLayernorm = new RMSNorm(weights.tensor(`${prefix}.input_layernorm.weight`), t.rmsNormEps);
+    this.postAttentionLayernorm = new RMSNorm(weights.tensor(`${prefix}.post_attention_layernorm.weight`), t.rmsNormEps);
+    // if (layer_idx not in mlp_only_layers) and (num_experts > 0 and
+    //     (layer_idx + 1) % decoder_sparse_step == 0): Qwen3MoeSparseMoeBlock else MLP
+    const isMoe =
+      !t.mlpOnlyLayers.includes(layerIdx) &&
+      t.numExperts > 0 &&
+      (layerIdx + 1) % t.decoderSparseStep === 0;
+    this.mlp = isMoe
+      ? new Qwen3MoeSparseMoeBlock(weights, config, `${prefix}.mlp`)
+      : new MLP(weights, config, `${prefix}.mlp`);
+  }
+
+  forward(x: MlxArray, mask: Mask, cache: Cache): MlxArray {
+    const xn = this.inputLayernorm.forward(x);                       // self.input_layernorm(x)
+    const r1 = this.selfAttn.forward(xn, mask, cache);              // r = self.self_attn(..., mask, cache)
+    xn.dispose();
+    const h = ops.add(x, r1); r1.dispose();                        // h = x + r
+    const hn = this.postAttentionLayernorm.forward(h);             // self.post_attention_layernorm(h)
+    const r2 = this.mlp.forward(hn); hn.dispose();                 // r = self.mlp(...)
+    const out = ops.add(h, r2); h.dispose(); r2.dispose();        // out = h + r
+    return out;
+  }
+}
+
+// ── qwen3_moe.py Qwen3MoeModel + Model ───────────────────────────────────────
+export class Qwen3MoeModel {
+  readonly config: ModelConfig;
+  readonly weightsBytes: number;
+  readonly prefixBase = "model";
+  readonly loraState = new LoraState();
+  readonly embed: QuantizedEmbedding;
+  readonly layers: Qwen3MoeDecoderLayer[];
+  readonly finalNorm: RMSNorm;
+  /** Untied lm_head (tie_word_embeddings=false for this family). */
+  readonly lmHead: QuantizedLinear | null;
+
+  constructor(weights: Weights, config: ModelConfig) {
+    this.config = config;
+    this.weightsBytes = [...weights.shards.files.values()]
+      .reduce((a, f) => a + f.mmap.size, 0);
+    this.embed = QuantizedEmbedding.load(weights, "model.embed_tokens", config); // nn.Embedding
+    this.layers = Array.from(
+      { length: config.text.numHiddenLayers },
+      (_, i) => new Qwen3MoeDecoderLayer(weights, config, `model.layers.${i}`, i),
+    );
+    this.finalNorm = new RMSNorm(weights.tensor("model.norm.weight"), config.text.rmsNormEps); // self.norm
+    this.lmHead = config.text.tieWordEmbeddings
+      ? null
+      : QuantizedLinear.load(weights, "lm_head", config);           // self.lm_head
+  }
+
+  loraTargets(): Map<string, QuantizedLinear> {
+    const out = new Map<string, QuantizedLinear>();
+    for (let i = 0; i < this.layers.length; i++) {
+      const l = this.layers[i]!;
+      const p = `model.layers.${i}`;
+      out.set(`${p}.self_attn.q_proj`, l.selfAttn.qProj);
+      out.set(`${p}.self_attn.k_proj`, l.selfAttn.kProj);
+      out.set(`${p}.self_attn.v_proj`, l.selfAttn.vProj);
+      out.set(`${p}.self_attn.o_proj`, l.selfAttn.oProj);
+      if (l.mlp instanceof MLP) {
+        out.set(`${p}.mlp.gate_proj`, l.mlp.gateProj);
+        out.set(`${p}.mlp.up_proj`, l.mlp.upProj);
+        out.set(`${p}.mlp.down_proj`, l.mlp.downProj);
+      }
+      // MoE switch_mlp / gate are QuantizedSwitchLinear (stacked experts) — not
+      // LoRA targets under our per-linear adapter machinery.
+    }
+    return out;
+  }
+
+  makeCache(): Cache[] {
+    return this.layers.map(() => new KVCache());
+  }
+
+  forwardHidden(ids: MlxArray, cache: Cache[]): MlxArray {
+    const h = this.embed.encode(ids);                               // self.embed_tokens(inputs)
+    return this.forwardLayers(h, cache);
+  }
+
+  forwardEmbeddings(_embeds: MlxArray, _cache: Cache[], _bidir: MlxArray | null): MlxArray {
+    throw new Error("qwen3_moe input-embedding path is not supported");
+  }
+
+  protected forwardLayers(h0: MlxArray, cache: Cache[]): MlxArray {
+    const L = h0.shape[1]!;
+    const mask = cache[0]!.makeMask(L, null);                       // create_attention_mask(h, cache[0])
+    let cur = h0;
+    for (let i = 0; i < this.layers.length; i++) {                  // for layer, c in zip(self.layers, cache)
+      const next = this.layers[i]!.forward(cur, mask, cache[i]!);   // h = layer(h, mask, c)
+      cur.dispose();
+      cur = next;
+    }
+    mask.arr?.dispose();
+    return disposing(cur, this.finalNorm.forward(cur));            // return self.norm(h)
+  }
+
+  logitsFromHidden(h: MlxArray): MlxArray {
+    return this.lmHead ? this.lmHead.forward(h) : this.embed.asLinear(h); // self.lm_head(out)
+  }
+
+  forward(tokens: number[] | MlxArray, cache: Cache[]): MlxArray {
+    const ids = Array.isArray(tokens)
+      ? ops.fromInt32(tokens, [1, tokens.length])
+      : tokens;
+    const h = this.forwardHidden(ids, cache);
+    if (Array.isArray(tokens)) ids.dispose();
+    const logits = this.logitsFromHidden(h);
+    h.dispose();
+    return logits;
+  }
+
+  generate(promptTokens: number[], maxTokens: number, eosIds: number[] = []): number[] {
+    const cache = this.makeCache();
+    const out: number[] = [];
+    try {
+      let tokens = promptTokens;
+      for (let step = 0; step < maxTokens; step++) {
+        const logits = this.forward(tokens, cache);
+        const next = argmaxLastPosition(logits);
+        logits.dispose();
+        if (eosIds.includes(next)) break;
+        out.push(next);
+        tokens = [next];
+      }
+    } finally {
+      for (const c of cache) c.dispose();
+    }
+    return out;
+  }
+}

package/src/model/qwen3.ts

@@ -24,8 +24,9 @@ import {
   type Cache,
   type Mask,
 } from "./gemma4-base";
+import { compiledSwiglu } from "./minicpm5";
 
-class Qwen3Attention {
+export class Qwen3Attention {
   readonly qProj: QuantizedLinear;
   readonly kProj: QuantizedLinear;
   readonly vProj: QuantizedLinear;
@@ -92,7 +93,7 @@ class Qwen3Attention {
   }
 }
 
-class Qwen3MLP {
+export class Qwen3MLP {
   readonly gate: QuantizedLinear;
   readonly up: QuantizedLinear;
   readonly down: QuantizedLinear;
@@ -106,12 +107,12 @@ class Qwen3MLP {
   forward(x: MlxArray): MlxArray {
     const gate = this.gate.forward(x);
     const up = this.up.forward(x);
-    const sig = ops.sigmoid(gate);
-    const silu = ops.mul(gate, sig);
+    // Oracle (mlx_lm/models/qwen3.py MLP): down_proj(swiglu(gate_proj(x), up_proj(x)))
+    // where swiglu is the `@mx.compile` closure (silu(gate)*up). Compiled = the
+    // same libmlx kernel mlx-lm dispatches (L1-faithful) AND fuses the three
+    // standalone ops into one dispatch. Bit-identical to the old unfused path.
+    const hidden = compiledSwiglu(gate, up);
     gate.dispose();
-    sig.dispose();
-    const hidden = ops.mul(silu, up);
-    silu.dispose();
     up.dispose();
     const out = this.down.forward(hidden);
     hidden.dispose();
@@ -119,7 +120,7 @@ class Qwen3MLP {
   }
 }
 
-class Qwen3Layer {
+export class Qwen3Layer {
   readonly attn: Qwen3Attention;
   readonly mlp: Qwen3MLP;
   readonly inputNorm: RMSNorm;

package/src/model/qwen3_5.ts

@@ -14,6 +14,7 @@ import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
+import { CompiledFunction } from "../mlx/compile";
 import {
   argmaxLastPosition,
   disposeTriple,
@@ -32,8 +33,76 @@ import { gatedDeltaUpdate, SSMCache } from "./qwen3-delta";
 
 const PREFIX = "language_model";
 
+// ── Compiled activations ─────────────────────────────────────────────────────
+// The oracle (mlx_lm/models/activations.py + qwen3_next.py) wraps BOTH swiglu
+// activations in `@partial(mx.compile, shapeless=True)`. We match it: every
+// activation site below (the MLP swiglu, the RMSNormGated `_precise_swiglu`, and
+// the conv `nn.silu`) runs through a compiled closure unconditionally, so the
+// dispatched kernel set matches the oracle op-for-op (= mlx-lm, bit-exact). Traced
+// once (shapeless), replayed thereafter; autograd-safe (mx.compile threads VJP
+// through the traced graph). `compiledSwiglu` is exported so the parity test can
+// assert the closure exists and the MLP actually uses it.
+
+/** activations.py: `@mx.compile def swiglu(gate, x): return nn.silu(gate) * x`.
+ *  nn.silu(g) == g * sigmoid(g); mx.compile fuses sigmoid+mul+mul → one kernel. */
+let _swigluClosure: CompiledFunction | null = null;
+export function compiledSwiglu(gate: MlxArray, up: MlxArray): MlxArray {
+  if (!_swigluClosure) {
+    _swigluClosure = new CompiledFunction((inputs) => {
+      const g = inputs[0]!, u = inputs[1]!;
+      const sig = ops.sigmoid(g);
+      const silu = ops.mul(g, sig); sig.dispose();
+      const out = ops.mul(silu, u); silu.dispose();
+      return [out];
+    });
+  }
+  return _swigluClosure.apply([gate, up])[0]!;
+}
+
+/** qwen3_next.py `_precise_swiglu(h, gate, x)`:
+ *    gate = nn.silu(gate.astype(f32)); x = x.astype(f32); (gate*x).astype(h.dtype)
+ *  Used by Qwen3NextRMSNormGated. Inputs: (h=hidden for the out dtype, gate=z,
+ *  x=rms_norm(hidden)). mx.compile fuses the two casts + silu + mul + cast. */
+let _preciseSwigluClosure: CompiledFunction | null = null;
+export function compiledPreciseSwiglu(h: MlxArray, gate: MlxArray, x: MlxArray): MlxArray {
+  if (!_preciseSwigluClosure) {
+    _preciseSwigluClosure = new CompiledFunction((inputs) => {
+      const hh = inputs[0]!, g = inputs[1]!, xx = inputs[2]!;
+      const gf = g.astype(Dtype.float32);
+      const sig = ops.sigmoid(gf);
+      const silu = ops.mul(gf, sig); gf.dispose(); sig.dispose();
+      const xf = xx.astype(Dtype.float32);
+      const prod = ops.mul(silu, xf); silu.dispose(); xf.dispose();
+      const out = prod.astype(hh.dtype); prod.dispose();
+      return [out];
+    });
+  }
+  return _preciseSwigluClosure.apply([h, gate, x])[0]!;
+}
+
+/** mlx.nn.silu is `@partial(mx.compile, shapeless=True) def silu(x): x*sigmoid(x)`
+ *  — a COMPILED kernel, not a standalone sigmoid+mul. Used for the conv
+ *  activation `nn.silu(conv1d(...))`, matching mlx-lm's fused BV2ISigmoid…_V_. */
+let _siluClosure: CompiledFunction | null = null;
+export function compiledSilu(x: MlxArray): MlxArray {
+  if (!_siluClosure) {
+    _siluClosure = new CompiledFunction((inputs) => {
+      const xx = inputs[0]!;
+      const sig = ops.sigmoid(xx);
+      const out = ops.mul(xx, sig); sig.dispose(); // x * sigmoid(x)
+      return [out];
+    });
+  }
+  return _siluClosure.apply([x])[0]!;
+}
+
+// NOTE: the attention output gate is `self.o_proj(output * mx.sigmoid(gate))`
+// INLINE in mlx-lm (qwen3_next.py) — NOT @mx.compile. It is intentionally kept
+// as a standalone sigmoid + multiply in Qwen3Attention.forward so the dispatched
+// kernel set matches the reference; there is no compiled-output-gate helper.
+
 /** Gated-DeltaNet linear-attention layer (mlx-lm GatedDeltaNet). */
-class GatedDeltaNet {
+export class GatedDeltaNet {
   readonly inProjQkv: QuantizedLinear;
   readonly inProjZ: QuantizedLinear;
   readonly inProjB: QuantizedLinear;
@@ -99,7 +168,7 @@ class GatedDeltaNet {
 
     const conv = ops.conv1d(convInput, this.convWeight, 1, 0, 1, convDim);
     convInput.dispose();
-    const convOut = ops.silu(conv); // [B,S,convDim]
+    const convOut = compiledSilu(conv); // nn.silu (mx.compile), matches mlx-lm
     conv.dispose();
 
     const [qFlat, kFlat, vFlat] = ops.split(
@@ -144,22 +213,16 @@ class GatedDeltaNet {
 
   private rmsNormGated(hidden: MlxArray, gate: MlxArray): MlxArray {
     const xn = ops.rmsNorm(hidden, this.normWeight, this.eps); // bf16
-    const gf = gate.astype(Dtype.float32);
-    const sg = ops.silu(gf); // silu(z) in f32
-    gf.dispose();
-    const xf = xn.astype(Dtype.float32);
+    // Oracle: _precise_swiglu(hidden, gate, xn) — ALWAYS one @mx.compile kernel
+    // (matches mlx-lm; no unfused silu+mul+cast).
+    const res = compiledPreciseSwiglu(hidden, gate, xn);
     xn.dispose();
-    const prod = ops.mul(sg, xf);
-    sg.dispose();
-    xf.dispose();
-    const res = prod.astype(hidden.dtype);
-    prod.dispose();
     return res;
   }
 }
 
 /** Full (softmax) attention with output gate + q/k norm + partial RoPE. */
-class Qwen3Attention {
+export class Qwen3Attention {
   readonly qProj: QuantizedLinear;
   readonly kProj: QuantizedLinear;
   readonly vProj: QuantizedLinear;
@@ -212,8 +275,17 @@ class Qwen3Attention {
     v = disposing(v, ops.reshape(v, [B, L, this.nKvHeads, this.headDim]));
     v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));
 
-    q = disposing(q, ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
-    k = disposing(k, ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
+    // 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.
+    const offArr = (cache as { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
+    q = disposing(q, offArr
+      ? ops.ropeDynamic(q, this.ropeDims, this.ropeBase, offArr, null)
+      : ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
+    k = disposing(k, offArr
+      ? ops.ropeDynamic(k, this.ropeDims, this.ropeBase, offArr, null)
+      : ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
 
     let attn: MlxArray;
     if (cache instanceof QuantizedKVCache) {
@@ -235,20 +307,22 @@ class Qwen3Attention {
 
     const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
     attn.dispose();
-    let merged = ops.reshape(attnT, [B, L, -1]);
+    const merged = ops.reshape(attnT, [B, L, -1]);
     attnT.dispose();
-    // Output gate: o_proj(output * sigmoid(gate)).
-    const sg = ops.sigmoid(gate);
+    // qwen3_next.py:158  return self.o_proj(output * mx.sigmoid(gate))
+    // INLINE, not @mx.compile — copy the exact ops.
+    const sig = ops.sigmoid(gate);
     gate.dispose();
-    merged = disposing(merged, ops.mul(merged, sg));
-    sg.dispose();
-    const out = this.oProj.forward(merged);
+    const gated = ops.mul(merged, sig);
     merged.dispose();
+    sig.dispose();
+    const out = this.oProj.forward(gated);
+    gated.dispose();
     return out;
   }
 }
 
-class Qwen3MLP {
+export class Qwen3MLP {
   readonly gate: QuantizedLinear;
   readonly up: QuantizedLinear;
   readonly down: QuantizedLinear;
@@ -262,10 +336,10 @@ class Qwen3MLP {
   forward(x: MlxArray): MlxArray {
     const g = this.gate.forward(x);
     const u = this.up.forward(x);
-    const silu = ops.silu(g);
+    // 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();
-    const hidden = ops.mul(silu, u);
-    silu.dispose();
     u.dispose();
     const out = this.down.forward(hidden);
     hidden.dispose();
@@ -273,7 +347,7 @@ class Qwen3MLP {
   }
 }
 
-class Qwen3Layer {
+export class Qwen3Layer {
   readonly isLinear: boolean;
   readonly linearAttn: GatedDeltaNet | null = null;
   readonly selfAttn: Qwen3Attention | null = null;

package/src/model/support.ts

@@ -30,6 +30,15 @@ export function isQwen3Config(config: ModelConfig): boolean {
   return config.modelType === "qwen3";
 }
 
+/** Qwen3-MoE (model_type `qwen3_moe`, Qwen3MoeForCausalLM) — plain-qwen3
+ *  attention (per-head q/k norm, GQA, RoPE) with a sparse MoE FFN
+ *  (gate Linear -> precise softmax -> top-k -> renormalize -> SwitchGLU experts
+ *  with compiled swiglu). Distinct from the qwen3 embedding backbone and the
+ *  qwen3_5 gated-DeltaNet hybrid. */
+export function isQwen3MoeConfig(config: ModelConfig): boolean {
+  return config.modelType === "qwen3_moe";
+}
+
 /** Speculative-decoding drafters (e.g. `gemma4_assistant`) are companion
  *  artifacts to a target model — Q-only, centroid-head, no standalone LM
  *  head. They are never servable/selectable on their own (the spec path
@@ -57,6 +66,7 @@ export function supportTier(modelType: string, repoId = ""): "targeted" | "gener
   // qwen3_5 / qwen3_5_text (dense hybrid). The MoE variant is a separate type.
   if (modelType === "qwen3_5" || modelType === "qwen3_5_text") return "targeted";
   if (modelType === "qwen3") return "targeted"; // plain Qwen3 (embedding backbone)
+  if (modelType === "qwen3_moe") return "targeted"; // Qwen3-MoE (sparse experts)
   if (modelType === "llama" && repoId.toLowerCase().includes("minicpm5-1b-optiq-4bit"))
     return "targeted";
   // Tier-0 generic fallback: the universal-dense descriptor table.
@@ -72,7 +82,8 @@ export function isSupportedModelConfig(config: ModelConfig): boolean {
   if (isDrafterModelType(config.modelType)) return false;
   if (
     config.modelType.startsWith("gemma4") || isMiniCPM5Config(config) ||
-    isQwen35Config(config) || isQwen3Config(config) || isDiffusionGemmaConfig(config)
+    isQwen35Config(config) || isQwen3Config(config) || isQwen3MoeConfig(config) ||
+    isDiffusionGemmaConfig(config)
   ) return true;
   // Generic tier: an arch descriptor exists for this model_type (and its
   // config parses — a malformed config is unsupported, not a crash).