diff --git a/.clang-format b/.clang-format index aec13e3762..7860641ede 100644 --- a/.clang-format +++ b/.clang-format @@ -261,7 +261,7 @@ SpacesInParensOptions: InEmptyParentheses: false Other: false SpacesInSquareBrackets: false -Standard: Auto +Standard: c++17 StatementAttributeLikeMacros: - Q_EMIT StatementMacros: @@ -277,4 +277,3 @@ WhitespaceSensitiveMacros: - PP_STRINGIZE - STRINGIZE ... - diff --git a/docs/envvars.rst b/docs/envvars.rst index 044a7f6a0d..912a1b40fb 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -267,7 +267,7 @@ Kernel Configuration :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. This can be useful in environments where NVRTC is not available or not desired. + :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. Existing transpose operations select their static fallback automatically. NVRTC-migrated fused softmax and normalization paths require their corresponding ``NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX`` or ``NVTE_BUILD_LEGACY_STATIC_NORM`` CMake option to have been enabled when the library was built; otherwise they report that no static fallback is available. .. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 1c4d86a3a8..98639386ad 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -32,6 +32,7 @@ add_executable(test_operator test_multi_cast_transpose.cu test_multi_padding.cu test_multi_unpadding.cu + test_softmax.cu test_causal_softmax.cu test_swizzle.cu test_multi_swizzle.cu diff --git a/tests/cpp/operator/test_softmax.cu b/tests/cpp/operator/test_softmax.cu new file mode 100644 index 0000000000..4fc5a294c8 --- /dev/null +++ b/tests/cpp/operator/test_softmax.cu @@ -0,0 +1,234 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include "../test_common.h" + +using namespace transformer_engine; + +namespace { + +template +void ref_softmax_row(Type *out, const Type *in, const uint8_t *mask, int cols, float scale) { + float max_value = -10000.0f; + bool has_unmasked = false; + for (int j = 0; j < cols; ++j) { + if (mask != nullptr && mask[j] == 1) continue; + max_value = std::max(max_value, static_cast(in[j]) * scale); + has_unmasked = true; + } + float sum = 0.0f; + for (int j = 0; j < cols; ++j) { + if (mask != nullptr && mask[j] == 1) { + out[j] = static_cast(0.0f); + continue; + } + const float val = has_unmasked ? std::exp(static_cast(in[j]) * scale - max_value) : 0.0f; + sum += val; + out[j] = static_cast(val); + } + for (int j = 0; j < cols; ++j) { + out[j] = static_cast(static_cast(out[j]) / sum); + } +} + +template +void ref_softmax_bwd(Type *grad_in, const Type *grad, const Type *softmax, int cols, float scale) { + float sum = 0.0f; + for (int j = 0; j < cols; ++j) { + sum += static_cast(grad[j]) * static_cast(softmax[j]); + } + for (int j = 0; j < cols; ++j) { + grad_in[j] = + static_cast(scale * (static_cast(grad[j]) - sum) * + static_cast(softmax[j])); + } +} + +template +void ref_upper_row(Type *out, const Type *in, int row, int cols, float scale) { + float max_value = -10000.0f; + for (int j = 0; j <= row; ++j) { + max_value = std::max(max_value, static_cast(in[j]) * scale); + } + float sum = 0.0f; + for (int j = 0; j < cols; ++j) { + if (j <= row) { + const float val = std::exp(static_cast(in[j]) * scale - max_value); + sum += val; + out[j] = static_cast(val); + } else { + out[j] = static_cast(0.0f); + } + } + for (int j = 0; j <= row; ++j) { + out[j] = static_cast(static_cast(out[j]) / sum); + } +} + +template +void test_scaled_softmax(DType dtype) { + using namespace test; + constexpr int batches = 2; + constexpr int heads = 2; + constexpr int rows = 8; + constexpr int cols = 32; + constexpr float scale = 0.7f; + constexpr size_t elements_total = batches * heads * rows * cols; + Tensor input("input", std::vector{batches, heads, rows, cols}, dtype); + Tensor softmax("softmax", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad("grad", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad_out("grad_out", std::vector{batches, heads, rows, cols}, dtype); + fillUniform(&input); + fillUniform(&grad); + nvte_scaled_softmax_forward(input.data(), softmax.data(), scale, 0); + nvte_scaled_softmax_backward(grad.data(), softmax.data(), grad_out.data(), scale, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::unique_ptr ref = std::make_unique(elements_total); + std::unique_ptr ref_grad = std::make_unique(elements_total); + const Type *input_cpu = input.rowwise_cpu_dptr(); + const Type *grad_cpu = grad.rowwise_cpu_dptr(); + for (size_t row = 0; row < elements_total / cols; ++row) { + ref_softmax_row(ref.get() + row * cols, input_cpu + row * cols, nullptr, cols, scale); + ref_softmax_bwd(ref_grad.get() + row * cols, grad_cpu + row * cols, ref.get() + row * cols, + cols, scale); + } + auto [atol, rtol] = getTolerances(dtype); + if (dtype == DType::kBFloat16) atol = 1e-3; + compareResults("scaled_softmax_fwd", softmax, ref.get(), true, atol, rtol); + compareResults("scaled_softmax_bwd", grad_out, ref_grad.get(), true, atol, rtol); +} + +template +void test_masked_softmax(DType dtype) { + using namespace test; + constexpr int batches = 2; + constexpr int heads = 2; + constexpr int rows = 8; + constexpr int cols = 32; + constexpr float scale = -0.3f; + constexpr size_t elements_total = batches * heads * rows * cols; + Tensor input("input", std::vector{batches, heads, rows, cols}, dtype); + Tensor mask("mask", std::vector{1, 1, rows, cols}, DType::kByte); + Tensor softmax("softmax", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad("grad", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad_out("grad_out", std::vector{batches, heads, rows, cols}, dtype); + fillUniform(&input); + fillUniform(&grad); + uint8_t *mask_cpu = mask.rowwise_cpu_dptr(); + for (size_t i = 0; i < rows * cols; ++i) { + mask_cpu[i] = (i % 7 == 0) ? 1 : 0; + } + mask.from_cpu(); + nvte_scaled_masked_softmax_forward(input.data(), mask.data(), softmax.data(), scale, 0); + nvte_scaled_masked_softmax_backward(grad.data(), softmax.data(), grad_out.data(), scale, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::unique_ptr ref = std::make_unique(elements_total); + std::unique_ptr ref_grad = std::make_unique(elements_total); + const Type *input_cpu = input.rowwise_cpu_dptr(); + const Type *grad_cpu = grad.rowwise_cpu_dptr(); + for (int row = 0; row < batches * heads * rows; ++row) { + const int mask_row = row % rows; + ref_softmax_row(ref.get() + row * cols, input_cpu + row * cols, mask_cpu + mask_row * cols, + cols, scale); + ref_softmax_bwd(ref_grad.get() + row * cols, grad_cpu + row * cols, ref.get() + row * cols, + cols, scale); + } + auto [atol, rtol] = getTolerances(dtype); + if (dtype == DType::kBFloat16) atol = 1e-3; + compareResults("masked_softmax_fwd", softmax, ref.get(), true, atol, rtol); + compareResults("masked_softmax_bwd", grad_out, ref_grad.get(), true, atol, rtol); +} + +template +void test_upper_softmax(DType dtype) { + using namespace test; + constexpr int attn_batches = 4; + constexpr int seq = 32; + constexpr float scale = 1.2f; + constexpr size_t elements_total = attn_batches * seq * seq; + Tensor input("input", std::vector{attn_batches, seq, seq}, dtype); + Tensor softmax("softmax", std::vector{attn_batches, seq, seq}, dtype); + Tensor grad("grad", std::vector{attn_batches, seq, seq}, dtype); + Tensor grad_out("grad_out", std::vector{attn_batches, seq, seq}, dtype); + fillUniform(&input); + fillUniform(&grad); + nvte_scaled_upper_triang_masked_softmax_forward(input.data(), softmax.data(), scale, 0); + nvte_scaled_upper_triang_masked_softmax_backward(grad.data(), softmax.data(), grad_out.data(), + scale, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::unique_ptr ref = std::make_unique(elements_total); + std::unique_ptr ref_grad = std::make_unique(elements_total); + const Type *input_cpu = input.rowwise_cpu_dptr(); + const Type *grad_cpu = grad.rowwise_cpu_dptr(); + for (int batch = 0; batch < attn_batches; ++batch) { + for (int row = 0; row < seq; ++row) { + const size_t offset = (batch * seq + row) * seq; + ref_upper_row(ref.get() + offset, input_cpu + offset, row, seq, scale); + ref_softmax_bwd(ref_grad.get() + offset, grad_cpu + offset, ref.get() + offset, seq, scale); + for (int col = row + 1; col < seq; ++col) { + ref_grad[offset + col] = static_cast(0.0f); + } + } + } + auto [atol, rtol] = getTolerances(dtype); + if (dtype == DType::kBFloat16) atol = 1e-3; + compareResults("upper_softmax_fwd", softmax, ref.get(), true, atol, rtol); + compareResults("upper_softmax_bwd", grad_out, ref_grad.get(), true, atol, rtol); +} + +} // namespace + +// Dispatch a 16-bit float dtype to a templated test body. Mirrors +// TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT but uses the test harness's own fp16/bf16 +// aliases so we don't have to include common.h here -- doing so would make the +// test's Tensor type ambiguous with transformer_engine::Tensor. +#define SOFTMAX_TEST_DISPATCH_16BIT(dtype, fn) \ + switch (dtype) { \ + case DType::kFloat16: \ + fn(dtype); \ + break; \ + case DType::kBFloat16: \ + fn(dtype); \ + break; \ + default: \ + GTEST_FAIL() << "Unsupported 16-bit dtype for test"; \ + } + +class SoftmaxApiTestSuite : public ::testing::TestWithParam {}; + +TEST_P(SoftmaxApiTestSuite, ScaledSoftmax) { + const DType dtype = GetParam(); + SOFTMAX_TEST_DISPATCH_16BIT(dtype, test_scaled_softmax); +} + +TEST_P(SoftmaxApiTestSuite, MaskedSoftmax) { + const DType dtype = GetParam(); + SOFTMAX_TEST_DISPATCH_16BIT(dtype, test_masked_softmax); +} + +TEST_P(SoftmaxApiTestSuite, UpperTriangularSoftmax) { + const DType dtype = GetParam(); + SOFTMAX_TEST_DISPATCH_16BIT(dtype, test_upper_softmax); +} + +INSTANTIATE_TEST_SUITE_P(OperatorTest, SoftmaxApiTestSuite, + ::testing::Values(DType::kFloat16, DType::kBFloat16), + [](const testing::TestParamInfo &info) { + return test::typeName(info.param); + }); diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index 6d70b7b84f..1dfd2fed4e 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -9,7 +9,15 @@ add_executable(test_util find_package(OpenMP REQUIRED) -target_link_libraries(test_util PUBLIC CUDA::cudart GTest::gtest_main ${TE_LIB} CUDA::nvrtc CUDNN::cudnn OpenMP::OpenMP_CXX) +find_package(Threads REQUIRED) +target_link_libraries(test_util PUBLIC + CUDA::cudart + GTest::gtest_main + ${TE_LIB} + CUDA::nvrtc + CUDNN::cudnn + OpenMP::OpenMP_CXX + Threads::Threads) target_compile_options(test_util PRIVATE -O2 -fopenmp) include(GoogleTest) diff --git a/tests/cpp/util/test_nvrtc.cpp b/tests/cpp/util/test_nvrtc.cpp index d41084449e..5228c60972 100644 --- a/tests/cpp/util/test_nvrtc.cpp +++ b/tests/cpp/util/test_nvrtc.cpp @@ -4,11 +4,14 @@ * See LICENSE for license information. ************************************************************************/ +#include + +#include +#include #include +#include #include -#include - #include "util/rtc.h" using namespace transformer_engine; @@ -19,10 +22,10 @@ TEST(UtilTest, NVRTC) { } // GPU data buffer - int *device_buffer; + int* device_buffer; std::vector host_buffer(2); - cudaMalloc((void**)&device_buffer, 2*sizeof(int)); // NOLINT(*) - cudaMemset(device_buffer, 0, 2*sizeof(int)); + cudaMalloc((void**)&device_buffer, 2 * sizeof(int)); // NOLINT(*) + cudaMemset(device_buffer, 0, 2 * sizeof(int)); // CUDA kernel implementations const char code1[] = R"code( @@ -38,45 +41,151 @@ __global__ void my_kernel(uint32_t *data) { data[0] = 789; data[1] = 12; } +)code"; + const char code3[] = R"code( +#ifndef NVTE_GTEST_RTC_VALUE +#error "NVTE_GTEST_RTC_VALUE must be provided" +#endif +__global__ void my_kernel(int *data) { + data[0] = NVTE_GTEST_RTC_VALUE; + data[1] = 34; +} +)code"; + const char header4[] = R"code( +#define NVTE_GTEST_RTC_HEADER_VALUE 78 +)code"; + const char code4[] = R"code( +#include "test_nvrtc_header.h" +__global__ void my_kernel(int *data) { + data[0] = NVTE_GTEST_RTC_HEADER_VALUE; + data[1] = 90; +} )code"; // Make sure kernels are not available auto& nvrtc_manager = rtc::KernelManager::instance(); EXPECT_FALSE(nvrtc_manager.is_compiled("my gtest kernel1")); EXPECT_FALSE(nvrtc_manager.is_compiled("my gtest kernel2")); - EXPECT_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, - device_buffer), + EXPECT_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, device_buffer), std::runtime_error); - EXPECT_THROW(nvrtc_manager.launch("my gtest kernel2", 1, 1, 0, 0, - device_buffer), + EXPECT_THROW(nvrtc_manager.launch("my gtest kernel2", 1, 1, 0, 0, device_buffer), std::runtime_error); // Compile and run first kernel - EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel1", - "my_kernel", - code1, - "test_nvrtc_kernel1.cu")); + EXPECT_NO_THROW( + nvrtc_manager.compile("my gtest kernel1", "my_kernel", code1, "test_nvrtc_kernel1.cu")); EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel1")); EXPECT_FALSE(nvrtc_manager.is_compiled("my gtest kernel2")); - EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, - device_buffer)); - EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2*sizeof(int), - cudaMemcpyDeviceToHost), + EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, device_buffer)); + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), cudaSuccess); EXPECT_EQ(host_buffer[0], 123); EXPECT_EQ(host_buffer[1], -456); // Compile and run second kernel - EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel2", - "my_kernel", - code2, - "test_nvrtc_kernel2.cu")); + EXPECT_NO_THROW( + nvrtc_manager.compile("my gtest kernel2", "my_kernel", code2, "test_nvrtc_kernel2.cu")); EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel1")); EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel2")); EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel2", 1, 1, 0, 0, device_buffer)); - EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2*sizeof(int), - cudaMemcpyDeviceToHost), + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), cudaSuccess); EXPECT_EQ(host_buffer[0], 789); EXPECT_EQ(host_buffer[1], 12); + + // Compile and run kernel with extra compile options + EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel3", "my_kernel", code3, + "test_nvrtc_kernel3.cu", {"-DNVTE_GTEST_RTC_VALUE=56"})); + EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel3")); + EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel3", 1, 1, 0, 0, device_buffer)); + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(host_buffer[0], 56); + EXPECT_EQ(host_buffer[1], 34); + + // Compile and run kernel with an extra in-memory header + EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel4", "my_kernel", code4, + "test_nvrtc_kernel4.cu", {}, + {{header4, "test_nvrtc_header.h"}})); + EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel4")); + EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel4", 1, 1, 0, 0, device_buffer)); + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(host_buffer[0], 78); + EXPECT_EQ(host_buffer[1], 90); + + EXPECT_EQ(cudaFree(device_buffer), cudaSuccess); +} + +TEST(UtilTest, NVRTCConcurrentCompile) { + if (!rtc::is_enabled()) { + GTEST_SKIP() << "NVRTC not enabled, skipping tests"; + } + + constexpr int num_threads = 8; + constexpr char kernel_label[] = "my concurrent gtest kernel"; + const char code[] = R"code( +__global__ void my_kernel(int *data) { + data[0] = 314; + data[1] = 159; +} +)code"; + + int device_id = 0; + ASSERT_EQ(cudaGetDevice(&device_id), cudaSuccess); + + auto& nvrtc_manager = rtc::KernelManager::instance(); + ASSERT_FALSE(nvrtc_manager.is_compiled(kernel_label)); + + std::atomic ready{0}; + std::atomic start{false}; + std::vector errors(num_threads); + std::vector threads; + threads.reserve(num_threads); + for (int thread_id = 0; thread_id < num_threads; ++thread_id) { + threads.emplace_back([&, thread_id] { + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + try { + const cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + throw std::runtime_error(cudaGetErrorString(status)); + } + (void)nvrtc_manager.is_compiled(kernel_label); + nvrtc_manager.compile(kernel_label, "my_kernel", code, "test_nvrtc_concurrent_kernel.cu"); + } catch (...) { + errors[thread_id] = std::current_exception(); + } + }); + } + + while (ready.load(std::memory_order_acquire) != num_threads) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (auto& thread : threads) { + thread.join(); + } + for (const auto& error : errors) { + if (error != nullptr) { + try { + std::rethrow_exception(error); + } catch (const std::exception& e) { + ADD_FAILURE() << e.what(); + } + } + } + + ASSERT_TRUE(nvrtc_manager.is_compiled(kernel_label)); + int* device_buffer = nullptr; + ASSERT_EQ(cudaMalloc(reinterpret_cast(&device_buffer), 2 * sizeof(int)), cudaSuccess); + ASSERT_NO_THROW(nvrtc_manager.launch(kernel_label, 1, 1, 0, 0, device_buffer)); + std::vector host_buffer(2); + ASSERT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(host_buffer[0], 314); + EXPECT_EQ(host_buffer[1], 159); + EXPECT_EQ(cudaFree(device_buffer), cudaSuccess); } diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index edb8c5e109..899eff7557 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -186,6 +186,7 @@ list(APPEND transformer_engine_cpp_sources fused_attn/fused_attn.cpp gemm/config.cpp normalization/common.cpp + normalization/rtc_dispatch.cpp normalization/layernorm/ln_api.cpp normalization/rmsnorm/rmsnorm_api.cpp util/cuda_driver.cpp @@ -554,13 +555,59 @@ make_string_header_from_file(transpose/rtc/transpose.cu string_code_transpose_rtc_transpose_cu) make_string_header_from_file(transpose/rtc/swap_first_dims.cu string_code_transpose_rtc_swap_first_dims_cu) +make_string_header_from_file(fused_softmax/scaled_masked_softmax.cu + string_code_fused_softmax_scaled_masked_softmax_cu) +make_string_header_from_file(fused_softmax/scaled_upper_triang_masked_softmax.cu + string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu) +make_string_header_from_file(fused_softmax/scaled_aligned_causal_masked_softmax.cu + string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu) make_string_header_from_file(utils.cuh string_code_utils_cuh) make_string_header_from_file(util/math.h string_code_util_math_h) + +# Norm NVRTC bundled headers + RTC source files +make_string_header_from_file(normalization/kernel_params.h + string_code_normalization_kernel_params_h) +make_string_header_from_file(normalization/kernel_traits.h + string_code_normalization_kernel_traits_h) +make_string_header_from_file(normalization/layernorm/ln_fwd_kernels.cuh + string_code_normalization_layernorm_ln_fwd_kernels_cuh) +make_string_header_from_file(normalization/layernorm/ln_bwd_kernels.cuh + string_code_normalization_layernorm_ln_bwd_kernels_cuh) +make_string_header_from_file(normalization/rmsnorm/rmsnorm_fwd_kernels.cuh + string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh) +make_string_header_from_file(normalization/rmsnorm/rmsnorm_bwd_kernels.cuh + string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh) +make_string_header_from_file(normalization/layernorm/rtc/ln_fwd_kernel.cu + string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu) +make_string_header_from_file(normalization/layernorm/rtc/ln_bwd_kernel.cu + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu) +make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu + string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu) +make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu) target_include_directories(transformer_engine PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/string_headers") +option(NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + "Also compile legacy static fused softmax kernels for NVTE_DISABLE_NVRTC fallback" + OFF) +target_compile_definitions(transformer_engine + PRIVATE + NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=$) + +# Default OFF: LayerNorm/RMSNorm require NVRTC. Set ON to additionally compile +# the legacy static template instantiations, which are selected only when +# NVTE_DISABLE_NVRTC=1 at runtime. This compiles all 556 +# REGISTER_NORM_LAUNCHER variants up front. +option(NVTE_BUILD_LEGACY_STATIC_NORM + "Also compile static norm kernels for NVTE_DISABLE_NVRTC fallback" + OFF) +target_compile_definitions(transformer_engine + PRIVATE + NVTE_BUILD_LEGACY_STATIC_NORM=$) + # Compiler options set(nvte_sources_with_fast_math) list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu diff --git a/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu index 6ea8017e07..1f8e748479 100644 --- a/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu @@ -4,9 +4,11 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include #include -#include #include #include #include @@ -19,10 +21,23 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/rtc.h" +#include "../util/string.h" #include "../utils.cuh" +#include "string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu.h" +#endif namespace transformer_engine { +#ifdef __CUDACC_RTC__ +using fp16 = half; +using bf16 = nv_bfloat16; +#endif + +#ifndef NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX +#define NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX 0 +#endif + template __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src); @@ -92,7 +107,7 @@ struct Max { template __device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) { -#if CUDA_VERSION >= 9000 +#if defined(__CUDACC_RTC__) || CUDA_VERSION >= 9000 return __shfl_xor_sync(mask, value, laneMask, width); #else return __shfl_xor(value, laneMask, width); @@ -345,6 +360,42 @@ __global__ void scaled_aligned_causal_masked_softmax_warp_backward( } } +#ifdef __CUDACC_RTC__ + +#else + +namespace { + +constexpr const char *kRtcSourceFile = + "transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu"; + +template +std::string make_softmax_rtc_code(int log2_elements) { + (void)log2_elements; + std::string code = string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu; + return code; +} + +template +std::string make_softmax_rtc_label(const char *direction, int log2_elements) { + return concat_strings("fused_softmax,variant=aligned_causal,direction=", direction, + ",type=", TypeInfo::name, ",log2=", log2_elements, ",fast_math=1"); +} + +template +std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); +} + +void throw_nvrtc_required(const char *direction) { + NVTE_ERROR("Fused aligned causal softmax RTC path is disabled for ", direction, + ". Set NVTE_DISABLE_NVRTC=0 or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=ON."); +} + +} // namespace + template void call_kernel_scaled_aligned_causal_masked_softmax_forward( dim3 grid_size, dim3 block_size, const int shmem_size, cudaStream_t stream, output_t *dst, @@ -454,6 +505,24 @@ void dispatch_scaled_aligned_causal_masked_softmax_forward(output_t *dst, const dim3 block_size(warp_width, warps_per_block); dim3 grid_size(blocks); + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_aligned_causal_masked_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, grid_size, block_size, 0, stream, dst, src, rtc_scale, + microbatches, query_seq_len, key_seq_len); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX // create an array of pointers to functions using ForwardFuncType = typename FunctionWrapper::ForwardType; static std::array forwardFunctionArray; @@ -466,6 +535,9 @@ void dispatch_scaled_aligned_causal_masked_softmax_forward(output_t *dst, const // Call the corresponding kernel forwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, dst, src, scale, microbatches, query_seq_len, key_seq_len); +#else + throw_nvrtc_required("forward"); +#endif } template @@ -499,6 +571,23 @@ void dispatch_scaled_aligned_causal_masked_softmax_backward( dim3 block_size(warp_width, warps_per_block); dim3 grid_size(blocks); + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("backward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_aligned_causal_masked_softmax_warp_backward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + rtc_manager.launch(kernel_label, grid_size, block_size, 0, stream, grad_input, grad, output, + scale, microbatches, query_seq_len, key_seq_len); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX // create an array of pointers to functions using BackwardFuncType = typename FunctionWrapper::BackwardType; static std::array backwardFunctionArray; @@ -511,6 +600,9 @@ void dispatch_scaled_aligned_causal_masked_softmax_backward( // Call the corresponding kernel backwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, grad_input, grad, output, scale, microbatches, query_seq_len, key_seq_len); +#else + throw_nvrtc_required("backward"); +#endif } void scaled_aligned_causal_masked_softmax_forward(const Tensor &input, Tensor *softmax_results, @@ -546,8 +638,13 @@ void scaled_aligned_causal_masked_softmax_backward(Tensor output_grads, const Te reinterpret_cast(softmax_results.data.dptr), scale_factor, query_seq_len, key_seq_len, batches, attn_heads, stream);); } + +#endif // __CUDACC_RTC__ + } // end namespace transformer_engine +#ifndef __CUDACC_RTC__ + void nvte_scaled_aligned_causal_masked_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream) { @@ -568,3 +665,5 @@ void nvte_scaled_aligned_causal_masked_softmax_backward(const NVTETensor incomin *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads), *convertNVTETensorCheck(softmax_results), scale_factor, stream); } + +#endif // __CUDACC_RTC__ diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu index 27f86673c5..6f295ad00a 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -4,9 +4,11 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include #include -#include #include #include #include @@ -17,10 +19,27 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/rtc.h" +#include "../util/string.h" #include "../utils.cuh" +#include "string_code_fused_softmax_scaled_masked_softmax_cu.h" +#endif namespace transformer_engine { +#ifdef __CUDACC_RTC__ +using bf16 = nv_bfloat16; +#endif + +#ifndef NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX +#define NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX 0 +#endif + +template +__device__ __forceinline__ T neg_infinity() { + return -static_cast(__int_as_float(0x7f800000)); +} + template __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src); @@ -67,7 +86,7 @@ struct Max { template __device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) { -#if CUDA_VERSION >= 9000 +#if defined(__CUDACC_RTC__) || CUDA_VERSION >= 9000 return __shfl_xor_sync(mask, value, laneMask, width); #else return __shfl_xor(value, laneMask, width); @@ -144,7 +163,7 @@ __global__ void scaled_softmax_warp_forward(output_t *dst, const input_t *src, c } else { #pragma unroll for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } @@ -167,7 +186,7 @@ __global__ void scaled_softmax_warp_forward(output_t *dst, const input_t *src, c for (int i = 0; i < WARP_BATCH; ++i) { #pragma unroll for (int it = 0; it < WARP_ITERATIONS; ++it) { - elements[i][it] = std::exp((elements[i][it] - max_value[i])); + elements[i][it] = expf((elements[i][it] - max_value[i])); sum[i] += elements[i][it]; } } @@ -269,7 +288,7 @@ __global__ void scaled_masked_softmax_warp_forward(output_t *dst, const input_t } else { #pragma unroll for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } @@ -299,7 +318,7 @@ __global__ void scaled_masked_softmax_warp_forward(output_t *dst, const input_t for (int i = 0; i < WARP_BATCH; ++i) { #pragma unroll for (int it = 0; it < WARP_ITERATIONS; ++it) { - elements[i][it] = std::exp((elements[i][it] - max_value[i])); + elements[i][it] = expf((elements[i][it] - max_value[i])); sum[i] += elements[i][it]; } } @@ -420,6 +439,42 @@ __global__ void scaled_masked_softmax_warp_backward(output_t *gradInput, const i } } +#ifdef __CUDACC_RTC__ + +#else + +namespace { + +constexpr const char *kRtcSourceFile = + "transformer_engine/common/fused_softmax/scaled_masked_softmax.cu"; + +template +std::string make_softmax_rtc_code(int log2_elements) { + (void)log2_elements; + std::string code = string_code_fused_softmax_scaled_masked_softmax_cu; + return code; +} + +template +std::string make_softmax_rtc_label(const char *variant, const char *direction, int log2_elements) { + return concat_strings("fused_softmax,variant=", variant, ",direction=", direction, + ",type=", TypeInfo::name, ",log2=", log2_elements, ",fast_math=1"); +} + +template +std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); +} + +void throw_nvrtc_required(const char *variant) { + NVTE_ERROR("Fused softmax RTC path is disabled for ", variant, + ". Set NVTE_DISABLE_NVRTC=0 or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=ON."); +} + +} // namespace + template void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const input_t scale, int query_seq_len, int key_seq_len, int batches, @@ -448,70 +503,89 @@ void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const in NVTE_CHECK(query_seq_len % batches_per_block == 0, "Unsupported shape."); dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 1: // 2 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 2: // 4 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 3: // 8 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 4: // 16 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 5: // 32 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 6: // 64 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 7: // 128 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 8: // 256 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 9: // 512 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 10: // 1024 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 11: // 2048 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 12: // 4096 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 13: // 8192 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 14: // 16384 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = + make_softmax_rtc_label("scaled", "forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile( + kernel_label, + make_softmax_rtc_kernel_name("scaled_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, rtc_scale, batch_count, + key_seq_len); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 1: // 2 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 2: // 4 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 3: // 8 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 4: // 16 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 5: // 32 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 6: // 64 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 7: // 128 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 8: // 256 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 9: // 512 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 10: // 1024 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 11: // 2048 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 12: // 4096 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 13: // 8192 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 14: // 16384 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + default: + break; + } +#else + throw_nvrtc_required("scaled softmax forward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -546,85 +620,105 @@ void dispatch_scaled_masked_softmax_forward(output_t *dst, const input_t *src, c NVTE_CHECK(query_seq_len % batches_per_block == 0, "Unsupported shape."); dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 1: // 2 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 2: // 4 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 3: // 8 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 4: // 16 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 5: // 32 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 6: // 64 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 7: // 128 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 8: // 256 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 9: // 512 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 10: // 1024 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 11: // 2048 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 12: // 4096 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 13: // 8192 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 14: // 16384 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = + make_softmax_rtc_label("masked", "forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_masked_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, mask, rtc_scale, + batch_count, key_seq_len, pad_batches); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 1: // 2 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 2: // 4 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 3: // 8 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 4: // 16 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 5: // 32 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 6: // 64 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 7: // 128 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 8: // 256 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 9: // 512 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 10: // 1024 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 11: // 2048 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 12: // 4096 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 13: // 8192 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 14: // 16384 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + default: + break; + } +#else + throw_nvrtc_required("scaled masked softmax forward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -658,85 +752,104 @@ void dispatch_scaled_masked_softmax_backward(output_t *grad_input, const input_t int batches_per_block = warps_per_block * batches_per_warp; int blocks = batch_count / batches_per_block; dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 1: // 2 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 2: // 4 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 3: // 8 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 4: // 16 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 5: // 32 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 6: // 64 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 7: // 128 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 8: // 256 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 9: // 512 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 10: // 1024 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 11: // 2048 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 12: // 4096 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 13: // 8192 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 14: // 16384 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = + make_softmax_rtc_label("masked", "backward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_masked_softmax_warp_backward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, grad_input, grad, output, scale, + batch_count, key_seq_len); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 1: // 2 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 2: // 4 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 3: // 8 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 4: // 16 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 5: // 32 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 6: // 64 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 7: // 128 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 8: // 256 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 9: // 512 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 10: // 1024 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 11: // 2048 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 12: // 4096 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 13: // 8192 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 14: // 16384 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + default: + break; + } +#else + throw_nvrtc_required("scaled softmax backward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -812,8 +925,12 @@ void scaled_masked_softmax_backward(Tensor output_grads, const Tensor incoming_g query_seq_len, key_seq_len, batches, attn_heads, stream);); } +#endif // __CUDACC_RTC__ + } // end namespace transformer_engine +#ifndef __CUDACC_RTC__ + void nvte_scaled_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream) { NVTE_API_CALL(nvte_scaled_softmax_forward); @@ -850,3 +967,5 @@ void nvte_scaled_masked_softmax_backward(const NVTETensor incoming_grads, *convertNVTETensorCheck(incoming_grads), *convertNVTETensorCheck(softmax_results), scale_factor, stream); } + +#endif // __CUDACC_RTC__ diff --git a/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu index 431148cd1d..c18444f646 100644 --- a/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu @@ -4,9 +4,11 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include #include -#include #include #include #include @@ -17,10 +19,28 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/rtc.h" +#include "../util/string.h" #include "../utils.cuh" +#include "string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu.h" +#endif namespace transformer_engine { +#ifdef __CUDACC_RTC__ +using fp16 = half; +using bf16 = nv_bfloat16; +#endif + +#ifndef NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX +#define NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX 0 +#endif + +template +__device__ __forceinline__ T neg_infinity() { + return -static_cast(__int_as_float(0x7f800000)); +} + template __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src); @@ -90,7 +110,7 @@ struct Max { template __device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) { -#if CUDA_VERSION >= 9000 +#if defined(__CUDACC_RTC__) || CUDA_VERSION >= 9000 return __shfl_xor_sync(mask, value, laneMask, width); #else return __shfl_xor(value, laneMask, width); @@ -166,13 +186,13 @@ __global__ void scaled_upper_triang_masked_softmax_warp_forward(output_t *dst, c if ((element_index + element) < batch_element_count) { elements[i][it + element] = (acc_t)temp_data[element] * scale; } else { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } else { #pragma unroll for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } @@ -196,7 +216,7 @@ __global__ void scaled_upper_triang_masked_softmax_warp_forward(output_t *dst, c #pragma unroll for (int it = 0; it < WARP_ITERATIONS; ++it) { if (it < warp_iteration_limit) { - elements[i][it] = std::exp((elements[i][it] - max_value[i])); + elements[i][it] = expf((elements[i][it] - max_value[i])); sum[i] += elements[i][it]; } } @@ -333,6 +353,42 @@ __global__ void scaled_upper_triang_masked_softmax_warp_backward(output_t *gradI } } +#ifdef __CUDACC_RTC__ + +#else + +namespace { + +constexpr const char *kRtcSourceFile = + "transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu"; + +template +std::string make_softmax_rtc_code(int log2_elements) { + (void)log2_elements; + std::string code = string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu; + return code; +} + +template +std::string make_softmax_rtc_label(const char *direction, int log2_elements) { + return concat_strings("fused_softmax,variant=upper_triang,direction=", direction, + ",type=", TypeInfo::name, ",log2=", log2_elements, ",fast_math=1"); +} + +template +std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); +} + +void throw_nvrtc_required(const char *direction) { + NVTE_ERROR("Fused upper-triangular softmax RTC path is disabled for ", direction, + ". Set NVTE_DISABLE_NVRTC=0 or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=ON."); +} + +} // namespace + template void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const input_t *src, const input_t scale, int softmax_elements, @@ -365,85 +421,104 @@ void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const in int blocks_per_seq = attn_batches / batches_per_block; dim3 blocks(seq_len, blocks_per_seq, 1); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 1: // 2 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 2: // 4 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 3: // 8 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 4: // 16 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 5: // 32 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 6: // 64 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 7: // 128 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 8: // 256 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 9: // 512 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 10: // 1024 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 11: // 2048 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 12: // 4096 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 13: // 8192 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 14: // 16384 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_upper_triang_masked_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, rtc_scale, batch_count, + softmax_elements_stride, softmax_elements); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 1: // 2 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 2: // 4 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 3: // 8 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 4: // 16 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 5: // 32 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 6: // 64 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 7: // 128 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 8: // 256 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 9: // 512 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 10: // 1024 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 11: // 2048 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 12: // 4096 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 13: // 8192 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 14: // 16384 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + default: + break; + } +#else + throw_nvrtc_required("forward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -482,85 +557,103 @@ void dispatch_scaled_upper_triang_masked_softmax_backward(output_t *grad_input, int blocks_per_seq = attn_batches / batches_per_block; dim3 blocks(seq_len, blocks_per_seq, 1); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 1: // 2 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 2: // 4 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 3: // 8 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 4: // 16 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 5: // 32 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 6: // 64 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 7: // 128 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 8: // 256 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 9: // 512 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 10: // 1024 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 11: // 2048 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 12: // 4096 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 13: // 8192 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 14: // 16384 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("backward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_upper_triang_masked_softmax_warp_backward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, grad_input, grad, output, scale, + batch_count, softmax_elements_stride, softmax_elements); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 1: // 2 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 2: // 4 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 3: // 8 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 4: // 16 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 5: // 32 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 6: // 64 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 7: // 128 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 8: // 256 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 9: // 512 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 10: // 1024 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 11: // 2048 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 12: // 4096 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 13: // 8192 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 14: // 16384 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + default: + break; + } +#else + throw_nvrtc_required("backward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -595,8 +688,12 @@ void scaled_upper_triang_masked_softmax_backward(Tensor output_grads, const Tens seq_len, attn_batches, stream);); } +#endif // __CUDACC_RTC__ + } // end namespace transformer_engine +#ifndef __CUDACC_RTC__ + void nvte_scaled_upper_triang_masked_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream) { @@ -615,3 +712,5 @@ void nvte_scaled_upper_triang_masked_softmax_backward(const NVTETensor incoming_ *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads), *convertNVTETensorCheck(softmax_results), scale_factor, stream); } + +#endif // __CUDACC_RTC__ diff --git a/transformer_engine/common/hadamard_transform/customized_pipeline.cuh b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh index bc46341e88..11754fbb50 100644 --- a/transformer_engine/common/hadamard_transform/customized_pipeline.cuh +++ b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh @@ -18,7 +18,7 @@ namespace detail { // by producer_commit. Use case, accumulator generation as // the result of MMA instructions. template , - class AtomThrShape_MNK_ = Shape<_1, _1, _1> > + class AtomThrShape_MNK_ = Shape<_1, _1, _1>> class CustomizedPipelineTmaUmmaAsync { public: static constexpr uint32_t Stages = Stages_; diff --git a/transformer_engine/common/normalization/common.h b/transformer_engine/common/normalization/common.h index f5dce64193..31a547f62c 100644 --- a/transformer_engine/common/normalization/common.h +++ b/transformer_engine/common/normalization/common.h @@ -24,6 +24,7 @@ #include "../common.h" #include "../cudnn_utils.h" #include "../util/system.h" +#include "kernel_params.h" namespace transformer_engine { @@ -51,97 +52,6 @@ struct LaunchParams { } }; -struct KernelParamsBase { - KernelParamsBase() - : ctas_per_col(0), - rows(0), - cols(0), - x(nullptr), - mu(nullptr), - rs(nullptr), - gamma(nullptr), - workspace(nullptr), - barrier(nullptr), - zero_centered_gamma(false) {} - - // For Multi-CTA, number of different CTA groups. Otherwise same as gridDim.x. - int ctas_per_col; - // Size of CTA group. - int ctas_per_row; - - // Input is interpreted as matrix. We normalize across columns. - int rows; - int cols; - - // Common data pointers. - void* x; - void* mu; - void* rs; - void* gamma; - - // Multi-CTA workspace in gmem. - void* workspace; - - // Multi-CTA sync barriers in gmem. - int* barrier; - - // Whether gamma is centered around 0 - bool zero_centered_gamma; -}; - -struct ForwardKernelParams : public KernelParamsBase { - ForwardKernelParams() - : KernelParamsBase(), z(nullptr), beta(nullptr), epsilon(0.f), fp8_out(false) {} - - // Output of LN FWD. - void* z; - void* beta; - float epsilon; - - // Scaling factor - void* scale; - int scale_byte_size; - - // Inverse of scaling factor - void* scale_inv; - - // AMax output - void* amax; - int amax_byte_size; - - // Whether to compute scale and amax - bool fp8_out; -}; - -struct BackwardKernelParams : public KernelParamsBase { - BackwardKernelParams() - : KernelParamsBase(), - dz(nullptr), - dbeta_part(nullptr), - dgamma_part(nullptr), - dx(nullptr), - dbeta(nullptr), - dgamma(nullptr) {} - - // Input: gradient wrt. LN FWD output. - void* dz; - - // Input: extra tensor to add for fused backward+add - void* add; - - // Workspace for Wgrad pre-reduction. - void* dbeta_part; - void* dgamma_part; - - // Output: Dgrad. - void* dx; - // Output: Wgrad. - void* dbeta; - void* dgamma; -}; - -using BackwardAddKernelParams = BackwardKernelParams; - enum class NVTE_Norm_Backend { Te, Cudnn }; enum class NVTE_Norm_Stage { Forward, Backward, BackwardAdd }; @@ -191,6 +101,16 @@ class TeNormalizationRegistry { return 0; } + // Overload for capturing-callable dispatchers (e.g. NVRTC closures). + static int registerFunction(TupleKeyType key, Function func) { + auto [general_key, batch_size, hidden_size, is_tuned] = key; + if (is_tuned) + getInstance().tuned_function_map.emplace(key, std::move(func)); + else + getInstance().general_function_map[general_key].emplace(hidden_size, std::move(func)); + return 0; + } + static Function getKernel(TupleKeyType key) { auto& instance = getInstance(); auto [general_key, batch_size, hidden_size, is_tuned] = key; diff --git a/transformer_engine/common/normalization/kernel_params.h b/transformer_engine/common/normalization/kernel_params.h new file mode 100644 index 0000000000..42dc423dd0 --- /dev/null +++ b/transformer_engine/common/normalization/kernel_params.h @@ -0,0 +1,88 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_PARAMS_H_ +#define TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_PARAMS_H_ + +// POD kernel-parameter structs shared between the host norm dispatchers and +// the NVRTC kernel sources. This header is intentionally free of host-only +// includes (no , no cudnn) so it can be compiled by +// NVRTC as well. + +namespace transformer_engine { +namespace normalization { + +struct KernelParamsBase { + // For Multi-CTA, number of different CTA groups. Otherwise same as gridDim.x. + int ctas_per_col = 0; + // Size of CTA group. + int ctas_per_row = 0; + + // Input is interpreted as matrix. We normalize across columns. + int rows = 0; + int cols = 0; + + // Common data pointers. + void* x = nullptr; + void* mu = nullptr; + void* rs = nullptr; + void* gamma = nullptr; + + // Multi-CTA workspace in gmem. + void* workspace = nullptr; + + // Multi-CTA sync barriers in gmem. + int* barrier = nullptr; + + // Whether gamma is centered around 0 + bool zero_centered_gamma = false; +}; + +struct ForwardKernelParams : public KernelParamsBase { + // Output of LN FWD. + void* z = nullptr; + void* beta = nullptr; + float epsilon = 0.f; + + // Scaling factor + void* scale = nullptr; + int scale_byte_size = 0; + + // Inverse of scaling factor + void* scale_inv = nullptr; + + // AMax output + void* amax = nullptr; + int amax_byte_size = 0; + + // Whether to compute scale and amax + bool fp8_out = false; +}; + +struct BackwardKernelParams : public KernelParamsBase { + // Input: gradient wrt. LN FWD output. + void* dz = nullptr; + + // Input: extra tensor to add for fused backward+add + void* add = nullptr; + + // Workspace for Wgrad pre-reduction. + void* dbeta_part = nullptr; + void* dgamma_part = nullptr; + + // Output: Dgrad. + void* dx = nullptr; + // Output: Wgrad. + void* dbeta = nullptr; + void* dgamma = nullptr; +}; + +using BackwardAddKernelParams = BackwardKernelParams; + +} // namespace normalization +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_PARAMS_H_ diff --git a/transformer_engine/common/normalization/kernel_traits.h b/transformer_engine/common/normalization/kernel_traits.h index 12fc095c38..41be82e742 100644 --- a/transformer_engine/common/normalization/kernel_traits.h +++ b/transformer_engine/common/normalization/kernel_traits.h @@ -7,12 +7,29 @@ #ifndef TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_TRAITS_H_ #define TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_TRAITS_H_ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include "../common.h" #include "../utils.cuh" +#endif namespace transformer_engine { namespace normalization { +#ifdef __CUDACC_RTC__ +// Under NVRTC the kernel sources do not include common.h (it pulls in cuDNN and +// other host-only headers), so the dtype aliases that the RTC name expressions +// reference are nototherwise visible (e.g. ::transformer_engine::normalization::fp16). +// Mirror the definitions from common.h here for the RTC. +// The underlying CUDA types come from utils.cuh (cuda_fp16/bf16/fp8 headers). +using fp32 = float; +using fp16 = half; +using bf16 = nv_bfloat16; +using fp8e4m3 = __nv_fp8_e4m3; +using fp8e5m2 = __nv_fp8_e5m2; +#endif // __CUDACC_RTC__ + template struct Kernel_traits_base { @@ -31,7 +48,7 @@ template > + compute_t_, index_t_, THREADS_PER_CTA_>> struct Kernel_traits_finalize : public Base { enum { ROWS_PER_CTA = Base::THREADS_PER_CTA / Base::THREADS_PER_WARP }; static_assert(static_cast(ROWS_PER_CTA) <= static_cast(Base::THREADS_PER_WARP)); @@ -69,7 +86,7 @@ template > + WARPS_M_ * WARPS_N_ * THREADS_PER_WARP>> struct Kernel_traits : public Base { using input_t = typename Base::input_t; using weight_t = typename Base::weight_t; diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh index c4b00b87c3..0480b588b9 100644 --- a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh @@ -7,8 +7,13 @@ #ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_BWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_BWD_KERNELS_CUH_ +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu b/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu index 68aa0942c1..2e3d5fd622 100644 --- a/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu @@ -7,6 +7,7 @@ #include "../../common.h" #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "ln_bwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -132,6 +133,30 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, NVTE_CHECK_CUDA(cudaGetLastError()); } +#define REGISTER_NORM_LAUNCHER_LN_BWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ + WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_bwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_bwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_LN_BWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_bwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_bwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -141,10 +166,16 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, launch_ln_bwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_LN_BWD_##LAUNCH_TYPE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ - } // namespace + } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_LN_BWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, ... diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu index 464df8d276..5be13357b1 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "ln_fwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -101,6 +102,33 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, } } +// Register a single RTC-first dispatcher. When the static fallback is enabled, +// its function pointer is passed to the dispatcher and selected only when +// NVTE_DISABLE_NVRTC=1. +#define REGISTER_NORM_LAUNCHER_LN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ + WARPS_M, WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_fwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_fwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, \ + BYTES_PER_LDG, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_LN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_fwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_fwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -110,10 +138,16 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, launch_ln_fwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_LN_FWD_##LAUNCH_TYPE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_LN_FWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh index 5a37cf46da..3e9c281f58 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh @@ -7,11 +7,16 @@ #ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include #include #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -147,7 +152,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_tuned_kernel( if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::rtc_detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } @@ -323,7 +328,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_general_kerne if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::rtc_detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } diff --git a/transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu b/transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu new file mode 100644 index 0000000000..f93f9e6edb --- /dev/null +++ b/transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu @@ -0,0 +1,8 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "kernel_traits.h" +#include "ln_bwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu b/transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu new file mode 100644 index 0000000000..5c34f05e8d --- /dev/null +++ b/transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu @@ -0,0 +1,12 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVRTC source file for LayerNorm forward kernels. The host bundles this as a +// string header; specific template instantiations are requested via +// nvrtcAddNameExpression at runtime. + +#include "kernel_traits.h" +#include "ln_fwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh index d620ee5260..4022ff97e7 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh @@ -7,10 +7,13 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ -#include - +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -18,19 +21,34 @@ namespace normalization { struct maybe_not_t {}; template -using maybe_t = std::conditional_t; - -template +using maybe_t = transformer_engine::rtc_detail::conditional_t; + +// dx and add share storage; `add` is positioned at the tail of the `dx` +// storage via leading padding. NeedsPadding is false when dx_t and add_t are +// the same size (or add_t is larger), in which case the padding array would be +// zero-length -- legal as a GNU extension under nvcc but rejected by NVRTC. The +// no-padding specialization below covers that case so both compilers are happy +// while keeping an identical layout. +template sizeof(maybe_t))> union dx_add_t { using add_t = maybe_t; using dx_t = Ivec; struct { - char _padding[sizeof(dx_t) > sizeof(add_t) ? sizeof(dx_t) - sizeof(add_t) : 0]; + char _padding[sizeof(dx_t) - sizeof(add_t)]; add_t add; }; dx_t dx; }; +template +union dx_add_t { + using add_t = maybe_t; + using dx_t = Ivec; + add_t add; + dx_t dx; +}; + template __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_bwd_tuned_kernel( BackwardKernelParams params) { diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu index 60238f256d..7e8a1cd624 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "rmsnorm_bwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -133,6 +134,60 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para NVTE_CHECK_CUDA(cudaGetLastError()); } +// Two-level dispatch: stage (Backward / BackwardAdd) and launch type +// (tuned / general). The optional static fallback is selected only when +// NVTE_DISABLE_NVRTC=1. +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_tuned_Backward(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_bwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, false, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_tuned_BackwardAdd(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, ADD_FLAG, STATIC_FALLBACK) \ + static_assert(ADD_FLAG, "RMSNorm BackwardAdd registrations require ADD_FLAG=true"); \ + [[maybe_unused]] static const int \ + _rmsn_bwd_tuned_add_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, ADD_FLAG, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_general_Backward( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_bwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + false, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_general_BackwardAdd(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, \ + CTYPE, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, ADD_FLAG, STATIC_FALLBACK) \ + static_assert(ADD_FLAG, "RMSNorm BackwardAdd registrations require ADD_FLAG=true"); \ + [[maybe_unused]] static const int \ + _rmsn_bwd_general_add_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + ADD_FLAG, STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -142,10 +197,16 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para launch_rmsnorm_bwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_RMSN_BWD_##LAUNCH_TYPE##_##NORM_STAGE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_RMSN_BWD_##LAUNCH_TYPE##_##NORM_STAGE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, \ + CTYPE, __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create rmsnorm bwd tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, ... diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu index 5522fd5c6b..c3308cbce0 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "rmsnorm_fwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -102,6 +103,31 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param } } +#define REGISTER_NORM_LAUNCHER_RMSN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_fwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_fwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, \ + BYTES_PER_LDG, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_fwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_fwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -111,10 +137,16 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param launch_rmsnorm_fwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_RMSN_FWD_##LAUNCH_TYPE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_RMSN_FWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create rmsnorm tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh index 900fb58be2..0f06b29a21 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh @@ -7,11 +7,16 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include #include #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -139,7 +144,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_fwd_tuned_ke if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::rtc_detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } @@ -298,7 +303,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_fwd_general_ if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::rtc_detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } diff --git a/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu new file mode 100644 index 0000000000..706e79870b --- /dev/null +++ b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu @@ -0,0 +1,8 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "kernel_traits.h" +#include "rmsnorm_bwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu new file mode 100644 index 0000000000..4d1bd10a7b --- /dev/null +++ b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu @@ -0,0 +1,8 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "kernel_traits.h" +#include "rmsnorm_fwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/rtc_dispatch.cpp b/transformer_engine/common/normalization/rtc_dispatch.cpp new file mode 100644 index 0000000000..5ae81f4233 --- /dev/null +++ b/transformer_engine/common/normalization/rtc_dispatch.cpp @@ -0,0 +1,742 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVRTC-backed registry registration for LayerNorm/RMSNorm forward + backward +// launchers. Each per-config REGISTER_NORM_LAUNCHER macro expands to a call +// into one of the register_*_tuned/general functions defined here. When +// NVTE_BUILD_LEGACY_STATIC_NORM=ON, that call also supplies a static fallback. +// The registered closure prefers NVRTC and selects the fallback only when +// NVTE_DISABLE_NVRTC=1. + +#include "rtc_dispatch.h" + +#include +#include +#include + +#include "../util/cuda_driver.h" +#include "../util/rtc.h" +#include "../util/string.h" +#include "common.h" + +// NVRTC source strings for the four kernel families. These are tiny stub +// files (each is a couple of #include lines that pull in kernel_traits + the +// matching kernel header); NVRTC then instantiates a specific +// (Kernel_traits) on demand via nvrtcAddNameExpression. +#include "string_code_normalization_kernel_params_h.h" +#include "string_code_normalization_kernel_traits_h.h" +#include "string_code_normalization_layernorm_ln_bwd_kernels_cuh.h" +#include "string_code_normalization_layernorm_ln_fwd_kernels_cuh.h" +#include "string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu.h" +#include "string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu.h" +#include "string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh.h" +#include "string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh.h" +#include "string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu.h" +#include "string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu.h" + +namespace transformer_engine { +namespace normalization { +namespace rtc_norm { + +namespace { + +const std::vector& norm_headers() { + static const std::vector headers = { + {string_code_normalization_kernel_params_h, "kernel_params.h"}, + {string_code_normalization_kernel_traits_h, "kernel_traits.h"}, + {string_code_normalization_layernorm_ln_fwd_kernels_cuh, "ln_fwd_kernels.cuh"}, + {string_code_normalization_layernorm_ln_bwd_kernels_cuh, "ln_bwd_kernels.cuh"}, + {string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh, "rmsnorm_fwd_kernels.cuh"}, + {string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh, "rmsnorm_bwd_kernels.cuh"}, + }; + return headers; +} + +void compile_norm_kernel(rtc::KernelManager& manager, const std::string& label, + const std::string& kernel_expr, const char* rtc_source, + const char* filename) { + manager.compile(label, kernel_expr, rtc_source, filename, {}, norm_headers()); +} + +// Map our DType enum onto the C++ type names used inside the norm RTC sources. +// Aliases come from normalization/common.h (`using bf16 = nv_bfloat16;` etc.). +const char* cpp_name_for(DType dt) { + switch (dt) { + case DType::kFloat32: + return "::transformer_engine::normalization::fp32"; + case DType::kFloat16: + return "::transformer_engine::normalization::fp16"; + case DType::kBFloat16: + return "::transformer_engine::normalization::bf16"; + case DType::kFloat8E4M3: + return "::transformer_engine::normalization::fp8e4m3"; + case DType::kFloat8E5M2: + return "::transformer_engine::normalization::fp8e5m2"; + default: + NVTE_ERROR("Unsupported DType for norm RTC dispatch"); + } +} + +int byte_size_of(DType dt) { + switch (dt) { + case DType::kFloat32: + return 4; + case DType::kFloat16: + case DType::kBFloat16: + return 2; + case DType::kFloat8E4M3: + case DType::kFloat8E5M2: + return 1; + default: + NVTE_ERROR("Unsupported DType for norm RTC dispatch"); + } +} + +// Build a Kernel_traits template-argument list as a string. The C++ argument +// list matches: +// Kernel_traits +std::string kernel_traits_expr(DType wt, DType it, DType ot, DType ct, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg) { + return concat_strings("::transformer_engine::normalization::Kernel_traits<", cpp_name_for(wt), + ", ", cpp_name_for(it), ", ", cpp_name_for(ot), ", ", cpp_name_for(ct), + ", uint32_t, ", hidden_size, ", ", ctas_per_row, ", ", warps_m, ", ", + warps_n, ", ", bytes_per_ldg, ">"); +} + +// Stats::SMEM_BYTES expressed in host code. +// stats_t = TypeToVec2::Type — for compute_t == fp32 that's float2 +// (8 bytes); for fp16 it's half2 (4 bytes); for bf16 it's nv_bfloat162 (4 +// bytes). Matches the formulas in utils.cuh: +// WARPS_N == 1 → 0 +// else → WARPS_M * WARPS_N * sizeof(stats_t) * 2 +int stats_smem_bytes(DType ctype, int warps_m, int warps_n) { + if (warps_n == 1) return 0; + int sizeof_stats_t; + switch (ctype) { + case DType::kFloat32: + sizeof_stats_t = 8; + break; + case DType::kFloat16: + case DType::kBFloat16: + sizeof_stats_t = 4; + break; + default: + NVTE_ERROR("Unsupported compute dtype for norm smem calc"); + } + return warps_m * warps_n * sizeof_stats_t * 2; +} + +// Reducer::SMEM_BYTES. +// reduce_t = TypeToVec2::Type — same sizes as stats_t. +int reducer_smem_bytes(DType ctype, int warps_m, int warps_n) { + return stats_smem_bytes(ctype, warps_m, warps_n); +} + +// Kernel_traits::SMEM_BYTES (used by backward launchers): +// SMEM_BYTES_DGRAD = Reducer::SMEM_BYTES +// SMEM_BYTES_WGRAD = (CTAS_PER_ROW > 1) ? 0 : WARPS_M * HIDDEN * sizeof(compute_t) +// SMEM_BYTES = DGRAD + WGRAD +int bwd_smem_bytes(DType ctype, int hidden_size, int ctas_per_row, int warps_m, int warps_n) { + const int dgrad = reducer_smem_bytes(ctype, warps_m, warps_n); + const int wgrad = (ctas_per_row > 1) ? 0 : warps_m * hidden_size * byte_size_of(ctype); + return dgrad + wgrad; +} + +template +bool try_static_fallback(StaticFallback static_fallback, + LaunchParams& launch_params, const bool configure_params) { + if (rtc::is_enabled()) { + return false; + } + NVTE_CHECK(static_fallback != nullptr, + "NVRTC is disabled and this normalization build has no static fallback. Rebuild with " + "NVTE_BUILD_LEGACY_STATIC_NORM=ON."); + static_fallback(launch_params, configure_params); + return true; +} + +// Common per-launch configure/launch helper. `kernel_expr` is the full +// templated kernel symbol (used as the nvrtcAddNameExpression argument); the +// closure captures everything needed. +template +void register_launcher(const std::string& label, const std::string& kernel_expr, + const char* rtc_source, const char* filename, TupleKeyType key, + int threads_per_cta, int dynamic_smem_bytes, int ctas_per_row, + bool needs_cooperative, int barrier_bytes_per_col, + int workspace_bytes_per_col, int dgamma_part_bytes_per_col, + StaticFallback static_fallback) { + auto closure = [label, kernel_expr, rtc_source, filename, threads_per_cta, dynamic_smem_bytes, + ctas_per_row, needs_cooperative, barrier_bytes_per_col, workspace_bytes_per_col, + dgamma_part_bytes_per_col, static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + compile_norm_kernel(mgr, label, kernel_expr, rtc_source, filename); + } + + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(label, threads_per_cta, dynamic_smem_bytes); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = + launch_params.multiprocessorCount * ctas_per_sm / ctas_per_row; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = barrier_bytes_per_col * launch_params.params.ctas_per_col; + launch_params.workspace_bytes = workspace_bytes_per_col * launch_params.params.ctas_per_col; + } + if (dgamma_part_bytes_per_col > 0) { + launch_params.dgamma_part_bytes = + dgamma_part_bytes_per_col * launch_params.params.ctas_per_col; + } + return; + } + + // Real launch. + if (dynamic_smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + dynamic_smem_bytes); + } + const auto stream = launch_params.stream; + const auto ctas_per_col = launch_params.params.ctas_per_col; + if (ctas_per_row == 1) { + mgr.launch(label, dim3(ctas_per_col), dim3(threads_per_cta), dynamic_smem_bytes, stream, + launch_params.params); + } else { + mgr.launch_cooperative(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + dynamic_smem_bytes, stream, launch_params.params); + } + (void)needs_cooperative; + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +} // namespace + +// ============================================================================ +// LayerNorm Forward +// ============================================================================ + +void register_ln_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl, StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Forward, wt, it, ot, ct, 0, hidden, false, true); + const std::string label = + concat_strings("ln_fwd_tuned,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",cr=", cr, ",wm=", wm, ",wn=", wn, ",bl=", bl); + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::ln_fwd_tuned_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = stats_smem_bytes(ct, wm, wn); + // tuned path multi-CTA workspace formula: + // barrier_bytes = 2 * ctas_per_col * sizeof(index_t == uint32_t == 4 bytes) + // workspace_bytes = ctas_per_col * WARPS_M * CTAS_PER_ROW * sizeof(stats_t) * 2 + const int sizeof_stats_t = (ct == DType::kFloat32) ? 8 : 4; // float2 vs half2/bf162 + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_stats_t * 2; + register_launcher( + label, kernel_expr, string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu, + "ln_fwd_kernel.cu", key, threads_per_cta, smem_bytes, cr, /*needs_cooperative=*/cr > 1, + barrier_per_col, workspace_per_col, /*dgamma_part_bytes_per_col=*/0, static_fallback); +} + +void register_ln_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, int wn, + int bl, StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Forward, wt, it, ot, ct, 0, hidden, false, false); + const std::string label = + concat_strings("ln_fwd_general,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",wm=", wm, ",wn=", wn, ",bl=", bl); + // "general" path always uses CTAS_PER_ROW=1 in the Kernel_traits. + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::ln_fwd_general_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const auto closure = [label, kernel_expr, threads_per_cta, hidden, wm, ct, static_fallback]( + LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + compile_norm_kernel(mgr, label, kernel_expr, + string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu, + "ln_fwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(label, threads_per_cta, /*smem=*/0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + // compute_t bytes + const int ctype_bytes = byte_size_of(ct); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * ctype_bytes * 2; + } + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, + stream, launch_params.params); + } + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +// ============================================================================ +// RMSNorm Forward (same shape as LayerNorm Forward) +// ============================================================================ + +void register_rmsnorm_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, NVTE_Norm_Stage::Forward, + wt, it, ot, ct, 0, hidden, false, true); + const std::string label = + concat_strings("rmsnorm_fwd_tuned,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",cr=", cr, ",wm=", wm, ",wn=", wn, ",bl=", bl); + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::rmsnorm_fwd_tuned_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = stats_smem_bytes(ct, wm, wn); + const int sizeof_stats_t = (ct == DType::kFloat32) ? 8 : 4; + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_stats_t * 2; + register_launcher( + label, kernel_expr, string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu, + "rmsnorm_fwd_kernel.cu", key, threads_per_cta, smem_bytes, cr, /*needs_cooperative=*/cr > 1, + barrier_per_col, workspace_per_col, 0, static_fallback); +} + +void register_rmsnorm_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, + int wn, int bl, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, NVTE_Norm_Stage::Forward, + wt, it, ot, ct, 0, hidden, false, false); + const std::string label = + concat_strings("rmsnorm_fwd_general,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",wm=", wm, ",wn=", wn, ",bl=", bl); + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::rmsnorm_fwd_general_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const auto closure = [label, kernel_expr, threads_per_cta, hidden, wm, ct, static_fallback]( + LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + compile_norm_kernel(mgr, label, kernel_expr, + string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu, + "rmsnorm_fwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = mgr.occupancy_max_active_blocks_per_sm(label, threads_per_cta, 0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + const int ctype_bytes = byte_size_of(ct); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * ctype_bytes * 2; + } + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, + stream, launch_params.params); + } + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +// ============================================================================ +// LayerNorm Backward (main kernel + finalize kernel) +// ============================================================================ + +void register_ln_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl_main, int bl_final, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Backward, wt, it, ot, ct, 0, hidden, false, true); + const std::string label = + concat_strings("ln_bwd_tuned,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",cr=", cr, ",wm=", wm, ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string main_traits = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl_main); + const std::string main_kexpr = concat_strings( + "&::transformer_engine::normalization::ln_bwd_tuned_kernel<", main_traits, ">"); + // Kernel_traits_finalize + const std::string finalize_traits = + concat_strings("::transformer_engine::normalization::Kernel_traits_finalize<", hidden, ", ", + cpp_name_for(wt), ", ", cpp_name_for(it), ", ", cpp_name_for(ot), ", ", + cpp_name_for(ct), ", uint32_t, 1024, ", bl_final, ">"); + const std::string finalize_kexpr = concat_strings( + "&::transformer_engine::normalization::ln_bwd_finalize_tuned_kernel<", finalize_traits, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = bwd_smem_bytes(ct, hidden, cr, wm, wn); + const int sizeof_reduce_t = (ct == DType::kFloat32) ? 8 : 4; + // tuned backward: workspace = ctas_per_col * WARPS_M * CTAS_PER_ROW * sizeof(reduce_t) * 2 + // dgamma_part_bytes = ctas_per_col * cols * sizeof(compute_t) + // barrier_bytes = 2 * ctas_per_col * sizeof(index_t) + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_reduce_t * 2; + const int dgamma_part_per_col = hidden * byte_size_of(ct); + + // Finalize kernel dims: Kernel_traits_finalize::THREADS_PER_CTA == 1024 + // Kernel_traits_finalize::CTAS = HIDDEN_SIZE / 32 (since COLS%32==0) + // COLS = HIDDEN_SIZE * sizeof(compute_t) / BYTES_PER_LDG_FINAL + // CTAS = COLS / 32 + const int colspass = hidden * byte_size_of(ct) / bl_final; + const int finalize_ctas = colspass / 32; + const int finalize_threads_per_cta = 1024; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, + smem_bytes, cr, barrier_per_col, workspace_per_col, dgamma_part_per_col, + finalize_ctas, finalize_threads_per_cta, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, smem_bytes); + launch_params.params.ctas_per_row = cr; + launch_params.params.ctas_per_col = launch_params.multiprocessorCount * ctas_per_sm / cr; + if (cr > 1) { + launch_params.barrier_bytes = barrier_per_col * launch_params.params.ctas_per_col; + launch_params.workspace_bytes = workspace_per_col * launch_params.params.ctas_per_col; + } + launch_params.dgamma_part_bytes = dgamma_part_per_col * launch_params.params.ctas_per_col; + return; + } + if (smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(main_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + smem_bytes); + } + const auto stream = launch_params.stream; + const auto ctas_per_col = launch_params.params.ctas_per_col; + if (cr == 1) { + mgr.launch(main_label, dim3(ctas_per_col), dim3(threads_per_cta), smem_bytes, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(cr * ctas_per_col), dim3(threads_per_cta), smem_bytes, + stream, launch_params.params); + } + mgr.launch(finalize_label, dim3(finalize_ctas), dim3(finalize_threads_per_cta), 0, stream, + launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +void register_ln_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, int wn, + int bl_main, int bl_final, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Backward, wt, it, ot, ct, 0, hidden, false, false); + const std::string label = + concat_strings("ln_bwd_general,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",wm=", wm, ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string traits = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl_main); + const std::string main_kexpr = + concat_strings("&::transformer_engine::normalization::ln_bwd_general_kernel<", traits, ">"); + // ln_bwd_finalize_general_kernel + const std::string finalize_kexpr = + concat_strings("&::transformer_engine::normalization::ln_bwd_finalize_general_kernel<", + cpp_name_for(wt), ", ", cpp_name_for(ct), ", 4, 1, ", bl_final, ", 32>"); + const int threads_per_cta = wm * wn * 32; + // general bwd uses ctas_per_row = ceil_div(cols, HIDDEN_SIZE); smem=0 for main kernel call. + const int ctype_bytes = byte_size_of(ct); + const int finalize_threads_per_warp = 32; + const int finalize_warps_n = 1; + const int finalize_warps_m = 4; + const int finalize_elts_n_per_cta = + finalize_threads_per_warp * finalize_warps_n * bl_final / ctype_bytes; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, hidden, + wm, ctype_bytes, finalize_elts_n_per_cta, finalize_warps_n, finalize_warps_m, + finalize_threads_per_warp, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, 0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * (ctype_bytes * 2) * 2; + } + launch_params.dgamma_part_bytes = ctas_per_col * cols * ctype_bytes; + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + 0, stream, launch_params.params); + } + const dim3 fin_block(finalize_threads_per_warp * finalize_warps_n, finalize_warps_m); + const dim3 fin_grid(ceil_div(cols, finalize_elts_n_per_cta), 1); + mgr.launch(finalize_label, fin_grid, fin_block, 0, stream, launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +// ============================================================================ +// RMSNorm Backward (main kernel + finalize kernel; same shape as LayerNorm bwd) +// ============================================================================ + +void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl_main, int bl_final, bool with_add, + StaticFallback static_fallback) { + const auto stage = with_add ? NVTE_Norm_Stage::BackwardAdd : NVTE_Norm_Stage::Backward; + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, stage, wt, it, ot, ct, 0, + hidden, false, true); + const std::string add_tag = with_add ? "_add" : ""; + const std::string label = concat_strings( + "rmsnorm_bwd_tuned", add_tag, ",w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, ",cr=", cr, + ",wm=", wm, ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string traits = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl_main); + const char* add_flag = with_add ? "true" : "false"; + const std::string main_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_tuned_kernel<", traits, + ", ", add_flag, ">"); + const std::string finalize_traits = + concat_strings("::transformer_engine::normalization::Kernel_traits_finalize<", hidden, ", ", + cpp_name_for(wt), ", ", cpp_name_for(it), ", ", cpp_name_for(ot), ", ", + cpp_name_for(ct), ", uint32_t, 1024, ", bl_final, ">"); + const std::string finalize_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_finalize_tuned_kernel<", + finalize_traits, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = bwd_smem_bytes(ct, hidden, cr, wm, wn); + const int sizeof_reduce_t = (ct == DType::kFloat32) ? 8 : 4; + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_reduce_t * 2; + const int dgamma_part_per_col = hidden * byte_size_of(ct); + const int colspass = hidden * byte_size_of(ct) / bl_final; + const int finalize_ctas = colspass / 32; + const int finalize_threads_per_cta = 1024; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, + smem_bytes, cr, barrier_per_col, workspace_per_col, dgamma_part_per_col, + finalize_ctas, finalize_threads_per_cta, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, smem_bytes); + launch_params.params.ctas_per_row = cr; + launch_params.params.ctas_per_col = launch_params.multiprocessorCount * ctas_per_sm / cr; + if (cr > 1) { + launch_params.barrier_bytes = barrier_per_col * launch_params.params.ctas_per_col; + launch_params.workspace_bytes = workspace_per_col * launch_params.params.ctas_per_col; + } + launch_params.dgamma_part_bytes = dgamma_part_per_col * launch_params.params.ctas_per_col; + return; + } + if (smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(main_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + smem_bytes); + } + const auto stream = launch_params.stream; + const auto ctas_per_col = launch_params.params.ctas_per_col; + if (cr == 1) { + mgr.launch(main_label, dim3(ctas_per_col), dim3(threads_per_cta), smem_bytes, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(cr * ctas_per_col), dim3(threads_per_cta), smem_bytes, + stream, launch_params.params); + } + mgr.launch(finalize_label, dim3(finalize_ctas), dim3(finalize_threads_per_cta), 0, stream, + launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +void register_rmsnorm_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, + int wn, int bl_main, int bl_final, bool with_add, + StaticFallback static_fallback) { + const auto stage = with_add ? NVTE_Norm_Stage::BackwardAdd : NVTE_Norm_Stage::Backward; + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, stage, wt, it, ot, ct, 0, + hidden, false, false); + const std::string add_tag = with_add ? "_add" : ""; + const std::string label = concat_strings( + "rmsnorm_bwd_general", add_tag, ",w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, ",wm=", wm, + ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string traits = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl_main); + const char* add_flag = with_add ? "true" : "false"; + const std::string main_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_general_kernel<", traits, + ", ", add_flag, ">"); + const std::string finalize_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_finalize_general_kernel<", + cpp_name_for(wt), ", ", cpp_name_for(ct), ", 4, 1, ", bl_final, ", 32>"); + const int threads_per_cta = wm * wn * 32; + const int ctype_bytes = byte_size_of(ct); + const int finalize_warps_m = 4; + const int finalize_warps_n = 1; + const int finalize_threads_per_warp = 32; + const int finalize_elts_n_per_cta = + finalize_threads_per_warp * finalize_warps_n * bl_final / ctype_bytes; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, hidden, + wm, ctype_bytes, finalize_elts_n_per_cta, finalize_warps_n, finalize_warps_m, + finalize_threads_per_warp, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, 0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * (ctype_bytes * 2) * 2; + } + launch_params.dgamma_part_bytes = ctas_per_col * cols * ctype_bytes; + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + 0, stream, launch_params.params); + } + const dim3 fin_block(finalize_threads_per_warp * finalize_warps_n, finalize_warps_m); + const dim3 fin_grid(ceil_div(cols, finalize_elts_n_per_cta), 1); + mgr.launch(finalize_label, fin_grid, fin_block, 0, stream, launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +} // namespace rtc_norm +} // namespace normalization +} // namespace transformer_engine diff --git a/transformer_engine/common/normalization/rtc_dispatch.h b/transformer_engine/common/normalization/rtc_dispatch.h new file mode 100644 index 0000000000..314f150f70 --- /dev/null +++ b/transformer_engine/common/normalization/rtc_dispatch.h @@ -0,0 +1,65 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_NORM_RTC_DISPATCH_H_ +#define TRANSFORMER_ENGINE_COMMON_NORM_RTC_DISPATCH_H_ + +#include + +#include "common.h" + +namespace transformer_engine { +namespace normalization { +namespace rtc_norm { + +template +using StaticFallback = void (*)(LaunchParams&, const bool); + +// Register an RTC-backed launcher for a single LayerNorm Forward "tuned" +// (multi-CTA-capable) config. Compiles and launches via NVRTC on first use. +void register_ln_fwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); + +// Register an RTC-backed launcher for a single LayerNorm Forward "general" +// (no multi-CTA) config. +void register_ln_fwd_general(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); + +// Register an RTC-backed launcher for a single LayerNorm Backward "tuned" config. +void register_ln_bwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, + StaticFallback static_fallback = nullptr); + +// Register an RTC-backed launcher for a single LayerNorm Backward "general" config. +void register_ln_bwd_general(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, + StaticFallback static_fallback = nullptr); + +// Same set for RMSNorm. +void register_rmsnorm_fwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); +void register_rmsnorm_fwd_general(DType wtype, DType itype, DType otype, DType ctype, + int hidden_size, int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); +void register_rmsnorm_bwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, bool with_add, + StaticFallback static_fallback = nullptr); +void register_rmsnorm_bwd_general(DType wtype, DType itype, DType otype, DType ctype, + int hidden_size, int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, bool with_add, + StaticFallback static_fallback = nullptr); + +} // namespace rtc_norm +} // namespace normalization +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_NORM_RTC_DISPATCH_H_ diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index 70024a202c..20616f8cb6 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -47,12 +47,7 @@ inline int max_supported_sm_arch() { } // namespace bool is_enabled() { - static bool is_enabled_ = false; - static bool need_to_check_env = true; - if (need_to_check_env) { - is_enabled_ = !getenv("NVTE_DISABLE_NVRTC"); - need_to_check_env = false; - } + static const bool is_enabled_ = !getenv("NVTE_DISABLE_NVRTC"); return is_enabled_; } @@ -132,6 +127,18 @@ void Kernel::set_function_cache_config(int device_id, CUfunc_cache cache_config) NVTE_CALL_CHECK_CUDA_DRIVER(cuFuncSetCacheConfig, get_function(device_id), cache_config); } +void Kernel::set_function_attribute(int device_id, CUfunction_attribute attr, int value) { + NVTE_CALL_CHECK_CUDA_DRIVER(cuFuncSetAttribute, get_function(device_id), attr, value); +} + +int Kernel::occupancy_max_active_blocks_per_sm(int device_id, int block_size, + std::size_t dynamic_smem_bytes) { + int num_blocks = 0; + NVTE_CALL_CHECK_CUDA_DRIVER(cuOccupancyMaxActiveBlocksPerMultiprocessor, &num_blocks, + get_function(device_id), block_size, dynamic_smem_bytes); + return num_blocks; +} + KernelManager& KernelManager::instance() { NVTE_CHECK(is_enabled(), "NVRTC support is not enabled"); static KernelManager instance_; @@ -139,11 +146,17 @@ KernelManager& KernelManager::instance() { } void KernelManager::compile(const std::string& kernel_label, const std::string& kernel_name, - const std::string& code, const std::string& filename) { - std::lock_guard lock_guard_(lock_); + const std::string& code, const std::string& filename, + const std::vector& extra_options, + const std::vector
& extra_headers) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::unique_lock lock_guard_(lock_); + if (kernel_cache_.count(key) > 0) { + return; + } // Choose whether to compile to PTX or cubin - const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); const int compile_sm_arch = std::min(sm_arch_, max_supported_sm_arch()); const bool compile_ptx = sm_arch_ != compile_sm_arch; @@ -160,6 +173,7 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& opts.push_back(concat_strings("--gpu-architecture=sm_", compile_sm_arch)); } opts.push_back(concat_strings("-I", cuda::include_directory(true))); + opts.insert(opts.end(), extra_options.begin(), extra_options.end()); std::vector opts_ptrs; for (const auto& opt : opts) { opts_ptrs.push_back(opt.c_str()); @@ -167,11 +181,19 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& // Compile source nvrtcProgram program; - constexpr int num_headers = 2; - constexpr const char* headers[num_headers] = {string_code_utils_cuh, string_code_util_math_h}; - constexpr const char* include_names[num_headers] = {"utils.cuh", "util/math.h"}; - NVTE_CHECK_NVRTC(nvrtcCreateProgram(&program, code.c_str(), filename.c_str(), num_headers, - headers, include_names)); + std::vector headers = {string_code_utils_cuh, string_code_util_math_h}; + std::vector include_names = {"utils.cuh", "util/math.h"}; + headers.reserve(headers.size() + extra_headers.size()); + include_names.reserve(include_names.size() + extra_headers.size()); + for (const auto& header : extra_headers) { + NVTE_CHECK(header.content != nullptr && header.include_name != nullptr, + "NVRTC header content and include name must not be null"); + headers.push_back(header.content); + include_names.push_back(header.include_name); + } + NVTE_CHECK_NVRTC(nvrtcCreateProgram(&program, code.c_str(), filename.c_str(), + static_cast(headers.size()), headers.data(), + include_names.data())); NVTE_CHECK_NVRTC(nvrtcAddNameExpression(program, kernel_name.c_str())); const nvrtcResult compile_result = nvrtcCompileProgram(program, opts_ptrs.size(), opts_ptrs.data()); @@ -239,7 +261,6 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& } // Cache compiled code - const auto key = get_kernel_cache_key(kernel_label, device_id); kernel_cache_.insert({key, Kernel(mangled_name, std::move(compiled_code))}); kernel_cache_.at(key).get_function(device_id); // Make sure kernel is available on device @@ -250,12 +271,34 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& void KernelManager::set_cache_config(const std::string& kernel_label, CUfunc_cache cache_config) { const int device_id = cuda::current_device(); const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to configure RTC kernel before compilation"); kernel_cache_.at(key).set_function_cache_config(device_id, cache_config); } +void KernelManager::set_function_attribute(const std::string& kernel_label, + CUfunction_attribute attr, int value) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); + NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to configure RTC kernel before compilation"); + kernel_cache_.at(key).set_function_attribute(device_id, attr, value); +} + +int KernelManager::occupancy_max_active_blocks_per_sm(const std::string& kernel_label, + int block_size, + std::size_t dynamic_smem_bytes) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); + NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to query occupancy before compilation"); + return kernel_cache_.at(key).occupancy_max_active_blocks_per_sm(device_id, block_size, + dynamic_smem_bytes); +} + bool KernelManager::is_compiled(const std::string& kernel_label, int device_id) const { const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); return kernel_cache_.count(key) > 0; } diff --git a/transformer_engine/common/util/rtc.h b/transformer_engine/common/util/rtc.h index 65faf7bcc2..745566ee60 100644 --- a/transformer_engine/common/util/rtc.h +++ b/transformer_engine/common/util/rtc.h @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -33,6 +34,12 @@ namespace rtc { */ bool is_enabled(); +/*! \brief Header made available to an NVRTC program */ +struct Header { + const char *content; + const char *include_name; +}; + /*! \brief Wrapper class for a runtime-compiled CUDA kernel */ class Kernel { public: @@ -79,6 +86,29 @@ class Kernel { */ void set_function_cache_config(int device_id, CUfunc_cache cache_config); + /*! \brief Set a kernel function attribute (driver-API wrapper of + * cuFuncSetAttribute, e.g. for dynamic shared memory size). + */ + void set_function_attribute(int device_id, CUfunction_attribute attr, int value); + + /*! \brief Wrapper of cuOccupancyMaxActiveBlocksPerMultiprocessor for a + * runtime-compiled function. + */ + int occupancy_max_active_blocks_per_sm(int device_id, int block_size, + std::size_t dynamic_smem_bytes); + + /*! \brief Cooperative launch of an RTC kernel via cuLaunchCooperativeKernel. + */ + template + void launch_cooperative(int device_id, const dim3 grid_dim, const dim3 block_dim, + unsigned int shared_mem_bytes, cudaStream_t stream, ArgTs &&...args) { + cuda_driver::ensure_context_exists(); + void *arg_ptrs[] = {const_cast(static_cast(&args))...}; + NVTE_CALL_CHECK_CUDA_DRIVER(cuLaunchCooperativeKernel, get_function(device_id), grid_dim.x, + grid_dim.y, grid_dim.z, block_dim.x, block_dim.y, block_dim.z, + shared_mem_bytes, static_cast(stream), arg_ptrs); + } + private: /*! \brief Mangled function name */ std::string mangled_name_; @@ -113,9 +143,13 @@ class KernelManager { * \param[in] code Kernel source code * \param[in] filename Path to associate with source code, * primarily for debugging + * \param[in] extra_options Additional NVRTC compiler options + * \param[in] extra_headers Additional in-memory headers available to the program */ void compile(const std::string &kernel_label, const std::string &kernel_name, - const std::string &code, const std::string &filename); + const std::string &code, const std::string &filename, + const std::vector &extra_options = {}, + const std::vector
&extra_headers = {}); /*! \brief Whether CUDA kernel has been compiled for CUDA device * @@ -143,6 +177,7 @@ class KernelManager { unsigned int shared_mem_bytes, cudaStream_t stream, ArgTs &&...args) { const int device_id = cuda::current_device(); const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to launch RTC kernel before compilation"); kernel_cache_.at(key).launch(device_id, grid_dim, block_dim, shared_mem_bytes, stream, std::forward(args)...); @@ -157,11 +192,32 @@ class KernelManager { */ void set_cache_config(const std::string &kernel_label, CUfunc_cache cache_config); + /*! \brief Set a function attribute (e.g. cuFuncAttributeMaxDynamicSharedMemorySize). */ + void set_function_attribute(const std::string &kernel_label, CUfunction_attribute attr, + int value); + + /*! \brief Query cuOccupancyMaxActiveBlocksPerMultiprocessor for a compiled kernel. */ + int occupancy_max_active_blocks_per_sm(const std::string &kernel_label, int block_size, + std::size_t dynamic_smem_bytes); + + /*! \brief Cooperative launch wrapper (cuLaunchCooperativeKernel). */ + template + void launch_cooperative(const std::string &kernel_label, const dim3 grid_dim, + const dim3 block_dim, unsigned int shared_mem_bytes, cudaStream_t stream, + ArgTs &&...args) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); + NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to launch RTC kernel before compilation"); + kernel_cache_.at(key).launch_cooperative(device_id, grid_dim, block_dim, shared_mem_bytes, + stream, std::forward(args)...); + } + private: /*! \brief Compiled kernels */ std::unordered_map kernel_cache_; - /*! \brief Mutex for thread-safe compilation */ - std::mutex lock_; + /*! \brief Mutex for thread-safe cache access */ + mutable std::shared_mutex lock_; KernelManager() = default; ~KernelManager() = default; diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 635c7a36d2..dd7dcdc4cc 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -30,6 +30,37 @@ static_assert(sizeof(uint32_t) == 4); static_assert(sizeof(uint64_t) == 8); #endif +// Minimal subset of used by RTC kernel headers. Keep these in a +// project-owned namespace because adding primary templates to std is undefined. +namespace transformer_engine { +namespace rtc_detail { + +template +struct is_same { + static constexpr bool value = false; +}; +template +struct is_same { + static constexpr bool value = true; +}; + +template +inline constexpr bool is_same_v = is_same::value; + +template +struct conditional { + using type = T; +}; +template +struct conditional { + using type = F; +}; +template +using conditional_t = typename conditional::type; + +} // namespace rtc_detail +} // namespace transformer_engine + //////////////////////////////////////////////////////////////////////////////////////////////////// constexpr uint32_t THREADS_PER_WARP = 32; @@ -940,13 +971,13 @@ __device__ __forceinline__ float ordered_uint_to_float(unsigned int u) { template __device__ __forceinline__ T abs_val(T val) { - if constexpr (std::is_same_v) { + if constexpr (rtc_detail::is_same_v) { #if __CUDA_ARCH__ >= 800 return __habs(val); #else return static_cast<__nv_bfloat16>(fabsf(static_cast(val))); #endif - } else if constexpr (std::is_same_v) { + } else if constexpr (rtc_detail::is_same_v) { return __habs(val); } else { return fabsf(val); @@ -955,13 +986,13 @@ __device__ __forceinline__ T abs_val(T val) { template __device__ __forceinline__ T max_val(T a, T b) { - if constexpr (std::is_same_v) { + if constexpr (rtc_detail::is_same_v) { #if __CUDA_ARCH__ >= 800 return __hmax(a, b); #else return static_cast<__nv_bfloat16>(fmaxf(static_cast(a), static_cast(b))); #endif - } else if constexpr (std::is_same_v) { + } else if constexpr (rtc_detail::is_same_v) { return __hmax(a, b); } else { return fmaxf(a, b);