@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/contracts/portable/graph.ts

@@ -0,0 +1,80 @@
+/** Describes an implementation's ABI, not a claim of oracle parity or speed. */
+export interface GraphDescriptor {
+  readonly id: string;
+  readonly backend: string;
+  readonly graphAbi: string;
+  readonly stateAbi: string;
+  readonly artifact: string;
+}
+
+/** How prepared media reaches the graph.
+ * - `embeddings`: spliced prompt embeddings (image/audio soft tokens, optional
+ *   bidirectional and multimodal masks) enter the ordinary forward.
+ * - `embeddings+positions`: spliced embeddings with request-owned multi-axis
+ *   rotary positions and a decode delta.
+ * - `pixels`: raw image pixels are the denoising method's own prefill input. */
+export type GraphMediaInput = "embeddings" | "embeddings+positions" | "pixels";
+
+/** Portable facts about a bound graph. Composition binds the
+ * operations these facts require and refuses what is genuinely missing;
+ * scheduling reads only these declarations plus request and dynamic state
+ * (batch membership, cache contents, cancellation), never a model's class or
+ * name. A declaration that promises an operation is checked where it is bound. */
+export interface GraphCapabilities {
+  /** The method the graph's forward serves. Token-level methods (speculation,
+   * grammar proposals, fill) bind only to `autoregressive`. */
+  readonly method: "autoregressive" | "denoising";
+  /** Prepared media the graph accepts, or null. Requests carrying prepared
+   * embeddings need a `bindMediaInput` operation; `video` says video frames
+   * are accepted input. */
+  readonly media: { readonly input: GraphMediaInput; readonly video: boolean } | null;
+  /** The graph owns a compiled decode step. The composition keeps it on only
+   * where its storage has a compiled form (`Composition.compiledDecode`). */
+  readonly compiledDecode: boolean;
+  /** Attention selects a sparse subset of positions (surfaced to clients). */
+  readonly sparseAttention: boolean;
+  /** Derived from the bound graph's pooled text-embedding operation. */
+  readonly embeddings: boolean;
+  /** Hidden-layer outputs can be tapped for drafts that read them. */
+  readonly hiddenLayerTaps: boolean;
+  /** A draft head the graph's own checkpoint carries and has loaded, or null.
+   * `kind` names the native draft provider that binds to this graph; the
+   * provider registry resolves it, so no consumer names the graph or its
+   * provider. `numDraftTokens` is the head's round width. */
+  readonly nativeDraft: { readonly kind: string; readonly numDraftTokens: number } | null;
+  /** How prefill is sized for this graph. */
+  readonly prefill: {
+    /** Long prompts prefill in smaller chunks so the attention score matrices
+     * a fallback (unfused) attention materializes stay within a fixed
+     * workspace. It changes work size, never admission. */
+    readonly boundedWorkspace: boolean;
+  };
+  /** What the graph's target verification qualifies beyond ordinary rounds.
+   * Each is a property of the graph's verifier, independent of the draft that
+   * supplies proposals and of how the composition's caches store keys and
+   * values (every cache lego holds the speculative transaction). */
+  readonly speculation: {
+    /** Verification under mounted adapters. */
+    readonly adapters: boolean;
+    /** Verification of requests that return logprobs. */
+    readonly logprobs: boolean;
+    /** Verification of externally supplied tokens (echo of a known continuation). */
+    readonly externalTokens: boolean;
+    /** Verified grammar-constraint proposals. */
+    readonly grammarProposals: boolean;
+  };
+}
+
+/** Facts supplied by a graph at construction. Pooled embedding support is the
+ * presence of its operation; only the bound portable description carries the
+ * derived boolean. `never` also rejects separately typed declaration objects
+ * that try to supply that flag. */
+export type GraphDeclaration = Omit<GraphCapabilities, "embeddings"> & {
+  readonly embeddings?: never;
+};
+
+/** A fused InferenceMethod can bypass this graph interface entirely. It need
+ * only satisfy the method's output, cancellation, and ownership contracts. */
+export interface GraphFactory<Artifact, Binding, Graph> {
+  open(artifact: Artifact, binding: Binding): Promise<{ graph: Graph; close(): Promise<void> }>;
+}

package/src/contracts/portable/index.ts

@@ -0,0 +1,15 @@
+export * from "./composition";
+export * from "./denoising";
+export * from "./execution";
+export * from "./generation";
+export * from "./graph";
+export * from "./memory-plan";
+export * from "./object-cache";
+export * from "./prefill";
+export * from "./prefix-cache";
+export * from "./prompt";
+export * from "./resources";
+export * from "./rollback";
+export * from "./sampling";
+export * from "./scheduling";
+export * from "./token-delivery";

package/src/contracts/portable/memory-plan.ts

@@ -0,0 +1,26 @@
+/** Common admission contract for every runtime. Representation-specific
+ * planners may expose richer fields, but serving consumes only this shape. */
+export interface MemoryPlan {
+  readonly schemaVersion: 1;
+  /** Names the accounting the plan comes from; informational, never branched on. */
+  readonly strategy: string;
+  readonly fits: boolean;
+  readonly contextTokens: number;
+  readonly maxSafeContext: number;
+  readonly weightsBytes: number;
+  readonly kvBytes: number;
+  readonly transientBytes: number;
+  readonly reserveBytes: number;
+  readonly totalBytes: number;
+  readonly usableBytes: number;
+  readonly predictedDecodeTps: number | null;
+  /** mlx allocator-cache limit, distinct from the process admission limit. */
+  readonly allocatorLimitBytes?: number;
+  /** Generated-token allowance reserved within `contextTokens`; absent when the plan reserves none. */
+  readonly maxGenerationTokens?: number;
+  /** Parts of `weightsBytes` held in bounded streaming tiers instead of resident
+   * memory, by display label. Resident weights are `weightsBytes` minus their sum. */
+  readonly streamedWeights?: Readonly<Record<string, number>>;
+  /** RAM the plan leaves to the rest of the machine. */
+  readonly machineHeadroomBytes?: number;
+}

package/src/contracts/portable/object-cache.ts

@@ -0,0 +1,8 @@
+/** Exact reusable work, independent of token-prefix matching. The same cache
+ * owner chooses residency and persistence for objects and conversation state. */
+export interface ObjectCache<Value> {
+  /** Borrow immutable data through an owned lease. A miss permits recompute. */
+  take(key: string): Promise<{ value: Value; dispose(): void } | null>;
+  /** Successful return transfers ownership, including immediate eviction. */
+  put(key: string, value: Value): void;
+}

package/src/contracts/portable/prefill.ts

@@ -0,0 +1,63 @@
+/** Request-level chunk policy. Execution consumes the selected size; model
+ * backends can account for their temporary workspace without changing admission. */
+export interface PrefillPolicy {
+  chunkSize(promptTokens: number): number;
+}
+
+export interface PrefillPosition {
+  readonly length: number;
+  readonly position: number;
+  readonly chunkSize: number;
+  readonly tailSplit: boolean;
+  readonly snapshotAt?: number | null;
+}
+
+/** The phase a prompt step runs in, named by the planner that cut it: a whole
+ * planned `chunk` (exactly the step's chunk size), a `tail` (any other number
+ * of positions above one: what is left after the whole chunks, or a piece a
+ * snapshot boundary cuts short), or `decode`, a single position, which appends
+ * through the one-position read. */
+export type PrefillPhase = "chunk" | "tail" | "decode";
+
+export interface PrefillStep {
+  readonly start: number;
+  readonly end: number;
+  readonly kind: "drain" | "final";
+  readonly snapshot: boolean;
+  /** Batch admission may finish a short last drain and token zero in one tick.
+   * Otherwise execution yields between every drain and the next forward. */
+  readonly batchYield: boolean;
+  readonly phase: PrefillPhase;
+}
+
+/** The phase of a step the planner made `whole` (a full chunk) or not, from
+ * `start` to `end`. */
+export function prefillPhase(start: number, end: number, whole: boolean): PrefillPhase {
+  return end - start === 1 ? "decode" : whole ? "chunk" : "tail";
+}
+
+/** The prefill program a row follows, alone or in a batch. No tensors, runtime policy reads,
+ * or sampling. Advance position only after executing the returned step. */
+export function nextPrefillStep(input: PrefillPosition): PrefillStep {
+  const { length, position, chunkSize, tailSplit, snapshotAt } = input;
+  if (!Number.isSafeInteger(length) || !Number.isSafeInteger(position) ||
+      !Number.isSafeInteger(chunkSize) || chunkSize < 1 || position < 0 || position >= length)
+    throw new Error("invalid prefill position or chunk size");
+  if (snapshotAt != null && (!Number.isSafeInteger(snapshotAt) || snapshotAt < 0 || snapshotAt > length))
+    throw new Error("invalid prefill snapshot boundary");
+  if (snapshotAt != null && snapshotAt > position && snapshotAt < length) {
+    const whole = position + chunkSize <= snapshotAt;
+    const end = whole ? position + chunkSize : snapshotAt;
+    return { start: position, end, kind: "drain", snapshot: end === snapshotAt, batchYield: true,
+      phase: prefillPhase(position, end, whole) };
+  }
+  const remaining = length - position;
+  if (remaining > (tailSplit ? 1 : chunkSize)) {
+    const whole = !tailSplit || chunkSize <= remaining - 1;
+    const end = position + (whole ? chunkSize : remaining - 1);
+    return { start: position, end, kind: "drain", snapshot: false,
+      batchYield: !tailSplit || remaining - 1 > chunkSize, phase: prefillPhase(position, end, whole) };
+  }
+  return { start: position, end: length, kind: "final", snapshot: false, batchYield: false,
+    phase: prefillPhase(position, length, remaining === chunkSize) };
+}

package/src/contracts/portable/prefix-cache.ts

@@ -0,0 +1,33 @@
+/** A reusable prefix and the backend/method state that covers its tokens.
+ * The caller owns a returned hit. After releasing or returning its state,
+ * it must release any retained backing through retain. */
+export interface PrefixCacheHit<State, Attachment = never> {
+  tokens: number[];
+  caches: State;
+  attachments?: Attachment;
+  retain?: () => void;
+}
+
+/** Storage is independent of request scheduling and the method producing
+ * state. Both prefill and decode can publish processed-token prefixes.
+ * State is opaque here; its backend/method defines layout and alignment.
+ * Implementations own prefix selection, retention and storage tiers. */
+export interface PrefixCache<State, Attachment = never> {
+  /** Give optional retained state back to the allocator at an execution
+   * boundary. This never rejects or changes the active request. */
+  reclaim?(): void;
+  /** Drop session affinity without deleting reusable checkpoints. */
+  closeSession?(sessionId: string): void;
+  /** Prepare a reusable prefix before execution. The returned release ends
+   * request interest; cache residency and IO remain implementation-owned. */
+  prefetch?(prompt: number[], ns?: string, sessionId?: string): Promise<() => void>;
+  /** Return owned state for a usable session continuation, otherwise the
+   * longest usable prefix, or null on a miss. Session identity is an affinity
+   * hint and never changes the content or numerical namespace.
+   * Namespace distinguishes incompatible artifacts, methods and adapters. */
+  take(prompt: number[], ns?: string, sessionId?: string): PrefixCacheHit<State, Attachment> | null;
+  /** Transfer state and its optional backing release to storage. If this
+   * throws, ownership stays with the caller; successful return transfers
+   * ownership even when storage immediately spills or disposes the entry. */
+  put(tokens: number[], caches: State, ns?: string, retain?: () => void, attachments?: Attachment, sessionId?: string): void;
+}

package/src/contracts/portable/prompt.ts

@@ -0,0 +1,2 @@
+/** Execute native prompt preparation under caller-owned resource coordination. */
+export type PromptNativeWork = <T>(work: () => Promise<T>) => Promise<T>;

package/src/contracts/portable/resources.ts

@@ -0,0 +1,25 @@
+/** A single owner. Borrowing does not extend lifetime; transfer invalidates
+ * the old owner. close is idempotent, including when release throws. */
+export interface ResourceOwner<T> {
+  borrow(): T;
+  transfer(): T;
+  close(): void;
+}
+
+/** State views may own temporary handles or borrow live state. Consumers
+ * always close the lease; only the provider knows which handles to release. */
+export interface StateView<T> {
+  borrow(): readonly T[];
+  close(): void;
+}
+
+export interface DisposableResource {
+  dispose(): void;
+}
+
+/** A prepared replacement for live state. Commit once between method rounds;
+ * it publishes already-built state without graph work or validation. Dispose
+ * releases the unpublished replacement, or the superseded state after commit. */
+export interface PreparedStateChange extends DisposableResource {
+  commit(): void;
+}

package/src/contracts/portable/rollback.ts

@@ -0,0 +1,7 @@
+/** A complete transaction over every participating target state layer.
+ * A failed begin/resolve invalidates the run; its state must be discarded. */
+export interface SpeculativeTransaction<Acceptance = number> {
+  canBegin(drafts: number): boolean;
+  begin(drafts: number): void;
+  resolve(accepted: Acceptance): void;
+}

package/src/contracts/portable/sampling.ts

@@ -0,0 +1,44 @@
+/** Sampling owns score processing, token selection, and request-local history.
+ * Score/token representations and sync/device execution belong to the backend;
+ * this contract has no scheduling, cache, or tensor-runtime dependency. */
+export interface SampledToken<Token, Metadata> {
+  token: Token;
+  extras: Metadata | null;
+}
+
+export interface SamplingSession<Scores, DeviceToken, Result> {
+  /** Optional stateless operation shared across compatible rows and positions.
+   * It borrows scores and returns owned device tokens, without advancing history. */
+  readonly independent?: { sample(scores: Scores): DeviceToken };
+  /** Optional per-position operation for ONE request whose selection depends on
+   * no grammar or probability capture: position p's draw depends only on its
+   * scores, its step and its processor history. Borrows [N,V] scores and returns
+   * owned device tokens [N] for the given steps, without advancing history.
+   * A request with processor history (`needsHistory`) also passes the window's
+   * inputs [N] (borrowed: the pending token, then the drafts): position p reads
+   * its history extended by inputs[1..p], the history sample() would hold there
+   * had every earlier position drawn its draft. Draws are those sample() would
+   * make at each step: seeded sampling keys by (seed, step), never by a shared
+   * stream, so unused positions cost nothing. `commit` takes the tokens a walk
+   * over the window used, in order, and advances history as sample() would have
+   * advanced it drawing each. */
+  readonly positional?: {
+    sample(scores: Scores, steps: readonly number[], inputs?: DeviceToken): DeviceToken;
+    commit(tokens: readonly number[]): void;
+  };
+  /** Optional draw left on the device for a request whose selection depends on
+   * its processor history but on no grammar or probability capture: borrows
+   * [1,V] scores and returns the owned device token sample() would draw at
+   * `step`, committing it to history unread when sample() would commit it. */
+  readonly unread?: { sample(scores: Scores, step: number): DeviceToken };
+  readonly isPlainGreedy: boolean;
+  readonly capturesLogprobs: boolean;
+  readonly needsHistory: boolean;
+  /** Borrows scores; transfers ownership of the result to the caller. */
+  sample(logits: Scores, step: number): Result;
+  seedHistory(tokens: readonly number[]): void;
+  /** Commit accepted tokens when configured for manual history updates. */
+  commitDevice(token: DeviceToken): void;
+  commitNumbers(tokens: readonly number[]): void;
+  dispose(): void;
+}

package/src/contracts/portable/scheduling.ts

@@ -0,0 +1,54 @@
+/** Backend-owned context shared by compatible requests. A key identifies
+ * equivalent execution state; enter returns its scoped release. Scheduling
+ * observes group readiness without inspecting this context. */
+export interface ExecutionContext {
+  readonly key: string;
+  enter(): () => void;
+}
+
+/** A backend execution group exposes bounded work units and readiness. Queue
+ * payloads, tensors, cache merging and sampling remain inside its adapter. */
+export interface ExecutionGroup {
+  /** Native preparation may borrow the execution lease between work units.
+   * Payloads and task failures belong to the backend's queue. */
+  readonly pendingTasks?: number;
+  advanceTask?(): Promise<void>;
+  readonly active: number;
+  readonly queued: number;
+  readonly preparing: boolean;
+  readonly preparingRows?: number;
+  /** The preparation owner can accept another row before its next work unit. */
+  readonly canPrepareMore?: boolean;
+  /** Work estimates affect grouping only. A request larger than this budget
+   * still runs alone; it is never rejected or sent to another executor. */
+  readonly preparingTokens?: number;
+  readonly nextPreparationTokens?: number;
+  readonly maxPreparationTokens?: number;
+  /** The latest preparation published output that can be flushed by yielding. */
+  readonly preparationPublishedOutput?: boolean;
+  readonly maxActive: number;
+  readonly admissionHeld: boolean;
+  readonly closed: boolean;
+  pruneCancelled(): void;
+  /** May reject/remove an impossible queued request before returning false. */
+  admitNext(): boolean;
+  /** Read-only budget check used to preserve admission before decode. */
+  canBurst(): boolean;
+  /** Method-owned token demand for a compatible mixed iteration. Absent when
+   * the backend has no packed implementation for the current work. */
+  readonly mixedPreparation?: { readonly runningTokens: number; readonly minimumPreparationTokens: number };
+  readonly maxIterationTokens?: number;
+  advanceMixed?(tokenBudget: number): Promise<void>;
+  advancePreparation(): Promise<void>;
+  advance(): Promise<void>;
+  failActive(error: unknown): void;
+  failAll(error: unknown): void;
+  reserveResidency(): () => void;
+  acquireExecution?(): Promise<() => void>;
+  waitForWork(): Promise<void>;
+}
+
+export interface SchedulingClock {
+  now(): number;
+  yield(): Promise<void>;
+}

package/src/contracts/portable/token-delivery.ts

@@ -0,0 +1,4 @@
+import type { TokenLogprobs } from "./generation";
+
+/** Returning false stops delivery after the current token. */
+export type OnToken = (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>;

package/src/contracts/portable/tsconfig.json

@@ -0,0 +1,14 @@
+{
+  "compilerOptions": {
+    "lib": ["ES2022"],
+    "target": "ES2022",
+    "module": "Preserve",
+    "moduleResolution": "bundler",
+    "types": [],
+    "strict": true,
+    "skipLibCheck": false,
+    "noUncheckedIndexedAccess": true,
+    "noEmit": true
+  },
+  "include": ["./*.ts"]
+}

package/src/embeddings/index.ts

@@ -0,0 +1 @@
+export * from "./text";

package/src/embeddings/text.ts

@@ -0,0 +1,85 @@
+// Text-embedding entry point. Wraps a graph's `embedPooled` (last-token hidden,
+// L2-normalized vector) with the input convention its profile declares, so the
+// CLI (`mlx-bun embed`), the server (`/v1/embeddings`), and in-repo experiments
+// all produce the SAME vectors — the ones verified bit-exact vs mlx-lm and main in
+// tests/parity/embedding-parity.test.ts (opt-in) for the Qwen3-Embedding backbone.
+
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { declaredGraph } from "../models/capabilities";
+import type { LoadedTokenizer } from "../input/tokenizer";
+import * as ops from "@mlx-bun/mlx/ops";
+
+/** The operation behind a graph's `embeddings` capability. */
+export interface EmbeddingGraph { embedPooled(ids: MlxArray): MlxArray }
+
+/** A graph declares the pooled-embedding path. */
+export function isEmbeddingModel(model: object): model is EmbeddingGraph {
+  return typeof declaredGraph(model).embedPooled === "function";
+}
+
+/** The token id of the profile's declared pooling terminator in this tokenizer. */
+export function embeddingTerminatorId(tok: Pick<LoadedTokenizer, "encode">, terminator: string): number {
+  const ids = tok.encode(terminator, false);
+  if (ids.length !== 1) throw new Error(`embedding terminator ${JSON.stringify(terminator)} is not a single token in this tokenizer`);
+  return ids[0]!;
+}
+
+/** Qwen3-Embedding query format: a task instruction steers WHICH similarity axis
+ *  the geometry reflects ("represent this for topic routing" vs "...for
+ *  sentiment"). Documents are embedded raw; only queries get an instruction. */
+export function withInstruction(text: string, instruction?: string): string {
+  return instruction ? `Instruct: ${instruction}\nQuery:${text}` : text;
+}
+
+/** Tripwire: counts embedOne invocations (embedMany accrues via embedOne) so
+ *  callers can assert the pipeline embeds exactly as many times as expected. */
+let embedCallCount = 0;
+
+/** Reset the embed tripwire counter to 0. */
+export function resetEmbedCounter(): void {
+  embedCallCount = 0;
+}
+
+/** Read the embed tripwire counter (total embedOne calls since last reset). */
+export function getEmbedCounter(): number {
+  return embedCallCount;
+}
+
+export interface EmbedResult {
+  /** L2-normalized embedding (length = model hidden size). */
+  vector: Float32Array;
+  /** Token count of the embedded input (incl. the pooling token). */
+  tokens: number;
+}
+
+/** Embed one text. `terminator` is the pooling token id from
+ * `embeddingTerminatorId`; `instruction` (optional) applies the query format. */
+export function embedOne(
+  model: EmbeddingGraph,
+  tok: LoadedTokenizer,
+  terminator: number,
+  text: string,
+  instruction?: string,
+): EmbedResult {
+  embedCallCount++;
+  // addSpecialTokens=false matches the bit-exact parity path (no BOS;
+  // we append the pooling token ourselves).
+  const ids = [...tok.encode(withInstruction(text, instruction), false), terminator];
+  const idArr = ops.fromInt32(ids, [1, ids.length]);
+  const vec = model.embedPooled(idArr);
+  idArr.dispose();
+  const out = vec.toFloat32();
+  vec.dispose();
+  return { vector: out, tokens: ids.length };
+}
+
+/** Embed many texts (one forward each — the runtime is single-sequence). */
+export function embedMany(
+  model: EmbeddingGraph,
+  tok: LoadedTokenizer,
+  terminator: number,
+  texts: string[],
+  instruction?: string,
+): EmbedResult[] {
+  return texts.map((t) => embedOne(model, tok, terminator, t, instruction));
+}

package/src/execution/admission.ts

@@ -0,0 +1,59 @@
+import type { Cancellation } from "../contracts/portable/generation";
+import type { DisposableResource } from "../contracts/portable/resources";
+import { throwIfCancelled } from "./cancellation";
+
+export class AdmissionRejected extends Error {
+  constructor() { super("admission queue is full"); this.name = "AdmissionRejected"; }
+}
+
+/** FIFO reservations. A cancelled waiter releases no active owner's capacity. */
+export class AdmissionPool {
+  #active = 0;
+  #closed = false;
+  #waiting: Array<{ resolve(lease: DisposableResource): void; reject(error: unknown): void; unsubscribe(): void }> = [];
+  constructor(readonly capacity: number, readonly maxQueued = 64) {
+    if (!Number.isSafeInteger(capacity) || capacity < 1 || !Number.isSafeInteger(maxQueued) || maxQueued < 0)
+      throw new Error("invalid admission capacity");
+  }
+  get active(): number { return this.#active; }
+  get queued(): number { return this.#waiting.length; }
+
+  acquire(cancellation?: Cancellation): Promise<DisposableResource> {
+    if (this.#closed) return Promise.reject(new Error("admission is closed"));
+    try { if (cancellation) throwIfCancelled(cancellation); }
+    catch (error) { return Promise.reject(error); }
+    if (this.#active < this.capacity) return Promise.resolve(this.#lease());
+    if (this.#waiting.length >= this.maxQueued) return Promise.reject(new AdmissionRejected());
+    return new Promise((resolve, reject) => {
+      const waiter = { resolve, reject, unsubscribe() {} };
+      this.#waiting.push(waiter);
+      waiter.unsubscribe = cancellation?.subscribe((reason) => {
+        const index = this.#waiting.indexOf(waiter);
+        if (index < 0) return;
+        this.#waiting.splice(index, 1);
+        waiter.unsubscribe();
+        reject(new Error(`admission cancelled: ${reason}`));
+      }) ?? (() => {});
+      if (!this.#waiting.includes(waiter)) waiter.unsubscribe();
+    });
+  }
+
+  #lease(): DisposableResource {
+    this.#active++;
+    let released = false;
+    return { dispose: () => {
+      if (released) return;
+      released = true;
+      this.#active--;
+      const next = this.#waiting.shift();
+      if (next) { next.unsubscribe(); next.resolve(this.#lease()); }
+    } };
+  }
+
+  close(): void {
+    this.#closed = true;
+    for (const waiter of this.#waiting.splice(0)) {
+      waiter.unsubscribe(); waiter.reject(new Error("admission is closed"));
+    }
+  }
+}