char-fork/llama.cpp

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

Xuan-Son Nguyenggml: support ggml_rope_set_offset on opencl, sycl, wgpu, hexagon (#27345)749f688fc

char/aliceai-t5-moe
24.5 KiB759 linesraw
1#pragma OPENCL EXTENSION cl_khr_fp16 : enable
2
3//------------------------------------------------------------------------------
4// kernel_rope
5//------------------------------------------------------------------------------
6float rope_yarn_ramp(float low, float high, int i0) {
7    const float y = (i0 / 2 - low) / max(0.001f, high - low);
8    return 1.0f - min(1.0f, max(0.0f, y));
9}
10
11// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn
12// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
13float2 rope_yarn(
14    float theta_extrap, float freq_scale, float2 corr_dims, int i0, float ext_factor, float mscale
15) {
16    // Get n-d rotational scaling corrected for extrapolation
17    float theta_interp = freq_scale * theta_extrap;
18    float theta = theta_interp;
19    if (ext_factor != 0.0f) {
20        float ramp_mix = rope_yarn_ramp(corr_dims.s0, corr_dims.s1, i0) * ext_factor;
21        theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
22
23        // Get n-d magnitude scaling corrected for interpolation
24        mscale *= 1.0f + 0.1f * log(1.0f / freq_scale);
25    }
26    return (float2)(cos(theta) * mscale, sin(theta) * mscale);
27}
28
29// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
30// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
31float rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) {
32    return n_dims * log(n_ctx_orig / (n_rot * 2 * M_PI_F)) / (2 * log(base));
33}
34
35float2 rope_yarn_corr_dims(
36    int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow
37) {
38    // start and end correction dims
39    return (float2)(
40        max(0.0f,         floor(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base))),
41        min(n_dims - 1.0f, ceil(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base)))
42    );
43}
44
45kernel void kernel_rope_norm_f32(
46        global void * src0,
47        ulong offset0,
48        global int * src1,
49        ulong offset1,
50        global float * src2,
51        ulong offset2,
52        global float * dst,
53        ulong offsetd,
54        int ne00,
55        int ne01,
56        int ne02,
57        int ne03,
58        ulong nb00,
59        ulong nb01,
60        ulong nb02,
61        ulong nb03,
62        int ne0,
63        int ne1,
64        int ne2,
65        int ne3,
66        ulong nb0,
67        ulong nb1,
68        ulong nb2,
69        ulong nb3,
70        int n_past,
71        int n_dims,
72        int n_ctx_orig,
73        float freq_base,
74        float freq_scale,
75        float ext_factor,
76        float attn_factor,
77        float beta_fast,
78        float beta_slow,
79        int n_offs
80) {
81    src0 = (global void*)((global char*)src0 + offset0);
82    src1 = (global int*)((global char*)src1 + offset1);
83    src2 = (global float*)((global char*)src2 + offset2);
84    dst = (global float*)((global char*)dst + offsetd);
85
86    int i3 = get_group_id(2);
87    int i2 = get_group_id(1);
88    int i1 = get_group_id(0);
89
90    float2 corr_dims = rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow);
91
92    global int * pos = src1;
93
94    float theta_base = (float) pos[i2];
95    float inv_ndims = -1.f/n_dims;
96
97    for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
98        if (i0 >= n_offs && i0 < n_offs + n_dims) {
99            int iw = i0 - n_offs; // relative idx
100            int ic = iw/2;
101
102            float theta = theta_base * pow(freq_base, inv_ndims*iw);
103
104            float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
105
106            float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
107
108            global float * src       = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
109            global float * dst_data  = (global float *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
110
111            float x0 = src[0];
112            float x1 = src[1];
113
114            dst_data[0] = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
115            dst_data[1] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
116        } else {
117            global float * src      = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
118            global float * dst_data = (global float *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
119
120            dst_data[0] = src[0];
121            dst_data[1] = src[1];
122        }
123    }
124}
125
126kernel void kernel_rope_norm_f16(
127        global void * src0,
128        ulong offset0,
129        global int * src1,
130        ulong offset1,
131        global float * src2,
132        ulong offset2,
133        global float * dst,
134        ulong offsetd,
135        int ne00,
136        int ne01,
137        int ne02,
138        int ne03,
139        ulong nb00,
140        ulong nb01,
141        ulong nb02,
142        ulong nb03,
143        int ne0,
144        int ne1,
145        int ne2,
146        int ne3,
147        ulong nb0,
148        ulong nb1,
149        ulong nb2,
150        ulong nb3,
151        int n_past,
152        int n_dims,
153        int n_ctx_orig,
154        float freq_base,
155        float freq_scale,
156        float ext_factor,
157        float attn_factor,
158        float beta_fast,
159        float beta_slow,
160        int n_offs
161) {
162    src0 = (global void*)((global char*)src0 + offset0);
163    src1 = (global int*)((global char*)src1 + offset1);
164    src2 = (global float*)((global char*)src2 + offset2);
165    dst = (global float*)((global char*)dst + offsetd);
166
167    int i3 = get_group_id(2);
168    int i2 = get_group_id(1);
169    int i1 = get_group_id(0);
170
171    float2 corr_dims = rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow);
172
173    global int * pos = src1;
174
175    float theta_base = (float) pos[i2];
176    float inv_ndims = -1.f/n_dims;
177
178    for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
179        if (i0 >= n_offs && i0 < n_offs + n_dims) {
180            int iw = i0 - n_offs; // relative idx
181            int ic = iw/2;
182
183            float theta = theta_base * pow(freq_base, inv_ndims*iw);
184
185            float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
186
187            float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
188
189            global half * src       = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
190            global half * dst_data  = (global half *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
191
192            float x0 = src[0];
193            float x1 = src[1];
194
195            dst_data[0] = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
196            dst_data[1] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
197        } else {
198            global half * src      = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
199            global half * dst_data = (global half *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
200
201            dst_data[0] = src[0];
202            dst_data[1] = src[1];
203        }
204    }
205}
206
207kernel void kernel_rope_neox_f32(
208        global void * src0,
209        ulong offset0,
210        global int * src1,
211        ulong offset1,
212        global float * src2,
213        ulong offset2,
214        global float * dst,
215        ulong offsetd,
216        int ne00,
217        int ne01,
218        int ne02,
219        int ne03,
220        ulong nb00,
221        ulong nb01,
222        ulong nb02,
223        ulong nb03,
224        int ne0,
225        int ne1,
226        int ne2,
227        int ne3,
228        ulong nb0,
229        ulong nb1,
230        ulong nb2,
231        ulong nb3,
232        int n_past,
233        int n_dims,
234        int n_ctx_orig,
235        float freq_base,
236        float freq_scale,
237        float ext_factor,
238        float attn_factor,
239        float beta_fast,
240        float beta_slow,
241        int n_offs
242) {
243    src0 = (global void*)((global char*)src0 + offset0);
244    src1 = (global int*)((global char*)src1 + offset1);
245    src2 = (global float*)((global char*)src2 + offset2);
246    dst = (global float*)((global char*)dst + offsetd);
247
248    int i3 = get_group_id(2);
249    int i2 = get_group_id(1);
250    int i1 = get_group_id(0);
251
252    float2 corr_dims = rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow);
253
254    global int * pos = src1;
255
256    float theta_base = (float) pos[i2];
257    float inv_ndims = -1.f/n_dims;
258
259    for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
260        if (i0 >= n_offs && i0 < n_offs + n_dims) {
261            int iw = i0 - n_offs; // relative idx
262            int ic = iw/2;
263
264            const float theta = theta_base * pow(freq_base, inv_ndims*iw);
265
266            const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
267
268            float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
269
270            global float * src      = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00);
271            global float * dst_data = (global float *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + (n_offs + ic)*nb0);
272
273            const float x0 = src[0];
274            const float x1 = src[n_dims/2];
275
276            dst_data[0]        = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
277            dst_data[n_dims/2] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
278        } else {
279            global float * const src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
280            global float * dst_data  = (global float *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
281
282            dst_data[0] = src[0];
283            dst_data[1] = src[1];
284        }
285    }
286}
287
288kernel void kernel_rope_neox_f16(
289        global void * src0,
290        ulong offset0,
291        global int * src1,
292        ulong offset1,
293        global float * src2,
294        ulong offset2,
295        global float * dst,
296        ulong offsetd,
297        int ne00,
298        int ne01,
299        int ne02,
300        int ne03,
301        ulong nb00,
302        ulong nb01,
303        ulong nb02,
304        ulong nb03,
305        int ne0,
306        int ne1,
307        int ne2,
308        int ne3,
309        ulong nb0,
310        ulong nb1,
311        ulong nb2,
312        ulong nb3,
313        int n_past,
314        int n_dims,
315        int n_ctx_orig,
316        float freq_base,
317        float freq_scale,
318        float ext_factor,
319        float attn_factor,
320        float beta_fast,
321        float beta_slow,
322        int n_offs
323) {
324    src0 = (global void*)((global char*)src0 + offset0);
325    src1 = (global int*)((global char*)src1 + offset1);
326    src2 = (global float*)((global char*)src2 + offset2);
327    dst = (global float*)((global char*)dst + offsetd);
328
329    int i3 = get_group_id(2);
330    int i2 = get_group_id(1);
331    int i1 = get_group_id(0);
332
333    float2 corr_dims = rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow);
334
335    global int * pos = src1;
336
337    float theta_base = (float) pos[i2];
338    float inv_ndims = -1.f/n_dims;
339
340    for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
341        if (i0 >= n_offs && i0 < n_offs + n_dims) {
342            int iw = i0 - n_offs; // relative idx
343            int ic = iw/2;
344
345            const float theta = theta_base * pow(freq_base, inv_ndims*iw);
346
347            const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
348
349            float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
350
351            global half * src       = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00);
352            global half * dst_data  = (global half *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + (n_offs + ic)*nb0);
353
354            const float x0 = src[0];
355            const float x1 = src[n_dims/2];
356
357            dst_data[0]        = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
358            dst_data[n_dims/2] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
359        } else {
360            global half * const src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
361            global half * dst_data  = (global half *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
362
363            dst_data[0] = src[0];
364            dst_data[1] = src[1];
365        }
366    }
367}
368
369kernel void kernel_rope_multi_f32(
370        global void * src0,
371        ulong offset0,
372        global int * src1,
373        ulong offset1,
374        global float * src2,
375        ulong offset2,
376        global float * dst,
377        ulong offsetd,
378        int ne00,
379        int ne01,
380        int ne02,
381        int ne03,
382        ulong nb00,
383        ulong nb01,
384        ulong nb02,
385        ulong nb03,
386        int ne0,
387        int ne1,
388        int ne2,
389        int ne3,
390        ulong nb0,
391        ulong nb1,
392        ulong nb2,
393        ulong nb3,
394        int n_past,
395        int n_dims,
396        int n_ctx_orig,
397        float freq_base,
398        float freq_scale,
399        float ext_factor,
400        float attn_factor,
401        float beta_fast,
402        float beta_slow,
403        int4 sections,
404        int  is_imrope,
405        int n_offs
406) {
407    src0 = (global void*)((global char*)src0 + offset0);
408    src1 = (global int*)((global char*)src1 + offset1);
409    src2 = (global float*)((global char*)src2 + offset2);
410    dst = (global float*)((global char*)dst + offsetd);
411
412    int i3 = get_group_id(2);
413    int i2 = get_group_id(1);
414    int i1 = get_group_id(0);
415
416    float2 corr_dims = rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow);
417
418    global int * pos = src1;
419
420    const int sect_dims = sections.s0 + sections.s1 + sections.s2 + sections.s3;
421    const int sec_w = sections.s1 + sections.s0;
422
423    float inv_ndims = -1.f/n_dims;
424
425    for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
426        if (i0 >= n_offs && i0 < n_offs + n_dims) {
427            int iw = i0 - n_offs; // relative idx
428            int ic = iw/2;
429
430            const int sector = ic % sect_dims;
431            float theta_base = 0.0f;
432
433            if (is_imrope) {
434                if (sector % 3 == 1 && sector < 3 * sections.s1) { // h
435                    theta_base = (float) pos[i2 + ne02 * 1];
436                } else if (sector % 3 == 2 && sector < 3 * sections.s2) { // w
437                    theta_base = (float) pos[i2 + ne02 * 2];
438                } else if (sector % 3 == 0 && sector < 3 * sections.s0) { // t
439                    theta_base = (float) pos[i2 + ne02 * 0];
440                } else { // e
441                    theta_base = (float) pos[i2 + ne02 * 3];
442                }
443            } else {
444                if (sector < sections.s0) {
445                    theta_base = pos[i2];
446                }
447                else if (sector >= sections.s0 && sector < sec_w) {
448                    theta_base = pos[i2 + ne2 * 1];
449                }
450                else if (sector >= sec_w && sector < sec_w + sections.s2) {
451                    theta_base = pos[i2 + ne2 * 2];
452                }
453                else if (sector >= sec_w + sections.s2) {
454                    theta_base = pos[i2 + ne2 * 3];
455                }
456            }
457
458            const float theta = theta_base * pow(freq_base, inv_ndims*iw);
459
460            const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
461
462            float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
463
464            global float * src      = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00);
465            global float * dst_data = (global float *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + (n_offs + ic)*nb0);
466
467            const float x0 = src[0];
468            const float x1 = src[n_dims/2];
469
470            dst_data[0]        = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
471            dst_data[n_dims/2] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
472        } else {
473            global float * const src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
474            global float * dst_data  = (global float *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
475
476            dst_data[0] = src[0];
477            dst_data[1] = src[1];
478        }
479    }
480}
481
482kernel void kernel_rope_multi_f16(
483        global void * src0,
484        ulong offset0,
485        global int * src1,
486        ulong offset1,
487        global float * src2,
488        ulong offset2,
489        global half * dst,
490        ulong offsetd,
491        int ne00,
492        int ne01,
493        int ne02,
494        int ne03,
495        ulong nb00,
496        ulong nb01,
497        ulong nb02,
498        ulong nb03,
499        int ne0,
500        int ne1,
501        int ne2,
502        int ne3,
503        ulong nb0,
504        ulong nb1,
505        ulong nb2,
506        ulong nb3,
507        int n_past,
508        int n_dims,
509        int n_ctx_orig,
510        float freq_base,
511        float freq_scale,
512        float ext_factor,
513        float attn_factor,
514        float beta_fast,
515        float beta_slow,
516        int4 sections,
517        int  is_imrope,
518        int n_offs
519) {
520    src0 = (global void*)((global char*)src0 + offset0);
521    src1 = (global int*)((global char*)src1 + offset1);
522    src2 = (global float*)((global char*)src2 + offset2);
523    dst = (global float*)((global char*)dst + offsetd);
524
525    int i3 = get_group_id(2);
526    int i2 = get_group_id(1);
527    int i1 = get_group_id(0);
528
529    float2 corr_dims = rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow);
530
531    global int * pos = src1;
532
533    const int sect_dims = sections.s0 + sections.s1 + sections.s2 + sections.s3;
534    const int sec_w = sections.s1 + sections.s0;
535
536    float inv_ndims = -1.f/n_dims;
537
538    for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
539        if (i0 >= n_offs && i0 < n_offs + n_dims) {
540            int iw = i0 - n_offs; // relative idx
541            int ic = iw/2;
542
543            const int sector = ic % sect_dims;
544            float theta_base = 0.0f;
545
546            if (is_imrope) {
547                if (sector % 3 == 1 && sector < 3 * sections.s1) { // h
548                    theta_base = (float) pos[i2 + ne02 * 1];
549                } else if (sector % 3 == 2 && sector < 3 * sections.s2) { // w
550                    theta_base = (float) pos[i2 + ne02 * 2];
551                } else if (sector % 3 == 0 && sector < 3 * sections.s0) { // t
552                    theta_base = (float) pos[i2 + ne02 * 0];
553                } else { // e
554                    theta_base = (float) pos[i2 + ne02 * 3];
555                }
556            } else {
557                if (sector < sections.s0) {
558                    theta_base = pos[i2];
559                }
560                else if (sector >= sections.s0 && sector < sec_w) {
561                    theta_base = pos[i2 + ne2 * 1];
562                }
563                else if (sector >= sec_w && sector < sec_w + sections.s2) {
564                    theta_base = pos[i2 + ne2 * 2];
565                }
566                else if (sector >= sec_w + sections.s2) {
567                    theta_base = pos[i2 + ne2 * 3];
568                }
569            }
570
571            const float theta = theta_base * pow(freq_base, inv_ndims*iw);
572
573            const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
574
575            float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
576
577            global half * src      = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00);
578            global half * dst_data = (global half *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + (n_offs + ic)*nb0);
579
580            const float x0 = src[0];
581            const float x1 = src[n_dims/2];
582
583            dst_data[0]        = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
584            dst_data[n_dims/2] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
585        } else {
586            global half * const src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
587            global half * dst_data  = (global half *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
588
589            dst_data[0] = src[0];
590            dst_data[1] = src[1];
591        }
592    }
593}
594
595kernel void kernel_rope_vision_f32(
596        global void * src0,
597        ulong offset0,
598        global int * src1,
599        ulong offset1,
600        global float * src2,
601        ulong offset2,
602        global float * dst,
603        ulong offsetd,
604        int ne00,
605        int ne01,
606        int ne02,
607        int ne03,
608        ulong nb00,
609        ulong nb01,
610        ulong nb02,
611        ulong nb03,
612        int ne0,
613        int ne1,
614        int ne2,
615        int ne3,
616        ulong nb0,
617        ulong nb1,
618        ulong nb2,
619        ulong nb3,
620        int n_past,
621        int n_dims,
622        int n_ctx_orig,
623        float freq_base,
624        float freq_scale,
625        float ext_factor,
626        float attn_factor,
627        float beta_fast,
628        float beta_slow,
629        int4 sections
630) {
631    src0 = (global void*)((global char*)src0 + offset0);
632    src1 = (global int*)((global char*)src1 + offset1);
633    src2 = (global float*)((global char*)src2 + offset2);
634    dst = (global float*)((global char*)dst + offsetd);
635
636    int i3 = get_group_id(2);
637    int i2 = get_group_id(1);
638    int i1 = get_group_id(0);
639
640    float2 corr_dims = rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow);
641
642    global int * pos = src1;
643
644    const int sect_dims = sections.s0 + sections.s1;
645    const int sec_w = sections.s1 + sections.s0;
646
647    float inv_ndims = -1.f/n_dims;
648
649    for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
650        int ic = i0/2;
651
652        const int sector = (i0/2) % sect_dims;
653        float theta_base = 0.0f;
654
655        if (sector < sections.s0) {
656            const int p = sector;
657            theta_base = pos[i2] * pow(freq_base, inv_ndims*2.0f*p);
658        } else if (sector >= sections.s0 && sector < sec_w) {
659            const int p = sector - sections.s0;
660            theta_base = pos[i2 + ne2] * pow(freq_base, inv_ndims*2.0f*p);
661        }
662
663        const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
664
665        float2 cos_sin_theta = rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor);
666
667        global float * src      = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00);
668        global float * dst_data = (global float *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + ic*nb0);
669
670        const float x0 = src[0];
671        const float x1 = src[n_dims];
672
673        dst_data[0]      = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
674        dst_data[n_dims] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
675    }
676}
677
678kernel void kernel_rope_vision_f16(
679        global void * src0,
680        ulong offset0,
681        global int * src1,
682        ulong offset1,
683        global float * src2,
684        ulong offset2,
685        global half * dst,
686        ulong offsetd,
687        int ne00,
688        int ne01,
689        int ne02,
690        int ne03,
691        ulong nb00,
692        ulong nb01,
693        ulong nb02,
694        ulong nb03,
695        int ne0,
696        int ne1,
697        int ne2,
698        int ne3,
699        ulong nb0,
700        ulong nb1,
701        ulong nb2,
702        ulong nb3,
703        int n_past,
704        int n_dims,
705        int n_ctx_orig,
706        float freq_base,
707        float freq_scale,
708        float ext_factor,
709        float attn_factor,
710        float beta_fast,
711        float beta_slow,
712        int4 sections
713) {
714    src0 = (global void*)((global char*)src0 + offset0);
715    src1 = (global int*)((global char*)src1 + offset1);
716    src2 = (global float*)((global char*)src2 + offset2);
717    dst = (global float*)((global char*)dst + offsetd);
718
719    int i3 = get_group_id(2);
720    int i2 = get_group_id(1);
721    int i1 = get_group_id(0);
722
723    float2 corr_dims = rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow);
724
725    global int * pos = src1;
726
727    const int sect_dims = sections.s0 + sections.s1;
728    const int sec_w = sections.s1 + sections.s0;
729
730    float inv_ndims = -1.f/n_dims;
731
732    for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
733        int ic = i0/2;
734
735        const int sector = (i0/2) % sect_dims;
736        float theta_base = 0.0f;
737
738        if (sector < sections.s0) {
739            const int p = sector;
740            theta_base = pos[i2] * pow(freq_base, inv_ndims*2.0f*p);
741        } else if (sector >= sections.s0 && sector < sec_w) {
742            const int p = sector - sections.s0;
743            theta_base = pos[i2 + ne2] * pow(freq_base, inv_ndims*2.0f*p);
744        }
745
746        const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
747
748        float2 cos_sin_theta = rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor);
749
750        global half * src      = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00);
751        global half * dst_data = (global half *)((global char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + ic*nb0);
752
753        const float x0 = src[0];
754        const float x1 = src[n_dims];
755
756        dst_data[0]      = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
757        dst_data[n_dims] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
758    }
759}