Qualcomm AI Engine Direct - GA Static Qwen3 (#13086)

Summary:
 - support Qwen3-0.6B
 - support Qwen3-1.7B
 - refactor HF model registration for static llama

Script
``` bash
python examples/qualcomm/oss_scripts/llama/llama.py -b build-android  -s $DEVICE -m SM8750 --prompt "I would like to learn python, could you teach me with a simple example?" --temperature 0 --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --ptq 16a4w --decoder_model qwen3
```

Stat
<img width="1709" height="789"
alt="ee13db20-f529-413f-95c8-b6ce4bfcb4f4"
src="https://github.com/user-attachments/assets/0bea9db8-f10c-4d5b-96fa-31cd775d0a74"
/>


### Test plan
Note: We only run Qwen3-0.6B for CI
``` bash
python backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_qwen3 --model SM8650 --build_folder build-android/ --executorch_root . -s $DEVICE --artifact ./qwen3
```

cc: @haowhsu-quic, @cccclai

Author: DannyYuyang-quic

backends/qualcomm/tests/test_qnn_delegate.py

@@ -4588,6 +4588,65 @@ def test_static_qwen2_5(self):
                         msg["inference_speed"], inference_speed_ref[self.model]
                     )
 
+    def test_qwen3(self):
+        if not self.required_envs():
+            self.skipTest("missing required envs")
+
+        prompt = "My favourite condiment is "
+        cmds = [
+            "python",
+            f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py",
+            "--artifact",
+            self.artifact_dir,
+            "--build_folder",
+            self.build_folder,
+            "--model",
+            self.model,
+            "--ip",
+            self.ip,
+            "--port",
+            str(self.port),
+            "--prompt",
+            f"{prompt}",
+            "--ptq",
+            "16a8w",
+            "--decoder_model",
+            "qwen3_0.6b",
+            "--model_mode",
+            "hybrid",
+            "--prefill_ar_len",
+            "32",
+            "--max_seq_len",
+            "128",
+        ]
+        if self.compile_only:
+            cmds.extend(["--compile_only"])
+        elif self.device:
+            cmds.extend(["--device", self.device])
+        if self.host:
+            cmds.extend(["--host", self.host])
+        elif self.enable_x86_64:
+            cmds.extend(["--enable_x86_64"])
+        if self.pre_gen_pte:
+            cmds.extend(["--pre_gen_pte", self.pre_gen_pte])
+
+        # Accuracy is bad for now. Just check user's prompt is returned.
+        golden_start_with = "My favourite condiment is "
+        p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL)
+        with Listener((self.ip, self.port)) as listener:
+            conn = listener.accept()
+            p.communicate()
+            msg = json.loads(conn.recv())
+            if "Error" in msg:
+                self.fail(msg["Error"])
+            else:
+                model_out = msg["result"][0]
+                self.assertTrue(
+                    model_out.startswith(golden_start_with),
+                    f"Expected Output: {golden_start_with}. Actual Output: {model_out}",
+                )
+                self.assertGreaterEqual(msg["inference_speed"], 70)  # Lanai
+
 
 class TestExampleOssScript(TestQNN):
     def test_albert(self):

examples/qualcomm/oss_scripts/llama/__init__.py

@@ -0,0 +1,73 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import Callable, Dict, Type
+
+from executorch.examples.models.qwen2_5 import (
+    convert_weights as convert_qwen2_5_weights,
+)
+from executorch.examples.models.qwen3 import convert_weights as convert_qwen3_weights
+
+from executorch.examples.qualcomm.oss_scripts.llama.decoder_constants import (
+    DECODER_MODEL_VERSION,
+)
+
+BASE_DIR = os.path.dirname(__file__)
+
+
+@dataclass(init=False, frozen=True)
+class HFModel(ABC):
+    repo_id: str
+    params_path: str
+    runner_version: str
+    convert_weights: Callable
+
+
+SUPPORTED_HF_MODELS: Dict[str, Type[HFModel]] = {}
+
+
+def register_hf_model(name: str):
+    def decorator(cls: Type[HFModel]):
+        SUPPORTED_HF_MODELS[name.lower()] = cls()
+        return cls()
+
+    return decorator
+
+
+@register_hf_model("qwen2_5")
+@dataclass(init=False, frozen=True)
+class Qwen2_5(HFModel):
+    repo_id: str = "Qwen/Qwen2.5-0.5B"
+    params_path: str = os.path.join(
+        BASE_DIR, "../../../models/qwen2_5/config/0_5b_config.json"
+    )
+    runner_version: str = field(default=DECODER_MODEL_VERSION["qwen2_5"])
+    convert_weights = convert_qwen2_5_weights
+
+
+@register_hf_model("qwen3_0_6b")
+@dataclass(init=False, frozen=True)
+class Qwen3_0_6B(HFModel):
+    repo_id: str = "Qwen/Qwen3-0.6B"
+    params_path: str = os.path.join(
+        BASE_DIR, "../../../models/qwen3/config/0_6b_config.json"
+    )
+    runner_version: str = field(default=DECODER_MODEL_VERSION["qwen2_5"])
+    convert_weights = convert_qwen3_weights
+
+
+@register_hf_model("qwen3_1_7b")
+@dataclass(init=False, frozen=True)
+class Qwen3_1_7B(HFModel):
+    repo_id: str = "Qwen/Qwen/Qwen3-1.7B"
+    params_path: str = os.path.join(
+        BASE_DIR, "../../../models/qwen3/config/1_7b_config.json"
+    )
+    runner_version: str = field(default=DECODER_MODEL_VERSION["qwen2_5"])
+    convert_weights = convert_qwen3_weights

examples/qualcomm/oss_scripts/llama/decoder_constants.py

@@ -4,8 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-HUGGING_FACE_REPO_IDS = {"qwen2_5": "Qwen/Qwen2.5-0.5B"}
-
 EVAL_MODE = {
     "kv": 0,
     "hybrid": 1,

examples/qualcomm/oss_scripts/llama/decoder_utils.py

@@ -19,7 +19,6 @@
     DECODER_MODEL_VERSION,
     EVAL_MODE,
 )
-
 from executorch.examples.qualcomm.utils import make_output_dir, SimpleADB
 from executorch.exir._serialize._program import deserialize_pte_binary
 from pytorch_tokenizers.hf_tokenizer import HuggingFaceTokenizer

examples/qualcomm/oss_scripts/llama/llama.py

@@ -64,10 +64,10 @@
 from executorch.examples.models.llama.source_transformation.quantize import (
     get_quant_embedding_transform,
 )
+from executorch.examples.qualcomm.oss_scripts.llama import SUPPORTED_HF_MODELS
 from executorch.examples.qualcomm.oss_scripts.llama.decoder_constants import (
     DECODER_MODEL_VERSION,
     EVAL_MODE,
-    HUGGING_FACE_REPO_IDS,
 )
 from executorch.examples.qualcomm.oss_scripts.llama.decoder_utils import (
     graph_module_inference,
@@ -227,7 +227,6 @@ def quantize(
 
         self.has_quant_io = True
         fx_graph_module = None
-
         with torch.no_grad():
             fx_graph_module = torch.export.export(
                 self.llama_graph_module, self.inputs, strict=True
@@ -351,14 +350,11 @@ def compile(args, pte_filename, tokenizer):
 
     kv_config, prefill_config = None, None
     if args.params:
-        with open(args.params) as f:
-            kv_config = ModelArgs(**json.load(f))
-    elif args.decoder_model == "qwen2_5":
-        from importlib.resources import files
-
-        data_dir = files("executorch").joinpath("examples/models/qwen2_5/config")
-        config_file = data_dir.joinpath("0_5b_config.json")
-        kv_config = ModelArgs(**json.loads(config_file.read_text()))
+        params_path = args.params
+    else:
+        params_path = SUPPORTED_HF_MODELS[args.decoder_model].params_path
+    with open(params_path) as f:
+        kv_config = ModelArgs(**json.load(f))
 
     # TODO: support batch inputs if necessary
     kv_config.max_batch_size = 1
@@ -430,13 +426,10 @@ def compile(args, pte_filename, tokenizer):
             raise RuntimeError(f"Unknown model_mode: {args.model_mode}.")
 
     if args.checkpoint is None:  # HF models
-        model_id = HUGGING_FACE_REPO_IDS[args.decoder_model]
-        if args.decoder_model == "qwen2_5":
-            from executorch.examples.models.qwen2_5 import (  # pyre-ignore[21]
-                convert_weights,
-            )
-
-            checkpoint = download_and_convert_hf_checkpoint(model_id, convert_weights)
+        checkpoint = download_and_convert_hf_checkpoint(
+            SUPPORTED_HF_MODELS[args.decoder_model].repo_id,
+            SUPPORTED_HF_MODELS[args.decoder_model].convert_weights,
+        )
         state_dict = torch.load(
             checkpoint, weights_only=True, map_location="cpu", mmap=True
         )
@@ -964,8 +957,9 @@ def _build_parser():
 
     parser.add_argument(
         "--decoder_model",
-        choices=["stories260k", "stories110m", "llama3_2", "qwen2_5"],
-        help="The Llama model to export. Current available options are: [stories260k, stories110m, llama3_2, qwen2_5]",
+        choices=["stories260k", "stories110m", "llama3_2"]
+        + list(SUPPORTED_HF_MODELS.keys()),
+        help=f"The Llama model to export. Current available options are: [stories260k, stories110m, llama3_2] + {SUPPORTED_HF_MODELS.keys()}",
         required=True,
     )
 
@@ -1176,11 +1170,19 @@ def export_llama(args) -> None:
             tokenizer, TiktokenTokenizer
         ), f"Wrong tokenizer provided for llama3_2."
         runtime_tokenizer_path = args.tokenizer_model
-    elif args.decoder_model == "qwen2_5":
-        model_id = HUGGING_FACE_REPO_IDS[args.decoder_model]
+    elif args.decoder_model in {"qwen2_5", "qwen3_0_6b", "qwen3_1_7b"}:
+        model_id = SUPPORTED_HF_MODELS[args.decoder_model].repo_id
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         runtime_tokenizer_path = tokenizer.save_pretrained(args.artifact)[-1]
         tokenizer = get_tokenizer(runtime_tokenizer_path)
+        with open(runtime_tokenizer_path, "r+") as file:
+            data = json.load(file)
+            # TODO: Encountered the following error during runtime, so switched behavior for now.
+            # Error: libc++abi: terminating due to uncaught exception of type std::runtime_error: Unsupported Normalizer type: NFC.
+            data.pop("normalizer")
+            file.seek(0)
+            json.dump(data, file, indent=4)
+            file.truncate()
     else:
         raise RuntimeError(f"Unknown decoder_model: {args.decoder_model}.")
 

examples/qualcomm/oss_scripts/llama/model/static_llama.py

@@ -15,7 +15,10 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from executorch.examples.models.llama.model_args import ModelArgs
-from executorch.examples.models.llama.rope import precompute_freqs_cis
+from executorch.examples.models.llama.rope import (
+    hf_precompute_freqs_cis,
+    precompute_freqs_cis,
+)
 
 
 def apply_rotary_emb_single(
@@ -48,6 +51,14 @@ def __init__(self, config: ModelArgs, output_new_cache_only=False):
         self.max_seq_len = config.max_seq_len
         self.output_new_cache_only = output_new_cache_only
         self.enable_masked_softmax = getattr(config, "enable_masked_softmax", False)
+        self.use_qk_norm = config.use_qk_norm
+        self.qk_norm_before_rope = config.qk_norm_before_rope
+
+        if self.use_qk_norm:
+            q_norm_dim = self.head_dim
+            k_norm_dim = self.head_dim
+            self.q_norm_fn = torch.nn.RMSNorm(q_norm_dim, eps=config.norm_eps)
+            self.k_norm_fn = torch.nn.RMSNorm(k_norm_dim, eps=config.norm_eps)
 
         self.wq = nn.Linear(
             self.dim,
@@ -151,7 +162,7 @@ def prepare_sha(self):
                 )
         self.wo_sha.weight.data.copy_(self.wo.weight[:, :, None, None])
 
-    def forward_sha(
+    def forward_sha(  # noqa: C901
         self,
         hidden_states: torch.Tensor,
         freqs_cos: torch.Tensor,
@@ -184,15 +195,23 @@ def forward_sha(
             .reshape(bsz, seq_len, self.head_dim)
             for wv_sha in self.wv_sha
         ]
+
         for i in range(len(q)):
+            if self.use_qk_norm and self.qk_norm_before_rope:
+                q[i] = self.q_norm_fn(q[i])
             q[i] = apply_rotary_emb_single(q[i], freqs_cos, freqs_sin)
             if hasattr(self.config, "enable_r3") and self.config.enable_r3:
                 q[i] = torch.matmul(q[i], self.r3_weight.T)
+            if self.use_qk_norm and not self.qk_norm_before_rope:
+                q[i] = self.q_norm_fn(q[i])
         for i in range(len(k)):
-            k[i] = apply_rotary_emb_single(k[i], freqs_cos, freqs_sin)
+            if self.use_qk_norm and self.qk_norm_before_rope:
+                k[i] = self.k_norm_fn(k[i])
+            k[i] = apply_rotary_emb_single(k[i], freqs_cos, freqs_sin).transpose(1, 2)
             if hasattr(self.config, "enable_r3") and self.config.enable_r3:
                 k[i] = torch.matmul(k[i], self.r3_weight.T)
-            k[i] = k[i].transpose(1, 2)
+            if self.use_qk_norm and not self.qk_norm_before_rope:
+                k[i] = self.k_norm_fn(k[i])
 
         output_y = []
         kh, vh = [], []
@@ -249,9 +268,17 @@ def forward(
         k = k.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
         v = v.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
 
+        if self.use_qk_norm and self.qk_norm_before_rope:
+            q = self.q_norm_fn(q)
+            k = self.k_norm_fn(k)
+
         q = apply_rotary_emb_single(q, freqs_cos, freqs_sin)
         k = apply_rotary_emb_single(k, freqs_cos, freqs_sin).permute(0, 2, 3, 1)
 
+        if self.use_qk_norm and not self.qk_norm_before_rope:
+            q = self.q_norm_fn(q)
+            k = self.k_norm_fn(k)
+
         output_kh, output_vh, output_y = [], [], []
         kh, vh = [], []
         # kv cache mode
@@ -403,13 +430,23 @@ def __init__(
         self.norm = torch.nn.RMSNorm(config.dim, eps=config.norm_eps)
         self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
         self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim)
-        freqs_cos, freqs_sin = precompute_freqs_cis(
-            config.head_dim,
-            config.max_seq_len,
-            config.rope_freq_base,
-            config.use_scaled_rope,
-            config.rope_scale_factor,
-        )
+        if config.use_hf_rope:
+            freqs_cos, freqs_sin = hf_precompute_freqs_cis(
+                config.head_dim,
+                config.max_seq_len,
+                config.rope_freq_base,
+                config.partial_rotary_factor,
+            )
+            freqs_cos = freqs_cos[:, : freqs_cos.shape[-1] // 2]
+            freqs_sin = freqs_sin[:, : freqs_sin.shape[-1] // 2]
+        else:
+            freqs_cos, freqs_sin = precompute_freqs_cis(
+                config.head_dim,
+                config.max_seq_len,
+                config.rope_freq_base,
+                config.use_scaled_rope,
+                config.rope_scale_factor,
+            )
         self.register_buffer("freqs_cos", freqs_cos, persistent=False)
         self.register_buffer("freqs_sin", freqs_sin, persistent=False)