From 8442c74f5673972eb6feea583cd16d93a6d02a65 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Wed, 19 Aug 2026 14:04:57 +0200 Subject: [PATCH] ggml: add ggml_rope_set_offset (+ metal support) (llama/27120) * add params * cpu kernel * metal kernel * add test backend ops * gate other backends * ggml: (cuda) support ggml_rope_set_offset (llama/27121) * rm cuda supports_op guard, fix webgpu clang-format * ggml: support ggml_rope_set_offset on vulkan (llama/27344) * ggml: support ggml_rope_set_offset on vulkan * remove inplace optimization --- ggml/include/ggml.h | 8 + ggml/src/ggml-cann/ggml-cann.cpp | 3 + ggml/src/ggml-cpu/ops.cpp | 17 +- ggml/src/ggml-cuda/ggml-cuda.cu | 6 + ggml/src/ggml-cuda/rope.cu | 152 +++++++++++------- ggml/src/ggml-et/ggml-et.cpp | 4 +- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 4 + ggml/src/ggml-metal/ggml-metal-impl.h | 2 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 7 + ggml/src/ggml-metal/ggml-metal.metal | 47 ++++-- ggml/src/ggml-opencl/ggml-opencl.cpp | 3 + ggml/src/ggml-openvino/ggml-openvino.cpp | 4 + ggml/src/ggml-sycl/ggml-sycl.cpp | 2 + ggml/src/ggml-vulkan/ggml-vulkan.cpp | 8 +- .../vulkan-shaders/rope_funcs.glsl | 64 ++++---- .../vulkan-shaders/rope_params.glsl | 1 + ggml/src/ggml-webgpu/ggml-webgpu.cpp | 4 +- ggml/src/ggml.c | 18 ++- 18 files changed, 244 insertions(+), 110 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index c2ccd9725..32462d79a 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -1981,6 +1981,14 @@ extern "C" { float beta_fast, float beta_slow); + // set the offset dims for RoPE + // a must be GGML_OP_ROPE or GGML_OP_ROPE_BACK + // vision RoPE is not supported + // example: (marking: x = rotated, 0 = unrotated) + // n_embd = 10, n_dims = 4, offset = 2 --> [00xxxx0000] + GGML_API struct ggml_tensor * ggml_rope_set_offset( + struct ggml_tensor * a, + int n_offs); // clamp // in-place, returns view(a) diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index ffa361af4..5e5541aac 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2534,6 +2534,9 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten } case GGML_OP_ROPE: { + if (((const int32_t *) op->op_params)[15] != 0) { + return false; // FIXME: support ggml_rope_set_offset + } if (op->src[0]->ne[0] > 896) { return false; } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 001e1ae85..2b5f68444 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -5979,6 +5979,8 @@ static void ggml_compute_forward_rope_flt( memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float)); memcpy(§ions, (int32_t *) dst->op_params + 11, sizeof(int)*4); + const int n_offs = ((int32_t *) dst->op_params)[15]; + GGML_TENSOR_UNARY_OP_LOCALS //printf("ne0: %d, ne1: %d, ne2: %d, ne3: %d\n", ne0, ne1, ne2, ne3); @@ -5995,6 +5997,10 @@ static void ggml_compute_forward_rope_flt( GGML_ASSERT(n_dims <= ne0); GGML_ASSERT(n_dims % 2 == 0); + GGML_ASSERT(n_offs >= 0); + GGML_ASSERT(n_offs % 2 == 0); + GGML_ASSERT(n_offs + n_dims <= ne0); + // rows per thread const int dr = (nr + nth - 1)/nth; @@ -6020,6 +6026,7 @@ static void ggml_compute_forward_rope_flt( if (is_vision) { GGML_ASSERT(n_dims == ne0/2); + GGML_ASSERT(n_offs == 0); } const float * freq_factors = NULL; @@ -6068,12 +6075,12 @@ static void ggml_compute_forward_rope_flt( switch (mode) { case GGML_ROPE_TYPE_NORMAL: - rotate_pairs(n_dims, 1, cache, src, dst_data, 1); + rotate_pairs(n_dims, 1, cache, src + n_offs, dst_data + n_offs, 1); break; case GGML_ROPE_TYPE_NEOX: case GGML_ROPE_TYPE_MROPE: case GGML_ROPE_TYPE_IMROPE: - rotate_pairs(n_dims, n_dims/2, cache, src, dst_data); + rotate_pairs(n_dims, n_dims/2, cache, src + n_offs, dst_data + n_offs); break; case GGML_ROPE_TYPE_VISION: rotate_pairs(ne0, n_dims, cache, src, dst_data); @@ -6084,7 +6091,11 @@ static void ggml_compute_forward_rope_flt( if (!is_vision) { // fill the remain channels with data from src tensor - for (int64_t i0 = n_dims; i0 < ne0; i0 += 2) { + for (int64_t i0 = 0; i0 < ne0; i0 += 2) { + if (i0 == n_offs) { + i0 += n_dims - 2; // skip the rotated channels + continue; + } const T * const src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); T * dst_data = (T *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index f2e381ee0..3b2a0ea85 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2723,6 +2723,12 @@ static bool ggml_cuda_should_fuse_rms_norm_mul_rope(const ggml_tensor * rms_norm return false; } + // ggml_rope_set_offset is not yet supported in the fused kernel + const int n_offs = ((const int32_t *) rope->op_params)[15]; + if (n_offs != 0) { + return false; + } + return true; } diff --git a/ggml/src/ggml-cuda/rope.cu b/ggml/src/ggml-cuda/rope.cu index 504c6b818..e546fb655 100644 --- a/ggml/src/ggml-cuda/rope.cu +++ b/ggml/src/ggml-cuda/rope.cu @@ -53,6 +53,7 @@ static __global__ void rope_norm(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int32_t * pos, const float freq_scale, const float ext_factor, @@ -61,7 +62,8 @@ static __global__ void rope_norm(const T * x, const float theta_scale, const float * freq_factors, const int64_t * row_indices, - const int set_rows_stride) { + const int set_rows_stride, + const bool inplace) { const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y); if (i0 >= ne00) { @@ -92,19 +94,24 @@ static __global__ void rope_norm(const T * x, ggml_cuda_memcpy_1<4>(dst + idst, &v); } }; - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { + if (inplace) { + return; + } store_coaelsced(x[ix + 0], x[ix + 1]); return; } - const float theta_base = pos[i2]*powf(theta_scale, i0/2.0f); + const int iw = i0 - n_offs; // relative idx - const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + const float theta_base = pos[i2]*powf(theta_scale, iw/2.0f); + + const float freq_factor = has_ff ? freq_factors[iw/2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta); const float x0 = x[ix + 0]; const float x1 = x[ix + 1]; @@ -125,6 +132,7 @@ static __global__ void rope_neox(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int32_t * pos, const float freq_scale, const float ext_factor, @@ -133,7 +141,8 @@ static __global__ void rope_neox(const T * x, const float theta_scale, const float * freq_factors, const int64_t * row_indices, - const int set_rows_stride) { + const int set_rows_stride, + const bool inplace) { ggml_cuda_pdl_lc(); const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y); @@ -158,27 +167,33 @@ static __global__ void rope_neox(const T * x, idst += row_indices[i2] * set_rows_stride; } - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { + if (inplace) { + return; + } dst[idst + i0 / 2 + 0] = ggml_cuda_cast(x[ix + i0 / 2 + 0]); dst[idst + i0 / 2 + 1] = ggml_cuda_cast(x[ix + i0 / 2 + 1]); return; } - const float theta_base = pos[i2]*powf(theta_scale, i0/2.0f); + const int iw = i0 - n_offs; // relative idx - const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + const float theta_base = pos[i2]*powf(theta_scale, iw/2.0f); + + const float freq_factor = has_ff ? freq_factors[iw/2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta); - const float x0 = x[ix + 0]; - const float x1 = x[ix + n_dims/2]; + // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2 + const float x0 = x[ix + n_offs/2 + 0]; + const float x1 = x[ix + n_offs/2 + n_dims/2]; - dst[idst + 0] = ggml_cuda_cast(x0 * cos_theta - x1 * sin_theta); - dst[idst + n_dims / 2] = ggml_cuda_cast(x0 * sin_theta + x1 * cos_theta); + dst[idst + n_offs/2 + 0] = ggml_cuda_cast(x0 * cos_theta - x1 * sin_theta); + dst[idst + n_offs/2 + n_dims / 2] = ggml_cuda_cast(x0 * sin_theta + x1 * cos_theta); } template @@ -194,6 +209,7 @@ static __global__ void rope_multi(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int32_t * pos, const float freq_scale, const float ext_factor, @@ -202,7 +218,8 @@ static __global__ void rope_multi(const T * x, const float theta_scale, const float * freq_factors, const mrope_sections sections, - const bool is_imrope) { + const bool is_imrope, + const bool inplace) { const int i0 = 2 * (blockDim.y * blockIdx.y + threadIdx.y); if (i0 >= ne00) { @@ -219,52 +236,58 @@ static __global__ void rope_multi(const T * x, const int ix = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03; ggml_cuda_pdl_sync(); - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { + if (inplace) { + return; + } dst[idst + i0/2 + 0] = x[ix + i0/2 + 0]; dst[idst + i0/2 + 1] = x[ix + i0/2 + 1]; return; } + const int iw = i0 - n_offs; // relative idx + const int sect_dims = sections.v[0] + sections.v[1] + sections.v[2] + sections.v[3]; const int sec_w = sections.v[1] + sections.v[0]; - const int sector = (i0 / 2) % sect_dims; + const int sector = (iw / 2) % sect_dims; float theta_base = 0.0; if (is_imrope) { if (sector % 3 == 1 && sector < 3 * sections.v[1]) { // h - theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, iw / 2.0f); } else if (sector % 3 == 2 && sector < 3 * sections.v[2]) { // w - theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, iw / 2.0f); } else if (sector % 3 == 0 && sector < 3 * sections.v[0]) { // t - theta_base = pos[i2] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2] * powf(theta_scale, iw / 2.0f); } else { - theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, iw / 2.0f); } } else { if (sector < sections.v[0]) { - theta_base = pos[i2] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2] * powf(theta_scale, iw / 2.0f); } else if (sector >= sections.v[0] && sector < sec_w) { - theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, iw / 2.0f); } else if (sector >= sec_w && sector < sec_w + sections.v[2]) { - theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, iw / 2.0f); } else if (sector >= sec_w + sections.v[2]) { - theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, iw / 2.0f); } } - const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + const float freq_factor = has_ff ? freq_factors[iw/2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta); - const float x0 = x[ix + 0]; - const float x1 = x[ix + n_dims/2]; + // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2 + const float x0 = x[ix + n_offs/2 + 0]; + const float x1 = x[ix + n_offs/2 + n_dims/2]; - dst[idst + 0] = x0*cos_theta - x1*sin_theta; - dst[idst + n_dims/2] = x0*sin_theta + x1*cos_theta; + dst[idst + n_offs/2 + 0] = x0*cos_theta - x1*sin_theta; + dst[idst + n_offs/2 + n_dims/2] = x0*sin_theta + x1*cos_theta; } template @@ -344,6 +367,7 @@ static void rope_norm_cuda(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int nr, const int32_t * pos, const float freq_scale, @@ -354,6 +378,7 @@ static void rope_norm_cuda(const T * x, const float * freq_factors, const int64_t * row_indices, const int set_rows_stride, + const bool inplace, cudaStream_t stream) { GGML_ASSERT(ne00 % 2 == 0); const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); @@ -364,12 +389,12 @@ static void rope_norm_cuda(const T * x, if (freq_factors == nullptr) { rope_norm<<>>( - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace); } else { rope_norm<<>>( - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace); } } @@ -386,6 +411,7 @@ static void rope_neox_cuda(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int nr, const int32_t * pos, const float freq_scale, @@ -396,6 +422,7 @@ static void rope_neox_cuda(const T * x, const float * freq_factors, const int64_t * row_indices, const int set_rows_stride, + const bool inplace, cudaStream_t stream) { GGML_ASSERT(ne00 % 2 == 0); const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); @@ -407,12 +434,12 @@ static void rope_neox_cuda(const T * x, if (freq_factors == nullptr) { ggml_cuda_kernel_launch(rope_neox, launch_params, - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace); } else { ggml_cuda_kernel_launch(rope_neox, launch_params, - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace); } } @@ -429,6 +456,7 @@ static void rope_multi_cuda(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int nr, const int32_t * pos, const float freq_scale, @@ -439,6 +467,7 @@ static void rope_multi_cuda(const T * x, const float * freq_factors, const mrope_sections sections, const bool is_imrope, + const bool inplace, cudaStream_t stream) { GGML_ASSERT(ne00 % 2 == 0); const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); @@ -450,13 +479,13 @@ static void rope_multi_cuda(const T * x, if (freq_factors == nullptr) { const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); ggml_cuda_kernel_launch(rope_multi, launch_params, - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace); } else { const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); ggml_cuda_kernel_launch(rope_multi, launch_params, - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace); } } @@ -552,8 +581,12 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, const int mode = ((int32_t *) dst->op_params)[2]; //const int n_ctx = ((int32_t *) dst->op_params)[3]; const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; + const int n_offs = ((int32_t *) dst->op_params)[15]; mrope_sections sections; + // when dst aliases src0, the channels outside the rotated window already hold the correct data + const bool inplace = dst_d == src0->data; + // RoPE alteration for extended context float freq_base; float freq_scale; @@ -581,6 +614,7 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, if (is_vision) { GGML_ASSERT(n_dims == ne00/2); + GGML_ASSERT(n_offs == 0); // offset not supported for vision, as the rotated pairs span the whole row } const int32_t * pos = (const int32_t *) src1_d; @@ -597,31 +631,31 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, if (is_neox) { if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) { rope_neox_cuda((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { rope_neox_cuda((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) { rope_neox_cuda((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else { GGML_ABORT("fatal error"); } } else if (is_mrope && !is_vision) { if (src0->type == GGML_TYPE_F32) { rope_multi_cuda((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1, - s2, s3, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, - corr_dims, freq_factors, sections, is_imrope, stream); + s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, + corr_dims, freq_factors, sections, is_imrope, inplace, stream); } else if (src0->type == GGML_TYPE_F16) { rope_multi_cuda((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1, - s2, s3, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, - corr_dims, freq_factors, sections, is_imrope, stream); + s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, + corr_dims, freq_factors, sections, is_imrope, inplace, stream); } else { GGML_ABORT("fatal error"); } @@ -640,19 +674,19 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, } else { if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) { rope_norm_cuda((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { rope_norm_cuda((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) { rope_norm_cuda((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else { GGML_ABORT("fatal error"); } diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp index e8482f734..b87b189a5 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1061,9 +1061,11 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm const bool zero_view_offset = op->src[0]->view_src == nullptr || op->src[0]->view_offs == 0; const bool has_sections = ggml_get_op_params_i32(op, 11) > 0 || ggml_get_op_params_i32(op, 12) > 0 || ggml_get_op_params_i32(op, 13) > 0; + // FIXME: support ggml_rope_set_offset + const bool zero_rot_offset = ggml_get_op_params_i32(op, 15) == 0; supported = - zero_view_offset && ndims <= 512 && + zero_view_offset && zero_rot_offset && ndims <= 512 && (is_normal || (is_neox && ndims % 16 == 0) || (is_imrope && ndims % 16 == 0 && has_sections)); } else { supported = false; diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index f80c60a50..b262a73d9 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3180,6 +3180,10 @@ static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * s static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { const int32_t * op_params = &op->op_params[0]; + if (op_params[15] != 0) { + return false; // FIXME: support ggml_rope_set_offset + } + int mode = op_params[2]; // n_dims == ne0/2, so the rotation spans the full row diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 1f6e8c48b..05ea7470e 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 b7f9b2d0d..d8435e957 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 c3a95dd6e..0537fa4cf 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -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); diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 257908605..733fab1c3 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7376,6 +7376,9 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te case GGML_OP_DIAG_MASK_INF: return op->ne[3] == 1; case GGML_OP_ROPE: { + if (((const int32_t *) op->op_params)[15] != 0) { + return false; // FIXME: support ggml_rope_set_offset + } const int mode = ((const int32_t *) op->op_params)[2]; const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE; const bool is_vision = mode == GGML_ROPE_TYPE_VISION; diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index cac83a1bd..e299e16c7 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -1227,6 +1227,10 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { const int32_t * op_params = op->op_params; const int n_dims = op_params[1]; const int mode = op_params[2]; + if (op_params[15] != 0) { + // FIXME: support ggml_rope_set_offset + return true; + } if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && mode != GGML_ROPE_TYPE_IMROPE) { // GGML_LOG_WARN("OpenVINO backend does not support ROPE with mode %d\n", mode); return true; diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index d31df611a..92c26839f 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6235,6 +6235,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons } case GGML_OP_ROPE: case GGML_OP_ROPE_BACK: + // FIXME: support ggml_rope_set_offset + return ((const int32_t *) op->op_params)[15] == 0; case GGML_OP_IM2COL: case GGML_OP_IM2COL_3D: case GGML_OP_UPSCALE: diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 356f9abd7..62ef7bb8a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1645,6 +1645,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; @@ -13144,6 +13145,7 @@ static uint32_t ggml_vk_rms_partials_size(ggml_backend_vk_context * ctx, const g 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 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]; @@ -13173,7 +13175,7 @@ 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, @@ -19219,6 +19221,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 033587931..feb55b203 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl @@ -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) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl index 3602485b9..b88a73fcc 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; diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index b79f71b68..4367f9a61 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4472,7 +4472,9 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const supports_op = (op->type == GGML_TYPE_F32 && src0->type == GGML_TYPE_F32) && ggml_is_contiguous_rows(src0); break; case GGML_OP_ROPE: - supports_op = op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16; + // FIXME: support ggml_rope_set_offset + supports_op = + (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && ((const int32_t *) op->op_params)[15] == 0; break; case GGML_OP_GLU: switch (ggml_get_glu_op(op)) { diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index d0d369c41..1a60fec79 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -4200,7 +4200,7 @@ static struct ggml_tensor * ggml_rope_impl( struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - int32_t params[15] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig }; + int32_t params[16] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig }; memcpy(params + 5, &freq_base, sizeof(float)); memcpy(params + 6, &freq_scale, sizeof(float)); memcpy(params + 7, &ext_factor, sizeof(float)); @@ -4212,6 +4212,8 @@ static struct ggml_tensor * ggml_rope_impl( } else { memset(params + 11, 0, sizeof(int32_t) * GGML_MROPE_SECTIONS); } + params[15] = 0; // n_offs, set via ggml_rope_set_offset() + ggml_set_op_params(result, params, sizeof(params)); result->op = GGML_OP_ROPE; @@ -4422,6 +4424,20 @@ struct ggml_tensor * ggml_rope_multi_back( result->op = GGML_OP_ROPE_BACK; return result; } + +struct ggml_tensor * ggml_rope_set_offset( + struct ggml_tensor * a, + int n_offs) { + GGML_ASSERT(a->op == GGML_OP_ROPE || a->op == GGML_OP_ROPE_BACK); + GGML_ASSERT(n_offs >= 0); + + const int32_t mode = ggml_get_op_params_i32(a, 2); + GGML_ASSERT(mode != GGML_ROPE_TYPE_VISION); + + ggml_set_op_params_i32(a, 15, n_offs); + return a; +} + // ggml_clamp struct ggml_tensor * ggml_clamp(