Qualcomm AI Engine Direct - GA Static Phi-4-mini

Summary:
 - Support Phi-4-mini-instruct for static llama path
 - add P-ROPE for phi-4-mini

Author: DannyYuyang-quic

examples/qualcomm/oss_scripts/llama/README.md

@@ -6,6 +6,8 @@ This file provides you the instructions to run LLM Decoder model with different
  2. LLAMA3.2 1B
  3. LLAMA3.2 3B
  4. QWEN2.5 0.5B
+ 5. QWEN3 0.6B / 1.7B
+ 6. Phi4-mini-instruct
 
 We offer the following modes to execute the model:
 

examples/qualcomm/oss_scripts/llama/__init__.py

@@ -9,6 +9,10 @@
 from dataclasses import dataclass, field
 from typing import Callable, Dict, Type
 
+from executorch.examples.models.phi_4_mini import (
+    convert_weights as convert_phi_4_mini_weights,
+)
+
 from executorch.examples.models.qwen2_5 import (
     convert_weights as convert_qwen2_5_weights,
 )
@@ -71,3 +75,14 @@ class Qwen3_1_7B(HFModel):
     )
     runner_version: str = field(default=DECODER_MODEL_VERSION["qwen2_5"])
     convert_weights = convert_qwen3_weights
+
+
+@register_hf_model("phi_4_mini")
+@dataclass(init=False, frozen=True)
+class Phi4Mini(HFModel):
+    repo_id: str = "microsoft/Phi-4-mini-instruct"
+    params_path: str = os.path.join(
+        BASE_DIR, "../../../models/phi_4_mini/config/config.json"
+    )
+    runner_version: str = field(default=DECODER_MODEL_VERSION["phi_4_mini"])
+    convert_weights = convert_phi_4_mini_weights

examples/qualcomm/oss_scripts/llama/decoder_constants.py

@@ -15,4 +15,5 @@
     "stories110m": "llama2",
     "llama3_2": "llama3",
     "qwen2_5": "qwen2_5",
+    "phi_4_mini": "phi_4_mini",
 }

examples/qualcomm/oss_scripts/llama/llama.py

@@ -362,6 +362,7 @@ def compile(args, pte_filename, tokenizer):
     kv_config.use_kv_cache = True
     kv_config.enable_masked_softmax = args.enable_masked_softmax
     kv_config.enable_r3 = args.r3
+    kv_config.base_model_name_or_path = args.decoder_model
 
     prefill_config = copy.copy(kv_config)
     prefill_config.use_kv_cache = (
@@ -565,6 +566,7 @@ def permute(w, heads):
         llama_instance_list[i] = SingleLlama(
             llama_instance_list[i].eval(), pte_filename
         )
+
         if args.embedding_quantize:
             llama_instance_list[i].passes_job[I64toI32][
                 QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY
@@ -1180,6 +1182,16 @@ def export_llama(args) -> None:
             # 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")
+    elif args.decoder_model == "phi_4_mini":
+        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: invert=true is not supported for Split PreTokenizer. Only invert=false is supported.
+            data["pre_tokenizer"]["pretokenizers"][-2]["invert"] = False
             file.seek(0)
             json.dump(data, file, indent=4)
             file.truncate()

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

@@ -34,11 +34,35 @@ def apply_rotary_emb_single(
         freqs_sin = freqs_sin[None, :, None, :]
     x_out_r = x_r * freqs_cos - x_i * freqs_sin
     x_out_i = x_r * freqs_sin + x_i * freqs_cos
-
     x_out = torch.cat([x_out_r, x_out_i], dim=-1)
     return x_out
 
 
+def apply_partial_rotary_emb_single(
+    x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor
+) -> torch.Tensor:
+
+    if x.dim() == 4:
+        freqs_cos = freqs_cos[None, :, None, :]
+        freqs_sin = freqs_sin[None, :, None, :]
+
+    rotary_dim = freqs_cos.shape[-1] * 2
+
+    x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:]
+    x_r, x_i = x_rot[..., : x_rot.shape[-1] // 2], x_rot[..., x_rot.shape[-1] // 2 :]
+    x_out_r = x_r * freqs_cos - x_i * freqs_sin
+    x_out_i = x_r * freqs_sin + x_i * freqs_cos
+    x_rotated = torch.cat([x_out_r, x_out_i], dim=-1)
+    return torch.cat([x_rotated, x_pass], dim=-1)
+
+
+APPLY_ROPE_EMBEDDING_FUNCTIONS = {
+    "phi_4_mini": apply_partial_rotary_emb_single,
+    "qwen2_5": apply_rotary_emb_single,
+    "llama3_2": apply_rotary_emb_single,
+}
+
+
 class LlamaAttention(nn.Module):
     def __init__(self, config: ModelArgs, output_new_cache_only=False):
         super().__init__()
@@ -59,6 +83,9 @@ def __init__(self, config: ModelArgs, output_new_cache_only=False):
             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.apply_rope_emb = APPLY_ROPE_EMBEDDING_FUNCTIONS[
+            config.base_model_name_or_path
+        ]
 
         self.wq = nn.Linear(
             self.dim,
@@ -199,15 +226,15 @@ def forward_sha(  # noqa: C901
         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)
+            q[i] = self.apply_rope_emb(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)):
             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)
+            k[i] = self.apply_rope_emb(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)
             if self.use_qk_norm and not self.qk_norm_before_rope:
@@ -272,8 +299,8 @@ def forward(
             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)
+        q = self.apply_rope_emb(q, freqs_cos, freqs_sin)
+        k = self.apply_rope_emb(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)
@@ -368,7 +395,8 @@ def __init__(self, config: ModelArgs, output_new_cache_only=False):
         super().__init__()
         self.dim = config.dim
         self.attention = LlamaAttention(
-            config=config, output_new_cache_only=output_new_cache_only
+            config=config,
+            output_new_cache_only=output_new_cache_only,
         )
         self.feed_forward = FeedForward(config)
         self.attention_norm = torch.nn.RMSNorm(config.dim, eps=config.norm_eps)

examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp

@@ -9,8 +9,8 @@
 /**
  * @file
  *
- * This tool can run Llama2 110M, Llama3.2 1B / 3B, Qwen2.5 0.5B with Qualcomm
- * AI Engine Direct.
+ * This tool can run Llama2 110M, Llama3.2 1B / 3B, Qwen2.5 0.5B, Qwen3 0.6B
+ * / 1.7B phi4-mini-instruct with Qualcomm AI Engine Direct.
  *
  */
 
@@ -104,6 +104,16 @@ std::string get_formatted_prompt(
     case example::DecoderModelVersion::kQwen2_5:
       formatted_prompt.append(prompt);
       break;
+    case example::DecoderModelVersion::kPhi4:
+      if (!system_prompt.empty()) {
+        formatted_prompt.append("<|system|>");
+        formatted_prompt.append(system_prompt);
+        formatted_prompt.append("<|end|>");
+      }
+      formatted_prompt.append("<|user|>");
+      formatted_prompt.append(prompt);
+      formatted_prompt.append("<|end|><|assistant|>");
+      break;
     case example::DecoderModelVersion::kLlama3:
       if (!system_prompt.empty()) {
         formatted_prompt.append(

examples/qualcomm/oss_scripts/llama/runner/runner.cpp

@@ -130,6 +130,8 @@ Runner::Runner(
     decoder_model_version_ = DecoderModelVersion::kLlama3;
   } else if (decoder_model_version == "qwen2_5") {
     decoder_model_version_ = DecoderModelVersion::kQwen2_5;
+  } else if (decoder_model_version == "phi_4_mini") {
+    decoder_model_version_ = DecoderModelVersion::kPhi4;
   } else {
     ET_CHECK_MSG(false, "Unsupported Decoder Model");
   }

examples/qualcomm/oss_scripts/llama/runner/runner.h

@@ -31,6 +31,7 @@ enum DecoderModelVersion {
   kLlama2 = 0,
   kLlama3,
   kQwen2_5,
+  kPhi4,
 };
 class Runner {
  public: