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

[dongjiang1989]

Summary

Add MLX backend support for Spark-X2.5 models (1.7B and 4B variants), enabling efficient inference on Apple Silicon.

• Add MLX 4-bit export config for Spark-X2.5 with KV cache metadata
• Add make spark_x2_5-mlx to build the shared Llama C++ runner with MLX enabled
• Add README instructions for export and C++/pybindings runs
• Add test coverage for Spark-X2.5 model registration and architecture configs
• Support both 1.7B and 4B model variants with hybrid attention architecture

Spark-X2.5 is a compact, efficient language model series with hybrid attention (3 sliding-window + 1 full-attention layer), supporting context windows up to 1M tokens and 200+ languages.

Test plan

• pytest -q examples/models/spark_x2_5/test_spark_x2_5_mlx.py
• make -n spark_x2_5-mlx
• make spark_x2_5-mlx
• Exported and runtime-loaded:
  • spark_x2_5_1_7b_mlx_4w.pte
  • spark_x2_5_4b_mlx_4w.pte
• Ran pybindings smoke generation for both 1.7B and 4B
• Ran C++ llama_main smoke generation with MLX exports

cc @mergennachin @iseeyuan @lucylq @helunwencser @tarun292 @kimishpatel @jackzhxng @nil-is-all @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23523

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 24 Awaiting Approval

As of commit 1221f15 with merge base 40d17e1 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Apple (gh)
• Build Aarch64 Linux CUDA Wheels (gh)
• Build Aarch64 Linux Wheels (gh)
• Build documentation (gh)
• Build Linux CUDA Wheels (gh)
• Build Linux Wheels (gh)
• Build macOS Wheels (gh)
• Build Presets (gh)
• Build Windows Wheels (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• MLX (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test ROCm AOTI (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[dongjiang1989]

cc @nil-is-all PTAL, thanks

[nil-is-all]

@dongjiang1989 could you rebase and ping for review? A branch update with main should resolve CI failures

[nil-is-all]

README and .yaml look good and simple enough to run. Hyperparams in test python file match 7B and 4B models.

[dongjiang1989]

@dongjiang1989 could you rebase and ping for review? A branch update with main should resolve CI failures

Rebase and updated.
@nil-is-all Please re-check it

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 4m 49s —— View job

---

This needs changes: the MLX attention swap drops Spark-X2.5's headwise output gate, so MLX exports give wrong output from the first token; it also ignores the 512-token sliding window, and the EOS/BOS metadata appears to have been copied from LFM2.

Detailed review

PR Review: #23523

Summary

This PR adds an MLX 4-bit export config, a make spark_x2_5-mlx target, README steps and config tests for Spark-X2.5. The MLX export path in this PR is not numerically correct for this architecture, and the new YAML's token metadata looks wrong.

Code Quality

• headwise_attn_output_gate is dropped under MLX, so output is wrong at every token. With backend.mlx.enabled, export_llama_lib.py:2271-2281 replaces every AttentionMHA with MLXAttentionMHA.from_attention_mha (backends/mlx/llm/et_attention.py:116-179).
  • That function copies wq/wk/wv/wo, rope and qk-norm. It does not copy og.
  • Its forward ends with return self.wo(output), None (et_attention.py:252). It never applies the per-head sigmoid(og(x)) gate that AttentionMHA._apply_output_transforms applies (examples/models/llama/attention.py:707-711).
  • Spark-X2.5 sets "headwise_attn_output_gate": true for all layers (spark_x2_5_1_7b_config.json:18). Every attention layer's output is therefore wrong, and the og weights are never used.
  • A smoke run can still produce readable text, so "it generated something" doesn't show the export is correct. Compare logits against eager or XNNPACK.
• The sliding window is ignored under MLX.
  • from_attention_mha never reads is_sliding/sliding_window. It gives every layer a full-length KVCache (et_attention.py:170) and calls custom_sdpa(..., is_causal=True) with no window mask (et_attention.py:242-248).
  • The 21 of 28 sliding layers (27 of 36 for 4B) therefore attend to the full history. Outputs diverge from eager once the position reaches 512, i.e. whenever max_context_length > 512.
  • The defaults (128) hide this. The "up to 1M tokens" framing in the description and README suggests long contexts will be used.
  • examples/models/gemma4_31b/mlx_source_transformations.py:56-100 already shows how to do this on MLX: RingBufferKVCache plus create_sliding_window_mask.
  • Suggested fix: teach MLXAttentionMHA.from_attention_mha to carry over og and apply the headwise gate, and to use a ring-buffer cache with a window mask for is_sliding layers. Have it fail loudly on AttentionMHA options it doesn't support (use_attn_o_gate, use_attn_o_norm, use_q_gate, scale_query_by, …) instead of silently dropping them. Until then, the README shouldn't present MLX as a supported path for Spark-X2.5.
    Fix this →
• The YAML's BOS/EOS metadata doesn't match the other Spark-X2.5 configs. config/spark_x2_5_mlx_4w.yaml:2 sets {"get_bos_id": 1, "get_eos_ids":[7]}, which is byte-for-byte LFM2's value (examples/models/lfm2/config/lfm2_mlx_4w.yaml:2). The existing Spark-X2.5 configs use {"get_bos_id": 0, "get_eos_ids":[1]} (spark_x2_5_xnnpack_q8da4w.yaml:2).
  • llama_main and the pybindings runner stop on these IDs. With the wrong EOS, generation either runs on to max_len or stops on an unrelated token 7.
  • Please confirm the IDs against the Spark-X2.5 tokenizer and use the same values as the XNNPACK configs.
• README.md:16-17 promises artifacts that "will be available on the Hugging Face Hub." Please add the link once they exist, or drop the sentence.
• The PR deletes the pointer to the mobile-app instructions ("To run the model on an example iOS or Android app…"), apparently by accident.
• README.md:89-98 hard-codes a personal env name (conda activate et-mlx). Point to the repo's MLX setup docs instead.
• The spark_x2_5-mlx Makefile target (Makefile:433-440) is identical to lfm_2_5-mlx. Both just build the generic llama_main with the llama-mlx preset. Adding a new alias for each model makes the help list longer without adding anything. Consider adding a single generic llama-mlx target and pointing the README to it.

Testing

• The new test file mostly repeats an existing one, and it doesn't test anything MLX-specific. test_spark_x2_5_mlx.py copies test_spark_x2_5.py:

  • The helpers are identical.
  • Two of its three tests already exist in test_spark_x2_5.py, which also checks head_dim and act_fn, so the new copies check less.
  • Only the YAML assertion is new, and it just restates the YAML file line by line. It would also have kept passing with the wrong BOS/EOS above.

  Move the one YAML test into test_spark_x2_5.py. More importantly, add a numerical test that compares eager AttentionMHA with the MLX-transformed attention on a tiny config that has headwise_attn_output_gate=True, sliding_window and layer_types. That test would have caught both correctness bugs above.

Recommendation

Request Changes

The MLX export is numerically wrong for this architecture: the output gate is missing at every layer, and sliding layers ignore the window past 512 tokens. The EOS/BOS metadata also appears to be copied from LFM2. Smoke generation doesn't catch either problem, so please validate the fix against eager logits.

[dongjiang1989]

@nil-is-all @metascroy Fix claude review. Please re-check it, thanks

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 15s —— View job

---

The headwise output gate is fixed now, but this still needs changes. On the real export path the sliding-window fix never takes effect, so Spark-X2.5 MLX exports still attend to the full history past 512 tokens. Other AttentionMHA options are still dropped without an error, and nothing tests the new attention code.

Detailed review

PR Review: #23523

Summary

The new commit adds three things to MLXAttentionMHA: the headwise output gate, a ring-buffer KV cache for sliding layers, and a per-layer partial-RoPE path. The gate now matches eager (et_attention.py:390-394 vs attention.py:707-711), and the BOS/EOS metadata, env name and Makefile duplication are fixed. Three problems remain:

• The sliding-window fix is dead on the export path.
• Options the class can't handle are still dropped silently.
• The new partial-RoPE branch is dead code.

Code Quality

• Sliding layers still get a full-length cache and a plain causal mask on the export path.
  • export_llama_lib goes through replace_mlx_attention → MLXAttentionMHA.from_attention_mha (backends/mlx/llm/source_transformation.py:306).
  • That function reads the window with getattr(other, "sliding_window", None) (et_attention.py:214). AttentionMHA never stores a sliding_window attribute. It only reads args.sliding_window inside __init__ (attention.py:461, :528), so the result is always None.
  • As a result, instance.is_sliding and instance.sliding_window is falsy, every layer gets a full KVCache, and forward uses is_causal=True without the window mask.
  • The effect is the same wrong output past 512 tokens as before this commit. The getattr fallback is why it fails silently (CLAUDE.md also asks to avoid getattr patterns like this).
  • Fix: take the window from the eager ring cache (other.kv_cache.window_size when isinstance(other.kv_cache, RingKVCache)), or store self.sliding_window on AttentionMHA. Read it as a plain attribute so a missing value raises.
  • When you fix this, also size the ring for prefill. RingBufferKVCache defaults max_write_len to the window (cache.py:332). With an int start_pos, nothing checks that a prefill chunk longer than 512 tokens fits. Gemma 4 passes max_write_len=min(max_seq_len, sliding_window) (gemma4_31b/mlx_source_transformations.py:196-207).
    Fix this →
• use_attn_o_gate, use_attn_o_norm and scale_query_by are still dropped without an error.
  • use_attn_o_gate and use_attn_o_norm are copied (et_attention.py:209-210), and og is even shared for use_attn_o_gate (:250-251). But forward only applies the headwise gate.
  • o_norm is never copied, and scale_query_by != 1.0 is never checked.
  • Any model using these options exports with wrong output and no error.
  • Fix: either implement them by reusing AttentionMHA._apply_output_transforms and the scale_query_by multiply, or add them to _unsupported (:176-184).
• The partial-RoPE branch never runs.
  • It depends on hasattr(rope, "inv_freq") (et_attention.py:134, :239-243). The llama Rope has no inv_freq attribute (examples/models/llama/rope.py:292-342). That pattern was copied from Gemma 4, whose rope has one.
  • So use_partial_rope is always False, and the mlx_freqs branch (:309-330) is dead. It also hard-codes traditional=False instead of self.rope_traditional.
  • The existing path already handles Spark-X2.5 correctly. construct_transformer passes each layer its own per-type Rope (llama_transformer.py:584-591), so rope_freq_base and partial_rotary_factor are already per layer. mlx.rope(dims=rope_dims, base=…) produces the same frequencies as HF partial rotary.
  • Remove the use_partial_rope/rotary_inv_freq state, the branch in forward, and the rope_parameters override in __init__ (:125-136), which repeats what the per-layer rope already provides.
• sdpa_start = max(0, self.kv_cache.buffer_size - seqlen) (et_attention.py:366) calls Python max on a value that is symbolic under dynamic seqlen, which can add an export guard. seqlen <= max_write_len <= window already keeps it non-negative. Use self.kv_cache.buffer_size - seqlen as Gemma 4 does (gemma4_31b/mlx_source_transformations.py:86).
• The README still deletes the pointer to the mobile-app instructions ("To run the model on an example iOS or Android app…").
• The .mlx-runner + $(MAKE) MODEL_NAME=… pattern is the only recursive-make use in the Makefile. Both targets still build the same generic llama_main. A plain shared target (or spark_x2_5-mlx: lfm_2_5-mlx) is simpler.

Testing

• None of the ~190 new lines in et_attention.py are tested. The PR's test only restates the YAML, and a smoke run producing text would not have caught the bug above.
  • Add a numerical test that compares eager AttentionMHA with MLXAttentionMHA.from_attention_mha on a tiny config: layer_types with both sliding and full layers, a small sliding_window (e.g. 4), headwise_attn_output_gate=True and rope_parameters with partial_rotary_factor<1.
  • Decode past the window. The mlx:: ops have eager reference implementations, so this runs without exporting.
  • Move the remaining YAML test into test_spark_x2_5.py rather than keeping a separate 17-line file.

Recommendation

Request Changes

On the actual export path, sliding layers still ignore the window: sliding_window is always None in from_attention_mha, so the main correctness fix does nothing. Please fix that and either raise on or implement the remaining unsupported options. Then add an eager-vs-MLX numerical test that covers positions past the window.

[dongjiang1989]

@claude review this code

[dongjiang1989]

Thx @nil-is-all Fixed. Please review it 😄

[nil-is-all]

Thx @nil-is-all Fixed. Please review it 😄

Thx for the ping. Could you update the out-of-date branch?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 14s —— View job

---

This still needs changes: backends/mlx/llm/et_attention.py has unresolved merge-conflict markers, so the module won't import, MLX export breaks for every model, and the new tests can't have run.

Detailed review

PR Review: #23523

Summary

This revision adds checks that reject unsupported options, takes the sliding window from the eager RingKVCache (which now works on the export path), removes the dead partial-RoPE branch, and replaces the YAML test with eager-vs-MLX attention tests. The file doesn't parse. The tests would fail even once it does, and use_attn_o_gate is still dropped without an error.

Code Quality

• Merge-conflict markers are committed in et_attention.py:335-342.
  • The file contains <<<<<<< HEAD … ======= … >>>>>>> 5810d8145, which is a SyntaxError.
  • replace_mlx_attention imports this module, so export_llm with backend.mlx.enabled now fails for every model, including LFM2, Qwen and Llama, not only Spark-X2.5.
  • test_spark_x2_5_mlx.py imports it too, so the test plan in the description can't have passed on this commit.
  • Keep the HEAD side (sdpa_start = self.kv_cache.buffer_size - seqlen) and delete the markers and the max(0, …) branch.
    Fix this →
• use_attn_o_gate is still dropped without an error.
  • from_attention_mha shares og when use_attn_o_gate is set (et_attention.py:243-244), but forward only applies headwise_attn_output_gate (:366-370).
  • The og weight is therefore carried along and never used, and the output differs from eager _apply_output_transforms (attention.py:703-705).
  • Either implement it (output * sigmoid(og(x)) per element, before wo) or add it to _unsupported and to the __init__ checks.
• The sliding-window ring cache is never sized for prefill.
  • Both RingBufferKVCache(...) calls (et_attention.py:136, :252) use the default max_write_len = window.
  • forward passes an int start_pos (from .item()), so RingBufferKVCache.update skips torch._check(seq_len <= self.max_write_len) (cache.py:365-372).
  • If an export uses max_seq_length > sliding_window (512), a long prefill chunk silently overwrites in-window K/V.
  • Raise at construction when max_seq_len > sliding_window, or pass max_write_len=min(max_seq_len, sliding_window) as Gemma 4 does.
• The new code still uses getattr/hasattr for fields that always exist (CLAUDE.md asks to avoid this).
  • Examples: getattr(args, "use_q_gate", False), getattr(args, "headwise_attn_output_gate", False), getattr(other, "is_sliding", False), getattr(other, "scale_query_by", 1.0) (et_attention.py:62-104, :172-217).
  • ModelArgs and AttentionMHA always define these fields (attention.py:417-431). The defaults would hide a rename the same way getattr(other, "sliding_window", None) hid the bug in the last round.
  • Use plain attribute access. For the window, use isinstance(other.kv_cache, RingKVCache) instead of hasattr(other.kv_cache, "window_size").
• is_kv_shared_layer is rejected in from_attention_mha but not in __init__, which contradicts the "must match from_attention_mha" comment at :61.
• The .mlx-runner + recursive $(MAKE) pattern is still the only one of its kind in the Makefile. spark_x2_5-mlx: lfm_2_5-mlx, or a plain shared target, does the same thing.

Testing

• The eager half of every numerical test fails.
  • The tests call eager_attn(x, freqs_cos=None, freqs_sin=None, ...).
  • Eager AttentionMHA._prepare_qkv calls self.rope.forward(q, k, freqs_cos, freqs_sin) (attention.py:596), which needs real tensors. MLX only avoids this because it always uses fused rope.
  • Pass freqs_cos, freqs_sin = rope.get_freqs(input_pos, seq_len) to both modules.
• test_mlx_attention_full_attention_layer doesn't use the per-layer rope_theta it claims to.
  • _create_eager_attention sets rope_freq_base from rope_parameters, but the copied dict still has rope_theta=10000.0.
  • ModelArgs.__post_init__ then overwrites rope_freq_base with it (model_args.py:207-208), so the "full attention" layer runs with base 10000, not 5000000.
  • Build the per-layer rope the way construct_transformer does, or override rope_theta as well.
• test_mlx_attention_matches_eager_past_window only decodes one token at a time. Add a multi-token prefill that crosses the window (e.g. 3 tokens at start_pos=3, window 4) to exercise buffer_size - seqlen and the mask.
• Use pytest.raises(NotImplementedError, match=...) instead of try/assert False/except. Add a use_attn_o_gate case once it's handled.
• Nothing in CI appears to run these tests: examples/models/spark_x2_5 isn't wired into mlx.yml. Confirm they run somewhere that has the mlx:: reference ops.

Recommendation

Request Changes

The conflict markers break MLX export for every model and must be removed. After that, fix the tests so they actually run (pass real rope freqs and the correct per-layer rope_theta), and either implement or reject use_attn_o_gate. Since the branch is also out of date (per @nil-is-all), please rebase and run pytest examples/models/spark_x2_5/test_spark_x2_5_mlx.py locally before re-requesting review.