diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 1f6e8c48bc..05ea7470ee 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -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 { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index b7f9b2d0d9..d8435e9577 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -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); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 243c997fc4..57a0873f97 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -4688,14 +4688,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); @@ -4706,6 +4707,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); @@ -4741,17 +4746,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]; @@ -4759,6 +4765,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); @@ -4793,8 +4803,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 @@ -4827,14 +4838,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]; @@ -4842,6 +4853,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); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 3349a64b17..675c884d5b 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5331,17 +5331,19 @@ 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 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;