mlx-bun 0.2.0 → 0.4.0

package/src/backends/mlx/speculative-round.ts

@@ -0,0 +1,161 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { Cancellation, GenerationOutput, TokenLogprobs } from "../../contracts/generation";
+import type { DisposableResource } from "../../contracts/resources";
+import { cleanupFailure, disposeResources } from "../../engine/resources";
+import type { SpeculativeTransaction } from "../../inference/rollback";
+import { DraftAcceptance } from "../../inference/draft-acceptance";
+import { deliverDraftOutputs, type DraftOutputResult } from "../../inference/draft-output";
+import { readStepExtras, type NumberStepSampler } from "../../sampler";
+
+/** Draft state owns proposal generation and its committed hidden/KV state.
+ * It borrows verification context and knows nothing about target storage. */
+export interface MlxDraftRows<Context = MlxArray> {
+  /** Each row returns 0..depth proposals. The method verifies the returned
+   * lengths; an empty row still consumes its pending token. */
+  draft(pending: readonly number[], depth: number, steps: readonly number[]): number[][] | Promise<number[][]>;
+  commit(accepted: readonly number[], context: Context): void | Promise<void>;
+}
+
+/** A bound graph transfers both tensors on success and releases partial
+ * allocations on failure. State retention is a separate injected interface. */
+export interface MlxRowVerification<Context extends DisposableResource = MlxArray> {
+  readonly transaction: SpeculativeTransaction<readonly number[]>;
+  forward(ids: MlxArray): Promise<{ logits: MlxArray; context: Context }>;
+}
+
+export interface MlxSpeculativeRow {
+  readonly pending: number;
+  readonly step: number;
+  readonly remaining: number;
+  readonly eosTokenIds: readonly number[];
+  readonly sampling: NumberStepSampler;
+  readonly grammarDone?: () => boolean;
+}
+
+export interface MlxSpeculativeRowResult {
+  readonly drafts: readonly number[];
+  readonly acceptance: DraftAcceptance;
+  /** One entry per content token; stopping EOS is excluded. */
+  readonly logprobs: readonly (TokenLogprobs | undefined)[];
+}
+
+export interface MlxSpeculativeOutputRow extends MlxSpeculativeRow {
+  readonly output: Pick<GenerationOutput, "commit">;
+  readonly cancellation?: Cancellation;
+}
+
+/** Execute one method step through output delivery and aligned state commit.
+ * The caller retires terminal rows and publishes eligible completed state
+ * before changing membership. Consumer failures are returned per request;
+ * a graph/state failure throws and requires discarding the whole group. */
+export async function advanceSpeculativeOutputs<Context extends DisposableResource = MlxArray>(
+  rows: readonly MlxSpeculativeOutputRow[], depth: number,
+  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
+): Promise<{ rounds: readonly MlxSpeculativeRowResult[]; outputs: readonly DraftOutputResult[] }> {
+  const round = await prepareSpeculativeRows(rows, depth, draft, target);
+  let outputs: DraftOutputResult[];
+  try {
+    outputs = await deliverDraftOutputs(round.results.map((result, row) => ({
+      acceptance: result.acceptance, logprobs: result.logprobs,
+      output: rows[row]!.output, cancellation: rows[row]!.cancellation,
+    })));
+    await round.commit(outputs.map(output => output.accepted));
+  } catch (error) { return cleanupFailure(error, () => round.dispose()); }
+  round.dispose();
+  return { rounds: round.results, outputs };
+}
+
+/** Retention stays unresolved while the method delivers output. A terminal
+ * row may keep fewer accepted inputs for an output-aligned checkpoint;
+ * continuing rows retain their complete accepted prefix. Commit exactly once,
+ * then dispose; disposing an uncommitted round requires discarding its state. */
+export interface MlxPreparedSpeculativeRound extends DisposableResource {
+  readonly results: readonly MlxSpeculativeRowResult[];
+  commit(keepAccepted?: readonly number[]): Promise<void>;
+}
+
+/** One committed method step at B=1 or B>1. The caller supplies live rows with
+ * a positive token budget and a common nonnegative draft depth. Admission,
+ * cancellation boundaries, output delivery and cache publication stay outside
+ * this step. On failure the caller discards the participating method state. */
+export async function advanceSpeculativeRows<Context extends DisposableResource = MlxArray>(
+  rows: readonly MlxSpeculativeRow[], depth: number,
+  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
+): Promise<readonly MlxSpeculativeRowResult[]> {
+  const round = await prepareSpeculativeRows(rows, depth, draft, target);
+  try { await round.commit(); }
+  catch (error) { return cleanupFailure(error, () => round.dispose()); }
+  round.dispose();
+  return round.results;
+}
+
+/** Draft, verify and sample without deciding which output-aligned state to
+ * retain. Graph-specific context is opaque to this operation. */
+export async function prepareSpeculativeRows<Context extends DisposableResource = MlxArray>(
+  rows: readonly MlxSpeculativeRow[], depth: number,
+  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
+): Promise<MlxPreparedSpeculativeRound> {
+  const proposals = await draft.draft(rows.map(row => row.pending), depth, rows.map(row => row.step));
+  const B = rows.length;
+  const verifyDepth = proposals.reduce((max, tokens) => Math.max(max, tokens.length), 0);
+  using ids = ops.fromInt32(proposals.flatMap((tokens, row) => {
+    const window = [rows[row]!.pending, ...tokens];
+    // Causal right padding only fills the rectangular graph. Sampling stops
+    // at this row's own bonus position; commit discards every padded input.
+    while (window.length <= verifyDepth) window.push(rows[row]!.pending);
+    return window;
+  }), [B, verifyDepth + 1]);
+  target.transaction.begin(verifyDepth);
+  const verification = await target.forward(ids);
+  const dispose = () => disposeResources([verification.context, verification.logits]);
+  try {
+    const results = await sampleSpeculativeRows(rows, proposals, verification.logits);
+    return {
+      results,
+      async commit(accepted = results.map(result => result.acceptance.accepted)) {
+        target.transaction.resolve(accepted);
+        await draft.commit(accepted, verification.context);
+      },
+      dispose,
+    };
+  } catch (error) { return cleanupFailure(error, dispose); }
+}
+
+/** Sample a borrowed verify window. Each request owns its processor history,
+ * grammar and RNG stream; completing one walk never samples its unused suffix.
+ * The existing request adapter and grouped method use this same operation. */
+export async function sampleSpeculativeRows(
+  rows: readonly MlxSpeculativeRow[], proposals: readonly (readonly number[])[], logits: MlxArray,
+): Promise<MlxSpeculativeRowResult[]> {
+  const V = logits.shape[2]!, width = logits.shape[1]!;
+  const independent = rows[0]?.sampling.independent;
+  let tokens: number[] | undefined;
+  if (independent && rows.every(row => row.sampling.independent === independent)) {
+    using flat = ops.reshape(logits, [rows.length * width, V]);
+    using sampled = independent.sample(flat);
+    tokens = sampled.toIntTokens();
+  }
+  const results: MlxSpeculativeRowResult[] = [];
+  for (let row = 0; row < rows.length; row++) {
+    const request = rows[row]!;
+    const acceptance = new DraftAcceptance(proposals[row]!, request.remaining, request.eosTokenIds);
+    const logprobs: (TokenLogprobs | undefined)[] = [];
+    while (!acceptance.done) {
+      const position = acceptance.position;
+      let token: number, metadata: TokenLogprobs | undefined;
+      if (tokens) token = tokens[row * width + position]!;
+      else {
+        using slice = logits.slice([row, position, 0], [row + 1, position + 1, V]);
+        using scores = ops.reshape(slice, [1, V]);
+        const sampled = await request.sampling.sample(scores, request.step + acceptance.emitted.length);
+        token = sampled.token;
+        metadata = readStepExtras(sampled.extras);
+      }
+      acceptance.accept(token, request.grammarDone?.() ?? false);
+      if (!acceptance.sawEos) logprobs.push(metadata);
+    }
+    results.push({ drafts: proposals[row]!, acceptance, logprobs });
+  }
+  return results;
+}

package/src/backends/mlx/speculative.ts

@@ -0,0 +1,114 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { Cache } from "../../model/gemma4";
+import type { RuntimeModel } from "../../model/factory";
+import type { Sampler } from "../../sampler";
+import type { DraftProvider, DraftSource } from "../../spec/source";
+import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
+import type { MlxModelMemory } from "./autoregressive";
+import type { SpeculativeTransaction } from "../../inference/rollback";
+import { bindCacheRollback } from "./rollback";
+import type { GraphDescriptor } from "../../inference/graph";
+import { bindLegacyDraftTarget } from "./draft-target";
+
+/** Bound target operations and draft construction for one speculative run.
+ * A replacement graph supplies this entire port, including any hidden taps.
+ * The verifier never needs a RuntimeModel or an artifact-family check. */
+export interface MlxSpeculativeBinding {
+  readonly runtime?: RuntimeConfig;
+  readonly memory?: MlxModelMemory;
+  readonly descriptor: GraphDescriptor;
+  readonly eosTokenIds: readonly number[];
+  readonly prefillTailSplit: boolean;
+  makeCache(): Cache[];
+  openDraft(sampler: Sampler, caches: Cache[]): DraftSource;
+  bindRollback(caches: Cache[]): SpeculativeTransaction;
+  forward(ids: MlxArray, caches: Cache[], tapLayers?: number[]):
+    Promise<{ hidden: MlxArray; ctxML: MlxArray | null }>;
+  projectLogits(hidden: MlxArray): MlxArray;
+  /** Establish the implementation's verify kernel context; restore on close. */
+  pinVerify?(): { close(): void };
+}
+
+/** Legacy mutable tap/kernel fields require the gateway's exclusive lease. */
+export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftProvider): MlxSpeculativeBinding {
+  const runtime = runtimeConfig();
+  return {
+    runtime,
+    memory: model,
+    descriptor: Object.freeze({ id: `legacy-spec:${model.config.modelType}`, backend: "mlx",
+      graphAbi: "mlx-hidden-bsh-v1", stateAbi: "legacy-cache-array-v1", artifact: "legacy-resident-model" }),
+    eosTokenIds: model.config.eosTokenIds,
+    prefillTailSplit: runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true),
+    makeCache: model.makeCache.bind(model),
+    openDraft: (sampler, caches) => provider.open({ sampler, target: bindLegacyDraftTarget(model, caches) }),
+    bindRollback: bindCacheRollback,
+    forward: (ids, caches, tapLayers) => legacyForwardWithTaps(model, ids, caches, tapLayers),
+    projectLogits: model.logitsFromHidden.bind(model),
+    ...("setSpecKernelPinned" in model ? {
+      pinVerify() {
+        model.setSpecKernelPinned(true);
+        return { close() { model.setSpecKernelPinned(false); } };
+      },
+    } : {}),
+  };
+}
+
+export function assertMlxSpeculativeBinding(binding: MlxSpeculativeBinding): void {
+  const descriptor = binding.descriptor;
+  if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-hidden-bsh-v1" ||
+      descriptor.stateAbi !== "legacy-cache-array-v1")
+    throw new Error(`speculative binding ${descriptor.id} has an incompatible backend, graph, or state ABI`);
+}
+
+/** Run one target forward, and when `tapLayers` is set (⟹ a Gemma4 target — a
+ *  DSpark-style source), also return the captured multi-layer context
+ *  [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. */
+async function legacyForwardWithTaps(
+  model: RuntimeModel,
+  ids: MlxArray,
+  caches: Cache[],
+  tapLayers: number[] | undefined,
+): 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 };
+  }
+  if (!("hiddenTap" in model)) throw new Error("target graph does not support hidden taps");
+  const m = model;
+  const previousTap = m.hiddenTap;
+  const cap = new Map<number, MlxArray>();
+  m.hiddenTap = { layers: new Set(tapLayers), captured: cap };
+  let hidden: MlxArray | null = null;
+  try {
+    hidden = model.forwardHidden(ids, caches);
+    const perLayer = tapLayers.map((li) => {
+      const a = cap.get(li);
+      if (!a) throw new Error(`spec tap: layer ${li} not captured`);
+      return a;
+    });
+    const ctxML = ops.concatAxis(perLayer, 2); // [1,L,m*H]
+    for (const [, a] of cap) a.dispose();
+    cap.clear(); // consumed — the finally must not double-dispose
+    const out = { hidden, ctxML };
+    hidden = null; // ownership returned to the caller
+    return out;
+  } finally {
+    // On any throw (forward mid-capture, missing layer, concat), free the
+    // partially-captured tap tensors and the orphaned hidden. On success both
+    // are already gone (cap cleared, hidden nulled).
+    hidden?.dispose();
+    for (const [, a] of cap) a.dispose();
+    m.hiddenTap = previousTap;
+  }
+}

package/src/backends/mlx/state-rows.ts

@@ -0,0 +1,104 @@
+import type { BatchableCache, Cache } from "../../model/gemma4-base";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../../engine/resources";
+import type { PreparedStateChange } from "../../contracts/resources";
+
+/** Preserve the graph's layout capabilities when an empty sibling is made. */
+export interface MlxRowLayout<Layout extends BatchableCache> extends BatchableCache {
+  makeEmptyBatch(): Layout;
+}
+
+/** Layer layouts own tensor geometry. This collection only applies the same
+ * row membership operation to every layer. It has no scheduling, sampling,
+ * checkpoint-key or persistence policy. */
+export class MlxStateRows<Layout extends MlxRowLayout<Layout> = BatchableCache> {
+  /** Borrowed by the bound target graph; collection ownership stays here. */
+  readonly caches: Layout[];
+  #rows = 0;
+
+  /** Takes ownership of empty, compatible layer layouts. */
+  constructor(layouts: readonly Layout[]) { this.caches = [...layouts]; }
+  get rowCount(): number { return this.#rows; }
+  /** Retain an immutable batch boundary without extracting or materializing
+   * individual rows. Layer layouts decide how to share their backing arrays. */
+  clone(): MlxStateRows<Layout> {
+    const layouts: Layout[] = [];
+    try {
+      for (const source of this.caches) {
+        const layout = source.makeEmptyBatch(); layouts.push(layout);
+        layout.mergeRows([source]);
+      }
+      const result = new MlxStateRows(layouts); result.#rows = this.#rows;
+      return result;
+    } catch (error) { return cleanupFailure(error, () => disposeResources(layouts)); }
+  }
+  /** Prepare an initial cohort with one merge per layer. */
+  mergeRows(rows: readonly (readonly Cache[])[]): void {
+    const next: Layout[] = [];
+    try {
+      for (let layer = 0; layer < this.caches.length; layer++) {
+        const layout = this.caches[layer]!.makeEmptyBatch(); next.push(layout);
+        layout.mergeRows(rows.map(row => row[layer]!));
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    const previous = this.caches.splice(0, this.caches.length, ...next);
+    this.#rows = rows.length; disposeResources(previous);
+  }
+
+  /** Borrow a prepared request's state. Build every replacement before
+   * releasing the active group, so a failed admission leaves it intact. */
+  append(row: readonly Cache[]): void {
+    applyStateChanges([() => this.prepareAppend(row)]);
+  }
+
+  /** Stage membership so a method can compose it with companion-state admission. */
+  prepareAppend(row: readonly Cache[]): PreparedStateChange {
+    return this.prepareAppendMany([row]);
+  }
+
+  /** Stage several admissions with one membership operation per layer. */
+  prepareAppendMany(rows: readonly (readonly Cache[])[]): PreparedStateChange {
+    const next: Layout[] = [];
+    try {
+      for (let layer = 0; layer < this.caches.length; layer++) {
+        const layout = this.caches[layer]!;
+        const merged = layout.makeEmptyBatch();
+        next.push(merged);
+        merged.mergeRows([...(this.#rows ? [layout] : []), ...rows.map(row => row[layer]!)]);
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    return {
+      commit: () => {
+        const previous = this.caches.splice(0, this.caches.length, ...next.splice(0));
+        this.#rows += rows.length;
+        next.push(...previous);
+      },
+      dispose: () => disposeResources(next.splice(0)),
+    };
+  }
+
+  /** The method calls this only between committed steps. A layer failure
+   * invalidates the group; its caller discards the state. */
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.clear(); return; }
+    for (const layout of this.caches) layout.filterRows(keep);
+    this.#rows = keep.length;
+  }
+
+  /** Returns independent request state for a cache publisher or a new group. */
+  extractRow(row: number): Cache[] {
+    const state: Cache[] = [];
+    try { for (const layout of this.caches) state.push(layout.extractRow(row)); return state; }
+    catch (error) { return cleanupFailure(error, () => disposeResources(state)); }
+  }
+
+  clear(): void {
+    const next: Layout[] = [];
+    try { for (const layout of this.caches) next.push(layout.makeEmptyBatch()); }
+    catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    const previous = this.caches.splice(0, this.caches.length, ...next);
+    this.#rows = 0;
+    disposeResources(previous);
+  }
+
+  dispose(): void { this.#rows = 0; disposeResources(this.caches); }
+}

package/src/backends/mlx/state-views.ts

@@ -0,0 +1,46 @@
+import type { StateView } from "../../contracts/resources";
+import { disposeResources, ownResource } from "../../engine/resources";
+import type { MlxArray } from "../../mlx/array";
+import { isBatchableCache, type Cache } from "../../model/gemma4-base";
+import { cloneKvCaches, type CacheCodecProvider } from "../../kv-store";
+
+/** Compatibility boundary for legacy state(). Numerical code consumes a lease,
+ * never the legacy ownership marker. No copy, evaluation or fence is added. */
+export function leaseCacheState(cache: Cache): StateView<MlxArray> {
+  return ownResource<readonly MlxArray[]>(cache.state(), cache.stateNeedsDispose ? disposeResources : () => {});
+}
+
+/** Preserve the existing prefill evaluation order: temporary views first,
+ * borrowed state second. Partial acquisition releases every acquired view. */
+export function leaseCacheStates(caches: readonly Cache[]): StateView<MlxArray> {
+  const leases: StateView<MlxArray>[] = [];
+  const close = () => disposeResources(leases.map((lease) => ({ dispose: () => lease.close() })));
+  try {
+    for (const owned of [true, false]) {
+      for (const cache of caches) if (!!cache.stateNeedsDispose === owned) leases.push(leaseCacheState(cache));
+    }
+    return ownResource(leases.flatMap((lease) => [...lease.borrow()]), close);
+  } catch (error) {
+    try { close(); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "state view acquisition failed"); }
+    throw error;
+  }
+}
+
+/** Precision/storage transitions constrain reuse independently of scheduling. */
+export function minimumReusableOffset(caches: readonly Pick<Cache, "minimumReusableOffset">[]): number {
+  return Math.max(0, ...caches.map(cache => cache.minimumReusableOffset ?? 0));
+}
+
+/** Snapshot an adopted single-row state into persistence-compatible caches.
+ * Row layouts own compact extraction; ordinary serial caches use their codec. */
+export function cloneSingleRowState(caches: readonly Cache[], codecs?: CacheCodecProvider): Cache[] {
+  const owned: Cache[] = [];
+  try {
+    for (const cache of caches) owned.push(...(isBatchableCache(cache)
+      ? [cache.extractRow(0)] : cloneKvCaches([cache], codecs)));
+    return owned;
+  } catch (error) {
+    try { disposeResources(owned); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "single-row snapshot failed"); }
+    throw error;
+  }
+}

package/src/backends/mlx/target-layout-capability.ts

@@ -0,0 +1,35 @@
+import { Glm52Cache } from "../../model/glm52-cache";
+import { SpeculativeRotatingKVCache } from "../../model/speculative-rotating-kv";
+import { RotatingAffineLayout, SpeculativeRotatingAffineLayout } from "../../model/rotating-kv-layout";
+import type { Cache } from "../../model/gemma4-base";
+import { KVCache, QuantizedKVCache, TurboQuantKVCache, RotatingKVCache, RotatingQuantizedKVCache } from "../../model/gemma4-base";
+import { DelayedRotatingQuantizedKVCache } from "../../model/delayed-rotating-quantized-kv";
+import { DelayedQuantizedKVCache } from "../../model/delayed-quantized-kv";
+import { DelayedTurboQuantKVCache } from "../../model/delayed-turboquant-kv";
+import { SSMCache } from "../../model/qwen3-delta";
+import { BatchedSSMCache } from "../../model/batched-ssm";
+import { BatchedKVCache } from "../../model/batched-kv";
+import { BatchedQuantizedKVCache } from "../../model/batched-quantized-kv";
+import { BatchedTurboQuantKVCache } from "../../model/batched-turboquant-kv";
+
+export type TargetRowLayout = Glm52Cache | DelayedRotatingQuantizedKVCache | RotatingAffineLayout | SpeculativeRotatingKVCache | BatchedKVCache | BatchedQuantizedKVCache | BatchedTurboQuantKVCache |
+  BatchedSSMCache | DelayedTurboQuantKVCache | DelayedQuantizedKVCache;
+
+/** The storage binding advertises a row transaction factory. The gateway can
+ * compose supported model graphs without identifying their model family. */
+export function targetRowLayoutFactory(cache: Cache): (() => TargetRowLayout) | undefined {
+  if (cache instanceof Glm52Cache) return () => cache.makeEmptyBatch();
+  if (cache instanceof DelayedRotatingQuantizedKVCache)
+    return () => new DelayedRotatingQuantizedKVCache(cache.maxSize, cache.groupSize, cache.bits, cache.start, cache.maintain, undefined, true);
+  if (cache instanceof RotatingAffineLayout || cache instanceof RotatingQuantizedKVCache)
+    return () => new SpeculativeRotatingAffineLayout(cache.maxSize, cache.groupSize, cache.bits);
+  if (cache instanceof RotatingKVCache || cache instanceof SpeculativeRotatingKVCache) return () => new SpeculativeRotatingKVCache(cache.maxSize);
+  if (cache instanceof DelayedTurboQuantKVCache || cache instanceof DelayedQuantizedKVCache ||
+      cache instanceof BatchedTurboQuantKVCache || cache instanceof BatchedKVCache ||
+      cache instanceof BatchedQuantizedKVCache || cache instanceof BatchedSSMCache) return () => cache.makeEmptyBatch();
+  if (cache instanceof TurboQuantKVCache) return () => new BatchedTurboQuantKVCache(cache.kBits, cache.vBits);
+  if (cache instanceof SSMCache) return () => new BatchedSSMCache();
+  if (cache instanceof QuantizedKVCache) return () => new BatchedQuantizedKVCache(cache.groupSize, cache.bits);
+  if (cache instanceof KVCache) return () => new BatchedKVCache();
+  return undefined;
+}

package/src/backends/mlx/text-engine.ts

@@ -0,0 +1,73 @@
+import { createCompletionMethod } from "../../serve/session-completion-engine";
+import type { ExecutionPlanner, InferenceEngine, Timer } from "../../contracts/generation";
+import { createInferenceEngine, type EngineOptions } from "../../engine/engine";
+import { throwIfCancelled } from "../../engine/cancellation";
+import type { GenerateOptions, GenerateStats } from "../../generate";
+import type { GenerationGateway, RequestShape } from "../../serve/generation-gateway";
+
+/** Data-only API for text requests whose policy/shape was already
+ * resolved by the caller. Native request resources still belong to the prepared
+ * completion API; they are deliberately absent from this token API. */
+export type TextGenerationOptions = Omit<GenerateOptions,
+  "signal" | "cache" | "grammar" | "fill" | "adapters" | "promptEmbeddings" |
+  "imageMask" | "multimodalMask" | "visionPixels" | "onPrefillDone" |
+  "onDecodeCheckpoint" | "initialPendingToken" | "initialGeneratedTokens" |
+  "originalPromptTokens" | "snapshotAt" | "checkpointEveryTokens"
+>;
+
+export interface TextGenerationRequest {
+  readonly promptIds: readonly number[];
+  readonly options: TextGenerationOptions;
+  readonly shape: RequestShape;
+}
+
+export const hostTimer: Timer = {
+  after(milliseconds, callback) {
+    const timer = setTimeout(callback, milliseconds);
+    return () => { clearTimeout(timer); };
+  },
+};
+
+/** Reuses the gateway's placement, exclusive lease, prompt cache, batching,
+ * and quant-specific execution. The portable session never imports them.
+ * Closing this engine cancels its sessions; it does not dispose a shared gateway. */
+export function createTextInferenceEngine(
+  gateway: Pick<GenerationGateway, "place" | "run">,
+  options: Omit<EngineOptions, "timer"> & { timer?: Timer } = {},
+): InferenceEngine<TextGenerationRequest, GenerateStats> {
+  const planner: ExecutionPlanner<TextGenerationRequest, GenerateStats> = {
+    async plan(request, cancellation) {
+      throwIfCancelled(cancellation);
+      if (request.shape.hasVision || request.shape.hasAdapters || request.shape.hasGrammar)
+        throw new Error("native request resources require the prepared completion API");
+      const placement = gateway.place(request.shape, request.options);
+      const maxTokens = request.options.maxTokens ?? 512;
+      return {
+        id: `text:${placement.mechanism}`, outputTokenLimit: maxTokens,
+        method: createCompletionMethod(gateway, {
+          prompt: [...request.promptIds], options: request.options,
+          shape: request.shape, placement,
+        }),
+      };
+    },
+  };
+  const engine = createInferenceEngine(planner, { ...options, timer: options.timer ?? hostTimer });
+  return {
+    open(request, control) {
+      // Snapshot before demand-start: callers may reuse/mutate request arrays.
+      // TextGenerationOptions contains data only; no native handles or callbacks.
+      const snapshot = structuredClone(request);
+      if (!snapshot.promptIds.length || snapshot.promptIds.some((id) => !Number.isInteger(id) || id < 0 || id > 0xffffffff))
+        return Promise.reject(new Error("prompt must contain valid token ids"));
+      const maxTokens = snapshot.options.maxTokens ?? 512;
+      if (!Number.isSafeInteger(maxTokens) || maxTokens <= 0)
+        return Promise.reject(new Error("text generation requires a positive token limit"));
+      if ((snapshot.options.topLogprobs ?? 0) > (options.maxTopLogprobs ?? 20))
+        return Promise.reject(new Error("top logprobs exceed session delivery capacity"));
+      if (control.output === "collect" && (snapshot.options.logprobs || snapshot.options.topLogprobs))
+        return Promise.reject(new Error("token logprobs require stream consumption"));
+      return engine.open(snapshot, control);
+    },
+    close: () => engine.close(),
+  };
+}