ggml: add ggml_rope_set_offset (+ metal support) (#27120)

* add params

* cpu kernel

* metal kernel

* add test backend ops

* gate other backends

* ggml: (cuda) support ggml_rope_set_offset (#27121)

* rm cuda supports_op guard, fix webgpu clang-format

* ggml: support ggml_rope_set_offset on vulkan (#27344)

* ggml: support ggml_rope_set_offset on vulkan

* remove inplace optimization
This commit is contained in:
Xuan-Son Nguyen
2026-08-19 14:04:57 +02:00
committed by GitHub
parent 77acca437f
commit fe8156f789
19 changed files with 272 additions and 116 deletions
@@ -50,19 +50,21 @@ void rope_norm(const uint i0, const uint i1, const uint i2, const uint i3, rope_
}
idst += p.d_offset;
if (i0 >= p.n_dims) {
if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) {
rope_data_d[idst + 0] = ROPE_D_TYPE(rope_data_a[ix + 0]);
rope_data_d[idst + 1] = ROPE_D_TYPE(rope_data_a[ix + 1]);
return;
}
const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, i0/2.0f);
const uint iw = i0 - p.n_offs; // relative idx
const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;
const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, iw/2.0f);
const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f;
float cos_theta, sin_theta;
rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p);
rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p);
const float x0 = float(rope_data_a[ix + 0]);
const float x1 = float(rope_data_a[ix + 1]);
@@ -87,25 +89,28 @@ void rope_neox(const uint i0, const uint i1, const uint i2, const uint i3, rope_
}
idst += p.d_offset;
if (i0 >= p.n_dims) {
if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) {
rope_data_d[idst + i0/2 + 0] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 0]);
rope_data_d[idst + i0/2 + 1] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 1]);
return;
}
const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, i0/2.0f);
const uint iw = i0 - p.n_offs; // relative idx
const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;
const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, iw/2.0f);
const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f;
float cos_theta, sin_theta;
rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p);
rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p);
const float x0 = float(rope_data_a[ix + 0]);
const float x1 = float(rope_data_a[ix + p.n_dims/2]);
// idst/ix point at channel i0/2; the first channel of the rotated pair is p.n_offs + iw/2 = i0/2 + p.n_offs/2
const float x0 = float(rope_data_a[ix + p.n_offs/2 + 0]);
const float x1 = float(rope_data_a[ix + p.n_offs/2 + p.n_dims/2]);
rope_data_d[idst + 0] = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
rope_data_d[idst + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
rope_data_d[idst + p.n_offs/2 + 0] = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
rope_data_d[idst + p.n_offs/2 + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
}
@@ -125,53 +130,56 @@ void rope_multi(const uint i0, const uint i1, const uint i2, const uint i3, rope
}
idst += p.d_offset;
if (i0 >= p.n_dims) {
if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) {
rope_data_d[idst + i0/2 + 0] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 0]);
rope_data_d[idst + i0/2 + 1] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 1]);
return;
}
const uint iw = i0 - p.n_offs; // relative idx
const int sect_dims = p.sections[0] + p.sections[1] + p.sections[2] + p.sections[3];
const int sec_w = p.sections[1] + p.sections[0];
const uint sector = (i0 / 2) % sect_dims;
const uint sector = (iw / 2) % sect_dims;
float theta_base = 0.0;
if (p.is_imrope != 0) {
if (sector % 3 == 1 && sector < 3 * p.sections[1]) {
theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, i0/2.0f);
theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, iw/2.0f);
} else if (sector % 3 == 2 && sector < 3 * p.sections[2]) {
theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, i0/2.0f);
theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, iw/2.0f);
} else if (sector % 3 == 0 && sector < 3 * p.sections[0]) {
theta_base = rope_data_pos[i2]*pow(p.theta_scale, i0/2.0f);
theta_base = rope_data_pos[i2]*pow(p.theta_scale, iw/2.0f);
} else {
theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, i0/2.0f);
theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, iw/2.0f);
}
} else {
if (sector < p.sections[0]) {
theta_base = rope_data_pos[i2]*pow(p.theta_scale, i0/2.0f);
theta_base = rope_data_pos[i2]*pow(p.theta_scale, iw/2.0f);
}
else if (sector >= p.sections[0] && sector < sec_w) {
theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, i0/2.0f);
theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, iw/2.0f);
}
else if (sector >= sec_w && sector < sec_w + p.sections[2]) {
theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, i0/2.0f);
theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, iw/2.0f);
}
else if (sector >= sec_w + p.sections[2]) {
theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, i0/2.0f);
theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, iw/2.0f);
}
}
const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;
const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f;
float cos_theta, sin_theta;
rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p);
rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p);
const float x0 = float(rope_data_a[ix + 0]);
const float x1 = float(rope_data_a[ix + p.n_dims/2]);
// idst/ix point at channel i0/2; the first channel of the rotated pair is p.n_offs + iw/2 = i0/2 + p.n_offs/2
const float x0 = float(rope_data_a[ix + p.n_offs/2 + 0]);
const float x1 = float(rope_data_a[ix + p.n_offs/2 + p.n_dims/2]);
rope_data_d[idst + 0] = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
rope_data_d[idst + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
rope_data_d[idst + p.n_offs/2 + 0] = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta);
rope_data_d[idst + p.n_offs/2 + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
}
void rope_vision(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {
@@ -5,6 +5,7 @@ struct rope_params {
uint rope_mode;
uint nrows;
uint n_dims;
uint n_offs;
float freq_scale;
float freq_base;
float ext_factor;