mlx-bun 0.4.0 → 1.0.0

package/docs/archive/investigations/diffusion-gemma-port.md

@@ -1,602 +0,0 @@
-# DiffusionGemma port — plan & parity strategy
-
-Porting **DiffusionGemma-26B-A4B-it** (`mlx-community/diffusiongemma-26B-A4B-it-OptiQ-4bit`,
-14 GB, model_type `diffusion_gemma`) into mlx-bun, holding **L2 parity with
-mlx-optiq**. This is the first **non-autoregressive** model in the codebase:
-it fills a fixed 256-token *canvas* and un-masks it over a few denoising steps
-rather than generating left-to-right. mlx-optiq docs:
-https://mlx-optiq.com/docs/diffusion
-
-> Status doc: STATUS.md. Durable phase log: PLAN.md. This file is the design +
-> phase plan for this one feature; promote the phase checklist into PLAN.md once
-> approved.
-
-## The parity oracle (resolved 2026-06-24)
-
-**Stock mlx-lm/mlx-vlm cannot load this model** — there is no bf16 ancestor in
-the usual L1 sense. OptiQ ships a *vendored, dependency-free decoder*, so **the
-oracle IS mlx-optiq itself**. The mixed-precision quant in the checkpoint is the
-model's native weight format (not an optional L3 layer), so **matching optiq's
-output = L2 = the correctness floor here.**
-
-- Reference is now installed: `mlx-optiq 0.2.7` in
-  `/Users/joshrossi/Code/mlx-lm/.venv` (upgraded from 0.2.1; the diffusion
-  decoder needs ≥0.2.3). `mlx`/`mlx-lm`/`mlx-metal` stayed pinned at
-  **0.31.2 / 0.31.3 / 0.31.2**, so the existing parity oracles (Gemma/CPM/Qwen)
-  are **unaffected** — verified by `pip install --dry-run`.
-- Reference source (read these; they are the contract):
-  - Model: `optiq/vlm/_mlxvlm/models/diffusion_gemma/`
-    — `language.py` (802 L, the core: Attention, Router/Experts MoE,
-    DecoderLayer, SelfConditioning, DecoderModel `_embed_canvas` /
-    `_make_decoder_masks`, EncoderModel vision merge), `diffusion_gemma.py`
-    (178 L, top `Model` + `sanitize` + `quant_predicate`), `config.py` (87 L),
-    `processing_diffusion_gemma.py` (137 L).
-  - Engine: `optiq/vlm/_mlxvlm/generate/diffusion.py` (1148 L) — canvas init,
-    linear temperature schedule (`t_min`→`t_max`), the un-mask loop, the two
-    samplers (`confidence-threshold` default in OptiQ, `entropy-bound` model
-    default), self-conditioning feedback, stability/EOS handling.
-  - OptiQ surface: `optiq/vlm/diffusion_gemma/{loader,generate,lora,
-    calibration,convert,eval_adapter}.py` (load/generate/train entry points).
-- Run the reference with `/Users/joshrossi/Code/mlx-lm/.venv/bin/python`. Golden
-  dumps for parity tests go through `optiq.vlm.diffusion_gemma.{load,generate}`.
-
-## What the checkpoint actually is (from config.json + weight map)
-
-- `architectures: ["DiffusionGemmaForBlockDiffusion"]`, `canvas_length: 256`,
-  `model_type: diffusion_gemma`, bf16 base, OptiQ mixed-precision quant
-  (group_size 64, bits ∈ {4,8} per-tensor; 8-bit budget on **early-layer
-  attention + routers**, 4-bit on dense MLPs — inverse of the hand-coded recipe).
-- **Language tower** `model.decoder`: 30 layers, hidden 2816, 16 heads. Each
-  layer runs a **parallel dense `mlp` + routed MoE** branch (128-expert top-8,
-  `router.proj` + `experts.{gate_up,down}_proj`) summed via 7 RMSNorms + a
-  `layer_scalar` — conceptually the gemma4 parallel-branch idea, but DIFFERENT
-  geometry/norm-placement than our gemma4-26B (verify, don't assume reuse).
-  Attention is **scale=1.0, QK/V-norm post-proj/pre-RoPE, NO attention softcap**
-  (only `final_logit_softcapping=30.0`, fp32); sliding layers hd=256/kv=8, full
-  layers (5,11,17,23,29) hd=512/kv=2 + partial-rotary 0.25. Plus
-  `model.decoder.self_conditioning` (gate/up/down MLP) — diffusion-specific.
-- **Vision tower** `model.encoder.vision_tower`: 27-layer SigLIP
-  (`patch_embedder.input_proj` + `encoder.layers.N.{self_attn,mlp}`). We already
-  have a SigLIP port in `src/vision/siglip.ts` (16-layer for e4b) — likely
-  reusable with a layer-count/config delta. image-text-to-text.
-- Generation config: `max_denoising_steps: 48`, `max_new_tokens: 256`,
-  `confidence_threshold: 0.005`, `t_min: 0.4`, `t_max: 0.8`,
-  `stability_threshold: 1`, default `EntropyBoundSamplerConfig(entropy_bound=0.1)`,
-  `eos_token_id: [1, 106, 50]`.
-- New tokenizer + chat template (`<|turn>role`, `<|channel>thought`,
-  `<|tool>`/`<tool|>` blocks, `<|image|>`). The canvas is initialized to
-  **uniform-random token IDs** (`randint(0, vocab_size)`), NOT a mask token —
-  no discrete mask-token path is used at inference (dossier §9.1).
-
-## Architecture deltas vs what we already have
-
-(Refined after D0 — see the dossier for exact shapes/math.)
-
-| Piece | Reuse / new |
-|---|---|
-| MoE Router + Experts (128-expert, top-8, `SwitchLinear` `_gather_sort`) | **Reuse/adapt** gemma4.ts — but router norm is **pre**-projection + `scale`·`hidden**-0.5`, weights `*per_expert_scale`; verify against gemma4 |
-| Attention — QK/V-norm **post-proj/pre-RoPE**, **scale=1.0, NO attn softcap**, partial-rotary 0.25 on full layers | **Adapt** gemma4.ts Attention (geometry differs: sliding hd=256/kv=8, full hd=512/kv=2; full layers reuse k as v) |
-| SigLIP vision tower (27 layers, hidden 1152, patch 16) | **Reuse** src/vision/siglip.ts with config delta + `MultimodalEmbedder` merge |
-| Quantized linear / mixed-precision load | **Reuse** loader; read per-tensor bits from `config.quantization` map (not just the predicate) |
-| Tied LM head via quantized `embed_tokens.as_linear()` | **New-ish** — need quantized transpose-matmul + `quantized_matmul(...,transpose=False)` for soft embeds |
-| `_embed_canvas` (hard embed + soft self-cond signal) | **New** |
-| Bidirectional / block decoder masks (`_make_decoder_masks`) | **New** — NOT causal; **the crux / highest bug-risk** |
-| `SelfConditioning` MLP + per-step soft-embedding feedback | **New** |
-| Denoising engine (uniform-random canvas init, linear temp schedule, un-mask loop, 2 samplers, stability/EOS) | **New** — a whole second generation path |
-| Cache mix: `KVCache` (full) + `RotatingKVCache(1024)` (sliding); decoder reads encoder K/V from cache state | **Reuse/adapt** — `_cache_offset`/`_cache_state` helpers |
-| Tokenizer + chat template + image processor (variable soft-token budget) | **New** |
-| ~~Canon depthwise-conv~~ | **N/A — absent in checkpoint** |
-
-## Phase plan (exit criteria)
-
-### D0 — Recon & reference dossier `[x]` (done 2026-06-24)
-Dossier appended below ("Reference dossier (D0)"). Headline corrections to the
-pre-recon assumptions: **no Canon/depthwise-conv tensors exist** (pure
-transformer); the **LM head is tied** (`embed_tokens.as_linear()`, a 4-bit
-`QuantizedEmbedding`); shapes are **hidden 2816 / 30 layers / 16 heads** (NOT the
-gemma4-26B geometry); each layer runs a **parallel dense-MLP + MoE** branch with
-**7 RMSNorms + a `layer_scalar`**; attention is **scale=1.0 with no attention
-softcap** (only `final_logit_softcapping=30.0`, fp32); the canvas is initialized
-to **uniform-random token IDs**, not a mask token. Full detail + 12 risk items
-below.
-
-#### Original D0 brief
-Read the entire reference and write it down so D1–D3 are mechanical. Done as a
-**parallel-reader workflow** (see Orchestration). Produces, appended to this doc:
-a component-by-component spec of `language.py` + `diffusion_gemma.py`, the
-weight-name → module map audited against the real checkpoint (0 missing / 0
-unused), the exact `_make_decoder_masks` construction, the denoising-loop
-pseudocode, both sampler algorithms, the self-conditioning data flow, and the
-tokenizer/chat-template/mask-token facts.
-**Exit:** dossier complete; weight-name audit clean; golden-dump harness sketched.
-
-### D1 — Single-forward parity (the static graph) `[x]` (done 2026-06-24 — BIT-EXACT)
-**DONE.** `src/model/diffusion-gemma.ts` (`DiffusionGemmaModel`) + factory detection
-(`isDiffusionGemmaConfig`; serving still throws "wired in D3" — the AR `RuntimeModel`
-union is left untouched so the 96 script-callers don't break). One full forward over
-the real 14 GB checkpoint — **encoder prefill → bidirectional decoder canvas pass
-(`_make_decoder_masks`) → parallel dense-MLP + 128-expert MoE → SelfConditioning →
-tied 4-bit head → fp32 softcap** — is **BIT-EXACT vs the optiq golden**: argmax 256/256,
-maxDiff 0.0, relRMSE 0.0, meanKL 0.0 (`tests/diffusion-parity.test.ts`,
-`MLX_BUN_TEST_DIFFUSION=1`). Per-stage sub-gates (enc_hidden / dec_hidden / presoftcap /
-per-layer / layer-0 attn·dense·MoE) all 0.0 (`scripts/experiments/diffusion-{stage,layer,l0}-diag.ts`,
-golden dumps in `scripts/oracle/gen-diffusion-golden.py`).
-- **The one bug found + fixed (per-model parity gotcha):** the Router pre-projection
-  norm must be the literal two-step `rms_norm(x, None, eps) * scale * hidden**-0.5`,
-  **NOT** the gemma4-style fold `rms_norm(x, scale*hidden**-0.5)`. The fold changes the
-  bf16 intermediate rounding → shifts the softmax routing weights by ~0.01 → 1.7% on the
-  expert output → 12.7% by the encoder output (argmax mostly survived, masking it). Same
-  experts selected (0/21 set mismatch); only the *weights* drifted. Copy-verbatim
-  methodology localized it: attn / h_mid / dense-MLP were already 0.0, only the MoE branch
-  diverged.
-- Vision merge deferred to D3 (golden is text-only; the 27-layer SigLIP tower is present
-  in the checkpoint and wired via `parseSiglipConfig` when D3 lands).
-
-### D1 — Single-forward parity (the static graph) — original brief `[x]`
-`src/model/diffusion-gemma.ts` + config detection (`support.ts`
-`isDiffusionGemmaConfig`) + factory wiring. Port embed_canvas, DecoderLayer
-(attn + dense MLP + MoE parallel + softcap), SelfConditioning, the bidirectional
-decoder masks, vision feature merge (reuse siglip), final norm + head.
-Sub-gates (model-free where possible, then on real weights): vision features
-bit-exact; one decoder layer bit-exact; MoE block bit-exact; self-conditioning
-bit-exact.
-**Exit:** **one full forward over the real 14 GB checkpoint produces logits
-bit-exact vs the optiq reference** (Python golden dump through
-`optiq.vlm.diffusion_gemma`), at a fixed canvas state. This is the L2 oracle for
-the graph; where the quantized GEMV is 1-ULP off (the known megakernel finding),
-gate by KL + argmax-agreement, not the bit-exact golden.
-
-### D2 — Denoising generation engine `[x]` (done 2026-06-24 — token-for-token)
-**DONE.** `src/diffusion/diffusion-generate.ts` (`diffusionGenerate`): prefill→cache
-reuse, linear temp schedule, the un-mask loop, BOTH samplers (confidence-threshold =
-OptiQ public default, entropy-bound = engine default), self-conditioning feedback,
-stability/EOS, canvas block loop. **Token-for-token parity vs the optiq engine at
-temp 0 on a fixed seed for BOTH samplers** (`tests/diffusion-gen-parity.test.ts`,
-`MLX_BUN_TEST_DIFFUSION=1`; golden `scripts/oracle/gen-diffusion-gen-golden.py`):
-confidence 17 tok / 7 steps, entropy 15 tok / 48 steps — both `.tokens`/`.steps`/
-`.finishReason` identical.
-- **RNG parity (the crux) SOLVED:** bound `mlx_random_randint` (+`mlx_random_seed`,
-  `mlx_cummax`, `logicalNot`/`equal`/`all`/`anyAxis`/`lessEqual`/`itemBool`) in
-  `ffi.ts`/`ops.ts`. `ops.randint(key=null)` threads the GLOBAL mlx key, so seeding +
-  calling randint in the same ORDER as the reference (1 canvas init + 1 re-noise per
-  non-final step) reproduces every draw bit-for-bit. Verified standalone: 0/256
-  mismatch across consecutive draws.
-- **Three bugs found + fixed:** (1) `processed = logits / schedT` must be a **division**,
-  not reciprocal-multiply — the f32 rounding differs and the 0.9 confidence threshold
-  is a hard cutoff that flips acceptance on 1 ULP, diverging the trajectory (was 8 vs 7
-  steps). (2) `history.push` needs an **independent copy** (add-zero) — a reshape/view
-  aliases the argmax buffer and reads freed data after disposal → false "stable" →
-  premature stop. (3) **The as-loaded oracle has `model.config.generation_config = None`**,
-  so `stable_and_confident` is a NO-OP (entropy runs all 48 steps) and the eos set is the
-  tokenizer's `stopping_criteria` = **{1, 106}** (50 is NOT eos despite generation_config.json).
-  L2 parity = match the oracle as it RUNS: stable-stop OFF unless explicitly configured.
-- Multi-block continuation (`extendPrefill`) is wired but only single-block is gated
-  (max_tokens 64 ≤ canvas). temperature>0 (categorical) + image input deferred to D3.
-
-### D2 — Denoising generation engine — original brief `[x]`
-`src/diffusion/` (or `src/generate-diffusion.ts`): canvas init, linear
-temperature schedule, the un-mask loop, `confidence-threshold` (OptiQ default,
-the L2 target) AND `entropy-bound` (model default) samplers, self-conditioning
-feedback, stability/EOS termination. Deterministic harness (same canvas init →
-same final tokens), parity vs optiq `generate(...)`.
-**Exit:** full deterministic generation matches optiq **token-for-token** for
-both samplers on a fixed seed/prompt set (text-only first). Step-count and
-final-canvas match.
-
-### D3 — Serving + CLI + image input `[x]` (DONE 2026-06-24 — text + image, parity-exact)
-**IMAGE-TEXT-TO-TEXT DONE — token-for-token parity.** A DEDICATED DiffusionGemma vision
-tower (`src/vision/diffusion-vision.ts`, `DiffusionVisionTower`) — its OWN module, NOT a
-reuse of the e4b SigLIP path (they're separate models: e4b ships a bf16
-`optiq_vision.safetensors` sidecar, DiffusionGemma's vision is inline-quantized). A
-parity-exact op-for-op port of optiq's gemma4 `VisionModel` at the DiffusionGemma geometry
-(hidden 1152, head_dim 72, 16 heads, 27 layers, position_embedding_size 10240, standardize ON).
-Wired through the encoder vision merge (`#embedInputsVision` masked-scatter + `#encoderVisionMask`
-causal|overlay) into the denoising engine; `preprocess()` (resize/rescale) + `spliceImageTokens`
-build the prompt. **Verified end-to-end** (`tests/diffusion-vision.test.ts`): spliced ids EXACT
-vs optiq + generation TOKEN-FOR-TOKEN ("This is a solid gray square." on grad-768) + **served
-live** via the OpenAI vision API. Server branch in `handleChat` (diffusion gets its own vision
-lane; v1 = single image); `generate()` threads `visionPixels`.
-- **Three vision bugs found + fixed (copy-verbatim per-stage diffing, dequant→bit-exact):**
-  (1) **input_proj is QUANTIZED (uint32 weight)**, so the reference's
-  `patches.astype(input_proj.weight.dtype)` TRUNCATES the `2*(x-0.5)` patches to integers
-  (mostly 0) before the matmul — a trained-in quirk e4b never hit (its input_proj is bf16). The
-  patchembed went bit-exact once replicated. (2) head_dim 72 isn't a fused-SDPA size; the
-  reference `ensure_fused_sdpa` PADS to 80 (zeros) and slices back — replicating it dropped the
-  per-layer error from 0.28% to ~bf16-noise. (3) `down_proj` is plain bf16 (not quantized) —
-  a `VisionLinear` that is quantized-or-plain. Residual: ~2.3% feature relRMSE from 27-layer bf16
-  accumulation (the e4b-class tier), which the denoising is robust to → identical tokens.
-
-### D3 — Serving + CLI + image input — earlier note `[x]`
-**TEXT serving DONE + verified live.** `createModel` returns `DiffusionGemmaModel`
-(added to the `RuntimeModel` union with the AR-only surface — `forwardHidden`/`forward`/
-`generate`/`forwardEmbeddings` — as throwing stubs, real `loraState`/`loraTargets`/
-`prefixBase`/`makeCache`, so the 96 AR script-callers still typecheck; baseline held at
-117). `generate()` detects the model and routes to the denoising engine
-(`generateDiffusionInner`), preserving the `Generation`/`GenerateStats` contract so the CLI
-and server stream it through the existing token machinery. The gateway keeps diffusion on
-the serial lane (`willBatch` returns false for it). CLI (`mlx-bun generate`) +
-`greedyDecodeBitExact` bypass wired (`src/eval/runner.ts`). **Verified live** against a
-running server (`mlx-bun serve diffusiongemma --port 8899`): OpenAI chat (stream +
-non-stream), Anthropic `/v1/messages`, `/v1/models`, single- AND multi-block — all coherent
-("The capital of France is Paris.", correct primary-colors answer). Gate:
-`tests/diffusion-serving.test.ts` (4/4, `MLX_BUN_TEST_DIFFUSION=1`). No AR regression
-(instanceof-guarded branches; typecheck baseline preserved).
-- **Streaming UX:** v1 runs the engine to completion then yields tokens (the existing SSE
-  machinery emits them as deltas). True per-block intra-stream + temperature>0 (categorical)
-  are noted follow-ups.
-- **REMAINING (image-text-to-text):** the 27-layer SigLIP tower IS in the checkpoint
-  (`parseSiglipConfig` handles it) but the diffusion vision path is NOT wired yet — needs the
-  encoder vision merge (`_embed_inputs` masked_scatter + `_vision_block_overlay` bidirectional
-  overlay), the `<|image|>`→`boi+image_token*N+eoi` splice + `mm_token_type_ids`, the SigLIP
-  weight-name adaptation (`.linear` suffix, `patch_embedder`), and an image-text-to-text golden.
-
-### D3 — Serving + CLI + image input — original brief `[~]`
-Wire DiffusionGemma into `generate.ts`/`cli.ts`/`server.ts` as a model that
-routes to the diffusion engine instead of the AR loop. Decide streaming
-semantics (canvas isn't left-to-right — stream per denoise-step or on
-stabilization). Image input through the existing vision path. Onboarding/registry
-entries; supported-model lists.
-**Exit:** `mlx-bun serve` + `mlx-bun generate` answer text **and** image prompts;
-output matches optiq `generate`; HTTP test green; no regression to AR models.
-
-### D4 — Performance (L3, optional, measured) `[ ]`
-`confidence-threshold` is already the fast path (optiq: ~58 tok/s code / ~9
-prose). Measure mlx-bun tok/s with `scripts/bench-serve.ts all`; fuse kernels only if a number
-says so. KV holds only the prompt (no per-step growth → no KV-quant needed).
-**Exit:** quotable tok/s within range of optiq on this machine; any fused kernel
-gated by KL + the 6-task quality suite (no bit-exact ancestor for novel fusions).
-
-### D5 — Diffusion-native LoRA `[x]` (done 2026-06-24 — trains, loss decreases)
-**DONE.** `src/train/diffusion-lora.ts` (`trainDiffusionLora` + `diffusionLoss`): port of
-optiq's denoising-objective LoRA (NOT AR cross-entropy). Corrupt the target canvas to a
-random noise level t∈[t_min,t_max] (uniform-random token per position with prob t), one
-encoder+decoder forward, CE on the CORRUPTED positions. LoRA mounts on the decoder blocks
-(DEFAULT_LORA_KEYS = attn q/k/v/o + dense MLP gate/up/down; the `loraTargets` map skips
-full-layer v_proj), reusing the existing training stack (`buildTrainableLora`/`ValueAndGrad`/
-`AdamW`/`saveAdapter`). **The MoE backward works** by `stop_gradient`-ing the router top-k
-indices (identity in the forward → D1/D2 stay bit-exact; lets the gather_qmm vjp differentiate
-only the activations). Autograd flows end-to-end through BOTH passes (encoder cache sliceUpdate
-+ decoder cross-attention concat + MoE). **Verified** (`tests/diffusion-lora.test.ts`): loss
-3.91 → 1.20 over 30 iters on one example, all finite, and the trained adapter changes the canvas
-logits (maxDiff 22.4). Mountable via the existing AdapterManager (`lora_a`/`lora_b` layout).
-
-### D5 — Diffusion-native LoRA — original brief `[x]`
-Port `optiq/vlm/diffusion_gemma/lora.py` (the denoising objective, not AR
-cross-entropy) into the training stack. Deferred unless prioritized.
-**Exit:** a LoRA trains end-to-end on `{prompt, completion}` data, loss
-decreases, adapter mounts and changes generation.
-
-## Risks / open questions (resolve in D0)
-
-- **Bidirectional masks** are the highest-risk port — block-diffusion attention
-  is not causal; `_make_decoder_masks` + `_vision_block_overlay` define it.
-- **"Canon depthwise-conv layers"** (per docs) — no conv tensors appeared in the
-  quant config scan; confirm whether this checkpoint has them and where.
-- **Self-conditioning ordering** in the step loop (when the previous prediction
-  is fed back) must match exactly or generation diverges.
-- **mask_token_id** + canvas init RNG must match optiq for deterministic parity.
-- **Vision tower delta**: 27 layers here vs 16 in the e4b siglip — confirm config
-  fields, pooling, and whether the same `SiglipVisionTower` covers it.
-- **Streaming UX**: define how a canvas maps to token-stream SSE (D3 decision).
-
-## Scope (confirmed with Josh 2026-06-24)
-
-- **Text + image together.** The 27-layer vision merge is built into D1 from the
-  start; v1 ships image-text-to-text. (Heavier first-parity, but the model is
-  natively multimodal and that's the target.)
-- **D5 LoRA is IN scope** — port optiq's denoising-objective `train_diffusion_lora`
-  after the inference path (D1–D4) is green.
-- **Generation parity is the bar**, not serving throughput — D4 is optional/L3.
-
-## Orchestration (sub-agents / workflows)
-
-- **D0** = one recon **workflow**: parallel Explore/reader agents, one per
-  reference file group — (a) `language.py`, (b) `diffusion_gemma.py`+`config.py`,
-  (c) `generate/diffusion.py`, (d) `processing`+tokenizer+chat_template, (e)
-  `loader`/`generate`/`lora`/`convert`. Each returns a structured spec; a
-  synthesis agent merges into the "Reference dossier" section here; a
-  weight-audit agent diffs module names vs the real `*.safetensors` index.
-- **D1** = pipeline per component (vision → decoder-layer → MoE →
-  self-conditioning → full-forward), each stage verify-gated against a Python
-  golden dump before the next builds on it.
-- **D2/D3** = mostly serial (single engine), with an adversarial parity-verify
-  agent per sampler.
-
----
-
-## Reference dossier (D0)
-
-This dossier consolidates six independent reverse-engineering passes over OptiQ's vendored DiffusionGemma (`optiq.vlm._mlxvlm.models.diffusion_gemma.*`, `optiq.vlm.diffusion_gemma.*`). It is the porting contract for the mlx-bun TypeScript/MLX implementation. Class/function/tensor names and constants are quoted as found; items the reports flagged as unverified are called out in §9.
-
-### 1. Weight schema & quant map
-
-**Checkpoint analyzed:** `DiffusionGemma-26B-A4B-it OptiQ-4bit`, ~14 GB safetensors (mixed 4/8-bit; BF16 full-precision ≈ 26–28 GB). OptiQ affine quant, `group_size=64` throughout. Sharded `model*.safetensors`, loaded lexicographically via `mx.load()` and kept lazy until eval.
-
-**Top-level config (`TextConfig` / `ModelConfig`):**
-
-| Param | Value |
-|---|---|
-| `model_type` (text / top) | `diffusion_gemma_text` / `diffusion_gemma` |
-| `vocab_size` | 262144 |
-| `hidden_size` | 2816 |
-| `intermediate_size` (dense MLP) | 2112 |
-| `moe_intermediate_size` (expert) | 704 |
-| `num_hidden_layers` | 30 |
-| `num_attention_heads` | 16 |
-| `num_key_value_heads` | 8 (sliding) |
-| `num_global_key_value_heads` | 2 (full_attention) |
-| `head_dim` | 256 (sliding) |
-| `global_head_dim` | 512 (full_attention) |
-| `hidden_activation` | `gelu_pytorch_tanh` |
-| `rms_norm_eps` | 1e-6 |
-| `max_position_embeddings` | 262144 |
-| `sliding_window` | 1024 |
-| `final_logit_softcapping` | 30.0 |
-| `num_experts` | 128 |
-| `top_k_experts` | 8 |
-| `attention_bias` | False |
-| `tie_word_embeddings` | True |
-| `use_bidirectional_attention` | `"vision"` |
-| `pad/eos/bos_token_id` | 0 / 1 / 2 |
-| `canvas_length` (ModelConfig) | 256 |
-| `boi_token_id` / `eoi_token_id` / `image_token_id` | 255999 / 258882 / 258880 |
-| `eos_token_id` (ModelConfig, list) | `[1, 106]` |
-
-**Unique tensor templates (22) and per-tensor quant:**
-
-Decoder embeddings:
-- `model.decoder.embed_tokens.weight` — (262144, 2816), **4-bit**. Tied LM head (no `lm_head.weight`; logits via `embed_tokens.as_linear()`).
-
-Per decoder layer `model.decoder.layers.N` (N=0..29):
-- `self_attn.q_proj` / `k_proj` / `o_proj` — **8-bit** (all 30 layers).
-- `self_attn.v_proj` — **8-bit**, present in **25 of 30** layers (absent in full_attention layers 5,11,17,23,29 — those reuse k as v).
-- `self_attn.q_norm.weight` / `k_norm.weight` — unquantized RMSNorm; `v_norm` — unquantized RMSNormNoScale.
-- `input_layernorm`, `post_attention_layernorm`, `pre_feedforward_layernorm`, `post_feedforward_layernorm`, `post_feedforward_layernorm_1`, `pre_feedforward_layernorm_2`, `post_feedforward_layernorm_2` — all unquantized RMSNorm (7 norms/layer).
-- `mlp.gate_proj` / `up_proj` / `down_proj` — **8-bit**.
-- `router.proj` — **8-bit** (Linear hidden→128, bias=False); `router.scale` (2816,), `router.per_expert_scale` (128,) — unquantized, init `mx.ones`.
-- `experts.gate_up_proj.weight` — (128, 1408, 2816) **4-bit** SwitchLinear; `experts.down_proj.weight` — (128, 2816, 704) **4-bit** SwitchLinear.
-- `layer_scalar` — (1,) unquantized, init `mx.ones`.
-
-Self-conditioning (`model.decoder.self_conditioning`):
-- `pre_norm` (RMSNorm) unquantized; `post_norm` (RMSNormNoScale) unquantized; `gate_proj` / `up_proj` / `down_proj` — **4-bit** (shapes mirror dense MLP: gate/up = (2112,2816), down = (2816,2112)).
-
-Final: `model.decoder.norm.weight` — unquantized RMSNorm.
-
-Encoder layer scalars: `model.encoder.language_model.layers.N.layer_scalar` — (1,) unquantized. (All other encoder *text* weights are skipped at load — the encoder text path reuses decoder weights via weakref.)
-
-Vision tower (SigLIP-style, 27 layers, hidden 1152, patch 16, 16 heads), all 4-bit projections, unquantized norms:
-- `model.encoder.vision_tower.patch_embedder.input_proj` — 4-bit.
-- `model.encoder.embed_vision.embedding_projection` — 4-bit.
-- `model.encoder.vision_tower.encoder.layers.N.self_attn.{q,k,v,o}_proj.linear` — 4-bit (all 27).
-- `model.encoder.vision_tower.encoder.layers.N.mlp.{gate,up}_proj.linear` — 4-bit; vision norms unquantized.
-
-**Quant predicate** (`LanguageModel.quant_predicate`, applied when no explicit per-module override in `config["quantization"]`):
-```python
-def predicate(path, m):
-    if not hasattr(m, "to_quantized"): return False
-    if "router" in path or path.endswith(("mlp.gate_proj","mlp.up_proj","mlp.down_proj")):
-        return {"group_size": 64, "bits": 8}
-    return True   # default → 4-bit, group_size 64
-```
-Loader-level `class_predicate` additionally returns False unless `f"{path}.scales" in weights` and `weight.size % 64 == 0`, and honors explicit `config["quantization"][path]` overrides first. Net effect: attention q/k/v/o + dense MLP + router = 8-bit; embeddings, experts, self-conditioning, vision = 4-bit.
-
-**`sanitize()` remapping:** skip `*rotary_emb*` and `lm_head.weight` (RoPE computed at runtime, head tied); skip `model.encoder.language_model.*` except `*.layer_scalar`; keep vision keys only if vision tower present; rename `*.experts.down_proj` → `*.experts.down_proj.weight` and `*.experts.gate_up_proj` → `*.experts.gate_up_proj.weight`.
-
-**Confirmed absences:** no depthwise/Canon conv tensors anywhere (pure transformer); no separate `lm_head`; no router bias.
-
-### 2. Decoder architecture (layer-by-layer)
-
-**`DecoderModel`:** `embed_tokens` (Embedding 262144→2816), `embed_scale = hidden_size**0.5` (≈53.07), 30× `DecoderLayer`, final `norm` (RMSNorm), `self_conditioning`.
-
-**`layer_types` pattern:** `[sliding_attention ×5, full_attention] ×5`, last layer forced full. Full-attention layers at indices **5, 11, 17, 23, 29**.
-
-**RoPE per layer type** (`initialize_rope`, `traditional=False`):
-- sliding_attention: `rope_theta = 10000.0`, default.
-- full_attention: `rope_theta = 1000000.0`, proportional scaling, `partial_rotary_factor = 0.25` (only 25% of `head_dim` rotated).
-
-**`Attention.__call__(x, mask, cache, *, decoder=False, offset=None)`** (`scale = 1.0`, **no attention softcap**):
-- `head_dim` = 512 if full_attention else 256; `n_kv_heads` = 2 if full else 8.
-- q = `q_norm(q_proj(x))` reshaped (B,L,n_heads,head_dim) → transpose → `rope(offset)`.
-- k = `k_norm(k_proj(x))` reshaped (B,L,n_kv_heads,head_dim) → transpose → `rope(offset)`.
-- v: sliding → `v_norm(v_proj(x))`; full → `v_norm(keys)` (v_proj is None, reuses k). **QK/V norm is post-projection, pre-RoPE.**
-- **decoder=True (canvas pass):** fetch `(encoder_keys, encoder_values)` from cache state; for sliding layers, if `offset >= encoder_len` and `encoder_len > window` (`window = sliding_window-1 = 1023`), slice encoder K/V to last 1023 and trim mask to `[..., -(window+L):]`; concat `[encoder_k, k]`, `[encoder_v, v]` on axis 2; `attn_cache=None` (no cache update on canvas pass).
-- **decoder=False (encoder/prefill):** `cache.update_and_fetch(k,v)`; `attn_cache=cache`.
-- `scaled_dot_product_attention(q,k,v,cache=attn_cache,scale=1.0,mask=mask)` → reshape → `o_proj`.
-
-**`DecoderLayer.__call__(x, mask, cache, *, decoder, offset, layer_scalar=None)`** — attention branch + **parallel dense-MLP + MoE** feedforward:
-```
-# attention
-residual = x
-h = input_layernorm(x); h = self_attn(h, mask, cache, decoder=decoder, offset=offset)
-h = post_attention_layernorm(h); h = residual + h
-# parallel FFN
-residual = h
-h1 = post_feedforward_layernorm_1(mlp(pre_feedforward_layernorm(h)))     # dense
-flat = residual.reshape(-1, hidden); idx, w = router(flat)
-h2 = experts(pre_feedforward_layernorm_2(flat), idx, w).reshape(residual.shape)
-h2 = post_feedforward_layernorm_2(h2)
-h = post_feedforward_layernorm(h1 + h2); h = residual + h
-return h * (layer_scalar if layer_scalar is not None else self.layer_scalar)
-```
-
-**`MLP`:** `down_proj(geglu(gate_proj(x), up_proj(x)))`, where `geglu(gate,x) = gelu_approx(gate) * x` (`@mx.compile(shapeless=True)`).
-
-**`Router`:** `x = rms_norm(x, None, eps); x = x * scale * hidden_size**-0.5; scores = proj(x)`; top-8 via `argpartition(scores, kth=-8)[...,-8:]`; `weights = softmax(take_along_axis(scores, idx), precise=True) * per_expert_scale[idx]`. Normalization is **pre**-projection.
-
-**`Experts`:** `gate_up_proj` / `down_proj` are `SwitchLinear` (weight (num_experts, out, in); `gather_mm(x, W.T, rhs_indices=idx, sorted_indices)`). `x = expand_dims(x,(-2,-3))`; if `idx.size>=64` → `_gather_sort` (group tokens by expert) then `_scatter_unsort`; gate_up split at 704; `y = down_proj(geglu(gate,up))`; final `(y * w[...,None]).sum(axis=-2)`.
-
-**`_make_decoder_masks(h, caches, decoder_attention_mask)`** — returns `{layer_type → mask}`, broadcast shape `(B, 1, canvas_length, key_len)`, `key_len = encoder_len + canvas_length`. `valid_encoder_len = min(_cache_offset(cache), encoder_len)`.
-- **full_attention:** canvas attends bidirectionally to all valid encoder positions + itself. If `encoder_len == valid_encoder_len` and no external mask → mask is `None` (all valid). Otherwise build a key-row: `concat([arange(encoder_len) < valid_encoder_len, ones(canvas_length, bool)])`, broadcast to `(B,1,canvas_length,key_len)`. External masks are sliced to `[..., -key_len:]` then broadcast.
-- **sliding_attention:** `window_prefix = max(sliding_window-1, 0) = 1023`. If `encoder_len == valid_encoder_len and encoder_len <= window_prefix` → `None`. Else keep encoder positions in `[max(0, valid_encoder_len - 1023), valid_encoder_len)`: `encoder_mask = (arange(encoder_len) >= start) & (arange < valid_encoder_len)`; `row = concat([encoder_mask, ones(canvas_length, bool)])`. With external mask, AND it elementwise with `keep` before broadcasting.
-
-Key invariant: **canvas tokens always see all kept encoder positions and are bidirectional among themselves (no causal mask on the canvas)** — diffusion decoding, not autoregressive. (Encoder pass `_make_encoder_masks` is the causal/sliding + vision-overlay one; see §4.)
-
-`DecoderModel.__call__(canvas_ids, cache, self_conditioning_logits, self_conditioning_embeddings, decoder_attention_mask)`: embed canvas (§3) → `masks = _make_decoder_masks(...)` → `offset = _cache_offset(cache[0])` → loop layers with `decoder=True, offset=offset` → `norm(h)`.
-
-### 3. Self-conditioning
-
-`SelfConditioning` module: `pre_norm` (RMSNorm), `gate_proj`/`up_proj`/`down_proj` (dense-MLP shapes), `post_norm` (RMSNormNoScale).
-```python
-def __call__(inputs_embeds, signal):
-    normed = pre_norm(signal)
-    s = down_proj(geglu(gate_proj(normed), up_proj(normed)))
-    return post_norm(inputs_embeds + s)
-```
-
-`DecoderModel._embed_canvas(canvas_ids, sc_logits=None, sc_embeddings=None)`:
-- Hard: `inputs_embeds = embed_tokens(canvas_ids) * embed_scale`.
-- Soft signal: if `sc_embeddings` given → cast to dtype; elif `sc_logits is None` → `zeros_like`; else `probs = softmax(sc_logits, precise=True)`, then `soft = probs @ embed_tokens.weight` (or `mx.quantized_matmul(probs, weight, scales, biases, transpose=False, group_size, bits)` if `QuantizedEmbedding`), `* embed_scale`.
-- Return `self_conditioning(inputs_embeds, soft)`.
-
-In the engine, after each denoising step except the last, soft embeddings are computed from the current logits (`_diffusion_entropy_and_soft_embeddings`) and fed as `self_conditioning_embeddings` to the next step.
-
-### 4. Vision tower & merge
-
-**`EncoderModel`** (holds `weakref` to decoder): `language_model` (`EncoderLanguageModel` — weakref + 30 `EncoderLayerScalar`, each `layer_scalar = mx.ones((1,))`), `vision_tower` (`VisionModel`, optional), `embed_vision` (`MultimodalEmbedder`, vision_hidden→text_hidden, eps=vision rms_norm_eps).
-
-**`make_cache(max_size=None)`** (30 entries by `layer_types`): `StaticPrefixKVCache(max_size)` if max_size given; else `KVCache()` (unbounded) for full_attention; else `RotatingKVCache(max_size=sliding_window=1024)` for sliding.
-
-**Patch embed** (`VisionPatchEmbedder`): reshape pixels to (B,C,pH,p,pW,p) → transpose (B,pH,pW,p,p,C) → reshape (B, pH·pW, C·p·p) → `2*(patches-0.5)` (normalize to [-1,1]) → `input_proj`. `get_image_features = embed_vision(vision_tower(pixel_values))`.
-
-**`_embed_inputs(input_ids, pixel_values)`:** `vision_mask = (ids==image_token_id) | (ids==video_token_id?)`; `llm_input_ids = where(vision_mask, pad_token_id, input_ids)`; `inputs_embeds = embed_tokens(llm_input_ids) * embed_scale`; if pixel_values: `masked_scatter(inputs_embeds, vision_mask_expanded, get_image_features(...))`.
-
-**`_vision_block_overlay(mm_token_type_ids, seq_len)`:** only when `use_bidirectional_attention == "vision"`. `is_vision = (mmtt==1)|(mmtt==2)`; contiguous-block IDs via `cumsum(starts)`; returns `(q_blocks != -1) & (q_blocks == k_blocks)` — tokens in the same image block attend bidirectionally.
-
-**`_make_encoder_masks(h, cache, attention_mask, mm_token_type_ids)`:** fast path uses `create_attention_mask(window_size=sliding_window if sliding else None)`. Slow path: base causal `q_pos >= k_pos`; AND sliding window `q_pos < k_pos + sliding_window` for sliding layers; OR vision overlay; AND key/padding mask. Overlay applied only in prompt (`key_len == N`), not continuation.
-
-**`EncoderModel.__call__`:** embed → masks → 30 layers with `decoder=False` and external `layer_scalar = language_model.layers[i].layer_scalar` → `decoder.norm(h), cache`.
-
-**Image preprocessing (`Gemma4ImageProcessor`):** `patch_size=16`, `max_soft_tokens=280` (allowed 70/140/280/560/1120), `pooling_kernel_size=3`, resize divisor `48 = pooling·patch`, aspect-preserving bicubic resize to multiples of 48 within `max_patches = max_soft_tokens·9` budget; rescale `1/255`; `do_normalize=False` by default (mean/std 0.5 if enabled); channel-first (C,H,W). `num_soft_tokens = (H//16)·(W//16) // 9`.
-
-### 5. Tokenizer / chat template / canvas
-
-**Special tokens** (Gemma4 tokenizer): `bos=2 <bos>`, `eos=[1,106]` (`<eos>`, `<turn|>`), `pad=0 <pad>`, `mask=4 <mask>`, `image_token=258880 <|image|>`, `boi=255999 <|image>`, `eoi=258882 <image|>`, `video=258884 <|video|>`, `audio=258881 <|audio|>`. (Token IDs 258880/255999/258882 confirmed against `ModelConfig`; mask=4 and video=258884 flagged in §9.)
-
-**Image splicing** (`Gemma4Processor`, `image_seq_length=280`): build `full_image_sequence = boi + image_token*280 + eoi`; per-image, replace each `<|image|>` placeholder with `boi + image_token*N + eoi` where N = that image's actual soft-token count (fallback to fixed 280). Then tokenize.
-
-**`mm_token_type_ids`** marks tokens: `0`=text/pad, `1`=image (258880), `2`=video (258884), `3`=audio (258881).
-
-**Chat template** (`chat_template.jinja`, Gemma4): turn markers `<|turn>user\n`, `<|turn>model\n`, `<turn|>\n`, optional `<|turn>system\n…<turn|>\n`; generation prompt continues from `<|turn>model\n`, optionally injecting `<|channel>thought\n<channel|>` when thinking enabled.
-
-**Canvas:** a fixed 256-token (`canvas_length`) block that the decoder denoises. Initialized to uniform-random token IDs per block (`mx.random.randint(0, vocab_size, (B, canvas_length))`), **not** the mask token. Block diffusion: multiple canvas blocks generated sequentially until `max_tokens` / EOS.
-
-### 6. Denoising engine pseudocode
-
-`stream_diffusion_generate(model, processor, tokenizer, input_ids, pixel_values=None, attention_mask=None, *, max_tokens, skip_special_token_ids, temperature=0.0, max_denoising_steps=48, diffusion_full_canvas=False, diffusion_min_canvas_length=64, diffusion_max_canvas_length=256, diffusion_static_cache=False, diffusion_sampler="entropy-bound", diffusion_threshold=0.9, diffusion_compile=False, diffusion_show_unmasking=False, diffusion_unmasking_interval=1, mm_token_type_ids=None, prefill_step_size=None)`:
-
-```
-# 1. PREFILL — encoder builds KV cache (optionally chunked by prefill_step_size)
-encoder(input_ids, attention_mask, cache, pixel_values, mm_token_type_ids)
-
-# 2. CANVAS LOOP — block diffusion
-while generated_tokens < max_new_tokens:
-  canvas_length = min(diffusion_max_canvas_length, remaining, model.config.canvas_length)
-  current_canvas = randint(0, vocab_size, (B, canvas_length)).astype(input_ids.dtype)
-  sc_embeddings = None
-  diffusion_history = []           # ring buffer for stability check
-  draft_reveal_mask = zeros(bool)  # confidence sampler
-  draft_canvas = current_canvas
-
-  # 3. DENOISING LOOP — high t-step → 1
-  for cur_step in reversed(range(1, max_denoising_steps + 1)):
-    logits = decoder_logits_without_sc(current_canvas) if sc_embeddings is None \
-             else decoder_logits_with_sc(current_canvas, sc_embeddings)
-    # logits already softcapped: tanh(x_fp32 / 30.0) * 30.0
-
-    sched_t = _diffusion_linear_temperature(cur_step, max_denoising_steps, schedule_cfg)
-    processed = logits / sched_t if sched_t is not None else logits   # t_min .4 → t_max .8
-
-    argmax_canvas = argmax(processed, -1)
-    if cur_step == 1 and not diffusion_show_unmasking: break          # skip final re-noise
-    denoiser_canvas = argmax_canvas if temperature <= 0 \
-                      else categorical(processed / temperature)
-
-    # 4. SAMPLER → acceptance_mask + next sc_embeddings   (see §7)
-    ... acceptance_mask, sc_embeddings(next) ...
-
-    # 5. FREEZE accepted, RE-NOISE the rest
-    #   entropy:    current_canvas = where(acceptance_mask, accepted, randint(...))
-    #   confidence: current_canvas = where(draft_reveal_mask | acceptance_mask, accepted, randint(...))
-
-    # 6. EARLY STOPS
-    if confidence sampler and cur_step>1 and all(draft_reveal_mask):
-        accepted_canvas = draft_canvas; break
-    if _diffusion_stable_and_confident(argmax_canvas, processed, history, stopping_cfg): break
-
-  # 7. EMIT tokens from final argmax_canvas; check tokenizer.stopping_criteria(token) for EOS
-```
-
-`_diffusion_stable_and_confident`: stable if last `stability_threshold` (default 1) canvases are identical; confident if `mean(token_entropy) < confidence_threshold` (default 0.005); returns `stable and confident`. (Distinct from the sampler's `entropy_bound=0.1`.)
-
-`_diffusion_linear_temperature`: `t_min + (t_max - t_min)*(cur_step/max_steps)`, defaults `t_min=0.4, t_max=0.8`; **divides** logits.
-
-### 7. The two samplers
-
-**Entropy-bound (default in engine):** `entropy_bound = 0.1`. Per step:
-- `cur_step>1`: `token_entropy, next_sc_embeddings = _diffusion_entropy_and_soft_embeddings(processed, soft_embedding_weight, embed_scale)`; else `token_entropy = _diffusion_token_entropy(processed)`, `next_sc=None`.
-- `_diffusion_entropy_transfer_mask(entropy, entropy_bound)`: sort positions by ascending entropy; `cum = cumsum(sorted_entropy)`, `cmax = cummax(sorted_entropy)`; `sorted_sel = (cum - cmax) <= entropy_bound`; unsort via `put_along_axis`. Accepts the lowest-entropy positions until cumulative-minus-running-max exceeds the bound.
-- `accepted = where(mask, denoiser_canvas, current_canvas)`; non-accepted re-noised next step; `draft_reveal_mask = mask`, `draft_canvas = argmax_canvas`.
-
-**Confidence-threshold:** `diffusion_threshold = 0.9` (CLI `args.threshold`, else 0.9). Per step:
-- `confidence = _diffusion_token_probability(processed, denoiser_canvas)` = `exp(token_logits - logsumexp(logits))` in fp32.
-- `_diffusion_confidence_transfer_mask(confidence, unrevealed_mask, threshold, force_all=(cur_step==1))`: if `force_all` → return `unrevealed_mask` (accept all on final step); else `transfer = unrevealed & (confidence >= threshold)`; for sequences with no qualifier, force the single highest-confidence unrevealed position (`argmax(where(unrevealed, confidence, -inf))`); return `transfer | forced`.
-- `accepted = where(mask, denoiser_canvas, draft_canvas)`; `current_canvas = where(draft_reveal_mask | mask, accepted, randint(...))`; `draft_reveal_mask |= mask`; `draft_canvas = where(mask, accepted, draft_canvas)`.
-
-Defaults discrepancy across reports: the engine signature default is `diffusion_sampler="entropy-bound"`, while the public `generate()` (`generate.py`) default is `DEFAULT_SAMPLER="confidence-threshold"`. The port should treat the public-API default as confidence-threshold and the engine internal default as entropy-bound.
-
-`_diffusion_sample_canvas(processed, dtype, temperature)`: fp32; `temperature<=0` → argmax; else `categorical(logits / temperature)`.
-
-### 8. OptiQ load / generate / lora API
-
-**`load_diffusion_gemma(...) → (model, config_dict)`** (`loader.py:55–96`): `load_config()` → `Model(ModelConfig.from_dict(config))` → `_gather_weights()` (lazy mmap) → `model.sanitize(weights)` → quantize via `nn.quantize(group_size, bits, mode, class_predicate)` (predicate honors explicit `config["quantization"][path]`, then `to_quantized` support, `weight.size % 64 == 0`, and `f"{path}.scales" in weights`) → `model.load_weights(...)` → `mx.eval` → `model.eval()`.
-
-**`generate(model, tokenizer, prompt, *, images=None, max_tokens=512, temperature=0.0, sampler="confidence-threshold", verbose=False) → str`** (`generate.py:169`): wraps `stream_generate()`, concatenates non-draft segments (`is_draft=False`), retries ≤3× if output collapses to empty.
-
-**`stream_generate()` → Generator[`GenerationResult`]** fields: `text`, `is_draft`, `draft_text`, `diffusion_step`, `diffusion_total_steps`, `diffusion_canvas_index`, plus metrics `prompt_tokens/generation_tokens/total_tokens`, `prompt_tps/generation_tps`, `diffusion_canvas_tokens`, `diffusion_denoising_steps`, `diffusion_work_tokens` (= canvas_tokens·steps), `finish_reason` ∈ {`stop`,`length`,`stable`}.
-
-**`_build_image_inputs(model, tokenizer, prompt, images)`** → `(input_ids, pixel_values, mm_token_type_ids)`: SigLIP preprocess + per-image `<|image|>` → `boi + image_token*N + eoi` expansion; `mm_token_type_ids = 1` where `t==image_token_id`.
-
-**LoRA** (`lora.py`):
-- `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"]` — applied to all decoder layers via `linear_to_lora_layers(model, len(model.layers), cfg)`. MoE experts and router stay frozen; encoder reuses decoder weights.
-- `diffusion_loss(model, prompt_ids, target_ids, *, t_min, t_max, vocab_size)`: sample `t = uniform(t_min, t_max)`; `corrupt = uniform((B,L)) < t`; `noise = randint(0, vocab_size, (B,L))`; `canvas = where(corrupt, noise, target_ids)`; `logits = model(input_ids=prompt_ids, canvas_ids=canvas).logits`; `ce = cross_entropy(logits, target_ids, reduction="none")`; return `(ce * corrupt).sum() / max(corrupt.sum(), 1)` — **CE on corrupted positions only**.
-- `train_diffusion_lora(model_path, data_dir, adapter_path, *, rank=8, scale=8.0, dropout=0.0, num_layers=-1, iters=200, learning_rate=1e-4, max_canvas=256, report_every=10, seed=0)`. Data: `{data_dir}/train.jsonl` with `{prompt,completion}` | `{messages:[...]}` | `{text}`. Scale 8.0 (vs AR default 20 which collapses).
-- Adapter out: `adapter_config.json` (`base_model`, `fine_tune_type="lora"`, `model_type="diffusion_gemma"`, `lora_parameters`{rank,scale,dropout,keys}) + `adapters.safetensors`.
-
-**Defaults table:**
-
-| Param | Default | Source |
-|---|---|---|
-| `DEFAULT_SAMPLER` | `confidence-threshold` | generate.py:26 |
-| `DEFAULT_MAX_TOKENS` | 512 | generate.py:27 |
-| `DEFAULT_DIFFUSION_MAX_DENOISING_STEPS` | 48 | diffusion.py:21 |
-| `DEFAULT_DIFFUSION_MIN_CANVAS_LENGTH` | 64 | diffusion.py:20 |
-| canvas_length | 256 | ModelConfig |
-| temperature schedule t_min / t_max | 0.4 / 0.8 | diffusion.py / generation_config.json |
-| entropy_bound (sampler) | 0.1 | diffusion.py:662 |
-| confidence threshold (sampler) | 0.9 | diffusion.py:597 |
-| stability_threshold / confidence_threshold (stop) | 1 / 0.005 | generation_config.json |
-| LoRA rank / scale | 8 / 8.0 | lora.py:99–100 |
-
-If `sampler_config._cls_name != "EntropyBoundSamplerConfig"`, the engine raises `NotImplementedError`.
-
-### 9. Open questions / risks for the TS port
-
-**Confirmed (low risk), called out because the brief asked:**
-- **Canon / depthwise conv: ABSENT.** No Conv tensors in the checkpoint; pure transformer (vision tower is attention-based SigLIP). Port needs no conv path.
-- **Tied head: YES.** No `lm_head.weight`; logits via `embed_tokens.as_linear()` (a `QuantizedEmbedding`). The port must implement quantized `as_linear` (transpose-matmul against the 4-bit embedding table) and the matching `quantized_matmul(probs, weight, scales, biases, transpose=False, group_size, bits)` path used for soft embeddings in `_embed_canvas`.
-
-**Unclear / verify before/while porting:**
-1. **`mask_token_id` value — UNCERTAIN.** `<mask>`=4 in `tokenizer_config.json`, but no engine path uses a mask token: the canvas is initialized to **uniform-random token IDs** (`randint(0, vocab_size, …)`), and re-noising also uses random IDs. Treat canvas init/re-noise as uniform-random; don't assume a discrete mask token unless a code path consuming id 4 is found.
-2. **`video_token_id`.** `ModelConfig.video_token_id = None` in the checkpoint, but tokenizer lists 258884 and `mm_token_type_ids` uses `==2` for video. Gate the video branch on a non-None config value; confirm whether this checkpoint supports video.
-3. **Default sampler conflict.** `generate.py` default = `confidence-threshold`; engine default = `entropy-bound`; `generation_config.json` ships `EntropyBoundSamplerConfig`. Public `generate()` → confidence-threshold; low-level engine → entropy-bound.
-4. **`eos_token_id` is a list `[1, 106]`** at ModelConfig level (106 = `<turn|>`) but `1` at TextConfig level. Stopping criteria must accept the full set.
-5. **Mask broadcasting precision in `_make_decoder_masks`.** Highest bug-risk area. The interaction between in-attention encoder K/V slicing (last 1023) and the precomputed mask trim (`[..., -(window+L):]`) must be replicated exactly. Build a per-op parity harness (copy-verbatim methodology) before optimizing.
-6. **`RotatingKVCache` (ring, max_size=1024) vs `KVCache` (unbounded) vs `StaticPrefixKVCache`.** `_cache_offset` / `_cache_state` helpers must be ported faithfully; the decoder reads encoder K/V out of cache state — offset/order bugs corrupt cross-attention silently.
-7. **`SwitchLinear` + `_gather_sort`/`_scatter_unsort`** with the `idx.size >= 64` sort threshold and `sorted_indices` flag. Match `mlx_lm.models.switch_layers` exactly; 4-bit quant amplifies tiny mismatches.
-8. **Quant predicate ordering.** Loader honors explicit `config["quantization"][path]` overrides before the predicate's blanket "8-bit attention." Read the real per-tensor map, not just the predicate.
-9. **`use_clipped_linears` / calibration-scalar gating** in `sanitize()` for vision/embed_vision keys — confirm the exact keep/drop condition against source.
-10. **`partial_rotary_factor=0.25` on full-attention RoPE** (only 25% of the 512-dim head rotated) + `traditional=False` + proportional scaling at theta 1e6 — verify rotary-dim split/interleaving against `initialize_rope` output.
-11. **Image processor budget math** (`max_patches = max_soft_tokens·9`, divisor 48, bicubic, `do_normalize=False`, patch pre-scale `2*(x-0.5)`) and per-image variable soft-token count feeding `<|image|>` expansion — replicate exactly so encoder seq lengths line up with `mm_token_type_ids`.
-12. **`final_logit_softcapping` is fp32**: `tanh(x.astype(float32)/30.0)*30.0` (`@mx.compile(shapeless=True)`); the upcast matters for parity.
-
-Source-line references are approximate locators from the six passes, not verified citations.