mlx-bun 0.4.0 → 1.0.0

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

@@ -1,317 +0,0 @@
-// 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 { Glm52MtpGraph } from "./glm52-mtp-graph";
-import type { DraftProvider, DraftSource, GroupedDraftProvider, TargetView, DraftRowSampling, DraftRowCheckpoint, DraftRowGroup, DraftPrefillGroup } from "./source";
-import { Glm52MtpRows, type Glm52MtpRowState } from "./glm52-mtp-rows";
-import { captureGlm52MtpState, restoreGlm52MtpState } from "./glm52-mtp-state";
-import { applyStateChanges, disposeResources } from "../engine/resources";
-
-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 grouped: GroupedDraftProvider = {
-    open: options => this.#openRows(options.target, options.sampling, options.checkpoints),
-    openPrefill: options => this.#openRows(options.target, null, options.checkpoints),
-  };
-  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.identity !== this.model)
-      throw new Error("native MTP provider was opened for a different target");
-    return new Glm52NativeMtpSource(this.model, this.#layer, opts.sampler);
-  }
-
-  #openRows(target: TargetView, sampling: DraftRowSampling | null,
-    checkpoints: readonly (DraftRowCheckpoint | null)[]): DraftRowGroup & DraftPrefillGroup {
-    if (target.identity !== this.model) throw new Error("native MTP provider was opened for a different target");
-    const rows = new Glm52MtpRows(new Glm52MtpGraph(this.model, this.#layer), sampling);
-    const prepareAppend = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => {
-      const states: (Glm52MtpRowState | null)[] = [];
-      try {
-        for (const checkpoint of checkpoints) states.push(checkpoint ? restoreGlm52MtpState(checkpoint, () => rows.makeCache()) : null);
-        return rows.prepareAppend(states);
-      } finally { disposeResources(states.flatMap(state => state ? [state.cache, state.hidden] : [])); }
-    };
-    const append = (states: readonly (DraftRowCheckpoint | null)[]) => applyStateChanges([() => prepareAppend(states)]);
-    try { append(checkpoints); } catch (error) { rows.dispose(); throw error; }
-    return {
-      namespace: "glm52-native-mtp-v1", prefillMode: "full", tapLayers: [],
-      get rowCount() { return rows.rowCount; },
-      append, prepareAppend, prefill: (tokens, context) => rows.prefill(tokens, context!),
-      materialize: rows.materialize.bind(rows), filterRows: rows.filterRows.bind(rows),
-      draft: rows.draft.bind(rows), commit: rows.commit.bind(rows),
-      capture(row) {
-        const state = rows.extractRow(row);
-        try { return captureGlm52MtpState(state); }
-        finally { state.cache.dispose(); state.hidden.dispose(); }
-      },
-      dispose: rows.dispose.bind(rows),
-    };
-  }
-
-  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;
-  readonly #graph: Glm52MtpGraph;
-  #roundStart = 0;
-  #lastDraftCount = 0;
-  #closed = false;
-
-  constructor(
-    readonly model: Glm52Model,
-    readonly layer: Glm52DecoderLayer,
-    readonly sampler: Sampler,
-  ) {
-    this.#graph = new Glm52MtpGraph(model, layer);
-    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 }> {
-    using ids = ops.fromInt32([token], [1, 1]);
-    using output = await this.#graph.forward(ids, hidden, this.#cache);
-    using logits = this.#graph.project(output);
-    // Sampling consumes [B,V], independently of the graph's token dimension.
-    using flat = ops.reshape(logits, [1, logits.shape.at(-1)!]);
-    using logprobs = toLogprobs(flat);
-    using sampled = this.sampler(logprobs, sampleStep);
-    return { hidden: ops.contiguous(output), token: ops.itemUint32(sampled) };
-  }
-
-  async #absorb(acceptedTokens: readonly number[], verifiedHidden: MlxArray): Promise<void> {
-    const count = acceptedTokens.length;
-    using ids = ops.fromInt32([...acceptedTokens], [1, count]);
-    using trueRows = verifiedHidden.slice([0, 0, 0], [1, count, this.model.glmConfig.hiddenSize]);
-    using output = await this.#graph.forward(ids, trueRows, this.#cache);
-    output.eval();
-  }
-
-  #checkOpen(): void {
-    if (this.#closed) throw new Error("native MTP source is closed");
-  }
-}

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

@@ -1,33 +0,0 @@
-import type { MlxArray } from "../mlx/array";
-import type { MLACache } from "../model/glm52-cache";
-import { cloneAttachments } from "../backends/mlx/checkpoint-state";
-import { disposeResources } from "../engine/resources";
-import type { DraftRowCheckpoint } from "./source";
-import type { Glm52MtpRowState } from "./glm52-mtp-rows";
-
-export function captureGlm52MtpState(state: Glm52MtpRowState): DraftRowCheckpoint {
-  const held: MlxArray[] = [];
-  try {
-    if (state.cache.offset) {
-      const cache = state.cache.fetch(); held.push(cache.latent, cache.rope);
-    }
-    held.push(state.hidden.slice([0, 0, 0], [...state.hidden.shape]));
-    return { processedTokens: state.processedTokens, attachment: {
-      schema: "glm52-native-mtp-v1", metadata: { processedTokens: state.processedTokens, draftOffset: state.cache.offset },
-      tensors: held.splice(0),
-    } };
-  } finally { disposeResources(held); }
-}
-
-export function restoreGlm52MtpState(checkpoint: DraftRowCheckpoint, makeCache: () => MLACache): Glm52MtpRowState {
-  const { attachment, processedTokens } = checkpoint;
-  if (attachment.schema !== "glm52-native-mtp-v1" || attachment.metadata.processedTokens !== processedTokens)
-    throw new Error("invalid native GLM MTP checkpoint alignment");
-  const held = cloneAttachments([attachment])[0]!.tensors;
-  let cache: MLACache | null = makeCache();
-  try {
-    const offset = attachment.metadata.draftOffset as number;
-    if (offset) { cache.restoreCompressedState(held[0]!, held[1]!, null, offset); held.splice(0, 2); }
-    const state = { cache, hidden: held.pop()!, processedTokens }; cache = null; return state;
-  } finally { disposeResources([...held, ...(cache ? [cache] : [])]); }
-}

package/src/spec/ngram-source.ts

@@ -1,204 +0,0 @@
-// NgramSource — MODEL-FREE prompt-lookup speculative drafting behind the
-// DraftSource seam (src/spec/source.ts). No weights, no caches, no tokenizer:
-// drafts are copied from the request's own token stream. When the trailing
-// k-gram of (prompt + emitted-so-far) has occurred earlier in that stream,
-// the tokens that followed the earlier occurrence are proposed as the draft;
-// the shared verify/accept executor (src/spec/serve-loop.ts) makes the result
-// LOSSLESS by construction (drafts are only proposals — the target's own
-// samples decide every emitted token, at any temperature; only the acceptance
-// rate moves). Best case: agentic/RAG/code-edit traffic that re-emits spans
-// already present in context.
-//
-// Prior art (ported, not invented): "Prompt Lookup Decoding" (Apoorv Saxena,
-// github.com/apoorvumang/prompt-lookup-decoding) and vLLM's `ngram` proposer
-// (vllm/v1 spec_decode ngram_proposer) — both match the LONGEST k-gram first
-// (k = max..min) and take the FIRST occurrence scanning left-to-right; we
-// follow that exactly. Defaults mirror Saxena (max_ngram_size=3,
-// num_pred_tokens=10 → our --num-draft-tokens default for this kind).
-//
-// Token-history reconstruction (the one subtle part): the seam hands sources
-// the prompt at prefill() and only the FEED tail each round — never the full
-// emitted stream — so this source rebuilds it from the feed/commit discipline
-// (same information two-model keeps in its draft KV):
-//   - prefill: history = prompt to len-1 under the oracle tail-split shape
-//     (the last prompt token arrives as the first round's feed = [pending]);
-//     full prompt under the legacy shape (token0 arrives as its feed).
-//   - draft(feed): history += feed. feed is [correction] after a rejected
-//     round, or [lastDraft, bonus] after an all-accept round (mlx-lm's
-//     re-feed rule, generate.py:645-648) — so the all-accept round's LAST
-//     draft is deliberately NOT pushed at commit (it arrives here instead).
-//   - commit(d, kAccept): history += drafts[0 .. min(kAccept, d-1)) — the
-//     accepted drafts except that re-fed last one. Rejected tips never enter
-//     history; the correction token arrives via the next feed.
-// Invariant at matching time (just after the feed push): history ==
-// prompt + every emitted token, ending at the pending/anchor token.
-//
-// A round with no match returns [] — the serve loop's d=0 semantics degrade
-// it to one plain target step (bit-equivalent to non-spec decode), so ngram
-// speculation is structurally never worse than plain decode by more than the
-// JS scan (~µs against 30k-token histories).
-
-import { flagOn } from "../runtime-config";
-import type { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import type { CheckpointAttachment } from "../backends/mlx/checkpoint-state";
-import type { DraftProvider, DraftSource, DraftRowCheckpoint, DraftRowGroup, DraftPrefillGroup, GroupedDraftProvider } from "./source";
-import { applyStateChanges } from "../engine/resources";
-
-/** The matching policy is independent of the executor's feed convention. */
-function proposeNgram(h: readonly number[], max: number, min: number, n: number): number[] {
-  const L = h.length;
-  for (let k = Math.min(max, L - 1); k >= min; k--) {
-    const tailAt = L - k;
-    search: for (let i = 0; i + k < L; i++) {
-      for (let j = 0; j < k; j++) if (h[i + j] !== h[tailAt + j]) continue search;
-      return h.slice(i + k, i + k + n);
-    }
-  }
-  return [];
-}
-
-const namespace = (max: number, min: number) => `ngram-history-v1:${max}:${min}`;
-const captureHistory = (history: readonly number[], max: number, min: number): CheckpointAttachment => ({
-  schema: "ngram-history-v1", metadata: { max, min }, tensors: [ops.fromInt32([...history], [history.length])],
-});
-function restoreHistory(state: DraftRowCheckpoint, max: number, min: number): number[] {
-  const { attachment, processedTokens } = state;
-  if (attachment.schema !== "ngram-history-v1" || attachment.metadata.max !== max || attachment.metadata.min !== min ||
-      attachment.tensors.length !== 1 || attachment.tensors[0]!.shape.length !== 1 || attachment.tensors[0]!.shape[0] !== processedTokens)
-    throw new Error("ngram checkpoint does not match its token coverage or lookup policy");
-  return attachment.tensors[0]!.toIntTokens();
-}
-
-/** Each row owns only its committed token history. Verification and sampling
- * remain in the shared method; target storage and scheduling are not needed. */
-class NgramRows implements DraftRowGroup, DraftPrefillGroup {
-  readonly prefillMode = "tail-split" as const;
-  readonly tapLayers: readonly number[] = [];
-  readonly namespace: string;
-  #histories: number[][] = [];
-  #drafts: number[][] = [];
-  constructor(readonly max: number, readonly min: number, checkpoints: readonly (DraftRowCheckpoint | null)[]) {
-    this.namespace = namespace(max, min);
-    this.append(checkpoints);
-  }
-  get rowCount() { return this.#histories.length; }
-  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]) {
-    let next = [...this.#histories, ...checkpoints.map(state => state ? restoreHistory(state, this.max, this.min) : [])];
-    return { commit: () => { const old = this.#histories; this.#histories = next; next = old; }, dispose: () => { next = []; } };
-  }
-  append(checkpoints: readonly (DraftRowCheckpoint | null)[]) { applyStateChanges([() => this.prepareAppend(checkpoints)]); }
-  filterRows(keep: readonly number[]) { this.#histories = keep.map(row => this.#histories[row]!); this.#drafts = []; }
-  prefill(tokens: MlxArray): void {
-    using packed = ops.contiguous(tokens);
-    const ids = packed.toIntTokens(), length = tokens.shape[1]!;
-    for (let row = 0; row < this.rowCount; row++) this.#histories[row]!.push(...ids.slice(row * length, (row + 1) * length));
-  }
-  materialize(): void {}
-  draft(pending: readonly number[], depth: number): number[][] {
-    this.#drafts = this.#histories.map((history, row) => {
-      history.push(pending[row]!);
-      return proposeNgram(history, this.max, this.min, depth);
-    });
-    return this.#drafts;
-  }
-  commit(accepted: readonly number[]) {
-    for (let row = 0; row < this.rowCount; row++) this.#histories[row]!.push(...this.#drafts[row]!.slice(0, accepted[row]!));
-    this.#drafts = [];
-  }
-  capture(row: number): DraftRowCheckpoint {
-    const history = this.#histories[row]!;
-    return { processedTokens: history.length, attachment: captureHistory(history, this.max, this.min) };
-  }
-  dispose() { this.#histories = []; this.#drafts = []; }
-}
-
-export interface NgramOptions {
-  /** Longest suffix k-gram tried first (Saxena max_ngram_size). */
-  max?: number;
-  /** Shortest k-gram tried before giving up (vLLM prompt_lookup_min). */
-  min?: number;
-}
-
-export class NgramProvider implements DraftProvider {
-  readonly id = "ngram";
-  readonly weightsBytes = 0;
-  readonly max: number;
-  readonly min: number;
-  readonly grouped: GroupedDraftProvider = {
-    supportsTargetAdapters: true,
-    open: options => new NgramRows(this.max, this.min, options.checkpoints),
-    openPrefill: options => new NgramRows(this.max, this.min, options.checkpoints),
-  };
-
-  constructor(opts: NgramOptions = {}) {
-    this.max = Math.max(1, opts.max ?? 3);
-    this.min = Math.max(1, Math.min(opts.min ?? 1, this.max));
-  }
-
-  open(_opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
-    // Sampler and target both ignored: proposals come from lookup, not a
-    // model, so drafting is sampler-independent (like the assistant source —
-    // correctness holds at any temperature, only acceptance moves).
-    return new NgramSource(this.max, this.min);
-  }
-
-  dispose(): void {}
-}
-
-class NgramSource implements DraftSource {
-  readonly weightsBytes = 0;
-  #hist: number[] = [];
-  #lastDrafts: number[] = [];
-  readonly #tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
-  readonly checkpoint: NonNullable<DraftSource["checkpoint"]>;
-
-  constructor(
-    private readonly max: number,
-    private readonly min: number,
-  ) {
-    this.checkpoint = {
-      namespace: namespace(max, min),
-      capture: () => captureHistory(this.#hist, this.max, this.min),
-      restore: (processedTokens, attachment) => {
-        this.#hist = restoreHistory({ processedTokens, attachment }, this.max, this.min);
-      },
-    };
-  }
-
-  /** Test hook — the reconstructed prompt+emitted stream (see header). */
-  get history(): readonly number[] {
-    return this.#hist;
-  }
-
-  prefill(promptIds: number[], _context?: Parameters<DraftSource["prefill"]>[1], processedTokens?: number): void {
-    // Mirror the serve loop's prefill shape (two-model.ts does the same): under
-    // the oracle tail split the last prompt token is never prefilled — it IS
-    // the first feed. Legacy shape (kill switch / 1-token prompt): full prompt,
-    // and the sampled token0 arrives as the first feed.
-    const upTo = processedTokens ?? (this.#tailSplit && promptIds.length > 1 ? promptIds.length - 1 : promptIds.length);
-    this.#hist = promptIds.slice(0, upTo);
-  }
-
-  draft(feed: number[], n: number, _stepBase: number): number[] {
-    this.#hist.push(...feed);
-    this.#lastDrafts = this.#propose(n);
-    return this.#lastDrafts;
-  }
-
-  commit(d: number, kAccept: number): void {
-    // Accepted drafts join history, EXCEPT an all-accept round's last draft —
-    // the serve loop re-feeds it ([lastDraft, bonus]) and the next draft()'s
-    // feed push would double it (see header). d=0 rounds push nothing.
-    const upTo = Math.max(0, Math.min(kAccept, d - 1));
-    for (let i = 0; i < upTo; i++) this.#hist.push(this.#lastDrafts[i]!);
-  }
-
-  /** Longest-k-first, first-occurrence prompt lookup (Saxena/vLLM order):
-   *  find the trailing k-gram earlier in history, propose what followed it. */
-  #propose(n: number): number[] {
-    return proposeNgram(this.#hist, this.max, this.min, n);
-  }
-
-  dispose(): void {}
-}

package/src/spec/projected-draft-rows.ts

@@ -1,73 +0,0 @@
-import type { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { KVCache, type Cache } from "../model/gemma4-base";
-import { MlxProjectedContextRows } from "../backends/mlx/projected-context-rows";
-import { captureCacheAttachment, restoreCacheAttachment } from "../backends/mlx/cache-attachment";
-import { applyStateChanges, disposeResources } from "../engine/resources";
-import type { PreparedStateChange } from "../contracts/resources";
-import type { DraftPrefillGroup, DraftRowGroup, DraftRowCheckpoint, GroupedDraftProvider, TargetView } from "./source";
-
-export interface ProjectedDraftMethod {
-  namespace: string;
-  schema: string;
-  layers: number;
-  tapLayers: readonly number[];
-  project(hidden: MlxArray, positions: number | MlxArray): { k: MlxArray; v: MlxArray }[];
-  draft(context: ReturnType<MlxProjectedContextRows["readAttention"]>, pending: readonly number[], positions: number | MlxArray, depth: number): number[][];
-}
-
-/** Provider numerics compose with shared context membership and persistence. */
-class ProjectedDraftRows implements DraftPrefillGroup, DraftRowGroup {
-  readonly prefillMode = "tail-split" as const;
-  readonly tapLayers: readonly number[];
-  readonly context: MlxProjectedContextRows;
-  constructor(readonly method: ProjectedDraftMethod,
-    checkpoints: readonly (DraftRowCheckpoint | null)[]) {
-    this.tapLayers = method.tapLayers;
-    this.context = new MlxProjectedContextRows(method.layers, { project: method.project });
-    try { this.append(checkpoints); } catch (error) { this.dispose(); throw error; }
-  }
-  get namespace(): string { return this.method.namespace; }
-  get rowCount(): number { return this.context.state.rowCount; }
-  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange {
-    const restored: Cache[][] = [];
-    const empty = () => Array.from({ length: this.method.layers }, () => new KVCache());
-    try {
-      for (const checkpoint of checkpoints) restored.push(checkpoint ? restoreCacheAttachment(checkpoint.attachment, empty) : empty());
-      return this.context.prepareAppend(restored);
-    } finally { disposeResources(restored.flat()); }
-  }
-  append(checkpoints: readonly (DraftRowCheckpoint | null)[]): void {
-    if (checkpoints.length) applyStateChanges([() => this.prepareAppend(checkpoints)]);
-  }
-  prefill(tokens: MlxArray, hidden?: MlxArray): void {
-    this.context.append(hidden!, Array(this.rowCount).fill(tokens.shape[1]!)); this.materialize();
-  }
-  materialize(): void { this.context.materialize(); }
-  filterRows(keep: readonly number[]): void { this.context.state.filterRows(keep); }
-  capture(row: number): DraftRowCheckpoint {
-    const caches = this.context.state.extractRow(row);
-    try {
-      const processedTokens = caches[0]!.offset;
-      return { processedTokens, attachment: captureCacheAttachment(this.method.schema, caches, { processedTokens }) };
-    } finally { disposeResources(caches); }
-  }
-  draft(pending: readonly number[], depth: number): number[][] {
-    if (!depth) return pending.map(() => []);
-    const positions = this.context.positions, donors = this.context.readAttention();
-    using offsets = positions.every(position => position === positions[0]) ? null : ops.fromInt32([...positions], [pending.length]);
-    try { return this.method.draft(donors, pending, offsets ?? positions[0]!, depth); }
-    finally { donors.dispose(); }
-  }
-  commit(accepted: readonly number[], hidden: MlxArray): void {
-    this.context.append(hidden, accepted.map(count => count + 1));
-  }
-  dispose(): void { this.context.dispose(); }
-}
-
-export function projectedDraftGroups(bind: (target: TargetView) => ProjectedDraftMethod): GroupedDraftProvider {
-  const open = ({ target, checkpoints }: {
-    target: TargetView; checkpoints: readonly (DraftRowCheckpoint | null)[];
-  }) => new ProjectedDraftRows(bind(target), checkpoints);
-  return { open, openPrefill: open };
-}