From 2bfc36155b2edcb38a261113fc5077b48bb9dd8f Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Wed, 12 Nov 2025 12:45:37 +0100 Subject: [PATCH 01/54] fix input_embeddings prefill bug in generate_step --- mlx_lm/generate.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index 6ca8a73fb..e9d48e74e 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -418,7 +418,8 @@ def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None): prompt_processed_tokens = 0 prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens) while total_prompt_tokens - prompt_processed_tokens > 1: - n_to_process = min(prefill_step_size, prompt.size - 1) + remaining = (total_prompt_tokens - prompt_processed_tokens) - 1 + n_to_process = min(prefill_step_size, remaining) _model_call( input_tokens=prompt[:n_to_process][None], input_embeddings=( From 5cf134cae379fe3b5771356d18d5d8b44deb9481 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Wed, 12 Nov 2025 12:56:42 +0100 Subject: [PATCH 02/54] format --- mlx_lm/generate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index e9d48e74e..1dda46c2f 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -418,7 +418,7 @@ def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None): prompt_processed_tokens = 0 prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens) while total_prompt_tokens - prompt_processed_tokens > 1: - remaining = (total_prompt_tokens - prompt_processed_tokens) - 1 + remaining = (total_prompt_tokens - prompt_processed_tokens) - 1 n_to_process = min(prefill_step_size, remaining) _model_call( input_tokens=prompt[:n_to_process][None], From 18c12cbb16090290af77b629bbaf949e04d73561 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 15:25:30 +0200 Subject: [PATCH 03/54] add DS4 --- mlx_lm/models/deepseek_v4.py | 1033 ++++++++++++++++++++++++++++++++++ mlx_lm/utils.py | 49 +- tests/test_models.py | 242 ++++++++ 3 files changed, 1323 insertions(+), 1 deletion(-) create mode 100644 mlx_lm/models/deepseek_v4.py diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py new file mode 100644 index 000000000..af367a8ef --- /dev/null +++ b/mlx_lm/models/deepseek_v4.py @@ -0,0 +1,1033 @@ +# Copyright © 2026 Apple Inc. + +import math +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +import mlx.core as mx +import mlx.nn as nn +from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients + +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .cache import RotatingKVCache +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" + quantization_config: Optional[Dict] = None + + 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 _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}") + + +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, + ): + super().__init__() + self.dims = dims + self.attention_scaling = 1.0 + + 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 + + attention_factor = scaling_config.get("attention_factor") + if attention_factor is None: + factor = factor or max_position_embeddings / original_max_position_embeddings + attention_factor = 1.0 if factor <= 1 else 0.1 * math.log(factor) + 1.0 + self.attention_scaling = attention_factor + elif rope_type not in (None, "default"): + raise ValueError(f"Unsupported DeepSeek-V4 RoPE type: {rope_type}") + + self._inv_freq = (inv_freq,) + + @property + def inv_freq(self): + return self._inv_freq[0] + + def __call__( + self, + x: mx.array, + offset: int = 0, + inverse: bool = False, + positions: Optional[mx.array] = None, + ): + dtype = x.dtype + L = x.shape[-2] + pos = ( + mx.arange(offset, offset + L, dtype=mx.float32) + if positions is None + else positions.astype(mx.float32) + ) + freqs = pos[:, None] * self.inv_freq[None, :] + emb = mx.concatenate([freqs, freqs], axis=-1) + cos = mx.cos(emb) * self.attention_scaling + sin = mx.sin(emb) * self.attention_scaling + if inverse: + sin = -sin + + broadcast_shape = (1,) * (x.ndim - 2) + cos.shape + cos = cos.reshape(broadcast_shape).astype(dtype) + sin = sin.reshape(broadcast_shape).astype(dtype) + + x1, x2 = mx.split(x, [x.shape[-1] // 2], axis=-1) + return mx.concatenate( + [ + x1 * cos[..., : x1.shape[-1]] - x2 * sin[..., : x2.shape[-1]], + x2 * cos[..., x1.shape[-1] :] + + x1 * sin[..., x1.shape[-1] :], + ], + axis=-1, + ) + + +def _apply_partial_rope( + x: mx.array, + rope: DeepseekV4RoPE, + offset: int = 0, + inverse: bool = False, + positions: Optional[mx.array] = None, +) -> mx.array: + rope_dim = rope.dims + if x.shape[-1] == rope_dim: + return rope(x, offset=offset, inverse=inverse, positions=positions) + nope, pe = mx.split(x, [x.shape[-1] - rope_dim], axis=-1) + pe = rope(pe, offset=offset, inverse=inverse, positions=positions) + return mx.concatenate([nope, pe], axis=-1) + + +def hc_split_sinkhorn( + 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 = ( + mx.sigmoid( + mixes[..., hc_mult : 2 * hc_mult] * post_scale + + base[hc_mult : 2 * hc_mult] + ) + + eps + ) + comb = ( + mx.sigmoid( + mixes[..., 2 * hc_mult :].reshape(*mixes.shape[:-1], hc_mult, hc_mult) + * comb_scale + + base[2 * hc_mult :].reshape(hc_mult, hc_mult) + ) + + eps + ) + for _ in range(sinkhorn_iters): + comb = comb / (comb.sum(axis=-1, keepdims=True) + eps) + comb = comb / (comb.sum(axis=-2, keepdims=True) + eps) + return pre, post, comb + + +class HyperConnection(nn.Module): + def __init__(self, config: ModelArgs): + 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 compute_weights(self, x: mx.array): + B, L, H, D = x.shape + flat = x.reshape(B, L, H * D).astype(mx.float32) + rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + self.norm_eps) + mixes = (flat @ self.fn.T) * rsqrt + return hc_split_sinkhorn( + mixes, + self.scale, + self.base, + self.hc_mult, + self.sinkhorn_iters, + self.hc_eps, + ) + + def collapse(self, x: mx.array): + pre, post, comb = self.compute_weights(x) + collapsed = (pre[..., None] * x.astype(mx.float32)).sum(axis=2) + return collapsed.astype(x.dtype), post, comb + + def expand( + self, + block_out: mx.array, + residual: mx.array, + post: mx.array, + comb: mx.array, + ): + y = post[..., None] * block_out[:, :, None, :].astype(mx.float32) + y = y + mx.einsum( + "bsij,bsjd->bsid", + comb.astype(mx.float32), + residual.astype(mx.float32), + ) + return y.astype(block_out.dtype) + + +class HyperHead(nn.Module): + def __init__(self, config: ModelArgs): + 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): + B, L, H, D = x.shape + flat = x.reshape(B, L, H * D).astype(mx.float32) + rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + self.norm_eps) + mixes = (flat @ self.fn.T) * rsqrt + pre = mx.sigmoid(mixes * self.scale[0] + self.base) + self.hc_eps + return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.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): + flat = x.reshape(-1, self.hidden_dim) + logits = flat.astype(mx.float32) @ self.weight.T.astype(mx.float32) + scores = _score_func(logits, self.scoring_func) + + if self.hash: + if input_ids is None: + raise ValueError("DeepSeek-V4 hash routing requires input_ids.") + inds = self.tid2eid[input_ids.reshape(-1)].astype(mx.int32) + else: + biased = scores + self.e_score_correction_bias + inds = mx.argpartition(-biased, kth=self.top_k - 1, axis=-1)[ + ..., : self.top_k + ] + + weights = mx.take_along_axis(scores, inds, axis=-1) + if self.scoring_func != "softmax" and self.norm_topk_prob: + weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) + weights = weights * self.routed_scaling_factor + 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]).sum(axis=-2).reshape(x.shape) + 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 DeepseekV4Cache: + def __init__(self, sliding_window: int): + self.local = RotatingKVCache(max_size=sliding_window, keep=0) + self.compressor_state = {"buffer_kv": None, "buffer_gate": None, "pooled": None} + self.indexer_state = {"buffer_kv": None, "buffer_gate": None, "pooled": None} + + @property + def offset(self): + return self.local.offset + + @property + def keys(self): + return self.local.keys + + @keys.setter + def keys(self, value): + self.local.keys = value + + def update_and_fetch(self, keys, values): + return self.local.update_and_fetch(keys, values) + + def make_mask(self, *args, **kwargs): + return self.local.make_mask(*args, **kwargs) + + def is_trimmable(self): + return self.local.is_trimmable() + + def trim(self, n): + return self.local.trim(n) + + def size(self): + return self.local.size() + + def empty(self): + return self.local.empty() + + @property + def nbytes(self): + total = self.local.nbytes + for state in (self.compressor_state, self.indexer_state): + for value in state.values(): + if value is not None: + total += value.nbytes + return total + + def _branch_state(self, state_key: str): + return self.indexer_state if state_key == "indexer_state" else self.compressor_state + + def accumulate_windows( + self, + kv: mx.array, + gate: mx.array, + state_key: str, + ratio: int, + start_pos: int, + ): + state = self._branch_state(state_key) + buf_kv, buf_gate = state["buffer_kv"], state["buffer_gate"] + if buf_kv is not None and buf_kv.shape[1]: + kv = mx.concatenate([buf_kv, kv], axis=1) + gate = mx.concatenate([buf_gate, gate], axis=1) + usable = (kv.shape[1] // ratio) * ratio + state["buffer_kv"] = kv[:, usable:] + state["buffer_gate"] = gate[:, usable:] + pool_base = max(0, start_pos) - (buf_kv.shape[1] if buf_kv is not None else 0) + return kv[:, :usable], gate[:, :usable], pool_base + + def update_pool(self, new_pooled: mx.array, state_key: str) -> mx.array: + state = self._branch_state(state_key) + pool = state["pooled"] + if new_pooled.shape[1] > 0: + pool = new_pooled if pool is None else mx.concatenate([pool, new_pooled], axis=1) + state["pooled"] = pool + if pool is None: + pool = mx.zeros((new_pooled.shape[0], 0, new_pooled.shape[-1]), new_pooled.dtype) + return pool + + +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) + + def _overlap_transform(self, x: mx.array, fill_value: float): + B, W, R, _ = x.shape + out = mx.full((B, W, 2 * R, self.head_dim), fill_value, dtype=x.dtype) + out[:, :, R:] = x[:, :, :, self.head_dim :] + out[:, 1:, :R] = x[:, :-1, :, : self.head_dim] + return out + + def __call__( + self, + x: mx.array, + rope: DeepseekV4RoPE, + cache: Optional[DeepseekV4Cache], + start_pos: int, + state_key: str = "compressor_state", + ) -> mx.array: + B, _, _ = x.shape + kv = self.wkv(x) + gate = self.wgate(x) + if cache is None: + usable = (kv.shape[1] // self.compress_ratio) * self.compress_ratio + ready_kv, ready_gate = kv[:, :usable], gate[:, :usable] + pool_base = start_pos + else: + ready_kv, ready_gate, pool_base = cache.accumulate_windows( + kv, gate, state_key, self.compress_ratio, start_pos + ) + + if ready_kv.shape[1] == 0: + new_pooled = mx.zeros((B, 0, self.head_dim), dtype=x.dtype) + else: + W = ready_kv.shape[1] // self.compress_ratio + kv = ready_kv.reshape(B, W, self.compress_ratio, self.out_dim) + gate = ( + ready_gate.reshape(B, W, self.compress_ratio, self.out_dim) + + self.ape.astype(ready_gate.dtype) + ) + if self.overlap: + kv = self._overlap_transform(kv, 0.0) + gate = self._overlap_transform(gate, -float("inf")) + weights = mx.softmax(gate.astype(mx.float32), axis=2, precise=True).astype(kv.dtype) + new_pooled = (kv * weights).sum(axis=2) + new_pooled = self.norm(new_pooled.astype(x.dtype)) + positions = ( + mx.arange(new_pooled.shape[1], dtype=mx.float32) * self.compress_ratio + + pool_base + ) + new_pooled = _apply_partial_rope( + new_pooled[:, None], rope, positions=positions + ).squeeze(1) + + if cache is not None: + return cache.update_pool(new_pooled, state_key) + 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, + rope: DeepseekV4RoPE, + position_rope: DeepseekV4RoPE, + cache: Optional[DeepseekV4Cache], + start_pos: int, + ): + B, L, _ = x.shape + pooled = self.compressor( + x, rope, cache, start_pos, state_key="indexer_state" + ) + if pooled.shape[1] == 0: + return None + + offset = start_pos + q = self.wq_b(q_residual).reshape(B, L, self.n_heads, self.head_dim) + q = q.transpose(0, 2, 1, 3) + q = _apply_partial_rope(q, position_rope, 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) + k = min(self.index_topk, pooled.shape[1]) + return mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k] + + +class V4Attention(nn.Module): + 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.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_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 = nn.Linear( + self.n_heads * self.head_dim // config.o_groups, + config.o_groups * config.o_lora_rank, + bias=False, + ) + 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, + config.rope_scaling, + config.max_position_embeddings, + ) + self.compress_rope = DeepseekV4RoPE( + config.qk_rope_head_dim, + config.compress_rope_theta, + config.rope_scaling, + config.max_position_embeddings, + ) + if self.compress_ratio: + self.compressor = Compressor(config, self.compress_ratio, self.head_dim) + if self.compress_ratio == 4: + self.indexer = Indexer(config, self.compress_ratio) + + def _grouped_output_projection(self, out: mx.array) -> mx.array: + B, L = out.shape[:2] + group_feat = (self.n_heads * self.head_dim) // self.o_groups + out = out.reshape(B, L, self.o_groups, group_feat) + + if isinstance(self.wo_a, nn.QuantizedLinear): + pieces = [] + for group_idx in range(self.o_groups): + rows = slice( + group_idx * self.o_lora_rank, + (group_idx + 1) * self.o_lora_rank, + ) + biases = self.wo_a.biases[rows] if self.wo_a.biases is not None else None + y = mx.quantized_matmul( + out[:, :, group_idx, :], + self.wo_a.weight[rows], + scales=self.wo_a.scales[rows], + biases=biases, + transpose=True, + group_size=self.wo_a.group_size, + bits=self.wo_a.bits, + mode=self.wo_a.mode, + ) + if "bias" in self.wo_a: + y = y + self.wo_a.bias[rows] + pieces.append(y) + return mx.concatenate(pieces, axis=-1) + + weight = self.wo_a.weight.reshape(self.o_groups, self.o_lora_rank, group_feat) + out = mx.einsum("bsgd,grd->bsgr", out, weight) + out = out.reshape(B, L, self.o_groups * self.o_lora_rank) + if "bias" in self.wo_a: + out = out + self.wo_a.bias + return out + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + local_cache = cache + if isinstance(cache, DeepseekV4Cache): + local_cache = cache + + offset = local_cache.offset if local_cache is not None else 0 + 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 = q.transpose(0, 2, 1, 3) + kv = self.kv_norm(self.wkv(x)).reshape(B, L, 1, self.head_dim) + kv = kv.transpose(0, 2, 1, 3) + + q = _apply_partial_rope(q, self.rope, offset) + kv = _apply_partial_rope(kv, self.rope, offset) + + if local_cache is not None: + kv, _ = local_cache.update_and_fetch(kv, kv) + full_kv = kv + + if self.compress_ratio: + v4_cache = cache if isinstance(cache, DeepseekV4Cache) else None + pooled = self.compressor(x, self.compress_rope, v4_cache, offset) + if hasattr(self, "indexer") and pooled.shape[1] > 0: + topk = self.indexer( + x, q_residual, self.compress_rope, self.rope, v4_cache, offset + ) + if topk is not None: + expanded = mx.broadcast_to( + pooled[:, None, None, :, :], + (B, 1, L, pooled.shape[1], self.head_dim), + ) + idx = topk[:, None, :, :, None] + pooled = mx.take_along_axis( + expanded, + mx.broadcast_to(idx, idx.shape[:-1] + (self.head_dim,)), + axis=3, + ).reshape(B, 1, -1, self.head_dim) + else: + pooled = pooled[:, None] + else: + pooled = pooled[:, None] + full_kv = mx.concatenate([full_kv, pooled], axis=2) + + if mask is not None and full_kv.shape[2] > mask.shape[-1]: + pad = mx.ones(mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],), dtype=mask.dtype) + mask = mx.concatenate([mask, pad], axis=-1) + + out = scaled_dot_product_attention( + q, + full_kv, + full_kv, + cache=local_cache, + scale=self.scale, + mask=mask, + sinks=self.attn_sink.astype(q.dtype), + ) + out = _apply_partial_rope(out, self.rope, offset, inverse=True) + out = out.transpose(0, 2, 1, 3).reshape(B, L, self.n_heads * self.head_dim) + out = self._grouped_output_projection(out) + return self.wo_b(out) + + +class DeepseekV4Block(nn.Module): + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__() + self.attn = V4Attention(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.collapse(h) + x = self.attn(self.attn_norm(x), mask=mask, cache=cache) + h = self.attn_hc.expand(x, residual, post, comb) + + residual = h + x, post, comb = self.ffn_hc.collapse(h) + x = self.ffn(self.ffn_norm(x), input_ids) + return self.ffn_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.local if isinstance(first_cache, DeepseekV4Cache) 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, DeepseekV4Cache): + cache_item = cache_item.local + 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: + if layer.attn.compress_ratio: + caches.append(DeepseekV4Cache(self.args.sliding_window)) + else: + caches.append(RotatingKVCache(max_size=self.args.sliding_window)) + 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 + + def scale_to_float(scale: mx.array) -> mx.array: + if scale.dtype == mx.uint8: + return mx.exp((scale.astype(mx.float32) - 127.0) * math.log(2.0)) + return scale.astype(mx.float32) + + def dequant_fp8(weight: mx.array, scale: mx.array, block_size: int = 128): + weight = mx.from_fp8(weight, dtype=mx.bfloat16) + scale = scale_to_float(scale) + m, n = weight.shape + pad_m = (-m) % block_size + pad_n = (-n) % block_size + weight = mx.pad(weight, ((0, pad_m), (0, pad_n))) + weight = weight.reshape( + (m + pad_m) // block_size, + block_size, + (n + pad_n) // block_size, + block_size, + ) + weight = (weight * scale[:, None, :, None]).reshape(m + pad_m, n + pad_n) + return weight[:m, :n].astype(mx.bfloat16) + + def dequant_fp4(weight: mx.array, scale: mx.array, block_size: int = 32): + table = mx.array( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + 0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=mx.float32, + ) + packed = weight.astype(mx.uint8) + low = packed & 0x0F + high = (packed >> 4) & 0x0F + unpacked = mx.stack([mx.take(table, low), mx.take(table, high)], axis=-1) + unpacked = unpacked.reshape(weight.shape[0], weight.shape[1] * 2) + scale = mx.repeat(scale_to_float(scale), block_size, axis=-1) + return (unpacked * scale).astype(mx.bfloat16) + + new_weights = {} + for k, v in weights.items(): + 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[wk] = dequant_fp4(weight, v) + elif weight.dtype == mx.uint8: + new_weights[wk] = dequant_fp8(weight, v) + 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"), + ): + key0 = f"{prefix}.0.{src}.weight" + if key0 in weights: + stacked = [ + weights.pop(f"{prefix}.{e}.{src}.weight") + for e in range(self.args.n_routed_experts) + ] + weights[ + f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.weight" + ] = mx.stack(stacked) + + return weights + + def shard(self, group: Optional[mx.distributed.Group] = None): + group = group or mx.distributed.init() + N = group.size() + for layer in self.model.layers: + layer.attn.wq_b = shard_linear( + layer.attn.wq_b, "all-to-sharded", group=group + ) + layer.attn.wo_b = shard_linear( + layer.attn.wo_b, "sharded-to-all", group=group + ) + 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/utils.py b/mlx_lm/utils.py index ef3d266b9..8569e24bc 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 ( @@ -279,6 +280,47 @@ def load_config(model_path: Path) -> dict: return config +def _load_safetensors(path: str, allow_unsupported_dtypes: bool = False) -> dict: + try: + return mx.load(path) + except RuntimeError as e: + if "F8_E8M0" not in str(e) or not allow_unsupported_dtypes: + raise + load_error = e + + with open(path, "r+b") as f: + header_len = struct.unpack(" header_len: + raise RuntimeError( + f"Cannot reinterpret F8_E8M0 safetensors header 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, @@ -319,8 +361,13 @@ def load_model( raise FileNotFoundError(f"No safetensors found in {model_path}") weights = {} + allow_unsupported_dtypes = config.get("model_type") == "deepseek_v4" for wf in weight_files: - weights.update(mx.load(wf)) + weights.update( + _load_safetensors( + wf, allow_unsupported_dtypes=allow_unsupported_dtypes + ) + ) if (model_file := config.get("model_file")) is not None: spec = importlib.util.spec_from_file_location( diff --git a/tests/test_models.py b/tests/test_models.py index 6e1fcd96e..35b8c0770 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1422,6 +1422,248 @@ def test_deepseek_v3(self): model, args.model_type, args.vocab_size, args.num_hidden_layers ) + def test_deepseek_v4_rope(self): + from mlx_lm.models.deepseek_v4 import DeepseekV4RoPE + + rope = 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, :] + emb = mx.concatenate([freqs, freqs], axis=-1) + cos = mx.cos(emb).reshape(1, 1, 3, 4) + sin = mx.sin(emb).reshape(1, 1, 3, 4) + x1, x2 = mx.split(x, [2], axis=-1) + expected = mx.concatenate( + [ + x1 * cos[..., :2] - x2 * sin[..., :2], + x2 * cos[..., 2:] + x1 * sin[..., 2:], + ], + axis=-1, + ) + self.assertTrue(mx.allclose(y, expected, rtol=1e-5, atol=1e-5)) + + positions = mx.array([1, 5, 9], dtype=mx.float32) + y = rope(x, positions=positions) + freqs = positions[:, None] * inv_freq[None, :] + emb = mx.concatenate([freqs, freqs], axis=-1) + cos = mx.cos(emb).reshape(1, 1, 3, 4) + sin = mx.sin(emb).reshape(1, 1, 3, 4) + expected = mx.concatenate( + [ + x1 * cos[..., :2] - x2 * sin[..., :2], + x2 * cos[..., 2:] + x1 * sin[..., 2:], + ], + axis=-1, + ) + self.assertTrue(mx.allclose(y, expected, rtol=1e-5, atol=1e-5)) + + def test_deepseek_v4(self): + from mlx_lm.models import deepseek_v4 + + 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) + inputs = mx.array([[0, 1, 2, 3, 4]], dtype=mx.int32) + + outputs = model(inputs) + mx.eval(outputs) + self.assertEqual(outputs.shape, (1, 5, args.vocab_size)) + + cache = model.make_cache() + self.assertIsInstance(cache[0], RotatingKVCache) + self.assertIsInstance(cache[2], deepseek_v4.DeepseekV4Cache) + outputs = model(inputs[:, :3], cache=cache) + mx.eval(outputs) + self.assertEqual(outputs.shape, (1, 3, args.vocab_size)) + outputs = model(inputs[:, 3:4], cache=cache) + mx.eval(outputs) + self.assertEqual(outputs.shape, (1, 1, args.vocab_size)) + + def test_deepseek_v4_sanitize_dequantizes_fp8_blocks(self): + from mlx_lm.models import deepseek_v4 + + args = deepseek_v4.ModelArgs( + model_type="deepseek_v4", + vocab_size=128, + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + q_lora_rank=16, + o_lora_rank=8, + o_groups=2, + head_dim=16, + qk_rope_head_dim=4, + moe_intermediate_size=2, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + hc_mult=2, + hc_sinkhorn_iters=2, + compress_ratios=[0], + ) + model = deepseek_v4.Model(args) + 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), + } + ) + key = "model.layers.0.attn.wkv.weight" + self.assertIn(key, converted) + self.assertTrue( + mx.allclose( + converted[key].astype(mx.float32), + mx.ones((128, 128), dtype=mx.float32), + rtol=1e-5, + atol=1e-5, + ) + ) + + def test_deepseek_v4_load_model_supports_e8m0_scales(self): + try: + import json + import tempfile + from pathlib import Path + + import torch + from safetensors.torch import save_file + except ImportError: + self.skipTest("torch and safetensors are required for this test") + + if not hasattr(torch, "float8_e4m3fn") or not hasattr( + torch, "float8_e8m0fnu" + ): + self.skipTest("torch build does not expose required float8 dtypes") + + from mlx_lm.utils import load_model + + config = { + "model_type": "deepseek_v4", + "vocab_size": 128, + "hidden_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "q_lora_rank": 16, + "o_lora_rank": 8, + "o_groups": 2, + "head_dim": 16, + "qk_rope_head_dim": 4, + "moe_intermediate_size": 2, + "n_routed_experts": 2, + "n_shared_experts": 1, + "num_experts_per_tok": 1, + "hc_mult": 2, + "hc_sinkhorn_iters": 2, + "compress_ratios": [0], + } + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) + with open(path / "config.json", "w") as f: + json.dump(config, f) + save_file( + { + "layers.0.attn.wkv.weight": torch.ones( + (16, 32), dtype=torch.float32 + ).to(torch.float8_e4m3fn), + "layers.0.attn.wkv.scale": torch.ones( + (1, 1), dtype=torch.float32 + ).to(torch.float8_e8m0fnu), + }, + str(path / "model.safetensors"), + ) + + model, _ = load_model(path, lazy=True, strict=False) + loaded = model.parameters()["model"]["layers"][0]["attn"]["wkv"][ + "weight" + ] + self.assertTrue( + mx.allclose( + loaded.astype(mx.float32), + mx.ones((16, 32), dtype=mx.float32), + rtol=1e-5, + atol=1e-5, + ) + ) + + def test_deepseek_v4_sanitize_unpacks_fp4_experts(self): + from mlx_lm.models import deepseek_v4 + + args = deepseek_v4.ModelArgs( + model_type="deepseek_v4", + vocab_size=128, + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + q_lora_rank=16, + o_lora_rank=8, + o_groups=2, + head_dim=16, + qk_rope_head_dim=4, + moe_intermediate_size=2, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + hc_mult=2, + hc_sinkhorn_iters=2, + compress_ratios=[0], + ) + model = deepseek_v4.Model(args) + packed = mx.array([[0x21] * 16, [0xFE] * 16], dtype=mx.uint8) + converted = model.sanitize( + { + "layers.0.ffn.experts.0.w1.weight": packed, + "layers.0.ffn.experts.0.w1.scale": mx.full( + (2, 1), 127, dtype=mx.uint8 + ), + "layers.0.ffn.experts.1.w1.weight": packed, + "layers.0.ffn.experts.1.w1.scale": mx.full( + (2, 1), 127, dtype=mx.uint8 + ), + } + ) + key = "model.layers.0.ffn.switch_mlp.gate_proj.weight" + self.assertIn(key, converted) + self.assertEqual(converted[key].shape, (2, 2, 32)) + self.assertTrue( + mx.array_equal( + converted[key][0, 0, :4].astype(mx.float32), + mx.array([0.5, 1.0, 0.5, 1.0], dtype=mx.float32), + ) + ) + self.assertTrue( + mx.array_equal( + converted[key][0, 1, :4].astype(mx.float32), + mx.array([-4.0, -6.0, -4.0, -6.0], dtype=mx.float32), + ) + ) + def test_gemma2(self): from mlx_lm.models import gemma2 From 28d991ff6668c0c119caca760589724f93c0e62c Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 16:02:08 +0200 Subject: [PATCH 04/54] Fix DeepSeek V4 Flash generation math --- mlx_lm/models/deepseek_v4.py | 93 +++++++++++++++++++++--------------- tests/test_models.py | 33 +++++-------- 2 files changed, 66 insertions(+), 60 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index af367a8ef..99c4ffcc7 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -111,7 +111,6 @@ def __init__( ): super().__init__() self.dims = dims - self.attention_scaling = 1.0 inv_freq = 1.0 / (base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)) rope_type = None @@ -145,11 +144,6 @@ def correction_dim(num_rotations): smooth = 1 - mx.clip(ramp, 0, 1) inv_freq = inv_freq / factor * (1 - smooth) + inv_freq * smooth - attention_factor = scaling_config.get("attention_factor") - if attention_factor is None: - factor = factor or max_position_embeddings / original_max_position_embeddings - attention_factor = 1.0 if factor <= 1 else 0.1 * math.log(factor) + 1.0 - self.attention_scaling = attention_factor elif rope_type not in (None, "default"): raise ValueError(f"Unsupported DeepSeek-V4 RoPE type: {rope_type}") @@ -174,9 +168,8 @@ def __call__( else positions.astype(mx.float32) ) freqs = pos[:, None] * self.inv_freq[None, :] - emb = mx.concatenate([freqs, freqs], axis=-1) - cos = mx.cos(emb) * self.attention_scaling - sin = mx.sin(emb) * self.attention_scaling + cos = mx.cos(freqs) + sin = mx.sin(freqs) if inverse: sin = -sin @@ -184,15 +177,10 @@ def __call__( cos = cos.reshape(broadcast_shape).astype(dtype) sin = sin.reshape(broadcast_shape).astype(dtype) - x1, x2 = mx.split(x, [x.shape[-1] // 2], axis=-1) - return mx.concatenate( - [ - x1 * cos[..., : x1.shape[-1]] - x2 * sin[..., : x2.shape[-1]], - x2 * cos[..., x1.shape[-1] :] - + x1 * sin[..., x1.shape[-1] :], - ], - axis=-1, - ) + x = x.reshape(*x.shape[:-1], x.shape[-1] // 2, 2) + x0, x1 = x[..., 0], x[..., 1] + out = mx.stack([x0 * cos - x1 * sin, x0 * sin + x1 * cos], axis=-1) + return out.reshape(*out.shape[:-2], out.shape[-2] * 2) def _apply_partial_rope( @@ -224,22 +212,18 @@ def hc_split_sinkhorn( 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 = ( - mx.sigmoid( - mixes[..., hc_mult : 2 * hc_mult] * post_scale - + base[hc_mult : 2 * hc_mult] - ) - + eps + post = 2 * mx.sigmoid( + mixes[..., hc_mult : 2 * hc_mult] * post_scale + + base[hc_mult : 2 * hc_mult] ) comb = ( - mx.sigmoid( - mixes[..., 2 * hc_mult :].reshape(*mixes.shape[:-1], hc_mult, hc_mult) - * comb_scale - + base[2 * hc_mult :].reshape(hc_mult, hc_mult) - ) - + eps + mixes[..., 2 * hc_mult :].reshape(*mixes.shape[:-1], hc_mult, hc_mult) + * comb_scale + + base[2 * hc_mult :].reshape(hc_mult, hc_mult) ) - for _ in range(sinkhorn_iters): + 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 @@ -423,6 +407,38 @@ def keys(self): def keys(self, value): self.local.keys = value + @property + def state(self): + local_state = None if self.local.empty() else self.local.state + return ( + local_state, + tuple(self.compressor_state[k] for k in ("buffer_kv", "buffer_gate", "pooled")), + tuple(self.indexer_state[k] for k in ("buffer_kv", "buffer_gate", "pooled")), + ) + + @state.setter + def state(self, value): + local_state, compressor_state, indexer_state = value + if local_state is None: + self.local.keys = None + self.local.values = None + else: + self.local.state = local_state + self.compressor_state = dict( + zip(("buffer_kv", "buffer_gate", "pooled"), compressor_state) + ) + self.indexer_state = dict( + zip(("buffer_kv", "buffer_gate", "pooled"), indexer_state) + ) + + @property + def meta_state(self): + return self.local.meta_state + + @meta_state.setter + def meta_state(self, value): + self.local.meta_state = value + def update_and_fetch(self, keys, values): return self.local.update_and_fetch(keys, values) @@ -627,18 +643,15 @@ def __init__(self, config: ModelArgs, layer_idx: int): ) self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) + rope_theta = config.compress_rope_theta if self.compress_ratio else config.rope_theta + rope_scaling = config.rope_scaling if self.compress_ratio else None self.rope = DeepseekV4RoPE( config.qk_rope_head_dim, - config.rope_theta, - config.rope_scaling, - config.max_position_embeddings, - ) - self.compress_rope = DeepseekV4RoPE( - config.qk_rope_head_dim, - config.compress_rope_theta, - config.rope_scaling, + rope_theta, + rope_scaling, config.max_position_embeddings, ) + self.compress_rope = self.rope if self.compress_ratio: self.compressor = Compressor(config, self.compress_ratio, self.head_dim) if self.compress_ratio == 4: @@ -693,6 +706,8 @@ def __call__( offset = local_cache.offset if local_cache is not None else 0 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 = q * mx.rsqrt((q.astype(mx.float32) ** 2).mean(axis=-1, keepdims=True) + self.config.rms_norm_eps) + q = q.astype(x.dtype) q = q.transpose(0, 2, 1, 3) kv = self.kv_norm(self.wkv(x)).reshape(B, L, 1, self.head_dim) kv = kv.transpose(0, 2, 1, 3) diff --git a/tests/test_models.py b/tests/test_models.py index 35b8c0770..0d220583a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1431,32 +1431,21 @@ def test_deepseek_v4_rope(self): 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, :] - emb = mx.concatenate([freqs, freqs], axis=-1) - cos = mx.cos(emb).reshape(1, 1, 3, 4) - sin = mx.sin(emb).reshape(1, 1, 3, 4) - x1, x2 = mx.split(x, [2], axis=-1) - expected = mx.concatenate( - [ - x1 * cos[..., :2] - x2 * sin[..., :2], - x2 * cos[..., 2:] + x1 * sin[..., 2:], - ], - axis=-1, - ) + 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)) positions = mx.array([1, 5, 9], dtype=mx.float32) y = rope(x, positions=positions) freqs = positions[:, None] * inv_freq[None, :] - emb = mx.concatenate([freqs, freqs], axis=-1) - cos = mx.cos(emb).reshape(1, 1, 3, 4) - sin = mx.sin(emb).reshape(1, 1, 3, 4) - expected = mx.concatenate( - [ - x1 * cos[..., :2] - x2 * sin[..., :2], - x2 * cos[..., 2:] + x1 * sin[..., 2:], - ], - axis=-1, - ) + 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, expected, rtol=1e-5, atol=1e-5)) def test_deepseek_v4(self): @@ -1500,9 +1489,11 @@ def test_deepseek_v4(self): outputs = model(inputs[:, :3], cache=cache) mx.eval(outputs) self.assertEqual(outputs.shape, (1, 3, args.vocab_size)) + mx.eval([c.state for c in cache]) outputs = model(inputs[:, 3:4], cache=cache) mx.eval(outputs) self.assertEqual(outputs.shape, (1, 1, args.vocab_size)) + mx.eval([c.state for c in cache]) def test_deepseek_v4_sanitize_dequantizes_fp8_blocks(self): from mlx_lm.models import deepseek_v4 From be714fd6faa24f91aff4dbaa2b9b377e29e26b09 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 16:10:09 +0200 Subject: [PATCH 05/54] Generalize safetensors dtype fallback --- mlx_lm/utils.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 8569e24bc..bda108747 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -57,6 +57,9 @@ 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} @@ -280,11 +283,11 @@ def load_config(model_path: Path) -> dict: return config -def _load_safetensors(path: str, allow_unsupported_dtypes: bool = False) -> dict: +def _load_safetensors(path: str) -> dict: try: return mx.load(path) except RuntimeError as e: - if "F8_E8M0" not in str(e) or not allow_unsupported_dtypes: + if not any(dtype in str(e) for dtype in SAFETENSORS_DTYPE_FALLBACKS): raise load_error = e @@ -295,8 +298,11 @@ def _load_safetensors(path: str, allow_unsupported_dtypes: bool = False) -> dict changed = False for tensor_info in header.values(): - if isinstance(tensor_info, dict) and tensor_info.get("dtype") == "F8_E8M0": - tensor_info["dtype"] = "U8" + if not isinstance(tensor_info, dict): + continue + dtype = tensor_info.get("dtype") + if dtype in SAFETENSORS_DTYPE_FALLBACKS: + tensor_info["dtype"] = SAFETENSORS_DTYPE_FALLBACKS[dtype] changed = True if not changed: @@ -305,7 +311,7 @@ def _load_safetensors(path: str, allow_unsupported_dtypes: bool = False) -> dict patched_header = json.dumps(header, separators=(",", ":")).encode("utf-8") if len(patched_header) > header_len: raise RuntimeError( - f"Cannot reinterpret F8_E8M0 safetensors header in {path}: " + f"Cannot reinterpret unsupported safetensors dtype in {path}: " "patched header is larger than the original header." ) @@ -361,13 +367,8 @@ def load_model( raise FileNotFoundError(f"No safetensors found in {model_path}") weights = {} - allow_unsupported_dtypes = config.get("model_type") == "deepseek_v4" for wf in weight_files: - weights.update( - _load_safetensors( - wf, allow_unsupported_dtypes=allow_unsupported_dtypes - ) - ) + weights.update(_load_safetensors(wf)) if (model_file := config.get("model_file")) is not None: spec = importlib.util.spec_from_file_location( From a833f261001c206ba6798b56b6002fdf6358378e Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 16:25:23 +0200 Subject: [PATCH 06/54] Remove deprecated DeepSeek V4 tests for sanitization and model loading --- tests/test_models.py | 141 ------------------------------------------- 1 file changed, 141 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 0d220583a..53bf8fab3 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1495,29 +1495,6 @@ def test_deepseek_v4(self): self.assertEqual(outputs.shape, (1, 1, args.vocab_size)) mx.eval([c.state for c in cache]) - def test_deepseek_v4_sanitize_dequantizes_fp8_blocks(self): - from mlx_lm.models import deepseek_v4 - - args = deepseek_v4.ModelArgs( - model_type="deepseek_v4", - vocab_size=128, - hidden_size=32, - num_hidden_layers=1, - num_attention_heads=4, - q_lora_rank=16, - o_lora_rank=8, - o_groups=2, - head_dim=16, - qk_rope_head_dim=4, - moe_intermediate_size=2, - n_routed_experts=2, - n_shared_experts=1, - num_experts_per_tok=1, - hc_mult=2, - hc_sinkhorn_iters=2, - compress_ratios=[0], - ) - model = deepseek_v4.Model(args) weight = mx.to_fp8(mx.ones((128, 128), dtype=mx.float32)) converted = model.sanitize( { @@ -1536,124 +1513,6 @@ def test_deepseek_v4_sanitize_dequantizes_fp8_blocks(self): ) ) - def test_deepseek_v4_load_model_supports_e8m0_scales(self): - try: - import json - import tempfile - from pathlib import Path - - import torch - from safetensors.torch import save_file - except ImportError: - self.skipTest("torch and safetensors are required for this test") - - if not hasattr(torch, "float8_e4m3fn") or not hasattr( - torch, "float8_e8m0fnu" - ): - self.skipTest("torch build does not expose required float8 dtypes") - - from mlx_lm.utils import load_model - - config = { - "model_type": "deepseek_v4", - "vocab_size": 128, - "hidden_size": 32, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "q_lora_rank": 16, - "o_lora_rank": 8, - "o_groups": 2, - "head_dim": 16, - "qk_rope_head_dim": 4, - "moe_intermediate_size": 2, - "n_routed_experts": 2, - "n_shared_experts": 1, - "num_experts_per_tok": 1, - "hc_mult": 2, - "hc_sinkhorn_iters": 2, - "compress_ratios": [0], - } - - with tempfile.TemporaryDirectory() as tmpdir: - path = Path(tmpdir) - with open(path / "config.json", "w") as f: - json.dump(config, f) - save_file( - { - "layers.0.attn.wkv.weight": torch.ones( - (16, 32), dtype=torch.float32 - ).to(torch.float8_e4m3fn), - "layers.0.attn.wkv.scale": torch.ones( - (1, 1), dtype=torch.float32 - ).to(torch.float8_e8m0fnu), - }, - str(path / "model.safetensors"), - ) - - model, _ = load_model(path, lazy=True, strict=False) - loaded = model.parameters()["model"]["layers"][0]["attn"]["wkv"][ - "weight" - ] - self.assertTrue( - mx.allclose( - loaded.astype(mx.float32), - mx.ones((16, 32), dtype=mx.float32), - rtol=1e-5, - atol=1e-5, - ) - ) - - def test_deepseek_v4_sanitize_unpacks_fp4_experts(self): - from mlx_lm.models import deepseek_v4 - - args = deepseek_v4.ModelArgs( - model_type="deepseek_v4", - vocab_size=128, - hidden_size=32, - num_hidden_layers=1, - num_attention_heads=4, - q_lora_rank=16, - o_lora_rank=8, - o_groups=2, - head_dim=16, - qk_rope_head_dim=4, - moe_intermediate_size=2, - n_routed_experts=2, - n_shared_experts=1, - num_experts_per_tok=1, - hc_mult=2, - hc_sinkhorn_iters=2, - compress_ratios=[0], - ) - model = deepseek_v4.Model(args) - packed = mx.array([[0x21] * 16, [0xFE] * 16], dtype=mx.uint8) - converted = model.sanitize( - { - "layers.0.ffn.experts.0.w1.weight": packed, - "layers.0.ffn.experts.0.w1.scale": mx.full( - (2, 1), 127, dtype=mx.uint8 - ), - "layers.0.ffn.experts.1.w1.weight": packed, - "layers.0.ffn.experts.1.w1.scale": mx.full( - (2, 1), 127, dtype=mx.uint8 - ), - } - ) - key = "model.layers.0.ffn.switch_mlp.gate_proj.weight" - self.assertIn(key, converted) - self.assertEqual(converted[key].shape, (2, 2, 32)) - self.assertTrue( - mx.array_equal( - converted[key][0, 0, :4].astype(mx.float32), - mx.array([0.5, 1.0, 0.5, 1.0], dtype=mx.float32), - ) - ) - self.assertTrue( - mx.array_equal( - converted[key][0, 1, :4].astype(mx.float32), - mx.array([-4.0, -6.0, -4.0, -6.0], dtype=mx.float32), - ) - ) def test_gemma2(self): from mlx_lm.models import gemma2 From 28a83ead52a5164cab3eb6776708613d5b2f2551 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 16:27:02 +0200 Subject: [PATCH 07/54] Refactor DeepSeek V4 tests to consolidate and clarify test cases --- tests/test_models.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 53bf8fab3..05fce0369 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1422,10 +1422,11 @@ def test_deepseek_v3(self): model, args.model_type, args.vocab_size, args.num_hidden_layers ) - def test_deepseek_v4_rope(self): - from mlx_lm.models.deepseek_v4 import DeepseekV4RoPE + def test_deepseek_v4(self): + from mlx_lm.models import deepseek_v4 - rope = DeepseekV4RoPE(4, 10000) + # RoPE test + rope = deepseek_v4.DeepseekV4RoPE(4, 10000) x = mx.random.uniform(shape=(1, 2, 3, 4)) y = rope(x, offset=1) @@ -1448,9 +1449,7 @@ def test_deepseek_v4_rope(self): expected = expected.reshape(*expected.shape[:-2], 4) self.assertTrue(mx.allclose(y, expected, rtol=1e-5, atol=1e-5)) - def test_deepseek_v4(self): - from mlx_lm.models import deepseek_v4 - + # Model test args = deepseek_v4.ModelArgs( model_type="deepseek_v4", vocab_size=128, @@ -1483,6 +1482,7 @@ def test_deepseek_v4(self): mx.eval(outputs) self.assertEqual(outputs.shape, (1, 5, args.vocab_size)) + # Cache test cache = model.make_cache() self.assertIsInstance(cache[0], RotatingKVCache) self.assertIsInstance(cache[2], deepseek_v4.DeepseekV4Cache) From bc1c2858d7d448e0d2b7b875f51bf6793f14a924 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 16:28:44 +0200 Subject: [PATCH 08/54] format --- mlx_lm/utils.py | 1 - tests/test_models.py | 1 - 2 files changed, 2 deletions(-) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index bda108747..36114f05e 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -60,7 +60,6 @@ SAFETENSORS_DTYPE_FALLBACKS = {"F8_E8M0": "U8"} - def _parse_size(x): sizes = {"M": 1e6, "G": 1e9, "MB": 1e6, "GB": 1e9, "": 1} split = 0 diff --git a/tests/test_models.py b/tests/test_models.py index 05fce0369..8e55968ec 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1513,7 +1513,6 @@ def test_deepseek_v4(self): ) ) - def test_gemma2(self): from mlx_lm.models import gemma2 From f4f7b4dd0df9f701b7e47e9204fb3ce6c3e00bb4 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 16:33:56 +0200 Subject: [PATCH 09/54] keep experts quantized --- mlx_lm/models/deepseek_v4.py | 107 +++++++++++++++++++++++------------ 1 file changed, 70 insertions(+), 37 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 99c4ffcc7..50d67d9b9 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -61,9 +61,11 @@ class ModelArgs(BaseModelArgs): 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 = ( + [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( @@ -129,8 +131,7 @@ def correction_dim(num_rotations): return ( dims * math.log( - original_max_position_embeddings - / (num_rotations * 2 * math.pi) + original_max_position_embeddings / (num_rotations * 2 * math.pi) ) / (2 * math.log(base)) ) @@ -213,14 +214,11 @@ def hc_split_sinkhorn( 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) + 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)): @@ -412,8 +410,12 @@ def state(self): local_state = None if self.local.empty() else self.local.state return ( local_state, - tuple(self.compressor_state[k] for k in ("buffer_kv", "buffer_gate", "pooled")), - tuple(self.indexer_state[k] for k in ("buffer_kv", "buffer_gate", "pooled")), + tuple( + self.compressor_state[k] for k in ("buffer_kv", "buffer_gate", "pooled") + ), + tuple( + self.indexer_state[k] for k in ("buffer_kv", "buffer_gate", "pooled") + ), ) @state.setter @@ -467,7 +469,11 @@ def nbytes(self): return total def _branch_state(self, state_key: str): - return self.indexer_state if state_key == "indexer_state" else self.compressor_state + return ( + self.indexer_state + if state_key == "indexer_state" + else self.compressor_state + ) def accumulate_windows( self, @@ -492,10 +498,16 @@ def update_pool(self, new_pooled: mx.array, state_key: str) -> mx.array: state = self._branch_state(state_key) pool = state["pooled"] if new_pooled.shape[1] > 0: - pool = new_pooled if pool is None else mx.concatenate([pool, new_pooled], axis=1) + pool = ( + new_pooled + if pool is None + else mx.concatenate([pool, new_pooled], axis=1) + ) state["pooled"] = pool if pool is None: - pool = mx.zeros((new_pooled.shape[0], 0, new_pooled.shape[-1]), new_pooled.dtype) + pool = mx.zeros( + (new_pooled.shape[0], 0, new_pooled.shape[-1]), new_pooled.dtype + ) return pool @@ -544,14 +556,15 @@ def __call__( else: W = ready_kv.shape[1] // self.compress_ratio kv = ready_kv.reshape(B, W, self.compress_ratio, self.out_dim) - gate = ( - ready_gate.reshape(B, W, self.compress_ratio, self.out_dim) - + self.ape.astype(ready_gate.dtype) - ) + gate = ready_gate.reshape( + B, W, self.compress_ratio, self.out_dim + ) + self.ape.astype(ready_gate.dtype) if self.overlap: kv = self._overlap_transform(kv, 0.0) gate = self._overlap_transform(gate, -float("inf")) - weights = mx.softmax(gate.astype(mx.float32), axis=2, precise=True).astype(kv.dtype) + weights = mx.softmax(gate.astype(mx.float32), axis=2, precise=True).astype( + kv.dtype + ) new_pooled = (kv * weights).sum(axis=2) new_pooled = self.norm(new_pooled.astype(x.dtype)) positions = ( @@ -590,9 +603,7 @@ def __call__( start_pos: int, ): B, L, _ = x.shape - pooled = self.compressor( - x, rope, cache, start_pos, state_key="indexer_state" - ) + pooled = self.compressor(x, rope, cache, start_pos, state_key="indexer_state") if pooled.shape[1] == 0: return None @@ -601,7 +612,9 @@ def __call__( q = q.transpose(0, 2, 1, 3) q = _apply_partial_rope(q, position_rope, offset) - scores = (q.astype(mx.float32) @ pooled[:, None].swapaxes(-1, -2).astype(mx.float32)) + 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) @@ -643,7 +656,9 @@ def __init__(self, config: ModelArgs, layer_idx: int): ) self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) - rope_theta = config.compress_rope_theta if self.compress_ratio else config.rope_theta + rope_theta = ( + config.compress_rope_theta if self.compress_ratio else config.rope_theta + ) rope_scaling = config.rope_scaling if self.compress_ratio else None self.rope = DeepseekV4RoPE( config.qk_rope_head_dim, @@ -669,7 +684,9 @@ def _grouped_output_projection(self, out: mx.array) -> mx.array: group_idx * self.o_lora_rank, (group_idx + 1) * self.o_lora_rank, ) - biases = self.wo_a.biases[rows] if self.wo_a.biases is not None else None + biases = ( + self.wo_a.biases[rows] if self.wo_a.biases is not None else None + ) y = mx.quantized_matmul( out[:, :, group_idx, :], self.wo_a.weight[rows], @@ -706,7 +723,10 @@ def __call__( offset = local_cache.offset if local_cache is not None else 0 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 = q * mx.rsqrt((q.astype(mx.float32) ** 2).mean(axis=-1, keepdims=True) + self.config.rms_norm_eps) + q = q * mx.rsqrt( + (q.astype(mx.float32) ** 2).mean(axis=-1, keepdims=True) + + self.config.rms_norm_eps + ) q = q.astype(x.dtype) q = q.transpose(0, 2, 1, 3) kv = self.kv_norm(self.wkv(x)).reshape(B, L, 1, self.head_dim) @@ -744,7 +764,9 @@ def __call__( full_kv = mx.concatenate([full_kv, pooled], axis=2) if mask is not None and full_kv.shape[2] > mask.shape[-1]: - pad = mx.ones(mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],), dtype=mask.dtype) + pad = mx.ones( + mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],), dtype=mask.dtype + ) mask = mx.concatenate([mask, pad], axis=-1) out = scaled_dot_product_attention( @@ -817,7 +839,11 @@ def __call__(self, inputs: mx.array, cache: Optional[Any] = None) -> mx.array: cache = [None] * len(self.pipeline_layers) first_cache = cache[0] - mask_cache = first_cache.local if isinstance(first_cache, DeepseekV4Cache) else first_cache + mask_cache = ( + first_cache.local + if isinstance(first_cache, DeepseekV4Cache) + else first_cache + ) mask = create_attention_mask( h[:, :, 0, :], mask_cache, @@ -873,6 +899,15 @@ def predicate(k): return predicate + @property + def quant_predicate(self): + def predicate(path, _): + if ".ffn.switch_mlp." in path: + return {"group_size": 32, "bits": 4} + return True + + return predicate + def make_cache(self): caches = [] for layer in self.layers: @@ -1013,9 +1048,9 @@ def dequant_fp4(weight: mx.array, scale: mx.array, block_size: int = 32): weights.pop(f"{prefix}.{e}.{src}.weight") for e in range(self.args.n_routed_experts) ] - weights[ - f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.weight" - ] = mx.stack(stacked) + weights[f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.weight"] = ( + mx.stack(stacked) + ) return weights @@ -1042,7 +1077,5 @@ def shard(self, group: Optional[mx.distributed.Group] = None): 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.down_proj, "sharded-to-all", group=group) shard_inplace(layer.ffn.switch_mlp.up_proj, "all-to-sharded", group=group) From e3221c45be84c5b133b99f0f1e05cb3b1509ce94 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 17:20:02 +0200 Subject: [PATCH 10/54] fix fp32 promotion (3 tps -> 11 tps) --- mlx_lm/models/deepseek_v4.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 50d67d9b9..a9dcc39fa 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -379,7 +379,7 @@ def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: inds, scores = self.gate(x, input_ids) y = self.switch_mlp(x, inds) - y = (y * scores[..., None]).sum(axis=-2).reshape(x.shape) + y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype).reshape(x.shape) y = y + self.shared_experts(x) if self.sharding_group is not None: From 16b268a940f89aeb0763fd1b4ba3c89c4a3e49a0 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 17:25:30 +0200 Subject: [PATCH 11/54] Fix DeepSeek V4 batched routing --- mlx_lm/models/deepseek_v4.py | 3 +++ tests/test_models.py | 21 ++++----------------- 2 files changed, 7 insertions(+), 17 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index a9dcc39fa..160ca7377 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -332,6 +332,9 @@ def __call__(self, x: mx.array, input_ids: Optional[mx.array] = None): if self.scoring_func != "softmax" and self.norm_topk_prob: weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) weights = weights * self.routed_scaling_factor + route_shape = (*x.shape[:-1], self.top_k) + inds = inds.reshape(route_shape) + weights = weights.reshape(route_shape) return inds, weights diff --git a/tests/test_models.py b/tests/test_models.py index 8e55968ec..d27b28d61 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1476,25 +1476,12 @@ def test_deepseek_v4(self): hc_sinkhorn_iters=2, ) model = deepseek_v4.Model(args) - inputs = mx.array([[0, 1, 2, 3, 4]], dtype=mx.int32) - outputs = model(inputs) - mx.eval(outputs) - self.assertEqual(outputs.shape, (1, 5, args.vocab_size)) - - # Cache test - cache = model.make_cache() - self.assertIsInstance(cache[0], RotatingKVCache) - self.assertIsInstance(cache[2], deepseek_v4.DeepseekV4Cache) - outputs = model(inputs[:, :3], cache=cache) - mx.eval(outputs) - self.assertEqual(outputs.shape, (1, 3, args.vocab_size)) - mx.eval([c.state for c in cache]) - outputs = model(inputs[:, 3:4], cache=cache) - mx.eval(outputs) - self.assertEqual(outputs.shape, (1, 1, args.vocab_size)) - mx.eval([c.state for c in cache]) + self.model_test_runner( + model, args.model_type, args.vocab_size, args.num_hidden_layers + ) + # Sanitize test weight = mx.to_fp8(mx.ones((128, 128), dtype=mx.float32)) converted = model.sanitize( { From 201267dc55f6954e99a56a372b57d0d9ba6ed287 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 19:08:52 +0200 Subject: [PATCH 12/54] compile _limited_swiglu, hc_split_sinkhorn and make robust predicate --- mlx_lm/models/deepseek_v4.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 160ca7377..44202e5ef 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -199,6 +199,7 @@ def _apply_partial_rope( return mx.concatenate([nope, pe], axis=-1) +@mx.compile def hc_split_sinkhorn( mixes: mx.array, scale: mx.array, From f26e519985d821d3972bbdc3b61da62412b748f7 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 19:12:48 +0200 Subject: [PATCH 13/54] Enhance quant_predicate to support new projection modes and add @mx.compile decorator to _limited_swiglu function --- mlx_lm/models/deepseek_v4.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 44202e5ef..07370109b 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -86,7 +86,7 @@ def _score_func(scores: mx.array, func: str) -> mx.array: return mx.sqrt(nn.softplus(scores)) raise ValueError(f"Unsupported DeepSeek-V4 scoring function: {func}") - +@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) @@ -198,7 +198,6 @@ def _apply_partial_rope( pe = rope(pe, offset=offset, inverse=inverse, positions=positions) return mx.concatenate([nope, pe], axis=-1) - @mx.compile def hc_split_sinkhorn( mixes: mx.array, @@ -906,8 +905,14 @@ def predicate(k): @property def quant_predicate(self): def predicate(path, _): - if ".ffn.switch_mlp." in path: - return {"group_size": 32, "bits": 4} + if path.endswith( + ( + ".ffn.switch_mlp.gate_proj", + ".ffn.switch_mlp.up_proj", + ".ffn.switch_mlp.down_proj", + ) + ): + return {"group_size": 32, "bits": 4, "mode": "mxfp4"} return True return predicate From 61e006ec6eba52960ef29a1867083dab518edb12 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 19:58:55 +0200 Subject: [PATCH 14/54] keep experts in mxfp4 --- mlx_lm/models/deepseek_v4.py | 71 ++++++++++++++++-------------------- mlx_lm/utils.py | 16 ++++++-- 2 files changed, 45 insertions(+), 42 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 07370109b..b2e20f9b9 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -14,6 +14,10 @@ from .switch_layers import SwitchGLU +def _default_quantization() -> Dict: + return {"group_size": 32, "bits": 4, "mode": "mxfp4"} + + @dataclass class ModelArgs(BaseModelArgs): model_type: str = "deepseek_v4" @@ -56,6 +60,7 @@ class ModelArgs(BaseModelArgs): num_nextn_predict_layers: int = 1 tie_word_embeddings: bool = False topk_method: str = "noaux_tc" + quantization: Optional[Dict] = field(default_factory=_default_quantization) quantization_config: Optional[Dict] = None def __post_init__(self): @@ -75,6 +80,15 @@ def __post_init__(self): 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}") + if self.quantization is None: + self.quantization = _default_quantization() + expert_quantization = _default_quantization() + for layer_idx in range(self.num_hidden_layers): + for proj in ("gate_proj", "up_proj", "down_proj"): + self.quantization.setdefault( + f"model.layers.{layer_idx}.ffn.switch_mlp.{proj}", + dict(expert_quantization), + ) def _score_func(scores: mx.array, func: str) -> mx.array: @@ -912,7 +926,7 @@ def predicate(path, _): ".ffn.switch_mlp.down_proj", ) ): - return {"group_size": 32, "bits": 4, "mode": "mxfp4"} + return _default_quantization() return True return predicate @@ -964,35 +978,12 @@ def dequant_fp8(weight: mx.array, scale: mx.array, block_size: int = 128): weight = (weight * scale[:, None, :, None]).reshape(m + pad_m, n + pad_n) return weight[:m, :n].astype(mx.bfloat16) - def dequant_fp4(weight: mx.array, scale: mx.array, block_size: int = 32): - table = mx.array( - [ - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, - 0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, - ], - dtype=mx.float32, - ) + def pack_fp4(weight: mx.array): packed = weight.astype(mx.uint8) - low = packed & 0x0F - high = (packed >> 4) & 0x0F - unpacked = mx.stack([mx.take(table, low), mx.take(table, high)], axis=-1) - unpacked = unpacked.reshape(weight.shape[0], weight.shape[1] * 2) - scale = mx.repeat(scale_to_float(scale), block_size, axis=-1) - return (unpacked * scale).astype(mx.bfloat16) + *dims, packed_in = packed.shape + packed = packed.reshape(*dims, packed_in // 4, 4).astype(mx.uint32) + shifts = mx.array([0, 8, 16, 24], dtype=mx.uint32) + return ((packed << shifts).sum(axis=-1)).astype(mx.uint32) new_weights = {} for k, v in weights.items(): @@ -1012,7 +1003,8 @@ def dequant_fp4(weight: mx.array, scale: mx.array, block_size: int = 32): and weight.dtype in (mx.int8, mx.uint8) and v.shape[-1] * 16 == weight.shape[-1] ): - new_weights[wk] = dequant_fp4(weight, v) + new_weights[wk] = pack_fp4(weight) + new_weights[wk[: -len(".weight")] + ".scales"] = v.astype(mx.uint8) elif weight.dtype == mx.uint8: new_weights[wk] = dequant_fp8(weight, v) else: @@ -1051,15 +1043,16 @@ def dequant_fp4(weight: mx.array, scale: mx.array, block_size: int = 32): ("w2", "down_proj"), ("w3", "up_proj"), ): - key0 = f"{prefix}.0.{src}.weight" - if key0 in weights: - stacked = [ - weights.pop(f"{prefix}.{e}.{src}.weight") - for e in range(self.args.n_routed_experts) - ] - weights[f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.weight"] = ( - mx.stack(stacked) - ) + 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) return weights diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 36114f05e..980b36601 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -392,14 +392,23 @@ def load_model( if hasattr(model, "sanitize"): weights = model.sanitize(weights) + if ( + config.get("quantization", None) is None + and getattr(model_args, "quantization", None) is not None + and any(k.endswith(".scales") for k in weights) + ): + config["quantization"] = model_args.quantization + def _quantize(quantization): def class_predicate(p, m): + if not hasattr(m, "to_quantized"): + return False + if f"{p}.scales" not in weights: + return False # Handle custom per layer quantizations if p in config["quantization"]: return config["quantization"][p] - if not hasattr(m, "to_quantized"): - return False - return f"{p}.scales" in weights + return True nn.quantize( model, @@ -863,6 +872,7 @@ def defaults_for_mode(mode, group_size, bits): # If the model is already partially quantized, return params so that # the config is set on a per-layer basis fine_grained_config = True + quantized_config["quantization"].update(quant_params) else: fine_grained_config = False quantized_config["quantization"] = quant_params From f2b4e9d9c629333ced61c123d14aa064b709c2c2 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 20:07:10 +0200 Subject: [PATCH 15/54] Refactor quantized linear layer --- mlx_lm/models/deepseek_v4.py | 53 ++++++++++++++++++++---------------- 1 file changed, 30 insertions(+), 23 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index b2e20f9b9..94081210a 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -695,29 +695,36 @@ def _grouped_output_projection(self, out: mx.array) -> mx.array: out = out.reshape(B, L, self.o_groups, group_feat) if isinstance(self.wo_a, nn.QuantizedLinear): - pieces = [] - for group_idx in range(self.o_groups): - rows = slice( - group_idx * self.o_lora_rank, - (group_idx + 1) * self.o_lora_rank, - ) - biases = ( - self.wo_a.biases[rows] if self.wo_a.biases is not None else None - ) - y = mx.quantized_matmul( - out[:, :, group_idx, :], - self.wo_a.weight[rows], - scales=self.wo_a.scales[rows], - biases=biases, - transpose=True, - group_size=self.wo_a.group_size, - bits=self.wo_a.bits, - mode=self.wo_a.mode, - ) - if "bias" in self.wo_a: - y = y + self.wo_a.bias[rows] - pieces.append(y) - return mx.concatenate(pieces, axis=-1) + out = out.transpose(2, 0, 1, 3) + weight = self.wo_a.weight.reshape(self.o_groups, self.o_lora_rank, -1)[ + :, None + ] + scales = self.wo_a.scales.reshape(self.o_groups, self.o_lora_rank, -1)[ + :, None + ] + biases = ( + None + if self.wo_a.biases is None + else self.wo_a.biases.reshape(self.o_groups, self.o_lora_rank, -1)[ + :, None + ] + ) + out = mx.quantized_matmul( + out, + weight, + scales=scales, + biases=biases, + transpose=True, + group_size=self.wo_a.group_size, + bits=self.wo_a.bits, + mode=self.wo_a.mode, + ) + out = out.transpose(1, 2, 0, 3).reshape( + B, L, self.o_groups * self.o_lora_rank + ) + if "bias" in self.wo_a: + out = out + self.wo_a.bias + return out weight = self.wo_a.weight.reshape(self.o_groups, self.o_lora_rank, group_feat) out = mx.einsum("bsgd,grd->bsgr", out, weight) From c0d9222d14ff0f21d0f88a16e9619767f2121030 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 20:27:44 +0200 Subject: [PATCH 16/54] Implement optimized Sinkhorn operations in DeepSeek V4 and add corresponding unit tests for HyperConnection --- mlx_lm/models/deepseek_v4.py | 139 ++++++++++++++++++++++++++++++++++- tests/test_models.py | 15 ++++ 2 files changed, 151 insertions(+), 3 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 94081210a..9b695bd12 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -213,7 +213,7 @@ def _apply_partial_rope( return mx.concatenate([nope, pe], axis=-1) @mx.compile -def hc_split_sinkhorn( +def _hc_split_sinkhorn_ops( mixes: mx.array, scale: mx.array, base: mx.array, @@ -241,12 +241,144 @@ def hc_split_sinkhorn( return pre, post, comb +def _make_hc_split_sinkhorn_kernel(): + if mx.default_device() != mx.gpu or not mx.metal.is_available(): + return None + + source = """ + uint idx = thread_position_in_grid.x; + constexpr int MIX = (2 + HC) * HC; + float epsv = static_cast(eps[0]); + + auto mix = mixes + idx * MIX; + auto pre_out = pre + idx * HC; + auto post_out = post + idx * HC; + auto comb_out = comb + idx * HC * HC; + + float pre_scale = static_cast(scale[0]); + float post_scale = static_cast(scale[1]); + float comb_scale = static_cast(scale[2]); + + for (int i = 0; i < HC; ++i) { + float z = static_cast(mix[i]) * pre_scale + + static_cast(base[i]); + pre_out[i] = 1.0f / (1.0f + metal::fast::exp(-z)) + epsv; + } + for (int i = 0; i < HC; ++i) { + int off = HC + i; + float z = static_cast(mix[off]) * post_scale + + static_cast(base[off]); + post_out[i] = 2.0f / (1.0f + metal::fast::exp(-z)); + } + + float c[HC * HC]; + for (int i = 0; i < HC; ++i) { + float row_max = -INFINITY; + for (int j = 0; j < HC; ++j) { + int cidx = i * HC + j; + int off = 2 * HC + cidx; + float v = static_cast(mix[off]) * comb_scale + + static_cast(base[off]); + c[cidx] = v; + row_max = metal::max(row_max, v); + } + float row_sum = 0.0f; + for (int j = 0; j < HC; ++j) { + int cidx = i * HC + j; + float v = metal::fast::exp(c[cidx] - row_max); + c[cidx] = v; + row_sum += v; + } + for (int j = 0; j < HC; ++j) { + int cidx = i * HC + j; + c[cidx] = c[cidx] / row_sum + epsv; + } + } + + for (int j = 0; j < HC; ++j) { + float col_sum = 0.0f; + for (int i = 0; i < HC; ++i) { + col_sum += c[i * HC + j]; + } + float denom = col_sum + epsv; + for (int i = 0; i < HC; ++i) { + c[i * HC + j] = c[i * HC + j] / denom; + } + } + + for (int iter = 1; iter < ITERS; ++iter) { + for (int i = 0; i < HC; ++i) { + float row_sum = 0.0f; + for (int j = 0; j < HC; ++j) { + row_sum += c[i * HC + j]; + } + float denom = row_sum + epsv; + for (int j = 0; j < HC; ++j) { + c[i * HC + j] = c[i * HC + j] / denom; + } + } + for (int j = 0; j < HC; ++j) { + float col_sum = 0.0f; + for (int i = 0; i < HC; ++i) { + col_sum += c[i * HC + j]; + } + float denom = col_sum + epsv; + for (int i = 0; i < HC; ++i) { + c[i * HC + j] = c[i * HC + j] / denom; + } + } + } + + for (int i = 0; i < HC * HC; ++i) { + comb_out[i] = c[i]; + } + """ + + return mx.fast.metal_kernel( + name="deepseek_v4_hc_split_sinkhorn", + input_names=["mixes", "scale", "base", "eps"], + output_names=["pre", "post", "comb"], + source=source, + ) + + +_hc_split_sinkhorn_kernel = _make_hc_split_sinkhorn_kernel() + + +def hc_split_sinkhorn( + mixes: mx.array, + scale: mx.array, + base: mx.array, + hc_mult: int, + sinkhorn_iters: int, + eps: float, +) -> Tuple[mx.array, mx.array, mx.array]: + if _hc_split_sinkhorn_kernel is None: + return _hc_split_sinkhorn_ops(mixes, scale, base, hc_mult, sinkhorn_iters, eps) + + if not isinstance(eps, mx.array): + eps = mx.array([eps], dtype=mx.float32) + return _hc_split_sinkhorn_kernel( + inputs=[mixes, scale, base, eps], + template=[("HC", hc_mult), ("ITERS", sinkhorn_iters)], + grid=(mixes.size // ((2 + hc_mult) * hc_mult), 1, 1), + threadgroup=(256, 1, 1), + output_shapes=[ + (*mixes.shape[:-1], hc_mult), + (*mixes.shape[:-1], hc_mult), + (*mixes.shape[:-1], hc_mult, hc_mult), + ], + output_dtypes=[mx.float32, mx.float32, mx.float32], + ) + + class HyperConnection(nn.Module): def __init__(self, config: ModelArgs): super().__init__() self.hc_mult = config.hc_mult self.sinkhorn_iters = config.hc_sinkhorn_iters self.hc_eps = config.hc_eps + self._hc_eps = (mx.array([config.hc_eps], dtype=mx.float32),) 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) @@ -258,13 +390,14 @@ def compute_weights(self, x: mx.array): flat = x.reshape(B, L, H * D).astype(mx.float32) rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + self.norm_eps) mixes = (flat @ self.fn.T) * rsqrt - return hc_split_sinkhorn( + split_sinkhorn = _hc_split_sinkhorn_ops if self.training else hc_split_sinkhorn + return split_sinkhorn( mixes, self.scale, self.base, self.hc_mult, self.sinkhorn_iters, - self.hc_eps, + self.hc_eps if self.training else self._hc_eps[0], ) def collapse(self, x: mx.array): diff --git a/tests/test_models.py b/tests/test_models.py index d27b28d61..4749f3afe 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1449,6 +1449,21 @@ def test_deepseek_v4(self): expected = expected.reshape(*expected.shape[:-2], 4) self.assertTrue(mx.allclose(y, 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) + expected = deepseek_v4._hc_split_sinkhorn_ops( + mixes, scale, base, hc_mult, 20, 1e-6 + ) + actual = deepseek_v4.hc_split_sinkhorn( + mixes, scale, base, hc_mult, 20, 1e-6 + ) + for x, y in zip(expected, actual): + self.assertTrue(mx.allclose(x, y, rtol=1e-5, atol=1e-5)) + # Model test args = deepseek_v4.ModelArgs( model_type="deepseek_v4", From ef8c95d6f91d777810938aedae585cec9da9b24d Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 21:15:45 +0200 Subject: [PATCH 17/54] Optimize matrix multiplication in HyperConnection by replacing einsum with matmul for improved performance --- mlx_lm/models/deepseek_v4.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 9b695bd12..af7f9ce8a 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -413,11 +413,7 @@ def expand( comb: mx.array, ): y = post[..., None] * block_out[:, :, None, :].astype(mx.float32) - y = y + mx.einsum( - "bsij,bsjd->bsid", - comb.astype(mx.float32), - residual.astype(mx.float32), - ) + y = y + mx.matmul(comb.astype(mx.float32), residual.astype(mx.float32)) return y.astype(block_out.dtype) From f7ff216231516d77b8c0d3a057743d7d6a71b0e6 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Fri, 24 Apr 2026 21:53:45 +0200 Subject: [PATCH 18/54] HC sinkhorn normalization --- mlx_lm/models/deepseek_v4.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index af7f9ce8a..c62993480 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -100,6 +100,7 @@ def _score_func(scores: mx.array, func: str) -> mx.array: return mx.sqrt(nn.softplus(scores)) raise ValueError(f"Unsupported DeepSeek-V4 scoring function: {func}") + @mx.compile def _limited_swiglu(gate: mx.array, up: mx.array, limit: float) -> mx.array: if limit and limit > 0: @@ -212,6 +213,7 @@ def _apply_partial_rope( pe = rope(pe, offset=offset, inverse=inverse, positions=positions) return mx.concatenate([nope, pe], axis=-1) + @mx.compile def _hc_split_sinkhorn_ops( mixes: mx.array, @@ -289,9 +291,10 @@ def _make_hc_split_sinkhorn_kernel(): c[cidx] = v; row_sum += v; } + float inv_sum = 1.0f / row_sum; for (int j = 0; j < HC; ++j) { int cidx = i * HC + j; - c[cidx] = c[cidx] / row_sum + epsv; + c[cidx] = c[cidx] * inv_sum + epsv; } } @@ -300,9 +303,9 @@ def _make_hc_split_sinkhorn_kernel(): for (int i = 0; i < HC; ++i) { col_sum += c[i * HC + j]; } - float denom = col_sum + epsv; + float inv_denom = 1.0f / (col_sum + epsv); for (int i = 0; i < HC; ++i) { - c[i * HC + j] = c[i * HC + j] / denom; + c[i * HC + j] *= inv_denom; } } @@ -312,9 +315,9 @@ def _make_hc_split_sinkhorn_kernel(): for (int j = 0; j < HC; ++j) { row_sum += c[i * HC + j]; } - float denom = row_sum + epsv; + float inv_denom = 1.0f / (row_sum + epsv); for (int j = 0; j < HC; ++j) { - c[i * HC + j] = c[i * HC + j] / denom; + c[i * HC + j] *= inv_denom; } } for (int j = 0; j < HC; ++j) { @@ -322,9 +325,9 @@ def _make_hc_split_sinkhorn_kernel(): for (int i = 0; i < HC; ++i) { col_sum += c[i * HC + j]; } - float denom = col_sum + epsv; + float inv_denom = 1.0f / (col_sum + epsv); for (int i = 0; i < HC; ++i) { - c[i * HC + j] = c[i * HC + j] / denom; + c[i * HC + j] *= inv_denom; } } } From c6a782867824237b918b0ca3b9be0513cfc6878f Mon Sep 17 00:00:00 2001 From: eauchs Date: Fri, 24 Apr 2026 23:27:41 +0200 Subject: [PATCH 19/54] fix(deepseek_v4): numerical stability and compressed attention mask padding - Replaced nn.softplus with logaddexp in _score_func for sqrtsoftplus to prevent numerical overflow in FP16/BF16. - Fixed compressed attention mask padding (changed mx.ones to mx.zeros) so the model attends to its compressed KV cache instead of masking it. Co-authored-by: machiabeli --- mlx_lm/models/deepseek_v4.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index c62993480..f02a1e612 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -97,7 +97,7 @@ def _score_func(scores: mx.array, func: str) -> mx.array: if func == "sigmoid": return mx.sigmoid(scores) if func == "sqrtsoftplus": - return mx.sqrt(nn.softplus(scores)) + return mx.sqrt(mx.logaddexp(scores, mx.zeros_like(scores))) raise ValueError(f"Unsupported DeepSeek-V4 scoring function: {func}") @@ -920,7 +920,7 @@ def __call__( full_kv = mx.concatenate([full_kv, pooled], axis=2) if mask is not None and full_kv.shape[2] > mask.shape[-1]: - pad = mx.ones( + pad = mx.zeros( mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],), dtype=mask.dtype ) mask = mx.concatenate([mask, pad], axis=-1) From 2a73e76f99a19062a5984361cd2c8422760be673 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 09:03:23 +0200 Subject: [PATCH 20/54] Refactor DeepSeek V4 scoring and compile rope full --- mlx_lm/models/deepseek_v4.py | 101 ++++++++++++++++++++++++----------- 1 file changed, 70 insertions(+), 31 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index f02a1e612..22d0e2228 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -97,7 +97,7 @@ def _score_func(scores: mx.array, func: str) -> mx.array: if func == "sigmoid": return mx.sigmoid(scores) if func == "sqrtsoftplus": - return mx.sqrt(mx.logaddexp(scores, mx.zeros_like(scores))) + return mx.sqrt(mx.logaddexp(scores, 0)) raise ValueError(f"Unsupported DeepSeek-V4 scoring function: {func}") @@ -176,27 +176,46 @@ def __call__( inverse: bool = False, positions: Optional[mx.array] = None, ): - dtype = x.dtype L = x.shape[-2] pos = ( mx.arange(offset, offset + L, dtype=mx.float32) if positions is None else positions.astype(mx.float32) ) - freqs = pos[:, None] * self.inv_freq[None, :] - cos = mx.cos(freqs) - sin = mx.sin(freqs) - if inverse: - sin = -sin + return _rope_full(x, self.inv_freq, pos, inverse, 0) - broadcast_shape = (1,) * (x.ndim - 2) + cos.shape - cos = cos.reshape(broadcast_shape).astype(dtype) - sin = sin.reshape(broadcast_shape).astype(dtype) - x = x.reshape(*x.shape[:-1], x.shape[-1] // 2, 2) - x0, x1 = x[..., 0], x[..., 1] - out = mx.stack([x0 * cos - x1 * sin, x0 * sin + x1 * cos], axis=-1) - return out.reshape(*out.shape[:-2], out.shape[-2] * 2) +@mx.compile +def _rope_full( + x: mx.array, + inv_freq: mx.array, + pos: mx.array, + inverse: bool, + nope_dim: int, +) -> mx.array: + freqs = pos[:, None] * inv_freq[None, :] + cos = mx.cos(freqs) + sin = mx.sin(freqs) + if inverse: + sin = -sin + + if nope_dim > 0: + pe = x[..., nope_dim:] + else: + pe = x + + broadcast_shape = (1,) * (pe.ndim - 2) + cos.shape + cos = cos.reshape(broadcast_shape).astype(pe.dtype) + sin = sin.reshape(broadcast_shape).astype(pe.dtype) + + pe = pe.reshape(*pe.shape[:-1], pe.shape[-1] // 2, 2) + x0, x1 = pe[..., 0], pe[..., 1] + out = mx.stack([x0 * cos - x1 * sin, x0 * sin + x1 * cos], axis=-1) + pe_out = out.reshape(*out.shape[:-2], out.shape[-2] * 2) + + if nope_dim > 0: + return mx.concatenate([x[..., :nope_dim], pe_out], axis=-1) + return pe_out def _apply_partial_rope( @@ -206,12 +225,14 @@ def _apply_partial_rope( inverse: bool = False, positions: Optional[mx.array] = None, ) -> mx.array: - rope_dim = rope.dims - if x.shape[-1] == rope_dim: - return rope(x, offset=offset, inverse=inverse, positions=positions) - nope, pe = mx.split(x, [x.shape[-1] - rope_dim], axis=-1) - pe = rope(pe, offset=offset, inverse=inverse, positions=positions) - return mx.concatenate([nope, pe], axis=-1) + nope_dim = x.shape[-1] - rope.dims if x.shape[-1] != rope.dims else 0 + L = x.shape[-2] + pos = ( + mx.arange(offset, offset + L, dtype=mx.float32) + if positions is None + else positions.astype(mx.float32) + ) + return _rope_full(x, rope.inv_freq, pos, inverse, nope_dim) @mx.compile @@ -375,6 +396,28 @@ def hc_split_sinkhorn( ) +@mx.compile +def _hc_collapse_op(pre: mx.array, x: mx.array) -> mx.array: + return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.dtype) + + +@mx.compile +def _hc_expand_op( + post: mx.array, + block_out: mx.array, + comb: mx.array, + residual: mx.array, +) -> mx.array: + y = post[..., None] * block_out[:, :, None, :].astype(mx.float32) + y = y + mx.matmul(comb, residual.astype(mx.float32)) + return y.astype(block_out.dtype) + + +@mx.compile +def _rms_rsqrt(flat: mx.array, eps: float) -> mx.array: + return mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + eps) + + class HyperConnection(nn.Module): def __init__(self, config: ModelArgs): super().__init__() @@ -391,7 +434,7 @@ def __init__(self, config: ModelArgs): def compute_weights(self, x: mx.array): B, L, H, D = x.shape flat = x.reshape(B, L, H * D).astype(mx.float32) - rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + self.norm_eps) + rsqrt = _rms_rsqrt(flat, self.norm_eps) mixes = (flat @ self.fn.T) * rsqrt split_sinkhorn = _hc_split_sinkhorn_ops if self.training else hc_split_sinkhorn return split_sinkhorn( @@ -405,8 +448,7 @@ def compute_weights(self, x: mx.array): def collapse(self, x: mx.array): pre, post, comb = self.compute_weights(x) - collapsed = (pre[..., None] * x.astype(mx.float32)).sum(axis=2) - return collapsed.astype(x.dtype), post, comb + return _hc_collapse_op(pre, x), post, comb def expand( self, @@ -415,9 +457,7 @@ def expand( post: mx.array, comb: mx.array, ): - y = post[..., None] * block_out[:, :, None, :].astype(mx.float32) - y = y + mx.matmul(comb.astype(mx.float32), residual.astype(mx.float32)) - return y.astype(block_out.dtype) + return _hc_expand_op(post, block_out, comb, residual) class HyperHead(nn.Module): @@ -435,7 +475,7 @@ def __init__(self, config: ModelArgs): def __call__(self, x: mx.array): B, L, H, D = x.shape flat = x.reshape(B, L, H * D).astype(mx.float32) - rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + self.norm_eps) + rsqrt = _rms_rsqrt(flat, self.norm_eps) mixes = (flat @ self.fn.T) * rsqrt pre = mx.sigmoid(mixes * self.scale[0] + self.base) + self.hc_eps return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.dtype) @@ -804,6 +844,7 @@ def __init__(self, config: ModelArgs, layer_idx: int): bias=config.attention_bias, ) self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) + self._q_l2_norm_weight = (mx.ones((self.head_dim,)),) rope_theta = ( config.compress_rope_theta if self.compress_ratio else config.rope_theta @@ -879,11 +920,9 @@ def __call__( offset = local_cache.offset if local_cache is not None else 0 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 = q * mx.rsqrt( - (q.astype(mx.float32) ** 2).mean(axis=-1, keepdims=True) - + self.config.rms_norm_eps + q = mx.fast.rms_norm( + q, self._q_l2_norm_weight[0].astype(q.dtype), self.config.rms_norm_eps ) - q = q.astype(x.dtype) q = q.transpose(0, 2, 1, 3) kv = self.kv_norm(self.wkv(x)).reshape(B, L, 1, self.head_dim) kv = kv.transpose(0, 2, 1, 3) From c2c9801c709d74f490a5a3933197a8035e74d325 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 10:49:39 +0200 Subject: [PATCH 21/54] Add DeepseekV4SwitchGLU and BatchRotatingKVCache; refactor position handling in DeepseekV4RoPE and DeepseekV4Cache --- mlx_lm/models/deepseek_v4.py | 302 ++++++++++++++++++++++++++++++++--- 1 file changed, 279 insertions(+), 23 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 22d0e2228..70635d5fb 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -9,9 +9,9 @@ from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention -from .cache import RotatingKVCache +from .cache import BatchRotatingKVCache, RotatingKVCache from .pipeline import PipelineMixin -from .switch_layers import SwitchGLU +from .switch_layers import SwitchGLU, _gather_sort, _scatter_unsort def _default_quantization() -> Dict: @@ -118,6 +118,33 @@ def __call__(self, x, gate): return _limited_swiglu(gate, x, self.limit) +class DeepseekV4SwitchGLU(SwitchGLU): + sort_threshold = 8 + + def __call__(self, x, indices) -> mx.array: + x = mx.expand_dims(x, (-2, -3)) + + do_sort = indices.size >= self.sort_threshold + idx = indices + inv_order = None + if do_sort: + x, idx, inv_order = _gather_sort(x, indices) + if self.training: + idx = mx.stop_gradient(idx) + x_up = self.up_proj(x, idx, sorted_indices=do_sort) + x_gate = self.gate_proj(x, idx, sorted_indices=do_sort) + x = self.down_proj( + self.activation(x_up, x_gate), + idx, + sorted_indices=do_sort, + ) + + if do_sort: + x = _scatter_unsort(x, inv_order, indices.shape) + + return x.squeeze(-2) + + class DeepseekV4RoPE(nn.Module): def __init__( self, @@ -172,16 +199,20 @@ def inv_freq(self): def __call__( self, x: mx.array, - offset: int = 0, + offset: Any = 0, inverse: bool = False, positions: Optional[mx.array] = None, ): L = x.shape[-2] - pos = ( - mx.arange(offset, offset + L, dtype=mx.float32) - if positions is None - else positions.astype(mx.float32) - ) + if positions is None: + pos = mx.arange(L, dtype=mx.float32) + if isinstance(offset, mx.array): + offset = offset.astype(mx.float32) + if offset.ndim > 0: + offset = offset[..., None] + pos = pos + offset + else: + pos = positions.astype(mx.float32) return _rope_full(x, self.inv_freq, pos, inverse, 0) @@ -193,7 +224,7 @@ def _rope_full( inverse: bool, nope_dim: int, ) -> mx.array: - freqs = pos[:, None] * inv_freq[None, :] + freqs = pos[..., None] * inv_freq cos = mx.cos(freqs) sin = mx.sin(freqs) if inverse: @@ -204,7 +235,12 @@ def _rope_full( else: pe = x - broadcast_shape = (1,) * (pe.ndim - 2) + cos.shape + if pos.ndim == 1: + broadcast_shape = (1,) * (pe.ndim - 2) + cos.shape + else: + broadcast_shape = ( + pos.shape[:-1] + (1,) * (pe.ndim - pos.ndim - 1) + cos.shape[-2:] + ) cos = cos.reshape(broadcast_shape).astype(pe.dtype) sin = sin.reshape(broadcast_shape).astype(pe.dtype) @@ -221,20 +257,23 @@ def _rope_full( def _apply_partial_rope( x: mx.array, rope: DeepseekV4RoPE, - offset: int = 0, + offset: Any = 0, inverse: bool = False, positions: Optional[mx.array] = None, ) -> mx.array: nope_dim = x.shape[-1] - rope.dims if x.shape[-1] != rope.dims else 0 L = x.shape[-2] - pos = ( - mx.arange(offset, offset + L, dtype=mx.float32) - if positions is None - else positions.astype(mx.float32) - ) + if positions is None: + pos = mx.arange(L, dtype=mx.float32) + if isinstance(offset, mx.array): + offset = offset.astype(mx.float32) + if offset.ndim > 0: + offset = offset[..., None] + pos = pos + offset + else: + pos = positions.astype(mx.float32) return _rope_full(x, rope.inv_freq, pos, inverse, nope_dim) - @mx.compile def _hc_split_sinkhorn_ops( mixes: mx.array, @@ -550,7 +589,7 @@ def __init__(self, config: ModelArgs, layer_idx: int): super().__init__() self.config = config self.gate = MoEGate(config, layer_idx) - self.switch_mlp = SwitchGLU( + self.switch_mlp = DeepseekV4SwitchGLU( config.hidden_size, config.moe_intermediate_size, config.n_routed_experts, @@ -577,6 +616,8 @@ def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: class DeepseekV4Cache: + _state_keys = ("buffer_kv", "buffer_gate", "pooled") + def __init__(self, sliding_window: int): self.local = RotatingKVCache(max_size=sliding_window, keep=0) self.compressor_state = {"buffer_kv": None, "buffer_gate": None, "pooled": None} @@ -680,7 +721,11 @@ def accumulate_windows( usable = (kv.shape[1] // ratio) * ratio state["buffer_kv"] = kv[:, usable:] state["buffer_gate"] = gate[:, usable:] - pool_base = max(0, start_pos) - (buf_kv.shape[1] if buf_kv is not None else 0) + if isinstance(start_pos, mx.array): + pool_base = mx.maximum(start_pos, 0) + else: + pool_base = max(0, start_pos) + pool_base = pool_base - (buf_kv.shape[1] if buf_kv is not None else 0) return kv[:, :usable], gate[:, :usable], pool_base def update_pool(self, new_pooled: mx.array, state_key: str) -> mx.array: @@ -699,6 +744,216 @@ def update_pool(self, new_pooled: mx.array, state_key: str) -> mx.array: ) return pool + def prepare(self, *args, **kwargs): + if hasattr(self.local, "prepare"): + self.local.prepare(*args, **kwargs) + return + + left_padding = kwargs.get("left_padding") + right_padding = kwargs.get("right_padding") + if left_padding is not None or ( + right_padding is not None and max(right_padding) > 0 + ): + batch_size = ( + len(left_padding) + if left_padding is not None + else len(kwargs.get("lengths")) + ) + self.local = BatchRotatingKVCache(self.local.max_size, [0] * batch_size) + self.local.prepare(*args, **kwargs) + + def finalize(self): + if hasattr(self.local, "finalize"): + self.local.finalize() + + def filter(self, batch_indices): + if hasattr(self.local, "filter"): + self.local.filter(batch_indices) + elif self.local.keys is not None: + self.local.keys = self.local.keys[batch_indices] + self.local.values = self.local.values[batch_indices] + for state in (self.compressor_state, self.indexer_state): + for key, value in state.items(): + if value is not None: + state[key] = value[batch_indices] + + def extend(self, other): + self_batch = self._cache_batch_size(self.local) + other_batch = self._cache_batch_size(other.local) + if hasattr(self.local, "extend"): + other_local = other.local + if not hasattr(other_local, "filter"): + other_local = self._batch_rotating_from_local(other_local) + self.local.extend(other_local) + elif ( + not hasattr(other.local, "filter") + and self.local.offset == other.local.offset + and self.local._idx == other.local._idx + ): + if self.local.keys is not None or other.local.keys is not None: + self.local.keys = self._concat_optional_local( + self.local.keys, other.local.keys + ) + self.local.values = self._concat_optional_local( + self.local.values, other.local.values + ) + else: + self.local = self._batch_rotating_from_local(self.local) + other_local = ( + other.local + if hasattr(other.local, "filter") + else self._batch_rotating_from_local(other.local) + ) + self.local.extend(other_local) + for self_state, other_state in ( + (self.compressor_state, other.compressor_state), + (self.indexer_state, other.indexer_state), + ): + for key in self._state_keys: + self_state[key] = self._concat_batch_state( + self_state[key], other_state[key], self_batch, other_batch + ) + + def extract(self, idx): + cache = DeepseekV4Cache(self.local.max_size) + cache.local = ( + self.local.extract(idx) + if hasattr(self.local, "extract") + else self._extract_local(self.local, idx) + ) + for cache_state, self_state in ( + (cache.compressor_state, self.compressor_state), + (cache.indexer_state, self.indexer_state), + ): + for key, value in self_state.items(): + cache_state[key] = ( + None if value is None else mx.contiguous(value[idx : idx + 1]) + ) + return cache + + @classmethod + def merge(cls, caches): + if not all(c.local.max_size == caches[0].local.max_size for c in caches): + raise ValueError( + "DeepseekV4Cache can only merge caches with the same sliding window" + ) + + cache = cls(caches[0].local.max_size) + cache.local = cls._merge_local([c.local for c in caches]) + for cache_state, state_name in ( + (cache.compressor_state, "compressor_state"), + (cache.indexer_state, "indexer_state"), + ): + for key in cls._state_keys: + cache_state[key] = cls._merge_batch_state( + [getattr(c, state_name)[key] for c in caches] + ) + return cache + + @staticmethod + def _cache_batch_size(cache): + offset = getattr(cache, "offset", 0) + if isinstance(offset, mx.array) and offset.ndim: + return offset.shape[0] + if cache.keys is not None: + return cache.keys.shape[0] + return 1 + + @staticmethod + def _extract_local(local, idx): + cache = RotatingKVCache(local.max_size, keep=getattr(local, "keep", 0)) + if local.keys is not None: + keys = local._temporal_order(local.keys) + values = local._temporal_order(local.values) + cache.keys = mx.contiguous(keys[idx : idx + 1]) + cache.values = mx.contiguous(values[idx : idx + 1]) + cache._idx = cache.keys.shape[2] + cache.offset = local.offset + return cache + + @classmethod + def _batch_rotating_from_local(cls, local): + batch_size = cls._cache_batch_size(local) + return BatchRotatingKVCache.merge( + [cls._extract_local(local, idx) for idx in range(batch_size)] + ) + + @staticmethod + def _concat_optional_local(a: Optional[mx.array], b: Optional[mx.array]): + if a is None: + return b + if b is None: + return a + return mx.concatenate([a, b], axis=0) + + @classmethod + def _merge_local(cls, locals): + offsets = [local.offset for local in locals] + sizes = [local.size() for local in locals] + use_fast_local = ( + all(not isinstance(offset, mx.array) for offset in offsets) + and all(offset == offsets[0] for offset in offsets) + and all(size == sizes[0] for size in sizes) + ) + if not use_fast_local: + if hasattr(locals[0], "merge"): + return locals[0].merge(locals) + return BatchRotatingKVCache.merge(locals) + + cache = RotatingKVCache(locals[0].max_size, keep=getattr(locals[0], "keep", 0)) + cache.offset = offsets[0] + if sizes[0] == 0: + return cache + + cache.keys = mx.concatenate( + [local._temporal_order(local.keys) for local in locals], axis=0 + ) + cache.values = mx.concatenate( + [local._temporal_order(local.values) for local in locals], axis=0 + ) + cache._idx = cache.keys.shape[2] + return cache + + @staticmethod + def _concat_batch_state( + a: Optional[mx.array], + b: Optional[mx.array], + a_batch: int, + b_batch: int, + ): + if a is None and b is None: + return None + if a is None: + shape = (a_batch, *b.shape[1:]) + a = mx.zeros(shape, dtype=b.dtype) + if b is None: + shape = (b_batch, *a.shape[1:]) + b = mx.zeros(shape, dtype=a.dtype) + if a.shape[1:] != b.shape[1:]: + raise ValueError( + "Cannot extend DeepseekV4Cache entries with different state shapes" + ) + return mx.concatenate([a, b], axis=0) + + @staticmethod + def _merge_batch_state(values: List[Optional[mx.array]]): + present = [v for v in values if v is not None] + if not present: + return None + if not all(v.shape[1:] == present[0].shape[1:] for v in present): + raise ValueError( + "Cannot batch DeepseekV4Cache entries with different state shapes" + ) + shape = present[0].shape + dtype = present[0].dtype + merged = [] + for value in values: + if value is None: + merged.append(mx.zeros((1, *shape[1:]), dtype=dtype)) + else: + merged.append(value) + return mx.concatenate(merged, axis=0) + class Compressor(nn.Module): def __init__(self, config: ModelArgs, compress_ratio: int, head_dim: int): @@ -756,10 +1011,11 @@ def __call__( ) new_pooled = (kv * weights).sum(axis=2) new_pooled = self.norm(new_pooled.astype(x.dtype)) - positions = ( - mx.arange(new_pooled.shape[1], dtype=mx.float32) * self.compress_ratio - + pool_base - ) + positions = mx.arange(new_pooled.shape[1], dtype=mx.float32) + positions = positions * self.compress_ratio + if isinstance(pool_base, mx.array) and pool_base.ndim > 0: + pool_base = pool_base[:, None] + positions = positions + pool_base new_pooled = _apply_partial_rope( new_pooled[:, None], rope, positions=positions ).squeeze(1) From 8ac37c0a76c9f16a15bf368ee2fe6399e04244df Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 11:02:40 +0200 Subject: [PATCH 22/54] Refactor DeepseekV4Cache to improve state management and add length handling for batching --- mlx_lm/models/deepseek_v4.py | 341 +++++++++++++++++++++++++++++++---- 1 file changed, 307 insertions(+), 34 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 70635d5fb..80b6048c7 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -274,6 +274,7 @@ def _apply_partial_rope( pos = positions.astype(mx.float32) return _rope_full(x, rope.inv_freq, pos, inverse, nope_dim) + @mx.compile def _hc_split_sinkhorn_ops( mixes: mx.array, @@ -617,11 +618,13 @@ def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: class DeepseekV4Cache: _state_keys = ("buffer_kv", "buffer_gate", "pooled") + _length_keys = ("buffer_lengths", "pooled_lengths") def __init__(self, sliding_window: int): self.local = RotatingKVCache(max_size=sliding_window, keep=0) - self.compressor_state = {"buffer_kv": None, "buffer_gate": None, "pooled": None} - self.indexer_state = {"buffer_kv": None, "buffer_gate": None, "pooled": None} + self.compressor_state = self._new_branch_state() + self.indexer_state = self._new_branch_state() + self._pending_lengths = None @property def offset(self): @@ -640,12 +643,8 @@ def state(self): local_state = None if self.local.empty() else self.local.state return ( local_state, - tuple( - self.compressor_state[k] for k in ("buffer_kv", "buffer_gate", "pooled") - ), - tuple( - self.indexer_state[k] for k in ("buffer_kv", "buffer_gate", "pooled") - ), + self._branch_state_tuple(self.compressor_state), + self._branch_state_tuple(self.indexer_state), ) @state.setter @@ -656,12 +655,8 @@ def state(self, value): self.local.values = None else: self.local.state = local_state - self.compressor_state = dict( - zip(("buffer_kv", "buffer_gate", "pooled"), compressor_state) - ) - self.indexer_state = dict( - zip(("buffer_kv", "buffer_gate", "pooled"), indexer_state) - ) + self.compressor_state = self._state_from_tuple(compressor_state) + self.indexer_state = self._state_from_tuple(indexer_state) @property def meta_state(self): @@ -694,7 +689,7 @@ def nbytes(self): total = self.local.nbytes for state in (self.compressor_state, self.indexer_state): for value in state.values(): - if value is not None: + if value is not None and hasattr(value, "nbytes"): total += value.nbytes return total @@ -705,6 +700,43 @@ def _branch_state(self, state_key: str): else self.compressor_state ) + @classmethod + def _new_branch_state(cls): + return {key: None for key in cls._state_keys + cls._length_keys} + + @classmethod + def _branch_state_tuple(cls, state): + return tuple(state[k] for k in cls._state_keys + cls._length_keys) + + @classmethod + def _state_from_tuple(cls, values): + keys = cls._state_keys + cls._length_keys + state = cls._new_branch_state() + state.update(zip(keys, values)) + return state + + @staticmethod + def _lengths_list(lengths, batch_size: int, default: Optional[int] = None): + if lengths is None: + if default is None: + return None + return [default] * batch_size + if isinstance(lengths, mx.array): + lengths = lengths.tolist() + return [int(length) for length in lengths] + + @staticmethod + def _filter_lengths(lengths, batch_indices): + if lengths is None: + return None + if isinstance(lengths, mx.array): + lengths = lengths.tolist() + if isinstance(batch_indices, mx.array): + batch_indices = batch_indices.tolist() + if len(lengths) == 1 and any(idx != 0 for idx in batch_indices): + lengths = lengths * (max(batch_indices) + 1) + return [int(lengths[idx]) for idx in batch_indices] + def accumulate_windows( self, kv: mx.array, @@ -715,22 +747,130 @@ def accumulate_windows( ): state = self._branch_state(state_key) buf_kv, buf_gate = state["buffer_kv"], state["buffer_gate"] - if buf_kv is not None and buf_kv.shape[1]: - kv = mx.concatenate([buf_kv, kv], axis=1) - gate = mx.concatenate([buf_gate, gate], axis=1) - usable = (kv.shape[1] // ratio) * ratio - state["buffer_kv"] = kv[:, usable:] - state["buffer_gate"] = gate[:, usable:] + B, L = kv.shape[:2] + buf_lengths = self._lengths_list(state["buffer_lengths"], B) + chunk_lengths = self._pending_lengths + if buf_lengths is None and chunk_lengths is None: + if buf_kv is not None and buf_kv.shape[1]: + kv = mx.concatenate([buf_kv, kv], axis=1) + gate = mx.concatenate([buf_gate, gate], axis=1) + usable = (kv.shape[1] // ratio) * ratio + state["buffer_kv"] = kv[:, usable:] + state["buffer_gate"] = gate[:, usable:] + state["buffer_lengths"] = None + state["_new_pooled_lengths"] = None + if isinstance(start_pos, mx.array): + pool_base = mx.maximum(start_pos, 0) + else: + pool_base = max(0, start_pos) + pool_base = pool_base - (buf_kv.shape[1] if buf_kv is not None else 0) + return kv[:, :usable], gate[:, :usable], pool_base + + buf_lengths = self._lengths_list( + state["buffer_lengths"], + B, + 0 if buf_kv is None else buf_kv.shape[1], + ) + chunk_lengths = self._lengths_list(chunk_lengths, B, L) + total_lengths = [ + buf_length + min(chunk_length, L) + for buf_length, chunk_length in zip(buf_lengths, chunk_lengths) + ] + usable_lengths = [(length // ratio) * ratio for length in total_lengths] + buffer_lengths = [ + length - usable for length, usable in zip(total_lengths, usable_lengths) + ] + max_total = max(total_lengths, default=0) + max_usable = max(usable_lengths, default=0) + max_buffer = max(buffer_lengths, default=0) + + combined_kv = mx.zeros((B, max_total, kv.shape[-1]), dtype=kv.dtype) + combined_gate = mx.zeros((B, max_total, gate.shape[-1]), dtype=gate.dtype) + for i, (buf_length, chunk_length, total_length) in enumerate( + zip(buf_lengths, chunk_lengths, total_lengths) + ): + parts_kv = [] + parts_gate = [] + if buf_length: + parts_kv.append(buf_kv[i : i + 1, :buf_length]) + parts_gate.append(buf_gate[i : i + 1, :buf_length]) + if chunk_length: + parts_kv.append(kv[i : i + 1, : min(chunk_length, L)]) + parts_gate.append(gate[i : i + 1, : min(chunk_length, L)]) + if parts_kv: + row_kv = ( + parts_kv[0] + if len(parts_kv) == 1 + else mx.concatenate(parts_kv, axis=1) + ) + row_gate = ( + parts_gate[0] + if len(parts_gate) == 1 + else mx.concatenate(parts_gate, axis=1) + ) + combined_kv[i : i + 1, :total_length] = row_kv + combined_gate[i : i + 1, :total_length] = row_gate + + ready_kv = combined_kv[:, :max_usable] + ready_gate = combined_gate[:, :max_usable] + state["buffer_kv"] = mx.zeros((B, max_buffer, kv.shape[-1]), dtype=kv.dtype) + state["buffer_gate"] = mx.zeros( + (B, max_buffer, gate.shape[-1]), dtype=gate.dtype + ) + for i, (usable, buffer_length) in enumerate( + zip(usable_lengths, buffer_lengths) + ): + if buffer_length: + state["buffer_kv"][i : i + 1, :buffer_length] = combined_kv[ + i : i + 1, usable : usable + buffer_length + ] + state["buffer_gate"][i : i + 1, :buffer_length] = combined_gate[ + i : i + 1, usable : usable + buffer_length + ] + state["buffer_lengths"] = buffer_lengths + state["_new_pooled_lengths"] = [ + usable // ratio for usable in usable_lengths + ] + + prev_lengths = mx.array(buf_lengths, dtype=mx.float32) if isinstance(start_pos, mx.array): - pool_base = mx.maximum(start_pos, 0) + pool_base = mx.maximum(start_pos, 0).astype(mx.float32) else: - pool_base = max(0, start_pos) - pool_base = pool_base - (buf_kv.shape[1] if buf_kv is not None else 0) - return kv[:, :usable], gate[:, :usable], pool_base + pool_base = mx.full((B,), max(0, start_pos), dtype=mx.float32) + return ready_kv, ready_gate, pool_base - prev_lengths def update_pool(self, new_pooled: mx.array, state_key: str) -> mx.array: state = self._branch_state(state_key) + new_lengths = state.pop("_new_pooled_lengths", None) pool = state["pooled"] + if new_lengths is not None: + B = new_pooled.shape[0] + pool_lengths = self._lengths_list( + state["pooled_lengths"], + B, + 0 if pool is None else pool.shape[1], + ) + total_lengths = [ + pool_length + new_length + for pool_length, new_length in zip(pool_lengths, new_lengths) + ] + max_total = max(total_lengths, default=0) + merged = mx.zeros( + (B, max_total, new_pooled.shape[-1]), dtype=new_pooled.dtype + ) + for i, (pool_length, new_length) in enumerate( + zip(pool_lengths, new_lengths) + ): + if pool is not None and pool_length: + merged[i : i + 1, :pool_length] = pool[i : i + 1, :pool_length] + if new_length: + merged[ + i : i + 1, pool_length : pool_length + new_length + ] = new_pooled[i : i + 1, :new_length] + state["pooled"] = merged + state["pooled_lengths"] = total_lengths + return merged + if new_pooled.shape[1] > 0: pool = ( new_pooled @@ -738,19 +878,29 @@ def update_pool(self, new_pooled: mx.array, state_key: str) -> mx.array: else mx.concatenate([pool, new_pooled], axis=1) ) state["pooled"] = pool + state["pooled_lengths"] = None if pool is None: pool = mx.zeros( (new_pooled.shape[0], 0, new_pooled.shape[-1]), new_pooled.dtype ) return pool + def pooled_lengths(self, state_key: str): + return self._branch_state(state_key)["pooled_lengths"] + def prepare(self, *args, **kwargs): + lengths = kwargs.get("lengths") + right_padding = kwargs.get("right_padding") + self._pending_lengths = ( + list(lengths) + if right_padding is not None and max(right_padding) > 0 + else None + ) if hasattr(self.local, "prepare"): self.local.prepare(*args, **kwargs) return left_padding = kwargs.get("left_padding") - right_padding = kwargs.get("right_padding") if left_padding is not None or ( right_padding is not None and max(right_padding) > 0 ): @@ -765,6 +915,7 @@ def prepare(self, *args, **kwargs): def finalize(self): if hasattr(self.local, "finalize"): self.local.finalize() + self._pending_lengths = None def filter(self, batch_indices): if hasattr(self.local, "filter"): @@ -773,9 +924,12 @@ def filter(self, batch_indices): self.local.keys = self.local.keys[batch_indices] self.local.values = self.local.values[batch_indices] for state in (self.compressor_state, self.indexer_state): - for key, value in state.items(): + for key in self._state_keys: + value = state[key] if value is not None: state[key] = value[batch_indices] + for key in self._length_keys: + state[key] = self._filter_lengths(state[key], batch_indices) def extend(self, other): self_batch = self._cache_batch_size(self.local) @@ -809,10 +963,21 @@ def extend(self, other): (self.compressor_state, other.compressor_state), (self.indexer_state, other.indexer_state), ): + self_tensors = {key: self_state[key] for key in self._state_keys} + other_tensors = {key: other_state[key] for key in self._state_keys} for key in self._state_keys: self_state[key] = self._concat_batch_state( self_state[key], other_state[key], self_batch, other_batch ) + for key in self._length_keys: + self_state[key] = self._concat_lengths( + self_state[key], + other_state[key], + self_tensors["pooled" if key.startswith("pooled") else "buffer_kv"], + other_tensors["pooled" if key.startswith("pooled") else "buffer_kv"], + self_batch, + other_batch, + ) def extract(self, idx): cache = DeepseekV4Cache(self.local.max_size) @@ -825,10 +990,19 @@ def extract(self, idx): (cache.compressor_state, self.compressor_state), (cache.indexer_state, self.indexer_state), ): - for key, value in self_state.items(): + for key in self._state_keys: + value = self_state[key] cache_state[key] = ( None if value is None else mx.contiguous(value[idx : idx + 1]) ) + for key in self._length_keys: + lengths = self_state[key] + if lengths is None: + cache_state[key] = None + elif isinstance(lengths, mx.array): + cache_state[key] = [int(lengths.tolist()[idx])] + else: + cache_state[key] = [int(lengths[idx])] return cache @classmethod @@ -848,6 +1022,12 @@ def merge(cls, caches): cache_state[key] = cls._merge_batch_state( [getattr(c, state_name)[key] for c in caches] ) + for key in cls._length_keys: + tensor_key = "pooled" if key.startswith("pooled") else "buffer_kv" + cache_state[key] = cls._merge_lengths( + [getattr(c, state_name)[key] for c in caches], + [getattr(c, state_name)[tensor_key] for c in caches], + ) return cache @staticmethod @@ -929,28 +1109,83 @@ def _concat_batch_state( if b is None: shape = (b_batch, *a.shape[1:]) b = mx.zeros(shape, dtype=a.dtype) - if a.shape[1:] != b.shape[1:]: + if a.shape[2:] != b.shape[2:]: raise ValueError( "Cannot extend DeepseekV4Cache entries with different state shapes" ) + if a.shape[1] != b.shape[1]: + seq_len = max(a.shape[1], b.shape[1]) + if a.shape[1] != seq_len: + padded = mx.zeros((a.shape[0], seq_len, *a.shape[2:]), dtype=a.dtype) + padded[:, : a.shape[1]] = a + a = padded + if b.shape[1] != seq_len: + padded = mx.zeros((b.shape[0], seq_len, *b.shape[2:]), dtype=b.dtype) + padded[:, : b.shape[1]] = b + b = padded return mx.concatenate([a, b], axis=0) + @staticmethod + def _full_lengths(lengths, value: Optional[mx.array], batch_size: int): + if lengths is not None: + if isinstance(lengths, mx.array): + lengths = lengths.tolist() + return [int(length) for length in lengths] + length = 0 if value is None else value.shape[1] + return [length] * batch_size + + @classmethod + def _concat_lengths( + cls, + a_lengths, + b_lengths, + a_value: Optional[mx.array], + b_value: Optional[mx.array], + a_batch: int, + b_batch: int, + ): + a_lengths = cls._full_lengths(a_lengths, a_value, a_batch) + b_lengths = cls._full_lengths(b_lengths, b_value, b_batch) + lengths = a_lengths + b_lengths + max_length = max(lengths, default=0) + return None if all(length == max_length for length in lengths) else lengths + + @classmethod + def _merge_lengths(cls, lengths, values): + batch_lengths = [ + cls._full_lengths(length, value, 1)[0] + for length, value in zip(lengths, values) + ] + max_length = max(batch_lengths, default=0) + return ( + None + if all(length == max_length for length in batch_lengths) + else batch_lengths + ) + @staticmethod def _merge_batch_state(values: List[Optional[mx.array]]): present = [v for v in values if v is not None] if not present: return None - if not all(v.shape[1:] == present[0].shape[1:] for v in present): + if not all(v.shape[2:] == present[0].shape[2:] for v in present): raise ValueError( "Cannot batch DeepseekV4Cache entries with different state shapes" ) + seq_len = max(v.shape[1] for v in present) shape = present[0].shape dtype = present[0].dtype merged = [] for value in values: if value is None: - merged.append(mx.zeros((1, *shape[1:]), dtype=dtype)) + merged.append(mx.zeros((1, seq_len, *shape[2:]), dtype=dtype)) else: + if value.shape[1] != seq_len: + padded = mx.zeros( + (value.shape[0], seq_len, *value.shape[2:]), dtype=value.dtype + ) + padded[:, : value.shape[1]] = value + value = padded merged.append(value) return mx.concatenate(merged, axis=0) @@ -1063,6 +1298,11 @@ def __call__( 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) + lengths = cache.pooled_lengths("indexer_state") if cache is not None else None + if lengths is not None: + lengths = mx.array(lengths) + valid = mx.arange(pooled.shape[1]) < lengths[:, None] + scores = mx.where(valid[:, 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] @@ -1174,6 +1414,8 @@ def __call__( local_cache = cache offset = local_cache.offset if local_cache is not None else 0 + if isinstance(offset, mx.array): + offset = offset + 0 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( @@ -1193,11 +1435,22 @@ def __call__( if self.compress_ratio: v4_cache = cache if isinstance(cache, DeepseekV4Cache) else None pooled = self.compressor(x, self.compress_rope, v4_cache, offset) + pooled_mask = None if hasattr(self, "indexer") and pooled.shape[1] > 0: topk = self.indexer( x, q_residual, self.compress_rope, self.rope, v4_cache, offset ) if topk is not None: + lengths = ( + v4_cache.pooled_lengths("compressor_state") + if v4_cache is not None + else None + ) + if lengths is not None: + lengths = mx.array(lengths) + pooled_mask = (topk < lengths[:, None, None]).reshape( + B, 1, 1, -1 + ) expanded = mx.broadcast_to( pooled[:, None, None, :, :], (B, 1, L, pooled.shape[1], self.head_dim), @@ -1211,13 +1464,33 @@ def __call__( else: pooled = pooled[:, None] else: + lengths = ( + v4_cache.pooled_lengths("compressor_state") + if v4_cache is not None + else None + ) + if lengths is not None: + lengths = mx.array(lengths) + pooled_mask = ( + mx.arange(pooled.shape[1]) < lengths[:, None] + ).reshape(B, 1, 1, -1) pooled = pooled[:, None] full_kv = mx.concatenate([full_kv, pooled], axis=2) if mask is not None and full_kv.shape[2] > mask.shape[-1]: - pad = mx.zeros( - mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],), dtype=mask.dtype - ) + pad_shape = mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],) + if mask.dtype == mx.bool_: + pad = mx.ones(pad_shape, dtype=mask.dtype) + if pooled_mask is not None: + pad = pad & pooled_mask + else: + pad = mx.zeros(pad_shape, dtype=mask.dtype) + if pooled_mask is not None: + pad = mx.where( + pooled_mask, + pad, + mx.full(pad_shape, mx.finfo(mask.dtype).min, dtype=mask.dtype), + ) mask = mx.concatenate([mask, pad], axis=-1) out = scaled_dot_product_attention( From 369fd37b753091318040eb8942cdd4ac1f68d478 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 11:02:59 +0200 Subject: [PATCH 23/54] format --- mlx_lm/models/deepseek_v4.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 80b6048c7..c101ecaf3 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -828,9 +828,7 @@ def accumulate_windows( i : i + 1, usable : usable + buffer_length ] state["buffer_lengths"] = buffer_lengths - state["_new_pooled_lengths"] = [ - usable // ratio for usable in usable_lengths - ] + state["_new_pooled_lengths"] = [usable // ratio for usable in usable_lengths] prev_lengths = mx.array(buf_lengths, dtype=mx.float32) if isinstance(start_pos, mx.array): @@ -864,9 +862,9 @@ def update_pool(self, new_pooled: mx.array, state_key: str) -> mx.array: if pool is not None and pool_length: merged[i : i + 1, :pool_length] = pool[i : i + 1, :pool_length] if new_length: - merged[ - i : i + 1, pool_length : pool_length + new_length - ] = new_pooled[i : i + 1, :new_length] + merged[i : i + 1, pool_length : pool_length + new_length] = ( + new_pooled[i : i + 1, :new_length] + ) state["pooled"] = merged state["pooled_lengths"] = total_lengths return merged @@ -974,7 +972,9 @@ def extend(self, other): self_state[key], other_state[key], self_tensors["pooled" if key.startswith("pooled") else "buffer_kv"], - other_tensors["pooled" if key.startswith("pooled") else "buffer_kv"], + other_tensors[ + "pooled" if key.startswith("pooled") else "buffer_kv" + ], self_batch, other_batch, ) From 166bba6cfce4a5a794e42293116f232380c672b8 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 12:49:11 +0200 Subject: [PATCH 24/54] Revert "keep experts in mxfp4" This reverts commit 61e006ec6eba52960ef29a1867083dab518edb12. --- mlx_lm/models/deepseek_v4.py | 71 ++++++++++++++++++++---------------- mlx_lm/utils.py | 16 ++------ 2 files changed, 42 insertions(+), 45 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index c101ecaf3..f0224aa22 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -14,10 +14,6 @@ from .switch_layers import SwitchGLU, _gather_sort, _scatter_unsort -def _default_quantization() -> Dict: - return {"group_size": 32, "bits": 4, "mode": "mxfp4"} - - @dataclass class ModelArgs(BaseModelArgs): model_type: str = "deepseek_v4" @@ -60,7 +56,6 @@ class ModelArgs(BaseModelArgs): num_nextn_predict_layers: int = 1 tie_word_embeddings: bool = False topk_method: str = "noaux_tc" - quantization: Optional[Dict] = field(default_factory=_default_quantization) quantization_config: Optional[Dict] = None def __post_init__(self): @@ -80,15 +75,6 @@ def __post_init__(self): 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}") - if self.quantization is None: - self.quantization = _default_quantization() - expert_quantization = _default_quantization() - for layer_idx in range(self.num_hidden_layers): - for proj in ("gate_proj", "up_proj", "down_proj"): - self.quantization.setdefault( - f"model.layers.{layer_idx}.ffn.switch_mlp.{proj}", - dict(expert_quantization), - ) def _score_func(scores: mx.array, func: str) -> mx.array: @@ -1633,7 +1619,7 @@ def predicate(path, _): ".ffn.switch_mlp.down_proj", ) ): - return _default_quantization() + return {"group_size": 32, "bits": 4, "mode": "mxfp4"} return True return predicate @@ -1685,12 +1671,35 @@ def dequant_fp8(weight: mx.array, scale: mx.array, block_size: int = 128): weight = (weight * scale[:, None, :, None]).reshape(m + pad_m, n + pad_n) return weight[:m, :n].astype(mx.bfloat16) - def pack_fp4(weight: mx.array): + def dequant_fp4(weight: mx.array, scale: mx.array, block_size: int = 32): + table = mx.array( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + 0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=mx.float32, + ) packed = weight.astype(mx.uint8) - *dims, packed_in = packed.shape - packed = packed.reshape(*dims, packed_in // 4, 4).astype(mx.uint32) - shifts = mx.array([0, 8, 16, 24], dtype=mx.uint32) - return ((packed << shifts).sum(axis=-1)).astype(mx.uint32) + low = packed & 0x0F + high = (packed >> 4) & 0x0F + unpacked = mx.stack([mx.take(table, low), mx.take(table, high)], axis=-1) + unpacked = unpacked.reshape(weight.shape[0], weight.shape[1] * 2) + scale = mx.repeat(scale_to_float(scale), block_size, axis=-1) + return (unpacked * scale).astype(mx.bfloat16) new_weights = {} for k, v in weights.items(): @@ -1710,8 +1719,7 @@ def pack_fp4(weight: mx.array): and weight.dtype in (mx.int8, mx.uint8) and v.shape[-1] * 16 == weight.shape[-1] ): - new_weights[wk] = pack_fp4(weight) - new_weights[wk[: -len(".weight")] + ".scales"] = v.astype(mx.uint8) + new_weights[wk] = dequant_fp4(weight, v) elif weight.dtype == mx.uint8: new_weights[wk] = dequant_fp8(weight, v) else: @@ -1750,16 +1758,15 @@ def pack_fp4(weight: mx.array): ("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) + key0 = f"{prefix}.0.{src}.weight" + if key0 in weights: + stacked = [ + weights.pop(f"{prefix}.{e}.{src}.weight") + for e in range(self.args.n_routed_experts) + ] + weights[f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.weight"] = ( + mx.stack(stacked) + ) return weights diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 980b36601..36114f05e 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -392,23 +392,14 @@ def load_model( if hasattr(model, "sanitize"): weights = model.sanitize(weights) - if ( - config.get("quantization", None) is None - and getattr(model_args, "quantization", None) is not None - and any(k.endswith(".scales") for k in weights) - ): - config["quantization"] = model_args.quantization - def _quantize(quantization): def class_predicate(p, m): - if not hasattr(m, "to_quantized"): - return False - if f"{p}.scales" not in weights: - return False # Handle custom per layer quantizations if p in config["quantization"]: return config["quantization"][p] - return True + if not hasattr(m, "to_quantized"): + return False + return f"{p}.scales" in weights nn.quantize( model, @@ -872,7 +863,6 @@ def defaults_for_mode(mode, group_size, bits): # If the model is already partially quantized, return params so that # the config is set on a per-layer basis fine_grained_config = True - quantized_config["quantization"].update(quant_params) else: fine_grained_config = False quantized_config["quantization"] = quant_params From cbb0b72dfe616d7ebe29dfcf912e60e74e3fe39d Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 12:55:05 +0200 Subject: [PATCH 25/54] Revert quant predicate (lower precision experts) --- mlx_lm/models/deepseek_v4.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index f0224aa22..2e9b8656d 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -1609,21 +1609,6 @@ def predicate(k): return predicate - @property - def quant_predicate(self): - def predicate(path, _): - if path.endswith( - ( - ".ffn.switch_mlp.gate_proj", - ".ffn.switch_mlp.up_proj", - ".ffn.switch_mlp.down_proj", - ) - ): - return {"group_size": 32, "bits": 4, "mode": "mxfp4"} - return True - - return predicate - def make_cache(self): caches = [] for layer in self.layers: From e3313d4cbfdcd7e1fa6dd733b2b477ed09b5919f Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 13:17:10 +0200 Subject: [PATCH 26/54] Remove quantization_config from ModelArgs in deepseek_v4.py --- mlx_lm/models/deepseek_v4.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 2e9b8656d..d24ba8575 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -56,7 +56,6 @@ class ModelArgs(BaseModelArgs): num_nextn_predict_layers: int = 1 tie_word_embeddings: bool = False topk_method: str = "noaux_tc" - quantization_config: Optional[Dict] = None def __post_init__(self): if not self.compress_ratios: From 16de966f7d42f4945957a33298fb0d410b4d23d8 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 16:12:32 +0200 Subject: [PATCH 27/54] Update deepseek_v4.py Co-authored-by: Tarjei Mandt --- mlx_lm/models/deepseek_v4.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index d24ba8575..79ba69677 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -434,7 +434,7 @@ def _hc_expand_op( residual: mx.array, ) -> mx.array: y = post[..., None] * block_out[:, :, None, :].astype(mx.float32) - y = y + mx.matmul(comb, residual.astype(mx.float32)) + y = y + mx.matmul(comb.swapaxes(-1, -2), residual.astype(mx.float32)) return y.astype(block_out.dtype) From 01bcd7d933863d91c0622f416ac75eb348d83d3b Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 19:51:40 +0200 Subject: [PATCH 28/54] Fix RoPE init Reported-by: Shinka-Man <183578543+Shinka-Man@users.noreply.github.com> --- mlx_lm/models/deepseek_v4.py | 15 ++++++++------- tests/test_models.py | 11 +++++++++++ 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 79ba69677..8f92ce847 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -1327,17 +1327,18 @@ def __init__(self, config: ModelArgs, layer_idx: int): self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) self._q_l2_norm_weight = (mx.ones((self.head_dim,)),) - rope_theta = ( - config.compress_rope_theta if self.compress_ratio else config.rope_theta - ) - rope_scaling = config.rope_scaling if self.compress_ratio else None self.rope = DeepseekV4RoPE( config.qk_rope_head_dim, - rope_theta, - rope_scaling, + config.rope_theta, + config.rope_scaling, + config.max_position_embeddings, + ) + self.compress_rope = DeepseekV4RoPE( + config.qk_rope_head_dim, + config.compress_rope_theta, + config.rope_scaling, config.max_position_embeddings, ) - self.compress_rope = self.rope if self.compress_ratio: self.compressor = Compressor(config, self.compress_ratio, self.head_dim) if self.compress_ratio == 4: diff --git a/tests/test_models.py b/tests/test_models.py index 4749f3afe..3c0c8ef1e 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1449,6 +1449,17 @@ def test_deepseek_v4(self): expected = expected.reshape(*expected.shape[:-2], 4) self.assertTrue(mx.allclose(y, expected, rtol=1e-5, atol=1e-5)) + # Compressed layers use a separate RoPE base only for the pooled branch. + compress_layer = model.model.layers[2].attn + freq_dims = mx.arange(0, args.qk_rope_head_dim, 2, dtype=mx.float32) + freq_dims = freq_dims / args.qk_rope_head_dim + main_expected = 1.0 / (args.rope_theta**freq_dims) + compress_expected = 1.0 / (args.compress_rope_theta**freq_dims) + self.assertTrue(mx.allclose(compress_layer.rope.inv_freq, main_expected)) + self.assertTrue( + mx.allclose(compress_layer.compress_rope.inv_freq, compress_expected) + ) + # HyperConnection Sinkhorn test for hc_mult in (2, 4): mix = (2 + hc_mult) * hc_mult From 8bd10c63ad3319857bdb48a3d3edfc28853de1be Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 20:23:36 +0200 Subject: [PATCH 29/54] Refactor RoPE initialization in V4Attention to streamline handling of compressed layers and update associated tests for consistency. --- mlx_lm/models/deepseek_v4.py | 15 +++++++-------- tests/test_models.py | 22 +++++++++++----------- 2 files changed, 18 insertions(+), 19 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 8f92ce847..79ba69677 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -1327,18 +1327,17 @@ def __init__(self, config: ModelArgs, layer_idx: int): self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) self._q_l2_norm_weight = (mx.ones((self.head_dim,)),) - self.rope = DeepseekV4RoPE( - config.qk_rope_head_dim, - config.rope_theta, - config.rope_scaling, - config.max_position_embeddings, + rope_theta = ( + config.compress_rope_theta if self.compress_ratio else config.rope_theta ) - self.compress_rope = DeepseekV4RoPE( + rope_scaling = config.rope_scaling if self.compress_ratio else None + self.rope = DeepseekV4RoPE( config.qk_rope_head_dim, - config.compress_rope_theta, - config.rope_scaling, + rope_theta, + rope_scaling, config.max_position_embeddings, ) + self.compress_rope = self.rope if self.compress_ratio: self.compressor = Compressor(config, self.compress_ratio, self.head_dim) if self.compress_ratio == 4: diff --git a/tests/test_models.py b/tests/test_models.py index 3c0c8ef1e..80d17d78a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1449,17 +1449,6 @@ def test_deepseek_v4(self): expected = expected.reshape(*expected.shape[:-2], 4) self.assertTrue(mx.allclose(y, expected, rtol=1e-5, atol=1e-5)) - # Compressed layers use a separate RoPE base only for the pooled branch. - compress_layer = model.model.layers[2].attn - freq_dims = mx.arange(0, args.qk_rope_head_dim, 2, dtype=mx.float32) - freq_dims = freq_dims / args.qk_rope_head_dim - main_expected = 1.0 / (args.rope_theta**freq_dims) - compress_expected = 1.0 / (args.compress_rope_theta**freq_dims) - self.assertTrue(mx.allclose(compress_layer.rope.inv_freq, main_expected)) - self.assertTrue( - mx.allclose(compress_layer.compress_rope.inv_freq, compress_expected) - ) - # HyperConnection Sinkhorn test for hc_mult in (2, 4): mix = (2 + hc_mult) * hc_mult @@ -1507,6 +1496,17 @@ def test_deepseek_v4(self): 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 + freq_dims = mx.arange(0, args.qk_rope_head_dim, 2, dtype=mx.float32) + freq_dims = freq_dims / args.qk_rope_head_dim + main_expected = 1.0 / (args.rope_theta**freq_dims) + compress_expected = 1.0 / (args.compress_rope_theta**freq_dims) + self.assertTrue(mx.allclose(main_layer.rope.inv_freq, main_expected)) + self.assertTrue(mx.allclose(compress_layer.rope.inv_freq, compress_expected)) + self.assertIs(compress_layer.compress_rope, compress_layer.rope) + # Sanitize test weight = mx.to_fp8(mx.ones((128, 128), dtype=mx.float32)) converted = model.sanitize( From 014a10b215cc119730979e0c9f33b9d92ce708f7 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sat, 25 Apr 2026 21:56:58 +0200 Subject: [PATCH 30/54] Improve DeepseekV4Cache and V4Attention with improved trimming logic and mask handling for compressed layers --- mlx_lm/models/deepseek_v4.py | 78 +++++++++++++++++++++++++++--------- 1 file changed, 59 insertions(+), 19 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 79ba69677..12e4add48 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -658,10 +658,29 @@ def make_mask(self, *args, **kwargs): return self.local.make_mask(*args, **kwargs) def is_trimmable(self): - return self.local.is_trimmable() + if not self.local.is_trimmable(): + return False + for state in (self.compressor_state, self.indexer_state): + pooled = state["pooled"] + if pooled is not None and pooled.shape[1] > 0: + return False + return True def trim(self, n): - return self.local.trim(n) + trimmed = self.local.trim(n) + batch_size = self._cache_batch_size(self.local) + for state in (self.compressor_state, self.indexer_state): + for key in ("buffer_kv", "buffer_gate"): + value = state[key] + if value is not None: + state[key] = value[:, : max(value.shape[1] - trimmed, 0)] + lengths = state["buffer_lengths"] + if lengths is not None: + state["buffer_lengths"] = [ + max(length - trimmed, 0) + for length in self._lengths_list(lengths, batch_size, 0) + ] + return trimmed def size(self): return self.local.size() @@ -1416,21 +1435,32 @@ def __call__( if local_cache is not None: kv, _ = local_cache.update_and_fetch(kv, kv) full_kv = kv + local_kv_len = kv.shape[2] + pooled_mask = None + pooled_bias = None if self.compress_ratio: v4_cache = cache if isinstance(cache, DeepseekV4Cache) else None pooled = self.compressor(x, self.compress_rope, v4_cache, offset) - pooled_mask = None - if hasattr(self, "indexer") and pooled.shape[1] > 0: + lengths = ( + v4_cache.pooled_lengths("compressor_state") + if v4_cache is not None + else None + ) + use_indexer = hasattr(self, "indexer") and pooled.shape[1] > 0 + select_all = ( + use_indexer + and lengths is None + and pooled.shape[1] <= self.indexer.index_topk + ) + if select_all: + pooled = pooled[:, None] + pooled_bias = math.log(L) + elif use_indexer: topk = self.indexer( x, q_residual, self.compress_rope, self.rope, v4_cache, offset ) if topk is not None: - lengths = ( - v4_cache.pooled_lengths("compressor_state") - if v4_cache is not None - else None - ) if lengths is not None: lengths = mx.array(lengths) pooled_mask = (topk < lengths[:, None, None]).reshape( @@ -1449,11 +1479,6 @@ def __call__( else: pooled = pooled[:, None] else: - lengths = ( - v4_cache.pooled_lengths("compressor_state") - if v4_cache is not None - else None - ) if lengths is not None: lengths = mx.array(lengths) pooled_mask = ( @@ -1462,17 +1487,32 @@ def __call__( pooled = pooled[:, None] full_kv = mx.concatenate([full_kv, pooled], axis=2) + if mask is not None and mask.shape[-1] > local_kv_len: + mask = mask[..., -local_kv_len:] + if mask is not None and full_kv.shape[2] > mask.shape[-1]: pad_shape = mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],) - if mask.dtype == mx.bool_: + pad_pooled_mask = pooled_mask + if pad_pooled_mask is not None and pad_pooled_mask.shape[-1] != pad_shape[-1]: + pad_pooled_mask = pad_pooled_mask[..., -pad_shape[-1] :] + if pooled_bias is not None: + dtype = q.dtype + if mask.dtype == mx.bool_: + mask = mx.where( + mask, + mx.array(0, dtype=dtype), + mx.full((), mx.finfo(dtype).min, dtype=dtype), + ) + pad = mx.full(pad_shape, pooled_bias, dtype=mask.dtype) + elif mask.dtype == mx.bool_: pad = mx.ones(pad_shape, dtype=mask.dtype) - if pooled_mask is not None: - pad = pad & pooled_mask + if pad_pooled_mask is not None: + pad = pad & pad_pooled_mask else: pad = mx.zeros(pad_shape, dtype=mask.dtype) - if pooled_mask is not None: + if pad_pooled_mask is not None: pad = mx.where( - pooled_mask, + pad_pooled_mask, pad, mx.full(pad_shape, mx.finfo(mask.dtype).min, dtype=mask.dtype), ) From 5d1a67aa366823d220e64da43dae11e239d251aa Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sun, 26 Apr 2026 01:12:13 +0200 Subject: [PATCH 31/54] Add scores parameter and updatie projection logic --- mlx_lm/models/deepseek_v4.py | 36 ++++++++++++++++++++---------------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 12e4add48..56cd0ed2c 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -11,7 +11,7 @@ from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention from .cache import BatchRotatingKVCache, RotatingKVCache from .pipeline import PipelineMixin -from .switch_layers import SwitchGLU, _gather_sort, _scatter_unsort +from .switch_layers import SwitchGLU, _scatter_unsort @dataclass @@ -106,28 +106,33 @@ def __call__(self, x, gate): class DeepseekV4SwitchGLU(SwitchGLU): sort_threshold = 8 - def __call__(self, x, indices) -> mx.array: + def __call__(self, x, indices, scores) -> mx.array: + out_shape = x.shape + route_shape = indices.shape x = mx.expand_dims(x, (-2, -3)) do_sort = indices.size >= self.sort_threshold - idx = indices inv_order = None if do_sort: - x, idx, inv_order = _gather_sort(x, indices) + flat_indices = indices.flatten() + order = mx.argsort(flat_indices) + inv_order = mx.argsort(order) + x = x.flatten(0, -3)[order // route_shape[-1]] + indices = flat_indices[order] + scores = scores.flatten()[order] if self.training: - idx = mx.stop_gradient(idx) - x_up = self.up_proj(x, idx, sorted_indices=do_sort) - x_gate = self.gate_proj(x, idx, sorted_indices=do_sort) - x = self.down_proj( - self.activation(x_up, x_gate), - idx, - sorted_indices=do_sort, - ) + indices = mx.stop_gradient(indices) + x_up = self.up_proj(x, indices, sorted_indices=do_sort) + x_gate = self.gate_proj(x, indices, sorted_indices=do_sort) + x = self.activation(x_up, x_gate) + x = x * scores.astype(x.dtype)[..., None, None] + x = self.down_proj(x, indices, sorted_indices=do_sort) if do_sort: - x = _scatter_unsort(x, inv_order, indices.shape) + x = _scatter_unsort(x, inv_order, route_shape) - return x.squeeze(-2) + x = x.squeeze(-2) + return x.sum(axis=-2).astype(x.dtype).reshape(out_shape) class DeepseekV4RoPE(nn.Module): @@ -592,8 +597,7 @@ def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: x = sum_gradients(self.sharding_group)(x) inds, scores = self.gate(x, input_ids) - y = self.switch_mlp(x, inds) - y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype).reshape(x.shape) + y = self.switch_mlp(x, inds, scores) y = y + self.shared_experts(x) if self.sharding_group is not None: From 26f49f5ec4e1b09bb167ac9980fdca9ba3c62cc4 Mon Sep 17 00:00:00 2001 From: "clandestine.eth" <96172957+0xClandestine@users.noreply.github.com> Date: Sun, 26 Apr 2026 11:28:08 +0200 Subject: [PATCH 32/54] Optimize DeepSeek V4 performance Squashes the perf/optimize-ds4 merge from 910b120 into a single linear commit. Co-authored-by: Prince Canuma --- mlx_lm/models/deepseek_v4.py | 566 +++++++++++++++++++++++++---------- 1 file changed, 408 insertions(+), 158 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 56cd0ed2c..67efa7831 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -86,6 +86,25 @@ def _score_func(scores: mx.array, func: str) -> mx.array: 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]: + 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 _limited_swiglu(gate: mx.array, up: mx.array, limit: float) -> mx.array: if limit and limit > 0: @@ -244,6 +263,7 @@ def _rope_full( return pe_out + def _apply_partial_rope( x: mx.array, rope: DeepseekV4RoPE, @@ -300,92 +320,75 @@ def _make_hc_split_sinkhorn_kernel(): source = """ uint idx = thread_position_in_grid.x; - constexpr int MIX = (2 + HC) * HC; - float epsv = static_cast(eps[0]); - - auto mix = mixes + idx * MIX; - auto pre_out = pre + idx * HC; - auto post_out = post + idx * HC; - auto comb_out = comb + idx * HC * HC; - - float pre_scale = static_cast(scale[0]); - float post_scale = static_cast(scale[1]); - float comb_scale = static_cast(scale[2]); - - for (int i = 0; i < HC; ++i) { - float z = static_cast(mix[i]) * pre_scale - + static_cast(base[i]); - pre_out[i] = 1.0f / (1.0f + metal::fast::exp(-z)) + epsv; - } - for (int i = 0; i < HC; ++i) { - int off = HC + i; - float z = static_cast(mix[off]) * post_scale - + static_cast(base[off]); - post_out[i] = 2.0f / (1.0f + metal::fast::exp(-z)); - } - - float c[HC * HC]; - for (int i = 0; i < HC; ++i) { - float row_max = -INFINITY; - for (int j = 0; j < HC; ++j) { - int cidx = i * HC + j; - int off = 2 * HC + cidx; - float v = static_cast(mix[off]) * comb_scale - + static_cast(base[off]); - c[cidx] = v; - row_max = metal::max(row_max, v); - } - float row_sum = 0.0f; - for (int j = 0; j < HC; ++j) { - int cidx = i * HC + j; - float v = metal::fast::exp(c[cidx] - row_max); - c[cidx] = v; - row_sum += v; - } - float inv_sum = 1.0f / row_sum; - for (int j = 0; j < HC; ++j) { - int cidx = i * HC + j; - c[cidx] = c[cidx] * inv_sum + epsv; - } + constexpr int MIX = (2 + HC) * HC; + constexpr int BASE = 2 * HC; + + const device float* mix = (const device float*)mixes + idx * MIX; + device float* pre_out = (device float*)pre + idx * HC; + device float* post_out = (device float*)post + idx * HC; + device float* comb_out = (device float*)comb + idx * HC * HC; + + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + const float epsv = eps[0]; + + // Pre-sigmoid + { + float4 z = *(const device float4*)mix * pre_scale + + *(const device float4*)base; + *(device float4*)pre_out = 1.0f / (1.0f + metal::fast::exp(-z)) + epsv; } - for (int j = 0; j < HC; ++j) { - float col_sum = 0.0f; - for (int i = 0; i < HC; ++i) { - col_sum += c[i * HC + j]; - } - float inv_denom = 1.0f / (col_sum + epsv); - for (int i = 0; i < HC; ++i) { - c[i * HC + j] *= inv_denom; - } + // Post-sigmoid + { + float4 z = *(const device float4*)(mix + HC) * post_scale + + *(const device float4*)(base + HC); + *(device float4*)post_out = 2.0f * 1.0f / (1.0f + metal::fast::exp(-z)); } + // Comb: four float4 loads — all independent, GPU issues in parallel + float4 v0 = *(const device float4*)(mix + BASE ) * comb_scale + *(const device float4*)(base + BASE ); + float4 v1 = *(const device float4*)(mix + BASE + 4) * comb_scale + *(const device float4*)(base + BASE + 4); + float4 v2 = *(const device float4*)(mix + BASE + 8) * comb_scale + *(const device float4*)(base + BASE + 8); + float4 v3 = *(const device float4*)(mix + BASE + 12) * comb_scale + *(const device float4*)(base + BASE + 12); + + // Per-row stable softmax: compute all maxes before any exp + float m0 = metal::max(metal::max(v0.x, v0.y), metal::max(v0.z, v0.w)); + float m1 = metal::max(metal::max(v1.x, v1.y), metal::max(v1.z, v1.w)); + float m2 = metal::max(metal::max(v2.x, v2.y), metal::max(v2.z, v2.w)); + float m3 = metal::max(metal::max(v3.x, v3.y), metal::max(v3.z, v3.w)); + + float4 e0 = metal::fast::exp(v0 - m0); + float4 e1 = metal::fast::exp(v1 - m1); + float4 e2 = metal::fast::exp(v2 - m2); + float4 e3 = metal::fast::exp(v3 - m3); + + // Explicit adds instead of dot(e, 1) — avoids unnecessary fmul + float4 r0 = e0 * 1.0f / (e0.x + e0.y + e0.z + e0.w) + epsv; + float4 r1 = e1 * 1.0f / (e1.x + e1.y + e1.z + e1.w) + epsv; + float4 r2 = e2 * 1.0f / (e2.x + e2.y + e2.z + e2.w) + epsv; + float4 r3 = e3 * 1.0f / (e3.x + e3.y + e3.z + e3.w) + epsv; + + // Initial column normalization + float4 col = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col; r1 *= col; r2 *= col; r3 *= col; + + // Sinkhorn iterations for (int iter = 1; iter < ITERS; ++iter) { - for (int i = 0; i < HC; ++i) { - float row_sum = 0.0f; - for (int j = 0; j < HC; ++j) { - row_sum += c[i * HC + j]; - } - float inv_denom = 1.0f / (row_sum + epsv); - for (int j = 0; j < HC; ++j) { - c[i * HC + j] *= inv_denom; - } - } - for (int j = 0; j < HC; ++j) { - float col_sum = 0.0f; - for (int i = 0; i < HC; ++i) { - col_sum += c[i * HC + j]; - } - float inv_denom = 1.0f / (col_sum + epsv); - for (int i = 0; i < HC; ++i) { - c[i * HC + j] *= inv_denom; - } - } + r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv); + r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv); + r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv); + r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv); + col = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col; r1 *= col; r2 *= col; r3 *= col; } - for (int i = 0; i < HC * HC; ++i) { - comb_out[i] = c[i]; - } + // Write comb output (four aligned 128-bit stores) + *(device float4*)(comb_out) = r0; + *(device float4*)(comb_out + 4) = r1; + *(device float4*)(comb_out + 8) = r2; + *(device float4*)(comb_out + 12) = r3; """ return mx.fast.metal_kernel( @@ -407,15 +410,16 @@ def hc_split_sinkhorn( sinkhorn_iters: int, eps: float, ) -> Tuple[mx.array, mx.array, mx.array]: - if _hc_split_sinkhorn_kernel is None: + if _hc_split_sinkhorn_kernel is None or hc_mult != 4: return _hc_split_sinkhorn_ops(mixes, scale, base, hc_mult, sinkhorn_iters, eps) if not isinstance(eps, mx.array): eps = mx.array([eps], dtype=mx.float32) + n_rows = mixes.size // ((2 + hc_mult) * hc_mult) return _hc_split_sinkhorn_kernel( inputs=[mixes, scale, base, eps], template=[("HC", hc_mult), ("ITERS", sinkhorn_iters)], - grid=(mixes.size // ((2 + hc_mult) * hc_mult), 1, 1), + grid=(n_rows, 1, 1), threadgroup=(256, 1, 1), output_shapes=[ (*mixes.shape[:-1], hc_mult), @@ -431,6 +435,159 @@ def _hc_collapse_op(pre: mx.array, x: mx.array) -> mx.array: return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.dtype) +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; + + 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 epsv = eps[0]; + + 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)) + epsv; + 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 + epsv)) + + epsv * 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) + ) + epsv); + 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 + epsv)) * 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) + ) + epsv); + 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, native bfloat4 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 bfloat16_t* x_row = (const device bfloat16_t*)x_in + + row * (HC * D); + device bfloat16_t* out_row = (device bfloat16_t*)collapsed + + row * D; + + // Native bfloat4 pointers: single 64-bit load per vector + using bf4 = vec; + const device bf4* x_row0 = (const device bf4*)(x_row + 0*D); + const device bf4* x_row1 = (const device bf4*)(x_row + 1*D); + const device bf4* x_row2 = (const device bf4*)(x_row + 2*D); + const device bf4* x_row3 = (const device bf4*)(x_row + 3*D); + device bf4* out4 = (device bf4*)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] = bf4(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] = (bfloat16_t)val; + } + #endif + """ + + return mx.fast.metal_kernel( + name="deepseek_v4_hc_sinkhorn_collapse", + input_names=["mixes", "scale", "base", "eps", "x_in"], + output_names=["post", "comb", "collapsed"], + source=source, + ) + + +_hc_sinkhorn_collapse_kernel = _make_hc_sinkhorn_collapse_kernel() + + @mx.compile def _hc_expand_op( post: mx.array, @@ -448,6 +605,13 @@ def _rms_rsqrt(flat: mx.array, eps: float) -> mx.array: return mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + eps) +@mx.compile +def _hc_mixes(flat: mx.array, fn_T: mx.array, norm_eps: float) -> mx.array: + """Fused RMS-rsqrt + matmul + scale into single compiled graph.""" + rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + norm_eps) + return (flat @ fn_T) * rsqrt + + class HyperConnection(nn.Module): def __init__(self, config: ModelArgs): super().__init__() @@ -460,12 +624,18 @@ def __init__(self, config: ModelArgs): 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) + self._fn_T = None def compute_weights(self, x: mx.array): B, L, H, D = x.shape flat = x.reshape(B, L, H * D).astype(mx.float32) - rsqrt = _rms_rsqrt(flat, self.norm_eps) - mixes = (flat @ self.fn.T) * rsqrt + if self._fn_T is None: + self._fn_T = self.fn.T + if self.training: + rsqrt = _rms_rsqrt(flat, self.norm_eps) + mixes = (flat @ self._fn_T) * rsqrt + else: + mixes = _hc_mixes(flat, self._fn_T, self.norm_eps) split_sinkhorn = _hc_split_sinkhorn_ops if self.training else hc_split_sinkhorn return split_sinkhorn( mixes, @@ -477,9 +647,42 @@ def compute_weights(self, x: mx.array): ) def collapse(self, x: mx.array): + if ( + not self.training + and _hc_sinkhorn_collapse_kernel is not None + and self.hc_mult == 4 + and x.dtype == mx.bfloat16 + ): + return self._fused_collapse(x) pre, post, comb = self.compute_weights(x) return _hc_collapse_op(pre, x), post, comb + def _fused_collapse(self, x: mx.array): + """Fused sinkhorn + collapse in a single Metal kernel dispatch.""" + B, L, H, D = x.shape + flat = x.reshape(B, L, H * D).astype(mx.float32) + if self._fn_T is None: + self._fn_T = self.fn.T + mixes = _hc_mixes(flat, self._fn_T, self.norm_eps) + + eps = self._hc_eps[0] + n_rows = B * L + x_flat = mx.contiguous(x.reshape(n_rows, H, D)) + + post, comb, collapsed = _hc_sinkhorn_collapse_kernel( + inputs=[mixes, self.scale, self.base, eps, x_flat], + template=[("HC", self.hc_mult), ("ITERS", self.sinkhorn_iters), ("D", D)], + grid=(n_rows * 256, 1, 1), + threadgroup=(256, 1, 1), + output_shapes=[ + (*mixes.shape[:-1], self.hc_mult), + (*mixes.shape[:-1], self.hc_mult, self.hc_mult), + (B, L, D), + ], + output_dtypes=[mx.float32, mx.float32, x.dtype], + ) + return collapsed, post, comb + def expand( self, block_out: mx.array, @@ -490,6 +693,24 @@ def expand( return _hc_expand_op(post, block_out, comb, residual) +@mx.compile +def _hyper_head_op( + x: mx.array, + fn: mx.array, + scale: mx.array, + base: mx.array, + norm_eps: float, + hc_eps: float, +) -> mx.array: + """Fused HyperHead: RMS-rsqrt + matmul + sigmoid + weighted sum.""" + B, L, H, D = x.shape + flat = x.reshape(B, L, H * D).astype(mx.float32) + rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + norm_eps) + mixes = (flat @ fn.T) * rsqrt + pre = mx.sigmoid(mixes * scale[0] + base) + hc_eps + return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.dtype) + + class HyperHead(nn.Module): def __init__(self, config: ModelArgs): super().__init__() @@ -503,6 +724,10 @@ def __init__(self, config: ModelArgs): self.scale = mx.ones((1,), dtype=mx.float32) def __call__(self, x: mx.array): + if not self.training: + return _hyper_head_op( + x, self.fn, self.scale, self.base, self.norm_eps, self.hc_eps + ) B, L, H, D = x.shape flat = x.reshape(B, L, H * D).astype(mx.float32) rsqrt = _rms_rsqrt(flat, self.norm_eps) @@ -522,6 +747,7 @@ def __init__(self, config: ModelArgs, layer_idx: int): 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)) + self._weight_T_f32 = None if self.hash: self.tid2eid = mx.zeros((config.vocab_size, self.top_k), dtype=mx.int32) else: @@ -531,23 +757,29 @@ def __init__(self, config: ModelArgs, layer_idx: int): def __call__(self, x: mx.array, input_ids: Optional[mx.array] = None): flat = x.reshape(-1, self.hidden_dim) - logits = flat.astype(mx.float32) @ self.weight.T.astype(mx.float32) - scores = _score_func(logits, self.scoring_func) + if self._weight_T_f32 is None: + self._weight_T_f32 = self.weight.T.astype(mx.float32) + logits = flat.astype(mx.float32) @ self._weight_T_f32 if self.hash: if input_ids is None: raise ValueError("DeepSeek-V4 hash routing requires input_ids.") + scores = _score_func(logits, self.scoring_func) inds = self.tid2eid[input_ids.reshape(-1)].astype(mx.int32) + weights = mx.take_along_axis(scores, inds, axis=-1) + if self.scoring_func != "softmax" and self.norm_topk_prob: + weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) + weights = weights * self.routed_scaling_factor else: - biased = scores + self.e_score_correction_bias - inds = mx.argpartition(-biased, kth=self.top_k - 1, axis=-1)[ - ..., : self.top_k - ] + inds, weights = _expert_select( + logits, + self.e_score_correction_bias, + self.top_k, + self.routed_scaling_factor, + self.norm_topk_prob, + self.scoring_func, + ) - weights = mx.take_along_axis(scores, inds, axis=-1) - if self.scoring_func != "softmax" and self.norm_topk_prob: - weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) - weights = weights * self.routed_scaling_factor route_shape = (*x.shape[:-1], self.top_k) inds = inds.reshape(route_shape) weights = weights.reshape(route_shape) @@ -1199,6 +1431,7 @@ def _merge_batch_state(values: List[Optional[mx.array]]): class Compressor(nn.Module): + def __init__(self, config: ModelArgs, compress_ratio: int, head_dim: int): super().__init__() self.compress_ratio = compress_ratio @@ -1213,10 +1446,12 @@ def __init__(self, config: ModelArgs, compress_ratio: int, head_dim: int): def _overlap_transform(self, x: mx.array, fill_value: float): B, W, R, _ = x.shape - out = mx.full((B, W, 2 * R, self.head_dim), fill_value, dtype=x.dtype) - out[:, :, R:] = x[:, :, :, self.head_dim :] - out[:, 1:, :R] = x[:, :-1, :, : self.head_dim] - return out + second_half = x[:, :, :, self.head_dim :] # (B, W, R, head_dim) + fill_row = mx.full((B, 1, R, self.head_dim), fill_value, dtype=x.dtype) + prev_first = mx.concatenate( + [fill_row, x[:, :-1, :, : self.head_dim]], axis=1 + ) # (B, W, R, head_dim) + return mx.concatenate([prev_first, second_half], axis=2) # (B, W, 2R, head_dim) def __call__( self, @@ -1315,6 +1550,7 @@ def __call__( return mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k] + class V4Attention(nn.Module): def __init__(self, config: ModelArgs, layer_idx: int): super().__init__() @@ -1349,6 +1585,7 @@ def __init__(self, config: ModelArgs, layer_idx: int): ) self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) self._q_l2_norm_weight = (mx.ones((self.head_dim,)),) + self._cached_dtype = None rope_theta = ( config.compress_rope_theta if self.compress_ratio else config.rope_theta @@ -1366,6 +1603,32 @@ def __init__(self, config: ModelArgs, layer_idx: int): if self.compress_ratio == 4: self.indexer = Indexer(config, self.compress_ratio) + def _ensure_cached(self, dtype): + if self._cached_dtype is not None and self._cached_dtype == dtype: + return + self._cached_dtype = dtype + self._attn_sink_cached = self.attn_sink.astype(dtype) + self._q_norm_weight_cached = self._q_l2_norm_weight[0].astype(dtype) + if isinstance(self.wo_a, nn.QuantizedLinear): + self._wo_a_weight = self.wo_a.weight.reshape( + self.o_groups, self.o_lora_rank, -1 + )[:, None] + self._wo_a_scales = self.wo_a.scales.reshape( + self.o_groups, self.o_lora_rank, -1 + )[:, None] + self._wo_a_biases = ( + None + if self.wo_a.biases is None + else self.wo_a.biases.reshape( + self.o_groups, self.o_lora_rank, -1 + )[:, None] + ) + else: + group_feat = (self.n_heads * self.head_dim) // self.o_groups + self._wo_a_weight_reshaped = self.wo_a.weight.reshape( + self.o_groups, self.o_lora_rank, group_feat + ) + def _grouped_output_projection(self, out: mx.array) -> mx.array: B, L = out.shape[:2] group_feat = (self.n_heads * self.head_dim) // self.o_groups @@ -1373,24 +1636,11 @@ def _grouped_output_projection(self, out: mx.array) -> mx.array: if isinstance(self.wo_a, nn.QuantizedLinear): out = out.transpose(2, 0, 1, 3) - weight = self.wo_a.weight.reshape(self.o_groups, self.o_lora_rank, -1)[ - :, None - ] - scales = self.wo_a.scales.reshape(self.o_groups, self.o_lora_rank, -1)[ - :, None - ] - biases = ( - None - if self.wo_a.biases is None - else self.wo_a.biases.reshape(self.o_groups, self.o_lora_rank, -1)[ - :, None - ] - ) out = mx.quantized_matmul( out, - weight, - scales=scales, - biases=biases, + self._wo_a_weight, + scales=self._wo_a_scales, + biases=self._wo_a_biases, transpose=True, group_size=self.wo_a.group_size, bits=self.wo_a.bits, @@ -1403,8 +1653,7 @@ def _grouped_output_projection(self, out: mx.array) -> mx.array: out = out + self.wo_a.bias return out - weight = self.wo_a.weight.reshape(self.o_groups, self.o_lora_rank, group_feat) - out = mx.einsum("bsgd,grd->bsgr", out, weight) + out = mx.einsum("bsgd,grd->bsgr", out, self._wo_a_weight_reshaped) out = out.reshape(B, L, self.o_groups * self.o_lora_rank) if "bias" in self.wo_a: out = out + self.wo_a.bias @@ -1426,8 +1675,9 @@ def __call__( offset = offset + 0 q_residual = self.q_norm(self.wq_a(x)) q = self.wq_b(q_residual).reshape(B, L, self.n_heads, self.head_dim) + self._ensure_cached(q.dtype) q = mx.fast.rms_norm( - q, self._q_l2_norm_weight[0].astype(q.dtype), self.config.rms_norm_eps + q, self._q_norm_weight_cached, self.config.rms_norm_eps ) q = q.transpose(0, 2, 1, 3) kv = self.kv_norm(self.wkv(x)).reshape(B, L, 1, self.head_dim) @@ -1446,50 +1696,51 @@ def __call__( if self.compress_ratio: v4_cache = cache if isinstance(cache, DeepseekV4Cache) else None pooled = self.compressor(x, self.compress_rope, v4_cache, offset) - lengths = ( - v4_cache.pooled_lengths("compressor_state") - if v4_cache is not None - else None - ) - use_indexer = hasattr(self, "indexer") and pooled.shape[1] > 0 - select_all = ( - use_indexer - and lengths is None - and pooled.shape[1] <= self.indexer.index_topk - ) - if select_all: - pooled = pooled[:, None] - pooled_bias = math.log(L) - elif use_indexer: - topk = self.indexer( - x, q_residual, self.compress_rope, self.rope, v4_cache, offset + if pooled.shape[1] > 0: + lengths = ( + v4_cache.pooled_lengths("compressor_state") + if v4_cache is not None + else None ) - if topk is not None: - if lengths is not None: - lengths = mx.array(lengths) - pooled_mask = (topk < lengths[:, None, None]).reshape( - B, 1, 1, -1 - ) - expanded = mx.broadcast_to( - pooled[:, None, None, :, :], - (B, 1, L, pooled.shape[1], self.head_dim), + use_indexer = hasattr(self, "indexer") + select_all = ( + use_indexer + and lengths is None + and pooled.shape[1] <= self.indexer.index_topk + ) + if select_all: + pooled = pooled[:, None] + pooled_bias = math.log(L) + elif use_indexer: + topk = self.indexer( + x, q_residual, self.compress_rope, self.rope, v4_cache, offset ) - idx = topk[:, None, :, :, None] - pooled = mx.take_along_axis( - expanded, - mx.broadcast_to(idx, idx.shape[:-1] + (self.head_dim,)), - axis=3, - ).reshape(B, 1, -1, self.head_dim) + if topk is not None: + if lengths is not None: + lengths = mx.array(lengths) + pooled_mask = (topk < lengths[:, None, None]).reshape( + B, 1, 1, -1 + ) + expanded = mx.broadcast_to( + pooled[:, None, None, :, :], + (B, 1, L, pooled.shape[1], self.head_dim), + ) + idx = topk[:, None, :, :, None] + pooled = mx.take_along_axis( + expanded, + mx.broadcast_to(idx, idx.shape[:-1] + (self.head_dim,)), + axis=3, + ).reshape(B, 1, -1, self.head_dim) + else: + pooled = pooled[:, None] else: + if lengths is not None: + lengths = mx.array(lengths) + pooled_mask = ( + mx.arange(pooled.shape[1]) < lengths[:, None] + ).reshape(B, 1, 1, -1) pooled = pooled[:, None] - else: - if lengths is not None: - lengths = mx.array(lengths) - pooled_mask = ( - mx.arange(pooled.shape[1]) < lengths[:, None] - ).reshape(B, 1, 1, -1) - pooled = pooled[:, None] - full_kv = mx.concatenate([full_kv, pooled], axis=2) + full_kv = mx.concatenate([full_kv, pooled], axis=2) if mask is not None and mask.shape[-1] > local_kv_len: mask = mask[..., -local_kv_len:] @@ -1521,7 +1772,6 @@ def __call__( mx.full(pad_shape, mx.finfo(mask.dtype).min, dtype=mask.dtype), ) mask = mx.concatenate([mask, pad], axis=-1) - out = scaled_dot_product_attention( q, full_kv, @@ -1529,7 +1779,7 @@ def __call__( cache=local_cache, scale=self.scale, mask=mask, - sinks=self.attn_sink.astype(q.dtype), + sinks=self._attn_sink_cached, ) out = _apply_partial_rope(out, self.rope, offset, inverse=True) out = out.transpose(0, 2, 1, 3).reshape(B, L, self.n_heads * self.head_dim) From 2ba4215de79018951a2565c926ab59a02eae0735 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sun, 26 Apr 2026 11:19:48 +0200 Subject: [PATCH 33/54] Refactor dtype handling in _ensure_cached method and improve pooled tensor processing logic (fix cache blowing up) - Changed dtype handling to use string representation for consistency. - Updated pooled tensor processing to include a mask when lengths are provided, improving handling of variable-length inputs. --- mlx_lm/models/deepseek_v4.py | 19 +++++++++++++++---- tests/test_models.py | 13 +++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 67efa7831..a39e62dee 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -1604,9 +1604,10 @@ def __init__(self, config: ModelArgs, layer_idx: int): self.indexer = Indexer(config, self.compress_ratio) def _ensure_cached(self, dtype): - if self._cached_dtype is not None and self._cached_dtype == dtype: + dtype_key = str(dtype) + if self._cached_dtype is not None and self._cached_dtype == dtype_key: return - self._cached_dtype = dtype + self._cached_dtype = dtype_key self._attn_sink_cached = self.attn_sink.astype(dtype) self._q_norm_weight_cached = self._q_l2_norm_weight[0].astype(dtype) if isinstance(self.wo_a, nn.QuantizedLinear): @@ -1705,12 +1706,22 @@ def __call__( use_indexer = hasattr(self, "indexer") select_all = ( use_indexer - and lengths is None - and pooled.shape[1] <= self.indexer.index_topk + and ( + L > 1 + or ( + lengths is None + and pooled.shape[1] <= self.indexer.index_topk + ) + ) ) if select_all: pooled = pooled[:, None] pooled_bias = math.log(L) + if lengths is not None: + lengths = mx.array(lengths) + pooled_mask = ( + mx.arange(pooled.shape[2]) < lengths[:, None] + ).reshape(B, 1, 1, -1) elif use_indexer: topk = self.indexer( x, q_residual, self.compress_rope, self.rope, v4_cache, offset diff --git a/tests/test_models.py b/tests/test_models.py index 80d17d78a..ad6f7296a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1499,6 +1499,19 @@ def test_deepseek_v4(self): # 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 + + class RaisingIndexer: + index_topk = 1 + + def __call__(self, *args, **kwargs): + raise AssertionError("Indexer should not run during prefill") + + compress_layer.indexer = RaisingIndexer() + 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]) + freq_dims = mx.arange(0, args.qk_rope_head_dim, 2, dtype=mx.float32) freq_dims = freq_dims / args.qk_rope_head_dim main_expected = 1.0 / (args.rope_theta**freq_dims) From 16f7205ff6131d035f78871ebd6ce7acb296fe9b Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sun, 26 Apr 2026 11:20:09 +0200 Subject: [PATCH 34/54] format --- mlx_lm/models/deepseek_v4.py | 29 +++++++++++------------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index a39e62dee..ec2d60f97 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -263,7 +263,6 @@ def _rope_full( return pe_out - def _apply_partial_rope( x: mx.array, rope: DeepseekV4RoPE, @@ -1550,7 +1549,6 @@ def __call__( return mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k] - class V4Attention(nn.Module): def __init__(self, config: ModelArgs, layer_idx: int): super().__init__() @@ -1620,9 +1618,9 @@ def _ensure_cached(self, dtype): self._wo_a_biases = ( None if self.wo_a.biases is None - else self.wo_a.biases.reshape( - self.o_groups, self.o_lora_rank, -1 - )[:, None] + else self.wo_a.biases.reshape(self.o_groups, self.o_lora_rank, -1)[ + :, None + ] ) else: group_feat = (self.n_heads * self.head_dim) // self.o_groups @@ -1677,9 +1675,7 @@ def __call__( q_residual = self.q_norm(self.wq_a(x)) q = self.wq_b(q_residual).reshape(B, L, self.n_heads, self.head_dim) self._ensure_cached(q.dtype) - q = mx.fast.rms_norm( - q, self._q_norm_weight_cached, self.config.rms_norm_eps - ) + q = mx.fast.rms_norm(q, self._q_norm_weight_cached, self.config.rms_norm_eps) q = q.transpose(0, 2, 1, 3) kv = self.kv_norm(self.wkv(x)).reshape(B, L, 1, self.head_dim) kv = kv.transpose(0, 2, 1, 3) @@ -1704,15 +1700,9 @@ def __call__( else None ) use_indexer = hasattr(self, "indexer") - select_all = ( - use_indexer - and ( - L > 1 - or ( - lengths is None - and pooled.shape[1] <= self.indexer.index_topk - ) - ) + select_all = use_indexer and ( + L > 1 + or (lengths is None and pooled.shape[1] <= self.indexer.index_topk) ) if select_all: pooled = pooled[:, None] @@ -1759,7 +1749,10 @@ def __call__( if mask is not None and full_kv.shape[2] > mask.shape[-1]: pad_shape = mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],) pad_pooled_mask = pooled_mask - if pad_pooled_mask is not None and pad_pooled_mask.shape[-1] != pad_shape[-1]: + if ( + pad_pooled_mask is not None + and pad_pooled_mask.shape[-1] != pad_shape[-1] + ): pad_pooled_mask = pad_pooled_mask[..., -pad_shape[-1] :] if pooled_bias is not None: dtype = q.dtype From 193aa77fe0aee4fe492c70037f4671522df5001f Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sun, 26 Apr 2026 14:32:15 +0200 Subject: [PATCH 35/54] Fix matrix multiplication in _hc_expand_op to use original comb tensor instead of swapped axes --- mlx_lm/models/deepseek_v4.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index ec2d60f97..be25dd722 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -595,7 +595,7 @@ def _hc_expand_op( residual: mx.array, ) -> mx.array: y = post[..., None] * block_out[:, :, None, :].astype(mx.float32) - y = y + mx.matmul(comb.swapaxes(-1, -2), residual.astype(mx.float32)) + y = y + mx.matmul(comb, residual.astype(mx.float32)) return y.astype(block_out.dtype) From 5a4aaa41db091c21ef9d04f135d6fa8d359261a8 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sun, 26 Apr 2026 22:44:24 +0200 Subject: [PATCH 36/54] Refactor output projection --- mlx_lm/models/deepseek_v4.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index be25dd722..e2c7d0397 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -1629,12 +1629,13 @@ def _ensure_cached(self, dtype): ) def _grouped_output_projection(self, out: mx.array) -> mx.array: - B, L = out.shape[:2] - group_feat = (self.n_heads * self.head_dim) // self.o_groups - out = out.reshape(B, L, self.o_groups, group_feat) + B, _, L, _ = out.shape + heads_per_group = self.n_heads // self.o_groups + out = out.reshape(B, self.o_groups, heads_per_group, L, self.head_dim) + out = out.transpose(1, 0, 3, 2, 4) + out = out.reshape(self.o_groups, B, L, heads_per_group * self.head_dim) if isinstance(self.wo_a, nn.QuantizedLinear): - out = out.transpose(2, 0, 1, 3) out = mx.quantized_matmul( out, self._wo_a_weight, @@ -1652,7 +1653,7 @@ def _grouped_output_projection(self, out: mx.array) -> mx.array: out = out + self.wo_a.bias return out - out = mx.einsum("bsgd,grd->bsgr", out, self._wo_a_weight_reshaped) + out = mx.einsum("gbsd,grd->bsgr", out, self._wo_a_weight_reshaped) out = out.reshape(B, L, self.o_groups * self.o_lora_rank) if "bias" in self.wo_a: out = out + self.wo_a.bias @@ -1786,7 +1787,6 @@ def __call__( sinks=self._attn_sink_cached, ) out = _apply_partial_rope(out, self.rope, offset, inverse=True) - out = out.transpose(0, 2, 1, 3).reshape(B, L, self.n_heads * self.head_dim) out = self._grouped_output_projection(out) return self.wo_b(out) From 2591b51f46056439fdea0eab6c0cf24e3dcc6b23 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sun, 26 Apr 2026 22:56:16 +0200 Subject: [PATCH 37/54] Refactor kv tensor reshaping in V4Attention for improved dimensionality handling --- mlx_lm/models/deepseek_v4.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index e2c7d0397..ba4131e6b 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -1678,8 +1678,7 @@ def __call__( self._ensure_cached(q.dtype) q = mx.fast.rms_norm(q, self._q_norm_weight_cached, self.config.rms_norm_eps) q = q.transpose(0, 2, 1, 3) - kv = self.kv_norm(self.wkv(x)).reshape(B, L, 1, self.head_dim) - kv = kv.transpose(0, 2, 1, 3) + kv = self.kv_norm(self.wkv(x)).reshape(B, 1, L, self.head_dim) q = _apply_partial_rope(q, self.rope, offset) kv = _apply_partial_rope(kv, self.rope, offset) From 22da01a1d4c70bd5dffa3fe4772ffab5195889e8 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sun, 26 Apr 2026 23:00:01 +0200 Subject: [PATCH 38/54] Remove redundant cache type check in V4Attention to streamline processing logic --- mlx_lm/models/deepseek_v4.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index ba4131e6b..804d36dc0 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -1667,9 +1667,6 @@ def __call__( ) -> mx.array: B, L, _ = x.shape local_cache = cache - if isinstance(cache, DeepseekV4Cache): - local_cache = cache - offset = local_cache.offset if local_cache is not None else 0 if isinstance(offset, mx.array): offset = offset + 0 From 8e8571a4680f14668a150ad01cbf89c71212171a Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Mon, 27 Apr 2026 01:08:35 +0200 Subject: [PATCH 39/54] Fix DeepSeek V4 sparse pooled prefill memory Keep pooled top-k attention grouped per query during prefill instead of flattening it into an L * top_k dense KV sequence. This avoids oversized SDPA score buffers on long prompts while preserving the dense path for small pooled caches and decode. Co-authored-by: adurham --- mlx_lm/models/deepseek_v4.py | 192 ++++++++++++++++++++++++----------- tests/test_models.py | 44 ++++++++ 2 files changed, 179 insertions(+), 57 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 804d36dc0..9a88ca79c 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -284,6 +284,55 @@ def _apply_partial_rope( return _rope_full(x, rope.inv_freq, pos, inverse, nope_dim) +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 _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) + pooled_scores = (q_scaled[:, :, :, None] * pooled).sum(axis=-1) + pooled_scores = _apply_score_mask(pooled_scores, pooled_mask) + + scores = mx.concatenate([local_scores, pooled_scores], axis=-1) + sink_offset = 0 + if sinks is not None: + sink_offset = 1 + sink_scores = mx.broadcast_to(sinks.reshape(1, H, 1, 1), (B, H, L, 1)) + scores = mx.concatenate([sink_scores.astype(scores.dtype), scores], axis=-1) + + weights = mx.softmax(scores, axis=-1, precise=True) + local_len = local_kv.shape[2] + local_weights = weights[..., sink_offset : sink_offset + local_len] + pooled_weights = weights[..., sink_offset + local_len :] + + out = local_weights @ local_kv + out = out + (pooled_weights[..., None] * pooled).sum(axis=-2) + return out.astype(q.dtype) + + @mx.compile def _hc_split_sinkhorn_ops( mixes: mx.array, @@ -1687,6 +1736,9 @@ def __call__( pooled_mask = None pooled_bias = None + sparse_pooled = None + sparse_topk = None + sparse_pooled_mask = None if self.compress_ratio: v4_cache = cache if isinstance(cache, DeepseekV4Cache) else None pooled = self.compressor(x, self.compress_rope, v4_cache, offset) @@ -1696,10 +1748,10 @@ def __call__( if v4_cache is not None else None ) - use_indexer = hasattr(self, "indexer") + use_indexer = isinstance(getattr(self, "indexer", None), Indexer) + max_pooled_length = pooled.shape[1] if lengths is None else max(lengths) select_all = use_indexer and ( - L > 1 - or (lengths is None and pooled.shape[1] <= self.indexer.index_topk) + max_pooled_length <= self.indexer.index_topk ) if select_all: pooled = pooled[:, None] @@ -1714,21 +1766,32 @@ def __call__( x, q_residual, self.compress_rope, self.rope, v4_cache, offset ) if topk is not None: - if lengths is not None: - lengths = mx.array(lengths) - pooled_mask = (topk < lengths[:, None, None]).reshape( - B, 1, 1, -1 + if L > 1: + sparse_pooled = pooled + sparse_topk = topk + if lengths is not None: + lengths = mx.array(lengths) + sparse_pooled_mask = ( + topk < lengths[:, None, None] + ).reshape(B, 1, L, -1) + else: + if lengths is not None: + lengths = mx.array(lengths) + pooled_mask = (topk < lengths[:, None, None]).reshape( + B, 1, 1, -1 + ) + expanded = mx.broadcast_to( + pooled[:, None, None, :, :], + (B, 1, L, pooled.shape[1], self.head_dim), ) - expanded = mx.broadcast_to( - pooled[:, None, None, :, :], - (B, 1, L, pooled.shape[1], self.head_dim), - ) - idx = topk[:, None, :, :, None] - pooled = mx.take_along_axis( - expanded, - mx.broadcast_to(idx, idx.shape[:-1] + (self.head_dim,)), - axis=3, - ).reshape(B, 1, -1, self.head_dim) + idx = topk[:, None, :, :, None] + pooled = mx.take_along_axis( + expanded, + mx.broadcast_to( + idx, idx.shape[:-1] + (self.head_dim,) + ), + axis=3, + ).reshape(B, 1, -1, self.head_dim) else: pooled = pooled[:, None] else: @@ -1738,50 +1801,65 @@ def __call__( mx.arange(pooled.shape[1]) < lengths[:, None] ).reshape(B, 1, 1, -1) pooled = pooled[:, None] - full_kv = mx.concatenate([full_kv, pooled], axis=2) + if sparse_topk is None: + full_kv = mx.concatenate([full_kv, pooled], axis=2) if mask is not None and mask.shape[-1] > local_kv_len: mask = mask[..., -local_kv_len:] - if mask is not None and full_kv.shape[2] > mask.shape[-1]: - pad_shape = mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],) - pad_pooled_mask = pooled_mask - if ( - pad_pooled_mask is not None - and pad_pooled_mask.shape[-1] != pad_shape[-1] - ): - pad_pooled_mask = pad_pooled_mask[..., -pad_shape[-1] :] - if pooled_bias is not None: - dtype = q.dtype - if mask.dtype == mx.bool_: - mask = mx.where( - mask, - mx.array(0, dtype=dtype), - mx.full((), mx.finfo(dtype).min, dtype=dtype), - ) - pad = mx.full(pad_shape, pooled_bias, dtype=mask.dtype) - elif mask.dtype == mx.bool_: - pad = mx.ones(pad_shape, dtype=mask.dtype) - if pad_pooled_mask is not None: - pad = pad & pad_pooled_mask - else: - pad = mx.zeros(pad_shape, dtype=mask.dtype) - if pad_pooled_mask is not None: - pad = mx.where( - pad_pooled_mask, - pad, - mx.full(pad_shape, mx.finfo(mask.dtype).min, dtype=mask.dtype), - ) - mask = mx.concatenate([mask, pad], axis=-1) - out = scaled_dot_product_attention( - q, - full_kv, - full_kv, - cache=local_cache, - scale=self.scale, - mask=mask, - sinks=self._attn_sink_cached, - ) + if sparse_topk is not None: + out = _sparse_pooled_attention( + q, + full_kv, + sparse_pooled, + sparse_topk, + mask, + sparse_pooled_mask, + self.scale, + self._attn_sink_cached, + ) + else: + if mask is not None and full_kv.shape[2] > mask.shape[-1]: + pad_shape = mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],) + pad_pooled_mask = pooled_mask + if ( + pad_pooled_mask is not None + and pad_pooled_mask.shape[-1] != pad_shape[-1] + ): + pad_pooled_mask = pad_pooled_mask[..., -pad_shape[-1] :] + if pooled_bias is not None: + dtype = q.dtype + if mask.dtype == mx.bool_: + mask = mx.where( + mask, + mx.array(0, dtype=dtype), + mx.full((), mx.finfo(dtype).min, dtype=dtype), + ) + pad = mx.full(pad_shape, pooled_bias, dtype=mask.dtype) + elif mask.dtype == mx.bool_: + pad = mx.ones(pad_shape, dtype=mask.dtype) + if pad_pooled_mask is not None: + pad = pad & pad_pooled_mask + else: + pad = mx.zeros(pad_shape, dtype=mask.dtype) + if pad_pooled_mask is not None: + pad = mx.where( + pad_pooled_mask, + pad, + mx.full( + pad_shape, mx.finfo(mask.dtype).min, dtype=mask.dtype + ), + ) + mask = mx.concatenate([mask, pad], axis=-1) + out = scaled_dot_product_attention( + q, + full_kv, + full_kv, + cache=local_cache, + scale=self.scale, + mask=mask, + sinks=self._attn_sink_cached, + ) out = _apply_partial_rope(out, self.rope, offset, inverse=True) out = self._grouped_output_projection(out) return self.wo_b(out) diff --git a/tests/test_models.py b/tests/test_models.py index ad6f7296a..c6c4f8271 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1464,6 +1464,50 @@ def test_deepseek_v4(self): for x, y in zip(expected, actual): self.assertTrue(mx.allclose(x, y, rtol=1e-5, atol=1e-5)) + # Sparse pooled attention keeps per-query top-k pooled KV grouped + # instead of flattening it into a dense L * top_k sequence. + B, H, L, D = 1, 2, 3, 4 + local_len, pooled_len, top_k = 5, 7, 2 + q = mx.random.normal((B, H, L, D)) + local_kv = mx.random.normal((B, 1, local_len, D)) + pooled = mx.random.normal((B, pooled_len, D)) + topk = mx.array([[[1, 3], [0, 4], [2, 5]]], dtype=mx.int32) + sinks = mx.random.normal((H,)) + idx = topk[:, None, :, :, None] + selected = mx.take_along_axis( + mx.broadcast_to(pooled[:, None, None], (B, 1, L, pooled_len, D)), + mx.broadcast_to(idx, idx.shape[:-1] + (D,)), + axis=3, + ) + dense_kv = mx.concatenate( + [local_kv, selected.reshape(B, 1, L * top_k, D)], axis=2 + ) + dense_mask = mx.zeros((B, H, L, local_len + L * top_k), dtype=mx.bool_) + dense_mask[..., :local_len] = True + for i in range(L): + start = local_len + i * top_k + dense_mask[:, :, i, start : start + top_k] = True + expected = deepseek_v4.scaled_dot_product_attention( + q, + dense_kv, + dense_kv, + cache=None, + scale=0.5, + mask=dense_mask, + sinks=sinks, + ) + actual = deepseek_v4._sparse_pooled_attention( + q, + local_kv, + pooled, + topk, + None, + None, + 0.5, + sinks, + ) + self.assertTrue(mx.allclose(actual, expected, rtol=1e-5, atol=1e-5)) + # Model test args = deepseek_v4.ModelArgs( model_type="deepseek_v4", From ddeffe33cc461957cc5c8b0b6a3ffca8f38379ee Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Mon, 27 Apr 2026 01:15:23 +0200 Subject: [PATCH 40/54] Remove sparse pooled attention test block --- tests/test_models.py | 44 -------------------------------------------- 1 file changed, 44 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index c6c4f8271..ad6f7296a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1464,50 +1464,6 @@ def test_deepseek_v4(self): for x, y in zip(expected, actual): self.assertTrue(mx.allclose(x, y, rtol=1e-5, atol=1e-5)) - # Sparse pooled attention keeps per-query top-k pooled KV grouped - # instead of flattening it into a dense L * top_k sequence. - B, H, L, D = 1, 2, 3, 4 - local_len, pooled_len, top_k = 5, 7, 2 - q = mx.random.normal((B, H, L, D)) - local_kv = mx.random.normal((B, 1, local_len, D)) - pooled = mx.random.normal((B, pooled_len, D)) - topk = mx.array([[[1, 3], [0, 4], [2, 5]]], dtype=mx.int32) - sinks = mx.random.normal((H,)) - idx = topk[:, None, :, :, None] - selected = mx.take_along_axis( - mx.broadcast_to(pooled[:, None, None], (B, 1, L, pooled_len, D)), - mx.broadcast_to(idx, idx.shape[:-1] + (D,)), - axis=3, - ) - dense_kv = mx.concatenate( - [local_kv, selected.reshape(B, 1, L * top_k, D)], axis=2 - ) - dense_mask = mx.zeros((B, H, L, local_len + L * top_k), dtype=mx.bool_) - dense_mask[..., :local_len] = True - for i in range(L): - start = local_len + i * top_k - dense_mask[:, :, i, start : start + top_k] = True - expected = deepseek_v4.scaled_dot_product_attention( - q, - dense_kv, - dense_kv, - cache=None, - scale=0.5, - mask=dense_mask, - sinks=sinks, - ) - actual = deepseek_v4._sparse_pooled_attention( - q, - local_kv, - pooled, - topk, - None, - None, - 0.5, - sinks, - ) - self.assertTrue(mx.allclose(actual, expected, rtol=1e-5, atol=1e-5)) - # Model test args = deepseek_v4.ModelArgs( model_type="deepseek_v4", From acf650c9df11b5ed3c2d9df5537627f73f6222c8 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Mon, 27 Apr 2026 08:50:02 +0200 Subject: [PATCH 41/54] Fix DeepSeek V4 HyperConnection expand orientation Restore the transpose in HyperConnection expand so the Sinkhorn combination matrix is applied in the same orientation as the original implementation. This fixes token dropping and repetition regressions seen in Japanese generation. Reported-by: https://github.com/Shinka-Man --- mlx_lm/models/deepseek_v4.py | 2 +- tests/test_models.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 9a88ca79c..acea50f6b 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -644,7 +644,7 @@ def _hc_expand_op( residual: mx.array, ) -> mx.array: y = post[..., None] * block_out[:, :, None, :].astype(mx.float32) - y = y + mx.matmul(comb, residual.astype(mx.float32)) + y = y + mx.matmul(comb.swapaxes(-1, -2), residual.astype(mx.float32)) return y.astype(block_out.dtype) diff --git a/tests/test_models.py b/tests/test_models.py index ad6f7296a..b890aea1c 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1464,6 +1464,19 @@ def test_deepseek_v4(self): for x, y in zip(expected, actual): self.assertTrue(mx.allclose(x, y, rtol=1e-5, atol=1e-5)) + # 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 = deepseek_v4._hc_expand_op(post, block_out, comb, residual) + self.assertTrue(mx.allclose(actual, expected, rtol=1e-5, atol=1e-5)) + # Model test args = deepseek_v4.ModelArgs( model_type="deepseek_v4", From 41ba61232dc29c99a1d34da3955ef0291fcdfa22 Mon Sep 17 00:00:00 2001 From: "clandestine.eth" <96172957+0xClandestine@users.noreply.github.com> Date: Mon, 27 Apr 2026 10:16:02 -0400 Subject: [PATCH 42/54] fix: replace broadcast multiply with matmul in sparse pooled attention The element-wise (q * pooled).sum() path broadcasts a (B,H,L,1,D) tensor against (B,1,L,topk,D), creating a (B,H,L,topk,D) intermediate. At 4k context with H=64, topk=512, D=512 this is ~137 GB per operation (x2). Replace with equivalent matmul: (B,L,H,D) @ (B,L,D,topk) which produces the (B,L,H,topk) result directly with ~0.25 GB peak memory. --- mlx_lm/models/deepseek_v4.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index acea50f6b..c41731bf7 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -313,7 +313,15 @@ def _sparse_pooled_attention( q_scaled = q * scale local_scores = q_scaled @ local_kv.swapaxes(-1, -2) local_scores = _apply_score_mask(local_scores, local_mask) - pooled_scores = (q_scaled[:, :, :, None] * pooled).sum(axis=-1) + + # Pooled scores via matmul instead of broadcast multiply + sum. + # The element-wise path creates a (B, H, L, topk, D) intermediate which + # at 4k context with H=64, topk=512, D=512 is ~137 GB. + # Matmul (B*L, H, D) @ (B*L, D, topk) → (B*L, H, topk) uses ~0.25 GB. + pooled_sq = pooled.squeeze(1) # (B, L, topk, D) + q_bl = q_scaled.transpose(0, 2, 1, 3) # (B, L, H, D) + pooled_scores = q_bl @ pooled_sq.swapaxes(-1, -2) # (B, L, H, topk) + pooled_scores = pooled_scores.transpose(0, 2, 1, 3) # (B, H, L, topk) pooled_scores = _apply_score_mask(pooled_scores, pooled_mask) scores = mx.concatenate([local_scores, pooled_scores], axis=-1) @@ -329,7 +337,9 @@ def _sparse_pooled_attention( pooled_weights = weights[..., sink_offset + local_len :] out = local_weights @ local_kv - out = out + (pooled_weights[..., None] * pooled).sum(axis=-2) + # Same matmul trick for weighted sum: (B*L, H, topk) @ (B*L, topk, D) + pw_bl = pooled_weights.transpose(0, 2, 1, 3) # (B, L, H, topk) + out = out + (pw_bl @ pooled_sq).transpose(0, 2, 1, 3) # (B, H, L, D) return out.astype(q.dtype) From f83460f15e0a2045eb696d93edf9e55e5d3189eb Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Tue, 28 Apr 2026 01:23:49 -0700 Subject: [PATCH 43/54] Enable loading the original checkpoint --- mlx_lm/models/deepseek_v4.py | 106 ++++++++++++++--------------------- mlx_lm/utils.py | 7 +++ 2 files changed, 48 insertions(+), 65 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index c41731bf7..ddfb13b8e 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -7,6 +7,7 @@ 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 BatchRotatingKVCache, RotatingKVCache @@ -76,6 +77,31 @@ def __post_init__(self): 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) @@ -1797,9 +1823,7 @@ def __call__( idx = topk[:, None, :, :, None] pooled = mx.take_along_axis( expanded, - mx.broadcast_to( - idx, idx.shape[:-1] + (self.head_dim,) - ), + mx.broadcast_to(idx, idx.shape[:-1] + (self.head_dim,)), axis=3, ).reshape(B, 1, -1, self.head_dim) else: @@ -2016,57 +2040,6 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: new_weights[k] = v weights = new_weights - def scale_to_float(scale: mx.array) -> mx.array: - if scale.dtype == mx.uint8: - return mx.exp((scale.astype(mx.float32) - 127.0) * math.log(2.0)) - return scale.astype(mx.float32) - - def dequant_fp8(weight: mx.array, scale: mx.array, block_size: int = 128): - weight = mx.from_fp8(weight, dtype=mx.bfloat16) - scale = scale_to_float(scale) - m, n = weight.shape - pad_m = (-m) % block_size - pad_n = (-n) % block_size - weight = mx.pad(weight, ((0, pad_m), (0, pad_n))) - weight = weight.reshape( - (m + pad_m) // block_size, - block_size, - (n + pad_n) // block_size, - block_size, - ) - weight = (weight * scale[:, None, :, None]).reshape(m + pad_m, n + pad_n) - return weight[:m, :n].astype(mx.bfloat16) - - def dequant_fp4(weight: mx.array, scale: mx.array, block_size: int = 32): - table = mx.array( - [ - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, - 0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, - ], - dtype=mx.float32, - ) - packed = weight.astype(mx.uint8) - low = packed & 0x0F - high = (packed >> 4) & 0x0F - unpacked = mx.stack([mx.take(table, low), mx.take(table, high)], axis=-1) - unpacked = unpacked.reshape(weight.shape[0], weight.shape[1] * 2) - scale = mx.repeat(scale_to_float(scale), block_size, axis=-1) - return (unpacked * scale).astype(mx.bfloat16) - new_weights = {} for k, v in weights.items(): if not k.endswith(".scale"): @@ -2085,9 +2058,11 @@ def dequant_fp4(weight: mx.array, scale: mx.array, block_size: int = 32): and weight.dtype in (mx.int8, mx.uint8) and v.shape[-1] * 16 == weight.shape[-1] ): - new_weights[wk] = dequant_fp4(weight, v) + new_weights[k + "s"] = v + new_weights[wk] = weight.view(mx.uint32) elif weight.dtype == mx.uint8: - new_weights[wk] = dequant_fp8(weight, v) + 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 @@ -2124,15 +2099,16 @@ def dequant_fp4(weight: mx.array, scale: mx.array, block_size: int = 32): ("w2", "down_proj"), ("w3", "up_proj"), ): - key0 = f"{prefix}.0.{src}.weight" - if key0 in weights: - stacked = [ - weights.pop(f"{prefix}.{e}.{src}.weight") - for e in range(self.args.n_routed_experts) - ] - weights[f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.weight"] = ( - mx.stack(stacked) - ) + 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) return weights diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 36114f05e..0f201a189 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -435,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): From 81a8c57689ec0d180f3fb599d27de72fe5745bc6 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Tue, 28 Apr 2026 02:28:07 -0700 Subject: [PATCH 44/54] Simplify GLU and gate remove intermediate castings --- mlx_lm/models/deepseek_v4.py | 85 +++++++++++++++--------------------- 1 file changed, 36 insertions(+), 49 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index ddfb13b8e..0e07a9846 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -108,7 +108,7 @@ def _score_func(scores: mx.array, func: str) -> mx.array: if func == "sigmoid": return mx.sigmoid(scores) if func == "sqrtsoftplus": - return mx.sqrt(mx.logaddexp(scores, 0)) + return mx.sqrt(nn.softplus(scores)) raise ValueError(f"Unsupported DeepSeek-V4 scoring function: {func}") @@ -121,6 +121,7 @@ def _expert_select( 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] @@ -131,6 +132,25 @@ def _expert_select( 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: @@ -148,38 +168,6 @@ def __call__(self, x, gate): return _limited_swiglu(gate, x, self.limit) -class DeepseekV4SwitchGLU(SwitchGLU): - sort_threshold = 8 - - def __call__(self, x, indices, scores) -> mx.array: - out_shape = x.shape - route_shape = indices.shape - x = mx.expand_dims(x, (-2, -3)) - - do_sort = indices.size >= self.sort_threshold - inv_order = None - if do_sort: - flat_indices = indices.flatten() - order = mx.argsort(flat_indices) - inv_order = mx.argsort(order) - x = x.flatten(0, -3)[order // route_shape[-1]] - indices = flat_indices[order] - scores = scores.flatten()[order] - if self.training: - indices = mx.stop_gradient(indices) - x_up = self.up_proj(x, indices, sorted_indices=do_sort) - x_gate = self.gate_proj(x, indices, sorted_indices=do_sort) - x = self.activation(x_up, x_gate) - x = x * scores.astype(x.dtype)[..., None, None] - x = self.down_proj(x, indices, sorted_indices=do_sort) - - if do_sort: - x = _scatter_unsort(x, inv_order, route_shape) - - x = x.squeeze(-2) - return x.sum(axis=-2).astype(x.dtype).reshape(out_shape) - - class DeepseekV4RoPE(nn.Module): def __init__( self, @@ -831,7 +819,6 @@ def __init__(self, config: ModelArgs, layer_idx: int): 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)) - self._weight_T_f32 = None if self.hash: self.tid2eid = mx.zeros((config.vocab_size, self.top_k), dtype=mx.int32) else: @@ -840,20 +827,19 @@ def __init__(self, config: ModelArgs, layer_idx: int): ) def __call__(self, x: mx.array, input_ids: Optional[mx.array] = None): - flat = x.reshape(-1, self.hidden_dim) - if self._weight_T_f32 is None: - self._weight_T_f32 = self.weight.T.astype(mx.float32) - logits = flat.astype(mx.float32) @ self._weight_T_f32 + logits = x @ self.weight.T if self.hash: if input_ids is None: raise ValueError("DeepSeek-V4 hash routing requires input_ids.") - scores = _score_func(logits, self.scoring_func) - inds = self.tid2eid[input_ids.reshape(-1)].astype(mx.int32) - weights = mx.take_along_axis(scores, inds, axis=-1) - if self.scoring_func != "softmax" and self.norm_topk_prob: - weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) - weights = weights * self.routed_scaling_factor + 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, @@ -864,9 +850,6 @@ def __call__(self, x: mx.array, input_ids: Optional[mx.array] = None): self.scoring_func, ) - route_shape = (*x.shape[:-1], self.top_k) - inds = inds.reshape(route_shape) - weights = weights.reshape(route_shape) return inds, weights @@ -896,7 +879,7 @@ def __init__(self, config: ModelArgs, layer_idx: int): super().__init__() self.config = config self.gate = MoEGate(config, layer_idx) - self.switch_mlp = DeepseekV4SwitchGLU( + self.switch_mlp = SwitchGLU( config.hidden_size, config.moe_intermediate_size, config.n_routed_experts, @@ -913,7 +896,8 @@ def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: x = sum_gradients(self.sharding_group)(x) inds, scores = self.gate(x, input_ids) - y = self.switch_mlp(x, inds, scores) + 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: @@ -2042,6 +2026,9 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: 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 From 4951496c6baaff87f216b2f9a28cd14a8477a377 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Tue, 28 Apr 2026 14:40:27 -0700 Subject: [PATCH 45/54] Fix RoPE to use the kernel by scaling freqs --- mlx_lm/models/deepseek_v4.py | 126 +++++++++++------------------------ tests/test_models.py | 45 +++++++------ 2 files changed, 61 insertions(+), 110 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 0e07a9846..47e9ebc4c 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -12,7 +12,7 @@ from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention from .cache import BatchRotatingKVCache, RotatingKVCache from .pipeline import PipelineMixin -from .switch_layers import SwitchGLU, _scatter_unsort +from .switch_layers import SwitchGLU @dataclass @@ -213,89 +213,42 @@ def correction_dim(num_rotations): elif rope_type not in (None, "default"): raise ValueError(f"Unsupported DeepSeek-V4 RoPE type: {rope_type}") - self._inv_freq = (inv_freq,) - - @property - def inv_freq(self): - return self._inv_freq[0] + self._freqs = 1.0 / inv_freq + self._freqs_cache = {} + + def _get_freqs(self, head_dim: int, inverse: bool, freq_scale: int): + key = (head_dim, inverse, freq_scale) + if key not in self._freqs_cache: + f = self._freqs + if freq_scale != 1: + f = f / 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, - positions: Optional[mx.array] = None, - ): - L = x.shape[-2] - if positions is None: - pos = mx.arange(L, dtype=mx.float32) - if isinstance(offset, mx.array): - offset = offset.astype(mx.float32) - if offset.ndim > 0: - offset = offset[..., None] - pos = pos + offset - else: - pos = positions.astype(mx.float32) - return _rope_full(x, self.inv_freq, pos, inverse, 0) - - -@mx.compile -def _rope_full( - x: mx.array, - inv_freq: mx.array, - pos: mx.array, - inverse: bool, - nope_dim: int, -) -> mx.array: - freqs = pos[..., None] * inv_freq - cos = mx.cos(freqs) - sin = mx.sin(freqs) - if inverse: - sin = -sin - - if nope_dim > 0: - pe = x[..., nope_dim:] - else: - pe = x - - if pos.ndim == 1: - broadcast_shape = (1,) * (pe.ndim - 2) + cos.shape - else: - broadcast_shape = ( - pos.shape[:-1] + (1,) * (pe.ndim - pos.ndim - 1) + cos.shape[-2:] + freq_scale: int = 1, + ) -> mx.array: + head_dim = x.shape[-1] + freqs = self._get_freqs(head_dim, inverse, freq_scale) + offset = offset // freq_scale if freq_scale != 1 else offset + return mx.fast.rope( + x, + head_dim, + traditional=True, + base=None, + scale=1.0, + offset=offset, + freqs=freqs, ) - cos = cos.reshape(broadcast_shape).astype(pe.dtype) - sin = sin.reshape(broadcast_shape).astype(pe.dtype) - - pe = pe.reshape(*pe.shape[:-1], pe.shape[-1] // 2, 2) - x0, x1 = pe[..., 0], pe[..., 1] - out = mx.stack([x0 * cos - x1 * sin, x0 * sin + x1 * cos], axis=-1) - pe_out = out.reshape(*out.shape[:-2], out.shape[-2] * 2) - - if nope_dim > 0: - return mx.concatenate([x[..., :nope_dim], pe_out], axis=-1) - return pe_out - - -def _apply_partial_rope( - x: mx.array, - rope: DeepseekV4RoPE, - offset: Any = 0, - inverse: bool = False, - positions: Optional[mx.array] = None, -) -> mx.array: - nope_dim = x.shape[-1] - rope.dims if x.shape[-1] != rope.dims else 0 - L = x.shape[-2] - if positions is None: - pos = mx.arange(L, dtype=mx.float32) - if isinstance(offset, mx.array): - offset = offset.astype(mx.float32) - if offset.ndim > 0: - offset = offset[..., None] - pos = pos + offset - else: - pos = positions.astype(mx.float32) - return _rope_full(x, rope.inv_freq, pos, inverse, nope_dim) def _apply_score_mask(scores: mx.array, mask: Optional[mx.array]) -> mx.array: @@ -1557,13 +1510,10 @@ def __call__( ) new_pooled = (kv * weights).sum(axis=2) new_pooled = self.norm(new_pooled.astype(x.dtype)) - positions = mx.arange(new_pooled.shape[1], dtype=mx.float32) - positions = positions * self.compress_ratio - if isinstance(pool_base, mx.array) and pool_base.ndim > 0: - pool_base = pool_base[:, None] - positions = positions + pool_base - new_pooled = _apply_partial_rope( - new_pooled[:, None], rope, positions=positions + new_pooled = rope( + new_pooled[:, None], + offset=pool_base, + freq_scale=self.compress_ratio, ).squeeze(1) if cache is not None: @@ -1601,7 +1551,7 @@ def __call__( offset = start_pos q = self.wq_b(q_residual).reshape(B, L, self.n_heads, self.head_dim) q = q.transpose(0, 2, 1, 3) - q = _apply_partial_rope(q, position_rope, offset) + q = position_rope(q, offset) scores = q.astype(mx.float32) @ pooled[:, None].swapaxes(-1, -2).astype( mx.float32 @@ -1746,8 +1696,8 @@ def __call__( q = q.transpose(0, 2, 1, 3) kv = self.kv_norm(self.wkv(x)).reshape(B, 1, L, self.head_dim) - q = _apply_partial_rope(q, self.rope, offset) - kv = _apply_partial_rope(kv, self.rope, offset) + q = self.rope(q, offset) + kv = self.rope(kv, offset) if local_cache is not None: kv, _ = local_cache.update_and_fetch(kv, kv) @@ -1878,7 +1828,7 @@ def __call__( mask=mask, sinks=self._attn_sink_cached, ) - out = _apply_partial_rope(out, self.rope, offset, inverse=True) + out = self.rope(out, offset, inverse=True) out = self._grouped_output_projection(out) return self.wo_b(out) diff --git a/tests/test_models.py b/tests/test_models.py index b890aea1c..ed321a76c 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) @@ -1440,14 +1445,21 @@ def test_deepseek_v4(self): expected = expected.reshape(*expected.shape[:-2], 4) self.assertTrue(mx.allclose(y, expected, rtol=1e-5, atol=1e-5)) - positions = mx.array([1, 5, 9], dtype=mx.float32) - y = rope(x, positions=positions) + # 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 + y_scaled = rope(x, offset=4, freq_scale=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, expected, rtol=1e-5, atol=1e-5)) + self.assertTrue(mx.allclose(y_scaled, expected, rtol=1e-5, atol=1e-5)) # HyperConnection Sinkhorn test for hc_mult in (2, 4): @@ -1525,14 +1537,6 @@ def __call__(self, *args, **kwargs): logits = model(inputs, cache=cache) mx.eval(logits, [c.state for c in cache]) - freq_dims = mx.arange(0, args.qk_rope_head_dim, 2, dtype=mx.float32) - freq_dims = freq_dims / args.qk_rope_head_dim - main_expected = 1.0 / (args.rope_theta**freq_dims) - compress_expected = 1.0 / (args.compress_rope_theta**freq_dims) - self.assertTrue(mx.allclose(main_layer.rope.inv_freq, main_expected)) - self.assertTrue(mx.allclose(compress_layer.rope.inv_freq, compress_expected)) - self.assertIs(compress_layer.compress_rope, compress_layer.rope) - # Sanitize test weight = mx.to_fp8(mx.ones((128, 128), dtype=mx.float32)) converted = model.sanitize( @@ -1541,16 +1545,13 @@ def __call__(self, *args, **kwargs): "layers.0.attn.wkv.scale": mx.full((1, 1), 127, dtype=mx.uint8), } ) - key = "model.layers.0.attn.wkv.weight" - self.assertIn(key, converted) - self.assertTrue( - mx.allclose( - converted[key].astype(mx.float32), - mx.ones((128, 128), dtype=mx.float32), - rtol=1e-5, - atol=1e-5, - ) - ) + 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 From 3cf528282f2db0811fb2b00f8d0982a840734021 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Tue, 28 Apr 2026 18:04:54 -0700 Subject: [PATCH 46/54] Start simplifying and speeding up the attention --- mlx_lm/models/deepseek_v4.py | 101 ++++++++++++++++++++++++++++++++++- 1 file changed, 100 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 47e9ebc4c..85cdc8659 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -11,6 +11,7 @@ from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention from .cache import BatchRotatingKVCache, RotatingKVCache +from .mla import MultiLinear from .pipeline import PipelineMixin from .switch_layers import SwitchGLU @@ -1568,6 +1569,94 @@ def __call__( 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, + ) + + 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 + + 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) + + return out + + +def v4_attention_factory(config: ModelArgs, layer_idx: int) -> nn.Module: + """Instantiate the appropriate attention module for a given layer.""" + if config.compress_ratios[layer_idx] == 0: + return LocalAttention(config, layer_idx) + return V4Attention(config, layer_idx) + + class V4Attention(nn.Module): def __init__(self, config: ModelArgs, layer_idx: int): super().__init__() @@ -1836,7 +1925,7 @@ def __call__( class DeepseekV4Block(nn.Module): def __init__(self, config: ModelArgs, layer_idx: int): super().__init__() - self.attn = V4Attention(config, layer_idx) + 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) @@ -2047,6 +2136,16 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.{suffix}" ] = mx.stack(stacked) + for layer_idx in range(n_layers): + if self.args.compress_ratios[layer_idx] != 0: + continue + 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): From ad9e3297c5f043294a55a072e64a68d57c291194 Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Wed, 29 Apr 2026 10:59:42 +0200 Subject: [PATCH 47/54] Attempt to fix IndexError on batching. --- mlx_lm/models/deepseek_v4.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 0e07a9846..267f0075f 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -1224,6 +1224,13 @@ def finalize(self): self._pending_lengths = None def filter(self, batch_indices): + if isinstance(batch_indices, mx.array): + idx_list = batch_indices.tolist() + else: + idx_list = list(batch_indices) + needed = max(idx_list) + 1 if idx_list else 0 + self._expand_state_to(needed) + if hasattr(self.local, "filter"): self.local.filter(batch_indices) elif self.local.keys is not None: @@ -1237,6 +1244,31 @@ def filter(self, batch_indices): for key in self._length_keys: state[key] = self._filter_lengths(state[key], batch_indices) + def _expand_state_to(self, batch_size: int): + """Ensure compressor/indexer state arrays/lists have at least `batch_size` + rows along dim 0. Slots without emitted positions are zero-padded.""" + for state in (self.compressor_state, self.indexer_state): + for key in self._state_keys: + value = state[key] + if value is None: + continue + cur = value.shape[0] + if cur < batch_size: + pad_shape = (batch_size - cur,) + value.shape[1:] + state[key] = mx.concatenate( + [value, mx.zeros(pad_shape, dtype=value.dtype)], axis=0 + ) + for key in self._length_keys: + lst = state[key] + if lst is None: + continue + if isinstance(lst, mx.array): + lst = lst.tolist() + if len(lst) < batch_size: + state[key] = lst + [0] * (batch_size - len(lst)) + else: + state[key] = lst + def extend(self, other): self_batch = self._cache_batch_size(self.local) other_batch = self._cache_batch_size(other.local) From 83d7e742ddc0be40d7849ef9bb5160a69948f499 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Wed, 29 Apr 2026 03:39:09 -0700 Subject: [PATCH 48/54] Refactor compressor and compressed non-sparse attention --- mlx_lm/models/deepseek_v4.py | 225 ++++++++++++++++++++++++++++------- tests/test_models.py | 3 +- 2 files changed, 185 insertions(+), 43 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 85cdc8659..499f9e581 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -2,6 +2,7 @@ import math from dataclasses import dataclass, field +from functools import partial from typing import Any, Dict, List, Optional, Tuple import mlx.core as mx @@ -176,9 +177,11 @@ def __init__( 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 @@ -217,12 +220,12 @@ def correction_dim(num_rotations): self._freqs = 1.0 / inv_freq self._freqs_cache = {} - def _get_freqs(self, head_dim: int, inverse: bool, freq_scale: int): - key = (head_dim, inverse, freq_scale) + def _get_freqs(self, head_dim: int, inverse: bool): + key = (head_dim, inverse) if key not in self._freqs_cache: f = self._freqs - if freq_scale != 1: - f = f / freq_scale + if self.freq_scale != 1: + f = f / self.freq_scale if inverse: f = -f nope_pairs = (head_dim - self.dims) // 2 @@ -236,11 +239,10 @@ def __call__( x: mx.array, offset: Any = 0, inverse: bool = False, - freq_scale: int = 1, ) -> mx.array: head_dim = x.shape[-1] - freqs = self._get_freqs(head_dim, inverse, freq_scale) - offset = offset // freq_scale if freq_scale != 1 else offset + 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, @@ -260,6 +262,33 @@ def _apply_score_mask(scores: mx.array, mask: Optional[mx.array]) -> mx.array: return scores + mask.astype(scores.dtype) +@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) + + def _sparse_pooled_attention( q: mx.array, local_kv: mx.array, @@ -1465,20 +1494,17 @@ def __init__(self, config: ModelArgs, compress_ratio: int, head_dim: int): 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) - - def _overlap_transform(self, x: mx.array, fill_value: float): - B, W, R, _ = x.shape - second_half = x[:, :, :, self.head_dim :] # (B, W, R, head_dim) - fill_row = mx.full((B, 1, R, self.head_dim), fill_value, dtype=x.dtype) - prev_first = mx.concatenate( - [fill_row, x[:, :-1, :, : self.head_dim]], axis=1 - ) # (B, W, R, head_dim) - return mx.concatenate([prev_first, second_half], axis=2) # (B, W, 2R, head_dim) + 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, - rope: DeepseekV4RoPE, cache: Optional[DeepseekV4Cache], start_pos: int, state_key: str = "compressor_state", @@ -1495,30 +1521,24 @@ def __call__( kv, gate, state_key, self.compress_ratio, start_pos ) - if ready_kv.shape[1] == 0: + if ready_kv.size == 0: new_pooled = mx.zeros((B, 0, self.head_dim), dtype=x.dtype) else: - W = ready_kv.shape[1] // self.compress_ratio - kv = ready_kv.reshape(B, W, self.compress_ratio, self.out_dim) - gate = ready_gate.reshape( - B, W, self.compress_ratio, self.out_dim - ) + self.ape.astype(ready_gate.dtype) - if self.overlap: - kv = self._overlap_transform(kv, 0.0) - gate = self._overlap_transform(gate, -float("inf")) - weights = mx.softmax(gate.astype(mx.float32), axis=2, precise=True).astype( - kv.dtype + compress_func = ( + _overlap_compress_kv if self.overlap else _simple_compress_kv ) - new_pooled = (kv * weights).sum(axis=2) - new_pooled = self.norm(new_pooled.astype(x.dtype)) - new_pooled = rope( + 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, - freq_scale=self.compress_ratio, ).squeeze(1) if cache is not None: - return cache.update_pool(new_pooled, state_key) + new_pooled = cache.update_pool(new_pooled, state_key) + return new_pooled @@ -1539,13 +1559,12 @@ def __call__( self, x: mx.array, q_residual: mx.array, - rope: DeepseekV4RoPE, position_rope: DeepseekV4RoPE, cache: Optional[DeepseekV4Cache], start_pos: int, ): B, L, _ = x.shape - pooled = self.compressor(x, rope, cache, start_pos, state_key="indexer_state") + pooled = self.compressor(x, cache, start_pos, state_key="indexer_state") if pooled.shape[1] == 0: return None @@ -1650,10 +1669,133 @@ def __call__( 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) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + v4_cache = cache if isinstance(cache, DeepseekV4Cache) else None + offset = cache.offset if cache is not None else 0 + + 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, kv) + local_kv_len = kv.shape[2] + + # Pool tokens into compressed KV and concatenate with local KV + pooled = self.compressor(x, v4_cache, offset) + pooled_mask = None + if pooled.shape[1] > 0: + lengths = ( + v4_cache.pooled_lengths("compressor_state") + if v4_cache is not None + else None + ) + if lengths is not None: + lengths = mx.array(lengths) + pooled_mask = (mx.arange(pooled.shape[1]) < lengths[:, None]).reshape( + B, 1, 1, -1 + ) + kv = mx.concatenate([kv, pooled[:, None]], axis=2) + + # Extend mask to cover pooled entries + if mask is not None and mask.shape[-1] > local_kv_len: + mask = mask[..., -local_kv_len:] + if mask is not None and kv.shape[2] > mask.shape[-1]: + pad_shape = mask.shape[:-1] + (kv.shape[2] - mask.shape[-1],) + if mask.dtype == mx.bool_: + pad = mx.ones(pad_shape, dtype=mask.dtype) + if pooled_mask is not None: + pad = pad & pooled_mask + else: + pad = mx.zeros(pad_shape, dtype=mask.dtype) + if pooled_mask is not None: + pad = mx.where( + pooled_mask, + pad, + mx.full(pad_shape, mx.finfo(mask.dtype).min, dtype=mask.dtype), + ) + mask = mx.concatenate([mask, pad], axis=-1) + + 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) + + return out + + def v4_attention_factory(config: ModelArgs, layer_idx: int) -> nn.Module: """Instantiate the appropriate attention module for a given layer.""" - if config.compress_ratios[layer_idx] == 0: + ratio = config.compress_ratios[layer_idx] + if ratio == 0: return LocalAttention(config, layer_idx) + if ratio == 128: + return CompressedAttention(config, layer_idx) return V4Attention(config, layer_idx) @@ -1703,7 +1845,6 @@ def __init__(self, config: ModelArgs, layer_idx: int): rope_scaling, config.max_position_embeddings, ) - self.compress_rope = self.rope if self.compress_ratio: self.compressor = Compressor(config, self.compress_ratio, self.head_dim) if self.compress_ratio == 4: @@ -1800,7 +1941,7 @@ def __call__( sparse_pooled_mask = None if self.compress_ratio: v4_cache = cache if isinstance(cache, DeepseekV4Cache) else None - pooled = self.compressor(x, self.compress_rope, v4_cache, offset) + pooled = self.compressor(x, v4_cache, offset) if pooled.shape[1] > 0: lengths = ( v4_cache.pooled_lengths("compressor_state") @@ -1821,9 +1962,7 @@ def __call__( mx.arange(pooled.shape[2]) < lengths[:, None] ).reshape(B, 1, 1, -1) elif use_indexer: - topk = self.indexer( - x, q_residual, self.compress_rope, self.rope, v4_cache, offset - ) + topk = self.indexer(x, q_residual, self.rope, v4_cache, offset) if topk is not None: if L > 1: sparse_pooled = pooled @@ -2136,8 +2275,10 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.{suffix}" ] = mx.stack(stacked) + # Reshape wo_a from nn.Linear (2D) to MultiLinear (3D) for + # LocalAttention (ratio=0) and CompressedAttention (ratio=128) layers for layer_idx in range(n_layers): - if self.args.compress_ratios[layer_idx] != 0: + if self.args.compress_ratios[layer_idx] == 4: continue prefix = f"model.layers.{layer_idx}.attn.wo_a" for key in (f"{prefix}.weight", f"{prefix}.scales", f"{prefix}.biases"): diff --git a/tests/test_models.py b/tests/test_models.py index ed321a76c..f494b4f31 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1452,7 +1452,8 @@ def test_deepseek_v4(self): # 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 - y_scaled = rope(x, offset=4, 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) From d8a47ae3f721ea5537931ad5f9368f0065558c39 Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Wed, 29 Apr 2026 16:13:47 +0200 Subject: [PATCH 49/54] Ensure dtype going into mx.fast.rope is int32 --- mlx_lm/models/deepseek_v4.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 499f9e581..53eba1c3a 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -1121,11 +1121,11 @@ def accumulate_windows( state["buffer_lengths"] = buffer_lengths state["_new_pooled_lengths"] = [usable // ratio for usable in usable_lengths] - prev_lengths = mx.array(buf_lengths, dtype=mx.float32) + prev_lengths = mx.array(buf_lengths, dtype=mx.int32) if isinstance(start_pos, mx.array): - pool_base = mx.maximum(start_pos, 0).astype(mx.float32) + pool_base = mx.maximum(start_pos, 0).astype(mx.int32) else: - pool_base = mx.full((B,), max(0, start_pos), dtype=mx.float32) + pool_base = mx.full((B,), max(0, start_pos), dtype=mx.int32) return ready_kv, ready_gate, pool_base - prev_lengths def update_pool(self, new_pooled: mx.array, state_key: str) -> mx.array: From 2120a1b1f7b25963944c61cdf67b47a45334c588 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Thu, 30 Apr 2026 16:56:48 -0700 Subject: [PATCH 50/54] Major cache refactor and attention simplification --- mlx_lm/generate.py | 4 + mlx_lm/models/cache.py | 533 ++++++++++++++++ mlx_lm/models/deepseek_v4.py | 1100 ++++++---------------------------- tests/test_models.py | 7 - 4 files changed, 714 insertions(+), 930 deletions(-) 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..82fdbfd37 100644 --- a/mlx_lm/models/cache.py +++ b/mlx_lm/models/cache.py @@ -900,6 +900,539 @@ 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: + new_buf_kv[i, :nr] = kv[i, vl - nr : vl] + new_buf_gate[i, :nr] = gate[i, vl - nr : 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 index f13705067..c2251f26a 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -3,7 +3,7 @@ import math from dataclasses import dataclass, field from functools import partial -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union import mlx.core as mx import mlx.nn as nn @@ -11,7 +11,7 @@ from mlx.utils import tree_flatten from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention -from .cache import BatchRotatingKVCache, RotatingKVCache +from .cache import CacheList, PoolingCache, RotatingKVCache from .mla import MultiLinear from .pipeline import PipelineMixin from .switch_layers import SwitchGLU @@ -262,6 +262,26 @@ def _apply_score_mask(scores: mx.array, mask: Optional[mx.array]) -> mx.array: 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) @@ -289,6 +309,15 @@ def _overlap_compress_kv(kv, gate, ape, head_dim): 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, @@ -310,33 +339,27 @@ def _sparse_pooled_attention( 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 scores via matmul instead of broadcast multiply + sum. - # The element-wise path creates a (B, H, L, topk, D) intermediate which - # at 4k context with H=64, topk=512, D=512 is ~137 GB. - # Matmul (B*L, H, D) @ (B*L, D, topk) → (B*L, H, topk) uses ~0.25 GB. - pooled_sq = pooled.squeeze(1) # (B, L, topk, D) - q_bl = q_scaled.transpose(0, 2, 1, 3) # (B, L, H, D) - pooled_scores = q_bl @ pooled_sq.swapaxes(-1, -2) # (B, L, H, topk) - pooled_scores = pooled_scores.transpose(0, 2, 1, 3) # (B, H, L, topk) + 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) + ) - scores = mx.concatenate([local_scores, pooled_scores], axis=-1) - sink_offset = 0 - if sinks is not None: - sink_offset = 1 - sink_scores = mx.broadcast_to(sinks.reshape(1, H, 1, 1), (B, H, L, 1)) - scores = mx.concatenate([sink_scores.astype(scores.dtype), scores], axis=-1) - - weights = mx.softmax(scores, axis=-1, precise=True) - local_len = local_kv.shape[2] - local_weights = weights[..., sink_offset : sink_offset + local_len] - pooled_weights = weights[..., sink_offset + local_len :] + 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 - # Same matmul trick for weighted sum: (B*L, H, topk) @ (B*L, topk, D) - pw_bl = pooled_weights.transpose(0, 2, 1, 3) # (B, L, H, topk) - out = out + (pw_bl @ pooled_sq).transpose(0, 2, 1, 3) # (B, H, L, D) + 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) @@ -888,631 +911,6 @@ def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: return y -class DeepseekV4Cache: - _state_keys = ("buffer_kv", "buffer_gate", "pooled") - _length_keys = ("buffer_lengths", "pooled_lengths") - - def __init__(self, sliding_window: int): - self.local = RotatingKVCache(max_size=sliding_window, keep=0) - self.compressor_state = self._new_branch_state() - self.indexer_state = self._new_branch_state() - self._pending_lengths = None - - @property - def offset(self): - return self.local.offset - - @property - def keys(self): - return self.local.keys - - @keys.setter - def keys(self, value): - self.local.keys = value - - @property - def state(self): - local_state = None if self.local.empty() else self.local.state - return ( - local_state, - self._branch_state_tuple(self.compressor_state), - self._branch_state_tuple(self.indexer_state), - ) - - @state.setter - def state(self, value): - local_state, compressor_state, indexer_state = value - if local_state is None: - self.local.keys = None - self.local.values = None - else: - self.local.state = local_state - self.compressor_state = self._state_from_tuple(compressor_state) - self.indexer_state = self._state_from_tuple(indexer_state) - - @property - def meta_state(self): - return self.local.meta_state - - @meta_state.setter - def meta_state(self, value): - self.local.meta_state = value - - def update_and_fetch(self, keys, values): - return self.local.update_and_fetch(keys, values) - - def make_mask(self, *args, **kwargs): - return self.local.make_mask(*args, **kwargs) - - def is_trimmable(self): - if not self.local.is_trimmable(): - return False - for state in (self.compressor_state, self.indexer_state): - pooled = state["pooled"] - if pooled is not None and pooled.shape[1] > 0: - return False - return True - - def trim(self, n): - trimmed = self.local.trim(n) - batch_size = self._cache_batch_size(self.local) - for state in (self.compressor_state, self.indexer_state): - for key in ("buffer_kv", "buffer_gate"): - value = state[key] - if value is not None: - state[key] = value[:, : max(value.shape[1] - trimmed, 0)] - lengths = state["buffer_lengths"] - if lengths is not None: - state["buffer_lengths"] = [ - max(length - trimmed, 0) - for length in self._lengths_list(lengths, batch_size, 0) - ] - return trimmed - - def size(self): - return self.local.size() - - def empty(self): - return self.local.empty() - - @property - def nbytes(self): - total = self.local.nbytes - for state in (self.compressor_state, self.indexer_state): - for value in state.values(): - if value is not None and hasattr(value, "nbytes"): - total += value.nbytes - return total - - def _branch_state(self, state_key: str): - return ( - self.indexer_state - if state_key == "indexer_state" - else self.compressor_state - ) - - @classmethod - def _new_branch_state(cls): - return {key: None for key in cls._state_keys + cls._length_keys} - - @classmethod - def _branch_state_tuple(cls, state): - return tuple(state[k] for k in cls._state_keys + cls._length_keys) - - @classmethod - def _state_from_tuple(cls, values): - keys = cls._state_keys + cls._length_keys - state = cls._new_branch_state() - state.update(zip(keys, values)) - return state - - @staticmethod - def _lengths_list(lengths, batch_size: int, default: Optional[int] = None): - if lengths is None: - if default is None: - return None - return [default] * batch_size - if isinstance(lengths, mx.array): - lengths = lengths.tolist() - return [int(length) for length in lengths] - - @staticmethod - def _filter_lengths(lengths, batch_indices): - if lengths is None: - return None - if isinstance(lengths, mx.array): - lengths = lengths.tolist() - if isinstance(batch_indices, mx.array): - batch_indices = batch_indices.tolist() - if len(lengths) == 1 and any(idx != 0 for idx in batch_indices): - lengths = lengths * (max(batch_indices) + 1) - return [int(lengths[idx]) for idx in batch_indices] - - def accumulate_windows( - self, - kv: mx.array, - gate: mx.array, - state_key: str, - ratio: int, - start_pos: int, - ): - state = self._branch_state(state_key) - buf_kv, buf_gate = state["buffer_kv"], state["buffer_gate"] - B, L = kv.shape[:2] - buf_lengths = self._lengths_list(state["buffer_lengths"], B) - chunk_lengths = self._pending_lengths - if buf_lengths is None and chunk_lengths is None: - if buf_kv is not None and buf_kv.shape[1]: - kv = mx.concatenate([buf_kv, kv], axis=1) - gate = mx.concatenate([buf_gate, gate], axis=1) - usable = (kv.shape[1] // ratio) * ratio - state["buffer_kv"] = kv[:, usable:] - state["buffer_gate"] = gate[:, usable:] - state["buffer_lengths"] = None - state["_new_pooled_lengths"] = None - if isinstance(start_pos, mx.array): - pool_base = mx.maximum(start_pos, 0) - else: - pool_base = max(0, start_pos) - pool_base = pool_base - (buf_kv.shape[1] if buf_kv is not None else 0) - return kv[:, :usable], gate[:, :usable], pool_base - - buf_lengths = self._lengths_list( - state["buffer_lengths"], - B, - 0 if buf_kv is None else buf_kv.shape[1], - ) - chunk_lengths = self._lengths_list(chunk_lengths, B, L) - total_lengths = [ - buf_length + min(chunk_length, L) - for buf_length, chunk_length in zip(buf_lengths, chunk_lengths) - ] - usable_lengths = [(length // ratio) * ratio for length in total_lengths] - buffer_lengths = [ - length - usable for length, usable in zip(total_lengths, usable_lengths) - ] - max_total = max(total_lengths, default=0) - max_usable = max(usable_lengths, default=0) - max_buffer = max(buffer_lengths, default=0) - - combined_kv = mx.zeros((B, max_total, kv.shape[-1]), dtype=kv.dtype) - combined_gate = mx.zeros((B, max_total, gate.shape[-1]), dtype=gate.dtype) - for i, (buf_length, chunk_length, total_length) in enumerate( - zip(buf_lengths, chunk_lengths, total_lengths) - ): - parts_kv = [] - parts_gate = [] - if buf_length: - parts_kv.append(buf_kv[i : i + 1, :buf_length]) - parts_gate.append(buf_gate[i : i + 1, :buf_length]) - if chunk_length: - parts_kv.append(kv[i : i + 1, : min(chunk_length, L)]) - parts_gate.append(gate[i : i + 1, : min(chunk_length, L)]) - if parts_kv: - row_kv = ( - parts_kv[0] - if len(parts_kv) == 1 - else mx.concatenate(parts_kv, axis=1) - ) - row_gate = ( - parts_gate[0] - if len(parts_gate) == 1 - else mx.concatenate(parts_gate, axis=1) - ) - combined_kv[i : i + 1, :total_length] = row_kv - combined_gate[i : i + 1, :total_length] = row_gate - - ready_kv = combined_kv[:, :max_usable] - ready_gate = combined_gate[:, :max_usable] - state["buffer_kv"] = mx.zeros((B, max_buffer, kv.shape[-1]), dtype=kv.dtype) - state["buffer_gate"] = mx.zeros( - (B, max_buffer, gate.shape[-1]), dtype=gate.dtype - ) - for i, (usable, buffer_length) in enumerate( - zip(usable_lengths, buffer_lengths) - ): - if buffer_length: - state["buffer_kv"][i : i + 1, :buffer_length] = combined_kv[ - i : i + 1, usable : usable + buffer_length - ] - state["buffer_gate"][i : i + 1, :buffer_length] = combined_gate[ - i : i + 1, usable : usable + buffer_length - ] - state["buffer_lengths"] = buffer_lengths - state["_new_pooled_lengths"] = [usable // ratio for usable in usable_lengths] - - prev_lengths = mx.array(buf_lengths, dtype=mx.int32) - if isinstance(start_pos, mx.array): - pool_base = mx.maximum(start_pos, 0).astype(mx.int32) - else: - pool_base = mx.full((B,), max(0, start_pos), dtype=mx.int32) - return ready_kv, ready_gate, pool_base - prev_lengths - - def update_pool(self, new_pooled: mx.array, state_key: str) -> mx.array: - state = self._branch_state(state_key) - new_lengths = state.pop("_new_pooled_lengths", None) - pool = state["pooled"] - if new_lengths is not None: - B = new_pooled.shape[0] - pool_lengths = self._lengths_list( - state["pooled_lengths"], - B, - 0 if pool is None else pool.shape[1], - ) - total_lengths = [ - pool_length + new_length - for pool_length, new_length in zip(pool_lengths, new_lengths) - ] - max_total = max(total_lengths, default=0) - merged = mx.zeros( - (B, max_total, new_pooled.shape[-1]), dtype=new_pooled.dtype - ) - for i, (pool_length, new_length) in enumerate( - zip(pool_lengths, new_lengths) - ): - if pool is not None and pool_length: - merged[i : i + 1, :pool_length] = pool[i : i + 1, :pool_length] - if new_length: - merged[i : i + 1, pool_length : pool_length + new_length] = ( - new_pooled[i : i + 1, :new_length] - ) - state["pooled"] = merged - state["pooled_lengths"] = total_lengths - return merged - - if new_pooled.shape[1] > 0: - pool = ( - new_pooled - if pool is None - else mx.concatenate([pool, new_pooled], axis=1) - ) - state["pooled"] = pool - state["pooled_lengths"] = None - if pool is None: - pool = mx.zeros( - (new_pooled.shape[0], 0, new_pooled.shape[-1]), new_pooled.dtype - ) - return pool - - def pooled_lengths(self, state_key: str): - return self._branch_state(state_key)["pooled_lengths"] - - def prepare(self, *args, **kwargs): - lengths = kwargs.get("lengths") - right_padding = kwargs.get("right_padding") - self._pending_lengths = ( - list(lengths) - if right_padding is not None and max(right_padding) > 0 - else None - ) - if hasattr(self.local, "prepare"): - self.local.prepare(*args, **kwargs) - return - - left_padding = kwargs.get("left_padding") - if left_padding is not None or ( - right_padding is not None and max(right_padding) > 0 - ): - batch_size = ( - len(left_padding) - if left_padding is not None - else len(kwargs.get("lengths")) - ) - self.local = BatchRotatingKVCache(self.local.max_size, [0] * batch_size) - self.local.prepare(*args, **kwargs) - - def finalize(self): - if hasattr(self.local, "finalize"): - self.local.finalize() - self._pending_lengths = None - - def filter(self, batch_indices): - if isinstance(batch_indices, mx.array): - idx_list = batch_indices.tolist() - else: - idx_list = list(batch_indices) - needed = max(idx_list) + 1 if idx_list else 0 - self._expand_state_to(needed) - - if hasattr(self.local, "filter"): - self.local.filter(batch_indices) - elif self.local.keys is not None: - self.local.keys = self.local.keys[batch_indices] - self.local.values = self.local.values[batch_indices] - for state in (self.compressor_state, self.indexer_state): - for key in self._state_keys: - value = state[key] - if value is not None: - state[key] = value[batch_indices] - for key in self._length_keys: - state[key] = self._filter_lengths(state[key], batch_indices) - - def _expand_state_to(self, batch_size: int): - """Ensure compressor/indexer state arrays/lists have at least `batch_size` - rows along dim 0. Slots without emitted positions are zero-padded.""" - for state in (self.compressor_state, self.indexer_state): - for key in self._state_keys: - value = state[key] - if value is None: - continue - cur = value.shape[0] - if cur < batch_size: - pad_shape = (batch_size - cur,) + value.shape[1:] - state[key] = mx.concatenate( - [value, mx.zeros(pad_shape, dtype=value.dtype)], axis=0 - ) - for key in self._length_keys: - lst = state[key] - if lst is None: - continue - if isinstance(lst, mx.array): - lst = lst.tolist() - if len(lst) < batch_size: - state[key] = lst + [0] * (batch_size - len(lst)) - else: - state[key] = lst - - def extend(self, other): - self_batch = self._cache_batch_size(self.local) - other_batch = self._cache_batch_size(other.local) - if hasattr(self.local, "extend"): - other_local = other.local - if not hasattr(other_local, "filter"): - other_local = self._batch_rotating_from_local(other_local) - self.local.extend(other_local) - elif ( - not hasattr(other.local, "filter") - and self.local.offset == other.local.offset - and self.local._idx == other.local._idx - ): - if self.local.keys is not None or other.local.keys is not None: - self.local.keys = self._concat_optional_local( - self.local.keys, other.local.keys - ) - self.local.values = self._concat_optional_local( - self.local.values, other.local.values - ) - else: - self.local = self._batch_rotating_from_local(self.local) - other_local = ( - other.local - if hasattr(other.local, "filter") - else self._batch_rotating_from_local(other.local) - ) - self.local.extend(other_local) - for self_state, other_state in ( - (self.compressor_state, other.compressor_state), - (self.indexer_state, other.indexer_state), - ): - self_tensors = {key: self_state[key] for key in self._state_keys} - other_tensors = {key: other_state[key] for key in self._state_keys} - for key in self._state_keys: - self_state[key] = self._concat_batch_state( - self_state[key], other_state[key], self_batch, other_batch - ) - for key in self._length_keys: - self_state[key] = self._concat_lengths( - self_state[key], - other_state[key], - self_tensors["pooled" if key.startswith("pooled") else "buffer_kv"], - other_tensors[ - "pooled" if key.startswith("pooled") else "buffer_kv" - ], - self_batch, - other_batch, - ) - - def extract(self, idx): - cache = DeepseekV4Cache(self.local.max_size) - cache.local = ( - self.local.extract(idx) - if hasattr(self.local, "extract") - else self._extract_local(self.local, idx) - ) - for cache_state, self_state in ( - (cache.compressor_state, self.compressor_state), - (cache.indexer_state, self.indexer_state), - ): - for key in self._state_keys: - value = self_state[key] - cache_state[key] = ( - None if value is None else mx.contiguous(value[idx : idx + 1]) - ) - for key in self._length_keys: - lengths = self_state[key] - if lengths is None: - cache_state[key] = None - elif isinstance(lengths, mx.array): - cache_state[key] = [int(lengths.tolist()[idx])] - else: - cache_state[key] = [int(lengths[idx])] - return cache - - @classmethod - def merge(cls, caches): - if not all(c.local.max_size == caches[0].local.max_size for c in caches): - raise ValueError( - "DeepseekV4Cache can only merge caches with the same sliding window" - ) - - cache = cls(caches[0].local.max_size) - cache.local = cls._merge_local([c.local for c in caches]) - for cache_state, state_name in ( - (cache.compressor_state, "compressor_state"), - (cache.indexer_state, "indexer_state"), - ): - for key in cls._state_keys: - cache_state[key] = cls._merge_batch_state( - [getattr(c, state_name)[key] for c in caches] - ) - for key in cls._length_keys: - tensor_key = "pooled" if key.startswith("pooled") else "buffer_kv" - cache_state[key] = cls._merge_lengths( - [getattr(c, state_name)[key] for c in caches], - [getattr(c, state_name)[tensor_key] for c in caches], - ) - return cache - - @staticmethod - def _cache_batch_size(cache): - offset = getattr(cache, "offset", 0) - if isinstance(offset, mx.array) and offset.ndim: - return offset.shape[0] - if cache.keys is not None: - return cache.keys.shape[0] - return 1 - - @staticmethod - def _extract_local(local, idx): - cache = RotatingKVCache(local.max_size, keep=getattr(local, "keep", 0)) - if local.keys is not None: - keys = local._temporal_order(local.keys) - values = local._temporal_order(local.values) - cache.keys = mx.contiguous(keys[idx : idx + 1]) - cache.values = mx.contiguous(values[idx : idx + 1]) - cache._idx = cache.keys.shape[2] - cache.offset = local.offset - return cache - - @classmethod - def _batch_rotating_from_local(cls, local): - batch_size = cls._cache_batch_size(local) - return BatchRotatingKVCache.merge( - [cls._extract_local(local, idx) for idx in range(batch_size)] - ) - - @staticmethod - def _concat_optional_local(a: Optional[mx.array], b: Optional[mx.array]): - if a is None: - return b - if b is None: - return a - return mx.concatenate([a, b], axis=0) - - @classmethod - def _merge_local(cls, locals): - offsets = [local.offset for local in locals] - sizes = [local.size() for local in locals] - use_fast_local = ( - all(not isinstance(offset, mx.array) for offset in offsets) - and all(offset == offsets[0] for offset in offsets) - and all(size == sizes[0] for size in sizes) - ) - if not use_fast_local: - if hasattr(locals[0], "merge"): - return locals[0].merge(locals) - return BatchRotatingKVCache.merge(locals) - - cache = RotatingKVCache(locals[0].max_size, keep=getattr(locals[0], "keep", 0)) - cache.offset = offsets[0] - if sizes[0] == 0: - return cache - - cache.keys = mx.concatenate( - [local._temporal_order(local.keys) for local in locals], axis=0 - ) - cache.values = mx.concatenate( - [local._temporal_order(local.values) for local in locals], axis=0 - ) - cache._idx = cache.keys.shape[2] - return cache - - @staticmethod - def _concat_batch_state( - a: Optional[mx.array], - b: Optional[mx.array], - a_batch: int, - b_batch: int, - ): - if a is None and b is None: - return None - if a is None: - shape = (a_batch, *b.shape[1:]) - a = mx.zeros(shape, dtype=b.dtype) - if b is None: - shape = (b_batch, *a.shape[1:]) - b = mx.zeros(shape, dtype=a.dtype) - if a.shape[2:] != b.shape[2:]: - raise ValueError( - "Cannot extend DeepseekV4Cache entries with different state shapes" - ) - if a.shape[1] != b.shape[1]: - seq_len = max(a.shape[1], b.shape[1]) - if a.shape[1] != seq_len: - padded = mx.zeros((a.shape[0], seq_len, *a.shape[2:]), dtype=a.dtype) - padded[:, : a.shape[1]] = a - a = padded - if b.shape[1] != seq_len: - padded = mx.zeros((b.shape[0], seq_len, *b.shape[2:]), dtype=b.dtype) - padded[:, : b.shape[1]] = b - b = padded - return mx.concatenate([a, b], axis=0) - - @staticmethod - def _full_lengths(lengths, value: Optional[mx.array], batch_size: int): - if lengths is not None: - if isinstance(lengths, mx.array): - lengths = lengths.tolist() - return [int(length) for length in lengths] - length = 0 if value is None else value.shape[1] - return [length] * batch_size - - @classmethod - def _concat_lengths( - cls, - a_lengths, - b_lengths, - a_value: Optional[mx.array], - b_value: Optional[mx.array], - a_batch: int, - b_batch: int, - ): - a_lengths = cls._full_lengths(a_lengths, a_value, a_batch) - b_lengths = cls._full_lengths(b_lengths, b_value, b_batch) - lengths = a_lengths + b_lengths - max_length = max(lengths, default=0) - return None if all(length == max_length for length in lengths) else lengths - - @classmethod - def _merge_lengths(cls, lengths, values): - batch_lengths = [ - cls._full_lengths(length, value, 1)[0] - for length, value in zip(lengths, values) - ] - max_length = max(batch_lengths, default=0) - return ( - None - if all(length == max_length for length in batch_lengths) - else batch_lengths - ) - - @staticmethod - def _merge_batch_state(values: List[Optional[mx.array]]): - present = [v for v in values if v is not None] - if not present: - return None - if not all(v.shape[2:] == present[0].shape[2:] for v in present): - raise ValueError( - "Cannot batch DeepseekV4Cache entries with different state shapes" - ) - seq_len = max(v.shape[1] for v in present) - shape = present[0].shape - dtype = present[0].dtype - merged = [] - for value in values: - if value is None: - merged.append(mx.zeros((1, seq_len, *shape[2:]), dtype=dtype)) - else: - if value.shape[1] != seq_len: - padded = mx.zeros( - (value.shape[0], seq_len, *value.shape[2:]), dtype=value.dtype - ) - padded[:, : value.shape[1]] = value - value = padded - merged.append(value) - return mx.concatenate(merged, axis=0) - - class Compressor(nn.Module): def __init__(self, config: ModelArgs, compress_ratio: int, head_dim: int): @@ -1537,20 +935,19 @@ def __init__(self, config: ModelArgs, compress_ratio: int, head_dim: int): def __call__( self, x: mx.array, - cache: Optional[DeepseekV4Cache], - start_pos: int, - state_key: str = "compressor_state", + pool_cache: Optional[PoolingCache], + offset: Union[int, mx.array], ) -> mx.array: B, _, _ = x.shape kv = self.wkv(x) gate = self.wgate(x) - if cache is None: + 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 = start_pos + pool_base = offset else: - ready_kv, ready_gate, pool_base = cache.accumulate_windows( - kv, gate, state_key, self.compress_ratio, start_pos + ready_kv, ready_gate, pool_base = pool_cache.accumulate_windows( + kv, gate, offset ) if ready_kv.size == 0: @@ -1568,8 +965,8 @@ def __call__( offset=pool_base, ).squeeze(1) - if cache is not None: - new_pooled = cache.update_pool(new_pooled, state_key) + if pool_cache is not None: + new_pooled = pool_cache.update_and_fetch(new_pooled) return new_pooled @@ -1592,15 +989,14 @@ def __call__( x: mx.array, q_residual: mx.array, position_rope: DeepseekV4RoPE, - cache: Optional[DeepseekV4Cache], - start_pos: int, + pool_cache: Optional[PoolingCache], + offset: Union[int, mx.array], ): B, L, _ = x.shape - pooled = self.compressor(x, cache, start_pos, state_key="indexer_state") + pooled = self.compressor(x, pool_cache, offset) if pooled.shape[1] == 0: return None - offset = start_pos 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) @@ -1611,11 +1007,13 @@ def __call__( 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) - lengths = cache.pooled_lengths("indexer_state") if cache is not None else None - if lengths is not None: - lengths = mx.array(lengths) - valid = mx.arange(pooled.shape[1]) < lengths[:, None] - scores = mx.where(valid[:, None], scores, mx.finfo(scores.dtype).min) + 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] @@ -1669,6 +1067,7 @@ def __call__( ) -> 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) @@ -1751,8 +1150,10 @@ def __call__( cache: Optional[Any] = None, ) -> mx.array: B, L, _ = x.shape - v4_cache = cache if isinstance(cache, DeepseekV4Cache) else None - offset = cache.offset if cache is not None else 0 + 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) @@ -1762,50 +1163,25 @@ def __call__( 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, kv) - local_kv_len = kv.shape[2] + 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, v4_cache, offset) + pooled = self.compressor(x, pool_cache, offset) pooled_mask = None if pooled.shape[1] > 0: - lengths = ( - v4_cache.pooled_lengths("compressor_state") - if v4_cache is not None - else None + pooled_mask = ( + pool_cache.make_mask(L, offset) if pool_cache is not None else None ) - if lengths is not None: - lengths = mx.array(lengths) - pooled_mask = (mx.arange(pooled.shape[1]) < lengths[:, None]).reshape( - B, 1, 1, -1 - ) kv = mx.concatenate([kv, pooled[:, None]], axis=2) - # Extend mask to cover pooled entries - if mask is not None and mask.shape[-1] > local_kv_len: - mask = mask[..., -local_kv_len:] - if mask is not None and kv.shape[2] > mask.shape[-1]: - pad_shape = mask.shape[:-1] + (kv.shape[2] - mask.shape[-1],) - if mask.dtype == mx.bool_: - pad = mx.ones(pad_shape, dtype=mask.dtype) - if pooled_mask is not None: - pad = pad & pooled_mask - else: - pad = mx.zeros(pad_shape, dtype=mask.dtype) - if pooled_mask is not None: - pad = mx.where( - pooled_mask, - pad, - mx.full(pad_shape, mx.finfo(mask.dtype).min, dtype=mask.dtype), - ) - mask = mx.concatenate([mask, pad], axis=-1) + mask = _extend_mask(mask, pooled_mask, kv.shape[2]) out = scaled_dot_product_attention( q, kv, kv, - cache=cache, + cache=local_cache, scale=self.scale, mask=mask, sinks=self.attn_sink.astype(q.dtype), @@ -1821,17 +1197,9 @@ def __call__( 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 V4Attention(config, layer_idx) - +class SparseCompressedAttention(nn.Module): + """DeepSeek V4 attention with sparse indexed pooled KV compression.""" -class V4Attention(nn.Module): def __init__(self, config: ModelArgs, layer_idx: int): super().__init__() self.config = config @@ -1840,8 +1208,6 @@ def __init__(self, config: ModelArgs, layer_idx: int): self.hidden_size = config.hidden_size self.n_heads = config.num_attention_heads self.head_dim = config.head_dim - self.rope_head_dim = config.qk_rope_head_dim - self.nope_head_dim = self.head_dim - self.rope_head_dim self.o_groups = config.o_groups self.o_lora_rank = config.o_lora_rank self.scale = self.head_dim**-0.5 @@ -1853,10 +1219,10 @@ def __init__(self, config: ModelArgs, layer_idx: int): ) 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 = nn.Linear( + self.wo_a = MultiLinear( self.n_heads * self.head_dim // config.o_groups, - config.o_groups * config.o_lora_rank, - bias=False, + config.o_lora_rank, + config.o_groups, ) self.wo_b = nn.Linear( config.o_groups * config.o_lora_rank, @@ -1864,81 +1230,15 @@ def __init__(self, config: ModelArgs, layer_idx: int): bias=config.attention_bias, ) self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) - self._q_l2_norm_weight = (mx.ones((self.head_dim,)),) - self._cached_dtype = None - rope_theta = ( - config.compress_rope_theta if self.compress_ratio else config.rope_theta - ) - rope_scaling = config.rope_scaling if self.compress_ratio else None self.rope = DeepseekV4RoPE( config.qk_rope_head_dim, - rope_theta, - rope_scaling, + config.compress_rope_theta, + config.rope_scaling, config.max_position_embeddings, ) - if self.compress_ratio: - self.compressor = Compressor(config, self.compress_ratio, self.head_dim) - if self.compress_ratio == 4: - self.indexer = Indexer(config, self.compress_ratio) - - def _ensure_cached(self, dtype): - dtype_key = str(dtype) - if self._cached_dtype is not None and self._cached_dtype == dtype_key: - return - self._cached_dtype = dtype_key - self._attn_sink_cached = self.attn_sink.astype(dtype) - self._q_norm_weight_cached = self._q_l2_norm_weight[0].astype(dtype) - if isinstance(self.wo_a, nn.QuantizedLinear): - self._wo_a_weight = self.wo_a.weight.reshape( - self.o_groups, self.o_lora_rank, -1 - )[:, None] - self._wo_a_scales = self.wo_a.scales.reshape( - self.o_groups, self.o_lora_rank, -1 - )[:, None] - self._wo_a_biases = ( - None - if self.wo_a.biases is None - else self.wo_a.biases.reshape(self.o_groups, self.o_lora_rank, -1)[ - :, None - ] - ) - else: - group_feat = (self.n_heads * self.head_dim) // self.o_groups - self._wo_a_weight_reshaped = self.wo_a.weight.reshape( - self.o_groups, self.o_lora_rank, group_feat - ) - - def _grouped_output_projection(self, out: mx.array) -> mx.array: - B, _, L, _ = out.shape - heads_per_group = self.n_heads // self.o_groups - out = out.reshape(B, self.o_groups, heads_per_group, L, self.head_dim) - out = out.transpose(1, 0, 3, 2, 4) - out = out.reshape(self.o_groups, B, L, heads_per_group * self.head_dim) - - if isinstance(self.wo_a, nn.QuantizedLinear): - out = mx.quantized_matmul( - out, - self._wo_a_weight, - scales=self._wo_a_scales, - biases=self._wo_a_biases, - transpose=True, - group_size=self.wo_a.group_size, - bits=self.wo_a.bits, - mode=self.wo_a.mode, - ) - out = out.transpose(1, 2, 0, 3).reshape( - B, L, self.o_groups * self.o_lora_rank - ) - if "bias" in self.wo_a: - out = out + self.wo_a.bias - return out - - out = mx.einsum("gbsd,grd->bsgr", out, self._wo_a_weight_reshaped) - out = out.reshape(B, L, self.o_groups * self.o_lora_rank) - if "bias" in self.wo_a: - out = out + self.wo_a.bias - return out + self.compressor = Compressor(config, self.compress_ratio, self.head_dim) + self.indexer = Indexer(config, self.compress_ratio) def __call__( self, @@ -1947,138 +1247,44 @@ def __call__( cache: Optional[Any] = None, ) -> mx.array: B, L, _ = x.shape - local_cache = cache + 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 - if isinstance(offset, mx.array): - offset = offset + 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) - self._ensure_cached(q.dtype) - q = mx.fast.rms_norm(q, self._q_norm_weight_cached, self.config.rms_norm_eps) + q = mx.fast.rms_norm(q, None, self.config.rms_norm_eps) q = q.transpose(0, 2, 1, 3) - kv = self.kv_norm(self.wkv(x)).reshape(B, 1, L, self.head_dim) - q = self.rope(q, offset) - kv = self.rope(kv, 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, kv) - full_kv = kv - local_kv_len = kv.shape[2] + kv, _ = local_cache.update_and_fetch(kv, mx.zeros((B, 1, L, 0))) - pooled_mask = None - pooled_bias = None - sparse_pooled = None - sparse_topk = None - sparse_pooled_mask = None - if self.compress_ratio: - v4_cache = cache if isinstance(cache, DeepseekV4Cache) else None - pooled = self.compressor(x, v4_cache, offset) - if pooled.shape[1] > 0: - lengths = ( - v4_cache.pooled_lengths("compressor_state") - if v4_cache is not None - else None - ) - use_indexer = isinstance(getattr(self, "indexer", None), Indexer) - max_pooled_length = pooled.shape[1] if lengths is None else max(lengths) - select_all = use_indexer and ( - max_pooled_length <= self.indexer.index_topk - ) - if select_all: - pooled = pooled[:, None] - pooled_bias = math.log(L) - if lengths is not None: - lengths = mx.array(lengths) - pooled_mask = ( - mx.arange(pooled.shape[2]) < lengths[:, None] - ).reshape(B, 1, 1, -1) - elif use_indexer: - topk = self.indexer(x, q_residual, self.rope, v4_cache, offset) - if topk is not None: - if L > 1: - sparse_pooled = pooled - sparse_topk = topk - if lengths is not None: - lengths = mx.array(lengths) - sparse_pooled_mask = ( - topk < lengths[:, None, None] - ).reshape(B, 1, L, -1) - else: - if lengths is not None: - lengths = mx.array(lengths) - pooled_mask = (topk < lengths[:, None, None]).reshape( - B, 1, 1, -1 - ) - expanded = mx.broadcast_to( - pooled[:, None, None, :, :], - (B, 1, L, pooled.shape[1], self.head_dim), - ) - idx = topk[:, None, :, :, None] - pooled = mx.take_along_axis( - expanded, - mx.broadcast_to(idx, idx.shape[:-1] + (self.head_dim,)), - axis=3, - ).reshape(B, 1, -1, self.head_dim) - else: - pooled = pooled[:, None] - else: - if lengths is not None: - lengths = mx.array(lengths) - pooled_mask = ( - mx.arange(pooled.shape[1]) < lengths[:, None] - ).reshape(B, 1, 1, -1) - pooled = pooled[:, None] - if sparse_topk is None: - full_kv = mx.concatenate([full_kv, pooled], axis=2) - - if mask is not None and mask.shape[-1] > local_kv_len: - mask = mask[..., -local_kv_len:] - - if sparse_topk is not None: - out = _sparse_pooled_attention( + 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, - full_kv, - sparse_pooled, - sparse_topk, - mask, - sparse_pooled_mask, - self.scale, - self._attn_sink_cached, + kv, + kv, + cache=local_cache, + scale=self.scale, + mask=mask, + sinks=sinks, ) - else: - if mask is not None and full_kv.shape[2] > mask.shape[-1]: - pad_shape = mask.shape[:-1] + (full_kv.shape[2] - mask.shape[-1],) - pad_pooled_mask = pooled_mask - if ( - pad_pooled_mask is not None - and pad_pooled_mask.shape[-1] != pad_shape[-1] - ): - pad_pooled_mask = pad_pooled_mask[..., -pad_shape[-1] :] - if pooled_bias is not None: - dtype = q.dtype - if mask.dtype == mx.bool_: - mask = mx.where( - mask, - mx.array(0, dtype=dtype), - mx.full((), mx.finfo(dtype).min, dtype=dtype), - ) - pad = mx.full(pad_shape, pooled_bias, dtype=mask.dtype) - elif mask.dtype == mx.bool_: - pad = mx.ones(pad_shape, dtype=mask.dtype) - if pad_pooled_mask is not None: - pad = pad & pad_pooled_mask - else: - pad = mx.zeros(pad_shape, dtype=mask.dtype) - if pad_pooled_mask is not None: - pad = mx.where( - pad_pooled_mask, - pad, - mx.full( - pad_shape, mx.finfo(mask.dtype).min, dtype=mask.dtype - ), - ) - mask = mx.concatenate([mask, pad], axis=-1) + + # 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, @@ -2086,11 +1292,48 @@ def __call__( cache=local_cache, scale=self.scale, mask=mask, - sinks=self._attn_sink_cached, + 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 = self._grouped_output_projection(out) - return self.wo_b(out) + + 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) + + 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): @@ -2149,9 +1392,7 @@ def __call__(self, inputs: mx.array, cache: Optional[Any] = None) -> mx.array: first_cache = cache[0] mask_cache = ( - first_cache.local - if isinstance(first_cache, DeepseekV4Cache) - else first_cache + first_cache[0] if isinstance(first_cache, CacheList) else first_cache ) mask = create_attention_mask( h[:, :, 0, :], @@ -2169,8 +1410,8 @@ def __call__(self, inputs: mx.array, cache: Optional[Any] = None) -> mx.array: if pipeline_rank != 0: h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size) cache_item = cache[-1] - if isinstance(cache_item, DeepseekV4Cache): - cache_item = cache_item.local + 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) @@ -2211,10 +1452,26 @@ def predicate(k): def make_cache(self): caches = [] for layer in self.layers: - if layer.attn.compress_ratio: - caches.append(DeepseekV4Cache(self.args.sliding_window)) - else: + 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]: @@ -2307,11 +1564,8 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.{suffix}" ] = mx.stack(stacked) - # Reshape wo_a from nn.Linear (2D) to MultiLinear (3D) for - # LocalAttention (ratio=0) and CompressedAttention (ratio=128) layers + # Reshape wo_a from nn.Linear (2D) to MultiLinear (3D) for all layers for layer_idx in range(n_layers): - if self.args.compress_ratios[layer_idx] == 4: - continue 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: diff --git a/tests/test_models.py b/tests/test_models.py index f494b4f31..ca48075d7 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1526,13 +1526,6 @@ def test_deepseek_v4(self): main_layer = model.model.layers[0].attn compress_layer = model.model.layers[2].attn - class RaisingIndexer: - index_topk = 1 - - def __call__(self, *args, **kwargs): - raise AssertionError("Indexer should not run during prefill") - - compress_layer.indexer = RaisingIndexer() inputs = mx.array([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=mx.int32) cache = model.make_cache() logits = model(inputs, cache=cache) From ca8b29905ce9642e35b6e8b4a4cadc8aefcbdccf Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Fri, 1 May 2026 02:07:52 -0700 Subject: [PATCH 51/54] Simplify HyperConnection --- mlx_lm/models/deepseek_v4.py | 460 +---------------------------------- tests/test_models.py | 18 +- 2 files changed, 17 insertions(+), 461 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index c2251f26a..036fe9e64 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -12,6 +12,7 @@ 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 @@ -363,457 +364,6 @@ def _sparse_pooled_attention( return out.astype(q.dtype) -@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 _make_hc_split_sinkhorn_kernel(): - if mx.default_device() != mx.gpu or not mx.metal.is_available(): - return None - - source = """ - uint idx = thread_position_in_grid.x; - constexpr int MIX = (2 + HC) * HC; - constexpr int BASE = 2 * HC; - - const device float* mix = (const device float*)mixes + idx * MIX; - device float* pre_out = (device float*)pre + idx * HC; - device float* post_out = (device float*)post + idx * HC; - device float* comb_out = (device float*)comb + idx * HC * HC; - - const float pre_scale = scale[0]; - const float post_scale = scale[1]; - const float comb_scale = scale[2]; - const float epsv = eps[0]; - - // Pre-sigmoid - { - float4 z = *(const device float4*)mix * pre_scale - + *(const device float4*)base; - *(device float4*)pre_out = 1.0f / (1.0f + metal::fast::exp(-z)) + epsv; - } - - // Post-sigmoid - { - float4 z = *(const device float4*)(mix + HC) * post_scale - + *(const device float4*)(base + HC); - *(device float4*)post_out = 2.0f * 1.0f / (1.0f + metal::fast::exp(-z)); - } - - // Comb: four float4 loads — all independent, GPU issues in parallel - float4 v0 = *(const device float4*)(mix + BASE ) * comb_scale + *(const device float4*)(base + BASE ); - float4 v1 = *(const device float4*)(mix + BASE + 4) * comb_scale + *(const device float4*)(base + BASE + 4); - float4 v2 = *(const device float4*)(mix + BASE + 8) * comb_scale + *(const device float4*)(base + BASE + 8); - float4 v3 = *(const device float4*)(mix + BASE + 12) * comb_scale + *(const device float4*)(base + BASE + 12); - - // Per-row stable softmax: compute all maxes before any exp - float m0 = metal::max(metal::max(v0.x, v0.y), metal::max(v0.z, v0.w)); - float m1 = metal::max(metal::max(v1.x, v1.y), metal::max(v1.z, v1.w)); - float m2 = metal::max(metal::max(v2.x, v2.y), metal::max(v2.z, v2.w)); - float m3 = metal::max(metal::max(v3.x, v3.y), metal::max(v3.z, v3.w)); - - float4 e0 = metal::fast::exp(v0 - m0); - float4 e1 = metal::fast::exp(v1 - m1); - float4 e2 = metal::fast::exp(v2 - m2); - float4 e3 = metal::fast::exp(v3 - m3); - - // Explicit adds instead of dot(e, 1) — avoids unnecessary fmul - float4 r0 = e0 * 1.0f / (e0.x + e0.y + e0.z + e0.w) + epsv; - float4 r1 = e1 * 1.0f / (e1.x + e1.y + e1.z + e1.w) + epsv; - float4 r2 = e2 * 1.0f / (e2.x + e2.y + e2.z + e2.w) + epsv; - float4 r3 = e3 * 1.0f / (e3.x + e3.y + e3.z + e3.w) + epsv; - - // Initial column normalization - float4 col = 1.0f / (r0 + r1 + r2 + r3 + epsv); - r0 *= col; r1 *= col; r2 *= col; r3 *= col; - - // Sinkhorn iterations - for (int iter = 1; iter < ITERS; ++iter) { - r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv); - r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv); - r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv); - r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv); - col = 1.0f / (r0 + r1 + r2 + r3 + epsv); - r0 *= col; r1 *= col; r2 *= col; r3 *= col; - } - - // Write comb output (four aligned 128-bit stores) - *(device float4*)(comb_out) = r0; - *(device float4*)(comb_out + 4) = r1; - *(device float4*)(comb_out + 8) = r2; - *(device float4*)(comb_out + 12) = r3; - """ - - return mx.fast.metal_kernel( - name="deepseek_v4_hc_split_sinkhorn", - input_names=["mixes", "scale", "base", "eps"], - output_names=["pre", "post", "comb"], - source=source, - ) - - -_hc_split_sinkhorn_kernel = _make_hc_split_sinkhorn_kernel() - - -def hc_split_sinkhorn( - mixes: mx.array, - scale: mx.array, - base: mx.array, - hc_mult: int, - sinkhorn_iters: int, - eps: float, -) -> Tuple[mx.array, mx.array, mx.array]: - if _hc_split_sinkhorn_kernel is None or hc_mult != 4: - return _hc_split_sinkhorn_ops(mixes, scale, base, hc_mult, sinkhorn_iters, eps) - - if not isinstance(eps, mx.array): - eps = mx.array([eps], dtype=mx.float32) - n_rows = mixes.size // ((2 + hc_mult) * hc_mult) - return _hc_split_sinkhorn_kernel( - inputs=[mixes, scale, base, eps], - template=[("HC", hc_mult), ("ITERS", sinkhorn_iters)], - grid=(n_rows, 1, 1), - threadgroup=(256, 1, 1), - output_shapes=[ - (*mixes.shape[:-1], hc_mult), - (*mixes.shape[:-1], hc_mult), - (*mixes.shape[:-1], hc_mult, hc_mult), - ], - output_dtypes=[mx.float32, mx.float32, mx.float32], - ) - - -@mx.compile -def _hc_collapse_op(pre: mx.array, x: mx.array) -> mx.array: - return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.dtype) - - -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; - - 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 epsv = eps[0]; - - 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)) + epsv; - 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 + epsv)) - + epsv * 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) - ) + epsv); - 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 + epsv)) * 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) - ) + epsv); - 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, native bfloat4 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 bfloat16_t* x_row = (const device bfloat16_t*)x_in - + row * (HC * D); - device bfloat16_t* out_row = (device bfloat16_t*)collapsed - + row * D; - - // Native bfloat4 pointers: single 64-bit load per vector - using bf4 = vec; - const device bf4* x_row0 = (const device bf4*)(x_row + 0*D); - const device bf4* x_row1 = (const device bf4*)(x_row + 1*D); - const device bf4* x_row2 = (const device bf4*)(x_row + 2*D); - const device bf4* x_row3 = (const device bf4*)(x_row + 3*D); - device bf4* out4 = (device bf4*)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] = bf4(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] = (bfloat16_t)val; - } - #endif - """ - - return mx.fast.metal_kernel( - name="deepseek_v4_hc_sinkhorn_collapse", - input_names=["mixes", "scale", "base", "eps", "x_in"], - output_names=["post", "comb", "collapsed"], - source=source, - ) - - -_hc_sinkhorn_collapse_kernel = _make_hc_sinkhorn_collapse_kernel() - - -@mx.compile -def _hc_expand_op( - post: mx.array, - block_out: mx.array, - comb: mx.array, - residual: mx.array, -) -> mx.array: - y = post[..., None] * block_out[:, :, None, :].astype(mx.float32) - y = y + mx.matmul(comb.swapaxes(-1, -2), residual.astype(mx.float32)) - return y.astype(block_out.dtype) - - -@mx.compile -def _rms_rsqrt(flat: mx.array, eps: float) -> mx.array: - return mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + eps) - - -@mx.compile -def _hc_mixes(flat: mx.array, fn_T: mx.array, norm_eps: float) -> mx.array: - """Fused RMS-rsqrt + matmul + scale into single compiled graph.""" - rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + norm_eps) - return (flat @ fn_T) * rsqrt - - -class HyperConnection(nn.Module): - def __init__(self, config: ModelArgs): - super().__init__() - self.hc_mult = config.hc_mult - self.sinkhorn_iters = config.hc_sinkhorn_iters - self.hc_eps = config.hc_eps - self._hc_eps = (mx.array([config.hc_eps], dtype=mx.float32),) - 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) - self._fn_T = None - - def compute_weights(self, x: mx.array): - B, L, H, D = x.shape - flat = x.reshape(B, L, H * D).astype(mx.float32) - if self._fn_T is None: - self._fn_T = self.fn.T - if self.training: - rsqrt = _rms_rsqrt(flat, self.norm_eps) - mixes = (flat @ self._fn_T) * rsqrt - else: - mixes = _hc_mixes(flat, self._fn_T, self.norm_eps) - split_sinkhorn = _hc_split_sinkhorn_ops if self.training else hc_split_sinkhorn - return split_sinkhorn( - mixes, - self.scale, - self.base, - self.hc_mult, - self.sinkhorn_iters, - self.hc_eps if self.training else self._hc_eps[0], - ) - - def collapse(self, x: mx.array): - if ( - not self.training - and _hc_sinkhorn_collapse_kernel is not None - and self.hc_mult == 4 - and x.dtype == mx.bfloat16 - ): - return self._fused_collapse(x) - pre, post, comb = self.compute_weights(x) - return _hc_collapse_op(pre, x), post, comb - - def _fused_collapse(self, x: mx.array): - """Fused sinkhorn + collapse in a single Metal kernel dispatch.""" - B, L, H, D = x.shape - flat = x.reshape(B, L, H * D).astype(mx.float32) - if self._fn_T is None: - self._fn_T = self.fn.T - mixes = _hc_mixes(flat, self._fn_T, self.norm_eps) - - eps = self._hc_eps[0] - n_rows = B * L - x_flat = mx.contiguous(x.reshape(n_rows, H, D)) - - post, comb, collapsed = _hc_sinkhorn_collapse_kernel( - inputs=[mixes, self.scale, self.base, eps, x_flat], - template=[("HC", self.hc_mult), ("ITERS", self.sinkhorn_iters), ("D", D)], - grid=(n_rows * 256, 1, 1), - threadgroup=(256, 1, 1), - output_shapes=[ - (*mixes.shape[:-1], self.hc_mult), - (*mixes.shape[:-1], self.hc_mult, self.hc_mult), - (B, L, D), - ], - output_dtypes=[mx.float32, mx.float32, x.dtype], - ) - return collapsed, post, comb - - def expand( - self, - block_out: mx.array, - residual: mx.array, - post: mx.array, - comb: mx.array, - ): - return _hc_expand_op(post, block_out, comb, residual) - - -@mx.compile -def _hyper_head_op( - x: mx.array, - fn: mx.array, - scale: mx.array, - base: mx.array, - norm_eps: float, - hc_eps: float, -) -> mx.array: - """Fused HyperHead: RMS-rsqrt + matmul + sigmoid + weighted sum.""" - B, L, H, D = x.shape - flat = x.reshape(B, L, H * D).astype(mx.float32) - rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + norm_eps) - mixes = (flat @ fn.T) * rsqrt - pre = mx.sigmoid(mixes * scale[0] + base) + hc_eps - return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.dtype) - - -class HyperHead(nn.Module): - def __init__(self, config: ModelArgs): - 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): - if not self.training: - return _hyper_head_op( - x, self.fn, self.scale, self.base, self.norm_eps, self.hc_eps - ) - B, L, H, D = x.shape - flat = x.reshape(B, L, H * D).astype(mx.float32) - rsqrt = _rms_rsqrt(flat, self.norm_eps) - mixes = (flat @ self.fn.T) * rsqrt - pre = mx.sigmoid(mixes * self.scale[0] + self.base) + self.hc_eps - return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.dtype) - - class MoEGate(nn.Module): def __init__(self, config: ModelArgs, layer_idx: int): super().__init__() @@ -1354,14 +904,14 @@ def __call__( input_ids: mx.array, ) -> mx.array: residual = h - x, post, comb = self.attn_hc.collapse(h) + x, post, comb = self.attn_hc(h) x = self.attn(self.attn_norm(x), mask=mask, cache=cache) - h = self.attn_hc.expand(x, residual, post, comb) + h = hc_expand(x, residual, post, comb) residual = h - x, post, comb = self.ffn_hc.collapse(h) + x, post, comb = self.ffn_hc(h) x = self.ffn(self.ffn_norm(x), input_ids) - return self.ffn_hc.expand(x, residual, post, comb) + return hc_expand(x, residual, post, comb) class DeepseekV4Model(PipelineMixin, nn.Module): diff --git a/tests/test_models.py b/tests/test_models.py index ca48075d7..23907b421 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1429,6 +1429,10 @@ def test_deepseek_v3(self): 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) @@ -1468,14 +1472,16 @@ def test_deepseek_v4(self): 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) - expected = deepseek_v4._hc_split_sinkhorn_ops( + pre, post, comb = _hc_split_sinkhorn_ops( mixes, scale, base, hc_mult, 20, 1e-6 ) - actual = deepseek_v4.hc_split_sinkhorn( - 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) ) - for x, y in zip(expected, actual): - self.assertTrue(mx.allclose(x, y, rtol=1e-5, atol=1e-5)) # Expand test post = mx.random.normal((2, 3, 4), dtype=mx.float32) @@ -1487,7 +1493,7 @@ def test_deepseek_v4(self): comb.swapaxes(-1, -2), residual.astype(mx.float32) ) expected = expected.astype(block_out.dtype) - actual = deepseek_v4._hc_expand_op(post, block_out, comb, residual) + actual = hc_expand(block_out, residual, post, comb) self.assertTrue(mx.allclose(actual, expected, rtol=1e-5, atol=1e-5)) # Model test From c58834b4c04b2cedb6f8e2b13462a0e73e262a58 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Fri, 1 May 2026 02:13:54 -0700 Subject: [PATCH 52/54] Add hyper_connection.py --- mlx_lm/models/hyper_connection.py | 285 ++++++++++++++++++++++++++++++ 1 file changed, 285 insertions(+) create mode 100644 mlx_lm/models/hyper_connection.py 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) From 7efb57fa570d1a2ebf449810a6c1ce0cfcb7a6a2 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Fri, 1 May 2026 03:04:13 -0700 Subject: [PATCH 53/54] Fix batch cache edge case --- mlx_lm/models/cache.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py index 82fdbfd37..b3c698ec1 100644 --- a/mlx_lm/models/cache.py +++ b/mlx_lm/models/cache.py @@ -1161,8 +1161,20 @@ def accumulate_windows(self, kv: mx.array, gate: mx.array, offset): # Fill new remainder buffer from the tail of the input if nr > 0: - new_buf_kv[i, :nr] = kv[i, vl - nr : vl] - new_buf_gate[i, :nr] = gate[i, vl - nr : vl] + 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 From 5c10538136b9038b9626c134612b08afc18d697a Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Fri, 1 May 2026 04:36:53 -0700 Subject: [PATCH 54/54] Fix tensor parallel distributed --- mlx_lm/models/deepseek_v4.py | 27 +++++++++++++++++++++++---- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py index 036fe9e64..49f1d8bd9 100644 --- a/mlx_lm/models/deepseek_v4.py +++ b/mlx_lm/models/deepseek_v4.py @@ -609,6 +609,8 @@ def __init__(self, config: ModelArgs, layer_idx: int): config.max_position_embeddings, ) + self.sharding_group = None + def __call__( self, x: mx.array, @@ -647,6 +649,9 @@ def __call__( 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 @@ -693,6 +698,8 @@ def __init__(self, config: ModelArgs, layer_idx: int): ) self.compressor = Compressor(config, self.compress_ratio, self.head_dim) + self.sharding_group = None + def __call__( self, x: mx.array, @@ -744,6 +751,9 @@ def __call__( 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 @@ -790,6 +800,8 @@ def __init__(self, config: ModelArgs, layer_idx: int): 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, @@ -873,6 +885,9 @@ def __call__( 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 @@ -1128,13 +1143,17 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: 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", group=group - ) - layer.attn.wo_b = shard_linear( - layer.attn.wo_b, "sharded-to-all", group=group + 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