allow gqa with dim3>1

jeffbolznv · jeffbolznv · commit 45792afd3a44 · 2025-07-03T10:39:04.000-05:00
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -6120,7 +6120,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
     }
 
     if (N == 1 && qk_ratio > 1 && qk_ratio <= max_gqa &&
-        qk_ratio * nek2 == neq2 && nek2 == nev2 && neq3 == 1 && nek3 == 1 && nev3 == 1 && nem2 <= 1 && nem3 <= 1) {
+        qk_ratio * nek2 == neq2 && nek2 == nev2 && nem2 <= 1) {
         // grouped query attention - make the N dimension equal to gqa_ratio, reduce
         // workgroups proportionally in y dimension. The shader will detect gqa_ratio > 1
         // and change addressing calculations to index Q's dimension 2.

Original file line number	Diff line number	Diff line change
`@@ -6120,7 +6120,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx`
`6120`	`6120`	`}`
`6121`	`6121`
`6122`	`6122`	`if (N == 1 && qk_ratio > 1 && qk_ratio <= max_gqa &&`
`6123`		`- qk_ratio * nek2 == neq2 && nek2 == nev2 && neq3 == 1 && nek3 == 1 && nev3 == 1 && nem2 <= 1 && nem3 <= 1) {`
	`6123`	`+ qk_ratio * nek2 == neq2 && nek2 == nev2 && nem2 <= 1) {`
`6124`	`6124`	`// grouped query attention - make the N dimension equal to gqa_ratio, reduce`
`6125`	`6125`	`// workgroups proportionally in y dimension. The shader will detect gqa_ratio > 1`
`6126`	`6126`	`// and change addressing calculations to index Q's dimension 2.`