metal : optimize multi-sequence FA vec kernel (ggml-org#13493)

ggerganov · web-flow · commit c252e0c4097b · 2025-05-13T18:04:00.000+03:00
* batched-bench : fix pp batch contents

* metal : optimize multi-sequence FA vec kernel

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`