update int8 sdpa

Author: Valentine233

test/prototype/inductor/test_int8_sdpa_fusion.py

@@ -1,5 +1,4 @@
 import itertools
-import os
 import unittest
 
 import torch
@@ -12,11 +11,12 @@
 from torch.testing._internal.inductor_utils import HAS_CPU
 
 import torchao
-from torchao.prototype.inductor.fx_passes.int8_sdpa_fusion import _int8_sdpa_init
+from torchao.prototype.inductor.fx_passes.int8_sdpa_fusion import (
+    _int8_sdpa_init,
+    custom_pass,
+)
 from torchao.utils import TORCH_VERSION_AT_LEAST_2_7
 
-use_cpp_avx512 = os.getenv("USE_AVX512", "0") == "1"
-
 
 class SelfAttnLikeModule(torch.nn.Module):
     def __init__(
@@ -147,7 +147,10 @@ def _check_common(
     @unittest.skipIf(
         not TORCH_VERSION_AT_LEAST_2_7, reason="int8 sdpa requires torch 2.7 or later"
     )
-    @unittest.skipIf(not use_cpp_avx512, reason="cpp kernels not built")
+    @unittest.skipIf(
+        "CPU" not in torch._C._dispatch_dump("torchao::qscaled_dot_product"),
+        reason="cpp kernels not built",
+    )
     @config.patch({"freezing": True})
     def _test_sdpa_int8_rewriter(self):
         from torch.export import export_for_training
@@ -184,6 +187,7 @@ def _test_sdpa_int8_rewriter(self):
                 torch.amp.autocast(
                     self.device, enabled=enable_autocast, dtype=torch.bfloat16
                 ),
+                config.patch(post_grad_custom_pre_pass=custom_pass),
             ):
                 _int8_sdpa_init()
                 quantizer = X86InductorQuantizer()

test/test_ops.py

@@ -5,7 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 import itertools
 import math
-import os
 import sys
 
 import pytest
@@ -32,8 +31,6 @@
     compute_max_diff,
 )
 
-use_cpp_avx512 = os.getenv("USE_AVX512", "0") == "1"
-
 if torch.version.hip is not None:
     pytest.skip("Skipping the test in ROCm", allow_module_level=True)
 
@@ -160,7 +157,10 @@ def _scaled_dot_product_int8_op_ref(
         not TORCH_VERSION_AT_LEAST_2_7, reason="int8 sdpa requires torch 2.7 or later"
     )
     @pytest.mark.skipif(not IS_LINUX, reason="only support on linux")
-    @pytest.mark.skipif(not use_cpp_avx512, reason="cpp kernels not built")
+    @pytest.mark.skipif(
+        "CPU" not in torch._C._dispatch_dump("torchao::qscaled_dot_product"),
+        reason="cpp kernels not built",
+    )
     @parametrize("batch_size", [56, 120])
     @parametrize("n_head", [2, 16])
     @parametrize("q_seq_len", [18, 89])

torchao/prototype/inductor/fx_passes/README.md

@@ -23,8 +23,9 @@ from torch._inductor.pattern_matcher import PatternMatcherPass
 
 # Example usage
 class _CustomPass(...): # create a custom pass class
-config.custom_pass = _CustomPass() # define the custom pass with the patterns
-_register_patterns(...) # register your own patterns
+custom_pass = _CustomPass() # create an instance of custom pass
+with config.patch(config.custom_pass=custom_pass):
+    _register_patterns(config.custom_pass) # register your own passes
 
 ```
 

torchao/prototype/inductor/fx_passes/int8_sdpa_fusion.py

@@ -365,6 +365,7 @@ def _register_int8_sdpa_lowerings(custom_pass_dict):
         )
 
 
+custom_pass = None
 if TORCH_VERSION_AT_LEAST_2_7:
     # TORCH_VERSION_AT_LEAST_2_7 is needed for custom graph pass
     from torch._inductor.custom_graph_pass import CustomGraphPass, get_hash_for_files
@@ -380,11 +381,12 @@ def __call__(self, g: torch.fx.graph.Graph):
         def uuid(self) -> bytes:
             return get_hash_for_files((__file__,))
 
+    custom_pass = _CustomPass()
+
 
 @functools.lru_cache(None)
 def _int8_sdpa_init():
     if TORCH_VERSION_AT_LEAST_2_7:
-        config.post_grad_custom_pre_pass = _CustomPass()
         _register_int8_sdpa_lowerings(config.post_grad_custom_pre_pass)
     else:
         pass