support scale_bias_last and quant_padding_float_type for cpu dequant kernel (#4943)

seanx92 · meta-codesync[bot] · commit eee643b94fea · 2025-10-02T15:34:39.000-07:00
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/1963 Pull Request resolved: #4943 scale_bias_last: decides whether scale/bias padding is at the front or end of the row quant_padding_float_type: decides if scale/bias is represented by fp32 or fp16 this is to match the cuda kernel implementation functionalities and allow cpu dequantization with front padded FP16 scale/bias Reviewed By: q10 Differential Revision: D83405212 fbshipit-source-id: 34628568cb26dc66de24a9f02e9fb1161f20ace9
diff --git a/fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h b/fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h
@@ -469,10 +469,19 @@ at::Tensor _fusednbitrowwise_to_float_or_half_gpu(
     const int64_t output_dtype);
 at::Tensor& _fused8bitrowwise_to_float_cpu_out(
     at::Tensor& output,
-    const at::Tensor& input);
+    const at::Tensor& input,
+    const bool scale_bias_last = true,
+    const bool quant_padding_float_type = true);
+at::Tensor& fused8bitrowwise_to_half_cpu_out(
+    at::Tensor& output,
+    const at::Tensor& input,
+    const bool scale_bias_last = true,
+    const bool quant_padding_float_type = true);
 at::Tensor& _fused8bitrowwise_to_bfloat16_cpu_out(
     at::Tensor& output,
-    const at::Tensor& input);
+    const at::Tensor& input,
+    const bool scale_bias_last = true,
+    const bool quant_padding_float_type = true);
 at::Tensor& _float_to_fused8bitrowwise_cpu_out(
     at::Tensor& output,
     const at::Tensor& input);
diff --git a/fbgemm_gpu/src/quantize_ops/quantize_ops_cpu.cpp b/fbgemm_gpu/src/quantize_ops/quantize_ops_cpu.cpp
@@ -58,18 +58,26 @@ Tensor& _float_to_fused8bitrowwise_cpu_out_t(
 template <typename output_t, bool is_uint16_t_of_type_bf16 = false>
 Tensor& _fused8bitrowwise_to_float_cpu_out_t(
     Tensor& output,
-    const Tensor& input) {
+    const Tensor& input,
+    const bool scale_bias_last,
+    const bool quant_padding_float_type) {
   TENSOR_ON_CPU(input);
   TORCH_CHECK(
       input.dim() >= 2,
       "Tensor 'input' must have >= 2 dimension(s). Found ",
       input.ndimension());
+  TORCH_CHECK(
+      quant_padding_float_type == true || scale_bias_last == false,
+      "2-byte padding (quant_padding_float_type=false) only works with scale_bias_last=false")
+
+  const int quant_padding_size =
+      (quant_padding_float_type) ? sizeof(float) : sizeof(fbgemm::float16);
 
   const auto input_sizes = input.sizes();
   const auto last_dim = input_sizes.size() - 1;
   const int64_t nrows = c10::size_to_dim_(last_dim, input_sizes);
   const int32_t ncols = input_sizes[last_dim];
-  const int32_t output_columns = ncols - 2 * sizeof(float);
+  const int32_t output_columns = ncols - 2 * quant_padding_size;
 
   auto output_dims = input_sizes.vec();
   output_dims[last_dim] = output_columns;
@@ -81,7 +89,12 @@ Tensor& _fused8bitrowwise_to_float_cpu_out_t(
   fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<
       output_t,
       is_uint16_t_of_type_bf16>(
-      input.data_ptr<uint8_t>(), nrows, ncols, output_data);
+      input.data_ptr<uint8_t>(),
+      nrows,
+      ncols,
+      output_data,
+      scale_bias_last,
+      quant_padding_float_type);
 
   return output;
 }
@@ -218,20 +231,29 @@ Tensor _fusednbitrowwise_sbfront_to_float_or_half_cpu(
 ///
 Tensor& _fused8bitrowwise_to_float_cpu_out(
     Tensor& output,
-    const Tensor& input) {
-  return _fused8bitrowwise_to_float_cpu_out_t<float, false>(output, input);
+    const Tensor& input,
+    const bool scale_bias_last,
+    const bool quant_padding_float_type) {
+  return _fused8bitrowwise_to_float_cpu_out_t<float, false>(
+      output, input, scale_bias_last, quant_padding_float_type);
 }
 
-Tensor& fused8bitrowwise_to_half_cpu_out(Tensor& output, const Tensor& input) {
+Tensor& fused8bitrowwise_to_half_cpu_out(
+    Tensor& output,
+    const Tensor& input,
+    const bool scale_bias_last,
+    const bool quant_padding_float_type) {
   return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::float16, false>(
-      output, input);
+      output, input, scale_bias_last, quant_padding_float_type);
 }
 
 Tensor& _fused8bitrowwise_to_bfloat16_cpu_out(
     Tensor& output,
-    const Tensor& input) {
+    const Tensor& input,
+    const bool scale_bias_last,
+    const bool quant_padding_float_type) {
   return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::bfloat16, true>(
-      output, input);
+      output, input, scale_bias_last, quant_padding_float_type);
 }
 
 /// @ingroup quantize-data-cpu
@@ -307,24 +329,27 @@ Tensor fused8bitrowwise_to_bfloat16_cpu(const Tensor& input) {
 Tensor fused8bitrowwise_to_float_or_half_cpu(
     const Tensor& input,
     const int64_t output_dtype,
-    [[maybe_unused]] const bool scale_bias_last,
-    [[maybe_unused]] const bool quant_padding_float_type) {
+    const bool scale_bias_last,
+    const bool quant_padding_float_type) {
   Tensor output;
   SparseType output_sparse_dtype = static_cast<SparseType>(output_dtype);
   switch (output_sparse_dtype) {
     case SparseType::FP32:
       output = at::empty({0}, input.options().dtype(at::kFloat));
 
-      output = _fused8bitrowwise_to_float_cpu_out(output, input);
+      output = _fused8bitrowwise_to_float_cpu_out(
+          output, input, scale_bias_last, quant_padding_float_type);
 
       break;
     case SparseType::FP16:
       output = at::empty({0}, input.options().dtype(at::kHalf));
-      output = fused8bitrowwise_to_half_cpu_out(output, input);
+      output = fused8bitrowwise_to_half_cpu_out(
+          output, input, scale_bias_last, quant_padding_float_type);
       break;
     case SparseType::BF16:
       output = at::empty({0}, input.options().dtype(at::kBFloat16));
-      output = _fused8bitrowwise_to_bfloat16_cpu_out(output, input);
+      output = _fused8bitrowwise_to_bfloat16_cpu_out(
+          output, input, scale_bias_last, quant_padding_float_type);
       break;
     default:
       TORCH_CHECK(false);
@@ -607,7 +632,7 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
   m.def(
       "Fused8BitRowwiseQuantizedToFloatOrHalf(Tensor input, int output_dtype=0, bool scale_bias_last=True, bool quant_padding_float_type=True) -> Tensor");
   m.def(
-      "Fused8BitRowwiseQuantizedToFloatOut(Tensor output, Tensor input) -> Tensor");
+      "Fused8BitRowwiseQuantizedToFloatOut(Tensor output, Tensor input, bool scale_bias_last=True, bool quant_padding_float_type=True) -> Tensor");
   m.def(
       "Fused8BitRowwiseQuantizedToFloatMixedDim(Tensor input, Tensor D_offsets, int output_dtype) -> Tensor");
   m.def(
diff --git a/fbgemm_gpu/test/quantize/fused_8bit_rowwise_test.py b/fbgemm_gpu/test/quantize/fused_8bit_rowwise_test.py
@@ -144,14 +144,37 @@ def quantize_and_dequantize_op_test_helper(  # noqa: C901
             # cpu path only supports bf16 dequantization
             if output_dtype == SparseType.BF16:
                 input_data = input_data.float()
+            if not test_generic_op and not quant_padding_float_type:
+                return
+            if not quant_padding_float_type and output_dtype == SparseType.FP32:
+                return
             if test_generic_op:
-                quantized_data = (
+                quantized_data_ref = (
                     torch.ops.fbgemm.FloatOrHalfToFused8BitRowwiseQuantized(input_data)
                 )
+                # fbgemm weight 2byte storages are scale_bias first layout
+                if quant_padding_float_type is False:
+                    scale_bias_last = False
+                    quant_pad = quantized_data_ref[:, -8:]
+                    quant_data = quantized_data_ref[:, :-8]
+                    quantized_data = torch.cat(
+                        [
+                            quant_pad.view(torch.float)
+                            .to(torch.half)
+                            .view(torch.uint8),
+                            quant_data,
+                        ],
+                        dim=1,
+                    )
+                else:
+                    scale_bias_last = True
+                    quantized_data = quantized_data_ref
                 dequantized_data = (
                     torch.ops.fbgemm.Fused8BitRowwiseQuantizedToFloatOrHalf(
                         quantized_data,
                         output_dtype.as_int(),
+                        quant_padding_float_type=quant_padding_float_type,
+                        scale_bias_last=scale_bias_last,
                     )
                 )
             else:
@@ -187,9 +210,17 @@ def quantize_and_dequantize_op_test_helper(  # noqa: C901
                 assert dequantized_data.numel() == 0
                 return
 
-            reference = torch.from_numpy(
-                fused_rowwise_8bit_dequantize_reference(quantized_data.numpy())
-            )
+            quantize_data_numpy = quantized_data.numpy()
+            if quant_padding_float_type:
+                reference = torch.from_numpy(
+                    fused_rowwise_8bit_dequantize_reference(quantize_data_numpy)
+                )
+            else:
+                reference = torch.from_numpy(
+                    fused_rowwise_8bit_dequantize_2bytes_padding_scale_bias_first_reference(
+                        quantize_data_numpy
+                    )
+                )
             if output_dtype == SparseType.FP32:
                 torch.testing.assert_close(dequantized_data.float(), reference.float())
             elif output_dtype == SparseType.FP16:
diff --git a/include/fbgemm/QuantUtils.h b/include/fbgemm/QuantUtils.h
@@ -336,7 +336,9 @@ FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
     const uint8_t* input,
     size_t input_rows,
     int input_columns,
-    OutputType* output);
+    OutputType* output,
+    const bool scale_bias_last = true,
+    const bool quant_padding_float_type = true);
 
 /**
  * Same as ToFusedNBitRowwiseQuantizedSBHalf but unoptimized.
@@ -383,6 +385,8 @@ FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef(
     const uint8_t* input,
     size_t input_rows,
     int input_columns,
-    OutputType* output);
+    OutputType* output,
+    const bool scale_bias_last = true,
+    const bool quant_padding_float_type = true);
 
 } // namespace fbgemm
diff --git a/include/fbgemm/QuantUtilsAvx2.h b/include/fbgemm/QuantUtilsAvx2.h
@@ -166,7 +166,10 @@ void FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfAvx2(
     int input_columns,
     OutputType* output);
 
-template <typename OutputType>
+template <
+    typename OutputType,
+    bool scale_bias_last = true,
+    bool quant_padding_float_type = true>
 void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfAvx2(
     const std::uint8_t* input,
     size_t input_rows,
diff --git a/include/fbgemm/QuantUtilsAvx512.h b/include/fbgemm/QuantUtilsAvx512.h
@@ -39,6 +39,7 @@ FBGEMM_API void requantizeOutputProcessingGConvAvx512(
     int ld_in,
     const requantizationParams_t<BIAS_TYPE>& r);
 
+template <bool scale_bias_last = true, bool quant_padding_float_type = true>
 void Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512(
     const std::uint8_t* input,
     size_t input_rows,
diff --git a/src/QuantUtils.cc b/src/QuantUtils.cc
diff --git a/src/QuantUtilsAvx2.cc b/src/QuantUtilsAvx2.cc
diff --git a/src/QuantUtilsAvx512.cc b/src/QuantUtilsAvx512.cc
