CUDA: remove unnecessary warp reduce in FA (ggml/1032)

mahorozte · ZhaoXiaoYu · ggerganov · commit e9e661bd5936 · 2024-12-03T20:04:49.000+02:00
* kqmax_new_j in every thread within warp is same after operate at line 199,this reduce can be omit

* same problem in vec32

---------

Co-authored-by: ZhaoXiaoYu &lt;zhao.xiaoyu@zte.com.cn&gt;
diff --git a/ggml/src/ggml-cuda/fattn-vec-f16.cuh b/ggml/src/ggml-cuda/fattn-vec-f16.cuh
@@ -220,7 +220,6 @@ static __global__ void flash_attn_vec_ext_f16(
         for (int j = 0; j < ncols; ++j) {
             half kqmax_new_j = ncols == 1 ? kqmax_new : kqmax_new_arr[j];
 
-            kqmax_new_j = warp_reduce_max(kqmax_new_j);
             if (threadIdx.x == 0) {
                 kqmax_shared[j][threadIdx.y] = kqmax_new_j;
             }
diff --git a/ggml/src/ggml-cuda/fattn-vec-f32.cuh b/ggml/src/ggml-cuda/fattn-vec-f32.cuh
@@ -206,7 +206,6 @@ static __global__ void flash_attn_vec_ext_f32(
         for (int j = 0; j < ncols; ++j) {
             float kqmax_new_j = kqmax_new_arr[j];
 
-            kqmax_new_j = warp_reduce_max(kqmax_new_j);
             if (threadIdx.x == 0) {
                 kqmax_shared[j][threadIdx.y] = kqmax_new_j;
             }

Original file line number	Diff line number	Diff line change
`@@ -220,7 +220,6 @@ static __global__ void flash_attn_vec_ext_f16(`
`220`	`220`	`for (int j = 0; j < ncols; ++j) {`
`221`	`221`	`half kqmax_new_j = ncols == 1 ? kqmax_new : kqmax_new_arr[j];`
`222`	`222`
`223`		`- kqmax_new_j = warp_reduce_max(kqmax_new_j);`
`224`	`223`	`if (threadIdx.x == 0) {`
`225`	`224`	`kqmax_shared[j][threadIdx.y] = kqmax_new_j;`
`226`	`225`	`}`
Original file line number	Diff line number	Diff line change
`@@ -206,7 +206,6 @@ static __global__ void flash_attn_vec_ext_f32(`
`206`	`206`	`for (int j = 0; j < ncols; ++j) {`
`207`	`207`	`float kqmax_new_j = kqmax_new_arr[j];`
`208`	`208`
`209`		`- kqmax_new_j = warp_reduce_max(kqmax_new_j);`
`210`	`209`	`if (threadIdx.x == 0) {`
`211`	`210`	`kqmax_shared[j][threadIdx.y] = kqmax_new_j;`
`212`	`211`	`}`