avx512 based int8 -> bf16 dequantization (#4912)

seanx92 · facebook-github-bot · commit fbfd59e54482 · 2025-09-22T13:40:45.000-07:00
Summary:

Use AVX512-bf16 intrinsics for int8 -&gt; bf16 dequantization

Differential Revision: D82507938
diff --git a/fbgemm_gpu/src/quantize_ops/quantize_ops_cpu.cpp b/fbgemm_gpu/src/quantize_ops/quantize_ops_cpu.cpp
@@ -220,10 +220,18 @@ Tensor& _fused8bitrowwise_to_float_cpu_out(
   return _fused8bitrowwise_to_float_cpu_out_t<float>(output, input);
 }
 
-Tensor& fused8bitrowwise_to_half_cpu_out(Tensor& output, const Tensor& input) {
+static Tensor& fused8bitrowwise_to_half_cpu_out(
+    Tensor& output,
+    const Tensor& input) {
   return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::float16>(output, input);
 }
 
+static Tensor& fused8bitrowwise_to_bfloat16_cpu_out(
+    Tensor& output,
+    const Tensor& input) {
+  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::bfloat16>(output, input);
+}
+
 /// @ingroup quantize-data-cpu
 ///
 Tensor& _float_to_fused8bitrowwise_cpu_out(
@@ -232,7 +240,9 @@ Tensor& _float_to_fused8bitrowwise_cpu_out(
   return _float_to_fused8bitrowwise_cpu_out_t<float>(output, input);
 }
 
-Tensor& _half_to_fused8bitrowwise_cpu_out(Tensor& output, const Tensor& input) {
+static Tensor& _half_to_fused8bitrowwise_cpu_out(
+    Tensor& output,
+    const Tensor& input) {
   return _float_to_fused8bitrowwise_cpu_out_t<fbgemm::float16>(output, input);
 }
 
@@ -285,6 +295,13 @@ Tensor fused8bitrowwise_to_half_cpu(const Tensor& input) {
   return fused8bitrowwise_to_half_cpu_out(output, input);
 }
 
+/// @ingroup quantize-data-cpu
+///
+Tensor fused8bitrowwise_to_bfloat16_cpu(const Tensor& input) {
+  auto output = at::empty({0}, input.options().dtype(at::kBFloat16));
+  return fused8bitrowwise_to_bfloat16_cpu_out(output, input);
+}
+
 /// @ingroup quantize-data-cpu
 ///
 Tensor fused8bitrowwise_to_float_or_half_cpu(
@@ -305,6 +322,10 @@ Tensor fused8bitrowwise_to_float_or_half_cpu(
       output = at::empty({0}, input.options().dtype(at::kHalf));
       output = fused8bitrowwise_to_half_cpu_out(output, input);
       break;
+    case SparseType::BF16:
+      output = at::empty({0}, input.options().dtype(at::kBFloat16));
+      output = fused8bitrowwise_to_bfloat16_cpu_out(output, input);
+      break;
     default:
       TORCH_CHECK(false);
   }
@@ -582,6 +603,7 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
       "FP8RowwiseQuantizedToFloat(Tensor input, bool forward, int output_dtype=0) -> Tensor",
       {PT2_COMPLIANT_TAG});
   m.def("Fused8BitRowwiseQuantizedToHalf(Tensor input) -> Tensor");
+  m.def("Fused8BitRowwiseQuantizedToBfloat16(Tensor input) -> Tensor");
   m.def(
       "Fused8BitRowwiseQuantizedToFloatOrHalf(Tensor input, int output_dtype=0, bool scale_bias_last=True, bool quant_padding_float_type=True) -> Tensor");
   m.def(
@@ -648,6 +670,9 @@ TORCH_LIBRARY_IMPL(fbgemm, CPU, m) {
   DISPATCH_TO_CPU(
       "Fused8BitRowwiseQuantizedToHalf",
       fbgemm_gpu::fused8bitrowwise_to_half_cpu);
+  DISPATCH_TO_CPU(
+      "Fused8BitRowwiseQuantizedToBfloat16",
+      fbgemm_gpu::fused8bitrowwise_to_bfloat16_cpu);
   DISPATCH_TO_CPU(
       "Fused8BitRowwiseQuantizedToFloatOrHalf",
       fbgemm_gpu::fused8bitrowwise_to_float_or_half_cpu);
diff --git a/fbgemm_gpu/test/quantize/fused_8bit_rowwise_test.py b/fbgemm_gpu/test/quantize/fused_8bit_rowwise_test.py
@@ -141,9 +141,9 @@ def quantize_and_dequantize_op_test_helper(  # noqa: C901
 
         assume(ncols % (2 * num_elem_per_byte) == 0)
         if not test_cuda:
-            # cpu path does not support bf16
+            # cpu path only supports bf16 dequantization
             if output_dtype == SparseType.BF16:
-                return
+                input_data = input_data.float()
             if test_generic_op:
                 quantized_data = (
                     torch.ops.fbgemm.FloatOrHalfToFused8BitRowwiseQuantized(input_data)
@@ -171,6 +171,15 @@ def quantize_and_dequantize_op_test_helper(  # noqa: C901
                     dequantized_data = torch.ops.fbgemm.Fused8BitRowwiseQuantizedToHalf(
                         quantized_data
                     )
+                elif output_dtype == SparseType.BF16:
+                    quantized_data = torch.ops.fbgemm.FloatToFused8BitRowwiseQuantized(
+                        input_data,
+                    )
+                    dequantized_data = (
+                        torch.ops.fbgemm.Fused8BitRowwiseQuantizedToBfloat16(
+                            quantized_data,
+                        )
+                    )
                 else:
                     raise NotImplementedError("Unsupported dtype")
 
@@ -185,6 +194,10 @@ def quantize_and_dequantize_op_test_helper(  # noqa: C901
                 torch.testing.assert_close(dequantized_data.float(), reference.float())
             elif output_dtype == SparseType.FP16:
                 torch.testing.assert_close(dequantized_data.half(), reference.half())
+            elif output_dtype == SparseType.BF16:
+                torch.testing.assert_close(
+                    dequantized_data.bfloat16(), reference.bfloat16()
+                )
         if test_cuda and gpu_available:
             if nrows == 0 or ncols == 0:
                 return
diff --git a/include/fbgemm/QuantUtils.h b/include/fbgemm/QuantUtils.h
@@ -10,6 +10,7 @@
 
 #include "./FbgemmBuild.h" // @manual
 #include "./QuantUtilsAvx2.h" // @manual
+#include "./QuantUtilsAvx512.h" // @manual
 #include "./QuantUtilsNeon.h" // @manual
 #include "./Types.h" // @manual
 #include "./Utils.h" // @manual
diff --git a/include/fbgemm/QuantUtilsAvx512.h b/include/fbgemm/QuantUtilsAvx512.h
@@ -37,6 +37,13 @@ FBGEMM_API void requantizeOutputProcessingGConvAvx512(
     int ld_out,
     int ld_in,
     const requantizationParams_t<BIAS_TYPE>& r);
+
+template <typename OutputType>
+void Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512(
+    const std::uint8_t* input,
+    size_t input_rows,
+    int input_columns,
+    OutputType* output);
 } // namespace fbgemm
 
 #endif
diff --git a/include/fbgemm/Utils.h b/include/fbgemm/Utils.h
@@ -177,6 +177,11 @@ FBGEMM_API bool fbgemmHasAvx2Support();
  */
 FBGEMM_API bool fbgemmHasAvx512VnniSupport();
 
+/**
+ * @brief Are we running on a AVX512_BF16 supported cpu?
+ */
+FBGEMM_API bool fbgemmHasAvx512Bf16Support();
+
 /**
  * @brief Are we running on a ARM Neon supported cpu?
  */
diff --git a/src/QuantUtils.cc b/src/QuantUtils.cc
@@ -825,6 +825,8 @@ void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef(
           input_row[col] * input_row_scale_bias[0] + input_row_scale_bias[1];
       if constexpr (std::is_same<OutputType, float>()) {
         output_row[col] = output_value;
+      } else if constexpr (std::is_same_v<OutputType, bfloat16>) {
+        output_row[col] = cpu_float2bfloat16(output_value);
       } else {
         output_row[col] = cpu_float2half_rn(output_value);
       }
@@ -842,15 +844,24 @@ void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
   Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfNeon<OutputType>(
       input, input_rows, input_columns, output);
 #else
-  if (cpuinfo_initialize() && fbgemmHasAvx2Support()) {
 #if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
-    Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfAvx2<OutputType>(
-        input, input_rows, input_columns, output);
-#endif
-  } else {
-    Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef<OutputType>(
-        input, input_rows, input_columns, output);
+  if (cpuinfo_initialize() && fbgemmHasAvx2Support()) {
+    if (fbgemmHasAvx512Bf16Support() && std::is_same_v<OutputType, bfloat16>) {
+      // Avx512 bfloat16 native support
+      Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512<OutputType>(
+          input, input_rows, input_columns, output);
+      return;
+    } else if (!std::is_same_v<OutputType, bfloat16>) {
+      // Avx2 does not support bfloat16
+      Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfAvx2<OutputType>(
+          input, input_rows, input_columns, output);
+      return;
+    }
   }
+#endif
+  // Fallback to ref kernel
+  Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef<OutputType>(
+      input, input_rows, input_columns, output);
 #endif
 }
 
diff --git a/src/QuantUtilsAvx512.cc b/src/QuantUtilsAvx512.cc
@@ -12,6 +12,7 @@
     (defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)))
 #include <immintrin.h>
 #endif
+#include <fbgemm/FloatConversion.h>
 #include <cassert>
 
 namespace fbgemm {
@@ -381,6 +382,46 @@ void requantizeOutputProcessingGConvAvx512(
   } // i loop
 }
 
+template <typename OutputType>
+void Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512(
+    const std::uint8_t* input,
+    size_t input_rows,
+    int input_columns,
+    OutputType* output) {
+  constexpr int VLEN = 8;
+  int output_columns = input_columns - 2 * sizeof(float);
+
+  for (size_t row = 0; row < input_rows; ++row) {
+    const std::uint8_t* input_row = input + row * input_columns;
+    const float* input_row_scale_bias =
+        reinterpret_cast<const float*>(input_row + output_columns);
+    OutputType* output_row = output + row * output_columns;
+
+    __m256 scale_v = _mm256_set1_ps(input_row_scale_bias[0]);
+    __m256 bias_v = _mm256_set1_ps(input_row_scale_bias[1]);
+
+    int col = 0;
+    for (col = 0; col < output_columns / VLEN * VLEN; col += VLEN) {
+      __m256 in_v = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(
+          _mm_loadl_epi64(reinterpret_cast<const __m128i*>(input_row + col))));
+#ifdef __FMA__
+      __m256 dequantzed_v = _mm256_fmadd_ps(in_v, scale_v, bias_v);
+#else
+      __m256 dequantzed_v = _mm256_add_ps(_mm256_mul_ps(in_v, scale_v), bias_v);
+#endif
+      _mm_storeu_si128(
+          reinterpret_cast<__m128i*>(output_row + col),
+          _mm256_cvtneps_pbh(dequantzed_v));
+    }
+
+    for (; col < output_columns; ++col) {
+      float output_value =
+          input_row[col] * input_row_scale_bias[0] + input_row_scale_bias[1];
+      output_row[col] = cpu_float2bfloat16(output_value);
+    }
+  } // for each row
+}
+
 #define INSTANTIATE_REQUANTIZE_BIAS_TYPE(              \
     A_SYM, B_SYM, Q_GRAN, BIAS, RELU, BIAS_TYPE)       \
   template void requantizeOutputProcessingGConvAvx512< \
@@ -468,4 +509,11 @@ INSTANTIATE_BIAS(false)
 #undef INSTANTIATE_B_SYM
 #undef INSTANTIATE_Q_GRANS
 #undef INSTANTIATE_BIAS
+
+template void Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512<bfloat16>(
+    const std::uint8_t* input,
+    size_t input_rows,
+    int input_columns,
+    bfloat16* output);
+
 } // namespace fbgemm
diff --git a/src/Utils.cc b/src/Utils.cc
@@ -319,6 +319,10 @@ bool fbgemmHasAvx512VnniSupport() {
   return cpuinfo_has_x86_avx512vnni();
 }
 
+bool fbgemmHasAvx512Bf16Support() {
+  return cpuinfo_has_x86_avx512bf16();
+}
+
 bool fbgemmHasArmNeonSupport() {
   return cpuinfo_has_arm_neon();
 }

Original file line number	Diff line number	Diff line change
`@@ -319,6 +319,10 @@ bool fbgemmHasAvx512VnniSupport() {`
`319`	`319`	`return cpuinfo_has_x86_avx512vnni();`
`320`	`320`	`}`
`321`	`321`
	`322`	`+bool fbgemmHasAvx512Bf16Support() {`
	`323`	`+ return cpuinfo_has_x86_avx512bf16();`
	`324`	`+}`
	`325`	`+`
`322`	`326`	`bool fbgemmHasArmNeonSupport() {`
`323`	`327`	`return cpuinfo_has_arm_neon();`
`324`	`328`	`}`