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
+28 -6
View File
@@ -5352,24 +5352,27 @@ struct test_rope : public test_case {
int v; // view (1 : non-contiguous a)
bool forward;
bool inplace;
int n_offs; // offset of the rotated dims window, set via ggml_rope_set_offset()
std::string vars() override {
// forward can be inferred from the op, does not need to be printed
return VARS_TO_STR11(type, ne_a, n_dims, mode, n_ctx, fs, ef, af, ff, v, inplace);
return VARS_TO_STR12(type, ne_a, n_dims, mode, n_ctx, fs, ef, af, ff, v, inplace, n_offs);
}
test_rope(ggml_type type = GGML_TYPE_F32,
std::array<int64_t, 4> ne_a = {10, 5, 3, 1},
int n_dims = 10, int mode = GGML_ROPE_TYPE_NORMAL, int n_ctx = 512, float fs = 1.0f,
float ef = 0.0f, float af = 0.0f, bool ff = false, int v = 0, bool forward = true, bool inplace = false)
: type(type), ne_a(ne_a), n_dims(n_dims), mode(mode), n_ctx(n_ctx), fs(fs), ef(ef), af(af), ff(ff), v(v), forward(forward), inplace(inplace) {}
float ef = 0.0f, float af = 0.0f, bool ff = false, int v = 0, bool forward = true, bool inplace = false,
int n_offs = 0)
: type(type), ne_a(ne_a), n_dims(n_dims), mode(mode), n_ctx(n_ctx), fs(fs), ef(ef), af(af), ff(ff), v(v), forward(forward), inplace(inplace), n_offs(n_offs) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
ggml_tensor * a;
if (v & 1) {
auto ne = ne_a; ne[0] *= 2; ne[1] *= 4; ne[2] *= 3;
a = ggml_new_tensor(ctx, type, 4, ne.data());
if (forward) {
if (forward && n_offs == 0) {
// FIXME: support gradients with n_offs > 0
ggml_set_param(a);
}
ggml_set_name(a, "a");
@@ -5382,7 +5385,8 @@ struct test_rope : public test_case {
// non-aligned buffer offset, which exercises backends' alignment paths.
auto ne = ne_a; ne[0] *= 2;
a = ggml_new_tensor(ctx, type, 4, ne.data());
if (forward) {
if (forward && n_offs == 0) {
// FIXME: support gradients with n_offs > 0
ggml_set_param(a);
}
ggml_set_name(a, "a");
@@ -5393,7 +5397,8 @@ struct test_rope : public test_case {
ggml_set_name(a, "view_of_a");
} else {
a = ggml_new_tensor(ctx, type, 4, ne_a.data());
if (forward) {
if (forward && n_offs == 0) {
// FIXME: support gradients with n_offs > 0
ggml_set_param(a);
}
ggml_set_name(a, "a");
@@ -5454,6 +5459,9 @@ struct test_rope : public test_case {
out = ggml_rope_ext_back(ctx, a, pos, freq, n_dims, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f);
}
}
if (n_offs != 0) {
out = ggml_rope_set_offset(out, n_offs);
}
ggml_set_name(out, "out");
return out;
@@ -9621,6 +9629,20 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
// rotated dims window at an offset (ggml_rope_set_offset), not supported for vision mode
for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16}) {
for (bool fw : {true, false}) { // fw == forward
for (bool ff : {false, true}) {
test_cases.emplace_back(new test_rope(type, {128, 32, 2, 1}, 32, GGML_ROPE_TYPE_NORMAL, 512, 1.4245f, 0.7465f, 1.4245f, ff, 0, fw, false, 32));
test_cases.emplace_back(new test_rope(type, {128, 32, 2, 1}, 32, GGML_ROPE_TYPE_NEOX, 512, 1.4245f, 0.7465f, 1.4245f, ff, 0, fw, false, 32));
test_cases.emplace_back(new test_rope(type, {128, 12, 2, 1}, 24, GGML_ROPE_TYPE_MROPE, 512, 1.4245f, 0.7465f, 1.4245f, ff, 0, fw, false, 32));
test_cases.emplace_back(new test_rope(type, {128, 12, 2, 1}, 24, GGML_ROPE_TYPE_IMROPE, 512, 1.4245f, 0.7465f, 1.4245f, ff, 0, fw, false, 32));
}
}
// inplace with an offset
test_cases.emplace_back(new test_rope(type, {128, 32, 2, 1}, 32, GGML_ROPE_TYPE_NEOX, 512, 1.4245f, 0.7465f, 1.4245f, false, 0, true, true, 32));
}
for (int v : { 0, 1, 2, 3 }) {
for (int dim : { 0, 1, 2, 3, }) {
test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {11, 12, 13, 14}, 7, dim, v));