| """MLX implementation of the DiffusionGemma text model (encoder + block-diffusion decoder). |
| |
| Shape conventions: activations are [B, L, hidden]; attention is [B, heads, L, head_dim]; the KV cache keeps |
| [B, heads, S, head_dim] per layer. Weight names follow the HF checkpoint (``model.decoder.*``), which is also |
| where the encoder's weights live (in the checkpoint the encoder shares them). A linear weight is either a float |
| array ``X.weight`` or NVFP4-packed ``X.weight`` (uint32, 8 nibbles per word) + ``X.scales`` (uint8, group 16). |
| Batch size is 1: the port serves one request at a time, so no padding/batch masks are needed. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from typing import Any |
|
|
| import mlx.core as mx |
|
|
| from .config import DiffusionGemmaConfig, TextConfig |
|
|
| NVFP4 = dict(group_size=16, bits=4, mode="nvfp4") |
|
|
|
|
| |
|
|
|
|
| def _packed(sd: dict[str, mx.array], name: str) -> tuple[mx.array, mx.array | None, mx.array | None]: |
| """(weight, scales, global_scale) for a checkpoint linear name ending in ``.weight``.""" |
| stem = name[: -len(".weight")] if name.endswith(".weight") else name |
| return sd[name], sd.get(stem + ".scales"), sd.get(stem + ".global_scale") |
|
|
|
|
| class Linear: |
| """``x @ W.T`` for a float or NVFP4-packed weight.""" |
|
|
| def __init__(self, weight: mx.array, scales: mx.array | None = None, global_scale: mx.array | None = None): |
| self.weight = weight |
| self.scales = scales |
| self.global_scale = global_scale |
|
|
| def __call__(self, x: mx.array) -> mx.array: |
| if self.scales is None: |
| return x @ self.weight.T |
| if self.global_scale is not None: |
| |
| |
| raise TypeError("packed Linear with a global scale is not supported") |
| return mx.quantized_matmul(x, self.weight, scales=self.scales, transpose=True, **NVFP4) |
|
|
|
|
| class RMSNorm: |
| """DiffusionGemma RMSNorm: ``x / sqrt(mean(x^2) + eps) * w``, computed in float32.""" |
|
|
| def __init__(self, weight: mx.array | None, eps: float): |
| self.weight = weight |
| self.eps = eps |
|
|
| def __call__(self, x: mx.array) -> mx.array: |
| dtype = x.dtype |
| x = x.astype(mx.float32) |
| x = x * mx.power(mx.mean(x * x, axis=-1, keepdims=True) + self.eps, -0.5) |
| if self.weight is not None: |
| x = x * self.weight.astype(mx.float32) |
| return x.astype(dtype) |
|
|
|
|
| def gelu_tanh(x: mx.array) -> mx.array: |
| return 0.5 * x * (1.0 + mx.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x))) |
|
|
|
|
| |
|
|
|
|
| def rope_tables(text: TextConfig, layer_type: str, positions: mx.array) -> tuple[mx.array, mx.array]: |
| """cos/sin of shape [L, head_dim/2]; both head halves rotate with the same angles (reference layout).""" |
| params = text.rope_parameters[layer_type] |
| head_dim = text.global_head_dim if layer_type == "full_attention" else text.head_dim |
| base = params["rope_theta"] |
| factor = params.get("partial_rotary_factor", 1.0) |
| if params.get("rope_type", "default") == "proportional": |
| |
| angles = int(factor * head_dim // 2) |
| exponents = mx.arange(0, 2 * angles, 2, dtype=mx.float32) / head_dim |
| inv_freq = mx.concatenate( |
| [1.0 / (base**exponents), mx.zeros(head_dim // 2 - angles, dtype=mx.float32)] |
| ) |
| else: |
| dim = int(factor * head_dim) |
| inv_freq = 1.0 / (base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)) |
| freqs = positions.astype(mx.float32)[:, None] * inv_freq[None, :] |
| return mx.cos(freqs), mx.sin(freqs) |
|
|
|
|
| def apply_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array: |
| """x [B, L, H, D], cos/sin [L, D/2].""" |
| half = x.shape[-1] // 2 |
| cos, sin = cos[None, :, None, :], sin[None, :, None, :] |
| x1, x2 = x[..., :half], x[..., half:] |
| return mx.concatenate([x1 * cos - x2 * sin, x1 * sin + x2 * cos], axis=-1) |
|
|
|
|
| class KVCache: |
| """Per-layer prefix cache. Sliding layers keep only ``sliding_window - 1`` tokens, like the reference.""" |
|
|
| def __init__(self, text: TextConfig): |
| n = text.num_hidden_layers |
| self.text = text |
| self.keys: list[mx.array | None] = [None] * n |
| self.values: list[mx.array | None] = [None] * n |
| self.seen = [0] * n |
|
|
| def append(self, i: int, k: mx.array, v: mx.array) -> tuple[mx.array, mx.array]: |
| """Append [B, H, L, D] states; returns all states visible to this forward.""" |
| self.seen[i] += k.shape[2] |
| if self.keys[i] is not None: |
| k = mx.concatenate([self.keys[i], k], axis=2) |
| v = mx.concatenate([self.values[i], v], axis=2) |
| if self.text.is_sliding(i): |
| keep = self.text.sliding_window - 1 |
| self.keys[i] = k[:, :, -keep:, :] if k.shape[2] > keep else k |
| self.values[i] = v[:, :, -keep:, :] if v.shape[2] > keep else v |
| else: |
| self.keys[i], self.values[i] = k, v |
| return k, v |
|
|
| def prefix(self, i: int) -> tuple[mx.array | None, mx.array | None, int]: |
| """Cached K/V and the absolute position of the first cached token.""" |
| if self.keys[i] is None: |
| return None, None, 0 |
| return self.keys[i], self.values[i], self.seen[i] - self.keys[i].shape[2] |
|
|
| def copy(self) -> KVCache: |
| """Shallow copy: ``append`` always replaces entries, so a copy can be extended without touching this one.""" |
| other = KVCache.__new__(KVCache) |
| other.text = self.text |
| other.keys = list(self.keys) |
| other.values = list(self.values) |
| other.seen = list(self.seen) |
| return other |
|
|
|
|
| class Attention: |
| def __init__(self, text: TextConfig, layer_idx: int, sd: dict[str, mx.array], prefix: str): |
| self.layer_idx = layer_idx |
| self.head_dim = text.layer_head_dim(layer_idx) |
| self.num_heads = text.num_attention_heads |
| self.num_kv_heads = text.layer_num_kv_heads(layer_idx) |
| self.scaling = 1.0 |
| self.q_proj = Linear(*_packed(sd, f"{prefix}.q_proj.weight")) |
| self.k_proj = Linear(*_packed(sd, f"{prefix}.k_proj.weight")) |
| v_proj = f"{prefix}.v_proj.weight" |
| |
| self.v_proj = Linear(*_packed(sd, v_proj)) if v_proj in sd else None |
| self.o_proj = Linear(*_packed(sd, f"{prefix}.o_proj.weight")) |
| eps = text.rms_norm_eps |
| self.q_norm = RMSNorm(sd[f"{prefix}.q_norm.weight"], eps) |
| self.k_norm = RMSNorm(sd[f"{prefix}.k_norm.weight"], eps) |
| self.v_norm = RMSNorm(None, eps) |
|
|
| def __call__(self, x, cos, sin, mask, cache: KVCache, write: bool) -> mx.array: |
| b, length, _ = x.shape |
| hd, nh, nkv = self.head_dim, self.num_heads, self.num_kv_heads |
| cos, sin = cos.astype(x.dtype), sin.astype(x.dtype) |
|
|
| q = apply_rope(self.q_norm(self.q_proj(x).reshape(b, length, nh, hd)), cos, sin).transpose(0, 2, 1, 3) |
| k = self.k_proj(x).reshape(b, length, nkv, hd) |
| v = k if self.v_proj is None else self.v_proj(x).reshape(b, length, nkv, hd) |
| k = apply_rope(self.k_norm(k), cos, sin) |
| v = self.v_norm(v) |
| k, v = k.transpose(0, 2, 1, 3), v.transpose(0, 2, 1, 3) |
|
|
| if write: |
| k, v = cache.append(self.layer_idx, k, v) |
| else: |
| pk, pv, _ = cache.prefix(self.layer_idx) |
| if pk is not None: |
| if pk.shape[0] != b: |
| pk = mx.broadcast_to(pk, (b,) + pk.shape[1:]) |
| pv = mx.broadcast_to(pv, (b,) + pv.shape[1:]) |
| k = mx.concatenate([pk, k], axis=2) |
| v = mx.concatenate([pv, v], axis=2) |
|
|
| out = mx.fast.scaled_dot_product_attention( |
| q, k, v, scale=self.scaling, mask=mask |
| ) |
| return self.o_proj(out.transpose(0, 2, 1, 3).reshape(b, length, nh * hd)) |
|
|
|
|
| |
|
|
|
|
| class Router: |
| def __init__(self, text: TextConfig, sd: dict[str, mx.array], prefix: str): |
| self.top_k = text.top_k_experts |
| self.scalar_root_size = text.hidden_size**-0.5 |
| self.norm = RMSNorm(None, text.rms_norm_eps) |
| self.proj = Linear(*_packed(sd, f"{prefix}.proj.weight")) |
| self.scale = sd[f"{prefix}.scale"] |
| self.per_expert_scale = sd[f"{prefix}.per_expert_scale"] |
|
|
| def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: |
| """x [T, hidden] -> (top-k expert ids, top-k weights), both [T, top_k].""" |
| h = self.norm(x) * self.scale * self.scalar_root_size |
| probs = mx.softmax(self.proj(h).astype(mx.float32), axis=-1) |
| idx = mx.argpartition(-probs, kth=self.top_k, axis=-1)[..., : self.top_k] |
| weights = mx.take_along_axis(probs, idx, axis=-1) |
| order = mx.argsort(-weights, axis=-1) |
| idx = mx.take_along_axis(idx, order, axis=-1) |
| weights = mx.take_along_axis(weights, order, axis=-1) |
| weights = weights / weights.sum(axis=-1, keepdims=True) |
| return idx, weights * self.per_expert_scale[idx] |
|
|
|
|
| def _expert_weights(sd: dict[str, mx.array], prefix: str, num_experts: int, kinds: list[str]): |
| """Stacked (weight, scales, global_scale) over experts, one entry per kind concatenated on the output axis. |
| |
| Takes the converted artifact's pre-stacked tensors as they are (``experts.gate_up_proj`` / |
| ``experts.down_proj``, with ``global_scale``) and falls back to stacking the checkpoint's per-expert |
| layout (``experts.{e}.gate_proj`` etc.) for float weights, as the tiny parity tests need. |
| """ |
| base = f"{prefix}.experts" |
| if f"{base}.{kinds[0]}.weight" in sd: |
| return _packed(sd, f"{base}.{kinds[0]}.weight") |
| if f"{base}.{kinds[0]}" in sd: |
| name = f"{base}.{kinds[0]}" |
| if sd[name].dtype in (mx.uint32, mx.uint8): |
| return sd[name], sd.get(name + ".scales"), sd.get(name + ".global_scale") |
| return sd[name], None, None |
| weights, scales = [], [] |
| for e in range(num_experts): |
| parts = [_packed(sd, f"{base}.{e}.{kind}.weight") for kind in kinds] |
| weights.append(mx.concatenate([p[0] for p in parts], axis=0)) |
| scales.append(None if parts[0][1] is None else mx.concatenate([p[1] for p in parts], axis=0)) |
| return mx.stack(weights), None if scales[0] is None else mx.stack(scales), None |
|
|
|
|
| class Experts: |
| """128-expert MLP; gate and up stacked on the output axis, as in the reference's ``gate_up_proj``.""" |
|
|
| def __init__(self, text: TextConfig, sd: dict[str, mx.array], prefix: str): |
| n = text.num_experts |
| kinds = ( |
| ["gate_up_proj"] |
| if f"{prefix}.experts.gate_up_proj.weight" in sd or f"{prefix}.experts.gate_up_proj" in sd |
| else ["gate_proj", "up_proj"] |
| ) |
| self.gate_up, self.gate_up_scales, self.gate_up_gs = _expert_weights(sd, prefix, n, kinds) |
| self.down, self.down_scales, self.down_gs = _expert_weights(sd, prefix, n, ["down_proj"]) |
|
|
| def __call__(self, x: mx.array, idx: mx.array, weights: mx.array) -> mx.array: |
| """x [T, hidden] with idx/weights [T, top_k] -> [T, hidden]. |
| |
| Expert assignments are sorted by expert before the gathers: unsorted nvfp4 gather_qmm is ~9x slower |
| here (measured on this Mac, 9 ms vs 1 ms per call), because the kernel cannot stream the weights. |
| """ |
| t, k = idx.shape |
| flat = idx.reshape(-1) |
| order = mx.argsort(flat) |
| idx_s = mx.take(flat, order) |
| xs = mx.take(x, order // k, axis=0) |
| gate_up = self._mm(xs, self.gate_up, self.gate_up_scales, self.gate_up_gs, idx_s) |
| gate, up = mx.split(gate_up, 2, axis=-1) |
| h = gelu_tanh(gate) * up |
| out = self._mm(h.reshape(t * k, -1), self.down, self.down_scales, self.down_gs, idx_s) |
| out = mx.take(out, mx.argsort(order), axis=0).reshape(t, k, -1) |
| |
| return (out * weights[..., None].astype(out.dtype)).sum(axis=-2) |
|
|
| @staticmethod |
| def _mm(x: mx.array, w: mx.array, scales: mx.array | None, gs: mx.array | None, idx: mx.array) -> mx.array: |
| """x [N, in] with sorted idx [N] expert ids -> [N, out].""" |
| x = mx.expand_dims(x, -2) |
| if scales is None: |
| out = mx.gather_mm(x, w.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=True) |
| else: |
| out = mx.gather_qmm( |
| x, w, scales, rhs_indices=idx, transpose=True, sorted_indices=True, global_scale=gs, **NVFP4 |
| ) |
| return out[:, 0, :] |
|
|
|
|
| |
|
|
|
|
| class Layer: |
| def __init__(self, text: TextConfig, layer_idx: int, sd: dict[str, mx.array], prefix: str): |
| self.self_attn = Attention(text, layer_idx, sd, f"{prefix}.self_attn") |
| self.gate_proj = Linear(*_packed(sd, f"{prefix}.mlp.gate_proj.weight")) |
| self.up_proj = Linear(*_packed(sd, f"{prefix}.mlp.up_proj.weight")) |
| self.down_proj = Linear(*_packed(sd, f"{prefix}.mlp.down_proj.weight")) |
| eps = text.rms_norm_eps |
| self.input_layernorm = RMSNorm(sd[f"{prefix}.input_layernorm.weight"], eps) |
| self.post_attention_layernorm = RMSNorm(sd[f"{prefix}.post_attention_layernorm.weight"], eps) |
| self.pre_feedforward_layernorm = RMSNorm(sd[f"{prefix}.pre_feedforward_layernorm.weight"], eps) |
| self.pre_feedforward_layernorm_2 = RMSNorm(sd[f"{prefix}.pre_feedforward_layernorm_2.weight"], eps) |
| self.post_feedforward_layernorm = RMSNorm(sd[f"{prefix}.post_feedforward_layernorm.weight"], eps) |
| self.post_feedforward_layernorm_1 = RMSNorm(sd[f"{prefix}.post_feedforward_layernorm_1.weight"], eps) |
| self.post_feedforward_layernorm_2 = RMSNorm(sd[f"{prefix}.post_feedforward_layernorm_2.weight"], eps) |
| self.layer_scalar = sd[f"{prefix}.layer_scalar"] |
| self.router = Router(text, sd, f"{prefix}.router") |
| self.experts = Experts(text, sd, prefix) |
|
|
| def __call__(self, x, cos, sin, mask, cache: KVCache, write: bool) -> mx.array: |
| attn = self.self_attn(self.input_layernorm(x), cos, sin, mask, cache, write) |
| residual = x + self.post_attention_layernorm(attn) |
|
|
| h = self.pre_feedforward_layernorm(residual) |
| dense = self.post_feedforward_layernorm_1(self.down_proj(gelu_tanh(self.gate_proj(h)) * self.up_proj(h))) |
|
|
| flat = residual.reshape(-1, residual.shape[-1]) |
| |
| idx, weights = self.router(flat) |
| moe = self.experts(self.pre_feedforward_layernorm_2(flat), idx, weights) |
| moe = self.post_feedforward_layernorm_2(moe.reshape(residual.shape)) |
|
|
| return (residual + self.post_feedforward_layernorm(dense + moe)) * self.layer_scalar |
|
|
|
|
| class SelfConditioning: |
| def __init__(self, text: TextConfig, sd: dict[str, mx.array], prefix: str): |
| eps = text.rms_norm_eps |
| self.pre_norm = RMSNorm(sd[f"{prefix}.pre_norm.weight"], eps) |
| self.post_norm = RMSNorm(None, eps) |
| self.gate_proj = Linear(*_packed(sd, f"{prefix}.gate_proj.weight")) |
| self.up_proj = Linear(*_packed(sd, f"{prefix}.up_proj.weight")) |
| self.down_proj = Linear(*_packed(sd, f"{prefix}.down_proj.weight")) |
|
|
| def __call__(self, inputs_embeds: mx.array, signal: mx.array) -> mx.array: |
| h = self.pre_norm(signal) |
| h = self.down_proj(gelu_tanh(self.gate_proj(h)) * self.up_proj(h)) |
| return self.post_norm(inputs_embeds + h) |
|
|
|
|
| |
|
|
|
|
| class DiffusionGemma: |
| """Text-only DiffusionGemma: an encoder that writes a prefix KV cache, and a diffusion decoder.""" |
|
|
| def __init__(self, config: DiffusionGemmaConfig, state_dict: dict[str, Any]): |
| text = config.text_config |
| self.config = config |
| self.text = text |
| sd = {k: (v if isinstance(v, mx.array) else mx.array(v)) for k, v in state_dict.items()} |
|
|
| self.embed_tokens = sd["model.decoder.embed_tokens.weight"] |
| self.lm_head = sd.get("lm_head.weight", self.embed_tokens) |
| |
| self.embed_scale = mx.array(math.sqrt(text.hidden_size)).astype(self.embed_tokens.dtype) |
|
|
| self.layers = [Layer(text, i, sd, f"model.decoder.layers.{i}") for i in range(text.num_hidden_layers)] |
| self.norm = RMSNorm(sd["model.decoder.norm.weight"], text.rms_norm_eps) |
| self.self_conditioning = SelfConditioning(text, sd, "model.decoder.self_conditioning") |
|
|
| |
|
|
| def _encoder_mask(self, layer_idx: int, positions: mx.array, kv_len: int, kv_pos0: int, block_ids): |
| """Causal (+ sliding window) mask, with bidirectional groups where ``block_ids`` agree.""" |
| text = self.text |
| kv_pos = mx.arange(kv_len) + kv_pos0 |
| allowed = kv_pos[None, :] <= positions[:, None] |
| if text.is_sliding(layer_idx): |
| allowed &= kv_pos[None, :] > (positions - text.sliding_window)[:, None] |
| if block_ids is not None: |
| |
| pad = kv_len - block_ids.shape[0] |
| keys_ids = mx.concatenate([mx.full((pad,), -1, mx.int32), block_ids]) if pad > 0 else block_ids |
| allowed |= (block_ids[:, None] == keys_ids[None, :]) & (block_ids[:, None] >= 0) |
| |
| return mx.where(allowed, 0.0, -mx.inf)[None, None].astype(self.embed_tokens.dtype) |
|
|
| def encode( |
| self, |
| input_ids: mx.array, |
| positions: mx.array, |
| cache: KVCache | None = None, |
| block_ids: mx.array | None = None, |
| ) -> tuple[mx.array, KVCache]: |
| """Encode ``input_ids`` into the prefix cache; returns (last_hidden_state, cache). |
| |
| ``block_ids`` (one int per token) marks groups that attend bidirectionally inside the group -- the |
| per-question blocks of the reference's ``create_masks_for_generate``; ``-1`` is plain causal text. |
| As in the reference, they apply to the prefill only (they are dropped for later canvases). |
| """ |
| text = self.text |
| cache = cache if cache is not None else KVCache(text) |
| if block_ids is not None and cache.seen[0] != 0 and block_ids.shape[0] != input_ids.shape[1]: |
| raise ValueError("block_ids must cover exactly the tokens being encoded") |
| x = self.embed_tokens[input_ids] * self.embed_scale |
| ropes, masks = {}, {} |
| for i, layer in enumerate(self.layers): |
| layer_type = text.layer_types[i] |
| if layer_type not in ropes: |
| ropes[layer_type] = rope_tables(text, layer_type, positions) |
| _, _, kv_pos0 = cache.prefix(i) |
| kv_len = cache.seen[i] - kv_pos0 + input_ids.shape[1] |
| masks[layer_type] = self._encoder_mask(i, positions, kv_len, kv_pos0, block_ids) |
| x = layer(x, *ropes[layer_type], masks[layer_type] if text.is_causal else None, cache, write=True) |
| return self.norm(x), cache |
|
|
| |
|
|
| def denoise( |
| self, |
| canvas_ids: mx.array, |
| cache: KVCache, |
| positions: mx.array, |
| self_conditioning_logits: mx.array | None = None, |
| read_positions: mx.array | None = None, |
| vocab_ids: mx.array | None = None, |
| ) -> mx.array: |
| """One denoiser forward over the canvas; returns softcapped logits [B, canvas, vocab] in float32. |
| |
| ``read_positions``/``vocab_ids`` project only those canvas rows onto those vocabulary ids, which is |
| what a constrained read needs and saves the full-vocabulary unembedding of every canvas position. |
| """ |
| text = self.text |
| x = self.embed_tokens[canvas_ids] * self.embed_scale |
| if self_conditioning_logits is not None: |
| probs = mx.softmax(self_conditioning_logits.astype(mx.float32), axis=-1) |
| soft = (probs.astype(self.embed_tokens.dtype) @ self.embed_tokens) * self.embed_scale |
| else: |
| soft = mx.zeros_like(x) |
| x = self.self_conditioning(x, soft) |
|
|
| ropes = {} |
| for i, layer in enumerate(self.layers): |
| layer_type = text.layer_types[i] |
| if layer_type not in ropes: |
| ropes[layer_type] = rope_tables(text, layer_type, positions) |
| x = layer(x, *ropes[layer_type], None, cache, write=False) |
| h = self.norm(x) |
| if read_positions is not None: |
| h = mx.take(h, read_positions, axis=1) |
| head = self.lm_head if vocab_ids is None else mx.take(self.lm_head, vocab_ids, axis=0) |
| logits = (h @ head.T).astype(mx.float32) |
| return mx.tanh(logits / text.final_logit_softcapping) * text.final_logit_softcapping |
|
|
