Enable Spark-X2.5 export and runner support for MLX

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 <target>\`. Available targets:"
@@ -125,19 +125,22 @@ 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"
 	@echo "  gemma4_31b-cuda     - Build Gemma 4 31B runner and worker with CUDA backend"
 	@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

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,20 +125,34 @@ 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
         self.use_fused_rope = self._can_use_fused_rope(rope.params)
         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,13 +210,31 @@ 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
         instance.wv = other.wv
         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