From 94265c0edad10dcea19deadc3340d6ff1f6e1869 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 13 Oct 2026 13:56:07 -0700 Subject: [PATCH 01/29] Read BF16 GeLU, SiLU and their derivatives from lookup tables in MXFP8 quantize A BF16 input takes only 2^16 values, so GeLU, SiLU, dGeLU and dSiLU of BF16 inputs can be read from small tables instead of computed. The tables are generated from the scalar util/math.h functions by nvte_build_activation_table. The framework builds each table once per device and passes it to the new _v2 activation entry points through kNVTEQuantizationConfigActivationTable. The generic MXFP8 quantize kernel reads it for BF16 inputs; other input types, recipes and activations are unchanged, and so is every call without a table. Lookups are branch-free. Forward tables hold BF16 values with closed forms outside the window (RN(x/2) below, x or -0 above). Derivative tables hold FP32 values; below the window the derivative is its linear Taylor term rounded once. Against util/math.h, table results differ only in GeLU/SiLU of -inf (-0 instead of NaN) and in the derivative below 2^-12 (dGeLU, within 1 ulp) or 2^-9 (dSiLU, within 2 ulp). PyTorch allocates the tables through its caching allocator; JAX allocates them with cudaMalloc in the FFI initialize stage, outside graph capture. Co-authored-by: Oleg Goncharov Signed-off-by: Teddy Do --- tests/cpp/operator/CMakeLists.txt | 1 + .../cpp/operator/test_cast_mxfp8_act_exact.cu | 470 ++++++++++++++++++ tests/cpp/test_common.h | 25 + tests/cpp/util/CMakeLists.txt | 1 + tests/cpp/util/test_activation_table.cu | 196 ++++++++ transformer_engine/common/CMakeLists.txt | 1 + .../common/activation/activation_table.cu | 102 ++++ .../common/activation/activation_template.h | 11 +- transformer_engine/common/activation/gelu.cu | 14 + .../common/activation/gelu_dbias.cu | 13 + .../common/activation/swiglu.cu | 14 + .../common/activation/swiglu_dbias.cu | 13 + .../common/cast/dispatch/quantize.cuh | 50 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 115 ++--- transformer_engine/common/common.h | 4 +- .../include/transformer_engine/activation.h | 99 ++++ .../common/include/transformer_engine/cast.h | 34 ++ .../transformer_engine/transformer_engine.h | 11 + .../common/transformer_engine.cpp | 6 + .../common/util/activation_table.cuh | 227 +++++++++ transformer_engine/jax/csrc/extensions.h | 5 + .../jax/csrc/extensions/activation.cpp | 49 +- .../jax/csrc/extensions/activation_table.cpp | 78 +++ transformer_engine/pytorch/csrc/common.h | 52 ++ .../pytorch/csrc/extensions/activation.cpp | 12 +- .../csrc/extensions/activation_table.cpp | 57 +++ .../pytorch/csrc/extensions/bias.cpp | 10 +- 27 files changed, 1578 insertions(+), 92 deletions(-) create mode 100644 tests/cpp/operator/test_cast_mxfp8_act_exact.cu create mode 100644 tests/cpp/util/test_activation_table.cu create mode 100644 transformer_engine/common/activation/activation_table.cu create mode 100644 transformer_engine/common/util/activation_table.cuh create mode 100644 transformer_engine/jax/csrc/extensions/activation_table.cpp create mode 100644 transformer_engine/pytorch/csrc/extensions/activation_table.cpp diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index aa3ec536712..e05919c2b53 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -12,6 +12,7 @@ add_executable(test_operator test_cast_mxfp8_gated_swiglu.cu test_qdq.cu test_cast_mxfp8.cu + test_cast_mxfp8_act_exact.cu test_cast_mxfp8_grouped.cu test_cast_mxfp8_grouped_scaled_swiglu.cu test_cast_nvfp4_transpose.cu diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu new file mode 100644 index 00000000000..dd79cf876e0 --- /dev/null +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -0,0 +1,470 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// The fused MXFP8 activation quantize must round exactly as if the activation were computed per +// element, truncated to the input type, and then quantized on its own. This compares the fused +// kernels bit for bit against that composition: the activation is computed on the GPU, and the +// result is quantized by the generic cast-only kernel (a noop tensor keeps the call off the +// specialized cast kernels). dbias is compared against a host sum in the order the kernel +// accumulates it. +// +// Without an activation table the activation is the scalar util/math.h function. With one, BF16 +// inputs follow the table specification (see NVTEActivationTableKind): -inf gives -0 in the +// forward functions, and the derivatives below the table window are the linear Taylor term. +// Other input types ignore the table. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../test_common.h" +#include "transformer_engine/transformer_engine.h" +#include "util/math.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +enum class Mode { GeLU, DGeLU, DBiasDGeLU, SiLU, DSiLU, DBiasDSiLU }; + +bool is_dbias(const Mode m) { return m == Mode::DBiasDGeLU || m == Mode::DBiasDSiLU; } + +const char *mode_name(const Mode m) { + switch (m) { + case Mode::GeLU: + return "GeLU"; + case Mode::DGeLU: + return "DGeLU"; + case Mode::DBiasDGeLU: + return "DBiasDGeLU"; + case Mode::SiLU: + return "SiLU"; + case Mode::DSiLU: + return "DSiLU"; + default: + return "DBiasDSiLU"; + } +} + +// act: y = OP(x); dact: y = grad * dOP(x), as the fused kernels form it in FP32. z keeps the FP32 +// value (dbias sums it before truncation), y the value truncated to the input type. +template +__global__ void scalar_activation_kernel(const Mode mode, const bool table_spec, const IType *x, + const IType *grad, IType *y, float *z, const size_t n) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + const float xv = static_cast(x[i]); + const Empty e{}; + // With a table, -inf gives -0 in the forward functions, and the derivatives below + // 2^-12 (dGeLU) and 2^-9 (dSiLU) are their linear Taylor terms. + const bool negative_inf = table_spec && xv == -INFINITY; + const float magnitude = fabsf(xv); + float v; + switch (mode) { + case Mode::GeLU: + v = negative_inf ? -0.0f : gelu(xv, e); + break; + case Mode::SiLU: + v = negative_inf ? -0.0f : silu(xv, e); + break; + case Mode::DGeLU: + case Mode::DBiasDGeLU: + v = static_cast(grad[i]); + v *= table_spec && magnitude < 0x1p-12f ? fmaf(0.79788456f, xv, 0.5f) + : dgelu(xv, e); + break; + default: + v = static_cast(grad[i]); + v *= table_spec && magnitude < 0x1p-9f ? fmaf(0.5f, xv, 0.5f) : dsilu(xv, e); + break; + } + z[i] = v; + y[i] = static_cast(v); +} + +enum class Fill { Typical, OrderedInputs, Specials, TinyGradients }; + +template +T from_bits(const uint32_t bits) { + if constexpr (std::is_same_v) { + float value; + std::memcpy(&value, &bits, sizeof(value)); + return value; + } else { + const uint16_t half_bits = static_cast(bits); + T value; + std::memcpy(&value, &half_bits, sizeof(value)); + return value; + } +} + +// Enumerate the entire 16-bit format in numerical order: -Inf, negative finite values, -0, +// +0, positive finite values, +Inf, then both signs of NaN (which have no numerical order). +// For FP32, enumerate the BF16 exponent/mantissa prefixes and vary the low mantissa bits. +template +uint32_t ordered_bits(const size_t i, const size_t n) { + // The high 16 bits select a format encoding; the low 16 sample FP32 mantissas. + const uint32_t position = static_cast((static_cast(i) << 32) / n); + const uint32_t rank = position >> 16; + const uint32_t pos_inf = std::is_same_v ? 0x7c00u : 0x7f80u; + const uint32_t non_nan_count = 2 * (pos_inf + 1); + uint32_t high; + if (rank <= pos_inf) { + high = 0x8000u | (pos_inf - rank); + } else if (rank < non_nan_count) { + high = rank - pos_inf - 1; + } else { + const uint32_t nan_rank = rank - non_nan_count; + const uint32_t nan_count_per_sign = 0x7fffu - pos_inf; + high = nan_rank < nan_count_per_sign ? 0x8000u | (pos_inf + 1 + nan_rank) + : pos_inf + 1 + nan_rank - nan_count_per_sign; + } + if constexpr (std::is_same_v) { + const uint32_t fraction = position & 0xffffu; + const uint32_t low = (high & 0x8000u) ? 0xffffu - fraction : fraction; + return (high << 16) | ((high & 0x7fffu) == pos_inf ? 0u : low); + } else { + return high; + } +} + +// Exactly representable edge cases in each input format, ordered except for the final NaNs. +template +std::array special_values() { + const float inf = Numeric_Traits::artifInf; + const float nan = std::numeric_limits::quiet_NaN(); + const float max_normal = Numeric_Traits::maxNorm; + const float min_normal = Numeric_Traits::minNorm; + const float min_subnormal = Numeric_Traits::minSubnorm; + return {static_cast(-inf), static_cast(-max_normal), + static_cast(-1.0f), static_cast(-min_normal), + static_cast(-min_subnormal), static_cast(-0.0f), + static_cast(0.0f), static_cast(min_subnormal), + static_cast(min_normal), static_cast(1.0f), + static_cast(max_normal), static_cast(inf), + static_cast(-nan), static_cast(nan)}; +} + +// Typical activations range over [-16, 16] and gradients over [-4, 4]. The ordered sweep +// covers every BF16/FP16 bit pattern, including subnormals, infinities, and NaNs. The special +// case set guarantees exact FP32 boundary values that a sampled mantissa sweep may miss. +template +void fill_inputs(const Fill fill, Tensor *x, Tensor *grad) { + const size_t n = product(x->rowwise_shape()); + T *xp = x->rowwise_cpu_dptr(); + T *gp = grad->rowwise_cpu_dptr(); + const auto specials = special_values(); + for (size_t i = 0; i < n; ++i) { + switch (fill) { + case Fill::Typical: + xp[i] = static_cast(-16.0f + 32.0f * static_cast(i) / (n - 1)); + gp[i] = static_cast(-4.0f + 8.0f * static_cast(i % 257) / 256.0f); + break; + case Fill::OrderedInputs: + xp[i] = from_bits(ordered_bits(i, n)); + gp[i] = static_cast(-4.0f + 8.0f * static_cast(i % 257) / 256.0f); + break; + case Fill::TinyGradients: { + // At x=0 both derivatives are 0.5. Cover subnormal gradients and a normal + // gradient whose product is subnormal, including both signs. + const float tiny_gradients[] = {static_cast(2 * Numeric_Traits::minSubnorm), + static_cast(Numeric_Traits::maxSubnorm), + static_cast(Numeric_Traits::minNorm)}; + xp[i] = static_cast(0.0f); + const float magnitude = tiny_gradients[(i / 2) % 3]; + gp[i] = static_cast(i % 2 == 0 ? magnitude : -magnitude); + break; + } + case Fill::Specials: { + const size_t index = i * specials.size() / n; + xp[i] = specials[index]; + gp[i] = specials[(index + 4) % specials.size()]; + break; + } + } + } + x->from_cpu(); + grad->from_cpu(); +} + +// dbias in the order the quantize kernel accumulates it: per 64-row tile, then over the tiles. +// With FP32 input, rowwise-only scaling and a fused dact, each of the tile's 32 threads first +// adds its two rows (r and r + 32), and the 32 thread sums are then added in order. +std::vector reference_dbias(const std::vector &z, const size_t rows, + const size_t cols, const bool rowwise_reduction) { + constexpr size_t kTileRows = 64; + constexpr size_t kThreadsY = 32; + const size_t tiles = (rows + kTileRows - 1) / kTileRows; + auto at = [&](const size_t r, const size_t c) { return r < rows ? z[r * cols + c] : 0.0f; }; + std::vector dbias(cols); + for (size_t c = 0; c < cols; ++c) { + float total = 0.0f; + for (size_t t = 0; t < tiles; ++t) { + const size_t r0 = t * kTileRows; + float partial = 0.0f; + if (rowwise_reduction) { + for (size_t i = 0; i < kThreadsY; ++i) { + float thread_sum = 0.0f; + thread_sum += at(r0 + i, c); + thread_sum += at(r0 + kThreadsY + i, c); + partial += thread_sum; + } + } else { + for (size_t r = r0; r < r0 + kTileRows; ++r) { + partial += at(r, c); + } + } + total += partial; + } + dbias[c] = total; + } + return dbias; +} + +NVTEActivationTableKind table_kind(const Mode m) { + switch (m) { + case Mode::GeLU: + return kNVTEActivationTableGelu; + case Mode::SiLU: + return kNVTEActivationTableSilu; + case Mode::DGeLU: + case Mode::DBiasDGeLU: + return kNVTEActivationTableDgelu; + default: + return kNVTEActivationTableDsilu; + } +} + +// Each table is built once for the whole test program, as a framework would per device. +NVTETensor activation_table(const NVTEActivationTableKind kind) { + static std::map> tables; + auto &table = tables[kind]; + if (!table) { + const size_t bytes = nvte_activation_table_size(kind); + table = std::make_unique("table", std::vector{bytes}, DType::kByte); + nvte_build_activation_table(kind, table->data(), 0); + } + return table->data(); +} + +size_t count_byte_mismatches(const uint8_t *a, const uint8_t *b, const size_t n) { + size_t m = 0; + for (size_t i = 0; i < n; ++i) m += (a[i] != b[i]); + return m; +} + +template +void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t cols, + const bool rowwise, const bool colwise, const bool swizzled, const bool use_table) { + const DType itype = TypeInfo::dtype; + const DType otype = DType::kFloat8E4M3; + const std::vector shape{rows, cols}; + const size_t n = rows * cols; + + Tensor x("x", shape, itype); + Tensor grad("grad", shape, itype); + fill_inputs(fill, &x, &grad); + + // Reference: scalar activation on the GPU, then the generic cast-only quantize. + Tensor y("y", shape, itype); + float *z_dev = nullptr; + ASSERT_EQ(cudaMalloc(&z_dev, n * sizeof(float)), cudaSuccess); + constexpr int block_size = 256; + const size_t grid_size = divide_round_up(n, block_size); + const bool table_spec = use_table && std::is_same_v; + scalar_activation_kernel + <<>>(mode, table_spec, static_cast(x.rowwise_dptr()), + static_cast(grad.rowwise_dptr()), + static_cast(y.rowwise_dptr()), z_dev, n); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::vector z(n); + ASSERT_EQ(cudaMemcpy(z.data(), z_dev, n * sizeof(float), cudaMemcpyDeviceToHost), cudaSuccess); + cudaFree(z_dev); + + Tensor ref("ref", shape, otype, rowwise, colwise, NVTE_MXFP8_1D_SCALING); + ref.set_with_gemm_swizzled_scales(swizzled); + Tensor noop("noop", std::vector{1}, DType::kFloat32); + nvte_quantize_noop(y.data(), ref.data(), noop.data(), 0); + + // Fused. + Tensor out("out", shape, otype, rowwise, colwise, NVTE_MXFP8_1D_SCALING); + out.set_with_gemm_swizzled_scales(swizzled); + Tensor dbias("dbias", std::vector{cols}, itype); + Tensor workspace; + if (use_table) { + QuantizationConfigWrapper config; + config.set_activation_table(activation_table(table_kind(mode))); + switch (mode) { + case Mode::GeLU: + nvte_gelu_v2(x.data(), out.data(), config, 0); + break; + case Mode::SiLU: + nvte_silu_v2(x.data(), out.data(), config, 0); + break; + case Mode::DGeLU: + nvte_dgelu_v2(grad.data(), x.data(), out.data(), config, 0); + break; + case Mode::DSiLU: + nvte_dsilu_v2(grad.data(), x.data(), out.data(), config, 0); + break; + default: { + auto fn = mode == Mode::DBiasDGeLU ? &nvte_quantize_dbias_dgelu_v2 + : &nvte_quantize_dbias_dsilu_v2; + fn(grad.data(), x.data(), out.data(), dbias.data(), workspace.data(), config, 0); + workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype()); + fn(grad.data(), x.data(), out.data(), dbias.data(), workspace.data(), config, 0); + } + } + } else { + switch (mode) { + case Mode::GeLU: + nvte_gelu(x.data(), out.data(), 0); + break; + case Mode::SiLU: + nvte_silu(x.data(), out.data(), 0); + break; + case Mode::DGeLU: + nvte_dgelu(grad.data(), x.data(), out.data(), 0); + break; + case Mode::DSiLU: + nvte_dsilu(grad.data(), x.data(), out.data(), 0); + break; + default: { + auto fn = + mode == Mode::DBiasDGeLU ? &nvte_quantize_dbias_dgelu : &nvte_quantize_dbias_dsilu; + fn(grad.data(), x.data(), out.data(), dbias.data(), workspace.data(), 0); + workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype()); + fn(grad.data(), x.data(), out.data(), dbias.data(), workspace.data(), 0); + } + } + } + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + out.to_cpu(); + ref.to_cpu(); + if (rowwise) { + EXPECT_EQ(count_byte_mismatches( + reinterpret_cast(out.rowwise_cpu_dptr()), + reinterpret_cast(ref.rowwise_cpu_dptr()), n), + 0u) + << "rowwise data"; + const size_t s = product(out.rowwise_scale_inv_shape()); + EXPECT_EQ(count_byte_mismatches(out.rowwise_cpu_scale_inv_ptr(), + ref.rowwise_cpu_scale_inv_ptr(), s), + 0u) + << "rowwise scales"; + } + if (colwise) { + EXPECT_EQ(count_byte_mismatches( + reinterpret_cast(out.columnwise_cpu_dptr()), + reinterpret_cast(ref.columnwise_cpu_dptr()), n), + 0u) + << "colwise data"; + const size_t s = product(out.columnwise_scale_inv_shape()); + EXPECT_EQ(count_byte_mismatches(out.columnwise_cpu_scale_inv_ptr(), + ref.columnwise_cpu_scale_inv_ptr(), s), + 0u) + << "colwise scales"; + } + if (is_dbias(mode)) { + const bool rowwise_reduction = std::is_same_v && !colwise; + const std::vector expected = reference_dbias(z, rows, cols, rowwise_reduction); + dbias.to_cpu(); + const IType *got = dbias.rowwise_cpu_dptr(); + size_t mismatches = 0; + for (size_t c = 0; c < cols; ++c) { + const IType e = static_cast(expected[c]); + mismatches += std::memcmp(&e, &got[c], sizeof(IType)) != 0 && + !(std::isnan(static_cast(e)) && std::isnan(static_cast(got[c]))); + } + EXPECT_EQ(mismatches, 0u) << "dbias"; + } +} + +using Params = std::tuple, bool /*use_table*/>; + +class CastMXFP8ActExactTestSuite : public ::testing::TestWithParam {}; + +} // namespace + +TEST_P(CastMXFP8ActExactTestSuite, MatchesScalarActivationThenCast) { + cudaDeviceProp prop; + ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); + if (prop.major < 10) { + GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; + } + const auto [mode, fill, itype, layout, swizzled, dims, use_table] = GetParam(); + const bool rowwise = layout != 1; + const bool colwise = layout != 0; + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( + itype, IType, + run_case(mode, fill, dims.first, dims.second, rowwise, colwise, swizzled, + use_table);); +} + +namespace { + +std::string case_name(const testing::TestParamInfo &info) { + const char *layouts[] = {"Rowwise", "Colwise", "Both"}; + const auto [mode, fill, itype, layout, swizzled, dims, use_table] = info.param; + const char *fills[] = {"XTypical", "XOrderedInputs", "XSpecials", "XTinyGradients"}; + return std::string(mode_name(mode)) + fills[static_cast(fill)] + "X" + + test::typeName(itype) + "X" + layouts[layout] + (swizzled ? "XSwizzled" : "XCompact") + + "X" + std::to_string(dims.first) + "X" + std::to_string(dims.second) + + (use_table ? "XTable" : "XNoTable"); +} + +} // namespace + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, CastMXFP8ActExactTestSuite, + ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::DBiasDGeLU, Mode::SiLU, + Mode::DSiLU, Mode::DBiasDSiLU), + ::testing::Values(Fill::Typical, Fill::OrderedInputs), + ::testing::Values(DType::kBFloat16, DType::kFloat16, DType::kFloat32), + ::testing::Values(0, 1, 2), ::testing::Bool(), + ::testing::Values(std::pair{1024, 2048}, + std::pair{544, 2080}), + ::testing::Bool()), + case_name); + +INSTANTIATE_TEST_SUITE_P( + SpecialValues, CastMXFP8ActExactTestSuite, + ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::DBiasDGeLU, Mode::SiLU, + Mode::DSiLU, Mode::DBiasDSiLU), + ::testing::Values(Fill::Specials), + ::testing::Values(DType::kBFloat16, DType::kFloat16, DType::kFloat32), + ::testing::Values(0, 1, 2), ::testing::Bool(), + ::testing::Values(std::pair{64, 128}), ::testing::Bool()), + case_name); + +INSTANTIATE_TEST_SUITE_P(TinyGradients, CastMXFP8ActExactTestSuite, + ::testing::Combine(::testing::Values(Mode::DGeLU, Mode::DBiasDGeLU, + Mode::DSiLU, Mode::DBiasDSiLU), + ::testing::Values(Fill::TinyGradients), + ::testing::Values(DType::kBFloat16, DType::kFloat16, + DType::kFloat32), + ::testing::Values(0, 1, 2), ::testing::Bool(), + ::testing::Values(std::pair{64, 128}), + ::testing::Bool()), + case_name); diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 58ea430d8a6..cd96d65beb4 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -7,6 +7,7 @@ #pragma once #include +#include #include #include #include @@ -452,6 +453,30 @@ struct Numeric_Traits { static constexpr int maxExpNorm = 1 << maxUnbiasedExponentAsFP32; }; +template <> +struct Numeric_Traits { + static constexpr double minSubnorm = 0x1p-24; + static constexpr double maxSubnorm = 0x1p-14 - minSubnorm; + static constexpr double minNorm = 0x1p-14; + static constexpr double maxNorm = 0x1.ffcp15; + static constexpr double artifInf = std::numeric_limits::infinity(); + static constexpr int maxBiasedExponentAsFP32 = 15 + FP32_EXPONENT_BIAS; + static constexpr int maxUnbiasedExponentAsFP32 = 15; + static constexpr int maxExpNorm = 1 << maxUnbiasedExponentAsFP32; +}; + +template <> +struct Numeric_Traits { + static constexpr double minSubnorm = 0x1p-133; + static constexpr double maxSubnorm = 0x1p-126 - minSubnorm; + static constexpr double minNorm = 0x1p-126; + static constexpr double maxNorm = 0x1.fep127; + static constexpr double artifInf = std::numeric_limits::infinity(); + static constexpr int maxBiasedExponentAsFP32 = 127 + FP32_EXPONENT_BIAS; + static constexpr int maxUnbiasedExponentAsFP32 = 127; + static constexpr double maxExpNorm = 0x1p127; +}; + template <> struct Numeric_Traits { static constexpr double minSubnorm = std::numeric_limits::denorm_min(); // std::pow(2.0, -149.0); diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index b80a588fe35..24acec658c6 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_activation_table.cu test_nvrtc.cpp test_string.cpp ../test_common.cu) diff --git a/tests/cpp/util/test_activation_table.cu b/tests/cpp/util/test_activation_table.cu new file mode 100644 index 00000000000..9cc81ec3805 --- /dev/null +++ b/tests/cpp/util/test_activation_table.cu @@ -0,0 +1,196 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// Every BF16 input read from each activation table, against the specification written out +// here independently of util/activation_table.cuh: +// - forward: util/math.h rounded to BF16, except that -inf gives -0; +// - derivative: util/math.h, except below the table window, where it is the linear Taylor +// term rounded once. +// The derivative is also checked against util/math.h with the ulp bound stated in the API. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "util/activation_table.cuh" +#include "util/math.h" + +namespace { + +using transformer_engine::Empty; + +constexpr int kNumInputs = 1 << 16; + +// Below-window bounds and linear forms of the specification. +constexpr uint32_t kDgeluWindowLow = 0x3980; // 2^-12 +constexpr uint32_t kDsiluWindowLow = 0x3B00; // 2^-9 + +__device__ float spec(const NVTEActivationTableKind kind, const uint32_t bits) { + const float x = __uint_as_float(bits << 16); + const uint32_t magnitude = bits & 0x7fffu; + switch (kind) { + case kNVTEActivationTableGelu: + if (bits == 0xff80u) return -0.0f; + return __bfloat162float(__float2bfloat16_rn(transformer_engine::gelu(x, {}))); + case kNVTEActivationTableSilu: + if (bits == 0xff80u) return -0.0f; + return __bfloat162float(__float2bfloat16_rn(transformer_engine::silu(x, {}))); + case kNVTEActivationTableDgelu: + if (magnitude < kDgeluWindowLow) return fmaf(0.79788456f, x, 0.5f); + return transformer_engine::dgelu(x, {}); + default: + if (magnitude < kDsiluWindowLow) return fmaf(0.5f, x, 0.5f); + return transformer_engine::dsilu(x, {}); + } +} + +__device__ float scalar_math(const NVTEActivationTableKind kind, const float x) { + switch (kind) { + case kNVTEActivationTableDgelu: + return transformer_engine::dgelu(x, {}); + default: + return transformer_engine::dsilu(x, {}); + } +} + +template +struct Op; +template <> +struct Op { + static constexpr float (*value)(float, const Empty &) = transformer_engine::gelu; +}; +template <> +struct Op { + static constexpr float (*value)(float, const Empty &) = transformer_engine::silu; +}; +template <> +struct Op { + static constexpr float (*value)(float, const Empty &) = transformer_engine::dgelu; +}; +template <> +struct Op { + static constexpr float (*value)(float, const Empty &) = transformer_engine::dsilu; +}; + +template +__global__ void lookup_kernel(const void *table, float *looked_up, float *expected, + float *math_h) { + const uint32_t bits = blockIdx.x * blockDim.x + threadIdx.x; + if (bits >= kNumInputs) return; + namespace at = transformer_engine::activation_table; + const float x = __uint_as_float(bits << 16); + looked_up[bits] = at::evaluate::value>(x, table); + expected[bits] = spec(kKind, bits); + math_h[bits] = (kKind == kNVTEActivationTableDgelu || kKind == kNVTEActivationTableDsilu) + ? scalar_math(kKind, x) + : 0.0f; +} + +uint32_t float_bits(const float f) { + uint32_t u; + std::memcpy(&u, &f, sizeof(u)); + return u; +} + +bool is_nan(const float f) { return (float_bits(f) & 0x7fffffffu) > 0x7f800000u; } + +// Monotone integer order of floats, for ulp distances. +int64_t ordered(const float f) { + const int32_t i = static_cast(float_bits(f)); + return i < 0 ? -static_cast(i & 0x7fffffff) : i; +} + +struct DeviceTable { + void *ptr = nullptr; + size_t bytes = 0; + explicit DeviceTable(const NVTEActivationTableKind kind) { + bytes = nvte_activation_table_size(kind); + EXPECT_EQ(cudaMalloc(&ptr, bytes), cudaSuccess); + transformer_engine::TensorWrapper table(ptr, std::vector{bytes}, + transformer_engine::DType::kByte); + nvte_build_activation_table(kind, table.data(), 0); + EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); + } + ~DeviceTable() { cudaFree(ptr); } +}; + +template +void check_table(const int max_ulp_below_window) { + DeviceTable table(kKind); + float *d = nullptr; + ASSERT_EQ(cudaMalloc(&d, 3 * kNumInputs * sizeof(float)), cudaSuccess); + lookup_kernel<<>>(table.ptr, d, d + kNumInputs, d + 2 * kNumInputs); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + std::vector h(3 * kNumInputs); + ASSERT_EQ(cudaMemcpy(h.data(), d, h.size() * sizeof(float), cudaMemcpyDeviceToHost), + cudaSuccess); + cudaFree(d); + const float *looked_up = h.data(), *expected = h.data() + kNumInputs, + *math_h = h.data() + 2 * kNumInputs; + + int mismatches = 0; + int64_t worst_ulp = 0; + for (uint32_t bits = 0; bits < kNumInputs; ++bits) { + const float got = looked_up[bits], want = expected[bits]; + const bool same = float_bits(got) == float_bits(want) || (is_nan(got) && is_nan(want)); + if (!same && mismatches++ < 8) { + ADD_FAILURE() << "input 0x" << std::hex << bits << ": table 0x" << float_bits(got) + << ", specification 0x" << float_bits(want); + } + if (max_ulp_below_window >= 0 && !is_nan(math_h[bits])) { + worst_ulp = std::max(worst_ulp, std::llabs(ordered(got) - ordered(math_h[bits]))); + } + } + EXPECT_EQ(mismatches, 0); + if (max_ulp_below_window >= 0) { + EXPECT_LE(worst_ulp, max_ulp_below_window) << "distance from util/math.h in ulp"; + } +} + +} // namespace + +TEST(ActivationTableTest, GeluMatchesSpecification) { + check_table(-1); +} +TEST(ActivationTableTest, SiluMatchesSpecification) { + check_table(-1); +} +TEST(ActivationTableTest, DgeluMatchesSpecification) { + check_table(1); +} +TEST(ActivationTableTest, DsiluMatchesSpecification) { + check_table(2); +} + +TEST(ActivationTableTest, SizesIdentifyKind) { + const size_t sizes[] = {nvte_activation_table_size(kNVTEActivationTableGelu), + nvte_activation_table_size(kNVTEActivationTableSilu), + nvte_activation_table_size(kNVTEActivationTableDgelu), + nvte_activation_table_size(kNVTEActivationTableDsilu)}; + for (int i = 0; i < 4; ++i) { + for (int j = i + 1; j < 4; ++j) EXPECT_NE(sizes[i], sizes[j]); + } +} + +TEST(ActivationTableTest, RejectsTableOfWrongSize) { + const size_t bytes = nvte_activation_table_size(kNVTEActivationTableDgelu); + void *ptr = nullptr; + ASSERT_EQ(cudaMalloc(&ptr, bytes), cudaSuccess); + transformer_engine::TensorWrapper table(ptr, std::vector{bytes}, + transformer_engine::DType::kByte); + EXPECT_THROW(nvte_build_activation_table(kNVTEActivationTableDsilu, table.data(), 0), + std::runtime_error); + cudaFree(ptr); +} diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index f90f74449c5..6fbe2111256 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -332,6 +332,7 @@ list(APPEND transformer_engine_cuda_sources list(APPEND transformer_engine_cuda_arch_specific_sources fused_attn/flash_attn.cu + activation/activation_table.cu activation/gelu.cu activation/gelu_dbias.cu activation/gelu_grouped.cu diff --git a/transformer_engine/common/activation/activation_table.cu b/transformer_engine/common/activation/activation_table.cu new file mode 100644 index 00000000000..b961a6ec832 --- /dev/null +++ b/transformer_engine/common/activation/activation_table.cu @@ -0,0 +1,102 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// Builds the activation tables from the scalar util/math.h functions. This file must not be +// compiled with --use_fast_math: the tables have to hold exactly what util/math.h computes. + +#include + +#include "../common.h" +#include "../util/activation_table.cuh" + +namespace transformer_engine { +namespace activation_table { +namespace { + +constexpr int kBlockSize = 256; + +template +__global__ void build_forward_kernel(__nv_bfloat16 *table) { + using T = Table; + const uint32_t row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= T::kRows) return; + const uint32_t magnitude = T::kLowSentinel + row; + __nv_bfloat16 positive, negative; + if (magnitude == T::kLowSentinel) { + positive = negative = __ushort_as_bfloat16(0xff80); // -inf + } else if (magnitude == T::kHighSentinel) { + positive = __ushort_as_bfloat16(0x7f80); // +inf + negative = __ushort_as_bfloat16(0x8000); // -0 + } else { + positive = __float2bfloat16_rn(T::op(__uint_as_float(magnitude << 16))); + negative = __float2bfloat16_rn(T::op(__uint_as_float((magnitude | 0x8000u) << 16))); + } + table[2 * row] = positive; + table[2 * row + 1] = negative; +} + +template +__global__ void build_derivative_kernel(float *table) { + using T = Table; + const uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= T::kSignStride + T::kEntriesPerSign) return; + if (i >= T::kEntriesPerSign && i < T::kSignStride) { + table[i] = 0.0f; // Unused gap between the two signs. + return; + } + const bool negative = i >= T::kSignStride; + const uint32_t magnitude = T::kLoBits + (negative ? i - T::kSignStride : i); + table[i] = T::op(__uint_as_float(((negative ? 0x8000u : 0u) | magnitude) << 16)); +} + +template +void build(void *table, cudaStream_t stream) { + using T = Table; + if constexpr (T::kForward) { + build_forward_kernel<<(T::kRows, kBlockSize), kBlockSize, 0, stream>>>( + static_cast<__nv_bfloat16 *>(table)); + } else { + constexpr uint32_t entries = T::kSignStride + T::kEntriesPerSign; + build_derivative_kernel<<(entries, kBlockSize), kBlockSize, 0, stream>>>( + static_cast(table)); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace +} // namespace activation_table +} // namespace transformer_engine + +size_t nvte_activation_table_size(NVTEActivationTableKind kind) { + NVTE_API_CALL(nvte_activation_table_size); + const size_t bytes = transformer_engine::activation_table::table_bytes(kind); + NVTE_CHECK(bytes != 0, "Invalid NVTEActivationTableKind (got ", static_cast(kind), ")"); + return bytes; +} + +void nvte_build_activation_table(NVTEActivationTableKind kind, NVTETensor table, + cudaStream_t stream) { + NVTE_API_CALL(nvte_build_activation_table); + using namespace transformer_engine; + using namespace transformer_engine::activation_table; + void *ptr = const_cast(checked_table_ptr(*convertNVTETensorCheck(table), kind)); + switch (kind) { + case kNVTEActivationTableGelu: + build(ptr, stream); + break; + case kNVTEActivationTableSilu: + build(ptr, stream); + break; + case kNVTEActivationTableDgelu: + build(ptr, stream); + break; + case kNVTEActivationTableDsilu: + build(ptr, stream); + break; + default: + NVTE_ERROR("Invalid NVTEActivationTableKind (got ", static_cast(kind), ")"); + } +} diff --git a/transformer_engine/common/activation/activation_template.h b/transformer_engine/common/activation/activation_template.h index ffbffafd1a1..c4f37b53d81 100644 --- a/transformer_engine/common/activation/activation_template.h +++ b/transformer_engine/common/activation/activation_template.h @@ -21,15 +21,16 @@ namespace transformer_engine { template -void act_fn(const NVTETensor input, NVTETensor output, cudaStream_t stream) { +void act_fn(const NVTETensor input, NVTETensor output, cudaStream_t stream, + const NVTEQuantizationConfig quant_config = nullptr) { using namespace detail; constexpr bool IS_ACT = true; - dispatch::quantize_fwd_helper(input, output, nullptr, stream); + dispatch::quantize_fwd_helper(input, output, quant_config, stream); } template -void dact_fn(const NVTETensor grad, const NVTETensor input, NVTETensor output, - cudaStream_t stream) { +void dact_fn(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream, + const NVTEQuantizationConfig quant_config = nullptr) { using namespace detail; constexpr bool IS_DBIAS = false; constexpr bool IS_DACT = true; @@ -37,7 +38,7 @@ void dact_fn(const NVTETensor grad, const NVTETensor input, NVTETensor output, constexpr NVTETensor workspace = nullptr; dispatch::quantize_bwd_helper(grad, input, output, dbias, workspace, - nullptr, stream); + quant_config, stream); } template diff --git a/transformer_engine/common/activation/gelu.cu b/transformer_engine/common/activation/gelu.cu index 6bd63672caa..9c4b1218f06 100644 --- a/transformer_engine/common/activation/gelu.cu +++ b/transformer_engine/common/activation/gelu.cu @@ -13,6 +13,13 @@ void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { act_fn>(input, output, stream); } +void nvte_gelu_v2(const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTE_API_CALL(nvte_gelu_v2); + using namespace transformer_engine; + act_fn>(input, output, stream, quant_config); +} + void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dgelu); @@ -20,6 +27,13 @@ void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, stream); } +void nvte_dgelu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTE_API_CALL(nvte_dgelu_v2); + using namespace transformer_engine; + dact_fn>(grad, input, output, stream, quant_config); +} + void nvte_geglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_geglu); using namespace transformer_engine; diff --git a/transformer_engine/common/activation/gelu_dbias.cu b/transformer_engine/common/activation/gelu_dbias.cu index 4eaa9e355b3..e3215931a65 100644 --- a/transformer_engine/common/activation/gelu_dbias.cu +++ b/transformer_engine/common/activation/gelu_dbias.cu @@ -20,6 +20,19 @@ void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor activati input, activation_input, output, dbias, workspace, nullptr, stream); } +void nvte_quantize_dbias_dgelu_v2(const NVTETensor input, const NVTETensor activation_input, + NVTETensor output, NVTETensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_dbias_dgelu_v2); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, quant_config, stream); +} + void nvte_quantize_dbias_dqgelu(const NVTETensor input, const NVTETensor activation_input, NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream) { diff --git a/transformer_engine/common/activation/swiglu.cu b/transformer_engine/common/activation/swiglu.cu index 0300ba1bd46..d4ae6edeb9e 100644 --- a/transformer_engine/common/activation/swiglu.cu +++ b/transformer_engine/common/activation/swiglu.cu @@ -15,6 +15,13 @@ void nvte_silu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { act_fn>(input, output, stream); } +void nvte_silu_v2(const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTE_API_CALL(nvte_silu_v2); + using namespace transformer_engine; + act_fn>(input, output, stream, quant_config); +} + void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dsilu); @@ -22,6 +29,13 @@ void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, stream); } +void nvte_dsilu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTE_API_CALL(nvte_dsilu_v2); + using namespace transformer_engine; + dact_fn>(grad, input, output, stream, quant_config); +} + void nvte_swiglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_swiglu); using namespace transformer_engine; diff --git a/transformer_engine/common/activation/swiglu_dbias.cu b/transformer_engine/common/activation/swiglu_dbias.cu index 0e532acc57d..4120b949e53 100644 --- a/transformer_engine/common/activation/swiglu_dbias.cu +++ b/transformer_engine/common/activation/swiglu_dbias.cu @@ -19,3 +19,16 @@ void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor activati dispatch::quantize_bwd_helper>( input, activation_input, output, dbias, workspace, nullptr, stream); } + +void nvte_quantize_dbias_dsilu_v2(const NVTETensor input, const NVTETensor activation_input, + NVTETensor output, NVTETensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_dbias_dsilu_v2); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, quant_config, stream); +} diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 13c3694d040..b748d65f260 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -38,6 +38,20 @@ namespace transformer_engine { namespace dispatch { +// Device pointer of the activation table to read for OP, or nullptr to compute OP. The table +// applies to BF16 inputs of the activations that have one. +template +const void *activation_lut_for(const QuantizationConfig &quant_config, const DType input_dtype) { + if constexpr (activation_table::kHasTable) { + if (quant_config.activation_table != nullptr && input_dtype == DType::kBFloat16) { + return activation_table::checked_table_ptr( + *convertNVTETensorCheck(quant_config.activation_table), + activation_table::kKindOf); + } + } + return nullptr; +} + template void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const NVTEQuantizationConfig quant_config, cudaStream_t stream) { @@ -95,18 +109,24 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const Tensor *dummy_input_tensor = nullptr; Tensor *dummy_dbias_tensor = nullptr; Tensor *dummy_workspace_tensor = nullptr; + const void *activation_lut = + IS_ACT ? activation_lut_for(quant_config_cpp, input_tensor->dtype()) + : nullptr; bool quantized_with_cutedsl = false; #ifdef NVTE_WITH_CUTEDSL - quantized_with_cutedsl = - cutedsl_backend::mxfp8_quantize_cutedsl( - input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, - dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); + // The CuTeDSL backend computes activations itself, so it cannot honor a table. + if (activation_lut == nullptr) { + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_cutedsl( + input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, + dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); + } #endif if (!quantized_with_cutedsl) { mxfp8::quantize( *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, - dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); + dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, activation_lut, stream); } break; } @@ -284,17 +304,25 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens break; } case NVTE_MXFP8_1D_SCALING: { + const void *activation_lut = nullptr; + if constexpr (IS_DACT) { + activation_lut = activation_lut_for(quant_config_cpp, input_tensor->dtype()); + } bool quantized_with_cutedsl = false; #ifdef NVTE_WITH_CUTEDSL - quantized_with_cutedsl = - cutedsl_backend::mxfp8_quantize_cutedsl( - grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, - quant_config_cpp.mxfp8_2d_quantization, stream); + // The CuTeDSL backend computes activations itself, so it cannot honor a table. + if (activation_lut == nullptr) { + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_cutedsl( + grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, + workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); + } #endif if (!quantized_with_cutedsl) { mxfp8::quantize( *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, - quant_config_cpp.mxfp8_2d_quantization, stream); + quant_config_cpp.mxfp8_2d_quantization, activation_lut, stream); } break; } diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index aaf03b0cb95..4d00bf62d7e 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -17,6 +17,7 @@ #include #include "../../common.h" +#include "../../util/activation_table.cuh" #include "../../util/cuda_runtime.h" #include "../../util/math.h" #include "../../util/ptx_arch_spec.cuh" @@ -48,7 +49,8 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 4 = 128 template + size_t CHUNK_DIM_X, size_t THREADS_PER_CHUNK, bool kIs2DBlockScaling, + bool kUseActivationTable> __global__ void __launch_bounds__(THREADS_PER_CHUNK) quantize_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_act_input, @@ -57,7 +59,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) e8m0_t *const scales_rowwise, e8m0_t *const scales_colwise, const float *noop, float *const dbias_workspace, float *const amax_ptr, const size_t rows, const size_t cols, const size_t scale_stride_rowwise, - const size_t scale_stride_colwise) { + const size_t scale_stride_colwise, const void *activation_lut) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) constexpr bool COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; @@ -266,11 +268,12 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float elt = static_cast(in_sh[shmem_offset_colwise]); if constexpr (IS_ACT) { - elt = OP(elt, {}); + elt = activation_table::evaluate(elt, activation_lut); } if constexpr (IS_DACT) { float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); - elt *= OP(act_in_elt, {}); + elt *= activation_table::evaluate(act_in_elt, + activation_lut); } if constexpr (DBIAS_REDUCTION_IN_COLWISE) { // Accumulate before the truncation below so the partial sums stay full precision @@ -356,11 +359,12 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float elt = static_cast(in_sh[shmem_offset_colwise]); if constexpr (IS_ACT) { - elt = OP(elt, {}); + elt = activation_table::evaluate(elt, activation_lut); } if constexpr (IS_DACT) { const float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); - elt *= OP(act_in_elt, {}); + elt *= activation_table::evaluate(act_in_elt, + activation_lut); } partial_dbias_colwise += elt; // Cache computed activations to avoid computing them again in the rowwise pass @@ -464,11 +468,13 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // Compute element float elt = static_cast(in.data.elt[e]); if constexpr (IS_ACT) { - elt = OP(elt, {}); + elt = + activation_table::evaluate(elt, activation_lut); } if constexpr (IS_DACT) { float act_in_elt = static_cast(act_in.data.elt[e]); - elt *= OP(act_in_elt, {}); + elt *= activation_table::evaluate(act_in_elt, + activation_lut); } // If DBIAS was computed in the 1st pass (COLWISE) then no need to compute it again @@ -689,7 +695,7 @@ template void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, // TODO (ksivamani) Tensor *output, Tensor *dbias, Tensor *workspace, const bool use_2d_quantization, - cudaStream_t stream) { + const void *activation_lut, cudaStream_t stream) { using namespace quantize_kernel; checkCuDriverContext(stream); @@ -1003,58 +1009,53 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - switch (scaling_type) { - case ScalingType::ROWWISE: { - TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { - auto kernel = - quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + // The table variant is instantiated only for BF16 inputs of the activations that + // have a table. + constexpr bool kTableCapable = + activation_table::kHasTable && std::is_same_v; + // Lambdas cannot capture the structured bindings rows and cols before C++20. + const size_t num_rows = rows; + const size_t num_cols = cols; + const auto launch = + [&](auto rowwise, auto colwise) { + TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { + const auto launch_kernel = [&](auto use_table) { + auto kernel = quantize_mxfp8_kernel< + IS_DBIAS, IS_DACT, IS_ACT, ParamOP, OP, IType, OType, + decltype(rowwise)::value, decltype(colwise)::value, + WITH_GEMM_SWIZZLED_SCALES, CHUNK_DIM_Y, CHUNK_DIM_X, THREADS_PER_CHUNK, + kIs2DBlockScaling, decltype(use_table)::value>; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, + noop_ptr, workspace_ptr, amax_ptr, num_rows, num_cols, + scale_stride_rowwise, scale_stride_colwise, activation_lut); + }; + if constexpr (kTableCapable) { + if (activation_lut != nullptr) { + launch_kernel(std::true_type{}); + } else { + launch_kernel(std::false_type{}); + } + } else { + launch_kernel(std::false_type{}); + } + }); + }; - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); - }); + switch (scaling_type) { + case ScalingType::ROWWISE: + launch(std::true_type{}, std::false_type{}); break; - } - case ScalingType::COLWISE: { - TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { - auto kernel = - quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); - }); + case ScalingType::COLWISE: + launch(std::false_type{}, std::true_type{}); break; - } - case ScalingType::BIDIMENSIONAL: { - TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { - auto kernel = - quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); - }); + case ScalingType::BIDIMENSIONAL: + launch(std::true_type{}, std::true_type{}); break; - } } NVTE_CHECK_CUDA(cudaGetLastError()); if constexpr (IS_DBIAS) { diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 7d201f6c8ba..621b6edd8c2 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -648,6 +648,7 @@ struct QuantizationConfig { NVTENVFP44Over6Mode nvfp4_4over6_mode = kNVTENVFP44Over6Disabled; bool nvfp4_4over6_err_use_fast_math = false; bool mxfp8_2d_quantization = false; + NVTETensor activation_table = nullptr; static constexpr size_t attr_sizes[] = { sizeof(uint8_t), // force_pow_2_scales @@ -660,7 +661,8 @@ struct QuantizationConfig { sizeof(uint8_t), // use_fast_math sizeof(uint8_t), // nvfp4_4over6_mode sizeof(uint8_t), // nvfp4_4over6_err_use_fast_math - sizeof(uint8_t) // mxfp8_2d_quantization + sizeof(uint8_t), // mxfp8_2d_quantization + sizeof(NVTETensor) // activation_table }; }; diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 31429cf89cf..dc742665d20 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -43,6 +43,47 @@ enum class NVTE_Activation_Type { CLAMPED_SWIGLU }; +/*! \enum NVTEActivationTableKind + * \brief Activation that can be read from a lookup table. + * + * A BF16 input takes only 2^16 values, so GeLU, SiLU and their derivatives of BF16 inputs + * can be read from a small table instead of computed. Build each table once per device + * with nvte_build_activation_table and pass it to the _v2 activation functions through + * kNVTEQuantizationConfigActivationTable. The table is read for BF16 inputs quantized to + * MXFP8; every other input type and recipe ignores it. + * + * Against util/math.h, table results differ only as follows: + * - GeLU and SiLU of -inf are -0 (the formulas give NaN). + * - For |x| < 2^-12 (dGeLU) or 2^-9 (dSiLU), the derivative is its linear Taylor term + * 0.5 + c*x rounded once, within 1 ulp (dGeLU) or 2 ulp (dSiLU) of util/math.h. + */ +enum NVTEActivationTableKind { + kNVTEActivationTableGelu = 0, /*!< GeLU (tanh form), BF16 entries */ + kNVTEActivationTableSilu = 1, /*!< SiLU, BF16 entries */ + kNVTEActivationTableDgelu = 2, /*!< GeLU derivative (tanh form), FP32 entries */ + kNVTEActivationTableDsilu = 3, /*!< SiLU derivative, FP32 entries */ +}; + +/*! \brief Size in bytes of an activation table. + * + * \param[in] kind Activation the table holds. + * \return Size of the table in bytes. Each kind has a different size. + */ +size_t nvte_activation_table_size(NVTEActivationTableKind kind); + +/*! \brief Fill an activation table on the current device. + * + * The table is computed from the scalar activation functions that TE uses when no table + * is given. Build it once per device and reuse it for every call; it does not depend on + * the input. + * + * \param[in] kind Activation the table holds. + * \param[in,out] table Byte tensor of nvte_activation_table_size(kind) elements. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_build_activation_table(NVTEActivationTableKind kind, NVTETensor table, + cudaStream_t stream); + /*! \brief Computes the GeLU activation of the input. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -53,6 +94,20 @@ enum class NVTE_Activation_Type { */ void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the GeLU activation of the input, with a quantization config. + * + * Same as nvte_gelu. With kNVTEQuantizationConfigActivationTable set to a + * kNVTEActivationTableGelu table, BF16 inputs quantized to MXFP8 read the activation + * from the table. + * + * \param[in] input Input tensor for activation. + * \param[in,out] output Output tensor. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_gelu_v2(const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream); + /*! \brief Computes the GeLU activation of the grouped input. * If the scaling mode of the grouped output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -74,6 +129,20 @@ void nvte_group_gelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cu */ void nvte_silu(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the SiLU activation of the input, with a quantization config. + * + * Same as nvte_silu. With kNVTEQuantizationConfigActivationTable set to a + * kNVTEActivationTableSilu table, BF16 inputs quantized to MXFP8 read the activation + * from the table. + * + * \param[in] input Input tensor for activation. + * \param[in,out] output Output tensor. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_silu_v2(const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream); + /*! \brief Computes the SiLU activation of the grouped input. * If the scaling mode of the grouped output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -205,6 +274,21 @@ void nvte_group_srelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, c void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the GeLU activation gradient, with a quantization config. + * + * Same as nvte_dgelu. With kNVTEQuantizationConfigActivationTable set to a + * kNVTEActivationTableDgelu table, BF16 inputs quantized to MXFP8 read the derivative + * from the table. + * + * \param[in] grad Incoming gradient. + * \param[in] input Input tensor for activation. + * \param[in,out] output Output tensor. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_dgelu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream); + /*! \brief Computes the GeLU activation gradient of the grouped input. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -230,6 +314,21 @@ void nvte_group_dgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor inpu void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the SiLU activation gradient, with a quantization config. + * + * Same as nvte_dsilu. With kNVTEQuantizationConfigActivationTable set to a + * kNVTEActivationTableDsilu table, BF16 inputs quantized to MXFP8 read the derivative + * from the table. + * + * \param[in] grad Incoming gradient. + * \param[in] input Input tensor for activation. + * \param[in,out] output Output tensor. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_dsilu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream); + /*! \brief Computes the SiLU activation gradient of the grouped input. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 1af93d42919..bec9847345e 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -192,6 +192,23 @@ void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor act_inpu NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); +/*! \brief Same as nvte_quantize_dbias_dgelu, with a quantization config. + * + * With kNVTEQuantizationConfigActivationTable set to a kNVTEActivationTableDgelu table, + * BF16 inputs quantized to MXFP8 read the derivative from the table. + * + * \param[in] input Input tensor to be cast. + * \param[in] act_input Activation input tensor. + * \param[in,out] output Output FP8/MXFP8 tensor. + * \param[out] dbias Result of the reduction of the input along columns. + * \param[out] workspace Workspace tensor. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_quantize_dbias_dgelu_v2(const NVTETensor input, const NVTETensor act_input, + NVTETensor output, NVTETensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream); + /*! \brief Computes backward of GeLU operation on the grouped input, then casts to FP8/MXFP8. * Additionally, reduces the result of the GeLU backward along columns. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, @@ -240,6 +257,23 @@ void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor act_inpu NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); +/*! \brief Same as nvte_quantize_dbias_dsilu, with a quantization config. + * + * With kNVTEQuantizationConfigActivationTable set to a kNVTEActivationTableDsilu table, + * BF16 inputs quantized to MXFP8 read the derivative from the table. + * + * \param[in] input Input tensor to be cast. + * \param[in] act_input Activation input tensor. + * \param[in,out] output Output FP8/MXFP8 tensor. + * \param[out] dbias Result of the reduction of the input along columns. + * \param[out] workspace Workspace tensor. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_quantize_dbias_dsilu_v2(const NVTETensor input, const NVTETensor act_input, + NVTETensor output, NVTETensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream); + /*! \brief Computes backward of SiLU operation on the grouped input, then casts to FP8/MXFP8. * Additionally, reduces the result of the SiLU backward along columns. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 9a08b15a8a8..08802873674 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -436,6 +436,11 @@ enum NVTEQuantizationConfigAttribute { kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath = 9, /*! Whether to use 2D block scaling for MXFP8 */ kNVTEQuantizationConfigMXFP82DQuantization = 10, + /*! Activation lookup table (NVTETensor), built with nvte_build_activation_table for the + * activation being computed. Read by the _v2 activation functions for BF16 inputs + * quantized to MXFP8, and ignored otherwise. See NVTEActivationTableKind. + */ + kNVTEQuantizationConfigActivationTable = 11, kNVTEQuantizationConfigNumAttributes }; @@ -1512,6 +1517,12 @@ class QuantizationConfigWrapper { &val, sizeof(val)); } + /*! \brief Set activation lookup table */ + void set_activation_table(NVTETensor activation_table) { + nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigActivationTable, + &activation_table, sizeof(NVTETensor)); + } + /*! \brief Set whether to use stochastic rounding */ void set_stochastic_rounding(bool stochastic_rounding) { const auto val = static_cast(stochastic_rounding); diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 66693dca3a1..3b9bf0b48d3 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -1143,6 +1143,9 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigMXFP82DQuantization: bool_to_uint8(config_.mxfp8_2d_quantization, buf); break; + case kNVTEQuantizationConfigActivationTable: + std::memcpy(buf, &config_.activation_table, attr_size); + break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } @@ -1213,6 +1216,9 @@ void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigMXFP82DQuantization: uint8_to_bool(buf, config_.mxfp8_2d_quantization); break; + case kNVTEQuantizationConfigActivationTable: + std::memcpy(&config_.activation_table, buf, attr_size); + break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } diff --git a/transformer_engine/common/util/activation_table.cuh b/transformer_engine/common/util/activation_table.cuh new file mode 100644 index 00000000000..e330b46fcf6 --- /dev/null +++ b/transformer_engine/common/util/activation_table.cuh @@ -0,0 +1,227 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file activation_table.cuh + * \brief GeLU, SiLU and their derivatives of BF16 inputs, read from lookup tables. + * + * A BF16 input takes only 2^16 values, so these activations can be read from a table + * generated once from the scalar util/math.h functions (nvte_build_activation_table). + * Outside each table's window the result has a closed form, so every lookup is + * branch-free. + * + * Forward (GeLU, SiLU): BF16 entries. Below the window the result is RN(x/2); above it + * it is x for x > 0 and -0 for x < 0. The forward result is truncated to BF16 by every + * caller, so BF16 entries lose nothing. One difference from util/math.h: -inf gives -0, + * where the formula gives NaN (-inf * 0). + * + * Derivatives (dGeLU, dSiLU): FP32 entries, because callers multiply the derivative by + * the incoming gradient in FP32. Below the window the derivative is its linear Taylor + * term 0.5 + c*x, rounded once; it differs from util/math.h by at most 1 ulp (dGeLU) or + * 2 ulp (dSiLU) there. Above the window the derivative has saturated, and the table's + * last entry of the same sign holds the saturated value. Inputs where util/math.h + * returns NaN (dGeLU from 0x6044 on, where 0.107*x*x overflows; +-inf; NaN) keep NaN. + * + * The tables follow TE's tanh-based GeLU. An erf-based variant would need its own + * tables, a different forward window, and an FP32 derivative Phi(x) + x*phi(x). + */ + +#ifndef TRANSFORMER_ENGINE_UTIL_ACTIVATION_TABLE_CUH_ +#define TRANSFORMER_ENGINE_UTIL_ACTIVATION_TABLE_CUH_ + +#include +#include + +#include +#include + +#include "../common.h" +#include "math.h" + +namespace transformer_engine { +namespace activation_table { + +/*! \brief Forward table over the BF16 magnitudes [kLo, kHi]. + * + * Rows are indexed by magnitude bits from kLowSentinel to kHighSentinel, two entries per + * row: op(+m) then op(-m). The kLowSentinel row holds -inf, so max(entry, RN(x/2)) picks + * RN(x/2) below the window. The kHighSentinel row holds +inf and -0, so min(entry, upper) + * picks x above the window for x > 0 and -0 for x < 0. + */ +template +struct ForwardLayout { + static constexpr bool kForward = true; + static constexpr uint32_t kLowSentinel = kLo - 1; + static constexpr uint32_t kHighSentinel = kHi + 1; + static constexpr uint32_t kRows = kHighSentinel - kLowSentinel + 1; + static constexpr size_t kBytes = kRows * 2 * sizeof(__nv_bfloat16); +}; + +/*! \brief Derivative table over the BF16 magnitudes [kLo, kHi]. + * + * Positive inputs at entries [0, n), negative ones at [kSignStride, kSignStride + n). + * The stride turns the BF16 sign bit, 0x8000, into a byte offset of 0x8000 >> 2, which + * lets a packed lookup form both lanes' offsets in one 32-bit word. + */ +template +struct DerivativeLayout { + static constexpr bool kForward = false; + static constexpr uint32_t kLoBits = kLo; + static constexpr uint32_t kHiBits = kHi; + //! Magnitudes from here on (and NaN) give NaN, as util/math.h does. + static constexpr uint32_t kSaturatedEndBits = kSaturatedEnd; + static constexpr uint32_t kEntriesPerSign = kHi - kLo + 1; + static constexpr uint32_t kSignStride = (0x8000 >> 2) / sizeof(float); + static_assert(kEntriesPerSign <= kSignStride, "The two signs' entries must not overlap."); + static constexpr size_t kBytes = (kSignStride + kEntriesPerSign) * sizeof(float); +}; + +template +struct Table; + +//! gelu(x) is RN(x/2) below 0x3b4b, x from 2.75 and -0 from -5.1875. +template <> +struct Table : ForwardLayout<0x3b4b, 0x40a5> { + __device__ static float op(const float x) { return gelu(x, {}); } +}; + +//! silu(x) is RN(x/2) below 0x3bb6, x from 5.96875 and -0 from -89. +template <> +struct Table : ForwardLayout<0x3bb6, 0x42b1> { + __device__ static float op(const float x) { return silu(x, {}); } +}; + +//! [2^-12, 8]. dgelu has saturated, to 1 and to +0, by |x| = 8, until 0.107*x*x overflows at +//! 0x6044. +template <> +struct Table : DerivativeLayout<0x3980, 0x4100, 0x6044> { + __device__ static float op(const float x) { return dgelu(x, {}); } + //! dgelu(x) = 0.5 + sqrt(2/pi)*x + O(x^3). + __device__ static float below_window(const float x) { return fmaf(0.79788456f, x, 0.5f); } +}; + +//! [2^-9, 89]. dsilu has saturated, to 1 and to +0, by |x| = 89, up to inf. +template <> +struct Table : DerivativeLayout<0x3B00, 0x42B2, 0x7F80> { + __device__ static float op(const float x) { return dsilu(x, {}); } + //! dsilu(x) = 0.5 + x/2 + O(x^3). + __device__ static float below_window(const float x) { return fmaf(0.5f, x, 0.5f); } +}; + +/*! \brief Table size in bytes. Each kind has a different size, which lets callers check that + * a table matches the activation it is passed with. + */ +constexpr size_t table_bytes(const NVTEActivationTableKind kind) { + switch (kind) { + case kNVTEActivationTableGelu: + return Table::kBytes; + case kNVTEActivationTableSilu: + return Table::kBytes; + case kNVTEActivationTableDgelu: + return Table::kBytes; + case kNVTEActivationTableDsilu: + return Table::kBytes; + default: + return 0; + } +} + +static_assert(table_bytes(kNVTEActivationTableGelu) != table_bytes(kNVTEActivationTableSilu) && + table_bytes(kNVTEActivationTableDgelu) != table_bytes(kNVTEActivationTableDsilu), + "Table sizes must identify their kind."); + +/*! \brief Device pointer of an activation table, checked against the kind it must hold. */ +inline const void *checked_table_ptr(const Tensor &table, const NVTEActivationTableKind kind) { + const size_t bytes = table_bytes(kind); + NVTE_CHECK(bytes != 0, "Invalid NVTEActivationTableKind (got ", static_cast(kind), ")"); + NVTE_CHECK(table.data.dptr != nullptr, "Activation table is not allocated."); + NVTE_CHECK(table.data.dtype == DType::kByte, "Activation table must be a byte tensor."); + NVTE_CHECK(table.data.numel() == bytes, "Activation table has ", table.data.numel(), + " bytes, but activation table kind ", static_cast(kind), " needs ", bytes, + " (was it built for another activation?)"); + return table.data.dptr; +} + +/*! \brief The table kind for OP, if OP is read from a table. */ +template +constexpr bool kHasTable = false; +template +constexpr NVTEActivationTableKind kKindOf = kNVTEActivationTableGelu; + +template <> +constexpr bool kHasTable> = true; +template <> +constexpr bool kHasTable> = true; +template <> +constexpr bool kHasTable> = true; +template <> +constexpr bool kHasTable> = true; +template <> +constexpr NVTEActivationTableKind kKindOf> = kNVTEActivationTableSilu; +template <> +constexpr NVTEActivationTableKind kKindOf> = kNVTEActivationTableDgelu; +template <> +constexpr NVTEActivationTableKind kKindOf> = kNVTEActivationTableDsilu; + +/*! \brief Forward activation of x, a BF16 value widened to FP32. + * + * Branch-free: the magnitude is clamped to the table rows with __hmax/__hmin, which + * ignore NaN, so NaN lands on the -inf row; __hmin_nan/__hmax_nan then propagate it. + */ +template +__device__ __forceinline__ float forward(const float x, const __nv_bfloat16 *__restrict__ table) { + using T = Table; + static_assert(T::kForward, "Not a forward table."); + const __nv_bfloat16 a = __float2bfloat16_rn(x); + const uint32_t sign = __float_as_uint(x) >> 31; + const __nv_bfloat16 magnitude = + __hmin(__hmax(__habs(a), __ushort_as_bfloat16(static_cast(T::kLowSentinel))), + __ushort_as_bfloat16(static_cast(T::kHighSentinel))); + const uint32_t row = __bfloat16_as_ushort(magnitude) - T::kLowSentinel; + const __nv_bfloat16 entry = __ldg(table + 2 * row + sign); + // RN(x/2), as util/math.h rounds x*0.5 below the window. + const __nv_bfloat16 half = __hmul(a, __ushort_as_bfloat16(0x3f00)); + // a for a > 0, -0 for a < 0, NaN for NaN. + const __nv_bfloat16 upper = __hmax_nan(a, __ushort_as_bfloat16(0x8000)); + return __bfloat162float(__hmin_nan(__hmax(entry, half), upper)); +} + +/*! \brief Derivative of x, a BF16 value widened to FP32. Branch-free. */ +template +__device__ __forceinline__ float derivative(const float x, const float *__restrict__ table) { + using T = Table; + static_assert(!T::kForward, "Not a derivative table."); + const uint32_t bits = __float_as_uint(x) >> 16; + const uint32_t magnitude = bits & 0x7fffu; + const uint32_t clamped = min(max(magnitude, T::kLoBits), T::kHiBits); + const uint32_t index = (clamped - T::kLoBits) + (bits >> 15) * T::kSignStride; + const float entry = __ldg(table + index); + const float value = magnitude < T::kLoBits ? T::below_window(x) : entry; + return magnitude >= T::kSaturatedEndBits ? __uint_as_float(0x7fffffffu) : value; +} + +/*! \brief OP(x), from the table when kUseTable, else from util/math.h. + * + * x must be a BF16 value widened to FP32 when kUseTable. + */ +template +__device__ __forceinline__ float evaluate(const float x, const void *table) { + if constexpr (kUseTable) { + static_assert(kHasTable, "OP has no table."); + constexpr NVTEActivationTableKind kKind = kKindOf; + if constexpr (Table::kForward) { + return forward(x, static_cast(table)); + } else { + return derivative(x, static_cast(table)); + } + } else { + return OP(x, {}); + } +} + +} // namespace activation_table +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_UTIL_ACTIVATION_TABLE_CUH_ diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 893374b3de3..6506c4c4cc2 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -243,6 +243,11 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(EpBootstrapBorrowedCommHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(TopkHandler); pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k); +// Activation lookup table to pass for act_type (or its derivative), built once per device; +// nullptr when the activation is not read from a table (see NVTEActivationTableKind). +NVTETensor activation_table_for(NVTE_Activation_Type act_type, bool derivative, DType input_dtype, + JAXX_Scaling_Mode scaling_mode, cudaStream_t stream); + } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/activation.cpp b/transformer_engine/jax/csrc/extensions/activation.cpp index 6325a700d15..b570c4c7ec4 100644 --- a/transformer_engine/jax/csrc/extensions/activation.cpp +++ b/transformer_engine/jax/csrc/extensions/activation.cpp @@ -104,18 +104,26 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal } switch (act_type) { - case NVTE_Activation_Type::GELU: - nvte_gelu(input_tensor.data(), output_tensor.data(), stream); + case NVTE_Activation_Type::GELU: { + QuantizationConfigWrapper config; + config.set_activation_table(activation_table_for( + act_type, /*derivative=*/false, static_cast(in_dtype), scaling_mode, stream)); + nvte_gelu_v2(input_tensor.data(), output_tensor.data(), config, stream); break; + } case NVTE_Activation_Type::GEGLU: nvte_geglu(input_tensor.data(), output_tensor.data(), stream); break; case NVTE_Activation_Type::GLU: nvte_glu(input_tensor.data(), output_tensor.data(), stream); break; - case NVTE_Activation_Type::SILU: - nvte_silu(input_tensor.data(), output_tensor.data(), stream); + case NVTE_Activation_Type::SILU: { + QuantizationConfigWrapper config; + config.set_activation_table(activation_table_for( + act_type, /*derivative=*/false, static_cast(in_dtype), scaling_mode, stream)); + nvte_silu_v2(input_tensor.data(), output_tensor.data(), config, stream); break; + } case NVTE_Activation_Type::SWIGLU: nvte_swiglu(input_tensor.data(), output_tensor.data(), stream); break; @@ -174,6 +182,12 @@ Error_Type ActLuInitializeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer int64_t act_enum, JAXX_Scaling_Mode scaling_mode, JAXX_Quantize_Layout quantize_layout, ActivationConfig act_params, bool output_amax_when_no_scaling) { + // Build the activation table outside the warm-up capture below, so that execution only + // reads it. + activation_table_for( + static_cast(act_enum), /*derivative=*/false, + static_cast(convert_ffi_datatype_to_te_dtype(input_buf.element_type())), scaling_mode, + stream); return wrapInStreamCapture(std::function(ActLuFFI), stream, input_buf, scale_buf, amax_buf, output_buf, colwise_output_buf, scale_inv_buf, colwise_scale_inv_buf, updated_amax_buf, act_enum, scaling_mode, quantize_layout, act_params, @@ -381,17 +395,20 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, is_quantize_2x2x(quantize_layout) && act_len == 2), "TE/common does not support delayed scaling for 2x with gated activations."); + QuantizationConfigWrapper table_config; + table_config.set_activation_table(activation_table_for( + act_type, /*derivative=*/true, static_cast(in_dtype), scaling_mode, stream)); if (is_dbias) { switch (act_type) { case NVTE_Activation_Type::GELU: - nvte_quantize_dbias_dgelu(input_tensor.data(), act_input_tensor.data(), - output_tensor.data(), dbias_tensor.data(), - workspace_tensor.data(), stream); + nvte_quantize_dbias_dgelu_v2(input_tensor.data(), act_input_tensor.data(), + output_tensor.data(), dbias_tensor.data(), + workspace_tensor.data(), table_config, stream); break; case NVTE_Activation_Type::SILU: - nvte_quantize_dbias_dsilu(input_tensor.data(), act_input_tensor.data(), - output_tensor.data(), dbias_tensor.data(), - workspace_tensor.data(), stream); + nvte_quantize_dbias_dsilu_v2(input_tensor.data(), act_input_tensor.data(), + output_tensor.data(), dbias_tensor.data(), + workspace_tensor.data(), table_config, stream); break; case NVTE_Activation_Type::RELU: nvte_quantize_dbias_drelu(input_tensor.data(), act_input_tensor.data(), @@ -415,10 +432,12 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, } else { switch (act_type) { case NVTE_Activation_Type::GELU: - nvte_dgelu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); + nvte_dgelu_v2(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), + table_config, stream); break; case NVTE_Activation_Type::SILU: - nvte_dsilu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); + nvte_dsilu_v2(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), + table_config, stream); break; case NVTE_Activation_Type::RELU: nvte_drelu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); @@ -489,6 +508,12 @@ Error_Type DActLuDBiasQuantizeInitializeFFI( Result_Type dbias_buf, Result_Type workspace_buf, JAXX_Scaling_Mode scaling_mode, int64_t act_enum, JAXX_Quantize_Layout quantize_layout, bool is_dbias, ActivationConfig act_params, bool output_amax_when_no_scaling) { + // Build the activation table outside the warm-up capture below, so that execution only + // reads it. + activation_table_for( + static_cast(act_enum), /*derivative=*/true, + static_cast(convert_ffi_datatype_to_te_dtype(input_buf.element_type())), scaling_mode, + stream); return wrapInStreamCapture(std::function(DActLuDBiasQuantizeFFI), stream, input_buf, act_input_buf, scale_buf, amax_buf, output_buf, colwise_output_buf, scale_inv_buf, colwise_scale_inv_buf, updated_amax_buf, dbias_buf, diff --git a/transformer_engine/jax/csrc/extensions/activation_table.cpp b/transformer_engine/jax/csrc/extensions/activation_table.cpp new file mode 100644 index 00000000000..a57999e3fb3 --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/activation_table.cpp @@ -0,0 +1,78 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include + +#include "../extensions.h" + +namespace transformer_engine { +namespace jax { + +namespace { + +struct CachedTable { + void *ptr = nullptr; + std::unique_ptr tensor; + bool ready = false; +}; + +NVTETensor activation_table(const NVTEActivationTableKind kind, cudaStream_t stream) { + static std::mutex mutex; + // Never destroyed: the tables live until the process exits. + static auto *tables = new std::map, CachedTable>(); + int device; + NVTE_CHECK_CUDA(cudaGetDevice(&device)); + + std::lock_guard lock(mutex); + CachedTable &table = (*tables)[{device, static_cast(kind)}]; + if (!table.tensor) { + // XLA gives custom calls no persistent device allocation, so the table is allocated + // directly. All tables together take tens of KB per device. + const size_t bytes = nvte_activation_table_size(kind); + NVTE_CHECK_CUDA(cudaMalloc(&table.ptr, bytes)); + table.tensor = + std::make_unique(table.ptr, std::vector{bytes}, DType::kByte); + } + if (!table.ready) { + nvte_build_activation_table(kind, table.tensor->data(), stream); + cudaStreamCaptureStatus capture; + NVTE_CHECK_CUDA(cudaStreamIsCapturing(stream, &capture)); + if (capture == cudaStreamCaptureStatusNone) { + // Later calls may come on other streams, so the build has to be complete. + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + table.ready = true; + } + // A build issued under graph capture only runs when the graph replays, so the next call + // outside capture builds again. Building is idempotent. + } + return table.tensor->data(); +} + +} // namespace + +NVTETensor activation_table_for(const NVTE_Activation_Type act_type, const bool derivative, + const DType input_dtype, const JAXX_Scaling_Mode scaling_mode, + cudaStream_t stream) { + if (input_dtype != DType::kBFloat16 || scaling_mode != JAXX_Scaling_Mode::MXFP8_1D_SCALING) { + return nullptr; + } + switch (act_type) { + case NVTE_Activation_Type::GELU: + return activation_table(derivative ? kNVTEActivationTableDgelu : kNVTEActivationTableGelu, + stream); + case NVTE_Activation_Type::SILU: + return activation_table(derivative ? kNVTEActivationTableDsilu : kNVTEActivationTableSilu, + stream); + default: + return nullptr; + } +} + +} // namespace jax +} // namespace transformer_engine diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 9a3c6314054..8a7ae66d5dd 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -629,6 +629,58 @@ void philox_unpack(at::PhiloxCudaState arg, int64_t* rng_state_ptr); // extract PhiloxCudaState from CUDA random number generator at::PhiloxCudaState init_philox_state(at::CUDAGeneratorImpl* gen, size_t elts_per_thread); +/*! \brief Activation table of the given kind for the current device. + * + * Built on first use and kept for the lifetime of the process, one per device. + */ +NVTETensor activation_table(NVTEActivationTableKind kind, cudaStream_t stream); + +/*! \brief Whether an activation of input written to output can be read from a table: BF16 + * inputs quantized to MXFP8. + */ +inline bool can_use_activation_table(const NVTETensor input, const NVTETensor output) { + return nvte_tensor_type(input) == kNVTEBFloat16 && + nvte_tensor_scaling_mode(output) == NVTE_MXFP8_1D_SCALING; +} + +/*! \brief An activation entry point that reads its table when it can. */ +template +void activation_with_table(const NVTETensor input, NVTETensor output, cudaStream_t stream) { + QuantizationConfigWrapper config; + if (can_use_activation_table(input, output)) { + config.set_activation_table(activation_table(kKind, stream)); + } + act_func_v2(input, output, config, stream); +} + +/*! \brief An activation-gradient entry point that reads its table when it can. */ +template +void dactivation_with_table(const NVTETensor grad, const NVTETensor input, NVTETensor output, + cudaStream_t stream) { + QuantizationConfigWrapper config; + if (can_use_activation_table(input, output)) { + config.set_activation_table(activation_table(kKind, stream)); + } + dact_func_v2(grad, input, output, config, stream); +} + +/*! \brief A fused dbias + activation-gradient entry point that reads its table when it can. */ +template +void dbias_dactivation_with_table(const NVTETensor grad, const NVTETensor input, NVTETensor output, + NVTETensor dbias, NVTETensor workspace, cudaStream_t stream) { + QuantizationConfigWrapper config; + if (can_use_activation_table(input, output)) { + config.set_activation_table(activation_table(kKind, stream)); + } + dact_dbias_func_v2(grad, input, output, dbias, workspace, config, stream); +} + } // namespace transformer_engine::pytorch namespace std { diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 6251872a036..9293cb3e077 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -247,11 +247,13 @@ py::object dactivation_helper(const at::Tensor& grad_output, const at::Tensor& i /* GELU and variants */ py::object gelu(const at::Tensor& input, py::handle quantizer) { - return activation_helper(input, quantizer); + return activation_helper, nullptr>( + input, quantizer); } py::object dgelu(const at::Tensor& grad, const at::Tensor& input, py::handle quantizer) { - return dactivation_helper(grad, input, quantizer); + return dactivation_helper, + nullptr>(grad, input, quantizer); } py::object glu(const at::Tensor& input, py::handle quantizer) { @@ -320,11 +322,13 @@ py::object dsreglu(const at::Tensor& grad, const at::Tensor& input, py::handle q } /* Silu and variants */ py::object silu(const at::Tensor& input, py::handle quantizer) { - return activation_helper(input, quantizer); + return activation_helper, nullptr>( + input, quantizer); } py::object dsilu(const at::Tensor& grad, const at::Tensor& input, py::handle quantizer) { - return dactivation_helper(grad, input, quantizer); + return dactivation_helper, + nullptr>(grad, input, quantizer); } py::object swiglu(const at::Tensor& input, py::handle quantizer) { diff --git a/transformer_engine/pytorch/csrc/extensions/activation_table.cpp b/transformer_engine/pytorch/csrc/extensions/activation_table.cpp new file mode 100644 index 00000000000..543242405c6 --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/activation_table.cpp @@ -0,0 +1,57 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include + +#include "../common.h" + +namespace transformer_engine::pytorch { + +namespace { + +struct CachedTable { + at::Tensor buffer; + std::unique_ptr tensor; + bool ready = false; +}; + +} // namespace + +NVTETensor activation_table(const NVTEActivationTableKind kind, cudaStream_t stream) { + static std::mutex mutex; + // Never destroyed: the tables live until the process exits, and destroying tensors after + // the CUDA context has been torn down is not safe. + static auto *tables = new std::map, CachedTable>(); + const int device = c10::cuda::current_device(); + + std::lock_guard lock(mutex); + CachedTable &table = (*tables)[{device, static_cast(kind)}]; + if (!table.tensor) { + const size_t bytes = nvte_activation_table_size(kind); + table.buffer = at::empty({static_cast(bytes)}, + at::TensorOptions().dtype(at::kByte).device(at::kCUDA, device)); + table.tensor = std::make_unique(table.buffer.data_ptr(), + std::vector{bytes}, DType::kByte); + } + if (!table.ready) { + nvte_build_activation_table(kind, table.tensor->data(), stream); + cudaStreamCaptureStatus capture; + NVTE_CHECK_CUDA(cudaStreamIsCapturing(stream, &capture)); + if (capture == cudaStreamCaptureStatusNone) { + // Later calls may come on other streams, so the build has to be complete. + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + table.ready = true; + } + // A build issued under graph capture only runs when the graph replays, so the next + // eager call builds again. Building is idempotent. + } + return table.tensor->data(); +} + +} // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/bias.cpp b/transformer_engine/pytorch/csrc/extensions/bias.cpp index 2dbeb965a1a..96178f34fff 100644 --- a/transformer_engine/pytorch/csrc/extensions/bias.cpp +++ b/transformer_engine/pytorch/csrc/extensions/bias.cpp @@ -266,12 +266,18 @@ std::vector dact_dbias( std::vector dbias_dgelu(const at::Tensor &grad_output, const at::Tensor &act_input, py::handle quantizer) { - return dact_dbias(nvte_quantize_dbias_dgelu, nvte_dgelu, grad_output, act_input, quantizer); + return dact_dbias( + dbias_dactivation_with_table, + dactivation_with_table, grad_output, act_input, + quantizer); } std::vector dbias_dsilu(const at::Tensor &grad_output, const at::Tensor &act_input, py::handle quantizer) { - return dact_dbias(nvte_quantize_dbias_dsilu, nvte_dsilu, grad_output, act_input, quantizer); + return dact_dbias( + dbias_dactivation_with_table, + dactivation_with_table, grad_output, act_input, + quantizer); } std::vector dbias_drelu(const at::Tensor &grad_output, const at::Tensor &act_input, From a16ee09837b845c31abe36e9a138db7fdefe33aa Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 13 Oct 2026 14:17:48 -0700 Subject: [PATCH 02/29] Add an experimental exact dSiLU activation table With NVTE_ACTIVATION_TABLE_EXACT_DSILU=1, the dSiLU table covers |x| >= 2^-25 instead of 2^-9 (32 KB instead of 16 KB). Below 2^-25 the linear term already equals util/math.h, so dSiLU read from the table matches util/math.h for every BF16 input. The variable is read once, when the table is built; kernels tell the two layouts apart by the table's size. Signed-off-by: Teddy Do --- docs/envvars.rst | 6 ++ .../cpp/operator/test_cast_mxfp8_act_exact.cu | 15 ++- tests/cpp/util/test_activation_table.cu | 43 ++++++--- .../common/activation/activation_table.cu | 34 ++++--- .../common/cast/dispatch/quantize.cuh | 24 ++--- .../common/cast/mxfp8/quantize_mxfp8.cuh | 40 +++++--- .../include/transformer_engine/activation.h | 5 +- .../common/util/activation_table.cuh | 95 ++++++++++++++----- 8 files changed, 180 insertions(+), 82 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index a0a64b3f637..02b1a2417a4 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -353,6 +353,12 @@ Kernel Configuration :Default: ``0`` :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. Existing transpose operations select their static fallback automatically. NVFP4 4over6 quantization, fused softmax, and normalization paths require their corresponding ``NVTE_BUILD_LEGACY_STATIC_NVFP4``, ``NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX``, or ``NVTE_BUILD_LEGACY_STATIC_NORM`` CMake option to have been enabled when the library was built; otherwise no static fallback is available. +.. envvar:: NVTE_ACTIVATION_TABLE_EXACT_DSILU + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Experimental. Choose the SiLU-derivative lookup table used when BF16 inputs are quantized to MXFP8 with an activation table. By default the 16 KB table covers ``|x| >= 2^-9`` and smaller inputs use the linear term ``0.5 + x/2``, which is within 2 ulp of the computed derivative. When set to ``1``, a 32 KB table covers ``|x| >= 2^-25``, and the derivative matches the computed one for every input. Read once per process, when the table is built. + .. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM :Type: ``int`` (0 or 1) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index dd79cf876e0..fa6a96aa729 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -67,14 +67,15 @@ const char *mode_name(const Mode m) { // act: y = OP(x); dact: y = grad * dOP(x), as the fused kernels form it in FP32. z keeps the FP32 // value (dbias sums it before truncation), y the value truncated to the input type. template -__global__ void scalar_activation_kernel(const Mode mode, const bool table_spec, const IType *x, +__global__ void scalar_activation_kernel(const Mode mode, const bool table_spec, + const float dsilu_window_low, const IType *x, const IType *grad, IType *y, float *z, const size_t n) { const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; if (i >= n) return; const float xv = static_cast(x[i]); const Empty e{}; // With a table, -inf gives -0 in the forward functions, and the derivatives below - // 2^-12 (dGeLU) and 2^-9 (dSiLU) are their linear Taylor terms. + // 2^-12 (dGeLU) and dsilu_window_low (dSiLU) are their linear Taylor terms. const bool negative_inf = table_spec && xv == -INFINITY; const float magnitude = fabsf(xv); float v; @@ -93,7 +94,8 @@ __global__ void scalar_activation_kernel(const Mode mode, const bool table_spec, break; default: v = static_cast(grad[i]); - v *= table_spec && magnitude < 0x1p-9f ? fmaf(0.5f, xv, 0.5f) : dsilu(xv, e); + v *= table_spec && magnitude < dsilu_window_low ? fmaf(0.5f, xv, 0.5f) + : dsilu(xv, e); break; } z[i] = v; @@ -290,8 +292,13 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t constexpr int block_size = 256; const size_t grid_size = divide_round_up(n, block_size); const bool table_spec = use_table && std::is_same_v; + // NVTE_ACTIVATION_TABLE_EXACT_DSILU=1 extends the dSiLU table down to 2^-25. + const char *exact_env = std::getenv("NVTE_ACTIVATION_TABLE_EXACT_DSILU"); + const float dsilu_window_low = + exact_env != nullptr && std::string(exact_env) == "1" ? 0x1p-25f : 0x1p-9f; scalar_activation_kernel - <<>>(mode, table_spec, static_cast(x.rowwise_dptr()), + <<>>(mode, table_spec, dsilu_window_low, + static_cast(x.rowwise_dptr()), static_cast(grad.rowwise_dptr()), static_cast(y.rowwise_dptr()), z_dev, n); ASSERT_EQ(cudaGetLastError(), cudaSuccess); diff --git a/tests/cpp/util/test_activation_table.cu b/tests/cpp/util/test_activation_table.cu index 9cc81ec3805..b94ee45ea83 100644 --- a/tests/cpp/util/test_activation_table.cu +++ b/tests/cpp/util/test_activation_table.cu @@ -10,6 +10,8 @@ // - derivative: util/math.h, except below the table window, where it is the linear Taylor // term rounded once. // The derivative is also checked against util/math.h with the ulp bound stated in the API. +// With NVTE_ACTIVATION_TABLE_EXACT_DSILU=1 the dSiLU window starts at 2^-25 and dSiLU must match +// util/math.h exactly. #include #include @@ -22,6 +24,7 @@ #include #include #include +#include #include #include "util/activation_table.cuh" @@ -34,10 +37,17 @@ using transformer_engine::Empty; constexpr int kNumInputs = 1 << 16; // Below-window bounds and linear forms of the specification. -constexpr uint32_t kDgeluWindowLow = 0x3980; // 2^-12 -constexpr uint32_t kDsiluWindowLow = 0x3B00; // 2^-9 +constexpr uint32_t kDgeluWindowLow = 0x3980; // 2^-12 +constexpr uint32_t kDsiluWindowLow = 0x3B00; // 2^-9 +constexpr uint32_t kExactDsiluWindowLow = 0x3300; // 2^-25 -__device__ float spec(const NVTEActivationTableKind kind, const uint32_t bits) { +bool exact_dsilu() { + const char *env = std::getenv("NVTE_ACTIVATION_TABLE_EXACT_DSILU"); + return env != nullptr && std::string(env) == "1"; +} + +__device__ float spec(const NVTEActivationTableKind kind, const uint32_t bits, + const uint32_t dsilu_window_low) { const float x = __uint_as_float(bits << 16); const uint32_t magnitude = bits & 0x7fffu; switch (kind) { @@ -51,7 +61,7 @@ __device__ float spec(const NVTEActivationTableKind kind, const uint32_t bits) { if (magnitude < kDgeluWindowLow) return fmaf(0.79788456f, x, 0.5f); return transformer_engine::dgelu(x, {}); default: - if (magnitude < kDsiluWindowLow) return fmaf(0.5f, x, 0.5f); + if (magnitude < dsilu_window_low) return fmaf(0.5f, x, 0.5f); return transformer_engine::dsilu(x, {}); } } @@ -84,15 +94,15 @@ struct Op { static constexpr float (*value)(float, const Empty &) = transformer_engine::dsilu; }; -template -__global__ void lookup_kernel(const void *table, float *looked_up, float *expected, - float *math_h) { +template +__global__ void lookup_kernel(const void *table, const uint32_t dsilu_window_low, float *looked_up, + float *expected, float *math_h) { const uint32_t bits = blockIdx.x * blockDim.x + threadIdx.x; if (bits >= kNumInputs) return; namespace at = transformer_engine::activation_table; const float x = __uint_as_float(bits << 16); - looked_up[bits] = at::evaluate::value>(x, table); - expected[bits] = spec(kKind, bits); + looked_up[bits] = at::evaluate::value>(x, table); + expected[bits] = spec(kKind, bits, dsilu_window_low); math_h[bits] = (kKind == kNVTEActivationTableDgelu || kKind == kNVTEActivationTableDsilu) ? scalar_math(kKind, x) : 0.0f; @@ -131,7 +141,18 @@ void check_table(const int max_ulp_below_window) { DeviceTable table(kKind); float *d = nullptr; ASSERT_EQ(cudaMalloc(&d, 3 * kNumInputs * sizeof(float)), cudaSuccess); - lookup_kernel<<>>(table.ptr, d, d + kNumInputs, d + 2 * kNumInputs); + namespace at = transformer_engine::activation_table; + const bool exact = kKind == kNVTEActivationTableDsilu && exact_dsilu(); + const uint32_t dsilu_window_low = exact ? kExactDsiluWindowLow : kDsiluWindowLow; + if (exact) { + if constexpr (kKind == kNVTEActivationTableDsilu) { + lookup_kernel<<>>( + table.ptr, dsilu_window_low, d, d + kNumInputs, d + 2 * kNumInputs); + } + } else { + lookup_kernel<<>>( + table.ptr, dsilu_window_low, d, d + kNumInputs, d + 2 * kNumInputs); + } ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); std::vector h(3 * kNumInputs); ASSERT_EQ(cudaMemcpy(h.data(), d, h.size() * sizeof(float), cudaMemcpyDeviceToHost), @@ -171,7 +192,7 @@ TEST(ActivationTableTest, DgeluMatchesSpecification) { check_table(1); } TEST(ActivationTableTest, DsiluMatchesSpecification) { - check_table(2); + check_table(exact_dsilu() ? 0 : 2); } TEST(ActivationTableTest, SizesIdentifyKind) { diff --git a/transformer_engine/common/activation/activation_table.cu b/transformer_engine/common/activation/activation_table.cu index b961a6ec832..63292605d6e 100644 --- a/transformer_engine/common/activation/activation_table.cu +++ b/transformer_engine/common/activation/activation_table.cu @@ -18,9 +18,8 @@ namespace { constexpr int kBlockSize = 256; -template +template __global__ void build_forward_kernel(__nv_bfloat16 *table) { - using T = Table; const uint32_t row = blockIdx.x * blockDim.x + threadIdx.x; if (row >= T::kRows) return; const uint32_t magnitude = T::kLowSentinel + row; @@ -38,9 +37,8 @@ __global__ void build_forward_kernel(__nv_bfloat16 *table) { table[2 * row + 1] = negative; } -template +template __global__ void build_derivative_kernel(float *table) { - using T = Table; const uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= T::kSignStride + T::kEntriesPerSign) return; if (i >= T::kEntriesPerSign && i < T::kSignStride) { @@ -52,15 +50,14 @@ __global__ void build_derivative_kernel(float *table) { table[i] = T::op(__uint_as_float(((negative ? 0x8000u : 0u) | magnitude) << 16)); } -template +template void build(void *table, cudaStream_t stream) { - using T = Table; if constexpr (T::kForward) { - build_forward_kernel<<(T::kRows, kBlockSize), kBlockSize, 0, stream>>>( + build_forward_kernel<<(T::kRows, kBlockSize), kBlockSize, 0, stream>>>( static_cast<__nv_bfloat16 *>(table)); } else { constexpr uint32_t entries = T::kSignStride + T::kEntriesPerSign; - build_derivative_kernel<<(entries, kBlockSize), kBlockSize, 0, stream>>>( + build_derivative_kernel<<(entries, kBlockSize), kBlockSize, 0, stream>>>( static_cast(table)); } NVTE_CHECK_CUDA(cudaGetLastError()); @@ -72,7 +69,7 @@ void build(void *table, cudaStream_t stream) { size_t nvte_activation_table_size(NVTEActivationTableKind kind) { NVTE_API_CALL(nvte_activation_table_size); - const size_t bytes = transformer_engine::activation_table::table_bytes(kind); + const size_t bytes = transformer_engine::activation_table::built_table_bytes(kind); NVTE_CHECK(bytes != 0, "Invalid NVTEActivationTableKind (got ", static_cast(kind), ")"); return bytes; } @@ -82,19 +79,28 @@ void nvte_build_activation_table(NVTEActivationTableKind kind, NVTETensor table, NVTE_API_CALL(nvte_build_activation_table); using namespace transformer_engine; using namespace transformer_engine::activation_table; - void *ptr = const_cast(checked_table_ptr(*convertNVTETensorCheck(table), kind)); + const TableRef ref = checked_table(*convertNVTETensorCheck(table), kind); + NVTE_CHECK((ref.use == kExactDsiluTable) == + (kind == kNVTEActivationTableDsilu && use_exact_dsilu_table()), + "Activation table size does not match NVTE_ACTIVATION_TABLE_EXACT_DSILU; allocate " + "nvte_activation_table_size(kind) bytes."); + void *ptr = const_cast(ref.ptr); switch (kind) { case kNVTEActivationTableGelu: - build(ptr, stream); + build>(ptr, stream); break; case kNVTEActivationTableSilu: - build(ptr, stream); + build>(ptr, stream); break; case kNVTEActivationTableDgelu: - build(ptr, stream); + build>(ptr, stream); break; case kNVTEActivationTableDsilu: - build(ptr, stream); + if (ref.use == kExactDsiluTable) { + build(ptr, stream); + } else { + build>(ptr, stream); + } break; default: NVTE_ERROR("Invalid NVTEActivationTableKind (got ", static_cast(kind), ")"); diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index b748d65f260..460e9a15907 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -38,18 +38,18 @@ namespace transformer_engine { namespace dispatch { -// Device pointer of the activation table to read for OP, or nullptr to compute OP. The table -// applies to BF16 inputs of the activations that have one. +// The activation table to read for OP, if any. Tables apply to BF16 inputs of the activations +// that have one. template -const void *activation_lut_for(const QuantizationConfig &quant_config, const DType input_dtype) { +activation_table::TableRef activation_lut_for(const QuantizationConfig &quant_config, + const DType input_dtype) { if constexpr (activation_table::kHasTable) { if (quant_config.activation_table != nullptr && input_dtype == DType::kBFloat16) { - return activation_table::checked_table_ptr( - *convertNVTETensorCheck(quant_config.activation_table), - activation_table::kKindOf); + return activation_table::checked_table(*convertNVTETensorCheck(quant_config.activation_table), + activation_table::kKindOf); } } - return nullptr; + return {}; } template @@ -109,13 +109,13 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const Tensor *dummy_input_tensor = nullptr; Tensor *dummy_dbias_tensor = nullptr; Tensor *dummy_workspace_tensor = nullptr; - const void *activation_lut = + const activation_table::TableRef activation_lut = IS_ACT ? activation_lut_for(quant_config_cpp, input_tensor->dtype()) - : nullptr; + : activation_table::TableRef{}; bool quantized_with_cutedsl = false; #ifdef NVTE_WITH_CUTEDSL // The CuTeDSL backend computes activations itself, so it cannot honor a table. - if (activation_lut == nullptr) { + if (activation_lut.ptr == nullptr) { quantized_with_cutedsl = cutedsl_backend::mxfp8_quantize_cutedsl( @@ -304,14 +304,14 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens break; } case NVTE_MXFP8_1D_SCALING: { - const void *activation_lut = nullptr; + activation_table::TableRef activation_lut; if constexpr (IS_DACT) { activation_lut = activation_lut_for(quant_config_cpp, input_tensor->dtype()); } bool quantized_with_cutedsl = false; #ifdef NVTE_WITH_CUTEDSL // The CuTeDSL backend computes activations itself, so it cannot honor a table. - if (activation_lut == nullptr) { + if (activation_lut.ptr == nullptr) { quantized_with_cutedsl = cutedsl_backend::mxfp8_quantize_cutedsl( diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 4d00bf62d7e..e36dcb152ae 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -50,7 +50,7 @@ template + int kActivationTableUse> __global__ void __launch_bounds__(THREADS_PER_CHUNK) quantize_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_act_input, @@ -268,11 +268,11 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float elt = static_cast(in_sh[shmem_offset_colwise]); if constexpr (IS_ACT) { - elt = activation_table::evaluate(elt, activation_lut); + elt = activation_table::evaluate(elt, activation_lut); } if constexpr (IS_DACT) { float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); - elt *= activation_table::evaluate(act_in_elt, + elt *= activation_table::evaluate(act_in_elt, activation_lut); } if constexpr (DBIAS_REDUCTION_IN_COLWISE) { @@ -359,11 +359,11 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float elt = static_cast(in_sh[shmem_offset_colwise]); if constexpr (IS_ACT) { - elt = activation_table::evaluate(elt, activation_lut); + elt = activation_table::evaluate(elt, activation_lut); } if constexpr (IS_DACT) { const float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); - elt *= activation_table::evaluate(act_in_elt, + elt *= activation_table::evaluate(act_in_elt, activation_lut); } partial_dbias_colwise += elt; @@ -469,11 +469,11 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float elt = static_cast(in.data.elt[e]); if constexpr (IS_ACT) { elt = - activation_table::evaluate(elt, activation_lut); + activation_table::evaluate(elt, activation_lut); } if constexpr (IS_DACT) { float act_in_elt = static_cast(act_in.data.elt[e]); - elt *= activation_table::evaluate(act_in_elt, + elt *= activation_table::evaluate(act_in_elt, activation_lut); } @@ -695,7 +695,7 @@ template void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, // TODO (ksivamani) Tensor *output, Tensor *dbias, Tensor *workspace, const bool use_2d_quantization, - const void *activation_lut, cudaStream_t stream) { + const activation_table::TableRef activation_lut, cudaStream_t stream) { using namespace quantize_kernel; checkCuDriverContext(stream); @@ -1019,12 +1019,12 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const auto launch = [&](auto rowwise, auto colwise) { TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { - const auto launch_kernel = [&](auto use_table) { + const auto launch_kernel = [&](auto table_use) { auto kernel = quantize_mxfp8_kernel< IS_DBIAS, IS_DACT, IS_ACT, ParamOP, OP, IType, OType, decltype(rowwise)::value, decltype(colwise)::value, WITH_GEMM_SWIZZLED_SCALES, CHUNK_DIM_Y, CHUNK_DIM_X, THREADS_PER_CHUNK, - kIs2DBlockScaling, decltype(use_table)::value>; + kIs2DBlockScaling, decltype(table_use)::value>; NVTE_CHECK_CUDA(cudaFuncSetAttribute( kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); @@ -1032,16 +1032,26 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr, num_rows, num_cols, - scale_stride_rowwise, scale_stride_colwise, activation_lut); + scale_stride_rowwise, scale_stride_colwise, activation_lut.ptr); }; + using activation_table::kDefaultTable; + using activation_table::kExactDsiluTable; + using activation_table::kNoTable; if constexpr (kTableCapable) { - if (activation_lut != nullptr) { - launch_kernel(std::true_type{}); + if (activation_lut.use == kDefaultTable) { + launch_kernel(std::integral_constant{}); + } else if (activation_lut.use == kExactDsiluTable) { + if constexpr (activation_table::kKindOf == + kNVTEActivationTableDsilu) { + launch_kernel(std::integral_constant{}); + } else { + NVTE_ERROR("Only dSiLU has an exact activation table."); + } } else { - launch_kernel(std::false_type{}); + launch_kernel(std::integral_constant{}); } } else { - launch_kernel(std::false_type{}); + launch_kernel(std::integral_constant{}); } }); }; diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index dc742665d20..aa25406689f 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -56,6 +56,8 @@ enum class NVTE_Activation_Type { * - GeLU and SiLU of -inf are -0 (the formulas give NaN). * - For |x| < 2^-12 (dGeLU) or 2^-9 (dSiLU), the derivative is its linear Taylor term * 0.5 + c*x rounded once, within 1 ulp (dGeLU) or 2 ulp (dSiLU) of util/math.h. + * With NVTE_ACTIVATION_TABLE_EXACT_DSILU=1 (experimental), the dSiLU table is 32 KB and + * matches util/math.h for every input. */ enum NVTEActivationTableKind { kNVTEActivationTableGelu = 0, /*!< GeLU (tanh form), BF16 entries */ @@ -67,7 +69,8 @@ enum NVTEActivationTableKind { /*! \brief Size in bytes of an activation table. * * \param[in] kind Activation the table holds. - * \return Size of the table in bytes. Each kind has a different size. + * \return Size of the table in bytes. Each kind has a different size. The dSiLU size + * depends on NVTE_ACTIVATION_TABLE_EXACT_DSILU. */ size_t nvte_activation_table_size(NVTEActivationTableKind kind); diff --git a/transformer_engine/common/util/activation_table.cuh b/transformer_engine/common/util/activation_table.cuh index e330b46fcf6..303e7df280e 100644 --- a/transformer_engine/common/util/activation_table.cuh +++ b/transformer_engine/common/util/activation_table.cuh @@ -20,7 +20,9 @@ * Derivatives (dGeLU, dSiLU): FP32 entries, because callers multiply the derivative by * the incoming gradient in FP32. Below the window the derivative is its linear Taylor * term 0.5 + c*x, rounded once; it differs from util/math.h by at most 1 ulp (dGeLU) or - * 2 ulp (dSiLU) there. Above the window the derivative has saturated, and the table's + * 2 ulp (dSiLU) there. With NVTE_ACTIVATION_TABLE_EXACT_DSILU=1 the dSiLU table extends + * down to 2^-25, below which the linear term equals util/math.h, so dSiLU matches + * util/math.h for every input (32 KB instead of 16 KB). Above the window the derivative has saturated, and the table's * last entry of the same sign holds the saturated value. Inputs where util/math.h * returns NaN (dGeLU from 0x6044 on, where 0.107*x*x overflows; +-inf; NaN) keep NaN. * @@ -39,6 +41,7 @@ #include "../common.h" #include "math.h" +#include "system.h" namespace transformer_engine { namespace activation_table { @@ -62,10 +65,9 @@ struct ForwardLayout { /*! \brief Derivative table over the BF16 magnitudes [kLo, kHi]. * * Positive inputs at entries [0, n), negative ones at [kSignStride, kSignStride + n). - * The stride turns the BF16 sign bit, 0x8000, into a byte offset of 0x8000 >> 2, which - * lets a packed lookup form both lanes' offsets in one 32-bit word. + * A power-of-two stride turns the BF16 sign bit into a byte offset with one shift. */ -template +template struct DerivativeLayout { static constexpr bool kForward = false; static constexpr uint32_t kLoBits = kLo; @@ -73,7 +75,7 @@ struct DerivativeLayout { //! Magnitudes from here on (and NaN) give NaN, as util/math.h does. static constexpr uint32_t kSaturatedEndBits = kSaturatedEnd; static constexpr uint32_t kEntriesPerSign = kHi - kLo + 1; - static constexpr uint32_t kSignStride = (0x8000 >> 2) / sizeof(float); + static constexpr uint32_t kSignStride = kStride; static_assert(kEntriesPerSign <= kSignStride, "The two signs' entries must not overlap."); static constexpr size_t kBytes = (kSignStride + kEntriesPerSign) * sizeof(float); }; @@ -110,6 +112,13 @@ struct Table : DerivativeLayout<0x3B00, 0x42B2, 0x7F8 __device__ static float below_window(const float x) { return fmaf(0.5f, x, 0.5f); } }; +//! dSiLU over [2^-25, 89]: below 2^-25 the linear term already equals util/math.h. Selected +//! with NVTE_ACTIVATION_TABLE_EXACT_DSILU=1. +struct ExactDsiluTable : DerivativeLayout<0x3300, 0x42B2, 0x7F80, 4096> { + __device__ static float op(const float x) { return dsilu(x, {}); } + __device__ static float below_window(const float x) { return fmaf(0.5f, x, 0.5f); } +}; + /*! \brief Table size in bytes. Each kind has a different size, which lets callers check that * a table matches the activation it is passed with. */ @@ -129,19 +138,54 @@ constexpr size_t table_bytes(const NVTEActivationTableKind kind) { } static_assert(table_bytes(kNVTEActivationTableGelu) != table_bytes(kNVTEActivationTableSilu) && - table_bytes(kNVTEActivationTableDgelu) != table_bytes(kNVTEActivationTableDsilu), + table_bytes(kNVTEActivationTableDgelu) != + table_bytes(kNVTEActivationTableDsilu) && + ExactDsiluTable::kBytes != table_bytes(kNVTEActivationTableDsilu), "Table sizes must identify their kind."); -/*! \brief Device pointer of an activation table, checked against the kind it must hold. */ -inline const void *checked_table_ptr(const Tensor &table, const NVTEActivationTableKind kind) { +//! Whether dSiLU tables are built in the exact layout (NVTE_ACTIVATION_TABLE_EXACT_DSILU=1). +inline bool use_exact_dsilu_table() { + static const bool exact = getenv("NVTE_ACTIVATION_TABLE_EXACT_DSILU"); + return exact; +} + +//! Size of the table that nvte_build_activation_table builds for kind. +inline size_t built_table_bytes(const NVTEActivationTableKind kind) { + if (kind == kNVTEActivationTableDsilu && use_exact_dsilu_table()) { + return ExactDsiluTable::kBytes; + } + return table_bytes(kind); +} + +//! Which table layout a kernel reads. +enum TableUse : int { + kNoTable = 0, //!< Compute OP with util/math.h. + kDefaultTable = 1, //!< Table. + kExactDsiluTable = 2, //!< ExactDsiluTable. +}; + +//! An activation table passed to a kernel. +struct TableRef { + const void *ptr = nullptr; + TableUse use = kNoTable; +}; + +/*! \brief An activation table, checked against the kind it must hold. + * + * A dSiLU table may have either layout; its size tells which. + */ +inline TableRef checked_table(const Tensor &table, const NVTEActivationTableKind kind) { const size_t bytes = table_bytes(kind); NVTE_CHECK(bytes != 0, "Invalid NVTEActivationTableKind (got ", static_cast(kind), ")"); NVTE_CHECK(table.data.dptr != nullptr, "Activation table is not allocated."); NVTE_CHECK(table.data.dtype == DType::kByte, "Activation table must be a byte tensor."); - NVTE_CHECK(table.data.numel() == bytes, "Activation table has ", table.data.numel(), - " bytes, but activation table kind ", static_cast(kind), " needs ", bytes, - " (was it built for another activation?)"); - return table.data.dptr; + const size_t got = table.data.numel(); + if (kind == kNVTEActivationTableDsilu && got == ExactDsiluTable::kBytes) { + return {table.data.dptr, kExactDsiluTable}; + } + NVTE_CHECK(got == bytes, "Activation table has ", got, " bytes, but activation table kind ", + static_cast(kind), " needs ", bytes, " (was it built for another activation?)"); + return {table.data.dptr, kDefaultTable}; } /*! \brief The table kind for OP, if OP is read from a table. */ @@ -170,9 +214,8 @@ constexpr NVTEActivationTableKind kKindOf> = kNVTEAct * Branch-free: the magnitude is clamped to the table rows with __hmax/__hmin, which * ignore NaN, so NaN lands on the -inf row; __hmin_nan/__hmax_nan then propagate it. */ -template +template __device__ __forceinline__ float forward(const float x, const __nv_bfloat16 *__restrict__ table) { - using T = Table; static_assert(T::kForward, "Not a forward table."); const __nv_bfloat16 a = __float2bfloat16_rn(x); const uint32_t sign = __float_as_uint(x) >> 31; @@ -189,9 +232,8 @@ __device__ __forceinline__ float forward(const float x, const __nv_bfloat16 *__r } /*! \brief Derivative of x, a BF16 value widened to FP32. Branch-free. */ -template +template __device__ __forceinline__ float derivative(const float x, const float *__restrict__ table) { - using T = Table; static_assert(!T::kForward, "Not a derivative table."); const uint32_t bits = __float_as_uint(x) >> 16; const uint32_t magnitude = bits & 0x7fffu; @@ -202,22 +244,25 @@ __device__ __forceinline__ float derivative(const float x, const float *__restri return magnitude >= T::kSaturatedEndBits ? __uint_as_float(0x7fffffffu) : value; } -/*! \brief OP(x), from the table when kUseTable, else from util/math.h. +/*! \brief OP(x), from the table given by kUse, or from util/math.h for kNoTable. * - * x must be a BF16 value widened to FP32 when kUseTable. + * x must be a BF16 value widened to FP32 when reading a table. */ -template +template __device__ __forceinline__ float evaluate(const float x, const void *table) { - if constexpr (kUseTable) { + if constexpr (kUse == kNoTable) { + return OP(x, {}); + } else { static_assert(kHasTable, "OP has no table."); constexpr NVTEActivationTableKind kKind = kKindOf; - if constexpr (Table::kForward) { - return forward(x, static_cast(table)); + if constexpr (kUse == kExactDsiluTable) { + static_assert(kKind == kNVTEActivationTableDsilu, "Only dSiLU has an exact table."); + return derivative(x, static_cast(table)); + } else if constexpr (Table::kForward) { + return forward>(x, static_cast(table)); } else { - return derivative(x, static_cast(table)); + return derivative>(x, static_cast(table)); } - } else { - return OP(x, {}); } } From 0bd0b675286793f0ac5ccf8d3e28e92245b0db9e Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 14 Oct 2026 11:51:45 -0700 Subject: [PATCH 03/29] [Common] Add a register-tiled MXFP8 activation kernel that reads the tables Port the register-tiled MXFP8 cast kernel from #3617 onto the shared BF16 activation tables. The kernel copies the table into shared memory with a bulk copy, evaluates the forward activations two lanes at a time with the branch-free lookup, and evaluates the derivatives from the same tables as the generic kernel, so both kernels produce bit-identical output. The kernel runs only when a table is passed. Pad table sizes to 16 bytes for the bulk copy, and add the packed bf16x2/fp32x2 PTX helpers it needs. Tests compare the register-tiled kernel against the generic kernel with a table, byte for byte. Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 95 +++ .../common/activation/activation_table.cu | 2 + .../common/cast/mxfp8/quantize_mxfp8.cuh | 17 + .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 639 ++++++++++++++++++ .../common/util/activation_table.cuh | 89 ++- transformer_engine/common/util/ptx.cuh | 52 ++ .../common/util/ptx_arch_spec.cuh | 14 + 7 files changed, 906 insertions(+), 2 deletions(-) create mode 100644 transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index fa6a96aa729..348a23fb321 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -407,6 +407,65 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t } } +// With a table, the register-resident kernel and the generic kernel must write the same bytes. +// The first call takes the default dispatch, which picks the register-resident kernel for these +// shapes; a noop tensor holding 0 keeps the second call on the generic kernel without skipping it. +void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, const size_t cols, + const bool swizzled) { + const std::vector shape{rows, cols}; + const size_t n = rows * cols; + Tensor x("x", shape, DType::kBFloat16); + Tensor grad("grad", shape, DType::kBFloat16); + fill_inputs(fill, &x, &grad); + Tensor noop("noop", std::vector{1}, DType::kFloat32); + + Tensor outs[2] = { + Tensor("default", shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING), + Tensor("generic", shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING)}; + for (int i = 0; i < 2; ++i) { + outs[i].set_with_gemm_swizzled_scales(swizzled); + QuantizationConfigWrapper config; + config.set_activation_table(activation_table(table_kind(mode))); + if (i == 1) config.set_noop_tensor(noop.data()); + switch (mode) { + case Mode::GeLU: + nvte_gelu_v2(x.data(), outs[i].data(), config, 0); + break; + case Mode::SiLU: + nvte_silu_v2(x.data(), outs[i].data(), config, 0); + break; + case Mode::DGeLU: + nvte_dgelu_v2(grad.data(), x.data(), outs[i].data(), config, 0); + break; + default: + nvte_dsilu_v2(grad.data(), x.data(), outs[i].data(), config, 0); + break; + } + } + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + for (auto &out : outs) out.to_cpu(); + EXPECT_EQ(count_byte_mismatches( + reinterpret_cast(outs[0].rowwise_cpu_dptr()), + reinterpret_cast(outs[1].rowwise_cpu_dptr()), n), + 0u) + << "rowwise data"; + EXPECT_EQ(count_byte_mismatches( + reinterpret_cast(outs[0].columnwise_cpu_dptr()), + reinterpret_cast(outs[1].columnwise_cpu_dptr()), n), + 0u) + << "colwise data"; + EXPECT_EQ(count_byte_mismatches(outs[0].rowwise_cpu_scale_inv_ptr(), + outs[1].rowwise_cpu_scale_inv_ptr(), + product(outs[0].rowwise_scale_inv_shape())), + 0u) + << "rowwise scales"; + EXPECT_EQ(count_byte_mismatches(outs[0].columnwise_cpu_scale_inv_ptr(), + outs[1].columnwise_cpu_scale_inv_ptr(), + product(outs[0].columnwise_scale_inv_shape())), + 0u) + << "colwise scales"; +} + using Params = std::tuple, bool /*use_table*/>; @@ -475,3 +534,39 @@ INSTANTIATE_TEST_SUITE_P(TinyGradients, CastMXFP8ActExactTestSuite, ::testing::Values(std::pair{64, 128}), ::testing::Bool()), case_name); + +namespace { + +using RegtileParams = std::tuple>; + +class CastMXFP8ActRegtileTestSuite : public ::testing::TestWithParam {}; + +std::string regtile_case_name(const testing::TestParamInfo &info) { + const char *fills[] = {"Typical", "OrderedInputs", "Specials", "TinyGradients"}; + const auto [mode, fill, swizzled, dims] = info.param; + return std::string(mode_name(mode)) + "X" + fills[static_cast(fill)] + + (swizzled ? "XSwizzled" : "XCompact") + "X" + std::to_string(dims.first) + "X" + + std::to_string(dims.second); +} + +} // namespace + +TEST_P(CastMXFP8ActRegtileTestSuite, MatchesGenericKernel) { + cudaDeviceProp prop; + ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); + if (prop.major < 10) { + GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; + } + const auto [mode, fill, swizzled, dims] = GetParam(); + run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, CastMXFP8ActRegtileTestSuite, + ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::SiLU, Mode::DSiLU), + ::testing::Values(Fill::Typical, Fill::OrderedInputs, Fill::Specials), + ::testing::Bool(), + ::testing::Values(std::pair{1024, 2048}, + std::pair{4160, 2048}, + std::pair{128, 256})), + regtile_case_name); diff --git a/transformer_engine/common/activation/activation_table.cu b/transformer_engine/common/activation/activation_table.cu index 63292605d6e..9402e4be2fb 100644 --- a/transformer_engine/common/activation/activation_table.cu +++ b/transformer_engine/common/activation/activation_table.cu @@ -52,6 +52,8 @@ __global__ void build_derivative_kernel(float *table) { template void build(void *table, cudaStream_t stream) { + // Zero the padding up to the 16-byte size, so the table's bytes are deterministic. + NVTE_CHECK_CUDA(cudaMemsetAsync(table, 0, T::kBytes, stream)); if constexpr (T::kForward) { build_forward_kernel<<(T::kRows, kBlockSize), kBlockSize, 0, stream>>>( static_cast<__nv_bfloat16 *>(table)); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index e36dcb152ae..c30e40d508f 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -23,6 +23,7 @@ #include "../../util/ptx_arch_spec.cuh" #include "../../utils.cuh" #include "../core/common.cuh" +#include "regtile/quantize_mxfp8_regtile.cuh" #include "specialized/cast_bidim.cuh" #include "specialized/cast_rowwise.cuh" #include "specialized/quantize_mxfp8.cuh" @@ -964,6 +965,22 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, return; } + // regtile::quantize replaces the generic kernel below for the requests it + // implements: bidimensional BF16 -> E4M3 with fused GeLU, SiLU, dGeLU or dSiLU read + // from an activation table, no dbias, rows % 64 == 0, cols % 256 == 0, 16-byte + // aligned inputs (see regtile::kImplements and regtile::can_use). It keeps the tile + // in registers instead of staging it through shared memory and writes the same + // bytes the generic kernel writes with the same table. + if constexpr (std::is_same_v && std::is_same_v && + regtile::kImplements) { + if (regtile::can_use(input, act_input, *output, *noop, + use_2d_quantization, activation_lut)) { + regtile::quantize( + input, act_input, output, activation_lut, stream); + return; + } + } + alignas(64) CUtensorMap tensor_map_input{}; alignas(64) CUtensorMap tensor_map_act_input{}; alignas(64) CUtensorMap tensor_map_output_rowwise{}; diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh new file mode 100644 index 00000000000..176821472e0 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -0,0 +1,639 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_mxfp8_regtile.cuh + * \brief Register-resident MXFP8 bidimensional quantize fused with GeLU or SiLU, or + * their derivatives. + * + * Rowwise and colwise MX scales need an amax over the same data in two + * directions. The generic kernel stages each tile through shared memory and + * walks it twice. Here a CTA covers whole 32-row bands, the colwise block + * height, so the tile is read once into registers and both directions are + * quantized from there; shared memory holds only the cross-warp colwise + * partials, the row scales and the activation table. + * + * The activation or derivative is read from the table the caller passes (see + * util/activation_table.cuh), copied into shared memory once per CTA, so this kernel + * produces exactly the bytes the generic kernel produces with the same table. + * + * It replaces the generic kernel, not the specialized ones: cast-only requests + * are served before the generic dispatch by the specialized kernels, one of + * which (specialized/cast_bidim.cuh) is also register-resident. Anything this + * kernel does not implement (see can_use) takes the generic kernel. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_MXFP8_REGTILE_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_MXFP8_REGTILE_CUH_ + +#include + +#include +#include +#include +#include + +#include "../../../common.h" +#include "../../../util/activation_table.cuh" +#include "../../../util/cuda_runtime.h" +#include "../../../util/math.h" +#include "../../../util/ptx.cuh" +#include "../../../util/ptx_arch_spec.cuh" +#include "../../../utils.cuh" +#include "../swizzle.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace regtile { + +namespace ptx = transformer_engine::ptx; + +namespace { + +// The datapath works on 32-bit words, each holding a BF16 column pair. + +/*! \brief Elementwise max(|a|, |b|) on a packed BF16 pair, as a raw word. */ +__device__ __forceinline__ unsigned abs_max_bf16x2(unsigned a, unsigned b) { + ptx::bf16x2 d; + ptx::abs_max_2x(d, reinterpret_cast(a), + reinterpret_cast(b)); + return reinterpret_cast(d); +} + +/*! \brief Widen a BF16 pair word to FP32, low half to x. Same result as ptx::up_cast, whose + * volatile asm the compiler cannot schedule as freely in the derivative loop. */ +__device__ __forceinline__ ptx::floatx2 widen_bf16x2(const unsigned v) { + return {__uint_as_float(v << 16), __uint_as_float(v & 0xffff0000u)}; +} + +/*! \brief Round an FP32 pair to a BF16 pair word, low lane in the low half. */ +__device__ __forceinline__ unsigned to_bf16x2(const ptx::floatx2& v) { + const __nv_bfloat162 r = __floats2bfloat162_rn(v.x, v.y); + return reinterpret_cast(r); +} + +/*! \brief The reciprocal of an E8M0 scale, 2^(127-e), broadcast to a BF16 pair. */ +__device__ __forceinline__ unsigned e8m0_to_bf16x2_reciprocal(unsigned e) { + const ptx::bf16x2 r = ptx::exp2f_rcp_2x(static_cast(e)); + return reinterpret_cast(r); +} + +/*! \brief E8M0 scales for two block amaxes, packed as ptx::float_to_e8m0_2x. + * + * The generic kernel's amax is a 0-seeded max that ignores NaN, so a block of + * only NaNs gets amax 0; the fmaxf reproduces that for the data-seeded maxima + * here. The conversion saturates, so the scale never reaches 255. + */ +__device__ __forceinline__ unsigned amax_to_e8m0_2x(float amax_hi, float amax_lo) { + constexpr float kMaxNormRcp = Quantized_Limits::max_norm_rcp; + return ptx::float_to_e8m0_2x(fmaxf(amax_hi * kMaxNormRcp, 0.0f), + fmaxf(amax_lo * kMaxNormRcp, 0.0f)); +} + +/*! \brief The derivative of a BF16 pair, from the table in shared memory. */ +template +__device__ __forceinline__ ptx::floatx2 dact_word(unsigned x, + const unsigned char* __restrict__ table) { + const float2 d = activation_table::derivative_2x(x, table); + return {d.x, d.y}; +} + +/*! \brief Start copying a table of kBytes into shared memory with one bulk copy. + * + * Called by every thread of the CTA; contains a __syncthreads. + */ +template +__device__ __forceinline__ void load_table_async(unsigned char* smem, uint64_t* barrier, + const void* table, const int tid) { + static_assert(kBytes % 16 == 0, "A bulk copy moves whole 16-byte units."); + if (tid == 0) { + ptx::mbarrier_init(barrier, 1); + ptx::fence_proxy_async_shared_cta(); + } + __syncthreads(); + if (tid == 0) { + ptx::mbarrier_arrive_expect_tx(barrier, kBytes); + ptx::cp_async_bulk_tensor_1d_global_to_shared(reinterpret_cast(smem), + reinterpret_cast(table), kBytes, + barrier); + } +} + +// Tile geometry. A CTA covers 256 columns and walks 32-row sub-tiles; 32 rows +// is the colwise MX block height, so the colwise reduction closes in the CTA. + +//! Columns a CTA covers. +constexpr int kTileCols = 256; +//! Elements in one MX block, and equivalently the colwise block height. +constexpr int kRowsPerMxBlock = 32; +//! MX groups (32 columns each) across the tile. +constexpr int kMxGroupsPerRow = kTileCols / 32; +//! A register-resident tile spreads its 256 columns over the 32 lanes of a warp. +constexpr int kColsPerLane = kTileCols / 32; +//! ...which a lane holds as BF16 pairs, one 32-bit word each. +constexpr int kWordsPerLane = kColsPerLane / 2; +//! ...read as 128-bit vectors, so one per row. +constexpr int kVecLoadsPerRow = kWordsPerLane / 4; +//! Lanes that have to cooperate to cover one 32-column MX group. +constexpr int kLanesPerMxGroup = 32 / kColsPerLane; + +/*! \brief Scale two BF16 pairs, each by its own reciprocal scale pair, to four E4M3. */ +__device__ __forceinline__ unsigned quantize_e4m3x4(unsigned in01, unsigned in23, unsigned scale01, + unsigned scale23) { + ptx::fp8e4m3x4 out; + ptx::mul_cvt_4x(out, reinterpret_cast(in01), + reinterpret_cast(scale01), + reinterpret_cast(in23), + reinterpret_cast(scale23)); + return reinterpret_cast(out); +} + +/*! \brief Store one lane's quantized row: kColsPerLane FP8 bytes, one STG.64. */ +__device__ __forceinline__ void store_quantized(unsigned char* p, const unsigned* v) { + static_assert(kColsPerLane == 8, "One STG.64 covers exactly eight FP8 bytes."); + *(uint2*)p = *(const uint2*)v; +} + +// Minimum resident CTAs per SM, the second __launch_bounds__ argument, capped +// at what the SM can hold: sm_75 and sm_107 have half the threads per SM of sm_100. +constexpr int kMinBlocksAct = 8; +constexpr int kMinBlocksDact = 6; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 750 || __CUDA_ARCH__ == 1070) +constexpr int kMaxThreadsPerSm = 1024; +#else +constexpr int kMaxThreadsPerSm = 2048; +#endif + +template +struct TileConfig { + static_assert(IS_ACT != IS_DACT, "Exactly one of the activation and its derivative is fused."); + static constexpr int kThreadsPerCta = IS_DACT ? 256 : 128; + static constexpr int kWarpsPerCta = kThreadsPerCta / 32; + static constexpr int kRowsPerWarp = 32 / kWarpsPerCta; + // The rowwise half is quantized before the cross-warp barrier and the row + // scales are drained inside the colwise barrier interval, saving one barrier + // per 32-row tile. A tile too narrow to leave idle threads hands the drain to + // the threads that just folded the colwise partials. + static constexpr int kDrainTidBase = + (kThreadsPerCta >= kTileCols / 2 + kTileCols / 4) ? kTileCols / 2 : 0; + static constexpr int kRowsInFlight = 2; + static constexpr int kRequestedMinBlocks = IS_DACT ? kMinBlocksDact : kMinBlocksAct; + static constexpr int kMinBlocksPerSm = kRequestedMinBlocks < kMaxThreadsPerSm / kThreadsPerCta + ? kRequestedMinBlocks + : kMaxThreadsPerSm / kThreadsPerCta; + static constexpr int kTableBytes = static_cast(TableT::kBytes); +}; + +// Register-resident tiling: the 32x256 tile stays in registers; only the +// colwise partials go through shared memory. +template +__device__ __forceinline__ void quantize_regtile( + const uint4* __restrict__ input, const uint4* __restrict__ act_input, + unsigned char* __restrict__ rowwise_out, unsigned char* __restrict__ rowwise_scales, + unsigned char* __restrict__ colwise_out, unsigned char* __restrict__ colwise_scales, int K, + int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* table) { + using C = TileConfig; + constexpr int kWarpsPerCta = C::kWarpsPerCta, kRowsPerWarp = C::kRowsPerWarp, + kRowsInFlight = C::kRowsInFlight; + + __shared__ __align__(16) unsigned col_amax_partials[kWarpsPerCta][kTileCols / 2]; + __shared__ __align__(16) unsigned col_scale_rcp_smem[kTileCols / 2]; + __shared__ __align__(8) unsigned char row_scale_bytes[kRowsPerMxBlock * kMxGroupsPerRow]; + __shared__ __align__(16) unsigned char table_smem[C::kTableBytes]; + __shared__ __align__(8) uint64_t table_barrier; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int col0 = blockIdx.x * kTileCols; + const int vecs_per_row = K >> 3; // uint4 (8 bf16 values) per row + + // The table copy overlaps the first input loads; see the wait below. + load_table_async(table_smem, &table_barrier, table, tid); + +#pragma unroll 1 + for (int it = 0; it < iters; ++it) { + // Row blocks are taken grid-strided, so the CTAs on the same step cover one + // contiguous row band. + const int row0 = (it * gridDim.y + blockIdx.y) * kRowsPerMxBlock; + // This thread's kRowsPerWarp x kColsPerLane sub-tile, as BF16 pairs. + unsigned tile_words[kRowsPerWarp * kWordsPerLane]; + const size_t load_offset = (size_t)(row0 + warp * kRowsPerWarp) * vecs_per_row + (col0 >> 3) + + (size_t)kVecLoadsPerRow * lane; + const uint4* input_ptr = input + load_offset; + const uint4* act_input_ptr = act_input + load_offset; + + // ---- load + activate, kRowsInFlight rows in flight at a time -------------------- +#pragma unroll + for (int h = 0; h < kRowsPerWarp / kRowsInFlight; ++h) { + uint4 x_vec[kRowsInFlight][kVecLoadsPerRow], grad_vec[kRowsInFlight][kVecLoadsPerRow]; + asm volatile("" ::: "memory"); +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) +#pragma unroll + for (int n = 0; n < kVecLoadsPerRow; ++n) + x_vec[t][n] = act_input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n]; + if constexpr (IS_DACT) { +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) +#pragma unroll + for (int n = 0; n < kVecLoadsPerRow; ++n) + grad_vec[t][n] = input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n]; + } + // Keep the load group's loads together, between the two empty asm statements: issued + // back to back, their latencies overlap. Otherwise register pressure leads the compiler + // to spread them between the previous group's table lookups, each then stalling on its + // own. + asm volatile("" ::: "memory"); + if (it == 0 && h == 0) ptx::mbarrier_wait_parity(&table_barrier, 0); + // Convert the load group into the BF16 pairs the epilogue quantizes. + if constexpr (!IS_DACT) { +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) { +#pragma unroll + for (int n = 0; n < kVecLoadsPerRow; ++n) { + const unsigned* x_words = (const unsigned*)&x_vec[t][n]; + const int tile_word_base = (h * kRowsInFlight + t) * kWordsPerLane + n * 4; +#pragma unroll + for (int m = 0; m < 4; ++m) { + tile_words[tile_word_base + m] = activation_table::forward_2x( + x_words[m], reinterpret_cast(table_smem)); + } + } + } + } else { +#define MXFP8_DACT_WORD(T, M) \ + { \ + const ptx::floatx2 product = \ + ptx::mul_2x(widen_bf16x2(grad_words[M]), dact_word(x_words[M], table_smem)); \ + tile_words[(h * kRowsInFlight + (T)) * kWordsPerLane + (M)] = to_bf16x2(product); \ + } +#define MXFP8_DACT_ROW(T) \ + { \ + const unsigned* x_words = (const unsigned*)&x_vec[T][0]; \ + const unsigned* grad_words = (const unsigned*)&grad_vec[T][0]; \ + MXFP8_DACT_WORD(T, 0) \ + MXFP8_DACT_WORD(T, 1) \ + MXFP8_DACT_WORD(T, 2) \ + MXFP8_DACT_WORD(T, 3) \ + } + static_assert(kRowsInFlight == 2, "A load group is two rows."); + MXFP8_DACT_ROW(0) + MXFP8_DACT_ROW(1) +#undef MXFP8_DACT_ROW +#undef MXFP8_DACT_WORD + } + } + + // The rowwise half does not depend on the colwise scales, so it runs before + // the cross-warp barrier (see TileConfig). + const size_t store_offset = + (size_t)(row0 + warp * kRowsPerWarp) * K + col0 + kColsPerLane * lane; + auto row_scale = [&](const unsigned* row, int j) { + unsigned amax = + abs_max_bf16x2(abs_max_bf16x2(row[0], row[1]), abs_max_bf16x2(row[2], row[3])); +#pragma unroll + for (int m = 4; m < kWordsPerLane; m += 2) + amax = abs_max_bf16x2(amax, abs_max_bf16x2(row[m], row[m + 1])); + // A 32-value rowwise block spans kLanesPerMxGroup lanes. +#pragma unroll + for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) + amax = abs_max_bf16x2(amax, __shfl_xor_sync(0xffffffffu, amax, lane_mask)); + // Fold the two halves with one PRMT and a packed max, which leaves the + // magnitude in both halves; the high half is then the FP32 bit pattern. + amax = abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); + const unsigned scale_byte = amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); + if ((lane & (kLanesPerMxGroup - 1)) == 0) + row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + (lane / kLanesPerMxGroup)] = + (unsigned char)scale_byte; + return e8m0_to_bf16x2_reciprocal(scale_byte); + }; + // On the 128-thread activation tile, two rows share one butterfly and conversion. + auto row_amax = [&](const unsigned* row) { + unsigned amax = + abs_max_bf16x2(abs_max_bf16x2(row[0], row[1]), abs_max_bf16x2(row[2], row[3])); +#pragma unroll + for (int m = 4; m < kWordsPerLane; m += 2) + amax = abs_max_bf16x2(amax, abs_max_bf16x2(row[m], row[m + 1])); + return abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); + }; + auto row_scale_pair = [&](const unsigned* row_a, const unsigned* row_b, int j, + unsigned& row_scale_rcp0, unsigned& row_scale_rcp1) { + unsigned amax_pair = __byte_perm(row_amax(row_a), row_amax(row_b), 0x5410) & 0x7fff7fffu; +#pragma unroll + for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) + amax_pair = abs_max_bf16x2(amax_pair, __shfl_xor_sync(0xffffffffu, amax_pair, lane_mask)); + const ptx::floatx2 amax_2x = widen_bf16x2(amax_pair); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + const unsigned scale_byte0 = scale_pair & 0xffu, scale_byte1 = (scale_pair >> 8) & 0xffu; + if ((lane & (kLanesPerMxGroup - 1)) == 0) { + unsigned char* scale_dst = row_scale_bytes + (warp * kRowsPerWarp + j) * kMxGroupsPerRow + + (lane / kLanesPerMxGroup); + scale_dst[0] = (unsigned char)scale_byte0; + scale_dst[kMxGroupsPerRow] = (unsigned char)scale_byte1; + } + row_scale_rcp0 = e8m0_to_bf16x2_reciprocal(scale_byte0); + row_scale_rcp1 = e8m0_to_bf16x2_reciprocal(scale_byte1); + }; + // Colwise amax: per-warp partial in registers, cross-warp fold in shared. + auto publish_col_amax_partials = [&]() { + unsigned col_amax[kWordsPerLane]; +#pragma unroll + for (int m = 0; m < kWordsPerLane; ++m) col_amax[m] = tile_words[m]; +#pragma unroll + for (int j = 1; j < kRowsPerWarp; ++j) +#pragma unroll + for (int m = 0; m < kWordsPerLane; ++m) + col_amax[m] = abs_max_bf16x2(col_amax[m], tile_words[j * kWordsPerLane + m]); +#pragma unroll + for (int n = 0; n < kVecLoadsPerRow; ++n) + *(uint4*)(&col_amax_partials[warp][kWordsPerLane * lane + 4 * n]) = + *(const uint4*)(col_amax + 4 * n); + }; + if constexpr (C::kThreadsPerCta == 128) { + unsigned char* row_out_ptr = rowwise_out + store_offset; +#pragma unroll + for (int j = 0; j < kRowsPerWarp; j += 2) { + const unsigned* row_a = tile_words + j * kWordsPerLane; + const unsigned* row_b = row_a + kWordsPerLane; + unsigned row_scale_rcp0, row_scale_rcp1; + row_scale_pair(row_a, row_b, j, row_scale_rcp0, row_scale_rcp1); + unsigned row_out[kWordsPerLane / 2]; +#pragma unroll + for (int m = 0; m < kWordsPerLane / 2; ++m) + row_out[m] = + quantize_e4m3x4(row_a[2 * m], row_a[2 * m + 1], row_scale_rcp0, row_scale_rcp0); + store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); +#pragma unroll + for (int m = 0; m < kWordsPerLane / 2; ++m) + row_out[m] = + quantize_e4m3x4(row_b[2 * m], row_b[2 * m + 1], row_scale_rcp1, row_scale_rcp1); + store_quantized(row_out_ptr + (unsigned)(j + 1) * (unsigned)K, row_out); + } + } else { + unsigned char* row_out_ptr = rowwise_out + store_offset; +#pragma unroll + for (int j = 0; j < kRowsPerWarp; ++j) { + const unsigned* v = tile_words + j * kWordsPerLane; + const unsigned row_scale_rcp = row_scale(v, j); + unsigned row_out[kWordsPerLane / 2]; +#pragma unroll + for (int m = 0; m < kWordsPerLane / 2; ++m) + row_out[m] = quantize_e4m3x4(v[2 * m], v[2 * m + 1], row_scale_rcp, row_scale_rcp); + store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); + } + } + + publish_col_amax_partials(); + __syncthreads(); + if (tid < kTileCols / 2) { + unsigned col_amax = col_amax_partials[0][tid]; +#pragma unroll + for (int w = 1; w < kWarpsPerCta; ++w) + col_amax = abs_max_bf16x2(col_amax, col_amax_partials[w][tid]); + col_amax &= 0x7fff7fffu; + const ptx::floatx2 amax_2x = widen_bf16x2(col_amax); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + // The colwise swizzle is the rowwise one with rows and columns swapped, + // so the two columns of this pair land in different 16-byte rows. The + // grid covers M exactly: gridDim.y * iters row blocks of 32. + const size_t row_block_tiles = + DIVUP(static_cast(gridDim.y) * iters * kRowsPerMxBlock, size_t{128}); + const size_t col = col0 + 2 * tid; + colwise_scales[swizzle::gemm_swizzled_scale_idx(col, row0 / 32, row_block_tiles)] = + static_cast(scale_pair); + colwise_scales[swizzle::gemm_swizzled_scale_idx(col + 1, row0 / 32, row_block_tiles)] = + static_cast(scale_pair >> 8); + } else { + *(unsigned short*)(colwise_scales + (size_t)(row0 / 32) * colwise_scale_stride + col0 + + 2 * tid) = (unsigned short)scale_pair; + } + const ptx::bf16x2 scale_rcp = ptx::exp2f_rcp_2x_per_lane(__byte_perm(scale_pair, 0, 0x4140)); + col_scale_rcp_smem[tid] = reinterpret_cast(scale_rcp); + } + auto drain_row_scales = [&]() { + constexpr int drain_tid_base = C::kDrainTidBase; + // A row's eight group scales are eight contiguous, 8-byte-aligned bytes of + // rowwise_scales, so one warp drains a 32-row tile with one STG.64 per lane. + // In the GEMM-swizzled layout they are two 4-byte runs in adjacent tiles. + if (tid >= drain_tid_base && tid < drain_tid_base + kRowsPerMxBlock) { + const int r = tid - drain_tid_base; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + constexpr size_t kTileBytes = + swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_X * swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_Y; + const size_t idx = swizzle::gemm_swizzled_scale_idx(row0 + r, col0 >> 5, K / 128); + const uint2 bytes = *(const uint2*)(row_scale_bytes + r * kMxGroupsPerRow); + *(unsigned*)(rowwise_scales + idx) = bytes.x; + *(unsigned*)(rowwise_scales + idx + kTileBytes) = bytes.y; + } else { + *(uint2*)(rowwise_scales + (size_t)(row0 + r) * rowwise_scale_stride + (col0 >> 5)) = + *(const uint2*)(row_scale_bytes + r * kMxGroupsPerRow); + } + } + }; + drain_row_scales(); + __syncthreads(); + + // ---- columnwise quantization -------------------------------------------- + { + uint4 col_scale_rcp_vec[kVecLoadsPerRow]; +#pragma unroll + for (int n = 0; n < kVecLoadsPerRow; ++n) + col_scale_rcp_vec[n] = *(const uint4*)(col_scale_rcp_smem + kWordsPerLane * lane + 4 * n); + const unsigned* col_scale_rcp = (const unsigned*)col_scale_rcp_vec; + unsigned char* col_out_ptr = colwise_out + store_offset; +#pragma unroll + for (int j = 0; j < kRowsPerWarp; ++j) { + const unsigned* v = tile_words + j * kWordsPerLane; + unsigned col_out[kWordsPerLane / 2]; +#pragma unroll + for (int m = 0; m < kWordsPerLane / 2; ++m) + col_out[m] = quantize_e4m3x4(v[2 * m], v[2 * m + 1], col_scale_rcp[2 * m], + col_scale_rcp[2 * m + 1]); + store_quantized(col_out_ptr + (unsigned)j * (unsigned)K, col_out); + } + } + } +} + +template +__global__ void __launch_bounds__(TileConfig::kThreadsPerCta, + TileConfig::kMinBlocksPerSm) + quantize_mxfp8_kernel(const unsigned* __restrict__ input, + const unsigned* __restrict__ act_input, + unsigned char* __restrict__ rowwise_out, + unsigned char* __restrict__ rowwise_scales, + unsigned char* __restrict__ colwise_out, + unsigned char* __restrict__ colwise_scales, int K, + int rowwise_scale_stride, int colwise_scale_stride, int iters, + const void* __restrict__ table) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + quantize_regtile( + (const uint4*)input, (const uint4*)act_input, rowwise_out, rowwise_scales, colwise_out, + colwise_scales, K, rowwise_scale_stride, colwise_scale_stride, iters, table); +#endif +} + +// Row blocks per CTA. A longer walk amortizes the per-CTA table copy; a +// shorter one keeps the grid deep enough to avoid a residency tail. +//! Grid size, in CTAs, that the walk length aims for. +constexpr long long kCtaTargetActivation = 16384; +//! Minimum walk for the instantiations that copy the table. +constexpr int kMinWalk = 2; +//! Maximum walk for the derivatives. +constexpr int kMaxWalkDact = 64; + +static int pick_walk_length(int row_blocks, int grid_cols, long long target, + int min_walk = kMinWalk) { + int walk = 1; + while ((row_blocks / walk) % 2 == 0 && (long long)grid_cols * (row_blocks / walk) > target) + walk *= 2; + if (walk < min_walk && row_blocks % min_walk == 0) walk = min_walk; + return walk; +} + +// Shared-memory carveout, in percent, for the activations, whose footprint leaves room: +// the rest of the unified array goes to L1. +constexpr int kCarveoutPercentAct = 40; + +template +static void set_carveout() { + constexpr int kPercent = IS_ACT ? kCarveoutPercentAct : 0; + if constexpr (kPercent > 0) { + // Function attributes are per device, like the tables. + static std::mutex mutex; + static std::vector done; + int device; + NVTE_CHECK_CUDA(cudaGetDevice(&device)); + std::lock_guard lock(mutex); + if (static_cast(device) >= done.size()) done.resize(device + 1, false); + if (!done[device]) { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + (const void*)quantize_mxfp8_kernel, + cudaFuncAttributePreferredSharedMemoryCarveout, kPercent)); + done[device] = true; + } + } +} + +// Row blocks per CTA walk: M / 32, walked in powers of two. +template +int walk_length(int M, int K) { + const int row_blocks = M / kRowsPerMxBlock; + const int grid_cols = K / kTileCols; + // The minimum walk amortizes the table copy, but not at the price of leaving + // SMs without a CTA. + int min_walk = kMinWalk; + if ((long long)grid_cols * (row_blocks / min_walk) < 2LL * cuda::sm_count()) min_walk = 1; + int iters = pick_walk_length(row_blocks, grid_cols, kCtaTargetActivation, min_walk); + if constexpr (IS_DACT) { + if (iters > kMaxWalkDact) iters = kMaxWalkDact; + } + return iters; +} + +template +static void launch(const void* input, const void* act_input, void* rowwise_out, + void* rowwise_scales, void* colwise_out, void* colwise_scales, int M, int K, + int rowwise_scale_stride, int colwise_scale_stride, const void* table, + cudaStream_t stream) { + using C = TileConfig; + set_carveout(); + const int iters = walk_length(M, K); + dim3 grid(K / kTileCols, M / kRowsPerMxBlock / iters); + quantize_mxfp8_kernel + <<>>( + (const unsigned*)input, (const unsigned*)act_input, (unsigned char*)rowwise_out, + (unsigned char*)rowwise_scales, (unsigned char*)colwise_out, + (unsigned char*)colwise_scales, K, rowwise_scale_stride, colwise_scale_stride, iters, + table); +} + +// The grid is derived by exact division, so columns must tile into 256. Rows +// must be a multiple of the generic kernel's 64-row activation tile: on a +// 32-row tail that kernel writes zero scales into the padding past the last +// row, which this one does not. Output offsets within a 32-row band are 32-bit. +constexpr size_t kGenericActivationTileRows = 64; + +inline bool shape_supported(size_t rows, size_t cols) { + constexpr size_t kMaxRows = std::numeric_limits::max(); + constexpr size_t kMaxCols = std::numeric_limits::max() / kRowsPerMxBlock; + return rows > 0 && cols > 0 && rows <= kMaxRows && cols <= kMaxCols && + rows % kGenericActivationTileRows == 0 && cols % kTileCols == 0; +} + +} // namespace + +//! Whether this kernel implements the request type: GeLU or SiLU, or the derivative of +//! either, without dbias. It reads the activation from a table. +template +constexpr bool kImplements = + !IS_DBIAS && (IS_ACT != IS_DACT) && activation_table::kHasTable && + (IS_ACT == activation_table::Table>::kForward); + +// Whether this kernel takes a request of a type in kImplements, with BF16 input, E4M3 +// output and an activation table. Anything else takes the generic kernel. +template +bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, const Tensor& noop, + bool use_2d_quantization, const activation_table::TableRef& table) { + if (table.use == activation_table::kNoTable) return false; + const auto [rows, cols] = input.flat_2d_dims(); + // Vector widths: 16-byte loads of both inputs, 8-byte stores of both outputs + // and of each row's scales. + const bool aligned = is_aligned_ptr(input.data.dptr, 16) && + (!IS_DACT || is_aligned_ptr(act_input->data.dptr, 16)) && + is_aligned_ptr(output.data.dptr, 8) && + is_aligned_ptr(output.columnwise_data.dptr, 8) && + is_aligned_ptr(output.scale_inv.dptr, 8) && + is_aligned_ptr(output.columnwise_scale_inv.dptr, 2); + if (!(transformer_engine::cuda::sm_arch() >= 100 && !use_2d_quantization && output.has_data() && + output.has_columnwise_data() && output.amax.dptr == nullptr && noop.data.dptr == nullptr && + shape_supported(rows, cols) && aligned)) { + return false; + } + constexpr size_t kMaxGridY = 65535; + const int M = static_cast(rows); + return static_cast(M / kRowsPerMxBlock / + walk_length(M, static_cast(cols))) <= kMaxGridY; +} + +template +void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, + const activation_table::TableRef& table, cudaStream_t stream) { + using namespace activation_table; + // Plain variables: lambdas cannot capture structured bindings before C++20. + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim(); + const void* act_ptr = IS_DACT ? act_input->data.dptr : input.data.dptr; + const auto run = [&](auto table_type) { + using TableT = typename decltype(table_type)::type; + launch( + input.data.dptr, act_ptr, output->data.dptr, output->scale_inv.dptr, + output->columnwise_data.dptr, output->columnwise_scale_inv.dptr, static_cast(rows), + static_cast(cols), static_cast(output->scale_inv.shape[1]), + static_cast(output->columnwise_scale_inv.shape[1]), table.ptr, stream); + }; + constexpr NVTEActivationTableKind kKind = kKindOf; + if constexpr (kKind == kNVTEActivationTableDsilu) { + if (table.use == kExactDsiluTable) { + run(std::common_type{}); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + } + run(std::common_type>{}); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace regtile +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_MXFP8_REGTILE_CUH_ diff --git a/transformer_engine/common/util/activation_table.cuh b/transformer_engine/common/util/activation_table.cuh index 303e7df280e..d01d310f17c 100644 --- a/transformer_engine/common/util/activation_table.cuh +++ b/transformer_engine/common/util/activation_table.cuh @@ -59,7 +59,8 @@ struct ForwardLayout { static constexpr uint32_t kLowSentinel = kLo - 1; static constexpr uint32_t kHighSentinel = kHi + 1; static constexpr uint32_t kRows = kHighSentinel - kLowSentinel + 1; - static constexpr size_t kBytes = kRows * 2 * sizeof(__nv_bfloat16); + //! Rounded up to 16 bytes, the granularity of a bulk copy into shared memory. + static constexpr size_t kBytes = (kRows * 2 * sizeof(__nv_bfloat16) + 15) / 16 * 16; }; /*! \brief Derivative table over the BF16 magnitudes [kLo, kHi]. @@ -77,7 +78,11 @@ struct DerivativeLayout { static constexpr uint32_t kEntriesPerSign = kHi - kLo + 1; static constexpr uint32_t kSignStride = kStride; static_assert(kEntriesPerSign <= kSignStride, "The two signs' entries must not overlap."); - static constexpr size_t kBytes = (kSignStride + kEntriesPerSign) * sizeof(float); + //! Rounded up to 16 bytes, the granularity of a bulk copy into shared memory. + static constexpr size_t kBytes = ((kSignStride + kEntriesPerSign) * sizeof(float) + 15) / 16 * 16; + //! The sign bit, 0x8000, shifted right by this is the byte offset of the negative half. + static constexpr int kSignShift = kStride == 2048 ? 2 : 1; + static_assert(kStride == 2048 || kStride == 4096, "The packed lookup assumes these strides."); }; template @@ -244,6 +249,86 @@ __device__ __forceinline__ float derivative(const float x, const float *__restri return magnitude >= T::kSaturatedEndBits ? __uint_as_float(0x7fffffffu) : value; } +/*! \brief Forward activation of a BF16 pair, low lane in the low half. Same results as + * forward() per lane; the table may be in shared memory. + */ +template +__device__ __forceinline__ uint32_t forward_2x(const uint32_t x, + const __nv_bfloat16 *__restrict__ table) { + static_assert(T::kForward, "Not a forward table."); + const __nv_bfloat162 a = *reinterpret_cast(&x); + // Clamp one lane's magnitude to the table rows. The +0 bound of the high lane zeroes it, + // so the word is exactly the clamped magnitude. + const auto row = [](const __nv_bfloat162 lane) { + constexpr uint32_t kLowBound = T::kLowSentinel * 0x10001u; + constexpr uint32_t kHighBound = T::kHighSentinel; + __nv_bfloat162 m = + __hmax2(__habs2(lane), *reinterpret_cast(&kLowBound)); + m = __hmin2(m, *reinterpret_cast(&kHighBound)); + return *reinterpret_cast(&m) - T::kLowSentinel; + }; + const __nv_bfloat16 entry_lo = table[2 * row(__low2bfloat162(a)) + ((x >> 15) & 1u)]; + const __nv_bfloat16 entry_hi = table[2 * row(__high2bfloat162(a)) + (x >> 31)]; + constexpr uint32_t kHalf = 0x3f003f00u, kNegativeZero = 0x80008000u; + const __nv_bfloat162 half = __hmul2(a, *reinterpret_cast(&kHalf)); + const __nv_bfloat162 upper = + __hmax2_nan(a, *reinterpret_cast(&kNegativeZero)); + const __nv_bfloat162 r = + __hmin2_nan(__hmax2(__halves2bfloat162(entry_lo, entry_hi), half), upper); + return *reinterpret_cast(&r); +} + +/*! \brief Derivative of a BF16 pair, low lane in .x. Same results as derivative() per lane; + * the table may be in shared memory. + * + * Both lanes' byte offsets are formed in one 32-bit word: the clamp runs on the BF16 pair, + * and each lane's offset stays below 0x10000, so the halves never carry into each other. + */ +template +__device__ __forceinline__ float2 derivative_2x(const uint32_t x, + const unsigned char *__restrict__ table) { + static_assert(!T::kForward, "Not a derivative table."); + static_assert((T::kSignStride + T::kEntriesPerSign) * sizeof(float) <= 0x10000, + "Byte offsets must fit in 16 bits."); + constexpr uint32_t kLo2 = T::kLoBits * 0x10001u, kHi2 = T::kHiBits * 0x10001u; + const uint32_t magnitude = x & 0x7fff7fffu; + const __nv_bfloat162 clamped = + __hmin2(__hmax2(*reinterpret_cast(&magnitude), + *reinterpret_cast(&kLo2)), + *reinterpret_cast(&kHi2)); + const uint32_t offsets = ((*reinterpret_cast(&clamped) - kLo2) << 2) + + ((x ^ magnitude) >> T::kSignShift); + const float entry_lo = *reinterpret_cast(table + (offsets & 0xffffu)); + const float entry_hi = *reinterpret_cast(table + (offsets >> 16)); + const float linear_lo = T::below_window(__uint_as_float(x << 16)); + const float linear_hi = T::below_window(__uint_as_float(x & 0xffff0000u)); + constexpr uint32_t kSaturatedEnd2 = T::kSaturatedEndBits * 0x10001u; + float lo, hi; +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // One BF16x2 compare per condition sets both lanes' predicates. The unordered >= also + // holds for NaN, which gives NaN like the saturated end. + asm("{\n\t" + ".reg .pred below_lo, below_hi, nan_lo, nan_hi;\n\t" + "setp.lt.bf16x2 below_lo|below_hi, %2, %3;\n\t" + "setp.geu.bf16x2 nan_lo|nan_hi, %2, %4;\n\t" + "selp.f32 %0, %5, %6, below_lo;\n\t" + "selp.f32 %1, %7, %8, below_hi;\n\t" + "selp.b32 %0, 0x7fffffff, %0, nan_lo;\n\t" + "selp.b32 %1, 0x7fffffff, %1, nan_hi;\n\t" + "}" + : "=f"(lo), "=f"(hi) + : "r"(magnitude), "r"(kLo2), "r"(kSaturatedEnd2), "f"(linear_lo), "f"(entry_lo), + "f"(linear_hi), "f"(entry_hi)); +#else + const uint32_t magnitude_lo = magnitude & 0xffffu, magnitude_hi = magnitude >> 16; + lo = magnitude_lo < T::kLoBits ? linear_lo : entry_lo; + hi = magnitude_hi < T::kLoBits ? linear_hi : entry_hi; + lo = magnitude_lo >= T::kSaturatedEndBits ? __uint_as_float(0x7fffffffu) : lo; + hi = magnitude_hi >= T::kSaturatedEndBits ? __uint_as_float(0x7fffffffu) : hi; +#endif + return make_float2(lo, hi); +} + /*! \brief OP(x), from the table given by kUse, or from util/math.h for kNoTable. * * x must be a BF16 value widened to FP32 when reading a table. diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 74fd2261f58..acba5f597fe 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -732,6 +732,47 @@ __device__ __forceinline__ void abs_max_2x(fp16x2 &dst, const fp16x2 &p1, const #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) } +// Packed FP32x2 arithmetic. Each lane rounds exactly like the scalar .rn operation. +__device__ __forceinline__ floatx2 add_2x(const floatx2 &a, const floatx2 &b) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + floatx2 d; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(reinterpret_cast(d)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); + return d; +#else + NVTE_DEVICE_ERROR("add_2x is only supported on SM 10.0+."); + return {}; +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ floatx2 mul_2x(const floatx2 &a, const floatx2 &b) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + floatx2 d; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(reinterpret_cast(d)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); + return d; +#else + NVTE_DEVICE_ERROR("mul_2x is only supported on SM 10.0+."); + return {}; +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ floatx2 fma_2x(const floatx2 &a, const floatx2 &b, const floatx2 &c) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + floatx2 d; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(reinterpret_cast(d)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); + return d; +#else + NVTE_DEVICE_ERROR("fma_2x is only supported on SM 10.0+."); + return {}; +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + __device__ __forceinline__ int32_t elect_one_sync(uint32_t mask = 0xFFFFFFFFu) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) int32_t pred = 0; @@ -836,6 +877,17 @@ __device__ __forceinline__ bf16x2 exp2f_rcp_2x(e8m0_t biased_exp) { return result; } +/*! \brief exp2f_rcp for two biased exponents at once, one in the low byte of each + * 16-bit half of \p biased_exps, returned as a BF16 pair. + * + * Exponent 254 gives the subnormal 2^-127, as exp2f_rcp_2x does. Exponent 255 is not handled. + */ +__device__ __forceinline__ bf16x2 exp2f_rcp_2x_per_lane(const uint32_t biased_exps) { + const uint32_t bits = + ((0x00FE00FEu - biased_exps) << 7) | (__vcmpeq2(biased_exps, 0x00FE00FEu) & 0x00400040u); + return reinterpret_cast(bits); +} + // Scale two BF16 pairs by independent scales and pack the four results into a // single FP8E4M3 word. The mul_cvt_4x overload below shares one scale across // all four elements, which is what a rowwise MX block wants; a colwise block diff --git a/transformer_engine/common/util/ptx_arch_spec.cuh b/transformer_engine/common/util/ptx_arch_spec.cuh index 14f390740e1..0b5599adb5a 100644 --- a/transformer_engine/common/util/ptx_arch_spec.cuh +++ b/transformer_engine/common/util/ptx_arch_spec.cuh @@ -103,6 +103,20 @@ __device__ __forceinline__ e8m0_t float_to_e8m0(float val) { } } +/*! \brief float_to_e8m0 for two values at once, returned as two packed bytes: + * \p hi in the upper byte, \p lo in the lower one. + */ +__device__ __forceinline__ uint32_t float_to_e8m0_2x(const float hi, const float lo) { + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + uint16_t out; + asm("cvt.rp.satfinite.ue8m0x2.f32 %0, %1, %2;" : "=h"(out) : "f"(hi), "f"(lo)); + return out; + } else { + return (static_cast(float_to_e8m0(hi)) << 8) | float_to_e8m0(lo); + } +} + __device__ __forceinline__ void reduce_sync_max_abs_f32(float &out, float const &in) { constexpr bool is_sm_100f = NVTE_CUDA_ARCH_MATCHES(ptx::FamilySpecific<100>); if constexpr (is_sm_100f) { From 8a2ffd349cb53c7d9bab1690dde41b2516d3d0c5 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 14 Oct 2026 13:09:52 -0700 Subject: [PATCH 04/29] [Common] Build the activation tables in TE common on first use Instead of having each framework allocate the tables and pass them through new _v2 entry points, TE common builds each table the first time it is needed and keeps it, once per device and process, as it does for library handles. The existing nvte_gelu, nvte_silu, nvte_dgelu, nvte_dsilu and nvte_quantize_dbias_d{gelu,silu} read the tables for BF16 inputs quantized to MXFP8, with no API change. The first use may happen while a CUDA graph is being captured. The build then runs in relaxed capture mode on a private stream and completes before the captured kernel is recorded, so the replayed graph reads a complete table. The tables are on by default. NVTE_ACTIVATION_TABLE=0 turns them off, and NVTE_MXFP8_REGTILE=0 keeps table requests on the generic kernel. When the CuTeDSL backend is enabled, it runs first and computes the activations itself. Remove the _v2 entry points, the activation-table quantization config attribute, the public table build and size functions, and the PyTorch and JAX table caches. Signed-off-by: Teddy Do --- docs/envvars.rst | 14 +- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 95 ++++-------- tests/cpp/util/test_activation_table.cu | 137 +++++++++++------- .../common/activation/activation_table.cu | 104 ++++++++----- .../common/activation/activation_template.h | 11 +- transformer_engine/common/activation/gelu.cu | 14 -- .../common/activation/gelu_dbias.cu | 13 -- .../common/activation/swiglu.cu | 14 -- .../common/activation/swiglu_dbias.cu | 13 -- .../common/cast/dispatch/quantize.cuh | 51 +++---- .../common/cast/mxfp8/quantize_mxfp8.cuh | 2 +- .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 8 +- transformer_engine/common/common.h | 4 +- .../include/transformer_engine/activation.h | 102 ------------- .../common/include/transformer_engine/cast.h | 34 ----- .../transformer_engine/transformer_engine.h | 11 -- .../common/libtransformer_engine.version | 1 + .../common/transformer_engine.cpp | 6 - .../common/util/activation_table.cuh | 94 ++++++------ transformer_engine/jax/csrc/extensions.h | 5 - .../jax/csrc/extensions/activation.cpp | 49 ++----- .../jax/csrc/extensions/activation_table.cpp | 78 ---------- transformer_engine/pytorch/csrc/common.h | 52 ------- .../pytorch/csrc/extensions/activation.cpp | 12 +- .../csrc/extensions/activation_table.cpp | 57 -------- .../pytorch/csrc/extensions/bias.cpp | 10 +- 26 files changed, 297 insertions(+), 694 deletions(-) delete mode 100644 transformer_engine/jax/csrc/extensions/activation_table.cpp delete mode 100644 transformer_engine/pytorch/csrc/extensions/activation_table.cpp diff --git a/docs/envvars.rst b/docs/envvars.rst index 02b1a2417a4..66f18018f38 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -353,11 +353,23 @@ Kernel Configuration :Default: ``0`` :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. Existing transpose operations select their static fallback automatically. NVFP4 4over6 quantization, fused softmax, and normalization paths require their corresponding ``NVTE_BUILD_LEGACY_STATIC_NVFP4``, ``NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX``, or ``NVTE_BUILD_LEGACY_STATIC_NORM`` CMake option to have been enabled when the library was built; otherwise no static fallback is available. +.. envvar:: NVTE_ACTIVATION_TABLE + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Read GeLU, SiLU and their derivatives from lookup tables when BF16 inputs are quantized to MXFP8. Each table is built on its first use and kept, once per device and process (under 64 KB per device in total), outside the framework's memory allocator. Results match the computed activations, except that GeLU and SiLU of ``-inf`` give ``-0`` instead of NaN, and derivatives of inputs with ``|x| < 2^-12`` (GeLU) or ``|x| < 2^-9`` (SiLU) are within 1 ulp (GeLU) or 2 ulp (SiLU) of the computed ones. When set to ``0``, the activations are computed. + .. envvar:: NVTE_ACTIVATION_TABLE_EXACT_DSILU :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Experimental. Choose the SiLU-derivative lookup table used when BF16 inputs are quantized to MXFP8 with an activation table. By default the 16 KB table covers ``|x| >= 2^-9`` and smaller inputs use the linear term ``0.5 + x/2``, which is within 2 ulp of the computed derivative. When set to ``1``, a 32 KB table covers ``|x| >= 2^-25``, and the derivative matches the computed one for every input. Read once per process, when the table is built. + :Description: Experimental. Choose the SiLU-derivative lookup table (see :envvar:`NVTE_ACTIVATION_TABLE`). By default the 16 KB table covers ``|x| >= 2^-9`` and smaller inputs use the linear term ``0.5 + x/2``, which is within 2 ulp of the computed derivative. When set to ``1``, a 32 KB table covers ``|x| >= 2^-25``, and the derivative matches the computed one for every input. Read once per process, when the table is built. + +.. envvar:: NVTE_MXFP8_REGTILE + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Allow the register-tiled MXFP8 kernel for activations read from lookup tables (see :envvar:`NVTE_ACTIVATION_TABLE`). When set to ``0``, those requests use the generic MXFP8 kernel, which gives the same results. .. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index 348a23fb321..b00b1649738 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -11,10 +11,10 @@ // specialized cast kernels). dbias is compared against a host sum in the order the kernel // accumulates it. // -// Without an activation table the activation is the scalar util/math.h function. With one, BF16 -// inputs follow the table specification (see NVTEActivationTableKind): -inf gives -0 in the -// forward functions, and the derivatives below the table window are the linear Taylor term. -// Other input types ignore the table. +// With NVTE_ACTIVATION_TABLE=0 the activation is the scalar util/math.h function. With the tables +// (the default), BF16 inputs follow the table specification (see util/activation_table.cuh): -inf +// gives -0 in the forward functions, and the derivatives below the table window are the linear +// Taylor term. Other input types do not use the tables. #include #include @@ -28,8 +28,8 @@ #include #include #include -#include #include +#include #include #include #include @@ -241,31 +241,26 @@ std::vector reference_dbias(const std::vector &z, const size_t row return dbias; } -NVTEActivationTableKind table_kind(const Mode m) { - switch (m) { - case Mode::GeLU: - return kNVTEActivationTableGelu; - case Mode::SiLU: - return kNVTEActivationTableSilu; - case Mode::DGeLU: - case Mode::DBiasDGeLU: - return kNVTEActivationTableDgelu; - default: - return kNVTEActivationTableDsilu; +// Sets an environment variable for one scope, and restores its previous value afterwards. +class ScopedEnv { + public: + ScopedEnv(const char *name, const char *value) : name_(name) { + const char *old = std::getenv(name); + if (old != nullptr) old_ = old; + setenv(name, value, 1); } -} - -// Each table is built once for the whole test program, as a framework would per device. -NVTETensor activation_table(const NVTEActivationTableKind kind) { - static std::map> tables; - auto &table = tables[kind]; - if (!table) { - const size_t bytes = nvte_activation_table_size(kind); - table = std::make_unique("table", std::vector{bytes}, DType::kByte); - nvte_build_activation_table(kind, table->data(), 0); + ~ScopedEnv() { + if (old_) { + setenv(name_, old_->c_str(), 1); + } else { + unsetenv(name_); + } } - return table->data(); -} + + private: + const char *name_; + std::optional old_; +}; size_t count_byte_mismatches(const uint8_t *a, const uint8_t *b, const size_t n) { size_t m = 0; @@ -316,31 +311,8 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t out.set_with_gemm_swizzled_scales(swizzled); Tensor dbias("dbias", std::vector{cols}, itype); Tensor workspace; - if (use_table) { - QuantizationConfigWrapper config; - config.set_activation_table(activation_table(table_kind(mode))); - switch (mode) { - case Mode::GeLU: - nvte_gelu_v2(x.data(), out.data(), config, 0); - break; - case Mode::SiLU: - nvte_silu_v2(x.data(), out.data(), config, 0); - break; - case Mode::DGeLU: - nvte_dgelu_v2(grad.data(), x.data(), out.data(), config, 0); - break; - case Mode::DSiLU: - nvte_dsilu_v2(grad.data(), x.data(), out.data(), config, 0); - break; - default: { - auto fn = mode == Mode::DBiasDGeLU ? &nvte_quantize_dbias_dgelu_v2 - : &nvte_quantize_dbias_dsilu_v2; - fn(grad.data(), x.data(), out.data(), dbias.data(), workspace.data(), config, 0); - workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype()); - fn(grad.data(), x.data(), out.data(), dbias.data(), workspace.data(), config, 0); - } - } - } else { + { + const ScopedEnv tables("NVTE_ACTIVATION_TABLE", use_table ? "1" : "0"); switch (mode) { case Mode::GeLU: nvte_gelu(x.data(), out.data(), 0); @@ -409,7 +381,7 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t // With a table, the register-resident kernel and the generic kernel must write the same bytes. // The first call takes the default dispatch, which picks the register-resident kernel for these -// shapes; a noop tensor holding 0 keeps the second call on the generic kernel without skipping it. +// shapes; NVTE_MXFP8_REGTILE=0 keeps the second call on the generic kernel. void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, const size_t cols, const bool swizzled) { const std::vector shape{rows, cols}; @@ -417,28 +389,25 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, Tensor x("x", shape, DType::kBFloat16); Tensor grad("grad", shape, DType::kBFloat16); fill_inputs(fill, &x, &grad); - Tensor noop("noop", std::vector{1}, DType::kFloat32); - Tensor outs[2] = { Tensor("default", shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING), Tensor("generic", shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING)}; for (int i = 0; i < 2; ++i) { outs[i].set_with_gemm_swizzled_scales(swizzled); - QuantizationConfigWrapper config; - config.set_activation_table(activation_table(table_kind(mode))); - if (i == 1) config.set_noop_tensor(noop.data()); + const ScopedEnv tables("NVTE_ACTIVATION_TABLE", "1"); + const ScopedEnv regtile("NVTE_MXFP8_REGTILE", i == 0 ? "1" : "0"); switch (mode) { case Mode::GeLU: - nvte_gelu_v2(x.data(), outs[i].data(), config, 0); + nvte_gelu(x.data(), outs[i].data(), 0); break; case Mode::SiLU: - nvte_silu_v2(x.data(), outs[i].data(), config, 0); + nvte_silu(x.data(), outs[i].data(), 0); break; case Mode::DGeLU: - nvte_dgelu_v2(grad.data(), x.data(), outs[i].data(), config, 0); + nvte_dgelu(grad.data(), x.data(), outs[i].data(), 0); break; default: - nvte_dsilu_v2(grad.data(), x.data(), outs[i].data(), config, 0); + nvte_dsilu(grad.data(), x.data(), outs[i].data(), 0); break; } } diff --git a/tests/cpp/util/test_activation_table.cu b/tests/cpp/util/test_activation_table.cu index b94ee45ea83..3b097ec8814 100644 --- a/tests/cpp/util/test_activation_table.cu +++ b/tests/cpp/util/test_activation_table.cu @@ -9,7 +9,8 @@ // - forward: util/math.h rounded to BF16, except that -inf gives -0; // - derivative: util/math.h, except below the table window, where it is the linear Taylor // term rounded once. -// The derivative is also checked against util/math.h with the ulp bound stated in the API. +// The derivative is also checked against util/math.h with the ulp bound stated in +// util/activation_table.cuh. The tables are the ones TE builds and caches per device. // With NVTE_ACTIVATION_TABLE_EXACT_DSILU=1 the dSiLU window starts at 2^-25 and dSiLU must match // util/math.h exactly. @@ -17,6 +18,7 @@ #include #include #include +#include #include #include @@ -27,12 +29,18 @@ #include #include +#include "../test_common.h" #include "util/activation_table.cuh" #include "util/math.h" namespace { using transformer_engine::Empty; +using transformer_engine::activation_table::kDgelu; +using transformer_engine::activation_table::kDsilu; +using transformer_engine::activation_table::kGelu; +using transformer_engine::activation_table::Kind; +using transformer_engine::activation_table::kSilu; constexpr int kNumInputs = 1 << 16; @@ -46,18 +54,18 @@ bool exact_dsilu() { return env != nullptr && std::string(env) == "1"; } -__device__ float spec(const NVTEActivationTableKind kind, const uint32_t bits, +__device__ float spec(const Kind kind, const uint32_t bits, const uint32_t dsilu_window_low) { const float x = __uint_as_float(bits << 16); const uint32_t magnitude = bits & 0x7fffu; switch (kind) { - case kNVTEActivationTableGelu: + case kGelu: if (bits == 0xff80u) return -0.0f; return __bfloat162float(__float2bfloat16_rn(transformer_engine::gelu(x, {}))); - case kNVTEActivationTableSilu: + case kSilu: if (bits == 0xff80u) return -0.0f; return __bfloat162float(__float2bfloat16_rn(transformer_engine::silu(x, {}))); - case kNVTEActivationTableDgelu: + case kDgelu: if (magnitude < kDgeluWindowLow) return fmaf(0.79788456f, x, 0.5f); return transformer_engine::dgelu(x, {}); default: @@ -66,35 +74,35 @@ __device__ float spec(const NVTEActivationTableKind kind, const uint32_t bits, } } -__device__ float scalar_math(const NVTEActivationTableKind kind, const float x) { +__device__ float scalar_math(const Kind kind, const float x) { switch (kind) { - case kNVTEActivationTableDgelu: + case kDgelu: return transformer_engine::dgelu(x, {}); default: return transformer_engine::dsilu(x, {}); } } -template +template struct Op; template <> -struct Op { +struct Op { static constexpr float (*value)(float, const Empty &) = transformer_engine::gelu; }; template <> -struct Op { +struct Op { static constexpr float (*value)(float, const Empty &) = transformer_engine::silu; }; template <> -struct Op { +struct Op { static constexpr float (*value)(float, const Empty &) = transformer_engine::dgelu; }; template <> -struct Op { +struct Op { static constexpr float (*value)(float, const Empty &) = transformer_engine::dsilu; }; -template +template __global__ void lookup_kernel(const void *table, const uint32_t dsilu_window_low, float *looked_up, float *expected, float *math_h) { const uint32_t bits = blockIdx.x * blockDim.x + threadIdx.x; @@ -103,7 +111,7 @@ __global__ void lookup_kernel(const void *table, const uint32_t dsilu_window_low const float x = __uint_as_float(bits << 16); looked_up[bits] = at::evaluate::value>(x, table); expected[bits] = spec(kKind, bits, dsilu_window_low); - math_h[bits] = (kKind == kNVTEActivationTableDgelu || kKind == kNVTEActivationTableDsilu) + math_h[bits] = (kKind == kDgelu || kKind == kDsilu) ? scalar_math(kKind, x) : 0.0f; } @@ -122,30 +130,18 @@ int64_t ordered(const float f) { return i < 0 ? -static_cast(i & 0x7fffffff) : i; } -struct DeviceTable { - void *ptr = nullptr; - size_t bytes = 0; - explicit DeviceTable(const NVTEActivationTableKind kind) { - bytes = nvte_activation_table_size(kind); - EXPECT_EQ(cudaMalloc(&ptr, bytes), cudaSuccess); - transformer_engine::TensorWrapper table(ptr, std::vector{bytes}, - transformer_engine::DType::kByte); - nvte_build_activation_table(kind, table.data(), 0); - EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); - } - ~DeviceTable() { cudaFree(ptr); } -}; - -template +template void check_table(const int max_ulp_below_window) { - DeviceTable table(kKind); + namespace at = transformer_engine::activation_table; + const at::TableRef table = at::device_table(kKind); + const bool exact = kKind == kDsilu && exact_dsilu(); + ASSERT_NE(table.ptr, nullptr); + ASSERT_EQ(table.use, exact ? at::kExactDsiluTable : at::kDefaultTable); float *d = nullptr; ASSERT_EQ(cudaMalloc(&d, 3 * kNumInputs * sizeof(float)), cudaSuccess); - namespace at = transformer_engine::activation_table; - const bool exact = kKind == kNVTEActivationTableDsilu && exact_dsilu(); const uint32_t dsilu_window_low = exact ? kExactDsiluWindowLow : kDsiluWindowLow; if (exact) { - if constexpr (kKind == kNVTEActivationTableDsilu) { + if constexpr (kKind == kDsilu) { lookup_kernel<<>>( table.ptr, dsilu_window_low, d, d + kNumInputs, d + 2 * kNumInputs); } @@ -183,35 +179,70 @@ void check_table(const int max_ulp_below_window) { } // namespace TEST(ActivationTableTest, GeluMatchesSpecification) { - check_table(-1); + check_table(-1); } TEST(ActivationTableTest, SiluMatchesSpecification) { - check_table(-1); + check_table(-1); } TEST(ActivationTableTest, DgeluMatchesSpecification) { - check_table(1); + check_table(1); } TEST(ActivationTableTest, DsiluMatchesSpecification) { - check_table(exact_dsilu() ? 0 : 2); + check_table(exact_dsilu() ? 0 : 2); } -TEST(ActivationTableTest, SizesIdentifyKind) { - const size_t sizes[] = {nvte_activation_table_size(kNVTEActivationTableGelu), - nvte_activation_table_size(kNVTEActivationTableSilu), - nvte_activation_table_size(kNVTEActivationTableDgelu), - nvte_activation_table_size(kNVTEActivationTableDsilu)}; - for (int i = 0; i < 4; ++i) { - for (int j = i + 1; j < 4; ++j) EXPECT_NE(sizes[i], sizes[j]); +TEST(ActivationTableTest, OneTablePerDevice) { + namespace at = transformer_engine::activation_table; + int num_devices = 0; + ASSERT_EQ(cudaGetDeviceCount(&num_devices), cudaSuccess); + int current = 0; + ASSERT_EQ(cudaGetDevice(¤t), cudaSuccess); + for (const Kind kind : {kGelu, kSilu, kDgelu, kDsilu}) { + std::vector tables; + for (int device = 0; device < num_devices; ++device) { + ASSERT_EQ(cudaSetDevice(device), cudaSuccess); + const void *first = at::device_table(kind).ptr; + EXPECT_EQ(at::device_table(kind).ptr, first) << "kind " << kind << ", device " << device; + for (const void *other : tables) EXPECT_NE(first, other); + tables.push_back(first); + } } + ASSERT_EQ(cudaSetDevice(current), cudaSuccess); } -TEST(ActivationTableTest, RejectsTableOfWrongSize) { - const size_t bytes = nvte_activation_table_size(kNVTEActivationTableDgelu); - void *ptr = nullptr; - ASSERT_EQ(cudaMalloc(&ptr, bytes), cudaSuccess); - transformer_engine::TensorWrapper table(ptr, std::vector{bytes}, - transformer_engine::DType::kByte); - EXPECT_THROW(nvte_build_activation_table(kNVTEActivationTableDsilu, table.data(), 0), - std::runtime_error); - cudaFree(ptr); +// The first use of a table, inside CUDA graph capture, must build the table outside the capture, +// and the replayed graph must read it. This runs in a fresh process (a death test), so that no +// earlier test has built the table. +TEST(ActivationTableDeathTest, FirstUseInsideGraphCapture) { + GTEST_FLAG_SET(death_test_style, "threadsafe"); + const auto body = [] { + const std::vector shape{256, 256}; + test::Tensor x("x", shape, transformer_engine::DType::kBFloat16); + test::fillUniform(&x); + test::Tensor captured("captured", shape, transformer_engine::DType::kFloat8E4M3, true, true, + NVTE_MXFP8_1D_SCALING); + test::Tensor eager("eager", shape, transformer_engine::DType::kFloat8E4M3, true, true, + NVTE_MXFP8_1D_SCALING); + cudaStream_t stream; + if (cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking) != cudaSuccess) std::exit(2); + cudaGraph_t graph; + cudaGraphExec_t exec; + if (cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal) != cudaSuccess) std::exit(3); + nvte_gelu(x.data(), captured.data(), stream); + if (cudaStreamEndCapture(stream, &graph) != cudaSuccess) std::exit(4); + if (cudaGraphInstantiate(&exec, graph, 0) != cudaSuccess) std::exit(5); + if (cudaGraphLaunch(exec, stream) != cudaSuccess) std::exit(6); + nvte_gelu(x.data(), eager.data(), stream); + if (cudaStreamSynchronize(stream) != cudaSuccess) std::exit(7); + captured.to_cpu(); + eager.to_cpu(); + const size_t n = shape[0] * shape[1]; + const bool same = + std::memcmp(captured.rowwise_cpu_dptr(), + eager.rowwise_cpu_dptr(), n) == 0 && + std::memcmp(captured.columnwise_cpu_dptr(), + eager.columnwise_cpu_dptr(), n) == 0; + std::exit(same ? 0 : 1); + }; + EXPECT_EXIT(body(), testing::ExitedWithCode(0), ""); } diff --git a/transformer_engine/common/activation/activation_table.cu b/transformer_engine/common/activation/activation_table.cu index 9402e4be2fb..f02b3c9075c 100644 --- a/transformer_engine/common/activation/activation_table.cu +++ b/transformer_engine/common/activation/activation_table.cu @@ -7,10 +7,12 @@ // Builds the activation tables from the scalar util/math.h functions. This file must not be // compiled with --use_fast_math: the tables have to hold exactly what util/math.h computes. -#include +#include +#include #include "../common.h" #include "../util/activation_table.cuh" +#include "../util/cuda_runtime.h" namespace transformer_engine { namespace activation_table { @@ -65,46 +67,80 @@ void build(void *table, cudaStream_t stream) { NVTE_CHECK_CUDA(cudaGetLastError()); } -} // namespace -} // namespace activation_table -} // namespace transformer_engine - -size_t nvte_activation_table_size(NVTEActivationTableKind kind) { - NVTE_API_CALL(nvte_activation_table_size); - const size_t bytes = transformer_engine::activation_table::built_table_bytes(kind); - NVTE_CHECK(bytes != 0, "Invalid NVTEActivationTableKind (got ", static_cast(kind), ")"); - return bytes; -} - -void nvte_build_activation_table(NVTEActivationTableKind kind, NVTETensor table, - cudaStream_t stream) { - NVTE_API_CALL(nvte_build_activation_table); - using namespace transformer_engine; - using namespace transformer_engine::activation_table; - const TableRef ref = checked_table(*convertNVTETensorCheck(table), kind); - NVTE_CHECK((ref.use == kExactDsiluTable) == - (kind == kNVTEActivationTableDsilu && use_exact_dsilu_table()), - "Activation table size does not match NVTE_ACTIVATION_TABLE_EXACT_DSILU; allocate " - "nvte_activation_table_size(kind) bytes."); - void *ptr = const_cast(ref.ptr); +void build(const Kind kind, void *table, cudaStream_t stream) { switch (kind) { - case kNVTEActivationTableGelu: - build>(ptr, stream); + case kGelu: + build>(table, stream); break; - case kNVTEActivationTableSilu: - build>(ptr, stream); + case kSilu: + build>(table, stream); break; - case kNVTEActivationTableDgelu: - build>(ptr, stream); + case kDgelu: + build>(table, stream); break; - case kNVTEActivationTableDsilu: - if (ref.use == kExactDsiluTable) { - build(ptr, stream); + case kDsilu: + if (use_exact_dsilu_table()) { + build(table, stream); } else { - build>(ptr, stream); + build>(table, stream); } break; default: - NVTE_ERROR("Invalid NVTEActivationTableKind (got ", static_cast(kind), ")"); + NVTE_ERROR("Invalid activation table kind (got ", static_cast(kind), ")"); } } + +// Sets this thread's stream capture mode, and restores the previous one when destroyed. +class CaptureModeGuard { + public: + explicit CaptureModeGuard(cudaStreamCaptureMode mode) : mode_(mode) { + NVTE_CHECK_CUDA(cudaThreadExchangeStreamCaptureMode(&mode_)); + } + ~CaptureModeGuard() { cudaThreadExchangeStreamCaptureMode(&mode_); } + CaptureModeGuard(const CaptureModeGuard &) = delete; + CaptureModeGuard &operator=(const CaptureModeGuard &) = delete; + + private: + cudaStreamCaptureMode mode_; +}; + +// Allocates and fills a table on the current device, synchronously. +void *allocate_and_build(const Kind kind) { + // The caller's stream, or another thread's, may be capturing a CUDA graph. In relaxed mode + // this thread may still allocate, and a private stream keeps the build out of the capture. + // The table is complete when this returns, so captured kernels can read it on replay. + CaptureModeGuard relaxed(cudaStreamCaptureModeRelaxed); + void *table = nullptr; + NVTE_CHECK_CUDA(cudaMalloc(&table, built_table_bytes(kind))); + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + build(kind, table, stream); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); + return table; +} + +// One set of tables per device. Tables are never freed: freeing them from a static destructor +// could run after the CUDA runtime has shut down. +struct DeviceTables { + std::once_flag built[kNumKinds]; + void *table[kNumKinds] = {}; +}; + +} // namespace + +TableRef device_table(const Kind kind) { + NVTE_CHECK(kind >= 0 && kind < kNumKinds, "Invalid activation table kind (got ", + static_cast(kind), ")"); + static const int num_devices = cuda::num_devices(); + static const std::unique_ptr tables(new DeviceTables[num_devices]); + const int device = cuda::current_device(); + NVTE_CHECK(device >= 0 && device < num_devices, "Invalid device (got ", device, ")"); + DeviceTables &mine = tables[device]; + std::call_once(mine.built[kind], [&] { mine.table[kind] = allocate_and_build(kind); }); + const bool exact = kind == kDsilu && use_exact_dsilu_table(); + return {mine.table[kind], exact ? kExactDsiluTable : kDefaultTable}; +} + +} // namespace activation_table +} // namespace transformer_engine diff --git a/transformer_engine/common/activation/activation_template.h b/transformer_engine/common/activation/activation_template.h index c4f37b53d81..ffbffafd1a1 100644 --- a/transformer_engine/common/activation/activation_template.h +++ b/transformer_engine/common/activation/activation_template.h @@ -21,16 +21,15 @@ namespace transformer_engine { template -void act_fn(const NVTETensor input, NVTETensor output, cudaStream_t stream, - const NVTEQuantizationConfig quant_config = nullptr) { +void act_fn(const NVTETensor input, NVTETensor output, cudaStream_t stream) { using namespace detail; constexpr bool IS_ACT = true; - dispatch::quantize_fwd_helper(input, output, quant_config, stream); + dispatch::quantize_fwd_helper(input, output, nullptr, stream); } template -void dact_fn(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream, - const NVTEQuantizationConfig quant_config = nullptr) { +void dact_fn(const NVTETensor grad, const NVTETensor input, NVTETensor output, + cudaStream_t stream) { using namespace detail; constexpr bool IS_DBIAS = false; constexpr bool IS_DACT = true; @@ -38,7 +37,7 @@ void dact_fn(const NVTETensor grad, const NVTETensor input, NVTETensor output, c constexpr NVTETensor workspace = nullptr; dispatch::quantize_bwd_helper(grad, input, output, dbias, workspace, - quant_config, stream); + nullptr, stream); } template diff --git a/transformer_engine/common/activation/gelu.cu b/transformer_engine/common/activation/gelu.cu index 9c4b1218f06..6bd63672caa 100644 --- a/transformer_engine/common/activation/gelu.cu +++ b/transformer_engine/common/activation/gelu.cu @@ -13,13 +13,6 @@ void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { act_fn>(input, output, stream); } -void nvte_gelu_v2(const NVTETensor input, NVTETensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - NVTE_API_CALL(nvte_gelu_v2); - using namespace transformer_engine; - act_fn>(input, output, stream, quant_config); -} - void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dgelu); @@ -27,13 +20,6 @@ void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, stream); } -void nvte_dgelu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - NVTE_API_CALL(nvte_dgelu_v2); - using namespace transformer_engine; - dact_fn>(grad, input, output, stream, quant_config); -} - void nvte_geglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_geglu); using namespace transformer_engine; diff --git a/transformer_engine/common/activation/gelu_dbias.cu b/transformer_engine/common/activation/gelu_dbias.cu index e3215931a65..4eaa9e355b3 100644 --- a/transformer_engine/common/activation/gelu_dbias.cu +++ b/transformer_engine/common/activation/gelu_dbias.cu @@ -20,19 +20,6 @@ void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor activati input, activation_input, output, dbias, workspace, nullptr, stream); } -void nvte_quantize_dbias_dgelu_v2(const NVTETensor input, const NVTETensor activation_input, - NVTETensor output, NVTETensor dbias, NVTETensor workspace, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_dbias_dgelu_v2); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, quant_config, stream); -} - void nvte_quantize_dbias_dqgelu(const NVTETensor input, const NVTETensor activation_input, NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream) { diff --git a/transformer_engine/common/activation/swiglu.cu b/transformer_engine/common/activation/swiglu.cu index d4ae6edeb9e..0300ba1bd46 100644 --- a/transformer_engine/common/activation/swiglu.cu +++ b/transformer_engine/common/activation/swiglu.cu @@ -15,13 +15,6 @@ void nvte_silu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { act_fn>(input, output, stream); } -void nvte_silu_v2(const NVTETensor input, NVTETensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - NVTE_API_CALL(nvte_silu_v2); - using namespace transformer_engine; - act_fn>(input, output, stream, quant_config); -} - void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dsilu); @@ -29,13 +22,6 @@ void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, stream); } -void nvte_dsilu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - NVTE_API_CALL(nvte_dsilu_v2); - using namespace transformer_engine; - dact_fn>(grad, input, output, stream, quant_config); -} - void nvte_swiglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_swiglu); using namespace transformer_engine; diff --git a/transformer_engine/common/activation/swiglu_dbias.cu b/transformer_engine/common/activation/swiglu_dbias.cu index 4120b949e53..0e532acc57d 100644 --- a/transformer_engine/common/activation/swiglu_dbias.cu +++ b/transformer_engine/common/activation/swiglu_dbias.cu @@ -19,16 +19,3 @@ void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor activati dispatch::quantize_bwd_helper>( input, activation_input, output, dbias, workspace, nullptr, stream); } - -void nvte_quantize_dbias_dsilu_v2(const NVTETensor input, const NVTETensor activation_input, - NVTETensor output, NVTETensor dbias, NVTETensor workspace, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_dbias_dsilu_v2); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, quant_config, stream); -} diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 460e9a15907..665af3823c6 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -39,14 +39,12 @@ namespace transformer_engine { namespace dispatch { // The activation table to read for OP, if any. Tables apply to BF16 inputs of the activations -// that have one. +// that have one, unless NVTE_ACTIVATION_TABLE=0. The first call per device builds the table. template -activation_table::TableRef activation_lut_for(const QuantizationConfig &quant_config, - const DType input_dtype) { +activation_table::TableRef activation_lut_for(const DType input_dtype) { if constexpr (activation_table::kHasTable) { - if (quant_config.activation_table != nullptr && input_dtype == DType::kBFloat16) { - return activation_table::checked_table(*convertNVTETensorCheck(quant_config.activation_table), - activation_table::kKindOf); + if (input_dtype == DType::kBFloat16 && activation_table::tables_enabled()) { + return activation_table::device_table(activation_table::kKindOf); } } return {}; @@ -109,21 +107,19 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const Tensor *dummy_input_tensor = nullptr; Tensor *dummy_dbias_tensor = nullptr; Tensor *dummy_workspace_tensor = nullptr; - const activation_table::TableRef activation_lut = - IS_ACT ? activation_lut_for(quant_config_cpp, input_tensor->dtype()) - : activation_table::TableRef{}; bool quantized_with_cutedsl = false; #ifdef NVTE_WITH_CUTEDSL - // The CuTeDSL backend computes activations itself, so it cannot honor a table. - if (activation_lut.ptr == nullptr) { - quantized_with_cutedsl = - cutedsl_backend::mxfp8_quantize_cutedsl( - input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, - dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); - } + // The CuTeDSL backend, when enabled, computes activations without the tables. + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_cutedsl( + input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, + dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); #endif if (!quantized_with_cutedsl) { + const activation_table::TableRef activation_lut = + IS_ACT ? activation_lut_for(input_tensor->dtype()) + : activation_table::TableRef{}; mxfp8::quantize( *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, activation_lut, stream); @@ -304,22 +300,19 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens break; } case NVTE_MXFP8_1D_SCALING: { - activation_table::TableRef activation_lut; - if constexpr (IS_DACT) { - activation_lut = activation_lut_for(quant_config_cpp, input_tensor->dtype()); - } bool quantized_with_cutedsl = false; #ifdef NVTE_WITH_CUTEDSL - // The CuTeDSL backend computes activations itself, so it cannot honor a table. - if (activation_lut.ptr == nullptr) { - quantized_with_cutedsl = - cutedsl_backend::mxfp8_quantize_cutedsl( - grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, - workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); - } + // The CuTeDSL backend, when enabled, computes activations without the tables. + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_cutedsl( + grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, + quant_config_cpp.mxfp8_2d_quantization, stream); #endif if (!quantized_with_cutedsl) { + activation_table::TableRef activation_lut; + if constexpr (IS_DACT) { + activation_lut = activation_lut_for(input_tensor->dtype()); + } mxfp8::quantize( *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, activation_lut, stream); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index c30e40d508f..f642ee9f5e6 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -1059,7 +1059,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, launch_kernel(std::integral_constant{}); } else if (activation_lut.use == kExactDsiluTable) { if constexpr (activation_table::kKindOf == - kNVTEActivationTableDsilu) { + activation_table::kDsilu) { launch_kernel(std::integral_constant{}); } else { NVTE_ERROR("Only dSiLU has an exact activation table."); diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh index 176821472e0..9a1c2a32f66 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -577,11 +577,13 @@ constexpr bool kImplements = (IS_ACT == activation_table::Table>::kForward); // Whether this kernel takes a request of a type in kImplements, with BF16 input, E4M3 -// output and an activation table. Anything else takes the generic kernel. +// output and an activation table. Anything else takes the generic kernel, as does every +// request with NVTE_MXFP8_REGTILE=0. template bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, const Tensor& noop, bool use_2d_quantization, const activation_table::TableRef& table) { if (table.use == activation_table::kNoTable) return false; + if (!getenv("NVTE_MXFP8_REGTILE", true)) return false; const auto [rows, cols] = input.flat_2d_dims(); // Vector widths: 16-byte loads of both inputs, 8-byte stores of both outputs // and of each row's scales. @@ -619,8 +621,8 @@ void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, static_cast(cols), static_cast(output->scale_inv.shape[1]), static_cast(output->columnwise_scale_inv.shape[1]), table.ptr, stream); }; - constexpr NVTEActivationTableKind kKind = kKindOf; - if constexpr (kKind == kNVTEActivationTableDsilu) { + constexpr Kind kKind = kKindOf; + if constexpr (kKind == kDsilu) { if (table.use == kExactDsiluTable) { run(std::common_type{}); NVTE_CHECK_CUDA(cudaGetLastError()); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 621b6edd8c2..7d201f6c8ba 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -648,7 +648,6 @@ struct QuantizationConfig { NVTENVFP44Over6Mode nvfp4_4over6_mode = kNVTENVFP44Over6Disabled; bool nvfp4_4over6_err_use_fast_math = false; bool mxfp8_2d_quantization = false; - NVTETensor activation_table = nullptr; static constexpr size_t attr_sizes[] = { sizeof(uint8_t), // force_pow_2_scales @@ -661,8 +660,7 @@ struct QuantizationConfig { sizeof(uint8_t), // use_fast_math sizeof(uint8_t), // nvfp4_4over6_mode sizeof(uint8_t), // nvfp4_4over6_err_use_fast_math - sizeof(uint8_t), // mxfp8_2d_quantization - sizeof(NVTETensor) // activation_table + sizeof(uint8_t) // mxfp8_2d_quantization }; }; diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index aa25406689f..31429cf89cf 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -43,50 +43,6 @@ enum class NVTE_Activation_Type { CLAMPED_SWIGLU }; -/*! \enum NVTEActivationTableKind - * \brief Activation that can be read from a lookup table. - * - * A BF16 input takes only 2^16 values, so GeLU, SiLU and their derivatives of BF16 inputs - * can be read from a small table instead of computed. Build each table once per device - * with nvte_build_activation_table and pass it to the _v2 activation functions through - * kNVTEQuantizationConfigActivationTable. The table is read for BF16 inputs quantized to - * MXFP8; every other input type and recipe ignores it. - * - * Against util/math.h, table results differ only as follows: - * - GeLU and SiLU of -inf are -0 (the formulas give NaN). - * - For |x| < 2^-12 (dGeLU) or 2^-9 (dSiLU), the derivative is its linear Taylor term - * 0.5 + c*x rounded once, within 1 ulp (dGeLU) or 2 ulp (dSiLU) of util/math.h. - * With NVTE_ACTIVATION_TABLE_EXACT_DSILU=1 (experimental), the dSiLU table is 32 KB and - * matches util/math.h for every input. - */ -enum NVTEActivationTableKind { - kNVTEActivationTableGelu = 0, /*!< GeLU (tanh form), BF16 entries */ - kNVTEActivationTableSilu = 1, /*!< SiLU, BF16 entries */ - kNVTEActivationTableDgelu = 2, /*!< GeLU derivative (tanh form), FP32 entries */ - kNVTEActivationTableDsilu = 3, /*!< SiLU derivative, FP32 entries */ -}; - -/*! \brief Size in bytes of an activation table. - * - * \param[in] kind Activation the table holds. - * \return Size of the table in bytes. Each kind has a different size. The dSiLU size - * depends on NVTE_ACTIVATION_TABLE_EXACT_DSILU. - */ -size_t nvte_activation_table_size(NVTEActivationTableKind kind); - -/*! \brief Fill an activation table on the current device. - * - * The table is computed from the scalar activation functions that TE uses when no table - * is given. Build it once per device and reuse it for every call; it does not depend on - * the input. - * - * \param[in] kind Activation the table holds. - * \param[in,out] table Byte tensor of nvte_activation_table_size(kind) elements. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_build_activation_table(NVTEActivationTableKind kind, NVTETensor table, - cudaStream_t stream); - /*! \brief Computes the GeLU activation of the input. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -97,20 +53,6 @@ void nvte_build_activation_table(NVTEActivationTableKind kind, NVTETensor table, */ void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); -/*! \brief Computes the GeLU activation of the input, with a quantization config. - * - * Same as nvte_gelu. With kNVTEQuantizationConfigActivationTable set to a - * kNVTEActivationTableGelu table, BF16 inputs quantized to MXFP8 read the activation - * from the table. - * - * \param[in] input Input tensor for activation. - * \param[in,out] output Output tensor. - * \param[in] quant_config Quantization configuration. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_gelu_v2(const NVTETensor input, NVTETensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream); - /*! \brief Computes the GeLU activation of the grouped input. * If the scaling mode of the grouped output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -132,20 +74,6 @@ void nvte_group_gelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cu */ void nvte_silu(const NVTETensor input, NVTETensor output, cudaStream_t stream); -/*! \brief Computes the SiLU activation of the input, with a quantization config. - * - * Same as nvte_silu. With kNVTEQuantizationConfigActivationTable set to a - * kNVTEActivationTableSilu table, BF16 inputs quantized to MXFP8 read the activation - * from the table. - * - * \param[in] input Input tensor for activation. - * \param[in,out] output Output tensor. - * \param[in] quant_config Quantization configuration. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_silu_v2(const NVTETensor input, NVTETensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream); - /*! \brief Computes the SiLU activation of the grouped input. * If the scaling mode of the grouped output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -277,21 +205,6 @@ void nvte_group_srelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, c void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); -/*! \brief Computes the GeLU activation gradient, with a quantization config. - * - * Same as nvte_dgelu. With kNVTEQuantizationConfigActivationTable set to a - * kNVTEActivationTableDgelu table, BF16 inputs quantized to MXFP8 read the derivative - * from the table. - * - * \param[in] grad Incoming gradient. - * \param[in] input Input tensor for activation. - * \param[in,out] output Output tensor. - * \param[in] quant_config Quantization configuration. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_dgelu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream); - /*! \brief Computes the GeLU activation gradient of the grouped input. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -317,21 +230,6 @@ void nvte_group_dgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor inpu void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); -/*! \brief Computes the SiLU activation gradient, with a quantization config. - * - * Same as nvte_dsilu. With kNVTEQuantizationConfigActivationTable set to a - * kNVTEActivationTableDsilu table, BF16 inputs quantized to MXFP8 read the derivative - * from the table. - * - * \param[in] grad Incoming gradient. - * \param[in] input Input tensor for activation. - * \param[in,out] output Output tensor. - * \param[in] quant_config Quantization configuration. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_dsilu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream); - /*! \brief Computes the SiLU activation gradient of the grouped input. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index bec9847345e..1af93d42919 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -192,23 +192,6 @@ void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor act_inpu NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); -/*! \brief Same as nvte_quantize_dbias_dgelu, with a quantization config. - * - * With kNVTEQuantizationConfigActivationTable set to a kNVTEActivationTableDgelu table, - * BF16 inputs quantized to MXFP8 read the derivative from the table. - * - * \param[in] input Input tensor to be cast. - * \param[in] act_input Activation input tensor. - * \param[in,out] output Output FP8/MXFP8 tensor. - * \param[out] dbias Result of the reduction of the input along columns. - * \param[out] workspace Workspace tensor. - * \param[in] quant_config Quantization configuration. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_quantize_dbias_dgelu_v2(const NVTETensor input, const NVTETensor act_input, - NVTETensor output, NVTETensor dbias, NVTETensor workspace, - const NVTEQuantizationConfig quant_config, cudaStream_t stream); - /*! \brief Computes backward of GeLU operation on the grouped input, then casts to FP8/MXFP8. * Additionally, reduces the result of the GeLU backward along columns. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, @@ -257,23 +240,6 @@ void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor act_inpu NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); -/*! \brief Same as nvte_quantize_dbias_dsilu, with a quantization config. - * - * With kNVTEQuantizationConfigActivationTable set to a kNVTEActivationTableDsilu table, - * BF16 inputs quantized to MXFP8 read the derivative from the table. - * - * \param[in] input Input tensor to be cast. - * \param[in] act_input Activation input tensor. - * \param[in,out] output Output FP8/MXFP8 tensor. - * \param[out] dbias Result of the reduction of the input along columns. - * \param[out] workspace Workspace tensor. - * \param[in] quant_config Quantization configuration. - * \param[in] stream CUDA stream used for the operation. - */ -void nvte_quantize_dbias_dsilu_v2(const NVTETensor input, const NVTETensor act_input, - NVTETensor output, NVTETensor dbias, NVTETensor workspace, - const NVTEQuantizationConfig quant_config, cudaStream_t stream); - /*! \brief Computes backward of SiLU operation on the grouped input, then casts to FP8/MXFP8. * Additionally, reduces the result of the SiLU backward along columns. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 08802873674..9a08b15a8a8 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -436,11 +436,6 @@ enum NVTEQuantizationConfigAttribute { kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath = 9, /*! Whether to use 2D block scaling for MXFP8 */ kNVTEQuantizationConfigMXFP82DQuantization = 10, - /*! Activation lookup table (NVTETensor), built with nvte_build_activation_table for the - * activation being computed. Read by the _v2 activation functions for BF16 inputs - * quantized to MXFP8, and ignored otherwise. See NVTEActivationTableKind. - */ - kNVTEQuantizationConfigActivationTable = 11, kNVTEQuantizationConfigNumAttributes }; @@ -1517,12 +1512,6 @@ class QuantizationConfigWrapper { &val, sizeof(val)); } - /*! \brief Set activation lookup table */ - void set_activation_table(NVTETensor activation_table) { - nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigActivationTable, - &activation_table, sizeof(NVTETensor)); - } - /*! \brief Set whether to use stochastic rounding */ void set_stochastic_rounding(bool stochastic_rounding) { const auto val = static_cast(stochastic_rounding); diff --git a/transformer_engine/common/libtransformer_engine.version b/transformer_engine/common/libtransformer_engine.version index ccd18fc153d..c608cfee04a 100644 --- a/transformer_engine/common/libtransformer_engine.version +++ b/transformer_engine/common/libtransformer_engine.version @@ -7,6 +7,7 @@ transformer_engine::cuda::supports_multicast*; transformer_engine::cuda::stream_priority_range*; transformer_engine::cuda::current_device*; + transformer_engine::activation_table::device_table*; transformer_engine::cuda::cublas_version*; transformer_engine::cuda_driver::get_symbol*; transformer_engine::cuda_driver::ensure_context_exists*; diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 3b9bf0b48d3..66693dca3a1 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -1143,9 +1143,6 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigMXFP82DQuantization: bool_to_uint8(config_.mxfp8_2d_quantization, buf); break; - case kNVTEQuantizationConfigActivationTable: - std::memcpy(buf, &config_.activation_table, attr_size); - break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } @@ -1216,9 +1213,6 @@ void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigMXFP82DQuantization: uint8_to_bool(buf, config_.mxfp8_2d_quantization); break; - case kNVTEQuantizationConfigActivationTable: - std::memcpy(&config_.activation_table, buf, attr_size); - break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } diff --git a/transformer_engine/common/util/activation_table.cuh b/transformer_engine/common/util/activation_table.cuh index d01d310f17c..5a68e9a4341 100644 --- a/transformer_engine/common/util/activation_table.cuh +++ b/transformer_engine/common/util/activation_table.cuh @@ -8,7 +8,9 @@ * \brief GeLU, SiLU and their derivatives of BF16 inputs, read from lookup tables. * * A BF16 input takes only 2^16 values, so these activations can be read from a table - * generated once from the scalar util/math.h functions (nvte_build_activation_table). + * generated once per device from the scalar util/math.h functions. TE builds each table + * on its first use and keeps it for the life of the process (device_table); kernels read + * them for BF16 inputs quantized to MXFP8. NVTE_ACTIVATION_TABLE=0 turns them off. * Outside each table's window the result has a closed form, so every lookup is * branch-free. * @@ -34,7 +36,6 @@ #define TRANSFORMER_ENGINE_UTIL_ACTIVATION_TABLE_CUH_ #include -#include #include #include @@ -46,6 +47,15 @@ namespace transformer_engine { namespace activation_table { +//! Activation read from a table. +enum Kind : int { + kGelu = 0, //!< GeLU (tanh form), BF16 entries + kSilu = 1, //!< SiLU, BF16 entries + kDgelu = 2, //!< GeLU derivative (tanh form), FP32 entries + kDsilu = 3, //!< SiLU derivative, FP32 entries + kNumKinds = 4, +}; + /*! \brief Forward table over the BF16 magnitudes [kLo, kHi]. * * Rows are indexed by magnitude bits from kLowSentinel to kHighSentinel, two entries per @@ -85,25 +95,25 @@ struct DerivativeLayout { static_assert(kStride == 2048 || kStride == 4096, "The packed lookup assumes these strides."); }; -template +template struct Table; //! gelu(x) is RN(x/2) below 0x3b4b, x from 2.75 and -0 from -5.1875. template <> -struct Table : ForwardLayout<0x3b4b, 0x40a5> { +struct Table : ForwardLayout<0x3b4b, 0x40a5> { __device__ static float op(const float x) { return gelu(x, {}); } }; //! silu(x) is RN(x/2) below 0x3bb6, x from 5.96875 and -0 from -89. template <> -struct Table : ForwardLayout<0x3bb6, 0x42b1> { +struct Table : ForwardLayout<0x3bb6, 0x42b1> { __device__ static float op(const float x) { return silu(x, {}); } }; //! [2^-12, 8]. dgelu has saturated, to 1 and to +0, by |x| = 8, until 0.107*x*x overflows at //! 0x6044. template <> -struct Table : DerivativeLayout<0x3980, 0x4100, 0x6044> { +struct Table : DerivativeLayout<0x3980, 0x4100, 0x6044> { __device__ static float op(const float x) { return dgelu(x, {}); } //! dgelu(x) = 0.5 + sqrt(2/pi)*x + O(x^3). __device__ static float below_window(const float x) { return fmaf(0.79788456f, x, 0.5f); } @@ -111,7 +121,7 @@ struct Table : DerivativeLayout<0x3980, 0x4100, 0x604 //! [2^-9, 89]. dsilu has saturated, to 1 and to +0, by |x| = 89, up to inf. template <> -struct Table : DerivativeLayout<0x3B00, 0x42B2, 0x7F80> { +struct Table : DerivativeLayout<0x3B00, 0x42B2, 0x7F80> { __device__ static float op(const float x) { return dsilu(x, {}); } //! dsilu(x) = 0.5 + x/2 + O(x^3). __device__ static float below_window(const float x) { return fmaf(0.5f, x, 0.5f); } @@ -124,39 +134,31 @@ struct ExactDsiluTable : DerivativeLayout<0x3300, 0x42B2, 0x7F80, 4096> { __device__ static float below_window(const float x) { return fmaf(0.5f, x, 0.5f); } }; -/*! \brief Table size in bytes. Each kind has a different size, which lets callers check that - * a table matches the activation it is passed with. - */ -constexpr size_t table_bytes(const NVTEActivationTableKind kind) { +//! Table size in bytes, in the default layout. +constexpr size_t table_bytes(const Kind kind) { switch (kind) { - case kNVTEActivationTableGelu: - return Table::kBytes; - case kNVTEActivationTableSilu: - return Table::kBytes; - case kNVTEActivationTableDgelu: - return Table::kBytes; - case kNVTEActivationTableDsilu: - return Table::kBytes; + case kGelu: + return Table::kBytes; + case kSilu: + return Table::kBytes; + case kDgelu: + return Table::kBytes; + case kDsilu: + return Table::kBytes; default: return 0; } } -static_assert(table_bytes(kNVTEActivationTableGelu) != table_bytes(kNVTEActivationTableSilu) && - table_bytes(kNVTEActivationTableDgelu) != - table_bytes(kNVTEActivationTableDsilu) && - ExactDsiluTable::kBytes != table_bytes(kNVTEActivationTableDsilu), - "Table sizes must identify their kind."); - //! Whether dSiLU tables are built in the exact layout (NVTE_ACTIVATION_TABLE_EXACT_DSILU=1). inline bool use_exact_dsilu_table() { static const bool exact = getenv("NVTE_ACTIVATION_TABLE_EXACT_DSILU"); return exact; } -//! Size of the table that nvte_build_activation_table builds for kind. -inline size_t built_table_bytes(const NVTEActivationTableKind kind) { - if (kind == kNVTEActivationTableDsilu && use_exact_dsilu_table()) { +//! Size of the table that device_table builds for kind. +inline size_t built_table_bytes(const Kind kind) { + if (kind == kDsilu && use_exact_dsilu_table()) { return ExactDsiluTable::kBytes; } return table_bytes(kind); @@ -175,29 +177,23 @@ struct TableRef { TableUse use = kNoTable; }; -/*! \brief An activation table, checked against the kind it must hold. +/*! \brief The current device's table for kind, built on first use. * - * A dSiLU table may have either layout; its size tells which. + * Each table is built once per device and process, then kept until the process exits + * (under 64 KB per device for all of them). The first call may come while the caller's + * stream is being captured into a CUDA graph: the table is then allocated and built outside + * the capture, and is ready before this returns. Thread-safe. */ -inline TableRef checked_table(const Tensor &table, const NVTEActivationTableKind kind) { - const size_t bytes = table_bytes(kind); - NVTE_CHECK(bytes != 0, "Invalid NVTEActivationTableKind (got ", static_cast(kind), ")"); - NVTE_CHECK(table.data.dptr != nullptr, "Activation table is not allocated."); - NVTE_CHECK(table.data.dtype == DType::kByte, "Activation table must be a byte tensor."); - const size_t got = table.data.numel(); - if (kind == kNVTEActivationTableDsilu && got == ExactDsiluTable::kBytes) { - return {table.data.dptr, kExactDsiluTable}; - } - NVTE_CHECK(got == bytes, "Activation table has ", got, " bytes, but activation table kind ", - static_cast(kind), " needs ", bytes, " (was it built for another activation?)"); - return {table.data.dptr, kDefaultTable}; -} +TableRef device_table(Kind kind); + +//! Whether kernels read the tables (NVTE_ACTIVATION_TABLE, default 1). Read on every call. +inline bool tables_enabled() { return getenv("NVTE_ACTIVATION_TABLE", true); } /*! \brief The table kind for OP, if OP is read from a table. */ template constexpr bool kHasTable = false; template -constexpr NVTEActivationTableKind kKindOf = kNVTEActivationTableGelu; +constexpr Kind kKindOf = kGelu; template <> constexpr bool kHasTable> = true; @@ -208,11 +204,11 @@ constexpr bool kHasTable> = true; template <> constexpr bool kHasTable> = true; template <> -constexpr NVTEActivationTableKind kKindOf> = kNVTEActivationTableSilu; +constexpr Kind kKindOf> = kSilu; template <> -constexpr NVTEActivationTableKind kKindOf> = kNVTEActivationTableDgelu; +constexpr Kind kKindOf> = kDgelu; template <> -constexpr NVTEActivationTableKind kKindOf> = kNVTEActivationTableDsilu; +constexpr Kind kKindOf> = kDsilu; /*! \brief Forward activation of x, a BF16 value widened to FP32. * @@ -339,9 +335,9 @@ __device__ __forceinline__ float evaluate(const float x, const void *table) { return OP(x, {}); } else { static_assert(kHasTable, "OP has no table."); - constexpr NVTEActivationTableKind kKind = kKindOf; + constexpr Kind kKind = kKindOf; if constexpr (kUse == kExactDsiluTable) { - static_assert(kKind == kNVTEActivationTableDsilu, "Only dSiLU has an exact table."); + static_assert(kKind == kDsilu, "Only dSiLU has an exact table."); return derivative(x, static_cast(table)); } else if constexpr (Table::kForward) { return forward>(x, static_cast(table)); diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 6506c4c4cc2..893374b3de3 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -243,11 +243,6 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(EpBootstrapBorrowedCommHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(TopkHandler); pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k); -// Activation lookup table to pass for act_type (or its derivative), built once per device; -// nullptr when the activation is not read from a table (see NVTEActivationTableKind). -NVTETensor activation_table_for(NVTE_Activation_Type act_type, bool derivative, DType input_dtype, - JAXX_Scaling_Mode scaling_mode, cudaStream_t stream); - } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/activation.cpp b/transformer_engine/jax/csrc/extensions/activation.cpp index b570c4c7ec4..6325a700d15 100644 --- a/transformer_engine/jax/csrc/extensions/activation.cpp +++ b/transformer_engine/jax/csrc/extensions/activation.cpp @@ -104,26 +104,18 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal } switch (act_type) { - case NVTE_Activation_Type::GELU: { - QuantizationConfigWrapper config; - config.set_activation_table(activation_table_for( - act_type, /*derivative=*/false, static_cast(in_dtype), scaling_mode, stream)); - nvte_gelu_v2(input_tensor.data(), output_tensor.data(), config, stream); + case NVTE_Activation_Type::GELU: + nvte_gelu(input_tensor.data(), output_tensor.data(), stream); break; - } case NVTE_Activation_Type::GEGLU: nvte_geglu(input_tensor.data(), output_tensor.data(), stream); break; case NVTE_Activation_Type::GLU: nvte_glu(input_tensor.data(), output_tensor.data(), stream); break; - case NVTE_Activation_Type::SILU: { - QuantizationConfigWrapper config; - config.set_activation_table(activation_table_for( - act_type, /*derivative=*/false, static_cast(in_dtype), scaling_mode, stream)); - nvte_silu_v2(input_tensor.data(), output_tensor.data(), config, stream); + case NVTE_Activation_Type::SILU: + nvte_silu(input_tensor.data(), output_tensor.data(), stream); break; - } case NVTE_Activation_Type::SWIGLU: nvte_swiglu(input_tensor.data(), output_tensor.data(), stream); break; @@ -182,12 +174,6 @@ Error_Type ActLuInitializeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer int64_t act_enum, JAXX_Scaling_Mode scaling_mode, JAXX_Quantize_Layout quantize_layout, ActivationConfig act_params, bool output_amax_when_no_scaling) { - // Build the activation table outside the warm-up capture below, so that execution only - // reads it. - activation_table_for( - static_cast(act_enum), /*derivative=*/false, - static_cast(convert_ffi_datatype_to_te_dtype(input_buf.element_type())), scaling_mode, - stream); return wrapInStreamCapture(std::function(ActLuFFI), stream, input_buf, scale_buf, amax_buf, output_buf, colwise_output_buf, scale_inv_buf, colwise_scale_inv_buf, updated_amax_buf, act_enum, scaling_mode, quantize_layout, act_params, @@ -395,20 +381,17 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, is_quantize_2x2x(quantize_layout) && act_len == 2), "TE/common does not support delayed scaling for 2x with gated activations."); - QuantizationConfigWrapper table_config; - table_config.set_activation_table(activation_table_for( - act_type, /*derivative=*/true, static_cast(in_dtype), scaling_mode, stream)); if (is_dbias) { switch (act_type) { case NVTE_Activation_Type::GELU: - nvte_quantize_dbias_dgelu_v2(input_tensor.data(), act_input_tensor.data(), - output_tensor.data(), dbias_tensor.data(), - workspace_tensor.data(), table_config, stream); + nvte_quantize_dbias_dgelu(input_tensor.data(), act_input_tensor.data(), + output_tensor.data(), dbias_tensor.data(), + workspace_tensor.data(), stream); break; case NVTE_Activation_Type::SILU: - nvte_quantize_dbias_dsilu_v2(input_tensor.data(), act_input_tensor.data(), - output_tensor.data(), dbias_tensor.data(), - workspace_tensor.data(), table_config, stream); + nvte_quantize_dbias_dsilu(input_tensor.data(), act_input_tensor.data(), + output_tensor.data(), dbias_tensor.data(), + workspace_tensor.data(), stream); break; case NVTE_Activation_Type::RELU: nvte_quantize_dbias_drelu(input_tensor.data(), act_input_tensor.data(), @@ -432,12 +415,10 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, } else { switch (act_type) { case NVTE_Activation_Type::GELU: - nvte_dgelu_v2(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), - table_config, stream); + nvte_dgelu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); break; case NVTE_Activation_Type::SILU: - nvte_dsilu_v2(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), - table_config, stream); + nvte_dsilu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); break; case NVTE_Activation_Type::RELU: nvte_drelu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); @@ -508,12 +489,6 @@ Error_Type DActLuDBiasQuantizeInitializeFFI( Result_Type dbias_buf, Result_Type workspace_buf, JAXX_Scaling_Mode scaling_mode, int64_t act_enum, JAXX_Quantize_Layout quantize_layout, bool is_dbias, ActivationConfig act_params, bool output_amax_when_no_scaling) { - // Build the activation table outside the warm-up capture below, so that execution only - // reads it. - activation_table_for( - static_cast(act_enum), /*derivative=*/true, - static_cast(convert_ffi_datatype_to_te_dtype(input_buf.element_type())), scaling_mode, - stream); return wrapInStreamCapture(std::function(DActLuDBiasQuantizeFFI), stream, input_buf, act_input_buf, scale_buf, amax_buf, output_buf, colwise_output_buf, scale_inv_buf, colwise_scale_inv_buf, updated_amax_buf, dbias_buf, diff --git a/transformer_engine/jax/csrc/extensions/activation_table.cpp b/transformer_engine/jax/csrc/extensions/activation_table.cpp deleted file mode 100644 index a57999e3fb3..00000000000 --- a/transformer_engine/jax/csrc/extensions/activation_table.cpp +++ /dev/null @@ -1,78 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include -#include -#include -#include - -#include "../extensions.h" - -namespace transformer_engine { -namespace jax { - -namespace { - -struct CachedTable { - void *ptr = nullptr; - std::unique_ptr tensor; - bool ready = false; -}; - -NVTETensor activation_table(const NVTEActivationTableKind kind, cudaStream_t stream) { - static std::mutex mutex; - // Never destroyed: the tables live until the process exits. - static auto *tables = new std::map, CachedTable>(); - int device; - NVTE_CHECK_CUDA(cudaGetDevice(&device)); - - std::lock_guard lock(mutex); - CachedTable &table = (*tables)[{device, static_cast(kind)}]; - if (!table.tensor) { - // XLA gives custom calls no persistent device allocation, so the table is allocated - // directly. All tables together take tens of KB per device. - const size_t bytes = nvte_activation_table_size(kind); - NVTE_CHECK_CUDA(cudaMalloc(&table.ptr, bytes)); - table.tensor = - std::make_unique(table.ptr, std::vector{bytes}, DType::kByte); - } - if (!table.ready) { - nvte_build_activation_table(kind, table.tensor->data(), stream); - cudaStreamCaptureStatus capture; - NVTE_CHECK_CUDA(cudaStreamIsCapturing(stream, &capture)); - if (capture == cudaStreamCaptureStatusNone) { - // Later calls may come on other streams, so the build has to be complete. - NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - table.ready = true; - } - // A build issued under graph capture only runs when the graph replays, so the next call - // outside capture builds again. Building is idempotent. - } - return table.tensor->data(); -} - -} // namespace - -NVTETensor activation_table_for(const NVTE_Activation_Type act_type, const bool derivative, - const DType input_dtype, const JAXX_Scaling_Mode scaling_mode, - cudaStream_t stream) { - if (input_dtype != DType::kBFloat16 || scaling_mode != JAXX_Scaling_Mode::MXFP8_1D_SCALING) { - return nullptr; - } - switch (act_type) { - case NVTE_Activation_Type::GELU: - return activation_table(derivative ? kNVTEActivationTableDgelu : kNVTEActivationTableGelu, - stream); - case NVTE_Activation_Type::SILU: - return activation_table(derivative ? kNVTEActivationTableDsilu : kNVTEActivationTableSilu, - stream); - default: - return nullptr; - } -} - -} // namespace jax -} // namespace transformer_engine diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 8a7ae66d5dd..9a3c6314054 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -629,58 +629,6 @@ void philox_unpack(at::PhiloxCudaState arg, int64_t* rng_state_ptr); // extract PhiloxCudaState from CUDA random number generator at::PhiloxCudaState init_philox_state(at::CUDAGeneratorImpl* gen, size_t elts_per_thread); -/*! \brief Activation table of the given kind for the current device. - * - * Built on first use and kept for the lifetime of the process, one per device. - */ -NVTETensor activation_table(NVTEActivationTableKind kind, cudaStream_t stream); - -/*! \brief Whether an activation of input written to output can be read from a table: BF16 - * inputs quantized to MXFP8. - */ -inline bool can_use_activation_table(const NVTETensor input, const NVTETensor output) { - return nvte_tensor_type(input) == kNVTEBFloat16 && - nvte_tensor_scaling_mode(output) == NVTE_MXFP8_1D_SCALING; -} - -/*! \brief An activation entry point that reads its table when it can. */ -template -void activation_with_table(const NVTETensor input, NVTETensor output, cudaStream_t stream) { - QuantizationConfigWrapper config; - if (can_use_activation_table(input, output)) { - config.set_activation_table(activation_table(kKind, stream)); - } - act_func_v2(input, output, config, stream); -} - -/*! \brief An activation-gradient entry point that reads its table when it can. */ -template -void dactivation_with_table(const NVTETensor grad, const NVTETensor input, NVTETensor output, - cudaStream_t stream) { - QuantizationConfigWrapper config; - if (can_use_activation_table(input, output)) { - config.set_activation_table(activation_table(kKind, stream)); - } - dact_func_v2(grad, input, output, config, stream); -} - -/*! \brief A fused dbias + activation-gradient entry point that reads its table when it can. */ -template -void dbias_dactivation_with_table(const NVTETensor grad, const NVTETensor input, NVTETensor output, - NVTETensor dbias, NVTETensor workspace, cudaStream_t stream) { - QuantizationConfigWrapper config; - if (can_use_activation_table(input, output)) { - config.set_activation_table(activation_table(kKind, stream)); - } - dact_dbias_func_v2(grad, input, output, dbias, workspace, config, stream); -} - } // namespace transformer_engine::pytorch namespace std { diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 9293cb3e077..6251872a036 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -247,13 +247,11 @@ py::object dactivation_helper(const at::Tensor& grad_output, const at::Tensor& i /* GELU and variants */ py::object gelu(const at::Tensor& input, py::handle quantizer) { - return activation_helper, nullptr>( - input, quantizer); + return activation_helper(input, quantizer); } py::object dgelu(const at::Tensor& grad, const at::Tensor& input, py::handle quantizer) { - return dactivation_helper, - nullptr>(grad, input, quantizer); + return dactivation_helper(grad, input, quantizer); } py::object glu(const at::Tensor& input, py::handle quantizer) { @@ -322,13 +320,11 @@ py::object dsreglu(const at::Tensor& grad, const at::Tensor& input, py::handle q } /* Silu and variants */ py::object silu(const at::Tensor& input, py::handle quantizer) { - return activation_helper, nullptr>( - input, quantizer); + return activation_helper(input, quantizer); } py::object dsilu(const at::Tensor& grad, const at::Tensor& input, py::handle quantizer) { - return dactivation_helper, - nullptr>(grad, input, quantizer); + return dactivation_helper(grad, input, quantizer); } py::object swiglu(const at::Tensor& input, py::handle quantizer) { diff --git a/transformer_engine/pytorch/csrc/extensions/activation_table.cpp b/transformer_engine/pytorch/csrc/extensions/activation_table.cpp deleted file mode 100644 index 543242405c6..00000000000 --- a/transformer_engine/pytorch/csrc/extensions/activation_table.cpp +++ /dev/null @@ -1,57 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include -#include -#include -#include - -#include "../common.h" - -namespace transformer_engine::pytorch { - -namespace { - -struct CachedTable { - at::Tensor buffer; - std::unique_ptr tensor; - bool ready = false; -}; - -} // namespace - -NVTETensor activation_table(const NVTEActivationTableKind kind, cudaStream_t stream) { - static std::mutex mutex; - // Never destroyed: the tables live until the process exits, and destroying tensors after - // the CUDA context has been torn down is not safe. - static auto *tables = new std::map, CachedTable>(); - const int device = c10::cuda::current_device(); - - std::lock_guard lock(mutex); - CachedTable &table = (*tables)[{device, static_cast(kind)}]; - if (!table.tensor) { - const size_t bytes = nvte_activation_table_size(kind); - table.buffer = at::empty({static_cast(bytes)}, - at::TensorOptions().dtype(at::kByte).device(at::kCUDA, device)); - table.tensor = std::make_unique(table.buffer.data_ptr(), - std::vector{bytes}, DType::kByte); - } - if (!table.ready) { - nvte_build_activation_table(kind, table.tensor->data(), stream); - cudaStreamCaptureStatus capture; - NVTE_CHECK_CUDA(cudaStreamIsCapturing(stream, &capture)); - if (capture == cudaStreamCaptureStatusNone) { - // Later calls may come on other streams, so the build has to be complete. - NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - table.ready = true; - } - // A build issued under graph capture only runs when the graph replays, so the next - // eager call builds again. Building is idempotent. - } - return table.tensor->data(); -} - -} // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/bias.cpp b/transformer_engine/pytorch/csrc/extensions/bias.cpp index 96178f34fff..2dbeb965a1a 100644 --- a/transformer_engine/pytorch/csrc/extensions/bias.cpp +++ b/transformer_engine/pytorch/csrc/extensions/bias.cpp @@ -266,18 +266,12 @@ std::vector dact_dbias( std::vector dbias_dgelu(const at::Tensor &grad_output, const at::Tensor &act_input, py::handle quantizer) { - return dact_dbias( - dbias_dactivation_with_table, - dactivation_with_table, grad_output, act_input, - quantizer); + return dact_dbias(nvte_quantize_dbias_dgelu, nvte_dgelu, grad_output, act_input, quantizer); } std::vector dbias_dsilu(const at::Tensor &grad_output, const at::Tensor &act_input, py::handle quantizer) { - return dact_dbias( - dbias_dactivation_with_table, - dactivation_with_table, grad_output, act_input, - quantizer); + return dact_dbias(nvte_quantize_dbias_dsilu, nvte_dsilu, grad_output, act_input, quantizer); } std::vector dbias_drelu(const at::Tensor &grad_output, const at::Tensor &act_input, From 44f8d05a326b94432946343a44ee15e9fdfa6a96 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 15 Oct 2026 11:28:21 -0700 Subject: [PATCH 05/29] [Common] Make the activation tables opt-in and compute them in the generic kernel Turn the tables off by default: NVTE_ACTIVATION_TABLE=1 enables them, as a new feature whose results differ from today's at -inf (-0 instead of NaN) and for tiny derivative inputs. The generic MXFP8 kernel no longer reads the table from global memory. A lookup there costs it more than the math, so it computes what the table holds from util/math.h instead: the formula with -inf clamped to give -0, and the linear term below the derivative window. Both are branch-free, so the compiler keeps overlapping the math across elements. The output stays bit-identical to the register-tiled kernel, which keeps reading the table from shared memory. Against TE main with the tables on, shapes and modes the register-tiled kernel does not take (dbias, rows not a multiple of 64) go from 0.75-0.96x to 0.94-1.00x. Tests check the computed form against the table specification for every BF16 input. Signed-off-by: Teddy Do --- docs/envvars.rst | 4 +-- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 4 +-- tests/cpp/util/test_activation_table.cu | 31 ++++++++++++++--- .../common/cast/dispatch/quantize.cuh | 2 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 14 +++----- .../common/util/activation_table.cuh | 33 +++++++++++++++++-- 6 files changed, 66 insertions(+), 22 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 66f18018f38..583f885b3a7 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -356,8 +356,8 @@ Kernel Configuration .. envvar:: NVTE_ACTIVATION_TABLE :Type: ``int`` (0 or 1) - :Default: ``1`` - :Description: Read GeLU, SiLU and their derivatives from lookup tables when BF16 inputs are quantized to MXFP8. Each table is built on its first use and kept, once per device and process (under 64 KB per device in total), outside the framework's memory allocator. Results match the computed activations, except that GeLU and SiLU of ``-inf`` give ``-0`` instead of NaN, and derivatives of inputs with ``|x| < 2^-12`` (GeLU) or ``|x| < 2^-9`` (SiLU) are within 1 ulp (GeLU) or 2 ulp (SiLU) of the computed ones. When set to ``0``, the activations are computed. + :Default: ``0`` + :Description: Experimental. When set to ``1``, read GeLU, SiLU and their derivatives from lookup tables when BF16 inputs are quantized to MXFP8. Each table is built on its first use and kept, once per device and process (under 64 KB per device in total), outside the framework's memory allocator. Results match the computed activations, except that GeLU and SiLU of ``-inf`` give ``-0`` instead of NaN, and derivatives of inputs with ``|x| < 2^-12`` (GeLU) or ``|x| < 2^-9`` (SiLU) are within 1 ulp (GeLU) or 2 ulp (SiLU) of the computed ones. When set to ``0``, the activations are computed. Read on every call. .. envvar:: NVTE_ACTIVATION_TABLE_EXACT_DSILU diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index b00b1649738..9996e4f4217 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -11,8 +11,8 @@ // specialized cast kernels). dbias is compared against a host sum in the order the kernel // accumulates it. // -// With NVTE_ACTIVATION_TABLE=0 the activation is the scalar util/math.h function. With the tables -// (the default), BF16 inputs follow the table specification (see util/activation_table.cuh): -inf +// Without the tables (NVTE_ACTIVATION_TABLE=0, the default) the activation is the scalar util/math.h +// function. With NVTE_ACTIVATION_TABLE=1, BF16 inputs follow the table specification (see util/activation_table.cuh): -inf // gives -0 in the forward functions, and the derivatives below the table window are the linear // Taylor term. Other input types do not use the tables. diff --git a/tests/cpp/util/test_activation_table.cu b/tests/cpp/util/test_activation_table.cu index 3b097ec8814..728a6570c45 100644 --- a/tests/cpp/util/test_activation_table.cu +++ b/tests/cpp/util/test_activation_table.cu @@ -110,6 +110,10 @@ __global__ void lookup_kernel(const void *table, const uint32_t dsilu_window_low namespace at = transformer_engine::activation_table; const float x = __uint_as_float(bits << 16); looked_up[bits] = at::evaluate::value>(x, table); + if constexpr (kUse == at::kComputedTable && (kKind == kGelu || kKind == kSilu)) { + // The computed forward is FP32; every caller rounds it to BF16, as the table already is. + looked_up[bits] = __bfloat162float(__float2bfloat16_rn(looked_up[bits])); + } expected[bits] = spec(kKind, bits, dsilu_window_low); math_h[bits] = (kKind == kDgelu || kKind == kDsilu) ? scalar_math(kKind, x) @@ -130,7 +134,9 @@ int64_t ordered(const float f) { return i < 0 ? -static_cast(i & 0x7fffffff) : i; } -template +// kComputed: check what the generic kernel computes in place of the table (kComputedTable, or +// util/math.h itself for the exact dSiLU table) instead of the table lookup. +template void check_table(const int max_ulp_below_window) { namespace at = transformer_engine::activation_table; const at::TableRef table = at::device_table(kKind); @@ -142,12 +148,14 @@ void check_table(const int max_ulp_below_window) { const uint32_t dsilu_window_low = exact ? kExactDsiluWindowLow : kDsiluWindowLow; if (exact) { if constexpr (kKind == kDsilu) { - lookup_kernel<<>>( - table.ptr, dsilu_window_low, d, d + kNumInputs, d + 2 * kNumInputs); + lookup_kernel + <<>>(table.ptr, dsilu_window_low, d, d + kNumInputs, + d + 2 * kNumInputs); } } else { - lookup_kernel<<>>( - table.ptr, dsilu_window_low, d, d + kNumInputs, d + 2 * kNumInputs); + lookup_kernel + <<>>(table.ptr, dsilu_window_low, d, d + kNumInputs, + d + 2 * kNumInputs); } ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); std::vector h(3 * kNumInputs); @@ -190,6 +198,18 @@ TEST(ActivationTableTest, DgeluMatchesSpecification) { TEST(ActivationTableTest, DsiluMatchesSpecification) { check_table(exact_dsilu() ? 0 : 2); } +TEST(ActivationTableTest, ComputedGeluMatchesSpecification) { + check_table(-1); +} +TEST(ActivationTableTest, ComputedSiluMatchesSpecification) { + check_table(-1); +} +TEST(ActivationTableTest, ComputedDgeluMatchesSpecification) { + check_table(1); +} +TEST(ActivationTableTest, ComputedDsiluMatchesSpecification) { + check_table(exact_dsilu() ? 0 : 2); +} TEST(ActivationTableTest, OneTablePerDevice) { namespace at = transformer_engine::activation_table; @@ -216,6 +236,7 @@ TEST(ActivationTableTest, OneTablePerDevice) { TEST(ActivationTableDeathTest, FirstUseInsideGraphCapture) { GTEST_FLAG_SET(death_test_style, "threadsafe"); const auto body = [] { + setenv("NVTE_ACTIVATION_TABLE", "1", 1); const std::vector shape{256, 256}; test::Tensor x("x", shape, transformer_engine::DType::kBFloat16); test::fillUniform(&x); diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 665af3823c6..80de7098ba5 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -39,7 +39,7 @@ namespace transformer_engine { namespace dispatch { // The activation table to read for OP, if any. Tables apply to BF16 inputs of the activations -// that have one, unless NVTE_ACTIVATION_TABLE=0. The first call per device builds the table. +// that have one, when NVTE_ACTIVATION_TABLE=1. The first call per device builds the table. template activation_table::TableRef activation_lut_for(const DType input_dtype) { if constexpr (activation_table::kHasTable) { diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index f642ee9f5e6..27178716630 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -1051,19 +1051,15 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, noop_ptr, workspace_ptr, amax_ptr, num_rows, num_cols, scale_stride_rowwise, scale_stride_colwise, activation_lut.ptr); }; + using activation_table::kComputedTable; using activation_table::kDefaultTable; - using activation_table::kExactDsiluTable; using activation_table::kNoTable; + // This kernel computes what the table holds rather than reading it: a + // lookup from global memory costs it more than the math (E3). The exact + // dSiLU table holds util/math.h itself. if constexpr (kTableCapable) { if (activation_lut.use == kDefaultTable) { - launch_kernel(std::integral_constant{}); - } else if (activation_lut.use == kExactDsiluTable) { - if constexpr (activation_table::kKindOf == - activation_table::kDsilu) { - launch_kernel(std::integral_constant{}); - } else { - NVTE_ERROR("Only dSiLU has an exact activation table."); - } + launch_kernel(std::integral_constant{}); } else { launch_kernel(std::integral_constant{}); } diff --git a/transformer_engine/common/util/activation_table.cuh b/transformer_engine/common/util/activation_table.cuh index 5a68e9a4341..405bce3fbd1 100644 --- a/transformer_engine/common/util/activation_table.cuh +++ b/transformer_engine/common/util/activation_table.cuh @@ -10,7 +10,7 @@ * A BF16 input takes only 2^16 values, so these activations can be read from a table * generated once per device from the scalar util/math.h functions. TE builds each table * on its first use and keeps it for the life of the process (device_table); kernels read - * them for BF16 inputs quantized to MXFP8. NVTE_ACTIVATION_TABLE=0 turns them off. + * them for BF16 inputs quantized to MXFP8 when NVTE_ACTIVATION_TABLE=1 (off by default). * Outside each table's window the result has a closed form, so every lookup is * branch-free. * @@ -169,6 +169,9 @@ enum TableUse : int { kNoTable = 0, //!< Compute OP with util/math.h. kDefaultTable = 1, //!< Table. kExactDsiluTable = 2, //!< ExactDsiluTable. + //! No table read: compute with util/math.h what Table holds. Gives the bytes of + //! kDefaultTable, without the memory latency of a lookup in global memory. + kComputedTable = 3, }; //! An activation table passed to a kernel. @@ -186,8 +189,8 @@ struct TableRef { */ TableRef device_table(Kind kind); -//! Whether kernels read the tables (NVTE_ACTIVATION_TABLE, default 1). Read on every call. -inline bool tables_enabled() { return getenv("NVTE_ACTIVATION_TABLE", true); } +//! Whether kernels read the tables (NVTE_ACTIVATION_TABLE, default 0). Read on every call. +inline bool tables_enabled() { return getenv("NVTE_ACTIVATION_TABLE", false); } /*! \brief The table kind for OP, if OP is read from a table. */ template @@ -333,6 +336,30 @@ template __device__ __forceinline__ float evaluate(const float x, const void *table) { if constexpr (kUse == kNoTable) { return OP(x, {}); + } else if constexpr (kUse == kComputedTable) { + // Branch-free on purpose: a branch around OP keeps the compiler from overlapping the + // latency of OP across elements, which costs more than computing OP for every element. + static_assert(kHasTable, "OP has no table."); + using T = Table>; + if constexpr (T::kForward) { + // Outside the window the table's closed forms equal util/math.h rounded to BF16, + // except at -inf, where the formula gives NaN and the table -0. util/math.h gives -0 + // for every finite x below -2^100, so clamp the input there; NaN fails the comparison + // and stays NaN. + return OP(x < -0x1p100f ? -0x1p100f : x, {}); + } else { + // Above the window util/math.h already gives the saturated value or NaN. x is a BF16 + // value, so |x| < the window's first magnitude is one float compare; NaN compares false. + constexpr uint32_t kWindowLow = T::kLoBits << 16; + const float computed = OP(x, {}); + const float linear = T::below_window(x); + float result; + asm("{\n\t.reg .pred below;\n\t.reg .f32 magnitude;\n\tabs.f32 magnitude, %3;\n\t" + "setp.lt.f32 below, magnitude, %4;\n\tselp.f32 %0, %1, %2, below;\n\t}" + : "=f"(result) + : "f"(linear), "f"(computed), "f"(x), "f"(__uint_as_float(kWindowLow))); + return result; + } } else { static_assert(kHasTable, "OP has no table."); constexpr Kind kKind = kKindOf; From cbbb29826cf74fe6c8a27f9b7a7586e54b95f860 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 15 Oct 2026 14:03:49 -0700 Subject: [PATCH 06/29] [Common] Extend the register-tiled MXFP8 activation kernel to row and column tails The register-tiled kernel took only rows % 64 == 0 and cols % 256 == 0; every other shape, such as MoE token counts, fell back to the generic kernel. It now takes any number of rows and any cols % 16 == 0, the shapes the generic kernel takes. Out-of-range elements load as 0, and the kernel writes the scales the generic kernel writes for its whole 64x64 tiles, so the output, scale padding included, stays byte-identical. Bands inside the tensor run without bounds checks; only bands at its edges pay for them. The derivatives use this kernel for every shape: with the band body behind the edge branch, the compiler schedules the loads better (GB200, 32768x7168 dGeLU 239 -> 220 us). GB200, tables on, normal data, against TE main: 993x7200 1.40-1.86x (was 0.95-1.00x through the generic kernel), 3517x7168 1.97-2.24x, regular shapes unchanged for the forward and 6-7% faster for the derivatives. The test compares the two kernels on tail shapes with every output byte pre-filled, so both must write the same set of bytes. Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 18 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 4 +- .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 499 ++++++++++-------- 3 files changed, 302 insertions(+), 219 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index 9996e4f4217..49c8a1a3118 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -394,6 +394,15 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, Tensor("generic", shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING)}; for (int i = 0; i < 2; ++i) { outs[i].set_with_gemm_swizzled_scales(swizzled); + // Fill every output byte, padding included, with a marker, so the two kernels must write the + // same set of bytes, not just the same values. + for (const auto param : {kNVTERowwiseData, kNVTEColumnwiseData, kNVTERowwiseScaleInv, + kNVTEColumnwiseScaleInv}) { + const NVTEBasicTensor b = nvte_get_tensor_param(outs[i].data(), param); + size_t bytes = b.data_ptr == nullptr ? 0 : 1; + for (size_t d = 0; d < b.shape.ndim; ++d) bytes *= b.shape.data[d]; + if (bytes != 0) ASSERT_EQ(cudaMemset(b.data_ptr, 0xA5, bytes), cudaSuccess); + } const ScopedEnv tables("NVTE_ACTIVATION_TABLE", "1"); const ScopedEnv regtile("NVTE_MXFP8_REGTILE", i == 0 ? "1" : "0"); switch (mode) { @@ -537,5 +546,12 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Bool(), ::testing::Values(std::pair{1024, 2048}, std::pair{4160, 2048}, - std::pair{128, 256})), + std::pair{128, 256}, + // Row and column tails; the last column tile ends + // 32 or 16 columns into the tile, mid MX group for 16. + std::pair{993, 7200}, + std::pair{1, 16}, + std::pair{33, 272}, + std::pair{100, 2064}, + std::pair{4127, 528})), regtile_case_name); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 27178716630..f04ac491bc9 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -967,8 +967,8 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, // regtile::quantize replaces the generic kernel below for the requests it // implements: bidimensional BF16 -> E4M3 with fused GeLU, SiLU, dGeLU or dSiLU read - // from an activation table, no dbias, rows % 64 == 0, cols % 256 == 0, 16-byte - // aligned inputs (see regtile::kImplements and regtile::can_use). It keeps the tile + // from an activation table, no dbias, cols % 16 == 0, 16-byte aligned inputs (see + // regtile::kImplements and regtile::can_use). It keeps the tile // in registers instead of staging it through shared memory and writes the same // bytes the generic kernel writes with the same table. if constexpr (std::is_same_v && std::is_same_v && diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh index 9a1c2a32f66..698095d7b51 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -189,12 +189,14 @@ struct TileConfig { // Register-resident tiling: the 32x256 tile stays in registers; only the // colwise partials go through shared memory. -template +// kTails: the grid may cover elements past M or K (see shape_supported). Without tails the +// bounds checks compile away, which matters at the kernel's register budget. +template __device__ __forceinline__ void quantize_regtile( const uint4* __restrict__ input, const uint4* __restrict__ act_input, unsigned char* __restrict__ rowwise_out, unsigned char* __restrict__ rowwise_scales, - unsigned char* __restrict__ colwise_out, unsigned char* __restrict__ colwise_scales, int K, - int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* table) { + unsigned char* __restrict__ colwise_out, unsigned char* __restrict__ colwise_scales, int M, + int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* table) { using C = TileConfig; constexpr int kWarpsPerCta = C::kWarpsPerCta, kRowsPerWarp = C::kRowsPerWarp, kRowsInFlight = C::kRowsInFlight; @@ -210,6 +212,14 @@ __device__ __forceinline__ void quantize_regtile( const int warp = tid >> 5; const int col0 = blockIdx.x * kTileCols; const int vecs_per_row = K >> 3; // uint4 (8 bf16 values) per row + // Tails: the last column tile may end inside the tile, and row bands may end past M. Out-of-range + // elements load as 0, which every activation maps to 0, so they never raise an amax. The + // generic kernel writes scales for its whole 64x64 tiles, and so does this one: rowwise scales + // for every row of the grid's bands, for the groups that start inside the tensor; colwise + // scales for the columns below K rounded up to 64. + const bool full_width = !kTails || col0 + kTileCols <= K; + const bool lane_in_cols = col0 + kColsPerLane * lane < K; + const int colwise_scale_cols = (K + 63) / 64 * 64; // The table copy overlaps the first input loads; see the wait below. load_table_async(table_smem, &table_barrier, table, tid); @@ -219,52 +229,63 @@ __device__ __forceinline__ void quantize_regtile( // Row blocks are taken grid-strided, so the CTAs on the same step cover one // contiguous row band. const int row0 = (it * gridDim.y + blockIdx.y) * kRowsPerMxBlock; - // This thread's kRowsPerWarp x kColsPerLane sub-tile, as BF16 pairs. - unsigned tile_words[kRowsPerWarp * kWordsPerLane]; - const size_t load_offset = (size_t)(row0 + warp * kRowsPerWarp) * vecs_per_row + (col0 >> 3) + - (size_t)kVecLoadsPerRow * lane; - const uint4* input_ptr = input + load_offset; - const uint4* act_input_ptr = act_input + load_offset; - - // ---- load + activate, kRowsInFlight rows in flight at a time -------------------- -#pragma unroll - for (int h = 0; h < kRowsPerWarp / kRowsInFlight; ++h) { - uint4 x_vec[kRowsInFlight][kVecLoadsPerRow], grad_vec[kRowsInFlight][kVecLoadsPerRow]; - asm volatile("" ::: "memory"); -#pragma unroll - for (int t = 0; t < kRowsInFlight; ++t) + // The band's work. A band inside the tensor runs it without bounds checks; only bands at + // the edges of the tensor (kEdge) pay for them. + const auto band = [&](auto edge) { + constexpr bool kEdge = decltype(edge)::value; + const int warp_row0 = row0 + warp * kRowsPerWarp; + // Whether this lane's elements of row warp_row0 + j are inside the tensor. + const auto in_bounds = [&](int j) { return !kEdge || (lane_in_cols && warp_row0 + j < M); }; + // This thread's kRowsPerWarp x kColsPerLane sub-tile, as BF16 pairs. + unsigned tile_words[kRowsPerWarp * kWordsPerLane]; + const size_t load_offset = (size_t)(row0 + warp * kRowsPerWarp) * vecs_per_row + (col0 >> 3) + + (size_t)kVecLoadsPerRow * lane; + const uint4* input_ptr = input + load_offset; + const uint4* act_input_ptr = act_input + load_offset; + + // ---- load + activate, kRowsInFlight rows in flight at a time -------------------- #pragma unroll - for (int n = 0; n < kVecLoadsPerRow; ++n) - x_vec[t][n] = act_input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n]; - if constexpr (IS_DACT) { + for (int h = 0; h < kRowsPerWarp / kRowsInFlight; ++h) { + uint4 x_vec[kRowsInFlight][kVecLoadsPerRow], grad_vec[kRowsInFlight][kVecLoadsPerRow]; + asm volatile("" ::: "memory"); #pragma unroll for (int t = 0; t < kRowsInFlight; ++t) #pragma unroll for (int n = 0; n < kVecLoadsPerRow; ++n) - grad_vec[t][n] = input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n]; - } - // Keep the load group's loads together, between the two empty asm statements: issued - // back to back, their latencies overlap. Otherwise register pressure leads the compiler - // to spread them between the previous group's table lookups, each then stalling on its - // own. - asm volatile("" ::: "memory"); - if (it == 0 && h == 0) ptx::mbarrier_wait_parity(&table_barrier, 0); - // Convert the load group into the BF16 pairs the epilogue quantizes. - if constexpr (!IS_DACT) { + x_vec[t][n] = in_bounds(h * kRowsInFlight + t) + ? act_input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n] + : make_uint4(0, 0, 0, 0); + if constexpr (IS_DACT) { +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) #pragma unroll - for (int t = 0; t < kRowsInFlight; ++t) { + for (int n = 0; n < kVecLoadsPerRow; ++n) + grad_vec[t][n] = in_bounds(h * kRowsInFlight + t) + ? input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n] + : make_uint4(0, 0, 0, 0); + } + // Keep the load group's loads together, between the two empty asm statements: issued + // back to back, their latencies overlap. Otherwise register pressure leads the compiler + // to spread them between the previous group's table lookups, each then stalling on its + // own. + asm volatile("" ::: "memory"); + if (it == 0 && h == 0) ptx::mbarrier_wait_parity(&table_barrier, 0); + // Convert the load group into the BF16 pairs the epilogue quantizes. + if constexpr (!IS_DACT) { +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) { #pragma unroll - for (int n = 0; n < kVecLoadsPerRow; ++n) { - const unsigned* x_words = (const unsigned*)&x_vec[t][n]; - const int tile_word_base = (h * kRowsInFlight + t) * kWordsPerLane + n * 4; + for (int n = 0; n < kVecLoadsPerRow; ++n) { + const unsigned* x_words = (const unsigned*)&x_vec[t][n]; + const int tile_word_base = (h * kRowsInFlight + t) * kWordsPerLane + n * 4; #pragma unroll - for (int m = 0; m < 4; ++m) { - tile_words[tile_word_base + m] = activation_table::forward_2x( - x_words[m], reinterpret_cast(table_smem)); + for (int m = 0; m < 4; ++m) { + tile_words[tile_word_base + m] = activation_table::forward_2x( + x_words[m], reinterpret_cast(table_smem)); + } } } - } - } else { + } else { #define MXFP8_DACT_WORD(T, M) \ { \ const ptx::floatx2 product = \ @@ -280,187 +301,218 @@ __device__ __forceinline__ void quantize_regtile( MXFP8_DACT_WORD(T, 2) \ MXFP8_DACT_WORD(T, 3) \ } - static_assert(kRowsInFlight == 2, "A load group is two rows."); - MXFP8_DACT_ROW(0) - MXFP8_DACT_ROW(1) + static_assert(kRowsInFlight == 2, "A load group is two rows."); + MXFP8_DACT_ROW(0) + MXFP8_DACT_ROW(1) #undef MXFP8_DACT_ROW #undef MXFP8_DACT_WORD + } } - } - // The rowwise half does not depend on the colwise scales, so it runs before - // the cross-warp barrier (see TileConfig). - const size_t store_offset = - (size_t)(row0 + warp * kRowsPerWarp) * K + col0 + kColsPerLane * lane; - auto row_scale = [&](const unsigned* row, int j) { - unsigned amax = - abs_max_bf16x2(abs_max_bf16x2(row[0], row[1]), abs_max_bf16x2(row[2], row[3])); + // The rowwise half does not depend on the colwise scales, so it runs before + // the cross-warp barrier (see TileConfig). + const size_t store_offset = + (size_t)(row0 + warp * kRowsPerWarp) * K + col0 + kColsPerLane * lane; + auto row_scale = [&](const unsigned* row, int j) { + unsigned amax = + abs_max_bf16x2(abs_max_bf16x2(row[0], row[1]), abs_max_bf16x2(row[2], row[3])); #pragma unroll - for (int m = 4; m < kWordsPerLane; m += 2) - amax = abs_max_bf16x2(amax, abs_max_bf16x2(row[m], row[m + 1])); - // A 32-value rowwise block spans kLanesPerMxGroup lanes. + for (int m = 4; m < kWordsPerLane; m += 2) + amax = abs_max_bf16x2(amax, abs_max_bf16x2(row[m], row[m + 1])); + // A 32-value rowwise block spans kLanesPerMxGroup lanes. #pragma unroll - for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) - amax = abs_max_bf16x2(amax, __shfl_xor_sync(0xffffffffu, amax, lane_mask)); - // Fold the two halves with one PRMT and a packed max, which leaves the - // magnitude in both halves; the high half is then the FP32 bit pattern. - amax = abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); - const unsigned scale_byte = amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); - if ((lane & (kLanesPerMxGroup - 1)) == 0) - row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + (lane / kLanesPerMxGroup)] = - (unsigned char)scale_byte; - return e8m0_to_bf16x2_reciprocal(scale_byte); - }; - // On the 128-thread activation tile, two rows share one butterfly and conversion. - auto row_amax = [&](const unsigned* row) { - unsigned amax = - abs_max_bf16x2(abs_max_bf16x2(row[0], row[1]), abs_max_bf16x2(row[2], row[3])); + for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) + amax = abs_max_bf16x2(amax, __shfl_xor_sync(0xffffffffu, amax, lane_mask)); + // Fold the two halves with one PRMT and a packed max, which leaves the + // magnitude in both halves; the high half is then the FP32 bit pattern. + amax = abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); + const unsigned scale_byte = amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); + if ((lane & (kLanesPerMxGroup - 1)) == 0) + row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + (lane / kLanesPerMxGroup)] = + (unsigned char)scale_byte; + return e8m0_to_bf16x2_reciprocal(scale_byte); + }; + // On the 128-thread activation tile, two rows share one butterfly and conversion. + auto row_amax = [&](const unsigned* row) { + unsigned amax = + abs_max_bf16x2(abs_max_bf16x2(row[0], row[1]), abs_max_bf16x2(row[2], row[3])); #pragma unroll - for (int m = 4; m < kWordsPerLane; m += 2) - amax = abs_max_bf16x2(amax, abs_max_bf16x2(row[m], row[m + 1])); - return abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); - }; - auto row_scale_pair = [&](const unsigned* row_a, const unsigned* row_b, int j, - unsigned& row_scale_rcp0, unsigned& row_scale_rcp1) { - unsigned amax_pair = __byte_perm(row_amax(row_a), row_amax(row_b), 0x5410) & 0x7fff7fffu; + for (int m = 4; m < kWordsPerLane; m += 2) + amax = abs_max_bf16x2(amax, abs_max_bf16x2(row[m], row[m + 1])); + return abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); + }; + auto row_scale_pair = [&](const unsigned* row_a, const unsigned* row_b, int j, + unsigned& row_scale_rcp0, unsigned& row_scale_rcp1) { + unsigned amax_pair = __byte_perm(row_amax(row_a), row_amax(row_b), 0x5410) & 0x7fff7fffu; #pragma unroll - for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) - amax_pair = abs_max_bf16x2(amax_pair, __shfl_xor_sync(0xffffffffu, amax_pair, lane_mask)); - const ptx::floatx2 amax_2x = widen_bf16x2(amax_pair); - const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); - const unsigned scale_byte0 = scale_pair & 0xffu, scale_byte1 = (scale_pair >> 8) & 0xffu; - if ((lane & (kLanesPerMxGroup - 1)) == 0) { - unsigned char* scale_dst = row_scale_bytes + (warp * kRowsPerWarp + j) * kMxGroupsPerRow + - (lane / kLanesPerMxGroup); - scale_dst[0] = (unsigned char)scale_byte0; - scale_dst[kMxGroupsPerRow] = (unsigned char)scale_byte1; - } - row_scale_rcp0 = e8m0_to_bf16x2_reciprocal(scale_byte0); - row_scale_rcp1 = e8m0_to_bf16x2_reciprocal(scale_byte1); - }; - // Colwise amax: per-warp partial in registers, cross-warp fold in shared. - auto publish_col_amax_partials = [&]() { - unsigned col_amax[kWordsPerLane]; + for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) + amax_pair = abs_max_bf16x2(amax_pair, __shfl_xor_sync(0xffffffffu, amax_pair, lane_mask)); + const ptx::floatx2 amax_2x = widen_bf16x2(amax_pair); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + const unsigned scale_byte0 = scale_pair & 0xffu, scale_byte1 = (scale_pair >> 8) & 0xffu; + if ((lane & (kLanesPerMxGroup - 1)) == 0) { + unsigned char* scale_dst = row_scale_bytes + (warp * kRowsPerWarp + j) * kMxGroupsPerRow + + (lane / kLanesPerMxGroup); + scale_dst[0] = (unsigned char)scale_byte0; + scale_dst[kMxGroupsPerRow] = (unsigned char)scale_byte1; + } + row_scale_rcp0 = e8m0_to_bf16x2_reciprocal(scale_byte0); + row_scale_rcp1 = e8m0_to_bf16x2_reciprocal(scale_byte1); + }; + // Colwise amax: per-warp partial in registers, cross-warp fold in shared. + auto publish_col_amax_partials = [&]() { + unsigned col_amax[kWordsPerLane]; #pragma unroll - for (int m = 0; m < kWordsPerLane; ++m) col_amax[m] = tile_words[m]; + for (int m = 0; m < kWordsPerLane; ++m) col_amax[m] = tile_words[m]; #pragma unroll - for (int j = 1; j < kRowsPerWarp; ++j) + for (int j = 1; j < kRowsPerWarp; ++j) #pragma unroll - for (int m = 0; m < kWordsPerLane; ++m) - col_amax[m] = abs_max_bf16x2(col_amax[m], tile_words[j * kWordsPerLane + m]); + for (int m = 0; m < kWordsPerLane; ++m) + col_amax[m] = abs_max_bf16x2(col_amax[m], tile_words[j * kWordsPerLane + m]); #pragma unroll - for (int n = 0; n < kVecLoadsPerRow; ++n) - *(uint4*)(&col_amax_partials[warp][kWordsPerLane * lane + 4 * n]) = - *(const uint4*)(col_amax + 4 * n); - }; - if constexpr (C::kThreadsPerCta == 128) { - unsigned char* row_out_ptr = rowwise_out + store_offset; + for (int n = 0; n < kVecLoadsPerRow; ++n) + *(uint4*)(&col_amax_partials[warp][kWordsPerLane * lane + 4 * n]) = + *(const uint4*)(col_amax + 4 * n); + }; + if constexpr (C::kThreadsPerCta == 128) { + unsigned char* row_out_ptr = rowwise_out + store_offset; #pragma unroll - for (int j = 0; j < kRowsPerWarp; j += 2) { - const unsigned* row_a = tile_words + j * kWordsPerLane; - const unsigned* row_b = row_a + kWordsPerLane; - unsigned row_scale_rcp0, row_scale_rcp1; - row_scale_pair(row_a, row_b, j, row_scale_rcp0, row_scale_rcp1); - unsigned row_out[kWordsPerLane / 2]; + for (int j = 0; j < kRowsPerWarp; j += 2) { + const unsigned* row_a = tile_words + j * kWordsPerLane; + const unsigned* row_b = row_a + kWordsPerLane; + unsigned row_scale_rcp0, row_scale_rcp1; + row_scale_pair(row_a, row_b, j, row_scale_rcp0, row_scale_rcp1); + unsigned row_out[kWordsPerLane / 2]; #pragma unroll - for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = - quantize_e4m3x4(row_a[2 * m], row_a[2 * m + 1], row_scale_rcp0, row_scale_rcp0); - store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); + for (int m = 0; m < kWordsPerLane / 2; ++m) + row_out[m] = + quantize_e4m3x4(row_a[2 * m], row_a[2 * m + 1], row_scale_rcp0, row_scale_rcp0); + if (in_bounds(j)) store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); #pragma unroll - for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = - quantize_e4m3x4(row_b[2 * m], row_b[2 * m + 1], row_scale_rcp1, row_scale_rcp1); - store_quantized(row_out_ptr + (unsigned)(j + 1) * (unsigned)K, row_out); - } - } else { - unsigned char* row_out_ptr = rowwise_out + store_offset; + for (int m = 0; m < kWordsPerLane / 2; ++m) + row_out[m] = + quantize_e4m3x4(row_b[2 * m], row_b[2 * m + 1], row_scale_rcp1, row_scale_rcp1); + if (in_bounds(j + 1)) + store_quantized(row_out_ptr + (unsigned)(j + 1) * (unsigned)K, row_out); + } + } else { + unsigned char* row_out_ptr = rowwise_out + store_offset; #pragma unroll - for (int j = 0; j < kRowsPerWarp; ++j) { - const unsigned* v = tile_words + j * kWordsPerLane; - const unsigned row_scale_rcp = row_scale(v, j); - unsigned row_out[kWordsPerLane / 2]; + for (int j = 0; j < kRowsPerWarp; ++j) { + const unsigned* v = tile_words + j * kWordsPerLane; + const unsigned row_scale_rcp = row_scale(v, j); + unsigned row_out[kWordsPerLane / 2]; #pragma unroll - for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = quantize_e4m3x4(v[2 * m], v[2 * m + 1], row_scale_rcp, row_scale_rcp); - store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); + for (int m = 0; m < kWordsPerLane / 2; ++m) + row_out[m] = quantize_e4m3x4(v[2 * m], v[2 * m + 1], row_scale_rcp, row_scale_rcp); + if (in_bounds(j)) store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); + } } - } - publish_col_amax_partials(); - __syncthreads(); - if (tid < kTileCols / 2) { - unsigned col_amax = col_amax_partials[0][tid]; + publish_col_amax_partials(); + __syncthreads(); + if (tid < kTileCols / 2) { + unsigned col_amax = col_amax_partials[0][tid]; #pragma unroll - for (int w = 1; w < kWarpsPerCta; ++w) - col_amax = abs_max_bf16x2(col_amax, col_amax_partials[w][tid]); - col_amax &= 0x7fff7fffu; - const ptx::floatx2 amax_2x = widen_bf16x2(col_amax); - const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); - if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - // The colwise swizzle is the rowwise one with rows and columns swapped, - // so the two columns of this pair land in different 16-byte rows. The - // grid covers M exactly: gridDim.y * iters row blocks of 32. - const size_t row_block_tiles = - DIVUP(static_cast(gridDim.y) * iters * kRowsPerMxBlock, size_t{128}); - const size_t col = col0 + 2 * tid; - colwise_scales[swizzle::gemm_swizzled_scale_idx(col, row0 / 32, row_block_tiles)] = - static_cast(scale_pair); - colwise_scales[swizzle::gemm_swizzled_scale_idx(col + 1, row0 / 32, row_block_tiles)] = - static_cast(scale_pair >> 8); - } else { - *(unsigned short*)(colwise_scales + (size_t)(row0 / 32) * colwise_scale_stride + col0 + - 2 * tid) = (unsigned short)scale_pair; - } - const ptx::bf16x2 scale_rcp = ptx::exp2f_rcp_2x_per_lane(__byte_perm(scale_pair, 0, 0x4140)); - col_scale_rcp_smem[tid] = reinterpret_cast(scale_rcp); - } - auto drain_row_scales = [&]() { - constexpr int drain_tid_base = C::kDrainTidBase; - // A row's eight group scales are eight contiguous, 8-byte-aligned bytes of - // rowwise_scales, so one warp drains a 32-row tile with one STG.64 per lane. - // In the GEMM-swizzled layout they are two 4-byte runs in adjacent tiles. - if (tid >= drain_tid_base && tid < drain_tid_base + kRowsPerMxBlock) { - const int r = tid - drain_tid_base; - if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - constexpr size_t kTileBytes = - swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_X * swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_Y; - const size_t idx = swizzle::gemm_swizzled_scale_idx(row0 + r, col0 >> 5, K / 128); - const uint2 bytes = *(const uint2*)(row_scale_bytes + r * kMxGroupsPerRow); - *(unsigned*)(rowwise_scales + idx) = bytes.x; - *(unsigned*)(rowwise_scales + idx + kTileBytes) = bytes.y; - } else { - *(uint2*)(rowwise_scales + (size_t)(row0 + r) * rowwise_scale_stride + (col0 >> 5)) = - *(const uint2*)(row_scale_bytes + r * kMxGroupsPerRow); + for (int w = 1; w < kWarpsPerCta; ++w) + col_amax = abs_max_bf16x2(col_amax, col_amax_partials[w][tid]); + col_amax &= 0x7fff7fffu; + const ptx::floatx2 amax_2x = widen_bf16x2(col_amax); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + if (!kEdge || col0 + 2 * tid < colwise_scale_cols) { + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + // The colwise swizzle is the rowwise one with rows and columns swapped, + // so the two columns of this pair land in different 16-byte rows. + const size_t row_block_tiles = DIVUP(static_cast(M), size_t{128}); + const size_t col = col0 + 2 * tid; + colwise_scales[swizzle::gemm_swizzled_scale_idx(col, row0 / 32, row_block_tiles)] = + static_cast(scale_pair); + colwise_scales[swizzle::gemm_swizzled_scale_idx(col + 1, row0 / 32, row_block_tiles)] = + static_cast(scale_pair >> 8); + } else { + *(unsigned short*)(colwise_scales + (size_t)(row0 / 32) * colwise_scale_stride + col0 + + 2 * tid) = (unsigned short)scale_pair; + } } + const ptx::bf16x2 scale_rcp = + ptx::exp2f_rcp_2x_per_lane(__byte_perm(scale_pair, 0, 0x4140)); + col_scale_rcp_smem[tid] = reinterpret_cast(scale_rcp); } - }; - drain_row_scales(); - __syncthreads(); + auto drain_row_scales = [&]() { + constexpr int drain_tid_base = C::kDrainTidBase; + // A row's eight group scales are eight contiguous, 4-byte-aligned bytes of + // rowwise_scales (the stride is a multiple of 4 groups), so one warp drains a + // 32-row tile with two STG.32 per lane. In the GEMM-swizzled layout they are two + // 4-byte runs in adjacent tiles. A tile that ends inside the tensor writes only + // the groups that start inside it. + if (tid >= drain_tid_base && tid < drain_tid_base + kRowsPerMxBlock) { + const int r = tid - drain_tid_base; + const unsigned char* bytes = row_scale_bytes + r * kMxGroupsPerRow; + const size_t col_tiles = DIVUP(static_cast(K), size_t{128}); + if (!kEdge || full_width) { + const uint2 words = *(const uint2*)bytes; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + constexpr size_t kTileBytes = + swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_X * swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_Y; + const size_t idx = swizzle::gemm_swizzled_scale_idx(row0 + r, col0 >> 5, col_tiles); + *(unsigned*)(rowwise_scales + idx) = words.x; + *(unsigned*)(rowwise_scales + idx + kTileBytes) = words.y; + } else { + unsigned* dst = (unsigned*)(rowwise_scales + + (size_t)(row0 + r) * rowwise_scale_stride + (col0 >> 5)); + dst[0] = words.x; + dst[1] = words.y; + } + } else { + const int groups = (K - col0 + 31) / 32; + for (int g = 0; g < groups; ++g) { + size_t idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + idx = swizzle::gemm_swizzled_scale_idx(row0 + r, (col0 >> 5) + g, col_tiles); + } else { + idx = (size_t)(row0 + r) * rowwise_scale_stride + (col0 >> 5) + g; + } + rowwise_scales[idx] = bytes[g]; + } + } + } + }; + drain_row_scales(); + __syncthreads(); - // ---- columnwise quantization -------------------------------------------- - { - uint4 col_scale_rcp_vec[kVecLoadsPerRow]; + // ---- columnwise quantization -------------------------------------------- + { + uint4 col_scale_rcp_vec[kVecLoadsPerRow]; #pragma unroll - for (int n = 0; n < kVecLoadsPerRow; ++n) - col_scale_rcp_vec[n] = *(const uint4*)(col_scale_rcp_smem + kWordsPerLane * lane + 4 * n); - const unsigned* col_scale_rcp = (const unsigned*)col_scale_rcp_vec; - unsigned char* col_out_ptr = colwise_out + store_offset; + for (int n = 0; n < kVecLoadsPerRow; ++n) + col_scale_rcp_vec[n] = *(const uint4*)(col_scale_rcp_smem + kWordsPerLane * lane + 4 * n); + const unsigned* col_scale_rcp = (const unsigned*)col_scale_rcp_vec; + unsigned char* col_out_ptr = colwise_out + store_offset; #pragma unroll - for (int j = 0; j < kRowsPerWarp; ++j) { - const unsigned* v = tile_words + j * kWordsPerLane; - unsigned col_out[kWordsPerLane / 2]; + for (int j = 0; j < kRowsPerWarp; ++j) { + const unsigned* v = tile_words + j * kWordsPerLane; + unsigned col_out[kWordsPerLane / 2]; #pragma unroll - for (int m = 0; m < kWordsPerLane / 2; ++m) - col_out[m] = quantize_e4m3x4(v[2 * m], v[2 * m + 1], col_scale_rcp[2 * m], - col_scale_rcp[2 * m + 1]); - store_quantized(col_out_ptr + (unsigned)j * (unsigned)K, col_out); + for (int m = 0; m < kWordsPerLane / 2; ++m) + col_out[m] = quantize_e4m3x4(v[2 * m], v[2 * m + 1], col_scale_rcp[2 * m], + col_scale_rcp[2 * m + 1]); + if (in_bounds(j)) store_quantized(col_out_ptr + (unsigned)j * (unsigned)K, col_out); + } + } + }; + if constexpr (kTails) { + if (full_width && row0 + kRowsPerMxBlock <= M) { + band(std::false_type{}); + } else { + band(std::true_type{}); } + } else { + band(std::false_type{}); } } } -template +template __global__ void __launch_bounds__(TileConfig::kThreadsPerCta, TileConfig::kMinBlocksPerSm) quantize_mxfp8_kernel(const unsigned* __restrict__ input, @@ -468,13 +520,13 @@ __global__ void __launch_bounds__(TileConfig::kThreadsP unsigned char* __restrict__ rowwise_out, unsigned char* __restrict__ rowwise_scales, unsigned char* __restrict__ colwise_out, - unsigned char* __restrict__ colwise_scales, int K, + unsigned char* __restrict__ colwise_scales, int M, int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* __restrict__ table) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - quantize_regtile( + quantize_regtile( (const uint4*)input, (const uint4*)act_input, rowwise_out, rowwise_scales, colwise_out, - colwise_scales, K, rowwise_scale_stride, colwise_scale_stride, iters, table); + colwise_scales, M, K, rowwise_scale_stride, colwise_scale_stride, iters, table); #endif } @@ -500,7 +552,7 @@ static int pick_walk_length(int row_blocks, int grid_cols, long long target, // the rest of the unified array goes to L1. constexpr int kCarveoutPercentAct = 40; -template +template static void set_carveout() { constexpr int kPercent = IS_ACT ? kCarveoutPercentAct : 0; if constexpr (kPercent > 0) { @@ -513,18 +565,25 @@ static void set_carveout() { if (static_cast(device) >= done.size()) done.resize(device + 1, false); if (!done[device]) { NVTE_CHECK_CUDA(cudaFuncSetAttribute( - (const void*)quantize_mxfp8_kernel, + (const void*) + quantize_mxfp8_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, kPercent)); done[device] = true; } } } -// Row blocks per CTA walk: M / 32, walked in powers of two. +//! 32-row bands the grid covers: whole 64-row tiles of the generic kernel, whose scales it +//! writes. +inline int row_bands(int M) { return (M + 63) / 64 * 2; } +//! 256-column tiles the grid covers. +inline int col_tiles(int K) { return (K + kTileCols - 1) / kTileCols; } + +// Row blocks per CTA walk, in powers of two. template int walk_length(int M, int K) { - const int row_blocks = M / kRowsPerMxBlock; - const int grid_cols = K / kTileCols; + const int row_blocks = row_bands(M); + const int grid_cols = col_tiles(K); // The minimum walk amortizes the table copy, but not at the price of leaving // SMs without a CTA. int min_walk = kMinWalk; @@ -542,28 +601,36 @@ static void launch(const void* input, const void* act_input, void* rowwise_out, int rowwise_scale_stride, int colwise_scale_stride, const void* table, cudaStream_t stream) { using C = TileConfig; - set_carveout(); const int iters = walk_length(M, K); - dim3 grid(K / kTileCols, M / kRowsPerMxBlock / iters); - quantize_mxfp8_kernel - <<>>( - (const unsigned*)input, (const unsigned*)act_input, (unsigned char*)rowwise_out, - (unsigned char*)rowwise_scales, (unsigned char*)colwise_out, - (unsigned char*)colwise_scales, K, rowwise_scale_stride, colwise_scale_stride, iters, - table); + dim3 grid(col_tiles(K), row_bands(M) / iters); + const auto run = [&](auto tails) { + constexpr bool kTails = decltype(tails)::value; + set_carveout(); + quantize_mxfp8_kernel + <<>>( + (const unsigned*)input, (const unsigned*)act_input, (unsigned char*)rowwise_out, + (unsigned char*)rowwise_scales, (unsigned char*)colwise_out, + (unsigned char*)colwise_scales, M, K, rowwise_scale_stride, colwise_scale_stride, iters, + table); + }; + // Whole 64-row tiles and 256-column tiles leave nothing past M or K. The derivatives take the + // tail-capable kernel for every shape: with the band body behind its edge branch, the compiler + // schedules the loads so that they wait less (GB200, 32768x7168 dGeLU: 239 -> 220 us, fewer + // long-scoreboard stalls, DRAM throughput 73% -> 79%). The forward kernels do not gain from it. + if (!IS_DACT && M % 64 == 0 && K % kTileCols == 0) { + run(std::false_type{}); + } else { + run(std::true_type{}); + } } -// The grid is derived by exact division, so columns must tile into 256. Rows -// must be a multiple of the generic kernel's 64-row activation tile: on a -// 32-row tail that kernel writes zero scales into the padding past the last -// row, which this one does not. Output offsets within a 32-row band are 32-bit. -constexpr size_t kGenericActivationTileRows = 64; - +// Any number of rows. Columns in multiples of 16, the shapes the generic kernel takes (its TMA +// store of the FP8 output needs 16-byte rows); this also keeps every input row 16-byte aligned. +// Output offsets within a 32-row band are 32-bit. inline bool shape_supported(size_t rows, size_t cols) { - constexpr size_t kMaxRows = std::numeric_limits::max(); + constexpr size_t kMaxRows = std::numeric_limits::max() - 64; constexpr size_t kMaxCols = std::numeric_limits::max() / kRowsPerMxBlock; - return rows > 0 && cols > 0 && rows <= kMaxRows && cols <= kMaxCols && - rows % kGenericActivationTileRows == 0 && cols % kTileCols == 0; + return rows > 0 && cols > 0 && rows <= kMaxRows && cols <= kMaxCols && cols % 16 == 0; } } // namespace @@ -600,8 +667,8 @@ bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, } constexpr size_t kMaxGridY = 65535; const int M = static_cast(rows); - return static_cast(M / kRowsPerMxBlock / - walk_length(M, static_cast(cols))) <= kMaxGridY; + return static_cast(row_bands(M) / walk_length(M, static_cast(cols))) <= + kMaxGridY; } template Date: Thu, 15 Oct 2026 17:00:17 -0700 Subject: [PATCH 07/29] [Common] Use the register-tiled MXFP8 kernel without activation tables With NVTE_ACTIVATION_TABLE=0, the default, GeLU, SiLU and their derivatives now also take the register-tiled kernel: it computes the activation with util/math.h instead of reading a table, so its output is byte-identical to the generic kernel's, and the default path gets the register tiling without any change in results. The derivatives take it at every size. The forward activations are compute-bound without a table, and on small tensors too few CTAs hide the latency, so they take it from 2^23 elements on. GB200, BF16 -> MXFP8 E4M3, against TE main: dGeLU 1.81-1.92x, dSiLU 1.97-2.10x, GeLU 1.62-1.76x, SiLU 1.62-1.67x on 8192x7168 to 16384x16384; 993x7200 dGeLU/dSiLU 1.22x/1.48x. The regtile-vs-generic tests run every case with and without a table. Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 22 ++++--- .../common/cast/mxfp8/quantize_mxfp8.cuh | 10 +-- .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 65 +++++++++++++++---- .../common/util/activation_table.cuh | 3 + 4 files changed, 74 insertions(+), 26 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index 49c8a1a3118..e459a80b453 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -383,7 +383,7 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t // The first call takes the default dispatch, which picks the register-resident kernel for these // shapes; NVTE_MXFP8_REGTILE=0 keeps the second call on the generic kernel. void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, const size_t cols, - const bool swizzled) { + const bool swizzled, const bool use_table) { const std::vector shape{rows, cols}; const size_t n = rows * cols; Tensor x("x", shape, DType::kBFloat16); @@ -403,7 +403,7 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, for (size_t d = 0; d < b.shape.ndim; ++d) bytes *= b.shape.data[d]; if (bytes != 0) ASSERT_EQ(cudaMemset(b.data_ptr, 0xA5, bytes), cudaSuccess); } - const ScopedEnv tables("NVTE_ACTIVATION_TABLE", "1"); + const ScopedEnv tables("NVTE_ACTIVATION_TABLE", use_table ? "1" : "0"); const ScopedEnv regtile("NVTE_MXFP8_REGTILE", i == 0 ? "1" : "0"); switch (mode) { case Mode::GeLU: @@ -515,16 +515,17 @@ INSTANTIATE_TEST_SUITE_P(TinyGradients, CastMXFP8ActExactTestSuite, namespace { -using RegtileParams = std::tuple>; +using RegtileParams = + std::tuple, bool /*use_table*/>; class CastMXFP8ActRegtileTestSuite : public ::testing::TestWithParam {}; std::string regtile_case_name(const testing::TestParamInfo &info) { const char *fills[] = {"Typical", "OrderedInputs", "Specials", "TinyGradients"}; - const auto [mode, fill, swizzled, dims] = info.param; + const auto [mode, fill, swizzled, dims, use_table] = info.param; return std::string(mode_name(mode)) + "X" + fills[static_cast(fill)] + (swizzled ? "XSwizzled" : "XCompact") + "X" + std::to_string(dims.first) + "X" + - std::to_string(dims.second); + std::to_string(dims.second) + (use_table ? "XTable" : "XComputed"); } } // namespace @@ -535,8 +536,8 @@ TEST_P(CastMXFP8ActRegtileTestSuite, MatchesGenericKernel) { if (prop.major < 10) { GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; } - const auto [mode, fill, swizzled, dims] = GetParam(); - run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled); + const auto [mode, fill, swizzled, dims, use_table] = GetParam(); + run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled, use_table); } INSTANTIATE_TEST_SUITE_P( @@ -553,5 +554,10 @@ INSTANTIATE_TEST_SUITE_P( std::pair{1, 16}, std::pair{33, 272}, std::pair{100, 2064}, - std::pair{4127, 528})), + std::pair{4127, 528}, + // Large enough for the computed forward path (2^23 + // elements), with a row and a column tail. + std::pair{4127, 2064}), + // With a table, or computed from util/math.h without one. + ::testing::Bool()), regtile_case_name); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index f04ac491bc9..4279be4d153 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -966,11 +966,11 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, } // regtile::quantize replaces the generic kernel below for the requests it - // implements: bidimensional BF16 -> E4M3 with fused GeLU, SiLU, dGeLU or dSiLU read - // from an activation table, no dbias, cols % 16 == 0, 16-byte aligned inputs (see - // regtile::kImplements and regtile::can_use). It keeps the tile - // in registers instead of staging it through shared memory and writes the same - // bytes the generic kernel writes with the same table. + // implements: bidimensional BF16 -> E4M3 with fused GeLU, SiLU, dGeLU or dSiLU, read + // from an activation table or computed, no dbias, cols % 16 == 0, 16-byte aligned + // inputs (see regtile::kImplements and regtile::can_use). It keeps the tile in + // registers instead of staging it through shared memory and writes the same bytes + // the generic kernel writes. if constexpr (std::is_same_v && std::is_same_v && regtile::kImplements) { if (regtile::can_use(input, act_input, *output, *noop, diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh index 698095d7b51..493588a8abb 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -16,8 +16,9 @@ * partials, the row scales and the activation table. * * The activation or derivative is read from the table the caller passes (see - * util/activation_table.cuh), copied into shared memory once per CTA, so this kernel - * produces exactly the bytes the generic kernel produces with the same table. + * util/activation_table.cuh), copied into shared memory once per CTA, or, without a + * table, computed with util/math.h. Either way this kernel produces exactly the bytes the + * generic kernel produces for the same request. * * It replaces the generic kernel, not the specialized ones: cast-only requests * are served before the generic dispatch by the specialized kernels, one of @@ -93,12 +94,27 @@ __device__ __forceinline__ unsigned amax_to_e8m0_2x(float amax_hi, float amax_lo fmaxf(amax_lo * kMaxNormRcp, 0.0f)); } +/*! \brief Table's activation computed with util/math.h instead of read from a table: + * the generic kernel's output without tables (NVTE_ACTIVATION_TABLE=0). No table is + * copied. + */ +template +struct ComputedTable : activation_table::Table { + static constexpr bool kComputed = true; + static constexpr size_t kBytes = 16; // unused shared-memory slot +}; + /*! \brief The derivative of a BF16 pair, from the table in shared memory. */ template __device__ __forceinline__ ptx::floatx2 dact_word(unsigned x, const unsigned char* __restrict__ table) { - const float2 d = activation_table::derivative_2x(x, table); - return {d.x, d.y}; + if constexpr (TableT::kComputed) { + const ptx::floatx2 v = widen_bf16x2(x); + return {TableT::op(v.x), TableT::op(v.y)}; + } else { + const float2 d = activation_table::derivative_2x(x, table); + return {d.x, d.y}; + } } /*! \brief Start copying a table of kBytes into shared memory with one bulk copy. @@ -222,7 +238,9 @@ __device__ __forceinline__ void quantize_regtile( const int colwise_scale_cols = (K + 63) / 64 * 64; // The table copy overlaps the first input loads; see the wait below. - load_table_async(table_smem, &table_barrier, table, tid); + if constexpr (!TableT::kComputed) { + load_table_async(table_smem, &table_barrier, table, tid); + } #pragma unroll 1 for (int it = 0; it < iters; ++it) { @@ -269,7 +287,9 @@ __device__ __forceinline__ void quantize_regtile( // to spread them between the previous group's table lookups, each then stalling on its // own. asm volatile("" ::: "memory"); - if (it == 0 && h == 0) ptx::mbarrier_wait_parity(&table_barrier, 0); + if constexpr (!TableT::kComputed) { + if (it == 0 && h == 0) ptx::mbarrier_wait_parity(&table_barrier, 0); + } // Convert the load group into the BF16 pairs the epilogue quantizes. if constexpr (!IS_DACT) { #pragma unroll @@ -280,8 +300,14 @@ __device__ __forceinline__ void quantize_regtile( const int tile_word_base = (h * kRowsInFlight + t) * kWordsPerLane + n * 4; #pragma unroll for (int m = 0; m < 4; ++m) { - tile_words[tile_word_base + m] = activation_table::forward_2x( - x_words[m], reinterpret_cast(table_smem)); + if constexpr (TableT::kComputed) { + // util/math.h, then rounded to BF16 as the generic kernel rounds it. + const ptx::floatx2 v = widen_bf16x2(x_words[m]); + tile_words[tile_word_base + m] = to_bf16x2({TableT::op(v.x), TableT::op(v.y)}); + } else { + tile_words[tile_word_base + m] = activation_table::forward_2x( + x_words[m], reinterpret_cast(table_smem)); + } } } } @@ -636,22 +662,30 @@ inline bool shape_supported(size_t rows, size_t cols) { } // namespace //! Whether this kernel implements the request type: GeLU or SiLU, or the derivative of -//! either, without dbias. It reads the activation from a table. +//! either, without dbias. It reads the activation from a table, or computes it with +//! util/math.h when there is none. template constexpr bool kImplements = !IS_DBIAS && (IS_ACT != IS_DACT) && activation_table::kHasTable && (IS_ACT == activation_table::Table>::kForward); -// Whether this kernel takes a request of a type in kImplements, with BF16 input, E4M3 -// output and an activation table. Anything else takes the generic kernel, as does every -// request with NVTE_MXFP8_REGTILE=0. +// Whether this kernel takes a request of a type in kImplements, with BF16 input and E4M3 +// output. Anything else takes the generic kernel, as does every request with +// NVTE_MXFP8_REGTILE=0. template bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, const Tensor& noop, bool use_2d_quantization, const activation_table::TableRef& table) { - if (table.use == activation_table::kNoTable) return false; if (!getenv("NVTE_MXFP8_REGTILE", true)) return false; const auto [rows, cols] = input.flat_2d_dims(); + // Without a table the forward activation is compute-bound here, and on small tensors too + // few CTAs hide its latency: GB200, 993x7200 GeLU is 0.93x of the generic kernel, 4160x2048 + // 1.26x. The derivatives stay memory-bound and win at every size. + constexpr size_t kComputedForwardMinElements = size_t{1} << 23; + if (table.use == activation_table::kNoTable && !IS_DACT && + rows * cols < kComputedForwardMinElements) { + return false; + } // Vector widths: 16-byte loads of both inputs, 8-byte stores of both outputs // and of each row's scales. const bool aligned = is_aligned_ptr(input.data.dptr, 16) && @@ -689,6 +723,11 @@ void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, static_cast(output->columnwise_scale_inv.shape[1]), table.ptr, stream); }; constexpr Kind kKind = kKindOf; + if (table.use == kNoTable) { + run(std::common_type>{}); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } if constexpr (kKind == kDsilu) { if (table.use == kExactDsiluTable) { run(std::common_type{}); diff --git a/transformer_engine/common/util/activation_table.cuh b/transformer_engine/common/util/activation_table.cuh index 405bce3fbd1..622cdfcd0e3 100644 --- a/transformer_engine/common/util/activation_table.cuh +++ b/transformer_engine/common/util/activation_table.cuh @@ -66,6 +66,8 @@ enum Kind : int { template struct ForwardLayout { static constexpr bool kForward = true; + //! Whether the kernel computes op() instead of reading the table (see the register-tiled kernel). + static constexpr bool kComputed = false; static constexpr uint32_t kLowSentinel = kLo - 1; static constexpr uint32_t kHighSentinel = kHi + 1; static constexpr uint32_t kRows = kHighSentinel - kLowSentinel + 1; @@ -81,6 +83,7 @@ struct ForwardLayout { template struct DerivativeLayout { static constexpr bool kForward = false; + static constexpr bool kComputed = false; static constexpr uint32_t kLoBits = kLo; static constexpr uint32_t kHiBits = kHi; //! Magnitudes from here on (and NaN) give NaN, as util/math.h does. From 17488a612c8d31d650eb906c9d364a1e1e597727 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 15 Oct 2026 16:37:15 -0700 Subject: [PATCH 08/29] [Common] Fuse dbias into the register-tiled MXFP8 activation kernel nvte_quantize_dbias_dgelu and nvte_quantize_dbias_dsilu took the generic kernel even with the activation tables on. The register-tiled kernel now handles them and writes the same dbias workspace, which common::reduce_dbias then reduces as before. The generic kernel sums each column of a 64-row chunk strictly in row order, in FP32, before rounding the products to BF16. To give the same bits, a CTA here takes the chunk's two 32-row bands back to back, stages each band's FP32 products in 32 KB of dynamic shared memory, and one thread per column adds them in row order, carrying the sum across the two bands. The staging buffer limits the kernel to four CTAs per SM, so its launch bounds allow 64 registers. GB200, tables on, against TE main: 1.29-1.75x on normal data and up to 2.47x on wide data (32768x7168 dbias_dgelu: 638 -> 391 us), where the generic kernel ran at 0.94-0.99x. The regtile-vs-generic test adds both dbias modes and compares the FP32 workspace and the BF16 dbias byte for byte. Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 37 +++- .../common/cast/mxfp8/quantize_mxfp8.cuh | 22 ++- .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 161 +++++++++++++----- 3 files changed, 163 insertions(+), 57 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index e459a80b453..3315bd716fa 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -392,6 +392,9 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, Tensor outs[2] = { Tensor("default", shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING), Tensor("generic", shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING)}; + Tensor dbias[2] = {Tensor("dbias_default", std::vector{cols}, DType::kBFloat16), + Tensor("dbias_generic", std::vector{cols}, DType::kBFloat16)}; + Tensor workspace[2]; for (int i = 0; i < 2; ++i) { outs[i].set_with_gemm_swizzled_scales(swizzled); // Fill every output byte, padding included, with a marker, so the two kernels must write the @@ -415,13 +418,42 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, case Mode::DGeLU: nvte_dgelu(grad.data(), x.data(), outs[i].data(), 0); break; - default: + case Mode::DSiLU: nvte_dsilu(grad.data(), x.data(), outs[i].data(), 0); break; + default: { + auto fn = + mode == Mode::DBiasDGeLU ? &nvte_quantize_dbias_dgelu : &nvte_quantize_dbias_dsilu; + fn(grad.data(), x.data(), outs[i].data(), dbias[i].data(), workspace[i].data(), 0); + workspace[i] = Tensor("workspace", workspace[i].rowwise_shape(), workspace[i].dtype()); + // The per-chunk column sums must match too, not just the BF16 dbias they reduce to. + ASSERT_EQ(cudaMemset(workspace[i].rowwise_dptr(), 0xA5, + product(workspace[i].rowwise_shape()) * sizeof(float)), + cudaSuccess); + fn(grad.data(), x.data(), outs[i].data(), dbias[i].data(), workspace[i].data(), 0); + } } } ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); for (auto &out : outs) out.to_cpu(); + if (is_dbias(mode)) { + for (int i = 0; i < 2; ++i) { + dbias[i].to_cpu(); + workspace[i].to_cpu(); + } + EXPECT_EQ(count_byte_mismatches( + reinterpret_cast(workspace[0].rowwise_cpu_dptr()), + reinterpret_cast(workspace[1].rowwise_cpu_dptr()), + product(workspace[0].rowwise_shape()) * sizeof(float)), + 0u) + << "dbias workspace"; + EXPECT_EQ(count_byte_mismatches( + reinterpret_cast(dbias[0].rowwise_cpu_dptr()), + reinterpret_cast(dbias[1].rowwise_cpu_dptr()), + cols * sizeof(bf16)), + 0u) + << "dbias"; + } EXPECT_EQ(count_byte_mismatches( reinterpret_cast(outs[0].rowwise_cpu_dptr()), reinterpret_cast(outs[1].rowwise_cpu_dptr()), n), @@ -542,7 +574,8 @@ TEST_P(CastMXFP8ActRegtileTestSuite, MatchesGenericKernel) { INSTANTIATE_TEST_SUITE_P( OperatorTest, CastMXFP8ActRegtileTestSuite, - ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::SiLU, Mode::DSiLU), + ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::DBiasDGeLU, Mode::SiLU, + Mode::DSiLU, Mode::DBiasDSiLU), ::testing::Values(Fill::Typical, Fill::OrderedInputs, Fill::Specials), ::testing::Bool(), ::testing::Values(std::pair{1024, 2048}, diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 4279be4d153..af32fd10443 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -967,16 +967,22 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, // regtile::quantize replaces the generic kernel below for the requests it // implements: bidimensional BF16 -> E4M3 with fused GeLU, SiLU, dGeLU or dSiLU, read - // from an activation table or computed, no dbias, cols % 16 == 0, 16-byte aligned - // inputs (see regtile::kImplements and regtile::can_use). It keeps the tile in - // registers instead of staging it through shared memory and writes the same bytes - // the generic kernel writes. + // from an activation table or computed, the derivatives optionally with dbias, + // cols % 16 == 0, 16-byte aligned inputs (see regtile::kImplements and + // regtile::can_use). It keeps the tile in registers instead of staging it through + // shared memory and writes the same bytes, and the same dbias workspace, the + // generic kernel writes. if constexpr (std::is_same_v && std::is_same_v && regtile::kImplements) { - if (regtile::can_use(input, act_input, *output, *noop, - use_2d_quantization, activation_lut)) { - regtile::quantize( - input, act_input, output, activation_lut, stream); + if (regtile::can_use( + input, act_input, *output, *noop, use_2d_quantization, activation_lut)) { + regtile::quantize( + input, act_input, output, activation_lut, workspace_ptr, stream); + if constexpr (IS_DBIAS) { + common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, + stream); + } return; } } diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh index 493588a8abb..65ee5e76e00 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -6,7 +6,7 @@ /*! \file quantize_mxfp8_regtile.cuh * \brief Register-resident MXFP8 bidimensional quantize fused with GeLU or SiLU, or - * their derivatives. + * their derivatives, optionally with dbias. * * Rowwise and colwise MX scales need an amax over the same data in two * directions. The generic kernel stages each tile through shared memory and @@ -183,7 +183,10 @@ constexpr int kMaxThreadsPerSm = 1024; constexpr int kMaxThreadsPerSm = 2048; #endif -template +// Minimum resident CTAs per SM with dbias: the 32 KB staging buffer allows about four. +constexpr int kMinBlocksDbias = 4; + +template struct TileConfig { static_assert(IS_ACT != IS_DACT, "Exactly one of the activation and its derivative is fused."); static constexpr int kThreadsPerCta = IS_DACT ? 256 : 128; @@ -196,24 +199,37 @@ struct TileConfig { static constexpr int kDrainTidBase = (kThreadsPerCta >= kTileCols / 2 + kTileCols / 4) ? kTileCols / 2 : 0; static constexpr int kRowsInFlight = 2; - static constexpr int kRequestedMinBlocks = IS_DACT ? kMinBlocksDact : kMinBlocksAct; + static constexpr int kRequestedMinBlocks = + kDbias ? kMinBlocksDbias : (IS_DACT ? kMinBlocksDact : kMinBlocksAct); static constexpr int kMinBlocksPerSm = kRequestedMinBlocks < kMaxThreadsPerSm / kThreadsPerCta ? kRequestedMinBlocks : kMaxThreadsPerSm / kThreadsPerCta; static constexpr int kTableBytes = static_cast(TableT::kBytes); + //! Dynamic shared memory: with dbias, the band's FP32 products before rounding. + static constexpr int kDbiasStageBytes = kDbias ? kRowsPerMxBlock * kTileCols * 4 : 0; + static_assert(!kDbias || (IS_DACT && kThreadsPerCta == kTileCols), + "dbias: one thread per column sums the column."); }; // Register-resident tiling: the 32x256 tile stays in registers; only the // colwise partials go through shared memory. // kTails: the grid may cover elements past M or K (see shape_supported). Without tails the // bounds checks compile away, which matters at the kernel's register budget. -template +// +// kDbias: also write the per-64-row-chunk column sums of the FP32 products (before rounding to +// BF16) into dbias_workspace[chunk][K], as the generic kernel does; common::reduce_dbias then +// adds the chunks. The generic kernel sums each column of a chunk strictly in row order, so a +// CTA here takes the chunk's two 32-row bands back to back, stages each band's products in +// shared memory, and one thread per column adds them in row order. +template __device__ __forceinline__ void quantize_regtile( const uint4* __restrict__ input, const uint4* __restrict__ act_input, unsigned char* __restrict__ rowwise_out, unsigned char* __restrict__ rowwise_scales, unsigned char* __restrict__ colwise_out, unsigned char* __restrict__ colwise_scales, int M, - int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* table) { - using C = TileConfig; + int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* table, + float* __restrict__ dbias_workspace) { + using C = TileConfig; constexpr int kWarpsPerCta = C::kWarpsPerCta, kRowsPerWarp = C::kRowsPerWarp, kRowsInFlight = C::kRowsInFlight; @@ -222,6 +238,8 @@ __device__ __forceinline__ void quantize_regtile( __shared__ __align__(8) unsigned char row_scale_bytes[kRowsPerMxBlock * kMxGroupsPerRow]; __shared__ __align__(16) unsigned char table_smem[C::kTableBytes]; __shared__ __align__(8) uint64_t table_barrier; + extern __shared__ __align__(16) float dbias_stage[]; // [kRowsPerMxBlock][kTileCols], dbias only + float dbias_partial = 0.0f; const int tid = threadIdx.x; const int lane = tid & 31; @@ -246,7 +264,13 @@ __device__ __forceinline__ void quantize_regtile( for (int it = 0; it < iters; ++it) { // Row blocks are taken grid-strided, so the CTAs on the same step cover one // contiguous row band. - const int row0 = (it * gridDim.y + blockIdx.y) * kRowsPerMxBlock; + // With dbias a CTA takes the two bands of a 64-row chunk on consecutive steps. + const int row0 = kDbias + ? (((it >> 1) * gridDim.y + blockIdx.y) * 2 + (it & 1)) * kRowsPerMxBlock + : (it * gridDim.y + blockIdx.y) * kRowsPerMxBlock; + if constexpr (kDbias) { + if ((it & 1) == 0) dbias_partial = 0.0f; + } // The band's work. A band inside the tensor runs it without bounds checks; only bands at // the edges of the tensor (kEdge) pay for them. const auto band = [&](auto edge) { @@ -317,15 +341,25 @@ __device__ __forceinline__ void quantize_regtile( const ptx::floatx2 product = \ ptx::mul_2x(widen_bf16x2(grad_words[M]), dact_word(x_words[M], table_smem)); \ tile_words[(h * kRowsInFlight + (T)) * kWordsPerLane + (M)] = to_bf16x2(product); \ + products[2 * (M)] = product.x; \ + products[2 * (M) + 1] = product.y; \ } -#define MXFP8_DACT_ROW(T) \ - { \ - const unsigned* x_words = (const unsigned*)&x_vec[T][0]; \ - const unsigned* grad_words = (const unsigned*)&grad_vec[T][0]; \ - MXFP8_DACT_WORD(T, 0) \ - MXFP8_DACT_WORD(T, 1) \ - MXFP8_DACT_WORD(T, 2) \ - MXFP8_DACT_WORD(T, 3) \ +#define MXFP8_DACT_ROW(T) \ + { \ + const unsigned* x_words = (const unsigned*)&x_vec[T][0]; \ + const unsigned* grad_words = (const unsigned*)&grad_vec[T][0]; \ + float products[2 * kWordsPerLane]; \ + MXFP8_DACT_WORD(T, 0) \ + MXFP8_DACT_WORD(T, 1) \ + MXFP8_DACT_WORD(T, 2) \ + MXFP8_DACT_WORD(T, 3) \ + if constexpr (kDbias) { \ + float4* stage = \ + (float4*)(dbias_stage + (warp * kRowsPerWarp + h * kRowsInFlight + (T)) * kTileCols + \ + kColsPerLane * lane); \ + stage[0] = make_float4(products[0], products[1], products[2], products[3]); \ + stage[1] = make_float4(products[4], products[5], products[6], products[7]); \ + } \ } static_assert(kRowsInFlight == 2, "A load group is two rows."); MXFP8_DACT_ROW(0) @@ -437,6 +471,14 @@ __device__ __forceinline__ void quantize_regtile( publish_col_amax_partials(); __syncthreads(); + if constexpr (kDbias) { + // The generic kernel's order: a strictly sequential FP32 sum down the column. + float sum = dbias_partial; +#pragma unroll 8 + for (int r = 0; r < kRowsPerMxBlock; ++r) + sum = __fadd_rn(sum, dbias_stage[r * kTileCols + tid]); + dbias_partial = sum; + } if (tid < kTileCols / 2) { unsigned col_amax = col_amax_partials[0][tid]; #pragma unroll @@ -535,12 +577,18 @@ __device__ __forceinline__ void quantize_regtile( } else { band(std::false_type{}); } + if constexpr (kDbias) { + if ((it & 1) == 1 && col0 + tid < K) { + dbias_workspace[(size_t)(row0 / (2 * kRowsPerMxBlock)) * K + col0 + tid] = dbias_partial; + } + } } } -template -__global__ void __launch_bounds__(TileConfig::kThreadsPerCta, - TileConfig::kMinBlocksPerSm) +template +__global__ void __launch_bounds__(TileConfig::kThreadsPerCta, + TileConfig::kMinBlocksPerSm) quantize_mxfp8_kernel(const unsigned* __restrict__ input, const unsigned* __restrict__ act_input, unsigned char* __restrict__ rowwise_out, @@ -548,11 +596,12 @@ __global__ void __launch_bounds__(TileConfig::kThreadsP unsigned char* __restrict__ colwise_out, unsigned char* __restrict__ colwise_scales, int M, int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, - const void* __restrict__ table) { + const void* __restrict__ table, float* __restrict__ dbias_workspace) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - quantize_regtile( + quantize_regtile( (const uint4*)input, (const uint4*)act_input, rowwise_out, rowwise_scales, colwise_out, - colwise_scales, M, K, rowwise_scale_stride, colwise_scale_stride, iters, table); + colwise_scales, M, K, rowwise_scale_stride, colwise_scale_stride, iters, table, + dbias_workspace); #endif } @@ -578,10 +627,12 @@ static int pick_walk_length(int row_blocks, int grid_cols, long long target, // the rest of the unified array goes to L1. constexpr int kCarveoutPercentAct = 40; -template +template static void set_carveout() { constexpr int kPercent = IS_ACT ? kCarveoutPercentAct : 0; - if constexpr (kPercent > 0) { + constexpr int kDynamicBytes = TileConfig::kDbiasStageBytes; + if constexpr (kPercent > 0 || kDynamicBytes > 0) { // Function attributes are per device, like the tables. static std::mutex mutex; static std::vector done; @@ -590,10 +641,18 @@ static void set_carveout() { std::lock_guard lock(mutex); if (static_cast(device) >= done.size()) done.resize(device + 1, false); if (!done[device]) { - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - (const void*) - quantize_mxfp8_kernel, - cudaFuncAttributePreferredSharedMemoryCarveout, kPercent)); + const void* kernel = (const void*) + quantize_mxfp8_kernel; + if constexpr (kPercent > 0) { + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, kPercent)); + } + if constexpr (kDynamicBytes > 0) { + // The staging buffer takes the kernel past the 48 KB of shared memory a CTA gets by + // default. + NVTE_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + kDynamicBytes)); + } done[device] = true; } } @@ -606,7 +665,7 @@ inline int row_bands(int M) { return (M + 63) / 64 * 2; } inline int col_tiles(int K) { return (K + kTileCols - 1) / kTileCols; } // Row blocks per CTA walk, in powers of two. -template +template int walk_length(int M, int K) { const int row_blocks = row_bands(M); const int grid_cols = col_tiles(K); @@ -618,26 +677,30 @@ int walk_length(int M, int K) { if constexpr (IS_DACT) { if (iters > kMaxWalkDact) iters = kMaxWalkDact; } + // With dbias a CTA walks whole 64-row chunks: two bands per chunk. + if constexpr (kDbias) { + if (iters % 2 != 0) iters *= 2; + } return iters; } -template +template static void launch(const void* input, const void* act_input, void* rowwise_out, void* rowwise_scales, void* colwise_out, void* colwise_scales, int M, int K, int rowwise_scale_stride, int colwise_scale_stride, const void* table, - cudaStream_t stream) { - using C = TileConfig; - const int iters = walk_length(M, K); + float* dbias_workspace, cudaStream_t stream) { + using C = TileConfig; + const int iters = walk_length(M, K); dim3 grid(col_tiles(K), row_bands(M) / iters); const auto run = [&](auto tails) { constexpr bool kTails = decltype(tails)::value; - set_carveout(); - quantize_mxfp8_kernel - <<>>( + set_carveout(); + quantize_mxfp8_kernel + <<>>( (const unsigned*)input, (const unsigned*)act_input, (unsigned char*)rowwise_out, (unsigned char*)rowwise_scales, (unsigned char*)colwise_out, (unsigned char*)colwise_scales, M, K, rowwise_scale_stride, colwise_scale_stride, iters, - table); + table, dbias_workspace); }; // Whole 64-row tiles and 256-column tiles leave nothing past M or K. The derivatives take the // tail-capable kernel for every shape: with the band body behind its edge branch, the compiler @@ -662,18 +725,18 @@ inline bool shape_supported(size_t rows, size_t cols) { } // namespace //! Whether this kernel implements the request type: GeLU or SiLU, or the derivative of -//! either, without dbias. It reads the activation from a table, or computes it with -//! util/math.h when there is none. +//! either, the derivative optionally with dbias. It reads the activation from a table, or +//! computes it with util/math.h when there is none. template constexpr bool kImplements = - !IS_DBIAS && (IS_ACT != IS_DACT) && activation_table::kHasTable && + (!IS_DBIAS || IS_DACT) && (IS_ACT != IS_DACT) && activation_table::kHasTable && (IS_ACT == activation_table::Table>::kForward); // Whether this kernel takes a request of a type in kImplements, with BF16 input and E4M3 // output. Anything else takes the generic kernel, as does every request with // NVTE_MXFP8_REGTILE=0. -template +template bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, const Tensor& noop, bool use_2d_quantization, const activation_table::TableRef& table) { if (!getenv("NVTE_MXFP8_REGTILE", true)) return false; @@ -701,14 +764,17 @@ bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, } constexpr size_t kMaxGridY = 65535; const int M = static_cast(rows); - return static_cast(row_bands(M) / walk_length(M, static_cast(cols))) <= - kMaxGridY; + return static_cast( + row_bands(M) / walk_length(M, static_cast(cols))) <= kMaxGridY; } -template +//! dbias_workspace: [ceil(rows / 64), cols] FP32 column sums for common::reduce_dbias, or +//! nullptr without dbias. +template void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, - const activation_table::TableRef& table, cudaStream_t stream) { + const activation_table::TableRef& table, float* dbias_workspace, + cudaStream_t stream) { using namespace activation_table; // Plain variables: lambdas cannot capture structured bindings before C++20. const size_t rows = input.flat_first_dim(); @@ -716,11 +782,12 @@ void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, const void* act_ptr = IS_DACT ? act_input->data.dptr : input.data.dptr; const auto run = [&](auto table_type) { using TableT = typename decltype(table_type)::type; - launch( + launch( input.data.dptr, act_ptr, output->data.dptr, output->scale_inv.dptr, output->columnwise_data.dptr, output->columnwise_scale_inv.dptr, static_cast(rows), static_cast(cols), static_cast(output->scale_inv.shape[1]), - static_cast(output->columnwise_scale_inv.shape[1]), table.ptr, stream); + static_cast(output->columnwise_scale_inv.shape[1]), table.ptr, dbias_workspace, + stream); }; constexpr Kind kKind = kKindOf; if (table.use == kNoTable) { From ccaa19697d1a487172d1b0b9ad74e04842c377aa Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 15 Oct 2026 17:02:45 -0700 Subject: [PATCH 09/29] [Common] Support amax, noop and E5M2 output in the register-tiled MXFP8 kernel The register-tiled kernel refused requests with an amax output, a noop tensor or E5M2 output, which then took the generic kernel. - E5M2: the kernel is templated on the FP8 output type; the scale uses that type's max norm and the packed conversion its E5M2 form. - amax: the tensor amax is the max of the column amaxes the kernel already computes, reduced per CTA and combined with one atomic max, the value the generic kernel reduces. It is a separate instantiation, so the kernels without amax keep their register budget. - noop: the generic kernel ignores the noop tensor for every request that computes an activation, and so does this kernel now. GB200, tables on: E4M3 without amax is unchanged; E5M2 runs at the E4M3 speed (2.3-2.8x TE main at 32768x7168); with amax, the forward costs about 1% and the derivatives about 5%. The register-tiled versus generic test gains the output type and amax as dimensions; amax must match bit for bit. Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 65 +++++++--- .../common/cast/mxfp8/quantize_mxfp8.cuh | 17 +-- .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 120 +++++++++++------- 3 files changed, 131 insertions(+), 71 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index 3315bd716fa..321ea6c33e8 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -381,22 +381,33 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t // With a table, the register-resident kernel and the generic kernel must write the same bytes. // The first call takes the default dispatch, which picks the register-resident kernel for these -// shapes; NVTE_MXFP8_REGTILE=0 keeps the second call on the generic kernel. +// shapes; NVTE_MXFP8_REGTILE=0 keeps the second call on the generic kernel. With with_amax, both +// also reduce the tensor amax, which must match bit for bit. void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, const size_t cols, - const bool swizzled, const bool use_table) { + const bool swizzled, const bool use_table, const DType otype, + const bool with_amax) { const std::vector shape{rows, cols}; const size_t n = rows * cols; Tensor x("x", shape, DType::kBFloat16); Tensor grad("grad", shape, DType::kBFloat16); fill_inputs(fill, &x, &grad); - Tensor outs[2] = { - Tensor("default", shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING), - Tensor("generic", shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING)}; + Tensor outs[2] = {Tensor("default", shape, otype, true, true, NVTE_MXFP8_1D_SCALING), + Tensor("generic", shape, otype, true, true, NVTE_MXFP8_1D_SCALING)}; Tensor dbias[2] = {Tensor("dbias_default", std::vector{cols}, DType::kBFloat16), Tensor("dbias_generic", std::vector{cols}, DType::kBFloat16)}; Tensor workspace[2]; + float *amax_dev = nullptr; + if (with_amax) { + ASSERT_EQ(cudaMalloc(&amax_dev, 2 * sizeof(float)), cudaSuccess); + ASSERT_EQ(cudaMemset(amax_dev, 0, 2 * sizeof(float)), cudaSuccess); + } for (int i = 0; i < 2; ++i) { outs[i].set_with_gemm_swizzled_scales(swizzled); + if (with_amax) { + const size_t one = 1; + const NVTEBasicTensor amax = {amax_dev + i, kNVTEFloat32, nvte_make_shape(&one, 1)}; + nvte_set_tensor_param_v2(outs[i].data(), kNVTEAmax, &amax, sizeof(amax)); + } // Fill every output byte, padding included, with a marker, so the two kernels must write the // same set of bytes, not just the same values. for (const auto param : {kNVTERowwiseData, kNVTEColumnwiseData, kNVTERowwiseScaleInv, @@ -435,6 +446,14 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, } } ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + if (with_amax) { + float amax[2]; + ASSERT_EQ(cudaMemcpy(amax, amax_dev, sizeof(amax), cudaMemcpyDeviceToHost), cudaSuccess); + cudaFree(amax_dev); + uint32_t bits[2]; + std::memcpy(bits, amax, sizeof(bits)); + EXPECT_EQ(bits[0], bits[1]) << "amax " << amax[0] << " vs " << amax[1]; + } for (auto &out : outs) out.to_cpu(); if (is_dbias(mode)) { for (int i = 0; i < 2; ++i) { @@ -454,15 +473,20 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, 0u) << "dbias"; } - EXPECT_EQ(count_byte_mismatches( - reinterpret_cast(outs[0].rowwise_cpu_dptr()), - reinterpret_cast(outs[1].rowwise_cpu_dptr()), n), - 0u) + // The FP8 data as raw bytes, whichever FP8 type the output holds. + const auto rowwise_bytes = [&](Tensor &t) { + return otype == DType::kFloat8E5M2 + ? reinterpret_cast(t.rowwise_cpu_dptr()) + : reinterpret_cast(t.rowwise_cpu_dptr()); + }; + const auto columnwise_bytes = [&](Tensor &t) { + return otype == DType::kFloat8E5M2 + ? reinterpret_cast(t.columnwise_cpu_dptr()) + : reinterpret_cast(t.columnwise_cpu_dptr()); + }; + EXPECT_EQ(count_byte_mismatches(rowwise_bytes(outs[0]), rowwise_bytes(outs[1]), n), 0u) << "rowwise data"; - EXPECT_EQ(count_byte_mismatches( - reinterpret_cast(outs[0].columnwise_cpu_dptr()), - reinterpret_cast(outs[1].columnwise_cpu_dptr()), n), - 0u) + EXPECT_EQ(count_byte_mismatches(columnwise_bytes(outs[0]), columnwise_bytes(outs[1]), n), 0u) << "colwise data"; EXPECT_EQ(count_byte_mismatches(outs[0].rowwise_cpu_scale_inv_ptr(), outs[1].rowwise_cpu_scale_inv_ptr(), @@ -548,16 +572,18 @@ INSTANTIATE_TEST_SUITE_P(TinyGradients, CastMXFP8ActExactTestSuite, namespace { using RegtileParams = - std::tuple, bool /*use_table*/>; + std::tuple, bool /*use_table*/, + DType /*output*/, bool /*amax*/>; class CastMXFP8ActRegtileTestSuite : public ::testing::TestWithParam {}; std::string regtile_case_name(const testing::TestParamInfo &info) { const char *fills[] = {"Typical", "OrderedInputs", "Specials", "TinyGradients"}; - const auto [mode, fill, swizzled, dims, use_table] = info.param; + const auto [mode, fill, swizzled, dims, use_table, otype, with_amax] = info.param; return std::string(mode_name(mode)) + "X" + fills[static_cast(fill)] + (swizzled ? "XSwizzled" : "XCompact") + "X" + std::to_string(dims.first) + "X" + - std::to_string(dims.second) + (use_table ? "XTable" : "XComputed"); + std::to_string(dims.second) + (use_table ? "XTable" : "XComputed") + + (otype == DType::kFloat8E5M2 ? "XE5M2" : "XE4M3") + (with_amax ? "XAmax" : ""); } } // namespace @@ -568,8 +594,9 @@ TEST_P(CastMXFP8ActRegtileTestSuite, MatchesGenericKernel) { if (prop.major < 10) { GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; } - const auto [mode, fill, swizzled, dims, use_table] = GetParam(); - run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled, use_table); + const auto [mode, fill, swizzled, dims, use_table, otype, with_amax] = GetParam(); + run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled, use_table, otype, + with_amax); } INSTANTIATE_TEST_SUITE_P( @@ -592,5 +619,7 @@ INSTANTIATE_TEST_SUITE_P( // elements), with a row and a column tail. std::pair{4127, 2064}), // With a table, or computed from util/math.h without one. + ::testing::Bool(), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), ::testing::Bool()), regtile_case_name); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index af32fd10443..ab98580a9fb 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -966,17 +966,18 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, } // regtile::quantize replaces the generic kernel below for the requests it - // implements: bidimensional BF16 -> E4M3 with fused GeLU, SiLU, dGeLU or dSiLU, read - // from an activation table or computed, the derivatives optionally with dbias, - // cols % 16 == 0, 16-byte aligned inputs (see regtile::kImplements and - // regtile::can_use). It keeps the tile in registers instead of staging it through - // shared memory and writes the same bytes, and the same dbias workspace, the - // generic kernel writes. - if constexpr (std::is_same_v && std::is_same_v && + // implements: bidimensional BF16 -> E4M3 or E5M2 with fused GeLU, SiLU, dGeLU or + // dSiLU, read from an activation table or computed, the derivatives optionally with + // dbias, with or without amax, cols % 16 == 0, 16-byte aligned inputs (see + // regtile::kImplements and regtile::can_use). It keeps the tile in registers instead + // of staging it through shared memory and writes the same bytes, the same amax and + // the same dbias workspace the generic kernel writes. + if constexpr (std::is_same_v && + (std::is_same_v || std::is_same_v) && regtile::kImplements) { if (regtile::can_use( input, act_input, *output, *noop, use_2d_quantization, activation_lut)) { - regtile::quantize( input, act_input, output, activation_lut, workspace_ptr, stream); if constexpr (IS_DBIAS) { diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh index 65ee5e76e00..236031e4fdd 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -88,8 +88,9 @@ __device__ __forceinline__ unsigned e8m0_to_bf16x2_reciprocal(unsigned e) { * only NaNs gets amax 0; the fmaxf reproduces that for the data-seeded maxima * here. The conversion saturates, so the scale never reaches 255. */ +template __device__ __forceinline__ unsigned amax_to_e8m0_2x(float amax_hi, float amax_lo) { - constexpr float kMaxNormRcp = Quantized_Limits::max_norm_rcp; + constexpr float kMaxNormRcp = Quantized_Limits::max_norm_rcp; return ptx::float_to_e8m0_2x(fmaxf(amax_hi * kMaxNormRcp, 0.0f), fmaxf(amax_lo * kMaxNormRcp, 0.0f)); } @@ -156,10 +157,12 @@ constexpr int kVecLoadsPerRow = kWordsPerLane / 4; //! Lanes that have to cooperate to cover one 32-column MX group. constexpr int kLanesPerMxGroup = 32 / kColsPerLane; -/*! \brief Scale two BF16 pairs, each by its own reciprocal scale pair, to four E4M3. */ -__device__ __forceinline__ unsigned quantize_e4m3x4(unsigned in01, unsigned in23, unsigned scale01, - unsigned scale23) { - ptx::fp8e4m3x4 out; +/*! \brief Scale two BF16 pairs, each by its own reciprocal scale pair, to four OType (E4M3 or + * E5M2). */ +template +__device__ __forceinline__ unsigned quantize_x4(unsigned in01, unsigned in23, unsigned scale01, + unsigned scale23) { + ptx::FPx4 out; ptx::mul_cvt_4x(out, reinterpret_cast(in01), reinterpret_cast(scale01), reinterpret_cast(in23), @@ -221,14 +224,16 @@ struct TileConfig { // adds the chunks. The generic kernel sums each column of a chunk strictly in row order, so a // CTA here takes the chunk's two 32-row bands back to back, stages each band's products in // shared memory, and one thread per column adds them in row order. -template +// +// kAmax: also fold the amax of the output into *amax_ptr. +template __device__ __forceinline__ void quantize_regtile( const uint4* __restrict__ input, const uint4* __restrict__ act_input, unsigned char* __restrict__ rowwise_out, unsigned char* __restrict__ rowwise_scales, unsigned char* __restrict__ colwise_out, unsigned char* __restrict__ colwise_scales, int M, int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* table, - float* __restrict__ dbias_workspace) { + float* __restrict__ dbias_workspace, float* __restrict__ amax_ptr) { using C = TileConfig; constexpr int kWarpsPerCta = C::kWarpsPerCta, kRowsPerWarp = C::kRowsPerWarp, kRowsInFlight = C::kRowsInFlight; @@ -260,6 +265,8 @@ __device__ __forceinline__ void quantize_regtile( load_table_async(table_smem, &table_barrier, table, tid); } + float thread_amax = 0.0f; + #pragma unroll 1 for (int it = 0; it < iters; ++it) { // Row blocks are taken grid-strided, so the CTAs on the same step cover one @@ -386,7 +393,8 @@ __device__ __forceinline__ void quantize_regtile( // Fold the two halves with one PRMT and a packed max, which leaves the // magnitude in both halves; the high half is then the FP32 bit pattern. amax = abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); - const unsigned scale_byte = amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); + const unsigned scale_byte = + amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); if ((lane & (kLanesPerMxGroup - 1)) == 0) row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + (lane / kLanesPerMxGroup)] = (unsigned char)scale_byte; @@ -408,7 +416,7 @@ __device__ __forceinline__ void quantize_regtile( for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) amax_pair = abs_max_bf16x2(amax_pair, __shfl_xor_sync(0xffffffffu, amax_pair, lane_mask)); const ptx::floatx2 amax_2x = widen_bf16x2(amax_pair); - const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); const unsigned scale_byte0 = scale_pair & 0xffu, scale_byte1 = (scale_pair >> 8) & 0xffu; if ((lane & (kLanesPerMxGroup - 1)) == 0) { unsigned char* scale_dst = row_scale_bytes + (warp * kRowsPerWarp + j) * kMxGroupsPerRow + @@ -446,12 +454,12 @@ __device__ __forceinline__ void quantize_regtile( #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) row_out[m] = - quantize_e4m3x4(row_a[2 * m], row_a[2 * m + 1], row_scale_rcp0, row_scale_rcp0); + quantize_x4(row_a[2 * m], row_a[2 * m + 1], row_scale_rcp0, row_scale_rcp0); if (in_bounds(j)) store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) row_out[m] = - quantize_e4m3x4(row_b[2 * m], row_b[2 * m + 1], row_scale_rcp1, row_scale_rcp1); + quantize_x4(row_b[2 * m], row_b[2 * m + 1], row_scale_rcp1, row_scale_rcp1); if (in_bounds(j + 1)) store_quantized(row_out_ptr + (unsigned)(j + 1) * (unsigned)K, row_out); } @@ -464,7 +472,7 @@ __device__ __forceinline__ void quantize_regtile( unsigned row_out[kWordsPerLane / 2]; #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = quantize_e4m3x4(v[2 * m], v[2 * m + 1], row_scale_rcp, row_scale_rcp); + row_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], row_scale_rcp, row_scale_rcp); if (in_bounds(j)) store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); } } @@ -486,7 +494,11 @@ __device__ __forceinline__ void quantize_regtile( col_amax = abs_max_bf16x2(col_amax, col_amax_partials[w][tid]); col_amax &= 0x7fff7fffu; const ptx::floatx2 amax_2x = widen_bf16x2(col_amax); - const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + // The tensor amax, when asked for: the max of these column amaxes. Like the generic + // kernel's per-element fmaxf, it ignores NaN (a column of only NaNs gives NaN here, + // which fmaxf drops) and is unaffected by the zeros loaded past M or K. + if constexpr (kAmax) thread_amax = fmaxf(thread_amax, fmaxf(amax_2x.x, amax_2x.y)); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); if (!kEdge || col0 + 2 * tid < colwise_scale_cols) { if constexpr (WITH_GEMM_SWIZZLED_SCALES) { // The colwise swizzle is the rowwise one with rows and columns swapped, @@ -562,8 +574,8 @@ __device__ __forceinline__ void quantize_regtile( unsigned col_out[kWordsPerLane / 2]; #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) - col_out[m] = quantize_e4m3x4(v[2 * m], v[2 * m + 1], col_scale_rcp[2 * m], - col_scale_rcp[2 * m + 1]); + col_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], col_scale_rcp[2 * m], + col_scale_rcp[2 * m + 1]); if (in_bounds(j)) store_quantized(col_out_ptr + (unsigned)j * (unsigned)K, col_out); } } @@ -583,10 +595,15 @@ __device__ __forceinline__ void quantize_regtile( } } } + + if constexpr (kAmax) { + const float block_amax = reduce_max(thread_amax, warp); + if (tid == 0) atomicMaxFloat(amax_ptr, block_amax); + } } -template +template __global__ void __launch_bounds__(TileConfig::kThreadsPerCta, TileConfig::kMinBlocksPerSm) quantize_mxfp8_kernel(const unsigned* __restrict__ input, @@ -596,12 +613,13 @@ __global__ void __launch_bounds__(TileConfig::k unsigned char* __restrict__ colwise_out, unsigned char* __restrict__ colwise_scales, int M, int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, - const void* __restrict__ table, float* __restrict__ dbias_workspace) { + const void* __restrict__ table, float* __restrict__ dbias_workspace, + float* __restrict__ amax_ptr) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - quantize_regtile( - (const uint4*)input, (const uint4*)act_input, rowwise_out, rowwise_scales, colwise_out, - colwise_scales, M, K, rowwise_scale_stride, colwise_scale_stride, iters, table, - dbias_workspace); + quantize_regtile((const uint4*)input, (const uint4*)act_input, rowwise_out, rowwise_scales, + colwise_out, colwise_scales, M, K, rowwise_scale_stride, + colwise_scale_stride, iters, table, dbias_workspace, amax_ptr); #endif } @@ -627,8 +645,8 @@ static int pick_walk_length(int row_blocks, int grid_cols, long long target, // the rest of the unified array goes to L1. constexpr int kCarveoutPercentAct = 40; -template +template static void set_carveout() { constexpr int kPercent = IS_ACT ? kCarveoutPercentAct : 0; constexpr int kDynamicBytes = TileConfig::kDbiasStageBytes; @@ -641,8 +659,9 @@ static void set_carveout() { std::lock_guard lock(mutex); if (static_cast(device) >= done.size()) done.resize(device + 1, false); if (!done[device]) { - const void* kernel = (const void*) - quantize_mxfp8_kernel; + const void* kernel = + (const void*)quantize_mxfp8_kernel; if constexpr (kPercent > 0) { NVTE_CHECK_CUDA( cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, kPercent)); @@ -684,23 +703,34 @@ int walk_length(int M, int K) { return iters; } -template +template static void launch(const void* input, const void* act_input, void* rowwise_out, void* rowwise_scales, void* colwise_out, void* colwise_scales, int M, int K, int rowwise_scale_stride, int colwise_scale_stride, const void* table, - float* dbias_workspace, cudaStream_t stream) { + float* dbias_workspace, float* amax, cudaStream_t stream) { using C = TileConfig; const int iters = walk_length(M, K); dim3 grid(col_tiles(K), row_bands(M) / iters); - const auto run = [&](auto tails) { + // kAmax: a separate instantiation, so that the kernels without amax keep their register + // budget (the derivatives run at the 40-register cap). + const auto run_amax = [&](auto tails, auto with_amax) { constexpr bool kTails = decltype(tails)::value; - set_carveout(); - quantize_mxfp8_kernel - <<>>( - (const unsigned*)input, (const unsigned*)act_input, (unsigned char*)rowwise_out, - (unsigned char*)rowwise_scales, (unsigned char*)colwise_out, - (unsigned char*)colwise_scales, M, K, rowwise_scale_stride, colwise_scale_stride, iters, - table, dbias_workspace); + constexpr bool kAmax = decltype(with_amax)::value; + set_carveout(); + quantize_mxfp8_kernel<<>>( + (const unsigned*)input, (const unsigned*)act_input, (unsigned char*)rowwise_out, + (unsigned char*)rowwise_scales, (unsigned char*)colwise_out, (unsigned char*)colwise_scales, + M, K, rowwise_scale_stride, colwise_scale_stride, iters, table, dbias_workspace, amax); + }; + const auto run = [&](auto tails) { + if (amax != nullptr) { + run_amax(tails, std::true_type{}); + } else { + run_amax(tails, std::false_type{}); + } }; // Whole 64-row tiles and 256-column tiles leave nothing past M or K. The derivatives take the // tail-capable kernel for every shape: with the band body behind its edge branch, the compiler @@ -733,9 +763,10 @@ constexpr bool kImplements = (!IS_DBIAS || IS_DACT) && (IS_ACT != IS_DACT) && activation_table::kHasTable && (IS_ACT == activation_table::Table>::kForward); -// Whether this kernel takes a request of a type in kImplements, with BF16 input and E4M3 -// output. Anything else takes the generic kernel, as does every request with -// NVTE_MXFP8_REGTILE=0. +// Whether this kernel takes a request of a type in kImplements, with BF16 input and E4M3 or +// E5M2 output. Anything else takes the generic kernel, as does every request with +// NVTE_MXFP8_REGTILE=0. A noop tensor is ignored, as the generic kernel ignores it for every +// request that computes an activation. template bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, const Tensor& noop, bool use_2d_quantization, const activation_table::TableRef& table) { @@ -758,8 +789,7 @@ bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, is_aligned_ptr(output.scale_inv.dptr, 8) && is_aligned_ptr(output.columnwise_scale_inv.dptr, 2); if (!(transformer_engine::cuda::sm_arch() >= 100 && !use_2d_quantization && output.has_data() && - output.has_columnwise_data() && output.amax.dptr == nullptr && noop.data.dptr == nullptr && - shape_supported(rows, cols) && aligned)) { + output.has_columnwise_data() && shape_supported(rows, cols) && aligned)) { return false; } constexpr size_t kMaxGridY = 65535; @@ -771,7 +801,7 @@ bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, //! dbias_workspace: [ceil(rows / 64), cols] FP32 column sums for common::reduce_dbias, or //! nullptr without dbias. template + float (*OP)(float, const ParamOP&), typename OType, bool WITH_GEMM_SWIZZLED_SCALES> void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, const activation_table::TableRef& table, float* dbias_workspace, cudaStream_t stream) { @@ -782,12 +812,12 @@ void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, const void* act_ptr = IS_DACT ? act_input->data.dptr : input.data.dptr; const auto run = [&](auto table_type) { using TableT = typename decltype(table_type)::type; - launch( + launch( input.data.dptr, act_ptr, output->data.dptr, output->scale_inv.dptr, output->columnwise_data.dptr, output->columnwise_scale_inv.dptr, static_cast(rows), static_cast(cols), static_cast(output->scale_inv.shape[1]), static_cast(output->columnwise_scale_inv.shape[1]), table.ptr, dbias_workspace, - stream); + static_cast(output->amax.dptr), stream); }; constexpr Kind kKind = kKindOf; if (table.use == kNoTable) { From 8dcd3adaa3aac1776d78dc49f6d1b3e8972a82bb Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 16 Oct 2026 10:45:02 -0700 Subject: [PATCH 10/29] [Common] Take FP16 inputs in the register-tiled MXFP8 kernel for dSiLU and dbias FP16 inputs of dSiLU, dbias_dGeLU and dbias_dSiLU now take the register-tiled kernel. FP16 has no table (the tables hold BF16 inputs), so the kernel computes the activation with util/math.h, and it scales every value in FP32 with mul_cvt_2x, as the generic kernel does, since the reciprocal scale is often outside FP16's range. The input-type-specific parts of the datapath are in Datapath; the BF16 path is unchanged. The FP32 scaling costs about twice the instructions of the BF16 path, which the compute-bound forward kernels and plain dGeLU cannot absorb (GB200: GeLU 0.74-0.91x, SiLU 0.86-1.26x, dGeLU 0.95-1.15x of TE main), so those keep the generic kernel. GB200 against TE main: dSiLU 1.26-1.74x, dbias_dGeLU 1.16-1.43x, dbias_dSiLU 1.29-1.51x. A new FP16 regtile-vs-generic suite covers every mode, output type, amax setting and tail shape. Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 58 ++++- .../common/cast/mxfp8/quantize_mxfp8.cuh | 8 +- .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 229 ++++++++++++------ 3 files changed, 212 insertions(+), 83 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index 321ea6c33e8..232ec2c56f2 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -385,16 +385,20 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t // also reduce the tensor amax, which must match bit for bit. void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, const size_t cols, const bool swizzled, const bool use_table, const DType otype, - const bool with_amax) { + const bool with_amax, const DType itype = DType::kBFloat16) { const std::vector shape{rows, cols}; const size_t n = rows * cols; - Tensor x("x", shape, DType::kBFloat16); - Tensor grad("grad", shape, DType::kBFloat16); - fill_inputs(fill, &x, &grad); + Tensor x("x", shape, itype); + Tensor grad("grad", shape, itype); + if (itype == DType::kFloat16) { + fill_inputs(fill, &x, &grad); + } else { + fill_inputs(fill, &x, &grad); + } Tensor outs[2] = {Tensor("default", shape, otype, true, true, NVTE_MXFP8_1D_SCALING), Tensor("generic", shape, otype, true, true, NVTE_MXFP8_1D_SCALING)}; - Tensor dbias[2] = {Tensor("dbias_default", std::vector{cols}, DType::kBFloat16), - Tensor("dbias_generic", std::vector{cols}, DType::kBFloat16)}; + Tensor dbias[2] = {Tensor("dbias_default", std::vector{cols}, itype), + Tensor("dbias_generic", std::vector{cols}, itype)}; Tensor workspace[2]; float *amax_dev = nullptr; if (with_amax) { @@ -466,11 +470,12 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, product(workspace[0].rowwise_shape()) * sizeof(float)), 0u) << "dbias workspace"; - EXPECT_EQ(count_byte_mismatches( - reinterpret_cast(dbias[0].rowwise_cpu_dptr()), - reinterpret_cast(dbias[1].rowwise_cpu_dptr()), - cols * sizeof(bf16)), - 0u) + const auto dbias_bytes = [&](Tensor &t) { + return itype == DType::kFloat16 + ? reinterpret_cast(t.rowwise_cpu_dptr()) + : reinterpret_cast(t.rowwise_cpu_dptr()); + }; + EXPECT_EQ(count_byte_mismatches(dbias_bytes(dbias[0]), dbias_bytes(dbias[1]), cols * 2), 0u) << "dbias"; } // The FP8 data as raw bytes, whichever FP8 type the output holds. @@ -623,3 +628,34 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), ::testing::Bool()), regtile_case_name); + +// FP16 inputs take the register-tiled kernel too, always computing the activation (the tables +// hold BF16 inputs). The table setting must not matter for them. +class CastMXFP8ActRegtileFP16TestSuite : public ::testing::TestWithParam {}; + +TEST_P(CastMXFP8ActRegtileFP16TestSuite, MatchesGenericKernel) { + cudaDeviceProp prop; + ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); + if (prop.major < 10) { + GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; + } + const auto [mode, fill, swizzled, dims, use_table, otype, with_amax] = GetParam(); + run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled, use_table, otype, + with_amax, DType::kFloat16); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, CastMXFP8ActRegtileFP16TestSuite, + ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::DBiasDGeLU, Mode::SiLU, + Mode::DSiLU, Mode::DBiasDSiLU), + ::testing::Values(Fill::Typical, Fill::OrderedInputs, Fill::Specials), + ::testing::Bool(), + ::testing::Values(std::pair{1024, 2048}, + std::pair{993, 7200}, + std::pair{33, 272}, + // Large enough for the computed forward path. + std::pair{4127, 2064}), + ::testing::Bool(), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), + ::testing::Bool()), + regtile_case_name); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index ab98580a9fb..c6f145c380e 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -966,18 +966,20 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, } // regtile::quantize replaces the generic kernel below for the requests it - // implements: bidimensional BF16 -> E4M3 or E5M2 with fused GeLU, SiLU, dGeLU or + // implements: bidimensional BF16 or FP16 -> E4M3 or E5M2 with fused GeLU, SiLU, dGeLU or // dSiLU, read from an activation table or computed, the derivatives optionally with // dbias, with or without amax, cols % 16 == 0, 16-byte aligned inputs (see // regtile::kImplements and regtile::can_use). It keeps the tile in registers instead // of staging it through shared memory and writes the same bytes, the same amax and // the same dbias workspace the generic kernel writes. - if constexpr (std::is_same_v && + if constexpr ((std::is_same_v || + (std::is_same_v && + regtile::kTakesFP16)) && (std::is_same_v || std::is_same_v) && regtile::kImplements) { if (regtile::can_use( input, act_input, *output, *noop, use_2d_quantization, activation_lut)) { - regtile::quantize( input, act_input, output, activation_lut, workspace_ptr, stream); if constexpr (IS_DBIAS) { diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh index 236031e4fdd..a1d90cacff1 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -54,7 +54,8 @@ namespace ptx = transformer_engine::ptx; namespace { -// The datapath works on 32-bit words, each holding a BF16 column pair. +// The datapath works on 32-bit words, each holding a column pair of the input type (BF16 +// or FP16; see Datapath). /*! \brief Elementwise max(|a|, |b|) on a packed BF16 pair, as a raw word. */ __device__ __forceinline__ unsigned abs_max_bf16x2(unsigned a, unsigned b) { @@ -105,12 +106,15 @@ struct ComputedTable : activation_table::Table { static constexpr size_t kBytes = 16; // unused shared-memory slot }; -/*! \brief The derivative of a BF16 pair, from the table in shared memory. */ -template +template +struct Datapath; // below + +/*! \brief The derivative of an input pair, from the table in shared memory or computed. */ +template __device__ __forceinline__ ptx::floatx2 dact_word(unsigned x, const unsigned char* __restrict__ table) { if constexpr (TableT::kComputed) { - const ptx::floatx2 v = widen_bf16x2(x); + const ptx::floatx2 v = Datapath::widen(x); return {TableT::op(v.x), TableT::op(v.y)}; } else { const float2 d = activation_table::derivative_2x(x, table); @@ -150,7 +154,7 @@ constexpr int kRowsPerMxBlock = 32; constexpr int kMxGroupsPerRow = kTileCols / 32; //! A register-resident tile spreads its 256 columns over the 32 lanes of a warp. constexpr int kColsPerLane = kTileCols / 32; -//! ...which a lane holds as BF16 pairs, one 32-bit word each. +//! ...which a lane holds as input-type pairs, one 32-bit word each. constexpr int kWordsPerLane = kColsPerLane / 2; //! ...read as 128-bit vectors, so one per row. constexpr int kVecLoadsPerRow = kWordsPerLane / 4; @@ -176,6 +180,77 @@ __device__ __forceinline__ void store_quantized(unsigned char* p, const unsigned *(uint2*)p = *(const uint2*)v; } +/*! \brief The input-type-specific part of the datapath. A 32-bit word holds a column pair of + * IType; Scale is the reciprocal scale of such a pair, in the form its quantize + * instruction takes. + */ +template +struct Datapath; + +template <> +struct Datapath { + //! A BF16 pair. 2^(127-e) is a power of two inside BF16's range, so the BF16 multiply in + //! mul_cvt_4x is exact, like the generic kernel's FP32 one. + using Scale = unsigned; + __device__ __forceinline__ static unsigned abs_max(unsigned a, unsigned b) { + return abs_max_bf16x2(a, b); + } + __device__ __forceinline__ static ptx::floatx2 widen(unsigned v) { return widen_bf16x2(v); } + __device__ __forceinline__ static unsigned narrow(const ptx::floatx2& v) { return to_bf16x2(v); } + //! The reciprocal of the scale e, for both columns of a pair. + __device__ __forceinline__ static Scale row_rcp(unsigned e) { + return e8m0_to_bf16x2_reciprocal(e); + } + //! The reciprocals of a column pair's scales, low byte for the low column. + __device__ __forceinline__ static Scale col_rcp(unsigned scale_pair) { + const ptx::bf16x2 r = ptx::exp2f_rcp_2x_per_lane(__byte_perm(scale_pair, 0, 0x4140)); + return reinterpret_cast(r); + } + template + __device__ __forceinline__ static unsigned quantize(unsigned in01, unsigned in23, Scale s01, + Scale s23) { + return quantize_x4(in01, in23, s01, s23); + } +}; + +template <> +struct Datapath { + //! An FP32 pair: 2^(127-e) is often outside FP16's range, so the multiply is done in FP32, + //! as the generic kernel does it (mul_cvt_2x with an FP32 scale pair). + using Scale = ptx::floatx2; + __device__ __forceinline__ static unsigned abs_max(unsigned a, unsigned b) { + ptx::fp16x2 d; + ptx::abs_max_2x(d, reinterpret_cast(a), + reinterpret_cast(b)); + return reinterpret_cast(d); + } + __device__ __forceinline__ static ptx::floatx2 widen(unsigned v) { + const float2 f = __half22float2(reinterpret_cast(v)); + return {f.x, f.y}; + } + __device__ __forceinline__ static unsigned narrow(const ptx::floatx2& v) { + const __half2 h = __floats2half2_rn(v.x, v.y); + return reinterpret_cast(h); + } + __device__ __forceinline__ static Scale row_rcp(unsigned e) { + const float r = ptx::exp2f_rcp(static_cast(e)); + return {r, r}; + } + __device__ __forceinline__ static Scale col_rcp(unsigned scale_pair) { + return {ptx::exp2f_rcp(static_cast(scale_pair & 0xffu)), + ptx::exp2f_rcp(static_cast((scale_pair >> 8) & 0xffu))}; + } + template + __device__ __forceinline__ static unsigned quantize(unsigned in01, unsigned in23, Scale s01, + Scale s23) { + ptx::FPx2 lo, hi; + ptx::mul_cvt_2x(lo, reinterpret_cast(in01), s01); + ptx::mul_cvt_2x(hi, reinterpret_cast(in23), s23); + return static_cast(reinterpret_cast(lo)) | + (static_cast(reinterpret_cast(hi)) << 16); + } +}; + // Minimum resident CTAs per SM, the second __launch_bounds__ argument, capped // at what the SM can hold: sm_75 and sm_107 have half the threads per SM of sm_100. constexpr int kMinBlocksAct = 8; @@ -226,7 +301,7 @@ struct TileConfig { // shared memory, and one thread per column adds them in row order. // // kAmax: also fold the amax of the output into *amax_ptr. -template __device__ __forceinline__ void quantize_regtile( const uint4* __restrict__ input, const uint4* __restrict__ act_input, @@ -235,11 +310,13 @@ __device__ __forceinline__ void quantize_regtile( int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* table, float* __restrict__ dbias_workspace, float* __restrict__ amax_ptr) { using C = TileConfig; + using DP = Datapath; + using Scale = typename DP::Scale; constexpr int kWarpsPerCta = C::kWarpsPerCta, kRowsPerWarp = C::kRowsPerWarp, kRowsInFlight = C::kRowsInFlight; __shared__ __align__(16) unsigned col_amax_partials[kWarpsPerCta][kTileCols / 2]; - __shared__ __align__(16) unsigned col_scale_rcp_smem[kTileCols / 2]; + __shared__ __align__(16) Scale col_scale_rcp_smem[kTileCols / 2]; __shared__ __align__(8) unsigned char row_scale_bytes[kRowsPerMxBlock * kMxGroupsPerRow]; __shared__ __align__(16) unsigned char table_smem[C::kTableBytes]; __shared__ __align__(8) uint64_t table_barrier; @@ -250,7 +327,7 @@ __device__ __forceinline__ void quantize_regtile( const int lane = tid & 31; const int warp = tid >> 5; const int col0 = blockIdx.x * kTileCols; - const int vecs_per_row = K >> 3; // uint4 (8 bf16 values) per row + const int vecs_per_row = K >> 3; // uint4 (8 BF16 or FP16 values) per row // Tails: the last column tile may end inside the tile, and row bands may end past M. Out-of-range // elements load as 0, which every activation maps to 0, so they never raise an amax. The // generic kernel writes scales for its whole 64x64 tiles, and so does this one: rowwise scales @@ -285,7 +362,7 @@ __device__ __forceinline__ void quantize_regtile( const int warp_row0 = row0 + warp * kRowsPerWarp; // Whether this lane's elements of row warp_row0 + j are inside the tensor. const auto in_bounds = [&](int j) { return !kEdge || (lane_in_cols && warp_row0 + j < M); }; - // This thread's kRowsPerWarp x kColsPerLane sub-tile, as BF16 pairs. + // This thread's kRowsPerWarp x kColsPerLane sub-tile, as input-type pairs. unsigned tile_words[kRowsPerWarp * kWordsPerLane]; const size_t load_offset = (size_t)(row0 + warp * kRowsPerWarp) * vecs_per_row + (col0 >> 3) + (size_t)kVecLoadsPerRow * lane; @@ -333,8 +410,8 @@ __device__ __forceinline__ void quantize_regtile( for (int m = 0; m < 4; ++m) { if constexpr (TableT::kComputed) { // util/math.h, then rounded to BF16 as the generic kernel rounds it. - const ptx::floatx2 v = widen_bf16x2(x_words[m]); - tile_words[tile_word_base + m] = to_bf16x2({TableT::op(v.x), TableT::op(v.y)}); + const ptx::floatx2 v = DP::widen(x_words[m]); + tile_words[tile_word_base + m] = DP::narrow({TableT::op(v.x), TableT::op(v.y)}); } else { tile_words[tile_word_base + m] = activation_table::forward_2x( x_words[m], reinterpret_cast(table_smem)); @@ -343,13 +420,13 @@ __device__ __forceinline__ void quantize_regtile( } } } else { -#define MXFP8_DACT_WORD(T, M) \ - { \ - const ptx::floatx2 product = \ - ptx::mul_2x(widen_bf16x2(grad_words[M]), dact_word(x_words[M], table_smem)); \ - tile_words[(h * kRowsInFlight + (T)) * kWordsPerLane + (M)] = to_bf16x2(product); \ - products[2 * (M)] = product.x; \ - products[2 * (M) + 1] = product.y; \ +#define MXFP8_DACT_WORD(T, M) \ + { \ + const ptx::floatx2 product = \ + ptx::mul_2x(DP::widen(grad_words[M]), dact_word(x_words[M], table_smem)); \ + tile_words[(h * kRowsInFlight + (T)) * kWordsPerLane + (M)] = DP::narrow(product); \ + products[2 * (M)] = product.x; \ + products[2 * (M) + 1] = product.y; \ } #define MXFP8_DACT_ROW(T) \ { \ @@ -381,41 +458,38 @@ __device__ __forceinline__ void quantize_regtile( const size_t store_offset = (size_t)(row0 + warp * kRowsPerWarp) * K + col0 + kColsPerLane * lane; auto row_scale = [&](const unsigned* row, int j) { - unsigned amax = - abs_max_bf16x2(abs_max_bf16x2(row[0], row[1]), abs_max_bf16x2(row[2], row[3])); + unsigned amax = DP::abs_max(DP::abs_max(row[0], row[1]), DP::abs_max(row[2], row[3])); #pragma unroll for (int m = 4; m < kWordsPerLane; m += 2) - amax = abs_max_bf16x2(amax, abs_max_bf16x2(row[m], row[m + 1])); + amax = DP::abs_max(amax, DP::abs_max(row[m], row[m + 1])); // A 32-value rowwise block spans kLanesPerMxGroup lanes. #pragma unroll for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) - amax = abs_max_bf16x2(amax, __shfl_xor_sync(0xffffffffu, amax, lane_mask)); + amax = DP::abs_max(amax, __shfl_xor_sync(0xffffffffu, amax, lane_mask)); // Fold the two halves with one PRMT and a packed max, which leaves the // magnitude in both halves; the high half is then the FP32 bit pattern. - amax = abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); - const unsigned scale_byte = - amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); + amax = DP::abs_max(amax, __byte_perm(amax, amax, 0x1032)); + const unsigned scale_byte = amax_to_e8m0_2x(0.f, DP::widen(amax & 0x7fff7fffu).y); if ((lane & (kLanesPerMxGroup - 1)) == 0) row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + (lane / kLanesPerMxGroup)] = (unsigned char)scale_byte; - return e8m0_to_bf16x2_reciprocal(scale_byte); + return DP::row_rcp(scale_byte); }; // On the 128-thread activation tile, two rows share one butterfly and conversion. auto row_amax = [&](const unsigned* row) { - unsigned amax = - abs_max_bf16x2(abs_max_bf16x2(row[0], row[1]), abs_max_bf16x2(row[2], row[3])); + unsigned amax = DP::abs_max(DP::abs_max(row[0], row[1]), DP::abs_max(row[2], row[3])); #pragma unroll for (int m = 4; m < kWordsPerLane; m += 2) - amax = abs_max_bf16x2(amax, abs_max_bf16x2(row[m], row[m + 1])); - return abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); + amax = DP::abs_max(amax, DP::abs_max(row[m], row[m + 1])); + return DP::abs_max(amax, __byte_perm(amax, amax, 0x1032)); }; auto row_scale_pair = [&](const unsigned* row_a, const unsigned* row_b, int j, - unsigned& row_scale_rcp0, unsigned& row_scale_rcp1) { + Scale& row_scale_rcp0, Scale& row_scale_rcp1) { unsigned amax_pair = __byte_perm(row_amax(row_a), row_amax(row_b), 0x5410) & 0x7fff7fffu; #pragma unroll for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) - amax_pair = abs_max_bf16x2(amax_pair, __shfl_xor_sync(0xffffffffu, amax_pair, lane_mask)); - const ptx::floatx2 amax_2x = widen_bf16x2(amax_pair); + amax_pair = DP::abs_max(amax_pair, __shfl_xor_sync(0xffffffffu, amax_pair, lane_mask)); + const ptx::floatx2 amax_2x = DP::widen(amax_pair); const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); const unsigned scale_byte0 = scale_pair & 0xffu, scale_byte1 = (scale_pair >> 8) & 0xffu; if ((lane & (kLanesPerMxGroup - 1)) == 0) { @@ -424,8 +498,8 @@ __device__ __forceinline__ void quantize_regtile( scale_dst[0] = (unsigned char)scale_byte0; scale_dst[kMxGroupsPerRow] = (unsigned char)scale_byte1; } - row_scale_rcp0 = e8m0_to_bf16x2_reciprocal(scale_byte0); - row_scale_rcp1 = e8m0_to_bf16x2_reciprocal(scale_byte1); + row_scale_rcp0 = DP::row_rcp(scale_byte0); + row_scale_rcp1 = DP::row_rcp(scale_byte1); }; // Colwise amax: per-warp partial in registers, cross-warp fold in shared. auto publish_col_amax_partials = [&]() { @@ -436,7 +510,7 @@ __device__ __forceinline__ void quantize_regtile( for (int j = 1; j < kRowsPerWarp; ++j) #pragma unroll for (int m = 0; m < kWordsPerLane; ++m) - col_amax[m] = abs_max_bf16x2(col_amax[m], tile_words[j * kWordsPerLane + m]); + col_amax[m] = DP::abs_max(col_amax[m], tile_words[j * kWordsPerLane + m]); #pragma unroll for (int n = 0; n < kVecLoadsPerRow; ++n) *(uint4*)(&col_amax_partials[warp][kWordsPerLane * lane + 4 * n]) = @@ -448,18 +522,18 @@ __device__ __forceinline__ void quantize_regtile( for (int j = 0; j < kRowsPerWarp; j += 2) { const unsigned* row_a = tile_words + j * kWordsPerLane; const unsigned* row_b = row_a + kWordsPerLane; - unsigned row_scale_rcp0, row_scale_rcp1; + Scale row_scale_rcp0, row_scale_rcp1; row_scale_pair(row_a, row_b, j, row_scale_rcp0, row_scale_rcp1); unsigned row_out[kWordsPerLane / 2]; #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = - quantize_x4(row_a[2 * m], row_a[2 * m + 1], row_scale_rcp0, row_scale_rcp0); + row_out[m] = DP::template quantize(row_a[2 * m], row_a[2 * m + 1], + row_scale_rcp0, row_scale_rcp0); if (in_bounds(j)) store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = - quantize_x4(row_b[2 * m], row_b[2 * m + 1], row_scale_rcp1, row_scale_rcp1); + row_out[m] = DP::template quantize(row_b[2 * m], row_b[2 * m + 1], + row_scale_rcp1, row_scale_rcp1); if (in_bounds(j + 1)) store_quantized(row_out_ptr + (unsigned)(j + 1) * (unsigned)K, row_out); } @@ -468,11 +542,12 @@ __device__ __forceinline__ void quantize_regtile( #pragma unroll for (int j = 0; j < kRowsPerWarp; ++j) { const unsigned* v = tile_words + j * kWordsPerLane; - const unsigned row_scale_rcp = row_scale(v, j); + const Scale row_scale_rcp = row_scale(v, j); unsigned row_out[kWordsPerLane / 2]; #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], row_scale_rcp, row_scale_rcp); + row_out[m] = + DP::template quantize(v[2 * m], v[2 * m + 1], row_scale_rcp, row_scale_rcp); if (in_bounds(j)) store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); } } @@ -491,9 +566,9 @@ __device__ __forceinline__ void quantize_regtile( unsigned col_amax = col_amax_partials[0][tid]; #pragma unroll for (int w = 1; w < kWarpsPerCta; ++w) - col_amax = abs_max_bf16x2(col_amax, col_amax_partials[w][tid]); + col_amax = DP::abs_max(col_amax, col_amax_partials[w][tid]); col_amax &= 0x7fff7fffu; - const ptx::floatx2 amax_2x = widen_bf16x2(col_amax); + const ptx::floatx2 amax_2x = DP::widen(col_amax); // The tensor amax, when asked for: the max of these column amaxes. Like the generic // kernel's per-element fmaxf, it ignores NaN (a column of only NaNs gives NaN here, // which fmaxf drops) and is unaffected by the zeros loaded past M or K. @@ -514,9 +589,7 @@ __device__ __forceinline__ void quantize_regtile( 2 * tid) = (unsigned short)scale_pair; } } - const ptx::bf16x2 scale_rcp = - ptx::exp2f_rcp_2x_per_lane(__byte_perm(scale_pair, 0, 0x4140)); - col_scale_rcp_smem[tid] = reinterpret_cast(scale_rcp); + col_scale_rcp_smem[tid] = DP::col_rcp(scale_pair); } auto drain_row_scales = [&]() { constexpr int drain_tid_base = C::kDrainTidBase; @@ -562,11 +635,14 @@ __device__ __forceinline__ void quantize_regtile( // ---- columnwise quantization -------------------------------------------- { - uint4 col_scale_rcp_vec[kVecLoadsPerRow]; + // This lane's kWordsPerLane column-pair scales, as 16-byte vectors. + constexpr int kScaleVecs = kWordsPerLane * sizeof(Scale) / sizeof(uint4); + uint4 col_scale_rcp_vec[kScaleVecs]; #pragma unroll - for (int n = 0; n < kVecLoadsPerRow; ++n) - col_scale_rcp_vec[n] = *(const uint4*)(col_scale_rcp_smem + kWordsPerLane * lane + 4 * n); - const unsigned* col_scale_rcp = (const unsigned*)col_scale_rcp_vec; + for (int n = 0; n < kScaleVecs; ++n) + col_scale_rcp_vec[n] = + reinterpret_cast(col_scale_rcp_smem + kWordsPerLane * lane)[n]; + const Scale* col_scale_rcp = reinterpret_cast(col_scale_rcp_vec); unsigned char* col_out_ptr = colwise_out + store_offset; #pragma unroll for (int j = 0; j < kRowsPerWarp; ++j) { @@ -574,8 +650,8 @@ __device__ __forceinline__ void quantize_regtile( unsigned col_out[kWordsPerLane / 2]; #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) - col_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], col_scale_rcp[2 * m], - col_scale_rcp[2 * m + 1]); + col_out[m] = DP::template quantize(v[2 * m], v[2 * m + 1], col_scale_rcp[2 * m], + col_scale_rcp[2 * m + 1]); if (in_bounds(j)) store_quantized(col_out_ptr + (unsigned)j * (unsigned)K, col_out); } } @@ -602,7 +678,7 @@ __device__ __forceinline__ void quantize_regtile( } } -template __global__ void __launch_bounds__(TileConfig::kThreadsPerCta, TileConfig::kMinBlocksPerSm) @@ -616,7 +692,7 @@ __global__ void __launch_bounds__(TileConfig::k const void* __restrict__ table, float* __restrict__ dbias_workspace, float* __restrict__ amax_ptr) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - quantize_regtile((const uint4*)input, (const uint4*)act_input, rowwise_out, rowwise_scales, colwise_out, colwise_scales, M, K, rowwise_scale_stride, colwise_scale_stride, iters, table, dbias_workspace, amax_ptr); @@ -645,7 +721,7 @@ static int pick_walk_length(int row_blocks, int grid_cols, long long target, // the rest of the unified array goes to L1. constexpr int kCarveoutPercentAct = 40; -template static void set_carveout() { constexpr int kPercent = IS_ACT ? kCarveoutPercentAct : 0; @@ -660,7 +736,7 @@ static void set_carveout() { if (static_cast(device) >= done.size()) done.resize(device + 1, false); if (!done[device]) { const void* kernel = - (const void*)quantize_mxfp8_kernel; if constexpr (kPercent > 0) { NVTE_CHECK_CUDA( @@ -703,7 +779,7 @@ int walk_length(int M, int K) { return iters; } -template static void launch(const void* input, const void* act_input, void* rowwise_out, void* rowwise_scales, void* colwise_out, void* colwise_scales, int M, int K, @@ -717,10 +793,10 @@ static void launch(const void* input, const void* act_input, void* rowwise_out, const auto run_amax = [&](auto tails, auto with_amax) { constexpr bool kTails = decltype(tails)::value; constexpr bool kAmax = decltype(with_amax)::value; - set_carveout(); - quantize_mxfp8_kernel<<>>( + quantize_mxfp8_kernel<<>>( (const unsigned*)input, (const unsigned*)act_input, (unsigned char*)rowwise_out, (unsigned char*)rowwise_scales, (unsigned char*)colwise_out, (unsigned char*)colwise_scales, M, K, rowwise_scale_stride, colwise_scale_stride, iters, table, dbias_workspace, amax); @@ -763,7 +839,16 @@ constexpr bool kImplements = (!IS_DBIAS || IS_DACT) && (IS_ACT != IS_DACT) && activation_table::kHasTable && (IS_ACT == activation_table::Table>::kForward); -// Whether this kernel takes a request of a type in kImplements, with BF16 input and E4M3 or +//! Whether FP16 inputs of this request type take this kernel. FP16 always computes the +//! activation and scales every value in FP32 (see Datapath), about twice the +//! instructions of the BF16 path, which the compute-bound kernels cannot absorb. GB200 against +//! TE main: dSiLU 1.25-1.71x, dbias_dGeLU 1.16-1.43x, dbias_dSiLU 1.27-1.50x, but GeLU +//! 0.74-0.91x, SiLU 0.86-1.26x and dGeLU 0.95-1.15x, which stay on the generic kernel. +template +constexpr bool kTakesFP16 = + IS_DBIAS || activation_table::kKindOf == activation_table::kDsilu; + +// Whether this kernel takes a request of a type in kImplements, with BF16 or FP16 input and E4M3 or // E5M2 output. Anything else takes the generic kernel, as does every request with // NVTE_MXFP8_REGTILE=0. A noop tensor is ignored, as the generic kernel ignores it for every // request that computes an activation. @@ -801,7 +886,8 @@ bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, //! dbias_workspace: [ceil(rows / 64), cols] FP32 column sums for common::reduce_dbias, or //! nullptr without dbias. template + float (*OP)(float, const ParamOP&), typename IType, typename OType, + bool WITH_GEMM_SWIZZLED_SCALES> void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, const activation_table::TableRef& table, float* dbias_workspace, cudaStream_t stream) { @@ -812,7 +898,7 @@ void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, const void* act_ptr = IS_DACT ? act_input->data.dptr : input.data.dptr; const auto run = [&](auto table_type) { using TableT = typename decltype(table_type)::type; - launch( + launch( input.data.dptr, act_ptr, output->data.dptr, output->scale_inv.dptr, output->columnwise_data.dptr, output->columnwise_scale_inv.dptr, static_cast(rows), static_cast(cols), static_cast(output->scale_inv.shape[1]), @@ -820,20 +906,25 @@ void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, static_cast(output->amax.dptr), stream); }; constexpr Kind kKind = kKindOf; - if (table.use == kNoTable) { + // The tables hold BF16 inputs; FP16 inputs always compute. + if (!std::is_same_v || table.use == kNoTable) { run(std::common_type>{}); NVTE_CHECK_CUDA(cudaGetLastError()); return; } - if constexpr (kKind == kDsilu) { + if constexpr (!std::is_same_v) { + return; + } else if constexpr (kKind == kDsilu) { if (table.use == kExactDsiluTable) { run(std::common_type{}); NVTE_CHECK_CUDA(cudaGetLastError()); return; } } - run(std::common_type>{}); - NVTE_CHECK_CUDA(cudaGetLastError()); + if constexpr (std::is_same_v) { + run(std::common_type>{}); + NVTE_CHECK_CUDA(cudaGetLastError()); + } } } // namespace regtile From 5186a431ba58980c110a59f3a3984b44ddd45df2 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 15 Oct 2026 18:06:59 -0700 Subject: [PATCH 11/29] [Common] Speed up the dbias reduction of the MXFP8 quantize common::reduce_dbias sums the FP32 dbias workspace column by column. It ran 256 threads per CTA with four columns each, so a 7168-column workspace got 7 CTAs, and each thread waited for one load per row. Use one thread per column, 64 per CTA, and copy each column into shared memory with cp.async, 32 rows per stage and four stages deep, adding a stage once it has landed. Each column is still added in row order from 0 and rounded once, so the result is unchanged; a new test checks it bit for bit against that sum on the host for FP32, BF16 and FP16 output, on data spanning 2^-30..2^30 with mixed signs. GB200, reduction of the 32768x7168 dbias_dgelu workspace (512x7168): 108.8 us -> 13.2 us. End to end against TE main: dbias_dgelu / dbias_dsilu / dbias 1.07-1.25x without tables and 1.00-1.31x with them across 8192x7168 to 993x7200. 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 | 84 ++++++++---- 3 files changed, 182 insertions(+), 29 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 24acec658c6..2ff5fa9ade3 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -4,6 +4,7 @@ add_executable(test_util test_activation_table.cu + 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..a5d26230131 --- /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 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: + // multiples of 4 (the output vector of the BF16/FP16 store), around the 64-thread CTA. + ::testing::Combine(::testing::Values(1, 2, 31, 32, 33, 63, 64, 65, 512, 1000, 2000), + ::testing::Values(4, 28, 60, 64, 68, 1028, 2052, 7168)), + [](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..94f3dbe124d 100644 --- a/transformer_engine/common/cast/core/grouped_layout.cuh +++ b/transformer_engine/common/cast/core/grouped_layout.cuh @@ -36,39 +36,66 @@ 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. + */ +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) { + __shared__ float stages[REDUCE_DBIAS_STAGES][REDUCE_DBIAS_STAGE_ROWS][REDUCE_DBIAS_THREADS]; + 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; + 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); + float acc = 0.0f; + 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); + dbias_output[col] = static_cast(acc); } template @@ -131,11 +158,10 @@ void reduce_dbias(const float *workspace_ptr, Tensor *dbias, const size_t rows, 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()); } From 1ff7b624b0a7473d935eb42743d50bba45366f3e Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 15 Oct 2026 18:05:05 -0700 Subject: [PATCH 12/29] [Common] Add a register-tiled MXFP8 kernel for SwiGLU and its gradient nvte_swiglu and nvte_dswiglu quantized to MXFP8 from BF16 to E4M3 in both directions now take a register-tiled kernel, like the one for GeLU and SiLU: a CTA owns 32-row bands of 256 output columns, keeps them in registers, and computes both scale directions from there. The backward quantizes the two gradient halves one after the other from the same band. The math, the rounding to BF16 and the scale writes (inside the tensor only) are those of the generic gated kernel, so the output is byte-identical. It takes any number of rows and K % 32 == 0, the generic kernel's requirement. NVTE_MXFP8_REGTILE=0 keeps the generic kernel. GB200, against TE main, normal data: SwiGLU 1.39-1.99x (5.6 TB/s on large shapes), dSwiGLU 1.39-1.59x; wide data up to 2.05x and 2.18x. The new test compares the two kernels byte for byte on regular and tail shapes with every output byte pre-filled. Signed-off-by: Teddy Do --- tests/cpp/operator/CMakeLists.txt | 1 + .../operator/test_cast_mxfp8_gated_regtile.cu | 176 ++++++++++ .../common/cast/mxfp8/gated_mxfp8.cuh | 15 + .../regtile/quantize_gated_mxfp8_regtile.cuh | 320 ++++++++++++++++++ 4 files changed, 512 insertions(+) create mode 100644 tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu create mode 100644 transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index e05919c2b53..0e540e5a03f 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -13,6 +13,7 @@ add_executable(test_operator test_qdq.cu test_cast_mxfp8.cu test_cast_mxfp8_act_exact.cu + test_cast_mxfp8_gated_regtile.cu test_cast_mxfp8_grouped.cu test_cast_mxfp8_grouped_scaled_swiglu.cu test_cast_nvfp4_transpose.cu diff --git a/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu b/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu new file mode 100644 index 00000000000..104a0406010 --- /dev/null +++ b/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu @@ -0,0 +1,176 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// The register-tiled SwiGLU / dSwiGLU MXFP8 kernel must write exactly the bytes of the generic +// gated kernel. Each case runs both (NVTE_MXFP8_REGTILE=0 keeps the generic kernel) into outputs +// pre-filled with a marker byte, so they must also write the same set of bytes. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../test_common.h" +#include "transformer_engine/transformer_engine.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +enum class GatedFill { Normal, Wide, Specials }; + +class ScopedEnv { + public: + ScopedEnv(const char *name, const char *value) : name_(name) { + const char *old = std::getenv(name); + if (old != nullptr) old_ = old; + setenv(name, value, 1); + } + ~ScopedEnv() { + if (old_) { + setenv(name_, old_->c_str(), 1); + } else { + unsetenv(name_); + } + } + + private: + const char *name_; + std::optional old_; +}; + +void fill(Tensor *t, const size_t rows, const size_t cols, const GatedFill kind, const int seed) { + std::mt19937 gen(seed); + std::normal_distribution normal(0.0f, 3.0f); + std::uniform_int_distribution exponent(-10, 9); + const float specials[] = {std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN(), 0.0f, -0.0f, 1e-30f, -1e30f}; + bf16 *p = t->rowwise_cpu_dptr(); + for (size_t r = 0; r < rows; ++r) { + // Wide: every 32x32 block gets its own power-of-two scale. + for (size_t c = 0; c < cols; ++c) { + float v = normal(gen); + if (kind == GatedFill::Wide) { + std::mt19937 block_gen(seed * 7919 + (r / 32) * 104729 + c / 32); + v *= std::ldexp(1.0f, exponent(block_gen)); + } else if (kind == GatedFill::Specials && gen() % 61 == 0) { + v = specials[gen() % 7]; + } + p[r * cols + c] = static_cast(v); + } + } + t->from_cpu(); +} + +size_t count_byte_mismatches(const uint8_t *a, const uint8_t *b, const size_t n) { + size_t m = 0; + for (size_t i = 0; i < n; ++i) m += (a[i] != b[i]); + return m; +} + +void run_case(const bool bwd, const GatedFill kind, const size_t rows, const size_t cols, + const bool swizzled) { + Tensor input("input", std::vector{rows, 2 * cols}, DType::kBFloat16); + Tensor grad("grad", std::vector{rows, cols}, DType::kBFloat16); + fill(&input, rows, 2 * cols, kind, 1); + fill(&grad, rows, cols, kind == GatedFill::Specials ? GatedFill::Normal : kind, 2); + const size_t out_cols = bwd ? 2 * cols : cols; + const std::vector out_shape{rows, out_cols}; + Tensor outs[2] = { + Tensor("regtile", out_shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING), + Tensor("generic", out_shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING)}; + for (int i = 0; i < 2; ++i) { + outs[i].set_with_gemm_swizzled_scales(swizzled); + for (const auto param : {kNVTERowwiseData, kNVTEColumnwiseData, kNVTERowwiseScaleInv, + kNVTEColumnwiseScaleInv}) { + const NVTEBasicTensor b = nvte_get_tensor_param(outs[i].data(), param); + size_t bytes = b.data_ptr == nullptr ? 0 : 1; + for (size_t d = 0; d < b.shape.ndim; ++d) bytes *= b.shape.data[d]; + if (bytes != 0) ASSERT_EQ(cudaMemset(b.data_ptr, 0xA5, bytes), cudaSuccess); + } + const ScopedEnv regtile("NVTE_MXFP8_REGTILE", i == 0 ? "1" : "0"); + if (bwd) { + nvte_dswiglu(grad.data(), input.data(), outs[i].data(), 0); + } else { + nvte_swiglu(input.data(), outs[i].data(), 0); + } + } + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + for (auto &out : outs) out.to_cpu(); + const size_t n = rows * out_cols; + EXPECT_EQ(count_byte_mismatches( + reinterpret_cast(outs[0].rowwise_cpu_dptr()), + reinterpret_cast(outs[1].rowwise_cpu_dptr()), n), + 0u) + << "rowwise data"; + EXPECT_EQ(count_byte_mismatches( + reinterpret_cast(outs[0].columnwise_cpu_dptr()), + reinterpret_cast(outs[1].columnwise_cpu_dptr()), n), + 0u) + << "colwise data"; + EXPECT_EQ(count_byte_mismatches(outs[0].rowwise_cpu_scale_inv_ptr(), + outs[1].rowwise_cpu_scale_inv_ptr(), + product(outs[0].rowwise_scale_inv_shape())), + 0u) + << "rowwise scales"; + EXPECT_EQ(count_byte_mismatches(outs[0].columnwise_cpu_scale_inv_ptr(), + outs[1].columnwise_cpu_scale_inv_ptr(), + product(outs[0].columnwise_scale_inv_shape())), + 0u) + << "colwise scales"; +} + +using Params = std::tuple>; + +class GatedRegtileTestSuite : public ::testing::TestWithParam {}; + +TEST_P(GatedRegtileTestSuite, MatchesGenericKernel) { + cudaDeviceProp prop; + ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); + if (prop.major < 10) { + GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; + } + const auto [bwd, kind, swizzled, dims] = GetParam(); + run_case(bwd, kind, dims.first, dims.second, swizzled); +} + +std::string case_name(const testing::TestParamInfo &info) { + const auto [bwd, kind, swizzled, dims] = info.param; + const char *fills[] = {"Normal", "Wide", "Specials"}; + return std::string(bwd ? "DSwiGLU" : "SwiGLU") + "X" + fills[static_cast(kind)] + "X" + + (swizzled ? "Swizzled" : "Compact") + "X" + std::to_string(dims.first) + "X" + + std::to_string(dims.second); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, GatedRegtileTestSuite, + ::testing::Combine(::testing::Bool(), + ::testing::Values(GatedFill::Normal, GatedFill::Wide, GatedFill::Specials), + ::testing::Bool(), + ::testing::Values(std::pair{1024, 2048}, + std::pair{4160, 2048}, + std::pair{128, 256}, + // Row and column tails (K % 32 == 0). + std::pair{993, 7200}, + std::pair{1, 32}, + std::pair{33, 288}, + std::pair{100, 2080})), + case_name); + +} // namespace diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 31e445d4fb4..17527c730c2 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -20,6 +20,7 @@ #include "../../util/math.h" #include "../../util/ptx_arch_spec.cuh" #include "../../utils.cuh" +#include "regtile/quantize_gated_mxfp8_regtile.cuh" #include "swizzle.cuh" namespace transformer_engine { @@ -850,6 +851,20 @@ void quantize_gated(const Tensor &gated_input, const Tensor &grad, Tensor *outpu } } + // regtile::gated replaces the kernel below for SwiGLU and its gradient from BF16 to + // E4M3 in both directions (see regtile::gated::can_use), writing the same bytes. + if constexpr (std::is_same_v && std::is_same_v && + std::is_same_v) { + if constexpr (ActOP == &silu && + (!IS_BWD || DActOP == &dsilu)) { + if (regtile::gated::can_use(gated_input, grad, *output, IS_BWD)) { + regtile::gated::quantize(gated_input, grad, + output, stream); + return; + } + } + } + switch (scaling_type) { case ScalingType::ROWWISE: { auto kernel = diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh new file mode 100644 index 00000000000..a7943f40a5c --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh @@ -0,0 +1,320 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_gated_mxfp8_regtile.cuh + * \brief Register-resident MXFP8 bidimensional quantize fused with SwiGLU or its gradient. + * + * The gated counterpart of quantize_mxfp8_regtile.cuh. Input [M, 2K] holds the activated half + * and the gate half side by side. The forward writes silu(act) * gate as [M, K]; the backward + * writes the gradients of both halves as [M, 2K], quantized independently, one after the other + * from the same 32-row band. A CTA covers 256 output columns and walks 32-row bands; the band + * stays in registers, and both scale directions are computed from there. + * + * It writes the bytes the generic gated kernel (gated_mxfp8.cuh) writes: the same math, the + * same rounding to BF16, and scales only inside the tensor. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_GATED_MXFP8_REGTILE_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_GATED_MXFP8_REGTILE_CUH_ + +#include "quantize_mxfp8_regtile.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace regtile { +namespace gated { + +constexpr int kThreadsPerCta = 256; +constexpr int kWarpsPerCta = kThreadsPerCta / 32; +constexpr int kRowsPerWarp = kRowsPerMxBlock / kWarpsPerCta; +constexpr int kRowsInFlight = 2; +// Resident CTAs per SM, the second __launch_bounds__ argument (GB200, 32768x7168): the forward +// is issue-bound and gains 16% from 6 CTAs over 4; the backward, which keeps more values live, +// loses 30% at 5. +template +constexpr int kMinBlocksPerSm = IS_BWD ? 4 : 6; +static_assert(kVecLoadsPerRow == 1, "One 16-byte vector per row and lane."); + +/*! \brief One output element of the gated activation, as the generic gated kernel computes it. + * + * kPart 0: the forward output, or the gradient of the activated half. kPart 1: the gradient of + * the gate half. Rounded to BF16 by the caller. + */ +template +__device__ __forceinline__ float gated_element(const float act_elt, const float gate_elt, + const float grad_elt) { + if constexpr (!IS_BWD) { + return silu(act_elt, {}) * gate_elt; + } else { + const float x = act_elt; + const float s = sigmoidf(x); + if constexpr (kPart == 0) { + const float dact_x = x * s * (1 - s) + s; + return dact_x * grad_elt * gate_elt; + } else { + const float act_x = x * s; + const float dgate_elt = 1.0f; + return dgate_elt * act_x * grad_elt; + } + } +} + +/*! \brief The gated activation of one BF16 pair of each input, rounded to a BF16 pair. */ +template +__device__ __forceinline__ unsigned gated_word(const unsigned act, const unsigned gate, + const unsigned grad) { + const ptx::floatx2 a = widen_bf16x2(act); + const ptx::floatx2 g = widen_bf16x2(gate); + const ptx::floatx2 d = widen_bf16x2(grad); + return to_bf16x2( + {gated_element(a.x, g.x, d.x), gated_element(a.y, g.y, d.y)}); +} + +template +__global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) + quantize_gated_kernel(const uint4* __restrict__ input, const uint4* __restrict__ grad, + unsigned char* __restrict__ rowwise_out, + unsigned char* __restrict__ rowwise_scales, + unsigned char* __restrict__ colwise_out, + unsigned char* __restrict__ colwise_scales, int M, int K, + int rowwise_scale_stride, int colwise_scale_stride, int iters) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + __shared__ __align__(16) unsigned col_amax_partials[kWarpsPerCta][kTileCols / 2]; + __shared__ __align__(16) unsigned col_scale_rcp_smem[kTileCols / 2]; + __shared__ __align__(8) unsigned char row_scale_bytes[kRowsPerMxBlock * kMxGroupsPerRow]; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int col0 = blockIdx.x * kTileCols; + // Output columns; the backward writes both halves. + const int out_cols = IS_BWD ? 2 * K : K; + const int k_vecs = K >> 3; // uint4 (8 BF16 values) per K columns + const int k_groups = (K + 31) / 32; + const bool full_width = !kTails || col0 + kTileCols <= K; + const bool lane_in_cols = col0 + kColsPerLane * lane < K; + +#pragma unroll 1 + for (int it = 0; it < iters; ++it) { + const int row0 = (it * gridDim.y + blockIdx.y) * kRowsPerMxBlock; + const auto band = [&](auto edge, auto part) { + constexpr bool kEdge = decltype(edge)::value; + constexpr int kPart = decltype(part)::value; + const int warp_row0 = row0 + warp * kRowsPerWarp; + const auto in_bounds = [&](int j) { return !kEdge || (lane_in_cols && warp_row0 + j < M); }; + unsigned tile_words[kRowsPerWarp * kWordsPerLane]; + const size_t in_offset = (size_t)warp_row0 * (2 * k_vecs) + (col0 >> 3) + lane; + const uint4* act_ptr = input + in_offset; + const uint4* gate_ptr = act_ptr + k_vecs; + const uint4* grad_ptr = grad + (size_t)warp_row0 * k_vecs + (col0 >> 3) + lane; + // The gate is read by the forward and by the gradient of the activated half. + constexpr bool kNeedsGate = !IS_BWD || kPart == 0; + +#pragma unroll + for (int h = 0; h < kRowsPerWarp / kRowsInFlight; ++h) { + uint4 act_vec[kRowsInFlight], gate_vec[kRowsInFlight], grad_vec[kRowsInFlight]; + asm volatile("" ::: "memory"); +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) { + const int j = h * kRowsInFlight + t; + const uint4 zero = make_uint4(0, 0, 0, 0); + act_vec[t] = in_bounds(j) ? act_ptr[(size_t)j * (2 * k_vecs)] : zero; + gate_vec[t] = kNeedsGate && in_bounds(j) ? gate_ptr[(size_t)j * (2 * k_vecs)] : zero; + grad_vec[t] = IS_BWD && in_bounds(j) ? grad_ptr[(size_t)j * k_vecs] : zero; + } + asm volatile("" ::: "memory"); +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) { + const unsigned* a = (const unsigned*)&act_vec[t]; + const unsigned* g = (const unsigned*)&gate_vec[t]; + const unsigned* d = (const unsigned*)&grad_vec[t]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + tile_words[(h * kRowsInFlight + t) * kWordsPerLane + m] = + gated_word(a[m], g[m], d[m]); + } + } + } + + // ---- rowwise ---------------------------------------------------------------------------- + const size_t store_offset = + (size_t)warp_row0 * out_cols + kPart * K + col0 + kColsPerLane * lane; + { + unsigned char* row_out_ptr = rowwise_out + store_offset; +#pragma unroll + for (int j = 0; j < kRowsPerWarp; ++j) { + const unsigned* v = tile_words + j * kWordsPerLane; + unsigned amax = abs_max_bf16x2(abs_max_bf16x2(v[0], v[1]), abs_max_bf16x2(v[2], v[3])); +#pragma unroll + for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) + amax = abs_max_bf16x2(amax, __shfl_xor_sync(0xffffffffu, amax, lane_mask)); + amax = abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); + const unsigned scale_byte = + amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); + if ((lane & (kLanesPerMxGroup - 1)) == 0) + row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + + (lane / kLanesPerMxGroup)] = (unsigned char)scale_byte; + const unsigned rcp = e8m0_to_bf16x2_reciprocal(scale_byte); + unsigned row_out[kWordsPerLane / 2]; +#pragma unroll + for (int m = 0; m < kWordsPerLane / 2; ++m) + row_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], rcp, rcp); + if (in_bounds(j)) store_quantized(row_out_ptr + (size_t)j * out_cols, row_out); + } + } + + // ---- colwise amax: per-warp partial in registers, cross-warp fold in shared --------------- + { + unsigned col_amax[kWordsPerLane]; +#pragma unroll + for (int m = 0; m < kWordsPerLane; ++m) col_amax[m] = tile_words[m]; +#pragma unroll + for (int j = 1; j < kRowsPerWarp; ++j) +#pragma unroll + for (int m = 0; m < kWordsPerLane; ++m) + col_amax[m] = abs_max_bf16x2(col_amax[m], tile_words[j * kWordsPerLane + m]); + *(uint4*)(&col_amax_partials[warp][kWordsPerLane * lane]) = *(const uint4*)col_amax; + } + __syncthreads(); + if (tid < kTileCols / 2) { + unsigned col_amax = col_amax_partials[0][tid]; +#pragma unroll + for (int w = 1; w < kWarpsPerCta; ++w) + col_amax = abs_max_bf16x2(col_amax, col_amax_partials[w][tid]); + col_amax &= 0x7fff7fffu; + const ptx::floatx2 amax_2x = widen_bf16x2(col_amax); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + const int col = col0 + 2 * tid; + // The generic gated kernel writes scales only inside the tensor. + if (!kEdge || (col < K && row0 < M)) { + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + const size_t row_block_tiles = DIVUP(static_cast(M), size_t{128}); + colwise_scales[swizzle::gemm_swizzled_scale_idx(kPart * K + col, row0 / 32, + row_block_tiles)] = + static_cast(scale_pair); + colwise_scales[swizzle::gemm_swizzled_scale_idx(kPart * K + col + 1, row0 / 32, + row_block_tiles)] = + static_cast(scale_pair >> 8); + } else { + *(unsigned short*)(colwise_scales + (size_t)(row0 / 32) * colwise_scale_stride + + kPart * K + col) = (unsigned short)scale_pair; + } + } + const ptx::bf16x2 scale_rcp = + ptx::exp2f_rcp_2x_per_lane(__byte_perm(scale_pair, 0, 0x4140)); + col_scale_rcp_smem[tid] = reinterpret_cast(scale_rcp); + } + // Drain the row scales: one thread per row, one byte per group inside the tensor. + if (tid >= kTileCols / 2 && tid < kTileCols / 2 + kRowsPerMxBlock) { + const int r = tid - kTileCols / 2; + if (!kEdge || row0 + r < M) { + const int groups = full_width ? kMxGroupsPerRow : (K - col0 + 31) / 32; + const size_t col_tiles = DIVUP(static_cast(out_cols), size_t{128}); + for (int g = 0; g < groups; ++g) { + const int group = kPart * k_groups + (col0 >> 5) + g; + size_t idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + idx = swizzle::gemm_swizzled_scale_idx(row0 + r, group, col_tiles); + } else { + idx = (size_t)(row0 + r) * rowwise_scale_stride + group; + } + rowwise_scales[idx] = row_scale_bytes[r * kMxGroupsPerRow + g]; + } + } + } + __syncthreads(); + + // ---- colwise quantization --------------------------------------------------------------- + { + const uint4 rcp_vec = *(const uint4*)(col_scale_rcp_smem + kWordsPerLane * lane); + const unsigned* rcp = (const unsigned*)&rcp_vec; + unsigned char* col_out_ptr = colwise_out + store_offset; +#pragma unroll + for (int j = 0; j < kRowsPerWarp; ++j) { + const unsigned* v = tile_words + j * kWordsPerLane; + unsigned col_out[kWordsPerLane / 2]; +#pragma unroll + for (int m = 0; m < kWordsPerLane / 2; ++m) + col_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], rcp[2 * m], rcp[2 * m + 1]); + if (in_bounds(j)) store_quantized(col_out_ptr + (size_t)j * out_cols, col_out); + } + } + }; + const auto parts = [&](auto edge) { + band(edge, std::integral_constant{}); + if constexpr (IS_BWD) band(edge, std::integral_constant{}); + }; + if constexpr (kTails) { + if (full_width && row0 + kRowsPerMxBlock <= M) { + parts(std::false_type{}); + } else { + parts(std::true_type{}); + } + } else { + parts(std::false_type{}); + } + } +#endif +} + +//! Whether the gated regtile kernel takes this request: BF16 [M, 2K] -> E4M3, both directions, +//! K % 32 == 0 (the generic gated kernel's requirement), aligned buffers. +inline bool can_use(const Tensor& input, const Tensor& grad, const Tensor& output, + const bool is_bwd) { + if (!getenv("NVTE_MXFP8_REGTILE", true)) return false; + if (transformer_engine::cuda::sm_arch() < 100) return false; + if (input.dtype() != DType::kBFloat16 || output.dtype() != DType::kFloat8E4M3) return false; + if (!output.has_data() || !output.has_columnwise_data()) return false; + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim() / 2; + constexpr size_t kMaxRows = std::numeric_limits::max() - 64; + if (rows == 0 || cols == 0 || cols % 32 != 0 || rows > kMaxRows || + 2 * cols > std::numeric_limits::max() / kRowsPerMxBlock) { + return false; + } + return is_aligned_ptr(input.data.dptr, 16) && (!is_bwd || is_aligned_ptr(grad.data.dptr, 16)) && + is_aligned_ptr(output.data.dptr, 8) && is_aligned_ptr(output.columnwise_data.dptr, 8) && + is_aligned_ptr(output.columnwise_scale_inv.dptr, 2); +} + +template +void quantize(const Tensor& input, const Tensor& grad, Tensor* output, cudaStream_t stream) { + const int M = static_cast(input.flat_first_dim()); + const int K = static_cast(input.flat_last_dim() / 2); + const int bands = (M + kRowsPerMxBlock - 1) / kRowsPerMxBlock; + const int grid_cols = (K + kTileCols - 1) / kTileCols; + int min_walk = kMinWalk; + if ((long long)grid_cols * (bands / min_walk) < 2LL * cuda::sm_count()) min_walk = 1; + int iters = pick_walk_length(bands, grid_cols, kCtaTargetActivation, min_walk); + if (iters > kMaxWalkDact) iters = kMaxWalkDact; + const dim3 grid(grid_cols, bands / iters); + const auto run = [&](auto tails) { + constexpr bool kTails = decltype(tails)::value; + quantize_gated_kernel + <<>>( + (const uint4*)input.data.dptr, (const uint4*)(IS_BWD ? grad.data.dptr : nullptr), + (unsigned char*)output->data.dptr, (unsigned char*)output->scale_inv.dptr, + (unsigned char*)output->columnwise_data.dptr, + (unsigned char*)output->columnwise_scale_inv.dptr, M, K, + static_cast(output->scale_inv.shape[1]), + static_cast(output->columnwise_scale_inv.shape[1]), iters); + }; + if (M % kRowsPerMxBlock == 0 && K % kTileCols == 0) { + run(std::false_type{}); + } else { + run(std::true_type{}); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace gated +} // namespace regtile +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_GATED_MXFP8_REGTILE_CUH_ From 5d860584a1704d6588c3f5d262cbdec31034e153 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 16 Oct 2026 14:05:01 -0700 Subject: [PATCH 13/29] [Common] Take FP32 inputs in the register-tiled MXFP8 activation kernel FP32 inputs of GeLU, SiLU, dGeLU and dSiLU, with or without dbias, now take the register-tiled kernel. A tile column pair is two FP32 registers instead of one 32-bit word, so an FP32 CTA has twice the threads (256 for the activations, 512 for the derivatives) and each thread holds half the rows: per-thread registers and threads per SM stay as for 16-bit inputs. A load group is one row, two 16-byte vectors per lane. FP32 inputs have no table; the activation is computed with util/math.h, and every value is scaled in FP32 and converted with mul_cvt_2x, as the generic kernel does, so the output stays byte-identical. GB200 against TE main, normal data: GeLU 1.12-1.35x, SiLU 1.24-1.35x, dGeLU 1.29-1.48x, dSiLU 1.48-1.92x, dbias_dGeLU 1.16-1.32x, dbias_dSiLU 1.35-1.61x (wide data up to 2.37x); forward tensors under 2^23 elements keep the generic kernel as for BF16. A new FP32 regtile-vs-generic suite covers every mode, output type, amax setting and tail shape. Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 39 ++- .../common/cast/mxfp8/quantize_mxfp8.cuh | 5 +- .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 240 ++++++++++++------ 3 files changed, 197 insertions(+), 87 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index 232ec2c56f2..ba0696a5384 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -392,6 +392,8 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, Tensor grad("grad", shape, itype); if (itype == DType::kFloat16) { fill_inputs(fill, &x, &grad); + } else if (itype == DType::kFloat32) { + fill_inputs(fill, &x, &grad); } else { fill_inputs(fill, &x, &grad); } @@ -473,9 +475,14 @@ void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, const auto dbias_bytes = [&](Tensor &t) { return itype == DType::kFloat16 ? reinterpret_cast(t.rowwise_cpu_dptr()) + : itype == DType::kFloat32 + ? reinterpret_cast(t.rowwise_cpu_dptr()) : reinterpret_cast(t.rowwise_cpu_dptr()); }; - EXPECT_EQ(count_byte_mismatches(dbias_bytes(dbias[0]), dbias_bytes(dbias[1]), cols * 2), 0u) + const size_t dbias_elt_bytes = itype == DType::kFloat32 ? 4 : 2; + EXPECT_EQ(count_byte_mismatches(dbias_bytes(dbias[0]), dbias_bytes(dbias[1]), + cols * dbias_elt_bytes), + 0u) << "dbias"; } // The FP8 data as raw bytes, whichever FP8 type the output holds. @@ -659,3 +666,33 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), ::testing::Bool()), regtile_case_name); + +// FP32 inputs: no table, twice the threads per CTA (see regtile::TileConfig). +class CastMXFP8ActRegtileFP32TestSuite : public ::testing::TestWithParam {}; + +TEST_P(CastMXFP8ActRegtileFP32TestSuite, MatchesGenericKernel) { + cudaDeviceProp prop; + ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); + if (prop.major < 10) { + GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; + } + const auto [mode, fill, swizzled, dims, use_table, otype, with_amax] = GetParam(); + run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled, use_table, otype, + with_amax, DType::kFloat32); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, CastMXFP8ActRegtileFP32TestSuite, + ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::DBiasDGeLU, Mode::SiLU, + Mode::DSiLU, Mode::DBiasDSiLU), + ::testing::Values(Fill::Typical, Fill::OrderedInputs, Fill::Specials), + ::testing::Bool(), + ::testing::Values(std::pair{1024, 2048}, + std::pair{993, 7200}, + std::pair{33, 272}, + // Large enough for the computed forward path. + std::pair{4127, 2064}), + ::testing::Bool(), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), + ::testing::Bool()), + regtile_case_name); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index c6f145c380e..c6f072f1259 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -966,7 +966,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, } // regtile::quantize replaces the generic kernel below for the requests it - // implements: bidimensional BF16 or FP16 -> E4M3 or E5M2 with fused GeLU, SiLU, dGeLU or + // implements: bidimensional BF16, FP16 or FP32 -> E4M3 or E5M2 with fused GeLU, SiLU, dGeLU or // dSiLU, read from an activation table or computed, the derivatives optionally with // dbias, with or without amax, cols % 16 == 0, 16-byte aligned inputs (see // regtile::kImplements and regtile::can_use). It keeps the tile in registers instead @@ -974,7 +974,8 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, // the same dbias workspace the generic kernel writes. if constexpr ((std::is_same_v || (std::is_same_v && - regtile::kTakesFP16)) && + regtile::kTakesFP16) || + std::is_same_v) && (std::is_same_v || std::is_same_v) && regtile::kImplements) { if (regtile::can_use( diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh index a1d90cacff1..24de19adebf 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -54,8 +54,8 @@ namespace ptx = transformer_engine::ptx; namespace { -// The datapath works on 32-bit words, each holding a column pair of the input type (BF16 -// or FP16; see Datapath). +// The datapath works on column pairs of the input type (Datapath::Word): one 32-bit word for +// BF16 and FP16, two FP32 registers for FP32. /*! \brief Elementwise max(|a|, |b|) on a packed BF16 pair, as a raw word. */ __device__ __forceinline__ unsigned abs_max_bf16x2(unsigned a, unsigned b) { @@ -110,8 +110,8 @@ template struct Datapath; // below /*! \brief The derivative of an input pair, from the table in shared memory or computed. */ -template -__device__ __forceinline__ ptx::floatx2 dact_word(unsigned x, +template +__device__ __forceinline__ ptx::floatx2 dact_word(const Word& x, const unsigned char* __restrict__ table) { if constexpr (TableT::kComputed) { const ptx::floatx2 v = Datapath::widen(x); @@ -154,9 +154,10 @@ constexpr int kRowsPerMxBlock = 32; constexpr int kMxGroupsPerRow = kTileCols / 32; //! A register-resident tile spreads its 256 columns over the 32 lanes of a warp. constexpr int kColsPerLane = kTileCols / 32; -//! ...which a lane holds as input-type pairs, one 32-bit word each. +//! ...which a lane holds as input-type pairs (Datapath::Word). constexpr int kWordsPerLane = kColsPerLane / 2; -//! ...read as 128-bit vectors, so one per row. +//! ...read, for 16-bit inputs, as one 128-bit vector per row (FP32 inputs take two; see +//! kVecsPerLaneRow in quantize_regtile). constexpr int kVecLoadsPerRow = kWordsPerLane / 4; //! Lanes that have to cooperate to cover one 32-column MX group. constexpr int kLanesPerMxGroup = 32 / kColsPerLane; @@ -187,8 +188,19 @@ __device__ __forceinline__ void store_quantized(unsigned char* p, const unsigned template struct Datapath; +/*! \brief Helpers on a packed pair of 16-bit values (one 32-bit word). */ +struct PackedPair16 { + using Word = unsigned; + //! The other value of the pair in each half. + __device__ __forceinline__ static Word swap_halves(Word v) { return __byte_perm(v, v, 0x1032); } + __device__ __forceinline__ static Word shfl_xor(Word v, int lane_mask) { + return __shfl_xor_sync(0xffffffffu, v, lane_mask); + } + __device__ __forceinline__ static Word clear_sign(Word v) { return v & 0x7fff7fffu; } +}; + template <> -struct Datapath { +struct Datapath : PackedPair16 { //! A BF16 pair. 2^(127-e) is a power of two inside BF16's range, so the BF16 multiply in //! mul_cvt_4x is exact, like the generic kernel's FP32 one. using Scale = unsigned; @@ -214,7 +226,7 @@ struct Datapath { }; template <> -struct Datapath { +struct Datapath : PackedPair16 { //! An FP32 pair: 2^(127-e) is often outside FP16's range, so the multiply is done in FP32, //! as the generic kernel does it (mul_cvt_2x with an FP32 scale pair). using Scale = ptx::floatx2; @@ -251,6 +263,48 @@ struct Datapath { } }; +template <> +struct Datapath { + //! A pair of FP32 values, two registers: the tile is twice as large as for 16-bit inputs, + //! so the CTA has twice the threads, each holding half the rows (see TileConfig). + using Word = ptx::floatx2; + using Scale = ptx::floatx2; + //! max(|a|, |b|) per lane. fmaxf ignores NaN, as the generic kernel's 0-seeded fmaxf does; + //! a block of only NaNs still gives NaN here, which amax_to_e8m0_2x maps to the scale of 0. + __device__ __forceinline__ static Word abs_max(const Word& a, const Word& b) { + return {fmaxf(fabsf(a.x), fabsf(b.x)), fmaxf(fabsf(a.y), fabsf(b.y))}; + } + __device__ __forceinline__ static Word swap_halves(const Word& v) { return {v.y, v.x}; } + __device__ __forceinline__ static Word shfl_xor(const Word& v, int lane_mask) { + return {__shfl_xor_sync(0xffffffffu, v.x, lane_mask), + __shfl_xor_sync(0xffffffffu, v.y, lane_mask)}; + } + __device__ __forceinline__ static Word clear_sign(const Word& v) { + return {fabsf(v.x), fabsf(v.y)}; + } + __device__ __forceinline__ static ptx::floatx2 widen(const Word& v) { return v; } + //! The generic kernel's rounding to the input type is a no-op for FP32. + __device__ __forceinline__ static Word narrow(const ptx::floatx2& v) { return v; } + __device__ __forceinline__ static Scale row_rcp(unsigned e) { + const float r = ptx::exp2f_rcp(static_cast(e)); + return {r, r}; + } + __device__ __forceinline__ static Scale col_rcp(unsigned scale_pair) { + return {ptx::exp2f_rcp(static_cast(scale_pair & 0xffu)), + ptx::exp2f_rcp(static_cast((scale_pair >> 8) & 0xffu))}; + } + //! FP32 multiply, then convert, as the generic kernel does (mul_cvt_2x on an FP32 pair). + template + __device__ __forceinline__ static unsigned quantize(const Word& in01, const Word& in23, Scale s01, + Scale s23) { + ptx::FPx2 lo, hi; + ptx::mul_cvt_2x(lo, in01, s01); + ptx::mul_cvt_2x(hi, in23, s23); + return static_cast(reinterpret_cast(lo)) | + (static_cast(reinterpret_cast(hi)) << 16); + } +}; + // Minimum resident CTAs per SM, the second __launch_bounds__ argument, capped // at what the SM can hold: sm_75 and sm_107 have half the threads per SM of sm_100. constexpr int kMinBlocksAct = 8; @@ -264,10 +318,13 @@ constexpr int kMaxThreadsPerSm = 2048; // Minimum resident CTAs per SM with dbias: the 32 KB staging buffer allows about four. constexpr int kMinBlocksDbias = 4; -template +template struct TileConfig { static_assert(IS_ACT != IS_DACT, "Exactly one of the activation and its derivative is fused."); - static constexpr int kThreadsPerCta = IS_DACT ? 256 : 128; + //! FP32 inputs take twice the registers per value, so twice the threads share a tile: each + //! thread holds the same number of registers of it, and an SM the same number of threads. + static constexpr int kThreadScale = sizeof(IType) == 4 ? 2 : 1; + static constexpr int kThreadsPerCta = (IS_DACT ? 256 : 128) * kThreadScale; static constexpr int kWarpsPerCta = kThreadsPerCta / 32; static constexpr int kRowsPerWarp = 32 / kWarpsPerCta; // The rowwise half is quantized before the cross-warp barrier and the row @@ -276,16 +333,17 @@ struct TileConfig { // the threads that just folded the colwise partials. static constexpr int kDrainTidBase = (kThreadsPerCta >= kTileCols / 2 + kTileCols / 4) ? kTileCols / 2 : 0; - static constexpr int kRowsInFlight = 2; + //! Rows loaded per group: one with FP32, whose row is two 16-byte vectors per lane. + static constexpr int kRowsInFlight = sizeof(IType) == 4 ? 1 : 2; static constexpr int kRequestedMinBlocks = - kDbias ? kMinBlocksDbias : (IS_DACT ? kMinBlocksDact : kMinBlocksAct); + (kDbias ? kMinBlocksDbias : (IS_DACT ? kMinBlocksDact : kMinBlocksAct)) / kThreadScale; static constexpr int kMinBlocksPerSm = kRequestedMinBlocks < kMaxThreadsPerSm / kThreadsPerCta ? kRequestedMinBlocks : kMaxThreadsPerSm / kThreadsPerCta; static constexpr int kTableBytes = static_cast(TableT::kBytes); //! Dynamic shared memory: with dbias, the band's FP32 products before rounding. static constexpr int kDbiasStageBytes = kDbias ? kRowsPerMxBlock * kTileCols * 4 : 0; - static_assert(!kDbias || (IS_DACT && kThreadsPerCta == kTileCols), + static_assert(!kDbias || (IS_DACT && kThreadsPerCta >= kTileCols), "dbias: one thread per column sums the column."); }; @@ -309,13 +367,18 @@ __device__ __forceinline__ void quantize_regtile( unsigned char* __restrict__ colwise_out, unsigned char* __restrict__ colwise_scales, int M, int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* table, float* __restrict__ dbias_workspace, float* __restrict__ amax_ptr) { - using C = TileConfig; + using C = TileConfig; using DP = Datapath; using Scale = typename DP::Scale; + using Word = typename DP::Word; constexpr int kWarpsPerCta = C::kWarpsPerCta, kRowsPerWarp = C::kRowsPerWarp, kRowsInFlight = C::kRowsInFlight; + //! Input values per 16-byte vector, words (column pairs) per vector, vectors per lane row. + constexpr int kValsPerVec = 16 / sizeof(IType); + constexpr int kWordsPerVec = 16 / sizeof(Word); + constexpr int kVecsPerLaneRow = kWordsPerLane / kWordsPerVec; - __shared__ __align__(16) unsigned col_amax_partials[kWarpsPerCta][kTileCols / 2]; + __shared__ __align__(16) Word col_amax_partials[kWarpsPerCta][kTileCols / 2]; __shared__ __align__(16) Scale col_scale_rcp_smem[kTileCols / 2]; __shared__ __align__(8) unsigned char row_scale_bytes[kRowsPerMxBlock * kMxGroupsPerRow]; __shared__ __align__(16) unsigned char table_smem[C::kTableBytes]; @@ -327,7 +390,7 @@ __device__ __forceinline__ void quantize_regtile( const int lane = tid & 31; const int warp = tid >> 5; const int col0 = blockIdx.x * kTileCols; - const int vecs_per_row = K >> 3; // uint4 (8 BF16 or FP16 values) per row + const int vecs_per_row = K / kValsPerVec; // Tails: the last column tile may end inside the tile, and row bands may end past M. Out-of-range // elements load as 0, which every activation maps to 0, so they never raise an amax. The // generic kernel writes scales for its whole 64x64 tiles, and so does this one: rowwise scales @@ -363,21 +426,21 @@ __device__ __forceinline__ void quantize_regtile( // Whether this lane's elements of row warp_row0 + j are inside the tensor. const auto in_bounds = [&](int j) { return !kEdge || (lane_in_cols && warp_row0 + j < M); }; // This thread's kRowsPerWarp x kColsPerLane sub-tile, as input-type pairs. - unsigned tile_words[kRowsPerWarp * kWordsPerLane]; - const size_t load_offset = (size_t)(row0 + warp * kRowsPerWarp) * vecs_per_row + (col0 >> 3) + - (size_t)kVecLoadsPerRow * lane; + Word tile_words[kRowsPerWarp * kWordsPerLane]; + const size_t load_offset = (size_t)(row0 + warp * kRowsPerWarp) * vecs_per_row + + col0 / kValsPerVec + (size_t)kVecsPerLaneRow * lane; const uint4* input_ptr = input + load_offset; const uint4* act_input_ptr = act_input + load_offset; // ---- load + activate, kRowsInFlight rows in flight at a time -------------------- #pragma unroll for (int h = 0; h < kRowsPerWarp / kRowsInFlight; ++h) { - uint4 x_vec[kRowsInFlight][kVecLoadsPerRow], grad_vec[kRowsInFlight][kVecLoadsPerRow]; + uint4 x_vec[kRowsInFlight][kVecsPerLaneRow], grad_vec[kRowsInFlight][kVecsPerLaneRow]; asm volatile("" ::: "memory"); #pragma unroll for (int t = 0; t < kRowsInFlight; ++t) #pragma unroll - for (int n = 0; n < kVecLoadsPerRow; ++n) + for (int n = 0; n < kVecsPerLaneRow; ++n) x_vec[t][n] = in_bounds(h * kRowsInFlight + t) ? act_input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n] : make_uint4(0, 0, 0, 0); @@ -385,7 +448,7 @@ __device__ __forceinline__ void quantize_regtile( #pragma unroll for (int t = 0; t < kRowsInFlight; ++t) #pragma unroll - for (int n = 0; n < kVecLoadsPerRow; ++n) + for (int n = 0; n < kVecsPerLaneRow; ++n) grad_vec[t][n] = in_bounds(h * kRowsInFlight + t) ? input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n] : make_uint4(0, 0, 0, 0); @@ -403,11 +466,11 @@ __device__ __forceinline__ void quantize_regtile( #pragma unroll for (int t = 0; t < kRowsInFlight; ++t) { #pragma unroll - for (int n = 0; n < kVecLoadsPerRow; ++n) { - const unsigned* x_words = (const unsigned*)&x_vec[t][n]; - const int tile_word_base = (h * kRowsInFlight + t) * kWordsPerLane + n * 4; + for (int n = 0; n < kVecsPerLaneRow; ++n) { + const Word* x_words = (const Word*)&x_vec[t][n]; + const int tile_word_base = (h * kRowsInFlight + t) * kWordsPerLane + n * kWordsPerVec; #pragma unroll - for (int m = 0; m < 4; ++m) { + for (int m = 0; m < kWordsPerVec; ++m) { if constexpr (TableT::kComputed) { // util/math.h, then rounded to BF16 as the generic kernel rounds it. const ptx::floatx2 v = DP::widen(x_words[m]); @@ -430,8 +493,8 @@ __device__ __forceinline__ void quantize_regtile( } #define MXFP8_DACT_ROW(T) \ { \ - const unsigned* x_words = (const unsigned*)&x_vec[T][0]; \ - const unsigned* grad_words = (const unsigned*)&grad_vec[T][0]; \ + const Word* x_words = (const Word*)&x_vec[T][0]; \ + const Word* grad_words = (const Word*)&grad_vec[T][0]; \ float products[2 * kWordsPerLane]; \ MXFP8_DACT_WORD(T, 0) \ MXFP8_DACT_WORD(T, 1) \ @@ -445,9 +508,9 @@ __device__ __forceinline__ void quantize_regtile( stage[1] = make_float4(products[4], products[5], products[6], products[7]); \ } \ } - static_assert(kRowsInFlight == 2, "A load group is two rows."); + static_assert(kRowsInFlight == 1 || kRowsInFlight == 2, "A load group is 1 or 2 rows."); MXFP8_DACT_ROW(0) - MXFP8_DACT_ROW(1) + if constexpr (kRowsInFlight == 2) MXFP8_DACT_ROW(1) #undef MXFP8_DACT_ROW #undef MXFP8_DACT_WORD } @@ -457,53 +520,27 @@ __device__ __forceinline__ void quantize_regtile( // the cross-warp barrier (see TileConfig). const size_t store_offset = (size_t)(row0 + warp * kRowsPerWarp) * K + col0 + kColsPerLane * lane; - auto row_scale = [&](const unsigned* row, int j) { - unsigned amax = DP::abs_max(DP::abs_max(row[0], row[1]), DP::abs_max(row[2], row[3])); + auto row_scale = [&](const Word* row, int j) { + Word amax = DP::abs_max(DP::abs_max(row[0], row[1]), DP::abs_max(row[2], row[3])); #pragma unroll for (int m = 4; m < kWordsPerLane; m += 2) amax = DP::abs_max(amax, DP::abs_max(row[m], row[m + 1])); // A 32-value rowwise block spans kLanesPerMxGroup lanes. #pragma unroll for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) - amax = DP::abs_max(amax, __shfl_xor_sync(0xffffffffu, amax, lane_mask)); - // Fold the two halves with one PRMT and a packed max, which leaves the - // magnitude in both halves; the high half is then the FP32 bit pattern. - amax = DP::abs_max(amax, __byte_perm(amax, amax, 0x1032)); - const unsigned scale_byte = amax_to_e8m0_2x(0.f, DP::widen(amax & 0x7fff7fffu).y); + amax = DP::abs_max(amax, DP::shfl_xor(amax, lane_mask)); + // Fold the two halves of the pair (for 16-bit pairs one PRMT and a packed max), which + // leaves the magnitude in both halves. + amax = DP::abs_max(amax, DP::swap_halves(amax)); + const unsigned scale_byte = amax_to_e8m0_2x(0.f, DP::widen(DP::clear_sign(amax)).y); if ((lane & (kLanesPerMxGroup - 1)) == 0) row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + (lane / kLanesPerMxGroup)] = (unsigned char)scale_byte; return DP::row_rcp(scale_byte); }; - // On the 128-thread activation tile, two rows share one butterfly and conversion. - auto row_amax = [&](const unsigned* row) { - unsigned amax = DP::abs_max(DP::abs_max(row[0], row[1]), DP::abs_max(row[2], row[3])); -#pragma unroll - for (int m = 4; m < kWordsPerLane; m += 2) - amax = DP::abs_max(amax, DP::abs_max(row[m], row[m + 1])); - return DP::abs_max(amax, __byte_perm(amax, amax, 0x1032)); - }; - auto row_scale_pair = [&](const unsigned* row_a, const unsigned* row_b, int j, - Scale& row_scale_rcp0, Scale& row_scale_rcp1) { - unsigned amax_pair = __byte_perm(row_amax(row_a), row_amax(row_b), 0x5410) & 0x7fff7fffu; -#pragma unroll - for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) - amax_pair = DP::abs_max(amax_pair, __shfl_xor_sync(0xffffffffu, amax_pair, lane_mask)); - const ptx::floatx2 amax_2x = DP::widen(amax_pair); - const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); - const unsigned scale_byte0 = scale_pair & 0xffu, scale_byte1 = (scale_pair >> 8) & 0xffu; - if ((lane & (kLanesPerMxGroup - 1)) == 0) { - unsigned char* scale_dst = row_scale_bytes + (warp * kRowsPerWarp + j) * kMxGroupsPerRow + - (lane / kLanesPerMxGroup); - scale_dst[0] = (unsigned char)scale_byte0; - scale_dst[kMxGroupsPerRow] = (unsigned char)scale_byte1; - } - row_scale_rcp0 = DP::row_rcp(scale_byte0); - row_scale_rcp1 = DP::row_rcp(scale_byte1); - }; // Colwise amax: per-warp partial in registers, cross-warp fold in shared. auto publish_col_amax_partials = [&]() { - unsigned col_amax[kWordsPerLane]; + Word col_amax[kWordsPerLane]; #pragma unroll for (int m = 0; m < kWordsPerLane; ++m) col_amax[m] = tile_words[m]; #pragma unroll @@ -512,16 +549,44 @@ __device__ __forceinline__ void quantize_regtile( for (int m = 0; m < kWordsPerLane; ++m) col_amax[m] = DP::abs_max(col_amax[m], tile_words[j * kWordsPerLane + m]); #pragma unroll - for (int n = 0; n < kVecLoadsPerRow; ++n) - *(uint4*)(&col_amax_partials[warp][kWordsPerLane * lane + 4 * n]) = - *(const uint4*)(col_amax + 4 * n); + for (int n = 0; n < kVecsPerLaneRow; ++n) + reinterpret_cast(&col_amax_partials[warp][kWordsPerLane * lane])[n] = + reinterpret_cast(col_amax)[n]; }; if constexpr (C::kThreadsPerCta == 128) { + static_assert(sizeof(Word) == 4, "Paired rows pack two 16-bit amaxes into one word."); + // On the 128-thread activation tile, two rows share one butterfly and conversion. + auto row_amax = [&](const unsigned* row) { + unsigned amax = DP::abs_max(DP::abs_max(row[0], row[1]), DP::abs_max(row[2], row[3])); +#pragma unroll + for (int m = 4; m < kWordsPerLane; m += 2) + amax = DP::abs_max(amax, DP::abs_max(row[m], row[m + 1])); + return DP::abs_max(amax, __byte_perm(amax, amax, 0x1032)); + }; + auto row_scale_pair = [&](const unsigned* row_a, const unsigned* row_b, int j, + Scale& row_scale_rcp0, Scale& row_scale_rcp1) { + unsigned amax_pair = __byte_perm(row_amax(row_a), row_amax(row_b), 0x5410) & 0x7fff7fffu; +#pragma unroll + for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) + amax_pair = DP::abs_max(amax_pair, __shfl_xor_sync(0xffffffffu, amax_pair, lane_mask)); + const ptx::floatx2 amax_2x = DP::widen(amax_pair); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + const unsigned scale_byte0 = scale_pair & 0xffu, scale_byte1 = (scale_pair >> 8) & 0xffu; + if ((lane & (kLanesPerMxGroup - 1)) == 0) { + unsigned char* scale_dst = row_scale_bytes + + (warp * kRowsPerWarp + j) * kMxGroupsPerRow + + (lane / kLanesPerMxGroup); + scale_dst[0] = (unsigned char)scale_byte0; + scale_dst[kMxGroupsPerRow] = (unsigned char)scale_byte1; + } + row_scale_rcp0 = DP::row_rcp(scale_byte0); + row_scale_rcp1 = DP::row_rcp(scale_byte1); + }; unsigned char* row_out_ptr = rowwise_out + store_offset; #pragma unroll for (int j = 0; j < kRowsPerWarp; j += 2) { - const unsigned* row_a = tile_words + j * kWordsPerLane; - const unsigned* row_b = row_a + kWordsPerLane; + const Word* row_a = tile_words + j * kWordsPerLane; + const Word* row_b = row_a + kWordsPerLane; Scale row_scale_rcp0, row_scale_rcp1; row_scale_pair(row_a, row_b, j, row_scale_rcp0, row_scale_rcp1); unsigned row_out[kWordsPerLane / 2]; @@ -541,7 +606,7 @@ __device__ __forceinline__ void quantize_regtile( unsigned char* row_out_ptr = rowwise_out + store_offset; #pragma unroll for (int j = 0; j < kRowsPerWarp; ++j) { - const unsigned* v = tile_words + j * kWordsPerLane; + const Word* v = tile_words + j * kWordsPerLane; const Scale row_scale_rcp = row_scale(v, j); unsigned row_out[kWordsPerLane / 2]; #pragma unroll @@ -556,19 +621,20 @@ __device__ __forceinline__ void quantize_regtile( __syncthreads(); if constexpr (kDbias) { // The generic kernel's order: a strictly sequential FP32 sum down the column. - float sum = dbias_partial; + if (tid < kTileCols) { + float sum = dbias_partial; #pragma unroll 8 - for (int r = 0; r < kRowsPerMxBlock; ++r) - sum = __fadd_rn(sum, dbias_stage[r * kTileCols + tid]); - dbias_partial = sum; + for (int r = 0; r < kRowsPerMxBlock; ++r) + sum = __fadd_rn(sum, dbias_stage[r * kTileCols + tid]); + dbias_partial = sum; + } } if (tid < kTileCols / 2) { - unsigned col_amax = col_amax_partials[0][tid]; + Word col_amax = col_amax_partials[0][tid]; #pragma unroll for (int w = 1; w < kWarpsPerCta; ++w) col_amax = DP::abs_max(col_amax, col_amax_partials[w][tid]); - col_amax &= 0x7fff7fffu; - const ptx::floatx2 amax_2x = DP::widen(col_amax); + const ptx::floatx2 amax_2x = DP::widen(DP::clear_sign(col_amax)); // The tensor amax, when asked for: the max of these column amaxes. Like the generic // kernel's per-element fmaxf, it ignores NaN (a column of only NaNs gives NaN here, // which fmaxf drops) and is unaffected by the zeros loaded past M or K. @@ -646,7 +712,7 @@ __device__ __forceinline__ void quantize_regtile( unsigned char* col_out_ptr = colwise_out + store_offset; #pragma unroll for (int j = 0; j < kRowsPerWarp; ++j) { - const unsigned* v = tile_words + j * kWordsPerLane; + const Word* v = tile_words + j * kWordsPerLane; unsigned col_out[kWordsPerLane / 2]; #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) @@ -666,7 +732,7 @@ __device__ __forceinline__ void quantize_regtile( band(std::false_type{}); } if constexpr (kDbias) { - if ((it & 1) == 1 && col0 + tid < K) { + if ((it & 1) == 1 && tid < kTileCols && col0 + tid < K) { dbias_workspace[(size_t)(row0 / (2 * kRowsPerMxBlock)) * K + col0 + tid] = dbias_partial; } } @@ -680,8 +746,9 @@ __device__ __forceinline__ void quantize_regtile( template -__global__ void __launch_bounds__(TileConfig::kThreadsPerCta, - TileConfig::kMinBlocksPerSm) +__global__ void __launch_bounds__( + TileConfig::kThreadsPerCta, + TileConfig::kMinBlocksPerSm) quantize_mxfp8_kernel(const unsigned* __restrict__ input, const unsigned* __restrict__ act_input, unsigned char* __restrict__ rowwise_out, @@ -725,7 +792,8 @@ template static void set_carveout() { constexpr int kPercent = IS_ACT ? kCarveoutPercentAct : 0; - constexpr int kDynamicBytes = TileConfig::kDbiasStageBytes; + constexpr int kDynamicBytes = + TileConfig::kDbiasStageBytes; if constexpr (kPercent > 0 || kDynamicBytes > 0) { // Function attributes are per device, like the tables. static std::mutex mutex; @@ -785,7 +853,7 @@ static void launch(const void* input, const void* act_input, void* rowwise_out, void* rowwise_scales, void* colwise_out, void* colwise_scales, int M, int K, int rowwise_scale_stride, int colwise_scale_stride, const void* table, float* dbias_workspace, float* amax, cudaStream_t stream) { - using C = TileConfig; + using C = TileConfig; const int iters = walk_length(M, K); dim3 grid(col_tiles(K), row_bands(M) / iters); // kAmax: a separate instantiation, so that the kernels without amax keep their register @@ -848,7 +916,11 @@ template constexpr bool kTakesFP16 = IS_DBIAS || activation_table::kKindOf == activation_table::kDsilu; -// Whether this kernel takes a request of a type in kImplements, with BF16 or FP16 input and E4M3 or +// FP32 inputs take this kernel for every request type: their values are already FP32, so the +// FP32 scaling costs nothing extra. GB200 against TE main: GeLU 1.12-1.35x, SiLU 1.24-1.43x, +// dGeLU 1.29-1.48x, dSiLU 1.48-2.37x, dbias_dGeLU 1.16-1.32x, dbias_dSiLU 1.35-2.00x. +// +// Whether this kernel takes a request of a type in kImplements, with BF16, FP16 or FP32 input and E4M3 or // E5M2 output. Anything else takes the generic kernel, as does every request with // NVTE_MXFP8_REGTILE=0. A noop tensor is ignored, as the generic kernel ignores it for every // request that computes an activation. From f5de75a60a067316c2e9de511de0950304505a89 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 16 Oct 2026 13:47:07 -0700 Subject: [PATCH 14/29] [Common] Compute both dSwiGLU gradient halves in one pass in the register-tiled kernel The register-tiled dSwiGLU kernel quantized the gradient of the activated half and the gradient of the gate half in two passes over each 32-row band, the second re-reading the activation and the incoming gradient and recomputing the sigmoid. It now reads the activation, gate and incoming gradient once, computes both halves' tiles together, and quantizes them one after the other. The math per element is unchanged, so the output is byte-identical to the generic gated kernel as before. GB200, BF16 -> MXFP8 E4M3, against TE main: dSwiGLU 1.74-2.16x on normal data (was 1.39-1.59x) and 1.87-2.76x on wide data; 32768x7168 goes from 550 to 405 us. The test adds an input whose activated half holds every BF16 value. Signed-off-by: Teddy Do --- .../operator/test_cast_mxfp8_gated_regtile.cu | 24 +++++- .../regtile/quantize_gated_mxfp8_regtile.cuh | 78 ++++++++++--------- 2 files changed, 61 insertions(+), 41 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu b/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu index 104a0406010..99f48fcd6ff 100644 --- a/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu +++ b/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu @@ -6,7 +6,8 @@ // The register-tiled SwiGLU / dSwiGLU MXFP8 kernel must write exactly the bytes of the generic // gated kernel. Each case runs both (NVTE_MXFP8_REGTILE=0 keeps the generic kernel) into outputs -// pre-filled with a marker byte, so they must also write the same set of bytes. +// pre-filled with a marker byte, so they must also write the same set of bytes. The AllBF16 fill +// puts every BF16 value in the activated half. #include #include @@ -16,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -31,7 +33,7 @@ using namespace test; namespace { -enum class GatedFill { Normal, Wide, Specials }; +enum class GatedFill { Normal, Wide, Specials, AllBF16 }; class ScopedEnv { public: @@ -65,6 +67,12 @@ void fill(Tensor *t, const size_t rows, const size_t cols, const GatedFill kind, // Wide: every 32x32 block gets its own power-of-two scale. for (size_t c = 0; c < cols; ++c) { float v = normal(gen); + if (kind == GatedFill::AllBF16 && c < cols / 2) { + // The activated half (the first cols / 2 columns) walks through every BF16 bit pattern. + const uint16_t bits = static_cast(r * (cols / 2) + c); + std::memcpy(&p[r * cols + c], &bits, sizeof(bits)); + continue; + } if (kind == GatedFill::Wide) { std::mt19937 block_gen(seed * 7919 + (r / 32) * 104729 + c / 32); v *= std::ldexp(1.0f, exponent(block_gen)); @@ -88,7 +96,8 @@ void run_case(const bool bwd, const GatedFill kind, const size_t rows, const siz Tensor input("input", std::vector{rows, 2 * cols}, DType::kBFloat16); Tensor grad("grad", std::vector{rows, cols}, DType::kBFloat16); fill(&input, rows, 2 * cols, kind, 1); - fill(&grad, rows, cols, kind == GatedFill::Specials ? GatedFill::Normal : kind, 2); + fill(&grad, rows, cols, + kind == GatedFill::Specials || kind == GatedFill::AllBF16 ? GatedFill::Normal : kind, 2); const size_t out_cols = bwd ? 2 * cols : cols; const std::vector out_shape{rows, out_cols}; Tensor outs[2] = { @@ -152,7 +161,7 @@ TEST_P(GatedRegtileTestSuite, MatchesGenericKernel) { std::string case_name(const testing::TestParamInfo &info) { const auto [bwd, kind, swizzled, dims] = info.param; - const char *fills[] = {"Normal", "Wide", "Specials"}; + const char *fills[] = {"Normal", "Wide", "Specials", "AllBF16"}; return std::string(bwd ? "DSwiGLU" : "SwiGLU") + "X" + fills[static_cast(kind)] + "X" + (swizzled ? "Swizzled" : "Compact") + "X" + std::to_string(dims.first) + "X" + std::to_string(dims.second); @@ -173,4 +182,11 @@ INSTANTIATE_TEST_SUITE_P( std::pair{100, 2080})), case_name); +// Every BF16 value in the activated half: 64 rows x 1024 columns. +INSTANTIATE_TEST_SUITE_P(AllBF16, GatedRegtileTestSuite, + ::testing::Combine(::testing::Bool(), ::testing::Values(GatedFill::AllBF16), + ::testing::Bool(), + ::testing::Values(std::pair{64, 1024})), + case_name); + } // namespace diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh index a7943f40a5c..df568ea11cc 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh @@ -101,44 +101,12 @@ __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) #pragma unroll 1 for (int it = 0; it < iters; ++it) { const int row0 = (it * gridDim.y + blockIdx.y) * kRowsPerMxBlock; - const auto band = [&](auto edge, auto part) { + const int warp_row0 = row0 + warp * kRowsPerWarp; + // Quantize one output half (kPart; the forward has only part 0) of the band from its tile. + const auto quantize_part = [&](auto edge, auto part, const unsigned* tile_words) { constexpr bool kEdge = decltype(edge)::value; constexpr int kPart = decltype(part)::value; - const int warp_row0 = row0 + warp * kRowsPerWarp; const auto in_bounds = [&](int j) { return !kEdge || (lane_in_cols && warp_row0 + j < M); }; - unsigned tile_words[kRowsPerWarp * kWordsPerLane]; - const size_t in_offset = (size_t)warp_row0 * (2 * k_vecs) + (col0 >> 3) + lane; - const uint4* act_ptr = input + in_offset; - const uint4* gate_ptr = act_ptr + k_vecs; - const uint4* grad_ptr = grad + (size_t)warp_row0 * k_vecs + (col0 >> 3) + lane; - // The gate is read by the forward and by the gradient of the activated half. - constexpr bool kNeedsGate = !IS_BWD || kPart == 0; - -#pragma unroll - for (int h = 0; h < kRowsPerWarp / kRowsInFlight; ++h) { - uint4 act_vec[kRowsInFlight], gate_vec[kRowsInFlight], grad_vec[kRowsInFlight]; - asm volatile("" ::: "memory"); -#pragma unroll - for (int t = 0; t < kRowsInFlight; ++t) { - const int j = h * kRowsInFlight + t; - const uint4 zero = make_uint4(0, 0, 0, 0); - act_vec[t] = in_bounds(j) ? act_ptr[(size_t)j * (2 * k_vecs)] : zero; - gate_vec[t] = kNeedsGate && in_bounds(j) ? gate_ptr[(size_t)j * (2 * k_vecs)] : zero; - grad_vec[t] = IS_BWD && in_bounds(j) ? grad_ptr[(size_t)j * k_vecs] : zero; - } - asm volatile("" ::: "memory"); -#pragma unroll - for (int t = 0; t < kRowsInFlight; ++t) { - const unsigned* a = (const unsigned*)&act_vec[t]; - const unsigned* g = (const unsigned*)&gate_vec[t]; - const unsigned* d = (const unsigned*)&grad_vec[t]; -#pragma unroll - for (int m = 0; m < 4; ++m) { - tile_words[(h * kRowsInFlight + t) * kWordsPerLane + m] = - gated_word(a[m], g[m], d[m]); - } - } - } // ---- rowwise ---------------------------------------------------------------------------- const size_t store_offset = @@ -244,9 +212,45 @@ __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) } } }; + // Load the band once and compute its tiles: one for the forward, both gradient halves for the + // backward, which then quantizes them one after the other. const auto parts = [&](auto edge) { - band(edge, std::integral_constant{}); - if constexpr (IS_BWD) band(edge, std::integral_constant{}); + constexpr bool kEdge = decltype(edge)::value; + const auto in_bounds = [&](int j) { return !kEdge || (lane_in_cols && warp_row0 + j < M); }; + unsigned tile0[kRowsPerWarp * kWordsPerLane]; + unsigned tile1[IS_BWD ? kRowsPerWarp * kWordsPerLane : 1]; + const size_t in_offset = (size_t)warp_row0 * (2 * k_vecs) + (col0 >> 3) + lane; + const uint4* act_ptr = input + in_offset; + const uint4* gate_ptr = act_ptr + k_vecs; + const uint4* grad_ptr = grad + (size_t)warp_row0 * k_vecs + (col0 >> 3) + lane; +#pragma unroll + for (int h = 0; h < kRowsPerWarp / kRowsInFlight; ++h) { + uint4 act_vec[kRowsInFlight], gate_vec[kRowsInFlight], grad_vec[kRowsInFlight]; + asm volatile("" ::: "memory"); +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) { + const int j = h * kRowsInFlight + t; + const uint4 zero = make_uint4(0, 0, 0, 0); + act_vec[t] = in_bounds(j) ? act_ptr[(size_t)j * (2 * k_vecs)] : zero; + gate_vec[t] = in_bounds(j) ? gate_ptr[(size_t)j * (2 * k_vecs)] : zero; + grad_vec[t] = IS_BWD && in_bounds(j) ? grad_ptr[(size_t)j * k_vecs] : zero; + } + asm volatile("" ::: "memory"); +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) { + const unsigned* a = (const unsigned*)&act_vec[t]; + const unsigned* g = (const unsigned*)&gate_vec[t]; + const unsigned* d = (const unsigned*)&grad_vec[t]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int w = (h * kRowsInFlight + t) * kWordsPerLane + m; + tile0[w] = gated_word(a[m], g[m], d[m]); + if constexpr (IS_BWD) tile1[w] = gated_word(a[m], g[m], d[m]); + } + } + } + quantize_part(edge, std::integral_constant{}, tile0); + if constexpr (IS_BWD) quantize_part(edge, std::integral_constant{}, tile1); }; if constexpr (kTails) { if (full_width && row0 + kRowsPerMxBlock <= M) { From b688e6139d09762d69b13dd4d26200a52894baec Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 16 Oct 2026 15:30:44 -0700 Subject: [PATCH 15/29] [Common] Support E5M2 output in the register-tiled SwiGLU kernel The register-tiled SwiGLU / dSwiGLU kernel wrote E4M3 only; requests for E5M2 took the generic gated kernel. The kernel is now templated on the FP8 output type, using the output-type-aware scale and convert helpers of the non-gated register-tiled kernel, and the gated dispatch admits E5M2. The output stays byte-identical to the generic gated kernel. GB200, BF16 -> MXFP8 E5M2, against TE main: SwiGLU 1.40-2.00x on normal data and 1.40-2.05x on wide data, dSwiGLU 1.74-2.13x and 1.89-2.75x, the same as E4M3, which is unchanged. The regtile-vs-generic test runs every gated case with both output types. Signed-off-by: Teddy Do --- .../operator/test_cast_mxfp8_gated_regtile.cu | 45 +++++++++++-------- .../common/cast/mxfp8/gated_mxfp8.cuh | 10 +++-- .../regtile/quantize_gated_mxfp8_regtile.cuh | 21 +++++---- 3 files changed, 45 insertions(+), 31 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu b/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu index 99f48fcd6ff..4322808dd57 100644 --- a/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu +++ b/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu @@ -92,7 +92,7 @@ size_t count_byte_mismatches(const uint8_t *a, const uint8_t *b, const size_t n) } void run_case(const bool bwd, const GatedFill kind, const size_t rows, const size_t cols, - const bool swizzled) { + const bool swizzled, const DType otype) { Tensor input("input", std::vector{rows, 2 * cols}, DType::kBFloat16); Tensor grad("grad", std::vector{rows, cols}, DType::kBFloat16); fill(&input, rows, 2 * cols, kind, 1); @@ -101,8 +101,8 @@ void run_case(const bool bwd, const GatedFill kind, const size_t rows, const siz const size_t out_cols = bwd ? 2 * cols : cols; const std::vector out_shape{rows, out_cols}; Tensor outs[2] = { - Tensor("regtile", out_shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING), - Tensor("generic", out_shape, DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING)}; + Tensor("regtile", out_shape, otype, true, true, NVTE_MXFP8_1D_SCALING), + Tensor("generic", out_shape, otype, true, true, NVTE_MXFP8_1D_SCALING)}; for (int i = 0; i < 2; ++i) { outs[i].set_with_gemm_swizzled_scales(swizzled); for (const auto param : {kNVTERowwiseData, kNVTEColumnwiseData, kNVTERowwiseScaleInv, @@ -123,15 +123,20 @@ void run_case(const bool bwd, const GatedFill kind, const size_t rows, const siz ASSERT_EQ(cudaGetLastError(), cudaSuccess); for (auto &out : outs) out.to_cpu(); const size_t n = rows * out_cols; - EXPECT_EQ(count_byte_mismatches( - reinterpret_cast(outs[0].rowwise_cpu_dptr()), - reinterpret_cast(outs[1].rowwise_cpu_dptr()), n), - 0u) + // The FP8 data as raw bytes, whichever FP8 type the output holds. + const auto rowwise_bytes = [&](Tensor &t) { + return otype == DType::kFloat8E5M2 + ? reinterpret_cast(t.rowwise_cpu_dptr()) + : reinterpret_cast(t.rowwise_cpu_dptr()); + }; + const auto columnwise_bytes = [&](Tensor &t) { + return otype == DType::kFloat8E5M2 + ? reinterpret_cast(t.columnwise_cpu_dptr()) + : reinterpret_cast(t.columnwise_cpu_dptr()); + }; + EXPECT_EQ(count_byte_mismatches(rowwise_bytes(outs[0]), rowwise_bytes(outs[1]), n), 0u) << "rowwise data"; - EXPECT_EQ(count_byte_mismatches( - reinterpret_cast(outs[0].columnwise_cpu_dptr()), - reinterpret_cast(outs[1].columnwise_cpu_dptr()), n), - 0u) + EXPECT_EQ(count_byte_mismatches(columnwise_bytes(outs[0]), columnwise_bytes(outs[1]), n), 0u) << "colwise data"; EXPECT_EQ(count_byte_mismatches(outs[0].rowwise_cpu_scale_inv_ptr(), outs[1].rowwise_cpu_scale_inv_ptr(), @@ -145,7 +150,8 @@ void run_case(const bool bwd, const GatedFill kind, const size_t rows, const siz << "colwise scales"; } -using Params = std::tuple>; +using Params = std::tuple, + DType /*output*/>; class GatedRegtileTestSuite : public ::testing::TestWithParam {}; @@ -155,16 +161,16 @@ TEST_P(GatedRegtileTestSuite, MatchesGenericKernel) { if (prop.major < 10) { GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; } - const auto [bwd, kind, swizzled, dims] = GetParam(); - run_case(bwd, kind, dims.first, dims.second, swizzled); + const auto [bwd, kind, swizzled, dims, otype] = GetParam(); + run_case(bwd, kind, dims.first, dims.second, swizzled, otype); } std::string case_name(const testing::TestParamInfo &info) { - const auto [bwd, kind, swizzled, dims] = info.param; + const auto [bwd, kind, swizzled, dims, otype] = info.param; const char *fills[] = {"Normal", "Wide", "Specials", "AllBF16"}; return std::string(bwd ? "DSwiGLU" : "SwiGLU") + "X" + fills[static_cast(kind)] + "X" + (swizzled ? "Swizzled" : "Compact") + "X" + std::to_string(dims.first) + "X" + - std::to_string(dims.second); + std::to_string(dims.second) + (otype == DType::kFloat8E5M2 ? "XE5M2" : "XE4M3"); } INSTANTIATE_TEST_SUITE_P( @@ -179,14 +185,17 @@ INSTANTIATE_TEST_SUITE_P( std::pair{993, 7200}, std::pair{1, 32}, std::pair{33, 288}, - std::pair{100, 2080})), + std::pair{100, 2080}), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2)), case_name); // Every BF16 value in the activated half: 64 rows x 1024 columns. INSTANTIATE_TEST_SUITE_P(AllBF16, GatedRegtileTestSuite, ::testing::Combine(::testing::Bool(), ::testing::Values(GatedFill::AllBF16), ::testing::Bool(), - ::testing::Values(std::pair{64, 1024})), + ::testing::Values(std::pair{64, 1024}), + ::testing::Values(DType::kFloat8E4M3, + DType::kFloat8E5M2)), case_name); } // namespace diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 17527c730c2..a5f215210b5 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -852,14 +852,16 @@ void quantize_gated(const Tensor &gated_input, const Tensor &grad, Tensor *outpu } // regtile::gated replaces the kernel below for SwiGLU and its gradient from BF16 to - // E4M3 in both directions (see regtile::gated::can_use), writing the same bytes. - if constexpr (std::is_same_v && std::is_same_v && + // E4M3 or E5M2 in both directions (see regtile::gated::can_use), writing the same + // bytes. + if constexpr (std::is_same_v && + (std::is_same_v || std::is_same_v) && std::is_same_v) { if constexpr (ActOP == &silu && (!IS_BWD || DActOP == &dsilu)) { if (regtile::gated::can_use(gated_input, grad, *output, IS_BWD)) { - regtile::gated::quantize(gated_input, grad, - output, stream); + regtile::gated::quantize( + gated_input, grad, output, stream); return; } } diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh index df568ea11cc..dd7c7cf98d3 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh @@ -74,7 +74,7 @@ __device__ __forceinline__ unsigned gated_word(const unsigned act, const unsigne {gated_element(a.x, g.x, d.x), gated_element(a.y, g.y, d.y)}); } -template +template __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) quantize_gated_kernel(const uint4* __restrict__ input, const uint4* __restrict__ grad, unsigned char* __restrict__ rowwise_out, @@ -122,7 +122,7 @@ __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) amax = abs_max_bf16x2(amax, __shfl_xor_sync(0xffffffffu, amax, lane_mask)); amax = abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); const unsigned scale_byte = - amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); + amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); if ((lane & (kLanesPerMxGroup - 1)) == 0) row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + (lane / kLanesPerMxGroup)] = (unsigned char)scale_byte; @@ -130,7 +130,7 @@ __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) unsigned row_out[kWordsPerLane / 2]; #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], rcp, rcp); + row_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], rcp, rcp); if (in_bounds(j)) store_quantized(row_out_ptr + (size_t)j * out_cols, row_out); } } @@ -155,7 +155,7 @@ __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) col_amax = abs_max_bf16x2(col_amax, col_amax_partials[w][tid]); col_amax &= 0x7fff7fffu; const ptx::floatx2 amax_2x = widen_bf16x2(col_amax); - const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); const int col = col0 + 2 * tid; // The generic gated kernel writes scales only inside the tensor. if (!kEdge || (col < K && row0 < M)) { @@ -207,7 +207,7 @@ __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) unsigned col_out[kWordsPerLane / 2]; #pragma unroll for (int m = 0; m < kWordsPerLane / 2; ++m) - col_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], rcp[2 * m], rcp[2 * m + 1]); + col_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], rcp[2 * m], rcp[2 * m + 1]); if (in_bounds(j)) store_quantized(col_out_ptr + (size_t)j * out_cols, col_out); } } @@ -265,13 +265,16 @@ __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) #endif } -//! Whether the gated regtile kernel takes this request: BF16 [M, 2K] -> E4M3, both directions, +//! Whether the gated regtile kernel takes this request: BF16 [M, 2K] -> E4M3 or E5M2, both directions, //! K % 32 == 0 (the generic gated kernel's requirement), aligned buffers. inline bool can_use(const Tensor& input, const Tensor& grad, const Tensor& output, const bool is_bwd) { if (!getenv("NVTE_MXFP8_REGTILE", true)) return false; if (transformer_engine::cuda::sm_arch() < 100) return false; - if (input.dtype() != DType::kBFloat16 || output.dtype() != DType::kFloat8E4M3) return false; + if (input.dtype() != DType::kBFloat16 || + (output.dtype() != DType::kFloat8E4M3 && output.dtype() != DType::kFloat8E5M2)) { + return false; + } if (!output.has_data() || !output.has_columnwise_data()) return false; const size_t rows = input.flat_first_dim(); const size_t cols = input.flat_last_dim() / 2; @@ -285,7 +288,7 @@ inline bool can_use(const Tensor& input, const Tensor& grad, const Tensor& outpu is_aligned_ptr(output.columnwise_scale_inv.dptr, 2); } -template +template void quantize(const Tensor& input, const Tensor& grad, Tensor* output, cudaStream_t stream) { const int M = static_cast(input.flat_first_dim()); const int K = static_cast(input.flat_last_dim() / 2); @@ -298,7 +301,7 @@ void quantize(const Tensor& input, const Tensor& grad, Tensor* output, cudaStrea const dim3 grid(grid_cols, bands / iters); const auto run = [&](auto tails) { constexpr bool kTails = decltype(tails)::value; - quantize_gated_kernel + quantize_gated_kernel <<>>( (const uint4*)input.data.dptr, (const uint4*)(IS_BWD ? grad.data.dptr : nullptr), (unsigned char*)output->data.dptr, (unsigned char*)output->scale_inv.dptr, From d18f434167b3f87d93ba132deef03696197a9cfd Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 16 Oct 2026 16:29:59 -0700 Subject: [PATCH 16/29] [Common] Build the activation tables with fast math in fast-math builds With NVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON the activation kernels compile with --use_fast_math, but activation_table.cu did not, so the tables held the accurate results while the kernels computed the fast-math ones. Compile it with the same flags. Fast-math builds remain approximate for GeLU. On GB200, with every BF16 input: SiLU, dSiLU, dGeLU, dbias_dGeLU, SwiGLU and the FP16/FP32 paths match between the register-tiled and generic kernels byte for byte, but GeLU does not, with or without tables. Without tables the register-tiled output matches the generic kernel's rowwise data and differs from its colwise data in 0.02% of bytes, so the generic kernel already disagrees with itself between its two passes: the compiler builds the fast-math GeLU differently in different code. With tables 0.8% of bytes differ. Default builds are unaffected. Signed-off-by: Teddy Do --- transformer_engine/common/CMakeLists.txt | 7 ++++++- transformer_engine/common/activation/activation_table.cu | 6 ++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 6fbe2111256..76d7b3a3bc0 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -773,7 +773,12 @@ list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu option(NVTE_BUILD_ACTIVATION_WITH_FAST_MATH "Compile activation kernels with --use_fast_math option" OFF) if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH) - list(APPEND nvte_sources_with_fast_math activation/gelu.cu + # activation_table.cu builds the activation tables from util/math.h, so it takes the same + # flags as the kernels that compute the activations. Fast-math builds are approximate: the + # compiler may build GeLU slightly differently in different kernels (even in the two passes of + # one kernel), so GeLU results can differ in a small fraction of values between kernels. + list(APPEND nvte_sources_with_fast_math activation/activation_table.cu + activation/gelu.cu activation/gelu_dbias.cu activation/gelu_grouped.cu activation/gelu_grouped_bwd.cu diff --git a/transformer_engine/common/activation/activation_table.cu b/transformer_engine/common/activation/activation_table.cu index f02b3c9075c..f9129cff470 100644 --- a/transformer_engine/common/activation/activation_table.cu +++ b/transformer_engine/common/activation/activation_table.cu @@ -4,8 +4,10 @@ * See LICENSE for license information. ************************************************************************/ -// Builds the activation tables from the scalar util/math.h functions. This file must not be -// compiled with --use_fast_math: the tables have to hold exactly what util/math.h computes. +// Builds the activation tables from the scalar util/math.h functions. The tables have to hold +// what the activation kernels compute, so this file is compiled with the same flags as those +// kernels: with --use_fast_math when NVTE_BUILD_ACTIVATION_WITH_FAST_MATH is on. In such builds +// GeLU is not reproducible between kernels (see CMakeLists.txt), with or without the tables. #include #include From abbea742735b971e8cfaeae7f48483c815dafa50 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 26 Oct 2026 16:16:02 -0700 Subject: [PATCH 17/29] Drop NVTE_MXFP8_REGTILE and the regtile-vs-generic tests NVTE_MXFP8_REGTILE was a development switch for comparing the register-tiled kernels against the generic ones. Remove it with the tests that used it: - The regtile-vs-generic suites in test_cast_mxfp8_act_exact.cu. The exact test already checks every kernel against the activation computed per element and then quantized; it now also covers E5M2 output, the tensor amax, the forward without tables from 2^23 elements (4127x2064, with row and column tails) and a dbias workspace of 130 rows (8257x544). - test_cast_mxfp8_gated_regtile.cu: CastMXFP8_GatedActTestSuite covers the same requests (BF16/FP16/FP32, E4M3/E5M2, forward and backward, tails). - test_reduce_dbias.cu: the exact test checks dbias bit for bit against a host sum in the kernels' order, now with a workspace past the reduction's 128-row pipeline. Signed-off-by: Teddy Do --- docs/envvars.rst | 6 - tests/cpp/operator/CMakeLists.txt | 1 - .../cpp/operator/test_cast_mxfp8_act_exact.cu | 350 ++++-------------- .../operator/test_cast_mxfp8_gated_regtile.cu | 201 ---------- tests/cpp/util/CMakeLists.txt | 1 - tests/cpp/util/test_reduce_dbias.cu | 126 ------- .../regtile/quantize_gated_mxfp8_regtile.cuh | 1 - .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 6 +- 8 files changed, 78 insertions(+), 614 deletions(-) delete mode 100644 tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu delete mode 100644 tests/cpp/util/test_reduce_dbias.cu diff --git a/docs/envvars.rst b/docs/envvars.rst index 583f885b3a7..1836e80e1bd 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -365,12 +365,6 @@ Kernel Configuration :Default: ``0`` :Description: Experimental. Choose the SiLU-derivative lookup table (see :envvar:`NVTE_ACTIVATION_TABLE`). By default the 16 KB table covers ``|x| >= 2^-9`` and smaller inputs use the linear term ``0.5 + x/2``, which is within 2 ulp of the computed derivative. When set to ``1``, a 32 KB table covers ``|x| >= 2^-25``, and the derivative matches the computed one for every input. Read once per process, when the table is built. -.. envvar:: NVTE_MXFP8_REGTILE - - :Type: ``int`` (0 or 1) - :Default: ``1`` - :Description: Allow the register-tiled MXFP8 kernel for activations read from lookup tables (see :envvar:`NVTE_ACTIVATION_TABLE`). When set to ``0``, those requests use the generic MXFP8 kernel, which gives the same results. - .. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM :Type: ``int`` (0 or 1) diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 0e540e5a03f..e05919c2b53 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -13,7 +13,6 @@ add_executable(test_operator test_qdq.cu test_cast_mxfp8.cu test_cast_mxfp8_act_exact.cu - test_cast_mxfp8_gated_regtile.cu test_cast_mxfp8_grouped.cu test_cast_mxfp8_grouped_scaled_swiglu.cu test_cast_nvfp4_transpose.cu diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index ba0696a5384..128dcb11fcc 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -6,10 +6,11 @@ // The fused MXFP8 activation quantize must round exactly as if the activation were computed per // element, truncated to the input type, and then quantized on its own. This compares the fused -// kernels bit for bit against that composition: the activation is computed on the GPU, and the -// result is quantized by the generic cast-only kernel (a noop tensor keeps the call off the -// specialized cast kernels). dbias is compared against a host sum in the order the kernel -// accumulates it. +// kernels (the register-tiled kernel where it applies, the generic kernel elsewhere) bit for bit +// against that composition: the activation is computed on the GPU, and the result is quantized by +// the generic cast-only kernel (a noop tensor keeps the call off the specialized cast kernels). +// dbias is compared against a host sum in the order the kernels accumulate it, and the tensor amax +// against the largest magnitude of the activation. // // Without the tables (NVTE_ACTIVATION_TABLE=0, the default) the activation is the scalar util/math.h // function. With NVTE_ACTIVATION_TABLE=1, BF16 inputs follow the table specification (see util/activation_table.cuh): -inf @@ -270,9 +271,9 @@ size_t count_byte_mismatches(const uint8_t *a, const uint8_t *b, const size_t n) template void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t cols, - const bool rowwise, const bool colwise, const bool swizzled, const bool use_table) { + const bool rowwise, const bool colwise, const bool swizzled, const bool use_table, + const DType otype, const bool with_amax) { const DType itype = TypeInfo::dtype; - const DType otype = DType::kFloat8E4M3; const std::vector shape{rows, cols}; const size_t n = rows * cols; @@ -309,6 +310,14 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t // Fused. Tensor out("out", shape, otype, rowwise, colwise, NVTE_MXFP8_1D_SCALING); out.set_with_gemm_swizzled_scales(swizzled); + float *amax_dev = nullptr; + if (with_amax) { + ASSERT_EQ(cudaMalloc(&amax_dev, sizeof(float)), cudaSuccess); + ASSERT_EQ(cudaMemset(amax_dev, 0, sizeof(float)), cudaSuccess); + const size_t one = 1; + const NVTEBasicTensor amax = {amax_dev, kNVTEFloat32, nvte_make_shape(&one, 1)}; + nvte_set_tensor_param_v2(out.data(), kNVTEAmax, &amax, sizeof(amax)); + } Tensor dbias("dbias", std::vector{cols}, itype); Tensor workspace; { @@ -340,11 +349,30 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t out.to_cpu(); ref.to_cpu(); + // The FP8 data as raw bytes, whichever FP8 type the output holds. + const auto rowwise_bytes = [&](Tensor &t) { + return otype == DType::kFloat8E5M2 + ? reinterpret_cast(t.rowwise_cpu_dptr()) + : reinterpret_cast(t.rowwise_cpu_dptr()); + }; + const auto columnwise_bytes = [&](Tensor &t) { + return otype == DType::kFloat8E5M2 + ? reinterpret_cast(t.columnwise_cpu_dptr()) + : reinterpret_cast(t.columnwise_cpu_dptr()); + }; + if (with_amax) { + // The largest magnitude of the activation as quantized (rounded to IType), ignoring NaN. + y.to_cpu(); + const IType *act = y.rowwise_cpu_dptr(); + float expected = 0.0f; + for (size_t i = 0; i < n; ++i) expected = std::fmax(expected, std::fabs(static_cast(act[i]))); + float got; + ASSERT_EQ(cudaMemcpy(&got, amax_dev, sizeof(float), cudaMemcpyDeviceToHost), cudaSuccess); + cudaFree(amax_dev); + EXPECT_EQ(std::memcmp(&got, &expected, sizeof(float)), 0) << "amax " << got << " vs " << expected; + } if (rowwise) { - EXPECT_EQ(count_byte_mismatches( - reinterpret_cast(out.rowwise_cpu_dptr()), - reinterpret_cast(ref.rowwise_cpu_dptr()), n), - 0u) + EXPECT_EQ(count_byte_mismatches(rowwise_bytes(out), rowwise_bytes(ref), n), 0u) << "rowwise data"; const size_t s = product(out.rowwise_scale_inv_shape()); EXPECT_EQ(count_byte_mismatches(out.rowwise_cpu_scale_inv_ptr(), @@ -353,10 +381,7 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t << "rowwise scales"; } if (colwise) { - EXPECT_EQ(count_byte_mismatches( - reinterpret_cast(out.columnwise_cpu_dptr()), - reinterpret_cast(ref.columnwise_cpu_dptr()), n), - 0u) + EXPECT_EQ(count_byte_mismatches(columnwise_bytes(out), columnwise_bytes(ref), n), 0u) << "colwise data"; const size_t s = product(out.columnwise_scale_inv_shape()); EXPECT_EQ(count_byte_mismatches(out.columnwise_cpu_scale_inv_ptr(), @@ -379,141 +404,9 @@ void run_case(const Mode mode, const Fill fill, const size_t rows, const size_t } } -// With a table, the register-resident kernel and the generic kernel must write the same bytes. -// The first call takes the default dispatch, which picks the register-resident kernel for these -// shapes; NVTE_MXFP8_REGTILE=0 keeps the second call on the generic kernel. With with_amax, both -// also reduce the tensor amax, which must match bit for bit. -void run_regtile_vs_generic(const Mode mode, const Fill fill, const size_t rows, const size_t cols, - const bool swizzled, const bool use_table, const DType otype, - const bool with_amax, const DType itype = DType::kBFloat16) { - const std::vector shape{rows, cols}; - const size_t n = rows * cols; - Tensor x("x", shape, itype); - Tensor grad("grad", shape, itype); - if (itype == DType::kFloat16) { - fill_inputs(fill, &x, &grad); - } else if (itype == DType::kFloat32) { - fill_inputs(fill, &x, &grad); - } else { - fill_inputs(fill, &x, &grad); - } - Tensor outs[2] = {Tensor("default", shape, otype, true, true, NVTE_MXFP8_1D_SCALING), - Tensor("generic", shape, otype, true, true, NVTE_MXFP8_1D_SCALING)}; - Tensor dbias[2] = {Tensor("dbias_default", std::vector{cols}, itype), - Tensor("dbias_generic", std::vector{cols}, itype)}; - Tensor workspace[2]; - float *amax_dev = nullptr; - if (with_amax) { - ASSERT_EQ(cudaMalloc(&amax_dev, 2 * sizeof(float)), cudaSuccess); - ASSERT_EQ(cudaMemset(amax_dev, 0, 2 * sizeof(float)), cudaSuccess); - } - for (int i = 0; i < 2; ++i) { - outs[i].set_with_gemm_swizzled_scales(swizzled); - if (with_amax) { - const size_t one = 1; - const NVTEBasicTensor amax = {amax_dev + i, kNVTEFloat32, nvte_make_shape(&one, 1)}; - nvte_set_tensor_param_v2(outs[i].data(), kNVTEAmax, &amax, sizeof(amax)); - } - // Fill every output byte, padding included, with a marker, so the two kernels must write the - // same set of bytes, not just the same values. - for (const auto param : {kNVTERowwiseData, kNVTEColumnwiseData, kNVTERowwiseScaleInv, - kNVTEColumnwiseScaleInv}) { - const NVTEBasicTensor b = nvte_get_tensor_param(outs[i].data(), param); - size_t bytes = b.data_ptr == nullptr ? 0 : 1; - for (size_t d = 0; d < b.shape.ndim; ++d) bytes *= b.shape.data[d]; - if (bytes != 0) ASSERT_EQ(cudaMemset(b.data_ptr, 0xA5, bytes), cudaSuccess); - } - const ScopedEnv tables("NVTE_ACTIVATION_TABLE", use_table ? "1" : "0"); - const ScopedEnv regtile("NVTE_MXFP8_REGTILE", i == 0 ? "1" : "0"); - switch (mode) { - case Mode::GeLU: - nvte_gelu(x.data(), outs[i].data(), 0); - break; - case Mode::SiLU: - nvte_silu(x.data(), outs[i].data(), 0); - break; - case Mode::DGeLU: - nvte_dgelu(grad.data(), x.data(), outs[i].data(), 0); - break; - case Mode::DSiLU: - nvte_dsilu(grad.data(), x.data(), outs[i].data(), 0); - break; - default: { - auto fn = - mode == Mode::DBiasDGeLU ? &nvte_quantize_dbias_dgelu : &nvte_quantize_dbias_dsilu; - fn(grad.data(), x.data(), outs[i].data(), dbias[i].data(), workspace[i].data(), 0); - workspace[i] = Tensor("workspace", workspace[i].rowwise_shape(), workspace[i].dtype()); - // The per-chunk column sums must match too, not just the BF16 dbias they reduce to. - ASSERT_EQ(cudaMemset(workspace[i].rowwise_dptr(), 0xA5, - product(workspace[i].rowwise_shape()) * sizeof(float)), - cudaSuccess); - fn(grad.data(), x.data(), outs[i].data(), dbias[i].data(), workspace[i].data(), 0); - } - } - } - ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); - if (with_amax) { - float amax[2]; - ASSERT_EQ(cudaMemcpy(amax, amax_dev, sizeof(amax), cudaMemcpyDeviceToHost), cudaSuccess); - cudaFree(amax_dev); - uint32_t bits[2]; - std::memcpy(bits, amax, sizeof(bits)); - EXPECT_EQ(bits[0], bits[1]) << "amax " << amax[0] << " vs " << amax[1]; - } - for (auto &out : outs) out.to_cpu(); - if (is_dbias(mode)) { - for (int i = 0; i < 2; ++i) { - dbias[i].to_cpu(); - workspace[i].to_cpu(); - } - EXPECT_EQ(count_byte_mismatches( - reinterpret_cast(workspace[0].rowwise_cpu_dptr()), - reinterpret_cast(workspace[1].rowwise_cpu_dptr()), - product(workspace[0].rowwise_shape()) * sizeof(float)), - 0u) - << "dbias workspace"; - const auto dbias_bytes = [&](Tensor &t) { - return itype == DType::kFloat16 - ? reinterpret_cast(t.rowwise_cpu_dptr()) - : itype == DType::kFloat32 - ? reinterpret_cast(t.rowwise_cpu_dptr()) - : reinterpret_cast(t.rowwise_cpu_dptr()); - }; - const size_t dbias_elt_bytes = itype == DType::kFloat32 ? 4 : 2; - EXPECT_EQ(count_byte_mismatches(dbias_bytes(dbias[0]), dbias_bytes(dbias[1]), - cols * dbias_elt_bytes), - 0u) - << "dbias"; - } - // The FP8 data as raw bytes, whichever FP8 type the output holds. - const auto rowwise_bytes = [&](Tensor &t) { - return otype == DType::kFloat8E5M2 - ? reinterpret_cast(t.rowwise_cpu_dptr()) - : reinterpret_cast(t.rowwise_cpu_dptr()); - }; - const auto columnwise_bytes = [&](Tensor &t) { - return otype == DType::kFloat8E5M2 - ? reinterpret_cast(t.columnwise_cpu_dptr()) - : reinterpret_cast(t.columnwise_cpu_dptr()); - }; - EXPECT_EQ(count_byte_mismatches(rowwise_bytes(outs[0]), rowwise_bytes(outs[1]), n), 0u) - << "rowwise data"; - EXPECT_EQ(count_byte_mismatches(columnwise_bytes(outs[0]), columnwise_bytes(outs[1]), n), 0u) - << "colwise data"; - EXPECT_EQ(count_byte_mismatches(outs[0].rowwise_cpu_scale_inv_ptr(), - outs[1].rowwise_cpu_scale_inv_ptr(), - product(outs[0].rowwise_scale_inv_shape())), - 0u) - << "rowwise scales"; - EXPECT_EQ(count_byte_mismatches(outs[0].columnwise_cpu_scale_inv_ptr(), - outs[1].columnwise_cpu_scale_inv_ptr(), - product(outs[0].columnwise_scale_inv_shape())), - 0u) - << "colwise scales"; -} - -using Params = std::tuple, bool /*use_table*/>; +using Params = + std::tuple, bool /*use_table*/, DType /*output*/, bool /*amax*/>; class CastMXFP8ActExactTestSuite : public ::testing::TestWithParam {}; @@ -525,25 +418,26 @@ TEST_P(CastMXFP8ActExactTestSuite, MatchesScalarActivationThenCast) { if (prop.major < 10) { GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; } - const auto [mode, fill, itype, layout, swizzled, dims, use_table] = GetParam(); + const auto [mode, fill, itype, layout, swizzled, dims, use_table, otype, with_amax] = GetParam(); const bool rowwise = layout != 1; const bool colwise = layout != 0; TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( itype, IType, - run_case(mode, fill, dims.first, dims.second, rowwise, colwise, swizzled, - use_table);); + run_case(mode, fill, dims.first, dims.second, rowwise, colwise, swizzled, use_table, + otype, with_amax);); } namespace { std::string case_name(const testing::TestParamInfo &info) { const char *layouts[] = {"Rowwise", "Colwise", "Both"}; - const auto [mode, fill, itype, layout, swizzled, dims, use_table] = info.param; + const auto [mode, fill, itype, layout, swizzled, dims, use_table, otype, with_amax] = info.param; const char *fills[] = {"XTypical", "XOrderedInputs", "XSpecials", "XTinyGradients"}; return std::string(mode_name(mode)) + fills[static_cast(fill)] + "X" + test::typeName(itype) + "X" + layouts[layout] + (swizzled ? "XSwizzled" : "XCompact") + "X" + std::to_string(dims.first) + "X" + std::to_string(dims.second) + - (use_table ? "XTable" : "XNoTable"); + (use_table ? "XTable" : "XNoTable") + (otype == DType::kFloat8E5M2 ? "XE5M2" : "") + + (with_amax ? "XAmax" : ""); } } // namespace @@ -556,8 +450,30 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kBFloat16, DType::kFloat16, DType::kFloat32), ::testing::Values(0, 1, 2), ::testing::Bool(), ::testing::Values(std::pair{1024, 2048}, - std::pair{544, 2080}), - ::testing::Bool()), + std::pair{544, 2080}, + // At least 2^23 elements (the register-tiled forward + // without tables), with a row and a column tail. + std::pair{4127, 2064}, + // A dbias workspace of 130 rows, past the dbias + // reduction's 128-row pipeline. + std::pair{8257, 544}), + ::testing::Bool(), ::testing::Values(DType::kFloat8E4M3), + ::testing::Values(false)), + case_name); + +// E5M2 output and the tensor amax, both directions. +INSTANTIATE_TEST_SUITE_P( + OutputTypeAndAmax, CastMXFP8ActExactTestSuite, + ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::DBiasDGeLU, Mode::SiLU, + Mode::DSiLU, Mode::DBiasDSiLU), + ::testing::Values(Fill::Typical, Fill::Specials), + ::testing::Values(DType::kBFloat16, DType::kFloat16, DType::kFloat32), + ::testing::Values(2), ::testing::Bool(), + ::testing::Values(std::pair{1024, 2048}, + std::pair{4127, 2064}), + ::testing::Bool(), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), + ::testing::Values(true)), case_name); INSTANTIATE_TEST_SUITE_P( @@ -567,7 +483,8 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(Fill::Specials), ::testing::Values(DType::kBFloat16, DType::kFloat16, DType::kFloat32), ::testing::Values(0, 1, 2), ::testing::Bool(), - ::testing::Values(std::pair{64, 128}), ::testing::Bool()), + ::testing::Values(std::pair{64, 128}), ::testing::Bool(), + ::testing::Values(DType::kFloat8E4M3), ::testing::Values(false)), case_name); INSTANTIATE_TEST_SUITE_P(TinyGradients, CastMXFP8ActExactTestSuite, @@ -578,121 +495,6 @@ INSTANTIATE_TEST_SUITE_P(TinyGradients, CastMXFP8ActExactTestSuite, DType::kFloat32), ::testing::Values(0, 1, 2), ::testing::Bool(), ::testing::Values(std::pair{64, 128}), - ::testing::Bool()), + ::testing::Bool(), ::testing::Values(DType::kFloat8E4M3), + ::testing::Values(false)), case_name); - -namespace { - -using RegtileParams = - std::tuple, bool /*use_table*/, - DType /*output*/, bool /*amax*/>; - -class CastMXFP8ActRegtileTestSuite : public ::testing::TestWithParam {}; - -std::string regtile_case_name(const testing::TestParamInfo &info) { - const char *fills[] = {"Typical", "OrderedInputs", "Specials", "TinyGradients"}; - const auto [mode, fill, swizzled, dims, use_table, otype, with_amax] = info.param; - return std::string(mode_name(mode)) + "X" + fills[static_cast(fill)] + - (swizzled ? "XSwizzled" : "XCompact") + "X" + std::to_string(dims.first) + "X" + - std::to_string(dims.second) + (use_table ? "XTable" : "XComputed") + - (otype == DType::kFloat8E5M2 ? "XE5M2" : "XE4M3") + (with_amax ? "XAmax" : ""); -} - -} // namespace - -TEST_P(CastMXFP8ActRegtileTestSuite, MatchesGenericKernel) { - cudaDeviceProp prop; - ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); - if (prop.major < 10) { - GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; - } - const auto [mode, fill, swizzled, dims, use_table, otype, with_amax] = GetParam(); - run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled, use_table, otype, - with_amax); -} - -INSTANTIATE_TEST_SUITE_P( - OperatorTest, CastMXFP8ActRegtileTestSuite, - ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::DBiasDGeLU, Mode::SiLU, - Mode::DSiLU, Mode::DBiasDSiLU), - ::testing::Values(Fill::Typical, Fill::OrderedInputs, Fill::Specials), - ::testing::Bool(), - ::testing::Values(std::pair{1024, 2048}, - std::pair{4160, 2048}, - std::pair{128, 256}, - // Row and column tails; the last column tile ends - // 32 or 16 columns into the tile, mid MX group for 16. - std::pair{993, 7200}, - std::pair{1, 16}, - std::pair{33, 272}, - std::pair{100, 2064}, - std::pair{4127, 528}, - // Large enough for the computed forward path (2^23 - // elements), with a row and a column tail. - std::pair{4127, 2064}), - // With a table, or computed from util/math.h without one. - ::testing::Bool(), - ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), - ::testing::Bool()), - regtile_case_name); - -// FP16 inputs take the register-tiled kernel too, always computing the activation (the tables -// hold BF16 inputs). The table setting must not matter for them. -class CastMXFP8ActRegtileFP16TestSuite : public ::testing::TestWithParam {}; - -TEST_P(CastMXFP8ActRegtileFP16TestSuite, MatchesGenericKernel) { - cudaDeviceProp prop; - ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); - if (prop.major < 10) { - GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; - } - const auto [mode, fill, swizzled, dims, use_table, otype, with_amax] = GetParam(); - run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled, use_table, otype, - with_amax, DType::kFloat16); -} - -INSTANTIATE_TEST_SUITE_P( - OperatorTest, CastMXFP8ActRegtileFP16TestSuite, - ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::DBiasDGeLU, Mode::SiLU, - Mode::DSiLU, Mode::DBiasDSiLU), - ::testing::Values(Fill::Typical, Fill::OrderedInputs, Fill::Specials), - ::testing::Bool(), - ::testing::Values(std::pair{1024, 2048}, - std::pair{993, 7200}, - std::pair{33, 272}, - // Large enough for the computed forward path. - std::pair{4127, 2064}), - ::testing::Bool(), - ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), - ::testing::Bool()), - regtile_case_name); - -// FP32 inputs: no table, twice the threads per CTA (see regtile::TileConfig). -class CastMXFP8ActRegtileFP32TestSuite : public ::testing::TestWithParam {}; - -TEST_P(CastMXFP8ActRegtileFP32TestSuite, MatchesGenericKernel) { - cudaDeviceProp prop; - ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); - if (prop.major < 10) { - GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; - } - const auto [mode, fill, swizzled, dims, use_table, otype, with_amax] = GetParam(); - run_regtile_vs_generic(mode, fill, dims.first, dims.second, swizzled, use_table, otype, - with_amax, DType::kFloat32); -} - -INSTANTIATE_TEST_SUITE_P( - OperatorTest, CastMXFP8ActRegtileFP32TestSuite, - ::testing::Combine(::testing::Values(Mode::GeLU, Mode::DGeLU, Mode::DBiasDGeLU, Mode::SiLU, - Mode::DSiLU, Mode::DBiasDSiLU), - ::testing::Values(Fill::Typical, Fill::OrderedInputs, Fill::Specials), - ::testing::Bool(), - ::testing::Values(std::pair{1024, 2048}, - std::pair{993, 7200}, - std::pair{33, 272}, - // Large enough for the computed forward path. - std::pair{4127, 2064}), - ::testing::Bool(), - ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), - ::testing::Bool()), - regtile_case_name); diff --git a/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu b/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu deleted file mode 100644 index 4322808dd57..00000000000 --- a/tests/cpp/operator/test_cast_mxfp8_gated_regtile.cu +++ /dev/null @@ -1,201 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -// The register-tiled SwiGLU / dSwiGLU MXFP8 kernel must write exactly the bytes of the generic -// gated kernel. Each case runs both (NVTE_MXFP8_REGTILE=0 keeps the generic kernel) into outputs -// pre-filled with a marker byte, so they must also write the same set of bytes. The AllBF16 fill -// puts every BF16 value in the activated half. - -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "../test_common.h" -#include "transformer_engine/transformer_engine.h" - -using namespace transformer_engine; -using namespace test; - -namespace { - -enum class GatedFill { Normal, Wide, Specials, AllBF16 }; - -class ScopedEnv { - public: - ScopedEnv(const char *name, const char *value) : name_(name) { - const char *old = std::getenv(name); - if (old != nullptr) old_ = old; - setenv(name, value, 1); - } - ~ScopedEnv() { - if (old_) { - setenv(name_, old_->c_str(), 1); - } else { - unsetenv(name_); - } - } - - private: - const char *name_; - std::optional old_; -}; - -void fill(Tensor *t, const size_t rows, const size_t cols, const GatedFill kind, const int seed) { - std::mt19937 gen(seed); - std::normal_distribution normal(0.0f, 3.0f); - std::uniform_int_distribution exponent(-10, 9); - const float specials[] = {std::numeric_limits::infinity(), - -std::numeric_limits::infinity(), - std::numeric_limits::quiet_NaN(), 0.0f, -0.0f, 1e-30f, -1e30f}; - bf16 *p = t->rowwise_cpu_dptr(); - for (size_t r = 0; r < rows; ++r) { - // Wide: every 32x32 block gets its own power-of-two scale. - for (size_t c = 0; c < cols; ++c) { - float v = normal(gen); - if (kind == GatedFill::AllBF16 && c < cols / 2) { - // The activated half (the first cols / 2 columns) walks through every BF16 bit pattern. - const uint16_t bits = static_cast(r * (cols / 2) + c); - std::memcpy(&p[r * cols + c], &bits, sizeof(bits)); - continue; - } - if (kind == GatedFill::Wide) { - std::mt19937 block_gen(seed * 7919 + (r / 32) * 104729 + c / 32); - v *= std::ldexp(1.0f, exponent(block_gen)); - } else if (kind == GatedFill::Specials && gen() % 61 == 0) { - v = specials[gen() % 7]; - } - p[r * cols + c] = static_cast(v); - } - } - t->from_cpu(); -} - -size_t count_byte_mismatches(const uint8_t *a, const uint8_t *b, const size_t n) { - size_t m = 0; - for (size_t i = 0; i < n; ++i) m += (a[i] != b[i]); - return m; -} - -void run_case(const bool bwd, const GatedFill kind, const size_t rows, const size_t cols, - const bool swizzled, const DType otype) { - Tensor input("input", std::vector{rows, 2 * cols}, DType::kBFloat16); - Tensor grad("grad", std::vector{rows, cols}, DType::kBFloat16); - fill(&input, rows, 2 * cols, kind, 1); - fill(&grad, rows, cols, - kind == GatedFill::Specials || kind == GatedFill::AllBF16 ? GatedFill::Normal : kind, 2); - const size_t out_cols = bwd ? 2 * cols : cols; - const std::vector out_shape{rows, out_cols}; - Tensor outs[2] = { - Tensor("regtile", out_shape, otype, true, true, NVTE_MXFP8_1D_SCALING), - Tensor("generic", out_shape, otype, true, true, NVTE_MXFP8_1D_SCALING)}; - for (int i = 0; i < 2; ++i) { - outs[i].set_with_gemm_swizzled_scales(swizzled); - for (const auto param : {kNVTERowwiseData, kNVTEColumnwiseData, kNVTERowwiseScaleInv, - kNVTEColumnwiseScaleInv}) { - const NVTEBasicTensor b = nvte_get_tensor_param(outs[i].data(), param); - size_t bytes = b.data_ptr == nullptr ? 0 : 1; - for (size_t d = 0; d < b.shape.ndim; ++d) bytes *= b.shape.data[d]; - if (bytes != 0) ASSERT_EQ(cudaMemset(b.data_ptr, 0xA5, bytes), cudaSuccess); - } - const ScopedEnv regtile("NVTE_MXFP8_REGTILE", i == 0 ? "1" : "0"); - if (bwd) { - nvte_dswiglu(grad.data(), input.data(), outs[i].data(), 0); - } else { - nvte_swiglu(input.data(), outs[i].data(), 0); - } - } - ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); - ASSERT_EQ(cudaGetLastError(), cudaSuccess); - for (auto &out : outs) out.to_cpu(); - const size_t n = rows * out_cols; - // The FP8 data as raw bytes, whichever FP8 type the output holds. - const auto rowwise_bytes = [&](Tensor &t) { - return otype == DType::kFloat8E5M2 - ? reinterpret_cast(t.rowwise_cpu_dptr()) - : reinterpret_cast(t.rowwise_cpu_dptr()); - }; - const auto columnwise_bytes = [&](Tensor &t) { - return otype == DType::kFloat8E5M2 - ? reinterpret_cast(t.columnwise_cpu_dptr()) - : reinterpret_cast(t.columnwise_cpu_dptr()); - }; - EXPECT_EQ(count_byte_mismatches(rowwise_bytes(outs[0]), rowwise_bytes(outs[1]), n), 0u) - << "rowwise data"; - EXPECT_EQ(count_byte_mismatches(columnwise_bytes(outs[0]), columnwise_bytes(outs[1]), n), 0u) - << "colwise data"; - EXPECT_EQ(count_byte_mismatches(outs[0].rowwise_cpu_scale_inv_ptr(), - outs[1].rowwise_cpu_scale_inv_ptr(), - product(outs[0].rowwise_scale_inv_shape())), - 0u) - << "rowwise scales"; - EXPECT_EQ(count_byte_mismatches(outs[0].columnwise_cpu_scale_inv_ptr(), - outs[1].columnwise_cpu_scale_inv_ptr(), - product(outs[0].columnwise_scale_inv_shape())), - 0u) - << "colwise scales"; -} - -using Params = std::tuple, - DType /*output*/>; - -class GatedRegtileTestSuite : public ::testing::TestWithParam {}; - -TEST_P(GatedRegtileTestSuite, MatchesGenericKernel) { - cudaDeviceProp prop; - ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); - if (prop.major < 10) { - GTEST_SKIP() << "This MXFP8 quantize kernel requires compute capability 10.0 or newer"; - } - const auto [bwd, kind, swizzled, dims, otype] = GetParam(); - run_case(bwd, kind, dims.first, dims.second, swizzled, otype); -} - -std::string case_name(const testing::TestParamInfo &info) { - const auto [bwd, kind, swizzled, dims, otype] = info.param; - const char *fills[] = {"Normal", "Wide", "Specials", "AllBF16"}; - return std::string(bwd ? "DSwiGLU" : "SwiGLU") + "X" + fills[static_cast(kind)] + "X" + - (swizzled ? "Swizzled" : "Compact") + "X" + std::to_string(dims.first) + "X" + - std::to_string(dims.second) + (otype == DType::kFloat8E5M2 ? "XE5M2" : "XE4M3"); -} - -INSTANTIATE_TEST_SUITE_P( - OperatorTest, GatedRegtileTestSuite, - ::testing::Combine(::testing::Bool(), - ::testing::Values(GatedFill::Normal, GatedFill::Wide, GatedFill::Specials), - ::testing::Bool(), - ::testing::Values(std::pair{1024, 2048}, - std::pair{4160, 2048}, - std::pair{128, 256}, - // Row and column tails (K % 32 == 0). - std::pair{993, 7200}, - std::pair{1, 32}, - std::pair{33, 288}, - std::pair{100, 2080}), - ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2)), - case_name); - -// Every BF16 value in the activated half: 64 rows x 1024 columns. -INSTANTIATE_TEST_SUITE_P(AllBF16, GatedRegtileTestSuite, - ::testing::Combine(::testing::Bool(), ::testing::Values(GatedFill::AllBF16), - ::testing::Bool(), - ::testing::Values(std::pair{64, 1024}), - ::testing::Values(DType::kFloat8E4M3, - DType::kFloat8E5M2)), - case_name); - -} // namespace diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index 2ff5fa9ade3..24acec658c6 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -4,7 +4,6 @@ add_executable(test_util test_activation_table.cu - 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 deleted file mode 100644 index a5d26230131..00000000000 --- a/tests/cpp/util/test_reduce_dbias.cu +++ /dev/null @@ -1,126 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -// The dbias reduction of the MXFP8 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: - // multiples of 4 (the output vector of the BF16/FP16 store), around the 64-thread CTA. - ::testing::Combine(::testing::Values(1, 2, 31, 32, 33, 63, 64, 65, 512, 1000, 2000), - ::testing::Values(4, 28, 60, 64, 68, 1028, 2052, 7168)), - [](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/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh index dd7c7cf98d3..ea3d056221e 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh @@ -269,7 +269,6 @@ __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) //! K % 32 == 0 (the generic gated kernel's requirement), aligned buffers. inline bool can_use(const Tensor& input, const Tensor& grad, const Tensor& output, const bool is_bwd) { - if (!getenv("NVTE_MXFP8_REGTILE", true)) return false; if (transformer_engine::cuda::sm_arch() < 100) return false; if (input.dtype() != DType::kBFloat16 || (output.dtype() != DType::kFloat8E4M3 && output.dtype() != DType::kFloat8E5M2)) { diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh index 24de19adebf..ebe2dba3900 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -921,13 +921,11 @@ constexpr bool kTakesFP16 = // dGeLU 1.29-1.48x, dSiLU 1.48-2.37x, dbias_dGeLU 1.16-1.32x, dbias_dSiLU 1.35-2.00x. // // Whether this kernel takes a request of a type in kImplements, with BF16, FP16 or FP32 input and E4M3 or -// E5M2 output. Anything else takes the generic kernel, as does every request with -// NVTE_MXFP8_REGTILE=0. A noop tensor is ignored, as the generic kernel ignores it for every -// request that computes an activation. +// E5M2 output. Anything else takes the generic kernel. A noop tensor is ignored, as the generic +// kernel ignores it for every request that computes an activation. template bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, const Tensor& noop, bool use_2d_quantization, const activation_table::TableRef& table) { - if (!getenv("NVTE_MXFP8_REGTILE", true)) return false; const auto [rows, cols] = input.flat_2d_dims(); // Without a table the forward activation is compute-bound here, and on small tensors too // few CTAs hide its latency: GB200, 993x7200 GeLU is 0.93x of the generic kernel, 4160x2048 From f36459c50475cbac3aa1aca1211ad3c14e511f9f Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 26 Oct 2026 16:16:10 -0700 Subject: [PATCH 18/29] Template the activation tables on the input type Table and ExactDsiluTable, with BF16 specializations only. The tables hold the activation, not its quantization, so the same tables serve any output format; FP16 tables would be further specializations. test_activation_table.cu describes each table by a per-kind Spec (activation, derivative window and slope), so the lookup and check helpers no longer carry dSiLU-specific parameters. Shorten the CMake comment on building activation_table.cu with fast math. Signed-off-by: Teddy Do --- tests/cpp/util/test_activation_table.cu | 132 +++++++++--------- transformer_engine/common/CMakeLists.txt | 5 +- .../common/activation/activation_table.cu | 10 +- .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 10 +- .../common/util/activation_table.cuh | 51 ++++--- 5 files changed, 105 insertions(+), 103 deletions(-) diff --git a/tests/cpp/util/test_activation_table.cu b/tests/cpp/util/test_activation_table.cu index 728a6570c45..009fa322eb5 100644 --- a/tests/cpp/util/test_activation_table.cu +++ b/tests/cpp/util/test_activation_table.cu @@ -27,6 +27,7 @@ #include #include #include +#include #include #include "../test_common.h" @@ -44,80 +45,72 @@ using transformer_engine::activation_table::kSilu; constexpr int kNumInputs = 1 << 16; -// Below-window bounds and linear forms of the specification. -constexpr uint32_t kDgeluWindowLow = 0x3980; // 2^-12 -constexpr uint32_t kDsiluWindowLow = 0x3B00; // 2^-9 -constexpr uint32_t kExactDsiluWindowLow = 0x3300; // 2^-25 - -bool exact_dsilu() { - const char *env = std::getenv("NVTE_ACTIVATION_TABLE_EXACT_DSILU"); - return env != nullptr && std::string(env) == "1"; -} - -__device__ float spec(const Kind kind, const uint32_t bits, - const uint32_t dsilu_window_low) { - const float x = __uint_as_float(bits << 16); - const uint32_t magnitude = bits & 0x7fffu; - switch (kind) { - case kGelu: - if (bits == 0xff80u) return -0.0f; - return __bfloat162float(__float2bfloat16_rn(transformer_engine::gelu(x, {}))); - case kSilu: - if (bits == 0xff80u) return -0.0f; - return __bfloat162float(__float2bfloat16_rn(transformer_engine::silu(x, {}))); - case kDgelu: - if (magnitude < kDgeluWindowLow) return fmaf(0.79788456f, x, 0.5f); - return transformer_engine::dgelu(x, {}); - default: - if (magnitude < dsilu_window_low) return fmaf(0.5f, x, 0.5f); - return transformer_engine::dsilu(x, {}); - } -} - -__device__ float scalar_math(const Kind kind, const float x) { - switch (kind) { - case kDgelu: - return transformer_engine::dgelu(x, {}); - default: - return transformer_engine::dsilu(x, {}); - } -} - +// The specification of each table, written out independently of util/activation_table.cuh: the +// activation and, for derivatives, the first magnitude of the table window and the slope of the +// linear term 0.5 + slope*x below it. template -struct Op; +struct Spec; template <> -struct Op { - static constexpr float (*value)(float, const Empty &) = transformer_engine::gelu; +struct Spec { + static constexpr float (*op)(float, const Empty &) = transformer_engine::gelu; + static constexpr bool kForward = true; }; template <> -struct Op { - static constexpr float (*value)(float, const Empty &) = transformer_engine::silu; +struct Spec { + static constexpr float (*op)(float, const Empty &) = transformer_engine::silu; + static constexpr bool kForward = true; }; template <> -struct Op { - static constexpr float (*value)(float, const Empty &) = transformer_engine::dgelu; +struct Spec { + static constexpr float (*op)(float, const Empty &) = transformer_engine::dgelu; + static constexpr bool kForward = false; + static constexpr float kSlope = 0.79788456f; // sqrt(2/pi) + static uint32_t window_low() { return 0x3980; } // 2^-12 }; template <> -struct Op { - static constexpr float (*value)(float, const Empty &) = transformer_engine::dsilu; +struct Spec { + static constexpr float (*op)(float, const Empty &) = transformer_engine::dsilu; + static constexpr bool kForward = false; + static constexpr float kSlope = 0.5f; + static constexpr uint32_t kExactWindowLow = 0x3300; // 2^-25 + // 2^-9, or kExactWindowLow for the exact table (NVTE_ACTIVATION_TABLE_EXACT_DSILU=1). + static uint32_t window_low() { + const char *env = std::getenv("NVTE_ACTIVATION_TABLE_EXACT_DSILU"); + return env != nullptr && std::string(env) == "1" ? kExactWindowLow : 0x3B00; + } }; +// The table's value for the BF16 input bits. window_low is unused for forward tables. +template +__device__ float spec(const uint32_t bits, const uint32_t window_low) { + using S = Spec; + const float x = __uint_as_float(bits << 16); + if constexpr (S::kForward) { + if (bits == 0xff80u) return -0.0f; + return __bfloat162float(__float2bfloat16_rn(S::op(x, {}))); + } else { + if ((bits & 0x7fffu) < window_low) return fmaf(S::kSlope, x, 0.5f); + return S::op(x, {}); + } +} + +// For every BF16 input: the activation as evaluate() reads or computes it with kUse, the +// specification, and (derivatives only) util/math.h. template -__global__ void lookup_kernel(const void *table, const uint32_t dsilu_window_low, float *looked_up, +__global__ void lookup_kernel(const void *table, const uint32_t window_low, float *looked_up, float *expected, float *math_h) { const uint32_t bits = blockIdx.x * blockDim.x + threadIdx.x; if (bits >= kNumInputs) return; namespace at = transformer_engine::activation_table; + using S = Spec; const float x = __uint_as_float(bits << 16); - looked_up[bits] = at::evaluate::value>(x, table); - if constexpr (kUse == at::kComputedTable && (kKind == kGelu || kKind == kSilu)) { + looked_up[bits] = at::evaluate(x, table); + if constexpr (kUse == at::kComputedTable && S::kForward) { // The computed forward is FP32; every caller rounds it to BF16, as the table already is. looked_up[bits] = __bfloat162float(__float2bfloat16_rn(looked_up[bits])); } - expected[bits] = spec(kKind, bits, dsilu_window_low); - math_h[bits] = (kKind == kDgelu || kKind == kDsilu) - ? scalar_math(kKind, x) - : 0.0f; + expected[bits] = spec(bits, window_low); + math_h[bits] = S::kForward ? 0.0f : S::op(x, {}); } uint32_t float_bits(const float f) { @@ -135,27 +128,30 @@ int64_t ordered(const float f) { } // kComputed: check what the generic kernel computes in place of the table (kComputedTable, or -// util/math.h itself for the exact dSiLU table) instead of the table lookup. +// util/math.h itself for the exact table) instead of the table lookup. template void check_table(const int max_ulp_below_window) { namespace at = transformer_engine::activation_table; const at::TableRef table = at::device_table(kKind); - const bool exact = kKind == kDsilu && exact_dsilu(); ASSERT_NE(table.ptr, nullptr); - ASSERT_EQ(table.use, exact ? at::kExactDsiluTable : at::kDefaultTable); + uint32_t window_low = 0; + if constexpr (!Spec::kForward) { + window_low = Spec::window_low(); + } float *d = nullptr; ASSERT_EQ(cudaMalloc(&d, 3 * kNumInputs * sizeof(float)), cudaSuccess); - const uint32_t dsilu_window_low = exact ? kExactDsiluWindowLow : kDsiluWindowLow; - if (exact) { + const auto launch = [&](auto use) { + lookup_kernel<<>>( + table.ptr, window_low, d, d + kNumInputs, d + 2 * kNumInputs); + }; + if (table.use == at::kExactDsiluTable) { + // The exact table's window ends where the linear term equals util/math.h. if constexpr (kKind == kDsilu) { - lookup_kernel - <<>>(table.ptr, dsilu_window_low, d, d + kNumInputs, - d + 2 * kNumInputs); + launch(std::integral_constant{}); } } else { - lookup_kernel - <<>>(table.ptr, dsilu_window_low, d, d + kNumInputs, - d + 2 * kNumInputs); + ASSERT_EQ(table.use, at::kDefaultTable); + launch(std::integral_constant{}); } ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); std::vector h(3 * kNumInputs); @@ -196,7 +192,7 @@ TEST(ActivationTableTest, DgeluMatchesSpecification) { check_table(1); } TEST(ActivationTableTest, DsiluMatchesSpecification) { - check_table(exact_dsilu() ? 0 : 2); + check_table(Spec::window_low() == Spec::kExactWindowLow ? 0 : 2); } TEST(ActivationTableTest, ComputedGeluMatchesSpecification) { check_table(-1); @@ -208,7 +204,7 @@ TEST(ActivationTableTest, ComputedDgeluMatchesSpecification) { check_table(1); } TEST(ActivationTableTest, ComputedDsiluMatchesSpecification) { - check_table(exact_dsilu() ? 0 : 2); + check_table(Spec::window_low() == Spec::kExactWindowLow ? 0 : 2); } TEST(ActivationTableTest, OneTablePerDevice) { diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 76d7b3a3bc0..9525e28b3b2 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -773,10 +773,7 @@ list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu option(NVTE_BUILD_ACTIVATION_WITH_FAST_MATH "Compile activation kernels with --use_fast_math option" OFF) if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH) - # activation_table.cu builds the activation tables from util/math.h, so it takes the same - # flags as the kernels that compute the activations. Fast-math builds are approximate: the - # compiler may build GeLU slightly differently in different kernels (even in the two passes of - # one kernel), so GeLU results can differ in a small fraction of values between kernels. + # activation_table.cu builds the tables from util/math.h, so it takes the kernels' flags. list(APPEND nvte_sources_with_fast_math activation/activation_table.cu activation/gelu.cu activation/gelu_dbias.cu diff --git a/transformer_engine/common/activation/activation_table.cu b/transformer_engine/common/activation/activation_table.cu index f9129cff470..d9a82e1164a 100644 --- a/transformer_engine/common/activation/activation_table.cu +++ b/transformer_engine/common/activation/activation_table.cu @@ -72,19 +72,19 @@ void build(void *table, cudaStream_t stream) { void build(const Kind kind, void *table, cudaStream_t stream) { switch (kind) { case kGelu: - build>(table, stream); + build>(table, stream); break; case kSilu: - build>(table, stream); + build>(table, stream); break; case kDgelu: - build>(table, stream); + build>(table, stream); break; case kDsilu: if (use_exact_dsilu_table()) { - build(table, stream); + build>(table, stream); } else { - build>(table, stream); + build>(table, stream); } break; default: diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh index ebe2dba3900..cbec0cfe140 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh @@ -96,12 +96,12 @@ __device__ __forceinline__ unsigned amax_to_e8m0_2x(float amax_hi, float amax_lo fmaxf(amax_lo * kMaxNormRcp, 0.0f)); } -/*! \brief Table's activation computed with util/math.h instead of read from a table: +/*! \brief Table's activation computed with util/math.h instead of read from a table: * the generic kernel's output without tables (NVTE_ACTIVATION_TABLE=0). No table is * copied. */ template -struct ComputedTable : activation_table::Table { +struct ComputedTable : activation_table::Table { static constexpr bool kComputed = true; static constexpr size_t kBytes = 16; // unused shared-memory slot }; @@ -905,7 +905,7 @@ template constexpr bool kImplements = (!IS_DBIAS || IS_DACT) && (IS_ACT != IS_DACT) && activation_table::kHasTable && - (IS_ACT == activation_table::Table>::kForward); + (IS_ACT == activation_table::Table>::kForward); //! Whether FP16 inputs of this request type take this kernel. FP16 always computes the //! activation and scales every value in FP32 (see Datapath), about twice the @@ -986,13 +986,13 @@ void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, return; } else if constexpr (kKind == kDsilu) { if (table.use == kExactDsiluTable) { - run(std::common_type{}); + run(std::common_type>{}); NVTE_CHECK_CUDA(cudaGetLastError()); return; } } if constexpr (std::is_same_v) { - run(std::common_type>{}); + run(std::common_type>{}); NVTE_CHECK_CUDA(cudaGetLastError()); } } diff --git a/transformer_engine/common/util/activation_table.cuh b/transformer_engine/common/util/activation_table.cuh index 622cdfcd0e3..87bfaf9426e 100644 --- a/transformer_engine/common/util/activation_table.cuh +++ b/transformer_engine/common/util/activation_table.cuh @@ -98,25 +98,31 @@ struct DerivativeLayout { static_assert(kStride == 2048 || kStride == 4096, "The packed lookup assumes these strides."); }; -template +/*! \brief The table of kKind for inputs of type IType. + * + * Only BF16 inputs have tables. The tables hold the activation, not its quantization, so the + * same tables serve any output format. FP16 inputs also take 2^16 values; tables for them + * would be specializations here, with a forward window of their own. + */ +template struct Table; //! gelu(x) is RN(x/2) below 0x3b4b, x from 2.75 and -0 from -5.1875. template <> -struct Table : ForwardLayout<0x3b4b, 0x40a5> { +struct Table<__nv_bfloat16, kGelu> : ForwardLayout<0x3b4b, 0x40a5> { __device__ static float op(const float x) { return gelu(x, {}); } }; //! silu(x) is RN(x/2) below 0x3bb6, x from 5.96875 and -0 from -89. template <> -struct Table : ForwardLayout<0x3bb6, 0x42b1> { +struct Table<__nv_bfloat16, kSilu> : ForwardLayout<0x3bb6, 0x42b1> { __device__ static float op(const float x) { return silu(x, {}); } }; //! [2^-12, 8]. dgelu has saturated, to 1 and to +0, by |x| = 8, until 0.107*x*x overflows at //! 0x6044. template <> -struct Table : DerivativeLayout<0x3980, 0x4100, 0x6044> { +struct Table<__nv_bfloat16, kDgelu> : DerivativeLayout<0x3980, 0x4100, 0x6044> { __device__ static float op(const float x) { return dgelu(x, {}); } //! dgelu(x) = 0.5 + sqrt(2/pi)*x + O(x^3). __device__ static float below_window(const float x) { return fmaf(0.79788456f, x, 0.5f); } @@ -124,7 +130,7 @@ struct Table : DerivativeLayout<0x3980, 0x4100, 0x6044> { //! [2^-9, 89]. dsilu has saturated, to 1 and to +0, by |x| = 89, up to inf. template <> -struct Table : DerivativeLayout<0x3B00, 0x42B2, 0x7F80> { +struct Table<__nv_bfloat16, kDsilu> : DerivativeLayout<0x3B00, 0x42B2, 0x7F80> { __device__ static float op(const float x) { return dsilu(x, {}); } //! dsilu(x) = 0.5 + x/2 + O(x^3). __device__ static float below_window(const float x) { return fmaf(0.5f, x, 0.5f); } @@ -132,22 +138,25 @@ struct Table : DerivativeLayout<0x3B00, 0x42B2, 0x7F80> { //! dSiLU over [2^-25, 89]: below 2^-25 the linear term already equals util/math.h. Selected //! with NVTE_ACTIVATION_TABLE_EXACT_DSILU=1. -struct ExactDsiluTable : DerivativeLayout<0x3300, 0x42B2, 0x7F80, 4096> { +template +struct ExactDsiluTable; +template <> +struct ExactDsiluTable<__nv_bfloat16> : DerivativeLayout<0x3300, 0x42B2, 0x7F80, 4096> { __device__ static float op(const float x) { return dsilu(x, {}); } __device__ static float below_window(const float x) { return fmaf(0.5f, x, 0.5f); } }; -//! Table size in bytes, in the default layout. +//! BF16 table size in bytes, in the default layout. constexpr size_t table_bytes(const Kind kind) { switch (kind) { case kGelu: - return Table::kBytes; + return Table<__nv_bfloat16, kGelu>::kBytes; case kSilu: - return Table::kBytes; + return Table<__nv_bfloat16, kSilu>::kBytes; case kDgelu: - return Table::kBytes; + return Table<__nv_bfloat16, kDgelu>::kBytes; case kDsilu: - return Table::kBytes; + return Table<__nv_bfloat16, kDsilu>::kBytes; default: return 0; } @@ -162,7 +171,7 @@ inline bool use_exact_dsilu_table() { //! Size of the table that device_table builds for kind. inline size_t built_table_bytes(const Kind kind) { if (kind == kDsilu && use_exact_dsilu_table()) { - return ExactDsiluTable::kBytes; + return ExactDsiluTable<__nv_bfloat16>::kBytes; } return table_bytes(kind); } @@ -170,9 +179,9 @@ inline size_t built_table_bytes(const Kind kind) { //! Which table layout a kernel reads. enum TableUse : int { kNoTable = 0, //!< Compute OP with util/math.h. - kDefaultTable = 1, //!< Table. - kExactDsiluTable = 2, //!< ExactDsiluTable. - //! No table read: compute with util/math.h what Table holds. Gives the bytes of + kDefaultTable = 1, //!< Table. + kExactDsiluTable = 2, //!< ExactDsiluTable. + //! No table read: compute with util/math.h what Table holds. Gives the bytes of //! kDefaultTable, without the memory latency of a lookup in global memory. kComputedTable = 3, }; @@ -183,7 +192,7 @@ struct TableRef { TableUse use = kNoTable; }; -/*! \brief The current device's table for kind, built on first use. +/*! \brief The current device's BF16 table for kind, built on first use. * * Each table is built once per device and process, then kept until the process exits * (under 64 KB per device for all of them). The first call may come while the caller's @@ -343,7 +352,7 @@ __device__ __forceinline__ float evaluate(const float x, const void *table) { // Branch-free on purpose: a branch around OP keeps the compiler from overlapping the // latency of OP across elements, which costs more than computing OP for every element. static_assert(kHasTable, "OP has no table."); - using T = Table>; + using T = Table<__nv_bfloat16, kKindOf>; if constexpr (T::kForward) { // Outside the window the table's closed forms equal util/math.h rounded to BF16, // except at -inf, where the formula gives NaN and the table -0. util/math.h gives -0 @@ -368,11 +377,11 @@ __device__ __forceinline__ float evaluate(const float x, const void *table) { constexpr Kind kKind = kKindOf; if constexpr (kUse == kExactDsiluTable) { static_assert(kKind == kDsilu, "Only dSiLU has an exact table."); - return derivative(x, static_cast(table)); - } else if constexpr (Table::kForward) { - return forward>(x, static_cast(table)); + return derivative>(x, static_cast(table)); + } else if constexpr (Table<__nv_bfloat16, kKind>::kForward) { + return forward>(x, static_cast(table)); } else { - return derivative>(x, static_cast(table)); + return derivative>(x, static_cast(table)); } } } From 67a83ea5c02142e65fc43de46807800d8f80b632 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 26 Oct 2026 23:09:07 -0700 Subject: [PATCH 19/29] [Common] Read the activation tables in the CuTeDSL MXFP8 kernels With NVTE_ACTIVATION_TABLE=1 the CuTeDSL backend now follows the table specification for BF16 GeLU, SiLU, dGeLU and dSiLU, as the CUDA kernels do: - The dispatcher passes the table and its layout ("default", or "exact" for the exact dSiLU table) to the CuTeDSL entry point. The layout is part of the kernel config; kernel names change only when a table is used. The table pointer is an always-present argument, null without a table, like the noop flag. - The register-resident bidimensional kernel copies the table into shared memory and reads it with packed BF16 lookups ported from util/activation_table.cuh. With a table it walks several row tiles per CTA so the table is copied once per CTA, uses 256-column tiles for the derivatives, and indexes tiles of a run-time row stride from a 64-bit base with 32-bit offsets. Without a table it is unchanged. - The other kernels compute what the table holds: -inf gives -0 in the forward, and the derivatives are the linear term below the table window. Every BF16 case of the exact-composition test that reaches the register- resident kernel matches with tables on and off; the remaining mismatches of the CuTeDSL backend are in its general kernel and do not depend on the tables. The test gains a 4096x2048 shape, on which the register-resident kernels walk. GB200, BF16, geometric mean speedup over TE main with tables: GeLU 2.12, SiLU 2.38, dGeLU 2.06, dSiLU 2.02 (CUDA backend: 2.09, 2.34, 2.14, 2.10). Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 3 + .../common/CuTeDSL/activations.py | 168 +++++++++++- .../cast/mxfp8/quantize_mxfp8_common.py | 54 +++- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 247 ++++++++++++++++-- .../mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py | 1 + .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 11 +- .../cast/mxfp8/quantize_mxfp8_g2s_bidim.py | 1 + .../cast/mxfp8/quantize_mxfp8_g2s_general.py | 23 +- .../common/cast/dispatch/quantize.cuh | 21 +- .../cast/mxfp8/quantize_mxfp8_cutedsl.cuh | 53 +++- 10 files changed, 519 insertions(+), 63 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index 128dcb11fcc..e4a5c3f8a17 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -450,6 +450,9 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kBFloat16, DType::kFloat16, DType::kFloat32), ::testing::Values(0, 1, 2), ::testing::Bool(), ::testing::Values(std::pair{1024, 2048}, + // Enough row tiles that the register-resident kernels + // walk several per CTA. + std::pair{4096, 2048}, std::pair{544, 2080}, // At least 2^23 elements (the register-tiled forward // without tables), with a row and a column tail. diff --git a/transformer_engine/common/CuTeDSL/activations.py b/transformer_engine/common/CuTeDSL/activations.py index 34772eb7257..c1de06139e0 100644 --- a/transformer_engine/common/CuTeDSL/activations.py +++ b/transformer_engine/common/CuTeDSL/activations.py @@ -8,9 +8,11 @@ # Each function is a 1:1 port of its math.h counterpart (see module docstring). # pylint: disable=missing-function-docstring import os +import struct +import cutlass from cutlass import cute -from cutlass import Float32 +from cutlass import Float32, Int32, Uint32 from transformer_engine.common.CuTeDSL.utils import fma_f32 @@ -80,3 +82,167 @@ def dact_dgelu(x: Float32) -> Float32: return Float32(0.5) * x * ( (Float32(1.0) - t * t) * (Float32(0.79788456) + Float32(0.1070322243) * x * x) ) + Float32(0.5) * (Float32(1.0) + t) + + +# BF16 activation tables (see common/util/activation_table.cuh, which builds them; the layouts +# below must match it). With a table, BF16 activations follow the table specification: GeLU and +# SiLU of -inf give -0, and the derivatives below the table window are the linear term +# 0.5 + slope*x. Kernels either read the table ("table_*_2x") or compute what it holds +# ("activation_op"/"dactivation_op"). + +# Forward tables: rows of (op(+m), op(-m)) BF16 entries for magnitudes m from LOW_SENTINEL to +# HIGH_SENTINEL; the LOW_SENTINEL row holds -inf, the HIGH_SENTINEL row +inf and -0. +FORWARD_TABLE_WINDOWS = {"gelu": (0x3B4B, 0x40A5), "silu": (0x3BB6, 0x42B1)} +# Derivative tables: FP32 entries for magnitudes [LO, HI], positive inputs from entry 0 and +# negative ones from entry STRIDE; from SATURATED_END on (and for NaN) the result is NaN. +# Keyed by (activation, layout); layout "exact" is the exact dSiLU table. +DERIVATIVE_TABLE_WINDOWS = { + ("dgelu", "default"): (0x3980, 0x4100, 0x6044, 2048), + ("dsilu", "default"): (0x3B00, 0x42B2, 0x7F80, 2048), + ("dsilu", "exact"): (0x3300, 0x42B2, 0x7F80, 4096), +} +DERIVATIVE_SLOPES = {"dgelu": 0.79788456, "dsilu": 0.5} + + +def activation_table_bytes(activation: str, layout: str) -> int: + """Size of the table for activation in layout, rounded up to 16 bytes as TE allocates it.""" + if activation in FORWARD_TABLE_WINDOWS: + lo, hi = FORWARD_TABLE_WINDOWS[activation] + rows = (hi + 1) - (lo - 1) + 1 + return (rows * 2 * 2 + 15) // 16 * 16 + lo, hi, _, stride = DERIVATIVE_TABLE_WINDOWS[(activation, layout)] + return ((stride + hi - lo + 1) * 4 + 15) // 16 * 16 + + +@cute.jit +def _select_below(x: Float32, bound: Float32, below: Float32, other: Float32) -> Float32: + """below if |x| < bound, else other. Branch-free; NaN compares false.""" + return cute.arch.inline_ptx( + "{ .reg .pred p; .reg .f32 m; abs.f32 m, {$r0}; setp.lt.f32 p, m, {$r1}; " + "selp.f32 {$w0}, {$r2}, {$r3}, p; }", + write_only_types=[Float32], + read_only_args=[x, bound, below, other], + ) + + +def forward_table_spec(op): + """op as the forward tables hold it (see the comment above), computed.""" + + @cute.jit + def table_spec(x: Float32) -> Float32: + # Outside the window the table's closed forms equal util/math.h rounded to BF16, except + # at -inf, where the formula gives NaN and the table -0. util/math.h gives -0 for every + # finite x below -2^100, so clamp the input there; NaN fails the comparison. + clamped = cute.arch.inline_ptx( + "{ .reg .pred p; setp.lt.f32 p, {$r0}, {$r1}; selp.f32 {$w0}, {$r1}, {$r0}, p; }", + write_only_types=[Float32], + read_only_args=[x, Float32(-(2.0**100))], + ) + return op(clamped) + + return table_spec + + +def derivative_table_spec(dop, activation: str, layout: str): + """dop as the derivative table of activation in layout holds it, computed.""" + # Below the exact dSiLU table's window the linear term equals util/math.h. + if layout == "exact": + return dop + lo = DERIVATIVE_TABLE_WINDOWS[(activation, layout)][0] + window_low = struct.unpack(" Float32: + # Above the window util/math.h already gives the saturated value or NaN. + linear = fma_f32(Float32(slope), x, Float32(0.5)) + return _select_below(x, Float32(window_low), linear, dop(x)) + + return table_spec + + +@cute.jit +def table_forward_2x( + x: Int32, + sTable_u16: cute.Tensor, + LOW_SENTINEL: cutlass.Constexpr[int], + HIGH_SENTINEL: cutlass.Constexpr[int], +) -> Int32: + """Forward activation of a BF16 pair (low lane in the low half) from a forward table in + shared memory, as BF16 bits. Same instructions as forward_2x in util/activation_table.cuh.""" + # Clamp both magnitudes to the table rows. bf16x2 max/min ignore NaN, which lands on the + # -inf row; min.NaN with the upper bound below then makes it NaN again. + clamped = Uint32( + cute.arch.inline_ptx( + "{ .reg .b32 m, lo, hi; " + f"mov.b32 lo, {LOW_SENTINEL * 0x10001}; mov.b32 hi, {HIGH_SENTINEL * 0x10001}; " + "and.b32 m, {$r0}, 0x7fff7fff; max.bf16x2 m, m, lo; min.bf16x2 {$w0}, m, hi; }", + write_only_types=[Int32], + read_only_args=[x], + ) + ) + # Both lanes' entry indices in one word: 2 * row + sign per half. Each stays below 0x10000, + # so the halves never carry into each other. + indices = ((clamped - Uint32(LOW_SENTINEL * 0x10001)) << Uint32(1)) + ( + (Uint32(x) >> Uint32(15)) & Uint32(0x10001) + ) + entry_lo = Uint32(sTable_u16[Int32(indices & Uint32(0xFFFF))]) + entry_hi = Uint32(sTable_u16[Int32(indices >> Uint32(16))]) + entries = Int32(entry_lo | (entry_hi << Uint32(16))) + # max(entry, RN(x/2)) picks RN(x/2) below the window; min.NaN(., max.NaN(x, -0)) picks x + # above it for x > 0, -0 for x < 0, and NaN for NaN. + return cute.arch.inline_ptx( + "{ .reg .b32 half, upper, neg_zero, h; " + "mov.b32 half, 0x3f003f00; mov.b32 neg_zero, 0x80008000; " + "mul.rn.bf16x2 h, {$r0}, half; max.NaN.bf16x2 upper, {$r0}, neg_zero; " + "max.bf16x2 h, {$r1}, h; min.NaN.bf16x2 {$w0}, h, upper; }", + write_only_types=[Int32], + read_only_args=[x, entries], + ) + + +@cute.jit +def table_derivative_2x( + x: Int32, + sTable_f32: cute.Tensor, + LO: cutlass.Constexpr[int], + HI: cutlass.Constexpr[int], + SATURATED_END: cutlass.Constexpr[int], + STRIDE: cutlass.Constexpr[int], + SLOPE: cutlass.Constexpr[float], +): + """Derivative of a BF16 pair (low lane in the low half) from a derivative table in shared + memory, as (lo, hi) FP32. Same results as derivative_2x in util/activation_table.cuh.""" + magnitude = Uint32(x) & Uint32(0x7FFF7FFF) + clamped = Uint32( + cute.arch.inline_ptx( + "{ .reg .b32 lo, hi; " + f"mov.b32 lo, {LO * 0x10001}; mov.b32 hi, {HI * 0x10001}; " + "max.bf16x2 {$w0}, {$r0}, lo; min.bf16x2 {$w0}, {$w0}, hi; }", + write_only_types=[Int32], + read_only_args=[Int32(magnitude)], + ) + ) + bits = Uint32(x) + # Both lanes' entry indices in one word: the sign bits (15 and 31), shifted down to STRIDE + # (a power of two), select the negative half. Each index stays below 0x10000, so the halves + # never carry into each other. + SIGN_SHIFT = 15 - (STRIDE.bit_length() - 1) + indices = (clamped - Uint32(LO * 0x10001)) + ((bits ^ magnitude) >> Uint32(SIGN_SHIFT)) + entry_lo = sTable_f32[Int32(indices & Uint32(0xFFFF))] + entry_hi = sTable_f32[Int32(indices >> Uint32(16))] + # Below the window: the linear term 0.5 + SLOPE*x, rounded once. + linear_lo = fma_f32(Float32(SLOPE), (bits << Uint32(16)).bitcast(Float32), Float32(0.5)) + linear_hi = fma_f32(Float32(SLOPE), (bits & Uint32(0xFFFF0000)).bitcast(Float32), Float32(0.5)) + # One BF16x2 compare per condition sets both lanes' predicates. The unordered >= also holds + # for NaN, which gives NaN like the saturated end. + return cute.arch.inline_ptx( + "{ .reg .pred below_lo, below_hi, nan_lo, nan_hi; .reg .b32 lo, end; mov.b32 lo," + f" {LO * 0x10001}; mov.b32 end, {SATURATED_END * 0x10001}; setp.lt.bf16x2" + " below_lo|below_hi, {$r0}, lo; setp.geu.bf16x2 nan_lo|nan_hi, {$r0}, end; selp.f32" + " {$w0}, {$r1}, {$r2}, below_lo; selp.f32 {$w1}, {$r3}, {$r4}, below_hi;" + " selp.b32 {$w0}, 0x7fffffff, {$w0}, nan_lo; selp.b32 {$w1}, 0x7fffffff, {$w1}," + " nan_hi; }", + write_only_types=[Float32, Float32], + read_only_args=[Int32(magnitude), linear_lo, entry_lo, linear_hi, entry_hi], + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index 8a66b55e48f..08e8478a88a 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -30,6 +30,8 @@ dact_dsilu, dact_dqgelu, dact_dgelu, + derivative_table_spec, + forward_table_spec, ) CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" @@ -67,6 +69,25 @@ "dsrelu": dact_dsrelu, } +# Activations with a BF16 lookup table (see util/activation_table.cuh), and the table layouts +# the C++ dispatcher passes: "none", "default", or "exact" (the exact dSiLU table). +ACTIVATIONS_WITH_TABLE = ("gelu", "silu", "dgelu", "dsilu") +ACTIVATION_TABLE_LAYOUTS = ("none", "default", "exact") + + +def activation_op(activation: str, activation_table: str): + """The forward activation a kernel computes: with a table, what the table holds.""" + op = SUPPORTED_ACTIVATIONS[activation] + return op if activation_table == "none" else forward_table_spec(op) + + +def dactivation_op(activation: str, activation_table: str): + """The activation derivative a kernel computes: with a table, what the table holds.""" + dop = SUPPORTED_DACTIVATIONS[activation] + if activation_table == "none": + return dop + return derivative_table_spec(dop, activation, activation_table) + @cute.jit def derive_swizzled_scale_layout( @@ -138,6 +159,7 @@ def __init__( with_act: bool = False, use_2d_quantization: bool = False, activation: Optional[str] = None, + activation_table: str = "none", ): if use_2d_quantization: raise ValueError("2D block scaling is not implemented by the CuTeDSL MXFP8 kernels") @@ -188,6 +210,22 @@ def __init__( ) self.WITH_DACT = with_dact self.WITH_ACT = with_act + # With a table, BF16 activations follow the table specification: the register-resident + # kernel reads the table, the others compute what it holds. + if activation_table not in ACTIVATION_TABLE_LAYOUTS: + raise ValueError( + f"unknown activation table {activation_table!r}; expected one of" + f" {ACTIVATION_TABLE_LAYOUTS}" + ) + if activation_table != "none" and ( + dtype != "BFloat16" + or self.ACTIVATION not in ACTIVATIONS_WITH_TABLE + or (activation_table == "exact" and self.ACTIVATION != "dsilu") + ): + raise ValueError( + f"no {activation_table!r} activation table for {self.ACTIVATION!r} with {dtype}" + ) + self.ACTIVATION_TABLE = activation_table # dbias is the column reduction of the (post-act/dact) element. With colwise # output each thread owns a full column (trivial reduction); rowwise-only # uses a cross-thread smem reduction over THREADS_Y. Both mirror the CUDA @@ -203,7 +241,8 @@ def __str__(self): f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, with_amax={self.WITH_AMAX}, " f"with_dbias={self.WITH_DBIAS}, with_dact={self.WITH_DACT}, " - f"with_act={self.WITH_ACT}, activation={self.ACTIVATION})" + f"with_act={self.WITH_ACT}, activation={self.ACTIVATION}, " + f"activation_table={self.ACTIVATION_TABLE})" ) __repr__ = __str__ @@ -247,6 +286,7 @@ def __call__( mNoop: cute.Pointer, mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], + mActTable: cute.Pointer, # activation table; null without one stream: CUstream, ): """ @@ -318,6 +358,7 @@ def __call__( mNoop: cute.Pointer, mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], + mActTable: cute.Pointer, # activation table; null without one stream: CUstream, ): dispatched_to_specialized = False @@ -337,6 +378,7 @@ def __call__( mNoop, mDActInput, mWorkspace, + mActTable, stream, ) # If not using a specialized kernel, fall back to the general kernel @@ -368,6 +410,7 @@ def __call__( mNoop, mDActInput, mWorkspace, + mActTable, stream, ) else: @@ -385,6 +428,7 @@ def __call__( mNoop, mDActInput, mWorkspace, + mActTable, stream, ) else: @@ -400,6 +444,7 @@ def __call__( mNoop, mDActInput, mWorkspace, + mActTable, stream, ) @@ -498,6 +543,10 @@ def compile_cutedsl_function_from_cfg(cfg): if cfg.WITH_DACT else None ) + # Like the noop flag, the activation table is an always-present pointer, null without a table. + act_table_fake = cute.runtime.nullptr( + cutlass.Uint8, mem_space=cute.AddressSpace.gmem, assumed_align=16 + ) workspace_fake = ( cute.runtime.make_fake_compact_tensor( Float32, ws_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4 @@ -517,6 +566,7 @@ def compile_cutedsl_function_from_cfg(cfg): noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) act_input_fake, # mDActInput (activation input, dact only) workspace_fake, # mWorkspace(backward slot, unused) + act_table_fake, # mActTable (activation table, may be null at launch) cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi # "handle" arg; C++ passes the CUDA stream # as void*) @@ -538,6 +588,7 @@ def get_mxfp8_quantization_function( with_act: bool, use_2d_quantization: bool, activation: str, + activation_table: str, ) -> bool: """Compile the MXFP8 quantize kernel for this config and register it in the TVM-FFI global registry under EXACTLY `fn_name` (the key the C++ dispatcher built; Python treats it as an opaque name). @@ -572,6 +623,7 @@ def get_mxfp8_quantization_function( with_act=with_act, use_2d_quantization=use_2d_quantization, activation=activation, + activation_table=activation_table, ) except ValueError as e: # The exception message states exactly why the config is unsupported diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 61598a7d37b..8cc2eaeddf0 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -8,7 +8,8 @@ import cutlass from cutlass import cute -from cutlass import Boolean, Int32, Int64, Uint32, Uint8 +from cutlass.utils import HardwareInfo +from cutlass import Boolean, Float32, Int32, Int64, Uint16, Uint32, Uint8 from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module from transformer_engine.common.CuTeDSL.utils import ( @@ -23,13 +24,21 @@ x2_lo_to_f32_bf16, cvt_f32x2_to_bf16x2, ) +from transformer_engine.common.CuTeDSL.activations import ( + DERIVATIVE_SLOPES, + DERIVATIVE_TABLE_WINDOWS, + FORWARD_TABLE_WINDOWS, + activation_table_bytes, + table_derivative_2x, + table_forward_2x, +) from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 from .quantize_mxfp8_common import ( MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, - SUPPORTED_ACTIVATIONS, - SUPPORTED_DACTIVATIONS, + activation_op, + dactivation_op, derive_swizzled_scale_layout, ) @@ -403,6 +412,8 @@ def quantize_mxfp8_gmem_tile_bidimensional( ROWWISE: cutlass.Constexpr[bool], COLWISE: cutlass.Constexpr[bool], ACTIVATION: cutlass.Constexpr[Optional[str]], + ACTIVATION_TABLE: cutlass.Constexpr[str], + sTable: Optional[cute.Tensor], WITH_ACT: cutlass.Constexpr[bool], WITH_DACT: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], @@ -497,17 +508,44 @@ def quantize_mxfp8_gmem_tile_bidimensional( rX_i32 = cute.recast_tensor(rX[None, row_idx], Int32) rA_i32 = cute.recast_tensor(rA[None, row_idx], Int32) for i in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): - lo = x2_lo_to_f32_bf16(rX_i32[i]) - hi = x2_hi_to_f32_bf16(rX_i32[i]) - if cutlass.const_expr(WITH_DACT): - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] - lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) - hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + if cutlass.const_expr(ACTIVATION_TABLE != "none"): + # Read the activation from the table, as the CUDA register-tiled kernel does. + if cutlass.const_expr(WITH_DACT): + low, high, end, stride = DERIVATIVE_TABLE_WINDOWS[ + (ACTIVATION, ACTIVATION_TABLE) + ] + d_lo, d_hi = table_derivative_2x( + rA_i32[i], + cute.recast_tensor(sTable, Float32), + LO=low, + HI=high, + SATURATED_END=end, + STRIDE=stride, + SLOPE=DERIVATIVE_SLOPES[ACTIVATION], + ) + rX_i32[i] = cvt_f32x2_to_bf16x2( + x2_lo_to_f32_bf16(rX_i32[i]) * d_lo, x2_hi_to_f32_bf16(rX_i32[i]) * d_hi + ) + else: + low, high = FORWARD_TABLE_WINDOWS[ACTIVATION] + rX_i32[i] = table_forward_2x( + rX_i32[i], + cute.recast_tensor(sTable, Uint16), + LOW_SENTINEL=low - 1, + HIGH_SENTINEL=high + 1, + ) else: - op = SUPPORTED_ACTIVATIONS[ACTIVATION] - lo = op(lo) - hi = op(hi) - rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + lo = x2_lo_to_f32_bf16(rX_i32[i]) + hi = x2_hi_to_f32_bf16(rX_i32[i]) + if cutlass.const_expr(WITH_DACT): + dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) + lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) + hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + else: + op = activation_op(ACTIVATION, ACTIVATION_TABLE) + lo = op(lo) + hi = op(hi) + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) if cutlass.const_expr(ROWWISE): gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] @@ -604,6 +642,9 @@ class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): _CTA_THREADS_Y = 8 _CTA_THREADS_X = 32 + # With a table, the grid holds about this many columns' worth of CTAs (16384 CTAs of 256 + # columns, as in the CUDA register-tiled kernel); each CTA walks the remaining row tiles. + _WALK_TARGET_COLS = 16384 * 256 @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: @@ -647,6 +688,14 @@ def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg # One CTA only quantize one columnwise MXFP8 block at a time self._ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // self._CTA_THREADS_Y # 4 + # SMs of the device compiled for, which sets the table path's walk length. + self._sm_count = 0 + if cfg.ACTIVATION_TABLE != "none": + from cuda.bindings import runtime as cudart # pylint: disable=no-name-in-module + + err, device_id = cudart.cudaGetDevice() + assert err == cudart.cudaError_t.cudaSuccess, f"CUDA runtime error {err}" + self._sm_count = HardwareInfo(device_id).get_device_multiprocessor_count() @cute.jit def __call__( @@ -660,12 +709,28 @@ def __call__( mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], + mActTable: cute.Pointer, # activation table; null without one stream: CUstream, ): M, N = mX.shape TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD WIDE_TILE_COLS = self._CTA_THREADS_X * 16 - if N % WIDE_TILE_COLS == 0: + # A derivative read from a table holds the gradient and the activation input in registers: + # 256-column tiles keep it under the register budget, and clusters do not help it. + if cutlass.const_expr(self.cfg.WITH_DACT and self.cfg.ACTIVATION_TABLE != "none"): + self._launch( + mX, + mDActInput, + mO_row, + mS_row, + mO_col, + mS_col, + mActTable, + stream, + ELEMENTS_X_PER_THREAD=8, + CLUSTER_WIDTH=1, + ) + elif N % WIDE_TILE_COLS == 0: if Int64(M // TILE_ROWS) * (N // WIDE_TILE_COLS) >= 4096: self._launch( mX, @@ -674,6 +739,7 @@ def __call__( mS_row, mO_col, mS_col, + mActTable, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=4, @@ -686,6 +752,7 @@ def __call__( mS_row, mO_col, mS_col, + mActTable, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=1, @@ -698,6 +765,7 @@ def __call__( mS_row, mO_col, mS_col, + mActTable, stream, ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=8, @@ -712,6 +780,7 @@ def _launch( mS_row: cute.Tensor, mO_col: cute.Tensor, mS_col: cute.Tensor, + mActTable: cute.Pointer, stream: CUstream, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], CLUSTER_WIDTH: cutlass.Constexpr[int], @@ -734,8 +803,24 @@ def _launch( # Adjust cluster size so that the CTAs in the column direction can be evenly divided into clusters. while cluster_cols > 1 and grid_cols % cluster_cols != 0: cluster_cols -= 1 + # With a table, each CTA walks several row tiles so that it copies the table once for + # all of them, as the CUDA register-tiled kernel does: halve the grid while it holds + # more than _WALK_TARGET_COLS / TILE_COLS CTAs and the row tiles divide evenly, and + # walk at least two tiles unless that leaves fewer than two CTAs per SM. + walk = Int32(1) + if cutlass.const_expr(self.cfg.ACTIVATION_TABLE != "none"): + while (grid_rows // walk) % 2 == 0 and grid_cols * (grid_rows // walk) > ( + self._WALK_TARGET_COLS // TILE_COLS + ): + walk *= 2 + if ( + walk == 1 + and grid_rows % 2 == 0 + and grid_cols * (grid_rows // 2) >= 2 * self._sm_count + ): + walk = Int32(2) # Keep column CTAs on grid.x so clusters read adjacent column tiles. - grid = [grid_cols, grid_rows] + grid = [grid_cols, grid_rows // walk] block = [THREADS_PER_CTA] cluster = [cluster_cols, 1, 1] @@ -778,6 +863,8 @@ def _launch( gS_row, gO_col, gS_col, + mActTable, + walk, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) elif N == 4096: @@ -818,6 +905,8 @@ def _launch( gS_row, gO_col, gS_col, + mActTable, + walk, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: @@ -834,6 +923,8 @@ def _launch( gS_row, mO_col, gS_col, + mActTable, + walk, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) @@ -846,6 +937,8 @@ def kernel( mS_row: cute.Tensor, mO_col: cute.Tensor, mS_col: cute.Tensor, + mActTable: cute.Pointer, + walk: Int32, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): """Run the kernel body. Like the CUDA specialized cast kernels it does not consume the @@ -857,9 +950,33 @@ def kernel( mS_row, mO_col, mS_col, + mActTable, + walk, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ) + @cute.jit + def _data_tile( + self, + mT: cute.Tensor, + row_tile: Int32, + cta_col: Int32, + TILER: cutlass.Constexpr, + DYNAMIC_STRIDE: cutlass.Constexpr[bool], + ALIGNMENT: cutlass.Constexpr[int], + ) -> cute.Tensor: + """The (row_tile, cta_col) tile of mT. With a run-time row stride, the tile's base takes the + 64-bit product once and offsets within the tile are 32-bit, as in the CUDA register-tiled + kernel: the full 64-bit layout costs registers and instructions on every access.""" + if cutlass.const_expr(DYNAMIC_STRIDE): + N = mT.shape[1] + base = Int64(row_tile) * Int64(TILER[0]) * Int64(N) + Int64(cta_col) * Int64(TILER[1]) + return cute.make_tensor( + (mT.iterator + base).align(ALIGNMENT), + cute.make_layout(TILER, stride=(Int32(N), 1)), + ) + return cute.local_tile(mT, TILER, (row_tile, cta_col)) + @cute.jit def _kernel_main( self, @@ -869,33 +986,109 @@ def _kernel_main( mS_row: cute.Tensor, mO_col: cute.Tensor, mS_col: cute.Tensor, + mActTable: cute.Pointer, + walk: Int32, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): - """Construct CTA tile views and run bidimensional quantization.""" - cta_col, cta_row, _ = cute.arch.block_idx() - cta_coord = (cta_row, cta_col) + """Run bidimensional quantization on the CTA's row tiles: blockIdx.y, and with a table + every gridDim.y-th row tile after it, walk tiles in all.""" + cta_col, first_cta_row, _ = cute.arch.block_idx() + _, grid_rows, _ = cute.arch.grid_dim() + + cache_policy = create_l2_policy(True) + + # With a table, copy it into shared memory once for all of the CTA's row tiles. + sTable = None + if cutlass.const_expr(self.cfg.ACTIVATION_TABLE != "none"): + TABLE_WORDS = ( + activation_table_bytes(self.cfg.ACTIVATION, self.cfg.ACTIVATION_TABLE) // 4 + ) + tidx, _, _ = cute.arch.thread_idx() + allocator = cutlass.memory.SmemAllocator() + sTable = allocator.allocate_tensor( + Uint32, cute.make_layout((TABLE_WORDS,)), byte_alignment=16 + ) + # 16-byte pieces: TE allocates the table 16-byte aligned and padded. + piece_layout = cute.make_layout((4, TABLE_WORDS // 4), stride=(1, 4)) + gTable = cute.make_tensor( + cute.recast_ptr(mActTable, dtype=Uint32).align(16), piece_layout + ) + sTable_pieces = cute.make_tensor(sTable.iterator, piece_layout) + for piece in cutlass.range( + tidx, TABLE_WORDS // 4, self._CTA_THREADS_Y * self._CTA_THREADS_X, unroll=1 + ): + cute.autovec_copy(gTable[None, piece], sTable_pieces[None, piece]) + cute.arch.sync_threads() + + if cutlass.const_expr(self.cfg.ACTIVATION_TABLE == "none"): + self._quantize_row_tile( + mX, + mA, + mO_row, + mS_row, + mO_col, + mS_col, + sTable, + first_cta_row, + cta_col, + cache_policy, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ) + else: + for step in cutlass.range(walk, unroll=1): + self._quantize_row_tile( + mX, + mA, + mO_row, + mS_row, + mO_col, + mS_col, + sTable, + step * grid_rows + first_cta_row, + cta_col, + cache_policy, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ) + @cute.jit + def _quantize_row_tile( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO_row: cute.Tensor, + mS_row: cute.Tensor, + mO_col: cute.Tensor, + mS_col: cute.Tensor, + sTable: Optional[cute.Tensor], + row_tile: Int32, + cta_col: Int32, + cache_policy: Int64, + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + ): + """Quantize the (row_tile, cta_col) tile.""" TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD - TILER = (TILE_ROWS, TILE_COLS) ROW_SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) - - gX_tile = cute.local_tile(mX, TILER, cta_coord) + # With a table and a row stride known only at run time, index each data tile from its own + # base address with 32-bit offsets (see _data_tile). + DYNAMIC_STRIDE = self.cfg.ACTIVATION_TABLE != "none" and not cute.is_static( + mX.layout.stride[0] + ) + cta_coord = (row_tile, cta_col) + gX_tile = self._data_tile(mX, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 32) gA_tile = None if cutlass.const_expr(self.cfg.WITH_DACT): - gA_tile = cute.local_tile(mA, TILER, cta_coord) + gA_tile = self._data_tile(mA, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 32) gO_row_tile = None gS_row_tile = None if cutlass.const_expr(self.cfg.ROWWISE): - gO_row_tile = cute.local_tile(mO_row, TILER, cta_coord) + gO_row_tile = self._data_tile(mO_row, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 16) gS_row_tile = cute.local_tile(mS_row, ROW_SCALE_TILER, cta_coord) - gO_col_tile = cute.local_tile(mO_col, TILER, cta_coord) + gO_col_tile = self._data_tile(mO_col, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 16) gS_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) - cache_policy = create_l2_policy(True) - quantize_mxfp8_gmem_tile_bidimensional( gX_tile, gA_tile, @@ -906,6 +1099,8 @@ def _kernel_main( ROWWISE=self.cfg.ROWWISE, COLWISE=self.cfg.COLWISE, ACTIVATION=self.cfg.ACTIVATION, + ACTIVATION_TABLE=self.cfg.ACTIVATION_TABLE, + sTable=sTable, WITH_ACT=self.cfg.WITH_ACT, WITH_DACT=self.cfg.WITH_DACT, DTYPE=self.cfg.DTYPE, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py index 0c287f02abe..0d2532954a3 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -261,6 +261,7 @@ def __call__( mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], # Unused, kept for API compatibility mWorkspace: Optional[cute.Tensor], # Unused, kept for API compatibility + mActTable: cute.Pointer, # activation table; null without one stream: CUstream, ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index 03d5d510ab8..62b4fe61f90 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -28,8 +28,8 @@ from .quantize_mxfp8_common import ( MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, - SUPPORTED_ACTIVATIONS, - SUPPORTED_DACTIVATIONS, + activation_op, + dactivation_op, derive_swizzled_scale_layout, ) @@ -73,6 +73,7 @@ def quantize_mxfp8_gmem_tile_rowwise_2lane( mS_tiles: cute.Tensor, N: Int32, ACTIVATION: cutlass.Constexpr[Optional[str]], + ACTIVATION_TABLE: cutlass.Constexpr[str], WITH_ACT: cutlass.Constexpr[bool], WITH_DACT: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], @@ -183,11 +184,11 @@ def quantize_mxfp8_gmem_tile_rowwise_2lane( lo = x2_lo_to_f32_bf16(rX_i32[i]) hi = x2_hi_to_f32_bf16(rX_i32[i]) if cutlass.const_expr(WITH_DACT): - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) else: - op = SUPPORTED_ACTIVATIONS[ACTIVATION] + op = activation_op(ACTIVATION, ACTIVATION_TABLE) lo = op(lo) hi = op(hi) rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) @@ -314,6 +315,7 @@ def __call__( mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], + mActTable: cute.Pointer, # activation table; null without one stream: CUstream, ): M, N = mX.shape @@ -658,6 +660,7 @@ def _kernel_main( mS_tiles, N, ACTIVATION=self.cfg.ACTIVATION, + ACTIVATION_TABLE=self.cfg.ACTIVATION_TABLE, WITH_ACT=self.cfg.WITH_ACT, WITH_DACT=self.cfg.WITH_DACT, DTYPE=self.cfg.DTYPE, diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py index 5f0b3d2aed5..02c35227053 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py @@ -297,6 +297,7 @@ def __call__( mNoop: cute.Pointer, # Unused: the dispatcher routes noop calls to the general kernel mDActInput: Optional[cute.Tensor], mWorkspace: Optional[cute.Tensor], + mActTable: cute.Pointer, # activation table; null without one stream: CUstream, ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py index 81148416fe9..325accdd41c 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -66,8 +66,8 @@ MXFP8QuantizeConfig, MXFP8QuantizeKernelBase, MXFP8_BLOCK_SCALING_SIZE, - SUPPORTED_ACTIVATIONS, - SUPPORTED_DACTIVATIONS, + activation_op, + dactivation_op, SYM_N_DIVISIBILITY, THREADS_PER_WARP, derive_swizzled_scale_layout, @@ -87,6 +87,7 @@ def _quantize_rowwise_mxfp8( M: Int32, N: Int32, ACTIVATION: cutlass.Constexpr[str | None], + ACTIVATION_TABLE: cutlass.Constexpr[str], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], TILE_X: cutlass.Constexpr[int], @@ -210,14 +211,14 @@ def _quantize_rowwise_mxfp8( if cutlass.const_expr(WITH_DACT): # Backward: out = grad · act'(act_input). sX is grad, sA is act_input. - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) sA_thread = cute.composition(sActInput_tile, tv_layout)[tidx, None] sA_thread_rw = cute.make_tensor( sA_thread.iterator, cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), ) elif cutlass.const_expr(WITH_ACT): - op = SUPPORTED_ACTIVATIONS[ACTIVATION] + op = activation_op(ACTIVATION, ACTIVATION_TABLE) if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 @@ -341,6 +342,7 @@ def _quantize_colwise_mxfp8( M: Int32, N: Int32, ACTIVATION: cutlass.Constexpr[str | None], + ACTIVATION_TABLE: cutlass.Constexpr[str], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], SWIZZLE: cutlass.Constexpr[bool], @@ -412,12 +414,12 @@ def _quantize_colwise_mxfp8( rX_thread_f32[i] = Float32(sX_thread[i]) # Apply activation (fwd) or grad·act'(act_input) (bwd dact) in f32. if cutlass.const_expr(WITH_DACT): - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) sA_thread = cute.composition(sActInput_tile, tv_layout)[tidx, None] for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): rX_thread_f32[i] = rX_thread_f32[i] * dop(Float32(sA_thread[i])) elif cutlass.const_expr(WITH_ACT): - op = SUPPORTED_ACTIVATIONS[ACTIVATION] + op = activation_op(ACTIVATION, ACTIVATION_TABLE) # Don't repeat these OOB masking computation in the loop thread_col_start = tile_col_start + tidx thread_col_oob = thread_col_start >= N @@ -515,6 +517,7 @@ def _dbias_reduction_colwise( sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16/fp32 smem view, post-TMA sA_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) activation-input smem tile (dact only) ACTIVATION: cutlass.Constexpr[str | None], + ACTIVATION_TABLE: cutlass.Constexpr[str], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], TILE_X: cutlass.Constexpr[int], WITH_DACT: cutlass.Constexpr[bool] = False, @@ -532,7 +535,7 @@ def _dbias_reduction_colwise( dbias_partial = Float32(0.0) if cutlass.const_expr(WITH_DACT): - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] + dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) sA_thread = cute.composition(sA_tile, tv_layout)[tidx, None] for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): x = Float32(sX_thread[i]) * dop(Float32(sA_thread[i])) @@ -564,6 +567,7 @@ def quantize_mxfp8_smem_tile_general( WITH_ACT: cutlass.Constexpr[bool], WITH_DACT: cutlass.Constexpr[bool], ACTIVATION: cutlass.Constexpr[str | None], + ACTIVATION_TABLE: cutlass.Constexpr[str], MAX_NORM_RCP: cutlass.Constexpr[float], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], @@ -597,6 +601,7 @@ def quantize_mxfp8_smem_tile_general( M, N, ACTIVATION=ACTIVATION, + ACTIVATION_TABLE=ACTIVATION_TABLE, DTYPE=DTYPE, FP8_DTYPE=FP8_DTYPE, SWIZZLE=WITH_GEMM_SWIZZLED_SCALES, @@ -621,6 +626,7 @@ def quantize_mxfp8_smem_tile_general( sX_tile, sActInput_tile, ACTIVATION=ACTIVATION, + ACTIVATION_TABLE=ACTIVATION_TABLE, DTYPE=DTYPE, TILE_X=TILE_X, WITH_DACT=WITH_DACT, @@ -642,6 +648,7 @@ def quantize_mxfp8_smem_tile_general( M, N, ACTIVATION=None if CACHE_ACTIVATION else ACTIVATION, + ACTIVATION_TABLE=ACTIVATION_TABLE, DTYPE=DTYPE, FP8_DTYPE=FP8_DTYPE, TILE_X=TILE_X, @@ -742,6 +749,7 @@ def __call__( mWorkspace: Optional[ cute.Tensor ], # Workspace for the dbias reduction, only used when WITH_DBIAS is True + mActTable: cute.Pointer, # activation table; null without one stream: CUstream, ): if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): @@ -1235,6 +1243,7 @@ class DactStorage: WITH_ACT=self.cfg.WITH_ACT, WITH_DACT=self.cfg.WITH_DACT, ACTIVATION=self.cfg.ACTIVATION, + ACTIVATION_TABLE=self.cfg.ACTIVATION_TABLE, MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, DTYPE=self.cfg.DTYPE, FP8_DTYPE=self.cfg.FP8_DTYPE, diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 80de7098ba5..999ce51f770 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -107,19 +107,19 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const Tensor *dummy_input_tensor = nullptr; Tensor *dummy_dbias_tensor = nullptr; Tensor *dummy_workspace_tensor = nullptr; + const activation_table::TableRef activation_lut = + IS_ACT ? activation_lut_for(input_tensor->dtype()) + : activation_table::TableRef{}; bool quantized_with_cutedsl = false; #ifdef NVTE_WITH_CUTEDSL - // The CuTeDSL backend, when enabled, computes activations without the tables. quantized_with_cutedsl = cutedsl_backend::mxfp8_quantize_cutedsl( input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, - dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); + dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, activation_lut, + stream); #endif if (!quantized_with_cutedsl) { - const activation_table::TableRef activation_lut = - IS_ACT ? activation_lut_for(input_tensor->dtype()) - : activation_table::TableRef{}; mxfp8::quantize( *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, activation_lut, stream); @@ -300,19 +300,18 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens break; } case NVTE_MXFP8_1D_SCALING: { + activation_table::TableRef activation_lut; + if constexpr (IS_DACT) { + activation_lut = activation_lut_for(input_tensor->dtype()); + } bool quantized_with_cutedsl = false; #ifdef NVTE_WITH_CUTEDSL - // The CuTeDSL backend, when enabled, computes activations without the tables. quantized_with_cutedsl = cutedsl_backend::mxfp8_quantize_cutedsl( grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, - quant_config_cpp.mxfp8_2d_quantization, stream); + quant_config_cpp.mxfp8_2d_quantization, activation_lut, stream); #endif if (!quantized_with_cutedsl) { - activation_table::TableRef activation_lut; - if constexpr (IS_DACT) { - activation_lut = activation_lut_for(input_tensor->dtype()); - } mxfp8::quantize( *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, activation_lut, stream); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh index 6380dda21da..458f3896039 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh @@ -17,6 +17,7 @@ #include "../../common.h" #include "../../tvm_ffi_bridge.h" +#include "../../util/activation_table.cuh" #include "../../util/cutedsl_utils.h" #include "../core/common.cuh" // dispatch::common::reduce_dbias #include "quantize_mxfp8.cuh" // dispatch::mxfp8::quantize_kernel::zero_scales_kernel @@ -41,16 +42,21 @@ struct MXFP8QuantConfig { bool with_act = false; // If an activation operation is fused bool use_2d_quantization = false; // If use 2D quantization Activation activation = Activation::kNone; + // The activation table layout the kernel reads (see util/activation_table.cuh): kNoTable, + // kDefaultTable or kExactDsiluTable. + activation_table::TableUse activation_table = activation_table::kNoTable; uint32_t sm_arch = static_cast(cuda::sm_arch()); // Bit layout: dtype [3:0] (4 used/reserved), fp8_dtype [7:4] (4 used/reserved), - // flags [15:8] (8 used/reserved), activation [21:16] (6 used/reserved), and - // SM architecture [30:22] (9 used/reserved). Bit 31 is unused. + // flags [15:8] (8 used/reserved), activation [19:16] (4 used/reserved), activation table + // [21:20] (2 used/reserved) and SM architecture [30:22] (9 used/reserved). Bit 31 is unused. uint32_t to_id() const { static_assert(static_cast(DType::kNumTypes) <= 16, "DType no longer fits in the 4 bits to_id() gives it."); - static_assert(static_cast(Activation::kNumTypes) <= 64, - "Activation no longer fits in the 6 bits to_id() gives it."); + static_assert(static_cast(Activation::kNumTypes) <= 16, + "Activation no longer fits in the 4 bits to_id() gives it."); + NVTE_CHECK(static_cast(activation_table) < 4, + "Activation table use no longer fits in the 2 bits to_id() gives it."); NVTE_CHECK(sm_arch < 512, "SM architecture no longer fits in the 9 bits to_id() gives it."); return static_cast(dtype) | (static_cast(fp8_dtype) << 4) | (static_cast(rowwise) << 8) | (static_cast(colwise) << 9) | @@ -58,7 +64,20 @@ struct MXFP8QuantConfig { (static_cast(with_dbias) << 12) | (static_cast(with_dact) << 13) | (static_cast(with_act) << 14) | (static_cast(use_2d_quantization) << 15) | - (static_cast(activation) << 16) | (sm_arch << 22); + (static_cast(activation) << 16) | + (static_cast(activation_table) << 20) | (sm_arch << 22); + } + + //! The activation table layout as the CuTeDSL config names it. + const char *activation_table_str() const { + switch (activation_table) { + case activation_table::kDefaultTable: + return "default"; + case activation_table::kExactDsiluTable: + return "exact"; + default: + return "none"; + } } std::optional get_kernel() const { @@ -95,6 +114,10 @@ struct MXFP8QuantConfig { .append(use_2d_quantization ? "1" : "0") .append("_") .append(activation_to_str(activation)); + // Kernels that read a table carry its layout in the name; the rest keep their names. + if (activation_table != activation_table::kNoTable) { + key.append("_table_").append(activation_table_str()); + } return key; } @@ -106,17 +129,19 @@ struct MXFP8QuantConfig { tvm::ffi::Any result = (*entrypoint)( tvm::ffi::String(fn_name), tvm::ffi::String(to_string(dtype)), tvm::ffi::String(to_string(fp8_dtype)), rowwise, colwise, swizzled, with_amax, with_dbias, - with_dact, with_act, use_2d_quantization, tvm::ffi::String(activation_to_str(activation))); + with_dact, with_act, use_2d_quantization, tvm::ffi::String(activation_to_str(activation)), + tvm::ffi::String(activation_table_str())); return result.try_cast().value_or(false); } }; // Signature mirrors mxfp8::quantize (input, act_input, noop, output, dbias, -// workspace, stream). Returns false to fall back to the CUDA kernel. +// workspace, activation table, stream). Returns false to fall back to the CUDA kernel. inline bool mxfp8_quantize_cutedsl(const MXFP8QuantConfig &config, const Tensor *input_tensor, const Tensor *act_input_tensor, const Tensor *noop_tensor, Tensor *output_tensor, Tensor *dbias_tensor, - Tensor *workspace_tensor, cudaStream_t stream) { + Tensor *workspace_tensor, const void *activation_table_ptr, + cudaStream_t stream) { const size_t flat_m = input_tensor->flat_first_dim(); const size_t flat_n = input_tensor->flat_last_dim(); if (flat_n % 16 != 0) { @@ -206,9 +231,10 @@ inline bool mxfp8_quantize_cutedsl(const MXFP8QuantConfig &config, const Tensor mActInput = tvm_ffi_bridge::DLTensorWrapper(act_input_tensor->data, true, device_index); if (workspace_tensor != nullptr && workspace_tensor->data.dptr != nullptr) mWorkspace = tvm_ffi_bridge::DLTensorWrapper(workspace_tensor->data, true, device_index); - // noop and stream are tvm-ffi opaque "handles"; pass them as void*. + // noop, the activation table and stream are tvm-ffi opaque "handles"; pass them as void*. (*mxfp8_quant_func_opt)(&mX, &mO_row, &mS_row, &mO_col, &mS_col, &mAmax, noop_ptr, &mActInput, - &mWorkspace, static_cast(stream)); + &mWorkspace, const_cast(activation_table_ptr), + static_cast(stream)); // If WITH_DBIAS, reduce the workspace partial dbias in CUDA C++ for now. // This will not be affected by the noop flag, which is aligned with the CUDA C++ implementation @@ -227,7 +253,7 @@ template ()}; + /*activation=*/activation_func_to_enum(), + /*activation_table=*/activation_lut.use}; // Sanity checks, mirroring mxfp8::quantize in quantize_mxfp8.cuh if (config.rowwise) { NVTE_CHECK(output_tensor->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); @@ -268,7 +295,7 @@ bool mxfp8_quantize_cutedsl(const Tensor *input_tensor, const Tensor *act_input_ NVTE_CHECK(workspace_tensor != nullptr, "Workspace must be a tensor."); } return mxfp8_quantize_cutedsl(config, input_tensor, act_input_tensor, noop_tensor, output_tensor, - dbias_tensor, workspace_tensor, stream); + dbias_tensor, workspace_tensor, activation_lut.ptr, stream); } } // namespace cutedsl_backend From 478a926f428c04ca6b46ad40c5650ce21b21e881 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 27 Oct 2026 12:07:32 -0700 Subject: [PATCH 20/29] [Common] Fuse dbias into the register-resident CuTeDSL MXFP8 kernel dGeLU and dSiLU with dbias, BF16 input, both directions or colwise only and row counts divisible by 64 now take the register-resident bidimensional kernel instead of the general one, with or without activation tables. dbias follows the CUDA kernels' order: each CTA takes the two 32-row tiles of a 64-row chunk back to back, stages the FP32 products in shared memory, and one thread per column adds them in row order, writing one workspace row per chunk for the C++ reduction. The general kernel adds each chunk as two 32-row partials, which differs from CUDA in the last bit; this path matches it. As with tables, the dbias path walks row tiles, uses 256-column tiles, and indexes tiles of a run-time row stride with 32-bit offsets. GB200, BF16, geometric mean speedup over TE main, tables off / on: dbias+dGeLU 1.69 / 1.87 (general kernel: 1.74 / -), dbias+dSiLU 1.78 / 1.90 (general: 1.42 / -); CUDA backend 1.80 / 2.01 and 1.82 / 2.04. Signed-off-by: Teddy Do --- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 105 +++++++++++++++--- 1 file changed, 92 insertions(+), 13 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 8cc2eaeddf0..5c8cfd10cf3 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -414,6 +414,7 @@ def quantize_mxfp8_gmem_tile_bidimensional( ACTIVATION: cutlass.Constexpr[Optional[str]], ACTIVATION_TABLE: cutlass.Constexpr[str], sTable: Optional[cute.Tensor], + sDbias: Optional[cute.Tensor], WITH_ACT: cutlass.Constexpr[bool], WITH_DACT: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], @@ -503,10 +504,15 @@ def quantize_mxfp8_gmem_tile_bidimensional( rA[None, row_idx], cache_policy=cache_policy, ) + # With dbias, the FP32 products are staged for the column sums (see _kernel_main). + if cutlass.const_expr(sDbias is not None): + sDbias_thread = cute.composition(sDbias, tv_layout)[tidx, None] # Apply activation row by row for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_i32 = cute.recast_tensor(rX[None, row_idx], Int32) rA_i32 = cute.recast_tensor(rA[None, row_idx], Int32) + if cutlass.const_expr(sDbias is not None): + rProducts = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), Float32) for i in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): if cutlass.const_expr(ACTIVATION_TABLE != "none"): # Read the activation from the table, as the CUDA register-tiled kernel does. @@ -523,9 +529,12 @@ def quantize_mxfp8_gmem_tile_bidimensional( STRIDE=stride, SLOPE=DERIVATIVE_SLOPES[ACTIVATION], ) - rX_i32[i] = cvt_f32x2_to_bf16x2( - x2_lo_to_f32_bf16(rX_i32[i]) * d_lo, x2_hi_to_f32_bf16(rX_i32[i]) * d_hi - ) + lo = x2_lo_to_f32_bf16(rX_i32[i]) * d_lo + hi = x2_hi_to_f32_bf16(rX_i32[i]) * d_hi + if cutlass.const_expr(sDbias is not None): + rProducts[2 * i] = lo + rProducts[2 * i + 1] = hi + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) else: low, high = FORWARD_TABLE_WINDOWS[ACTIVATION] rX_i32[i] = table_forward_2x( @@ -541,11 +550,16 @@ def quantize_mxfp8_gmem_tile_bidimensional( dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + if cutlass.const_expr(sDbias is not None): + rProducts[2 * i] = lo + rProducts[2 * i + 1] = hi else: op = activation_op(ACTIVATION, ACTIVATION_TABLE) lo = op(lo) hi = op(hi) rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + if cutlass.const_expr(sDbias is not None): + cute.autovec_copy(rProducts, sDbias_thread[(None, row_idx),]) if cutlass.const_expr(ROWWISE): gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] @@ -649,9 +663,10 @@ class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. + # dbias only fused with an activation derivative, like the CUDA register-tiled kernel. return ( not cfg.WITH_AMAX - and not cfg.WITH_DBIAS + and (not cfg.WITH_DBIAS or cfg.WITH_DACT) and cfg.DTYPE is cutlass.BFloat16 and cfg.COLWISE ) @@ -674,8 +689,10 @@ def eligible_to_dispatch( act_input_aligned = mDActInput.iterator.toint() % 32 == 0 # Whole 32-row by 256-column tiles, grid.y within limits, and aligned input and # colwise-scale pointers. + # With dbias a CTA sums whole 64-row chunks. + rows_per_chunk = 64 if cutlass.const_expr(cfg.WITH_DBIAS) else 32 return ( - M % 32 == 0 + M % rows_per_chunk == 0 and N % 256 == 0 and M // 32 <= 65535 and mX.iterator.toint() % 32 == 0 @@ -688,9 +705,9 @@ def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg # One CTA only quantize one columnwise MXFP8 block at a time self._ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // self._CTA_THREADS_Y # 4 - # SMs of the device compiled for, which sets the table path's walk length. + # SMs of the device compiled for, which set the walk length with a table or dbias. self._sm_count = 0 - if cfg.ACTIVATION_TABLE != "none": + if cfg.ACTIVATION_TABLE != "none" or cfg.WITH_DBIAS: from cuda.bindings import runtime as cudart # pylint: disable=no-name-in-module err, device_id = cudart.cudaGetDevice() @@ -717,7 +734,10 @@ def __call__( WIDE_TILE_COLS = self._CTA_THREADS_X * 16 # A derivative read from a table holds the gradient and the activation input in registers: # 256-column tiles keep it under the register budget, and clusters do not help it. - if cutlass.const_expr(self.cfg.WITH_DACT and self.cfg.ACTIVATION_TABLE != "none"): + # With dbias, the staged FP32 products of a 256-column tile take 32 KB of shared memory. + if cutlass.const_expr( + self.cfg.WITH_DACT and (self.cfg.ACTIVATION_TABLE != "none" or self.cfg.WITH_DBIAS) + ): self._launch( mX, mDActInput, @@ -726,6 +746,7 @@ def __call__( mO_col, mS_col, mActTable, + mWorkspace, stream, ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=1, @@ -740,6 +761,7 @@ def __call__( mO_col, mS_col, mActTable, + mWorkspace, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=4, @@ -753,6 +775,7 @@ def __call__( mO_col, mS_col, mActTable, + mWorkspace, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=1, @@ -766,6 +789,7 @@ def __call__( mO_col, mS_col, mActTable, + mWorkspace, stream, ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=8, @@ -781,6 +805,7 @@ def _launch( mO_col: cute.Tensor, mS_col: cute.Tensor, mActTable: cute.Pointer, + mWorkspace: Optional[cute.Tensor], stream: CUstream, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], CLUSTER_WIDTH: cutlass.Constexpr[int], @@ -808,7 +833,7 @@ def _launch( # more than _WALK_TARGET_COLS / TILE_COLS CTAs and the row tiles divide evenly, and # walk at least two tiles unless that leaves fewer than two CTAs per SM. walk = Int32(1) - if cutlass.const_expr(self.cfg.ACTIVATION_TABLE != "none"): + if cutlass.const_expr(self.cfg.ACTIVATION_TABLE != "none" or self.cfg.WITH_DBIAS): while (grid_rows // walk) % 2 == 0 and grid_cols * (grid_rows // walk) > ( self._WALK_TARGET_COLS // TILE_COLS ): @@ -819,6 +844,12 @@ def _launch( and grid_cols * (grid_rows // 2) >= 2 * self._sm_count ): walk = Int32(2) + # With dbias the grid and the walk count 64-row chunks, whose two row tiles a CTA takes + # back to back: a CTA walks at least one chunk. + if cutlass.const_expr(self.cfg.WITH_DBIAS): + grid_rows = grid_rows // 2 + if walk > 1: + walk = walk // 2 # Keep column CTAs on grid.x so clusters read adjacent column tiles. grid = [grid_cols, grid_rows // walk] block = [THREADS_PER_CTA] @@ -865,6 +896,7 @@ def _launch( gS_col, mActTable, walk, + mWorkspace, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) elif N == 4096: @@ -907,6 +939,7 @@ def _launch( gS_col, mActTable, walk, + mWorkspace, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: @@ -925,6 +958,7 @@ def _launch( gS_col, mActTable, walk, + mWorkspace, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) @@ -939,6 +973,7 @@ def kernel( mS_col: cute.Tensor, mActTable: cute.Pointer, walk: Int32, + mWorkspace: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): """Run the kernel body. Like the CUDA specialized cast kernels it does not consume the @@ -952,6 +987,7 @@ def kernel( mS_col, mActTable, walk, + mWorkspace, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ) @@ -988,6 +1024,7 @@ def _kernel_main( mS_col: cute.Tensor, mActTable: cute.Pointer, walk: Int32, + mWorkspace: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): """Run bidimensional quantization on the CTA's row tiles: blockIdx.y, and with a table @@ -1020,7 +1057,47 @@ def _kernel_main( cute.autovec_copy(gTable[None, piece], sTable_pieces[None, piece]) cute.arch.sync_threads() - if cutlass.const_expr(self.cfg.ACTIVATION_TABLE == "none"): + if cutlass.const_expr(self.cfg.WITH_DBIAS): + # dbias in the CUDA kernels' order: each 64-row chunk's FP32 products summed down the + # column in row order, one workspace row per chunk (reduced by the caller). The CTA + # takes the chunk's two row tiles back to back; each stages its products in shared + # memory, and one thread per column adds them to the column's running sum. + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD + tidx, _, _ = cute.arch.thread_idx() + sDbias = cutlass.memory.SmemAllocator().allocate_tensor( + Float32, + cute.make_layout((TILE_ROWS, TILE_COLS), stride=(TILE_COLS, 1)), + byte_alignment=16, + ) + for step in cutlass.range(walk, unroll=1): + chunk = step * grid_rows + first_cta_row + column_sum = Float32(0.0) + # A run-time loop, so the tile body (and its shared memory) is emitted once. + for half in cutlass.range(2, unroll=1): + self._quantize_row_tile( + mX, + mA, + mO_row, + mS_row, + mO_col, + mS_col, + sTable, + 2 * chunk + half, + cta_col, + cache_policy, + sDbias=sDbias, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ) + cute.arch.sync_threads() + if tidx < TILE_COLS: + for row in cutlass.range_constexpr(TILE_ROWS): + column_sum = column_sum + sDbias[row, tidx] + # The next tile overwrites the staged products. + cute.arch.sync_threads() + if tidx < TILE_COLS: + mWorkspace[chunk, cta_col * TILE_COLS + tidx] = column_sum + elif cutlass.const_expr(self.cfg.ACTIVATION_TABLE == "none"): self._quantize_row_tile( mX, mA, @@ -1064,6 +1141,7 @@ def _quantize_row_tile( cta_col: Int32, cache_policy: Int64, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + sDbias: Optional[cute.Tensor] = None, ): """Quantize the (row_tile, cta_col) tile.""" TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD @@ -1073,9 +1151,9 @@ def _quantize_row_tile( COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) # With a table and a row stride known only at run time, index each data tile from its own # base address with 32-bit offsets (see _data_tile). - DYNAMIC_STRIDE = self.cfg.ACTIVATION_TABLE != "none" and not cute.is_static( - mX.layout.stride[0] - ) + DYNAMIC_STRIDE = ( + self.cfg.ACTIVATION_TABLE != "none" or self.cfg.WITH_DBIAS + ) and not cute.is_static(mX.layout.stride[0]) cta_coord = (row_tile, cta_col) gX_tile = self._data_tile(mX, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 32) gA_tile = None @@ -1101,6 +1179,7 @@ def _quantize_row_tile( ACTIVATION=self.cfg.ACTIVATION, ACTIVATION_TABLE=self.cfg.ACTIVATION_TABLE, sTable=sTable, + sDbias=sDbias, WITH_ACT=self.cfg.WITH_ACT, WITH_DACT=self.cfg.WITH_DACT, DTYPE=self.cfg.DTYPE, From e0b2deaf9c24b438aea73b94a5c214f91463d32c Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 27 Oct 2026 12:14:01 -0700 Subject: [PATCH 21/29] [Common] Reduce amax in the register-resident CuTeDSL MXFP8 kernel Fused activations (forward, derivative, with or without dbias) that ask for the tensor amax now stay on the register-resident bidimensional kernel. Each thread takes the max of the BF16 values it quantizes, ignoring NaN like the CUDA kernels' fmaxf; the CTA reduces them and updates the amax with an atomic max, as the general kernel does. Plain casts with amax keep the general kernel. The general kernel's scales differ from CUDA for blocks holding Inf or NaN; on this path the BF16 amax cases of the exact-composition test now match, including those with special values. Signed-off-by: Teddy Do --- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 71 ++++++++++++++++--- 1 file changed, 63 insertions(+), 8 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 5c8cfd10cf3..00a386526fc 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -14,6 +14,7 @@ from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, + fabs_f32, bf16_pair_magnitude, bf16_mx_scale_reciprocal, bf16_mx_scale_reciprocal_pair, @@ -415,6 +416,7 @@ def quantize_mxfp8_gmem_tile_bidimensional( ACTIVATION_TABLE: cutlass.Constexpr[str], sTable: Optional[cute.Tensor], sDbias: Optional[cute.Tensor], + WITH_AMAX: cutlass.Constexpr[bool], WITH_ACT: cutlass.Constexpr[bool], WITH_DACT: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], @@ -561,6 +563,19 @@ def quantize_mxfp8_gmem_tile_bidimensional( if cutlass.const_expr(sDbias is not None): cute.autovec_copy(rProducts, sDbias_thread[(None, row_idx),]) + # The thread's amax of the values it quantizes (the activation rounded to BF16). Like the + # CUDA kernels' fmaxf, it ignores NaN: the BF16 max returns the other operand of a NaN, and + # the FP32 fmax drops a pair of NaNs. + thread_amax = Float32(0.0) + if cutlass.const_expr(WITH_AMAX): + rX_all_i32 = cute.recast_tensor(rX, Int32) + pair_amax = rX_all_i32[0] + for word in cutlass.range_constexpr(1, cute.size(rX_all_i32)): + pair_amax = abs_max_x2_bf16(pair_amax, rX_all_i32[word]) + thread_amax = cute.arch.fmax( + fabs_f32(x2_lo_to_f32_bf16(pair_amax)), fabs_f32(x2_hi_to_f32_bf16(pair_amax)) + ) + if cutlass.const_expr(ROWWISE): gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] @@ -642,6 +657,8 @@ def quantize_mxfp8_gmem_tile_bidimensional( FP8_DTYPE=FP8_DTYPE, ) + return thread_amax + class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): """Load one CTA tile into registers and quantize it along both axes. @@ -663,9 +680,10 @@ class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: # BF16 packed math. - # dbias only fused with an activation derivative, like the CUDA register-tiled kernel. + # dbias only fused with an activation derivative, like the CUDA register-tiled kernel; + # amax only with a fused activation; plain casts with amax take the general kernel. return ( - not cfg.WITH_AMAX + (not cfg.WITH_AMAX or cfg.WITH_ACT or cfg.WITH_DACT) and (not cfg.WITH_DBIAS or cfg.WITH_DACT) and cfg.DTYPE is cutlass.BFloat16 and cfg.COLWISE @@ -747,6 +765,7 @@ def __call__( mS_col, mActTable, mWorkspace, + mAmax, stream, ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=1, @@ -762,6 +781,7 @@ def __call__( mS_col, mActTable, mWorkspace, + mAmax, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=4, @@ -776,6 +796,7 @@ def __call__( mS_col, mActTable, mWorkspace, + mAmax, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=1, @@ -790,6 +811,7 @@ def __call__( mS_col, mActTable, mWorkspace, + mAmax, stream, ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=8, @@ -806,6 +828,7 @@ def _launch( mS_col: cute.Tensor, mActTable: cute.Pointer, mWorkspace: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], stream: CUstream, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], CLUSTER_WIDTH: cutlass.Constexpr[int], @@ -897,6 +920,7 @@ def _launch( mActTable, walk, mWorkspace, + mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) elif N == 4096: @@ -940,6 +964,7 @@ def _launch( mActTable, walk, mWorkspace, + mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) else: @@ -959,6 +984,7 @@ def _launch( mActTable, walk, mWorkspace, + mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) @@ -974,6 +1000,7 @@ def kernel( mActTable: cute.Pointer, walk: Int32, mWorkspace: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): """Run the kernel body. Like the CUDA specialized cast kernels it does not consume the @@ -988,6 +1015,7 @@ def kernel( mActTable, walk, mWorkspace, + mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ) @@ -1025,6 +1053,7 @@ def _kernel_main( mActTable: cute.Pointer, walk: Int32, mWorkspace: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], ): """Run bidimensional quantization on the CTA's row tiles: blockIdx.y, and with a table @@ -1057,6 +1086,7 @@ def _kernel_main( cute.autovec_copy(gTable[None, piece], sTable_pieces[None, piece]) cute.arch.sync_threads() + thread_amax = Float32(0.0) if cutlass.const_expr(self.cfg.WITH_DBIAS): # dbias in the CUDA kernels' order: each 64-row chunk's FP32 products summed down the # column in row order, one workspace row per chunk (reduced by the caller). The CTA @@ -1075,7 +1105,7 @@ def _kernel_main( column_sum = Float32(0.0) # A run-time loop, so the tile body (and its shared memory) is emitted once. for half in cutlass.range(2, unroll=1): - self._quantize_row_tile( + tile_amax = self._quantize_row_tile( mX, mA, mO_row, @@ -1089,6 +1119,7 @@ def _kernel_main( sDbias=sDbias, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ) + thread_amax = cute.arch.fmax(thread_amax, tile_amax) cute.arch.sync_threads() if tidx < TILE_COLS: for row in cutlass.range_constexpr(TILE_ROWS): @@ -1098,7 +1129,7 @@ def _kernel_main( if tidx < TILE_COLS: mWorkspace[chunk, cta_col * TILE_COLS + tidx] = column_sum elif cutlass.const_expr(self.cfg.ACTIVATION_TABLE == "none"): - self._quantize_row_tile( + tile_amax = self._quantize_row_tile( mX, mA, mO_row, @@ -1111,9 +1142,10 @@ def _kernel_main( cache_policy, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ) + thread_amax = cute.arch.fmax(thread_amax, tile_amax) else: for step in cutlass.range(walk, unroll=1): - self._quantize_row_tile( + tile_amax = self._quantize_row_tile( mX, mA, mO_row, @@ -1126,6 +1158,27 @@ def _kernel_main( cache_policy, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, ) + thread_amax = cute.arch.fmax(thread_amax, tile_amax) + + if cutlass.const_expr(self.cfg.WITH_AMAX): + # The CTA's amax, then an atomic max on the float's bits (amaxes are not negative), + # as the general kernel does. + tidx, _, _ = cute.arch.thread_idx() + warp_amax = cute.arch.warp_redux_sync(thread_amax, kind="fmax") + NUM_WARPS = self._CTA_THREADS_Y * self._CTA_THREADS_X // 32 + sAmax = cutlass.memory.SmemAllocator().allocate_tensor( + Float32, cute.make_layout((NUM_WARPS,)), byte_alignment=4 + ) + if tidx % 32 == 0: + sAmax[tidx // 32] = warp_amax + cute.arch.sync_threads() + if tidx == 0: + cta_amax = Float32(0.0) + for w in cutlass.range_constexpr(NUM_WARPS): + cta_amax = cute.arch.fmax(cta_amax, sAmax[w]) + cute.arch.atomic_max( + cute.recast_ptr(mAmax.iterator, dtype=Int32), cta_amax.bitcast(Int32) + ) @cute.jit def _quantize_row_tile( @@ -1142,8 +1195,9 @@ def _quantize_row_tile( cache_policy: Int64, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], sDbias: Optional[cute.Tensor] = None, - ): - """Quantize the (row_tile, cta_col) tile.""" + ) -> Float32: + """Quantize the (row_tile, cta_col) tile; return the thread's amax of it (0 without + WITH_AMAX).""" TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD TILER = (TILE_ROWS, TILE_COLS) @@ -1167,7 +1221,7 @@ def _quantize_row_tile( gO_col_tile = self._data_tile(mO_col, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 16) gS_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) - quantize_mxfp8_gmem_tile_bidimensional( + return quantize_mxfp8_gmem_tile_bidimensional( gX_tile, gA_tile, gO_row_tile, @@ -1180,6 +1234,7 @@ def _quantize_row_tile( ACTIVATION_TABLE=self.cfg.ACTIVATION_TABLE, sTable=sTable, sDbias=sDbias, + WITH_AMAX=self.cfg.WITH_AMAX, WITH_ACT=self.cfg.WITH_ACT, WITH_DACT=self.cfg.WITH_DACT, DTYPE=self.cfg.DTYPE, From 03903aaeaded3880b050bc586dabe5c9925e35f1 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 27 Oct 2026 13:38:25 -0700 Subject: [PATCH 22/29] [Common] Take any row count in the register-resident CuTeDSL MXFP8 kernel Requests with an activation table, dbias or amax now take the register- resident bidimensional kernel for any number of rows, as the CUDA register- tiled kernel does (the others keep whole 32-row tiles). A tensor whose rows end inside a tile launches a predicated kernel: rows past the tensor load as zeros, which leave the column amaxes and dbias unchanged, and are not stored. Tensors of whole tiles launch the unpredicated kernel, so their speed is unchanged; one kernel holding both variants would double the shared memory of the tile body. With dbias, the last 64-row chunk of an odd number of row tiles has no second tile; skipping it gives the same sum as adding its zero rows. The exact-composition test gains a 976x2048 shape: 31 row tiles, the last of 16 rows. Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 3 + .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 343 +++++++++++++----- 2 files changed, 261 insertions(+), 85 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index e4a5c3f8a17..7daa1fc7928 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -453,6 +453,9 @@ INSTANTIATE_TEST_SUITE_P( // Enough row tiles that the register-resident kernels // walk several per CTA. std::pair{4096, 2048}, + // A row tail on 256-column tiles: 31 row tiles, the + // last of 16 rows, and an odd number for dbias. + std::pair{976, 2048}, std::pair{544, 2080}, // At least 2^23 elements (the register-tiled forward // without tables), with a row and a column tail. diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 00a386526fc..095d8d16dc5 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -61,7 +61,10 @@ def _quantize_rowwise( ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + ROW_EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, ): + """With ROW_EDGE, only the thread's first rows_left rows are stored.""" tidx, _, _ = cute.arch.thread_idx() MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 @@ -95,17 +98,39 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: reciprocal_pair, ) # rO_row_u32 and rO_row share the same memory - cute.copy( - output_store_atom, rO_row, gO_row_thread[(None, row_idx),], cache_policy=cache_policy - ) - if is_writer_thread: - rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_pair)).bitcast(SCALE_DTYPE) + if cutlass.const_expr(ROW_EDGE): + if row_idx < rows_left: + cute.copy( + output_store_atom, + rO_row, + gO_row_thread[(None, row_idx),], + cache_policy=cache_policy, + ) + if is_writer_thread: + rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_pair)).bitcast( + SCALE_DTYPE + ) + cute.copy( + row_scale_store_atom, + rS_row, + cute.local_tile(gS_row_thread, (1,), (row_idx,)), + cache_policy=cache_policy, + ) + else: cute.copy( - row_scale_store_atom, - rS_row, - cute.local_tile(gS_row_thread, (1,), (row_idx,)), + output_store_atom, + rO_row, + gO_row_thread[(None, row_idx),], cache_policy=cache_policy, ) + if is_writer_thread: + rS_row[0] = Uint8(bf16_mx_scale_bytes_pair(reciprocal_pair)).bitcast(SCALE_DTYPE) + cute.copy( + row_scale_store_atom, + rS_row, + cute.local_tile(gS_row_thread, (1,), (row_idx,)), + cache_policy=cache_policy, + ) @cute.jit @@ -304,8 +329,12 @@ def _apply_colwise_scales( ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + ROW_EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, ): """ + With ROW_EDGE, only the thread's first rows_left rows are stored. + rX: RMEM fragment of this thread's rows with shape (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) gO_col_thread: GMEM fragment of this thread's columnwise quantized output with shape @@ -330,7 +359,15 @@ def _apply_colwise_scales( rX_row_i32[2 * pack_idx + 1], rS_col_rcp_u32[2 * pack_idx + 1], ) - cute.copy(store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy) + if cutlass.const_expr(ROW_EDGE): + if row_idx < rows_left: + cute.copy( + store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy + ) + else: + cute.copy( + store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy + ) @cute.jit @@ -352,6 +389,8 @@ def _quantize_colwise( SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + ROW_EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, ): """Reduce columnwise amaxes across the CTA, write the columnwise scales, and quantize rX with them.""" tidx, _, _ = cute.arch.thread_idx() @@ -399,6 +438,8 @@ def _quantize_colwise( ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, FP8_DTYPE=FP8_DTYPE, + ROW_EDGE=ROW_EDGE, + rows_left=rows_left, ) @@ -416,6 +457,8 @@ def quantize_mxfp8_gmem_tile_bidimensional( ACTIVATION_TABLE: cutlass.Constexpr[str], sTable: Optional[cute.Tensor], sDbias: Optional[cute.Tensor], + ROW_EDGE: cutlass.Constexpr[bool], + rows_valid: Int32, WITH_AMAX: cutlass.Constexpr[bool], WITH_ACT: cutlass.Constexpr[bool], WITH_DACT: cutlass.Constexpr[bool], @@ -479,11 +522,32 @@ def quantize_mxfp8_gmem_tile_bidimensional( scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) + # With ROW_EDGE the tile ends at rows_valid: the thread's rows from rows_left on are past the + # tensor. They load as zeros, which leave the column amaxes and dbias unchanged, and are not + # stored. + rows_left = Int32(ELEMENTS_Y_PER_THREAD) + if cutlass.const_expr(ROW_EDGE): + thread_rows = cute.composition(cute.make_identity_tensor(gX_tile.shape), tv_layout)[ + tidx, None + ] + rows_left = rows_valid - thread_rows[((0, 0),)][0] + # Copy all rows a thread sweeps into the RMEM fragment at once for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - cute.copy( - load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy - ) + if cutlass.const_expr(ROW_EDGE): + if row_idx < rows_left: + cute.copy( + load_atom, + gX_thread[(None, row_idx),], + rX[None, row_idx], + cache_policy=cache_policy, + ) + else: + rX[None, row_idx].fill(0) + else: + cute.copy( + load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy + ) if cutlass.const_expr(WITH_ACT or WITH_DACT): rA = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), DTYPE) @@ -500,12 +564,23 @@ def quantize_mxfp8_gmem_tile_bidimensional( # Activation tile has the same layout as the input tile so we can reuse the tv layout gA_thread = cute.composition(gA_tile, tv_layout)[tidx, None] for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - cute.copy( - act_load_atom, - gA_thread[(None, row_idx),], - rA[None, row_idx], - cache_policy=cache_policy, - ) + if cutlass.const_expr(ROW_EDGE): + if row_idx < rows_left: + cute.copy( + act_load_atom, + gA_thread[(None, row_idx),], + rA[None, row_idx], + cache_policy=cache_policy, + ) + else: + rA[None, row_idx].fill(0) + else: + cute.copy( + act_load_atom, + gA_thread[(None, row_idx),], + rA[None, row_idx], + cache_policy=cache_policy, + ) # With dbias, the FP32 products are staged for the column sums (see _kernel_main). if cutlass.const_expr(sDbias is not None): sDbias_thread = cute.composition(sDbias, tv_layout)[tidx, None] @@ -610,6 +685,8 @@ def quantize_mxfp8_gmem_tile_bidimensional( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, SCALE_DTYPE=SCALE_DTYPE, FP8_DTYPE=FP8_DTYPE, + ROW_EDGE=ROW_EDGE, + rows_left=rows_left, ) if cutlass.const_expr(COLWISE): @@ -655,6 +732,8 @@ def quantize_mxfp8_gmem_tile_bidimensional( SCALE_DTYPE=SCALE_DTYPE, DTYPE=DTYPE, FP8_DTYPE=FP8_DTYPE, + ROW_EDGE=ROW_EDGE, + rows_left=rows_left, ) return thread_amax @@ -705,20 +784,27 @@ def eligible_to_dispatch( act_input_aligned = Boolean(True) if cutlass.const_expr(mDActInput is not None): act_input_aligned = mDActInput.iterator.toint() % 32 == 0 - # Whole 32-row by 256-column tiles, grid.y within limits, and aligned input and - # colwise-scale pointers. - # With dbias a CTA sums whole 64-row chunks. - rows_per_chunk = 64 if cutlass.const_expr(cfg.WITH_DBIAS) else 32 + # Whole 256-column tiles, whole 32-row tiles (any row count with a table, dbias or amax, + # see takes_row_tails), grid.y within limits, and aligned input and colwise-scale pointers. + rows_ok = Boolean(True) + if cutlass.const_expr(not MXFP8QuantizeBidimensionalG2RKernel.takes_row_tails(cfg)): + rows_ok = M % 32 == 0 return ( - M % rows_per_chunk == 0 + rows_ok and N % 256 == 0 - and M // 32 <= 65535 + and (M + 31) // 32 <= 65535 and mX.iterator.toint() % 32 == 0 and act_input_aligned and mS_col.iterator.toint() % 16 == 0 and mNoop.toint() == 0 ) + @staticmethod + def takes_row_tails(cfg: MXFP8QuantizeConfig) -> bool: + """Whether this config takes any row count: the requests with a table, dbias or amax, as + the CUDA register-tiled kernel does. The others keep whole 32-row tiles.""" + return cfg.ACTIVATION_TABLE != "none" or cfg.WITH_DBIAS or cfg.WITH_AMAX + def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg # One CTA only quantize one columnwise MXFP8 block at a time @@ -847,6 +933,8 @@ def _launch( grid_cols = N // TILE_COLS grid_rows = M // TILE_ROWS + if cutlass.const_expr(self.takes_row_tails(self.cfg)): + grid_rows = (M + TILE_ROWS - 1) // TILE_ROWS cluster_cols = Int32(CLUSTER_WIDTH) # Adjust cluster size so that the CTAs in the column direction can be evenly divided into clusters. while cluster_cols > 1 and grid_cols % cluster_cols != 0: @@ -870,7 +958,7 @@ def _launch( # With dbias the grid and the walk count 64-row chunks, whose two row tiles a CTA takes # back to back: a CTA walks at least one chunk. if cutlass.const_expr(self.cfg.WITH_DBIAS): - grid_rows = grid_rows // 2 + grid_rows = (grid_rows + 1) // 2 # the last chunk may hold one tile if walk > 1: walk = walk // 2 # Keep column CTAs on grid.x so clusters read adjacent column tiles. @@ -910,19 +998,41 @@ def _launch( mS_col.iterator, cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), ) - self.kernel( - gX, - gA, - gO_row, - gS_row, - gO_col, - gS_col, - mActTable, - walk, - mWorkspace, - mAmax, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) + # A tensor with a row tail takes the kernel whose tiles are predicated. + if cutlass.const_expr(self.takes_row_tails(self.cfg)) and M % TILE_ROWS != 0: + self.kernel( + gX, + gA, + gO_row, + gS_row, + gO_col, + gS_col, + mActTable, + walk, + mWorkspace, + mAmax, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=True, + ).launch( + grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + ) + else: + self.kernel( + gX, + gA, + gO_row, + gS_row, + gO_col, + gS_col, + mActTable, + walk, + mWorkspace, + mAmax, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=False, + ).launch( + grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + ) elif N == 4096: payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) gX = cute.make_tensor(mX.iterator, payload_layout) @@ -954,19 +1064,41 @@ def _launch( mS_col.iterator, cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), ) - self.kernel( - gX, - gA, - gO_row, - gS_row, - gO_col, - gS_col, - mActTable, - walk, - mWorkspace, - mAmax, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) + # A tensor with a row tail takes the kernel whose tiles are predicated. + if cutlass.const_expr(self.takes_row_tails(self.cfg)) and M % TILE_ROWS != 0: + self.kernel( + gX, + gA, + gO_row, + gS_row, + gO_col, + gS_col, + mActTable, + walk, + mWorkspace, + mAmax, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=True, + ).launch( + grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + ) + else: + self.kernel( + gX, + gA, + gO_row, + gS_row, + gO_col, + gS_col, + mActTable, + walk, + mWorkspace, + mAmax, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=False, + ).launch( + grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + ) else: if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): gS_row, gS_col = derive_swizzled_scale_layout( @@ -974,19 +1106,41 @@ def _launch( ) else: gS_row, gS_col = mS_row, mS_col - self.kernel( - mX, - mA, - mO_row, - gS_row, - mO_col, - gS_col, - mActTable, - walk, - mWorkspace, - mAmax, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ).launch(grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream) + # A tensor with a row tail takes the kernel whose tiles are predicated. + if cutlass.const_expr(self.takes_row_tails(self.cfg)) and M % TILE_ROWS != 0: + self.kernel( + mX, + mA, + mO_row, + gS_row, + mO_col, + gS_col, + mActTable, + walk, + mWorkspace, + mAmax, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=True, + ).launch( + grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + ) + else: + self.kernel( + mX, + mA, + mO_row, + gS_row, + mO_col, + gS_col, + mActTable, + walk, + mWorkspace, + mAmax, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=False, + ).launch( + grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + ) @cute.kernel def kernel( @@ -1002,6 +1156,7 @@ def kernel( mWorkspace: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + ROW_TAILS: cutlass.Constexpr[bool] = False, ): """Run the kernel body. Like the CUDA specialized cast kernels it does not consume the cast_noop flag: the dispatcher routes calls with one to the general kernel.""" @@ -1017,6 +1172,7 @@ def kernel( mWorkspace, mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=ROW_TAILS, ) @cute.jit @@ -1055,6 +1211,7 @@ def _kernel_main( mWorkspace: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + ROW_TAILS: cutlass.Constexpr[bool] = False, ): """Run bidimensional quantization on the CTA's row tiles: blockIdx.y, and with a table every gridDim.y-th row tile after it, walk tiles in all.""" @@ -1105,27 +1262,35 @@ def _kernel_main( column_sum = Float32(0.0) # A run-time loop, so the tile body (and its shared memory) is emitted once. for half in cutlass.range(2, unroll=1): - tile_amax = self._quantize_row_tile( - mX, - mA, - mO_row, - mS_row, - mO_col, - mS_col, - sTable, - 2 * chunk + half, - cta_col, - cache_policy, - sDbias=sDbias, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ) - thread_amax = cute.arch.fmax(thread_amax, tile_amax) - cute.arch.sync_threads() - if tidx < TILE_COLS: - for row in cutlass.range_constexpr(TILE_ROWS): - column_sum = column_sum + sDbias[row, tidx] - # The next tile overwrites the staged products. - cute.arch.sync_threads() + # With a row tail, the last chunk of an odd row-tile count has no second + # tile; its rows would add zeros, which leave a sum that starts at +0 + # unchanged. + tile_in_tensor = Boolean(True) + if cutlass.const_expr(ROW_TAILS): + tile_in_tensor = (2 * chunk + half) * TILE_ROWS < mX.shape[0] + if tile_in_tensor: + tile_amax = self._quantize_row_tile( + mX, + mA, + mO_row, + mS_row, + mO_col, + mS_col, + sTable, + 2 * chunk + half, + cta_col, + cache_policy, + sDbias=sDbias, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=ROW_TAILS, + ) + thread_amax = cute.arch.fmax(thread_amax, tile_amax) + cute.arch.sync_threads() + if tidx < TILE_COLS: + for row in cutlass.range_constexpr(TILE_ROWS): + column_sum = column_sum + sDbias[row, tidx] + # The next tile overwrites the staged products. + cute.arch.sync_threads() if tidx < TILE_COLS: mWorkspace[chunk, cta_col * TILE_COLS + tidx] = column_sum elif cutlass.const_expr(self.cfg.ACTIVATION_TABLE == "none"): @@ -1141,6 +1306,7 @@ def _kernel_main( cta_col, cache_policy, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=ROW_TAILS, ) thread_amax = cute.arch.fmax(thread_amax, tile_amax) else: @@ -1157,6 +1323,7 @@ def _kernel_main( cta_col, cache_policy, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + ROW_TAILS=ROW_TAILS, ) thread_amax = cute.arch.fmax(thread_amax, tile_amax) @@ -1195,6 +1362,7 @@ def _quantize_row_tile( cache_policy: Int64, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], sDbias: Optional[cute.Tensor] = None, + ROW_TAILS: cutlass.Constexpr[bool] = False, ) -> Float32: """Quantize the (row_tile, cta_col) tile; return the thread's amax of it (0 without WITH_AMAX).""" @@ -1221,6 +1389,9 @@ def _quantize_row_tile( gO_col_tile = self._data_tile(mO_col, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 16) gS_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) + # With ROW_TAILS (a tensor with a row tail) every tile is predicated: rows past the tensor + # are zero-filled and not stored. A second, unpredicated variant for the whole tiles in + # the same kernel would double the shared memory of the tile body. return quantize_mxfp8_gmem_tile_bidimensional( gX_tile, gA_tile, @@ -1234,6 +1405,8 @@ def _quantize_row_tile( ACTIVATION_TABLE=self.cfg.ACTIVATION_TABLE, sTable=sTable, sDbias=sDbias, + ROW_EDGE=ROW_TAILS, + rows_valid=mX.shape[0] - row_tile * TILE_ROWS, WITH_AMAX=self.cfg.WITH_AMAX, WITH_ACT=self.cfg.WITH_ACT, WITH_DACT=self.cfg.WITH_DACT, From 6e09bf8c693401b3a45cc6424e32d90f23b987fc Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 27 Oct 2026 14:21:15 -0700 Subject: [PATCH 23/29] [Common] Use 256-column tiles for derivatives in the CuTeDSL register kernel A derivative holds the gradient and the activation input in registers. With 512-column tiles (16 columns per thread) the register-resident kernel spills at its 64-register budget; 256-column tiles without clusters, and 32-bit offsets within tiles of a run-time row stride, avoid that, as they already did for derivatives read from a table. Derivatives now also take any row count, like the CUDA register-tiled kernel. Only tiling and addressing change; the arithmetic is the same. GB200, BF16 without tables, geometric mean speedup over TE main: dGeLU 1.69 -> 1.87, dSiLU 1.55 -> 1.87 (CUDA backend: 1.68, 1.83). Signed-off-by: Teddy Do --- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 095d8d16dc5..6d3f909a0da 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -801,9 +801,9 @@ def eligible_to_dispatch( @staticmethod def takes_row_tails(cfg: MXFP8QuantizeConfig) -> bool: - """Whether this config takes any row count: the requests with a table, dbias or amax, as - the CUDA register-tiled kernel does. The others keep whole 32-row tiles.""" - return cfg.ACTIVATION_TABLE != "none" or cfg.WITH_DBIAS or cfg.WITH_AMAX + """Whether this config takes any row count: the derivatives and the requests with a table + or amax, as the CUDA register-tiled kernel does. The others keep whole 32-row tiles.""" + return cfg.ACTIVATION_TABLE != "none" or cfg.WITH_DACT or cfg.WITH_AMAX def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg @@ -836,12 +836,10 @@ def __call__( M, N = mX.shape TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD WIDE_TILE_COLS = self._CTA_THREADS_X * 16 - # A derivative read from a table holds the gradient and the activation input in registers: - # 256-column tiles keep it under the register budget, and clusters do not help it. + # A derivative holds the gradient and the activation input in registers: 256-column tiles + # keep it under the register budget (wider tiles spill), and clusters do not help it. # With dbias, the staged FP32 products of a 256-column tile take 32 KB of shared memory. - if cutlass.const_expr( - self.cfg.WITH_DACT and (self.cfg.ACTIVATION_TABLE != "none" or self.cfg.WITH_DBIAS) - ): + if cutlass.const_expr(self.cfg.WITH_DACT): self._launch( mX, mDActInput, @@ -1371,10 +1369,10 @@ def _quantize_row_tile( TILER = (TILE_ROWS, TILE_COLS) ROW_SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) COL_SCALE_TILER = (TILE_ROWS // 32, TILE_COLS) - # With a table and a row stride known only at run time, index each data tile from its own - # base address with 32-bit offsets (see _data_tile). + # With a table or a derivative and a row stride known only at run time, index each data + # tile from its own base address with 32-bit offsets (see _data_tile). DYNAMIC_STRIDE = ( - self.cfg.ACTIVATION_TABLE != "none" or self.cfg.WITH_DBIAS + self.cfg.ACTIVATION_TABLE != "none" or self.cfg.WITH_DACT ) and not cute.is_static(mX.layout.stride[0]) cta_coord = (row_tile, cta_col) gX_tile = self._data_tile(mX, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 32) From a6504b84af57b3a328eb222906476a88a84e9648 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 27 Oct 2026 15:01:05 -0700 Subject: [PATCH 24/29] [Common] Take column tails in the register-resident CuTeDSL MXFP8 kernel Requests with tails (derivatives, tables, amax) now also take any column count with N % 16 == 0, as the CUDA register-tiled kernel does. A thread's 8 or 16 columns then lie all inside or all past the tensor, so one flag per thread covers them: threads past the tensor load zeros and store no data or scales, which also covers a 16-column partial MX block. With dbias, workspace columns past the tensor are not written. The predicated kernel loads with predicated copies instead of branches, so a thread's row loads still issue together. GB200, BF16, geometric mean speedup over TE main on 993x7200, 4127x2064, 16384x2080 and 32768x7184, general kernel -> register kernel: tables off dSiLU 1.21 -> 1.63, dbias+dSiLU 1.41 -> 1.60; tables on GeLU 1.27 -> 1.59, dSiLU 1.19 -> 1.48, dbias+dSiLU 1.38 -> 1.71. The predicated kernel spills at the register budget, so tail shapes stay behind whole-tile ones. Signed-off-by: Teddy Do --- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 182 +++++++++++------- 1 file changed, 108 insertions(+), 74 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 6d3f909a0da..0485abf023e 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -61,10 +61,10 @@ def _quantize_rowwise( ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - ROW_EDGE: cutlass.Constexpr[bool] = False, + EDGE: cutlass.Constexpr[bool] = False, rows_left: Int32 = 0, ): - """With ROW_EDGE, only the thread's first rows_left rows are stored.""" + """With EDGE, only the thread's first rows_left rows are stored.""" tidx, _, _ = cute.arch.thread_idx() MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 @@ -98,7 +98,7 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: reciprocal_pair, ) # rO_row_u32 and rO_row share the same memory - if cutlass.const_expr(ROW_EDGE): + if cutlass.const_expr(EDGE): if row_idx < rows_left: cute.copy( output_store_atom, @@ -292,9 +292,12 @@ def _write_colwise_scale( CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + EDGE: cutlass.Constexpr[bool] = False, + cols_ok: Boolean = True, ): """ rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) + With EDGE, a thread whose columns are past the tensor (not cols_ok) writes none. """ tidx, _, _ = cute.arch.thread_idx() @@ -312,7 +315,11 @@ def _write_colwise_scale( bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx + 1]) << Uint32(16) ) # Use RMEM buffer to stash the columnwise scale and flush it to GMEM using vectorized store at once - cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + if cutlass.const_expr(EDGE): + if cols_ok: + cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + else: + cute.copy(col_scale_atom, rS_col, gS_col_thread, cache_policy=cache_policy) @cute.jit @@ -329,11 +336,11 @@ def _apply_colwise_scales( ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - ROW_EDGE: cutlass.Constexpr[bool] = False, + EDGE: cutlass.Constexpr[bool] = False, rows_left: Int32 = 0, ): """ - With ROW_EDGE, only the thread's first rows_left rows are stored. + With EDGE, only the thread's first rows_left rows are stored. rX: RMEM fragment of this thread's rows with shape (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) rS_col_rcp: RMEM fragment of this thread's columnwise scale reciprocals with shape (ELEMENTS_X_PER_THREAD,) @@ -359,7 +366,7 @@ def _apply_colwise_scales( rX_row_i32[2 * pack_idx + 1], rS_col_rcp_u32[2 * pack_idx + 1], ) - if cutlass.const_expr(ROW_EDGE): + if cutlass.const_expr(EDGE): if row_idx < rows_left: cute.copy( store_atom, rO_col, gO_col_thread[(None, row_idx),], cache_policy=cache_policy @@ -389,8 +396,9 @@ def _quantize_colwise( SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - ROW_EDGE: cutlass.Constexpr[bool] = False, + EDGE: cutlass.Constexpr[bool] = False, rows_left: Int32 = 0, + cols_ok: Boolean = True, ): """Reduce columnwise amaxes across the CTA, write the columnwise scales, and quantize rX with them.""" tidx, _, _ = cute.arch.thread_idx() @@ -426,6 +434,8 @@ def _quantize_colwise( CTA_THREADS_X=CTA_THREADS_X, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, SCALE_DTYPE=SCALE_DTYPE, + EDGE=EDGE, + cols_ok=cols_ok, ) _apply_colwise_scales( @@ -438,7 +448,7 @@ def _quantize_colwise( ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, FP8_DTYPE=FP8_DTYPE, - ROW_EDGE=ROW_EDGE, + EDGE=EDGE, rows_left=rows_left, ) @@ -457,8 +467,9 @@ def quantize_mxfp8_gmem_tile_bidimensional( ACTIVATION_TABLE: cutlass.Constexpr[str], sTable: Optional[cute.Tensor], sDbias: Optional[cute.Tensor], - ROW_EDGE: cutlass.Constexpr[bool], + EDGE: cutlass.Constexpr[bool], rows_valid: Int32, + cols_valid: Int32, WITH_AMAX: cutlass.Constexpr[bool], WITH_ACT: cutlass.Constexpr[bool], WITH_DACT: cutlass.Constexpr[bool], @@ -522,28 +533,37 @@ def quantize_mxfp8_gmem_tile_bidimensional( scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) - # With ROW_EDGE the tile ends at rows_valid: the thread's rows from rows_left on are past the - # tensor. They load as zeros, which leave the column amaxes and dbias unchanged, and are not - # stored. + # With EDGE the tile ends at rows_valid rows and cols_valid columns: the thread's rows from + # rows_left on are past the tensor, and all of them when its columns are (a thread's columns + # lie all inside or all past the tensor, as N % 16 == 0). They load as zeros, which leave + # the column amaxes and dbias unchanged, and are not stored. rows_left = Int32(ELEMENTS_Y_PER_THREAD) - if cutlass.const_expr(ROW_EDGE): - thread_rows = cute.composition(cute.make_identity_tensor(gX_tile.shape), tv_layout)[ + cols_ok = Boolean(True) + if cutlass.const_expr(EDGE): + thread_origin = cute.composition(cute.make_identity_tensor(gX_tile.shape), tv_layout)[ tidx, None - ] - rows_left = rows_valid - thread_rows[((0, 0),)][0] + ][((0, 0),)] + cols_ok = thread_origin[1] < cols_valid + rows_left = rows_valid - thread_origin[0] + if not cols_ok: + rows_left = Int32(0) # Copy all rows a thread sweeps into the RMEM fragment at once for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - if cutlass.const_expr(ROW_EDGE): - if row_idx < rows_left: - cute.copy( - load_atom, - gX_thread[(None, row_idx),], - rX[None, row_idx], - cache_policy=cache_policy, - ) - else: - rX[None, row_idx].fill(0) + if cutlass.const_expr(EDGE): + # Predicated loads rather than branches, so the rows' loads still issue together. + rX[None, row_idx].fill(0) + row_pred = cute.make_rmem_tensor( + cute.make_layout((ELEMENTS_X_PER_THREAD, (1,)), stride=(0, (0,))), Boolean + ) + row_pred[0] = row_idx < rows_left + cute.copy( + load_atom, + gX_thread[(None, row_idx),], + rX[None, row_idx], + pred=row_pred, + cache_policy=cache_policy, + ) else: cute.copy( load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy @@ -564,16 +584,19 @@ def quantize_mxfp8_gmem_tile_bidimensional( # Activation tile has the same layout as the input tile so we can reuse the tv layout gA_thread = cute.composition(gA_tile, tv_layout)[tidx, None] for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - if cutlass.const_expr(ROW_EDGE): - if row_idx < rows_left: - cute.copy( - act_load_atom, - gA_thread[(None, row_idx),], - rA[None, row_idx], - cache_policy=cache_policy, - ) - else: - rA[None, row_idx].fill(0) + if cutlass.const_expr(EDGE): + rA[None, row_idx].fill(0) + act_row_pred = cute.make_rmem_tensor( + cute.make_layout((ELEMENTS_X_PER_THREAD, (1,)), stride=(0, (0,))), Boolean + ) + act_row_pred[0] = row_idx < rows_left + cute.copy( + act_load_atom, + gA_thread[(None, row_idx),], + rA[None, row_idx], + pred=act_row_pred, + cache_policy=cache_policy, + ) else: cute.copy( act_load_atom, @@ -685,7 +708,7 @@ def quantize_mxfp8_gmem_tile_bidimensional( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, SCALE_DTYPE=SCALE_DTYPE, FP8_DTYPE=FP8_DTYPE, - ROW_EDGE=ROW_EDGE, + EDGE=EDGE, rows_left=rows_left, ) @@ -732,8 +755,9 @@ def quantize_mxfp8_gmem_tile_bidimensional( SCALE_DTYPE=SCALE_DTYPE, DTYPE=DTYPE, FP8_DTYPE=FP8_DTYPE, - ROW_EDGE=ROW_EDGE, + EDGE=EDGE, rows_left=rows_left, + cols_ok=cols_ok, ) return thread_amax @@ -784,14 +808,15 @@ def eligible_to_dispatch( act_input_aligned = Boolean(True) if cutlass.const_expr(mDActInput is not None): act_input_aligned = mDActInput.iterator.toint() % 32 == 0 - # Whole 256-column tiles, whole 32-row tiles (any row count with a table, dbias or amax, - # see takes_row_tails), grid.y within limits, and aligned input and colwise-scale pointers. - rows_ok = Boolean(True) - if cutlass.const_expr(not MXFP8QuantizeBidimensionalG2RKernel.takes_row_tails(cfg)): - rows_ok = M % 32 == 0 + # Whole 32-row by 256-column tiles (any shape with N % 16 == 0 for the requests in + # takes_tails), grid.y within limits, and aligned input and colwise-scale pointers. + shape_ok = Boolean(True) + if cutlass.const_expr(MXFP8QuantizeBidimensionalG2RKernel.takes_tails(cfg)): + shape_ok = N % 16 == 0 + else: + shape_ok = M % 32 == 0 and N % 256 == 0 return ( - rows_ok - and N % 256 == 0 + shape_ok and (M + 31) // 32 <= 65535 and mX.iterator.toint() % 32 == 0 and act_input_aligned @@ -800,9 +825,10 @@ def eligible_to_dispatch( ) @staticmethod - def takes_row_tails(cfg: MXFP8QuantizeConfig) -> bool: - """Whether this config takes any row count: the derivatives and the requests with a table - or amax, as the CUDA register-tiled kernel does. The others keep whole 32-row tiles.""" + def takes_tails(cfg: MXFP8QuantizeConfig) -> bool: + """Whether this config takes row and column tails (any row count, N % 16 == 0): the + derivatives and the requests with a table or amax, as the CUDA register-tiled kernel + does. The others keep whole 32-row by 256-column tiles.""" return cfg.ACTIVATION_TABLE != "none" or cfg.WITH_DACT or cfg.WITH_AMAX def __init__(self, cfg: MXFP8QuantizeConfig): @@ -931,7 +957,8 @@ def _launch( grid_cols = N // TILE_COLS grid_rows = M // TILE_ROWS - if cutlass.const_expr(self.takes_row_tails(self.cfg)): + if cutlass.const_expr(self.takes_tails(self.cfg)): + grid_cols = (N + TILE_COLS - 1) // TILE_COLS grid_rows = (M + TILE_ROWS - 1) // TILE_ROWS cluster_cols = Int32(CLUSTER_WIDTH) # Adjust cluster size so that the CTAs in the column direction can be evenly divided into clusters. @@ -996,8 +1023,10 @@ def _launch( mS_col.iterator, cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), ) - # A tensor with a row tail takes the kernel whose tiles are predicated. - if cutlass.const_expr(self.takes_row_tails(self.cfg)) and M % TILE_ROWS != 0: + # A tensor with a row or column tail takes the kernel whose tiles are predicated. + if cutlass.const_expr(self.takes_tails(self.cfg)) and ( + M % TILE_ROWS != 0 or N % TILE_COLS != 0 + ): self.kernel( gX, gA, @@ -1010,7 +1039,7 @@ def _launch( mWorkspace, mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=True, + TAILS=True, ).launch( grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream ) @@ -1027,7 +1056,7 @@ def _launch( mWorkspace, mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=False, + TAILS=False, ).launch( grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream ) @@ -1062,8 +1091,10 @@ def _launch( mS_col.iterator, cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), ) - # A tensor with a row tail takes the kernel whose tiles are predicated. - if cutlass.const_expr(self.takes_row_tails(self.cfg)) and M % TILE_ROWS != 0: + # A tensor with a row or column tail takes the kernel whose tiles are predicated. + if cutlass.const_expr(self.takes_tails(self.cfg)) and ( + M % TILE_ROWS != 0 or N % TILE_COLS != 0 + ): self.kernel( gX, gA, @@ -1076,7 +1107,7 @@ def _launch( mWorkspace, mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=True, + TAILS=True, ).launch( grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream ) @@ -1093,7 +1124,7 @@ def _launch( mWorkspace, mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=False, + TAILS=False, ).launch( grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream ) @@ -1104,8 +1135,10 @@ def _launch( ) else: gS_row, gS_col = mS_row, mS_col - # A tensor with a row tail takes the kernel whose tiles are predicated. - if cutlass.const_expr(self.takes_row_tails(self.cfg)) and M % TILE_ROWS != 0: + # A tensor with a row or column tail takes the kernel whose tiles are predicated. + if cutlass.const_expr(self.takes_tails(self.cfg)) and ( + M % TILE_ROWS != 0 or N % TILE_COLS != 0 + ): self.kernel( mX, mA, @@ -1118,7 +1151,7 @@ def _launch( mWorkspace, mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=True, + TAILS=True, ).launch( grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream ) @@ -1135,7 +1168,7 @@ def _launch( mWorkspace, mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=False, + TAILS=False, ).launch( grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream ) @@ -1154,7 +1187,7 @@ def kernel( mWorkspace: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - ROW_TAILS: cutlass.Constexpr[bool] = False, + TAILS: cutlass.Constexpr[bool] = False, ): """Run the kernel body. Like the CUDA specialized cast kernels it does not consume the cast_noop flag: the dispatcher routes calls with one to the general kernel.""" @@ -1170,7 +1203,7 @@ def kernel( mWorkspace, mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=ROW_TAILS, + TAILS=TAILS, ) @cute.jit @@ -1209,7 +1242,7 @@ def _kernel_main( mWorkspace: Optional[cute.Tensor], mAmax: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], - ROW_TAILS: cutlass.Constexpr[bool] = False, + TAILS: cutlass.Constexpr[bool] = False, ): """Run bidimensional quantization on the CTA's row tiles: blockIdx.y, and with a table every gridDim.y-th row tile after it, walk tiles in all.""" @@ -1264,7 +1297,7 @@ def _kernel_main( # tile; its rows would add zeros, which leave a sum that starts at +0 # unchanged. tile_in_tensor = Boolean(True) - if cutlass.const_expr(ROW_TAILS): + if cutlass.const_expr(TAILS): tile_in_tensor = (2 * chunk + half) * TILE_ROWS < mX.shape[0] if tile_in_tensor: tile_amax = self._quantize_row_tile( @@ -1280,7 +1313,7 @@ def _kernel_main( cache_policy, sDbias=sDbias, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=ROW_TAILS, + TAILS=TAILS, ) thread_amax = cute.arch.fmax(thread_amax, tile_amax) cute.arch.sync_threads() @@ -1289,7 +1322,7 @@ def _kernel_main( column_sum = column_sum + sDbias[row, tidx] # The next tile overwrites the staged products. cute.arch.sync_threads() - if tidx < TILE_COLS: + if tidx < TILE_COLS and cta_col * TILE_COLS + tidx < mX.shape[1]: mWorkspace[chunk, cta_col * TILE_COLS + tidx] = column_sum elif cutlass.const_expr(self.cfg.ACTIVATION_TABLE == "none"): tile_amax = self._quantize_row_tile( @@ -1304,7 +1337,7 @@ def _kernel_main( cta_col, cache_policy, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=ROW_TAILS, + TAILS=TAILS, ) thread_amax = cute.arch.fmax(thread_amax, tile_amax) else: @@ -1321,7 +1354,7 @@ def _kernel_main( cta_col, cache_policy, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - ROW_TAILS=ROW_TAILS, + TAILS=TAILS, ) thread_amax = cute.arch.fmax(thread_amax, tile_amax) @@ -1360,7 +1393,7 @@ def _quantize_row_tile( cache_policy: Int64, ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], sDbias: Optional[cute.Tensor] = None, - ROW_TAILS: cutlass.Constexpr[bool] = False, + TAILS: cutlass.Constexpr[bool] = False, ) -> Float32: """Quantize the (row_tile, cta_col) tile; return the thread's amax of it (0 without WITH_AMAX).""" @@ -1387,8 +1420,8 @@ def _quantize_row_tile( gO_col_tile = self._data_tile(mO_col, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 16) gS_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) - # With ROW_TAILS (a tensor with a row tail) every tile is predicated: rows past the tensor - # are zero-filled and not stored. A second, unpredicated variant for the whole tiles in + # With TAILS (a tensor with a row or column tail) every tile is predicated: rows and + # columns past the tensor are zero-filled and not stored. A second, unpredicated variant for the whole tiles in # the same kernel would double the shared memory of the tile body. return quantize_mxfp8_gmem_tile_bidimensional( gX_tile, @@ -1403,8 +1436,9 @@ def _quantize_row_tile( ACTIVATION_TABLE=self.cfg.ACTIVATION_TABLE, sTable=sTable, sDbias=sDbias, - ROW_EDGE=ROW_TAILS, + EDGE=TAILS, rows_valid=mX.shape[0] - row_tile * TILE_ROWS, + cols_valid=mX.shape[1] - cta_col * TILE_COLS, WITH_AMAX=self.cfg.WITH_AMAX, WITH_ACT=self.cfg.WITH_ACT, WITH_DACT=self.cfg.WITH_DACT, From 98dfb671a7a5673356c49d47936217433b49be47 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 27 Oct 2026 16:13:57 -0700 Subject: [PATCH 25/29] [Common] Take FP16 and FP32 inputs in the CuTeDSL register-resident MXFP8 kernel Fused activations with FP16 or FP32 input now take the register-resident bidimensional kernel where it is faster than the general one. As in the CUDA register-tiled kernel, these inputs are scaled in FP32: a block's scale is e8m0(max(amax * max_norm_rcp, 0)) and each value is multiplied in FP32 and converted to FP8. The activation is computed in FP32 and rounded to the input type; dbias takes the FP32 products. FP32 tiles use twice the warps with half the rows each. The BF16 datapath is unchanged. FP16 colwise and both-direction outputs match the CUDA kernels byte for byte (mxh2 parity against TE main, special values included). FP32 SiLU still differs from CUDA, as in the general kernel: CuTeDSL's exp is not expf. GB200, geometric mean speedup over TE main, general kernel -> this kernel: FP16 GeLU 1.46 -> 1.62, SiLU 1.41 -> 1.69, dGeLU 1.56 -> 1.69, dSiLU 1.28 -> 1.76, dbias+dSiLU 1.48 -> 1.74; FP32 SiLU 1.32 -> 1.54, dGeLU 1.38 -> 1.52, dSiLU 1.10 -> 1.85, dbias+dSiLU 1.18 -> 1.53. dbias+dGeLU and FP32 GeLU are faster on the general kernel and stay there. Signed-off-by: Teddy Do --- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 479 ++++++++++++++---- 1 file changed, 383 insertions(+), 96 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index 0485abf023e..a4bfe3921bd 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -14,6 +14,7 @@ from transformer_engine.common.CuTeDSL.utils import ( abs_max_x2_bf16, + exp2f_rcp, fabs_f32, bf16_pair_magnitude, bf16_mx_scale_reciprocal, @@ -33,7 +34,11 @@ table_derivative_2x, table_forward_2x, ) -from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + get_cvt_f32x2_to_fp8x2_func, + mul_bf16x2x2_cvt_bf16x4_to_fp8x4, +) from .quantize_mxfp8_common import ( MXFP8_BLOCK_SCALING_SIZE, MXFP8QuantizeConfig, @@ -453,6 +458,170 @@ def _quantize_colwise( ) +# FP16 and FP32 inputs scale in FP32, as the CUDA register-tiled kernel and the generic kernel do +# (2^(127-e) is often outside FP16's range): a block's scale is +# e8m0(max(amax * MAX_NORM_RCP, 0)), so a block of only NaNs gets the scale of 0 like the +# generic kernel's 0-seeded fmaxf, and each value is multiplied in FP32 by 2^(127-e) and then +# converted to FP8. + + +@cute.jit +def _scale_from_amax_f32(amax: Float32, MAX_NORM_RCP: cutlass.Constexpr[float]): + """The E8M0 scale of a block with this amax, and its FP32 reciprocal.""" + biased = cvt_f32_to_fp8e8m0fnu(cute.arch.fmax(amax * Float32(MAX_NORM_RCP), Float32(0.0))) + return biased, exp2f_rcp(biased) + + +@cute.jit +def _quantize_rowwise_f32_scaled( + rX: cute.Tensor, + gO_row_thread: cute.Tensor, + gS_row_thread: cute.Tensor, + output_store_atom: cute.CopyAtom, + row_scale_store_atom: cute.CopyAtom, + cache_policy: Int64, + THREADS_X_PER_MX_BLOCK: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, +): + """_quantize_rowwise for FP16 and FP32 values, scaled in FP32.""" + tidx, _, _ = cute.arch.thread_idx() + cvt_x2 = get_cvt_f32x2_to_fp8x2_func(FP8_DTYPE) + rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) + rO_row = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_row_u16 = cute.recast_tensor(rO_row, Uint16) + is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + amax = Float32(0.0) + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + amax = cute.arch.fmax(amax, fabs_f32(Float32(rX[col, row_idx]))) + # The block's amax across the lanes that share it. + for stage in cutlass.range_constexpr(THREADS_X_PER_MX_BLOCK.bit_length() - 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) + biased, reciprocal = _scale_from_amax_f32(amax, MAX_NORM_RCP) + for pair in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): + rO_row_u16[pair] = cvt_x2( + Float32(rX[2 * pair + 1, row_idx]) * reciprocal, + Float32(rX[2 * pair, row_idx]) * reciprocal, + ) + row_ok = Boolean(True) + if cutlass.const_expr(EDGE): + row_ok = row_idx < rows_left + if row_ok: + cute.copy( + output_store_atom, + rO_row, + gO_row_thread[(None, row_idx),], + cache_policy=cache_policy, + ) + if is_writer_thread: + rS_row[0] = biased + cute.copy( + row_scale_store_atom, + rS_row, + cute.local_tile(gS_row_thread, (1,), (row_idx,)), + cache_policy=cache_policy, + ) + + +@cute.jit +def _quantize_colwise_f32_scaled( + rX: cute.Tensor, + gO_col_thread: cute.Tensor, + gS_col_thread: cute.Tensor, + output_store_atom: cute.CopyAtom, + col_scale_store_atom: cute.CopyAtom, + cache_policy: Int64, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, + cols_ok: Boolean = True, +): + """_quantize_colwise for FP16 and FP32 values, scaled in FP32. Thread x of each warp row + holds columns [x * ELEMENTS_X_PER_THREAD, (x + 1) * ELEMENTS_X_PER_THREAD).""" + tidx, _, _ = cute.arch.thread_idx() + THREADS = CTA_THREADS_Y * CTA_THREADS_X + TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + cvt_x2 = get_cvt_f32x2_to_fp8x2_func(FP8_DTYPE) + allocator = cutlass.memory.SmemAllocator() + # Each thread's column maxima over its rows, contiguous per thread. + sPartial = allocator.allocate_tensor( + Float32, + cute.make_layout((ELEMENTS_X_PER_THREAD, THREADS), stride=(1, ELEMENTS_X_PER_THREAD)), + byte_alignment=16, + ) + sScale = allocator.allocate_tensor( + SCALE_DTYPE, cute.make_layout((TILE_COLS,)), byte_alignment=16 + ) + sReciprocal = allocator.allocate_tensor( + Float32, cute.make_layout((TILE_COLS,)), byte_alignment=16 + ) + + rPartial = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), Float32) + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + amax = Float32(0.0) + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + amax = cute.arch.fmax(amax, fabs_f32(Float32(rX[col, row_idx]))) + rPartial[col] = amax + cute.autovec_copy(rPartial, sPartial[None, tidx]) + cute.arch.sync_threads() + + # One thread per column reduces the warps' maxima: column c is thread (c / X) of each warp + # row, element c % X. + if tidx < TILE_COLS: + owner = tidx // ELEMENTS_X_PER_THREAD + element = tidx % ELEMENTS_X_PER_THREAD + amax = Float32(0.0) + for warp_row in cutlass.range_constexpr(CTA_THREADS_Y): + amax = cute.arch.fmax(amax, sPartial[element, warp_row * CTA_THREADS_X + owner]) + biased, reciprocal = _scale_from_amax_f32(amax, MAX_NORM_RCP) + sScale[tidx] = biased + sReciprocal[tidx] = reciprocal + cute.arch.sync_threads() + + column0 = (tidx % CTA_THREADS_X) * ELEMENTS_X_PER_THREAD + # The first warp row writes the scales. + if tidx < CTA_THREADS_X: + rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + rS_col[col] = sScale[column0 + col] + if cols_ok: + cute.copy(col_scale_store_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + + rReciprocal = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), Float32) + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + rReciprocal[col] = sReciprocal[column0 + col] + rO_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_col_u16 = cute.recast_tensor(rO_col, Uint16) + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + for pair in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): + rO_col_u16[pair] = cvt_x2( + Float32(rX[2 * pair + 1, row_idx]) * rReciprocal[2 * pair + 1], + Float32(rX[2 * pair, row_idx]) * rReciprocal[2 * pair], + ) + row_ok = Boolean(True) + if cutlass.const_expr(EDGE): + row_ok = row_idx < rows_left + if row_ok: + cute.copy( + output_store_atom, + rO_col, + gO_col_thread[(None, row_idx),], + cache_policy=cache_policy, + ) + + @cute.jit def quantize_mxfp8_gmem_tile_bidimensional( gX_tile: cute.Tensor, @@ -475,6 +644,7 @@ def quantize_mxfp8_gmem_tile_bidimensional( WITH_DACT: cutlass.Constexpr[bool], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], @@ -608,64 +778,87 @@ def quantize_mxfp8_gmem_tile_bidimensional( if cutlass.const_expr(sDbias is not None): sDbias_thread = cute.composition(sDbias, tv_layout)[tidx, None] # Apply activation row by row - for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - rX_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - rA_i32 = cute.recast_tensor(rA[None, row_idx], Int32) - if cutlass.const_expr(sDbias is not None): - rProducts = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), Float32) - for i in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): - if cutlass.const_expr(ACTIVATION_TABLE != "none"): - # Read the activation from the table, as the CUDA register-tiled kernel does. - if cutlass.const_expr(WITH_DACT): - low, high, end, stride = DERIVATIVE_TABLE_WINDOWS[ - (ACTIVATION, ACTIVATION_TABLE) - ] - d_lo, d_hi = table_derivative_2x( - rA_i32[i], - cute.recast_tensor(sTable, Float32), - LO=low, - HI=high, - SATURATED_END=end, - STRIDE=stride, - SLOPE=DERIVATIVE_SLOPES[ACTIVATION], - ) - lo = x2_lo_to_f32_bf16(rX_i32[i]) * d_lo - hi = x2_hi_to_f32_bf16(rX_i32[i]) * d_hi - if cutlass.const_expr(sDbias is not None): - rProducts[2 * i] = lo - rProducts[2 * i + 1] = hi - rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) - else: - low, high = FORWARD_TABLE_WINDOWS[ACTIVATION] - rX_i32[i] = table_forward_2x( - rX_i32[i], - cute.recast_tensor(sTable, Uint16), - LOW_SENTINEL=low - 1, - HIGH_SENTINEL=high + 1, - ) - else: - lo = x2_lo_to_f32_bf16(rX_i32[i]) - hi = x2_hi_to_f32_bf16(rX_i32[i]) + if cutlass.const_expr(DTYPE is not cutlass.BFloat16): + # FP16 and FP32: per value in FP32, rounded to the input type, as the generic + # kernel does; dbias takes the FP32 products before rounding. + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + if cutlass.const_expr(sDbias is not None): + rProducts = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), Float32) + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + x = Float32(rX[col, row_idx]) if cutlass.const_expr(WITH_DACT): dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) - lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) - hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + x = x * dop(Float32(rA[col, row_idx])) if cutlass.const_expr(sDbias is not None): - rProducts[2 * i] = lo - rProducts[2 * i + 1] = hi + rProducts[col] = x else: - op = activation_op(ACTIVATION, ACTIVATION_TABLE) - lo = op(lo) - hi = op(hi) - rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) - if cutlass.const_expr(sDbias is not None): - cute.autovec_copy(rProducts, sDbias_thread[(None, row_idx),]) + x = activation_op(ACTIVATION, ACTIVATION_TABLE)(x) + rX[col, row_idx] = DTYPE(x) + if cutlass.const_expr(sDbias is not None): + cute.autovec_copy(rProducts, sDbias_thread[(None, row_idx),]) + else: + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + rA_i32 = cute.recast_tensor(rA[None, row_idx], Int32) + if cutlass.const_expr(sDbias is not None): + rProducts = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), Float32) + for i in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): + if cutlass.const_expr(ACTIVATION_TABLE != "none"): + # Read the activation from the table, as the CUDA register-tiled kernel does. + if cutlass.const_expr(WITH_DACT): + low, high, end, stride = DERIVATIVE_TABLE_WINDOWS[ + (ACTIVATION, ACTIVATION_TABLE) + ] + d_lo, d_hi = table_derivative_2x( + rA_i32[i], + cute.recast_tensor(sTable, Float32), + LO=low, + HI=high, + SATURATED_END=end, + STRIDE=stride, + SLOPE=DERIVATIVE_SLOPES[ACTIVATION], + ) + lo = x2_lo_to_f32_bf16(rX_i32[i]) * d_lo + hi = x2_hi_to_f32_bf16(rX_i32[i]) * d_hi + if cutlass.const_expr(sDbias is not None): + rProducts[2 * i] = lo + rProducts[2 * i + 1] = hi + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + else: + low, high = FORWARD_TABLE_WINDOWS[ACTIVATION] + rX_i32[i] = table_forward_2x( + rX_i32[i], + cute.recast_tensor(sTable, Uint16), + LOW_SENTINEL=low - 1, + HIGH_SENTINEL=high + 1, + ) + else: + lo = x2_lo_to_f32_bf16(rX_i32[i]) + hi = x2_hi_to_f32_bf16(rX_i32[i]) + if cutlass.const_expr(WITH_DACT): + dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) + lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) + hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + if cutlass.const_expr(sDbias is not None): + rProducts[2 * i] = lo + rProducts[2 * i + 1] = hi + else: + op = activation_op(ACTIVATION, ACTIVATION_TABLE) + lo = op(lo) + hi = op(hi) + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + if cutlass.const_expr(sDbias is not None): + cute.autovec_copy(rProducts, sDbias_thread[(None, row_idx),]) # The thread's amax of the values it quantizes (the activation rounded to BF16). Like the # CUDA kernels' fmaxf, it ignores NaN: the BF16 max returns the other operand of a NaN, and # the FP32 fmax drops a pair of NaNs. thread_amax = Float32(0.0) - if cutlass.const_expr(WITH_AMAX): + if cutlass.const_expr(WITH_AMAX and DTYPE is not cutlass.BFloat16): + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + thread_amax = cute.arch.fmax(thread_amax, fabs_f32(Float32(rX[col, row_idx]))) + elif cutlass.const_expr(WITH_AMAX): rX_all_i32 = cute.recast_tensor(rX, Int32) pair_amax = rX_all_i32[0] for word in cutlass.range_constexpr(1, cute.size(rX_all_i32)): @@ -695,22 +888,40 @@ def quantize_mxfp8_gmem_tile_bidimensional( cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8 ) - _quantize_rowwise( - rX, - gO_row_thread, - gS_row_thread, - output_store_atom=store_atom, - row_scale_store_atom=row_scale_atom, - cache_policy=cache_policy, - scale_and_pack=scale_and_pack, - THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - FP8_DTYPE=FP8_DTYPE, - EDGE=EDGE, - rows_left=rows_left, - ) + if cutlass.const_expr(DTYPE is not cutlass.BFloat16): + _quantize_rowwise_f32_scaled( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + MAX_NORM_RCP=MAX_NORM_RCP, + EDGE=EDGE, + rows_left=rows_left, + ) + else: + _quantize_rowwise( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + EDGE=EDGE, + rows_left=rows_left, + ) if cutlass.const_expr(COLWISE): gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] @@ -740,25 +951,45 @@ def quantize_mxfp8_gmem_tile_bidimensional( num_bits_per_copy=8 if WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, ) - _quantize_colwise( - rX, - gO_col_thread, - gS_col_thread, - output_store_atom=store_atom, - col_scale_store_atom=col_scale_atom, - cache_policy=cache_policy, - scale_and_pack=scale_and_pack, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - DTYPE=DTYPE, - FP8_DTYPE=FP8_DTYPE, - EDGE=EDGE, - rows_left=rows_left, - cols_ok=cols_ok, - ) + if cutlass.const_expr(DTYPE is not cutlass.BFloat16): + _quantize_colwise_f32_scaled( + rX, + gO_col_thread, + gS_col_thread, + output_store_atom=store_atom, + col_scale_store_atom=col_scale_atom, + cache_policy=cache_policy, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + MAX_NORM_RCP=MAX_NORM_RCP, + EDGE=EDGE, + rows_left=rows_left, + cols_ok=cols_ok, + ) + else: + _quantize_colwise( + rX, + gO_col_thread, + gS_col_thread, + output_store_atom=store_atom, + col_scale_store_atom=col_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + EDGE=EDGE, + rows_left=rows_left, + cols_ok=cols_ok, + ) return thread_amax @@ -782,16 +1013,33 @@ class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): @staticmethod def supports(cfg: MXFP8QuantizeConfig) -> bool: - # BF16 packed math. + # BF16 packed math; FP16 and FP32 scaled in FP32, only with a fused activation. # dbias only fused with an activation derivative, like the CUDA register-tiled kernel; # amax only with a fused activation; plain casts with amax take the general kernel. return ( (not cfg.WITH_AMAX or cfg.WITH_ACT or cfg.WITH_DACT) and (not cfg.WITH_DBIAS or cfg.WITH_DACT) - and cfg.DTYPE is cutlass.BFloat16 + and ( + cfg.DTYPE is cutlass.BFloat16 + or MXFP8QuantizeBidimensionalG2RKernel.takes_scaled_in_f32(cfg) + ) and cfg.COLWISE ) + @staticmethod + def takes_scaled_in_f32(cfg: MXFP8QuantizeConfig) -> bool: + """Whether an FP16 or FP32 request takes this kernel: the fused activations where it is + faster than the general kernel. GB200, geometric mean speedup over TE main, general -> + this kernel: FP16 GeLU 1.46 -> 1.62, SiLU 1.41 -> 1.69, dGeLU 1.56 -> 1.69, dSiLU + 1.28 -> 1.76, dbias_dSiLU 1.48 -> 1.74; FP32 SiLU 1.32 -> 1.54, dGeLU 1.38 -> 1.52, + dSiLU 1.10 -> 1.85, dbias_dSiLU 1.18 -> 1.53. dbias_dGeLU (FP16 1.75 -> 1.50, FP32 + 1.47 -> 1.25) and FP32 GeLU (1.20 -> 1.16) keep the general kernel.""" + if not (cfg.WITH_ACT or cfg.WITH_DACT): + return False + if cfg.WITH_DBIAS and cfg.ACTIVATION == "dgelu": + return False + return not (cfg.DTYPE is cutlass.Float32 and cfg.ACTIVATION == "gelu") + @staticmethod @cute.jit def eligible_to_dispatch( @@ -829,10 +1077,21 @@ def takes_tails(cfg: MXFP8QuantizeConfig) -> bool: """Whether this config takes row and column tails (any row count, N % 16 == 0): the derivatives and the requests with a table or amax, as the CUDA register-tiled kernel does. The others keep whole 32-row by 256-column tiles.""" - return cfg.ACTIVATION_TABLE != "none" or cfg.WITH_DACT or cfg.WITH_AMAX + return ( + cfg.ACTIVATION_TABLE != "none" + or cfg.WITH_DACT + or cfg.WITH_AMAX + or cfg.DTYPE is not cutlass.BFloat16 + ) def __init__(self, cfg: MXFP8QuantizeConfig): self.cfg = cfg + # FP32 values take twice the registers: twice the warps, each with half the rows, and + # half the resident CTAs per SM, as in the CUDA register-tiled kernel. + self._MIN_BLOCKS_PER_MP = 4 + if cfg.DTYPE is cutlass.Float32: + self._CTA_THREADS_Y = 16 + self._MIN_BLOCKS_PER_MP = 2 # One CTA only quantize one columnwise MXFP8 block at a time self._ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // self._CTA_THREADS_Y # 4 # SMs of the device compiled for, which set the walk length with a table or dbias. @@ -865,7 +1124,8 @@ def __call__( # A derivative holds the gradient and the activation input in registers: 256-column tiles # keep it under the register budget (wider tiles spill), and clusters do not help it. # With dbias, the staged FP32 products of a 256-column tile take 32 KB of shared memory. - if cutlass.const_expr(self.cfg.WITH_DACT): + # FP16 and FP32 values are scaled in FP32 and take the same tiles. + if cutlass.const_expr(self.cfg.WITH_DACT or self.cfg.DTYPE is not cutlass.BFloat16): self._launch( mX, mDActInput, @@ -1041,7 +1301,11 @@ def _launch( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, TAILS=True, ).launch( - grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, ) else: self.kernel( @@ -1058,7 +1322,11 @@ def _launch( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, TAILS=False, ).launch( - grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, ) elif N == 4096: payload_layout = cute.make_layout((M, 4096), stride=(4096, 1)) @@ -1109,7 +1377,11 @@ def _launch( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, TAILS=True, ).launch( - grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, ) else: self.kernel( @@ -1126,7 +1398,11 @@ def _launch( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, TAILS=False, ).launch( - grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, ) else: if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): @@ -1153,7 +1429,11 @@ def _launch( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, TAILS=True, ).launch( - grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, ) else: self.kernel( @@ -1170,7 +1450,11 @@ def _launch( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, TAILS=False, ).launch( - grid=grid, block=block, cluster=cluster, min_blocks_per_mp=4, stream=stream + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, ) @cute.kernel @@ -1405,7 +1689,9 @@ def _quantize_row_tile( # With a table or a derivative and a row stride known only at run time, index each data # tile from its own base address with 32-bit offsets (see _data_tile). DYNAMIC_STRIDE = ( - self.cfg.ACTIVATION_TABLE != "none" or self.cfg.WITH_DACT + self.cfg.ACTIVATION_TABLE != "none" + or self.cfg.WITH_DACT + or self.cfg.DTYPE is not cutlass.BFloat16 ) and not cute.is_static(mX.layout.stride[0]) cta_coord = (row_tile, cta_col) gX_tile = self._data_tile(mX, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 32) @@ -1444,6 +1730,7 @@ def _quantize_row_tile( WITH_DACT=self.cfg.WITH_DACT, DTYPE=self.cfg.DTYPE, FP8_DTYPE=self.cfg.FP8_DTYPE, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, CTA_THREADS_Y=self._CTA_THREADS_Y, CTA_THREADS_X=self._CTA_THREADS_X, From 3c96842d343763311341f03d72ac4ae976f89498 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 27 Oct 2026 18:06:35 -0700 Subject: [PATCH 26/29] [Common] Match the CUDA kernels in CuTeDSL SiLU, dbias order and NaN-block scales - SiLU is x * sigmoid(x), as util/math.h computes it; x / (1 + exp(-x)) rounds differently, so FP32 SiLU differed from CUDA. - The general kernel carries one dbias sum down the column across its stages, as the CUDA kernels do, instead of adding 32-row partial sums. - With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's generic kernel, in the register-resident, 2-lane rowwise and general kernels. Plain casts keep the NaN scale of CUDA's cast kernels. - FP32 dbias+dGeLU takes the register-resident kernel, whose dbias matches CUDA (the general kernel's still differs in the last bit for FP32). mxh2 parity against TE main (CUDA), tables off, 702 cases over BF16/FP16/FP32, all modes, layouts and fills including special values: 678 match (PR #3642 alone: 464). The rest are FP32 dbias with rowwise-only output, where CUDA adds row pairs (r, r + 32) before summing. Signed-off-by: Teddy Do --- .../common/CuTeDSL/activations.py | 3 +- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 65 +++++++++++++++++-- .../mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py | 4 ++ .../cast/mxfp8/quantize_mxfp8_g2s_general.py | 44 +++++++++---- 4 files changed, 98 insertions(+), 18 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/activations.py b/transformer_engine/common/CuTeDSL/activations.py index c1de06139e0..5f2440c9299 100644 --- a/transformer_engine/common/CuTeDSL/activations.py +++ b/transformer_engine/common/CuTeDSL/activations.py @@ -38,7 +38,8 @@ def act_gelu(x: Float32) -> Float32: def act_silu(x: Float32) -> Float32: - return x / (Float32(1.0) + exp(-x)) + # x * sigmoid(x), as math.h computes it: x / (1 + exp(-x)) rounds differently in FP32. + return x * sigmoid(x) def act_qgelu(x: Float32) -> Float32: diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index a4bfe3921bd..c4eed9f1040 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -49,6 +49,41 @@ ) +@cute.jit +def _bf16_nan_magnitudes_to_zero(pair: Uint32) -> Uint32: + """A packed pair of BF16 block amaxes with each NaN half replaced by +0. + + With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's + generic kernel (0-seeded fmaxf); CUDA's cast kernels, like the plain cast here, give it the NaN + scale (cast_bidim.cuh's kNaNReciprocal). + """ + magnitudes = pair & Uint32(0x7FFF7FFF) + lo = magnitudes & Uint32(0xFFFF) + hi = magnitudes >> Uint32(16) + if lo > Uint32(0x7F80): + lo = Uint32(0) + if hi > Uint32(0x7F80): + hi = Uint32(0) + return lo | (hi << Uint32(16)) + + +@cute.jit +def _bf16_mx_scale_reciprocal_fused( + magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int] +) -> Uint32: + """bf16_mx_scale_reciprocal with a fused activation: a NaN block amax gives the reciprocal of + amax 0 (see _bf16_nan_magnitudes_to_zero). Same instructions, so it costs nothing extra.""" + rounded = (magnitude + Uint32(31)) & Uint32(0x7F80) + if rounded < Uint32(MAX_EXPONENT << 7): # pylint: disable=consider-using-max-builtin + rounded = Uint32(MAX_EXPONENT << 7) + reciprocal = Uint32((254 + MAX_EXPONENT) << 7) - rounded + if magnitude >= Uint32(0x7F80): + reciprocal = Uint32(0x0040) + if magnitude != Uint32(0x7F80): + reciprocal = Uint32(254 << 7) + return reciprocal + + @cute.jit def _quantize_rowwise( rX: cute.Tensor, @@ -68,8 +103,10 @@ def _quantize_rowwise( FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], EDGE: cutlass.Constexpr[bool] = False, rows_left: Int32 = 0, + FUSED: cutlass.Constexpr[bool] = False, ): - """With EDGE, only the thread's first rows_left rows are stored.""" + """With EDGE, only the thread's first rows_left rows are stored. FUSED: an activation is fused + (see _bf16_nan_magnitudes_to_zero).""" tidx, _, _ = cute.arch.thread_idx() MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 @@ -91,7 +128,12 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) - reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) + if cutlass.const_expr(FUSED): + reciprocal = _bf16_mx_scale_reciprocal_fused( + extract_mx_block_amax(rX_row_i32), MAX_EXPONENT + ) + else: + reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) reciprocal_pair = reciprocal | (reciprocal << Uint32(16)) for pack_idx in cutlass.range_constexpr( cute.size(rO_row_u32) @@ -202,8 +244,11 @@ def _reduce_partial_colwise_amaxes( ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FUSED: cutlass.Constexpr[bool] = False, ) -> cute.Tensor: """ + FUSED: an activation is fused (see _bf16_nan_magnitudes_to_zero). + sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) that contains each threads' own columnwise amax of its ELEMENTS_X_PER_THREAD columns @@ -278,6 +323,8 @@ def _reduce_partial_colwise_amaxes( for scratch_row_idx in cutlass.range_constexpr(1, CTA_THREADS_Y): amaxX2 = abs_max_x2_bf16(amaxX2, sAmaxs_col_frag_i32_rows[scratch_row_idx]) + if cutlass.const_expr(FUSED): + amaxX2 = Int32(_bf16_nan_magnitudes_to_zero(Uint32(amaxX2))) # Obtain the columnwise scale reciprocal for this column pair sS_col_i32[tidx, pair_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) cute.arch.sync_threads() @@ -404,8 +451,10 @@ def _quantize_colwise( EDGE: cutlass.Constexpr[bool] = False, rows_left: Int32 = 0, cols_ok: Boolean = True, + FUSED: cutlass.Constexpr[bool] = False, ): - """Reduce columnwise amaxes across the CTA, write the columnwise scales, and quantize rX with them.""" + """Reduce columnwise amaxes across the CTA, write the columnwise scales, and quantize rX with them. + FUSED: an activation is fused (see _bf16_nan_magnitudes_to_zero).""" tidx, _, _ = cute.arch.thread_idx() sAmaxs_col_scratch = _gather_partial_colwise_amaxes( @@ -424,6 +473,7 @@ def _quantize_colwise( ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, DTYPE=DTYPE, FP8_DTYPE=FP8_DTYPE, + FUSED=FUSED, ) # Copy the reduced columnwise scale reciprocals to RMEM. @@ -921,6 +971,7 @@ def quantize_mxfp8_gmem_tile_bidimensional( FP8_DTYPE=FP8_DTYPE, EDGE=EDGE, rows_left=rows_left, + FUSED=WITH_ACT or WITH_DACT, ) if cutlass.const_expr(COLWISE): @@ -989,6 +1040,7 @@ def quantize_mxfp8_gmem_tile_bidimensional( EDGE=EDGE, rows_left=rows_left, cols_ok=cols_ok, + FUSED=WITH_ACT or WITH_DACT, ) return thread_amax @@ -1032,11 +1084,12 @@ def takes_scaled_in_f32(cfg: MXFP8QuantizeConfig) -> bool: faster than the general kernel. GB200, geometric mean speedup over TE main, general -> this kernel: FP16 GeLU 1.46 -> 1.62, SiLU 1.41 -> 1.69, dGeLU 1.56 -> 1.69, dSiLU 1.28 -> 1.76, dbias_dSiLU 1.48 -> 1.74; FP32 SiLU 1.32 -> 1.54, dGeLU 1.38 -> 1.52, - dSiLU 1.10 -> 1.85, dbias_dSiLU 1.18 -> 1.53. dbias_dGeLU (FP16 1.75 -> 1.50, FP32 - 1.47 -> 1.25) and FP32 GeLU (1.20 -> 1.16) keep the general kernel.""" + dSiLU 1.10 -> 1.85, dbias_dSiLU 1.18 -> 1.53, dbias_dGeLU 1.47 -> 1.25 (CUDA 1.30; the + general kernel's FP32 dbias differs from CUDA's in the last bit, this kernel's matches). + FP16 dbias_dGeLU (1.75 -> 1.50) and FP32 GeLU (1.20 -> 1.16) keep the general kernel.""" if not (cfg.WITH_ACT or cfg.WITH_DACT): return False - if cfg.WITH_DBIAS and cfg.ACTIVATION == "dgelu": + if cfg.WITH_DBIAS and cfg.ACTIVATION == "dgelu" and cfg.DTYPE is cutlass.Float16: return False return not (cfg.DTYPE is cutlass.Float32 and cfg.ACTIVATION == "gelu") diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py index 62b4fe61f90..d0ff6fcbcd3 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -220,6 +220,10 @@ def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: rX_tile_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) rO_u32 = cute.recast_tensor(rO, Uint32) amax = extract_mx_block_amax(rX_tile_i32) + # With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's + # generic kernel; cast_rowwise.cuh, like the plain cast here, gives it the NaN scale. + if cutlass.const_expr(ACTIVATION is not None): + amax = cute.arch.fmax(amax, Float32(0.0)) exponent = cvt_f32_to_fp8e8m0fnu(amax * MAX_NORM_RCP) # The even lane of each pair stores its block's scale byte, as cast_rowwise.cuh does. diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py index 325accdd41c..b12576f0e8a 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -96,6 +96,8 @@ def _quantize_rowwise_mxfp8( SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, WITH_ACT: cutlass.Constexpr[bool] = False, WITH_DACT: cutlass.Constexpr[bool] = False, + # An activation is fused, even if this pass reads it cached (ACTIVATION is then None). + FUSED: cutlass.Constexpr[bool] = False, WITH_DBIAS: cutlass.Constexpr[bool] = False, dbias_acc: Optional[cute.Tensor] = None, # only needed when WITH_DBIAS is True ): @@ -289,6 +291,11 @@ def _quantize_rowwise_mxfp8( if cutlass.const_expr(FUSE_RELU): amax_r = cute.arch.fmax(amax_r, Float32(0.0)) # If relu, the amax is at least 0 + # With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's + # generic kernel (0-seeded fmaxf); CUDA's cast kernels, like the plain cast here, give it the + # NaN scale. + if cutlass.const_expr(ACTIVATION is not None or FUSED): + amax_r = cute.arch.fmax(amax_r, Float32(0.0)) biased_exp_r = cvt_f32_to_fp8e8m0fnu(amax_r * MAX_NORM_RCP) # mS_row_stage has logical shape (32, 2) and we have 64 threads where each is mapped to one scale factor @@ -354,8 +361,10 @@ def _quantize_colwise_mxfp8( WITH_DACT: cutlass.Constexpr[bool] = False, WITH_DBIAS: cutlass.Constexpr[bool] = False, CACHE_ACTIVATION: cutlass.Constexpr[bool] = False, # cache post-activation values to sX_tile + dbias_running: Float32 = 0.0, # the column's dbias sum over the CTA's earlier stages ): - """Quantize one SMEM tile colwise to MXFP8 (per-column 32-elt block scales); returns (amax, dbias_partial).""" + """Quantize one SMEM tile colwise to MXFP8 (per-column 32-elt block scales); returns (amax, dbias), + dbias being dbias_running plus this stage's rows, added in row order.""" tidx, _, _ = cute.arch.thread_idx() _, tv_layout = cute.make_layout_tv( @@ -374,7 +383,8 @@ def _quantize_colwise_mxfp8( FUSE_RELU = cutlass.const_expr(ACTIVATION == "relu") and not WITH_DBIAS # Keep input in half precision format if possible USE_HALF_PRECISION = is_packed16(DTYPE) and (ACTIVATION is None or FUSE_RELU) - dbias_partial = Float32(0.0) + # One running sum down the column across the stages, as the CUDA kernels add the chunk's rows. + dbias_partial = Float32(dbias_running) if cutlass.const_expr(USE_HALF_PRECISION): max_scalar = max_scalar_f16 if DTYPE is cutlass.Float16 else max_scalar_bf16 @@ -465,6 +475,11 @@ def _quantize_colwise_mxfp8( # Irregular shapes: skip when this stage's row range or this thread's # column lies past the input extents. TILE_Y == MXFP8_BLOCK_SCALING_SIZE so each stage # is exactly one scale-row; valid iff `tile_row_start < M`. + # With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's + # generic kernel (0-seeded fmaxf); CUDA's cast kernels, like the plain cast here, give it the + # NaN scale. + if cutlass.const_expr(ACTIVATION is not None): + amax_c = cute.arch.fmax(amax_c, Float32(0.0)) biased_exp_c = cvt_f32_to_fp8e8m0fnu(amax_c * MAX_NORM_RCP) # If the input shape is divisible by the tile size * tile num, then we won't access OOB regions for sure if cutlass.const_expr(SKIP_SCALE_BOUNDS): @@ -522,8 +537,10 @@ def _dbias_reduction_colwise( TILE_X: cutlass.Constexpr[int], WITH_DACT: cutlass.Constexpr[bool] = False, CACHE_ACTIVATION: cutlass.Constexpr[bool] = False, # cache post-activation values to sX_tile + dbias_running: Float32 = 0.0, # the column's dbias sum over the CTA's earlier stages ): - """Reduce one SMEM tile along rows (the dbias direction) and possibly cache dact values, with no quantization.""" + """Reduce one SMEM tile along rows (the dbias direction) and possibly cache dact values, with no quantization. + Returns dbias_running plus this stage's rows, added in row order.""" tidx, _, _ = cute.arch.thread_idx() # Same TV layout as _quantize_colwise_mxfp8: thread tidx owns column tidx, all TILE_Y rows. @@ -533,7 +550,7 @@ def _dbias_reduction_colwise( ) sX_thread = cute.composition(sX_tile, tv_layout)[tidx, None] - dbias_partial = Float32(0.0) + dbias_partial = Float32(dbias_running) if cutlass.const_expr(WITH_DACT): dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) sA_thread = cute.composition(sA_tile, tv_layout)[tidx, None] @@ -580,11 +597,13 @@ def quantize_mxfp8_smem_tile_general( DBIAS_REDUCTION_IN_COLWISE: cutlass.Constexpr[bool], COLWISE_DBIAS_REDUCTION_ONLY: cutlass.Constexpr[bool], DBIAS_REDUCTION_IN_ROWWISE: cutlass.Constexpr[bool], + dbias_running: Float32 = 0.0, ): """Quantize one SMEM tile in the directions enabled by ROWWISE / COLWISE; returns (amax, dbias). amax is this thread's amax over the tile (None unless WITH_AMAX). dbias is this thread's - column partial sum (None unless DBIAS_REDUCTION_IN_COLWISE or COLWISE_DBIAS_REDUCTION_ONLY); + column sum so far, dbias_running plus this tile's rows (None unless DBIAS_REDUCTION_IN_COLWISE + or COLWISE_DBIAS_REDUCTION_ONLY); with DBIAS_REDUCTION_IN_ROWWISE the partial sums are accumulated into rowwise_dbias_acc instead. """ amax = None @@ -613,6 +632,7 @@ def quantize_mxfp8_smem_tile_general( WITH_DACT=WITH_DACT, WITH_DBIAS=DBIAS_REDUCTION_IN_COLWISE, CACHE_ACTIVATION=CACHE_ACTIVATION, + dbias_running=dbias_running, ) # Either direction's amax already spans every element of the tile, so favor colwise # and skip the redundant comparison in the rowwise pass. @@ -631,6 +651,7 @@ def quantize_mxfp8_smem_tile_general( TILE_X=TILE_X, WITH_DACT=WITH_DACT, CACHE_ACTIVATION=CACHE_ACTIVATION, + dbias_running=dbias_running, ) # If we cache the activation in shared memory, we need to ensure that all threads have finished writing to the shared memory # from the columnwise pass before any thread reads from it in the rowwise pass. @@ -657,6 +678,7 @@ def quantize_mxfp8_smem_tile_general( SKIP_SCALE_BOUNDS=SKIP_SCALE_BOUNDS, WITH_ACT=WITH_ACT and not CACHE_ACTIVATION, WITH_DACT=WITH_DACT and not CACHE_ACTIVATION, + FUSED=WITH_ACT or WITH_DACT, WITH_DBIAS=DBIAS_REDUCTION_IN_ROWWISE, dbias_acc=rowwise_dbias_acc, ) @@ -1180,11 +1202,10 @@ class DactStorage: rowwise_dbias_acc[c] = Float32(0.0) block_dbias = Float32(0.0) # Prepare thread-level register for columnwise dbias reduction. - # Each thread will process two (MXFP8_BLOCK_SCALING_SIZE, 1) columns in two stages, and in each stage the thread will reduce the - # (after dact applied) column to (1,) and add to this register. - # Then this partial sum scalar will be written to the GMEM workspace buffer directly. - if cutlass.const_expr(self.DBIAS_REDUCTION_IN_COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY): - block_dbias = Float32(0.0) + # Each thread will process two (MXFP8_BLOCK_SCALING_SIZE, 1) columns in two stages; each stage continues + # this running sum with its (after dact applied) column, in row order, as the CUDA kernels do. + # Then this sum will be written to the GMEM workspace buffer directly. + block_dbias = Float32(0.0) # Consumer: all warps fetch from the pipeline, process its tile, and issue a new load to the tile buffer it just consumed for tile_idx in cutlass.range(num_tiles, unroll=1): @@ -1256,13 +1277,14 @@ class DactStorage: DBIAS_REDUCTION_IN_COLWISE=self.DBIAS_REDUCTION_IN_COLWISE, COLWISE_DBIAS_REDUCTION_ONLY=self.COLWISE_DBIAS_REDUCTION_ONLY, DBIAS_REDUCTION_IN_ROWWISE=self.DBIAS_REDUCTION_IN_ROWWISE, + dbias_running=block_dbias, ) if cutlass.const_expr(cfg.WITH_AMAX): per_thread_amax = cute.arch.fmax(per_thread_amax, amax) if cutlass.const_expr( self.DBIAS_REDUCTION_IN_COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY ): - block_dbias += dbias + block_dbias = dbias # Make the shared-memory writes visible to the TMA's async proxy before the TMA reads them. cute.arch.fence_proxy( From 9fb999ef122fc2553e70f853c1bfad932bbaab6e Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 28 Oct 2026 14:15:59 -0700 Subject: [PATCH 27/29] [Common] Drop the CUDA register-tiled activation kernel; keep the gated one The CuTeDSL backend's register-resident kernel now covers GeLU, SiLU, dGeLU and dSiLU (tables, dbias, amax, tails, BF16/FP16/FP32), so the CUDA copy of it in cast/mxfp8/regtile is removed. The CUDA backend computes these in the generic kernel again; with NVTE_ACTIVATION_TABLE=1 it still follows the table specification (kComputedTable), so both backends give the same bytes. The table helpers only that kernel used (packed lookups, kComputed, kSignShift) go with it. CuTeDSL has no gated path yet, so the register-tiled SwiGLU/dSwiGLU kernel stays, moved with the helpers it needs to cast/mxfp8/specialized/gated_bidim.cuh next to the other register-resident MXFP8 kernels. Its speed is unchanged. Signed-off-by: Teddy Do --- .../cpp/operator/test_cast_mxfp8_act_exact.cu | 14 +- .../common/cast/mxfp8/gated_mxfp8.cuh | 10 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 27 - .../mxfp8/regtile/quantize_mxfp8_regtile.cuh | 1005 ----------------- .../gated_bidim.cuh} | 142 ++- .../common/util/activation_table.cuh | 100 +- 6 files changed, 145 insertions(+), 1153 deletions(-) delete mode 100644 transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh rename transformer_engine/common/cast/mxfp8/{regtile/quantize_gated_mxfp8_regtile.cuh => specialized/gated_bidim.cuh} (72%) diff --git a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu index 7daa1fc7928..d29ebd98d0d 100644 --- a/tests/cpp/operator/test_cast_mxfp8_act_exact.cu +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -6,9 +6,9 @@ // The fused MXFP8 activation quantize must round exactly as if the activation were computed per // element, truncated to the input type, and then quantized on its own. This compares the fused -// kernels (the register-tiled kernel where it applies, the generic kernel elsewhere) bit for bit -// against that composition: the activation is computed on the GPU, and the result is quantized by -// the generic cast-only kernel (a noop tensor keeps the call off the specialized cast kernels). +// kernels (CUDA, or CuTeDSL with NVTE_ENABLE_CUTEDSL_BACKEND=1) bit for bit against that +// composition: the activation is computed on the GPU, and the result is quantized by the generic +// cast-only kernel (a noop tensor keeps the call off the specialized cast kernels). // dbias is compared against a host sum in the order the kernels accumulate it, and the tensor amax // against the largest magnitude of the activation. // @@ -450,15 +450,15 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kBFloat16, DType::kFloat16, DType::kFloat32), ::testing::Values(0, 1, 2), ::testing::Bool(), ::testing::Values(std::pair{1024, 2048}, - // Enough row tiles that the register-resident kernels - // walk several per CTA. + // Enough row tiles that the CuTeDSL register-resident + // kernel walks several per CTA. std::pair{4096, 2048}, // A row tail on 256-column tiles: 31 row tiles, the // last of 16 rows, and an odd number for dbias. std::pair{976, 2048}, std::pair{544, 2080}, - // At least 2^23 elements (the register-tiled forward - // without tables), with a row and a column tail. + // A row tail and a column tail ending inside an MX + // block. std::pair{4127, 2064}, // A dbias workspace of 130 rows, past the dbias // reduction's 128-row pipeline. diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index a5f215210b5..c8d5615967f 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -20,7 +20,7 @@ #include "../../util/math.h" #include "../../util/ptx_arch_spec.cuh" #include "../../utils.cuh" -#include "regtile/quantize_gated_mxfp8_regtile.cuh" +#include "specialized/gated_bidim.cuh" #include "swizzle.cuh" namespace transformer_engine { @@ -851,16 +851,16 @@ void quantize_gated(const Tensor &gated_input, const Tensor &grad, Tensor *outpu } } - // regtile::gated replaces the kernel below for SwiGLU and its gradient from BF16 to - // E4M3 or E5M2 in both directions (see regtile::gated::can_use), writing the same + // gated_bidim replaces the kernel below for SwiGLU and its gradient from BF16 to + // E4M3 or E5M2 in both directions (see gated_bidim::can_use), writing the same // bytes. if constexpr (std::is_same_v && (std::is_same_v || std::is_same_v) && std::is_same_v) { if constexpr (ActOP == &silu && (!IS_BWD || DActOP == &dsilu)) { - if (regtile::gated::can_use(gated_input, grad, *output, IS_BWD)) { - regtile::gated::quantize( + if (gated_bidim::can_use(gated_input, grad, *output, IS_BWD)) { + gated_bidim::quantize( gated_input, grad, output, stream); return; } diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index c6f072f1259..d2126a7da57 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -23,7 +23,6 @@ #include "../../util/ptx_arch_spec.cuh" #include "../../utils.cuh" #include "../core/common.cuh" -#include "regtile/quantize_mxfp8_regtile.cuh" #include "specialized/cast_bidim.cuh" #include "specialized/cast_rowwise.cuh" #include "specialized/quantize_mxfp8.cuh" @@ -965,32 +964,6 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, return; } - // regtile::quantize replaces the generic kernel below for the requests it - // implements: bidimensional BF16, FP16 or FP32 -> E4M3 or E5M2 with fused GeLU, SiLU, dGeLU or - // dSiLU, read from an activation table or computed, the derivatives optionally with - // dbias, with or without amax, cols % 16 == 0, 16-byte aligned inputs (see - // regtile::kImplements and regtile::can_use). It keeps the tile in registers instead - // of staging it through shared memory and writes the same bytes, the same amax and - // the same dbias workspace the generic kernel writes. - if constexpr ((std::is_same_v || - (std::is_same_v && - regtile::kTakesFP16) || - std::is_same_v) && - (std::is_same_v || std::is_same_v) && - regtile::kImplements) { - if (regtile::can_use( - input, act_input, *output, *noop, use_2d_quantization, activation_lut)) { - regtile::quantize( - input, act_input, output, activation_lut, workspace_ptr, stream); - if constexpr (IS_DBIAS) { - common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, - stream); - } - return; - } - } - alignas(64) CUtensorMap tensor_map_input{}; alignas(64) CUtensorMap tensor_map_act_input{}; alignas(64) CUtensorMap tensor_map_output_rowwise{}; diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh deleted file mode 100644 index cbec0cfe140..00000000000 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_mxfp8_regtile.cuh +++ /dev/null @@ -1,1005 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file quantize_mxfp8_regtile.cuh - * \brief Register-resident MXFP8 bidimensional quantize fused with GeLU or SiLU, or - * their derivatives, optionally with dbias. - * - * Rowwise and colwise MX scales need an amax over the same data in two - * directions. The generic kernel stages each tile through shared memory and - * walks it twice. Here a CTA covers whole 32-row bands, the colwise block - * height, so the tile is read once into registers and both directions are - * quantized from there; shared memory holds only the cross-warp colwise - * partials, the row scales and the activation table. - * - * The activation or derivative is read from the table the caller passes (see - * util/activation_table.cuh), copied into shared memory once per CTA, or, without a - * table, computed with util/math.h. Either way this kernel produces exactly the bytes the - * generic kernel produces for the same request. - * - * It replaces the generic kernel, not the specialized ones: cast-only requests - * are served before the generic dispatch by the specialized kernels, one of - * which (specialized/cast_bidim.cuh) is also register-resident. Anything this - * kernel does not implement (see can_use) takes the generic kernel. - */ - -#ifndef TRANSFORMER_ENGINE_QUANTIZE_MXFP8_REGTILE_CUH_ -#define TRANSFORMER_ENGINE_QUANTIZE_MXFP8_REGTILE_CUH_ - -#include - -#include -#include -#include -#include - -#include "../../../common.h" -#include "../../../util/activation_table.cuh" -#include "../../../util/cuda_runtime.h" -#include "../../../util/math.h" -#include "../../../util/ptx.cuh" -#include "../../../util/ptx_arch_spec.cuh" -#include "../../../utils.cuh" -#include "../swizzle.cuh" - -namespace transformer_engine { -namespace dispatch { -namespace mxfp8 { -namespace regtile { - -namespace ptx = transformer_engine::ptx; - -namespace { - -// The datapath works on column pairs of the input type (Datapath::Word): one 32-bit word for -// BF16 and FP16, two FP32 registers for FP32. - -/*! \brief Elementwise max(|a|, |b|) on a packed BF16 pair, as a raw word. */ -__device__ __forceinline__ unsigned abs_max_bf16x2(unsigned a, unsigned b) { - ptx::bf16x2 d; - ptx::abs_max_2x(d, reinterpret_cast(a), - reinterpret_cast(b)); - return reinterpret_cast(d); -} - -/*! \brief Widen a BF16 pair word to FP32, low half to x. Same result as ptx::up_cast, whose - * volatile asm the compiler cannot schedule as freely in the derivative loop. */ -__device__ __forceinline__ ptx::floatx2 widen_bf16x2(const unsigned v) { - return {__uint_as_float(v << 16), __uint_as_float(v & 0xffff0000u)}; -} - -/*! \brief Round an FP32 pair to a BF16 pair word, low lane in the low half. */ -__device__ __forceinline__ unsigned to_bf16x2(const ptx::floatx2& v) { - const __nv_bfloat162 r = __floats2bfloat162_rn(v.x, v.y); - return reinterpret_cast(r); -} - -/*! \brief The reciprocal of an E8M0 scale, 2^(127-e), broadcast to a BF16 pair. */ -__device__ __forceinline__ unsigned e8m0_to_bf16x2_reciprocal(unsigned e) { - const ptx::bf16x2 r = ptx::exp2f_rcp_2x(static_cast(e)); - return reinterpret_cast(r); -} - -/*! \brief E8M0 scales for two block amaxes, packed as ptx::float_to_e8m0_2x. - * - * The generic kernel's amax is a 0-seeded max that ignores NaN, so a block of - * only NaNs gets amax 0; the fmaxf reproduces that for the data-seeded maxima - * here. The conversion saturates, so the scale never reaches 255. - */ -template -__device__ __forceinline__ unsigned amax_to_e8m0_2x(float amax_hi, float amax_lo) { - constexpr float kMaxNormRcp = Quantized_Limits::max_norm_rcp; - return ptx::float_to_e8m0_2x(fmaxf(amax_hi * kMaxNormRcp, 0.0f), - fmaxf(amax_lo * kMaxNormRcp, 0.0f)); -} - -/*! \brief Table's activation computed with util/math.h instead of read from a table: - * the generic kernel's output without tables (NVTE_ACTIVATION_TABLE=0). No table is - * copied. - */ -template -struct ComputedTable : activation_table::Table { - static constexpr bool kComputed = true; - static constexpr size_t kBytes = 16; // unused shared-memory slot -}; - -template -struct Datapath; // below - -/*! \brief The derivative of an input pair, from the table in shared memory or computed. */ -template -__device__ __forceinline__ ptx::floatx2 dact_word(const Word& x, - const unsigned char* __restrict__ table) { - if constexpr (TableT::kComputed) { - const ptx::floatx2 v = Datapath::widen(x); - return {TableT::op(v.x), TableT::op(v.y)}; - } else { - const float2 d = activation_table::derivative_2x(x, table); - return {d.x, d.y}; - } -} - -/*! \brief Start copying a table of kBytes into shared memory with one bulk copy. - * - * Called by every thread of the CTA; contains a __syncthreads. - */ -template -__device__ __forceinline__ void load_table_async(unsigned char* smem, uint64_t* barrier, - const void* table, const int tid) { - static_assert(kBytes % 16 == 0, "A bulk copy moves whole 16-byte units."); - if (tid == 0) { - ptx::mbarrier_init(barrier, 1); - ptx::fence_proxy_async_shared_cta(); - } - __syncthreads(); - if (tid == 0) { - ptx::mbarrier_arrive_expect_tx(barrier, kBytes); - ptx::cp_async_bulk_tensor_1d_global_to_shared(reinterpret_cast(smem), - reinterpret_cast(table), kBytes, - barrier); - } -} - -// Tile geometry. A CTA covers 256 columns and walks 32-row sub-tiles; 32 rows -// is the colwise MX block height, so the colwise reduction closes in the CTA. - -//! Columns a CTA covers. -constexpr int kTileCols = 256; -//! Elements in one MX block, and equivalently the colwise block height. -constexpr int kRowsPerMxBlock = 32; -//! MX groups (32 columns each) across the tile. -constexpr int kMxGroupsPerRow = kTileCols / 32; -//! A register-resident tile spreads its 256 columns over the 32 lanes of a warp. -constexpr int kColsPerLane = kTileCols / 32; -//! ...which a lane holds as input-type pairs (Datapath::Word). -constexpr int kWordsPerLane = kColsPerLane / 2; -//! ...read, for 16-bit inputs, as one 128-bit vector per row (FP32 inputs take two; see -//! kVecsPerLaneRow in quantize_regtile). -constexpr int kVecLoadsPerRow = kWordsPerLane / 4; -//! Lanes that have to cooperate to cover one 32-column MX group. -constexpr int kLanesPerMxGroup = 32 / kColsPerLane; - -/*! \brief Scale two BF16 pairs, each by its own reciprocal scale pair, to four OType (E4M3 or - * E5M2). */ -template -__device__ __forceinline__ unsigned quantize_x4(unsigned in01, unsigned in23, unsigned scale01, - unsigned scale23) { - ptx::FPx4 out; - ptx::mul_cvt_4x(out, reinterpret_cast(in01), - reinterpret_cast(scale01), - reinterpret_cast(in23), - reinterpret_cast(scale23)); - return reinterpret_cast(out); -} - -/*! \brief Store one lane's quantized row: kColsPerLane FP8 bytes, one STG.64. */ -__device__ __forceinline__ void store_quantized(unsigned char* p, const unsigned* v) { - static_assert(kColsPerLane == 8, "One STG.64 covers exactly eight FP8 bytes."); - *(uint2*)p = *(const uint2*)v; -} - -/*! \brief The input-type-specific part of the datapath. A 32-bit word holds a column pair of - * IType; Scale is the reciprocal scale of such a pair, in the form its quantize - * instruction takes. - */ -template -struct Datapath; - -/*! \brief Helpers on a packed pair of 16-bit values (one 32-bit word). */ -struct PackedPair16 { - using Word = unsigned; - //! The other value of the pair in each half. - __device__ __forceinline__ static Word swap_halves(Word v) { return __byte_perm(v, v, 0x1032); } - __device__ __forceinline__ static Word shfl_xor(Word v, int lane_mask) { - return __shfl_xor_sync(0xffffffffu, v, lane_mask); - } - __device__ __forceinline__ static Word clear_sign(Word v) { return v & 0x7fff7fffu; } -}; - -template <> -struct Datapath : PackedPair16 { - //! A BF16 pair. 2^(127-e) is a power of two inside BF16's range, so the BF16 multiply in - //! mul_cvt_4x is exact, like the generic kernel's FP32 one. - using Scale = unsigned; - __device__ __forceinline__ static unsigned abs_max(unsigned a, unsigned b) { - return abs_max_bf16x2(a, b); - } - __device__ __forceinline__ static ptx::floatx2 widen(unsigned v) { return widen_bf16x2(v); } - __device__ __forceinline__ static unsigned narrow(const ptx::floatx2& v) { return to_bf16x2(v); } - //! The reciprocal of the scale e, for both columns of a pair. - __device__ __forceinline__ static Scale row_rcp(unsigned e) { - return e8m0_to_bf16x2_reciprocal(e); - } - //! The reciprocals of a column pair's scales, low byte for the low column. - __device__ __forceinline__ static Scale col_rcp(unsigned scale_pair) { - const ptx::bf16x2 r = ptx::exp2f_rcp_2x_per_lane(__byte_perm(scale_pair, 0, 0x4140)); - return reinterpret_cast(r); - } - template - __device__ __forceinline__ static unsigned quantize(unsigned in01, unsigned in23, Scale s01, - Scale s23) { - return quantize_x4(in01, in23, s01, s23); - } -}; - -template <> -struct Datapath : PackedPair16 { - //! An FP32 pair: 2^(127-e) is often outside FP16's range, so the multiply is done in FP32, - //! as the generic kernel does it (mul_cvt_2x with an FP32 scale pair). - using Scale = ptx::floatx2; - __device__ __forceinline__ static unsigned abs_max(unsigned a, unsigned b) { - ptx::fp16x2 d; - ptx::abs_max_2x(d, reinterpret_cast(a), - reinterpret_cast(b)); - return reinterpret_cast(d); - } - __device__ __forceinline__ static ptx::floatx2 widen(unsigned v) { - const float2 f = __half22float2(reinterpret_cast(v)); - return {f.x, f.y}; - } - __device__ __forceinline__ static unsigned narrow(const ptx::floatx2& v) { - const __half2 h = __floats2half2_rn(v.x, v.y); - return reinterpret_cast(h); - } - __device__ __forceinline__ static Scale row_rcp(unsigned e) { - const float r = ptx::exp2f_rcp(static_cast(e)); - return {r, r}; - } - __device__ __forceinline__ static Scale col_rcp(unsigned scale_pair) { - return {ptx::exp2f_rcp(static_cast(scale_pair & 0xffu)), - ptx::exp2f_rcp(static_cast((scale_pair >> 8) & 0xffu))}; - } - template - __device__ __forceinline__ static unsigned quantize(unsigned in01, unsigned in23, Scale s01, - Scale s23) { - ptx::FPx2 lo, hi; - ptx::mul_cvt_2x(lo, reinterpret_cast(in01), s01); - ptx::mul_cvt_2x(hi, reinterpret_cast(in23), s23); - return static_cast(reinterpret_cast(lo)) | - (static_cast(reinterpret_cast(hi)) << 16); - } -}; - -template <> -struct Datapath { - //! A pair of FP32 values, two registers: the tile is twice as large as for 16-bit inputs, - //! so the CTA has twice the threads, each holding half the rows (see TileConfig). - using Word = ptx::floatx2; - using Scale = ptx::floatx2; - //! max(|a|, |b|) per lane. fmaxf ignores NaN, as the generic kernel's 0-seeded fmaxf does; - //! a block of only NaNs still gives NaN here, which amax_to_e8m0_2x maps to the scale of 0. - __device__ __forceinline__ static Word abs_max(const Word& a, const Word& b) { - return {fmaxf(fabsf(a.x), fabsf(b.x)), fmaxf(fabsf(a.y), fabsf(b.y))}; - } - __device__ __forceinline__ static Word swap_halves(const Word& v) { return {v.y, v.x}; } - __device__ __forceinline__ static Word shfl_xor(const Word& v, int lane_mask) { - return {__shfl_xor_sync(0xffffffffu, v.x, lane_mask), - __shfl_xor_sync(0xffffffffu, v.y, lane_mask)}; - } - __device__ __forceinline__ static Word clear_sign(const Word& v) { - return {fabsf(v.x), fabsf(v.y)}; - } - __device__ __forceinline__ static ptx::floatx2 widen(const Word& v) { return v; } - //! The generic kernel's rounding to the input type is a no-op for FP32. - __device__ __forceinline__ static Word narrow(const ptx::floatx2& v) { return v; } - __device__ __forceinline__ static Scale row_rcp(unsigned e) { - const float r = ptx::exp2f_rcp(static_cast(e)); - return {r, r}; - } - __device__ __forceinline__ static Scale col_rcp(unsigned scale_pair) { - return {ptx::exp2f_rcp(static_cast(scale_pair & 0xffu)), - ptx::exp2f_rcp(static_cast((scale_pair >> 8) & 0xffu))}; - } - //! FP32 multiply, then convert, as the generic kernel does (mul_cvt_2x on an FP32 pair). - template - __device__ __forceinline__ static unsigned quantize(const Word& in01, const Word& in23, Scale s01, - Scale s23) { - ptx::FPx2 lo, hi; - ptx::mul_cvt_2x(lo, in01, s01); - ptx::mul_cvt_2x(hi, in23, s23); - return static_cast(reinterpret_cast(lo)) | - (static_cast(reinterpret_cast(hi)) << 16); - } -}; - -// Minimum resident CTAs per SM, the second __launch_bounds__ argument, capped -// at what the SM can hold: sm_75 and sm_107 have half the threads per SM of sm_100. -constexpr int kMinBlocksAct = 8; -constexpr int kMinBlocksDact = 6; -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 750 || __CUDA_ARCH__ == 1070) -constexpr int kMaxThreadsPerSm = 1024; -#else -constexpr int kMaxThreadsPerSm = 2048; -#endif - -// Minimum resident CTAs per SM with dbias: the 32 KB staging buffer allows about four. -constexpr int kMinBlocksDbias = 4; - -template -struct TileConfig { - static_assert(IS_ACT != IS_DACT, "Exactly one of the activation and its derivative is fused."); - //! FP32 inputs take twice the registers per value, so twice the threads share a tile: each - //! thread holds the same number of registers of it, and an SM the same number of threads. - static constexpr int kThreadScale = sizeof(IType) == 4 ? 2 : 1; - static constexpr int kThreadsPerCta = (IS_DACT ? 256 : 128) * kThreadScale; - static constexpr int kWarpsPerCta = kThreadsPerCta / 32; - static constexpr int kRowsPerWarp = 32 / kWarpsPerCta; - // The rowwise half is quantized before the cross-warp barrier and the row - // scales are drained inside the colwise barrier interval, saving one barrier - // per 32-row tile. A tile too narrow to leave idle threads hands the drain to - // the threads that just folded the colwise partials. - static constexpr int kDrainTidBase = - (kThreadsPerCta >= kTileCols / 2 + kTileCols / 4) ? kTileCols / 2 : 0; - //! Rows loaded per group: one with FP32, whose row is two 16-byte vectors per lane. - static constexpr int kRowsInFlight = sizeof(IType) == 4 ? 1 : 2; - static constexpr int kRequestedMinBlocks = - (kDbias ? kMinBlocksDbias : (IS_DACT ? kMinBlocksDact : kMinBlocksAct)) / kThreadScale; - static constexpr int kMinBlocksPerSm = kRequestedMinBlocks < kMaxThreadsPerSm / kThreadsPerCta - ? kRequestedMinBlocks - : kMaxThreadsPerSm / kThreadsPerCta; - static constexpr int kTableBytes = static_cast(TableT::kBytes); - //! Dynamic shared memory: with dbias, the band's FP32 products before rounding. - static constexpr int kDbiasStageBytes = kDbias ? kRowsPerMxBlock * kTileCols * 4 : 0; - static_assert(!kDbias || (IS_DACT && kThreadsPerCta >= kTileCols), - "dbias: one thread per column sums the column."); -}; - -// Register-resident tiling: the 32x256 tile stays in registers; only the -// colwise partials go through shared memory. -// kTails: the grid may cover elements past M or K (see shape_supported). Without tails the -// bounds checks compile away, which matters at the kernel's register budget. -// -// kDbias: also write the per-64-row-chunk column sums of the FP32 products (before rounding to -// BF16) into dbias_workspace[chunk][K], as the generic kernel does; common::reduce_dbias then -// adds the chunks. The generic kernel sums each column of a chunk strictly in row order, so a -// CTA here takes the chunk's two 32-row bands back to back, stages each band's products in -// shared memory, and one thread per column adds them in row order. -// -// kAmax: also fold the amax of the output into *amax_ptr. -template -__device__ __forceinline__ void quantize_regtile( - const uint4* __restrict__ input, const uint4* __restrict__ act_input, - unsigned char* __restrict__ rowwise_out, unsigned char* __restrict__ rowwise_scales, - unsigned char* __restrict__ colwise_out, unsigned char* __restrict__ colwise_scales, int M, - int K, int rowwise_scale_stride, int colwise_scale_stride, int iters, const void* table, - float* __restrict__ dbias_workspace, float* __restrict__ amax_ptr) { - using C = TileConfig; - using DP = Datapath; - using Scale = typename DP::Scale; - using Word = typename DP::Word; - constexpr int kWarpsPerCta = C::kWarpsPerCta, kRowsPerWarp = C::kRowsPerWarp, - kRowsInFlight = C::kRowsInFlight; - //! Input values per 16-byte vector, words (column pairs) per vector, vectors per lane row. - constexpr int kValsPerVec = 16 / sizeof(IType); - constexpr int kWordsPerVec = 16 / sizeof(Word); - constexpr int kVecsPerLaneRow = kWordsPerLane / kWordsPerVec; - - __shared__ __align__(16) Word col_amax_partials[kWarpsPerCta][kTileCols / 2]; - __shared__ __align__(16) Scale col_scale_rcp_smem[kTileCols / 2]; - __shared__ __align__(8) unsigned char row_scale_bytes[kRowsPerMxBlock * kMxGroupsPerRow]; - __shared__ __align__(16) unsigned char table_smem[C::kTableBytes]; - __shared__ __align__(8) uint64_t table_barrier; - extern __shared__ __align__(16) float dbias_stage[]; // [kRowsPerMxBlock][kTileCols], dbias only - float dbias_partial = 0.0f; - - const int tid = threadIdx.x; - const int lane = tid & 31; - const int warp = tid >> 5; - const int col0 = blockIdx.x * kTileCols; - const int vecs_per_row = K / kValsPerVec; - // Tails: the last column tile may end inside the tile, and row bands may end past M. Out-of-range - // elements load as 0, which every activation maps to 0, so they never raise an amax. The - // generic kernel writes scales for its whole 64x64 tiles, and so does this one: rowwise scales - // for every row of the grid's bands, for the groups that start inside the tensor; colwise - // scales for the columns below K rounded up to 64. - const bool full_width = !kTails || col0 + kTileCols <= K; - const bool lane_in_cols = col0 + kColsPerLane * lane < K; - const int colwise_scale_cols = (K + 63) / 64 * 64; - - // The table copy overlaps the first input loads; see the wait below. - if constexpr (!TableT::kComputed) { - load_table_async(table_smem, &table_barrier, table, tid); - } - - float thread_amax = 0.0f; - -#pragma unroll 1 - for (int it = 0; it < iters; ++it) { - // Row blocks are taken grid-strided, so the CTAs on the same step cover one - // contiguous row band. - // With dbias a CTA takes the two bands of a 64-row chunk on consecutive steps. - const int row0 = kDbias - ? (((it >> 1) * gridDim.y + blockIdx.y) * 2 + (it & 1)) * kRowsPerMxBlock - : (it * gridDim.y + blockIdx.y) * kRowsPerMxBlock; - if constexpr (kDbias) { - if ((it & 1) == 0) dbias_partial = 0.0f; - } - // The band's work. A band inside the tensor runs it without bounds checks; only bands at - // the edges of the tensor (kEdge) pay for them. - const auto band = [&](auto edge) { - constexpr bool kEdge = decltype(edge)::value; - const int warp_row0 = row0 + warp * kRowsPerWarp; - // Whether this lane's elements of row warp_row0 + j are inside the tensor. - const auto in_bounds = [&](int j) { return !kEdge || (lane_in_cols && warp_row0 + j < M); }; - // This thread's kRowsPerWarp x kColsPerLane sub-tile, as input-type pairs. - Word tile_words[kRowsPerWarp * kWordsPerLane]; - const size_t load_offset = (size_t)(row0 + warp * kRowsPerWarp) * vecs_per_row + - col0 / kValsPerVec + (size_t)kVecsPerLaneRow * lane; - const uint4* input_ptr = input + load_offset; - const uint4* act_input_ptr = act_input + load_offset; - - // ---- load + activate, kRowsInFlight rows in flight at a time -------------------- -#pragma unroll - for (int h = 0; h < kRowsPerWarp / kRowsInFlight; ++h) { - uint4 x_vec[kRowsInFlight][kVecsPerLaneRow], grad_vec[kRowsInFlight][kVecsPerLaneRow]; - asm volatile("" ::: "memory"); -#pragma unroll - for (int t = 0; t < kRowsInFlight; ++t) -#pragma unroll - for (int n = 0; n < kVecsPerLaneRow; ++n) - x_vec[t][n] = in_bounds(h * kRowsInFlight + t) - ? act_input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n] - : make_uint4(0, 0, 0, 0); - if constexpr (IS_DACT) { -#pragma unroll - for (int t = 0; t < kRowsInFlight; ++t) -#pragma unroll - for (int n = 0; n < kVecsPerLaneRow; ++n) - grad_vec[t][n] = in_bounds(h * kRowsInFlight + t) - ? input_ptr[(size_t)(h * kRowsInFlight + t) * vecs_per_row + n] - : make_uint4(0, 0, 0, 0); - } - // Keep the load group's loads together, between the two empty asm statements: issued - // back to back, their latencies overlap. Otherwise register pressure leads the compiler - // to spread them between the previous group's table lookups, each then stalling on its - // own. - asm volatile("" ::: "memory"); - if constexpr (!TableT::kComputed) { - if (it == 0 && h == 0) ptx::mbarrier_wait_parity(&table_barrier, 0); - } - // Convert the load group into the BF16 pairs the epilogue quantizes. - if constexpr (!IS_DACT) { -#pragma unroll - for (int t = 0; t < kRowsInFlight; ++t) { -#pragma unroll - for (int n = 0; n < kVecsPerLaneRow; ++n) { - const Word* x_words = (const Word*)&x_vec[t][n]; - const int tile_word_base = (h * kRowsInFlight + t) * kWordsPerLane + n * kWordsPerVec; -#pragma unroll - for (int m = 0; m < kWordsPerVec; ++m) { - if constexpr (TableT::kComputed) { - // util/math.h, then rounded to BF16 as the generic kernel rounds it. - const ptx::floatx2 v = DP::widen(x_words[m]); - tile_words[tile_word_base + m] = DP::narrow({TableT::op(v.x), TableT::op(v.y)}); - } else { - tile_words[tile_word_base + m] = activation_table::forward_2x( - x_words[m], reinterpret_cast(table_smem)); - } - } - } - } - } else { -#define MXFP8_DACT_WORD(T, M) \ - { \ - const ptx::floatx2 product = \ - ptx::mul_2x(DP::widen(grad_words[M]), dact_word(x_words[M], table_smem)); \ - tile_words[(h * kRowsInFlight + (T)) * kWordsPerLane + (M)] = DP::narrow(product); \ - products[2 * (M)] = product.x; \ - products[2 * (M) + 1] = product.y; \ - } -#define MXFP8_DACT_ROW(T) \ - { \ - const Word* x_words = (const Word*)&x_vec[T][0]; \ - const Word* grad_words = (const Word*)&grad_vec[T][0]; \ - float products[2 * kWordsPerLane]; \ - MXFP8_DACT_WORD(T, 0) \ - MXFP8_DACT_WORD(T, 1) \ - MXFP8_DACT_WORD(T, 2) \ - MXFP8_DACT_WORD(T, 3) \ - if constexpr (kDbias) { \ - float4* stage = \ - (float4*)(dbias_stage + (warp * kRowsPerWarp + h * kRowsInFlight + (T)) * kTileCols + \ - kColsPerLane * lane); \ - stage[0] = make_float4(products[0], products[1], products[2], products[3]); \ - stage[1] = make_float4(products[4], products[5], products[6], products[7]); \ - } \ - } - static_assert(kRowsInFlight == 1 || kRowsInFlight == 2, "A load group is 1 or 2 rows."); - MXFP8_DACT_ROW(0) - if constexpr (kRowsInFlight == 2) MXFP8_DACT_ROW(1) -#undef MXFP8_DACT_ROW -#undef MXFP8_DACT_WORD - } - } - - // The rowwise half does not depend on the colwise scales, so it runs before - // the cross-warp barrier (see TileConfig). - const size_t store_offset = - (size_t)(row0 + warp * kRowsPerWarp) * K + col0 + kColsPerLane * lane; - auto row_scale = [&](const Word* row, int j) { - Word amax = DP::abs_max(DP::abs_max(row[0], row[1]), DP::abs_max(row[2], row[3])); -#pragma unroll - for (int m = 4; m < kWordsPerLane; m += 2) - amax = DP::abs_max(amax, DP::abs_max(row[m], row[m + 1])); - // A 32-value rowwise block spans kLanesPerMxGroup lanes. -#pragma unroll - for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) - amax = DP::abs_max(amax, DP::shfl_xor(amax, lane_mask)); - // Fold the two halves of the pair (for 16-bit pairs one PRMT and a packed max), which - // leaves the magnitude in both halves. - amax = DP::abs_max(amax, DP::swap_halves(amax)); - const unsigned scale_byte = amax_to_e8m0_2x(0.f, DP::widen(DP::clear_sign(amax)).y); - if ((lane & (kLanesPerMxGroup - 1)) == 0) - row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + (lane / kLanesPerMxGroup)] = - (unsigned char)scale_byte; - return DP::row_rcp(scale_byte); - }; - // Colwise amax: per-warp partial in registers, cross-warp fold in shared. - auto publish_col_amax_partials = [&]() { - Word col_amax[kWordsPerLane]; -#pragma unroll - for (int m = 0; m < kWordsPerLane; ++m) col_amax[m] = tile_words[m]; -#pragma unroll - for (int j = 1; j < kRowsPerWarp; ++j) -#pragma unroll - for (int m = 0; m < kWordsPerLane; ++m) - col_amax[m] = DP::abs_max(col_amax[m], tile_words[j * kWordsPerLane + m]); -#pragma unroll - for (int n = 0; n < kVecsPerLaneRow; ++n) - reinterpret_cast(&col_amax_partials[warp][kWordsPerLane * lane])[n] = - reinterpret_cast(col_amax)[n]; - }; - if constexpr (C::kThreadsPerCta == 128) { - static_assert(sizeof(Word) == 4, "Paired rows pack two 16-bit amaxes into one word."); - // On the 128-thread activation tile, two rows share one butterfly and conversion. - auto row_amax = [&](const unsigned* row) { - unsigned amax = DP::abs_max(DP::abs_max(row[0], row[1]), DP::abs_max(row[2], row[3])); -#pragma unroll - for (int m = 4; m < kWordsPerLane; m += 2) - amax = DP::abs_max(amax, DP::abs_max(row[m], row[m + 1])); - return DP::abs_max(amax, __byte_perm(amax, amax, 0x1032)); - }; - auto row_scale_pair = [&](const unsigned* row_a, const unsigned* row_b, int j, - Scale& row_scale_rcp0, Scale& row_scale_rcp1) { - unsigned amax_pair = __byte_perm(row_amax(row_a), row_amax(row_b), 0x5410) & 0x7fff7fffu; -#pragma unroll - for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) - amax_pair = DP::abs_max(amax_pair, __shfl_xor_sync(0xffffffffu, amax_pair, lane_mask)); - const ptx::floatx2 amax_2x = DP::widen(amax_pair); - const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); - const unsigned scale_byte0 = scale_pair & 0xffu, scale_byte1 = (scale_pair >> 8) & 0xffu; - if ((lane & (kLanesPerMxGroup - 1)) == 0) { - unsigned char* scale_dst = row_scale_bytes + - (warp * kRowsPerWarp + j) * kMxGroupsPerRow + - (lane / kLanesPerMxGroup); - scale_dst[0] = (unsigned char)scale_byte0; - scale_dst[kMxGroupsPerRow] = (unsigned char)scale_byte1; - } - row_scale_rcp0 = DP::row_rcp(scale_byte0); - row_scale_rcp1 = DP::row_rcp(scale_byte1); - }; - unsigned char* row_out_ptr = rowwise_out + store_offset; -#pragma unroll - for (int j = 0; j < kRowsPerWarp; j += 2) { - const Word* row_a = tile_words + j * kWordsPerLane; - const Word* row_b = row_a + kWordsPerLane; - Scale row_scale_rcp0, row_scale_rcp1; - row_scale_pair(row_a, row_b, j, row_scale_rcp0, row_scale_rcp1); - unsigned row_out[kWordsPerLane / 2]; -#pragma unroll - for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = DP::template quantize(row_a[2 * m], row_a[2 * m + 1], - row_scale_rcp0, row_scale_rcp0); - if (in_bounds(j)) store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); -#pragma unroll - for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = DP::template quantize(row_b[2 * m], row_b[2 * m + 1], - row_scale_rcp1, row_scale_rcp1); - if (in_bounds(j + 1)) - store_quantized(row_out_ptr + (unsigned)(j + 1) * (unsigned)K, row_out); - } - } else { - unsigned char* row_out_ptr = rowwise_out + store_offset; -#pragma unroll - for (int j = 0; j < kRowsPerWarp; ++j) { - const Word* v = tile_words + j * kWordsPerLane; - const Scale row_scale_rcp = row_scale(v, j); - unsigned row_out[kWordsPerLane / 2]; -#pragma unroll - for (int m = 0; m < kWordsPerLane / 2; ++m) - row_out[m] = - DP::template quantize(v[2 * m], v[2 * m + 1], row_scale_rcp, row_scale_rcp); - if (in_bounds(j)) store_quantized(row_out_ptr + (unsigned)j * (unsigned)K, row_out); - } - } - - publish_col_amax_partials(); - __syncthreads(); - if constexpr (kDbias) { - // The generic kernel's order: a strictly sequential FP32 sum down the column. - if (tid < kTileCols) { - float sum = dbias_partial; -#pragma unroll 8 - for (int r = 0; r < kRowsPerMxBlock; ++r) - sum = __fadd_rn(sum, dbias_stage[r * kTileCols + tid]); - dbias_partial = sum; - } - } - if (tid < kTileCols / 2) { - Word col_amax = col_amax_partials[0][tid]; -#pragma unroll - for (int w = 1; w < kWarpsPerCta; ++w) - col_amax = DP::abs_max(col_amax, col_amax_partials[w][tid]); - const ptx::floatx2 amax_2x = DP::widen(DP::clear_sign(col_amax)); - // The tensor amax, when asked for: the max of these column amaxes. Like the generic - // kernel's per-element fmaxf, it ignores NaN (a column of only NaNs gives NaN here, - // which fmaxf drops) and is unaffected by the zeros loaded past M or K. - if constexpr (kAmax) thread_amax = fmaxf(thread_amax, fmaxf(amax_2x.x, amax_2x.y)); - const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); - if (!kEdge || col0 + 2 * tid < colwise_scale_cols) { - if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - // The colwise swizzle is the rowwise one with rows and columns swapped, - // so the two columns of this pair land in different 16-byte rows. - const size_t row_block_tiles = DIVUP(static_cast(M), size_t{128}); - const size_t col = col0 + 2 * tid; - colwise_scales[swizzle::gemm_swizzled_scale_idx(col, row0 / 32, row_block_tiles)] = - static_cast(scale_pair); - colwise_scales[swizzle::gemm_swizzled_scale_idx(col + 1, row0 / 32, row_block_tiles)] = - static_cast(scale_pair >> 8); - } else { - *(unsigned short*)(colwise_scales + (size_t)(row0 / 32) * colwise_scale_stride + col0 + - 2 * tid) = (unsigned short)scale_pair; - } - } - col_scale_rcp_smem[tid] = DP::col_rcp(scale_pair); - } - auto drain_row_scales = [&]() { - constexpr int drain_tid_base = C::kDrainTidBase; - // A row's eight group scales are eight contiguous, 4-byte-aligned bytes of - // rowwise_scales (the stride is a multiple of 4 groups), so one warp drains a - // 32-row tile with two STG.32 per lane. In the GEMM-swizzled layout they are two - // 4-byte runs in adjacent tiles. A tile that ends inside the tensor writes only - // the groups that start inside it. - if (tid >= drain_tid_base && tid < drain_tid_base + kRowsPerMxBlock) { - const int r = tid - drain_tid_base; - const unsigned char* bytes = row_scale_bytes + r * kMxGroupsPerRow; - const size_t col_tiles = DIVUP(static_cast(K), size_t{128}); - if (!kEdge || full_width) { - const uint2 words = *(const uint2*)bytes; - if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - constexpr size_t kTileBytes = - swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_X * swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_Y; - const size_t idx = swizzle::gemm_swizzled_scale_idx(row0 + r, col0 >> 5, col_tiles); - *(unsigned*)(rowwise_scales + idx) = words.x; - *(unsigned*)(rowwise_scales + idx + kTileBytes) = words.y; - } else { - unsigned* dst = (unsigned*)(rowwise_scales + - (size_t)(row0 + r) * rowwise_scale_stride + (col0 >> 5)); - dst[0] = words.x; - dst[1] = words.y; - } - } else { - const int groups = (K - col0 + 31) / 32; - for (int g = 0; g < groups; ++g) { - size_t idx; - if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - idx = swizzle::gemm_swizzled_scale_idx(row0 + r, (col0 >> 5) + g, col_tiles); - } else { - idx = (size_t)(row0 + r) * rowwise_scale_stride + (col0 >> 5) + g; - } - rowwise_scales[idx] = bytes[g]; - } - } - } - }; - drain_row_scales(); - __syncthreads(); - - // ---- columnwise quantization -------------------------------------------- - { - // This lane's kWordsPerLane column-pair scales, as 16-byte vectors. - constexpr int kScaleVecs = kWordsPerLane * sizeof(Scale) / sizeof(uint4); - uint4 col_scale_rcp_vec[kScaleVecs]; -#pragma unroll - for (int n = 0; n < kScaleVecs; ++n) - col_scale_rcp_vec[n] = - reinterpret_cast(col_scale_rcp_smem + kWordsPerLane * lane)[n]; - const Scale* col_scale_rcp = reinterpret_cast(col_scale_rcp_vec); - unsigned char* col_out_ptr = colwise_out + store_offset; -#pragma unroll - for (int j = 0; j < kRowsPerWarp; ++j) { - const Word* v = tile_words + j * kWordsPerLane; - unsigned col_out[kWordsPerLane / 2]; -#pragma unroll - for (int m = 0; m < kWordsPerLane / 2; ++m) - col_out[m] = DP::template quantize(v[2 * m], v[2 * m + 1], col_scale_rcp[2 * m], - col_scale_rcp[2 * m + 1]); - if (in_bounds(j)) store_quantized(col_out_ptr + (unsigned)j * (unsigned)K, col_out); - } - } - }; - if constexpr (kTails) { - if (full_width && row0 + kRowsPerMxBlock <= M) { - band(std::false_type{}); - } else { - band(std::true_type{}); - } - } else { - band(std::false_type{}); - } - if constexpr (kDbias) { - if ((it & 1) == 1 && tid < kTileCols && col0 + tid < K) { - dbias_workspace[(size_t)(row0 / (2 * kRowsPerMxBlock)) * K + col0 + tid] = dbias_partial; - } - } - } - - if constexpr (kAmax) { - const float block_amax = reduce_max(thread_amax, warp); - if (tid == 0) atomicMaxFloat(amax_ptr, block_amax); - } -} - -template -__global__ void __launch_bounds__( - TileConfig::kThreadsPerCta, - TileConfig::kMinBlocksPerSm) - quantize_mxfp8_kernel(const unsigned* __restrict__ input, - const unsigned* __restrict__ act_input, - unsigned char* __restrict__ rowwise_out, - unsigned char* __restrict__ rowwise_scales, - unsigned char* __restrict__ colwise_out, - unsigned char* __restrict__ colwise_scales, int M, int K, - int rowwise_scale_stride, int colwise_scale_stride, int iters, - const void* __restrict__ table, float* __restrict__ dbias_workspace, - float* __restrict__ amax_ptr) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - quantize_regtile((const uint4*)input, (const uint4*)act_input, rowwise_out, rowwise_scales, - colwise_out, colwise_scales, M, K, rowwise_scale_stride, - colwise_scale_stride, iters, table, dbias_workspace, amax_ptr); -#endif -} - -// Row blocks per CTA. A longer walk amortizes the per-CTA table copy; a -// shorter one keeps the grid deep enough to avoid a residency tail. -//! Grid size, in CTAs, that the walk length aims for. -constexpr long long kCtaTargetActivation = 16384; -//! Minimum walk for the instantiations that copy the table. -constexpr int kMinWalk = 2; -//! Maximum walk for the derivatives. -constexpr int kMaxWalkDact = 64; - -static int pick_walk_length(int row_blocks, int grid_cols, long long target, - int min_walk = kMinWalk) { - int walk = 1; - while ((row_blocks / walk) % 2 == 0 && (long long)grid_cols * (row_blocks / walk) > target) - walk *= 2; - if (walk < min_walk && row_blocks % min_walk == 0) walk = min_walk; - return walk; -} - -// Shared-memory carveout, in percent, for the activations, whose footprint leaves room: -// the rest of the unified array goes to L1. -constexpr int kCarveoutPercentAct = 40; - -template -static void set_carveout() { - constexpr int kPercent = IS_ACT ? kCarveoutPercentAct : 0; - constexpr int kDynamicBytes = - TileConfig::kDbiasStageBytes; - if constexpr (kPercent > 0 || kDynamicBytes > 0) { - // Function attributes are per device, like the tables. - static std::mutex mutex; - static std::vector done; - int device; - NVTE_CHECK_CUDA(cudaGetDevice(&device)); - std::lock_guard lock(mutex); - if (static_cast(device) >= done.size()) done.resize(device + 1, false); - if (!done[device]) { - const void* kernel = - (const void*)quantize_mxfp8_kernel; - if constexpr (kPercent > 0) { - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, kPercent)); - } - if constexpr (kDynamicBytes > 0) { - // The staging buffer takes the kernel past the 48 KB of shared memory a CTA gets by - // default. - NVTE_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - kDynamicBytes)); - } - done[device] = true; - } - } -} - -//! 32-row bands the grid covers: whole 64-row tiles of the generic kernel, whose scales it -//! writes. -inline int row_bands(int M) { return (M + 63) / 64 * 2; } -//! 256-column tiles the grid covers. -inline int col_tiles(int K) { return (K + kTileCols - 1) / kTileCols; } - -// Row blocks per CTA walk, in powers of two. -template -int walk_length(int M, int K) { - const int row_blocks = row_bands(M); - const int grid_cols = col_tiles(K); - // The minimum walk amortizes the table copy, but not at the price of leaving - // SMs without a CTA. - int min_walk = kMinWalk; - if ((long long)grid_cols * (row_blocks / min_walk) < 2LL * cuda::sm_count()) min_walk = 1; - int iters = pick_walk_length(row_blocks, grid_cols, kCtaTargetActivation, min_walk); - if constexpr (IS_DACT) { - if (iters > kMaxWalkDact) iters = kMaxWalkDact; - } - // With dbias a CTA walks whole 64-row chunks: two bands per chunk. - if constexpr (kDbias) { - if (iters % 2 != 0) iters *= 2; - } - return iters; -} - -template -static void launch(const void* input, const void* act_input, void* rowwise_out, - void* rowwise_scales, void* colwise_out, void* colwise_scales, int M, int K, - int rowwise_scale_stride, int colwise_scale_stride, const void* table, - float* dbias_workspace, float* amax, cudaStream_t stream) { - using C = TileConfig; - const int iters = walk_length(M, K); - dim3 grid(col_tiles(K), row_bands(M) / iters); - // kAmax: a separate instantiation, so that the kernels without amax keep their register - // budget (the derivatives run at the 40-register cap). - const auto run_amax = [&](auto tails, auto with_amax) { - constexpr bool kTails = decltype(tails)::value; - constexpr bool kAmax = decltype(with_amax)::value; - set_carveout(); - quantize_mxfp8_kernel<<>>( - (const unsigned*)input, (const unsigned*)act_input, (unsigned char*)rowwise_out, - (unsigned char*)rowwise_scales, (unsigned char*)colwise_out, (unsigned char*)colwise_scales, - M, K, rowwise_scale_stride, colwise_scale_stride, iters, table, dbias_workspace, amax); - }; - const auto run = [&](auto tails) { - if (amax != nullptr) { - run_amax(tails, std::true_type{}); - } else { - run_amax(tails, std::false_type{}); - } - }; - // Whole 64-row tiles and 256-column tiles leave nothing past M or K. The derivatives take the - // tail-capable kernel for every shape: with the band body behind its edge branch, the compiler - // schedules the loads so that they wait less (GB200, 32768x7168 dGeLU: 239 -> 220 us, fewer - // long-scoreboard stalls, DRAM throughput 73% -> 79%). The forward kernels do not gain from it. - if (!IS_DACT && M % 64 == 0 && K % kTileCols == 0) { - run(std::false_type{}); - } else { - run(std::true_type{}); - } -} - -// Any number of rows. Columns in multiples of 16, the shapes the generic kernel takes (its TMA -// store of the FP8 output needs 16-byte rows); this also keeps every input row 16-byte aligned. -// Output offsets within a 32-row band are 32-bit. -inline bool shape_supported(size_t rows, size_t cols) { - constexpr size_t kMaxRows = std::numeric_limits::max() - 64; - constexpr size_t kMaxCols = std::numeric_limits::max() / kRowsPerMxBlock; - return rows > 0 && cols > 0 && rows <= kMaxRows && cols <= kMaxCols && cols % 16 == 0; -} - -} // namespace - -//! Whether this kernel implements the request type: GeLU or SiLU, or the derivative of -//! either, the derivative optionally with dbias. It reads the activation from a table, or -//! computes it with util/math.h when there is none. -template -constexpr bool kImplements = - (!IS_DBIAS || IS_DACT) && (IS_ACT != IS_DACT) && activation_table::kHasTable && - (IS_ACT == activation_table::Table>::kForward); - -//! Whether FP16 inputs of this request type take this kernel. FP16 always computes the -//! activation and scales every value in FP32 (see Datapath), about twice the -//! instructions of the BF16 path, which the compute-bound kernels cannot absorb. GB200 against -//! TE main: dSiLU 1.25-1.71x, dbias_dGeLU 1.16-1.43x, dbias_dSiLU 1.27-1.50x, but GeLU -//! 0.74-0.91x, SiLU 0.86-1.26x and dGeLU 0.95-1.15x, which stay on the generic kernel. -template -constexpr bool kTakesFP16 = - IS_DBIAS || activation_table::kKindOf == activation_table::kDsilu; - -// FP32 inputs take this kernel for every request type: their values are already FP32, so the -// FP32 scaling costs nothing extra. GB200 against TE main: GeLU 1.12-1.35x, SiLU 1.24-1.43x, -// dGeLU 1.29-1.48x, dSiLU 1.48-2.37x, dbias_dGeLU 1.16-1.32x, dbias_dSiLU 1.35-2.00x. -// -// Whether this kernel takes a request of a type in kImplements, with BF16, FP16 or FP32 input and E4M3 or -// E5M2 output. Anything else takes the generic kernel. A noop tensor is ignored, as the generic -// kernel ignores it for every request that computes an activation. -template -bool can_use(const Tensor& input, const Tensor* act_input, const Tensor& output, const Tensor& noop, - bool use_2d_quantization, const activation_table::TableRef& table) { - const auto [rows, cols] = input.flat_2d_dims(); - // Without a table the forward activation is compute-bound here, and on small tensors too - // few CTAs hide its latency: GB200, 993x7200 GeLU is 0.93x of the generic kernel, 4160x2048 - // 1.26x. The derivatives stay memory-bound and win at every size. - constexpr size_t kComputedForwardMinElements = size_t{1} << 23; - if (table.use == activation_table::kNoTable && !IS_DACT && - rows * cols < kComputedForwardMinElements) { - return false; - } - // Vector widths: 16-byte loads of both inputs, 8-byte stores of both outputs - // and of each row's scales. - const bool aligned = is_aligned_ptr(input.data.dptr, 16) && - (!IS_DACT || is_aligned_ptr(act_input->data.dptr, 16)) && - is_aligned_ptr(output.data.dptr, 8) && - is_aligned_ptr(output.columnwise_data.dptr, 8) && - is_aligned_ptr(output.scale_inv.dptr, 8) && - is_aligned_ptr(output.columnwise_scale_inv.dptr, 2); - if (!(transformer_engine::cuda::sm_arch() >= 100 && !use_2d_quantization && output.has_data() && - output.has_columnwise_data() && shape_supported(rows, cols) && aligned)) { - return false; - } - constexpr size_t kMaxGridY = 65535; - const int M = static_cast(rows); - return static_cast( - row_bands(M) / walk_length(M, static_cast(cols))) <= kMaxGridY; -} - -//! dbias_workspace: [ceil(rows / 64), cols] FP32 column sums for common::reduce_dbias, or -//! nullptr without dbias. -template -void quantize(const Tensor& input, const Tensor* act_input, Tensor* output, - const activation_table::TableRef& table, float* dbias_workspace, - cudaStream_t stream) { - using namespace activation_table; - // Plain variables: lambdas cannot capture structured bindings before C++20. - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); - const void* act_ptr = IS_DACT ? act_input->data.dptr : input.data.dptr; - const auto run = [&](auto table_type) { - using TableT = typename decltype(table_type)::type; - launch( - input.data.dptr, act_ptr, output->data.dptr, output->scale_inv.dptr, - output->columnwise_data.dptr, output->columnwise_scale_inv.dptr, static_cast(rows), - static_cast(cols), static_cast(output->scale_inv.shape[1]), - static_cast(output->columnwise_scale_inv.shape[1]), table.ptr, dbias_workspace, - static_cast(output->amax.dptr), stream); - }; - constexpr Kind kKind = kKindOf; - // The tables hold BF16 inputs; FP16 inputs always compute. - if (!std::is_same_v || table.use == kNoTable) { - run(std::common_type>{}); - NVTE_CHECK_CUDA(cudaGetLastError()); - return; - } - if constexpr (!std::is_same_v) { - return; - } else if constexpr (kKind == kDsilu) { - if (table.use == kExactDsiluTable) { - run(std::common_type>{}); - NVTE_CHECK_CUDA(cudaGetLastError()); - return; - } - } - if constexpr (std::is_same_v) { - run(std::common_type>{}); - NVTE_CHECK_CUDA(cudaGetLastError()); - } -} - -} // namespace regtile -} // namespace mxfp8 -} // namespace dispatch -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_QUANTIZE_MXFP8_REGTILE_CUH_ diff --git a/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh b/transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh similarity index 72% rename from transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh rename to transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh index ea3d056221e..512cb393911 100644 --- a/transformer_engine/common/cast/mxfp8/regtile/quantize_gated_mxfp8_regtile.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh @@ -4,29 +4,138 @@ * See LICENSE for license information. ************************************************************************/ -/*! \file quantize_gated_mxfp8_regtile.cuh +/*! \file gated_bidim.cuh * \brief Register-resident MXFP8 bidimensional quantize fused with SwiGLU or its gradient. * - * The gated counterpart of quantize_mxfp8_regtile.cuh. Input [M, 2K] holds the activated half - * and the gate half side by side. The forward writes silu(act) * gate as [M, K]; the backward - * writes the gradients of both halves as [M, 2K], quantized independently, one after the other - * from the same 32-row band. A CTA covers 256 output columns and walks 32-row bands; the band - * stays in registers, and both scale directions are computed from there. + * Input [M, 2K] holds the activated half and the gate half side by side. The forward writes + * silu(act) * gate as [M, K]; the backward writes the gradients of both halves as [M, 2K], + * quantized independently, one after the other from the same 32-row band. A CTA covers 256 + * output columns and walks 32-row bands; the band stays in registers, and both scale + * directions are computed from there, instead of being staged through shared memory as in the + * generic gated kernel (gated_mxfp8.cuh). * - * It writes the bytes the generic gated kernel (gated_mxfp8.cuh) writes: the same math, the - * same rounding to BF16, and scales only inside the tensor. + * It writes the bytes the generic gated kernel writes: the same math, the same rounding to + * BF16, and scales only inside the tensor. */ -#ifndef TRANSFORMER_ENGINE_QUANTIZE_GATED_MXFP8_REGTILE_CUH_ -#define TRANSFORMER_ENGINE_QUANTIZE_GATED_MXFP8_REGTILE_CUH_ +#ifndef TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_GATED_BIDIM_CUH_ +#define TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_GATED_BIDIM_CUH_ -#include "quantize_mxfp8_regtile.cuh" +#include + +#include +#include + +#include "../../../common.h" +#include "../../../util/cuda_runtime.h" +#include "../../../util/math.h" +#include "../../../util/ptx.cuh" +#include "../../../utils.cuh" +#include "../swizzle.cuh" namespace transformer_engine { namespace dispatch { namespace mxfp8 { -namespace regtile { -namespace gated { +namespace gated_bidim { + +namespace ptx = transformer_engine::ptx; + +namespace { + +/*! \brief Elementwise max(|a|, |b|) on a packed BF16 pair, as a raw word. */ +__device__ __forceinline__ unsigned abs_max_bf16x2(unsigned a, unsigned b) { + ptx::bf16x2 d; + ptx::abs_max_2x(d, reinterpret_cast(a), + reinterpret_cast(b)); + return reinterpret_cast(d); +} + +/*! \brief Widen a BF16 pair word to FP32, low half to x. Same result as ptx::up_cast, whose + * volatile asm the compiler cannot schedule as freely in the derivative loop. */ +__device__ __forceinline__ ptx::floatx2 widen_bf16x2(const unsigned v) { + return {__uint_as_float(v << 16), __uint_as_float(v & 0xffff0000u)}; +} + +/*! \brief Round an FP32 pair to a BF16 pair word, low lane in the low half. */ +__device__ __forceinline__ unsigned to_bf16x2(const ptx::floatx2& v) { + const __nv_bfloat162 r = __floats2bfloat162_rn(v.x, v.y); + return reinterpret_cast(r); +} + +/*! \brief The reciprocal of an E8M0 scale, 2^(127-e), broadcast to a BF16 pair. */ +__device__ __forceinline__ unsigned e8m0_to_bf16x2_reciprocal(unsigned e) { + const ptx::bf16x2 r = ptx::exp2f_rcp_2x(static_cast(e)); + return reinterpret_cast(r); +} + +/*! \brief E8M0 scales for two block amaxes, packed as ptx::float_to_e8m0_2x. + * + * The generic kernel's amax is a 0-seeded max that ignores NaN, so a block of + * only NaNs gets amax 0; the fmaxf reproduces that for the data-seeded maxima + * here. The conversion saturates, so the scale never reaches 255. + */ +template +__device__ __forceinline__ unsigned amax_to_e8m0_2x(float amax_hi, float amax_lo) { + constexpr float kMaxNormRcp = Quantized_Limits::max_norm_rcp; + return ptx::float_to_e8m0_2x(fmaxf(amax_hi * kMaxNormRcp, 0.0f), + fmaxf(amax_lo * kMaxNormRcp, 0.0f)); +} + +// Tile geometry. A CTA covers 256 columns and walks 32-row sub-tiles; 32 rows +// is the colwise MX block height, so the colwise reduction closes in the CTA. + +//! Columns a CTA covers. +constexpr int kTileCols = 256; +//! Elements in one MX block, and equivalently the colwise block height. +constexpr int kRowsPerMxBlock = 32; +//! MX groups (32 columns each) across the tile. +constexpr int kMxGroupsPerRow = kTileCols / 32; +//! A register-resident tile spreads its 256 columns over the 32 lanes of a warp. +constexpr int kColsPerLane = kTileCols / 32; +//! ...which a lane holds as BF16 pairs, one 32-bit word each. +constexpr int kWordsPerLane = kColsPerLane / 2; +//! ...read as one 128-bit vector per row. +constexpr int kVecLoadsPerRow = kWordsPerLane / 4; +//! Lanes that have to cooperate to cover one 32-column MX group. +constexpr int kLanesPerMxGroup = 32 / kColsPerLane; + +/*! \brief Scale two BF16 pairs, each by its own reciprocal scale pair, to four OType (E4M3 or + * E5M2). */ +template +__device__ __forceinline__ unsigned quantize_x4(unsigned in01, unsigned in23, unsigned scale01, + unsigned scale23) { + ptx::FPx4 out; + ptx::mul_cvt_4x(out, reinterpret_cast(in01), + reinterpret_cast(scale01), + reinterpret_cast(in23), + reinterpret_cast(scale23)); + return reinterpret_cast(out); +} + +/*! \brief Store one lane's quantized row: kColsPerLane FP8 bytes, one STG.64. */ +__device__ __forceinline__ void store_quantized(unsigned char* p, const unsigned* v) { + static_assert(kColsPerLane == 8, "One STG.64 covers exactly eight FP8 bytes."); + *(uint2*)p = *(const uint2*)v; +} + + +// Row bands per CTA. Walking several bands deepens each CTA's work without starving SMs. +//! Grid size, in CTAs, that the walk length aims for. +constexpr long long kCtaTargetActivation = 16384; +//! Minimum walk, unless that leaves fewer than two CTAs per SM. +constexpr int kMinWalk = 2; +//! Maximum walk. +constexpr int kMaxWalkDact = 64; + +int pick_walk_length(int row_blocks, int grid_cols, long long target, int min_walk = kMinWalk) { + int walk = 1; + while ((row_blocks / walk) % 2 == 0 && (long long)grid_cols * (row_blocks / walk) > target) + walk *= 2; + if (walk < min_walk && row_blocks % min_walk == 0) walk = min_walk; + return walk; +} + +} // namespace constexpr int kThreadsPerCta = 256; constexpr int kWarpsPerCta = kThreadsPerCta / 32; @@ -265,7 +374,7 @@ __global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) #endif } -//! Whether the gated regtile kernel takes this request: BF16 [M, 2K] -> E4M3 or E5M2, both directions, +//! Whether this kernel takes the request: BF16 [M, 2K] -> E4M3 or E5M2, both directions, //! K % 32 == 0 (the generic gated kernel's requirement), aligned buffers. inline bool can_use(const Tensor& input, const Tensor& grad, const Tensor& output, const bool is_bwd) { @@ -317,10 +426,9 @@ void quantize(const Tensor& input, const Tensor& grad, Tensor* output, cudaStrea NVTE_CHECK_CUDA(cudaGetLastError()); } -} // namespace gated -} // namespace regtile +} // namespace gated_bidim } // namespace mxfp8 } // namespace dispatch } // namespace transformer_engine -#endif // TRANSFORMER_ENGINE_QUANTIZE_GATED_MXFP8_REGTILE_CUH_ +#endif // TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_GATED_BIDIM_CUH_ diff --git a/transformer_engine/common/util/activation_table.cuh b/transformer_engine/common/util/activation_table.cuh index 87bfaf9426e..550e1c7f1f5 100644 --- a/transformer_engine/common/util/activation_table.cuh +++ b/transformer_engine/common/util/activation_table.cuh @@ -9,10 +9,11 @@ * * A BF16 input takes only 2^16 values, so these activations can be read from a table * generated once per device from the scalar util/math.h functions. TE builds each table - * on its first use and keeps it for the life of the process (device_table); kernels read - * them for BF16 inputs quantized to MXFP8 when NVTE_ACTIVATION_TABLE=1 (off by default). - * Outside each table's window the result has a closed form, so every lookup is - * branch-free. + * on its first use and keeps it for the life of the process (device_table). With + * NVTE_ACTIVATION_TABLE=1 (off by default), BF16 inputs quantized to MXFP8 follow the tables: + * the CuTeDSL register-resident kernel reads them, and the other kernels compute what they + * hold (kComputedTable). Outside each table's window the result has a closed form, so every + * lookup is branch-free. * * Forward (GeLU, SiLU): BF16 entries. Below the window the result is RN(x/2); above it * it is x for x > 0 and -0 for x < 0. The forward result is truncated to BF16 by every @@ -66,12 +67,10 @@ enum Kind : int { template struct ForwardLayout { static constexpr bool kForward = true; - //! Whether the kernel computes op() instead of reading the table (see the register-tiled kernel). - static constexpr bool kComputed = false; static constexpr uint32_t kLowSentinel = kLo - 1; static constexpr uint32_t kHighSentinel = kHi + 1; static constexpr uint32_t kRows = kHighSentinel - kLowSentinel + 1; - //! Rounded up to 16 bytes, the granularity of a bulk copy into shared memory. + //! Rounded up to 16 bytes, so kernels can copy it into shared memory in 16-byte pieces. static constexpr size_t kBytes = (kRows * 2 * sizeof(__nv_bfloat16) + 15) / 16 * 16; }; @@ -83,7 +82,6 @@ struct ForwardLayout { template struct DerivativeLayout { static constexpr bool kForward = false; - static constexpr bool kComputed = false; static constexpr uint32_t kLoBits = kLo; static constexpr uint32_t kHiBits = kHi; //! Magnitudes from here on (and NaN) give NaN, as util/math.h does. @@ -91,11 +89,9 @@ struct DerivativeLayout { static constexpr uint32_t kEntriesPerSign = kHi - kLo + 1; static constexpr uint32_t kSignStride = kStride; static_assert(kEntriesPerSign <= kSignStride, "The two signs' entries must not overlap."); - //! Rounded up to 16 bytes, the granularity of a bulk copy into shared memory. + //! Rounded up to 16 bytes, so kernels can copy it into shared memory in 16-byte pieces. static constexpr size_t kBytes = ((kSignStride + kEntriesPerSign) * sizeof(float) + 15) / 16 * 16; - //! The sign bit, 0x8000, shifted right by this is the byte offset of the negative half. - static constexpr int kSignShift = kStride == 2048 ? 2 : 1; - static_assert(kStride == 2048 || kStride == 4096, "The packed lookup assumes these strides."); + static_assert(kStride == 2048 || kStride == 4096, "The packed lookups assume these strides."); }; /*! \brief The table of kKind for inputs of type IType. @@ -260,86 +256,6 @@ __device__ __forceinline__ float derivative(const float x, const float *__restri return magnitude >= T::kSaturatedEndBits ? __uint_as_float(0x7fffffffu) : value; } -/*! \brief Forward activation of a BF16 pair, low lane in the low half. Same results as - * forward() per lane; the table may be in shared memory. - */ -template -__device__ __forceinline__ uint32_t forward_2x(const uint32_t x, - const __nv_bfloat16 *__restrict__ table) { - static_assert(T::kForward, "Not a forward table."); - const __nv_bfloat162 a = *reinterpret_cast(&x); - // Clamp one lane's magnitude to the table rows. The +0 bound of the high lane zeroes it, - // so the word is exactly the clamped magnitude. - const auto row = [](const __nv_bfloat162 lane) { - constexpr uint32_t kLowBound = T::kLowSentinel * 0x10001u; - constexpr uint32_t kHighBound = T::kHighSentinel; - __nv_bfloat162 m = - __hmax2(__habs2(lane), *reinterpret_cast(&kLowBound)); - m = __hmin2(m, *reinterpret_cast(&kHighBound)); - return *reinterpret_cast(&m) - T::kLowSentinel; - }; - const __nv_bfloat16 entry_lo = table[2 * row(__low2bfloat162(a)) + ((x >> 15) & 1u)]; - const __nv_bfloat16 entry_hi = table[2 * row(__high2bfloat162(a)) + (x >> 31)]; - constexpr uint32_t kHalf = 0x3f003f00u, kNegativeZero = 0x80008000u; - const __nv_bfloat162 half = __hmul2(a, *reinterpret_cast(&kHalf)); - const __nv_bfloat162 upper = - __hmax2_nan(a, *reinterpret_cast(&kNegativeZero)); - const __nv_bfloat162 r = - __hmin2_nan(__hmax2(__halves2bfloat162(entry_lo, entry_hi), half), upper); - return *reinterpret_cast(&r); -} - -/*! \brief Derivative of a BF16 pair, low lane in .x. Same results as derivative() per lane; - * the table may be in shared memory. - * - * Both lanes' byte offsets are formed in one 32-bit word: the clamp runs on the BF16 pair, - * and each lane's offset stays below 0x10000, so the halves never carry into each other. - */ -template -__device__ __forceinline__ float2 derivative_2x(const uint32_t x, - const unsigned char *__restrict__ table) { - static_assert(!T::kForward, "Not a derivative table."); - static_assert((T::kSignStride + T::kEntriesPerSign) * sizeof(float) <= 0x10000, - "Byte offsets must fit in 16 bits."); - constexpr uint32_t kLo2 = T::kLoBits * 0x10001u, kHi2 = T::kHiBits * 0x10001u; - const uint32_t magnitude = x & 0x7fff7fffu; - const __nv_bfloat162 clamped = - __hmin2(__hmax2(*reinterpret_cast(&magnitude), - *reinterpret_cast(&kLo2)), - *reinterpret_cast(&kHi2)); - const uint32_t offsets = ((*reinterpret_cast(&clamped) - kLo2) << 2) + - ((x ^ magnitude) >> T::kSignShift); - const float entry_lo = *reinterpret_cast(table + (offsets & 0xffffu)); - const float entry_hi = *reinterpret_cast(table + (offsets >> 16)); - const float linear_lo = T::below_window(__uint_as_float(x << 16)); - const float linear_hi = T::below_window(__uint_as_float(x & 0xffff0000u)); - constexpr uint32_t kSaturatedEnd2 = T::kSaturatedEndBits * 0x10001u; - float lo, hi; -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - // One BF16x2 compare per condition sets both lanes' predicates. The unordered >= also - // holds for NaN, which gives NaN like the saturated end. - asm("{\n\t" - ".reg .pred below_lo, below_hi, nan_lo, nan_hi;\n\t" - "setp.lt.bf16x2 below_lo|below_hi, %2, %3;\n\t" - "setp.geu.bf16x2 nan_lo|nan_hi, %2, %4;\n\t" - "selp.f32 %0, %5, %6, below_lo;\n\t" - "selp.f32 %1, %7, %8, below_hi;\n\t" - "selp.b32 %0, 0x7fffffff, %0, nan_lo;\n\t" - "selp.b32 %1, 0x7fffffff, %1, nan_hi;\n\t" - "}" - : "=f"(lo), "=f"(hi) - : "r"(magnitude), "r"(kLo2), "r"(kSaturatedEnd2), "f"(linear_lo), "f"(entry_lo), - "f"(linear_hi), "f"(entry_hi)); -#else - const uint32_t magnitude_lo = magnitude & 0xffffu, magnitude_hi = magnitude >> 16; - lo = magnitude_lo < T::kLoBits ? linear_lo : entry_lo; - hi = magnitude_hi < T::kLoBits ? linear_hi : entry_hi; - lo = magnitude_lo >= T::kSaturatedEndBits ? __uint_as_float(0x7fffffffu) : lo; - hi = magnitude_hi >= T::kSaturatedEndBits ? __uint_as_float(0x7fffffffu) : hi; -#endif - return make_float2(lo, hi); -} - /*! \brief OP(x), from the table given by kUse, or from util/math.h for kNoTable. * * x must be a BF16 value widened to FP32 when reading a table. From 47e947fda49e0418575f2b1d633e4e46942e0e90 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:17:18 +0000 Subject: [PATCH 28/29] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh b/transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh index 512cb393911..12684f093b8 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh @@ -118,7 +118,6 @@ __device__ __forceinline__ void store_quantized(unsigned char* p, const unsigned *(uint2*)p = *(const uint2*)v; } - // Row bands per CTA. Walking several bands deepens each CTA's work without starving SMs. //! Grid size, in CTAs, that the walk length aims for. constexpr long long kCtaTargetActivation = 16384; From 8dae8054429a39eb06a616750a0c7c46e0705284 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 2 Nov 2026 17:04:23 -0800 Subject: [PATCH 29/29] [Common] Add CuTeDSL gated MXFP8 kernels, including a register-resident one Carries the CuTeDSL port of the gated MXFP8 kernel (gated_mxfp8.cuh) from kainzhong/TransformerEngine#13 (head 4687cc18), which is not merged yet, and adds a register-resident gated kernel on top of it. Revert this commit to drop both until that port lands. From kainzhong/TransformerEngine#13, unchanged except for formatting (black) and one argument: - quantize_mxfp8_gated_g2s.py: the gated kernel on the general kernel's TMA pipeline, for all gated activations, forward and backward. - MXFP8GatedQuantizeConfig, gated_mxfp8_cutedsl.cuh, cutedsl_utils.h and the dispatch in cast/dispatch/gated.cuh: CuTeDSL first, the CUDA kernel as the fallback. - The gated kernel passes ACTIVATION_TABLE="none" to the general shared-memory quantizer, which takes the activation table layout on this branch. Register-resident gated kernel (quantize_mxfp8_gated_g2r_bidim.py), a port of cast/mxfp8/specialized/gated_bidim.cuh for every gated activation: - BF16 input with columnwise output (rowwise optional). A CTA covers 256 output columns and walks 32-row bands. Each thread loads its rows of the act and gate halves (and of the gradient, in backward) into registers, computes the gated op there with the elementwise math of the shared-memory port, rounds it to BF16 as the CUDA kernels do, and quantizes the band in both directions from registers. The backward quantizes its two output halves one after the other. - The register-tile quantizer of the bidimensional kernel is factored out (quantize_mxfp8_rmem_tile_bidimensional, with its layout, bounds and row-load helpers) and shared by both kernels; the bidimensional kernel is unchanged. - Only the bands that cross a tensor's edge take the predicated body. - MXFP8GatedQuantizeEntry selects it at run time, and the shared-memory port otherwise (FP16/FP32, rowwise-only output, misaligned columnwise scales). mxh2 parity against TE main (CUDA), 1230 gated cases over SwiGLU, GeGLU, ReGLU, QGeGLU and SReGLU in both directions, all layouts, compact and swizzled scales and special-value fills: every case the register-resident kernel takes matches. The shared-memory port differs only in scale bytes of blocks of only NaNs (255 where the CUDA gated kernel writes 0). CastMXFP8_GatedAct passes with the CuTeDSL backend on and off. GB200, BF16, both directions, geometric mean speedup over TE main on 11 shapes: SwiGLU 1.58, dSwiGLU 1.82, GeGLU 1.26, dGeGLU 1.42, ReGLU 1.35, dReGLU 1.13, QGeGLU 1.59, dQGeGLU 2.27, SReGLU 1.42, dSReGLU 1.07 (CUDA backend: SwiGLU 1.65, dSwiGLU 1.87 with the register-tiled kernel, about 1.0 for the others). Co-authored-by: Kaining Zhong Signed-off-by: Teddy Do --- tests/jax/test_mxfp8_cutedsl_backend.py | 121 ++- .../mxfp8/test_mxfp8_cutedsl_backend.py | 127 ++- transformer_engine/common/CuTeDSL/__init__.py | 13 +- .../common/CuTeDSL/activations.py | 54 ++ .../cast/mxfp8/quantize_mxfp8_common.py | 555 ++++++++---- .../cast/mxfp8/quantize_mxfp8_g2r_bidim.py | 480 +++++----- .../mxfp8/quantize_mxfp8_gated_g2r_bidim.py | 611 +++++++++++++ .../cast/mxfp8/quantize_mxfp8_gated_g2s.py | 821 ++++++++++++++++++ .../common/cast/dispatch/gated.cuh | 28 +- .../common/cast/mxfp8/gated_mxfp8_cutedsl.cuh | 203 +++++ .../common/util/cutedsl_utils.h | 70 ++ 11 files changed, 2705 insertions(+), 378 deletions(-) create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2r_bidim.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py create mode 100644 transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 2ee82f8c5d5..c44b93b5614 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -23,6 +23,7 @@ from transformer_engine.common import _get_shared_object_file from transformer_engine.common.CuTeDSL.utils import device_compute_capability from transformer_engine.jax import cpp_extensions as tex +from transformer_engine.jax.cpp_extensions.activation import ActivationParams from transformer_engine.jax.quantize import ( QuantizerFactory, QuantizeLayout, @@ -68,8 +69,8 @@ # Every activation the CuTeDSL backend implements, as (JAX activation_type, name it goes by # in the config key). The two differ for qgelu/srelu: JAX spells them quick_gelu/squared_relu # (ActivationEnum in jax/cpp_extensions/activation.py), the backend uses the C++ spelling -# (Activation in common/util/cutedsl_utils.h). Gated variants are absent because they go to a -# separate TE/common kernel that the backend does not cover. +# (Activation in common/util/cutedsl_utils.h). Gated variants go to a separate kernel and are +# covered by GATED_ACTIVATIONS below. ACTIVATIONS = [ (("gelu",), "gelu"), (("relu",), "relu"), @@ -316,3 +317,119 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): ) def test_specialized_kernels(shape, q_layout, in_dtype): run_test_case("CAST_ONLY", *NO_ACT, shape, q_layout, in_dtype, jnp.float8_e4m3fn) + + +# Gated activations (gated_mxfp8.cuh), as (JAX activation_type, name in the config key +# (GatedActivation in common/util/cutedsl_utils.h), activation parameters). The clamped SwiGLU +# parameters are chosen so that its clamps are active on the [-2, 2] inputs generate_inputs draws. +GATED_ACTIVATIONS = [ + (("gelu", "linear"), "geglu", {}), + (("quick_gelu", "linear"), "qgeglu", {}), + (("relu", "linear"), "reglu", {}), + (("squared_relu", "linear"), "sreglu", {}), + (("silu", "linear"), "swiglu", {}), + (("sigmoid", "linear"), "glu", {}), + ( + ("clamped_silu", "clamped_linear"), + "clamped_swiglu", + {"limit": 1.0, "alpha": 1.702, "glu_linear_offset": 0.5}, + ), +] +GATED_IDS = [desc for _, desc, _ in GATED_ACTIVATIONS] +get_direction_id = lambda is_bwd: "bwd" if is_bwd else "fwd" + + +def get_gated_cfg_key(act_desc, in_dtype, fp8_dtype, q_layout, is_bwd): + """Mirror of MXFP8GatedQuantConfig::to_key (gated_mxfp8_cutedsl.cuh).""" + # JAX quantize emits scales in the linear (non-swizzled) layout + flags = (True, q_layout.has_colwise, False, is_bwd) + major, minor = device_compute_capability() + return ( + f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_" + + "_".join("1" if f else "0" for f in flags) + + f"_{act_desc}" + ) + + +def run_gated_test_case( + act_type, act_desc, act_kwargs, is_bwd, shape, q_layout, in_dtype, fp8_dtype +): + """Assert the CuTeDSL and CUDA backends produce bit-identical gated activation outputs. + shape is (M, N): the gated input is (M, 2, N), [act | gate] along axis -2. + """ + M, N = shape + # generate_inputs returns (M, N) and (M, 1, N) draws; stack them into the (M, 2, N) input + act, gate = generate_inputs(M, N, in_dtype) + x = jnp.concatenate([jnp.expand_dims(act, axis=-2), gate], axis=-2) + grad, _ = generate_inputs(M, N, in_dtype, seed=1) + act_params = ActivationParams.create(act_type, **act_kwargs) + + def run(): + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, q_dtype=fp8_dtype, q_layout=q_layout + ) + if is_bwd: + out, _ = tex.quantize_dact_dbias( + grad, x, act_type, is_dbias=False, quantizer=quantizer, act_params=act_params + ) + else: + out = tex.act_lu(x, act_type, quantizer=quantizer, act_params=act_params) + return extract_quantized_output(out, None)[0] + + set_cutedsl_backend(False) + cuda_output = run() + set_cutedsl_backend(True) + try: + cutedsl_output = run() + finally: + set_cutedsl_backend(False) + + key = get_gated_cfg_key(act_desc, in_dtype, fp8_dtype, q_layout, is_bwd) + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + tag = ( + f"{act_desc}/{get_direction_id(is_bwd)}/{get_layout_id(q_layout)}/{M}x{N}/" + f"{DTYPE_TO_STR[in_dtype]}/{FP8_TO_STR[fp8_dtype]}" + ) + for name, cuda_bytes in cuda_output.items(): + assert np.array_equal( + cutedsl_output[name], cuda_bytes + ), f"{tag}: {name} differ between backends" + + +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("q_layout", Q_LAYOUTS, ids=get_layout_id) +@pytest.mark.parametrize("act_type,act_desc,act_kwargs", GATED_ACTIVATIONS, ids=GATED_IDS) +@pytest.mark.parametrize("is_bwd", [False, True], ids=get_direction_id) +def test_gated_sizes(is_bwd, act_type, act_desc, act_kwargs, q_layout, shape): + run_gated_test_case( + act_type, + act_desc, + act_kwargs, + is_bwd, + shape, + q_layout, + jnp.bfloat16, + jnp.float8_e4m3fn, + ) + + +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("act_type,act_desc,act_kwargs", GATED_ACTIVATIONS, ids=GATED_IDS) +@pytest.mark.parametrize("is_bwd", [False, True], ids=get_direction_id) +def test_gated_dtypes(is_bwd, act_type, act_desc, act_kwargs, fp8_dtype, in_dtype): + run_gated_test_case( + act_type, + act_desc, + act_kwargs, + is_bwd, + (256, 384), + QuantizeLayout.ROWWISE_COLWISE, + in_dtype, + fp8_dtype, + ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 5d7feeb6ce6..fa51d292cb7 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -63,8 +63,8 @@ BLOCK_SIZES = [ROW, COL, BOTH] # Every activation the CuTeDSL backend implements; "desc" is the name it goes by in the -# config key (Activation in common/util/cutedsl_utils.h). Gated variants are absent because -# they go to a separate TE/common kernel that the backend does not cover. +# config key (Activation in common/util/cutedsl_utils.h). Gated variants go to a separate kernel +# and are covered by GATED_ACTIVATIONS below. IDENTITY = {"name": "Identity", "act": None, "dact": None, "dbias_dact": None, "desc": "none"} ACTIVATIONS = [ { @@ -431,3 +431,126 @@ def test_specialized_kernels(shape, block_size, in_dtype, swizzled, input_offset swizzled, input_offset, ) + + +# Gated activations (gated_mxfp8.cuh). Input is [M, 2N] = [act | gate]; the forward writes +# [M, N] and the backward, given grad [M, N], writes [M, 2N]. "desc" is the name the activation +# goes by in the config key (GatedActivation in common/util/cutedsl_utils.h). The parameters of +# the parameterized ones are chosen so that their clamps and tanh terms are active on the +# [-2, 2] inputs generate_inputs draws. +GATED_ACTIVATIONS = [ + {"name": "GeGLU", "act": tex.geglu, "dact": tex.dgeglu, "desc": "geglu", "kwargs": {}}, + {"name": "QGeGLU", "act": tex.qgeglu, "dact": tex.dqgeglu, "desc": "qgeglu", "kwargs": {}}, + {"name": "ReGLU", "act": tex.reglu, "dact": tex.dreglu, "desc": "reglu", "kwargs": {}}, + {"name": "SReGLU", "act": tex.sreglu, "dact": tex.dsreglu, "desc": "sreglu", "kwargs": {}}, + {"name": "SwiGLU", "act": tex.swiglu, "dact": tex.dswiglu, "desc": "swiglu", "kwargs": {}}, + {"name": "GLU", "act": tex.glu, "dact": tex.dglu, "desc": "glu", "kwargs": {}}, + { + "name": "ClampedSwiGLU", + "act": tex.clamped_swiglu, + "dact": tex.clamped_dswiglu, + "desc": "clamped_swiglu", + "kwargs": {"limit": 1.0, "alpha": 1.702, "glu_linear_offset": 0.5}, + }, + { + "name": "SiTUGLU", + "act": tex.situglu, + "dact": tex.dsituglu, + "desc": "situglu", + "kwargs": {"beta1": 1.0, "beta2": 2.0}, + }, +] +get_gated_id = lambda a: a["name"] +GATED_DIRECTIONS = [False, True] +get_direction_id = lambda is_bwd: "bwd" if is_bwd else "fwd" + + +def run_on_both_backends(run): + """Return (CUDA result, CuTeDSL result) of run(), with run() responsible for materializing + whatever it returns before the backend is switched.""" + set_cutedsl_backend(False) + cuda_result = run() + set_cutedsl_backend(True) + try: + cutedsl_result = run() + finally: + set_cutedsl_backend(False) + return cuda_result, cutedsl_result + + +def assert_cutedsl_kernel_registered(key): + """Guard against a silent CUDA fallback, which would make a cross-backend comparison compare + CUDA against itself.""" + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + +def get_gated_cfg_key(act, in_dtype, fp8_dtype, rowwise, colwise, swizzled, is_bwd): + """Mirror of MXFP8GatedQuantConfig::to_key (gated_mxfp8_cutedsl.cuh).""" + major, minor = device_compute_capability() + flags = (rowwise, colwise, swizzled, is_bwd) + return ( + f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_" + + "_".join("1" if f else "0" for f in flags) + + f"_{act['desc']}" + ) + + +def run_gated_test_case(act, is_bwd, shape, block_size, in_dtype, fp8_dtype, swizzled=False): + """Assert the CuTeDSL and CUDA backends produce bit-identical gated activation outputs. + shape is (M, N): the gated input is [M, 2N]. + """ + M, N = shape + if is_bwd and swizzled: + pytest.skip("CuTeDSL hands backward gated activations with GEMM-swizzled scales to CUDA") + rowwise = block_size[1] != 1 + columnwise = block_size[0] != 1 + x, _ = generate_inputs(M, 2 * N, in_dtype) + grad, _ = generate_inputs(M, N, in_dtype, seed=1) + + def run(): + q = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=rowwise, columnwise=columnwise) + q.optimize_for_gemm = swizzled + if is_bwd: + out = act["dact"](grad, x, q, **act["kwargs"]) + else: + out = act["act"](x, q, **act["kwargs"]) + return extract_quantized_output(out, rowwise, columnwise, swizzled) + + cuda_output, cutedsl_output = run_on_both_backends(run) + assert_cutedsl_kernel_registered( + get_gated_cfg_key(act, in_dtype, fp8_dtype, rowwise, columnwise, swizzled, is_bwd) + ) + + layout = "swizzled" if swizzled else "linear" + tag = ( + f"{act['name']}/{get_direction_id(is_bwd)}/{M}x{N}/{DTYPE_TO_STR[in_dtype]}/" + f"{FP8_TO_STR[fp8_dtype]}/{layout}" + ) + for name, cuda_bytes in cuda_output.items(): + assert torch.equal( + cutedsl_output[name], cuda_bytes + ), f"{tag}: {name} differ between backends" + + +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +@pytest.mark.parametrize("act", GATED_ACTIVATIONS, ids=get_gated_id) +@pytest.mark.parametrize("is_bwd", GATED_DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_gated_sizes(swizzled, is_bwd, act, block_size, shape): + run_gated_test_case( + act, is_bwd, shape, block_size, torch.bfloat16, tex.DType.kFloat8E4M3, swizzled + ) + + +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("act", GATED_ACTIVATIONS, ids=get_gated_id) +@pytest.mark.parametrize("is_bwd", GATED_DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_gated_dtypes(swizzled, is_bwd, act, fp8_dtype, in_dtype): + run_gated_test_case(act, is_bwd, (256, 384), BOTH, in_dtype, fp8_dtype, swizzled) diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 15aae1aba2d..aac7f232ab0 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -12,6 +12,7 @@ from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8_common import ( get_mxfp8_quantization_function, + get_mxfp8_gated_quantization_function, ) from transformer_engine.common.CuTeDSL.cast.mxfp8.requantize_mxfp8 import ( get_mxfp8_requantization_function, @@ -22,9 +23,9 @@ def register_cutedsl_backends(): """Register all available CuTeDSL backends for on-demand compilation via TVM-FFI. The C++ dispatcher retrieves them by the names defined here. """ - tvm_ffi.register_global_func( - "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True - ) - tvm_ffi.register_global_func( - "get_mxfp8_requantization_function", get_mxfp8_requantization_function, override=True - ) + for entrypoint in ( + get_mxfp8_quantization_function, + get_mxfp8_gated_quantization_function, + get_mxfp8_requantization_function, + ): + tvm_ffi.register_global_func(entrypoint.__name__, entrypoint, override=True) diff --git a/transformer_engine/common/CuTeDSL/activations.py b/transformer_engine/common/CuTeDSL/activations.py index 5f2440c9299..cf3a42b49b7 100644 --- a/transformer_engine/common/CuTeDSL/activations.py +++ b/transformer_engine/common/CuTeDSL/activations.py @@ -85,6 +85,60 @@ def dact_dgelu(x: Float32) -> Float32: ) + Float32(0.5) * (Float32(1.0) + t) +# ---- Gated activations (gated_mxfp8.cuh) ---- +# The parameterized ones take their ClampedSwiGLUParam / SiTUGLUParam fields as +# trailing runtime arguments. + +LOG2E = 1.4426950408889634 # log2(e), what __expf scales its argument by + + +def act_sigmoid(x: Float32) -> Float32: + return sigmoid(x) + + +def dact_dsigmoid(x: Float32) -> Float32: + return dsigmoid(x) + + +def act_clamped_silu(x: Float32, limit: Float32, alpha: Float32) -> Float32: + c = cute.arch.fmin(limit, x) + return c * sigmoid(alpha * c) + + +def act_situ_gate(x: Float32, beta1: Float32) -> Float32: + return beta1 * tanh(x / beta1) * sigmoid(x) + + +def dact_dsitu_gate(x: Float32, beta1: Float32) -> Float32: + t = tanh(x / beta1) + s = sigmoid(x) + return (Float32(1.0) - t * t) * s + beta1 * t * s * (Float32(1.0) - s) + + +def situ_up(x: Float32, beta2: Float32) -> Float32: + return beta2 * tanh(x / beta2) + + +def dsitu_up(x: Float32, beta2: Float32) -> Float32: + t = tanh(x / beta2) + return Float32(1.0) - t * t + + +@cute.jit +def sigmoidf(x: Float32) -> Float32: + """math.h's sigmoidf: __frcp_rn(1.0f + __expf(-x)).""" + e = cute.arch.inline_ptx( + "ex2.approx.f32 {$w0}, {$r0};", + write_only_types=[Float32], + read_only_args=[-x * Float32(LOG2E)], + ) + return cute.arch.inline_ptx( + "rcp.rn.f32 {$w0}, {$r0};", + write_only_types=[Float32], + read_only_args=[Float32(1.0) + e], + ) + + # BF16 activation tables (see common/util/activation_table.cuh, which builds them; the layouts # below must match it). With a table, BF16 activations follow the table specification: GeLU and # SiLU of -inf give -0, and the derivatives below the table window are the linear term diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index 08e8478a88a..e7d498c68d3 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -30,6 +30,8 @@ dact_dsilu, dact_dqgelu, dact_dgelu, + act_sigmoid, + dact_dsigmoid, derivative_table_spec, forward_table_spec, ) @@ -89,6 +91,21 @@ def dactivation_op(activation: str, activation_table: str): return derivative_table_spec(dop, activation, activation_table) +# Gated activations of gated_mxfp8.cuh, as (ActOP, DActOP) ports of the math.h functions the +# CUDA dispatcher instantiates them with. swiglu and clamped_swiglu backward, and situglu in +# both directions, do not go through this table: like the CUDA kernel they take a special path. +GATED_ACTIVATIONS = { + "geglu": (act_gelu, dact_dgelu), + "qgeglu": (act_qgelu, dact_dqgelu), + "reglu": (act_relu, dact_drelu), + "sreglu": (act_srelu, dact_dsrelu), + "glu": (act_sigmoid, dact_dsigmoid), + "swiglu": (act_silu, None), + "clamped_swiglu": (None, None), + "situglu": (None, None), +} + + @cute.jit def derive_swizzled_scale_layout( M: Int32, @@ -251,6 +268,294 @@ def is_plain_cast(self) -> bool: """Whether cfg is a plain cast: no amax, dbias or (d)activation fused in.""" return not (self.WITH_AMAX or self.WITH_DBIAS or self.WITH_DACT or self.WITH_ACT) + def compile(self): + """ + Return the compiled CuTeDSL function object for the given MXFP8 quantization config. + """ + + kernel_obj = MXFP8QuantizeEntry(self) + sym_M = cute.sym_int32() + sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) + in_shape = out_shape = (sym_M, sym_N) + # TE allocates scale tensors at a padded shape (see + # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): + # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) + # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) + # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the + # scales their own fresh syms carrying the divisibility the padding + # guarantees (rowwise: 128 x 4; colwise: 4 x 128). + scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) + scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) + ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N + # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ + # bridge sends). + out_dtype = self.FP8_DTYPE + scale_dtype = cutlass.Float8E8M0FNU + + in_fake = cute.runtime.make_fake_compact_tensor( + self.DTYPE, + in_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + out_row_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.ROWWISE + else None + ) + scale_row_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_rowwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if self.ROWWISE + else None + ) + out_col_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.COLWISE + else None + ) + scale_col_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_colwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if self.COLWISE + else None + ) + amax_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 + ) + if self.WITH_AMAX + else None + ) + # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that + # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). + # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher + # passes at launch. + noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) + act_input_fake = ( + cute.runtime.make_fake_compact_tensor( + self.DTYPE, + in_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.WITH_DACT + else None + ) + # Like the noop flag, the activation table is an always-present pointer, null without a + # table. + act_table_fake = cute.runtime.nullptr( + cutlass.Uint8, mem_space=cute.AddressSpace.gmem, assumed_align=16 + ) + workspace_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, + ws_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if self.WITH_DBIAS + else None + ) + + compiled = cute.compile( + kernel_obj, + in_fake, # mX + out_row_fake, + scale_row_fake, # mO_row, mS_row + out_col_fake, + scale_col_fake, # mO_col, mS_col + amax_fake, # mAmax + noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) + act_input_fake, # mDActInput (activation input, dact only) + workspace_fake, # mWorkspace(backward slot, unused) + act_table_fake, # mActTable (activation table, may be null at launch) + cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi + # "handle" arg; C++ passes the CUDA stream + # as void*) + options="--enable-tvm-ffi", + ) + return compiled + + +def _fake_gmem_tensor(dtype, shape, align): + """Row-major fake gmem tensor for cute.compile.""" + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=tuple(reversed(range(len(shape)))), + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + +class MXFP8GatedQuantizeConfig(MXFP8QuantizeConfig): + """Compile-time config of gated_mxfp8.cuh's port: a plain MXFP8 cast config plus the gated op.""" + + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + is_bwd: bool, + activation: str, + ): + super().__init__( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=False, + activation="none", + ) + if activation not in GATED_ACTIVATIONS: + raise ValueError( + f"unknown gated activation {activation!r}; expected one of" + f" {sorted(GATED_ACTIVATIONS)}" + ) + if is_bwd and with_gemm_swizzled_scales: + # TODO(kainingz): technically this is supported but I need to come up with a less disgusting way + # to implement it in CuTe + raise ValueError("backward gated activations with GEMM-swizzled scales") + self.IS_BWD = is_bwd + self.GATED_ACTIVATION = activation + + def __str__(self): + return ( + f"MXFP8GatedQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " + f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, is_bwd={self.IS_BWD}, " + f"activation={self.GATED_ACTIVATION})" + ) + + __repr__ = __str__ + + def compile(self): + """Return the compiled CuTeDSL function object for this config.""" + # The kernel module imports this one + # pylint: disable-next=import-outside-toplevel + + sym_M = cute.sym_int32() + # The C++ dispatcher requires N % 32 == 0, like the CUDA kernel. + sym_N = cute.sym_int32(divisibility=32) + sym_2N = cute.sym_int32(divisibility=64) + out_cols = sym_2N if self.IS_BWD else sym_N + # Padded scale shapes, see MXFP8QuantizeConfig.compile + scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) + scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) + # Only the activation's own params are runtime arguments; the others are compiled as None + # and the C++ launcher passes None for them. + clamped = self.GATED_ACTIVATION == "clamped_swiglu" + situ = self.GATED_ACTIVATION == "situglu" + params = [Float32(0.0) if clamped else None] * 3 # limit, alpha, glu_linear_offset + params += [Float32(0.0) if situ else None] * 2 # beta1, beta2 + return cute.compile( + MXFP8GatedQuantizeEntry(self), + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_2N), 16), # mGatedInput + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_N), 16) if self.IS_BWD else None, # mGrad + # mO_row, mS_row + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.ROWWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_rowwise_shape, 4) + if self.ROWWISE + else None + ), + # mO_col, mS_col + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.COLWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_colwise_shape, 4) + if self.COLWISE + else None + ), + *params, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + + +class MXFP8GatedQuantizeEntry: + """Select the gated MXFP8 kernel: the register-resident one where it applies, otherwise the + shared-memory port of gated_mxfp8.cuh.""" + + def __init__(self, cfg: MXFP8GatedQuantizeConfig): + # The kernel modules import this one. + # pylint: disable=import-outside-toplevel + from .quantize_mxfp8_gated_g2s import MXFP8GatedQuantizeG2SKernel + from .quantize_mxfp8_gated_g2r_bidim import MXFP8GatedQuantizeG2RKernel + + # pylint: enable=import-outside-toplevel + self.cfg = cfg + self.g2s_kernel = MXFP8GatedQuantizeG2SKernel(cfg) + self.g2r_kernel = ( + MXFP8GatedQuantizeG2RKernel(cfg) if MXFP8GatedQuantizeG2RKernel.supports(cfg) else None + ) + + @cute.jit + def __call__( + self, + mGatedInput: cute.Tensor, + mGrad: Optional[cute.Tensor], + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + stream: CUstream, + ): + args = ( + mGatedInput, + mGrad, + mO_row, + mS_row, + mO_col, + mS_col, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + stream, + ) + if cutlass.const_expr(self.g2r_kernel is not None): + if self.g2r_kernel.eligible_to_dispatch(mGatedInput, mS_col): + self.g2r_kernel(*args) + else: + self.g2s_kernel(*args) + else: + self.g2s_kernel(*args) + class MXFP8QuantizeKernelBase(abc.ABC): """Base class for MXFP8 quantize kernels.""" @@ -449,148 +754,8 @@ def __call__( ) -def compile_cutedsl_function_from_cfg(cfg): - """ - Return the compiled CuTeDSL function object for the given MXFP8 quantization config. - """ - - kernel_obj = MXFP8QuantizeEntry(cfg) - sym_M = cute.sym_int32() - sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) - in_shape = out_shape = (sym_M, sym_N) - # TE allocates scale tensors at a padded shape (see - # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): - # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) - # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) - # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the - # scales their own fresh syms carrying the divisibility the padding - # guarantees (rowwise: 128 x 4; colwise: 4 x 128). - scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) - scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) - ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N - # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ - # bridge sends). - out_dtype = cfg.FP8_DTYPE - scale_dtype = cutlass.Float8E8M0FNU - - in_fake = cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, in_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=16 - ) - out_row_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.ROWWISE - else None - ) - scale_row_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_rowwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.ROWWISE - else None - ) - out_col_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.COLWISE - else None - ) - scale_col_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_colwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.COLWISE - else None - ) - amax_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_AMAX - else None - ) - # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that - # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). - # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher - # passes at launch. - noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) - act_input_fake = ( - cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, - in_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.WITH_DACT - else None - ) - # Like the noop flag, the activation table is an always-present pointer, null without a table. - act_table_fake = cute.runtime.nullptr( - cutlass.Uint8, mem_space=cute.AddressSpace.gmem, assumed_align=16 - ) - workspace_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, ws_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_DBIAS - else None - ) - - compiled = cute.compile( - kernel_obj, - in_fake, # mX - out_row_fake, - scale_row_fake, # mO_row, mS_row - out_col_fake, - scale_col_fake, # mO_col, mS_col - amax_fake, # mAmax - noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) - act_input_fake, # mDActInput (activation input, dact only) - workspace_fake, # mWorkspace(backward slot, unused) - act_table_fake, # mActTable (activation table, may be null at launch) - cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi - # "handle" arg; C++ passes the CUDA stream - # as void*) - options="--enable-tvm-ffi", - ) - return compiled - - -def get_mxfp8_quantization_function( - fn_name: str, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - with_gemm_swizzled_scales: bool, - with_amax: bool, - with_dbias: bool, - with_dact: bool, - with_act: bool, - use_2d_quantization: bool, - activation: str, - activation_table: str, -) -> bool: - """Compile the MXFP8 quantize kernel for this config and register it in the TVM-FFI global registry +def compile_and_register_function(fn_name: str, cfg: MXFP8QuantizeConfig) -> bool: + """Compile the CuTeDSL kernel for this config and register it in the TVM-FFI global registry under EXACTLY `fn_name` (the key the C++ dispatcher built; Python treats it as an opaque name). Returns True if a kernel is successfully registered under `fn_name` (the C++ side then fetches it with GetGlobal(fn_name)); False if the config is unsupported, so the caller caches the negative result and falls back to the CUDA C++ kernel. @@ -610,34 +775,8 @@ def get_mxfp8_quantization_function( ) return False - try: - cfg = MXFP8QuantizeConfig( - dtype=dtype, - fp8_dtype=fp8_dtype, - rowwise=rowwise, - colwise=colwise, - with_gemm_swizzled_scales=with_gemm_swizzled_scales, - with_amax=with_amax, - with_dbias=with_dbias, - with_dact=with_dact, - with_act=with_act, - use_2d_quantization=use_2d_quantization, - activation=activation, - activation_table=activation_table, - ) - except ValueError as e: - # The exception message states exactly why the config is unsupported - # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a - # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. - logger.warning( - "CuTeDSL MXFP8 backend does not support this config, " - "falling back to the CUDA C++ kernel: %s", - e, - ) - return False - logger.debug("Compiling CuTeDSL MXFP8 quantization kernel for %s", cfg) - compiled = compile_cutedsl_function_from_cfg(cfg) + compiled = cfg.compile() # The returned compiled object is not neccessarily the compiled function itself. It could be a Python # wrapper that parses arguments and calls the underlying function. This is needed in the general case # as TVM-FFI is positional-only, however, for us this is unnecessary launch overhead as we pass the @@ -659,3 +798,81 @@ def get_mxfp8_quantization_function( ) # Unconditionally fallback to CUDA path because we can't tell if this exception is transient or permanent. return False + + +def get_mxfp8_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + with_amax: bool, + with_dbias: bool, + with_dact: bool, + with_act: bool, + use_2d_quantization: bool, + activation: str, + activation_table: str, +) -> bool: + """Compile and register the MXFP8 quantize kernel for this config under `fn_name`.""" + try: + cfg = MXFP8QuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=with_amax, + with_dbias=with_dbias, + with_dact=with_dact, + with_act=with_act, + use_2d_quantization=use_2d_quantization, + activation=activation, + activation_table=activation_table, + ) + except ValueError as e: + # The exception message states exactly why the config is unsupported + # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a + # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. + logger.warning( + "CuTeDSL MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + return compile_and_register_function(fn_name, cfg) + + +def get_mxfp8_gated_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + is_bwd: bool, + activation: str, +) -> bool: + """Compile and register the gated MXFP8 quantize kernel for this config under `fn_name`.""" + try: + cfg = MXFP8GatedQuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + is_bwd=is_bwd, + activation=activation, + ) + except ValueError as e: + # The exception message states exactly why the config is unsupported + # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a + # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. + logger.warning( + "CuTeDSL MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + return compile_and_register_function(fn_name, cfg) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py index c4eed9f1040..82f006233e6 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -672,6 +672,246 @@ def _quantize_colwise_f32_scaled( ) +def thread_tile_layout( + CTA_THREADS_Y: int, CTA_THREADS_X: int, ELEMENTS_Y_PER_THREAD: int, ELEMENTS_X_PER_THREAD: int +): + """The (thread, value) layout of a register-resident (CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD, + CTA_THREADS_X * ELEMENTS_X_PER_THREAD) tile: each warp owns ELEMENTS_Y_PER_THREAD consecutive + rows, and each lane ELEMENTS_X_PER_THREAD adjacent columns.""" + thread_layout = cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)) + _, tv_layout = cute.make_layout_tv( + thr_layout=thread_layout, + val_layout=cute.make_layout( + (ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD), stride=(ELEMENTS_X_PER_THREAD, 1) + ), + ) + return tv_layout + + +@cute.jit +def thread_tile_bounds( + tile_shape: cute.Shape, + tv_layout: cute.Layout, + rows_valid: Int32, + cols_valid: Int32, +): + """(rows_left, cols_ok) of the calling thread in a tile that ends at rows_valid rows and + cols_valid columns: its rows from rows_left on are past the tensor, and all of them when its + columns are (a thread's columns lie all inside or all past the tensor).""" + tidx, _, _ = cute.arch.thread_idx() + thread_origin = cute.composition(cute.make_identity_tensor(tile_shape), tv_layout)[tidx, None][ + ((0, 0),) + ] + cols_ok = thread_origin[1] < cols_valid + rows_left = rows_valid - thread_origin[0] + if not cols_ok: + rows_left = Int32(0) + return rows_left, cols_ok + + +@cute.jit +def load_thread_rows( + gT_thread: cute.Tensor, + rT: cute.Tensor, + load_atom: cute.CopyAtom, + cache_policy: Int64, + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, +): + """Copy all rows a thread sweeps into its (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) + register fragment. With EDGE, rows from rows_left on load as zeros.""" + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + if cutlass.const_expr(EDGE): + # Predicated loads rather than branches, so the rows' loads still issue together. + rT[None, row_idx].fill(0) + row_pred = cute.make_rmem_tensor( + cute.make_layout((ELEMENTS_X_PER_THREAD, (1,)), stride=(0, (0,))), Boolean + ) + row_pred[0] = row_idx < rows_left + cute.copy( + load_atom, + gT_thread[(None, row_idx),], + rT[None, row_idx], + pred=row_pred, + cache_policy=cache_policy, + ) + else: + cute.copy( + load_atom, gT_thread[(None, row_idx),], rT[None, row_idx], cache_policy=cache_policy + ) + + +@cute.jit +def quantize_mxfp8_rmem_tile_bidimensional( + rX: cute.Tensor, + gO_row_tile: Optional[cute.Tensor], + gS_row_tile: Optional[cute.Tensor], + gO_col_tile: Optional[cute.Tensor], + gS_col_tile: Optional[cute.Tensor], + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + FUSED: cutlass.Constexpr[bool], + EDGE: cutlass.Constexpr[bool], + rows_left: Int32, + cols_ok: Boolean, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + cache_policy: Int64, +): + """Quantize a register-resident 32-row tile, the CTA's rX fragments (see thread_tile_layout), + to the (TILE_ROWS, TILE_COLS) output tiles of the directions enabled by ROWWISE / COLWISE. + FUSED: the values come from a fused activation (see _bf16_nan_magnitudes_to_zero). With EDGE, + rows_left and cols_ok come from thread_tile_bounds.""" + TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD + THREADS_X_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_X_PER_THREAD + THREADS_Y_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_Y_PER_THREAD + SCALE_DTYPE = cutlass.Float8E8M0FNU + tidx, _, _ = cute.arch.thread_idx() + tv_layout = thread_tile_layout( + CTA_THREADS_Y, CTA_THREADS_X, ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD + ) + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + FP8_DTYPE, + num_bits_per_copy=ELEMENTS_X_PER_THREAD * FP8_DTYPE.width, + ) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) + + if cutlass.const_expr(ROWWISE): + gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_row_scale = cute.make_layout( + ( + (THREADS_X_PER_MX_BLOCK, (CTA_THREADS_X // THREADS_X_PER_MX_BLOCK, CTA_THREADS_Y)), + ELEMENTS_Y_PER_THREAD, + ), + stride=( + (0, (TILE_ROWS, ELEMENTS_Y_PER_THREAD)), + 1, + ), + ) + gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] + + row_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8 + ) + + if cutlass.const_expr(DTYPE is not cutlass.BFloat16): + _quantize_rowwise_f32_scaled( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + MAX_NORM_RCP=MAX_NORM_RCP, + EDGE=EDGE, + rows_left=rows_left, + ) + else: + _quantize_rowwise( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + EDGE=EDGE, + rows_left=rows_left, + FUSED=FUSED, + ) + + if cutlass.const_expr(COLWISE): + gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_col_scale = cute.make_layout( + ( + (CTA_THREADS_X, (THREADS_Y_PER_MX_BLOCK, CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK)), + ELEMENTS_X_PER_THREAD, + ), + stride=( + (ELEMENTS_X_PER_THREAD, (0, 1)), + 1, + ), + ) + # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment + # Otherwise coalesce will not do anything which is still fine + tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) + gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] + + # GEMM-swizzled colwise scales are not contiguous (adjacent columns land 16 bytes apart), + # so they are stored a byte at a time, as the general kernel does. + col_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + SCALE_DTYPE, + num_bits_per_copy=8 if WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, + ) + + if cutlass.const_expr(DTYPE is not cutlass.BFloat16): + _quantize_colwise_f32_scaled( + rX, + gO_col_thread, + gS_col_thread, + output_store_atom=store_atom, + col_scale_store_atom=col_scale_atom, + cache_policy=cache_policy, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + MAX_NORM_RCP=MAX_NORM_RCP, + EDGE=EDGE, + rows_left=rows_left, + cols_ok=cols_ok, + ) + else: + _quantize_colwise( + rX, + gO_col_thread, + gS_col_thread, + output_store_atom=store_atom, + col_scale_store_atom=col_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + EDGE=EDGE, + rows_left=rows_left, + cols_ok=cols_ok, + FUSED=FUSED, + ) + + @cute.jit def quantize_mxfp8_gmem_tile_bidimensional( gX_tile: cute.Tensor, @@ -710,23 +950,10 @@ def quantize_mxfp8_gmem_tile_bidimensional( Load the tile once, then reuse its registers for each enabled quantization pass. The tensors of a disabled direction are unused and may be None. """ - # Shape of the tile - TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD - - # How many threads cooperate to process a single MXFP8 block (32 elements) for both quantized directions - THREADS_X_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_X_PER_THREAD - THREADS_Y_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_Y_PER_THREAD - - SCALE_DTYPE = cutlass.Float8E8M0FNU - tidx, _, _ = cute.arch.thread_idx() - thread_layout = cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)) - _, tv_layout = cute.make_layout_tv( - thr_layout=thread_layout, - val_layout=cute.make_layout( - (ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD), stride=(ELEMENTS_X_PER_THREAD, 1) - ), + tv_layout = thread_tile_layout( + CTA_THREADS_Y, CTA_THREADS_X, ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD ) gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] @@ -745,14 +972,6 @@ def quantize_mxfp8_gmem_tile_bidimensional( # invariant=True here will drop the cache policy invariant=False, ) - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), - FP8_DTYPE, - num_bits_per_copy=ELEMENTS_X_PER_THREAD * FP8_DTYPE.width, - ) - - scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) - # With EDGE the tile ends at rows_valid rows and cols_valid columns: the thread's rows from # rows_left on are past the tensor, and all of them when its columns are (a thread's columns # lie all inside or all past the tensor, as N % 16 == 0). They load as zeros, which leave @@ -760,34 +979,19 @@ def quantize_mxfp8_gmem_tile_bidimensional( rows_left = Int32(ELEMENTS_Y_PER_THREAD) cols_ok = Boolean(True) if cutlass.const_expr(EDGE): - thread_origin = cute.composition(cute.make_identity_tensor(gX_tile.shape), tv_layout)[ - tidx, None - ][((0, 0),)] - cols_ok = thread_origin[1] < cols_valid - rows_left = rows_valid - thread_origin[0] - if not cols_ok: - rows_left = Int32(0) + rows_left, cols_ok = thread_tile_bounds(gX_tile.shape, tv_layout, rows_valid, cols_valid) # Copy all rows a thread sweeps into the RMEM fragment at once - for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - if cutlass.const_expr(EDGE): - # Predicated loads rather than branches, so the rows' loads still issue together. - rX[None, row_idx].fill(0) - row_pred = cute.make_rmem_tensor( - cute.make_layout((ELEMENTS_X_PER_THREAD, (1,)), stride=(0, (0,))), Boolean - ) - row_pred[0] = row_idx < rows_left - cute.copy( - load_atom, - gX_thread[(None, row_idx),], - rX[None, row_idx], - pred=row_pred, - cache_policy=cache_policy, - ) - else: - cute.copy( - load_atom, gX_thread[(None, row_idx),], rX[None, row_idx], cache_policy=cache_policy - ) + load_thread_rows( + gX_thread, + rX, + load_atom, + cache_policy, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + EDGE=EDGE, + rows_left=rows_left, + ) if cutlass.const_expr(WITH_ACT or WITH_DACT): rA = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), DTYPE) @@ -803,27 +1007,16 @@ def quantize_mxfp8_gmem_tile_bidimensional( ) # Activation tile has the same layout as the input tile so we can reuse the tv layout gA_thread = cute.composition(gA_tile, tv_layout)[tidx, None] - for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): - if cutlass.const_expr(EDGE): - rA[None, row_idx].fill(0) - act_row_pred = cute.make_rmem_tensor( - cute.make_layout((ELEMENTS_X_PER_THREAD, (1,)), stride=(0, (0,))), Boolean - ) - act_row_pred[0] = row_idx < rows_left - cute.copy( - act_load_atom, - gA_thread[(None, row_idx),], - rA[None, row_idx], - pred=act_row_pred, - cache_policy=cache_policy, - ) - else: - cute.copy( - act_load_atom, - gA_thread[(None, row_idx),], - rA[None, row_idx], - cache_policy=cache_policy, - ) + load_thread_rows( + gA_thread, + rA, + act_load_atom, + cache_policy, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + EDGE=EDGE, + rows_left=rows_left, + ) # With dbias, the FP32 products are staged for the column sums (see _kernel_main). if cutlass.const_expr(sDbias is not None): sDbias_thread = cute.composition(sDbias, tv_layout)[tidx, None] @@ -917,131 +1110,28 @@ def quantize_mxfp8_gmem_tile_bidimensional( fabs_f32(x2_lo_to_f32_bf16(pair_amax)), fabs_f32(x2_hi_to_f32_bf16(pair_amax)) ) - if cutlass.const_expr(ROWWISE): - gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] - - # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. - # And assign stride 0 to that submode - tv_row_scale = cute.make_layout( - ( - (THREADS_X_PER_MX_BLOCK, (CTA_THREADS_X // THREADS_X_PER_MX_BLOCK, CTA_THREADS_Y)), - ELEMENTS_Y_PER_THREAD, - ), - stride=( - (0, (TILE_ROWS, ELEMENTS_Y_PER_THREAD)), - 1, - ), - ) - gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] - - row_scale_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8 - ) - - if cutlass.const_expr(DTYPE is not cutlass.BFloat16): - _quantize_rowwise_f32_scaled( - rX, - gO_row_thread, - gS_row_thread, - output_store_atom=store_atom, - row_scale_store_atom=row_scale_atom, - cache_policy=cache_policy, - THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - FP8_DTYPE=FP8_DTYPE, - MAX_NORM_RCP=MAX_NORM_RCP, - EDGE=EDGE, - rows_left=rows_left, - ) - else: - _quantize_rowwise( - rX, - gO_row_thread, - gS_row_thread, - output_store_atom=store_atom, - row_scale_store_atom=row_scale_atom, - cache_policy=cache_policy, - scale_and_pack=scale_and_pack, - THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - FP8_DTYPE=FP8_DTYPE, - EDGE=EDGE, - rows_left=rows_left, - FUSED=WITH_ACT or WITH_DACT, - ) - - if cutlass.const_expr(COLWISE): - gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] - - # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. - # And assign stride 0 to that submode - tv_col_scale = cute.make_layout( - ( - (CTA_THREADS_X, (THREADS_Y_PER_MX_BLOCK, CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK)), - ELEMENTS_X_PER_THREAD, - ), - stride=( - (ELEMENTS_X_PER_THREAD, (0, 1)), - 1, - ), - ) - # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment - # Otherwise coalesce will not do anything which is still fine - tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) - gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] - - # GEMM-swizzled colwise scales are not contiguous (adjacent columns land 16 bytes apart), - # so they are stored a byte at a time, as the general kernel does. - col_scale_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), - SCALE_DTYPE, - num_bits_per_copy=8 if WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, - ) - - if cutlass.const_expr(DTYPE is not cutlass.BFloat16): - _quantize_colwise_f32_scaled( - rX, - gO_col_thread, - gS_col_thread, - output_store_atom=store_atom, - col_scale_store_atom=col_scale_atom, - cache_policy=cache_policy, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - FP8_DTYPE=FP8_DTYPE, - MAX_NORM_RCP=MAX_NORM_RCP, - EDGE=EDGE, - rows_left=rows_left, - cols_ok=cols_ok, - ) - else: - _quantize_colwise( - rX, - gO_col_thread, - gS_col_thread, - output_store_atom=store_atom, - col_scale_store_atom=col_scale_atom, - cache_policy=cache_policy, - scale_and_pack=scale_and_pack, - CTA_THREADS_Y=CTA_THREADS_Y, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - DTYPE=DTYPE, - FP8_DTYPE=FP8_DTYPE, - EDGE=EDGE, - rows_left=rows_left, - cols_ok=cols_ok, - FUSED=WITH_ACT or WITH_DACT, - ) + quantize_mxfp8_rmem_tile_bidimensional( + rX, + gO_row_tile, + gS_row_tile, + gO_col_tile, + gS_col_tile, + ROWWISE=ROWWISE, + COLWISE=COLWISE, + FUSED=WITH_ACT or WITH_DACT, + EDGE=EDGE, + rows_left=rows_left, + cols_ok=cols_ok, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + MAX_NORM_RCP=MAX_NORM_RCP, + WITH_GEMM_SWIZZLED_SCALES=WITH_GEMM_SWIZZLED_SCALES, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + cache_policy=cache_policy, + ) return thread_amax diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2r_bidim.py new file mode 100644 index 00000000000..69282fb4fc3 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2r_bidim.py @@ -0,0 +1,611 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Register-resident gated-activation MXFP8 cast in CuTeDSL, ported from the CUDA +register-tiled gated kernel (cast/mxfp8/specialized/gated_bidim.cuh). + +A CTA covers 256 output columns and walks 32-row bands. Each thread loads its rows of the act and +gate halves (and of the incoming gradient, in backward) into registers, computes the gated op +there with the elementwise math of the shared-memory port (quantize_mxfp8_gated_g2s.py, which +mirrors gated_mxfp8.cuh), rounds it to BF16 as the CUDA kernels do, and quantizes the band in +both directions from registers with the bidimensional register-tile quantizer +(quantize_mxfp8_g2r_bidim.py). The backward quantizes its two output halves one after the other +from the same band. Both write the bytes gated_mxfp8.cuh writes. +""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass.utils import HardwareInfo +from cutlass import Boolean, Float32, Int32, Int64 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + create_l2_policy, + cvt_f32x2_to_bf16x2, + max_num_bits_per_copy_ld_global, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8GatedQuantizeConfig, + derive_swizzled_scale_layout, +) +from .quantize_mxfp8_g2r_bidim import ( + quantize_mxfp8_rmem_tile_bidimensional, + thread_tile_bounds, + thread_tile_layout, +) +from .quantize_mxfp8_gated_g2s import _gated_element, _split_gated_scales + + +@cute.jit +def _row_tile( + mT: cute.Tensor, + row_tile: Int32, + cta_col: Int32, + TILER: cutlass.Constexpr, + ALIGNMENT: cutlass.Constexpr[int], +) -> cute.Tensor: + """The (row_tile, cta_col) tile of a row-major mT, which may be a column slice of a wider + tensor (row stride 2N): the tile's base takes the 64-bit product once and offsets within the + tile are 32-bit, as in the CUDA register-tiled kernel. Every row stride is N or 2N, and the C++ + dispatcher requires N % 32 == 0.""" + row_stride = cute.assume(Int32(mT.layout.stride[0]), divby=32) + base = Int64(row_tile) * Int64(TILER[0]) * Int64(row_stride) + Int64(cta_col) * Int64(TILER[1]) + return cute.make_tensor( + (mT.iterator + base).align(ALIGNMENT), + cute.make_layout(TILER, stride=(row_stride, 1)), + ) + + +class MXFP8GatedQuantizeG2RKernel: + """Register-resident gated MXFP8 cast: BF16 [M, 2N] -> E4M3 or E5M2, both directions or + columnwise only, any M, N % 32 == 0 (the C++ dispatcher's requirement).""" + + _CTA_THREADS_Y = 8 + _CTA_THREADS_X = 32 + _ELEMENTS_Y_PER_THREAD = 4 # 8 warps x 4 rows = one 32-row band + _ELEMENTS_X_PER_THREAD = 8 # 32 lanes x 8 columns = 256 columns + _ROWS_IN_FLIGHT = 4 # input rows loaded at a time + # Row bands per CTA, as in gated_bidim.cuh: halve the grid while it holds more than + # _WALK_TARGET_CTAS CTAs, walk at least _MIN_WALK bands unless that leaves fewer than two CTAs + # per SM, and at most _MAX_WALK. + _WALK_TARGET_CTAS = 16384 + _MIN_WALK = 2 + _MAX_WALK = 64 + + @staticmethod + def supports(cfg: MXFP8GatedQuantizeConfig) -> bool: + """BF16 with columnwise output (rowwise optional); GEMM-swizzled scales in the forward + only, like the shared-memory port.""" + return ( + cfg.DTYPE is cutlass.BFloat16 + and cfg.COLWISE + and not (cfg.IS_BWD and cfg.WITH_GEMM_SWIZZLED_SCALES) + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + mGatedInput: cute.Tensor, + mS_col: cute.Tensor, + ) -> Boolean: + # The C++ dispatcher checks 16-byte alignment of the data pointers. Columnwise scales are + # stored eight bytes per thread, and grid.y stays within limits. + M = mGatedInput.shape[0] + return (M + 31) // 32 <= 65535 and mS_col.iterator.toint() % 8 == 0 + + def __init__(self, cfg: MXFP8GatedQuantizeConfig): + assert self.supports(cfg), f"{cfg} is not supported by {type(self).__name__}" + self.cfg = cfg + # 4 resident CTAs per SM (64 registers), loading all 4 input rows at once. GB200, BF16, + # against gated_bidim.cuh, geometric mean over 5 shapes: forward 6 CTAs with 2 rows in + # flight (40 registers, spills) 0.91, 4 CTAs with 4 rows 0.93; backward 2 rows 0.92, + # 4 rows 0.94. + self._MIN_BLOCKS_PER_MP = 4 + from cuda.bindings import runtime as cudart # pylint: disable=no-name-in-module + + err, device_id = cudart.cudaGetDevice() + assert err == cudart.cudaError_t.cudaSuccess, f"CUDA runtime error {err}" + self._sm_count = HardwareInfo(device_id).get_device_multiprocessor_count() + + @cute.jit + def __call__( + self, + mGatedInput: cute.Tensor, # [M, 2N] = [act | gate] + mGrad: Optional[cute.Tensor], # [M, N], backward only + mO_row: Optional[cute.Tensor], # [M, C], C = N (forward) or 2N (backward) + mS_row: Optional[cute.Tensor], + mO_col: cute.Tensor, # [M, C], not transposed + mS_col: cute.Tensor, + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTUGLU only + beta2: Optional[Float32], # SiTUGLU only + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] MXFP8GatedQuantizeG2RKernel.__call__() with config: {self.cfg}\n" + ) + M = mGatedInput.shape[0] + N = mO_col.shape[1] if not self.cfg.IS_BWD else mGrad.shape[1] + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_X_PER_THREAD + + # eligible_to_dispatch checks this alignment: each thread stores eight columnwise scales. + mS_col = cute.make_tensor(mS_col.iterator.align(8), mS_col.layout) + mAct = cute.local_tile(mGatedInput, (M, N), (0, 0)) + mGate = cute.local_tile(mGatedInput, (M, N), (0, 1)) + + # Output parts: the forward output, or the two halves [d act | d gate] of the backward one. + mO_row1 = mS_row1 = mO_col1 = mS_col1 = None + if cutlass.const_expr(self.cfg.IS_BWD): + mS_row0, mS_row1, mS_col0, mS_col1 = _split_gated_scales( + mS_row, mS_col, M, N, self.cfg.ROWWISE, True, False + ) + mO_row0 = None + if cutlass.const_expr(self.cfg.ROWWISE): + mO_row0 = cute.local_tile(mO_row, (M, N), (0, 0)) + mO_row1 = cute.local_tile(mO_row, (M, N), (0, 1)) + mO_col0 = cute.local_tile(mO_col, (M, N), (0, 0)) + mO_col1 = cute.local_tile(mO_col, (M, N), (0, 1)) + else: + mO_row0, mO_col0 = mO_row, mO_col + mS_row0, mS_col0 = mS_row, mS_col + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row0, mS_col0 = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, True, mS_row, mS_col + ) + + grid_cols = (N + TILE_COLS - 1) // TILE_COLS + bands = (M + TILE_ROWS - 1) // TILE_ROWS + min_walk = Int32(self._MIN_WALK) + if grid_cols * (bands // self._MIN_WALK) < 2 * self._sm_count: + min_walk = Int32(1) + walk = Int32(1) + while (bands // walk) % 2 == 0 and grid_cols * (bands // walk) > self._WALK_TARGET_CTAS: + walk *= 2 + if walk < min_walk and bands % min_walk == 0: + walk = min_walk + if walk > self._MAX_WALK: + walk = Int32(self._MAX_WALK) + + # A tensor with a row or column tail takes the kernel whose tiles are predicated. + if M % TILE_ROWS != 0 or N % TILE_COLS != 0: + self.kernel( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + walk, + TAILS=True, + ).launch( + grid=[grid_cols, bands // walk], + block=[self._CTA_THREADS_Y * self._CTA_THREADS_X], + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) + else: + self.kernel( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + walk, + TAILS=False, + ).launch( + grid=[grid_cols, bands // walk], + block=[self._CTA_THREADS_Y * self._CTA_THREADS_X], + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mAct: cute.Tensor, + mGate: cute.Tensor, + mGrad: Optional[cute.Tensor], + mO_row0: Optional[cute.Tensor], + mS_row0: Optional[cute.Tensor], + mO_col0: cute.Tensor, + mS_col0: cute.Tensor, + mO_row1: Optional[cute.Tensor], + mS_row1: Optional[cute.Tensor], + mO_col1: Optional[cute.Tensor], + mS_col1: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + walk: Int32, + TAILS: cutlass.Constexpr[bool], + ): + """Quantize the CTA's row bands: blockIdx.y and every gridDim.y-th band after it, walk + bands in all.""" + cta_col, first_cta_row, _ = cute.arch.block_idx() + _, grid_rows, _ = cute.arch.grid_dim() + cache_policy = create_l2_policy(True) + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_X_PER_THREAD + for step in cutlass.range(walk, unroll=1): + row_tile = step * grid_rows + first_cta_row + # With TAILS, only the bands that cross the tensor's edge take the predicated body. + if cutlass.const_expr(TAILS): + if (row_tile + 1) * TILE_ROWS > mAct.shape[0] or ( + cta_col + 1 + ) * TILE_COLS > mAct.shape[1]: + self._gated_row_tile( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + row_tile, + cta_col, + cache_policy, + TAILS=True, + ) + else: + self._gated_row_tile( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + row_tile, + cta_col, + cache_policy, + TAILS=False, + ) + else: + self._gated_row_tile( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + row_tile, + cta_col, + cache_policy, + TAILS=False, + ) + + @cute.jit + def _quantize_part( + self, + rX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: cute.Tensor, + mS_col: cute.Tensor, + row_tile: Int32, + cta_col: Int32, + rows_left: Int32, + cols_ok: Boolean, + cache_policy: Int64, + TAILS: cutlass.Constexpr[bool], + ): + """Quantize one output part of the band from its register tile.""" + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_X_PER_THREAD + TILER = (TILE_ROWS, TILE_COLS) + cta_coord = (row_tile, cta_col) + gO_row_tile = gS_row_tile = None + if cutlass.const_expr(self.cfg.ROWWISE): + gO_row_tile = _row_tile(mO_row, row_tile, cta_col, TILER, 16) + gS_row_tile = cute.local_tile(mS_row, (TILE_ROWS, TILE_COLS // 32), cta_coord) + gO_col_tile = _row_tile(mO_col, row_tile, cta_col, TILER, 16) + gS_col_tile = cute.local_tile(mS_col, (TILE_ROWS // 32, TILE_COLS), cta_coord) + quantize_mxfp8_rmem_tile_bidimensional( + rX, + gO_row_tile, + gS_row_tile, + gO_col_tile, + gS_col_tile, + ROWWISE=self.cfg.ROWWISE, + COLWISE=True, + # The gated kernel's block amaxes ignore NaN like the generic kernel's. + FUSED=True, + EDGE=TAILS, + rows_left=rows_left, + cols_ok=cols_ok, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, + WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=self._ELEMENTS_X_PER_THREAD, + cache_policy=cache_policy, + ) + + @cute.jit + def _thread_view(self, mT: cute.Tensor, row_tile: Int32, cta_col: Int32) -> cute.Tensor: + """The calling thread's (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) values of the + (row_tile, cta_col) band of mT, in global memory.""" + ELEMENTS_Y = self._ELEMENTS_Y_PER_THREAD + ELEMENTS_X = self._ELEMENTS_X_PER_THREAD + TILER = (self._CTA_THREADS_Y * ELEMENTS_Y, self._CTA_THREADS_X * ELEMENTS_X) + tidx, _, _ = cute.arch.thread_idx() + tv_layout = thread_tile_layout( + self._CTA_THREADS_Y, self._CTA_THREADS_X, ELEMENTS_Y, ELEMENTS_X + ) + return cute.composition(_row_tile(mT, row_tile, cta_col, TILER, 16), tv_layout)[tidx, None] + + @cute.jit + def _load_rows( + self, + gT_thread: cute.Tensor, + first_row: cutlass.Constexpr[int], + rows_left: Int32, + cache_policy: Int64, + TAILS: cutlass.Constexpr[bool], + ) -> cute.Tensor: + """Rows first_row .. first_row + _ROWS_IN_FLIGHT - 1 of the thread's values, in + registers; with TAILS, rows from rows_left on load as zeros.""" + DTYPE = self.cfg.DTYPE + ELEMENTS_X = self._ELEMENTS_X_PER_THREAD + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=min(ELEMENTS_X * DTYPE.width, max_num_bits_per_copy_ld_global()), + invariant=False, + ) + rT = cute.make_rmem_tensor((ELEMENTS_X, self._ROWS_IN_FLIGHT), DTYPE) + for t in cutlass.range_constexpr(self._ROWS_IN_FLIGHT): + if cutlass.const_expr(TAILS): + rT[None, t].fill(0) + row_pred = cute.make_rmem_tensor( + cute.make_layout((ELEMENTS_X, (1,)), stride=(0, (0,))), Boolean + ) + row_pred[0] = first_row + t < rows_left + cute.copy( + load_atom, + gT_thread[(None, first_row + t),], + rT[None, t], + pred=row_pred, + cache_policy=cache_policy, + ) + else: + cute.copy( + load_atom, + gT_thread[(None, first_row + t),], + rT[None, t], + cache_policy=cache_policy, + ) + return rT + + @cute.jit + def _gated_row( + self, + rAct_row: cute.Tensor, + rGate_row: cute.Tensor, + rGrad_row: Optional[cute.Tensor], + rX0_row: cute.Tensor, + rX1_row: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + ): + """The gated op of one row of the thread's values, a BF16 pair at a time.""" + IS_BWD = self.cfg.IS_BWD + rA_i32 = cute.recast_tensor(rAct_row, Int32) + rG_i32 = cute.recast_tensor(rGate_row, Int32) + rX0_i32 = cute.recast_tensor(rX0_row, Int32) + if cutlass.const_expr(IS_BWD): + rD_i32 = cute.recast_tensor(rGrad_row, Int32) + rX1_i32 = cute.recast_tensor(rX1_row, Int32) + for i in cutlass.range_constexpr(self._ELEMENTS_X_PER_THREAD // 2): + grad_lo = grad_hi = None + if cutlass.const_expr(IS_BWD): + grad_lo = x2_lo_to_f32_bf16(rD_i32[i]) + grad_hi = x2_hi_to_f32_bf16(rD_i32[i]) + out_lo, gate_out_lo = _gated_element( + x2_lo_to_f32_bf16(rA_i32[i]), + x2_lo_to_f32_bf16(rG_i32[i]), + grad_lo, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=self.cfg.GATED_ACTIVATION, + IS_BWD=IS_BWD, + ) + out_hi, gate_out_hi = _gated_element( + x2_hi_to_f32_bf16(rA_i32[i]), + x2_hi_to_f32_bf16(rG_i32[i]), + grad_hi, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=self.cfg.GATED_ACTIVATION, + IS_BWD=IS_BWD, + ) + rX0_i32[i] = cvt_f32x2_to_bf16x2(out_lo, out_hi) + if cutlass.const_expr(IS_BWD): + rX1_i32[i] = cvt_f32x2_to_bf16x2(gate_out_lo, gate_out_hi) + + @cute.jit + def _gated_row_tile( + self, + mAct: cute.Tensor, + mGate: cute.Tensor, + mGrad: Optional[cute.Tensor], + mO_row0: Optional[cute.Tensor], + mS_row0: Optional[cute.Tensor], + mO_col0: cute.Tensor, + mS_col0: cute.Tensor, + mO_row1: Optional[cute.Tensor], + mS_row1: Optional[cute.Tensor], + mO_col1: Optional[cute.Tensor], + mS_col1: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + row_tile: Int32, + cta_col: Int32, + cache_policy: Int64, + TAILS: cutlass.Constexpr[bool], + ): + """Load one band's inputs, compute the gated op in registers and quantize its parts.""" + IS_BWD = self.cfg.IS_BWD + DTYPE = self.cfg.DTYPE + ELEMENTS_Y = self._ELEMENTS_Y_PER_THREAD + ELEMENTS_X = self._ELEMENTS_X_PER_THREAD + TILE_ROWS = self._CTA_THREADS_Y * ELEMENTS_Y + TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X + TILER = (TILE_ROWS, TILE_COLS) + tv_layout = thread_tile_layout( + self._CTA_THREADS_Y, self._CTA_THREADS_X, ELEMENTS_Y, ELEMENTS_X + ) + + # With TAILS every band is predicated: rows and columns past the tensor load as zeros, + # which every gated op maps to zero, and are not stored. + rows_left = Int32(ELEMENTS_Y) + cols_ok = Boolean(True) + if cutlass.const_expr(TAILS): + rows_left, cols_ok = thread_tile_bounds( + TILER, + tv_layout, + mAct.shape[0] - row_tile * TILE_ROWS, + mAct.shape[1] - cta_col * TILE_COLS, + ) + + gAct = self._thread_view(mAct, row_tile, cta_col) + gGate = self._thread_view(mGate, row_tile, cta_col) + gGrad = None + if cutlass.const_expr(IS_BWD): + gGrad = self._thread_view(mGrad, row_tile, cta_col) + + # The gated op in FP32, rounded to BF16 as gated_mxfp8.cuh does before taking amaxes: + # rX0 is the forward output or the act half's gradient, rX1 the gate half's gradient. + # Inputs are loaded _ROWS_IN_FLIGHT rows at a time, as in gated_bidim.cuh, which keeps + # fewer of them live. + rX0 = cute.make_rmem_tensor((ELEMENTS_X, ELEMENTS_Y), DTYPE) + rX1 = None + if cutlass.const_expr(IS_BWD): + rX1 = cute.make_rmem_tensor((ELEMENTS_X, ELEMENTS_Y), DTYPE) + for first_row in cutlass.range_constexpr(0, ELEMENTS_Y, self._ROWS_IN_FLIGHT): + rAct = self._load_rows(gAct, first_row, rows_left, cache_policy, TAILS=TAILS) + rGate = self._load_rows(gGate, first_row, rows_left, cache_policy, TAILS=TAILS) + rGrad = None + if cutlass.const_expr(IS_BWD): + rGrad = self._load_rows(gGrad, first_row, rows_left, cache_policy, TAILS=TAILS) + for t in cutlass.range_constexpr(self._ROWS_IN_FLIGHT): + rGrad_row = rX1_row = None + if cutlass.const_expr(IS_BWD): + rGrad_row = rGrad[None, t] + rX1_row = rX1[None, first_row + t] + self._gated_row( + rAct[None, t], + rGate[None, t], + rGrad_row, + rX0[None, first_row + t], + rX1_row, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ) + + self._quantize_part( + rX0, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + row_tile, + cta_col, + rows_left, + cols_ok, + cache_policy, + TAILS=TAILS, + ) + if cutlass.const_expr(IS_BWD): + self._quantize_part( + rX1, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + row_tile, + cta_col, + rows_left, + cols_ok, + cache_policy, + TAILS=TAILS, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py new file mode 100644 index 00000000000..46c635e3815 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -0,0 +1,821 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Gated-activation MXFP8 casts in CuTeDSL, ported from gated_mxfp8.cuh. + +The kernel reuses MXFP8QuantizeGeneralG2SKernel's TMA pipeline and its plain-cast SMEM tile +quantizer. What it adds is an elementwise pass that runs on each staged tile before it is +quantized: it combines the act and gate halves of the input (and the incoming gradient, in +backward) and writes the result back into the act (and gate) input buffer in the input dtype. +That is exactly the truncation the CUDA kernel applies before computing block amaxes, so the +plain cast that follows reproduces its output. + +Gated input is [M, 2N] with [act | gate] along the last dim. Every operand is addressed through +(M, N) views (stride 2N, column offset 0 or N), mirroring the CUDA kernel's tensor maps: + forward : out[M, N] = act(x[:, :N]) * gate(x[:, N:]) + backward: out[M, 2N] = [d/d act | d/d gate] given grad[M, N] +""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Boolean, Float32, Int32 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + fabs_f32, + is_packed16, + truncate_f32_bf16, + truncate_f32_f16, +) +from transformer_engine.common.CuTeDSL.activations import ( + act_clamped_silu, + act_situ_gate, + dact_dsitu_gate, + situ_up, + dsitu_up, + sigmoidf, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + GATED_ACTIVATIONS, + MXFP8GatedQuantizeConfig, + MXFP8_BLOCK_SCALING_SIZE, + derive_swizzled_scale_layout, +) +from .quantize_mxfp8_g2s_general import ( + MXFP8QuantizeGeneralG2SKernel, + quantize_mxfp8_smem_tile_general, +) + + +@cute.jit +def _select_f32(pred: Boolean, if_true: Float32, if_false: Float32) -> Float32: + result = if_false + if pred: + result = if_true + return result + + +@cute.jit +def _gated_element( + act_elt: Float32, + gate_elt: Float32, + grad_elt: Optional[Float32], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + ACTIVATION: cutlass.Constexpr[str], + IS_BWD: cutlass.Constexpr[bool], +): + """One element of the gated op in FP32; returns (act output, gate output or None). + + Mirrors the elementwise math of gated_mxfp8.cuh's columnwise pass (forward and backward), + operation by operation. + """ + out_gate = None + # Transform the gate. Gating is ideally the identity, with derivative 1. + gate = gate_elt + dgate = None + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + # |g| > limit is gated_mxfp8.cuh's (g > limit || g < -limit), NaN included + dgate = _select_f32(fabs_f32(gate_elt) > limit, Float32(0.0), Float32(1.0)) + gate = cute.arch.fmin(cute.arch.fmax(-limit, gate_elt), limit) + glu_linear_offset + elif cutlass.const_expr(ACTIVATION == "situglu"): + dgate = dsitu_up(gate_elt, beta2) + gate = situ_up(gate_elt, beta2) + + if cutlass.const_expr(not IS_BWD): + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + act = act_clamped_silu(act_elt, limit, alpha) + elif cutlass.const_expr(ACTIVATION == "situglu"): + act = act_situ_gate(act_elt, beta1) + else: + act = GATED_ACTIVATIONS[ACTIVATION][0](act_elt) + out_act = act * gate + else: + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + x = cute.arch.fmin(act_elt, limit) + s = sigmoidf(alpha * x) + act_x = x * s + dact_x = _select_f32( + act_elt <= limit, s + s * (Float32(1.0) - s) * alpha * x, Float32(0.0) + ) + elif cutlass.const_expr(ACTIVATION == "situglu"): + act_x = act_situ_gate(act_elt, beta1) + dact_x = dact_dsitu_gate(act_elt, beta1) + elif cutlass.const_expr(ACTIVATION == "swiglu"): + s = sigmoidf(act_elt) + act_x = act_elt * s + dact_x = act_elt * s * (Float32(1.0) - s) + s + else: + act_fn, dact_fn = GATED_ACTIVATIONS[ACTIVATION] + act_x = act_fn(act_elt) + dact_x = dact_fn(act_elt) + out_act = dact_x * grad_elt * gate + if cutlass.const_expr(dgate is None): + out_gate = act_x * grad_elt + else: + out_gate = dgate * act_x * grad_elt + return out_act, out_gate + + +@cute.jit +def gated_elementwise_smem_tile( + sAct_tile: cute.Tensor, + sGate_tile: cute.Tensor, + sGrad_tile: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + ACTIVATION: cutlass.Constexpr[str], + IS_BWD: cutlass.Constexpr[bool], + DTYPE: cutlass.Constexpr[type], + TILE_X: cutlass.Constexpr[int], +): + """Compute the gated op on one staged tile, in place, rounded to the input dtype. + + Thread tidx owns column tidx of the tile, which is also the element set it owns in the + colwise quantization pass, so a colwise-only consumer needs no barrier after this. + """ + tidx, _, _ = cute.arch.thread_idx() + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), + val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), + ) + sAct_thread = cute.composition(sAct_tile, tv_layout)[tidx, None] + sGate_thread = cute.composition(sGate_tile, tv_layout)[tidx, None] + if cutlass.const_expr(IS_BWD): + sGrad_thread = cute.composition(sGrad_tile, tv_layout)[tidx, None] + if cutlass.const_expr(is_packed16(DTYPE)): + truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 + + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + grad_elt = None + if cutlass.const_expr(IS_BWD): + grad_elt = Float32(sGrad_thread[i]) + out_act, out_gate = _gated_element( + Float32(sAct_thread[i]), + Float32(sGate_thread[i]), + grad_elt, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=ACTIVATION, + IS_BWD=IS_BWD, + ) + # Numerical truncation: round to the input dtype, like the CUDA kernels do before + # taking the amax. The cast below is then exact. + if cutlass.const_expr(is_packed16(DTYPE)): + out_act = truncate_f32(out_act) + if cutlass.const_expr(IS_BWD): + out_gate = truncate_f32(out_gate) + sAct_thread[i] = DTYPE(out_act) + if cutlass.const_expr(IS_BWD): + sGate_thread[i] = DTYPE(out_gate) + + +@cute.jit +def quantize_mxfp8_smem_tile_general_partial( + sX_tile: cute.Tensor, + sO_row_tile: Optional[cute.Tensor], + mS_row_stage: Optional[cute.Tensor], + sO_col_tile: Optional[cute.Tensor], + mS_col_stage: Optional[cute.Tensor], + tile_row_start: Int32, + tile_col_start: Int32, + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[type], + FP8_DTYPE: cutlass.Constexpr[type], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], +): + """A partial quantize_mxfp8_smem_tile_general function with some args fixed.""" + quantize_mxfp8_smem_tile_general( + sX_tile, + None, + sO_row_tile, + mS_row_stage, + sO_col_tile, + mS_col_stage, + None, + tile_row_start, + tile_col_start, + M, + N, + ROWWISE=ROWWISE, + COLWISE=COLWISE, + WITH_AMAX=False, + WITH_ACT=False, + WITH_DACT=False, + ACTIVATION=None, + ACTIVATION_TABLE="none", + MAX_NORM_RCP=MAX_NORM_RCP, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + # The gated op maps the zeros TMA fills out-of-bounds elements with to zero. + SKIP_INPUT_MASKING=True, + SKIP_SCALE_BOUNDS=False, + CACHE_ACTIVATION=False, + DBIAS_REDUCTION_IN_COLWISE=False, + COLWISE_DBIAS_REDUCTION_ONLY=False, + DBIAS_REDUCTION_IN_ROWWISE=False, + ) + + +@cute.jit +def _split_gated_scales( + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + SWIZZLED: cutlass.Constexpr[bool], +): + """Act and gate views of the backward output's scales; returns + (act_row, gate_row, act_col, gate_col). + + The [M, 2N] output is [d act | d gate] and its scales are laid out for the whole output. Each + view has the shape of an [M, N] tensor's scales and the strides of the whole output's, and + starts where the whole output's layout places that half's first scale. For swizzled scales + this needs the halves to be whole 128x4 tiles (N % 128 == 0, required by the C++ dispatcher). + """ + act_row = gate_row = act_col = gate_col = None + if cutlass.const_expr(SWIZZLED): + mS_row, mS_col = derive_swizzled_scale_layout(M, 2 * N, ROWWISE, COLWISE, mS_row, mS_col) + half_tiles = N // 128 + if cutlass.const_expr(ROWWISE): + row_tiler = ( + cute.make_layout((32, 4, cute.ceil_div(M, 128))), # all rows + cute.make_layout((4, half_tiles)), # one half's scale columns + ) + act_row = cute.local_tile(mS_row, row_tiler, (0, 0)) + gate_row = cute.local_tile(mS_row, row_tiler, (0, 1)) + if cutlass.const_expr(COLWISE): + col_tiler = ( + cute.make_layout((4, cute.ceil_div(cute.ceil_div(M, 32), 4))), # all scale rows + cute.make_layout((32, 4, half_tiles)), # one half's columns + ) + act_col = cute.local_tile(mS_col, col_tiler, (0, 0)) + gate_col = cute.local_tile(mS_col, col_tiler, (0, 1)) + else: + # Compact scales are row-major, so each half is a block of scale columns + if cutlass.const_expr(ROWWISE): + row_tiler = (mS_row.shape[0], N // MXFP8_BLOCK_SCALING_SIZE) + act_row = cute.local_tile(mS_row, row_tiler, (0, 0)) + gate_row = cute.local_tile(mS_row, row_tiler, (0, 1)) + if cutlass.const_expr(COLWISE): + col_tiler = (mS_col.shape[0], N) + act_col = cute.local_tile(mS_col, col_tiler, (0, 0)) + gate_col = cute.local_tile(mS_col, col_tiler, (0, 1)) + return act_row, gate_row, act_col, gate_col + + +class MXFP8GatedQuantizeG2SKernel(MXFP8QuantizeGeneralG2SKernel): + """gated_mxfp8.cuh in CuTeDSL. The general kernel's pipeline, loading the act and gate halves + (and grad, backward) each stage and storing one (forward) or two (backward) quantized halves. + + Inherits the plain-cast tile config, which is gated_mxfp8.cuh's: 64 threads per CTA, each CTA + covering two vertically stacked 32x64 tiles. + """ + + def __init__(self, cfg: MXFP8GatedQuantizeConfig): + super().__init__(cfg) + # Same occupancy hint as the general kernel with a fused activation + self._MIN_BLOCKS_PER_MP = 1 + # Every pipeline stage has its own SMEM buffers, so a CTA with fewer tiles than stages + # would allocate buffers it never fills + assert ( + self._NUM_STAGES <= self._NUM_TILES + ), f"{self._NUM_STAGES} pipeline stages for {self._NUM_TILES} tiles per CTA" + + @cute.jit + def __call__( + self, + mGatedInput: cute.Tensor, # [M, 2N] = [act | gate] + mGrad: Optional[cute.Tensor], # [M, N], bwd only + # Outputs, with C = N (FWD) or 2N (BWD) + mO_row: Optional[cute.Tensor], # [M, C] + mS_row: Optional[cute.Tensor], # [roundup(M, 128), roundup(C / 32, 4)] + mO_col: Optional[cute.Tensor], # [M, C], not transposed + mS_col: Optional[cute.Tensor], # [roundup(ceil(M / 32), 4), roundup(C, 128)] + # Special activation params, None if not used by the activation + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTUGLU only + beta2: Optional[Float32], # SiTUGLU only + stream: CUstream, + ): + """ + mGatedInput: [a | g] with shape M x 2N, where a is the activation half and g is the gate half. + mGrad: None for FWD, [dy] with shape [M, N] for BWD + FWD mode: + out = quantize(act(a) * gate(g)) with shape [M, N] + BWD mode: + out = quantize([ dy * gate(g) * dact(a) | dy * act(a) * dgate(g) ]) with shape [M, 2N] + """ + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] MXFP8GatedQuantizeG2SKernel.__call__() with config: {self.cfg}\n" + ) + + M = mGatedInput.shape[0] + # Pick any mO that is used just to get N + mO = mO_row if mO_row is not None else mO_col + # FWD: mO is [M, N]; BWD: mGrad is [M, N] + N = mO.shape[1] if not self.cfg.IS_BWD else mGrad.shape[1] + + # The act and gate halves of the [M, 2N] input + mAct = cute.local_tile(mGatedInput, (M, N), (0, 0)) + mGate = cute.local_tile(mGatedInput, (M, N), (0, 1)) + + mO_row_dact = mO_row_dgate = mO_col_dact = mO_col_dgate = None # BWD only + mS_row_dact = mS_row_dgate = mS_col_dact = mS_col_dgate = None # BWD only + if cutlass.const_expr(self.cfg.IS_BWD): + # The [M, 2N] output is [d act | d gate]; the kernel takes its two halves + # TODO(kainingz): how to implement this in a more elegant way + mS_row_dact, mS_row_dgate, mS_col_dact, mS_col_dgate = _split_gated_scales( + mS_row, + mS_col, + M, + N, + self.cfg.ROWWISE, + self.cfg.COLWISE, + self.cfg.WITH_GEMM_SWIZZLED_SCALES, + ) + if cutlass.const_expr(self.cfg.ROWWISE): + mO_row_dact = cute.local_tile(mO_row, (M, N), (0, 0)) + mO_row_dgate = cute.local_tile(mO_row, (M, N), (0, 1)) + if cutlass.const_expr(self.cfg.COLWISE): + mO_col_dact = cute.local_tile(mO_col, (M, N), (0, 0)) + mO_col_dgate = cute.local_tile(mO_col, (M, N), (0, 1)) + # In BWD mode we won't use these, and we should set them to None so we don't build TMA atoms for them + mO_row = mO_col = mS_row = mS_col = None + elif cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, mS_col = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col + ) + + atoms = self._make_tma_atoms( + mGrad, + mAct, + mGate, + mO_row, + mO_col, + mO_row_dact, + mO_row_dgate, + mO_col_dact, + mO_col_dgate, + ) + # The TMA atoms are spread from a list, which pylint cannot count + self.kernel( # pylint: disable=no-value-for-parameter + M, + N, + mS_row, + mS_col, + mS_row_dact, + mS_row_dgate, + mS_col_dact, + mS_col_dgate, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + *atoms, + ).launch( + grid=[ + cute.ceil_div(N, self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), + ], + block=[self._THREADS_PER_CTA], + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) + + @cute.jit + def _make_tma_atoms( + self, + mGrad: Optional[cute.Tensor], + mAct: cute.Tensor, + mGate: cute.Tensor, + mO_row: Optional[cute.Tensor], # FWD only + mO_col: Optional[cute.Tensor], # FWD only + mO_row_dact: Optional[cute.Tensor], # BWD only + mO_row_dgate: Optional[cute.Tensor], # BWD only + mO_col_dact: Optional[cute.Tensor], # BWD only + mO_col_dgate: Optional[cute.Tensor], # BWD only + ): + """Build TMA atoms for input / quantized output; scales will not be loaded / stored by TMA.""" + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + atoms = [] + for op, mT in ( + (op_load, mGrad), + (op_load, mAct), + (op_load, mGate), + (op_store, mO_row), + (op_store, mO_col), + (op_store, mO_row_dact), + (op_store, mO_row_dgate), + (op_store, mO_col_dact), + (op_store, mO_col_dgate), + ): + if cutlass.const_expr(mT is None): + atoms.extend([None, None]) + else: + atoms.extend( + cute.nvgpu.cpasync.make_tiled_tma_atom( + op, mT, smem_tile_layout, cta_tiler, num_multicast=1 + ) + ) + return atoms + + @cute.kernel + def kernel( + self, + M: Int32, + N: Int32, + mS_row: Optional[cute.Tensor], # FWD only + mS_col: Optional[cute.Tensor], # FWD only + mS_row_dact: Optional[cute.Tensor], # BWD only + mS_row_dgate: Optional[cute.Tensor], # BWD only + mS_col_dact: Optional[cute.Tensor], # BWD only + mS_col_dgate: Optional[cute.Tensor], # BWD only + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + # TMA for mGrad, only used in BWD + tma_atom_grad: Optional[cute.CopyAtom], + tma_src_grad: Optional[cute.Tensor], + # TMA for mAct and mGate, used in both FWD and BWD + tma_atom_act: cute.CopyAtom, + tma_src_act: cute.Tensor, + tma_atom_gate: cute.CopyAtom, + tma_src_gate: cute.Tensor, + # TMA for mO_row/col, only used in FWD + tma_atom_row: Optional[cute.CopyAtom], + tma_dst_row: Optional[cute.Tensor], + tma_atom_col: Optional[cute.CopyAtom], + tma_dst_col: Optional[cute.Tensor], + # TMA for mO_row/col_dact/dgate, only used in BWD + tma_atom_row_dact: Optional[cute.CopyAtom], + tma_dst_row_dact: Optional[cute.Tensor], + tma_atom_row_dgate: Optional[cute.CopyAtom], + tma_dst_row_dgate: Optional[cute.Tensor], + tma_atom_col_dact: Optional[cute.CopyAtom], + tma_dst_col_dact: Optional[cute.Tensor], + tma_atom_col_dgate: Optional[cute.CopyAtom], + tma_dst_col_dgate: Optional[cute.Tensor], + ): + """ + FWD mode: + mAct: activation half input + mGate: gate half input + mO_row/col: quantized output for row / column directions + mS_row/col: output scales for row / column directions + + BWD mode: + mGrad: incoming gradient + mAct: activation half input from FWD pass + mGate: gate half input from FWD pass + mO_row/col_dact: quantized activation gradient for row / column directions + mO_row/col_dgate: quantized gate gradient for row / column directions + mS_row/col_dact: output scales for activation gradient row / column directions + mS_row/col_dgate: output scales for gate gradient row / column directions + mO_row/col: quantized gradient (both activation and gate) for row / column directions + mS_row/col: output scales for gradient (both activation and gate) row / column directions + + mAct: both FWD and BWD + mGate: both FWD and BWD + mGrad: BWD only + mO_row/col: both FWD (N) and BWD (2N) + mS_row/col: both FWD (N) and BWD (2N) + mO_row/col_dact: BWD only + mO_row/col_dgate: BWD only + mS_row/col_dact: BWD only + mS_row/col_dgate: BWD only + """ + + cfg = self.cfg + DTYPE = cfg.DTYPE + TILE_ROWS = self._TILE_ROWS + TILE_COLS = self._TILE_COLS + NUM_STAGES = self._NUM_STAGES + NUM_TILES = self._NUM_TILES + + # The scales of one (TILE_ROWS, TILE_COLS) tile + row_scale_tile = (TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + col_scale_tile = (TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) + + # Shared memory: one ((TILE_ROWS, TILE_COLS), NUM_STAGES) ring buffer per operand + smem = cutlass.memory.SmemAllocator() + + @cute.struct + class BarrierStorage: # pylint: disable=missing-class-docstring + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * NUM_STAGES] + + storage = smem.allocate(BarrierStorage) + stages_layout = cute.make_layout( + ((TILE_ROWS, TILE_COLS), NUM_STAGES), + stride=((TILE_COLS, 1), TILE_ROWS * TILE_COLS), + ) + + # Allocate SMEM buffers for inputs + sAct = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + sGate = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + sGrad = None + if cutlass.const_expr(cfg.IS_BWD): + sGrad = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + + # Allocate SMEM buffers for outputs + sO_row = sO_col = None + sO_row_dact = sO_row_dgate = sO_col_dact = sO_col_dgate = None + if cutlass.const_expr(cfg.ROWWISE): + if cutlass.const_expr(cfg.IS_BWD): + sO_row_dact = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + sO_row_dgate = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + else: + sO_row = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + if cutlass.const_expr(cfg.COLWISE): + if cutlass.const_expr(cfg.IS_BWD): + sO_col_dact = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + sO_col_dgate = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + else: + sO_col = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bidx, bidy, _ = cute.arch.block_idx() + + # Issue TMA descriptor prefetch using the TMA warp (warp 0) + if warp_idx == 0: + for load_atom in (tma_atom_act, tma_atom_gate, tma_atom_grad): + if cutlass.const_expr(load_atom is not None): + cute.nvgpu.cpasync.prefetch_descriptor(load_atom) + + # TMA warp (warp 0)'s elected thread produces (issues TMA), every warp consumes + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) + # Input has 2 tiles (act + gate) in FWD, 3 tiles (grad + act + gate) in BWD + num_inputs = 3 if cfg.IS_BWD else 2 + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_storage.data_ptr(), + num_stages=NUM_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + # Every input tile of a stage arrives on the same barrier + tx_count=num_inputs * TILE_ROWS * TILE_COLS * DTYPE.width // 8, + cta_layout_vmnk=None, + ) + prod_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, NUM_STAGES) + cons_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, NUM_STAGES) + + cta_coord = (bidy, bidx) + # Shape of each individual tile. A CTA processes one tile at a time + tile_shape = (TILE_ROWS, TILE_COLS) + # The block processed by a CTA, consisting of NUM_TILES vertically stacked tiles + cta_block = (NUM_TILES * TILE_ROWS, TILE_COLS) + + # Prepare parameters to pass to the SMEM tile quantize function by dact / dgate / output, where + # for each group, input, quantized output are SMEM tiles, and scales are GMEM tiles + cta_block_row_scales = (NUM_TILES * TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + cta_block_col_scales = (NUM_TILES * TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) + if cutlass.const_expr(cfg.IS_BWD): + if cutlass.const_expr(cfg.ROWWISE): + mS_row_dact = cute.local_tile(mS_row_dact, cta_block_row_scales, cta_coord) + mS_row_dgate = cute.local_tile(mS_row_dgate, cta_block_row_scales, cta_coord) + if cutlass.const_expr(cfg.COLWISE): + mS_col_dact = cute.local_tile(mS_col_dact, cta_block_col_scales, cta_coord) + mS_col_dgate = cute.local_tile(mS_col_dgate, cta_block_col_scales, cta_coord) + # In BWD mode we output gradient for both the act and gate halves, which are stored in sAct and sGate respectively + tile_quantizer_params = [ + ( + sAct, + sO_row_dact, + mS_row_dact, + sO_col_dact, + mS_col_dact, + ), # We will write unquantized dact to sAct in-place + ( + sGate, + sO_row_dgate, + mS_row_dgate, + sO_col_dgate, + mS_col_dgate, + ), # We will write unquantized dgate to sGate in-place + ] + else: + if cutlass.const_expr(cfg.ROWWISE): + mS_row = cute.local_tile(mS_row, cta_block_row_scales, cta_coord) + if cutlass.const_expr(cfg.COLWISE): + mS_col = cute.local_tile(mS_col, cta_block_col_scales, cta_coord) + # In FWD mode we only output the result after the gated op, which is stored in sAct + tile_quantizer_params = [ + (sAct, sO_row, mS_row, sO_col, mS_col) + ] # We will write unquantized output to sAct in-place + + # Prepare input TMA (atom, smem partition, gmem partition) triples: act and gate, and BWD's grad + inputs = [] + for atom_in, src_in, sX_in in ( + (tma_atom_act, tma_src_act, sAct), + (tma_atom_gate, tma_src_gate, sGate), + (tma_atom_grad, tma_src_grad, sGrad), + ): + if cutlass.const_expr(atom_in is not None): + tXsX_in, tXgX_in = cute.nvgpu.cpasync.tma_partition( + atom_in, + 0, + cute.make_layout(1), + sX_in, + cute.zipped_divide(cute.local_tile(src_in, cta_block, cta_coord), tile_shape), + ) + inputs.append((atom_in, tXsX_in, tXgX_in)) + + # Prepare output TMA (atom, smem partition, gmem partition) triples of the outputs this config has + TMA_atom_and_partitions = [] + for atom_out, dst_out, sO_out in ( + (tma_atom_row, tma_dst_row, sO_row), # FWD only + (tma_atom_col, tma_dst_col, sO_col), # FWD only + (tma_atom_row_dact, tma_dst_row_dact, sO_row_dact), # BWD only + (tma_atom_row_dgate, tma_dst_row_dgate, sO_row_dgate), # BWD only + (tma_atom_col_dact, tma_dst_col_dact, sO_col_dact), # BWD only + (tma_atom_col_dgate, tma_dst_col_dgate, sO_col_dgate), # BWD only + ): + if cutlass.const_expr(atom_out is not None): + tXsO_out, tXgO_out = cute.nvgpu.cpasync.tma_partition( + atom_out, + 0, + cute.make_layout(1), + sO_out, + cute.zipped_divide(cute.local_tile(dst_out, cta_block, cta_coord), tile_shape), + ) + TMA_atom_and_partitions.append((atom_out, tXsO_out, tXgO_out)) + + # Barrier init must be visible to every thread + cute.arch.sync_threads() + + # Prologue: fill the pipeline + if warp_idx == 0: + for s in cutlass.range_constexpr(NUM_STAGES): + mainloop_pipeline.producer_acquire(prod_state) + for load_atom, load_sX, load_gX in inputs: + cute.copy( + load_atom, + load_gX[(None, (s, 0))], + load_sX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + for tile_idx in cutlass.range(NUM_TILES, unroll=1): + mainloop_pipeline.consumer_wait(cons_state) + # We only reuse previously used SMEM buffers when we have more tiles to process than pipeline stages, and only + # after we've already finished one round and wrap around to the start of the ring buffer + if cutlass.const_expr(NUM_TILES > NUM_STAGES): + if tile_idx >= NUM_STAGES: + # Only allow at most NUM_STAGES - 1 store groups reading SMEM on the fly, so the tile NUM_STAGES + # ago must finish before we reuse its SMEM buffer + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + # The tile should be ready to use now + + stage_idx = cons_state.index + tile_coord = (tile_idx, 0) # within this CTA's block + + # Select current stage's SMEM tile + sAct_tile = sAct[(None, stage_idx)] + sGate_tile = sGate[(None, stage_idx)] + sGrad_tile = None + if cutlass.const_expr(cfg.IS_BWD): + sGrad_tile = sGrad[(None, stage_idx)] + + # Compute the gated op on this tile in fp32, cast back to the original dtype, + # and store the result in-place in the SMEM tile without quantization. + # FWD: write the output into sAct_tile + # BWD: write the act half's gradient to sAct_tile, and write the gate half's gradient to sGate_tile + gated_elementwise_smem_tile( + sAct_tile, # we also store the FWD output / BWD act half's gradient here + sGate_tile, # we also store the BWD gate half's gradient here + sGrad_tile, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=cfg.GATED_ACTIVATION, + IS_BWD=cfg.IS_BWD, + DTYPE=DTYPE, + TILE_X=TILE_COLS, + ) + + # The rowwise quantization pass reads elements other threads wrote so we need to sync here + if cutlass.const_expr(cfg.ROWWISE): + cute.arch.sync_threads() + + # Quantize every output of this tile into its SMEM buffers and scales + for sX_out, sO_row_out, mS_row_cta, sO_col_out, mS_col_cta in tile_quantizer_params: + # Slice the SMEM / GMEM tile for the current iteration + sO_row_tile = None + mS_row_stage = None + if cutlass.const_expr(cfg.ROWWISE): + sO_row_tile = sO_row_out[(None, stage_idx)] + mS_row_stage = cute.flatten( + cute.local_tile(mS_row_cta, row_scale_tile, tile_coord) + ) + sO_col_tile = None + mS_col_stage = None + if cutlass.const_expr(cfg.COLWISE): + sO_col_tile = sO_col_out[(None, stage_idx)] + mS_col_stage = cute.flatten( + cute.local_tile(mS_col_cta, col_scale_tile, tile_coord) + ) + # Use the general MXFP8 SMEM tile quantize function to quantize it + quantize_mxfp8_smem_tile_general_partial( + sX_out[(None, stage_idx)], # The SMEM tile to quantize + sO_row_tile, # SMEM tile for rowwise quantized output + mS_row_stage, # GMEM tile for rowwise quantized scales + sO_col_tile, # SMEM tile for columnwise quantized output + mS_col_stage, # GMEM tile for columnwise quantized scales + # The tile's first row and column in the whole tensor, for bounds checks + (bidy * NUM_TILES + tile_idx) * TILE_ROWS, + bidx * TILE_COLS, + M, + N, + ROWWISE=cfg.ROWWISE, + COLWISE=cfg.COLWISE, + MAX_NORM_RCP=cfg.MAX_NORM_RCP, + DTYPE=DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_COLS, + TILE_Y=TILE_ROWS, + ) + + # Make the SMEM writes visible to the TMA's async proxy before it reads them + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + mainloop_pipeline.consumer_release(cons_state) + + if warp_idx == 0: + # The tile quantize function has already written the quantized outputs to SMEM, + # so we can issue TMA stores from SMEM to GMEM + for store_atom, store_sO, store_gO in TMA_atom_and_partitions: + cute.copy( + store_atom, + store_sO[(None, stage_idx)], + store_gO[(None, (tile_idx, 0))], + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # Refill the stage just released with the tile NUM_STAGES ahead + if warp_idx == 0: + next_tile_idx = tile_idx + NUM_STAGES + if next_tile_idx < NUM_TILES: + mainloop_pipeline.producer_acquire(prod_state) + for load_atom, load_sX, load_gX in inputs: + cute.copy( + load_atom, + load_gX[(None, (next_tile_idx, 0))], + load_sX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # Only wait until the in-flight TMA stores have read SMEM, so the CTA retires early + cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/transformer_engine/common/cast/dispatch/gated.cuh b/transformer_engine/common/cast/dispatch/gated.cuh index 0c1761264a1..02a0740f635 100644 --- a/transformer_engine/common/cast/dispatch/gated.cuh +++ b/transformer_engine/common/cast/dispatch/gated.cuh @@ -19,6 +19,10 @@ #include "../fp8/gated_fp8.cuh" #include "../mxfp8/gated_mxfp8.cuh" +#ifdef NVTE_WITH_CUTEDSL +#include "../mxfp8/gated_mxfp8_cutedsl.cuh" +#endif + namespace transformer_engine { namespace dispatch { @@ -89,8 +93,16 @@ void quantize_gated_fwd_helper(const NVTETensor nvte_input, NVTETensor nvte_outp NVTE_CHECK(is_supported_by_CC_100(), "Gated FWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); Tensor dummy_grad_tensor; - mxfp8::quantize_gated(input, dummy_grad_tensor, - output, p, stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_gated_cutedsl( + input, dummy_grad_tensor, output, p, stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::quantize_gated(input, dummy_grad_tensor, + output, p, stream); + } break; } default: @@ -182,8 +194,16 @@ void quantize_gated_bwd_helper(const NVTETensor nvte_grad, const NVTETensor nvte NVTE_CHECK(is_supported_by_CC_100(), "Gated BWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); - mxfp8::quantize_gated(gated_input, grad, output, p, - stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_gated_cutedsl( + gated_input, grad, output, p, stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::quantize_gated(gated_input, grad, output, p, + stream); + } break; } default: diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh new file mode 100644 index 00000000000..485c04c5452 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh @@ -0,0 +1,203 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file gated_mxfp8_cutedsl.cuh + * \brief CuTeDSL backend of gated_mxfp8.cuh. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ + +#include +#include + +#include +#include +#include +#include +#include + +#include "../../common.h" +#include "../../tvm_ffi_bridge.h" +#include "../../util/cutedsl_utils.h" + +namespace transformer_engine { +namespace cutedsl_backend { + +using namespace tvm_ffi_bridge; + +namespace detail { + +// TMA needs 16-byte aligned global addresses. +inline bool is_16B_aligned(const void *ptr) { return reinterpret_cast(ptr) % 16 == 0; } + +} // namespace detail + +struct MXFP8GatedQuantConfig { + static constexpr const char *kEntrypointName = "get_mxfp8_gated_quantization_function"; + + DType dtype; // The input format + DType fp8_dtype; // The fp8 output format + bool rowwise; // If quantize rowwisely + bool colwise; // If quantize columnwisely + bool swizzled; // If the scale output is used for cudnn's swizzled layout + bool is_bwd; // Backward (dgated) rather than forward + GatedActivation activation; + uint32_t sm_arch = static_cast(cuda::sm_arch()); + + // Bit layout: dtype [3:0], fp8_dtype [7:4], flags [11:8], activation [17:12] and + // SM architecture [26:18]. + uint32_t to_id() const { + static_assert(static_cast(DType::kNumTypes) <= 16, + "DType no longer fits in the 4 bits to_id() gives it."); + static_assert(static_cast(GatedActivation::kNumTypes) <= 64, + "GatedActivation no longer fits in the 6 bits to_id() gives it."); + NVTE_CHECK(sm_arch < 512, "SM architecture no longer fits in the 9 bits to_id() gives it."); + return static_cast(dtype) | (static_cast(fp8_dtype) << 4) | + (static_cast(rowwise) << 8) | (static_cast(colwise) << 9) | + (static_cast(swizzled) << 10) | (static_cast(is_bwd) << 11) | + (static_cast(activation) << 12) | (sm_arch << 18); + } + + std::optional get_kernel() const { + static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); + return cache.get_or_load(*this); + } + + std::string to_key() const { + std::string key; + key.reserve(80); + key.append("cutedsl_mxfp8_gated_sm") + .append(std::to_string(sm_arch)) + .append("_") + .append(to_string(dtype)) + .append("_") + .append(to_string(fp8_dtype)) + .append("_") + .append(rowwise ? "1" : "0") + .append("_") + .append(colwise ? "1" : "0") + .append("_") + .append(swizzled ? "1" : "0") + .append("_") + .append(is_bwd ? "1" : "0") + .append("_") + .append(gated_activation_to_str(activation)); + return key; + } + + bool retrieve_func_from_python(const std::string &fn_name) const { + auto entrypoint = tvm::ffi::Function::GetGlobal(kEntrypointName); + if (!entrypoint.has_value()) { + return false; + } + tvm::ffi::Any result = + (*entrypoint)(tvm::ffi::String(fn_name), tvm::ffi::String(to_string(dtype)), + tvm::ffi::String(to_string(fp8_dtype)), rowwise, colwise, swizzled, is_bwd, + tvm::ffi::String(gated_activation_to_str(activation))); + return result.try_cast().value_or(false); + } +}; + +// Signature mirrors mxfp8::quantize_gated. Returns false to fall back to the CUDA kernel. +template +bool mxfp8_quantize_gated_cutedsl(const Tensor &gated_input, const Tensor &grad, Tensor *output, + const ParamOP &p, cudaStream_t stream) { + if (!TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { + return false; + } + constexpr GatedActivation activation = gated_activation_func_to_enum(); + if constexpr (activation == GatedActivation::kUnsupported) { + report_cutedsl_not_chosen("the gated activation is not supported."); + return false; + } else { + const size_t rows = gated_input.flat_first_dim(); + const size_t cols = gated_input.flat_last_dim() / 2; + const MXFP8GatedQuantConfig config{/*dtype=*/gated_input.dtype(), + /*fp8_dtype=*/output->dtype(), + /*rowwise=*/output->has_data(), + /*colwise=*/output->has_columnwise_data(), + /*swizzled=*/output->with_gemm_swizzled_scales, + /*is_bwd=*/IS_BWD, + /*activation=*/activation}; + // TODO: backward gated activations with GEMM-swizzled scales are not supported in CuTeDSL yet + // (MXFP8GatedQuantizeConfig rejects them); the CUDA kernel handles them. + if (IS_BWD && config.swizzled) { + return false; + } + bool aligned = detail::is_16B_aligned(gated_input.data.dptr) && + (!config.rowwise || detail::is_16B_aligned(output->data.dptr)) && + (!config.colwise || detail::is_16B_aligned(output->columnwise_data.dptr)); + if constexpr (IS_BWD) { + aligned = aligned && detail::is_16B_aligned(grad.data.dptr); + } + if (!aligned) { + report_cutedsl_not_chosen("an input or output is not 16-byte aligned."); + return false; + } + + std::optional fn = config.get_kernel(); + if (!fn.has_value()) { + return false; + } + checkCuDriverContext(stream); + + // Sanity checks and swizzled-scale padding, mirroring mxfp8::quantize_gated + if (config.rowwise) { + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); + } + if (config.colwise) { + NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); + } + if (config.swizzled && (cols % 128 != 0 || rows % 128 != 0)) { + if (config.rowwise) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->scale_inv.dptr, 0, + output->scale_inv.buffer_size_bytes(), stream)); + } + if (config.colwise) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->columnwise_scale_inv.dptr, 0, + output->columnwise_scale_inv.buffer_size_bytes(), stream)); + } + } + + // The activation's own parameters are runtime arguments; the kernel is compiled with None in + // the slots of the others, so they are passed as None (nullptr). + tvm::ffi::Any limit = nullptr, alpha = nullptr, glu_linear_offset = nullptr, beta1 = nullptr, + beta2 = nullptr; + if constexpr (std::is_same_v) { + limit = p.limit; + alpha = p.alpha; + glu_linear_offset = p.glu_linear_offset; + } else if constexpr (std::is_same_v) { + beta1 = p.beta1; + beta2 = p.beta2; + } + + const int32_t device_index = transformer_engine::cuda::current_device(); + DLTensorWrapper mGatedInput(gated_input.data, true, device_index); + DLTensorWrapper mGrad, mO_row, mS_row, mO_col, mS_col; + if constexpr (IS_BWD) { + mGrad = DLTensorWrapper(grad.data, true, device_index); + } + if (config.rowwise) { + mO_row = DLTensorWrapper(output->data, true, device_index); + mS_row = DLTensorWrapper(output->scale_inv, true, device_index); + } + if (config.colwise) { + mO_col = DLTensorWrapper(output->columnwise_data, true, device_index); + mS_col = DLTensorWrapper(output->columnwise_scale_inv, true, device_index); + } + (*fn)(&mGatedInput, &mGrad, &mO_row, &mS_row, &mO_col, &mS_col, limit, alpha, glu_linear_offset, + beta1, beta2, static_cast(stream)); + return true; + } +} + +} // namespace cutedsl_backend +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ diff --git a/transformer_engine/common/util/cutedsl_utils.h b/transformer_engine/common/util/cutedsl_utils.h index 4fd4ee56f69..4af0d78e906 100644 --- a/transformer_engine/common/util/cutedsl_utils.h +++ b/transformer_engine/common/util/cutedsl_utils.h @@ -90,6 +90,76 @@ NVTE_SPECIALIZE_CUTEDSL_ACTIVATION_FUNC_TO_ENUM(dsrelu, kDSReLU) #undef NVTE_SPECIALIZE_CUTEDSL_ACTIVATION_FUNC_TO_ENUM +// Encodes the gated activation selected by the gated dispatcher (gated_mxfp8.cuh), identified by +// its forward ActOP and, for backward, its DActOP. +enum class GatedActivation { + kGeGLU, + kQGeGLU, + kReGLU, + kSReGLU, + kSwiGLU, + kGLU, + kSiTUGLU, + kClampedSwiGLU, + kUnsupported, + kNumTypes +}; + +inline const char *gated_activation_to_str(GatedActivation act) { + switch (act) { + case GatedActivation::kGeGLU: + return "geglu"; + case GatedActivation::kQGeGLU: + return "qgeglu"; + case GatedActivation::kReGLU: + return "reglu"; + case GatedActivation::kSReGLU: + return "sreglu"; + case GatedActivation::kSwiGLU: + return "swiglu"; + case GatedActivation::kGLU: + return "glu"; + case GatedActivation::kSiTUGLU: + return "situglu"; + case GatedActivation::kClampedSwiGLU: + return "clamped_swiglu"; + case GatedActivation::kUnsupported: + case GatedActivation::kNumTypes: + return "unsupported"; + } + return "unsupported"; +} + +// DActOP is nullptr for the forward. +template +constexpr GatedActivation gated_activation_func_to_enum() { + return GatedActivation::kUnsupported; +} + +#define NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(PARAM, ACT, DACT, ACTIVATION) \ + template <> \ + constexpr GatedActivation gated_activation_func_to_enum, nullptr>() { \ + return GatedActivation::ACTIVATION; \ + } \ + template <> \ + constexpr GatedActivation \ + gated_activation_func_to_enum, DACT>() { \ + return GatedActivation::ACTIVATION; \ + } + +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, gelu, dgelu, kGeGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, qgelu, dqgelu, kQGeGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, relu, drelu, kReGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, srelu, dsrelu, kSReGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, silu, dsilu, kSwiGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, sigmoid, dsigmoid, kGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(SiTUGLUParam, situ_gate, dsitu_gate, kSiTUGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(ClampedSwiGLUParam, clamped_silu, + clamped_dsilu, kClampedSwiGLU) + +#undef NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM + } // namespace cutedsl_backend } // namespace transformer_engine