From 96a959d53b78d9d658ae47d674f0f937c16a0430 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 10 Nov 2026 12:17:28 -0800 Subject: [PATCH] [Common] Pipeline the dbias reduction of the MXFP8 and FP8 quantize common::reduce_dbias sums the FP32 dbias workspace column by column for the MXFP8 quantize (CUDA and CuTeDSL backends) and the FP8 tensor-scaling quantize. It ran 256 threads per CTA with four columns each, so a 7168-column workspace got 7 CTAs, and each thread waited on one load per row. Use one thread per column, 64 per CTA. On sm_80 and newer, copy each column into shared memory with cp.async, 32 rows per stage and 4 stages deep, and add each stage once it has landed; below sm_80, load the column directly. Each column is still added in row order from 0 and rounded once, so dbias does not change. Any column count is now accepted. Add test_reduce_dbias, which checks the reduction bit for bit against that sum on the host for FP32, BF16 and FP16 output. Signed-off-by: Teddy Do --- tests/cpp/util/CMakeLists.txt | 1 + tests/cpp/util/test_reduce_dbias.cu | 126 ++++++++++++++++++ .../common/cast/core/grouped_layout.cuh | 99 +++++++++----- 3 files changed, 190 insertions(+), 36 deletions(-) create mode 100644 tests/cpp/util/test_reduce_dbias.cu diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index b80a588fe35..669f3d64f0c 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/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) diff --git a/tests/cpp/util/test_reduce_dbias.cu b/tests/cpp/util/test_reduce_dbias.cu new file mode 100644 index 00000000000..0d48fab96d6 --- /dev/null +++ b/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 +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "cast/core/grouped_layout.cuh" + +namespace { + +using transformer_engine::DType; + +template +T round_to(float x); +template <> +float round_to(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 +DType dtype_of(); +template <> +DType dtype_of() { + return DType::kFloat32; +} +template <> +DType dtype_of<__nv_bfloat16>() { + return DType::kBFloat16; +} +template <> +DType dtype_of<__half>() { + return DType::kFloat16; +} + +template +void check(const size_t rows, const size_t cols, const uint32_t seed) { + std::mt19937 gen(seed); + std::uniform_real_distribution mantissa(1.0f, 2.0f); + std::uniform_int_distribution exponent(-30, 30); + std::bernoulli_distribution negative(0.5); + std::vector workspace(rows * cols); + for (float &w : workspace) { + w = std::ldexp(mantissa(gen), exponent(gen)) * (negative(gen) ? -1.0f : 1.0f); + } + std::vector 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(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()); + transformer_engine::dispatch::common::reduce_dbias(workspace_dev, &dbias, rows, cols, 0); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + std::vector 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; + +class ReduceDbiasTest : public ::testing::TestWithParam {}; + +TEST_P(ReduceDbiasTest, AddsEachColumnInRowOrder) { + const auto [rows, cols] = GetParam(); + check(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 &info) { + return std::to_string(std::get<0>(info.param)) + "X" + + std::to_string(std::get<1>(info.param)); + }); + +} // namespace diff --git a/transformer_engine/common/cast/core/grouped_layout.cuh b/transformer_engine/common/cast/core/grouped_layout.cuh index 8337e448158..68e3eb2fdf1 100644 --- a/transformer_engine/common/cast/core/grouped_layout.cuh +++ b/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 -__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; - using OutputVec = Vec; - 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 +__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(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(min(static_cast(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(__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(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(min(static_cast(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(acc); } template @@ -127,15 +159,10 @@ template 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(REDUCE_DBIAS_THREADS)); - reduce_dbias_kernel - <<>>( - reinterpret_cast(dbias->data.dptr), workspace_ptr, rows, cols); + reduce_dbias_kernel<<>>( + reinterpret_cast(dbias->data.dptr), workspace_ptr, rows, cols); NVTE_CHECK_CUDA(cudaGetLastError()); }