Qualcomm AI Engine Direct - GA Static Qwen3

DannyYuyang-quic · DannyYuyang-quic · commit d0e1d1f8d397 · 2025-08-04T11:02:25.000+08:00
Summary:
 - support Qwen3-0.6B
 - support Qwen3-1.7B
 - refactor HF model registration for static llama
diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py
@@ -4372,6 +4372,65 @@ def test_qwen2_5(self):
                 )
                 self.assertGreaterEqual(msg["inference_speed"], 95)  # Lanai
 
+    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):
diff --git a/examples/qualcomm/oss_scripts/llama/__init__.py b/examples/qualcomm/oss_scripts/llama/__init__.py
@@ -0,0 +1,63 @@
+import os
+from abc import ABC
+from dataclasses import dataclass
+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
+
+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 = "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 = "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 = "qwen2_5"
+    convert_weights = convert_qwen3_weights
diff --git a/examples/qualcomm/oss_scripts/llama/llama.py b/examples/qualcomm/oss_scripts/llama/llama.py
@@ -65,6 +65,7 @@
 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.model.static_llama import (
     LlamaModel,
     ModelArgs,
@@ -103,8 +104,6 @@
 logging.basicConfig(level=logging.INFO, format=FORMAT)
 logging.getLogger().setLevel(logging.INFO)
 
-HUGGING_FACE_REPO_IDS = {"qwen2_5": "Qwen/Qwen2.5-0.5B"}
-
 
 def next_power_of_two(n):
     if n == 0:
@@ -409,7 +408,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
@@ -517,18 +515,7 @@ def compile(args, pte_filename, tokenizer):
     if args.params:
         params_path = args.params
     else:
-        if args.decoder_model == "qwen2_5":
-            cur_dir = os.path.dirname(__file__)
-            params_path = os.path.join(
-                cur_dir,
-                "..",
-                "..",
-                "..",
-                "models",
-                "qwen2_5",
-                "config",
-                "0_5b_config.json",
-            )
+        params_path = SUPPORTED_HF_MODELS[args.decoder_model].params_path
     with open(params_path) as f:
         kv_config = ModelArgs(**json.load(f))
 
@@ -601,13 +588,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
         )
@@ -1041,8 +1025,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,
     )
 
@@ -1244,12 +1229,12 @@ def export_llama(args) -> None:
         ), f"Wrong tokenizer provided for llama3_2."
         runtime_tokenizer_path = args.tokenizer_model
         decoder_model_version = "llama3"
-    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)
-        decoder_model_version = args.decoder_model
+        decoder_model_version = SUPPORTED_HF_MODELS[args.decoder_model].runner_version
 
         with open(runtime_tokenizer_path, "r+") as file:
             data = json.load(file)
diff --git a/examples/qualcomm/oss_scripts/llama/model/static_llama.py b/examples/qualcomm/oss_scripts/llama/model/static_llama.py
@@ -13,7 +13,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(
@@ -46,6 +49,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,
@@ -170,10 +181,19 @@ 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 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)):
+            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 self.use_qk_norm and not self.qk_norm_before_rope:
+                k[i] = self.k_norm_fn(k[i])
 
         output_y = []
         kh, vh = [], []
@@ -230,9 +250,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
@@ -384,13 +412,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)
 
