diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index e7aaf78f6c..9541db4b19 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -250,18 +250,27 @@ list(APPEND transformer_engine_cuda_arch_specific_sources activation/gelu.cu activation/gelu_dbias.cu activation/gelu_grouped.cu + activation/gelu_grouped_bwd.cu activation/gelu_grouped_dbias.cu + activation/qgelu_grouped.cu + activation/qgelu_grouped_bwd.cu + activation/qgelu_grouped_dbias.cu activation/glu.cu activation/relu.cu activation/relu_dbias.cu activation/relu_grouped.cu + activation/relu_grouped_bwd.cu activation/relu_grouped_dbias.cu + activation/srelu_grouped.cu + activation/srelu_grouped_bwd.cu + activation/srelu_grouped_dbias.cu activation/scaled_activation.cu activation/scaled_srelu.cu activation/scaled_swiglu.cu activation/swiglu.cu activation/swiglu_dbias.cu activation/swiglu_grouped.cu + activation/swiglu_grouped_bwd.cu activation/swiglu_grouped_dbias.cu cast/cast.cu cast/cast_dbias.cu @@ -631,18 +640,27 @@ if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH) list(APPEND nvte_sources_with_fast_math activation/gelu.cu activation/gelu_dbias.cu activation/gelu_grouped.cu + activation/gelu_grouped_bwd.cu activation/gelu_grouped_dbias.cu + activation/qgelu_grouped.cu + activation/qgelu_grouped_bwd.cu + activation/qgelu_grouped_dbias.cu activation/glu.cu activation/relu.cu activation/relu_dbias.cu activation/relu_grouped.cu + activation/relu_grouped_bwd.cu activation/relu_grouped_dbias.cu + activation/srelu_grouped.cu + activation/srelu_grouped_bwd.cu + activation/srelu_grouped_dbias.cu activation/scaled_activation.cu activation/scaled_srelu.cu activation/scaled_swiglu.cu activation/swiglu.cu activation/swiglu_dbias.cu activation/swiglu_grouped.cu + activation/swiglu_grouped_bwd.cu activation/swiglu_grouped_dbias.cu) endif() diff --git a/transformer_engine/common/activation/gelu_grouped.cu b/transformer_engine/common/activation/gelu_grouped.cu index c3267356f8..82ab374aaa 100644 --- a/transformer_engine/common/activation/gelu_grouped.cu +++ b/transformer_engine/common/activation/gelu_grouped.cu @@ -14,40 +14,3 @@ void nvte_group_gelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cu dispatch::group_quantize_fwd_helper>(input, output, nullptr, stream); } - -void nvte_group_dgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_dgelu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} - -void nvte_group_qgelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, - cudaStream_t stream) { - NVTE_API_CALL(nvte_group_qgelu); - using namespace transformer_engine; - constexpr bool IS_ACT = true; - dispatch::group_quantize_fwd_helper>(input, output, nullptr, - stream); -} - -void nvte_group_dqgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_dqgelu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} diff --git a/transformer_engine/common/activation/gelu_grouped_bwd.cu b/transformer_engine/common/activation/gelu_grouped_bwd.cu new file mode 100644 index 0000000000..8d3c35a093 --- /dev/null +++ b/transformer_engine/common/activation/gelu_grouped_bwd.cu @@ -0,0 +1,22 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_dgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_dgelu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/gelu_grouped_dbias.cu b/transformer_engine/common/activation/gelu_grouped_dbias.cu index e8b549f692..cf19b17971 100644 --- a/transformer_engine/common/activation/gelu_grouped_dbias.cu +++ b/transformer_engine/common/activation/gelu_grouped_dbias.cu @@ -20,17 +20,3 @@ void nvte_group_quantize_dbias_dgelu(const NVTEGroupedTensor input, dispatch::group_quantize_bwd_helper>( input, activation_input, output, dbias, workspace, nullptr, stream); } - -void nvte_group_quantize_dbias_dqgelu(const NVTEGroupedTensor input, - const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTEGroupedTensor dbias, - NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_quantize_dbias_dqgelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} diff --git a/transformer_engine/common/activation/qgelu_grouped.cu b/transformer_engine/common/activation/qgelu_grouped.cu new file mode 100644 index 0000000000..e241fa3163 --- /dev/null +++ b/transformer_engine/common/activation/qgelu_grouped.cu @@ -0,0 +1,17 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_qgelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_qgelu); + using namespace transformer_engine; + constexpr bool IS_ACT = true; + dispatch::group_quantize_fwd_helper>(input, output, nullptr, + stream); +} diff --git a/transformer_engine/common/activation/qgelu_grouped_bwd.cu b/transformer_engine/common/activation/qgelu_grouped_bwd.cu new file mode 100644 index 0000000000..8558bee0bd --- /dev/null +++ b/transformer_engine/common/activation/qgelu_grouped_bwd.cu @@ -0,0 +1,22 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_dqgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_dqgelu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/qgelu_grouped_dbias.cu b/transformer_engine/common/activation/qgelu_grouped_dbias.cu new file mode 100644 index 0000000000..26d735cd71 --- /dev/null +++ b/transformer_engine/common/activation/qgelu_grouped_dbias.cu @@ -0,0 +1,22 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_quantize_dbias_dqgelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize_dbias_dqgelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/relu_grouped.cu b/transformer_engine/common/activation/relu_grouped.cu index 93ce6b82fe..fe70ba2a48 100644 --- a/transformer_engine/common/activation/relu_grouped.cu +++ b/transformer_engine/common/activation/relu_grouped.cu @@ -14,40 +14,3 @@ void nvte_group_relu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cu dispatch::group_quantize_fwd_helper>(input, output, nullptr, stream); } - -void nvte_group_drelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_drelu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} - -void nvte_group_srelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, - cudaStream_t stream) { - NVTE_API_CALL(nvte_group_srelu); - using namespace transformer_engine; - constexpr bool IS_ACT = true; - dispatch::group_quantize_fwd_helper>(input, output, nullptr, - stream); -} - -void nvte_group_dsrelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_dsrelu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} diff --git a/transformer_engine/common/activation/relu_grouped_bwd.cu b/transformer_engine/common/activation/relu_grouped_bwd.cu new file mode 100644 index 0000000000..f4143f3b9b --- /dev/null +++ b/transformer_engine/common/activation/relu_grouped_bwd.cu @@ -0,0 +1,22 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_drelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_drelu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/relu_grouped_dbias.cu b/transformer_engine/common/activation/relu_grouped_dbias.cu index 2b9dcd35d4..c733fcee01 100644 --- a/transformer_engine/common/activation/relu_grouped_dbias.cu +++ b/transformer_engine/common/activation/relu_grouped_dbias.cu @@ -20,17 +20,3 @@ void nvte_group_quantize_dbias_drelu(const NVTEGroupedTensor input, dispatch::group_quantize_bwd_helper>( input, activation_input, output, dbias, workspace, nullptr, stream); } - -void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, - const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTEGroupedTensor dbias, - NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_quantize_dbias_dsrelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} diff --git a/transformer_engine/common/activation/srelu_grouped.cu b/transformer_engine/common/activation/srelu_grouped.cu new file mode 100644 index 0000000000..d7f43bf2c8 --- /dev/null +++ b/transformer_engine/common/activation/srelu_grouped.cu @@ -0,0 +1,17 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_srelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_srelu); + using namespace transformer_engine; + constexpr bool IS_ACT = true; + dispatch::group_quantize_fwd_helper>(input, output, nullptr, + stream); +} diff --git a/transformer_engine/common/activation/srelu_grouped_bwd.cu b/transformer_engine/common/activation/srelu_grouped_bwd.cu new file mode 100644 index 0000000000..1793cc02fc --- /dev/null +++ b/transformer_engine/common/activation/srelu_grouped_bwd.cu @@ -0,0 +1,22 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_dsrelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_dsrelu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/srelu_grouped_dbias.cu b/transformer_engine/common/activation/srelu_grouped_dbias.cu new file mode 100644 index 0000000000..207c4ee6e5 --- /dev/null +++ b/transformer_engine/common/activation/srelu_grouped_dbias.cu @@ -0,0 +1,22 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize_dbias_dsrelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/swiglu_grouped.cu b/transformer_engine/common/activation/swiglu_grouped.cu index 160ab66288..8cb927afbd 100644 --- a/transformer_engine/common/activation/swiglu_grouped.cu +++ b/transformer_engine/common/activation/swiglu_grouped.cu @@ -14,17 +14,3 @@ void nvte_group_silu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cu dispatch::group_quantize_fwd_helper>(input, output, nullptr, stream); } - -void nvte_group_dsilu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_dsilu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} diff --git a/transformer_engine/common/activation/swiglu_grouped_bwd.cu b/transformer_engine/common/activation/swiglu_grouped_bwd.cu new file mode 100644 index 0000000000..b02c47ab21 --- /dev/null +++ b/transformer_engine/common/activation/swiglu_grouped_bwd.cu @@ -0,0 +1,22 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_dsilu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_dsilu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +}