natas_hw
/
openvino
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[GPU] Added RoPE support for Llama2, Llama3 and Phi-3 (#24829) 

### Tickets:
 - *[119150](https://jira.devtools.intel.com/browse/CVS-119150)*
 - *[141777](https://jira.devtools.intel.com/browse/CVS-141777)*
This commit is contained in:
Roman Lyamin 2024-06-05 02:16:48 +04:00 committed by GitHub
parent c45a46edca
commit adb590d5f4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 169 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
src/plugins/intel_gpu/include/intel_gpu/primitives/rope.hpp
View File

@ -19,14 +19,18 @@ struct rope : public primitive_base<rope> {
/// @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<input_info>& 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<rope> {
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<rope> {
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<rope> {
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<rope> {
ib >> config.rotary_ndims;
ib >> config.slice_start;
ib >> config.slice_stop;
ib >> gather_rank;
}
};
} // namespace cldnn

35
src/plugins/intel_gpu/src/graph/impls/ocl/rope.cpp
View File

@ -39,6 +39,7 @@ struct rope_impl : typed_primitive_impl_ocl<rope> {
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<rope> {
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<rope> {
return params;
}
static kernel_impl_params static_canonicalize_shapes(const kernel_impl_params& impl_params) {
const auto& primitive = impl_params.typed_desc<rope>();
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<size_t> 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);

67
src/plugins/intel_gpu/src/kernel_selector/cl_kernels/rope_ref.cl
View File

@ -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

60
src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.cpp
View File

@ -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<std::vector<Tensor::DataChannelName>> 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<std::vector<Tensor::DataChannelName>> 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;

2
src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_base.h
View File

@ -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;
};
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

1
src/plugins/intel_gpu/src/kernel_selector/kernels/rope/rope_kernel_ref.cpp
View File

@ -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);

8
src/plugins/intel_gpu/src/plugin/ops/rope.cpp
View File

@ -24,9 +24,15 @@ static void CreateRoPEOp(ProgramBuilder& p, const std::shared_ptr<op::internal::
auto inputs = p.GetInputInfo(op);
const auto& config = op->get_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);
}

3
src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp
View File

@ -816,11 +816,8 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
manager.register_pass<ov::intel_gpu::BroadcastAndPadZeroPointBuffers>(zp_pad_size);
manager.register_pass<ov::pass::RoPEFusion>();
pass_config->disable<ov::pass::RoPEFusionGPTNEOX>();
pass_config->disable<ov::pass::RoPEFusionGPTJ>();
pass_config->disable<ov::pass::RoPEFusionCosSinPreprocess>();
pass_config->disable<ov::pass::RoPEFusionIOSlicing>();
pass_config->disable<ov::pass::RoPEFusionPreprocess>();
pass_config->disable<ov::pass::RoPEShareCosSin>();
// This is supposed to be the last pass to ensure that we don't have name collisions until

5
src/plugins/intel_gpu/tests/functional/shared_tests_instances/subgraph_tests/rotary_pos_emb.cpp
View File

@ -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