support int8 scaled_embedding_bag pattern_match

Author: shiyang-weng

torchao/quantization/pt2e/inductor_passes/x86.py

@@ -2860,7 +2860,11 @@ def _register_scaled_embedding_bag_pass(pattern, pass_number, dtype=torch.float3
     def scaled_embedding_bag(match: Match, *args, **kwargs):
         assert dtype in [torch.float32, torch.bfloat16]
 
-        getitem_node = match.output_node()
+        if "o_dtype" in kwargs:
+            quant_node = match.output_node()
+            getitem_node = quant_node.args[0]
+        else:
+            getitem_node = match.output_node()
         embedding_bag_node = getitem_node.args[0]
         assert embedding_bag_node.target is aten._embedding_bag_forward_only.default
 
@@ -2889,11 +2893,25 @@ def scaled_embedding_bag(match: Match, *args, **kwargs):
             kwargs["mode"],
             kwargs["include_last_offset"],
         )
-        # only support fp32 output, next step to support more dtype
+        # only support fp32 and int8 output on kernel
+        # next step to support more output_type
+        output_type = torch.float
         o_scale = 1.0
+        if "o_dtype" in kwargs:
+            output_type = torch.int8
+            o_scale = kwargs["o_inv_scale"]
 
         graph = match.graph
         with graph.inserting_before(getitem_node):
+            # scale type is float on int8 q/dq
+            # Not support float scale yet on scaled_embedding_bag
+            # convert scale from float into tensor
+            if output_type == torch.int8:
+                w_scale = graph.call_function(
+                    torch.ops.aten.full.default,
+                    args=([1], w_scale),
+                    kwargs={"dtype": torch.float},
+                )
             new_args: tuple[Any, ...] = (
                 qw,
                 indices,
@@ -2902,13 +2920,18 @@ def scaled_embedding_bag(match: Match, *args, **kwargs):
                 o_scale,
                 mode,
                 include_last_offset,
-                torch.float,
+                output_type,
             )
 
             new_embedding_bag_node = graph.call_function(
                 torch.ops.torchao._scaled_embedding_bag.default, args=new_args
             )
 
+            # remove quant node
+            if output_type == torch.int8:
+                quant_node.replace_all_uses_with(getitem_node)
+                getitem_node.meta.update(quant_node.meta)
+                graph.erase_node(quant_node)
             getitem_node.replace_all_uses_with(new_embedding_bag_node)
             new_embedding_bag_node.meta.update(embedding_bag_node.meta)
 
@@ -2943,20 +2966,37 @@ def _generate_scaled_embedding_bag_patterns(dq_pattern):
 
 
 def _register_quantization_embeddingbag_pass():
-    for dtype in [torch.float32, torch.bfloat16]:
-        _register_scaled_embedding_bag_pass(
-            _generate_scaled_embedding_bag_patterns(
+    for is_fp8 in [True, False]:
+        for dtype in [torch.float32, torch.bfloat16]:
+            embeddingbag_pattern = _generate_scaled_embedding_bag_patterns(
                 _may_generate_pattern_with_dtype_convert(
                     get_dequantize_per_tensor_activation_pattern(
-                        is_tensor_overload=False, is_fp8=True
+                        is_tensor_overload=False, is_fp8=is_fp8
                     ),
                     KeywordArg("autocast_act_dtype"),
                     dtype == torch.bfloat16,
                 ),
-            ),
-            pass_number=1,
-            dtype=dtype,
-        )  # pass_number=0 to run before weight prepack
+            )
+
+            _register_scaled_embedding_bag_pass(
+                embeddingbag_pattern,
+                pass_number=1,
+                dtype=dtype
+            )
+
+            # will support fp8 output later
+            if not is_fp8:
+                embeddingbag_with_qoutput_pattern = generate_pattern_with_output_quant(
+                    embeddingbag_pattern,
+                    dtype == torch.bfloat16,
+                    is_fp8,
+                )
+
+                _register_scaled_embedding_bag_pass(
+                    embeddingbag_with_qoutput_pattern,
+                    pass_number=0,
+                    dtype=dtype,
+                )
 
 
 @functools.lru_cache(None)