From: Johannes Gäßler <redacted>
Date: Wed, 12 Jun 2024 15:41:51 +0000 (+0200)
Subject: CUDA: fix broken oob check for FA vec f32 kernel (llama/7904)
X-Git-Tag: upstream/0.0.1642~597
X-Git-Url: https://git.djapps.eu/?a=commitdiff_plain;h=91f2ca8bb10a6051150b6a72fb7cf10f348ab419;p=pkg%2Fggml%2Fsources%2Fggml

CUDA: fix broken oob check for FA vec f32 kernel (llama/7904)
---

diff --git a/src/ggml-cuda/fattn-vec-f32.cuh b/src/ggml-cuda/fattn-vec-f32.cuh
index ddf0c837..11a5e355 100644
--- a/src/ggml-cuda/fattn-vec-f32.cuh
+++ b/src/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;
             }