mlx-bun 0.3.0 → 0.5.0

package/src/spec/qwen-mtp-module.ts

@@ -0,0 +1,171 @@
+// 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/array";
+import * as ops from "../mlx/ops";
+import type { ModelConfig } from "../config";
+import type { Weights } from "../weights";
+import { disposing, QuantizedLinear, RMSNorm, type Cache, type Mask } from "../model/gemma4-base";
+import { DenseLinear } from "../model/universal/modules";
+import { compiledSwiglu } from "../model/qwen3_5";
+
+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;
+}
+
+/** Qwen3Attention.forward with the companion's dense or quantized projections.
+ *  Attention operations follow src/model/qwen3_5.ts verbatim. */
+class MtpAttention {
+  readonly qProj: MtpLinear;
+  readonly kProj: MtpLinear;
+  readonly vProj: MtpLinear;
+  readonly oProj: MtpLinear;
+  readonly qNorm: RMSNorm;
+  readonly kNorm: RMSNorm;
+  readonly nHeads: number;
+  readonly nKvHeads: number;
+  readonly headDim: number;
+  readonly scale: number;
+  readonly ropeDims: number;
+  readonly ropeBase: number;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string, resources: DisposableStack) {
+    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.ropeDims = Math.trunc(this.headDim * t.partialRotaryFactor);
+    this.ropeBase = t.ropeParameters.full_attention?.ropeTheta ?? 10000;
+    this.qProj = loadMtpLinear(weights, `${prefix}.q_proj`, config, resources);
+    this.kProj = loadMtpLinear(weights, `${prefix}.k_proj`, config, resources);
+    this.vProj = loadMtpLinear(weights, `${prefix}.v_proj`, config, resources);
+    this.oProj = loadMtpLinear(weights, `${prefix}.o_proj`, config, resources);
+    this.qNorm = new RMSNorm(weights.tensor(`${prefix}.q_norm.weight`), t.rmsNormEps);
+    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];
+    const qp = this.qProj.forward(x);
+    const qpr = disposing(qp, ops.reshape(qp, [B, L, this.nHeads, this.headDim * 2]));
+    const [qHeads, gateHeads] = ops.split(qpr, [this.headDim], -1) as [MlxArray, MlxArray];
+    qpr.dispose();
+    const gate = disposing(gateHeads, ops.reshape(gateHeads, [B, L, this.nHeads * this.headDim]));
+
+    let k = this.kProj.forward(x);
+    let v = this.vProj.forward(x);
+
+    let q = this.qNorm.forward(qHeads);
+    qHeads.dispose();
+    q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
+    k = disposing(k, ops.reshape(k, [B, L, this.nKvHeads, this.headDim]));
+    k = disposing(k, this.kNorm.forward(k));
+    k = disposing(k, ops.transposeAxes(k, [0, 2, 1, 3]));
+    v = disposing(v, ops.reshape(v, [B, L, this.nKvHeads, this.headDim]));
+    v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));
+
+    const offsets = cache.ropeOffsetArr;
+    q = disposing(q, offsets
+      ? ops.ropeDynamic(q, this.ropeDims, this.ropeBase, offsets, null)
+      : ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
+    k = disposing(k, offsets
+      ? ops.ropeDynamic(k, this.ropeDims, this.ropeBase, offsets, null)
+      : ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
+
+    const [keys, values] = cache.updateAndFetch(k, v);
+    k.dispose();
+    v.dispose();
+    const attn = ops.sdpa(q, keys, values, this.scale, mask.mode, mask.arr);
+    keys.dispose();
+    values.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 sig = ops.sigmoid(gate);
+    gate.dispose();
+    const gated = ops.mul(merged, sig);
+    merged.dispose();
+    sig.dispose();
+    const out = this.oProj.forward(gated);
+    gated.dispose();
+    return out;
+  }
+}
+
+/** The one MTP decoder block: fc-merge → attention → swiglu MLP → norm. */
+export class MtpModule {
+  readonly fc: MtpLinear;
+  readonly preFcNormEmbedding: RMSNorm;
+  readonly preFcNormHidden: RMSNorm;
+  readonly attn: MtpAttention;
+  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);
+    this.attn = new MtpAttention(weights, config, "layers.0.self_attn", 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 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). 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`). */
+  forward(tokenEmbeds: MlxArray, hiddens: MlxArray, cache: Cache): 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 L = x.shape[1]!;
+    const mask = cache.makeMask(L, null);
+    const xn = this.inputNorm.forward(x);
+    const r = this.attn.forward(xn, mask, cache);
+    xn.dispose();
+    mask.arr?.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));
+  }
+}

package/src/spec/qwen-mtp-rows.ts

@@ -0,0 +1,230 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { materializeCopy } from "../mlx/materialize";
+import { toLogprobs } from "../sampler";
+import { BatchedKVCache } from "../model/batched-kv";
+import { KVCache } from "../model/gemma4-base";
+import type { DraftRowSampling, QwenMtpTarget } from "./source";
+import type { MtpModule } from "./qwen-mtp-module";
+import type { PreparedStateChange } from "../contracts/resources";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../engine/resources";
+
+/** A committed request's draft KV and the preceding true target hidden row.
+ * Construction borrows these arrays; extraction returns owned state. */
+export interface MtpRowState {
+  readonly cache: KVCache;
+  readonly hidden: MlxArray;
+}
+
+/** Sampling owns each request's RNG policy. Scores are [B,V] log-probabilities;
+ * the returned owned [B] token array stays on device through the draft chain. */
+export type MtpRowSampling = DraftRowSampling;
+
+/** The same draft graph and state lifecycle for one or several requests.
+ * Admission, target verification, output publication and persistence belong to
+ * the caller. Rows remain fixed within a draft/commit round. */
+export class QwenMtpRows {
+  #cache = new BatchedKVCache();
+  #hidden: MlxArray | null = null;
+  #prefilled: boolean[] = [];
+  #drafts: number[][] = [];
+  #depth = 0;
+
+  constructor(
+    readonly target: QwenMtpTarget,
+    readonly module: MtpModule,
+    readonly sampling: MtpRowSampling | null,
+    states: readonly (MtpRowState | null)[],
+  ) {
+    this.append(states);
+  }
+
+  get rowCount(): number { return this.#cache.rowOffsets.length; }
+
+  /** Add committed state at a method boundary. Inputs remain caller-owned. */
+  append(states: readonly (MtpRowState | null)[]): void {
+    if (!states.length) return;
+    applyStateChanges([() => this.prepareAppend(states)]);
+  }
+
+  prepareAppend(states: readonly (MtpRowState | null)[]): PreparedStateChange {
+    if (!states.length) return { commit() {}, dispose() {} };
+    const prefilled = [...this.#prefilled, ...states.map(state => state !== null)];
+    const example = this.#hidden ?? states.find(state => state !== null)?.hidden;
+    let hidden: MlxArray | null = null;
+    let cache: BatchedKVCache | null = new BatchedKVCache();
+    const held: { dispose(): void }[] = [];
+    try {
+      if (example) {
+        const zero = (rows: number) => {
+          const array = ops.zeros([rows, 1, this.target.hiddenSize], example.dtype);
+          held.push(array); return array;
+        };
+        hidden = ops.concatAxis([
+          ...(this.rowCount ? [this.#hidden ?? zero(this.rowCount)] : []),
+          ...states.map(state => state?.hidden ?? zero(1)),
+        ], 0);
+      }
+      cache.mergeRows([...(this.rowCount ? [this.#cache] : []), ...states.map(state => {
+        if (state) return state.cache;
+        const empty = new KVCache(); held.push(empty); return empty;
+      })]);
+    } catch (error) { return cleanupFailure(error, () => disposeResources([cache!, ...(hidden ? [hidden] : [])])); }
+    finally { disposeResources(held); }
+    return {
+      commit: () => {
+        const previousCache = this.#cache, previousHidden = this.#hidden;
+        this.#cache = cache!; this.#hidden = hidden;
+        this.#prefilled = prefilled;
+        cache = previousCache; hidden = previousHidden;
+      },
+      dispose: () => {
+        const held = [cache, hidden].filter(value => value !== null);
+        cache = null; hidden = null;
+        disposeResources(held);
+      },
+    };
+  }
+
+  /** Retire/reorder committed requests without involving their sampler policy. */
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.dispose(); return; }
+    using indices = ops.fromInt32([...keep], [keep.length]);
+    const hidden = this.#hidden ? ops.takeAxis(this.#hidden, indices, 0) : null;
+    try { this.#cache.filterRows(keep); }
+    catch (error) { hidden?.dispose(); throw error; }
+    this.#hidden?.dispose(); this.#hidden = hidden;
+    this.#prefilled = keep.map(row => this.#prefilled[row]!);
+  }
+
+  /** Build companion KV from borrowed new target tokens/taps at the same B.
+   * The first token bridges to each restored row's preceding true hidden.
+   * Cold rows discard that bridge. All later rows pair token[p+1] with h[p].
+   * Only KV and the final true hidden are evaluated; unused attention/MLP
+   * outputs remain outside the evaluated prefill graph. */
+  prefill(tokens: MlxArray, context: MlxArray): void {
+    const [B, length] = tokens.shape as [number, number];
+    const H = context.shape[2]!;
+    if (this.#prefilled.some(Boolean)) {
+      using first = tokens.slice([0, 0], [B, 1]);
+      using embeds = this.target.embed(first);
+      const mixed = this.#prefilled.some(value => !value);
+      if (mixed) this.#cache.specRoundBegin();
+      this.module.forward(embeds, this.#hidden!, this.#cache).dispose();
+      if (mixed) this.#cache.specRoundRollback(this.#prefilled.map(Number));
+    }
+    for (let pos = 0; pos + 1 < length; pos += 2048) {
+      const end = Math.min(pos + 2048, length - 1);
+      using ids = tokens.slice([0, pos + 1], [B, end + 1]);
+      using embeds = this.target.embed(ids);
+      using hidden = context.slice([0, pos, 0], [B, end, H]);
+      this.module.forward(embeds, hidden, this.#cache).dispose();
+      ops.evalAll(this.#cache.state());
+    }
+    using tail = context.slice([0, length - 1, 0], [B, length, H]);
+    const hidden = materializeCopy(tail);
+    this.#hidden?.dispose(); this.#hidden = hidden;
+    ops.evalAll([hidden, ...this.#cache.state()]);
+    this.#prefilled.fill(true);
+  }
+
+  materialize(): void {
+    const state = [...this.#cache.state(), ...(this.#hidden ? [this.#hidden] : [])];
+    if (state.length) ops.evalAll(state);
+  }
+
+  /** Align the companion to verified external continuations through the same
+   * true-hidden prefill graph. Different rows retain different prefix lengths. */
+  consume(tokens: MlxArray, context: MlxArray, lengths: readonly number[]): void {
+    this.#cache.specRoundBegin();
+    this.prefill(tokens, context);
+    this.#cache.specRoundRollback([...lengths]);
+    using positions = ops.fromInt32(lengths.map(length => length - 1), [lengths.length, 1, 1]);
+    const hidden = ops.takeAlongAxis(context, positions, 1);
+    this.#hidden!.dispose(); this.#hidden = hidden;
+  }
+
+  /** Nonnegative depth, with every request's first pending token already known.
+   * Build the complete device dependency chain before reading proposal IDs. */
+  draft(pending: readonly number[], depth: number, steps: readonly number[]): number[][] {
+    const B = pending.length;
+    const tokens: MlxArray[] = [];
+    let ids: MlxArray | null = ops.fromInt32([...pending], [B, 1]);
+    let chained: MlxArray | null = null;
+    this.#cache.specRoundBegin();
+    try {
+      // A zero-proposal round still processes the pending target token. Keep
+      // its companion KV aligned so a later round can resume speculation.
+      if (depth === 0) {
+        using embeds = this.target.embed(ids!);
+        this.module.forward(embeds, this.#hidden!, this.#cache).dispose();
+        this.#drafts = Array.from({ length: B }, () => []);
+        this.#depth = 0;
+        this.#hidden!.dispose(); this.#hidden = null;
+        return this.#drafts;
+      }
+      for (let position = 0; position < depth; position++) {
+        using embeds = this.target.embed(ids!);
+        const output = this.module.forward(embeds, chained ?? this.#hidden!, this.#cache);
+        chained?.dispose(); chained = output;
+        using logits = this.target.logitsFromHidden(output);
+        using flat = ops.reshape(logits, [B, logits.shape.at(-1)!]);
+        using logprobs = toLogprobs(flat);
+        const token = this.sampling!.sample(logprobs, steps.map(step => step + position));
+        tokens.push(token);
+        ids!.dispose(); ids = ops.reshape(token, [B, 1]);
+      }
+      using packed = ops.concatAxis(tokens, 0);
+      const read = packed.toIntTokens();
+      this.#drafts = Array.from({ length: B }, (_, row) =>
+        Array.from({ length: depth }, (_, position) => read[position * B + row]!));
+      this.#depth = depth;
+      this.#hidden!.dispose(); this.#hidden = null;
+      return this.#drafts;
+    } finally {
+      ids?.dispose(); chained?.dispose();
+      for (const token of tokens) token.dispose();
+    }
+  }
+
+  /** Resolve draft storage using each request's accepted prefix. The verify
+   * context is borrowed [B,depth+1,H], before the target's final norm. */
+  commit(accepted: readonly number[], context: MlxArray): void {
+    const B = accepted.length, depth = this.#depth;
+    const appended = Math.max(1, depth);
+    const kept = accepted.map(count => Math.min(count + 1, appended));
+    if (kept.every(count => count === appended)) this.#cache.specRoundCommit();
+    else this.#cache.specRoundRollback(kept);
+
+    // Full acceptance needs the last accepted token's missing draft row,
+    // using the TRUE target hidden at its preceding position. Other requests
+    // keep their current coverage; their temporary append is never retained.
+    const extra = accepted.map(count => Number(depth > 0 && count === depth));
+    if (extra.some(Boolean)) {
+      using ids = ops.fromInt32(this.#drafts.map(row => row[depth - 1]!), [B, 1]);
+      using embeds = this.target.embed(ids);
+      using hidden = context.slice([0, depth - 1, 0], [B, depth, context.shape[2]!]);
+      this.#cache.specRoundBegin();
+      this.module.forward(embeds, hidden, this.#cache).dispose();
+      this.#cache.specRoundRollback(extra);
+    }
+    using positions = ops.fromInt32([...accepted], [B, 1, 1]);
+    this.#hidden = ops.takeAlongAxis(context, positions, 1);
+    this.#drafts = []; this.#depth = 0;
+  }
+
+  /** Only between committed rounds. Returned state is independent of the
+   * group's storage and uses the existing single-request cache representation. */
+  extractRow(row: number): MtpRowState {
+    const cache = this.#cache.extractRow(row);
+    try {
+      using view = this.#hidden!.slice([row, 0, 0], [row + 1, 1, this.#hidden!.shape[2]!]);
+      return { cache, hidden: materializeCopy(view) };
+    } catch (error) { cache.dispose(); throw error; }
+  }
+
+  dispose(): void {
+    this.#cache.dispose(); this.#hidden?.dispose(); this.#hidden = null;
+    this.#drafts = []; this.#depth = 0; this.#prefilled = [];
+  }
+}