mlx-bun 0.0.11 → 0.0.12

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");
+  }
+}

package/src/spec/serve-loop.ts

@@ -57,13 +57,24 @@ const PREFILL_CHUNK = 2048;
  *  [1,L,m*H] (tapLayers concatenated on the feature axis; index nLayers is the
  *  post-finalNorm sentinel). Non-tapping sources get ctxML=null and never
  *  touch model.hiddenTap. Mirrors generate-dflash.ts forwardTapped. */
-function forwardMaybeTap(
+async function forwardMaybeTap(
   model: RuntimeModel,
   ids: MlxArray,
   caches: Cache[],
   tapLayers: number[] | undefined,
-): { hidden: MlxArray; ctxML: MlxArray | null } {
-  if (!tapLayers) return { hidden: model.forwardHidden(ids, caches), ctxML: null };
+): Promise<{ hidden: MlxArray; ctxML: MlxArray | null }> {
+  if (!tapLayers) {
+    const asyncModel = model as RuntimeModel & {
+      forwardHiddenAsync?: (
+        ids: MlxArray,
+        caches: Cache[],
+      ) => Promise<MlxArray>;
+    };
+    const hidden = typeof asyncModel.forwardHiddenAsync === "function"
+      ? await asyncModel.forwardHiddenAsync(ids, caches)
+      : model.forwardHidden(ids, caches);
+    return { hidden, ctxML: null };
+  }
   const m = model as unknown as {
     hiddenTap: { layers: Set<number>; captured: Map<number, MlxArray> } | null;
   };
@@ -96,7 +107,12 @@ function forwardMaybeTap(
 export interface SpecServeExtras {
   drafted: number;
   accepted: number;
+  rejected: number;
   targetCalls: number;
+  rounds: number;
+  acceptanceLengths: number[];
+  tokensPerForward: number;
+  forwardsSaved: number;
   /** Per-draft-position counters (index = position within a round's block,
    *  0..γ-1): how many rounds drafted/accepted at that position. Drives the
    *  Phase-1c per-position acceptance report (scripts/dspark-drafter-ab.ts)
@@ -149,7 +165,16 @@ export async function specServeRun(
     processors.length > 0 ? ops.fromInt32(promptIds, [promptIds.length]) : null;
 
   const extras: SpecServeExtras = {
-    drafted: 0, accepted: 0, targetCalls: 0, draftedByPos: [], acceptedByPos: [],
+    drafted: 0,
+    accepted: 0,
+    rejected: 0,
+    targetCalls: 0,
+    rounds: 0,
+    acceptanceLengths: [],
+    tokensPerForward: 0,
+    forwardsSaved: 0,
+    draftedByPos: [],
+    acceptedByPos: [],
   };
   const stats: GenerateStats = {
     promptTokens: promptIds.length,
@@ -263,7 +288,9 @@ export async function specServeRun(
     // assistant source's first-round anchor hidden = "the hidden that
     // produced pending" = position len-2, the last drain chunk's final row.
     const t0 = performance.now();
-    const tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+    const tailSplit =
+      src.prefillMode !== "full" &&
+      flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
     const oracleShape = tailSplit && promptIds.length > 1;
     if (oracleShape) {
       let pos = 0;
@@ -272,7 +299,12 @@ export async function specServeRun(
         const n = Math.min(PREFILL_CHUNK, end - pos);
         const chunk = promptIds.slice(pos, pos + n);
         const ids = ops.fromInt32(chunk, [1, chunk.length]);
-        const { hidden: h, ctxML } = forwardMaybeTap(model, ids, caches, tapLayers);
+        const { hidden: h, ctxML } = await forwardMaybeTap(
+          model,
+          ids,
+          caches,
+          tapLayers,
+        );
         ids.dispose();
         if (ctxML) ctxParts.push(ctxML);
         if (pos + n >= end) {
@@ -290,7 +322,12 @@ export async function specServeRun(
       for (let off = 0; off < promptIds.length; off += PREFILL_CHUNK) {
         const chunk = promptIds.slice(off, off + PREFILL_CHUNK);
         const ids = ops.fromInt32(chunk, [1, chunk.length]);
-        const { hidden: h, ctxML } = forwardMaybeTap(model, ids, caches, tapLayers);
+        const { hidden: h, ctxML } = await forwardMaybeTap(
+          model,
+          ids,
+          caches,
+          tapLayers,
+        );
         ids.dispose();
         if (ctxML) ctxParts.push(ctxML);
         if (off + PREFILL_CHUNK >= promptIds.length) {
@@ -317,10 +354,10 @@ export async function specServeRun(
         for (const p of ctxParts) p.dispose();
       }
       ctxParts.length = 0; // parts consumed (transferred as prefillCtx or disposed)
-      src.prefill(promptIds, prefillCtx ?? undefined); // takes ownership of prefillCtx
+      await src.prefill(promptIds, prefillCtx ?? undefined); // takes ownership of prefillCtx
       prefillCtx = null; // ownership transferred
     } else {
-      src.prefill(promptIds);
+      await src.prefill(promptIds);
     }
     stats.prefillMs = performance.now() - t0;
     stats.prefillTps = (promptIds.length / Math.max(stats.prefillMs, 1e-6)) * 1000;
@@ -363,7 +400,12 @@ export async function specServeRun(
       if (!speculating) {
         // plain single-token continuation (post-ring-wrap fallback)
         const ids = ops.fromInt32([pending], [1, 1]);
-        const h = model.forwardHidden(ids, caches);
+        const { hidden: h } = await forwardMaybeTap(
+          model,
+          ids,
+          caches,
+          undefined,
+        );
         ids.dispose();
         const lg = model.logitsFromHidden(h);
         h.dispose();
@@ -406,11 +448,22 @@ export async function specServeRun(
 
       // (a) draft UP TO n tokens (request sampler, mlx-lm parity). The return
       // length d is authoritative — DSpark's confidence scheduler may prune the
-      // block short (source.ts contract; never zero). The assistant source
+      // block short or skip it entirely. The assistant source
       // borrows the target's anchor hidden; two-model/dflash ignore it.
-      const drafts = src.draft(feed, n, stats.generatedTokens, anchorHidden ?? undefined);
+      const drafts = await src.draft(
+        feed,
+        n,
+        stats.generatedTokens,
+        anchorHidden ?? undefined,
+      );
       const d = drafts.length;
       if (d < 0 || d > n) throw new Error(`DraftSource returned ${d} drafts (contract: 0..${n})`);
+      if (Bun.env.MLX_BUN_SPEC_TRACE === "1") {
+        console.error(
+          `[SPEC_TRACE] round=${extras.rounds + 1} ` +
+          `feed=${feed.join(",")} drafts=${drafts.join(",")}`,
+        );
+      }
       // d === 0 (a confidence scheduler skipping the round, DeepSpec ℓ=0
       // semantics) needs NO special case: the verify window degenerates to
       // [pending] alone, the accept walk to the single bonus position, commit
@@ -424,12 +477,24 @@ export async function specServeRun(
       // DSpark's H_ctx). vHidden is retained past the accept walk: its slice at
       // the emitted position is the next round's anchor hidden.
       const vIds = ops.fromInt32([pending, ...drafts], [1, d + 1]);
-      const ft = forwardMaybeTap(model, vIds, caches, tapLayers);
-      vHidden = ft.hidden;
-      vCtxML = ft.ctxML;
-      vIds.dispose();
+      const pinTarget = src.pinTargetKernelFamily === true &&
+        "setSpecKernelPinned" in model;
+      if (pinTarget) model.setSpecKernelPinned(true);
+      try {
+        const ft = await forwardMaybeTap(model, vIds, caches, tapLayers);
+        // Assign these before the lm-head call so the outer cleanup owns them
+        // even if logits construction throws.
+        vHidden = ft.hidden;
+        vCtxML = ft.ctxML;
+        vLogits = model.logitsFromHidden(vHidden);
+      } finally {
+        vIds.dispose();
+        if (pinTarget) model.setSpecKernelPinned(false);
+      }
       extras.targetCalls++;
-      vLogits = model.logitsFromHidden(vHidden); // batched lm-head, ONE matmul
+      // The lm-head remains one logical batched operation. Native GLM MTP's
+      // SPEC_PIN implementation evaluates its rows through the same M=1
+      // quantized family before concatenating them.
 
       // (c) per-position accept walk: sample the target at each position,
       // accept while it reproduces the draft. Sampling is sequential because
@@ -471,6 +536,9 @@ export async function specServeRun(
       vLogits!.dispose();
       vLogits = null;
       extras.accepted += kAccept;
+      extras.rejected = extras.drafted - extras.accepted;
+      extras.rounds++;
+      extras.acceptanceLengths.push(kAccept);
       for (let i = 0; i < kAccept; i++)
         extras.acceptedByPos[i] = (extras.acceptedByPos[i] ?? 0) + 1;
       clearCache();
@@ -481,6 +549,13 @@ export async function specServeRun(
         if ((await onToken(tok)) === false) { halted = true; break; }
         if (stats.generatedTokens >= maxTokens) break;
       }
+      extras.tokensPerForward = extras.rounds > 0
+        ? stats.generatedTokens / extras.rounds
+        : 0;
+      extras.forwardsSaved = Math.max(
+        0,
+        Math.max(0, stats.generatedTokens - 1) - extras.rounds,
+      );
       const stop = sawEos || halted || grammarDone || stats.generatedTokens >= maxTokens;
 
       // (e) roll back the rejected suffix + commit to the source (skipped on
@@ -491,7 +566,13 @@ export async function specServeRun(
       // outlives its parent's dispose).
       if (!stop) {
         if (kAccept < d) for (const c of caches) c.trim(d - kAccept);
-        src.commit(d, kAccept, vCtxML ?? undefined); // takes ownership of vCtxML
+        await src.commit(
+          d,
+          kAccept,
+          vCtxML ?? undefined,
+          vHidden,
+          drafts.slice(0, kAccept),
+        ); // takes ownership of vCtxML; verifiedHidden stays caller-owned
         vCtxML = null;
         const H = vHidden!.shape[2]!;
         anchorHidden?.dispose();

package/src/spec/source.ts

@@ -1,6 +1,6 @@
 // DraftSource — the seam between the serve-time speculative verify loop
 // (src/spec/serve-loop.ts) and whatever produces draft tokens. Designed in
-// docs/design/mlx-lm-tool-parity-plan.md §7 so three drafters share ONE
+// docs/design/mlx-lm-tool-parity-plan.md §7 so every drafter shares ONE
 // verify/accept executor:
 //   - TwoModelSource (src/spec/two-model.ts) — mlx-lm parity (L1 oracle:
 //     mlx_lm.server --draft-model), a full second model. Ignores `target`.
@@ -10,8 +10,11 @@
 //   - DflashSource (src/spec/dflash-source.ts) — DSpark (L3, KL/quality-gated).
 //     Taps the target's multi-layer hiddens (prefill + verify) into a growing
 //     H_ctx (docs/design/dspark-speculative-decoding.md).
+//   - NgramSource (src/spec/ngram-source.ts) — model-free prompt lookup.
+//   - Glm52NativeMtpSource (src/spec/glm52-mtp-source.ts) — the target
+//     artifact's native Colibri MTP row.
 //
-// The three differ ONLY in what fills the draft; the serve loop, admission
+// The sources differ ONLY in what fills the draft; the serve loop, admission
 // accounting, and stats never change. The KV-borrowing sources need target
 // state the two-model source doesn't — carried by the `target` view at open()
 // and the optional prefill/draft/commit arguments below (all no-ops for
@@ -33,6 +36,14 @@ export interface TargetView {
 /** A per-request draft-token producer. Created per generation (owns its own
  *  draft-side state), disposed by the serve loop's finally. */
 export interface DraftSource {
+  /** Target prefill shape required by this source's oracle. Most mlx-lm
+   *  sources leave the final prompt token pending; native Colibri MTP starts
+   *  from a full-prompt target forward. */
+  readonly prefillMode?: "tail-split" | "full";
+  /** Request one fixed target kernel family across the speculative verify
+   * batch. Native GLM MTP uses the direct-Colibri SPEC_PIN contract. */
+  readonly pinTargetKernelFamily?: boolean;
+
   /** Multi-layer target tap the source needs captured on the target's prefill
    *  AND every verify forward (DSpark's H_ctx; e4b {20,31,41,42}). When set,
    *  the serve loop sets model.hiddenTap around those forwards and passes the
@@ -44,7 +55,7 @@ export interface DraftSource {
    *  assistant/dflash: read the target's state — mostly a no-op, but DSpark
    *  seeds H_ctx from `ctxML`, the tapped prefill context [1,Lp,m*H], present
    *  iff tapLayers is set). */
-  prefill(promptIds: number[], ctxML?: MlxArray): void;
+  prefill(promptIds: number[], ctxML?: MlxArray): void | Promise<void>;
 
   /** Propose 0..n tokens (RETURN LENGTH IS AUTHORITATIVE — a source may
    *  return fewer than n, e.g. DSpark's confidence-scheduled draft-length
@@ -59,15 +70,29 @@ export interface DraftSource {
    *  `anchorHidden` is the target's final hidden [1,1,H] at the pending/anchor
    *  position — the assistant source borrows it for its first draft step;
    *  two-model and dflash ignore it. */
-  draft(feed: number[], n: number, stepBase: number, anchorHidden?: MlxArray): number[];
+  draft(
+    feed: number[],
+    n: number,
+    stepBase: number,
+    anchorHidden?: MlxArray,
+  ): number[] | Promise<number[]>;
 
   /** Verify outcome for the last round: kAccept of d accepted, where d is the
    *  length draft() actually RETURNED (≤ n). Two-model: trim the draft cache
    *  by max(d - kAccept - 1, 0) — mlx-lm's rewind rule (generate.py:589-591).
    *  DSpark: grow H_ctx by the accepted window from `vCtxML`, the verified
    *  window's tapped context [1,d+1,m*H] (present iff tapLayers is set), and
-   *  drop the rejected tips. */
-  commit(d: number, kAccept: number, vCtxML?: MlxArray): void;
+   *  drop the rejected tips. Native GLM MTP receives the target's verified
+   *  hidden window [1,d+1,H] in `verifiedHidden` and the kAccept accepted token
+   *  ids in `acceptedTokens`; it uses them to rebuild accepted MTP KV rows
+   *  from target state rather than recursively drafted hidden state. */
+  commit(
+    d: number,
+    kAccept: number,
+    vCtxML?: MlxArray,
+    verifiedHidden?: MlxArray,
+    acceptedTokens?: readonly number[],
+  ): void | Promise<void>;
 
   /** Resident per-request draft weights, for admission accounting (0 when the
    *  provider owns the weights). */