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).
@@ -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()
@@ -353,6 +357,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,
@@ -361,6 +366,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
@@ -379,7 +385,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 (
@@ -421,10 +434,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:
@@ -463,11 +489,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}. "
@@ -517,12 +567,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",
+    ],
+)