Skip to content

Commit 0eca4f5

Browse files
qnixsynapseqnixsynapse
committed
SYCL: Add mrope kernel
1 parent f5cd27b commit 0eca4f5

File tree

2 files changed

+116
-8
lines changed

2 files changed

+116
-8
lines changed

‎ggml/src/ggml-sycl/ggml-sycl.cpp‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4204,14 +4204,6 @@ static bool ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, const g
42044204
case GGML_OP_SOFT_MAX:
42054205
return true;
42064206
case GGML_OP_ROPE:
4207-
{
4208-
const int mode = ((const int32_t *) op->op_params)[2];
4209-
// mode is not used as a bitmask in practice, the various rope type modes are independent implementations
4210-
if (mode == GGML_ROPE_TYPE_MROPE) {
4211-
return false;
4212-
}
4213-
return true;
4214-
}
42154207
case GGML_OP_IM2COL:
42164208
return true;
42174209
case GGML_OP_UPSCALE:

‎ggml/src/ggml-sycl/rope.cpp‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,66 @@ static void rope_neox(const T * x, T * dst, const int ne0, const int ne1, const
122122
dst[i + n_dims / 2] = x0 * sin_theta + x1 * cos_theta;
123123
}
124124

125+
template <typename T, bool has_ff>
126+
static void rope_multi(const T * x, T * dst, const int ne0, const int ne1, const int ne2, const size_t s1,
127+
const size_t s2, const int n_dims, const int32_t * pos, const float freq_scale,
128+
const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims,
129+
const float theta_scale, const float * freq_factors, const mrope_sections sections,
130+
const sycl::nd_item<3> & item_ct1) {
131+
// get index pos
132+
const int i0 = 2 * (item_ct1.get_group(1) * item_ct1.get_local_range(1) + item_ct1.get_local_id(1));
133+
if (i0 >= ne0) {
134+
return;
135+
}
136+
const int row_dst = (item_ct1.get_group(2) * item_ct1.get_local_range(2)) + item_ct1.get_local_id(2);
137+
138+
if (i0 >= n_dims) {
139+
const int i = row_dst*ne0 + i0;
140+
141+
dst[i + 0] = x[i + 0];
142+
dst[i + 1] = x[i + 1];
143+
144+
return;
145+
}
146+
147+
const int row_x = row_dst % ne1;
148+
const int channel_x = row_dst / ne1;
149+
const int idst = (row_dst * ne0) + (i0 / 2);
150+
const size_t ix = ((size_t) channel_x * s2) + ((size_t) row_x * s1) + (i0 / 2);
151+
152+
const int sect_dims = sections.v[0] + sections.v[1] + sections.v[2] + sections.v[3];
153+
const int sec_w = sections.v[1] + sections.v[0];
154+
const int sector = (i0 / 2) % sect_dims;
155+
156+
157+
float theta_base = 0.0;
158+
if (sector < sections.v[0]) {
159+
theta_base = pos[channel_x]*sycl::pow(theta_scale, i0/2.0f);
160+
}
161+
else if (sector >= sections.v[0] && sector < sec_w) {
162+
theta_base = pos[channel_x + ne2 * 1]*sycl::pow(theta_scale, i0/2.0f);
163+
}
164+
else if (sector >= sec_w && sector < sec_w + sections.v[2]) {
165+
theta_base = pos[channel_x + ne2 * 2]*sycl::pow(theta_scale, i0/2.0f);
166+
}
167+
else if (sector >= sec_w + sections.v[2]) {
168+
theta_base = pos[channel_x + ne2 * 3]*sycl::pow(theta_scale, i0/2.0f);
169+
}
170+
171+
const float freq_factor = has_ff ? freq_factors[i0 / 2] : 1.0f;
172+
float cos_theta;
173+
float sin_theta;
174+
rope_yarn(theta_base / freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, &cos_theta, &sin_theta);
175+
const float x0 = x[ix + 0];
176+
const float x1 = x[ix + n_dims/2];
177+
178+
// store results in dst
179+
dst[idst + 0] = x0 * cos_theta - x1 * sin_theta;
180+
dst[idst + n_dims/2] = x0 * sin_theta + x1 * cos_theta;
181+
}
182+
183+
184+
125185
template <typename T, bool has_ff>
126186
static void rope_vision(const T * x, T * dst, const int ne0, const int ne1, const int ne2, const size_t s1,
127187
const size_t s2, const int n_dims, const int32_t * pos, const float freq_scale,
@@ -228,6 +288,40 @@ static void rope_neox_sycl(const T * x, T * dst, const int ne0, const int ne1, c
228288
}
229289
}
230290

291+
template <typename T>
292+
static void rope_multi_sycl(const T * x, T * dst, const int ne0, const int ne1, const int ne2, const size_t s1,
293+
const size_t s2, const int n_dims, const int nr, const int32_t * pos,
294+
const float freq_scale, const float freq_base, const float ext_factor,
295+
const float attn_factor, const rope_corr_dims corr_dims, const float * freq_factors,
296+
const mrope_sections sections, queue_ptr stream) {
297+
GGML_ASSERT(ne0 % 2 == 0);
298+
const sycl::range<3> block_dims(1, SYCL_ROPE_BLOCK_SIZE, 1);
299+
const int n_blocks_y = (ne0 + 2 * SYCL_ROPE_BLOCK_SIZE - 1) / (2 * SYCL_ROPE_BLOCK_SIZE);
300+
const sycl::range<3> grid_dims(1, n_blocks_y, nr);
301+
const sycl::nd_range<3> nd_range(grid_dims * block_dims, block_dims);
302+
303+
const float theta_scale = std::pow(freq_base, -2.0f / n_dims);
304+
// Add FP16 capability check if T could be sycl::half
305+
if constexpr (std::is_same_v<T, sycl::half>) {
306+
dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 });
307+
}
308+
// launch kernel
309+
if (freq_factors == nullptr) {
310+
stream->parallel_for(nd_range, [=](sycl::nd_item<3> item_ct1) {
311+
rope_multi<T, false>(x, dst, ne0, ne1, ne2, s1, s2, n_dims, pos, freq_scale, ext_factor, attn_factor,
312+
corr_dims, theta_scale, freq_factors, sections, item_ct1);
313+
});
314+
} else {
315+
stream->parallel_for(nd_range, [=](sycl::nd_item<3> item_ct1) {
316+
rope_multi<T, true>(x, dst, ne0, ne1, ne2, s1, s2, n_dims, pos, freq_scale, ext_factor, attn_factor,
317+
corr_dims, theta_scale, freq_factors, sections, item_ct1);
318+
});
319+
}
320+
}
321+
322+
323+
324+
231325
// rope vision
232326
template <typename T>
233327
static void rope_vision_sycl(const T * x, T * dst, const int ne0, const int ne1, const int ne2, const size_t s1,
@@ -298,8 +392,17 @@ inline void ggml_sycl_op_rope(ggml_backend_sycl_context & ctx, ggml_tensor *dst)
298392
memcpy(&sections.v, (int32_t *) dst->op_params + 11, sizeof(int)*4);
299393

300394
const bool is_neox = mode & GGML_ROPE_TYPE_NEOX;
395+
const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE;
301396
const bool is_vision = mode == GGML_ROPE_TYPE_VISION;
302397

398+
if (is_mrope) {
399+
GGML_ASSERT(sections.v[0] > 0 || sections.v[1] > 0 || sections.v[2] > 0);
400+
}
401+
402+
if (is_vision) {
403+
GGML_ASSERT(n_dims == ne00/2);
404+
}
405+
303406
const int32_t * pos = (const int32_t *) dst->src[1]->data;
304407

305408
const float * freq_factors = nullptr;
@@ -326,6 +429,19 @@ inline void ggml_sycl_op_rope(ggml_backend_sycl_context & ctx, ggml_tensor *dst)
326429
} else {
327430
GGML_ABORT("fatal error");
328431
}
432+
} else if (is_mrope && !is_vision) {
433+
GGML_SYCL_DEBUG("%s: mrope path\n", __func__);
434+
if (dst->src[0]->type == GGML_TYPE_F16) {
435+
rope_multi_sycl((const sycl::half *)dst->src[0]->data, (sycl::half *)dst->data, ne00, ne01, ne02, s01,
436+
s02, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
437+
freq_factors, sections, main_stream);
438+
} else if (dst->src[0]->type == GGML_TYPE_F32) {
439+
rope_multi_sycl((const float *) dst->src[0]->data, (float *) dst->data, ne00, ne01, ne02, s01, s02, n_dims,
440+
nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, sections,
441+
main_stream);
442+
} else {
443+
GGML_ABORT("Fatal error: Tensor type unsupported!");
444+
}
329445
} else if (is_vision) {
330446
GGML_SYCL_DEBUG("%s: vision path\n", __func__);
331447
if (dst->src[0]->type == GGML_TYPE_F16) {

0 commit comments

Comments
 (0)