char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Neo Zhangsycl : fix error Error OP FLASH_ATTN_EXT on arc770 (#26441)eef5f3e34

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
27.9 KiB684 linesraw
1#ifndef GGML_SYCL_FATTN_VEC_HPP
2#define GGML_SYCL_FATTN_VEC_HPP
3
4#include <sycl/sycl.hpp>
5#include <sycl/ext/oneapi/work_group_static.hpp>
6#include <iostream>
7#include <iomanip>
8
9#include "dpct/helper.hpp"
10#include "common.hpp"
11#include "ggml.h"
12#include "fattn-common.hpp"
13#include <cmath>
14#include <float.h>
15
16namespace syclex = sycl::ext::oneapi::experimental;
17
18static int ggml_sycl_fattn_vec_get_nthreads_device(gpu_arch arch) {
19    // Xe2 (Battlemage, Lunar Lake) runs the flash-attention vec kernel best with a 256-thread work group.
20    return (arch == gpu_arch::intel_gpu_bmg_g21 ||
21            arch == gpu_arch::intel_gpu_bmg_g31 ||
22            arch == gpu_arch::intel_gpu_lnl_m) ? 256 : 128;
23}
24
25// Currenlty llvm with the amdgcn target dose not support unrolling loops
26// that contain a break that can not be resolved at compile time.
27#ifdef __clang__
28#pragma clang diagnostic push
29#pragma clang diagnostic ignored "-Wpass-failed"
30#endif // __clang__
31
32template <int D,
33          int ncols,
34          int type_K,
35          int type_V,
36          bool use_logit_softcap,
37          int warp_size,
38          int nthreads>  // D == head size
39static void flash_attn_ext_vec(const char* __restrict__ Q,
40                        const char* __restrict__ K,
41                        const char* __restrict__ V,
42                        const char* __restrict__ mask,
43                        const char* __restrict__ sinks,
44                        const int* __restrict__ KV_max,
45                        float* __restrict__ dst,
46                        sycl::float2* __restrict__ dst_meta,
47                        const float scale,
48                        const float max_bias,
49                        const float m0,
50                        const float m1,
51                        const uint32_t n_head_log2,
52                        const float logit_softcap,
53                        const int32_t ne00,
54                        const sycl::uint3 ne01,
55                        const int32_t ne02,
56                        const int32_t ne03,
57                        const int32_t nb01,
58                        const int32_t nb02,
59                        const int32_t nb03,
60                        const int32_t ne10,
61                        const int32_t ne11,
62                        const int32_t ne12,
63                        const int32_t ne13,
64                        const int32_t nb11,
65                        const int32_t nb12,
66                        const int64_t nb13,
67                        const int32_t nb21,
68                        const int32_t nb22,
69                        const int64_t nb23,
70                        const int32_t ne31,
71                        const int32_t ne32,
72                        const int32_t ne33,
73                        const int32_t nb31,
74                        const int32_t nb32,
75                        const int64_t nb33) {
76
77#ifdef SYCL_FLASH_ATTN
78    // Skip unused kernel variants for faster compilation:
79
80    auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
81    if (use_logit_softcap && !(D == 128 || D == 256)) {
82        GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale,
83            max_bias, m0, m1, n_head_log2, logit_softcap,
84            ne00, ne01, ne02, ne03,
85                  nb01, nb02, nb03,
86            ne10, ne11, ne12, ne13,
87                  nb11, nb12, nb13,
88                  nb21, nb22, nb23,
89                  ne31, ne32, ne33,
90                  nb31, nb32, nb33);
91        return;
92    }
93
94    //In this kernel Q, K, V are matrices while i, j, k are matrix indices.
95
96    constexpr int cpy_nb = ggml_sycl_get_max_cpy_bytes();
97    constexpr int cpy_ne = cpy_nb / 4;
98
99    constexpr int nthreads_KQ_q = (D/4 < warp_size ? D/4 : warp_size);
100    constexpr int nthreads_V_q  = (D/4 < warp_size ? D/4 : warp_size);
101
102    constexpr int nthreads_KQ = type_K == GGML_TYPE_F16 ? 128 / cpy_nb : nthreads_KQ_q;
103    constexpr int nthreads_V  = type_V == GGML_TYPE_F16 ? 128 / cpy_nb : nthreads_V_q;
104
105    static_assert(warp_size % nthreads_KQ == 0, "bad nthreads_K");
106    static_assert(warp_size % nthreads_V  == 0, "bad nthreads_V");
107
108    constexpr int V_rows_per_thread = type_V == GGML_TYPE_F16 ? 2*cpy_ne : 4;
109    constexpr int V_cols_per_iter   = warp_size / nthreads_V;
110
111    constexpr vec_dot_KQ_t vec_dot_KQ = get_vec_dot_KQ<type_K, D, nthreads_KQ, warp_size>();
112    constexpr bool Q_q8_1 = type_K != GGML_TYPE_F16;
113#ifdef GGML_SYCL_F16
114    constexpr dequantize_V_t dequantize_V = get_dequantize_V<type_V, sycl::half, V_rows_per_thread>();
115#else
116    constexpr dequantize_V_t dequantize_V = get_dequantize_V<type_V, float, V_rows_per_thread>();
117#endif // GGML_SYCL_F16
118
119    const int ic0 = item_ct1.get_group(2) * ncols;  // Index of the Q/QKV column to work on.
120
121    const int sequence  = item_ct1.get_group(0) / ne02;
122    const int head      = item_ct1.get_group(0) - sequence * ne02;
123    const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix.
124    Q += nb03*sequence + nb02* head              + nb01*ic0;
125    K += nb13*sequence + nb12*(head / gqa_ratio);
126    V += nb23*sequence + nb22*(head / gqa_ratio);
127
128    const sycl::half * maskh = (const sycl::half *) (mask + nb33 * (sequence % ne33) + nb31 * ic0);
129
130    const float slope = get_alibi_slope(max_bias, head, n_head_log2, m0, m1);
131
132    static_assert(D % (2*warp_size) == 0, "D not divisible by 2*warp_size == 64.");
133    constexpr int nwarps = nthreads / warp_size;
134    const int     tid    = warp_size * item_ct1.get_local_id(1) + item_ct1.get_local_id(2);
135    __builtin_assume(tid < nthreads);
136
137    constexpr int ne_KQ      = ncols*D;
138    constexpr int ne_combine = nwarps*V_cols_per_iter*D;
139
140    constexpr size_t lsm_size1 = ncols * warp_size;
141    constexpr size_t lsm_size2 = ncols * warp_size;
142#ifdef GGML_SYCL_F16
143    sycl::half2 VKQ[ncols][(D / 2) / nthreads_V] = { { { 0.0f, 0.0f } } };
144    constexpr size_t lsm_size3 = (ne_KQ > ne_combine ? ne_KQ : ne_combine);
145    constexpr size_t local_share_mem_size = (lsm_size1 + lsm_size2)*sizeof(float) + lsm_size3*sizeof(sycl::half);
146
147    syclex::work_group_static<char[local_share_mem_size]> lsm;
148
149    float *KQ_max_shared = (float *)&lsm;
150    float *KQ_sum_shared = KQ_max_shared+lsm_size1;
151    sycl::half* KQ = (sycl::half*)(KQ_sum_shared + lsm_size2);
152
153
154#else
155    sycl::float2 VKQ[ncols][(D/2)/nthreads_V] = {{{0.0f, 0.0f}}};
156
157    constexpr size_t lsm_size3 = (ne_KQ > ne_combine ? ne_KQ : ne_combine);
158    constexpr size_t local_share_mem_size = (lsm_size1 + lsm_size2 + lsm_size3)*sizeof(float);
159
160
161    syclex::work_group_static<char[local_share_mem_size]> lsm;
162    float *KQ_max_shared = (float *)&lsm;
163    float *KQ_sum_shared = KQ_max_shared+lsm_size1;
164    float* KQ = KQ_sum_shared + lsm_size2;
165
166#endif // GGML_SYCL_F16
167
168    float KQ_max[ncols];
169    float KQ_sum[ncols];
170#pragma unroll
171    for (int j = 0; j < ncols; ++j) {
172        KQ_max[j] = -FLT_MAX/2.0f;
173        KQ_sum[j] = 0.0f;
174    }
175
176    // Convert Q to float2 (f16 K) or q8_1 (quantized K) and store in registers:
177#ifdef GGML_SYCL_F16
178    sycl::half2 Q_reg[ncols][(D / 2) / nthreads_KQ] = {{{0.0f, 0.0f}}};  // Will be initialized completely.
179#else
180    sycl::float2 Q_reg[ncols][(D/2)/nthreads_KQ] = {{{0.0f, 0.0f}}}; // May be only partially initialized.
181#endif // GGML_SYCL_F16
182    int    Q_i32[ncols][1 > D/(sizeof(int)*nthreads_KQ) ? 1 : D/(sizeof(int)*nthreads_KQ)];
183    sycl::float2 Q_ds[ncols][1 > D / (sizeof(int) * nthreads_KQ) ? 1 : D / (sizeof(int) * nthreads_KQ)];
184    if constexpr (Q_q8_1) {
185#pragma unroll
186        for (int j0 = 0; j0 < ncols; j0 += nwarps) {
187            const int j = j0 + item_ct1.get_local_id(1);
188
189            if (j0 + nwarps > ncols && j >= ncols) {
190                break;
191            }
192
193            // Reuse KQ as temporary storage for converting Q to q8_1:
194            int    * tmp_q_i32 = (int    *) &KQ[j*D];
195            sycl::float2 * tmp_q_ds  = (sycl::float2 *) (tmp_q_i32 + D / sizeof(int));
196
197            // Set memory to zero if out of bounds:
198            if (ncols > 1 && ic0 + j >= int(ne01.z())) {
199#pragma unroll
200                for (int i0 = 0; i0 < int(D/sizeof(int)); i0 += warp_size) {
201                    const int i = i0 + item_ct1.get_local_id(2);
202
203                    if (i0 + warp_size <= int(D/sizeof(int)) || i < int(D/sizeof(int))) {
204                        tmp_q_i32[i] = 0;
205                    }
206                }
207                if (item_ct1.get_local_id(2) < D/QK8_1) {
208                    tmp_q_ds[item_ct1.get_local_id(2)] = sycl::float2(0.0f, 0.0f);
209                }
210            } else {
211                const float * Q_f = (const float *) (Q + j*nb01);
212                constexpr int nthreads_quantize = D/sizeof(int) < warp_size ? D/sizeof(int) : warp_size;
213#pragma unroll
214                for (int i0 = 0; i0 < int(D/sizeof(int)); i0 += nthreads_quantize) {
215                    quantize_q8_1_to_shared<sycl::float2, nthreads_quantize, warp_size>
216                        (Q_f + i0*sizeof(int), scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1);
217                }
218            }
219        }
220
221
222        item_ct1.barrier(sycl::access::fence_space::local_space);
223
224#pragma unroll
225        for (int j = 0; j < ncols; ++j) {
226            int    * tmp_q_i32 = (int    *) &KQ[j*D];
227            sycl::float2 * tmp_q_ds  = (sycl::float2 *) (tmp_q_i32 + D / sizeof(int));
228
229#pragma unroll
230            for (int i0 = 0; i0 < int(D/sizeof(int)); i0 += nthreads_KQ) {
231                const int i =
232                    i0 + (nthreads_KQ == warp_size ? item_ct1.get_local_id(2) : item_ct1.get_local_id(2) % nthreads_KQ);
233
234                Q_i32[j][i0/nthreads_KQ] = tmp_q_i32[i];
235                Q_ds[j][i0/nthreads_KQ]  = tmp_q_ds[i/QI8_1];
236            }
237        }
238
239        item_ct1.barrier(sycl::access::fence_space::local_space);
240
241    } else {
242#ifdef GGML_SYCL_F16
243        const sycl::half2 scale_h2 = sycl::half2(scale, scale);
244#pragma unroll
245        for (int j = 0; j < ncols; ++j) {
246            const sycl::float2 * Q_j = (const sycl::float2 *) (Q + j * nb01);
247#pragma unroll
248            for (int i0 = 0; i0 < D/2; i0 += nthreads_KQ*cpy_ne) {
249                const int i = i0 + (nthreads_KQ == warp_size ? item_ct1.get_local_id(2) :
250                                                               item_ct1.get_local_id(2) % nthreads_KQ) *
251                                       cpy_ne;
252
253                sycl::float2 tmp[cpy_ne] = {
254                    { 0.0f, 0.0f }
255                };
256                if (ncols == 1 || ic0 + j < int(ne01.z())) {
257                    ggml_sycl_memcpy_1<cpy_nb>(tmp,            &Q_j[i]);
258                    ggml_sycl_memcpy_1<cpy_nb>(tmp + cpy_ne/2, &Q_j[i + cpy_ne/2]);
259                }
260#pragma unroll
261                for (int i1 = 0; i1 < cpy_ne; ++i1) {
262                    Q_reg[j][i0 / nthreads_KQ + i1] = sycl::half2(tmp[i1].x(), tmp[i1].y());
263                }
264            }
265#pragma unroll
266            for (int k = 0; k < (D/2)/nthreads_KQ; ++k) {
267                Q_reg[j][k] *= scale_h2;
268            }
269        }
270#else
271#pragma unroll
272        for (int j = 0; j < ncols; ++j) {
273            const sycl::float2 * Q_j = (const sycl::float2 *) (Q + j*nb01);
274#pragma unroll
275            for (int i0 = 0; i0 < D/2; i0 += nthreads_KQ*cpy_ne) {
276                const int i = i0 + (nthreads_KQ == warp_size ? item_ct1.get_local_id(2) : item_ct1.get_local_id(2) % nthreads_KQ)*cpy_ne;
277                if (ncols == 1 || ic0 + j < int(ne01.z())) {
278                    ggml_sycl_memcpy_1<cpy_nb>(&Q_reg[j][i0/nthreads_KQ],            &Q_j[i]);
279                    ggml_sycl_memcpy_1<cpy_nb>(&Q_reg[j][i0/nthreads_KQ + cpy_ne/2], &Q_j[i + cpy_ne/2]);
280                }
281            }
282#pragma unroll
283            for (int k = 0; k < (D/2)/nthreads_KQ; ++k) {
284                Q_reg[j][k].x() *= scale;
285                Q_reg[j][k].y() *= scale;
286            }
287        }
288#endif // GGML_SYCL_F16
289    }
290
291    const int k_VKQ_max = KV_max ? KV_max[sequence * item_ct1.get_group_range(2) + item_ct1.get_group(2)] : ne11;
292    K += item_ct1.get_group(1) * nthreads * nb11;
293    V += item_ct1.get_group(1) * nthreads * nb21;
294    maskh += item_ct1.get_group(1) * nthreads;
295    for (int k_VKQ_0 = item_ct1.get_group(1) * nthreads; k_VKQ_0 < k_VKQ_max;
296         k_VKQ_0 += item_ct1.get_group_range(1) * nthreads,
297             // Increment pointers after each loop:
298         K += item_ct1.get_group_range(1) * nthreads * nb11, V += item_ct1.get_group_range(1) * nthreads * nb21,
299             maskh += item_ct1.get_group_range(1) * nthreads) {
300        // Calculate KQ tile and keep track of new maximum KQ values:
301        float KQ_reg[ncols]={}; // KQ in registers.
302        float KQ_max_new[ncols]={};
303
304
305#pragma unroll
306        for (int j = 0; j < ncols; ++j) {
307            KQ_max_new[j] = KQ_max[j];
308        }
309
310#pragma unroll
311        for (int i_KQ_0 = 0; i_KQ_0 < nthreads_KQ; ++i_KQ_0) {
312            const int i_KQ = item_ct1.get_local_id(1) * warp_size +
313                             (nthreads_KQ == warp_size ? 0 : (item_ct1.get_local_id(2) & ~(nthreads_KQ - 1))) + i_KQ_0;
314
315#pragma unroll
316            for (int j = 0; j < ncols; ++j) {
317                float sum = vec_dot_KQ(K + i_KQ*nb11, Q_reg[j], Q_i32[j], Q_ds[j]);
318                sum = warp_reduce_sum<nthreads_KQ>(sum);
319
320                if (use_logit_softcap) {
321                    sum = logit_softcap * sycl::tanh(sum);
322                }
323                if (mask) {
324                    sum += slope * sycl::vec<sycl::half, 1>(maskh[j * ne11 + i_KQ])
325                                       .convert<float, sycl::rounding_mode::automatic>()[0];
326                }
327
328                KQ_max_new[j] = sycl::fmax((float) KQ_max_new[j], sum);
329
330                if (int(nthreads_KQ == warp_size ? item_ct1.get_local_id(2)
331                                                 : item_ct1.get_local_id(2) %
332                                                       nthreads_KQ) == i_KQ_0) {
333                  KQ_reg[j] = sum;
334                }
335            }
336        }
337
338#pragma unroll
339        for (int j = 0; j < ncols; ++j) {
340#pragma unroll
341            for (int offset = nthreads_KQ; offset < warp_size; offset <<= 1) {
342               KQ_max_new[j] = sycl::fmax(
343                  (float)KQ_max_new[j],
344                  (float)dpct::permute_sub_group_by_xor(
345                      sycl::ext::oneapi::this_work_item::get_sub_group(),
346                      KQ_max_new[j],
347                      offset,
348                      warp_size));
349            }
350            const float KQ_max_scale = sycl::native::exp((float) (KQ_max[j] - KQ_max_new[j]));
351            KQ_max[j] = KQ_max_new[j];
352
353            KQ_reg[j]            = sycl::native::exp((float) (KQ_reg[j] - KQ_max[j]));
354            KQ_sum[j] = KQ_sum[j]*KQ_max_scale + KQ_reg[j];
355            KQ[j*nthreads + tid] = KQ_reg[j];
356
357#ifdef GGML_SYCL_F16
358            const sycl::half2 KQ_max_scale_h2 = sycl::half2(KQ_max_scale, KQ_max_scale);
359#pragma unroll
360            for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) {
361                VKQ[j][i_VKQ_0/nthreads_V] *= KQ_max_scale_h2;
362            }
363#else
364#pragma unroll
365            for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) {
366                VKQ[j][i_VKQ_0/nthreads_V].x() *= KQ_max_scale;
367                VKQ[j][i_VKQ_0/nthreads_V].y() *= KQ_max_scale;
368            }
369#endif // GGML_SYCL_F16
370        }
371
372        sycl::group_barrier(sycl::ext::oneapi::this_work_item::get_sub_group());
373
374#pragma unroll
375        for (int k0 = 0; k0 < warp_size; k0 += V_cols_per_iter) {
376            const int k = item_ct1.get_local_id(1) * warp_size + k0 +
377                          (nthreads_V == warp_size ? 0 : item_ct1.get_local_id(2) / nthreads_V);
378
379#ifdef GGML_SYCL_F16
380            sycl::half2 KQ_k[ncols];
381#pragma unroll
382            for (int j = 0; j < ncols; ++j) {
383                KQ_k[j] = sycl::half2(KQ[j * nthreads + k]);
384            }
385#pragma unroll
386            for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) {
387                sycl::half2 tmp[V_rows_per_thread / 2];
388                dequantize_V(V + k * nb21, tmp,
389                             2 * i_VKQ_0 + (nthreads_V == warp_size ? item_ct1.get_local_id(2) :
390                                                                      item_ct1.get_local_id(2) % nthreads_V) *
391                                               V_rows_per_thread);
392#pragma unroll
393                for (int i_VKQ_1 = 0; i_VKQ_1 < V_rows_per_thread/2; ++i_VKQ_1) {
394#pragma unroll
395                    for (int j = 0; j < ncols; ++j) {
396                        VKQ[j][i_VKQ_0/nthreads_V + i_VKQ_1] += tmp[i_VKQ_1]*KQ_k[j];
397                    }
398                }
399            }
400#else
401            float KQ_k[ncols];
402#pragma unroll
403            for (int j = 0; j < ncols; ++j) {
404                KQ_k[j] = KQ[j*nthreads + k];
405            }
406#pragma unroll
407            for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) {
408                sycl::float2 tmp[V_rows_per_thread/2];
409                dequantize_V(V + k*nb21, tmp,
410                    2*i_VKQ_0 + (nthreads_V == warp_size ? item_ct1.get_local_id(2) : item_ct1.get_local_id(2) % nthreads_V)*V_rows_per_thread);
411#pragma unroll
412                for (int i_VKQ_1 = 0; i_VKQ_1 < V_rows_per_thread/2; ++i_VKQ_1) {
413#pragma unroll
414                    for (int j = 0; j < ncols; ++j) {
415                        VKQ[j][i_VKQ_0/nthreads_V + i_VKQ_1].x() += tmp[i_VKQ_1].x()*KQ_k[j];
416                        VKQ[j][i_VKQ_0/nthreads_V + i_VKQ_1].y() += tmp[i_VKQ_1].y()*KQ_k[j];
417                    }
418                }
419            }
420#endif // GGML_SYCL_F16
421        }
422    }
423
424    if (sinks && item_ct1.get_group(1) == 0) {
425        const float sink = ((const float *) sinks)[head];
426
427#pragma unroll
428        for (int j0 = 0; j0 < ncols; j0 += nwarps) {
429            const int j = j0 + item_ct1.get_local_id(1);
430
431            if (j0 + nwarps > ncols && j >= ncols) {
432                break;
433            }
434            const float kqmax_new_j  = sycl::fmax(sink, (float) KQ_max[j]);
435            const float KQ_max_scale = sycl::native::exp((float) (KQ_max[j] - kqmax_new_j));
436            KQ_max[j] = kqmax_new_j;
437
438            KQ_sum[j] = KQ_sum[j] * KQ_max_scale +
439                        (item_ct1.get_local_id(2) == 0 ? sycl::native::exp((float) (sink - KQ_max[j])) : 0.0f);
440#ifdef GGML_SYCL_F16
441            const sycl::half2 KQ_max_scale_h2 = sycl::half2(KQ_max_scale, KQ_max_scale);
442#pragma unroll
443            for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) {
444                VKQ[j][i_VKQ_0/nthreads_V] *= KQ_max_scale_h2;
445            }
446#else
447#pragma unroll
448            for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) {
449                VKQ[j][i_VKQ_0/nthreads_V].x() *= KQ_max_scale;
450                VKQ[j][i_VKQ_0/nthreads_V].y() *= KQ_max_scale;
451            }
452#endif // GGML_SYCL_F16
453        }
454    }
455
456#pragma unroll
457    for (int j = 0; j < ncols; ++j) {
458        if (item_ct1.get_local_id(1) == 0) {
459            KQ_max_shared[j*warp_size+item_ct1.get_local_id(2)] = -FLT_MAX / 2.0f;
460            KQ_sum_shared[j*warp_size+item_ct1.get_local_id(2)] = 0.0f;
461        }
462    }
463
464    item_ct1.barrier(sycl::access::fence_space::local_space);
465
466#pragma unroll
467    for (int j = 0; j < ncols; ++j) {
468        if (item_ct1.get_local_id(2) == 0) {
469            KQ_max_shared[j*warp_size+item_ct1.get_local_id(1)] = KQ_max[j];
470        }
471    }
472
473    item_ct1.barrier(sycl::access::fence_space::local_space);
474
475#pragma unroll
476    for (int j_VKQ = 0; j_VKQ < ncols; ++j_VKQ) {
477        if (ncols > 1 && ic0 + j_VKQ >= int(ne01.z())) {
478            break;
479        }
480
481        float kqmax_new         = KQ_max_shared[j_VKQ*warp_size+item_ct1.get_local_id(2)];
482        kqmax_new = warp_reduce_max<warp_size>(kqmax_new);
483        const float kqmax_scale = sycl::native::exp((float) (KQ_max[j_VKQ] - kqmax_new));
484        KQ_max[j_VKQ] = kqmax_new;
485
486#ifdef GGML_SYCL_F16
487        sycl::half2 * VKQ_tmp = (sycl::half2 *) KQ + item_ct1.get_local_id(1) * (V_cols_per_iter * D / 2) +
488                                (nthreads_V == warp_size ? 0 : item_ct1.get_local_id(2) / nthreads_V) * (D / 2);
489
490        const sycl::half2 kqmax_scale_h2 = sycl::half2(kqmax_scale, kqmax_scale);
491#pragma unroll
492        for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) {
493            VKQ[j_VKQ][i_VKQ_0/nthreads_V] *= kqmax_scale_h2;
494        }
495#pragma unroll
496        for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) {
497            const int i_VKQ =
498                i_VKQ_0 + (nthreads_V == warp_size ? item_ct1.get_local_id(2) : item_ct1.get_local_id(2) % nthreads_V) *
499                              (V_rows_per_thread / 2);
500
501            ggml_sycl_memcpy_1<V_rows_per_thread * sizeof(sycl::half)>(VKQ_tmp + i_VKQ,
502                                                                       &VKQ[j_VKQ][i_VKQ_0 / nthreads_V]);
503        }
504#else
505        sycl::float2 * VKQ_tmp = (sycl::float2 *) KQ + item_ct1.get_local_id(1)*(V_cols_per_iter*D/2)
506            + (nthreads_V == warp_size ? 0 : item_ct1.get_local_id(2) / nthreads_V)*(D/2);
507#pragma unroll
508        for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) {
509            VKQ[j_VKQ][i_VKQ_0/nthreads_V].x() *= kqmax_scale;
510            VKQ[j_VKQ][i_VKQ_0/nthreads_V].y() *= kqmax_scale;
511        }
512#pragma unroll
513        for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) {
514            const int i_VKQ = i_VKQ_0 + (nthreads_V == warp_size ? item_ct1.get_local_id(2) : item_ct1.get_local_id(2) % nthreads_V)*(V_rows_per_thread/2);
515
516            ggml_sycl_memcpy_1<V_rows_per_thread/2*sizeof(float)>(VKQ_tmp + i_VKQ,                       &VKQ[j_VKQ][i_VKQ_0/nthreads_V]);
517            ggml_sycl_memcpy_1<V_rows_per_thread/2*sizeof(float)>(VKQ_tmp + i_VKQ + V_rows_per_thread/4, &VKQ[j_VKQ][i_VKQ_0/nthreads_V + V_rows_per_thread/4]);
518        }
519#endif // GGML_SYCL_F16
520
521        KQ_sum[j_VKQ] *= kqmax_scale;
522        KQ_sum[j_VKQ] = warp_reduce_sum<warp_size>(KQ_sum[j_VKQ]);
523        if (item_ct1.get_local_id(2) == 0) {
524            KQ_sum_shared[j_VKQ*warp_size+item_ct1.get_local_id(1)] = KQ_sum[j_VKQ];
525        }
526
527        item_ct1.barrier(sycl::access::fence_space::local_space);
528
529
530        if (nthreads <= D || tid < D) {
531            KQ_sum[j_VKQ] = KQ_sum_shared[j_VKQ*warp_size+item_ct1.get_local_id(2)];
532            KQ_sum[j_VKQ] = warp_reduce_sum<warp_size>(KQ_sum[j_VKQ]);
533
534#pragma unroll
535            for (int i0 = 0; i0 < D; i0 += nthreads) {
536                float dst_val = 0;
537#pragma unroll
538                for (int w = 0; w < nwarps; ++w) {
539#pragma unroll
540                    for (int v = 0; v < V_cols_per_iter; ++v) {
541                        dst_val += float(KQ[w*V_cols_per_iter*D + v*D + i0 + tid]);
542                    }
543                }
544                if (item_ct1.get_group_range(1) == 1) {
545                    dst_val /= KQ_sum[j_VKQ];
546                }
547                dst[(((sequence * int(ne01.z()) + ic0 + j_VKQ) * ne02 + head) * item_ct1.get_group_range(1) +
548                     item_ct1.get_group(1)) *
549                        D +
550                    i0 + tid] = dst_val;
551            }
552        }
553
554        if (j_VKQ < ncols-1) {
555            item_ct1.barrier(sycl::access::fence_space::local_space);
556        }
557
558    }
559
560    if (item_ct1.get_group_range(1) != 1 && tid < ncols && (ncols == 1 || ic0 + tid < int(ne01.z()))) {
561        dst_meta[((sequence * int(ne01.z()) + ic0 + tid) * ne02 + head) * item_ct1.get_group_range(1) +
562                 item_ct1.get_group(1)] = make_float2(KQ_max[tid], KQ_sum[tid]);
563    }
564#else
565    GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale,
566        max_bias, m0, m1, n_head_log2, logit_softcap,
567        ne00, ne01, ne02, ne03,
568              nb01, nb02, nb03,
569        ne10, ne11, ne12, ne13,
570              nb11, nb12, nb13,
571              nb21, nb22, nb23,
572              ne31, ne32, ne33,
573              nb31, nb32, nb33);
574
575#endif // SYCL_FLASH_ATTN
576}
577#ifdef __clang__
578#pragma clang diagnostic pop
579#endif // __clang__
580
581
582
583template <int D, int cols_per_block, int type_K, int type_V, bool use_logit_softcap>
584void ggml_sycl_flash_attn_ext_vec_case_impl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
585
586    constexpr int warp_size = WARP_16_SIZE; //better performance than WARP_32_SIZE
587
588    const bool need_f16_K = type_K == GGML_TYPE_F16;
589    const bool need_f16_V = type_V == GGML_TYPE_F16;
590    constexpr size_t nbytes_shared = 0;
591
592    const auto arch = ggml_sycl_info().devices[ctx.device].hw_info.arch;
593    const int nthreads = ggml_sycl_fattn_vec_get_nthreads_device(arch);
594    if constexpr (D <= 256) {
595        if (nthreads == 256) {
596            constexpr int nthreads_hw = 256;
597            constexpr int nwarps = nthreads_hw / warp_size;
598            launch_fattn<D, cols_per_block, 1,
599                         flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
600                                            use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
601                ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
602            return;
603        }
604    }
605
606    constexpr int nthreads_hw = 128;
607    constexpr int nwarps = nthreads_hw / warp_size;
608    launch_fattn<D, cols_per_block, 1,
609                 flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
610                                    use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
611        ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
612}
613
614template <int D, int type_K, int type_V>
615void ggml_sycl_flash_attn_ext_vec_case(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
616    const ggml_tensor * KQV = dst;
617    const ggml_tensor * Q   = dst->src[0];
618
619    float logit_softcap;
620    memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
621
622    if (Q->ne[1] == 1) {
623        constexpr int cols_per_block = 1;
624        if (logit_softcap == 0.0f) {
625            constexpr bool use_logit_softcap = false;
626            ggml_sycl_flash_attn_ext_vec_case_impl<D, cols_per_block, type_K, type_V, use_logit_softcap>(ctx, dst);
627        } else {
628            constexpr bool use_logit_softcap = true;
629            ggml_sycl_flash_attn_ext_vec_case_impl<D, cols_per_block, type_K, type_V, use_logit_softcap>(ctx, dst);
630        }
631        return;
632    }
633
634    constexpr int cols_per_block = 2;
635    if (logit_softcap == 0.0f) {
636        constexpr bool use_logit_softcap = false;
637        ggml_sycl_flash_attn_ext_vec_case_impl<D, cols_per_block, type_K, type_V, use_logit_softcap>(ctx, dst);
638    } else {
639        constexpr bool use_logit_softcap = true;
640        ggml_sycl_flash_attn_ext_vec_case_impl<D, cols_per_block, type_K, type_V, use_logit_softcap>(ctx, dst);
641    }
642}
643
644#define DECL_FATTN_VEC_CASE(D, type_K, type_V)                              \
645    template void ggml_sycl_flash_attn_ext_vec_case                         \
646    <D, type_K, type_V>(ggml_backend_sycl_context & ctx, ggml_tensor * dst) \
647
648#define EXTERN_DECL_FATTN_VEC_CASES(D, type_K)             \
649    extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_F16);  \
650    extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q4_0); \
651    extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q4_1); \
652    extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_0); \
653    extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_1); \
654    extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q8_0); \
655
656EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_F16)
657EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q4_0)
658EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q4_1)
659EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q5_0)
660EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q5_1)
661EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q8_0)
662
663EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_F16)
664EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q4_0)
665EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q4_1)
666EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q5_0)
667EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q5_1)
668EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q8_0)
669
670EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_F16)
671EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q4_0)
672EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q4_1)
673EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q5_0)
674EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q5_1)
675EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q8_0)
676
677EXTERN_DECL_FATTN_VEC_CASES(512, GGML_TYPE_F16)
678EXTERN_DECL_FATTN_VEC_CASES(512, GGML_TYPE_Q4_0)
679EXTERN_DECL_FATTN_VEC_CASES(512, GGML_TYPE_Q4_1)
680EXTERN_DECL_FATTN_VEC_CASES(512, GGML_TYPE_Q5_0)
681EXTERN_DECL_FATTN_VEC_CASES(512, GGML_TYPE_Q5_1)
682EXTERN_DECL_FATTN_VEC_CASES(512, GGML_TYPE_Q8_0)
683
684#endif // GGML_SYCL_FATTN_VEC_HPP