diff --git a/docs/envvars.rst b/docs/envvars.rst index 46b70bbe46a..23a94175f74 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -150,7 +150,7 @@ General :Type: ``int`` (0 or 1) :Default: ``0`` :Description: Use CuTeDSL kernel implementations when available. This requires that Transformer Engine was built with ``NVTE_WITH_CUTEDSL=1``. - It also requires ``apache-tvm-ffi``, ``nvidia-cutlass-dsl`` to be installed and you are using CPython 3.10 or later. + It also requires ``apache-tvm-ffi`` and ``nvidia-cutlass-dsl>=4.8.0`` to be installed and you are using CPython 3.10 or later. If these requirements are not met, Transformer Engine will fall back to using CUDA C++ kernels if possible. .. envvar:: NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN @@ -342,6 +342,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 c5acd043363..901a5ff51f9 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.cu b/tests/cpp/operator/test_cast_mxfp8.cu index 9368d4030b5..bc62d333b5f 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -699,6 +699,9 @@ std::vector> matrix_sizes = { {1, 16}, {16, 48}, {128, 128}, + // This case will cover cast_rowwise's last row being partial case, in which + // the kernel will reshape the input tensor with last dim being 1024 so the last row + // is not a full row due to 993*512 % 1024 != 0 {993, 512}, {1024}, {8, 32, 1024}, 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..d29ebd98d0d --- /dev/null +++ b/tests/cpp/operator/test_cast_mxfp8_act_exact.cu @@ -0,0 +1,506 @@ +/************************************************************************* + * 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 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)) { + const bool rowwise_reduction = std::is_same_v && !colwise; + const std::vector expected = reference_dbias(z, rows, cols, rowwise_reduction); + dbias.to_cpu(); + const IType *got = dbias.rowwise_cpu_dptr(); + size_t mismatches = 0; + for (size_t c = 0; c < cols; ++c) { + const IType e = static_cast(expected[c]); + mismatches += std::memcmp(&e, &got[c], sizeof(IType)) != 0 && + !(std::isnan(static_cast(e)) && std::isnan(static_cast(got[c]))); + } + EXPECT_EQ(mismatches, 0u) << "dbias"; + } +} + +using Params = + std::tuple, bool /*use_table*/, 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 1b5b95d7314..32fd6bb7f0d 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -7,6 +7,7 @@ #pragma once #include +#include #include #include #include @@ -449,6 +450,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 1dfd2fed4e4..c5d30942aa9 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 4671ea06869..c44b93b5614 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -23,6 +23,7 @@ from transformer_engine.common import _get_shared_object_file from transformer_engine.common.CuTeDSL.utils import device_compute_capability from transformer_engine.jax import cpp_extensions as tex +from transformer_engine.jax.cpp_extensions.activation import ActivationParams from transformer_engine.jax.quantize import ( QuantizerFactory, QuantizeLayout, @@ -68,8 +69,8 @@ # Every activation the CuTeDSL backend implements, as (JAX activation_type, name it goes by # in the config key). The two differ for qgelu/srelu: JAX spells them quick_gelu/squared_relu # (ActivationEnum in jax/cpp_extensions/activation.py), the backend uses the C++ spelling -# (Activation in common/util/cutedsl_utils.h). Gated variants are absent because they go to a -# separate TE/common kernel that the backend does not cover. +# (Activation in common/util/cutedsl_utils.h). Gated variants go to a separate kernel and are +# covered by GATED_ACTIVATIONS below. ACTIVATIONS = [ (("gelu",), "gelu"), (("relu",), "relu"), @@ -197,12 +198,15 @@ def extract_quantized_output(out, dbias): ScaledTensor1x carries scale_inv already trimmed to the unpadded shape (see its __post_init__), so there is no uninitialized scale padding to exclude. + FP8 -0.0 (0x80) is folded into +0.0: both backends round to zero there and may differ only + in its sign, which training does not depend on. """ tensors = [out] if isinstance(out, ScaledTensor1x) else [out.rowwise_tensor, out.colwise_tensor] parts = {} for t in tensors: name = "colwise" if t.is_colwise else "rowwise" - parts[f"{name} data"] = np.asarray(t.data.view(jnp.uint8)) + data = np.asarray(t.data.view(jnp.uint8)) + parts[f"{name} data"] = np.where(data == 0x80, 0, data).astype(np.uint8) parts[f"{name} scales"] = np.asarray(t.scale_inv.view(jnp.uint8)) return parts, None if dbias is None else np.asarray(dbias) @@ -283,3 +287,149 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): in_dtype, fp8_dtype, ) + + +# Cover the specialized kernels' shapes +# Note that JAX only emits linear scales and passes no noop tensor +SPECIALIZED_CASES = { + # G2RRowwise2Lane: its four launch tiers (by output size). + "g2r_rowwise2lane_tier0": ((1024, 4096), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier1": ((4064, 4480), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier2": ((8192, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), + "g2r_rowwise2lane_tier3": ((16384, 8192), QuantizeLayout.ROWWISE, jnp.bfloat16), + # G2RRowwise1Lane + "g2r_rowwise1lane_fp16": ((1024, 4096), QuantizeLayout.ROWWISE, jnp.float16), + # RegisterBidimensional: static N == 2048 / 4096, the wide tile with and without the + # deep cluster, and the narrow tile (N % 512 != 0). + "g2r_bidim_n2048": ((256, 2048), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_n4096": ((256, 4096), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_wide": ((256, 6144), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_wide_cluster": ((4096, 16384), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + "g2r_bidim_narrow": ((160, 2304), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), + # SpecializedBidimensional (g2s): fp16, or bf16 with N % 256 != 0. + "g2s_bidim_fp16": ((256, 1024), QuantizeLayout.ROWWISE_COLWISE, jnp.float16), + "g2s_bidim_bf16_n384": ((256, 384), QuantizeLayout.ROWWISE_COLWISE, jnp.bfloat16), +} + + +@pytest.mark.parametrize( + "shape,q_layout,in_dtype", SPECIALIZED_CASES.values(), ids=SPECIALIZED_CASES.keys() +) +def test_specialized_kernels(shape, q_layout, in_dtype): + run_test_case("CAST_ONLY", *NO_ACT, shape, q_layout, in_dtype, jnp.float8_e4m3fn) + + +# Gated activations (gated_mxfp8.cuh), as (JAX activation_type, name in the config key +# (GatedActivation in common/util/cutedsl_utils.h), activation parameters). The clamped SwiGLU +# parameters are chosen so that its clamps are active on the [-2, 2] inputs generate_inputs draws. +GATED_ACTIVATIONS = [ + (("gelu", "linear"), "geglu", {}), + (("quick_gelu", "linear"), "qgeglu", {}), + (("relu", "linear"), "reglu", {}), + (("squared_relu", "linear"), "sreglu", {}), + (("silu", "linear"), "swiglu", {}), + (("sigmoid", "linear"), "glu", {}), + ( + ("clamped_silu", "clamped_linear"), + "clamped_swiglu", + {"limit": 1.0, "alpha": 1.702, "glu_linear_offset": 0.5}, + ), +] +GATED_IDS = [desc for _, desc, _ in GATED_ACTIVATIONS] +get_direction_id = lambda is_bwd: "bwd" if is_bwd else "fwd" + + +def get_gated_cfg_key(act_desc, in_dtype, fp8_dtype, q_layout, is_bwd): + """Mirror of MXFP8GatedQuantConfig::to_key (gated_mxfp8_cutedsl.cuh).""" + # JAX quantize emits scales in the linear (non-swizzled) layout + flags = (True, q_layout.has_colwise, False, is_bwd) + major, minor = device_compute_capability() + return ( + f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_" + + "_".join("1" if f else "0" for f in flags) + + f"_{act_desc}" + ) + + +def run_gated_test_case( + act_type, act_desc, act_kwargs, is_bwd, shape, q_layout, in_dtype, fp8_dtype +): + """Assert the CuTeDSL and CUDA backends produce bit-identical gated activation outputs. + shape is (M, N): the gated input is (M, 2, N), [act | gate] along axis -2. + """ + M, N = shape + # generate_inputs returns (M, N) and (M, 1, N) draws; stack them into the (M, 2, N) input + act, gate = generate_inputs(M, N, in_dtype) + x = jnp.concatenate([jnp.expand_dims(act, axis=-2), gate], axis=-2) + grad, _ = generate_inputs(M, N, in_dtype, seed=1) + act_params = ActivationParams.create(act_type, **act_kwargs) + + def run(): + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, q_dtype=fp8_dtype, q_layout=q_layout + ) + if is_bwd: + out, _ = tex.quantize_dact_dbias( + grad, x, act_type, is_dbias=False, quantizer=quantizer, act_params=act_params + ) + else: + out = tex.act_lu(x, act_type, quantizer=quantizer, act_params=act_params) + return extract_quantized_output(out, None)[0] + + set_cutedsl_backend(False) + cuda_output = run() + set_cutedsl_backend(True) + try: + cutedsl_output = run() + finally: + set_cutedsl_backend(False) + + key = get_gated_cfg_key(act_desc, in_dtype, fp8_dtype, q_layout, is_bwd) + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + tag = ( + f"{act_desc}/{get_direction_id(is_bwd)}/{get_layout_id(q_layout)}/{M}x{N}/" + f"{DTYPE_TO_STR[in_dtype]}/{FP8_TO_STR[fp8_dtype]}" + ) + for name, cuda_bytes in cuda_output.items(): + assert np.array_equal( + cutedsl_output[name], cuda_bytes + ), f"{tag}: {name} differ between backends" + + +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("q_layout", Q_LAYOUTS, ids=get_layout_id) +@pytest.mark.parametrize("act_type,act_desc,act_kwargs", GATED_ACTIVATIONS, ids=GATED_IDS) +@pytest.mark.parametrize("is_bwd", [False, True], ids=get_direction_id) +def test_gated_sizes(is_bwd, act_type, act_desc, act_kwargs, q_layout, shape): + run_gated_test_case( + act_type, + act_desc, + act_kwargs, + is_bwd, + shape, + q_layout, + jnp.bfloat16, + jnp.float8_e4m3fn, + ) + + +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("act_type,act_desc,act_kwargs", GATED_ACTIVATIONS, ids=GATED_IDS) +@pytest.mark.parametrize("is_bwd", [False, True], ids=get_direction_id) +def test_gated_dtypes(is_bwd, act_type, act_desc, act_kwargs, fp8_dtype, in_dtype): + run_gated_test_case( + act_type, + act_desc, + act_kwargs, + is_bwd, + (256, 384), + QuantizeLayout.ROWWISE_COLWISE, + in_dtype, + fp8_dtype, + ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index f409f76993d..fa51d292cb7 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -59,11 +59,12 @@ (96, 128), ] # (block_rows, block_cols): (1,32)=rowwise, (32,1)=colwise, (32,32)=both. -BLOCK_SIZES = [(1, 32), (32, 1), (32, 32)] +ROW, COL, BOTH = (1, 32), (32, 1), (32, 32) +BLOCK_SIZES = [ROW, COL, BOTH] # Every activation the CuTeDSL backend implements; "desc" is the name it goes by in the -# config key (Activation in common/util/cutedsl_utils.h). Gated variants are absent because -# they go to a separate TE/common kernel that the backend does not cover. +# config key (Activation in common/util/cutedsl_utils.h). Gated variants go to a separate kernel +# and are covered by GATED_ACTIVATIONS below. IDENTITY = {"name": "Identity", "act": None, "dact": None, "dbias_dact": None, "desc": "none"} ACTIVATIONS = [ { @@ -250,31 +251,39 @@ def extract_quantized_output(out, rowwise, columnwise, swizzled): compared. Swizzled layout: the meaningful scales are scattered by the swizzle, so the top-left slice is meaningless; both backends zero the padding (see zero_scales_kernel), so compare the whole buffer instead -- which also covers the padding-zeroing itself. + FP8 -0.0 (0x80) is folded into +0.0: both backends round to zero there and may differ only + in its sign, which training does not depend on. """ parts = {} if rowwise: d = out._rowwise_data.view(torch.uint8) M, N = d.shape - parts["rowwise data"] = d.clone() + parts["rowwise data"] = d.masked_fill(d == 0x80, 0) s = out._rowwise_scale_inv parts["rowwise scales"] = (s if swizzled else s[:M, : (N + 31) // 32]).clone() if columnwise: d = out._columnwise_data.view(torch.uint8) M, N = d.shape - parts["colwise data"] = d.clone() + parts["colwise data"] = d.masked_fill(d == 0x80, 0) s = out._columnwise_scale_inv parts["colwise scales"] = (s if swizzled else s[: (M + 31) // 32, :N]).clone() return parts -def run_test_case(method, act, shape, block_size, in_dtype, fp8_dtype, swizzled=False): +def run_test_case( + method, act, shape, block_size, in_dtype, fp8_dtype, swizzled=False, input_offset=0 +): """Assert the CuTeDSL and CUDA backends produce bit-identical outputs for the - same input and config. + same input and config. input_offset (in elements) shifts x's data pointer past the + start of its allocation, to test misaligned inputs. """ M, N = shape rowwise = block_size[1] != 1 columnwise = block_size[0] != 1 x, act_input = generate_inputs(M, N, in_dtype) + if input_offset: + buf = torch.empty(input_offset + M * N, dtype=in_dtype, device="cuda") + x = buf[input_offset:].view(M, N).copy_(x) set_cutedsl_backend(False) out_cuda, dbias_cuda = run_quantize( @@ -370,3 +379,178 @@ def test_sizes(swizzled, method, act, block_size, shape): @pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) def test_dtypes(swizzled, method, act, fp8_dtype, in_dtype): run_test_case(method, act, (256, 384), (32, 32), in_dtype, fp8_dtype, swizzled) + + +# Cover the specialized kernels' shapes + +SPECIALIZED_PATH_CASES = [ + # G2RRowwise2Lane: the four launch tiers (by output size) on the flat packed-scale view. + ("g2r_rowwise2lane_tier0", (1024, 4096), ROW, torch.bfloat16, False, 0), + ("g2r_rowwise2lane_tier1", (4064, 4480), ROW, torch.bfloat16, False, 0), + ("g2r_rowwise2lane_tier2", (8192, 8192), ROW, torch.bfloat16, False, 0), + ("g2r_rowwise2lane_tier3", (16384, 8192), ROW, torch.bfloat16, False, 0), + # Swizzled scales use the 2-D grid instead of the flat view. + ("g2r_rowwise2lane_swizzled", (4096, 4096), ROW, torch.bfloat16, True, 0), + # Partial row and column tiles, and scale padding. + ("g2r_rowwise2lane_swizzled_partial", (160, 384), ROW, torch.bfloat16, True, 0), + # Row tiles past grid.y's limit spill into grid.z. + ("g2r_rowwise2lane_swizzled_grid_z", (524352, 128), ROW, torch.bfloat16, True, 0), + # G2RRowwise1Lane: input only 16-byte aligned, or fp16. + ("g2r_rowwise1lane_misaligned_bf16", (1024, 4096), ROW, torch.bfloat16, False, 8), + ("g2r_rowwise1lane_fp16", (1024, 4096), ROW, torch.float16, False, 0), + # G2RBidim: static N == 2048 / 4096, the wide tile with and without the + # deep cluster, and the narrow tile (N % 512 != 0). + ("g2r_bidim_n2048", (256, 2048), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_n2048_swizzled", (256, 2048), BOTH, torch.bfloat16, True, 0), + ("g2r_bidim_n4096_swizzled", (256, 4096), BOTH, torch.bfloat16, True, 0), + ("g2r_bidim_wide", (256, 6144), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_wide_cluster", (4096, 16384), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_narrow", (160, 2304), BOTH, torch.bfloat16, False, 0), + ("g2r_bidim_narrow_swizzled", (160, 2304), BOTH, torch.bfloat16, True, 0), + # M % 128 != 0: swizzled scale padding. + ("g2r_bidim_wide_swizzled_pad", (4128, 6144), BOTH, torch.bfloat16, True, 0), + # G2SBidim: fp16, or bf16 with N % 256 != 0. + ("g2s_bidim_fp16", (256, 1024), BOTH, torch.float16, False, 0), + ("g2s_bidim_bf16_n384", (256, 384), BOTH, torch.bfloat16, False, 0), +] + + +@pytest.mark.parametrize( + "shape,block_size,in_dtype,swizzled,input_offset", + [case[1:] for case in SPECIALIZED_PATH_CASES], + ids=[case[0] for case in SPECIALIZED_PATH_CASES], +) +def test_specialized_kernels(shape, block_size, in_dtype, swizzled, input_offset): + run_test_case( + "CAST_ONLY", + IDENTITY, + shape, + block_size, + in_dtype, + tex.DType.kFloat8E4M3, + swizzled, + input_offset, + ) + + +# Gated activations (gated_mxfp8.cuh). Input is [M, 2N] = [act | gate]; the forward writes +# [M, N] and the backward, given grad [M, N], writes [M, 2N]. "desc" is the name the activation +# goes by in the config key (GatedActivation in common/util/cutedsl_utils.h). The parameters of +# the parameterized ones are chosen so that their clamps and tanh terms are active on the +# [-2, 2] inputs generate_inputs draws. +GATED_ACTIVATIONS = [ + {"name": "GeGLU", "act": tex.geglu, "dact": tex.dgeglu, "desc": "geglu", "kwargs": {}}, + {"name": "QGeGLU", "act": tex.qgeglu, "dact": tex.dqgeglu, "desc": "qgeglu", "kwargs": {}}, + {"name": "ReGLU", "act": tex.reglu, "dact": tex.dreglu, "desc": "reglu", "kwargs": {}}, + {"name": "SReGLU", "act": tex.sreglu, "dact": tex.dsreglu, "desc": "sreglu", "kwargs": {}}, + {"name": "SwiGLU", "act": tex.swiglu, "dact": tex.dswiglu, "desc": "swiglu", "kwargs": {}}, + {"name": "GLU", "act": tex.glu, "dact": tex.dglu, "desc": "glu", "kwargs": {}}, + { + "name": "ClampedSwiGLU", + "act": tex.clamped_swiglu, + "dact": tex.clamped_dswiglu, + "desc": "clamped_swiglu", + "kwargs": {"limit": 1.0, "alpha": 1.702, "glu_linear_offset": 0.5}, + }, + { + "name": "SiTUGLU", + "act": tex.situglu, + "dact": tex.dsituglu, + "desc": "situglu", + "kwargs": {"beta1": 1.0, "beta2": 2.0}, + }, +] +get_gated_id = lambda a: a["name"] +GATED_DIRECTIONS = [False, True] +get_direction_id = lambda is_bwd: "bwd" if is_bwd else "fwd" + + +def run_on_both_backends(run): + """Return (CUDA result, CuTeDSL result) of run(), with run() responsible for materializing + whatever it returns before the backend is switched.""" + set_cutedsl_backend(False) + cuda_result = run() + set_cutedsl_backend(True) + try: + cutedsl_result = run() + finally: + set_cutedsl_backend(False) + return cuda_result, cutedsl_result + + +def assert_cutedsl_kernel_registered(key): + """Guard against a silent CUDA fallback, which would make a cross-backend comparison compare + CUDA against itself.""" + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + +def get_gated_cfg_key(act, in_dtype, fp8_dtype, rowwise, colwise, swizzled, is_bwd): + """Mirror of MXFP8GatedQuantConfig::to_key (gated_mxfp8_cutedsl.cuh).""" + major, minor = device_compute_capability() + flags = (rowwise, colwise, swizzled, is_bwd) + return ( + f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_" + + "_".join("1" if f else "0" for f in flags) + + f"_{act['desc']}" + ) + + +def run_gated_test_case(act, is_bwd, shape, block_size, in_dtype, fp8_dtype, swizzled=False): + """Assert the CuTeDSL and CUDA backends produce bit-identical gated activation outputs. + shape is (M, N): the gated input is [M, 2N]. + """ + M, N = shape + if is_bwd and swizzled: + pytest.skip("CuTeDSL hands backward gated activations with GEMM-swizzled scales to CUDA") + rowwise = block_size[1] != 1 + columnwise = block_size[0] != 1 + x, _ = generate_inputs(M, 2 * N, in_dtype) + grad, _ = generate_inputs(M, N, in_dtype, seed=1) + + def run(): + q = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=rowwise, columnwise=columnwise) + q.optimize_for_gemm = swizzled + if is_bwd: + out = act["dact"](grad, x, q, **act["kwargs"]) + else: + out = act["act"](x, q, **act["kwargs"]) + return extract_quantized_output(out, rowwise, columnwise, swizzled) + + cuda_output, cutedsl_output = run_on_both_backends(run) + assert_cutedsl_kernel_registered( + get_gated_cfg_key(act, in_dtype, fp8_dtype, rowwise, columnwise, swizzled, is_bwd) + ) + + layout = "swizzled" if swizzled else "linear" + tag = ( + f"{act['name']}/{get_direction_id(is_bwd)}/{M}x{N}/{DTYPE_TO_STR[in_dtype]}/" + f"{FP8_TO_STR[fp8_dtype]}/{layout}" + ) + for name, cuda_bytes in cuda_output.items(): + assert torch.equal( + cutedsl_output[name], cuda_bytes + ), f"{tag}: {name} differ between backends" + + +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +@pytest.mark.parametrize("act", GATED_ACTIVATIONS, ids=get_gated_id) +@pytest.mark.parametrize("is_bwd", GATED_DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_gated_sizes(swizzled, is_bwd, act, block_size, shape): + run_gated_test_case( + act, is_bwd, shape, block_size, torch.bfloat16, tex.DType.kFloat8E4M3, swizzled + ) + + +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("act", GATED_ACTIVATIONS, ids=get_gated_id) +@pytest.mark.parametrize("is_bwd", GATED_DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_gated_dtypes(swizzled, is_bwd, act, fp8_dtype, in_dtype): + run_gated_test_case(act, is_bwd, (256, 384), BOTH, in_dtype, fp8_dtype, swizzled) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 843ec412c3d..4af4369d50f 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -331,6 +331,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 @@ -769,7 +770,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/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 50d536852c2..f9fabe72e5b 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -10,8 +10,9 @@ import tvm_ffi -from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( +from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8_common import ( get_mxfp8_quantization_function, + get_mxfp8_gated_quantization_function, ) @@ -19,6 +20,8 @@ def register_cutedsl_backends(): """Register all available CuTeDSL backends for on-demand compilation via TVM-FFI. The C++ dispatcher retrieves them by the names defined here. """ - tvm_ffi.register_global_func( - "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True - ) + for entrypoint in ( + get_mxfp8_quantization_function, + get_mxfp8_gated_quantization_function, + ): + tvm_ffi.register_global_func(entrypoint.__name__, entrypoint, override=True) diff --git a/transformer_engine/common/CuTeDSL/activations.py b/transformer_engine/common/CuTeDSL/activations.py index 34772eb7257..cf3a42b49b7 100644 --- a/transformer_engine/common/CuTeDSL/activations.py +++ b/transformer_engine/common/CuTeDSL/activations.py @@ -8,9 +8,11 @@ # Each function is a 1:1 port of its math.h counterpart (see module docstring). # pylint: disable=missing-function-docstring import os +import struct +import cutlass from cutlass import cute -from cutlass import Float32 +from cutlass import Float32, Int32, Uint32 from transformer_engine.common.CuTeDSL.utils import fma_f32 @@ -36,7 +38,8 @@ def act_gelu(x: Float32) -> Float32: def act_silu(x: Float32) -> Float32: - return x / (Float32(1.0) + exp(-x)) + # x * sigmoid(x), as math.h computes it: x / (1 + exp(-x)) rounds differently in FP32. + return x * sigmoid(x) def act_qgelu(x: Float32) -> Float32: @@ -80,3 +83,221 @@ 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) + + +# ---- Gated activations (gated_mxfp8.cuh) ---- +# The parameterized ones take their ClampedSwiGLUParam / SiTUGLUParam fields as +# trailing runtime arguments. + +LOG2E = 1.4426950408889634 # log2(e), what __expf scales its argument by + + +def act_sigmoid(x: Float32) -> Float32: + return sigmoid(x) + + +def dact_dsigmoid(x: Float32) -> Float32: + return dsigmoid(x) + + +def act_clamped_silu(x: Float32, limit: Float32, alpha: Float32) -> Float32: + c = cute.arch.fmin(limit, x) + return c * sigmoid(alpha * c) + + +def act_situ_gate(x: Float32, beta1: Float32) -> Float32: + return beta1 * tanh(x / beta1) * sigmoid(x) + + +def dact_dsitu_gate(x: Float32, beta1: Float32) -> Float32: + t = tanh(x / beta1) + s = sigmoid(x) + return (Float32(1.0) - t * t) * s + beta1 * t * s * (Float32(1.0) - s) + + +def situ_up(x: Float32, beta2: Float32) -> Float32: + return beta2 * tanh(x / beta2) + + +def dsitu_up(x: Float32, beta2: Float32) -> Float32: + t = tanh(x / beta2) + return Float32(1.0) - t * t + + +@cute.jit +def sigmoidf(x: Float32) -> Float32: + """math.h's sigmoidf: __frcp_rn(1.0f + __expf(-x)).""" + e = cute.arch.inline_ptx( + "ex2.approx.f32 {$w0}, {$r0};", + write_only_types=[Float32], + read_only_args=[-x * Float32(LOG2E)], + ) + return cute.arch.inline_ptx( + "rcp.rn.f32 {$w0}, {$r0};", + write_only_types=[Float32], + read_only_args=[Float32(1.0) + e], + ) + + +# BF16 activation tables (see common/util/activation_table.cuh, which builds them; the layouts +# below must match it). With a table, BF16 activations follow the table specification: GeLU and +# SiLU of -inf give -0, and the derivatives below the table window are the linear term +# 0.5 + slope*x. Kernels either read the table ("table_*_2x") or compute what it holds +# ("activation_op"/"dactivation_op"). + +# Forward tables: rows of (op(+m), op(-m)) BF16 entries for magnitudes m from LOW_SENTINEL to +# HIGH_SENTINEL; the LOW_SENTINEL row holds -inf, the HIGH_SENTINEL row +inf and -0. +FORWARD_TABLE_WINDOWS = {"gelu": (0x3B4B, 0x40A5), "silu": (0x3BB6, 0x42B1)} +# Derivative tables: FP32 entries for magnitudes [LO, HI], positive inputs from entry 0 and +# negative ones from entry STRIDE; from SATURATED_END on (and for NaN) the result is NaN. +# Keyed by (activation, layout); layout "exact" is the exact dSiLU table. +DERIVATIVE_TABLE_WINDOWS = { + ("dgelu", "default"): (0x3980, 0x4100, 0x6044, 2048), + ("dsilu", "default"): (0x3B00, 0x42B2, 0x7F80, 2048), + ("dsilu", "exact"): (0x3300, 0x42B2, 0x7F80, 4096), +} +DERIVATIVE_SLOPES = {"dgelu": 0.79788456, "dsilu": 0.5} + + +def activation_table_bytes(activation: str, layout: str) -> int: + """Size of the table for activation in layout, rounded up to 16 bytes as TE allocates it.""" + if activation in FORWARD_TABLE_WINDOWS: + lo, hi = FORWARD_TABLE_WINDOWS[activation] + rows = (hi + 1) - (lo - 1) + 1 + return (rows * 2 * 2 + 15) // 16 * 16 + lo, hi, _, stride = DERIVATIVE_TABLE_WINDOWS[(activation, layout)] + return ((stride + hi - lo + 1) * 4 + 15) // 16 * 16 + + +@cute.jit +def _select_below(x: Float32, bound: Float32, below: Float32, other: Float32) -> Float32: + """below if |x| < bound, else other. Branch-free; NaN compares false.""" + return cute.arch.inline_ptx( + "{ .reg .pred p; .reg .f32 m; abs.f32 m, {$r0}; setp.lt.f32 p, m, {$r1}; " + "selp.f32 {$w0}, {$r2}, {$r3}, p; }", + write_only_types=[Float32], + read_only_args=[x, bound, below, other], + ) + + +def forward_table_spec(op): + """op as the forward tables hold it (see the comment above), computed.""" + + @cute.jit + def table_spec(x: Float32) -> Float32: + # Outside the window the table's closed forms equal util/math.h rounded to BF16, except + # at -inf, where the formula gives NaN and the table -0. util/math.h gives -0 for every + # finite x below -2^100, so clamp the input there; NaN fails the comparison. + clamped = cute.arch.inline_ptx( + "{ .reg .pred p; setp.lt.f32 p, {$r0}, {$r1}; selp.f32 {$w0}, {$r1}, {$r0}, p; }", + write_only_types=[Float32], + read_only_args=[x, Float32(-(2.0**100))], + ) + return op(clamped) + + return table_spec + + +def derivative_table_spec(dop, activation: str, layout: str): + """dop as the derivative table of activation in layout holds it, computed.""" + # Below the exact dSiLU table's window the linear term equals util/math.h. + if layout == "exact": + return dop + lo = DERIVATIVE_TABLE_WINDOWS[(activation, layout)][0] + window_low = struct.unpack(" Float32: + # Above the window util/math.h already gives the saturated value or NaN. + linear = fma_f32(Float32(slope), x, Float32(0.5)) + return _select_below(x, Float32(window_low), linear, dop(x)) + + return table_spec + + +@cute.jit +def table_forward_2x( + x: Int32, + sTable_u16: cute.Tensor, + LOW_SENTINEL: cutlass.Constexpr[int], + HIGH_SENTINEL: cutlass.Constexpr[int], +) -> Int32: + """Forward activation of a BF16 pair (low lane in the low half) from a forward table in + shared memory, as BF16 bits. Same instructions as forward_2x in util/activation_table.cuh.""" + # Clamp both magnitudes to the table rows. bf16x2 max/min ignore NaN, which lands on the + # -inf row; min.NaN with the upper bound below then makes it NaN again. + clamped = Uint32( + cute.arch.inline_ptx( + "{ .reg .b32 m, lo, hi; " + f"mov.b32 lo, {LOW_SENTINEL * 0x10001}; mov.b32 hi, {HIGH_SENTINEL * 0x10001}; " + "and.b32 m, {$r0}, 0x7fff7fff; max.bf16x2 m, m, lo; min.bf16x2 {$w0}, m, hi; }", + write_only_types=[Int32], + read_only_args=[x], + ) + ) + # Both lanes' entry indices in one word: 2 * row + sign per half. Each stays below 0x10000, + # so the halves never carry into each other. + indices = ((clamped - Uint32(LOW_SENTINEL * 0x10001)) << Uint32(1)) + ( + (Uint32(x) >> Uint32(15)) & Uint32(0x10001) + ) + entry_lo = Uint32(sTable_u16[Int32(indices & Uint32(0xFFFF))]) + entry_hi = Uint32(sTable_u16[Int32(indices >> Uint32(16))]) + entries = Int32(entry_lo | (entry_hi << Uint32(16))) + # max(entry, RN(x/2)) picks RN(x/2) below the window; min.NaN(., max.NaN(x, -0)) picks x + # above it for x > 0, -0 for x < 0, and NaN for NaN. + return cute.arch.inline_ptx( + "{ .reg .b32 half, upper, neg_zero, h; " + "mov.b32 half, 0x3f003f00; mov.b32 neg_zero, 0x80008000; " + "mul.rn.bf16x2 h, {$r0}, half; max.NaN.bf16x2 upper, {$r0}, neg_zero; " + "max.bf16x2 h, {$r1}, h; min.NaN.bf16x2 {$w0}, h, upper; }", + write_only_types=[Int32], + read_only_args=[x, entries], + ) + + +@cute.jit +def table_derivative_2x( + x: Int32, + sTable_f32: cute.Tensor, + LO: cutlass.Constexpr[int], + HI: cutlass.Constexpr[int], + SATURATED_END: cutlass.Constexpr[int], + STRIDE: cutlass.Constexpr[int], + SLOPE: cutlass.Constexpr[float], +): + """Derivative of a BF16 pair (low lane in the low half) from a derivative table in shared + memory, as (lo, hi) FP32. Same results as derivative_2x in util/activation_table.cuh.""" + magnitude = Uint32(x) & Uint32(0x7FFF7FFF) + clamped = Uint32( + cute.arch.inline_ptx( + "{ .reg .b32 lo, hi; " + f"mov.b32 lo, {LO * 0x10001}; mov.b32 hi, {HI * 0x10001}; " + "max.bf16x2 {$w0}, {$r0}, lo; min.bf16x2 {$w0}, {$w0}, hi; }", + write_only_types=[Int32], + read_only_args=[Int32(magnitude)], + ) + ) + bits = Uint32(x) + # Both lanes' entry indices in one word: the sign bits (15 and 31), shifted down to STRIDE + # (a power of two), select the negative half. Each index stays below 0x10000, so the halves + # never carry into each other. + SIGN_SHIFT = 15 - (STRIDE.bit_length() - 1) + indices = (clamped - Uint32(LO * 0x10001)) + ((bits ^ magnitude) >> Uint32(SIGN_SHIFT)) + entry_lo = sTable_f32[Int32(indices & Uint32(0xFFFF))] + entry_hi = sTable_f32[Int32(indices >> Uint32(16))] + # Below the window: the linear term 0.5 + SLOPE*x, rounded once. + linear_lo = fma_f32(Float32(SLOPE), (bits << Uint32(16)).bitcast(Float32), Float32(0.5)) + linear_hi = fma_f32(Float32(SLOPE), (bits & Uint32(0xFFFF0000)).bitcast(Float32), Float32(0.5)) + # One BF16x2 compare per condition sets both lanes' predicates. The unordered >= also holds + # for NaN, which gives NaN like the saturated end. + return cute.arch.inline_ptx( + "{ .reg .pred below_lo, below_hi, nan_lo, nan_hi; .reg .b32 lo, end; mov.b32 lo," + f" {LO * 0x10001}; mov.b32 end, {SATURATED_END * 0x10001}; setp.lt.bf16x2" + " below_lo|below_hi, {$r0}, lo; setp.geu.bf16x2 nan_lo|nan_hi, {$r0}, end; selp.f32" + " {$w0}, {$r1}, {$r2}, below_lo; selp.f32 {$w1}, {$r3}, {$r4}, below_hi;" + " selp.b32 {$w0}, 0x7fffffff, {$w0}, nan_lo; selp.b32 {$w1}, 0x7fffffff, {$w1}," + " nan_hi; }", + write_only_types=[Float32, Float32], + read_only_args=[Int32(magnitude), linear_lo, entry_lo, linear_hi, entry_hi], + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py index 6491573690f..444d85234ec 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py @@ -4,4 +4,4 @@ """CuTeDSL MXFP8 quantization kernels.""" -from . import quantize_mxfp8 +from . import quantize_mxfp8_common diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py deleted file mode 100644 index 824ed18aa59..00000000000 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py +++ /dev/null @@ -1,2846 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""MXFP8 quantization kernel implemented in CuTeDSL. - -Replicates the core logic of quantize_mxfp8.cuh: given a 2D tensor of BF16/FP16 -values, quantize to MXFP8 format (FP8E4M3 data + E8M0 per-block scales). - -Matches the C++ kernel's tile dimensions and thread layout: - CHUNK_DIM_Y = 64, CHUNK_DIM_X = 64, THREADS_PER_CTA = 64 - BUFF_DIM_Y = 32, BUFF_DIM_X = 64, STAGES = 2 - MXFP8_BLOCK_SCALING_SIZE = 32 (elements per MXFP8 scaling block) - -Grid: (ceil(N / 64), ceil(M / 64)) -Each block processes a 64x64 chunk in 2 stages of 32x64 tiles loaded into -shared memory. -""" - -# Local @cute.struct classes are SMEM-layout descriptors that need no docstrings. -# pylint: disable=missing-class-docstring - -import abc -import logging -import os -from typing import Optional, Type - -import cutlass -from cutlass import cute -from cutlass import pipeline -from cutlass import Boolean, Float32, Float8E8M0FNU, Int32, Int64, Uint16, Uint32, Uint8 -from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module -import tvm_ffi - -from transformer_engine.common.CuTeDSL.utils import ( - device_compute_capability, - str_to_cutlass_dtype, - is_packed16, - abs_max_x2_bf16, - abs_max_x2_f16, - max_x2_bf16, - max_x2_f16, - max_scalar_bf16, - max_scalar_f16, - abs_max_scalar_bf16, - abs_max_scalar_f16, - to_f32_bf16, - to_f32_f16, - x2_lo_to_f32_bf16, - x2_lo_to_f32_f16, - x2_hi_to_f32_bf16, - x2_hi_to_f32_f16, - truncate_f32_bf16, - truncate_f32_f16, - fabs_f32, - exp2f_rcp, - pack_f32x2, - unpack_i64_to_i32x2, -) -from transformer_engine.common.CuTeDSL.activations import ( - act_relu, - act_gelu, - act_silu, - act_qgelu, - act_srelu, - dact_drelu, - dact_dsrelu, - dact_dsilu, - dact_dqgelu, - dact_dgelu, -) -from transformer_engine.common.CuTeDSL.utils_fp8 import ( - get_cvt_f32x2_to_fp8x2_func, - cvt_f32_to_fp8e8m0fnu, - mul_f32x2_cvt_f32x4_to_fp8x4, - mul_f32x4_cvt_f32x4_to_fp8x4, - mul_f32x2_cvt_packed16x4_to_fp8x4, -) - -CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" - -logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") - -# Number of elements per MXFP8 scale block. They will share the same E8M0 scale factor -MXFP8_BLOCK_SCALING_SIZE = 32 -# How many threads are in one warp -THREADS_PER_WARP = 32 - -# FP8E4M3 max representable value -FP8E4M3_MAX_NORM = 448.0 -FP8E4M3_MAX_NORM_RCP = 1.0 / FP8E4M3_MAX_NORM -FP8E5M2_MAX_NORM = 57344.0 -FP8E5M2_MAX_NORM_RCP = 1.0 / FP8E5M2_MAX_NORM - -# If N is not divisible by 16, writing quantized FP8 output with N bytes per row will fail because -# TMA requires 16-byte alignment for each row -SYM_N_DIVISIBILITY = 16 - -SUPPORTED_ACTIVATIONS = { - "relu": act_relu, - "gelu": act_gelu, - "silu": act_silu, - "qgelu": act_qgelu, - "srelu": act_srelu, -} - -SUPPORTED_DACTIVATIONS = { - "drelu": dact_drelu, - "dgelu": dact_dgelu, - "dsilu": dact_dsilu, - "dqgelu": dact_dqgelu, - "dsrelu": dact_dsrelu, -} - - -@cute.jit -def derive_swizzled_scale_layout( - M: Int32, - N: Int32, - ROWWISE: cutlass.Constexpr[bool], - COLWISE: cutlass.Constexpr[bool], - mS_row: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], -): - """Derive the swizzled layout for the rowwise and colwise scale tensors.""" - num_scale_cols = cute.ceil_div(N, MXFP8_BLOCK_SCALING_SIZE) - num_scale_rows = cute.ceil_div(M, MXFP8_BLOCK_SCALING_SIZE) - - num_tiles_M = cute.ceil_div(M, 128) - num_tiles_SC = cute.ceil_div(num_scale_cols, 4) - num_tiles_SR = cute.ceil_div(num_scale_rows, 4) - num_tiles_N = cute.ceil_div(N, 128) - - if cutlass.const_expr(ROWWISE): - mS_row = cute.make_tensor( - mS_row.iterator, - cute.make_layout( - ((32, 4, num_tiles_M), (4, num_tiles_SC)), - stride=((16, 4, num_tiles_SC * 512), (1, 512)), - ), - ) - if cutlass.const_expr(COLWISE): - mS_col = cute.make_tensor( - mS_col.iterator, - cute.make_layout( - ((4, num_tiles_SR), (32, 4, num_tiles_N)), - stride=((1, 512), (16, 4, num_tiles_SR * 512)), - ), - ) - return mS_row, mS_col - - -@cute.jit -def quantize_rowwise_mxfp8( - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA - sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) - sO_row_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) - mS_row_stage: cute.Tensor, # rowwise scale tensor (1D swizzled, or 2D linear) - MAX_NORM_RCP: cutlass.Constexpr[float], - tile_row_start: Int32, # global row index of this stage's row 0 - tile_col_start: Int32, # global col index of this CTA's col 0 - M: Int32, - N: Int32, - ACTIVATION: cutlass.Constexpr[str | None], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - TILE_X: cutlass.Constexpr[int], - TILE_Y: cutlass.Constexpr[int], - WAVES: cutlass.Constexpr[int], - THREADS_PER_BANK: cutlass.Constexpr[int], - PACK_SIZE: cutlass.Constexpr[int], - SKIP_INPUT_MASKING: cutlass.Constexpr[bool] = False, - SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, - WITH_ACT: cutlass.Constexpr[bool] = False, - WITH_DACT: cutlass.Constexpr[bool] = False, - WITH_DBIAS: cutlass.Constexpr[bool] = False, - dbias_acc: Optional[cute.Tensor] = None, # only needed when WITH_DBIAS is True -): - """Quantize one SMEM tile rowwise to MXFP8 (per-row 32-elt block scales); returns the tile amax.""" - tidx, _, _ = cute.arch.thread_idx() - - CTA_THREADS_Y = TILE_Y # threads per column (rows per tile) - CTA_THREADS_X = TILE_X // MXFP8_BLOCK_SCALING_SIZE # threads per row (chunks per row) - - _, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), - val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), - ) - - sX_tv = cute.composition(sX_tile, tv_layout) - sO_tv = cute.composition(sO_row_tile, tv_layout) - - # I/O Elements that belong to this thread - sX_thread = sX_tv[tidx, None] # shape (32,) bf16 - sO_thread = sO_tv[tidx, None] # shape (32,) fp8 - - sO_thread_u32_ptr = cute.recast_ptr(sO_thread.iterator, dtype=Uint32) - sO_thread_u32 = cute.make_tensor( - sO_thread_u32_ptr, - cute.make_layout( - (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) - ), # 1 uint32 is 4 fp8 elements - ) - - # PTX allows to fuse relu activation in `cvt.rn.satfinite` unless we need to reduction for dbias - FUSE_RELU = cutlass.const_expr(ACTIVATION == "relu") and not WITH_DBIAS - # For this fast path we can read in pack of 2 instead of reading individual f16 / bf16 element, - # and keep the elements in half precision instead of upcasting every one of them to f32. - USE_HALF_PRECISION = is_packed16(DTYPE) and (ACTIVATION is None or FUSE_RELU) - - amax_r = Float32(0.0) - - # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks - # to avoid bank conflict. - bank_group = (tidx % THREADS_PER_WARP) // THREADS_PER_BANK - # The offset this thread should start reading from based on what's its first bank to access. - offset = bank_group * PACK_SIZE - if cutlass.const_expr(USE_HALF_PRECISION): - # If no activation, f16 / bf16 and rowwise quantization, we can read 2 f16 / bf16 at once in a pack - # and use max.xorsign.abs.f16x2 / max.xorsign.abs.bf16x2 to compute - max_x2 = max_x2_f16 if DTYPE is cutlass.Float16 else max_x2_bf16 - abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 - x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 - x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 - sX_thread_rw_i64 = cute.make_tensor( - cute.recast_ptr(sX_thread.iterator, dtype=Int64), - cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE // 4), stride=(0, 1) - ), # 1 int64 is 4 fp16/bf16 elements - ) - # Each wave reads its 4 elements (PACK_SIZE) as one 8-byte vectorized load - in_r = [[None, None] for _ in range(WAVES)] - for w in cutlass.range_constexpr(WAVES): - idx = (w + offset // 4) % (MXFP8_BLOCK_SCALING_SIZE // 4) - in_r[w][0], in_r[w][1] = unpack_i64_to_i32x2(sX_thread_rw_i64[0, idx]) - - if cutlass.const_expr(WITH_DBIAS): - for w in cutlass.range_constexpr(WAVES): - dbias_acc[w * PACK_SIZE + 0] += x2_lo_to_f32(in_r[w][0]) - dbias_acc[w * PACK_SIZE + 1] += x2_hi_to_f32(in_r[w][0]) - dbias_acc[w * PACK_SIZE + 2] += x2_lo_to_f32(in_r[w][1]) - dbias_acc[w * PACK_SIZE + 3] += x2_hi_to_f32(in_r[w][1]) - - amax_2x = in_r[0][0] - # Each wave will use max.xorsign.abs.f16x2 or max.xorsign.abs.bf16x2 to compare 2 packed elements in parallel - for w in cutlass.range_constexpr(WAVES): - if cutlass.const_expr(FUSE_RELU): - # Skip the first iteration because we assigned in_r[0][0] to amax_2x already - if cutlass.const_expr(w > 0): - # If we fuse relu then we don't want to do abs since negative value will be set to 0 and they will lose comparison automatically - amax_2x = max_x2(amax_2x, in_r[w][0]) - amax_2x = max_x2(amax_2x, in_r[w][1]) - else: - # Skip the first iteration because we assigned in_r[0][0] to amax_2x already - if cutlass.const_expr(w > 0): - amax_2x = abs_max_x2(amax_2x, in_r[w][0]) - amax_2x = abs_max_x2(amax_2x, in_r[w][1]) - if cutlass.const_expr(FUSE_RELU): - # Compare the 2 packed max without abs - amax_r = cute.arch.fmax( - x2_lo_to_f32(amax_2x), - x2_hi_to_f32(amax_2x), - ) - # For relu the max is at least 0 - amax_r = cute.arch.fmax(amax_r, Float32(0.0)) - else: - # Compare the 2 packed abs max - amax_r = cute.arch.fmax( - fabs_f32(x2_lo_to_f32(amax_2x)), - fabs_f32(x2_hi_to_f32(amax_2x)), - ) - else: - # Since we need to do computation on individual f16 / bf16 elements, we can't read in pack - sX_thread_rw = cute.make_tensor( - sX_thread.iterator, - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), - ) - - if cutlass.const_expr(WITH_DACT): - # Backward: out = grad · act'(act_input). sX is grad, sA is act_input. - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] - 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] - - 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 - - # Each wave we read PACK_SIZE elements, and we have WAVES waves, so we read WAVES * PACK_SIZE (= MXFP8_BLOCK_SCALING_SIZE) elements in total. - in_r = [[None] * PACK_SIZE for _ in range(WAVES)] - # Don't repeat these OOB masking computation in the loop - thread_row_start = tile_row_start + tidx // CTA_THREADS_X - thread_col_start = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE - thread_row_oob = thread_row_start >= M - # PACK_SIZE must divide every term of `thread_col_start + start` and N itself, or a - # pack could straddle the N boundary and the check below would be wrong. TILE_X is a - # tunable, so it is asserted rather than assumed. - assert MXFP8_BLOCK_SCALING_SIZE % PACK_SIZE == 0 - assert SYM_N_DIVISIBILITY % PACK_SIZE == 0 - assert TILE_X % PACK_SIZE == 0 - for w in cutlass.range_constexpr(WAVES): - start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE - # thread_col_start + start >= N means "if the first element in this pack is OOB" - # If the first element is OOB, then (obviously) elements after it within this pack must be OOB - # If the first element is not OOB, then elements after it within this pack must be not OOB either. - # Proof: - # (1) when the first element is not OOB, we have thread_col_start + start < N, - # which implies N - (thread_col_start + start) > 0 - # (2) with the assertions above: - # - N is divisible by SYM_N_DIVISIBILITY, so N is also divisible by PACK_SIZE - # - thread_col_start = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE is divisible by PACK_SIZE - # because TILE_X and MXFP8_BLOCK_SCALING_SIZE are both divisible by PACK_SIZE - # - start = (w * PACK_SIZE + bank_group * PACK_SIZE) % MXFP8_BLOCK_SCALING_SIZE is also divisible by PACK_SIZE - # because MXFP8_BLOCK_SCALING_SIZE is a multiple of PACK_SIZE - # so N - (thread_col_start + start) must be divisible by PACK_SIZE - # Therefore, N - (thread_col_start + start) >= PACK_SIZE because PACK_SIZE is the minimal number - # for (1) and (2) to hold, which implies thread_col_start + start + PACK_SIZE -1 < N - # Hence, we only need to check if the first element is OOB to determine if the whole pack is OOB - thread_col_oob = thread_col_start + start >= N - for i in cutlass.range_constexpr(PACK_SIZE): - x = Float32(sX_thread_rw[0, start + i]) - if cutlass.const_expr(WITH_DACT): - # out = grad · act'(act_input) - x = x * dop(Float32(sA_thread_rw[0, start + i])) - # If IS_ACT, apply activation function to x in f32 - elif cutlass.const_expr(WITH_ACT): - # If it's relu, we can handle it later - if not cutlass.const_expr(FUSE_RELU): - x = op(x) - if not cutlass.const_expr(SKIP_INPUT_MASKING): - # If the input shape is not divisible by the tile size, - # TMA would zero-fills the input tile outside its logical MxN bounds. - # This is fine for non-activation cases, but for activation cases, - # op(0) might not be 0 which will pollute the amax and dbias. - # So we must manually mask the OOB region here. - if thread_row_oob or thread_col_oob: - x = Float32(0.0) - # Accumulate to the per-thread dbias register buffer for this tile if WITH_DBIAS - if cutlass.const_expr(WITH_DBIAS): - # dbias_acc is register buffer so we can just write without bank conflict - dbias_acc[w * PACK_SIZE + i] += x - # If 16-bit input with activation, truncate to IType - if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): - x = truncate_f32(x) - in_r[w][i] = x - if cutlass.const_expr(FUSE_RELU): - amax_r = cute.arch.fmax( - amax_r, x - ) # For relu cases, we don't need abs since negative values will be 0 so they lose comparison automatically - else: - amax_r = cute.arch.fmax(amax_r, fabs_f32(x)) - if cutlass.const_expr(FUSE_RELU): - amax_r = cute.arch.fmax(amax_r, Float32(0.0)) # If relu, the amax is at least 0 - - biased_exp_r = cvt_f32_to_fp8e8m0fnu(amax_r * MAX_NORM_RCP) - - # mS_row_stage has logical shape (32, 2) and we have 64 threads where each is mapped to one scale factor - # The TV layout is equivalent to TV layout with thr_layout=(32, 2):(2, 1), val_layout=(1,) - # but it's too trival so let's just index it directly without using layout - # Note this is the logical layout, which is on top of the swizzled / non-swizzled scale factor layout - # that mappes the logical index to the physical offset - - # For irregular shapes, skip the scale store if this thread's logical row / col-block lies past the input's actual extents. - # TMA already zero-fills OOB input reads and drops OOB output writes; only the direct scale-byte gmem store needs an explicit guard. - # If the input shape is divisible by the tile size, then we won't access OOB regions because - # we never we only access tiles we actually need (num_tiles) which are never OOB - if cutlass.const_expr(SKIP_SCALE_BOUNDS): - mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = biased_exp_r - else: - scale_row = tile_row_start + tidx // CTA_THREADS_X - scale_col_first_elt = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE - if scale_row < M and scale_col_first_elt < N: - mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = biased_exp_r - - inv_scale_r = exp2f_rcp(biased_exp_r) # f32 reciprocal of the scale - scale_2x = pack_f32x2(inv_scale_r, inv_scale_r) - if cutlass.const_expr(USE_HALF_PRECISION): - mul_cvt_x4_func = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, FP8_DTYPE, FUSE_RELU) - else: - mul_cvt_x4_func = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE, FUSE_RELU) - - for w in cutlass.range_constexpr(WAVES): - idx = (w * 4 + offset) % MXFP8_BLOCK_SCALING_SIZE - idx = idx // 4 - if cutlass.const_expr(USE_HALF_PRECISION): - # Convert 2 packed f16/bf16 pairs to 4 fp8 in one fused op - sO_thread_u32[idx] = mul_cvt_x4_func(in_r[w][0], in_r[w][1], scale_2x) - else: - # Convert 4 f32 to 4 fp8 in one fused op - sO_thread_u32[idx] = mul_cvt_x4_func( - in_r[w][0], in_r[w][1], in_r[w][2], in_r[w][3], scale_2x - ) - - return amax_r - - -@cute.jit -def quantize_colwise_mxfp8( - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA - sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) - sO_col_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) - mS_col_stage: cute.Tensor, # colwise scale tensor (1D swizzled, or 2D linear) - MAX_NORM_RCP: cutlass.Constexpr[float], - tile_row_start: Int32, # global row index of this stage's row 0 - tile_col_start: Int32, # global col index of this CTA's col 0 - M: Int32, - N: Int32, - ACTIVATION: cutlass.Constexpr[str | None], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - SWIZZLE: cutlass.Constexpr[bool], - TILE_X: cutlass.Constexpr[int], - TILE_Y: cutlass.Constexpr[int], # pylint: disable=unused-argument # kept for API consistency - SKIP_INPUT_MASKING: cutlass.Constexpr[bool] = False, - SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, - WITH_ACT: cutlass.Constexpr[bool] = False, - 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 -): - """Quantize one SMEM tile colwise to MXFP8 (per-column 32-elt block scales); returns (amax, dbias_partial).""" - tidx, _, _ = cute.arch.thread_idx() - - _, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), - val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), - ) - - sX_tv = cute.composition(sX_tile, tv_layout) - sO_tv = cute.composition(sO_col_tile, tv_layout) - - # I/O Elements that belong to this thread - sX_thread = sX_tv[tidx, None] - sO_thread = sO_tv[tidx, None] - - # PTX allows to fuse relu activation in `cvt.rn.satfinite` unless we need to reduction for dbias - 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) - - if cutlass.const_expr(USE_HALF_PRECISION): - max_scalar = max_scalar_f16 if DTYPE is cutlass.Float16 else max_scalar_bf16 - abs_max_scalar = abs_max_scalar_f16 if DTYPE is cutlass.Float16 else abs_max_scalar_bf16 - to_f32 = to_f32_f16 if DTYPE is cutlass.Float16 else to_f32_bf16 - # If we can use the half precision format, then use the input tile directly since there is no need to upcast - sX_thread_packed16 = cute.make_tensor( - sX_thread.iterator, - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(TILE_X,)), - ) - # Stash the strided column reads in registers (CUDA's in_colwise_IType): - # the cvt loop below reuses them instead of re-reading smem. - in_c = [None] * MXFP8_BLOCK_SCALING_SIZE - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - in_c[i] = sX_thread_packed16[i] - - amax_value = in_c[0] - # Skip the first iteration because we assigned in_c[0] to amax_value already - for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE): - if cutlass.const_expr(FUSE_RELU): - # If we fuse relu then we don't want to do abs since negative value will be set to 0 - # and they will lose comparison automatically - amax_value = max_scalar(amax_value, in_c[i]) - else: - amax_value = abs_max_scalar(amax_value, in_c[i]) - if cutlass.const_expr(FUSE_RELU): - amax_c = cute.arch.fmax(to_f32(amax_value), Float32(0.0)) - else: - amax_c = fabs_f32(to_f32(amax_value)) - else: - # Otherwise we need to case input values to fp32. Allocate the register tensor and load from SMEM input tiles. - rX_thread_f32 = cute.make_rmem_tensor( - layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), - dtype=Float32, - ) - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - 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] - 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] - # Don't repeat these OOB masking computation in the loop - thread_col_start = tile_col_start + tidx - thread_col_oob = thread_col_start >= N - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - # If it's relu, we can handle it later in the cvt - if not cutlass.const_expr(FUSE_RELU): - rX_thread_f32[i] = op(rX_thread_f32[i]) - # If the input shape is not divisible by the tile size, - # TMA would zero-fills the input tile outside its logical MxN bounds. - # This is fine for non-activation cases, but for activation cases, - # op(0) might not be 0 which will pollute the amax and dbias. - # So we must manually mask the OOB region here. - if not cutlass.const_expr(SKIP_INPUT_MASKING): - if tile_row_start + i >= M or thread_col_oob: - rX_thread_f32[i] = Float32(0.0) - # Accumulate fp32 activations to DBIAS before we truncate to half precision when the input is half precision - if cutlass.const_expr(WITH_DBIAS): - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - dbias_partial += rX_thread_f32[i] - # Truncate the activation (after we apply op) back to the half precision type if input is also half precision. - 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 - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - rX_thread_f32[i] = truncate_f32(rX_thread_f32[i]) - # Columnwise is the preferred direction so it runs first. If it needs to cache the activation in the input tile - # to let the rowwise pass read it, we need to cast and overwrite the input data in-place here - if cutlass.const_expr(CACHE_ACTIVATION): - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - sX_thread[i] = DTYPE(rX_thread_f32[i]) - amax_c = Float32(0.0) - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - if cutlass.const_expr(FUSE_RELU): - amax_c = cute.arch.fmax(amax_c, rX_thread_f32[i]) - else: - amax_c = cute.arch.fmax(amax_c, fabs_f32(rX_thread_f32[i])) - - # Irregular shapes: skip when this stage's row range or this thread's - # column lies past the input extents. TILE_Y == MXFP8_BLOCK_SCALING_SIZE so each stage - # is exactly one scale-row; valid iff `tile_row_start < M`. - biased_exp_c = cvt_f32_to_fp8e8m0fnu(amax_c * MAX_NORM_RCP) - # If the input shape is divisible by the tile size, then we won't access OOB regions because - # we never we only access tiles we actually need (num_tiles) which are never OOB - if cutlass.const_expr(SKIP_SCALE_BOUNDS): - if cutlass.const_expr(SWIZZLE): - mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c - else: - mS_col_stage[(0, tidx)] = biased_exp_c - else: - scale_col = tile_col_start + tidx - if tile_row_start < M and scale_col < N: - if cutlass.const_expr(SWIZZLE): - mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c - else: - mS_col_stage[(0, tidx)] = biased_exp_c - - inv_scale_c = exp2f_rcp(biased_exp_c) - # cvt.rn.satfinite can be vectorized to convert 2 f32 to 2 fp8 in one instruction - cvt_x2_func = get_cvt_f32x2_to_fp8x2_func(FP8_DTYPE, FUSE_RELU) - sO_thread_fp8 = cute.make_tensor( - cute.recast_ptr(sO_thread.iterator, dtype=FP8_DTYPE), sO_thread.layout - ) - if cutlass.const_expr(USE_HALF_PRECISION): - to_f32 = to_f32_f16 if DTYPE is cutlass.Float16 else to_f32_bf16 - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): - lo, hi = 2 * j, 2 * j + 1 - v_lo = to_f32(in_c[lo]) - v_hi = to_f32(in_c[hi]) - # Accumulate the per-thread column partial for dbias if WITH_DBIAS. - # Kept as two adds in element order: f32 addition is not associative - if cutlass.const_expr(WITH_DBIAS): - dbias_partial += v_lo - dbias_partial += v_hi - pair = cvt_x2_func(v_hi * inv_scale_c, v_lo * inv_scale_c) - sO_thread_fp8[lo] = Uint8(pair & Uint16(0xFF)).bitcast(FP8_DTYPE) - sO_thread_fp8[hi] = Uint8(pair >> Uint16(8)).bitcast(FP8_DTYPE) - else: - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): - lo, hi = 2 * j, 2 * j + 1 - pair = cvt_x2_func(rX_thread_f32[hi] * inv_scale_c, rX_thread_f32[lo] * inv_scale_c) - sO_thread_fp8[lo] = Uint8(pair & Uint16(0xFF)).bitcast(FP8_DTYPE) - sO_thread_fp8[hi] = Uint8(pair >> Uint16(8)).bitcast(FP8_DTYPE) - - # Return this stage's per-column partial alongside amax; the caller accumulates - # it across stages (a scalar can't be updated in-place through the arg). - return amax_c, dbias_partial - - -@cute.jit -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], - 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 -): - """Reduce one SMEM tile along rows (the dbias direction) and possibly cache dact values, with no quantization.""" - tidx, _, _ = cute.arch.thread_idx() - - # Same TV layout as quantize_colwise_mxfp8: thread tidx owns column tidx, all TILE_Y rows. - _, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), - val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), - ) - sX_thread = cute.composition(sX_tile, tv_layout)[tidx, None] - - dbias_partial = Float32(0.0) - if cutlass.const_expr(WITH_DACT): - dop = SUPPORTED_DACTIVATIONS[ACTIVATION] - 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])) - dbias_partial += x - if cutlass.const_expr(CACHE_ACTIVATION): - sX_thread[i] = DTYPE(x) - else: - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - dbias_partial += Float32(sX_thread[i]) - return dbias_partial - - -@cute.jit -def quantize_bidimensional_mxfp8_swizzled( - sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage), Sw<3,4,3> swizzled - sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile - sO_col_tile: cute.Tensor, # 32x64 colwise-output tile - sS_row_tile: cute.Tensor, # (32, 2) smem rowwise-scale staging tile (flushed in the epilogue) - sS_col_tile: cute.Tensor, # (1, 64) smem colwise-scale staging tile (flushed in the epilogue) - sColReduce_warp: cute.Tensor, # (32,) SMEM fp32 columnwise scale reduction buffer for this warp - WARPS_PER_CTA: cutlass.Constexpr[int], - MAX_NORM_RCP: cutlass.Constexpr[float], - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], -): - """Quantize a pre-sliced 32x64 tile -> both rowwise and colwise MXFP8. Elements are - addressed through TV layouts: a warp = 32 lanes = 32 rows (lane == row), and each - thread owns one 32-col row segment (its 32-col block) and one output column. - - No bounds check: the caller only loops over valid tiles (`num_tiles`), and - M % 32 == 0 means a tile can only be partial along N. The partial tile's - out-of-bounds warp still runs, but harmlessly -- it reads TMA zero-filled smem, - its output columns are masked by the caller's TMA store, and its scale writes - (rowwise and colwise alike) go to staging slots whose flush targets are past-N - padding columns of the respective scale tensors. - """ - mul_cvt4 = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE) - mul_cvt4_elemwise = mul_f32x4_cvt_f32x4_to_fp8x4(FP8_DTYPE) - - _, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) - val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE)), - ) - _, tv_layout_rowwise_scale = cute.make_layout_tv( - thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) - val_layout=cute.make_layout((1, 1)), - ) - _, tv_layout_colwise_scale = cute.make_layout_tv( - thr_layout=cute.make_layout((1, (MXFP8_BLOCK_SCALING_SIZE, WARPS_PER_CTA))), # (1, (32, 2)) - val_layout=cute.make_layout((1, 1)), - ) - - tidx, _, _ = cute.arch.thread_idx() - lane = tidx % MXFP8_BLOCK_SCALING_SIZE - - # Each composed [tidx, None] slice is 1-D (the value mode is flattened): the data - # slices are size 32 (this thread's row segment), the scale slices size 1. - tXsX = cute.composition(sX_tile, tv_layout)[tidx, None] # (32,) input row segment - tXsO_row = cute.composition(sO_row_tile, tv_layout)[tidx, None] # (32,) rowwise out - tXsO_col = cute.composition(sO_col_tile, tv_layout)[tidx, None] # (32,) colwise out - tSsS_row_tile = cute.composition(sS_row_tile, tv_layout_rowwise_scale)[tidx, None] # (1,) - tSsS_col_tile = cute.composition(sS_col_tile, tv_layout_colwise_scale)[tidx, None] # (1,) - - rO_row = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) - rO_col = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) - rO_row_u32 = cute.make_tensor( - cute.recast_ptr(rO_row.iterator, dtype=Uint32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), - ) - rO_col_u32 = cute.make_tensor( - cute.recast_ptr(rO_col.iterator, dtype=Uint32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), - ) - sColReduce = cute.make_tensor( - sColReduce_warp.iterator, - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4, 4), stride=(4, 1)), - ) - - if cutlass.const_expr(is_packed16(DTYPE)): - # If the input is bf16 / fp16, take this fast path and process 2 elements at a time in a packed i32 - abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 - x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 - x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 - rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, DTYPE) - # Do a vectorized load from SMEM to RMEM and unswizzle in the meantime. - cute.autovec_copy(tXsX, rX) - rX_2x = cute.make_tensor( - cute.recast_ptr(rX.iterator, dtype=Int32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - sColReduce_2x = cute.make_tensor( - cute.recast_ptr(sColReduce.iterator, dtype=Int64), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - - row_amax2 = rX_2x[0] - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): - pair = rX_2x[i] - # Skip the first iteration because we assigned rX_2x[0] to row_amax2 already - if cutlass.const_expr(i > 0): - row_amax2 = abs_max_x2(row_amax2, pair) - a_lo = fabs_f32(x2_lo_to_f32(pair)) - a_hi = fabs_f32(x2_hi_to_f32(pair)) - col_lo = cute.arch.warp_redux_sync(a_lo, kind="fmax") - col_hi = cute.arch.warp_redux_sync(a_hi, kind="fmax") - with cute.arch.elect_one(): - sColReduce_2x[i] = pack_f32x2(col_lo, col_hi) - - # Compute the rowwise scale factor - row_amax = cute.arch.fmax( - fabs_f32(x2_lo_to_f32(row_amax2)), fabs_f32(x2_hi_to_f32(row_amax2)) - ) - row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) - row_inv = exp2f_rcp(row_exp) - tSsS_row_tile[0] = row_exp - cute.arch.sync_warp() - - # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) - col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) - tSsS_col_tile[0] = col_exp - sColReduce_warp[lane] = exp2f_rcp(col_exp) - cute.arch.sync_warp() - - row_scale_2x = pack_f32x2(row_inv, row_inv) - # Vectorized multiply-and-convert: 4 f32 → 4 fp8 - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - # Vectorized load for 4 columnwise scale factors from SMEM to RMEM - col_inv4 = cute.make_rmem_tensor(4, Float32) - cute.autovec_copy(sColReduce[j, None], col_inv4) # LDS.128, warp-broadcast - p01 = rX_2x[2 * j] - p23 = rX_2x[2 * j + 1] - f0 = x2_lo_to_f32(p01) - f1 = x2_hi_to_f32(p01) - f2 = x2_lo_to_f32(p23) - f3 = x2_hi_to_f32(p23) - # For rowwise quantized values, they use the same scale for all 4 elements, - # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. - rO_row_u32[j] = mul_cvt4(f0, f1, f2, f3, row_scale_2x) - # For columnwise quantized values, each element has its own scale; the - # elementwise variant fuses the per-element multiply into the cvt sequence. - rO_col_u32[j] = mul_cvt4_elemwise( - f0, f1, f2, f3, col_inv4[0], col_inv4[1], col_inv4[2], col_inv4[3] - ) - cute.autovec_copy(rO_row, tXsO_row) - cute.autovec_copy(rO_col, tXsO_col) - else: - # If input is fp32, take this slow path and process it normally without packing - rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, Float32) - cute.autovec_copy(tXsX, rX) - - row_amax = Float32(0.0) - for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - a = fabs_f32(rX[c]) - row_amax = cute.arch.fmax(row_amax, a) - col_amax = cute.arch.warp_redux_sync(a, kind="fmax") - with cute.arch.elect_one(): - sColReduce_warp[c] = col_amax - - # Compute the rowwise scale factor - row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) - row_inv = exp2f_rcp(row_exp) - tSsS_row_tile[0] = row_exp # rowwise scale (this row-block, staged in smem) - cute.arch.sync_warp() - - # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) - col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) - tSsS_col_tile[0] = col_exp # colwise scale (this thread's column) - sColReduce_warp[lane] = exp2f_rcp(col_exp) - cute.arch.sync_warp() - - row_scale_2x = pack_f32x2(row_inv, row_inv) - # Vectorized multiply-and-convert: 4 f32 → 4 fp8 - for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - # Vectorized load for 4 columnwise scale factors from SMEM to RMEM - col_inv4 = cute.make_rmem_tensor(4, Float32) - cute.autovec_copy(sColReduce[j, None], col_inv4) - offset = 4 * j - # For rowwise quantized values, they use the same scale for all 4 elements, - # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. - rO_row_u32[j] = mul_cvt4( - rX[offset], rX[offset + 1], rX[offset + 2], rX[offset + 3], row_scale_2x - ) - # For columnwise quantized values, each element has its own scale; the - # elementwise variant fuses the per-element multiply into the cvt sequence. - rO_col_u32[j] = mul_cvt4_elemwise( - rX[offset], - rX[offset + 1], - rX[offset + 2], - rX[offset + 3], - col_inv4[0], - col_inv4[1], - col_inv4[2], - col_inv4[3], - ) - cute.autovec_copy(rO_row, tXsO_row) - cute.autovec_copy(rO_col, tXsO_col) - - -@cute.jit -def noop_flag_is_set(mNoop: cute.Pointer) -> Boolean: - """Whether the cast_noop flag says this quantization is a no-op and must be skipped. - - mNoop is a pointer rather than a tensor so that one compiled kernel serves both a present and - an absent flag, hence the address is checked before it is dereferenced, exactly like the CUDA - C++ kernel's `noop != nullptr && noop[0] == 1.0f`. The two checks cannot be joined with `and`, - which the DSL lowers to a non-short-circuiting op that would load from the null pointer. - """ - flag_is_set = Boolean(False) - if mNoop.toint() != Int64(0): - flag_is_set = cute.make_tensor(mNoop, cute.make_layout((1,)))[0] == Float32(1.0) - return flag_is_set - - -class MXFP8QuantizeConfig: - """Configs for the compiled CuTeDSL kernel. These will be fixed once the kernel is compiled and - they will behave as const expressions. - """ - - def __init__( - self, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - with_gemm_swizzled_scales: bool, - with_amax: bool, - with_dbias: bool = False, - with_dact: bool = False, - with_act: bool = False, - use_2d_quantization: bool = False, - activation: Optional[str] = None, - ): - if use_2d_quantization: - raise ValueError("2D block scaling is not implemented by the CuTeDSL MXFP8 kernels") - if dtype is None or dtype not in ("Float32", "Float16", "BFloat16"): - raise ValueError(f"unknown input dtype {dtype!r}; expected Float32|Float16|BFloat16") - self.DTYPE = str_to_cutlass_dtype(dtype) - self.DTYPE_STR = dtype # readable input-dtype token, for __str__ - if fp8_dtype not in ("Float8E4M3", "Float8E5M2"): - raise ValueError( - f"unknown FP8 dtype {fp8_dtype!r}; expected 'Float8E4M3' or 'Float8E5M2'" - ) - self.FP8_DTYPE = str_to_cutlass_dtype(fp8_dtype) - self.FP8_DTYPE_STR = fp8_dtype # readable token, for __str__ - self.ROWWISE = rowwise - self.COLWISE = colwise - if not (rowwise or colwise): - raise ValueError("at least one of rowwise or colwise must be true") - self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales - self.WITH_AMAX = with_amax - if not with_dact and not with_act: - if activation == "none": - self.ACTIVATION = None - else: - raise ValueError( - "activation must be none when with_dact and with_act are both False" - ) - else: - if with_dact and with_act: - raise ValueError( - "with_dact and with_act cannot be true at the same time since they are used for" - " different paths (bwd vs fwd)" - ) - if with_dact: - if activation in SUPPORTED_DACTIVATIONS: - self.ACTIVATION = activation - else: - raise ValueError( - f"unknown activation {activation!r} for with_dact=True; expected one of" - f" {sorted(SUPPORTED_DACTIVATIONS)}" - ) - elif with_act: - if activation in SUPPORTED_ACTIVATIONS: - self.ACTIVATION = activation - else: - raise ValueError( - f"unknown activation {activation!r} for with_act=True; expected one of" - f" {sorted(SUPPORTED_ACTIVATIONS)}" - ) - self.WITH_DACT = with_dact - self.WITH_ACT = with_act - # 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 - # kernel's COLWISE_SCALING / rowwise dbias branches. - self.WITH_DBIAS = with_dbias - self.MAX_NORM_RCP = ( - FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP - ) - - def __str__(self): - return ( - f"MXFP8QuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " - 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})" - ) - - __repr__ = __str__ - - -class MXFP8QuantizeKernelBase(abc.ABC): - """Base class for MXFP8 quantize kernels.""" - - @abc.abstractmethod - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, - mDActInput: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - stream: CUstream, - ): - """ - Compiled kernel entrypoint (decorate with @cute.jit). - All MXFP8 quantize kernels must implement this interface because this is our C++ call site's contract in `quantize_mxfp8_cutedsl.cuh`. - C++ call site will pass the arguments in this exact order via tvm-ffi and the kernel must accept them even if they are not used. - """ - - -class MXFP8QuantizeKernel(MXFP8QuantizeKernelBase): - """The MXFP8 quantization kernel that mirrors the standard (non-specialized) MXFP8 CUDA C++ quantization kernel - with multiple fusions (activation, dbias, etc.). - `__call__` method is the entrypoint which is AOT compiled. `self` will be captured so it's fixed per compiled kernel - """ - - # Vectorised access constants for bank-conflict avoidance (rowwise pass) - _PACK_SIZE = 4 # Elements per vector load - # Each thread reads 8 waves with each wave reads 4 packed bf16, so it reads a whole MXFP8 block in total - _WAVES = MXFP8_BLOCK_SCALING_SIZE // _PACK_SIZE - _TOTAL_BANKS_WIDTH = (32 * 4) // 1 # 32 banks × 4 bytes, in bytes - _THREADS_PER_BANK = _TOTAL_BANKS_WIDTH // MXFP8_BLOCK_SCALING_SIZE # 4 threads per bank - _NUM_STAGES = 2 # The pipeline depth is always 2 - - def __init__( - self, - cfg: MXFP8QuantizeConfig, - SKIP_INPUT_MASKING: bool = False, - SKIP_SCALE_BOUNDS: bool = False, - ): - self.cfg = cfg - # If the input shape is divisible by the tile size or f(0)=0 holds for activaions, - # we can skip masking inputs with zero in the kernel and save some instructions. - self.SKIP_INPUT_MASKING = SKIP_INPUT_MASKING - # If the input shape is divisible by the tile size, we can skip bounds check for scale writes and save some instructions. - self.SKIP_SCALE_BOUNDS = SKIP_SCALE_BOUNDS - # Only honor the noop flag when no activation or dbias is fused to match CUDA C++'s implementation - self.CHECK_NOOP_FLAG: cutlass.const_expr = ( - not self.cfg.WITH_ACT and not self.cfg.WITH_DACT and not self.cfg.WITH_DBIAS - ) - cast_dbias_only = cfg.WITH_DBIAS and not cfg.WITH_DACT and not cfg.WITH_ACT - # Use a different tile size for dbias only config - # No matter what tile size we use, each thread always handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk - if cast_dbias_only: - self._NUM_TILES = 4 # Each CTA handles 4 tiles stacked vertically - self._THREADS_PER_CTA = 128 - else: - self._NUM_TILES = 2 # Each CTA handles 2 tiles stacked vertically - self._THREADS_PER_CTA = 64 - # Each thread handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk - self._TILE_COLS = self._THREADS_PER_CTA - self._TILE_ROWS = MXFP8_BLOCK_SCALING_SIZE - self._NUM_WARPS = self._THREADS_PER_CTA // 32 - # We prefer to do dbias reduction in colwise which is easier (no cross-thread reduction needed). - # Only do rowwise reduction when we don't quantize columnwisely when WITH_DBIAS is True. - - # If columnwise is not quantized, and dbias is fused, we will do a columnwise reduction only pass - # (may also cache dact values) without quantization because doing dbias in rowwise requires a SMEM shuffle - # which is even slower - self.COLWISE_DBIAS_REDUCTION_ONLY = cfg.WITH_DBIAS and not cfg.COLWISE - # If columnwise is quantized and dbias is fused, do dbias reduction in colwise which is easier - # (no cross-thread reduction needed). - self.DBIAS_REDUCTION_IN_COLWISE = cfg.WITH_DBIAS and cfg.COLWISE - # Never do dbias reduction in rowwise because it is slow. Leave this so in case it's enabled in the future - self.DBIAS_REDUCTION_IN_ROWWISE = False - - # Cache activation in-place in the SMEM input tile when we process both rowwise and colwise passes - # so the activation is only computed once in the direction we favor (columnwise) and the other direction (rowwise) - # reads the cached value instead of recomputing it. - # Note: if activation is relu, there is no standalong relu applied because it's already fused into `cvt.rn.satfinite` - # so it should be treated as "no activation" - self.CACHE_ACTIVATION = ( - (cfg.WITH_ACT or cfg.WITH_DACT) - and cfg.ROWWISE - # Always cache activation in the colwise quantization or dbias reduction only pass - and (cfg.COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY) - and cfg.ACTIVATION != "relu" - ) - # The global tensor amax (mAmax) is the max over ALL elements. Each direction's - # per-block amaxes already span every element, so when both passes run we only - # fold the global amax from one of them — favor colwise (matches the flags - # above). The per-block *scale* amax is still computed in each pass for its own - # scale; this only skips the redundant global comparison in the other pass. - self.AMAX_FROM_COLWISE = cfg.WITH_AMAX and cfg.COLWISE - self.AMAX_FROM_ROWWISE = cfg.WITH_AMAX and not cfg.COLWISE - - @cute.jit - def __call__( - self, - mX: cute.Tensor, # Input tensor to quantize - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], # Rowwise output and scale tensors - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], # Colwise output and scale tensors - mAmax: Optional[cute.Tensor], # Global amax accumulator, only used when WITH_AMAX is True - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device - mDActInput: Optional[ - cute.Tensor - ], # Activation input for activation derivative fusion, only used when WITH_DACT is True - mWorkspace: Optional[ - cute.Tensor - ], # Workspace for the dbias reduction, only used when WITH_DBIAS is True - stream: CUstream, - ): - if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf(f"[CuTeDSL] MXFP8QuantizeKernel.__call__() with config: {self.cfg}\n") - - M = mX.shape[0] - N = mX.shape[1] - cfg = self.cfg - - # If WITH_GEMM_SWIZZLED_SCALES is enabled, the output must satisfy cublas's swizzled layout - # This is expressed as a CuTe layout applied to the output tensor so it can be transparent throughout the kernel implementation. - # See https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout for more details. - if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_row, mS_col = derive_swizzled_scale_layout( - M, N, cfg.ROWWISE, cfg.COLWISE, mS_row, mS_col - ) - - # We have 2 stages in our pipeline where each stage loads / computes a (TILE_Y, TILE_X) tile - smem_tile_layout = cute.make_ordered_layout( - (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) - ) - cta_tiler = (self._TILE_ROWS, self._TILE_COLS) - - # Input TMA atoms - op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() - tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_load, - mX, - smem_tile_layout, - cta_tiler, - num_multicast=1, - ) - - # Activation input TMA atoms for activation derivative fusion - tma_atom_act = None - tma_src_act = None - if cutlass.const_expr(cfg.WITH_DACT): - tma_atom_act, tma_src_act = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_load, - mDActInput, - smem_tile_layout, - cta_tiler, - num_multicast=1, - ) - - # Output TMA atoms - op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() - out_smem_layout = cute.make_ordered_layout((self._TILE_ROWS, self._TILE_COLS), order=(1, 0)) - tma_atom_out_row = None - tma_dst_out_row = None - tma_atom_out_col = None - tma_dst_out_col = None - if cutlass.const_expr(cfg.ROWWISE): - tma_atom_out_row, tma_dst_out_row = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_store, - mO_row, - out_smem_layout, - cta_tiler, - num_multicast=1, - ) - if cutlass.const_expr(cfg.COLWISE): - tma_atom_out_col, tma_dst_out_col = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_store, - mO_col, - out_smem_layout, - cta_tiler, - num_multicast=1, - ) - - grid = [ - cute.ceil_div(Int32(N), self._TILE_COLS), - cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), - ] - block = [ - self._THREADS_PER_CTA, - ] - - self.kernel( - mX, - mS_row, - mS_col, - mAmax, - mNoop, - mWorkspace, - mX.element_type, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - tma_atom_act, - tma_src_act, - ).launch( - grid=grid, - block=block, - stream=stream, - ) - - @cute.kernel - def kernel( - self, - mX: cute.Tensor, - mS_row: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, - mWorkspace: Optional[cute.Tensor], - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, # Input TMA atoms - tma_atom_out_row: Optional[cute.CopyAtom], - tma_dst_out_row: Optional[cute.Tensor], # Rowwise output TMA atoms - tma_atom_out_col: Optional[cute.CopyAtom], - tma_dst_out_col: Optional[cute.Tensor], # Colwise output TMA atoms - tma_atom_act: Optional[cute.CopyAtom], - tma_src_act: Optional[ - cute.Tensor - ], # Activation derivative TMA atoms, None unless WITH_DACT - ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - - skip_execution = False - if cutlass.const_expr(self.CHECK_NOOP_FLAG): - skip_execution = noop_flag_is_set(mNoop) - if not skip_execution: - self._kernel_main( - mX, - mS_row, - mS_col, - mAmax, - mWorkspace, - dtype, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - tma_atom_act, - tma_src_act, - ) - - @cute.jit - def _kernel_main( - self, - mX: cute.Tensor, - mS_row: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, # Input TMA atoms - tma_atom_out_row: Optional[cute.CopyAtom], - tma_dst_out_row: Optional[cute.Tensor], # Rowwise output TMA atoms - tma_atom_out_col: Optional[cute.CopyAtom], - tma_dst_out_col: Optional[cute.Tensor], # Colwise output TMA atoms - tma_atom_act: Optional[cute.CopyAtom], - tma_src_act: Optional[ - cute.Tensor - ], # Activation derivative TMA atoms, None unless WITH_DACT - ): - cfg = self.cfg - - if cutlass.const_expr(cfg.ROWWISE): - mS_row = cute.zipped_divide( - mS_row, (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) - ) - if cutlass.const_expr(cfg.COLWISE): - mS_col = cute.zipped_divide( - mS_col, (self._TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, self._TILE_COLS) - ) - - # Allocate shared memory for the input and rowwise / columnwise outputs - FP8_DTYPE = cfg.FP8_DTYPE - - if cutlass.const_expr(cfg.ROWWISE and cfg.COLWISE): - - @cute.struct - class SharedStorage: - mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] - sX: cute.struct.Align[ - cute.struct.MemRange[ - dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sO_row: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sO_col: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] - - elif cutlass.const_expr(cfg.ROWWISE and not cfg.COLWISE): - - @cute.struct - class SharedStorage: - mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] - sX: cute.struct.Align[ - cute.struct.MemRange[ - dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sO_row: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] - - else: - - @cute.struct - class SharedStorage: - mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] - sX: cute.struct.Align[ - cute.struct.MemRange[ - dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sO_col: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] - - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - # Apply the layout to the allocated shared memory buffers so the first rank is the tile (nested layout) - # and the second rank is the pipeline stage - sX = storage.sX.get_tensor( - cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - ) - if cutlass.const_expr(cfg.ROWWISE): - sO_row = storage.sO_row.get_tensor( - cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - ) - if cutlass.const_expr(cfg.COLWISE): - sO_col = storage.sO_col.get_tensor( - cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - ) - - # Allocate shared memory for the activation input used for the activation derivative fusion. - if cutlass.const_expr(cfg.WITH_DACT): - - @cute.struct - class DactStorage: - sActInput: cute.struct.Align[ - cute.struct.MemRange[ - dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - - dact_storage = smem.allocate(DactStorage) - # Apply the same layout as the input - sActInput = dact_storage.sActInput.get_tensor( - cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - ) - - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) - - # Prefetch TMA descriptors - if warp_idx == 0: - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) - if cutlass.const_expr(cfg.WITH_DACT): - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_act) - - tidx, _, _ = cute.arch.thread_idx() - bidx, bidy, _ = cute.arch.block_idx() - - # Only warp 0 is the producer (issues TMA) - producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) - # Every warp is the consumer (reads the data loaded by TMA) - consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) - - # Bytes transferred per TMA copy: one (TILE_Y, TILE_X) tile of dtype. - tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 - # dact loads two tiles (grad + act_input) under the same per-stage barrier, - # so the barrier must expect both copies' bytes. - if cutlass.const_expr(cfg.WITH_DACT): - tx_count *= 2 - - mainloop_pipeline = pipeline.PipelineTmaAsync.create( - barrier_storage=storage.mbar_storage.data_ptr(), - num_stages=self._NUM_STAGES, - producer_group=producer_group, - consumer_group=consumer_group, - tx_count=tx_count, - cta_layout_vmnk=None, # single-CTA, no cluster/multicast - ) - - prod_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self._NUM_STAGES - ) - cons_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self._NUM_STAGES - ) - - M = mX.shape[0] - N = mX.shape[1] - - num_tiles = cutlass.min( - self._NUM_TILES, - cute.ceil_div(M - bidy * self._TILE_ROWS * self._NUM_TILES, self._TILE_ROWS), - ) - - # Tile the TMA gmem view: ((TILE_Y, TILE_X), (M/TILE_Y, N/TILE_X)). - gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) - - # Partition sX/gX for the TMA atom (single-CTA, no cluster/multicast). - tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( - tma_atom, - 0, # Use the only CTA to do the TMA copy - cute.make_layout(1), # This cluster only has 1 CTAs - sX, - gX_tiled, - ) - - # If WITH_DACT, partition the activation input for TMA as well in the same way - if cutlass.const_expr(cfg.WITH_DACT): - gA_tiled = cute.zipped_divide(tma_src_act, (self._TILE_ROWS, self._TILE_COLS)) - tXsA, tXgA = cute.nvgpu.cpasync.tma_partition( - tma_atom_act, - 0, - cute.make_layout(1), - sActInput, - gA_tiled, - ) - - # Partitioning for rowwise / columnwise outputs - if cutlass.const_expr(cfg.ROWWISE): - gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self._TILE_ROWS, self._TILE_COLS)) - tXsO_row, tXgO_row = cute.nvgpu.cpasync.tma_partition( - tma_atom_out_row, - 0, - cute.make_layout(1), - sO_row, - gO_row_tiled, - ) - if cutlass.const_expr(cfg.COLWISE): - gO_col_tiled = cute.zipped_divide(tma_dst_out_col, (self._TILE_ROWS, self._TILE_COLS)) - tXsO_col, tXgO_col = cute.nvgpu.cpasync.tma_partition( - tma_atom_out_col, - 0, - cute.make_layout(1), - sO_col, - gO_col_tiled, - ) - - # Ensure barrier init is visible to all threads before the pipeline is used. - cute.arch.sync_threads() - - # Prologue: warp 0 prefetches up to NUM_STAGES tiles to fully fill the pipeline - if warp_idx == 0: - for s in cutlass.range_constexpr(self._NUM_STAGES): - if s < num_tiles: - mainloop_pipeline.producer_acquire(prod_state) - tile_y = bidy * self._NUM_TILES + s - cute.copy( - tma_atom, - tXgX[(None, (tile_y, bidx))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - if cutlass.const_expr(cfg.WITH_DACT): - cute.copy( - tma_atom_act, - tXgA[(None, (tile_y, bidx))], - tXsA[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - - # Per-thread amax accumulator - if cutlass.const_expr(cfg.WITH_AMAX): - per_thread_amax = Float32(0.0) - - # Prepare thread-level register accumulators for rowwise dbias reduction. - # Each thread will process two (1, MXFP8_BLOCK_SCALING_SIZE) rows in two stages, and in each stage the thread will add the - # (after dact applied) value to this register array with the same shape so it carries the the two stages' partial sum. - # Then it will be written to a SMEM buffer to let the whole CTA do the reduction separately to yield - # the final (1, TILE_X) dbias workspace output. - rowwise_dbias_acc = None - if cutlass.const_expr(self.DBIAS_REDUCTION_IN_ROWWISE): - rowwise_dbias_acc = cute.make_rmem_tensor( - layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), - dtype=Float32, - ) - # Zero the accumulator registers. - for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): - 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) - - # 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): - mainloop_pipeline.consumer_wait(cons_state) - # Only allow at most _NUM_STAGES-1 stages to be in-flight, because this iteration will reuse the ring buffer - # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish - if warp_idx == 0: - cute.arch.cp_async_bulk_wait_group(self._NUM_STAGES - 1, read=True) - cute.arch.sync_threads() - # The current pipeline stage index, which is the tile index modulo the number of stages. - # This is used to index into the shared memory ring buffers that are wrapped around the number of stages. - stage_idx = cons_state.index - sX_tile = sX[(None, stage_idx)] - # Also fetch the activation input if WITH_DACT - sActInput_tile = None - if cutlass.const_expr(cfg.WITH_DACT): - sActInput_tile = sActInput[(None, stage_idx)] - # Each CTA handles `NUM_TILES` tiles stacked vertically, so tile_idx_x is just the block index along X dimension - # and tile_idx_y is the tile that this stage handles out of the `NUM_TILES` tiles - tile_idx_x = bidx - tile_idx_y = bidy * self._NUM_TILES + tile_idx - # Process rowwise and colwise quantization separately - if cutlass.const_expr(cfg.COLWISE): - # The first row that belongs to this CTA. Each CTA handles NUM_TILES of (TILE_Y, TILE_X) tiles stacked vertically, - # and each stage handles one of them. - sO_col_tile = sO_col[(None, stage_idx)] - mS_col_stage = cute.flatten(mS_col[(None, (tile_idx_y, tile_idx_x))]) - - amax_c, dbias_c = self._process_colwise( - sX_tile, - sO_col_tile, - mS_col_stage, - tile_idx_y * self._TILE_ROWS, - bidx * self._TILE_COLS, - M, - N, - sActInput_tile, - ) - if cutlass.const_expr(self.AMAX_FROM_COLWISE): - per_thread_amax = cute.arch.fmax(per_thread_amax, amax_c) - if cutlass.const_expr(self.DBIAS_REDUCTION_IN_COLWISE): - block_dbias += dbias_c - # If we don't quantize columnwise but we fuse dbias, do a dbias reduction only pass in the columnwise direction without quantization - if cutlass.const_expr(self.COLWISE_DBIAS_REDUCTION_ONLY): - block_dbias += self._dbias_only_colwise(sX_tile, sActInput_tile) - # 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. - if cutlass.const_expr(self.CACHE_ACTIVATION): - cute.arch.sync_threads() - if cutlass.const_expr(cfg.ROWWISE): - sO_row_tile = sO_row[(None, stage_idx)] - # mS_row is ((SCALE_TILE), (GRID)) where SCALE_TILE = (32, 2). - # Each CTA owns NUM_TILES consecutive row-tiles of GRID. cute - # auto-decomposes the flat row coord `bidy * NUM_TILES + tile_idx` - # onto GRID's hierarchical row modes — which is the - # (i_hi, tile_Y) tile-major order for swizzled, and the plain - # row-tile order for compact. Same source, both layouts correct. - mS_row_stage = cute.flatten(mS_row[(None, (tile_idx_y, tile_idx_x))]) - amax_r = self._process_rowwise( - sX_tile, - sO_row_tile, - mS_row_stage, - tile_idx_y * self._TILE_ROWS, - bidx * self._TILE_COLS, - M, - N, - sActInput_tile, - rowwise_dbias_acc, - ) - - if cutlass.const_expr(self.AMAX_FROM_ROWWISE): - per_thread_amax = cute.arch.fmax(per_thread_amax, amax_r) - - # Make the shared-memory writes visible to the TMA's async proxy before the TMA reads them. - cute.arch.fence_proxy( - "async.shared", - space="cta", - ) - cute.arch.sync_threads() - - # We are done with this input pipeline SMEM buffer, signal the producer that it can write to this buffer - mainloop_pipeline.consumer_release(cons_state) - - # Warp 0 issues TMA copy to write the quantized output tile from shared memory to global memory and then commits - if warp_idx == 0: - tile_y = bidy * self._NUM_TILES + tile_idx - if cutlass.const_expr(cfg.ROWWISE): - cute.copy( - tma_atom_out_row, - tXsO_row[(None, stage_idx)], - tXgO_row[(None, (tile_y, bidx))], - ) - if cutlass.const_expr(cfg.COLWISE): - cute.copy( - tma_atom_out_col, - tXsO_col[(None, stage_idx)], - tXgO_col[(None, (tile_y, bidx))], - ) - cute.arch.cp_async_bulk_commit_group() - - cons_state.advance() - - # The pipeline is no longer fully filled after we consume this tile, so we fetch a new tile to fill the pipeline. - # The next _NUM_STAGES-1 tiles are already in-flight, so the next tile to fetch is after _NUM_STAGES tiles. - if warp_idx == 0: - next_tile_idx = tile_idx + self._NUM_STAGES - if next_tile_idx < num_tiles: - mainloop_pipeline.producer_acquire(prod_state) - tile_y = bidy * self._NUM_TILES + next_tile_idx - cute.copy( - tma_atom, - tXgX[(None, (tile_y, bidx))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - if cutlass.const_expr(cfg.WITH_DACT): - cute.copy( - tma_atom_act, - tXgA[(None, (tile_y, bidx))], - tXsA[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - # End of the main pipeline loop - - # Complete the cross-thread dbias reduction after each thread has its own per-thread partial sum after the rowwise quantization. - if cutlass.const_expr(self.DBIAS_REDUCTION_IN_ROWWISE): - # If we do the dbias reduction in the rowwise pass, each thread will have a (1, MXFP8_BLOCK_SCALING_SIZE) partial sum - # and we need to write these to a SMEM buffer and let each thread reduce it in the columnwise direction - block_dbias = self._dbias_reduction_rowwise_epilouge(smem, tidx, rowwise_dbias_acc) - - # Write the per-tile reduced dbias to the global workspace. - if cutlass.const_expr(cfg.WITH_DBIAS): - dbias_col = bidx * self._TILE_COLS + tidx - if dbias_col < N: - mWorkspace[(bidy, dbias_col)] = block_dbias - - if cutlass.const_expr(cfg.WITH_AMAX): - sAmax = storage.sAmax.get_tensor(cute.make_layout(self._NUM_WARPS)) - self._amax_epilogue(sAmax, mAmax, tidx, warp_idx, per_thread_amax) - - # Wait for in-flight TMA stores so data is visible to the host - # before the kernel returns. - cute.arch.cp_async_bulk_wait_group(0, read=False) - - @cute.jit - def _dbias_only_colwise( - self, - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA - sActInput_tile: Optional[ - cute.Tensor - ] = None, # (TILE_Y, TILE_X) act_input tile (dact only), - ): - return dbias_reduction_colwise( - sX_tile, - sActInput_tile, - ACTIVATION=self.cfg.ACTIVATION, - DTYPE=self.cfg.DTYPE, - TILE_X=self._TILE_COLS, - WITH_DACT=self.cfg.WITH_DACT, - CACHE_ACTIVATION=self.CACHE_ACTIVATION, - ) - - @cute.jit - def _dbias_reduction_rowwise_epilouge( - self, smem: cutlass.utils.SmemAllocator, tidx: Int32, rowwise_dbias_acc: cute.Tensor - ): - # Pad the buffer to avoid bank conflicts. The logical shape is still the same. Only the stride is different. - DBIAS_BUFF_WIDTH = ( - self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE * (MXFP8_BLOCK_SCALING_SIZE + 1) - ) - - # Allocate the SMEM buffer that all threads use to reduce the two-stage partial sum (per thread) to the - # partial sum (per block). - @cute.struct - class DbiasStorage: - sDbias: cute.struct.MemRange[Float32, self._TILE_ROWS * DBIAS_BUFF_WIDTH] - - dbias_storage = smem.allocate(DbiasStorage) - sDbias = dbias_storage.sDbias.get_tensor( - cute.make_layout((self._TILE_ROWS, self._TILE_COLS), stride=(DBIAS_BUFF_WIDTH, 1)), - ) - _, tv_layout_dbias_write = cute.make_layout_tv( - thr_layout=cute.make_layout( - (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE), - stride=(self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, 1), - ), - val_layout=cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) - ), - ) - sDbias_write = cute.composition(sDbias, tv_layout_dbias_write) - # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks - # to avoid bank conflict. - bank_group = (tidx % THREADS_PER_WARP) // self._THREADS_PER_BANK - # The offset this thread should start reading from based on what's its first bank to access. - offset = bank_group * self._PACK_SIZE - for w in cutlass.range_constexpr( - self._WAVES - ): # Each thread starts from this offset when writing into SMEM to avoid bank conflict - start = (w * self._PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE - for i in cutlass.range_constexpr(self._PACK_SIZE): - # All threads write their per-thread partial sum results to the shared buffer. - sDbias_write[(tidx, start + i)] = rowwise_dbias_acc[w * self._PACK_SIZE + i] - cute.arch.sync_threads() - # All threads reduce the cross-thread partial sums to the per-block partial sum. - _, tv_layout_dbias_reduce = cute.make_layout_tv( - thr_layout=cute.make_layout((1, self._TILE_COLS), stride=(self._TILE_COLS, 1)), - val_layout=cute.make_layout((self._TILE_ROWS, 1), stride=(1, 1)), - ) - sDbias_reduce = cute.composition(sDbias, tv_layout_dbias_reduce) - # make_layout_tv yields a (thread, value) layout: thread=tidx -> column tidx, - # value=i -> row i. So index [tidx, i] (thread first), summing the column's rows. - block_dbias = Float32(0.0) - for i in cutlass.range_constexpr(self._TILE_ROWS): - block_dbias += sDbias_reduce[tidx, i] - return block_dbias - - @cute.jit - def _amax_epilogue( - self, - sAmax: cute.Tensor, - mAmax: cute.Tensor, - tidx: Int32, - warp_idx: Int32, - per_thread_amax: Float32, - ): - # Reduce and get the per-warp amax. - warp_amax = cute.arch.warp_redux_sync(per_thread_amax, kind="fmax") - # Write the per-warp amax to shared memory - lane_idx = tidx % 32 - if lane_idx == 0: - sAmax[warp_idx] = warp_amax - cute.arch.sync_threads() - if tidx == 0: - cta_amax = Float32(0.0) - # The first thread reduces all the per-warp amax to the per-CTA amax - for w in cutlass.range_constexpr(self._NUM_WARPS): - cta_amax = cute.arch.fmax(cta_amax, sAmax[w]) - amax_i32 = cute.make_tensor( - cute.recast_ptr(mAmax.iterator, dtype=Int32), - cute.make_layout(1), - ) - # The first thread updates the global amax with an atomic max on the bitcasted float value - cute.arch.atomic_max( - amax_i32.iterator, - cta_amax.bitcast(Int32), - ) - - @cute.jit - def _process_rowwise( - self, - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA - sO_row_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) - mS_row_stage: cute.Tensor, # rowwise scale tensor (1D swizzled, or 2D linear) - tile_row_start: Int32, # global row of this stage's row 0 - tile_col_start: Int32, # global col of this CTA's col 0 - M: Int32, - N: Int32, # full input extents, for OOB masking - sActInput_tile: Optional[cute.Tensor] = None, # (TILE_Y, TILE_X) act_input tile (dact only) - dbias_acc: Optional[ - cute.Tensor - ] = None, # rmem Float32[32] dbias accumulator (rowwise-only dbias) - ): - cfg = self.cfg - return quantize_rowwise_mxfp8( - sX_tile, - None if self.CACHE_ACTIVATION else sActInput_tile, - sO_row_tile, - mS_row_stage, - self.cfg.MAX_NORM_RCP, - tile_row_start, - tile_col_start, - M, - N, - ACTIVATION=None if self.CACHE_ACTIVATION else cfg.ACTIVATION, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - TILE_X=self._TILE_COLS, - TILE_Y=self._TILE_ROWS, - WAVES=self._WAVES, - THREADS_PER_BANK=self._THREADS_PER_BANK, - PACK_SIZE=self._PACK_SIZE, - SKIP_INPUT_MASKING=self.SKIP_INPUT_MASKING, - SKIP_SCALE_BOUNDS=self.SKIP_SCALE_BOUNDS, - WITH_ACT=cfg.WITH_ACT and not self.CACHE_ACTIVATION, - WITH_DACT=cfg.WITH_DACT and not self.CACHE_ACTIVATION, - WITH_DBIAS=self.DBIAS_REDUCTION_IN_ROWWISE, - dbias_acc=dbias_acc, - ) - - @cute.jit - def _process_colwise( - self, - sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA - sO_col_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) - mS_col_stage: cute.Tensor, # colwise scale tensor (1D swizzled, or 2D linear) - tile_row_start: Int32, # global row of this stage's row 0 - tile_col_start: Int32, # global col of this CTA's col 0 - M: Int32, - N: Int32, # full input extents, for OOB masking - sActInput_tile: Optional[cute.Tensor] = None, # (TILE_Y, TILE_X) act_input tile (dact only) - ): - cfg = self.cfg - return quantize_colwise_mxfp8( - sX_tile, - sActInput_tile, - sO_col_tile, - mS_col_stage, - self.cfg.MAX_NORM_RCP, - tile_row_start, - tile_col_start, - M, - N, - ACTIVATION=cfg.ACTIVATION, - DTYPE=cfg.DTYPE, - FP8_DTYPE=cfg.FP8_DTYPE, - SWIZZLE=cfg.WITH_GEMM_SWIZZLED_SCALES, - TILE_X=self._TILE_COLS, - TILE_Y=self._TILE_ROWS, - SKIP_INPUT_MASKING=self.SKIP_INPUT_MASKING, - SKIP_SCALE_BOUNDS=self.SKIP_SCALE_BOUNDS, - WITH_ACT=cfg.WITH_ACT, - WITH_DACT=cfg.WITH_DACT, - WITH_DBIAS=self.DBIAS_REDUCTION_IN_COLWISE, - CACHE_ACTIVATION=self.CACHE_ACTIVATION, - ) - - -class MXFP8QuantizeSpecializedRowwiseKernel(MXFP8QuantizeKernelBase): - """Specialized cast-only ROWWISE-only MXFP8 kernel. - - Requires N % 128 == 0 (full vectorizable column chunks). - - Plain rowwise-only quantize. Each thread owns one 32-element MXFP8 chunk and - uses vectorized global loads/stores (no TMA used).""" - - _TILE_ROWS = 4 - _TILE_COLS = 1024 - _THREADS_PER_CTA = 128 - - def __init__(self, cfg: MXFP8QuantizeConfig): - self.cfg = cfg - # If True, then this kernel will first write each thread's scale byte to a shared memory buffer, - # then utilize vectorized store to flush the buffer to global memory. - self._STASH_SCALE_TO_SMEM = True # Hardcode to true for now - - @cute.jit - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], # Unused, kept for API compatibility - mAmax: Optional[cute.Tensor], # Unused, kept for API compatibility - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device - mDActInput: Optional[cute.Tensor], # Unused, kept for API compatibility - mWorkspace: Optional[cute.Tensor], # Unused, kept for API compatibility - stream: CUstream, - ): - if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf( - "[CuTeDSL] MXFP8QuantizeSpecializedRowwiseKernel.__call__() with config:" - f" {self.cfg}\n" - ) - - M = mX.shape[0] - N = mX.shape[1] - - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_row, _ = derive_swizzled_scale_layout(M, N, True, False, mS_row, None) - - grid = [ - cute.ceil_div(Int32(N), self._TILE_COLS), - cute.ceil_div(M, self._TILE_ROWS), - ] - block = [self._THREADS_PER_CTA] - - self.kernel( - mX, - mO_row, - mS_row, - mNoop, - mX.element_type, - ).launch(grid=grid, block=block, stream=stream) - - @cute.kernel - def kernel( - self, - mX: cute.Tensor, - mO_row: cute.Tensor, - mS_row: cute.Tensor, - mNoop: cute.Pointer, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main(mX, mO_row, mS_row, DTYPE) - - @cute.jit - def _kernel_main( - self, - mX: cute.Tensor, - mO_row: cute.Tensor, - mS_row: cute.Tensor, - DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], - ): - """Device entry for the specialized rowwise-only cast kernel (vectorized global loads/stores, no TMA).""" - tidx, _, _ = cute.arch.thread_idx() - bidx, bidy, _ = cute.arch.block_idx() - M = mX.shape[0] - N = mX.shape[1] - - # Each thread handles one 32-element MXFP8 chunk (= one scale block). - # The 128 threads in the CTA are grouped as (4, 32), so they cover a - # (4, 1024) input tile and the matching (4, 32) scale tile. - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - tiler, tv_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout( - (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) - ), - ) - tiler_scale, tv_layout_scale = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), - val_layout=cute.make_layout((1, 1), stride=(1, 1)), - ) - - # Select the tile that belongs to this CTA, then the fragment per thread. - mX_tile = cute.local_tile(mX, tiler, (bidy, bidx)) - mO_tile = cute.local_tile(mO_row, tiler, (bidy, bidx)) - mS_tile = cute.local_tile(mS_row, tiler_scale, (bidy, bidx)) - mX_thread = cute.composition(mX_tile, tv_layout)[tidx, None] - mO_thread = cute.composition(mO_tile, tv_layout)[tidx, None] - mS_thread = cute.composition(mS_tile, tv_layout_scale)[tidx, None] - - rX_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=DTYPE, - ) - abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 - x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 - x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 - mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, self.cfg.FP8_DTYPE) - rX_i32 = cute.make_tensor( - cute.recast_ptr(rX_thread.iterator, dtype=Int32), - cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), - ) - rO_thread = cute.make_rmem_tensor( - cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), - dtype=self.cfg.FP8_DTYPE, - ) - rO_u32 = cute.make_tensor( - cute.recast_ptr(rO_thread.iterator, dtype=Uint32), - cute.make_layout( - (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) - ), # Unit is Uint32, divide by 4 here - ) - - sS_thread = None - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - - @cute.struct - class SharedStorage: - buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] - - storage = cutlass.utils.SmemAllocator().allocate(SharedStorage) - sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) - # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale - # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. - sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] - # Zero first so padding columns (cols past N/32 in the padded scale - # matrix) flush as 0 and we never read uninitialized smem. - # Raw 0x00: e8m0 has no zero value to convert from. - sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) - cute.arch.sync_threads() - - row = bidy * self._TILE_ROWS + tidx // CTA_X - col = bidx * self._TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE - if row < M and col < N: - cute.autovec_copy(mX_thread, rX_thread) - amax_2x = rX_i32[0] - # Skip the first iteration because we assigned rX_i32[0] to amax_2x already - for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): - amax_2x = abs_max_x2(amax_2x, rX_i32[i]) - amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) - - biased_exp = cvt_f32_to_fp8e8m0fnu(amax * self.cfg.MAX_NORM_RCP) - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - sS_thread[0] = biased_exp - else: - mS_thread[0] = biased_exp - - # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), - # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. - inv_scale = exp2f_rcp(biased_exp) - scale_2x = pack_f32x2(inv_scale, inv_scale) - for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): - rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) - cute.autovec_copy(rO_thread, mO_thread) - - # Cooperative wide flush of the staged scales where padding columns flush as 0. - if cutlass.const_expr(self._STASH_SCALE_TO_SMEM): - cute.arch.sync_threads() - # Use cute.size instead of .shape[1] because under the swizzled layout mode 1 is - # a nested tuple ((4, num_tiles_SC)), not a plain int. - padded_cols = cute.size(mS_row, mode=[1]) - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) - # so we have at most 4 elements continuous which is our vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - elif padded_cols % 16 == 0: - # If columns is divisible by 16, use 16 bytes as the vectorized store width - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 16) - else: - # Otherwise use 4 bytes as the vectorized store width. - # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here - self._flush_scales_to_gmem(sScale, mS_tile, tidx, bidx, bidy, M, padded_cols, 4) - - @cute.jit - def _flush_scales_to_gmem( - self, - sScale: cute.Tensor, - mS_tile: cute.Tensor, - tidx: Int32, - bidx: Int32, - bidy: Int32, - M: Int32, - padded_cols: Int32, - width: cutlass.Constexpr[int], - ): - """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" - CTA_Y = self._TILE_ROWS - CTA_X = self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - # Previously each threads has 1 byte, but now we are doing vectorized store, - # which means only a subset of threads will need to issue the store while other threads are not used. - active_threads = CTA_X // width - _, tv_flush = cute.make_layout_tv( - thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), - val_layout=cute.make_layout((1, width), stride=(width, 1)), - ) - # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write - # so if the thread is outside of this subset, it will remain inactive - if tidx < CTA_Y * active_threads: - # Absolute position of the scale vector to write in the GMEM buffer - thread_y = bidy * CTA_Y + tidx // active_threads - thread_x = bidx * CTA_X + (tidx % active_threads) * width - if thread_y < M and thread_x < padded_cols: - cute.autovec_copy( - cute.composition(sScale, tv_flush)[tidx, None], - cute.composition(mS_tile, tv_flush)[tidx, None], - ) - - -class MXFP8QuantizeSpecializedBidimensionalKernel(MXFP8QuantizeKernelBase): - """Specialized cast-only rowwise+colwise MXFP8 kernel (swizzled TMA, one 32x32 tile per warp).""" - - _WARPS_PER_CTA = 2 - _THREADS_PER_CTA = _WARPS_PER_CTA * THREADS_PER_WARP - # A warp handles a 32x32 tile - _WARP_ROWS = MXFP8_BLOCK_SCALING_SIZE - _WARP_COLS = MXFP8_BLOCK_SCALING_SIZE - # A CTA handles a tile consisted two 32x32 warp subtiles side-by-side at a time, each handled by a warp - _TILE_ROWS = _WARP_ROWS - _TILE_COLS = _WARP_COLS * _WARPS_PER_CTA - _NUM_STAGES = 2 - _NUM_TILES_X = 4 - _NUM_TILES_Y = 1 - _NUM_TILES = _NUM_TILES_X * _NUM_TILES_Y - - # Rows and columns for the rowwise / columnwise scale factor tensor - _SCALE_ROWS = _TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE - _SCALE_COLS = _TILE_COLS // MXFP8_BLOCK_SCALING_SIZE - - def __init__(self, cfg: MXFP8QuantizeConfig): - self.cfg = cfg - if cfg.WITH_GEMM_SWIZZLED_SCALES: - self._NUM_TILES_X = 2 - self._NUM_TILES_Y = 2 - self._NUM_TILES = self._NUM_TILES_X * self._NUM_TILES_Y - - @cute.jit - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device - mDActInput: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - stream: CUstream, - ): - if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): - cute.printf( - "[CuTeDSL] MXFP8QuantizeSpecializedBidimensionalKernel.__call__() with config:" - f" {self.cfg}\n" - ) - M = mX.shape[0] - N = mX.shape[1] - - if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): - mS_row, mS_col = derive_swizzled_scale_layout( - M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col - ) - - smem_tile_layout = cute.make_ordered_layout( - (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) - ) - cta_tiler = (self._TILE_ROWS, self._TILE_COLS) - # Apply 128B input Swizzle<3,4,3> to input tiles - in_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, smem_tile_layout) - # Apply 64B output swizzle<2,4,3> for output tiles - out_smem_layout = cute.make_composed_layout(cute.make_swizzle(2, 4, 3), 0, smem_tile_layout) - op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() - op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() - - # Input TMA atom - tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_load, - mX, - in_smem_layout, - cta_tiler, - num_multicast=1, - ) - # Rowwise output TMA atoms - tma_atom_out_row, tma_dst_out_row = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_store, - mO_row, - out_smem_layout, - cta_tiler, - num_multicast=1, - ) - # Colwise output TMA atoms - tma_atom_out_col, tma_dst_out_col = cute.nvgpu.cpasync.make_tiled_tma_atom( - op_store, - mO_col, - out_smem_layout, - cta_tiler, - num_multicast=1, - ) - - grid = [ - # Each CTA covers a (TILE_ROWS * NUM_TILES_Y, TILE_COLS * NUM_TILES_X) area. - cute.ceil_div(Int32(N), self._TILE_COLS * self._NUM_TILES_X), - cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES_Y), - ] - block = [self._THREADS_PER_CTA] - self.kernel( - mX, - mS_row, - mS_col, - mNoop, - mX.element_type, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ).launch(grid=grid, block=block, stream=stream) - - @cute.kernel - def kernel( - self, - mX: cute.Tensor, - mS_row: cute.Tensor, - mS_col: cute.Tensor, - mNoop: cute.Pointer, - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, - tma_atom_out_row: cute.CopyAtom, - tma_dst_out_row: cute.Tensor, - tma_atom_out_col: cute.CopyAtom, - tma_dst_out_col: cute.Tensor, - ): - """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" - if not noop_flag_is_set(mNoop): - self._kernel_main( - mX, - mS_row, - mS_col, - dtype, - tma_atom, - tma_src, - tma_atom_out_row, - tma_dst_out_row, - tma_atom_out_col, - tma_dst_out_col, - ) - - @cute.jit - def _kernel_main( - self, - mX: cute.Tensor, - mS_row: cute.Tensor, - mS_col: cute.Tensor, - dtype: cutlass.Constexpr[Type[cutlass.Numeric]], - tma_atom: cute.CopyAtom, - tma_src: cute.Tensor, - tma_atom_out_row: cute.CopyAtom, - tma_dst_out_row: cute.Tensor, - tma_atom_out_col: cute.CopyAtom, - tma_dst_out_col: cute.Tensor, - ): - """Device entry for the specialized bidimensional (rowwise+colwise) cast kernel.""" - - bidx, bidy, _ = cute.arch.block_idx() - - M = mX.shape[0] - N = mX.shape[1] - - # The CTA owns a NUM_TILES_Y x NUM_TILES_X arrangement of tiles, starting here. - tile_x_base = bidx * self._NUM_TILES_X - tile_y_base = bidy * self._NUM_TILES_Y - - num_tiles_x = cutlass.min( - self._NUM_TILES_X, - # Valid 64-col tiles remaining from this CTA's start column (bidx * span). - cute.ceil_div(Int32(N) - bidx * self._TILE_COLS * self._NUM_TILES_X, self._TILE_COLS), - ) - num_tiles_y = cutlass.min( - self._NUM_TILES_Y, - # Valid 32-row tiles remaining from this CTA's start row (bidy * span). - cute.ceil_div(Int32(M) - bidy * self._TILE_ROWS * self._NUM_TILES_Y, self._TILE_ROWS), - ) - num_tiles = num_tiles_x * num_tiles_y - - FP8_DTYPE = self.cfg.FP8_DTYPE - - @cute.struct - class SharedStorage: - mbar_storage: cute.struct.MemRange[ - cute.Int64, 2 * self._NUM_STAGES - ] # (full,empty) per stage - sX: cute.struct.Align[ - cute.struct.MemRange[dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES], - 128, - ] - sO_row: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - sO_col: cute.struct.Align[ - cute.struct.MemRange[ - FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES - ], - 128, - ] - # Per-warp colwise-reduce scratchpad - sColReduce: cute.struct.Align[ - cute.struct.MemRange[Float32, THREADS_PER_WARP * self._WARPS_PER_CTA], 16 - ] - # Staged rowwise scales for the whole CTA - sScaleRow: cute.struct.Align[ - cute.struct.MemRange[ - Float8E8M0FNU, - self._TILE_ROWS * self._NUM_TILES * self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, - ], - 16, - ] - # Staged colwise scales for the whole CTA - sScaleCol: cute.struct.Align[ - cute.struct.MemRange[Float8E8M0FNU, self._NUM_TILES * self._TILE_COLS], 16 - ] - - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - - tile_layout = cute.make_layout( - ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), - stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), - ) - # SMEM input tile should have the same Swizzle<3,4,3> as the input TMA atom - sX = storage.sX.get_tensor(tile_layout, swizzle=cute.make_swizzle(3, 4, 3)) - # SMEM output tiles should have the same Swizzle<2,4,3> as the output TMA atom - sO_row = storage.sO_row.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) - sO_col = storage.sO_col.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) - # Reshape the per-warp colwise-reduce scratchpad to a (threads, warps) layout for easier indexing. - sColReduce = storage.sColReduce.get_tensor( - cute.make_layout((THREADS_PER_WARP, self._WARPS_PER_CTA), stride=(1, THREADS_PER_WARP)) - ) - # The whole rowwise scale factor SMEM tensor covered by a CTA - sScaleRow2D = storage.sScaleRow.get_tensor( - cute.make_ordered_layout( - ( - self._TILE_ROWS * self._NUM_TILES_Y, - self._SCALE_COLS * self._NUM_TILES_X, - ), - order=(1, 0), - ) - ) - # Divide by a single tile's shape, so it becomes ((TILE_ROWS, SCALE_COLS), TILES) for easier indexing - # where TILES = NUM_TILES_X * NUM_TILES_Y - sScaleRow = cute.zipped_divide(sScaleRow2D, (self._TILE_ROWS, self._SCALE_COLS)) - # The whole colwise scale factor SMEM tensor covered by a CTA - sScaleCol2D = storage.sScaleCol.get_tensor( - cute.make_ordered_layout( - ( - self._SCALE_ROWS * self._NUM_TILES_Y, - self._TILE_COLS * self._NUM_TILES_X, - ), - order=(1, 0), - ) - ) - # Divide by a single tile's shape, so it becomes ((SCALE_ROWS, TILE_COLS), TILES) for easier indexing - # where TILES = NUM_TILES_X * NUM_TILES_Y - sScaleCol = cute.zipped_divide(sScaleCol2D, (self._SCALE_ROWS, self._TILE_COLS)) - - # Zero the scale SMEM buffers so partial tiles / padding columns flush as 0. - tidx, _, _ = cute.arch.thread_idx() - # View rowwise staging buffers as flat uint32 and stride the CTA over them. - _ROW_SCALE_WORDS = self._TILE_ROWS * self._SCALE_COLS * self._NUM_TILES // 4 - sScaleRow_u32 = cute.make_tensor( - cute.recast_ptr(sScaleRow2D.iterator, dtype=Uint32), - cute.make_layout((_ROW_SCALE_WORDS,), stride=(1,)), - ) - for i in cutlass.range_constexpr(cute.ceil_div(_ROW_SCALE_WORDS, self._THREADS_PER_CTA)): - slot = i * self._THREADS_PER_CTA + tidx - if slot < _ROW_SCALE_WORDS: - sScaleRow_u32[slot] = Uint32(0) - # View colwise staging buffers as flat uint32 and stride the CTA over them. - _COL_SCALE_WORDS = self._SCALE_ROWS * self._TILE_COLS * self._NUM_TILES // 4 - sScaleCol_u32 = cute.make_tensor( - cute.recast_ptr(sScaleCol2D.iterator, dtype=Uint32), - cute.make_layout((_COL_SCALE_WORDS,), stride=(1,)), - ) - for i in cutlass.range_constexpr(cute.ceil_div(_COL_SCALE_WORDS, self._THREADS_PER_CTA)): - slot = i * self._THREADS_PER_CTA + tidx - if slot < _COL_SCALE_WORDS: - sScaleCol_u32[slot] = Uint32(0) - - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - if warp_idx == 0: - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) - - # Only warp 0 is the producer (issues TMA) - producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) - # Every warp is the consumer - consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._WARPS_PER_CTA) - # One TMA loads a tile - tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 - - mainloop_pipeline = pipeline.PipelineTmaAsync.create( - barrier_storage=storage.mbar_storage.data_ptr(), - num_stages=self._NUM_STAGES, - producer_group=producer_group, - consumer_group=consumer_group, - tx_count=tx_count, - cta_layout_vmnk=None, - ) - - prod_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self._NUM_STAGES - ) - cons_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self._NUM_STAGES - ) - - gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) - tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( - tma_atom, 0, cute.make_layout(1), sX, gX_tiled - ) - gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self._TILE_ROWS, self._TILE_COLS)) - tXsO_row, tXgO_row = cute.nvgpu.cpasync.tma_partition( - tma_atom_out_row, 0, cute.make_layout(1), sO_row, gO_row_tiled - ) - gO_col_tiled = cute.zipped_divide(tma_dst_out_col, (self._TILE_ROWS, self._TILE_COLS)) - tXsO_col, tXgO_col = cute.nvgpu.cpasync.tma_partition( - tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled - ) - - cute.arch.sync_threads() - - # Prologue: warp 0 prefetches up to NUM_STAGES tiles to fully fill the pipeline - if warp_idx == 0: - for s in cutlass.range_constexpr(self._NUM_STAGES): - if s < num_tiles: - # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. - tile_y = s // num_tiles_x - tile_x = s % num_tiles_x - mainloop_pipeline.producer_acquire(prod_state) - cute.copy( - tma_atom, - tXgX[(None, (tile_y_base + tile_y, tile_x_base + tile_x))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - - # 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): - mainloop_pipeline.consumer_wait(cons_state) - # Only allow at most _NUM_STAGES-1 stages to be in-flight, because this iteration will reuse the ring buffer - # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish - if warp_idx == 0: - cute.arch.cp_async_bulk_wait_group(self._NUM_STAGES - 1, read=True) - cute.arch.sync_threads() - # The current pipeline stage index, which is the tile index modulo the number of stages. - # This is used to index into the shared memory ring buffers that are wrapped around the number of stages. - stage_idx = cons_state.index - - # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. - tile_y = tile_idx // num_tiles_x - tile_x = tile_idx % num_tiles_x - - # Process the 32x64 SMEM tile for this stage - quantize_bidimensional_mxfp8_swizzled( - sX[None, stage_idx], - sO_row[None, stage_idx], - sO_col[None, stage_idx], - sScaleRow[(None, (tile_y, tile_x))], - sScaleCol[(None, (tile_y, tile_x))], - sColReduce[ - None, warp_idx - ], # Pick the per-warp colwise-reduce scratchpad for this warp - self._WARPS_PER_CTA, - self.cfg.MAX_NORM_RCP, - dtype, - self.cfg.FP8_DTYPE, - ) - - # Make the smem output writes visible to the TMA async proxy, then store. - cute.arch.fence_proxy("async.shared", space="cta") - cute.arch.sync_threads() - - # We are done with this input pipeline SMEM buffer, signal the producer that it can write to this buffer - mainloop_pipeline.consumer_release(cons_state) - - # Issue TMA and write the output to GMEM - if warp_idx == 0: - tile_coord = (tile_y_base + tile_y, tile_x_base + tile_x) - cute.copy( - tma_atom_out_row, tXsO_row[(None, stage_idx)], tXgO_row[(None, tile_coord)] - ) - cute.copy( - tma_atom_out_col, tXsO_col[(None, stage_idx)], tXgO_col[(None, tile_coord)] - ) - cute.arch.cp_async_bulk_commit_group() - - cons_state.advance() - - # Producer: refill the buffer the consumer just freed with the tile - # NUM_STAGES ahead. - next_tile_idx = tile_idx + self._NUM_STAGES - if next_tile_idx < num_tiles: - if warp_idx == 0: - next_tile_y = next_tile_idx // num_tiles_x - next_tile_x = next_tile_idx % num_tiles_x - mainloop_pipeline.producer_acquire(prod_state) - cute.copy( - tma_atom, - tXgX[(None, (tile_y_base + next_tile_y, tile_x_base + next_tile_x))], - tXsX[(None, prod_state.index)], - tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), - ) - mainloop_pipeline.producer_commit(prod_state) - prod_state.advance() - - # TODO(kainingz): the rowwise WIDTH is hardcoded to 4. Consider using other width if we can convince - # cute.autovec_copy that wider stores are safe (e.g. 16B) and the runtime row pitch allows it. This might - # require us to change `scale_rowwise_shape`'s divisibility when we compile with fake tensors - self._flush_scales_to_gmem( - sScaleRow2D, - mS_row, - tidx, - bidx, - bidy, - ROWS=self._TILE_ROWS * self._NUM_TILES_Y, - COLS=self._SCALE_COLS * self._NUM_TILES_X, - WIDTH=4, - ) - self._flush_scales_to_gmem( - sScaleCol2D, - mS_col, - tidx, - bidx, - bidy, - ROWS=self._SCALE_ROWS * self._NUM_TILES_Y, - COLS=self._TILE_COLS * self._NUM_TILES_X, - WIDTH=1 if self.cfg.WITH_GEMM_SWIZZLED_SCALES else 16, - ) - - # Drain the final stores (their gmem writes must complete before the CTA exits). - cute.arch.cp_async_bulk_wait_group(0, read=False) - - @cute.jit - def _flush_scales_to_gmem( - self, - sScale2D: cute.Tensor, - mS: cute.Tensor, - tidx: Int32, - bidx: Int32, - bidy: Int32, - ROWS: cutlass.Constexpr[int], - COLS: cutlass.Constexpr[int], - WIDTH: cutlass.Constexpr[int], - ): - """Flush a staged (ROWS, COLS) SMEM scale block (a plain row-major 2D tensor) to its (bidy, bidx) slice of the gmem scale tensor, - where each GMEM slice has the same shape as this SMEM tile. Use `WIDTH` bytes per vectorized store. - """ - # Use cute.size instead of cute.shape because under the swizzled layout each mode is a nested tuple ((32, 4, num_tiles_M) etc.) - mS_M = cute.size(mS, mode=[0]) - mS_N = cute.size(mS, mode=[1]) - # Obtain the GMEM slice for the output scale factor block of this CTA - mS_tile = cute.local_tile(mS, (ROWS, COLS), (bidy, bidx)) - - ACTIVE_THREAD_COLS = COLS // WIDTH - _, tv_flush_layout = cute.make_layout_tv( - thr_layout=cute.make_layout((ROWS, ACTIVE_THREAD_COLS), stride=(ACTIVE_THREAD_COLS, 1)), - val_layout=cute.make_layout((1, WIDTH), stride=(WIDTH, 1)), - ) - TOTAL_ACTIVE = ROWS * ACTIVE_THREAD_COLS - - # We may have more total active threads than threads in a CTA, so we do multiple waves of stores to flush the whole buffer - for wave in cutlass.range_constexpr(cute.ceil_div(TOTAL_ACTIVE, self._THREADS_PER_CTA)): - thread_idx = wave * self._THREADS_PER_CTA + tidx - if thread_idx < TOTAL_ACTIVE: - # Find the position for this slot's vectorized store in the GMEM scale factor buffer - thread_y = bidy * ROWS + thread_idx // ACTIVE_THREAD_COLS - thread_x = bidx * COLS + (thread_idx % ACTIVE_THREAD_COLS) * WIDTH - # For rowwise we have N divisible by 4 and WIDTH=4, and for colwise we have N divisible by 128 and WIDTH=16, - # so `thread_x < mS_N` with vectorized store is safe here. - # A thread only writes to a single row so `thread_y < mS_M` is also safe here. - if thread_y < mS_M and thread_x < mS_N: - cute.autovec_copy( - cute.composition(sScale2D, tv_flush_layout)[thread_idx, None], - cute.composition(mS_tile, tv_flush_layout)[thread_idx, None], - ) - - -class MXFP8QuantizeEntry(MXFP8QuantizeKernelBase): - """Select the appropriate MXFP8 quantization kernel based on the configuration and runtime shapes.""" - - def __init__(self, cfg: MXFP8QuantizeConfig): - self.cfg = cfg - # These activation functions satisfy f(0) = 0 so we don't need to mask with OOB regions - # (zeros filled by TMA are still zeros without applying activation to them) - # Note that as the time of writing, all activations supported by MXFP8 satisfy this property, so !ACT_NEED_MASKING is never reachable. - # This is only a precaution for future use cases if someone wants to add a new activation function. - self.ACT_NEED_MASKING = cfg.WITH_ACT and cfg.ACTIVATION not in ( - "relu", - "gelu", - "silu", - "qgelu", - "srelu", - ) - # Instantiate all possible kernels at compile time, - # and we will pick the right one at runtime based on the input shape and config. - self.general_skip_input_masking_kernel = MXFP8QuantizeKernel(cfg, SKIP_INPUT_MASKING=True) - self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeKernel( - cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True - ) - # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics - self.general_kernel = MXFP8QuantizeKernel(cfg) if self.ACT_NEED_MASKING else None - self.specialized_rowwise = ( - MXFP8QuantizeSpecializedRowwiseKernel(cfg) if cfg.ROWWISE else None - ) - self.specialized_bidim = ( - MXFP8QuantizeSpecializedBidimensionalKernel(cfg) - if cfg.ROWWISE and cfg.COLWISE - else None - ) - - @cute.jit - def __call__( - self, - mX: cute.Tensor, - mO_row: Optional[cute.Tensor], - mS_row: Optional[cute.Tensor], - mO_col: Optional[cute.Tensor], - mS_col: Optional[cute.Tensor], - mAmax: Optional[cute.Tensor], - mNoop: cute.Pointer, - mDActInput: Optional[cute.Tensor], - mWorkspace: Optional[cute.Tensor], - stream: CUstream, - ): - M = mX.shape[0] - N = mX.shape[1] - # Select specialized kernels if possible - plain_cast_only = ( - not self.cfg.WITH_AMAX - and not self.cfg.WITH_DBIAS - and not self.cfg.WITH_DACT - and not self.cfg.WITH_ACT - ) - dispatched_to_specialized = False - # Only dispatch to the specialized kernels for packed16 types (bf16/fp16) - if cutlass.const_expr(plain_cast_only and is_packed16(self.cfg.DTYPE)): - if cutlass.const_expr(self.cfg.ROWWISE and not self.cfg.COLWISE): - rowwise_specialized_grid_fits = ( - cute.ceil_div(M, self.specialized_rowwise._TILE_ROWS) <= 65535 - ) - # The rowwise specialized kernel requires N divisible by 128 for vectorized stores - if N % 128 == 0 and rowwise_specialized_grid_fits: - dispatched_to_specialized = True - self.specialized_rowwise( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # The bidimensional specialized kernel supports swizzled scales but we don't dispatch to it for now - if cutlass.const_expr( - self.cfg.ROWWISE and self.cfg.COLWISE and not self.cfg.WITH_GEMM_SWIZZLED_SCALES - ): - bidim_specialized_grid_fits = ( - cute.ceil_div( - M, - self.specialized_bidim._TILE_ROWS * self.specialized_bidim._NUM_TILES_Y, - ) - <= 65535 - ) - if bidim_specialized_grid_fits: - dispatched_to_specialized = True - self.specialized_bidim( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - # If not using a specialized kernel, fall back to the general kernel - if not dispatched_to_specialized: - # If the input shape can be perfectly tiled by the general kernel's tile size, we can skip some boundary check because - # we know we will not touch any out-of-bounds region. - shape_is_divisible = ( - mX.shape[0] % self.general_skip_input_masking_kernel._TILE_ROWS == 0 - and mX.shape[1] % self.general_skip_input_masking_kernel._TILE_COLS == 0 - ) - # If shape is divisible, then we won't access OOB regions regardless whatever - if shape_is_divisible: - self.general_skip_input_masking_and_scale_bounds_kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - else: - # Otherwise we can't skip scale bound check, but we may still skip masking activation inputs with zero - if cutlass.const_expr(self.ACT_NEED_MASKING): - # Masking and the scale guards now share one condition, so this is a plain - # two-way choice: either the shape tiles exactly (drop both) or it does not. - self.general_kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - else: - # We still need to check the scale bounds, but we can skip some activations masking because their output is 0 - # for OOB regions where TMA fills with zeros - self.general_skip_input_masking_kernel( - mX, - mO_row, - mS_row, - mO_col, - mS_col, - mAmax, - mNoop, - mDActInput, - mWorkspace, - stream, - ) - - -def compile_cutedsl_function_from_cfg(cfg): - """ - Return the compiled CuTeDSL function object for the given MXFP8 quantization config. - """ - - kernel_obj = MXFP8QuantizeEntry(cfg) - sym_M = cute.sym_int32() - sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) - in_shape = out_shape = (sym_M, sym_N) - # TE allocates scale tensors at a padded shape (see - # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): - # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) - # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) - # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the - # scales their own fresh syms carrying the divisibility the padding - # guarantees (rowwise: 128 x 4; colwise: 4 x 128). - scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) - scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) - ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N - # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ - # bridge sends). - out_dtype = cfg.FP8_DTYPE - scale_dtype = cutlass.Float8E8M0FNU - - in_fake = cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, in_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=16 - ) - out_row_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.ROWWISE - else None - ) - scale_row_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_rowwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.ROWWISE - else None - ) - out_col_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.COLWISE - else None - ) - scale_col_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_colwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.COLWISE - else None - ) - amax_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_AMAX - else None - ) - # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that - # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). - # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher - # passes at launch. - noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) - act_input_fake = ( - cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, - in_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.WITH_DACT - else None - ) - workspace_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, ws_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_DBIAS - else None - ) - - compiled = cute.compile( - kernel_obj, - in_fake, # mX - out_row_fake, - scale_row_fake, # mO_row, mS_row - out_col_fake, - scale_col_fake, # mO_col, mS_col - amax_fake, # mAmax - noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) - act_input_fake, # mDActInput (backward slot, unused) - workspace_fake, # mWorkspace(backward slot, unused) - cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi - # "handle" arg; C++ passes the CUDA stream - # as void*) - options="--enable-tvm-ffi", - ) - return compiled - - -def get_mxfp8_quantization_function( - fn_name: str, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - with_gemm_swizzled_scales: bool, - with_amax: bool, - with_dbias: bool, - with_dact: bool, - with_act: bool, - use_2d_quantization: bool, - activation: str, -) -> 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). - Returns True if a kernel is successfully registered under `fn_name` (the C++ side then fetches it with GetGlobal(fn_name)); - False if the config is unsupported, so the caller caches the negative result and falls back to the CUDA C++ kernel. - """ - try: - # Already registered (e.g. by a prior call) -> supported. - if tvm_ffi.get_global_func(fn_name, allow_missing=True) is not None: - return True - - major, minor = device_compute_capability() - if major < 10: - logger.warning( - "CuTeDSL MXFP8 backend requires compute capability >= 10.0 (Blackwell), " - "but detected %d.%d; falling back to the CUDA C++ kernel.", - major, - minor, - ) - return False - - try: - cfg = MXFP8QuantizeConfig( - dtype=dtype, - fp8_dtype=fp8_dtype, - rowwise=rowwise, - colwise=colwise, - with_gemm_swizzled_scales=with_gemm_swizzled_scales, - with_amax=with_amax, - with_dbias=with_dbias, - with_dact=with_dact, - with_act=with_act, - use_2d_quantization=use_2d_quantization, - activation=activation, - ) - except ValueError as e: - # The exception message states exactly why the config is unsupported - # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a - # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. - logger.warning( - "CuTeDSL MXFP8 backend does not support this config, " - "falling back to the CUDA C++ kernel: %s", - e, - ) - return False - - logger.debug("Compiling CuTeDSL MXFP8 quantization kernel for %s", cfg) - compiled = compile_cutedsl_function_from_cfg(cfg) - # The returned compiled object is not neccessarily the compiled function itself. It could be a Python - # wrapper that parses arguments and calls the underlying function. This is needed in the general case - # as TVM-FFI is positional-only, however, for us this is unnecessary launch overhead as we pass the - # arguments in the correct order on the C++ side. - # Relevant issue: https://github.com/NVIDIA/cutlass/issues/3527 - # https://github.com/NVIDIA/cutlass/pull/3589 recommends using an explicit positional only marker (/) - # in the function signature, however it is not merged and it is something easy to miss and silently - # introduce unnecessarry overhead. - native = getattr(compiled, "__tvm_ffi_object__", lambda: None)() - tvm_ffi.register_global_func( - fn_name, native if native is not None else compiled, override=True - ) - return True - except Exception as e: # pylint: disable=broad-exception-caught - logger.error( - "CuTeDSL MXFP8 kernel compilation & registration failed, falling back to the CUDA" - " C++ kernel: %s", - e, - ) - # Unconditionally fallback to CUDA path because we can't tell if this exception is transient or permanent. - return False diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py new file mode 100644 index 00000000000..c2207574aea --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -0,0 +1,878 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared definitions, kernel dispatch and TVM-FFI registration for CuTeDSL MXFP8.""" + +import abc +import logging +import os +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import Boolean, Float32, Int32, Int64 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module +import tvm_ffi + +from transformer_engine.common.CuTeDSL.utils import ( + str_to_cutlass_dtype, + device_compute_capability, +) +from transformer_engine.common.CuTeDSL.activations import ( + act_relu, + act_gelu, + act_silu, + act_qgelu, + act_srelu, + dact_drelu, + dact_dsrelu, + dact_dsilu, + dact_dqgelu, + dact_dgelu, + act_sigmoid, + dact_dsigmoid, + derivative_table_spec, + forward_table_spec, +) + +CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" + +logger = logging.getLogger("transformer_engine.cutedsl.mxfp8") + +# Number of elements per MXFP8 scale block. They will share the same E8M0 scale factor +MXFP8_BLOCK_SCALING_SIZE = 32 +# How many threads are in one warp +THREADS_PER_WARP = 32 + +# FP8E4M3 max representable value +FP8E4M3_MAX_NORM = 448.0 +FP8E4M3_MAX_NORM_RCP = 1.0 / FP8E4M3_MAX_NORM +FP8E5M2_MAX_NORM = 57344.0 +FP8E5M2_MAX_NORM_RCP = 1.0 / FP8E5M2_MAX_NORM + +# If N is not divisible by 16, writing quantized FP8 output with N bytes per row will fail because +# TMA requires 16-byte alignment for each row +SYM_N_DIVISIBILITY = 16 + +SUPPORTED_ACTIVATIONS = { + "relu": act_relu, + "gelu": act_gelu, + "silu": act_silu, + "qgelu": act_qgelu, + "srelu": act_srelu, +} + +SUPPORTED_DACTIVATIONS = { + "drelu": dact_drelu, + "dgelu": dact_dgelu, + "dsilu": dact_dsilu, + "dqgelu": dact_dqgelu, + "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) + + +# Gated activations of gated_mxfp8.cuh, as (ActOP, DActOP) ports of the math.h functions the +# CUDA dispatcher instantiates them with. swiglu and clamped_swiglu backward, and situglu in +# both directions, do not go through this table: like the CUDA kernel they take a special path. +GATED_ACTIVATIONS = { + "geglu": (act_gelu, dact_dgelu), + "qgeglu": (act_qgelu, dact_dqgelu), + "reglu": (act_relu, dact_drelu), + "sreglu": (act_srelu, dact_dsrelu), + "glu": (act_sigmoid, dact_dsigmoid), + "swiglu": (act_silu, None), + "clamped_swiglu": (None, None), + "situglu": (None, None), +} + + +@cute.jit +def derive_swizzled_scale_layout( + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], +): + """Derive the swizzled layout for the rowwise and colwise scale tensors.""" + num_scale_cols = cute.ceil_div(N, MXFP8_BLOCK_SCALING_SIZE) + num_scale_rows = cute.ceil_div(M, MXFP8_BLOCK_SCALING_SIZE) + + num_tiles_M = cute.ceil_div(M, 128) + num_tiles_SC = cute.ceil_div(num_scale_cols, 4) + num_tiles_SR = cute.ceil_div(num_scale_rows, 4) + num_tiles_N = cute.ceil_div(N, 128) + + if cutlass.const_expr(ROWWISE): + mS_row = cute.make_tensor( + mS_row.iterator, + cute.make_layout( + ((32, 4, num_tiles_M), (4, num_tiles_SC)), + stride=((16, 4, num_tiles_SC * 512), (1, 512)), + ), + ) + if cutlass.const_expr(COLWISE): + mS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout( + ((4, num_tiles_SR), (32, 4, num_tiles_N)), + stride=((1, 512), (16, 4, num_tiles_SR * 512)), + ), + ) + return mS_row, mS_col + + +@cute.jit +def noop_flag_is_set(mNoop: cute.Pointer) -> Boolean: + """Whether the cast_noop flag says this quantization is a no-op and must be skipped. + + mNoop is a pointer rather than a tensor so that one compiled kernel serves both a present and + an absent flag, hence the address is checked before it is dereferenced, exactly like the CUDA + C++ kernel's `noop != nullptr && noop[0] == 1.0f`. The two checks cannot be joined with `and`, + which the DSL lowers to a non-short-circuiting op that would load from the null pointer. + """ + flag_is_set = Boolean(False) + if mNoop.toint() != Int64(0): + flag_is_set = cute.make_tensor(mNoop, cute.make_layout((1,)))[0] == Float32(1.0) + return flag_is_set + + +class MXFP8QuantizeConfig: + """Configs for the compiled CuTeDSL kernel. These will be fixed once the kernel is compiled and + they will behave as const expressions. + """ + + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + with_amax: bool, + with_dbias: bool = False, + with_dact: bool = False, + 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") + if dtype is None or dtype not in ("Float32", "Float16", "BFloat16"): + raise ValueError(f"unknown input dtype {dtype!r}; expected Float32|Float16|BFloat16") + self.DTYPE = str_to_cutlass_dtype(dtype) + self.DTYPE_STR = dtype # readable input-dtype token, for __str__ + if fp8_dtype not in ("Float8E4M3", "Float8E5M2"): + raise ValueError( + f"unknown FP8 dtype {fp8_dtype!r}; expected 'Float8E4M3' or 'Float8E5M2'" + ) + self.FP8_DTYPE = str_to_cutlass_dtype(fp8_dtype) + self.FP8_DTYPE_STR = fp8_dtype # readable token, for __str__ + self.ROWWISE = rowwise + self.COLWISE = colwise + if not (rowwise or colwise): + raise ValueError("at least one of rowwise or colwise must be true") + self.WITH_GEMM_SWIZZLED_SCALES = with_gemm_swizzled_scales + self.WITH_AMAX = with_amax + if not with_dact and not with_act: + if activation == "none": + self.ACTIVATION = None + else: + raise ValueError( + "activation must be none when with_dact and with_act are both False" + ) + else: + if with_dact and with_act: + raise ValueError( + "with_dact and with_act cannot be true at the same time since they are used for" + " different paths (bwd vs fwd)" + ) + if with_dact: + if activation in SUPPORTED_DACTIVATIONS: + self.ACTIVATION = activation + else: + raise ValueError( + f"unknown activation {activation!r} for with_dact=True; expected one of" + f" {sorted(SUPPORTED_DACTIVATIONS)}" + ) + elif with_act: + if activation in SUPPORTED_ACTIVATIONS: + self.ACTIVATION = activation + else: + raise ValueError( + f"unknown activation {activation!r} for with_act=True; expected one of" + f" {sorted(SUPPORTED_ACTIVATIONS)}" + ) + 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 + # kernel's COLWISE_SCALING / rowwise dbias branches. + self.WITH_DBIAS = with_dbias + self.MAX_NORM_RCP = ( + FP8E4M3_MAX_NORM_RCP if fp8_dtype == "Float8E4M3" else FP8E5M2_MAX_NORM_RCP + ) + + def __str__(self): + return ( + f"MXFP8QuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + 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"activation_table={self.ACTIVATION_TABLE})" + ) + + __repr__ = __str__ + + def is_plain_cast(self) -> bool: + """Whether cfg is a plain cast: no amax, dbias or (d)activation fused in.""" + return not (self.WITH_AMAX or self.WITH_DBIAS or self.WITH_DACT or self.WITH_ACT) + + def compile(self): + """ + Return the compiled CuTeDSL function object for the given MXFP8 quantization config. + """ + + kernel_obj = MXFP8QuantizeEntry(self) + sym_M = cute.sym_int32() + sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) + in_shape = out_shape = (sym_M, sym_N) + # TE allocates scale tensors at a padded shape (see + # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): + # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) + # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) + # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the + # scales their own fresh syms carrying the divisibility the padding + # guarantees (rowwise: 128 x 4; colwise: 4 x 128). + scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) + scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) + ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N + # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ + # bridge sends). + out_dtype = self.FP8_DTYPE + scale_dtype = cutlass.Float8E8M0FNU + + in_fake = cute.runtime.make_fake_compact_tensor( + self.DTYPE, + in_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + out_row_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.ROWWISE + else None + ) + scale_row_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_rowwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if self.ROWWISE + else None + ) + out_col_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.COLWISE + else None + ) + scale_col_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_colwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if self.COLWISE + else None + ) + amax_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 + ) + if self.WITH_AMAX + else None + ) + # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that + # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). + # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher + # passes at launch. + noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) + act_input_fake = ( + cute.runtime.make_fake_compact_tensor( + self.DTYPE, + in_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.WITH_DACT + else None + ) + # Like the noop flag, the activation table is an always-present pointer, null without a + # table. + act_table_fake = cute.runtime.nullptr( + cutlass.Uint8, mem_space=cute.AddressSpace.gmem, assumed_align=16 + ) + workspace_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, + ws_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if self.WITH_DBIAS + else None + ) + + compiled = cute.compile( + kernel_obj, + in_fake, # mX + out_row_fake, + scale_row_fake, # mO_row, mS_row + out_col_fake, + scale_col_fake, # mO_col, mS_col + amax_fake, # mAmax + noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) + act_input_fake, # mDActInput (activation input, dact only) + workspace_fake, # mWorkspace(backward slot, unused) + act_table_fake, # mActTable (activation table, may be null at launch) + cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi + # "handle" arg; C++ passes the CUDA stream + # as void*) + options="--enable-tvm-ffi", + ) + return compiled + + +def _fake_gmem_tensor(dtype, shape, align): + """Row-major fake gmem tensor for cute.compile.""" + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=tuple(reversed(range(len(shape)))), + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + +class MXFP8GatedQuantizeConfig(MXFP8QuantizeConfig): + """Compile-time config of gated_mxfp8.cuh's port: a plain MXFP8 cast config plus the gated op.""" + + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + is_bwd: bool, + activation: str, + ): + super().__init__( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=False, + activation="none", + ) + if activation not in GATED_ACTIVATIONS: + raise ValueError( + f"unknown gated activation {activation!r}; expected one of" + f" {sorted(GATED_ACTIVATIONS)}" + ) + if is_bwd and with_gemm_swizzled_scales: + # TODO(kainingz): technically this is supported but I need to come up with a less disgusting way + # to implement it in CuTe + raise ValueError("backward gated activations with GEMM-swizzled scales") + self.IS_BWD = is_bwd + self.GATED_ACTIVATION = activation + + def __str__(self): + return ( + f"MXFP8GatedQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " + f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, is_bwd={self.IS_BWD}, " + f"activation={self.GATED_ACTIVATION})" + ) + + __repr__ = __str__ + + def compile(self): + """Return the compiled CuTeDSL function object for this config.""" + # The kernel module imports this one + # pylint: disable-next=import-outside-toplevel + + sym_M = cute.sym_int32() + # The C++ dispatcher requires N % 32 == 0, like the CUDA kernel. + sym_N = cute.sym_int32(divisibility=32) + sym_2N = cute.sym_int32(divisibility=64) + out_cols = sym_2N if self.IS_BWD else sym_N + # Padded scale shapes, see MXFP8QuantizeConfig.compile + scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) + scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) + # Only the activation's own params are runtime arguments; the others are compiled as None + # and the C++ launcher passes None for them. + clamped = self.GATED_ACTIVATION == "clamped_swiglu" + situ = self.GATED_ACTIVATION == "situglu" + params = [Float32(0.0) if clamped else None] * 3 # limit, alpha, glu_linear_offset + params += [Float32(0.0) if situ else None] * 2 # beta1, beta2 + return cute.compile( + MXFP8GatedQuantizeEntry(self), + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_2N), 16), # mGatedInput + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_N), 16) if self.IS_BWD else None, # mGrad + # mO_row, mS_row + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.ROWWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_rowwise_shape, 4) + if self.ROWWISE + else None + ), + # mO_col, mS_col + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.COLWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_colwise_shape, 4) + if self.COLWISE + else None + ), + *params, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + + +class MXFP8GatedQuantizeEntry: + """Select the gated MXFP8 kernel: the register-resident one where it applies, otherwise the + shared-memory port of gated_mxfp8.cuh.""" + + def __init__(self, cfg: MXFP8GatedQuantizeConfig): + # The kernel modules import this one. + # pylint: disable=import-outside-toplevel + from .quantize_mxfp8_gated_g2s import MXFP8GatedQuantizeG2SKernel + from .quantize_mxfp8_gated_g2r_bidim import MXFP8GatedQuantizeG2RKernel + + # pylint: enable=import-outside-toplevel + self.cfg = cfg + self.g2s_kernel = MXFP8GatedQuantizeG2SKernel(cfg) + self.g2r_kernel = ( + MXFP8GatedQuantizeG2RKernel(cfg) if MXFP8GatedQuantizeG2RKernel.supports(cfg) else None + ) + + @cute.jit + def __call__( + self, + mGatedInput: cute.Tensor, + mGrad: Optional[cute.Tensor], + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + stream: CUstream, + ): + args = ( + mGatedInput, + mGrad, + mO_row, + mS_row, + mO_col, + mS_col, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + stream, + ) + if cutlass.const_expr(self.g2r_kernel is not None): + if self.g2r_kernel.eligible_to_dispatch(mGatedInput, mS_col): + self.g2r_kernel(*args) + else: + self.g2s_kernel(*args) + else: + self.g2s_kernel(*args) + + +class MXFP8QuantizeKernelBase(abc.ABC): + """Base class for MXFP8 quantize kernels.""" + + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: # pylint: disable=unused-argument + """Compile time check: whether this specialized kernel handles cfg at all.""" + return False + + @staticmethod + def eligible_to_dispatch( + cfg: MXFP8QuantizeConfig, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + ) -> Boolean: + """Runtime check: whether we can dispatch to the instantiated kernel with this call's shapes and noops.""" + raise NotImplementedError + + @abc.abstractmethod + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + mActTable: cute.Pointer, # activation table; null without one + stream: CUstream, + ): + """ + Compiled kernel entrypoint (decorate with @cute.jit). + All MXFP8 quantize kernels must implement this interface because this is our C++ call site's contract in `quantize_mxfp8_cutedsl.cuh`. + C++ call site will pass the arguments in this exact order via tvm-ffi and the kernel must accept them even if they are not used. + """ + + +class MXFP8QuantizeEntry(MXFP8QuantizeKernelBase): + """Select the appropriate MXFP8 quantization kernel based on the configuration and runtime shapes.""" + + def __init__(self, cfg: MXFP8QuantizeConfig): + # Kernels import shared definitions from this module. Resolve their + # implementations only when constructing the compile-time entry object. + # pylint: disable=import-outside-toplevel + from .quantize_mxfp8_g2s_general import MXFP8QuantizeGeneralG2SKernel + from .quantize_mxfp8_g2r_rowwise_1lane import MXFP8QuantizeRowwise1LaneG2RKernel + from .quantize_mxfp8_g2s_bidim import MXFP8QuantizeBidimensionalG2SKernel + from .quantize_mxfp8_g2r_rowwise_2lane import MXFP8QuantizeRowwise2LaneG2RKernel + from .quantize_mxfp8_g2r_bidim import MXFP8QuantizeBidimensionalG2RKernel + + # pylint: enable=import-outside-toplevel + + self.cfg = cfg + # These activation functions satisfy f(0) = 0 so we don't need to mask with OOB regions + # (zeros filled by TMA are still zeros without applying activation to them) + # Note that as the time of writing, all activations supported by MXFP8 satisfy this property, so !ACT_NEED_MASKING is never reachable. + # This is only a precaution for future use cases if someone wants to add a new activation function. + self.ACT_NEED_MASKING = cfg.WITH_ACT and cfg.ACTIVATION not in ( + "relu", + "gelu", + "silu", + "qgelu", + "srelu", + ) + # Instantiate all possible kernels at compile time, + # and we will pick the right one at runtime based on the input shape and config. + self.general_skip_input_masking_kernel = MXFP8QuantizeGeneralG2SKernel( + cfg, SKIP_INPUT_MASKING=True + ) + self.general_skip_input_masking_and_scale_bounds_kernel = MXFP8QuantizeGeneralG2SKernel( + cfg, SKIP_INPUT_MASKING=True, SKIP_SCALE_BOUNDS=True + ) + # We only need to mask when WITH_ACT is enabled because with activations applied zeros filled by TMA affect block statistics + self.general_kernel = MXFP8QuantizeGeneralG2SKernel(cfg) if self.ACT_NEED_MASKING else None + + # Specialized cast-only kernels in a specific dispatch priority order so we prefer to dispatch to kernels that come first in the list. + self.specialized_kernels = [ + kernel_cls(cfg) + for kernel_cls in ( + MXFP8QuantizeRowwise2LaneG2RKernel, + MXFP8QuantizeBidimensionalG2RKernel, + MXFP8QuantizeRowwise1LaneG2RKernel, + MXFP8QuantizeBidimensionalG2SKernel, + ) + if kernel_cls.supports(cfg) + ] + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + 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 + for kernel in self.specialized_kernels: + if not dispatched_to_specialized: + if kernel.eligible_to_dispatch( + self.cfg, mX, mO_row, mS_row, mO_col, mS_col, mNoop, mDActInput + ): + dispatched_to_specialized = True + kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + mActTable, + stream, + ) + # If not using a specialized kernel, fall back to the general kernel + if not dispatched_to_specialized: + # With _USE_STATIC_NUM_TILES, each CTA always processes _NUM_TILES tiles, including OOB ones, + # so the scale-bounds check can only be dropped if M divides TILE_ROWS * _NUM_TILES. + num_tiles_per_cta = ( + self.general_skip_input_masking_kernel._NUM_TILES + if self.general_skip_input_masking_kernel._USE_STATIC_NUM_TILES + else 1 + ) + # If the input shape can be perfectly tiled by the general kernel's tile size, we can skip some boundary check because + # we know we will not touch any out-of-bounds region. + shape_is_divisible = ( + mX.shape[0] + % (self.general_skip_input_masking_kernel._TILE_ROWS * num_tiles_per_cta) + == 0 + and mX.shape[1] % self.general_skip_input_masking_kernel._TILE_COLS == 0 + ) + # If shape is divisible, then we won't access OOB regions regardless whatever + if shape_is_divisible: + self.general_skip_input_masking_and_scale_bounds_kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + mActTable, + stream, + ) + else: + # Otherwise we can't skip scale bound check, but we may still skip masking activation inputs with zero + if cutlass.const_expr(self.ACT_NEED_MASKING): + # Masking and the scale guards now share one condition, so this is a plain + # two-way choice: either the shape tiles exactly (drop both) or it does not. + self.general_kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + mActTable, + stream, + ) + else: + # We still need to check the scale bounds, but we can skip some activations masking because their output is 0 + # for OOB regions where TMA fills with zeros + self.general_skip_input_masking_kernel( + mX, + mO_row, + mS_row, + mO_col, + mS_col, + mAmax, + mNoop, + mDActInput, + mWorkspace, + mActTable, + stream, + ) + + +def compile_and_register_function(fn_name: str, cfg: MXFP8QuantizeConfig) -> bool: + """Compile the CuTeDSL kernel for this config and register it in the TVM-FFI global registry + under EXACTLY `fn_name` (the key the C++ dispatcher built; Python treats it as an opaque name). + Returns True if a kernel is successfully registered under `fn_name` (the C++ side then fetches it with GetGlobal(fn_name)); + False if the config is unsupported, so the caller caches the negative result and falls back to the CUDA C++ kernel. + """ + try: + # Already registered (e.g. by a prior call) -> supported. + if tvm_ffi.get_global_func(fn_name, allow_missing=True) is not None: + return True + + major, minor = device_compute_capability() + if major < 10: + logger.warning( + "CuTeDSL MXFP8 backend requires compute capability >= 10.0 (Blackwell), " + "but detected %d.%d; falling back to the CUDA C++ kernel.", + major, + minor, + ) + return False + + logger.debug("Compiling CuTeDSL MXFP8 quantization kernel for %s", cfg) + compiled = cfg.compile() + # The returned compiled object is not neccessarily the compiled function itself. It could be a Python + # wrapper that parses arguments and calls the underlying function. This is needed in the general case + # as TVM-FFI is positional-only, however, for us this is unnecessary launch overhead as we pass the + # arguments in the correct order on the C++ side. + # Relevant issue: https://github.com/NVIDIA/cutlass/issues/3527 + # https://github.com/NVIDIA/cutlass/pull/3589 recommends using an explicit positional only marker (/) + # in the function signature, however it is not merged and it is something easy to miss and silently + # introduce unnecessarry overhead. + native = getattr(compiled, "__tvm_ffi_object__", lambda: None)() + tvm_ffi.register_global_func( + fn_name, native if native is not None else compiled, override=True + ) + return True + except Exception as e: # pylint: disable=broad-exception-caught + logger.error( + "CuTeDSL MXFP8 kernel compilation & registration failed, falling back to the CUDA" + " C++ kernel: %s", + e, + ) + # Unconditionally fallback to CUDA path because we can't tell if this exception is transient or permanent. + return False + + +def get_mxfp8_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + with_amax: bool, + with_dbias: bool, + with_dact: bool, + with_act: bool, + use_2d_quantization: bool, + activation: str, + activation_table: str, +) -> bool: + """Compile and register the MXFP8 quantize kernel for this config under `fn_name`.""" + try: + cfg = MXFP8QuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=with_amax, + with_dbias=with_dbias, + with_dact=with_dact, + with_act=with_act, + use_2d_quantization=use_2d_quantization, + activation=activation, + activation_table=activation_table, + ) + except ValueError as e: + # The exception message states exactly why the config is unsupported + # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a + # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. + logger.warning( + "CuTeDSL MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + return compile_and_register_function(fn_name, cfg) + + +def get_mxfp8_gated_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + is_bwd: bool, + activation: str, +) -> bool: + """Compile and register the gated MXFP8 quantize kernel for this config under `fn_name`.""" + try: + cfg = MXFP8GatedQuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + is_bwd=is_bwd, + activation=activation, + ) + except ValueError as e: + # The exception message states exactly why the config is unsupported + # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a + # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. + logger.warning( + "CuTeDSL MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + return compile_and_register_function(fn_name, cfg) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py new file mode 100644 index 00000000000..5bd1ad9269b --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py @@ -0,0 +1,1883 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Register-resident CuTeDSL bidimensional MXFP8 cast, ported from cast_bidim.cuh.""" + +from typing import Callable, Optional, Type + +import cutlass +from cutlass import cute +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, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, + cvt_f32x2_to_bf16x2, +) +from transformer_engine.common.CuTeDSL.activations import ( + DERIVATIVE_SLOPES, + DERIVATIVE_TABLE_WINDOWS, + FORWARD_TABLE_WINDOWS, + activation_table_bytes, + table_derivative_2x, + table_forward_2x, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + 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, + activation_op, + dactivation_op, + derive_swizzled_scale_layout, +) + + +@cute.jit +def _bf16_nan_magnitudes_to_zero(pair: Uint32) -> Uint32: + """A packed pair of BF16 block amaxes with each NaN half replaced by +0. + + With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's + generic kernel (0-seeded fmaxf); CUDA's cast kernels, like the plain cast here, give it the NaN + scale (cast_bidim.cuh's kNaNReciprocal). + """ + magnitudes = pair & Uint32(0x7FFF7FFF) + lo = magnitudes & Uint32(0xFFFF) + hi = magnitudes >> Uint32(16) + if lo > Uint32(0x7F80): + lo = Uint32(0) + if hi > Uint32(0x7F80): + hi = Uint32(0) + return lo | (hi << Uint32(16)) + + +@cute.jit +def _bf16_mx_scale_reciprocal_fused( + magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int] +) -> Uint32: + """bf16_mx_scale_reciprocal with a fused activation: a NaN block amax gives the reciprocal of + amax 0 (see _bf16_nan_magnitudes_to_zero). Same instructions, so it costs nothing extra.""" + rounded = (magnitude + Uint32(31)) & Uint32(0x7F80) + if rounded < Uint32(MAX_EXPONENT << 7): # pylint: disable=consider-using-max-builtin + rounded = Uint32(MAX_EXPONENT << 7) + reciprocal = Uint32((254 + MAX_EXPONENT) << 7) - rounded + if magnitude >= Uint32(0x7F80): + reciprocal = Uint32(0x0040) + if magnitude != Uint32(0x7F80): + reciprocal = Uint32(254 << 7) + return reciprocal + + +@cute.jit +def _quantize_rowwise( + rX: cute.Tensor, + gO_row_thread: cute.Tensor, + gS_row_thread: cute.Tensor, + # Atoms + output_store_atom: cute.CopyAtom, + row_scale_store_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], + # Constexprs + 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]], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, + FUSED: cutlass.Constexpr[bool] = False, +): + """With EDGE, only the thread's first rows_left rows are stored. FUSED: an activation is fused + (see _bf16_nan_magnitudes_to_zero).""" + tidx, _, _ = cute.arch.thread_idx() + MAX_EXPONENT = 8 if FP8_DTYPE is cutlass.Float8E4M3FN else 15 + + @cute.jit + def extract_mx_block_amax(values_i32: cute.Tensor) -> Uint32: + """Reduce a 32-column MXFP8 block across its two or four lane owners.""" + amax = values_i32[0] + for bf16x2_idx in cutlass.range_constexpr(1, cute.size(values_i32)): + amax = abs_max_x2_bf16(amax, values_i32[bf16x2_idx]) + for stage in cutlass.range_constexpr(THREADS_X_PER_MX_BLOCK.bit_length() - 1): + amax = abs_max_x2_bf16(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) + return bf16_pair_magnitude(amax) + + rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) + rO_row = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_row_u32 = cute.recast_tensor(rO_row, Uint32) + + # Only one thread from threads that share the MXFP8 block writes the rowwise scale + is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + if cutlass.const_expr(FUSED): + reciprocal = _bf16_mx_scale_reciprocal_fused( + extract_mx_block_amax(rX_row_i32), MAX_EXPONENT + ) + else: + reciprocal = bf16_mx_scale_reciprocal(extract_mx_block_amax(rX_row_i32), MAX_EXPONENT) + reciprocal_pair = reciprocal | (reciprocal << Uint32(16)) + for pack_idx in cutlass.range_constexpr( + cute.size(rO_row_u32) + ): # 4 output bytes in one pack + rO_row_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + reciprocal_pair, + rX_row_i32[2 * pack_idx + 1], + reciprocal_pair, + ) + # rO_row_u32 and rO_row share the same memory + 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( + 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 +def _gather_partial_colwise_amaxes( + rX: cute.Tensor, + # Constexprs + 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], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], +) -> cute.Tensor: + """ + rX is a RMEM fragment of layout (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) that contains this thread's data + + It returns a SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)), where each thread + reduces their partial columnwise amax of shape (ELEMENTS_X_PER_THREAD,) and contributes it to the SMEM buffer + """ + tidx, _, _ = cute.arch.thread_idx() + + # We first reduce each thread's rows into one maximum per column before reduce across different threads in the CTA + rAmax_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), DTYPE) + + # Gather the amax for each column that this thread sweeps + rX_i32 = cute.recast_tensor(rX, Int32) + rAmax_col_i32 = cute.recast_tensor(rAmax_col, Int32) + for bf16x2_idx in cutlass.range_constexpr(cute.size(rX_i32, mode=[0])): + amax = rX_i32[bf16x2_idx, 0] + for row_idx in cutlass.range_constexpr(1, ELEMENTS_Y_PER_THREAD): + amax = abs_max_x2_bf16(amax, rX_i32[bf16x2_idx, row_idx]) + rAmax_col_i32[bf16x2_idx] = amax + + # Keep eight BF16 values (16 bytes) contiguous. XOR BF16 offset bit 6 into + # bit 3 so each 128-bit warp transaction uses all 32 banks without overlap. + scratch_swizzle = cute.make_swizzle(1, 3, 3) + + # Prepare amax reduction scratch space in SMEM. + sAmaxs_col_scratch_layout = cute.make_composed_layout( + scratch_swizzle, + 0, + cute.zipped_product( + rAmax_col.layout, # Each thread contributes (ELEMENTS_X_PER_THREAD,) + cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X)), # Total threads in the CTA + ), + ) + allocator = cutlass.memory.SmemAllocator() + sAmaxs_col_scratch = allocator.allocate_tensor( + DTYPE, sAmaxs_col_scratch_layout, byte_alignment=16 + ) + + # Every thread writes its partial maxima using its flat thread ID. + cute.autovec_copy(rAmax_col, sAmaxs_col_scratch[None, tidx]) + cute.arch.sync_threads() + + return sAmaxs_col_scratch + + +@cute.jit +def _reduce_partial_colwise_amaxes( + sAmaxs_col_scratch: cute.Tensor, + # Constexprs + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FUSED: cutlass.Constexpr[bool] = False, +) -> cute.Tensor: + """ + FUSED: an activation is fused (see _bf16_nan_magnitudes_to_zero). + + sAmaxs_col_scratch: the SMEM buffer of layout (ELEMENTS_X_PER_THREAD, (CTA_THREADS_Y, CTA_THREADS_X)) + that contains each threads' own columnwise amax of its ELEMENTS_X_PER_THREAD columns + + 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), + which is the return value + """ + tidx, _, _ = cute.arch.thread_idx() + + TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + + # Each reducer owns one column pair and visits its partial maximum from + # every warp. Column pairs precede warp rows in the scratch's logical order. + REDUCTION_THREADS_NUM = CTA_THREADS_X * ELEMENTS_X_PER_THREAD // 2 + is_reduction_thread = tidx < REDUCTION_THREADS_NUM + # How many bf16x2 pairs we need to reduce + REDUCTION_i32_PAIRS = TILE_COLS // REDUCTION_THREADS_NUM // 2 + + # Use the same 16-byte swizzle as the partial-max scratch. Consecutive + # column-pair reducers still access distinct banks through this layout. + scratch_swizzle = cute.make_swizzle(1, 3, 3) + # SMEM buffer for colwise scale reduction where mode 0 is how many reciprocals a thread has, + # and mode 1 is how many threads we have in X axis, and the layout size is just TILE_COLS + sS_col_layout = cute.make_composed_layout( + scratch_swizzle, + 0, + cute.zipped_product( + cute.make_layout((ELEMENTS_X_PER_THREAD,)), + cute.make_layout((CTA_THREADS_X,)), + ), + ) + allocator = cutlass.memory.SmemAllocator() + sS_col = allocator.allocate_tensor(DTYPE, sS_col_layout, byte_alignment=16) + sS_col_i32 = cute.recast_tensor(sS_col, Int32) + + # Colwise amaxes layout: each reduction thread owns CTA_THREADS_Y rows, with REDUCTION_i32_PAIRS column pairs + _, tv_col_amaxes_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), + val_layout=cute.make_layout( + (REDUCTION_i32_PAIRS, CTA_THREADS_Y), stride=(1, REDUCTION_i32_PAIRS) + ), + ) + # Colwise rcp layout: each reduction thread owns only 1 row, with REDUCTION_i32_PAIRS column pairs + _, tv_col_rcp_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((REDUCTION_THREADS_NUM, 1), stride=(1, REDUCTION_THREADS_NUM)), + val_layout=cute.make_layout((REDUCTION_i32_PAIRS, 1), stride=(1, REDUCTION_i32_PAIRS)), + ) + + sS_col_i32 = cute.composition(sS_col_i32, tv_col_rcp_layout) + + # Each reduction process `REDUCTION_i32_PAIRS` bf16x2 pairs by visiting all rows (CTA_THREADS_Y) + # 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]) + + if cutlass.const_expr(FUSED): + amaxX2 = Int32(_bf16_nan_magnitudes_to_zero(Uint32(amaxX2))) + # Obtain the columnwise scale reciprocal for this column pair + sS_col_i32[tidx, pair_idx] = bf16_mx_scale_reciprocal_pair(amaxX2, MAX_EXPONENT) + cute.arch.sync_threads() + + return sS_col + + +@cute.jit +def _write_colwise_scale( + rS_col_rcp: cute.Tensor, + gS_col_thread: cute.Tensor, + # Atoms + col_scale_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + # Constexprs + 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() + + rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) + + # Only threads from CTA's the first row writes the scale fragments + if tidx < CTA_THREADS_X: + rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) + rS_col_u32 = cute.recast_tensor(rS_col, Uint32) + # Get the colwise scale from the reciprocal + for pack_idx in cutlass.range_constexpr( + cute.size(rS_col_u32) + ): # 4 output bytes in one pack + rS_col_u32[pack_idx] = bf16_mx_scale_bytes_pair(rS_col_rcp_u32[2 * pack_idx]) | ( + 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 + 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 +def _apply_colwise_scales( + rX: cute.Tensor, + rS_col_rcp: cute.Tensor, + gO_col_thread: cute.Tensor, + # Atoms + store_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], + # Constexprs + 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 + (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) + + This function quantizes its fragments using the input and columnwise scales, and write to its output fragment + """ + rS_col_rcp_u32 = cute.recast_tensor(rS_col_rcp, Uint32) + + rO_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_col_u32 = cute.recast_tensor(rO_col, Uint32) + + # Each thread quantizes its own rows with the columnwise reciprocals and write its output to GMEM. + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_row_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + for pack_idx in cutlass.range_constexpr( + cute.size(rO_col_u32) + ): # 4 output bytes in one pack + rO_col_u32[pack_idx] = scale_and_pack( + rX_row_i32[2 * pack_idx], + rS_col_rcp_u32[2 * pack_idx], + rX_row_i32[2 * pack_idx + 1], + rS_col_rcp_u32[2 * pack_idx + 1], + ) + 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 +def _quantize_colwise( + rX: cute.Tensor, + gO_col_thread: cute.Tensor, + gS_col_thread: cute.Tensor, + # Atoms + output_store_atom: cute.CopyAtom, + col_scale_store_atom: cute.CopyAtom, + # Miscellaneous + cache_policy: Int64, + scale_and_pack: cutlass.Constexpr[Callable[..., Uint32]], + # Constexprs + 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]], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, + cols_ok: Boolean = True, + FUSED: cutlass.Constexpr[bool] = False, +): + """Reduce columnwise amaxes across the CTA, write the columnwise scales, and quantize rX with them. + FUSED: an activation is fused (see _bf16_nan_magnitudes_to_zero).""" + tidx, _, _ = cute.arch.thread_idx() + + sAmaxs_col_scratch = _gather_partial_colwise_amaxes( + rX, + 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, + DTYPE=DTYPE, + ) + + sS_col = _reduce_partial_colwise_amaxes( + sAmaxs_col_scratch, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + FUSED=FUSED, + ) + + # Copy the reduced columnwise scale reciprocals to RMEM. + # Threads that quantize different rows share the same columnwise scale + 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, + EDGE=EDGE, + cols_ok=cols_ok, + ) + + _apply_colwise_scales( + rX, + rS_col_rcp, + gO_col_thread, + store_atom=output_store_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + 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 register-tiled kernel and the generic kernel do +# (2^(127-e) is often outside FP16's range): a block's scale is +# e8m0(max(amax * MAX_NORM_RCP, 0)), so a block of only NaNs gets the scale of 0 like the +# generic kernel's 0-seeded fmaxf, and each value is multiplied in FP32 by 2^(127-e) and then +# converted to FP8. + + +@cute.jit +def _scale_from_amax_f32(amax: Float32, MAX_NORM_RCP: cutlass.Constexpr[float]): + """The E8M0 scale of a block with this amax, and its FP32 reciprocal.""" + biased = cvt_f32_to_fp8e8m0fnu(cute.arch.fmax(amax * Float32(MAX_NORM_RCP), Float32(0.0))) + return biased, exp2f_rcp(biased) + + +@cute.jit +def _quantize_rowwise_f32_scaled( + rX: cute.Tensor, + gO_row_thread: cute.Tensor, + gS_row_thread: cute.Tensor, + output_store_atom: cute.CopyAtom, + row_scale_store_atom: cute.CopyAtom, + cache_policy: Int64, + THREADS_X_PER_MX_BLOCK: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, +): + """_quantize_rowwise for FP16 and FP32 values, scaled in FP32.""" + tidx, _, _ = cute.arch.thread_idx() + cvt_x2 = get_cvt_f32x2_to_fp8x2_func(FP8_DTYPE) + rS_row = cute.make_rmem_tensor((1,), SCALE_DTYPE) + rO_row = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_row_u16 = cute.recast_tensor(rO_row, Uint16) + is_writer_thread = tidx % THREADS_X_PER_MX_BLOCK == 0 + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + amax = Float32(0.0) + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + amax = cute.arch.fmax(amax, fabs_f32(Float32(rX[col, row_idx]))) + # The block's amax across the lanes that share it. + for stage in cutlass.range_constexpr(THREADS_X_PER_MX_BLOCK.bit_length() - 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, 1 << stage)) + biased, reciprocal = _scale_from_amax_f32(amax, MAX_NORM_RCP) + for pair in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): + rO_row_u16[pair] = cvt_x2( + Float32(rX[2 * pair + 1, row_idx]) * reciprocal, + Float32(rX[2 * pair, row_idx]) * reciprocal, + ) + row_ok = Boolean(True) + if cutlass.const_expr(EDGE): + row_ok = row_idx < rows_left + if row_ok: + cute.copy( + output_store_atom, + rO_row, + gO_row_thread[(None, row_idx),], + cache_policy=cache_policy, + ) + if is_writer_thread: + rS_row[0] = biased + cute.copy( + row_scale_store_atom, + rS_row, + cute.local_tile(gS_row_thread, (1,), (row_idx,)), + cache_policy=cache_policy, + ) + + +@cute.jit +def _quantize_colwise_f32_scaled( + rX: cute.Tensor, + gO_col_thread: cute.Tensor, + gS_col_thread: cute.Tensor, + output_store_atom: cute.CopyAtom, + col_scale_store_atom: cute.CopyAtom, + cache_policy: Int64, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + SCALE_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, + cols_ok: Boolean = True, +): + """_quantize_colwise for FP16 and FP32 values, scaled in FP32. Thread x of each warp row + holds columns [x * ELEMENTS_X_PER_THREAD, (x + 1) * ELEMENTS_X_PER_THREAD).""" + tidx, _, _ = cute.arch.thread_idx() + THREADS = CTA_THREADS_Y * CTA_THREADS_X + TILE_COLS = CTA_THREADS_X * ELEMENTS_X_PER_THREAD + cvt_x2 = get_cvt_f32x2_to_fp8x2_func(FP8_DTYPE) + allocator = cutlass.memory.SmemAllocator() + # Each thread's column maxima over its rows, contiguous per thread. + sPartial = allocator.allocate_tensor( + Float32, + cute.make_layout((ELEMENTS_X_PER_THREAD, THREADS), stride=(1, ELEMENTS_X_PER_THREAD)), + byte_alignment=16, + ) + sScale = allocator.allocate_tensor( + SCALE_DTYPE, cute.make_layout((TILE_COLS,)), byte_alignment=16 + ) + sReciprocal = allocator.allocate_tensor( + Float32, cute.make_layout((TILE_COLS,)), byte_alignment=16 + ) + + rPartial = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), Float32) + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + amax = Float32(0.0) + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + amax = cute.arch.fmax(amax, fabs_f32(Float32(rX[col, row_idx]))) + rPartial[col] = amax + cute.autovec_copy(rPartial, sPartial[None, tidx]) + cute.arch.sync_threads() + + # One thread per column reduces the warps' maxima: column c is thread (c / X) of each warp + # row, element c % X. + if tidx < TILE_COLS: + owner = tidx // ELEMENTS_X_PER_THREAD + element = tidx % ELEMENTS_X_PER_THREAD + amax = Float32(0.0) + for warp_row in cutlass.range_constexpr(CTA_THREADS_Y): + amax = cute.arch.fmax(amax, sPartial[element, warp_row * CTA_THREADS_X + owner]) + biased, reciprocal = _scale_from_amax_f32(amax, MAX_NORM_RCP) + sScale[tidx] = biased + sReciprocal[tidx] = reciprocal + cute.arch.sync_threads() + + column0 = (tidx % CTA_THREADS_X) * ELEMENTS_X_PER_THREAD + # The first warp row writes the scales. + if tidx < CTA_THREADS_X: + rS_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), SCALE_DTYPE) + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + rS_col[col] = sScale[column0 + col] + if cols_ok: + cute.copy(col_scale_store_atom, rS_col, gS_col_thread, cache_policy=cache_policy) + + rReciprocal = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), Float32) + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + rReciprocal[col] = sReciprocal[column0 + col] + rO_col = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), FP8_DTYPE) + rO_col_u16 = cute.recast_tensor(rO_col, Uint16) + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + for pair in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): + rO_col_u16[pair] = cvt_x2( + Float32(rX[2 * pair + 1, row_idx]) * rReciprocal[2 * pair + 1], + Float32(rX[2 * pair, row_idx]) * rReciprocal[2 * pair], + ) + row_ok = Boolean(True) + if cutlass.const_expr(EDGE): + row_ok = row_idx < rows_left + if row_ok: + cute.copy( + output_store_atom, + rO_col, + gO_col_thread[(None, row_idx),], + cache_policy=cache_policy, + ) + + +def thread_tile_layout( + CTA_THREADS_Y: int, CTA_THREADS_X: int, ELEMENTS_Y_PER_THREAD: int, ELEMENTS_X_PER_THREAD: int +): + """The (thread, value) layout of a register-resident (CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD, + CTA_THREADS_X * ELEMENTS_X_PER_THREAD) tile: each warp owns ELEMENTS_Y_PER_THREAD consecutive + rows, and each lane ELEMENTS_X_PER_THREAD adjacent columns.""" + thread_layout = cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)) + _, tv_layout = cute.make_layout_tv( + thr_layout=thread_layout, + val_layout=cute.make_layout( + (ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD), stride=(ELEMENTS_X_PER_THREAD, 1) + ), + ) + return tv_layout + + +@cute.jit +def thread_tile_bounds( + tile_shape: cute.Shape, + tv_layout: cute.Layout, + rows_valid: Int32, + cols_valid: Int32, +): + """(rows_left, cols_ok) of the calling thread in a tile that ends at rows_valid rows and + cols_valid columns: its rows from rows_left on are past the tensor, and all of them when its + columns are (a thread's columns lie all inside or all past the tensor).""" + tidx, _, _ = cute.arch.thread_idx() + thread_origin = cute.composition(cute.make_identity_tensor(tile_shape), tv_layout)[tidx, None][ + ((0, 0),) + ] + cols_ok = thread_origin[1] < cols_valid + rows_left = rows_valid - thread_origin[0] + if not cols_ok: + rows_left = Int32(0) + return rows_left, cols_ok + + +@cute.jit +def load_thread_rows( + gT_thread: cute.Tensor, + rT: cute.Tensor, + load_atom: cute.CopyAtom, + cache_policy: Int64, + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + EDGE: cutlass.Constexpr[bool] = False, + rows_left: Int32 = 0, +): + """Copy all rows a thread sweeps into its (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) + register fragment. With EDGE, rows from rows_left on load as zeros.""" + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + if cutlass.const_expr(EDGE): + # Predicated loads rather than branches, so the rows' loads still issue together. + rT[None, row_idx].fill(0) + row_pred = cute.make_rmem_tensor( + cute.make_layout((ELEMENTS_X_PER_THREAD, (1,)), stride=(0, (0,))), Boolean + ) + row_pred[0] = row_idx < rows_left + cute.copy( + load_atom, + gT_thread[(None, row_idx),], + rT[None, row_idx], + pred=row_pred, + cache_policy=cache_policy, + ) + else: + cute.copy( + load_atom, gT_thread[(None, row_idx),], rT[None, row_idx], cache_policy=cache_policy + ) + + +@cute.jit +def quantize_mxfp8_rmem_tile_bidimensional( + rX: cute.Tensor, + gO_row_tile: Optional[cute.Tensor], + gS_row_tile: Optional[cute.Tensor], + gO_col_tile: Optional[cute.Tensor], + gS_col_tile: Optional[cute.Tensor], + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + FUSED: cutlass.Constexpr[bool], + EDGE: cutlass.Constexpr[bool], + rows_left: Int32, + cols_ok: Boolean, + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + cache_policy: Int64, +): + """Quantize a register-resident 32-row tile, the CTA's rX fragments (see thread_tile_layout), + to the (TILE_ROWS, TILE_COLS) output tiles of the directions enabled by ROWWISE / COLWISE. + FUSED: the values come from a fused activation (see _bf16_nan_magnitudes_to_zero). With EDGE, + rows_left and cols_ok come from thread_tile_bounds.""" + TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD + THREADS_X_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_X_PER_THREAD + THREADS_Y_PER_MX_BLOCK = MXFP8_BLOCK_SCALING_SIZE // ELEMENTS_Y_PER_THREAD + SCALE_DTYPE = cutlass.Float8E8M0FNU + tidx, _, _ = cute.arch.thread_idx() + tv_layout = thread_tile_layout( + CTA_THREADS_Y, CTA_THREADS_X, ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD + ) + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + FP8_DTYPE, + num_bits_per_copy=ELEMENTS_X_PER_THREAD * FP8_DTYPE.width, + ) + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) + + if cutlass.const_expr(ROWWISE): + gO_row_thread = cute.composition(gO_row_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_X_PER_MX_BLOCK since they share the same rowwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_row_scale = cute.make_layout( + ( + (THREADS_X_PER_MX_BLOCK, (CTA_THREADS_X // THREADS_X_PER_MX_BLOCK, CTA_THREADS_Y)), + ELEMENTS_Y_PER_THREAD, + ), + stride=( + (0, (TILE_ROWS, ELEMENTS_Y_PER_THREAD)), + 1, + ), + ) + gS_row_thread = cute.composition(gS_row_tile, tv_row_scale)[tidx, None] + + row_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), SCALE_DTYPE, num_bits_per_copy=8 + ) + + if cutlass.const_expr(DTYPE is not cutlass.BFloat16): + _quantize_rowwise_f32_scaled( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + MAX_NORM_RCP=MAX_NORM_RCP, + EDGE=EDGE, + rows_left=rows_left, + ) + else: + _quantize_rowwise( + rX, + gO_row_thread, + gS_row_thread, + output_store_atom=store_atom, + row_scale_store_atom=row_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + THREADS_X_PER_MX_BLOCK=THREADS_X_PER_MX_BLOCK, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + EDGE=EDGE, + rows_left=rows_left, + FUSED=FUSED, + ) + + if cutlass.const_expr(COLWISE): + gO_col_thread = cute.composition(gO_col_tile, tv_layout)[tidx, None] + + # Group CTA threads by THREADS_Y_PER_MX_BLOCK since they share the same columnwise scale for a MXFP8 block. + # And assign stride 0 to that submode + tv_col_scale = cute.make_layout( + ( + (CTA_THREADS_X, (THREADS_Y_PER_MX_BLOCK, CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK)), + ELEMENTS_X_PER_THREAD, + ), + stride=( + (ELEMENTS_X_PER_THREAD, (0, 1)), + 1, + ), + ) + # If CTA_THREADS_Y // THREADS_Y_PER_MX_BLOCK == 1, this will help compiler derive alignment + # Otherwise coalesce will not do anything which is still fine + tv_col_scale = cute.coalesce(tv_col_scale, target_profile=(1, 1)) + gS_col_thread = cute.composition(gS_col_tile, tv_col_scale)[tidx, None] + + # GEMM-swizzled colwise scales are not contiguous (adjacent columns land 16 bytes apart), + # so they are stored a byte at a time, as the general kernel does. + col_scale_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + SCALE_DTYPE, + num_bits_per_copy=8 if WITH_GEMM_SWIZZLED_SCALES else ELEMENTS_X_PER_THREAD * 8, + ) + + if cutlass.const_expr(DTYPE is not cutlass.BFloat16): + _quantize_colwise_f32_scaled( + rX, + gO_col_thread, + gS_col_thread, + output_store_atom=store_atom, + col_scale_store_atom=col_scale_atom, + cache_policy=cache_policy, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + FP8_DTYPE=FP8_DTYPE, + MAX_NORM_RCP=MAX_NORM_RCP, + EDGE=EDGE, + rows_left=rows_left, + cols_ok=cols_ok, + ) + else: + _quantize_colwise( + rX, + gO_col_thread, + gS_col_thread, + output_store_atom=store_atom, + col_scale_store_atom=col_scale_atom, + cache_policy=cache_policy, + scale_and_pack=scale_and_pack, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + SCALE_DTYPE=SCALE_DTYPE, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + EDGE=EDGE, + rows_left=rows_left, + cols_ok=cols_ok, + FUSED=FUSED, + ) + + +@cute.jit +def quantize_mxfp8_gmem_tile_bidimensional( + gX_tile: cute.Tensor, + gA_tile: Optional[cute.Tensor], + gO_row_tile: Optional[cute.Tensor], + gS_row_tile: Optional[cute.Tensor], + gO_col_tile: Optional[cute.Tensor], + gS_col_tile: Optional[cute.Tensor], + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + 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], + ELEMENTS_Y_PER_THREAD: cutlass.Constexpr[int], + ELEMENTS_X_PER_THREAD: cutlass.Constexpr[int], + cache_policy: Int64, +): + """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile in the directions enabled by ROWWISE / COLWISE. + + TILE_ROWS = CTA_THREADS_Y * ELEMENTS_Y_PER_THREAD = 32, so every columnwise + MXFP8 block is contained in the tile. Each warp owns ELEMENTS_Y_PER_THREAD + consecutive rows, and each lane owns ELEMENTS_X_PER_THREAD adjacent columns. + Load the tile once, then reuse its registers for each enabled quantization pass. + The tensors of a disabled direction are unused and may be None. + """ + tidx, _, _ = cute.arch.thread_idx() + + tv_layout = thread_tile_layout( + CTA_THREADS_Y, CTA_THREADS_X, ELEMENTS_Y_PER_THREAD, ELEMENTS_X_PER_THREAD + ) + + gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] + + # Input fragment contains all rows it handles; output fragment contains only one row and we write it back row by row + rX = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), DTYPE) + + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=min( + ELEMENTS_X_PER_THREAD * DTYPE.width, + max_num_bits_per_copy_ld_global(), + ), + # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using + # invariant=True here will drop the cache policy + invariant=False, + ) + # 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): + rows_left, cols_ok = thread_tile_bounds(gX_tile.shape, tv_layout, rows_valid, cols_valid) + + # Copy all rows a thread sweeps into the RMEM fragment at once + load_thread_rows( + gX_thread, + rX, + load_atom, + cache_policy, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + EDGE=EDGE, + rows_left=rows_left, + ) + + if cutlass.const_expr(WITH_ACT or WITH_DACT): + rA = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD), DTYPE) + if cutlass.const_expr(WITH_DACT): + act_load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=min( + ELEMENTS_X_PER_THREAD * DTYPE.width, + max_num_bits_per_copy_ld_global(), + ), + invariant=False, + ) + # 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] + load_thread_rows( + gA_thread, + rA, + act_load_atom, + cache_policy, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + EDGE=EDGE, + rows_left=rows_left, + ) + # With dbias, the FP32 products are staged for the column sums (see _kernel_main). + if cutlass.const_expr(sDbias is not None): + sDbias_thread = cute.composition(sDbias, tv_layout)[tidx, None] + # 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),]) + else: + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + rX_i32 = cute.recast_tensor(rX[None, row_idx], Int32) + rA_i32 = cute.recast_tensor(rA[None, row_idx], Int32) + if cutlass.const_expr(sDbias is not None): + rProducts = cute.make_rmem_tensor((ELEMENTS_X_PER_THREAD,), Float32) + for i in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD // 2): + if cutlass.const_expr(ACTIVATION_TABLE != "none"): + # Read the activation from the table, as the CUDA register-tiled kernel does. + if cutlass.const_expr(WITH_DACT): + low, high, end, stride = DERIVATIVE_TABLE_WINDOWS[ + (ACTIVATION, ACTIVATION_TABLE) + ] + d_lo, d_hi = table_derivative_2x( + rA_i32[i], + cute.recast_tensor(sTable, Float32), + LO=low, + HI=high, + SATURATED_END=end, + STRIDE=stride, + SLOPE=DERIVATIVE_SLOPES[ACTIVATION], + ) + lo = x2_lo_to_f32_bf16(rX_i32[i]) * d_lo + hi = x2_hi_to_f32_bf16(rX_i32[i]) * d_hi + if cutlass.const_expr(sDbias is not None): + rProducts[2 * i] = lo + rProducts[2 * i + 1] = hi + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + else: + low, high = FORWARD_TABLE_WINDOWS[ACTIVATION] + rX_i32[i] = table_forward_2x( + rX_i32[i], + cute.recast_tensor(sTable, Uint16), + LOW_SENTINEL=low - 1, + HIGH_SENTINEL=high + 1, + ) + else: + lo = x2_lo_to_f32_bf16(rX_i32[i]) + hi = x2_hi_to_f32_bf16(rX_i32[i]) + if cutlass.const_expr(WITH_DACT): + dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) + lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) + hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + if cutlass.const_expr(sDbias is not None): + rProducts[2 * i] = lo + rProducts[2 * i + 1] = hi + else: + op = activation_op(ACTIVATION, ACTIVATION_TABLE) + lo = op(lo) + hi = op(hi) + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + if cutlass.const_expr(sDbias is not None): + cute.autovec_copy(rProducts, sDbias_thread[(None, row_idx),]) + + # The thread's amax of the values it quantizes (the activation rounded to BF16). Like the + # CUDA kernels' fmaxf, it ignores NaN: the BF16 max returns the other operand of a NaN, and + # the FP32 fmax drops a pair of NaNs. + thread_amax = Float32(0.0) + if cutlass.const_expr(WITH_AMAX and DTYPE is not cutlass.BFloat16): + for row_idx in cutlass.range_constexpr(ELEMENTS_Y_PER_THREAD): + for col in cutlass.range_constexpr(ELEMENTS_X_PER_THREAD): + thread_amax = cute.arch.fmax(thread_amax, fabs_f32(Float32(rX[col, row_idx]))) + elif cutlass.const_expr(WITH_AMAX): + rX_all_i32 = cute.recast_tensor(rX, Int32) + pair_amax = rX_all_i32[0] + for word in cutlass.range_constexpr(1, cute.size(rX_all_i32)): + 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)) + ) + + quantize_mxfp8_rmem_tile_bidimensional( + rX, + gO_row_tile, + gS_row_tile, + gO_col_tile, + gS_col_tile, + ROWWISE=ROWWISE, + COLWISE=COLWISE, + FUSED=WITH_ACT or WITH_DACT, + EDGE=EDGE, + rows_left=rows_left, + cols_ok=cols_ok, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + MAX_NORM_RCP=MAX_NORM_RCP, + WITH_GEMM_SWIZZLED_SCALES=WITH_GEMM_SWIZZLED_SCALES, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + cache_policy=cache_policy, + ) + + return thread_amax + + +class MXFP8QuantizeBidimensionalG2RKernel(MXFP8QuantizeKernelBase): + """Load one CTA tile into registers and quantize it along both axes. + + TILE_ROWS = _CTA_THREADS_Y * _ELEMENTS_Y_PER_THREAD = 32. + TILE_COLS = _CTA_THREADS_X * ELEMENTS_X_PER_THREAD = 256 or 512. + The 32-row tile contains complete columnwise MXFP8 blocks; shared memory + combines the partial column maxima from the eight warps. Requires aligned + BF16 input, M divisible by 32 and N divisible by 256; writes packed or GEMM-swizzled + scales. + """ + + _CTA_THREADS_Y = 8 + _CTA_THREADS_X = 32 + # With a table, the grid holds about this many columns' worth of CTAs (16384 CTAs of 256 + # columns, as in the CUDA register-tiled kernel); each CTA walks the remaining row tiles. + _WALK_TARGET_COLS = 16384 * 256 + + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # BF16 packed math; FP16 and FP32 scaled in FP32, only with a fused activation. + # dbias only fused with an activation derivative, like the CUDA register-tiled kernel; + # amax only with a fused activation; plain casts with amax take the general kernel. + return ( + (not cfg.WITH_AMAX or cfg.WITH_ACT or cfg.WITH_DACT) + and (not cfg.WITH_DBIAS or cfg.WITH_DACT) + and ( + cfg.DTYPE is cutlass.BFloat16 + or MXFP8QuantizeBidimensionalG2RKernel.takes_scaled_in_f32(cfg) + ) + and cfg.COLWISE + ) + + @staticmethod + def takes_scaled_in_f32(cfg: MXFP8QuantizeConfig) -> bool: + """Whether an FP16 or FP32 request takes this kernel: the fused activations where it is + faster than the general kernel. GB200, geometric mean speedup over TE main, general -> + this kernel: FP16 GeLU 1.46 -> 1.62, SiLU 1.41 -> 1.69, dGeLU 1.56 -> 1.69, dSiLU + 1.28 -> 1.76, dbias_dSiLU 1.48 -> 1.74; FP32 SiLU 1.32 -> 1.54, dGeLU 1.38 -> 1.52, + dSiLU 1.10 -> 1.85, dbias_dSiLU 1.18 -> 1.53, dbias_dGeLU 1.47 -> 1.25 (CUDA 1.30; the + general kernel's FP32 dbias differs from CUDA's in the last bit, this kernel's matches). + FP16 dbias_dGeLU (1.75 -> 1.50) and FP32 GeLU (1.20 -> 1.16) keep the general kernel.""" + if not (cfg.WITH_ACT or cfg.WITH_DACT): + return False + if cfg.WITH_DBIAS and cfg.ACTIVATION == "dgelu" 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( + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + ) -> Boolean: + M, N = mX.shape + 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 (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 ( + 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, as the CUDA register-tiled kernel + does. The others keep whole 32-row by 256-column tiles.""" + return ( + cfg.ACTIVATION_TABLE != "none" + or cfg.WITH_DACT + or cfg.WITH_AMAX + or cfg.DTYPE is not cutlass.BFloat16 + ) + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + # FP32 values take twice the registers: twice the warps, each with half the rows, and + # half the resident CTAs per SM, as in the CUDA register-tiled kernel. + self._MIN_BLOCKS_PER_MP = 4 + if cfg.DTYPE is cutlass.Float32: + self._CTA_THREADS_Y = 16 + self._MIN_BLOCKS_PER_MP = 2 + # One CTA only quantize one columnwise MXFP8 block at a time + self._ELEMENTS_Y_PER_THREAD = MXFP8_BLOCK_SCALING_SIZE // self._CTA_THREADS_Y # 4 + # SMs of the device compiled for, which set the walk length with a table or dbias. + self._sm_count = 0 + if cfg.ACTIVATION_TABLE != "none" or cfg.WITH_DBIAS: + from cuda.core import ( + Device, + ) # pylint: disable=import-outside-toplevel,no-name-in-module + + self._sm_count = HardwareInfo(Device().device_id).get_device_multiprocessor_count() + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + 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 + # 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, + mDActInput, + mO_row, + mS_row, + mO_col, + mS_col, + mActTable, + mWorkspace, + mAmax, + stream, + ELEMENTS_X_PER_THREAD=16, + CLUSTER_WIDTH=4, + ) + else: + self._launch( + mX, + mDActInput, + mO_row, + mS_row, + mO_col, + mS_col, + mActTable, + mWorkspace, + mAmax, + stream, + ELEMENTS_X_PER_THREAD=16, + CLUSTER_WIDTH=1, + ) + else: + self._launch( + mX, + mDActInput, + mO_row, + mS_row, + mO_col, + mS_col, + mActTable, + mWorkspace, + mAmax, + stream, + ELEMENTS_X_PER_THREAD=8, + CLUSTER_WIDTH=8, + ) + + @cute.jit + def _launch( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO_row: cute.Tensor, + 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], + ): + assert ELEMENTS_X_PER_THREAD in (8, 16) + + # Dispatch checks these alignments before selecting this kernel. + mX = cute.make_tensor(mX.iterator.align(32), mX.layout) + mS_col = cute.make_tensor(mS_col.iterator.align(16), mS_col.layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + mA = cute.make_tensor(mA.iterator.align(32), mA.layout) + M, N = mX.shape + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X_PER_THREAD + THREADS_PER_CTA = self._CTA_THREADS_Y * self._CTA_THREADS_X + + 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) + # Adjust cluster size so that the CTAs in the column direction can be evenly divided into clusters. + while cluster_cols > 1 and grid_cols % cluster_cols != 0: + cluster_cols -= 1 + # With a table, each CTA walks several row tiles so that it copies the table once for + # all of them, as the CUDA register-tiled kernel does: halve the grid while it holds + # more than _WALK_TARGET_COLS / TILE_COLS CTAs and the row tiles divide evenly, and + # walk at least two tiles unless that leaves fewer than two CTAs per SM. + walk = Int32(1) + if cutlass.const_expr(self.cfg.ACTIVATION_TABLE != "none" 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 // walk] + block = [THREADS_PER_CTA] + cluster = [cluster_cols, 1, 1] + + # Hardcode the layout using static N for some of the commonly used shapes + if N == 2048: + payload_layout = cute.make_layout((M, 2048), stride=(2048, 1)) + gX = cute.make_tensor(mX.iterator, payload_layout) + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) + gO_row = ( + cute.make_tensor(mO_row.iterator, payload_layout) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) + gO_col = cute.make_tensor(mO_col.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + gS_row, gS_col = derive_swizzled_scale_layout( + M, 2048, self.cfg.ROWWISE, True, mS_row, mS_col + ) + else: + gS_row = ( + cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 64), stride=(64, 1)), + ) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 2048), stride=(2048, 1)), + ) + # 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) + gA = ( + cute.make_tensor(mA.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_DACT) + else None + ) + gO_row = ( + cute.make_tensor(mO_row.iterator, payload_layout) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) + gO_col = cute.make_tensor(mO_col.iterator, payload_layout) + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + gS_row, gS_col = derive_swizzled_scale_layout( + M, 4096, self.cfg.ROWWISE, True, mS_row, mS_col + ) + else: + gS_row = ( + cute.make_tensor( + mS_row.iterator, + cute.make_layout((cute.size(mS_row, mode=[0]), 128), stride=(128, 1)), + ) + if cutlass.const_expr(self.cfg.ROWWISE) + else None + ) + gS_col = cute.make_tensor( + mS_col.iterator, + cute.make_layout((cute.size(mS_col, mode=[0]), 4096), stride=(4096, 1)), + ) + # 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( + M, N, self.cfg.ROWWISE, True, mS_row, mS_col + ) + else: + gS_row, gS_col = mS_row, mS_col + # 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( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO_row: cute.Tensor, + 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.""" + self._kernel_main( + mX, + mA, + mO_row, + 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, as in the CUDA register-tiled + kernel: the full 64-bit layout costs registers and instructions on every access.""" + if cutlass.const_expr(DYNAMIC_STRIDE): + N = mT.shape[1] + base = Int64(row_tile) * Int64(TILER[0]) * Int64(N) + Int64(cta_col) * Int64(TILER[1]) + return cute.make_tensor( + (mT.iterator + base).align(ALIGNMENT), + cute.make_layout(TILER, stride=(Int32(N), 1)), + ) + return cute.local_tile(mT, TILER, (row_tile, cta_col)) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO_row: cute.Tensor, + 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 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) + # 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 = 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 = 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 = self._data_tile(mO_col, row_tile, cta_col, TILER, DYNAMIC_STRIDE, 16) + gS_col_tile = cute.local_tile(mS_col, COL_SCALE_TILER, cta_coord) + + # With 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, + gS_row_tile, + gO_col_tile, + gS_col_tile, + 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, + ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=ELEMENTS_X_PER_THREAD, + cache_policy=cache_policy, + ) 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 new file mode 100644 index 00000000000..cabe7cf3d73 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_1lane.py @@ -0,0 +1,335 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""One-lane GMEM-to-RMEM rowwise MXFP8 cast with vectorized global I/O.""" + +# Local @cute.struct classes are shared-memory storage descriptors. +# pylint: disable=missing-class-docstring + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import Boolean, Float8E8M0FNU, Int32, Uint32, Uint8 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + is_packed16, + abs_max_x2_bf16, + abs_max_x2_f16, + x2_lo_to_f32_bf16, + x2_lo_to_f32_f16, + x2_hi_to_f32_bf16, + x2_hi_to_f32_f16, + fabs_f32, + exp2f_rcp, + pack_f32x2, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + mul_f32x2_cvt_packed16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + MXFP8_BLOCK_SCALING_SIZE, + derive_swizzled_scale_layout, +) + + +@cute.jit +def _flush_scales_to_gmem( + sScale: cute.Tensor, + gS_tile: cute.Tensor, + cta_row: Int32, + cta_col: Int32, + M: Int32, + padded_cols: Int32, + CTA_Y: cutlass.Constexpr[int], + CTA_X: cutlass.Constexpr[int], + width: cutlass.Constexpr[int], +): + """Flush the staged (CTA_Y, CTA_X) scale tile to gmem with vectorized stores.""" + tidx, _, _ = cute.arch.thread_idx() + # Previously each threads has 1 byte, but now we are doing vectorized store, + # which means only a subset of threads will need to issue the store while other threads are not used. + active_threads = CTA_X // width + _, tv_flush = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, active_threads), stride=(active_threads, 1)), + val_layout=cute.make_layout((1, width), stride=(width, 1)), + ) + # We only need to use a subset of threads with shape (CTA_Y, active_threads) to write + # so if the thread is outside of this subset, it will remain inactive + if tidx < CTA_Y * active_threads: + # Absolute position of the scale vector to write in the GMEM buffer + thread_y = cta_row * CTA_Y + tidx // active_threads + thread_x = cta_col * CTA_X + (tidx % active_threads) * width + if thread_y < M and thread_x < padded_cols: + cute.autovec_copy( + cute.composition(sScale, tv_flush)[tidx, None], + cute.composition(gS_tile, tv_flush)[tidx, None], + ) + + +@cute.jit +def quantize_mxfp8_gmem_tile_rowwise_1lane( + gX_tile: cute.Tensor, + gO_tile: cute.Tensor, + gS_tile: cute.Tensor, + M: Int32, + N: Int32, + padded_cols: Int32, + cta_row: Int32, + cta_col: Int32, + 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], + TILE_ROWS: cutlass.Constexpr[int], + TILE_COLS: cutlass.Constexpr[int], + STASH_SCALE_TO_SMEM: cutlass.Constexpr[bool], +): + """Quantize one (TILE_ROWS, TILE_COLS) GMEM tile rowwise with TILE_ROWS * TILE_COLS / 32 threads. + + Each thread handles one 32-element MXFP8 chunk (= one scale block). (cta_row, cta_col) is the + tile coordinate, used with M and N for bounds checks; padded_cols is the column count of the + (padded) scale tensor that gS_tile is a tile of. + """ + tidx, _, _ = cute.arch.thread_idx() + + # The CTA threads are grouped as (CTA_Y, CTA_X), so they cover a (TILE_ROWS, TILE_COLS) + # input tile and the matching (CTA_Y, CTA_X) scale tile. + CTA_Y = TILE_ROWS + CTA_X = TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) + ), + ) + _, tv_layout_scale = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1)), + val_layout=cute.make_layout((1, 1), stride=(1, 1)), + ) + + gX_thread = cute.composition(gX_tile, tv_layout)[tidx, None] + gO_thread = cute.composition(gO_tile, tv_layout)[tidx, None] + gS_thread = cute.composition(gS_tile, tv_layout_scale)[tidx, None] + + rX_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=DTYPE, + ) + abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 + x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 + x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 + mul_cvt4 = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, FP8_DTYPE) + rX_i32 = cute.make_tensor( + cute.recast_ptr(rX_thread.iterator, dtype=Int32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + rO_thread = cute.make_rmem_tensor( + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1)), + dtype=FP8_DTYPE, + ) + rO_u32 = cute.make_tensor( + cute.recast_ptr(rO_thread.iterator, dtype=Uint32), + cute.make_layout( + (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) + ), # Unit is Uint32, divide by 4 here + ) + + sS_thread = None + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + + @cute.struct + class SharedStorage: + buf: cute.struct.Align[cute.struct.MemRange[Float8E8M0FNU, CTA_Y * CTA_X], 16] + + storage = cutlass.memory.SmemAllocator().allocate(SharedStorage) + sScale = storage.buf.get_tensor(cute.make_layout((CTA_Y, CTA_X), stride=(CTA_X, 1))) + # sScale is (CTA_Y, CTA_X):(CTA_X, 1), which is the same layout as tv_layout_scale + # so sS_thread is really just a 1-element e8m0 buffer for this thread's scale. + sS_thread = cute.composition(sScale, tv_layout_scale)[tidx, None] + # Zero first so padding columns (cols past N/32 in the padded scale + # matrix) flush as 0 and we never read uninitialized smem. + # Raw 0x00: e8m0 has no zero value to convert from. + sS_thread[0] = Uint8(0).bitcast(Float8E8M0FNU) + cute.arch.sync_threads() + + row = cta_row * TILE_ROWS + tidx // CTA_X + col = cta_col * TILE_COLS + (tidx % CTA_X) * MXFP8_BLOCK_SCALING_SIZE + if row < M and col < N: + cute.autovec_copy(gX_thread, rX_thread) + amax_2x = rX_i32[0] + # Skip the first iteration because we assigned rX_i32[0] to amax_2x already + for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE // 2): + amax_2x = abs_max_x2(amax_2x, rX_i32[i]) + amax = cute.arch.fmax(fabs_f32(x2_lo_to_f32(amax_2x)), fabs_f32(x2_hi_to_f32(amax_2x))) + + biased_exp = cvt_f32_to_fp8e8m0fnu(amax * MAX_NORM_RCP) + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + sS_thread[0] = biased_exp + else: + gS_thread[0] = biased_exp + + # Rescale + FP8 cast, 4 elements per fused mul_cvt (one uint32 out), + # then a vectorized store. Mirrors CUDA's _use_cvt_4x path. + inv_scale = exp2f_rcp(biased_exp) + scale_2x = pack_f32x2(inv_scale, inv_scale) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + rO_u32[i] = mul_cvt4(rX_i32[2 * i], rX_i32[2 * i + 1], scale_2x) + cute.autovec_copy(rO_thread, gO_thread) + + # Cooperative wide flush of the staged scales where padding columns flush as 0. + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + cute.arch.sync_threads() + if cutlass.const_expr(WITH_GEMM_SWIZZLED_SCALES): + # Swizzled rowwise scale layout is ((32, 4, num_tiles_M), (4, num_tiles_SC)):((16, 4, num_tiles_SC * 512), (1, 512)) + # so we have at most 4 elements continuous which is our vectorized store width + _flush_scales_to_gmem( + sScale, gS_tile, cta_row, cta_col, M, padded_cols, CTA_Y, CTA_X, 4 + ) + elif padded_cols % 16 == 0: + # If columns is divisible by 16, use 16 bytes as the vectorized store width + _flush_scales_to_gmem( + sScale, gS_tile, cta_row, cta_col, M, padded_cols, CTA_Y, CTA_X, 16 + ) + else: + # Otherwise use 4 bytes as the vectorized store width. + # Note our fake tensor requires 4 divisibility so this is enforced as long as you can get here + _flush_scales_to_gmem( + sScale, gS_tile, cta_row, cta_col, M, padded_cols, CTA_Y, CTA_X, 4 + ) + + +class MXFP8QuantizeRowwise1LaneG2RKernel(MXFP8QuantizeKernelBase): + """Specialized cast-only ROWWISE-only MXFP8 kernel. + + Requires N % 128 == 0 (full vectorizable column chunks). + + Plain rowwise-only quantize. Each thread owns one 32-element MXFP8 chunk and + uses vectorized global loads/stores (no TMA used).""" + + _TILE_ROWS = 4 + _TILE_COLS = 1024 + _THREADS_PER_CTA = 128 + + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + return cfg.is_plain_cast() and is_packed16(cfg.DTYPE) and cfg.ROWWISE and not cfg.COLWISE + + @staticmethod + @cute.jit + def eligible_to_dispatch( + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + ) -> Boolean: + M, N = mX.shape + # N divisible by 128 for vectorized stores, and grid.y within limits. + return ( + N % 128 == 0 + and cute.ceil_div(M, MXFP8QuantizeRowwise1LaneG2RKernel._TILE_ROWS) <= 65535 + and mNoop.toint() == 0 + ) + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + # If True, then this kernel will first write each thread's scale byte to a shared memory buffer, + # then utilize vectorized store to flush the buffer to global memory. + self._STASH_SCALE_TO_SMEM = True # Hardcode to true for now + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], # Unused, kept for API compatibility + mAmax: Optional[cute.Tensor], # Unused, kept for API compatibility + 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): + cute.printf( + f"[CuTeDSL] MXFP8QuantizeRowwise1LaneG2RKernel.__call__() with config: {self.cfg}\n" + ) + + M = mX.shape[0] + N = mX.shape[1] + + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, _ = derive_swizzled_scale_layout(M, N, True, False, mS_row, None) + + grid = [ + cute.ceil_div(Int32(N), self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS), + ] + block = [self._THREADS_PER_CTA] + + self.kernel( + mX, + mO_row, + mS_row, + ).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + ): + """Device entry. Like the CUDA specialized cast kernels it does not consume the + cast_noop flag: the dispatcher routes calls with one to the general kernel.""" + self._kernel_main(mX, mO_row, mS_row) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mO_row: cute.Tensor, + mS_row: cute.Tensor, + ): + """Device entry for the specialized rowwise-only cast kernel (vectorized global loads/stores, no TMA).""" + bidx, bidy, _ = cute.arch.block_idx() + M = mX.shape[0] + N = mX.shape[1] + + TILER = (self._TILE_ROWS, self._TILE_COLS) + SCALE_TILER = (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + gX_tile = cute.local_tile(mX, TILER, (bidy, bidx)) + gO_tile = cute.local_tile(mO_row, TILER, (bidy, bidx)) + gS_tile = cute.local_tile(mS_row, SCALE_TILER, (bidy, bidx)) + + quantize_mxfp8_gmem_tile_rowwise_1lane( + gX_tile, + gO_tile, + gS_tile, + M, + N, + # Use cute.size instead of .shape[1] because under the swizzled layout mode 1 is + # a nested tuple ((4, num_tiles_SC)), not a plain int. + cute.size(mS_row, mode=[1]), + bidy, + bidx, + 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, + TILE_ROWS=self._TILE_ROWS, + TILE_COLS=self._TILE_COLS, + STASH_SCALE_TO_SMEM=self._STASH_SCALE_TO_SMEM, + ) 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 new file mode 100644 index 00000000000..d0ff6fcbcd3 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_rowwise_2lane.py @@ -0,0 +1,696 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Two-lane GMEM-to-RMEM rowwise MXFP8 cast, ported from cast_rowwise.cuh.""" + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import Boolean, Float32, Int32, Int64, Uint32 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + abs_max_x2_bf16, + bf16_pair_magnitude, + create_l2_policy, + exp2_bf16x2_rcp, + max_num_bits_per_copy_ld_global, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, + cvt_f32x2_to_bf16x2, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + mul_bf16x2x2_cvt_bf16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + activation_op, + dactivation_op, + derive_swizzled_scale_layout, +) + + +@cute.jit +def _flush_scales_to_gmem( + sScale: cute.Tensor, + gS_tile: cute.Tensor, + cta_row: Int32, + cta_col: Int32, + M: Int32, + padded_cols: Int32, + CTA_Y: cutlass.Constexpr[int], + CTA_X: cutlass.Constexpr[int], +): + """Flush each 4-row by 4-column scale block with a 16-byte store. + + sScale and gS_tile must have byte strides (4, 1) within each aligned (4, 4) row-column block. + """ + tidx, _, _ = cute.arch.thread_idx() + active_threads = CTA_X // 4 + _, tv_flush = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_Y // 4, active_threads), stride=(active_threads, 1)), + val_layout=cute.make_layout((4, 4), stride=(4, 1)), + ) + if tidx < (CTA_Y // 4) * active_threads: + thread_y = cta_row * CTA_Y + (tidx // active_threads) * 4 + thread_x = cta_col * CTA_X + (tidx % active_threads) * 4 + if thread_y < M and thread_x < padded_cols: + cute.autovec_copy( + cute.composition(sScale, tv_flush)[tidx, None], + cute.composition(gS_tile, tv_flush)[tidx, None], + ) + + +@cute.jit +def quantize_mxfp8_gmem_tile_rowwise_2lane( + gX_tiles: cute.Tensor, + gA_tiles: Optional[cute.Tensor], + gO_tiles: cute.Tensor, + 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]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + MAX_NORM_RCP: cutlass.Constexpr[float], + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + ELEMENTS_PER_THREAD: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + CHECK_BOUNDS: cutlass.Constexpr[bool], + input_policy: Int64, + output_policy: Int64, + cta_row: Int32, + cta_col: Int32, + num_elems: Int64, +): + """Quantize a CTA's (TILE_ROWS, TILE_COLS, NUM_TILES) GMEM tiles. + + CTA_THREADS_X threads own each tile row. Each thread loads 16 BF16 values, and + two adjacent lanes share the maximum and scale for a 32-value MXFP8 block. + Load all tiles before reducing to preserve independent memory requests. + """ + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * ELEMENTS_PER_THREAD + tidx, _, _ = cute.arch.thread_idx() + + _, tv_payload = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + val_layout=cute.make_layout((1, ELEMENTS_PER_THREAD), stride=(ELEMENTS_PER_THREAD, 1)), + ) + # Mode 0.0: two adjacent threads share one scale (stride 0). + # Mode 0.1: thread pairs are arranged as (scale columns, tile rows). + # Mode 1: each thread addresses one scale element. + tv_scale = cute.make_layout( + ((2, (CTA_THREADS_X // 2, CTA_THREADS_Y)), 1), stride=((0, (TILE_ROWS, 1)), 0) + ) + + # We load fragments from all GMEM tiles at once + rX = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), DTYPE) + # Backward: the activation input, loaded alongside the incoming gradient in rX + rA = cute.make_rmem_tensor((ELEMENTS_PER_THREAD, NUM_TILES), DTYPE) + # We store a fragment of one tile at a time + rO = cute.make_rmem_tensor((ELEMENTS_PER_THREAD,), FP8_DTYPE) + # RMEM flags to indicate whether each tile is out of bounds or not + # With CHECK_BOUNDS=False DCE will remove it so we don't pay extra cost when we don't use it + tile_valid_flags = cute.make_rmem_tensor((NUM_TILES,), Boolean) + + # 16 * 2 bytes = 32 bytes, one 256-bit load where the toolchain supports it + LOAD_WIDTH = min( + ELEMENTS_PER_THREAD * DTYPE.width, + max_num_bits_per_copy_ld_global(), + ) + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=LOAD_WIDTH, + # TODO(kainingz): due to https://github.com/NVIDIA/cutlass/issues/3726 using + # invariant=True here will drop the cache policy + invariant=False, + ) + act_load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=LOAD_WIDTH, + invariant=False, + ) + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + FP8_DTYPE, + num_bits_per_copy=ELEMENTS_PER_THREAD * FP8_DTYPE.width, + ) + + scale_and_pack = mul_bf16x2x2_cvt_bf16x4_to_fp8x4(FP8_DTYPE) + + # Copy all tiles from GMEM to RMEM and potentially check their bounds + for tile_idx in cutlass.range_constexpr(NUM_TILES): + # gX_tiles is (TILE_ROWS, TILE_COLS, NUM_TILES) + gX_tile = gX_tiles[None, None, tile_idx] + gX_thread = cute.composition(gX_tile, tv_payload)[tidx, None] + valid = Boolean(True) + if cutlass.const_expr(CHECK_BOUNDS): + row = cta_row * TILE_ROWS + tidx // CTA_THREADS_X + col = (cta_col * NUM_TILES + tile_idx) * TILE_COLS + ( + tidx % CTA_THREADS_X + ) * ELEMENTS_PER_THREAD + # When we apply the (-1, 1024) manually constructed layout to the input tensor, + # the last row might not be a full row so we need to check OOB with element number + # instead of just M and N + valid = col < N and Int64(row) * N + col < num_elems + tile_valid_flags[tile_idx] = valid + rX[None, tile_idx].fill(0) + if cutlass.const_expr(WITH_DACT): + rA[None, tile_idx].fill(0) + if valid: + cute.copy(load_atom, gX_thread, rX[None, tile_idx], cache_policy=input_policy) + if cutlass.const_expr(WITH_DACT): + gA_thread = cute.composition(gA_tiles[None, None, tile_idx], tv_payload)[tidx, None] + cute.copy(act_load_atom, gA_thread, rA[None, tile_idx], cache_policy=input_policy) + + # Fused activation: widen to FP32, apply it, and round back to BF16 as the general kernel + # does (it truncates the activated value to the input type before the amax). The values are + # then exactly BF16, so the packed BF16 amax and scaling below apply unchanged. + if cutlass.const_expr(WITH_ACT or WITH_DACT): + for tile_idx in cutlass.range_constexpr(NUM_TILES): + rX_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) + rA_i32 = cute.recast_tensor(rA[None, tile_idx], Int32) + for i in cutlass.range_constexpr(ELEMENTS_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 = dactivation_op(ACTIVATION, ACTIVATION_TABLE) + lo = lo * dop(x2_lo_to_f32_bf16(rA_i32[i])) + hi = hi * dop(x2_hi_to_f32_bf16(rA_i32[i])) + else: + op = activation_op(ACTIVATION, ACTIVATION_TABLE) + lo = op(lo) + hi = op(hi) + rX_i32[i] = cvt_f32x2_to_bf16x2(lo, hi) + + @cute.jit + def extract_mx_block_amax(values_i32: cute.Tensor) -> Float32: + """Return the 32-value MXFP8 block amax shared by two adjacent threads. + + Balanced eight-pair tree, followed by one butterfly between lane partners. + """ + level = [ + abs_max_x2_bf16(values_i32[2 * i], values_i32[2 * i + 1]) + for i in range(cute.size(values_i32) // 2) + ] + amax_pair = abs_max_x2_bf16( + abs_max_x2_bf16(level[0], level[2]), abs_max_x2_bf16(level[1], level[3]) + ) + # Compare with the neighbor thread's amax (2 threads cooperate to process a MXFP8 block). + partner = cute.arch.shuffle_sync_bfly(amax_pair, 1) + amax_pair = abs_max_x2_bf16(amax_pair, partner) + return (bf16_pair_magnitude(amax_pair) << Uint32(16)).bitcast(Float32) + + for tile_idx in cutlass.range_constexpr(NUM_TILES): + gO_tile = gO_tiles[None, None, tile_idx] + mS_tile = mS_tiles[None, None, tile_idx] + gO_thread = cute.composition(gO_tile, tv_payload)[tidx, None] + mS_thread = cute.composition(mS_tile, tv_scale)[tidx, None] + # Reinterpret the native tensors only for the packed BF16/FP8 instructions. + rX_tile_i32 = cute.recast_tensor(rX[None, tile_idx], Int32) + rO_u32 = cute.recast_tensor(rO, Uint32) + amax = extract_mx_block_amax(rX_tile_i32) + # With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's + # generic kernel; cast_rowwise.cuh, like the plain cast here, gives it the NaN scale. + if cutlass.const_expr(ACTIVATION is not None): + amax = cute.arch.fmax(amax, Float32(0.0)) + exponent = cvt_f32_to_fp8e8m0fnu(amax * MAX_NORM_RCP) + + # The even lane of each pair stores its block's scale byte, as cast_rowwise.cuh does. + # A warp's stores coalesce in either layout: one 16-byte run when packed, four + # 4-byte runs when GEMM-swizzled. Gathering the bytes into 32-bit words first + # touches the same sectors but adds shuffles and a branch (+12% instructions). + # TODO(kainingz): not using vectorized word store seems to be faster in CuTeDSL. + # Check if we should port this to CUDA + if tidx % 2 == 0 and tile_valid_flags[tile_idx]: + mS_thread[0] = exponent + + # Compute and write the quantized MXFP8 values to GMEM + reciprocal_packed = exp2_bf16x2_rcp(exponent) + for i in cutlass.range_constexpr(cute.size(rO_u32)): + rO_u32[i] = scale_and_pack( + rX_tile_i32[2 * i], reciprocal_packed, rX_tile_i32[2 * i + 1], reciprocal_packed + ) + if tile_valid_flags[tile_idx]: + cute.copy(store_atom, rO, gO_thread, cache_policy=output_policy) + + +class MXFP8QuantizeRowwise2LaneG2RKernel(MXFP8QuantizeKernelBase): + """Two lanes per MXFP8 block, with a CTA sweeping NUM_TILES horizontally. + + One CTA tile has TILE_ROWS = CTA_THREADS_Y and + TILE_COLS = CTA_THREADS_X * _ELEMENTS_PER_THREAD. The CTA + processes NUM_TILES = 1 or 2 horizontally. With staged swizzled scales, + tensor views group rows 32 apart to collect complete scale vectors. + _L2_CACHED_CTA_PERCENT tunes the share of CTAs whose input loads use normal + L2 retention: 0 streams all input and 100 caches all input normally. + """ + + _CTA_THREADS_X = 32 + _ELEMENTS_PER_THREAD = 16 + # cast_rowwise.cuh's tier table: the two smaller tiers stream all of their input. + _L2_CACHED_CTA_PERCENT = 40 + _SMALL_TIER_CACHED_PERCENT = 0 + _MAX_GRID_Y = 65535 + + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # BF16 packed math. + return ( + not cfg.WITH_AMAX + and not cfg.WITH_DBIAS + and cfg.DTYPE is cutlass.BFloat16 + and cfg.ROWWISE + and not cfg.COLWISE + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + ) -> Boolean: + N = mX.shape[1] + act_input_aligned = Boolean(True) + if cutlass.const_expr(mDActInput is not None): + act_input_aligned = mDActInput.iterator.toint() % 32 == 0 + scale_stride_matches = Boolean(True) + if cutlass.const_expr(not cfg.WITH_GEMM_SWIZZLED_SCALES): + # The flat view needs the packed scales to be contiguous. + scale_stride_matches = mS_row.shape[1] == N // 32 + # Column tiles use grid.x and row tiles spill from grid.y into grid.z, so no shape + # exceeds the grid limits. + return ( + N % 128 == 0 + and mX.iterator.toint() % 32 == 0 + and act_input_aligned + and scale_stride_matches + and mNoop.toint() == 0 + ) + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + 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 + num_elems = Int64(M) * N + # View the input tensor as a layout with last dimension = 1024 when not swizzled and the number of elements can be expressed using i32 + # However, M*N might not be divisible by 1024, so the last row of this layout might be a partial row with fewer than 1024 elements, + # which is why we need to pass num_elems to the kernel so OOB checks will identify these regions as invalid offsets correctly + if cutlass.const_expr(not self.cfg.WITH_GEMM_SWIZZLED_SCALES): + if num_elems <= 2**31 - 1: + rows = cute.ceil_div(num_elems, 1024) + flat = cute.make_layout((rows, 1024), stride=(1024, 1)) + mA = None + if cutlass.const_expr(self.cfg.WITH_DACT): + mA = cute.make_tensor(mDActInput.iterator, flat) + self._launch_tiers( + cute.make_tensor(mX.iterator, flat), + mA, + cute.make_tensor(mO_row.iterator, flat), + cute.make_tensor(mS_row.iterator, cute.make_layout((rows, 32), stride=(32, 1))), + stream, + num_elems, + ) + else: + self._launch_tiers(mX, mDActInput, mO_row, mS_row, stream, num_elems) + else: + self._launch_tiers(mX, mDActInput, mO_row, mS_row, stream, num_elems) + + @cute.jit + def _launch_tiers( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO_row: cute.Tensor, + mS_row: cute.Tensor, + stream: CUstream, + num_elems: Int64, + ): + """Pick the launch configuration: the size tiers, except for swizzled scales with a fused + activation when N does not fill the tiers' 1024-column CTAs.""" + if cutlass.const_expr( + self.cfg.WITH_GEMM_SWIZZLED_SCALES and (self.cfg.WITH_ACT or self.cfg.WITH_DACT) + ): + if mX.shape[1] % 1024 != 0: + # The swizzled path tiles the real (M, N) shape, so a 1024-column CTA overhangs + # N and its idle lanes still pay for the activation math. 256-column CTAs divide + # every eligible N (N % 128 == 0 and the tile is 2 x 128 columns). + self._launch( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + CTA_THREADS_Y=16, + CTA_THREADS_X=8, + NUM_TILES=2, + L2_CACHED_CTA_PERCENT=0, + ) + else: + self._launch_size_tiers(mX, mA, mO_row, mS_row, stream, num_elems) + else: + self._launch_size_tiers(mX, mA, mO_row, mS_row, stream, num_elems) + + @cute.jit + def _launch_size_tiers( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO_row: cute.Tensor, + mS_row: cute.Tensor, + stream: CUstream, + num_elems: Int64, + ): + """Pick the launch configuration for this many elements (cast_rowwise.cuh's tiers).""" + output_bytes = num_elems + if output_bytes <= 12 << 20: + self._launch( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + CTA_THREADS_Y=8, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=1, + L2_CACHED_CTA_PERCENT=self._SMALL_TIER_CACHED_PERCENT, + ) + elif output_bytes <= 48 << 20: + self._launch( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + CTA_THREADS_Y=8, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=2, + L2_CACHED_CTA_PERCENT=self._SMALL_TIER_CACHED_PERCENT, + ) + elif output_bytes <= 96 << 20: + self._launch( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + CTA_THREADS_Y=4, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=2, + L2_CACHED_CTA_PERCENT=self._L2_CACHED_CTA_PERCENT, + ) + else: + self._launch( + mX, + mA, + mO_row, + mS_row, + stream, + num_elems, + CTA_THREADS_Y=8, + CTA_THREADS_X=self._CTA_THREADS_X, + NUM_TILES=2, + L2_CACHED_CTA_PERCENT=self._L2_CACHED_CTA_PERCENT, + ) + + @cute.jit + def _launch( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO: cute.Tensor, + mS: cute.Tensor, + stream: CUstream, + num_elems: Int64, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], + ): + assert 0 <= L2_CACHED_CTA_PERCENT <= 100 + assert self._ELEMENTS_PER_THREAD == 16 + assert CTA_THREADS_X % 8 == 0 + assert CTA_THREADS_Y * CTA_THREADS_X % 32 == 0 + M, N = mX.shape + # Each CTA processes NUM_TILES horizontal tiles of (TILE_ROWS, TILE_COLS) values. + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD + THREADS_PER_CTA = CTA_THREADS_Y * CTA_THREADS_X + # Total columns covered by a CTA; its row count stays TILE_ROWS. + CTA_ROWS = TILE_ROWS + CTA_COLS = TILE_COLS * NUM_TILES + + grid_rows = cute.ceil_div(M, CTA_ROWS) + grid_cols = cute.ceil_div(N, CTA_COLS) + check_bounds = M % CTA_ROWS != 0 or N % CTA_COLS != 0 or num_elems != Int64(M) * N + grid = Int64(grid_rows) * grid_cols + first_streaming_cta = grid * L2_CACHED_CTA_PERCENT // 100 + + # Column tiles vary fastest, so consecutive CTAs read consecutive memory along a + # row, as cast_rowwise.cuh's flat grid does. Putting row tiles on x instead + # sends concurrent CTAs down a column strip, one 2 KiB piece per row, and costs + # up to 9% at large shapes. Row tiles past grid.y's limit spill into grid.z. + grid_y = cutlass.min(grid_rows, self._MAX_GRID_Y) + grid_z = cute.ceil_div(grid_rows, self._MAX_GRID_Y) + launch_grid = [grid_cols, grid_y, grid_z] + kernel_args = ( + mX, + mA, + mO, + mS, + first_streaming_cta, + num_elems, + CTA_THREADS_Y, + CTA_THREADS_X, + NUM_TILES, + L2_CACHED_CTA_PERCENT, + ) + launch_args = {"grid": launch_grid, "block": [THREADS_PER_CTA], "stream": stream} + stash_scale_to_smem = False + if cutlass.const_expr( + self.cfg.WITH_GEMM_SWIZZLED_SCALES and CTA_THREADS_Y % 4 == 0 and not self.cfg.WITH_DACT + ): + # Staged scale vectors need complete 128-row bands. + stash_scale_to_smem = M % 128 == 0 + + if grid > 0: + if check_bounds: + if stash_scale_to_smem: + self._launch_kernel(kernel_args, launch_args, True, True) + else: + self._launch_kernel(kernel_args, launch_args, True, False) + else: + if stash_scale_to_smem: + self._launch_kernel(kernel_args, launch_args, False, True) + else: + self._launch_kernel(kernel_args, launch_args, False, False) + + @cute.jit + def _launch_kernel( + self, + kernel_args, + launch_args, + CHECK_BOUNDS: cutlass.Constexpr[bool], + STASH_SCALE_TO_SMEM: cutlass.Constexpr[bool], + ): + self.kernel( + *kernel_args, CHECK_BOUNDS=CHECK_BOUNDS, STASH_SCALE_TO_SMEM=STASH_SCALE_TO_SMEM + ).launch(**launch_args) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO: cute.Tensor, + mS: cute.Tensor, + first_streaming_cta: Int64, + num_elems: Int64, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], + CHECK_BOUNDS: cutlass.Constexpr[bool], + STASH_SCALE_TO_SMEM: cutlass.Constexpr[bool], + ): + """Run the kernel body. Like cast_rowwise.cuh, there is no noop handling: the + dispatcher only selects this kernel when no cast_noop flag is passed.""" + self._kernel_main( + mX, + mA, + mO, + mS, + first_streaming_cta, + num_elems, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + NUM_TILES=NUM_TILES, + L2_CACHED_CTA_PERCENT=L2_CACHED_CTA_PERCENT, + CHECK_BOUNDS=CHECK_BOUNDS, + STASH_SCALE_TO_SMEM=STASH_SCALE_TO_SMEM, + ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mA: Optional[cute.Tensor], + mO: cute.Tensor, + mS: cute.Tensor, + first_streaming_cta: Int64, + num_elems: Int64, + CTA_THREADS_Y: cutlass.Constexpr[int], + CTA_THREADS_X: cutlass.Constexpr[int], + NUM_TILES: cutlass.Constexpr[int], + L2_CACHED_CTA_PERCENT: cutlass.Constexpr[int], + CHECK_BOUNDS: cutlass.Constexpr[bool], + STASH_SCALE_TO_SMEM: cutlass.Constexpr[bool], + ): + """Construct CTA tile views and run the specialized quantization.""" + # Dispatch checks 32-byte input alignment before selecting this kernel. + gX = cute.make_tensor(mX.iterator.align(32), mX.layout) + if cutlass.const_expr(self.cfg.WITH_DACT): + mA = cute.make_tensor(mA.iterator.align(32), mA.layout) + # grid.x walks column tiles; grid.y and grid.z together walk row tiles. + bidx, bidy, bidz = cute.arch.block_idx() + cta_row = bidz * self._MAX_GRID_Y + bidy + cta_coord = (cta_row, bidx) + M, N = mX.shape + TILE_ROWS = CTA_THREADS_Y + TILE_COLS = CTA_THREADS_X * self._ELEMENTS_PER_THREAD + CTA_ROWS = TILE_ROWS + CTA_COLS = TILE_COLS * NUM_TILES + + gS = mS + # Use the swizzled layout as is. Composing it with (M, N // 32) layouts to flatten + # its shape made CuTe emit runtime integer div/mod on every scale address, + # about 7x this kernel's instruction count. + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + gS, _ = derive_swizzled_scale_layout(M, N, True, False, mS, None) + + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + # Apply the same permutation to input, output and scale so it becomes transparent and + # other code is unaware of this permutation pattern + # However, this this permutation, the scale bytes emitted by a CTA are all in this CTA's scale + # tile partition instead of scatter everywhere, so we can vectorize the flush to GMEM + # For example, with a 128x128 input: + # gX & mO: (128,128):(128,1) -> ((4,32),128):((4096,128),1) + # gS: ((32,4,1),(4,1)):((16,4,512),(1,512)) -> ((4,32), 4):((4, 16),1) + row_perm = cute.make_layout((4, 32, M // 128), stride=(32, 1, 128)) + gX = cute.composition(gX, (row_perm, None)) + mO = cute.composition(mO, (row_perm, None)) + gS = cute.composition(gS, (row_perm, None)) + + CTA_TILER = (CTA_ROWS, CTA_COLS) + CTA_SCALE_TILER = (CTA_ROWS, CTA_COLS // 32) + gX_cta = cute.local_tile(gX, CTA_TILER, cta_coord) + gO_cta = cute.local_tile(mO, CTA_TILER, cta_coord) + gS_cta = cute.local_tile(gS, CTA_SCALE_TILER, cta_coord) + + fraction = Float32(1.0) + if cutlass.const_expr(L2_CACHED_CTA_PERCENT > 0): + fraction = Float32(0.0) + # A linear index is needed only for the L2 policy split. It follows launch + # order, so the cached CTAs are the first ones scheduled. + gdx, _, _ = cute.arch.grid_dim() + cta_id = Int64(cta_row) * gdx + bidx + if cta_id >= first_streaming_cta: + fraction = Float32(1.0) + input_policy = create_l2_policy(False, fraction) + output_policy = create_l2_policy(True) + + # Divide CTA_TILER to multiple TILER shape TILES + TILER = (TILE_ROWS, TILE_COLS) + SCALE_TILER = (TILE_ROWS, TILE_COLS // 32) + # There is only one tile on the row dimension because CTA_ROWS == TILE_ROWS. + gX_tiles = cute.local_tile(gX_cta, TILER, (0, None)) + gA_tiles = None + if cutlass.const_expr(self.cfg.WITH_DACT): + gA_tiles = cute.local_tile(cute.local_tile(mA, CTA_TILER, cta_coord), TILER, (0, None)) + gO_tiles = cute.local_tile(gO_cta, TILER, (0, None)) + + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + # Preserve gS_cta's (4 rows, 4 columns):(4, 1) block ordering for the flush. + sScale_scratch = cutlass.memory.SmemAllocator().allocate_tensor( + mS.element_type, cute.make_layout_like(gS_cta), byte_alignment=16 + ) + # The tile quantizer writes scales to SMEM for the cooperative vectorized flush. + mS_tiles = cute.local_tile(sScale_scratch, SCALE_TILER, (0, None)) + else: + # The tile quantizer writes scales directly to GMEM. + mS_tiles = cute.local_tile(gS_cta, SCALE_TILER, (0, None)) + + # The last grid.z slice can overshoot the row tiles; those CTAs have no work. + if cta_row * CTA_ROWS < M: + quantize_mxfp8_gmem_tile_rowwise_2lane( + gX_tiles, + gA_tiles, + gO_tiles, + 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, + FP8_DTYPE=self.cfg.FP8_DTYPE, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, + CTA_THREADS_Y=CTA_THREADS_Y, + CTA_THREADS_X=CTA_THREADS_X, + ELEMENTS_PER_THREAD=self._ELEMENTS_PER_THREAD, + NUM_TILES=NUM_TILES, + CHECK_BOUNDS=CHECK_BOUNDS, + input_policy=input_policy, + output_policy=output_policy, + cta_row=cta_coord[0], + cta_col=cta_coord[1], + num_elems=num_elems, + ) + + if cutlass.const_expr(STASH_SCALE_TO_SMEM): + cute.arch.sync_threads() + _flush_scales_to_gmem( + sScale_scratch, + cute.make_tensor(gS_cta.iterator.align(16), gS_cta.layout), + cta_row, + bidx, + M, + N // 32, + CTA_ROWS, + CTA_COLS // 32, + ) 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 new file mode 100644 index 00000000000..1dfaf911c7e --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_bidim.py @@ -0,0 +1,738 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Specialized CuTeDSL bidimensional MXFP8 cast with staged TMA I/O.""" + +# Local @cute.struct classes are shared-memory storage descriptors. +# pylint: disable=missing-class-docstring + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Boolean, Float32, Float8E8M0FNU, Int32, Int64, Uint32 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + is_packed16, + abs_max_x2_bf16, + abs_max_x2_f16, + x2_lo_to_f32_bf16, + x2_lo_to_f32_f16, + x2_hi_to_f32_bf16, + x2_hi_to_f32_f16, + fabs_f32, + exp2f_rcp, + pack_f32x2, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + cvt_f32_to_fp8e8m0fnu, + mul_f32x2_cvt_f32x4_to_fp8x4, + mul_f32x4_cvt_f32x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + MXFP8_BLOCK_SCALING_SIZE, + THREADS_PER_WARP, + derive_swizzled_scale_layout, +) + + +@cute.jit +def quantize_mxfp8_smem_tile_bidimensional( + sX_tile: cute.Tensor, # 32x64 input tile (already sliced to this stage); any SMEM layout, + # Sw<3,4,3> avoids bank conflicts + sO_row_tile: cute.Tensor, # 32x64 rowwise-output tile + sO_col_tile: cute.Tensor, # 32x64 colwise-output tile + sS_row_tile: cute.Tensor, # (32, 2) smem rowwise-scale staging tile (flushed in the epilogue) + sS_col_tile: cute.Tensor, # (1, 64) smem colwise-scale staging tile (flushed in the epilogue) + sColReduce_warp: cute.Tensor, # (32,) SMEM fp32 columnwise scale reduction buffer for this warp + WARPS_PER_CTA: cutlass.Constexpr[int], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], +): + """Quantize a pre-sliced 32x64 tile -> both rowwise and colwise MXFP8. Elements are + addressed through TV layouts: a warp = 32 lanes = 32 rows (lane == row), and each + thread owns one 32-col row segment (its 32-col block) and one output column. + + No bounds check: the caller only loops over valid tiles (`num_tiles`), and + M % 32 == 0 means a tile can only be partial along N. The partial tile's + out-of-bounds warp still runs, but harmlessly -- it reads TMA zero-filled smem, + its output columns are masked by the caller's TMA store, and its scale writes + (rowwise and colwise alike) go to staging slots whose flush targets are past-N + padding columns of the respective scale tensors. + """ + mul_cvt4 = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE) + mul_cvt4_elemwise = mul_f32x4_cvt_f32x4_to_fp8x4(FP8_DTYPE) + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) + val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE)), + ) + _, tv_layout_rowwise_scale = cute.make_layout_tv( + thr_layout=cute.make_layout(((MXFP8_BLOCK_SCALING_SIZE, 1), WARPS_PER_CTA)), # ((32, 1) 2) + val_layout=cute.make_layout((1, 1)), + ) + _, tv_layout_colwise_scale = cute.make_layout_tv( + thr_layout=cute.make_layout((1, (MXFP8_BLOCK_SCALING_SIZE, WARPS_PER_CTA))), # (1, (32, 2)) + val_layout=cute.make_layout((1, 1)), + ) + + tidx, _, _ = cute.arch.thread_idx() + lane = tidx % MXFP8_BLOCK_SCALING_SIZE + + # Each composed [tidx, None] slice is 1-D (the value mode is flattened): the data + # slices are size 32 (this thread's row segment), the scale slices size 1. + tXsX = cute.composition(sX_tile, tv_layout)[tidx, None] # (32,) input row segment + tXsO_row = cute.composition(sO_row_tile, tv_layout)[tidx, None] # (32,) rowwise out + tXsO_col = cute.composition(sO_col_tile, tv_layout)[tidx, None] # (32,) colwise out + tSsS_row_tile = cute.composition(sS_row_tile, tv_layout_rowwise_scale)[tidx, None] # (1,) + tSsS_col_tile = cute.composition(sS_col_tile, tv_layout_colwise_scale)[tidx, None] # (1,) + + rO_row = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) + rO_col = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, FP8_DTYPE) + rO_row_u32 = cute.make_tensor( + cute.recast_ptr(rO_row.iterator, dtype=Uint32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), + ) + rO_col_u32 = cute.make_tensor( + cute.recast_ptr(rO_col.iterator, dtype=Uint32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,)), + ) + sColReduce = cute.make_tensor( + sColReduce_warp.iterator, + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 4, 4), stride=(4, 1)), + ) + + if cutlass.const_expr(is_packed16(DTYPE)): + # If the input is bf16 / fp16, take this fast path and process 2 elements at a time in a packed i32 + abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 + x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 + x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 + rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, DTYPE) + # Do a vectorized load from SMEM to RMEM and unswizzle in the meantime. + cute.autovec_copy(tXsX, rX) + rX_2x = cute.make_tensor( + cute.recast_ptr(rX.iterator, dtype=Int32), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + sColReduce_2x = cute.make_tensor( + cute.recast_ptr(sColReduce.iterator, dtype=Int64), + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE // 2,), stride=(1,)), + ) + + row_amax2 = rX_2x[0] + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): + pair = rX_2x[i] + # Skip the first iteration because we assigned rX_2x[0] to row_amax2 already + if cutlass.const_expr(i > 0): + row_amax2 = abs_max_x2(row_amax2, pair) + a_lo = fabs_f32(x2_lo_to_f32(pair)) + a_hi = fabs_f32(x2_hi_to_f32(pair)) + col_lo = cute.arch.warp_redux_sync(a_lo, kind="fmax") + col_hi = cute.arch.warp_redux_sync(a_hi, kind="fmax") + with cute.arch.elect_one(): + sColReduce_2x[i] = pack_f32x2(col_lo, col_hi) + + # Compute the rowwise scale factor + row_amax = cute.arch.fmax( + fabs_f32(x2_lo_to_f32(row_amax2)), fabs_f32(x2_hi_to_f32(row_amax2)) + ) + row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) + row_inv = exp2f_rcp(row_exp) + tSsS_row_tile[0] = row_exp + cute.arch.sync_warp() + + # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) + col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) + tSsS_col_tile[0] = col_exp + sColReduce_warp[lane] = exp2f_rcp(col_exp) + cute.arch.sync_warp() + + row_scale_2x = pack_f32x2(row_inv, row_inv) + # Vectorized multiply-and-convert: 4 f32 → 4 fp8 + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + # Vectorized load for 4 columnwise scale factors from SMEM to RMEM + col_inv4 = cute.make_rmem_tensor(4, Float32) + cute.autovec_copy(sColReduce[j, None], col_inv4) # LDS.128, warp-broadcast + p01 = rX_2x[2 * j] + p23 = rX_2x[2 * j + 1] + f0 = x2_lo_to_f32(p01) + f1 = x2_hi_to_f32(p01) + f2 = x2_lo_to_f32(p23) + f3 = x2_hi_to_f32(p23) + # For rowwise quantized values, they use the same scale for all 4 elements, + # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. + rO_row_u32[j] = mul_cvt4(f0, f1, f2, f3, row_scale_2x) + # For columnwise quantized values, each element has its own scale; the + # elementwise variant fuses the per-element multiply into the cvt sequence. + rO_col_u32[j] = mul_cvt4_elemwise( + f0, f1, f2, f3, col_inv4[0], col_inv4[1], col_inv4[2], col_inv4[3] + ) + cute.autovec_copy(rO_row, tXsO_row) + cute.autovec_copy(rO_col, tXsO_col) + else: + # If input is fp32, take this slow path and process it normally without packing + rX = cute.make_rmem_tensor(MXFP8_BLOCK_SCALING_SIZE, Float32) + cute.autovec_copy(tXsX, rX) + + row_amax = Float32(0.0) + for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + a = fabs_f32(rX[c]) + row_amax = cute.arch.fmax(row_amax, a) + col_amax = cute.arch.warp_redux_sync(a, kind="fmax") + with cute.arch.elect_one(): + sColReduce_warp[c] = col_amax + + # Compute the rowwise scale factor + row_exp = cvt_f32_to_fp8e8m0fnu(row_amax * MAX_NORM_RCP) + row_inv = exp2f_rcp(row_exp) + tSsS_row_tile[0] = row_exp # rowwise scale (this row-block, staged in smem) + cute.arch.sync_warp() + + # Compute the colwise scale factor (only handle the one that belongs to this thread / lane) + col_exp = cvt_f32_to_fp8e8m0fnu(sColReduce_warp[lane] * MAX_NORM_RCP) + tSsS_col_tile[0] = col_exp # colwise scale (this thread's column) + sColReduce_warp[lane] = exp2f_rcp(col_exp) + cute.arch.sync_warp() + + row_scale_2x = pack_f32x2(row_inv, row_inv) + # Vectorized multiply-and-convert: 4 f32 → 4 fp8 + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 4): + # Vectorized load for 4 columnwise scale factors from SMEM to RMEM + col_inv4 = cute.make_rmem_tensor(4, Float32) + cute.autovec_copy(sColReduce[j, None], col_inv4) + offset = 4 * j + # For rowwise quantized values, they use the same scale for all 4 elements, + # so we can just pass two row_inv to mul_cvt4 to apply it to all 4 elements at once. + rO_row_u32[j] = mul_cvt4( + rX[offset], rX[offset + 1], rX[offset + 2], rX[offset + 3], row_scale_2x + ) + # For columnwise quantized values, each element has its own scale; the + # elementwise variant fuses the per-element multiply into the cvt sequence. + rO_col_u32[j] = mul_cvt4_elemwise( + rX[offset], + rX[offset + 1], + rX[offset + 2], + rX[offset + 3], + col_inv4[0], + col_inv4[1], + col_inv4[2], + col_inv4[3], + ) + cute.autovec_copy(rO_row, tXsO_row) + cute.autovec_copy(rO_col, tXsO_col) + + +class MXFP8QuantizeBidimensionalG2SKernel(MXFP8QuantizeKernelBase): + """Specialized cast-only rowwise+colwise MXFP8 kernel (swizzled TMA, one 32x32 tile per warp).""" + + _WARPS_PER_CTA = 2 + _THREADS_PER_CTA = _WARPS_PER_CTA * THREADS_PER_WARP + # A warp handles a 32x32 tile + _WARP_ROWS = MXFP8_BLOCK_SCALING_SIZE + _WARP_COLS = MXFP8_BLOCK_SCALING_SIZE + # A CTA handles a tile consisted two 32x32 warp subtiles side-by-side at a time, each handled by a warp + _TILE_ROWS = _WARP_ROWS + _TILE_COLS = _WARP_COLS * _WARPS_PER_CTA + _NUM_STAGES = 2 + _NUM_TILES_X = 4 + _NUM_TILES_Y = 1 + _NUM_TILES = _NUM_TILES_X * _NUM_TILES_Y + + # Rows and columns for the rowwise / columnwise scale factor tensor + _SCALE_ROWS = _TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE + _SCALE_COLS = _TILE_COLS // MXFP8_BLOCK_SCALING_SIZE + + @staticmethod + def supports(cfg: MXFP8QuantizeConfig) -> bool: + # This kernel supports swizzled scales but we don't dispatch to it for now. + return ( + cfg.is_plain_cast() + and is_packed16(cfg.DTYPE) + and cfg.ROWWISE + and cfg.COLWISE + and not cfg.WITH_GEMM_SWIZZLED_SCALES + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + cfg: cutlass.Constexpr[MXFP8QuantizeConfig], + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mNoop: cute.Pointer, + mDActInput: Optional[cute.Tensor], + ) -> Boolean: + cls = MXFP8QuantizeBidimensionalG2SKernel + rows_per_cta = cls._TILE_ROWS * cls._NUM_TILES_Y # packed scales: one tile row + return cute.ceil_div(mX.shape[0], rows_per_cta) <= 65535 and mNoop.toint() == 0 + + def __init__(self, cfg: MXFP8QuantizeConfig): + self.cfg = cfg + if cfg.WITH_GEMM_SWIZZLED_SCALES: + self._NUM_TILES_X = 2 + self._NUM_TILES_Y = 2 + self._NUM_TILES = self._NUM_TILES_X * self._NUM_TILES_Y + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + 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): + cute.printf( + "[CuTeDSL] MXFP8QuantizeBidimensionalG2SKernel.__call__() with config:" + f" {self.cfg}\n" + ) + M = mX.shape[0] + N = mX.shape[1] + + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, mS_col = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col + ) + + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + # Apply 128B input Swizzle<3,4,3> to input tiles + in_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, smem_tile_layout) + # Apply 64B output swizzle<2,4,3> for output tiles + out_smem_layout = cute.make_composed_layout(cute.make_swizzle(2, 4, 3), 0, smem_tile_layout) + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + # Input TMA atom + tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, + mX, + in_smem_layout, + cta_tiler, + num_multicast=1, + ) + # Rowwise output TMA atoms + tma_atom_out_row, tma_dst_out_row = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_row, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + # Colwise output TMA atoms + tma_atom_out_col, tma_dst_out_col = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_col, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + + grid = [ + # Each CTA covers a (TILE_ROWS * NUM_TILES_Y, TILE_COLS * NUM_TILES_X) area. + cute.ceil_div(Int32(N), self._TILE_COLS * self._NUM_TILES_X), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES_Y), + ] + block = [self._THREADS_PER_CTA] + self.kernel( + mX, + mS_row, + mS_col, + mX.element_type, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mS_row: cute.Tensor, + mS_col: cute.Tensor, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, + tma_atom_out_row: cute.CopyAtom, + tma_dst_out_row: cute.Tensor, + tma_atom_out_col: cute.CopyAtom, + tma_dst_out_col: cute.Tensor, + ): + """Device entry. Like the CUDA specialized cast kernels it does not consume the + cast_noop flag: the dispatcher routes calls with one to the general kernel.""" + self._kernel_main( + mX, + mS_row, + mS_col, + dtype, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mS_row: cute.Tensor, + mS_col: cute.Tensor, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, + tma_atom_out_row: cute.CopyAtom, + tma_dst_out_row: cute.Tensor, + tma_atom_out_col: cute.CopyAtom, + tma_dst_out_col: cute.Tensor, + ): + """Device entry for the specialized bidimensional (rowwise+colwise) cast kernel.""" + + bidx, bidy, _ = cute.arch.block_idx() + + M = mX.shape[0] + N = mX.shape[1] + + # The CTA owns a NUM_TILES_Y x NUM_TILES_X arrangement of tiles, starting here. + tile_x_base = bidx * self._NUM_TILES_X + tile_y_base = bidy * self._NUM_TILES_Y + + num_tiles_x = cutlass.min( + self._NUM_TILES_X, + # Valid 64-col tiles remaining from this CTA's start column (bidx * span). + cute.ceil_div(Int32(N) - bidx * self._TILE_COLS * self._NUM_TILES_X, self._TILE_COLS), + ) + num_tiles_y = cutlass.min( + self._NUM_TILES_Y, + # Valid 32-row tiles remaining from this CTA's start row (bidy * span). + cute.ceil_div(Int32(M) - bidy * self._TILE_ROWS * self._NUM_TILES_Y, self._TILE_ROWS), + ) + num_tiles = num_tiles_x * num_tiles_y + + FP8_DTYPE = self.cfg.FP8_DTYPE + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[ + cute.Int64, 2 * self._NUM_STAGES + ] # (full,empty) per stage + sX: cute.struct.Align[ + cute.struct.MemRange[dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + # Per-warp colwise-reduce scratchpad + sColReduce: cute.struct.Align[ + cute.struct.MemRange[Float32, THREADS_PER_WARP * self._WARPS_PER_CTA], 16 + ] + # Staged rowwise scales for the whole CTA + sScaleRow: cute.struct.Align[ + cute.struct.MemRange[ + Float8E8M0FNU, + self._TILE_ROWS * self._NUM_TILES * self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, + ], + 16, + ] + # Staged colwise scales for the whole CTA + sScaleCol: cute.struct.Align[ + cute.struct.MemRange[Float8E8M0FNU, self._NUM_TILES * self._TILE_COLS], 16 + ] + + smem = cutlass.memory.SmemAllocator() + storage = smem.allocate(SharedStorage) + + tile_layout = cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + # SMEM input tile should have the same Swizzle<3,4,3> as the input TMA atom + sX = storage.sX.get_tensor(tile_layout, swizzle=cute.make_swizzle(3, 4, 3)) + # SMEM output tiles should have the same Swizzle<2,4,3> as the output TMA atom + sO_row = storage.sO_row.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) + sO_col = storage.sO_col.get_tensor(tile_layout, swizzle=cute.make_swizzle(2, 4, 3)) + # Reshape the per-warp colwise-reduce scratchpad to a (threads, warps) layout for easier indexing. + sColReduce = storage.sColReduce.get_tensor( + cute.make_layout((THREADS_PER_WARP, self._WARPS_PER_CTA), stride=(1, THREADS_PER_WARP)) + ) + # The whole rowwise scale factor SMEM tensor covered by a CTA + sScaleRow2D = storage.sScaleRow.get_tensor( + cute.make_ordered_layout( + ( + self._TILE_ROWS * self._NUM_TILES_Y, + self._SCALE_COLS * self._NUM_TILES_X, + ), + order=(1, 0), + ) + ) + # Divide by a single tile's shape, so it becomes ((TILE_ROWS, SCALE_COLS), TILES) for easier indexing + # where TILES = NUM_TILES_X * NUM_TILES_Y + sScaleRow = cute.zipped_divide(sScaleRow2D, (self._TILE_ROWS, self._SCALE_COLS)) + # The whole colwise scale factor SMEM tensor covered by a CTA + sScaleCol2D = storage.sScaleCol.get_tensor( + cute.make_ordered_layout( + ( + self._SCALE_ROWS * self._NUM_TILES_Y, + self._TILE_COLS * self._NUM_TILES_X, + ), + order=(1, 0), + ) + ) + # Divide by a single tile's shape, so it becomes ((SCALE_ROWS, TILE_COLS), TILES) for easier indexing + # where TILES = NUM_TILES_X * NUM_TILES_Y + sScaleCol = cute.zipped_divide(sScaleCol2D, (self._SCALE_ROWS, self._TILE_COLS)) + + # Zero the scale SMEM buffers so partial tiles / padding columns flush as 0. + tidx, _, _ = cute.arch.thread_idx() + # View rowwise staging buffers as flat uint32 and stride the CTA over them. + _ROW_SCALE_WORDS = self._TILE_ROWS * self._SCALE_COLS * self._NUM_TILES // 4 + sScaleRow_u32 = cute.make_tensor( + cute.recast_ptr(sScaleRow2D.iterator, dtype=Uint32), + cute.make_layout((_ROW_SCALE_WORDS,), stride=(1,)), + ) + for i in cutlass.range_constexpr(cute.ceil_div(_ROW_SCALE_WORDS, self._THREADS_PER_CTA)): + slot = i * self._THREADS_PER_CTA + tidx + if slot < _ROW_SCALE_WORDS: + sScaleRow_u32[slot] = Uint32(0) + # View colwise staging buffers as flat uint32 and stride the CTA over them. + _COL_SCALE_WORDS = self._SCALE_ROWS * self._TILE_COLS * self._NUM_TILES // 4 + sScaleCol_u32 = cute.make_tensor( + cute.recast_ptr(sScaleCol2D.iterator, dtype=Uint32), + cute.make_layout((_COL_SCALE_WORDS,), stride=(1,)), + ) + for i in cutlass.range_constexpr(cute.ceil_div(_COL_SCALE_WORDS, self._THREADS_PER_CTA)): + slot = i * self._THREADS_PER_CTA + tidx + if slot < _COL_SCALE_WORDS: + sScaleCol_u32[slot] = Uint32(0) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) + + # Only warp 0 is the producer (issues TMA) + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + # Every warp is the consumer + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._WARPS_PER_CTA) + # One TMA loads a tile + tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 + + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_storage.data_ptr(), + num_stages=self._NUM_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + tx_count=tx_count, + cta_layout_vmnk=None, + ) + + prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._NUM_STAGES + ) + cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._NUM_STAGES + ) + + gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) + tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( + tma_atom, 0, cute.make_layout(1), sX, gX_tiled + ) + gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_row, tXgO_row = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_row, 0, cute.make_layout(1), sO_row, gO_row_tiled + ) + gO_col_tiled = cute.zipped_divide(tma_dst_out_col, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_col, tXgO_col = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_col, 0, cute.make_layout(1), sO_col, gO_col_tiled + ) + + cute.arch.sync_threads() + + # Prologue: warp 0 prefetches up to NUM_STAGES tiles to fully fill the pipeline + if warp_idx == 0: + for s in cutlass.range_constexpr(self._NUM_STAGES): + if s < num_tiles: + # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. + tile_y = s // num_tiles_x + tile_x = s % num_tiles_x + mainloop_pipeline.producer_acquire(prod_state) + cute.copy( + tma_atom, + tXgX[(None, (tile_y_base + tile_y, tile_x_base + tile_x))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # 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): + mainloop_pipeline.consumer_wait(cons_state) + # Only allow at most _NUM_STAGES-1 stages to be in-flight, because this iteration will reuse the ring buffer + # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(self._NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + # The current pipeline stage index, which is the tile index modulo the number of stages. + # This is used to index into the shared memory ring buffers that are wrapped around the number of stages. + stage_idx = cons_state.index + + # This tile's position in the CTA's NUM_TILES_Y x NUM_TILES_X arrangement. + tile_y = tile_idx // num_tiles_x + tile_x = tile_idx % num_tiles_x + + # Process the 32x64 SMEM tile for thiHs stage + quantize_mxfp8_smem_tile_bidimensional( + sX[None, stage_idx], + sO_row[None, stage_idx], + sO_col[None, stage_idx], + sScaleRow[(None, (tile_y, tile_x))], + sScaleCol[(None, (tile_y, tile_x))], + sColReduce[ + None, warp_idx + ], # Pick the per-warp colwise-reduce scratchpad for this warp + self._WARPS_PER_CTA, + self.cfg.MAX_NORM_RCP, + dtype, + self.cfg.FP8_DTYPE, + ) + + # Make the smem output writes visible to the TMA async proxy, then store. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + + # We are done with this input pipeline SMEM buffer, signal the producer that it can write to this buffer + mainloop_pipeline.consumer_release(cons_state) + + # Issue TMA and write the output to GMEM + if warp_idx == 0: + tile_coord = (tile_y_base + tile_y, tile_x_base + tile_x) + cute.copy( + tma_atom_out_row, tXsO_row[(None, stage_idx)], tXgO_row[(None, tile_coord)] + ) + cute.copy( + tma_atom_out_col, tXsO_col[(None, stage_idx)], tXgO_col[(None, tile_coord)] + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # Producer: refill the buffer the consumer just freed with the tile + # NUM_STAGES ahead. + next_tile_idx = tile_idx + self._NUM_STAGES + if next_tile_idx < num_tiles: + if warp_idx == 0: + next_tile_y = next_tile_idx // num_tiles_x + next_tile_x = next_tile_idx % num_tiles_x + mainloop_pipeline.producer_acquire(prod_state) + cute.copy( + tma_atom, + tXgX[(None, (tile_y_base + next_tile_y, tile_x_base + next_tile_x))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # TODO(kainingz): the rowwise WIDTH is hardcoded to 4. Consider using other width if we can convince + # cute.autovec_copy that wider stores are safe (e.g. 16B) and the runtime row pitch allows it. This might + # require us to change `scale_rowwise_shape`'s divisibility when we compile with fake tensors + self._flush_scales_to_gmem( + sScaleRow2D, + mS_row, + tidx, + bidx, + bidy, + ROWS=self._TILE_ROWS * self._NUM_TILES_Y, + COLS=self._SCALE_COLS * self._NUM_TILES_X, + WIDTH=4, + ) + self._flush_scales_to_gmem( + sScaleCol2D, + mS_col, + tidx, + bidx, + bidy, + ROWS=self._SCALE_ROWS * self._NUM_TILES_Y, + COLS=self._TILE_COLS * self._NUM_TILES_X, + WIDTH=1 if self.cfg.WITH_GEMM_SWIZZLED_SCALES else 16, + ) + + # Drain the final stores (their gmem writes must complete before the CTA exits). + cute.arch.cp_async_bulk_wait_group(0, read=False) + + @cute.jit + def _flush_scales_to_gmem( + self, + sScale2D: cute.Tensor, + mS: cute.Tensor, + tidx: Int32, + bidx: Int32, + bidy: Int32, + ROWS: cutlass.Constexpr[int], + COLS: cutlass.Constexpr[int], + WIDTH: cutlass.Constexpr[int], + ): + """Flush a staged (ROWS, COLS) SMEM scale block (a plain row-major 2D tensor) to its (bidy, bidx) slice of the gmem scale tensor, + where each GMEM slice has the same shape as this SMEM tile. Use `WIDTH` bytes per vectorized store. + """ + # Use cute.size instead of cute.shape because under the swizzled layout each mode is a nested tuple ((32, 4, num_tiles_M) etc.) + mS_M = cute.size(mS, mode=[0]) + mS_N = cute.size(mS, mode=[1]) + # Obtain the GMEM slice for the output scale factor block of this CTA + mS_tile = cute.local_tile(mS, (ROWS, COLS), (bidy, bidx)) + + ACTIVE_THREAD_COLS = COLS // WIDTH + _, tv_flush_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((ROWS, ACTIVE_THREAD_COLS), stride=(ACTIVE_THREAD_COLS, 1)), + val_layout=cute.make_layout((1, WIDTH), stride=(WIDTH, 1)), + ) + TOTAL_ACTIVE = ROWS * ACTIVE_THREAD_COLS + + # We may have more total active threads than threads in a CTA, so we do multiple waves of stores to flush the whole buffer + for wave in cutlass.range_constexpr(cute.ceil_div(TOTAL_ACTIVE, self._THREADS_PER_CTA)): + thread_idx = wave * self._THREADS_PER_CTA + tidx + if thread_idx < TOTAL_ACTIVE: + # Find the position for this slot's vectorized store in the GMEM scale factor buffer + thread_y = bidy * ROWS + thread_idx // ACTIVE_THREAD_COLS + thread_x = bidx * COLS + (thread_idx % ACTIVE_THREAD_COLS) * WIDTH + # For rowwise we have N divisible by 4 and WIDTH=4, and for colwise we have N divisible by 128 and WIDTH=16, + # so `thread_x < mS_N` with vectorized store is safe here. + # A thread only writes to a single row so `thread_y < mS_M` is also safe here. + if thread_y < mS_M and thread_x < mS_N: + cute.autovec_copy( + cute.composition(sScale2D, tv_flush_layout)[thread_idx, None], + cute.composition(mS_tile, tv_flush_layout)[thread_idx, None], + ) 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 new file mode 100644 index 00000000000..9ea619b06bd --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_g2s_general.py @@ -0,0 +1,1442 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP8 quantization kernel implemented in CuTeDSL. + +Replicates the core logic of quantize_mxfp8.cuh: given a 2D tensor of BF16/FP16 +values, quantize to MXFP8 format (FP8E4M3 data + E8M0 per-block scales). + +This is the general kernel: it handles every config. MXFP8QuantizeEntry in +quantize_mxfp8_common.py dispatches plain casts to the specialized kernels when they +apply and falls back to this kernel otherwise. + +Matches the C++ kernel's tile dimensions and thread layout: + CHUNK_DIM_Y = 64, CHUNK_DIM_X = 64, THREADS_PER_CTA = 64 + BUFF_DIM_Y = 32, BUFF_DIM_X = 64, STAGES = 2 + MXFP8_BLOCK_SCALING_SIZE = 32 (elements per MXFP8 scaling block) + +Grid: (ceil(N / 64), ceil(M / 64)) +Each block processes a 64x64 chunk in 2 stages of 32x64 tiles loaded into +shared memory. +""" + +# Local @cute.struct classes are SMEM-layout descriptors that need no docstrings. +# pylint: disable=missing-class-docstring + +from typing import Optional, Type + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import 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 ( + is_packed16, + abs_max_x2_bf16, + abs_max_x2_f16, + max_x2_bf16, + max_x2_f16, + max_scalar_bf16, + max_scalar_f16, + abs_max_scalar_bf16, + abs_max_scalar_f16, + to_f32_bf16, + to_f32_f16, + x2_lo_to_f32_bf16, + x2_lo_to_f32_f16, + x2_hi_to_f32_bf16, + x2_hi_to_f32_f16, + truncate_f32_bf16, + truncate_f32_f16, + fabs_f32, + exp2f_rcp, + pack_f32x2, + unpack_i64_to_i32x2, +) +from transformer_engine.common.CuTeDSL.utils_fp8 import ( + get_cvt_f32x2_to_fp8x2_func, + cvt_f32_to_fp8e8m0fnu, + mul_f32x2_cvt_f32x4_to_fp8x4, + mul_f32x2_cvt_packed16x4_to_fp8x4, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8QuantizeConfig, + MXFP8QuantizeKernelBase, + MXFP8_BLOCK_SCALING_SIZE, + activation_op, + dactivation_op, + SYM_N_DIVISIBILITY, + THREADS_PER_WARP, + derive_swizzled_scale_layout, + noop_flag_is_set, +) + + +@cute.jit +def _quantize_rowwise_mxfp8( + sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA + sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) + sO_row_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) + mS_row_stage: cute.Tensor, # rowwise scale tensor (1D swizzled, or 2D linear) + MAX_NORM_RCP: cutlass.Constexpr[float], + tile_row_start: Int32, # global row index of this stage's row 0 + tile_col_start: Int32, # global col index of this CTA's col 0 + 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], + TILE_Y: cutlass.Constexpr[int], + SKIP_INPUT_MASKING: cutlass.Constexpr[bool] = False, + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, + WITH_ACT: cutlass.Constexpr[bool] = False, + WITH_DACT: cutlass.Constexpr[bool] = False, + # An activation is fused, even if this pass reads it cached (ACTIVATION is then None). + FUSED: cutlass.Constexpr[bool] = False, + WITH_DBIAS: cutlass.Constexpr[bool] = False, + dbias_acc: Optional[cute.Tensor] = None, # only needed when WITH_DBIAS is True +): + """Quantize one SMEM tile rowwise to MXFP8 (per-row 32-elt block scales); returns the tile amax.""" + tidx, _, _ = cute.arch.thread_idx() + + # Vectorised SMEM access constants for bank-conflict avoidance. The loads and the fused 4-wide + # mul_cvt below assume PACK_SIZE == 4, so these are not tunable. + PACK_SIZE = 4 # Elements per vector load + # Each thread reads 8 waves with each wave reads 4 packed bf16, so it reads a whole MXFP8 block in total + WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE + TOTAL_BANKS_WIDTH = 32 * 4 # 32 banks × 4 bytes, in bytes + THREADS_PER_BANK = TOTAL_BANKS_WIDTH // MXFP8_BLOCK_SCALING_SIZE # 4 threads per bank + + CTA_THREADS_Y = TILE_Y # threads per column (rows per tile) + CTA_THREADS_X = TILE_X // MXFP8_BLOCK_SCALING_SIZE # threads per row (chunks per row) + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((CTA_THREADS_Y, CTA_THREADS_X), stride=(CTA_THREADS_X, 1)), + val_layout=cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + + sX_tv = cute.composition(sX_tile, tv_layout) + sO_tv = cute.composition(sO_row_tile, tv_layout) + + # I/O Elements that belong to this thread + sX_thread = sX_tv[tidx, None] # shape (32,) bf16 + sO_thread = sO_tv[tidx, None] # shape (32,) fp8 + + sO_thread_u32_ptr = cute.recast_ptr(sO_thread.iterator, dtype=Uint32) + sO_thread_u32 = cute.make_tensor( + sO_thread_u32_ptr, + cute.make_layout( + (MXFP8_BLOCK_SCALING_SIZE // 4,), stride=(1,) + ), # 1 uint32 is 4 fp8 elements + ) + + # PTX allows to fuse relu activation in `cvt.rn.satfinite` unless we need to reduction for dbias + FUSE_RELU = cutlass.const_expr(ACTIVATION == "relu") and not WITH_DBIAS + # For this fast path we can read in pack of 2 instead of reading individual f16 / bf16 element, + # and keep the elements in half precision instead of upcasting every one of them to f32. + USE_HALF_PRECISION = is_packed16(DTYPE) and (ACTIVATION is None or FUSE_RELU) + + amax_r = Float32(0.0) + + # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks + # to avoid bank conflict. + bank_group = (tidx % THREADS_PER_WARP) // THREADS_PER_BANK + # The offset this thread should start reading from based on what's its first bank to access. + offset = bank_group * PACK_SIZE + if cutlass.const_expr(USE_HALF_PRECISION): + # If no activation, f16 / bf16 and rowwise quantization, we can read 2 f16 / bf16 at once in a pack + # and use max.xorsign.abs.f16x2 / max.xorsign.abs.bf16x2 to compute + max_x2 = max_x2_f16 if DTYPE is cutlass.Float16 else max_x2_bf16 + abs_max_x2 = abs_max_x2_f16 if DTYPE is cutlass.Float16 else abs_max_x2_bf16 + x2_lo_to_f32 = x2_lo_to_f32_f16 if DTYPE is cutlass.Float16 else x2_lo_to_f32_bf16 + x2_hi_to_f32 = x2_hi_to_f32_f16 if DTYPE is cutlass.Float16 else x2_hi_to_f32_bf16 + sX_thread_rw_i64 = cute.make_tensor( + cute.recast_ptr(sX_thread.iterator, dtype=Int64), + cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE // 4), stride=(0, 1) + ), # 1 int64 is 4 fp16/bf16 elements + ) + # Each wave reads its 4 elements (PACK_SIZE) as one 8-byte vectorized load + in_r = [[None, None] for _ in range(WAVES)] + for w in cutlass.range_constexpr(WAVES): + idx = (w + offset // 4) % (MXFP8_BLOCK_SCALING_SIZE // 4) + in_r[w][0], in_r[w][1] = unpack_i64_to_i32x2(sX_thread_rw_i64[0, idx]) + + if cutlass.const_expr(WITH_DBIAS): + for w in cutlass.range_constexpr(WAVES): + dbias_acc[w * PACK_SIZE + 0] += x2_lo_to_f32(in_r[w][0]) + dbias_acc[w * PACK_SIZE + 1] += x2_hi_to_f32(in_r[w][0]) + dbias_acc[w * PACK_SIZE + 2] += x2_lo_to_f32(in_r[w][1]) + dbias_acc[w * PACK_SIZE + 3] += x2_hi_to_f32(in_r[w][1]) + + amax_2x = in_r[0][0] + # Each wave will use max.xorsign.abs.f16x2 or max.xorsign.abs.bf16x2 to compare 2 packed elements in parallel + for w in cutlass.range_constexpr(WAVES): + if cutlass.const_expr(FUSE_RELU): + # Skip the first iteration because we assigned in_r[0][0] to amax_2x already + if cutlass.const_expr(w > 0): + # If we fuse relu then we don't want to do abs since negative value will be set to 0 and they will lose comparison automatically + amax_2x = max_x2(amax_2x, in_r[w][0]) + amax_2x = max_x2(amax_2x, in_r[w][1]) + else: + # Skip the first iteration because we assigned in_r[0][0] to amax_2x already + if cutlass.const_expr(w > 0): + amax_2x = abs_max_x2(amax_2x, in_r[w][0]) + amax_2x = abs_max_x2(amax_2x, in_r[w][1]) + if cutlass.const_expr(FUSE_RELU): + # Compare the 2 packed max without abs + amax_r = cute.arch.fmax( + x2_lo_to_f32(amax_2x), + x2_hi_to_f32(amax_2x), + ) + # For relu the max is at least 0 + amax_r = cute.arch.fmax(amax_r, Float32(0.0)) + else: + # Compare the 2 packed abs max + amax_r = cute.arch.fmax( + fabs_f32(x2_lo_to_f32(amax_2x)), + fabs_f32(x2_hi_to_f32(amax_2x)), + ) + else: + # Since we need to do computation on individual f16 / bf16 elements, we can't read in pack + sX_thread_rw = cute.make_tensor( + sX_thread.iterator, + cute.make_layout((1, MXFP8_BLOCK_SCALING_SIZE), stride=(0, 1)), + ) + + if cutlass.const_expr(WITH_DACT): + # Backward: out = grad · act'(act_input). sX is grad, sA is act_input. + 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 = 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 + + # Each wave we read PACK_SIZE elements, and we have WAVES waves, so we read WAVES * PACK_SIZE (= MXFP8_BLOCK_SCALING_SIZE) elements in total. + in_r = [[None] * PACK_SIZE for _ in range(WAVES)] + # Don't repeat these OOB masking computation in the loop + thread_row_start = tile_row_start + tidx // CTA_THREADS_X + thread_col_start = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE + thread_row_oob = thread_row_start >= M + # PACK_SIZE must divide every term of `thread_col_start + start` and N itself, or a + # pack could straddle the N boundary and the check below would be wrong. TILE_X is a + # tunable, so it is asserted rather than assumed. + assert MXFP8_BLOCK_SCALING_SIZE % PACK_SIZE == 0 + assert SYM_N_DIVISIBILITY % PACK_SIZE == 0 + assert TILE_X % PACK_SIZE == 0 + for w in cutlass.range_constexpr(WAVES): + start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + # thread_col_start + start >= N means "if the first element in this pack is OOB" + # If the first element is OOB, then (obviously) elements after it within this pack must be OOB + # If the first element is not OOB, then elements after it within this pack must be not OOB either. + # Proof: + # (1) when the first element is not OOB, we have thread_col_start + start < N, + # which implies N - (thread_col_start + start) > 0 + # (2) with the assertions above: + # - N is divisible by SYM_N_DIVISIBILITY, so N is also divisible by PACK_SIZE + # - thread_col_start = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE is divisible by PACK_SIZE + # because TILE_X and MXFP8_BLOCK_SCALING_SIZE are both divisible by PACK_SIZE + # - start = (w * PACK_SIZE + bank_group * PACK_SIZE) % MXFP8_BLOCK_SCALING_SIZE is also divisible by PACK_SIZE + # because MXFP8_BLOCK_SCALING_SIZE is a multiple of PACK_SIZE + # so N - (thread_col_start + start) must be divisible by PACK_SIZE + # Therefore, N - (thread_col_start + start) >= PACK_SIZE because PACK_SIZE is the minimal number + # for (1) and (2) to hold, which implies thread_col_start + start + PACK_SIZE -1 < N + # Hence, we only need to check if the first element is OOB to determine if the whole pack is OOB + thread_col_oob = thread_col_start + start >= N + for i in cutlass.range_constexpr(PACK_SIZE): + x = Float32(sX_thread_rw[0, start + i]) + if cutlass.const_expr(WITH_DACT): + # out = grad · act'(act_input) + x = x * dop(Float32(sA_thread_rw[0, start + i])) + # If IS_ACT, apply activation function to x in f32 + elif cutlass.const_expr(WITH_ACT): + # If it's relu, we can handle it later + if not cutlass.const_expr(FUSE_RELU): + x = op(x) + if not cutlass.const_expr(SKIP_INPUT_MASKING): + # If the input shape is not divisible by the tile size, + # TMA would zero-fills the input tile outside its logical MxN bounds. + # This is fine for non-activation cases, but for activation cases, + # op(0) might not be 0 which will pollute the amax and dbias. + # So we must manually mask the OOB region here. + if thread_row_oob or thread_col_oob: + x = Float32(0.0) + # Accumulate to the per-thread dbias register buffer for this tile if WITH_DBIAS + if cutlass.const_expr(WITH_DBIAS): + # dbias_acc is register buffer so we can just write without bank conflict + dbias_acc[w * PACK_SIZE + i] += x + # If 16-bit input with activation, truncate to IType + if cutlass.const_expr(is_packed16(DTYPE) and ACTIVATION is not None): + x = truncate_f32(x) + in_r[w][i] = x + if cutlass.const_expr(FUSE_RELU): + amax_r = cute.arch.fmax( + amax_r, x + ) # For relu cases, we don't need abs since negative values will be 0 so they lose comparison automatically + else: + amax_r = cute.arch.fmax(amax_r, fabs_f32(x)) + if cutlass.const_expr(FUSE_RELU): + amax_r = cute.arch.fmax(amax_r, Float32(0.0)) # If relu, the amax is at least 0 + + # With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's + # generic kernel (0-seeded fmaxf); CUDA's cast kernels, like the plain cast here, give it the + # NaN scale. + if cutlass.const_expr(ACTIVATION is not None or FUSED): + amax_r = cute.arch.fmax(amax_r, Float32(0.0)) + biased_exp_r = cvt_f32_to_fp8e8m0fnu(amax_r * MAX_NORM_RCP) + + # mS_row_stage has logical shape (32, 2) and we have 64 threads where each is mapped to one scale factor + # The TV layout is equivalent to TV layout with thr_layout=(32, 2):(2, 1), val_layout=(1,) + # but it's too trival so let's just index it directly without using layout + # Note this is the logical layout, which is on top of the swizzled / non-swizzled scale factor layout + # that mappes the logical index to the physical offset + + # For irregular shapes, skip the scale store if this thread's logical row / col-block lies past the input's actual extents. + # TMA already zero-fills OOB input reads and drops OOB output writes; only the direct scale-byte gmem store needs an explicit guard. + # If the input shape is divisible by the tile size * tile num, then we won't access OOB regions for sure + if cutlass.const_expr(SKIP_SCALE_BOUNDS): + mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = biased_exp_r + else: + scale_row = tile_row_start + tidx // CTA_THREADS_X + scale_col_first_elt = tile_col_start + (tidx % CTA_THREADS_X) * MXFP8_BLOCK_SCALING_SIZE + if scale_row < M and scale_col_first_elt < N: + mS_row_stage[(tidx // CTA_THREADS_X, tidx % CTA_THREADS_X)] = biased_exp_r + + inv_scale_r = exp2f_rcp(biased_exp_r) # f32 reciprocal of the scale + scale_2x = pack_f32x2(inv_scale_r, inv_scale_r) + if cutlass.const_expr(USE_HALF_PRECISION): + mul_cvt_x4_func = mul_f32x2_cvt_packed16x4_to_fp8x4(DTYPE, FP8_DTYPE, FUSE_RELU) + else: + mul_cvt_x4_func = mul_f32x2_cvt_f32x4_to_fp8x4(FP8_DTYPE, FUSE_RELU) + + for w in cutlass.range_constexpr(WAVES): + idx = (w * 4 + offset) % MXFP8_BLOCK_SCALING_SIZE + idx = idx // 4 + if cutlass.const_expr(USE_HALF_PRECISION): + # Convert 2 packed f16/bf16 pairs to 4 fp8 in one fused op + sO_thread_u32[idx] = mul_cvt_x4_func(in_r[w][0], in_r[w][1], scale_2x) + else: + # Convert 4 f32 to 4 fp8 in one fused op + sO_thread_u32[idx] = mul_cvt_x4_func( + in_r[w][0], in_r[w][1], in_r[w][2], in_r[w][3], scale_2x + ) + + return amax_r + + +@cute.jit +def _quantize_colwise_mxfp8( + sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA + sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) + sO_col_tile: cute.Tensor, # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) + mS_col_stage: cute.Tensor, # colwise scale tensor (1D swizzled, or 2D linear) + MAX_NORM_RCP: cutlass.Constexpr[float], + tile_row_start: Int32, # global row index of this stage's row 0 + tile_col_start: Int32, # global col index of this CTA's col 0 + 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], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], # pylint: disable=unused-argument # kept for API consistency + SKIP_INPUT_MASKING: cutlass.Constexpr[bool] = False, + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool] = False, + WITH_ACT: cutlass.Constexpr[bool] = False, + 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), + dbias being dbias_running plus this stage's rows, added in row order.""" + tidx, _, _ = cute.arch.thread_idx() + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), + val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), + ) + + sX_tv = cute.composition(sX_tile, tv_layout) + sO_tv = cute.composition(sO_col_tile, tv_layout) + + # I/O Elements that belong to this thread + sX_thread = sX_tv[tidx, None] + sO_thread = sO_tv[tidx, None] + + # PTX allows to fuse relu activation in `cvt.rn.satfinite` unless we need to reduction for dbias + 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) + # 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 + abs_max_scalar = abs_max_scalar_f16 if DTYPE is cutlass.Float16 else abs_max_scalar_bf16 + to_f32 = to_f32_f16 if DTYPE is cutlass.Float16 else to_f32_bf16 + # If we can use the half precision format, then use the input tile directly since there is no need to upcast + sX_thread_packed16 = cute.make_tensor( + sX_thread.iterator, + cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(TILE_X,)), + ) + # Stash the strided column reads in registers (CUDA's in_colwise_IType): + # the cvt loop below reuses them instead of re-reading smem. + in_c = [None] * MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + in_c[i] = sX_thread_packed16[i] + + amax_value = in_c[0] + # Skip the first iteration because we assigned in_c[0] to amax_value already + for i in cutlass.range_constexpr(1, MXFP8_BLOCK_SCALING_SIZE): + if cutlass.const_expr(FUSE_RELU): + # If we fuse relu then we don't want to do abs since negative value will be set to 0 + # and they will lose comparison automatically + amax_value = max_scalar(amax_value, in_c[i]) + else: + amax_value = abs_max_scalar(amax_value, in_c[i]) + if cutlass.const_expr(FUSE_RELU): + amax_c = cute.arch.fmax(to_f32(amax_value), Float32(0.0)) + else: + amax_c = fabs_f32(to_f32(amax_value)) + else: + # Otherwise we need to case input values to fp32. Allocate the register tensor and load from SMEM input tiles. + rX_thread_f32 = cute.make_rmem_tensor( + layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), + dtype=Float32, + ) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + 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 = 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 = 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 + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + # If it's relu, we can handle it later in the cvt + if not cutlass.const_expr(FUSE_RELU): + rX_thread_f32[i] = op(rX_thread_f32[i]) + # If the input shape is not divisible by the tile size, + # TMA would zero-fills the input tile outside its logical MxN bounds. + # This is fine for non-activation cases, but for activation cases, + # op(0) might not be 0 which will pollute the amax and dbias. + # So we must manually mask the OOB region here. + if not cutlass.const_expr(SKIP_INPUT_MASKING): + if tile_row_start + i >= M or thread_col_oob: + rX_thread_f32[i] = Float32(0.0) + # Accumulate fp32 activations to DBIAS before we truncate to half precision when the input is half precision + if cutlass.const_expr(WITH_DBIAS): + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + dbias_partial += rX_thread_f32[i] + # Truncate the activation (after we apply op) back to the half precision type if input is also half precision. + 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 + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + rX_thread_f32[i] = truncate_f32(rX_thread_f32[i]) + # Columnwise is the preferred direction so it runs first. If it needs to cache the activation in the input tile + # to let the rowwise pass read it, we need to cast and overwrite the input data in-place here + if cutlass.const_expr(CACHE_ACTIVATION): + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + sX_thread[i] = DTYPE(rX_thread_f32[i]) + amax_c = Float32(0.0) + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + if cutlass.const_expr(FUSE_RELU): + amax_c = cute.arch.fmax(amax_c, rX_thread_f32[i]) + else: + amax_c = cute.arch.fmax(amax_c, fabs_f32(rX_thread_f32[i])) + + # Irregular shapes: skip when this stage's row range or this thread's + # column lies past the input extents. TILE_Y == MXFP8_BLOCK_SCALING_SIZE so each stage + # is exactly one scale-row; valid iff `tile_row_start < M`. + # With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's + # generic kernel (0-seeded fmaxf); CUDA's cast kernels, like the plain cast here, give it the + # NaN scale. + if cutlass.const_expr(ACTIVATION is not None): + amax_c = cute.arch.fmax(amax_c, Float32(0.0)) + biased_exp_c = cvt_f32_to_fp8e8m0fnu(amax_c * MAX_NORM_RCP) + # If the input shape is divisible by the tile size * tile num, then we won't access OOB regions for sure + if cutlass.const_expr(SKIP_SCALE_BOUNDS): + if cutlass.const_expr(SWIZZLE): + mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c + else: + mS_col_stage[(0, tidx)] = biased_exp_c + else: + scale_col = tile_col_start + tidx + if tile_row_start < M and scale_col < N: + if cutlass.const_expr(SWIZZLE): + mS_col_stage[(0, tidx % 32, tidx // 32)] = biased_exp_c + else: + mS_col_stage[(0, tidx)] = biased_exp_c + + inv_scale_c = exp2f_rcp(biased_exp_c) + # cvt.rn.satfinite can be vectorized to convert 2 f32 to 2 fp8 in one instruction + cvt_x2_func = get_cvt_f32x2_to_fp8x2_func(FP8_DTYPE, FUSE_RELU) + sO_thread_fp8 = cute.make_tensor( + cute.recast_ptr(sO_thread.iterator, dtype=FP8_DTYPE), sO_thread.layout + ) + if cutlass.const_expr(USE_HALF_PRECISION): + to_f32 = to_f32_f16 if DTYPE is cutlass.Float16 else to_f32_bf16 + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): + lo, hi = 2 * j, 2 * j + 1 + v_lo = to_f32(in_c[lo]) + v_hi = to_f32(in_c[hi]) + # Accumulate the per-thread column partial for dbias if WITH_DBIAS. + # Kept as two adds in element order: f32 addition is not associative + if cutlass.const_expr(WITH_DBIAS): + dbias_partial += v_lo + dbias_partial += v_hi + pair = cvt_x2_func(v_hi * inv_scale_c, v_lo * inv_scale_c) + sO_thread_fp8[lo] = Uint8(pair & Uint16(0xFF)).bitcast(FP8_DTYPE) + sO_thread_fp8[hi] = Uint8(pair >> Uint16(8)).bitcast(FP8_DTYPE) + else: + for j in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE // 2): + lo, hi = 2 * j, 2 * j + 1 + pair = cvt_x2_func(rX_thread_f32[hi] * inv_scale_c, rX_thread_f32[lo] * inv_scale_c) + sO_thread_fp8[lo] = Uint8(pair & Uint16(0xFF)).bitcast(FP8_DTYPE) + sO_thread_fp8[hi] = Uint8(pair >> Uint16(8)).bitcast(FP8_DTYPE) + + # Return this stage's per-column partial alongside amax; the caller accumulates + # it across stages (a scalar can't be updated in-place through the arg). + return amax_c, dbias_partial + + +@cute.jit +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. + 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. + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), + val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), + ) + sX_thread = cute.composition(sX_tile, tv_layout)[tidx, None] + + dbias_partial = Float32(dbias_running) + if cutlass.const_expr(WITH_DACT): + dop = dactivation_op(ACTIVATION, ACTIVATION_TABLE) + sA_thread = cute.composition(sA_tile, tv_layout)[tidx, None] + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + x = Float32(sX_thread[i]) * dop(Float32(sA_thread[i])) + dbias_partial += x + if cutlass.const_expr(CACHE_ACTIVATION): + sX_thread[i] = DTYPE(x) + else: + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + dbias_partial += Float32(sX_thread[i]) + return dbias_partial + + +@cute.jit +def quantize_mxfp8_smem_tile_general( + sX_tile: cute.Tensor, # (TILE_Y, TILE_X) bf16/fp16 smem view, post-TMA + sActInput_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) act-input smem tile (dact only) + sO_row_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) fp8 smem view (rowwise FP8 output) + mS_row_stage: Optional[cute.Tensor], # rowwise scale tensor (1D swizzled, or 2D linear) + sO_col_tile: Optional[cute.Tensor], # (TILE_Y, TILE_X) fp8 smem view (colwise FP8 output) + mS_col_stage: Optional[cute.Tensor], # colwise scale tensor (1D swizzled, or 2D linear) + rowwise_dbias_acc: Optional[cute.Tensor], # rmem Float32[32] (DBIAS_REDUCTION_IN_ROWWISE only) + tile_row_start: Int32, # global row index of this tile's row 0 + tile_col_start: Int32, # global col index of this tile's col 0 + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + WITH_AMAX: cutlass.Constexpr[bool], + WITH_ACT: cutlass.Constexpr[bool], + WITH_DACT: cutlass.Constexpr[bool], + ACTIVATION: cutlass.Constexpr[str | None], + ACTIVATION_TABLE: cutlass.Constexpr[str], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + FP8_DTYPE: cutlass.Constexpr[Type[cutlass.Numeric]], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], + SKIP_INPUT_MASKING: cutlass.Constexpr[bool], + SKIP_SCALE_BOUNDS: cutlass.Constexpr[bool], + CACHE_ACTIVATION: cutlass.Constexpr[bool], + DBIAS_REDUCTION_IN_COLWISE: cutlass.Constexpr[bool], + COLWISE_DBIAS_REDUCTION_ONLY: cutlass.Constexpr[bool], + DBIAS_REDUCTION_IN_ROWWISE: cutlass.Constexpr[bool], + dbias_running: Float32 = 0.0, +): + """Quantize one SMEM tile in the directions enabled by ROWWISE / COLWISE; returns (amax, dbias). + + amax is this thread's amax over the tile (None unless WITH_AMAX). dbias is this thread's + column 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 + dbias = None + if cutlass.const_expr(COLWISE): + amax_c, dbias_c = _quantize_colwise_mxfp8( + sX_tile, + sActInput_tile, + sO_col_tile, + mS_col_stage, + MAX_NORM_RCP, + tile_row_start, + tile_col_start, + M, + N, + ACTIVATION=ACTIVATION, + ACTIVATION_TABLE=ACTIVATION_TABLE, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + SWIZZLE=WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + SKIP_INPUT_MASKING=SKIP_INPUT_MASKING, + SKIP_SCALE_BOUNDS=SKIP_SCALE_BOUNDS, + WITH_ACT=WITH_ACT, + 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. + if cutlass.const_expr(WITH_AMAX): + amax = amax_c + if cutlass.const_expr(DBIAS_REDUCTION_IN_COLWISE): + dbias = dbias_c + # If we don't quantize columnwise but we fuse dbias, do a dbias reduction only pass in the columnwise direction without quantization + if cutlass.const_expr(COLWISE_DBIAS_REDUCTION_ONLY): + dbias = _dbias_reduction_colwise( + 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. + if cutlass.const_expr(CACHE_ACTIVATION): + cute.arch.sync_threads() + if cutlass.const_expr(ROWWISE): + amax_r = _quantize_rowwise_mxfp8( + sX_tile, + None if CACHE_ACTIVATION else sActInput_tile, + sO_row_tile, + mS_row_stage, + MAX_NORM_RCP, + tile_row_start, + tile_col_start, + M, + N, + ACTIVATION=None if CACHE_ACTIVATION else ACTIVATION, + ACTIVATION_TABLE=ACTIVATION_TABLE, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + SKIP_INPUT_MASKING=SKIP_INPUT_MASKING, + SKIP_SCALE_BOUNDS=SKIP_SCALE_BOUNDS, + WITH_ACT=WITH_ACT and not CACHE_ACTIVATION, + WITH_DACT=WITH_DACT and not CACHE_ACTIVATION, + FUSED=WITH_ACT or WITH_DACT, + WITH_DBIAS=DBIAS_REDUCTION_IN_ROWWISE, + dbias_acc=rowwise_dbias_acc, + ) + if cutlass.const_expr(WITH_AMAX and not COLWISE): + amax = amax_r + return amax, dbias + + +class MXFP8QuantizeGeneralG2SKernel(MXFP8QuantizeKernelBase): + """The MXFP8 quantization kernel that mirrors the standard (non-specialized) MXFP8 CUDA C++ quantization kernel + with multiple fusions (activation, dbias, etc.). + `__call__` method is the entrypoint which is AOT compiled. `self` will be captured so it's fixed per compiled kernel + """ + + _NUM_STAGES = 2 # The pipeline depth is always 2 + + def __init__( + self, + cfg: MXFP8QuantizeConfig, + SKIP_INPUT_MASKING: bool = False, + SKIP_SCALE_BOUNDS: bool = False, + ): + self.cfg = cfg + # If the input shape is divisible by the tile size or f(0)=0 holds for activaions, + # we can skip masking inputs with zero in the kernel and save some instructions. + self.SKIP_INPUT_MASKING = SKIP_INPUT_MASKING + # If the input shape is divisible by the tile size, we can skip bounds check for scale writes and save some instructions. + self.SKIP_SCALE_BOUNDS = SKIP_SCALE_BOUNDS + # Only honor the noop flag when no activation or dbias is fused to match CUDA C++'s implementation + self.CHECK_NOOP_FLAG: cutlass.const_expr = ( + not self.cfg.WITH_ACT and not self.cfg.WITH_DACT and not self.cfg.WITH_DBIAS + ) + cast_dbias_only = cfg.WITH_DBIAS and not cfg.WITH_DACT and not cfg.WITH_ACT + # Use a different tile size for dbias only config + # No matter what tile size we use, each thread always handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk + if cast_dbias_only: + self._NUM_TILES = 4 # Each CTA handles 4 tiles stacked vertically + self._THREADS_PER_CTA = 128 + else: + self._NUM_TILES = 2 # Each CTA handles 2 tiles stacked vertically + self._THREADS_PER_CTA = 64 + # If _USE_STATIC_NUM_TILES, then each CTA will always process _NUM_TILES tiles without checking if those tiles + # are OOB (which will be zero padded by TMA). Hardcoded to true because benchmark shows this is always better + self._USE_STATIC_NUM_TILES = True + # A hint that is tuned based on the benchmark which helps when ACT/DACT is fused + self._MIN_BLOCKS_PER_MP = 1 if (cfg.WITH_ACT or cfg.WITH_DACT) else 0 + # Each thread handles a (1, MXFP8_BLOCK_SCALING_SIZE) chunk + self._TILE_COLS = self._THREADS_PER_CTA + self._TILE_ROWS = MXFP8_BLOCK_SCALING_SIZE + self._NUM_WARPS = self._THREADS_PER_CTA // 32 + # We prefer to do dbias reduction in colwise which is easier (no cross-thread reduction needed). + # Only do rowwise reduction when we don't quantize columnwisely when WITH_DBIAS is True. + + # If columnwise is not quantized, and dbias is fused, we will do a columnwise reduction only pass + # (may also cache dact values) without quantization because doing dbias in rowwise requires a SMEM shuffle + # which is even slower + self.COLWISE_DBIAS_REDUCTION_ONLY = cfg.WITH_DBIAS and not cfg.COLWISE + # If columnwise is quantized and dbias is fused, do dbias reduction in colwise which is easier + # (no cross-thread reduction needed). + self.DBIAS_REDUCTION_IN_COLWISE = cfg.WITH_DBIAS and cfg.COLWISE + # Never do dbias reduction in rowwise because it is slow. Leave this so in case it's enabled in the future + self.DBIAS_REDUCTION_IN_ROWWISE = False + + # Cache activation in-place in the SMEM input tile when we process both rowwise and colwise passes + # so the activation is only computed once in the direction we favor (columnwise) and the other direction (rowwise) + # reads the cached value instead of recomputing it. + # Note: if activation is relu, there is no standalong relu applied because it's already fused into `cvt.rn.satfinite` + # so it should be treated as "no activation" + self.CACHE_ACTIVATION = ( + (cfg.WITH_ACT or cfg.WITH_DACT) + and cfg.ROWWISE + # Always cache activation in the colwise quantization or dbias reduction only pass + and (cfg.COLWISE or self.COLWISE_DBIAS_REDUCTION_ONLY) + and cfg.ACTIVATION != "relu" + ) + + @cute.jit + def __call__( + self, + mX: cute.Tensor, # Input tensor to quantize + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], # Rowwise output and scale tensors + mO_col: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], # Colwise output and scale tensors + mAmax: Optional[cute.Tensor], # Global amax accumulator, only used when WITH_AMAX is True + mNoop: cute.Pointer, # f32 cast_noop flag; may be null, checked on device + mDActInput: Optional[ + cute.Tensor + ], # Activation input for activation derivative fusion, only used when WITH_DACT is True + 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): + cute.printf( + f"[CuTeDSL] MXFP8QuantizeGeneralG2SKernel.__call__() with config: {self.cfg}\n" + ) + + M = mX.shape[0] + N = mX.shape[1] + cfg = self.cfg + + # If WITH_GEMM_SWIZZLED_SCALES is enabled, the output must satisfy cublas's swizzled layout + # This is expressed as a CuTe layout applied to the output tensor so it can be transparent throughout the kernel implementation. + # See https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout for more details. + if cutlass.const_expr(cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, mS_col = derive_swizzled_scale_layout( + M, N, cfg.ROWWISE, cfg.COLWISE, mS_row, mS_col + ) + + # We have 2 stages in our pipeline where each stage loads / computes a (TILE_Y, TILE_X) tile + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + + # Input TMA atoms + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + tma_atom, tma_src = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, + mX, + smem_tile_layout, + cta_tiler, + num_multicast=1, + ) + + # Activation input TMA atoms for activation derivative fusion + tma_atom_act = None + tma_src_act = None + if cutlass.const_expr(cfg.WITH_DACT): + tma_atom_act, tma_src_act = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_load, + mDActInput, + smem_tile_layout, + cta_tiler, + num_multicast=1, + ) + + # Output TMA atoms + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + out_smem_layout = cute.make_ordered_layout((self._TILE_ROWS, self._TILE_COLS), order=(1, 0)) + tma_atom_out_row = None + tma_dst_out_row = None + tma_atom_out_col = None + tma_dst_out_col = None + if cutlass.const_expr(cfg.ROWWISE): + tma_atom_out_row, tma_dst_out_row = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_row, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + if cutlass.const_expr(cfg.COLWISE): + tma_atom_out_col, tma_dst_out_col = cute.nvgpu.cpasync.make_tiled_tma_atom( + op_store, + mO_col, + out_smem_layout, + cta_tiler, + num_multicast=1, + ) + + grid = [ + cute.ceil_div(Int32(N), self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), + ] + block = [ + self._THREADS_PER_CTA, + ] + + self.kernel( + mX, + mS_row, + mS_col, + mAmax, + mNoop, + mWorkspace, + mX.element_type, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + tma_atom_act, + tma_src_act, + ).launch( + grid=grid, + block=block, + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mNoop: cute.Pointer, + mWorkspace: Optional[cute.Tensor], + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, # Input TMA atoms + tma_atom_out_row: Optional[cute.CopyAtom], + tma_dst_out_row: Optional[cute.Tensor], # Rowwise output TMA atoms + tma_atom_out_col: Optional[cute.CopyAtom], + tma_dst_out_col: Optional[cute.Tensor], # Colwise output TMA atoms + tma_atom_act: Optional[cute.CopyAtom], + tma_src_act: Optional[ + cute.Tensor + ], # Activation derivative TMA atoms, None unless WITH_DACT + ): + """Device entry: no-op the CTA when the noop flag is set, else run the quantize main loop.""" + + skip_execution = False + if cutlass.const_expr(self.CHECK_NOOP_FLAG): + skip_execution = noop_flag_is_set(mNoop) + if not skip_execution: + self._kernel_main( + mX, + mS_row, + mS_col, + mAmax, + mWorkspace, + dtype, + tma_atom, + tma_src, + tma_atom_out_row, + tma_dst_out_row, + tma_atom_out_col, + tma_dst_out_col, + tma_atom_act, + tma_src_act, + ) + + @cute.jit + def _kernel_main( + self, + mX: cute.Tensor, + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + mAmax: Optional[cute.Tensor], + mWorkspace: Optional[cute.Tensor], + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + tma_atom: cute.CopyAtom, + tma_src: cute.Tensor, # Input TMA atoms + tma_atom_out_row: Optional[cute.CopyAtom], + tma_dst_out_row: Optional[cute.Tensor], # Rowwise output TMA atoms + tma_atom_out_col: Optional[cute.CopyAtom], + tma_dst_out_col: Optional[cute.Tensor], # Colwise output TMA atoms + tma_atom_act: Optional[cute.CopyAtom], + tma_src_act: Optional[ + cute.Tensor + ], # Activation derivative TMA atoms, None unless WITH_DACT + ): + cfg = self.cfg + + if cutlass.const_expr(cfg.ROWWISE): + mS_row = cute.zipped_divide( + mS_row, (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + ) + if cutlass.const_expr(cfg.COLWISE): + mS_col = cute.zipped_divide( + mS_col, (self._TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, self._TILE_COLS) + ) + + # Allocate shared memory for the input and rowwise / columnwise outputs + FP8_DTYPE = cfg.FP8_DTYPE + + if cutlass.const_expr(cfg.ROWWISE and cfg.COLWISE): + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + elif cutlass.const_expr(cfg.ROWWISE and not cfg.COLWISE): + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_row: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + else: + + @cute.struct + class SharedStorage: + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * self._NUM_STAGES] + sX: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sO_col: cute.struct.Align[ + cute.struct.MemRange[ + FP8_DTYPE, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + sAmax: cute.struct.MemRange[Float32, self._NUM_WARPS] + + smem = cutlass.memory.SmemAllocator() + storage = smem.allocate(SharedStorage) + # Apply the layout to the allocated shared memory buffers so the first rank is the tile (nested layout) + # and the second rank is the pipeline stage + sX = storage.sX.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + if cutlass.const_expr(cfg.ROWWISE): + sO_row = storage.sO_row.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + if cutlass.const_expr(cfg.COLWISE): + sO_col = storage.sO_col.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + + # Allocate shared memory for the activation input used for the activation derivative fusion. + if cutlass.const_expr(cfg.WITH_DACT): + + @cute.struct + class DactStorage: + sActInput: cute.struct.Align[ + cute.struct.MemRange[ + dtype, self._TILE_ROWS * self._TILE_COLS * self._NUM_STAGES + ], + 128, + ] + + dact_storage = smem.allocate(DactStorage) + # Apply the same layout as the input + sActInput = dact_storage.sActInput.get_tensor( + cute.make_layout( + ((self._TILE_ROWS, self._TILE_COLS), self._NUM_STAGES), + stride=((self._TILE_COLS, 1), self._TILE_ROWS * self._TILE_COLS), + ) + ) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Prefetch TMA descriptors + if warp_idx == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) + if cutlass.const_expr(cfg.WITH_DACT): + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_act) + + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + + # Only warp 0 is the producer (issues TMA) + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + # Every warp is the consumer (reads the data loaded by TMA) + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) + + # Bytes transferred per TMA copy: one (TILE_Y, TILE_X) tile of dtype. + tx_count = self._TILE_ROWS * self._TILE_COLS * dtype.width // 8 + # dact loads two tiles (grad + act_input) under the same per-stage barrier, + # so the barrier must expect both copies' bytes. + if cutlass.const_expr(cfg.WITH_DACT): + tx_count *= 2 + + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_storage.data_ptr(), + num_stages=self._NUM_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + tx_count=tx_count, + cta_layout_vmnk=None, # single-CTA, no cluster/multicast + ) + + prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._NUM_STAGES + ) + cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._NUM_STAGES + ) + + M = mX.shape[0] + N = mX.shape[1] + + # With static num tiles, always process _NUM_TILES tiles like the CUDA kernel: TMA zero-fills + # OOB tiles, so the loop trip count is a compile-time constant. + if cutlass.const_expr(self._USE_STATIC_NUM_TILES): + num_tiles = self._NUM_TILES + else: + num_tiles = cutlass.min( + self._NUM_TILES, + cute.ceil_div(M - bidy * self._TILE_ROWS * self._NUM_TILES, self._TILE_ROWS), + ) + + # Tile the TMA gmem view: ((TILE_Y, TILE_X), (M/TILE_Y, N/TILE_X)). + gX_tiled = cute.zipped_divide(tma_src, (self._TILE_ROWS, self._TILE_COLS)) + + # Partition sX/gX for the TMA atom (single-CTA, no cluster/multicast). + tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( + tma_atom, + 0, # Use the only CTA to do the TMA copy + cute.make_layout(1), # This cluster only has 1 CTAs + sX, + gX_tiled, + ) + + # If WITH_DACT, partition the activation input for TMA as well in the same way + if cutlass.const_expr(cfg.WITH_DACT): + gA_tiled = cute.zipped_divide(tma_src_act, (self._TILE_ROWS, self._TILE_COLS)) + tXsA, tXgA = cute.nvgpu.cpasync.tma_partition( + tma_atom_act, + 0, + cute.make_layout(1), + sActInput, + gA_tiled, + ) + + # Partitioning for rowwise / columnwise outputs + if cutlass.const_expr(cfg.ROWWISE): + gO_row_tiled = cute.zipped_divide(tma_dst_out_row, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_row, tXgO_row = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_row, + 0, + cute.make_layout(1), + sO_row, + gO_row_tiled, + ) + if cutlass.const_expr(cfg.COLWISE): + gO_col_tiled = cute.zipped_divide(tma_dst_out_col, (self._TILE_ROWS, self._TILE_COLS)) + tXsO_col, tXgO_col = cute.nvgpu.cpasync.tma_partition( + tma_atom_out_col, + 0, + cute.make_layout(1), + sO_col, + gO_col_tiled, + ) + + # Ensure barrier init is visible to all threads before the pipeline is used. + cute.arch.sync_threads() + + # Prologue: warp 0 prefetches up to NUM_STAGES tiles to fully fill the pipeline + if warp_idx == 0: + for s in cutlass.range_constexpr(self._NUM_STAGES): + if s < num_tiles: + mainloop_pipeline.producer_acquire(prod_state) + tile_y = bidy * self._NUM_TILES + s + cute.copy( + tma_atom, + tXgX[(None, (tile_y, bidx))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + if cutlass.const_expr(cfg.WITH_DACT): + cute.copy( + tma_atom_act, + tXgA[(None, (tile_y, bidx))], + tXsA[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # Per-thread amax accumulator + if cutlass.const_expr(cfg.WITH_AMAX): + per_thread_amax = Float32(0.0) + + # Prepare thread-level register accumulators for rowwise dbias reduction. + # Each thread will process two (1, MXFP8_BLOCK_SCALING_SIZE) rows in two stages, and in each stage the thread will add the + # (after dact applied) value to this register array with the same shape so it carries the the two stages' partial sum. + # Then it will be written to a SMEM buffer to let the whole CTA do the reduction separately to yield + # the final (1, TILE_X) dbias workspace output. + rowwise_dbias_acc = None + if cutlass.const_expr(self.DBIAS_REDUCTION_IN_ROWWISE): + rowwise_dbias_acc = cute.make_rmem_tensor( + layout_or_shape=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE,), stride=(1,)), + dtype=Float32, + ) + # Zero the accumulator registers. + for c in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + 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; each stage continues + # this running sum with its (after dact applied) column, in row order, as the CUDA kernels do. + # Then this sum will be written to the GMEM workspace buffer directly. + block_dbias = Float32(0.0) + + # Consumer: all warps fetch from the pipeline, process its tile, and issue a new load to the tile buffer it just consumed + for tile_idx in cutlass.range(num_tiles, unroll=1): + mainloop_pipeline.consumer_wait(cons_state) + # Only allow at most _NUM_STAGES-1 stages to be in-flight, because this iteration will reuse the ring buffer + # that is read _NUM_STAGES iterations ago so we must wait for whoever is reading that buffer to finish. + # Note: no sO buffer is reused when _NUM_TILES <= _NUM_STAGES, so we can always skip waiting when this holds + if cutlass.const_expr(self._NUM_TILES > self._NUM_STAGES): + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(self._NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + # The current pipeline stage index, which is the tile index modulo the number of stages. + # This is used to index into the shared memory ring buffers that are wrapped around the number of stages. + stage_idx = cons_state.index + sX_tile = sX[(None, stage_idx)] + # Also fetch the activation input if WITH_DACT + sActInput_tile = None + if cutlass.const_expr(cfg.WITH_DACT): + sActInput_tile = sActInput[(None, stage_idx)] + # Each CTA handles `NUM_TILES` tiles stacked vertically, so tile_idx_x is just the block index along X dimension + # and tile_idx_y is the tile that this stage handles out of the `NUM_TILES` tiles + tile_idx_x = bidx + tile_idx_y = bidy * self._NUM_TILES + tile_idx + sO_row_tile = None + mS_row_stage = None + if cutlass.const_expr(cfg.ROWWISE): + sO_row_tile = sO_row[(None, stage_idx)] + # mS_row is ((SCALE_TILE), (GRID)) where SCALE_TILE = (32, 2). + # Each CTA owns NUM_TILES consecutive row-tiles of GRID. cute + # auto-decomposes the flat row coord `bidy * NUM_TILES + tile_idx` + # onto GRID's hierarchical row modes — which is the + # (i_hi, tile_Y) tile-major order for swizzled, and the plain + # row-tile order for compact. Same source, both layouts correct. + mS_row_stage = cute.flatten(mS_row[(None, (tile_idx_y, tile_idx_x))]) + sO_col_tile = None + mS_col_stage = None + if cutlass.const_expr(cfg.COLWISE): + sO_col_tile = sO_col[(None, stage_idx)] + mS_col_stage = cute.flatten(mS_col[(None, (tile_idx_y, tile_idx_x))]) + + amax, dbias = quantize_mxfp8_smem_tile_general( + sX_tile, + sActInput_tile, + sO_row_tile, + mS_row_stage, + sO_col_tile, + mS_col_stage, + rowwise_dbias_acc, + tile_idx_y * self._TILE_ROWS, + tile_idx_x * self._TILE_COLS, + M, + N, + ROWWISE=self.cfg.ROWWISE, + COLWISE=self.cfg.COLWISE, + WITH_AMAX=self.cfg.WITH_AMAX, + 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, + WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_X=self._TILE_COLS, + TILE_Y=self._TILE_ROWS, + SKIP_INPUT_MASKING=self.SKIP_INPUT_MASKING, + SKIP_SCALE_BOUNDS=self.SKIP_SCALE_BOUNDS, + CACHE_ACTIVATION=self.CACHE_ACTIVATION, + 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 + + # Make the shared-memory writes visible to the TMA's async proxy before the TMA reads them. + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + cute.arch.sync_threads() + + # We are done with this input pipeline SMEM buffer, signal the producer that it can write to this buffer + mainloop_pipeline.consumer_release(cons_state) + + # Warp 0 issues TMA copy to write the quantized output tile from shared memory to global memory and then commits + if warp_idx == 0: + tile_y = bidy * self._NUM_TILES + tile_idx + if cutlass.const_expr(cfg.ROWWISE): + cute.copy( + tma_atom_out_row, + tXsO_row[(None, stage_idx)], + tXgO_row[(None, (tile_y, bidx))], + ) + if cutlass.const_expr(cfg.COLWISE): + cute.copy( + tma_atom_out_col, + tXsO_col[(None, stage_idx)], + tXgO_col[(None, (tile_y, bidx))], + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # The pipeline is no longer fully filled after we consume this tile, so we fetch a new tile to fill the pipeline. + # The next _NUM_STAGES-1 tiles are already in-flight, so the next tile to fetch is after _NUM_STAGES tiles. + if warp_idx == 0: + next_tile_idx = tile_idx + self._NUM_STAGES + if next_tile_idx < num_tiles: + mainloop_pipeline.producer_acquire(prod_state) + tile_y = bidy * self._NUM_TILES + next_tile_idx + cute.copy( + tma_atom, + tXgX[(None, (tile_y, bidx))], + tXsX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + if cutlass.const_expr(cfg.WITH_DACT): + cute.copy( + tma_atom_act, + tXgA[(None, (tile_y, bidx))], + tXsA[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + # End of the main pipeline loop + + # Complete the cross-thread dbias reduction after each thread has its own per-thread partial sum after the rowwise quantization. + if cutlass.const_expr(self.DBIAS_REDUCTION_IN_ROWWISE): + # If we do the dbias reduction in the rowwise pass, each thread will have a (1, MXFP8_BLOCK_SCALING_SIZE) partial sum + # and we need to write these to a SMEM buffer and let each thread reduce it in the columnwise direction + block_dbias = self._dbias_reduction_rowwise_epilouge(smem, tidx, rowwise_dbias_acc) + + # Write the per-tile reduced dbias to the global workspace. + if cutlass.const_expr(cfg.WITH_DBIAS): + dbias_col = bidx * self._TILE_COLS + tidx + if dbias_col < N: + mWorkspace[(bidy, dbias_col)] = block_dbias + + if cutlass.const_expr(cfg.WITH_AMAX): + sAmax = storage.sAmax.get_tensor(cute.make_layout(self._NUM_WARPS)) + self._amax_epilogue(sAmax, mAmax, tidx, warp_idx, per_thread_amax) + + # Wait only until in-flight TMA stores have read SMEM so the CTA retires (frees its SMEM) early; + # the grid's completion still guarantees the global writes are done. + cute.arch.cp_async_bulk_wait_group(0, read=True) + + @cute.jit + def _dbias_reduction_rowwise_epilouge( + self, smem: cutlass.memory.SmemAllocator, tidx: Int32, rowwise_dbias_acc: cute.Tensor + ): + # Pad the buffer to avoid bank conflicts. The logical shape is still the same. Only the stride is different. + DBIAS_BUFF_WIDTH = ( + self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE * (MXFP8_BLOCK_SCALING_SIZE + 1) + ) + + # Allocate the SMEM buffer that all threads use to reduce the two-stage partial sum (per thread) to the + # partial sum (per block). + @cute.struct + class DbiasStorage: + sDbias: cute.struct.MemRange[Float32, self._TILE_ROWS * DBIAS_BUFF_WIDTH] + + dbias_storage = smem.allocate(DbiasStorage) + sDbias = dbias_storage.sDbias.get_tensor( + cute.make_layout((self._TILE_ROWS, self._TILE_COLS), stride=(DBIAS_BUFF_WIDTH, 1)), + ) + _, tv_layout_dbias_write = cute.make_layout_tv( + thr_layout=cute.make_layout( + (self._TILE_ROWS, self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE), + stride=(self._TILE_COLS // MXFP8_BLOCK_SCALING_SIZE, 1), + ), + val_layout=cute.make_layout( + (1, MXFP8_BLOCK_SCALING_SIZE), stride=(MXFP8_BLOCK_SCALING_SIZE, 1) + ), + ) + sDbias_write = cute.composition(sDbias, tv_layout_dbias_write) + # rowwise_dbias_acc is in _quantize_rowwise_mxfp8's rotated wave order; these must match its constants. + PACK_SIZE = 4 + WAVES = MXFP8_BLOCK_SCALING_SIZE // PACK_SIZE + THREADS_PER_BANK = 32 * 4 // MXFP8_BLOCK_SCALING_SIZE + # Each thread start reading from the specfic bank based on its thread ID so they can do their best to access different banks + # to avoid bank conflict. + bank_group = (tidx % THREADS_PER_WARP) // THREADS_PER_BANK + # The offset this thread should start reading from based on what's its first bank to access. + offset = bank_group * PACK_SIZE + for w in cutlass.range_constexpr( + WAVES + ): # Each thread starts from this offset when writing into SMEM to avoid bank conflict + start = (w * PACK_SIZE + offset) % MXFP8_BLOCK_SCALING_SIZE + for i in cutlass.range_constexpr(PACK_SIZE): + # All threads write their per-thread partial sum results to the shared buffer. + sDbias_write[(tidx, start + i)] = rowwise_dbias_acc[w * PACK_SIZE + i] + cute.arch.sync_threads() + # All threads reduce the cross-thread partial sums to the per-block partial sum. + _, tv_layout_dbias_reduce = cute.make_layout_tv( + thr_layout=cute.make_layout((1, self._TILE_COLS), stride=(self._TILE_COLS, 1)), + val_layout=cute.make_layout((self._TILE_ROWS, 1), stride=(1, 1)), + ) + sDbias_reduce = cute.composition(sDbias, tv_layout_dbias_reduce) + # make_layout_tv yields a (thread, value) layout: thread=tidx -> column tidx, + # value=i -> row i. So index [tidx, i] (thread first), summing the column's rows. + block_dbias = Float32(0.0) + for i in cutlass.range_constexpr(self._TILE_ROWS): + block_dbias += sDbias_reduce[tidx, i] + return block_dbias + + @cute.jit + def _amax_epilogue( + self, + sAmax: cute.Tensor, + mAmax: cute.Tensor, + tidx: Int32, + warp_idx: Int32, + per_thread_amax: Float32, + ): + # Reduce and get the per-warp amax. + warp_amax = cute.arch.warp_redux_sync(per_thread_amax, kind="fmax") + # Write the per-warp amax to shared memory + lane_idx = tidx % 32 + if lane_idx == 0: + sAmax[warp_idx] = warp_amax + cute.arch.sync_threads() + if tidx == 0: + cta_amax = Float32(0.0) + # The first thread reduces all the per-warp amax to the per-CTA amax + for w in cutlass.range_constexpr(self._NUM_WARPS): + cta_amax = cute.arch.fmax(cta_amax, sAmax[w]) + amax_i32 = cute.make_tensor( + cute.recast_ptr(mAmax.iterator, dtype=Int32), + cute.make_layout(1), + ) + # The first thread updates the global amax with an atomic max on the bitcasted float value + cute.arch.atomic_max( + amax_i32.iterator, + cta_amax.bitcast(Int32), + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2r_bidim.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2r_bidim.py new file mode 100644 index 00000000000..b1ce6599c26 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2r_bidim.py @@ -0,0 +1,609 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Register-resident gated-activation MXFP8 cast in CuTeDSL, ported from the CUDA +register-tiled gated kernel (cast/mxfp8/specialized/gated_bidim.cuh). + +A CTA covers 256 output columns and walks 32-row bands. Each thread loads its rows of the act and +gate halves (and of the incoming gradient, in backward) into registers, computes the gated op +there with the elementwise math of the shared-memory port (quantize_mxfp8_gated_g2s.py, which +mirrors gated_mxfp8.cuh), rounds it to BF16 as the CUDA kernels do, and quantizes the band in +both directions from registers with the bidimensional register-tile quantizer +(quantize_mxfp8_g2r_bidim.py). The backward quantizes its two output halves one after the other +from the same band. Both write the bytes gated_mxfp8.cuh writes. +""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass.utils import HardwareInfo +from cutlass import Boolean, Float32, Int32, Int64 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + create_l2_policy, + cvt_f32x2_to_bf16x2, + max_num_bits_per_copy_ld_global, + x2_hi_to_f32_bf16, + x2_lo_to_f32_bf16, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + MXFP8GatedQuantizeConfig, + derive_swizzled_scale_layout, +) +from .quantize_mxfp8_g2r_bidim import ( + quantize_mxfp8_rmem_tile_bidimensional, + thread_tile_bounds, + thread_tile_layout, +) +from .quantize_mxfp8_gated_g2s import _gated_element, _split_gated_scales + + +@cute.jit +def _row_tile( + mT: cute.Tensor, + row_tile: Int32, + cta_col: Int32, + TILER: cutlass.Constexpr, + ALIGNMENT: cutlass.Constexpr[int], +) -> cute.Tensor: + """The (row_tile, cta_col) tile of a row-major mT, which may be a column slice of a wider + tensor (row stride 2N): the tile's base takes the 64-bit product once and offsets within the + tile are 32-bit, as in the CUDA register-tiled kernel. Every row stride is N or 2N, and the C++ + dispatcher requires N % 32 == 0.""" + row_stride = cute.assume(Int32(mT.layout.stride[0]), divby=32) + base = Int64(row_tile) * Int64(TILER[0]) * Int64(row_stride) + Int64(cta_col) * Int64(TILER[1]) + return cute.make_tensor( + (mT.iterator + base).align(ALIGNMENT), + cute.make_layout(TILER, stride=(row_stride, 1)), + ) + + +class MXFP8GatedQuantizeG2RKernel: + """Register-resident gated MXFP8 cast: BF16 [M, 2N] -> E4M3 or E5M2, both directions or + columnwise only, any M, N % 32 == 0 (the C++ dispatcher's requirement).""" + + _CTA_THREADS_Y = 8 + _CTA_THREADS_X = 32 + _ELEMENTS_Y_PER_THREAD = 4 # 8 warps x 4 rows = one 32-row band + _ELEMENTS_X_PER_THREAD = 8 # 32 lanes x 8 columns = 256 columns + _ROWS_IN_FLIGHT = 4 # input rows loaded at a time + # Row bands per CTA, as in gated_bidim.cuh: halve the grid while it holds more than + # _WALK_TARGET_CTAS CTAs, walk at least _MIN_WALK bands unless that leaves fewer than two CTAs + # per SM, and at most _MAX_WALK. + _WALK_TARGET_CTAS = 16384 + _MIN_WALK = 2 + _MAX_WALK = 64 + + @staticmethod + def supports(cfg: MXFP8GatedQuantizeConfig) -> bool: + """BF16 with columnwise output (rowwise optional); GEMM-swizzled scales in the forward + only, like the shared-memory port.""" + return ( + cfg.DTYPE is cutlass.BFloat16 + and cfg.COLWISE + and not (cfg.IS_BWD and cfg.WITH_GEMM_SWIZZLED_SCALES) + ) + + @staticmethod + @cute.jit + def eligible_to_dispatch( + mGatedInput: cute.Tensor, + mS_col: cute.Tensor, + ) -> Boolean: + # The C++ dispatcher checks 16-byte alignment of the data pointers. Columnwise scales are + # stored eight bytes per thread, and grid.y stays within limits. + M = mGatedInput.shape[0] + return (M + 31) // 32 <= 65535 and mS_col.iterator.toint() % 8 == 0 + + def __init__(self, cfg: MXFP8GatedQuantizeConfig): + assert self.supports(cfg), f"{cfg} is not supported by {type(self).__name__}" + self.cfg = cfg + # 4 resident CTAs per SM (64 registers), loading all 4 input rows at once. GB200, BF16, + # against gated_bidim.cuh, geometric mean over 5 shapes: forward 6 CTAs with 2 rows in + # flight (40 registers, spills) 0.91, 4 CTAs with 4 rows 0.93; backward 2 rows 0.92, + # 4 rows 0.94. + self._MIN_BLOCKS_PER_MP = 4 + from cuda.core import Device # pylint: disable=import-outside-toplevel,no-name-in-module + + self._sm_count = HardwareInfo(Device().device_id).get_device_multiprocessor_count() + + @cute.jit + def __call__( + self, + mGatedInput: cute.Tensor, # [M, 2N] = [act | gate] + mGrad: Optional[cute.Tensor], # [M, N], backward only + mO_row: Optional[cute.Tensor], # [M, C], C = N (forward) or 2N (backward) + mS_row: Optional[cute.Tensor], + mO_col: cute.Tensor, # [M, C], not transposed + mS_col: cute.Tensor, + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTUGLU only + beta2: Optional[Float32], # SiTUGLU only + stream: CUstream, + ): + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] MXFP8GatedQuantizeG2RKernel.__call__() with config: {self.cfg}\n" + ) + M = mGatedInput.shape[0] + N = mO_col.shape[1] if not self.cfg.IS_BWD else mGrad.shape[1] + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_X_PER_THREAD + + # eligible_to_dispatch checks this alignment: each thread stores eight columnwise scales. + mS_col = cute.make_tensor(mS_col.iterator.align(8), mS_col.layout) + mAct = cute.local_tile(mGatedInput, (M, N), (0, 0)) + mGate = cute.local_tile(mGatedInput, (M, N), (0, 1)) + + # Output parts: the forward output, or the two halves [d act | d gate] of the backward one. + mO_row1 = mS_row1 = mO_col1 = mS_col1 = None + if cutlass.const_expr(self.cfg.IS_BWD): + mS_row0, mS_row1, mS_col0, mS_col1 = _split_gated_scales( + mS_row, mS_col, M, N, self.cfg.ROWWISE, True, False + ) + mO_row0 = None + if cutlass.const_expr(self.cfg.ROWWISE): + mO_row0 = cute.local_tile(mO_row, (M, N), (0, 0)) + mO_row1 = cute.local_tile(mO_row, (M, N), (0, 1)) + mO_col0 = cute.local_tile(mO_col, (M, N), (0, 0)) + mO_col1 = cute.local_tile(mO_col, (M, N), (0, 1)) + else: + mO_row0, mO_col0 = mO_row, mO_col + mS_row0, mS_col0 = mS_row, mS_col + if cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row0, mS_col0 = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, True, mS_row, mS_col + ) + + grid_cols = (N + TILE_COLS - 1) // TILE_COLS + bands = (M + TILE_ROWS - 1) // TILE_ROWS + min_walk = Int32(self._MIN_WALK) + if grid_cols * (bands // self._MIN_WALK) < 2 * self._sm_count: + min_walk = Int32(1) + walk = Int32(1) + while (bands // walk) % 2 == 0 and grid_cols * (bands // walk) > self._WALK_TARGET_CTAS: + walk *= 2 + if walk < min_walk and bands % min_walk == 0: + walk = min_walk + if walk > self._MAX_WALK: + walk = Int32(self._MAX_WALK) + + # A tensor with a row or column tail takes the kernel whose tiles are predicated. + if M % TILE_ROWS != 0 or N % TILE_COLS != 0: + self.kernel( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + walk, + TAILS=True, + ).launch( + grid=[grid_cols, bands // walk], + block=[self._CTA_THREADS_Y * self._CTA_THREADS_X], + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) + else: + self.kernel( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + walk, + TAILS=False, + ).launch( + grid=[grid_cols, bands // walk], + block=[self._CTA_THREADS_Y * self._CTA_THREADS_X], + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mAct: cute.Tensor, + mGate: cute.Tensor, + mGrad: Optional[cute.Tensor], + mO_row0: Optional[cute.Tensor], + mS_row0: Optional[cute.Tensor], + mO_col0: cute.Tensor, + mS_col0: cute.Tensor, + mO_row1: Optional[cute.Tensor], + mS_row1: Optional[cute.Tensor], + mO_col1: Optional[cute.Tensor], + mS_col1: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + walk: Int32, + TAILS: cutlass.Constexpr[bool], + ): + """Quantize the CTA's row bands: blockIdx.y and every gridDim.y-th band after it, walk + bands in all.""" + cta_col, first_cta_row, _ = cute.arch.block_idx() + _, grid_rows, _ = cute.arch.grid_dim() + cache_policy = create_l2_policy(True) + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_X_PER_THREAD + for step in cutlass.range(walk, unroll=1): + row_tile = step * grid_rows + first_cta_row + # With TAILS, only the bands that cross the tensor's edge take the predicated body. + if cutlass.const_expr(TAILS): + if (row_tile + 1) * TILE_ROWS > mAct.shape[0] or ( + cta_col + 1 + ) * TILE_COLS > mAct.shape[1]: + self._gated_row_tile( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + row_tile, + cta_col, + cache_policy, + TAILS=True, + ) + else: + self._gated_row_tile( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + row_tile, + cta_col, + cache_policy, + TAILS=False, + ) + else: + self._gated_row_tile( + mAct, + mGate, + mGrad, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + row_tile, + cta_col, + cache_policy, + TAILS=False, + ) + + @cute.jit + def _quantize_part( + self, + rX: cute.Tensor, + mO_row: Optional[cute.Tensor], + mS_row: Optional[cute.Tensor], + mO_col: cute.Tensor, + mS_col: cute.Tensor, + row_tile: Int32, + cta_col: Int32, + rows_left: Int32, + cols_ok: Boolean, + cache_policy: Int64, + TAILS: cutlass.Constexpr[bool], + ): + """Quantize one output part of the band from its register tile.""" + TILE_ROWS = self._CTA_THREADS_Y * self._ELEMENTS_Y_PER_THREAD + TILE_COLS = self._CTA_THREADS_X * self._ELEMENTS_X_PER_THREAD + TILER = (TILE_ROWS, TILE_COLS) + cta_coord = (row_tile, cta_col) + gO_row_tile = gS_row_tile = None + if cutlass.const_expr(self.cfg.ROWWISE): + gO_row_tile = _row_tile(mO_row, row_tile, cta_col, TILER, 16) + gS_row_tile = cute.local_tile(mS_row, (TILE_ROWS, TILE_COLS // 32), cta_coord) + gO_col_tile = _row_tile(mO_col, row_tile, cta_col, TILER, 16) + gS_col_tile = cute.local_tile(mS_col, (TILE_ROWS // 32, TILE_COLS), cta_coord) + quantize_mxfp8_rmem_tile_bidimensional( + rX, + gO_row_tile, + gS_row_tile, + gO_col_tile, + gS_col_tile, + ROWWISE=self.cfg.ROWWISE, + COLWISE=True, + # The gated kernel's block amaxes ignore NaN like the generic kernel's. + FUSED=True, + EDGE=TAILS, + rows_left=rows_left, + cols_ok=cols_ok, + DTYPE=self.cfg.DTYPE, + FP8_DTYPE=self.cfg.FP8_DTYPE, + MAX_NORM_RCP=self.cfg.MAX_NORM_RCP, + WITH_GEMM_SWIZZLED_SCALES=self.cfg.WITH_GEMM_SWIZZLED_SCALES, + CTA_THREADS_Y=self._CTA_THREADS_Y, + CTA_THREADS_X=self._CTA_THREADS_X, + ELEMENTS_Y_PER_THREAD=self._ELEMENTS_Y_PER_THREAD, + ELEMENTS_X_PER_THREAD=self._ELEMENTS_X_PER_THREAD, + cache_policy=cache_policy, + ) + + @cute.jit + def _thread_view(self, mT: cute.Tensor, row_tile: Int32, cta_col: Int32) -> cute.Tensor: + """The calling thread's (ELEMENTS_X_PER_THREAD, ELEMENTS_Y_PER_THREAD) values of the + (row_tile, cta_col) band of mT, in global memory.""" + ELEMENTS_Y = self._ELEMENTS_Y_PER_THREAD + ELEMENTS_X = self._ELEMENTS_X_PER_THREAD + TILER = (self._CTA_THREADS_Y * ELEMENTS_Y, self._CTA_THREADS_X * ELEMENTS_X) + tidx, _, _ = cute.arch.thread_idx() + tv_layout = thread_tile_layout( + self._CTA_THREADS_Y, self._CTA_THREADS_X, ELEMENTS_Y, ELEMENTS_X + ) + return cute.composition(_row_tile(mT, row_tile, cta_col, TILER, 16), tv_layout)[tidx, None] + + @cute.jit + def _load_rows( + self, + gT_thread: cute.Tensor, + first_row: cutlass.Constexpr[int], + rows_left: Int32, + cache_policy: Int64, + TAILS: cutlass.Constexpr[bool], + ) -> cute.Tensor: + """Rows first_row .. first_row + _ROWS_IN_FLIGHT - 1 of the thread's values, in + registers; with TAILS, rows from rows_left on load as zeros.""" + DTYPE = self.cfg.DTYPE + ELEMENTS_X = self._ELEMENTS_X_PER_THREAD + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + DTYPE, + num_bits_per_copy=min(ELEMENTS_X * DTYPE.width, max_num_bits_per_copy_ld_global()), + invariant=False, + ) + rT = cute.make_rmem_tensor((ELEMENTS_X, self._ROWS_IN_FLIGHT), DTYPE) + for t in cutlass.range_constexpr(self._ROWS_IN_FLIGHT): + if cutlass.const_expr(TAILS): + rT[None, t].fill(0) + row_pred = cute.make_rmem_tensor( + cute.make_layout((ELEMENTS_X, (1,)), stride=(0, (0,))), Boolean + ) + row_pred[0] = first_row + t < rows_left + cute.copy( + load_atom, + gT_thread[(None, first_row + t),], + rT[None, t], + pred=row_pred, + cache_policy=cache_policy, + ) + else: + cute.copy( + load_atom, + gT_thread[(None, first_row + t),], + rT[None, t], + cache_policy=cache_policy, + ) + return rT + + @cute.jit + def _gated_row( + self, + rAct_row: cute.Tensor, + rGate_row: cute.Tensor, + rGrad_row: Optional[cute.Tensor], + rX0_row: cute.Tensor, + rX1_row: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + ): + """The gated op of one row of the thread's values, a BF16 pair at a time.""" + IS_BWD = self.cfg.IS_BWD + rA_i32 = cute.recast_tensor(rAct_row, Int32) + rG_i32 = cute.recast_tensor(rGate_row, Int32) + rX0_i32 = cute.recast_tensor(rX0_row, Int32) + if cutlass.const_expr(IS_BWD): + rD_i32 = cute.recast_tensor(rGrad_row, Int32) + rX1_i32 = cute.recast_tensor(rX1_row, Int32) + for i in cutlass.range_constexpr(self._ELEMENTS_X_PER_THREAD // 2): + grad_lo = grad_hi = None + if cutlass.const_expr(IS_BWD): + grad_lo = x2_lo_to_f32_bf16(rD_i32[i]) + grad_hi = x2_hi_to_f32_bf16(rD_i32[i]) + out_lo, gate_out_lo = _gated_element( + x2_lo_to_f32_bf16(rA_i32[i]), + x2_lo_to_f32_bf16(rG_i32[i]), + grad_lo, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=self.cfg.GATED_ACTIVATION, + IS_BWD=IS_BWD, + ) + out_hi, gate_out_hi = _gated_element( + x2_hi_to_f32_bf16(rA_i32[i]), + x2_hi_to_f32_bf16(rG_i32[i]), + grad_hi, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=self.cfg.GATED_ACTIVATION, + IS_BWD=IS_BWD, + ) + rX0_i32[i] = cvt_f32x2_to_bf16x2(out_lo, out_hi) + if cutlass.const_expr(IS_BWD): + rX1_i32[i] = cvt_f32x2_to_bf16x2(gate_out_lo, gate_out_hi) + + @cute.jit + def _gated_row_tile( + self, + mAct: cute.Tensor, + mGate: cute.Tensor, + mGrad: Optional[cute.Tensor], + mO_row0: Optional[cute.Tensor], + mS_row0: Optional[cute.Tensor], + mO_col0: cute.Tensor, + mS_col0: cute.Tensor, + mO_row1: Optional[cute.Tensor], + mS_row1: Optional[cute.Tensor], + mO_col1: Optional[cute.Tensor], + mS_col1: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + row_tile: Int32, + cta_col: Int32, + cache_policy: Int64, + TAILS: cutlass.Constexpr[bool], + ): + """Load one band's inputs, compute the gated op in registers and quantize its parts.""" + IS_BWD = self.cfg.IS_BWD + DTYPE = self.cfg.DTYPE + ELEMENTS_Y = self._ELEMENTS_Y_PER_THREAD + ELEMENTS_X = self._ELEMENTS_X_PER_THREAD + TILE_ROWS = self._CTA_THREADS_Y * ELEMENTS_Y + TILE_COLS = self._CTA_THREADS_X * ELEMENTS_X + TILER = (TILE_ROWS, TILE_COLS) + tv_layout = thread_tile_layout( + self._CTA_THREADS_Y, self._CTA_THREADS_X, ELEMENTS_Y, ELEMENTS_X + ) + + # With TAILS every band is predicated: rows and columns past the tensor load as zeros, + # which every gated op maps to zero, and are not stored. + rows_left = Int32(ELEMENTS_Y) + cols_ok = Boolean(True) + if cutlass.const_expr(TAILS): + rows_left, cols_ok = thread_tile_bounds( + TILER, + tv_layout, + mAct.shape[0] - row_tile * TILE_ROWS, + mAct.shape[1] - cta_col * TILE_COLS, + ) + + gAct = self._thread_view(mAct, row_tile, cta_col) + gGate = self._thread_view(mGate, row_tile, cta_col) + gGrad = None + if cutlass.const_expr(IS_BWD): + gGrad = self._thread_view(mGrad, row_tile, cta_col) + + # The gated op in FP32, rounded to BF16 as gated_mxfp8.cuh does before taking amaxes: + # rX0 is the forward output or the act half's gradient, rX1 the gate half's gradient. + # Inputs are loaded _ROWS_IN_FLIGHT rows at a time, as in gated_bidim.cuh, which keeps + # fewer of them live. + rX0 = cute.make_rmem_tensor((ELEMENTS_X, ELEMENTS_Y), DTYPE) + rX1 = None + if cutlass.const_expr(IS_BWD): + rX1 = cute.make_rmem_tensor((ELEMENTS_X, ELEMENTS_Y), DTYPE) + for first_row in cutlass.range_constexpr(0, ELEMENTS_Y, self._ROWS_IN_FLIGHT): + rAct = self._load_rows(gAct, first_row, rows_left, cache_policy, TAILS=TAILS) + rGate = self._load_rows(gGate, first_row, rows_left, cache_policy, TAILS=TAILS) + rGrad = None + if cutlass.const_expr(IS_BWD): + rGrad = self._load_rows(gGrad, first_row, rows_left, cache_policy, TAILS=TAILS) + for t in cutlass.range_constexpr(self._ROWS_IN_FLIGHT): + rGrad_row = rX1_row = None + if cutlass.const_expr(IS_BWD): + rGrad_row = rGrad[None, t] + rX1_row = rX1[None, first_row + t] + self._gated_row( + rAct[None, t], + rGate[None, t], + rGrad_row, + rX0[None, first_row + t], + rX1_row, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ) + + self._quantize_part( + rX0, + mO_row0, + mS_row0, + mO_col0, + mS_col0, + row_tile, + cta_col, + rows_left, + cols_ok, + cache_policy, + TAILS=TAILS, + ) + if cutlass.const_expr(IS_BWD): + self._quantize_part( + rX1, + mO_row1, + mS_row1, + mO_col1, + mS_col1, + row_tile, + cta_col, + rows_left, + cols_ok, + cache_policy, + TAILS=TAILS, + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py new file mode 100644 index 00000000000..46c635e3815 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -0,0 +1,821 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Gated-activation MXFP8 casts in CuTeDSL, ported from gated_mxfp8.cuh. + +The kernel reuses MXFP8QuantizeGeneralG2SKernel's TMA pipeline and its plain-cast SMEM tile +quantizer. What it adds is an elementwise pass that runs on each staged tile before it is +quantized: it combines the act and gate halves of the input (and the incoming gradient, in +backward) and writes the result back into the act (and gate) input buffer in the input dtype. +That is exactly the truncation the CUDA kernel applies before computing block amaxes, so the +plain cast that follows reproduces its output. + +Gated input is [M, 2N] with [act | gate] along the last dim. Every operand is addressed through +(M, N) views (stride 2N, column offset 0 or N), mirroring the CUDA kernel's tensor maps: + forward : out[M, N] = act(x[:, :N]) * gate(x[:, N:]) + backward: out[M, 2N] = [d/d act | d/d gate] given grad[M, N] +""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Boolean, Float32, Int32 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + fabs_f32, + is_packed16, + truncate_f32_bf16, + truncate_f32_f16, +) +from transformer_engine.common.CuTeDSL.activations import ( + act_clamped_silu, + act_situ_gate, + dact_dsitu_gate, + situ_up, + dsitu_up, + sigmoidf, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + GATED_ACTIVATIONS, + MXFP8GatedQuantizeConfig, + MXFP8_BLOCK_SCALING_SIZE, + derive_swizzled_scale_layout, +) +from .quantize_mxfp8_g2s_general import ( + MXFP8QuantizeGeneralG2SKernel, + quantize_mxfp8_smem_tile_general, +) + + +@cute.jit +def _select_f32(pred: Boolean, if_true: Float32, if_false: Float32) -> Float32: + result = if_false + if pred: + result = if_true + return result + + +@cute.jit +def _gated_element( + act_elt: Float32, + gate_elt: Float32, + grad_elt: Optional[Float32], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + ACTIVATION: cutlass.Constexpr[str], + IS_BWD: cutlass.Constexpr[bool], +): + """One element of the gated op in FP32; returns (act output, gate output or None). + + Mirrors the elementwise math of gated_mxfp8.cuh's columnwise pass (forward and backward), + operation by operation. + """ + out_gate = None + # Transform the gate. Gating is ideally the identity, with derivative 1. + gate = gate_elt + dgate = None + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + # |g| > limit is gated_mxfp8.cuh's (g > limit || g < -limit), NaN included + dgate = _select_f32(fabs_f32(gate_elt) > limit, Float32(0.0), Float32(1.0)) + gate = cute.arch.fmin(cute.arch.fmax(-limit, gate_elt), limit) + glu_linear_offset + elif cutlass.const_expr(ACTIVATION == "situglu"): + dgate = dsitu_up(gate_elt, beta2) + gate = situ_up(gate_elt, beta2) + + if cutlass.const_expr(not IS_BWD): + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + act = act_clamped_silu(act_elt, limit, alpha) + elif cutlass.const_expr(ACTIVATION == "situglu"): + act = act_situ_gate(act_elt, beta1) + else: + act = GATED_ACTIVATIONS[ACTIVATION][0](act_elt) + out_act = act * gate + else: + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + x = cute.arch.fmin(act_elt, limit) + s = sigmoidf(alpha * x) + act_x = x * s + dact_x = _select_f32( + act_elt <= limit, s + s * (Float32(1.0) - s) * alpha * x, Float32(0.0) + ) + elif cutlass.const_expr(ACTIVATION == "situglu"): + act_x = act_situ_gate(act_elt, beta1) + dact_x = dact_dsitu_gate(act_elt, beta1) + elif cutlass.const_expr(ACTIVATION == "swiglu"): + s = sigmoidf(act_elt) + act_x = act_elt * s + dact_x = act_elt * s * (Float32(1.0) - s) + s + else: + act_fn, dact_fn = GATED_ACTIVATIONS[ACTIVATION] + act_x = act_fn(act_elt) + dact_x = dact_fn(act_elt) + out_act = dact_x * grad_elt * gate + if cutlass.const_expr(dgate is None): + out_gate = act_x * grad_elt + else: + out_gate = dgate * act_x * grad_elt + return out_act, out_gate + + +@cute.jit +def gated_elementwise_smem_tile( + sAct_tile: cute.Tensor, + sGate_tile: cute.Tensor, + sGrad_tile: Optional[cute.Tensor], + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + ACTIVATION: cutlass.Constexpr[str], + IS_BWD: cutlass.Constexpr[bool], + DTYPE: cutlass.Constexpr[type], + TILE_X: cutlass.Constexpr[int], +): + """Compute the gated op on one staged tile, in place, rounded to the input dtype. + + Thread tidx owns column tidx of the tile, which is also the element set it owns in the + colwise quantization pass, so a colwise-only consumer needs no barrier after this. + """ + tidx, _, _ = cute.arch.thread_idx() + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), + val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), + ) + sAct_thread = cute.composition(sAct_tile, tv_layout)[tidx, None] + sGate_thread = cute.composition(sGate_tile, tv_layout)[tidx, None] + if cutlass.const_expr(IS_BWD): + sGrad_thread = cute.composition(sGrad_tile, tv_layout)[tidx, None] + if cutlass.const_expr(is_packed16(DTYPE)): + truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 + + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + grad_elt = None + if cutlass.const_expr(IS_BWD): + grad_elt = Float32(sGrad_thread[i]) + out_act, out_gate = _gated_element( + Float32(sAct_thread[i]), + Float32(sGate_thread[i]), + grad_elt, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=ACTIVATION, + IS_BWD=IS_BWD, + ) + # Numerical truncation: round to the input dtype, like the CUDA kernels do before + # taking the amax. The cast below is then exact. + if cutlass.const_expr(is_packed16(DTYPE)): + out_act = truncate_f32(out_act) + if cutlass.const_expr(IS_BWD): + out_gate = truncate_f32(out_gate) + sAct_thread[i] = DTYPE(out_act) + if cutlass.const_expr(IS_BWD): + sGate_thread[i] = DTYPE(out_gate) + + +@cute.jit +def quantize_mxfp8_smem_tile_general_partial( + sX_tile: cute.Tensor, + sO_row_tile: Optional[cute.Tensor], + mS_row_stage: Optional[cute.Tensor], + sO_col_tile: Optional[cute.Tensor], + mS_col_stage: Optional[cute.Tensor], + tile_row_start: Int32, + tile_col_start: Int32, + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[type], + FP8_DTYPE: cutlass.Constexpr[type], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], +): + """A partial quantize_mxfp8_smem_tile_general function with some args fixed.""" + quantize_mxfp8_smem_tile_general( + sX_tile, + None, + sO_row_tile, + mS_row_stage, + sO_col_tile, + mS_col_stage, + None, + tile_row_start, + tile_col_start, + M, + N, + ROWWISE=ROWWISE, + COLWISE=COLWISE, + WITH_AMAX=False, + WITH_ACT=False, + WITH_DACT=False, + ACTIVATION=None, + ACTIVATION_TABLE="none", + MAX_NORM_RCP=MAX_NORM_RCP, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + # The gated op maps the zeros TMA fills out-of-bounds elements with to zero. + SKIP_INPUT_MASKING=True, + SKIP_SCALE_BOUNDS=False, + CACHE_ACTIVATION=False, + DBIAS_REDUCTION_IN_COLWISE=False, + COLWISE_DBIAS_REDUCTION_ONLY=False, + DBIAS_REDUCTION_IN_ROWWISE=False, + ) + + +@cute.jit +def _split_gated_scales( + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + SWIZZLED: cutlass.Constexpr[bool], +): + """Act and gate views of the backward output's scales; returns + (act_row, gate_row, act_col, gate_col). + + The [M, 2N] output is [d act | d gate] and its scales are laid out for the whole output. Each + view has the shape of an [M, N] tensor's scales and the strides of the whole output's, and + starts where the whole output's layout places that half's first scale. For swizzled scales + this needs the halves to be whole 128x4 tiles (N % 128 == 0, required by the C++ dispatcher). + """ + act_row = gate_row = act_col = gate_col = None + if cutlass.const_expr(SWIZZLED): + mS_row, mS_col = derive_swizzled_scale_layout(M, 2 * N, ROWWISE, COLWISE, mS_row, mS_col) + half_tiles = N // 128 + if cutlass.const_expr(ROWWISE): + row_tiler = ( + cute.make_layout((32, 4, cute.ceil_div(M, 128))), # all rows + cute.make_layout((4, half_tiles)), # one half's scale columns + ) + act_row = cute.local_tile(mS_row, row_tiler, (0, 0)) + gate_row = cute.local_tile(mS_row, row_tiler, (0, 1)) + if cutlass.const_expr(COLWISE): + col_tiler = ( + cute.make_layout((4, cute.ceil_div(cute.ceil_div(M, 32), 4))), # all scale rows + cute.make_layout((32, 4, half_tiles)), # one half's columns + ) + act_col = cute.local_tile(mS_col, col_tiler, (0, 0)) + gate_col = cute.local_tile(mS_col, col_tiler, (0, 1)) + else: + # Compact scales are row-major, so each half is a block of scale columns + if cutlass.const_expr(ROWWISE): + row_tiler = (mS_row.shape[0], N // MXFP8_BLOCK_SCALING_SIZE) + act_row = cute.local_tile(mS_row, row_tiler, (0, 0)) + gate_row = cute.local_tile(mS_row, row_tiler, (0, 1)) + if cutlass.const_expr(COLWISE): + col_tiler = (mS_col.shape[0], N) + act_col = cute.local_tile(mS_col, col_tiler, (0, 0)) + gate_col = cute.local_tile(mS_col, col_tiler, (0, 1)) + return act_row, gate_row, act_col, gate_col + + +class MXFP8GatedQuantizeG2SKernel(MXFP8QuantizeGeneralG2SKernel): + """gated_mxfp8.cuh in CuTeDSL. The general kernel's pipeline, loading the act and gate halves + (and grad, backward) each stage and storing one (forward) or two (backward) quantized halves. + + Inherits the plain-cast tile config, which is gated_mxfp8.cuh's: 64 threads per CTA, each CTA + covering two vertically stacked 32x64 tiles. + """ + + def __init__(self, cfg: MXFP8GatedQuantizeConfig): + super().__init__(cfg) + # Same occupancy hint as the general kernel with a fused activation + self._MIN_BLOCKS_PER_MP = 1 + # Every pipeline stage has its own SMEM buffers, so a CTA with fewer tiles than stages + # would allocate buffers it never fills + assert ( + self._NUM_STAGES <= self._NUM_TILES + ), f"{self._NUM_STAGES} pipeline stages for {self._NUM_TILES} tiles per CTA" + + @cute.jit + def __call__( + self, + mGatedInput: cute.Tensor, # [M, 2N] = [act | gate] + mGrad: Optional[cute.Tensor], # [M, N], bwd only + # Outputs, with C = N (FWD) or 2N (BWD) + mO_row: Optional[cute.Tensor], # [M, C] + mS_row: Optional[cute.Tensor], # [roundup(M, 128), roundup(C / 32, 4)] + mO_col: Optional[cute.Tensor], # [M, C], not transposed + mS_col: Optional[cute.Tensor], # [roundup(ceil(M / 32), 4), roundup(C, 128)] + # Special activation params, None if not used by the activation + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTUGLU only + beta2: Optional[Float32], # SiTUGLU only + stream: CUstream, + ): + """ + mGatedInput: [a | g] with shape M x 2N, where a is the activation half and g is the gate half. + mGrad: None for FWD, [dy] with shape [M, N] for BWD + FWD mode: + out = quantize(act(a) * gate(g)) with shape [M, N] + BWD mode: + out = quantize([ dy * gate(g) * dact(a) | dy * act(a) * dgate(g) ]) with shape [M, 2N] + """ + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] MXFP8GatedQuantizeG2SKernel.__call__() with config: {self.cfg}\n" + ) + + M = mGatedInput.shape[0] + # Pick any mO that is used just to get N + mO = mO_row if mO_row is not None else mO_col + # FWD: mO is [M, N]; BWD: mGrad is [M, N] + N = mO.shape[1] if not self.cfg.IS_BWD else mGrad.shape[1] + + # The act and gate halves of the [M, 2N] input + mAct = cute.local_tile(mGatedInput, (M, N), (0, 0)) + mGate = cute.local_tile(mGatedInput, (M, N), (0, 1)) + + mO_row_dact = mO_row_dgate = mO_col_dact = mO_col_dgate = None # BWD only + mS_row_dact = mS_row_dgate = mS_col_dact = mS_col_dgate = None # BWD only + if cutlass.const_expr(self.cfg.IS_BWD): + # The [M, 2N] output is [d act | d gate]; the kernel takes its two halves + # TODO(kainingz): how to implement this in a more elegant way + mS_row_dact, mS_row_dgate, mS_col_dact, mS_col_dgate = _split_gated_scales( + mS_row, + mS_col, + M, + N, + self.cfg.ROWWISE, + self.cfg.COLWISE, + self.cfg.WITH_GEMM_SWIZZLED_SCALES, + ) + if cutlass.const_expr(self.cfg.ROWWISE): + mO_row_dact = cute.local_tile(mO_row, (M, N), (0, 0)) + mO_row_dgate = cute.local_tile(mO_row, (M, N), (0, 1)) + if cutlass.const_expr(self.cfg.COLWISE): + mO_col_dact = cute.local_tile(mO_col, (M, N), (0, 0)) + mO_col_dgate = cute.local_tile(mO_col, (M, N), (0, 1)) + # In BWD mode we won't use these, and we should set them to None so we don't build TMA atoms for them + mO_row = mO_col = mS_row = mS_col = None + elif cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, mS_col = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col + ) + + atoms = self._make_tma_atoms( + mGrad, + mAct, + mGate, + mO_row, + mO_col, + mO_row_dact, + mO_row_dgate, + mO_col_dact, + mO_col_dgate, + ) + # The TMA atoms are spread from a list, which pylint cannot count + self.kernel( # pylint: disable=no-value-for-parameter + M, + N, + mS_row, + mS_col, + mS_row_dact, + mS_row_dgate, + mS_col_dact, + mS_col_dgate, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + *atoms, + ).launch( + grid=[ + cute.ceil_div(N, self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), + ], + block=[self._THREADS_PER_CTA], + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) + + @cute.jit + def _make_tma_atoms( + self, + mGrad: Optional[cute.Tensor], + mAct: cute.Tensor, + mGate: cute.Tensor, + mO_row: Optional[cute.Tensor], # FWD only + mO_col: Optional[cute.Tensor], # FWD only + mO_row_dact: Optional[cute.Tensor], # BWD only + mO_row_dgate: Optional[cute.Tensor], # BWD only + mO_col_dact: Optional[cute.Tensor], # BWD only + mO_col_dgate: Optional[cute.Tensor], # BWD only + ): + """Build TMA atoms for input / quantized output; scales will not be loaded / stored by TMA.""" + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + cta_tiler = (self._TILE_ROWS, self._TILE_COLS) + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + atoms = [] + for op, mT in ( + (op_load, mGrad), + (op_load, mAct), + (op_load, mGate), + (op_store, mO_row), + (op_store, mO_col), + (op_store, mO_row_dact), + (op_store, mO_row_dgate), + (op_store, mO_col_dact), + (op_store, mO_col_dgate), + ): + if cutlass.const_expr(mT is None): + atoms.extend([None, None]) + else: + atoms.extend( + cute.nvgpu.cpasync.make_tiled_tma_atom( + op, mT, smem_tile_layout, cta_tiler, num_multicast=1 + ) + ) + return atoms + + @cute.kernel + def kernel( + self, + M: Int32, + N: Int32, + mS_row: Optional[cute.Tensor], # FWD only + mS_col: Optional[cute.Tensor], # FWD only + mS_row_dact: Optional[cute.Tensor], # BWD only + mS_row_dgate: Optional[cute.Tensor], # BWD only + mS_col_dact: Optional[cute.Tensor], # BWD only + mS_col_dgate: Optional[cute.Tensor], # BWD only + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + # TMA for mGrad, only used in BWD + tma_atom_grad: Optional[cute.CopyAtom], + tma_src_grad: Optional[cute.Tensor], + # TMA for mAct and mGate, used in both FWD and BWD + tma_atom_act: cute.CopyAtom, + tma_src_act: cute.Tensor, + tma_atom_gate: cute.CopyAtom, + tma_src_gate: cute.Tensor, + # TMA for mO_row/col, only used in FWD + tma_atom_row: Optional[cute.CopyAtom], + tma_dst_row: Optional[cute.Tensor], + tma_atom_col: Optional[cute.CopyAtom], + tma_dst_col: Optional[cute.Tensor], + # TMA for mO_row/col_dact/dgate, only used in BWD + tma_atom_row_dact: Optional[cute.CopyAtom], + tma_dst_row_dact: Optional[cute.Tensor], + tma_atom_row_dgate: Optional[cute.CopyAtom], + tma_dst_row_dgate: Optional[cute.Tensor], + tma_atom_col_dact: Optional[cute.CopyAtom], + tma_dst_col_dact: Optional[cute.Tensor], + tma_atom_col_dgate: Optional[cute.CopyAtom], + tma_dst_col_dgate: Optional[cute.Tensor], + ): + """ + FWD mode: + mAct: activation half input + mGate: gate half input + mO_row/col: quantized output for row / column directions + mS_row/col: output scales for row / column directions + + BWD mode: + mGrad: incoming gradient + mAct: activation half input from FWD pass + mGate: gate half input from FWD pass + mO_row/col_dact: quantized activation gradient for row / column directions + mO_row/col_dgate: quantized gate gradient for row / column directions + mS_row/col_dact: output scales for activation gradient row / column directions + mS_row/col_dgate: output scales for gate gradient row / column directions + mO_row/col: quantized gradient (both activation and gate) for row / column directions + mS_row/col: output scales for gradient (both activation and gate) row / column directions + + mAct: both FWD and BWD + mGate: both FWD and BWD + mGrad: BWD only + mO_row/col: both FWD (N) and BWD (2N) + mS_row/col: both FWD (N) and BWD (2N) + mO_row/col_dact: BWD only + mO_row/col_dgate: BWD only + mS_row/col_dact: BWD only + mS_row/col_dgate: BWD only + """ + + cfg = self.cfg + DTYPE = cfg.DTYPE + TILE_ROWS = self._TILE_ROWS + TILE_COLS = self._TILE_COLS + NUM_STAGES = self._NUM_STAGES + NUM_TILES = self._NUM_TILES + + # The scales of one (TILE_ROWS, TILE_COLS) tile + row_scale_tile = (TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + col_scale_tile = (TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) + + # Shared memory: one ((TILE_ROWS, TILE_COLS), NUM_STAGES) ring buffer per operand + smem = cutlass.memory.SmemAllocator() + + @cute.struct + class BarrierStorage: # pylint: disable=missing-class-docstring + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * NUM_STAGES] + + storage = smem.allocate(BarrierStorage) + stages_layout = cute.make_layout( + ((TILE_ROWS, TILE_COLS), NUM_STAGES), + stride=((TILE_COLS, 1), TILE_ROWS * TILE_COLS), + ) + + # Allocate SMEM buffers for inputs + sAct = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + sGate = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + sGrad = None + if cutlass.const_expr(cfg.IS_BWD): + sGrad = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + + # Allocate SMEM buffers for outputs + sO_row = sO_col = None + sO_row_dact = sO_row_dgate = sO_col_dact = sO_col_dgate = None + if cutlass.const_expr(cfg.ROWWISE): + if cutlass.const_expr(cfg.IS_BWD): + sO_row_dact = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + sO_row_dgate = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + else: + sO_row = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + if cutlass.const_expr(cfg.COLWISE): + if cutlass.const_expr(cfg.IS_BWD): + sO_col_dact = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + sO_col_dgate = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + else: + sO_col = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bidx, bidy, _ = cute.arch.block_idx() + + # Issue TMA descriptor prefetch using the TMA warp (warp 0) + if warp_idx == 0: + for load_atom in (tma_atom_act, tma_atom_gate, tma_atom_grad): + if cutlass.const_expr(load_atom is not None): + cute.nvgpu.cpasync.prefetch_descriptor(load_atom) + + # TMA warp (warp 0)'s elected thread produces (issues TMA), every warp consumes + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) + # Input has 2 tiles (act + gate) in FWD, 3 tiles (grad + act + gate) in BWD + num_inputs = 3 if cfg.IS_BWD else 2 + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_storage.data_ptr(), + num_stages=NUM_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + # Every input tile of a stage arrives on the same barrier + tx_count=num_inputs * TILE_ROWS * TILE_COLS * DTYPE.width // 8, + cta_layout_vmnk=None, + ) + prod_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, NUM_STAGES) + cons_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, NUM_STAGES) + + cta_coord = (bidy, bidx) + # Shape of each individual tile. A CTA processes one tile at a time + tile_shape = (TILE_ROWS, TILE_COLS) + # The block processed by a CTA, consisting of NUM_TILES vertically stacked tiles + cta_block = (NUM_TILES * TILE_ROWS, TILE_COLS) + + # Prepare parameters to pass to the SMEM tile quantize function by dact / dgate / output, where + # for each group, input, quantized output are SMEM tiles, and scales are GMEM tiles + cta_block_row_scales = (NUM_TILES * TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + cta_block_col_scales = (NUM_TILES * TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) + if cutlass.const_expr(cfg.IS_BWD): + if cutlass.const_expr(cfg.ROWWISE): + mS_row_dact = cute.local_tile(mS_row_dact, cta_block_row_scales, cta_coord) + mS_row_dgate = cute.local_tile(mS_row_dgate, cta_block_row_scales, cta_coord) + if cutlass.const_expr(cfg.COLWISE): + mS_col_dact = cute.local_tile(mS_col_dact, cta_block_col_scales, cta_coord) + mS_col_dgate = cute.local_tile(mS_col_dgate, cta_block_col_scales, cta_coord) + # In BWD mode we output gradient for both the act and gate halves, which are stored in sAct and sGate respectively + tile_quantizer_params = [ + ( + sAct, + sO_row_dact, + mS_row_dact, + sO_col_dact, + mS_col_dact, + ), # We will write unquantized dact to sAct in-place + ( + sGate, + sO_row_dgate, + mS_row_dgate, + sO_col_dgate, + mS_col_dgate, + ), # We will write unquantized dgate to sGate in-place + ] + else: + if cutlass.const_expr(cfg.ROWWISE): + mS_row = cute.local_tile(mS_row, cta_block_row_scales, cta_coord) + if cutlass.const_expr(cfg.COLWISE): + mS_col = cute.local_tile(mS_col, cta_block_col_scales, cta_coord) + # In FWD mode we only output the result after the gated op, which is stored in sAct + tile_quantizer_params = [ + (sAct, sO_row, mS_row, sO_col, mS_col) + ] # We will write unquantized output to sAct in-place + + # Prepare input TMA (atom, smem partition, gmem partition) triples: act and gate, and BWD's grad + inputs = [] + for atom_in, src_in, sX_in in ( + (tma_atom_act, tma_src_act, sAct), + (tma_atom_gate, tma_src_gate, sGate), + (tma_atom_grad, tma_src_grad, sGrad), + ): + if cutlass.const_expr(atom_in is not None): + tXsX_in, tXgX_in = cute.nvgpu.cpasync.tma_partition( + atom_in, + 0, + cute.make_layout(1), + sX_in, + cute.zipped_divide(cute.local_tile(src_in, cta_block, cta_coord), tile_shape), + ) + inputs.append((atom_in, tXsX_in, tXgX_in)) + + # Prepare output TMA (atom, smem partition, gmem partition) triples of the outputs this config has + TMA_atom_and_partitions = [] + for atom_out, dst_out, sO_out in ( + (tma_atom_row, tma_dst_row, sO_row), # FWD only + (tma_atom_col, tma_dst_col, sO_col), # FWD only + (tma_atom_row_dact, tma_dst_row_dact, sO_row_dact), # BWD only + (tma_atom_row_dgate, tma_dst_row_dgate, sO_row_dgate), # BWD only + (tma_atom_col_dact, tma_dst_col_dact, sO_col_dact), # BWD only + (tma_atom_col_dgate, tma_dst_col_dgate, sO_col_dgate), # BWD only + ): + if cutlass.const_expr(atom_out is not None): + tXsO_out, tXgO_out = cute.nvgpu.cpasync.tma_partition( + atom_out, + 0, + cute.make_layout(1), + sO_out, + cute.zipped_divide(cute.local_tile(dst_out, cta_block, cta_coord), tile_shape), + ) + TMA_atom_and_partitions.append((atom_out, tXsO_out, tXgO_out)) + + # Barrier init must be visible to every thread + cute.arch.sync_threads() + + # Prologue: fill the pipeline + if warp_idx == 0: + for s in cutlass.range_constexpr(NUM_STAGES): + mainloop_pipeline.producer_acquire(prod_state) + for load_atom, load_sX, load_gX in inputs: + cute.copy( + load_atom, + load_gX[(None, (s, 0))], + load_sX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + for tile_idx in cutlass.range(NUM_TILES, unroll=1): + mainloop_pipeline.consumer_wait(cons_state) + # We only reuse previously used SMEM buffers when we have more tiles to process than pipeline stages, and only + # after we've already finished one round and wrap around to the start of the ring buffer + if cutlass.const_expr(NUM_TILES > NUM_STAGES): + if tile_idx >= NUM_STAGES: + # Only allow at most NUM_STAGES - 1 store groups reading SMEM on the fly, so the tile NUM_STAGES + # ago must finish before we reuse its SMEM buffer + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + # The tile should be ready to use now + + stage_idx = cons_state.index + tile_coord = (tile_idx, 0) # within this CTA's block + + # Select current stage's SMEM tile + sAct_tile = sAct[(None, stage_idx)] + sGate_tile = sGate[(None, stage_idx)] + sGrad_tile = None + if cutlass.const_expr(cfg.IS_BWD): + sGrad_tile = sGrad[(None, stage_idx)] + + # Compute the gated op on this tile in fp32, cast back to the original dtype, + # and store the result in-place in the SMEM tile without quantization. + # FWD: write the output into sAct_tile + # BWD: write the act half's gradient to sAct_tile, and write the gate half's gradient to sGate_tile + gated_elementwise_smem_tile( + sAct_tile, # we also store the FWD output / BWD act half's gradient here + sGate_tile, # we also store the BWD gate half's gradient here + sGrad_tile, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=cfg.GATED_ACTIVATION, + IS_BWD=cfg.IS_BWD, + DTYPE=DTYPE, + TILE_X=TILE_COLS, + ) + + # The rowwise quantization pass reads elements other threads wrote so we need to sync here + if cutlass.const_expr(cfg.ROWWISE): + cute.arch.sync_threads() + + # Quantize every output of this tile into its SMEM buffers and scales + for sX_out, sO_row_out, mS_row_cta, sO_col_out, mS_col_cta in tile_quantizer_params: + # Slice the SMEM / GMEM tile for the current iteration + sO_row_tile = None + mS_row_stage = None + if cutlass.const_expr(cfg.ROWWISE): + sO_row_tile = sO_row_out[(None, stage_idx)] + mS_row_stage = cute.flatten( + cute.local_tile(mS_row_cta, row_scale_tile, tile_coord) + ) + sO_col_tile = None + mS_col_stage = None + if cutlass.const_expr(cfg.COLWISE): + sO_col_tile = sO_col_out[(None, stage_idx)] + mS_col_stage = cute.flatten( + cute.local_tile(mS_col_cta, col_scale_tile, tile_coord) + ) + # Use the general MXFP8 SMEM tile quantize function to quantize it + quantize_mxfp8_smem_tile_general_partial( + sX_out[(None, stage_idx)], # The SMEM tile to quantize + sO_row_tile, # SMEM tile for rowwise quantized output + mS_row_stage, # GMEM tile for rowwise quantized scales + sO_col_tile, # SMEM tile for columnwise quantized output + mS_col_stage, # GMEM tile for columnwise quantized scales + # The tile's first row and column in the whole tensor, for bounds checks + (bidy * NUM_TILES + tile_idx) * TILE_ROWS, + bidx * TILE_COLS, + M, + N, + ROWWISE=cfg.ROWWISE, + COLWISE=cfg.COLWISE, + MAX_NORM_RCP=cfg.MAX_NORM_RCP, + DTYPE=DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_COLS, + TILE_Y=TILE_ROWS, + ) + + # Make the SMEM writes visible to the TMA's async proxy before it reads them + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + mainloop_pipeline.consumer_release(cons_state) + + if warp_idx == 0: + # The tile quantize function has already written the quantized outputs to SMEM, + # so we can issue TMA stores from SMEM to GMEM + for store_atom, store_sO, store_gO in TMA_atom_and_partitions: + cute.copy( + store_atom, + store_sO[(None, stage_idx)], + store_gO[(None, (tile_idx, 0))], + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # Refill the stage just released with the tile NUM_STAGES ahead + if warp_idx == 0: + next_tile_idx = tile_idx + NUM_STAGES + if next_tile_idx < NUM_TILES: + mainloop_pipeline.producer_acquire(prod_state) + for load_atom, load_sX, load_gX in inputs: + cute.copy( + load_atom, + load_gX[(None, (next_tile_idx, 0))], + load_sX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # Only wait until the in-flight TMA stores have read SMEM, so the CTA retires early + cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index ed05aa1c3c6..90c61d03052 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -27,6 +27,16 @@ logger = logging.getLogger("transformer_engine.cutedsl.utils") +@cute.jit +def max_num_bits_per_copy_ld_global() -> int: + """Widest global load to issue: 256-bit loads (`ld.global.v8.b32`) need PTX ISA 8.8, + i.e. CUDA 12.9. Older toolchains split each one into two 128-bit loads, as + cast_rowwise.cuh / cast_bidim.cuh do via util/ptx.cuh's ld_global_nc_b32x8.""" + if cutlass.const_expr(cutlass.target_version(min_version="12.9")): + return 256 + return 128 + + def device_compute_capability(device_id: Optional[int] = None) -> tuple: """(major, minor) compute capability of a CUDA device (current by default), or (0, 0) if it can't be queried.""" from cuda.core import Device # pylint: disable=no-name-in-module @@ -83,6 +93,72 @@ def exp2f_rcp(scale_e8m0) -> Float32: return result +@cute.jit +def exp2_bf16x2_rcp(scale_e8m0) -> Uint32: + """Pack two BF16 copies of the E8M0 reciprocal computed by exp2f_rcp.""" + # These reciprocal powers of two are exact in BF16. Taking the upper bits + # also preserves exp2f_rcp's chosen subnormal and NaN encodings. + bf16_bits = exp2f_rcp(scale_e8m0).bitcast(Uint32) >> Uint32(16) + return bf16_bits | (bf16_bits << Uint32(16)) + + +@cute.jit +def bf16_pair_magnitude(pair: Int32) -> Uint32: + """Fold both BF16 halves and clear the sign set by max.xorsign.abs.""" + swapped = cute.arch.inline_ptx( + "prmt.b32 {$w0}, {$r0}, {$r0}, 0x1032;", + write_only_types=[Int32], + read_only_args=[pair], + ) + return Uint32(abs_max_x2_bf16(pair, swapped)) & Uint32(0x7FFF) + + +@cute.jit +def bf16_mx_scale_reciprocal(magnitude: Uint32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Exact BF16 reciprocal scale bits, including E8M0 zero, Inf and NaN cases.""" + rounded = (magnitude + Uint32(31)) & Uint32(0x7F80) + if rounded < Uint32(MAX_EXPONENT << 7): # pylint: disable=consider-using-max-builtin + rounded = Uint32(MAX_EXPONENT << 7) + reciprocal = Uint32((254 + MAX_EXPONENT) << 7) - rounded + if magnitude >= Uint32(0x7F80): + reciprocal = Uint32(0x0040) + if magnitude != Uint32(0x7F80): + reciprocal = Uint32(0x7FFF) + return reciprocal + + +@cute.jit +def bf16_mx_scale_reciprocal_pair(pair: Int32, MAX_EXPONENT: cutlass.Constexpr[int]) -> Uint32: + """Compute two reciprocal scales without borrowing across packed BF16 halves.""" + magnitudes = Uint32(pair) & Uint32(0x7FFF7FFF) + rounded = (magnitudes + Uint32(0x001F001F)) & Uint32(0xFF80FF80) + floor = (MAX_EXPONENT << 23) | (MAX_EXPONENT << 7) + # Packed BF16 max accepts registers, so materialize the constant inside PTX. + clamped = cute.arch.inline_ptx( + "{ .reg.b32 floor_pair; " + f"mov.b32 floor_pair, {floor}; " + "max.xorsign.abs.bf16x2 {$w0}, {$r0}, floor_pair; }", + write_only_types=[Int32], + read_only_args=[rounded.bitcast(Int32)], + ) + bias = ((254 + MAX_EXPONENT) << 23) | ((254 + MAX_EXPONENT) << 7) + reciprocal = Uint32(bias) - Uint32(clamped) + if ((magnitudes + Uint32(0x00800080)) & Uint32(0x80008000)) != Uint32(0): + lo = bf16_mx_scale_reciprocal(magnitudes & Uint32(0x7FFF), MAX_EXPONENT) + hi = bf16_mx_scale_reciprocal(magnitudes >> Uint32(16), MAX_EXPONENT) + reciprocal = lo | (hi << Uint32(16)) + return reciprocal + + +@cute.jit +def bf16_mx_scale_bytes_pair(reciprocal: Uint32) -> Uint32: + """Gather the E8M0 bytes of two packed BF16 reciprocals into the low half.""" + bytes_pair = (Uint32(0x01FE01FE) - ((reciprocal >> Uint32(7)) & Uint32(0x00FF00FF))) & Uint32( + 0x00FF00FF + ) + return (bytes_pair & Uint32(0xFF)) | ((bytes_pair >> Uint32(8)) & Uint32(0xFF00)) + + @cute.jit def pack_f32x2(lo: Float32, hi: Float32) -> Int64: """Pack two f32 scalars into a single 64-bit register (`floatx2` layout). @@ -210,6 +286,16 @@ def x2_hi_to_f32_bf16(bits: Int32) -> Float32: return ((bits >> Int32(16)) << Int32(16)).bitcast(Float32) +@cute.jit +def cvt_f32x2_to_bf16x2(lo: Float32, hi: Float32) -> Int32: + """Round two FP32 values to BF16 (round-to-nearest-even) and pack them, lo in the low half.""" + return cute.arch.inline_ptx( + "cvt.rn.bf16x2.f32 {$w0}, {$r0}, {$r1};", + write_only_types=[Int32], + read_only_args=[hi, lo], + ) + + @cute.jit def truncate_f32_bf16(val: Float32) -> Float32: """Round FP32 to BF16 precision (round-to-nearest-even), keeping FP32 storage.""" @@ -258,3 +344,15 @@ def truncate_f32_f16(val: Float32) -> Float32: def is_packed16(dtype) -> bool: """True if `dtype` is one of the 16-bit packed input formats.""" return dtype is cutlass.BFloat16 or dtype is cutlass.Float16 + + +# TODO(kainingz): remove this once https://github.com/NVIDIA/cutlass/pull/3717 is merged and released +@cute.jit +def create_l2_policy(evict_last: cutlass.Constexpr[bool], fraction: Float32 = 1.0) -> Int64: + """Create the L2 eviction policy used by register-resident casts.""" + priority = "evict_last" if evict_last else "evict_first" + return cute.arch.inline_ptx( + f"createpolicy.fractional.L2::{priority}.b64 {{$w0}}, {{$r0}};", + write_only_types=[Int64], + read_only_args=[Float32(fraction)], + ) diff --git a/transformer_engine/common/CuTeDSL/utils_fp8.py b/transformer_engine/common/CuTeDSL/utils_fp8.py index b3b73e7998d..9a972e39469 100644 --- a/transformer_engine/common/CuTeDSL/utils_fp8.py +++ b/transformer_engine/common/CuTeDSL/utils_fp8.py @@ -180,6 +180,49 @@ def mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu: bool = False) -> Callable[..., return _build_mul_f32x4_cvt_f32x4_to_fp8x4(fp8_dtype, relu) +def mul_bf16x2x2_cvt_bf16x4_to_fp8x4(fp8_dtype) -> Callable[..., Uint32]: + """Return a multiply/cast for two BF16 pairs, producing four packed FP8 bytes. + + CUDA 13.1+ supports direct BF16-to-FP8 conversion. Older compilers use + FP32 fused multiply-add, matching the corresponding ptx::mul_cvt_4x overload. + """ + out_op = "e5m2x2" if fp8_dtype is Float8E5M2 else "e4m3x2" + if cutlass.target_version(min_version="13.1"): + body = ( + ".reg.b32 y0, y1;\n\t" + "mul.rn.bf16x2 y0, {$r0}, {$r1};\n\t" + "mul.rn.bf16x2 y1, {$r2}, {$r3};\n\t" + f"cvt.rn.satfinite.{out_op}.bf16x2 z0, y0;\n\t" + f"cvt.rn.satfinite.{out_op}.bf16x2 z1, y1;\n\t" + ) + else: + body = ( + ".reg.b16 x0,x1,x2,x3,s0,s1,s2,s3;\n\t" + ".reg.f32 y0,y1,y2,y3;\n\t" + "mov.b32 {x0,x1}, {$r0};\n\t" + "mov.b32 {s0,s1}, {$r1};\n\t" + "mov.b32 {x2,x3}, {$r2};\n\t" + "mov.b32 {s2,s3}, {$r3};\n\t" + "fma.rn.f32.bf16 y0, x0, s0, 0f00000000;\n\t" + "fma.rn.f32.bf16 y1, x1, s1, 0f00000000;\n\t" + "fma.rn.f32.bf16 y2, x2, s2, 0f00000000;\n\t" + "fma.rn.f32.bf16 y3, x3, s3, 0f00000000;\n\t" + f"cvt.rn.satfinite.{out_op}.f32 z0, y1, y0;\n\t" + f"cvt.rn.satfinite.{out_op}.f32 z1, y3, y2;\n\t" + ) + asm = "{\n.reg.b16 z0,z1;\n\t" + body + "mov.b32 {$w0}, {z0,z1};\n}" + + @cute.jit + def fn(in0: Int32, scale0: Uint32, in1: Int32, scale1: Uint32) -> Uint32: + return cute.arch.inline_ptx( + asm, + write_only_types=[Uint32], + read_only_args=[in0, scale0, in1, scale1], + ) + + return fn + + def _build_mul_f32x2_cvt_packed16x4_to_fp8x4( in_dtype, fp8_dtype, relu: bool = False ) -> Callable[..., Uint32]: 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/core/grouped_layout.cuh b/transformer_engine/common/cast/core/grouped_layout.cuh index 8337e448158..94f3dbe124d 100644 --- a/transformer_engine/common/cast/core/grouped_layout.cuh +++ b/transformer_engine/common/cast/core/grouped_layout.cuh @@ -36,39 +36,66 @@ inline bool full_tile_1D_tensor(const Tensor *const t, const size_t elems_per_bl namespace kernel { constexpr size_t THREADS_PER_BLOCK = 256; -template -__global__ void __launch_bounds__(THREADS_PER_BLOCK) - reduce_dbias_kernel(OType *const dbias_output, const float *const dbias_partial, - const size_t rows, const size_t cols) { - using ComputeVec = Vec; - using OutputVec = Vec; - const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; +//! Threads per CTA of reduce_dbias_kernel, one column each. +constexpr int REDUCE_DBIAS_THREADS = 64; +//! Rows of a column copied into shared memory per pipeline stage. +constexpr int REDUCE_DBIAS_STAGE_ROWS = 32; +//! Pipeline stages: while one stage is added up, the others are in flight. +constexpr int REDUCE_DBIAS_STAGES = 4; - if (thread_id * nvec >= cols) { +/*! \brief dbias[c] = sum over rows r of partial[r][c], added in row order from 0. + * + * The adds of a column form one sequential chain, so the result does not depend on how the + * work is split; what the split decides is how many loads are in flight. One thread per + * column copies its column into shared memory with cp.async, REDUCE_DBIAS_STAGE_ROWS rows per + * stage and REDUCE_DBIAS_STAGES stages deep, and adds each stage in row order once it has + * landed. A thread reads back only the rows it copied itself, so no barrier is needed. + */ +template +__global__ void __launch_bounds__(REDUCE_DBIAS_THREADS) + reduce_dbias_kernel(OType *const dbias_output, const float *const dbias_partial, + const size_t rows, const size_t cols) { + __shared__ float stages[REDUCE_DBIAS_STAGES][REDUCE_DBIAS_STAGE_ROWS][REDUCE_DBIAS_THREADS]; + const size_t col = static_cast(blockIdx.x) * REDUCE_DBIAS_THREADS + threadIdx.x; + if (col >= cols) { return; } - - const float *const thread_in_base = dbias_partial + thread_id * nvec; - OType *const thread_out_base = dbias_output + thread_id * nvec; - - ComputeVec ldg_vec; - ComputeVec acc_vec; - acc_vec.clear(); - for (int i = 0; i < rows; ++i) { - ldg_vec.load_from(thread_in_base + i * cols); -#pragma unroll - for (int e = 0; e < nvec; ++e) { - acc_vec.data.elt[e] += ldg_vec.data.elt[e]; + const float *in = dbias_partial + col; + const size_t num_stages = (rows + REDUCE_DBIAS_STAGE_ROWS - 1) / REDUCE_DBIAS_STAGE_ROWS; + + // Copies stage k of this thread's column into its slot of stages[k % REDUCE_DBIAS_STAGES], and + // commits it as one cp.async group; past the last stage it commits an empty group, so that + // every iteration below can wait for the same number of groups. + const auto issue = [&](const size_t k) { + if (k < num_stages) { + const size_t r0 = k * REDUCE_DBIAS_STAGE_ROWS; + const int n = static_cast(min(static_cast(REDUCE_DBIAS_STAGE_ROWS), rows - r0)); + float(*dst)[REDUCE_DBIAS_THREADS] = stages[k % REDUCE_DBIAS_STAGES]; + for (int i = 0; i < n; ++i) { + const uint32_t dst_addr = + static_cast(__cvta_generic_to_shared(&dst[i][threadIdx.x])); + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" ::"r"(dst_addr), + "l"(in + (r0 + i) * cols) + : "memory"); + } } - } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + }; - OutputVec stg_vec; #pragma unroll - for (int e = 0; e < nvec; ++e) { - stg_vec.data.elt[e] = static_cast(acc_vec.data.elt[e]); + for (int k = 0; k < REDUCE_DBIAS_STAGES - 1; ++k) issue(k); + float acc = 0.0f; + for (size_t k = 0; k < num_stages; ++k) { + issue(k + REDUCE_DBIAS_STAGES - 1); + // All but the REDUCE_DBIAS_STAGES - 1 most recent groups are complete: stage k has landed. + asm volatile("cp.async.wait_group %0;\n" ::"n"(REDUCE_DBIAS_STAGES - 1) : "memory"); + const size_t r0 = k * REDUCE_DBIAS_STAGE_ROWS; + const int n = static_cast(min(static_cast(REDUCE_DBIAS_STAGE_ROWS), rows - r0)); + const float(*src)[REDUCE_DBIAS_THREADS] = stages[k % REDUCE_DBIAS_STAGES]; + for (int i = 0; i < n; ++i) acc += src[i][threadIdx.x]; } - stg_vec.store_to(thread_out_base); + dbias_output[col] = static_cast(acc); } template @@ -131,11 +158,10 @@ void reduce_dbias(const float *workspace_ptr, Tensor *dbias, const size_t rows, constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType); NVTE_CHECK(cols % reduce_dbias_nvec == 0, "Unsupported shape."); - const size_t reduce_dbias_num_blocks = DIVUP(cols, THREADS_PER_BLOCK * reduce_dbias_nvec); + const size_t num_blocks = DIVUP(cols, static_cast(REDUCE_DBIAS_THREADS)); - reduce_dbias_kernel - <<>>( - reinterpret_cast(dbias->data.dptr), workspace_ptr, rows, cols); + reduce_dbias_kernel<<>>( + reinterpret_cast(dbias->data.dptr), workspace_ptr, rows, cols); NVTE_CHECK_CUDA(cudaGetLastError()); } diff --git a/transformer_engine/common/cast/dispatch/gated.cuh b/transformer_engine/common/cast/dispatch/gated.cuh index 06e8f0e306c..aa4a5f56636 100644 --- a/transformer_engine/common/cast/dispatch/gated.cuh +++ b/transformer_engine/common/cast/dispatch/gated.cuh @@ -19,6 +19,10 @@ #include "../fp8/gated_fp8.cuh" #include "../mxfp8/gated_mxfp8.cuh" +#ifdef NVTE_WITH_CUTEDSL +#include "../mxfp8/gated_mxfp8_cutedsl.cuh" +#endif + namespace transformer_engine { namespace dispatch { @@ -86,8 +90,16 @@ void quantize_gated_fwd_helper(const NVTETensor nvte_input, NVTETensor nvte_outp NVTE_CHECK(is_supported_by_CC_100(), "Gated FWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); Tensor dummy_grad_tensor; - mxfp8::quantize_gated(input, dummy_grad_tensor, - output, p, stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_gated_cutedsl( + input, dummy_grad_tensor, output, p, stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::quantize_gated(input, dummy_grad_tensor, + output, p, stream); + } break; } default: @@ -176,8 +188,16 @@ void quantize_gated_bwd_helper(const NVTETensor nvte_grad, const NVTETensor nvte NVTE_CHECK(is_supported_by_CC_100(), "Gated BWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); - mxfp8::quantize_gated(gated_input, grad, output, p, - stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_gated_cutedsl( + gated_input, grad, output, p, stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::quantize_gated(gated_input, grad, output, p, + stream); + } break; } default: diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index aa1d16fd894..bed5a31e76e 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/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 44dea1a7c76..3efa8819cbc 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -20,6 +20,7 @@ #include "../../util/math.h" #include "../../util/ptx_arch_spec.cuh" #include "../../utils.cuh" +#include "specialized/gated_bidim.cuh" #include "swizzle.cuh" namespace transformer_engine { @@ -845,6 +846,22 @@ void quantize_gated(const Tensor &gated_input, const Tensor &grad, Tensor *outpu } } + // gated_bidim replaces the kernel below for SwiGLU and its gradient from BF16 to + // E4M3 or E5M2 in both directions (see gated_bidim::can_use), writing the same + // bytes. + if constexpr (std::is_same_v && + (std::is_same_v || std::is_same_v) && + std::is_same_v) { + if constexpr (ActOP == &silu && + (!IS_BWD || DActOP == &dsilu)) { + if (gated_bidim::can_use(gated_input, grad, *output, IS_BWD)) { + gated_bidim::quantize( + gated_input, grad, output, stream); + return; + } + } + } + switch (scaling_type) { case ScalingType::ROWWISE: { auto kernel = diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh new file mode 100644 index 00000000000..cb6a3a44d07 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh @@ -0,0 +1,207 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file gated_mxfp8_cutedsl.cuh + * \brief CuTeDSL backend of gated_mxfp8.cuh. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ + +#include +#include + +#include +#include +#include +#include +#include + +#include "../../common.h" +#include "../../tvm_ffi_bridge.h" +#include "../../util/cutedsl_utils.h" + +namespace transformer_engine { +namespace cutedsl_backend { + +using namespace tvm_ffi_bridge; + +namespace detail { + +// TMA needs 16-byte aligned global addresses. +inline bool is_16B_aligned(const void *ptr) { return reinterpret_cast(ptr) % 16 == 0; } + +} // namespace detail + +struct MXFP8GatedQuantConfig { + static constexpr const char *kEntrypointName = "get_mxfp8_gated_quantization_function"; + + DType dtype; // The input format + DType fp8_dtype; // The fp8 output format + bool rowwise; // If quantize rowwisely + bool colwise; // If quantize columnwisely + bool swizzled; // If the scale output is used for cudnn's swizzled layout + bool is_bwd; // Backward (dgated) rather than forward + GatedActivation activation; + uint32_t sm_arch = static_cast(cuda::sm_arch()); + + // Bit layout: dtype [3:0], fp8_dtype [7:4], flags [11:8], activation [17:12] and + // SM architecture [26:18]. + uint32_t to_id() const { + static_assert(static_cast(DType::kNumTypes) <= 16, + "DType no longer fits in the 4 bits to_id() gives it."); + static_assert(static_cast(GatedActivation::kNumTypes) <= 64, + "GatedActivation no longer fits in the 6 bits to_id() gives it."); + NVTE_CHECK(sm_arch < 512, "SM architecture no longer fits in the 9 bits to_id() gives it."); + return static_cast(dtype) | (static_cast(fp8_dtype) << 4) | + (static_cast(rowwise) << 8) | (static_cast(colwise) << 9) | + (static_cast(swizzled) << 10) | (static_cast(is_bwd) << 11) | + (static_cast(activation) << 12) | (sm_arch << 18); + } + + std::optional get_kernel() const { + static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); + return cache.get_or_load(*this); + } + + std::string to_key() const { + std::string key; + key.reserve(80); + key.append("cutedsl_mxfp8_gated_sm") + .append(std::to_string(sm_arch)) + .append("_") + .append(to_string(dtype)) + .append("_") + .append(to_string(fp8_dtype)) + .append("_") + .append(rowwise ? "1" : "0") + .append("_") + .append(colwise ? "1" : "0") + .append("_") + .append(swizzled ? "1" : "0") + .append("_") + .append(is_bwd ? "1" : "0") + .append("_") + .append(gated_activation_to_str(activation)); + return key; + } + + bool retrieve_func_from_python(const std::string &fn_name) const { + auto entrypoint = tvm::ffi::Function::GetGlobal(kEntrypointName); + if (!entrypoint.has_value()) { + return false; + } + tvm::ffi::Any result = + (*entrypoint)(tvm::ffi::String(fn_name), tvm::ffi::String(to_string(dtype)), + tvm::ffi::String(to_string(fp8_dtype)), rowwise, colwise, swizzled, is_bwd, + tvm::ffi::String(gated_activation_to_str(activation))); + return result.try_cast().value_or(false); + } +}; + +// Signature mirrors mxfp8::quantize_gated. Returns false to fall back to the CUDA kernel. +template +bool mxfp8_quantize_gated_cutedsl(const Tensor &gated_input, const Tensor &grad, Tensor *output, + const ParamOP &p, cudaStream_t stream) { + if (!TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { + maybe_warn_cutedsl_not_chosen("the CuTeDSL backend is disabled."); + return false; + } + constexpr GatedActivation activation = gated_activation_func_to_enum(); + if constexpr (activation == GatedActivation::kUnsupported) { + maybe_warn_cutedsl_not_chosen("the gated activation is not supported."); + return false; + } else { + const size_t rows = gated_input.flat_first_dim(); + const size_t cols = gated_input.flat_last_dim() / 2; + const MXFP8GatedQuantConfig config{/*dtype=*/gated_input.dtype(), + /*fp8_dtype=*/output->dtype(), + /*rowwise=*/output->has_data(), + /*colwise=*/output->has_columnwise_data(), + /*swizzled=*/output->with_gemm_swizzled_scales, + /*is_bwd=*/IS_BWD, + /*activation=*/activation}; + // The kernel offsets the gate half's swizzled scales by a fixed amount, which only matches + // the swizzled layout when the half spans whole 128-column scale tiles. + if (IS_BWD && config.swizzled && cols % 128 != 0) { + maybe_warn_cutedsl_not_chosen( + "backward gated activations with GEMM-swizzled scales need the last dimension of the " + "output halves to be divisible by 128."); + return false; + } + bool aligned = detail::is_16B_aligned(gated_input.data.dptr) && + (!config.rowwise || detail::is_16B_aligned(output->data.dptr)) && + (!config.colwise || detail::is_16B_aligned(output->columnwise_data.dptr)); + if constexpr (IS_BWD) { + aligned = aligned && detail::is_16B_aligned(grad.data.dptr); + } + if (!aligned) { + maybe_warn_cutedsl_not_chosen("an input or output is not 16-byte aligned."); + return false; + } + + std::optional fn = config.get_kernel(); + if (!fn.has_value()) { + return false; + } + checkCuDriverContext(stream); + + // Sanity checks and swizzled-scale padding, mirroring mxfp8::quantize_gated + if (config.rowwise) { + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); + } + if (config.colwise) { + NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); + } + if (config.swizzled && (cols % 128 != 0 || rows % 128 != 0)) { + if (config.rowwise) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->scale_inv.dptr, 0, + output->scale_inv.buffer_size_bytes(), stream)); + } + if (config.colwise) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->columnwise_scale_inv.dptr, 0, + output->columnwise_scale_inv.buffer_size_bytes(), stream)); + } + } + + // The activation's own parameters are runtime arguments; the kernel is compiled with None in + // the slots of the others, so they are passed as None (nullptr). + tvm::ffi::Any limit = nullptr, alpha = nullptr, glu_linear_offset = nullptr, beta1 = nullptr, + beta2 = nullptr; + if constexpr (std::is_same_v) { + limit = p.limit; + alpha = p.alpha; + glu_linear_offset = p.glu_linear_offset; + } else if constexpr (std::is_same_v) { + beta1 = p.beta1; + beta2 = p.beta2; + } + + const int32_t device_index = transformer_engine::cuda::current_device(); + DLTensorWrapper mGatedInput(gated_input.data, true, device_index); + DLTensorWrapper mGrad, mO_row, mS_row, mO_col, mS_col; + if constexpr (IS_BWD) { + mGrad = DLTensorWrapper(grad.data, true, device_index); + } + if (config.rowwise) { + mO_row = DLTensorWrapper(output->data, true, device_index); + mS_row = DLTensorWrapper(output->scale_inv, true, device_index); + } + if (config.colwise) { + mO_col = DLTensorWrapper(output->columnwise_data, true, device_index); + mS_col = DLTensorWrapper(output->columnwise_scale_inv, true, device_index); + } + (*fn)(&mGatedInput, &mGrad, &mO_row, &mS_row, &mO_col, &mS_col, limit, alpha, glu_linear_offset, + beta1, beta2, static_cast(stream)); + return true; + } +} + +} // namespace cutedsl_backend +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index f7749f65dd7..5b84445d465 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; @@ -256,11 +258,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 @@ -342,11 +345,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 @@ -441,11 +445,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 @@ -655,7 +661,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); @@ -969,58 +975,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 (E3). 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 431812db143..e7b9df5807b 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"); @@ -269,7 +296,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/cast/mxfp8/specialized/gated_bidim.cuh b/transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh new file mode 100644 index 00000000000..12684f093b8 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/gated_bidim.cuh @@ -0,0 +1,433 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file gated_bidim.cuh + * \brief Register-resident MXFP8 bidimensional quantize fused with SwiGLU or its gradient. + * + * Input [M, 2K] holds the activated half and the gate half side by side. The forward writes + * silu(act) * gate as [M, K]; the backward writes the gradients of both halves as [M, 2K], + * quantized independently, one after the other from the same 32-row band. A CTA covers 256 + * output columns and walks 32-row bands; the band stays in registers, and both scale + * directions are computed from there, instead of being staged through shared memory as in the + * generic gated kernel (gated_mxfp8.cuh). + * + * It writes the bytes the generic gated kernel writes: the same math, the same rounding to + * BF16, and scales only inside the tensor. + */ + +#ifndef TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_GATED_BIDIM_CUH_ +#define TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_GATED_BIDIM_CUH_ + +#include + +#include +#include + +#include "../../../common.h" +#include "../../../util/cuda_runtime.h" +#include "../../../util/math.h" +#include "../../../util/ptx.cuh" +#include "../../../utils.cuh" +#include "../swizzle.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace gated_bidim { + +namespace ptx = transformer_engine::ptx; + +namespace { + +/*! \brief Elementwise max(|a|, |b|) on a packed BF16 pair, as a raw word. */ +__device__ __forceinline__ unsigned abs_max_bf16x2(unsigned a, unsigned b) { + ptx::bf16x2 d; + ptx::abs_max_2x(d, reinterpret_cast(a), + reinterpret_cast(b)); + return reinterpret_cast(d); +} + +/*! \brief Widen a BF16 pair word to FP32, low half to x. Same result as ptx::up_cast, whose + * volatile asm the compiler cannot schedule as freely in the derivative loop. */ +__device__ __forceinline__ ptx::floatx2 widen_bf16x2(const unsigned v) { + return {__uint_as_float(v << 16), __uint_as_float(v & 0xffff0000u)}; +} + +/*! \brief Round an FP32 pair to a BF16 pair word, low lane in the low half. */ +__device__ __forceinline__ unsigned to_bf16x2(const ptx::floatx2& v) { + const __nv_bfloat162 r = __floats2bfloat162_rn(v.x, v.y); + return reinterpret_cast(r); +} + +/*! \brief The reciprocal of an E8M0 scale, 2^(127-e), broadcast to a BF16 pair. */ +__device__ __forceinline__ unsigned e8m0_to_bf16x2_reciprocal(unsigned e) { + const ptx::bf16x2 r = ptx::exp2f_rcp_2x(static_cast(e)); + return reinterpret_cast(r); +} + +/*! \brief E8M0 scales for two block amaxes, packed as ptx::float_to_e8m0_2x. + * + * The generic kernel's amax is a 0-seeded max that ignores NaN, so a block of + * only NaNs gets amax 0; the fmaxf reproduces that for the data-seeded maxima + * here. The conversion saturates, so the scale never reaches 255. + */ +template +__device__ __forceinline__ unsigned amax_to_e8m0_2x(float amax_hi, float amax_lo) { + constexpr float kMaxNormRcp = Quantized_Limits::max_norm_rcp; + return ptx::float_to_e8m0_2x(fmaxf(amax_hi * kMaxNormRcp, 0.0f), + fmaxf(amax_lo * kMaxNormRcp, 0.0f)); +} + +// Tile geometry. A CTA covers 256 columns and walks 32-row sub-tiles; 32 rows +// is the colwise MX block height, so the colwise reduction closes in the CTA. + +//! Columns a CTA covers. +constexpr int kTileCols = 256; +//! Elements in one MX block, and equivalently the colwise block height. +constexpr int kRowsPerMxBlock = 32; +//! MX groups (32 columns each) across the tile. +constexpr int kMxGroupsPerRow = kTileCols / 32; +//! A register-resident tile spreads its 256 columns over the 32 lanes of a warp. +constexpr int kColsPerLane = kTileCols / 32; +//! ...which a lane holds as BF16 pairs, one 32-bit word each. +constexpr int kWordsPerLane = kColsPerLane / 2; +//! ...read as one 128-bit vector per row. +constexpr int kVecLoadsPerRow = kWordsPerLane / 4; +//! Lanes that have to cooperate to cover one 32-column MX group. +constexpr int kLanesPerMxGroup = 32 / kColsPerLane; + +/*! \brief Scale two BF16 pairs, each by its own reciprocal scale pair, to four OType (E4M3 or + * E5M2). */ +template +__device__ __forceinline__ unsigned quantize_x4(unsigned in01, unsigned in23, unsigned scale01, + unsigned scale23) { + ptx::FPx4 out; + ptx::mul_cvt_4x(out, reinterpret_cast(in01), + reinterpret_cast(scale01), + reinterpret_cast(in23), + reinterpret_cast(scale23)); + return reinterpret_cast(out); +} + +/*! \brief Store one lane's quantized row: kColsPerLane FP8 bytes, one STG.64. */ +__device__ __forceinline__ void store_quantized(unsigned char* p, const unsigned* v) { + static_assert(kColsPerLane == 8, "One STG.64 covers exactly eight FP8 bytes."); + *(uint2*)p = *(const uint2*)v; +} + +// Row bands per CTA. Walking several bands deepens each CTA's work without starving SMs. +//! Grid size, in CTAs, that the walk length aims for. +constexpr long long kCtaTargetActivation = 16384; +//! Minimum walk, unless that leaves fewer than two CTAs per SM. +constexpr int kMinWalk = 2; +//! Maximum walk. +constexpr int kMaxWalkDact = 64; + +int pick_walk_length(int row_blocks, int grid_cols, long long target, int min_walk = kMinWalk) { + int walk = 1; + while ((row_blocks / walk) % 2 == 0 && (long long)grid_cols * (row_blocks / walk) > target) + walk *= 2; + if (walk < min_walk && row_blocks % min_walk == 0) walk = min_walk; + return walk; +} + +} // namespace + +constexpr int kThreadsPerCta = 256; +constexpr int kWarpsPerCta = kThreadsPerCta / 32; +constexpr int kRowsPerWarp = kRowsPerMxBlock / kWarpsPerCta; +constexpr int kRowsInFlight = 2; +// Resident CTAs per SM, the second __launch_bounds__ argument (GB200, 32768x7168): the forward +// is issue-bound and gains 16% from 6 CTAs over 4; the backward, which keeps more values live, +// loses 30% at 5. +template +constexpr int kMinBlocksPerSm = IS_BWD ? 4 : 6; +static_assert(kVecLoadsPerRow == 1, "One 16-byte vector per row and lane."); + +/*! \brief One output element of the gated activation, as the generic gated kernel computes it. + * + * kPart 0: the forward output, or the gradient of the activated half. kPart 1: the gradient of + * the gate half. Rounded to BF16 by the caller. + */ +template +__device__ __forceinline__ float gated_element(const float act_elt, const float gate_elt, + const float grad_elt) { + if constexpr (!IS_BWD) { + return silu(act_elt, {}) * gate_elt; + } else { + const float x = act_elt; + const float s = sigmoidf(x); + if constexpr (kPart == 0) { + const float dact_x = x * s * (1 - s) + s; + return dact_x * grad_elt * gate_elt; + } else { + const float act_x = x * s; + const float dgate_elt = 1.0f; + return dgate_elt * act_x * grad_elt; + } + } +} + +/*! \brief The gated activation of one BF16 pair of each input, rounded to a BF16 pair. */ +template +__device__ __forceinline__ unsigned gated_word(const unsigned act, const unsigned gate, + const unsigned grad) { + const ptx::floatx2 a = widen_bf16x2(act); + const ptx::floatx2 g = widen_bf16x2(gate); + const ptx::floatx2 d = widen_bf16x2(grad); + return to_bf16x2( + {gated_element(a.x, g.x, d.x), gated_element(a.y, g.y, d.y)}); +} + +template +__global__ void __launch_bounds__(kThreadsPerCta, kMinBlocksPerSm) + quantize_gated_kernel(const uint4* __restrict__ input, const uint4* __restrict__ grad, + unsigned char* __restrict__ rowwise_out, + unsigned char* __restrict__ rowwise_scales, + unsigned char* __restrict__ colwise_out, + unsigned char* __restrict__ colwise_scales, int M, int K, + int rowwise_scale_stride, int colwise_scale_stride, int iters) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + __shared__ __align__(16) unsigned col_amax_partials[kWarpsPerCta][kTileCols / 2]; + __shared__ __align__(16) unsigned col_scale_rcp_smem[kTileCols / 2]; + __shared__ __align__(8) unsigned char row_scale_bytes[kRowsPerMxBlock * kMxGroupsPerRow]; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int col0 = blockIdx.x * kTileCols; + // Output columns; the backward writes both halves. + const int out_cols = IS_BWD ? 2 * K : K; + const int k_vecs = K >> 3; // uint4 (8 BF16 values) per K columns + const int k_groups = (K + 31) / 32; + const bool full_width = !kTails || col0 + kTileCols <= K; + const bool lane_in_cols = col0 + kColsPerLane * lane < K; + +#pragma unroll 1 + for (int it = 0; it < iters; ++it) { + const int row0 = (it * gridDim.y + blockIdx.y) * kRowsPerMxBlock; + const int warp_row0 = row0 + warp * kRowsPerWarp; + // Quantize one output half (kPart; the forward has only part 0) of the band from its tile. + const auto quantize_part = [&](auto edge, auto part, const unsigned* tile_words) { + constexpr bool kEdge = decltype(edge)::value; + constexpr int kPart = decltype(part)::value; + const auto in_bounds = [&](int j) { return !kEdge || (lane_in_cols && warp_row0 + j < M); }; + + // ---- rowwise ---------------------------------------------------------------------------- + const size_t store_offset = + (size_t)warp_row0 * out_cols + kPart * K + col0 + kColsPerLane * lane; + { + unsigned char* row_out_ptr = rowwise_out + store_offset; +#pragma unroll + for (int j = 0; j < kRowsPerWarp; ++j) { + const unsigned* v = tile_words + j * kWordsPerLane; + unsigned amax = abs_max_bf16x2(abs_max_bf16x2(v[0], v[1]), abs_max_bf16x2(v[2], v[3])); +#pragma unroll + for (int lane_mask = 1; lane_mask < kLanesPerMxGroup; lane_mask <<= 1) + amax = abs_max_bf16x2(amax, __shfl_xor_sync(0xffffffffu, amax, lane_mask)); + amax = abs_max_bf16x2(amax, __byte_perm(amax, amax, 0x1032)); + const unsigned scale_byte = + amax_to_e8m0_2x(0.f, __uint_as_float(amax & 0x7fff0000u)); + if ((lane & (kLanesPerMxGroup - 1)) == 0) + row_scale_bytes[(warp * kRowsPerWarp + j) * kMxGroupsPerRow + + (lane / kLanesPerMxGroup)] = (unsigned char)scale_byte; + const unsigned rcp = e8m0_to_bf16x2_reciprocal(scale_byte); + unsigned row_out[kWordsPerLane / 2]; +#pragma unroll + for (int m = 0; m < kWordsPerLane / 2; ++m) + row_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], rcp, rcp); + if (in_bounds(j)) store_quantized(row_out_ptr + (size_t)j * out_cols, row_out); + } + } + + // ---- colwise amax: per-warp partial in registers, cross-warp fold in shared --------------- + { + unsigned col_amax[kWordsPerLane]; +#pragma unroll + for (int m = 0; m < kWordsPerLane; ++m) col_amax[m] = tile_words[m]; +#pragma unroll + for (int j = 1; j < kRowsPerWarp; ++j) +#pragma unroll + for (int m = 0; m < kWordsPerLane; ++m) + col_amax[m] = abs_max_bf16x2(col_amax[m], tile_words[j * kWordsPerLane + m]); + *(uint4*)(&col_amax_partials[warp][kWordsPerLane * lane]) = *(const uint4*)col_amax; + } + __syncthreads(); + if (tid < kTileCols / 2) { + unsigned col_amax = col_amax_partials[0][tid]; +#pragma unroll + for (int w = 1; w < kWarpsPerCta; ++w) + col_amax = abs_max_bf16x2(col_amax, col_amax_partials[w][tid]); + col_amax &= 0x7fff7fffu; + const ptx::floatx2 amax_2x = widen_bf16x2(col_amax); + const unsigned scale_pair = amax_to_e8m0_2x(amax_2x.y, amax_2x.x); + const int col = col0 + 2 * tid; + // The generic gated kernel writes scales only inside the tensor. + if (!kEdge || (col < K && row0 < M)) { + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + const size_t row_block_tiles = DIVUP(static_cast(M), size_t{128}); + colwise_scales[swizzle::gemm_swizzled_scale_idx(kPart * K + col, row0 / 32, + row_block_tiles)] = + static_cast(scale_pair); + colwise_scales[swizzle::gemm_swizzled_scale_idx(kPart * K + col + 1, row0 / 32, + row_block_tiles)] = + static_cast(scale_pair >> 8); + } else { + *(unsigned short*)(colwise_scales + (size_t)(row0 / 32) * colwise_scale_stride + + kPart * K + col) = (unsigned short)scale_pair; + } + } + const ptx::bf16x2 scale_rcp = + ptx::exp2f_rcp_2x_per_lane(__byte_perm(scale_pair, 0, 0x4140)); + col_scale_rcp_smem[tid] = reinterpret_cast(scale_rcp); + } + // Drain the row scales: one thread per row, one byte per group inside the tensor. + if (tid >= kTileCols / 2 && tid < kTileCols / 2 + kRowsPerMxBlock) { + const int r = tid - kTileCols / 2; + if (!kEdge || row0 + r < M) { + const int groups = full_width ? kMxGroupsPerRow : (K - col0 + 31) / 32; + const size_t col_tiles = DIVUP(static_cast(out_cols), size_t{128}); + for (int g = 0; g < groups; ++g) { + const int group = kPart * k_groups + (col0 >> 5) + g; + size_t idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + idx = swizzle::gemm_swizzled_scale_idx(row0 + r, group, col_tiles); + } else { + idx = (size_t)(row0 + r) * rowwise_scale_stride + group; + } + rowwise_scales[idx] = row_scale_bytes[r * kMxGroupsPerRow + g]; + } + } + } + __syncthreads(); + + // ---- colwise quantization --------------------------------------------------------------- + { + const uint4 rcp_vec = *(const uint4*)(col_scale_rcp_smem + kWordsPerLane * lane); + const unsigned* rcp = (const unsigned*)&rcp_vec; + unsigned char* col_out_ptr = colwise_out + store_offset; +#pragma unroll + for (int j = 0; j < kRowsPerWarp; ++j) { + const unsigned* v = tile_words + j * kWordsPerLane; + unsigned col_out[kWordsPerLane / 2]; +#pragma unroll + for (int m = 0; m < kWordsPerLane / 2; ++m) + col_out[m] = quantize_x4(v[2 * m], v[2 * m + 1], rcp[2 * m], rcp[2 * m + 1]); + if (in_bounds(j)) store_quantized(col_out_ptr + (size_t)j * out_cols, col_out); + } + } + }; + // Load the band once and compute its tiles: one for the forward, both gradient halves for the + // backward, which then quantizes them one after the other. + const auto parts = [&](auto edge) { + constexpr bool kEdge = decltype(edge)::value; + const auto in_bounds = [&](int j) { return !kEdge || (lane_in_cols && warp_row0 + j < M); }; + unsigned tile0[kRowsPerWarp * kWordsPerLane]; + unsigned tile1[IS_BWD ? kRowsPerWarp * kWordsPerLane : 1]; + const size_t in_offset = (size_t)warp_row0 * (2 * k_vecs) + (col0 >> 3) + lane; + const uint4* act_ptr = input + in_offset; + const uint4* gate_ptr = act_ptr + k_vecs; + const uint4* grad_ptr = grad + (size_t)warp_row0 * k_vecs + (col0 >> 3) + lane; +#pragma unroll + for (int h = 0; h < kRowsPerWarp / kRowsInFlight; ++h) { + uint4 act_vec[kRowsInFlight], gate_vec[kRowsInFlight], grad_vec[kRowsInFlight]; + asm volatile("" ::: "memory"); +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) { + const int j = h * kRowsInFlight + t; + const uint4 zero = make_uint4(0, 0, 0, 0); + act_vec[t] = in_bounds(j) ? act_ptr[(size_t)j * (2 * k_vecs)] : zero; + gate_vec[t] = in_bounds(j) ? gate_ptr[(size_t)j * (2 * k_vecs)] : zero; + grad_vec[t] = IS_BWD && in_bounds(j) ? grad_ptr[(size_t)j * k_vecs] : zero; + } + asm volatile("" ::: "memory"); +#pragma unroll + for (int t = 0; t < kRowsInFlight; ++t) { + const unsigned* a = (const unsigned*)&act_vec[t]; + const unsigned* g = (const unsigned*)&gate_vec[t]; + const unsigned* d = (const unsigned*)&grad_vec[t]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const int w = (h * kRowsInFlight + t) * kWordsPerLane + m; + tile0[w] = gated_word(a[m], g[m], d[m]); + if constexpr (IS_BWD) tile1[w] = gated_word(a[m], g[m], d[m]); + } + } + } + quantize_part(edge, std::integral_constant{}, tile0); + if constexpr (IS_BWD) quantize_part(edge, std::integral_constant{}, tile1); + }; + if constexpr (kTails) { + if (full_width && row0 + kRowsPerMxBlock <= M) { + parts(std::false_type{}); + } else { + parts(std::true_type{}); + } + } else { + parts(std::false_type{}); + } + } +#endif +} + +//! Whether this kernel takes the request: BF16 [M, 2K] -> E4M3 or E5M2, both directions, +//! K % 32 == 0 (the generic gated kernel's requirement), aligned buffers. +inline bool can_use(const Tensor& input, const Tensor& grad, const Tensor& output, + const bool is_bwd) { + if (transformer_engine::cuda::sm_arch() < 100) return false; + if (input.dtype() != DType::kBFloat16 || + (output.dtype() != DType::kFloat8E4M3 && output.dtype() != DType::kFloat8E5M2)) { + return false; + } + if (!output.has_data() || !output.has_columnwise_data()) return false; + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim() / 2; + constexpr size_t kMaxRows = std::numeric_limits::max() - 64; + if (rows == 0 || cols == 0 || cols % 32 != 0 || rows > kMaxRows || + 2 * cols > std::numeric_limits::max() / kRowsPerMxBlock) { + return false; + } + return is_aligned_ptr(input.data.dptr, 16) && (!is_bwd || is_aligned_ptr(grad.data.dptr, 16)) && + is_aligned_ptr(output.data.dptr, 8) && is_aligned_ptr(output.columnwise_data.dptr, 8) && + is_aligned_ptr(output.columnwise_scale_inv.dptr, 2); +} + +template +void quantize(const Tensor& input, const Tensor& grad, Tensor* output, cudaStream_t stream) { + const int M = static_cast(input.flat_first_dim()); + const int K = static_cast(input.flat_last_dim() / 2); + const int bands = (M + kRowsPerMxBlock - 1) / kRowsPerMxBlock; + const int grid_cols = (K + kTileCols - 1) / kTileCols; + int min_walk = kMinWalk; + if ((long long)grid_cols * (bands / min_walk) < 2LL * cuda::sm_count()) min_walk = 1; + int iters = pick_walk_length(bands, grid_cols, kCtaTargetActivation, min_walk); + if (iters > kMaxWalkDact) iters = kMaxWalkDact; + const dim3 grid(grid_cols, bands / iters); + const auto run = [&](auto tails) { + constexpr bool kTails = decltype(tails)::value; + quantize_gated_kernel + <<>>( + (const uint4*)input.data.dptr, (const uint4*)(IS_BWD ? grad.data.dptr : nullptr), + (unsigned char*)output->data.dptr, (unsigned char*)output->scale_inv.dptr, + (unsigned char*)output->columnwise_data.dptr, + (unsigned char*)output->columnwise_scale_inv.dptr, M, K, + static_cast(output->scale_inv.shape[1]), + static_cast(output->columnwise_scale_inv.shape[1]), iters); + }; + if (M % kRowsPerMxBlock == 0 && K % kTileCols == 0) { + run(std::false_type{}); + } else { + run(std::true_type{}); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace gated_bidim +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_GATED_BIDIM_CUH_ 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..550e1c7f1f5 --- /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 +constexpr bool kHasTable = false; +template +constexpr Kind kKindOf = kGelu; + +template <> +constexpr bool kHasTable> = true; +template <> +constexpr bool kHasTable> = true; +template <> +constexpr bool kHasTable> = true; +template <> +constexpr bool kHasTable> = true; +template <> +constexpr Kind kKindOf> = kSilu; +template <> +constexpr Kind kKindOf> = kDgelu; +template <> +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_ diff --git a/transformer_engine/common/util/cutedsl_utils.h b/transformer_engine/common/util/cutedsl_utils.h index 4fd4ee56f69..4af0d78e906 100644 --- a/transformer_engine/common/util/cutedsl_utils.h +++ b/transformer_engine/common/util/cutedsl_utils.h @@ -90,6 +90,76 @@ NVTE_SPECIALIZE_CUTEDSL_ACTIVATION_FUNC_TO_ENUM(dsrelu, kDSReLU) #undef NVTE_SPECIALIZE_CUTEDSL_ACTIVATION_FUNC_TO_ENUM +// Encodes the gated activation selected by the gated dispatcher (gated_mxfp8.cuh), identified by +// its forward ActOP and, for backward, its DActOP. +enum class GatedActivation { + kGeGLU, + kQGeGLU, + kReGLU, + kSReGLU, + kSwiGLU, + kGLU, + kSiTUGLU, + kClampedSwiGLU, + kUnsupported, + kNumTypes +}; + +inline const char *gated_activation_to_str(GatedActivation act) { + switch (act) { + case GatedActivation::kGeGLU: + return "geglu"; + case GatedActivation::kQGeGLU: + return "qgeglu"; + case GatedActivation::kReGLU: + return "reglu"; + case GatedActivation::kSReGLU: + return "sreglu"; + case GatedActivation::kSwiGLU: + return "swiglu"; + case GatedActivation::kGLU: + return "glu"; + case GatedActivation::kSiTUGLU: + return "situglu"; + case GatedActivation::kClampedSwiGLU: + return "clamped_swiglu"; + case GatedActivation::kUnsupported: + case GatedActivation::kNumTypes: + return "unsupported"; + } + return "unsupported"; +} + +// DActOP is nullptr for the forward. +template +constexpr GatedActivation gated_activation_func_to_enum() { + return GatedActivation::kUnsupported; +} + +#define NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(PARAM, ACT, DACT, ACTIVATION) \ + template <> \ + constexpr GatedActivation gated_activation_func_to_enum, nullptr>() { \ + return GatedActivation::ACTIVATION; \ + } \ + template <> \ + constexpr GatedActivation \ + gated_activation_func_to_enum, DACT>() { \ + return GatedActivation::ACTIVATION; \ + } + +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, gelu, dgelu, kGeGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, qgelu, dqgelu, kQGeGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, relu, drelu, kReGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, srelu, dsrelu, kSReGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, silu, dsilu, kSwiGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, sigmoid, dsigmoid, kGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(SiTUGLUParam, situ_gate, dsitu_gate, kSiTUGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(ClampedSwiGLUParam, clamped_silu, + clamped_dsilu, kClampedSwiGLU) + +#undef NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM + } // namespace cutedsl_backend } // namespace transformer_engine diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 49686437c63..bf4d89b7e60 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -732,6 +732,47 @@ __device__ __forceinline__ void abs_max_2x(fp16x2 &dst, const fp16x2 &p1, const #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) } +// Packed FP32x2 arithmetic. Each lane rounds exactly like the scalar .rn operation. +__device__ __forceinline__ floatx2 add_2x(const floatx2 &a, const floatx2 &b) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + floatx2 d; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(reinterpret_cast(d)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); + return d; +#else + NVTE_DEVICE_ERROR("add_2x is only supported on SM 10.0+."); + return {}; +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ floatx2 mul_2x(const floatx2 &a, const floatx2 &b) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + floatx2 d; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(reinterpret_cast(d)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); + return d; +#else + NVTE_DEVICE_ERROR("mul_2x is only supported on SM 10.0+."); + return {}; +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ floatx2 fma_2x(const floatx2 &a, const floatx2 &b, const floatx2 &c) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + floatx2 d; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(reinterpret_cast(d)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); + return d; +#else + NVTE_DEVICE_ERROR("fma_2x is only supported on SM 10.0+."); + return {}; +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + __device__ __forceinline__ int32_t elect_one_sync(uint32_t mask = 0xFFFFFFFFu) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) int32_t pred = 0; @@ -836,6 +877,17 @@ __device__ __forceinline__ bf16x2 exp2f_rcp_2x(e8m0_t biased_exp) { return result; } +/*! \brief exp2f_rcp for two biased exponents at once, one in the low byte of each + * 16-bit half of \p biased_exps, returned as a BF16 pair. + * + * Exponent 254 gives the subnormal 2^-127, as exp2f_rcp_2x does. Exponent 255 is not handled. + */ +__device__ __forceinline__ bf16x2 exp2f_rcp_2x_per_lane(const uint32_t biased_exps) { + const uint32_t bits = + ((0x00FE00FEu - biased_exps) << 7) | (__vcmpeq2(biased_exps, 0x00FE00FEu) & 0x00400040u); + return reinterpret_cast(bits); +} + // Scale two BF16 pairs by independent scales and pack the four results into a // single FP8E4M3 word. The mul_cvt_4x overload below shares one scale across // all four elements, which is what a rowwise MX block wants; a colwise block diff --git a/transformer_engine/common/util/ptx_arch_spec.cuh b/transformer_engine/common/util/ptx_arch_spec.cuh index 14f390740e1..0b5599adb5a 100644 --- a/transformer_engine/common/util/ptx_arch_spec.cuh +++ b/transformer_engine/common/util/ptx_arch_spec.cuh @@ -103,6 +103,20 @@ __device__ __forceinline__ e8m0_t float_to_e8m0(float val) { } } +/*! \brief float_to_e8m0 for two values at once, returned as two packed bytes: + * \p hi in the upper byte, \p lo in the lower one. + */ +__device__ __forceinline__ uint32_t float_to_e8m0_2x(const float hi, const float lo) { + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + uint16_t out; + asm("cvt.rp.satfinite.ue8m0x2.f32 %0, %1, %2;" : "=h"(out) : "f"(hi), "f"(lo)); + return out; + } else { + return (static_cast(float_to_e8m0(hi)) << 8) | float_to_e8m0(lo); + } +} + __device__ __forceinline__ void reduce_sync_max_abs_f32(float &out, float const &in) { constexpr bool is_sm_100f = NVTE_CUDA_ARCH_MATCHES(ptx::FamilySpecific<100>); if constexpr (is_sm_100f) {