diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index 3573b2640..c301e8d33 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -30,9 +30,11 @@ from .models.cache import ( ArraysCache, BatchKVCache, + BatchPoolingCache, BatchRotatingKVCache, CacheList, KVCache, + PoolingCache, QuantizedKVCache, RotatingKVCache, TokenBuffer, @@ -847,6 +849,8 @@ def to_batch_cache(c): elif isinstance(c, ArraysCache): c.left_padding = mx.array(left_padding) return c + elif isinstance(c, PoolingCache): + return BatchPoolingCache(c.ratio, left_padding) elif isinstance(c, RotatingKVCache): if c.keep > 0: raise ValueError("RotatingKVCache with keep tokens is not supported.") diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py index b84c9d650..b3c698ec1 100644 --- a/mlx_lm/models/cache.py +++ b/mlx_lm/models/cache.py @@ -900,6 +900,551 @@ def from_state(cls, state, meta_state): return obj +class PoolingCache(_BaseCache): + """Cache for pooled (compressed) KV tokens with a remainder buffer. + + Stores two things: + 1. A growing pool of compressed tokens (step-allocated). + 2. A small remainder buffer of tokens not yet forming a full window. + """ + + def __init__(self, ratio: int): + self.ratio = ratio + + self.buf_kv = None + self.buf_gate = None + self.remainder = 0 + + self.pooled = None + + @property + def offset(self): + return 0 if self.pooled is None else self.pooled.shape[1] + + def accumulate_windows(self, kv: mx.array, gate: mx.array, offset): + B, L, D1 = kv.shape + _, _, D2 = gate.shape + + if self.buf_kv is None: + self.buf_kv = mx.zeros((B, self.ratio, D1), dtype=kv.dtype) + self.buf_gate = mx.zeros((B, self.ratio, D2), dtype=gate.dtype) + + # Prompt mode + if L > 1: + total = L + self.remainder + usable = (total // self.ratio) * self.ratio + new_remainder = total % self.ratio + + if usable > 0: + r_kv = mx.concatenate( + [ + self.buf_kv[:, : self.remainder], + kv[:, : (usable - self.remainder)], + ], + axis=1, + ) + r_gate = mx.concatenate( + [ + self.buf_gate[:, : self.remainder], + gate[:, : (usable - self.remainder)], + ], + axis=1, + ) + r_base = offset - self.remainder + self.remainder = 0 + else: + r_kv = mx.zeros((B, 0, D1), dtype=kv.dtype) + r_gate = mx.zeros((B, 0, D2), dtype=gate.dtype) + r_base = 0 + + if new_remainder > 0: + self.buf_kv[:, self.remainder : new_remainder] = kv[:, -new_remainder:] + self.buf_gate[:, self.remainder : new_remainder] = gate[ + :, -new_remainder: + ] + self.remainder = new_remainder + + return r_kv, r_gate, r_base + + # Decode mode + else: + self.buf_kv[:, self.remainder : self.remainder + 1] = kv + self.buf_gate[:, self.remainder : self.remainder + 1] = gate + self.remainder = (self.remainder + 1) % self.ratio + + if self.remainder == 0: + r_kv = self.buf_kv + r_gate = self.buf_gate + r_base = offset - self.ratio + 1 + else: + r_kv = mx.zeros((B, 0, D1), dtype=kv.dtype) + r_gate = mx.zeros((B, 0, D2), dtype=gate.dtype) + r_base = 0 + + return r_kv, r_gate, r_base + + def update_and_fetch(self, px: mx.array): + if px.shape[1] == 0: + if self.pooled is None: + return mx.zeros((px.shape[0], 0, px.shape[-1]), dtype=px.dtype) + return self.pooled + + if self.pooled is None: + self.pooled = px + else: + self.pooled = mx.concatenate([self.pooled, px], axis=1) + return self.pooled + + def make_mask(self, L: int = 1, offset: int = 0): + """Build a causal validity mask for pooled positions. + + Query at absolute position ``offset + j`` can attend to pooled token + ``i`` iff ``i < (offset + j) // ratio``. + + Returns ``(N, P)`` bool mask, or ``None`` when every pooled position + is visible to every query (common during decode). + """ + if self.pooled is None or L == 1: + return None + + pool_idx = mx.arange(self.pooled.shape[1]) + query_idx = mx.arange(offset + 1, offset + L + 1) + return pool_idx < query_idx[:, None] // self.ratio + + @property + def state(self): + buf_kv = self.buf_kv[:, : self.remainder] if self.remainder > 0 else None + buf_gate = self.buf_gate[:, : self.remainder] if self.remainder > 0 else None + return (buf_kv, buf_gate, self.pooled) + + @state.setter + def state(self, v): + buf_kv, buf_gate, pooled = v + self.remainder = 0 + self.buf_kv = self.buf_gate = None + if buf_kv is not None: + self.accumulate_windows(buf_kv, buf_gate, 0) + self.pooled = pooled + + @property + def meta_state(self): + return self.ratio + + @meta_state.setter + def meta_state(self, v): + self.ratio = v + + def is_trimmable(self): + return self.pooled is None + + def trim(self, n): + n = min(self.remainder, n) + self.remainder -= n + return n + + def size(self): + return 0 if self.pooled is None else self.pooled.shape[1] + + def empty(self): + return self.pooled is None and self.remainder == 0 + + @property + def nbytes(self): + total = 0 + if self.buf_kv is not None: + total += self.buf_kv.nbytes + self.buf_gate.nbytes + if self.pooled is not None: + total += self.pooled.nbytes + return total + + @classmethod + def merge(cls, caches): + return BatchPoolingCache.merge(caches) + + +class BatchPoolingCache(_BaseCache): + """Batched pooling cache with per-element variable-length tracking.""" + + def __init__(self, ratio: int, left_padding: List[int]): + self.ratio = ratio + + if not all(p == 0 for p in left_padding): + raise RuntimeError("BatchPoolingCache does not support left padding") + + batch_size = len(left_padding) + + self.buf_kv = None + self.buf_gate = None + self.remainder = [0] * batch_size + + self.pooled = None + self._pool_lengths = [0] * batch_size + + self._lengths = [2**31] * batch_size + self._processed = [0] * batch_size + + @property + def offset(self): + return mx.array(self._pool_lengths, dtype=mx.int32) + + def prepare(self, *, lengths=None, right_padding=None, left_padding=None): + if left_padding is not None: + raise RuntimeError("BatchPoolingCache does not support left padding") + if lengths is not None: + self._lengths = [p + l for p, l in zip(self._processed, lengths)] + + def finalize(self): + self._lengths = [2**31] * len(self._pool_lengths) + + def accumulate_windows(self, kv: mx.array, gate: mx.array, offset): + B, L, D1 = kv.shape + _, _, D2 = gate.shape + ratio = self.ratio + + if self.buf_kv is None: + self.buf_kv = mx.zeros((B, ratio, D1), dtype=kv.dtype) + self.buf_gate = mx.zeros((B, ratio, D2), dtype=gate.dtype) + + valid_lengths = [min(l - p, L) for l, p in zip(self._lengths, self._processed)] + if max(valid_lengths) != L: + raise RuntimeError() + for i in range(B): + self._processed[i] += valid_lengths[i] + + totals = [vl + r for vl, r in zip(valid_lengths, self.remainder)] + usable = [(t // ratio) * ratio for t in totals] + max_usable = max(usable) + new_remainder = [t % ratio for t in totals] + + # No sequence produced a full window yet + if max_usable == 0: + for i in range(B): + r = self.remainder[i] + vl = valid_lengths[i] + self.buf_kv[i, r : r + vl] = kv[i, :vl] + self.buf_gate[i, r : r + vl] = gate[i, :vl] + self.remainder = new_remainder + + r_kv = mx.zeros((B, 0, D1), dtype=kv.dtype) + r_gate = mx.zeros((B, 0, D2), dtype=gate.dtype) + r_base = 0 + return r_kv, r_gate, r_base + + # At least one sequence completed a window + r_kv = mx.zeros((B, max_usable, D1), dtype=kv.dtype) + r_gate = mx.zeros((B, max_usable, D2), dtype=gate.dtype) + r_base = [0] * B + + new_buf_kv = mx.zeros_like(self.buf_kv) + new_buf_gate = mx.zeros_like(self.buf_gate) + + for i in range(B): + r = self.remainder[i] + vl = valid_lengths[i] + u = usable[i] + nr = new_remainder[i] + + if u > 0: + # Tokens from the buffer (the leftover from last call) + if r > 0: + r_kv[i, :r] = self.buf_kv[i, :r] + r_gate[i, :r] = self.buf_gate[i, :r] + + # Tokens from the new input that complete full windows + consume = u - r + r_kv[i, r : r + consume] = kv[i, :consume] + r_gate[i, r : r + consume] = gate[i, :consume] + + r_base[i] = ( + offset[i] - r if isinstance(offset, mx.array) else offset - r + ) + + # Fill new remainder buffer from the tail of the input + if nr > 0: + if u > 0: + # Old remainder was consumed into usable output; + # new remainder is purely from the tail of new input. + new_buf_kv[i, :nr] = kv[i, vl - nr : vl] + new_buf_gate[i, :nr] = gate[i, vl - nr : vl] + else: + # No full window produced: carry over old buffer and + # append any new valid tokens. + if r > 0: + new_buf_kv[i, :r] = self.buf_kv[i, :r] + new_buf_gate[i, :r] = self.buf_gate[i, :r] + if vl > 0: + new_buf_kv[i, r : r + vl] = kv[i, :vl] + new_buf_gate[i, r : r + vl] = gate[i, :vl] + + self.buf_kv = new_buf_kv + self.buf_gate = new_buf_gate + self.remainder = new_remainder + + r_base = mx.array(r_base) + return r_kv, r_gate, r_base + + def update_and_fetch(self, px: mx.array): + B, N, D = px.shape + + if N == 0: + if self.pooled is None: + return mx.zeros((B, 0, D), dtype=px.dtype) + return self.pooled + + # Derive how many new pooled tokens each sequence actually produced. + new_counts = [ + (self._processed[i] - self.remainder[i]) // self.ratio + - self._pool_lengths[i] + for i in range(B) + ] + max_new = max(new_counts) + if max_new == 0: + if self.pooled is None: + return mx.zeros((B, 0, D), dtype=px.dtype) + return self.pooled + + max_pool = max(self._pool_lengths) + max_new + + if self.pooled is None: + self.pooled = mx.zeros((B, max_pool, D), dtype=px.dtype) + elif self.pooled.shape[1] < max_pool: + pad = mx.zeros((B, max_pool - self.pooled.shape[1], D), dtype=px.dtype) + self.pooled = mx.concatenate([self.pooled, pad], axis=1) + + for i in range(B): + nc = new_counts[i] + if nc > 0: + pl = self._pool_lengths[i] + self.pooled[i, pl : pl + nc] = px[i, :nc] + self._pool_lengths[i] = pl + nc + + return self.pooled + + def make_mask(self, L: int = 1, offset=0): + if self.pooled is None: + return None + + B, P, _ = self.pooled.shape + pool_lengths = mx.array(self._pool_lengths) + + # Length based mask + pool_idx = mx.arange(P)[None, None, :] + valid = pool_idx < pool_lengths[:, None, None] + + # Decode so no need for causal masking + if L == 1: + if all(pl == P for pl in self._pool_lengths): + return None + return valid + + # Prompt so we need to combine with causal + if isinstance(offset, mx.array): + query_pos = offset[:, None] + mx.arange(1, L + 1) + else: + query_pos = offset + mx.arange(offset + 1, offset + L + 1)[None] + + causal = pool_idx < (query_pos[..., None] // self.ratio) + mask = causal & valid + return mask + + @property + def state(self): + return (self.buf_kv, self.buf_gate, self.pooled) + + @state.setter + def state(self, v): + self.buf_kv, self.buf_gate, self.pooled = v + + @property + def meta_state(self): + return (self.ratio, self.remainder, self._pool_lengths, self._processed) + + @meta_state.setter + def meta_state(self, v): + self.ratio, self.remainder, self._pool_lengths, self._processed = v + + def is_trimmable(self): + return self.pooled is None + + def trim(self, n): + n = min(min(self.remainder), n) + for i in range(len(self.remainder)): + self.remainder[i] -= n + self._processed[i] -= n + return n + + def size(self): + return 0 if self.pooled is None else self.pooled.shape[1] + + def empty(self): + return self.pooled is None and all(r == 0 for r in self.remainder) + + @property + def nbytes(self): + total = 0 + if self.buf_kv is not None: + total += self.buf_kv.nbytes + self.buf_gate.nbytes + if self.pooled is not None: + total += self.pooled.nbytes + return total + + def filter(self, batch_indices): + if isinstance(batch_indices, mx.array): + idx_list = batch_indices.tolist() + else: + idx_list = list(batch_indices) + + if self.buf_kv is not None: + self.buf_kv = self.buf_kv[batch_indices] + self.buf_gate = self.buf_gate[batch_indices] + if self.pooled is not None: + self.pooled = self.pooled[batch_indices] + + self.remainder = [self.remainder[i] for i in idx_list] + self._pool_lengths = [self._pool_lengths[i] for i in idx_list] + self._lengths = [self._lengths[i] for i in idx_list] + self._processed = [self._processed[i] for i in idx_list] + + def extend(self, other): + # Merge the remainder buffers + if self.buf_kv is None and other.buf_kv is None: + pass + elif self.buf_kv is not None and other.buf_kv is not None: + self.buf_kv = mx.concatenate([self.buf_kv, other.buf_kv], axis=0) + self.buf_gate = mx.concatenate([self.buf_gate, other.buf_gate], axis=0) + elif self.buf_kv is None: + B = len(self.remainder) + D1 = other.buf_kv.shape[2] + D2 = other.buf_gate.shape[2] + self.buf_kv = mx.concatenate( + [mx.zeros((B, self.ratio, D1), dtype=other.buf_kv.dtype), other.buf_kv], + axis=0, + ) + self.buf_gate = mx.concatenate( + [ + mx.zeros((B, self.ratio, D2), dtype=other.buf_gate.dtype), + other.buf_gate, + ], + axis=0, + ) + else: + B2 = len(other.remainder) + D1 = self.buf_kv.shape[2] + D2 = self.buf_gate.shape[2] + self.buf_kv = mx.concatenate( + [self.buf_kv, mx.zeros((B2, self.ratio, D1), dtype=self.buf_kv.dtype)], + axis=0, + ) + self.buf_gate = mx.concatenate( + [ + self.buf_gate, + mx.zeros((B2, self.ratio, D2), dtype=self.buf_gate.dtype), + ], + axis=0, + ) + + # Merge the pooled buffers + if self.pooled is None and other.pooled is None: + pass + else: + B1 = len(self.remainder) + B2 = len(other.remainder) + P1 = 0 if self.pooled is None else self.pooled.shape[1] + P2 = 0 if other.pooled is None else other.pooled.shape[1] + max_P = max(P1, P2) + + if max_P > 0: + if self.pooled is not None: + D = self.pooled.shape[2] + else: + D = other.pooled.shape[2] + dt = (self.pooled if self.pooled is not None else other.pooled).dtype + + def pad_pool(pooled, B, P): + if pooled is None: + return mx.zeros((B, max_P, D), dtype=dt) + if P < max_P: + pad = mx.zeros((pooled.shape[0], max_P - P, D), dtype=dt) + return mx.concatenate([pooled, pad], axis=1) + return pooled + + self.pooled = mx.concatenate( + [pad_pool(self.pooled, B1, P1), pad_pool(other.pooled, B2, P2)], + axis=0, + ) + + self.remainder = self.remainder + other.remainder + self._pool_lengths = self._pool_lengths + other._pool_lengths + self._lengths = self._lengths + other._lengths + self._processed = self._processed + other._processed + + def extract(self, idx): + cache = PoolingCache(self.ratio) + pl = self._pool_lengths[idx] + r = self.remainder[idx] + + if self.pooled is not None and pl > 0: + cache.pooled = mx.contiguous(self.pooled[idx : idx + 1, :pl]) + + if self.buf_kv is not None and r > 0: + cache.buf_kv = mx.contiguous(self.buf_kv[idx : idx + 1]) + cache.buf_gate = mx.contiguous(self.buf_gate[idx : idx + 1]) + cache.remainder = r + + return cache + + @classmethod + def merge(cls, caches): + """Merge a list of PoolingCache instances into a BatchPoolingCache.""" + B = len(caches) + if not all(c.ratio == caches[0].ratio for c in caches): + raise ValueError( + "BatchPoolingCache can only merge caches with the same ratio" + ) + ratio = caches[0].ratio + batch_cache = cls(ratio, [0] * B) + + # Check if all caches are empty + if all(c.empty() for c in caches): + return batch_cache + + # Merge pooled buffers + pool_sizes = [c.size() for c in caches] + max_pool = max(pool_sizes) + if max_pool > 0: + D = next(c.pooled.shape[2] for c in caches if c.pooled is not None) + dt = next(c.pooled.dtype for c in caches if c.pooled is not None) + pooled = mx.zeros((B, max_pool, D), dtype=dt) + for i, c in enumerate(caches): + if c.pooled is not None: + ps = c.pooled.shape[1] + pooled[i, :ps] = c.pooled[0] + batch_cache.pooled = pooled + + batch_cache._pool_lengths = pool_sizes + batch_cache.remainder = [c.remainder for c in caches] + batch_cache._processed = [ + c.remainder + ps * ratio for c, ps in zip(caches, pool_sizes) + ] + + # Merge remainder buffers + has_buf = any(c.buf_kv is not None for c in caches) + if has_buf: + D1 = next(c.buf_kv.shape[2] for c in caches if c.buf_kv is not None) + D2 = next(c.buf_gate.shape[2] for c in caches if c.buf_gate is not None) + dt = next(c.buf_kv.dtype for c in caches if c.buf_kv is not None) + buf_kv = mx.zeros((B, ratio, D1), dtype=dt) + buf_gate = mx.zeros((B, ratio, D2), dtype=dt) + for i, c in enumerate(caches): + if c.buf_kv is not None and c.remainder > 0: + buf_kv[i, : c.remainder] = c.buf_kv[0, : c.remainder] + buf_gate[i, : c.remainder] = c.buf_gate[0, : c.remainder] + batch_cache.buf_kv = buf_kv + batch_cache.buf_gate = buf_gate + + return batch_cache + + def dynamic_roll(x, shifts, axis): n = x.shape[axis] expand_shifts = (...,) + (None,) * (x.ndim - axis) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py new file mode 100644 index 000000000..49f1d8bd9 --- /dev/null +++ b/mlx_lm/models/deepseek_v4.py @@ -0,0 +1,1171 @@ +# Copyright © 2026 Apple Inc. + +import math +from dataclasses import dataclass, field +from functools import partial +from typing import Any, Dict, List, Optional, Tuple, Union + +import mlx.core as mx +import mlx.nn as nn +from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients +from mlx.utils import tree_flatten + +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .cache import CacheList, PoolingCache, RotatingKVCache +from .hyper_connection import HyperConnection, HyperHead, hc_expand +from .mla import MultiLinear +from .pipeline import PipelineMixin +from .switch_layers import SwitchGLU + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str = "deepseek_v4" + vocab_size: int = 129280 + hidden_size: int = 4096 + intermediate_size: int = 18432 + moe_intermediate_size: int = 2048 + num_hidden_layers: int = 43 + num_attention_heads: int = 64 + num_key_value_heads: int = 1 + n_shared_experts: int = 1 + n_routed_experts: int = 256 + routed_scaling_factor: float = 1.5 + q_lora_rank: int = 1024 + qk_rope_head_dim: int = 64 + num_experts_per_tok: int = 6 + norm_topk_prob: bool = True + hidden_act: str = "silu" + max_position_embeddings: int = 1048576 + rms_norm_eps: float = 1e-6 + rope_theta: float = 10000.0 + rope_scaling: Optional[Dict] = None + attention_bias: bool = False + attention_dropout: float = 0.0 + head_dim: int = 512 + scoring_func: str = "sqrtsoftplus" + compress_ratios: List[int] = field(default_factory=list) + compress_rope_theta: float = 160000.0 + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + num_hash_layers: int = 3 + swiglu_limit: float = 10.0 + sliding_window: int = 128 + o_groups: int = 8 + o_lora_rank: int = 1024 + index_n_heads: int = 64 + index_head_dim: int = 128 + index_topk: int = 512 + num_nextn_predict_layers: int = 1 + tie_word_embeddings: bool = False + topk_method: str = "noaux_tc" + + def __post_init__(self): + if not self.compress_ratios: + n = self.num_hidden_layers + self.compress_ratios = ( + [0] + + [4 if i % 2 else 128 for i in range(max(n - 2, 0))] + + ([0] if n >= 2 else []) + ) + self.compress_ratios = list(self.compress_ratios[: self.num_hidden_layers]) + if len(self.compress_ratios) != self.num_hidden_layers: + raise ValueError( + "`compress_ratios` must have one entry per hidden layer, " + f"got {len(self.compress_ratios)} for {self.num_hidden_layers} layers." + ) + bad = [r for r in self.compress_ratios if r not in (0, 4, 128)] + if bad: + raise ValueError(f"Unsupported DeepSeek-V4 compress ratios: {bad}") + + +def make_quantization_config(model): + mxfp4 = {"group_size": 32, "bits": 4, "mode": "mxfp4"} + mxfp8 = {"group_size": 32, "bits": 8, "mode": "mxfp8"} + + flat_modules = tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module) + experts = { + k: mxfp4 + for k, _ in flat_modules + if ".ffn.switch_mlp." in k and k.endswith("_proj") + } + shared_experts = {k: mxfp8 for k, _ in flat_modules if ".ffn.shared_experts." in k} + attn = { + k: mxfp8 for k, _ in flat_modules if ".attn.w" in k or ".attn.indexer.wq" in k + } + + return { + "group_size": 64, + "bits": 8, + "mode": "affine", + **experts, + **shared_experts, + **attn, + } + + +def _score_func(scores: mx.array, func: str) -> mx.array: + if func == "softmax": + return mx.softmax(scores, axis=-1, precise=True) + if func == "sigmoid": + return mx.sigmoid(scores) + if func == "sqrtsoftplus": + return mx.sqrt(nn.softplus(scores)) + raise ValueError(f"Unsupported DeepSeek-V4 scoring function: {func}") + + +@mx.compile +def _expert_select( + logits: mx.array, + e_score_correction_bias: mx.array, + top_k: int, + routed_scaling_factor: float, + norm_topk_prob: bool, + scoring_func: str, +) -> Tuple[mx.array, mx.array]: + logits = logits.astype(mx.float32) + scores = _score_func(logits, scoring_func) + biased = scores + e_score_correction_bias + inds = mx.argpartition(-biased, kth=top_k - 1, axis=-1)[..., :top_k] + weights = mx.take_along_axis(scores, inds, axis=-1) + if scoring_func != "softmax" and norm_topk_prob: + weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) + weights = weights * routed_scaling_factor + return inds, weights + + +@mx.compile +def _hash_expert_select( + input_ids: mx.array, + logits: mx.array, + tid2eid: mx.array, + routed_scaling_factor: float, + norm_topk_prob: bool, + scoring_func: str, +) -> Tuple[mx.array, mx.array]: + logits = logits.astype(mx.float32) + scores = _score_func(logits, scoring_func) + inds = tid2eid[input_ids] + weights = mx.take_along_axis(scores, inds, axis=-1) + if scoring_func != "softmax" and norm_topk_prob: + weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) + weights = weights * routed_scaling_factor + return inds, weights + + +@mx.compile +def _limited_swiglu(gate: mx.array, up: mx.array, limit: float) -> mx.array: + if limit and limit > 0: + gate = mx.minimum(gate, limit) + up = mx.clip(up, -limit, limit) + return nn.silu(gate) * up + + +class LimitedSwiGLU(nn.Module): + def __init__(self, limit: float): + super().__init__() + self.limit = limit + + def __call__(self, x, gate): + return _limited_swiglu(gate, x, self.limit) + + +class DeepseekV4RoPE(nn.Module): + def __init__( + self, + dims: int, + base: float, + scaling_config: Optional[Dict] = None, + max_position_embeddings: int = 1048576, + freq_scale: int = 1, + ): + super().__init__() + self.dims = dims + self.freq_scale = freq_scale + + inv_freq = 1.0 / (base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)) + rope_type = None + if scaling_config is not None: + rope_type = scaling_config.get("type") or scaling_config.get("rope_type") + + if rope_type in ("yarn", "deepseek_yarn"): + factor = scaling_config["factor"] + original_max_position_embeddings = scaling_config[ + "original_max_position_embeddings" + ] + beta_fast = scaling_config.get("beta_fast", 32) + beta_slow = scaling_config.get("beta_slow", 1) + + def correction_dim(num_rotations): + return ( + dims + * math.log( + original_max_position_embeddings / (num_rotations * 2 * math.pi) + ) + / (2 * math.log(base)) + ) + + low = max(math.floor(correction_dim(beta_fast)), 0) + high = min(math.ceil(correction_dim(beta_slow)), dims - 1) + if low == high: + high += 0.001 + + ramp = (mx.arange(dims // 2, dtype=mx.float32) - low) / (high - low) + smooth = 1 - mx.clip(ramp, 0, 1) + inv_freq = inv_freq / factor * (1 - smooth) + inv_freq * smooth + + elif rope_type not in (None, "default"): + raise ValueError(f"Unsupported DeepSeek-V4 RoPE type: {rope_type}") + + self._freqs = 1.0 / inv_freq + self._freqs_cache = {} + + def _get_freqs(self, head_dim: int, inverse: bool): + key = (head_dim, inverse) + if key not in self._freqs_cache: + f = self._freqs + if self.freq_scale != 1: + f = f / self.freq_scale + if inverse: + f = -f + nope_pairs = (head_dim - self.dims) // 2 + if nope_pairs > 0: + f = mx.concatenate([mx.full((nope_pairs,), mx.inf), f]) + self._freqs_cache[key] = f + return self._freqs_cache[key] + + def __call__( + self, + x: mx.array, + offset: Any = 0, + inverse: bool = False, + ) -> mx.array: + head_dim = x.shape[-1] + freqs = self._get_freqs(head_dim, inverse) + offset = offset // self.freq_scale if self.freq_scale != 1 else offset + return mx.fast.rope( + x, + head_dim, + traditional=True, + base=None, + scale=1.0, + offset=offset, + freqs=freqs, + ) + + +def _apply_score_mask(scores: mx.array, mask: Optional[mx.array]) -> mx.array: + if mask is None: + return scores + if mask.dtype == mx.bool_: + return mx.where(mask, scores, mx.finfo(scores.dtype).min) + return scores + mask.astype(scores.dtype) + + +def _extend_mask(mask: Optional[mx.array], pool_mask: Optional[mx.array], N: int): + if mask is None: + return None + + if mask.ndim == 2: + mask = mask[None, None] + B, H, L, S = mask.shape + + if pool_mask is None: + pool_mask = mx.ones((B, H, L, N - S), dtype=mx.bool_) + elif pool_mask.ndim == 2: + pool_mask = mx.broadcast_to(pool_mask, (B, H, L, N - S)) + elif pool_mask.ndim == 3: + pool_mask = mx.broadcast_to(pool_mask[:, None], (B, H, L, N - S)) + + full_mask = mx.concatenate([mask, pool_mask], axis=-1) + + return full_mask + + +@partial(mx.compile, shapeless=True) +def _simple_compress_kv(kv, gate, ape, head_dim): + weights = mx.softmax(gate.astype(mx.float32) + ape, axis=-2) + weights = weights.astype(kv.dtype) + return (kv * weights).sum(axis=-2) + + +@mx.compile +def _overlap_compress_kv(kv, gate, ape, head_dim): + B, L, R, D = kv.shape + + gate = gate + ape.astype(gate.dtype) + + kv_0 = mx.zeros((B, 1, R, D // 2), dtype=kv.dtype) + kv_a, kv_b = mx.split(kv, 2, axis=-1) + kv_a = mx.concatenate([kv_0, kv_a[:, :-1]], axis=1) + kv = mx.concatenate([kv_a, kv_b], axis=2) + + gate_0 = mx.full((B, 1, R, D // 2), -mx.inf, dtype=kv.dtype) + gate_a, gate_b = mx.split(gate, 2, axis=-1) + gate_a = mx.concatenate([gate_0, gate_a[:, :-1]], axis=1) + gate = mx.concatenate([gate_a, gate_b], axis=2) + + weights = mx.softmax(gate, axis=-2, precise=True) + return (kv * weights).sum(axis=-2) + + +@partial(mx.compile, shapeless=True) +def _split_softmax(log_normalizer, logits_a, logits_b, sinks=None): + if sinks is not None: + log_normalizer = mx.logaddexp(log_normalizer, sinks) + weights_a = mx.exp(logits_a - log_normalizer) + weights_b = mx.exp(logits_b - log_normalizer) + return weights_a, weights_b + + +def _sparse_pooled_attention( + q: mx.array, + local_kv: mx.array, + pooled: mx.array, + topk: mx.array, + local_mask: Optional[mx.array], + pooled_mask: Optional[mx.array], + scale: float, + sinks: Optional[mx.array], +) -> mx.array: + B, H, L, D = q.shape + idx = topk[:, None, :, :, None] + pooled = mx.take_along_axis( + mx.broadcast_to(pooled[:, None, None], (B, 1, L, pooled.shape[1], D)), + mx.broadcast_to(idx, idx.shape[:-1] + (D,)), + axis=3, + ) + + q_scaled = q * scale + local_scores = q_scaled @ local_kv.swapaxes(-1, -2) + local_scores = _apply_score_mask(local_scores, local_mask) + normalizer = mx.logsumexp(local_scores, -1, keepdims=True) + + pooled_sq = pooled.squeeze(1) + q_bl = q_scaled.transpose(0, 2, 1, 3) + pooled_scores = q_bl @ pooled_sq.swapaxes(-1, -2) + pooled_scores = pooled_scores.transpose(0, 2, 1, 3) + pooled_scores = _apply_score_mask(pooled_scores, pooled_mask) + normalizer = mx.logaddexp( + normalizer, mx.logsumexp(pooled_scores, -1, keepdims=True) + ) + + local_weights, pooled_weights = _split_softmax( + normalizer, + local_scores, + pooled_scores, + sinks[None, :, None, None] if sinks is not None else None, + ) + + out = local_weights @ local_kv + pw_bl = pooled_weights.transpose(0, 2, 1, 3) + out = out + (pw_bl @ pooled_sq).transpose(0, 2, 1, 3) + return out.astype(q.dtype) + + +class MoEGate(nn.Module): + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.n_routed_experts + self.hidden_dim = config.hidden_size + self.hash = layer_idx < config.num_hash_layers + self.scoring_func = config.scoring_func + self.routed_scaling_factor = config.routed_scaling_factor + self.norm_topk_prob = config.norm_topk_prob + self.weight = mx.zeros((self.num_experts, self.hidden_dim)) + if self.hash: + self.tid2eid = mx.zeros((config.vocab_size, self.top_k), dtype=mx.int32) + else: + self.e_score_correction_bias = mx.zeros( + (self.num_experts,), dtype=mx.float32 + ) + + def __call__(self, x: mx.array, input_ids: Optional[mx.array] = None): + logits = x @ self.weight.T + + if self.hash: + if input_ids is None: + raise ValueError("DeepSeek-V4 hash routing requires input_ids.") + inds, weights = _hash_expert_select( + input_ids, + logits, + self.tid2eid, + self.routed_scaling_factor, + self.norm_topk_prob, + self.scoring_func, + ) + else: + inds, weights = _expert_select( + logits, + self.e_score_correction_bias, + self.top_k, + self.routed_scaling_factor, + self.norm_topk_prob, + self.scoring_func, + ) + + return inds, weights + + +class DeepseekV4MLP(nn.Module): + def __init__( + self, + config: ModelArgs, + intermediate_size: Optional[int] = None, + swiglu_limit: float = 0.0, + ): + super().__init__() + hidden_size = config.hidden_size + intermediate_size = intermediate_size or config.intermediate_size + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + self.swiglu_limit = swiglu_limit + + def __call__(self, x: mx.array) -> mx.array: + return self.down_proj( + _limited_swiglu(self.gate_proj(x), self.up_proj(x), self.swiglu_limit) + ) + + +class DeepseekV4MoE(nn.Module): + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__() + self.config = config + self.gate = MoEGate(config, layer_idx) + self.switch_mlp = SwitchGLU( + config.hidden_size, + config.moe_intermediate_size, + config.n_routed_experts, + activation=LimitedSwiGLU(config.swiglu_limit), + ) + self.shared_experts = DeepseekV4MLP( + config, + intermediate_size=config.moe_intermediate_size * config.n_shared_experts, + ) + self.sharding_group = None + + def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: + if self.sharding_group is not None: + x = sum_gradients(self.sharding_group)(x) + + inds, scores = self.gate(x, input_ids) + y = self.switch_mlp(x, inds) + y = (y * scores[..., None].astype(y.dtype)).sum(-2) + y = y + self.shared_experts(x) + + if self.sharding_group is not None: + y = mx.distributed.all_sum(y, group=self.sharding_group) + return y + + +class Compressor(nn.Module): + + def __init__(self, config: ModelArgs, compress_ratio: int, head_dim: int): + super().__init__() + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.overlap = compress_ratio == 4 + self.out_dim = head_dim * (2 if self.overlap else 1) + self.wkv = nn.Linear(config.hidden_size, self.out_dim, bias=False) + self.wgate = nn.Linear(config.hidden_size, self.out_dim, bias=False) + self.ape = mx.zeros((compress_ratio, self.out_dim), dtype=mx.float32) + self.norm = nn.RMSNorm(head_dim, eps=config.rms_norm_eps) + self.rope = DeepseekV4RoPE( + config.qk_rope_head_dim, + config.compress_rope_theta, + config.rope_scaling, + config.max_position_embeddings, + freq_scale=compress_ratio, + ) + + def __call__( + self, + x: mx.array, + pool_cache: Optional[PoolingCache], + offset: Union[int, mx.array], + ) -> mx.array: + B, _, _ = x.shape + kv = self.wkv(x) + gate = self.wgate(x) + if pool_cache is None: + usable = (kv.shape[1] // self.compress_ratio) * self.compress_ratio + ready_kv, ready_gate = kv[:, :usable], gate[:, :usable] + pool_base = offset + else: + ready_kv, ready_gate, pool_base = pool_cache.accumulate_windows( + kv, gate, offset + ) + + if ready_kv.size == 0: + new_pooled = mx.zeros((B, 0, self.head_dim), dtype=x.dtype) + else: + compress_func = ( + _overlap_compress_kv if self.overlap else _simple_compress_kv + ) + kv = mx.unflatten(ready_kv, 1, (-1, self.compress_ratio)) + gate = mx.unflatten(ready_gate, 1, (-1, self.compress_ratio)) + new_pooled = compress_func(kv, gate, self.ape, self.head_dim) + new_pooled = self.norm(new_pooled) + new_pooled = self.rope( + new_pooled[:, None], + offset=pool_base, + ).squeeze(1) + + if pool_cache is not None: + new_pooled = pool_cache.update_and_fetch(new_pooled) + + return new_pooled + + +class Indexer(nn.Module): + def __init__(self, config: ModelArgs, compress_ratio: int): + super().__init__() + self.n_heads = config.index_n_heads + self.head_dim = config.index_head_dim + self.index_topk = config.index_topk + self.wq_b = nn.Linear( + config.q_lora_rank, self.n_heads * self.head_dim, bias=False + ) + self.weights_proj = nn.Linear(config.hidden_size, self.n_heads, bias=False) + self.compressor = Compressor(config, compress_ratio, self.head_dim) + self.scale = self.head_dim**-0.5 + + def __call__( + self, + x: mx.array, + q_residual: mx.array, + position_rope: DeepseekV4RoPE, + pool_cache: Optional[PoolingCache], + offset: Union[int, mx.array], + ): + B, L, _ = x.shape + pooled = self.compressor(x, pool_cache, offset) + if pooled.shape[1] == 0: + return None + + q = self.wq_b(q_residual).reshape(B, L, self.n_heads, self.head_dim) + q = q.transpose(0, 2, 1, 3) + q = position_rope(q, offset) + + scores = q.astype(mx.float32) @ pooled[:, None].swapaxes(-1, -2).astype( + mx.float32 + ) + scores = mx.maximum(scores, 0) * self.scale + weights = self.weights_proj(x).astype(mx.float32) * (self.n_heads**-0.5) + scores = (scores * weights.swapaxes(-1, -2)[..., None]).sum(axis=1) + pmask = pool_cache.make_mask(L, offset) if pool_cache is not None else None + if pmask is not None: + scores = mx.where( + pmask if pmask.ndim == 3 else pmask[None], + scores, + mx.finfo(scores.dtype).min, + ) + k = min(self.index_topk, pooled.shape[1]) + return mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k] + + +class LocalAttention(nn.Module): + """DeepSeek V4 attention with no KV compression.""" + + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.compress_ratio = 0 + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.o_groups = config.o_groups + self.o_lora_rank = config.o_lora_rank + self.scale = self.head_dim**-0.5 + + self.wq_a = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False) + self.q_norm = nn.RMSNorm(config.q_lora_rank, eps=config.rms_norm_eps) + self.wq_b = nn.Linear( + config.q_lora_rank, self.n_heads * self.head_dim, bias=False + ) + self.wkv = nn.Linear(config.hidden_size, self.head_dim, bias=False) + self.kv_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.wo_a = MultiLinear( + self.n_heads * self.head_dim // config.o_groups, + config.o_lora_rank, + config.o_groups, + ) + self.wo_b = nn.Linear( + config.o_groups * config.o_lora_rank, + config.hidden_size, + bias=config.attention_bias, + ) + self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) + + self.rope = DeepseekV4RoPE( + config.qk_rope_head_dim, + config.rope_theta, + None, + config.max_position_embeddings, + ) + + self.sharding_group = None + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + offset = cache.offset if cache is not None else 0 + offset = mx.array(offset) if isinstance(offset, mx.array) else offset + + q = self.wq_b(self.q_norm(self.wq_a(x))) + q = q.reshape(B, L, self.n_heads, self.head_dim) + q = mx.fast.rms_norm(q, None, self.config.rms_norm_eps) + q = q.transpose(0, 2, 1, 3) + q = self.rope(q, offset) + + kv = self.kv_norm(self.wkv(x)).reshape(B, 1, L, self.head_dim) + kv = self.rope(kv, offset) + if cache is not None: + kv, _ = cache.update_and_fetch(kv, mx.zeros((B, 1, L, 0))) + + out = scaled_dot_product_attention( + q, + kv, + kv, + cache=cache, + scale=self.scale, + mask=mask, + sinks=self.attn_sink.astype(q.dtype), + ) + out = self.rope(out, offset, inverse=True) + + out = out.reshape(B, self.o_groups, -1, L, self.head_dim) + out = out.transpose(0, 1, 3, 2, 4).flatten(-2) + out = self.wo_a(out) + out = out.transpose(0, 2, 1, 3).flatten(-2) + out = self.wo_b(out) + + if self.sharding_group is not None: + out = mx.distributed.all_sum(out, group=self.sharding_group) + + return out + + +class CompressedAttention(nn.Module): + """DeepSeek V4 attention with pooled KV compression.""" + + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.compress_ratio = config.compress_ratios[layer_idx] + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.o_groups = config.o_groups + self.o_lora_rank = config.o_lora_rank + self.scale = self.head_dim**-0.5 + + self.wq_a = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False) + self.q_norm = nn.RMSNorm(config.q_lora_rank, eps=config.rms_norm_eps) + self.wq_b = nn.Linear( + config.q_lora_rank, self.n_heads * self.head_dim, bias=False + ) + self.wkv = nn.Linear(config.hidden_size, self.head_dim, bias=False) + self.kv_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.wo_a = MultiLinear( + self.n_heads * self.head_dim // config.o_groups, + config.o_lora_rank, + config.o_groups, + ) + self.wo_b = nn.Linear( + config.o_groups * config.o_lora_rank, + config.hidden_size, + bias=config.attention_bias, + ) + self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) + + # Compressed layers use Yarn-scaled RoPE + self.rope = DeepseekV4RoPE( + config.qk_rope_head_dim, + config.compress_rope_theta, + config.rope_scaling, + config.max_position_embeddings, + ) + self.compressor = Compressor(config, self.compress_ratio, self.head_dim) + + self.sharding_group = None + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + local_cache = cache[0] if cache is not None else None + pool_cache = cache[1] if cache is not None else None + offset = local_cache.offset if local_cache is not None else 0 + offset = mx.array(offset) if isinstance(offset, mx.array) else offset + + q = self.wq_b(self.q_norm(self.wq_a(x))) + q = q.reshape(B, L, self.n_heads, self.head_dim) + q = mx.fast.rms_norm(q, None, self.config.rms_norm_eps) + q = q.transpose(0, 2, 1, 3) + q = self.rope(q, offset) + + kv = self.kv_norm(self.wkv(x)).reshape(B, 1, L, self.head_dim) + kv = self.rope(kv, offset) + if local_cache is not None: + kv, _ = local_cache.update_and_fetch(kv, mx.zeros((B, 1, L, 0))) + + # Pool tokens into compressed KV and concatenate with local KV + pooled = self.compressor(x, pool_cache, offset) + pooled_mask = None + if pooled.shape[1] > 0: + pooled_mask = ( + pool_cache.make_mask(L, offset) if pool_cache is not None else None + ) + kv = mx.concatenate([kv, pooled[:, None]], axis=2) + + mask = _extend_mask(mask, pooled_mask, kv.shape[2]) + + out = scaled_dot_product_attention( + q, + kv, + kv, + cache=local_cache, + scale=self.scale, + mask=mask, + sinks=self.attn_sink.astype(q.dtype), + ) + out = self.rope(out, offset, inverse=True) + + out = out.reshape(B, self.o_groups, -1, L, self.head_dim) + out = out.transpose(0, 1, 3, 2, 4).flatten(-2) + out = self.wo_a(out) + out = out.transpose(0, 2, 1, 3).flatten(-2) + out = self.wo_b(out) + + if self.sharding_group is not None: + out = mx.distributed.all_sum(out, group=self.sharding_group) + + return out + + +class SparseCompressedAttention(nn.Module): + """DeepSeek V4 attention with sparse indexed pooled KV compression.""" + + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.compress_ratio = config.compress_ratios[layer_idx] + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.o_groups = config.o_groups + self.o_lora_rank = config.o_lora_rank + self.scale = self.head_dim**-0.5 + + self.wq_a = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False) + self.q_norm = nn.RMSNorm(config.q_lora_rank, eps=config.rms_norm_eps) + self.wq_b = nn.Linear( + config.q_lora_rank, self.n_heads * self.head_dim, bias=False + ) + self.wkv = nn.Linear(config.hidden_size, self.head_dim, bias=False) + self.kv_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.wo_a = MultiLinear( + self.n_heads * self.head_dim // config.o_groups, + config.o_lora_rank, + config.o_groups, + ) + self.wo_b = nn.Linear( + config.o_groups * config.o_lora_rank, + config.hidden_size, + bias=config.attention_bias, + ) + self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) + + self.rope = DeepseekV4RoPE( + config.qk_rope_head_dim, + config.compress_rope_theta, + config.rope_scaling, + config.max_position_embeddings, + ) + self.compressor = Compressor(config, self.compress_ratio, self.head_dim) + self.indexer = Indexer(config, self.compress_ratio) + + self.sharding_group = None + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + local_cache = cache[0] if cache is not None else None + comp_cache = cache[1] if cache is not None else None + idx_cache = cache[2] if cache is not None else None + offset = local_cache.offset if local_cache is not None else 0 + offset = mx.array(offset) if isinstance(offset, mx.array) else offset + + q_residual = self.q_norm(self.wq_a(x)) + q = self.wq_b(q_residual).reshape(B, L, self.n_heads, self.head_dim) + q = mx.fast.rms_norm(q, None, self.config.rms_norm_eps) + q = q.transpose(0, 2, 1, 3) + q = self.rope(q, offset) + + kv = self.kv_norm(self.wkv(x)).reshape(B, 1, L, self.head_dim) + kv = self.rope(kv, offset) + if local_cache is not None: + kv, _ = local_cache.update_and_fetch(kv, mx.zeros((B, 1, L, 0))) + + pooled = self.compressor(x, comp_cache, offset) + pmask = comp_cache.make_mask(L, offset) if comp_cache is not None else None + topk = self.indexer(x, q_residual, self.rope, idx_cache, offset) + sinks = self.attn_sink.astype(q.dtype) + + # Local attention + if pooled.shape[1] == 0: + out = scaled_dot_product_attention( + q, + kv, + kv, + cache=local_cache, + scale=self.scale, + mask=mask, + sinks=sinks, + ) + + # Compressed attention + elif pooled.shape[1] <= self.indexer.index_topk: + full_kv = mx.concatenate([kv, pooled[:, None]], axis=2) + mask = _extend_mask(mask, pmask, full_kv.shape[2]) + out = scaled_dot_product_attention( + q, + full_kv, + full_kv, + cache=local_cache, + scale=self.scale, + mask=mask, + sinks=sinks, + ) + + # Sparse compressed attention + else: + sparse_mask = None + if pmask is not None: + sparse_mask = mx.take_along_axis( + pmask[None] if pmask.ndim == 2 else pmask, + topk, + axis=2, + )[:, None] + out = _sparse_pooled_attention( + q, + kv, + pooled, + topk, + mask, + sparse_mask, + self.scale, + sinks, + ) + + out = self.rope(out, offset, inverse=True) + + out = out.reshape(B, self.o_groups, -1, L, self.head_dim) + out = out.transpose(0, 1, 3, 2, 4).flatten(-2) + out = self.wo_a(out) + out = out.transpose(0, 2, 1, 3).flatten(-2) + out = self.wo_b(out) + + if self.sharding_group is not None: + out = mx.distributed.all_sum(out, group=self.sharding_group) + + return out + + +def v4_attention_factory(config: ModelArgs, layer_idx: int) -> nn.Module: + """Instantiate the appropriate attention module for a given layer.""" + ratio = config.compress_ratios[layer_idx] + if ratio == 0: + return LocalAttention(config, layer_idx) + if ratio == 128: + return CompressedAttention(config, layer_idx) + return SparseCompressedAttention(config, layer_idx) + + +class DeepseekV4Block(nn.Module): + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__() + self.attn = v4_attention_factory(config, layer_idx) + self.ffn = DeepseekV4MoE(config, layer_idx) + self.attn_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.attn_hc = HyperConnection(config) + self.ffn_hc = HyperConnection(config) + + def __call__( + self, + h: mx.array, + mask: Optional[mx.array], + cache: Optional[Any], + input_ids: mx.array, + ) -> mx.array: + residual = h + x, post, comb = self.attn_hc(h) + x = self.attn(self.attn_norm(x), mask=mask, cache=cache) + h = hc_expand(x, residual, post, comb) + + residual = h + x, post, comb = self.ffn_hc(h) + x = self.ffn(self.ffn_norm(x), input_ids) + return hc_expand(x, residual, post, comb) + + +class DeepseekV4Model(PipelineMixin, nn.Module): + def __init__(self, config: ModelArgs): + super().__init__() + self.args = config + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = [ + DeepseekV4Block(config, idx) for idx in range(config.num_hidden_layers) + ] + self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hc_head = HyperHead(config) + + def __call__(self, inputs: mx.array, cache: Optional[Any] = None) -> mx.array: + h = self.embed_tokens(inputs) + h = mx.broadcast_to( + h[:, :, None, :], + (h.shape[0], h.shape[1], self.args.hc_mult, h.shape[2]), + ) + h = mx.contiguous(h) + + pipeline_rank = self.pipeline_rank + pipeline_size = self.pipeline_size + + if cache is None: + cache = [None] * len(self.pipeline_layers) + + first_cache = cache[0] + mask_cache = ( + first_cache[0] if isinstance(first_cache, CacheList) else first_cache + ) + mask = create_attention_mask( + h[:, :, 0, :], + mask_cache, + window_size=self.args.sliding_window, + return_array=True, + ) + + if pipeline_rank < pipeline_size - 1: + h = mx.distributed.recv_like(h, (pipeline_rank + 1)) + + for layer, layer_cache in zip(self.pipeline_layers, cache): + h = layer(h, mask, layer_cache, inputs) + + if pipeline_rank != 0: + h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size) + cache_item = cache[-1] + if isinstance(cache_item, CacheList): + cache_item = cache_item[0] + if cache_item is not None: + cache_item.keys = mx.depends(cache_item.keys, h) + + if pipeline_size > 1: + h = mx.distributed.all_gather(h)[: h.shape[0]] + + return self.norm(self.hc_head(h)) + + +class Model(nn.Module): + def __init__(self, config: ModelArgs): + super().__init__() + self.args = config + self.model_type = config.model_type + self.model = DeepseekV4Model(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def __call__(self, inputs: mx.array, cache: Optional[Any] = None): + return self.lm_head(self.model(inputs, cache)) + + @property + def layers(self): + return self.model.pipeline_layers + + @property + def cast_predicate(self): + def predicate(k): + return not ( + "attn_sink" in k + or "e_score_correction_bias" in k + or ".attn_hc." in k + or ".ffn_hc." in k + or ".hc_head." in k + ) + + return predicate + + def make_cache(self): + caches = [] + for layer in self.layers: + ratio = layer.attn.compress_ratio + if ratio == 0: + caches.append(RotatingKVCache(max_size=self.args.sliding_window)) + elif isinstance(layer.attn, SparseCompressedAttention): + # local + compressor pool + indexer pool + caches.append( + CacheList( + RotatingKVCache(max_size=self.args.sliding_window), + PoolingCache(ratio), + PoolingCache(ratio), + ) + ) + else: + # local + compressor pool + caches.append( + CacheList( + RotatingKVCache(max_size=self.args.sliding_window), + PoolingCache(ratio), + ) + ) + return caches + + def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: + n_layers = self.args.num_hidden_layers + + new_weights = {} + for k, v in weights.items(): + if k.startswith("mtp."): + continue + parts = k.split(".") + if len(parts) >= 2 and parts[0] == "layers": + try: + if int(parts[1]) >= n_layers: + continue + except ValueError: + pass + new_weights[k] = v + weights = new_weights + + new_weights = {} + for k, v in weights.items(): + if "tid2eid" in k: + new_weights[k] = v.astype(mx.int32) + + if not k.endswith(".scale"): + if k not in new_weights: + new_weights[k] = v + continue + + wk = k[: -len(".scale")] + ".weight" + weight = weights.get(wk) + if weight is None: + new_weights[k] = v + continue + if ( + ".ffn.experts." in wk + and ".shared_experts." not in wk + and weight.dtype in (mx.int8, mx.uint8) + and v.shape[-1] * 16 == weight.shape[-1] + ): + new_weights[k + "s"] = v + new_weights[wk] = weight.view(mx.uint32) + elif weight.dtype == mx.uint8: + new_weights[k + "s"] = mx.repeat(mx.repeat(v, 4, -1), 128, 0) + new_weights[wk] = weight.view(mx.uint32) + else: + new_weights[k] = v + weights = new_weights + + top_remap = { + "embed.weight": "model.embed_tokens.weight", + "norm.weight": "model.norm.weight", + "head.weight": "lm_head.weight", + "hc_head_fn": "model.hc_head.fn", + "hc_head_base": "model.hc_head.base", + "hc_head_scale": "model.hc_head.scale", + } + for old, new in top_remap.items(): + if old in weights: + weights[new] = weights.pop(old) + + remapped = {} + w_remap = {"w1": "gate_proj", "w2": "down_proj", "w3": "up_proj"} + for k, v in weights.items(): + nk = "model." + k if k.startswith("layers.") else k + nk = nk.replace(".ffn.gate.bias", ".ffn.gate.e_score_correction_bias") + for sub in ("attn", "ffn"): + for param in ("fn", "base", "scale"): + nk = nk.replace(f".hc_{sub}_{param}", f".{sub}_hc.{param}") + for old, new in w_remap.items(): + nk = nk.replace(f".shared_experts.{old}.", f".shared_experts.{new}.") + remapped[nk] = v + weights = remapped + + for layer_idx in range(n_layers): + prefix = f"model.layers.{layer_idx}.ffn.experts" + for src, dst in ( + ("w1", "gate_proj"), + ("w2", "down_proj"), + ("w3", "up_proj"), + ): + for suffix in ("weight", "scales"): + key0 = f"{prefix}.0.{src}.{suffix}" + if key0 in weights: + stacked = [ + weights.pop(f"{prefix}.{e}.{src}.{suffix}") + for e in range(self.args.n_routed_experts) + ] + weights[ + f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.{suffix}" + ] = mx.stack(stacked) + + # Reshape wo_a from nn.Linear (2D) to MultiLinear (3D) for all layers + for layer_idx in range(n_layers): + prefix = f"model.layers.{layer_idx}.attn.wo_a" + for key in (f"{prefix}.weight", f"{prefix}.scales", f"{prefix}.biases"): + if key in weights and weights[key].ndim == 2: + weights[key] = weights[key].reshape( + self.args.o_groups, self.args.o_lora_rank, -1 + ) + + return weights + + def shard(self, group: Optional[mx.distributed.Group] = None): + group = group or mx.distributed.init() + N = group.size() + rank = group.rank() + for layer in self.model.layers: + layer.attn.sharding_group = group + layer.attn.wq_b = shard_linear( + layer.attn.wq_b, + "all-to-sharded", + segments=self.args.o_groups, + group=group, + ) + shard_inplace(layer.attn.wo_a, "sharded-to-all", group=group) + layer.attn.attn_sink = mx.split(layer.attn.attn_sink, N)[rank] + layer.attn.n_heads //= N + + layer.ffn.sharding_group = group + shard_inplace( + layer.ffn.shared_experts.gate_proj, "all-to-sharded", group=group + ) + shard_inplace( + layer.ffn.shared_experts.down_proj, "sharded-to-all", group=group + ) + shard_inplace( + layer.ffn.shared_experts.up_proj, "all-to-sharded", group=group + ) + shard_inplace(layer.ffn.switch_mlp.gate_proj, "all-to-sharded", group=group) + shard_inplace(layer.ffn.switch_mlp.down_proj, "sharded-to-all", group=group) + shard_inplace(layer.ffn.switch_mlp.up_proj, "all-to-sharded", group=group) diff --git a/mlx_lm/models/hyper_connection.py b/mlx_lm/models/hyper_connection.py new file mode 100644 index 000000000..45ee78d39 --- /dev/null +++ b/mlx_lm/models/hyper_connection.py @@ -0,0 +1,285 @@ +# Copyright © 2026 Apple Inc. + +from typing import Tuple + +import mlx.core as mx +import mlx.nn as nn + + +def _make_hc_sinkhorn_collapse_kernel(): + """Fused sinkhorn + collapse: eliminates one dispatch per HC cycle. + + 1. BRANCHLESS SINKHORN: all 32 lanes in simd group 0 execute identical + instructions. Lanes >= HC use multiplicative mask (active=0) instead + of divergent branches — eliminates SIMD serialization. + 2. PARALLEL SINKHORN: lanes 0-3 each own one comb row. Column norm + via simd_sum() — free SIMD shuffle. + 3. NATIVE bfloat4 LOADS: single 64-bit load yields 4 bfloat16 values; + cast to float4 is a free hardware conversion. + 4. FMA CHAINS: collapse uses fused multiply-add for 3 of 4 terms. + """ + if mx.default_device() != mx.gpu or not mx.metal.is_available(): + return None + + source = """ + uint tid = thread_position_in_threadgroup.x; + uint row = threadgroup_position_in_grid.x; + uint lane = tid % 32; + uint sg = tid / 32; + + constexpr int MIX = (2 + HC) * HC; + constexpr int BASE_OFF = 2 * HC; + constexpr float EPS = EPS_INT * 1e-9; + + const device float* mix = (const device float*)mixes + row * MIX; + device float* post_out = (device float*)post + row * HC; + device float* comb_out = (device float*)comb + row * HC * HC; + + threadgroup float pre_shared[HC]; + + // ================================================================ + // PHASE 1: Branchless sinkhorn on simd group 0 + // All 32 lanes execute identical instructions. Lanes >= HC + // compute on clamped indices but multiply by active=0, so they + // contribute zero to simd_sum. No divergent branches in the loop. + // ================================================================ + if (sg == 0) { + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + + const float active = (lane < (uint)HC) ? 1.0f : 0.0f; + const uint llane = metal::min(lane, (uint)(HC - 1)); + + // Pre/post sigmoids: all lanes compute, only active lanes write + float pre_z = mix[llane] * pre_scale + base[llane]; + float post_z = mix[HC + llane] * post_scale + base[HC + llane]; + float pre_v = 1.0f / (1.0f + metal::fast::exp(-pre_z)) + EPS; + float post_v = 2.0f / (1.0f + metal::fast::exp(-post_z)); + + if (lane < (uint)HC) { + pre_shared[lane] = pre_v; + post_out[lane] = post_v; + } + + // Comb softmax: load + mask. Inactive lanes load row 0 (safe) + // but multiply by active=0 so they hold zeros. + float4 v = (*(const device float4*)(mix + BASE_OFF + llane * HC) + * comb_scale + + *(const device float4*)(base + BASE_OFF + llane * HC)) + * active; + + float row_max = metal::max(metal::max(v.x, v.y), + metal::max(v.z, v.w)); + float4 e = metal::fast::exp(v - row_max) * active; + float4 r = e * (1.0f / (e.x + e.y + e.z + e.w + EPS)) + + EPS * active; + + // Initial column normalization + float4 col_inv = 1.0f / (float4( + simd_sum(r.x), simd_sum(r.y), + simd_sum(r.z), simd_sum(r.w) + ) + EPS); + r *= col_inv; + + // Sinkhorn iterations: zero branches in the loop body + for (int iter = 1; iter < ITERS; ++iter) { + // Row norm + re-clamp inactive lanes + r *= (1.0f / (r.x + r.y + r.z + r.w + EPS)) * active; + + // Col norm via simd_sum + col_inv = 1.0f / (float4( + simd_sum(r.x), simd_sum(r.y), + simd_sum(r.z), simd_sum(r.w) + ) + EPS); + r *= col_inv; + } + + if (lane < (uint)HC) { + *(device float4*)(comb_out + lane * HC) = r; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // ================================================================ + // PHASE 2: Collapse — all 256 threads, vectorized + // ================================================================ + const float p0 = pre_shared[0]; + const float p1 = pre_shared[1]; + const float p2 = pre_shared[2]; + const float p3 = pre_shared[3]; + + const device T* x_row = (const device T*)x_in + + row * (HC * D); + device U* out_row = (device U*)collapsed + + row * D; + + using T4 = vec; + using U4 = vec; + const device T4* x_row0 = (const device T4*)(x_row + 0*D); + const device T4* x_row1 = (const device T4*)(x_row + 1*D); + const device T4* x_row2 = (const device T4*)(x_row + 2*D); + const device T4* x_row3 = (const device T4*)(x_row + 3*D); + device U4* out4 = (device U4*)out_row; + + constexpr uint D4 = (uint)D / 4; + + for (uint d4 = tid; d4 < D4; d4 += 256) { + float4 x0 = float4(x_row0[d4]); + float4 x1 = float4(x_row1[d4]); + float4 x2 = float4(x_row2[d4]); + float4 x3 = float4(x_row3[d4]); + + float4 result = fma(float4(p0), x0, + fma(float4(p1), x1, + fma(float4(p2), x2, float4(p3) * x3))); + + out4[d4] = U4(result); + } + + // Scalar tail for D not divisible by 4 + #if (D % 4) != 0 + for (uint d = D4 * 4 + tid; d < (uint)D; d += 256) { + float val = p0*(float)x_row[0*D+d] + p1*(float)x_row[1*D+d] + + p2*(float)x_row[2*D+d] + p3*(float)x_row[3*D+d]; + out_row[d] = (U)val; + } + #endif + """ + + return mx.fast.metal_kernel( + name="hc_sinkhorn_collapse", + input_names=["x_in", "mixes", "scale", "base"], + output_names=["collapsed", "post", "comb"], + source=source, + ensure_row_contiguous=True, + ) + + +_hc_sinkhorn_collapse_kernel = _make_hc_sinkhorn_collapse_kernel() + + +def _hc_kernel(x, y, mixes, scale, base, hc_mult, sinkhorn_iters, eps): + B, L, H, D = x.shape + + return _hc_sinkhorn_collapse_kernel( + inputs=[x, mixes, scale, base], + template=[ + ("T", x.dtype), + ("U", x.dtype), + ("HC", hc_mult), + ("ITERS", sinkhorn_iters), + ("D", D), + ("EPS_INT", round(eps / 1e-9)), + ], + grid=(B * L * 256, 1, 1), + threadgroup=(256, 1, 1), + output_shapes=[(B, L, D), (B, L, hc_mult), (B, L, hc_mult, hc_mult)], + output_dtypes=[x.dtype, mx.float32, mx.float32], + ) + + +@mx.compile +def _hc_split_sinkhorn_ops( + mixes: mx.array, + scale: mx.array, + base: mx.array, + hc_mult: int, + sinkhorn_iters: int, + eps: float, +) -> Tuple[mx.array, mx.array, mx.array]: + mixes = mixes.astype(mx.float32) + scale = scale.astype(mx.float32) + base = base.astype(mx.float32) + pre_scale, post_scale, comb_scale = scale[0], scale[1], scale[2] + + pre = mx.sigmoid(mixes[..., :hc_mult] * pre_scale + base[:hc_mult]) + eps + post = 2 * mx.sigmoid( + mixes[..., hc_mult : 2 * hc_mult] * post_scale + base[hc_mult : 2 * hc_mult] + ) + comb = mixes[..., 2 * hc_mult :].reshape( + *mixes.shape[:-1], hc_mult, hc_mult + ) * comb_scale + base[2 * hc_mult :].reshape(hc_mult, hc_mult) + comb = mx.softmax(comb, axis=-1, precise=True) + eps + comb = comb / (comb.sum(axis=-2, keepdims=True) + eps) + for _ in range(max(sinkhorn_iters - 1, 0)): + comb = comb / (comb.sum(axis=-1, keepdims=True) + eps) + comb = comb / (comb.sum(axis=-2, keepdims=True) + eps) + return pre, post, comb + + +def _hc_ops(x, y, mixes, scale, base, hc_mult, sinkhorn_iters, eps): + pre, post, comb = _hc_split_sinkhorn_ops( + mixes, scale, base, hc_mult, sinkhorn_iters, eps + ) + return (pre[..., None] * y).sum(axis=2).astype(x.dtype), post, comb + + +class HyperConnection(nn.Module): + def __init__(self, config): + super().__init__() + self.hc_mult = config.hc_mult + self.sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.norm_eps = config.rms_norm_eps + + mix = (2 + self.hc_mult) * self.hc_mult + self.fn = mx.zeros((mix, self.hc_mult * config.hidden_size), dtype=mx.float32) + self.base = mx.zeros((mix,), dtype=mx.float32) + self.scale = mx.ones((3,), dtype=mx.float32) + + def __call__(self, x: mx.array): + B, L, H, D = x.shape + y = x.astype(mx.float32) + z = mx.fast.rms_norm(y.flatten(-2), None, self.norm_eps) + mixes = z @ self.fn.T + + use_ops = ( + self.training + or mx.default_device() != mx.gpu + or not mx.metal.is_available() + ) + hc_func = _hc_ops if use_ops else _hc_kernel + + return hc_func( + x, + y, + mixes, + self.scale, + self.base, + self.hc_mult, + self.sinkhorn_iters, + self.hc_eps, + ) + + +@mx.compile +def _hc_expand_op(x, residual, post, comb): + y = post[..., None] * x[:, :, None, :].astype(mx.float32) + y = y + mx.matmul(comb.swapaxes(-1, -2), residual.astype(mx.float32)) + return y.astype(x.dtype) + + +def hc_expand(x, residual, post, comb): + return _hc_expand_op(x, residual, post, comb) + + +class HyperHead(nn.Module): + def __init__(self, config): + super().__init__() + self.hc_mult = config.hc_mult + self.norm_eps = config.rms_norm_eps + self.hc_eps = config.hc_eps + self.fn = mx.zeros( + (self.hc_mult, self.hc_mult * config.hidden_size), dtype=mx.float32 + ) + self.base = mx.zeros((self.hc_mult,), dtype=mx.float32) + self.scale = mx.ones((1,), dtype=mx.float32) + + def __call__(self, x: mx.array): + y = x.astype(mx.float32) + z = mx.fast.rms_norm(y.flatten(-2), None, self.norm_eps) + mixes = z @ self.fn.T + pre = mx.sigmoid(mixes * self.scale + self.base) + self.hc_eps + return (pre[..., None] * y).sum(axis=2).astype(x.dtype) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index ef3d266b9..0f201a189 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -8,6 +8,7 @@ import os import resource import shutil +import struct from pathlib import Path from textwrap import dedent from typing import ( @@ -56,6 +57,8 @@ MAX_FILE_SIZE_GB = 5 +SAFETENSORS_DTYPE_FALLBACKS = {"F8_E8M0": "U8"} + def _parse_size(x): sizes = {"M": 1e6, "G": 1e9, "MB": 1e6, "GB": 1e9, "": 1} @@ -279,6 +282,50 @@ def load_config(model_path: Path) -> dict: return config +def _load_safetensors(path: str) -> dict: + try: + return mx.load(path) + except RuntimeError as e: + if not any(dtype in str(e) for dtype in SAFETENSORS_DTYPE_FALLBACKS): + raise + load_error = e + + with open(path, "r+b") as f: + header_len = struct.unpack(" header_len: + raise RuntimeError( + f"Cannot reinterpret unsupported safetensors dtype in {path}: " + "patched header is larger than the original header." + ) + + try: + f.seek(8) + f.write(patched_header) + f.write(b" " * (header_len - len(patched_header))) + f.flush() + return mx.load(path) + finally: + f.seek(8) + f.write(original_header) + f.flush() + + def load_model( model_path: Path, lazy: bool = False, @@ -320,7 +367,7 @@ def load_model( weights = {} for wf in weight_files: - weights.update(mx.load(wf)) + weights.update(_load_safetensors(wf)) if (model_file := config.get("model_file")) is not None: spec = importlib.util.spec_from_file_location( @@ -388,6 +435,13 @@ def class_predicate(p, m): config["quantization"] = quantization config["quantization_config"] = quantization _quantize(quantization) + elif quant_method == "fp8" and config.get("model_type", None) == "deepseek_v4": + from .models.deepseek_v4 import make_quantization_config + + quantization = make_quantization_config(model) + config["quantization"] = quantization + config["quantization_config"] = quantization + _quantize(quantization) if config.get("quantize_activations", False): diff --git a/tests/test_models.py b/tests/test_models.py index 6e1fcd96e..23907b421 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -344,7 +344,12 @@ def model_test_runner(self, model, model_type, vocab_size, num_layers): self.assertEqual(model.model_type, model_type) for t in [mx.float32, mx.float16]: - model.update(tree_map(lambda p: p.astype(t), model.parameters())) + model.update( + tree_map( + lambda p: p.astype(t) if mx.issubdtype(p.dtype, mx.floating) else p, + model.parameters(), + ) + ) inputs = mx.array([[0, 1]]) outputs = model(inputs) @@ -1422,6 +1427,132 @@ def test_deepseek_v3(self): model, args.model_type, args.vocab_size, args.num_hidden_layers ) + def test_deepseek_v4(self): + from mlx_lm.models import deepseek_v4 + from mlx_lm.models.hyper_connection import ( + _hc_split_sinkhorn_ops, + hc_expand, + ) + + # RoPE test + rope = deepseek_v4.DeepseekV4RoPE(4, 10000) + x = mx.random.uniform(shape=(1, 2, 3, 4)) + y = rope(x, offset=1) + + inv_freq = 1.0 / (10000 ** (mx.arange(0, 4, 2, dtype=mx.float32) / 4)) + freqs = mx.arange(1, 4, dtype=mx.float32)[:, None] * inv_freq[None, :] + cos = mx.cos(freqs).reshape(1, 1, 3, 2) + sin = mx.sin(freqs).reshape(1, 1, 3, 2) + pairs = x.reshape(*x.shape[:-1], 2, 2) + x0, x1 = pairs[..., 0], pairs[..., 1] + expected = mx.stack([x0 * cos - x1 * sin, x0 * sin + x1 * cos], axis=-1) + expected = expected.reshape(*expected.shape[:-2], 4) + self.assertTrue(mx.allclose(y, expected, rtol=1e-5, atol=1e-5)) + + # Inverse RoPE round-trip test + y_inv = rope(y, offset=1, inverse=True) + self.assertTrue(mx.allclose(y_inv, x, rtol=1e-5, atol=1e-5)) + + # freq_scale test (compressor strided positions: pool_base + t * R) + # With pool_base=4, R=4, positions are [4, 8, 12] = arange(3)*4 + 4 + # Equivalent to offset=4//4=1, freq_scale=4 + scaled_rope = deepseek_v4.DeepseekV4RoPE(4, 10000, freq_scale=4) + y_scaled = scaled_rope(x, offset=4) + positions = mx.array([4, 8, 12], dtype=mx.float32) + freqs = positions[:, None] * inv_freq[None, :] + cos = mx.cos(freqs).reshape(1, 1, 3, 2) + sin = mx.sin(freqs).reshape(1, 1, 3, 2) + expected = mx.stack([x0 * cos - x1 * sin, x0 * sin + x1 * cos], axis=-1) + expected = expected.reshape(*expected.shape[:-2], 4) + self.assertTrue(mx.allclose(y_scaled, expected, rtol=1e-5, atol=1e-5)) + + # HyperConnection Sinkhorn test + for hc_mult in (2, 4): + mix = (2 + hc_mult) * hc_mult + mixes = mx.random.normal((2, 3, mix), dtype=mx.float32) + scale = mx.array([1.2, 0.7, 1.1], dtype=mx.float32) + base = mx.random.normal((mix,), dtype=mx.float32) + pre, post, comb = _hc_split_sinkhorn_ops( + mixes, scale, base, hc_mult, 20, 1e-6 + ) + # Verify comb is approximately doubly stochastic + self.assertTrue( + mx.allclose(comb.sum(-1), mx.ones_like(comb.sum(-1)), atol=0.1) + ) + self.assertTrue( + mx.allclose(comb.sum(-2), mx.ones_like(comb.sum(-2)), atol=0.1) + ) + + # Expand test + post = mx.random.normal((2, 3, 4), dtype=mx.float32) + block_out = mx.random.normal((2, 3, 8), dtype=mx.bfloat16) + comb = mx.random.normal((2, 3, 4, 4), dtype=mx.float32) + residual = mx.random.normal((2, 3, 4, 8), dtype=mx.bfloat16) + expected = post[..., None] * block_out[:, :, None, :].astype(mx.float32) + expected = expected + mx.matmul( + comb.swapaxes(-1, -2), residual.astype(mx.float32) + ) + expected = expected.astype(block_out.dtype) + actual = hc_expand(block_out, residual, post, comb) + self.assertTrue(mx.allclose(actual, expected, rtol=1e-5, atol=1e-5)) + + # Model test + args = deepseek_v4.ModelArgs( + model_type="deepseek_v4", + vocab_size=128, + hidden_size=64, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=1, + q_lora_rank=16, + o_lora_rank=8, + o_groups=2, + head_dim=16, + qk_rope_head_dim=4, + sliding_window=16, + compress_ratios=[0, 0, 4, 0], + index_n_heads=4, + index_head_dim=8, + index_topk=4, + moe_intermediate_size=16, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + num_hash_layers=1, + hc_mult=2, + hc_sinkhorn_iters=2, + ) + model = deepseek_v4.Model(args) + + self.model_test_runner( + model, args.model_type, args.vocab_size, args.num_hidden_layers + ) + + # Compressed layers use the compressed RoPE base for local and pooled KV. + main_layer = model.model.layers[0].attn + compress_layer = model.model.layers[2].attn + + inputs = mx.array([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=mx.int32) + cache = model.make_cache() + logits = model(inputs, cache=cache) + mx.eval(logits, [c.state for c in cache]) + + # Sanitize test + weight = mx.to_fp8(mx.ones((128, 128), dtype=mx.float32)) + converted = model.sanitize( + { + "layers.0.attn.wkv.weight": weight, + "layers.0.attn.wkv.scale": mx.full((1, 1), 127, dtype=mx.uint8), + } + ) + wkey = "model.layers.0.attn.wkv.weight" + skey = "model.layers.0.attn.wkv.scales" + self.assertIn(wkey, converted) + self.assertIn(skey, converted) + self.assertTrue(mx.all(converted[wkey] == weight.view(mx.uint32))) + self.assertTrue(mx.all(converted[skey] == 127)) + self.assertEqual(converted[skey].shape, (128, 4)) + def test_gemma2(self): from mlx_lm.models import gemma2