@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/generation/fill/session.ts

@@ -0,0 +1,464 @@
+// Token fast-forwarding (K3) — the per-request proposal table.
+//
+// DOCTRINE. This is PURE CONTEXT EXTENSION, not speculation. A transformer is
+// a next-token function over a token sequence; it has no memory of who wrote
+// which token, so an APPENDED token is indistinguishable from a sampled one.
+// When the engine already knows the next m tokens (a chat template's tool-call
+// scaffold is a deterministic function of the request's `tools`), it appends
+// them itself in one append and resumes sampling after them. The append can
+// use several execution chunks while retaining the complete span. No
+// draft, no verify, no rollback, and no comparison against what the model
+// "would have" produced — nothing here routes through
+// packages/inference/src/generation/speculative/.
+//
+// The echo tier (K3c) adds a weaker claim over the SAME mechanism: a span
+// copied from earlier in the session, held under policy "verify" — the engine
+// keeps the prefix the model agrees with (read from the free logits of the
+// very same forward) and rewinds the rest. See ./proposal.ts for the two
+// policies and ./echo-index.ts for the branch-point rule.
+//
+// This module is bookkeeping only: it owns the running history, runs the
+// sources in priority order, clamps what they return, and counts. It never
+// touches MLX, the tokenizer, or the template. Row COMPILATION is
+// ./schema-rows.ts.
+import { runtimeFlag, runtimeNumber, runtimeValue } from "../../runtime/config";
+import { EchoSource, type EchoConfig } from "./echo-index";
+import type {
+  Proposal, ProposalOrigin, ProposalPolicy, ProposalSource, TokenView,
+} from "./proposal";
+
+export type { Proposal, ProposalPolicy, ProposalSource, TokenView };
+
+/** Which producer a strict row came from (telemetry + the mismatch policy). */
+export type FillKind = "scaffold" | "name" | "key" | "close" | "turn-end";
+
+export interface FillRow {
+  /** Token ids that must be the TAIL of the generated history for this row to
+   *  fire (matched as an exact suffix ending at the token just pushed). */
+  trigger: number[];
+  /** Ids appended immediately after the trigger. Sliced out of a full
+   *  template rendering — never `encode(fragment)` in isolation. */
+  emit: number[];
+  kind: FillKind;
+  /** Template-derived rows need a parser-owned context before asserting. */
+  requiresContext?: true;
+}
+
+export interface StrictFillContext {
+  observe(ids: readonly number[]): void;
+  allows(row: FillRow): boolean;
+}
+
+/** Echo-tier (K3c) configuration; null disables the source entirely. */
+export interface FillEchoConfig {
+  k: number;
+  maxSpan: number;
+  maxCandidates: number;
+  indexMax: number;
+}
+
+export interface FillPlan {
+  rows: FillRow[];
+  /** Echo tier: null = strict rows only (MLX_BUN_FILL=strict). */
+  echo: FillEchoConfig | null;
+  /** Token ids that legitimately END an injected span — the closing quote /
+   *  markup that follows an argument value, derived from the template by the
+   *  strict-row compiler. A span is cut so it ends AT the first delimiter it
+   *  contains, never past one. */
+  delimiters?: ReadonlySet<number>;
+  /** EOS ids. An injected span is cut BEFORE the first one: end-of-generation
+   *  is the model's decision, never ours. (The grammar jump burst yields its
+   *  ids with NO eos check — that bug is deliberately not inherited.) */
+  eos: number[];
+}
+
+export interface FillStats {
+  /** Fills that put at least one token into the stream. */
+  events: number;
+  /** Total injected tokens (billed as generated). */
+  injected: number;
+  /** …from strict schema/template rows (always policy "assert"). */
+  strict: number;
+  /** …from the echo index. */
+  echo: number;
+  /** Length of each injected span, in event order. */
+  spanLens: number[];
+  /** In-flight sampled tokens dropped unexamined at an ASSERT fill — a
+   *  discarded pipeline dispatch, NOT a rejected draft. */
+  wastedSamples: number;
+  /** Tool-call parses that failed after strict rows had fired. */
+  parseFallback: number;
+  /** Spans shortened by a clamp (EOS / delimiter / max span / budget). */
+  indexTruncated: number;
+  /** Tokens the model actually sampled (the real decode steps). */
+  decodeSteps: number;
+  /** Verify-policy proposals that reached the forward. */
+  verifyEvents: number;
+  /** Tokens accepted under verify (the model's own argmax agreed). */
+  verifyAccepted: number;
+  /** Tokens proposed under verify and dropped (incl. rejected-at-position-0,
+   *  which costs nothing — that check happens before the forward). */
+  verifyRejected: number;
+  /** Verify proposals dropped because this model's caches cannot rewind. */
+  verifyUnsupported: number;
+  /** Wall time in cache checkpoint / rollback / commit (verify only). */
+  checkpointMs: number;
+  /** Echo spans that stopped because the session history forked. */
+  branchStops: number;
+}
+
+const emptyStats = (): FillStats => ({
+  events: 0, injected: 0, strict: 0, echo: 0, spanLens: [],
+  wastedSamples: 0, parseFallback: 0, indexTruncated: 0, decodeSteps: 0,
+  verifyEvents: 0, verifyAccepted: 0, verifyRejected: 0, verifyUnsupported: 0,
+  checkpointMs: 0, branchStops: 0,
+});
+
+export type FillMode = "off" | "strict" | "echo";
+
+/** MLX_BUN_FILL=off|strict|echo (default off). `echo` is additive: the strict
+ *  rows still apply (and stay policy "assert"); the echo index joins as a
+ *  second, weaker source. Lab tier — default off, A/B before any default. */
+export function resolveFillMode(raw = runtimeValue("MLX_BUN_FILL")): FillMode {
+  if (raw === undefined || raw === "" || raw === "off" || raw === "0") return "off";
+  if (raw === "strict" || raw === "1") return "strict";
+  if (raw === "echo") return "echo";
+  throw new Error(`MLX_BUN_FILL=${raw}: expected off|strict|echo`);
+}
+
+/** Hard cap on one injected span (MLX_BUN_FILL_MAX_SPAN, default 32). */
+export function fillMaxSpan(): number {
+  const n = Math.floor(runtimeNumber("MLX_BUN_FILL_MAX_SPAN", 32));
+  return n >= 2 ? n : 2;
+}
+
+/** Assert-append chunk cap; zero selects the model's qualified limit. */
+export function fillAppendChunkSize(): number {
+  return Math.max(0, Math.floor(runtimeNumber("MLX_BUN_FILL_APPEND_CHUNK_SIZE", 0)));
+}
+
+/** Echo-tier knobs. k: anchor length (MLX_BUN_FILL_K, default 8 — the corpus
+ *  study's token-level threshold). candidates: nearest-occurrence bucket cap
+ *  (MLX_BUN_FILL_CANDIDATES, default 24). indexMax: token cap on the growing
+ *  index (MLX_BUN_FILL_INDEX_MAX, default 131072). */
+export function fillEchoConfig(): FillEchoConfig {
+  return {
+    k: Math.max(2, Math.floor(runtimeNumber("MLX_BUN_FILL_K", 8))),
+    maxSpan: fillMaxSpan(),
+    maxCandidates: Math.max(1, Math.floor(runtimeNumber("MLX_BUN_FILL_CANDIDATES", 24))),
+    indexMax: Math.max(1024, Math.floor(runtimeNumber("MLX_BUN_FILL_INDEX_MAX", 131072))),
+  };
+}
+
+/** MLX_BUN_FILL_TRACE=1 — cache-alignment assertions + per-event logging on
+ *  the engine's fill path. Off by default (the invariant it checks is the
+ *  easiest bug in the feature, and the check costs a JS compare). */
+export function fillTraceEnabled(): boolean {
+  return runtimeFlag("MLX_BUN_FILL_TRACE", false) || fillTracePath() !== null;
+}
+
+/** MLX_BUN_FILL_TRACE=<file.jsonl> — additionally append one record per
+ *  proposal: the proposed ids/text next to what the model's OWN logits said
+ *  at every span position (position 0 = the in-flight sample, j>0 = argmax
+ *  after ids[j-1]), for BOTH policies. Under `assert` the readback is
+ *  trace-only (one lm_head over the span) — the served path never pays it.
+ *  This is the list that answers "would the model have produced this span?"
+ *  without a verify pass in production. */
+export function fillTracePath(): string | null {
+  const raw = runtimeValue("MLX_BUN_FILL_TRACE");
+  if (raw === undefined || raw === "" || raw === "0" || raw === "1" || raw === "true") return null;
+  return raw;
+}
+
+export interface FillTraceRecord {
+  ts: string;
+  origin: string;
+  policy: string;
+  generated: number;
+  proposedLen: number;
+  accepted: number;
+  /** Index of the first position where the model disagreed, or -1. */
+  firstMismatch: number;
+  proposed: number[];
+  actual: number[];
+  proposedText?: string;
+  actualText?: string;
+}
+
+/** Strict schema/template rows as a ProposalSource. Always policy "assert":
+ *  the spans are sliced out of the model's own chat template, so they are
+ *  determined by construction, not guessed. */
+export class StrictRowSource implements ProposalSource {
+  readonly name = "strict";
+  readonly #byLastTrigger = new Map<number, FillRow[]>();
+  /** Longest trigger — the tail this source needs to match against. */
+  readonly maxTrigger: number = 1;
+  get windowNeeded(): number { return this.maxTrigger; }
+  #enabled = true;
+
+  constructor(rows: readonly FillRow[], readonly context?: StrictFillContext) {
+    for (const row of rows) {
+      if (row.trigger.length === 0 || row.emit.length === 0) continue;
+      this.maxTrigger = Math.max(this.maxTrigger, row.trigger.length);
+      const last = row.trigger[row.trigger.length - 1]!;
+      const bucket = this.#byLastTrigger.get(last);
+      if (bucket) bucket.push(row);
+      else this.#byLastTrigger.set(last, [row]);
+    }
+    // Longest trigger wins (most specific row); plan order breaks ties.
+    for (const bucket of this.#byLastTrigger.values())
+      bucket.sort((a, b) => b.trigger.length - a.trigger.length);
+  }
+
+  get enabled(): boolean { return this.#enabled; }
+  disable(): void { this.#enabled = false; }
+
+  observe(ids: readonly number[]): void { this.context?.observe(ids); }
+
+  propose(view: TokenView): Proposal | null {
+    if (!this.#enabled) return null;
+    const tail = view.tail(this.maxTrigger);
+    const last = tail[tail.length - 1];
+    if (last === undefined) return null;
+    const bucket = this.#byLastTrigger.get(last);
+    if (!bucket) return null;
+    for (const row of bucket) {
+      const t = row.trigger;
+      if (t.length > tail.length) continue;
+      const base = tail.length - t.length;
+      let ok = true;
+      for (let i = 0; i < t.length; i++) {
+        if (tail[base + i] !== t[i]) { ok = false; break; }
+      }
+      if (ok && (!row.requiresContext || this.context?.allows(row))) {
+        return {
+          ids: [...row.emit],
+          policy: "assert",
+          // The scaffold and the close come from the TEMPLATE; the name and
+          // key spans are what the request's SCHEMA determines.
+          origin: row.kind === "name" || row.kind === "key" ? "schema" : "template",
+        };
+      }
+    }
+    return null;
+  }
+}
+
+/** One request's proposal sources plus the running match state.
+ *
+ *  CONTRACT. push() returns a proposal; the caller MUST then call commit()
+ *  exactly once with the number of ids it actually emitted (all of them for an
+ *  assert, the accepted prefix for a verify, 0 when it declined). Nothing
+ *  enters the session's history until commit — so a partial accept needs no
+ *  rollback here. */
+export class FillSession {
+  readonly stats: FillStats = emptyStats();
+  /** Execution chunk size, independent of the number of tokens injected. */
+  readonly appendChunkSize: number;
+  readonly #sources: ProposalSource[];
+  readonly #strict: StrictRowSource;
+  readonly #eos: Set<number>;
+  readonly #delimiters: ReadonlySet<number>;
+  readonly #maxSpan: number;
+  readonly #window: number;
+  /** Bounded tail of the token history — all a source can see directly (the
+   *  echo index keeps its own full copy via observe()). */
+  #tail: number[] = [];
+  #length = 0;
+  #budget = Number.POSITIVE_INFINITY;
+
+  /** Optional token→text decoder for the proposal trace (set by the serve
+   *  layer, which owns the tokenizer). */
+  readonly decode: ((ids: readonly number[]) => string) | null;
+
+  constructor(
+    readonly plan: FillPlan,
+    promptIds: readonly number[],
+    options: {
+      maxSpan?: number;
+      appendChunkSize?: number;
+      sources?: ProposalSource[];
+      decode?: (ids: readonly number[]) => string;
+      strictContext?: StrictFillContext;
+    } = {},
+  ) {
+    this.decode = options.decode ?? null;
+    this.appendChunkSize = options.appendChunkSize ?? fillAppendChunkSize();
+    if (!Number.isSafeInteger(this.appendChunkSize) || this.appendChunkSize < 0)
+      throw new Error("fill appendChunkSize must be a nonnegative safe integer");
+    this.#eos = new Set(plan.eos);
+    this.#delimiters = plan.delimiters ?? new Set<number>();
+    this.#maxSpan = Math.max(2, Math.floor(options.maxSpan ?? fillMaxSpan()));
+    this.#strict = new StrictRowSource(plan.rows, options.strictContext);
+    const echo = plan.echo
+      ? new EchoSource(
+        { ...plan.echo, maxSpan: this.#maxSpan, delimiters: this.#delimiters } as EchoConfig,
+        promptIds,
+      )
+      : null;
+    // Priority: determined before guessed, then anything the caller injected
+    // (the seam the spec lane's draft rows would adapt onto — see
+    // proposal.ts MIGRATION NOTE).
+    this.#sources = [
+      this.#strict,
+      ...(echo ? [echo] : []),
+      ...(options.sources ?? []),
+    ];
+    // Every source declares the tail it needs; the session retains the max.
+    this.#window = Math.max(
+      1, ...this.#sources.map((s) => s.windowNeeded ?? 1),
+    );
+    // A row may be triggered by the prompt's own tail (a re-rendered turn that
+    // ends mid-scaffold), so the history starts seeded.
+    this.#tail = promptIds.slice(Math.max(0, promptIds.length - this.#window));
+    this.#length = promptIds.length;
+  }
+
+  /** Are strict rows still armed? (Cleared by noteParseFailure.) */
+  get strictEnabled(): boolean { return this.#strict.enabled; }
+
+  /** Feed one SAMPLED token, then ask the sources what follows. `budget` is
+   *  the caller's remaining allowance (max_tokens − generated). */
+  push(tokenId: number, budget: number = Number.POSITIVE_INFINITY): Proposal | null {
+    this.observe(tokenId);
+    return this.propose(budget);
+  }
+
+  /** Record sampled output independently of when the method requests proposals. */
+  observe(tokenId: number): void {
+    this.stats.decodeSteps++;
+    this.#append([tokenId]);
+  }
+
+  /** Select a continuation after a committed output boundary. Methods may
+   * request only verified proposals while retaining ordinary sampled policy. */
+  propose(budget: number = Number.POSITIVE_INFINITY, policy?: ProposalPolicy): Proposal | null {
+    this.#budget = budget;
+    for (const source of this.#sources) {
+      const proposal = source.propose(this);
+      if (!proposal || (policy && proposal.policy !== policy)) continue;
+      if (proposal.branchStop) this.stats.branchStops++;
+      const clamped = this.#clamp(proposal.ids, budget, proposal.origin === "echo");
+      if (!clamped) continue;
+      // An assert claim that a clamp cut short for any reason OTHER than
+      // ending at a delimiter is no longer the claim the source made. Strict
+      // rows are exempt: they are determined by construction, and a budget
+      // clamp shortens the span without weakening it.
+      const weakened = proposal.origin === "echo" &&
+        clamped.ids.length !== proposal.ids.length &&
+        !this.#delimiters.has(clamped.ids[clamped.ids.length - 1]!);
+      return {
+        ...proposal,
+        ids: clamped.ids,
+        policy: weakened ? "verify" : proposal.policy,
+      };
+    }
+    return null;
+  }
+
+  /** Record what the engine actually emitted from `proposal` (all of it for an
+   *  assert; the accepted prefix for a verify; 0 when the engine declined).
+   *  MUST be called exactly once per proposal returned by push(). */
+  commit(proposal: Proposal, accepted: number): void {
+    const n = Math.max(0, Math.min(accepted, proposal.ids.length));
+    const rejected = proposal.ids.length - n;
+    if (proposal.policy === "verify") {
+      this.stats.verifyAccepted += n;
+      this.stats.verifyRejected += rejected;
+    }
+    if (n === 0) return;
+    const ids = proposal.ids.slice(0, n);
+    this.#append(ids);
+    this.stats.events++;
+    this.stats.injected += n;
+    if (proposal.origin === "echo") this.stats.echo += n;
+    else this.stats.strict += n;
+    this.stats.spanLens.push(n);
+  }
+
+  /** A verify proposal reached the forward (the checkpoint was taken). */
+  noteVerifyEvent(checkpointMs: number): void {
+    this.stats.verifyEvents++;
+    this.stats.checkpointMs += checkpointMs;
+  }
+
+  /** A verify proposal was dropped: this model's caches cannot rewind. */
+  noteVerifyUnsupported(): void {
+    this.stats.verifyUnsupported++;
+  }
+
+  /** The engine dropped an in-flight sampled token to take an ASSERT fill. */
+  noteWastedSample(): void {
+    this.stats.wastedSamples++;
+  }
+
+  /** The tool-call parser rejected this request's markup. Strict rows are
+   *  disarmed for the rest of the request — the template's own rendering
+   *  disagreed with what the model actually emits, so every further row is
+   *  suspect. (Echo rows are unaffected: they never claimed to be schema.)
+   *
+   *  REALITY NOTE: the served parse (ToolAwareStream.toolCalls →
+   *  parseGeneratedToolCalls) runs at sink FLUSH, i.e. after the generation
+   *  has ended, so within one request this is telemetry and the disarm has no
+   *  remaining tokens to act on. The seam is here so a future incremental
+   *  parse (or a plan cached across requests) disarms for real. */
+  noteParseFailure(): void {
+    if (this.stats.strict === 0) return; // no strict row armed this request
+    this.stats.parseFallback++;
+    this.#strict.disable();
+  }
+
+  // --- TokenView ---------------------------------------------------------
+  get length(): number { return this.#length; }
+  get budget(): number { return this.#budget; }
+  tail(n: number): readonly number[] {
+    return n >= this.#tail.length ? this.#tail : this.#tail.slice(this.#tail.length - n);
+  }
+
+  #append(ids: readonly number[]): void {
+    this.#length += ids.length;
+    for (const source of this.#sources) source.observe?.(ids);
+    this.#tail.push(...ids);
+    if (this.#tail.length > this.#window)
+      this.#tail.splice(0, this.#tail.length - this.#window);
+  }
+
+  /** Clamp order: EOS → delimiter → max span / caller budget → reject <2.
+   *  A one-token span saves no forward (the engine already had that token's
+   *  successor in flight), so it is never worth the discarded dispatch.
+   *
+   *  `applyDelimiters` is ECHO-ONLY. A delimiter marks where an argument VALUE
+   *  ends, which is exactly where a copied span must stop — but a template
+   *  scaffold legitimately contains the same tokens as structure (`{"name": "`
+   *  is three quotes deep), and it is determined by construction. Clamping
+   *  strict rows on delimiters would cut every scaffold at its first quote. */
+  #clamp(
+    emit: readonly number[], budget: number, applyDelimiters: boolean,
+  ): { ids: number[]; truncated: boolean } | null {
+    let ids = emit as number[];
+    let truncated = false;
+    for (let i = 0; i < ids.length; i++) {
+      if (this.#eos.has(ids[i]!)) { ids = ids.slice(0, i); truncated = true; break; }
+    }
+    if (applyDelimiters && this.#delimiters.size) {
+      for (let i = 0; i < ids.length; i++) {
+        // A span may END at a delimiter — injecting the closing quote is the
+        // point — but never continue past one.
+        if (this.#delimiters.has(ids[i]!)) {
+          if (i + 1 < ids.length) { ids = ids.slice(0, i + 1); truncated = true; }
+          break;
+        }
+      }
+    }
+    const cap = Math.min(
+      this.#maxSpan,
+      Number.isFinite(budget) ? Math.max(0, Math.floor(budget)) : this.#maxSpan,
+    );
+    if (ids.length > cap) { ids = ids.slice(0, cap); truncated = true; }
+    if (ids.length < 2) return null;
+    if (truncated) this.stats.indexTruncated++;
+    return { ids: ids === emit ? [...ids] : ids, truncated };
+  }
+}
+
+export type { ProposalOrigin };

package/src/generation/fill/tool-boundary.ts

@@ -0,0 +1,103 @@
+import type { ToolDefinition } from "../../input/chat-template";
+import { parseStrictToolCall } from "../../input/tool-call";
+import type { FillRow, StrictFillContext } from "./session";
+
+const OPEN = "<tool_call>", CLOSE = "</tool_call>";
+const THINK_OPEN = "<think>", THINK_CLOSE = "</think>";
+const MAX_CONTEXT_TOKENS = 65_536;
+
+/** Conservative for raw XML argument text. An unfinished literal can still
+ * contain a closing-tag example; leaving it to generation loses only a fill. */
+function literalsClosed(text: string): boolean {
+  let quote = "", escaped = false;
+  for (const c of text) {
+    if (escaped) { escaped = false; continue; }
+    if (c === "\\" && quote !== "'") { escaped = true; continue; }
+    if (quote) { if (c === quote) quote = ""; }
+    else if (c === '"' || c === "'" || c === "`") quote = c;
+  }
+  return quote === "" && !escaped;
+}
+
+function reasoningClosed(text: string): boolean {
+  if (text.includes(THINK_OPEN)) return false;
+  let fence = "", prose = "";
+  for (const line of text.split("\n")) {
+    const marker = /^\s*(`{3,}|~{3,})(.*)$/.exec(line);
+    if (fence) {
+      if (marker && marker[1]![0] === fence[0] &&
+          marker[1]!.length >= fence.length && !marker[2]!.trim()) fence = "";
+    } else if (marker) fence = marker[1]!;
+    else prose += line + "\n";
+  }
+  return !fence && literalsClosed(prose);
+}
+
+/** Request-local proof for template-derived rows. Only a tool-only assistant
+ * turn, optionally following its reasoning block, qualifies. Prose, quoted
+ * examples and unsupported formats retain ordinary generation. Decode runs
+ * only when a full token trigger matches, not for every generated token. */
+export class ToolCallFillContext implements StrictFillContext {
+  #ids: number[] = [];
+  #overflow = false;
+  readonly #startsInReasoning: boolean;
+
+  constructor(
+    readonly decode: (ids: readonly number[]) => string,
+    readonly tools: ToolDefinition[],
+    promptText: string,
+  ) {
+    this.#startsInReasoning = /<think>\s*$/.test(promptText);
+  }
+
+  observe(ids: readonly number[]): void {
+    if (this.#overflow) return;
+    if (this.#ids.length + ids.length > MAX_CONTEXT_TOKENS) {
+      this.#overflow = true;
+      this.#ids = [];
+      return;
+    }
+    for (const id of ids) this.#ids.push(id);
+  }
+
+  allows(row: FillRow): boolean {
+    if (this.#overflow) return false;
+    let text = this.decode(this.#ids);
+    if (this.#startsInReasoning) {
+      const end = text.indexOf(THINK_CLOSE);
+      if (end < 0 || !reasoningClosed(text.slice(0, end))) return false;
+      text = text.slice(end + THINK_CLOSE.length);
+    }
+    text = text.trimStart();
+    if (text.startsWith(THINK_OPEN)) {
+      const end = text.indexOf(THINK_CLOSE, THINK_OPEN.length);
+      if (end < 0 || !reasoningClosed(text.slice(THINK_OPEN.length, end))) return false;
+      text = text.slice(end + THINK_CLOSE.length).trimStart();
+    }
+    // Earlier calls must be complete and valid before another can open.
+    for (;;) {
+      if (!text.startsWith(OPEN)) return false;
+      const end = text.indexOf(CLOSE, OPEN.length);
+      if (end < 0) break;
+      const stop = end + CLOSE.length;
+      if (!parseStrictToolCall(text.slice(0, stop), this.tools)) return false;
+      text = text.slice(stop).trimStart();
+    }
+    if (text.indexOf(OPEN, OPEN.length) !== -1) return false;
+    if (row.kind === "scaffold" || row.kind === "name" || row.kind === "key") {
+      const trigger = this.decode(row.trigger);
+      const start = trigger.indexOf(OPEN);
+      return start >= 0 && text === trigger.slice(start);
+    }
+    if (row.kind !== "close") return false;
+    const call = parseStrictToolCall(text + this.decode(row.emit), this.tools);
+    if (!call) return false;
+    const tool = this.tools.find(tool => tool.function?.name === call.name);
+    const schema = tool?.function?.parameters;
+    const required = schema?.required;
+    if (!Array.isArray(required) || required.length !== 1 ||
+        Object.keys(call.arguments).length !== 1 || !Object.hasOwn(call.arguments, required[0])) return false;
+    const value = call.arguments[required[0]];
+    return typeof value !== "string" || literalsClosed(value);
+  }
+}

package/src/generation/fill-append.ts

@@ -0,0 +1,58 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache } from "../contracts/mlx/cache";
+import { withResource } from "../runtime/resources";
+import { leaseCacheStates } from "../state/leases";
+
+/** Append an already accepted span without intermediate vocabulary heads.
+ * Chunking changes the graph shape, so callers must qualify its numerics for
+ * their model. Verification uses its separate round operation because
+ * recurrent rollback records one forward per round. Append precedes emission/cancellation
+ * boundaries, preserving the cache-covered prefix on a consumer break.
+ */
+export async function appendFillHidden(
+  forward: (ids: MlxArray, cache: Cache[], positions: number) => MlxArray | Promise<MlxArray>,
+  cache: Cache[],
+  ids: number[],
+  chunkSize: (cache: readonly Cache[]) => number,
+): Promise<MlxArray> {
+  const input = ops.fromInt32(ids, [1, ids.length]);
+  try { return await appendHiddenRows(forward, cache, input, chunkSize); }
+  finally { input.dispose(); }
+}
+
+/** Append equal-length committed token rows [B,L]. The caller owns the input;
+ * the returned hidden rows belong to the caller. Model policy receives B so
+ * a larger cohort cannot silently cross its qualified projection geometry. */
+export async function appendHiddenRows(
+  forward: (ids: MlxArray, cache: Cache[], positions: number) => MlxArray | Promise<MlxArray>,
+  cache: Cache[],
+  input: MlxArray,
+  chunkSize: (cache: readonly Cache[], rows: number) => number,
+): Promise<MlxArray> {
+  const [rows, length] = input.shape as [number, number];
+  const hidden: MlxArray[] = [];
+  try {
+    for (let start = 0; start < length;) {
+      const limit = chunkSize(cache, rows);
+      if (!Number.isSafeInteger(limit) || limit < 1)
+        throw new Error("Model append chunk limit must be a positive integer");
+      const end = Math.min(length, start + limit);
+      if (start === 0 && end === length) return await forward(input, cache, length);
+      const part = input.slice([0, start], [rows, end]);
+      try {
+        const h = await forward(part, cache, end - start);
+        hidden.push(h);
+        // Bound transient graphs while retaining all positions for an optional
+        // diagnostic trace. Normal generation projects only the last position.
+        withResource(leaseCacheStates(cache), (state) => ops.evalAll([h, ...state]));
+      } finally {
+        part.dispose();
+      }
+      start = end;
+    }
+    return ops.concatAxis(hidden, 1);
+  } finally {
+    for (const h of hidden) h.dispose();
+  }
+}

package/src/generation/generate.ts

@@ -0,0 +1,44 @@
+// 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 {
+bindMlxAutoregressiveGraph
+} from "./bindings/autoregressive";
+import { bindLegacyDenoisingModel } from "./bindings/denoising";
+import { declaredGraph } from "../models/capabilities";
+import type { MlxTokenGraph } from "../models/graph";
+import { generateAutoregressive } from "./autoregressive";
+import { generateDenoising } from "./denoising";
+import { Generation } from "./result";
+import { GenerateDiagnostics,GenerateOptions } from "./types";
+
+export function generate(
+  model: MlxTokenGraph,
+  promptTokens: number[],
+  options: GenerateOptions = {},
+  diagnostics: GenerateDiagnostics = {},
+): Generation {
+  return bindGeneration(model)(promptTokens, options, diagnostics);
+}
+
+/** Resolve the compatibility model once at host construction. Every request
+ * borrows the same binding and immutable runtime snapshot. */
+export function bindGeneration(model: MlxTokenGraph): (
+  promptTokens: number[], options?: GenerateOptions, diagnostics?: GenerateDiagnostics,
+) => Generation {
+  // DiffusionGemma is non-autoregressive: route to the denoising engine instead
+  // of the AR decode loop. Same Generation/GenerateStats contract so the CLI and
+  // server stream it through the existing token machinery.
+  if (declaredGraph(model).graphCapabilities.method !== "denoising") {
+    const binding = bindMlxAutoregressiveGraph(model);
+    return (prompt, options, diagnostics) => generateAutoregressive(binding, prompt, options, diagnostics);
+  }
+  const binding = bindLegacyDenoisingModel(model);
+  return (prompt, options) => generateDenoising(binding, prompt, options);
+}