@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/models/qwen/checkpoint.ts

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

package/src/models/qwen/families.ts

@@ -0,0 +1,57 @@
+import { L1, defineFamily } from "../family";
+
+/** Qwen3.5 hybrid gated-DeltaNet family (model_type qwen3_5 / qwen3_5_text).
+ * Dense MLP only for now: the MoE variant (qwen3_5_moe) is deferred. */
+export const QWEN35_FAMILY = defineFamily({
+  graph: "qwen3.5",
+  accepts: config => (config.modelType === "qwen3_5" || config.modelType === "qwen3_5_text") &&
+    !config.text.enableMoeBlock &&
+    config.text.numExperts === 0 &&
+    config.text.linearNumValueHeads > 0 &&
+    config.text.fullAttentionInterval > 0,
+  hasModelType: modelType => modelType === "qwen3_5" || modelType === "qwen3_5_text",
+  tier: "targeted",
+  label: "qwen3_5",
+  loader: "safetensors",
+  loop: "autoregressive",
+  fidelity: L1,
+  specialization: "dedicated",
+  profileId: "qwen3.5-dedicated",
+  capabilities: ["qwen3.5-graph", "recurrent-state"],
+});
+
+/** Plain Qwen3 (model_type `qwen3`, Qwen3ForCausalLM): a standard dense decoder
+ * with per-head q/k norm and tied embeddings. Also the text-embedding backbone
+ * (mlx-community Qwen3-Embedding-*). Distinct from the qwen3_5 hybrid. */
+export const QWEN3_FAMILY = defineFamily({
+  graph: "qwen3",
+  accepts: config => config.modelType === "qwen3",
+  hasModelType: modelType => modelType === "qwen3",
+  tier: "targeted",
+  label: "qwen3",
+  loader: "safetensors",
+  loop: "autoregressive",
+  fidelity: L1,
+  specialization: "dedicated",
+  profileId: "qwen3-dedicated",
+  capabilities: ["qwen3-graph"],
+  embedding: { terminator: "<|endoftext|>" },
+});
+
+/** Qwen3-MoE (model_type `qwen3_moe`): plain-qwen3 attention 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 const QWEN3_MOE_FAMILY = defineFamily({
+  graph: "qwen3-moe",
+  accepts: config => config.modelType === "qwen3_moe",
+  hasModelType: modelType => modelType === "qwen3_moe",
+  tier: "targeted",
+  label: "qwen3_moe",
+  loader: "safetensors",
+  loop: "autoregressive",
+  fidelity: L1,
+  specialization: "dedicated",
+  profileId: "qwen3-moe-dedicated",
+  capabilities: ["qwen3-moe-graph"],
+});

package/src/models/qwen/media-input.ts

@@ -0,0 +1,161 @@
+import { rowOperationFor } from "../../runtime/row-operations";
+import type { ModelConfig } from "../../artifacts/config";
+import { TURBOQUANT_FOLD_FILE, readTurboQuantFold } from "../../artifacts/turboquant-fold";
+import type { QuantizedEmbedding } from "../../layers/quantized-embedding";
+import { QuantizedLinear } from "../../layers/quantized-linear";
+import { readDraftVocabulary } from "../../artifacts/draft-vocabulary";
+import type { DraftProjection } from "../../contracts/mlx/draft-projection";
+import { mropePositionIds } from "../../layers/qwen-mrope";
+import { type MropeRequestState } from "../../contracts/mlx/positions";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { Cache } from "../../contracts/mlx/cache";
+import type { MediaEncoders, MlxDecodeState, MlxPromptInput, VisionEncoder, MlxPositionedEmbeddingsRange, MlxPositionedRange, MlxPositionedForward } from "../../contracts/mlx/media";
+import type { MlxHeadForward, MlxRowTable } from "../../contracts/mlx/graph";
+import type { DraftProjectionKernels, ResidualBasis, TargetView } from "../../contracts/mlx/draft-target";
+import type { DraftHeadOperations } from "../../contracts/mlx/draft-projection";
+import { existsSync } from "node:fs";
+import { join } from "node:path";
+import * as ops from "@mlx-bun/mlx/ops";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { bindEmbeddingsInput } from "../media-input";
+
+/** What the Qwen graph's media and draft ports read from it. */
+export interface QwenMediaGraph {
+  readonly config: ModelConfig;
+  readonly embed: QuantizedEmbedding;
+  readonly lmHead: QuantizedLinear | null;
+  /** The full-attention layer whose cache supplies logical positions. */
+  readonly faIdx: number;
+  /** The vision prefill phase: a prompt's prepared embeddings (borrowed) at request positions. */
+  prefillEmbeddingsAtPositions?(embeddings: MlxArray, caches: Cache[], positions: MlxArray, width: number): MlxArray;
+  /** A prompt's token ids after a reused prefix, at request positions. */
+  prefillAtPositions?(ids: MlxArray, caches: Cache[], positions: MlxArray, width: number): MlxArray;
+  /** One new token per row at request positions. */
+  decodeAtPositions?(ids: MlxArray, caches: Cache[], positions: MlxArray, width: number): MlxArray;
+  readonly mediaPrefill?: readonly MlxPositionedEmbeddingsRange[];
+  readonly mediaPrompt?: readonly MlxPositionedRange[];
+  readonly mediaDecode?: MlxRowTable<MlxPositionedForward>;
+  readonly logits: MlxRowTable<MlxHeadForward>;
+}
+
+class QwenMediaDecodeState implements MlxDecodeState {
+  readonly key = "qwen-interleaved-mrope";
+  constructor(readonly model: QwenMediaGraph, readonly delta: number) {}
+
+  decode(ids: MlxArray, caches: Cache[], rows: readonly MlxDecodeState[], width: number): MlxArray {
+    const B = rows.length, L = ids.shape[1]!;
+    const cache = caches[this.model.faIdx]!;
+    // The cache supplies logical positions, including left padding, on device.
+    // No offset readback or request-global position mutation is needed.
+    const offsets = (cache as Cache & { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
+    using deltas = ops.fromInt32(rows.map(row => (row as QwenMediaDecodeState).delta), [B]);
+    using scalar = offsets ? null : ops.fromInt32([cache.offset], []);
+    using starts = ops.add(offsets ?? scalar!, deltas);
+    using startRows = ops.reshape(starts, [B, 1]);
+    using indices = ops.arange(0, L, 1, Dtype.int32);
+    using positions = ops.add(startRows, indices);
+    using axis = ops.reshape(positions, [1, B, L]);
+    using grid = ops.concatAxis([axis, axis, axis], 0);
+    return this.model.mediaDecode ? rowOperationFor(this.model.mediaDecode, width)(ids, caches, grid)
+      : this.model.decodeAtPositions!(ids, caches, grid, width);
+  }
+}
+
+/** The prompt owns its grid; only its continuation delta survives into decode. */
+export function bindQwenMediaInput(model: QwenMediaGraph, embeddings: MlxArray,
+  state: MropeRequestState): MlxPromptInput {
+  const input = bindEmbeddingsInput((ids, caches, start, phase, width) => {
+    using positions = mropePositionIds(state, start, ids.shape[1]!);
+    if (start === 0) return model.mediaPrefill
+      ? model.mediaPrefill.find(range => width >= range.from && width <= range.to)!.operation(embeddings, caches, positions)
+      : model.prefillEmbeddingsAtPositions!(embeddings, caches, positions, width);
+    if (phase === "decode") return model.mediaDecode ? rowOperationFor(model.mediaDecode, width)(ids, caches, positions)
+      : model.decodeAtPositions!(ids, caches, positions, width);
+    return model.mediaPrompt
+      ? model.mediaPrompt.find(range => width >= range.from && width <= range.to)!.operation(ids, caches, positions)
+      : model.prefillAtPositions!(ids, caches, positions, width);
+  });
+  return { ...input, decodeState: new QwenMediaDecodeState(model, state.delta) };
+}
+
+/** Tower weights arrive either as the OptiQ-convention sidecar or in the main
+ * shards (mlx-vlm convention); the tower loader handles both. The encoder is a
+ * Qwen3VLVisionTower riding the shared lazy slot; the Qwen media route is its
+ * only consumer and reads it through the structural contract it defines. */
+export async function qwenMediaEncoders(model: QwenMediaGraph, modelDir: string): Promise<MediaEncoders> {
+  const hasSidecar = existsSync(`${modelDir}/optiq/optiq_vision.safetensors`);
+  if (!hasSidecar && model.config.raw.vision_config === undefined) return { vision: null, audio: null };
+  const { Qwen3VLVisionTower } = await import("../vision/qwen3vl");
+  return { vision: () => Qwen3VLVisionTower.load(modelDir) as unknown as VisionEncoder, audio: null };
+}
+
+/** The residual basis a rotated Qwen artifact declares in its fold record
+ * (`turboquant_fold.json`), for its graph's constructor to read once. A record
+ * whose stream carries no R1 rotation declares none. A missing record, or one
+ * without `final_gain`, declares none either: DFlash 2 then drafts in its
+ * trained basis (output exact, acceptance collapsed), and one line names the
+ * file to update. */
+export function qwenResidualBasis(graph: string, modelDir: string): ResidualBasis | undefined {
+  const fold = readTurboQuantFold(modelDir);
+  if (fold && !fold.r1) return undefined;
+  if (fold?.finalGain) return Object.freeze({ r1Seed: fold.seed, finalGain: fold.finalGain });
+  console.error(`[graph] ${graph}: ${join(modelDir, TURBOQUANT_FOLD_FILE)} ${fold ? "records no final_gain" : "is missing"}; ` +
+    "DFlash 2 drafts unrotated (output exact, acceptance collapses) until it records the final-norm gain the fold moved into lm_head");
+  return undefined;
+}
+
+/** Ports draft sources read from this graph: its embedding, head, layer taps,
+ * the MTP geometry a checkpoint-native draft head needs, from a graph whose
+ * residual stream carries a rotation fold that basis, and from a graph whose
+ * composition measured them the kernels of a drafter's own projections (each the
+ * same object on every view) and the head over its draft vocabulary (`vocabulary`,
+ * read when a drafter first reads the projection's). */
+export function qwenDraftTarget(model: QwenMediaGraph, residualBasis?: ResidualBasis,
+  draftProjections?: DraftProjectionKernels, vocabulary?: () => DraftProjection["vocabulary"]): TargetView {
+  const projection: DraftProjection = { embed: model.embed, logits: model.logits };
+  // The vocabulary's rows are gathered when a drafter first reads them, not for drafters that never do.
+  if (vocabulary) Object.defineProperty(projection, "vocabulary", { get: vocabulary, enumerable: true });
+  return Object.freeze({ identity: model, ...(residualBasis ? { residualBasis } : {}), ...(draftProjections ? { draftProjections } : {}),
+    hiddenLayerTaps: Object.freeze({
+    layerCount: model.config.text.numHiddenLayers,
+    projection: Object.freeze(projection),
+  }), recurrentMtp: Object.freeze({
+    hiddenSize: model.config.text.hiddenSize, layerCount: model.config.text.numHiddenLayers,
+    embed: model.embed.encode.bind(model.embed), logits: model.logits,
+    vocabularyHead: model.lmHead ? Object.freeze({ w: model.lmHead.w, scales: model.lmHead.scales,
+      biases: model.lmHead.biases, spec: model.lmHead.spec, vocabSize: model.config.text.vocabSize }) : undefined,
+  }) });
+}
+
+/** The head's rows of the artifact's frequency-ranked draft vocabulary (`file`, a path inside the artifact) and their
+ *  projection operations on `kernels` (the head's own layers over those rows), for a graph whose composition
+ *  measured a drafter selecting over them. A row gather of an affine head is exact: the rows' logits are the full
+ *  head's arithmetic per row. The gathered arrays go to `owned`, which the graph frees. */
+export function qwenDraftVocabulary(head: QuantizedLinear, modelDir: string, file: string, vocabSize: number,
+  kernels: (rows: QuantizedLinear) => DraftHeadOperations,
+  owned: { dispose(): void }[]): NonNullable<DraftProjection["vocabulary"]> {
+  const listed = readDraftVocabulary(join(modelDir, file), vocabSize);
+  using index = ops.fromInt32([...listed], [listed.length]);
+  const ids = index.astype(Dtype.uint32);
+  owned.push(ids);
+  const gather = (array: MlxArray) => {
+    using rows = ops.takeAxis(array, index, 0);
+    const copy = ops.contiguous(rows);
+    owned.push(copy);
+    return copy;
+  };
+  const rows = new QuantizedLinear(gather(head.w), gather(head.scales), head.biases ? gather(head.biases) : null, head.spec);
+  ops.evalAll([ids, rows.w, rows.scales, ...(rows.biases ? [rows.biases] : [])]);
+  return Object.freeze({ ids, logits: kernels(rows) });
+}
+
+
+/** Gather a draft vocabulary only when a source first reads it, then retain
+ * the same port and owned weights for every target view. The graph supplies
+ * the concrete projection operations over the gathered rows. */
+export function qwenDraftVocabularyMemo(head: QuantizedLinear, modelDir: string, file: string, vocabSize: number,
+  kernels: (rows: QuantizedLinear) => DraftHeadOperations,
+  owned: { dispose(): void }[]): () => NonNullable<DraftProjection["vocabulary"]> {
+  let vocabulary: NonNullable<DraftProjection["vocabulary"]> | undefined;
+  return () => vocabulary ??= qwenDraftVocabulary(head, modelDir, file, vocabSize, kernels, owned);
+}

package/src/models/qwen/mtp.ts

@@ -0,0 +1,101 @@
+// Qwen MTP numerical graph. Request state, sampling and checkpoint retention
+// belong to its callers; the cache supplies positions and attention storage.
+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 { disposing } from "../../layers/helpers";
+import { QuantizedLinear } from "../../layers/quantized-linear";
+import { RMSNorm } from "../../layers/normalization";
+import { type Cache } from "../../contracts/mlx/cache";
+import type { RecurrentMtpForward, RecurrentMtpModule } from "../../contracts/mlx/drafter";
+import type { MlxRowTable } from "../../contracts/mlx/graph";
+import { DenseLinear } from "../../layers/dense-linear";
+import { compiledSwiglu } from "../../layers/swiglu";
+import { Qwen3Attention } from "./blocks";
+import { rowTable } from "../graph";
+
+type MtpLinear = DenseLinear | QuantizedLinear;
+
+/** Projection tensors belong to Weights; dense transpose views belong to
+ *  the provider's resource stack. Quantized heads keep no dense copy. */
+function loadMtpLinear(
+  weights: Weights, path: string, config: ModelConfig, resources: DisposableStack,
+): MtpLinear {
+  if (weights.has(`${path}.scales`)) return QuantizedLinear.load(weights, path, config);
+  const layer = new DenseLinear(weights.tensor(`${path}.weight`), null);
+  resources.use(layer.wT);
+  return layer;
+}
+
+/** The one MTP decoder block: fc-merge → attention → swiglu MLP → norm. */
+export class MtpModule implements RecurrentMtpModule {
+  readonly fc: MtpLinear;
+  readonly preFcNormEmbedding: RMSNorm;
+  readonly preFcNormHidden: RMSNorm;
+  readonly attn: Qwen3Attention<MtpLinear>;
+  readonly mlpGate: MtpLinear;
+  readonly mlpUp: MtpLinear;
+  readonly mlpDown: MtpLinear;
+  readonly inputNorm: RMSNorm;
+  readonly postAttnNorm: RMSNorm;
+  readonly finalNorm: RMSNorm;
+
+  constructor(weights: Weights, config: ModelConfig, resources: DisposableStack) {
+    const eps = config.text.rmsNormEps;
+    this.fc = loadMtpLinear(weights, "fc", config, resources);
+    this.preFcNormEmbedding = new RMSNorm(weights.tensor("pre_fc_norm_embedding.weight"), eps);
+    this.preFcNormHidden = new RMSNorm(weights.tensor("pre_fc_norm_hidden.weight"), eps);
+    // The target's attention block (blocks.ts) with the companion's dense or quantized heads.
+    this.attn = new Qwen3Attention<MtpLinear>(weights, config, "layers.0.self_attn",
+      (w, path, c) => loadMtpLinear(w, path, c, resources));
+    this.mlpGate = loadMtpLinear(weights, "layers.0.mlp.gate_proj", config, resources);
+    this.mlpUp = loadMtpLinear(weights, "layers.0.mlp.up_proj", config, resources);
+    this.mlpDown = loadMtpLinear(weights, "layers.0.mlp.down_proj", config, resources);
+    this.inputNorm = new RMSNorm(weights.tensor("layers.0.input_layernorm.weight"), eps);
+    this.postAttnNorm = new RMSNorm(weights.tensor("layers.0.post_attention_layernorm.weight"), eps);
+    this.finalNorm = new RMSNorm(weights.tensor("norm.weight"), eps);
+  }
+
+  /** One position per row; the module has one draft kernel set, at every width. */
+  readonly decode: MlxRowTable<RecurrentMtpForward> = rowTable<RecurrentMtpForward>(
+    (tokenEmbeds, hiddens, cache) => this.#block(tokenEmbeds, hiddens, cache, "decode"));
+  /** Several positions per row. */
+  readonly window: RecurrentMtpForward = (tokenEmbeds, hiddens, cache) => this.#block(tokenEmbeds, hiddens, cache, "window");
+
+  /** One block forward over [B,S,·]: token embeddings ([B,S,H], target
+   *  embed_tokens output) paired with hiddens ([B,S,H], target pre-final-norm
+   *  or the module's own chained output), the attention reading its cache for
+   *  `phase`. Appends S rows to `cache`; returns the module output [B,S,H]
+   *  (post final norm — what the target lm_head consumes AND what chains into
+   *  the next step's `hidden`). */
+  #block(tokenEmbeds: MlxArray, hiddens: MlxArray, cache: Cache, phase: "decode" | "window"): MlxArray {
+    const embNorm = this.preFcNormEmbedding.forward(tokenEmbeds);
+    const hidNorm = this.preFcNormHidden.forward(hiddens);
+    const joined = ops.concatAxis([embNorm, hidNorm], 2);
+    embNorm.dispose();
+    hidNorm.dispose();
+    const x = this.fc.forward(joined);
+    joined.dispose();
+
+    // Decoder layer (Qwen3Layer.forward shape).
+    const xn = this.inputNorm.forward(x);
+    const r = this.attn.forward(xn, cache, phase);
+    xn.dispose();
+    const h = ops.add(x, r);
+    x.dispose();
+    r.dispose();
+    const hn = this.postAttnNorm.forward(h);
+    const g = this.mlpGate.forward(hn);
+    const u = this.mlpUp.forward(hn);
+    hn.dispose();
+    const act = compiledSwiglu(g, u);
+    g.dispose(); u.dispose();
+    const m = this.mlpDown.forward(act);
+    act.dispose();
+    const out = ops.add(h, m);
+    h.dispose();
+    m.dispose();
+    return disposing(out, this.finalNorm.forward(out));
+  }
+}