char-fork/llama.cpp

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

Xuan-Son Nguyenggml: add ggml_rope_set_offset (+ metal support) (#27120)fe8156f78

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
8.0 KiB218 linesraw
1
2float rope_yarn_ramp(const float low, const float high, const uint i0) {
3    const float y = (i0 / 2 - low) / max(0.001f, high - low);
4    return 1.0f - min(1.0f, max(0.0f, y));
5}
6
7uint rope_a_coord(const uint i0, const uint i01, const uint i02, const uint i03, rope_params p) {
8#if RMS_NORM_ROPE_FUSION
9    // Per-row offset in shared memory
10    const uint ix = i0;
11#else
12    const uint ix = p.a_offset + i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i0;
13#endif
14    return ix;
15}
16
17void rope_yarn(const float theta_extrap, const uint i0, out float cos_theta, out float sin_theta, rope_params p) {
18    float mscale = p.attn_factor;
19    // Get n-d rotational scaling corrected for extrapolation
20    float theta_interp = p.freq_scale * theta_extrap;
21    float theta = theta_interp;
22    if (p.ext_factor != 0.0f) {
23        float ramp_mix = rope_yarn_ramp(p.corr_dims[0], p.corr_dims[1], i0) * p.ext_factor;
24        theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
25
26        // Get n-d magnitude scaling corrected for interpolation
27        mscale *= 1.0f + 0.1f * log(1.0f / p.freq_scale);
28    }
29    // Backprogagation uses inverted rotation
30    if (p.is_back != 0) {
31        theta = -theta;
32    }
33    cos_theta = cos(theta) * mscale;
34    sin_theta = sin(theta) * mscale;
35}
36
37void rope_norm(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {
38    if (i0 >= p.ne00) {
39        return;
40    }
41
42    uint idst = i0 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
43    const uint ix = rope_a_coord(i0, i1, i2, i3, p);
44
45    // Fusion optimization: ROPE + VIEW + SET_ROWS.
46    // The rope output is viewed as a 1D tensor and offset based on a row index in rope_data_i.
47    if (p.set_rows_stride != 0) {
48        idst = i1*p.nb11 + i0;
49        idst += rope_data_i[i2].x * p.set_rows_stride;
50    }
51    idst += p.d_offset;
52
53    if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) {
54        rope_data_d[idst + 0] = ROPE_D_TYPE(rope_data_a[ix + 0]);
55        rope_data_d[idst + 1] = ROPE_D_TYPE(rope_data_a[ix + 1]);
56
57        return;
58    }
59
60    const uint iw = i0 - p.n_offs; // relative idx
61
62    const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, iw/2.0f);
63
64    const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f;
65
66    float cos_theta, sin_theta;
67    rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p);
68
69    const float x0 = float(rope_data_a[ix + 0]);
70    const float x1 = float(rope_data_a[ix + 1]);
71
72    rope_data_d[idst + 0] = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
73    rope_data_d[idst + 1] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
74}
75
76void rope_neox(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {
77    if (i0 >= p.ne00) {
78        return;
79    }
80
81    uint idst = i0/2 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
82    const uint ix = rope_a_coord(i0/2, i1, i2, i3, p);
83
84    // Fusion optimization: ROPE + VIEW + SET_ROWS.
85    // The rope output is viewed as a 1D tensor and offset based on a row index in rope_data_i.
86    if (p.set_rows_stride != 0) {
87        idst = i1*p.nb11 + i0/2;
88        idst += rope_data_i[i2].x * p.set_rows_stride;
89    }
90    idst += p.d_offset;
91
92    if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) {
93        rope_data_d[idst + i0/2 + 0] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 0]);
94        rope_data_d[idst + i0/2 + 1] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 1]);
95
96        return;
97    }
98
99    const uint iw = i0 - p.n_offs; // relative idx
100
101    const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, iw/2.0f);
102
103    const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f;
104
105    float cos_theta, sin_theta;
106    rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p);
107
108    // idst/ix point at channel i0/2; the first channel of the rotated pair is p.n_offs + iw/2 = i0/2 + p.n_offs/2
109    const float x0 = float(rope_data_a[ix + p.n_offs/2 + 0]);
110    const float x1 = float(rope_data_a[ix + p.n_offs/2 + p.n_dims/2]);
111
112    rope_data_d[idst + p.n_offs/2 + 0]          = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
113    rope_data_d[idst + p.n_offs/2 + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
114}
115
116
117void rope_multi(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {
118    if (i0 >= p.ne00) {
119        return;
120    }
121
122    uint idst = i0/2 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
123    const uint ix = rope_a_coord(i0/2, i1, i2, i3, p);
124
125    // Fusion optimization: ROPE + VIEW + SET_ROWS.
126    // The rope output is viewed as a 1D tensor and offset based on a row index in rope_data_i.
127    if (p.set_rows_stride != 0) {
128        idst = i1*p.nb11 + i0/2;
129        idst += rope_data_i[i2].x * p.set_rows_stride;
130    }
131    idst += p.d_offset;
132
133    if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) {
134        rope_data_d[idst + i0/2 + 0] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 0]);
135        rope_data_d[idst + i0/2 + 1] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 1]);
136
137        return;
138    }
139
140    const uint iw = i0 - p.n_offs; // relative idx
141
142    const int sect_dims = p.sections[0] + p.sections[1] + p.sections[2] + p.sections[3];
143    const int sec_w = p.sections[1] + p.sections[0];
144    const uint sector = (iw / 2) % sect_dims;
145
146    float theta_base = 0.0;
147    if (p.is_imrope != 0) {
148        if (sector % 3 == 1 && sector < 3 * p.sections[1]) {
149            theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, iw/2.0f);
150        } else if (sector % 3 == 2 && sector < 3 * p.sections[2]) {
151            theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, iw/2.0f);
152        } else if (sector % 3 == 0 && sector < 3 * p.sections[0]) {
153            theta_base = rope_data_pos[i2]*pow(p.theta_scale, iw/2.0f);
154        } else {
155            theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, iw/2.0f);
156        }
157    } else {
158        if (sector < p.sections[0]) {
159            theta_base = rope_data_pos[i2]*pow(p.theta_scale, iw/2.0f);
160        }
161        else if (sector >= p.sections[0] && sector < sec_w) {
162            theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, iw/2.0f);
163        }
164        else if (sector >= sec_w && sector < sec_w + p.sections[2]) {
165            theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, iw/2.0f);
166        }
167        else if (sector >= sec_w + p.sections[2]) {
168            theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, iw/2.0f);
169        }
170    }
171
172    const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f;
173
174    float cos_theta, sin_theta;
175    rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p);
176
177    // idst/ix point at channel i0/2; the first channel of the rotated pair is p.n_offs + iw/2 = i0/2 + p.n_offs/2
178    const float x0 = float(rope_data_a[ix + p.n_offs/2 + 0]);
179    const float x1 = float(rope_data_a[ix + p.n_offs/2 + p.n_dims/2]);
180
181    rope_data_d[idst + p.n_offs/2 + 0]          = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
182    rope_data_d[idst + p.n_offs/2 + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
183}
184
185void rope_vision(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {
186    if (i0 >= p.ne00) {
187        return;
188    }
189
190    const uint idst = p.d_offset + i0/2 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
191    const uint ix = rope_a_coord(i0/2, i1, i2, i3, p);
192
193    const int sect_dims = p.sections[0] + p.sections[1];
194    const int sec_w = p.sections[1] + p.sections[0];
195    const uint sector = (i0 / 2) % sect_dims;
196
197    float theta_base = 0.0;
198    if (sector < p.sections[0]) {
199        const uint p0 = sector;
200        theta_base = rope_data_pos[i2]*pow(p.theta_scale, p0);
201    }
202    else if (sector >= p.sections[0] && sector < sec_w) {
203        const uint p0 = sector - p.sections[0];
204        theta_base = rope_data_pos[i2 + p.ne02]*pow(p.theta_scale, p0);
205    }
206
207    const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;
208
209    float cos_theta, sin_theta;
210    rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p);
211
212    const float x0 = float(rope_data_a[ix + 0]);
213    const float x1 = float(rope_data_a[ix + p.n_dims]);
214
215    rope_data_d[idst + 0]        = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
216    rope_data_d[idst + p.n_dims] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
217}
218