apply GPU transformation to avoid KV cache broadcast

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 <openvino/util/pp.hpp>
+
+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<ov::op::v1::Reshape>(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<ov::Node>& 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<ov::Node>& output) {
+            auto concat = ov::as_type<ov::op::v0::Concat>(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<ov::op::v0::Concat>(broadcast_concat);
+
+        auto reshape2_pattern_key_m = any_input();
+        auto reshape_5d_to_4d_key_m = wrap_type<ov::op::v1::Reshape>({ concat_key_m, reshape2_pattern_key_m }, reshape_5d_to_4d);
+
+        auto concat_value_m = wrap_type<ov::op::v0::Concat>(broadcast_concat);
+
+        auto reshape2_pattern_value_m = any_input();
+        auto reshape_5d_to_4d_value_m = wrap_type<ov::op::v1::Reshape>({ concat_value_m, reshape2_pattern_value_m }, reshape_5d_to_4d);
+
+        auto sdpa_without_attn_mask_m = wrap_type<op::SDPA>({ input_a_m, reshape_5d_to_4d_key_m, reshape_5d_to_4d_value_m });
+        auto sdpa_with_attn_mask_m =
+            wrap_type<op::SDPA>({ 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<op::SDPA>({ 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<ov::pass::pattern::op::Or>(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<op::SDPA>(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<ov::Node>& 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<ov::pass::pattern::Matcher>(sdpa_m, "KVCacheGQABroadcastElimination");
+        this->register_matcher(m, callback);
+    }
+
+}  // namespace ov::intel_gpu

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

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<ov::Model> func) {
             manager.register_pass<ov::intel_gpu::UnsqueezeBroadcastReshapeMatmulFusion>();
         }
         manager.register_pass<ov::intel_gpu::UnsqueezeBroadcastReshapeSDPAFusion>();
+        manager.register_pass<ov::intel_gpu::KVCacheGQABroadcastElimination>();
+
 
         manager.register_pass<ov::pass::GLUFusion>();
         manager.register_pass<ov::intel_gpu::IndirectKVCache>();

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 <memory>
+
+using namespace testing;
+using namespace ov::intel_gpu;
+
+namespace ov {
+namespace test {
+namespace intel_gpu {
+
+TEST_F(TransformationTestsF, KVCacheGQABroadcastEliminationTest1) {
+    std::vector<int64_t> in0_order = {0, 1, 2, 3};
+    std::vector<int64_t> in1_order = {0, 1, 2, 3};
+    std::vector<int64_t> in2_order = {0, 1, 2, 3};
+    std::vector<int64_t> out_order = {0, 1, 2, 3};
+    const bool is_causal = false;
+    {
+        auto input_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::Shape{1, 4, 8, 16});
+        auto input_k = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::Shape{1, 2, 8, 16});
+        auto input_v = std::make_shared<ov::op::v0::Parameter>(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<ov::op::v3::ScatterUpdate>(input_k, indices, k_updates, axis);
+        auto scatter_v = std::make_shared<ov::op::v3::ScatterUpdate>(input_v, indices, v_updates, axis);
+        auto split_k = std::make_shared<ov::op::v1::VariadicSplit>(scatter_k, axis, split_lengths);
+        auto split_v = std::make_shared<ov::op::v1::VariadicSplit>(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<ov::op::v1::Reshape>(split_k_out, reshape_k_5d_pattern, false);
+        auto reshape_v_5d = std::make_shared<ov::op::v1::Reshape>(split_v_out, reshape_v_5d_pattern, false);
+
+        auto concat_k = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{reshape_k_5d, reshape_k_5d, reshape_k_5d, reshape_k_5d},2);
+        auto concat_v = std::make_shared<ov::op::v0::Concat>(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<ov::op::v1::Reshape>(concat_k, reshape_k_4d_pattern, false);
+        auto reshape_v_4d = std::make_shared<ov::op::v1::Reshape>(concat_v, reshape_v_4d_pattern, false);
+
+        auto inputs = ov::OutputVector{input_q, reshape_k_4d, reshape_v_4d};
+        auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(inputs, is_causal, in0_order, in1_order, in2_order, out_order);
+
+        model = std::make_shared<ov::Model>(ov::OutputVector{ sdpa }, ov::ParameterVector{input_q, input_k, input_v});
+        manager.register_pass<KVCacheGQABroadcastElimination>();
+    }
+    {
+        auto input_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::Shape{1, 4, 8, 16});
+        auto input_k = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::Shape{1, 2, 8, 16});
+        auto input_v = std::make_shared<ov::op::v0::Parameter>(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<ov::op::v3::ScatterUpdate>(input_k, indices, k_updates, axis);
+        auto scatter_v = std::make_shared<ov::op::v3::ScatterUpdate>(input_v, indices, v_updates, axis);
+        auto split_k = std::make_shared<ov::op::v1::VariadicSplit>(scatter_k, axis, split_lengths);
+        auto split_v = std::make_shared<ov::op::v1::VariadicSplit>(scatter_v, axis, split_lengths);
+
+        auto inputs = ov::OutputVector{ input_q, split_k->output(0), split_v->output(0) };
+        auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(inputs, is_causal, in0_order, in1_order, in2_order, out_order);
+
+        model_ref = std::make_shared<ov::Model>(ov::OutputVector{ sdpa }, ov::ParameterVector{ input_q, input_k, input_v });
+        comparator.enable(FunctionsComparator::ATTRIBUTES);
+    }
+}
+
+}  // namespace intel_gpu
+}  // namespace test
+}  // namespace ov