2025-09-27 nightly release (dbc93d4)

Author: pytorchbot

fbgemm_gpu/cmake/Fbgemm.cmake

@@ -26,7 +26,8 @@ set(fbgemm_sources_avx2
   "${FBGEMM}/src/QuantUtilsAvx2.cc")
 
 set(fbgemm_sources_avx512
-  "${FBGEMM}/src/EmbeddingSpMDMAvx512.cc")
+  "${FBGEMM}/src/EmbeddingSpMDMAvx512.cc"
+  "${FBGEMM}/src/QuantUtilsAvx512.cc")
 
 if(CXX_AVX2_FOUND)
   set_source_files_properties(${fbgemm_sources_avx2}
@@ -46,7 +47,7 @@ if(CXX_AVX2_FOUND)
     ${fbgemm_sources}
     ${fbgemm_sources_avx2})
 endif()
-if((NOT FBGEMM_BUILD_VARIANT STREQUAL BUILD_VARIANT_ROCM) AND CXX_AVX512_FOUND)
+if(CXX_AVX512_FOUND)
   set(fbgemm_sources
     ${fbgemm_sources}
     ${fbgemm_sources_avx2}

fbgemm_gpu/codegen/inference/embedding_forward_quantized_cpu_template.cpp

@@ -18,7 +18,6 @@
 #include "fbgemm_gpu/embedding_common.h"
 #include "fbgemm/FbgemmEmbedding.h"
 #include "fbgemm_gpu/utils/tensor_utils.h"
-#include "fbgemm_gpu/config/feature_gates.h"
 
 #if defined(__x86_64__) || defined(__i386__) || (defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)))
 #include <immintrin.h>
@@ -191,9 +190,8 @@ Tensor int_nbit_split_embedding{{ "_nobag" if nobag else "" }}_codegen_forward_{
     {% else %}
     TORCH_CHECK(D > 0);
     {% endif %}
-    const static bool disablePinnedMemory = fbgemm_gpu::config::is_feature_enabled_from_env(fbgemm_gpu::config::FeatureGateName::TBE_CPU_OUTPUT_DISABLE_PINNED_MEMORY);
     bool pinned_memory = false;
-    if (!disablePinnedMemory && at::Context::hasCUDA() && at::getNumGPUs() > 0) {
+    if (at::Context::hasCUDA() && at::getNumGPUs() > 0) {
       pinned_memory = true;
     }
 

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/device/fmha_device_bwd.hpp

@@ -267,14 +267,15 @@ class Sm100FmhaBwd {
     auto [Q_, K, D, D_VO, HB] = args.problem_shape;
     auto [H, B] = product_each(HB);
     D = cutlass::round_up(D, 8);  // Alignment
-    int Q = cutlass::round_up(static_cast<int>(Q_), 8);  // Alignment
+    size_t Q = cutlass::round_up(static_cast<int>(Q_), 8);  // Alignment
     size_t workspace_bytes = 0;
+    size_t accum_size = sizeof(ElementAccumulator);
     // OdO vector
-    workspace_bytes += B*H*Q * sizeof(ElementAccumulator);
+    workspace_bytes += static_cast<size_t>(B)*static_cast<size_t>(H)*Q * accum_size;
     // scaled LSE vector
-    workspace_bytes += B*H*Q * sizeof(ElementAccumulator);
+    workspace_bytes += static_cast<size_t>(B)*static_cast<size_t>(H)*Q * accum_size;
     // FP32 versions of outputs that are churned (start off with Q only)
-    workspace_bytes += B*H*Q*D * sizeof(ElementAccumulator);
+    workspace_bytes += static_cast<size_t>(B)*static_cast<size_t>(H)*Q*static_cast<size_t>(D) * accum_size;
     return workspace_bytes;
   }
 

fbgemm_gpu/include/fbgemm_gpu/config/feature_gates.h

@@ -62,8 +62,7 @@ namespace fbgemm_gpu::config {
   X(TBE_ROCM_INFERENCE_PACKED_BAGS) \
   X(TBE_ROCM_HIP_BACKWARD_KERNEL)   \
   X(BOUNDS_CHECK_INDICES_V2)        \
-  X(TBE_REPORT_INPUT_PARAMS)        \
-  X(TBE_CPU_OUTPUT_DISABLE_PINNED_MEMORY)
+  X(TBE_REPORT_INPUT_PARAMS)
 // X(EXAMPLE_FEATURE_FLAG)
 
 /// @ingroup fbgemm-gpu-config
@@ -92,13 +91,6 @@ bool check_feature_gate_key(const std::string& key);
 /// is enabled.
 bool is_feature_enabled(const FeatureGateName& feature);
 
-/// @ingroup fbgemm-gpu-config
-///
-/// @brief For the given `FeatureGateName`, check if the corresponding inference
-/// feature is enabled in the env vars only. Only applicable for inference
-/// features suitable for env var rollouts
-bool is_feature_enabled_from_env(const FeatureGateName& feature);
-
 #ifdef FBGEMM_FBCODE
 bool is_feature_enabled(const FbFeatureGateName& feature);
 #endif

fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h

@@ -411,6 +411,7 @@ at::Tensor FP8rowwise_to_float_cpu(
     const bool forward = true,
     const int64_t output_dtype = 0);
 at::Tensor fused8bitrowwise_to_half_cpu(const at::Tensor& input);
+at::Tensor fused8bitrowwise_to_bfloat16_cpu(const at::Tensor& input);
 at::Tensor fused8bitrowwise_to_float_or_half_cpu(
     const at::Tensor& input,
     const int64_t output_dtype,
@@ -469,6 +470,9 @@ at::Tensor _fusednbitrowwise_to_float_or_half_gpu(
 at::Tensor& _fused8bitrowwise_to_float_cpu_out(
     at::Tensor& output,
     const at::Tensor& input);
+at::Tensor& _fused8bitrowwise_to_bfloat16_cpu_out(
+    at::Tensor& output,
+    const at::Tensor& input);
 at::Tensor& _float_to_fused8bitrowwise_cpu_out(
     at::Tensor& output,
     const at::Tensor& input);

fbgemm_gpu/src/config/feature_gates.cpp

@@ -45,50 +45,39 @@ bool ev_check_key(const std::string& key) {
   }
 }
 
-static bool check_feature_gate_key_impl(
-    const std::string& key,
-    bool check_env_vars_only) {
+DLL_PUBLIC bool check_feature_gate_key(const std::string& key) {
   // Cache feature flags to avoid repeated JK and env var checks
   static std::map<std::string, bool> feature_flags_cache;
-  if (const auto search = feature_flags_cache.find(key);
-      search != feature_flags_cache.end()) {
-    return search->second;
-  }
 #ifdef FBGEMM_FBCODE
-  const auto value =
-      check_env_vars_only ? ev_check_key(key) : jk_check_key(key);
-#else
-  const auto value = ev_check_key(key);
+  static const auto no_jk = ev_check_key("NO_JK");
 #endif
 
-  feature_flags_cache.insert({key, value});
-  return value;
-}
+  if (const auto search = feature_flags_cache.find(key);
+      search != feature_flags_cache.end()) {
+    return search->second;
 
-DLL_PUBLIC bool check_feature_gate_key(const std::string& key) {
+  } else {
+    const auto value =
 #ifdef FBGEMM_FBCODE
-  static const auto no_jk = ev_check_key("NO_JK");
+        (no_jk) ? ev_check_key(key) : jk_check_key(key);
 #else
-  static const auto no_jk = false;
+        ev_check_key(key);
 #endif
 
-  return check_feature_gate_key_impl(key, no_jk);
+    feature_flags_cache.insert({key, value});
+    return value;
+  }
 }
 
 DLL_PUBLIC bool is_feature_enabled(const FeatureGateName& feature) {
   return check_feature_gate_key(to_string(feature));
 }
 
-DLL_PUBLIC bool is_feature_enabled_from_env(const FeatureGateName& feature) {
-  return check_feature_gate_key_impl(
-      to_string(feature), /* check_env_vars_only */ true);
-}
-
 #ifdef FBGEMM_FBCODE
 DLL_PUBLIC bool is_feature_enabled(const FbFeatureGateName& feature) {
   return check_feature_gate_key(to_string(feature));
 }
-#endif // FBGEMM_FBCODE
+#endif
 
 } // namespace fbgemm_gpu::config
 

fbgemm_gpu/src/quantize_ops/quantize_ops_cpu.cpp

@@ -55,7 +55,7 @@ Tensor& _float_to_fused8bitrowwise_cpu_out_t(
   return output;
 }
 
-template <typename output_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) {
@@ -78,7 +78,9 @@ Tensor& _fused8bitrowwise_to_float_cpu_out_t(
   auto output_data = static_cast<output_t*>(
       output.data_ptr()); // output.data_ptr<output_t>(); -> Yields
                           // unresolved data_ptr symbol.
-  fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<output_t>(
+  fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<
+      output_t,
+      is_uint16_t_of_type_bf16>(
       input.data_ptr<uint8_t>(), nrows, ncols, output_data);
 
   return output;
@@ -217,11 +219,19 @@ 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>(output, input);
+  return _fused8bitrowwise_to_float_cpu_out_t<float, false>(output, input);
 }
 
 Tensor& fused8bitrowwise_to_half_cpu_out(Tensor& output, const Tensor& input) {
-  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::float16>(output, input);
+  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::float16, false>(
+      output, input);
+}
+
+Tensor& _fused8bitrowwise_to_bfloat16_cpu_out(
+    Tensor& output,
+    const Tensor& input) {
+  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::bfloat16, true>(
+      output, input);
 }
 
 /// @ingroup quantize-data-cpu
@@ -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);

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

fbgemm_gpu/test/tbe/inference/nbit_forward_test.py

@@ -8,7 +8,6 @@
 # pyre-strict
 # pyre-ignore-all-errors[56]
 
-import os
 import random
 import unittest
 from typing import Any, Callable, Optional, Union
@@ -124,12 +123,6 @@
 
 @optests.generate_opcheck_tests(fast=True, additional_decorators=additional_decorators)
 class NBitFowardTest(NBitFowardTestCommon):
-    def _is_cpu_output_on_pinned_memory(self) -> bool:
-        return (
-            os.getenv("FBGEMM_TBE_CPU_OUTPUT_DISABLE_PINNED_MEMORY") != "1"
-            and torch.cuda.is_available()
-        )
-
     def execute_nbit_forward_fused_pooled_emb_quant_(
         self,
         T: int,
@@ -905,9 +898,6 @@ def test_nbit_forward_cpu_seq_int8(
         lengths = torch.cat(lengths_list, 0)
         offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths)
         quant_cc_output = quant_cc(indices.int(), offsets.int())
-        self.assertEqual(
-            quant_cc_output.is_pinned(), self._is_cpu_output_on_pinned_memory()
-        )
         tables_rows = [
             T for T, _, _ in quant_cc.split_embedding_weights_with_scale_bias(0)
         ]

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
@@ -330,7 +331,7 @@ FBGEMM_API void FloatOrHalfToFused8BitRowwiseQuantizedSBFloat(
  * This version intentionally supports only 8-bit because
  * the corresponding quantize version only supports 8-bit.
  */
-template <typename OutputType>
+template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
 FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
     const uint8_t* input,
     size_t input_rows,
@@ -377,7 +378,7 @@ FBGEMM_API void FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef(
  * Same as Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf but unoptimized.
  * This should not be called directly except in testing.
  */
-template <typename OutputType>
+template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
 FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef(
     const uint8_t* input,
     size_t input_rows,