leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

Allow changing epsilon parameter in RMS norm kernel (#1112)

Browse Source
This commit is contained in:
masahi
2023-10-02 20:40:28 -04:00
committed by GitHub
parent 26986bbc60
commit ff61a49dd1
2 changed files with 13 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/util/include/cutlass/util/device_rmsnorm.h
+8 -7
View File
@@ -43,7 +43,7 @@ namespace cutlass {
__global__ void rmsnorm_twoPassAlgo_e8(float4 *output, const float4 *input,
const float4 *weight,
const int m, const int n) {
const int m, const int n, float epsilon) {
const int m_idx = blockIdx.x;
const int tid = threadIdx.x;
const int bdimx = blockDim.x;
@@ -76,7 +76,7 @@ __global__ void rmsnorm_twoPassAlgo_e8(float4 *output, const float4 *input,
blockReduceSum<float, 1>(local_sums);
}
if (threadIdx.x == 0) {
s_mean = rsqrtf(local_sums[0] / n + 1e-6);
s_mean = rsqrtf(local_sums[0] / n + epsilon);
}
__syncthreads();
@@ -117,7 +117,8 @@ template<typename T>
__global__ void rmsnorm_twoPassAlgo_e1(T* output,
const T* input,
const T* weight,
const int m, const int n)
const int m, const int n,
float epsilon)
{
const int m_idx = blockIdx.x;
const int tid = threadIdx.x;
@@ -139,7 +140,7 @@ __global__ void rmsnorm_twoPassAlgo_e1(T* output,
blockReduceSum<float, 1>(local_sums);
}
if (threadIdx.x == 0) {
s_mean = rsqrtf(local_sums[0] / n + 1e-6);
s_mean = rsqrtf(local_sums[0] / n + epsilon);
}
__syncthreads();
@@ -155,7 +156,7 @@ void rmsnorm(cutlass::MatrixCoord tensor_size,
TensorRef<T, layout::RowMajor> ref_output,
TensorRef<T, layout::RowMajor> ref_input,
TensorRef<T, layout::RowMajor> ref_weight,
cudaStream_t stream){
cudaStream_t stream, float epsilon = 1e-5){
const int m = tensor_size.row();
const int n = tensor_size.column();
T* output = ref_output.data();
@@ -167,12 +168,12 @@ void rmsnorm(cutlass::MatrixCoord tensor_size,
dim3 block(min(1024, (n / 8 + 31) / 32 * 32));
rmsnorm_twoPassAlgo_e8<<<grid, block, 0, stream>>>(
(float4 *)output, (const float4 *)input, (const float4 *)weight, m, n);
(float4 *)output, (const float4 *)input, (const float4 *)weight, m, n, epsilon);
} else {
dim3 block(min(1024, ((n + 31)/32 + 31)/32*32));
rmsnorm_twoPassAlgo_e1<<<grid, block, 0, stream>>>(
output, input, weight, m, n);
output, input, weight, m, n, epsilon);
}
auto result = cudaGetLastError();