feat(backend): update mlx-lm version to 0.28.4 (#322)

Author: yuhao-zh

pyproject.toml

@@ -43,19 +43,19 @@ parallax = "parallax.cli:main"
 
 mac = [
   "torch==2.8.0",
-  "mlx-lm==0.28.0",
+  "mlx-lm==0.28.4",
   "mlx==0.30.0",
 ]
 
 gpu = [
   "sglang[all]==0.5.5",
-  "mlx-lm==0.28.0",
+  "mlx-lm==0.28.4",
   "mlx[cpu]==0.30.0",
 ]
 
 vllm = [
   "vllm==0.11.0",
-  "mlx-lm==0.28.0",
+  "mlx-lm==0.28.4",
   "mlx[cpu]==0.30.0",
 ]
 

src/parallax/models/deepseek_v32.py

@@ -1,81 +1,18 @@
 # Copyright © 2025 Apple Inc.
-import math
-from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 import mlx.core as mx
-import mlx.nn as nn
-from mlx_lm.models.base import BaseModelArgs, scaled_dot_product_attention
-from mlx_lm.models.rope_utils import initialize_rope
-from mlx_lm.models.switch_layers import SwitchGLU
+from mlx_lm.models.base import scaled_dot_product_attention
+from mlx_lm.models.deepseek_v32 import DeepseekV32Attention as MLXDeepseekV32Attention
+from mlx_lm.models.deepseek_v32 import DeepseekV32DecoderLayer as MLXDeepseekV32Block
+from mlx_lm.models.deepseek_v32 import Indexer as MLXDeepseekV32Indexer
+from mlx_lm.models.deepseek_v32 import ModelArgs
 
 from parallax.metal.indexer.kernel import q_dot_k, store_indexer_cache
 from parallax.metal.paged_attention.kernel import paged_attention, reshape_and_cache
 
 
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str = "deepseek_v32"
-    vocab_size: int = 102400
-    hidden_size: int = 4096
-    index_n_heads: int = 64
-    index_head_dim: int = 128
-    index_topk: int = 2048
-    intermediate_size: int = 11008
-    moe_intermediate_size: int = 1407
-    num_hidden_layers: int = 30
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 32
-    n_shared_experts: Optional[int] = None
-    n_routed_experts: Optional[int] = None
-    routed_scaling_factor: float = 1.0
-    kv_lora_rank: int = 512
-    q_lora_rank: int = 1536
-    qk_rope_head_dim: int = 64
-    v_head_dim: int = 128
-    qk_nope_head_dim: int = 128
-    topk_method: str = "noaux_tc"
-    scoring_func: str = "sigmoid"
-    norm_topk_prob: bool = True
-    n_group: int = 1
-    topk_group: int = 1
-    num_experts_per_tok: int = 1
-    moe_layer_freq: int = 1
-    first_k_dense_replace: int = 0
-    max_position_embeddings: int = 2048
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 10000.0
-    rope_scaling: Dict = None
-    attention_bias: bool = False
-
-
-class Indexer(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.dim = args.hidden_size
-        self.n_heads = args.index_n_heads
-        self.head_dim = args.index_head_dim
-        self.rope_head_dim = args.qk_rope_head_dim
-        self.index_topk = args.index_topk
-        self.q_lora_rank = args.q_lora_rank
-        self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)
-        self.wk = nn.Linear(self.dim, self.head_dim, bias=False)
-        self.k_norm = nn.LayerNorm(self.head_dim)
-        self.weights_proj = nn.Linear(self.dim, self.n_heads, bias=False)
-        self.softmax_scale = self.head_dim**-0.5
-        self.rope = nn.RoPE(
-            dims=self.rope_head_dim,
-            base=args.rope_theta,
-            traditional=False,  # Non-interleaved
-        )
-        self.rope = initialize_rope(
-            dims=args.qk_rope_head_dim,
-            base=args.rope_theta,
-            traditional=False,
-            max_position_embeddings=args.max_position_embeddings,
-            scaling_config=args.rope_scaling,
-        )
-
+class ParallaxDeepSeekV32Indexer(MLXDeepseekV32Indexer):
     def __call__(
         self,
         x: mx.array,
@@ -168,207 +105,11 @@ def __call__(
             return mx.argpartition(scores, kth=-self.index_topk, axis=-1)[..., -self.index_topk :]
 
 
-class DeepseekV32Attention(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.max_position_embeddings = config.max_position_embeddings
-        self.rope_theta = config.rope_theta
-        self.q_lora_rank = config.q_lora_rank
-        self.qk_rope_head_dim = config.qk_rope_head_dim
-        self.kv_lora_rank = config.kv_lora_rank
-        self.v_head_dim = config.v_head_dim
-        self.qk_nope_head_dim = config.qk_nope_head_dim
-        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
-
-        self.scale = self.q_head_dim**-0.5
-
-        if self.q_lora_rank is None:
-            self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.q_head_dim, bias=False)
-        else:
-            self.q_a_proj = nn.Linear(
-                self.hidden_size, self.q_lora_rank, bias=config.attention_bias
-            )
-            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
-            self.q_b_proj = nn.Linear(
-                self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
-            )
-
-        self.kv_a_proj_with_mqa = nn.Linear(
-            self.hidden_size,
-            self.kv_lora_rank + self.qk_rope_head_dim,
-            bias=config.attention_bias,
-        )
-        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
-            bias=False,
-        )
-
-        self.o_proj = nn.Linear(
-            self.num_heads * self.v_head_dim,
-            self.hidden_size,
-            bias=config.attention_bias,
-        )
-
-        if self.config.rope_scaling is not None:
-            mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
-            if mscale_all_dim:
-                scaling_factor = self.config.rope_scaling["factor"]
-                if scaling_factor > 1:
-                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
-                    self.scale = self.scale * s * s
-
-        self.indexer = Indexer(config)
-        self.rope = initialize_rope(
-            dims=self.qk_rope_head_dim,
-            base=self.rope_theta,
-            traditional=True,
-            max_position_embeddings=self.max_position_embeddings,
-            scaling_config=self.config.rope_scaling,
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        pass
-
-
-class DeepseekV32MLP(nn.Module):
-    def __init__(self, config: ModelArgs, hidden_size: int = None, intermediate_size: int = None):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
-        self.intermediate_size = (
-            config.intermediate_size if intermediate_size is None else intermediate_size
-        )
-
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-
-    def __call__(self, x):
-        down_proj = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
-
-
-@mx.compile
-def group_expert_select(
-    gates,
-    e_score_correction_bias,
-    top_k,
-    n_group,
-    topk_group,
-    routed_scaling_factor,
-    norm_topk_prob,
-):
-
-    scores = mx.sigmoid(gates.astype(mx.float32))
-    orig_scores = scores
-    scores = scores + e_score_correction_bias
-    if n_group > 1:
-        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
-        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
-        k = n_group - topk_group
-        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
-        scores = mx.put_along_axis(scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2)
-        scores = mx.flatten(scores, -2, -1)
-
-    k = top_k
-    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
-    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
-    if top_k > 1 and norm_topk_prob:
-        denominator = scores.sum(axis=-1, keepdims=True)
-        scores = scores / denominator
-    scores = scores * routed_scaling_factor
-
-    return inds, scores
-
-
-class MoEGate(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.top_k = config.num_experts_per_tok
-        self.norm_topk_prob = config.norm_topk_prob
-        self.n_routed_experts = config.n_routed_experts
-        self.routed_scaling_factor = config.routed_scaling_factor
-        self.n_group = config.n_group
-        self.topk_group = config.topk_group
-        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
-        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
-        assert config.topk_method == "noaux_tc", "Unsupported topk method."
-
-    def __call__(self, x):
-        return group_expert_select(
-            x @ self.weight.T,
-            self.e_score_correction_bias,
-            self.top_k,
-            self.n_group,
-            self.topk_group,
-            self.routed_scaling_factor,
-            self.norm_topk_prob,
-        )
-
-
-class DeepseekV32MoE(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.num_experts_per_tok = config.num_experts_per_tok
-        self.switch_mlp = SwitchGLU(
-            config.hidden_size,
-            config.moe_intermediate_size,
-            config.n_routed_experts,
-        )
-
-        self.gate = MoEGate(config)
-        if config.n_shared_experts is not None:
-            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
-            self.shared_experts = DeepseekV32MLP(config=config, intermediate_size=intermediate_size)
+class ParallaxDeepSeekV32Attention(MLXDeepseekV32Attention):
 
-    def __call__(self, x):
-        inds, scores = self.gate(x)
-        y = self.switch_mlp(x, inds)
-        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
-        if self.config.n_shared_experts is not None:
-            y = y + self.shared_experts(x)
-
-        return y
-
-
-class DeepseekV32DecoderLayer(nn.Module):
-    def __init__(self, config: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.self_attn = DeepseekV32Attention(config)
-        self.mlp = (
-            DeepseekV32MoE(config)
-            if (
-                config.n_routed_experts is not None
-                and layer_idx >= config.first_k_dense_replace
-                and layer_idx % config.moe_layer_freq == 0
-            )
-            else DeepseekV32MLP(config)
-        )
-        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        pass
-
-
-class ParallaxDeepSeekV32Attention(DeepseekV32Attention):
+    def __init__(self, args: ModelArgs):
+        super().__init__(args)
+        self.indexer = ParallaxDeepSeekV32Indexer(args)
 
     def __call__(
         self,
@@ -481,7 +222,7 @@ def __call__(
         return self.o_proj(output)
 
 
-class ParallaxDeepSeekV32Block(DeepseekV32DecoderLayer):
+class ParallaxDeepSeekV32Block(MLXDeepseekV32Block):
     def __init__(self, args: ModelArgs, layer_idx: int):
         super().__init__(args, layer_idx=layer_idx)
         self.self_attn = ParallaxDeepSeekV32Attention(args)