[Common] Pipeline the dbias reduction of the MXFP8 and FP8 quantize

tests/cpp/util/CMakeLists.txt

@@ -3,6 +3,7 @@
 # See LICENSE for license information.
 
 add_executable(test_util
+               test_reduce_dbias.cu
                test_nvrtc.cpp
                test_string.cpp
                ../test_common.cu)

tests/cpp/util/test_reduce_dbias.cu

@@ -0,0 +1,126 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+// The dbias reduction of the MXFP8 and FP8 quantize must add each column of the FP32 workspace in
+// row order, starting from 0, and round once to the output type. This checks it bit for bit against
+// that sum on the host, on data whose magnitudes span 2^-30..2^30 with mixed signs, so that any
+// other order or grouping of the adds would change the result.
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <random>
+#include <tuple>
+#include <vector>
+
+#include "cast/core/grouped_layout.cuh"
+
+namespace {
+
+using transformer_engine::DType;
+
+template <typename T>
+T round_to(float x);
+template <>
+float round_to<float>(float x) {
+  return x;
+}
+template <>
+__nv_bfloat16 round_to<__nv_bfloat16>(float x) {
+  return __float2bfloat16_rn(x);
+}
+template <>
+__half round_to<__half>(float x) {
+  return __float2half_rn(x);
+}
+
+template <typename T>
+DType dtype_of();
+template <>
+DType dtype_of<float>() {
+  return DType::kFloat32;
+}
+template <>
+DType dtype_of<__nv_bfloat16>() {
+  return DType::kBFloat16;
+}
+template <>
+DType dtype_of<__half>() {
+  return DType::kFloat16;
+}
+
+template <typename T>
+void check(const size_t rows, const size_t cols, const uint32_t seed) {
+  std::mt19937 gen(seed);
+  std::uniform_real_distribution<float> mantissa(1.0f, 2.0f);
+  std::uniform_int_distribution<int> exponent(-30, 30);
+  std::bernoulli_distribution negative(0.5);
+  std::vector<float> workspace(rows * cols);
+  for (float &w : workspace) {
+    w = std::ldexp(mantissa(gen), exponent(gen)) * (negative(gen) ? -1.0f : 1.0f);
+  }
+  std::vector<T> expected(cols);
+  for (size_t c = 0; c < cols; ++c) {
+    float acc = 0.0f;
+    for (size_t r = 0; r < rows; ++r) acc += workspace[r * cols + c];
+    expected[c] = round_to<T>(acc);
+  }
+
+  float *workspace_dev = nullptr;
+  T *dbias_dev = nullptr;
+  ASSERT_EQ(cudaMalloc(&workspace_dev, workspace.size() * sizeof(float)), cudaSuccess);
+  ASSERT_EQ(cudaMalloc(&dbias_dev, cols * sizeof(T)), cudaSuccess);
+  ASSERT_EQ(cudaMemcpy(workspace_dev, workspace.data(), workspace.size() * sizeof(float),
+                       cudaMemcpyHostToDevice),
+            cudaSuccess);
+  ASSERT_EQ(cudaMemset(dbias_dev, 0xA5, cols * sizeof(T)), cudaSuccess);
+  transformer_engine::Tensor dbias;
+  dbias.data = transformer_engine::SimpleTensor(dbias_dev, {cols}, dtype_of<T>());
+  transformer_engine::dispatch::common::reduce_dbias<T>(workspace_dev, &dbias, rows, cols, 0);
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+  std::vector<T> got(cols);
+  ASSERT_EQ(cudaMemcpy(got.data(), dbias_dev, cols * sizeof(T), cudaMemcpyDeviceToHost),
+            cudaSuccess);
+  cudaFree(workspace_dev);
+  cudaFree(dbias_dev);
+
+  size_t mismatches = 0;
+  for (size_t c = 0; c < cols; ++c) {
+    if (std::memcmp(&got[c], &expected[c], sizeof(T)) != 0 && mismatches++ < 4) {
+      ADD_FAILURE() << "column " << c << " of " << rows << "x" << cols << " differs";
+    }
+  }
+  EXPECT_EQ(mismatches, 0u);
+}
+
+using Params = std::tuple<size_t /*rows*/, size_t /*cols*/>;
+
+class ReduceDbiasTest : public ::testing::TestWithParam<Params> {};
+
+TEST_P(ReduceDbiasTest, AddsEachColumnInRowOrder) {
+  const auto [rows, cols] = GetParam();
+  check<float>(rows, cols, 1);
+  check<__nv_bfloat16>(rows, cols, 2);
+  check<__half>(rows, cols, 3);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    OperatorTest, ReduceDbiasTest,
+    // Rows: around the 32-row load batch and up to the workspace of 128K-row tensors. Columns:
+    // around the 64-thread CTA, including counts that are not multiples of 4.
+    ::testing::Combine(::testing::Values(1, 2, 31, 32, 33, 63, 64, 65, 512, 1000, 2000),
+                       ::testing::Values(1, 3, 4, 28, 60, 63, 64, 65, 68, 1028, 2052, 7168, 7170)),
+    [](const testing::TestParamInfo<Params> &info) {
+      return std::to_string(std::get<0>(info.param)) + "X" +
+             std::to_string(std::get<1>(info.param));
+    });
+
+}  // namespace

transformer_engine/common/cast/core/grouped_layout.cuh

@@ -9,9 +9,9 @@
  *
  *  Maps CTA coordinates to (tensor, tile) work-items for grouped tensors of
  *  varying shapes, plus the dbias reductions. Contains no arch-specific PTX or
- *  TMA machinery, so it can be included by sources that compile for the generic
- *  (non-suffixed) Blackwell architectures. The TMA descriptor management and
- *  bulk-copy staging live in grouped_tma.cuh.
+ *  TMA machinery (cp.async is used only on sm_80 and newer), so it can be included
+ *  by sources that compile for the generic (non-suffixed) Blackwell architectures.
+ *  The TMA descriptor management and bulk-copy staging live in grouped_tma.cuh.
  */
 
 #ifndef TRANSFORMER_ENGINE_QUANTIZE_CORE_GROUPED_LAYOUT_CUH_
@@ -36,39 +36,71 @@ inline bool full_tile_1D_tensor(const Tensor *const t, const size_t elems_per_bl
 namespace kernel {
 
 constexpr size_t THREADS_PER_BLOCK = 256;
-template <int nvec, typename OType>
-__global__ void __launch_bounds__(THREADS_PER_BLOCK)
-    reduce_dbias_kernel(OType *const dbias_output, const float *const dbias_partial,
-                        const size_t rows, const size_t cols) {
-  using ComputeVec = Vec<float, nvec>;
-  using OutputVec = Vec<OType, nvec>;
 
-  const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
+//! Threads per CTA of reduce_dbias_kernel, one column each.
+constexpr int REDUCE_DBIAS_THREADS = 64;
+//! Rows of a column copied into shared memory per pipeline stage.
+constexpr int REDUCE_DBIAS_STAGE_ROWS = 32;
+//! Pipeline stages: while one stage is added up, the others are in flight.
+constexpr int REDUCE_DBIAS_STAGES = 4;
 
-  if (thread_id * nvec >= cols) {
+/*! \brief dbias[c] = sum over rows r of partial[r][c], added in row order from 0.
+ *
+ *  The adds of a column form one sequential chain, so the result does not depend on how the
+ *  work is split; what the split decides is how many loads are in flight. One thread per
+ *  column copies its column into shared memory with cp.async, REDUCE_DBIAS_STAGE_ROWS rows per
+ *  stage and REDUCE_DBIAS_STAGES stages deep, and adds each stage in row order once it has
+ *  landed. A thread reads back only the rows it copied itself, so no barrier is needed. Before
+ *  sm_80, which has no cp.async, each thread loads its column directly, in the same order.
+ */
+template <typename OType>
+__global__ void __launch_bounds__(REDUCE_DBIAS_THREADS)
+    reduce_dbias_kernel(OType *const dbias_output, const float *const dbias_partial,
+                        const size_t rows, const size_t cols) {
+  const size_t col = static_cast<size_t>(blockIdx.x) * REDUCE_DBIAS_THREADS + threadIdx.x;
+  if (col >= cols) {
     return;
   }
-
-  const float *const thread_in_base = dbias_partial + thread_id * nvec;
-  OType *const thread_out_base = dbias_output + thread_id * nvec;
-
-  ComputeVec ldg_vec;
-  ComputeVec acc_vec;
-  acc_vec.clear();
-  for (int i = 0; i < rows; ++i) {
-    ldg_vec.load_from(thread_in_base + i * cols);
-#pragma unroll
-    for (int e = 0; e < nvec; ++e) {
-      acc_vec.data.elt[e] += ldg_vec.data.elt[e];
+  const float *in = dbias_partial + col;
+  float acc = 0.0f;
+#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+  __shared__ float stages[REDUCE_DBIAS_STAGES][REDUCE_DBIAS_STAGE_ROWS][REDUCE_DBIAS_THREADS];
+  const size_t num_stages = (rows + REDUCE_DBIAS_STAGE_ROWS - 1) / REDUCE_DBIAS_STAGE_ROWS;
+
+  // Copies stage k of this thread's column into its slot of stages[k % REDUCE_DBIAS_STAGES], and
+  // commits it as one cp.async group; past the last stage it commits an empty group, so that
+  // every iteration below can wait for the same number of groups.
+  const auto issue = [&](const size_t k) {
+    if (k < num_stages) {
+      const size_t r0 = k * REDUCE_DBIAS_STAGE_ROWS;
+      const int n = static_cast<int>(min(static_cast<size_t>(REDUCE_DBIAS_STAGE_ROWS), rows - r0));
+      float(*dst)[REDUCE_DBIAS_THREADS] = stages[k % REDUCE_DBIAS_STAGES];
+      for (int i = 0; i < n; ++i) {
+        const uint32_t dst_addr =
+            static_cast<uint32_t>(__cvta_generic_to_shared(&dst[i][threadIdx.x]));
+        asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" ::"r"(dst_addr),
+                     "l"(in + (r0 + i) * cols)
+                     : "memory");
+      }
     }
-  }
+    asm volatile("cp.async.commit_group;\n" ::: "memory");
+  };
 
-  OutputVec stg_vec;
 #pragma unroll
-  for (int e = 0; e < nvec; ++e) {
-    stg_vec.data.elt[e] = static_cast<OType>(acc_vec.data.elt[e]);
+  for (int k = 0; k < REDUCE_DBIAS_STAGES - 1; ++k) issue(k);
+  for (size_t k = 0; k < num_stages; ++k) {
+    issue(k + REDUCE_DBIAS_STAGES - 1);
+    // All but the REDUCE_DBIAS_STAGES - 1 most recent groups are complete: stage k has landed.
+    asm volatile("cp.async.wait_group %0;\n" ::"n"(REDUCE_DBIAS_STAGES - 1) : "memory");
+    const size_t r0 = k * REDUCE_DBIAS_STAGE_ROWS;
+    const int n = static_cast<int>(min(static_cast<size_t>(REDUCE_DBIAS_STAGE_ROWS), rows - r0));
+    const float(*src)[REDUCE_DBIAS_THREADS] = stages[k % REDUCE_DBIAS_STAGES];
+    for (int i = 0; i < n; ++i) acc += src[i][threadIdx.x];
   }
-  stg_vec.store_to(thread_out_base);
+#else
+  for (size_t r = 0; r < rows; ++r) acc += in[r * cols];
+#endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+  dbias_output[col] = static_cast<OType>(acc);
 }
 
 template <int nvec, typename OType>
@@ -127,15 +159,10 @@ template <typename IType>
 void reduce_dbias(const float *workspace_ptr, Tensor *dbias, const size_t rows, const size_t cols,
                   cudaStream_t stream) {
   using namespace kernel;
-  constexpr size_t reduce_dbias_store_bytes = 8;  // stg.64
-  constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType);
-
-  NVTE_CHECK(cols % reduce_dbias_nvec == 0, "Unsupported shape.");
-  const size_t reduce_dbias_num_blocks = DIVUP(cols, THREADS_PER_BLOCK * reduce_dbias_nvec);
+  const size_t num_blocks = DIVUP(cols, static_cast<size_t>(REDUCE_DBIAS_THREADS));
 
-  reduce_dbias_kernel<reduce_dbias_nvec, IType>
-      <<<reduce_dbias_num_blocks, THREADS_PER_BLOCK, 0, stream>>>(
-          reinterpret_cast<IType *>(dbias->data.dptr), workspace_ptr, rows, cols);
+  reduce_dbias_kernel<IType><<<num_blocks, REDUCE_DBIAS_THREADS, 0, stream>>>(
+      reinterpret_cast<IType *>(dbias->data.dptr), workspace_ptr, rows, cols);
   NVTE_CHECK_CUDA(cudaGetLastError());
 }