@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/generation/bindings/autoregressive.ts

@@ -0,0 +1,104 @@
+import { pagedDecodeCompiles } from "../../state/paged/cache";
+import type { GraphDescriptor } from "../../contracts/portable/graph";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { Cache } from "../../contracts/mlx/cache";
+import type { MlxCompiledDecodeStep, MlxKnownSpans, MlxModelMemory } from "../../contracts/mlx/graph";
+import { compiledDecodeStepOf, tokenPhasesOf } from "../../models/capabilities";
+import { runtimeConfig, type RuntimeConfig } from "../../runtime/config";
+import type { MlxTokenGraph } from "../../models/graph";
+import type { PrefillPolicy } from "../../contracts/portable/prefill";
+import { resolveMlxPrefillPolicy } from "./prefill-policy";
+import { COMPOSITION_DEFAULTS, type KvStorageScheme, type PagedKvComposition } from "../../contracts/portable/composition";
+import { MlxPhaseDispatch } from "../phase-dispatch";
+
+export type { MlxModelMemory };
+
+export interface MlxDecodeStep {
+  /** Consume one pending token and advance state once. Returned arrays belong
+   * to the caller. A step that throws fails the generation. */
+  step(token: MlxArray, state: Cache[]): { logits: MlxArray; evalWith: MlxArray[] };
+  /** Release per-run workspaces after pending outputs are dropped and before
+   * cache disposal. Wait for native work if releasing requires completion. */
+  close(): void | Promise<void>;
+}
+
+/** One replaceable binding owns all graph-specific operations used by the AR
+ * loop: the graph's named phases and the facts the loop reads once. The
+ * current method requires MLX tensors and the legacy Cache[] ABI. Weights are
+ * borrowed for the binding's lifetime; generated caches are owned by the run.
+ * Caller-provided caches and media remain borrowed. */
+export interface MlxAutoregressiveBinding {
+  readonly runtime?: RuntimeConfig;
+  readonly prefillPolicy?: PrefillPolicy;
+  readonly descriptor: GraphDescriptor;
+  /** The graph's phase operations. */
+  readonly phases: MlxPhaseDispatch;
+  readonly eosTokenIds: readonly number[];
+  readonly memory: MlxModelMemory;
+  readonly adapters?: { active: string[] };
+  /** The composition's KV scheme, which `makeCache` composes; absent: bf16. */
+  readonly kv?: KvStorageScheme;
+  /** The composition's paging (`pagedKv`); absent: none. */
+  readonly pagedKv?: PagedKvComposition | null;
+  /** The composition's `prefillTailSplit`; absent: the default. */
+  readonly prefillTailSplit?: boolean;
+  /** The composition's `grammarJump`; absent: the default. */
+  readonly grammarJump?: boolean;
+  /** The graph's committed known-token spans, when it declares them. */
+  readonly knownSpans?: MlxKnownSpans;
+  makeCache(): Cache[];
+  /** The vision prefill phase over prepared embeddings (borrowed). */
+  prefillEmbeddings?(
+    embeddings: MlxArray, state: Cache[], imageMask: MlxArray | null,
+    ids: MlxArray, multimodalMask: MlxArray | null, width: number,
+  ): MlxArray;
+  /** The graph's compiled decode step for this run, chosen once after
+   * prefill: present when the graph declares one, the run allows replay and
+   * its composition kept compiled decode on. A compiled/fused decoder belongs to
+   * THIS graph, never to a model inferred from a name or inherited by a
+   * replacement. */
+  createDecode?(policy: { hasAdapters: boolean }): MlxDecodeStep | null;
+}
+
+/** Bind a resident graph to the autoregressive method: its phases, caches and
+ * memory, and the compiled decode step it declares. Nothing here reads which
+ * model the graph is. */
+export function bindMlxAutoregressiveGraph(model: MlxTokenGraph): MlxAutoregressiveBinding {
+  const runtime = runtimeConfig();
+  const facts = model.composition ?? COMPOSITION_DEFAULTS;
+  const compiledOn = (!model.composition?.pagedKv || pagedDecodeCompiles(model.composition.pagedKv, model.composition.kv));
+  return {
+    runtime,
+    prefillPolicy: resolveMlxPrefillPolicy(model, runtime),
+    descriptor: Object.freeze({ id: `legacy:${model.config.modelType}`, artifact: "legacy-resident-model",
+      stateAbi: "legacy-cache-array-v1", backend: "mlx", graphAbi: "mlx-hidden-bsh-v1" }),
+    phases: new MlxPhaseDispatch(tokenPhasesOf(model)),
+    eosTokenIds: model.config.eosTokenIds,
+    memory: model,
+    adapters: model.loraState,
+    ...(model.composition ? { kv: model.composition.kv } : {}),
+    pagedKv: facts.pagedKv ?? null, prefillTailSplit: facts.prefillTailSplit, grammarJump: facts.grammarJump,
+    ...(model.knownSpans ? { knownSpans: model.knownSpans } : {}),
+    makeCache: model.makeCache.bind(model),
+    prefillEmbeddings: model.embeddingPrompt ? (embeddings, caches, imageMask, ids, multimodalMask, rows) =>
+      model.embeddingPrompt!.find(range => rows >= range.from && rows <= range.to)!.operation(embeddings, caches, imageMask, ids, multimodalMask)
+      : model.prefillEmbeddings?.bind(model),
+    createDecode(policy) {
+      // Only a graph that declares a compiled step has one, and only a
+      // composition whose storage the step drives keeps it on. Adapters would
+      // bake residuals into the tape.
+      if (!compiledOn || policy.hasAdapters) return null;
+      const compiled: MlxCompiledDecodeStep | null = compiledDecodeStepOf(model);
+      if (!compiled) return null;
+      // The graph owns the compiled closures; its unload releases them.
+      return { step: (token, caches) => compiled.step(token, caches), close() {} };
+    },
+  };
+}
+
+export function assertMlxAutoregressiveBinding(binding: MlxAutoregressiveBinding): void {
+  const descriptor = binding.descriptor;
+  if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-hidden-bsh-v1" ||
+      descriptor.stateAbi !== "legacy-cache-array-v1")
+    throw new Error(`AR binding ${descriptor.id} has an incompatible backend, graph, or state ABI`);
+}

package/src/generation/bindings/denoising.ts

@@ -0,0 +1,60 @@
+import type { DenoisingGraph } from "../../contracts/portable/denoising";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { Cache } from "../../contracts/mlx/cache";
+import type { MlxModelMemory } from "./autoregressive";
+import type { MlxDeclaredGraph, MlxDenoisingOperations } from "../../contracts/mlx/graph";
+import { declaredGraph } from "../../models/capabilities";
+import { disposeResources } from "../../runtime/resources";
+import { runtimeConfig, type RuntimeConfig } from "../../runtime/config";
+
+export interface MlxDenoisingBinding<State = Cache[]> {
+  readonly runtime?: RuntimeConfig;
+  readonly graph: DenoisingGraph<MlxArray, State>;
+  readonly memory: MlxModelMemory;
+  readonly adapters?: { active: string[] };
+}
+
+/** A resident graph as the denoising binding takes it: its declaration, the denoising operations the
+ * declared method promises, its weights for the wired-limit scope and its mounted adapters. */
+export type MlxDenoisingModel = MlxDeclaredGraph & MlxModelMemory & Partial<MlxDenoisingOperations> &
+  { readonly loraState?: { active: string[] } };
+
+/** Only a graph that declares the denoising method binds here, and it must provide the operations. */
+export function bindLegacyDenoisingModel(resident: MlxDenoisingModel): MlxDenoisingBinding {
+  if (declaredGraph(resident).graphCapabilities.method !== "denoising")
+    throw new TypeError("the graph does not declare the denoising method");
+  const model = resident as MlxDenoisingOperations;
+  for (const operation of ["extendPrefill", "decoderLogits", "dequantEmbedWeight", "softEmbeddings"] as const)
+    if (typeof model[operation] !== "function")
+      throw new TypeError(`the graph declares the denoising method but provides no ${operation}`);
+  if (!model.denoisingPrefill && typeof model.prefill !== "function")
+    throw new TypeError("the graph declares the denoising method but provides no prefill");
+  return {
+    runtime: runtimeConfig(),
+    memory: resident, adapters: resident.loraState,
+    graph: {
+      descriptor: Object.freeze({ id: "legacy-diffusion-gemma", backend: "mlx",
+        graphAbi: "mlx-denoising-v1", stateAbi: "legacy-cache-array-v1", artifact: "legacy-resident-model" }),
+      vocabSize: model.config.text.vocabSize, canvasLength: model.canvasLength, embedScale: model.embedScale,
+      prefill: (ids, vision) => model.denoisingPrefill
+        ? model.denoisingPrefill.find(range => ids.length >= range.from && ids.length <= range.to)!.operation(ids, vision ?? null)
+        : vision ? model.prefillVision!(ids, vision) : model.prefill!(ids),
+      extendPrefill: (tokens,state) => model.denoisingAppend
+        ? model.denoisingAppend.find(range=>tokens.size>=range.from && tokens.size<=range.to)!.operation(tokens,state)
+        : model.extendPrefill(tokens,state),
+      decoderLogits: (canvas,state,feedback) => model.denoisingDecode
+        ? model.denoisingDecode.find(range=>canvas.size>=range.from && canvas.size<=range.to)!.operation(canvas,state,feedback)
+        : model.decoderLogits(canvas,state,feedback),
+      dequantEmbedWeight: model.dequantEmbedWeight.bind(model),
+      softEmbeddings: model.softEmbeddings.bind(model),
+      closeState: disposeResources,
+    },
+  };
+}
+
+export function assertMlxDenoisingGraph<State>(graph: DenoisingGraph<MlxArray, State>): void {
+  const descriptor = graph.descriptor;
+  if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-denoising-v1" ||
+      !descriptor.stateAbi)
+    throw new Error(`denoising graph ${descriptor.id} has an incompatible backend, graph, or state ABI`);
+}

package/src/generation/bindings/prefill-policy.ts

@@ -0,0 +1,36 @@
+import type { ModelConfig } from "../../artifacts/config";
+import type { GraphDeclaration } from "../../contracts/portable/graph";
+import { declaredGraph } from "../../models/capabilities";
+import type { PrefillPolicy } from "../../contracts/portable/prefill";
+import { sdpaFallbackBytes } from "../../state/kv-scheme";
+import type { RuntimeConfig } from "../../runtime/config";
+
+/** The graph facts prefill sizing reads: its geometry and its declaration. */
+export interface PrefillSizedGraph {
+  readonly config: ModelConfig;
+  readonly graphCapabilities: GraphDeclaration;
+}
+
+/** Retain large chunks for short prompts, but where the graph declares a
+ * bounded prefill workspace, bound the materialized attention score workspace
+ * at long context. This changes work size, never request admission. Explicit
+ * chunk settings remain exact. */
+export function resolveMlxPrefillPolicy(graph: PrefillSizedGraph, runtime: RuntimeConfig,
+  override?: number): PrefillPolicy {
+  const explicit = override ?? (runtime.value("MLX_BUN_RD_PREFILL_CHUNK") === undefined
+    ? undefined : runtime.number("MLX_BUN_RD_PREFILL_CHUNK", 2048));
+  if (explicit !== undefined) {
+    const chunk = Math.max(1, Math.floor(explicit));
+    return { chunkSize: () => chunk };
+  }
+  if (!declaredGraph(graph).graphCapabilities.prefill.boundedWorkspace) return { chunkSize: () => 2048 };
+  const config = graph.config;
+  // Capture geometry now, independently of later configuration mutation.
+  const geometry = { ...config, text: { ...config.text, layerTypes: [...(config.text.layerTypes ?? [])] } };
+  const workspaceBytes = 1024 ** 3;
+  return { chunkSize(promptTokens) {
+    let chunk = 2048;
+    while (chunk > 8 && sdpaFallbackBytes(geometry, chunk, promptTokens) > workspaceBytes) chunk /= 2;
+    return chunk;
+  } };
+}

package/src/generation/checkpoint-identity.ts

@@ -0,0 +1,60 @@
+import type { GenerateOptions } from "./types";
+import type { ResolvedExecution } from "../contracts/portable/execution";
+import { createHash } from "node:crypto";
+
+/** Sort object fields only; array order is meaningful for adapter composition,
+ * stop handling and quantization recipes. */
+function canonical(value: unknown): unknown {
+  if (Array.isArray(value)) return value.map(canonical);
+  if (value && typeof value === "object")
+    return Object.fromEntries(Object.entries(value).sort(([a], [b]) => a.localeCompare(b))
+      .map(([key, item]) => [key, canonical(item)]));
+  return value;
+}
+
+/** Stable identity for an exactly resumable generation. The prompt already
+ * includes rendered messages and tool schemas; the remaining fields cover
+ * every sampling policy that can change the continuation. The KV scheme is
+ * the composition's: the store holding the checkpoint is keyed by it. */
+export function generationCheckpointKey(
+  promptIds: number[], options: GenerateOptions & { stopSequences?: readonly string[] }, cacheNs = "",
+  execution?: ResolvedExecution, bindingIdentity?: unknown,
+): string {
+  const policy = {
+    version: 5,
+    execution: execution && { method: execution.method, mechanism: execution.mechanism,
+      fill: execution.fill,
+      compiledDecode: execution.compiledDecode, grammarJump: execution.grammarJump },
+    bindingIdentity,
+    stopSequences: options.stopSequences ?? [],
+    cacheNs,
+    promptIds,
+    maxTokens: options.maxTokens,
+    eosTokenIds: options.eosTokenIds,
+    temperature: options.temperature,
+    topP: options.topP,
+    topK: options.topK,
+    minP: options.minP,
+    minTokensToKeep: options.minTokensToKeep,
+    xtcProbability: options.xtcProbability,
+    xtcThreshold: options.xtcThreshold,
+    xtcSpecialTokens: options.xtcSpecialTokens,
+    // A missing client seed means "fresh randomness" for a new request, but
+    // an identical RETRY must find the durable checkpoint and restore the
+    // original random stream. Keep the generated seed out of request identity;
+    // it travels in checkpoint metadata instead. Explicit seeds remain part of
+    // the policy, so changing one cannot select an incompatible checkpoint.
+    seedMode: options.seedWasExplicit ? "explicit" : "server-default",
+    seed: options.seedWasExplicit ? options.seed : undefined,
+    hlg: options.hlg,
+    curve: options.curve,
+    logitBias: options.logitBias,
+    repetitionPenalty: options.repetitionPenalty,
+    repetitionContextSize: options.repetitionContextSize,
+    presencePenalty: options.presencePenalty,
+    presenceContextSize: options.presenceContextSize,
+    frequencyPenalty: options.frequencyPenalty,
+    frequencyContextSize: options.frequencyContextSize,
+  };
+  return createHash("sha256").update(JSON.stringify(canonical(policy))).digest("hex");
+}

package/src/generation/denoising.ts

@@ -0,0 +1,85 @@
+// Denoising generation: adapts a non-autoregressive denoising graph to the
+// Generation contract. The denoising engine runs its own prefill and canvas
+// loop to completion (`denoiseAsync`); only then are the finished canvas's
+// tokens yielded one by one, so nothing streams while it runs. There is no
+// chunked prefill, decode pipelining or early token-zero yield here (those are
+// the autoregressive path), and prefill/cache stats are zero.
+
+import { runtimeConfig } from "../runtime/config";
+import { type MlxDenoisingBinding } from "./bindings/denoising";
+import { adapterScoped,modelNeedsWiredLimit,usageScoped,wiredScoped } from "./scopes";
+import { denoiseAsync, type DiffusionGenOptions } from "./diffusion";
+import { Generation } from "./result";
+import { GenerateOptions,GenerateStats,GeneratedToken } from "./types";
+
+/** Denoising bindings use their own state and feedback graph. Each request
+ * draws from its own key sequence; callers still serialize model access. */
+export function generateDenoising<State>(
+  binding: MlxDenoisingBinding<State>, promptTokens: number[], options: GenerateOptions = {},
+): Generation {
+  let inner = generateDiffusionInner(binding, promptTokens, options);
+  if (options.adapters?.length && binding.adapters)
+    inner = adapterScoped({ loraState: binding.adapters }, options.adapters, inner);
+  if (binding.memory.expertRuntime?.finishUsage || binding.memory.expertRuntime?.flushUsage)
+    inner = usageScoped(binding.memory, inner);
+  const runtime = binding.runtime ?? runtimeConfig();
+  return new Generation(modelNeedsWiredLimit(binding.memory, undefined,
+    runtime.value("MLX_BUN_FORCE_WIRE") === "1") ? wiredScoped(inner) : inner, runtime);
+}
+
+/** The served denoising request: greedy confidence-threshold, a fresh canvas
+ *  seed unless the caller pins one, the checkpoint's stopping set {1, 106}
+ *  united with any caller EOS, and 256 tokens unless the caller sets a limit. */
+export function denoisingRequestOptions(options: GenerateOptions): DiffusionGenOptions {
+  const seed =
+    options.seed !== undefined
+      ? BigInt(options.seed)
+      : BigInt(Math.floor(Math.random() * 0x7fffffff));
+  return {
+    maxTokens: options.maxTokens ?? 256,
+    sampler: "confidence-threshold",
+    temperature: 0,
+    eosTokenIds: [...new Set([1, 106, ...(options.eosTokenIds ?? [])])],
+    seed,
+    visionPixels: options.visionPixels,
+  };
+}
+
+/** Non-autoregressive diffusion generation, adapted to the AR Generation
+ *  contract. Runs the denoising engine (its own prefill + canvas loop) and
+ *  streams the emitted tokens. v1: greedy (temperature 0, confidence-threshold
+ *  sampler — the OptiQ default); per-block intra-stream + temperature>0
+ *  (categorical) are follow-ups. */
+async function* generateDiffusionInner<State>(
+  binding: MlxDenoisingBinding<State>,
+  promptTokens: number[],
+  options: GenerateOptions,
+): AsyncGenerator<GeneratedToken, GenerateStats> {
+  options.signal?.throwIfAborted();
+  const t0 = performance.now();
+  const result = await denoiseAsync(binding.graph, promptTokens, denoisingRequestOptions(options), options.signal);
+  const decodeMs = performance.now() - t0;
+  let index = 0;
+  let failure: { error: unknown } | undefined;
+  try {
+    for (const token of result.tokens) {
+      options.signal?.throwIfAborted();
+      yield { token, index: index++ };
+    }
+  } catch (error) { failure = { error }; }
+  finally {
+    // A stop parser may close delivery before the finished canvas is exhausted.
+    // Settle the emitted count on return(), without suppressing cancellation/errors.
+    if (failure) throw failure.error;
+    return {
+      promptTokens: promptTokens.length,
+      cachedTokens: 0,
+      generatedTokens: index,
+      prefillTps: 0,
+      decodeTps: index / Math.max(decodeMs / 1000, 1e-9),
+      prefillMs: 0,
+      decodeMs,
+      cacheTokens: [],
+    };
+  }
+}