1#pragma OPENCL EXTENSION cl_khr_fp16 : enable
2
3#define ACC_TYPE float
4#define ACC_TYPE4 float4
5#define DATA_TYPE float
6#define DATA_TYPE4 float4
7#define MASK_DATA_TYPE half
8#define CONVERT_ACC4(x) (x)
9#define CONVERT_DATA4(x) (x)
10
11#define DK_VEC (DK/4)
12#define DV_VEC (DV/4)
13#define WG_SIZE (BLOCK_M)
14#define Q1_WG_SIZE 64
15
16inline float get_alibi_slope(
17 const float max_bias, const uint h, const uint n_head_log2, const float m0, const float m1
18) {
19 if (max_bias <= 0.0f) {
20 return 1.0f;
21 }
22 const float base = h < n_head_log2 ? m0 : m1;
23 const int exph = h < n_head_log2 ? h + 1 : 2*(h - n_head_log2) + 1;
24
25 return pow(base, exph);
26}
27__kernel void flash_attn_f32(
28 const global void * q_void, ulong q_offset,
29 const global void * k_void, ulong k_offset,
30 const global void * v_void, ulong v_offset,
31 global void * o_void, ulong o_offset,
32 const float scale,
33 const int n_q,
34 const int n_kv,
35 const int is_causal,
36 const int n_head,
37 const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
38 const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
39 const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
40 const ulong o_nb1, const ulong o_nb2, const ulong o_nb3,
41 const float max_bias,
42 const float m0,
43 const float m1,
44 const int n_head_log2,
45 const float logit_softcap,
46 const int n_head_kv,
47 const global void* mask_void,
48 const ulong mask_offset,
49 const ulong mask_nb1,
50 const ulong mask_nb2,
51 const ulong mask_nb3,
52 const int mask_ne2,
53 const int mask_ne3,
54 const global void* sinks_void,
55 const ulong sinks_offset
56) {
57 const int tid = get_local_id(0);
58 const int block_q_idx = get_group_id(0);
59 const int head_batch_idx = get_global_id(1);
60
61 const int my_query_row = block_q_idx * BLOCK_M + tid;
62
63 const int batch_idx = head_batch_idx / n_head;
64 const int head_idx = head_batch_idx % n_head;
65
66 const int gqa_ratio = n_head / n_head_kv;
67 const int head_kv_idx = head_idx / gqa_ratio;
68
69 const global char* q_base = (const global char*)q_void + q_offset;
70 const global char* k_base = (const global char*)k_void + k_offset;
71 const global char* v_base = (const global char*)v_void + v_offset;
72 global char* o_base = (global char*)o_void + o_offset;
73
74 const global char* mask_base = NULL;
75 if (mask_void != NULL) {
76 const int mask_head_idx = head_idx % mask_ne2;
77 const int mask_batch_idx = batch_idx % mask_ne3;
78 mask_base = (const global char*)mask_void + mask_offset + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
79 }
80
81 ACC_TYPE4 q_priv[DK_VEC];
82 if (my_query_row < n_q) {
83 const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + my_query_row * q_nb1;
84 const global DATA_TYPE4* q_ptr = (const global DATA_TYPE4*)(q_base + q_row_offset);
85 #pragma unroll
86 for (int i = 0; i < DK_VEC; ++i) {
87 q_priv[i] = CONVERT_ACC4(q_ptr[i]);
88 }
89 }
90
91 ACC_TYPE4 o_acc[DV_VEC];
92 #pragma unroll
93 for (int i = 0; i < DV_VEC; ++i) {
94 o_acc[i] = (ACC_TYPE4)(0.0f);
95 }
96 ACC_TYPE m_i = -INFINITY;
97 ACC_TYPE l_i = 0.0f;
98
99 float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
100
101 __local DATA_TYPE4 l_k[BLOCK_N][DK_VEC];
102 __local DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
103
104 for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) {
105 for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) {
106 const int row = i / DK_VEC;
107 const int col = i % DK_VEC;
108 const int k_row_idx = k_start + row;
109 if (k_row_idx < n_kv) {
110 const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_row_idx * k_nb1;
111 l_k[row][col] = ((__global DATA_TYPE4*)(k_base + k_row_offset))[col];
112 }
113 }
114 for (int i = tid; i < BLOCK_N * DV_VEC; i += WG_SIZE) {
115 const int row = i / DV_VEC;
116 const int col = i % DV_VEC;
117 const int v_row_idx = k_start + row;
118 if (v_row_idx < n_kv) {
119 const ulong v_row_offset = batch_idx * v_nb3 + head_kv_idx * v_nb2 + v_row_idx * v_nb1;
120 l_v[row][col] = ((__global DATA_TYPE4*)(v_base + v_row_offset))[col];
121 }
122 }
123 barrier(CLK_LOCAL_MEM_FENCE);
124
125 if (my_query_row >= n_q) {
126 continue;
127 }
128
129 for (int j = 0; j < BLOCK_N; j += 2) {
130 const int k_row0 = k_start + j;
131 const int k_row1 = k_start + j + 1;
132
133 ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f);
134 ACC_TYPE4 dot_acc1 = (ACC_TYPE4)(0.0f);
135 #pragma unroll
136 for (int k = 0; k < DK_VEC; k++) {
137 dot_acc0 = mad(q_priv[k], CONVERT_ACC4(l_k[j][k]), dot_acc0);
138 dot_acc1 = mad(q_priv[k], CONVERT_ACC4(l_k[j+1][k]), dot_acc1);
139 }
140 ACC_TYPE score0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale;
141 ACC_TYPE score1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;
142
143 if (is_causal) {
144 if (k_row0 > (n_kv - n_q + my_query_row)) score0 = -INFINITY;
145 if (k_row1 > (n_kv - n_q + my_query_row)) score1 = -INFINITY;
146 }
147
148 if (k_row0 >= n_kv) score0 = -INFINITY;
149 if (k_row1 >= n_kv) score1 = -INFINITY;
150
151 if (mask_base != NULL) {
152 const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base + my_query_row * mask_nb1);
153 if (k_row0 < n_kv) score0 += slope * (ACC_TYPE)mask_ptr[k_row0];
154 if (k_row1 < n_kv) score1 += slope * (ACC_TYPE)mask_ptr[k_row1];
155 }
156
157 if (logit_softcap > 0.0f) {
158 score0 = logit_softcap * tanh(score0 / logit_softcap);
159 score1 = logit_softcap * tanh(score1 / logit_softcap);
160 }
161
162 const ACC_TYPE m_new = max(m_i, max(score0, score1));
163 const ACC_TYPE p0 = exp(score0 - m_new);
164 const ACC_TYPE p1 = exp(score1 - m_new);
165 const ACC_TYPE scale_prev = exp(m_i - m_new);
166
167 #pragma unroll
168 for (int i = 0; i < DV_VEC; ++i) {
169 o_acc[i] = o_acc[i] * scale_prev + p0 * CONVERT_ACC4(l_v[j][i]) + p1 * CONVERT_ACC4(l_v[j+1][i]);
170 }
171 l_i = l_i * scale_prev + p0 + p1;
172 m_i = m_new;
173 }
174 }
175
176 if (my_query_row < n_q) {
177 if (sinks_void != NULL) {
178 const global ACC_TYPE* sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset);
179 const ACC_TYPE m_sink = sinks_ptr[head_idx];
180 const ACC_TYPE m_final = max(m_i, m_sink);
181
182 const ACC_TYPE scale_o = exp(m_i - m_final);
183 #pragma unroll
184 for (int i = 0; i < DV_VEC; ++i) {
185 o_acc[i] *= scale_o;
186 }
187
188 l_i = l_i * exp(m_i - m_final) + exp(m_sink - m_final);
189 }
190
191 const ulong o_row_offset = batch_idx * o_nb3 + my_query_row * o_nb2 + head_idx * o_nb1;
192 global DATA_TYPE4 *o_row = (global DATA_TYPE4 *)(o_base + o_row_offset);
193 if (l_i > 0.0f) {
194 const ACC_TYPE l_inv = 1.0f / l_i;
195 #pragma unroll
196 for (int i = 0; i < DV_VEC; ++i) {
197 o_row[i] = CONVERT_DATA4(o_acc[i] * l_inv);
198 }
199 } else {
200 #pragma unroll
201 for (int i = 0; i < DV_VEC; ++i) {
202 o_row[i] = (DATA_TYPE4)(0.0f);
203 }
204 }
205 }
206}
207
208__kernel void flash_attn_f32_q1(
209 const global void * q_void, ulong q_offset,
210 const global void * k_void, ulong k_offset,
211 const global void * v_void, ulong v_offset,
212 global void * o_void, ulong o_offset,
213 const float scale,
214 const int n_q,
215 const int n_kv,
216 const int is_causal,
217 const int n_head,
218 const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
219 const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
220 const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
221 const ulong o_nb1, const ulong o_nb2, const ulong o_nb3,
222 const float max_bias,
223 const float m0,
224 const float m1,
225 const int n_head_log2,
226 const float logit_softcap,
227 const int n_head_kv,
228 const global void* mask_void,
229 const ulong mask_offset,
230 const ulong mask_nb1,
231 const ulong mask_nb2,
232 const ulong mask_nb3,
233 const int mask_ne2,
234 const int mask_ne3,
235 const global void* sinks_void,
236 const ulong sinks_offset
237) {
238 const int tid = get_local_id(0);
239 const int head_batch_idx = get_global_id(1);
240
241 const int batch_idx = head_batch_idx / n_head;
242 const int head_idx = head_batch_idx % n_head;
243
244 const int gqa_ratio = n_head / n_head_kv;
245 const int head_kv_idx = head_idx / gqa_ratio;
246
247 const global char* q_base = (const global char*)q_void + q_offset;
248 const global char* k_base = (const global char*)k_void + k_offset;
249 const global char* v_base = (const global char*)v_void + v_offset;
250 global char* o_base = (global char*)o_void + o_offset;
251
252 const global char* mask_base = NULL;
253 if (mask_void != NULL) {
254 const int mask_head_idx = head_idx % mask_ne2;
255 const int mask_batch_idx = batch_idx % mask_ne3;
256 mask_base = (const global char*)mask_void + mask_offset + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
257 }
258
259 ACC_TYPE4 q_priv[DK_VEC];
260 const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2;
261 const global DATA_TYPE4* q_ptr = (const global DATA_TYPE4*)(q_base + q_row_offset);
262 #pragma unroll
263 for (int i = 0; i < DK_VEC; ++i) {
264 q_priv[i] = CONVERT_ACC4(q_ptr[i]);
265 }
266
267 float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
268
269 const global ACC_TYPE* sinks_ptr = NULL;
270 if (sinks_void != NULL) {
271 sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset);
272 }
273
274 ACC_TYPE m_i = (sinks_ptr != NULL) ? sinks_ptr[head_idx] : -INFINITY;
275 for (int k_idx = tid; k_idx < n_kv; k_idx += Q1_WG_SIZE) {
276 const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
277 const global DATA_TYPE4* k_ptr = (const global DATA_TYPE4*)(k_base + k_row_offset);
278 ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f);
279 #pragma unroll
280 for (int k = 0; k < DK_VEC; k++) {
281 dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
282 }
283 ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
284 if (mask_base != NULL) {
285 const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
286 score += slope * (ACC_TYPE)mask_ptr[k_idx];
287 }
288 if (logit_softcap > 0.0f) {
289 score = logit_softcap * tanh(score / logit_softcap);
290 }
291 m_i = max(m_i, score);
292 }
293
294 __local ACC_TYPE local_m[Q1_WG_SIZE];
295 local_m[tid] = m_i;
296 barrier(CLK_LOCAL_MEM_FENCE);
297 #pragma unroll
298 for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
299 if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]);
300 barrier(CLK_LOCAL_MEM_FENCE);
301 }
302 const ACC_TYPE m_final = local_m[0];
303
304 ACC_TYPE4 o_acc[DV_VEC];
305 #pragma unroll
306 for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
307 ACC_TYPE l_i = 0.0f;
308
309 for (int k_idx = tid; k_idx < n_kv; k_idx += Q1_WG_SIZE) {
310 const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
311 const ulong v_row_offset = batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
312 const global DATA_TYPE4* k_ptr = (const global DATA_TYPE4*)(k_base + k_row_offset);
313 const global DATA_TYPE4* v_ptr = (const global DATA_TYPE4*)(v_base + v_row_offset);
314 ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f);
315 #pragma unroll
316 for (int k = 0; k < DK_VEC; k++) {
317 dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
318 }
319 ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
320 if (mask_base != NULL) {
321 const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
322 score += slope * (ACC_TYPE)mask_ptr[k_idx];
323 }
324 if (logit_softcap > 0.0f) {
325 score = logit_softcap * tanh(score / logit_softcap);
326 }
327 const ACC_TYPE p = exp(score - m_final);
328 l_i += p;
329 #pragma unroll
330 for (int i = 0; i < DV_VEC; i++) {
331 o_acc[i] = mad(p, CONVERT_ACC4(v_ptr[i]), o_acc[i]);
332 }
333 }
334
335 __local ACC_TYPE local_l[Q1_WG_SIZE];
336 __local ACC_TYPE4 local_o_comp[Q1_WG_SIZE];
337 local_l[tid] = l_i;
338 barrier(CLK_LOCAL_MEM_FENCE);
339 #pragma unroll
340 for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
341 if (tid < s) local_l[tid] += local_l[tid + s];
342 barrier(CLK_LOCAL_MEM_FENCE);
343 }
344
345 const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1;
346 global DATA_TYPE4 *o_row = (global DATA_TYPE4 *)(o_base + o_row_offset);
347 ACC_TYPE l_final = local_l[0];
348
349 if (sinks_ptr != NULL) {
350 l_final += exp(sinks_ptr[head_idx] - m_final);
351 }
352
353 if (l_final > 0.0f) {
354 const ACC_TYPE l_inv = 1.0f / l_final;
355 for (int i = 0; i < DV_VEC; i++) {
356 local_o_comp[tid] = o_acc[i];
357 barrier(CLK_LOCAL_MEM_FENCE);
358 #pragma unroll
359 for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
360 if (tid < s) local_o_comp[tid] += local_o_comp[tid + s];
361 barrier(CLK_LOCAL_MEM_FENCE);
362 }
363 if (tid == 0) {
364 o_row[i] = CONVERT_DATA4(local_o_comp[0] * l_inv);
365 }
366 }
367 } else if (tid == 0) {
368 #pragma unroll
369 for (int i = 0; i < DV_VEC; ++i) o_row[i] = (DATA_TYPE4)(0.0f);
370 }
371}