mlx-bun 0.5.0 → 1.0.0

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

@@ -1,171 +0,0 @@
-// 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

@@ -1,230 +0,0 @@
-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 = [];
-  }
-}