diff --git a/src/plugins/intel_gpu/include/intel_gpu/op/indirect_sdpa.hpp b/src/plugins/intel_gpu/include/intel_gpu/op/indirect_sdpa.hpp index 18c41cf2c12..b4d34a3975a 100644 --- a/src/plugins/intel_gpu/include/intel_gpu/op/indirect_sdpa.hpp +++ b/src/plugins/intel_gpu/include/intel_gpu/op/indirect_sdpa.hpp @@ -19,36 +19,7 @@ public: IndirectSDPA() = default; - IndirectSDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, - const ov::Output& beam_table, - const bool is_causal, - const int64_t indirect_axis, - const std::vector& order_q, - const std::vector& order_k, - const std::vector& order_v, - const std::vector& order_out, - const ov::element::Type output_type = ov::element::undefined); - - IndirectSDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, - const ov::Output& attn_mask, - const ov::Output& beam_table, - const bool is_causal, - const int64_t indirect_axis, - const std::vector& order_q, - const std::vector& order_k, - const std::vector& order_v, - const std::vector& order_out, - const ov::element::Type output_type = ov::element::undefined); - - IndirectSDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, - const ov::Output& attn_mask, - const ov::Output& scale, + IndirectSDPA(const OutputVector& data_inputs, const ov::Output& beam_table, const bool is_causal, const int64_t indirect_axis, diff --git a/src/plugins/intel_gpu/include/intel_gpu/op/sdpa.hpp b/src/plugins/intel_gpu/include/intel_gpu/op/sdpa.hpp index 45416b4e538..9f5d4dad16e 100644 --- a/src/plugins/intel_gpu/include/intel_gpu/op/sdpa.hpp +++ b/src/plugins/intel_gpu/include/intel_gpu/op/sdpa.hpp @@ -19,37 +19,12 @@ public: SDPA() = default; - SDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, + SDPA(const OutputVector& inputs, + const bool is_causal, const std::vector& order_q, const std::vector& order_k, const std::vector& order_v, const std::vector& order_out, - const bool is_causal, - const ov::element::Type output_type = ov::element::undefined); - - SDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, - const ov::Output& attn_mask, - const std::vector& order_q, - const std::vector& order_k, - const std::vector& order_v, - const std::vector& order_out, - const bool is_causal, - const ov::element::Type output_type = ov::element::undefined); - - SDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, - const ov::Output& attn_mask, - const ov::Output& scale, - const std::vector& order_q, - const std::vector& order_k, - const std::vector& order_v, - const std::vector& order_out, - const bool is_causal, const ov::element::Type output_type = ov::element::undefined); bool visit_attributes(ov::AttributeVisitor &visitor) override; @@ -73,11 +48,11 @@ public: } protected: + bool m_is_causal; std::vector m_order_q; std::vector m_order_k; std::vector m_order_v; std::vector m_order_out; - bool m_is_causal; ov::element::Type m_output_type; }; diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl/scaled_dot_product_attention.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl/scaled_dot_product_attention.cpp index 364c9418f10..681237b2bcd 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl/scaled_dot_product_attention.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl/scaled_dot_product_attention.cpp @@ -72,7 +72,6 @@ protected: } static size_t get_beam_table_id(std::shared_ptr primitive) { - GPU_DEBUG_TRACE << "get_beam_table_id " << primitive->input_size() - 1 << "\n"; return primitive->input_size() - 1; } diff --git a/src/plugins/intel_gpu/src/graph/scaled_dot_product_attention.cpp b/src/plugins/intel_gpu/src/graph/scaled_dot_product_attention.cpp index 42e5aeb9f13..1e9cd9a2230 100644 --- a/src/plugins/intel_gpu/src/graph/scaled_dot_product_attention.cpp +++ b/src/plugins/intel_gpu/src/graph/scaled_dot_product_attention.cpp @@ -20,7 +20,27 @@ layout scaled_dot_product_attention_inst::calc_output_layout(scaled_dot_product_ kernel_impl_params const& impl_param) { auto desc = impl_param.typed_desc(); - return impl_param.get_input_layout(0); + auto transpose_shape = [&desc](const ov::PartialShape& shape, const std::vector& order) { + if (desc->input_q_transpose_order.empty()) + return shape; + + auto shape_transposed = ov::PartialShape(shape); + auto rank_diff = shape.size() - order.size(); + for (size_t i = 0; i < order.size(); i++) { + size_t idx = static_cast(order[i]); + shape_transposed[i + rank_diff] = shape[idx + rank_diff]; + } + + return shape_transposed; + }; + + auto input0_layout = impl_param.get_input_layout(0); + auto default_out_dt = data_type_traits::is_floating_point(input0_layout.data_type) ? input0_layout.data_type : data_types::f32; + auto output_type = desc->output_data_types[0].value_or(default_out_dt); + auto output_format = input0_layout.format; + auto output_shape = transpose_shape(input0_layout.get_partial_shape(), desc->input_q_transpose_order); // output shape matches with Q input shape + + return { layout{output_shape, output_type, output_format, desc->output_paddings[0]} }; } template diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_opt.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_opt.cl index 0f355fff1af..635aa4d796d 100644 --- a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_opt.cl +++ b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_opt.cl @@ -20,16 +20,16 @@ inline uint FUNC(get_input0_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w, uint z, uint y, uint x) { #if INPUT0_SIMPLE - return GET_DATA_INDEX_6D_SAFE(INPUT0, b, f, w, z, y, x); + return GET_DATA_INDEX_6D(INPUT0, b, f, w, z, y, x); #else #if INPUT0_DIMS == 4 - return INPUT0_GET_INDEX_SAFE(b, f, y, x); + return INPUT0_GET_INDEX(b, f, y, x); #elif INPUT0_DIMS == 5 - return INPUT0_GET_INDEX_SAFE(b, f, z, y, x); + return INPUT0_GET_INDEX(b, f, z, y, x); #elif INPUT0_DIMS == 6 - return INPUT0_GET_INDEX_SAFE(b, f, w, z, y, x); + return INPUT0_GET_INDEX(b, f, w, z, y, x); #else -# error sdpa_ref.cl : Unsupported input 0 format +# error sdpa_opt.cl : Unsupported input 0 format #endif #endif } @@ -47,16 +47,16 @@ inline uint FUNC(get_input1_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, ui DO_BROADCAST_KEY_VALUE; #endif #if INPUT1_SIMPLE - return GET_DATA_INDEX_6D_SAFE(INPUT1, b, f, w, z, y, x); + return GET_DATA_INDEX_6D(INPUT1, b, f, w, z, y, x); #else #if INPUT1_DIMS == 4 - return INPUT1_GET_INDEX_SAFE(b, f, y, x); + return INPUT1_GET_INDEX(b, f, y, x); #elif INPUT1_DIMS == 5 - return INPUT1_GET_INDEX_SAFE(b, f, z, y, x); + return INPUT1_GET_INDEX(b, f, z, y, x); #elif INPUT1_DIMS == 6 - return INPUT1_GET_INDEX_SAFE(b, f, w, z, y, x); + return INPUT1_GET_INDEX(b, f, w, z, y, x); #else -# error sdpa_ref.cl : Unsupported input 1 format +# error sdpa_opt.cl : Unsupported input 1 format #endif #endif } @@ -77,13 +77,13 @@ inline uint FUNC(get_input2_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, ui return GET_DATA_INDEX_6D_SAFE(INPUT2, b, f, w, z, y, x); #else #if INPUT2_DIMS == 4 - return INPUT2_GET_INDEX_SAFE(b, f, y, x); + return INPUT2_GET_INDEX(b, f, y, x); #elif INPUT2_DIMS == 5 - return INPUT2_GET_INDEX_SAFE(b, f, z, y, x); + return INPUT2_GET_INDEX(b, f, z, y, x); #elif INPUT2_DIMS == 6 - return INPUT2_GET_INDEX_SAFE(b, f, w, z, y, x); + return INPUT2_GET_INDEX(b, f, w, z, y, x); #else -# error sdpa_ref.cl : Unsupported input 1 format +# error sdpa_opt.cl : Unsupported input 1 format #endif #endif } @@ -101,7 +101,7 @@ inline uint FUNC(get_bt_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w #if BEAM_TABLE_SIMPLE return GET_DATA_INDEX_6D_SAFE(BEAM_TABLE, b, f, w, z, y, x); #else -# error sdpa_ref.cl : Unsupported beam table format +# error sdpa_opt.cl : Unsupported beam table format #endif } @@ -466,8 +466,8 @@ KERNEL(sdpa_opt)( for (uint seq_len = 0; seq_len < partition_seq_len / SUBGROUP_SIZE; seq_len++) { #ifdef BEAM_TABLE_TYPE - uint b_idx = beam_table[FUNC_CALL(get_bt_index_value)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, (head_size_idx / SUBGROUP_SIZE) * SUBGROUP_SIZE)]; - uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, (head_size_idx / SUBGROUP_SIZE) * SUBGROUP_SIZE); + uint b_idx = beam_table[FUNC_CALL(get_bt_index_value)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, sgid * SUBGROUP_SIZE)]; + uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, sgid * SUBGROUP_SIZE); #else #ifdef INPUT2_DIMS_ORDER uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE), head_size_idx); @@ -715,7 +715,7 @@ KERNEL(sdpa_opt)( #endif #if INPUT0_TYPE_SIZE == 2 /* Adding this clamp improves performance for some reason */ - acc[i] = SOFTMAX_ACCUMULATOR_MIN_FUNC(SOFTMAX_ACCUMULATOR_MAX_FUNC(acc[i], INPUT0_VAL_MIN), INPUT0_VAL_MAX); + acc[i] = INPUT0_MIN_FUNC(INPUT0_MAX_FUNC(acc[i], INPUT0_VAL_MIN), INPUT0_VAL_MAX); #endif if (seq_len + i >= partition_seq_len) { acc[i] = INPUT0_VAL_MIN; @@ -749,7 +749,7 @@ KERNEL(sdpa_opt)( #define QK_ITERS_END QK_MAX_NUMS_PER_SG #endif - OUTPUT_TYPE qk_max[QK_MAX_NUMS_PER_SG]; + SOFTMAX_ACCUMULATOR_TYPE qk_max[QK_MAX_NUMS_PER_SG]; for (uint i = 0; i < QK_MAX_NUMS_PER_SG; i++) qk_max[i] = SOFTMAX_ACCUMULATOR_VAL_MIN; diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_opt.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_opt.cpp index 359e8696cbe..23dc964eb18 100644 --- a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_opt.cpp +++ b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_opt.cpp @@ -23,12 +23,15 @@ static size_t get_target_seq_len_block_size() { return block_size; } - static size_t get_seq_len_partition_size() { const size_t seq_len = 256; return seq_len; } +static Datatype get_softmax_acc_type() { + return Datatype::F32; +} + static std::string GetKernelName(std::string base_name, KernelsTypes type, bool is_indirect) { auto kernel_name = base_name; if (is_indirect) @@ -80,7 +83,7 @@ bool SDPAKernelOpt::Validate(const Params& p) const { JitConstants SDPAKernelOpt::GetJitConstants(const sdpa_params& params, size_t kernel_idx) const { auto jit = SDPAKernelBase::GetJitConstants(params); - const auto softmax_acc_dt = params.inputs[0].GetDType(); + const auto softmax_acc_dt = get_softmax_acc_type(); jit.Merge(MakeTypeJitConstants(softmax_acc_dt, "SOFTMAX_ACCUMULATOR")); const auto& config = params.conf; @@ -232,7 +235,7 @@ void SDPAKernelOpt::GetUpdateDispatchDataFunc(KernelData& kd) const { auto num_of_partitions = CeilDiv(source_seq_len, get_seq_len_partition_size()); - auto buf_dt_size = output.ElementSize(); + auto buf_dt_size = BytesPerElement(get_softmax_acc_type()); auto buf_elements_count = (num_of_partitions == 1) ? 1 : output.LogicalSize() / head_size * num_of_partitions; auto buf_size = buf_elements_count * buf_dt_size; @@ -241,26 +244,26 @@ void SDPAKernelOpt::GetUpdateDispatchDataFunc(KernelData& kd) const { auto tmp_out_size = tmp_out_elements_count * tmp_out_dt_size; auto dispatch_data1 = SetDefault(prim_params, 0); - kernel_data.kernels[0].params.workGroups.global = dispatch_data1.gws; - kernel_data.kernels[0].params.workGroups.local = dispatch_data1.lws; - kernel_data.kernels[0].skip_execution = target_seq_len > 1; + kernel_data.kernels[KernelsTypes::SINGLE_TOKEN].params.workGroups.global = dispatch_data1.gws; + kernel_data.kernels[KernelsTypes::SINGLE_TOKEN].params.workGroups.local = dispatch_data1.lws; + kernel_data.kernels[KernelsTypes::SINGLE_TOKEN].skip_execution = target_seq_len > 1; auto dispatch_data2 = SetDefault(prim_params, 1); - kernel_data.kernels[1].params.workGroups.global = dispatch_data2.gws; - kernel_data.kernels[1].params.workGroups.local = dispatch_data2.lws; - kernel_data.kernels[1].skip_execution = target_seq_len == 1; + kernel_data.kernels[KernelsTypes::MULTI_TOKENS].params.workGroups.global = dispatch_data2.gws; + kernel_data.kernels[KernelsTypes::MULTI_TOKENS].params.workGroups.local = dispatch_data2.lws; + kernel_data.kernels[KernelsTypes::MULTI_TOKENS].skip_execution = target_seq_len == 1; ScalarDescriptor num_of_partitions_scalar; num_of_partitions_scalar.t = ScalarDescriptor::Types::UINT32; num_of_partitions_scalar.v.u32 = static_cast(num_of_partitions); auto dispatch_data3 = SetDefault(prim_params, 2); - kernel_data.kernels[2].params.workGroups.global = dispatch_data3.gws; - kernel_data.kernels[2].params.workGroups.local = dispatch_data3.lws; - kernel_data.kernels[2].skip_execution = num_of_partitions == 1; + kernel_data.kernels[KernelsTypes::FINALIZATION].params.workGroups.global = dispatch_data3.gws; + kernel_data.kernels[KernelsTypes::FINALIZATION].params.workGroups.local = dispatch_data3.lws; + kernel_data.kernels[KernelsTypes::FINALIZATION].skip_execution = num_of_partitions == 1; - kernel_data.kernels[2].params.scalars.clear(); - kernel_data.kernels[2].params.scalars.push_back(num_of_partitions_scalar); + kernel_data.kernels[KernelsTypes::FINALIZATION].params.scalars.clear(); + kernel_data.kernels[KernelsTypes::FINALIZATION].params.scalars.push_back(num_of_partitions_scalar); kernel_data.internalBufferSizes.clear(); kernel_data.internalBufferSizes.push_back(buf_size); diff --git a/src/plugins/intel_gpu/src/plugin/transformations/indirect_kv_cache.cpp b/src/plugins/intel_gpu/src/plugin/transformations/indirect_kv_cache.cpp index e6f58aaa25f..429fa671fdd 100644 --- a/src/plugins/intel_gpu/src/plugin/transformations/indirect_kv_cache.cpp +++ b/src/plugins/intel_gpu/src/plugin/transformations/indirect_kv_cache.cpp @@ -190,45 +190,27 @@ IndirectSDPAOpt::IndirectSDPAOpt() { auto order_out = sdpa->get_output_transpose_order(); auto is_causal = sdpa->get_causal(); - std::shared_ptr indirect_sdpa; - if (pattern_map.find(sdpa_without_attn_mask_m) != pattern_map.end()) { - indirect_sdpa = std::make_shared(sdpa->get_input_node_shared_ptr(0), - sdpa->get_input_node_shared_ptr(1), - sdpa->get_input_node_shared_ptr(2), - indirect_kv_cache_0->output(1), // beam table - is_causal, - gather_axis_1, - order_in0, - order_in1, - order_in2, - order_out); - } else if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) { - indirect_sdpa = std::make_shared(sdpa->get_input_node_shared_ptr(0), - sdpa->get_input_node_shared_ptr(1), - sdpa->get_input_node_shared_ptr(2), - sdpa->get_input_node_shared_ptr(3), - indirect_kv_cache_0->output(1), // beam table - is_causal, - gather_axis_1, - order_in0, - order_in1, - order_in2, - order_out); + OutputVector data_inputs; + data_inputs.push_back(sdpa->get_input_node_shared_ptr(0)); // Q + data_inputs.push_back(sdpa->get_input_node_shared_ptr(1)); // K + data_inputs.push_back(sdpa->get_input_node_shared_ptr(2)); // V + + if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) { + data_inputs.push_back(sdpa->get_input_source_output(3)); } else if (pattern_map.find(sdpa_with_attn_mask_and_scale_m) != pattern_map.end()) { - indirect_sdpa = std::make_shared(sdpa->get_input_node_shared_ptr(0), - sdpa->get_input_node_shared_ptr(1), - sdpa->get_input_node_shared_ptr(2), - sdpa->get_input_node_shared_ptr(3), - sdpa->get_input_node_shared_ptr(4), - indirect_kv_cache_0->output(1), // beam table - is_causal, - gather_axis_1, - order_in0, - order_in1, - order_in2, - order_out); + data_inputs.push_back(sdpa->get_input_source_output(3)); + data_inputs.push_back(sdpa->get_input_source_output(4)); } + auto indirect_sdpa = std::make_shared(data_inputs, + indirect_kv_cache_0->output(1), // beam table + is_causal, + gather_axis_1, + order_in0, + order_in1, + order_in2, + order_out); + OPENVINO_ASSERT(indirect_sdpa != nullptr); indirect_sdpa->set_friendly_name(sdpa->get_friendly_name()); diff --git a/src/plugins/intel_gpu/src/plugin/transformations/op/indirect_sdpa.cpp b/src/plugins/intel_gpu/src/plugin/transformations/op/indirect_sdpa.cpp index 9b36bfcb3d3..681c88119ef 100644 --- a/src/plugins/intel_gpu/src/plugin/transformations/op/indirect_sdpa.cpp +++ b/src/plugins/intel_gpu/src/plugin/transformations/op/indirect_sdpa.cpp @@ -9,9 +9,7 @@ namespace ov { namespace intel_gpu { namespace op { -IndirectSDPA::IndirectSDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, +IndirectSDPA::IndirectSDPA(const OutputVector& data_inputs, const ov::Output& beam_table, const bool is_causal, const int64_t indirect_axis, @@ -20,62 +18,28 @@ IndirectSDPA::IndirectSDPA(const ov::Output& Q, const std::vector& order_v, const std::vector& order_out, const ov::element::Type output_type) - : ov::intel_gpu::op::SDPA(Q, K, V, order_q, order_k, order_v, order_out, is_causal, output_type) + : ov::intel_gpu::op::SDPA(data_inputs, is_causal, order_q, order_k, order_v, order_out, output_type) , m_indirect_axis(indirect_axis) { - set_argument(3, beam_table); - validate_and_infer_types(); -} - -IndirectSDPA::IndirectSDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, - const ov::Output& attn_mask, - const ov::Output& beam_table, - const bool is_causal, - const int64_t indirect_axis, - const std::vector& order_q, - const std::vector& order_k, - const std::vector& order_v, - const std::vector& order_out, - const ov::element::Type output_type) - : ov::intel_gpu::op::SDPA(Q, K, V, attn_mask, order_q, order_k, order_v, order_out, is_causal, output_type) - , m_indirect_axis(indirect_axis) { - set_argument(4, beam_table); - validate_and_infer_types(); -} - -IndirectSDPA::IndirectSDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, - const ov::Output& attn_mask, - const ov::Output& scale, - const ov::Output& beam_table, - const bool is_causal, - const int64_t indirect_axis, - const std::vector& order_q, - const std::vector& order_k, - const std::vector& order_v, - const std::vector& order_out, - const ov::element::Type output_type) - : ov::intel_gpu::op::SDPA(Q, K, V, attn_mask, scale, order_q, order_k, order_v, order_out, is_causal, output_type) - , m_indirect_axis(indirect_axis) { - set_argument(5, beam_table); + auto beam_table_idx = data_inputs.size(); + set_argument(beam_table_idx, beam_table); validate_and_infer_types(); } std::shared_ptr IndirectSDPA::clone_with_new_inputs(const ov::OutputVector& new_args) const { check_new_args_count(this, new_args); - if (new_args.size() == 4) { - return std::make_shared(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3), - m_is_causal, m_indirect_axis, m_order_q, m_order_k, m_order_v, m_order_out, m_output_type); - } else if (new_args.size() == 5) { - return std::make_shared(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3), new_args.at(4), - m_is_causal, m_indirect_axis, m_order_q, m_order_k, m_order_v, m_order_out, m_output_type); - } else { - return std::make_shared(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3), new_args.at(4), new_args.at(5), - m_is_causal, m_indirect_axis, m_order_q, m_order_k, m_order_v, m_order_out, m_output_type); - } + // Exclude beam_table input + OutputVector data_inputs(new_args.begin(), new_args.end() - 1); + + return std::make_shared(data_inputs, + new_args.back(), + m_is_causal, + m_indirect_axis, + m_order_q, + m_order_k, + m_order_v, + m_order_out, + m_output_type); } void IndirectSDPA::validate_and_infer_types() { diff --git a/src/plugins/intel_gpu/src/plugin/transformations/op/sdpa.cpp b/src/plugins/intel_gpu/src/plugin/transformations/op/sdpa.cpp index 5965cb5e991..1006ad280df 100644 --- a/src/plugins/intel_gpu/src/plugin/transformations/op/sdpa.cpp +++ b/src/plugins/intel_gpu/src/plugin/transformations/op/sdpa.cpp @@ -14,65 +14,20 @@ namespace ov { namespace intel_gpu { namespace op { -SDPA::SDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, +SDPA::SDPA(const OutputVector& inputs, + const bool is_causal, const std::vector& order_q, const std::vector& order_k, const std::vector& order_v, const std::vector& order_out, - const bool is_causal, const ov::element::Type output_type) - : m_order_q(order_q) + : m_is_causal(is_causal) + , m_order_q(order_q) , m_order_k(order_k) , m_order_v(order_v) , m_order_out(order_out) - , m_is_causal(is_causal) , m_output_type(output_type) { - set_arguments({Q, K, V}); - set_causal(is_causal); - validate_and_infer_types(); -} - -SDPA::SDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, - const ov::Output& attn_mask, - const std::vector& order_q, - const std::vector& order_k, - const std::vector& order_v, - const std::vector& order_out, - const bool is_causal, - const ov::element::Type output_type) - : m_order_q(order_q) - , m_order_k(order_k) - , m_order_v(order_v) - , m_order_out(order_out) - , m_is_causal(is_causal) - , m_output_type(output_type) { - set_arguments({Q, K, V, attn_mask}); - set_causal(is_causal); - validate_and_infer_types(); -} - -SDPA::SDPA(const ov::Output& Q, - const ov::Output& K, - const ov::Output& V, - const ov::Output& attn_mask, - const ov::Output& scale, - const std::vector& order_q, - const std::vector& order_k, - const std::vector& order_v, - const std::vector& order_out, - const bool is_causal, - const ov::element::Type output_type) - : m_order_q(order_q) - , m_order_k(order_k) - , m_order_v(order_v) - , m_order_out(order_out) - , m_is_causal(is_causal) - , m_output_type(output_type) { - set_arguments({Q, K, V, attn_mask, scale}); + set_arguments(inputs); set_causal(is_causal); validate_and_infer_types(); } @@ -80,16 +35,13 @@ SDPA::SDPA(const ov::Output& Q, std::shared_ptr SDPA::clone_with_new_inputs(const ov::OutputVector& new_args) const { check_new_args_count(this, new_args); - if (new_args.size() == 3) { - return std::make_shared(new_args.at(0), new_args.at(1), new_args.at(2), - m_order_q, m_order_k, m_order_v, m_order_out, m_is_causal, m_output_type); - } else if (new_args.size() == 4) { - return std::make_shared(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3), - m_order_q, m_order_k, m_order_v, m_order_out, m_is_causal, m_output_type); - } else { - return std::make_shared(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3), new_args.at(4), - m_order_q, m_order_k, m_order_v, m_order_out, m_is_causal, m_output_type); - } + return std::make_shared(new_args, + m_is_causal, + m_order_q, + m_order_k, + m_order_v, + m_order_out, + m_output_type); } void SDPA::validate_and_infer_types() { diff --git a/src/plugins/intel_gpu/src/plugin/transformations/transpose_fusion.cpp b/src/plugins/intel_gpu/src/plugin/transformations/transpose_fusion.cpp index f418376d145..3a631597857 100644 --- a/src/plugins/intel_gpu/src/plugin/transformations/transpose_fusion.cpp +++ b/src/plugins/intel_gpu/src/plugin/transformations/transpose_fusion.cpp @@ -134,18 +134,20 @@ TransposeSDPAMatcher::TransposeSDPAMatcher() { auto input_k = ov::Output(pattern_map.at(input_k_m).get_node_shared_ptr(), input_k_output_idx); auto input_v = ov::Output(pattern_map.at(input_v_m).get_node_shared_ptr(), input_v_output_idx); - std::shared_ptr sdpa_new; - if (pattern_map.find(sdpa_without_attn_mask_m) != pattern_map.end()) { - sdpa_new = std::make_shared(input_q, input_k, input_v, order_q, order_k, order_v, order_output, sdpa->get_causal()); - } else if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) { - auto attn_mask = sdpa->get_input_source_output(3); - sdpa_new = std::make_shared(input_q, input_k, input_v, attn_mask, order_q, order_k, order_v, order_output, sdpa->get_causal()); + OutputVector inputs; + inputs.push_back(ov::Output(pattern_map.at(input_q_m).get_node_shared_ptr(), input_q_output_idx)); + inputs.push_back(ov::Output(pattern_map.at(input_k_m).get_node_shared_ptr(), input_k_output_idx)); + inputs.push_back(ov::Output(pattern_map.at(input_v_m).get_node_shared_ptr(), input_v_output_idx)); + + if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) { + inputs.push_back(sdpa->get_input_source_output(3)); } else if (pattern_map.find(sdpa_with_attn_mask_and_scale_m) != pattern_map.end()) { - auto attn_mask = sdpa->get_input_source_output(3); - auto scale = sdpa->get_input_source_output(4); - sdpa_new = std::make_shared(input_q, input_k, input_v, attn_mask, scale, order_q, order_k, order_v, order_output, sdpa->get_causal()); + inputs.push_back(sdpa->get_input_source_output(3)); + inputs.push_back(sdpa->get_input_source_output(4)); } + auto sdpa_new = std::make_shared(inputs, sdpa->get_causal(), order_q, order_k, order_v, order_output); + sdpa_new->set_friendly_name(sdpa->get_friendly_name()); ov::copy_runtime_info(m.get_matched_nodes(), sdpa_new); ov::replace_node(sdpa, sdpa_new); diff --git a/src/plugins/intel_gpu/src/plugin/transformations/unsqueeze_broadcast_reshape_sdpa_fusion.cpp b/src/plugins/intel_gpu/src/plugin/transformations/unsqueeze_broadcast_reshape_sdpa_fusion.cpp index 3fdb3794585..d525792ccd8 100644 --- a/src/plugins/intel_gpu/src/plugin/transformations/unsqueeze_broadcast_reshape_sdpa_fusion.cpp +++ b/src/plugins/intel_gpu/src/plugin/transformations/unsqueeze_broadcast_reshape_sdpa_fusion.cpp @@ -97,27 +97,25 @@ UnsqueezeBroadcastReshapeSDPAFusion::UnsqueezeBroadcastReshapeSDPAFusion() { return false; } - auto input_a = pattern_map.at(input_a_m).get_node_shared_ptr(); - auto input_b = pattern_map.at(input_b_m).get_node_shared_ptr(); - auto input_c = pattern_map.at(input_c_m).get_node_shared_ptr(); + OutputVector data_inputs; + data_inputs.push_back(pattern_map.at(input_a_m).get_node_shared_ptr()); // Q + data_inputs.push_back(pattern_map.at(input_b_m).get_node_shared_ptr()); // K + data_inputs.push_back(pattern_map.at(input_c_m).get_node_shared_ptr()); // V auto sdpa = std::dynamic_pointer_cast(m.get_match_root()); + if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) { + data_inputs.push_back(sdpa->get_input_source_output(3)); // attn_mask + } else if (pattern_map.find(sdpa_with_attn_mask_and_scale_m) != pattern_map.end()) { + data_inputs.push_back(sdpa->get_input_source_output(3)); // attn_mask + data_inputs.push_back(sdpa->get_input_source_output(4)); // scale + } + auto order_a = sdpa->get_input0_transpose_order(); auto order_b = sdpa->get_input1_transpose_order(); auto order_c = sdpa->get_input2_transpose_order(); auto order_d = sdpa->get_output_transpose_order(); - std::shared_ptr sdpa_new; - if (pattern_map.find(sdpa_without_attn_mask_m) != pattern_map.end()) { - sdpa_new = std::make_shared(input_a, input_b, input_c, order_a, order_b, order_c, order_d, sdpa->get_causal()); - } else if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) { - auto attn_mask = sdpa->get_input_source_output(3); - sdpa_new = std::make_shared(input_a, input_b, input_c, attn_mask, order_a, order_b, order_c, order_d, sdpa->get_causal()); - } else if (pattern_map.find(sdpa_with_attn_mask_and_scale_m) != pattern_map.end()) { - auto attn_mask = sdpa->get_input_source_output(3); - auto scale = sdpa->get_input_source_output(4); - sdpa_new = std::make_shared(input_a, input_b, input_c, attn_mask, scale, order_a, order_b, order_c, order_d, sdpa->get_causal()); - } + auto sdpa_new = std::make_shared(data_inputs, sdpa->get_causal(), order_a, order_b, order_c, order_d); sdpa_new->set_friendly_name(sdpa->get_friendly_name()); ov::copy_runtime_info(m.get_matched_nodes(), sdpa_new); diff --git a/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/scaled_dot_product_attention.cpp b/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/scaled_dot_product_attention.cpp index 15203e9c5f2..5abfe681ac9 100644 --- a/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/scaled_dot_product_attention.cpp +++ b/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/scaled_dot_product_attention.cpp @@ -25,8 +25,7 @@ typedef std::tuple, // shape bool, // is_causal bool, // has_attn - bool, // has_scale - std::string // targetDevice + bool // has_scale > ScaledAttnGPUTestParams; class ScaledAttnLayerGPUTest : public testing::WithParamInterface, @@ -48,8 +47,7 @@ std::string ScaledAttnLayerGPUTest::getTestCaseName(const testing::TestParamInfo bool is_causal; bool has_attn; bool has_scale; - std::string targetDevice; - std::tie(inType, inputShapes, is_causal, has_attn, has_scale, targetDevice) = obj.param; + std::tie(inType, inputShapes, is_causal, has_attn, has_scale) = obj.param; std::ostringstream result; result << "netPRC=" << inType << "_"; @@ -67,7 +65,6 @@ std::string ScaledAttnLayerGPUTest::getTestCaseName(const testing::TestParamInfo result << "is_causal=" << is_causal << "_"; result << "has_attn=" << has_attn << "_"; result << "has_scale=" << has_scale << "_"; - result << "trgDev=" << targetDevice; return result.str(); } @@ -75,7 +72,10 @@ std::string ScaledAttnLayerGPUTest::getTestCaseName(const testing::TestParamInfo void ScaledAttnLayerGPUTest::SetUp() { ov::element::Type inType; std::vector inputShapes; - std::tie(inType, inputShapes, is_causal, has_attn, has_scale, targetDevice) = this->GetParam(); + + targetDevice = ov::test::utils::DEVICE_GPU; + + std::tie(inType, inputShapes, is_causal, has_attn, has_scale) = this->GetParam(); init_input_shapes(inputShapes); ov::ParameterVector inputParams; @@ -166,8 +166,7 @@ TEST_P(ScaledAttnLayerGPUTest, CompareWithRefs) { bool is_causal; bool has_attn; bool has_scale; - std::string targetDevice; - std::tie(inType, inputShapes, is_causal, has_attn, has_scale, targetDevice) = this->GetParam(); + std::tie(inType, inputShapes, is_causal, has_attn, has_scale) = this->GetParam(); run(); } @@ -190,15 +189,15 @@ const std::vector> shapes{ { // q shape {ov::test::InputShape{ov::PartialShape{-1, 5, -1, 128}, - {ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 384, 128}}} + {ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 387, 128}}} }, // kv shape {ov::test::InputShape{ov::PartialShape{-1, 5, -1, 128}, - {ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 384, 128}}} + {ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 387, 128}}} }, // attn shape: [B, 1, -1, L0+L1] {ov::test::InputShape{ov::PartialShape{-1, 1, -1, -1}, - {ov::Shape{1, 1, 100, 100}, ov::Shape{1, 1, 1, 1}, ov::Shape{2, 1, 384, 384}}} + {ov::Shape{1, 1, 100, 100}, ov::Shape{1, 1, 1, 1}, ov::Shape{2, 1, 387, 387}}} }, }, // heads number of kv is 1, attn mask: [B, H, L1, L0+L1] @@ -222,8 +221,7 @@ const auto params = testing::Combine(testing::Values(ov::element::f16 /*, ov::el testing::ValuesIn(shapes), testing::Values(true, false), testing::Values(true, false), - testing::Values(true, false), - testing::Values(ov::test::utils::DEVICE_GPU)); + testing::Values(true, false)); INSTANTIATE_TEST_SUITE_P(smoke_ScaledAttn_GPU, ScaledAttnLayerGPUTest, diff --git a/src/plugins/intel_gpu/tests/unit/transformations/indirect_kv_cache_test.cpp b/src/plugins/intel_gpu/tests/unit/transformations/indirect_kv_cache_test.cpp index b23c3064a56..dc2203d7d34 100644 --- a/src/plugins/intel_gpu/tests/unit/transformations/indirect_kv_cache_test.cpp +++ b/src/plugins/intel_gpu/tests/unit/transformations/indirect_kv_cache_test.cpp @@ -22,6 +22,8 @@ #include "intel_gpu/op/indirect_gemm.hpp" #include "intel_gpu/op/gemm.hpp" +#include "intel_gpu/op/indirect_sdpa.hpp" +#include "intel_gpu/op/sdpa.hpp" #include "intel_gpu/op/read_value.hpp" #include "intel_gpu/op/kv_cache.hpp" @@ -132,3 +134,51 @@ TEST_F(TransformationTestsF, IndirectKVCache3) { comparator.enable(FunctionsComparator::ATTRIBUTES); } } + +TEST_F(TransformationTestsF, IndirectKVCache4) { + std::vector in0_order = {0, 1, 2, 3}; + std::vector in1_order = {0, 2, 1, 3}; + std::vector in2_order = {0, 2, 1, 3}; + std::vector out_order = {0, 1, 2, 3}; + const bool is_causal = true; + { + auto key_variable = std::make_shared(ov::op::util::VariableInfo{{1, -1, 32, 128}, ov::element::f32, "v0"}); + auto value_variable = std::make_shared(ov::op::util::VariableInfo{{1, -1, 32, 128}, ov::element::f32, "v1"}); + auto parameter_key = std::make_shared(ov::element::f32, ov::PartialShape{1, -1, 32, 128}); + auto parameter_value = std::make_shared(ov::element::f32, ov::PartialShape{1, -1, 32, 128}); + auto beam_idx = std::make_shared(ov::element::i32, ov::PartialShape{1}); + auto key_past = std::make_shared(key_variable); + auto value_past = std::make_shared(value_variable); + auto axis = std::make_shared(ov::element::i64, ov::Shape{}, 1); + auto key_gather_past = std::make_shared(key_past, beam_idx, axis); + auto value_gather_past = std::make_shared(value_past, beam_idx, axis); + auto key_cache = std::make_shared(key_gather_past, parameter_key, key_variable, 2, ov::element::f32); + auto value_cache = std::make_shared(value_gather_past, parameter_value, value_variable, 2, ov::element::f32); + auto sdpa_q = std::make_shared(ov::element::f32, ov::PartialShape{1, 32, -1, 128}); + auto inputs = ov::OutputVector{sdpa_q, key_cache, value_cache}; + auto sdpa = std::make_shared(inputs, is_causal, in0_order, in1_order, in2_order, out_order); + auto result = std::make_shared(sdpa); + + model = std::make_shared(ov::ResultVector{result}, ov::ParameterVector{parameter_key, parameter_value, beam_idx, sdpa_q}); + manager.register_pass(); + } + { + auto indirect_axis = 1; + auto key_variable = std::make_shared(ov::op::util::VariableInfo{{1, -1, 32, 128}, ov::element::f32, "v0"}); + auto value_variable = std::make_shared(ov::op::util::VariableInfo{{1, -1, 32, 128}, ov::element::f32, "v1"}); + auto parameter_key = std::make_shared(ov::element::f32, ov::PartialShape{1, -1, 32, 128}); + auto parameter_value = std::make_shared(ov::element::f32, ov::PartialShape{1, -1, 32, 128}); + auto beam_idx = std::make_shared(ov::element::i32, ov::PartialShape{1}); + auto key_past = std::make_shared(key_variable); + auto value_past = std::make_shared(value_variable); + auto key_cache = std::make_shared(key_past, parameter_key, beam_idx, key_variable, 2, indirect_axis, ov::element::f32); + auto value_cache = std::make_shared(value_past, parameter_value, beam_idx, value_variable, 2, indirect_axis, ov::element::f32); + auto sdpa_q = std::make_shared(ov::element::f32, ov::PartialShape{1, 32, -1, 128}); + auto inputs = ov::OutputVector{sdpa_q, key_cache, value_cache}; + auto sdpa = std::make_shared(inputs, key_cache->output(1), is_causal, indirect_axis, in0_order, in1_order, in2_order, out_order); + auto result = std::make_shared(sdpa); + + model_ref = std::make_shared(ov::ResultVector{result}, ov::ParameterVector{parameter_key, parameter_value, beam_idx, sdpa_q}); + comparator.enable(FunctionsComparator::ATTRIBUTES); + } +} diff --git a/src/plugins/intel_gpu/tests/unit/transformations/transpose_sdpa_fusion_test.cpp b/src/plugins/intel_gpu/tests/unit/transformations/transpose_sdpa_fusion_test.cpp index ebe15f4d806..f6d21681767 100644 --- a/src/plugins/intel_gpu/tests/unit/transformations/transpose_sdpa_fusion_test.cpp +++ b/src/plugins/intel_gpu/tests/unit/transformations/transpose_sdpa_fusion_test.cpp @@ -25,11 +25,12 @@ namespace test { namespace intel_gpu { TEST_F(TransformationTestsF, TranposeSDPAFusion1) { + const bool is_causal = true; { auto input_a = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_b = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_c = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); - auto sdpa = std::make_shared(input_a, input_b, input_c, true); + auto sdpa = std::make_shared(input_a, input_b, input_c, is_causal); model = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); manager.register_pass(); @@ -42,21 +43,22 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion1) { auto input_a = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_b = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_c = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); - auto sdpa = std::make_shared(input_a, input_b, input_c, order_a, order_b, order_c, order_output, true, ov::element::undefined ); + auto sdpa = std::make_shared(ov::OutputVector{input_a, input_b, input_c}, is_causal, order_a, order_b, order_c, order_output, ov::element::undefined ); model_ref = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); comparator.enable(FunctionsComparator::ATTRIBUTES); } } -TEST_F(TransformationTestsF, TranposeSDPAFusion2) { +TEST_F(TransformationTestsF, TransformationTestsF) { + const bool is_causal = true; { auto input_a = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto tranpose_a_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3}); auto tranpose_a = std::make_shared(input_a, tranpose_a_const); auto input_b = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_c = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); - auto sdpa = std::make_shared(tranpose_a, input_b, input_c, true); + auto sdpa = std::make_shared(tranpose_a, input_b, input_c, is_causal); model = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); manager.register_pass(); @@ -69,7 +71,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion2) { auto input_a = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_b = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_c = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); - auto sdpa = std::make_shared(input_a, input_b, input_c, order_a, order_b, order_c, order_output, true, ov::element::undefined); + auto sdpa = std::make_shared(ov::OutputVector{input_a, input_b, input_c}, is_causal, order_a, order_b, order_c, order_output, ov::element::undefined); model_ref = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); comparator.enable(FunctionsComparator::ATTRIBUTES); @@ -77,6 +79,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion2) { } TEST_F(TransformationTestsF, TranposeSDPAFusion3) { + const bool is_causal = false; { auto input_a = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto tranpose_a_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3}); @@ -86,7 +89,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion3) { auto tranpose_b = std::make_shared(input_b, tranpose_b_const); auto input_c = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); - auto sdpa = std::make_shared(tranpose_a, tranpose_b, input_c, false); + auto sdpa = std::make_shared(tranpose_a, tranpose_b, input_c, is_causal); model = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); manager.register_pass(); @@ -99,7 +102,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion3) { auto input_a = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_b = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_c = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); - auto sdpa = std::make_shared(input_a, input_b, input_c, order_a, order_b, order_c, order_output, false, ov::element::undefined); + auto sdpa = std::make_shared(ov::OutputVector{input_a, input_b, input_c}, is_causal, order_a, order_b, order_c, order_output, ov::element::undefined); model_ref = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); comparator.enable(FunctionsComparator::ATTRIBUTES); @@ -107,6 +110,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion3) { } TEST_F(TransformationTestsF, TranposeSDPAFusion4) { + const bool is_causal = false; { auto input_a = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto tranpose_a_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3}); @@ -118,7 +122,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion4) { auto tranpose_c_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3}); auto tranpose_c = std::make_shared(input_c, tranpose_c_const); - auto sdpa = std::make_shared(tranpose_a, tranpose_b, tranpose_c, false); + auto sdpa = std::make_shared(tranpose_a, tranpose_b, tranpose_c, is_causal); model = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); manager.register_pass(); @@ -131,7 +135,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion4) { auto input_a = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_b = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto input_c = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); - auto sdpa = std::make_shared(input_a, input_b, input_c, order_a, order_b, order_c, order_output, false, ov::element::undefined); + auto sdpa = std::make_shared(ov::OutputVector{input_a, input_b, input_c}, is_causal, order_a, order_b, order_c, order_output, ov::element::undefined); model_ref = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); comparator.enable(FunctionsComparator::ATTRIBUTES); @@ -139,6 +143,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion4) { } TEST_F(TransformationTestsF, TranposeSDPAFusion5) { + const bool is_causal = false; { auto input_a = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); auto tranpose_a_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3}); @@ -150,7 +155,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion5) { auto tranpose_c_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {3, 2, 1, 0}); auto tranpose_c = std::make_shared(input_c, tranpose_c_const); - auto sdpa = std::make_shared(tranpose_a, tranpose_b, tranpose_c, false); + auto sdpa = std::make_shared(tranpose_a, tranpose_b, tranpose_c, is_causal); model = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); manager.register_pass(); @@ -166,7 +171,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion5) { auto tranpose_c_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {3, 2, 1, 0}); auto tranpose_c = std::make_shared(input_c, tranpose_c_const); - auto sdpa = std::make_shared(tranpose_a, tranpose_b, tranpose_c, false); + auto sdpa = std::make_shared(tranpose_a, tranpose_b, tranpose_c, is_causal); model_ref = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c }); comparator.enable(FunctionsComparator::ATTRIBUTES); diff --git a/src/plugins/intel_gpu/tests/unit/transformations/unsqueeze_broadcast_reshape_sdpa_fusion_test.cpp b/src/plugins/intel_gpu/tests/unit/transformations/unsqueeze_broadcast_reshape_sdpa_fusion_test.cpp new file mode 100644 index 00000000000..b24124e679f --- /dev/null +++ b/src/plugins/intel_gpu/tests/unit/transformations/unsqueeze_broadcast_reshape_sdpa_fusion_test.cpp @@ -0,0 +1,233 @@ +// Copyright (C) 2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "common_test_utils/ov_test_utils.hpp" + +#include "openvino/core/model.hpp" +#include "openvino/pass/manager.hpp" +#include "openvino/op/abs.hpp" +#include "openvino/op/broadcast.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/gather.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/unsqueeze.hpp" + +#include "intel_gpu/op/sdpa.hpp" +#include "intel_gpu/op/read_value.hpp" +#include "intel_gpu/op/kv_cache.hpp" + +#include "plugin/transformations/unsqueeze_broadcast_reshape_sdpa_fusion.hpp" + +#include + +using namespace testing; +using namespace ov::intel_gpu; + +namespace ov { +namespace test { +namespace intel_gpu { + +TEST_F(TransformationTestsF, UnsqueezeBroadReshapeSDPAFusion1) { + std::vector in0_order = {0, 2, 1, 3}; + std::vector in1_order = {0, 2, 1, 3}; + std::vector in2_order = {0, 2, 1, 3}; + std::vector out_order = {0, 1, 2, 3}; + std::vector axes_val = {-2}; + std::vector target_shape_kv = {1, 1, 1, 16, 1}; + std::vector pattern_shape = {0, 0, 32, 32}; + const bool is_causal = true; + { + auto input_q = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); + auto key_variable = std::make_shared(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v0"}); + auto value_variable = std::make_shared(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v1"}); + auto key_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, -1, 2, 32}); + auto value_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, -1, 2, 32}); + auto beam_idx = std::make_shared(ov::element::i32, ov::PartialShape{1}); + auto key_past = std::make_shared(key_variable); + auto value_past = std::make_shared(key_variable); + auto axis = std::make_shared(ov::element::i64, ov::Shape{}, 0); + auto key_gather_past = std::make_shared(key_past, beam_idx, axis); + auto value_gather_past = std::make_shared(value_past, beam_idx, axis); + auto key_cache = std::make_shared(key_gather_past, key_token_param, key_variable, 2, ov::element::f32); + auto value_cache = std::make_shared(value_gather_past, value_token_param, value_variable, 2, ov::element::f32); + auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, axes_val); + auto key_unsqueeze = std::make_shared(key_cache, axes); + auto value_unsqueeze = std::make_shared(value_cache, axes); + auto target_shape = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{5}, target_shape_kv); + auto key_broadcast = std::make_shared(key_unsqueeze, target_shape, ov::op::BroadcastType::BIDIRECTIONAL); + auto value_broadcast = std::make_shared(value_unsqueeze, target_shape, ov::op::BroadcastType::BIDIRECTIONAL); + auto pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, pattern_shape); + auto key_reshape = std::make_shared(key_broadcast, pattern, true); + auto value_reshape = std::make_shared(value_broadcast, pattern, true); + auto inputs = ov::OutputVector{input_q, key_reshape, value_reshape}; + auto sdpa = std::make_shared(inputs, is_causal, in0_order, in1_order, in2_order, out_order); + + model = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx }); + manager.register_pass(); + } + { + auto input_q = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); + auto key_variable = std::make_shared(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v0"}); + auto value_variable = std::make_shared(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v1"}); + auto key_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, -1, 2, 32}); + auto value_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, -1, 2, 32}); + auto beam_idx = std::make_shared(ov::element::i32, ov::PartialShape{1}); + auto key_past = std::make_shared(key_variable); + auto value_past = std::make_shared(key_variable); + auto axis = std::make_shared(ov::element::i64, ov::Shape{}, 0); + auto key_gather_past = std::make_shared(key_past, beam_idx, axis); + auto value_gather_past = std::make_shared(value_past, beam_idx, axis); + auto key_cache = std::make_shared(key_gather_past, key_token_param, key_variable, 2, ov::element::f32); + auto value_cache = std::make_shared(value_gather_past, value_token_param, value_variable, 2, ov::element::f32); + auto inputs = ov::OutputVector{input_q, key_cache, value_cache}; + auto sdpa = std::make_shared(inputs, is_causal, in0_order, in1_order, in2_order, out_order); + + model_ref = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx }); + comparator.enable(FunctionsComparator::ATTRIBUTES); + } +} + +TEST_F(TransformationTestsF, UnsqueezeBroadReshapeSDPAFusion2) { + std::vector in0_order = {0, 2, 1, 3}; + std::vector in1_order = {0, 2, 1, 3}; + std::vector in2_order = {0, 2, 1, 3}; + std::vector out_order = {0, 1, 2, 3}; + std::vector axes_val = {2}; + std::vector pattern_shape = {0, 32, -1, 32}; + const bool is_causal = true; + { + auto input_q = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); + auto key_variable = std::make_shared(ov::op::util::VariableInfo{{-1, 8, -1, 32}, ov::element::f32, "v0"}); + auto value_variable = std::make_shared(ov::op::util::VariableInfo{{-1, 8, -1, 32}, ov::element::f32, "v1"}); + auto key_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, 8, -1, 32}); + auto value_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, 8, -1, 32}); + auto beam_idx = std::make_shared(ov::element::i32, ov::PartialShape{1}); + auto key_past = std::make_shared(key_variable); + auto value_past = std::make_shared(key_variable); + auto axis = std::make_shared(ov::element::i64, ov::Shape{}, 0); + auto key_gather_past = std::make_shared(key_past, beam_idx, axis); + auto value_gather_past = std::make_shared(value_past, beam_idx, axis); + auto key_cache = std::make_shared(key_gather_past, key_token_param, key_variable, 2, ov::element::f32); + auto value_cache = std::make_shared(value_gather_past, value_token_param, value_variable, 2, ov::element::f32); + auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, axes_val); + auto key_unsqueeze = std::make_shared(key_cache, axes); + auto value_unsqueeze = std::make_shared(value_cache, axes); + auto abs_param = std::make_shared(ov::element::i32, ov::Shape{5}); + auto abs = std::make_shared(abs_param); + auto key_broadcast = std::make_shared(key_unsqueeze, abs, ov::op::BroadcastType::BIDIRECTIONAL); + auto value_broadcast = std::make_shared(value_unsqueeze, abs, ov::op::BroadcastType::BIDIRECTIONAL); + auto pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, pattern_shape); + auto key_reshape = std::make_shared(key_broadcast, pattern, true); + auto value_reshape = std::make_shared(value_broadcast, pattern, true); + auto inputs = ov::OutputVector{input_q, key_cache, value_cache}; + auto sdpa = std::make_shared(inputs, is_causal, in0_order, in1_order, in2_order, out_order); + + model = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx }); + manager.register_pass(); + } + { + auto input_q = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); + auto key_variable = std::make_shared(ov::op::util::VariableInfo{{-1, 8, -1, 32}, ov::element::f32, "v0"}); + auto value_variable = std::make_shared(ov::op::util::VariableInfo{{-1, 8, -1, 32}, ov::element::f32, "v1"}); + auto key_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, 8, -1, 32}); + auto value_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, 8, -1, 32}); + auto beam_idx = std::make_shared(ov::element::i32, ov::PartialShape{1}); + auto key_past = std::make_shared(key_variable); + auto value_past = std::make_shared(key_variable); + auto axis = std::make_shared(ov::element::i64, ov::Shape{}, 0); + auto key_gather_past = std::make_shared(key_past, beam_idx, axis); + auto value_gather_past = std::make_shared(value_past, beam_idx, axis); + auto key_cache = std::make_shared(key_gather_past, key_token_param, key_variable, 2, ov::element::f32); + auto value_cache = std::make_shared(value_gather_past, value_token_param, value_variable, 2, ov::element::f32); + auto inputs = ov::OutputVector{input_q, key_cache, value_cache}; + auto sdpa = std::make_shared(inputs, is_causal, in0_order, in1_order, in2_order, out_order); + + model_ref = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx }); + comparator.enable(FunctionsComparator::ATTRIBUTES); + } +} + +TEST_F(TransformationTestsF, UnsqueezeBroadReshapeSDPAFusion3) { + std::vector in0_order = {0, 2, 1, 3}; + std::vector in1_order = {0, 2, 1, 3}; + std::vector in2_order = {0, 2, 1, 3}; + std::vector out_order = {0, 1, 2, 3}; + std::vector axes_val = {-2}; + std::vector target_shape_kv = {1, 1, 1, 16, 1}; + std::vector pattern_shape = {0, 0, -1, 32}; + const bool is_causal = true; + { + auto input_q = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); + auto input_k = std::make_shared(ov::element::f32, ov::PartialShape{-1, -1, -1, 32}); + auto input_v = std::make_shared(ov::element::f32, ov::PartialShape{-1, -1, -1, 32}); + auto unsqueeze_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, axes_val); + auto key_unsqueeze = std::make_shared(input_k, unsqueeze_const); + auto value_unsqueeze = std::make_shared(input_v, unsqueeze_const); + auto broadcast_const = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{5}, target_shape_kv); + auto key_broadcast = std::make_shared(key_unsqueeze, broadcast_const, ov::op::BroadcastType::BIDIRECTIONAL); + auto value_broadcast = std::make_shared(value_unsqueeze, broadcast_const, ov::op::BroadcastType::BIDIRECTIONAL); + auto reshape_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, pattern_shape); + auto key_reshape = std::make_shared(key_broadcast, reshape_const, true); + auto value_reshape = std::make_shared(value_broadcast, reshape_const, true); + auto inputs = ov::OutputVector{input_q, key_reshape, value_reshape}; + auto sdpa = std::make_shared(inputs, is_causal, in0_order, in1_order, in2_order, out_order); + + model = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, input_k, input_v }); + manager.register_pass(); + } + { + model_ref = model->clone(); + comparator.enable(FunctionsComparator::ATTRIBUTES); + } +} + +TEST_F(TransformationTestsF, UnsqueezeBroadReshapeSDPAFusion4) { + std::vector in0_order = {0, 2, 1, 3}; + std::vector in1_order = {0, 2, 1, 3}; + std::vector in2_order = {0, 2, 1, 3}; + std::vector out_order = {0, 1, 2, 3}; + std::vector axes_val = {-2}; + std::vector target_shape_k = {1, 1, 1, 16, 1}; + std::vector target_shape_v = {1, 1, 1, 32, 1}; + std::vector pattern_shape = {0, 0, 32, 32}; + const bool is_causal = true; + { + auto input_q = std::make_shared(ov::element::f32, ov::PartialShape::dynamic(4)); + auto key_variable = std::make_shared(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v0"}); + auto value_variable = std::make_shared(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v1"}); + auto key_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, -1, 2, 32}); + auto value_token_param = std::make_shared(ov::element::f32, ov::PartialShape{-1, -1, 2, 32}); + auto beam_idx = std::make_shared(ov::element::i32, ov::PartialShape{1}); + auto key_past = std::make_shared(key_variable); + auto value_past = std::make_shared(key_variable); + auto axis = std::make_shared(ov::element::i64, ov::Shape{}, 0); + auto key_gather_past = std::make_shared(key_past, beam_idx, axis); + auto value_gather_past = std::make_shared(value_past, beam_idx, axis); + auto key_cache = std::make_shared(key_gather_past, key_token_param, key_variable, 2, ov::element::f32); + auto value_cache = std::make_shared(value_gather_past, value_token_param, value_variable, 2, ov::element::f32); + auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, axes_val); + auto key_unsqueeze = std::make_shared(key_cache, axes); + auto value_unsqueeze = std::make_shared(value_cache, axes); + auto target_shape_key = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{5}, target_shape_k); + auto target_shape_value = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{5}, target_shape_v); + auto key_broadcast = std::make_shared(key_unsqueeze, target_shape_key, ov::op::BroadcastType::BIDIRECTIONAL); + auto value_broadcast = std::make_shared(value_unsqueeze, target_shape_value, ov::op::BroadcastType::BIDIRECTIONAL); + auto pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, pattern_shape); + auto key_reshape = std::make_shared(key_broadcast, pattern, true); + auto value_reshape = std::make_shared(value_broadcast, pattern, true); + auto inputs = ov::OutputVector{input_q, key_reshape, value_reshape}; + auto sdpa = std::make_shared(inputs, is_causal, in0_order, in1_order, in2_order, out_order); + + model = std::make_shared(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx }); + manager.register_pass(); + } + { + model_ref = model->clone(); + comparator.enable(FunctionsComparator::ATTRIBUTES); + } +} + +} // namespace intel_gpu +} // namespace test +} // namespace ov