ggml: support ggml_rope_set_offset on opencl, sycl, wgpu, hexagon (#27345)

* ggml: support ggml_rope_set_offset on opencl, sycl, wgpu, hexagon

* rm inplace optimization
This commit is contained in:
Xuan-Son Nguyen
2026-08-21 00:36:57 +02:00
committed by GitHub
parent 0e1d9185c5
commit 749f688fca
8 changed files with 152 additions and 112 deletions
-2
View File
@@ -6242,8 +6242,6 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
}
case GGML_OP_ROPE:
case GGML_OP_ROPE_BACK:
// FIXME: support ggml_rope_set_offset
return ((const int32_t *) op->op_params)[15] == 0;
case GGML_OP_IM2COL:
case GGML_OP_IM2COL_3D:
case GGML_OP_UPSCALE:
+58 -48
View File
@@ -41,7 +41,7 @@ template <bool forward, bool has_ff, typename T, typename D>
static void rope_norm(const T *x, D *dst, const int ne00, const int ne01,
const int ne02, const int s01, const int s02,
const int s03, const int s1, const int s2, const int s3,
const int n_dims, const int32_t *pos,
const int n_dims, const int n_offs, const int32_t *pos,
const float freq_scale, const float ext_factor,
const float attn_factor, const rope_corr_dims corr_dims,
const float theta_scale, const float *freq_factors,
@@ -78,19 +78,21 @@ static void rope_norm(const T *x, D *dst, const int ne00, const int ne01,
ggml_sycl_memcpy_1<4>(dst + idst, &v);
}
};
if (i0 >= n_dims) {
if (i0 < n_offs || i0 >= n_offs + n_dims) {
store_coaelsced(x[ix + 0], x[ix + 1]);
return;
}
const float theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f);
const int iw = i0 - n_offs; // relative idx
const float freq_factor = has_ff ? freq_factors[i0 / 2] : 1.0f;
const float theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f);
const float freq_factor = has_ff ? freq_factors[iw / 2] : 1.0f;
float cos_theta;
float sin_theta;
rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, i0,
rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, iw,
ext_factor, attn_factor, cos_theta, sin_theta);
const float x0 = x[ix + 0];
@@ -104,7 +106,7 @@ template <bool forward, bool has_ff, typename T, typename D>
static void rope_neox(const T *x, D *dst, const int ne00, const int ne01,
const int ne02, const int s01, const int s02,
const int s03, const int s1, const int s2, const int s3,
const int n_dims, const int32_t *pos,
const int n_dims, const int n_offs, const int32_t *pos,
const float freq_scale, const float ext_factor,
const float attn_factor, const rope_corr_dims corr_dims,
const float theta_scale, const float *freq_factors,
@@ -132,35 +134,38 @@ static void rope_neox(const T *x, D *dst, const int ne00, const int ne01,
idst += row_indices[i2] * set_rows_stride;
}
if (i0 >= n_dims) {
if (i0 < n_offs || i0 >= n_offs + n_dims) {
dst[idst + i0 / 2 + 0] = ggml_sycl_cast<D>(x[ix + i0 / 2 + 0]);
dst[idst + i0 / 2 + 1] = ggml_sycl_cast<D>(x[ix + i0 / 2 + 1]);
return;
}
const float theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f);
const int iw = i0 - n_offs; // relative idx
const float freq_factor = has_ff ? freq_factors[i0 / 2] : 1.0f;
const float theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f);
const float freq_factor = has_ff ? freq_factors[iw / 2] : 1.0f;
float cos_theta;
float sin_theta;
rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, i0,
rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, iw,
ext_factor, attn_factor, cos_theta, sin_theta);
const float x0 = x[ix + 0];
const float x1 = x[ix + n_dims / 2];
// idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2
const float x0 = x[ix + n_offs / 2 + 0];
const float x1 = x[ix + n_offs / 2 + n_dims / 2];
dst[idst + 0] = ggml_sycl_cast<D>(x0 * cos_theta - x1 * sin_theta);
dst[idst + n_dims / 2] = ggml_sycl_cast<D>(x0 * sin_theta + x1 * cos_theta);
dst[idst + n_offs / 2 + 0] = ggml_sycl_cast<D>(x0 * cos_theta - x1 * sin_theta);
dst[idst + n_offs / 2 + n_dims / 2] = ggml_sycl_cast<D>(x0 * sin_theta + x1 * cos_theta);
}
template <bool forward, bool has_ff, typename T>
static void rope_multi(const T *x, T *dst, const int ne00, const int ne01,
const int ne02, const int s01, const int s02,
const int s03, const int s1, const int s2, const int s3,
const int n_dims, const int32_t *pos,
const int n_dims, const int n_offs, const int32_t *pos,
const float freq_scale, const float ext_factor,
const float attn_factor, const rope_corr_dims corr_dims,
const float theta_scale, const float *freq_factors,
@@ -183,54 +188,57 @@ static void rope_multi(const T *x, T *dst, const int ne00, const int ne01,
int idst = i0 / 2 + i1 * s1 + i2 * s2 + i3 * s3;
const int ix = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03;
if (i0 >= n_dims) {
if (i0 < n_offs || i0 >= n_offs + n_dims) {
dst[idst + i0 / 2 + 0] = x[ix + i0 / 2 + 0];
dst[idst + i0 / 2 + 1] = x[ix + i0 / 2 + 1];
return;
}
const int iw = i0 - n_offs; // relative idx
const int sect_dims =
sections.v[0] + sections.v[1] + sections.v[2] + sections.v[3];
const int sec_w = sections.v[1] + sections.v[0];
const int sector = (i0 / 2) % sect_dims;
const int sector = (iw / 2) % sect_dims;
float theta_base = 0.0;
if (is_imrope) {
if (sector % 3 == 1 && sector < 3 * sections.v[1]) { // h
theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, i0 / 2.0f);
theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, iw / 2.0f);
} else if (sector % 3 == 2 && sector < 3 * sections.v[2]) { // w
theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, i0 / 2.0f);
theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, iw / 2.0f);
} else if (sector % 3 == 0 && sector < 3 * sections.v[0]) { // t
theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f);
theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f);
} else {
theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, i0 / 2.0f);
theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, iw / 2.0f);
}
} else {
if (sector < sections.v[0]) {
theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f);
theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f);
} else if (sector >= sections.v[0] && sector < sec_w) {
theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, i0 / 2.0f);
theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, iw / 2.0f);
} else if (sector >= sec_w && sector < sec_w + sections.v[2]) {
theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, i0 / 2.0f);
theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, iw / 2.0f);
} else if (sector >= sec_w + sections.v[2]) {
theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, i0 / 2.0f);
theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, iw / 2.0f);
}
}
const float freq_factor = has_ff ? freq_factors[i0 / 2] : 1.0f;
const float freq_factor = has_ff ? freq_factors[iw / 2] : 1.0f;
float cos_theta;
float sin_theta;
rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, i0,
rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, iw,
ext_factor, attn_factor, cos_theta, sin_theta);
const float x0 = x[ix + 0];
const float x1 = x[ix + n_dims / 2];
// idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2
const float x0 = x[ix + n_offs / 2 + 0];
const float x1 = x[ix + n_offs / 2 + n_dims / 2];
dst[idst + 0] = x0 * cos_theta - x1 * sin_theta;
dst[idst + n_dims / 2] = x0 * sin_theta + x1 * cos_theta;
dst[idst + n_offs / 2 + 0] = x0 * cos_theta - x1 * sin_theta;
dst[idst + n_offs / 2 + n_dims / 2] = x0 * sin_theta + x1 * cos_theta;
}
template <bool forward, bool has_ff, typename T>
@@ -293,7 +301,7 @@ static void
rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01,
const int ne02, const int s01, const int s02, const int s03,
const int s1, const int s2, const int s3, const int n_dims,
const int nr, const int32_t *pos, const float freq_scale,
const int n_offs, const int nr, const int32_t *pos, const float freq_scale,
const float freq_base, const float ext_factor,
const float attn_factor, const rope_corr_dims corr_dims,
const float *freq_factors, const int64_t *row_indices,
@@ -313,7 +321,7 @@ rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01,
GGML_UNUSED(item_ct1);
rope_norm<forward, false>(
x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
pos, freq_scale, ext_factor, attn_factor, corr_dims,
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, row_indices, set_rows_stride);
});
} else {
@@ -323,7 +331,7 @@ rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01,
GGML_UNUSED(item_ct1);
rope_norm<forward, true>(
x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
pos, freq_scale, ext_factor, attn_factor, corr_dims,
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, row_indices, set_rows_stride);
});
}
@@ -334,7 +342,7 @@ static void
rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01,
const int ne02, const int s01, const int s02, const int s03,
const int s1, const int s2, const int s3, const int n_dims,
const int nr, const int32_t *pos, const float freq_scale,
const int n_offs, const int nr, const int32_t *pos, const float freq_scale,
const float freq_base, const float ext_factor,
const float attn_factor, const rope_corr_dims corr_dims,
const float *freq_factors, const int64_t *row_indices,
@@ -354,7 +362,7 @@ rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01,
GGML_UNUSED(item_ct1);
rope_neox<forward, false>(
x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
pos, freq_scale, ext_factor, attn_factor, corr_dims,
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, row_indices, set_rows_stride);
});
} else {
@@ -364,7 +372,7 @@ rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01,
GGML_UNUSED(item_ct1);
rope_neox<forward, true>(
x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
pos, freq_scale, ext_factor, attn_factor, corr_dims,
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, row_indices, set_rows_stride);
});
}
@@ -375,7 +383,7 @@ static void
rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01,
const int ne02, const int s01, const int s02, const int s03,
const int s1, const int s2, const int s3, const int n_dims,
const int nr, const int32_t *pos, const float freq_scale,
const int n_offs, const int nr, const int32_t *pos, const float freq_scale,
const float freq_base, const float ext_factor,
const float attn_factor, const rope_corr_dims corr_dims,
const float *freq_factors, const mrope_sections sections,
@@ -395,7 +403,7 @@ rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01,
GGML_UNUSED(item_ct1);
rope_multi<forward, false, T>(
x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
pos, freq_scale, ext_factor, attn_factor, corr_dims,
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, sections, is_imrope);
});
} else {
@@ -405,7 +413,7 @@ rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01,
GGML_UNUSED(item_ct1);
rope_multi<forward, true, T>(
x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
pos, freq_scale, ext_factor, attn_factor, corr_dims,
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, sections, is_imrope);
});
}
@@ -497,6 +505,7 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst,
const int n_dims = ((int32_t *)dst->op_params)[1];
const int mode = ((int32_t *)dst->op_params)[2];
const int n_ctx_orig = ((int32_t *)dst->op_params)[4];
const int n_offs = ((int32_t *)dst->op_params)[15];
mrope_sections sections;
float freq_base;
@@ -526,6 +535,7 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst,
if (is_vision) {
GGML_ASSERT(n_dims == ne00 / 2);
GGML_ASSERT(n_offs == 0); // offset not supported for vision, as the rotated pairs span the whole row
}
const int32_t *pos = (const int32_t *)src1_d;
@@ -545,19 +555,19 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst,
if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
rope_neox_sycl<forward, float, float>(
(const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01,
s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
set_rows_stride, stream);
} else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
rope_neox_sycl<forward, float, sycl::half>(
(const float *)src0_d, (sycl::half *)dst_d, ne00, ne01, ne02,
s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale,
s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
row_indices, set_rows_stride, stream);
} else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
rope_neox_sycl<forward, sycl::half, sycl::half>(
(const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01,
ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale,
ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
row_indices, set_rows_stride, stream);
} else {
@@ -568,13 +578,13 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst,
if (src0->type == GGML_TYPE_F32) {
rope_multi_sycl<forward>((const float *)src0_d, (float *)dst_d,
ne00, ne01, ne02, s01, s02, s03, s1, s2,
s3, n_dims, nr, pos, freq_scale, freq_base,
s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
ext_factor, attn_factor, corr_dims,
freq_factors, sections, is_imrope, stream);
} else if (src0->type == GGML_TYPE_F16) {
rope_multi_sycl<forward>(
(const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01,
ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale,
ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
sections, is_imrope, stream);
} else {
@@ -602,19 +612,19 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst,
if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
rope_norm_sycl<forward, float, float>(
(const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01,
s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
set_rows_stride, stream);
} else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
rope_norm_sycl<forward, float, sycl::half>(
(const float *)src0_d, (sycl::half *)dst_d, ne00, ne01, ne02,
s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale,
s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
row_indices, set_rows_stride, stream);
} else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
rope_norm_sycl<forward, sycl::half, sycl::half>(
(const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01,
ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale,
ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
row_indices, set_rows_stride, stream);
} else {