Explorar o código

cuda : fix FA Q src index (1 -> 0) (#9374)

Georgi Gerganov hai 1 ano
pai
achega
e079bffb66
Modificáronse 1 ficheiros con 2 adicións e 2 borrados
  1. 2 2
      ggml/src/ggml-cuda/fattn.cu

+ 2 - 2
ggml/src/ggml-cuda/fattn.cu

@@ -152,7 +152,7 @@ static void ggml_cuda_flash_attn_ext_wmma_f16(ggml_backend_cuda_context & ctx, g
     }                                                                       \
     }                                                                       \
 
 
 static void ggml_cuda_flash_attn_ext_vec_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
 static void ggml_cuda_flash_attn_ext_vec_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
-    ggml_tensor * Q = dst->src[1];
+    ggml_tensor * Q = dst->src[0];
     ggml_tensor * K = dst->src[1];
     ggml_tensor * K = dst->src[1];
     ggml_tensor * V = dst->src[2];
     ggml_tensor * V = dst->src[2];
 
 
@@ -227,7 +227,7 @@ static void ggml_cuda_flash_attn_ext_vec_f16(ggml_backend_cuda_context & ctx, gg
     }                                                                       \
     }                                                                       \
 
 
 static void ggml_cuda_flash_attn_ext_vec_f32(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
 static void ggml_cuda_flash_attn_ext_vec_f32(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
-    ggml_tensor * Q = dst->src[1];
+    ggml_tensor * Q = dst->src[0];
     ggml_tensor * K = dst->src[1];
     ggml_tensor * K = dst->src[1];
     ggml_tensor * V = dst->src[2];
     ggml_tensor * V = dst->src[2];