|
2 | 2 | // Licensed under the MIT License. |
3 | 3 |
|
4 | 4 | #include "core/providers/cuda/cuda_common.h" |
5 | | -#include "core/providers/cuda/nn/layer_norm_impl.h" |
6 | 5 | #include "core/common/narrow.h" |
7 | 6 | #include "skip_layer_norm.h" |
8 | 7 | #include "skip_layer_norm_impl.h" |
@@ -42,26 +41,14 @@ template <typename T, bool Simplified> |
42 | 41 | SkipLayerNorm<T, Simplified>::SkipLayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) { |
43 | 42 | ORT_ENFORCE(op_kernel_info.GetAttr<float>("epsilon", &epsilon_).IsOK()); |
44 | 43 | ORT_ENFORCE(epsilon_ >= 0); |
45 | | - |
46 | | -#ifdef BUILD_CUDA_EP_AS_PLUGIN |
47 | | - // Plugin adapter cannot static_cast to CUDAExecutionProvider directly. |
48 | | - // Use the adapter shim that reads the config from the per-EP runtime map. |
49 | | - strict_ = onnxruntime::cuda::GetCudaKernelAdapterSkipLayerNormStrictMode(op_kernel_info.GetExecutionProvider()); |
50 | | -#else |
51 | | - const CUDAExecutionProvider* cuda_ep = static_cast<const CUDAExecutionProvider*>(op_kernel_info.GetExecutionProvider()); |
52 | | - strict_ = cuda_ep->IsSkipLayerNormInStrictMode(); |
53 | | -#endif |
| 44 | + // Note: the enable_skip_layer_norm_strict_mode provider option is deprecated and ignored. |
| 45 | + // The kernel always accumulates in fp32, so the previous strict-mode path is no longer needed. |
54 | 46 | } |
55 | 47 |
|
56 | 48 | template <typename T, bool Simplified> |
57 | 49 | Status SkipLayerNorm<T, Simplified>::ComputeInternal(OpKernelContext* ctx) const { |
58 | 50 | const Tensor* input = ctx->Input<Tensor>(0); |
59 | 51 | const Tensor* skip = ctx->Input<Tensor>(1); |
60 | | - if (strict_ && skip->Shape() != input->Shape()) { |
61 | | - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, |
62 | | - "'input' and 'skip' shall have same shape when enable_skip_layer_norm_strict_mode is True"); |
63 | | - } |
64 | | - |
65 | 52 | const Tensor* gamma = ctx->Input<Tensor>(2); |
66 | 53 |
|
67 | 54 | const Tensor* beta = Simplified ? nullptr : ctx->Input<Tensor>(3); |
@@ -94,53 +81,34 @@ Status SkipLayerNorm<T, Simplified>::ComputeInternal(OpKernelContext* ctx) const |
94 | 81 |
|
95 | 82 | const int skip_size = onnxruntime::narrow<int>(skip->Shape().Size()); |
96 | 83 |
|
97 | | - if (strict_) { |
98 | | - HostApplyLayerNorm<CudaT, float, CudaT, Simplified>( |
99 | | - GetDeviceProp(), |
| 84 | + if constexpr (std::is_same_v<T, BFloat16>) { |
| 85 | + LaunchSkipLayerNormKernel<nv_bfloat16, Simplified>( |
100 | 86 | Stream(ctx), |
101 | | - reinterpret_cast<CudaT*>(output->MutableData<T>()), // Y_data |
102 | | - nullptr, // mean_data |
103 | | - nullptr, // inv_var_data |
104 | | - reinterpret_cast<const CudaT*>(input->Data<T>()), // X_data |
105 | | - row_count, // n1 |
106 | | - hidden_size, // n2 |
107 | | - (double)epsilon_, // epsilon |
108 | | - reinterpret_cast<const CudaT*>(gamma->Data<T>()), // gamma |
109 | | - (beta != nullptr) ? reinterpret_cast<const CudaT*>(beta->Data<T>()) : nullptr, // beta |
110 | | - 0, // no broadcast for gamma/beta |
111 | | - reinterpret_cast<const CudaT*>(skip->Data<T>()), // skip or residual to add |
112 | | - (bias != nullptr) ? reinterpret_cast<const CudaT*>(bias->Data<T>()) : nullptr, // bias to add |
113 | | - sum_output != nullptr ? reinterpret_cast<CudaT*>(sum_output->MutableData<T>()) : nullptr); |
| 87 | + reinterpret_cast<nv_bfloat16*>(output->MutableData<T>()), |
| 88 | + sum_output != nullptr ? reinterpret_cast<nv_bfloat16*>(sum_output->MutableData<T>()) : nullptr, |
| 89 | + reinterpret_cast<const nv_bfloat16*>(input->Data<T>()), |
| 90 | + reinterpret_cast<const nv_bfloat16*>(skip->Data<T>()), |
| 91 | + (bias != nullptr) ? reinterpret_cast<const nv_bfloat16*>(bias->Data<T>()) : nullptr, |
| 92 | + reinterpret_cast<const nv_bfloat16*>(gamma->Data<T>()), |
| 93 | + (beta != nullptr) ? reinterpret_cast<const nv_bfloat16*>(beta->Data<T>()) : nullptr, |
| 94 | + epsilon_, |
| 95 | + hidden_size, |
| 96 | + row_count, |
| 97 | + skip_size); |
114 | 98 | } else { |
115 | | - if constexpr (std::is_same_v<T, BFloat16>) { |
116 | | - LaunchSkipLayerNormKernel<nv_bfloat16, Simplified>( |
117 | | - Stream(ctx), |
118 | | - reinterpret_cast<nv_bfloat16*>(output->MutableData<T>()), |
119 | | - sum_output != nullptr ? reinterpret_cast<nv_bfloat16*>(sum_output->MutableData<T>()) : nullptr, |
120 | | - reinterpret_cast<const nv_bfloat16*>(input->Data<T>()), |
121 | | - reinterpret_cast<const nv_bfloat16*>(skip->Data<T>()), |
122 | | - (bias != nullptr) ? reinterpret_cast<const nv_bfloat16*>(bias->Data<T>()) : nullptr, |
123 | | - reinterpret_cast<const nv_bfloat16*>(gamma->Data<T>()), |
124 | | - (beta != nullptr) ? reinterpret_cast<const nv_bfloat16*>(beta->Data<T>()) : nullptr, |
125 | | - epsilon_, |
126 | | - hidden_size, |
127 | | - row_count, |
128 | | - skip_size); |
129 | | - } else { |
130 | | - LaunchSkipLayerNormKernel<CudaT, Simplified>( |
131 | | - Stream(ctx), |
132 | | - reinterpret_cast<CudaT*>(output->MutableData<T>()), |
133 | | - sum_output != nullptr ? reinterpret_cast<CudaT*>(sum_output->MutableData<T>()) : nullptr, |
134 | | - reinterpret_cast<const CudaT*>(input->Data<T>()), |
135 | | - reinterpret_cast<const CudaT*>(skip->Data<T>()), |
136 | | - (bias != nullptr) ? reinterpret_cast<const CudaT*>(bias->Data<T>()) : nullptr, |
137 | | - reinterpret_cast<const CudaT*>(gamma->Data<T>()), |
138 | | - (beta != nullptr) ? reinterpret_cast<const CudaT*>(beta->Data<T>()) : nullptr, |
139 | | - epsilon_, |
140 | | - hidden_size, |
141 | | - row_count, |
142 | | - skip_size); |
143 | | - } |
| 99 | + LaunchSkipLayerNormKernel<CudaT, Simplified>( |
| 100 | + Stream(ctx), |
| 101 | + reinterpret_cast<CudaT*>(output->MutableData<T>()), |
| 102 | + sum_output != nullptr ? reinterpret_cast<CudaT*>(sum_output->MutableData<T>()) : nullptr, |
| 103 | + reinterpret_cast<const CudaT*>(input->Data<T>()), |
| 104 | + reinterpret_cast<const CudaT*>(skip->Data<T>()), |
| 105 | + (bias != nullptr) ? reinterpret_cast<const CudaT*>(bias->Data<T>()) : nullptr, |
| 106 | + reinterpret_cast<const CudaT*>(gamma->Data<T>()), |
| 107 | + (beta != nullptr) ? reinterpret_cast<const CudaT*>(beta->Data<T>()) : nullptr, |
| 108 | + epsilon_, |
| 109 | + hidden_size, |
| 110 | + row_count, |
| 111 | + skip_size); |
144 | 112 | } |
145 | 113 |
|
146 | 114 | CUDA_RETURN_IF_ERROR(cudaGetLastError()); |
|
0 commit comments