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