From 0d5466040117242f1ea43357d3429564a9e6c47f Mon Sep 17 00:00:00 2001 From: Jan Hilgard Date: Mon, 16 Feb 2026 19:43:22 +0100 Subject: [PATCH] Add Step 3.5 Flash model support with MTP weight addition Step 3.5 Flash is a 196B MoE model (288 experts, top-8 routing, ~11B active params) with 3 MTP prediction layers. The MLX community 4-bit conversion strips MTP weights and lacks MTP-aware modeling code. This adds: - scripts/add_mtp_weights_step3p5.py: Downloads BF16 MTP shards from the original model, extracts layers 45-47, remaps to mtp.layers.*, quantizes to 4-bit, and installs the MTP modeling file - scripts/modeling_step3p5_mtp.py: Full MLX-native model implementation with MTP support (Step3p5MTP, Step3p5MTPLayer, Step3p5SharedHead) - Reasoning parser alias "step3p5" (reuses deepseek_r1 parser) - Documentation updates in README.md and docs/reference/models.md Note: The custom modeling file is a workaround until ml-explore/mlx-lm#901 is merged upstream. Co-Authored-By: Claude Opus 4.6 --- README.md | 3 +- docs/reference/models.md | 1 + scripts/add_mtp_weights_step3p5.py | 377 +++++++++++++++++ scripts/modeling_step3p5_mtp.py | 655 +++++++++++++++++++++++++++++ vllm_mlx/reasoning/__init__.py | 1 + 5 files changed, 1036 insertions(+), 1 deletion(-) create mode 100644 scripts/add_mtp_weights_step3p5.py create mode 100644 scripts/modeling_step3p5_mtp.py diff --git a/README.md b/README.md index b77d116dd..9765a7604 100644 --- a/README.md +++ b/README.md @@ -164,7 +164,7 @@ python examples/tts_multilingual.py --list-languages ### Reasoning Models -Extract the thinking process from reasoning models like Qwen3 and DeepSeek-R1: +Extract the thinking process from reasoning models like Qwen3, DeepSeek-R1, and Step 3.5 Flash: ```bash # Start server with reasoning parser @@ -187,6 +187,7 @@ print("Answer:", response.choices[0].message.content) |--------|--------|-------------| | `qwen3` | Qwen3 series | Requires both `` and `` tags | | `deepseek_r1` | DeepSeek-R1 | Handles implicit `` tag | +| `step3p5` | Step 3.5 Flash | Alias for `deepseek_r1`, `` tags | ### Embeddings diff --git a/docs/reference/models.md b/docs/reference/models.md index a45550e4d..5f9ad540a 100644 --- a/docs/reference/models.md +++ b/docs/reference/models.md @@ -18,6 +18,7 @@ Browse thousands of pre-optimized models at: **https://huggingface.co/mlx-commun | Phi-3 | 3.8B, 14B | 4-bit | | Granite 3.x, 4.x | Various | 4-bit | | Nemotron | 3 Nano 30B | 6-bit | +| Step 3.5 | Flash 196B (MoE, 11B active) | 4-bit | ### Recommended Models diff --git a/scripts/add_mtp_weights_step3p5.py b/scripts/add_mtp_weights_step3p5.py new file mode 100644 index 000000000..6071f1b06 --- /dev/null +++ b/scripts/add_mtp_weights_step3p5.py @@ -0,0 +1,377 @@ +#!/usr/bin/env python3 +""" +Add MTP (Multi-Token Prediction) weights to the MLX Step-3.5-Flash-4bit model. + +Step 3.5 Flash has 3 MTP layers (layers 45, 46, 47 in the original model) +that were stripped during the MLX community 4-bit conversion. This script: + +1. Downloads the BF16 shards containing MTP weights from the original model +2. Extracts layers 45-47, remaps to mtp.layers.{0,1,2}.* +3. Quantizes to 4-bit (matching the MLX model's quantization) +4. Saves as model-mtp.safetensors +5. Updates model.safetensors.index.json and config.json + +Usage: + python add_mtp_weights_step3p5.py [--mlx-model-path PATH] [--download-dir DIR] + +Requirements: + pip install mlx +""" + +import argparse +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +# Default paths +DEFAULT_MLX_MODEL = os.path.expanduser( + "~/.cache/huggingface/hub/models--mlx-community--Step-3.5-Flash-4bit" +) +DEFAULT_SOURCE_MODEL = "stepfun-ai/Step-3.5-Flash" + +# MTP layers are in these two shards (layer 45 shared_head.norm in shard 1, rest in shard 2) +MTP_SHARDS = { + "model-00001.safetensors": "https://huggingface.co/{model}/resolve/main/model-00001.safetensors", + "model-00002.safetensors": "https://huggingface.co/{model}/resolve/main/model-00002.safetensors", +} + +# Source layer indices for MTP +MTP_SOURCE_LAYERS = [45, 46, 47] +NUM_MTP_LAYERS = 3 + + +def find_snapshot_dir(model_path: str) -> Path: + """Find the latest snapshot directory in HF cache structure.""" + snapshots_dir = Path(model_path) / "snapshots" + if not snapshots_dir.exists(): + if (Path(model_path) / "config.json").exists(): + return Path(model_path) + raise FileNotFoundError(f"No snapshots found in {model_path}") + + snapshots = sorted(snapshots_dir.iterdir(), key=lambda p: p.stat().st_mtime) + if not snapshots: + raise FileNotFoundError(f"No snapshots in {snapshots_dir}") + return snapshots[-1] + + +def download_mtp_shards(dest_path: Path, source_model: str) -> list[Path]: + """Download the BF16 shards containing MTP weights using curl with resume.""" + shard_paths = [] + for shard_name, url_template in MTP_SHARDS.items(): + url = url_template.format(model=source_model) + shard_path = dest_path / shard_name + + if shard_path.exists(): + size_gb = shard_path.stat().st_size / 1e9 + print(f"Shard {shard_name} already exists ({size_gb:.2f} GB)") + if size_gb < 1.0: + print(" WARNING: File seems too small, re-downloading...") + else: + shard_paths.append(shard_path) + continue + + print(f"Downloading {shard_name}...") + print(f" URL: {url}") + result = subprocess.run( + ["curl", "-L", "-C", "-", "-o", str(shard_path), url], + check=False, + ) + if result.returncode != 0: + raise RuntimeError( + f"Download of {shard_name} failed with code {result.returncode}" + ) + + size_gb = shard_path.stat().st_size / 1e9 + print(f" Downloaded: {size_gb:.2f} GB") + shard_paths.append(shard_path) + + return shard_paths + + +def extract_and_quantize_mtp_weights( + shard_paths: list[Path], snapshot_dir: Path, quantization_bits: int = 4 +): + """Extract MTP weights from BF16 shards, remap, quantize, and save.""" + import mlx.core as mx + + mx.set_default_device(mx.cpu) + + print(f"\nExtracting MTP weights from {len(shard_paths)} shards...") + + # Load and merge MTP weights from both shards + mtp_weights = {} + for shard_path in shard_paths: + print(f" Loading {shard_path.name}...") + all_weights = mx.load(str(shard_path)) + for k, v in all_weights.items(): + # Only keep layers 45, 46, 47 + for src_layer in MTP_SOURCE_LAYERS: + if k.startswith(f"model.layers.{src_layer}."): + mtp_weights[k] = v + del all_weights + + print(f"Found {len(mtp_weights)} MTP weight keys") + + # Remap keys: model.layers.{45,46,47}.* -> mtp.layers.{0,1,2}.* + remapped = {} + for key, value in mtp_weights.items(): + new_key = key + for i, src_layer in enumerate(MTP_SOURCE_LAYERS): + src_prefix = f"model.layers.{src_layer}." + if key.startswith(src_prefix): + suffix = key[len(src_prefix) :] + # transformer.shared_head.* -> shared_head.* + if suffix.startswith("transformer.shared_head."): + suffix = suffix.replace("transformer.shared_head.", "shared_head.") + new_key = f"mtp.layers.{i}.{suffix}" + break + remapped[new_key] = value + del mtp_weights + + print(f"Remapped to {len(remapped)} keys with mtp.layers.* prefix") + + # Read quantization config + config_path = snapshot_dir / "config.json" + with open(config_path) as f: + config = json.load(f) + + quant_config = config.get("quantization", {}) + bits = quant_config.get("bits", quantization_bits) + group_size = quant_config.get("group_size", 64) + print(f"Target quantization: {bits}-bit, group_size={group_size}") + + # Norm suffixes that need +1.0 adjustment (HF zero_centered RMSNorm -> MLX standard) + norm_suffixes = ( + ".input_layernorm.weight", + ".post_attention_layernorm.weight", + ".q_norm.weight", + ".k_norm.weight", + ".enorm.weight", + ".hnorm.weight", + ".shared_head.norm.weight", + ) + + # Keys to keep in FP (not quantize) + def should_skip_quantize(key): + skip_patterns = ( + "layernorm.weight", + "q_norm.weight", + "k_norm.weight", + "enorm.weight", + "hnorm.weight", + "shared_head.norm.weight", + ) + return any(p in key for p in skip_patterns) + + quantized_weights = {} + for key in sorted(remapped.keys()): + weight = remapped.pop(key) + + # Norm adjustment: +1.0 for zero-centered RMSNorm weights (HF -> MLX convention) + if any(key.endswith(s) for s in norm_suffixes) and weight.ndim == 1: + weight = weight + 1.0 + mx.eval(weight) + print(f" Adjusted norm: {key}") + + if should_skip_quantize(key): + print(f" Keep FP: {key} {weight.shape}") + quantized_weights[key] = weight + elif weight.ndim >= 2 and weight.shape[-1] >= group_size: + q_w, q_s, q_b = mx.quantize(weight, group_size=group_size, bits=bits) + mx.eval(q_w, q_s, q_b) + print(f" Quantize {bits}-bit: {key} {q_w.shape}") + quantized_weights[key] = q_w + quantized_weights[key.replace(".weight", ".scales")] = q_s + quantized_weights[key.replace(".weight", ".biases")] = q_b + else: + print(f" Keep FP (small): {key} {weight.shape}") + quantized_weights[key] = weight + del weight + del remapped + + # Save MTP weights + mtp_output_file = snapshot_dir / "model-mtp.safetensors" + print( + f"\nSaving {len(quantized_weights)} quantized MTP weights to {mtp_output_file}" + ) + mx.save_safetensors(str(mtp_output_file), quantized_weights) + + total_bytes = sum(v.nbytes for v in quantized_weights.values()) + print(f"MTP weights size: {total_bytes / 1e6:.1f} MB (quantized)") + + return mtp_output_file, list(quantized_weights.keys()) + + +def update_model_index(snapshot_dir: Path, mtp_keys: list): + """Update model.safetensors.index.json to include MTP weight keys.""" + index_path = snapshot_dir / "model.safetensors.index.json" + if not index_path.exists(): + print(f"WARNING: No index file found at {index_path}, skipping") + return + + with open(index_path) as f: + index = json.load(f) + + weight_map = index.get("weight_map", {}) + + for key in mtp_keys: + weight_map[key] = "model-mtp.safetensors" + + index["weight_map"] = weight_map + + with open(index_path, "w") as f: + json.dump(index, f, indent=2) + + print(f"Updated {index_path} with {len(mtp_keys)} MTP weight entries") + + +def update_config(snapshot_dir: Path): + """Ensure config.json has num_nextn_predict_layers=3.""" + config_path = snapshot_dir / "config.json" + with open(config_path) as f: + config = json.load(f) + + config["num_nextn_predict_layers"] = NUM_MTP_LAYERS + + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + print(f"Updated {config_path}: num_nextn_predict_layers={NUM_MTP_LAYERS}") + + +def install_modeling_file(snapshot_dir: Path): + """Copy modeling_step3p5_mtp.py into the snapshot directory as modeling_step3p5.py. + + The stock MLX community model lacks MTP support in its modeling file. + This replaces it with the MTP-enabled version bundled alongside this script. + + Note: Once https://github.com/ml-explore/mlx-lm/pull/901 is merged upstream, + this workaround will no longer be necessary. + """ + script_dir = Path(__file__).resolve().parent + src = script_dir / "modeling_step3p5_mtp.py" + dst = snapshot_dir / "modeling_step3p5.py" + + if not src.exists(): + print(f"WARNING: {src} not found, skipping modeling file install") + return + + import shutil + + shutil.copy2(src, dst) + print(f"Installed MTP modeling file: {dst}") + + +def main(): + parser = argparse.ArgumentParser( + description="Add MTP weights to MLX Step-3.5-Flash-4bit model" + ) + parser.add_argument( + "--mlx-model-path", + type=str, + default=DEFAULT_MLX_MODEL, + help=f"Path to MLX model directory (default: {DEFAULT_MLX_MODEL})", + ) + parser.add_argument( + "--source-model", + type=str, + default=DEFAULT_SOURCE_MODEL, + help=f"HuggingFace model for MTP shards (default: {DEFAULT_SOURCE_MODEL})", + ) + parser.add_argument( + "--download-dir", + type=str, + default=None, + help="Directory to download MTP shards to (default: temp dir)", + ) + parser.add_argument( + "--bits", + type=int, + default=4, + help="Quantization bits (default: 4, matching 4-bit model)", + ) + parser.add_argument( + "--skip-download", + action="store_true", + help="Skip download (use existing shards)", + ) + args = parser.parse_args() + + print("=" * 60) + print("MTP Weight Addition for Step-3.5-Flash MLX Model") + print("=" * 60) + + # Find snapshot directory + snapshot_dir = find_snapshot_dir(args.mlx_model_path) + print(f"\nMLX model snapshot: {snapshot_dir}") + + # Verify config exists + config_path = snapshot_dir / "config.json" + if not config_path.exists(): + print(f"ERROR: No config.json found in {snapshot_dir}") + sys.exit(1) + + with open(config_path) as f: + config = json.load(f) + print(f"Model type: {config.get('model_type', 'unknown')}") + print(f"Hidden size: {config.get('hidden_size', '?')}") + print(f"Num hidden layers: {config.get('num_hidden_layers', '?')}") + print(f"MoE experts: {config.get('moe_num_experts', '?')}") + + # Check if MTP weights already exist + existing_mtp = config.get("num_nextn_predict_layers", 0) + if existing_mtp > 0: + index_path = snapshot_dir / "model.safetensors.index.json" + if index_path.exists(): + with open(index_path) as f: + index = json.load(f) + mtp_keys = [k for k in index.get("weight_map", {}) if k.startswith("mtp.")] + if mtp_keys: + print(f"\nFound {len(mtp_keys)} existing MTP weight keys") + print("MTP weights already added. Nothing to do.") + sys.exit(0) + + # Download MTP shards + download_dir = ( + Path(args.download_dir) if args.download_dir else Path(tempfile.mkdtemp()) + ) + print(f"\nDownload directory: {download_dir}") + + if not args.skip_download: + shard_paths = download_mtp_shards(download_dir, args.source_model) + else: + shard_paths = [download_dir / name for name in MTP_SHARDS] + for p in shard_paths: + if not p.exists(): + print(f"ERROR: Shard not found at {p}") + sys.exit(1) + + # Extract, quantize, and save + mtp_file, mtp_keys = extract_and_quantize_mtp_weights( + shard_paths, snapshot_dir, quantization_bits=args.bits + ) + + # Update model index + update_model_index(snapshot_dir, mtp_keys) + + # Update config + update_config(snapshot_dir) + + # Install MTP-enabled modeling file + install_modeling_file(snapshot_dir) + + print("\n" + "=" * 60) + print("SUCCESS! MTP weights added to MLX model.") + print("=" * 60) + print(f"\nMTP weight file: {mtp_file}") + print(f"Total MTP keys: {len(mtp_keys)}") + print("\nTo use MTP, start the server with --enable-mtp:") + print(" vllm-mlx serve mlx-community/Step-3.5-Flash-4bit \\") + print(" --enable-mtp --port 1340") + + +if __name__ == "__main__": + main() diff --git a/scripts/modeling_step3p5_mtp.py b/scripts/modeling_step3p5_mtp.py new file mode 100644 index 000000000..477146273 --- /dev/null +++ b/scripts/modeling_step3p5_mtp.py @@ -0,0 +1,655 @@ +# Copyright 2025 StepFun AI / MLX Community +# MLX-native implementation of Step3p5 with MTP (Multi-Token Prediction) support. +# +# Based on the HuggingFace PyTorch modeling_step3p5.py, rewritten for mlx_lm. +# MTP architecture: 3 prediction layers (originally layers 45-47) with dense MLP, +# per-layer shared_head, eh_proj fusion, and sliding_attention. + +from dataclasses import dataclass, field +from typing import Any + +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models.base import ( + BaseModelArgs, + create_attention_mask, + scaled_dot_product_attention, +) +from mlx_lm.models.cache import KVCache +from mlx_lm.models.rope_utils import initialize_rope +from mlx_lm.models.switch_layers import SwitchGLU + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str = "step3p5" + hidden_size: int = 4096 + num_hidden_layers: int = 45 + intermediate_size: int = 11264 + num_attention_heads: int = 64 + num_attention_groups: int = 8 # num_key_value_heads for full attention + head_dim: int = 128 + vocab_size: int = 128896 + rms_norm_eps: float = 1e-5 + rope_theta: Any = 10000.0 # float or list[float] + rope_scaling: dict[str, Any] | None = None + max_position_embeddings: int = 262144 + max_seq_len: int = 262144 + sliding_window: int = 512 + # MoE config + use_moe: bool = True + moe_num_experts: int = 288 + moe_top_k: int = 8 + moe_intermediate_size: int = 1280 + moe_every_n_layer: int = 1 + moe_layer_offset: int = 0 + moe_layers_enum: str = "" + share_expert_dim: int = 1280 + moe_router_activation: str = "sigmoid" + moe_router_scaling_factor: float = 3.0 + norm_expert_weight: bool = True + need_fp32_gate: bool = True + use_moe_router_bias: bool = True + # Attention config + att_impl_type: str = "GQA" + layer_types: list[str] = field(default_factory=list) + attention_other_setting: dict[str, Any] | None = None + use_head_wise_attn_gate: bool = True + use_qk_norm: bool = True + use_rope_layers: list[bool] = field(default_factory=list) + partial_rotary_factors: list[float] | None = None + yarn_only_types: list[str] = field(default_factory=list) + # SwiGLU clamp limits + swiglu_limits: list[float | None] = field(default_factory=list) + swiglu_limits_shared: list[float | None] = field(default_factory=list) + # MTP config + num_nextn_predict_layers: int = 0 + # Zero-centered RMSNorm (MLX community already converted to standard format) + zero_centered: bool = True + tie_word_embeddings: bool = False + + def __post_init__(self): + # Parse moe_layers_enum string to list + if isinstance(self.moe_layers_enum, str) and self.moe_layers_enum.strip(): + self._moe_layer_indices = set( + int(x) for x in self.moe_layers_enum.strip().split(",") + ) + else: + self._moe_layer_indices = set(range(1, self.num_hidden_layers)) + + +class Step3p5RMSNorm(nn.Module): + """RMSNorm — MLX community weights already have +1 baked in (standard format).""" + + def __init__(self, hidden_size: int, eps: float = 1e-5): + super().__init__() + self.weight = mx.ones(hidden_size) + self.eps = eps + + def __call__(self, x: mx.array) -> mx.array: + return mx.fast.rms_norm(x, self.weight, self.eps) + + +class Step3p5MLP(nn.Module): + """Standard dense SwiGLU MLP.""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + swiglu_limit: float | None = None, + ): + 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.limit = swiglu_limit + + def __call__(self, x: mx.array) -> mx.array: + gate = nn.silu(self.gate_proj(x)) + up = self.up_proj(x) + if self.limit is not None and self.limit > 0: + gate = mx.clip(gate, a_min=None, a_max=self.limit) + up = mx.clip(up, a_min=-self.limit, a_max=self.limit) + return self.down_proj(gate * up) + + +class Step3p5Router(nn.Module): + """MoE router with sigmoid gating and optional router bias. + Weight keys: gate.gate.{weight,scales,biases}, gate.router_bias + """ + + def __init__(self, args: ModelArgs): + super().__init__() + self.num_experts = args.moe_num_experts + self.top_k = args.moe_top_k + self.routed_scaling_factor = args.moe_router_scaling_factor + self.norm_expert_weight = args.norm_expert_weight + self.use_moe_router_bias = args.use_moe_router_bias + self.need_fp32_gate = args.need_fp32_gate + + self.gate = nn.Linear(args.hidden_size, self.num_experts, bias=False) + + if self.use_moe_router_bias: + self.router_bias = mx.zeros(self.num_experts) + + def __call__(self, x: mx.array): + # Always use self.gate(x) — weight may be quantized (packed format) + router_logits = self.gate(x) + if self.need_fp32_gate: + router_logits = router_logits.astype(mx.float32) + + gate_prob = mx.sigmoid(router_logits.astype(mx.float32)) + + if self.use_moe_router_bias: + gate_prob_biased = gate_prob + self.router_bias + inds = mx.argpartition(gate_prob_biased, kth=-self.top_k, axis=-1)[ + ..., -self.top_k : + ] + else: + inds = mx.argpartition(gate_prob, kth=-self.top_k, axis=-1)[ + ..., -self.top_k : + ] + + scores = mx.take_along_axis(gate_prob, inds, axis=-1) + if self.norm_expert_weight: + scores = scores / (scores.sum(axis=-1, keepdims=True) + 1e-20) + + scores = scores * self.routed_scaling_factor + return inds, scores + + +class Step3p5MoEBlock(nn.Module): + """MoE block matching weight key structure: mlp.gate, mlp.switch_mlp, mlp.share_expert.""" + + def __init__( + self, + args: ModelArgs, + swiglu_limit: float | None = None, + swiglu_limit_shared: float | None = None, + ): + super().__init__() + self.gate = Step3p5Router(args) + self.switch_mlp = SwitchGLU( + args.hidden_size, args.moe_intermediate_size, args.moe_num_experts + ) + self.share_expert = Step3p5MLP( + args.hidden_size, args.share_expert_dim, swiglu_limit=swiglu_limit_shared + ) + self.limit = swiglu_limit + + def __call__(self, x: mx.array) -> mx.array: + inds, scores = self.gate(x) + y = self.switch_mlp(x, inds) + if self.limit is not None and self.limit > 0: + y = mx.clip(y, a_min=-self.limit, a_max=self.limit) + y = (y * scores[..., None]).sum(axis=-2) + y = y + self.share_expert(x) + return y + + +class Step3p5Attention(nn.Module): + """Multi-head attention with optional head-wise gating, QK norm, and sliding window.""" + + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.head_dim = args.head_dim + + # Determine attention type + if args.layer_types and layer_idx < len(args.layer_types): + self.is_sliding = args.layer_types[layer_idx] == "sliding_attention" + else: + self.is_sliding = layer_idx % 2 == 0 + + # Set head counts based on attention type + if self.is_sliding and args.attention_other_setting: + self.num_heads = args.attention_other_setting.get( + "num_attention_heads", args.num_attention_heads + ) + self.num_kv_heads = args.attention_other_setting.get( + "num_attention_groups", args.num_attention_groups + ) + else: + self.num_heads = args.num_attention_heads + self.num_kv_heads = args.num_attention_groups + + self.scale = self.head_dim**-0.5 + + self.q_proj = nn.Linear( + args.hidden_size, self.num_heads * self.head_dim, bias=False + ) + self.k_proj = nn.Linear( + args.hidden_size, self.num_kv_heads * self.head_dim, bias=False + ) + self.v_proj = nn.Linear( + args.hidden_size, self.num_kv_heads * self.head_dim, bias=False + ) + self.o_proj = nn.Linear( + self.num_heads * self.head_dim, args.hidden_size, bias=False + ) + + self.q_norm = Step3p5RMSNorm(self.head_dim, eps=args.rms_norm_eps) + self.k_norm = Step3p5RMSNorm(self.head_dim, eps=args.rms_norm_eps) + + self.use_head_wise_attn_gate = args.use_head_wise_attn_gate + if self.use_head_wise_attn_gate: + self.g_proj = nn.Linear(args.hidden_size, self.num_heads, bias=False) + + # RoPE — determine partial_rotary_factor and theta for this layer + partial_rotary_factor = 1.0 + if args.partial_rotary_factors and layer_idx < len(args.partial_rotary_factors): + partial_rotary_factor = args.partial_rotary_factors[layer_idx] + + rope_theta = args.rope_theta + if isinstance(args.rope_theta, list): + rope_theta = ( + args.rope_theta[layer_idx] + if layer_idx < len(args.rope_theta) + else 10000.0 + ) + + # Determine if this layer uses yarn/rope scaling + rope_scaling = None + if args.yarn_only_types and args.layer_types: + if ( + layer_idx < len(args.layer_types) + and args.layer_types[layer_idx] in args.yarn_only_types + ): + rope_scaling = args.rope_scaling + elif args.rope_scaling: + rope_scaling = args.rope_scaling + + rotary_dim = int(self.head_dim * partial_rotary_factor) + self.rope = initialize_rope( + rotary_dim, + base=rope_theta, + traditional=False, + scaling_config=rope_scaling, + max_position_embeddings=args.max_position_embeddings, + ) + + def __call__( + self, + x: mx.array, + mask: mx.array | None = None, + cache: Any | None = None, + ) -> mx.array: + B, L, _ = x.shape + + queries = self.q_proj(x) + keys = self.k_proj(x) + values = self.v_proj(x) + + # Reshape and apply QK norm + queries = self.q_norm( + queries.reshape(B, L, self.num_heads, self.head_dim) + ).transpose(0, 2, 1, 3) + keys = self.k_norm( + keys.reshape(B, L, self.num_kv_heads, self.head_dim) + ).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.num_kv_heads, self.head_dim).transpose( + 0, 2, 1, 3 + ) + + # Head-wise gate + gate = None + if self.use_head_wise_attn_gate: + gate = self.g_proj(x) # [B, L, num_heads] + + # RoPE + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + keys, values = cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + + # Apply head-wise gating + if gate is not None: + output = ( + output.reshape(B, L, self.num_heads, self.head_dim) + * mx.sigmoid(gate)[..., None] + ).reshape(B, L, -1) + + return self.o_proj(output) + + +class Step3p5DecoderLayer(nn.Module): + """Single transformer decoder layer with attention + MLP/MoE.""" + + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.self_attn = Step3p5Attention(args, layer_idx) + + self.input_layernorm = Step3p5RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = Step3p5RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + + # Determine swiglu limits for this layer + swiglu_limit = None + if args.swiglu_limits and layer_idx < len(args.swiglu_limits): + val = args.swiglu_limits[layer_idx] + if val is not None and val != 0: + swiglu_limit = val + + swiglu_limit_shared = None + if args.swiglu_limits_shared and layer_idx < len(args.swiglu_limits_shared): + val = args.swiglu_limits_shared[layer_idx] + if val is not None and val != 0: + swiglu_limit_shared = val + + # MoE or dense MLP — both use self.mlp to match weight key naming + self.is_moe = layer_idx in args._moe_layer_indices + if self.is_moe: + self.mlp = Step3p5MoEBlock( + args, swiglu_limit=swiglu_limit, swiglu_limit_shared=swiglu_limit_shared + ) + else: + self.mlp = Step3p5MLP( + args.hidden_size, + args.intermediate_size, + swiglu_limit=swiglu_limit_shared, + ) + + def __call__( + self, + x: mx.array, + mask: mx.array | None = None, + cache: Any | None = None, + ) -> mx.array: + # Self-attention + residual = x + x = self.input_layernorm(x) + x = self.self_attn(x, mask=mask, cache=cache) + x = residual + x + + # FFN + residual = x + x = self.post_attention_layernorm(x) + x = self.mlp(x) + x = residual + x + + return x + + +class Step3p5Model(nn.Module): + """Step 3.5 backbone transformer.""" + + def __init__(self, args: ModelArgs): + super().__init__() + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + Step3p5DecoderLayer(args, layer_idx=i) + for i in range(args.num_hidden_layers) + ] + self.norm = Step3p5RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + + def __call__( + self, + inputs: mx.array, + cache: Any | None = None, + return_prenorm: bool = False, + ) -> mx.array: + hidden_states = self.embed_tokens(inputs) + + if cache is None: + cache = [None] * len(self.layers) + + mask = create_attention_mask(hidden_states, cache[0]) + + for layer, c in zip(self.layers, cache): + hidden_states = layer(hidden_states, mask=mask, cache=c) + + if return_prenorm: + return self.norm(hidden_states), hidden_states + return self.norm(hidden_states) + + +class Step3p5SharedHead(nn.Module): + """Per-MTP-layer prediction head: norm + linear output projection.""" + + def __init__(self, args: ModelArgs): + super().__init__() + self.norm = Step3p5RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.output = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + return self.output(self.norm(x)) + + +class Step3p5MTPLayer(nn.Module): + """Single MTP prediction layer. + + Architecture: + 1. Normalize hidden_states (hnorm) and token embedding (enorm) separately + 2. Concatenate and project: [B, L, 2H] -> [B, L, H] via eh_proj + 3. Standard decoder block: attention + dense MLP (NOT MoE) + 4. Per-layer shared_head for logit prediction + """ + + def __init__(self, args: ModelArgs, layer_idx: int = 0): + super().__init__() + self.hnorm = Step3p5RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.enorm = Step3p5RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.eh_proj = nn.Linear(args.hidden_size * 2, args.hidden_size, bias=False) + + # MTP uses sliding_attention type — pick a sliding layer_idx for RoPE config + # Find a sliding_attention layer index for correct RoPE params + mtp_layer_idx = 1 # default sliding layer + if args.layer_types: + for i, lt in enumerate(args.layer_types): + if lt == "sliding_attention": + mtp_layer_idx = i + break + + self.self_attn = Step3p5Attention(args, layer_idx=mtp_layer_idx) + self.mlp = Step3p5MLP(args.hidden_size, args.intermediate_size) + self.input_layernorm = Step3p5RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = Step3p5RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + self.shared_head = Step3p5SharedHead(args) + + def __call__( + self, + hidden_states: mx.array, + input_embeds: mx.array, + mask: mx.array | None = None, + cache: Any | None = None, + ) -> tuple[mx.array, mx.array]: + """ + Returns: + (mtp_hidden, logits) — hidden for chaining to next MTP layer, logits for this layer + """ + h = self.hnorm(hidden_states) + e = self.enorm(input_embeds) + x = self.eh_proj(mx.concatenate([e, h], axis=-1)) + + # Standard decoder: attention + dense MLP + residual = x + x = self.input_layernorm(x) + x = self.self_attn(x, mask=mask, cache=cache) + residual + + residual = x + x = self.post_attention_layernorm(x) + x = self.mlp(x) + residual + + logits = self.shared_head(x) + return x, logits + + +class Step3p5MTP(nn.Module): + """MTP module with multiple prediction layers.""" + + def __init__(self, args: ModelArgs): + super().__init__() + self.layers = [ + Step3p5MTPLayer(args, layer_idx=i) + for i in range(args.num_nextn_predict_layers) + ] + + +class Model(nn.Module): + """Step3p5ForCausalLM — MLX-native with MTP support.""" + + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = Step3p5Model(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + # MTP head + self._mtp_num_layers = args.num_nextn_predict_layers + if self._mtp_num_layers > 0: + self.mtp = Step3p5MTP(args) + else: + self.mtp = None + + def __call__( + self, + inputs: mx.array, + cache: Any | None = None, + return_hidden: bool = False, + ) -> mx.array: + if return_hidden: + hidden_states, prenorm_hidden = self.model( + inputs, cache, return_prenorm=True + ) + else: + hidden_states = self.model(inputs, cache) + + if self.args.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(hidden_states) + else: + out = self.lm_head(hidden_states) + + if return_hidden: + return out, prenorm_hidden + return out + + def mtp_forward( + self, + hidden_states: mx.array, + next_token_ids: mx.array, + mtp_cache: Any | None = None, + ) -> mx.array: + """Run MTP head to predict token n+2 given hidden states and token n+1. + + Uses only the first MTP layer (vllm-mlx compatible single-draft-token mode). + + Args: + hidden_states: [B, 1, H] prenorm hidden states from main model + next_token_ids: [B, 1] token IDs for position n+1 + mtp_cache: list of KVCache for MTP layers + + Returns: + logits: [B, 1, V] logits for token n+2 + """ + if self.mtp is None: + raise RuntimeError("MTP head not loaded (num_nextn_predict_layers=0)") + + input_embeds = self.model.embed_tokens(next_token_ids) + + layer = self.mtp.layers[0] + cache_entry = mtp_cache[0] if mtp_cache else None + mask = create_attention_mask(input_embeds, cache_entry) + _, logits = layer(hidden_states, input_embeds, mask=mask, cache=cache_entry) + return logits + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + return [KVCache() for _ in self.layers] + + def make_mtp_cache(self): + """Create KV cache for MTP layers.""" + if self.mtp is None: + return None + return [KVCache() for _ in self.mtp.layers] + + def sanitize(self, weights): + """Remap weight keys from stored format to module structure.""" + # Check for MTP weights + has_mtp_weights = any(k.startswith("mtp.") for k in weights) + + # Filter out original model.layers.{45,46,47} if they somehow survived + weights = { + k: v + for k, v in weights.items() + if not any(k.startswith(f"model.layers.{i}.") for i in [45, 46, 47]) + } + + if not has_mtp_weights or self._mtp_num_layers == 0: + weights = {k: v for k, v in weights.items() if not k.startswith("mtp.")} + + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + + # Stack per-expert weights if stored individually (HF PyTorch format) + per_expert_key = "model.layers.3.mlp.up_proj.0.weight" + if per_expert_key in weights: + for layer_idx in range(self.args.num_hidden_layers): + if layer_idx not in self.args._moe_layer_indices: + continue + prefix = f"model.layers.{layer_idx}.mlp" + for proj in ["up_proj", "down_proj", "gate_proj"]: + expert_keys = [ + f"{prefix}.{proj}.{e}.weight" + for e in range(self.args.moe_num_experts) + ] + if all(k in weights for k in expert_keys): + stacked = mx.stack([weights.pop(k) for k in expert_keys]) + weights[f"{prefix}.switch_mlp.{proj}.weight"] = stacked + + # Handle MoELinear bulk format (mlp.up_proj.weight [E, out, in] → mlp.switch_mlp) + for layer_idx in range(self.args.num_hidden_layers): + if layer_idx not in self.args._moe_layer_indices: + continue + prefix = f"model.layers.{layer_idx}.mlp" + for proj in ["up_proj", "down_proj", "gate_proj"]: + hf_key = f"{prefix}.{proj}.weight" + mlx_key = f"{prefix}.switch_mlp.{proj}.weight" + if hf_key in weights and mlx_key not in weights: + weights[mlx_key] = weights.pop(hf_key) + + # RMSNorm: MLX community already added +1 to zero-centered weights (standard format). + # add_mtp_weights_step3p5.py also added +1 to MTP norms. No adjustment needed. + + return weights + + @property + def quant_predicate(self): + def predicate(path, _): + # MoE gate routing — use 8-bit + if "mlp.gate" in path and path.endswith(".gate"): + return {"group_size": 64, "bits": 8} + # MTP norms and projections — keep FP + if "mtp." in path and any( + x in path + for x in [ + ".enorm.", + ".hnorm.", + ".shared_head.norm.", + ".input_layernorm.", + ".post_attention_layernorm.", + ".q_norm.", + ".k_norm.", + ] + ): + return False + return True + + return predicate diff --git a/vllm_mlx/reasoning/__init__.py b/vllm_mlx/reasoning/__init__.py index f138796ff..c47599278 100644 --- a/vllm_mlx/reasoning/__init__.py +++ b/vllm_mlx/reasoning/__init__.py @@ -84,6 +84,7 @@ def _register_builtin_parsers(): register_parser("deepseek_r1", DeepSeekR1ReasoningParser) register_parser("gpt_oss", GptOssReasoningParser) register_parser("harmony", HarmonyReasoningParser) + register_parser("step3p5", DeepSeekR1ReasoningParser) # Register built-in parsers on module load