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
+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)) {
+11 -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,22 @@ 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) {
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;
}
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 +207,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 +216,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);