hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

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

Akarshan Biswassycl : add flash-attn support for head size 512 (#21654)b54cb2e3d2026-04-09

1#include <sycl/sycl.hpp>
2#include <sycl/ext/oneapi/work_group_static.hpp>
3#include "dpct/helper.hpp"
4#include "common.hpp"
5#include "fattn-common.hpp"
6
7#include <cmath>
8#include <float.h>
9
10namespace syclex = sycl::ext::oneapi::experimental;
11
12#define GGML_SYCL_FATTN_TILE_CONFIG_CASE(DKQ_, DV_, ncols_, nthreads, occupancy, nbatch_fa, nbatch_K) \
13    if (DKQ == (DKQ_) && DV == (DV_) && ncols == (ncols_)) {                                          \
14        static_assert((nthreads)          <= 512, "bad nthreads");                                    \
15        static_assert((occupancy)         <=   8, "bad occupancy");                                   \
16        static_assert((nbatch_fa)         <= 256, "bad nbatch_fa");                                   \
17        static_assert((nbatch_K)          <= 256, "bad nbatch_K");                                    \
18        return ((nthreads) << 0) | ((occupancy) << 10) | ((nbatch_fa) << 14) | ((nbatch_K) << 23);    \
19    }                                                                                                 \
20
21static constexpr uint32_t ggml_sycl_fattn_tile_get_config_fp16(const int DKQ, const int DV, const int ncols) {
22    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40,  2,  64, 2,  64,  40)
23    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40,  4, 128, 2,  64,  40)
24    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40,  8, 256, 2,  64,  40)
25    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40, 16, 256, 2,  64,  40)
26    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40, 32, 256, 2,  64,  40)
27
28    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64,  2,  64, 2,  64,  64)
29    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64,  4, 128, 2,  64,  64)
30    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64,  8, 256, 2,  64,  64)
31    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64, 16, 256, 2,  64,  64)
32    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64, 32, 256, 2,  64,  64)
33
34    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72,  2,  64, 2,  64,  72)
35    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72,  4, 128, 2,  64,  72)
36    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72,  8, 256, 2,  64,  72)
37    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72, 16, 256, 2,  64,  72)
38    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72, 32, 256, 2,  64,  72)
39
40    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80,  2,  64, 2,  64,  40)
41    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80,  4, 128, 2,  64,  40)
42    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80,  8, 256, 2,  64,  40)
43    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80, 16, 256, 2,  64,  40)
44    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80, 32, 256, 2,  64,  40)
45
46    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96,  2,  64, 2,  64,  48)
47    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96,  4, 128, 2,  64,  48)
48    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96,  8, 256, 2,  64,  48)
49    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96, 16, 256, 2,  64,  48)
50    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96, 32, 256, 2,  64,  48)
51
52    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112,  2,  64, 2,  64,  56)
53    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112,  4, 128, 2,  64,  56)
54    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112,  8, 256, 2,  64,  56)
55    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112, 16, 256, 2,  64,  56)
56    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112, 32, 256, 2,  64,  56)
57
58    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128,  2,  64, 2,  64,  64)
59    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128,  4, 128, 2,  64,  64)
60    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128,  8, 256, 2,  64,  64)
61    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128, 16, 256, 2,  64,  64)
62    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128, 32, 256, 2,  64,  64)
63
64    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256,  2,  64, 2,  64,  64)
65    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256,  4, 128, 2,  64,  64)
66    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256,  8, 256, 2,  64,  64)
67    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256, 16, 256, 2,  64,  64)
68    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256, 32, 256, 2,  64,  64)
69
70    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512,  2,  64, 2,  64,  64)
71    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512,  4, 128, 2,  64,  64)
72    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512,  8, 256, 2,  64,  64)
73    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512, 16, 256, 2,  64,  64)
74    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512, 32, 256, 2,  64,  64)
75
76    GGML_SYCL_FATTN_TILE_CONFIG_CASE(576, 512,  4, 128, 2,  64,  64)
77    GGML_SYCL_FATTN_TILE_CONFIG_CASE(576, 512,  8, 256, 2,  64,  64)
78    GGML_SYCL_FATTN_TILE_CONFIG_CASE(576, 512, 16, 256, 2,  64,  64)
79    GGML_SYCL_FATTN_TILE_CONFIG_CASE(576, 512, 32, 256, 2,  64,  64)
80
81    return 0;
82}
83
84static constexpr uint32_t ggml_sycl_fattn_tile_get_config_fp32(const int DKQ, const int DV, const int ncols) {
85    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40,  2,  64, 2,  32,  40)
86    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40,  4, 128, 2,  32,  40)
87    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40,  8, 256, 2,  32,  40)
88    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40, 16, 256, 2,  32,  40)
89    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 40,  40, 32, 256, 2,  32,  40)
90
91    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64,  2, 128, 3,  64,  64)
92    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64,  4, 128, 3,  32,  64)
93    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64,  8, 128, 3,  32,  64)
94    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64, 16, 128, 3,  64,  64)
95    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 64,  64, 32, 256, 2,  64,  64)
96
97    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72,  2,  64, 2,  32,  72)
98    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72,  4, 128, 2,  32,  72)
99    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72,  8, 256, 2,  32,  72)
100    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72, 16, 256, 2,  32,  72)
101    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 72,  72, 32, 256, 2,  32,  72)
102
103    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80,  2,  64, 2,  32,  40)
104    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80,  4, 128, 2,  32,  40)
105    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80,  8, 256, 2,  32,  40)
106    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80, 16, 256, 2,  32,  40)
107    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 80,  80, 32, 256, 2,  32,  40)
108
109    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96,  2,  64, 2,  32,  48)
110    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96,  4, 128, 2,  32,  48)
111    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96,  8, 256, 2,  32,  48)
112    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96, 16, 256, 2,  32,  48)
113    GGML_SYCL_FATTN_TILE_CONFIG_CASE( 96,  96, 32, 256, 2,  32,  48)
114
115    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112,  2,  64, 2,  32,  56)
116    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112,  4, 128, 2,  32,  56)
117    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112,  8, 256, 2,  32,  56)
118    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112, 16, 256, 2,  32,  56)
119    GGML_SYCL_FATTN_TILE_CONFIG_CASE(112, 112, 32, 256, 2,  32,  56)
120
121    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128,  2, 128, 3,  64,  64)
122    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128,  4, 128, 3,  32, 128)
123    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128,  8, 128, 3,  64, 128)
124    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128, 16, 128, 3,  32, 128)
125    GGML_SYCL_FATTN_TILE_CONFIG_CASE(128, 128, 32, 256, 2,  64,  64)
126
127    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256,  2, 128, 3,  64,  64)
128    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256,  4, 128, 3,  32,  64)
129    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256,  8, 256, 2,  32, 256)
130    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256, 16, 256, 2,  32, 128)
131    GGML_SYCL_FATTN_TILE_CONFIG_CASE(256, 256, 32, 256, 2,  32,  64)
132
133    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512,  2, 128, 2,  64,  64)
134    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512,  4, 128, 2,  64,  64)
135    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512,  8, 256, 2,  64,  64)
136    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512, 16, 256, 2,  64,  64)
137    GGML_SYCL_FATTN_TILE_CONFIG_CASE(512, 512, 32, 256, 2,  64,  64)
138
139    GGML_SYCL_FATTN_TILE_CONFIG_CASE(576, 512,  4, 128, 2,  32,  64)
140    GGML_SYCL_FATTN_TILE_CONFIG_CASE(576, 512,  8, 256, 2,  32,  64)
141    GGML_SYCL_FATTN_TILE_CONFIG_CASE(576, 512, 16, 256, 2,  32,  64)
142
143    return 0;
144}
145
146static constexpr uint32_t ggml_sycl_fattn_tile_get_config(const int DKQ, const int DV, const int ncols, const int cc) {
147    if(fast_fp16_available(cc))
148        return ggml_sycl_fattn_tile_get_config_fp16(DKQ, DV, ncols);
149    else
150        return ggml_sycl_fattn_tile_get_config_fp32(DKQ, DV, ncols);
151}
152
153static constexpr uint32_t ggml_sycl_fattn_tile_get_config(const int DKQ, const int DV, const int ncols) {
154#ifdef SYCL_FAST_FP16
155    return ggml_sycl_fattn_tile_get_config_fp16(DKQ, DV, ncols);
156#else
157    return ggml_sycl_fattn_tile_get_config_fp32(DKQ, DV, ncols);
158#endif // SYCL_FAST_FP16
159}
160
161static int ggml_sycl_fattn_tile_get_nthreads(const int DKQ, const int DV, const int ncols, const int cc) {
162    return (ggml_sycl_fattn_tile_get_config(DKQ, DV, ncols, cc) >> 0) & ((1 << 10) - 1);
163}
164
165static constexpr int ggml_sycl_fattn_tile_get_nthreads(const int DKQ, const int DV, const int ncols) {
166    return (ggml_sycl_fattn_tile_get_config(DKQ, DV, ncols) >> 0) & ((1 << 10) - 1);
167}
168
169static int ggml_sycl_fattn_tile_get_occupancy(const int DKQ, const int DV, const int ncols, const int cc) {
170    return (ggml_sycl_fattn_tile_get_config(DKQ, DV, ncols, cc) >> 10) & ((1 << 4) - 1);
171}
172
173static constexpr int ggml_sycl_fattn_tile_get_occupancy(const int DKQ, const int DV, const int ncols) {
174    return (ggml_sycl_fattn_tile_get_config(DKQ, DV, ncols) >> 10) & ((1 << 4) - 1);
175}
176
177static int ggml_sycl_fattn_tile_get_nbatch_fa(const int DKQ, const int DV, const int ncols, const int cc) {
178    return (ggml_sycl_fattn_tile_get_config(DKQ, DV, ncols, cc) >> 14) & ((1 << 9) - 1);
179}
180
181static constexpr int ggml_sycl_fattn_tile_get_nbatch_fa(const int DKQ, const int DV, const int ncols) {
182    return (ggml_sycl_fattn_tile_get_config(DKQ, DV, ncols) >> 14) & ((1 << 9) - 1);
183}
184
185static int ggml_sycl_fattn_tile_get_nbatch_K(const int DKQ, const int DV, const int ncols, const int cc) {
186    return (ggml_sycl_fattn_tile_get_config(DKQ, DV, ncols, cc) >> 23) & ((1 << 9) - 1);
187}
188
189static constexpr int ggml_sycl_fattn_tile_get_nbatch_K(const int DKQ, const int DV, const int ncols) {
190    return (ggml_sycl_fattn_tile_get_config(DKQ, DV, ncols) >> 23) & ((1 << 9) - 1);
191}
192
193template <int warp_size, int nwarps, int I, int J, int J_padding, bool oob_check>
194static __dpct_inline__ void flash_attn_tile_load_tile(const sycl::half2 * const __restrict__ KV,
195                                                      sycl::half2 * const __restrict__ tile_KV,
196                                                      const int stride_KV,
197                                                      const int i_sup) {
198    auto      item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
199    constexpr int cpy_nb = ggml_sycl_get_max_cpy_bytes();
200    constexpr int cpy_ne = cpy_nb / 4;
201
202    auto load = [&] (const int n) {
203        const int stride_j = warp_size >> n;
204
205        if (stride_j == 0) {
206            return;
207        }
208
209        const int j0_start = stride_j == warp_size ? 0 : ((J/2)/cpy_ne) - ((J/2)/cpy_ne) % (2*stride_j);
210        const int j0_stop  =                             ((J/2)/cpy_ne) - ((J/2)/cpy_ne) % (1*stride_j);
211        const int stride_i = warp_size / stride_j;
212
213        if (j0_start == j0_stop) {
214            return;
215        }
216
217#pragma unroll
218        for (int i0 = 0; i0 < I; i0 += nwarps*stride_i) {
219            const int i = i0 + item_ct1.get_local_id(1) * stride_i +
220                          (stride_j == warp_size ? 0 : item_ct1.get_local_id(2) / stride_j);
221
222            if (i0 + nwarps*stride_i <= I || i < I) {
223#pragma unroll
224                for (int j0 = j0_start; j0 < j0_stop; j0 += stride_j) {
225                    const int j = j0 * cpy_ne + (stride_j == warp_size ? item_ct1.get_local_id(2) :
226                                                                         item_ct1.get_local_id(2) % stride_j) *
227                                                    cpy_ne;
228
229                    const __dpct_align__(16) sycl::half2 zero[cpy_ne] = {
230                        { 0.0f, 0.0f }
231                    };
232                    ggml_sycl_memcpy_1<cpy_nb>(
233                        tile_KV + i*(J/2 + J_padding) + j,
234                        !oob_check || i < i_sup ? KV + i*stride_KV + j : zero);
235                }
236            }
237        }
238    };
239    // 1: max 64*16=512 bytes, 512 half
240    // 2: max 32*16=512 bytes, 256 half
241    // 3: max 16*16=256 bytes, 128 half
242    // 4: max  8*16=128 bytes,  64 half
243    // 5: max  4*16= 64 bytes,  32 half
244    // 6: max  2*16= 32 bytes,  16 half
245    // 7: max  1*16= 16 bytes,   8 half
246    static_assert(J % 8 == 0, "bad J");
247    static_assert((J/2) % cpy_ne == 0, "bad J");
248    ggml_sycl_unroll<7>{}(load);
249}
250
251template <int warp_size, int nwarps, int I, int J, int J_padding, bool oob_check>
252static __dpct_inline__ void flash_attn_tile_load_tile(const sycl::half2 * const __restrict__ KV,
253                                                      float * const __restrict__ tile_KV,
254                                                      const int stride_KV,
255                                                      const int i_sup) {
256    constexpr int cpy_nb = ggml_sycl_get_max_cpy_bytes();
257    constexpr int cpy_ne = cpy_nb / 4;
258
259    auto load = [&] (const int n) {
260        auto      item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
261        const int stride_j = warp_size >> n;
262
263        if (stride_j == 0) {
264            return;
265        }
266
267        const int j0_start = stride_j == warp_size ? 0 : (J/cpy_ne) - (J/cpy_ne) % (2*stride_j);
268        const int j0_stop  =                             (J/cpy_ne) - (J/cpy_ne) % (1*stride_j);
269        const int stride_i = warp_size / stride_j;
270
271        if (j0_start == j0_stop) {
272            return;
273        }
274
275#pragma unroll
276        for (int i0 = 0; i0 < I; i0 += nwarps*stride_i) {
277            const int i = i0 + item_ct1.get_local_id(1) * stride_i +
278                          (stride_j == warp_size ? 0 : item_ct1.get_local_id(2) / stride_j);
279
280            if (i0 + nwarps*stride_i <= I || i < I) {
281#pragma unroll
282                for (int j0 = j0_start; j0 < j0_stop; j0 += stride_j) {
283                    const int j = j0 * (cpy_ne / 2) + (stride_j == warp_size ? item_ct1.get_local_id(2) :
284                                                                               item_ct1.get_local_id(2) % stride_j) *
285                                                          (cpy_ne / 2);
286
287                    const sycl::half2 zero[cpy_ne / 2] = {
288                        { 0.0f, 0.0f }
289                    };
290                    __dpct_align__(16) sycl::half2 tmp_h2[cpy_ne / 2];
291                    ggml_sycl_memcpy_1<sizeof(tmp_h2)>(
292                        tmp_h2, !oob_check || i < i_sup ? KV + i*stride_KV + j : zero);
293
294                    __dpct_align__(16) sycl::float2 tmp_f2[cpy_ne / 2];
295#pragma unroll
296                    for (int l = 0; l < cpy_ne/2; ++l) {
297                        tmp_f2[l] = tmp_h2[l].template convert<float, sycl::rounding_mode::automatic>();
298                    }
299                    ggml_sycl_memcpy_1<sizeof(tmp_f2)>(tile_KV + i*(J + J_padding) + 2*j, tmp_f2);
300                }
301            }
302        }
303    };
304    // 1: max 32*16=512 bytes, 128 float
305    // 2: max 16*16=256 bytes,  64 float
306    // 3: max  8*16=128 bytes,  32 float
307    // 4: max  4*16= 64 bytes,  16 float
308    // 5: max  2*16= 32 bytes,   8 float
309    static_assert(J % 8 == 0, "bad J");
310    static_assert(J % cpy_ne == 0, "bad J");
311    ggml_sycl_unroll<5>{}(load);
312}
313
314// Function that performs a single iteration in for the KQ matrix multiplication:
315template <int  warp_size,
316          int  nwarps,
317          int  ncols1,
318          int  ncols2,
319          int  DKQ,
320          int  nbatch_fa,
321          int  nbatch_K,
322          bool use_logit_softcap,
323          bool oob_check,
324          typename T_vec_dot>
325static __dpct_inline__ void flash_attn_tile_iter_KQ(T_vec_dot * const Q_tmp,
326                                                    const sycl::half2 * const __restrict__ K_h2,
327                                                    T_vec_dot * const KV_tmp,
328                                                    const int         stride_K2,
329                                                    const int         k_VKQ_0,
330                                                    const int         k_VKQ_sup,
331                                                    const int         k_KQ_0,
332                                                    float *           KQ_acc) {
333    auto          item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
334    constexpr int cpy_nb   = ggml_sycl_get_max_cpy_bytes();
335    constexpr int cpy_ne = cpy_nb / 4;
336
337    constexpr int ncols = ncols1*ncols2;
338    constexpr int cpw   = ncols > nwarps ? ncols/nwarps : 1; // Q columns per warp
339    constexpr int np    = nwarps > ncols ? nwarps/ncols : 1; // number of parallel warps per Q column
340
341    flash_attn_tile_load_tile<warp_size, nwarps, nbatch_fa, nbatch_K, cpy_ne, oob_check>
342        (K_h2 + int64_t(k_VKQ_0)*stride_K2 + k_KQ_0/2, KV_tmp, stride_K2, k_VKQ_sup);
343    item_ct1.barrier(sycl::access::fence_space::local_space);
344
345#ifdef SYCL_FAST_FP16
346    static_assert((nbatch_K/2) % cpy_ne == 0, "bad nbatch_K");
347#pragma unroll
348    for (int k_KQ_1 = 0; k_KQ_1 < nbatch_K/2; k_KQ_1 += cpy_ne) {
349        __dpct_align__(16) sycl::half2 K_k[nbatch_fa / (np * warp_size)][cpy_ne];
350        __dpct_align__(16) sycl::half2 Q_k[cpw][cpy_ne];
351#else
352    static_assert(nbatch_K % cpy_ne == 0, "bad nbatch_K");
353#pragma unroll
354    for (int k_KQ_1 = 0; k_KQ_1 < nbatch_K; k_KQ_1 += cpy_ne) {
355        __dpct_align__(16) float K_k[nbatch_fa/(np*warp_size)][cpy_ne];
356        __dpct_align__(16) float Q_k[cpw][cpy_ne];
357#endif // SYCL_FAST_FP16
358
359#pragma unroll
360        for (int i_KQ_0 = 0; i_KQ_0 < nbatch_fa; i_KQ_0 += np*warp_size) {
361            const int i_KQ = i_KQ_0 + (item_ct1.get_local_id(1) % np) * warp_size + item_ct1.get_local_id(2);
362
363#ifdef SYCL_FAST_FP16
364            ggml_sycl_memcpy_1<cpy_nb>(&K_k[i_KQ_0/(np*warp_size)], &KV_tmp[i_KQ*(nbatch_K/2 + cpy_ne) + k_KQ_1]);
365#else
366            ggml_sycl_memcpy_1<cpy_nb>(&K_k[i_KQ_0/(np*warp_size)], &KV_tmp[i_KQ*(nbatch_K   + cpy_ne) + k_KQ_1]);
367#endif // SYCL_FAST_FP16
368        }
369#pragma unroll
370        for (int jc0 = 0; jc0 < cpw; ++jc0) {
371            const int jc = jc0 + (item_ct1.get_local_id(1) / np) * cpw;
372
373#ifdef SYCL_FAST_FP16
374            ggml_sycl_memcpy_1<cpy_nb>(&Q_k[jc0], &Q_tmp[jc*(DKQ/2) + k_KQ_0/2 + k_KQ_1]);
375#else
376            ggml_sycl_memcpy_1<cpy_nb>(&Q_k[jc0], &Q_tmp[jc* DKQ    + k_KQ_0   + k_KQ_1]);
377#endif // SYCL_FAST_FP16
378        }
379
380#pragma unroll
381        for (int i_KQ_0 = 0; i_KQ_0 < nbatch_fa; i_KQ_0 += np*warp_size) {
382#pragma unroll
383            for (int jc0 = 0; jc0 < cpw; ++jc0) {
384#pragma unroll
385                for (int k = 0; k < cpy_ne; ++k) {
386                    ggml_sycl_mad(KQ_acc[i_KQ_0/(np*warp_size)*cpw + jc0], K_k[i_KQ_0/(np*warp_size)][k], Q_k[jc0][k]);
387                }
388            }
389        }
390    }
391
392    if (k_KQ_0 + nbatch_K < DKQ) {
393        item_ct1.barrier(sycl::access::fence_space::local_space);  // Sync not needed on last iteration.
394    }
395}
396
397// Function that performs a single iteration of the main loop over up to nbatch_fa tokens.
398template <int  warp_size,
399          int  nwarps,
400          int  ncols1,
401          int  ncols2,
402          int  DKQ,
403          int  DV,
404          int  nbatch_fa,
405          int  nbatch_K,
406          bool use_logit_softcap,
407          bool oob_check,
408          typename T_vec_dot,
409          typename T_KQ,
410          typename T_acc>
411/*
412The total declared local variable size in device function flash_attn_tile_iter exceeds 128 bytes and may cause high register pressure. Consult with your hardware vendor to find the total register size available and adjust the code, or use smaller sub-group size to avoid high register pressure.
413*/
414static __dpct_inline__ void flash_attn_tile_iter(T_vec_dot * const Q_tmp,
415                                                 const sycl::half2 * const __restrict__ K_h2,
416                                                 const sycl::half2 * const __restrict__ V_h2,
417                                                 const sycl::half * const __restrict__ mask,
418                                                 const sycl::uint3 ne01,
419                                                 const float       logit_softcap,
420                                                 const float       slope,
421                                                 T_KQ * const      KQ,
422                                                 T_vec_dot * const KV_tmp,
423                                                 const int         stride_K2,
424                                                 const int         stride_V2,
425                                                 const int         stride_mask,
426                                                 float * const     KQ_max,
427                                                 float * const     KQ_sum,
428                                                 T_acc * const     VKQ,
429                                                 const int         k_VKQ_0,
430                                                 const int         k_VKQ_max,
431                                                 const int         col_Q_0,
432                                                 float *           KQ_max_new_shared) {
433    auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
434    constexpr int cpy_nb   = ggml_sycl_get_max_cpy_bytes();
435    constexpr int cpy_ne = cpy_nb / 4;
436
437    constexpr int ncols = ncols1*ncols2;
438    constexpr int cpw   = ncols > nwarps ? ncols/nwarps : 1; // Q columns per warp
439    constexpr int np    = nwarps > ncols ? nwarps/ncols : 1; // number of parallel warps per Q column
440
441    constexpr int DVp = (DV + 2*warp_size - 1) & ~(2*warp_size - 1); // DV padded to multiple of 2*warp_size.
442
443#ifdef SYCL_FAST_FP16
444    constexpr int KQ_cs = cpw < 2*cpy_ne ? cpw : 2*cpy_ne;
445#else
446    constexpr int KQ_cs = cpw < 1*cpy_ne ? cpw : 1*cpy_ne;
447#endif // SYCL_FAST_FP16
448    static_assert(cpw % KQ_cs == 0, "bad KQ_cs");
449    const int k_VKQ_sup = k_VKQ_max - k_VKQ_0; // k supremum, only smaller k values have valid KV data
450
451    float KQ_max_new[cpw];
452#pragma unroll
453    for (int jc0 = 0; jc0 < cpw; ++jc0) {
454        KQ_max_new[jc0] = KQ_max[jc0];
455    }
456
457    float KQ_acc[nbatch_fa/(np*warp_size) * cpw] = {0.0f}; // Accumulators for KQ matrix multiplication.
458
459    // KQ = K @ Q matrix multiplication:
460    constexpr int nbatch_K_last = DKQ % nbatch_K;
461#pragma unroll
462    for (int k_KQ_0 = 0; k_KQ_0 < DKQ - nbatch_K_last; k_KQ_0 += nbatch_K) {
463        flash_attn_tile_iter_KQ<warp_size, nwarps, ncols1, ncols2, DKQ, nbatch_fa, nbatch_K, use_logit_softcap, oob_check>(
464            Q_tmp, K_h2, KV_tmp, stride_K2, k_VKQ_0, k_VKQ_sup, k_KQ_0, KQ_acc);
465    }
466    if (nbatch_K_last > 0) {
467        constexpr int k_KQ_0 = DKQ - nbatch_K_last;
468        flash_attn_tile_iter_KQ<warp_size, nwarps, ncols1, ncols2, DKQ, nbatch_fa, nbatch_K_last, use_logit_softcap, oob_check>(
469            Q_tmp, K_h2, KV_tmp, stride_K2, k_VKQ_0, k_VKQ_sup, k_KQ_0, KQ_acc);
470    }
471
472    // Apply logit softcap + mask, update KQ_max:
473#pragma unroll
474    for (int jc0 = 0; jc0 < cpw; ++jc0) {
475        const int j = fastmodulo(col_Q_0 + (jc0 + (item_ct1.get_local_id(1) / np) * cpw) / ncols2, ne01);
476
477#pragma unroll
478        for (int i_KQ_0 = 0; i_KQ_0 < nbatch_fa; i_KQ_0 += np*warp_size) {
479            const int i_KQ = i_KQ_0 + (item_ct1.get_local_id(1) % np) * warp_size + item_ct1.get_local_id(2);
480
481#if defined(SYCL_FAST_FP16) && !defined(GGML_SYCL_F16)
482            // Without the v_dot2_f32_f16 instruction there is a higher risk of numerical overflow in the KQ calculation.
483            // Therefore, scale down Q values and apply the inverse scale the FP32 KQ values afterwards again.
484            KQ_acc[i_KQ_0/(np*warp_size)*cpw + jc0] *= 4.0f;
485#endif // defined(SYCL_FAST_FP16) && !defined(GGML_SYCL_F16)
486
487            if (use_logit_softcap) {
488                KQ_acc[(i_KQ_0 / (np * warp_size)) * cpw + jc0] =
489                    logit_softcap * sycl::tanh((float) KQ_acc[(i_KQ_0 / (np * warp_size)) * cpw + jc0]);
490            }
491
492            if (!oob_check || i_KQ < k_VKQ_sup) {
493                KQ_acc[(i_KQ_0 / (np * warp_size)) * cpw + jc0] +=
494                    (ncols2 > 1 || mask) ? slope * sycl::vec<sycl::half, 1>(mask[j * stride_mask + k_VKQ_0 + i_KQ])
495                                                       .convert<float, sycl::rounding_mode::automatic>()[0] :
496                                           0.0f;
497
498                KQ_max_new[jc0] =
499                    sycl::fmax((float) KQ_max_new[jc0],
500                               (float) (KQ_acc[(i_KQ_0 / (np * warp_size)) * cpw + jc0] + FATTN_KQ_MAX_OFFSET));
501            }
502        }
503
504        KQ_max_new[jc0] = warp_reduce_max<warp_size>(KQ_max_new[jc0]);
505    }
506
507    if constexpr (np == 1) {
508        item_ct1.barrier(sycl::access::fence_space::local_space);
509    } else {
510        static_assert(cpw == 1, "bad cpw");
511
512        if (item_ct1.get_local_id(2) == 0) {
513            KQ_max_new_shared[item_ct1.get_local_id(1)] = KQ_max_new[0];
514        }
515        item_ct1.barrier(sycl::access::fence_space::local_space);
516        KQ_max_new[0] = KQ_max_new_shared[(item_ct1.get_local_id(1) & ~(np - 1)) + item_ct1.get_local_id(2) % np];
517        KQ_max_new[0] = warp_reduce_max<np>(KQ_max_new[0]);
518    }
519
520    // Calculate KQ softmax, write to shared KQ buffer, re-scale VKQ accumulators:
521#pragma unroll
522    for (int jc0 = 0; jc0 < cpw; jc0 += KQ_cs) {
523#ifdef SYCL_FAST_FP16
524        __dpct_align__(16) sycl::half tmp[nbatch_fa / (np * warp_size)][KQ_cs];
525#else
526        __dpct_align__(16) float tmp[nbatch_fa/(np*warp_size)][KQ_cs];
527#endif // SYCL_FAST_FP16
528
529#pragma unroll
530        for (int jc1 = 0; jc1 < KQ_cs; ++jc1) {
531            const int jc = jc0 + jc1;
532
533            const float KQ_max_scale = sycl::native::exp((float) (KQ_max[jc] - KQ_max_new[jc]));
534            KQ_max[jc] = KQ_max_new[jc];
535
536            float KQ_sum_add = 0.0f;
537#pragma unroll
538            for (int i0 = 0; i0 < nbatch_fa; i0 += np*warp_size) {
539                const float val =
540                    !oob_check || i0 + (item_ct1.get_local_id(1) % np) * warp_size + item_ct1.get_local_id(2) <
541                                      static_cast<uint32_t>(k_VKQ_sup) ?
542                        sycl::native::exp((float) (KQ_acc[(i0 / (np * warp_size)) * cpw + jc] - KQ_max[jc])) :
543                        0.0f;
544                KQ_sum_add += val;
545                tmp[i0/(np*warp_size)][jc1] = val;
546            }
547            KQ_sum[jc] = KQ_sum[jc]*KQ_max_scale + KQ_sum_add;
548
549#ifdef SYCL_FAST_FP16
550            const sycl::half2 KQ_max_scale_h2 = sycl::half2(KQ_max_scale, KQ_max_scale);
551#pragma unroll
552            for (int i0 = 0; i0 < DVp/2; i0 += warp_size) {
553                VKQ[jc*((DVp/2)/warp_size) + i0/warp_size].x() *= KQ_max_scale_h2.x();
554                VKQ[jc*((DVp/2)/warp_size) + i0/warp_size].y() *= KQ_max_scale_h2.y();
555            }
556#else
557#pragma unroll
558            for (int i0 = 0; i0 < DVp/2; i0 += warp_size) {
559                VKQ[jc*((DVp/2)/warp_size) + i0/warp_size].x() *= KQ_max_scale;
560                VKQ[jc*((DVp/2)/warp_size) + i0/warp_size].y() *= KQ_max_scale;
561            }
562#endif // SYCL_FAST_FP16
563        }
564
565#pragma unroll
566        for (int i0 = 0; i0 < nbatch_fa; i0 += np*warp_size) {
567            const int i = i0 + (item_ct1.get_local_id(1) % np) * warp_size + item_ct1.get_local_id(2);
568
569            ggml_sycl_memcpy_1<sizeof(tmp[0])>(
570                KQ + (jc0 / KQ_cs + (item_ct1.get_local_id(1) / np) * (cpw / KQ_cs)) * (nbatch_fa * KQ_cs) + i * KQ_cs,
571                tmp[i0 / (np * warp_size)]);
572        }
573    }
574
575    // VKQ = V @ KQ matrix multiplication:
576    static_assert(DV <= DKQ, "bad DV");
577    static_assert(DV % nbatch_K == 0 || (nbatch_K % 3 == 0 && DV % (nbatch_K*2/3) == 0), "bad nbatch_K");
578    constexpr int nbatch_V = (DV % nbatch_K == 0 ? nbatch_K : nbatch_K*2/3) * nbatch_fa / DV; // Number of V columns that fit in SRAM for K.
579    static_assert(nbatch_fa % nbatch_V == 0, "bad nbatch_V");
580    static_assert(nbatch_V % np == 0, "bad nbatch_V");
581#pragma unroll
582    for (int k0 = 0; k0 < nbatch_fa; k0 += nbatch_V) {
583        flash_attn_tile_load_tile<warp_size, nwarps, nbatch_V, DV, 0, oob_check>
584            (V_h2 + int64_t(k_VKQ_0 + k0)*stride_V2, KV_tmp, stride_V2, k_VKQ_sup - k0);
585        item_ct1.barrier(sycl::access::fence_space::local_space);
586
587#ifdef SYCL_FAST_FP16
588#pragma unroll
589        for (int k1 = 0; k1 < nbatch_V; k1 += np) {
590            __dpct_align__(16) sycl::half2 V_k[(DVp / 2) / warp_size];
591            __dpct_align__(16) sycl::half2 KQ_k[cpw];
592
593            constexpr int cpy_ne_D = cpy_ne/2 < (DVp/2)/warp_size ? cpy_ne/2 : (DVp/2)/warp_size;
594#pragma unroll
595            for (int i0 = 0; i0 < DVp/2; i0 += warp_size*cpy_ne_D) {
596                ggml_sycl_memcpy_1<cpy_ne_D * 4>(&V_k[i0 / warp_size],
597                                                 &KV_tmp[(k1 + item_ct1.get_local_id(1) % np) * (DV / 2) + i0 +
598                                                         item_ct1.get_local_id(2) * cpy_ne_D]);
599            }
600#pragma unroll
601            for (int jc_VKQ_0 = 0; jc_VKQ_0 < cpw; jc_VKQ_0 += KQ_cs) {
602                const int jc_KQ = jc_VKQ_0 / KQ_cs + (item_ct1.get_local_id(1) / np) * (cpw / KQ_cs);
603
604                __dpct_align__(16) sycl::half tmp[KQ_cs];
605                ggml_sycl_memcpy_1<KQ_cs * sizeof(sycl::half)>(
606                    &tmp, KQ + jc_KQ * (nbatch_fa * KQ_cs) + (k0 + k1 + item_ct1.get_local_id(1) % np) * KQ_cs);
607#pragma unroll
608                for (int jc_VKQ_1 = 0; jc_VKQ_1 < KQ_cs; ++jc_VKQ_1) {
609                    KQ_k[jc_VKQ_0 + jc_VKQ_1] = sycl::half2(tmp[jc_VKQ_1]);
610                }
611            }
612
613#pragma unroll
614            for (int i0 = 0; i0 < DVp/2; i0 += warp_size) {
615#pragma unroll
616                for (int jc_VKQ_0 = 0; jc_VKQ_0 < cpw; ++jc_VKQ_0) {
617                    VKQ[jc_VKQ_0*((DVp/2)/warp_size) + i0/warp_size].x() +=
618                        V_k[i0/warp_size].x()*KQ_k[jc_VKQ_0].x();
619                    VKQ[jc_VKQ_0*((DVp/2)/warp_size) + i0/warp_size].y() +=
620                        V_k[i0/warp_size].y()*KQ_k[jc_VKQ_0].y();
621                }
622            }
623        }
624#else
625#pragma unroll
626        for (int k1 = 0; k1 < nbatch_V; k1 += np) {
627            __dpct_align__(16) sycl::float2 V_k[(DVp/2)/warp_size];
628            __dpct_align__(16) float  KQ_k[cpw];
629
630            constexpr int cpy_ne_D = cpy_ne < DVp/warp_size ? cpy_ne : DVp/warp_size;
631#pragma unroll
632            for (int i0 = 0; i0 < DVp; i0 += warp_size*cpy_ne_D) {
633                ggml_sycl_memcpy_1<cpy_ne_D*4>(&V_k[i0/(2*warp_size)], &KV_tmp[(k1 + item_ct1.get_local_id(1) % np)*DV + i0 + item_ct1.get_local_id(2)*cpy_ne_D]);
634            }
635#pragma unroll
636            for (int jc_VKQ_0 = 0; jc_VKQ_0 < cpw; jc_VKQ_0 += KQ_cs) {
637                const int jc_KQ = jc_VKQ_0/KQ_cs + (item_ct1.get_local_id(1) / np)*(cpw/KQ_cs);
638
639                ggml_sycl_memcpy_1<KQ_cs*sizeof(float)>(
640                    &KQ_k[jc_VKQ_0], KQ + jc_KQ*(nbatch_fa*KQ_cs) + (k0 + k1 + item_ct1.get_local_id(1) % np)*KQ_cs);
641            }
642
643#pragma unroll
644            for (int i0 = 0; i0 < DVp/2; i0 += warp_size) {
645#pragma unroll
646                for (int jc_VKQ_0 = 0; jc_VKQ_0 < cpw; ++jc_VKQ_0) {
647                    VKQ[jc_VKQ_0*((DVp/2)/warp_size) + i0/warp_size].x() += V_k[i0/warp_size].x()*KQ_k[jc_VKQ_0];
648                    VKQ[jc_VKQ_0*((DVp/2)/warp_size) + i0/warp_size].y() += V_k[i0/warp_size].y()*KQ_k[jc_VKQ_0];
649                }
650            }
651        }
652#endif // SYCL_FAST_FP16
653        item_ct1.barrier(sycl::access::fence_space::local_space);
654    }
655}
656
657template <int DKQ, int DV, int ncols1, int ncols2, bool use_logit_softcap, int warp_size>  // D == head size
658/*
659The total declared local variable size in device function flash_attn_tile exceeds 128 bytes and may cause high register pressure. Consult with your hardware vendor to find the total register size available and adjust the code, or use smaller sub-group size to avoid high register pressure.
660*/
661static void flash_attn_tile(const char *  Q,
662                            const char *  K,
663                            const char *  V,
664                            const char *  mask,
665                            const char *  sinks,
666                            const int *  KV_max,
667                            float *  dst,
668                            sycl::float2 *  dst_meta,
669                            const float          scale,
670                            const float          max_bias,
671                            const float          m0,
672                            const float          m1,
673                            const uint32_t       n_head_log2,
674                            const float          logit_softcap,
675                            const int32_t        ne00,
676                            const sycl::uint3    ne01,
677                            const int32_t        ne02,
678                            const int32_t        ne03,
679                            const int32_t        nb01,
680                            const int32_t        nb02,
681                            const int32_t        nb03,
682                            const int32_t        ne10,
683                            const int32_t        ne11,
684                            const int32_t        ne12,
685                            const int32_t        ne13,
686                            const int32_t        nb11,
687                            const int32_t        nb12,
688                            const int64_t        nb13,
689                            const int32_t        nb21,
690                            const int32_t        nb22,
691                            const int64_t        nb23,
692                            const int32_t        ne31,
693                            const int32_t        ne32,
694                            const int32_t        ne33,
695                            const int32_t        nb31,
696                            const int32_t        nb32,
697                            const int64_t        nb33) {
698#ifdef SYCL_FLASH_ATTN
699    // Skip unused kernel variants for faster compilation:
700    auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
701    if ((use_logit_softcap && !(DV == 128 || DV == 256))) {
702        GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale,
703            max_bias, m0, m1, n_head_log2, logit_softcap,
704            ne00, ne01, ne02, ne03,
705                  nb01, nb02, nb03,
706            ne10, ne11, ne12, ne13,
707                  nb11, nb12, nb13,
708                  nb21, nb22, nb23,
709                  ne31, ne32, ne33,
710                  nb31, nb32, nb33);
711        return;
712    }
713
714    static_assert(ggml_sycl_fattn_tile_get_config(DKQ, DV, ncols1*ncols2) != 0, "kernel config not defined");
715
716    constexpr int ncols     = ncols1*ncols2;
717
718    constexpr int nwarps    = ggml_sycl_fattn_tile_get_nthreads (DKQ, DV, ncols1*ncols2) / warp_size;
719    constexpr int nbatch_fa = ggml_sycl_fattn_tile_get_nbatch_fa(DKQ, DV, ncols1*ncols2);
720    constexpr int nbatch_K  = ggml_sycl_fattn_tile_get_nbatch_K (DKQ, DV, ncols1*ncols2);
721
722    // In this kernel Q, K, V are matrices while i, j, k are matrix indices.
723
724    const int col_Q_0 = item_ct1.get_group(2) * ncols1;  // Index of the first Q column for this SYCL block to work on.
725
726    const int           sequence  = item_ct1.get_group(0) / (ne02 / ncols2);
727    const int           head0     = item_ct1.get_group(0) * ncols2 - sequence * ne02;  // == item_ct1.get_group(0) % (ne02/ncols2)
728    const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix.
729    const float * Q_f  = (const float *) (Q + nb03*sequence + nb02* head0);
730    const sycl::half2 * K_h2      = (const sycl::half2 *) (K + nb13 * sequence + nb12 * (head0 / gqa_ratio));
731    const sycl::half2 * V_h2 =
732        (const sycl::half2 *) (V + nb23 * sequence + nb22 * (head0 / gqa_ratio));  // K and V have same shape
733
734    const sycl::half * maskh = mask ? (const sycl::half *) (mask + nb33 * (sequence % ne33)) : nullptr;
735
736    const int stride_K2   = nb11 / sizeof(sycl::half2);
737    const int stride_V2   = nb21 / sizeof(sycl::half2);
738    const int stride_mask = nb31 / sizeof(sycl::half);
739
740    const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, head0, n_head_log2, m0, m1) : 1.0f;
741
742    constexpr int cpy_nb = ggml_sycl_get_max_cpy_bytes();
743    constexpr int cpy_ne = cpy_nb / 4;
744
745    constexpr int cpw = ncols > nwarps ? ncols/nwarps : 1; // Q columns per warp.
746    constexpr int np  = nwarps > ncols ? nwarps/ncols : 1; // Number of parallel warps per Q column.
747
748    static_assert(cpw == 1 || np == 1, "bad cpw / np");
749    static_assert(nbatch_fa % (np*warp_size) == 0, "nbatch_fa % (np*warp_size) != 0");
750
751    constexpr int DKQp = (DKQ + 2*warp_size - 1) & ~(2*warp_size - 1); // DKQ padded to multiple of 2*warp_size.
752    constexpr int DVp  = (DV  + 2*warp_size - 1) & ~(2*warp_size - 1); // DV  padded to multiple of 2*warp_size.
753
754    // Q_tmp == SRAM buffer to hold Q data for the entire lifetime of the kernel.
755    // KV_tmp == SRAM buffer to hold fragments of K/V data while iterating over ne11.
756    //     KV_tmp is padded to avoid memory conflicts for K (cpy_ne) and OOB accesses for V (DVp-DV).
757    // KQ == SRAM buffer to hold KQ fragments between KQ and VKQ matrix multiplications.
758    // VKQ == Accumulators in registers for the final VKQ result.
759
760
761#ifdef SYCL_FAST_FP16
762    constexpr size_t lsm_size1 = ncols * DKQ/2 ;
763    constexpr size_t lsm_size2 = nbatch_fa * (nbatch_K/2 + cpy_ne) + DVp-DV ;
764    constexpr size_t lsm_size3 = ncols * nbatch_fa;
765    constexpr size_t lsm_size4 = nwarps;
766
767    constexpr size_t local_share_mem_size = lsm_size1 * sizeof(sycl::half2) +
768                                            lsm_size2 * sizeof(sycl::half2) +
769                                            lsm_size3 * sizeof(sycl::half) +
770                                            lsm_size4 * sizeof(float);
771
772    syclex::work_group_static<char[local_share_mem_size]> lsm;
773
774    sycl::half2 *Q_tmp = (sycl::half2 *)&lsm;
775    sycl::half2 *KV_tmp = (sycl::half2*)(Q_tmp +lsm_size1);
776    sycl::half *KQ = (sycl::half *)(KV_tmp+lsm_size2);
777    float *KQ_max_new_shared = (float *)(KQ+lsm_size3);
778
779    __dpct_align__(16) sycl::half2 VKQ[cpw * ((DVp / 2) / warp_size)] = {
780        { 0.0f, 0.0f }
781    };
782#else
783    constexpr size_t lsm_size1 = ncols * DKQ ;
784    constexpr size_t lsm_size2 = nbatch_fa * (nbatch_K + cpy_ne) + DVp-DV;
785    constexpr size_t lsm_size3 = ncols * nbatch_fa;
786    constexpr size_t lsm_size4 = nwarps;
787
788    constexpr size_t local_share_mem_size = (lsm_size1 + lsm_size2 +lsm_size3 + lsm_size4) * sizeof(float);
789
790    syclex::work_group_static<char[local_share_mem_size]> lsm;
791
792    float *Q_tmp = (float *)&lsm;
793    float *KV_tmp = Q_tmp +lsm_size1;
794    float *KQ = KV_tmp+lsm_size2;
795    float *KQ_max_new_shared = KQ+lsm_size3;
796
797    __dpct_align__(16) sycl::float2 VKQ[cpw * ((DVp/2)/warp_size)] = {{0.0f, 0.0f}};
798
799
800#endif // SYCL_FAST_FP16
801
802    float KQ_max[cpw] = {};
803
804#pragma unroll
805    for (int j0 = 0; j0 < ncols; j0 += nwarps) {
806        KQ_max[j0/nwarps] = -FLT_MAX/2.0f;
807    }
808    float KQ_sum[cpw] = {0.0f};
809
810    // Load Q data, convert to FP16 if fast:
811#pragma unroll
812    for (int jc0 = 0; jc0 < cpw; ++jc0) {
813        const int jc = jc0 + (item_ct1.get_local_id(1) / np) * cpw;
814
815        const int j = jc / ncols2;
816        const int c = jc % ncols2;
817
818        constexpr int cpy_ne_D = cpy_ne < DKQp/warp_size ? cpy_ne : DKQp/warp_size;
819
820#pragma unroll
821        for (int i0 = 0; i0 < DKQp; i0 += np*warp_size*cpy_ne_D) {
822            if (i0 + np * warp_size * cpy_ne_D <= DKQ ||
823                i0 + (item_ct1.get_local_id(1) % np) * (warp_size * cpy_ne_D) + item_ct1.get_local_id(2) * cpy_ne_D <
824                    DKQ) {
825                __dpct_align__(16) float tmp_f[cpy_ne_D] = { 0.0f };
826                ggml_sycl_memcpy_1<sizeof(tmp_f)>(
827                    tmp_f, &Q_f[c * (nb02 / sizeof(float)) + fastmodulo(col_Q_0 + j, ne01) * (nb01 / sizeof(float)) +
828                                i0 + (item_ct1.get_local_id(1) % np) * (warp_size * cpy_ne_D) +
829                                item_ct1.get_local_id(2) * cpy_ne_D]);
830
831#pragma unroll
832                for (int i1 = 0; i1 < cpy_ne_D; ++i1) {
833                    tmp_f[i1] *= scale;
834                }
835
836#ifdef SYCL_FAST_FP16
837                __dpct_align__(16) sycl::half2 tmp_h2[cpy_ne_D / 2];
838#pragma unroll
839                for (int i1 = 0; i1 < cpy_ne_D; i1 += 2) {
840                    tmp_h2[i1/2] = make_half2(tmp_f[i1 + 0], tmp_f[i1 + 1]);
841#if defined(SYCL_FAST_FP16) && !defined(GGML_SYCL_F16)
842                    // Without the v_dot2_f32_f16 instruction there is a higher risk of numerical overflow in the KQ calculation.
843                    // Therefore, scale down Q values and apply the inverse scale the FP32 KQ values afterwards again.
844                    tmp_h2[i1 / 2] *= sycl::half2(0.25f, 0.25f);
845#endif // defined(SYCL_FAST_FP16) && !defined(GGML_SYCL_F16)
846                }
847                ggml_sycl_memcpy_1<sizeof(tmp_h2)>(
848                    &Q_tmp[jc * (DKQ / 2) + i0 / 2 + (item_ct1.get_local_id(1) % np) * (warp_size * cpy_ne_D / 2) +
849                           item_ct1.get_local_id(2) * (cpy_ne_D / 2)],
850                    tmp_h2);
851#else
852                ggml_sycl_memcpy_1<sizeof(tmp_f)>(
853                    &Q_tmp[jc* DKQ    + i0   + (item_ct1.get_local_id(1) % np)*(warp_size*cpy_ne_D)   + item_ct1.get_local_id(2)* cpy_ne_D],
854                    tmp_f);
855#endif // SYCL_FAST_FP16
856            }
857        }
858    }
859
860    item_ct1.barrier(sycl::access::fence_space::local_space);
861
862    // Main loop over KV cache:
863    const int k_VKQ_max = KV_max ? KV_max[sequence * item_ct1.get_group_range(2) + item_ct1.get_group(2)] : ne11;
864    if (ncols2 == 1) {
865        // Branch with out-of-bounds checks.
866        int k_VKQ_0 = item_ct1.get_group(1) * nbatch_fa;
867        while (k_VKQ_0 < k_VKQ_max - nbatch_fa) {
868            constexpr bool oob_check = false;
869            flash_attn_tile_iter<warp_size, nwarps, ncols1, ncols2, DKQ, DV, nbatch_fa, nbatch_K, use_logit_softcap,
870                                 oob_check>(Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, stride_K2,
871                                            stride_V2, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0,
872                                            KQ_max_new_shared);
873            k_VKQ_0 += item_ct1.get_group_range(1) * nbatch_fa;
874        }
875        if (k_VKQ_0 < k_VKQ_max) {
876            constexpr bool oob_check = true;
877            flash_attn_tile_iter<warp_size, nwarps, ncols1, ncols2, DKQ, DV, nbatch_fa, nbatch_K, use_logit_softcap,
878                                 oob_check>(Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, stride_K2,
879                                            stride_V2, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0,
880                                            KQ_max_new_shared);
881        }
882    } else {
883        // Branch without out-of-bounds checks.
884        for (int k_VKQ_0 = item_ct1.get_group(1) * nbatch_fa; k_VKQ_0 < k_VKQ_max;
885             k_VKQ_0 += item_ct1.get_group_range(1) * nbatch_fa) {
886
887            constexpr bool oob_check = false;
888            flash_attn_tile_iter<warp_size, nwarps, ncols1, ncols2, DKQ, DV, nbatch_fa, nbatch_K, use_logit_softcap,
889                                 oob_check>(Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, stride_K2,
890                                            stride_V2, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0,
891                                            KQ_max_new_shared);
892        }
893    }
894
895#pragma unroll
896    for (int jc0 = 0; jc0 < cpw; ++jc0) {
897        KQ_sum[jc0] = warp_reduce_sum<warp_size>(KQ_sum[jc0]);
898    }
899
900    if constexpr (np > 1) {
901        static_assert(cpw == 1, "bad cpw");
902        static_assert(nbatch_fa*nbatch_K >= nwarps*DVp, "KV_tmp too small");
903
904#ifdef SYCL_FAST_FP16
905        sycl::half2 * VKQ_combine = (sycl::half2 *) KV_tmp;
906#else
907        float * VKQ_combine    = (float *) KV_tmp;
908#endif // SYCL_FAST_FP16
909
910        float * KQ_sum_combine = (float *) Q_tmp;
911
912        if (item_ct1.get_local_id(1) % np != 0) {
913
914#ifdef SYCL_FAST_FP16
915            constexpr int cpy_ne_D = cpy_ne < (DVp/2)/warp_size ? cpy_ne : (DVp/2)/warp_size;
916#pragma unroll
917            for (int i0 = 0; i0 < DVp/2; i0 += warp_size*cpy_ne_D) {
918                ggml_sycl_memcpy_1<cpy_ne_D * 4>(
919                    &VKQ_combine[item_ct1.get_local_id(1) * (DVp / 2) + i0 + item_ct1.get_local_id(2) * cpy_ne_D],
920                    &VKQ[i0 / warp_size]);
921            }
922#else
923
924            constexpr int cpy_ne_D = cpy_ne < DVp/warp_size ? cpy_ne : DVp/warp_size;
925
926#pragma unroll
927            for (int i0 = 0; i0 < DVp; i0 += warp_size*cpy_ne_D) {
928                ggml_sycl_memcpy_1<cpy_ne_D*4>(
929                    &VKQ_combine[item_ct1.get_local_id(1)*DVp + i0 + item_ct1.get_local_id(2)*cpy_ne_D], ((const float *) VKQ) + i0/warp_size);
930            }
931#endif // SYCL_FAST_FP16
932
933            if (item_ct1.get_local_id(2) == 0) {
934                KQ_sum_combine[item_ct1.get_local_id(1)] = KQ_sum[0];
935            }
936            return;
937        }
938
939        item_ct1.barrier(sycl::access::fence_space::local_space);
940
941#pragma unroll
942        for (int ip = 1; ip < np; ++ip) {
943#ifdef SYCL_FAST_FP16
944            constexpr int cpy_ne_D = cpy_ne < (DVp/2)/warp_size ? cpy_ne : (DVp/2)/warp_size;
945#pragma unroll
946            for (int i0 = 0; i0 < DVp/2; i0 += warp_size*cpy_ne_D) {
947                __dpct_align__(16) sycl::half2 tmp[cpy_ne_D];
948                ggml_sycl_memcpy_1<cpy_ne_D * 4>(tmp, &VKQ_combine[(item_ct1.get_local_id(1) + ip) * (DVp / 2) + i0 +
949                                                                   item_ct1.get_local_id(2) * cpy_ne_D]);
950#pragma unroll
951                for (int i1 = 0; i1 < cpy_ne_D; ++i1) {
952                    VKQ[i0/warp_size + i1] += tmp[i1];
953                }
954            }
955#else
956            constexpr int cpy_ne_D = cpy_ne < DVp/warp_size ? cpy_ne : DVp/warp_size;
957#pragma unroll
958            for (int i0 = 0; i0 < DVp; i0 += warp_size*cpy_ne_D) {
959                __dpct_align__(16) float tmp[cpy_ne_D];
960                ggml_sycl_memcpy_1<cpy_ne_D*4>(tmp, &VKQ_combine[(item_ct1.get_local_id(1) + ip)*DVp + i0 + item_ct1.get_local_id(2)*cpy_ne_D]);
961#pragma unroll
962                for (int i1 = 0; i1 < cpy_ne_D; ++i1) {
963                    ((float *)VKQ)[i0/warp_size + i1] += tmp[i1];
964                }
965            }
966#endif // SYCL_FAST_FP16
967
968            KQ_sum[0] += KQ_sum_combine[item_ct1.get_local_id(1) + ip];
969        }
970    }
971
972    // Attention sink: adjust KQ max and sum only for the first of all parallel blocks:
973    if (sinks && item_ct1.get_group(1) == 0) {
974#pragma unroll
975        for (int jc0 = 0; jc0 < cpw; ++jc0) {
976            const int   jc   = jc0 + (item_ct1.get_local_id(1) / np) * cpw;
977            const float sink = ((const float *) sinks)[head0 + jc % ncols2];
978
979            float       KQ_max_new_j = sycl::fmax((float) KQ_max[jc0], sink);
980            const float KQ_max_scale = sycl::native::exp((float) (KQ_max[jc0] - KQ_max_new_j));
981            KQ_max[jc0] = KQ_max_new_j;
982
983            const float val = sycl::native::exp((float) (sink - KQ_max[jc0]));
984            KQ_sum[jc0] = KQ_sum[jc0]*KQ_max_scale + val;
985
986#ifdef SYCL_FAST_FP16
987            const sycl::half2 KQ_max_scale_h2 = sycl::half2(KQ_max_scale, KQ_max_scale);
988#pragma unroll
989            for (int i0 = 0; i0 < DVp/2; i0 += warp_size) {
990                VKQ[jc0*((DVp/2)/warp_size) + i0/warp_size] *= KQ_max_scale_h2;
991            }
992#else
993#pragma unroll
994            for (int i0 = 0; i0 < DVp/2; i0 += warp_size) {
995                VKQ[jc0*((DVp/2)/warp_size) + i0/warp_size].x() *= KQ_max_scale;
996                VKQ[jc0*((DVp/2)/warp_size) + i0/warp_size].y() *= KQ_max_scale;
997            }
998#endif // SYCL_FAST_FP16
999        }
1000    }
1001
1002    // Write back results:
1003#pragma unroll
1004    for (int jc0 = 0; jc0 < cpw; ++jc0) {
1005        const int jc = jc0 + (item_ct1.get_local_id(1) / np) * cpw;
1006
1007        const int j = jc / ncols2;
1008        const int c = jc % ncols2;
1009
1010        if (ncols1 > 1 && col_Q_0 + j >= int(ne01.z())) {
1011            return;
1012        }
1013
1014        const float scale = item_ct1.get_group_range(1) == 1 ? 1.0f / KQ_sum[jc0] : 1.0f;
1015
1016        const int j_dst_unrolled =
1017            ((sequence * int(ne01.z()) + col_Q_0 + j) * ne02 + head0 + c) * item_ct1.get_group_range(1) +
1018            item_ct1.get_group(1);
1019
1020#ifdef SYCL_FAST_FP16
1021        constexpr int cpy_ne_D = cpy_ne/2 < (DVp/2)/warp_size ? cpy_ne/2 : (DVp/2)/warp_size;
1022#pragma unroll
1023        for (int i0 = 0; i0 < DVp/2; i0 += warp_size*cpy_ne_D) {
1024            __dpct_align__(16) sycl::float2 tmp[cpy_ne_D];
1025#pragma unroll
1026            for (int i1 = 0; i1 < cpy_ne_D; ++i1) {
1027                tmp[i1] = VKQ[jc0 * ((DVp / 2) / warp_size) + i0 / warp_size + i1]
1028                              .template convert<float, sycl::rounding_mode::automatic>();
1029                tmp[i1].x() *= scale;
1030                tmp[i1].y() *= scale;
1031            }
1032            if (i0 + warp_size * cpy_ne_D <= DV / 2 || i0 + item_ct1.get_local_id(2) * cpy_ne_D < DV / 2) {
1033                ggml_sycl_memcpy_1<sizeof(tmp)>(
1034                    &dst[j_dst_unrolled * DV + 2 * i0 + item_ct1.get_local_id(2) * (2 * cpy_ne_D)], tmp);
1035            }
1036        }
1037#else
1038        constexpr int cpy_ne_D = cpy_ne < DVp/warp_size ? cpy_ne : DVp/warp_size;
1039#pragma unroll
1040        for (int i0 = 0; i0 < DVp; i0 += warp_size*cpy_ne_D) {
1041            if (i0 + warp_size*cpy_ne_D <= DV || i0 + item_ct1.get_local_id(2)*cpy_ne_D < DV) {
1042#pragma unroll
1043                for (int i1 = 0; i1 < cpy_ne_D/2; ++i1) {
1044                    VKQ[jc0*((DVp/2)/warp_size) + i0/(2*warp_size) + i1].x() *= scale;
1045                    VKQ[jc0*((DVp/2)/warp_size) + i0/(2*warp_size) + i1].y() *= scale;
1046                }
1047                ggml_sycl_memcpy_1<cpy_ne_D*4>(
1048                    &dst[j_dst_unrolled*DV + i0 + item_ct1.get_local_id(2)*cpy_ne_D],
1049                    &VKQ[jc0*((DVp/2)/warp_size) + i0/(2*warp_size)]);
1050            }
1051        }
1052#endif // SYCL_FAST_FP16
1053
1054        if (item_ct1.get_group_range(1) != 1 && item_ct1.get_local_id(2) == 0) {
1055            dst_meta[j_dst_unrolled] = make_float2(KQ_max[jc0], KQ_sum[jc0]);
1056        }
1057    }
1058#else
1059    GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale,
1060        max_bias, m0, m1, n_head_log2, logit_softcap,
1061        ne00, ne01, ne02, ne03,
1062              nb01, nb02, nb03,
1063        ne10, ne11, ne12, ne13,
1064              nb11, nb12, nb13,
1065              nb21, nb22, nb23,
1066              ne31, ne32, ne33,
1067              nb31, nb32, nb33);
1068#endif // SYCL_FLASH_ATTN
1069}
1070
1071template <int DKQ, int DV, int ncols2, bool use_logit_softcap>
1072static void launch_fattn_tile_switch_ncols1(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
1073    const ggml_tensor * Q = dst->src[0];
1074
1075    const int id        = ggml_sycl_get_device();
1076    const int cc        = ggml_sycl_info().devices[id].cc;
1077    const int warp_size = WARP_32_SIZE; //can't support WARP_16_SIZE
1078
1079    constexpr size_t nbytes_shared = 0;
1080
1081    if (DV < 512 && Q->ne[1] < 32) {
1082        if constexpr (ncols2 <= 32) {
1083            if (Q->ne[1] > 16/ncols2) {
1084                constexpr int cols_per_block = 32;
1085                const int nwarps    = ggml_sycl_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size;
1086                const int nbatch_fa = ggml_sycl_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
1087                launch_fattn<DV, cols_per_block/ncols2, ncols2,
1088                    flash_attn_tile<DKQ, DV, cols_per_block / ncols2, ncols2, use_logit_softcap, warp_size>, warp_size>
1089                    (ctx, dst, nwarps, nbytes_shared, nbatch_fa, true, true, false);
1090                return;
1091            }
1092        }
1093        if constexpr (ncols2 <= 16) {
1094            if (Q->ne[1] > 8/ncols2) {
1095                constexpr int cols_per_block = 16;
1096                const int nwarps    = ggml_sycl_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size;
1097                const int nbatch_fa = ggml_sycl_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
1098                launch_fattn<DV, cols_per_block/ncols2, ncols2,
1099                    flash_attn_tile<DKQ, DV, cols_per_block / ncols2, ncols2, use_logit_softcap, warp_size>, warp_size>
1100                    (ctx, dst, nwarps, nbytes_shared, nbatch_fa, true, true, false);
1101                return;
1102            }
1103        }
1104        if constexpr (ncols2 <= 8) {
1105            if (Q->ne[1] > 4/ncols2) {
1106                constexpr int cols_per_block = 8;
1107                const int nwarps    = ggml_sycl_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size;
1108                const int nbatch_fa = ggml_sycl_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
1109                launch_fattn<DV, cols_per_block/ncols2, ncols2,
1110                    flash_attn_tile<DKQ, DV, cols_per_block / ncols2, ncols2, use_logit_softcap, warp_size>, warp_size>
1111                    (ctx, dst, nwarps, nbytes_shared, nbatch_fa, true, true, false);
1112                return;
1113            }
1114        }
1115    }
1116
1117    if constexpr (ncols2 <= 4) {
1118        if (Q->ne[1] > 2/ncols2) {
1119            constexpr int cols_per_block = 4;
1120            const int nwarps    = ggml_sycl_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size;
1121            const int nbatch_fa = ggml_sycl_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
1122            launch_fattn<DV, cols_per_block/ncols2, ncols2,
1123                flash_attn_tile<DKQ, DV, cols_per_block / ncols2, ncols2, use_logit_softcap, warp_size>, warp_size>
1124                (ctx, dst, nwarps, nbytes_shared, nbatch_fa, true, true, false);
1125            return;
1126        }
1127    }
1128
1129    if constexpr (ncols2 <= 2) {
1130        constexpr int cols_per_block = 2;
1131        const int nwarps    = ggml_sycl_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size;
1132        const int nbatch_fa = ggml_sycl_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
1133        launch_fattn<DV, cols_per_block/ncols2, ncols2,
1134            flash_attn_tile<DKQ, DV, cols_per_block / ncols2, ncols2, use_logit_softcap, warp_size>, warp_size>
1135            (ctx, dst, nwarps, nbytes_shared, nbatch_fa, true, true, false);
1136        return;
1137    }
1138
1139    {
1140        constexpr int cols_per_block = ncols2*2;
1141        const int nwarps    = ggml_sycl_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size;
1142        const int nbatch_fa = ggml_sycl_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
1143        launch_fattn<DV, cols_per_block/ncols2, ncols2,
1144            flash_attn_tile<DKQ, DV, cols_per_block / ncols2, ncols2, use_logit_softcap, warp_size>, warp_size>
1145            (ctx, dst, nwarps, nbytes_shared, nbatch_fa, true, true, false);
1146        return;
1147    }
1148
1149    GGML_ABORT("fatal error");
1150}
1151
1152template <int DKQ, int DV, bool use_logit_softcap>
1153static void launch_fattn_tile_switch_ncols2(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
1154    const ggml_tensor * KQV  = dst;
1155    const ggml_tensor * Q    = dst->src[0];
1156    const ggml_tensor * K    = dst->src[1];
1157    const ggml_tensor * mask = dst->src[3];
1158
1159    float max_bias = 0.0f;
1160    memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float));
1161
1162    GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);
1163    const int gqa_ratio = Q->ne[2] / K->ne[2];
1164
1165    // On NVIDIA (Pascal and older) the GQA optimizations seem to be detrimental in some cases.
1166    // However, for DKQ == 576, DV == 512 only the kernel variant with GQA optimizations is implemented.
1167    //const bool nvidia = GGML_SYCL_CC_IS_NVIDIA(ggml_sycl_info().devices[ggml_sycl_get_device()].cc);
1168    const int gqa_limit = gqa_ratio <= 4 && DV <= 256 ? 16 : INT_MAX;
1169    const bool use_gqa_opt = mask && max_bias == 0.0f && Q->ne[1] <= gqa_limit && K->ne[1] % FATTN_KQ_STRIDE == 0;
1170
1171    if constexpr (DV == 512) {
1172        if (use_gqa_opt && gqa_ratio % 16 == 0) {
1173            launch_fattn_tile_switch_ncols1<DKQ, DV, 16, use_logit_softcap>(ctx, dst);
1174            return;
1175        }
1176        if (use_gqa_opt && gqa_ratio % 4 == 0) {
1177            launch_fattn_tile_switch_ncols1<DKQ, DV, 4, use_logit_softcap>(ctx, dst);
1178            return;
1179        }
1180        // ncols2=2 and ncols2=1 fallbacks only for cases where ncols=2 config exists (DKQ == DV).
1181        // For DKQ == 576, DV == 512 only GQA-optimized variants are implemented.
1182        if constexpr (DKQ == DV) {
1183            if (use_gqa_opt && gqa_ratio % 2 == 0) {
1184                launch_fattn_tile_switch_ncols1<DKQ, DV, 2, use_logit_softcap>(ctx, dst);
1185                return;
1186            }
1187            launch_fattn_tile_switch_ncols1<DKQ, DV, 1, use_logit_softcap>(ctx, dst);
1188            return;
1189        }
1190    }
1191
1192    if constexpr (DV <= 256) {
1193        if (use_gqa_opt && gqa_ratio % 8 == 0) {
1194            launch_fattn_tile_switch_ncols1<DKQ, DV, 8, use_logit_softcap>(ctx, dst);
1195            return;
1196        }
1197
1198        if (use_gqa_opt && gqa_ratio % 4 == 0) {
1199            launch_fattn_tile_switch_ncols1<DKQ, DV, 4, use_logit_softcap>(ctx, dst);
1200            return;
1201        }
1202
1203        if (use_gqa_opt && gqa_ratio % 2 == 0) {
1204            launch_fattn_tile_switch_ncols1<DKQ, DV, 2, use_logit_softcap>(ctx, dst);
1205            return;
1206        }
1207
1208        launch_fattn_tile_switch_ncols1<DKQ, DV, 1, use_logit_softcap>(ctx, dst);
1209        return;
1210    }
1211    GGML_ABORT("fatal error");
1212}
1213
1214template <int DKQ, int DV>
1215void ggml_sycl_flash_attn_ext_tile_case(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
1216    const ggml_tensor * KQV = dst;
1217
1218    float logit_softcap;
1219    memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
1220
1221    if (logit_softcap == 0.0f) {
1222        constexpr bool use_logit_softcap = false;
1223        launch_fattn_tile_switch_ncols2<DKQ, DV, use_logit_softcap>(ctx, dst);
1224    } else {
1225        constexpr bool use_logit_softcap = true;
1226        launch_fattn_tile_switch_ncols2<DKQ, DV, use_logit_softcap>(ctx, dst);
1227    }
1228}
1229
1230void ggml_sycl_flash_attn_ext_tile(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
1231
1232#define DECL_FATTN_TILE_CASE(DKQ, DV)                             \
1233    template void ggml_sycl_flash_attn_ext_tile_case              \
1234    <DKQ, DV>(ggml_backend_sycl_context & ctx, ggml_tensor * dst) \
1235
1236extern DECL_FATTN_TILE_CASE( 40,  40);
1237extern DECL_FATTN_TILE_CASE( 64,  64);
1238extern DECL_FATTN_TILE_CASE( 72,  72);
1239extern DECL_FATTN_TILE_CASE( 80,  80);
1240extern DECL_FATTN_TILE_CASE( 96,  96);
1241extern DECL_FATTN_TILE_CASE(112, 112);
1242extern DECL_FATTN_TILE_CASE(128, 128);
1243extern DECL_FATTN_TILE_CASE(256, 256);
1244extern DECL_FATTN_TILE_CASE(512, 512);
1245extern DECL_FATTN_TILE_CASE(576, 512);
1246