Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-20 01:31:31 +02:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Xuan Son Nguyen
2026-08-19 00:29:53 +02:00
parent fe8156f789
commit efee18d262
8 changed files with 224 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-hexagon/ggml-hexagon.cpp
+3 -2
View File
@@ -3180,8 +3180,9 @@ static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * s
static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
const int32_t * op_params = &op->op_params[0];
if (op_params[15] != 0) {
return false; // FIXME: support ggml_rope_set_offset
// ggml_rope_set_offset: HVX kernels need a VLEN-aligned window start (32 f32 elems)
if (op_params[15] % 32 != 0) {
return false;
}
int mode = op_params[2];
ggml/src/ggml-hexagon/htp/rope-ops.c
+16 -6
View File
@@ -53,6 +53,7 @@
struct htp_rope_context {
int32_t n_dims;
int32_t n_offs;
int32_t mode;
int32_t n_ctx_orig;
int32_t sections[4];
@@ -405,32 +406,40 @@ static inline void hvx_rope_f32_aa(float * restrict dst, const float * restrict
static void inline rope_basic_f32(struct htp_rope_context * rctx, uint8_t * restrict dst, uint8_t * restrict src,
uint32_t nr, uint32_t ne0, const float * restrict theta_cache) {
const uint32_t n_offs = rctx->n_offs; // VLEN-aligned (enforced by supports_op)
#pragma unroll(4)
for (uint32_t i = 0; i < nr; i++) {
float * d = (float *) (dst + i * rctx->dst_row_size_aligned);
float * s = (float *) (src + i * rctx->src0_row_size_aligned);
hvx_rope_f32_aa(d, s, rctx->n_dims, theta_cache);
hvx_rope_f32_aa(d + n_offs, s + n_offs, rctx->n_dims, theta_cache);
// fill the remain channels with data from src tensor
if (rctx->n_dims < ne0) {
hvx_copy_f32_uu((uint8_t *)(d + rctx->n_dims), (uint8_t *)(s + rctx->n_dims), ne0 - rctx->n_dims);
if (n_offs > 0) {
hvx_copy_f32_uu((uint8_t *) d, (uint8_t *) s, n_offs);
}
if (n_offs + rctx->n_dims < ne0) {
hvx_copy_f32_uu((uint8_t *)(d + n_offs + rctx->n_dims), (uint8_t *)(s + n_offs + rctx->n_dims), ne0 - n_offs - rctx->n_dims);
}
}
}
static void inline rope_neox_f32(struct htp_rope_context * rctx, uint8_t * restrict dst, uint8_t * restrict src,
uint32_t nr, uint32_t ne0, const float * restrict theta_cache) {
const uint32_t n_offs = rctx->n_offs; // VLEN-aligned (enforced by supports_op)
#pragma unroll(4)
for (uint32_t i = 0; i < nr; i++) {
float * d = (float *) (dst + i * rctx->dst_row_size_aligned);
float * s = (float *) (src + i * rctx->src0_row_size_aligned);
hvx_rope_neox_f32_aa(d, s, rctx->n_dims, theta_cache);
hvx_rope_neox_f32_aa(d + n_offs, s + n_offs, rctx->n_dims, theta_cache);
// fill the remain channels with data from src tensor
if (rctx->n_dims < ne0) {
hvx_copy_f32_uu((uint8_t *)(d + rctx->n_dims), (uint8_t *)(s + rctx->n_dims), ne0 - rctx->n_dims);
if (n_offs > 0) {
hvx_copy_f32_uu((uint8_t *) d, (uint8_t *) s, n_offs);
}
if (n_offs + rctx->n_dims < ne0) {
hvx_copy_f32_uu((uint8_t *)(d + n_offs + rctx->n_dims), (uint8_t *)(s + n_offs + rctx->n_dims), ne0 - n_offs - rctx->n_dims);
}
}
}
@@ -673,6 +682,7 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
rctx.n_dims = ((const int32_t *) op_params)[1];
rctx.mode = ((const int32_t *) op_params)[2];
rctx.n_ctx_orig = ((const int32_t *) op_params)[4];
rctx.n_offs = ((const int32_t *) op_params)[15];
memcpy(&rctx.freq_base, (int32_t *) op_params + 5, sizeof(float));
memcpy(&rctx.freq_scale, (int32_t *) op_params + 6, sizeof(float));
ggml/src/ggml-opencl/ggml-opencl.cpp
+13 -3
View File
@@ -7376,9 +7376,6 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
case GGML_OP_DIAG_MASK_INF:
return op->ne[3] == 1;
case GGML_OP_ROPE: {
if (((const int32_t *) op->op_params)[15] != 0) {
return false; // FIXME: support ggml_rope_set_offset
}
const int mode = ((const int32_t *) op->op_params)[2];
const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE;
const bool is_vision = mode == GGML_ROPE_TYPE_VISION;
@@ -23706,6 +23703,10 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const
const int n_dims = ((int *) dst->op_params)[1];
const int mode = ((int *) 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];
// when dst aliases src0, the channels outside the rotated window already hold the correct data
const int inplace = extrad->data_device == extra0->data_device && offsetd == offset0;
float freq_base;
float freq_scale;
@@ -23734,6 +23735,7 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const
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
}
cl_kernel kernel;
@@ -23825,6 +23827,14 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const
if (is_mrope && !is_vision) {
CL_CHECK(clSetKernelArg(kernel, 34, sizeof(int), &is_imrope));
}
// norm and neox have n_offs and inplace after beta_slow, mrope has them after is_imrope
if (!is_mrope && !is_vision) {
CL_CHECK(clSetKernelArg(kernel, 33, sizeof(int), &n_offs));
CL_CHECK(clSetKernelArg(kernel, 34, sizeof(int), &inplace));
} else if (is_mrope && !is_vision) {
CL_CHECK(clSetKernelArg(kernel, 35, sizeof(int), &n_offs));
CL_CHECK(clSetKernelArg(kernel, 36, sizeof(int), &inplace));
}
size_t global_work_size[] = {(size_t)ne01*nth, (size_t)ne02, (size_t)ne03};
size_t local_work_size[] = {(size_t)nth, 1, 1};
ggml/src/ggml-opencl/kernels/rope.cl
+76 -40
View File
@@ -75,7 +75,9 @@ kernel void kernel_rope_norm_f32(
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow
float beta_slow,
int n_offs,
int inplace
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -94,14 +96,15 @@ kernel void kernel_rope_norm_f32(
float inv_ndims = -1.f/n_dims;
for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
if (i0 < n_dims) {
int ic = i0/2;
if (i0 >= n_offs && i0 < n_offs + n_dims) {
int iw = i0 - n_offs; // relative idx
int ic = iw/2;
float theta = theta_base * pow(freq_base, inv_ndims*i0);
float theta = theta_base * pow(freq_base, inv_ndims*iw);
float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor);
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
@@ -112,6 +115,9 @@ kernel void kernel_rope_norm_f32(
dst_data[0] = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
dst_data[1] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
} else {
if (inplace) {
continue;
}
global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
@@ -154,7 +160,9 @@ kernel void kernel_rope_norm_f16(
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow
float beta_slow,
int n_offs,
int inplace
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -173,14 +181,15 @@ kernel void kernel_rope_norm_f16(
float inv_ndims = -1.f/n_dims;
for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
if (i0 < n_dims) {
int ic = i0/2;
if (i0 >= n_offs && i0 < n_offs + n_dims) {
int iw = i0 - n_offs; // relative idx
int ic = iw/2;
float theta = theta_base * pow(freq_base, inv_ndims*i0);
float theta = theta_base * pow(freq_base, inv_ndims*iw);
float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor);
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
@@ -191,6 +200,9 @@ kernel void kernel_rope_norm_f16(
dst_data[0] = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
dst_data[1] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
} else {
if (inplace) {
continue;
}
global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
@@ -233,7 +245,9 @@ kernel void kernel_rope_neox_f32(
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow
float beta_slow,
int n_offs,
int inplace
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -252,17 +266,18 @@ kernel void kernel_rope_neox_f32(
float inv_ndims = -1.f/n_dims;
for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
if (i0 < n_dims) {
int ic = i0/2;
if (i0 >= n_offs && i0 < n_offs + n_dims) {
int iw = i0 - n_offs; // relative idx
int ic = iw/2;
const float theta = theta_base * pow(freq_base, inv_ndims*i0);
const float theta = theta_base * pow(freq_base, inv_ndims*iw);
const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor);
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00);
global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0);
global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00);
global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0);
const float x0 = src[0];
const float x1 = src[n_dims/2];
@@ -270,6 +285,9 @@ kernel void kernel_rope_neox_f32(
dst_data[0] = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
dst_data[n_dims/2] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
} else {
if (inplace) {
continue;
}
global float * const src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
@@ -312,7 +330,9 @@ kernel void kernel_rope_neox_f16(
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow
float beta_slow,
int n_offs,
int inplace
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -331,17 +351,18 @@ kernel void kernel_rope_neox_f16(
float inv_ndims = -1.f/n_dims;
for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
if (i0 < n_dims) {
int ic = i0/2;
if (i0 >= n_offs && i0 < n_offs + n_dims) {
int iw = i0 - n_offs; // relative idx
int ic = iw/2;
const float theta = theta_base * pow(freq_base, inv_ndims*i0);
const float theta = theta_base * pow(freq_base, inv_ndims*iw);
const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor);
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00);
global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0);
global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00);
global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0);
const float x0 = src[0];
const float x1 = src[n_dims/2];
@@ -349,6 +370,9 @@ kernel void kernel_rope_neox_f16(
dst_data[0] = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
dst_data[n_dims/2] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
} else {
if (inplace) {
continue;
}
global half * const src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
@@ -393,7 +417,9 @@ kernel void kernel_rope_multi_f32(
float beta_fast,
float beta_slow,
int4 sections,
int is_imrope
int is_imrope,
int n_offs,
int inplace
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -414,10 +440,11 @@ kernel void kernel_rope_multi_f32(
float inv_ndims = -1.f/n_dims;
for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
if (i0 < n_dims) {
int ic = i0/2;
if (i0 >= n_offs && i0 < n_offs + n_dims) {
int iw = i0 - n_offs; // relative idx
int ic = iw/2;
const int sector = (i0 / 2) % sect_dims;
const int sector = ic % sect_dims;
float theta_base = 0.0f;
if (is_imrope) {
@@ -445,14 +472,14 @@ kernel void kernel_rope_multi_f32(
}
}
const float theta = theta_base * pow(freq_base, inv_ndims*i0);
const float theta = theta_base * pow(freq_base, inv_ndims*iw);
const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor);
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00);
global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0);
global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00);
global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0);
const float x0 = src[0];
const float x1 = src[n_dims/2];
@@ -460,6 +487,9 @@ kernel void kernel_rope_multi_f32(
dst_data[0] = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
dst_data[n_dims/2] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
} else {
if (inplace) {
continue;
}
global float * const src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
@@ -504,7 +534,9 @@ kernel void kernel_rope_multi_f16(
float beta_fast,
float beta_slow,
int4 sections,
int is_imrope
int is_imrope,
int n_offs,
int inplace
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -525,10 +557,11 @@ kernel void kernel_rope_multi_f16(
float inv_ndims = -1.f/n_dims;
for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) {
if (i0 < n_dims) {
int ic = i0/2;
if (i0 >= n_offs && i0 < n_offs + n_dims) {
int iw = i0 - n_offs; // relative idx
int ic = iw/2;
const int sector = (i0 / 2) % sect_dims;
const int sector = ic % sect_dims;
float theta_base = 0.0f;
if (is_imrope) {
@@ -556,14 +589,14 @@ kernel void kernel_rope_multi_f16(
}
}
const float theta = theta_base * pow(freq_base, inv_ndims*i0);
const float theta = theta_base * pow(freq_base, inv_ndims*iw);
const float freq_factor = src2 != src0 ? src2[ic] : 1.0f;
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor);
float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor);
global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00);
global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0);
global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00);
global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0);
const float x0 = src[0];
const float x1 = src[n_dims/2];
@@ -571,6 +604,9 @@ kernel void kernel_rope_multi_f16(
dst_data[0] = x0*cos_sin_theta.s0 - x1*cos_sin_theta.s1;
dst_data[n_dims/2] = x0*cos_sin_theta.s1 + x1*cos_sin_theta.s0;
} else {
if (inplace) {
continue;
}
global half * const src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
ggml/src/ggml-sycl/ggml-sycl.cpp
-2
View File
@@ -6235,8 +6235,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:
ggml/src/ggml-sycl/rope.cpp
+97 -68
View File
@@ -41,11 +41,12 @@ 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,
const int64_t *row_indices, const int set_rows_stride) {
const int64_t *row_indices, const int set_rows_stride,
const bool inplace) {
auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
const int i0 = 2 * (item_ct1.get_local_range(1) * item_ct1.get_group(1) +
item_ct1.get_local_id(1));
@@ -78,19 +79,24 @@ 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) {
if (inplace) {
return;
}
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,11 +110,12 @@ 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,
const int64_t *row_indices, const int set_rows_stride) {
const int64_t *row_indices, const int set_rows_stride,
const bool inplace) {
auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
const int i0 = 2 * (item_ct1.get_local_range(1) * item_ct1.get_group(1) +
item_ct1.get_local_id(1));
@@ -132,39 +139,46 @@ 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) {
if (inplace) {
return;
}
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,
const mrope_sections sections, const bool is_imrope) {
const mrope_sections sections, const bool is_imrope,
const bool inplace) {
auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
const int i0 = 2 * (item_ct1.get_local_range(1) * item_ct1.get_group(1) +
item_ct1.get_local_id(1));
@@ -183,54 +197,60 @@ 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) {
if (inplace) {
return;
}
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,11 +313,11 @@ 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,
const int set_rows_stride, dpct::queue_ptr stream) {
const int set_rows_stride, const bool inplace, dpct::queue_ptr stream) {
GGML_ASSERT(ne00 % 2 == 0);
const dpct::dim3 block_dims(1, SYCL_ROPE_BLOCK_SIZE, 1);
const int n_blocks_x =
@@ -313,8 +333,9 @@ 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,
theta_scale, freq_factors, row_indices, set_rows_stride);
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, row_indices, set_rows_stride,
inplace);
});
} else {
stream->parallel_for(
@@ -323,8 +344,9 @@ 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,
theta_scale, freq_factors, row_indices, set_rows_stride);
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, row_indices, set_rows_stride,
inplace);
});
}
}
@@ -334,11 +356,11 @@ 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,
const int set_rows_stride, dpct::queue_ptr stream) {
const int set_rows_stride, const bool inplace, dpct::queue_ptr stream) {
GGML_ASSERT(ne00 % 2 == 0);
const dpct::dim3 block_dims(1, SYCL_ROPE_BLOCK_SIZE, 1);
const int n_blocks_x =
@@ -354,8 +376,9 @@ 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,
theta_scale, freq_factors, row_indices, set_rows_stride);
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, row_indices, set_rows_stride,
inplace);
});
} else {
stream->parallel_for(
@@ -364,8 +387,9 @@ 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,
theta_scale, freq_factors, row_indices, set_rows_stride);
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, row_indices, set_rows_stride,
inplace);
});
}
}
@@ -375,11 +399,11 @@ 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,
const bool is_imrope, dpct::queue_ptr stream) {
const bool is_imrope, const bool inplace, dpct::queue_ptr stream) {
GGML_ASSERT(ne00 % 2 == 0);
const dpct::dim3 block_dims(1, SYCL_ROPE_BLOCK_SIZE, 1);
const int n_blocks_x =
@@ -395,8 +419,8 @@ 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,
theta_scale, freq_factors, sections, is_imrope);
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, sections, is_imrope, inplace);
});
} else {
stream->parallel_for(
@@ -405,8 +429,8 @@ 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,
theta_scale, freq_factors, sections, is_imrope);
n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
theta_scale, freq_factors, sections, is_imrope, inplace);
});
}
}
@@ -497,8 +521,12 @@ 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;
// when dst aliases src0, the channels outside the rotated window already hold the correct data
const bool inplace = dst_d == src0->data;
float freq_base;
float freq_scale;
float ext_factor;
@@ -526,6 +554,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,21 +574,21 @@ 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);
set_rows_stride, inplace, 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);
row_indices, set_rows_stride, inplace, 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);
row_indices, set_rows_stride, inplace, stream);
} else {
GGML_ABORT("Fatal error: Tensor type unsupported!");
}
@@ -568,15 +597,15 @@ 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);
freq_factors, sections, is_imrope, inplace, 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);
sections, is_imrope, inplace, stream);
} else {
GGML_ABORT("Fatal error: Tensor type unsupported!");
}
@@ -602,21 +631,21 @@ 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);
set_rows_stride, inplace, 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);
row_indices, set_rows_stride, inplace, 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);
row_indices, set_rows_stride, inplace, stream);
} else {
GGML_ABORT("Fatal error: Tensor type unsupported!");
}
ggml/src/ggml-webgpu/ggml-webgpu.cpp
+4 -4
View File
@@ -2714,6 +2714,7 @@ static webgpu_encoded_op ggml_webgpu_rope(webgpu_context & ctx,
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];
float freq_base;
float freq_scale;
@@ -2762,7 +2763,8 @@ static webgpu_encoded_op ggml_webgpu_rope(webgpu_context & ctx,
(uint32_t) sections[0],
(uint32_t) sections[1],
(uint32_t) sections[2],
(uint32_t) sections[3]
(uint32_t) sections[3],
(uint32_t) n_offs
};
std::vector<wgpu::BindGroupEntry> entries = { ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0),
@@ -4472,9 +4474,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
supports_op = (op->type == GGML_TYPE_F32 && src0->type == GGML_TYPE_F32) && ggml_is_contiguous_rows(src0);
break;
case GGML_OP_ROPE:
// FIXME: support ggml_rope_set_offset
supports_op =
(op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && ((const int32_t *) op->op_params)[15] == 0;
supports_op = op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16;
break;
case GGML_OP_GLU:
switch (ggml_get_glu_op(op)) {
ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl
+15 -7
View File
@@ -38,7 +38,8 @@ struct Params {
sections0: u32,
sections1: u32,
sections2: u32,
sections3: u32
sections3: u32,
n_offs: u32
};
@group(0) @binding(0)
@@ -126,7 +127,8 @@ fn rope_yarn(theta_extrap: f32, i: u32) -> vec2<f32> {
fn pair_base(i0: u32, div_2: bool) -> u32 {
if (div_2) {
return i0 / 2;
// first channel of the rotated pair: n_offs + (i0 - n_offs)/2
return i0 / 2 + params.n_offs / 2;
} else {
return i0;
}
@@ -165,20 +167,26 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let i_src_row = params.offset_src0 + i3 * params.stride_src03 + i2 * params.stride_src02 + i1 * params.stride_src01;
let i_dst_row = params.offset_dst + i3 * params.stride_dst3 + i2 * params.stride_dst2 + i1 * params.stride_dst1;
if (i0 >= params.n_dims && !is_vision) {
if ((i0 < params.n_offs || i0 >= params.n_offs + params.n_dims) && !is_vision) {
#ifdef INPLACE
return;
#else
let i_src = i_src_row + i0;
let i_dst = i_dst_row + i0;
rotate(i_dst, i_dst + 1, f32(src0[i_src]), f32(src0[i_src + 1]));
return;
#endif
}
let iw = i0 - params.n_offs; // relative idx
var theta_base_mult: u32 = 0;
var theta_scale_pwr: u32 = i0 / 2;
var theta_scale_pwr: u32 = iw / 2;
if (is_mrope) {
let sect_dims = params.sections0 + params.sections1 + params.sections2 + params.sections3;
let sec_w = params.sections1 + params.sections0;
let sec_e = params.sections2 + sec_w;
let sector = (i0 / 2) % sect_dims;
let sector = (iw / 2) % sect_dims;
if (is_imrope) {
if (sector % 3 == 1 && sector < 3 * params.sections1) {
theta_base_mult = 1;
@@ -203,7 +211,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
} else if (sector >= sec_e) {
if (is_vision) {
theta_scale_pwr = sector - sec_e;
theta_scale_pwr = (i0 / 2) % sec_e;
theta_scale_pwr = (iw / 2) % sec_e;
}
theta_base_mult = 3;
} else if (is_vision) {
@@ -212,7 +220,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
}
}
let theta_base = f32(src1[params.offset_src1 + i2 + params.ne2 * theta_base_mult]) * pow(params.theta_scale, f32(theta_scale_pwr));
let thetas = rope_yarn(theta_base/freq_factor(i0), i0);
let thetas = rope_yarn(theta_base/freq_factor(iw), iw);
let i_src = i_src_row + pair_base(i0, is_neox || is_mrope || is_vision);
let i_dst = i_dst_row + pair_base(i0, is_neox || is_mrope || is_vision);