diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index b262a73d950e..e8a5009b381b 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3180,8 +3180,9 @@ 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 + // ggml_rope_set_offset: HVX kernels need a VLEN-aligned window start (32 f32 elems) + if (op_params[15] % 32 != 0) { + return false; } int mode = op_params[2]; diff --git a/ggml/src/ggml-hexagon/htp/rope-ops.c b/ggml/src/ggml-hexagon/htp/rope-ops.c index 5bc7d74f5e21..6c689824934f 100644 --- a/ggml/src/ggml-hexagon/htp/rope-ops.c +++ b/ggml/src/ggml-hexagon/htp/rope-ops.c @@ -53,6 +53,7 @@ struct htp_rope_context { int32_t n_dims; + int32_t n_offs; int32_t mode; int32_t n_ctx_orig; int32_t sections[4]; @@ -405,32 +406,40 @@ static inline void hvx_rope_f32_aa(float * restrict dst, const float * restrict static void inline rope_basic_f32(struct htp_rope_context * rctx, uint8_t * restrict dst, uint8_t * restrict src, uint32_t nr, uint32_t ne0, const float * restrict theta_cache) { + const uint32_t n_offs = rctx->n_offs; // VLEN-aligned (enforced by supports_op) #pragma unroll(4) for (uint32_t i = 0; i < nr; i++) { float * d = (float *) (dst + i * rctx->dst_row_size_aligned); float * s = (float *) (src + i * rctx->src0_row_size_aligned); - hvx_rope_f32_aa(d, s, rctx->n_dims, theta_cache); + hvx_rope_f32_aa(d + n_offs, s + n_offs, rctx->n_dims, theta_cache); // fill the remain channels with data from src tensor - if (rctx->n_dims < ne0) { - hvx_copy_f32_uu((uint8_t *)(d + rctx->n_dims), (uint8_t *)(s + rctx->n_dims), ne0 - rctx->n_dims); + if (n_offs > 0) { + hvx_copy_f32_uu((uint8_t *) d, (uint8_t *) s, n_offs); + } + if (n_offs + rctx->n_dims < ne0) { + hvx_copy_f32_uu((uint8_t *)(d + n_offs + rctx->n_dims), (uint8_t *)(s + n_offs + rctx->n_dims), ne0 - n_offs - rctx->n_dims); } } } static void inline rope_neox_f32(struct htp_rope_context * rctx, uint8_t * restrict dst, uint8_t * restrict src, uint32_t nr, uint32_t ne0, const float * restrict theta_cache) { + const uint32_t n_offs = rctx->n_offs; // VLEN-aligned (enforced by supports_op) #pragma unroll(4) for (uint32_t i = 0; i < nr; i++) { float * d = (float *) (dst + i * rctx->dst_row_size_aligned); float * s = (float *) (src + i * rctx->src0_row_size_aligned); - hvx_rope_neox_f32_aa(d, s, rctx->n_dims, theta_cache); + hvx_rope_neox_f32_aa(d + n_offs, s + n_offs, rctx->n_dims, theta_cache); // fill the remain channels with data from src tensor - if (rctx->n_dims < ne0) { - hvx_copy_f32_uu((uint8_t *)(d + rctx->n_dims), (uint8_t *)(s + rctx->n_dims), ne0 - rctx->n_dims); + if (n_offs > 0) { + hvx_copy_f32_uu((uint8_t *) d, (uint8_t *) s, n_offs); + } + if (n_offs + rctx->n_dims < ne0) { + hvx_copy_f32_uu((uint8_t *)(d + n_offs + rctx->n_dims), (uint8_t *)(s + n_offs + rctx->n_dims), ne0 - n_offs - rctx->n_dims); } } } @@ -673,6 +682,7 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) { rctx.n_dims = ((const int32_t *) op_params)[1]; rctx.mode = ((const int32_t *) op_params)[2]; rctx.n_ctx_orig = ((const int32_t *) op_params)[4]; + rctx.n_offs = ((const int32_t *) op_params)[15]; memcpy(&rctx.freq_base, (int32_t *) op_params + 5, sizeof(float)); memcpy(&rctx.freq_scale, (int32_t *) op_params + 6, sizeof(float)); diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 733fab1c34a1..e8a49e9624bc 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7376,9 +7376,6 @@ 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; @@ -23706,6 +23703,7 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const const int n_dims = ((int *) dst->op_params)[1]; const int mode = ((int *) dst->op_params)[2]; const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; + const int n_offs = ((int32_t *) dst->op_params)[15]; float freq_base; float freq_scale; @@ -23734,6 +23732,7 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const 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 } cl_kernel kernel; @@ -23825,6 +23824,12 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const if (is_mrope && !is_vision) { CL_CHECK(clSetKernelArg(kernel, 34, sizeof(int), &is_imrope)); } + // norm and neox have n_offs after beta_slow, mrope has it after is_imrope + if (!is_mrope && !is_vision) { + CL_CHECK(clSetKernelArg(kernel, 33, sizeof(int), &n_offs)); + } else if (is_mrope && !is_vision) { + CL_CHECK(clSetKernelArg(kernel, 35, sizeof(int), &n_offs)); + } size_t global_work_size[] = {(size_t)ne01*nth, (size_t)ne02, (size_t)ne03}; size_t local_work_size[] = {(size_t)nth, 1, 1}; diff --git a/ggml/src/ggml-opencl/kernels/rope.cl b/ggml/src/ggml-opencl/kernels/rope.cl index 82f4cd87407d..27fdbbbc4ff1 100644 --- a/ggml/src/ggml-opencl/kernels/rope.cl +++ b/ggml/src/ggml-opencl/kernels/rope.cl @@ -75,7 +75,8 @@ kernel void kernel_rope_norm_f32( float ext_factor, float attn_factor, float beta_fast, - float beta_slow + float beta_slow, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -94,14 +95,15 @@ kernel void kernel_rope_norm_f32( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - float theta = theta_base * pow(freq_base, inv_ndims*i0); + float theta = theta_base * pow(freq_base, inv_ndims*iw); float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); @@ -154,7 +156,8 @@ kernel void kernel_rope_norm_f16( float ext_factor, float attn_factor, float beta_fast, - float beta_slow + float beta_slow, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -173,14 +176,15 @@ kernel void kernel_rope_norm_f16( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - float theta = theta_base * pow(freq_base, inv_ndims*i0); + float theta = theta_base * pow(freq_base, inv_ndims*iw); float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); @@ -233,7 +237,8 @@ kernel void kernel_rope_neox_f32( float ext_factor, float attn_factor, float beta_fast, - float beta_slow + float beta_slow, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -252,17 +257,18 @@ kernel void kernel_rope_neox_f32( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - const float theta = theta_base * pow(freq_base, inv_ndims*i0); + const float theta = theta_base * pow(freq_base, inv_ndims*iw); const float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); - global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00); - global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0); + global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00); + global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0); const float x0 = src[0]; const float x1 = src[n_dims/2]; @@ -312,7 +318,8 @@ kernel void kernel_rope_neox_f16( float ext_factor, float attn_factor, float beta_fast, - float beta_slow + float beta_slow, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -331,17 +338,18 @@ kernel void kernel_rope_neox_f16( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - const float theta = theta_base * pow(freq_base, inv_ndims*i0); + const float theta = theta_base * pow(freq_base, inv_ndims*iw); const float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); - global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00); - global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0); + global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00); + global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0); const float x0 = src[0]; const float x1 = src[n_dims/2]; @@ -393,7 +401,8 @@ kernel void kernel_rope_multi_f32( float beta_fast, float beta_slow, int4 sections, - int is_imrope + int is_imrope, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -414,10 +423,11 @@ kernel void kernel_rope_multi_f32( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - const int sector = (i0 / 2) % sect_dims; + const int sector = ic % sect_dims; float theta_base = 0.0f; if (is_imrope) { @@ -445,14 +455,14 @@ kernel void kernel_rope_multi_f32( } } - const float theta = theta_base * pow(freq_base, inv_ndims*i0); + const float theta = theta_base * pow(freq_base, inv_ndims*iw); const float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); - global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00); - global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0); + global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00); + global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0); const float x0 = src[0]; const float x1 = src[n_dims/2]; @@ -504,7 +514,8 @@ kernel void kernel_rope_multi_f16( float beta_fast, float beta_slow, int4 sections, - int is_imrope + int is_imrope, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -525,10 +536,11 @@ kernel void kernel_rope_multi_f16( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - const int sector = (i0 / 2) % sect_dims; + const int sector = ic % sect_dims; float theta_base = 0.0f; if (is_imrope) { @@ -556,14 +568,14 @@ kernel void kernel_rope_multi_f16( } } - const float theta = theta_base * pow(freq_base, inv_ndims*i0); + const float theta = theta_base * pow(freq_base, inv_ndims*iw); const float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); - global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00); - global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0); + global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00); + global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0); const float x0 = src[0]; const float x1 = src[n_dims/2]; diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 92c26839fce7..d31df611a4d8 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6235,8 +6235,6 @@ 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-sycl/rope.cpp b/ggml/src/ggml-sycl/rope.cpp index 9d83a1e9fa09..b6d22559d18c 100644 --- a/ggml/src/ggml-sycl/rope.cpp +++ b/ggml/src/ggml-sycl/rope.cpp @@ -41,7 +41,7 @@ template static void rope_norm(const T *x, D *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, - const int n_dims, const int32_t *pos, + const int n_dims, const int n_offs, const int32_t *pos, const float freq_scale, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float theta_scale, const float *freq_factors, @@ -78,19 +78,21 @@ static void rope_norm(const T *x, D *dst, const int ne00, const int ne01, ggml_sycl_memcpy_1<4>(dst + idst, &v); } }; - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { store_coaelsced(x[ix + 0], x[ix + 1]); return; } - const float theta_base = pos[i2] * dpct::pow(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] * dpct::pow(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, + 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]; @@ -104,7 +106,7 @@ template static void rope_neox(const T *x, D *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, - const int n_dims, const int32_t *pos, + const int n_dims, const int n_offs, const int32_t *pos, const float freq_scale, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float theta_scale, const float *freq_factors, @@ -132,35 +134,38 @@ static void rope_neox(const T *x, D *dst, const int ne00, const int ne01, idst += row_indices[i2] * set_rows_stride; } - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { dst[idst + i0 / 2 + 0] = ggml_sycl_cast(x[ix + i0 / 2 + 0]); dst[idst + i0 / 2 + 1] = ggml_sycl_cast(x[ix + i0 / 2 + 1]); return; } - const float theta_base = pos[i2] * dpct::pow(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] * dpct::pow(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, + 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_sycl_cast(x0 * cos_theta - x1 * sin_theta); - dst[idst + n_dims / 2] = ggml_sycl_cast(x0 * sin_theta + x1 * cos_theta); + dst[idst + n_offs / 2 + 0] = ggml_sycl_cast(x0 * cos_theta - x1 * sin_theta); + dst[idst + n_offs / 2 + n_dims / 2] = ggml_sycl_cast(x0 * sin_theta + x1 * cos_theta); } template static void rope_multi(const T *x, T *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, - const int n_dims, const int32_t *pos, + const int n_dims, const int n_offs, const int32_t *pos, const float freq_scale, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float theta_scale, const float *freq_factors, @@ -183,54 +188,57 @@ static void rope_multi(const T *x, T *dst, const int ne00, const int ne01, int idst = i0 / 2 + i1 * s1 + i2 * s2 + i3 * s3; const int ix = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03; - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { 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] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector % 3 == 2 && sector < 3 * sections.v[2]) { // w - theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector % 3 == 0 && sector < 3 * sections.v[0]) { // t - theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f); } else { - theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, iw / 2.0f); } } else { if (sector < sections.v[0]) { - theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector >= sections.v[0] && sector < sec_w) { - theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector >= sec_w && sector < sec_w + sections.v[2]) { - theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector >= sec_w + sections.v[2]) { - theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 3] * dpct::pow(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, + 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 @@ -293,7 +301,7 @@ static void rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, const int n_dims, - const int nr, const int32_t *pos, const float freq_scale, + const int n_offs, const int nr, const int32_t *pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float *freq_factors, const int64_t *row_indices, @@ -313,7 +321,7 @@ rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_norm( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); }); } else { @@ -323,7 +331,7 @@ rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_norm( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); }); } @@ -334,7 +342,7 @@ static void rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, const int n_dims, - const int nr, const int32_t *pos, const float freq_scale, + const int n_offs, const int nr, const int32_t *pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float *freq_factors, const int64_t *row_indices, @@ -354,7 +362,7 @@ rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_neox( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); }); } else { @@ -364,7 +372,7 @@ rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_neox( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); }); } @@ -375,7 +383,7 @@ static void rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, const int n_dims, - const int nr, const int32_t *pos, const float freq_scale, + const int n_offs, const int nr, const int32_t *pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float *freq_factors, const mrope_sections sections, @@ -395,7 +403,7 @@ rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_multi( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope); }); } else { @@ -405,7 +413,7 @@ rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_multi( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope); }); } @@ -497,6 +505,7 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, const int n_dims = ((int32_t *)dst->op_params)[1]; const int mode = ((int32_t *)dst->op_params)[2]; const int n_ctx_orig = ((int32_t *)dst->op_params)[4]; + const int n_offs = ((int32_t *)dst->op_params)[15]; mrope_sections sections; float freq_base; @@ -526,6 +535,7 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, 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; @@ -545,19 +555,19 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) { rope_neox_sycl( (const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01, - s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s02, 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); } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { rope_neox_sycl( (const float *)src0_d, (sycl::half *)dst_d, ne00, ne01, ne02, - s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + s01, s02, 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); } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) { rope_neox_sycl( (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01, - ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + ne02, s01, s02, 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); } else { @@ -568,13 +578,13 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, if (src0->type == GGML_TYPE_F32) { rope_multi_sycl((const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01, s02, s03, s1, s2, - s3, n_dims, nr, pos, freq_scale, freq_base, + s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope, stream); } else if (src0->type == GGML_TYPE_F16) { rope_multi_sycl( (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01, - ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope, stream); } else { @@ -602,19 +612,19 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) { rope_norm_sycl( (const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01, - s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s02, 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); } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { rope_norm_sycl( (const float *)src0_d, (sycl::half *)dst_d, ne00, ne01, ne02, - s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + s01, s02, 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); } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) { rope_norm_sycl( (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01, - ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + ne02, s01, s02, 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); } else { diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 4367f9a6109e..2434848a55a8 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -2714,6 +2714,7 @@ static webgpu_encoded_op ggml_webgpu_rope(webgpu_context & ctx, const int n_dims = ((int32_t *) dst->op_params)[1]; const int mode = ((int32_t *) dst->op_params)[2]; const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; + const int n_offs = ((int32_t *) dst->op_params)[15]; float freq_base; float freq_scale; @@ -2762,7 +2763,8 @@ static webgpu_encoded_op ggml_webgpu_rope(webgpu_context & ctx, (uint32_t) sections[0], (uint32_t) sections[1], (uint32_t) sections[2], - (uint32_t) sections[3] + (uint32_t) sections[3], + (uint32_t) n_offs }; std::vector entries = { ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0), @@ -4472,9 +4474,7 @@ 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: - // 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; + supports_op = op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16; break; case GGML_OP_GLU: switch (ggml_get_glu_op(op)) { diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl index 1c874e14240e..6ff53088c461 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl @@ -38,7 +38,8 @@ struct Params { sections0: u32, sections1: u32, sections2: u32, - sections3: u32 + sections3: u32, + n_offs: u32 }; @group(0) @binding(0) @@ -126,7 +127,8 @@ fn rope_yarn(theta_extrap: f32, i: u32) -> vec2 { fn pair_base(i0: u32, div_2: bool) -> u32 { if (div_2) { - return i0 / 2; + // first channel of the rotated pair: n_offs + (i0 - n_offs)/2 + return i0 / 2 + params.n_offs / 2; } else { return i0; } @@ -165,20 +167,22 @@ fn main(@builtin(global_invocation_id) gid: vec3) { let i_src_row = params.offset_src0 + i3 * params.stride_src03 + i2 * params.stride_src02 + i1 * params.stride_src01; let i_dst_row = params.offset_dst + i3 * params.stride_dst3 + i2 * params.stride_dst2 + i1 * params.stride_dst1; - if (i0 >= params.n_dims && !is_vision) { + if ((i0 < params.n_offs || i0 >= params.n_offs + params.n_dims) && !is_vision) { let i_src = i_src_row + i0; let i_dst = i_dst_row + i0; rotate(i_dst, i_dst + 1, f32(src0[i_src]), f32(src0[i_src + 1])); return; } + let iw = i0 - params.n_offs; // relative idx + var theta_base_mult: u32 = 0; - var theta_scale_pwr: u32 = i0 / 2; + var theta_scale_pwr: u32 = iw / 2; if (is_mrope) { let sect_dims = params.sections0 + params.sections1 + params.sections2 + params.sections3; let sec_w = params.sections1 + params.sections0; let sec_e = params.sections2 + sec_w; - let sector = (i0 / 2) % sect_dims; + let sector = (iw / 2) % sect_dims; if (is_imrope) { if (sector % 3 == 1 && sector < 3 * params.sections1) { theta_base_mult = 1; @@ -203,7 +207,7 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } else if (sector >= sec_e) { if (is_vision) { theta_scale_pwr = sector - sec_e; - theta_scale_pwr = (i0 / 2) % sec_e; + theta_scale_pwr = (iw / 2) % sec_e; } theta_base_mult = 3; } else if (is_vision) { @@ -212,7 +216,7 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } } let theta_base = f32(src1[params.offset_src1 + i2 + params.ne2 * theta_base_mult]) * pow(params.theta_scale, f32(theta_scale_pwr)); - let thetas = rope_yarn(theta_base/freq_factor(i0), i0); + let thetas = rope_yarn(theta_base/freq_factor(iw), iw); let i_src = i_src_row + pair_base(i0, is_neox || is_mrope || is_vision); let i_dst = i_dst_row + pair_base(i0, is_neox || is_mrope || is_vision);