CUDA: fix broken oob check for FA vec f32 kernel (ggml-org#7904)

JohannesGaessler · web-flow · commit 963552903f51 · 2024-06-12T17:41:51.000+02:00
diff --git a/ggml-cuda/fattn-vec-f32.cuh b/ggml-cuda/fattn-vec-f32.cuh
@@ -149,7 +149,7 @@ static __global__ void flash_attn_vec_ext_f32(
             for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) {
                 const int i = i0 + threadIdx.x;
 
-                Q_f2[j][i0/WARP_SIZE]    = ncols <= 2 || ic0 + j ? Q_f2_j[i] : make_float2(0.0f, 0.0f);
+                Q_f2[j][i0/WARP_SIZE]    = ncols <= 2 || ic0 + j < ne01 ? Q_f2_j[i] : make_float2(0.0f, 0.0f);
                 Q_f2[j][i0/WARP_SIZE].x *= scale;
                 Q_f2[j][i0/WARP_SIZE].y *= scale;
             }

Original file line number	Diff line number	Diff line change
`@@ -149,7 +149,7 @@ static __global__ void flash_attn_vec_ext_f32(`
`149`	`149`	`for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) {`
`150`	`150`	`const int i = i0 + threadIdx.x;`
`151`	`151`
`152`		`- Q_f2[j][i0/WARP_SIZE] = ncols <= 2 \|\| ic0 + j ? Q_f2_j[i] : make_float2(0.0f, 0.0f);`
	`152`	`+ Q_f2[j][i0/WARP_SIZE] = ncols <= 2 \|\| ic0 + j < ne01 ? Q_f2_j[i] : make_float2(0.0f, 0.0f);`
`153`	`153`	`Q_f2[j][i0/WARP_SIZE].x *= scale;`
`154`	`154`	`Q_f2[j][i0/WARP_SIZE].y *= scale;`
`155`	`155`	`}`