diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 5ce83d7e97..c981c55e51 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1644,6 +1644,7 @@ struct vk_op_rope_push_constants { uint32_t rope_mode; uint32_t nrows; uint32_t n_dims; + uint32_t n_offs; float freq_scale; float freq_base; float ext_factor; @@ -1655,6 +1656,7 @@ struct vk_op_rope_push_constants { uint32_t is_imrope; uint32_t is_back; uint32_t set_rows_stride; + uint32_t inplace; uint32_t ne00; uint32_t ne01; uint32_t ne02; @@ -13117,9 +13119,10 @@ static uint32_t ggml_vk_rms_partials_size(ggml_backend_vk_context * ctx, const g return num_bytes; } -static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor *dst, const ggml_tensor *src0, const bool has_ff, bool backprop, const uint32_t set_rows_stride) { +static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor *dst, const ggml_tensor *src0, const bool has_ff, bool backprop, const uint32_t set_rows_stride, const bool inplace) { const int n_dims = ((const int32_t *) dst->op_params)[1]; const int mode = ((const int32_t *) dst->op_params)[2]; + const int n_offs = ((const int32_t *) dst->op_params)[15]; // const int n_ctx = ((const int32_t *) dst->op_params)[3]; const int n_ctx_orig = ((const int32_t *) dst->op_params)[4]; const float freq_base = ((const float *) dst->op_params)[5]; @@ -13149,9 +13152,9 @@ static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor * uint32_t nb13 = dst->nb[3] / ggml_type_size(dst->type); vk_op_rope_push_constants rope { - (uint32_t)mode, (uint32_t)ggml_nrows(src0), (uint32_t)n_dims, freq_scale, + (uint32_t)mode, (uint32_t)ggml_nrows(src0), (uint32_t)n_dims, (uint32_t)n_offs, freq_scale, freq_base, ext_factor, attn_factor, {corr_dims[0], corr_dims[1]}, theta_scale, has_ff, - { sections[0], sections[1], sections[2], sections[3] }, is_imrope, backprop, set_rows_stride, + { sections[0], sections[1], sections[2], sections[3] }, is_imrope, backprop, set_rows_stride, inplace, (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], @@ -13220,7 +13223,7 @@ static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, vk_op_rms_norm_mul_rope_push_constants pc; pc.bin = bin; - pc.rope = ggml_vk_make_rope_constants(rope, rope->src[0], tensors[4] != nullptr, false, set_rows_stride); + pc.rope = ggml_vk_make_rope_constants(rope, rope->src[0], tensors[4] != nullptr, false, set_rows_stride, false); vk_pipeline pipeline = tensors[5]->type == GGML_TYPE_F16 ? ctx->device->pipeline_rms_norm_mul_rope_f32_f16 : ctx->device->pipeline_rms_norm_mul_rope_f32_f32; @@ -13568,8 +13571,11 @@ static void ggml_vk_rope(ggml_backend_vk_context * ctx, vk_context& subctx, cons dst = cgraph->nodes[node_idx + 2]; } + // when dst aliases src0, the channels outside the rotated window already hold the correct data + const bool inplace = dst->data == src0->data; + ggml_vk_op_f32(ctx, subctx, src0, src1, src2, src3, dst, GGML_OP_ROPE, - ggml_vk_make_rope_constants(cgraph->nodes[node_idx], src0, src2 != nullptr, backprop, set_rows_stride)); + ggml_vk_make_rope_constants(cgraph->nodes[node_idx], src0, src2 != nullptr, backprop, set_rows_stride, inplace)); } static void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) { @@ -18253,15 +18259,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_REPEAT_BACK: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32; case GGML_OP_ROPE: - if (((const int32_t *) op->op_params)[15] != 0) { - return false; // FIXME: support ggml_rope_set_offset - } return ggml_is_contiguous_rows(op) && ggml_is_contiguous_rows(op->src[0]); case GGML_OP_ROPE_BACK: - if (((const int32_t *) op->op_params)[15] != 0) { - return false; // FIXME: support ggml_rope_set_offset - } - return true; case GGML_OP_NONE: case GGML_OP_RESHAPE: case GGML_OP_VIEW: @@ -19202,6 +19201,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * tensor_clone = ggml_rope_ext_back(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], n_dims, mode, n_ctx_orig_ggml, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); } } + const int n_offs = ((int32_t *) tensor->op_params)[15]; + if (n_offs != 0) { + tensor_clone = ggml_rope_set_offset(tensor_clone, n_offs); + } } else if (tensor->op == GGML_OP_UNARY) { switch (ggml_get_unary_op(tensor)) { case GGML_UNARY_OP_EXP: diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl index 0335879314..009e04dce9 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl @@ -50,19 +50,24 @@ 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) { + if (p.inplace != 0) { + return; + } 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; // channel index relative to the rotated window - 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 +92,31 @@ 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) { + if (p.inplace != 0) { + return; + } 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; // channel index relative to the rotated window - 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 +136,59 @@ 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) { + if (p.inplace != 0) { + return; + } 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; // channel index relative to the rotated window + 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) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl index 3602485b94..7aecbc2ef1 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl @@ -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; @@ -16,6 +17,7 @@ struct rope_params { uint is_imrope; uint is_back; uint set_rows_stride; + uint inplace; uint ne00; uint ne01;