mlx-bun 0.0.7 → 0.0.8

package/src/model/factory.ts

@@ -11,11 +11,17 @@ import { configFingerprint } from "./fingerprint";
 import { GENERATED } from "./generated";
 import { MiniCPM5Model } from "./minicpm5";
 import { Qwen35Model } from "./qwen3_5";
-import { isMiniCPM5Config, isQwen35Config } from "./support";
+import { DiffusionGemmaModel } from "./diffusion-gemma";
+import { isDiffusionGemmaConfig, isMiniCPM5Config, isQwen35Config } from "./support";
 
-export type RuntimeModel = Gemma4Model | MiniCPM5Model | Qwen35Model;
+export type RuntimeModel = Gemma4Model | MiniCPM5Model | Qwen35Model | DiffusionGemmaModel;
 
 export function createModel(weights: Weights, config: ModelConfig): RuntimeModel {
+  // DiffusionGemma is non-autoregressive — generate() detects it and routes to
+  // the denoising engine instead of the AR decode loop. It exposes the shared
+  // RuntimeModel surface (config/weightsBytes/loraState/makeCache) with the
+  // AR-only methods as throwing stubs (never called on this model).
+  if (isDiffusionGemmaConfig(config)) return new DiffusionGemmaModel(weights, config);
   if (isMiniCPM5Config(config)) return new MiniCPM5Model(weights, config);
   if (isQwen35Config(config)) return new Qwen35Model(weights, config);
   if (config.modelType === "llama")

package/src/model/support.ts

@@ -31,9 +31,17 @@ export function isDrafterModelType(modelType: string): boolean {
   return modelType.endsWith("_assistant");
 }
 
+/** DiffusionGemma (model_type `diffusion_gemma`): block/masked-diffusion canvas
+ *  model. Non-autoregressive — routed through the diffusion engine, not the AR
+ *  loop. See docs/design/diffusion-gemma-port.md. */
+export function isDiffusionGemmaConfig(config: ModelConfig): boolean {
+  return config.modelType === "diffusion_gemma";
+}
+
 export function isSupportedModelRecord(modelType: string, repoId = ""): boolean {
   if (isDrafterModelType(modelType)) return false;
   if (modelType.startsWith("gemma4")) return true;
+  if (modelType === "diffusion_gemma") return true;
   // qwen3_5 / qwen3_5_text (dense hybrid). The MoE variant is a separate type.
   if (modelType === "qwen3_5" || modelType === "qwen3_5_text") return true;
   return modelType === "llama" && repoId.toLowerCase().includes("minicpm5-1b-optiq-4bit");
@@ -42,5 +50,5 @@ export function isSupportedModelRecord(modelType: string, repoId = ""): boolean
 export function isSupportedModelConfig(config: ModelConfig): boolean {
   if (isDrafterModelType(config.modelType)) return false;
   return config.modelType.startsWith("gemma4") || isMiniCPM5Config(config) ||
-    isQwen35Config(config);
+    isQwen35Config(config) || isDiffusionGemmaConfig(config);
 }

package/src/serve/generation-gateway.ts

@@ -27,6 +27,7 @@
 
 import { MlxArray } from "../mlx/array";
 import type { RuntimeModel } from "../model/factory";
+import { DiffusionGemmaModel } from "../model/diffusion-gemma";
 import type { GenerateOptions, GenerateStats } from "../generate";
 import { makeSampler, toLogprobs } from "../sampler";
 import { BatchScheduler } from "./batch-scheduler";
@@ -93,6 +94,9 @@ export class GenerationGateway {
 
   /** Decide whether a request joins the batch or runs serially. */
   willBatch(shape: RequestShape): boolean {
+    // DiffusionGemma is non-autoregressive — the batch scheduler assumes the AR
+    // KV-cache decode path, so it always runs serially through generate().
+    if (this.model instanceof DiffusionGemmaModel) return false;
     return (
       this.batchingEnabled &&
       !shape.hasVision &&

package/src/server.ts

@@ -23,6 +23,8 @@ const pkgVersion = (pkgJson as { version: string }).version;
 import { loadModelConfig, type KvQuantSpec, type ModelConfig } from "./config";
 import { Weights } from "./weights";
 import { Gemma4Model } from "./model/gemma4";
+import { DiffusionGemmaModel } from "./model/diffusion-gemma";
+import { spliceImageTokens } from "./vision/diffusion-vision";
 import { createModel, type RuntimeModel } from "./model/factory";
 import { isMiniCPM5Config, isQwen35Config, isSupportedModelRecord } from "./model/support";
 import { generate, type GenerateOptions } from "./generate";
@@ -839,6 +841,7 @@ export function createServer(
     } finally {
       vision?.embeddings.dispose();
       vision?.imageMask.dispose();
+      options.visionPixels?.dispose();
     }
   };
 
@@ -1402,8 +1405,33 @@ export function createServer(
         // Vision is Gemma4-only (no switchable thinking channel), so it's false.
         let startInThinking = false;
         let vision: Parameters<typeof runGeneration>[3];
+        let diffusionPixels: import("./mlx/array").MlxArray | null = null;
         try {
-          if (hasImages) {
+          if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
+            // DiffusionGemma image-text-to-text: its OWN dedicated SigLIP tower +
+            // encoder vision merge feed the denoising engine (NOT the AR
+            // forwardEmbeddings path). v1 supports a single image.
+            const dm = ctx.model;
+            if (!dm.visionTower)
+              return Response.json(
+                { error: { message: "this checkpoint has no vision tower" } }, { status: 400 },
+              );
+            const { messages, images } = await extractImages(normalizeMessages(body.messages));
+            if (images.length !== 1)
+              return Response.json(
+                { error: { message: "DiffusionGemma image input supports exactly one image" } },
+                { status: 400 },
+              );
+            const rendered = ctx.template.render(messages, { tools, addGenerationPrompt: true });
+            const rawIds = ctx.tokenizer.encode(rendered, /* addSpecialTokens */ false);
+            const { pixels, softTokens } = await dm.visionTower.preprocess(images[0]!);
+            promptIds = spliceImageTokens(rawIds, [softTokens], {
+              image: ctx.visionTokenIds.imageTokenId,
+              boi: ctx.visionTokenIds.boiTokenId,
+              eoi: ctx.visionTokenIds.eoiTokenId,
+            });
+            diffusionPixels = pixels;
+          } else if (hasImages) {
             // Loads (and caches) the tower on first image request — text-only
             // sessions never pay for it.
             const tower = getVisionTower(ctx);
@@ -1432,6 +1460,7 @@ export function createServer(
           );
         }
         const options = toOptions(body);
+        if (diffusionPixels) options.visionPixels = diffusionPixels;
         // Admission: reject what cannot finish within the memory budget
         // (the GPU OOM it would otherwise hit is uncatchable and kills
         // the process — Phase 6 finding).
@@ -1439,6 +1468,7 @@ export function createServer(
         if (requiredCtx > admission.maxSafeContext) {
           vision?.embeddings.dispose();
           vision?.imageMask.dispose();
+          diffusionPixels?.dispose();
           return Response.json(
             {
               error: {

package/src/train/diffusion-lora.ts

@@ -0,0 +1,216 @@
+// DiffusionGemma LoRA fine-tuning (port D5). DiffusionGemma is a discrete
+// block-diffusion LM, so the AR cross-entropy trainer doesn't apply — train LoRA
+// with the model's NATIVE denoising objective (port of optiq
+// vlm/diffusion_gemma/lora.py `diffusion_loss` + `train_diffusion_lora`):
+//   corrupt the target canvas to a random noise level t∈[t_min,t_max] (replace
+//   each token with prob t by a uniform-random token — the same corruption the
+//   inference canvas starts from), run one encoder+decoder forward, and minimise
+//   cross-entropy on the CORRUPTED positions (predict the clean token).
+//
+// LoRA is injected on the decoder blocks (DEFAULT_LORA_KEYS = attn q/k/v/o +
+// dense MLP gate/up/down). The encoder reuses those same blocks (weight-tied,
+// only per-layer scalars differ), so one injection trains both paths. Experts/
+// router stay frozen (their routing indices are stop_gradient'd so the MoE
+// backward differentiates only the activations). Adapter saved in the
+// AdapterManager.mount layout (lora_a/lora_b), mountable for inference.
+
+import * as ops from "../mlx/ops";
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import { ValueAndGrad } from "../mlx/autograd";
+import { DiffusionGemmaModel } from "../model/diffusion-gemma";
+import { AdamW } from "./optimizer";
+import {
+  attachForTraining,
+  buildTrainableLora,
+  detachTraining,
+  flatParams,
+  type TrainableLora,
+} from "./lora-params";
+
+export const DEFAULT_LORA_KEYS = [
+  "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj",
+  "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj",
+];
+
+export interface DiffusionPair {
+  promptIds: number[];
+  targetIds: number[];
+}
+
+export interface DiffusionLoraConfig {
+  rank?: number; // default 8
+  scale?: number; // default 8.0 (diffusion is scale-sensitive; AR's 20 collapses)
+  iters?: number; // default 200
+  learningRate?: number; // default 1e-4
+  tMin?: number; // default 0.4
+  tMax?: number; // default 0.8
+  seed?: number; // default 0
+  reportEvery?: number; // default 10
+  onReport?: (iter: number, avgLoss: number) => void;
+}
+
+/** swap/restore the differentiated primals into the LoRA leaves (port of the
+ *  trainer.ts private helpers — primals order is [...A, ...B]). */
+function swapPrimals(lora: TrainableLora, primals: MlxArray[]): MlxArray[] {
+  const n = lora.targets.length;
+  const saved: MlxArray[] = [];
+  for (let i = 0; i < n; i++) {
+    saved.push(lora.targets[i]!.lw.a, lora.targets[i]!.lw.b);
+    lora.targets[i]!.lw.a = primals[i]!;
+    lora.targets[i]!.lw.b = primals[n + i]!;
+  }
+  return saved;
+}
+function restorePrimals(lora: TrainableLora, saved: MlxArray[]): void {
+  for (let i = 0; i < lora.targets.length; i++) {
+    lora.targets[i]!.lw.a = saved[2 * i]!;
+    lora.targets[i]!.lw.b = saved[2 * i + 1]!;
+  }
+}
+function writeParam(lora: TrainableLora, i: number, p: MlxArray): void {
+  const n = lora.targets.length;
+  if (i < n) lora.targets[i]!.lw.a = p;
+  else lora.targets[i - n]!.lw.b = p;
+}
+
+/** sum all elements -> scalar []. */
+function sumAll(a: MlxArray): MlxArray {
+  const flat = ops.reshape(a, [a.size]);
+  const s = ops.sumAxis(flat, 0, false);
+  flat.dispose();
+  return s;
+}
+
+/** Denoising cross-entropy on corrupted canvas positions (the reference
+ *  diffusion_loss). Returns a scalar loss MlxArray, differentiable w.r.t. the
+ *  mounted LoRA. */
+export function diffusionLoss(
+  model: DiffusionGemmaModel,
+  promptIds: number[],
+  targetIds: number[],
+  tMin: number,
+  tMax: number,
+  vocab: number,
+): MlxArray {
+  const L = targetIds.length;
+  // t ~ U(t_min, t_max) scalar; corrupt[i] = U(0,1) < t ; noise ~ randint
+  const t = ops.randomUniform([1], Dtype.float32, tMin, tMax, null);
+  const u = ops.randomUniform([1, L], Dtype.float32, 0, 1, null);
+  const corrupt = ops.less(u, t); // bool [1, L]
+  u.dispose();
+  t.dispose();
+  const noise = ops.randint(0, vocab, [1, L], Dtype.int32);
+  const target = MlxArray.fromInt32(Int32Array.from(targetIds), [1, L]);
+  const canvas = ops.where(corrupt, noise, target);
+  noise.dispose();
+
+  const logits = model.forwardCanvasLogitsArr(promptIds, canvas); // [1, L, vocab] f32
+  canvas.dispose();
+
+  // cross_entropy(reduction=none) = logsumexp(logits) - logit_at_target
+  const lse = ops.logsumexpAxis(logits, -1, false); // [1, L]
+  const tExp = ops.expandDims(target, -1); // [1, L, 1]
+  const gathered = ops.takeAlongAxis(logits, tExp, -1); // [1, L, 1]
+  tExp.dispose();
+  logits.dispose();
+  const tgt = ops.reshape(gathered, [1, L]);
+  gathered.dispose();
+  const ce = ops.sub(lse, tgt); // [1, L]
+  lse.dispose();
+  tgt.dispose();
+  target.dispose();
+
+  const corruptF = corrupt.astype(Dtype.float32);
+  corrupt.dispose();
+  const masked = ops.mul(ce, corruptF);
+  ce.dispose();
+  const num = sumAll(masked);
+  masked.dispose();
+  const cnt = sumAll(corruptF);
+  corruptF.dispose();
+  const one = ops.scalarLike(1, cnt);
+  const denom = ops.maximum(cnt, one);
+  cnt.dispose();
+  one.dispose();
+  const loss = ops.div(num, denom);
+  num.dispose();
+  denom.dispose();
+  return loss;
+}
+
+/** Train a LoRA adapter on DiffusionGemma with the denoising objective. Returns
+ *  the TrainableLora (caller saves it via saveAdapter) + the loss history. */
+export function trainDiffusionLora(
+  model: DiffusionGemmaModel,
+  pairs: DiffusionPair[],
+  cfg: DiffusionLoraConfig = {},
+): { lora: TrainableLora; losses: number[] } {
+  const rank = cfg.rank ?? 8;
+  const scale = cfg.scale ?? 8.0;
+  const iters = cfg.iters ?? 200;
+  const lr = cfg.learningRate ?? 1e-4;
+  const tMin = cfg.tMin ?? 0.4;
+  const tMax = cfg.tMax ?? 0.8;
+  const seed = cfg.seed ?? 0;
+  const reportEvery = cfg.reportEvery ?? 10;
+  const vocab = model.config.text.vocabSize;
+  if (pairs.length === 0) throw new Error("no training pairs");
+
+  ops.randomSeed(BigInt(seed >>> 0));
+
+  // LoRA on every decoder block's attn + dense MLP (skips full-layer v_proj
+  // automatically — buildTrainableLora filters to existing loraTargets).
+  const ranks = new Map<string, number>();
+  for (let i = 0; i < model.layers.length; i++)
+    for (const key of DEFAULT_LORA_KEYS) ranks.set(`model.decoder.layers.${i}.${key}`, rank);
+  const lora = buildTrainableLora(model, ranks, scale, seed);
+  attachForTraining(model, lora, "train");
+  model.loraState.active = ["train"];
+
+  const params = flatParams(lora);
+  const opt = new AdamW(
+    params,
+    { lr, betas: [0.9, 0.999], eps: 1e-8, weightDecay: 0.01 },
+    (i, p) => writeParam(lora, i, p),
+  );
+
+  let current: DiffusionPair = pairs[0]!;
+  const vag = new ValueAndGrad((primals) => {
+    const saved = swapPrimals(lora, primals);
+    try {
+      return diffusionLoss(model, current.promptIds, current.targetIds, tMin, tMax, vocab);
+    } finally {
+      restorePrimals(lora, saved);
+    }
+  }, params.map((_, i) => i));
+
+  const losses: number[] = [];
+  try {
+    for (let it = 0; it < iters; it++) {
+      current = pairs[it % pairs.length]!;
+      if (current.targetIds.length === 0) continue;
+      const { value, grads } = vag.apply(flatParams(lora));
+      opt.step(grads); // takes ownership of grads + writes via writeParam
+      ops.evalAll([...flatParams(lora), value]);
+      const lossNum = value.astype(Dtype.float32).toFloat32()[0]!;
+      value.dispose();
+      losses.push(lossNum);
+      if ((it + 1) % reportEvery === 0) {
+        const window = losses.slice(-reportEvery);
+        cfg.onReport?.(it + 1, window.reduce((a, b) => a + b, 0) / window.length);
+      }
+    }
+  } catch (e) {
+    // On failure, undo the mount so the model is left clean.
+    detachTraining(model, lora);
+    model.loraState.active = [];
+    vag.dispose();
+    throw e;
+  }
+  // On success leave the trained LoRA ATTACHED + ACTIVE so the model is
+  // immediately usable; the caller saves it (saveAdapter) and may detach.
+  vag.dispose();
+  model.loraState.active = ["train"];
+  return { lora, losses };
+}