Added soc model param to get_qnn_partitioner

Author: Guang Yang

examples/models/llama2/export_llama_lib.py

@@ -510,11 +510,18 @@ def _export_llama(modelname, args) -> LLMEdgeManager:  # noqa: C901
         modelname = f"coreml_{modelname}"
 
     if args.qnn:
+        # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.qualcomm.serialization.qnn_compile_spec_schema`
+        from executorch.backends.qualcomm.serialization.qnn_compile_spec_schema import (
+            QcomChipset,
+        )
         from executorch.extension.llm.custom_ops import model_sharding
 
         partitioners.append(
             get_qnn_partitioner(
-                args.use_kv_cache, args.pt2e_quantize, args.num_sharding
+                QcomChipset.SM8650,  # Llama 2 works only on SM8650
+                args.use_kv_cache,
+                args.pt2e_quantize,
+                args.num_sharding,
             )
         )
         # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.qualcomm.utils.utils`

extension/llm/export/partitioner_lib.py

@@ -105,6 +105,7 @@ def get_coreml_partitioner(
 
 
 def get_qnn_partitioner(
+    soc_model,
     use_kv_cache: bool = False,
     pt2e_quantize: Optional[str] = None,
     num_sharding: int = 0,
@@ -118,11 +119,6 @@ def get_qnn_partitioner(
             QnnPartitioner,
         )
 
-        # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.qualcomm.serialization.qnn_compile_spec_schema`
-        from executorch.backends.qualcomm.serialization.qnn_compile_spec_schema import (
-            QcomChipset,
-        )
-
         # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.qualcomm.utils.utils`
         from executorch.backends.qualcomm.utils.utils import (
             generate_htp_compiler_spec,
@@ -140,7 +136,7 @@ def get_qnn_partitioner(
 
     return QnnPartitioner(  # pyre-fixme[16]
         generate_qnn_executorch_compiler_spec(  # pyre-fixme[16]
-            soc_model=QcomChipset.SM8450,  # default to SM8450  # pyre-fixme[16]
+            soc_model=soc_model,  # pyre-fixme[16]
             # pyre-fixme[16]
             backend_options=generate_htp_compiler_spec(
                 use_fp16=use_fp16,