From 5b686a2a1f03871a77ac653971606673cbe19582 Mon Sep 17 00:00:00 2001 From: Carlos Gomes Date: Wed, 3 Jun 2026 14:51:55 +0000 Subject: [PATCH 1/9] [Common] NVRTC for fused softmax and normalization (Phase 0) Move the fused-softmax and LayerNorm/RMSNorm kernels from build-time template instantiation to runtime NVRTC compilation, with full coverage of the existing kernel set so the NVRTC path is the default. Fused softmax: - RTC compile/launch path for scaled / scaled-masked / scaled-upper-triangular / scaled-aligned-causal softmax, keyed by dtype, shape and mask/causal mode. - NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX (default OFF) restores the static template dispatch. Normalization (LayerNorm + RMSNorm, forward + backward): - Replace the static REGISTER_NORM_LAUNCHER template fanout with an NVRTC registry that compiles the selected (norm type, direction, dtypes, hidden size, CTA config) kernel on first use and caches it. - NVTE_BUILD_LEGACY_STATIC_NORM (default OFF) restores the static launchers. - NVRTC-safe kernel sources: kernel sources/headers avoid common.h under __CUDACC_RTC__; add the dtype aliases and a minimal std::is_same/conditional_t in the RTC build, and replace a zero-length padding array (a GNU extension nvcc accepts but NVRTC rejects) with a no-padding union specialization. KernelManager (util/rtc.{h,cpp}) gains occupancy / function-attribute / cooperative-launch helpers needed by the norm launchers. Validated on sm_89 (RTX 6000 Ada): full normalization operator suite 192/192, softmax + NVRTC unit tests pass; libtransformer_engine.so shrinks ~72 MB -> ~65 MB. On sm_100a the NVRTC norm forward kernel builds where the static instantiation crashed the compiler. Signed-off-by: CarlosGomes98 --- tests/cpp/operator/CMakeLists.txt | 1 + tests/cpp/operator/test_softmax.cu | 234 ++++++ tests/cpp/util/test_nvrtc.cpp | 23 + transformer_engine/common/CMakeLists.txt | 47 ++ .../scaled_aligned_causal_masked_softmax.cu | 100 +++ .../fused_softmax/scaled_masked_softmax.cu | 128 ++- .../scaled_upper_triang_masked_softmax.cu | 106 ++- .../common/normalization/common.h | 104 +-- .../common/normalization/kernel_params.h | 88 ++ .../common/normalization/kernel_traits.h | 17 + .../layernorm/ln_bwd_kernels.cuh | 5 + .../layernorm/ln_bwd_semi_cuda_kernel.cu | 38 + .../layernorm/ln_fwd_cuda_kernel.cu | 42 + .../layernorm/ln_fwd_kernels.cuh | 7 + .../layernorm/rtc/ln_bwd_kernel.cu | 8 + .../layernorm/rtc/ln_fwd_kernel.cu | 12 + .../rmsnorm/rmsnorm_bwd_kernels.cuh | 26 +- .../rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu | 73 ++ .../rmsnorm/rmsnorm_fwd_cuda_kernel.cu | 38 + .../rmsnorm/rmsnorm_fwd_kernels.cuh | 7 + .../rmsnorm/rtc/rmsnorm_bwd_kernel.cu | 8 + .../rmsnorm/rtc/rmsnorm_fwd_kernel.cu | 8 + .../common/normalization/rtc_dispatch.cpp | 751 ++++++++++++++++++ .../common/normalization/rtc_dispatch.h | 54 ++ transformer_engine/common/util/rtc.cpp | 64 +- transformer_engine/common/util/rtc.h | 44 +- transformer_engine/common/utils.cuh | 23 + 27 files changed, 1952 insertions(+), 104 deletions(-) create mode 100644 tests/cpp/operator/test_softmax.cu create mode 100644 transformer_engine/common/normalization/kernel_params.h create mode 100644 transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu create mode 100644 transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu create mode 100644 transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu create mode 100644 transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu create mode 100644 transformer_engine/common/normalization/rtc_dispatch.cpp create mode 100644 transformer_engine/common/normalization/rtc_dispatch.h 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..23d72fdd65 --- /dev/null +++ b/tests/cpp/operator/test_softmax.cu @@ -0,0 +1,234 @@ +/************************************************************************* + * Copyright (c) 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/test_nvrtc.cpp b/tests/cpp/util/test_nvrtc.cpp index d41084449e..385c60377b 100644 --- a/tests/cpp/util/test_nvrtc.cpp +++ b/tests/cpp/util/test_nvrtc.cpp @@ -38,6 +38,15 @@ __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"; // Make sure kernels are not available @@ -79,4 +88,18 @@ __global__ void my_kernel(uint32_t *data) { 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); } diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index edb8c5e109..a692e6ea32 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 + "Compile legacy static fused softmax template dispatch instead of requiring NVRTC fallback" + OFF) +target_compile_definitions(transformer_engine + PRIVATE + NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=$) + +# Default OFF: use NVRTC dispatch for LayerNorm/RMSNorm forward/backward. +# Set ON to fall back to the legacy static template instantiations (compiles +# all 556 REGISTER_NORM_LAUNCHER variants up front; only needed if NVRTC is +# disabled via NVTE_DISABLE_NVRTC=1 at runtime). +option(NVTE_BUILD_LEGACY_STATIC_NORM + "Compile legacy static norm REGISTER_NORM_LAUNCHER template dispatch" + 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..7bac152b68 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,6 +4,10 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include +#include "utils.cuh" +#else #include #include #include @@ -19,10 +23,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); @@ -345,6 +362,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("&", 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 +507,23 @@ 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"}); + } + rtc_manager.launch(kernel_label, grid_size, block_size, 0, stream, dst, src, 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 +536,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 +572,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 +601,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 +639,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 +666,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..0bb3d8aa99 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -4,6 +4,10 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include +#include "utils.cuh" +#else #include #include #include @@ -17,10 +21,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); @@ -144,7 +165,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 +188,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 +290,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 +320,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 +441,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("&", 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,6 +505,21 @@ 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); + 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"}); + } + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, 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 @@ -513,6 +585,10 @@ void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const in default: break; } +#else + throw_nvrtc_required("scaled softmax forward"); +#endif + } NVTE_CHECK_CUDA(cudaGetLastError()); } } @@ -546,6 +622,21 @@ 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); + 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"}); + } + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, mask, 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 @@ -626,6 +717,10 @@ void dispatch_scaled_masked_softmax_forward(output_t *dst, const input_t *src, c default: break; } +#else + throw_nvrtc_required("scaled masked softmax forward"); +#endif + } NVTE_CHECK_CUDA(cudaGetLastError()); } } @@ -658,6 +753,21 @@ 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); + 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 @@ -738,6 +848,10 @@ void dispatch_scaled_masked_softmax_backward(output_t *grad_input, const input_t default: break; } +#else + throw_nvrtc_required("scaled softmax backward"); +#endif + } NVTE_CHECK_CUDA(cudaGetLastError()); } } @@ -812,8 +926,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 +968,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..677ac196cf 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,6 +4,10 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include +#include "utils.cuh" +#else #include #include #include @@ -17,10 +21,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); @@ -166,13 +188,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 +218,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 +355,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("&", 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,6 +423,20 @@ 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); + 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"}); + } + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, 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 @@ -445,6 +517,10 @@ void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const in default: break; } +#else + throw_nvrtc_required("forward"); +#endif + } NVTE_CHECK_CUDA(cudaGetLastError()); } } @@ -482,6 +558,20 @@ 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); + 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 @@ -562,6 +652,10 @@ void dispatch_scaled_upper_triang_masked_softmax_backward(output_t *grad_input, default: break; } +#else + throw_nvrtc_required("backward"); +#endif + } NVTE_CHECK_CUDA(cudaGetLastError()); } } @@ -595,8 +689,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 +713,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/normalization/common.h b/transformer_engine/common/normalization/common.h index f5dce64193..14ec51fda5 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,96 +52,9 @@ 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; +// KernelParamsBase / ForwardKernelParams / BackwardKernelParams / +// BackwardAddKernelParams are defined in kernel_params.h, which is also +// included unchanged by the NVRTC kernel sources for the norm RTC migration. enum class NVTE_Norm_Backend { Te, Cudnn }; enum class NVTE_Norm_Stage { Forward, Backward, BackwardAdd }; @@ -191,6 +105,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..63914249d9 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 -- e.g. ::transformer_engine::normalization::fp16 -- are not +// otherwise visible. Mirror the definitions from common.h here for the RTC TU. +// 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 { diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh index c4b00b87c3..f4edfa540b 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 "utils.cuh" +#include "kernel_params.h" +#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..21d74f6a2c 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,10 +7,18 @@ #include "../../common.h" #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "ln_bwd_kernels.cuh" using namespace transformer_engine::normalization; +#define NVTE_NORM_DT_fp32 ::transformer_engine::DType::kFloat32 +#define NVTE_NORM_DT_fp16 ::transformer_engine::DType::kFloat16 +#define NVTE_NORM_DT_bf16 ::transformer_engine::DType::kBFloat16 +#define NVTE_NORM_DT_fp8e4m3 ::transformer_engine::DType::kFloat8E4M3 +#define NVTE_NORM_DT_fp8e5m2 ::transformer_engine::DType::kFloat8E5M2 +#define NVTE_NORM_DT(tok) NVTE_NORM_DT_##tok + template @@ -132,6 +140,7 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, NVTE_CHECK_CUDA(cudaGetLastError()); } +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -145,6 +154,35 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ 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_BWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__) +#define REGISTER_NORM_LAUNCHER_LN_BWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ + WARPS_M, WARPS_N, BL_MAIN, BL_FINAL) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL); \ + return 0; \ + })(); \ + } +#define REGISTER_NORM_LAUNCHER_LN_BWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BL_MAIN, BL_FINAL) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL); \ + return 0; \ + })(); \ + } +#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..609f778632 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu @@ -6,10 +6,19 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "ln_fwd_kernels.cuh" using namespace transformer_engine::normalization; +// Token → DType mapping used by the RTC variant of REGISTER_NORM_LAUNCHER. +#define NVTE_NORM_DT_fp32 ::transformer_engine::DType::kFloat32 +#define NVTE_NORM_DT_fp16 ::transformer_engine::DType::kFloat16 +#define NVTE_NORM_DT_bf16 ::transformer_engine::DType::kBFloat16 +#define NVTE_NORM_DT_fp8e4m3 ::transformer_engine::DType::kFloat8E4M3 +#define NVTE_NORM_DT_fp8e5m2 ::transformer_engine::DType::kFloat8E5M2 +#define NVTE_NORM_DT(tok) NVTE_NORM_DT_##tok + template @@ -101,6 +110,8 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, } } +#if NVTE_BUILD_LEGACY_STATIC_NORM +// Legacy: emit a uniquely-named launcher and register the static template instantiation. #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -114,6 +125,37 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +// NVRTC path: dispatch through token-pasting on LAUNCH_TYPE so we can route +// "tuned" and "general" registrations to the right rtc_norm helper. +#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__) +#define REGISTER_NORM_LAUNCHER_LN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ + WARPS_M, WARPS_N, BYTES_PER_LDG) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG); \ + return 0; \ + })(); \ + } +#define REGISTER_NORM_LAUNCHER_LN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BYTES_PER_LDG) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG); \ + return 0; \ + })(); \ + } +#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..af73799f75 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh @@ -7,11 +7,18 @@ #ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ +#ifndef __CUDACC_RTC__ #include #include +#endif +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#include "kernel_params.h" +#else #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { 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..3afbe8dbe3 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh @@ -7,10 +7,17 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ +#ifndef __CUDACC_RTC__ #include +#endif +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#include "kernel_params.h" +#else #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -20,17 +27,32 @@ struct maybe_not_t {}; template using maybe_t = std::conditional_t; -template +// 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..ac01b6c89a 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,10 +6,18 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "rmsnorm_bwd_kernels.cuh" using namespace transformer_engine::normalization; +#define NVTE_NORM_DT_fp32 ::transformer_engine::DType::kFloat32 +#define NVTE_NORM_DT_fp16 ::transformer_engine::DType::kFloat16 +#define NVTE_NORM_DT_bf16 ::transformer_engine::DType::kBFloat16 +#define NVTE_NORM_DT_fp8e4m3 ::transformer_engine::DType::kFloat8E4M3 +#define NVTE_NORM_DT_fp8e5m2 ::transformer_engine::DType::kFloat8E5M2 +#define NVTE_NORM_DT(tok) NVTE_NORM_DT_##tok + template @@ -133,6 +141,12 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para NVTE_CHECK_CUDA(cudaGetLastError()); } +// Map NORM_STAGE token onto with_add boolean. +#define NVTE_RMSN_BWD_ADD_Backward false +#define NVTE_RMSN_BWD_ADD_BackwardAdd true +#define NVTE_RMSN_BWD_ADD(stage) NVTE_RMSN_BWD_ADD_##stage + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -146,6 +160,65 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +// Two-level dispatch: stage (Backward / BackwardAdd) AND launch type (tuned / general). +#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__) + +#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) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, false); \ + 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) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true); \ + return 0; \ + })(); \ + } +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_general_Backward(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + WARPS_M, WARPS_N, BL_MAIN, BL_FINAL) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, false); \ + 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) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true); \ + return 0; \ + })(); \ + } +#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..a892e1a9e7 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu @@ -6,10 +6,18 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "rmsnorm_fwd_kernels.cuh" using namespace transformer_engine::normalization; +#define NVTE_NORM_DT_fp32 ::transformer_engine::DType::kFloat32 +#define NVTE_NORM_DT_fp16 ::transformer_engine::DType::kFloat16 +#define NVTE_NORM_DT_bf16 ::transformer_engine::DType::kBFloat16 +#define NVTE_NORM_DT_fp8e4m3 ::transformer_engine::DType::kFloat8E4M3 +#define NVTE_NORM_DT_fp8e5m2 ::transformer_engine::DType::kFloat8E5M2 +#define NVTE_NORM_DT(tok) NVTE_NORM_DT_##tok + template @@ -102,6 +110,7 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param } } +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -115,6 +124,35 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ 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__) +#define REGISTER_NORM_LAUNCHER_RMSN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ + WARPS_M, WARPS_N, BYTES_PER_LDG) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG); \ + return 0; \ + })(); \ + } +#define REGISTER_NORM_LAUNCHER_RMSN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BYTES_PER_LDG) \ + namespace { \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG); \ + return 0; \ + })(); \ + } +#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..141ac87449 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh @@ -7,11 +7,18 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ +#ifndef __CUDACC_RTC__ #include #include +#endif +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#include "kernel_params.h" +#else #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { 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..f7e0dfb994 --- /dev/null +++ b/transformer_engine/common/normalization/rtc_dispatch.cpp @@ -0,0 +1,751 @@ +/************************************************************************* + * 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. When NVTE_BUILD_LEGACY_STATIC_NORM=OFF, the per-config +// REGISTER_NORM_LAUNCHER macro expands to a call into one of the +// register_*_tuned/general functions defined here. Each registers a closure +// in the TeNormalizationRegistry that, on first invocation, compiles the +// matching template instantiation via NVRTC and from then on dispatches via +// the cached CUfunction. + +#include "rtc_dispatch.h" + +#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_layernorm_rtc_ln_fwd_kernel_cu.h" +#include "string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu.h" +#include "string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu.h" +#include "string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu.h" + +namespace transformer_engine { +namespace normalization { +namespace rtc_norm { + +namespace { + +// 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; +} + +// 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) { + 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]( + LaunchParams& launch_params, const bool configure_params) { + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + mgr.compile(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)); +} + +// Backward "general" launcher additionally launches a finalize kernel after +// the main kernel — host-side replicating the static path needs awareness of +// that. Same dispatch shape but with a second kernel symbol. +template +void register_bwd_general_launcher(const std::string& main_label, + const std::string& main_kernel_expr, + const std::string& finalize_label, + const std::string& finalize_kernel_expr, + const char* rtc_source, const char* filename, TupleKeyType key, + int threads_per_cta, int dynamic_smem_bytes, + int finalize_threads_per_cta, int finalize_ctas, + int hidden_size, int ctype_bytes, int warps_m) { + auto closure = [main_label, main_kernel_expr, finalize_label, finalize_kernel_expr, rtc_source, + filename, threads_per_cta, dynamic_smem_bytes, finalize_threads_per_cta, + finalize_ctas, hidden_size, ctype_bytes, warps_m]( + LaunchParams& launch_params, const bool configure_params) { + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + mgr.compile(main_label, main_kernel_expr, rtc_source, filename); + } + if (!mgr.is_compiled(finalize_label)) { + mgr.compile(finalize_label, finalize_kernel_expr, rtc_source, filename); + } + + // ctas_per_col + ctas_per_row for "general" backward is computed exactly + // like the static launcher: ceil_div on hidden_size and on rows, capped by + // SM count × occupancy. + 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, + dynamic_smem_bytes); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden_size); + ctas_per_col = std::min(ceil_div(rows, warps_m), 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 * warps_m * ctas_per_row * (ctype_bytes * 2) * 2; + } + launch_params.dgamma_part_bytes = ctas_per_col * cols * ctype_bytes; + return; + } + + if (dynamic_smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(main_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + dynamic_smem_bytes); + } + 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), + dynamic_smem_bytes, stream, launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + dynamic_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)); +} + +} // 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) { + 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); +} + +void register_ln_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, int wn, + int bl) { + 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]( + LaunchParams& launch_params, + const bool configure_params) { + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + mgr.compile(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) { + 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); +} + +void register_rmsnorm_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, + int wn, int bl) { + 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]( + LaunchParams& launch_params, + const bool configure_params) { + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + mgr.compile(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) { + 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]( + LaunchParams& launch_params, + const bool configure_params) { + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + mgr.compile(main_label, main_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, "ln_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + mgr.compile(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) { + 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](LaunchParams& launch_params, + const bool configure_params) { + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + mgr.compile(main_label, main_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, "ln_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + mgr.compile(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) { + 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]( + LaunchParams& launch_params, + const bool configure_params) { + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + mgr.compile(main_label, main_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + mgr.compile(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) { + 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](LaunchParams& launch_params, + const bool configure_params) { + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + mgr.compile(main_label, main_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + mgr.compile(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..2a659a97b0 --- /dev/null +++ b/transformer_engine/common/normalization/rtc_dispatch.h @@ -0,0 +1,54 @@ +/************************************************************************* + * 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 { + +// 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); + +// 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); + +// 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); + +// 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); + +// 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); +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); +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); +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); + +} // 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..bf8f633f2c 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -25,6 +25,12 @@ namespace { // Strings with headers for RTC kernels #include "string_code_util_math_h.h" #include "string_code_utils_cuh.h" +#include "string_code_normalization_kernel_params_h.h" +#include "string_code_normalization_kernel_traits_h.h" +#include "string_code_normalization_layernorm_ln_fwd_kernels_cuh.h" +#include "string_code_normalization_layernorm_ln_bwd_kernels_cuh.h" +#include "string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh.h" +#include "string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh.h" /*! \brief Latest compute capability that NVRTC supports * @@ -132,6 +138,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,7 +157,8 @@ KernelManager& KernelManager::instance() { } void KernelManager::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) { std::lock_guard lock_guard_(lock_); // Choose whether to compile to PTX or cubin @@ -160,6 +179,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,9 +187,27 @@ 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"}; + constexpr int num_headers = 8; + constexpr const char* headers[num_headers] = { + string_code_utils_cuh, + string_code_util_math_h, + string_code_normalization_kernel_params_h, + string_code_normalization_kernel_traits_h, + string_code_normalization_layernorm_ln_fwd_kernels_cuh, + string_code_normalization_layernorm_ln_bwd_kernels_cuh, + string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh, + string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh, + }; + constexpr const char* include_names[num_headers] = { + "utils.cuh", + "util/math.h", + "kernel_params.h", + "kernel_traits.h", + "ln_fwd_kernels.cuh", + "ln_bwd_kernels.cuh", + "rmsnorm_fwd_kernels.cuh", + "rmsnorm_bwd_kernels.cuh", + }; NVTE_CHECK_NVRTC(nvrtcCreateProgram(&program, code.c_str(), filename.c_str(), num_headers, headers, include_names)); NVTE_CHECK_NVRTC(nvrtcAddNameExpression(program, kernel_name.c_str())); @@ -254,6 +292,24 @@ void KernelManager::set_cache_config(const std::string& kernel_label, CUfunc_cac 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); + 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); + 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); return kernel_cache_.count(key) > 0; diff --git a/transformer_engine/common/util/rtc.h b/transformer_engine/common/util/rtc.h index 65faf7bcc2..f879849f57 100644 --- a/transformer_engine/common/util/rtc.h +++ b/transformer_engine/common/util/rtc.h @@ -79,6 +79,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_; @@ -115,7 +138,8 @@ class KernelManager { * primarily for debugging */ 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 = {}); /*! \brief Whether CUDA kernel has been compiled for CUDA device * @@ -157,6 +181,24 @@ 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); + 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_; diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 635c7a36d2..d79f155fdc 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -28,6 +28,29 @@ static_assert(sizeof(uint8_t) == 1); static_assert(sizeof(uint16_t) == 2); static_assert(sizeof(uint32_t) == 4); static_assert(sizeof(uint64_t) == 8); +// Minimal subset of used by the kernel headers. The full header +// is not reliably available under NVRTC, so define just what the norm/transpose +// device code references (std::is_same, std::conditional[_t]). +namespace std { +template +struct is_same { + static constexpr bool value = false; +}; +template +struct is_same { + static constexpr bool value = true; +}; +template +struct conditional { + using type = T; +}; +template +struct conditional { + using type = F; +}; +template +using conditional_t = typename conditional::type; +} // namespace std #endif //////////////////////////////////////////////////////////////////////////////////////////////////// From 5fa3573f04a60e822f248f4e7957fbe2296f8206 Mon Sep 17 00:00:00 2001 From: Carlos Gomes Date: Tue, 30 Jun 2026 11:11:04 +0200 Subject: [PATCH 2/9] Add fully qualified name to softmax kernels Signed-off-by: CarlosGomes98 --- .../fused_softmax/scaled_aligned_causal_masked_softmax.cu | 6 +++--- .../common/fused_softmax/scaled_masked_softmax.cu | 6 +++--- .../fused_softmax/scaled_upper_triang_masked_softmax.cu | 6 +++--- transformer_engine/common/normalization/common.h | 4 ---- transformer_engine/common/normalization/kernel_traits.h | 4 ++-- 5 files changed, 11 insertions(+), 15 deletions(-) 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 7bac152b68..65750b9e93 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 @@ -109,7 +109,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); @@ -386,8 +386,8 @@ std::string make_softmax_rtc_label(const char *direction, int log2_elements) { template std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { - return concat_strings("&", kernel_name, "<", TypeInfo::name, ",", TypeInfo::name, - ",float,", log2_elements, ">"); + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); } void throw_nvrtc_required(const char *direction) { diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu index 0bb3d8aa99..2fb81dd961 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -88,7 +88,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); @@ -465,8 +465,8 @@ std::string make_softmax_rtc_label(const char *variant, const char *direction, i template std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { - return concat_strings("&", kernel_name, "<", TypeInfo::name, ",", TypeInfo::name, - ",float,", log2_elements, ">"); + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); } void throw_nvrtc_required(const char *variant) { 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 677ac196cf..e637e7a945 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 @@ -112,7 +112,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); @@ -379,8 +379,8 @@ std::string make_softmax_rtc_label(const char *direction, int log2_elements) { template std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { - return concat_strings("&", kernel_name, "<", TypeInfo::name, ",", TypeInfo::name, - ",float,", log2_elements, ">"); + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); } void throw_nvrtc_required(const char *direction) { diff --git a/transformer_engine/common/normalization/common.h b/transformer_engine/common/normalization/common.h index 14ec51fda5..31a547f62c 100644 --- a/transformer_engine/common/normalization/common.h +++ b/transformer_engine/common/normalization/common.h @@ -52,10 +52,6 @@ struct LaunchParams { } }; -// KernelParamsBase / ForwardKernelParams / BackwardKernelParams / -// BackwardAddKernelParams are defined in kernel_params.h, which is also -// included unchanged by the NVRTC kernel sources for the norm RTC migration. - enum class NVTE_Norm_Backend { Te, Cudnn }; enum class NVTE_Norm_Stage { Forward, Backward, BackwardAdd }; diff --git a/transformer_engine/common/normalization/kernel_traits.h b/transformer_engine/common/normalization/kernel_traits.h index 63914249d9..598eccef5a 100644 --- a/transformer_engine/common/normalization/kernel_traits.h +++ b/transformer_engine/common/normalization/kernel_traits.h @@ -20,8 +20,8 @@ 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 -- e.g. ::transformer_engine::normalization::fp16 -- are not -// otherwise visible. Mirror the definitions from common.h here for the RTC TU. +// 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; From 918a6db25f18be5a8d12d4e68ff38848701b59ed Mon Sep 17 00:00:00 2001 From: CarlosGomes98 Date: Tue, 30 Jun 2026 11:58:41 +0200 Subject: [PATCH 3/9] Add static fallback option, fix softmax acc_t dtype Signed-off-by: CarlosGomes98 --- docs/envvars.rst | 2 +- transformer_engine/common/CMakeLists.txt | 12 +- .../scaled_aligned_causal_masked_softmax.cu | 3 +- .../fused_softmax/scaled_masked_softmax.cu | 8 +- .../scaled_upper_triang_masked_softmax.cu | 5 +- .../layernorm/ln_bwd_semi_cuda_kernel.cu | 54 +++--- .../layernorm/ln_fwd_cuda_kernel.cu | 58 +++---- .../rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu | 112 ++++++------ .../rmsnorm/rmsnorm_fwd_cuda_kernel.cu | 52 +++--- .../common/normalization/rtc_dispatch.cpp | 161 ++++++++---------- .../common/normalization/rtc_dispatch.h | 27 ++- transformer_engine/common/utils.cuh | 3 +- 12 files changed, 235 insertions(+), 262 deletions(-) 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/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index a692e6ea32..899eff7557 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -591,18 +591,18 @@ target_include_directories(transformer_engine PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/string_headers") option(NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX - "Compile legacy static fused softmax template dispatch instead of requiring NVRTC fallback" + "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: use NVRTC dispatch for LayerNorm/RMSNorm forward/backward. -# Set ON to fall back to the legacy static template instantiations (compiles -# all 556 REGISTER_NORM_LAUNCHER variants up front; only needed if NVRTC is -# disabled via NVTE_DISABLE_NVRTC=1 at runtime). +# 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 - "Compile legacy static norm REGISTER_NORM_LAUNCHER template dispatch" + "Also compile static norm kernels for NVTE_DISABLE_NVRTC fallback" OFF) target_compile_definitions(transformer_engine PRIVATE 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 65750b9e93..2d62ddbfb0 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 @@ -517,7 +517,8 @@ void dispatch_scaled_aligned_causal_masked_softmax_forward(output_t *dst, const make_softmax_rtc_code(log2_elements), kRtcSourceFile, {"--use_fast_math"}); } - rtc_manager.launch(kernel_label, grid_size, block_size, 0, stream, dst, src, scale, + 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; diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu index 2fb81dd961..39496b1662 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -516,8 +516,9 @@ void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const in make_softmax_rtc_code(log2_elements), kRtcSourceFile, {"--use_fast_math"}); } - rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, scale, batch_count, - key_seq_len); + 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 @@ -633,7 +634,8 @@ void dispatch_scaled_masked_softmax_forward(output_t *dst, const input_t *src, c make_softmax_rtc_code(log2_elements), kRtcSourceFile, {"--use_fast_math"}); } - rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, mask, scale, + 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 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 e637e7a945..18949a7287 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 @@ -433,8 +433,9 @@ void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const in make_softmax_rtc_code(log2_elements), kRtcSourceFile, {"--use_fast_math"}); } - rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, scale, batch_count, - softmax_elements_stride, softmax_elements); + 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 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 21d74f6a2c..bffd7fc7ea 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 @@ -140,6 +140,27 @@ 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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + 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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + 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, ...) \ @@ -150,38 +171,15 @@ 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__) -#define REGISTER_NORM_LAUNCHER_LN_BWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ - WARPS_M, WARPS_N, BL_MAIN, BL_FINAL) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL); \ - return 0; \ - })(); \ - } -#define REGISTER_NORM_LAUNCHER_LN_BWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ - WARPS_N, BL_MAIN, BL_FINAL) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL); \ - return 0; \ - })(); \ - } + 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: 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 609f778632..f7ced7419b 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu @@ -110,8 +110,31 @@ 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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + 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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK); \ + return 0; \ + })() + #if NVTE_BUILD_LEGACY_STATIC_NORM -// Legacy: emit a uniquely-named launcher and register the static template instantiation. #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -121,40 +144,15 @@ 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 -// NVRTC path: dispatch through token-pasting on LAUNCH_TYPE so we can route -// "tuned" and "general" registrations to the right rtc_norm helper. #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__) -#define REGISTER_NORM_LAUNCHER_LN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ - WARPS_M, WARPS_N, BYTES_PER_LDG) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG); \ - return 0; \ - })(); \ - } -#define REGISTER_NORM_LAUNCHER_LN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ - WARPS_N, BYTES_PER_LDG) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG); \ - return 0; \ - })(); \ - } + 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: 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 ac01b6c89a..cca13b48f8 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 @@ -141,10 +141,53 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para NVTE_CHECK_CUDA(cudaGetLastError()); } -// Map NORM_STAGE token onto with_add boolean. -#define NVTE_RMSN_BWD_ADD_Backward false -#define NVTE_RMSN_BWD_ADD_BackwardAdd true -#define NVTE_RMSN_BWD_ADD(stage) NVTE_RMSN_BWD_ADD_##stage +// 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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + 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) \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true, 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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + 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) \ + [[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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true, STATIC_FALLBACK); \ + return 0; \ + })() #if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ @@ -156,68 +199,15 @@ 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 -// Two-level dispatch: stage (Backward / BackwardAdd) AND launch type (tuned / general). #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__) - -#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) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, false); \ - 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) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true); \ - return 0; \ - })(); \ - } -#define REGISTER_NORM_LAUNCHER_RMSN_BWD_general_Backward(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ - WARPS_M, WARPS_N, BL_MAIN, BL_FINAL) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, false); \ - 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) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true); \ - return 0; \ - })(); \ - } + 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: 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 a892e1a9e7..d2c951da3a 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu @@ -110,6 +110,27 @@ 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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + 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( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ + 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, ...) \ @@ -120,38 +141,15 @@ 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__) -#define REGISTER_NORM_LAUNCHER_RMSN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ - WARPS_M, WARPS_N, BYTES_PER_LDG) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG); \ - return 0; \ - })(); \ - } -#define REGISTER_NORM_LAUNCHER_RMSN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ - WARPS_N, BYTES_PER_LDG) \ - namespace { \ - [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG); \ - return 0; \ - })(); \ - } + 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: diff --git a/transformer_engine/common/normalization/rtc_dispatch.cpp b/transformer_engine/common/normalization/rtc_dispatch.cpp index f7e0dfb994..78265d9337 100644 --- a/transformer_engine/common/normalization/rtc_dispatch.cpp +++ b/transformer_engine/common/normalization/rtc_dispatch.cpp @@ -5,12 +5,11 @@ ************************************************************************/ // NVRTC-backed registry registration for LayerNorm/RMSNorm forward + backward -// launchers. When NVTE_BUILD_LEGACY_STATIC_NORM=OFF, the per-config -// REGISTER_NORM_LAUNCHER macro expands to a call into one of the -// register_*_tuned/general functions defined here. Each registers a closure -// in the TeNormalizationRegistry that, on first invocation, compiles the -// matching template instantiation via NVRTC and from then on dispatches via -// the cached CUfunction. +// 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" @@ -123,6 +122,19 @@ int bwd_smem_bytes(DType ctype, int hidden_size, int ctas_per_row, int warps_m, 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. @@ -131,11 +143,15 @@ 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) { + 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]( + 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)) { mgr.compile(label, kernel_expr, rtc_source, filename); @@ -179,73 +195,6 @@ void register_launcher(const std::string& label, const std::string& kernel_expr, TeNormalizationRegistry::registerFunction(key, std::move(closure)); } -// Backward "general" launcher additionally launches a finalize kernel after -// the main kernel — host-side replicating the static path needs awareness of -// that. Same dispatch shape but with a second kernel symbol. -template -void register_bwd_general_launcher(const std::string& main_label, - const std::string& main_kernel_expr, - const std::string& finalize_label, - const std::string& finalize_kernel_expr, - const char* rtc_source, const char* filename, TupleKeyType key, - int threads_per_cta, int dynamic_smem_bytes, - int finalize_threads_per_cta, int finalize_ctas, - int hidden_size, int ctype_bytes, int warps_m) { - auto closure = [main_label, main_kernel_expr, finalize_label, finalize_kernel_expr, rtc_source, - filename, threads_per_cta, dynamic_smem_bytes, finalize_threads_per_cta, - finalize_ctas, hidden_size, ctype_bytes, warps_m]( - LaunchParams& launch_params, const bool configure_params) { - auto& mgr = rtc::KernelManager::instance(); - if (!mgr.is_compiled(main_label)) { - mgr.compile(main_label, main_kernel_expr, rtc_source, filename); - } - if (!mgr.is_compiled(finalize_label)) { - mgr.compile(finalize_label, finalize_kernel_expr, rtc_source, filename); - } - - // ctas_per_col + ctas_per_row for "general" backward is computed exactly - // like the static launcher: ceil_div on hidden_size and on rows, capped by - // SM count × occupancy. - 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, - dynamic_smem_bytes); - const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; - ctas_per_row = ceil_div(cols, hidden_size); - ctas_per_col = std::min(ceil_div(rows, warps_m), 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 * warps_m * ctas_per_row * (ctype_bytes * 2) * 2; - } - launch_params.dgamma_part_bytes = ctas_per_col * cols * ctype_bytes; - return; - } - - if (dynamic_smem_bytes >= 48 * 1024) { - mgr.set_function_attribute(main_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, - dynamic_smem_bytes); - } - 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), - dynamic_smem_bytes, stream, launch_params.params); - } else { - mgr.launch_cooperative(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), - dynamic_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)); -} - } // namespace // ============================================================================ @@ -253,7 +202,7 @@ void register_bwd_general_launcher(const std::string& main_label, // ============================================================================ void register_ln_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, - int wn, int bl) { + 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 = @@ -276,11 +225,11 @@ void register_ln_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, i 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); + 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) { + 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 = @@ -293,9 +242,12 @@ void register_ln_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, 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]( + 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)) { mgr.compile(label, kernel_expr, string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu, @@ -339,7 +291,8 @@ void register_ln_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, // ============================================================================ void register_rmsnorm_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, - int wn, int bl) { + 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 = @@ -357,11 +310,12 @@ void register_rmsnorm_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd 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); + 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) { + 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 = @@ -372,9 +326,12 @@ void register_rmsnorm_fwd_general(DType wt, DType it, DType ot, DType ct, int hi 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]( + 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)) { mgr.compile(label, kernel_expr, @@ -418,7 +375,8 @@ void register_rmsnorm_fwd_general(DType wt, DType it, DType ot, DType ct, int hi // ============================================================================ 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) { + 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 = @@ -458,9 +416,12 @@ void register_ln_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, i 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]( + 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)) { mgr.compile(main_label, main_kexpr, @@ -503,7 +464,8 @@ void register_ln_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, i } 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) { + 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 = @@ -531,8 +493,12 @@ void register_ln_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, 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](LaunchParams& launch_params, - const bool configure_params) { + 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)) { mgr.compile(main_label, main_kexpr, @@ -582,7 +548,8 @@ void register_ln_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, // ============================================================================ 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) { + 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, @@ -620,9 +587,12 @@ void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd 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]( + 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)) { mgr.compile(main_label, main_kexpr, @@ -667,7 +637,8 @@ void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd } 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) { + 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, @@ -698,8 +669,12 @@ void register_rmsnorm_bwd_general(DType wt, DType it, DType ot, DType ct, int hi 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](LaunchParams& launch_params, - const bool configure_params) { + 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)) { mgr.compile(main_label, main_kexpr, diff --git a/transformer_engine/common/normalization/rtc_dispatch.h b/transformer_engine/common/normalization/rtc_dispatch.h index 2a659a97b0..314f150f70 100644 --- a/transformer_engine/common/normalization/rtc_dispatch.h +++ b/transformer_engine/common/normalization/rtc_dispatch.h @@ -15,37 +15,48 @@ 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); + 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); + 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); + 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); + 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); + 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); + 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); + 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); + int bytes_per_ldg_final, bool with_add, + StaticFallback static_fallback = nullptr); } // namespace rtc_norm } // namespace normalization diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index d79f155fdc..bd5ea59be5 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -29,8 +29,7 @@ static_assert(sizeof(uint16_t) == 2); static_assert(sizeof(uint32_t) == 4); static_assert(sizeof(uint64_t) == 8); // Minimal subset of used by the kernel headers. The full header -// is not reliably available under NVRTC, so define just what the norm/transpose -// device code references (std::is_same, std::conditional[_t]). +// is not reliably available under NVRTC namespace std { template struct is_same { From 52357238613f2b415b571c8e041d9688e0205648 Mon Sep 17 00:00:00 2001 From: CarlosGomes98 Date: Tue, 30 Jun 2026 14:59:59 +0200 Subject: [PATCH 4/9] add missing license Signed-off-by: CarlosGomes98 --- tests/cpp/operator/test_softmax.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cpp/operator/test_softmax.cu b/tests/cpp/operator/test_softmax.cu index 23d72fdd65..4fc5a294c8 100644 --- a/tests/cpp/operator/test_softmax.cu +++ b/tests/cpp/operator/test_softmax.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ From 476dc9a544c72a01cdde572288e36e3529ecf18e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 30 Jun 2026 13:01:57 +0000 Subject: [PATCH 5/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../scaled_aligned_causal_masked_softmax.cu | 9 +- .../fused_softmax/scaled_masked_softmax.cu | 464 +++++++++--------- .../scaled_upper_triang_masked_softmax.cu | 325 ++++++------ .../common/normalization/kernel_traits.h | 2 +- .../layernorm/ln_bwd_kernels.cuh | 2 +- .../layernorm/ln_bwd_semi_cuda_kernel.cu | 46 +- .../layernorm/ln_fwd_cuda_kernel.cu | 46 +- .../layernorm/ln_fwd_kernels.cuh | 2 +- .../rmsnorm/rmsnorm_bwd_kernels.cuh | 2 +- .../rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu | 87 ++-- .../rmsnorm/rmsnorm_fwd_cuda_kernel.cu | 47 +- .../rmsnorm/rmsnorm_fwd_kernels.cuh | 2 +- .../common/normalization/rtc_dispatch.cpp | 127 +++-- transformer_engine/common/util/rtc.cpp | 12 +- transformer_engine/common/util/rtc.h | 8 +- 15 files changed, 590 insertions(+), 591 deletions(-) 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 2d62ddbfb0..0f8d5e8a9a 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 @@ -6,6 +6,7 @@ #ifdef __CUDACC_RTC__ #include + #include "utils.cuh" #else #include @@ -404,8 +405,8 @@ void call_kernel_scaled_aligned_causal_masked_softmax_forward( const input_t *src, const acc_t scale, const int microbatches, const int query_seq_len, const int key_seq_len) { scaled_aligned_causal_masked_softmax_warp_forward - <<>>(dst, src, scale, microbatches, query_seq_len, - key_seq_len); + << > >(dst, src, scale, microbatches, + query_seq_len, key_seq_len); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -415,8 +416,8 @@ void call_kernel_scaled_aligned_causal_masked_softmax_backward( const input_t *grad, const input_t *output, const acc_t scale, const int microbatches, const int query_seq_len, const int key_seq_len) { scaled_aligned_causal_masked_softmax_warp_backward - <<>>(gradInput, grad, output, scale, microbatches, - query_seq_len, key_seq_len); + << > >(gradInput, grad, output, scale, microbatches, + query_seq_len, key_seq_len); NVTE_CHECK_CUDA(cudaGetLastError()); } diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu index 39496b1662..4acb4ecb3e 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -6,6 +6,7 @@ #ifdef __CUDACC_RTC__ #include + #include "utils.cuh" #else #include @@ -510,82 +511,81 @@ void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const in 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"}); + 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); + 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; - } + // 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 @@ -639,86 +639,86 @@ void dispatch_scaled_masked_softmax_forward(output_t *dst, const input_t *src, c 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; - } + // 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 @@ -770,86 +770,86 @@ void dispatch_scaled_masked_softmax_backward(output_t *grad_input, const input_t 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; - } + // 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 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 18949a7287..3b5502f067 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 @@ -6,6 +6,7 @@ #ifdef __CUDACC_RTC__ #include + #include "utils.cuh" #else #include @@ -434,90 +435,90 @@ void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const in {"--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); + 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; - } + // 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 @@ -573,86 +574,86 @@ void dispatch_scaled_upper_triang_masked_softmax_backward(output_t *grad_input, 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; - } + // 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 diff --git a/transformer_engine/common/normalization/kernel_traits.h b/transformer_engine/common/normalization/kernel_traits.h index 598eccef5a..dd30c1c22a 100644 --- a/transformer_engine/common/normalization/kernel_traits.h +++ b/transformer_engine/common/normalization/kernel_traits.h @@ -20,7 +20,7 @@ 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). +// 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; diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh index f4edfa540b..0480b588b9 100644 --- a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh @@ -8,8 +8,8 @@ #define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_BWD_KERNELS_CUH_ #ifdef __CUDACC_RTC__ -#include "utils.cuh" #include "kernel_params.h" +#include "utils.cuh" #else #include "../../utils.cuh" #include "../common.h" 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 bffd7fc7ea..ade9a6e535 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 @@ -140,25 +140,27 @@ 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 \ +#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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_bwd_tuned( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), 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 \ +#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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_bwd_general( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + STATIC_FALLBACK); \ + return 0; \ })() #if NVTE_BUILD_LEGACY_STATIC_NORM @@ -171,15 +173,15 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, launch_ln_bwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_LAUNCHER_LN_BWD_##LAUNCH_TYPE( \ - HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ + 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 #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) +#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: 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 f7ced7419b..c4d0af67b0 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu @@ -113,25 +113,27 @@ 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 \ +#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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_fwd_tuned( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), 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 \ +#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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_fwd_general( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK); \ + return 0; \ })() #if NVTE_BUILD_LEGACY_STATIC_NORM @@ -144,15 +146,15 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, launch_ln_fwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_LAUNCHER_LN_FWD_##LAUNCH_TYPE( \ - HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ + 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) +#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: diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh index af73799f75..2fdeef7d1b 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh @@ -13,8 +13,8 @@ #endif #ifdef __CUDACC_RTC__ -#include "utils.cuh" #include "kernel_params.h" +#include "utils.cuh" #else #include "../../utils.cuh" #include "../common.h" diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh index 3afbe8dbe3..5966024e23 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh @@ -12,8 +12,8 @@ #endif #ifdef __CUDACC_RTC__ -#include "utils.cuh" #include "kernel_params.h" +#include "utils.cuh" #else #include "../../utils.cuh" #include "../common.h" 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 cca13b48f8..d199a06e9a 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 @@ -144,49 +144,52 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para // 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 \ +#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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, false, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_tuned( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), 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) \ - [[maybe_unused]] static const int \ +#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) \ + [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_tuned( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, true, 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 \ +#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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, false, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_general( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), 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) \ - [[maybe_unused]] static const int \ +#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) \ + [[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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_general( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true, \ + STATIC_FALLBACK); \ + return 0; \ })() #if NVTE_BUILD_LEGACY_STATIC_NORM @@ -199,15 +202,15 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para launch_rmsnorm_bwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_LAUNCHER_RMSN_BWD_##LAUNCH_TYPE##_##NORM_STAGE( \ - HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ + 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) +#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: 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 d2c951da3a..ab7a24aac5 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu @@ -110,25 +110,28 @@ 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 \ +#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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_fwd_tuned( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), 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 \ +#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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), NVTE_NORM_DT(CTYPE), \ - HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK); \ - return 0; \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_fwd_general( \ + NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ + NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK); \ + return 0; \ })() #if NVTE_BUILD_LEGACY_STATIC_NORM @@ -141,15 +144,15 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param launch_rmsnorm_fwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_LAUNCHER_RMSN_FWD_##LAUNCH_TYPE( \ - HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ + 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) +#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: diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh index 141ac87449..4d83352089 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh @@ -13,8 +13,8 @@ #endif #ifdef __CUDACC_RTC__ -#include "utils.cuh" #include "kernel_params.h" +#include "utils.cuh" #else #include "../../utils.cuh" #include "../common.h" diff --git a/transformer_engine/common/normalization/rtc_dispatch.cpp b/transformer_engine/common/normalization/rtc_dispatch.cpp index 78265d9337..4bef2fd722 100644 --- a/transformer_engine/common/normalization/rtc_dispatch.cpp +++ b/transformer_engine/common/normalization/rtc_dispatch.cpp @@ -25,10 +25,10 @@ // 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_layernorm_rtc_ln_fwd_kernel_cu.h" #include "string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu.h" -#include "string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu.h" +#include "string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu.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 { @@ -117,8 +117,7 @@ int reducer_smem_bytes(DType ctype, int warps_m, int warps_n) { // 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); + const int wgrad = (ctas_per_row > 1) ? 0 : warps_m * hidden_size * byte_size_of(ctype); return dgrad + wgrad; } @@ -146,9 +145,9 @@ void register_launcher(const std::string& label, const std::string& kernel_expr, 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) { + 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; } @@ -164,10 +163,8 @@ void register_launcher(const std::string& label, const std::string& kernel_expr, 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; + 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 = @@ -210,16 +207,14 @@ void register_ln_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, i ",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 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 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( @@ -238,9 +233,8 @@ void register_ln_fwd_general(DType wt, DType it, DType ot, DType ct, int 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 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, @@ -293,8 +287,8 @@ void register_ln_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, 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 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, @@ -316,8 +310,8 @@ void register_rmsnorm_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd 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 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, @@ -334,8 +328,7 @@ void register_rmsnorm_fwd_general(DType wt, DType it, DType ot, DType ct, int hi } auto& mgr = rtc::KernelManager::instance(); if (!mgr.is_compiled(label)) { - mgr.compile(label, kernel_expr, - string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu, + mgr.compile(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; }; @@ -344,8 +337,7 @@ void register_rmsnorm_fwd_general(DType wt, DType it, DType ot, DType ct, int hi 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 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)); @@ -416,16 +408,16 @@ void register_ln_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, i 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) { + 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)) { - mgr.compile(main_label, main_kexpr, - string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, "ln_bwd_kernel.cu"); + mgr.compile(main_label, main_kexpr, string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); } if (!mgr.is_compiled(finalize_label)) { mgr.compile(finalize_label, finalize_kexpr, @@ -435,8 +427,7 @@ void register_ln_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, i 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; + 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; @@ -454,8 +445,8 @@ void register_ln_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, i 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_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); @@ -475,8 +466,8 @@ void register_ln_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, 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, ">"); + 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 = @@ -493,16 +484,16 @@ void register_ln_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, 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) { + 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)) { - mgr.compile(main_label, main_kexpr, - string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, "ln_bwd_kernel.cu"); + mgr.compile(main_label, main_kexpr, string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); } if (!mgr.is_compiled(finalize_label)) { mgr.compile(finalize_label, finalize_kexpr, @@ -550,17 +541,14 @@ void register_ln_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, 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 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 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); @@ -572,9 +560,9 @@ void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd 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 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; @@ -587,9 +575,9 @@ void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd 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) { + 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; } @@ -608,8 +596,7 @@ void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd 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; + 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; @@ -627,8 +614,8 @@ void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd 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_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); @@ -639,16 +626,14 @@ void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd 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 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 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); @@ -669,9 +654,9 @@ void register_rmsnorm_bwd_general(DType wt, DType it, DType ot, DType ct, int hi 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) { + finalize_threads_per_warp, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { if (try_static_fallback(static_fallback, launch_params, configure_params)) { return; } diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index bf8f633f2c..7a4eca2e12 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -23,14 +23,14 @@ namespace rtc { namespace { // Strings with headers for RTC kernels -#include "string_code_util_math_h.h" -#include "string_code_utils_cuh.h" #include "string_code_normalization_kernel_params_h.h" #include "string_code_normalization_kernel_traits_h.h" -#include "string_code_normalization_layernorm_ln_fwd_kernels_cuh.h" #include "string_code_normalization_layernorm_ln_bwd_kernels_cuh.h" -#include "string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh.h" +#include "string_code_normalization_layernorm_ln_fwd_kernels_cuh.h" #include "string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh.h" +#include "string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh.h" +#include "string_code_util_math_h.h" +#include "string_code_utils_cuh.h" /*! \brief Latest compute capability that NVRTC supports * @@ -301,8 +301,8 @@ void KernelManager::set_function_attribute(const std::string& kernel_label, } int KernelManager::occupancy_max_active_blocks_per_sm(const std::string& kernel_label, - int block_size, - std::size_t dynamic_smem_bytes) { + 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); NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to query occupancy before compilation"); diff --git a/transformer_engine/common/util/rtc.h b/transformer_engine/common/util/rtc.h index f879849f57..67754cf452 100644 --- a/transformer_engine/common/util/rtc.h +++ b/transformer_engine/common/util/rtc.h @@ -182,7 +182,8 @@ 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); + 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, @@ -190,8 +191,9 @@ class KernelManager { /*! \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) { + 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); NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to launch RTC kernel before compilation"); From b909589a708f1ced2771aa62e5065ddd645ddbb8 Mon Sep 17 00:00:00 2001 From: CarlosGomes98 Date: Wed, 1 Jul 2026 15:53:16 +0200 Subject: [PATCH 6/9] greptile changes Signed-off-by: CarlosGomes98 --- tests/cpp/util/CMakeLists.txt | 10 +- tests/cpp/util/test_nvrtc.cpp | 146 ++++++++++++++---- .../layernorm/ln_bwd_semi_cuda_kernel.cu | 15 +- .../layernorm/ln_fwd_cuda_kernel.cu | 18 +-- .../layernorm/ln_fwd_kernels.cuh | 4 +- .../rmsnorm/rmsnorm_bwd_kernels.cuh | 6 +- .../rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu | 31 ++-- .../rmsnorm/rmsnorm_fwd_cuda_kernel.cu | 17 +- .../rmsnorm/rmsnorm_fwd_kernels.cuh | 4 +- .../common/normalization/rtc_dispatch.cpp | 81 +++++++--- transformer_engine/common/util/rtc.cpp | 65 ++++---- transformer_engine/common/util/rtc.h | 18 ++- transformer_engine/common/utils.cuh | 27 ++-- 13 files changed, 272 insertions(+), 170 deletions(-) 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 385c60377b..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( @@ -47,59 +50,142 @@ __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_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), + 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/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu b/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu index ade9a6e535..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 @@ -12,13 +12,6 @@ using namespace transformer_engine::normalization; -#define NVTE_NORM_DT_fp32 ::transformer_engine::DType::kFloat32 -#define NVTE_NORM_DT_fp16 ::transformer_engine::DType::kFloat16 -#define NVTE_NORM_DT_bf16 ::transformer_engine::DType::kBFloat16 -#define NVTE_NORM_DT_fp8e4m3 ::transformer_engine::DType::kFloat8E4M3 -#define NVTE_NORM_DT_fp8e5m2 ::transformer_engine::DType::kFloat8E5M2 -#define NVTE_NORM_DT(tok) NVTE_NORM_DT_##tok - template @@ -146,8 +139,8 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, _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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + 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; \ })() @@ -157,8 +150,8 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, _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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ STATIC_FALLBACK); \ return 0; \ })() 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 c4d0af67b0..5be13357b1 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu @@ -11,14 +11,6 @@ using namespace transformer_engine::normalization; -// Token → DType mapping used by the RTC variant of REGISTER_NORM_LAUNCHER. -#define NVTE_NORM_DT_fp32 ::transformer_engine::DType::kFloat32 -#define NVTE_NORM_DT_fp16 ::transformer_engine::DType::kFloat16 -#define NVTE_NORM_DT_bf16 ::transformer_engine::DType::kBFloat16 -#define NVTE_NORM_DT_fp8e4m3 ::transformer_engine::DType::kFloat8E4M3 -#define NVTE_NORM_DT_fp8e5m2 ::transformer_engine::DType::kFloat8E5M2 -#define NVTE_NORM_DT(tok) NVTE_NORM_DT_##tok - template @@ -119,9 +111,9 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, _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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, \ - STATIC_FALLBACK); \ + 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, \ @@ -130,8 +122,8 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, _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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ STATIC_FALLBACK); \ return 0; \ })() diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh index 2fdeef7d1b..8724b443d5 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh @@ -154,7 +154,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); } } @@ -330,7 +330,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/rmsnorm/rmsnorm_bwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh index 5966024e23..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,6 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ -#ifndef __CUDACC_RTC__ -#include -#endif - #ifdef __CUDACC_RTC__ #include "kernel_params.h" #include "utils.cuh" @@ -25,7 +21,7 @@ namespace normalization { struct maybe_not_t {}; template -using maybe_t = std::conditional_t; +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 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 d199a06e9a..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 @@ -11,13 +11,6 @@ using namespace transformer_engine::normalization; -#define NVTE_NORM_DT_fp32 ::transformer_engine::DType::kFloat32 -#define NVTE_NORM_DT_fp16 ::transformer_engine::DType::kFloat16 -#define NVTE_NORM_DT_bf16 ::transformer_engine::DType::kBFloat16 -#define NVTE_NORM_DT_fp8e4m3 ::transformer_engine::DType::kFloat8E4M3 -#define NVTE_NORM_DT_fp8e5m2 ::transformer_engine::DType::kFloat8E5M2 -#define NVTE_NORM_DT(tok) NVTE_NORM_DT_##tok - template @@ -151,21 +144,22 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para _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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + 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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ - BL_FINAL, true, STATIC_FALLBACK); \ + 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( \ @@ -174,21 +168,22 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para _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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, false, \ - STATIC_FALLBACK); \ + 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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, true, \ - STATIC_FALLBACK); \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + ADD_FLAG, STATIC_FALLBACK); \ return 0; \ })() 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 ab7a24aac5..c3308cbce0 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu @@ -11,13 +11,6 @@ using namespace transformer_engine::normalization; -#define NVTE_NORM_DT_fp32 ::transformer_engine::DType::kFloat32 -#define NVTE_NORM_DT_fp16 ::transformer_engine::DType::kFloat16 -#define NVTE_NORM_DT_bf16 ::transformer_engine::DType::kBFloat16 -#define NVTE_NORM_DT_fp8e4m3 ::transformer_engine::DType::kFloat8E4M3 -#define NVTE_NORM_DT_fp8e5m2 ::transformer_engine::DType::kFloat8E5M2 -#define NVTE_NORM_DT(tok) NVTE_NORM_DT_##tok - template @@ -117,9 +110,9 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param _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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, \ - STATIC_FALLBACK); \ + 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, \ @@ -128,8 +121,8 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param _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( \ - NVTE_NORM_DT(WTYPE), NVTE_NORM_DT(ITYPE), NVTE_NORM_DT(OTYPE), \ - NVTE_NORM_DT(CTYPE), HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ STATIC_FALLBACK); \ return 0; \ })() diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh index 4d83352089..0f5341cffb 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh @@ -146,7 +146,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); } } @@ -305,7 +305,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/rtc_dispatch.cpp b/transformer_engine/common/normalization/rtc_dispatch.cpp index 4bef2fd722..5ae81f4233 100644 --- a/transformer_engine/common/normalization/rtc_dispatch.cpp +++ b/transformer_engine/common/normalization/rtc_dispatch.cpp @@ -15,6 +15,7 @@ #include #include +#include #include "../util/cuda_driver.h" #include "../util/rtc.h" @@ -25,8 +26,14 @@ // 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" @@ -36,6 +43,24 @@ 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) { @@ -153,7 +178,7 @@ void register_launcher(const std::string& label, const std::string& kernel_expr, } auto& mgr = rtc::KernelManager::instance(); if (!mgr.is_compiled(label)) { - mgr.compile(label, kernel_expr, rtc_source, filename); + compile_norm_kernel(mgr, label, kernel_expr, rtc_source, filename); } if (configure_params) { @@ -244,8 +269,9 @@ void register_ln_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, } auto& mgr = rtc::KernelManager::instance(); if (!mgr.is_compiled(label)) { - mgr.compile(label, kernel_expr, string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu, - "ln_fwd_kernel.cu"); + 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; @@ -328,8 +354,9 @@ void register_rmsnorm_fwd_general(DType wt, DType it, DType ot, DType ct, int hi } auto& mgr = rtc::KernelManager::instance(); if (!mgr.is_compiled(label)) { - mgr.compile(label, kernel_expr, string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu, - "rmsnorm_fwd_kernel.cu"); + 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; @@ -416,12 +443,14 @@ void register_ln_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, i } auto& mgr = rtc::KernelManager::instance(); if (!mgr.is_compiled(main_label)) { - mgr.compile(main_label, main_kexpr, string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, - "ln_bwd_kernel.cu"); + 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)) { - mgr.compile(finalize_label, finalize_kexpr, - string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, "ln_bwd_kernel.cu"); + 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 = @@ -492,12 +521,14 @@ void register_ln_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, } auto& mgr = rtc::KernelManager::instance(); if (!mgr.is_compiled(main_label)) { - mgr.compile(main_label, main_kexpr, string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, - "ln_bwd_kernel.cu"); + 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)) { - mgr.compile(finalize_label, finalize_kexpr, - string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, "ln_bwd_kernel.cu"); + 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; @@ -583,14 +614,14 @@ void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidd } auto& mgr = rtc::KernelManager::instance(); if (!mgr.is_compiled(main_label)) { - mgr.compile(main_label, main_kexpr, - string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, - "rmsnorm_bwd_kernel.cu"); + 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)) { - mgr.compile(finalize_label, finalize_kexpr, - string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, - "rmsnorm_bwd_kernel.cu"); + 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 = @@ -662,14 +693,14 @@ void register_rmsnorm_bwd_general(DType wt, DType it, DType ot, DType ct, int hi } auto& mgr = rtc::KernelManager::instance(); if (!mgr.is_compiled(main_label)) { - mgr.compile(main_label, main_kexpr, - string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, - "rmsnorm_bwd_kernel.cu"); + 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)) { - mgr.compile(finalize_label, finalize_kexpr, - string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, - "rmsnorm_bwd_kernel.cu"); + 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; diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index 7a4eca2e12..20616f8cb6 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -23,12 +23,6 @@ namespace rtc { namespace { // Strings with headers for RTC kernels -#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_rmsnorm_rmsnorm_bwd_kernels_cuh.h" -#include "string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh.h" #include "string_code_util_math_h.h" #include "string_code_utils_cuh.h" @@ -53,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_; } @@ -158,11 +147,16 @@ KernelManager& KernelManager::instance() { void KernelManager::compile(const std::string& kernel_label, const std::string& kernel_name, const std::string& code, const std::string& filename, - const std::vector& extra_options) { - std::lock_guard lock_guard_(lock_); + 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; @@ -187,29 +181,19 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& // Compile source nvrtcProgram program; - constexpr int num_headers = 8; - constexpr const char* headers[num_headers] = { - string_code_utils_cuh, - string_code_util_math_h, - string_code_normalization_kernel_params_h, - string_code_normalization_kernel_traits_h, - string_code_normalization_layernorm_ln_fwd_kernels_cuh, - string_code_normalization_layernorm_ln_bwd_kernels_cuh, - string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh, - string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh, - }; - constexpr const char* include_names[num_headers] = { - "utils.cuh", - "util/math.h", - "kernel_params.h", - "kernel_traits.h", - "ln_fwd_kernels.cuh", - "ln_bwd_kernels.cuh", - "rmsnorm_fwd_kernels.cuh", - "rmsnorm_bwd_kernels.cuh", - }; - 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()); @@ -277,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 @@ -288,6 +271,7 @@ 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); } @@ -296,6 +280,7 @@ 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); } @@ -305,6 +290,7 @@ int KernelManager::occupancy_max_active_blocks_per_sm(const std::string& kernel_ 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); @@ -312,6 +298,7 @@ int KernelManager::occupancy_max_active_blocks_per_sm(const std::string& kernel_ 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 67754cf452..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: @@ -136,10 +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::vector &extra_options = {}); + const std::vector &extra_options = {}, + const std::vector
&extra_headers = {}); /*! \brief Whether CUDA kernel has been compiled for CUDA device * @@ -167,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)...); @@ -196,6 +207,7 @@ class KernelManager { 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)...); @@ -204,8 +216,8 @@ class KernelManager { 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 bd5ea59be5..dd7dcdc4cc 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -28,9 +28,13 @@ static_assert(sizeof(uint8_t) == 1); static_assert(sizeof(uint16_t) == 2); static_assert(sizeof(uint32_t) == 4); static_assert(sizeof(uint64_t) == 8); -// Minimal subset of used by the kernel headers. The full header -// is not reliably available under NVRTC -namespace std { +#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; @@ -39,6 +43,10 @@ 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; @@ -49,8 +57,9 @@ struct conditional { }; template using conditional_t = typename conditional::type; -} // namespace std -#endif + +} // namespace rtc_detail +} // namespace transformer_engine //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -962,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); @@ -977,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); From 7157e36aa3a3525dfcffbd22d9bc838850a50aff Mon Sep 17 00:00:00 2001 From: CarlosGomes98 Date: Wed, 1 Jul 2026 16:24:07 +0200 Subject: [PATCH 7/9] fix formatting, .clang-format Signed-off-by: CarlosGomes98 --- .clang-format | 3 +- .../scaled_aligned_causal_masked_softmax.cu | 8 +- .../fused_softmax/scaled_masked_softmax.cu | 150 +++++++++--------- .../scaled_upper_triang_masked_softmax.cu | 120 +++++++------- 4 files changed, 140 insertions(+), 141 deletions(-) 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/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu index 0f8d5e8a9a..9d2756ec28 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 @@ -405,8 +405,8 @@ void call_kernel_scaled_aligned_causal_masked_softmax_forward( const input_t *src, const acc_t scale, const int microbatches, const int query_seq_len, const int key_seq_len) { scaled_aligned_causal_masked_softmax_warp_forward - << > >(dst, src, scale, microbatches, - query_seq_len, key_seq_len); + <<>>(dst, src, scale, microbatches, query_seq_len, + key_seq_len); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -416,8 +416,8 @@ void call_kernel_scaled_aligned_causal_masked_softmax_backward( const input_t *grad, const input_t *output, const acc_t scale, const int microbatches, const int query_seq_len, const int key_seq_len) { scaled_aligned_causal_masked_softmax_warp_backward - << > >(gradInput, grad, output, scale, microbatches, - query_seq_len, key_seq_len); + <<>>(gradInput, grad, output, scale, microbatches, + query_seq_len, key_seq_len); NVTE_CHECK_CUDA(cudaGetLastError()); } diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu index 4acb4ecb3e..d3029e5a4c 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -525,63 +525,63 @@ void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const in switch (log2_elements) { case 0: // 1 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 1: // 2 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 2: // 4 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 3: // 8 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 4: // 16 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 5: // 32 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 6: // 64 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 7: // 128 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 8: // 256 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 9: // 512 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 10: // 1024 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 11: // 2048 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 12: // 4096 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 13: // 8192 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; case 14: // 16384 scaled_softmax_warp_forward - << > >(dst, src, scale, batch_count, key_seq_len); + <<>>(dst, src, scale, batch_count, key_seq_len); break; default: break; @@ -643,78 +643,78 @@ void dispatch_scaled_masked_softmax_forward(output_t *dst, const input_t *src, c switch (log2_elements) { case 0: // 1 scaled_masked_softmax_warp_forward - << > >(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); break; default: break; @@ -774,78 +774,78 @@ void dispatch_scaled_masked_softmax_backward(output_t *grad_input, const input_t switch (log2_elements) { case 0: // 1 scaled_masked_softmax_warp_backward - << > >(grad_input, grad, output, scale, batch_count, - key_seq_len); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); break; default: break; 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 3b5502f067..7c30c38850 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 @@ -443,78 +443,78 @@ void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const in switch (log2_elements) { case 0: // 1 scaled_upper_triang_masked_softmax_warp_forward - << > >(dst, src, scale, batch_count, - softmax_elements_stride, softmax_elements); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); break; default: break; @@ -578,78 +578,78 @@ void dispatch_scaled_upper_triang_masked_softmax_backward(output_t *grad_input, 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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(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); + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); break; default: break; From 329945fee6113e5071d3c6ac93584a28580c7e90 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 14:25:27 +0000 Subject: [PATCH 8/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/hadamard_transform/customized_pipeline.cuh | 2 +- transformer_engine/common/normalization/kernel_traits.h | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) 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/kernel_traits.h b/transformer_engine/common/normalization/kernel_traits.h index dd30c1c22a..41be82e742 100644 --- a/transformer_engine/common/normalization/kernel_traits.h +++ b/transformer_engine/common/normalization/kernel_traits.h @@ -48,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)); @@ -86,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; From dbfb9e925933bff0924cdc6bde74cf85b6646f16 Mon Sep 17 00:00:00 2001 From: CarlosGomes98 Date: Thu, 2 Jul 2026 15:45:09 +0200 Subject: [PATCH 9/9] import cleanup Signed-off-by: CarlosGomes98 --- .../fused_softmax/scaled_aligned_causal_masked_softmax.cu | 3 --- .../common/fused_softmax/scaled_masked_softmax.cu | 3 --- .../fused_softmax/scaled_upper_triang_masked_softmax.cu | 3 --- .../common/normalization/layernorm/ln_fwd_kernels.cuh | 8 +++----- .../common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh | 8 +++----- 5 files changed, 6 insertions(+), 19 deletions(-) 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 9d2756ec28..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 @@ -5,13 +5,10 @@ ************************************************************************/ #ifdef __CUDACC_RTC__ -#include - #include "utils.cuh" #else #include #include -#include #include #include #include diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu index d3029e5a4c..6f295ad00a 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -5,13 +5,10 @@ ************************************************************************/ #ifdef __CUDACC_RTC__ -#include - #include "utils.cuh" #else #include #include -#include #include #include #include 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 7c30c38850..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 @@ -5,13 +5,10 @@ ************************************************************************/ #ifdef __CUDACC_RTC__ -#include - #include "utils.cuh" #else #include #include -#include #include #include #include diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh index 8724b443d5..3e9c281f58 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh @@ -7,15 +7,13 @@ #ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ -#ifndef __CUDACC_RTC__ -#include -#include -#endif - #ifdef __CUDACC_RTC__ #include "kernel_params.h" #include "utils.cuh" #else +#include +#include + #include "../../utils.cuh" #include "../common.h" #endif diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh index 0f5341cffb..0f06b29a21 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh @@ -7,15 +7,13 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ -#ifndef __CUDACC_RTC__ -#include -#include -#endif - #ifdef __CUDACC_RTC__ #include "kernel_params.h" #include "utils.cuh" #else +#include +#include + #include "../../utils.cuh" #include "../common.h" #endif