vulkan: Fix data race/hang in scalar/cm1 flash attention (llama/17887)

diff --git a/src/ggml-vulkan/vulkan-shaders/flash_attn.comp b/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
index 4bef48b006ce78da839e758ae7e3a389c3f5e612..0379e5d50240f8d31cfc63434670c606a5cd45ad 100644 (file)
--- a/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
+++ b/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
@@ -256,6 +256,9 @@ void main() {
         barrier();
     }
 
+    // prevent race on tmpsh
+    barrier();
+
     // reduce across threads
 
     [[unroll]] for (uint32_t r = 0; r < Br; ++r) {

diff --git a/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp
index cd82e4abfabc46d1733610bb954f9764d8046b93..c995ab140ee65776183e38b59071a8cd43b42469 100644 (file)
--- a/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp
+++ b/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp
@@ -302,6 +302,9 @@ void main() {
         barrier();
     }
 
+    // prevent race on tmpsh
+    barrier();
+
     // reduce across threads
 
     float rowmaxf[rows_per_thread], eMf[rows_per_thread], Moldf[rows_per_thread];