mlx-bun 0.0.11 → 0.0.12

package/src/chat-template.ts

@@ -42,6 +42,126 @@ export interface RenderOptions {
   enableThinking?: boolean;
 }
 
+type ChatRenderer = (
+  messages: ChatMessage[],
+  options: RenderOptions,
+) => string;
+
+function glm52ContentText(content: ChatMessage["content"], label: string): string {
+  if (typeof content === "string") return content;
+  if (!Array.isArray(content))
+    throw new Error(`${label} must be a string or an array of text parts`);
+  return content.map((part, index) => {
+    if (
+      !part ||
+      (part.type !== "text" && part.type !== "input_text") ||
+      typeof part.text !== "string"
+    ) {
+      throw new Error(`${label}.${index}: GLM-5.2 supports text content only`);
+    }
+    return part.text;
+  }).join("");
+}
+
+/** Text/tool subset of the pinned GLM-5.2 template used by Colibri. */
+export function renderGlm52Chat(
+  messages: ChatMessage[],
+  options: RenderOptions = {},
+): string {
+  if (!Array.isArray(messages) || messages.length === 0)
+    throw new Error("GLM-5.2 messages must be a non-empty array");
+  const {
+    addGenerationPrompt = true,
+    tools = null,
+    enableThinking = false,
+  } = options;
+  const prompt = ["[gMASK]<sop>"];
+  if (enableThinking)
+    prompt.push("<|system|>Reasoning Effort: Max");
+
+  const normalizedTools = normalizeToolSchemas(tools);
+  if (normalizedTools && normalizedTools.length > 0) {
+    prompt.push(
+      "<|system|>\n# Tools\n\nYou may call one or more functions to assist with the " +
+      "user query.\n\nYou are provided with function signatures within <tools></tools> " +
+      "XML tags:\n<tools>\n",
+    );
+    for (const tool of normalizedTools) {
+      const fn = { ...tool.function } as Record<string, unknown>;
+      delete fn.defer_loading;
+      delete fn.strict;
+      prompt.push(`${JSON.stringify(fn)}\n`);
+    }
+    prompt.push(
+      "</tools>\n\nFor each function call, output the function name and arguments " +
+      "within the following XML format:\n<tool_call>{function-name}" +
+      "<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value>" +
+      "<arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call>",
+    );
+  }
+
+  let previousWasTool = false;
+  for (let index = 0; index < messages.length; index++) {
+    const message = messages[index]!;
+    const role = message.role;
+    if (role === "system" || role === "developer") {
+      prompt.push(
+        `<|system|>${glm52ContentText(message.content, `messages.${index}.content`)}`,
+      );
+    } else if (role === "user") {
+      prompt.push(
+        `<|user|>${glm52ContentText(message.content, `messages.${index}.content`)}`,
+      );
+    } else if (role === "assistant") {
+      const content = message.content === null || message.content === undefined
+        ? ""
+        : glm52ContentText(message.content, `messages.${index}.content`);
+      prompt.push(`<|assistant|><think></think>${content.trim()}`);
+      for (const call of message.tool_calls ?? []) {
+        const raw = call.function.arguments;
+        let args: Record<string, unknown> = {};
+        if (typeof raw === "string") {
+          try {
+            const parsed = JSON.parse(raw);
+            if (parsed && typeof parsed === "object" && !Array.isArray(parsed))
+              args = parsed as Record<string, unknown>;
+          } catch {
+            // Match the pinned template path: malformed historical arguments
+            // contribute an empty call body rather than changing the prompt.
+          }
+        } else {
+          args = raw;
+        }
+        prompt.push(`<tool_call>${call.function.name}`);
+        for (const [key, value] of Object.entries(args)) {
+          prompt.push(
+            `<arg_key>${key}</arg_key><arg_value>` +
+            `${typeof value === "string" ? value : JSON.stringify(value)}</arg_value>`,
+          );
+        }
+        prompt.push("</tool_call>");
+      }
+    } else if (role === "tool") {
+      if (!previousWasTool) prompt.push("<|observation|>");
+      prompt.push(
+        `<tool_response>${glm52ContentText(message.content, `messages.${index}.content`)}` +
+        "</tool_response>",
+      );
+    } else {
+      throw new Error(`GLM-5.2 unsupported message role ${JSON.stringify(role)}`);
+    }
+    previousWasTool = role === "tool";
+  }
+  if (addGenerationPrompt) {
+    prompt.push(
+      enableThinking
+        ? "<|assistant|><think>"
+        : "<|assistant|><think></think>",
+    );
+  }
+  return prompt.join("");
+}
+
 /** JSON-schema type implied by a literal value. */
 function jsonTypeOf(v: unknown): string {
   if (typeof v === "number") return Number.isInteger(v) ? "integer" : "number";
@@ -116,7 +236,8 @@ export function normalizeToolSchemas(tools: ToolDefinition[] | null): ToolDefini
 }
 
 export class ChatTemplate {
-  readonly #template: Template;
+  readonly #template: Template | null;
+  readonly #renderer: ChatRenderer | null;
   readonly #bosToken: string | null;
   readonly #eosToken: string | null;
   /** Which reasoning format the template's `enable_thinking` channel uses, or
@@ -132,15 +253,20 @@ export class ChatTemplate {
   readonly supportsThinking: boolean;
 
   private constructor(
-    source: string,
+    source: string | null,
     bosToken: string | null,
     eosToken: string | null,
     forceNoThinking = false,
+    renderer: ChatRenderer | null = null,
+    explicitThinkingFormat: "think-tag" | "gemma-channel" | null = null,
   ) {
-    this.#template = new Template(source);
+    if (source === null && renderer === null)
+      throw new Error("ChatTemplate needs a Jinja source or renderer");
+    this.#template = source === null ? null : new Template(source);
+    this.#renderer = renderer;
     this.#bosToken = bosToken;
     this.#eosToken = eosToken;
-    const gatesThinking = source.includes("enable_thinking");
+    const gatesThinking = source?.includes("enable_thinking") ?? false;
     // `forceNoThinking` suppresses the switchable channel even when the template
     // carries the gemma markers. DiffusionGemma ships the shared gemma-family
     // template (with the `<|channel>thought…<channel|>` reasoning channel), but
@@ -149,13 +275,13 @@ export class ChatTemplate {
     // reasoning, no answer). The OptiQ reference never enables thinking for this
     // model — it always renders the template's `default(false)` pre-closed empty
     // channel and decodes the whole canvas as plain text. We match that.
-    this.thinkingFormat = forceNoThinking
+    this.thinkingFormat = explicitThinkingFormat ?? (forceNoThinking
       ? null
-      : gatesThinking && source.includes("<think>")
+      : gatesThinking && source!.includes("<think>")
         ? "think-tag"
-        : gatesThinking && source.includes("<|channel>")
+        : gatesThinking && source!.includes("<|channel>")
           ? "gemma-channel"
-          : null;
+          : null);
     this.supportsThinking = this.thinkingFormat !== null;
   }
 
@@ -169,7 +295,25 @@ export class ChatTemplate {
       const jinjaFile = Bun.file(`${modelDir}/chat_template.jinja`);
       if (await jinjaFile.exists()) source = await jinjaFile.text();
     }
-    if (!source) throw new Error(`${modelDir}: no chat template found`);
+    if (!source) {
+      const modelConfigFile = Bun.file(`${modelDir}/config.json`);
+      const modelConfig = await modelConfigFile.exists()
+        ? await modelConfigFile.json() as Record<string, unknown>
+        : null;
+      if (modelConfig?.model_type === "glm_moe_dsa") {
+        const tokenText = (t: unknown): string | null =>
+          t == null ? null : typeof t === "string" ? t : (t as any).content;
+        return new ChatTemplate(
+          null,
+          tokenText(config.bos_token),
+          tokenText(config.eos_token),
+          false,
+          renderGlm52Chat,
+          "think-tag",
+        );
+      }
+      throw new Error(`${modelDir}: no chat template found`);
+    }
     // @huggingface/jinja lacks the `min`/`max` array filters that real
     // Jinja2 has. MiniCPM5's template uses `[a, b]|min` in its assistant
     // tool-call history branch, so without this rewrite every multi-turn
@@ -189,8 +333,9 @@ export class ChatTemplate {
   }
 
   render(messages: ChatMessage[], options: RenderOptions = {}): string {
+    if (this.#renderer) return this.#renderer(messages, options);
     const { addGenerationPrompt = true, tools = null, enableThinking } = options;
-    return this.#template.render({
+    return this.#template!.render({
       messages,
       add_generation_prompt: addGenerationPrompt,
       // Guarantee every tool param schema has a `type` so templates that do

package/src/cli.ts

@@ -17,7 +17,15 @@ import { fit, skuMatrix, thisMachine } from "./fit";
 import { EvalDB } from "./evaldb";
 import pkg from "../package.json" with { type: "json" };
 import { renderHelp } from "./tui";
-import { isSupportedModelRecord } from "./model/support";
+import { isGlm52Config, isSupportedModelRecord } from "./model/support";
+import {
+  GLM52_G5_ASPIRATIONAL_DECODE_TPS,
+  GLM52_G5_DEFAULT_MAX_GENERATION_TOKENS,
+  GLM52_G5_DEFAULT_PROCESS_LIMIT_BYTES,
+  GLM52_G5_DIRECT_ORACLE_WARM_DECODE_TPS,
+  GLM52_G5_MEASURED_WARM_DECODE_TPS,
+  planGlm52MemoryForArtifact,
+} from "./model/glm52-memory";
 import type { JobEvent } from "./jobs/types";
 
 const argv = process.argv.slice(2);
@@ -153,6 +161,11 @@ Model & quality:
                             DSpark pins to its trained block width]
   --ngram-max <k>           Prompt-lookup longest match window  [default: 3]
   --ngram-min <k>           Prompt-lookup shortest match window  [default: 1]
+  --mtp <on|off>            GLM-5.2 checkpoint-native MTP. On by default;
+                            routes requests through the exact serial verify
+                            lane. off enables ordinary continuous batching.
+  --context-length <n>      GLM-5.2 resource-plan context reservation
+                            [default: 4096]
   --kv-quant <mode>         KV cache quantization: config (per-layer
                             kv_config.json when the model ships one), off
                             (bf16), 4 / 8 (uniform bits), or turbo[:k<bits>v
@@ -190,7 +203,8 @@ Model & quality:
                             [default: the model's generation_config.json]
   --max-tokens <n>          Completion cap when a request omits max_tokens
                             (mlx_lm.server flag; its default is 512 there)
-                            [default: 65536 chat / 512 raw completion]
+                            [default: 128 GLM-5.2; otherwise 65536 chat /
+                            512 raw completion]
   --hlg-sampling on|off     Piecewise tone-curve sampling (HLG): rolls off the
                             top, boosts the mids, gentles the tail. Gain folds
                             from --temperature. [default: off]
@@ -285,7 +299,9 @@ Options:
 
 Resumable (Range requests against partial blobs) and verified (sha256
 for LFS blobs). Re-running after an interruption continues where it
-stopped. Uses HF_TOKEN / hf auth login credentials when present.
+stopped. Before payload transfer, complete blobs and partial prefixes are
+credited and the target volume must fit the remainder plus a 1 GiB reserve.
+Uses HF_TOKEN / hf auth login credentials when present.
 When upstream pushed a new revision, the previous snapshot stays on
 disk — \`mlx-bun gc\` reclaims it.`,
 
@@ -1007,6 +1023,39 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
   return { port: Number(opt("port", "8080")), serverOptions };
 }
 
+/** Direct Colibri runtime defaults shared by `serve` and the embedded `pi`
+ * server. Other model families ignore this load-only option block. */
+function glm52RuntimeFlags(
+  serverOptions: import("./server").ServerOptions,
+): import("./model/factory").Glm52RuntimeOpenOptions {
+  const mtpRaw = opt("mtp");
+  let enableMtp: boolean | undefined;
+  if (mtpRaw !== null) {
+    if (["on", "1", "true"].includes(mtpRaw)) enableMtp = true;
+    else if (["off", "0", "false"].includes(mtpRaw)) enableMtp = false;
+    else {
+      console.error(`--mtp expects on|off (got "${mtpRaw}")`);
+      process.exit(1);
+    }
+  }
+  const contextRaw = opt("context-length");
+  let contextTokens: number | undefined;
+  if (contextRaw !== null) {
+    contextTokens = Number(contextRaw);
+    if (!Number.isSafeInteger(contextTokens) || contextTokens < 1) {
+      console.error(`--context-length expects an integer >= 1 (got "${contextRaw}")`);
+      process.exit(1);
+    }
+  }
+  return {
+    memoryBudgetBytes: serverOptions.memoryBudgetBytes,
+    batchSize: serverOptions.batch ?? 8,
+    maxGenerationTokens: serverOptions.defaultMaxTokens ?? 128,
+    ...(contextTokens !== undefined ? { contextTokens } : {}),
+    ...(enableMtp !== undefined ? { enableMtp } : {}),
+  };
+}
+
 /** `--adapter <dir>` (alias `--adapter-path`, mlx_lm.server's spelling):
  *  mount a LoRA adapter at startup — the same machinery as POST /v1/adapters —
  *  and make it the default for requests that don't select one (a request's
@@ -1332,9 +1381,69 @@ switch (cmd) {
     if (reg.list().length === 0) await reg.scan();
     const m = reg.resolve(query);
     const config = await loadModelConfig(m.path);
-    const ctx = Number(opt("ctx", "8192"));
-    const r = fit(config, m.sizeBytes, ctx, thisMachine(), undefined, m.expertsBytes);
     const { box, table, style, h1, gradient } = await import("./tui");
+    const glm = isGlm52Config(config);
+    const ctx = Number(opt("ctx", glm ? "4096" : "8192"));
+    if (glm) {
+      const machine = thisMachine();
+      let plan;
+      try {
+        plan = await planGlm52MemoryForArtifact(m.path, {
+          machineBytes: machine.ramBytes,
+          processLimitBytes: Math.min(
+            GLM52_G5_DEFAULT_PROCESS_LIMIT_BYTES,
+            machine.ramBytes,
+          ),
+          contextTokens: ctx,
+          maxGenerationTokens: Math.min(
+            ctx,
+            GLM52_G5_DEFAULT_MAX_GENERATION_TOKENS,
+          ),
+          batchSize: 1,
+          enableMtp: true,
+        });
+      } catch (error) {
+        console.error(
+          `GLM-5.2 resource plan refused: ` +
+          `${error instanceof Error ? error.message : String(error)}`,
+        );
+        process.exitCode = 1;
+        break;
+      }
+      const gib = (bytes: number): string => `${(bytes / 2 ** 30).toFixed(2)} GiB`;
+      const li = plan.lineItems;
+      const kvBytes = li.targetKvBytes + li.mtpKvBytes;
+      const otherBytes = plan.plannedProcessBytes - li.residentWeightsBytes -
+        li.mainExpertSlabBytes - li.mtpExpertSlabBytes - kvBytes;
+      h1("will it fit?");
+      console.log(`  ${style.bold(m.repoId)} ${style.dim(`@ ${ctx.toLocaleString()} context · streamed GLM plan`)}`);
+      console.log();
+      box([
+        `artifact on disk  ${gib(m.sizeBytes).padStart(10)} ${style.dim("(streamed, not resident)")}`,
+        `resident weights  ${gib(li.residentWeightsBytes).padStart(10)}`,
+        `main expert slab  ${gib(li.mainExpertSlabBytes).padStart(10)}`,
+        `MTP expert slab   ${gib(li.mtpExpertSlabBytes).padStart(10)}`,
+        `target + MTP KV   ${gib(kvBytes).padStart(10)}`,
+        `runtime reserves  ${gib(otherBytes).padStart(10)}`,
+        `process plan      ${gib(plan.plannedProcessBytes).padStart(10)} ${style.dim(`of ${gib(plan.processLimitBytes)}`)}  ${style.green(style.bold("FITS"))}`,
+        `macOS headroom    ${gib(plan.machineHeadroomBytes).padStart(10)}`,
+        "",
+        `max safe context  ${style.bold(plan.contextTokens.toLocaleString())} tokens`,
+        `max generation    ${style.bold(plan.maxGenerationTokens.toLocaleString())} tokens`,
+        `measured warm     ${gradient(`${GLM52_G5_MEASURED_WARM_DECODE_TPS.toFixed(3)} tok/s`)} ${style.dim("(quality-preserving, MTP on)")}`,
+        `direct oracle     ${GLM52_G5_DIRECT_ORACLE_WARM_DECODE_TPS.toFixed(2)} tok/s`,
+        `aspirational      ${GLM52_G5_ASPIRATIONAL_DECODE_TPS.toFixed(1)} tok/s ${style.dim("(not a fit prediction or release gate)")}`,
+      ]);
+      if (flag("skus")) {
+        console.log(style.dim(
+          "  GLM-5.2's direct-container preset is validated on the M1 Max 32 GB; " +
+          "the generic resident-weight SKU projection does not apply.",
+        ));
+      }
+      console.log();
+      break;
+    }
+    const r = fit(config, m.sizeBytes, ctx, thisMachine(), undefined, m.expertsBytes);
     h1("will it fit?");
     console.log(`  ${style.bold(m.repoId)} ${style.dim(`@ ${ctx.toLocaleString()} context · this machine`)}`);
     console.log();
@@ -1529,6 +1638,7 @@ switch (cmd) {
     const t0 = performance.now();
     const ctx = await loadContext(m.path, m.repoId, {
       memoryBudgetBytes: rt.serverOptions.memoryBudgetBytes,
+      glm: glm52RuntimeFlags(rt.serverOptions),
       ...(draftModelDir || draftKind === "ngram"
         ? {
             ...(draftModelDir ? { draftModelDir } : {}),
@@ -1585,7 +1695,12 @@ switch (cmd) {
       process.exit(1);
     }
     const { loadTaskModel, generateText } = await import("./eval/runner");
-    const tm = await loadTaskModel(positional(0) ?? opt("query") ?? "");
+    const maxTokens = Number(opt("max-tokens", "256"));
+    const tm = await loadTaskModel(
+      positional(0) ?? opt("query") ?? "",
+      undefined,
+      { maxGenerationTokens: maxTokens, enableMtp: false },
+    );
     // Resolve the tier/--kv-quant KV scheme exactly the way serve does
     // (server.ts kvScheme): "config" = the model's kv_config.json (absent →
     // bf16, same silent fallback as serve), N = uniform bits from decode
@@ -1601,7 +1716,7 @@ switch (cmd) {
       : undefined;
     const num = (n: string): number | undefined => { const v = opt(n); return v == null ? undefined : Number(v); };
     const text = await generateText(tm, prompt, {
-      maxTokens: Number(opt("max-tokens", "256")),
+      maxTokens,
       useChat: !flag("raw"),
       ...(kvScheme !== undefined ? { kvScheme } : {}),
       sampler: { temperature: num("temperature") ?? num("temp"), topP: num("top-p"), topK: num("top-k"), seed: num("seed") },
@@ -1883,6 +1998,7 @@ switch (cmd) {
       "--draft-model", "--num-draft-tokens", "--draft-kind", "--ngram-max", "--ngram-min",
       "--compiled-decode", "--compiled-activations", "--fused-sdpa", "--thinking",
       "--temperature", "--temp", "--top-p", "--top-k", "--max-tokens",
+      "--mtp", "--context-length",
       "--hlg-sampling", "--hlg-width", "--hlg-shoulder", "--hlg-toe", "--hlg-pivot-offset",
       "--paged-kv-block-size",
     ]);
@@ -1948,7 +2064,10 @@ switch (cmd) {
       const { createServer, loadContext } = await import("./server");
       const sLoad = step(`loading ${m.repoId}`);
       const t0 = performance.now();
-      const ctx = await loadContext(m.path, m.repoId, { memoryBudgetBytes: rt.serverOptions.memoryBudgetBytes });
+      const ctx = await loadContext(m.path, m.repoId, {
+        memoryBudgetBytes: rt.serverOptions.memoryBudgetBytes,
+        glm: glm52RuntimeFlags(rt.serverOptions),
+      });
       await mountStartupAdapter(ctx, rt.serverOptions);
       createServer(ctx, port, { ...rt.serverOptions, owner: "pi-session" });
       startedServer = true;

package/src/config.ts

@@ -2,6 +2,11 @@
 // Parses HF config.json (Gemma 4 unified layout: text_config nested),
 // the OptiQ per-layer quantization map, and kv_config.json when present.
 
+import {
+  parseGlm52Config,
+  type Glm52Config,
+} from "./model/glm52-config";
+
 export interface RopeParams {
   ropeTheta: number;
   ropeType: string;
@@ -223,6 +228,18 @@ export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
   const raw = (await Bun.file(`${modelDir}/config.json`).json()) as Record<string, any>;
   const modelType = raw.model_type;
   const isDiffusion = modelType === "diffusion_gemma";
+  const isGlm52 = modelType === "glm_moe_dsa";
+  // The converted Colibri artifact keeps stop metadata split across the two
+  // HF files. Preserve the old config-only behavior for every other family;
+  // GLM-5.2 explicitly unions both lists through its dedicated parser.
+  let glm52: Glm52Config | null = null;
+  if (isGlm52) {
+    const generationFile = Bun.file(`${modelDir}/generation_config.json`);
+    const generation = await generationFile.exists()
+      ? await generationFile.json() as Record<string, any>
+      : null;
+    glm52 = parseGlm52Config(modelDir, raw, generation);
+  }
   // Gemma 4 unified nests the LM config; plain text models keep it at top level.
   // DiffusionGemma keeps it flat too, but omits all dims — backfill from defaults.
   const baseT = (raw.text_config ?? raw) as Record<string, any>;
@@ -238,32 +255,44 @@ export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
   const qwenRope = (t.rope_parameters ?? {}) as Record<string, any>;
 
   const text: TextConfig = {
-    hiddenSize: t.hidden_size,
-    numHiddenLayers: t.num_hidden_layers,
-    numAttentionHeads: t.num_attention_heads,
-    numKeyValueHeads: t.num_key_value_heads,
-    headDim: t.head_dim ?? Math.floor(t.hidden_size / t.num_attention_heads),
-    numGlobalKeyValueHeads: t.num_global_key_value_heads ?? t.num_key_value_heads,
-    globalHeadDim: t.global_head_dim ?? t.head_dim ?? Math.floor(t.hidden_size / t.num_attention_heads),
+    hiddenSize: glm52?.hiddenSize ?? t.hidden_size,
+    numHiddenLayers: glm52?.numHiddenLayers ?? t.num_hidden_layers,
+    numAttentionHeads: glm52?.numAttentionHeads ?? t.num_attention_heads,
+    numKeyValueHeads: glm52?.numKeyValueHeads ?? t.num_key_value_heads,
+    headDim: glm52?.qkHeadDim ?? t.head_dim ?? Math.floor(t.hidden_size / t.num_attention_heads),
+    numGlobalKeyValueHeads:
+      glm52?.numKeyValueHeads ?? t.num_global_key_value_heads ?? t.num_key_value_heads,
+    globalHeadDim:
+      glm52?.qkHeadDim ??
+      t.global_head_dim ??
+      t.head_dim ??
+      Math.floor(t.hidden_size / t.num_attention_heads),
     attentionKEqV: t.attention_k_eq_v ?? false,
-    intermediateSize: t.intermediate_size,
+    intermediateSize: glm52?.intermediateSize ?? t.intermediate_size,
     hiddenActivation: t.hidden_activation ?? t.hidden_act ?? "gelu_pytorch_tanh",
-    rmsNormEps: t.rms_norm_eps,
-    vocabSize: t.vocab_size,
-    maxPositionEmbeddings: t.max_position_embeddings,
+    rmsNormEps: glm52?.rmsNormEps ?? t.rms_norm_eps,
+    vocabSize: glm52?.vocabSize ?? t.vocab_size,
+    maxPositionEmbeddings:
+      glm52?.maxPositionEmbeddings ?? t.max_position_embeddings,
     slidingWindow: t.sliding_window ?? 0,
-    layerTypes: t.layer_types ?? (isLlama ? Array(t.num_hidden_layers).fill("full_attention") : []),
+    layerTypes: glm52
+      ? Array(glm52.numHiddenLayers).fill("full_attention")
+      : t.layer_types ?? (isLlama ? Array(t.num_hidden_layers).fill("full_attention") : []),
     hiddenSizePerLayerInput: t.hidden_size_per_layer_input ?? 0,
-    vocabSizePerLayerInput: t.vocab_size_per_layer_input ?? t.vocab_size,
+    vocabSizePerLayerInput:
+      t.vocab_size_per_layer_input ?? glm52?.vocabSize ?? t.vocab_size,
     numKvSharedLayers: t.num_kv_shared_layers ?? 0,
-    enableMoeBlock: t.enable_moe_block ?? false,
-    numExperts: t.num_experts ?? 0,
+    enableMoeBlock: glm52 ? true : t.enable_moe_block ?? false,
+    numExperts: glm52?.numRoutedExperts ?? t.num_experts ?? 0,
     // gemma4 uses `top_k_experts`; qwen3_moe uses `num_experts_per_tok`.
-    topKExperts: t.top_k_experts ?? t.num_experts_per_tok ?? 0,
-    moeIntermediateSize: t.moe_intermediate_size ?? 0,
-    decoderSparseStep: t.decoder_sparse_step ?? 1,
-    mlpOnlyLayers: t.mlp_only_layers ?? [],
-    normTopkProb: t.norm_topk_prob ?? false,
+    topKExperts:
+      glm52?.numExpertsPerToken ?? t.top_k_experts ?? t.num_experts_per_tok ?? 0,
+    moeIntermediateSize: glm52?.moeIntermediateSize ?? t.moe_intermediate_size ?? 0,
+    decoderSparseStep: glm52 ? 1 : t.decoder_sparse_step ?? 1,
+    mlpOnlyLayers: glm52
+      ? Array.from({ length: glm52.firstKDenseReplace }, (_, layer) => layer)
+      : t.mlp_only_layers ?? [],
+    normTopkProb: glm52?.normTopkProb ?? t.norm_topk_prob ?? false,
     linearNumValueHeads: t.linear_num_value_heads ?? 0,
     linearNumKeyHeads: t.linear_num_key_heads ?? 0,
     linearKeyHeadDim: t.linear_key_head_dim ?? 0,
@@ -271,9 +300,19 @@ export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
     linearConvKernelDim: t.linear_conv_kernel_dim ?? 0,
     fullAttentionInterval: t.full_attention_interval ?? 0,
     attnOutputGate: t.attn_output_gate ?? false,
-    partialRotaryFactor:
-      qwenRope.partial_rotary_factor ?? t.partial_rotary_factor ?? 1.0,
-    ropeParameters: isQwen35
+    partialRotaryFactor: glm52
+      ? glm52.qkRopeHeadDim / glm52.qkHeadDim
+      : qwenRope.partial_rotary_factor ?? t.partial_rotary_factor ?? 1.0,
+    ropeParameters: glm52
+      ? {
+          full_attention: {
+            ropeTheta: glm52.ropeTheta,
+            ropeType: "default",
+            partialRotaryFactor: glm52.qkRopeHeadDim / glm52.qkHeadDim,
+            factor: 1.0,
+          },
+        }
+      : isQwen35
       ? {
           full_attention: {
             ropeTheta: qwenRope.rope_theta ?? t.rope_theta ?? 10000,
@@ -297,7 +336,7 @@ export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
     finalLogitSoftcapping: t.final_logit_softcapping ?? null,
     tieWordEmbeddings: t.tie_word_embeddings ?? raw.tie_word_embeddings ?? false,
     bosTokenId: t.bos_token_id ?? raw.bos_token_id,
-    eosTokenId: t.eos_token_id ?? raw.eos_token_id,
+    eosTokenId: glm52?.eosTokenIds ?? t.eos_token_id ?? raw.eos_token_id,
     canvasLength: raw.canvas_length ?? t.canvas_length,
   };
 
@@ -311,14 +350,20 @@ export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
     }));
   }
 
-  const eos = raw.eos_token_id ?? text.eosTokenId;
+  const eos = glm52?.eosTokenIds ?? raw.eos_token_id ?? text.eosTokenId;
   return {
     modelDir,
     modelType,
     architectures: raw.architectures ?? [],
     dtype: raw.dtype ?? "bfloat16",
     text,
-    quantization: parseQuantization(raw.quantization ?? raw.quantization_config),
+    // Colibri's config retains the SOURCE FP8 quantization_config, while the
+    // serving shards are converted U8 + `.qs` tensors whose int4/int8/group
+    // geometry is inferred and validated by Glm52Container. Advertising the
+    // source FP8 block as mlx affine quantization would be actively wrong.
+    quantization: glm52
+      ? null
+      : parseQuantization(raw.quantization ?? raw.quantization_config),
     kvQuant,
     hasVisionSidecar: await Bun.file(`${modelDir}/optiq_vision.safetensors`).exists(),
     eosTokenIds: Array.isArray(eos) ? eos : [eos],

package/src/diffusion/diffusion-generate.ts

@@ -245,9 +245,7 @@ export function diffusionGenerate(
 
       let currentCanvas = ops.randint(0, vocab, [1, canvasLength], I32); // RNG: init
       let draftReveal = ops.zeros([1, canvasLength], Dtype.bool);
-      let draftCanvas: MlxArray = ops.logicalNot(ops.logicalNot(draftReveal)); // placeholder, replaced below
-      draftCanvas.dispose();
-      draftCanvas = ops.reshape(currentCanvas, [1, canvasLength]); // alias copy
+      let draftCanvas = ops.reshape(currentCanvas, [1, canvasLength]); // alias copy
       let argmaxCanvas = ops.reshape(currentCanvas, [1, canvasLength]);
       let scEmbeddings: MlxArray | null = null;
       const history: MlxArray[] = [];
@@ -266,7 +264,9 @@ export function diffusionGenerate(
         logits.dispose();
 
         argmaxCanvas.dispose();
-        argmaxCanvas = ops.argmaxAxis(processed, -1).astype(I32); // [1, L]
+        const argmax = ops.argmaxAxis(processed, -1);
+        argmaxCanvas = argmax.astype(I32); // [1, L]
+        argmax.dispose();
 
         if (curStep === 1) {
           processed.dispose();
@@ -296,7 +296,9 @@ export function diffusionGenerate(
           currentCanvas.dispose();
           currentCanvas = newCurrent;
           draftReveal.dispose();
-          draftReveal = ops.logicalNot(ops.logicalNot(acceptance)); // = acceptance
+          const notAcceptance = ops.logicalNot(acceptance);
+          draftReveal = ops.logicalNot(notAcceptance); // = acceptance
+          notAcceptance.dispose();
           draftCanvas.dispose();
           draftCanvas = ops.reshape(argmaxCanvas, [1, canvasLength]);
           accepted.dispose();
@@ -341,6 +343,10 @@ export function diffusionGenerate(
           stableStop &&
           stableAndConfident(argmaxCanvas, processed, history, stabilityThreshold, stopConfidence)
         ) {
+          // Entropy sampling precomputes the next step's self-conditioning
+          // tensor before evaluating this stop. The break means ownership is
+          // never transferred to scEmbeddings, so release the pending value.
+          nextSc?.dispose();
           processed.dispose();
           break;
         }
@@ -361,9 +367,9 @@ export function diffusionGenerate(
       draftCanvas.dispose();
 
       // Emit from the final argmax canvas.
-      const blockTokens = [...argmaxCanvas.astype(Dtype.float32).toFloat32()].map((x) =>
-        Math.round(x),
-      );
+      const argmaxFloat = argmaxCanvas.astype(Dtype.float32);
+      const blockTokens = [...argmaxFloat.toFloat32()].map((x) => Math.round(x));
+      argmaxFloat.dispose();
       argmaxCanvas.dispose();
       blocks.push(blockTokens);
       totalSteps += stepsThisCanvas;