@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/models/cache-composition.ts

@@ -0,0 +1,44 @@
+import type { Cache } from "../contracts/mlx/cache";
+import type { MlxCacheLayer } from "../contracts/mlx/graph";
+import type { Composition } from "../contracts/portable/composition";
+import { AttentionMasks, foldedWindowRead, sdpaRead } from "../state/attention-read";
+import { Bf16FirstKv4Head256Cache, composedFullAttentionCache, composedSlidingAttentionCache } from "../state/bf16-first-kv";
+import { composedAffineKernels } from "../state/affine-attention";
+import { SSMCache } from "../state/ssm";
+import { composedLatentCache, latentLegoCodecs } from "../state/latent-legos";
+import { maybePageKv } from "../state/request-policy";
+import { disposeResources } from "../runtime/resources";
+import { denseTurboQuantKernels, deferredTurboQuantKernels } from "../state/turboquant-attention";
+
+/** The loader composes a request's cache legos from the file's declaration. */
+export function composeGraphCaches(graph: { readonly cacheLayers: readonly MlxCacheLayer[]; readonly cachePaging?: boolean },
+  composition: {readonly kv: Composition["kv"]; readonly pagedKv: {readonly blockSize: number; readonly direct?: boolean} | null}): Cache[] {
+  const masks = new AttentionMasks();
+  const { kv, pagedKv } = composition;
+  const caches = graph.cacheLayers.map((layer, index) => {
+    switch (layer.kind) {
+      case "recurrent": return new SSMCache();
+      case "sliding": return composedSlidingAttentionCache(kv, index, layer.window, masks);
+      case "latent": return composedLatentCache(kv, pagedKv, index, layer.geometry);
+      case "full":
+        if (layer.kv4Head256 && kv.kind === "affine" && kv.bits === 4 && kv.groupSize === 64)
+          return new Bf16FirstKv4Head256Cache(kv.quantizedKvStart ?? 0, composedAffineKernels(kv.kernels), masks, layer.kv4DecodeBlock);
+        return composedFullAttentionCache(kv, index, masks,
+          layer.turboQuant === "deferred" ? deferredTurboQuantKernels : denseTurboQuantKernels,
+          layer.windowRead === "folded" ? foldedWindowRead : sdpaRead);
+    }
+  });
+  if (graph.cachePaging) {
+    try { maybePageKv(caches, kv, pagedKv ? {...pagedKv,direct:pagedKv.direct??false} : null, pagedKv?.blockSize ?? 0); }
+    catch (error) { disposeResources(caches); throw error; }
+  }
+  return caches;
+}
+
+/** Bind request-state construction at load; the graph only declares layout. */
+export function bindGraphCaches<Graph extends { readonly cacheLayers: readonly MlxCacheLayer[]; readonly cachePaging?: boolean }>(graph: Graph,
+  composition: {readonly kv: Composition["kv"]; readonly pagedKv: {readonly blockSize: number; readonly direct?: boolean} | null}): Graph & { makeCache(): Cache[]; readonly stateCodecs?: import("../state/persistence-types").CacheCodecProvider } {
+  return Object.assign(graph, { makeCache: () => composeGraphCaches(graph, composition),
+    ...(graph.cacheLayers.some(layer => layer.kind === "latent") && (composition.kv.kind !== "bf16" || composition.pagedKv)
+      ? { stateCodecs: latentLegoCodecs } : {}) });
+}

package/src/models/capabilities.ts

@@ -0,0 +1,89 @@
+import type { GraphCapabilities, GraphDeclaration } from "../contracts/portable/graph";
+import type { MlxCompiledDecodeStep, MlxDeclaredGraph, MlxTokenPhases } from "../contracts/mlx/graph";
+import type { Composition } from "../contracts/portable/composition";
+import type { CacheCodecProvider } from "../state/persistence-types";
+
+type Declaration = Partial<Omit<GraphDeclaration, "prefill" | "speculation">> & {
+  readonly prefill?: Partial<GraphCapabilities["prefill"]>;
+  readonly speculation?: Partial<GraphCapabilities["speculation"]>;
+};
+
+/** Build a graph's declaration from the defaults an ordinary autoregressive
+ * graph has (batched and mountable adapters, verification qualified for every
+ * request shape, no media, no compiled step, unbounded prefill chunks)
+ * and what the graph does differently. */
+export function declareGraph(declared: Declaration = {}): GraphDeclaration {
+  if ("embeddings" in declared)
+    throw new TypeError("pooled embedding support is declared by embedPooled, not an embeddings flag");
+  return Object.freeze({
+    method: "autoregressive", media: null, compiledDecode: false,
+    sparseAttention: false, hiddenLayerTaps: false, nativeDraft: null,
+    ...declared,
+    prefill: Object.freeze({ boundedWorkspace: false, ...declared.prefill }),
+    speculation: Object.freeze({ adapters: true, logprobs: true, externalTokens: true, grammarProposals: true,
+      ...declared.speculation }),
+  });
+}
+
+/** Verification qualified for ordinary rounds only: no adapters, logprobs,
+ * supplied external tokens or verified grammar proposals (a graph whose
+ * attention softcaps its scores). */
+export const ORDINARY_VERIFICATION = Object.freeze({
+  adapters: false, logprobs: false, externalTokens: false, grammarProposals: false,
+});
+
+/** A graph's declaration with the operations it promises, checked once where
+ * composition binds it. A graph that declares nothing is refused, never read as
+ * "no capabilities". */
+export function declaredGraph(model: object): MlxDeclaredGraph {
+  const graph = model as Partial<MlxDeclaredGraph>;
+  const capabilities = graph.graphCapabilities;
+  if (!capabilities) throw new TypeError("the graph must declare its capabilities (graphCapabilities)");
+  if ("embeddings" in capabilities)
+    throw new TypeError("pooled embedding support is declared by embedPooled, not an embeddings flag");
+  if (graph.embedPooled !== undefined && typeof graph.embedPooled !== "function")
+    throw new TypeError("the graph's embedPooled operation must be a function");
+  const media = capabilities.media;
+  if (media && media.input !== "pixels" && typeof graph.bindMediaInput !== "function")
+    throw new TypeError(`the graph declares ${media.input} media input but provides no bindMediaInput`);
+  if (media?.input === "pixels" && typeof graph.pixelInput !== "function")
+    throw new TypeError("the graph declares pixel media input but provides no pixelInput");
+  return graph as MlxDeclaredGraph;
+}
+
+/** Describe a validated graph for portable consumers. Operation presence is
+ * the single source of pooled embedding support; retain the operation on the
+ * graph so invocation preserves its receiver and resource ownership. */
+export function graphCapabilitiesOf(graph: MlxDeclaredGraph): GraphCapabilities {
+  return Object.freeze({ ...graph.graphCapabilities, embeddings: typeof graph.embedPooled === "function" });
+}
+
+/** The named token phases of a graph whose method is autoregressive, checked
+ * once where composition binds it; null for a denoising graph, whose canvas
+ * method runs none. An autoregressive graph without them is refused. */
+export function tokenPhasesOf(model: object): MlxTokenPhases & { readonly composition?: Composition } | null {
+  const graph = declaredGraph(model);
+  if (graph.graphCapabilities.method !== "autoregressive") return null;
+  const phases = model as Partial<MlxTokenPhases>;
+  for (const name of ["prompt", "multiSequencePrompt", "decode", "verify", "appendKnown", "logits"] as const)
+    if (!phases[name]) throw new TypeError(`the graph declares the autoregressive method but provides no ${name} phase`);
+  return model as MlxTokenPhases & { readonly composition?: Composition };
+}
+
+/** The compiled decode step of a graph that declares one, created on first use
+ * and kept by the graph; null for every other graph. A graph that promises a
+ * step and provides none is refused where composition binds it. */
+export function compiledDecodeStepOf(model: object): MlxCompiledDecodeStep | null {
+  const graph = declaredGraph(model);
+  if (!graph.graphCapabilities.compiledDecode) return null;
+  if (typeof graph.compiledDecodeStep !== "function")
+    throw new TypeError("the graph declares a compiled decode step but provides no compiledDecodeStep");
+  return graph.compiledDecodeStep();
+}
+
+/** The persistence codecs a graph's caches need beyond the default set (the
+ * graph's `stateCodecs`: GLM-5.2's under a quantized KV scheme or paging, whose
+ * latent legos the default codecs do not save); undefined: the default set. */
+export function stateCodecsOf(model: object): CacheCodecProvider | undefined {
+  return (model as { readonly stateCodecs?: CacheCodecProvider }).stateCodecs;
+}

package/src/models/chat-template.ts

@@ -0,0 +1,10 @@
+import { ChatTemplate } from "../input/chat-template";
+import { chatTemplateFallbackFor, type ResolvedModelProfile } from "./profile";
+
+/** Load the chat template of a resolved model: the artifact's own, else the
+ * renderer its profile declares for artifacts that ship none. */
+export function loadModelChatTemplate(
+  modelDir: string, resolved: ResolvedModelProfile, options: { disableThinking?: boolean } = {},
+): Promise<ChatTemplate> {
+  return ChatTemplate.load(modelDir, { ...options, fallback: chatTemplateFallbackFor(resolved) ?? undefined });
+}

package/src/models/composition.ts

@@ -0,0 +1,107 @@
+// The one resolver of a composition record. The application calls it once per
+// load, before the graph is built; the convenience load API calls it for
+// callers that supply no record.
+import { deviceArchitecture } from "@mlx-bun/mlx/ffi";
+import type { ModelConfig, TurboQuantScheme } from "../artifacts/config";
+import { COMPOSITION_DEFAULTS, type Composition } from "../contracts/portable/composition";
+import { runtimeConfig, type RuntimeConfig } from "../runtime/config";
+import { resolveKvScheme, type KvQuantOverride } from "../state/kv-scheme";
+import { PagedKVCache } from "../state/paged/cache";
+
+/** Rows per forward when the caller states no cap (the app's `--batch` default). */
+export const DEFAULT_MAX_ROWS = 8;
+/** The prefill chunk when `MLX_BUN_RD_PREFILL_CHUNK` is unset, the prefill
+ * policy's ceiling (`generation/bindings/prefill-policy.ts`). */
+export const PREFILL_CHUNK_TOKENS = 2048;
+
+/** The chunk the prefill policy starts from: its `MLX_BUN_RD_PREFILL_CHUNK`
+ * override, read through the runtime configuration exactly as it reads it
+ * (a set, non-numeric value counts as 2048), else {@link PREFILL_CHUNK_TOKENS}. */
+function resolveChunkTokens(runtime: RuntimeConfig): number {
+  const explicit = runtime.value("MLX_BUN_RD_PREFILL_CHUNK") === undefined
+    ? undefined : runtime.number("MLX_BUN_RD_PREFILL_CHUNK", 2048);
+  return explicit === undefined ? PREFILL_CHUNK_TOKENS : Math.max(1, Math.floor(explicit));
+}
+
+/** What the caller decided; each absent value takes the default named on it. */
+export interface CompositionRequest {
+  /** The KV flags: `--kv-quant` (`override`), its TurboQuant spelling,
+   * `--quantized-kv-start` and `--fused-sdpa` (the affine read kernels: on,
+   * tiled-causal; off, unfused; absent, by the scheme's kind). Absent: bf16.
+   * A `config` override reads the artifact's `kv_config.json` from the model
+   * config. */
+  readonly kv?: {
+    readonly override?: KvQuantOverride;
+    readonly turboQuant?: TurboQuantScheme;
+    readonly quantizedKvStart?: number;
+    readonly fusedSdpa?: boolean;
+  };
+  /** Paged KV (`--paged-kv`): full-attention storage in pages of `blockSize`
+   * positions (default 256), read by the direct kernel where `direct` (default
+   * false; see `PagedKvComposition`). Absent or null: unpaged. Pages store what
+   * the KV scheme names for each layer. */
+  readonly pagedKv?: { readonly blockSize?: number; readonly direct?: boolean } | null;
+  /** The loaded drafter's tokens per round, and whether the scheduler adapts
+   * each round's count up to it. Absent or null: no drafter. */
+  readonly draft?: { readonly numDraftTokens: number; readonly adaptive?: boolean } | null;
+  /** The most rows one forward carries. Default {@link DEFAULT_MAX_ROWS}. */
+  readonly maxRows?: number;
+  /** One argmax over a decode step whose live rows are all plain greedy. Default true. */
+  readonly vectorGreedySampling?: boolean;
+  /** One forward carries a decode group and a prefill group. Default false. */
+  readonly mixedPrefill?: boolean;
+  /** Prompt tokens a mixed forward carries beside its decode rows. Default 256. */
+  readonly mixedTokenBudget?: number;
+  /** A prompt's last token prefills alone, as the first one-position step. Default true. */
+  readonly prefillTailSplit?: boolean;
+  /** Grammar-forced continuations enter the KV without one forward per token. Default false. */
+  readonly grammarJump?: boolean;
+  /** Speculative rows draft by argmax. Default false. */
+  readonly greedyDraft?: boolean;
+}
+
+/** Resolve the record for `config` on this machine. Reads the GPU architecture
+ * and the prefill chunk override from the current runtime configuration. The
+ * result is frozen. */
+export function resolveComposition(config: ModelConfig, request: CompositionRequest = {}): Composition {
+  return compose(request, resolveKvScheme({ ...request.kv, config: config.kvQuant }).storage);
+}
+
+/** The record a load resolves on this machine without an artifact's config:
+ * `resolveComposition` with bf16 KV, which is what a load with no KV flag
+ * resolves. The install step compiles the Neural Engine programs of this record
+ * (`compileAnePrograms`), so it needs no downloaded model. Frozen. */
+export function machineComposition(request: Omit<CompositionRequest, "kv"> = {}): Composition {
+  return compose(request, resolveKvScheme({}).storage);
+}
+
+function compose(request: Omit<CompositionRequest, "kv">, kv: Composition["kv"]): Composition {
+  const maxRows = request.maxRows ?? DEFAULT_MAX_ROWS;
+  if (!Number.isSafeInteger(maxRows) || maxRows < 1) throw new Error("composition maxRows must be a positive integer");
+  const draft = request.draft;
+  if (draft && (!Number.isSafeInteger(draft.numDraftTokens) || draft.numDraftTokens < 1))
+    throw new Error("composition draft tokens must be a positive integer");
+  const mixedTokenBudget = request.mixedTokenBudget ?? COMPOSITION_DEFAULTS.mixedTokenBudget;
+  if (!Number.isSafeInteger(mixedTokenBudget) || mixedTokenBudget < 1)
+    throw new Error("composition mixedTokenBudget must be a positive integer");
+  const pagedKv = request.pagedKv ? Object.freeze({ blockSize: request.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE,
+    direct: request.pagedKv.direct ?? false }) : null;
+  if (pagedKv && (!Number.isSafeInteger(pagedKv.blockSize) || pagedKv.blockSize < 1))
+    throw new Error("composition paged KV block size must be a positive integer");
+  const runtime = runtimeConfig();
+  return Object.freeze({
+    device: deviceArchitecture(),
+    kv,
+    pagedKv,
+    draftDepth: !draft ? 0 : draft.adaptive ? Object.freeze({ adaptive: true as const, max: draft.numDraftTokens }) : draft.numDraftTokens,
+    prefillChunkTokens: resolveChunkTokens(runtime),
+    maxRows,
+    vectorGreedySampling: request.vectorGreedySampling ?? COMPOSITION_DEFAULTS.vectorGreedySampling,
+    // The caches' declaration: direct pages have no compiled form.
+    mixedPrefill: request.mixedPrefill ?? COMPOSITION_DEFAULTS.mixedPrefill,
+    mixedTokenBudget,
+    prefillTailSplit: request.prefillTailSplit ?? COMPOSITION_DEFAULTS.prefillTailSplit,
+    grammarJump: request.grammarJump ?? COMPOSITION_DEFAULTS.grammarJump,
+    greedyDraft: request.greedyDraft ?? COMPOSITION_DEFAULTS.greedyDraft,
+  });
+}

package/src/models/diffusion-gemma/experts.ts

@@ -0,0 +1,59 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { CompiledGeglu } from "../../layers/geglu";
+import type { QuantizedSwitchLinear } from "../../layers/quantized-switch-linear";
+
+export interface DiffExpertWeights {
+  readonly gateUp: QuantizedSwitchLinear;
+  readonly down: QuantizedSwitchLinear;
+  readonly hiddenDims: number;
+}
+
+/** Fused gate/up experts in their original token order. */
+export class UnsortedDiffExperts {
+  readonly activation = new CompiledGeglu();
+  constructor(readonly base: DiffExpertWeights) {}
+  forward(x: MlxArray, indices: MlxArray, weights: MlxArray): MlxArray {
+    using expanded = ops.expandDims(x,-2);
+    using h = ops.expandDims(expanded,-3);
+    using gateUp = this.base.gateUp.forward(h,indices,false);
+    const [gate,up] = ops.split(gateUp,[this.base.hiddenDims],-1);
+    using mid = this.activation.forward(gate!,up!);
+    using projected = this.base.down.forward(mid,indices,false);
+    const shape = projected.shape; shape.splice(shape.length-2,1);
+    using y = ops.reshape(projected,shape);
+    using w = ops.expandDims(weights,-1);
+    using weighted = ops.mul(y,w);
+    return ops.sumAxis(weighted,-2,false);
+  }
+}
+
+/** Expert-sorted rows and inverse permutation, bound by the graph. */
+export class SortedDiffExperts {
+  readonly activation = new CompiledGeglu();
+  constructor(readonly base: DiffExpertWeights) {}
+  forward(x: MlxArray, indices: MlxArray, weights: MlxArray): MlxArray {
+    using expanded = ops.expandDims(x,-2);
+    using h = ops.expandDims(expanded,-3);
+    const k = indices.shape.at(-1)!;
+    using idxFlat = ops.reshape(indices,[indices.size]);
+    using order = ops.argsortAxis(idxFlat,0);
+    using inverse = ops.argsortAxis(order,0);
+    using kScalar = ops.scalarLike(k,order);
+    using rows = ops.floorDivide(order,kScalar);
+    using flat = ops.reshape(h,[-1,1,h.shape.at(-1)!]);
+    using sorted = ops.takeAxis(flat,rows,0);
+    using idx = ops.takeAxis(idxFlat,order,0);
+    using gateUp = this.base.gateUp.forward(sorted,idx,true);
+    const [gate,up] = ops.split(gateUp,[this.base.hiddenDims],-1);
+    using mid = this.activation.forward(gate!,up!);
+    using projected = this.base.down.forward(mid,idx,true);
+    using restored = ops.takeAxis(projected,inverse,0);
+    using unsorted = ops.reshape(restored,[...indices.shape,1,restored.shape.at(-1)!]);
+    const shape = unsorted.shape; shape.splice(shape.length-2,1);
+    using y = ops.reshape(unsorted,shape);
+    using w = ops.expandDims(weights,-1);
+    using weighted = ops.mul(y,w);
+    return ops.sumAxis(weighted,-2,false);
+  }
+}

package/src/models/diffusion-gemma/family.ts

@@ -0,0 +1,21 @@
+import { L2, configuredMoe, defineFamily } from "../family";
+import { GEMMA_MEDIA_TOKENS, GEMMA_TRAINING_DEFAULTS } from "../gemma4/family";
+
+/** DiffusionGemma (model_type `diffusion_gemma`): block/masked-diffusion canvas
+ * model. Non-autoregressive, routed through the diffusion engine, not the AR loop. */
+export const DIFFUSION_GEMMA_FAMILY = defineFamily({
+  graph: "diffusion-gemma",
+  accepts: config => config.modelType === "diffusion_gemma",
+  hasModelType: modelType => modelType === "diffusion_gemma",
+  tier: "targeted",
+  label: "diffusion_gemma",
+  loader: "safetensors",
+  loop: "diffusion",
+  fidelity: L2,
+  specialization: "dedicated",
+  profileId: "diffusion-gemma-dedicated",
+  capabilities: ["diffusion-gemma-graph"],
+  moe: configuredMoe,
+  trainingDefaults: GEMMA_TRAINING_DEFAULTS,
+  mediaTokens: GEMMA_MEDIA_TOKENS,
+});