vllm-omni/vllm_omni/edge/qnn_export.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Export the Qwen3-TTS talker / code-predictor **decode step** for phone accelerators (QNN, TFLite, ONNX Runtime).

Companion to :mod:`vllm_omni.edge.decoder_export`, which exports the Code2Wav vocoder. Together the two cover the
whole on-device path served by the codec-token stream (WP5): the talker and its code predictor produce codec tokens,
the vocoder turns them into audio.

Layout
------
One token in, one token out, with an explicit static KV cache::

    inputs   x [1, 1, H], cos [1, 1, 1, D], sin [1, 1, 1, D], mask [1, 1, 1, L+1],
             k_cache_i / v_cache_i [1, KV, L, D] for each layer i
    outputs  logits [1, vocab], k_new_i / v_new_i [1, KV, 1, D] for each layer i

Everything that is cheap on a CPU and awkward on an accelerator stays on the host: embedding lookups, the rotary
angles (passed in as ``cos``/``sin``, which also keeps M-RoPE's interleaved sections out of the graph), sampling, and
appending ``k_new``/``v_new`` into the caches.

Two details are what make this fast, and both were measured on real devices (Snapdragon 8 Gen 3 / 8 Elite via
Qualcomm AI Hub, 2026-09-10; see ``analysis/mobile_htp_simulation.md`` §5):

* **per-layer cache tensors**, not one stacked ``[num_layers, ...]`` tensor that each layer slices. The stacked form
  spent ~70 % of its NPU cycles in cache ``select``/``Reshape`` and fp32→fp16 conversion of the whole cache;
* **grouped-query attention by reshaping the query** to ``[1, KV, G, D]`` instead of expanding keys and values to
  ``[1, H, L+1, D]``. Expanding multiplies the attention tensors by ``G`` before the matmul.

On a Snapdragon 8 Elite this took the 0.6B talker step from 36.9 ms to 15.8 ms (fp16, 256-token cache); the code
predictor step is 2.25 ms, so a full 80 ms frame (talker + 15 predictor sub-steps) costs ≈50 ms on the NPU.

Placement
---------
Measured, same phone, same graphs (per 80 ms frame, fp16):

===================  ===========  ==========  =========
component            Hexagon NPU  Adreno GPU  phone CPU
===================  ===========  ==========  =========
talker step            15.8 ms      29.9 ms    29.8 ms
predictor step x15     33.9 ms      65.0 ms    72.3 ms
Code2Wav vocoder      182 ms        21.7 ms   109 ms
===================  ===========  ==========  =========

So run **this** graph on the NPU and the **vocoder on the GPU**: 8.4x for the vocoder, ~2x for the transformer steps,
and the two stages pipeline the way they already do in the server deployment. Sending everything to the NPU is 3.3x
slower than the split. Integer quantization is *not* the lever it appears to be: every automatic int8 build measured
was numerically broken (predictor 3.1 dB, vocoder -4.3 dB vs the fp32 reference), while fp16 is faithful
(39-48 dB on the NPU).

ONNX for the Qualcomm toolchains
--------------------------------
:func:`sanitize_onnx` fixes the three things the dynamo exporter emits that ``qairt-converter`` and AI Hub reject:
``IsNaN``/``Where`` guards, ``Reshape`` with ``allowzero=1``, and rank-0 ``axes`` initializers on
``Unsqueeze``/``Squeeze`` (which crash AI Hub's shape inference).

CLI::

    python -m vllm_omni.edge.qnn_export --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \\
        --component talker --cache-len 256 --out-dir ./qnn_export
"""

from __future__ import annotations

import argparse
import json
import math
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import torch
from torch import nn

__all__ = [
    "DecodeStepConfig",
    "DecodeStep",
    "config_from_qwen3_tts",
    "load_qwen3_tts_weights",
    "example_inputs",
    "input_names",
    "output_names",
    "export_onnx",
    "sanitize_onnx",
    "PLACEMENT_MS_PER_FRAME",
]

#: Measured on Samsung Galaxy S25 (Snapdragon 8 Elite), fp16, per 80 ms audio frame. See the module docstring.
PLACEMENT_MS_PER_FRAME = {
    "talker_step": {"npu": 15.8, "gpu": 29.9, "cpu": 29.8},
    "predictor_step_x15": {"npu": 33.9, "gpu": 65.0, "cpu": 72.3},
    "code2wav_chunk": {"npu": 182.0, "gpu": 21.7, "cpu": 109.0},
}


@dataclass(frozen=True)
class DecodeStepConfig:
    """Shape of one autoregressive step. Mirrors the fields of a Qwen3 decoder config that the graph needs."""

    hidden_size: int
    num_hidden_layers: int
    num_attention_heads: int
    num_key_value_heads: int
    head_dim: int
    intermediate_size: int
    vocab_size: int
    rms_norm_eps: float = 1e-6
    rope_theta: float = 1e6

    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 of num_key_value_heads {self.num_key_value_heads}")

    @property
    def group_size(self) -> int:
        return self.num_attention_heads // self.num_key_value_heads


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:
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return x * cos + torch.cat([-x2, x1], dim=-1) * sin


class _Layer(nn.Module):
    def __init__(self, cfg: DecodeStepConfig):
        super().__init__()
        self.cfg = cfg
        h, d = cfg.hidden_size, cfg.head_dim
        self.input_layernorm = nn.Parameter(torch.ones(h))
        self.post_attention_layernorm = nn.Parameter(torch.ones(h))
        self.q_norm = nn.Parameter(torch.ones(d))
        self.k_norm = nn.Parameter(torch.ones(d))
        self.q_proj = nn.Linear(h, cfg.num_attention_heads * d, bias=False)
        self.k_proj = nn.Linear(h, cfg.num_key_value_heads * d, bias=False)
        self.v_proj = nn.Linear(h, cfg.num_key_value_heads * d, bias=False)
        self.o_proj = nn.Linear(cfg.num_attention_heads * d, 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
        kv, d, g = cfg.num_key_value_heads, cfg.head_dim, cfg.group_size
        h = _rms(x, self.input_layernorm, cfg.rms_norm_eps)
        q = _rms(self.q_proj(h).view(1, cfg.num_attention_heads, 1, d), self.q_norm, cfg.rms_norm_eps)
        k = _rms(self.k_proj(h).view(1, kv, 1, d), self.k_norm, cfg.rms_norm_eps)
        v = self.v_proj(h).view(1, kv, 1, d)
        q, k = _rope(q, cos, sin), _rope(k, cos, sin)
        keys = torch.cat([k_cache, k], dim=2)
        vals = torch.cat([v_cache, v], dim=2)
        # GQA without expanding keys/values: group the query heads that share a kv head.
        scores = torch.matmul(q.view(1, kv, g, d), keys.transpose(2, 3)) * (1.0 / math.sqrt(d)) + mask
        attn = torch.matmul(torch.softmax(scores, dim=-1), vals).view(1, 1, cfg.num_attention_heads * d)
        x = x + self.o_proj(attn)
        h = _rms(x, self.post_attention_layernorm, cfg.rms_norm_eps)
        x = x + self.down_proj(torch.nn.functional.silu(self.gate_proj(h)) * self.up_proj(h))
        return x, k, v


class DecodeStep(nn.Module):
    """One decode step with an explicit per-layer KV cache; see the module docstring for the tensor contract."""

    def __init__(self, cfg: DecodeStepConfig):
        super().__init__()
        self.cfg = cfg
        self.layers = nn.ModuleList(_Layer(cfg) for _ in range(cfg.num_hidden_layers))
        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, sin, mask, *caches):
        if len(caches) != 2 * self.cfg.num_hidden_layers:
            raise ValueError(f"expected {2 * self.cfg.num_hidden_layers} cache tensors, got {len(caches)}")
        new: list[torch.Tensor] = []
        for i, layer in enumerate(self.layers):
            x, k, v = layer(x, cos, sin, caches[2 * i], caches[2 * i + 1], mask)
            new += [k, v]
        logits = self.lm_head(_rms(x, self.norm, self.cfg.rms_norm_eps))[:, 0]
        return (logits, *new)


def input_names(cfg: DecodeStepConfig) -> list[str]:
    return ["x", "cos", "sin", "mask"] + [f"{kv}_cache_{i}" for i in range(cfg.num_hidden_layers) for kv in ("k", "v")]


def output_names(cfg: DecodeStepConfig) -> list[str]:
    return ["logits"] + [f"{kv}_new_{i}" for i in range(cfg.num_hidden_layers) for kv in ("k", "v")]


def example_inputs(cfg: DecodeStepConfig, cache_len: int, position: int | None = None, dtype: torch.dtype = torch.float32) -> tuple[torch.Tensor, ...]:
    """Concrete inputs for tracing: a random hidden state, the rotary angles for ``position`` and a causal mask."""
    pos = cache_len // 2 if position is None else position
    d = cfg.head_dim
    inv = 1.0 / (cfg.rope_theta ** (torch.arange(0, d, 2, dtype=torch.float32) / d))
    ang = pos * inv
    cos = torch.cat([ang.cos(), ang.cos()]).view(1, 1, 1, d).to(dtype)
    sin = torch.cat([ang.sin(), ang.sin()]).view(1, 1, 1, d).to(dtype)
    mask = torch.zeros(1, 1, 1, cache_len + 1, dtype=dtype)
    mask[..., pos:cache_len] = torch.finfo(torch.float16).min / 4  # positions not yet written
    x = (torch.randn(1, 1, cfg.hidden_size) * 0.1).to(dtype)
    caches = [torch.randn(1, cfg.num_key_value_heads, cache_len, d, dtype=dtype) * 0.5 for _ in range(2 * cfg.num_hidden_layers)]
    return (x, cos, sin, mask, *caches)


def config_from_qwen3_tts(hf_config: dict[str, Any], component: str) -> DecodeStepConfig:
    """Build a :class:`DecodeStepConfig` from a Qwen3-TTS ``config.json`` dict. ``component``: talker | predictor."""
    talker = hf_config["talker_config"] if "talker_config" in hf_config else hf_config
    cfg = talker if component == "talker" else talker["code_predictor_config"]
    return DecodeStepConfig(
        hidden_size=cfg["hidden_size"],
        num_hidden_layers=cfg["num_hidden_layers"],
        num_attention_heads=cfg["num_attention_heads"],
        num_key_value_heads=cfg["num_key_value_heads"],
        head_dim=cfg.get("head_dim", cfg["hidden_size"] // cfg["num_attention_heads"]),
        intermediate_size=cfg["intermediate_size"],
        vocab_size=cfg["vocab_size"],
        rms_norm_eps=cfg.get("rms_norm_eps", 1e-6),
        rope_theta=cfg.get("rope_theta", 1e6),
    )


def load_qwen3_tts_weights(module: DecodeStep, safetensors_path: str | Path, component: str, lm_head_index: int = 0) -> int:
    """Copy real weights into ``module``. ``component``: ``talker`` (codec head) or ``predictor`` (``lm_head.<i>``)."""
    from safetensors import safe_open

    prefix = "talker." if component == "talker" else "talker.code_predictor."
    head_key = "codec_head.weight" if component == "talker" else f"lm_head.{lm_head_index}.weight"
    loaded = 0
    with safe_open(str(safetensors_path), "pt") as st:
        get = lambda key: st.get_tensor(prefix + key).float()  # noqa: E731
        for i, layer in enumerate(module.layers):
            p = f"model.layers.{i}."
            for name, param in (("input_layernorm", layer.input_layernorm), ("post_attention_layernorm", layer.post_attention_layernorm)):
                param.data.copy_(get(p + name + ".weight"))
                loaded += 1
            for name, param in (("q_norm", layer.q_norm), ("k_norm", layer.k_norm)):
                param.data.copy_(get(p + "self_attn." + name + ".weight"))
                loaded += 1
            for name in ("q_proj", "k_proj", "v_proj", "o_proj"):
                getattr(layer, name).weight.data.copy_(get(p + "self_attn." + name + ".weight"))
                loaded += 1
            for name in ("gate_proj", "up_proj", "down_proj"):
                getattr(layer, name).weight.data.copy_(get(p + "mlp." + name + ".weight"))
                loaded += 1
        module.norm.data.copy_(get("model.norm.weight"))
        module.lm_head.weight.data.copy_(get(head_key))
        loaded += 2
    return loaded


def sanitize_onnx(path: str | Path) -> dict[str, int]:
    """Make a dynamo-exported ONNX file digestible by ``qairt-converter`` and Qualcomm AI Hub, in place.

    Removes ``IsNaN``/``Where`` guard pairs (the converter has no ``IsNaN``), clears ``Reshape``'s ``allowzero``
    attribute (rejected), and promotes rank-0 ``axes`` initializers on ``Unsqueeze``/``Squeeze`` to rank 1 (AI Hub's
    shape inference raises "Inferred shape and existing shape differ in rank" on them). Returns the counts.
    """
    import onnx

    model = onnx.load(str(path))
    graph = model.graph
    produced = {out: node for node in graph.node for out in node.output}
    renamed: dict[str, str] = {}
    for node in list(graph.node):
        if node.op_type != "Where" or node.input[0] not in produced:
            continue
        guard = produced[node.input[0]]
        if guard.op_type != "IsNaN":
            continue
        source = guard.input[0]
        keep = source if source in (node.input[1], node.input[2]) else None
        if keep is None:
            continue
        renamed[node.output[0]] = keep
        graph.node.remove(node)
        graph.node.remove(guard)
    initializers = {init.name: init for init in graph.initializer}
    constants = {node.output[0]: node for node in graph.node if node.op_type == "Constant"}
    stats = {"nan_guards_removed": len(renamed), "allowzero_cleared": 0, "scalar_axes_fixed": 0}
    scalar_axes: list[str] = []
    for node in graph.node:
        for i, name in enumerate(node.input):
            if name in renamed:
                node.input[i] = renamed[name]
        if node.op_type == "Reshape":
            for attr in node.attribute:
                if attr.name == "allowzero" and attr.i:
                    attr.i = 0
                    stats["allowzero_cleared"] += 1
        if node.op_type in ("Unsqueeze", "Squeeze") and len(node.input) > 1:
            axes = node.input[1]
            if axes in initializers and not initializers[axes].dims:
                initializers[axes].dims.extend([1])
                scalar_axes.append(axes)
            elif axes in constants:
                for attr in constants[axes].attribute:
                    if attr.name == "value" and not attr.t.dims:
                        attr.t.dims.extend([1])
                        scalar_axes.append(axes)
    stats["scalar_axes_fixed"] = len(scalar_axes)
    for collection in (graph.value_info, graph.input):
        for info in list(collection):
            if info.name in scalar_axes:
                collection.remove(info)  # stale rank-0 shape annotation
    for out in graph.output:
        if out.name in renamed:
            out.name = renamed[out.name]
    onnx.save(model, str(path))
    return stats


def export_onnx(module: DecodeStep, cache_len: int, path: str | Path, opset: int = 18) -> dict[str, Any]:
    """Export ``module`` to ONNX at ``path`` (dynamo exporter) and sanitize it for the Qualcomm toolchains."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    cfg = module.cfg
    args = example_inputs(cfg, cache_len)
    program = torch.onnx.export(
        module,
        args,
        input_names=input_names(cfg),
        output_names=output_names(cfg),
        opset_version=opset,
        dynamo=True,
        optimize=True,
    )
    program.save(str(path))
    stats = sanitize_onnx(path)
    return {"path": str(path), "cache_len": cache_len, "config": asdict(cfg), "sanitize": stats}


def _main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice")
    parser.add_argument("--component", choices=("talker", "predictor"), default="talker")
    parser.add_argument("--cache-len", type=int, default=256)
    parser.add_argument("--lm-head-index", type=int, default=0, help="predictor only: which residual head to attach")
    parser.add_argument("--out-dir", default="./qnn_export")
    parser.add_argument("--no-weights", action="store_true", help="export the graph shape only (random weights)")
    args = parser.parse_args(argv)

    from huggingface_hub import snapshot_download

    snapshot = Path(snapshot_download(args.model, allow_patterns=["config.json", "model.safetensors"]))
    cfg = config_from_qwen3_tts(json.loads((snapshot / "config.json").read_text()), args.component)
    module = DecodeStep(cfg).eval()
    if not args.no_weights:
        loaded = load_qwen3_tts_weights(module, snapshot / "model.safetensors", args.component, args.lm_head_index)
        print(f"loaded {loaded} tensors ({sum(p.numel() for p in module.parameters()) / 1e6:.0f}M params)", flush=True)
    out = Path(args.out_dir) / f"{args.component}_step_L{args.cache_len}.onnx"
    manifest = export_onnx(module, args.cache_len, out)
    manifest["model"] = args.model
    manifest["component"] = args.component
    (Path(args.out_dir) / f"{args.component}_step_L{args.cache_len}.json").write_text(json.dumps(manifest, indent=2))
    print(json.dumps(manifest, indent=2), flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(_main())