[X86] Bug fixes and refinements in da8w4/float8 linear kernels (#4301)

* [X86] Refine da8w4/float8 linear kernels

* Fix correctness bugs in fallback path

* Refine code

---------

Co-authored-by: Copilot <copilot@github.com>

Author: Xia-Weiwen

torchao/csrc/cpu/aten_kernels/da8w4_linear.cpp

@@ -9,16 +9,14 @@ namespace {
 
 #define BLOCK_N 32
 
-static bool cpublas_checked = false;
+static std::once_flag cpublas_once;
 static bool cpublas_can_pack = false;
 
-bool cpublas_could_pack() {
+static inline bool cpublas_could_pack() {
   // the could_pack check requires AMX support implicitly
-  if (cpublas_checked) {
-    return cpublas_can_pack;
-  }
-  cpublas_can_pack = at::native::cpublas::could_pack(at::kByte);
-  cpublas_checked = true;
+  std::call_once(cpublas_once, []() {
+    cpublas_can_pack = at::native::cpublas::could_pack(at::kByte);
+  });
   return cpublas_can_pack;
 }
 
@@ -135,7 +133,6 @@ struct ActDtype<false> {
   using type = uint8_t;
 };
 
-
 #if defined(CPU_CAPABILITY_AVX512)
 inline std::array<__m256i, 2> load_zps_4vnni(const int8_t* __restrict__ zps) {
   // broadcast 01234567 to
@@ -286,8 +283,8 @@ void _dequant_weight_zp_only(
   for (int k = 0; k < K; ++k) {
     for (int n = 0; n < N / 2; ++n) {
       int32_t b = (int32_t)B[k * ldb + n];
-      dqB[k * N + n * 2] = (b & 0xf) - qzeros[n];
-      dqB[k * N + n * 2 + 1] = (b >> 4) - qzeros[n];
+      dqB[k * N + n * 2] = (b & 0xf) - qzeros[n * 2];
+      dqB[k * N + n * 2 + 1] = ((b >> 4) & 0xf) - qzeros[n * 2 + 1];
     }
   }
 }
@@ -407,7 +404,6 @@ void _dequant_gemm_accum_small_M(
     _mm512_storeu_ps(C + row * ldc + col * 16, vc_float);
   };
   c10::ForcedUnroll<M * COLS>{}(store);
-
 }
 
 #define call_dequant_gemm_accum_small_M(M) \

torchao/csrc/cpu/aten_kernels/float8_linear.cpp

@@ -19,20 +19,17 @@ namespace {
 #define PER_ROW 2
 #define PER_GROUP 3
 
-static bool cpublas_checked = false;
+static std::once_flag cpublas_flag;
 static bool cpublas_can_pack = false;
 
-bool cpublas_could_pack() {
-  // the could_pack check requires AMX support implicitly
-  if (cpublas_checked) {
-    return cpublas_can_pack;
-  }
+static inline bool cpublas_could_pack() {
+  std::call_once(cpublas_flag, []() {
 #ifdef CPUBLAS_BRGEMM_F8F8F32
-  cpublas_can_pack = at::native::cpublas::could_pack(at::kFloat8_e4m3fn);
+    cpublas_can_pack = at::native::cpublas::could_pack(at::kFloat8_e4m3fn);
 #else
-  cpublas_can_pack = at::native::cpublas::could_pack(at::kBFloat16);
+    cpublas_can_pack = at::native::cpublas::could_pack(at::kBFloat16);
 #endif
-  cpublas_checked = true;
+  });
   return cpublas_can_pack;
 }
 
@@ -124,59 +121,34 @@ float8_linear_prepack_impl(
 }
 
 #if defined(CPU_CAPABILITY_AVX512)
-// this doesn't handle NaN.
-inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) {
-  const __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
-
-  const __m512i mant = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x07)), 4);
-  const __m512i raw_exp = _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x78)), 3);
-  const __m512i exp = _mm512_slli_epi16(_mm512_add_epi16(raw_exp, _mm512_set1_epi16(120)), 7);
-  const __m512i nonsign = _mm512_or_si512(exp, mant);
-
-  const __m512i sign = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x80)), 8);
-  const __m512i combined = _mm512_or_si512(nonsign, sign);
-
-  const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512());
-  return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined);
-}
-
 static void cvt_f8e4m3_to_bf16(
-    const at::Float8_e4m3fn* __restrict__ in,
-    at::BFloat16* out,
-    int64_t rows,
-    int64_t cols,
-    int64_t stride) {
-  if (stride == cols) {
-    // A contiguous buffer
-    size_t len = rows * cols;
-    size_t i = 0;
-    for (; i < len; i += 32) {
-      __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[i]);
-      __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec);
-      _mm512_storeu_si512((__m512i*)(out + i), (__m512i)bf16_vec);
-    }
-    for (; i < len; ++i) {
-      out[i] = (at::BFloat16)in[i];
-    }
-  } else {
-    // Non-contiguous. Access each row with stride
-    TORCH_CHECK(stride > cols);
-    for (int r = 0; r < rows; ++r) {
-      size_t i = 0;
-      size_t vec_len = cols / 32 * 32;
-      for (; i < vec_len; i += 32) {
-        __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[r * stride + i]);
-        __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec);
-        _mm512_storeu_si512((__m512i*)(out + r * cols + i), (__m512i)bf16_vec);
-      }
-      for (; i < cols; ++i) {
-        out[r * cols + i] = (at::BFloat16)in[r * stride + i];
-      }
-    }
+  const at::Float8_e4m3fn* __restrict__ in,
+  at::BFloat16* out,
+  int64_t rows,
+  int64_t cols,
+  int64_t stride) {
+  constexpr int64_t vec_len = 32; // 256 bit = 32 fp8 values
+  __m512 fp32_vec_0, fp32_vec_1;
+  for (int r = 0; r < rows; ++r) {
+  size_t i = 0;
+  size_t vec_len_aligned = cols / vec_len * vec_len;
+  for (; i < vec_len_aligned; i += vec_len) {
+    __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[r * stride + i]);
+    // Convert fp8 to fp32
+    at::vec::CPU_CAPABILITY::cvtfp8e4m3_fp32(_mm256_castsi256_si128(fp8_vec), fp32_vec_0);
+    at::vec::CPU_CAPABILITY::cvtfp8e4m3_fp32(_mm256_extracti128_si256(fp8_vec, 1), fp32_vec_1);
+    // Convert to bf16 and store
+    __m256i bf16_vec_0 = at::vec::cvtfp32_bf16(fp32_vec_0);
+    __m256i bf16_vec_1 = at::vec::cvtfp32_bf16(fp32_vec_1);
+    __m512i bf16_vec = _mm512_inserti32x8(_mm512_castsi256_si512(bf16_vec_0), bf16_vec_1, 1);
+    _mm512_storeu_si512((__m512i*)(out + r * cols + i), bf16_vec);
+  }
+  for (; i < cols; ++i) {
+    out[r * cols + i] = (at::BFloat16)in[r * stride + i];
+  }
   }
 }
 
-
 // accumulate and store result to buffer
 // if act/wei are per_group quantized, apply scales
 template <bool accum, int64_t N, int act_quant_mode, int wei_quant_mode>
@@ -294,7 +266,8 @@ inline void store_out(
       if constexpr (wei_quant_mode == PER_ROW) {
         b_scale = scales_b[j];
       }
-      c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale);
+      float bias_val = bias ? bias[j] : 0.0f;
+      c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale + bias_val);
     }
   } // for M
 }
@@ -341,7 +314,8 @@ inline void store_out(
       if constexpr (wei_quant_mode == PER_ROW) {
         b_scale = scales_b[j];
       }
-      c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale);
+      float bias_val = bias ? bias[j] : 0.0f;
+      c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale + bias_val);
     }
   } // for M
 }
@@ -515,7 +489,7 @@ void _float8_linear_impl(
 
       for (int mci = mc; mci < mc_end; ++mci) {
         int64_t m_size = mci * block_m + block_m > M ? M - mci * block_m : block_m;
-        zero_buffer(y_buf, m_size * block_n);
+        memset(y_buf, 0, sizeof(float) * m_size * block_n);
         for (int kci = 0; kci < Kc; ++kci) {
           auto scales_a = a_scales_ptr + mci * block_m * num_groups + kci / block_per_group;
           auto scales_b = b_scales_ptr + nc * block_n * num_groups + kci / block_per_group * block_n;

torchao/csrc/cpu/aten_kernels/utils.h

@@ -35,24 +35,6 @@ get_m_blocking(int64_t M) {
 // Cached check for AVX-512F support in this process, for use by CPU kernels
 // that include this header and are compiled with CPU_CAPABILITY_AVX512.
 inline const bool kHasAVX512 = __builtin_cpu_supports("avx512f");
-
-template<typename T>
-void zero_buffer(T* data, int64_t size) {
-  const int32_t vec_size = at::vec::Vectorized<T>::size();
-  auto zero_vec = at::vec::Vectorized<T>(0);
-  int64_t d = 0;
-  for (; d < size - (size % vec_size); d += vec_size) {
-    zero_vec.store(data + d);
-  }
-  if (d < size) {
-    zero_vec.store(data + d, size - d);
-  }
-}
-#else
-template<typename T>
-void zero_buffer(T* data, int64_t size) {
-  memset(data, 0, sizeof(T) * size);
-}
 #endif
 
 template <typename T> struct vnni_traits;