Selaa lähdekoodia

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

* batched-bench : fix pp batch contents

* metal : optimize multi-sequence FA vec kernel

ggml-ci
Georgi Gerganov 8 kuukautta sitten
vanhempi
sitoutus
c252e0c409
1 muutettua tiedostoa jossa 5 lisäystä ja 0 poistoa
  1. 5 0
      ggml/src/ggml-metal/ggml-metal.metal

+ 5 - 0
ggml/src/ggml-metal/ggml-metal.metal

@@ -3887,6 +3887,11 @@ kernel void kernel_flash_attn_ext_vec(
                 sm[tiisg] = pm[ic + tiisg];
                 sm[tiisg] = pm[ic + tiisg];
             }
             }
 
 
+            // skip -INF blocks
+            if (simd_max(sm[tiisg]) == -INFINITY) {
+                continue;
+            }
+
             // Q*K^T
             // Q*K^T
             {
             {
                 // each simdgroup processes 1 query and NE (NW/NL) head elements
                 // each simdgroup processes 1 query and NE (NW/NL) head elements