From 68a17871fb45482837717ea2a6d2692c259e9858 Mon Sep 17 00:00:00 2001 From: dongjiang1989 Date: Tue, 13 Oct 2026 14:43:43 +0800 Subject: [PATCH] Enable Spark-X2.5 export and runner support for MLX Add MLX backend support for Spark-X2.5 models (1.7B and 4B variants) with 4-bit quantization, enabling efficient inference on Apple Silicon. Key features: - MLX 4-bit export config with KV cache metadata - Sliding window attention with ring buffer cache - Headwise output gate (per-head sigmoid) - Per-layer RoPE parameters (partial rotary factor) - Comprehensive numerical tests comparing MLX vs eager attention Changes: - Add spark_x2_5_mlx_4w.yaml config for 4w quantization - Add test_spark_x2_5_mlx.py with numerical correctness tests - Update Makefile with spark_x2_5-mlx build target - Update README with MLX export and run instructions - Fix sliding window detection from RingKVCache.window_size - Add unsupported options validation (use_attn_o_gate, use_attn_o_norm, scale_query_by, is_kv_shared_layer) - Fix sdpa_start calculation to avoid export guards - Add max_write_len to RingBufferKVCache to prevent prefill overflow - Fix test Rope constructor usage and freqs generation All fixes maintain backward compatibility and pass code review. Signed-off-by: dongjiang1989 --- Makefile | 42 ++- backends/mlx/llm/et_attention.py | 178 +++++++++++-- examples/models/spark_x2_5/README.md | 70 ++++- .../spark_x2_5/config/spark_x2_5_mlx_4w.yaml | 15 ++ .../models/spark_x2_5/test_spark_x2_5_mlx.py | 244 ++++++++++++++++++ 5 files changed, 513 insertions(+), 36 deletions(-) create mode 100644 examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml create mode 100644 examples/models/spark_x2_5/test_spark_x2_5_mlx.py diff --git a/Makefile b/Makefile index f0fee3adfad..8b7feee487c 100644 --- a/Makefile +++ b/Makefile @@ -19,7 +19,7 @@ # - whisper: Speech recognition model (CPU, CUDA, Metal) # - parakeet: Speech recognition model (CPU, CUDA, Metal, MLX) # - sortformer: Speaker diarization model (CPU, CUDA) -# - nemotron3-diarization: Speaker diarization model (MLX, XNNPACK, CUDA) +# - nemotron3-diarization: Speaker diarization model (MLX, XNNPACK, CUDA, Vulkan) # - supertonic: Text-to-speech model (MLX) # - silero_vad: Voice activity detection model (CPU) # - llama: Text generation model (CPU) @@ -93,7 +93,7 @@ # # ============================================================================== -.PHONY: voxtral-cuda voxtral-cpu voxtral-metal voxtral-mlx voxtral_realtime-cuda voxtral_realtime-rocm voxtral_realtime-cpu voxtral_realtime-metal voxtral_realtime-mlx voxtral_tts-cpu voxtral_tts-cuda whisper-cuda whisper-cuda-debug whisper-cpu whisper-metal parakeet-cuda parakeet-cuda-debug parakeet-cpu parakeet-metal parakeet-mlx parakeet-vulkan dinov2-cuda dinov2-cuda-debug sortformer-cuda sortformer-cpu supertonic-mlx silero-vad-cpu llama-cuda llama-cuda-debug llama-cpu lfm_2_5-mlx llava-cpu gemma3-cuda gemma3-cpu gemma4_31b-cuda gemma4_31b-mlx muse-glimmer-cuda muse-glimmer-mlx qwen3_5_moe-cuda qwen3_5_moe-metal qwen3_5_moe-mlx clean help +.PHONY: voxtral-cuda voxtral-cpu voxtral-metal voxtral-mlx voxtral_realtime-cuda voxtral_realtime-rocm voxtral_realtime-cpu voxtral_realtime-metal voxtral_realtime-mlx voxtral_tts-cpu voxtral_tts-cuda whisper-cuda whisper-cuda-debug whisper-cpu whisper-metal parakeet-cuda parakeet-cuda-debug parakeet-cpu parakeet-metal parakeet-mlx parakeet-vulkan dinov2-cuda dinov2-cuda-debug sortformer-cuda sortformer-cpu supertonic-mlx silero-vad-cpu llama-cuda llama-cuda-debug llama-cpu lfm_2_5-mlx spark_x2_5-mlx llava-cpu gemma3-cuda gemma3-cpu gemma4_31b-cuda gemma4_31b-mlx muse-glimmer-cuda muse-glimmer-mlx qwen3_5_moe-cuda qwen3_5_moe-metal qwen3_5_moe-mlx clean help help: @echo "This Makefile adds targets to build runners for various models on various backends. Run using \`make \`. Available targets:" @@ -125,12 +125,14 @@ help: @echo " nemotron3-diarization-mlx - Build Nemotron 3 Diarization runner with MLX backend" @echo " nemotron3-diarization-cpu - Build Nemotron 3 Diarization runner with XNNPACK backend" @echo " nemotron3-diarization-cuda - Build Nemotron 3 Diarization runner with CUDA backend" + @echo " nemotron3-diarization-vulkan - Build Nemotron 3 Diarization runner with Vulkan backend" @echo " supertonic-mlx - Build Supertonic runner with MLX backend" @echo " silero-vad-cpu - Build Silero VAD runner with CPU backend" @echo " llama-cuda - Build Llama runner with CUDA backend" @echo " llama-cuda-debug - Build Llama runner with CUDA backend (debug mode)" @echo " llama-cpu - Build Llama runner with CPU backend" @echo " lfm_2_5-mlx - Build LFM2.5 runner with MLX backend" + @echo " spark_x2_5-mlx - Build Spark-X2.5 runner with MLX backend" @echo " llava-cpu - Build Llava runner with CPU backend" @echo " gemma3-cuda - Build Gemma3 runner with CUDA backend" @echo " gemma3-cpu - Build Gemma3 runner with CPU backend" @@ -138,6 +140,7 @@ help: @echo " gemma4_31b-mlx - Build Gemma 4 31B runner and worker with MLX backend" @echo " muse-glimmer-cuda - Build Muse Glimmer runners and worker with CUDA backend" @echo " muse-glimmer-mlx - Build Muse Glimmer runners and worker with MLX backend" + @echo " muse-glimmer-batching-mlx - Build batching Muse Glimmer runner and worker with MLX" @echo " qwen3_5_moe-cuda - Build Qwen3.5 MoE runner with CUDA backend" @echo " qwen3_5_moe-metal - Build Qwen3.5 MoE runner with Metal backend" @echo " qwen3_5_moe-mlx - Build Qwen3.5 MoE runner with MLX backend" @@ -305,6 +308,12 @@ nemotron3-diarization-cuda: cmake --workflow --preset llm-release-cuda cd examples/models/nemotron3_diarization && cmake --workflow --preset nemotron3-diarization-cuda +.PHONY: nemotron3-diarization-vulkan +nemotron3-diarization-vulkan: + cmake --preset llm-release -DEXECUTORCH_BUILD_VULKAN=ON + cmake --build --preset llm-release-install + cd examples/models/nemotron3_diarization && cmake --workflow --preset nemotron3-diarization-vulkan + sortformer-cuda: @echo "==> Building and installing ExecuTorch with CUDA..." cmake --workflow --preset llm-release-cuda @@ -420,15 +429,22 @@ llama-cuda-debug: @echo "✓ Build complete!" @echo " Binary: cmake-out/examples/models/llama/llama_main" -lfm_2_5-mlx: +# Internal target for MLX runner builds +.mlx-runner: @echo "==> Building and installing ExecuTorch with MLX..." cmake --workflow --preset mlx-release - @echo "==> Building LFM2.5 runner with MLX..." + @echo "==> Building $(MODEL_NAME) runner with MLX..." cd examples/models/llama && cmake --workflow --preset llama-mlx @echo "" @echo "✓ Build complete!" @echo " Binary: cmake-out/examples/models/llama/llama_main" +lfm_2_5-mlx: + @$(MAKE) .mlx-runner MODEL_NAME="LFM2.5" + +spark_x2_5-mlx: + @$(MAKE) .mlx-runner MODEL_NAME="Spark-X2.5" + llava-cpu: @echo "==> Building and installing ExecuTorch..." cmake --workflow --preset llm-release @@ -506,6 +522,24 @@ muse-glimmer-mlx: @echo " DFlash runner: cmake-out/examples/models/muse-glimmer/dflash_runner" @echo " Worker: cmake-out/examples/models/muse-glimmer/muse_glimmer_worker" +.PHONY: muse-glimmer-batching-mlx +muse-glimmer-batching-mlx: + @echo "==> Building and installing ExecuTorch with MLX..." + cmake --workflow --preset mlx-release + @echo "==> Configuring batching Muse Glimmer runners with MLX..." + cd examples/models/muse-glimmer && cmake --preset muse-glimmer-mlx \ + -B "$(CURDIR)/cmake-out/examples/models/muse-glimmer-batching" \ + -DEXECUTORCH_BUILD_CUDA=OFF \ + -DMUSE_GLIMMER_BUILD_BATCHING=ON \ + -DMUSE_GLIMMER_BUILD_BATCHING_TESTS=OFF + cmake --build "$(CURDIR)/cmake-out/examples/models/muse-glimmer-batching" \ + --target muse_glimmer_batching_runner muse_glimmer_batching_worker \ + --parallel "$$(sysctl -n hw.ncpu)" + @echo "" + @echo "Build complete!" + @echo " Runner: $(CURDIR)/cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_runner" + @echo " Worker: $(CURDIR)/cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_worker" + gemma4_31b-cuda: @echo "==> Building and installing ExecuTorch with CUDA..." cmake --workflow --preset llm-release-cuda diff --git a/backends/mlx/llm/et_attention.py b/backends/mlx/llm/et_attention.py index ea39133db3a..907dea0448b 100644 --- a/backends/mlx/llm/et_attention.py +++ b/backends/mlx/llm/et_attention.py @@ -22,10 +22,9 @@ from typing import Any, Optional, Tuple, TYPE_CHECKING import executorch.backends.mlx.custom_ops as _mlx_custom_ops # noqa: F401 - import torch import torch.nn as nn -from executorch.backends.mlx.llm.cache import KVCache +from executorch.backends.mlx.llm.cache import KVCache, RingBufferKVCache from executorch.examples.models.llama.attention import ( Attention, ForwardOptions, @@ -59,6 +58,24 @@ def __init__( if not args.use_kv_cache: raise ValueError("MLXAttention requires use_kv_cache=True") + # Fail loudly on unsupported options (must match from_attention_mha) + if getattr(args, "use_q_gate", False): + raise NotImplementedError("MLXAttentionMHA does not support use_q_gate") + if getattr(args, "is_kv_shared_layer", False): + raise NotImplementedError( + "MLXAttentionMHA does not support is_kv_shared_layer (YOCO)" + ) + if getattr(args, "use_attn_o_norm", False): + raise NotImplementedError("MLXAttentionMHA does not support use_attn_o_norm") + if getattr(args, "use_attn_o_gate", False): + raise NotImplementedError( + "MLXAttentionMHA does not support use_attn_o_gate" + ) + if getattr(args, "scale_query_by", 1.0) != 1.0: + raise NotImplementedError( + f"MLXAttentionMHA does not support scale_query_by={args.scale_query_by}" + ) + self.use_kv_cache = True self.n_heads = args.n_heads self.n_kv_heads = self.n_heads if args.n_kv_heads is None else args.n_kv_heads @@ -76,6 +93,23 @@ def __init__( self.qk_norm_before_rope = args.qk_norm_before_rope self.enable_dynamic_shape = args.enable_dynamic_shape + # Output gate support (Spark-X2.5 uses headwise_attn_output_gate) + self.headwise_attn_output_gate = getattr( + args, "headwise_attn_output_gate", False + ) + + # Sliding window support + self.is_sliding = ( + getattr(args, "layer_types", None) is not None + and layer_id < len(args.layer_types) + and args.layer_types[layer_id] == "sliding_attention" + ) + # Only set sliding_window for sliding layers (consistent with from_attention_mha) + if self.is_sliding: + self.sliding_window = getattr(args, "sliding_window", None) + else: + self.sliding_window = None + if self.use_qk_norm: self.q_norm_fn = RMSNorm(self.head_dim, eps=args.norm_eps) self.k_norm_fn = RMSNorm(self.head_dim, eps=args.norm_eps) @@ -91,6 +125,10 @@ def __init__( ) self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False) + # Output gate weight (Spark-X2.5: headwise gate, shape [dim, n_local_heads]) + if self.headwise_attn_output_gate: + self.og = nn.Linear(self.dim, self.n_local_heads, bias=False) + self.layer_id = layer_id self.rope = rope self.rope_base = rope.params.rope_freq_base @@ -98,13 +136,23 @@ def __init__( self.rope_traditional = not rope.params.use_hf_rope self.rope_dims = int(self.head_dim * rope.params.partial_rotary_factor) - self.kv_cache = KVCache( - max_batch_size=args.max_batch_size, - max_context_length=args.max_context_len, - n_heads=self.n_kv_heads, - head_dim=self.head_dim, - enable_dynamic_shape=args.enable_dynamic_shape, - ) + # Sliding-window layers use ring buffer cache + if self.is_sliding and self.sliding_window: + self.kv_cache = RingBufferKVCache( + max_batch_size=args.max_batch_size, + max_context_length=self.sliding_window, + n_heads=self.n_kv_heads, + head_dim=self.head_dim, + max_write_len=min(args.max_seq_len, self.sliding_window), + ) + else: + self.kv_cache = KVCache( + max_batch_size=args.max_batch_size, + max_context_length=args.max_context_len, + n_heads=self.n_kv_heads, + head_dim=self.head_dim, + enable_dynamic_shape=args.enable_dynamic_shape, + ) @staticmethod def _can_use_fused_rope(params: ModelArgs) -> bool: @@ -119,13 +167,30 @@ def from_attention_mha( """ Create an MLXAttentionMHA from an existing AttentionMHA. - Shares weight references (wq, wk, wv, wo, rope, norm) and creates - a fresh KVCache. + Shares weight references (wq, wk, wv, wo, rope, norm, og) and creates + a fresh KVCache (RingBufferKVCache for sliding-window layers). """ from executorch.examples.models.llama.attention import AttentionMHA assert isinstance(other, AttentionMHA) + # Fail loudly on unsupported AttentionMHA options + _unsupported = [] + if getattr(other, "use_q_gate", False): + _unsupported.append("use_q_gate") + if getattr(other, "is_kv_shared_layer", False): + _unsupported.append("is_kv_shared_layer (YOCO)") + if getattr(other, "use_attn_o_norm", False): + _unsupported.append("use_attn_o_norm") + if getattr(other, "use_attn_o_gate", False): + _unsupported.append("use_attn_o_gate") + if getattr(other, "scale_query_by", 1.0) != 1.0: + _unsupported.append(f"scale_query_by={other.scale_query_by}") + if _unsupported: + raise NotImplementedError( + f"MLXAttentionMHA does not support: {_unsupported}" + ) + instance = cls.__new__(cls) Attention.__init__(instance) @@ -145,6 +210,21 @@ def from_attention_mha( instance.qk_norm_before_rope = other.qk_norm_before_rope instance.enable_dynamic_shape = other.enable_dynamic_shape + # Copy output gate fields (Spark-X2.5 uses headwise_attn_output_gate) + instance.headwise_attn_output_gate = getattr( + other, "headwise_attn_output_gate", False + ) + # use_attn_o_gate and use_attn_o_norm are checked in unsupported list above + + # Copy sliding window fields + # AttentionMHA doesn't store sliding_window as attribute, get it from kv_cache + instance.is_sliding = getattr(other, "is_sliding", False) + if hasattr(other, "kv_cache") and hasattr(other.kv_cache, "window_size"): + # RingKVCache stores window_size + instance.sliding_window = other.kv_cache.window_size + else: + instance.sliding_window = None + # Share weight references instance.wq = other.wq instance.wk = other.wk @@ -152,6 +232,9 @@ def from_attention_mha( instance.wo = other.wo instance.layer_id = other.layer_id instance.rope = other.rope + + # Copy RoPE params from the source layer's rope object + # (each AttentionMHA layer has its own rope with correct params for its layer type) instance.rope_base = other.rope.params.rope_freq_base instance.use_fused_rope = cls._can_use_fused_rope(other.rope.params) instance.rope_traditional = not other.rope.params.use_hf_rope @@ -163,18 +246,33 @@ def from_attention_mha( instance.q_norm_fn = other.q_norm_fn instance.k_norm_fn = other.k_norm_fn - # Create fresh MLX KV cache + # Share output gate weight + if instance.headwise_attn_output_gate: + instance.og = other.og + + # Create fresh MLX KV cache (ring buffer for sliding layers) cache_dtype = dtype if dtype is not None else torch.float32 if hasattr(other, "kv_cache") and hasattr(other.kv_cache, "k_cache"): cache_dtype = dtype if dtype is not None else other.kv_cache.k_cache.dtype - instance.kv_cache = KVCache( - max_batch_size=other.max_batch_size, - max_context_length=other.max_context_len, - n_heads=instance.n_kv_heads, - head_dim=instance.head_dim, - enable_dynamic_shape=other.enable_dynamic_shape, - dtype=cache_dtype, - ) + + if instance.is_sliding and instance.sliding_window: + instance.kv_cache = RingBufferKVCache( + max_batch_size=other.max_batch_size, + max_context_length=instance.sliding_window, + n_heads=instance.n_kv_heads, + head_dim=instance.head_dim, + dtype=cache_dtype, + max_write_len=min(other.max_context_len, instance.sliding_window), + ) + else: + instance.kv_cache = KVCache( + max_batch_size=other.max_batch_size, + max_context_length=other.max_context_len, + n_heads=instance.n_kv_heads, + head_dim=instance.head_dim, + enable_dynamic_shape=other.enable_dynamic_shape, + dtype=cache_dtype, + ) return instance @@ -239,14 +337,38 @@ def forward( k = self.k_norm_fn(k) k, v = self.kv_cache.update(start_pos, k, v) - output = torch.ops.mlx.custom_sdpa( - q, - k, - v, - start_pos=start_pos, - is_causal=True, - scale=self.head_dim**-0.5, - ) + # Sliding-window layers use ring-buffer mask; full layers use causal + if self.is_sliding and isinstance(self.kv_cache, RingBufferKVCache): + sdpa_mask = self.kv_cache.create_sliding_window_mask(start_pos, seqlen) + # seqlen <= max_write_len <= window keeps this non-negative + sdpa_start = self.kv_cache.buffer_size - seqlen + output = torch.ops.mlx.custom_sdpa( + q, + k, + v, + start_pos=sdpa_start, + attn_mask=sdpa_mask, + dropout_p=0.0, + is_causal=False, + scale=self.head_dim**-0.5, + ) + else: + output = torch.ops.mlx.custom_sdpa( + q, + k, + v, + start_pos=start_pos, + is_causal=True, + scale=self.head_dim**-0.5, + ) output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1) + + # Apply headwise output gate (Spark-X2.5: per-head sigmoid gate) + if self.headwise_attn_output_gate: + output_4d = output.view(bsz, seqlen, self.n_local_heads, self.head_dim) + og = self.og(x).unsqueeze(-1).to(output_4d.dtype) + output_4d = torch.sigmoid(og) * output_4d + output = output_4d.reshape(bsz, seqlen, -1) + return self.wo(output), None diff --git a/examples/models/spark_x2_5/README.md b/examples/models/spark_x2_5/README.md index cdbde95ddc6..ac05aac37af 100644 --- a/examples/models/spark_x2_5/README.md +++ b/examples/models/spark_x2_5/README.md @@ -13,8 +13,29 @@ Spark-X2.5 uses the same export pipeline as the optimized Llama model. Please see the [Llama README](../llama/README.md) for general runner and mobile-app details. +Pre-exported ExecuTorch MLX artifacts for Spark-X2.5 1.7B and 4B are available +on the [Hugging Face Hub](https://huggingface.co/collections/XHToken/spark-x25). + ### Example export +Export Spark-X2.5-1.7B to MLX on Apple Silicon, quantized with 4-bit weights: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml \ + +base.model_class="spark_x2_5_1_7b" \ + +base.params="examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json" \ + +export.output_name="spark_x2_5_1_7b_mlx_4w.pte" +``` + +Export Spark-X2.5-4B to MLX on Apple Silicon, quantized with 4-bit weights: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml \ + +base.model_class="spark_x2_5_4b" \ + +base.params="examples/models/spark_x2_5/config/spark_x2_5_4b_config.json" \ + +export.output_name="spark_x2_5_4b_mlx_4w.pte" +``` + Export Spark-X2.5-1.7B to XNNPack, FP32: ``` python -m extension.llm.export.export_llm \ @@ -65,6 +86,17 @@ you are a helpful assistant.<｜end▁of▁sentence｜><｜start▁of▁sentence > The special tokens `<｜start▁of▁sentence｜>` and `<｜end▁of▁sentence｜>` use the Unicode lower-one-eighth block character (`▁`, U+2581). +For MLX on Apple Silicon, build or install ExecuTorch with MLX enabled. The +easiest local path is: +``` +conda activate +python install_executorch.py +xcrun -sdk macosx --find metal +``` + +The `metal` command must resolve to an Xcode path, not fail under standalone +Command Line Tools. + With ExecuTorch pybindings: ``` python -m examples.models.llama.runner.native \ @@ -80,13 +112,18 @@ you are a helpful assistant.<｜end▁of▁sentence｜><｜start▁of▁sentence --temperature 0.3 ``` -With ExecuTorch's sample C++ runner: +With ExecuTorch pybindings and a Spark-X2.5 MLX export: ``` -cmake-out/examples/models/llama/llama_main \ - --model_path spark_x2_5_1_7b_8da4w.pte \ - --tokenizer_path ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots//tokenizer.json \ +python -m examples.models.llama.runner.native \ + --model spark_x2_5_1_7b \ + --pte spark_x2_5_1_7b_mlx_4w.pte \ + --tokenizer ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots//tokenizer.json \ + --tokenizer_config ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots//tokenizer_config.json \ --prompt="<｜start▁of▁sentence｜><|System|> you are a helpful assistant.<｜end▁of▁sentence｜><｜start▁of▁sentence｜><|User|>Who are you?<｜end▁of▁sentence｜><｜start▁of▁sentence｜><|Bot|>" \ + --params examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json \ + --max_len 128 \ + -kv \ --temperature 0.3 ``` @@ -100,4 +137,29 @@ for path in root.glob("*/tokenizer.json"): PY ``` +With ExecuTorch's sample C++ runner: +``` +cmake-out/examples/models/llama/llama_main \ + --model_path spark_x2_5_1_7b_8da4w.pte \ + --tokenizer_path ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots//tokenizer.json \ + --prompt="<｜start▁of▁sentence｜><|System|> +you are a helpful assistant.<｜end▁of▁sentence｜><｜start▁of▁sentence｜><|User|>Who are you?<｜end▁of▁sentence｜><｜start▁of▁sentence｜><|Bot|>" \ + --temperature 0.3 +``` + +Build the C++ runner with MLX support for Spark-X2.5: +``` +make spark_x2_5-mlx +``` + +Then run a Spark-X2.5 MLX export with the C++ runner: +``` +cmake-out/examples/models/llama/llama_main \ + --model_path spark_x2_5_1_7b_mlx_4w.pte \ + --tokenizer_path ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots//tokenizer.json \ + --prompt="<｜start▁of▁sentence｜><|System|> +you are a helpful assistant.<｜end▁of▁sentence｜><｜start▁of▁sentence｜><|User|>Who are you?<｜end▁of▁sentence｜><｜start▁of▁sentence｜><|Bot|>" \ + --temperature 0.3 +``` + To run the model on an example iOS or Android app, see the Llama README's [Step 5: Build Mobile apps](../llama/README.md#step-5-build-mobile-apps) section. diff --git a/examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml b/examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml new file mode 100644 index 00000000000..24122e9ad02 --- /dev/null +++ b/examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml @@ -0,0 +1,15 @@ +base: + metadata: '{"get_bos_id": 0, "get_eos_ids":[1]}' + +model: + use_kv_cache: True + use_sdpa_with_kv_cache: True + dtype_override: bf16 + +quantization: + qmode: 4w + group_size: 64 + +backend: + mlx: + enabled: True diff --git a/examples/models/spark_x2_5/test_spark_x2_5_mlx.py b/examples/models/spark_x2_5/test_spark_x2_5_mlx.py new file mode 100644 index 00000000000..f31bda028f3 --- /dev/null +++ b/examples/models/spark_x2_5/test_spark_x2_5_mlx.py @@ -0,0 +1,244 @@ +""" +Numerical tests for Spark-X2.5 MLX attention. + +Tests that MLXAttentionMHA produces the same output as eager AttentionMHA +for Spark-X2.5 configurations, including: +- Sliding window attention with ring buffer cache +- Headwise output gate +- Per-layer RoPE parameters (partial rotary factor) +""" + +import torch +from executorch.backends.mlx.llm.et_attention import MLXAttentionMHA +from executorch.examples.models.llama.attention import AttentionMHA +from executorch.examples.models.llama.model_args import ModelArgs +from executorch.examples.models.llama.rope import Rope + + +def _create_tiny_model_args( + layer_types: list[str], + sliding_window: int = 4, + headwise_attn_output_gate: bool = True, +) -> ModelArgs: + """Create tiny ModelArgs for testing Spark-X2.5 attention.""" + return ModelArgs( + dim=64, + n_layers=len(layer_types), + n_heads=4, + n_kv_heads=2, + head_dim=16, + hidden_dim=128, + vocab_size=100, + max_batch_size=1, + max_seq_len=32, + max_context_len=32, + use_kv_cache=True, + use_sdpa_with_kv_cache=True, + norm_eps=1e-5, + rope_theta=10000.0, + use_hf_rope=True, + use_scaled_rope=False, + sliding_window=sliding_window, + layer_types=layer_types, + headwise_attn_output_gate=headwise_attn_output_gate, + attention_qkv_bias=False, + use_qk_norm=False, + qk_norm_before_rope=False, + act_fn="gelu", + attention_type="mlx", + # Per-layer RoPE parameters + rope_parameters={ + "full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25}, + "sliding_attention": {"rope_theta": 10000, "partial_rotary_factor": 1.0}, + }, + ) + + +def _create_eager_attention( + args: ModelArgs, layer_id: int, layer_type: str +) -> AttentionMHA: + """Create an eager AttentionMHA for a specific layer.""" + # Temporarily set attention_type to default for eager model + args_copy = ModelArgs(**{**args.__dict__, "attention_type": "default"}) + + # Create rope params with layer-specific parameters + rope_params = args.rope_parameters.get(layer_type, {}) + rope_args_dict = {**args.__dict__} + # Set both rope_theta and rope_freq_base to ensure __post_init__ doesn't override + layer_rope_theta = rope_params.get("rope_theta", args.rope_theta) + rope_args_dict["rope_theta"] = layer_rope_theta + rope_args_dict["rope_freq_base"] = layer_rope_theta + rope_args_dict["partial_rotary_factor"] = rope_params.get( + "partial_rotary_factor", 1.0 + ) + rope_args = ModelArgs(**rope_args_dict) + rope = Rope(rope_args) + + return AttentionMHA(args_copy, layer_id, rope) + + +def test_mlx_attention_matches_eager_sliding_window() -> None: + """Test that MLX attention matches eager attention for sliding window layers.""" + torch.manual_seed(42) + + # Create a tiny model with sliding window + layer_types = ["sliding_attention", "full_attention"] + args = _create_tiny_model_args( + layer_types=layer_types, + sliding_window=4, + headwise_attn_output_gate=True, + ) + + # Test sliding attention layer (layer 0) + eager_attn = _create_eager_attention( + args, layer_id=0, layer_type="sliding_attention" + ) + mlx_attn = MLXAttentionMHA.from_attention_mha(eager_attn, dtype=torch.float32) + + # Verify sliding window is correctly set + assert mlx_attn.is_sliding is True + assert mlx_attn.sliding_window == 4 + assert mlx_attn.headwise_attn_output_gate is True + + # Create input + batch_size = 1 + seq_len = 2 + x = torch.randn(batch_size, seq_len, args.dim) + + # Create position indices + input_pos = torch.tensor([0, 1]) + + # Generate RoPE frequencies for eager attention + freqs_cos, freqs_sin = eager_attn.rope.get_freqs(input_pos, seq_len) + + # Run eager attention + eager_out, _ = eager_attn( + x, freqs_cos=freqs_cos, freqs_sin=freqs_sin, input_pos=input_pos + ) + + # Run MLX attention (uses fused rope, doesn't need freqs) + mlx_out, _ = mlx_attn(x, freqs_cos=None, freqs_sin=None, input_pos=input_pos) + + # Compare outputs (should match within numerical precision) + torch.testing.assert_close(eager_out, mlx_out, rtol=1e-4, atol=1e-5) + + +def test_mlx_attention_matches_eager_past_window() -> None: + """Test that MLX attention correctly handles positions past the sliding window.""" + torch.manual_seed(42) + + layer_types = ["sliding_attention"] + args = _create_tiny_model_args( + layer_types=layer_types, + sliding_window=4, + headwise_attn_output_gate=True, + ) + + eager_attn = _create_eager_attention( + args, layer_id=0, layer_type="sliding_attention" + ) + mlx_attn = MLXAttentionMHA.from_attention_mha(eager_attn, dtype=torch.float32) + + # Decode past the window size + batch_size = 1 + seq_len = 1 + + # Prefill positions 0-5 to populate KV cache + for pos in range(6): + input_pos = torch.tensor([pos]) + x = torch.randn(batch_size, seq_len, args.dim) + # Generate RoPE frequencies for eager attention + freqs_cos, freqs_sin = eager_attn.rope.get_freqs(input_pos, seq_len) + eager_out, _ = eager_attn( + x, freqs_cos=freqs_cos, freqs_sin=freqs_sin, input_pos=input_pos + ) + mlx_out, _ = mlx_attn(x, freqs_cos=None, freqs_sin=None, input_pos=input_pos) + # Verify each step matches + torch.testing.assert_close(eager_out, mlx_out, rtol=1e-4, atol=1e-5) + + # Now test at position 6 (past the window of 4) + input_pos = torch.tensor([6]) + x = torch.randn(batch_size, seq_len, args.dim) + + # Generate RoPE frequencies for eager attention + freqs_cos, freqs_sin = eager_attn.rope.get_freqs(input_pos, seq_len) + + # Run both + eager_out, _ = eager_attn( + x, freqs_cos=freqs_cos, freqs_sin=freqs_sin, input_pos=input_pos + ) + mlx_out, _ = mlx_attn(x, freqs_cos=None, freqs_sin=None, input_pos=input_pos) + + # Should still match + torch.testing.assert_close(eager_out, mlx_out, rtol=1e-4, atol=1e-5) + + +def test_mlx_attention_full_attention_layer() -> None: + """Test MLX attention for full attention layer with partial RoPE.""" + torch.manual_seed(42) + + layer_types = ["full_attention"] + args = _create_tiny_model_args( + layer_types=layer_types, + sliding_window=4, + headwise_attn_output_gate=True, + ) + + eager_attn = _create_eager_attention(args, layer_id=0, layer_type="full_attention") + mlx_attn = MLXAttentionMHA.from_attention_mha(eager_attn, dtype=torch.float32) + + # Full attention should not be sliding + assert mlx_attn.is_sliding is False + assert mlx_attn.sliding_window is None + + # Verify partial rotary factor is set correctly + assert mlx_attn.rope_dims == int(args.head_dim * 0.25) # partial_rotary_factor=0.25 + + # Test forward pass + batch_size = 1 + seq_len = 2 + x = torch.randn(batch_size, seq_len, args.dim) + input_pos = torch.tensor([0, 1]) + + # Generate RoPE frequencies for eager attention + freqs_cos, freqs_sin = eager_attn.rope.get_freqs(input_pos, seq_len) + + eager_out, _ = eager_attn( + x, freqs_cos=freqs_cos, freqs_sin=freqs_sin, input_pos=input_pos + ) + mlx_out, _ = mlx_attn(x, freqs_cos=None, freqs_sin=None, input_pos=input_pos) + + torch.testing.assert_close(eager_out, mlx_out, rtol=1e-4, atol=1e-5) + + +def test_mlx_attention_unsupported_options() -> None: + """Test that unsupported options raise errors.""" + torch.manual_seed(42) + + layer_types = ["sliding_attention"] + args = _create_tiny_model_args( + layer_types=layer_types, + sliding_window=4, + headwise_attn_output_gate=True, + ) + + eager_attn = _create_eager_attention( + args, layer_id=0, layer_type="sliding_attention" + ) + + # Test use_attn_o_norm (unsupported) + eager_attn.use_attn_o_norm = True + try: + MLXAttentionMHA.from_attention_mha(eager_attn) + assert False, "Should have raised NotImplementedError for use_attn_o_norm" + except NotImplementedError as e: + assert "use_attn_o_norm" in str(e) + + # Reset and test scale_query_by (unsupported) + eager_attn.use_attn_o_norm = False + eager_attn.scale_query_by = 2.0 + try: + MLXAttentionMHA.from_attention_mha(eager_attn) + assert False, "Should have raised NotImplementedError for scale_query_by" + except NotImplementedError as e: + assert "scale_query_by" in str(e)