@@ -182,6 +182,7 @@ def is_constant(
 def deduce_dtype(
     tensor: torch.Tensor, quant_infos: Optional[Dict] = None
 ) -> torch.dtype:
+    # Please be extra careful when modifying the logic here as HF Repo uses it too.
     if quant_infos:
         quant_range = quant_infos["quant_max"] - quant_infos["quant_min"]
         unsigned = quant_infos["quant_min"] >= 0
 
@@ -0,0 +1,46 @@
+# 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 logging
+from typing import Optional
+
+from executorch.backends.qualcomm.hf_transformers.causal_lm.hf_llm_quant_recipe import (
+    DefaultQuantRecipe,
+    Granite_3_3_2B_Instruct_HFQuantRecipe,
+    Llama3_2_1B_HFQuantRecipe,
+    LLMQuantRecipe,
+    Qwen2_5_0_5B_HFQuantRecipe,
+    Qwen2_5_1_5B_HFQuantRecipe,
+    Qwen3_0_6B_HFQuantRecipe,
+    Smollm2_HFQuantRecipe,
+)
+
+FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
+logging.basicConfig(level=logging.INFO, format=FORMAT)
+
+HUGGING_FACE_QUANT_RECIPES = {
+    "NousResearch/Llama-3.2-1B": Llama3_2_1B_HFQuantRecipe,
+    "Qwen/Qwen2.5-0.5B": Qwen2_5_0_5B_HFQuantRecipe,
+    "Qwen/Qwen2.5-0.5B-Instruct": Qwen2_5_0_5B_HFQuantRecipe,
+    "Qwen/Qwen2.5-1.5B-Instruct": Qwen2_5_1_5B_HFQuantRecipe,
+    "Qwen/Qwen3-0.6B": Qwen3_0_6B_HFQuantRecipe,
+    "HuggingFaceTB/SmolLM2-135M": Smollm2_HFQuantRecipe,
+    "ibm-granite/granite-3.3-2b-instruct": Granite_3_3_2B_Instruct_HFQuantRecipe,
+}
+
+
+def get_quant_recipe(model_id, verbose=True):
+    """
+    model_id = The Hugging Face Model ID.
+    verbose: If True, it will the quant config tagged for each node.
+    """
+    recipe_cls = HUGGING_FACE_QUANT_RECIPES.get(model_id, DefaultQuantRecipe)
+    if recipe_cls == DefaultQuantRecipe:
+        logging.warning(
+            f"{model_id} does not have customized quant recipe using default quant recipe."
+        )
+    quant_recipe: Optional[LLMQuantRecipe] = recipe_cls(verbose) if recipe_cls else None
+    return quant_recipe
@@ -187,24 +187,30 @@ def forward(self, x, position_ids):
 
 
 class QnnCausalLMExportableModule(torch.nn.Module):
-    def __init__(self, model):
+    def __init__(self, model: torch.nn.Module, max_seq_len, ar_len=1):
         super().__init__()
         self.model = model
-        self.config = model.config
         self._metadata = save_config_to_constant_methods(
             model.config, model.generation_config
         )
         logging.info(f"Metadata to be recorded in PTE: {self._metadata}")
 
-        self.num_layers = self.config.num_hidden_layers
+        self.num_layers = self.model.config.num_hidden_layers
         self.num_kv_heads = getattr(
-            self.config, "num_key_value_heads", self.config.num_attention_heads
+            self.model.config,
+            "num_key_value_heads",
+            self.model.config.num_attention_heads,
         )
-        self.head_dim = self.config.head_dim
-        self.max_seq_len = self.config.max_seq_len
-        self.ar_len = self.config.ar_len
+        self.ar_len = ar_len
+        self.max_seq_len = max_seq_len
         self.past_len = self.max_seq_len - self.ar_len
-
+        # Some config has head_dim provided that is different from equation below(e.g., qwen2)
+        if not hasattr(self.model.config, "head_dim"):
+            self.head_dim = (
+                self.model.config.hidden_size // self.model.config.num_attention_heads
+            )
+        else:
+            self.head_dim = self.model.config.head_dim
         self._register_attention_mask_for_4_53()
         self._use_precomputed_rope()
 
@@ -261,7 +267,13 @@ def get_example_inputs(self):
             torch.zeros(1, self.num_kv_heads, self.past_len, self.head_dim)
             for _ in range(self.num_layers)
         ]
-        return (input_tokens, atten_mask, pos_ids, past_k, past_v)
+        return {
+            "input_tokens": input_tokens,
+            "atten_mask": atten_mask,
+            "pos_ids": pos_ids,
+            "past_k": past_k,
+            "past_v": past_v,
+        }
 
     def forward(self, input_tokens, atten_mask, pos_ids, past_k, past_v):
         cache = QnnCustomStaticCache(past_k, past_v, max_cache_len=self.max_seq_len)
 
@@ -9760,11 +9760,11 @@ def test_hf_causal_lm(self):
 
         # TODO: Robust testing framework to check accuracy and performance metrics.
         golden_start_with = {
-            "llama3_2-1b": "Simply put, the theory of relativity states that the speed of light",
-            "qwen2_5-0_5b": "Simply put, the theory of relativity states that the laws of physics",
-            "qwen3-0_6b": "Simply put, the theory of relativity states that the laws of physics",
-            "smollm2_135m": "Simply put, the theory of relativity states that the speed of light",
-            "granite-3_3-2b": "Simply put, the theory of relativity states that the laws of physics",
+            "NousResearch/Llama-3.2-1B": "Simply put, the theory of relativity states that the speed of light",
+            "Qwen/Qwen2.5-0.5B": "Simply put, the theory of relativity states that the laws of physics",
+            "Qwen/Qwen3-0.6B": "Simply put, the theory of relativity states that the laws of physics",
+            "HuggingFaceTB/SmolLM2-135M": "Simply put, the theory of relativity states that the speed of light",
+            "ibm-granite/granite-3.3-2b-instruct": "Simply put, the theory of relativity states that the laws of physics",
         }
         assert (
             self.model_name in golden_start_with
@@ -9775,7 +9775,7 @@ def test_hf_causal_lm(self):
             f"{self.executorch_root}/examples/qualcomm/oss_scripts/hf_causal_lm.py",
             "--prompt",
             prompt,
-            "--decoder_model",
+            "--decoder_model_id",
             self.model_name,
             "--max_seq_len",
             "128",
 
@@ -504,12 +504,21 @@ def ensure_graph_specific_dict(value, graph_names):
                 )
             ep = m
         else:
-            ep = torch.export.export(
-                m,
-                inputs[graph_name],
-                dynamic_shapes=dynamic_shapes[graph_name],
-                strict=True,
-            )
+            if isinstance(inputs[graph_name], dict):
+                ep = torch.export.export(
+                    m,
+                    args=(),
+                    kwargs=inputs[graph_name],
+                    dynamic_shapes=dynamic_shapes[graph_name],
+                    strict=True,
+                )
+            else:
+                ep = torch.export.export(
+                    m,
+                    args=inputs[graph_name],
+                    dynamic_shapes=dynamic_shapes[graph_name],
+                    strict=True,
+                )
         option = generate_qnn_executorch_option(compiler_specs[graph_name])
         python_options = flatbuffer_to_option(option)
         backend_type = python_options.backend_options.backend_type
 
@@ -20,10 +20,8 @@
     setup_common_args_and_variables,
     SimpleADB,
 )
-
 from executorch.examples.qualcomm.oss_scripts.llm_utils.qnn_decoder_model_manager import (
     get_qnn_llm_edge_manager,
-    HUGGING_FACE_REPO_IDS,
 )
 from executorch.examples.qualcomm.utils import make_output_dir
 
@@ -35,25 +33,15 @@
 
 PTE_FILENAME = "hf_causal_lm_qnn"
 
-# Map the HF decoder_model keys (HUGGING_FACE_REPO_IDS) to the version strings
-# the shared qnn_llama_runner understands (see runner.cpp Runner()).
-DECODER_MODEL_VERSION = {
-    "llama3_2-1b": "llama3",
-    "qwen2_5-0_5b": "qwen2_5",
-    "qwen2_5-1_5b_instruct": "qwen2_5",
-    "qwen2_5-0_5b_instruct": "qwen2_5",
-    "qwen3-0_6b": "qwen3",
-    "smollm2_135m": "smollm2_135m",
-    "granite-3_3-2b": "granite",
-}
-
 
-def compile(args: argparse.Namespace, qnn_config: QnnConfig):  # noqa: C901
+def compile(
+    args: argparse.Namespace, qnn_config: QnnConfig, pte_name: str
+):  # noqa: C901
 
     # ensure the working directory exist.
     os.makedirs(args.artifact, exist_ok=True)
 
-    manager = get_qnn_llm_edge_manager(args.decoder_model, args.max_seq_len)
+    manager = get_qnn_llm_edge_manager(args.decoder_model_id, args.max_seq_len)
 
     fixed_point_type = {}
     if not args.use_fp16:
@@ -71,14 +59,11 @@ def compile(args: argparse.Namespace, qnn_config: QnnConfig):  # noqa: C901
         else:
             raise ValueError("Only support uint16 logits output for quantized hf llm.")
 
-        model_id = HUGGING_FACE_REPO_IDS[args.decoder_model]
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        tokenizer = AutoTokenizer.from_pretrained(args.decoder_model_id)
         tokenizer_json_path = tokenizer.save_pretrained(args.artifact)[-1]
 
         manager.pt2e_quantize(
             fixed_point_type,
-            args.calibration_tasks,
-            args.calibration_limit,
             args.prompt,
             tokenizer_json_path,
             qnn_config.backend,
@@ -97,15 +82,15 @@ def compile(args: argparse.Namespace, qnn_config: QnnConfig):  # noqa: C901
                 "scale": logits_quant_attrs["scale"],
                 "zero_point": logits_quant_attrs["zero_point"],
             },
-            open(f"{args.artifact}/{PTE_FILENAME}_quant_attrs.txt", "w"),
+            open(f"{args.artifact}/{pte_name}_quant_attrs.txt", "w"),
         )
 
-    manager.to_executorch(args.artifact, PTE_FILENAME)
+    manager.to_executorch(args.artifact, pte_name)
 
 
-def inference(args: argparse, qnn_config: QnnConfig):
-    workspace = f"/data/local/tmp/{getpass.getuser()}/executorch/{PTE_FILENAME}"
-    pte_path = f"{args.artifact}/{PTE_FILENAME}.pte"
+def inference(args: argparse, qnn_config: QnnConfig, pte_name: str):
+    workspace = f"/data/local/tmp/{getpass.getuser()}/executorch/{pte_name}"
+    pte_path = f"{args.artifact}/{pte_name}.pte"
     # collect output data
     output_data_folder = f"{args.artifact}/outputs"
     make_output_dir(output_data_folder)
@@ -121,12 +106,10 @@ def post_process():
             text = text[len(prefix) :]
         outputs.append(args.prompt + text)
 
-    model_id = HUGGING_FACE_REPO_IDS[args.decoder_model]
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer = AutoTokenizer.from_pretrained(args.decoder_model_id)
     tokenizer_json_path = tokenizer.save_pretrained(args.artifact)[-1]
     seq_len = args.max_seq_len
     runner_bin = "examples/qualcomm/oss_scripts/llama/qnn_llama_runner"
-    decoder_model_version = DECODER_MODEL_VERSION[args.decoder_model]
 
     # The base (non-instruct) HF models were not trained on the runner's chat
     # template. Tokenize the raw prompt here (matching the Python calibration
@@ -146,7 +129,6 @@ def post_process():
                 f"export LD_LIBRARY_PATH={qnn_sdk}/lib/{target}/:{args.build_folder}/lib &&",
                 f"{args.build_folder}/{runner_bin}",
                 f"--tokenized_prompt {tokenized_prompt_path}",
-                f"--decoder_model_version {decoder_model_version}",
                 "--eval_mode 0",
                 f"--tokenizer_path {tokenizer_json_path}",
                 f"--model_path {pte_path}",
@@ -168,10 +150,9 @@ def post_process():
                 f"cd {workspace} &&",
                 "./qnn_llama_runner",
                 "--tokenized_prompt tokenized_prompt.raw",
-                f"--decoder_model_version {decoder_model_version}",
                 "--eval_mode 0",
                 "--tokenizer_path tokenizer.json",
-                f"--model_path {PTE_FILENAME}.pte",
+                f"--model_path {pte_name}.pte",
                 f"--seq_len {seq_len}",
                 "--temperature 0",
                 "--output_path outputs/result.txt",
@@ -207,12 +188,12 @@ def main(args):
     qnn_config = QnnConfig.load_config(args.config_file if args.config_file else args)
 
     if args.compile_only:
-        compile(args, qnn_config)
+        compile(args, qnn_config, PTE_FILENAME)
     elif args.pre_gen_pte:
-        inference(args, qnn_config)
+        inference(args, qnn_config, PTE_FILENAME)
     else:
-        compile(args, qnn_config)
-        inference(args, qnn_config)
+        compile(args, qnn_config, PTE_FILENAME)
+        inference(args, qnn_config, PTE_FILENAME)
 
 
 if __name__ == "__main__":
@@ -242,10 +223,9 @@ def main(args):
     )
 
     parser.add_argument(
-        "--decoder_model",
-        choices=list(HUGGING_FACE_REPO_IDS.keys()),
-        help=f"The Hugging Face decoder model to export. Available options are: {list(HUGGING_FACE_REPO_IDS.keys())}",
-        required=True,
+        "--decoder_model_id",
+        help="The Hugging Face ID to export, e.g., 'NousResearch/Llama-3.2-1B'",
+        default="NousResearch/Llama-3.2-1B",
     )
 
     parser.add_argument(
@@ -254,31 +234,9 @@ def main(args):
         default=128,
         type=int,
     )
-    parser.add_argument(
-        "--calibration_tasks",
-        nargs="+",
-        type=str,
-        default=None,
-        help="Tasks for GPTQ calibration from lm_eval. Currently unsupported.",
-    )
-    parser.add_argument(
-        "--calibration_limit",
-        type=int,
-        default=None,
-        help="number of samples used for calibration from lm_eval. Currently unsupported.",
-    )
 
     try:
         args = parser.parse_args()
-
-        if args.calibration_tasks is not None or args.calibration_limit is not None:
-            parser.error(
-                "--calibration_tasks/--calibration_limit are not supported yet. "
-                "Calibration uses --prompt instead."
-            )
-
-        if args.artifact is None:
-            args.artifact = args.decoder_model
         main(args)
     except Exception as e:
         if args.ip and args.port != -1: