wip

ggml-metal.m

@@ -2183,6 +2183,7 @@ static bool ggml_metal_graph_compute(
                         struct ggml_tensor * src3 = gf->nodes[i]->src[3];
 
                         GGML_ASSERT(ggml_are_same_shape(src1, src2));
+                        GGML_ASSERT(src3);
 
                         size_t offs_src2 = 0;
                         size_t offs_src3 = 0;
@@ -2252,15 +2253,19 @@ static bool ggml_metal_graph_compute(
                         [encoder setBytes:&ne3     length:sizeof( int64_t) atIndex:26];
                         [encoder setBytes:&scale   length:sizeof(   float) atIndex:27];
 
-                        const int64_t nwarps = 32;
-                        const int64_t nhptg  = 2; // heads per threadgroup
+                        const int64_t nsg   = 16; // simdgroups per threadgroup (a.k.a. warps)
+                        const int64_t nhptg = 2; // heads   per threadgroup !! sync with kernel template arguments !!
+                        const int64_t nqptg = 2; // queries per threadgroup !! sync with kernel template arguments !!
+                        const int64_t ncpsg = 8;
 
-                        const size_t smem = (nhptg*ne00 + nwarps*(nhptg*ne00 + 32))*(sizeof(float)/2);
+                      //const size_t smem = nqptg*(nhptg*ne00 + nsg*(nhptg*ne00 + 256))*(sizeof(float)/2);
+                        const size_t smem = nqptg*(nhptg*ne00 + nsg*(32*ncpsg))*(sizeof(float)/2);
 
+                        //printf("smem: %zu, max: %zu\n", smem, ctx->device.maxThreadgroupMemoryLength);
                         GGML_ASSERT(smem <= ctx->device.maxThreadgroupMemoryLength);
                         [encoder setThreadgroupMemoryLength:smem atIndex:0];
 
-                        [encoder dispatchThreadgroups:MTLSizeMake(ne01, (ne02 + nhptg - 1)/(nhptg), ne03) threadsPerThreadgroup:MTLSizeMake(32, nwarps, 1)];
+                        [encoder dispatchThreadgroups:MTLSizeMake((ne01 + nqptg - 1)/nqptg, (ne02 + nhptg - 1)/(nhptg), ne03) threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
                     } break;
                 case GGML_OP_DUP:
                 case GGML_OP_CPY:

ggml-metal.metal

@@ -1995,7 +1995,7 @@ typedef void (flash_attn_ext_f16_t)(
         uint  tiisg[[thread_index_in_simdgroup]],
         uint  sgitg[[simdgroup_index_in_threadgroup]]);
 
-template<int64_t D, int64_t R> // head size, rows per threadgroup
+template<int64_t D, int64_t H, int64_t Q, int64_t C> // head size, heads per threadgroup, queries per threadgroup
 kernel void kernel_flash_attn_ext_f16(
         device const  char * q,
         device const  char * k,
@@ -2032,177 +2032,276 @@ kernel void kernel_flash_attn_ext_f16(
         uint  tiisg[[thread_index_in_simdgroup]],
         uint  sgitg[[simdgroup_index_in_threadgroup]]) {
     const uint nsg = ntg.y;         // number of simdgroups
-    const uint tph = N_SIMDWIDTH/R; // threads per head
+    const uint tph = N_SIMDWIDTH/H; // threads per head
 
     const int64_t iq3 = tgpig[2];
-    const int64_t iq2 = tgpig[1]*R + tiisg/tph;
-    const int64_t iq1 = tgpig[0];
+    const int64_t iq2 = tgpig[1]*H + tiisg/tph;
+    const int64_t iq1 = tgpig[0]*Q;
 
     if (iq2 >= ne02) {
         return;
     }
 
-    // assume K and V are same shape
-    const int64_t ne22 = ne12;
-    const int64_t ne23 = ne13;
-
-    const uint64_t nb21 = nb11;
-    const uint64_t nb22 = nb12;
-    const uint64_t nb23 = nb13;
-
-    // broadcast
-    const int64_t rk2 = ne02/ne12;
-    const int64_t rk3 = ne03/ne13;
-
-    const int64_t rv2 = ne02/ne22;
-    const int64_t rv3 = ne03/ne23;
-
-    // k indices
-    const int64_t ik2 = iq2 / rk2;
-    const int64_t ik3 = iq3 / rk3;
-
-    // v indices
-    const int64_t iv2 = iq2 / rv2;
-    const int64_t iv3 = iq3 / rv3;
-
-    const int64_t ir = iq3*ne02*ne01 + iq2*ne01 + iq1;
-
-    device const float * mp = mask ? (device const float *) (mask + (ir%ne31)*nb31) : nullptr;
-
     const int64_t D4 = D/4;
 
-    threadgroup half4 * pq4 = (threadgroup half4 *) (shared +                    0*R*D);
-    threadgroup half4 * ps4 = (threadgroup half4 *) (shared + sgitg*(R*D + 32) + 1*R*D);
-    threadgroup half  * ss  = (threadgroup half  *) (shared + sgitg*(R*D + 32) + 2*R*D);
+    const int64_t T  = (H*D + nsg*(32*C)); // shared memory size per query in half
+    const int64_t T4 = T/4;                // shared memory size per query in half4
+
+    threadgroup half4 * pq4 = (threadgroup half4 *) (shared +                0*H*D);
+    threadgroup half  * ss  = (threadgroup half  *) (shared + sgitg*(32*C) + 1*H*D);
+    threadgroup half4 * ss4 = (threadgroup half4 *) (shared + sgitg*(32*C) + 1*H*D);
 
     const uint tiih  = tiisg%tph; // thread index in head
     const uint hiisg = tiisg/tph; // head index in simdgroup
 
-    // load R heads from Q to shared memory
+    half4 ps4[Q][D4/tph];
+
+    // load H heads from Q to shared memory
     for (int64_t i = 0; i < D4/tph; ++i) {
-        if (sgitg == 0) {
-            pq4[hiisg*D4 + tph*i + tiih] = ((device const half4 *) ((device const char *) q + (iq1*nb01 + iq2*nb02 + iq3*nb03)))[tph*i + tiih];
+        for (int64_t j = sgitg; j < Q; j += nsg) {
+            if (iq1 + j < ne01) {
+                pq4[j*T4 + hiisg*D4 + tph*i + tiih] = ((device const half4 *) ((device const char *) q + ((iq1 + j)*nb01 + iq2*nb02 + iq3*nb03)))[tph*i + tiih];
+            } else {
+                pq4[j*T4 + hiisg*D4 + tph*i + tiih] = 0.0h;
+            }
         }
 
-        ps4[hiisg*D4 + tph*i + tiih] = 0.0h;
+        for (int64_t j = 0; j < Q; ++j) {
+            //ps4[j*T4 + hiisg*D4 + tph*i + tiih] = 0.0h;
+            ps4[j][i] = 0.0h;
+        }
+    }
+
+    for (int64_t j = 0; j < Q; ++j) {
+        ss[j*T + hiisg*tph + tiih] = 0.0h;
+        ss[j*T + hiisg*tph + tiih] = 0.0h;
     }
 
     threadgroup_barrier(mem_flags::mem_threadgroup);
 
-    half S = 0.0h;
-    half M = -INFINITY;
+    half S = { 0.0h };
+    half M = { -INFINITY };
 
-    for (int64_t ic = sgitg; ic < ne11; ic += nsg) {
-        const half mv = mp ? mp[ic] : 0.0h;
-        if (mv == -INFINITY) {
-            continue;
+    {
+        // assume K and V are same shape
+        const int64_t ne22 = ne12;
+        const int64_t ne23 = ne13;
+
+        const uint64_t nb21 = nb11;
+        const uint64_t nb22 = nb12;
+        const uint64_t nb23 = nb13;
+
+        // broadcast
+        const int64_t rk2 = ne02/ne12;
+        const int64_t rk3 = ne03/ne13;
+
+        const int64_t rv2 = ne02/ne22;
+        const int64_t rv3 = ne03/ne23;
+
+        // k indices
+        const int64_t ik2 = iq2 / rk2;
+        const int64_t ik3 = iq3 / rk3;
+
+        // v indices
+        const int64_t iv2 = iq2 / rv2;
+        const int64_t iv3 = iq3 / rv3;
+
+        device const float * mp[Q];
+
+        {
+            const int64_t ir = iq3*ne02*ne01 + iq2*ne01 + iq1;
+
+            for (int64_t j = 0; j < Q; ++j) {
+                if (iq1 + j < ne01) {
+                    mp[j] = (device const float *) (mask + ((ir + j)%ne31)*nb31);
+                } else {
+                    mp[j] = nullptr;
+                }
+            }
         }
 
-        device const half4 * pk4 = (device const half4 *) ((device char *) k + (ic*nb11 + ik2*nb12 + ik3*nb13));
-        device const half4 * pv4 = (device const half4 *) ((device char *) v + (ic*nb21 + iv2*nb22 + iv3*nb23));
-
-        half4 s4 = 0.0h;
-
-#pragma unroll
-        for (int64_t i = 0; i < D4/tph; ++i) {
-            s4 += pq4[hiisg*D4 + tph*i + tiih] * pk4[tph*i + tiih];
-        }
-
-        ss[hiisg*tph + tiih] = (s4.x + s4.y + s4.z + s4.w);
-
-        simdgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiih == 0) {
-            half s = 0.0h;
-
-#pragma unroll
-            for (int64_t i = 0; i < tph; ++i) {
-                s += ss[hiisg*tph + i];
+        for (int64_t iic = C*sgitg; iic < ne11; iic += C*nsg) {
+            {
+                bool skip = true;
+                for (int64_t j = 0; j < Q; ++j) {
+                    skip = skip && (mp[j][iic] == -INFINITY);
+                }
+                if (skip) {
+                    continue;
+                }
             }
 
-            s = s*scale + mv;
+            for (int p = 0; p < C; ++p) {
+                const int64_t ic = iic + p;
 
-            const half m = M;
+                device const half4 * pk4 = (device const half4 *) ((device char *) k + (ic*nb11 + ik2*nb12 + ik3*nb13));
 
-            M = max(M, s);
+                for (int64_t j = 0; j < Q; ++j) {
+                    half4 s4 = 0.0h;
 
-            const half ms = exp(m - M);
-            const half vs = exp(s - M);
+                    for (int64_t i = 0; i < D4/tph; ++i) {
+                        s4 += pq4[j*T4 + hiisg*D4 + tph*i + tiih]*pk4[tph*i + tiih];
+                    }
 
-            S = S*ms + vs;
+                    ss[j*T + 32*p + hiisg*tph + tiih] = s4.x + s4.y + s4.z + s4.w;
+                }
+            }
 
-            ss[2*hiisg + 0] = ms;
-            ss[2*hiisg + 1] = vs;
+            simdgroup_barrier(mem_flags::mem_none);
+
+            if (tiih < Q) {
+                const int64_t j = tiih;
+
+                for (int p = 0; p < C; ++p) {
+                    half4 s4 = 0.0h;
+
+                    for (int64_t i = 0; i < tph/4; ++i) {
+                        s4 += ss4[j*T4 + 8*p + hiisg*tph/4 + i];
+                    }
+
+                    half s = (s4.x + s4.y + s4.z + s4.w)*scale + mp[j][iic + p];
+
+                    const half m = M;
+
+                    M = max(M, s);
+
+                    const half ms = m == -INFINITY ? 0.0h : exp(m - M);
+                    const half vs = s == -INFINITY ? 0.0h : exp(s - M);
+
+                    S = S*ms + vs;
+
+                    ss[j*T + 32*p + 2*hiisg + 0] = ms;
+                    ss[j*T + 32*p + 2*hiisg + 1] = vs;
+                }
+            }
+
+            simdgroup_barrier(mem_flags::mem_none);
+
+            for (int p = 0; p < C; ++p) {
+                const int64_t ic = iic + p;
+                device const half4 * pv4 = (device const half4 *) ((device char *) v + (ic*nb21 + iv2*nb22 + iv3*nb23));
+
+                for (int64_t j = 0; j < Q; ++j) {
+                    const half ms = ss[j*T + 32*p + 2*hiisg + 0];
+                    const half vs = ss[j*T + 32*p + 2*hiisg + 1];
+
+                    for (int64_t i = 0; i < D4/tph; ++i) {
+                        ps4[j][i] = ps4[j][i]*ms + pv4[tph*i + tiih]*vs;
+                    }
+                }
+            }
         }
 
-        simdgroup_barrier(mem_flags::mem_threadgroup);
+        if (tiih < Q) {
+            const int64_t j = tiih;
 
-        const half ms = ss[2*hiisg + 0];
-        const half vs = ss[2*hiisg + 1];
-
-#pragma unroll
-        for (int64_t i = 0; i < D4/tph; ++i) {
-            ps4[hiisg*D4 + tph*i + tiih] = ps4[hiisg*D4 + tph*i + tiih]*ms + pv4[tph*i + tiih]*vs;
+            ss[j*T + 2*hiisg + 0] = S;
+            ss[j*T + 2*hiisg + 1] = M;
         }
     }
 
-    if (tiih == 0) {
-        ss[2*hiisg + 0] = S;
-        ss[2*hiisg + 1] = M;
-    }
-
     threadgroup_barrier(mem_flags::mem_threadgroup);
 
     // reduce the warps
-    if (sgitg == 0) {
-        for (int64_t sg = 1; sg < nsg; ++sg) {
-            const half S0 = ss[                2*hiisg + 0];
-            const half S1 = ss[sg*(R*D + 32) + 2*hiisg + 0];
+    //if (sgitg == 0) {
+    //    for (int64_t j = 0; j < Q; ++j) {
+    //        for (int64_t sg = 1; sg < nsg; ++sg) {
 
-            const half M0 = ss[                2*hiisg + 1];
-            const half M1 = ss[sg*(R*D + 32) + 2*hiisg + 1];
+    //            const half S0 = ss[j*T +            2*hiisg + 0];
+    //            const half S1 = ss[j*T + sg*(256) + 2*hiisg + 0];
 
-            M = max(M0, M1);
+    //            const half M0 = ss[j*T +            2*hiisg + 1];
+    //            const half M1 = ss[j*T + sg*(256) + 2*hiisg + 1];
 
-            const half ms0 = exp(M0 - M);
-            const half ms1 = exp(M1 - M);
+    //            M = max(M0, M1);
 
-            S = S0*ms0 + S1*ms1;
+    //            const half ms0 = exp(M0 - M);
+    //            const half ms1 = exp(M1 - M);
 
-            if (tiih == 0) {
-                ss[2*hiisg + 0] = S;
-                ss[2*hiisg + 1] = M;
-            }
+    //            S = S0*ms0 + S1*ms1;
 
-            for (int64_t i = 0; i < D4/tph; ++i) {
-                ps4[hiisg*D4 + tph*i + tiih] = ps4[hiisg*D4 + tph*i + tiih]*ms0 + ps4[sg*(R*D + 32)/4 + hiisg*D4 + tph*i + tiih]*ms1;
+    //            if (tiih == 0) {
+    //                ss[j*T + 2*hiisg + 0] = S;
+    //                ss[j*T + 2*hiisg + 1] = M;
+    //            }
+
+    //            for (int64_t i = 0; i < D4/tph; ++i) {
+    //                ps4[j*T4 + hiisg*D4 + tph*i + tiih] = ps4[j*T4 + hiisg*D4 + tph*i + tiih]*ms0 + ps4[j*T4 + sg*(256)/4 + hiisg*D4 + tph*i + tiih]*ms1;
+    //            }
+    //        }
+
+    //        for (int64_t i = 0; i < D4/tph; ++i) {
+    //            ps4[j*T4 + hiisg*D4 + tph*i + tiih] = ps4[j*T4 + hiisg*D4 + tph*i + tiih]/S;
+    //        }
+    //    }
+    //}
+
+    for (int64_t sg = 1; sg < nsg; ++sg) {
+        if (sgitg == sg) {
+            // store heads to shared memory - reuse pq4
+            for (int64_t j = 0; j < Q; ++j) {
+                for (int64_t i = 0; i < D4/tph; ++i) {
+                    pq4[j*T4 + hiisg*D4 + tph*i + tiih] = ps4[j][i];
+                }
             }
         }
 
-        for (int64_t i = 0; i < D4/tph; ++i) {
-            ps4[hiisg*D4 + tph*i + tiih] = ps4[hiisg*D4 + tph*i + tiih]/S;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (sgitg == 0) {
+            for (int64_t j = 0; j < Q; ++j) {
+                const half S0 = ss[j*T +             2*hiisg + 0];
+                const half S1 = ss[j*T + sg*(32*C) + 2*hiisg + 0];
+
+                const half M0 = ss[j*T +             2*hiisg + 1];
+                const half M1 = ss[j*T + sg*(32*C) + 2*hiisg + 1];
+
+                M = max(M0, M1);
+
+                const half ms0 = exp(M0 - M);
+                const half ms1 = exp(M1 - M);
+
+                S = S0*ms0 + S1*ms1;
+
+                if (tiih == 0) {
+                    ss[j*T + 2*hiisg + 0] = S;
+                    ss[j*T + 2*hiisg + 1] = M;
+                }
+
+                for (int64_t i = 0; i < D4/tph; ++i) {
+                    ps4[j][i] = ps4[j][i]*ms0 + pq4[j*T4 + hiisg*D4 + tph*i + tiih]*ms1;
+                }
+            }
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    if (sgitg == 0) {
+        for (int64_t j = 0; j < Q; ++j) {
+            S = ss[j*T + 2*hiisg + 0];
+            for (int64_t i = 0; i < D4/tph; ++i) {
+                ps4[j][i] = ps4[j][i]/S;
+            }
         }
     }
 
     simdgroup_barrier(mem_flags::mem_threadgroup);
 
-    // dst indices
-    const int64_t i1 = iq1;
-    const int64_t i2 = iq2;
-    const int64_t i3 = iq3;
-
     device float4 * dst4 = (device float4 *) dst;
 
     if (sgitg == 0) {
-        for (int64_t i = 0; i < D4/tph; ++i) {
-            dst4[(i3*ne2*ne1 + i2 + i1*ne1)*D4 + tph*i + tiih] = (float4) ps4[hiisg*D4 + tph*i + tiih];
+        for (int64_t j = 0; j < Q && iq1 + j < ne01; ++j) {
+            for (int64_t i = 0; i < D4/tph; ++i) {
+                dst4[(iq3*ne2*ne1 + iq2 + (iq1 + j)*ne1)*D4 + tph*i + tiih] = (float4) ps4[j][i];
+            }
         }
     }
 }
 
-template [[host_name("kernel_flash_attn_ext_f16_h64" )]] kernel flash_attn_ext_f16_t kernel_flash_attn_ext_f16<64,  2>;
-template [[host_name("kernel_flash_attn_ext_f16_h80" )]] kernel flash_attn_ext_f16_t kernel_flash_attn_ext_f16<80,  2>;
-template [[host_name("kernel_flash_attn_ext_f16_h128")]] kernel flash_attn_ext_f16_t kernel_flash_attn_ext_f16<128, 2>;
+template [[host_name("kernel_flash_attn_ext_f16_h64" )]] kernel flash_attn_ext_f16_t kernel_flash_attn_ext_f16<64,  2, 2, 8>;
+template [[host_name("kernel_flash_attn_ext_f16_h80" )]] kernel flash_attn_ext_f16_t kernel_flash_attn_ext_f16<80,  2, 2, 8>;
+template [[host_name("kernel_flash_attn_ext_f16_h128")]] kernel flash_attn_ext_f16_t kernel_flash_attn_ext_f16<128, 2, 2, 8>;
 
 kernel void kernel_cpy_f16_f16(
         device  const half * src0,

tests/test-backend-ops.cpp

@@ -1397,7 +1397,7 @@ struct test_flash_attn_ext : public test_case {
     }
 
     double max_nmse_err() override {
-        return 5e-4;
+        return 5e-5;
     }
 
     test_flash_attn_ext(ggml_type typeq = GGML_TYPE_F16,
@@ -1680,7 +1680,9 @@ static bool test_backend(ggml_backend_t backend, test_mode mode, const char * op
     test_cases.emplace_back(new test_pad());
     test_cases.emplace_back(new test_leaky_relu());
 
-    test_cases.emplace_back(new test_flash_attn_ext(GGML_TYPE_F16, 128, 32, 96, 8));
+    test_cases.emplace_back(new test_flash_attn_ext(GGML_TYPE_F16, 128, 32, 256, 8));
+    test_cases.emplace_back(new test_flash_attn_ext(GGML_TYPE_F16, 128, 32, 256, 7));
+    test_cases.emplace_back(new test_flash_attn_ext(GGML_TYPE_F16, 128, 32, 256, 1));
 
 #if !defined(__SANITIZE_THREAD__)
     // FIXME: these tests use too much memory with thread sanitizer