[Bug fix][CPU] Fix fp8 sdpa compiling issue with latest PyTorch (#2991)

* [CPU] Fix fp8 sdpa compiling issue with latest pytorch

* disable fp8 fusion

Author: Valentine233

torchao/csrc/cpu/aten_kernels/quantized_sdpa.cpp

@@ -1775,6 +1775,7 @@ int8_sdpa_fused_kernel_impl(
   at::native::cpublas::brgemm_release();
 }
 
+#if defined(CPUBLAS_BRGEMM_F8F8F32)
 // FP8 - kernel with f8f8f8 GEMM
 template <typename scalar_t, typename mask_t,
           int64_t q_split_size, int64_t kv_split_size>
@@ -2136,6 +2137,7 @@ fp8_sdpa_fused_kernel_impl(
     at::native::cpublas::brgemm_release();
   });
 }
+#endif // CPUBLAS_BRGEMM_F8F8F32
 
 template <typename scalar_t, typename mask_t, int64_t q_split_size, int64_t kv_split_size>
 inline typename std::enable_if_t<std::is_same_v<scalar_t, unsigned char>, void>
@@ -2304,6 +2306,7 @@ void int8_sdpa_fused_kernel(
   }
 }
 
+#if defined(CPUBLAS_BRGEMM_F8F8F32)
 void fp8_sdpa_fused_kernel(
     const at::Tensor& output,
     const at::Tensor& query,
@@ -2380,6 +2383,7 @@ void fp8_sdpa_fused_kernel(
     });
   }
 }
+#endif // CPUBLAS_BRGEMM_F8F8F32
 #endif // CPU_CAPABILITY_AVX512
 
 at::Tensor int8_sdpa_math_kernel(

torchao/prototype/inductor/fx_passes/qsdpa_fusion.py

@@ -28,7 +28,7 @@
 ]
 
 aten = torch.ops.aten
-quantize_dtypes = [torch.uint8, torch.float8_e4m3fn]
+quantize_dtypes = [torch.uint8]
 
 
 def _is_valid_qsdpa_pattern():
@@ -121,53 +121,31 @@ def qsdpa(match: Match, *args, **kwargs):
 def _generate_dequant_pattern(
     input_pattern, qtype, is_reduced_type, scale: str, zp: str = None
 ):
-    if qtype == torch.uint8:
-        assert zp is not None, "Zero point must be provided for uint8 dequantization"
-        return CallFunction(
-            torch.ops.quantized_decomposed.dequantize_per_tensor.default,
-            input_pattern,
-            KeywordArg(scale),
-            KeywordArg(zp),
-            Arg(),
-            Arg(),
-            Arg(),
-        )
-    else:
-        assert zp is None, "Fp8 dequantization does not support zero point"
-        if is_reduced_type:
-            return CallFunction(
-                torch.ops.torchao.dequantize_affine_float8.default,
-                input_pattern,
-                KeywordArg(scale),
-                Arg(),
-            )
-        else:
-            return CallFunction(
-                torch.ops.torchao.dequantize_affine_float8.default,
-                input_pattern,
-                KeywordArg(scale),
-            )
+    assert qtype is torch.uint8, "QSDPA expects type to be uint8"
+    assert zp is not None, "Zero point must be provided for uint8 dequantization"
+    return CallFunction(
+        torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+        input_pattern,
+        KeywordArg(scale),
+        KeywordArg(zp),
+        Arg(),
+        Arg(),
+        Arg(),
+    )
 
 
 def _generate_quant_pattern(input_pattern, qtype, scale: str, zp: str = None):
-    if qtype == torch.uint8:
-        assert zp is not None, "Zero point must be provided for uint8 quantization"
-        return CallFunction(
-            torch.ops.quantized_decomposed.quantize_per_tensor.default,
-            input_pattern,
-            KeywordArg(scale),
-            KeywordArg(zp),
-            Arg(),
-            Arg(),
-            Arg(),
-        )
-    else:
-        assert zp is None, "Fp8 quantization does not support zero point"
-        return CallFunction(
-            torch.ops.torchao.quantize_affine_float8.default,
-            input_pattern,
-            KeywordArg(scale),
-        )
+    assert qtype is torch.uint8, "QSDPA expects type to be uint8"
+    assert zp is not None, "Zero point must be provided for uint8 quantization"
+    return CallFunction(
+        torch.ops.quantized_decomposed.quantize_per_tensor.default,
+        input_pattern,
+        KeywordArg(scale),
+        KeywordArg(zp),
+        Arg(),
+        Arg(),
+        Arg(),
+    )
 
 
 def _get_qsdpa_qkv_pattern(