vulkan: fix data race in mul_mat_id shader (llama/19790)

diff --git a/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
index 717d124e01933cf7f44ac136260ffc219cf40545..497a18ff8a7c73f214d23ba4f88f5176fb1a43fd 100644 (file)
--- a/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
+++ b/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
@@ -167,7 +167,9 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) {
         uint id = ids[iter++];
         uvec4 ballot = subgroupBallot(in_range && id == expert_idx);
 
-        ballots_sh[gl_SubgroupID] = ballot;
+        if (gl_SubgroupInvocationID == 0) {
+            ballots_sh[gl_SubgroupID] = ballot;
+        }
         barrier();
 
         uint subgroup_base = 0;

diff --git a/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl b/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl
index 743004ff8ad1e092574229fd88f94c0775ef7966..26c5c12a49a226085921d1019eb0776e4a6a109a 100644 (file)
--- a/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl
+++ b/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl
@@ -43,7 +43,9 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) {
         uint id = ids[iter++];
         uvec4 ballot = subgroupBallot(in_range && id == expert_idx);
 
-        ballots_sh[gl_SubgroupID] = ballot;
+        if (gl_SubgroupInvocationID == 0) {
+            ballots_sh[gl_SubgroupID] = ballot;
+        }
         barrier();
 
         uint subgroup_base = 0;