@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/models/qwen/qwen3-moe.ts

@@ -0,0 +1,355 @@
+import type { MlxCacheLayer, MlxMixedRange, MlxSequenceRange } from "../../contracts/mlx/graph";
+// Concrete model graph for Qwen3-MoE (model_type `qwen3_moe`,
+// Qwen3MoeForCausalLM). A STRAIGHT op-for-op transcription of the oracle
+// (`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 "../../artifacts/config";
+import type { Weights } from "../../artifacts/weights";
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { disposing } from "../../layers/helpers";
+import { compiledSwiglu } from "../../layers/swiglu";
+import { LoraState } from "../../layers/lora";
+import type { GraphDeclaration } from "../../contracts/portable/graph";
+import type { MlxRowTable, MlxCapturingForward, MlxDeclaredGraph, MlxHeadForward, MlxPhaseForward, MlxPromptRange, MlxTokenPhases, MlxVerifyForward } from "../../contracts/mlx/graph";
+import { declareGraph } from "../capabilities";
+import { DEFAULT_PHASE_ROWS, rowTableFrom, rowTable, type ReadPhase   } from "../graph";
+import { QuantizedEmbedding } from "../../layers/quantized-embedding";
+import { QuantizedLinear, QuantizedLinearPlain } from "../../layers/quantized-linear";
+import { QuantizedSwitchLinear } from "../../layers/quantized-switch-linear";
+import { RMSNorm } from "../../layers/normalization";
+import { type Cache } from "../../contracts/mlx/cache";
+import { Qwen3Attention } from "./qwen3";
+
+// ── qwen3_moe.py Attention (== plain qwen3, rope_theta default 1e7) ───────────
+const MOE_ROPE_THETA = 10000000;                                        // args.rope_theta
+
+// ── 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 kind = "dense";
+  readonly gateProj: QuantizedLinear;
+  readonly upProj: QuantizedLinear;
+  readonly downProj: QuantizedLinear;
+
+  readonly #gate: QuantizedLinearPlain;
+  readonly #up: QuantizedLinearPlain;
+  readonly #down: QuantizedLinearPlain;
+
+  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);
+    this.#gate = new QuantizedLinearPlain(this.gateProj);
+    this.#up = new QuantizedLinearPlain(this.upProj);
+    this.#down = new QuantizedLinearPlain(this.downProj);
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const g = this.#gate.forward(x);                                // self.gate_proj(x)
+    const u = this.#up.forward(x);                                  // self.up_proj(x)
+    const s = compiledSwiglu(g, u);                                    // swiglu(gate, x)  (mx.compile)
+    g.dispose(); u.dispose();
+    const y = this.#down.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.
+export class Qwen3SwitchWeights {
+  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);
+  }
+}
+
+interface ExpertForward { forward(x: MlxArray, indices: MlxArray): MlxArray }
+
+/** Gather projections over the original token order. */
+export class Qwen3UnsortedExperts implements ExpertForward {
+  constructor(readonly weights: Qwen3SwitchWeights) {}
+  forward(x: MlxArray, indices: MlxArray): MlxArray {
+    using rows = ops.expandDims(x, -2);
+    using expanded = ops.expandDims(rows, -3);
+    using up = this.weights.upProj.forward(expanded, indices, false);
+    using gate = this.weights.gateProj.forward(expanded, indices, false);
+    using hidden = compiledSwiglu(gate, up);
+    const output = this.weights.downProj.forward(hidden, indices, false);
+    const shape = output.shape;
+    shape.splice(shape.length - 2, 1);
+    return disposing(output, ops.reshape(output, shape));
+  }
+}
+
+/** Sorted gather followed by the reference inverse permutation. */
+export class Qwen3SortedExperts implements ExpertForward {
+  constructor(readonly weights: Qwen3SwitchWeights) {}
+  forward(x: MlxArray, indices: MlxArray): MlxArray {
+    const topK = indices.shape[indices.ndim - 1]!;
+    using rows = ops.expandDims(x, -2);
+    using expanded = ops.expandDims(rows, -3);
+    using flatIndices = ops.reshape(indices, [indices.size]);
+    using order = ops.argsortAxis(flatIndices, 0);
+    using inverse = ops.argsortAxis(order, 0);
+    using count = ops.scalarLike(topK, order);
+    using rowIndices = ops.floorDivide(order, count);
+    const width = expanded.shape[4]!;
+    using flat = ops.reshape(expanded, [-1, 1, width]);
+    using sortedRows = ops.takeAxis(flat, rowIndices, 0);
+    using sortedIndices = ops.takeAxis(flatIndices, order, 0);
+    using up = this.weights.upProj.forward(sortedRows, sortedIndices, true);
+    using gate = this.weights.gateProj.forward(sortedRows, sortedIndices, true);
+    using hidden = compiledSwiglu(gate, up);
+    using sortedOutput = this.weights.downProj.forward(hidden, sortedIndices, true);
+    using output = ops.takeAxis(sortedOutput, inverse, 0);
+    const [batch, length, selected] = indices.shape;
+    const outputWidth = output.shape[output.ndim - 1]!;
+    using unflattened = ops.reshape(output, [batch!, length!, selected!, 1, outputWidth]);
+    const shape = unflattened.shape;
+    shape.splice(shape.length - 2, 1);
+    return ops.reshape(unflattened, shape);
+  }
+}
+
+// ── qwen3_moe.py Qwen3MoeSparseMoeBlock ──────────────────────────────────────
+class Qwen3MoeSparseMoeBlock {
+  readonly kind = "moe";
+  readonly gate: QuantizedLinear;
+  readonly switchMlp: Qwen3SwitchWeights;
+  readonly numExperts: number;
+  readonly topK: number;
+  readonly normTopkProb: boolean;
+  readonly #gate: QuantizedLinearPlain;
+  readonly #normalize: (scores: MlxArray) => MlxArray;
+
+  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 Qwen3SwitchWeights(weights, config, `${prefix}.switch_mlp`); // SwitchGLU(dim, moe_inter, num_experts)
+    this.#gate = new QuantizedLinearPlain(this.gate);
+    this.#normalize = this.normTopkProb ? scores => {
+      using denominator = ops.sumAxis(scores, -1, true);
+      return disposing(scores, ops.div(scores, denominator));
+    } : scores => scores;
+  }
+
+  forward(x: MlxArray, experts: ExpertForward): 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)
+    const inds = part.slice([0, 0, this.numExperts - k], [B, L, this.numExperts]);
+    part.dispose();
+
+    let scores = ops.takeAlongAxis(gates, inds, -1);                 // scores = mx.take_along_axis(gates, inds, axis=-1)
+    gates.dispose();
+    scores = this.#normalize(scores);
+
+    const y = experts.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: Qwen3Attention;
+  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 Qwen3Attention(weights, config, `${prefix}.self_attn`, MOE_ROPE_THETA);
+    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`);
+  }
+}
+
+class FixedMoeBlock {
+  constructor(readonly weights: Qwen3MoeSparseMoeBlock, readonly experts: ExpertForward) {}
+  forward(x: MlxArray): MlxArray { return this.weights.forward(x, this.experts); }
+}
+class Qwen3MoeLayerView {
+  readonly #mlp: MLP | FixedMoeBlock;
+  constructor(readonly base: Qwen3MoeDecoderLayer, Experts: new (weights: Qwen3SwitchWeights) => ExpertForward) {
+    this.#mlp = base.mlp.kind === "dense" ? base.mlp : new FixedMoeBlock(base.mlp, new Experts(base.mlp.switchMlp));
+  }
+  forward(x: MlxArray, cache: Cache, phase: ReadPhase): MlxArray {
+    const xn = this.base.inputLayernorm.forward(x);                       // self.input_layernorm(x)
+    const r1 = this.base.selfAttn.forward(xn, cache, phase);             // r = self.self_attn(..., mask, cache)
+    xn.dispose();
+    const h = ops.add(x, r1); r1.dispose();                        // h = x + r
+    const hn = this.base.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 implements MlxDeclaredGraph, MlxTokenPhases {
+  readonly config: ModelConfig;
+  readonly weightsBytes: number;
+  readonly prefixBase = "model";
+  readonly loraState = new LoraState();
+  /** An ordinary autoregressive graph: every method over every cache lego. */
+  get graphCapabilities(): GraphDeclaration { return declareGraph(); }
+  readonly embed: QuantizedEmbedding;
+  readonly layers: Qwen3MoeDecoderLayer[];
+  readonly finalNorm: RMSNorm;
+  /** Untied lm_head (tie_word_embeddings=false for this family). */
+  readonly lmHead: QuantizedLinear | null;
+  readonly #head: (hidden: MlxArray) => MlxArray;
+
+  constructor(weights: Weights, config: ModelConfig, maxPhaseRows = DEFAULT_PHASE_ROWS) {
+    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
+
+    const head = this.lmHead && new QuantizedLinearPlain(this.lmHead);
+    this.#head = head ? hidden => head.forward(hidden) : hidden => this.embed.asLinear(hidden);
+    const unsorted = this.layers.map(layer => new Qwen3MoeLayerView(layer, Qwen3UnsortedExperts));
+    const sorted = this.layers.map(layer => new Qwen3MoeLayerView(layer, Qwen3SortedExperts));
+    const sortFrom = Math.ceil(64 / config.text.topKExperts);
+    const decodeUnsorted = (ids: MlxArray, cache: Cache[]) => this.forwardLayers(this.embed.encode(ids), cache, "decode", unsorted);
+    const decodeSorted = (ids: MlxArray, cache: Cache[]) => this.forwardLayers(this.embed.encode(ids), cache, "decode", sorted);
+    const windowUnsorted = (ids: MlxArray, cache: Cache[]) => this.forwardLayers(this.embed.encode(ids), cache, "window", unsorted);
+    const windowSorted = (ids: MlxArray, cache: Cache[]) => this.forwardLayers(this.embed.encode(ids), cache, "window", sorted);
+    this.prompt = [{ from: 1, to: sortFrom - 1, operation: windowUnsorted }, { from: sortFrom, to: Infinity, operation: windowSorted }]
+      .filter(range => range.from <= range.to);
+    this.promptChunk = this.prompt;
+    this.multiSequencePrompt = this.prompt;
+    this.multiSequencePromptChunk = this.multiSequencePrompt;
+    this.decode = rowTableFrom(rows => rows < sortFrom ? decodeUnsorted : decodeSorted, maxPhaseRows);
+    const windows = rowTableFrom(rows => rows < sortFrom ? windowUnsorted : windowSorted, maxPhaseRows);
+    this.verify = windows;
+    this.logits = rowTable(hidden => this.#head(hidden), maxPhaseRows);
+    const sequence: MlxSequenceRange[] = [
+      { from: 1, to: sortFrom - 1, operation: (ids, cache) => windowUnsorted(ids, cache) },
+      { from: sortFrom, to: Infinity, operation: (ids, cache) => windowSorted(ids, cache) },
+    ];
+    this.sequence = sequence.filter(range => range.from <= range.to);
+    this.appendKnown = this.sequence;
+
+    this.cacheLayers = this.layers.map(() => ({kind:"full" as const,turboQuant:"dense" as const}));
+  }
+
+  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.kind === "dense") {
+        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);
+      }
+      else out.set(`${p}.mlp.gate`, l.mlp.gate);
+      // Stacked expert weights are separate from the two-dimensional router target.
+    }
+    return out;
+  }
+
+  /** Builds no constants and opens no weights: nothing to release. */
+  dispose(): void {}
+
+  /** One cache per layer, sharing one mask memo (create_attention_mask's one
+   * mask per forward, built by the caches). Their TurboQuant read is dense
+   * (`denseTurboQuantKernels`): the read decodes keys and values, then runs the
+   * fused SDPA. */
+  readonly cacheLayers: readonly MlxCacheLayer[];
+  readonly sequence: readonly MlxSequenceRange[];
+
+
+  /** One whole chunk of a prompt; rows are sequences times the chunk size. The GPU is this graph's one unit. */
+  readonly prompt: readonly MlxPromptRange[];
+  readonly promptChunk: typeof this.prompt;
+  readonly multiSequencePromptChunk: typeof this.multiSequencePrompt;
+  readonly multiSequencePrompt: readonly MlxPromptRange[];
+
+  /** A prompt's positions that are not a whole chunk: the window read at every width. */
+
+  /** Known tokens: this graph declares no committed spans, so they read as a window. */
+  readonly appendKnown: import("../../contracts/mlx/graph").MlxKnownOperations;
+
+  /** The one output head at every width. */
+  readonly logits: MlxRowTable<MlxHeadForward>;
+
+  /** One new token per sequence: `ids` [sequences, 1]. */
+  readonly decode: MlxRowTable<MlxCapturingForward>;
+
+  /** A speculative window of depth + 1 positions per sequence. This graph
+   * declares no hidden-layer taps, so `capture` is always null. */
+  readonly verify: MlxRowTable<MlxCapturingForward>;
+
+  /** The decoder over `h0` (consumed), every block reading its cache for `phase`. */
+  protected forwardLayers(h0: MlxArray, cache: Cache[], phase: ReadPhase, layers: readonly Qwen3MoeLayerView[]): MlxArray {
+    let cur = h0;
+    for (let i = 0; i < layers.length; i++) {                  // for layer, c in zip(self.layers, cache)
+      const next = layers[i]!.forward(cur, cache[i]!, phase);  // h = layer(h, mask, c)
+      cur.dispose();
+      cur = next;
+    }
+    return disposing(cur, this.finalNorm.forward(cur));            // return self.norm(h)
+  }
+
+  logitsFromHidden(h: MlxArray): MlxArray {
+    return this.#head(h); // self.lm_head(out)
+  }
+}

package/src/models/qwen/qwen3.ts

@@ -0,0 +1,323 @@
+import type { MlxCacheLayer, MlxMixedRange, MlxSequenceRange } from "../../contracts/mlx/graph";
+// Concrete model graph for plain Qwen3 (Qwen3ForCausalLM).
+// Port target: mlx_lm.models.qwen3. A standard dense decoder — GQA softmax
+// attention with per-head q/k RMSNorm (applied over head_dim BEFORE RoPE),
+// full-head RoPE (theta 1e6), SwiGLU MLP, RMSNorm, tied embeddings.
+//
+// Primary use here is as a TEXT-EMBEDDING backbone (mlx-community Qwen3-Embedding-*):
+// `embedPooled` returns the last-token hidden, L2-normalized — the pooling
+// Qwen3-Embedding was trained with (an <|endoftext|> token is appended by the
+// caller; its hidden is the sentence vector). The full causal-LM surface
+// (forward/generate via the tied head) is wired too so it conforms to RuntimeModel.
+//
+// Each attention block reads its cache through the cache contract for the
+// phase its caller names (`ReadPhase`); the cache was composed with its read.
+
+import type { ModelConfig } from "../../artifacts/config";
+import type { Weights } from "../../artifacts/weights";
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { disposing } from "../../layers/helpers";
+import { composedFullAttentionCache } from "../../state/bf16-first-kv";
+import { AttentionMasks } from "../../state/attention-read";
+import { denseTurboQuantKernels } from "../../state/turboquant-attention";
+import { LoraState } from "../../layers/lora";
+import type { GraphDeclaration } from "../../contracts/portable/graph";
+import type { MlxRowTable, MlxCapturingForward, MlxDeclaredGraph, MlxHeadForward, MlxPhaseForward, MlxPromptRange, MlxTokenPhases, MlxVerifyForward } from "../../contracts/mlx/graph";
+import { declareGraph } from "../capabilities";
+import { DEFAULT_PHASE_ROWS,  rowTable, type ReadPhase   } from "../graph";
+import { QuantizedEmbedding } from "../../layers/quantized-embedding";
+import { QuantizedLinear, QuantizedLinearPlain } from "../../layers/quantized-linear";
+import { RMSNorm } from "../../layers/normalization";
+import { type Cache } from "../../contracts/mlx/cache";
+import { compiledSwiglu } from "../../layers/swiglu";
+import { cacheRead, type AttentionReadStrategy } from "./blocks";
+
+/** Qwen3 attention (per-head q/k RMSNorm over head_dim before full-head RoPE,
+ *  GQA). qwen3 and qwen3_moe share it; they differ only in the `rope_theta`
+ *  default used when the config omits one. */
+export class Qwen3Attention {
+  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;
+
+  readonly #q: QuantizedLinearPlain;
+  readonly #k: QuantizedLinearPlain;
+  readonly #v: QuantizedLinearPlain;
+  readonly #out: QuantizedLinearPlain;
+
+  /** `reads` is how the block reads its cache: the cache's own read, fixed when
+   * the block is built. */
+  constructor(weights: Weights, config: ModelConfig, prefix: string, defaultRopeBase = 1000000,
+    readonly reads: AttentionReadStrategy = cacheRead) {
+    const t = config.text;
+    this.nHeads = t.numAttentionHeads;
+    this.nKvHeads = t.numKeyValueHeads;
+    this.headDim = t.headDim;
+    this.scale = Math.pow(this.headDim, -0.5);
+    this.ropeBase = t.ropeParameters.full_attention?.ropeTheta ?? defaultRopeBase;
+    this.qProj = QuantizedLinear.load(weights, `${prefix}.q_proj`, config);
+    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);
+    this.kNorm = new RMSNorm(weights.tensor(`${prefix}.k_norm.weight`), t.rmsNormEps);
+    this.#q = new QuantizedLinearPlain(this.qProj);
+    this.#k = new QuantizedLinearPlain(this.kProj);
+    this.#v = new QuantizedLinearPlain(this.vProj);
+    this.#out = new QuantizedLinearPlain(this.oProj);
+  }
+
+  /** `x` [B, L, hidden] (borrowed) to an owned [B, L, hidden], appending to
+   *  `cache` and attending through the read `phase` names. */
+  forward(x: MlxArray, cache: Cache, phase: ReadPhase): MlxArray {
+    const [B, L] = x.shape as [number, number, number];
+    let q = this.#q.forward(x);
+    let k = this.#k.forward(x);
+    let v = this.#v.forward(x);
+
+    // Reshape to heads, then q/k RMSNorm over head_dim BEFORE transpose+RoPE
+    // (mlx-lm qwen3 reference order).
+    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));
+    k = disposing(k, this.kNorm.forward(k));
+
+    q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
+    k = disposing(k, ops.transposeAxes(k, [0, 2, 1, 3]));
+    v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));
+
+    q = disposing(q, ops.rope(q, this.headDim, this.ropeBase, cache.offset, null));
+    k = disposing(k, ops.rope(k, this.headDim, this.ropeBase, cache.offset, null));
+
+    // The cache attends with the kernel it was composed with: mlx-lm's
+    // scaled_dot_product_attention reads quantized caches through quantized SDPA.
+    let attn: MlxArray;
+    try {
+      const read = this.reads.read(cache, phase, k, v);
+      try { attn = read.attend(q, this.scale); } finally { read.dispose(); }
+    } finally { k.dispose(); v.dispose(); q.dispose(); }
+
+    const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+    attn.dispose();
+    const merged = ops.reshape(attnT, [B, L, -1]);
+    attnT.dispose();
+    const out = this.#out.forward(merged);
+    merged.dispose();
+    return out;
+  }
+}
+
+export class Qwen3MLP {
+  readonly gate: QuantizedLinear;
+  readonly up: QuantizedLinear;
+  readonly down: QuantizedLinear;
+
+  readonly #gate: QuantizedLinearPlain;
+  readonly #up: QuantizedLinearPlain;
+  readonly #down: QuantizedLinearPlain;
+
+  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 = new QuantizedLinearPlain(this.gate);
+    this.#up = new QuantizedLinearPlain(this.up);
+    this.#down = new QuantizedLinearPlain(this.down);
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const gate = this.#gate.forward(x);
+    const up = this.#up.forward(x);
+    // 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();
+    up.dispose();
+    const out = this.#down.forward(hidden);
+    hidden.dispose();
+    return out;
+  }
+}
+
+export class Qwen3Layer {
+  readonly attn: Qwen3Attention;
+  readonly mlp: Qwen3MLP;
+  readonly inputNorm: RMSNorm;
+  readonly postAttnNorm: RMSNorm;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    this.attn = new Qwen3Attention(weights, config, `${prefix}.self_attn`);
+    this.mlp = new Qwen3MLP(weights, config, `${prefix}.mlp`);
+    this.inputNorm = new RMSNorm(weights.tensor(`${prefix}.input_layernorm.weight`), config.text.rmsNormEps);
+    this.postAttnNorm = new RMSNorm(weights.tensor(`${prefix}.post_attention_layernorm.weight`), config.text.rmsNormEps);
+  }
+
+  forward(x: MlxArray, cache: Cache, phase: ReadPhase): MlxArray {
+    const xn = this.inputNorm.forward(x);
+    const a = this.attn.forward(xn, cache, phase);
+    xn.dispose();
+    const h = ops.add(x, a);
+    a.dispose();
+    const hn = this.postAttnNorm.forward(h);
+    const m = this.mlp.forward(hn);
+    hn.dispose();
+    const out = ops.add(h, m);
+    h.dispose();
+    m.dispose();
+    return out;
+  }
+}
+
+export class Qwen3Model implements MlxDeclaredGraph, MlxTokenPhases {
+  readonly config: ModelConfig;
+  readonly weightsBytes: number;
+  readonly prefixBase = "model";
+  readonly loraState = new LoraState();
+  /** An ordinary autoregressive graph that also exposes the pooled embedding path. */
+  get graphCapabilities(): GraphDeclaration { return declareGraph(); }
+  readonly embed: QuantizedEmbedding;
+  readonly layers: Qwen3Layer[];
+  readonly finalNorm: RMSNorm;
+  /** Separate lm_head when NOT tied; otherwise the tied embedding's asLinear. */
+  readonly lmHead: QuantizedLinear | null;
+  readonly #head: (hidden: MlxArray) => MlxArray;
+
+  constructor(weights: Weights, config: ModelConfig, maxPhaseRows = DEFAULT_PHASE_ROWS) {
+    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);
+    this.layers = Array.from(
+      { length: config.text.numHiddenLayers },
+      (_, i) => new Qwen3Layer(weights, config, `model.layers.${i}`),
+    );
+    this.finalNorm = new RMSNorm(weights.tensor("model.norm.weight"), config.text.rmsNormEps);
+    this.lmHead = config.text.tieWordEmbeddings
+      ? null
+      : QuantizedLinear.load(weights, "lm_head", config);
+    const head = this.lmHead && new QuantizedLinearPlain(this.lmHead);
+    this.#head = head ? hidden => head.forward(hidden) : hidden => this.embed.asLinear(hidden);
+
+    const multiSequencePrompt: MlxCapturingForward = (ids, cache) =>
+    this.forwardLayers(this.embed.encode(ids), cache, "window");
+    this.prompt = [{ from: 1, to: Infinity, operation: multiSequencePrompt }];
+    this.promptChunk = this.prompt;
+    this.multiSequencePrompt = this.prompt;
+    this.multiSequencePromptChunk = this.multiSequencePrompt;
+    this.appendKnown = [{from:1,to:Infinity,operation:(ids, cache) =>
+    this.forwardLayers(this.embed.encode(ids), cache, "window")}];
+    this.logits = rowTable<MlxHeadForward>(hidden => this.logitsFromHidden(hidden), maxPhaseRows);
+    this.decode = rowTable<MlxPhaseForward>((ids, cache) =>
+    this.forwardLayers(this.embed.encode(ids), cache, "decode"), maxPhaseRows);
+    this.verify = rowTable<MlxVerifyForward>((ids, cache) =>
+    this.forwardLayers(this.embed.encode(ids), cache, "window"), maxPhaseRows);
+
+
+    this.sequence = [{from:1,to:Infinity,operation:(ids,cache)=>this.sequenceHidden(ids,cache,1)}];
+
+    this.cacheLayers = this.layers.map(() => ({kind:"full" as const,turboQuant:"dense" as const}));
+  }
+
+  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.attn.qProj);
+      out.set(`${p}.self_attn.k_proj`, l.attn.kProj);
+      out.set(`${p}.self_attn.v_proj`, l.attn.vProj);
+      out.set(`${p}.self_attn.o_proj`, l.attn.oProj);
+      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);
+    }
+    return out;
+  }
+
+  /** Builds no constants and opens no weights: nothing to release. */
+  dispose(): void {}
+
+  /** One cache per layer, sharing one mask memo. Their TurboQuant read is dense
+   * (`denseTurboQuantKernels`): the read decodes keys and values, then runs the
+   * fused SDPA. */
+  readonly cacheLayers: readonly MlxCacheLayer[];
+  readonly sequence: readonly MlxSequenceRange[];
+
+
+  /** The scoring forward (`MlxSequenceForward`): the window read at every width. */
+  sequenceHidden(ids: MlxArray, cache: Cache[], _width: number): MlxArray {
+    return this.forwardLayers(this.embed.encode(ids), cache, "window");
+  }
+
+  /** One whole chunk of a prompt; rows are sequences times the chunk size. The GPU is this graph's one unit. */
+  readonly prompt: readonly MlxPromptRange[];
+  readonly promptChunk: typeof this.prompt;
+  readonly multiSequencePromptChunk: typeof this.multiSequencePrompt;
+  readonly multiSequencePrompt: readonly MlxPromptRange[];
+
+  /** A prompt's positions that are not a whole chunk: the window read at every width. */
+
+  /** Known tokens: this graph declares no committed spans, so they read as a window. */
+  readonly appendKnown: import("../../contracts/mlx/graph").MlxKnownOperations;
+
+  /** The one output head at every width. */
+  readonly logits: MlxRowTable<MlxHeadForward>;
+
+  /** One new token per sequence: `ids` [sequences, 1]. */
+  readonly decode: MlxRowTable<MlxCapturingForward>;
+
+  /** A speculative window of depth + 1 positions per sequence. This graph
+   * declares no hidden-layer taps, so `capture` is always null. */
+  readonly verify: MlxRowTable<MlxCapturingForward>;
+
+  /** The decoder over `h0` (consumed), every block reading its cache for `phase`. */
+  protected forwardLayers(h0: MlxArray, cache: Cache[], phase: ReadPhase): MlxArray {
+    let cur = h0;
+    for (let i = 0; i < this.layers.length; i++) {
+      const next = this.layers[i]!.forward(cur, cache[i]!, phase);
+      cur.dispose();
+      cur = next;
+    }
+    return disposing(cur, this.finalNorm.forward(cur));
+  }
+
+  /** Last-token pooled, L2-normalized sentence embedding for one sequence.
+   *  ids [1, L] → vector [1, hidden]. The caller appends the EOS/eod token
+   *  (Qwen3-Embedding's pooling token) to `ids`. */
+  embedPooled(this: Qwen3Model & { makeCache(): Cache[] }, ids: MlxArray): MlxArray {
+    const cache = this.makeCache();
+    try {
+      const h = this.forwardLayers(this.embed.encode(ids), cache, "window"); // [1, L, hidden]
+      const [B, L, H] = h.shape as [number, number, number];
+      const last = h.slice([0, L - 1, 0], [B, L, H]); // [1, 1, hidden]
+      h.dispose();
+      const vec = disposing(last, ops.reshape(last, [B, H])); // [1, hidden]
+      // L2 normalize: x / max(||x||_2, eps)
+      const sq = ops.square(vec);
+      const ss = disposing(sq, ops.sumAxis(sq, -1, true));
+      const norm = disposing(ss, ops.sqrt(ss));
+      const out = ops.div(vec, norm);
+      vec.dispose();
+      norm.dispose();
+      return out;
+    } finally {
+      for (const c of cache) c.dispose();
+    }
+  }
+
+  logitsFromHidden(h: MlxArray): MlxArray {
+    return this.#head(h);
+  }
+}