metal : optimize multi-sequence FA vec kernel

ggerganov · ggerganov · commit fdfc7de7fcff · 2025-05-13T08:03:27.000+03:00
ggml-ci
diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal
@@ -3887,6 +3887,11 @@ kernel void kernel_flash_attn_ext_vec(
                 sm[tiisg] = pm[ic + tiisg];
             }
 
+            // skip -INF blocks
+            if (simd_max(sm[tiisg]) == -INFINITY) {
+                continue;
+            }
+
             // Q*K^T
             {
                 // each simdgroup processes 1 query and NE (NW/NL) head elements

Original file line number	Diff line number	Diff line change
`@@ -3887,6 +3887,11 @@ kernel void kernel_flash_attn_ext_vec(`
`3887`	`3887`	`sm[tiisg] = pm[ic + tiisg];`
`3888`	`3888`	`}`
`3889`	`3889`
	`3890`	`+ // skip -INF blocks`
	`3891`	`+ if (simd_max(sm[tiisg]) == -INFINITY) {`
	`3892`	`+ continue;`
	`3893`	`+ }`
	`3894`	`+`
`3890`	`3895`	`// Q*K^T`
`3891`	`3896`	`{`
`3892`	`3897`	`// each simdgroup processes 1 query and NE (NW/NL) head elements`