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

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).
@@ -481,6 +491,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 +521,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 +700,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

@@ -1,3 +1,3 @@
 # @lint-ignore-every LICENSELINT
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
@@ -215,7 +215,9 @@
         ):
             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()
@@ -353,6 +355,7 @@
         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,
@@ -361,6 +364,7 @@
         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
@@ -379,7 +383,14 @@
                 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 (
@@ -421,10 +432,19 @@
             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:
@@ -463,11 +483,35 @@
         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}. "
@@ -517,12 +561,20 @@
             )
             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",
+    ],
+)

examples/models/spark_x2_5/README.md

@@ -0,0 +1,100 @@
+## Summary
+[Spark-X2.5](https://huggingface.co/collections/XHToken/spark-x25) is a compact, efficient general-purpose language model series developed by [XHToken](https://github.com/XHToken/Spark-X2.5), available in two variants — 1.7B and 4B. The models use a hybrid attention architecture that combines one full-attention layer with three sliding-window attention layers, natively supporting context windows of up to 1M tokens and 200+ languages.
+
+## Architecture highlights
+- **Hybrid attention**: 3 sliding-window attention layers + 1 full-attention layer (repeating)
+- **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
+- **Sliding window**: 512 tokens for sliding-attention layers
+- **GELU activation** in the MLP
+- **Tied word embeddings**
+
+## Instructions
+
+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.
+
+### Example export
+
+Export Spark-X2.5-1.7B to XNNPack, FP32:
+```
+python -m extension.llm.export.export_llm \
+  --config examples/models/spark_x2_5/config/spark_x2_5_xnnpack_fp32.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_fp32.pte"
+```
+
+Export Spark-X2.5-1.7B to XNNPack, quantized with 8da4w:
+```
+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"
+```
+
+Export Spark-X2.5-4B to XNNPack, quantized with 8da4w:
+```
+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_4b" \
+  +base.params="examples/models/spark_x2_5/config/spark_x2_5_4b_config.json" \
+  +export.output_name="spark_x2_5_4b_8da4w.pte"
+```
+
+Export Spark-X2.5-1.7B to MLX on Apple Silicon, 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"
+```
+
+To export with extended context (e.g., 2048 tokens):
+```
+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.max_seq_length=2048 \
+  +export.max_context_length=2048 \
+  +export.output_name="spark_x2_5_1_7b_8da4w.pte"
+```
+
+### Example run
+
+With ExecuTorch pybindings:
+```
+python -m examples.models.llama.runner.native \
+  --model spark_x2_5_1_7b \
+  --pte spark_x2_5_1_7b_8da4w.pte \
+  --tokenizer ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots/<snapshot>/tokenizer.json \
+  --tokenizer_config ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots/<snapshot>/tokenizer_config.json \
+  --prompt="<|user|>\nWho are you?<|end|>\n<|assistant|>\n" \
+  --params examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json \
+  --max_len 128 \
+  -kv \
+  --temperature 0.3
+```
+
+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/<snapshot>/tokenizer.json \
+  --prompt="<|user|>\nWho are you?<|end|>\n<|assistant|>\n" \
+  --temperature 0.3
+```
+
+Find the Hugging Face cache snapshot directory with:
+```
+python - <<'PY'
+from pathlib import Path
+root = Path.home() / ".cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots"
+for path in root.glob("*/tokenizer.json"):
+    print(path.parent)
+PY
+```
+
+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.

examples/models/spark_x2_5/__init__.py

@@ -0,0 +1,5 @@
+from executorch.examples.models.spark_x2_5.convert_weights import convert_weights
+
+__all__ = [
+    "convert_weights",
+]