feat(example): Add Spark-X2.5 model example with hybrid attention support (pytorch#22865)

## Summary

Add `examples/models/spark_x2_5/` with support for the [Spark-X2.5-1.7B
and Spark-X2.5-4B
models](https://huggingface.co/collections/XHToken/spark-x25) from
XHToken. These are compact, general-purpose language models with a
hybrid attention architecture (3 sliding-window attention layers + 1
full-attention layer, repeating) that natively supports context windows
up to 1M tokens.

### Architecture highlights
- **Hybrid attention**: pattern of 3 sliding-window + 1 full-attention
layer (28 layers for 1.7B, 36 for 4B)
- **Per-layer-type RoPE**: full-attention layers use `rope_theta=5M,
partial_rotary_factor=0.25`; sliding-attention layers use
`rope_theta=10K, partial_rotary_factor=1.0`
- **Headwise attention output gate**: per-head sigmoid gate broadcast
over head dim (a lighter variant of `use_attn_o_gate`)
- **Sliding window KV cache** (`RingKVCache`, window=512 tokens)
- **GELU activation** in MLP (vs default SiLU)
- **Tied word embeddings**

### Files added (in `examples/models/spark_x2_5/`)
- `convert_weights.py` — HF safetensors → Meta format conversion with
fused QKV split, sharded checkpoint support, and tied embeddings
handling
- `config/` — JSON configs for 1.7B and 4B variants; XNNPack (fp32,
q8da4w), CoreML, and MLX export configs
- `test_spark_x2_5.py` — config validation and model registration tests
- `BUCK`, `README.md` — build target and usage documentation

### Core changes to existing files
- `examples/models/llama/model_args.py` — add
`headwise_attn_output_gate` and `rope_parameters` fields
- `examples/models/llama/feed_forward.py` — `FeedForward` now accepts
`act_fn` parameter (was hardcoded to SiLU)
- `examples/models/llama/attention.py` — headwise attention output gate
(dim→n_heads, broadcast over head_dim), per-layer `is_sliding`
detection, `RingKVCache` for sliding-window layers, mutual-exclusion
validation for gate flags
- `examples/models/llama/llama_transformer.py` — per-layer-type RoPE via
`_build_ropes()` helper, `freqs_by_type` dispatch in
`_forward_layers()`, `act_fn` plumbed to FeedForward
- `examples/models/llama/export_llama_lib.py` — register
`spark_x2_5_1_7b` and `spark_x2_5_4b` in `EXECUTORCH_DEFINED_MODELS`,
`HUGGING_FACE_REPO_IDS`, and weight-conversion dispatch
- `extension/llm/export/config/llm_config.py` — add `spark_x2_5_1_7b`
and `spark_x2_5_4b` to `ModelType` enum

### Example export
```
python -m extension.llm.export.export_llm \
  --config examples/models/spark_x2_5/config/spark_x2_5_xnnpack_q8da4w.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_8da4w.pte"
```

## Test plan

### Unit tests
```
pytest examples/models/spark_x2_5/test_spark_x2_5.py -v
# 3 passed
```

### Model construction + forward pass (both variants)
```python
from executorch.examples.models.llama.model_args import ModelArgs
from executorch.examples.models.llama.llama_transformer import construct_transformer

# 1.7B: 28 layers, dim=2048
# 4B:   36 layers, dim=2560
# Both: prefill + multi-step decode with KV cache succeed
# Layer types: sliding→RingKVCache, full→KVCache
# Per-layer RoPE: 2 distinct RoPE instances (full_attention, sliding_attention)
```

### Sharded checkpoint conversion
```
# Spark-X2.5-1.7B uses 2 shards, 4B uses 5 shards
# Verified: 283 keys converted correctly, QKV split matches, tied embeddings preserved
```

### Lint
```
flake8 examples/models/spark_x2_5/ examples/models/llama/model_args.py \
    examples/models/llama/feed_forward.py examples/models/llama/attention.py \
    examples/models/llama/llama_transformer.py examples/models/llama/export_llama_lib.py \
    extension/llm/export/config/llm_config.py --max-line-length=120
# No warnings
```

This PR was authored with AI assistance (Claude Code).

cc @mergennachin @iseeyuan @lucylq @helunwencser @tarun292 @kimishpatel
@jackzhxng @larryliu0820 @cccclai @digantdesai

---------

Signed-off-by: dongjiang1989 <dongjiang1989@126.com>
Co-authored-by: Mergen Nachin <mnachin@meta.com>

Author: dongjiang1989

examples/models/llama/BUCK

@@ -32,6 +32,7 @@ fbcode_target(_kind = runtime.python_library,
         ":transformer_modules",
         "//caffe2:torch",
         "//executorch/examples/models/lfm2:lfm2",
+        "//executorch/examples/models/spark_x2_5:spark_x2_5",
     ],
 )
 

examples/models/llama/attention.py

@@ -418,7 +418,17 @@ def __init__(
         self.enable_dynamic_shape = args.enable_dynamic_shape
         self.scale_query_by = args.scale_query_by
         self.use_attn_o_gate = args.use_attn_o_gate
+        self.headwise_attn_output_gate = args.headwise_attn_output_gate
+        if self.use_attn_o_gate and self.headwise_attn_output_gate:
+            raise ValueError(
+                "use_attn_o_gate and headwise_attn_output_gate are mutually exclusive"
+            )
         self.use_attn_o_norm = args.use_attn_o_norm
+        self.is_sliding = (
+            args.layer_types is not None
+            and layer_id < len(args.layer_types)
+            and args.layer_types[layer_id] == "sliding_attention"
+        )
         q_out_dim = self.n_heads * self.head_dim * (2 if self.use_q_gate else 1)
 
         # YOCO: Determine if this is a KV shared layer (receives shared KV from donor).
@@ -447,6 +457,10 @@ def __init__(
                 device="cpu",
             )
         )
+        # Sliding-window layers: additionally mask out positions outside the window.
+        if self.is_sliding and args.sliding_window:
+            window = args.sliding_window
+            causal_mask = causal_mask.triu(diagonal=1 - window)
         self.register_buffer("mask", causal_mask, persistent=False)
 
         if self.use_kv_cache:
@@ -481,6 +495,8 @@ def _init_norms(self, args: ModelArgs) -> None:
             self.o_norm = ScalelessRMSNorm(self.head_dim, eps=args.norm_eps)
         if self.use_attn_o_gate:
             self.og = nn.Linear(args.dim, self.n_heads * self.head_dim, bias=False)
+        if self.headwise_attn_output_gate:
+            self.og = nn.Linear(args.dim, self.n_local_heads, bias=False)
 
     def _init_projections(self, args: ModelArgs, q_out_dim: int) -> None:
         """Initialize Q/K/V/O projection layers."""
@@ -509,13 +525,22 @@ def _init_projections(self, args: ModelArgs, q_out_dim: int) -> None:
     def _init_kv_cache(self, args: ModelArgs) -> None:
         """Initialize KV cache (only for non-shared layers)."""
         if self.has_kv_weights:
-            self.kv_cache = KVCache(
-                args.max_batch_size,
-                args.max_context_len,
-                self.n_kv_heads,
-                self.head_dim,
-                args.enable_dynamic_shape,
-            )
+            if self.is_sliding and args.sliding_window:
+                self.kv_cache = RingKVCache(
+                    args.max_batch_size,
+                    args.sliding_window,
+                    self.n_kv_heads,
+                    self.head_dim,
+                    args.enable_dynamic_shape,
+                )
+            else:
+                self.kv_cache = KVCache(
+                    args.max_batch_size,
+                    args.max_context_len,
+                    self.n_kv_heads,
+                    self.head_dim,
+                    args.enable_dynamic_shape,
+                )
         else:
             self.kv_cache = None
 
@@ -679,6 +704,11 @@ def _apply_output_transforms(
                 og = self.og(x).view(bsz, seqlen, self.n_local_heads, self.head_dim)
                 output_4d = torch.sigmoid(og) * output_4d
             output = output_4d.reshape(bsz, seqlen, -1)
+        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)
         if gate is not None:
             output = output * torch.sigmoid(gate)
         return output

examples/models/llama/export_llama_lib.py

@@ -121,6 +121,8 @@
     "lfm2_1_2b",  # hybrid
     "lfm2_5_350m",  # hybrid
     "lfm2_5_1_2b",  # hybrid
+    "spark_x2_5_1_7b",  # hybrid
+    "spark_x2_5_4b",  # hybrid
 ]
 TORCHTUNE_DEFINED_MODELS = ["llama3_2_vision"]
 HUGGING_FACE_REPO_IDS = {
@@ -141,6 +143,8 @@
     "lfm2_1_2b": "LiquidAI/LFM2-1.2B",
     "lfm2_5_350m": "LiquidAI/LFM2.5-350M",
     "lfm2_5_1_2b": "LiquidAI/LFM2.5-1.2B-Instruct",
+    "spark_x2_5_1_7b": "XHToken/Spark-X2.5-1.7B",
+    "spark_x2_5_4b": "XHToken/Spark-X2.5-4B",
 }
 
 
@@ -718,6 +722,8 @@ def export_llama(  # noqa: C901
             from executorch.examples.models.smollm2 import convert_weights
         elif model_name.startswith("lfm2"):
             from executorch.examples.models.lfm2 import convert_weights
+        elif model_name.startswith("spark_x2_5"):
+            from executorch.examples.models.spark_x2_5 import convert_weights
         else:
             raise ValueError(
                 f"Converting weights to meta format for {model_name} is not yet supported"

examples/models/llama/feed_forward.py

@@ -6,14 +6,15 @@
 
 
 class FeedForward(nn.Module):
-    def __init__(self, dim: int, hidden_dim: int):
+    def __init__(self, dim: int, hidden_dim: int, act_fn=F.silu):
         super().__init__()
         self.w1 = nn.Linear(dim, hidden_dim, bias=False)
         self.w2 = nn.Linear(hidden_dim, dim, bias=False)
         self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+        self.act_fn = act_fn
 
     def forward(self, x):
-        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+        return self.w2(self.act_fn(self.w1(x)) * self.w3(x))
 
 
 class LoRAFeedForward(nn.Module):

examples/models/llama/llama_transformer.py

@@ -215,7 +215,11 @@ def __init__(
         ):
             self.feed_forward = LoRAFeedForward(args.dim, args.hidden_dim, args)
         else:
-            self.feed_forward = FeedForward(dim=args.dim, hidden_dim=args.hidden_dim)
+            self.feed_forward = FeedForward(
+                dim=args.dim,
+                hidden_dim=args.hidden_dim,
+                act_fn=args.act_fn.get_function(),
+            )
 
         if isinstance(self.attention, AttentionSkip):
             self.attention_norm = nn.Identity()
@@ -357,6 +361,7 @@ def __init__(self, params: ModelArgs, layers: nn.ModuleList, rope: Rope):
         self.output_prune_map = params.output_prune_map
         # YOCO (You Only Cache Once) KV sharing configuration.
         self.num_kv_shared_layers = params.num_kv_shared_layers
+        self.layer_types = params.layer_types
 
     def _forward_layers(
         self,
@@ -365,6 +370,7 @@ def _forward_layers(
         freqs_sin: torch.Tensor,
         attn_options_: Dict,
         seqlen: int,
+        freqs_by_type: Optional[Dict[str, Tuple[torch.Tensor, torch.Tensor]]] = None,
     ) -> Tuple[torch.Tensor, Optional[Any]]:
         """Run transformer layers with YOCO KV sharing support."""
         attn_options_update = None
@@ -383,7 +389,14 @@ def _forward_layers(
                 if donor_idx in shared_kv:
                     attn_options_["shared_kv"] = shared_kv[donor_idx]
 
-            h, attn_options_update = layer(h, freqs_cos, freqs_sin, attn_options_)
+            # Per-layer-type RoPE: select freqs based on layer type when available.
+            l_cos, l_sin = freqs_cos, freqs_sin
+            if freqs_by_type is not None and self.layer_types is not None:
+                layer_type = self.layer_types[layer_idx]
+                if layer_type in freqs_by_type:
+                    l_cos, l_sin = freqs_by_type[layer_type]
+
+            h, attn_options_update = layer(h, l_cos, l_sin, attn_options_)
 
             if _is_kv_donor_layer(layer_idx, self.n_layers, self.num_kv_shared_layers):
                 assert (
@@ -425,10 +438,23 @@ def forward(
             attn_options.get("input_pos"), seqlen
         )
 
+        # Compute per-layer-type freqs when per-layer RoPE is configured.
+        freqs_by_type = None
+        if hasattr(self, "ropes"):
+            input_pos = attn_options.get("input_pos")
+            freqs_by_type = {
+                lt: r.get_freqs(input_pos, seqlen) for lt, r in self.ropes.items()
+            }
+
         attn_options_ = attn_options.copy() if attn_options is not None else {}
 
         h, attn_options_update = self._forward_layers(
-            h, freqs_cos, freqs_sin, attn_options_, seqlen
+            h,
+            freqs_cos,
+            freqs_sin,
+            attn_options_,
+            seqlen,
+            freqs_by_type=freqs_by_type,
         )
 
         if not self.generate_full_logits:
@@ -467,11 +493,35 @@ def forward(
         return logits
 
 
+def _build_ropes(model_args: ModelArgs) -> Tuple[Rope, Dict[str, Rope]]:
+    """Build Rope instances, creating per-layer-type ropes when rope_parameters is set.
+
+    Returns (default_rope, ropes_by_type). ropes_by_type is empty when no
+    per-layer-type configuration is provided.
+    """
+    import copy as _copy
+
+    if not model_args.rope_parameters:
+        return Rope(model_args), {}
+
+    ropes: Dict[str, Rope] = {}
+    for layer_type, rope_params in model_args.rope_parameters.items():
+        rope_args = _copy.copy(model_args)
+        if "rope_theta" in rope_params:
+            rope_args.rope_theta = rope_params["rope_theta"]
+            rope_args.rope_freq_base = rope_params["rope_theta"]
+        if "partial_rotary_factor" in rope_params:
+            rope_args.partial_rotary_factor = rope_params["partial_rotary_factor"]
+        ropes[layer_type] = Rope(rope_args)
+    return next(iter(ropes.values())), ropes
+
+
 def construct_transformer(model_args: ModelArgs) -> Transformer:
     """
     Construct a Transformer model from the given model arguments.
     """
-    rope = Rope(model_args)
+    rope, ropes = _build_ropes(model_args)
+
     if model_args.attention_type not in ATTENTION_REGISTRY:
         raise ValueError(
             f"Unknown attention type: {model_args.attention_type}. "
@@ -521,12 +571,20 @@ def construct_transformer(model_args: ModelArgs) -> Transformer:
             )
             layers.append(transformer_block)
         else:
+            # Select per-layer-type RoPE when available.
+            layer_rope = rope
+            if ropes and model_args.layer_types:
+                layer_type = model_args.layer_types[layer_id]
+                layer_rope = ropes.get(layer_type, rope)
             attention = cls(
-                model_args, layer_id, rope, **model_args.attention_kwargs
+                model_args, layer_id, layer_rope, **model_args.attention_kwargs
             )  # pyre-ignore[45]
             transformer_block = TransformerBlock(
                 model_args, attention, layer_id=layer_id
             )
             layers.append(transformer_block)
 
-    return Transformer(model_args, layers, rope)
+    transformer = Transformer(model_args, layers, rope)
+    if ropes:
+        transformer.ropes = torch.nn.ModuleDict(ropes)
+    return transformer

examples/models/llama/model_args.py

@@ -126,6 +126,9 @@ class ModelArgs:
     local_rope_theta: Optional[float] = (
         None  # For sliding window attention. e.g., gemma3-1b
     )
+    rope_parameters: Optional[Dict[str, Dict[str, Any]]] = (
+        None  # Per-layer-type RoPE configs. e.g., {"full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25}}
+    )
     rope_freq_base: float = 10000.0  # The base frequency for RoPE. Keep it for BC.
     use_scaled_rope: bool = False  # Use scaled RoPE, introduced in llama3.1.
     rope_scale_factor: int = 8
@@ -184,6 +187,7 @@ class ModelArgs:
     normalize_tok_embeddings: bool = False
     scale_query_by: float = 1.0
     use_attn_o_gate: bool = False
+    headwise_attn_output_gate: bool = False
     use_attn_o_norm: bool = False
     use_residual_gate: bool = False
     use_ffn_learnable_scales: bool = False

examples/models/spark_x2_5/BUCK

@@ -0,0 +1,22 @@
+load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target")
+# Any targets that should be shared between fbcode and xplat must be defined in
+# targets.bzl. This file can contain fbcode-only targets.
+
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+fbcode_target(_kind = runtime.python_library,
+    name = "spark_x2_5",
+    srcs = [
+        "__init__.py",
+        "convert_weights.py",
+    ],
+    base_module = "executorch.examples.models.spark_x2_5",
+    visibility = ["PUBLIC"],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/examples/models/llama:transformer_modules",
+        "fbsource//third-party/pypi/safetensors:safetensors",
+    ],
+)