diff --git a/mlx_lm/convert.py b/mlx_lm/convert.py index ab3fc62ac..6e4f51f0c 100644 --- a/mlx_lm/convert.py +++ b/mlx_lm/convert.py @@ -63,14 +63,26 @@ def mixed_quant_predicate( or index >= 7 * num_layers // 8 or (index - num_layers // 8) % 3 == 2 ) + always_more_bits = ( + "lm_head" in path + or "embed_tokens" in path + or "wq_a" in path + or "wq_b" in path + or "wkv" in path + or "wo_a" in path + or "wo_b" in path + or "compressor" in path + or "indexer" in path + or "shared_experts" in path + ) + if always_more_bits: + return {"group_size": group_size, "bits": high_bits, "mode": mode} if ( "v_proj" in path or "v_a_proj" in path or "v_b_proj" in path ) and use_more_bits: return {"group_size": group_size, "bits": high_bits, "mode": mode} if "down_proj" in path and use_more_bits: return {"group_size": group_size, "bits": high_bits, "mode": mode} - if "lm_head" in path: - return {"group_size": group_size, "bits": high_bits, "mode": mode} return {"group_size": group_size, "bits": low_bits, "mode": mode} diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index 3573b2640..649063809 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -5,8 +5,11 @@ import copy import functools import json +import math +import random import sys import time +import warnings from collections import deque from dataclasses import dataclass from functools import partial @@ -219,6 +222,12 @@ def setup_arg_parser(): help="Number of tokens to draft when using speculative decoding.", default=3, ) + parser.add_argument( + "--mtp", + action="store_true", + help="Use native Multi-Token Prediction for speculative decoding " + "(requires a model with an MTP head, e.g. DeepSeek-V4).", + ) return parser @@ -654,12 +663,186 @@ def _draft_generate(y, num_draft): _rewind_cache(num_draft, n) +def mtp_generate_step( + prompt: mx.array, + model: nn.Module, + *, + max_tokens: int = 256, + sampler: Optional[Callable[[mx.array], mx.array]] = None, + logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None, + prompt_cache: Optional[Any] = None, + prefill_step_size: int = 2048, + kv_bits: Optional[int] = None, + kv_group_size: int = 64, + quantized_kv_start: int = 0, +) -> Generator[Tuple[mx.array, mx.array, bool], None, None]: + """A generator that uses the model's native MTP head for speculative decoding. + + Each iteration runs one backbone forward pass over the current token and its + pending draft, then one MTP forward pass to propose the next draft. Up to 2 + tokens are emitted per backbone step: one always-accepted backbone token and + one conditionally-accepted draft token. + + The model must implement ``mtp_forward(hidden, next_tok, mtp_cache)`` and + support ``return_hidden=True`` in its ``__call__``. + + Yields: + Tuple[mx.array, mx.array, bool]: (token, log-probabilities, from_draft). + ``from_draft`` is ``True`` when the token came from the MTP head. + """ + y = prompt.astype(mx.uint32) + prev_tokens = None + + if prompt_cache is None: + model_cache = cache.make_prompt_cache(model) + mtp_cache = model.make_mtp_cache() + else: + n_main = len(model.layers) + model_cache = prompt_cache[:n_main] + mtp_cache = prompt_cache[n_main:] or model.make_mtp_cache() + + _is_greedy = sampler is None + sampler = sampler or (lambda x: mx.argmax(x, axis=-1)) + + quantize_cache_fn = functools.partial( + maybe_quantize_kv_cache, + quantized_kv_start=quantized_kv_start, + kv_group_size=kv_group_size, + kv_bits=kv_bits, + ) + + def _process_and_sample(tokens, logits): + if logits_processors: + for processor in logits_processors: + logits = processor(tokens, logits) + logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) + return sampler(logprobs), logprobs + + def _step_backbone(y, prev_tokens, n_predict=1): + with mx.stream(generation_stream): + logits, hidden = model( + y[None], cache=model_cache, return_hidden=True + ) + logits = logits[:, -n_predict:, :] + quantize_cache_fn(model_cache) + toks, lps = [], [] + for i in range(n_predict): + if logits_processors: + prev_tokens = ( + mx.concatenate([prev_tokens, y[i : i + 1]]) + if prev_tokens is not None + else y[i : i + 1] + ) + tok, lp = _process_and_sample(prev_tokens, logits[:, i, :].squeeze(0)) + toks.append(tok) + lps.append(lp) + return mx.stack(toks), mx.stack(lps), hidden, prev_tokens + + def _step_mtp(hidden_last, main_tok, prev_tokens): + next_ids = main_tok.reshape(1, 1) + with mx.stream(generation_stream): + mtp_logits = model.mtp_forward(hidden_last, next_ids, mtp_cache) + quantize_cache_fn(mtp_cache) + mtp_logits = mtp_logits[:, -1, :].squeeze(0) + if logits_processors: + tokens_for_proc = ( + mx.concatenate([prev_tokens, main_tok.reshape(-1)]) + if prev_tokens is not None + else main_tok.reshape(-1) + ) + else: + tokens_for_proc = prev_tokens + draft_tok, draft_lp = _process_and_sample(tokens_for_proc, mtp_logits) + return draft_tok, draft_lp + + def _prefill(y): + while y.size > 1: + n = min(prefill_step_size, y.size - 1) + model(y[:n][None], cache=model_cache) + quantize_cache_fn(model_cache) + mx.eval([c.state for c in model_cache if hasattr(c, "state")]) + y = y[n:] + mx.clear_cache() + return y + + with mx.stream(generation_stream): + y = _prefill(y) + + ntoks = 0 + draft_tok = draft_lp = None + + while ntoks < max_tokens: + if draft_tok is None: + toks, lps, hidden, prev_tokens = _step_backbone(y, prev_tokens, n_predict=1) + mx.eval(toks) + main_tok, main_lp = toks[0], lps[0] + ntoks += 1 + yield main_tok.item(), main_lp, False + if ntoks >= max_tokens: + return + hidden_at_main = hidden[:, -1:, :] + draft_tok, draft_lp = _step_mtp(hidden_at_main, main_tok, prev_tokens) + mx.eval(draft_tok) + y = mx.array([main_tok.item()], mx.uint32) + else: + y_with_draft = mx.concatenate([y, mx.array([draft_tok.item()], mx.uint32)]) + toks, lps, hidden, prev_tokens = _step_backbone( + y_with_draft, prev_tokens, n_predict=2 + ) + mx.eval(toks, draft_tok) + + verify_pred, bonus_tok = toks[0], toks[1] + verify_lp, bonus_lp = lps[0], lps[1] + draft_tok_id = draft_tok.item() + + if _is_greedy: + accept = verify_pred.item() == draft_tok_id + else: + log_accept = (verify_lp[draft_tok_id] - draft_lp[draft_tok_id]).item() + accept = log_accept >= 0 or random.random() < math.exp(log_accept) + + hidden_at_confirmed = hidden[:, 0:1, :] + hidden_at_draft = hidden[:, 1:2, :] + + if accept: + ntoks += 1 + yield draft_tok_id, draft_lp, True + if ntoks >= max_tokens: + return + ntoks += 1 + yield bonus_tok.item(), bonus_lp, False + if ntoks >= max_tokens: + return + draft_tok, draft_lp = _step_mtp(hidden_at_draft, bonus_tok, prev_tokens) + mx.eval(draft_tok) + y = mx.array([bonus_tok.item()], mx.uint32) + else: + # Reject draft: trim the draft token from both caches + for c in model_cache: + c.trim(1) + for c in mtp_cache: + c.trim(1) + if logits_processors and prev_tokens is not None: + prev_tokens = prev_tokens[:-1] + verify_tok_id = verify_pred.item() + ntoks += 1 + yield verify_tok_id, verify_lp, False + if ntoks >= max_tokens: + return + draft_tok, draft_lp = _step_mtp( + hidden_at_confirmed, verify_pred, prev_tokens + ) + mx.eval(draft_tok) + y = mx.array([verify_tok_id], mx.uint32) + + def stream_generate( model: nn.Module, tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper], prompt: Union[str, mx.array, List[int]], max_tokens: int = 256, draft_model: Optional[nn.Module] = None, + mtp: bool = False, **kwargs, ) -> Generator[GenerationResponse, None, None]: """ @@ -675,6 +858,8 @@ def stream_generate( draft_model (Optional[nn.Module]): An optional draft model. If provided then speculative decoding is used. The draft model must use the same tokenizer as the main model. Default: ``None``. + mtp (bool): Use native Multi-Token Prediction for speculative + decoding. Requires a model with an MTP head. Default: ``False``. kwargs: The remaining options get passed to :func:`generate_step`. See :func:`generate_step` for more details. @@ -698,19 +883,30 @@ def stream_generate( kwargs["max_tokens"] = max_tokens - if draft_model is None: + if draft_model is not None: + kwargs.pop("max_kv_size", None) + kwargs.pop("prompt_progress_callback", None) + token_generator = speculative_generate_step( + prompt, model, draft_model, **kwargs + ) + elif mtp and hasattr(model, "mtp"): + kwargs.pop("max_kv_size", None) + kwargs.pop("prompt_progress_callback", None) + kwargs.pop("num_draft_tokens", None) + token_generator = mtp_generate_step(prompt, model, **kwargs) + else: + if mtp: + warnings.warn( + "--mtp flag ignored: model does not have an MTP head. " + "Falling back to standard generation.", + stacklevel=2, + ) kwargs.pop("num_draft_tokens", None) token_generator = generate_step(prompt, model, **kwargs) # from_draft always false for non-speculative generation token_generator = ( (token, logprobs, False) for token, logprobs in token_generator ) - else: - kwargs.pop("max_kv_size", None) - kwargs.pop("prompt_progress_callback", None) - token_generator = speculative_generate_step( - prompt, model, draft_model, **kwargs - ) with wired_limit(model, [generation_stream]): tic = time.perf_counter() for n, (token, logprobs, from_draft) in enumerate(token_generator): @@ -2083,6 +2279,7 @@ def main(): quantized_kv_start=args.quantized_kv_start, draft_model=draft_model, num_draft_tokens=args.num_draft_tokens, + mtp=args.mtp, ) if not args.verbose: print(response) diff --git a/mlx_lm/models/deepseek_v4.py b/mlx_lm/models/deepseek_v4.py new file mode 100644 index 000000000..16d3eb352 --- /dev/null +++ b/mlx_lm/models/deepseek_v4.py @@ -0,0 +1,1571 @@ +# Copyright © 2026 Apple Inc. / mlx-community +# +# DeepSeek-V4 (Pro / Flash) for mlx-lm. +# Architecture: Multi-head Latent Attention (num_kv_heads=1) + grouped low-rank output, +# sliding-window + compressed KV + indexer topk (sparse attention), hash-routed MoE +# with sqrtsoftplus scoring, Manifold-constrained Hyper-Connections (mHC) replacing +# residuals. Weights are native FP8 (e4m3) with 128x128 block scaling (ue8m0). +# +# Reference: deepseek-ai/DeepSeek-V4 (Apr 2026). mHC: arXiv:2512.24880. + +import math +from dataclasses import dataclass, field +from typing import Dict, List, Optional + +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 KVCache, RotatingKVCache +from .hyper_connection import HyperConnection, HyperHead +from .pipeline import PipelineMixin +from .switch_layers import SwitchGLU + + +# --------------------------------------------------------------------------- # +# Config # +# --------------------------------------------------------------------------- # + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str = "deepseek_v4" + vocab_size: int = 129280 + hidden_size: int = 4096 + num_hidden_layers: int = 43 + num_attention_heads: int = 64 + num_key_value_heads: int = 1 + + # Attention (MLA-style with single shared KV head) + q_lora_rank: int = 1024 + o_lora_rank: int = 1024 + o_groups: int = 8 + head_dim: int = 512 + qk_rope_head_dim: int = 64 + attention_bias: bool = False + sliding_window: int = 128 + compress_ratios: List[int] = field(default_factory=list) + + # Compressor / Indexer + index_n_heads: int = 64 + index_head_dim: int = 128 + index_topk: int = 512 + compress_rope_theta: float = 160000.0 + + # MoE + moe_intermediate_size: int = 2048 + n_routed_experts: int = 256 + n_shared_experts: int = 1 + num_experts_per_tok: int = 6 + num_hash_layers: int = 3 + scoring_func: str = "sqrtsoftplus" + topk_method: str = "noaux_tc" + norm_topk_prob: bool = True + routed_scaling_factor: float = 1.5 + swiglu_limit: float = 10.0 + + # Hyper-Connections + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + + # MTP (multi-token prediction) — present in checkpoint but unused at inference + num_nextn_predict_layers: int = 1 + + # RoPE / YaRN + max_position_embeddings: int = 1048576 + rope_theta: float = 10000.0 + rope_scaling: Optional[Dict] = None + rms_norm_eps: float = 1e-6 + + # Quantization (FP8 block) + quantization_config: Optional[Dict] = None + + def __post_init__(self): + # Auto-fill compress_ratios with V4 defaults if not specified, and + # validate length / values. Adapted from @eauchs c6a7828 (#1192). + 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 []) + ) + total_layers = self.num_hidden_layers + self.num_nextn_predict_layers + self.compress_ratios = list(self.compress_ratios[:total_layers]) + # MTP layers default to compress_ratio=0 (no compression) + while len(self.compress_ratios) < total_layers: + self.compress_ratios.append(0) + 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}") + + +# --------------------------------------------------------------------------- # +# Fused partial-RoPE Metal kernel # +# --------------------------------------------------------------------------- # +# +# Decode dispatch reduction: the scalar-Python rotation (slice -> reshape -> +# index x0/x1 -> 4 muls + add/sub -> stack -> reshape) issues ~5 graph ops per +# rope call, and DeepseekV4 invokes rope ~3x per attention layer (q_pe, k_pe, +# inverse on attention output) -> 129 calls/token at L=1. Collapsing the chain +# into a single Metal kernel removes ~600 dispatches/token on the decode path. +# +# Adapted from @0xClandestine's optimization PR +# (https://github.com/Blaizzy/mlx-lm/pull/13) targeting Blaizzy's V4 branch. +# We use the rope-only signature (V4Attention splits nope/rope outside the +# rope call), so the nope passthrough loop is dropped from the source. +# +# One SIMD-group per (b, h, l) work item; lane t handles the interleaved +# pair (x[2t], x[2t+1]). + +def _make_partial_rope_kernel(): + # Env-var escape hatch so benchmarks can A/B kernel ON vs OFF without + # monkey-patching: MLX_LM_DISABLE_PARTIAL_ROPE_KERNEL=1 -> falls back + # to the pure-MLX path used pre-2026-04-25. + import os + if os.environ.get("MLX_LM_DISABLE_PARTIAL_ROPE_KERNEL", "0") == "1": + return None + if mx.default_device() != mx.gpu or not mx.metal.is_available(): + return None + + source = """ + uint tid = thread_position_in_threadgroup.x; + uint gid = threadgroup_position_in_grid.x; + + constexpr int DRH = D_ROPE / 2; + int L_v = dims[0]; + int H_v = dims[1]; + uint l = gid % (uint)L_v; + uint tmp = gid / (uint)L_v; + uint h = tmp % (uint)H_v; + uint b = tmp / (uint)H_v; + + const auto xp = x + ((uint64_t)b * H_v * L_v + h * L_v + l) * D_ROPE; + auto yp = y + ((uint64_t)b * H_v * L_v + h * L_v + l) * D_ROPE; + const auto cp = cos_s + l * DRH; + const auto sp = sin_s + l * DRH; + + // Lane t handles one interleaved pair (x[2t], x[2t+1]). + if ((int)tid < DRH) { + float x0 = float(xp[2 * tid]); + float x1 = float(xp[2 * tid + 1]); + float c = float(cp[tid]); + float s = float(sp[tid]); + if (INVERSE) { + store_elem(yp[2 * tid], fma( x1, s, x0 * c)); // x0*c + x1*s + store_elem(yp[2 * tid + 1], fma(-x0, s, x1 * c)); // -x0*s + x1*c + } else { + store_elem(yp[2 * tid], fma(-x1, s, x0 * c)); // x0*c - x1*s + store_elem(yp[2 * tid + 1], fma( x0, s, x1 * c)); // x0*s + x1*c + } + } + """ + return mx.fast.metal_kernel( + name="ds4_partial_rope", + input_names=["x", "cos_s", "sin_s", "dims"], + output_names=["y"], + header="template inline void store_elem(device T& dst, float v) { dst = T(v); }", + source=source, + ) + + +_partial_rope_kernel = _make_partial_rope_kernel() + + +class DeepseekV4RoPE(nn.Module): + """DeepSeek-V4 rotary embedding. + + The reference implementation applies RoPE to the KV tensor before attention + and applies the conjugate rotation to the attention output. The generic MLX + RoPE layers do not expose an inverse path, so keep the small DeepSeek-specific + implementation here. + """ + + def __init__( + self, + dims: int, + base: float, + scaling_config: Optional[Dict] = None, + ): + super().__init__() + self.dims = dims + + 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 = math.floor(correction_dim(beta_fast)) + high = math.ceil(correction_dim(beta_slow)) + low = max(low, 0) + high = min(high, dims - 1) + if low == high: + high += 0.001 + + ramp = (mx.arange(dims // 2, dtype=mx.float32) - low) / (high - low) + smooth = 1 - mx.clip(ramp, 0, 1) + inv_freq = inv_freq / factor * (1 - smooth) + inv_freq * smooth + elif rope_type not in (None, "default", "linear"): + raise ValueError(f"Unsupported DeepSeek-V4 RoPE type {rope_type}") + + # This is derived from config, not a checkpoint parameter. + 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): + dtype = x.dtype + T = x.shape[-2] + if isinstance(offset, mx.array): + if offset.size == 1: + offset = offset.item() + else: + B = offset.shape[0] + pos = offset[:, None] + mx.arange(T, dtype=mx.float32)[None, :] + theta = pos[..., None] * self.inv_freq[None, None, :] + if inverse: + theta = -theta + # theta: [B, T, dims//2]. Reshape for x dims: [B,H,T,D] or [B,1,T,D] + target_shape = (B,) + (1,) * (x.ndim - 3) + (T, self.dims // 2) + cos = mx.cos(theta).reshape(target_shape).astype(dtype) + sin = mx.sin(theta).reshape(target_shape).astype(dtype) + rot = x[..., : self.dims].reshape(*x.shape[:-1], self.dims // 2, 2) + x0 = rot[..., 0] + x1 = rot[..., 1] + r0 = x0 * cos - x1 * sin + r1 = x0 * sin + x1 * cos + rotated = mx.stack([r0, r1], axis=-1).reshape(*x.shape[:-1], self.dims) + if self.dims < x.shape[-1]: + return mx.concatenate([rotated, x[..., self.dims:]], axis=-1) + return rotated + # Fast path: fused Metal kernel for the rope-only 4D case used by + # V4Attention. Falls through to the pure-MLX path on CPU, on Mode-B + # (x has a nope tail), or on non-4D inputs (e.g. Indexer rope). + # The kernel itself handles inverse via formula sign-flip; theta is + # always forward-direction (do NOT negate it here). + if ( + _partial_rope_kernel is not None + and x.shape[-1] == self.dims + and x.ndim == 4 + ): + B, H, L, _ = x.shape + pos = mx.arange(offset, offset + T, dtype=mx.float32) + theta = pos[:, None] * self.inv_freq[None, :] + cos = mx.cos(theta).astype(mx.float32) + sin = mx.sin(theta).astype(mx.float32) + dims_arr = mx.array([L, H], dtype=mx.int32) + return _partial_rope_kernel( + inputs=[x, cos, sin, dims_arr], + template=[("D_ROPE", self.dims), ("INVERSE", 1 if inverse else 0)], + grid=(B * H * L * 32, 1, 1), + threadgroup=(32, 1, 1), + output_shapes=[x.shape], + output_dtypes=[x.dtype], + )[0] + + pos = mx.arange(offset, offset + T, dtype=mx.float32) + theta = pos[:, None] * self.inv_freq[None, :] + if inverse: + theta = -theta + + broadcast_shape = (1,) * (x.ndim - 2) + theta.shape + cos = mx.cos(theta).reshape(broadcast_shape).astype(dtype) + sin = mx.sin(theta).reshape(broadcast_shape).astype(dtype) + + rot = x[..., : self.dims].reshape(*x.shape[:-1], self.dims // 2, 2) + x0 = rot[..., 0] + x1 = rot[..., 1] + y = mx.stack((x0 * cos - x1 * sin, x0 * sin + x1 * cos), axis=-1) + y = y.reshape(*x.shape[:-1], self.dims) + if x.shape[-1] == self.dims: + return y + return mx.concatenate([y, x[..., self.dims :]], axis=-1) + + +# --------------------------------------------------------------------------- # +# Gate (hash + score-based) # +# --------------------------------------------------------------------------- # + +# Pre-allocated scalar zero for sqrtsoftplus: avoids mx.zeros_like() allocation per call. +_SCORE_ZERO = mx.array(0.0) + + +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) + # sqrtsoftplus: sqrt(softplus(x)) — used by V4 + # Scalar broadcast avoids allocating a zeros tensor every call. + return mx.sqrt(mx.logaddexp(scores, _SCORE_ZERO)) + + +class MoEGate(nn.Module): + """Routing gate. First `num_hash_layers` layers use a deterministic hash + (token-id -> expert-id table) instead of learned score-based topk. Remaining + layers run sqrtsoftplus scoring + e_score_correction_bias + topk, with + post-softmax renormalization if score_func != 'softmax'.""" + + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.n_routed = args.n_routed_experts + self.top_k = args.num_experts_per_tok + self.hash = layer_idx < args.num_hash_layers + self.score_func = args.scoring_func + self.route_scale = args.routed_scaling_factor + self.norm_topk_prob = args.norm_topk_prob + + self.weight = mx.zeros((self.n_routed, args.hidden_size)) + # Cache transposed weight to avoid recomputing .T every forward call. + self._weight_t = None + if self.hash: + # tid2eid: [vocab, top_k] int32 — predetermined expert routing per token id + self.tid2eid = mx.zeros((args.vocab_size, self.top_k), dtype=mx.int32) + else: + self.e_score_correction_bias = mx.zeros((self.n_routed,), dtype=mx.float32) + + @property + def weight_t(self): + if self._weight_t is None: + self._weight_t = self.weight.T + return self._weight_t + + def __call__(self, x: mx.array, input_ids: Optional[mx.array] = None): + # x: [B, S, D] or [N, D] + if self.hash: + # x shape -> [B*S, D]; input_ids -> [B, S] flattened to [B*S] + flat = x.reshape(-1, x.shape[-1]) + scores = flat.astype(mx.float32) @ self.weight_t.astype(mx.float32) + scores = _score_func(scores, self.score_func) + ids = input_ids.reshape(-1) + inds = self.tid2eid[ids].astype(mx.int32) + weights = mx.take_along_axis(scores, inds, axis=-1) + # Reshape inds/weights back to match x's leading dims so SwitchGLU + # can broadcast against x: [B, S, top_k] (mirrors non-hash branch). + inds = inds.reshape(*x.shape[:-1], self.top_k) + weights = weights.reshape(*x.shape[:-1], self.top_k) + else: + scores = x.astype(mx.float32) @ self.weight_t.astype(mx.float32) + scores = _score_func(scores, self.score_func) + orig = scores + 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(orig, inds, axis=-1) + + if self.score_func != "softmax" and self.norm_topk_prob: + weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) + weights = weights * self.route_scale + return inds, weights + + +# --------------------------------------------------------------------------- # +# MoE # +# --------------------------------------------------------------------------- # + +def _swiglu_limited(gate: mx.array, up: mx.array, limit: float) -> mx.array: + if limit and limit > 0: + up = mx.clip(up, -limit, limit) + gate = mx.minimum(gate, limit) + return nn.silu(gate) * up + + +class DeepseekV4MLP(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int, swiglu_limit: float = 0.0): + super().__init__() + 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(_swiglu_limited(self.gate_proj(x), self.up_proj(x), self.swiglu_limit)) + + +class DeepseekV4MoE(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.num_experts_per_tok = args.num_experts_per_tok + self.switch_mlp = SwitchGLU( + args.hidden_size, + args.moe_intermediate_size, + args.n_routed_experts, + ) + self.gate = MoEGate(args, layer_idx) + if args.n_shared_experts: + self.shared_experts = DeepseekV4MLP( + args.hidden_size, + args.moe_intermediate_size * args.n_shared_experts, + swiglu_limit=0.0, + ) + 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, weights = self.gate(x, input_ids) + # Compute shared_experts before switch_mlp so MLX can overlap both + # on the GPU — shared_experts doesn't depend on routing results. + shared_y = self.shared_experts(x) if hasattr(self, "shared_experts") else None + y = self.switch_mlp(x, inds) + y = (y * weights[..., None]).sum(axis=-2).astype(y.dtype) + if shared_y is not None: + y = y + shared_y + if self.sharding_group is not None: + y = mx.distributed.all_sum(y, group=self.sharding_group) + return y + + +# --------------------------------------------------------------------------- # +# Attention: MLA (num_kv_heads=1) + sliding window + optional compressed KV # +# --------------------------------------------------------------------------- # + +class CompressedKVCache(KVCache): + """Cache for compressed-attention layers: sliding-window local cache + compressed KV pool. + + During prefill, the compressor produces all compressed rows at once. + During decode, tokens accumulate in a buffer; every `ratio` tokens the + buffer is compressed and the result is appended to the pool. + + Inherits from KVCache so external engines (vllm-mlx) recognize it via + isinstance checks. All state is proxied through self.local (RotatingKVCache). + """ + + def __init__(self, max_size: int = 128): + # Skip KVCache.__init__ — we proxy everything through self.local + self.local = RotatingKVCache(max_size=max_size, keep=0) + self._pool = None + self._state_kv = None + self._state_score = None + self._buf = None + self._buf_count = 0 + self._abs_pos = 0 + + @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 values(self): + return self.local.values + + @values.setter + def values(self, value): + self.local.values = value + + @property + def pool(self): + return self._pool + + def update_and_fetch(self, keys, values): + return self.local.update_and_fetch(keys, values) + + @property + def state(self): + return self.local.state + + @state.setter + def state(self, value): + self.local.state = value + + @property + def nbytes(self): + n = self.local.nbytes + if self._pool is not None: + n += self._pool.nbytes + if self._buf is not None: + n += self._buf.nbytes + return n + + @property + def meta_state(self): + return self.local.meta_state + + @meta_state.setter + def meta_state(self, value): + self.local.meta_state = value + + def is_trimmable(self): + return self.local.is_trimmable() + + def trim(self, n): + return self.local.trim(n) + + @classmethod + def merge(cls, caches): + """Merge multiple CompressedKVCaches into a single batched cache.""" + merged = cls.__new__(cls) + + # Merge local rotating caches (delegates to BatchRotatingKVCache) + merged.local = caches[0].local.merge([c.local for c in caches]) + + # Merge compressed pools: pad to max length, stack along B + pools = [c._pool for c in caches] + if all(p is None for p in pools): + merged._pool = None + else: + head_dim = next(p.shape[-1] for p in pools if p is not None) + dtype = next(p.dtype for p in pools if p is not None) + max_len = max(p.shape[1] if p is not None else 0 for p in pools) + padded = [] + for p in pools: + if p is None: + padded.append(mx.zeros((1, max_len, head_dim), dtype=dtype)) + elif p.shape[1] < max_len: + pad = mx.zeros((1, max_len - p.shape[1], head_dim), dtype=dtype) + padded.append(mx.concatenate([p, pad], axis=1)) + else: + padded.append(p) + merged._pool = mx.concatenate(padded, axis=0) + + # Merge buffers: pad to max buf_count, stack along B + bufs = [c._buf for c in caches] + buf_counts = [c._buf_count for c in caches] + if all(b is None for b in bufs): + merged._buf = None + merged._buf_count = 0 + else: + D = next(b.shape[-1] for b in bufs if b is not None) + dtype = next(b.dtype for b in bufs if b is not None) + max_bc = max(buf_counts) + padded = [] + for b, bc in zip(bufs, buf_counts): + if b is None: + padded.append(mx.zeros((1, max_bc, D), dtype=dtype)) + elif b.shape[1] < max_bc: + pad = mx.zeros((1, max_bc - b.shape[1], D), dtype=dtype) + padded.append(mx.concatenate([b, pad], axis=1)) + else: + padded.append(b) + merged._buf = mx.concatenate(padded, axis=0) + merged._buf_count = max_bc + + return merged + + def filter(self, batch_indices): + if hasattr(self.local, 'filter'): + self.local.filter(batch_indices) + if self._pool is not None: + self._pool = self._pool[batch_indices] + if self._buf is not None: + self._buf = self._buf[batch_indices] + + def extend(self, other): + if hasattr(self.local, 'extend'): + self.local.extend(other.local) + # Extend pools + if self._pool is None and other._pool is None: + pass + elif self._pool is None: + self._pool = other._pool + elif other._pool is None: + pass + else: + max_len = max(self._pool.shape[1], other._pool.shape[1]) + def pad_pool(p, target): + if p.shape[1] < target: + pad = mx.zeros((p.shape[0], target - p.shape[1], p.shape[2]), dtype=p.dtype) + return mx.concatenate([p, pad], axis=1) + return p + self._pool = mx.concatenate([pad_pool(self._pool, max_len), pad_pool(other._pool, max_len)], axis=0) + # Extend buffers + if self._buf is None and other._buf is None: + pass + elif self._buf is None: + self._buf = other._buf + self._buf_count = other._buf_count + elif other._buf is None: + pass + else: + max_bc = max(self._buf.shape[1], other._buf.shape[1]) + def pad_buf(b, target): + if b.shape[1] < target: + pad = mx.zeros((b.shape[0], target - b.shape[1], b.shape[2]), dtype=b.dtype) + return mx.concatenate([b, pad], axis=1) + return b + self._buf = mx.concatenate([pad_buf(self._buf, max_bc), pad_buf(other._buf, max_bc)], axis=0) + self._buf_count = max_bc + + def finalize(self): + if hasattr(self.local, 'finalize'): + self.local.finalize() + + def extract(self, idx): + extracted = CompressedKVCache.__new__(CompressedKVCache) + extracted.local = self.local.extract(idx) if hasattr(self.local, 'extract') else self.local + extracted._pool = self._pool[idx:idx+1] if self._pool is not None else None + extracted._buf = self._buf[idx:idx+1] if self._buf is not None else None + extracted._buf_count = self._buf_count + return extracted + + @property + def batch_size(self): + if hasattr(self.local, 'batch_size'): + return self.local.batch_size + return 1 + + def accumulate(self, x: mx.array, compressor: 'Compressor') -> Optional[mx.array]: + """Buffer tokens and compress when a full window is ready. + + Uses ds4-style rolling state: each token is immediately projected through + wkv/wgate with correct absolute-position APE, stored in rolling state + buffers, and pooled at ratio boundaries via softmax-weighted gating. + + This fixes the decode-time cache divergence bug where buffer-relative APE + indices caused wrong KV at decode S=1 (see: anerjy's report on PR #1189, + cross-validated against antirez/ds4 reference at ds4.c:6970-7034). + + Args: + x: [B, S, D] hidden states for current step(s) + compressor: the Compressor module to apply + + Returns: + The full compressed pool [B, N_compressed, head_dim], or None if empty. + """ + B, S, D = x.shape + r = compressor.ratio + + if S > 1: + ckv = compressor(x) + if ckv.shape[1] > 0: + self._pool = ckv if self._pool is None else mx.concatenate([self._pool, ckv], axis=1) + remainder = S % r + if remainder > 0: + self._buf = x[:, -remainder:] + self._buf_count = remainder + else: + self._buf = None + self._buf_count = 0 + self._abs_pos = S + self._state_kv = None + self._state_score = None + return self._pool + + coff = 2 if compressor.overlap else 1 + width = coff * compressor.head_dim + pos = self._abs_pos + pos_mod = pos % r + + xf = x.astype(mx.float32) + kv_cur = compressor.wkv(xf) + sc_cur = compressor.wgate(xf) + sc_cur = sc_cur + compressor.ape[pos_mod:pos_mod+1] + + if self._state_kv is None: + self._state_kv = mx.zeros((B, r if not compressor.overlap else 2 * r, width), dtype=mx.float32) + self._state_score = mx.full((B, r if not compressor.overlap else 2 * r, width), float('-inf'), dtype=mx.float32) + + row = (r + pos_mod) if compressor.overlap else pos_mod + self._state_kv[:, row:row+1, :] = kv_cur + self._state_score[:, row:row+1, :] = sc_cur + + self._abs_pos = pos + 1 + + if (pos + 1) % r == 0: + weights = mx.softmax(self._state_score, axis=1, precise=True) + pooled = (self._state_kv * weights).sum(axis=1, keepdims=True) + pooled = pooled[:, :, :compressor.head_dim] + ckv = compressor.norm(pooled.astype(x.dtype)) + self._pool = ckv if self._pool is None else mx.concatenate([self._pool, ckv], axis=1) + self._state_kv = None + self._state_score = None + + return self._pool + + +class Compressor(nn.Module): + """Learned gated pooling over `ratio` consecutive tokens for KV compression. + + At prefill, produces ~ seq/ratio compressed KV rows. At decode, accumulates + tokens in a state buffer and emits a compressed row every `ratio` steps. + Pure-MLX; a fused Metal kernel may replace this in a follow-up. + """ + + def __init__(self, args: ModelArgs, compress_ratio: int, head_dim: int): + super().__init__() + self.dim = args.hidden_size + self.head_dim = head_dim + self.rope_head_dim = args.qk_rope_head_dim + self.ratio = compress_ratio + self.overlap = compress_ratio == 4 + out_dim = head_dim * (2 if self.overlap else 1) + self.wkv = nn.Linear(self.dim, out_dim, bias=False) + self.wgate = nn.Linear(self.dim, out_dim, bias=False) + self.ape = mx.zeros((compress_ratio, out_dim), dtype=mx.float32) + self.norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps) + + def _overlap_transform(self, tensor: mx.array, value: float) -> mx.array: + B, S, R, _ = tensor.shape + D = self.head_dim + out = mx.full((B, S, 2 * R, D), value, dtype=tensor.dtype) + out[:, :, R:] = tensor[:, :, :, D:] + out[:, 1:, :R] = tensor[:, :-1, :, :D] + return out + + def __call__(self, x: mx.array) -> mx.array: + # Prefill-only MVP: chunk x into windows of `ratio` tokens. Ratio-4 + # layers use the overlapping layout from the reference implementation. + # Returns compressed KV: [B, S//ratio, head_dim] (bf16). + B, S, _ = x.shape + r = self.ratio + keep = (S // r) * r + if keep == 0: + return mx.zeros((B, 0, self.head_dim), dtype=x.dtype) + xf = x[:, :keep].astype(mx.float32) + kv = self.wkv(xf).reshape(B, keep // r, r, -1) + score = self.wgate(xf).reshape(B, keep // r, r, -1) + self.ape + if self.overlap: + kv = self._overlap_transform(kv, 0.0) + score = self._overlap_transform(score, float("-inf")) + weights = mx.softmax(score, axis=2, precise=True) + kv = (kv * weights).sum(axis=2) + return self.norm(kv.astype(x.dtype)) + + +class V4Attention(nn.Module): + """V4 attention block. + + Checkpoint shapes (Flash): + n_heads=64, head_dim=512, rope_head_dim=64 (nope=448) + q_lora_rank=1024, wq_a: [dim, 1024], wq_b: [1024, n_heads*head_dim] + wkv: [dim, head_dim] (single shared K=V head, MQA-style) + attn_sink: [n_heads] fp32 + wo_a: [n_heads*head_dim/n_groups, n_groups*o_lora_rank] + wo_b: [n_groups*o_lora_rank, dim] + For compress_ratio != 0: compressor.wkv/wgate/ape/norm; and if ratio==4, indexer.* + + Forward path (MVP): + - Project Q (64 heads), K=V (1 head); apply RoPE to last `rope_head_dim` dims. + - For ratio=0 layers: sliding window mask of size `sliding_window`. + - For ratio!=0 layers: append compressed KV rows to attend to (no topk filtering + yet — full compressed cache). Use attn_sink via SDPA `sinks=` argument. + - Grouped low-rank output projection: wo_a per group -> concat -> wo_b. + """ + + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.args = args + self.layer_idx = layer_idx + self.dim = args.hidden_size + self.n_heads = args.num_attention_heads + self.head_dim = args.head_dim + self.rope_head_dim = args.qk_rope_head_dim + self.nope_head_dim = args.head_dim - args.qk_rope_head_dim + self.n_groups = args.o_groups + self.q_lora_rank = args.q_lora_rank + self.o_lora_rank = args.o_lora_rank + self.window = args.sliding_window + self.eps = args.rms_norm_eps + + ratios = args.compress_ratios or [] + self.compress_ratio = ratios[layer_idx] if layer_idx < len(ratios) else 0 + + self.scale = self.head_dim ** -0.5 + + # q path + self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=args.attention_bias) + self.q_norm = nn.RMSNorm(self.q_lora_rank, eps=self.eps) + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False) + + # kv path (single shared head) + self.wkv = nn.Linear(self.dim, self.head_dim, bias=False) + self.kv_norm = nn.RMSNorm(self.head_dim, eps=self.eps) + + # attention sink (per-head learnable bias added in softmax denominator) + self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32) + + # grouped low-rank output projection + group_feat = (self.n_heads * self.head_dim) // self.n_groups + self.wo_a = nn.Linear(group_feat, self.n_groups * self.o_lora_rank, bias=False) + self.wo_b = nn.Linear(self.n_groups * self.o_lora_rank, self.dim, bias=args.attention_bias) + + # RoPE: main Q/K always rotate with rope_theta. Compressed-pool RoPE + # (when present) uses compress_rope_theta. Reference DeepSeek-V4 + # initializes both as separate instances — sharing them ties the main + # attention rotation to the wrong base on compressed layers, manifesting + # as periodic token drops in CJK (cf. Shinka-Man's report on #1192, + # fixed there in @Blaizzy/mlx-lm@b78ccb1). + self.rope = DeepseekV4RoPE(self.rope_head_dim, args.rope_theta, args.rope_scaling) + self.compress_rope = DeepseekV4RoPE( + self.rope_head_dim, args.compress_rope_theta, args.rope_scaling, + ) + + # Compressor / Indexer — present only when compress_ratio > 0 + if self.compress_ratio: + self.compressor = Compressor(args, self.compress_ratio, self.head_dim) + if self.compress_ratio == 4: + self.indexer = Indexer(args, self.compress_ratio) + + def _grouped_output_projection(self, out: mx.array) -> mx.array: + B, S = out.shape[:2] + group_feat = (self.n_heads * self.head_dim) // self.n_groups + out = out.reshape(B, S, self.n_groups, group_feat) + + if isinstance(self.wo_a, nn.QuantizedLinear): + # Batched grouped quantized matmul: collapse the per-group Python + # loop (8 dispatches) into a single mx.quantized_matmul call by + # treating the group dim as a broadcast batch dim. Adapted from + # @Blaizzy's pc/add-deepseekv4flash-model branch. + # + # Shapes: + # out (after transpose): [G, B, S, group_feat] + # weight (reshaped): [G, 1, o_lora_rank, group_feat / pack_factor] + # scales: [G, 1, o_lora_rank, group_feat / group_size] + # Single dispatch returns [G, B, S, o_lora_rank], then transpose + # back to [B, S, G, o_lora_rank] -> [B, S, G * o_lora_rank]. + out_g = out.transpose(2, 0, 1, 3) + weight = self.wo_a.weight.reshape(self.n_groups, self.o_lora_rank, -1)[:, None] + scales = self.wo_a.scales.reshape(self.n_groups, self.o_lora_rank, -1)[:, None] + biases = ( + None + if self.wo_a.biases is None + else self.wo_a.biases.reshape(self.n_groups, self.o_lora_rank, -1)[:, None] + ) + out_g = mx.quantized_matmul( + out_g, + 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_g.transpose(1, 2, 0, 3).reshape(B, S, self.n_groups * self.o_lora_rank) + if "bias" in self.wo_a: + out = out + self.wo_a.bias + return out + + wa = self.wo_a.weight.reshape(self.n_groups, self.o_lora_rank, group_feat) + out = mx.einsum("bsgd,grd->bsgr", out, wa) + out = out.reshape(B, S, self.n_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=None, cache=None): + B, S, _ = x.shape + + # --- Q (shared intermediate reused by indexer) --- + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).reshape(B, S, self.n_heads, self.head_dim).transpose(0, 2, 1, 3) + q = mx.fast.rms_norm(q, weight=None, eps=self.eps) + + # --- K = V (shared single-head) --- + kv = self.kv_norm(self.wkv(x)) + kv = kv.reshape(B, 1, S, self.head_dim) + + offset = cache.offset if cache is not None else 0 + + # Apply RoPE only to the last rope_head_dim dims + q_nope, q_pe = mx.split(q, [self.nope_head_dim], axis=-1) + k_nope, k_pe = mx.split(kv, [self.nope_head_dim], axis=-1) + q_pe = self.rope(q_pe, offset=offset) + k_pe = self.rope(k_pe, offset=offset) + q = mx.concatenate([q_nope, q_pe], axis=-1) + k = v = mx.concatenate([k_nope, k_pe], axis=-1) + + # --- Compressed sparse attention --- + compressed_k = compressed_v = None + if self.compress_ratio: + comp_cache = cache if isinstance(cache, CompressedKVCache) else None + if comp_cache is not None: + pool = comp_cache.accumulate(x, self.compressor) + elif S > 1: + pool = self.compressor(x) + pool = pool if pool.shape[1] > 0 else None + else: + pool = None + + if pool is not None: + ckv = pool + if hasattr(self, "indexer") and ckv.shape[1] > self.args.index_topk: + topk_idx = self.indexer(x, qr) + if topk_idx is not None: + idx = mx.broadcast_to( + topk_idx[:, :, None], + (B, topk_idx.shape[1], self.head_dim), + ) + ckv = mx.take_along_axis(ckv, idx, axis=1) + compressed_k = ckv[:, None, :, :] + compressed_v = compressed_k + + # Update KV cache + if cache is not None: + k, v = cache.update_and_fetch(k, v) + + # Prepend compressed KV to cached KV for sparse attention + if compressed_k is not None: + k = mx.concatenate([compressed_k, k], axis=2) + v = mx.concatenate([compressed_v, v], axis=2) + n_comp = compressed_k.shape[2] + if mask is not None: + comp_shape = list(mask.shape) + comp_shape[-1] = n_comp + comp_mask = mx.zeros(comp_shape, dtype=mask.dtype) + mask = mx.concatenate([comp_mask, mask], axis=-1) + + out = scaled_dot_product_attention( + q, + k, + v, + cache=cache, + scale=self.scale, + mask=mask, + sinks=self.attn_sink.astype(q.dtype), + ) + + out_nope, out_pe = mx.split(out, [self.nope_head_dim], axis=-1) + out_pe = self.rope(out_pe, offset=offset, inverse=True) + out = mx.concatenate([out_nope, out_pe], axis=-1) + + # Grouped low-rank projection: [B, n_heads, S, head_dim] -> [B, S, n_heads*head_dim] + out = out.transpose(0, 2, 1, 3).reshape(B, S, self.n_heads * self.head_dim) + out = self._grouped_output_projection(out) + return self.wo_b(out) + + +class Indexer(nn.Module): + """Top-k selector over compressed KV rows for ratio-4 sparse attention. + + Two-pass design: this module uses a lightweight compressor (index_head_dim, + typically 128) to score all compressed rows cheaply, then returns topk + indices used to gather from the main attention compressor's output + (head_dim, typically 512). This reduces per-layer attention from O(S/4) + to O(topk) compressed rows — 500x at 1M context with topk=512. + + Checkpoint params: + wq_b: [q_lora_rank, n_heads * index_head_dim] + weights_proj: [hidden_size, n_heads] + compressor.{wkv, wgate, ape, norm} + """ + + def __init__(self, args: ModelArgs, compress_ratio: int): + super().__init__() + self.dim = args.hidden_size + self.n_heads = args.index_n_heads + self.head_dim = args.index_head_dim + self.index_topk = args.index_topk + self.q_lora_rank = args.q_lora_rank + self.scale = args.index_head_dim ** -0.5 + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False) + self.weights_proj = nn.Linear(self.dim, self.n_heads, bias=False) + self.compressor = Compressor(args, compress_ratio, self.head_dim) + + def __call__( + self, + x: mx.array, + q_intermediate: mx.array, + ) -> Optional[mx.array]: + """Score compressed rows and return topk indices. + + Args: + x: [B, S, D] hidden state (fed to the lightweight compressor). + q_intermediate: [B, S, q_lora_rank] post wq_a+q_norm (shared with main attn). + + Returns: + topk_indices [B, topk] or None when there are too few compressed rows. + Indices are shared across heads (head-weighted scores are aggregated). + """ + B, S, _ = x.shape + + ck = self.compressor(x) + n_compressed = ck.shape[1] + if n_compressed == 0: + return None + + q = self.wq_b(q_intermediate) + q = q.reshape(B, S, self.n_heads, self.head_dim) + q = q.transpose(0, 2, 1, 3) + + scores = (q @ ck[:, None].transpose(0, 1, 3, 2)) * self.scale + + hw = mx.sigmoid(self.weights_proj(x)) + hw = hw.transpose(0, 2, 1)[..., None] + scores = scores * hw + + agg = scores.sum(axis=2).mean(axis=1) + + topk = min(self.index_topk, n_compressed) + return mx.argpartition(-agg, kth=topk - 1, axis=-1)[:, :topk] + + +# --------------------------------------------------------------------------- # +# Block # +# --------------------------------------------------------------------------- # + +class DeepseekV4Block(nn.Module): + """V4 block: mHC-wrapped (attention-norm -> attention), mHC-wrapped (moe-norm -> moe). + + The block maintains `hc_mult` parallel hidden-state copies. Each sub-layer + reduces them to 1 via hc_pre, applies its block, then expands back via hc_post. + """ + + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.attn_norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.attn = V4Attention(args, layer_idx) + self.hc_attn = HyperConnection( + args.hidden_size, args.hc_mult, + args.rms_norm_eps, args.hc_sinkhorn_iters, args.hc_eps, + ) + + self.ffn_norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.ffn = DeepseekV4MoE(args, layer_idx) + self.hc_ffn = HyperConnection( + args.hidden_size, args.hc_mult, + args.rms_norm_eps, args.hc_sinkhorn_iters, args.hc_eps, + ) + + def __call__(self, h: mx.array, mask, cache, input_ids: mx.array) -> mx.array: + # h: [B, S, hc, D] + # Attention half + residual = h + y, post, comb = self.hc_attn.hc_pre(h) + y = self.attn_norm(y) + y = self.attn(y, mask=mask, cache=cache) + h = self.hc_attn.hc_post(y, residual, post, comb) + + # FFN half + residual = h + y, post, comb = self.hc_ffn.hc_pre(h) + y = self.ffn_norm(y) + y = self.ffn(y, input_ids) + h = self.hc_ffn.hc_post(y, residual, post, comb) + return h + + +# --------------------------------------------------------------------------- # +# MTP Block (next-N-token prediction head, from Blaizzy/mlx-lm PR #15) # +# --------------------------------------------------------------------------- # + +class MTPBlock(nn.Module): + """Next-N-token prediction head. Each MTP block predicts one extra future + token by re-mixing the previous hidden state with the embedded "next" token, + then running it through a copy of the V4 transformer block + hc_head. + + Adapted from Blaizzy/mlx-lm PR #15. HyperHead signature matches our fork + (hidden_size, hc_mult, rms_norm_eps, hc_eps) instead of PR's HyperHead(config). + """ + + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + dim = args.hidden_size + self.block = DeepseekV4Block(args, layer_idx) + self.e_proj = nn.Linear(dim, dim, bias=False) + self.h_proj = nn.Linear(dim, dim, bias=False) + self.enorm = nn.RMSNorm(dim, eps=args.rms_norm_eps) + self.hnorm = nn.RMSNorm(dim, eps=args.rms_norm_eps) + self.norm = nn.RMSNorm(dim, eps=args.rms_norm_eps) + self.hc_head = HyperHead( + args.hidden_size, args.hc_mult, args.rms_norm_eps, args.hc_eps + ) + + def __call__( + self, + h: mx.array, + embed_tokens: nn.Embedding, + input_ids: mx.array, + mask: Optional[mx.array], + cache: Optional[Any], + ) -> mx.array: + e = embed_tokens(input_ids) + e = self.enorm(e) + h_norm = self.hnorm(h) + x = self.e_proj(e)[:, :, None, :] + self.h_proj(h_norm) + x = mx.contiguous(x) + x = self.block(x, mask, cache, input_ids) + return x + + +# --------------------------------------------------------------------------- # +# Model # +# --------------------------------------------------------------------------- # + +class DeepseekV4Model(nn.Module, PipelineMixin): + def __init__(self, args: ModelArgs): + super().__init__() + PipelineMixin.__init__(self) + self.args = args + self.vocab_size = args.vocab_size + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [DeepseekV4Block(args, i) for i in range(args.num_hidden_layers)] + self.start_idx = 0 + self.end_idx = len(self.layers) + self.num_layers = self.end_idx + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + # Final HC head (reduces hc copies -> 1 before lm_head) + self.hc_head = HyperHead( + args.hidden_size, args.hc_mult, args.rms_norm_eps, args.hc_eps + ) + + def __call__(self, inputs: mx.array, cache=None, return_raw_hidden: bool = False): + h = self.embed_tokens(inputs) # [B, S, D] + # Expand to hc_mult parallel copies + h = mx.broadcast_to(h[:, :, None, :], (h.shape[0], h.shape[1], self.args.hc_mult, h.shape[2])) + # Make it contiguous — broadcast_to gives a view + h = mx.contiguous(h) + + if cache is None: + cache = [None] * self.num_layers + + first_cache = cache[0] + if isinstance(first_cache, CompressedKVCache): + first_cache = first_cache.local + elif isinstance(first_cache, (list, tuple)): + first_cache = first_cache[0] + mask = create_attention_mask( + h[:, :, 0, :], + first_cache if first_cache is not None else None, + return_array=True, + ) + + pipeline_rank = self.pipeline_rank + pipeline_size = self.pipeline_size + if pipeline_rank < pipeline_size - 1: + h = mx.distributed.recv_like(h, (pipeline_rank + 1)) + + for i in range(self.num_layers): + h = self.layers[self.start_idx + i](h, mask, cache[i], inputs) + + if pipeline_rank != 0: + h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size) + last_cache = cache[-1] + if last_cache is not None: + lc = last_cache.local if isinstance(last_cache, CompressedKVCache) else last_cache + if hasattr(lc, 'keys') and lc.keys is not None: + lc.keys = mx.depends(lc.keys, h) + + if pipeline_size > 1: + h = mx.distributed.all_gather(h)[: h.shape[0]] + + # Reduce [B,S,hc,D] -> [B,S,D] then RMSNorm + out = self.norm(self.hc_head(h)) + if return_raw_hidden: + return out, h + return out + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = DeepseekV4Model(args) + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + if getattr(args, "num_nextn_predict_layers", 0) > 0: + n = args.num_hidden_layers + self.mtp = [ + MTPBlock(args, n + i) + for i in range(args.num_nextn_predict_layers) + ] + + def __call__( + self, + inputs: mx.array, + cache=None, + return_hidden: bool = False, + ): + if return_hidden: + h, h_raw = self.model(inputs, cache, return_raw_hidden=True) + return self.lm_head(h), h_raw + h = self.model(inputs, cache) + return self.lm_head(h) + + @property + def layers(self): + return self.model.layers[self.model.start_idx : self.model.end_idx] + + @property + def cast_predicate(self): + def pred(k: str): + # Keep mHC params and gate biases in fp32 + if "hc_" in k or "e_score_correction_bias" in k or "attn_sink" in k: + return False + if k.endswith(".fn") or k.endswith(".base") or k.endswith(".scale"): + return False + return True + return pred + + def make_cache(self): + caches = [] + for layer in self.layers: + if layer.attn.compress_ratio: + caches.append(CompressedKVCache(max_size=self.args.sliding_window)) + else: + caches.append(RotatingKVCache(max_size=self.args.sliding_window)) + return caches + + def make_mtp_cache(self): + if not hasattr(self, "mtp"): + return None + caches = [] + for mtp_block in self.mtp: + attn = mtp_block.block.attn + if attn.compress_ratio: + caches.append(CompressedKVCache(max_size=self.args.sliding_window)) + else: + caches.append(RotatingKVCache(max_size=self.args.sliding_window)) + return caches + + def mtp_forward( + self, + h: mx.array, + input_ids: mx.array, + cache: Optional[List[Any]] = None, + ) -> mx.array: + if cache is None: + cache = [None] * len(self.mtp) + + first_cache = cache[0] + mask_cache = ( + first_cache.local + if isinstance(first_cache, CompressedKVCache) + else first_cache + ) + mask = create_attention_mask( + h[:, :, 0, :] if h.ndim == 4 else h, + mask_cache, + window_size=self.args.sliding_window, + return_array=True, + ) + + for mtp_block, layer_cache in zip(self.mtp, cache): + h = mtp_block( + h, self.model.embed_tokens, input_ids, mask, layer_cache + ) + + out = mtp_block.hc_head(h) + out = mtp_block.norm(out) + return self.lm_head(out) + + # ------------------------------------------------------------------- # + # Weight loading # + # ------------------------------------------------------------------- # + + def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: + """Handle DeepSeek-V4 checkpoint conversion. + + Supports both raw HF checkpoints (FP8/FP4 block scales) and + pre-quantized MLX checkpoints (e.g. mlx-community 8-bit). + + Checkpoint naming (from HF): + layers.N.attn.{wq_a,wq_b,wkv,wo_a,wo_b}.{weight,scale} + layers.N.attn.{q_norm,kv_norm,attn_sink} + layers.N.attn.compressor.{wkv,wgate,ape,norm} + layers.N.attn.indexer.{wq_b,weights_proj,compressor.*} + layers.N.ffn.gate.{weight,bias,tid2eid} + layers.N.ffn.experts.E.w{1,2,3}.{weight,scale} + layers.N.ffn.shared_experts.w{1,2,3}.{weight,scale} + layers.N.{attn_norm,ffn_norm}.weight + layers.N.hc_{attn,ffn}_{fn,base,scale} + embed.weight, head.weight, hc_head_{fn,base,scale} + mtp.0.* (dropped) + + MLX-quantized naming (community 8-bit): + embed.{weight,biases,scales}, head.{weight,biases,scales} + layers.N.attn.wo_a.G.{weight,biases,scales} (per-group) + layers.N.ffn.experts.w{1,2,3}.{weight,biases,scales} (pre-stacked) + """ + n_layers = self.args.num_hidden_layers + + # 1) Keep MTP weights only when self.mtp exists; drop layers beyond n_layers + has_mtp = hasattr(self, "mtp") + has_mtp_weights = any(k.startswith("mtp.") for k in weights) + # Disable MTP module if weights are absent (e.g. quantized checkpoints) + if has_mtp and not has_mtp_weights: + del self.mtp + has_mtp = False + new_weights = {} + for k, v in weights.items(): + if k.startswith("mtp."): + if not has_mtp: + continue + new_weights[k] = v + continue + parts = k.split(".") + if len(parts) >= 2 and parts[0] == "layers": + try: + idx = int(parts[1]) + except ValueError: + new_weights[k] = v + continue + if idx >= n_layers: + continue + 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) + + # 2) FP8/FP4 block dequant: + # `X.weight` + `X.scale` -> dequantized bf16 `X.weight` + # Routed experts in Flash are FP4-packed int8; other scaled matrices + # are FP8 e4m3 with 128x128 block scales. + def _dequant_fp8_block(weight: mx.array, scale: mx.array, bs: int = 128) -> mx.array: + weight = mx.from_fp8(weight, dtype=mx.bfloat16) + scale = _scale_to_float(scale) + m, n = weight.shape + pad_b = (-m) % bs + pad_s = (-n) % bs + weight = mx.pad(weight, ((0, pad_b), (0, pad_s))) + weight = weight.reshape(((m + pad_b) // bs, bs, (n + pad_s) // bs, bs)) + weight = (weight * scale[:, None, :, None]).reshape(m + pad_b, n + pad_s) + return weight[:m, :n].astype(mx.bfloat16) + + def _dequant_fp4_block(weight: mx.array, scale: mx.array, bs: int = 32) -> mx.array: + 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), bs, axis=-1) + return (unpacked * scale).astype(mx.bfloat16) + + new = {} + for k, v in weights.items(): + if k.endswith(".scale"): + wk = k[:-len(".scale")] + ".weight" + weight = weights.get(wk) + if ( + weight is not None + and ".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[wk] = _dequant_fp4_block(weight, v) + elif weight is not None and weight.dtype in (mx.uint8,): + new[wk] = _dequant_fp8_block(weights[wk], v) + else: + new[k] = v + elif k not in new: + new[k] = v + weights = new + + # 3) Remap top-level names to our module structure + # Prefix-based remap handles both raw (.weight) and quantized + # (.weight, .biases, .scales) checkpoints. + top_prefix_remap = { + "embed.": "model.embed_tokens.", + "head.": "lm_head.", + } + top_exact_remap = { + "norm.weight": "model.norm.weight", + "hc_head_fn": "model.hc_head.fn", + "hc_head_base": "model.hc_head.base", + "hc_head_scale": "model.hc_head.scale", + } + new = {} + for k, v in weights.items(): + nk = k + for old_pfx, new_pfx in top_prefix_remap.items(): + if nk.startswith(old_pfx): + nk = new_pfx + nk[len(old_pfx):] + break + if nk in top_exact_remap: + nk = top_exact_remap[nk] + new[nk] = v + weights = new + + # 4) Remap layer-level names: layers.N.X -> model.layers.N.X + # Also remap gate.bias -> gate.e_score_correction_bias, + # hc_{attn,ffn}_{fn,base,scale} -> hc_{attn,ffn}.{fn,base,scale}, + # shared_experts.w{1,2,3} -> shared_experts.{gate,down,up}_proj + new = {} + w_remap = {"w1": "gate_proj", "w2": "down_proj", "w3": "up_proj"} + mtp_block_subs = ( + "attn.", "ffn.", "attn_norm.", "ffn_norm.", + "hc_attn_", "hc_ffn_", + ) + for k, v in weights.items(): + nk = k + # Add model. prefix for main-model layers + if nk.startswith("layers."): + nk = "model." + nk + + # MTP block: nest block-internal weights under .block. + if nk.startswith("mtp."): + parts = nk.split(".", 2) # ["mtp", "0", "rest"] + if len(parts) == 3: + rest = parts[2] + if any(rest.startswith(s) for s in mtp_block_subs): + nk = f"mtp.{parts[1]}.block.{rest}" + # HC head weights for MTP block + for param in ("fn", "base", "scale"): + if rest == f"hc_head_{param}": + nk = f"mtp.{parts[1]}.hc_head.{param}" + + # gate.bias -> gate.e_score_correction_bias + nk = nk.replace(".ffn.gate.bias", ".ffn.gate.e_score_correction_bias") + + # hc_attn_fn -> hc_attn.fn (etc.) — raw HF checkpoint underscores + for sub in ("attn", "ffn"): + for param in ("fn", "base", "scale"): + nk = nk.replace(f".hc_{sub}_{param}", f".hc_{sub}.{param}") + + # attn_hc.X -> hc_attn.X (mlx-community/DeepSeek-V4-Flash-8bit + # naming order: per-layer hyper-connections stored as _hc. + # rather than hc_.). Apply after the underscore rename so + # both naming orders converge to the model's hc_. layout. + for sub in ("attn", "ffn"): + nk = nk.replace(f".{sub}_hc.", f".hc_{sub}.") + + # shared_experts.w1 -> shared_experts.gate_proj (etc.) + for w_old, w_new in w_remap.items(): + nk = nk.replace(f".shared_experts.{w_old}.", f".shared_experts.{w_new}.") + + new[nk] = v + weights = new + + # 5) Stack expert weights: experts.E.w{1,2,3}.weight -> switch_mlp.{gate,down,up}_proj.weight + # Also handle pre-stacked experts (community quants): experts.w{1,2,3}.X -> switch_mlp.{proj}.X + expert_remap = {"w1": "gate_proj", "w2": "down_proj", "w3": "up_proj"} + for l in range(n_layers): + prefix = f"model.layers.{l}.ffn.experts" + for src, dst in expert_remap.items(): + # Case A: per-expert weights need stacking (raw HF checkpoint) + key0 = f"{prefix}.0.{src}.weight" + if key0 in weights: + stack = [weights.pop(f"{prefix}.{e}.{src}.weight") + for e in range(self.args.n_routed_experts)] + weights[f"model.layers.{l}.ffn.switch_mlp.{dst}.weight"] = mx.stack(stack) + # Case B: already-stacked (community quant) — rename experts.w1.X -> switch_mlp.gate_proj.X + for suffix in ("weight", "biases", "scales"): + old = f"{prefix}.{src}.{suffix}" + if old in weights: + weights[f"model.layers.{l}.ffn.switch_mlp.{dst}.{suffix}"] = weights.pop(old) + + # 6) Fuse split wo_a: community quants store wo_a.G.{weight,biases,scales} + # per-group; our model uses a single QuantizedLinear with grouped dequant. + n_groups = self.args.o_groups + for l in range(n_layers): + prefix = f"model.layers.{l}.attn.wo_a" + if f"{prefix}.0.weight" in weights: + for suffix in ("weight", "biases", "scales"): + parts = [] + for g in range(n_groups): + key = f"{prefix}.{g}.{suffix}" + if key in weights: + parts.append(weights.pop(key)) + if parts: + weights[f"{prefix}.{suffix}"] = mx.concatenate(parts, axis=0) + + # Stack routed expert weights for MTP layers + if has_mtp: + for mtp_idx in range(self.args.num_nextn_predict_layers): + prefix = f"mtp.{mtp_idx}.block.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"mtp.{mtp_idx}.block.ffn.switch_mlp.{dst}.weight" + ] = mx.stack(stacked) + + return weights + + # ------------------------------------------------------------------- # + # Distributed sharding # + # ------------------------------------------------------------------- # + + def shard(self, group: Optional[mx.distributed.Group] = None): + group = group or mx.distributed.init() + N = group.size() + R = group.rank() + for layer in self.model.layers: + a = layer.attn + a.wq_b = shard_linear(a.wq_b, "all-to-sharded", group=group) + a.wo_b = shard_linear(a.wo_b, "sharded-to-all", group=group) + a.n_heads //= N + # Slice attn_sink to local heads (mirrors gpt_oss.py:308-312). + # Order matters: must run AFTER `a.n_heads //= N` so the stride is + # the post-division (local) head count. + a.attn_sink = a.attn_sink[a.n_heads * R : a.n_heads * (R + 1)] + + # wo_a: shape (n_groups * o_lora_rank, group_feat). + # group_feat = n_heads * v_head_dim / n_groups. After sharding, + # n_heads //= N and n_groups //= N cancel in the ratio, so group_feat + # stays constant. Only the OUTPUT dim (n_groups axis) gets sharded — + # each rank owns n_groups//N consecutive groups. + # wo_b is "sharded-to-all" so its input = n_groups_local * o_lora_rank. + old_n_groups = a.n_groups + new_n_groups = old_n_groups // N + gs = new_n_groups * R + ge = new_n_groups * (R + 1) + if isinstance(a.wo_a, nn.QuantizedLinear): + gf = a.wo_a.weight.shape[-1] + w = a.wo_a.weight.reshape(old_n_groups, a.o_lora_rank, gf) + a.wo_a.weight = w[gs:ge].reshape(new_n_groups * a.o_lora_rank, gf) + sc_gf = a.wo_a.scales.shape[-1] + s = a.wo_a.scales.reshape(old_n_groups, a.o_lora_rank, sc_gf) + a.wo_a.scales = s[gs:ge].reshape(new_n_groups * a.o_lora_rank, sc_gf) + if getattr(a.wo_a, "biases", None) is not None: + b_gf = a.wo_a.biases.shape[-1] + b = a.wo_a.biases.reshape(old_n_groups, a.o_lora_rank, b_gf) + a.wo_a.biases = b[gs:ge].reshape(new_n_groups * a.o_lora_rank, b_gf) + else: + gf = a.wo_a.weight.shape[-1] + w = a.wo_a.weight.reshape(old_n_groups, a.o_lora_rank, gf) + a.wo_a.weight = w[gs:ge].reshape(new_n_groups * a.o_lora_rank, gf) + a.n_groups = new_n_groups + + if isinstance(layer.ffn, DeepseekV4MoE): + layer.ffn.sharding_group = group + if hasattr(layer.ffn, "shared_experts"): + shard_inplace(layer.ffn.shared_experts.gate_proj, "all-to-sharded", group=group) + shard_inplace(layer.ffn.shared_experts.down_proj, "sharded-to-all", group=group) + shard_inplace(layer.ffn.shared_experts.up_proj, "all-to-sharded", group=group) + shard_inplace(layer.ffn.switch_mlp.gate_proj, "all-to-sharded", group=group) + shard_inplace(layer.ffn.switch_mlp.down_proj, "sharded-to-all", group=group) + shard_inplace(layer.ffn.switch_mlp.up_proj, "all-to-sharded", group=group) diff --git a/mlx_lm/models/hyper_connection.py b/mlx_lm/models/hyper_connection.py new file mode 100644 index 000000000..95f3a7feb --- /dev/null +++ b/mlx_lm/models/hyper_connection.py @@ -0,0 +1,116 @@ +# Copyright © 2026 Apple Inc. + +"""Manifold-constrained Hyper-Connections (mHC) — shared building blocks. + +DeepSeek-V4 introduced multi-HyperConnection (mHC) as a residual replacement: +expand the hidden state into `hc_mult` parallel copies, mix them via a +doubly-stochastic matrix on the Birkhoff polytope, apply the block, and +recombine. This module hosts the `nn.Module` layers; the Sinkhorn projection +that produces the doubly-stochastic mixing matrix lives in +`mlx_lm.models.sinkhorn`. + +Two layers: + - HyperConnection — per-block mHC: hc_pre reduces hc_mult -> 1; block F runs; + hc_post expands 1 -> hc_mult via (post * f_out + comb @ residual). + - HyperHead — final-layer head variant: sigmoid-weighted reduction + hc_mult -> 1 with no Sinkhorn (simpler than HyperConnection). + +References: + - mHC: arXiv:2512.24880 (DeepSeek, Dec 2025) + - HC base: arXiv:2409.19606 (Sep 2024) +""" + +import mlx.core as mx +import mlx.nn as nn + +from .sinkhorn import hc_split_sinkhorn + + +class HyperConnection(nn.Module): + """Per-block mHC parameters: projects x -> (pre, post, comb) used in hc_pre/hc_post. + + Paper/ref stores the weights as: + hc_fn : [(2+hc)*hc, hc*dim] + hc_scale : [3] + hc_base : [(2+hc)*hc] + + hc_pre reduces `hc_mult` parallel hidden states to 1 via `pre`. + Block F is applied to the reduced state. hc_post expands 1 -> hc via `post` (the new + contribution) added to `comb @ residual` (where `comb` is a doubly-stochastic mix + that recombines the input `hc_mult` copies to stay on the Birkhoff manifold). + """ + + def __init__(self, dim: int, hc_mult: int, norm_eps: float, sinkhorn_iters: int, hc_eps: float): + super().__init__() + self.dim = dim + self.hc_mult = hc_mult + self.norm_eps = norm_eps + self.sinkhorn_iters = sinkhorn_iters + self.hc_eps = hc_eps + mix_hc = (2 + hc_mult) * hc_mult + hc_dim = hc_mult * dim + # All mHC params are fp32 in the checkpoint. + self.fn = mx.zeros((mix_hc, hc_dim), dtype=mx.float32) + self.base = mx.zeros((mix_hc,), dtype=mx.float32) + self.scale = mx.zeros((3,), dtype=mx.float32) + self._fn_t = None # lazy transpose cache (avoids 86 .T calls/token) + + def hc_pre(self, x: mx.array): + B, S, hc, D = x.shape + dtype = x.dtype + xf = x.reshape(B, S, hc * D).astype(mx.float32) + xf_norm = mx.fast.rms_norm(xf, weight=None, eps=self.norm_eps) + if self._fn_t is None: + self._fn_t = self.fn.T + mixes = (xf_norm @ self._fn_t).reshape(B * S, -1) + pre, post, comb = hc_split_sinkhorn( + mixes, self.scale, self.base, hc, self.sinkhorn_iters, self.hc_eps + ) + pre = pre.reshape(B, S, hc) + post = post.reshape(B, S, hc) + comb = comb.reshape(B, S, hc, hc) + y = (pre[..., None] * x.astype(mx.float32)).sum(axis=2) + return y.astype(dtype), post, comb + + def hc_post(self, f_out: mx.array, residual: mx.array, post: mx.array, comb: mx.array): + # f_out [B,S,D] (block output, reduced state) + # residual [B,S,hc,D] (input to hc_pre) + # post [B,S,hc] + # comb [B,S,hc,hc] + # returns [B,S,hc,D] + dtype = f_out.dtype + # post.unsqueeze(-1) * f_out.unsqueeze(-2) -> [B,S,hc,D] + term_new = post[..., None] * f_out[:, :, None, :].astype(mx.float32) + # comb @ residual: [B,S,hc,hc] @ [B,S,hc,D] -> [B,S,hc,D] + term_res = comb.astype(mx.float32) @ residual.astype(mx.float32) + y = term_new + term_res + return y.astype(dtype) + + +class HyperHead(nn.Module): + """Final (head) mHC projection: reduces [B,S,hc,D] -> [B,S,D] via sigmoid-weighted sum. + No Sinkhorn here — this is the simpler head variant from `ParallelHead.hc_head`. + """ + + def __init__(self, dim: int, hc_mult: int, norm_eps: float, hc_eps: float): + super().__init__() + self.dim = dim + self.hc_mult = hc_mult + self.norm_eps = norm_eps + self.hc_eps = hc_eps + self.fn = mx.zeros((hc_mult, hc_mult * dim), dtype=mx.float32) + self.base = mx.zeros((hc_mult,), dtype=mx.float32) + self.scale = mx.zeros((1,), dtype=mx.float32) + self._fn_t = None # lazy transpose cache + + def __call__(self, x: mx.array): + B, S, hc, D = x.shape + dtype = x.dtype + xf = x.reshape(B, S, hc * D).astype(mx.float32) + inv = mx.rsqrt((xf * xf).mean(axis=-1, keepdims=True) + self.norm_eps) + if self._fn_t is None: + self._fn_t = self.fn.T + mixes = (xf @ self._fn_t) * inv # [B,S,hc] + pre = mx.sigmoid(mixes * self.scale[0] + self.base) + self.hc_eps + y = (pre[..., None] * x.astype(mx.float32)).sum(axis=2) + return y.astype(dtype) diff --git a/mlx_lm/models/sinkhorn.py b/mlx_lm/models/sinkhorn.py new file mode 100644 index 000000000..017238954 --- /dev/null +++ b/mlx_lm/models/sinkhorn.py @@ -0,0 +1,272 @@ +# Copyright © 2025 Apple Inc. + +"""Sinkhorn normalization for doubly-stochastic matrices (HyperConnection mHC). + +Shared module for models that use learned mixing weights on the Birkhoff +polytope (e.g., DeepSeek-V4's multi-HyperConnection). Provides both a +fused Metal kernel (register-resident, one thread per token) and a pure-MLX +reference path. +""" + +import mlx.core as mx + +_SINKHORN_KERNELS: dict = {} + + +def _make_sinkhorn_kernel(hc: int, iters: int, eps: float): + """Build a fully-unrolled, register-resident Sinkhorn Metal kernel. + + Each thread owns one token's [hc, hc] matrix in registers (hc^2 floats). + No threadgroup memory, no atomics, no global memory traffic between iters + — kernel reads input once and writes output once. + """ + key = (hc, iters) + if key in _SINKHORN_KERNELS: + return _SINKHORN_KERNELS[key] + + n_elem = hc * hc + eps_lit = f"{eps:.8e}f" + + def row_softmax(): + out = [] + for r in range(hc): + base = r * hc + out.append(f" {{ float mx = m[{base}];") + for c in range(1, hc): + out.append(f" mx = metal::max(mx, m[{base + c}]);") + out.append(f" float s = 0.0f;") + for c in range(hc): + out.append(f" m[{base + c}] = metal::exp(m[{base + c}] - mx); s += m[{base + c}];") + out.append(f" float inv = 1.0f / s;") + for c in range(hc): + out.append(f" m[{base + c}] *= inv;") + out.append(" }") + return "\n".join(out) + + def row_norm(): + out = [] + for r in range(hc): + base = r * hc + terms = " + ".join(f"m[{base + c}]" for c in range(hc)) + out.append(f" {{ float inv = 1.0f / ({terms} + {eps_lit});") + for c in range(hc): + out.append(f" m[{base + c}] *= inv;") + out.append(" }") + return "\n".join(out) + + def col_norm(): + out = [] + for c in range(hc): + terms = " + ".join(f"m[{r * hc + c}]" for r in range(hc)) + out.append(f" {{ float inv = 1.0f / ({terms} + {eps_lit});") + for r in range(hc): + out.append(f" m[{r * hc + c}] *= inv;") + out.append(" }") + return "\n".join(out) + + def add_eps(): + return "\n".join(f" m[{i}] += {eps_lit};" for i in range(n_elem)) + + iter_body = "\n".join([row_norm(), col_norm()]) + inner_iters = "\n".join([iter_body] * (iters - 1)) + + source = f""" + uint n = thread_position_in_grid.x; + if (n >= n_tokens[0]) return; + + uint base = n * {n_elem}; + float m[{n_elem}]; +{chr(10).join(f" m[{i}] = comb_log[base + {i}];" for i in range(n_elem))} + + // Row softmax +{row_softmax()} + + // Add eps +{add_eps()} + + // Initial column normalization (matches reference: cols first after softmax) +{col_norm()} + + // Remaining (iters - 1) rounds of (row_norm, col_norm) +{inner_iters} + +{chr(10).join(f" comb[base + {i}] = m[{i}];" for i in range(n_elem))} + """ + + kernel = mx.fast.metal_kernel( + name=f"sinkhorn_hc{hc}_it{iters}", + input_names=["comb_log", "n_tokens"], + output_names=["comb"], + source=source, + ) + _SINKHORN_KERNELS[key] = kernel + return kernel + + +def _make_hc_split_sinkhorn_fused_kernel(): + """All-fused HC split + sinkhorn kernel (hc_mult=4 only). + + Combines pre-sigmoid, post-sigmoid, comb scaling, softmax, and Sinkhorn + iterations into a single Metal dispatch. Replaces 4 dispatches with 1 on + the hc_mult=4 fast path. Source ported from Blaizzy/mlx-lm#1192. + """ + 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]; + + { + 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; + } + { + 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)); + } + + 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); + + 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); + + 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; + + float4 col = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col; r1 *= col; r2 *= col; r3 *= col; + + 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; + } + + *(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_fused", + input_names=["mixes", "scale", "base", "eps"], + output_names=["pre", "post", "comb"], + source=source, + ) + + +_hc_split_sinkhorn_fused_kernel = _make_hc_split_sinkhorn_fused_kernel() + + +def hc_split_sinkhorn( + mixes: mx.array, # [B*S, (2+hc)*hc] fp32 + hc_scale: mx.array, # [3] fp32 + hc_base: mx.array, # [(2+hc)*hc] fp32 + hc_mult: int = 4, + sinkhorn_iters: int = 20, + eps: float = 1e-6, +): + """Split `mixes` into (pre, post, comb_logits); Sinkhorn-normalize comb to doubly stochastic. + + Returns: + pre [N, hc] — sigmoid(mixes[:,:hc] * s0 + base[:hc]) + eps + post [N, hc] — 2*sigmoid(mixes[:,hc:2hc] * s1 + base[hc:2hc]) + comb [N, hc, hc] — Sinkhorn-normalized (rows & cols ~= 1) from the last hc*hc logits. + """ + # Fast path: all-fused single-dispatch kernel for hc_mult=4 (DeepSeek-V4 default). + # Combines pre/post sigmoid, comb scaling+softmax, and Sinkhorn iters into 1 GPU dispatch + # (was 4 dispatches: pre sigmoid, post sigmoid, comb scaling, sinkhorn kernel). + if ( + hc_mult == 4 + and _hc_split_sinkhorn_fused_kernel is not None + and mx.metal.is_available() + and mixes.size > 0 + ): + n_rows = mixes.size // ((2 + hc_mult) * hc_mult) + eps_arr = mx.array([eps], dtype=mx.float32) + return _hc_split_sinkhorn_fused_kernel( + inputs=[mixes, hc_scale, hc_base, eps_arr], + 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], + ) + + # Fallback: original split path (general hc_mult, or no Metal). + n = mixes.shape[0] + mix = mixes + s0, s1, s2 = hc_scale[0], hc_scale[1], hc_scale[2] + + pre_log = mix[:, :hc_mult] * s0 + hc_base[:hc_mult] + post_log = mix[:, hc_mult:2 * hc_mult] * s1 + hc_base[hc_mult:2 * hc_mult] + comb_log = ( + mix[:, 2 * hc_mult:].reshape(n, hc_mult, hc_mult) * s2 + + hc_base[2 * hc_mult:].reshape(hc_mult, hc_mult) + ) + + pre = mx.sigmoid(pre_log) + eps + post = 2 * mx.sigmoid(post_log) + + use_kernel = ( + hc_mult <= 8 + and mx.metal.is_available() + and comb_log.size > 0 + ) + if use_kernel: + kernel = _make_sinkhorn_kernel(hc_mult, sinkhorn_iters, eps) + flat = comb_log.reshape(n, hc_mult * hc_mult).astype(mx.float32) + n_tokens = mx.array([n], dtype=mx.uint32) + (comb_flat,) = kernel( + inputs=[flat, n_tokens], + output_shapes=[(n, hc_mult * hc_mult)], + output_dtypes=[mx.float32], + grid=(n, 1, 1), + threadgroup=(min(n, 256) or 1, 1, 1), + ) + comb = comb_flat.reshape(n, hc_mult, hc_mult) + else: + comb = mx.softmax(comb_log, axis=-1, precise=True) + eps + col_sum = comb.sum(axis=1, keepdims=True) + eps + comb = comb / col_sum + for _ in range(sinkhorn_iters - 1): + row_sum = comb.sum(axis=2, keepdims=True) + eps + comb = comb / row_sum + col_sum = comb.sum(axis=1, keepdims=True) + eps + comb = comb / col_sum + + return pre, post, comb diff --git a/mlx_lm/server.py b/mlx_lm/server.py index ce8d95817..5084cd7bc 100644 --- a/mlx_lm/server.py +++ b/mlx_lm/server.py @@ -367,10 +367,17 @@ def _load(self, model_path, adapter_path=None, draft_model_path=None): "Speculative decoding may not work as expected." ) - # Compute batchability + # Compute batchability. + # NOTE: DeepSeek-V4's hybrid cache (CompressedKVCache + RotatingKVCache) + # corrupts generation under BatchGenerator even though both implement + # `merge`. Force non-batch when CompressedKVCache is present. is_batchable = draft_model is None + _prompt_cache_probe = make_prompt_cache(model) is_batchable = is_batchable and all( - hasattr(c, "merge") for c in make_prompt_cache(model) + hasattr(c, "merge") for c in _prompt_cache_probe + ) + is_batchable = is_batchable and not any( + type(c).__name__ == "CompressedKVCache" for c in _prompt_cache_probe ) # Update the member variables @@ -983,6 +990,7 @@ def progress(tokens_processed, tokens_total): prompt_cache=cache, draft_model=draft_model, num_draft_tokens=args.num_draft_tokens, + mtp=getattr(self.cli_args, "mtp", False), prompt_progress_callback=progress, prefill_step_size=self.cli_args.prefill_step_size, ): @@ -1790,6 +1798,12 @@ def main(): help="Number of tokens to draft when using speculative decoding.", default=3, ) + parser.add_argument( + "--mtp", + action="store_true", + help="Use native Multi-Token Prediction for speculative decoding " + "(requires a model with an MTP head, e.g. DeepSeek-V4).", + ) parser.add_argument( "--trust-remote-code", action="store_true", diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py index c7e50fbe7..f4ba58706 100644 --- a/mlx_lm/tokenizer_utils.py +++ b/mlx_lm/tokenizer_utils.py @@ -5,7 +5,7 @@ from json import JSONDecodeError from typing import Any, Dict, List, Optional -from transformers import AutoTokenizer, PreTrainedTokenizerFast +from transformers import AutoTokenizer, PreTrainedConfig, PreTrainedTokenizerFast class StreamingDetokenizer: @@ -611,9 +611,30 @@ def load( tokenizer_config_file = model_path / "tokenizer_config.json" chat_template = None - tokenizer = AutoTokenizer.from_pretrained( - model_path, **(tokenizer_config_extra or {}) - ) + tokenizer_config_extra = tokenizer_config_extra or {} + try: + tokenizer = AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra) + except (AttributeError, ValueError) as e: + message = str(e) + if ( + "config" in tokenizer_config_extra + or ( + "deepseek_v4" not in message + and "max_position_embeddings" not in message + ) + ): + raise + warnings.warn( + "Falling back to generic tokenizer config because Transformers does " + f"not recognize this model config: {e}", + RuntimeWarning, + stacklevel=2, + ) + tokenizer = AutoTokenizer.from_pretrained( + model_path, + config=PreTrainedConfig(), + **tokenizer_config_extra, + ) tokenizer_config = tokenizer.init_kwargs diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index ef3d266b9..952544a37 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -8,6 +8,8 @@ import os import resource import shutil +import struct +import warnings from pathlib import Path from textwrap import dedent from typing import ( @@ -279,6 +281,64 @@ def load_config(model_path: Path) -> dict: return config +# TODO(mlx#3448): drop this helper once https://github.com/ml-explore/mlx/pull/3448 +# lands and MLX's safetensors loader recognizes F8_E8M0 natively. +def _reinterpret_safetensor_e8m0_scales_as_uint8(path: str) -> bool: + """Rewrite safetensors E8M0 scale metadata to U8 in-place. + + DeepSeek-V4 stores FP8/FP4 block scales as float8_e8m0fnu. The payload is + one byte per element, and the model sanitizer decodes those exponent bytes. + MLX currently rejects the safetensors dtype before the sanitizer can run, so + reinterpret the header as uint8 while leaving tensor bytes untouched. + """ + with open(path, "r+b") as f: + header_len = struct.unpack(" header_len: + raise RuntimeError( + f"Cannot reinterpret F8_E8M0 safetensors header in {path}: " + "rewritten header is larger than original header." + ) + + f.seek(8) + f.write(new_header) + f.write(b" " * (header_len - len(new_header))) + + return True + + +def _load_safetensors(path: str, *, allow_e8m0_uint8: bool = False) -> dict: + try: + return mx.load(path) + except RuntimeError as e: + if "F8_E8M0" not in str(e) or not allow_e8m0_uint8: + raise + + if _reinterpret_safetensor_e8m0_scales_as_uint8(path): + warnings.warn( + f"Reinterpreted F8_E8M0 scale metadata as uint8 in {path}. " + "Tensor bytes were not changed.", + RuntimeWarning, + stacklevel=2, + ) + return mx.load(path) + + def load_model( model_path: Path, lazy: bool = False, @@ -319,8 +379,9 @@ def load_model( raise FileNotFoundError(f"No safetensors found in {model_path}") weights = {} + allow_e8m0_uint8 = config.get("model_type") == "deepseek_v4" for wf in weight_files: - weights.update(mx.load(wf)) + weights.update(_load_safetensors(wf, allow_e8m0_uint8=allow_e8m0_uint8)) 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..1083e5381 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,6 +1,7 @@ # Copyright © 2024 Apple Inc. import copy import importlib +import math import unittest import mlx.core as mx @@ -1422,6 +1423,477 @@ def test_deepseek_v3(self): model, args.model_type, args.vocab_size, args.num_hidden_layers ) + def test_deepseek_v4_rope_inverse(self): + from mlx_lm.models.deepseek_v4 import DeepseekV4RoPE + + scaling = { + "type": "yarn", + "factor": 16, + "original_max_position_embeddings": 65536, + "beta_fast": 32, + "beta_slow": 1, + } + rope = DeepseekV4RoPE(8, 160000, scaling) + x = mx.random.uniform(shape=(1, 2, 4, 8)) + + y = rope(x, offset=3) + z = rope(y, offset=3, inverse=True) + self.assertTrue(mx.allclose(x, z, rtol=1e-5, atol=1e-5)) + + inv_freq = 1.0 / (160000 ** (mx.arange(0, 8, 2, dtype=mx.float32) / 8)) + + def correction_dim(num_rotations): + return ( + 8 + * math.log(65536 / (num_rotations * 2 * math.pi)) + / (2 * math.log(160000)) + ) + + low = max(math.floor(correction_dim(32)), 0) + high = min(math.ceil(correction_dim(1)), 7) + if low == high: + high += 0.001 + ramp = (mx.arange(4, dtype=mx.float32) - low) / (high - low) + smooth = 1 - mx.clip(ramp, 0, 1) + inv_freq = inv_freq / 16 * (1 - smooth) + inv_freq * smooth + + theta = mx.arange(3, 7, dtype=mx.float32)[:, None] * inv_freq[None, :] + cos = mx.cos(theta).reshape(1, 1, 4, 4) + sin = mx.sin(theta).reshape(1, 1, 4, 4) + rot = x.reshape(1, 2, 4, 4, 2) + expected = mx.stack( + ( + rot[..., 0] * cos - rot[..., 1] * sin, + rot[..., 0] * sin + rot[..., 1] * cos, + ), + axis=-1, + ).reshape(1, 2, 4, 8) + 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=1024, + hidden_size=128, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=1, + q_lora_rank=32, + o_lora_rank=16, + o_groups=2, + head_dim=32, + qk_rope_head_dim=8, + sliding_window=16, + compress_ratios=[0, 0, 4, 0], + index_n_heads=4, + index_head_dim=16, + index_topk=8, + moe_intermediate_size=32, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + num_hash_layers=1, + hc_mult=2, + hc_sinkhorn_iters=2, + max_position_embeddings=256, + rope_scaling={ + "beta_fast": 32, + "beta_slow": 1, + "factor": 2, + "original_max_position_embeddings": 128, + "type": "yarn", + }, + ) + model = deepseek_v4.Model(args) + self.assertEqual(len(model.layers), args.num_hidden_layers) + self.assertEqual(model.model_type, args.model_type) + parameter_names = {name for name, _ in tree_flatten(model.parameters())} + self.assertNotIn("model.layers.0.attn.rope.inv_freq", parameter_names) + self.assertEqual( + model.layers[2].attn.compressor.wkv.weight.shape, + (2 * args.head_dim, args.hidden_size), + ) + self.assertEqual( + model.layers[2].attn.indexer.compressor.wkv.weight.shape, + (2 * args.index_head_dim, args.hidden_size), + ) + + for dtype in [mx.float32, mx.float16, mx.bfloat16]: + model.update( + tree_map( + lambda p: p.astype(dtype) + if mx.issubdtype(p.dtype, mx.floating) + else p, + model.parameters(), + ) + ) + for layer in model.model.layers: + layer.attn.attn_sink = layer.attn.attn_sink.astype(mx.float32) + + inputs = mx.array([[0, 1, 2, 3, 4]], dtype=mx.int32) + outputs = model(inputs) + self.assertEqual(outputs.shape, (1, 5, args.vocab_size)) + self.assertEqual(outputs.dtype, dtype) + + cache = model.make_cache() + self.assertIsInstance(cache[0], RotatingKVCache) + self.assertIsInstance(cache[2], deepseek_v4.CompressedKVCache) + outputs = model(inputs[:, :3], cache=cache) + self.assertEqual(outputs.shape, (1, 3, args.vocab_size)) + self.assertEqual(outputs.dtype, dtype) + outputs = model(inputs[:, 3:4], cache=cache) + self.assertEqual(outputs.shape, (1, 1, args.vocab_size)) + self.assertEqual(outputs.dtype, dtype) + + def test_mixed_quant_preserves_deepseek_v4_attention_paths(self): + from mlx_lm.convert import mixed_quant_predicate_builder + 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, + 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) + modules = dict(model.named_modules()) + predicate = mixed_quant_predicate_builder("mixed_3_6", model, group_size=32) + + high = {"group_size": 32, "bits": 6, "mode": "affine"} + low = {"group_size": 32, "bits": 3, "mode": "affine"} + for path in [ + "model.layers.0.attn.wq_a", + "model.layers.0.attn.wq_b", + "model.layers.0.attn.wkv", + "model.layers.0.attn.wo_a", + "model.layers.0.attn.wo_b", + "model.layers.2.attn.compressor.wkv", + "model.layers.2.attn.indexer.wq_b", + "model.layers.0.ffn.shared_experts.down_proj", + "model.embed_tokens", + "lm_head", + ]: + self.assertEqual(predicate(path, modules[path]), high) + + self.assertEqual( + predicate( + "model.layers.0.ffn.switch_mlp.gate_proj", + modules["model.layers.0.ffn.switch_mlp.gate_proj"], + ), + low, + ) + + def test_deepseek_v4_quantized_grouped_output_projection(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=1, + num_attention_heads=4, + 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], + 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, + ) + attn = deepseek_v4.V4Attention(args, layer_idx=0) + attn.wo_a = nn.QuantizedLinear.from_linear( + attn.wo_a, + group_size=32, + bits=6, + mode="affine", + ) + + out = mx.random.uniform(shape=(1, 3, args.num_attention_heads * args.head_dim)) + y = attn._grouped_output_projection(out) + mx.eval(y) + self.assertEqual(y.shape, (1, 3, args.o_groups * args.o_lora_rank)) + + def test_deepseek_v4_indexer_topk(self): + from mlx_lm.models import deepseek_v4 + + args = deepseek_v4.ModelArgs( + model_type="deepseek_v4", + vocab_size=1024, + hidden_size=128, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=1, + q_lora_rank=32, + o_lora_rank=16, + o_groups=2, + head_dim=32, + qk_rope_head_dim=8, + sliding_window=16, + compress_ratios=[0, 0, 4, 0], + index_n_heads=4, + index_head_dim=16, + index_topk=4, + moe_intermediate_size=32, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + num_hash_layers=1, + hc_mult=2, + hc_sinkhorn_iters=2, + max_position_embeddings=256, + ) + model = deepseek_v4.Model(args) + + # Indexer should be on ratio-4 layer (layer 2) + self.assertTrue(hasattr(model.layers[2].attn, "indexer")) + indexer = model.layers[2].attn.indexer + self.assertEqual(indexer.index_topk, 4) + self.assertEqual(indexer.n_heads, 4) + self.assertEqual(indexer.head_dim, 16) + + # Test indexer forward: 32 tokens gives 8 compressed rows (32/4), + # indexer should select topk=4 of them + B, S, D = 1, 32, args.hidden_size + x = mx.random.normal((B, S, D)) + q_inter = mx.random.normal((B, S, args.q_lora_rank)) + topk_idx = indexer(x, q_inter) + mx.eval(topk_idx) + self.assertIsNotNone(topk_idx) + self.assertEqual(topk_idx.shape, (B, 4)) + # All indices must be valid (< n_compressed = 32/4 = 8) + self.assertTrue((topk_idx < 8).all().item()) + self.assertTrue((topk_idx >= 0).all().item()) + + # Full model forward should work with enough tokens to trigger indexer + inputs = mx.array([list(range(32))], dtype=mx.int32) + outputs = model(inputs) + mx.eval(outputs) + self.assertEqual(outputs.shape, (1, 32, args.vocab_size)) + + # Prefill + decode should also work + cache = model.make_cache() + outputs = model(inputs[:, :24], cache=cache) + mx.eval(outputs) + outputs = model(inputs[:, 24:25], cache=cache) + mx.eval(outputs) + self.assertEqual(outputs.shape, (1, 1, args.vocab_size)) + + 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, + ) + model = deepseek_v4.Model(args) + + packed = mx.array( + [ + [0x21] * 16, + [0xFE] * 16, + ], + dtype=mx.int8, + ) + weights = { + "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), + } + + converted = model.sanitize(weights) + 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_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, + ) + 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), + "norm.weight": mx.ones((32,), dtype=mx.float32), + } + ) + key = "model.layers.0.attn.wkv.weight" + self.assertIn(key, converted) + self.assertIn("model.norm.weight", 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_loads_e8m0_scales_as_uint8(self): + try: + 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_safetensors + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "model.safetensors" + save_file( + { + "weight": torch.tensor([1.0, -2.0], dtype=torch.float32).to( + torch.float8_e4m3fn + ), + "scale": torch.tensor([[1.0, 2.0]], dtype=torch.float32).to( + torch.float8_e8m0fnu + ), + }, + str(path), + ) + + with self.assertRaisesRegex(RuntimeError, "F8_E8M0"): + mx.load(str(path)) + + loaded = _load_safetensors(str(path), allow_e8m0_uint8=True) + self.assertEqual(loaded["scale"].dtype, mx.uint8) + self.assertEqual(loaded["weight"].dtype, mx.uint8) + self.assertTrue( + mx.array_equal( + loaded["scale"], + mx.array([[127, 128]], dtype=mx.uint8), + ) + ) + self.assertTrue( + mx.allclose( + mx.from_fp8(loaded["weight"], dtype=mx.float32), + mx.array([1.0, -2.0], dtype=mx.float32), + ) + ) + + def test_deepseek_v4_sanitize_renames_attn_hc_order(self): + """Community quants (mlx-community/DeepSeek-V4-Flash-8bit) store the + per-layer hyper-connections as `attn_hc.X` / `ffn_hc.X` rather than the + `hc_attn.X` / `hc_ffn.X` order used by the raw HF checkpoint and our + model. The sanitizer must converge both to the model's order.""" + 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, + ) + model = deepseek_v4.Model(args) + + weights = { + # community-quant order: _hc. + "layers.0.attn_hc.fn": mx.zeros((4,), dtype=mx.float32), + "layers.0.attn_hc.base": mx.zeros((6,), dtype=mx.float32), + "layers.0.attn_hc.scale": mx.ones((3,), dtype=mx.float32), + "layers.0.ffn_hc.fn": mx.zeros((4,), dtype=mx.float32), + "layers.0.ffn_hc.base": mx.zeros((6,), dtype=mx.float32), + "layers.0.ffn_hc.scale": mx.ones((3,), dtype=mx.float32), + } + + converted = model.sanitize(weights) + + # Verify every input key has been renamed to the model's order. + for sub in ("attn", "ffn"): + for param in ("fn", "base", "scale"): + self.assertIn(f"model.layers.0.hc_{sub}.{param}", converted) + self.assertNotIn(f"model.layers.0.{sub}_hc.{param}", converted) + def test_gemma2(self): from mlx_lm.models import gemma2 diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index 54906af1c..e18be957d 100644 --- a/tests/test_tokenizers.py +++ b/tests/test_tokenizers.py @@ -2,13 +2,17 @@ import unittest from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch from huggingface_hub import snapshot_download +from transformers import PreTrainedConfig from mlx_lm.tokenizer_utils import ( BPEStreamingDetokenizer, NaiveStreamingDetokenizer, SPMStreamingDetokenizer, + load as load_tokenizer_impl, ) from mlx_lm.utils import load_tokenizer @@ -109,6 +113,40 @@ def test_thinking(self): self.assertIsNone(tokenizer.think_start_id) self.assertIsNone(tokenizer.think_end_id) + def test_unknown_model_config_tokenizer_fallback(self): + class MockTokenizer: + eos_token_id = 1 + chat_template = None + init_kwargs = {} + + def get_vocab(self): + return {} + + calls = [] + + def from_pretrained(*args, **kwargs): + calls.append(kwargs) + if len(calls) == 1: + raise AttributeError( + "'PreTrainedConfig' object has no attribute " + "'max_position_embeddings'" + ) + return MockTokenizer() + + with TemporaryDirectory() as tmpdir: + tokenizer_json = Path(tmpdir) / "tokenizer.json" + tokenizer_json.write_text("{}", encoding="utf-8") + + with patch( + "mlx_lm.tokenizer_utils.AutoTokenizer.from_pretrained", + side_effect=from_pretrained, + ): + tokenizer = load_tokenizer_impl(Path(tmpdir)) + + self.assertEqual(tokenizer.eos_token_id, 1) + self.assertEqual(len(calls), 2) + self.assertIsInstance(calls[1]["config"], PreTrainedConfig) + if __name__ == "__main__": unittest.main()