Add Spark2.5 model support

[simple-zheng]

Summary

This PR adds native Spark2.5 inference support to MLX LM.

The implementation supports checkpoints with:

• model_type: spark2_5
• Fused QKV projection
• Grouped-query attention
• Head-wise sigmoid attention output gating
• Mixed sliding-window and full-attention layers
• Layer-specific partial RoPE dimensions and frequency bases
• Parallel GELU feed-forward layers
• Tied and untied output embeddings
• Standard and rotating KV caches

Implementation

A new model implementation has been added at:

mlx_lm/models/spark2_5.py

Since the filename matches the checkpoint's model_type, Spark2.5 checkpoints are discovered through the existing MLX LM model-loading mechanism without additional loader registration.

The implementation includes:

• Configuration validation for layer types, attention heads, RoPE parameters, activation functions, and attention gates
• Separate masks and cache implementations for sliding-window and full-attention layers
• Weight sanitization for tied embeddings and serialized RoPE frequency tensors
• Materialization of fused QKV views before RoPE and attention for backend compatibility

Supported Checkpoint

The implementation was verified with a local safetensors checkpoint of:

Spark-X2.5-1.7B

The tested checkpoint uses 28 decoder layers, mixed sliding/full attention, a vocabulary size of 131,072, and tied input/output embeddings.

Test Coverage

A Spark2.5 model test has been added to tests/test_models.py.

The test covers:

• Model construction from a minimal configuration
• Forward-pass output shape and dtype
• FP16 and FP32 execution
• Batch sizes greater than one
• Cached generation
• Model copying
• RotatingKVCache for sliding-attention layers
• KVCache for full-attention layers

Targeted test command:

python -m unittest discover \
  -s tests \
  -p 'test_models.py' \
  -k spark2_5

Result: passed.

Formatting check:

pre-commit run --files \
  mlx_lm/models/spark2_5.py \
  tests/test_models.py

Result: passed.

Real Checkpoint Verification

Command:

mlx_lm.generate \
  --model <local-spark2.5-checkpoint> \
  -p "只输出53乘以42的结果，不要解释。" \
  -m 64 \
  --temp 0

Output:

53乘以42的结果是2226。

Prompt: 24 tokens, 37.214 tokens-per-sec
Generation: 13 tokens, 8.128 tokens-per-sec
Peak memory: 5.608 GB

The checkpoint loaded successfully, generated the expected deterministic answer, and stopped normally.

Scope and Limitations

• This PR adds inference support for the Spark2.5 architecture.
• Validation was performed with the Spark-X2.5-1.7B checkpoint.
• No model weights are included in this PR.
• LoRA-specific configuration changes are outside the scope of this PR.

AI Usage Disclosure

Codex was used to assist with adapting the standalone Spark2.5 implementation to the MLX LM model interface, adding unit-test coverage.The implementation and verification results were reviewed and validated by the contributor.

[DXXS]

Seems to be working here, integrated with 0.31.3 (though would work better if my mlx_lm had rotating kv support/4bit!).

[DXXS]

After further checking of mixed quant conversion, GLM-5.3 indicates that there is a "spark2_5.py:63 sanity check that wrongly assumes hidden_size == n_heads * head_dim" and recommends either just removing it (checked fine on my system), or perhaps replacing the check instead with some sort of check of kv width vs q width.

[michalk8]

superseded by #1891