diff --git a/docs/envvars.rst b/docs/envvars.rst index a0a64b3f637..1836e80e1bd 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -353,6 +353,18 @@ 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: ``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 + + :Type: ``int`` (0 or 1) + :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_USE_CUTLASS_GROUPED_GEMM :Type: ``int`` (0 or 1) diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index ecb746d7b55..7e3f9a91fdb 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..1b523cc4bab --- /dev/null +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -0,0 +1,510 @@ +/************************************************************************* + * 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 (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. +// +// 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. + +#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 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 dsilu_window_low (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 < dsilu_window_low ? 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 rowwise_reduction (CUDA, for 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; +} + +// 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); + } + ~ScopedEnv() { + if (old_) { + setenv(name_, old_->c_str(), 1); + } else { + unsetenv(name_); + } + } + + 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; + 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 otype, const bool with_amax) { + const DType itype = TypeInfo::dtype; + 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; + // 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, 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); + 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); + 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; + { + const ScopedEnv tables("NVTE_ACTIVATION_TABLE", use_table ? "1" : "0"); + 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(); + // 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(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(), + ref.rowwise_cpu_scale_inv_ptr(), s), + 0u) + << "rowwise scales"; + } + if (colwise) { + 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(), + ref.columnwise_cpu_scale_inv_ptr(), s), + 0u) + << "colwise scales"; + } + if (is_dbias(mode)) { + // CuTeDSL's general kernel adds FP32 rows in row order too (no rowwise dbias pass). + const char *cutedsl = std::getenv("NVTE_ENABLE_CUTEDSL_BACKEND"); + const bool cutedsl_on = cutedsl != nullptr && std::string(cutedsl) == "1"; + const bool rowwise_reduction = std::is_same_v && !colwise && !cutedsl_on; + 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*/, DType /*output*/, bool /*amax*/>; + +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, 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, + 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, 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") + (otype == DType::kFloat8E5M2 ? "XE5M2" : "") + + (with_amax ? "XAmax" : ""); +} + +} // 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}, + // 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}, + // 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. + 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( + 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(), + ::testing::Values(DType::kFloat8E4M3), ::testing::Values(false)), + 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(), ::testing::Values(DType::kFloat8E4M3), + ::testing::Values(false)), + 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..009fa322eb5 --- /dev/null +++ b/tests/cpp/util/test_activation_table.cu @@ -0,0 +1,265 @@ +/************************************************************************* + * 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 +// 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. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#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; + +// 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 Spec; +template <> +struct Spec { + static constexpr float (*op)(float, const Empty &) = transformer_engine::gelu; + static constexpr bool kForward = true; +}; +template <> +struct Spec { + static constexpr float (*op)(float, const Empty &) = transformer_engine::silu; + static constexpr bool kForward = true; +}; +template <> +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 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 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(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(bits, window_low); + math_h[bits] = S::kForward ? 0.0f : S::op(x, {}); +} + +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; +} + +// kComputed: check what the generic kernel computes in place of the table (kComputedTable, or +// 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); + ASSERT_NE(table.ptr, nullptr); + 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 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) { + launch(std::integral_constant{}); + } + } else { + ASSERT_EQ(table.use, at::kDefaultTable); + launch(std::integral_constant{}); + } + 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(Spec::window_low() == Spec::kExactWindowLow ? 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(Spec::window_low() == Spec::kExactWindowLow ? 0 : 2); +} + +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); +} + +// 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 = [] { + setenv("NVTE_ACTIVATION_TABLE", "1", 1); + 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/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 2ee82f8c5d5..313bdfe1be5 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -176,7 +176,7 @@ def get_cfg_key(method, act_desc, in_dtype, fp8_dtype, q_layout): # trailing False is use_2d_quantization; JAX never requests 2D block scaling flags = (True, q_layout.has_colwise, False, False, with_dbias, with_dact, with_act, False) major, minor = device_compute_capability() - return ( + key = ( f"cutedsl_mxfp8_sm{major * 10 + minor}_" + DTYPE_TO_STR[in_dtype] + "_" @@ -186,6 +186,15 @@ def get_cfg_key(method, act_desc, in_dtype, fp8_dtype, q_layout): + "_" + desc ) + # Kernels that read a BF16 activation table carry its layout in the name. + if ( + in_dtype == jnp.bfloat16 + and desc in ("gelu", "silu", "dgelu", "dsilu") + and int(os.getenv("NVTE_ACTIVATION_TABLE") or 0) + ): + exact = desc == "dsilu" and int(os.getenv("NVTE_ACTIVATION_TABLE_EXACT_DSILU") or 0) + key += "_table_exact" if exact else "_table_default" + return key def extract_quantized_output(out, dbias): diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 334c8b7bcbc..b48bcb68328 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -232,7 +232,7 @@ def get_cfg_key(method, act, in_dtype, fp8_dtype, rowwise, colwise, swizzled): # The last False is for use_2d_quantization; these cases never request 2D block scaling flags = (rowwise, colwise, swizzled, False, with_dbias, with_dact, with_act, False) major, minor = device_compute_capability() - return ( + key = ( f"cutedsl_mxfp8_sm{major * 10 + minor}_" + DTYPE_TO_STR[in_dtype] + "_" @@ -242,6 +242,15 @@ def get_cfg_key(method, act, in_dtype, fp8_dtype, rowwise, colwise, swizzled): + "_" + desc ) + # Kernels that read a BF16 activation table carry its layout in the name. + if ( + in_dtype == torch.bfloat16 + and desc in ("gelu", "silu", "dgelu", "dsilu") + and int(os.getenv("NVTE_ACTIVATION_TABLE") or 0) + ): + exact = desc == "dsilu" and int(os.getenv("NVTE_ACTIVATION_TABLE_EXACT_DSILU") or 0) + key += "_table_exact" if exact else "_table_default" + return key def extract_quantized_output(out, rowwise, columnwise, swizzled): @@ -271,16 +280,18 @@ def extract_quantized_output(out, rowwise, columnwise, swizzled): def run_test_case( - method, act, shape, block_size, in_dtype, fp8_dtype, swizzled=False, input_offset=0 + method, act, shape, block_size, in_dtype, fp8_dtype, swizzled=False, input_offset=0, x=None ): """Assert the CuTeDSL and CUDA backends produce bit-identical outputs for the same input and config. input_offset (in elements) shifts x's data pointer past the - start of its allocation, to test misaligned inputs. + start of its allocation, to test misaligned inputs. x, if given, replaces the generated input. """ M, N = shape rowwise = block_size[1] != 1 columnwise = block_size[0] != 1 - x, act_input = generate_inputs(M, N, in_dtype) + generated_x, act_input = generate_inputs(M, N, in_dtype) + if x is None: + x = generated_x if input_offset: buf = torch.empty(input_offset + M * N, dtype=in_dtype, device="cuda") x = buf[input_offset:].view(M, N).copy_(x) @@ -433,6 +444,27 @@ def test_specialized_kernels(shape, block_size, in_dtype, swizzled, input_offset ) +@pytest.mark.parametrize("shape", [(512, 1024), (512, 768)], ids=get_shape_id) +@pytest.mark.parametrize("block_size", [BOTH, COL], ids=get_block_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_silu_fast_reciprocal_rows(monkeypatch, swizzled, block_size, shape): + """SiLU forward without a table on the register-resident kernel: a thread row takes the + sigmoid's fast reciprocal only if all its values are at least -80, and -100 changes the result + on that path. Each 16 columns of a row hold one -100, at another position in each of 16 + consecutive rows, so the row check must read every value. N = 1024 runs 512-column tiles, + N = 768 256-column tiles.""" + monkeypatch.setenv("NVTE_ACTIVATION_TABLE", "0") + M, N = shape + x, _ = generate_inputs(M, N, torch.bfloat16) + rows = torch.arange(M, device="cuda").unsqueeze(1) + cols = torch.arange(N, device="cuda") + x = x.masked_fill(cols % 16 == rows % 16, -100.0) + silu = next(a for a in ACTIVATIONS if a["name"] == "SiLU") + run_test_case( + "CAST_ACT", silu, shape, block_size, torch.bfloat16, tex.DType.kFloat8E4M3, swizzled, x=x + ) + + OUT_DTYPES = [torch.float32, torch.bfloat16, torch.float16] DEQUANT_SHAPES = MATRIX_SIZES + [(100, 48)] diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index f90f74449c5..9525e28b3b2 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 @@ -772,7 +773,9 @@ 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 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 activation/gelu_grouped.cu activation/gelu_grouped_bwd.cu diff --git a/transformer_engine/common/CuTeDSL/activations.py b/transformer_engine/common/CuTeDSL/activations.py index 34772eb7257..88585972932 100644 --- a/transformer_engine/common/CuTeDSL/activations.py +++ b/transformer_engine/common/CuTeDSL/activations.py @@ -8,14 +8,20 @@ # 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 +from transformer_engine.common.CuTeDSL.utils import fma_f32, rcp_f32_in_range USE_FAST_MATH = os.environ.get("NVTE_USE_FAST_MATH", "0") == "1" +# For x >= SIGMOID_FAST_RCP_MIN, 1 <= 1 + exp(-x) < 2^126 (also with fast_math), where +# rcp_f32_in_range gives the same result as the division. +SIGMOID_FAST_RCP_MIN = -80.0 + def tanh(x: Float32) -> Float32: return cute.math.tanh(x, fastmath=USE_FAST_MATH) @@ -36,7 +42,7 @@ def act_gelu(x: Float32) -> Float32: def act_silu(x: Float32) -> Float32: - return x / (Float32(1.0) + exp(-x)) + return x * sigmoid(x) def act_qgelu(x: Float32) -> Float32: @@ -57,8 +63,13 @@ def dact_dsrelu(x: Float32) -> Float32: return cute.arch.fmax(Float32(2.0) * x, Float32(0.0)) -def sigmoid(x: Float32) -> Float32: - return Float32(1.0) / (Float32(1.0) + exp(-x)) +def sigmoid(x: Float32, fast_rcp: bool = False) -> Float32: + """1 / (1 + exp(-x)); fast_rcp skips the reciprocal's range check and requires + x >= SIGMOID_FAST_RCP_MIN.""" + y = Float32(1.0) + exp(-x) + if fast_rcp: + return rcp_f32_in_range(y) + return Float32(1.0) / y def dsigmoid(x: Float32) -> Float32: @@ -80,3 +91,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. The table is ForwardLayout 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. The table is DerivativeLayout 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..9100ea82f82 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,28 +8,46 @@ 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 ( abs_max_x2_bf16, + exp2f_rcp, + fabs_f32, bf16_pair_magnitude, bf16_mx_scale_reciprocal, bf16_mx_scale_reciprocal_pair, bf16_mx_scale_bytes_pair, create_l2_policy, max_num_bits_per_copy_ld_global, + min_nan_bf16x2, x2_hi_to_f32_bf16, x2_lo_to_f32_bf16, cvt_f32x2_to_bf16x2, ) -from transformer_engine.common.CuTeDSL.utils_fp8 import mul_bf16x2x2_cvt_bf16x4_to_fp8x4 +from transformer_engine.common.CuTeDSL.activations import ( + DERIVATIVE_SLOPES, + DERIVATIVE_TABLE_WINDOWS, + FORWARD_TABLE_WINDOWS, + SIGMOID_FAST_RCP_MIN, + activation_table_bytes, + sigmoid, + table_derivative_2x, + table_forward_2x, +) +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, MXFP8QuantizeKernelBase, - SUPPORTED_ACTIVATIONS, - SUPPORTED_DACTIVATIONS, + activation_op, + dactivation_op, derive_swizzled_scale_layout, ) @@ -51,7 +69,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]], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, ): + """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 @@ -85,17 +106,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(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 @@ -153,19 +196,89 @@ def _gather_partial_colwise_amaxes( return sAmaxs_col_scratch +@cute.jit +def _reduce_column_pairs( + sAmaxs_col_frags_i32: cute.Tensor, + sS_col_i32: cute.Tensor, + gS_col_tile: Optional[cute.Tensor], + col_scale_atom: Optional[cute.CopyAtom], + cache_policy: Int64, + cols_valid: Int32, + # Constexprs + REDUCTION_THREADS_NUM: cutlass.Constexpr[int], + REDUCTION_i32_PAIRS: cutlass.Constexpr[int], + CTA_THREADS_Y: cutlass.Constexpr[int], + MAX_EXPONENT: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + EDGE: cutlass.Constexpr[bool], +): + """The calling reducer's part of _reduce_partial_colwise_amaxes: the reciprocals of its column + pairs to sS_col_i32 and, with gS_col_tile, their scales to gS_col_tile.""" + tidx, _, _ = cute.arch.thread_idx() + WRITE_SCALES = gS_col_tile is not None + sAmaxs_col_frag_i32 = sAmaxs_col_frags_i32[tidx, None] + # Expose the thread's value mode as (partial rows, column pairs). + sAmaxs_col_frag_i32 = cute.composition( + sAmaxs_col_frag_i32, + cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PAIRS), stride=(REDUCTION_i32_PAIRS, 1)), + ) + # The reducer's scales: its 2 * REDUCTION_i32_PAIRS adjacent columns of the scale tile. + REDUCER_COLS = 2 * REDUCTION_i32_PAIRS + if cutlass.const_expr(WRITE_SCALES): + rS_col = cute.make_rmem_tensor((REDUCER_COLS,), SCALE_DTYPE) + rS_col_u16 = cute.recast_tensor(rS_col, Uint16) + + # Iterative over all column pairs + for pair_idx in cutlass.range_constexpr(REDUCTION_i32_PAIRS): + # Select the pair column + sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pair_idx] + + # Reduce across all rows' partial amaxes to obtain the columnwise amax for this column pair + amaxX2 = sAmaxs_col_frag_i32_rows[0] + 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]) + + # Obtain the columnwise scale reciprocal for this column pair + reciprocal = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + sS_col_i32[tidx, pair_idx] = reciprocal + if cutlass.const_expr(WRITE_SCALES): + rS_col_u16[pair_idx] = Uint16(bf16_mx_scale_bytes_pair(reciprocal)) + if cutlass.const_expr(WRITE_SCALES): + gS_col_reducer = cute.composition( + gS_col_tile, + cute.make_layout((REDUCTION_THREADS_NUM, REDUCER_COLS), stride=(REDUCER_COLS, 1)), + )[tidx, None] + if cutlass.const_expr(EDGE): + if tidx * REDUCER_COLS < cols_valid: + cute.copy(col_scale_atom, rS_col, gS_col_reducer, cache_policy=cache_policy) + else: + cute.copy(col_scale_atom, rS_col, gS_col_reducer, cache_policy=cache_policy) + + @cute.jit def _reduce_partial_colwise_amaxes( sAmaxs_col_scratch: cute.Tensor, + gS_col_tile: Optional[cute.Tensor], + # Atoms + col_scale_atom: Optional[cute.CopyAtom], + # Miscellaneous + cache_policy: Int64, # Constexprs CTA_THREADS_Y: cutlass.Constexpr[int], CTA_THREADS_X: cutlass.Constexpr[int], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + EDGE: cutlass.Constexpr[bool] = False, + cols_valid: Int32 = 0, ) -> cute.Tensor: """ 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 + gS_col_tile: None, or the (1, TILE_COLS) columnwise scales of the tile in GMEM, which each + reducer then writes for its own columns. With EDGE, only the tile's first cols_valid + columns are inside the tensor and written. This function reduces the partial amax buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) along the Y axis, to a single columnwise scale SMEM buffer of layout (ELEMENTS_X_PER_THREAD * CTA_THREADS_X), @@ -218,28 +331,47 @@ def _reduce_partial_colwise_amaxes( # to extract the amax of that column sAmaxs_col_scratch_i32 = cute.recast_tensor(sAmaxs_col_scratch, Int32) MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 - if is_reduction_thread: - sAmaxs_col_frag_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_amaxes_layout)[ - tidx, None - ] - # Expose the thread's value mode as (partial rows, column pairs). - sAmaxs_col_frag_i32 = cute.composition( - sAmaxs_col_frag_i32, - cute.make_layout((CTA_THREADS_Y, REDUCTION_i32_PAIRS), stride=(REDUCTION_i32_PAIRS, 1)), - ) - - # Iterative over all column pairs - for pair_idx in cutlass.range_constexpr(REDUCTION_i32_PAIRS): - # Select the pair column - sAmaxs_col_frag_i32_rows = sAmaxs_col_frag_i32[None, pair_idx] - # Reduce across all rows' partial amaxes to obtain the columnwise amax for this column pair - amaxX2 = sAmaxs_col_frag_i32_rows[0] - 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]) - - # Obtain the columnwise scale reciprocal for this column pair - sS_col_i32[tidx, pair_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + sAmaxs_col_frags_i32 = cute.composition(sAmaxs_col_scratch_i32, tv_col_amaxes_layout) + if cutlass.const_expr(REDUCTION_THREADS_NUM == CTA_THREADS_Y * CTA_THREADS_X): + if is_reduction_thread: + _reduce_column_pairs( + sAmaxs_col_frags_i32, + sS_col_i32, + gS_col_tile, + col_scale_atom, + cache_policy, + cols_valid, + REDUCTION_THREADS_NUM=REDUCTION_THREADS_NUM, + REDUCTION_i32_PAIRS=REDUCTION_i32_PAIRS, + CTA_THREADS_Y=CTA_THREADS_Y, + MAX_EXPONENT=MAX_EXPONENT, + SCALE_DTYPE=SCALE_DTYPE, + EDGE=EDGE, + ) + else: + # The reducers fill whole warps. Each runs the body once and the other warps skip it: a + # loop of one trip or none, because ptxas predicates the body of an `if` in every warp. + for _ in cutlass.range( + cute.arch.make_warp_uniform(cute.arch.warp_idx()), + REDUCTION_THREADS_NUM // 32, + CTA_THREADS_Y * CTA_THREADS_X // 32, + unroll=1, + ): + _reduce_column_pairs( + sAmaxs_col_frags_i32, + sS_col_i32, + gS_col_tile, + col_scale_atom, + cache_policy, + cols_valid, + REDUCTION_THREADS_NUM=REDUCTION_THREADS_NUM, + REDUCTION_i32_PAIRS=REDUCTION_i32_PAIRS, + CTA_THREADS_Y=CTA_THREADS_Y, + MAX_EXPONENT=MAX_EXPONENT, + SCALE_DTYPE=SCALE_DTYPE, + EDGE=EDGE, + ) cute.arch.sync_threads() return sS_col @@ -257,9 +389,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() @@ -277,7 +412,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 @@ -294,8 +433,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]], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, ): """ + 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,) gO_col_thread: GMEM fragment of this thread's columnwise quantized output with shape @@ -320,7 +463,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(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 @@ -328,6 +479,7 @@ def _quantize_colwise( rX: cute.Tensor, gO_col_thread: cute.Tensor, gS_col_thread: cute.Tensor, + gS_col_tile: cute.Tensor, # Atoms output_store_atom: cute.CopyAtom, col_scale_store_atom: cute.CopyAtom, @@ -342,8 +494,15 @@ def _quantize_colwise( SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, + cols_ok: Boolean = True, + cols_valid: Int32 = 0, ): - """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. + The scales go to gS_col_thread (the thread's columns, written by the first warp row) or + gS_col_tile; with EDGE, the tile's columns from cols_valid on are past the tensor.""" tidx, _, _ = cute.arch.thread_idx() sAmaxs_col_scratch = _gather_partial_colwise_amaxes( @@ -355,13 +514,34 @@ def _quantize_colwise( DTYPE=DTYPE, ) + # The reducers write the scales of their own columns when the scales are GEMM-swizzled outside + # the predicated (EDGE) tiles (a byte at a time, spread over all reducer warps) or the column + # pairs are fewer than the threads (the reducers then fill only some warps, and the other warps + # skip both). Otherwise, as on 512-column EDGE tiles, the first warp row writes the scales. + REDUCERS_WRITE_SCALES = (WITH_GEMM_SWIZZLED_SCALES and not EDGE) or ( + CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 < CTA_THREADS_Y * CTA_THREADS_X + ) + gS_col_reducers = reducer_scale_atom = None + if cutlass.const_expr(REDUCERS_WRITE_SCALES): + gS_col_reducers = gS_col_tile + reducer_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + SCALE_DTYPE, + num_bits_per_copy=8, # byte stores: 16-bit compact stores slow the tile's row loads + ) sS_col = _reduce_partial_colwise_amaxes( sAmaxs_col_scratch, + gS_col_reducers, + col_scale_atom=reducer_scale_atom, + cache_policy=cache_policy, CTA_THREADS_Y=CTA_THREADS_Y, CTA_THREADS_X=CTA_THREADS_X, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, DTYPE=DTYPE, FP8_DTYPE=FP8_DTYPE, + EDGE=EDGE, + cols_valid=cols_valid, ) # Copy the reduced columnwise scale reciprocals to RMEM. @@ -369,15 +549,18 @@ def _quantize_colwise( rS_col_rcp = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) cute.autovec_copy(sS_col[None, tidx % CTA_THREADS_X], rS_col_rcp) - _write_colwise_scale( - rS_col_rcp, - gS_col_thread, - col_scale_atom=col_scale_store_atom, - cache_policy=cache_policy, - CTA_THREADS_X=CTA_THREADS_X, - ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, - SCALE_DTYPE=SCALE_DTYPE, - ) + if cutlass.const_expr(not REDUCERS_WRITE_SCALES): + _write_colwise_scale( + rS_col_rcp, + gS_col_thread, + col_scale_atom=col_scale_store_atom, + cache_policy=cache_policy, + 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( rX, @@ -389,9 +572,200 @@ def _quantize_colwise( ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, FP8_DTYPE=FP8_DTYPE, + EDGE=EDGE, + rows_left=rows_left, + ) + + +# FP16 and FP32 inputs scale in FP32, as the CUDA generic kernel does (2^(127-e) is often outside +# FP16's range): a block's scale is e8m0(amax * MAX_NORM_RCP), the amax being NaN-seeded like the +# generic kernel's so that a block of only NaNs gets the NaN scale, 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(amax * Float32(MAX_NORM_RCP)) + 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(float("nan")) + 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(float("nan")) + 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(float("nan")) + 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 _row_takes_fast_rcp(rX_row: cute.Tensor) -> Int32: + """1 if every BF16 value of the thread row is at least SIGMOID_FAST_RCP_MIN (so none is NaN), + else 0.""" + rX_i32 = cute.recast_tensor(rX_row, Int32) + x_min = rX_i32[0] + for i in cutlass.range_constexpr(1, cute.size(rX_i32)): + x_min = min_nan_bf16x2(x_min, rX_i32[i]) + return Int32(x2_lo_to_f32_bf16(x_min) >= Float32(SIGMOID_FAST_RCP_MIN)) & Int32( + x2_hi_to_f32_bf16(x_min) >= Float32(SIGMOID_FAST_RCP_MIN) ) +@cute.jit +def _silu_row(rX_row: cute.Tensor, FAST_RCP: cutlass.Constexpr[bool]): + """act_silu of the thread row's BF16 values, in place. FAST_RCP: the sigmoid takes its fast + reciprocal (sigmoid's fast_rcp), for a row where _row_takes_fast_rcp holds.""" + rX_i32 = cute.recast_tensor(rX_row, Int32) + for i in cutlass.range_constexpr(cute.size(rX_i32)): + lo = x2_lo_to_f32_bf16(rX_i32[i]) + hi = x2_hi_to_f32_bf16(rX_i32[i]) + rX_i32[i] = cvt_f32x2_to_bf16x2( + lo * sigmoid(lo, fast_rcp=FAST_RCP), hi * sigmoid(hi, fast_rcp=FAST_RCP) + ) + + @cute.jit def quantize_mxfp8_gmem_tile_bidimensional( gX_tile: cute.Tensor, @@ -403,10 +777,18 @@ 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], + sDbias: Optional[cute.Tensor], + 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], 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], @@ -465,11 +847,41 @@ def quantize_mxfp8_gmem_tile_bidimensional( 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 + # the column amaxes and dbias unchanged, and are not stored. + 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) + # 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(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 + ) if cutlass.const_expr(WITH_ACT or WITH_DACT): rA = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), DTYPE) @@ -486,28 +898,127 @@ 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(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, + ) + # 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 + 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) + x = x * dop(Float32(rA[col, row_idx])) + if cutlass.const_expr(sDbias is not None): + rProducts[col] = x + else: + 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),]) + elif cutlass.const_expr(WITH_ACT and ACTIVATION == "silu" and ACTIVATION_TABLE == "none"): + # The sigmoid takes its fast reciprocal on rows whose values allow it. Each path is a + # loop of one trip or none, so a warp whose rows all take one path skips the other. + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + fast = _row_takes_fast_rcp(rX[None, row_idx]) + for _ in cutlass.range(fast, unroll=1): + _silu_row(rX[None, row_idx], FAST_RCP=True) + for _ in cutlass.range(Int32(1) - fast, unroll=1): + _silu_row(rX[None, row_idx], FAST_RCP=False) + 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. + 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 and DTYPE is not cutlass.BFloat16): 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) - 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])) - else: - op = SUPPORTED_ACTIVATIONS[ACTIVATION] - lo = op(lo) - hi = op(hi) - rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + 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)): + 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] @@ -530,20 +1041,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, - ) + 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] @@ -573,22 +1104,50 @@ 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, - ) + 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, + gS_col_tile, + 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, + WITH_GEMM_SWIZZLED_SCALES=WITH_GEMM_SWIZZLED_SCALES, + EDGE=EDGE, + rows_left=rows_left, + cols_ok=cols_ok, + cols_valid=cols_valid, + ) + + return thread_amax class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): @@ -604,17 +1163,37 @@ 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); each CTA walks the remaining row tiles. + _WALK_TARGET_COLS = 16384 * 256 @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; + # amax only with a fused activation; plain casts with amax take the general kernel. return ( - not cfg.WITH_AMAX - and not cfg.WITH_DBIAS - and cfg.DTYPE is cutlass.BFloat16 + (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 + 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, and FP32 dbias+dGeLU, whose dbias matches the CUDA + kernels' bits here and not in the general kernel. FP16 dbias+dGeLU and FP32 GeLU keep the + general kernel.""" + if not (cfg.WITH_ACT or cfg.WITH_DACT): + return False + 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") + @staticmethod @cute.jit def eligible_to_dispatch( @@ -631,22 +1210,52 @@ 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. + # 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 ( - M % 32 == 0 - and N % 256 == 0 - and M // 32 <= 65535 + shape_ok + 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_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. The others keep whole 32-row by + 256-column tiles.""" + 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. + 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. + self._sm_count = 0 + 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() + 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 +1269,32 @@ 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 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. + # 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, + mO_row, + mS_row, + mO_col, + mS_col, + mActTable, + mWorkspace, + mAmax, + 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 +1303,9 @@ def __call__( mS_row, mO_col, mS_col, + mActTable, + mWorkspace, + mAmax, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=4, @@ -686,6 +1318,9 @@ def __call__( mS_row, mO_col, mS_col, + mActTable, + mWorkspace, + mAmax, stream, ELEMENTS_X_PER_THREAD=16, CLUSTER_WIDTH=1, @@ -698,6 +1333,9 @@ def __call__( mS_row, mO_col, mS_col, + mActTable, + mWorkspace, + mAmax, stream, ELEMENTS_X_PER_THREAD=8, CLUSTER_WIDTH=8, @@ -712,6 +1350,9 @@ def _launch( mS_row: cute.Tensor, mO_col: cute.Tensor, 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], @@ -730,12 +1371,42 @@ def _launch( grid_cols = N // TILE_COLS grid_rows = M // TILE_ROWS + 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) + # Clusters of CTAs on adjacent column tiles speed up the memory reads that bound a cast. + # With a table, shared-memory lookups bound the tile instead and a cluster can leave CTA + # slots idle, so it launches unclustered. + if cutlass.const_expr(self.cfg.ACTIVATION_TABLE != "none"): + cluster_cols = Int32(1) # 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: 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" or self.cfg.WITH_DBIAS): + 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) + # 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 + 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. - grid = [grid_cols, grid_rows] + grid = [grid_cols, grid_rows // walk] block = [THREADS_PER_CTA] cluster = [cluster_cols, 1, 1] @@ -771,15 +1442,51 @@ 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, - 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 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, + gO_row, + gS_row, + gO_col, + gS_col, + mActTable, + walk, + mWorkspace, + mAmax, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + TAILS=True, + ).launch( + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + 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, + TAILS=False, + ).launch( + 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)) gX = cute.make_tensor(mX.iterator, payload_layout) @@ -811,15 +1518,51 @@ 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, - 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 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, + gO_row, + gS_row, + gO_col, + gS_col, + mActTable, + walk, + mWorkspace, + mAmax, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + TAILS=True, + ).launch( + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + 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, + TAILS=False, + ).launch( + 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): gS_row, gS_col = derive_swizzled_scale_layout( @@ -827,15 +1570,51 @@ def _launch( ) else: gS_row, gS_col = mS_row, mS_col - self.kernel( - mX, - mA, - mO_row, - gS_row, - mO_col, - gS_col, - 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 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, + mO_row, + gS_row, + mO_col, + gS_col, + mActTable, + walk, + mWorkspace, + mAmax, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + TAILS=True, + ).launch( + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + 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, + TAILS=False, + ).launch( + grid=grid, + block=block, + cluster=cluster, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) @cute.kernel def kernel( @@ -846,7 +1625,12 @@ def kernel( mS_row: cute.Tensor, mO_col: cute.Tensor, mS_col: cute.Tensor, + mActTable: cute.Pointer, + walk: Int32, + mWorkspace: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + 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.""" @@ -857,9 +1641,36 @@ def kernel( mS_row, mO_col, mS_col, + mActTable, + walk, + mWorkspace, + mAmax, ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + TAILS=TAILS, ) + @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: 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,34 +1680,195 @@ def _kernel_main( mS_row: cute.Tensor, mO_col: cute.Tensor, mS_col: cute.Tensor, + mActTable: cute.Pointer, + walk: Int32, + mWorkspace: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + TAILS: cutlass.Constexpr[bool] = False, ): - """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() + + 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 + # 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): + # 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(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, + TAILS=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 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( + 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, + TAILS=TAILS, + ) + thread_amax = cute.arch.fmax(thread_amax, tile_amax) + else: + for step in cutlass.range(walk, unroll=1): + tile_amax = 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, + TAILS=TAILS, + ) + 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( + 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], + sDbias: Optional[cute.Tensor] = None, + TAILS: cutlass.Constexpr[bool] = False, + ) -> 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) 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 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 + 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) 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( + # 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, gA_tile, gO_row_tile, @@ -906,10 +1878,18 @@ def _kernel_main( ROWWISE=self.cfg.ROWWISE, COLWISE=self.cfg.COLWISE, ACTIVATION=self.cfg.ACTIVATION, + ACTIVATION_TABLE=self.cfg.ACTIVATION_TABLE, + sTable=sTable, + sDbias=sDbias, + 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, 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, 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..b56477ee005 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], @@ -352,8 +354,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( @@ -372,7 +376,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 @@ -412,12 +417,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,12 +520,15 @@ 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, 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. @@ -530,9 +538,9 @@ 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 = 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])) @@ -576,11 +584,14 @@ 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], + ACTIVATION_TABLE: cutlass.Constexpr[str] = "none", + 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 @@ -597,6 +608,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, @@ -608,6 +620,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. @@ -621,10 +634,12 @@ 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, 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. @@ -642,6 +657,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 +758,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): @@ -1172,11 +1189,11 @@ 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' columnwise reduction does. 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): @@ -1235,6 +1252,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, @@ -1247,13 +1265,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( diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index e38112689c5..f87c83e0903 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -93,6 +93,28 @@ def fma_f32(a: Float32, b: Float32, c: Float32) -> Float32: ) +@cute.jit +def rcp_f32_in_range(y: Float32) -> Float32: + """1 / y rounded to nearest, for 1 <= y < 2^126: the approximate reciprocal refined by one + Newton step, as rcp.rn.f32 computes it in that range, without its range check.""" + return cute.arch.inline_ptx( + "{ .reg .f32 r, e; rcp.approx.ftz.f32 r, {$r0}; fma.rn.f32 e, {$r0}, r, 0fBF800000; " + "neg.f32 e, e; fma.rn.f32 {$w0}, r, e, r; }", + write_only_types=[Float32], + read_only_args=[y], + ) + + +@cute.jit +def min_nan_bf16x2(a: Int32, b: Int32) -> Int32: + """Lane-wise minimum of two packed BF16 pairs; a lane is NaN if either input lane is.""" + return cute.arch.inline_ptx( + "min.NaN.bf16x2 {$w0}, {$r0}, {$r1};", + write_only_types=[Int32], + read_only_args=[a, b], + ) + + @cute.jit def exp2f_rcp(scale_e8m0) -> Float32: """2^(127 - biased_exp) with special-case handling, for an e8m0 scale.""" diff --git a/transformer_engine/common/activation/activation_table.cu b/transformer_engine/common/activation/activation_table.cu new file mode 100644 index 00000000000..d9a82e1164a --- /dev/null +++ b/transformer_engine/common/activation/activation_table.cu @@ -0,0 +1,148 @@ +/************************************************************************* + * 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. 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 + +#include "../common.h" +#include "../util/activation_table.cuh" +#include "../util/cuda_runtime.h" + +namespace transformer_engine { +namespace activation_table { +namespace { + +constexpr int kBlockSize = 256; + +template +__global__ void build_forward_kernel(__nv_bfloat16 *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) { + 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) { + // 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)); + } else { + constexpr uint32_t entries = T::kSignStride + T::kEntriesPerSign; + build_derivative_kernel<<(entries, kBlockSize), kBlockSize, 0, stream>>>( + static_cast(table)); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void build(const Kind kind, void *table, cudaStream_t stream) { + switch (kind) { + case kGelu: + build>(table, stream); + break; + case kSilu: + build>(table, stream); + break; + case kDgelu: + build>(table, stream); + break; + case kDsilu: + if (use_exact_dsilu_table()) { + build>(table, stream); + } else { + build>(table, stream); + } + break; + default: + 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/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 13c3694d040..999ce51f770 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -38,6 +38,18 @@ 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, 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) { + if (input_dtype == DType::kBFloat16 && activation_table::tables_enabled()) { + return activation_table::device_table(activation_table::kKindOf); + } + } + return {}; +} + template void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const NVTEQuantizationConfig quant_config, cudaStream_t stream) { @@ -95,18 +107,22 @@ 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 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) { 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 +300,21 @@ 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 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) { 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..ce291a348d5 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, + 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, @@ -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 activation_table::TableRef activation_lut, cudaStream_t stream) { using namespace quantize_kernel; checkCuDriverContext(stream); @@ -1003,58 +1009,59 @@ 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 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(table_use)::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.ptr); + }; + using activation_table::kComputedTable; + using activation_table::kDefaultTable; + 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. The exact dSiLU + // table holds util/math.h itself. + if constexpr (kTableCapable) { + if (activation_lut.use == kDefaultTable) { + launch_kernel(std::integral_constant{}); + } else { + launch_kernel(std::integral_constant{}); + } + } else { + launch_kernel(std::integral_constant{}); + } + }); + }; - 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/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 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/util/activation_table.cuh b/transformer_engine/common/util/activation_table.cuh new file mode 100644 index 00000000000..709e09587c5 --- /dev/null +++ b/transformer_engine/common/util/activation_table.cuh @@ -0,0 +1,308 @@ +/************************************************************************* + * 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 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). 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 + * 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. 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. + * + * 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 "../common.h" +#include "math.h" +#include "system.h" + +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 + * 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; + //! 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; +}; + +/*! \brief Derivative table over the BF16 magnitudes [kLo, kHi]. + * + * Positive inputs at entries [0, n), negative ones at [kSignStride, kSignStride + n). + * A power-of-two stride turns the BF16 sign bit into a byte offset with one shift. + */ +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 = kStride; + static_assert(kEntriesPerSign <= kSignStride, "The two signs' entries must not overlap."); + //! 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; + static_assert(kStride == 2048 || kStride == 4096, "The packed lookups assume these strides."); +}; + +/*! \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<__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<__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<__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); } +}; + +//! [2^-9, 89]. dsilu has saturated, to 1 and to +0, by |x| = 89, up to inf. +template <> +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); } +}; + +//! dSiLU over [2^-25, 89]: below 2^-25 the linear term already equals util/math.h. Selected +//! with NVTE_ACTIVATION_TABLE_EXACT_DSILU=1. +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); } +}; + +//! BF16 table size in bytes, in the default layout. +constexpr size_t table_bytes(const Kind kind) { + switch (kind) { + case kGelu: + return Table<__nv_bfloat16, kGelu>::kBytes; + case kSilu: + return Table<__nv_bfloat16, kSilu>::kBytes; + case kDgelu: + return Table<__nv_bfloat16, kDgelu>::kBytes; + case kDsilu: + return Table<__nv_bfloat16, kDsilu>::kBytes; + default: + return 0; + } +} + +//! 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 device_table builds for kind. +inline size_t built_table_bytes(const Kind kind) { + if (kind == kDsilu && use_exact_dsilu_table()) { + return ExactDsiluTable<__nv_bfloat16>::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. + //! 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. +struct TableRef { + const void *ptr = nullptr; + TableUse use = kNoTable; +}; + +/*! \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 + * 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. + */ +TableRef device_table(Kind kind); + +//! 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 +inline constexpr bool kHasTable = false; +template +inline constexpr Kind kKindOf = kGelu; + +template <> +inline constexpr bool kHasTable> = true; +template <> +inline constexpr bool kHasTable> = true; +template <> +inline constexpr bool kHasTable> = true; +template <> +inline constexpr bool kHasTable> = true; +template <> +inline constexpr Kind kKindOf> = kSilu; +template <> +inline constexpr Kind kKindOf> = kDgelu; +template <> +inline constexpr Kind kKindOf> = kDsilu; + +/*! \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) { + 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) { + 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 given by kUse, or from util/math.h for kNoTable. + * + * x must be a BF16 value widened to FP32 when reading a table. + */ +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<__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 + // 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; + if constexpr (kUse == kExactDsiluTable) { + static_assert(kKind == kDsilu, "Only dSiLU has an exact 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)); + } + } +} + +} // namespace activation_table +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_UTIL_ACTIVATION_TABLE_CUH_