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
+8 -3
View File
@@ -7434,9 +7434,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;
@@ -23910,6 +23907,7 @@ 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];
float freq_base;
float freq_scale;
@@ -23938,6 +23936,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;
@@ -24029,6 +24028,12 @@ 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 after beta_slow, mrope has it after is_imrope
if (!is_mrope && !is_vision) {
CL_CHECK(clSetKernelArg(kernel, 33, sizeof(int), &n_offs));
} else if (is_mrope && !is_vision) {
CL_CHECK(clSetKernelArg(kernel, 35, sizeof(int), &n_offs));
}
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};
+52 -40
View File
@@ -75,7 +75,8 @@ kernel void kernel_rope_norm_f32(
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow
float beta_slow,
int n_offs
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -94,14 +95,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);
@@ -154,7 +156,8 @@ kernel void kernel_rope_norm_f16(
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow
float beta_slow,
int n_offs
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -173,14 +176,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);
@@ -233,7 +237,8 @@ kernel void kernel_rope_neox_f32(
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow
float beta_slow,
int n_offs
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -252,17 +257,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];
@@ -312,7 +318,8 @@ kernel void kernel_rope_neox_f16(
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow
float beta_slow,
int n_offs
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -331,17 +338,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];
@@ -393,7 +401,8 @@ kernel void kernel_rope_multi_f32(
float beta_fast,
float beta_slow,
int4 sections,
int is_imrope
int is_imrope,
int n_offs
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -414,10 +423,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 +455,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];
@@ -504,7 +514,8 @@ kernel void kernel_rope_multi_f16(
float beta_fast,
float beta_slow,
int4 sections,
int is_imrope
int is_imrope,
int n_offs
) {
src0 = (global void*)((global char*)src0 + offset0);
src1 = (global int*)((global char*)src1 + offset1);
@@ -525,10 +536,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 +568,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];