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

uvosHIP: fattn-mma: use fp32 accumulation on MFMA devices (#28576)bfdc321832026-09-14

1#include "common.cuh"
2#include "cp-async.cuh"
3#include "mma.cuh"
4#include "fattn-common.cuh"
5#include "fattn-swizzle.cuh"
6
7using namespace ggml_cuda_mma;
8
9// Config options for the MMA kernel.
10// Should not affect results, only speed/register pressure/shared memory use.
11struct fattn_mma_config {
12    int  nthreads;       // Number of threads per CUDA block.
13    int  occupancy;      // Targeted occupancy for the MMA kernel.
14    int  nbatch_fa;      // Number of KV rows per softmax rescaling of KQ rowsums and VKQ accumulators.
15    int  nbatch_K2;      // Number of K half2 values in direction of DKQ to load in parallel.
16    int  nbatch_V2;      // Number of V half2 values in direction of DV to load in parallel.
17    int  nbatch_combine; // Number of VKQ half2 values in direction of DV to combine in parallel.
18    int  nstages_target; // Number of pipeline stages to use ideally, 1 == always load data synchronously, 2 == preload data if there is hardware support.
19    bool Q_in_reg;       // Whether the Q values should be kept permanently in registers.
20
21    constexpr __host__ __device__ fattn_mma_config(
22            int nthreads, int occupancy, int nbatch_fa, int nbatch_K2, int nbatch_V2, int nbatch_combine, int nstages_target, bool Q_in_reg) :
23        nthreads(nthreads), occupancy(occupancy), nbatch_fa(nbatch_fa), nbatch_K2(nbatch_K2), nbatch_V2(nbatch_V2), nbatch_combine(nbatch_combine),
24        nstages_target(nstages_target), Q_in_reg(Q_in_reg) {}
25};
26
27#define GGML_CUDA_FATTN_MMA_CONFIG_CASE(DKQ_, DV_, ncols_, nthreads_, occupancy_, nbatch_fa_, nbatch_K2_, nbatch_V2_, nbatch_combine_, nstages_target_, Q_in_reg_) \
28    if (DKQ == (DKQ_) && DV == (DV_) && ncols == (ncols_)) {                                                                                                       \
29        static_assert((nthreads_)       % 32 == 0 && (nthreads_)       <= 512, "bad nthreads");                                                                    \
30        static_assert(                               (occupancy_)      <=   8, "bad occupancy");                                                                   \
31        static_assert((nbatch_fa_)      % 32 == 0 && (nbatch_fa_)      <= 256, "bad nbatch_fa");                                                                   \
32        static_assert((nbatch_K2_)      %  4 == 0 && (nbatch_K2_)      <= 512, "bad nbatch_K2");                                                                   \
33        static_assert((nbatch_V2_)      %  4 == 0 && (nbatch_V2_)      <= 256, "bad nbatch_V2");                                                                   \
34        static_assert((nbatch_combine_) %  4 == 0 && (nbatch_combine_) <= 128, "bad nbatch_combine");                                                              \
35        static_assert((nstages_target_)      >= 1 && (nstages_target_) <=   2, "bad nstages_target");                                                              \
36        return fattn_mma_config{(nthreads_), (occupancy_), (nbatch_fa_), (nbatch_K2_), (nbatch_V2_), (nbatch_combine_), (nstages_target_), (Q_in_reg_)};           \
37    }                                                                                                                                                              \
38
39static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_config_ampere(const int DKQ, const int DV, const int ncols) {
40    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64,  8, 128, 2, 128,  32,  32,  32, 2, true);
41    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64, 16, 128, 2,  64,  32,  32,  32, 2, true);
42    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64, 32, 128, 2,  64,  32,  32,  32, 2, true);
43    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64, 64, 128, 2,  64,  32,  32,  32, 2, true);
44
45    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80,  8, 128, 2, 128,  40,  40,  40, 2, true);
46    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80, 16, 128, 2,  64,  40,  40,  40, 2, true);
47    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80, 32, 128, 2,  64,  40,  40,  40, 2, true);
48    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80, 64, 128, 2,  64,  40,  40,  40, 2, true);
49
50    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96,  8, 128, 2, 128,  48,  48,  48, 2, true);
51    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96, 16, 128, 2,  64,  48,  48,  48, 2, true);
52    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96, 32, 128, 2,  64,  48,  48,  48, 2, true);
53    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96, 64, 128, 2,  64,  48,  48,  48, 2, true);
54
55    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112,  8, 128, 2, 128,  56,  56,  56, 2, true);
56    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112, 16, 128, 2,  64,  56,  56,  56, 2, true);
57    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112, 32, 128, 2,  64,  56,  56,  56, 2, true);
58    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112, 64, 128, 2,  64,  56,  56,  56, 2, true);
59
60    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128,  8, 128, 2, 128,  64,  64,  64, 2, true);
61    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 16, 128, 2,  64,  64,  64,  64, 2, true);
62    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 32, 128, 2,  64,  64,  64,  64, 2, true);
63    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 64, 128, 2,  64,  64,  64,  64, 2, true);
64
65    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128,  8,  64, 4,  64,  96,  64,  64, 2, true);
66    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 16,  64, 4,  32,  96,  64,  64, 2, true);
67    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 32, 128, 2,  32,  96,  64,  64, 2, true);
68    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 64, 128, 2,  32,  96,  64,  64, 2, true);
69
70    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8, 128, 2,  64, 128, 128, 128, 2, true);
71    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16,  64, 4,  32, 128, 128, 128, 2, true);
72    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2,  32, 128, 128, 128, 2, true);
73    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2,  32, 128, 128, 128, 2, true);
74
75    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2,  32, 128, 128, 128, 1, false);
76    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 256, 1,  32, 128, 128, 128, 1, false);
77
78    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512,  8,  64, 4,  32, 256, 256, 128, 1, false);
79    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16,  64, 4,  32, 256, 256, 128, 1, false);
80    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 128, 2,  32, 128, 128, 128, 1, false);
81    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 64, 256, 1,  32, 128, 128, 128, 1, false);
82
83    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512,  8,  64, 4,  32, 288, 256, 128, 1, false);
84    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 16,  64, 4,  32, 288, 256, 128, 1, false);
85    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 32, 128, 2,  32, 160, 128, 128, 1, false);
86    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 64, 256, 1,  32, 160, 128, 128, 1, false);
87
88    return fattn_mma_config(32, 1, 0, 0, 0, 0, 0, false);
89}
90
91static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_config_turing(const int DKQ, const int DV, const int ncols) {
92    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8, 128, 2,  64, 128, 128, 128, 2, true);
93    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16, 128, 2,  64, 128, 128, 128, 2, true);
94    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2,  64, 128, 128,  64, 2, true);
95    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2,  64, 128, 128,  64, 2, true);
96
97    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2,  32, 128, 128, 128, 1, false);
98    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 256, 1,  32, 128, 128, 128, 1, false);
99
100    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512,  8,  64, 4,  32,  96,  64, 128, 1, false);
101    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16,  64, 4,  32,  96,  64, 128, 1, false);
102    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 128, 2,  32, 128, 128, 128, 1, false);
103    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 64, 256, 1,  32, 128, 128, 128, 1, false);
104
105    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512,  8,  64, 4,  32,  96,  64, 128, 1, false);
106    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 16,  64, 4,  32,  96,  64, 128, 1, false);
107    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 32, 128, 2,  32, 160, 128, 128, 1, false);
108    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 64, 256, 1,  32, 160, 128, 128, 1, false);
109
110    return ggml_cuda_fattn_mma_get_config_ampere(DKQ, DV, ncols);
111}
112
113static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_config_volta(const int DKQ, const int DV, const int ncols) {
114    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512,  8,  64, 4,  32, 256, 256,  64, 1, false);
115    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16,  64, 4,  32, 256, 256,  64, 1, false);
116    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 128, 2,  32, 128, 128,  64, 1, false);
117    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 64, 256, 1,  32, 128, 128,  64, 1, false);
118
119    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512,  8,  64, 4,  32, 288, 256,  64, 1, false);
120    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 16,  64, 4,  32, 288, 256,  64, 1, false);
121    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 32, 128, 2,  32, 160, 128,  64, 1, false);
122    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 64, 256, 1,  32, 160, 128,  64, 1, false);
123
124    // TODO tune specifically for Volta
125    return ggml_cuda_fattn_mma_get_config_ampere(DKQ, DV, ncols);
126}
127
128static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_config_rdna(const int DKQ, const int DV, const int ncols) {
129    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64,  8, 128, 2,  64,  32,  32,  32, 1, true);
130    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64, 16, 128, 2,  64,  32,  32,  32, 1, true);
131    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64, 32, 128, 2,  64,  32,  32,  32, 1, true);
132    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64, 64, 128, 2,  64,  32,  32,  32, 1, true);
133
134    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80,  8,  64, 2,  32,  40,  40,  40, 1, true);
135    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80, 16,  64, 2,  32,  40,  40,  40, 1, true);
136    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80, 32, 128, 2,  64,  40,  40,  40, 1, true);
137    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80, 64, 128, 2,  64,  40,  40,  40, 1, true);
138
139    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96,  8,  64, 2,  32,  48,  48,  48, 1, true);
140    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96, 16,  64, 2,  32,  48,  48,  48, 1, true);
141    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96, 32, 128, 2,  64,  48,  48,  48, 1, true);
142    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96, 64, 128, 2,  64,  48,  48,  48, 1, true);
143
144    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112,  8,  64, 2,  32,  56,  56,  56, 1, true);
145    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112, 16,  64, 2,  32,  56,  56,  56, 1, true);
146    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112, 32, 128, 2,  64,  56,  56,  56, 1, true);
147    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112, 64, 128, 2,  64,  56,  56,  56, 1, true);
148
149    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128,  8,  64, 2,  32,  64,  64,  64, 1, true);
150    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 16,  64, 2,  32,  64,  64,  64, 1, true);
151    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 32, 128, 2,  64,  64,  64,  64, 1, true);
152    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 64, 128, 2,  64,  64,  64,  64, 1, true);
153
154    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128,  8,  64, 2,  32,  96,  64,  64, 1, true);
155    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 16,  64, 2,  32,  96,  64,  64, 1, true);
156    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 32, 128, 2,  64,  96,  64,  64, 1, true);
157    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 64, 128, 2,  64,  96,  64,  64, 1, true);
158
159    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8,  64, 2,  32, 128, 128, 128, 1, true);
160    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16,  64, 2,  32, 128, 128, 128, 1, true);
161    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 256, 2,  64, 128, 128,  64, 1, true);
162    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 256, 2,  64, 128, 128,  64, 1, true);
163
164    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2,  32, 160, 128, 128, 1, true);
165    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 128, 2,  32, 160, 128, 128, 1, true);
166
167    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512,  8, 128, 3,  64,  96,  64, 128, 1, true);
168    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16, 128, 3,  64,  96,  64, 128, 1, true);
169    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 128, 2,  32, 128, 128, 128, 1, true);
170    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 64, 128, 2,  32, 128, 128, 128, 1, true);
171
172    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512,  8, 128, 3,  64,  96,  64, 128, 1, true);
173    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 16, 128, 3,  64,  96,  64, 128, 1, true);
174    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 32, 128, 2,  32, 160, 128, 128, 1, true);
175    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 64, 128, 2,  32, 160, 128, 128, 1, true);
176
177    return fattn_mma_config(32, 1, 0, 0, 0, 0, 0, false);
178}
179
180static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_config_cdna(const int DKQ, const int DV, const int ncols) {
181    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64,  8, 128, 1,  64,  32,  32,  32, 1, true);
182    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64, 16, 256, 2,  64,  32,  32,  32, 1, true);
183    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64, 32, 256, 2,  64,  32,  32,  32, 1, true);
184    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64,  64, 64, 256, 3,  64,  32,  32,  32, 1, true);
185
186    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80,  8, 256, 2,  64,  40,  40,  40, 1, true);
187    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80, 16, 256, 2,  64,  40,  40,  40, 1, true);
188    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80, 32, 256, 2,  64,  40,  40,  40, 1, true);
189    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80,  80, 64, 256, 2,  64,  40,  40,  40, 1, true);
190
191    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96,  8, 256, 2,  64,  48,  48,  48, 1, true);
192    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96, 16, 256, 2,  64,  48,  48,  48, 1, true);
193    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96, 32, 256, 2,  64,  48,  48,  48, 1, true);
194    GGML_CUDA_FATTN_MMA_CONFIG_CASE( 96,  96, 64, 256, 2,  64,  48,  48,  48, 1, true);
195
196    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112,  8, 256, 2,  64,  56,  56,  56, 1, true);
197    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112, 16, 256, 2,  64,  56,  56,  56, 1, true);
198    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112, 32, 256, 2,  64,  56,  56,  56, 1, true);
199    GGML_CUDA_FATTN_MMA_CONFIG_CASE(112, 112, 64, 256, 2,  64,  56,  56,  56, 1, true);
200
201    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128,  8, 256, 2,  64,  64,  64,  64, 1, true);
202    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 16, 256, 2,  64,  64,  64,  64, 1, true);
203    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 32, 256, 2,  64,  64,  64,  64, 1, true);
204    GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 64, 256, 2,  64,  64,  64,  64, 1, true);
205
206    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128,  8, 256, 1,  64,  64,  64,  64, 1, true);
207    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 16, 256, 1,  64,  64,  64,  64, 1, true);
208    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 32, 256, 1,  64,  64,  64,  64, 1, true);
209    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 64, 512, 1,  64,  64,  64,  64, 1, true);
210
211    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8, 256, 1,  64, 128, 128, 128, 1, true);
212    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16, 256, 1,  64, 128, 128, 128, 1, true);
213    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 256, 1,  64, 128, 128, 128, 1, true);
214    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 512, 1,  64, 128, 128,  64, 1, true);
215
216    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 256, 1,  64, 160, 128, 128, 1, true);
217    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 256, 1,  64, 160, 128, 128, 1, true);
218
219    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512,  8, 256, 1,  64, 128, 128, 128, 1, true);
220    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16, 256, 1,  64, 128, 128, 128, 1, true);
221    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 256, 1,  64, 128, 128, 128, 1, true);
222    GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 64, 256, 1,  64, 128, 128, 128, 1, true);
223
224    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512,  8, 256, 1,  64, 128, 128, 128, 1, true);
225    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 16, 256, 1,  64, 128, 128, 128, 1, true);
226    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 32, 256, 1,  64, 160, 128, 128, 1, true);
227    GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 64, 256, 1,  64, 160, 128, 128, 1, true);
228
229    return fattn_mma_config(32, 1, 0, 0, 0, 0, 0, false);
230}
231
232static __host__ fattn_mma_config ggml_cuda_fattn_mma_get_config(const int DKQ, const int DV, const int ncols, const int cc) {
233    if (ampere_mma_available(cc)) {
234        return ggml_cuda_fattn_mma_get_config_ampere(DKQ, DV, ncols);
235    }
236    if (turing_mma_available(cc)) {
237        return ggml_cuda_fattn_mma_get_config_turing(DKQ, DV, ncols);
238    }
239    if (amd_mfma_available(cc)) {
240        return ggml_cuda_fattn_mma_get_config_cdna(DKQ, DV, ncols);
241    }
242    if (amd_wmma_available(cc)) {
243        return ggml_cuda_fattn_mma_get_config_rdna(DKQ, DV, ncols);
244    }
245    GGML_ASSERT(volta_mma_available(cc));
246    return ggml_cuda_fattn_mma_get_config_volta(DKQ, DV, ncols);
247}
248
249static constexpr __device__ fattn_mma_config ggml_cuda_fattn_mma_get_config(const int DKQ, const int DV, const int ncols) {
250#if defined(AMPERE_MMA_AVAILABLE)
251    return ggml_cuda_fattn_mma_get_config_ampere(DKQ, DV, ncols);
252#elif defined(TURING_MMA_AVAILABLE)
253    return ggml_cuda_fattn_mma_get_config_turing(DKQ, DV, ncols);
254#elif defined(AMD_MFMA_AVAILABLE)
255    return ggml_cuda_fattn_mma_get_config_cdna(DKQ, DV, ncols);
256#elif defined(VOLTA_MMA_AVAILABLE)
257    return ggml_cuda_fattn_mma_get_config_volta(DKQ, DV, ncols);
258#elif defined(AMD_WMMA_AVAILABLE)
259    return ggml_cuda_fattn_mma_get_config_rdna(DKQ, DV, ncols);
260#else
261    GGML_UNUSED_VARS(DKQ, DV, ncols);
262    return fattn_mma_config(32, 1, 0, 0, 0, 0, 0, false);
263#endif // defined(AMPERE_MMA_AVAILABLE)
264}
265
266static __host__ int ggml_cuda_fattn_mma_get_nthreads(const int DKQ, const int DV, const int ncols, const int cc) {
267    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols, cc).nthreads;
268}
269
270static constexpr __device__ int ggml_cuda_fattn_mma_get_nthreads(const int DKQ, const int DV, const int ncols) {
271    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).nthreads;
272}
273
274static __host__ int ggml_cuda_fattn_mma_get_occupancy(const int DKQ, const int DV, const int ncols, const int cc) {
275    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols, cc).occupancy;
276}
277
278static constexpr __device__ int ggml_cuda_fattn_mma_get_occupancy(const int DKQ, const int DV, const int ncols) {
279    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).occupancy;
280}
281
282static __host__ int ggml_cuda_fattn_mma_get_nbatch_fa(const int DKQ, const int DV, const int ncols, const int cc) {
283    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols, cc).nbatch_fa;
284}
285
286static constexpr __device__ int ggml_cuda_fattn_mma_get_nbatch_fa(const int DKQ, const int DV, const int ncols) {
287    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).nbatch_fa;
288}
289
290static __host__ int ggml_cuda_fattn_mma_get_nbatch_K2(const int DKQ, const int DV, const int ncols, const int cc) {
291    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols, cc).nbatch_K2;
292}
293
294static constexpr __device__ int ggml_cuda_fattn_mma_get_nbatch_K2(const int DKQ, const int DV, const int ncols) {
295    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).nbatch_K2;
296}
297
298static __host__ int ggml_cuda_fattn_mma_get_nbatch_V2(const int DKQ, const int DV, const int ncols, const int cc) {
299    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols, cc).nbatch_V2;
300}
301
302static constexpr __device__ int ggml_cuda_fattn_mma_get_nbatch_V2(const int DKQ, const int DV, const int ncols) {
303    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).nbatch_V2;
304}
305
306static __host__ int ggml_cuda_fattn_mma_get_nbatch_combine(const int DKQ, const int DV, const int ncols, const int cc) {
307    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols, cc).nbatch_combine;
308}
309
310static constexpr __device__ int ggml_cuda_fattn_mma_get_nbatch_combine(const int DKQ, const int DV, const int ncols) {
311    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).nbatch_combine;
312}
313
314static __host__ int ggml_cuda_fattn_mma_get_nstages_target(const int DKQ, const int DV, const int ncols, const int cc) {
315    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols, cc).nstages_target;
316}
317
318static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages_target(const int DKQ, const int DV, const int ncols) {
319    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).nstages_target;
320}
321
322static __host__ bool ggml_cuda_fattn_mma_get_Q_in_reg(const int DKQ, const int DV, const int ncols, const int cc) {
323    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols, cc).Q_in_reg;
324}
325
326static constexpr __device__ bool ggml_cuda_fattn_mma_get_Q_in_reg(const int DKQ, const int DV, const int ncols) {
327    return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).Q_in_reg;
328}
329
330static constexpr __device__ int get_cols_per_thread() {
331#if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
332    return 1; // AMD has a single column per thread.
333#else
334    return 2; // This is specifically KQ columns, Volta only has a single VKQ column.
335#endif // defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
336}
337
338static __host__ int get_cols_per_warp(const int cc) {
339    if (turing_mma_available(cc) || amd_wmma_available(cc) || amd_mfma_available(cc)) {
340        return 16;
341    } else {
342        // Volta
343        return 32;
344    }
345}
346
347// ------------------------------------------------------------------------------------------------------------------
348
349static __host__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) {
350    return cp_async_available(cc) && ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2, cc) : 0;
351}
352
353static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(
354        const int DKQ, const int DV, const int ncols1, const int ncols2, const bool use_sparse) {
355#ifdef CP_ASYNC_AVAILABLE
356    const int nstages_target = ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0;
357    // sparse gather is not implemented for multi-stage loading
358    return use_sparse && nstages_target > 1 ? 1 : nstages_target;
359#else
360    GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2, use_sparse);
361    return 0;
362#endif // CP_ASYNC_AVAILABLE
363}
364
365// ------------------------------------------------------------------------------------------------------------------
366
367template<int stride_tile, bool swz, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check, bool use_sparse>
368static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
369        const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV,
370        const int k_VKQ_0, const int i_sup, const int32_t * const __restrict__ indices) {
371    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
372    // K/V data is loaded with decreasing granularity for D for better memory bandwidth.
373    // The minimum granularity is 16 bytes.
374    constexpr int h2_per_chunk = 16/sizeof(half2);
375    const int chunks_per_row = D2 / h2_per_chunk;
376    if constexpr (use_cp_async) {
377        static_assert(warp_size == 32, "bad warp_size");
378        static_assert(!oob_check || use_sparse, "OOB check not compatible with cp_async");
379        constexpr int preload = 64;
380
381        const unsigned int tile_KV_32 = ggml_cuda_cvta_generic_to_shared(tile_KV);
382
383        auto load = [&] __device__ (auto n) {
384            const int stride_k = warp_size >> n;
385            const int k0_start = stride_k == warp_size ? 0 : chunks_per_row - chunks_per_row % (2*stride_k);
386            const int k0_stop  =                             chunks_per_row - chunks_per_row % (1*stride_k);
387            const int stride_i = warp_size / stride_k;
388
389            if (k0_start == k0_stop) {
390                return;
391            }
392
393#pragma unroll
394            for (int i0 = 0; i0 < nbatch_fa; i0 += nwarps*stride_i) {
395                const int i = i0 + threadIdx.y*stride_i + (stride_k == warp_size ? 0 : threadIdx.x / stride_k);
396
397                if (i0 + nwarps*stride_i > nbatch_fa && i >= nbatch_fa) {
398                    break;
399                }
400
401                int64_t i_KV;
402                if constexpr (use_sparse) {
403                    // padded slots gather row 0, the -inf mask removes their contribution
404                    const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : 0;
405                    i_KV = index >= 0 ? index : 0;
406                } else {
407                    i_KV = k_VKQ_0 + i;
408                }
409
410#pragma unroll
411                for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
412                    const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);
413
414                    if constexpr (swz) {
415                        const int smem_offs_b = ggml_cuda_fattn_smem_swizzle::bytes_rc<stride_tile>(i, k*h2_per_chunk);
416                        cp_async_cg_16<preload>(tile_KV_32 + smem_offs_b, KV + i_KV*stride_KV + k*h2_per_chunk);
417                    } else {
418                        cp_async_cg_16<preload>(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i_KV*stride_KV + k*h2_per_chunk);
419                    }
420                }
421            }
422        };
423        // 1: max 32*16=512 bytes, 256 half
424        // 2: max 16*16=256 bytes, 128 half
425        // 3: max  8*16=128 bytes,  64 half
426        // 4: max  4*16= 64 bytes,  32 half
427        // 5: max  2*16= 32 bytes,  16 half
428        // 6: max  1*16= 16 bytes,   8 half
429        ggml_cuda_unroll<6>{}(load);
430    } else {
431        const half2 zero[4] = {{0.0f, 0.0f}, {0.0f, 0.0f}, {0.0f, 0.0f}, {0.0f, 0.0f}};
432        auto load = [&] __device__ (const int n) {
433            const int stride_k = 32 >> n;
434            const int k0_start = stride_k == 32 ? 0 : chunks_per_row - chunks_per_row % (2*stride_k);
435            const int k0_stop  =                      chunks_per_row - chunks_per_row % (1*stride_k);
436            const int stride_i = warp_size / stride_k;
437
438            if (k0_start == k0_stop) {
439                return;
440            }
441
442#pragma unroll
443            for (int i0 = 0; i0 < nbatch_fa; i0 += nwarps*stride_i) {
444                const int i = i0 + threadIdx.y*stride_i + (stride_k == warp_size ? 0 : threadIdx.x / stride_k);
445
446                if (i0 + nwarps*stride_i > nbatch_fa && i >= nbatch_fa) {
447                    break;
448                }
449
450#pragma unroll
451                for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
452                    const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);
453
454                    const half2 * src;
455                    if constexpr (use_sparse) {
456                        const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1;
457                        src = index >= 0 ? KV + int64_t(index)*stride_KV + k*h2_per_chunk : zero;
458                    } else {
459                        src = !oob_check || i < i_sup ? KV + int64_t(k_VKQ_0 + i)*stride_KV + k*h2_per_chunk : zero;
460                    }
461                    if constexpr (swz) {
462                        ggml_cuda_memcpy_1<16>((char *) tile_KV + ggml_cuda_fattn_smem_swizzle::bytes_rc<stride_tile>(i, k*h2_per_chunk), src);
463                    } else {
464                        ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, src);
465                    }
466                }
467            }
468        };
469        // 1: max 32*16=512 bytes, 256 half
470        // 2: max 16*16=256 bytes, 128 half
471        // 3: max  8*16=128 bytes,  64 half
472        // 4: max  4*16= 64 bytes,  32 half
473        // 5: max  2*16= 32 bytes,  16 half
474        // 6: max  1*16= 16 bytes,   8 half
475        ggml_cuda_unroll<6>{}(load);
476    }
477}
478
479template<int ncols1, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check, bool use_sparse>
480static __device__ __forceinline__ void flash_attn_ext_f16_load_mask(
481        const half * const __restrict__ mask_h, half * const __restrict__ tile_mask,
482        const int stride_mask, const int k_VKQ_0, const int i_sup, const int j0, const uint3 ne01,
483        const int32_t * const __restrict__ indices) {
484    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
485    if constexpr (use_cp_async) {
486        static_assert(nbatch_fa <= 8*warp_size && nbatch_fa % 8 == 0, "bad nbatch_fa");
487        static_assert(!oob_check, "OOB check incompatible with cp_async");
488        static_assert(!use_sparse, "sparse gather incompatible with cp_async");
489        constexpr int preload = nbatch_fa >= 32 ? nbatch_fa * sizeof(half) : 64;
490        constexpr int cols_per_warp = 8*warp_size/nbatch_fa;
491        constexpr int stride_j = nwarps * cols_per_warp;
492
493        const unsigned int tile_mask_32 = ggml_cuda_cvta_generic_to_shared(tile_mask);
494
495#pragma unroll
496        for (int j1 = 0; j1 < ncols1; j1 += stride_j) {
497            const int j_sram = j1 + threadIdx.y*cols_per_warp + threadIdx.x / (warp_size/cols_per_warp);
498            const int j_vram = fastmodulo(j0 + j_sram, ne01);
499
500            if (j1 + stride_j > ncols1 && j_sram >= ncols1) {
501                break;
502            }
503
504            const int i = 8 * (threadIdx.x % (nbatch_fa/8));
505
506            cp_async_cg_16<preload>(tile_mask_32 + j_sram*(nbatch_fa*sizeof(half) + 16) + i*sizeof(half), mask_h + int64_t(j_vram)*stride_mask + k_VKQ_0 + i);
507        }
508    } else if constexpr (oob_check || use_sparse) {
509#pragma unroll
510        for (int j1 = 0; j1 < ncols1; j1 += nwarps) {
511            const int j_sram = j1 + threadIdx.y;
512            const int j_vram = fastmodulo(j0 + j_sram, ne01);
513
514            if (j1 + nwarps > ncols1 && j_sram >= ncols1) {
515                break;
516            }
517
518#pragma unroll
519            for (int i0 = 0; i0 < nbatch_fa; i0 += warp_size) {
520                const int i = i0 + threadIdx.x;
521
522                if constexpr (use_sparse) {
523                    const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1;
524                    tile_mask[j_sram*(nbatch_fa + 8) + i] = index >= 0 ? mask_h[int64_t(j_vram)*stride_mask + index] : half(-INFINITY);
525                } else {
526                    tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + k_VKQ_0 + i] : half(0.0f);
527                }
528            }
529        }
530    } else if constexpr (nbatch_fa < 2*warp_size) {
531        constexpr int cols_per_warp = 2*warp_size/nbatch_fa;
532        constexpr int stride_j = nwarps * cols_per_warp;
533#pragma unroll
534        for (int j1 = 0; j1 < ncols1; j1 += stride_j) {
535            const int j_sram = j1 + threadIdx.y*cols_per_warp + threadIdx.x / (warp_size/cols_per_warp);
536            const int j_vram = fastmodulo(j0 + j_sram, ne01);
537
538            if (j1 + stride_j > ncols1 && j_sram >= ncols1) {
539                break;
540            }
541
542            const int i = threadIdx.x % (warp_size/cols_per_warp);
543
544            ggml_cuda_memcpy_1<sizeof(half2)>(tile_mask + j_sram*(nbatch_fa + 8) + 2*i, mask_h + int64_t(j_vram)*stride_mask + k_VKQ_0 + 2*i);
545        }
546    } else {
547#pragma unroll
548        for (int j1 = 0; j1 < ncols1; j1 += nwarps) {
549            const int j_sram = j1 + threadIdx.y;
550            const int j_vram = fastmodulo(j0 + j_sram, ne01);
551
552            if (j1 + nwarps > ncols1 && j_sram >= ncols1) {
553                break;
554            }
555
556#pragma unroll
557            for (int i0 = 0; i0 < nbatch_fa; i0 += 2*warp_size) {
558                const int i = i0 + 2*threadIdx.x;
559
560                ggml_cuda_memcpy_1<sizeof(half2)>(tile_mask + j_sram*(nbatch_fa + 8) + i, mask_h + int64_t(j_vram)*stride_mask + k_VKQ_0 + i);
561            }
562        }
563    }
564}
565
566template<int DKQ, int DV, int ncols1, int ncols2, int nwarps,
567    bool use_logit_softcap, bool V_is_K_view, bool use_sparse, bool needs_fixup, bool is_fixup, bool last_iter, bool oob_check,
568    typename T_A_KQ, typename T_B_KQ, typename T_C_KQ, typename T_A_VKQ, typename T_B_VKQ, typename T_C_VKQ>
569static __device__ __forceinline__ void flash_attn_ext_f16_iter(
570        const float2 * const __restrict__ Q_f2,
571        const half2  * const __restrict__ K_h2,
572        const half2  * const __restrict__ V_h2,
573        const half   * const __restrict__ mask_h,
574        const int32_t * const __restrict__ indices,
575        float2       * const __restrict__ dstk,
576        float2       * const __restrict__ dstk_fixup,
577        const float scale,
578        const float slope,
579        const float logit_softcap,
580        const uint3 ne01,
581        const int ne02,
582        const int stride_K,
583        const int stride_V,
584        const int stride_mask,
585        half2        * const __restrict__ tile_Q,
586        half2        * const __restrict__ tile_K,
587        half2        * const __restrict__ tile_V,
588        half         * const __restrict__ tile_mask,
589        T_B_KQ       * const __restrict__ Q_B,
590        T_C_VKQ      * const __restrict__ VKQ_C,
591        float        * const __restrict__ KQ_max,
592        float        * const __restrict__ KQ_rowsum,
593        const int jt,
594        const int kb0,
595        const int k_VKQ_sup) {
596#if defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
597    constexpr int  warp_size       = ggml_cuda_get_physical_warp_size();
598    constexpr int  ncols           = ncols1 * ncols2;
599    constexpr int  cols_per_warp   = T_B_KQ::I;
600    constexpr int  cols_per_thread = get_cols_per_thread();
601    constexpr int  np              = cols_per_warp > ncols ? nwarps : nwarps * cols_per_warp/ncols; // Number of parallel CUDA warps per Q column.
602    constexpr int  nbatch_fa       = ggml_cuda_fattn_mma_get_nbatch_fa(DKQ, DV, ncols);
603    constexpr int  nbatch_K2       = ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols);
604    constexpr int  nbatch_V2       = ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols);
605    constexpr bool Q_in_reg        = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols);
606    constexpr int  nstages         = ggml_cuda_fattn_mma_get_nstages  (DKQ, DV, ncols1, ncols2, use_sparse);
607
608    // swizzle the tile stride for K and V based on the batch size.
609    constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2);
610    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2);
611    constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2);
612    constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2);
613
614    const int k_VKQ_0 = kb0 * nbatch_fa;
615#if defined(TURING_MMA_AVAILABLE)
616    T_C_KQ KQ_C[nbatch_fa/(np*(cols_per_warp == 8 ? T_C_KQ::I : T_C_KQ::J))];
617#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
618    T_C_KQ KQ_C[nbatch_fa/(np*T_C_KQ::J)];
619#else // Volta
620    T_C_KQ KQ_C[nbatch_fa/(np*T_C_KQ::J)];
621#endif // defined(TURING_MMA_AVAILABLE)
622
623    if constexpr (nstages > 1) {
624        static_assert(!oob_check, "OOB check incompatible with multi-stage pipeline");
625        static_assert(!V_is_K_view, "K data reuse not implemented multi-stage loading");
626        static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi stage loading");
627        constexpr bool use_cp_async = true;
628        cp_async_wait_all();
629        __syncthreads();
630        flash_attn_ext_f16_load_tile<stride_tile_V, swz_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
631            (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, nullptr);
632    } else {
633        // the sparse mask values are gathered per element, always load them synchronously
634        constexpr bool use_cp_async = nstages == 1 && !use_sparse;
635        if (ncols2 > 1 || mask_h) {
636            flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
637                (mask_h, tile_mask, stride_mask, k_VKQ_0, k_VKQ_sup, jt*ncols1, ne01, indices);
638        }
639    }
640
641    // For MLA K and V have the same data.
642    // Therefore, iterate over K in reverse and later re-use the data if possible.
643#pragma unroll
644    for (int k0_start = (DKQ/2-1) - (DKQ/2-1) % nbatch_K2; k0_start >= 0; k0_start -= nbatch_K2) {
645        const int k0_stop = k0_start + nbatch_K2 < DKQ/2 ? k0_start + nbatch_K2 : DKQ/2;
646
647        if constexpr (nstages <= 1) {
648            const int k0_diff = k0_stop - k0_start;
649            constexpr bool use_cp_async = nstages == 1;
650            flash_attn_ext_f16_load_tile<stride_tile_K, swz_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
651                (K_h2 + k0_start, tile_K, k0_diff, stride_K, k_VKQ_0, k_VKQ_sup, indices);
652            if (use_cp_async) {
653                cp_async_wait_all();
654            }
655            __syncthreads();
656        }
657
658        // Calculate tile of KQ:
659        if constexpr (Q_in_reg) {
660#pragma unroll
661            for (int i_KQ_00 = 0; i_KQ_00 < nbatch_fa; i_KQ_00 += np*T_A_KQ::I) {
662                const int i_KQ_0 = i_KQ_00 + (threadIdx.y % np)*T_A_KQ::I;
663#pragma unroll
664                for (int k_KQ_0 = k0_start; k_KQ_0 < k0_stop; k_KQ_0 += T_A_KQ::J) {
665                    T_A_KQ K_A;
666                    ggml_cuda_fattn_smem_swizzle::load_ldmatrix<stride_tile_K, swz_K>(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start);
667                    if constexpr (cols_per_warp == 8) {
668                        mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[k_KQ_0/T_A_KQ::J]);
669                    } else {
670                        // Wide version of KQ_C is column-major
671#if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
672                        // AMD matrix C is column-major.
673                        mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[k_KQ_0/T_A_KQ::J]);
674#else
675                        // swap A and B for CUDA.
676                        mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], Q_B[k_KQ_0/T_A_KQ::J], K_A);
677#endif // defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
678                    }
679                }
680            }
681        } else {
682            constexpr int stride_tile_Q = DKQ/2 + 4;
683#pragma unroll
684            for (int k_KQ_0 = k0_start; k_KQ_0 < k0_stop; k_KQ_0 += T_A_KQ::J) {
685                load_ldmatrix(Q_B[0], tile_Q + (threadIdx.y / np)*(T_B_KQ::I*stride_tile_Q) + k_KQ_0, stride_tile_Q);
686
687#pragma unroll
688                for (int i_KQ_00 = 0; i_KQ_00 < nbatch_fa; i_KQ_00 += np*T_A_KQ::I) {
689                    const int i_KQ_0 = i_KQ_00 + (threadIdx.y % np)*T_A_KQ::I;
690
691                    T_A_KQ K_A;
692                    ggml_cuda_fattn_smem_swizzle::load_ldmatrix<stride_tile_K, swz_K>(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start);
693
694                    if constexpr (cols_per_warp == 8) {
695                        mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[0]);
696                    } else {
697                        // Wide version of KQ_C is column-major
698#if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
699                        // AMD matrix C is column-major.
700                        mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[0]);
701#else
702                        // swap A and B for CUDA.
703                        mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], Q_B[0], K_A);
704#endif // defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
705                    }
706                }
707            }
708        }
709
710        if constexpr (nstages <= 1) {
711            __syncthreads(); // Only needed if tile_K == tile_V.
712        }
713    }
714
715    if (use_logit_softcap) {
716        constexpr int stride = cols_per_warp == 8 ? np*T_C_KQ::I : np*T_C_KQ::J;
717        static_assert(nbatch_fa % stride == 0, "bad loop size");
718#pragma unroll
719        for (int i = 0; i < nbatch_fa/stride; ++i) {
720#pragma unroll
721            for (int l = 0; l < T_C_KQ::ne; ++l) {
722                KQ_C[i].x[l] = logit_softcap*tanhf(KQ_C[i].x[l]);
723            }
724        }
725    }
726
727    float KQ_max_new[cols_per_thread];
728#pragma unroll
729    for (int col = 0; col < cols_per_thread; ++col) {
730        KQ_max_new[col] = KQ_max[col];
731    }
732    float KQ_rowsum_add[cols_per_thread] = {0.0f};
733
734    if constexpr (cols_per_warp == 8) {
735        if (ncols2 > 1 || mask_h) {
736#pragma unroll
737            for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::I) {
738                const int i0 = i00 + (threadIdx.y % np)*T_C_KQ::I;
739#pragma unroll
740                for (int l = 0; l < T_C_KQ::ne; ++l) {
741                    const int i = i0 + T_C_KQ::get_i(l);
742                    const int j = ((threadIdx.y / np)*T_C_KQ::J + T_C_KQ::get_j(l)) / ncols2;
743
744                    KQ_C[i00/(np*T_C_KQ::I)].x[l] += slope * __half2float(tile_mask[j*(nbatch_fa + 8) + i]);
745                }
746            }
747        }
748
749        // Calculate softmax for each KQ column using the current max. value.
750        // The divisor is stored in KQ_rowsum and will be applied at the end.
751        static_assert(nbatch_fa % (np*T_C_KQ::I) == 0, "bad loop size");
752#pragma unroll
753        for (int k0 = 0; k0 < nbatch_fa; k0 += np*T_C_KQ::I) {
754#pragma unroll
755            for (int l = 0; l < T_C_KQ::ne; ++l) {
756                if (!oob_check || k0 + (threadIdx.y % np)*T_C_KQ::I + T_C_KQ::get_i(l) < k_VKQ_sup) {
757#if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
758                    constexpr int KQ_idx = 0;
759#else
760                    // Turing + Volta:
761                    const int KQ_idx = l % 2;
762#endif // defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
763                    KQ_max_new[KQ_idx] = fmaxf(KQ_max_new[KQ_idx], KQ_C[k0/(np*T_C_KQ::I)].x[l] + FATTN_KQ_MAX_OFFSET);
764                }
765            }
766        }
767
768        // Values per KQ column are spread across 8 threads:
769#pragma unroll
770        for (int col = 0; col < cols_per_thread; ++col) {
771#pragma unroll
772            for (int offset = 16; offset >= 4; offset >>= 1) {
773                KQ_max_new[col] = fmaxf(KQ_max_new[col], __shfl_xor_sync(0xFFFFFFFF, KQ_max_new[col], offset, warp_size));
774            }
775        }
776
777        static_assert(nbatch_fa % (np*T_C_KQ::I) == 0, "bad loop size");
778#pragma unroll
779        for (int k0 = 0; k0 < nbatch_fa; k0 += np*T_C_KQ::I) {
780#pragma unroll
781            for (int l = 0; l < T_C_KQ::ne; ++l) {
782                if (!oob_check || k0 + (threadIdx.y % np)*T_C_KQ::I + T_C_KQ::get_i(l) < k_VKQ_sup) {
783#if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
784                    constexpr int KQ_idx = 0;
785#else
786                    // Turing + Volta:
787                    const int KQ_idx = l % 2;
788#endif // defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
789                    KQ_C[k0/(np*T_C_KQ::I)].x[l] = expf(KQ_C[k0/(np*T_C_KQ::I)].x[l] - KQ_max_new[KQ_idx]);
790                    KQ_rowsum_add[KQ_idx] += KQ_C[k0/(np*T_C_KQ::I)].x[l];
791                } else {
792                    KQ_C[k0/(np*T_C_KQ::I)].x[l] = 0.0f;
793                }
794            }
795        }
796    } else { // not Turing mma or T_B_KQ::I > 8
797        if (ncols2 > 1 || mask_h) {
798#pragma unroll
799            for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::J) {
800                const int i0 = i00 + (threadIdx.y % np)*T_C_KQ::J;
801
802                // The mask is stored as 16 bit half values, loading them as 32 bit half2 values is preferred in terms of speed.
803                // However, this is not possible for RDNA3 where 2 consecutive l indices are not consecutive in the mask memory layout.
804#ifdef RDNA3
805#pragma unroll
806                for (int l = 0; l < T_C_KQ::ne; ++l) {
807                    const int i = i0 + T_C_KQ::get_j(l);
808                    const int j = ((threadIdx.y / np)*cols_per_warp + T_C_KQ::get_i(l)) / ncols2;
809
810                    KQ_C[i00/(np*T_C_KQ::J)].x[l] += __half2float(tile_mask[j*(nbatch_fa + 8) + i]);
811                }
812#else
813#pragma unroll
814                for (int l0 = 0; l0 < T_C_KQ::ne; l0 += 2) {
815                    const int i = (i0 + T_C_KQ::get_j(l0)) / 2;
816                    const int j = ((threadIdx.y / np)*cols_per_warp + T_C_KQ::get_i(l0)) / ncols2;
817
818                    const float2 tmp = __half22float2(((const half2 *)tile_mask)[j*(nbatch_fa/2 + 4) + i]);
819                    KQ_C[i00/(np*T_C_KQ::J)].x[l0 + 0] += slope*tmp.x;
820                    KQ_C[i00/(np*T_C_KQ::J)].x[l0 + 1] += slope*tmp.y;
821                }
822#endif // RDNA3
823            }
824        }
825
826        // Calculate softmax for each KQ column using the current max. value.
827        // The divisor is stored in KQ_rowsum and will be applied at the end.
828        static_assert(nbatch_fa % (np*T_C_KQ::J) == 0, "bad loop size");
829#pragma unroll
830        for (int k0 = 0; k0 < nbatch_fa; k0 += np*T_C_KQ::J) {
831#pragma unroll
832            for (int l = 0; l < T_C_KQ::ne; ++l) {
833                if (!oob_check || k0 + (threadIdx.y % np)*T_C_KQ::J + T_C_KQ::get_j(l) < k_VKQ_sup) {
834#if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
835                    constexpr int KQ_idx = 0;
836#else
837                    // Turing + Volta:
838                    const int KQ_idx = (l/2) % 2;
839#endif // defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
840                    KQ_max_new[KQ_idx] = fmaxf(KQ_max_new[KQ_idx], KQ_C[(k0/(np*T_C_KQ::J))].x[l] + FATTN_KQ_MAX_OFFSET);
841                }
842            }
843        }
844
845#pragma unroll
846        for (int col = 0; col < cols_per_thread; ++col) {
847#if defined(TURING_MMA_AVAILABLE)
848            // Values per KQ column are spread across 4 threads:
849            constexpr int offset_first = 2;
850            constexpr int offset_last  = 1;
851#elif defined(AMD_MFMA_AVAILABLE)
852            // MFMA: 4 threads per Q column (threadIdx.x % 16 == col, spaced by 16).
853            constexpr int offset_first = 32;
854            constexpr int offset_last  = 16;
855#elif defined(AMD_WMMA_AVAILABLE)
856            // Values per KQ column are spread across 2 threads:
857            constexpr int offset_first = 16;
858            constexpr int offset_last  = 16;
859#else // Volta
860            // Values per KQ column are spread across 2 threads:
861            constexpr int offset_first = 2;
862            constexpr int offset_last  = 2;
863#endif // defined(TURING_MMA_AVAILABLE)
864#pragma unroll
865            for (int offset = offset_first; offset >= offset_last; offset >>= 1) {
866                KQ_max_new[col] = fmaxf(KQ_max_new[col], __shfl_xor_sync(0xFFFFFFFF, KQ_max_new[col], offset, warp_size));
867            }
868        }
869
870        static_assert(nbatch_fa % (np*T_C_KQ::J) == 0, "bad loop size");
871#pragma unroll
872        for (int k0 = 0; k0 < nbatch_fa; k0 += np*T_C_KQ::J) {
873#pragma unroll
874            for (int l = 0; l < T_C_KQ::ne; ++l) {
875                if (!oob_check || k0 + (threadIdx.y % np)*T_C_KQ::J + T_C_KQ::get_j(l) < k_VKQ_sup) {
876#if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
877                    constexpr int KQ_idx = 0;
878#else
879                    // Turing + Volta:
880                    const int KQ_idx = (l/2) % 2;
881#endif // defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
882                    KQ_C[(k0/(np*T_C_KQ::J))].x[l] = expf(KQ_C[(k0/(np*T_C_KQ::J))].x[l] - KQ_max_new[KQ_idx]);
883                    KQ_rowsum_add[KQ_idx] += KQ_C[(k0/(np*T_C_KQ::J))].x[l];
884                } else {
885                    KQ_C[(k0/(np*T_C_KQ::J))].x[l] = 0.0f;
886                }
887            }
888        }
889    }
890
891    {
892        float KQ_max_scale[cols_per_thread];
893#pragma unroll
894        for (int col = 0; col < cols_per_thread; ++col) {
895            const float KQ_max_diff = KQ_max[col] - KQ_max_new[col];
896            KQ_max_scale[col] = expf(KQ_max_diff);
897            KQ_max[col] = KQ_max_new[col];
898
899            *((uint32_t *) &KQ_max_scale[col]) *= KQ_max_diff >= SOFTMAX_FTZ_THRESHOLD;
900
901            // Scale previous KQ_rowsum to account for a potential increase in KQ_max:
902            KQ_rowsum[col] = KQ_max_scale[col]*KQ_rowsum[col] + KQ_rowsum_add[col];
903        }
904
905#if defined(TURING_MMA_AVAILABLE)
906        if constexpr (cols_per_warp == 8) {
907            const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[cols_per_thread - 1]);
908#pragma unroll
909            for (int i = 0; i < DV/T_C_VKQ::I; ++i) {
910#pragma unroll
911                for (int l = 0; l < T_C_VKQ::ne; ++l) {
912                    VKQ_C[i].x[l] *= KQ_max_scale_h2;
913                }
914            }
915        } else {
916#pragma unroll
917            for (int col = 0; col < cols_per_thread; ++col) {
918                const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[col], KQ_max_scale[col]);
919#pragma unroll
920                for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
921#pragma unroll
922                    for (int l0 = 0; l0 < T_C_VKQ::ne; l0 += 2) {
923                        VKQ_C[i].x[l0 + col] *= KQ_max_scale_h2;
924                    }
925                }
926            }
927        }
928#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
929        if constexpr (std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>) {
930            const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[0]);
931#pragma unroll
932            for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
933#pragma unroll
934                for (int l = 0; l < T_C_VKQ::ne; ++l) {
935                    VKQ_C[i].x[l] *= KQ_max_scale_h2;
936                }
937            }
938        } else {
939            static_assert(std::is_same_v<decltype(T_C_VKQ::x), float[T_C_VKQ::ne]>, "bad VKQ type");
940#pragma unroll
941            for (int i = 0; i < DV/T_C_VKQ::J; ++i) {
942#pragma unroll
943                for (int l = 0; l < T_C_VKQ::ne; ++l) {
944                    VKQ_C[i].x[l] *= KQ_max_scale[0];
945                }
946            }
947        }
948#else // Volta
949        const half2 KQ_max_scale_h2 = make_half2(
950            KQ_max_scale[(threadIdx.x / 2) % 2], KQ_max_scale[(threadIdx.x / 2) % 2]);
951#pragma unroll
952        for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
953#pragma unroll
954            for (int l = 0; l < T_C_VKQ::ne; ++l) {
955                VKQ_C[i].x[l] *= KQ_max_scale_h2;
956            }
957        }
958#endif // defined(TURING_MMA_AVAILABLE)
959    }
960
961    // Convert KQ C tiles into B tiles for VKQ calculation:
962    T_B_VKQ B[nbatch_fa/(np*2*T_B_VKQ::J)];
963    static_assert(nbatch_fa % (np*2*T_B_VKQ::J) == 0, "bad loop size");
964    if constexpr (cols_per_warp == 8) {
965#pragma unroll
966        for (int k = 0; k < nbatch_fa/(np*2*T_B_VKQ::J); ++k) {
967            B[k] = get_transposed(get_half2(KQ_C[k]));
968        }
969    } else {
970        for (int k = 0; k < nbatch_fa/(np*2*T_B_VKQ::J); ++k) {
971            B[k] = get_half2(KQ_C[k]);
972        }
973    }
974
975    if constexpr (nstages > 1) {
976        static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading");
977        static_assert(!V_is_K_view, "K data reuse not implemented multi-stage loading");
978        // Preload K tile for next iteration:
979        constexpr bool use_cp_async = true;
980        cp_async_wait_all();
981        __syncthreads();
982        if (!last_iter) {
983            if (ncols2 > 1 || mask_h) {
984                flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
985                    (mask_h, tile_mask, stride_mask, k_VKQ_0 + nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
986            }
987            flash_attn_ext_f16_load_tile<stride_tile_K, swz_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
988                (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, nullptr);
989        }
990    }
991
992
993    // Calculate VKQ tile, need to use logical rather than physical elements for i0 due to transposition of V:
994#pragma unroll
995    for (int i0_start = 0; i0_start < DV; i0_start += 2*nbatch_V2) {
996        static_assert(DV % (2*nbatch_V2) == 0, "bad loop size");
997        const int i0_stop = i0_start + 2*nbatch_V2;
998
999        if constexpr (nstages <= 1) {
1000            const int i0_diff = i0_stop - i0_start;
1001            if (!V_is_K_view || i0_stop > 2*nbatch_K2) {
1002                constexpr bool use_cp_async = nstages == 1;
1003                flash_attn_ext_f16_load_tile<stride_tile_V, swz_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
1004                    (V_h2 + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_0, k_VKQ_sup, indices);
1005                if (use_cp_async) {
1006                    cp_async_wait_all();
1007                }
1008                __syncthreads();
1009            }
1010        }
1011        const half2 * tile_V_i = !V_is_K_view || i0_stop > 2*nbatch_K2 ? tile_V : tile_V + i0_start/2;
1012
1013#if defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
1014#pragma unroll
1015        for (int i_VKQ_0 = i0_start; i_VKQ_0 < i0_stop; i_VKQ_0 += T_A_VKQ::I) {
1016            static_assert((nbatch_fa/2) % (np*T_A_VKQ::J) == 0, "bad loop size");
1017#pragma unroll
1018            for (int k00 = 0; k00 < nbatch_fa/2; k00 += np*T_A_VKQ::J) {
1019                const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::J;
1020
1021                T_A_VKQ A; // Transposed in SRAM but not in registers, gets transposed on load.
1022                ggml_cuda_fattn_smem_swizzle::load_ldmatrix_trans<stride_tile_V, swz_V>(A, tile_V, (int)(tile_V_i - tile_V) + 2*k0*stride_tile_V + (i_VKQ_0 - i0_start)/2);
1023                if constexpr (T_B_KQ::I == 8) {
1024                    mma(VKQ_C[i_VKQ_0/T_A_VKQ::I], A, B[k00/(np*T_A_VKQ::J)]);
1025                } else {
1026                    // Wide version of VKQ_C is column-major.
1027#if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
1028                    // AMD matrix C is column-major.
1029                    mma(VKQ_C[i_VKQ_0/T_A_VKQ::I], A, B[k00/(np*T_A_VKQ::J)]);
1030#else
1031                    // swap A and B for CUDA.
1032                    mma(VKQ_C[i_VKQ_0/T_A_VKQ::I], B[k00/(np*T_A_VKQ::J)], A);
1033#endif // defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
1034                }
1035            }
1036        }
1037#else // Volta
1038        constexpr int i0_stride = 2*T_C_VKQ::J;
1039#pragma unroll
1040        for (int i_VKQ_0 = i0_start; i_VKQ_0 < i0_stop; i_VKQ_0 += i0_stride) {
1041            static_assert(nbatch_fa % (np*T_A_VKQ::I) == 0, "bad loop size");
1042            static_assert(2*T_B_VKQ::J == T_A_VKQ::I, "bad tile sizes");
1043#pragma unroll
1044            for (int k00 = 0; k00 < nbatch_fa; k00 += np*T_A_VKQ::I) {
1045                const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::I;
1046
1047                T_A_VKQ A; // Transposed in both SRAM and registers, load normally.
1048                ggml_cuda_fattn_smem_swizzle::load_ldmatrix<stride_tile_V, swz_V>(A, tile_V, (int)(tile_V_i - tile_V) + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2);
1049                mma(VKQ_C[i_VKQ_0/i0_stride], B[k00/(np*T_A_VKQ::I)], A);
1050            }
1051        }
1052#endif // defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
1053
1054        if constexpr (nstages <= 1) {
1055            __syncthreads(); // Only needed if tile_K == tile_V.
1056        }
1057    }
1058#else
1059    GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup,
1060        scale, slope, logit_softcap, ne01, ne02,
1061        stride_K, stride_V, stride_mask,
1062        tile_Q, tile_K, tile_V, tile_mask,
1063        Q_B, VKQ_C, KQ_max, KQ_rowsum, kb0);
1064    NO_DEVICE_CODE;
1065#endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
1066}
1067
1068#if defined(TURING_MMA_AVAILABLE)
1069template<int DV, int ncols> struct mma_tile_sizes {
1070    using T_A_KQ  = tile<16,  8, half2>; // row-major
1071    using T_B_KQ  = tile<16,  8, half2>; // column-major
1072    using T_C_KQ  = tile<16, 16, float>; // column-major
1073    using T_A_VKQ = tile<16,  8, half2>; // row-major
1074    using T_B_VKQ = tile<16,  8, half2>; // column-major
1075    using T_C_VKQ = tile<16,  8, half2>; // column-major
1076};
1077template<int DV> struct mma_tile_sizes<DV, 8> {
1078    using T_A_KQ  = tile<16,  8, half2>; // row-major
1079    using T_B_KQ  = tile< 8,  8, half2>; // column-major
1080    using T_C_KQ  = tile<16,  8, float>; // row-major
1081    using T_A_VKQ = tile<16,  8, half2>; // row-major
1082    using T_B_VKQ = tile< 8,  8, half2>; // column-major
1083    using T_C_VKQ = tile<16,  4, half2>; // row-major
1084};
1085#elif defined(AMD_WMMA_AVAILABLE)
1086#ifdef RDNA3
1087template<int DV, int ncols> struct mma_tile_sizes {
1088    using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
1089    using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
1090    using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
1091    using T_A_VKQ = tile<32,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
1092    using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
1093    using T_C_VKQ = tile<16, 16, half2, DATA_LAYOUT_I_MAJOR>;          // column-major
1094};
1095template<int ncols> struct mma_tile_sizes<80, ncols> {
1096    using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
1097    using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
1098    using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
1099    using T_A_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
1100    using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
1101    using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
1102};
1103template<int ncols> struct mma_tile_sizes<112, ncols> {
1104    using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
1105    using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
1106    using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
1107    using T_A_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
1108    using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
1109    using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
1110};
1111#else
1112template<int DV, int ncols> struct mma_tile_sizes {
1113    using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR>;           // row-major
1114    using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR>;           // column-major
1115    using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;           // column-major
1116    using T_A_VKQ = tile<32,  8, half2, DATA_LAYOUT_I_MAJOR>;           // row-major
1117    using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR>;           // column-major
1118    using T_C_VKQ = tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED>; // column-major
1119};
1120template<int ncols> struct mma_tile_sizes<80, ncols> {
1121    using T_A_KQ  = tile<16,  8, half2>; // row-major
1122    using T_B_KQ  = tile<16,  8, half2>; // column-major
1123    using T_C_KQ  = tile<16, 16, float>; // column-major
1124    using T_A_VKQ = tile<16,  8, half2>; // row-major
1125    using T_B_VKQ = tile<16,  8, half2>; // column-major
1126    using T_C_VKQ = tile<16,  8, half2>; // column-major
1127};
1128template<int ncols> struct mma_tile_sizes<112, ncols> {
1129    using T_A_KQ  = tile<16,  8, half2>; // row-major
1130    using T_B_KQ  = tile<16,  8, half2>; // column-major
1131    using T_C_KQ  = tile<16, 16, float>; // column-major
1132    using T_A_VKQ = tile<16,  8, half2>; // row-major
1133    using T_B_VKQ = tile<16,  8, half2>; // column-major
1134    using T_C_VKQ = tile<16,  8, half2>; // column-major
1135};
1136#endif // RDNA3
1137#elif defined(AMD_MFMA_AVAILABLE)
1138template<int DV, int ncols> struct mma_tile_sizes {
1139    using T_A_KQ  = tile<16,  8, half2>; // row-major
1140    using T_B_KQ  = tile<16,  8, half2>; // column-major
1141    using T_C_KQ  = tile<16, 16, float>; // column-major
1142    using T_A_VKQ = tile<16,  8, half2>; // row-major
1143    using T_B_VKQ = tile<16,  8, half2>; // column-major
1144    using T_C_VKQ = tile<16, 16, float>; // column-major
1145};
1146#else // Volta
1147template<int DV, int ncols> struct mma_tile_sizes {
1148    using T_A_KQ  = tile< 8,  4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
1149    using T_B_KQ  = tile<32,  4, half2, DATA_LAYOUT_I_MAJOR>;          // column-major
1150    using T_C_KQ  = tile<32,  8, float, DATA_LAYOUT_I_MAJOR>;          // column-major
1151    using T_A_VKQ = tile< 8,  4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED>; // column-major
1152    using T_B_VKQ = tile<32,  4, half2, DATA_LAYOUT_I_MAJOR>;          // column-major
1153    using T_C_VKQ = tile<32,  4, half2, DATA_LAYOUT_I_MAJOR>;          // column-major
1154};
1155#endif // defined(TURING_MMA_AVAILABLE)
1156
1157template<int DKQ, int DV, int ncols1, int ncols2, int nwarps, bool use_logit_softcap, bool V_is_K_view, bool use_sparse, bool needs_fixup, bool is_fixup>
1158static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
1159        const float2 * const __restrict__ Q_f2,
1160        const half2  * const __restrict__ K_h2,
1161        const half2  * const __restrict__ V_h2,
1162        const half   * const __restrict__ mask_h,
1163        const int32_t * const __restrict__ indices,
1164        const float  * const __restrict__ sinks_f,
1165        float2       * const __restrict__ dstk,
1166        float2       * const __restrict__ dstk_fixup,
1167        const float scale,
1168        const float slope,
1169        const float logit_softcap,
1170        const uint3 ne01,
1171        const int ne02,
1172        const int gqa_ratio,
1173        const int ne11,
1174        const int stride_Q1,
1175        const int stride_Q2,
1176        const int stride_K,
1177        const int stride_V,
1178        const int stride_mask,
1179        const int jt,
1180        const int zt_gqa,
1181        const int kb0_start,
1182        const int kb0_stop) {
1183#if defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
1184    //In this kernel Q, K, V are matrices while i, j, k are matrix indices.
1185
1186    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
1187    constexpr int ncols = ncols1 * ncols2;
1188    using     T_A_KQ    = typename mma_tile_sizes<DV, ncols>::T_A_KQ;
1189    using     T_B_KQ    = typename mma_tile_sizes<DV, ncols>::T_B_KQ;
1190    using     T_C_KQ    = typename mma_tile_sizes<DV, ncols>::T_C_KQ;
1191    using     T_A_VKQ   = typename mma_tile_sizes<DV, ncols>::T_A_VKQ;
1192    using     T_B_VKQ   = typename mma_tile_sizes<DV, ncols>::T_B_VKQ;
1193    using     T_C_VKQ   = typename mma_tile_sizes<DV, ncols>::T_C_VKQ;
1194
1195    constexpr int  cols_per_warp   = T_B_KQ::I;
1196    constexpr int  cols_per_thread = get_cols_per_thread();
1197    constexpr int  np              = cols_per_warp > ncols ? nwarps : nwarps * cols_per_warp/ncols; // Number of parallel CUDA warps per Q column.
1198    constexpr int  nbatch_fa       = ggml_cuda_fattn_mma_get_nbatch_fa     (DKQ, DV, ncols);
1199    constexpr int  nbatch_K2       = ggml_cuda_fattn_mma_get_nbatch_K2     (DKQ, DV, ncols);
1200    constexpr int  nbatch_V2       = ggml_cuda_fattn_mma_get_nbatch_V2     (DKQ, DV, ncols);
1201    constexpr int  nbatch_combine  = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols);
1202    constexpr bool Q_in_reg        = ggml_cuda_fattn_mma_get_Q_in_reg      (DKQ, DV, ncols);
1203    constexpr int  nstages         = ggml_cuda_fattn_mma_get_nstages       (DKQ, DV, ncols1, ncols2, use_sparse);
1204
1205    if (cols_per_warp > ncols) {
1206        NO_DEVICE_CODE;
1207        return;
1208    }
1209
1210    static_assert(nwarps * (cols_per_warp/ncols2) % ncols1 == 0, "bad nwarps");
1211
1212    constexpr int stride_tile_Q = DKQ/2     + 4;
1213    // swizzle the tile stride for K and V based on the batch size.
1214    constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2);
1215    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2);
1216    constexpr int stride_tile_KV_max = stride_tile_K > stride_tile_V ? stride_tile_K : stride_tile_V;
1217    constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2);
1218    constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2);
1219
1220    extern __shared__ half2 tile_Q[];
1221    half2 * tile_K    = Q_in_reg              ? tile_Q                             : tile_Q + ncols     * stride_tile_Q;
1222    half2 * tile_V    =           nstages > 1 ? tile_K + nbatch_fa * stride_tile_K : tile_K;
1223    half  * tile_mask = (half *) (nstages > 1 ? tile_V + nbatch_fa * stride_tile_V : tile_V + nbatch_fa * stride_tile_KV_max);
1224
1225    T_B_KQ    Q_B[(Q_in_reg ? DKQ/(2*T_B_KQ::J) : 1)];
1226#if defined(TURING_MMA_AVAILABLE)
1227    T_C_VKQ VKQ_C[cols_per_warp == 8 ? DV/T_C_VKQ::I : DV/(2*T_C_VKQ::J)];
1228#elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
1229    T_C_VKQ VKQ_C[DV % 32 != 0       ? DV/T_C_VKQ::J : DV/(2*T_C_VKQ::J)];
1230#elif defined(AMD_MFMA_AVAILABLE)
1231    T_C_VKQ VKQ_C[                                     DV/T_C_VKQ::J];
1232#elif defined(AMD_WMMA_AVAILABLE)
1233    T_C_VKQ VKQ_C[                                     DV/(2*T_C_VKQ::J)];
1234#else // Volta
1235    T_C_VKQ VKQ_C[                                     DV/(2*T_C_VKQ::J)];
1236#endif // defined(TURING_MMA_AVAILABLE)
1237
1238    float KQ_rowsum[cols_per_thread] = {0.0f};
1239    float KQ_max[cols_per_thread];
1240#pragma unroll
1241    for (int col = 0; col < cols_per_thread; ++col) {
1242        KQ_max[col] = -FLT_MAX/2.0f;
1243    }
1244
1245    // Load Q data into tile_Q, either temporarily or permanently.
1246    // Q in registers is faster, but register pressure is the biggest bottleneck.
1247    // The loading is done with decreasing granularity for D for better memory bandwidth.
1248    const half2 scale_h2 = make_half2(scale, scale);
1249#pragma unroll
1250    for (int stride_k : {warp_size, warp_size/2, warp_size/4, warp_size/8}) {
1251        const int k0_start  = stride_k == warp_size ? 0 : DKQ/2 - (DKQ/2) % (2*stride_k);
1252        const int k0_stop   =                             DKQ/2 - (DKQ/2) % (1*stride_k);
1253        const int stride_jc = warp_size / stride_k;
1254
1255        if (k0_start == k0_stop) {
1256            continue;
1257        }
1258
1259#pragma unroll
1260        for (int jc0 = 0; jc0 < ncols; jc0 += nwarps*stride_jc) {
1261            const int jc = jc0 + threadIdx.y*stride_jc + (stride_k == warp_size ? 0 : threadIdx.x / stride_k);
1262
1263            if (jc0 + nwarps*stride_jc > ncols && jc >= ncols) {
1264                break;
1265            }
1266
1267            const int j = jc / ncols2;
1268            const int c = jc % ncols2;
1269
1270            if ((ncols1 == 1 || jt*ncols1 + j < int(ne01.z)) && (ncols2 == 1 || zt_gqa*ncols2 + c < gqa_ratio)) {
1271#pragma unroll
1272                for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
1273                    const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);
1274
1275                    const float2 tmp = Q_f2[(jt*ncols1 + j)*stride_Q1 + c*stride_Q2 + k];
1276                    tile_Q[jc*stride_tile_Q + k] = scale_h2 * make_half2(tmp.x, tmp.y);
1277                }
1278            } else {
1279#pragma unroll
1280                for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
1281                    const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);
1282
1283                    tile_Q[jc*stride_tile_Q + k] = make_half2(0.0f, 0.0f);
1284                }
1285            }
1286        }
1287    }
1288
1289    __syncthreads();
1290
1291    if (Q_in_reg) {
1292        const int j0 = (threadIdx.y / np) * cols_per_warp;
1293
1294#pragma unroll
1295        for (int k0 = 0; k0 < DKQ/2; k0 += T_B_KQ::J) {
1296            load_ldmatrix(Q_B[k0/T_B_KQ::J], tile_Q + j0*stride_tile_Q + k0, stride_tile_Q);
1297        }
1298    }
1299
1300    __syncthreads();
1301
1302    int kb0 = kb0_start;
1303
1304    // Preload mask and K data for first iteration when using cp_async with multiple stages:
1305    if constexpr (nstages > 1) {
1306        static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading");
1307        static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline");
1308        constexpr bool use_cp_async = true;
1309        constexpr bool oob_check    = false;
1310        constexpr int  k_VKQ_sup    = nbatch_fa;
1311        if (ncols2 > 1 || mask_h) {
1312            flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
1313                (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
1314        }
1315        flash_attn_ext_f16_load_tile<stride_tile_K, swz_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
1316            (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr);
1317    }
1318
1319    // kb0_start is always < kb0_stop so the last iter can be executed unconditionally.
1320    if constexpr (ncols2 == 1 || use_sparse) {
1321        constexpr bool oob_check = true;
1322        for (; kb0 < kb0_stop-1; ++kb0) {
1323            constexpr bool last_iter = false;
1324            constexpr int  k_VKQ_sup = nbatch_fa;
1325            flash_attn_ext_f16_iter
1326                <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup, last_iter, oob_check,
1327                 T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ>
1328                (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap,
1329                 ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C,
1330                 KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup);
1331        }
1332        constexpr bool last_iter = true;
1333        const     int  k_VKQ_sup = ne11 - kb0*nbatch_fa;
1334        flash_attn_ext_f16_iter
1335            <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup, last_iter, oob_check,
1336              T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ>
1337            (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap,
1338             ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C,
1339             KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup);
1340    } else {
1341        constexpr bool oob_check = false;
1342        for (; kb0 < kb0_stop-1; ++kb0) {
1343            constexpr bool last_iter = false;
1344            constexpr int  k_VKQ_sup = nbatch_fa;
1345            flash_attn_ext_f16_iter
1346                <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup, last_iter, oob_check,
1347                 T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ>
1348                (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap,
1349                 ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C,
1350                 KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup);
1351        }
1352        constexpr bool last_iter = true;
1353        constexpr int  k_VKQ_sup = nbatch_fa;
1354        flash_attn_ext_f16_iter
1355            <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup, last_iter, oob_check,
1356             T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ>
1357            (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap,
1358             ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C,
1359             KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup);
1360    }
1361
1362    // With multi-stage loading there is no __syncthreads at the end of the iter,
1363    //     there can be a race condition on shared memory access for combining/writing back results.
1364    if constexpr (nstages > 1 && nwarps*cols_per_warp > nbatch_fa) {
1365        __syncthreads();
1366    }
1367
1368    // Finally, sum up partial KQ rowsums.
1369    {
1370#if defined(TURING_MMA_AVAILABLE)
1371        // The partial sums are spread across 8/4 threads.
1372        constexpr int offset_first = cols_per_warp == 8 ? 16 : 2;
1373        constexpr int offset_last  = cols_per_warp == 8 ?  4 : 1;
1374#elif defined(AMD_MFMA_AVAILABLE)
1375        // The partial sums are spread across 4 threads (wavefront64, 16 cols).
1376        constexpr int offset_first = 32;
1377        constexpr int offset_last  = 16;
1378#elif defined(AMD_WMMA_AVAILABLE)
1379        // The partial sums are spread across 2 threads.
1380        constexpr int offset_first = 16;
1381        constexpr int offset_last  = 16;
1382#else // Volta
1383        // The partial sums are spread across 2 threads.
1384        constexpr int offset_first = 2;
1385        constexpr int offset_last  = 2;
1386#endif // defined(TURING_MMA_AVAILABLE)
1387#pragma unroll
1388        for (int col = 0; col < cols_per_thread; ++col) {
1389#pragma unroll
1390            for (int offset = offset_first; offset >= offset_last; offset >>= 1) {
1391                KQ_rowsum[col] += __shfl_xor_sync(0xFFFFFFFF, KQ_rowsum[col], offset, warp_size);
1392            }
1393        }
1394    }
1395
1396    // If attention sinks are used, potentially re-scale if KQ_max is small.
1397    // Also add the sink as a value to KQ_rowsum, this is done after synchronization of KQ_rowsum
1398    //     so it's being done unconditionally for every thread.
1399    if (!is_fixup && (np == 1 || threadIdx.y % np == 0) && sinks_f) {
1400        float KQ_max_scale[cols_per_thread];
1401#pragma unroll
1402        for (int col = 0; col < cols_per_thread; ++col) {
1403            const int jc = (threadIdx.y/np)*cols_per_warp + (cols_per_warp == 8 ? T_C_KQ::get_j(col) : T_C_KQ::get_i(2*col));
1404            const float sink = sinks_f[jc % ncols2];
1405
1406            const float KQ_max_new = fmaxf(KQ_max[col], sink);
1407            const float KQ_max_diff = KQ_max[col] - KQ_max_new;
1408            KQ_max_scale[col] = expf(KQ_max_diff);
1409            KQ_max[col] = KQ_max_new;
1410
1411            *((uint32_t *) &KQ_max_scale[col]) *= KQ_max_diff >= SOFTMAX_FTZ_THRESHOLD;
1412
1413            const float KQ_max_add = expf(sink - KQ_max_new);
1414            KQ_rowsum[col] = KQ_max_scale[col]*KQ_rowsum[col] + KQ_max_add;
1415        }
1416
1417#if defined(TURING_MMA_AVAILABLE)
1418        if constexpr (cols_per_warp == 8) {
1419            const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[cols_per_thread - 1]);
1420#pragma unroll
1421            for (int i = 0; i < DV/T_C_VKQ::I; ++i) {
1422#pragma unroll
1423                for (int l = 0; l < T_C_VKQ::ne; ++l) {
1424                    VKQ_C[i].x[l] *= KQ_max_scale_h2;
1425                }
1426            }
1427        } else {
1428#pragma unroll
1429            for (int col = 0; col < cols_per_thread; ++col) {
1430                const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[col], KQ_max_scale[col]);
1431#pragma unroll
1432                for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
1433#pragma unroll
1434                    for (int l0 = 0; l0 < T_C_VKQ::ne; l0 += 2) {
1435                        VKQ_C[i].x[l0 + col] *= KQ_max_scale_h2;
1436                    }
1437                }
1438            }
1439        }
1440#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
1441        if constexpr (std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>) {
1442            const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[0]);
1443#pragma unroll
1444            for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
1445#pragma unroll
1446                for (int l = 0; l < T_C_VKQ::ne; ++l) {
1447                    VKQ_C[i].x[l] *= KQ_max_scale_h2;
1448                }
1449            }
1450        } else {
1451            static_assert(std::is_same_v<decltype(T_C_VKQ::x), float[T_C_VKQ::ne]>, "bad VKQ type");
1452#pragma unroll
1453            for (int i = 0; i < DV/T_C_VKQ::J; ++i) {
1454#pragma unroll
1455                for (int l = 0; l < T_C_VKQ::ne; ++l) {
1456                    VKQ_C[i].x[l] *= KQ_max_scale[0];
1457                }
1458            }
1459        }
1460#else // Volta
1461        const int col = (threadIdx.x / 2) % 2;
1462        const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[col], KQ_max_scale[col]);
1463#pragma unroll
1464        for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
1465#pragma unroll
1466            for (int l = 0; l < T_C_VKQ::ne; ++l) {
1467                VKQ_C[i].x[l] *= KQ_max_scale_h2;
1468            }
1469        }
1470#endif // defined(TURING_MMA_AVAILABLE)
1471    }
1472
1473    // Combine VKQ accumulator values if np > 1.
1474    // It's also faster to do small writes to shared memory, then large write to VRAM than to do small writes to VRAM.
1475    // So also write VKQ accumulators to shared memory in column-major format if np == 1.
1476
1477    constexpr int tile_stride = nbatch_combine + 4;
1478    static_assert((DV/2) % nbatch_combine == 0, "bad nbatch_combine");
1479
1480    constexpr bool combine_needs_sync = swz_K || swz_V;
1481
1482    if constexpr (cols_per_warp == 8) {
1483        const int jc_cwmo = (threadIdx.x % (2*T_C_VKQ::J)) / T_C_VKQ::J; // jc combine write meta offset
1484        const int jc_cwm = threadIdx.y*(2*T_C_VKQ::J) + 2*T_C_VKQ::get_j(-1) + jc_cwmo; // jc combine write meta
1485        const float2 KQ_cmr = make_float2(KQ_max[jc_cwmo], KQ_rowsum[jc_cwmo]); // KQ combine max rowsum
1486
1487        if constexpr (combine_needs_sync) {
1488            __syncthreads();
1489        }
1490
1491        if (((!needs_fixup && !is_fixup) || np > 1) && threadIdx.x < 2*T_C_VKQ::J) {
1492            // Use the 16 bytes of padding in each row to store the meta data: KQ max, KQ rowsum, KQ max scale.
1493            ((float2 *) tile_Q)[jc_cwm*(tile_stride/2) + nbatch_combine/2] = KQ_cmr;
1494        }
1495
1496        __syncthreads();
1497
1498        if (np == 1) {
1499            // No combination is needed, the meta data can be directly written from registers to VRAM.
1500            if (needs_fixup && threadIdx.x < T_B_KQ::I) {
1501                float2 * dstk_fixup_meta = dstk_fixup + blockIdx.x*ncols;
1502                dstk_fixup_meta[jc_cwm] = KQ_cmr;
1503            }
1504            if (is_fixup && threadIdx.x < T_B_KQ::I) {
1505                float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols;
1506                dstk_fixup_meta[jc_cwm] = KQ_cmr;
1507            }
1508        }
1509    } else {
1510        // jc_cwm = jc combine write meta
1511        // KQ_cmr = KQ combine max rowsum
1512        // Use the 16 bytes of padding in each Q column to store the meta data: KQ max, KQ rowsum, KQ max scale.
1513#if defined(TURING_MMA_AVAILABLE)
1514        const int jc_cwm = threadIdx.y*cols_per_warp + T_C_VKQ::get_i(threadIdx.x % 4);
1515        const float2 KQ_cmr = make_float2(KQ_max[threadIdx.x % cols_per_thread], KQ_rowsum[threadIdx.x % cols_per_thread]);
1516        const bool thread_should_write = threadIdx.x % 4 < cols_per_thread;
1517#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
1518        const int jc_cwm = threadIdx.y*cols_per_warp + T_C_VKQ::get_i(0);
1519        const float2 KQ_cmr = make_float2(KQ_max[0], KQ_rowsum[0]);
1520        const bool thread_should_write = threadIdx.x / 16 < cols_per_thread;
1521#else // Volta
1522        const int jc_cwm = threadIdx.y*cols_per_warp + T_C_KQ::get_i(threadIdx.x & 2);
1523        const float2 KQ_cmr = make_float2(KQ_max[(threadIdx.x & 2) / 2], KQ_rowsum[(threadIdx.x & 2) / 2]);
1524        const bool thread_should_write = T_C_KQ::J == 8 || T_C_KQ::get_j(threadIdx.x & 2) < 8;
1525#endif // defined(TURING_MMA_AVAILABLE)
1526
1527        if constexpr (combine_needs_sync) {
1528            __syncthreads();
1529        }
1530
1531        if (((!needs_fixup && !is_fixup) || np > 1) && thread_should_write) {
1532            ((float2 *) tile_Q)[jc_cwm*(tile_stride/2) + nbatch_combine/2] = KQ_cmr;
1533        }
1534
1535        __syncthreads();
1536
1537        if (np == 1) {
1538            // No combination is needed, the meta data can be directly written from registers to VRAM.
1539            if (needs_fixup && thread_should_write) {
1540                float2 * dstk_fixup_meta = dstk_fixup + blockIdx.x*ncols;
1541                dstk_fixup_meta[jc_cwm] = KQ_cmr;
1542            }
1543            if (is_fixup && thread_should_write) {
1544                float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols;
1545                dstk_fixup_meta[jc_cwm] = KQ_cmr;
1546            }
1547        }
1548    }
1549
1550    if (np > 1) {
1551        constexpr int nmeta = np*cols_per_warp >= warp_size ? np*cols_per_warp/warp_size : 1;
1552
1553        float KQ_cmn;
1554        float KQ_cms[nmeta];
1555        float KQ_crs;
1556
1557        const int jc_meta = threadIdx.y*cols_per_warp + (np*cols_per_warp < warp_size ? threadIdx.x % (np*cols_per_warp) : threadIdx.x);
1558        float2 * const meta_ptr = ((float2 *) tile_Q) + jc_meta*(tile_stride/2) + nbatch_combine/2;
1559
1560        if (threadIdx.y % np == 0) {
1561            // Combine the meta data for parallel warps via shared memory.
1562            float2 meta[nmeta];
1563#pragma unroll
1564            for (int imeta = 0; imeta < nmeta; ++imeta) {
1565                meta[imeta] = meta_ptr[imeta * warp_size * tile_stride/2];
1566            }
1567
1568            KQ_cmn = meta[0].x; // KQ combine max new, max between all parallel warps.
1569#pragma unroll
1570            for (int imeta = 1; imeta < nmeta; ++imeta) {
1571                KQ_cmn = fmaxf(KQ_cmn, meta[imeta].x);
1572            }
1573#pragma unroll
1574            for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) {
1575                if (offset < warp_size) {
1576                    KQ_cmn = fmaxf(KQ_cmn, __shfl_xor_sync(0xFFFFFFFF, KQ_cmn, offset, warp_size));
1577                }
1578            }
1579
1580#pragma unroll
1581            for (int imeta = 0; imeta < nmeta; ++imeta) {
1582                KQ_cms[imeta] = expf(meta[imeta].x - KQ_cmn);
1583            }
1584
1585            KQ_crs = KQ_cms[0]*meta[0].y; // KQ combine rowsum, scaled sum of all parallel warps.
1586#pragma unroll
1587            for (int imeta = 1; imeta < nmeta; ++imeta) {
1588                KQ_crs += KQ_cms[imeta]*meta[imeta].y;
1589            }
1590#pragma unroll
1591            for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) {
1592                if (offset < warp_size) {
1593                    KQ_crs += __shfl_xor_sync(0xFFFFFFFF, KQ_crs, offset, warp_size);
1594                }
1595            }
1596        }
1597
1598        __syncthreads();
1599
1600        if (threadIdx.y % np == 0) {
1601            // Write back combined meta data:
1602#pragma unroll
1603            for (int imeta = 0; imeta < nmeta; ++imeta) {
1604                if (np*cols_per_warp >= warp_size || threadIdx.x < np*cols_per_warp) {
1605                    // Combined KQ max scale + rowsum.
1606                    meta_ptr[imeta * warp_size * tile_stride/2] = make_float2(KQ_cms[imeta], KQ_crs);
1607                }
1608            }
1609
1610            // Combined KQ max + rowsum.
1611            static_assert(cols_per_warp <= warp_size);
1612            if (needs_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) {
1613                float2 * dstk_fixup_meta = dstk_fixup + blockIdx.x*ncols;
1614                dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs);
1615            }
1616            if (is_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) {
1617                float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols;
1618                dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs);
1619            }
1620        }
1621    }
1622
1623#pragma unroll
1624    for (int k00 = 0; k00 < DV/2; k00 += nbatch_combine) {
1625        if constexpr (cols_per_warp == 8) {
1626            static_assert(std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>, "bad VKQ type");
1627            const int jc_cwd = threadIdx.y*T_B_KQ::I + T_B_KQ::get_i(-1); // jc combine write data
1628#pragma unroll
1629            for (int k1 = 0; k1 < nbatch_combine; k1 += T_B_KQ::J) {
1630                const T_B_KQ B = get_transposed(VKQ_C[(k00 + k1)/T_B_KQ::J]); // Conversion of C to B matrix puts it in column-major format.
1631
1632#pragma unroll
1633                for (int l = 0; l < T_B_KQ::ne; ++l) {
1634                    const int k = k1 + T_B_KQ::get_j(l);
1635
1636                    tile_Q[jc_cwd*tile_stride + k] = B.x[l];
1637                }
1638            }
1639        } else {
1640            const int j0 = threadIdx.y*cols_per_warp;
1641            if constexpr (std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>) {
1642                if constexpr (T_C_VKQ::dl == DATA_LAYOUT_I_MAJOR) {
1643#pragma unroll
1644                    for (int k1 = 0; k1 < nbatch_combine; k1 += T_C_VKQ::J) {
1645#pragma unroll
1646                        for (int l = 0; l < T_C_VKQ::ne; ++l) {
1647                            const int j = j0 + T_C_VKQ::get_i(l);
1648                            const int k = k1 + T_C_VKQ::get_j(l);
1649
1650                            tile_Q[j*tile_stride + k] = VKQ_C[(k00 + k1)/T_C_VKQ::J].x[l];
1651                        }
1652                    }
1653                } else {
1654                    static_assert(T_C_VKQ::dl == DATA_LAYOUT_I_MAJOR_SCRAMBLED, "bad T_C_VKQ data layout");
1655                    using T_C_VKQ_us = tile<T_C_VKQ::I, T_C_VKQ::J, half2, DATA_LAYOUT_I_MAJOR>; // us == unscrambled
1656#pragma unroll
1657                    for (int k1 = 0; k1 < nbatch_combine; k1 += T_C_VKQ::J) {
1658                        const T_C_VKQ_us VKQ_C_us = unscramble(VKQ_C[(k00 + k1)/T_C_VKQ::J]);
1659#pragma unroll
1660                        for (int l = 0; l < T_C_VKQ_us::ne; ++l) {
1661                            const int j = j0 + T_C_VKQ_us::get_i(l);
1662                            const int k = k1 + T_C_VKQ_us::get_j(l);
1663
1664                            tile_Q[j*tile_stride + k] = VKQ_C_us.x[l];
1665                        }
1666                    }
1667                }
1668            } else {
1669                static_assert(std::is_same_v<decltype(T_C_VKQ::x), float[T_C_VKQ::ne]>, "bad VKQ type");
1670                half * tile_Q_h = (half *) tile_Q;
1671#pragma unroll
1672                for (int k1 = 0; k1 < nbatch_combine; k1 += T_C_VKQ::J/2) {
1673#pragma unroll
1674                    for (int l = 0; l < T_C_VKQ::ne; ++l) {
1675                        const int j = j0 + T_C_VKQ::get_i(l);
1676                        const int k = 2*k1 + T_C_VKQ::get_j(l);
1677
1678                        tile_Q_h[j*(2*tile_stride) + k] = VKQ_C[(k00 + k1)/(T_C_VKQ::J/2)].x[l];
1679                    }
1680                }
1681            }
1682        }
1683
1684        __syncthreads();
1685
1686        if (np == 1 || threadIdx.y % np == 0) {
1687            // The first 2*2*gridDim.x*ncols floats in dstk_fixup are for storing max. values and row sums.
1688            // The values after that are for the partial results of the individual blocks.
1689            float2 * dstk_fixup_data = dstk_fixup + gridDim.x*(2*ncols) + blockIdx.x*(ncols*(DV/2));
1690
1691#pragma unroll
1692            for (int stride_k : {warp_size, warp_size/2, warp_size/4, warp_size/8}) {
1693                const int k0_start  = stride_k == warp_size ? 0 : nbatch_combine - nbatch_combine % (2*stride_k);
1694                const int k0_stop   =                             nbatch_combine - nbatch_combine % (1*stride_k);
1695                const int stride_jc = warp_size / stride_k;
1696
1697                if (k0_start == k0_stop) {
1698                    continue;
1699                }
1700
1701#pragma unroll
1702                for (int jc0_dst = 0; jc0_dst < ncols; jc0_dst += (nwarps/np)*stride_jc) {
1703                    const int jc_dst = jc0_dst + (threadIdx.y/np)*stride_jc + (stride_k == warp_size ? 0 : threadIdx.x / stride_k);
1704
1705                    if (jc0_dst + (nwarps/np)*stride_jc > ncols && jc_dst >= ncols) {
1706                        break;
1707                    }
1708
1709                    const int jc_tile_K = (jc_dst/cols_per_warp)*(np*cols_per_warp) + jc_dst % cols_per_warp;
1710
1711                    const int j_dst = jc_dst / ncols2;
1712                    const int c_dst = jc_dst % ncols2;
1713
1714                    if (!is_fixup && ((ncols1 > 1 && jt*ncols1 + j_dst >= int(ne01.z)) || (ncols2 > 1 && zt_gqa*ncols2 + c_dst >= gqa_ratio))) {
1715                        continue;
1716                    }
1717
1718                    const float * meta_j = (const float *) tile_Q + jc_tile_K*tile_stride + nbatch_combine;
1719#pragma unroll
1720                    for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
1721                        const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);
1722
1723                        float2 dstk_val = make_float2(0.0f, 0.0f);
1724#pragma unroll
1725                        for (int ip = 0; ip < np; ++ip) {
1726                            const float KQ_crs = np == 1 ? 1.0f : meta_j[ip*cols_per_warp * tile_stride + 0];
1727                            const float2 dstk_val_add = __half22float2(tile_Q[(jc_tile_K + ip*cols_per_warp) * tile_stride + k]);
1728                            dstk_val.x += dstk_val_add.x*KQ_crs;
1729                            dstk_val.y += dstk_val_add.y*KQ_crs;
1730                        }
1731
1732                        if (!needs_fixup && !is_fixup) {
1733                            const float KQ_rowsum_j = meta_j[1];
1734                            dstk_val.x /= KQ_rowsum_j;
1735                            dstk_val.y /= KQ_rowsum_j;
1736                        }
1737
1738                        if (is_fixup) {
1739                            dstk_fixup_data[jc_dst*(DV/2) + k00 + k] = dstk_val;
1740                        } else {
1741                            dstk[((jt*ncols1 + j_dst)*ne02 + c_dst)*(DV/2) + k00 + k] = dstk_val;
1742                        }
1743                    }
1744                }
1745            }
1746        }
1747        if (np > 1) {
1748            __syncthreads();
1749        }
1750    }
1751#else
1752    GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dstk_fixup,
1753        scale, slope, logit_softcap, ne01, ne02, gqa_ratio,
1754        stride_Q1, stride_Q2, stride_K, stride_V, stride_mask,
1755        jt, kb0_start, kb0_stop);
1756    NO_DEVICE_CODE;
1757#endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
1758}
1759
1760static constexpr __host__ __device__ bool ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(
1761        const int DKQ, const int DV, const int ncols1, const int ncols2) {
1762    return (DKQ == 512 && DV == 512 && ncols1 == 1 && ncols2 == 8) ||
1763           (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16);
1764}
1765
1766template<int DKQ, int DV, int ncols1, int ncols2, bool use_logit_softcap, bool V_is_K_view, bool use_sparse>
1767__launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2))
1768static __global__ void flash_attn_ext_f16(
1769        const char * Q_ptr,
1770        const char * K_ptr,
1771        const char * V_ptr,
1772        const char * mask_ptr,
1773        const char * sinks_ptr,
1774        const int  * KV_max_ptr,
1775        float      * dst_ptr,
1776        float2     * dst_meta_ptr,
1777        const float scale,
1778        const float max_bias,
1779        const float m0,
1780        const float m1,
1781        const uint32_t n_head_log2,
1782        const float logit_softcap,
1783        const int32_t ne00, const uint3   ne01, const int32_t ne02, const int32_t ne03,
1784                            const int32_t nb01, const int32_t nb02, const int32_t nb03,
1785        const int32_t ne10, const int32_t ne11, const int32_t ne12, const int32_t ne13,
1786                            const int32_t nb11, const int32_t nb12, const int64_t nb13,
1787                            const int32_t nb21, const int32_t nb22, const int64_t nb23,
1788                            const int32_t ne31, const int32_t ne32, const int32_t ne33,
1789                            const int32_t nb31, const int32_t nb32, const int64_t nb33) {
1790    ggml_cuda_pdl_sync(); // TODO optimize placement
1791#if defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE))
1792    const char * GGML_CUDA_RESTRICT Q              = Q_ptr;
1793    const char * GGML_CUDA_RESTRICT K              = K_ptr;
1794    const char * GGML_CUDA_RESTRICT V              = V_ptr;
1795    const char * GGML_CUDA_RESTRICT mask           = mask_ptr;
1796    const char * GGML_CUDA_RESTRICT sinks          = sinks_ptr;
1797    const int  * GGML_CUDA_RESTRICT KV_max         = use_sparse ? nullptr : KV_max_ptr;
1798    const int  * GGML_CUDA_RESTRICT sparse_indices = use_sparse ? KV_max_ptr : nullptr;
1799    float      * GGML_CUDA_RESTRICT dst            = dst_ptr;
1800    float2     * GGML_CUDA_RESTRICT dst_meta       = dst_meta_ptr;
1801
1802    // Skip unused kernel variants for faster compilation:
1803    if (use_logit_softcap && !(DKQ == 128 || DKQ == 256 || DKQ == 512)) {
1804        NO_DEVICE_CODE;
1805        return;
1806    }
1807    if (DKQ == 192 && ncols2 != 8 && ncols2 != 16) {
1808        NO_DEVICE_CODE;
1809        return;
1810    }
1811
1812    if (!ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2) && use_sparse) {
1813        NO_DEVICE_CODE;
1814        return;
1815    }
1816#ifdef VOLTA_MMA_AVAILABLE
1817    if (ncols1*ncols2 < 32) {
1818        NO_DEVICE_CODE;
1819        return;
1820    }
1821#endif // VOLTA_MMA_AVAILABLE
1822
1823#if __CUDA_ARCH__ == GGML_CUDA_CC_TURING
1824    if (ncols1*ncols2 > 32) {
1825        NO_DEVICE_CODE;
1826        return;
1827    }
1828#endif // __CUDA_ARCH__ == GGML_CUDA_CC_TURING
1829
1830#if defined(AMD_WMMA_AVAILABLE)
1831    if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) {
1832        NO_DEVICE_CODE;
1833        return;
1834    }
1835#endif // defined(AMD_WMMA_AVAILABLE)
1836
1837#if defined(AMD_MFMA_AVAILABLE)
1838    if (ncols1*ncols2 < 16 || DKQ > 256) {
1839        NO_DEVICE_CODE;
1840        return;
1841    }
1842#endif // defined(AMD_MFMA_AVAILABLE)
1843
1844    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
1845    constexpr int ncols     = ncols1 * ncols2;
1846    constexpr int nbatch_fa = ggml_cuda_fattn_mma_get_nbatch_fa(DKQ, DV, ncols);
1847    constexpr int nthreads  = ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols);
1848    constexpr int nwarps    = nthreads / warp_size;
1849
1850    const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix.
1851
1852    const int stride_Q1   = nb01 / sizeof(float2);
1853    const int stride_Q2   = nb02 / sizeof(float2);
1854    const int stride_K    = nb11 / sizeof(half2);
1855    const int stride_mask = nb31 / sizeof(half);
1856
1857    const int stride_V = V_is_K_view ? stride_K : nb21 / sizeof(half2);
1858
1859    const int iter_k     = (ne11      + (nbatch_fa - 1)) / nbatch_fa;
1860    const int iter_j     = (ne01.z    + (ncols1    - 1)) / ncols1;
1861    const int iter_z_gqa = (gqa_ratio + (ncols2    - 1)) / ncols2;
1862
1863    // kbc == k block continuous, current index in continuous ijk space.
1864    int       kbc      = int64_t(blockIdx.x + 0)*(iter_k*iter_j*iter_z_gqa*ne12*ne03) / gridDim.x;
1865    const int kbc_stop = int64_t(blockIdx.x + 1)*(iter_k*iter_j*iter_z_gqa*ne12*ne03) / gridDim.x;
1866
1867    // If the seams of 2 CUDA blocks fall within an output tile their results need to be combined.
1868    // For this we need to track both the block that starts the tile (needs_fixup) and the block that finishes the tile (is_fixup).
1869    // In the most general case >2 seams can fall into the same tile.
1870
1871    // kb0 == k start index when in the output tile.
1872    int kb0_start = kbc % iter_k;
1873    int kb0_stop  = min(iter_k, kb0_start + kbc_stop - kbc);
1874
1875    while (kbc < kbc_stop && kb0_stop == iter_k) {
1876        // z_KV == K/V head index, zt_gqa = Q head start index per K/V head, jt = token position start index
1877        const int sequence =  kbc /(iter_k*iter_j*iter_z_gqa*ne12);
1878        const int z_KV     = (kbc - iter_k*iter_j*iter_z_gqa*ne12 * sequence)/(iter_k*iter_j*iter_z_gqa);
1879        const int zt_gqa   = (kbc - iter_k*iter_j*iter_z_gqa*ne12 * sequence - iter_k*iter_j*iter_z_gqa * z_KV)/(iter_k*iter_j);
1880        const int jt       = (kbc - iter_k*iter_j*iter_z_gqa*ne12 * sequence - iter_k*iter_j*iter_z_gqa * z_KV - iter_k*iter_j * zt_gqa) / iter_k;
1881
1882        const int zt_Q = z_KV*gqa_ratio + zt_gqa*ncols2; // Global Q head start index.
1883
1884        const float2 * Q_f2   = (const float2 *) (Q + nb03*sequence + nb02*zt_Q);
1885        const half2  * K_h2   = (const half2  *) (K + nb13*sequence + nb12*z_KV);
1886        const half   * mask_h = ncols2 == 1 && !mask ? nullptr :
1887            (const half *) (mask + nb33*(sequence % ne33));
1888        float2       * dstk   = ((float2 *) dst) + (sequence*ne01.z*ne02 + zt_Q) * (DV/2);
1889
1890        const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
1891        const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
1892        const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr;
1893
1894        const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
1895
1896        if (KV_max) {
1897            kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa);
1898        }
1899        constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
1900        if (kb0_start == 0) {
1901            constexpr bool needs_fixup = false; // CUDA block is working on an entire tile.
1902            flash_attn_ext_f16_process_tile<DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup>
1903                (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
1904                 ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
1905        } else {
1906            constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile.
1907            flash_attn_ext_f16_process_tile<DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup>
1908                (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
1909                 ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
1910        }
1911
1912        kbc += iter_k;
1913        kbc -= kbc % iter_k;
1914
1915        kb0_start = 0;
1916        kb0_stop  = min(iter_k, kbc_stop - kbc);
1917    }
1918
1919    if (kbc >= kbc_stop) {
1920        return;
1921    }
1922
1923    // z_KV == K/V head index, zt_gqa = Q head start index per K/V head, jt = token position start index.
1924    const int sequence =  kbc /(iter_k*iter_j*iter_z_gqa*ne12);
1925    const int z_KV     = (kbc - iter_k*iter_j*iter_z_gqa*ne12 * sequence)/(iter_k*iter_j*iter_z_gqa);
1926    const int zt_gqa   = (kbc - iter_k*iter_j*iter_z_gqa*ne12 * sequence - iter_k*iter_j*iter_z_gqa * z_KV)/(iter_k*iter_j);
1927    const int jt       = (kbc - iter_k*iter_j*iter_z_gqa*ne12 * sequence - iter_k*iter_j*iter_z_gqa * z_KV - iter_k*iter_j * zt_gqa) / iter_k;
1928
1929    const int zt_Q = z_KV*gqa_ratio + zt_gqa*ncols2; // Global Q head start index.
1930
1931    const float2 * Q_f2   = (const float2 *) (Q + nb03*sequence + nb02*zt_Q);
1932    const half2  * K_h2   = (const half2  *) (K + nb13*sequence + nb12*z_KV);
1933    const half   * mask_h = ncols2 == 1 && !mask ? nullptr :
1934        (const half *) (mask + nb33*(sequence % ne33));
1935    float2       * dstk   = ((float2 *) dst) + (sequence*ne01.z*ne02 + zt_Q) * (DV/2);
1936
1937    const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
1938    const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
1939    const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr;
1940
1941    const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
1942
1943    if (KV_max) {
1944        kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa);
1945    }
1946
1947    constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
1948    constexpr bool needs_fixup = false;
1949    flash_attn_ext_f16_process_tile<DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup>
1950        (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
1951         ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
1952#else
1953    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
1954        max_bias, m0, m1, n_head_log2, logit_softcap,
1955        ne00, ne01, ne02, ne03,
1956              nb01, nb02, nb03,
1957        ne10, ne11, ne12, ne13,
1958              nb11, nb12, nb13,
1959              nb21, nb22, nb23,
1960              ne31, ne32, ne33,
1961              nb31, nb32, nb33);
1962    NO_DEVICE_CODE;
1963#endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE))
1964}
1965
1966bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
1967
1968template <int DKQ, int DV, int ncols1, int ncols2>
1969void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
1970    const ggml_tensor * KQV = dst;
1971    const int id = ggml_cuda_get_device();
1972    const int cc = ggml_cuda_info().devices[id].cc;
1973
1974    constexpr int ncols = ncols1 * ncols2;
1975
1976    const int  nthreads       = ggml_cuda_fattn_mma_get_nthreads      (DKQ, DV, ncols, cc);
1977    const int  nbatch_fa      = ggml_cuda_fattn_mma_get_nbatch_fa     (DKQ, DV, ncols, cc);
1978    const int  nbatch_K2      = ggml_cuda_fattn_mma_get_nbatch_K2     (DKQ, DV, ncols, cc);
1979    const int  nbatch_V2      = ggml_cuda_fattn_mma_get_nbatch_V2     (DKQ, DV, ncols, cc);
1980    const int  nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols, cc);
1981    const bool Q_in_reg       = ggml_cuda_fattn_mma_get_Q_in_reg      (DKQ, DV, ncols, cc);
1982    const int  nstages        = ggml_cuda_fattn_mma_get_nstages       (DKQ, DV, ncols1, ncols2, cc);
1983
1984    const int cols_per_warp = std::min(ncols, get_cols_per_warp(cc));
1985    const int warp_size_host = ggml_cuda_info().devices[ctx.device].warp_size;
1986    const int nwarps         = nthreads / warp_size_host;
1987
1988    constexpr bool V_is_K_view = DKQ == 576; // Guaranteed by the kernel selection logic in fattn.cu
1989
1990    // KV tile strides must match flash_attn_ext_f16_iter / _process_tile.
1991    const int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2, cc);
1992    const int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2, cc);
1993    const size_t nbytes_shared_KV_1stage = nbatch_fa            * std::max(stride_tile_K,  stride_tile_V) * sizeof(half2);
1994    const size_t nbytes_shared_KV_2stage = nbatch_fa            *         (stride_tile_K + stride_tile_V) * sizeof(half2);
1995    const size_t nbytes_shared_Q         = ncols                * (DKQ/2 + 4)                             * sizeof(half2);
1996    const size_t nbytes_shared_mask      = ncols1               * (nbatch_fa/2 + 4)                       * sizeof(half2);
1997    const size_t nbytes_shared_combine   = nwarps*cols_per_warp * (nbatch_combine + 4)                    * sizeof(half2);
1998
1999    const size_t nbytes_shared_KV = nstages <= 1 ? nbytes_shared_KV_1stage : nbytes_shared_KV_2stage;
2000
2001    const size_t nbytes_shared_total = std::max(nbytes_shared_combine, Q_in_reg ?
2002        std::max(nbytes_shared_Q,  nbytes_shared_KV + nbytes_shared_mask) :
2003                 nbytes_shared_Q + nbytes_shared_KV + nbytes_shared_mask);
2004
2005    float logit_softcap;
2006    memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
2007
2008#if defined(GGML_USE_HIP)
2009    using fattn_kernel_ptr_t = const void*;
2010#else
2011    using fattn_kernel_ptr_t = fattn_kernel_t;
2012#endif // defined(GGML_USE_HIP)
2013    fattn_kernel_t fattn_kernel;
2014    bool use_sparse = false;
2015    if (logit_softcap == 0.0f) {
2016        constexpr bool use_logit_softcap = false;
2017#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
2018        if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) {
2019            if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) {
2020                constexpr bool use_sparse_kernel = true;
2021                fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
2022                use_sparse = true;
2023
2024                static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
2025                if (!shared_memory_limit_raised[id]) {
2026                    CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
2027                    shared_memory_limit_raised[id] = true;
2028                }
2029            } else {
2030                constexpr bool use_sparse_kernel = false;
2031                fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
2032
2033                static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
2034                if (!shared_memory_limit_raised[id]) {
2035                    CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
2036                    shared_memory_limit_raised[id] = true;
2037                }
2038            }
2039        } else
2040#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
2041        {
2042            constexpr bool use_sparse_kernel = false;
2043            fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
2044
2045#if !defined(GGML_USE_MUSA)
2046            static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
2047            if (!shared_memory_limit_raised[id]) {
2048                CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
2049                shared_memory_limit_raised[id] = true;
2050            }
2051#endif // !defined(GGML_USE_MUSA)
2052        }
2053    } else {
2054        constexpr bool use_logit_softcap = true;
2055        constexpr bool use_sparse_kernel = false;
2056        fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
2057
2058#if !defined(GGML_USE_MUSA)
2059        static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
2060        if (!shared_memory_limit_raised[id]) {
2061            CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
2062            shared_memory_limit_raised[id] = true;
2063        }
2064#endif // !defined(GGML_USE_MUSA)
2065    }
2066
2067    launch_fattn<DV, ncols1, ncols2>
2068        (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host);
2069}
2070
2071
2072#define DECL_FATTN_MMA_F16_CASE(DKQ, DV, ncols1, ncols2)                          \
2073    template void ggml_cuda_flash_attn_ext_mma_f16_case                           \
2074    <DKQ, DV, ncols1, ncols2>(ggml_backend_cuda_context & ctx, ggml_tensor * dst) \
2075
2076#define DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(DKQ, DV, ncols)   \
2077    extern DECL_FATTN_MMA_F16_CASE(DKQ, DV, (ncols)/ 1,  1); \
2078    extern DECL_FATTN_MMA_F16_CASE(DKQ, DV, (ncols)/ 2,  2); \
2079    extern DECL_FATTN_MMA_F16_CASE(DKQ, DV, (ncols)/ 4,  4); \
2080    extern DECL_FATTN_MMA_F16_CASE(DKQ, DV, (ncols)/ 8,  8); \
2081    extern DECL_FATTN_MMA_F16_CASE(DKQ, DV, (ncols)/16, 16); \
2082
2083DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 64,  64,   8)
2084DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 80,  80,   8)
2085DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 96,  96,   8)
2086DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112,   8)
2087DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128,   8)
2088DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256,   8)
2089
2090DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 64,  64,  16)
2091DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 80,  80,  16)
2092DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 96,  96,  16)
2093DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112,  16)
2094DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128,  16)
2095DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256,  16)
2096
2097DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 64,  64,  32)
2098DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 80,  80,  32)
2099DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 96,  96,  32)
2100DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112,  32)
2101DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128,  32)
2102DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256,  32)
2103
2104DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 64,  64,  64)
2105DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 80,  80,  64)
2106DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 96,  96,  64)
2107DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112,  64)
2108DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128,  64)
2109DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256,  64)
2110
2111extern DECL_FATTN_MMA_F16_CASE(512, 512,  4,  2);
2112extern DECL_FATTN_MMA_F16_CASE(512, 512,  8,  2);
2113extern DECL_FATTN_MMA_F16_CASE(512, 512, 16,  2);
2114extern DECL_FATTN_MMA_F16_CASE(512, 512, 32,  2);
2115extern DECL_FATTN_MMA_F16_CASE(512, 512,  2,  4);
2116extern DECL_FATTN_MMA_F16_CASE(512, 512,  4,  4);
2117extern DECL_FATTN_MMA_F16_CASE(512, 512,  8,  4);
2118extern DECL_FATTN_MMA_F16_CASE(512, 512, 16,  4);
2119extern DECL_FATTN_MMA_F16_CASE(512, 512,  1,  8);
2120extern DECL_FATTN_MMA_F16_CASE(512, 512,  2,  8);
2121extern DECL_FATTN_MMA_F16_CASE(512, 512,  4,  8);
2122extern DECL_FATTN_MMA_F16_CASE(512, 512,  8,  8);
2123
2124// The number of viable configurations for Deepseek is very limited:
2125extern DECL_FATTN_MMA_F16_CASE(576, 512, 1, 16);
2126extern DECL_FATTN_MMA_F16_CASE(576, 512, 2, 16);
2127extern DECL_FATTN_MMA_F16_CASE(576, 512, 4, 16);
2128
2129// Mistral Small 4 (DKQ=320, DV=256), GQA=32-only build:
2130extern DECL_FATTN_MMA_F16_CASE(320, 256,  1, 32);
2131extern DECL_FATTN_MMA_F16_CASE(320, 256,  2, 32);
2132
2133// For GLM 4.7 Flash
2134extern DECL_FATTN_MMA_F16_CASE(576, 512,  4,  4);
2135extern DECL_FATTN_MMA_F16_CASE(576, 512,  8,  4);
2136extern DECL_FATTN_MMA_F16_CASE(576, 512, 16,  4);
2137extern DECL_FATTN_MMA_F16_CASE(576, 512,  1, 32);
2138extern DECL_FATTN_MMA_F16_CASE(576, 512,  2, 32);