@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/contracts/mlx/drafter.ts

@@ -0,0 +1,176 @@
+// Drafter models as the draft sources see them. A drafter is a model: it is built
+// and loaded in `models/` (`models/drafters.ts`), and the sources in
+// `generation/speculative/` consume only these ports, never a concrete drafter.
+
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { AssistantDonors } from "./attention";
+import type { Cache } from "./cache";
+import type { DraftProjection } from "./draft-projection";
+import type { DraftProjectionKernels, ResidualBasis } from "./draft-target";
+import type { MlxRowTable } from "./row-table";
+
+export interface DrafterOperationRange<Operation> { readonly from: number; readonly to: number; readonly operation: Operation }
+export type DrafterOperations<Operation> = MlxRowTable<Operation> | readonly DrafterOperationRange<Operation>[];
+
+/** One drafted step of a KV-borrowing assistant: owned token ids and the post-projected hidden. */
+export interface DrafterRowsStep { tokens: MlxArray; nextHidden: MlxArray }
+
+/** One assistant draft step over B rows: the last token's embedding and the hidden it follows [B, 1, H]
+ *  (borrowed), the donors, and the rows' position (one shared, or [B]); owned tokens and next hidden out. */
+export type AssistantDraftStep = (lastTokenEmb: MlxArray, targetHidden: MlxArray, donors: AssistantDonors,
+  position: number | MlxArray) => DrafterRowsStep;
+
+/** The optiq KV-borrowing assistant: no cache of its own, it reads the target's donor attention. */
+export interface AssistantDrafterModel {
+  /** The draft step per row count or range of the rows it drafts for, each with the token selection its head
+   *  bound for that class at construction; the caller indexes it by its rows (actual positive row count). */
+  readonly forwardRows: MlxRowTable<AssistantDraftStep>;
+  dispose(): void;
+}
+
+/** What a projected-context drafter (DSpark, DeepSpec) attends over the target's tapped hidden states. */
+export interface ContextKV {
+  k: MlxArray; // [1, nKvHeads, ctxLen, headDim]
+  v: MlxArray; // [1, nKvHeads, ctxLen, headDim]
+}
+
+/** A projected-context drafter's context projection over B × L rows: the target's tapped hiddens
+ *  [B, L, taps × H] (borrowed) to context rows [B, L, H] (owned). */
+export type DrafterContextProjection = (targetHiddens: MlxArray) => MlxArray;
+
+/** Each drafter layer's context K/V [B, kvHeads, L, headDim] (owned) of B × L context rows [B, L, H] (borrowed),
+ *  at absolute positions from `position` (one shared, or [B] when the rows start apart). */
+export type DrafterContextKV = (context: MlxArray, position: number | MlxArray) => ContextKV[];
+
+/** Confidence-scheduled draft-length pruning and logits collection of one DSpark block. */
+export interface DsparkDraftOptions {
+  thresholds?: number[];
+  minConf?: number;
+  /** Concatenate the per-position draft logits into the result; default true. */
+  collectLogits?: boolean;
+  collectConfidence?: boolean;
+}
+
+/** Storage owns context validity; the DSpark graph supplies only Q and the block's K/V. */
+export interface DsparkContextAttention {
+  attend(layer: number, query: MlxArray, keys: MlxArray, values: MlxArray, scale: number): MlxArray;
+}
+
+/** Our trained DSpark module: block drafting over the target's multi-layer hidden context. */
+export interface DsparkDrafterModel {
+  readonly cfg: {
+    readonly gamma: number;
+    readonly nLayers: number;
+    readonly tapLayers: number[];
+    readonly sts?: { readonly thresholds: number[] };
+  };
+  projectContextRows(hidden: MlxArray): { k: MlxArray; v: MlxArray }[];
+  forwardRows(model: DraftProjection, hCtx: MlxArray | null, anchors: readonly number[], gamma: number,
+    opts?: DsparkDraftOptions, context?: DsparkContextAttention): { tokens: number[][] };
+  dispose(): void;
+}
+
+/** Storage owns context validity; the DeepSpec graph supplies only Q and new block KV. */
+export interface DeepspecContextAttention {
+  attend(layer: number, query: MlxArray, blockKeys: MlxArray, blockValues: MlxArray): MlxArray;
+}
+
+export interface DeepspecDraftRows {
+  tokens: number[][];
+  conf: number[][];
+  baseLogits?: MlxArray;
+}
+
+/** DeepSeek's released DSpark drafter (DeepSpec): a block forward over cached, already projected context K/V. */
+export interface DeepspecDrafterModel {
+  readonly cfg: { readonly num_target_layers: number; readonly num_hidden_layers: number };
+  /** The target layers whose hiddens the drafter conditions on. */
+  readonly tapLayers: number[];
+  /** The trained block width. */
+  readonly gamma: number;
+  readonly hidden: number;
+  /** The context projection per row count or range of the rows projected (B × L), indexed by the caller
+   *  (`B × L`): MLX's quantized matmul (matmul for a BF16 checkpoint) at every class. */
+  readonly projectContext: MlxRowTable<DrafterContextProjection>;
+  /** Each layer's context K/V per row count or range of the context rows (B × L), indexed like `projectContext`:
+   *  the same kernel at every class. */
+  readonly projectContextKVRows: MlxRowTable<DrafterContextKV>;
+  draftRows(context: DeepspecContextAttention, anchors: readonly number[], position: number | MlxArray,
+    collectLogits?: boolean): DeepspecDraftRows;
+  dispose(): void;
+}
+
+/** One recurrent multi-token-prediction forward: token embeddings [B, S, H] paired with hiddens
+ * [B, S, H] (both borrowed) in, the module output (post final norm) [B, S, H] out, owned; appends S
+ * positions per row to `cache`, which supplies positions and storage. */
+export type RecurrentMtpForward = (tokenEmbeds: MlxArray, hiddens: MlxArray, cache: Cache) => MlxArray;
+
+/** The recurrent multi-token-prediction block split from a Qwen release, its phases named by the
+ * caller like a graph's: the draft step indexed by the rows it drafts for. */
+export interface RecurrentMtpModule {
+  /** One position per row (a draft step, the round's bridge or full-accept row), by the rows' row count. */
+  readonly decode: MlxRowTable<RecurrentMtpForward>;
+  /** Several positions per row (prompt rows, a verified continuation). */
+  readonly window: RecurrentMtpForward;
+}
+
+/** The checkpoint-native multi-token-prediction row of a graph that declares `nativeDraft`, over
+ * compressed (MLA) attention state. The graph owns the numerics and the head's weights; the draft
+ * source owns sampling and state retention. */
+export interface NativeMtpHead {
+  readonly hiddenSize: number;
+  /** Geometry of the compressed cache the head appends to. */
+  readonly cache: { readonly kvLoraRank: number; readonly ropeHeadDim: number; readonly maxTokens: number };
+  /** One draft step: `ids` [B, 1] (one token per sequence) over each row's anchor
+   * `hidden` [B, 1, hidden], appended to `cache`; the caller owns the output
+   * [B, 1, hidden]. */
+  decode(ids: MlxArray, hidden: MlxArray, cache: Cache): Promise<MlxArray>;
+  /** A window of two or more tokens per sequence, `ids` [B, L] over `hidden`
+   * [B, L, hidden], appended to `cache`; the caller owns the output. The caller
+   * names the phase from the window it built: one token per sequence is `decode`. */
+  window(ids: MlxArray, hidden: MlxArray, cache: Cache): Promise<MlxArray>;
+  /** Logits [B, positions, V] of the head's output. */
+  project(hidden: MlxArray): MlxArray;
+}
+
+/** DFlash 2's projected-context attention: block keys follow the context rows. */
+export interface Dflash2ContextAttention {
+  attend(layer: number, query: MlxArray, blockKeys: MlxArray, blockValues: MlxArray, scale?: number): MlxArray;
+}
+
+/** The DFlash 2 block drafter: target taps → projected context K/V; one greedy block per anchor.
+ *  It loads in the basis it was trained on and learns the target's when a draft group binds. */
+export interface Dflash2DrafterModel {
+  readonly cfg: { readonly layers: number; readonly numTargetLayers: number };
+  readonly tapLayers: number[];
+  /** Draft tokens per block (block size minus the anchor). */
+  readonly gamma: number;
+  /** Apply the target's declarations before any projection: its residual basis (null: the trained basis) and
+   *  the kernels its composition names for the drafter's 4-bit projections (null: the drafter's own rule). The
+   *  first call applies them; every later call must pass the same declarations. */
+  bindTarget(basis: ResidualBasis | null, projections: DraftProjectionKernels | null): void;
+  /** The context projection (fc over the target's taps, then the hidden norm) per row count or range of the rows
+   *  projected (B × L: a prompt's rows at prefill, a commit's accepted rows), indexed by the caller
+   *  (`B × L`). A 4-bit fc runs the kernel the target declares for the class or, by the drafter's
+   *  own rule, the simdgroup matrix unit (`Affine4MmaLinear`, where the weight's shape fits it) at 5..8
+   *  and MLX's quantized matmul (`QuantizedLinear`) at every other class; a BF16 fc runs one matmul at every
+   *  class. No entry runs before `bindTarget`. */
+  readonly projectContext: DrafterOperations<DrafterContextProjection>;
+  /** Each layer's k and v over the context rows (k norm, RoPE), per row count or range of those rows (B × L), indexed
+   *  like `projectContext` and on the kernel it names for the class. */
+  readonly projectContextKVRows: DrafterOperations<DrafterContextKV>;
+  /** The columns of a block drafting `depth` tokens: its anchor and up to `depth` mask positions, at most the
+   *  trained block size. */
+  blockColumns(depth: number): number;
+  /** The draft block per row count or range of its rows (anchors times columns), indexed by the scheduler: every 4-bit
+   *  projection on the kernel `projectContext` runs for the class; BF16 projections run one matmul at every
+   *  class. */
+  readonly draft: DrafterOperations<Dflash2DraftBlock>;
+  dispose(): void;
+}
+
+/** One DFlash 2 draft block: one row per anchor, each the anchor and `columns - 1` mask positions, over the
+ *  projected context; the greedy proposals left on the device, [anchors, columns - 1] uint32 row-major. Every
+ *  projection in the block runs the kernel its row count names. */
+export type Dflash2DraftBlock = (context: Dflash2ContextAttention, projection: DraftProjection, anchors: readonly number[],
+  position: number | MlxArray, columns: number) => MlxArray;

package/src/contracts/mlx/forward-work.ts

@@ -0,0 +1,23 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { Cache } from "./cache";
+import type { MlxLayerCapture, MlxPhaseGroup } from "./graph";
+
+/** One forward a method hands to the mixed iteration: the phase its rows run
+ * in, their count and width class, the drafter's capture, and the forward
+ * those rows run on their own, which the iteration calls when it does not pack
+ * them with another group. */
+export interface MlxForwardStep {
+  readonly phase: MlxPhaseGroup["phase"];
+  /** Rows: sequences times positions. */
+  readonly rows: number;
+  readonly capture: MlxLayerCapture | null;
+  own(): Promise<MlxArray>;
+}
+
+export type MlxForwardWork = (ids: MlxArray, cache: Cache[], step: MlxForwardStep) => Promise<MlxArray>;
+
+export interface MlxPreparationWork {
+  /** Real prompt tokens across all rows, not a padded sequence width. */
+  readonly maxTokens?: number;
+  readonly forward?: MlxForwardWork;
+}

package/src/contracts/mlx/glm52-weights.ts

@@ -0,0 +1,25 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+
+/** Structural seam used by tiny fixtures as well as ColibriGlm52Weights. */
+export interface Glm52MlaWeightSource {
+  tensor(name: string): MlxArray;
+  dequantized(name: string, outputRows: number, inputColumns: number): MlxArray;
+  linear(
+    x: MlxArray,
+    name: string,
+    outputRows: number,
+    inputColumns: number,
+  ): MlxArray;
+}
+
+export interface Glm52WeightSource extends Glm52MlaWeightSource {
+  readonly weightsBytes: number;
+  has(name: string): boolean;
+  embedding(
+    ids: MlxArray,
+    name: string,
+    vocabSize: number,
+    hiddenSize: number,
+  ): MlxArray;
+  dispose(): void;
+}

package/src/contracts/mlx/graph.ts

@@ -0,0 +1,309 @@
+import type { MlxHeadForward, MlxHeadRange } from "./head";
+export type { MlxHeadForward, MlxHeadRange } from "./head";
+import type { GraphDeclaration } from "../portable/graph";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { Cache } from "./cache";
+import type { TargetView } from "./draft-target";
+import type { NativeMtpHead } from "./drafter";
+import type { TrainableGraph } from "./trainable";
+import type { MediaEncoders, MediaSidecarProbes, MlxPromptInput, PixelInput, TextEmbeddingModel, Vision } from "./media";
+import type { MlxRowTable } from "./row-table";
+export type { MlxRowTable } from "./row-table";
+
+/** Weight and expert-residency facts the wired-limit scope reads. */
+export interface MlxModelMemory {
+  readonly weightsBytes: number;
+  /** The graph's composition measured that its weights are wired for a generation whatever their share of the
+   *  recommended working set; absent: the share decides. */
+  readonly wiredWeights?: boolean;
+  readonly expertRuntime?: {
+    readonly plan: { readonly plannedBytes: number };
+    flushUsage?: () => void;
+    finishUsage?: () => Promise<void>;
+  } | null;
+}
+
+/** A verify forward's hidden-layer capture. The graph calls it once per
+ * decoder layer, in layer order, with that layer's residual output, then once
+ * with index `layerCount` (the decoder's layer count) and the post-final-norm
+ * output, all before the forward returns. `hidden` [sequences, tokens per
+ * sequence, hidden-width] is borrowed and possibly unevaluated: it is valid
+ * only during the call. The callback copies the layers it keeps (for example
+ * with `ops.contiguous`) and owns and disposes those copies. */
+export type MlxLayerCapture = (layer: number, hidden: MlxArray) => void;
+
+/** A prefill or decode forward. `ids` is integer token ids [sequences, tokens
+ * per sequence]. The caller owns `ids` and may dispose them once the call
+ * returns (or its promise settles); the result does not depend on the caller's
+ * handle. The graph borrows `state` and appends every position of `ids` to
+ * each layer's cache (keys and values, or recurrent state), advancing each
+ * sequence by its token count; it never rolls back. The returned post-final-norm
+ * hidden [sequences, tokens per sequence, hidden-width] is owned by the caller,
+ * who disposes it, and may be unevaluated, as may the appended cache state.
+ * A forward that throws may leave a partial append, which the caller discards
+ * or rolls back. */
+export type MlxPhaseForward = (ids: MlxArray, state: Cache[]) => MlxArray | Promise<MlxArray>;
+
+/** A forward that also hands the drafter the hidden layers it taps:
+ * `MlxPhaseForward`'s ownership and append, and `capture` receives each
+ * decoder layer's residual output and then the post-final-norm output (see
+ * `MlxLayerCapture`). `capture` is null when the composition's drafter reads
+ * none; only a graph that declares `hiddenLayerTaps` receives a non-null
+ * capture. Prompt chunks, prompt tails and verify windows take it. */
+export type MlxCapturingForward = (ids: MlxArray, state: Cache[], capture: MlxLayerCapture | null) =>
+  MlxArray | Promise<MlxArray>;
+
+/** A speculative verify forward (`MlxCapturingForward`). */
+export type MlxVerifyForward = MlxCapturingForward;
+
+
+/** The scoring and training forward, not a serving phase: every position of
+ * `ids` [sequences, tokens] as one causal window over the caller's caches (the
+ * stateless training caches, which own the mask of an unpadded or right-padded
+ * batch and the training attention kernel), on the kernels of `width`, the
+ * actual number of its rows. Returns the post-final-norm hidden the caller owns. */
+export type MlxSequenceForward = (ids: MlxArray, state: Cache[], width: number) => MlxArray;
+
+/** One contiguous range of positions in a single sequence's prompt piece.
+ * The scheduler selects the operation; the graph never inspects its shape. */
+export interface MlxPromptRange {
+  readonly from: number;
+  readonly to: number;
+  readonly operation: MlxCapturingForward;
+  readonly unit?: "gpu" | "ane";
+}
+
+/** A whole accepted span uses one declared operation; callers must not split
+ * it to fit a verification table because that can change attention arithmetic. */
+export interface MlxKnownRange { readonly from: number; readonly to: number; readonly operation: MlxPhaseForward }
+export type MlxKnownOperations = MlxRowTable<MlxPhaseForward> | readonly MlxKnownRange[];
+
+/** Operations selected by exact row count or prompt position range. */
+export interface MlxTokenPhases {
+  readonly prompt: readonly MlxPromptRange[];
+  /** Optional chunk bindings when a file distinguishes chunks from prompt tails. */
+  readonly promptChunk?: readonly MlxPromptRange[];
+  readonly multiSequencePrompt: readonly MlxPromptRange[];
+  readonly multiSequencePromptChunk?: readonly MlxPromptRange[];
+  readonly decode: MlxRowTable<MlxCapturingForward>;
+  readonly verify: MlxRowTable<MlxVerifyForward>;
+  readonly appendKnown: MlxKnownOperations;
+  readonly logits: MlxRowTable<MlxHeadForward>;
+}
+
+/** How a graph whose composition qualifies committed known-token spans takes
+ * them (`MlxTokenPhases.appendKnown` over its committed read): how many
+ * positions one call takes. Every cache lego's committed read keeps the
+ * one-position arithmetic (the cache contract holds it to the same positions
+ * read one at a time), whatever the composition's KV scheme. A graph without
+ * it takes known tokens one position at a time. */
+export interface MlxKnownSpans {
+  /** Positions one sequence's `appendKnown` call takes at the caches' current
+   * length. Recheck after each call: native arithmetic can change at a
+   * cache-length boundary. */
+  maxChunkSize(state: readonly Cache[]): number;
+}
+
+/** One group of a forward that serves several groups of rows at once, each in
+ * its own phase over its own caches. */
+export interface MlxPhaseGroup {
+  /** `decode` (one position per sequence), `window` (a prompt piece) or
+   * `verify` (a speculative window, which keeps its own projection geometry). */
+  readonly phase: "decode" | "window" | "verify";
+  /** The actual number of rows in this group. */
+  readonly rows: number;
+  readonly ids: MlxArray;
+  readonly state: Cache[];
+  readonly capture: MlxLayerCapture | null;
+}
+
+/** Declared by a composition that serves mixed rows: decode, prompt and verify
+ * groups in one forward, each group's attention and recurrence over its own
+ * caches for its own phase, the row-wise layers over the groups' rows packed
+ * together unless a verify group keeps its geometry. `totalRows` is the full
+ * round's row count, including its separate verification groups. Borrows inputs and state; returns one
+ * owned post-final-norm hidden per group. */
+export type MlxMixedForward = (groups: readonly MlxPhaseGroup[], totalRows: number) => MlxArray[];
+
+/** A graph's compiled single-token decode step: the whole forward replayed from
+ * one recorded graph instead of rebuilt per token. The scheduler chooses it
+ * once, for a row whose caches all declare the compiled-decode capability and
+ * hold no batched layout; it has no fallback. */
+export interface MlxCompiledDecodeStep {
+  /** Consume one pending uint32 [1] token (unevaluated is fine) and advance
+   * every cache one position. The returned logits [1,1,V] belong to the
+   * caller; `evalWith` holds cache updates that must ride the same async
+   * evaluation. A step that throws leaves `state` exactly as it was. */
+  step(token: MlxArray, state: Cache[]): { logits: MlxArray; evalWith: MlxArray[] };
+}
+
+/** method `denoising`: what the canvas denoiser drives over a graph's weights, with the prefill state
+ * an array of caches. A graph that declares the method implements these; binding refuses one that does not. */
+export interface MlxDenoisingPrefillRange {
+  readonly from: number;
+  readonly to: number;
+  readonly operation: (promptIds: number[], pixels: MlxArray | null) => Cache[];
+}
+
+export interface MlxDenoisingDecodeRange {
+  readonly from: number;
+  readonly to: number;
+  readonly operation: (canvas: MlxArray, state: Cache[], feedback: MlxArray | null) => MlxArray;
+}
+export interface MlxDenoisingAppendRange {
+  readonly from: number;
+  readonly to: number;
+  readonly operation: (tokens: MlxArray, state: Cache[]) => void;
+}
+
+export interface MlxDenoisingOperations {
+  readonly denoisingDecode?: readonly MlxDenoisingDecodeRange[];
+  readonly denoisingAppend?: readonly MlxDenoisingAppendRange[];
+  readonly denoisingPrefill?: readonly MlxDenoisingPrefillRange[];
+  readonly config: { readonly text: { readonly vocabSize: number } };
+  readonly canvasLength: number;
+  readonly embedScale: number;
+  prefill?(promptIds: number[]): Cache[];
+  /** Image-conditioned prefill from channel-first pixels. */
+  prefillVision?(promptIds: number[], pixels: MlxArray): Cache[];
+  extendPrefill(tokens: MlxArray, state: Cache[]): void;
+  decoderLogits(canvas: MlxArray, state: Cache[], feedback: MlxArray | null): MlxArray;
+  dequantEmbedWeight(): MlxArray;
+  softEmbeddings(logits: MlxArray, weight: MlxArray): MlxArray;
+}
+
+/** The adaptive draft gate's constants for one drafter on one machine
+ *  (`MlxDeclaredGraph.adaptiveDraftGates`); each absent one keeps the gate's
+ *  default (`AdaptiveDraftGateOptions`). */
+export interface AdaptiveDraftGateSettings {
+  /** Accepted and rejected counts every draft position starts from: their ratio
+   *  is the acceptance the gate assumes before it observes the text, their sum
+   *  how many rounds of evidence that assumption weighs. */
+  readonly priorAccepted?: number;
+  readonly priorRejected?: number;
+  /** The factor a position's counts keep per round that observes it, and every
+   *  count per round without drafts. */
+  readonly decay?: number;
+  readonly idleDecay?: number;
+  /** How much a round's expected rate must exceed a step without drafts. */
+  readonly margin?: number;
+  /** Rounds without drafts before a probe, doubling after each probe that
+   *  leaves the gate drafting nothing, up to `maxProbeInterval`. */
+  readonly probeInterval?: number;
+  readonly maxProbeInterval?: number;
+  /** The depth of a probe a drafter proposes without a verify round. */
+  readonly probeDepth?: number;
+}
+
+/** The graph's operations and construction facts. Pooled embedding support is
+ * declared by its operation alone. Other operations accompany the promises in
+ * `graphCapabilities`; composition checks their pairing when it binds
+ * (`declaredGraph`). */
+export interface MlxDeclaredGraph {
+  readonly graphCapabilities: GraphDeclaration;
+  /** Release what the graph built at construction (computed constants, Neural
+   *  Engine programs, its own weight source where it opened one) once every
+   *  execution and compiled runner has stopped. A graph that built nothing
+   *  releases nothing. Its owner (the host, a draft provider) calls it before
+   *  releasing the weights it supplied. */
+  dispose(): void;
+  /** Unscaled token embeddings, for media prompts that splice embeddings. */
+  readonly embed?: TextEmbeddingModel["embed"];
+  /** media `embeddings`/`embeddings+positions`: bind one request's prepared
+   * embeddings into an input whose prefill is indivisible. */
+  bindMediaInput?(input: Vision): MlxPromptInput;
+  /** Presence declares pooled embeddings; there is no separate support flag.
+   * The pooled, L2-normalized sentence vector of one sequence,
+   * ids [1, L] to [1, hidden]. The caller appends the pooling token the
+   * profile declares and owns the result. */
+  embedPooled?(ids: MlxArray): MlxArray;
+  /** media: the lazy encoders this graph loads from the checkpoint at `modelDir`. */
+  mediaEncoders?(modelDir: string, probes: MediaSidecarProbes): Promise<MediaEncoders>;
+  /** media `pixels`: the pixel input a denoising method consumes; null when the checkpoint has none. */
+  pixelInput?(): PixelInput | null;
+  /** Ports over this graph's live caches that draft sources may consume. */
+  draftTarget?(caches: Cache[]): TargetView;
+  /** `nativeDraft`: the checkpoint's own multi-token-prediction head, refused when its tier is not loaded. */
+  nativeDraftHead?(): NativeMtpHead;
+
+  /** What the trainer consumes: the head, segmented backward, prefix sharing,
+   * gradient checkpointing, flash-attention constraint and denoising objective
+   * this graph implements. Undeclared parts make the trainer refuse. */
+  readonly trainable?: TrainableGraph;
+  /** `compiledDecode`: the graph's compiled decode step, created on first use
+   * and owned by the graph until `releaseCompiledDecode`. */
+  compiledDecodeStep?(): MlxCompiledDecodeStep;
+  /** Retire the compiled step before the graph unloads: its recorded graph
+   * borrows the graph's weights and constants. Safe to call when none exists. */
+  releaseCompiledDecode?(): void;
+  /** Snapshot of the streamed expert residency, for diagnostics; null when the graph streams none. */
+  expertResidency?(): Record<string, unknown> | null;
+  /** Measured cost of a verify forward by its row count, relative to one
+   *  single-row decode step (rows 1 = 1). Declared only by a graph measured on
+   *  the hardware it serves; speculative depth scheduling reads it. */
+  readonly verifyRoundCosts?: ReadonlyMap<number, number>;
+  /** The widest adaptive round per draft provider kind, measured on the hardware
+   *  it serves: an adaptive drafter of a listed kind drafts up to it instead of
+   *  its own width, unless the request names a width. Absent, or a kind it does
+   *  not list: the drafter's own width. */
+  readonly adaptiveDraftWidths?: ReadonlyMap<string, number>;
+  /** The adaptive draft gate's settings per draft provider kind, measured on the
+   *  hardware it serves (an adaptive drafter of a listed kind gates with them).
+   *  Absent, or a kind it does not list: the gate's defaults. */
+  readonly adaptiveDraftGates?: ReadonlyMap<string, AdaptiveDraftGateSettings>;
+  /** Committed known-token spans, declared when the composition qualifies them
+   *  (`appendKnown` then reads committed). */
+  readonly knownSpans?: MlxKnownSpans;
+  /** Mixed rows in one forward, declared by a composition that serves them. */
+  readonly mixed?: readonly MlxMixedRange[];
+  /** Mixed groups that retain their own projection geometry. The scheduler
+   * routes each group independently when a verify group is present. */
+  readonly mixedSeparate?: MlxSeparateMixed;
+  /** Full-sequence scoring operations, indexed by total token rows. */
+  readonly sequence?: readonly MlxSequenceRange[];
+  readonly sequenceLogits?: readonly MlxHeadRange[];
+  /** Prepared embeddings without explicit positions, indexed by total rows. */
+  readonly embeddingPrompt?: readonly import("./media").MlxEmbeddingsRange[];
+}
+
+/** Cache layout written by the graph, independent of request storage policy. */
+export type MlxCacheLayer =
+  | { readonly kind: "full"; readonly turboQuant: "dense" | "deferred"; readonly kv4Head256?: boolean; readonly kv4DecodeBlock?: number; readonly windowRead?: "folded" }
+  | { readonly kind: "sliding"; readonly window: number }
+  | { readonly kind: "recurrent" }
+  | { readonly kind: "latent"; readonly geometry: {
+      readonly kvLoraRank: number; readonly ropeHeadDim: number; readonly maxTokens?: number; readonly indexHeadDim?: number;
+    } };
+
+/** Fixed mixed operation for a contiguous range of packed token rows. */
+export interface MlxMixedRange {
+  readonly from: number;
+  readonly to: number;
+  /** Optional bounds on all round rows, including independently routed verify groups. */
+  readonly totalFrom?: number;
+  readonly totalTo?: number;
+  readonly operation: MlxMixedForward;
+}
+
+
+/** A full-sequence forward fixed for a contiguous range of token rows. */
+export interface MlxSequenceRange {
+  readonly from: number;
+  readonly to: number;
+  readonly operation: (ids: MlxArray, state: Cache[]) => MlxArray;
+}
+
+/** One mixed group using its own fixed projection geometry. */
+export interface MlxMixedGroupRange {
+  readonly from: number;
+  readonly to: number;
+  readonly operation: (group: MlxPhaseGroup) => MlxArray;
+}
+
+export interface MlxSeparateMixed {
+  /** Optional packing of multiple nonverify groups alongside separately routed verification. */
+  readonly packed?: readonly MlxMixedRange[];
+  readonly decode: readonly MlxMixedGroupRange[];
+  readonly window: readonly MlxMixedGroupRange[];
+  readonly verify: readonly MlxMixedGroupRange[];
+}

package/src/contracts/mlx/head.ts

@@ -0,0 +1,6 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+
+/** The output head over a post-final-norm hidden [.., hidden-width] (borrowed)
+ * to logits [.., vocabulary] the caller owns, possibly unevaluated. */
+export type MlxHeadForward = (hidden: MlxArray) => MlxArray;
+export interface MlxHeadRange { readonly from: number; readonly to: number; readonly operation: MlxHeadForward }

package/src/contracts/mlx/index.ts

@@ -0,0 +1,15 @@
+export * from "./attention";
+export * from "./cache";
+export * from "./checkpoint";
+export * from "./draft-checkpoint";
+export * from "./draft-projection";
+export * from "./draft-target";
+export * from "./drafter";
+export * from "./forward-work";
+export * from "./glm52-weights";
+export * from "./graph";
+export * from "./kv-codec";
+export * from "./media";
+export * from "./positions";
+export * from "./trainable";
+export * from "./whisper";

package/src/contracts/mlx/kv-codec.ts

@@ -0,0 +1,9 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+
+/** A backend KV representation independent of row storage and scheduling.
+ * Inputs are borrowed. Encoding and decoding return caller-owned arrays.
+ * A codec may retain values in a rotated domain for attention to invert once. */
+export interface KvCodec<Encoded> {
+  encode(keys: MlxArray, values: MlxArray): Encoded;
+  decode(encoded: Encoded, length: number, headDim: number, deferValues?: boolean): [MlxArray, MlxArray];
+}