[ONLY FOR TEST] test macos whl issue

setup.py

@@ -55,6 +55,28 @@ def read_version(file_path="version.txt"):
     and platform.system() == "Darwin"
 )
 
+use_cpp_avx512 = os.getenv("USE_AVX512", "1") == "1" and platform.system() == "Linux"
+
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_7
+import subprocess
+
+# Run 'uname -a'
+uname_output = subprocess.run(['uname', '-a'], capture_output=True, text=True)
+print("Output of 'uname -a':")
+print(uname_output.stdout)
+
+if platform.system() == "Darwin":
+    sysctl_output = subprocess.run(['sysctl', '-a'], capture_output=True, text=True)
+    print("\nFiltered CPU info from 'sysctl -a':")
+    for line in sysctl_output.stdout.splitlines():
+        if 'machdep.cpu' in line:
+            print(line)
+else:
+    # Run 'lscpu'
+    lscpu_output = subprocess.run(['lscpu'], capture_output=True, text=True)
+    print("\nOutput of 'lscpu':")
+    print(lscpu_output.stdout)
+
 version_prefix = read_version()
 # Version is version.dev year month date if using nightlies and version if not
 version = (
@@ -284,6 +306,17 @@ def get_extensions():
             ["-O3" if not debug_mode else "-O0", "-fdiagnostics-color=always"]
         )
 
+        if use_cpp_avx512 and TORCH_VERSION_AT_LEAST_2_7:
+            if torch._C._cpu._is_avx512_supported():
+                extra_compile_args["cxx"].extend(
+                    [
+                        "-DCPU_CAPABILITY_AVX512",
+                        "-march=native",
+                        "-mfma",
+                        "-fopenmp",
+                    ]
+                )
+
         if debug_mode:
             extra_compile_args["cxx"].append("-g")
             if "nvcc" in extra_compile_args:

torchao/csrc/cpu/int8_sdpa.cpp

@@ -0,0 +1,12 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/Dispatch.h>
+#include <ATen/core/Tensor.h>
+#include <ATen/cpu/vec/functional.h>
+#include <ATen/cpu/vec/vec.h>
+#include <ATen/cpu/Utils.h>
+#include <ATen/native/cpu/utils.h>
+#include <ATen/native/CPUBlas.h>
+#include <ATen/Parallel.h>
+#include <ATen/Tensor.h>
+#include <c10/util/irange.h>