mlx-bun 0.0.4

package/src/spec/drafter.ts

@@ -0,0 +1,261 @@
+// Gemma-4 "-assistant" drafter — port of optiq's
+// runtime/spec/drafters/gemma_assistant.py (itself a port of Ollama's
+// x/models/gemma4/assistant.go, MIT).
+//
+// Q-only 4-layer transformer: K/V come from the TARGET's last sliding-
+// and full-attention donor caches at call time; pre/post projections
+// bridge target hidden space (2560) <-> drafter space (256); output via
+// the centroid-clustered head (2048 centroids, top-32 × 128 tokens).
+//
+// Implementation deviation (argmax-equivalent): the reference scatters
+// candidate scores into full-vocab logits initialised to -1e30 and
+// argmaxes; we argmax over the 4096 candidate scores and map through
+// candidate ids directly — same winner, no 262k materialisation. The
+// drafter's numeric details only influence acceptance RATE, never
+// output correctness (the target's verify decides every token).
+
+import type { ModelConfig } from "../config";
+import { Weights } from "../weights";
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import type { SharedKv } from "../model/gemma4";
+
+export interface DrafterStep {
+  /** Drafted token id. */
+  token: number;
+  /** Post-projected hidden (1,1,backbone) — feeds the next draft step. */
+  nextHidden: MlxArray;
+}
+
+interface Block {
+  layerType: string;
+  headDim: number;
+  ropeBase: number | null;
+  ropeFreqs: MlxArray | null;
+  qProj: MlxArray;
+  oProj: MlxArray;
+  qNorm: MlxArray;
+  inputNorm: MlxArray;
+  postAttnNorm: MlxArray;
+  preFfNorm: MlxArray;
+  postFfNorm: MlxArray;
+  gateProj: MlxArray;
+  upProj: MlxArray;
+  downProj: MlxArray;
+  layerScalar: MlxArray;
+}
+
+export class GemmaAssistantDrafter {
+  readonly hidden: number;
+  readonly nHeads: number;
+  readonly slidingWindow: number;
+  readonly eps: number;
+  readonly numCentroids: number;
+  readonly topK: number;
+  readonly vocabPerCentroid: number;
+  #w: Weights;
+  #blocks: Block[] = [];
+  #prePT: MlxArray; // pre_projection transposed [in, out]
+  #postPT: MlxArray;
+  #centroidsT: MlxArray; // [hidden, numCentroids]
+  #ordering: MlxArray;   // [numCentroids, vocabPerCentroid] int32
+  #embed: MlxArray;      // drafter token embeddings [vocab, hidden]
+  #norm: MlxArray;
+
+  private constructor(w: Weights, cfg: Record<string, any>) {
+    const t = cfg.text_config ?? cfg;
+    this.hidden = t.hidden_size;
+    this.nHeads = t.num_attention_heads;
+    this.slidingWindow = t.sliding_window;
+    this.eps = t.rms_norm_eps;
+    this.numCentroids = cfg.num_centroids;
+    this.topK = cfg.centroid_intermediate_top_k;
+    this.vocabPerCentroid = Math.floor(t.vocab_size / this.numCentroids);
+    this.#w = w;
+
+    const T = (name: string) => w.tensor(name);
+    const transposed = (name: string): MlxArray => {
+      const a = ops.transposeAxes(T(name), [1, 0]);
+      const c = ops.contiguous(a);
+      a.dispose();
+      c.eval();
+      return c;
+    };
+    this.#prePT = transposed("pre_projection.weight");
+    this.#postPT = transposed("post_projection.weight");
+    this.#centroidsT = transposed("masked_embedding.centroids.weight");
+    this.#embed = T("model.embed_tokens.weight");
+    this.#norm = T("model.norm.weight");
+    const ord64 = T("masked_embedding.token_ordering");
+    const ord32 = ord64.astype(Dtype.int32);
+    this.#ordering = ops.reshape(ord32, [this.numCentroids, this.vocabPerCentroid]);
+    this.#ordering.eval();
+    ord32.dispose();
+
+    const ropeP = t.rope_parameters ?? {};
+    for (let i = 0; i < t.num_hidden_layers; i++) {
+      const layerType = t.layer_types[i];
+      const isFull = layerType === "full_attention";
+      const headDim = isFull ? t.global_head_dim : t.head_dim;
+      const rp = ropeP[layerType] ?? {};
+      let ropeBase: number | null = rp.rope_theta ?? (isFull ? 1e6 : 1e4);
+      let ropeFreqs: MlxArray | null = null;
+      const partial = rp.partial_rotary_factor ?? 1.0;
+      if (isFull && partial < 1.0) {
+        // ProportionalRoPE freqs, matching the target's donor-K rope
+        const rotated = Math.floor(headDim * partial);
+        const n = headDim / 2;
+        const freqs = new Float32Array(n).fill(Infinity);
+        for (let k = 0; k < rotated / 2; k++)
+          freqs[k] = Math.pow(ropeBase!, (2 * k) / headDim);
+        ropeFreqs = MlxArray.fromFloat32(freqs, [n]);
+        ropeBase = null;
+      }
+      const p = `model.layers.${i}`;
+      this.#blocks.push({
+        layerType, headDim, ropeBase, ropeFreqs,
+        qProj: transposed(`${p}.self_attn.q_proj.weight`),
+        oProj: transposed(`${p}.self_attn.o_proj.weight`),
+        qNorm: T(`${p}.self_attn.q_norm.weight`),
+        inputNorm: T(`${p}.input_layernorm.weight`),
+        postAttnNorm: T(`${p}.post_attention_layernorm.weight`),
+        preFfNorm: T(`${p}.pre_feedforward_layernorm.weight`),
+        postFfNorm: T(`${p}.post_feedforward_layernorm.weight`),
+        gateProj: transposed(`${p}.mlp.gate_proj.weight`),
+        upProj: transposed(`${p}.mlp.up_proj.weight`),
+        downProj: transposed(`${p}.mlp.down_proj.weight`),
+        layerScalar: T(`${p}.layer_scalar`),
+      });
+    }
+  }
+
+  static async load(modelDir: string): Promise<GemmaAssistantDrafter> {
+    const cfg = (await Bun.file(`${modelDir}/config.json`).json()) as Record<string, any>;
+    return new GemmaAssistantDrafter(await Weights.open(modelDir), cfg);
+  }
+
+  /** f32-internal RMSNorm (reference _RMSNorm casts through float32). */
+  #rms(x: MlxArray, weight: MlxArray): MlxArray {
+    const f = x.astype(Dtype.float32);
+    const n = ops.rmsNorm(f, weight, this.eps);
+    f.dispose();
+    const back = n.astype(x.dtype);
+    n.dispose();
+    return back;
+  }
+
+  /** One draft step. All inputs (1,1,·); sharedKv holds the target's
+   *  donor K/V in chronological order. position = target's last position. */
+  forward(
+    lastTokenEmb: MlxArray, targetHidden: MlxArray,
+    sharedKv: { sliding: [MlxArray, MlxArray]; full: [MlxArray, MlxArray] },
+    position: number,
+  ): DrafterStep {
+    let h = ops.concatAxis([lastTokenEmb, targetHidden], 2); // (1,1,2*backbone)
+    h = disposing(h, ops.matmul(h, this.#prePT));            // (1,1,256)
+
+    for (const blk of this.#blocks) {
+      const [k, v] = blk.layerType === "sliding_attention" ? sharedKv.sliding : sharedKv.full;
+      const kLen = k.shape[2]!;
+
+      // sliding-window additive mask over cached K positions
+      let mask: MlxArray | null = null;
+      if (blk.layerType === "sliding_attention") {
+        const firstCached = position - kLen + 1;
+        const windowStart = Math.max(0, position - this.slidingWindow + 1);
+        const allowedFrom = Math.max(0, windowStart - firstCached);
+        if (allowedFrom > 0) {
+          const m = new Float32Array(kLen);
+          for (let i = 0; i < allowedFrom; i++) m[i] = -1e9;
+          const mf = MlxArray.fromFloat32(m, [1, 1, 1, kLen]);
+          mask = mf.astype(k.dtype);
+          mf.dispose();
+        }
+      }
+
+      const residual = h;
+      let x = this.#rms(h, blk.inputNorm);
+      // Q-only attention
+      let q = ops.matmul(x, blk.qProj);
+      x.dispose();
+      q = disposing(q, ops.reshape(q, [1, 1, this.nHeads, blk.headDim]));
+      q = disposing(q, this.#rms(q, blk.qNorm));
+      q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
+      q = disposing(q, ops.rope(q, blk.headDim, blk.ropeBase, position, blk.ropeFreqs));
+      let attn = ops.sdpa(q, k, v, Math.pow(blk.headDim, -0.5), mask ? "array" : "", mask);
+      q.dispose();
+      mask?.dispose();
+      attn = disposing(attn, ops.transposeAxes(attn, [0, 2, 1, 3]));
+      attn = disposing(attn, ops.reshape(attn, [1, 1, this.nHeads * blk.headDim]));
+      attn = disposing(attn, ops.matmul(attn, blk.oProj));
+      attn = disposing(attn, this.#rms(attn, blk.postAttnNorm));
+      h = ops.add(residual, attn);
+      residual.dispose();
+      attn.dispose();
+
+      const res2 = h;
+      let f = this.#rms(h, blk.preFfNorm);
+      const g = ops.matmul(f, blk.gateProj);
+      const u = ops.matmul(f, blk.upProj);
+      f.dispose();
+      const act = ops.geluApprox(g);
+      g.dispose();
+      let mlp = ops.mul(act, u);
+      act.dispose();
+      u.dispose();
+      mlp = disposing(mlp, ops.matmul(mlp, blk.downProj));
+      mlp = disposing(mlp, this.#rms(mlp, blk.postFfNorm));
+      h = ops.add(res2, mlp);
+      res2.dispose();
+      mlp.dispose();
+
+      h = disposing(h, ops.mul(h, blk.layerScalar));
+    }
+
+    h = disposing(h, this.#rms(h, this.#norm));
+
+    // centroid decode → drafted token (argmax-equivalent shortcut)
+    const cScores = ops.matmul(h, this.#centroidsT);          // (1,1,2048)
+    const flat = ops.reshape(cScores, [1, this.numCentroids]);
+    cScores.dispose();
+    const neg = ops.neg(flat);
+    flat.dispose();
+    const part = ops.argpartitionAxis(neg, this.topK - 1, -1);
+    neg.dispose();
+    const topIdx = part.slice([0, 0], [1, this.topK]);        // (1, topK)
+    part.dispose();
+    const topFlat = ops.reshape(topIdx, [this.topK]);
+    topIdx.dispose();
+    const candIds = ops.takeAxis(this.#ordering, topFlat, 0); // (topK, vpc) int32
+    topFlat.dispose();
+    const candFlat = ops.reshape(candIds, [this.topK * this.vocabPerCentroid]);
+    candIds.dispose();
+    const candEmb = ops.takeAxis(this.#embed, candFlat, 0);   // (topK*vpc, hidden)
+    const hVec = ops.reshape(h, [this.hidden, 1]);
+    const scores = ops.matmul(candEmb, hVec);                  // (topK*vpc, 1)
+    candEmb.dispose();
+    hVec.dispose();
+    const scoresFlat = ops.reshape(scores, [1, this.topK * this.vocabPerCentroid]);
+    scores.dispose();
+    const am = ops.argmaxAxis(scoresFlat, -1);
+    scoresFlat.dispose();
+    const winner = ops.itemUint32(am);
+    am.dispose();
+    const tokArr = candFlat.slice([winner], [winner + 1]);
+    candFlat.dispose();
+    const tokU32 = tokArr.astype(Dtype.uint32);
+    tokArr.dispose();
+    const token = ops.itemUint32(tokU32);
+    tokU32.dispose();
+
+    const nextHidden = ops.matmul(h, this.#postPT);           // (1,1,2560)
+    h.dispose();
+    return { token, nextHidden };
+  }
+}
+
+function disposing(old: MlxArray, next: MlxArray): MlxArray {
+  old.dispose();
+  return next;
+}

package/src/spec/generate.ts

@@ -0,0 +1,221 @@
+// Speculative decoding loop — port of optiq runtime/spec/runtime.py
+// (γ-draft / verify-in-one-pass / accept-longest-prefix / rollback).
+//
+// Correctness contract: greedy spec output is token-for-token identical
+// to greedy non-spec decode — the verify step re-derives every emitted
+// token with the SAME pick function as the stock path (logprobs argmax
+// over [1, V]); drafts only ever skip target forwards, never change
+// what the target would emit.
+//
+// Limitation (same as the reference): partial-accept rollback requires
+// trimmable caches; RotatingKVCache loses trimability once its ring
+// wraps (offset ≥ sliding_window). We throw, as the reference does.
+
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { Gemma4Model, type Cache } from "../model/gemma4";
+import { toLogprobs } from "../sampler";
+import type { GemmaAssistantDrafter } from "./drafter";
+
+export interface SpecOptions {
+  gamma?: number;
+  maxTokens?: number;
+  eosTokenIds?: number[];
+}
+
+export interface SpecStats {
+  emitted: number;
+  drafted: number;
+  accepted: number;
+  targetCalls: number;
+  decodeMs: number;
+  prefillMs: number;
+  acceptanceRate: number;
+  decodeTps: number;
+}
+
+export interface SpecResult {
+  tokens: number[];
+  stats: SpecStats;
+}
+
+/** The stock path's greedy pick, bit-identical shapes: slice ONE hidden
+ *  position [1,1,H], lm-head on it (exactly what generate() does for a
+ *  decode step), then logprobs argmax over [1, V]. Computing the lm-head
+ *  per position (instead of batched over the verify window) keeps every
+ *  pick in the same kernel shapes as stock decode — batched lm-head
+ *  rounds differently and flips bf16 knife-edges. */
+function pickFromHidden(model: Gemma4Model, hidden: MlxArray, pos: number): number {
+  const H = hidden.shape[2]!;
+  const hSl = hidden.slice([0, pos, 0], [1, pos + 1, H]);
+  const logits = model.logitsFromHidden(hSl);
+  hSl.dispose();
+  const V = logits.shape[2]!;
+  const flat = ops.reshape(logits, [1, V]);
+  logits.dispose();
+  const lp = toLogprobs(flat);
+  flat.dispose();
+  const am = ops.argmaxAxis(lp, -1);
+  lp.dispose();
+  const t = ops.itemUint32(am);
+  am.dispose();
+  return t;
+}
+
+/** Donor caches: the target's LAST sliding and LAST full cache owners
+ *  (port of optiq kv_view.find_donor_layers). */
+function donorIndices(model: Gemma4Model): { sliding: number; full: number } {
+  let sliding = -1;
+  let full = -1;
+  for (let i = 0; i < model.numDonors; i++) {
+    if (model.layers[i]!.layerType === "sliding_attention") sliding = i;
+    else full = i;
+  }
+  if (sliding < 0 || full < 0) throw new Error("missing donor layer type");
+  return { sliding, full };
+}
+
+export function specGenerate(
+  model: Gemma4Model,
+  drafter: GemmaAssistantDrafter,
+  promptTokens: number[],
+  options: SpecOptions = {},
+): SpecResult {
+  const { gamma = 2, maxTokens = 256, eosTokenIds = model.config.eosTokenIds } = options;
+  if (gamma < 1) throw new Error("gamma must be >= 1");
+
+  const caches = model.makeCache();
+  const donors = donorIndices(model);
+  const stats: SpecStats = {
+    emitted: 0, drafted: 0, accepted: 0, targetCalls: 0,
+    decodeMs: 0, prefillMs: 0, acceptanceRate: 0, decodeTps: 0,
+  };
+  const out: number[] = [];
+
+  const embedScaled = (token: number): MlxArray => {
+    const ids = ops.fromInt32([token], [1, 1]);
+    const e = model.embed.encode(ids);
+    ids.dispose();
+    const s = ops.mulScalar(e, model.embedScale);
+    e.dispose();
+    return s;
+  };
+
+  const readDonors = () => ({
+    sliding: (caches[donors.sliding] as never as { temporalView(): [MlxArray, MlxArray] }).temporalView(),
+    full: (caches[donors.full] as never as { temporalView(): [MlxArray, MlxArray] }).temporalView(),
+  });
+
+  try {
+    // 1. prefill
+    const t0 = performance.now();
+    const ids = ops.fromInt32(promptTokens, [1, promptTokens.length]);
+    const hidden = model.forwardHidden(ids, caches);
+    ids.dispose();
+    stats.targetCalls++;
+    const H = hidden.shape[2]!;
+    const Lp = hidden.shape[1]!;
+    let next = pickFromHidden(model, hidden, Lp - 1);
+    let lastHidden = hidden.slice([0, Lp - 1, 0], [1, Lp, H]);
+    hidden.dispose();
+    stats.prefillMs = performance.now() - t0;
+
+    // EOS convention: an EOS id stops generation and counts toward
+    // stats.emitted (reference-faithful: optiq runtime.py counts it) but
+    // is NOT part of the returned content — matching our generate(),
+    // which never yields EOS. optiq's spec yields it only as a stream
+    // EVENT (clients see the stop); its token array role here is content.
+    const nextIsEos = eosTokenIds.includes(next);
+    if (!nextIsEos) out.push(next);
+    stats.emitted++;
+    if (nextIsEos) return { tokens: out, stats };
+
+    // 2. outer loop
+    const tDecode = performance.now();
+    outer: while (stats.emitted < maxTokens) {
+      const position = caches[0]!.offset - 1;
+      const shared = readDonors();
+
+      // 2a. draft γ tokens against the frozen donor views
+      const drafts: number[] = [];
+      let dTok = next;
+      let dHid = lastHidden; // borrowed for k=0
+      const ownedHiddens: MlxArray[] = [];
+      for (let k = 0; k < gamma; k++) {
+        const emb = embedScaled(dTok);
+        const step = drafter.forward(emb, dHid, shared, position + k);
+        emb.dispose();
+        drafts.push(step.token);
+        stats.drafted++;
+        ownedHiddens.push(step.nextHidden);
+        dTok = step.token;
+        dHid = step.nextHidden;
+      }
+      for (const a of ownedHiddens) a.dispose();
+      for (const [k, v] of [shared.sliding, shared.full]) {
+        k.dispose();
+        v.dispose();
+      }
+
+      // 2b. verify all γ drafts (+ the pending token) in one forward
+      const vIds = ops.fromInt32([next, ...drafts], [1, gamma + 1]);
+      const vHidden = model.forwardHidden(vIds, caches);
+      vIds.dispose();
+      stats.targetCalls++;
+
+      // 2c. accept the longest matching prefix (per-position lm-head —
+      // see pickFromHidden)
+      const gt: number[] = [];
+      for (let k = 0; k <= gamma; k++) gt.push(pickFromHidden(model, vHidden, k));
+      let kAccept = 0;
+      while (kAccept < gamma && drafts[kAccept] === gt[kAccept]) kAccept++;
+      stats.accepted += kAccept;
+
+      // 2d. emit accepted drafts
+      for (let k = 0; k < kAccept; k++) {
+        const isEos = eosTokenIds.includes(drafts[k]!);
+        if (!isEos) out.push(drafts[k]!);
+        stats.emitted++;
+        if (isEos || stats.emitted >= maxTokens) {
+          vHidden.dispose();
+          break outer;
+        }
+      }
+
+      // 2e. emit the correction (or bonus when all accepted)
+      const emit = gt[kAccept]!;
+      const emitIsEos = eosTokenIds.includes(emit);
+      if (!emitIsEos) out.push(emit);
+      stats.emitted++;
+
+      // 2f. roll back rejected entries
+      if (kAccept < gamma) {
+        const n = gamma - kAccept;
+        for (const c of caches) {
+          if (!c.isTrimmable())
+            throw new Error(`${c.constructor.name} not trimmable — context exceeded the sliding window during spec decode`);
+          c.trim(n);
+        }
+      }
+
+      if (emitIsEos || stats.emitted >= maxTokens) {
+        vHidden.dispose();
+        break;
+      }
+
+      // 2g. chain state
+      next = emit;
+      lastHidden.dispose();
+      lastHidden = vHidden.slice([0, kAccept, 0], [1, kAccept + 1, H]);
+      vHidden.dispose();
+    }
+    stats.decodeMs = performance.now() - tDecode;
+    lastHidden.dispose();
+  } finally {
+    for (const c of caches) c.dispose();
+  }
+
+  stats.acceptanceRate = stats.drafted ? stats.accepted / stats.drafted : 0;
+  stats.decodeTps = stats.decodeMs ? ((stats.emitted - 1) / stats.decodeMs) * 1000 : 0;
+  return { tokens: out, stats };
+}