Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ struct OrtCUDAProviderOptionsV2 {
int tunable_op_enable = 0; // flag specifying if TunableOp is enabled.
int tunable_op_tuning_enable = 0; // flag specifying if TunableOp is enabled for tuning, this relies on TunableOp is enabled.
int tunable_op_max_tuning_duration_ms = 0; // Max tuning duration time limit for TunableOp.
int enable_skip_layer_norm_strict_mode = 0; // flag specifying if SkipLayerNorm is in strict mode. If true, use LayerNormalization kernel.
// The strict mode has better accuracy but lower performance.
int enable_skip_layer_norm_strict_mode = 0; // [Deprecated] No longer has any effect. SkipLayerNorm always accumulates in fp32.
// Kept for backward compatibility; the value is ignored.
Comment thread
tianleiwu marked this conversation as resolved.
Outdated
int prefer_nhwc = 0; // make the CUDA EP NHWC preferred
int use_ep_level_unified_stream = 0; // flag specifying if ep level stream is used or not
int use_tf32 = 1; // use TF32
Expand Down
88 changes: 28 additions & 60 deletions onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
// Licensed under the MIT License.

#include "core/providers/cuda/cuda_common.h"
#include "core/providers/cuda/nn/layer_norm_impl.h"
#include "core/common/narrow.h"
#include "skip_layer_norm.h"
#include "skip_layer_norm_impl.h"
Expand Down Expand Up @@ -42,26 +41,14 @@ template <typename T, bool Simplified>
SkipLayerNorm<T, Simplified>::SkipLayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {
ORT_ENFORCE(op_kernel_info.GetAttr<float>("epsilon", &epsilon_).IsOK());
ORT_ENFORCE(epsilon_ >= 0);

#ifdef BUILD_CUDA_EP_AS_PLUGIN
// Plugin adapter cannot static_cast to CUDAExecutionProvider directly.
// Use the adapter shim that reads the config from the per-EP runtime map.
strict_ = onnxruntime::cuda::GetCudaKernelAdapterSkipLayerNormStrictMode(op_kernel_info.GetExecutionProvider());
#else
const CUDAExecutionProvider* cuda_ep = static_cast<const CUDAExecutionProvider*>(op_kernel_info.GetExecutionProvider());
strict_ = cuda_ep->IsSkipLayerNormInStrictMode();
#endif
// Note: the enable_skip_layer_norm_strict_mode provider option is deprecated and ignored.
// The kernel always accumulates in fp32, so the previous strict-mode path is no longer needed.
}

template <typename T, bool Simplified>
Status SkipLayerNorm<T, Simplified>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor* input = ctx->Input<Tensor>(0);
const Tensor* skip = ctx->Input<Tensor>(1);
if (strict_ && skip->Shape() != input->Shape()) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"'input' and 'skip' shall have same shape when enable_skip_layer_norm_strict_mode is True");
}

const Tensor* gamma = ctx->Input<Tensor>(2);

const Tensor* beta = Simplified ? nullptr : ctx->Input<Tensor>(3);
Expand Down Expand Up @@ -94,53 +81,34 @@ Status SkipLayerNorm<T, Simplified>::ComputeInternal(OpKernelContext* ctx) const

const int skip_size = onnxruntime::narrow<int>(skip->Shape().Size());

if (strict_) {
HostApplyLayerNorm<CudaT, float, CudaT, Simplified>(
GetDeviceProp(),
if constexpr (std::is_same_v<T, BFloat16>) {
LaunchSkipLayerNormKernel<nv_bfloat16, Simplified>(
Stream(ctx),
reinterpret_cast<CudaT*>(output->MutableData<T>()), // Y_data
nullptr, // mean_data
nullptr, // inv_var_data
reinterpret_cast<const CudaT*>(input->Data<T>()), // X_data
row_count, // n1
hidden_size, // n2
(double)epsilon_, // epsilon
reinterpret_cast<const CudaT*>(gamma->Data<T>()), // gamma
(beta != nullptr) ? reinterpret_cast<const CudaT*>(beta->Data<T>()) : nullptr, // beta
0, // no broadcast for gamma/beta
reinterpret_cast<const CudaT*>(skip->Data<T>()), // skip or residual to add
(bias != nullptr) ? reinterpret_cast<const CudaT*>(bias->Data<T>()) : nullptr, // bias to add
sum_output != nullptr ? reinterpret_cast<CudaT*>(sum_output->MutableData<T>()) : nullptr);
reinterpret_cast<nv_bfloat16*>(output->MutableData<T>()),
sum_output != nullptr ? reinterpret_cast<nv_bfloat16*>(sum_output->MutableData<T>()) : nullptr,
reinterpret_cast<const nv_bfloat16*>(input->Data<T>()),
reinterpret_cast<const nv_bfloat16*>(skip->Data<T>()),
(bias != nullptr) ? reinterpret_cast<const nv_bfloat16*>(bias->Data<T>()) : nullptr,
reinterpret_cast<const nv_bfloat16*>(gamma->Data<T>()),
(beta != nullptr) ? reinterpret_cast<const nv_bfloat16*>(beta->Data<T>()) : nullptr,
epsilon_,
hidden_size,
row_count,
skip_size);
} else {
if constexpr (std::is_same_v<T, BFloat16>) {
LaunchSkipLayerNormKernel<nv_bfloat16, Simplified>(
Stream(ctx),
reinterpret_cast<nv_bfloat16*>(output->MutableData<T>()),
sum_output != nullptr ? reinterpret_cast<nv_bfloat16*>(sum_output->MutableData<T>()) : nullptr,
reinterpret_cast<const nv_bfloat16*>(input->Data<T>()),
reinterpret_cast<const nv_bfloat16*>(skip->Data<T>()),
(bias != nullptr) ? reinterpret_cast<const nv_bfloat16*>(bias->Data<T>()) : nullptr,
reinterpret_cast<const nv_bfloat16*>(gamma->Data<T>()),
(beta != nullptr) ? reinterpret_cast<const nv_bfloat16*>(beta->Data<T>()) : nullptr,
epsilon_,
hidden_size,
row_count,
skip_size);
} else {
LaunchSkipLayerNormKernel<CudaT, Simplified>(
Stream(ctx),
reinterpret_cast<CudaT*>(output->MutableData<T>()),
sum_output != nullptr ? reinterpret_cast<CudaT*>(sum_output->MutableData<T>()) : nullptr,
reinterpret_cast<const CudaT*>(input->Data<T>()),
reinterpret_cast<const CudaT*>(skip->Data<T>()),
(bias != nullptr) ? reinterpret_cast<const CudaT*>(bias->Data<T>()) : nullptr,
reinterpret_cast<const CudaT*>(gamma->Data<T>()),
(beta != nullptr) ? reinterpret_cast<const CudaT*>(beta->Data<T>()) : nullptr,
epsilon_,
hidden_size,
row_count,
skip_size);
}
LaunchSkipLayerNormKernel<CudaT, Simplified>(
Stream(ctx),
reinterpret_cast<CudaT*>(output->MutableData<T>()),
sum_output != nullptr ? reinterpret_cast<CudaT*>(sum_output->MutableData<T>()) : nullptr,
reinterpret_cast<const CudaT*>(input->Data<T>()),
reinterpret_cast<const CudaT*>(skip->Data<T>()),
(bias != nullptr) ? reinterpret_cast<const CudaT*>(bias->Data<T>()) : nullptr,
reinterpret_cast<const CudaT*>(gamma->Data<T>()),
(beta != nullptr) ? reinterpret_cast<const CudaT*>(beta->Data<T>()) : nullptr,
epsilon_,
hidden_size,
row_count,
skip_size);
}

CUDA_RETURN_IF_ERROR(cudaGetLastError());
Expand Down
1 change: 0 additions & 1 deletion onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ class SkipLayerNorm final : public CudaKernel {

private:
float epsilon_;
bool strict_;
};

} // namespace cuda
Expand Down
1 change: 0 additions & 1 deletion onnxruntime/core/providers/cuda/cuda_execution_provider.h
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,6 @@ class CUDAExecutionProvider : public IExecutionProvider {
bool DoCopyOnDefaultStream() const { return info_.do_copy_in_default_stream; }
bool GetCudnnConvUseMaxWorkspace() const { return info_.cudnn_conv_use_max_workspace; }
bool GetCudnnConv1dPadToNc1d() const { return info_.cudnn_conv1d_pad_to_nc1d; }
bool IsSkipLayerNormInStrictMode() const { return info_.enable_skip_layer_norm_strict_mode; }
bool IsNHWCPreferred() const { return info_.prefer_nhwc; }
bool IsFuseConvBias() const { return info_.fuse_conv_bias; }
bool UseTF32() const { return info_.use_tf32; }
Expand Down
5 changes: 0 additions & 5 deletions onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h
Original file line number Diff line number Diff line change
Expand Up @@ -738,11 +738,6 @@ inline void SetCudaKernelAdapterRuntimeConfigForProvider(
PL_CUDA_CALL_THROW(cudaGetDeviceProperties(&config->device_prop, config->device_id));
}
Comment thread
tianleiwu marked this conversation as resolved.

inline bool GetCudaKernelAdapterSkipLayerNormStrictMode(const void* provider) {
const auto config = detail::GetCudaKernelAdapterRuntimeConfigForProvider(provider);
return config->skip_layer_norm_strict_mode;
}

// Global aliases and shims
using Status = onnxruntime::common::Status;
using MLFloat16 = onnxruntime::MLFloat16;
Expand Down
28 changes: 3 additions & 25 deletions onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ namespace test {
constexpr float epsilon_ = 1e-12f;

static void RunOneTest(
bool strict,
const std::vector<float>& input_data,
const std::vector<float>& skip_data,
const std::vector<float>& gamma_data,
Expand Down Expand Up @@ -136,14 +135,7 @@ static void RunOneTest(
ToBFloat16(sum_output_data));
}

if (strict) {
Ort::CUDAProviderOptions cuda_options;
std::unordered_map<std::string, std::string> options = {{"enable_skip_layer_norm_strict_mode", "1"}};
cuda_options.Update(options);
execution_providers.push_back(CudaExecutionProviderWithOptions(std::move(cuda_options)));
} else {
execution_providers.push_back(DefaultCudaExecutionProvider());
}
execution_providers.push_back(DefaultCudaExecutionProvider());

test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
} else if (HasCudaEnvironment(530 /*min_cuda_architecture*/) ||
Expand Down Expand Up @@ -187,14 +179,7 @@ static void RunOneTest(
} else if (dml_ep != nullptr) {
execution_providers.push_back(DefaultDmlExecutionProvider());
} else {
if (strict) {
Ort::CUDAProviderOptions cuda_options;
std::unordered_map<std::string, std::string> options = {{"enable_skip_layer_norm_strict_mode", "1"}};
cuda_options.Update(options);
execution_providers.push_back(CudaExecutionProviderWithOptions(std::move(cuda_options)));
} else {
execution_providers.push_back(DefaultCudaExecutionProvider());
}
execution_providers.push_back(DefaultCudaExecutionProvider());
}
Comment thread
tianleiwu marked this conversation as resolved.

test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
Expand All @@ -220,16 +205,9 @@ static void RunTest(
bool use_token_count = false,
bool broadcast_skip = false,
bool no_batch_size = false) {
RunOneTest(false, input_data, skip_data, gamma_data, beta_data, bias_data, output_data, sum_output_data,
RunOneTest(input_data, skip_data, gamma_data, beta_data, bias_data, output_data, sum_output_data,
epsilon, batch_size, sequence_length, hidden_size, use_float16, use_bfloat16, no_beta, simplified,
use_token_count, broadcast_skip, no_batch_size);

// strict mode does not support skip broadcasting.
if (!broadcast_skip) {
RunOneTest(true, input_data, skip_data, gamma_data, beta_data, bias_data, output_data, sum_output_data,
epsilon, batch_size, sequence_length, hidden_size, use_float16, use_bfloat16, no_beta, simplified,
use_token_count, broadcast_skip, no_batch_size);
}
}

TEST(SkipLayerNormTest, SkipLayerNormPrePack) {
Expand Down
Loading