diff --git a/src/plugins/intel_gpu/src/plugin/transformations/kvcache_gqa_broadcast_elimination.cpp b/src/plugins/intel_gpu/src/plugin/transformations/kvcache_gqa_broadcast_elimination.cpp new file mode 100644 index 00000000000000..1061e0a6a6c256 --- /dev/null +++ b/src/plugins/intel_gpu/src/plugin/transformations/kvcache_gqa_broadcast_elimination.cpp @@ -0,0 +1,91 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "kvcache_gqa_broadcast_elimination.hpp" +#include "openvino/pass/pattern/matcher.hpp" +#include "openvino/pass/pattern/op/or.hpp" +#include "openvino/pass/pattern/op/wrap_type.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/concat.hpp" +#include "openvino/core/partial_shape.hpp" +#include "intel_gpu/op/sdpa.hpp" +#include + +namespace ov::intel_gpu { + KVCacheGQABroadcastElimination::KVCacheGQABroadcastElimination() { + using namespace ov::pass::pattern; + auto is_reshape_4d_to_5d = [](const ov::Node* node) { + if (!ov::is_type(node)) { + return false; + } + auto in_ps = node->get_input_partial_shape(0); + auto out_ps = node->get_output_partial_shape(0); + return in_ps.rank().is_static() && out_ps.rank().is_static() && in_ps.size() == 4 && out_ps.size() == 5; + }; + + auto reshape_5d_to_4d = [](const ov::Output& output) { + auto in_ps = output.get_node()->get_input_partial_shape(0); + auto out_ps = output.get_node()->get_output_partial_shape(0); + return in_ps.rank().is_static() && out_ps.rank().is_static() && in_ps.size() == 5 && out_ps.size() == 4; + }; + + auto broadcast_concat = [is_reshape_4d_to_5d](const ov::Output& output) { + auto concat = ov::as_type(output.get_node()); + if (concat == nullptr || concat->get_input_size() < 2) { + return false; + } + auto first = concat->get_input_node_ptr(0); + if (!is_reshape_4d_to_5d(first)) { + return false; + } + for (size_t i = 1; i < concat->get_input_size(); ++i) { + if (concat->get_input_node_ptr(i) != first) { + return false; + + } + } + return true; + }; + + auto input_a_m = any_input(); + auto input_attn_mask_m = any_input(); + auto input_scale_m = any_input(); + + auto concat_key_m = wrap_type(broadcast_concat); + + auto reshape2_pattern_key_m = any_input(); + auto reshape_5d_to_4d_key_m = wrap_type({ concat_key_m, reshape2_pattern_key_m }, reshape_5d_to_4d); + + auto concat_value_m = wrap_type(broadcast_concat); + + auto reshape2_pattern_value_m = any_input(); + auto reshape_5d_to_4d_value_m = wrap_type({ concat_value_m, reshape2_pattern_value_m }, reshape_5d_to_4d); + + auto sdpa_without_attn_mask_m = wrap_type({ input_a_m, reshape_5d_to_4d_key_m, reshape_5d_to_4d_value_m }); + auto sdpa_with_attn_mask_m = + wrap_type({ input_a_m, reshape_5d_to_4d_key_m, reshape_5d_to_4d_value_m, any_input() }); + auto sdpa_with_attn_mask_and_scale_m = + wrap_type({ input_a_m, reshape_5d_to_4d_key_m, reshape_5d_to_4d_value_m, input_attn_mask_m, input_scale_m }); + + auto sdpa_m = std::make_shared(OutputVector{ sdpa_without_attn_mask_m, sdpa_with_attn_mask_m, sdpa_with_attn_mask_and_scale_m }); + + ov::matcher_pass_callback callback = [OV_CAPTURE_CPY_AND_THIS](ov::pass::pattern::Matcher& m) { + const auto& pattern_map = m.get_pattern_value_map(); + auto sdpa = ov::as_type_ptr(m.get_match_root()); + if (pattern_map.count(concat_key_m) == 0 || pattern_map.count(concat_value_m) == 0) { + return false; + } + // Concat -> Reshape(4d->5d) -> original source that was broadcast. + auto broadcast_source = [](const ov::Output& concat_out) { + return concat_out.get_node()->input_value(0).get_node()->input_value(0); + }; + sdpa->input(1).replace_source_output(broadcast_source(pattern_map.at(concat_key_m))); + sdpa->input(2).replace_source_output(broadcast_source(pattern_map.at(concat_value_m))); + return true; + }; + auto m = std::make_shared(sdpa_m, "KVCacheGQABroadcastElimination"); + this->register_matcher(m, callback); + } + +} // namespace ov::intel_gpu diff --git a/src/plugins/intel_gpu/src/plugin/transformations/kvcache_gqa_broadcast_elimination.hpp b/src/plugins/intel_gpu/src/plugin/transformations/kvcache_gqa_broadcast_elimination.hpp new file mode 100644 index 00000000000000..ef598d326d8fcd --- /dev/null +++ b/src/plugins/intel_gpu/src/plugin/transformations/kvcache_gqa_broadcast_elimination.hpp @@ -0,0 +1,17 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include "openvino/pass/graph_rewrite.hpp" + +namespace ov::intel_gpu { + +class KVCacheGQABroadcastElimination : public ov::pass::MatcherPass { +public: + OPENVINO_MATCHER_PASS_RTTI("KVCacheGQABroadcastElimination"); + KVCacheGQABroadcastElimination(); +}; + +} // namespace ov::intel_gpu diff --git a/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp b/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp index c626f8c876cb6c..47576230b9c72b 100644 --- a/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp +++ b/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp @@ -114,6 +114,7 @@ #include "plugin/transformations/transpose_fusion.hpp" #include "plugin/transformations/unsqueeze_broadcast_reshape_matmul_fusion.hpp" #include "plugin/transformations/unsqueeze_broadcast_reshape_sdpa_fusion.hpp" +#include "plugin/transformations/kvcache_gqa_broadcast_elimination.hpp" #include "plugin/transformations/disable_fp16_comp_rms.hpp" #include "plugin/transformations/swiglu_fusion_with_clamp.hpp" #include "plugin/transformations/disable_fp16_comp_cumsum_sin_gen.hpp" @@ -1609,6 +1610,8 @@ void TransformationsPipeline::apply(std::shared_ptr func) { manager.register_pass(); } manager.register_pass(); + manager.register_pass(); + manager.register_pass(); manager.register_pass(); diff --git a/src/plugins/intel_gpu/tests/unit/transformations/kvcache_gqa_broadcast_elimination_test.cpp b/src/plugins/intel_gpu/tests/unit/transformations/kvcache_gqa_broadcast_elimination_test.cpp new file mode 100644 index 00000000000000..183a57a4ffd66d --- /dev/null +++ b/src/plugins/intel_gpu/tests/unit/transformations/kvcache_gqa_broadcast_elimination_test.cpp @@ -0,0 +1,95 @@ +// Copyright (C) 2018-2026 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/constant.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/concat.hpp" +#include "openvino/op/scatter_update.hpp" +#include "openvino/op/variadic_split.hpp" +#include "intel_gpu/op/sdpa.hpp" +#include "plugin/transformations/kvcache_gqa_broadcast_elimination.hpp" +#include + +using namespace testing; +using namespace ov::intel_gpu; + +namespace ov { +namespace test { +namespace intel_gpu { + +TEST_F(TransformationTestsF, KVCacheGQABroadcastEliminationTest1) { + std::vector in0_order = {0, 1, 2, 3}; + std::vector in1_order = {0, 1, 2, 3}; + std::vector in2_order = {0, 1, 2, 3}; + std::vector out_order = {0, 1, 2, 3}; + const bool is_causal = false; + { + auto input_q = std::make_shared(ov::element::f16, ov::Shape{1, 4, 8, 16}); + auto input_k = std::make_shared(ov::element::f16, ov::Shape{1, 2, 8, 16}); + auto input_v = std::make_shared(ov::element::f16, ov::Shape{1, 2, 8, 16}); + + auto indices = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {0}); + auto axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {1}); + auto split_lengths = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{2}, {1, 1}); + auto k_updates = ov::op::v0::Constant::create(ov::element::f16, ov::Shape{1, 1, 8, 16}, {0.0f}); + auto v_updates = ov::op::v0::Constant::create(ov::element::f16, ov::Shape{1, 1, 8, 16}, {0.0f}); + + auto scatter_k = std::make_shared(input_k, indices, k_updates, axis); + auto scatter_v = std::make_shared(input_v, indices, v_updates, axis); + auto split_k = std::make_shared(scatter_k, axis, split_lengths); + auto split_v = std::make_shared(scatter_v, axis, split_lengths); + auto split_k_out = split_k->output(0); + auto split_v_out = split_v->output(0); + + auto reshape_k_5d_pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{5}, {1, 1, 1, 8, 16}); + auto reshape_v_5d_pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{5}, {1, 1, 1, 8, 16}); + + auto reshape_k_5d = std::make_shared(split_k_out, reshape_k_5d_pattern, false); + auto reshape_v_5d = std::make_shared(split_v_out, reshape_v_5d_pattern, false); + + auto concat_k = std::make_shared(ov::OutputVector{reshape_k_5d, reshape_k_5d, reshape_k_5d, reshape_k_5d},2); + auto concat_v = std::make_shared(ov::OutputVector{reshape_v_5d, reshape_v_5d, reshape_v_5d, reshape_v_5d},2); + + auto reshape_k_4d_pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{ 4 }, { 1, 4, 8, 16 }); + auto reshape_v_4d_pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{ 4 }, { 1, 4, 8, 16 }); + + auto reshape_k_4d = std::make_shared(concat_k, reshape_k_4d_pattern, false); + auto reshape_v_4d = std::make_shared(concat_v, reshape_v_4d_pattern, false); + + auto inputs = ov::OutputVector{input_q, reshape_k_4d, reshape_v_4d}; + auto sdpa = std::make_shared(inputs, is_causal, in0_order, in1_order, in2_order, out_order); + + model = std::make_shared(ov::OutputVector{ sdpa }, ov::ParameterVector{input_q, input_k, input_v}); + manager.register_pass(); + } + { + auto input_q = std::make_shared(ov::element::f16, ov::Shape{1, 4, 8, 16}); + auto input_k = std::make_shared(ov::element::f16, ov::Shape{1, 2, 8, 16}); + auto input_v = std::make_shared(ov::element::f16, ov::Shape{1, 2, 8, 16}); + + auto indices = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {0}); + auto axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {1}); + auto split_lengths = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{2}, {1, 1}); + auto k_updates = ov::op::v0::Constant::create(ov::element::f16, ov::Shape{1, 1, 8, 16}, {0.0f}); + auto v_updates = ov::op::v0::Constant::create(ov::element::f16, ov::Shape{1, 1, 8, 16}, {0.0f}); + + auto scatter_k = std::make_shared(input_k, indices, k_updates, axis); + auto scatter_v = std::make_shared(input_v, indices, v_updates, axis); + auto split_k = std::make_shared(scatter_k, axis, split_lengths); + auto split_v = std::make_shared(scatter_v, axis, split_lengths); + + auto inputs = ov::OutputVector{ input_q, split_k->output(0), split_v->output(0) }; + auto sdpa = std::make_shared(inputs, is_causal, in0_order, in1_order, in2_order, out_order); + + model_ref = std::make_shared(ov::OutputVector{ sdpa }, ov::ParameterVector{ input_q, input_k, input_v }); + comparator.enable(FunctionsComparator::ATTRIBUTES); + } +} + +} // namespace intel_gpu +} // namespace test +} // namespace ov