[Edge][Mobile] Spark-X2.5 decode-step export, and a compile-safe head gate

Exports one Spark decode step with the KV cache as explicit tensors, sized
per layer type: the 21 sliding layers hold a fixed 512 entries at any
context, the 7 full-attention layers grow. That asymmetry is the whole
reason the architecture is interesting on a phone, so the two types get
separately sized caches, masks and rotary angles.

Verified against the HF reference at 1024 context: 138.5 dB logits SNR,
identical top-1 and top-5.

Profiled on a real Galaxy S25 (Snapdragon 8 Elite) through AI Hub, QNN
context binary on the NPU, ms per layer:

  sliding (512 KV, any ctx)  fp16 1.689  w8a16 0.912  w4a16 0.917
  full @ 1k                  fp16 1.559  w8a16 1.129  w4a16 1.127
  full @ 4k                  fp16 2.964  w8a16 1.681  w4a16 1.662

Composed into a step (21 sliding + 7 full), w4a16 gives 27.2 ms at 1k and
30.9 ms at 4k, against 31.6 and 46.5 ms for the same model with every layer
full-attention -- 1.16x and 1.51x. Worth noting the hybrid is 0.94x at 1k in
fp16: sliding layers rotate all 256 head dims where full layers rotate 64,
and below roughly 2k that RoPE cost outweighs the shorter cache.

Also moves the head-wise attention gate off vLLM's CPU linear dispatch. The
gate is [2048, 8], and any linear that narrow is routed to the sgl-kernel
packed GEMM, whose fake-tensor rule assumes the packed weight keeps N in dim
0; VNNI packing does not at that size, so it reports K and torch.compile
fails tracing the gate multiply. A plain parameter with F.linear gains
nothing less -- a 2048x8 matmul has no use for packed AMX tiles -- and it
keeps the model compilable on CPU and GPU alike.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ

tests/edge/test_spark_export.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for the Spark-X2.5 decode-step exporter.
+
+The point of the export is the hybrid cache layout, so that is what these
+check: sliding layers stay at a fixed 512 entries however long the context
+gets, full layers grow, and the window actually rolls.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.edge.spark_export import (
+    FULL,
+    SLIDING,
+    SparkDecodeStep,
+    SparkStepConfig,
+    config_from_spark,
+    example_inputs,
+    input_names,
+    kv_cache_bytes,
+    output_names,
+)
+
+HF_CONFIG = {
+    "hidden_size": 64,
+    "num_attention_heads": 4,
+    "num_key_value_heads": 2,
+    "head_dim": 16,
+    "intermediate_size": 128,
+    "vocab_size": 256,
+    "sliding_window": 8,
+    "rms_norm_eps": 1e-6,
+    "layer_types": [SLIDING, SLIDING, SLIDING, FULL],
+    "rope_parameters": {
+        "full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25},
+        "sliding_attention": {"rope_theta": 10000, "partial_rotary_factor": 1.0},
+    },
+}
+
+
+def _cfg(**kw) -> SparkStepConfig:
+    return config_from_spark(HF_CONFIG, **kw)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_sliding_cache_is_capped_but_full_cache_grows():
+    cfg = _cfg()
+    assert cfg.cache_len(SLIDING, 4) == 4  # shorter than the window
+    assert cfg.cache_len(SLIDING, 1024) == 8  # capped at sliding_window
+    assert cfg.cache_len(SLIDING, 1_000_000) == 8
+    assert cfg.cache_len(FULL, 1024) == 1024
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_rotary_dim_differs_by_layer_type():
+    cfg = _cfg()
+    assert cfg.rotary_dim(SLIDING) == 16  # all head dims
+    assert cfg.rotary_dim(FULL) == 4  # first quarter only
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_kv_bytes_beats_uniform_attention_and_the_gap_grows():
+    cfg = _cfg()
+    near = kv_cache_bytes(cfg, 16)
+    far = kv_cache_bytes(cfg, 4096)
+    assert near["hybrid"] < near["uniform"]
+    assert far["hybrid"] < far["uniform"]
+    # The saving is a growing fraction, not a constant one.
+    assert far["saved"] / far["uniform"] > near["saved"] / near["uniform"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_example_inputs_match_the_declared_signature():
+    cfg = _cfg()
+    args = example_inputs(cfg, context=1024)
+    names = input_names(cfg)
+    assert len(args) == len(names)
+    caches = args[7:]
+    assert len(caches) == 2 * cfg.num_layers
+    for i, layer_type in enumerate(cfg.layer_types):
+        expected = cfg.cache_len(layer_type, 1024)
+        assert caches[2 * i].shape == (1, cfg.num_key_value_heads, expected, cfg.head_dim)
+    assert len(output_names(cfg)) == 1 + 2 * cfg.num_layers
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_sliding_layer_rolls_the_window_and_full_layer_appends():
+    cfg = _cfg()
+    step = SparkDecodeStep(cfg).eval()
+    context = 1024
+    args = example_inputs(cfg, context)
+    with torch.no_grad():
+        out = step(*args)
+    logits, caches = out[0], out[1:]
+    assert logits.shape == (1, cfg.vocab_size)
+    for i, layer_type in enumerate(cfg.layer_types):
+        k_new = caches[2 * i]
+        if layer_type == SLIDING:
+            # A sliding layer returns the whole rolled window, same size in.
+            assert k_new.shape[2] == cfg.cache_len(SLIDING, context)
+        else:
+            # A full layer returns only the new entry for the runtime to append.
+            assert k_new.shape[2] == 1
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_sliding_window_drops_the_oldest_entry():
+    """The rolled window must equal the old cache shifted by one, with the new
+    key appended -- that is what makes a masked pad slot safe to prepend."""
+    cfg = _cfg(layer_slice=slice(0, 1))
+    step = SparkDecodeStep(cfg).eval()
+    args = list(example_inputs(cfg, context=1024))
+    k_cache = args[7]
+    with torch.no_grad():
+        out = step(*args)
+    k_new = out[1]
+    assert torch.allclose(k_new[:, :, :-1], k_cache[:, :, 1:])
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_layer_slice_selects_one_layer_type():
+    types = HF_CONFIG["layer_types"]
+    sliding = _cfg(layer_slice=slice(0, 1))
+    assert sliding.layer_types == (SLIDING,)
+    full = _cfg(layer_slice=slice(types.index(FULL), types.index(FULL) + 1))
+    assert full.layer_types == (FULL,)
+    assert full.first_layer == types.index(FULL)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_unknown_layer_type_is_rejected():
+    bad = dict(HF_CONFIG, layer_types=[SLIDING, "linear_attention", SLIDING, FULL])
+    with pytest.raises(ValueError, match="unknown layer types"):
+        config_from_spark(bad)

vllm_omni/edge/spark_export.py

@@ -0,0 +1,413 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Export a Spark-X2.5 decode step for phone accelerators (QNN / TFLite / ORT).
+
+The graph is one autoregressive step with the KV cache passed in and out as
+explicit tensors, which is the shape every mobile NPU runtime wants: static
+shapes, no control flow, no in-place cache mutation.
+
+What makes Spark worth exporting this way is its hybrid attention. Three of
+every four layers attend to a 512-token sliding window, so their cache is a
+fixed 512 entries *no matter how long the conversation is*; only the seven
+full-attention layers grow with context. At a 4k context that is 81 MiB of
+KV instead of the 235 MiB a uniform-attention model of the same shape would
+carry, and the gap widens linearly from there. The two layer types are
+therefore exported with separately sized caches, masks and rotary angles --
+sliding layers rotate all 256 head dims with theta 1e4, full layers only the
+first 64 with theta 5e6.
+
+``python -m vllm_omni.edge.spark_export --help`` exports one layer of each
+type, or the whole 28-layer step, to ONNX.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+
+from vllm_omni.edge.qnn_export import sanitize_onnx
+
+__all__ = [
+    "SparkStepConfig",
+    "SparkDecodeStep",
+    "config_from_spark",
+    "example_inputs",
+    "input_names",
+    "output_names",
+    "export_onnx",
+    "load_spark_weights",
+    "kv_cache_bytes",
+]
+
+SLIDING = "sliding_attention"
+FULL = "full_attention"
+
+
+@dataclass
+class SparkStepConfig:
+    """Shape of one Spark-X2.5 decode step."""
+
+    hidden_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    intermediate_size: int
+    vocab_size: int
+    sliding_window: int
+    layer_types: tuple[str, ...]
+    rope_theta: dict[str, float]
+    partial_rotary_factor: dict[str, float]
+    rms_norm_eps: float = 1e-6
+    include_lm_head: bool = True
+    first_layer: int = 0
+    """Index of the first exported layer, so a subset can be profiled on its own."""
+
+    def __post_init__(self) -> None:
+        if self.num_attention_heads % self.num_key_value_heads:
+            raise ValueError(
+                f"num_attention_heads {self.num_attention_heads} is not a multiple "
+                f"of num_key_value_heads {self.num_key_value_heads}"
+            )
+        bad = set(self.layer_types) - {SLIDING, FULL}
+        if bad:
+            raise ValueError(f"unknown layer types: {sorted(bad)}")
+
+    @property
+    def group_size(self) -> int:
+        return self.num_attention_heads // self.num_key_value_heads
+
+    @property
+    def num_layers(self) -> int:
+        return len(self.layer_types)
+
+    def rotary_dim(self, layer_type: str) -> int:
+        return int(self.head_dim * self.partial_rotary_factor[layer_type])
+
+    def cache_len(self, layer_type: str, context: int) -> int:
+        """KV entries a layer of this type holds at ``context`` tokens."""
+        return min(context, self.sliding_window) if layer_type == SLIDING else context
+
+
+def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    return weight * (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps))
+
+
+def _rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Neox-style rotation of the leading ``cos.shape[-1]`` dims; the rest pass through."""
+    rot = cos.shape[-1]
+    if rot < x.shape[-1]:
+        x_rot, x_pass = x[..., :rot], x[..., rot:]
+    else:
+        x_rot, x_pass = x, None
+    half = rot // 2
+    x1, x2 = x_rot[..., :half], x_rot[..., half:]
+    out = x_rot * cos + torch.cat([-x2, x1], dim=-1) * sin
+    return out if x_pass is None else torch.cat([out, x_pass], dim=-1)
+
+
+class _Layer(nn.Module):
+    def __init__(self, cfg: SparkStepConfig, layer_type: str):
+        super().__init__()
+        self.cfg = cfg
+        self.layer_type = layer_type
+        h, d = cfg.hidden_size, cfg.head_dim
+        q_dim = cfg.num_attention_heads * d
+        kv_dim = cfg.num_key_value_heads * d
+        self.input_layernorm = nn.Parameter(torch.ones(h))
+        self.post_attention_layernorm = nn.Parameter(torch.ones(h))
+        self.q_k_v_proj = nn.Linear(h, q_dim + 2 * kv_dim, bias=False)
+        self.g_proj = nn.Linear(h, cfg.num_attention_heads, bias=False)
+        self.out_proj = nn.Linear(q_dim, h, bias=False)
+        self.gate_proj = nn.Linear(h, cfg.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(h, cfg.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(cfg.intermediate_size, h, bias=False)
+
+    def forward(self, x, cos, sin, k_cache, v_cache, mask):
+        cfg = self.cfg
+        heads, kv, d, g = (
+            cfg.num_attention_heads,
+            cfg.num_key_value_heads,
+            cfg.head_dim,
+            cfg.group_size,
+        )
+        h = _rms(x, self.input_layernorm, cfg.rms_norm_eps)
+        qkv = self.q_k_v_proj(h)
+        q = qkv[..., : heads * d].view(1, heads, 1, d)
+        k = qkv[..., heads * d : (heads + kv) * d].view(1, kv, 1, d)
+        v = qkv[..., (heads + kv) * d :].view(1, kv, 1, d)
+        gate = torch.sigmoid(self.g_proj(h)).view(1, heads, 1, 1)
+        q, k = _rope(q, cos, sin), _rope(k, cos, sin)
+
+        # Roll the window: the oldest entry falls off the front. Full-attention
+        # layers keep everything, so their cache simply grows.
+        keys = torch.cat([k_cache, k], dim=2)
+        vals = torch.cat([v_cache, v], dim=2)
+        if self.layer_type == SLIDING:
+            keys, vals = keys[:, :, 1:], vals[:, :, 1:]
+
+        # GQA without materialising expanded K/V: fold the query heads that
+        # share a kv head into their own axis.
+        scores = torch.matmul(q.view(1, kv, g, d), keys.transpose(2, 3))
+        scores = scores * (1.0 / math.sqrt(d)) + mask
+        attn = torch.matmul(torch.softmax(scores, dim=-1), vals)
+        attn = (attn.view(1, heads, 1, d) * gate).view(1, 1, heads * d)
+
+        x = x + self.out_proj(attn)
+        h = _rms(x, self.post_attention_layernorm, cfg.rms_norm_eps)
+        x = x + self.down_proj(
+            torch.nn.functional.gelu(self.gate_proj(h)) * self.up_proj(h)
+        )
+        return x, keys if self.layer_type == SLIDING else k, vals if self.layer_type == SLIDING else v
+
+
+class SparkDecodeStep(nn.Module):
+    """One decode step. Sliding layers return their whole rolled window; full
+    layers return only the new entry, which the runtime appends."""
+
+    def __init__(self, cfg: SparkStepConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.layers = nn.ModuleList(_Layer(cfg, t) for t in cfg.layer_types)
+        if cfg.include_lm_head:
+            self.norm = nn.Parameter(torch.ones(cfg.hidden_size))
+            self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
+
+    def forward(self, x, cos_sw, sin_sw, cos_full, sin_full, mask_sw, mask_full, *caches):
+        cfg = self.cfg
+        if len(caches) != 2 * cfg.num_layers:
+            raise ValueError(
+                f"expected {2 * cfg.num_layers} cache tensors, got {len(caches)}"
+            )
+        new: list[torch.Tensor] = []
+        for i, layer in enumerate(self.layers):
+            sliding = cfg.layer_types[i] == SLIDING
+            cos, sin = (cos_sw, sin_sw) if sliding else (cos_full, sin_full)
+            mask = mask_sw if sliding else mask_full
+            x, k, v = layer(x, cos, sin, caches[2 * i], caches[2 * i + 1], mask)
+            new += [k, v]
+        if not cfg.include_lm_head:
+            return (x, *new)
+        logits = self.lm_head(_rms(x, self.norm, cfg.rms_norm_eps))[:, 0]
+        return (logits, *new)
+
+
+def config_from_spark(
+    hf_config: dict[str, Any],
+    layer_slice: slice | None = None,
+    include_lm_head: bool = True,
+) -> SparkStepConfig:
+    """Build a step config from a Spark ``config.json``.
+
+    ``layer_slice`` selects a contiguous subset of layers, which is how a
+    single sliding or full layer gets profiled on its own.
+    """
+    types = list(hf_config["layer_types"])
+    first = 0
+    if layer_slice is not None:
+        first = layer_slice.start or 0
+        types = types[layer_slice]
+    rope = hf_config["rope_parameters"]
+    return SparkStepConfig(
+        hidden_size=hf_config["hidden_size"],
+        num_attention_heads=hf_config["num_attention_heads"],
+        num_key_value_heads=hf_config["num_key_value_heads"],
+        head_dim=hf_config["head_dim"],
+        intermediate_size=hf_config["intermediate_size"],
+        vocab_size=hf_config["vocab_size"],
+        sliding_window=hf_config["sliding_window"],
+        layer_types=tuple(types),
+        rope_theta={k: float(v.get("rope_theta", 10000)) for k, v in rope.items()},
+        partial_rotary_factor={
+            k: float(v.get("partial_rotary_factor", 1.0)) for k, v in rope.items()
+        },
+        rms_norm_eps=hf_config.get("rms_norm_eps", 1e-6),
+        include_lm_head=include_lm_head,
+        first_layer=first,
+    )
+
+
+def input_names(cfg: SparkStepConfig) -> list[str]:
+    return [
+        "x",
+        "cos_sw",
+        "sin_sw",
+        "cos_full",
+        "sin_full",
+        "mask_sw",
+        "mask_full",
+    ] + [f"{kv}_cache_{i}" for i in range(cfg.num_layers) for kv in ("k", "v")]
+
+
+def output_names(cfg: SparkStepConfig) -> list[str]:
+    head = "logits" if cfg.include_lm_head else "hidden"
+    return [head] + [
+        f"{kv}_new_{i}" for i in range(cfg.num_layers) for kv in ("k", "v")
+    ]
+
+
+def _angles(position: int, rotary_dim: int, theta: float, dtype: torch.dtype):
+    inv = 1.0 / (
+        theta ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
+    )
+    ang = position * inv
+    cos = torch.cat([ang.cos(), ang.cos()]).view(1, 1, 1, rotary_dim).to(dtype)
+    sin = torch.cat([ang.sin(), ang.sin()]).view(1, 1, 1, rotary_dim).to(dtype)
+    return cos, sin
+
+
+def example_inputs(
+    cfg: SparkStepConfig,
+    context: int,
+    position: int | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, ...]:
+    """Concrete inputs for tracing at a given context length."""
+    pos = context if position is None else position
+    sw_len = cfg.cache_len(SLIDING, context)
+    full_len = cfg.cache_len(FULL, context)
+    cos_sw, sin_sw = _angles(pos, cfg.rotary_dim(SLIDING), cfg.rope_theta[SLIDING], dtype)
+    cos_full, sin_full = _angles(pos, cfg.rotary_dim(FULL), cfg.rope_theta[FULL], dtype)
+    # Sliding layers drop the oldest slot, so their score row is sw_len wide;
+    # full layers score against everything cached plus the new token.
+    mask_sw = torch.zeros(1, 1, 1, sw_len, dtype=dtype)
+    mask_full = torch.zeros(1, 1, 1, full_len + 1, dtype=dtype)
+    x = (torch.randn(1, 1, cfg.hidden_size) * 0.1).to(dtype)
+    caches: list[torch.Tensor] = []
+    for t in cfg.layer_types:
+        n = cfg.cache_len(t, context)
+        for _ in range(2):
+            caches.append(
+                torch.randn(1, cfg.num_key_value_heads, n, cfg.head_dim, dtype=dtype)
+                * 0.5
+            )
+    return (x, cos_sw, sin_sw, cos_full, sin_full, mask_sw, mask_full, *caches)
+
+
+def kv_cache_bytes(cfg: SparkStepConfig, context: int, bytes_per_elem: int = 2) -> dict[str, int]:
+    """KV footprint of the hybrid layout versus an all-full-attention model."""
+    per_entry = 2 * cfg.num_key_value_heads * cfg.head_dim * bytes_per_elem
+    hybrid = sum(cfg.cache_len(t, context) for t in cfg.layer_types) * per_entry
+    uniform = cfg.num_layers * context * per_entry
+    return {"hybrid": hybrid, "uniform": uniform, "saved": uniform - hybrid}
+
+
+def load_spark_weights(
+    module: SparkDecodeStep, weights_dir: str | Path, dtype: torch.dtype = torch.float32
+) -> int:
+    """Load real Spark weights into the exported step. Returns tensors copied."""
+    from safetensors.torch import load_file
+
+    weights_dir = Path(weights_dir)
+    shards = sorted(weights_dir.glob("*.safetensors"))
+    if not shards:
+        raise FileNotFoundError(f"no safetensors under {weights_dir}")
+    state: dict[str, torch.Tensor] = {}
+    for shard in shards:
+        state.update(load_file(str(shard)))
+
+    cfg = module.cfg
+    copied = 0
+    with torch.no_grad():
+        for i, layer in enumerate(module.layers):
+            src = f"model.layers.{cfg.first_layer + i}."
+            for dst_name, key in (
+                ("input_layernorm", "input_layernorm.weight"),
+                ("post_attention_layernorm", "post_attention_layernorm.weight"),
+            ):
+                getattr(layer, dst_name).copy_(state[src + key].to(dtype))
+                copied += 1
+            for mod, key in (
+                (layer.q_k_v_proj, "self_attn.q_k_v_proj.weight"),
+                (layer.g_proj, "self_attn.g_proj.weight"),
+                (layer.out_proj, "self_attn.out_proj.weight"),
+                (layer.gate_proj, "mlp.gate_proj.weight"),
+                (layer.up_proj, "mlp.up_proj.weight"),
+                (layer.down_proj, "mlp.down_proj.weight"),
+            ):
+                mod.weight.copy_(state[src + key].to(dtype))
+                copied += 1
+        if cfg.include_lm_head:
+            module.norm.copy_(state["model.norm.weight"].to(dtype))
+            # Spark ties the output projection to the input embedding.
+            emb = state.get("lm_head.weight", state["model.embedding.weight"])
+            module.lm_head.weight.copy_(emb.to(dtype))
+            copied += 2
+    return copied
+
+
+def export_onnx(
+    module: SparkDecodeStep, context: int, path: str | Path, opset: int = 18
+) -> dict[str, Any]:
+    """Trace to ONNX at a fixed context length and sanitize for Qualcomm tools."""
+    cfg = module.cfg
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    args = example_inputs(cfg, context)
+    module.eval()
+    with torch.no_grad():
+        torch.onnx.export(
+            module,
+            args,
+            str(path),
+            input_names=input_names(cfg),
+            output_names=output_names(cfg),
+            opset_version=opset,
+            dynamo=False,
+        )
+    stats = sanitize_onnx(path)
+    return {
+        "path": str(path),
+        "context": context,
+        "layers": cfg.num_layers,
+        "layer_types": list(cfg.layer_types),
+        "include_lm_head": cfg.include_lm_head,
+        "sanitized": stats,
+        "kv_bytes": kv_cache_bytes(cfg, context),
+    }
+
+
+def _main(argv: list[str] | None = None) -> int:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--config", required=True, help="path to Spark config.json")
+    ap.add_argument("--weights", default=None, help="directory of safetensors shards")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--context", type=int, default=1024)
+    ap.add_argument(
+        "--layers",
+        default="all",
+        help="'all', 'sliding', 'full', or START:STOP",
+    )
+    ap.add_argument("--no-lm-head", action="store_true")
+    a = ap.parse_args(argv)
+
+    hf = json.loads(Path(a.config).read_text())
+    types = hf["layer_types"]
+    if a.layers == "all":
+        sl = None
+    elif a.layers == "sliding":
+        i = types.index(SLIDING)
+        sl = slice(i, i + 1)
+    elif a.layers == "full":
+        i = types.index(FULL)
+        sl = slice(i, i + 1)
+    else:
+        start, stop = (int(v) for v in a.layers.split(":"))
+        sl = slice(start, stop)
+
+    cfg = config_from_spark(hf, sl, include_lm_head=not a.no_lm_head)
+    step = SparkDecodeStep(cfg)
+    if a.weights:
+        print(f"loaded {load_spark_weights(step, a.weights)} tensors")
+    print(json.dumps(export_onnx(step, a.context, a.out), indent=2))
+    return 0
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(_main())

vllm_omni/model_executor/models/spark2_5/spark2_5.py

@@ -28,12 +28,15 @@ from torch import nn
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from vllm.distributed import (
+    get_pp_group,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
 from vllm.model_executor.layers.activation import get_act_and_mul_fn
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
     MergedColumnParallelLinear,
     QKVParallelLinear,
     RowParallelLinear,
@@ -46,6 +49,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
+from vllm.model_executor.utils import set_weight_attrs
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     WeightsMapper,
@@ -140,24 +144,34 @@ class Spark2_5Attention(nn.Module):
             prefix=f"{prefix}.out_proj",
         )
 
-        # One scalar gate per attention head. Sharded column-wise so each rank
-        # gets exactly the gates for the query heads it owns.
+        # One scalar gate per attention head, so this projection is only
+        # [hidden, num_heads] wide. It is deliberately a plain parameter
+        # rather than a ColumnParallelLinear: vLLM's CPU backend routes any
+        # narrow linear to the sgl-kernel packed GEMM, whose fake tensor rule
+        # assumes the packed weight keeps N in dim 0. VNNI packing does not
+        # for N this small, so the shape it reports is K rather than N and
+        # torch.compile dies tracing the gate multiply. A 2048x8 matmul has
+        # nothing to gain from the packed AMX kernel anyway.
         self.headwise_attn_output_gate = config.headwise_attn_output_gate
         if self.headwise_attn_output_gate:
-            self.g_proj = ColumnParallelLinear(
-                config.hidden_size,
-                self.total_num_heads,
-                bias=config.attention_bias,
-                quant_config=quant_config,
-                prefix=f"{prefix}.g_proj",
-            )
+            if config.attention_bias:
+                raise NotImplementedError(
+                    "Spark-X2.5 with attention_bias and a head-wise gate is "
+                    "not supported"
+                )
             if config.gate_attn_act_mode not in ("sigmoid", "silu"):
                 raise ValueError(
                     f"Unsupported gate_attn_act_mode: {config.gate_attn_act_mode}"
                 )
             self.gate_attn_act_mode = config.gate_attn_act_mode
+            self.g_weight = nn.Parameter(
+                torch.empty(
+                    self.num_heads, config.hidden_size, dtype=torch.get_default_dtype()
+                )
+            )
+            set_weight_attrs(self.g_weight, {"weight_loader": self._load_g_weight})
         else:
-            self.g_proj = None
+            self.g_weight = None
 
         layer_idx = extract_layer_index(prefix)
         layer_type = config.layer_types[layer_idx]
@@ -186,6 +200,12 @@ class Spark2_5Attention(nn.Module):
             prefix=f"{prefix}.attn",
         )
 
+    def _load_g_weight(self, param: nn.Parameter, loaded_weight: torch.Tensor) -> None:
+        """Take just the gates belonging to this rank's query heads."""
+        tp_rank = get_tensor_model_parallel_rank()
+        shard = loaded_weight.narrow(0, tp_rank * self.num_heads, self.num_heads)
+        param.data.copy_(shard)
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -196,17 +216,19 @@ class Spark2_5Attention(nn.Module):
         q, k = self.rotary_emb(positions, q, k)
         attn_output = self.attn(q, k, v)
 
-        if self.g_proj is not None:
-            gate_score, _ = self.g_proj(hidden_states)
+        if self.g_weight is not None:
+            gate_score = torch.nn.functional.linear(hidden_states, self.g_weight)
             if self.gate_attn_act_mode == "sigmoid":
                 gate = torch.sigmoid(gate_score.float())
             else:
                 gate = nn.functional.silu(gate_score.float())
             gate = gate.to(attn_output.dtype)
-            attn_output = (
-                attn_output.view(-1, self.num_heads, self.head_dim)
-                * gate.view(-1, self.num_heads, 1)
-            ).view(-1, self.num_heads * self.head_dim)
+            # Reshape only the feature dim: a view(-1, ...) here would make
+            # torch.compile re-derive the symbolic token count and fail.
+            attn_output = attn_output.unflatten(
+                -1, (self.num_heads, self.head_dim)
+            ) * gate.unsqueeze(-1)
+            attn_output = attn_output.flatten(-2)
 
         output, _ = self.out_proj(attn_output)
         return output
@@ -254,7 +276,10 @@ class Spark2_5DecoderLayer(nn.Module):
 @support_torch_compile
 class Spark2_5Model(nn.Module):
     hf_to_vllm_mapper = WeightsMapper(
-        orig_to_new_substr={"model.embedding.": "model.embed_tokens."},
+        orig_to_new_substr={
+            "model.embedding.": "model.embed_tokens.",
+            ".self_attn.g_proj.weight": ".self_attn.g_weight",
+        },
         orig_to_new_stacked={
             ".q_proj": (".qkv_proj", "q"),
             ".k_proj": (".qkv_proj", "k"),