Ver Fonte

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 há 8 meses atrás
pai
commit
c252e0c409
1 ficheiros alterados com 5 adições e 0 exclusões
  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];
             }
 
+            // skip -INF blocks
+            if (simd_max(sm[tiisg]) == -INFINITY) {
+                continue;
+            }
+
             // Q*K^T
             {
                 // each simdgroup processes 1 query and NE (NW/NL) head elements