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

Aman GuptaCUDA + ggml: add sparse-fa for DSV4/GLM (#27970)8e93a97732026-09-02

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