Quantize embedding op and 16x8 matmul

backends/qualcomm/quantizer/custom_annotation.py

@@ -118,3 +118,29 @@ def annotate_matmul_input1(node: Node, quantization_config: QuantizationConfig):
                 if "SDPA" in full_qualified_name:
                     annotate_matmul(node, quantization_config_16a8w)
                     annotate_matmul_input1(node.args[1], quantization_config_8a8w)
+
+
+def custom_annotate_matmul_16a8w(gm: torch.fx.GraphModule):
+    """
+    Annotate matmul op with 16a8w quantization config
+    """
+
+    def annotate_matmul(node: Node, quantization_config: QuantizationConfig):
+        input_qspec_map = {}
+        input_act = node.args[0]
+        input_spec = quantization_config.input_activation
+        input_qspec_map[input_act] = input_spec
+        input_act1 = node.args[1]
+        input_spec1 = quantization_config.weight
+        input_qspec_map[input_act1] = input_spec1
+        node.meta[QUANT_ANNOTATION_KEY] = QuantizationAnnotation(
+            input_qspec_map=input_qspec_map,
+            output_qspec=quantization_config.output_activation,
+            _annotated=True,
+        )
+
+    # Annotate 16a8w for matmul op to get better performance
+    quantization_config_16a8w = get_16a8w_qnn_ptq_config()
+    for node in gm.graph.nodes:
+        if node.op == "call_function" and node.target == torch.ops.aten.matmul.default:
+            annotate_matmul(node, quantization_config_16a8w)

examples/models/llama2/export_llama_lib.py

@@ -53,21 +53,23 @@
     get_quant_embedding_transform,
     get_quant_weight_transform,
 )
-from .source_transformation.quantized_kv_cache import (
-    replace_kv_cache_with_quantized_kv_cache,
-)
+
+# from .source_transformation.quantized_kv_cache import (
+#     replace_kv_cache_with_quantized_kv_cache,
+# )
 from .source_transformation.rms_norm import replace_rms_norm_with_native_rms_norm
 
 from .source_transformation.rope import materialze_broadcast_of_rope_freq_cis
-from .source_transformation.sdpa import (
-    replace_causal_mask,
-    replace_kv_cache_with_coreml_kv_cache,
-    replace_kv_cache_with_simple_kv_cache,
-    replace_sdpa_with_coreml_sdpa,
-    replace_sdpa_with_custom_op,
-    replace_sdpa_with_flex_sdpa,
-    replace_sdpa_with_simple_sdpa,
-)
+
+# from .source_transformation.sdpa import (
+#     replace_causal_mask,
+#     replace_kv_cache_with_coreml_kv_cache,
+#     replace_kv_cache_with_simple_kv_cache,
+#     replace_sdpa_with_coreml_sdpa,
+#     replace_sdpa_with_custom_op,
+#     replace_sdpa_with_flex_sdpa,
+#     replace_sdpa_with_simple_sdpa,
+# )
 
 IS_FBCODE = True  #  os.environ.get("FBCODE_PLATFORM", False)
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
@@ -524,7 +526,7 @@ def get_quantizer_and_quant_params(args):
     if args.qnn and args.pt2e_quantize:
         assert len(quantizers) == 0, "Should not enable both xnnpack and qnn"
         qnn_quantizer, quant_dtype = get_qnn_quantizer(
-            args.pt2e_quantize, args.quantization_mode
+            args.pt2e_quantize, args.use_kv_cache, args.quantization_mode
         )
         quantizers.append(qnn_quantizer)
     if args.coreml and args.pt2e_quantize:
@@ -893,23 +895,20 @@ def _get_source_transforms(  # noqa
         assert args.use_kv_cache, "quantize_kv_cache requires use_kv_cache=True"
         transforms.append(replace_kv_cache_with_quantized_kv_cache)
 
+    if args.qnn:
+        # 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 convert_linear_to_conv2d
+
+        # transforms.append(replace_kv_cache_with_simple_kv_cache)
+        # transforms.append(replace_sdpa_with_flex_sdpa)
+        # transforms.append(replace_causal_mask)
+        transforms.append(replace_rms_norm_with_native_rms_norm)
+        if args.optimized_rotation_path:
+            transforms.append(fuse_layer_norms)
+            transforms.append(get_model_with_r1_r2(args.optimized_rotation_path))
+        transforms.append(convert_linear_to_conv2d)
     if args.use_kv_cache:
-        if args.qnn:
-            # 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 (
-                convert_linear_to_conv2d,
-            )
-
-            transforms.append(replace_kv_cache_with_simple_kv_cache)
-            transforms.append(replace_sdpa_with_flex_sdpa)
-            transforms.append(replace_causal_mask)
-            transforms.append(replace_rms_norm_with_native_rms_norm)
-            if args.optimized_rotation_path:
-                transforms.append(fuse_layer_norms)
-                transforms.append(get_model_with_r1_r2(args.optimized_rotation_path))
-            transforms.append(convert_linear_to_conv2d)
-
-        elif args.mps:
+        if args.mps:
             # Currently mps doesn't support sdpa op, use the simpler decomposition
             # to get free perf gain.
             transforms.append(replace_sdpa_with_simple_sdpa)