mlx-bun 0.0.10 → 0.0.12

package/src/spec/dspark/quantize-drafter.ts

@@ -0,0 +1,91 @@
+// Quantize a DeepSpec-format drafter checkpoint (Phase 1a of the DSpark
+// serving program, docs/design/dspark-serving-program.md).
+//
+// Thin policy layer over the house quantizer (src/quantize): the generic
+// shape-eligibility rule already captures every 2-D matmul weight — the
+// layer projections, fc, lm_head, markov_w2, and the two gather tables
+// (embed_tokens, markov_w1: quantized-embedding gather-dequant at load,
+// mlx QuantizedEmbedding semantics). The drafter-specific overlay:
+//
+//   KEEP bf16: confidence_head.* — threshold-sensitive (the truncation
+//   comparison runs on its sigmoid; borderline flips change proposal
+//   lengths) and tiny (~4k params), so quantizing it buys nothing.
+//   Norms / layer_scalar / biases are 1-D and never shape-eligible.
+//
+// Drafter numerics only move ACCEPTANCE, never correctness (the target
+// verifies every draft) — the quality gate for any scheme swept through
+// here is the Phase-1c acceptance A/B, not a KL battery.
+
+import { quantizeModelDir, type QuantizeResult } from "../../quantize";
+import { DeepspecDrafter } from "./deepspec-module";
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+
+/** Module bases kept full-precision regardless of shape eligibility. */
+const KEEP_BF16 = [/^confidence_head\./];
+
+/** The drafter per-tensor policy: quantize everything shape-eligible
+ *  except the KEEP_BF16 set. */
+export function drafterQuantizePredicate(base: string): boolean {
+  return !KEEP_BF16.some((re) => re.test(base));
+}
+
+export interface QuantizeDrafterOptions {
+  bits: 4 | 8;
+  groupSize: 32 | 64;
+  /** Skip the post-write load smoke (tests exercise loading separately). */
+  skipLoadSmoke?: boolean;
+  onProgress?: (stage: string, message: string, progress: number) => void;
+}
+
+/**
+ * Quantize the DeepSpec drafter at `srcDir` into `outDir` and smoke the
+ * result: reload through DeepspecDrafter (quantized detection + graph
+ * build) and eval one projectContext row through the quantized fc path.
+ */
+export async function quantizeDrafterDir(
+  srcDir: string,
+  outDir: string,
+  opts: QuantizeDrafterOptions,
+): Promise<QuantizeResult> {
+  const raw = (await Bun.file(`${srcDir}/config.json`).json()) as Record<string, any>;
+  if (raw.architectures?.[0] !== "Gemma4DSparkModel")
+    throw new Error(
+      `quantizeDrafterDir: expected a DeepSpec checkpoint (architectures[0]==="Gemma4DSparkModel"), got ${JSON.stringify(raw.architectures)} — for regular models use \`mlx-bun quantize\``,
+    );
+  if (raw.quantization ?? raw.quantization_config)
+    throw new Error("quantizeDrafterDir: source is already quantized — start from the bf16 checkpoint");
+
+  const result = await quantizeModelDir(
+    srcDir,
+    outDir,
+    {
+      bits: opts.bits,
+      groupSize: opts.groupSize,
+      quantizePredicate: drafterQuantizePredicate,
+    },
+    opts.onProgress ? (e) => opts.onProgress!(e.stage, e.message, e.progress) : undefined,
+  );
+
+  if (!opts.skipLoadSmoke) {
+    // Load smoke: quantized detection + one real quantized matmul (fc) so a
+    // wrong spec/layout fails HERE, not at serve time.
+    const d = await DeepspecDrafter.load(outDir);
+    try {
+      const width = d.tapLayers.length * d.hidden;
+      const zeros = MlxArray.fromFloat32(new Float32Array(width), [1, 1, width]);
+      const x = zeros.astype(Dtype.bfloat16);
+      zeros.dispose();
+      const proj = d.projectContext(x);
+      x.dispose();
+      const got = proj.shape;
+      proj.dispose();
+      if (got.length !== 3 || got[2] !== d.hidden)
+        throw new Error(`load smoke: projectContext shape ${JSON.stringify(got)}, expected [1,1,${d.hidden}]`);
+    } finally {
+      d.dispose();
+    }
+  }
+
+  return result;
+}

package/src/spec/generate.ts

@@ -128,18 +128,36 @@ export function specGenerate(
     full: (caches[donors.full] as never as { temporalView(): [MlxArray, MlxArray] }).temporalView(),
   });
 
+  // Round/prefill scratch tensors, hoisted so the finally frees them on a
+  // mid-round throw (the 2f not-trimmable throw, or any mlx op) — try-body
+  // locals are invisible to the finally. Each is nulled at its normal
+  // disposal, so the finally never double-frees. Same discipline as
+  // src/spec/serve-loop.ts / generate-dflash.ts (2026-07-06 review fix);
+  // see memory [[mlx-inline-slice-leak-pattern]].
+  let prefillH: MlxArray | null = null;
+  let lastHidden: MlxArray | null = null;
+  let vHidden: MlxArray | null = null;
+
   try {
-    // 1. prefill
+    // 1. prefill — FULL-prompt single forward, deliberately. Per-oracle
+    // convention (2026-07-07 re-anchor): mlx-lm's speculative path drains to
+    // len-1 with no step-0 (the SERVE loop, whose oracle it is, now mirrors
+    // that — src/spec/serve-loop.ts); THIS loop's oracle is optiq
+    // spec_generate, which prefills the whole prompt in one forward and
+    // samples token0 from it (runtime/spec/runtime.py:162-170, read from the
+    // installed oracle venv) — the 2026-06-14 bit-exact gate was established
+    // against exactly that shape. Do not tail-split here.
     const t0 = performance.now();
     const ids = ops.fromInt32(promptTokens, [1, promptTokens.length]);
-    const hidden = model.forwardHidden(ids, caches);
+    prefillH = 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();
+    const H = prefillH.shape[2]!;
+    const Lp = prefillH.shape[1]!;
+    let next = pickFromHidden(model, prefillH, Lp - 1);
+    lastHidden = prefillH.slice([0, Lp - 1, 0], [1, Lp, H]);
+    prefillH.dispose();
+    prefillH = null;
     stats.prefillMs = performance.now() - t0;
 
     // EOS convention: an EOS id stops generation and counts toward
@@ -158,30 +176,38 @@ export function specGenerate(
       const position = caches[0]!.offset - 1;
       const shared = readDonors();
 
-      // 2a. draft γ tokens against the frozen donor views
+      // 2a. draft γ tokens against the frozen donor views. try/finally so a
+      // mid-draft throw still frees the chained hiddens + donor views + the
+      // in-flight emb (assistant-source.ts draft() discipline).
       const drafts: number[] = [];
       let dTok = next;
-      let dHid = lastHidden; // borrowed for k=0
+      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();
+      let emb: MlxArray | null = null;
+      try {
+        for (let k = 0; k < gamma; k++) {
+          emb = embedScaled(dTok);
+          const step = drafter.forward(emb, dHid, shared, position + k);
+          emb.dispose();
+          emb = null;
+          drafts.push(step.token);
+          stats.drafted++;
+          ownedHiddens.push(step.nextHidden);
+          dTok = step.token;
+          dHid = step.nextHidden;
+        }
+      } finally {
+        emb?.dispose();
+        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);
+      vHidden = model.forwardHidden(vIds, caches);
       vIds.dispose();
       stats.targetCalls++;
 
@@ -199,6 +225,7 @@ export function specGenerate(
         stats.emitted++;
         if (isEos || stats.emitted >= maxTokens) {
           vHidden.dispose();
+          vHidden = null;
           break outer;
         }
       }
@@ -221,18 +248,25 @@ export function specGenerate(
 
       if (emitIsEos || stats.emitted >= maxTokens) {
         vHidden.dispose();
+        vHidden = null;
         break;
       }
 
       // 2g. chain state
       next = emit;
-      lastHidden.dispose();
+      lastHidden!.dispose();
       lastHidden = vHidden.slice([0, kAccept, 0], [1, kAccept + 1, H]);
       vHidden.dispose();
+      vHidden = null;
     }
     stats.decodeMs = performance.now() - tDecode;
     lastHidden.dispose();
+    lastHidden = null;
   } finally {
+    // Survivors of a mid-round throw (each nulled at its normal disposal).
+    prefillH?.dispose();
+    lastHidden?.dispose();
+    vHidden?.dispose();
     for (const c of caches) c.dispose();
   }
 

package/src/spec/glm52-mtp-source.ts

@@ -0,0 +1,392 @@
+// Native GLM-5.2 MTP draft source.
+//
+// The target and MTP row share embeddings, the output head, and the resident
+// dense target weights. Only the MTP layer's own dense tensors and int8 routed
+// experts add model weight residency. The generic speculative serve loop owns
+// verification, grammar composition, target-cache rollback, sampling, and
+// telemetry. Prompt lookup remains the model-free alternative DraftProvider;
+// selecting one provider per request prevents two independent draft histories
+// from advancing the target or grammar state twice.
+
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { toLogprobs } from "../sampler";
+import { Glm52DecoderLayer, Glm52Model } from "../model/glm52";
+import { MLACache } from "../model/glm52-cache";
+import { rmsNormF32Mlx } from "../model/glm52-mla";
+import type { DraftProvider, DraftSource } from "./source";
+
+type Sampler = (logprobs: MlxArray, step: number) => MlxArray;
+
+function lastToken(feed: readonly number[]): number {
+  const token = feed.at(-1);
+  if (!Number.isSafeInteger(token) || token! < 0)
+    throw new Error("native MTP requires a non-empty non-negative token feed");
+  return token!;
+}
+
+export class Glm52NativeMtpProvider implements DraftProvider {
+  readonly id = "glm52-native-mtp";
+  readonly weightsBytes = 0;
+  readonly #layer: Glm52DecoderLayer;
+
+  constructor(readonly model: Glm52Model) {
+    if (!model.capabilities.mtpMetadata)
+      throw new Error("GLM-5.2 artifact does not contain a complete MTP row");
+    if (model.capabilities.mtpEnabled === false)
+      throw new Error("native GLM-5.2 MTP is disabled for this model instance");
+    if (model.glmConfig.numNextnPredictLayers !== 1) {
+      throw new Error(
+        `native GLM-5.2 MTP requires exactly one next-token layer; got ` +
+        `${model.glmConfig.numNextnPredictLayers}`,
+      );
+    }
+    const mtpBackend = model.expertRuntime?.mtpExecutor ?? null;
+    if (model.expertBackend && !mtpBackend) {
+      throw new Error(
+        "streamed GLM-5.2 MTP requires the bounded int8 MTP expert tier",
+      );
+    }
+    this.#layer = new Glm52DecoderLayer(
+      model.glmConfig,
+      model.weights,
+      model.glmConfig.numHiddenLayers,
+      false,
+      mtpBackend,
+    );
+  }
+
+  open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
+    if (opts.target.model !== this.model)
+      throw new Error("native MTP provider was opened for a different target");
+    return new Glm52NativeMtpSource(this.model, this.#layer, opts.sampler);
+  }
+
+  dispose(): void {
+    // The target model owns every shared/MTP weight and the expert tier.
+  }
+}
+
+export class Glm52NativeMtpSource implements DraftSource {
+  readonly prefillMode = "full" as const;
+  readonly pinTargetKernelFamily = true;
+  readonly weightsBytes = 0;
+  #cache: MLACache;
+  #roundStart = 0;
+  #lastDraftCount = 0;
+  #closed = false;
+
+  constructor(
+    readonly model: Glm52Model,
+    readonly layer: Glm52DecoderLayer,
+    readonly sampler: Sampler,
+  ) {
+    this.#cache = new MLACache({
+      kvLoraRank: model.glmConfig.kvLoraRank,
+      ropeHeadDim: model.glmConfig.qkRopeHeadDim,
+      maxTokens: model.glmConfig.maxPositionEmbeddings,
+      role: "mtp",
+    });
+  }
+
+  get cacheOffset(): number {
+    return this.#cache.offset;
+  }
+
+  get cacheBytes(): number {
+    return this.#cache.byteLength;
+  }
+
+  /** Caller-owned zero-copy snapshot for the v3 KV persistence writer. */
+  clonePersistentCache(): MLACache {
+    this.#checkOpen();
+    if (this.#lastDraftCount !== 0)
+      throw new Error("native MTP cache cannot snapshot an uncommitted round");
+    if (this.#cache.offset === 0)
+      throw new Error("native MTP cache cannot snapshot empty state");
+    const state = this.#cache.fetch();
+    const clone = new MLACache({
+      kvLoraRank: this.#cache.kvLoraRank,
+      ropeHeadDim: this.#cache.ropeHeadDim,
+      maxTokens: this.#cache.maxTokens,
+      role: "mtp",
+    });
+    try {
+      clone.restoreCompressedState(state.latent, state.rope, null, this.#cache.offset);
+      return clone;
+    } catch (error) {
+      state.latent.dispose();
+      state.rope.dispose();
+      clone.dispose();
+      throw error;
+    }
+  }
+
+  /**
+   * Adopt one restored `mtp-mla` cache. Ownership transfers only after every
+   * role/geometry/offset check passes; the source disposes it thereafter.
+   */
+  restorePersistentCache(cache: MLACache): void {
+    this.#checkOpen();
+    if (this.#lastDraftCount !== 0)
+      throw new Error("native MTP cache cannot restore during a draft round");
+    if (this.#cache.offset !== 0)
+      throw new Error("native MTP cache restore requires an empty source");
+    if (cache.role !== "mtp" || cache.dsa)
+      throw new Error("native MTP restore requires one mtp-mla cache");
+    if (
+      cache.kvLoraRank !== this.#cache.kvLoraRank ||
+      cache.ropeHeadDim !== this.#cache.ropeHeadDim ||
+      cache.maxTokens !== this.#cache.maxTokens
+    ) {
+      throw new Error("native MTP restored cache geometry does not match model");
+    }
+    if (cache.offset <= 0 || cache.batchSize !== 1)
+      throw new Error("native MTP restored cache must contain one non-empty row");
+    this.#cache.dispose();
+    this.#cache = cache;
+    this.#roundStart = cache.offset;
+  }
+
+  prefill(_promptIds: number[]): void {
+    this.#checkOpen();
+    if (this.#cache.offset !== 0)
+      throw new Error("native MTP source cannot be prefilled twice");
+    // Direct Colibri intentionally leaves MTP KV empty here ("niente
+    // prefill"): its first draft opens a decode-only window from the target
+    // anchor hidden and pending token. Seeding this cache with the prompt
+    // would change both the oracle state and the first speculative round.
+  }
+
+  async draft(
+    feed: number[],
+    n: number,
+    stepBase: number,
+    anchorHidden?: MlxArray,
+  ): Promise<number[]> {
+    this.#checkOpen();
+    if (!anchorHidden)
+      throw new Error("native MTP drafting requires the target anchor hidden");
+    if (!Number.isSafeInteger(n) || n < 0)
+      throw new Error(`native MTP draft length must be non-negative; got ${n}`);
+    if (this.#lastDraftCount !== 0)
+      throw new Error("native MTP draft called before the prior round committed");
+    const [batch, tokens, hidden] = anchorHidden.shape;
+    if (batch !== 1 || tokens !== 1 || hidden !== this.model.glmConfig.hiddenSize) {
+      throw new Error(
+        `native MTP anchor must be [1,1,${this.model.glmConfig.hiddenSize}], ` +
+        `got [${anchorHidden.shape.join(",")}]`,
+      );
+    }
+
+    this.#roundStart = this.#cache.offset;
+    let token = lastToken(feed);
+    let state: MlxArray | null = ops.contiguous(anchorHidden);
+    const drafts: number[] = [];
+    try {
+      for (let position = 0; position < n; position++) {
+        const nextState = await this.#step(token, state, stepBase + position);
+        state.dispose();
+        state = nextState.hidden;
+        token = nextState.token;
+        drafts.push(token);
+      }
+      this.#lastDraftCount = drafts.length;
+      return drafts;
+    } catch (error) {
+      const appended = this.#cache.offset - this.#roundStart;
+      if (appended > 0) this.#cache.trim(appended);
+      throw error;
+    } finally {
+      state?.dispose();
+    }
+  }
+
+  async commit(
+    d: number,
+    kAccept: number,
+    vCtxML?: MlxArray,
+    verifiedHidden?: MlxArray,
+    acceptedTokens: readonly number[] = [],
+  ): Promise<void> {
+    this.#checkOpen();
+    vCtxML?.dispose();
+    if (d !== this.#lastDraftCount)
+      throw new Error(`native MTP commit d=${d}, expected ${this.#lastDraftCount}`);
+    if (!Number.isSafeInteger(kAccept) || kAccept < 0 || kAccept > d)
+      throw new Error(`native MTP accepted ${kAccept} outside [0,${d}]`);
+    if (acceptedTokens.length !== kAccept)
+      throw new Error(
+        `native MTP received ${acceptedTokens.length} accepted tokens for k=${kAccept}`,
+      );
+    if (kAccept > 0 && !verifiedHidden)
+      throw new Error("native MTP absorption requires verified target hidden rows");
+
+    try {
+      // Only the first speculative MTP row is conditioned on a true target
+      // hidden. Later rows use earlier MTP hiddens, so accepted positions are
+      // rebuilt below from the target's verified hidden window.
+      const speculativeTail = Math.max(d - 1, 0);
+      if (speculativeTail > 0) this.#cache.trim(speculativeTail);
+      if (kAccept > 0)
+        await this.#absorb(acceptedTokens, verifiedHidden!);
+      const expected = this.#roundStart + (d > 0 ? 1 : 0) + kAccept;
+      if (this.#cache.offset !== expected) {
+        throw new Error(
+          `native MTP cache committed to ${this.#cache.offset}, expected ${expected}`,
+        );
+      }
+    } catch (error) {
+      const appended = this.#cache.offset - this.#roundStart;
+      if (appended > 0) this.#cache.trim(appended);
+      throw error;
+    } finally {
+      this.#lastDraftCount = 0;
+    }
+  }
+
+  dispose(): void {
+    if (this.#closed) return;
+    this.#closed = true;
+    this.#cache.dispose();
+  }
+
+  async #step(
+    token: number,
+    hidden: MlxArray,
+    sampleStep: number,
+  ): Promise<{ hidden: MlxArray; token: number }> {
+    const config = this.model.glmConfig;
+    const ids = ops.fromInt32([token], [1, 1]);
+    let embedded: MlxArray | null = null;
+    let embeddedNorm: MlxArray | null = null;
+    let hiddenNorm: MlxArray | null = null;
+    let joined: MlxArray | null = null;
+    let projected: MlxArray | null = null;
+    let output: MlxArray | null = null;
+    let headInput: MlxArray | null = null;
+    let logits: MlxArray | null = null;
+    let logprobs: MlxArray | null = null;
+    let sampled: MlxArray | null = null;
+    try {
+      embedded = this.model.weights.embedding(
+        ids,
+        "model.embed_tokens.weight",
+        config.vocabSize,
+        config.hiddenSize,
+      );
+      embeddedNorm = rmsNormF32Mlx(
+        embedded,
+        this.model.weights.tensor(
+          `model.layers.${config.numHiddenLayers}.enorm.weight`,
+        ),
+        config.rmsNormEps,
+      );
+      hiddenNorm = rmsNormF32Mlx(
+        hidden,
+        this.model.weights.tensor(
+          `model.layers.${config.numHiddenLayers}.hnorm.weight`,
+        ),
+        config.rmsNormEps,
+      );
+      joined = ops.concatAxis([embeddedNorm, hiddenNorm], 2);
+      projected = this.model.weights.linear(
+        joined,
+        `model.layers.${config.numHiddenLayers}.eh_proj.weight`,
+        config.hiddenSize,
+        2 * config.hiddenSize,
+      );
+      output = await this.layer.forwardAsync(projected, this.#cache, null);
+      headInput = rmsNormF32Mlx(
+        output,
+        this.model.weights.tensor(
+          `model.layers.${config.numHiddenLayers}.shared_head.norm.weight`,
+        ),
+        config.rmsNormEps,
+      );
+      logits = this.model.logitsFromHidden(headInput);
+      logprobs = toLogprobs(logits);
+      sampled = this.sampler(logprobs, sampleStep);
+      return {
+        hidden: ops.contiguous(output),
+        token: ops.itemUint32(sampled),
+      };
+    } finally {
+      ids.dispose();
+      embedded?.dispose();
+      embeddedNorm?.dispose();
+      hiddenNorm?.dispose();
+      joined?.dispose();
+      projected?.dispose();
+      output?.dispose();
+      headInput?.dispose();
+      logits?.dispose();
+      logprobs?.dispose();
+      sampled?.dispose();
+    }
+  }
+
+  async #absorb(
+    acceptedTokens: readonly number[],
+    verifiedHidden: MlxArray,
+  ): Promise<void> {
+    const config = this.model.glmConfig;
+    const k = acceptedTokens.length;
+    const ids = ops.fromInt32([...acceptedTokens], [1, k]);
+    let embedded: MlxArray | null = null;
+    let embeddedNorm: MlxArray | null = null;
+    let trueRows: MlxArray | null = null;
+    let hiddenNorm: MlxArray | null = null;
+    let joined: MlxArray | null = null;
+    let projected: MlxArray | null = null;
+    let output: MlxArray | null = null;
+    try {
+      embedded = this.model.weights.embedding(
+        ids,
+        "model.embed_tokens.weight",
+        config.vocabSize,
+        config.hiddenSize,
+      );
+      embeddedNorm = rmsNormF32Mlx(
+        embedded,
+        this.model.weights.tensor(
+          `model.layers.${config.numHiddenLayers}.enorm.weight`,
+        ),
+        config.rmsNormEps,
+      );
+      trueRows = verifiedHidden.slice(
+        [0, 0, 0],
+        [1, k, config.hiddenSize],
+      );
+      hiddenNorm = rmsNormF32Mlx(
+        trueRows,
+        this.model.weights.tensor(
+          `model.layers.${config.numHiddenLayers}.hnorm.weight`,
+        ),
+        config.rmsNormEps,
+      );
+      joined = ops.concatAxis([embeddedNorm, hiddenNorm], 2);
+      projected = this.model.weights.linear(
+        joined,
+        `model.layers.${config.numHiddenLayers}.eh_proj.weight`,
+        config.hiddenSize,
+        2 * config.hiddenSize,
+      );
+      output = await this.layer.forwardAsync(projected, this.#cache, null);
+      output.eval();
+    } finally {
+      ids.dispose();
+      embedded?.dispose();
+      embeddedNorm?.dispose();
+      trueRows?.dispose();
+      hiddenNorm?.dispose();
+      joined?.dispose();
+      projected?.dispose();
+      output?.dispose();
+    }
+  }
+
+  #checkOpen(): void {
+    if (this.#closed) throw new Error("native MTP source is closed");
+  }
+}