Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal-impl.h
+2
View File
@@ -329,6 +329,7 @@ typedef struct {
uint64_t nb3;
int32_t n_past;
int32_t n_dims;
int32_t n_offs;
int32_t n_ctx_orig;
float freq_base;
float freq_scale;
@@ -341,6 +342,7 @@ typedef struct {
int32_t sect_2;
int32_t sect_3;
bool src2;
bool inplace;
} ggml_metal_kargs_rope;
typedef struct {
ggml/src/ggml-metal/ggml-metal-ops.cpp
+7
View File
@@ -3884,6 +3884,11 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) {
const int sect_2 = ((const int32_t *) op->op_params)[13];
const int sect_3 = ((const int32_t *) op->op_params)[14];
const int n_offs = ((const int32_t *) op->op_params)[15];
// when dst aliases src0, the channels outside the rotated window already hold the correct data
const bool inplace = op->data == op->src[0]->data;
ggml_metal_kargs_rope args = {
/*.ne00 =*/ ne00,
/*.ne01 =*/ ne01,
@@ -3903,6 +3908,7 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) {
/*.nb3 =*/ nb3,
/*.n_past =*/ n_past,
/*.n_dims =*/ n_dims,
/*.n_offs =*/ n_offs,
/*.n_ctx_orig =*/ n_ctx_orig,
/*.freq_base =*/ freq_base,
/*.freq_scale =*/ freq_scale,
@@ -3915,6 +3921,7 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) {
/* sect_2 =*/ sect_2,
/* sect_3 =*/ sect_3,
/* src2 =*/ op->src[2] != nullptr,
/* inplace =*/ inplace,
};
auto pipeline = ggml_metal_library_get_pipeline_rope(lib, op);
ggml/src/ggml-metal/ggml-metal.metal
+31 -16
View File
@@ -4686,14 +4686,15 @@ kernel void kernel_rope_norm(
float sin_theta;
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
if (i0 < args.n_dims) {
const int ic = i0/2;
if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
const int iw = i0 - args.n_offs; // relative idx
const int ic = iw/2;
const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
@@ -4704,6 +4705,10 @@ kernel void kernel_rope_norm(
dst_data[0] = x0*cos_theta - x1*sin_theta;
dst_data[1] = x0*sin_theta + x1*cos_theta;
} else {
if (args.inplace) {
continue;
}
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
@@ -4739,17 +4744,18 @@ kernel void kernel_rope_neox(
float sin_theta;
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
if (i0 < args.n_dims) {
const int ic = i0/2;
if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
const int iw = i0 - args.n_offs; // relative idx
const int ic = iw/2;
const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + ic*args.nb0);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + (args.n_offs + ic)*args.nb0);
const float x0 = src[0];
const float x1 = src[args.n_dims/2];
@@ -4757,6 +4763,10 @@ kernel void kernel_rope_neox(
dst_data[0] = x0*cos_theta - x1*sin_theta;
dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
} else {
if (args.inplace) {
continue;
}
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
@@ -4791,8 +4801,9 @@ kernel void kernel_rope_multi(
float sin_theta;
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
if (i0 < args.n_dims) {
const int ic = i0/2;
if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) {
const int iw = i0 - args.n_offs; // relative idx
const int ic = iw/2;
// mrope theta calculations
// note: the rest is the same as kernel_rope_neox
@@ -4825,14 +4836,14 @@ kernel void kernel_rope_multi(
}
// end of mrope
const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
const float theta = theta_base * pow(args.freq_base, inv_ndims*iw);
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + ic*args.nb0);
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + (args.n_offs + ic)*args.nb0);
const float x0 = src[0];
const float x1 = src[args.n_dims/2];
@@ -4840,6 +4851,10 @@ kernel void kernel_rope_multi(
dst_data[0] = x0*cos_theta - x1*sin_theta;
dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
} else {
if (args.inplace) {
continue;
}
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);