@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/models/gemma4/graphs.ts

@@ -0,0 +1,147 @@
+import type { Weights } from "../../artifacts/weights";
+import type { GraphEntry, ServedComposition } from "../implementation";
+import * as e4b from "./gemma4-e4b";
+import * as twelve from "./gemma4-12b";
+import * as moe from "./gemma4-26b";
+
+export type Gemma4Graph = e4b.Gemma4E4b | twelve.Gemma4TwelveB | moe.Gemma4Moe;
+
+// Loader lookup keys. The graph files contain only their written operations.
+const e4bKey: ServedComposition = Object.freeze({
+  fingerprint: e4b.FINGERPRINT,
+  kv: Object.freeze([Object.freeze({
+    kind: "affine-layers" as const,
+    layers: Object.freeze([
+      { layer: 0, bits: 4, groupSize: 64 },
+      { layer: 1, bits: 4, groupSize: 64 },
+      { layer: 2, bits: 4, groupSize: 64 },
+      { layer: 3, bits: 8, groupSize: 64 },
+      { layer: 4, bits: 8, groupSize: 64 },
+      { layer: 5, bits: 4, groupSize: 64 },
+      { layer: 6, bits: 4, groupSize: 64 },
+      { layer: 7, bits: 4, groupSize: 64 },
+      { layer: 8, bits: 4, groupSize: 64 },
+      { layer: 9, bits: 8, groupSize: 64 },
+      { layer: 10, bits: 8, groupSize: 64 },
+      { layer: 11, bits: 8, groupSize: 64 },
+      { layer: 12, bits: 4, groupSize: 64 },
+      { layer: 13, bits: 4, groupSize: 64 },
+      { layer: 14, bits: 8, groupSize: 64 },
+      { layer: 15, bits: 4, groupSize: 64 },
+      { layer: 16, bits: 4, groupSize: 64 },
+      { layer: 17, bits: 4, groupSize: 64 },
+      { layer: 18, bits: 4, groupSize: 64 },
+      { layer: 19, bits: 4, groupSize: 64 },
+      { layer: 20, bits: 4, groupSize: 64 },
+      { layer: 21, bits: 4, groupSize: 64 },
+      { layer: 22, bits: 4, groupSize: 64 },
+      { layer: 23, bits: 4, groupSize: 64 },
+    ]),
+  })]),
+});
+
+const twelveKey: ServedComposition = Object.freeze({
+  fingerprint: twelve.FINGERPRINT,
+  kv: Object.freeze([Object.freeze({
+    kind: "affine-layers" as const,
+    layers: Object.freeze([
+      { layer: 0, bits: 4, groupSize: 64 },
+      { layer: 1, bits: 4, groupSize: 64 },
+      { layer: 2, bits: 4, groupSize: 64 },
+      { layer: 3, bits: 8, groupSize: 64 },
+      { layer: 4, bits: 4, groupSize: 64 },
+      { layer: 5, bits: 4, groupSize: 64 },
+      { layer: 6, bits: 8, groupSize: 64 },
+      { layer: 7, bits: 8, groupSize: 64 },
+      { layer: 8, bits: 8, groupSize: 64 },
+      { layer: 9, bits: 4, groupSize: 64 },
+      { layer: 10, bits: 8, groupSize: 64 },
+      { layer: 11, bits: 4, groupSize: 64 },
+      { layer: 12, bits: 8, groupSize: 64 },
+      { layer: 13, bits: 4, groupSize: 64 },
+      { layer: 14, bits: 4, groupSize: 64 },
+      { layer: 15, bits: 4, groupSize: 64 },
+      { layer: 16, bits: 4, groupSize: 64 },
+      { layer: 17, bits: 4, groupSize: 64 },
+      { layer: 18, bits: 4, groupSize: 64 },
+      { layer: 19, bits: 4, groupSize: 64 },
+      { layer: 20, bits: 4, groupSize: 64 },
+      { layer: 21, bits: 4, groupSize: 64 },
+      { layer: 22, bits: 4, groupSize: 64 },
+      { layer: 23, bits: 4, groupSize: 64 },
+      { layer: 24, bits: 4, groupSize: 64 },
+      { layer: 25, bits: 4, groupSize: 64 },
+      { layer: 26, bits: 4, groupSize: 64 },
+      { layer: 27, bits: 4, groupSize: 64 },
+      { layer: 28, bits: 4, groupSize: 64 },
+      { layer: 29, bits: 4, groupSize: 64 },
+      { layer: 30, bits: 4, groupSize: 64 },
+      { layer: 31, bits: 4, groupSize: 64 },
+      { layer: 32, bits: 4, groupSize: 64 },
+      { layer: 33, bits: 4, groupSize: 64 },
+      { layer: 34, bits: 4, groupSize: 64 },
+      { layer: 35, bits: 4, groupSize: 64 },
+      { layer: 36, bits: 4, groupSize: 64 },
+      { layer: 37, bits: 4, groupSize: 64 },
+      { layer: 38, bits: 4, groupSize: 64 },
+      { layer: 39, bits: 4, groupSize: 64 },
+      { layer: 40, bits: 4, groupSize: 64 },
+      { layer: 41, bits: 4, groupSize: 64 },
+      { layer: 42, bits: 4, groupSize: 64 },
+      { layer: 43, bits: 4, groupSize: 64 },
+      { layer: 44, bits: 4, groupSize: 64 },
+      { layer: 45, bits: 4, groupSize: 64 },
+      { layer: 46, bits: 4, groupSize: 64 },
+      { layer: 47, bits: 4, groupSize: 64 },
+    ]),
+  })]),
+});
+
+const moeKey: ServedComposition = Object.freeze({
+  fingerprint: moe.FINGERPRINT,
+  kv: Object.freeze([Object.freeze({
+    kind: "affine-layers" as const,
+    layers: Object.freeze([
+      { layer: 0, bits: 8, groupSize: 64 },
+      { layer: 1, bits: 4, groupSize: 64 },
+      { layer: 2, bits: 4, groupSize: 64 },
+      { layer: 3, bits: 8, groupSize: 64 },
+      { layer: 4, bits: 8, groupSize: 64 },
+      { layer: 5, bits: 8, groupSize: 64 },
+      { layer: 6, bits: 4, groupSize: 64 },
+      { layer: 7, bits: 4, groupSize: 64 },
+      { layer: 8, bits: 8, groupSize: 64 },
+      { layer: 9, bits: 8, groupSize: 64 },
+      { layer: 10, bits: 4, groupSize: 64 },
+      { layer: 11, bits: 8, groupSize: 64 },
+      { layer: 12, bits: 8, groupSize: 64 },
+      { layer: 13, bits: 4, groupSize: 64 },
+      { layer: 14, bits: 4, groupSize: 64 },
+      { layer: 15, bits: 4, groupSize: 64 },
+      { layer: 16, bits: 4, groupSize: 64 },
+      { layer: 17, bits: 4, groupSize: 64 },
+      { layer: 18, bits: 4, groupSize: 64 },
+      { layer: 19, bits: 4, groupSize: 64 },
+      { layer: 20, bits: 4, groupSize: 64 },
+      { layer: 21, bits: 4, groupSize: 64 },
+      { layer: 22, bits: 4, groupSize: 64 },
+      { layer: 23, bits: 4, groupSize: 64 },
+      { layer: 24, bits: 4, groupSize: 64 },
+      { layer: 25, bits: 4, groupSize: 64 },
+      { layer: 26, bits: 4, groupSize: 64 },
+      { layer: 27, bits: 4, groupSize: 64 },
+      { layer: 28, bits: 4, groupSize: 64 },
+      { layer: 29, bits: 4, groupSize: 64 },
+    ]),
+  })]),
+});
+
+export const GEMMA_GRAPHS: readonly GraphEntry<Weights, Gemma4Graph>[] = Object.freeze([
+  { id: "gemma4-e4b", serves: e4bKey, create: (weights, config) => new e4b.Gemma4E4b(weights, config) },
+  { id: "gemma4-12b", serves: twelveKey, create: (weights, config) => new twelve.Gemma4TwelveB(weights, config) },
+  { id: "gemma4-26b", serves: moeKey, create: (weights, config) => new moe.Gemma4Moe(weights, config) },
+]);
+export { Gemma4E4b } from "./gemma4-e4b";
+export { Gemma4TwelveB } from "./gemma4-12b";
+export { Gemma4Moe } from "./gemma4-26b";
+export * from "./written-blocks";

package/src/models/gemma4/media.ts

@@ -0,0 +1,88 @@
+import { rowOperationFor } from "../../runtime/row-operations";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache } from "../../contracts/mlx/cache";
+import type { TargetView } from "../../contracts/mlx/draft-target";
+import type { MediaEncoders, MediaSidecarProbes, MlxPromptInput, Vision, MlxEmbeddingsRange } from "../../contracts/mlx/media";
+import type { MlxHeadForward, MlxLayerCapture, MlxRowTable, MlxPromptRange } from "../../contracts/mlx/graph";
+import { bindEmbeddingsInput } from "../media-input";
+import { readAssistantDonors } from "./assistant-target";
+import type { ModelConfig } from "../../artifacts/config";
+import type { QuantizedEmbedding } from "../../layers/quantized-embedding";
+
+/** What the Gemma4 graph's media and draft ports read from it. */
+export interface Gemma4MediaGraph {
+  readonly config: ModelConfig;
+  readonly embed: QuantizedEmbedding;
+  readonly embedScale: number;
+  readonly numDonors: number;
+  readonly layers: readonly { readonly layerType: string }[];
+  /** The vision prefill phase: a prompt's prepared embeddings (borrowed), image tokens attending each other both
+   *  ways where `imageMask` marks them; `ids` feed the per-layer inputs. */
+  prefillEmbeddings?(embeddings: MlxArray, caches: Cache[], imageMask: MlxArray | null,
+    ids: MlxArray, multimodalMask: MlxArray | null): MlxArray;
+  /** A prompt's token ids after a reused prefix: the window read, or one position the decode read. */
+  readonly mediaPrefill?: readonly MlxEmbeddingsRange[];
+  readonly prompt: readonly MlxPromptRange[];
+  readonly decode: MlxRowTable<import("../../contracts/mlx/graph").MlxCapturingForward>;
+  readonly logits: MlxRowTable<MlxHeadForward>;
+}
+
+/** Prepared embeddings enter the vision prefill phase; a prompt continuing after a reused prefix and decode
+ *  continue from token ids, in the phase the scheduler names. */
+export function bindGemma4MediaInput(model: Gemma4MediaGraph, input: Vision): MlxPromptInput {
+  return bindEmbeddingsInput((ids, caches, start, phase, width) => start === 0
+    ? model.mediaPrefill
+      ? model.mediaPrefill.find(range => width >= range.from && width <= range.to)!.operation(input.embeddings, caches, input.imageMask ?? null, ids, input.multimodalMask ?? null)
+      : model.prefillEmbeddings!(input.embeddings, caches, input.imageMask ?? null, ids, input.multimodalMask ?? null)
+    : (phase === "decode" ? rowOperationFor(model.decode, width) : model.prompt.find(range => width >= range.from && width <= range.to)!.operation)(ids, caches, null) as MlxArray);
+}
+
+/** The graph's ports for draft sources over its live caches: donor KV for the
+ * assistant drafter and the layer taps for hidden-state drafters. */
+export function gemma4DraftTarget(model: Gemma4MediaGraph, caches: Cache[]): TargetView {
+  let sliding = -1, full = -1;
+  for (let i = 0; i < model.numDonors; i++) {
+    if (model.layers[i]!.layerType === "sliding_attention") sliding = i;
+    else full = i;
+  }
+  return Object.freeze({ identity: model, assistantRows: sliding >= 0 && full >= 0 ? {
+    hiddenSize: model.config.text.hiddenSize,
+    embed(ids: MlxArray) { using embedded = model.embed.encode(ids); return ops.mulScalar(embedded, model.embedScale); },
+    readDonors: () => readAssistantDonors(caches[sliding]!, caches[full]!),
+  } : undefined, hiddenLayerTaps: Object.freeze({
+    layerCount: model.layers.length,
+    projection: Object.freeze({ embed: model.embed, logits: model.logits }),
+  }) });
+}
+
+/** The tower a checkpoint's sidecar ships: the encoder-free gemma4_unified
+ * (12B) tower vs the SigLIP encoder (gemma4_vision: e2b/e4b/26B/31B), chosen by
+ * the sidecar's vision_config.model_type; and the Conformer audio tower when the
+ * sidecar header names its tensors (the local 12B pairs audio_config with a stub
+ * sidecar, and a non-null loader advertises audio on every capability surface).
+ * Towers load on first use, so text-only sessions never pay for them. */
+export async function gemma4MediaEncoders(
+  model: Gemma4MediaGraph, modelDir: string, probes: MediaSidecarProbes,
+): Promise<MediaEncoders> {
+  const config = model.config;
+  if (!config.hasVisionSidecar) return { vision: null, audio: null };
+  const vc = config.raw.vision_config as Record<string, any> | undefined;
+  let vision: MediaEncoders["vision"];
+  if (vc?.model_type === "gemma4_vision") {
+    const { SiglipVisionTower, parseSiglipConfig } = await import("../vision/siglip");
+    const sigCfg = parseSiglipConfig(vc);
+    vision = () => SiglipVisionTower.load(modelDir, sigCfg, model.embedScale);
+  } else {
+    // gemma4_unified_vision (or unlabelled): the encoder-free patch embedder.
+    const { VisionTower } = await import("../vision/unified");
+    vision = () => VisionTower.load(modelDir, model.embedScale, config.text.rmsNormEps);
+  }
+  let audio: MediaEncoders["audio"] = null;
+  if (config.raw.audio_config && probes.shipsAudioTower(`${modelDir}/optiq_vision.safetensors`)) {
+    const { AudioTower, parseAudioConfig } = await import("../audio/conformer");
+    const audioCfg = parseAudioConfig(config.raw.audio_config as Record<string, any>);
+    audio = () => AudioTower.load(modelDir, audioCfg, model.embedScale);
+  }
+  return { vision, audio };
+}