bugfix, read and write race at fast_ln_fwd (#56435)

jeng1220 · web-flow · commit 1f987a7591f4 · 2023-08-21T16:01:39.000+08:00
diff --git a/paddle/phi/kernels/funcs/layer_norm_impl.cu.h b/paddle/phi/kernels/funcs/layer_norm_impl.cu.h
@@ -268,10 +268,10 @@ __global__ __launch_bounds__(THREADS_PER_CTA) void fast_ln_fwd_kernel(
         for (int it = 0; it < WARPS_N; ++it) {
           mu_local += smem[warp_m * WARPS_N + it];
         }
-        smem[warp_m] = mu_local;
+        smem[warp_m * WARPS_N] = mu_local;
       }
       __syncthreads();
-      mu_local = smem[warp_m];
+      mu_local = smem[warp_m * WARPS_N];
     }
 
     mu_local *= rn;
@@ -295,6 +295,7 @@ __global__ __launch_bounds__(THREADS_PER_CTA) void fast_ln_fwd_kernel(
     }
 
     if (WARPS_N > 1) {
+      __syncthreads();
       if (lane == 0) {
         smem[warp_m * WARPS_N + warp_n] = var_local;
       }
@@ -305,10 +306,10 @@ __global__ __launch_bounds__(THREADS_PER_CTA) void fast_ln_fwd_kernel(
         for (int it = 0; it < WARPS_N; ++it) {
           var_local += smem[warp_m * WARPS_N + it];
         }
-        smem[warp_m] = var_local;
+        smem[warp_m * WARPS_N] = var_local;
       }
       __syncthreads();
-      var_local = smem[warp_m];
+      var_local = smem[warp_m * WARPS_N];
     }
 
     // Note: to assure if it is right for double
diff --git a/paddle/phi/kernels/gpu/layer_norm_kernel.cu b/paddle/phi/kernels/gpu/layer_norm_kernel.cu
@@ -606,7 +606,7 @@ void LayerNormKernel(const Context &dev_ctx,
   if ((feature_size >= 768 && feature_size <= 2048 && feature_size % 256 == 0 ||
        feature_size == 4096) &&
       scale != nullptr && bias != nullptr) {
-    can_call_fast_kernel = false;
+    can_call_fast_kernel = true;
   }
 
   if (can_call_fast_kernel) {

Original file line number	Diff line number	Diff line change
`@@ -268,10 +268,10 @@ __global__ __launch_bounds__(THREADS_PER_CTA) void fast_ln_fwd_kernel(`
`268`	`268`	`for (int it = 0; it < WARPS_N; ++it) {`
`269`	`269`	`mu_local += smem[warp_m * WARPS_N + it];`
`270`	`270`	`}`
`271`		`- smem[warp_m] = mu_local;`
	`271`	`+ smem[warp_m * WARPS_N] = mu_local;`
`272`	`272`	`}`
`273`	`273`	`__syncthreads();`
`274`		`- mu_local = smem[warp_m];`
	`274`	`+ mu_local = smem[warp_m * WARPS_N];`
`275`	`275`	`}`
`276`	`276`
`277`	`277`	`mu_local *= rn;`
`@@ -295,6 +295,7 @@ __global__ __launch_bounds__(THREADS_PER_CTA) void fast_ln_fwd_kernel(`
`295`	`295`	`}`
`296`	`296`
`297`	`297`	`if (WARPS_N > 1) {`
	`298`	`+ __syncthreads();`
`298`	`299`	`if (lane == 0) {`
`299`	`300`	`smem[warp_m * WARPS_N + warp_n] = var_local;`
`300`	`301`	`}`
`@@ -305,10 +306,10 @@ __global__ __launch_bounds__(THREADS_PER_CTA) void fast_ln_fwd_kernel(`
`305`	`306`	`for (int it = 0; it < WARPS_N; ++it) {`
`306`	`307`	`var_local += smem[warp_m * WARPS_N + it];`
`307`	`308`	`}`
`308`		`- smem[warp_m] = var_local;`
	`309`	`+ smem[warp_m * WARPS_N] = var_local;`
`309`	`310`	`}`
`310`	`311`	`__syncthreads();`
`311`		`- var_local = smem[warp_m];`
	`312`	`+ var_local = smem[warp_m * WARPS_N];`
`312`	`313`	`}`
`313`	`314`
`314`	`315`	`// Note: to assure if it is right for double`
Original file line number	Diff line number	Diff line change
`@@ -606,7 +606,7 @@ void LayerNormKernel(const Context &dev_ctx,`
`606`	`606`	`if ((feature_size >= 768 && feature_size <= 2048 && feature_size % 256 == 0 \|\|`
`607`	`607`	`feature_size == 4096) &&`
`608`	`608`	`scale != nullptr && bias != nullptr) {`
`609`		`- can_call_fast_kernel = false;`
	`609`	`+ can_call_fast_kernel = true;`
`610`	`610`	`}`
`611`	`611`
`612`	`612`	`if (can_call_fast_kernel) {`