@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/generation/speculative/sources/ngram-source.ts

@@ -0,0 +1,165 @@
+// NgramProvider — MODEL-FREE prompt-lookup speculative drafting behind the
+// grouped draft rows (packages/inference/src/generation/speculative/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 speculative group 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: each row keeps the committed stream it has seen. Prefill
+// appends the target's prompt chunks (to len-1 under the tail-split shape, the
+// last prompt token being the first round's pending token); draft appends the
+// row's pending token, then looks up; commit appends the accepted drafts.
+// Invariant at matching time: history == prompt + every emitted token, ending
+// at the pending token. A round with no match drafts nothing, which the group
+// verifies as one plain target step, so ngram speculation is structurally never
+// worse than plain decode by more than the host lookup and verification overhead.
+
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { CheckpointAttachment } from "../../../contracts/mlx/checkpoint";
+import type { DraftProvider, DraftRowGroup, DraftPrefillGroup, GroupedDraftProvider, DraftRowConstraints } from "../source";
+import type { DraftRowCheckpoint } from "../../../contracts/mlx/draft-checkpoint";
+import { applyStateChanges } from "../../../runtime/resources";
+
+// Histories are append-only between prefill/restore replacements. Weak keys
+// release derived indexes with retired rows; checkpoints still contain tokens
+// only. Ordered positions preserve the original first-occurrence policy.
+const ngramPositions = new WeakMap<readonly number[], { length: number; positions: Map<number, number[]> }>();
+function proposeNgram(h: readonly number[], max: number, min: number, n: number): number[] {
+  let index = ngramPositions.get(h);
+  if (!index || h.length < index.length) {
+    index = { length: 0, positions: new Map() };
+    ngramPositions.set(h, index);
+  }
+  for (let i = index.length; i < h.length; i++) {
+    const token = h[i]!;
+    const positions = index.positions.get(token);
+    if (positions) positions.push(i); else index.positions.set(token, [i]);
+  }
+  index.length = h.length;
+  const L = h.length;
+  for (let k = Math.min(max, L - 1); k >= min; k--) {
+    const tailAt = L - k;
+    const positions = index.positions.get(h[tailAt]!) ?? [];
+    search: for (const i of positions) {
+      if (i + k >= L) break;
+      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)[],
+    readonly constraints?: DraftRowConstraints) {
+    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[][] | Promise<number[][]> {
+    if (this.constraints) {
+      for (const [row, history] of this.#histories.entries()) history.push(pending[row]!);
+      return Promise.all(this.#histories.map((_, row) => this.constraints!.propose(row, depth)))
+        .then(proposals => this.#drafts = proposals);
+    }
+    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 = {
+    checkpointNamespace: () => namespace(this.max, this.min),
+    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));
+  }
+
+  dispose(): void {}
+}
+
+/** Constraint candidates share committed-history ownership with prompt lookup.
+ * They never perform ngram lookup. */
+export function constraintDraftProvider(): { id: string; grouped: GroupedDraftProvider } {
+  return {
+    id: "grammar",
+    grouped: {
+      checkpointNamespace: () => namespace(3, 1),
+      supportsTargetAdapters: true,
+      openPrefill: options => new NgramRows(3, 1, options.checkpoints),
+      open: options => new NgramRows(3, 1, options.checkpoints,
+        options.constraints ?? { propose: async () => [] }),
+    },
+  };
+}

package/src/generation/speculative/sources/qwen-mtp-source.ts

@@ -0,0 +1,154 @@
+// Native Qwen3.8 MTP draft provider — the Qwen-TRAINED multi-step prediction
+// head, split from the raw release's last shard and published as
+// mlx-community/Qwen3.8-27B-MTP-* (model_type "qwen3_5_mtp", block_size 3).
+//
+// Reference implementation (locate + port + gate, don't design):
+// mlx_vlm/speculative/drafters/qwen3_5_mtp/qwen3_5_mtp.py. Mechanism, per
+// step: h = fc(concat(rms_emb(embed(token)), rms_hidden(hidden))) → ONE
+// full-attention decoder layer (own KVCache, positions continue the target
+// sequence) → norm → the TARGET's lm_head → next draft token; recursive for
+// the block ("trained with multiple steps"). The drafter is NOT standalone:
+// it binds the target's embed_tokens + lm_head (mtp_use_dedicated_embeddings
+// false) and consumes the target's PRE-final-norm last-layer hidden (mlx-vlm
+// captures it with skip_final_norm=True) — carried here by the provider's tap
+// layer, NOT the target's post-norm output.
+//
+// Row convention (predict-2-ahead): drafter KV row at position p is built
+// from (embed(token at position p+1), hidden at position p). Invariant
+// across rounds: after commit, the drafter offset equals the next pending
+// token's position minus zero — i.e. a draft always opens by building the
+// pending token's row from the TARGET's TRUE hidden at the emitted position
+// (held over from the verify tap). Computationally identical to mlx-vlm's
+// commit-time "seed" append, just performed at the top of the next round.
+//
+// Companion projections may be dense or quantized according to their metadata.
+// Published drafter norms are already sanitized to runtime layout; no +1.0 shift.
+
+import { loadModelConfig, type ModelConfig } from "../../../artifacts/config";
+import { Weights } from "../../../artifacts/weights";
+import type { RecurrentMtpModule } from "../../../contracts/mlx/drafter";
+import { DRAFT_CHECKPOINT_SCHEMA } from "../../../contracts/mlx/draft-checkpoint";
+import { loadQwenMtpModule } from "../../../models/drafter-loaders";
+import type { DraftProvider, DraftRowGroup, DraftPrefillGroup, DraftRowSampling, GroupedDraftProvider, RecurrentMtpTarget, TargetView } from "../source";
+import { targetLacks } from "../source";
+import type { DraftRowCheckpoint } from "../../../contracts/mlx/draft-checkpoint";
+import { QwenMtpRows } from "../bindings/qwen-mtp-rows";
+import { configuredDraftVocabulary } from "../draft-vocab";
+import { type MtpRowState } from "../../../state/speculative/qwen-mtp-state";
+import { captureQwenMtpState, restoreQwenMtpState } from "../../../state/speculative/qwen-mtp-state";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../../../runtime/resources";
+import { artifactIdentity } from "../../../artifacts/identity";
+import { configFingerprint } from "../../../artifacts/fingerprint";
+import type { PreparedStateChange } from "../../../contracts/portable/resources";
+
+/** MTP drafts from the target's last hidden layer, before its final norm. */
+const mtpTapLayers = (target: RecurrentMtpTarget): readonly number[] => Object.freeze([target.layerCount - 1]);
+
+export class QwenMtpProvider implements DraftProvider {
+  readonly grouped: GroupedDraftProvider = {
+    supportsExternalTokens: true,
+    targetTapLayers: view => mtpTapLayers(this.#target(view)),
+    // A list shipped beside the companion (or named by MLX_BUN_SPEC_DRAFT_VOCAB),
+    // for a target whose quantized head the rows can project a row subset of.
+    draftVocabulary: view => {
+      const head = this.#target(view).vocabularyHead;
+      return head ? configuredDraftVocabulary(head.vocabSize, this.#dir) : null;
+    },
+    checkpointNamespace: () => this.#checkpointNamespace,
+    open: options => this.#openRows(options.target, options.sampling, options.checkpoints),
+    openPrefill: options => this.#openRows(options.target, null, options.checkpoints),
+  };
+  readonly id: string;
+  readonly weightsBytes: number;
+  readonly #module: RecurrentMtpModule;
+  readonly #config: ModelConfig;
+  readonly #resources: DisposableStack;
+  readonly #checkpointNamespace: string;
+  /** The companion's directory, where its optional draft vocabulary ships. */
+  readonly #dir: string;
+
+  private constructor(id: string, dir: string, config: ModelConfig, weightsBytes: number, module: RecurrentMtpModule, resources: DisposableStack, checkpointNamespace: string) {
+    this.id = id;
+    this.#dir = dir;
+    this.#config = config;
+    this.#module = module;
+    this.weightsBytes = weightsBytes;
+    this.#resources = resources;
+    this.#checkpointNamespace = checkpointNamespace;
+  }
+
+  static async load(dir: string): Promise<QwenMtpProvider> {
+    const config = await loadModelConfig(dir);
+    const weights = await Weights.open(dir);
+    using resources = new DisposableStack();
+    resources.defer(() => weights.dispose());
+    const weightsBytes = [...weights.shards.files.values()]
+      .reduce((a, f) => a + f.mmap.size, 0);
+    const module = await loadQwenMtpModule(weights, config, resources);
+    // Stable across restarts and distinct for differently folded/quantized
+    // companions. Hash once at provider load, outside inference execution.
+    const identity = await artifactIdentity(configFingerprint(config),
+      [...weights.shards.files].map(([name, shard]) => ({ name, path: shard.path })));
+    return new QwenMtpProvider(
+      dir.split("/").filter(Boolean).at(-1) ?? "qwen-mtp", dir,
+      config, weightsBytes, module, resources.move(), `${DRAFT_CHECKPOINT_SCHEMA.qwenMtp}:${identity}`,
+    );
+  }
+
+  #target(view: TargetView): RecurrentMtpTarget {
+    if (this.#resources.disposed)
+      throw new Error("qwen MTP provider is disposed");
+    const target = view.recurrentMtp;
+    if (!target) throw targetLacks("recurrentMtp");
+    if (target.hiddenSize !== this.#config.text.hiddenSize) {
+      throw new Error(
+        `qwen MTP drafter hidden ${this.#config.text.hiddenSize} != target ` +
+        `${target.hiddenSize} — split from a different checkpoint?`,
+      );
+    }
+    return target;
+  }
+
+  #openRows(view: TargetView, sampling: DraftRowSampling | null,
+    checkpoints: readonly (DraftRowCheckpoint | null)[]): DraftRowGroup & DraftPrefillGroup {
+    const target = this.#target(view);
+    const rows = new QwenMtpRows(target, this.#module, sampling, []);
+    const prepareAppend = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => {
+      const states: Array<MtpRowState | null> = [];
+      let change: PreparedStateChange | undefined;
+      try {
+        for (const checkpoint of checkpoints) states.push(checkpoint ? restoreQwenMtpState(checkpoint) : null);
+        change = rows.prepareAppend(states);
+        disposeResources(states.splice(0).flatMap(state => state ? [state.cache, state.hidden] : []));
+        return change;
+      } catch (error) {
+        return cleanupFailure(error, () => disposeResources([
+          ...states.flatMap(state => state ? [state.cache, state.hidden] : []), ...(change ? [change] : []),
+        ]));
+      }
+    };
+    const append = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => applyStateChanges([() => prepareAppend(checkpoints)]);
+    try { append(checkpoints); }
+    catch (error) { return cleanupFailure(error, () => rows.dispose()); }
+    return {
+      namespace: this.#checkpointNamespace, prefillMode: "full", tapLayers: mtpTapLayers(target),
+      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), draftDevice: rows.draftDevice.bind(rows), preview: rows.preview.bind(rows), commit: rows.commit.bind(rows),
+      consume: rows.consume.bind(rows),
+      capture(row) {
+        const state = rows.extractRow(row);
+        try { return captureQwenMtpState(state); }
+        finally { disposeResources([state.cache, state.hidden]); }
+      },
+      dispose: rows.dispose.bind(rows),
+    };
+  }
+
+  dispose(): void {
+    // Release cached transpose views before their native weight maps. MLX
+    // retains buffers needed by outstanding GPU commands until completion.
+    this.#resources.dispose();
+  }
+}

package/src/generation/speculative/sources/two-model.ts

@@ -0,0 +1,75 @@
+// Two-model speculative drafting — mlx-lm parity (`mlx_lm.server
+// --draft-model`). A full second model drafts autoregressively in the grouped
+// draft rows (bindings/two-model-rows.ts); the shared speculative group
+// verifies with the target. L1 oracle:
+// mlx_lm.server with the same target/draft pair, greedy, token-for-token
+// (spec-vs-spec — both batch the verify lm-head, so neither is bit-exact to
+// stock decode at knife-edges).
+//
+// Faithfulness notes (read from the oracle venv's generate.py):
+//  - drafts are sampled with the REQUEST sampler (generate.py:593-601) —
+//    greedy drafting under a temperature>0 request is not parity;
+//  - draft-cache rewind is max(n - kAccept - 1, 0) (generate.py:589-591),
+//    with the all-accept re-feed of the last draft kept pending in the rows
+//    and joined with the next target token (generate.py:645-648);
+//  - logits processors do NOT run on draft steps (they're target-side).
+
+import { loadModelConfig } from "../../../artifacts/config";
+import { Weights } from "../../../artifacts/weights";
+import { createModel, type RuntimeModel } from "../../../models/factory";
+import type { DraftProvider, GroupedDraftProvider } from "../source";
+import { standaloneDraftGroups } from "../bindings/two-model-rows";
+import { bindMlxAutoregressiveGraph } from "../../bindings/autoregressive";
+import { artifactIdentity } from "../../../artifacts/identity";
+import { configFingerprint } from "../../../artifacts/fingerprint";
+import { cleanupFailure, disposeResources } from "../../../runtime/resources";
+
+export class TwoModelProvider implements DraftProvider {
+  readonly id: string;
+  readonly weightsBytes: number;
+  readonly grouped: GroupedDraftProvider;
+
+  private constructor(
+    readonly model: RuntimeModel,
+    private readonly weights: Weights,
+    id: string,
+    namespace: string,
+  ) {
+    this.id = id;
+    this.grouped = standaloneDraftGroups(bindMlxAutoregressiveGraph(model), namespace);
+    this.weightsBytes = [...weights.shards.files.values()].reduce(
+      (a, f) => a + f.mmap.size,
+      0,
+    );
+  }
+
+  /** Load the draft model from a snapshot dir. Vocab-size mismatch with the
+   *  target is a WARNING (mlx-lm parity, server.py:363-368); the tokenizer-
+   *  family hard check lives at the server (it owns both tokenizers). */
+  static async load(modelDir: string, targetVocabSize?: number): Promise<TwoModelProvider> {
+    const config = await loadModelConfig(modelDir);
+    const weights = await Weights.open(modelDir);
+    let model: RuntimeModel | null = null;
+    try {
+      model = createModel(weights, config);
+      if (
+        targetVocabSize !== undefined &&
+        config.text.vocabSize !== targetVocabSize
+      )
+        console.warn(
+          `draft model vocab size ${config.text.vocabSize} != target ${targetVocabSize} — ` +
+            `speculation will accept nothing useful (mlx-lm warns identically)`,
+        );
+      const identity = await artifactIdentity(configFingerprint(config),
+        [...weights.shards.files].map(([name, shard]) => ({ name, path: shard.path })));
+      return new TwoModelProvider(model, weights, modelDir.split("/").filter(Boolean).at(-1)!,
+        `standalone-draft-v1:${identity}`);
+    } catch (error) { return cleanupFailure(error, () => disposeResources([...(model ? [model] : []), weights])); }
+  }
+
+  /** The draft graph releases what it built (constants, Neural Engine programs),
+   *  then its weights go. */
+  dispose(): void {
+    disposeResources([this.model, this.weights]);
+  }
+}

package/src/generation/types.ts

@@ -0,0 +1,194 @@
+// Streaming generation — port of mlx-lm's generate_step:
+// - prefill in chunks; cache state evaluated per chunk (bounded transient
+//   memory), logits never computed for non-final prefill positions
+// - decode pipelining via mx.async_eval: step n+1's graph is built and
+//   dispatched before step n's token is read back, so the GPU never idles
+//   on the JS round-trip
+// - opt-in early token-zero yield reduces latency before that pipeline starts
+// - sampling stays on-device; only the chosen token id crosses to JS
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { type Cache } from "../contracts/mlx/cache";
+import type { TokenLogprobs } from "../contracts/portable/generation";
+import type { PromptResponseTrace } from "../runtime/trace";
+import type { GrammarController } from "../sampling/grammar";
+import {
+type LogitsProcessorOptions,
+type SamplerOptions
+} from "../sampling/index";
+import {
+type FillSession,
+type FillStats
+} from "./fill/session";
+
+export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions {
+  /** Application session affinity; independent of numerical cache identity. */
+  cacheSessionId?: string;
+  /** Resolved host policy. Direct compatibility calls resolve from their
+   * captured binding when this is absent. */
+  decodePolicy?: Readonly<Pick<import("../contracts/portable/execution").ResolvedExecution, "compiledDecode" | "grammarJump">>;
+  /** Cooperatively cancel at prefill/decode boundaries; pending native work
+   *  finishes before owned resources are released. */
+  signal?: AbortSignal;
+  maxTokens?: number;
+  eosTokenIds?: number[];
+  prefillChunkSize?: number;
+  /** Whether the request supplied its seed explicitly. Durable generation
+   *  checkpoints use this to distinguish a reproducible user policy from the
+   *  server-generated seed that must be recovered from checkpoint metadata. */
+  seedWasExplicit?: boolean;
+  /** Fired once during prefill, at `snapshotAt` tokens (default: the full
+   *  prompt) and BEFORE any further KV is written — the one moment the
+   *  caches hold exactly that prefix (post-wrap rings can't be rewound
+   *  later). The server's prompt-boundary cache snapshot hangs here. */
+  onPrefillDone?: () => void;
+  /** Token count at which onPrefillDone fires (the STABLE cache boundary —
+   *  chat prompts end in a generation primer, e.g. 12B's thought-channel
+   *  tokens, that the NEXT turn's re-render does not contain; snapshotting
+   *  there would make the entry untrimmable-and-divergent). Prefill is
+   *  split at this index when it falls mid-prompt. */
+  snapshotAt?: number;
+  /** Resume an interrupted autoregressive generation from a durable KV
+   *  checkpoint. `cache` must cover every token in `promptTokens` exactly;
+   *  `initialPendingToken` is the already-sampled next token, so no cache
+   *  rewind is required (important for recurrent/untrimmable layers). */
+  initialPendingToken?: number;
+  /** Number of completion tokens already emitted before this resumed run. */
+  initialGeneratedTokens?: number;
+  /** Original request prompt length used for protocol usage accounting when
+   *  `promptTokens` also contains replayed completion tokens. */
+  originalPromptTokens?: number;
+  /** Periodic, safe decode-boundary checkpoint hook. The cache covers
+   *  `cacheTokens`; `pendingToken` has been sampled but not emitted. */
+  checkpointEveryTokens?: number;
+  onDecodeCheckpoint?: (state: {
+    cacheTokens: number[];
+    caches: Cache[];
+    generatedTokens: number;
+    pendingToken: number;
+  }) => void | Promise<void>;
+  /** Pre-warmed KV caches (e.g. from the prompt cache). cache[0].offset
+   *  prompt tokens are treated as already prefilled; only the suffix is
+   *  forwarded. Caller keeps ownership — generate() will not dispose. */
+  cache?: Cache[];
+  /** Vision path: pre-merged (unscaled) input embeddings [1, L, hidden]
+   *  covering the whole prompt; prefilled in one shot (no chunking).
+   *  Caller keeps ownership. */
+  promptEmbeddings?: MlxArray;
+  /** bool [L] marking image tokens (bidirectional attention among them).
+   *  MUST be unset when the prompt contains any audio — audio prompts run
+   *  fully causal (`02d723a:docs/design/generic-model-support.md` §6.6). */
+  imageMask?: MlxArray;
+  /** bool [L] marking ALL multimodal soft tokens (image | audio) for
+   *  per-layer-input id zeroing (e2b/e4b), decoupled from imageMask so
+   *  audio-only prompts (no bidirectional mask) still zero their positions.
+   *  When unset, zeroing falls back to imageMask (the legacy vision shape).
+   *  Caller keeps ownership. */
+  multimodalMask?: MlxArray;
+  /** Mounted LoRA adapter ids to apply (resolved/validated by
+   *  AdapterManager.resolveSpec). Residuals sum in order. Set on the
+   *  model's LoraState for exactly the duration of this generation —
+   *  a plain field, safe because the generation queue is serialized. */
+  adapters?: string[];
+  /** DiffusionGemma image-text-to-text: channel-first pixel_values [1,3,H,W].
+   *  When set, `promptTokens` are the spliced (<|image|>-expanded) ids and the
+   *  denoising engine prefills the merged vision features. Caller keeps ownership. */
+  visionPixels?: MlxArray;
+  /** Capture each emitted token's log-probability (mlx_lm.server `logprobs`).
+   *  The distribution matches mlx-lm generate_step exactly: full-vocab
+   *  log-softmax of the logits AFTER logits processors, BEFORE the sampler's
+   *  temperature/top-p/top-k/min-p/XTC (generate.py L409-422). Off by default —
+   *  the hot path pays nothing when unset. */
+  logprobs?: boolean;
+  /** Capture the top-k (token id, logprob) pairs per emitted token from the
+   *  same distribution (mlx_lm.server `top_logprobs`). 0/unset = off. */
+  topLogprobs?: number;
+  /** Grammar-constrained decoding (packages/inference/src/sampling/grammar.ts): a compiled GrammarController
+   *  that masks invalid tokens to -inf each step. L2-class (oMLX oracle). When
+   *  set, the decode loop takes a slightly different shape: it eager-reads the
+   *  token id (acceptToken needs a JS number, which the pipelined loop defers),
+   *  advances the matcher, and awaits the async mask precompute (which overlaps
+   *  the GPU forward). Non-grammar requests keep the fast pipelined loop. */
+  grammar?: GrammarController;
+  /** Token fast-forwarding (packages/inference/src/generation/fill/,
+   *  `02d723a:docs/design/speculative-decoding.md` §7 "Token fast-forwarding"):
+   *  a per-request table of DETERMINED token spans
+   *  (compiled from the request's `tools` + the chat template). When the
+   *  stream enters one, the engine appends the whole span itself with ONE
+   *  multi-token forward and resumes sampling after it.
+   *
+   *  This is context extension, NOT speculation: nothing is drafted, verified,
+   *  or rolled back, and no comparison is made against what the model would
+   *  have produced. Injection bypasses the sampler, which is a documented
+   *  behavior-policy deviation at temperature > 0 — hence opt-in
+   *  (MLX_BUN_FILL=strict) and the request refusals in shouldUseFill.
+   *  Serial lane only: the batch scheduler and the spec loop never read it. */
+  fill?: FillSession;
+}
+
+export interface GenerateStats {
+  promptTokens: number;
+  /** An observed terminal cause takes precedence over the token-count limit. */
+  finishReason?: "stop" | "length";
+  /** Prompt tokens skipped via a pre-warmed cache. */
+  cachedTokens: number;
+  generatedTokens: number;
+  prefillTps: number;
+  decodeTps: number;
+  prefillMs: number;
+  decodeMs: number;
+  /** Exact token sequence whose KV now lives in the cache (prompt + every
+   *  decoded token that was forwarded, including a trailing EOS the
+   *  pipeline forwarded before reading it). For PromptCache.put(). */
+  cacheTokens: number[];
+  /** Speculative-decoding telemetry (serve --draft-model path only).
+   *  draftedByPos/acceptedByPos: per-draft-position round counts (index =
+   *  position within a round's block) — the per-position
+   *  acceptance signal. */
+  spec?: {
+    drafted: number; accepted: number; targetCalls: number;
+    draftedByPos?: number[]; acceptedByPos?: number[];
+    rejected?: number; rounds?: number; acceptanceLengths?: number[];
+    tokensPerForward?: number; forwardsSaved?: number;
+    /** Round phase wall time, summed over the request's rounds. Present only
+     *  under `MLX_BUN_SPEC_PHASE_TIMING=1`, which forces an evaluation at each
+     *  phase boundary (diagnostic: it removes the overlap production decode
+     *  relies on, so the sum is not production round time). */
+    phaseMs?: SpecPhaseMs;
+  };
+  /** Token fast-forwarding telemetry (MLX_BUN_FILL=strict). Present only when
+   *  a fill table was actually armed for this generation. */
+  fill?: FillStats;
+}
+
+/** Wall milliseconds per speculative round phase (see GenerateStats.spec). */
+export interface SpecPhaseMs {
+  /** draft(): the drafter's chain including its proposal readback. */
+  draft: number;
+  /** Target verify forward, forced to completion before sampling. */
+  verify: number;
+  /** Window sampling and its readback, plus the host accept walk. */
+  sample: number;
+  /** Target transaction resolve plus draft commit. Includes retained-state
+   * evaluation when the target/drafter supply diagnostic materialize hooks. */
+  commit: number;
+  rounds: number;
+  /** Target-forward component wall ms (`MLX_BUN_SPEC_LAYER_PROFILE=1`): per-op
+   *  evaluation barriers inside the verify forward, so `verify` grows while set. */
+  layers?: Record<string, number>;
+  /** Graph nodes built per op name (`MLX_BUN_SPEC_OP_INVENTORY=1`), summed over rounds. */
+  verifyOps?: Record<string, number>;
+  draftOps?: Record<string, number>;
+}
+
+export interface GenerateDiagnostics {
+  trace?: PromptResponseTrace;
+  mechanism?: "serial" | "continuous";
+}
+
+export interface GeneratedToken {
+  token: number;
+  index: number;
+  /** Present only when GenerateOptions.logprobs / topLogprobs requested it. */
+  logprobs?: TokenLogprobs;
+}

package/src/index.ts

@@ -0,0 +1,20 @@
+// Common library entry points. Specialized components have explicit subpath imports.
+export { loadModelConfig } from "./artifacts/config";
+export { Weights } from "./artifacts/weights";
+export { createModel, openModel, openGlm52RuntimeModel } from "./models/factory";
+export type { RuntimeModel, ModelOpenOptions, Glm52RuntimeOpenOptions } from "./models/factory";
+export { resolveComposition } from "./models/composition";
+export type { CompositionRequest } from "./models/composition";
+export type { Composition } from "./contracts/portable/composition";
+export { ModelImplementationRegistry } from "./models/implementation";
+export { loadTokenizer } from "./input/tokenizer";
+export type { LoadedTokenizer } from "./input/tokenizer";
+export { ChatTemplate } from "./input/chat-template";
+export { generate, bindGeneration } from "./generation/generate";
+export { generateAutoregressive } from "./generation/autoregressive";
+export { generateDenoising } from "./generation/denoising";
+export { Generation } from "./generation/result";
+export type { GenerateOptions, GenerateStats, GeneratedToken } from "./generation/types";
+export { embedOne, embedMany, embeddingTerminatorId, isEmbeddingModel, withInstruction } from "./embeddings/text";
+export type { EmbedResult, EmbeddingGraph } from "./embeddings/text";
+export { AdapterManager } from "./adapters/manager";