diff --git a/src/plugins/intel_gpu/include/intel_gpu/primitives/rope.hpp b/src/plugins/intel_gpu/include/intel_gpu/primitives/rope.hpp index b9568c76641..51bd0e3bfdf 100644 --- a/src/plugins/intel_gpu/include/intel_gpu/primitives/rope.hpp +++ b/src/plugins/intel_gpu/include/intel_gpu/primitives/rope.hpp @@ -19,14 +19,18 @@ struct rope : public primitive_base { /// @param id This primitive id /// @param inputs Inputs primitive ids /// @param config Specific RoPE config + /// @param gather_rank Required for correct processing of gather input (if applicable) rope(const primitive_id& id, const std::vector& inputs, const RoPE::Config& config, + size_t gather_rank = 0, const padding& output_padding = padding()) : primitive_base(id, inputs, {output_padding}), - config(config) {} + config(config), + gather_rank(gather_rank) {} RoPE::Config config; + size_t gather_rank; size_t hash() const override { size_t seed = primitive::hash(); @@ -40,6 +44,7 @@ struct rope : public primitive_base { seed = hash_combine(seed, config.rotary_ndims); seed = hash_combine(seed, config.slice_start); seed = hash_combine(seed, config.slice_stop); + seed = hash_combine(seed, gather_rank); return seed; } @@ -58,7 +63,8 @@ struct rope : public primitive_base { config.is_qwen == rhs_casted.config.is_qwen && config.rotary_ndims == rhs_casted.config.rotary_ndims && config.slice_start == rhs_casted.config.slice_start && - config.slice_stop == rhs_casted.config.slice_stop; + config.slice_stop == rhs_casted.config.slice_stop && + gather_rank == rhs_casted.gather_rank; } void save(BinaryOutputBuffer& ob) const override { @@ -73,6 +79,7 @@ struct rope : public primitive_base { ob << config.rotary_ndims; ob << config.slice_start; ob << config.slice_stop; + ob << gather_rank; } void load(BinaryInputBuffer& ib) override { @@ -87,6 +94,7 @@ struct rope : public primitive_base { ib >> config.rotary_ndims; ib >> config.slice_start; ib >> config.slice_stop; + ib >> gather_rank; } }; } // namespace cldnn diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl/rope.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl/rope.cpp index 02ef732daaa..4706227e35c 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl/rope.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl/rope.cpp @@ -39,6 +39,7 @@ struct rope_impl : typed_primitive_impl_ocl { params.head_cnt = primitive->config.head_cnt; params.head_size = primitive->config.head_size; params.rotary_ndims = primitive->config.rotary_ndims; + params.gather_rank = primitive->gather_rank; params.slice_start = primitive->config.slice_start; params.slice_stop = primitive->config.slice_stop; @@ -46,8 +47,13 @@ struct rope_impl : typed_primitive_impl_ocl { params.axis = primitive->config.is_qwen || primitive->config.is_chatglm ? 2 : 3; params.num_of_inputs = primitive->config.is_chatglm || primitive->config.is_interleaved ? 2 : 3; + if (params.gather_rank > 0) { + params.num_of_inputs++; + } + params.is_qwen = primitive->config.is_qwen; params.is_chatglm = primitive->config.is_chatglm; + params.transposed_input = primitive->config.input_trans0213; for (size_t i = 1; i < impl_param.input_layouts.size(); ++i) { params.inputs.push_back(convert_data_tensor(impl_param.get_input_layout(i))); @@ -55,6 +61,35 @@ struct rope_impl : typed_primitive_impl_ocl { return params; } + static kernel_impl_params static_canonicalize_shapes(const kernel_impl_params& impl_params) { + const auto& primitive = impl_params.typed_desc(); + + if (primitive->config.is_chatglm || primitive->config.is_qwen) { + return primitive_impl::static_canonicalize_shapes(impl_params); + } else { + auto updated_impl_params = canonicalize_fused_shapes(impl_params); + + std::set canonicalize_from_begin = { 1, 2 }; + for (size_t i = 0; i < updated_impl_params.input_layouts.size(); ++i) { + auto& input_layout = updated_impl_params.input_layouts[i]; + if (canonicalize_from_begin.count(i) != 0) { + input_layout.set_partial_shape(extend_shape_to_rank_from_begin(input_layout.get_partial_shape())); + } else { + input_layout.set_partial_shape(extend_shape_to_rank_from_end(input_layout.get_partial_shape())); + } + } + + auto& output_layout = updated_impl_params.output_layouts[0]; + output_layout.set_partial_shape(extend_shape_to_rank_from_end(output_layout.get_partial_shape())); + + return updated_impl_params; + } + } + + kernel_impl_params canonicalize_shapes(const kernel_impl_params& impl_params) const override { + return static_canonicalize_shapes(impl_params); + } + void update_dispatch_data(const kernel_impl_params& impl_param) override { auto kernel_params = get_kernel_params(impl_param, true); (_kernel_data.update_dispatch_data_func)(kernel_params, _kernel_data); diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/rope_ref.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/rope_ref.cl index 1b6b9fe6549..e79b801770d 100644 --- a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/rope_ref.cl +++ b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/rope_ref.cl @@ -13,8 +13,8 @@ KERNEL(rope_ref)( { const uint p = get_global_id(0); const uint b = get_global_id(1); - const uint h = get_global_id(2) % HEAD_COUNT; - const uint rf = get_global_id(2) / HEAD_COUNT; + const uint h = (uint)get_global_id(2) % HEAD_COUNT; + const uint rf = (uint)get_global_id(2) / HEAD_COUNT; uint r = rf < HALF_ROTARY_NDIMS ? rf * 2 : 0; uint f = rf < HEAD_SIZE - ROTARY_NDIMS ? rf : 0; @@ -52,13 +52,13 @@ KERNEL(rope_ref)( OPTIONAL_SHAPE_INFO_ARG const __global INPUT0_TYPE* input, const __global INPUT1_TYPE* cos, - const __global INPUT1_TYPE* sin, + const __global INPUT2_TYPE* sin, __global OUTPUT_TYPE* output) { const uint b = get_global_id(0); const uint p = get_global_id(1); - const uint h = get_global_id(2) / HALF_ROTARY_NDIMS; - const uint r = get_global_id(2) % HALF_ROTARY_NDIMS; + const uint h = (uint)get_global_id(2) / HALF_ROTARY_NDIMS; + const uint r = (uint)get_global_id(2) % HALF_ROTARY_NDIMS; #ifdef ENABLE_SLICE uint input_idx = GET_DATA_INDEX(SLICED_INPUT0, b, p, h * HEAD_SIZE, 0); @@ -84,3 +84,60 @@ KERNEL(rope_ref)( sin[cos_sin_idx + HALF_ROTARY_NDIMS + r] * in1; } #endif + +#ifdef RotateHalf +KERNEL(rope_ref)( + OPTIONAL_SHAPE_INFO_ARG + const __global INPUT0_TYPE* input, + const __global INPUT1_TYPE* cos, + const __global INPUT2_TYPE* sin, +#ifdef ENABLE_GATHER + const __global INPUT3_TYPE* gather, +#endif + __global OUTPUT_TYPE* output) +{ + const uint b = get_global_id(0); + const uint h = get_global_id(1); + const uint p = (uint)get_global_id(2) / HALF_ROTARY_NDIMS; + const uint r = (uint)get_global_id(2) % HALF_ROTARY_NDIMS; + +#ifdef ENABLE_SLICE + uint input_idx = GET_DATA_INDEX(SLICED_INPUT0, b, h, p, 0); + + input_idx += SLICED_FROM_START * (b * INPUT0_FEATURE_NUM + h + 1) + + SLICED_FROM_END * (b * INPUT0_FEATURE_NUM + h); +#elif ENABLE_TRANSPOSE + uint input_idx = GET_DATA_INDEX(TRANSPOSED_INPUT0, b, h, p, 0); +#else + uint input_idx = INPUT0_GET_INDEX(b, h, p, 0); +#endif + + uint cos_sin_b = b < INPUT1_BATCH_NUM ? b : 0; + uint cos_sin_h = h < INPUT1_FEATURE_NUM ? h : 0; + uint cos_sin_p = p; +#ifdef ENABLE_GATHER + uint gather_b = b < INPUT3_BATCH_NUM ? b : 0; +#if GATHER_RANK == 4 + uint gather_h = h < INPUT3_FEATURE_NUM ? h : 0; + uint gather_p = p < INPUT3_SIZE_Y ? p : 0; + uint gather_idx = INPUT3_GET_INDEX(gather_b, gather_h, gather_p, 0); +#else + uint gather_p = p < INPUT3_FEATURE_NUM ? p : 0; + uint gather_idx = INPUT3_GET_INDEX(gather_b, gather_p, 0, 0); +#endif + cos_sin_p = gather[gather_idx]; +#endif + cos_sin_p = cos_sin_p < INPUT1_SIZE_Y ? cos_sin_p : 0; + uint cos_sin_idx = INPUT1_GET_INDEX(cos_sin_b, cos_sin_h, cos_sin_p, 0); + + uint output_idx = OUTPUT_GET_INDEX(b, h, p, 0); + + INPUT0_TYPE in1 = input[input_idx + r]; + INPUT0_TYPE in2 = input[input_idx + HALF_ROTARY_NDIMS + r]; + + output[output_idx + r] = cos[cos_sin_idx + r] * in1 - sin[cos_sin_idx + r] * in2; + + output[output_idx + HALF_ROTARY_NDIMS + r] = cos[cos_sin_idx + HALF_ROTARY_NDIMS + r] * in2 + + sin[cos_sin_idx + HALF_ROTARY_NDIMS + r] * in1; +} +#endif diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.cpp index afc5a322a52..f6f3414c2f7 100644 --- a/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.cpp +++ b/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.cpp @@ -22,28 +22,53 @@ JitConstants RoPEKernelBase::GetJitConstants(const rope_params& params, RoPEKern jit.AddConstant(MakeJitConstant("ENABLE_IO_COPY", true)); } - if (params.slice_stop - params.slice_start > 0) { + if (params.gather_rank > 0) { + jit.AddConstant(MakeJitConstant("ENABLE_GATHER", true)); + jit.AddConstant(MakeJitConstant("GATHER_RANK", params.gather_rank)); + } + + if (params.slice_stop > params.slice_start) { jit.AddConstant(MakeJitConstant("ENABLE_SLICE", true)); auto f = toCodeString(params.inputs[0].Feature(), 1); auto x = toCodeString(params.inputs[0].X(), 2); auto y = toCodeString(params.inputs[0].Y(), 3); - auto sliced_y = toCodeString(params.slice_stop - params.slice_start); + + auto sliced_val = toCodeString(params.slice_stop - params.slice_start); + auto sliced_x = params.axis == 3 ? sliced_val : x; + auto sliced_y = params.axis == 2 ? sliced_val : y; jit.AddConstant(MakeJitConstant("SLICED_INPUT0_X_PITCH", 1)); - jit.AddConstant(MakeJitConstant("SLICED_INPUT0_Y_PITCH", x)); - jit.AddConstant(MakeJitConstant("SLICED_INPUT0_FEATURE_PITCH", x + "*" + sliced_y)); - jit.AddConstant(MakeJitConstant("SLICED_INPUT0_BATCH_PITCH", x + "*" + sliced_y + "*" + f)); + jit.AddConstant(MakeJitConstant("SLICED_INPUT0_Y_PITCH", sliced_x)); + jit.AddConstant(MakeJitConstant("SLICED_INPUT0_FEATURE_PITCH", sliced_x + "*" + sliced_y)); + jit.AddConstant(MakeJitConstant("SLICED_INPUT0_BATCH_PITCH", sliced_x + "*" + sliced_y + "*" + f)); jit.AddConstant(MakeJitConstant("SLICED_INPUT0_OFFSET", 0)); - jit.AddConstant(MakeJitConstant("SLICED_FROM_START", toCodeString(params.slice_start))); - jit.AddConstant(MakeJitConstant("SLICED_FROM_END", "(" + y + "-" + toCodeString(params.slice_stop) + ")")); + + if (params.axis == 2) { + jit.AddConstant(MakeJitConstant("SLICED_FROM_END", "(" + y + "-" + toCodeString(params.slice_stop) + ")")); + } else if (params.axis == 3) { + jit.AddConstant(MakeJitConstant("SLICED_FROM_END", "(" + x + "-" + toCodeString(params.slice_stop) + ")")); + } else { + OPENVINO_THROW("[GPU] Invalid axis value for RoPE operation"); + } + } + + if (params.transposed_input) { + jit.AddConstant(MakeJitConstant("ENABLE_TRANSPOSE", true)); + jit.AddConstant(MakeJitConstant("TRANSPOSED_INPUT0_OFFSET", 0)); + jit.AddConstant(MakeJitConstant("TRANSPOSED_INPUT0_X_PITCH", 1)); + jit.AddConstant(MakeJitConstant("TRANSPOSED_INPUT0_Y_PITCH", "INPUT0_FEATURE_PITCH")); + jit.AddConstant(MakeJitConstant("TRANSPOSED_INPUT0_FEATURE_PITCH", "INPUT0_Y_PITCH")); + jit.AddConstant(MakeJitConstant("TRANSPOSED_INPUT0_BATCH_PITCH", "INPUT0_BATCH_PITCH")); } if (params.is_qwen) { jit.AddConstant(MakeJitConstant("QWEN", true)); } else if (params.is_chatglm) { jit.AddConstant(MakeJitConstant("CHATGLM", true)); + } else { + jit.AddConstant(MakeJitConstant("RotateHalf", true)); } return jit; @@ -54,12 +79,21 @@ RoPEKernelBase::DispatchData RoPEKernelBase::SetDefault(const rope_params& param const auto& input = params.inputs[0]; const auto& output = params.outputs[0]; - std::vector> dims_by_gws = {{ Tensor::DataChannelName::BATCH }, - { Tensor::DataChannelName::FEATURE }, - { Tensor::DataChannelName::Y, Tensor::DataChannelName::X }}; - dispatchData.gws = {input.Batch().v, - input.Feature().v, - params.head_cnt * std::max(params.rotary_ndims / 2ul, params.head_size - params.rotary_ndims)}; + std::vector> dims_by_gws; + if (params.is_chatglm || params.is_qwen) { + dims_by_gws = {{ Tensor::DataChannelName::BATCH }, { Tensor::DataChannelName::FEATURE }, + { Tensor::DataChannelName::Y, Tensor::DataChannelName::X }}; + dispatchData.gws = {input.Batch().v, + input.Feature().v, + params.head_cnt * std::max(params.rotary_ndims / 2ul, params.head_size - params.rotary_ndims)}; + } else { + dims_by_gws = {{ Tensor::DataChannelName::BATCH }, { Tensor::DataChannelName::Y }, + { Tensor::DataChannelName::FEATURE, Tensor::DataChannelName::X }}; + dispatchData.gws = {input.Batch().v, + input.Y().v, + input.Feature().v * params.rotary_ndims / 2ul}; + } + dispatchData.lws = GetOptimalLocalWorkGroupSizes(dispatchData.gws, params.engineInfo, input.GetLayout(), output.GetLayout(), dims_by_gws); return dispatchData; diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.h b/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.h index f4a92183c5d..dde691bfb43 100644 --- a/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.h +++ b/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.h @@ -20,9 +20,11 @@ struct rope_params : public base_params { size_t slice_stop; size_t axis; size_t num_of_inputs; + size_t gather_rank; bool is_qwen; bool is_chatglm; + bool transposed_input; }; //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_ref.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_ref.cpp index 27d7efeb525..5ec125ef6f0 100644 --- a/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_ref.cpp +++ b/src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_ref.cpp @@ -9,6 +9,7 @@ namespace kernel_selector { ParamsKey RoPEKernelRef::GetSupportedKey() const { ParamsKey k; + k.EnableInputDataType(Datatype::INT32); k.EnableInputDataType(Datatype::F16); k.EnableInputDataType(Datatype::F32); k.EnableOutputDataType(Datatype::F16); diff --git a/src/plugins/intel_gpu/src/plugin/ops/rope.cpp b/src/plugins/intel_gpu/src/plugin/ops/rope.cpp index 2b299890fe9..321342b3395 100644 --- a/src/plugins/intel_gpu/src/plugin/ops/rope.cpp +++ b/src/plugins/intel_gpu/src/plugin/ops/rope.cpp @@ -24,9 +24,15 @@ static void CreateRoPEOp(ProgramBuilder& p, const std::shared_ptrget_config(); + size_t gather_rank = 0; + if (config.gather_position_arg_id > 0) { + gather_rank = op->get_input_partial_shape(config.gather_position_arg_id).size(); + } + auto rope = cldnn::rope(layer_type_name_ID(op), inputs, - config); + config, + gather_rank); p.add_primitive(*op, rope); } diff --git a/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp b/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp index 581fb42a109..faac47298f5 100644 --- a/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp +++ b/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp @@ -816,11 +816,8 @@ void TransformationsPipeline::apply(std::shared_ptr func) { manager.register_pass(zp_pad_size); manager.register_pass(); - pass_config->disable(); pass_config->disable(); - pass_config->disable(); pass_config->disable(); - pass_config->disable(); pass_config->disable(); // This is supposed to be the last pass to ensure that we don't have name collisions until diff --git a/src/plugins/intel_gpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp b/src/plugins/intel_gpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp index bed957cc35f..5bfcadb10c8 100644 --- a/src/plugins/intel_gpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp +++ b/src/plugins/intel_gpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp @@ -18,5 +18,10 @@ INSTANTIATE_TEST_SUITE_P(smoke_RoPETestQwen7b, ::testing::Values(ov::test::utils::DEVICE_GPU)), RoPETestQwen7b::getTestCaseName); +INSTANTIATE_TEST_SUITE_P(smoke_RoPETestLlama2, + RoPETestLlama2, + ::testing::Values(ov::test::utils::DEVICE_GPU), + RoPETestLlama2::getTestCaseName); + } // namespace test } // namespace ov