Simplify broadcast op in attention

Author: wine99

ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp

@@ -2,9 +2,11 @@
 #include "../op_table.hpp"
 #include "../utils.hpp"
 
+#include <cstdint>
 #include <memory>
 #include <openvino/op/broadcast.hpp>
 #include <openvino/op/concat.hpp>
+#include <openvino/op/constant.hpp>
 #include <openvino/op/convert.hpp>
 #include <openvino/op/reshape.hpp>
 #include <openvino/op/scaled_dot_product_attention.hpp>
@@ -58,45 +60,40 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
         mask_sliced = std::make_shared<ov::op::v0::Convert>(mask_sliced, ov::element::f16);
     }
 
-    auto tile_kv = [](int64_t q_batch, int64_t kv_batch, ov::Output<Node> kv, bool is_static) {
-        int64_t factor = q_batch / kv_batch;
+    auto tile_kv = [&](int64_t num_heads, int64_t num_heads_kv, int64_t head_size, ov::Output<Node> kv,
+                       bool is_static) {
+        int64_t factor = num_heads / num_heads_kv;
         if (factor > 1) {
-            auto q_batch_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector<int64_t>{q_batch});
-            auto kv_batch_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector<int64_t>{kv_batch});
-            auto factor_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector<int64_t>{factor});
-
             ov::Output<ov::Node> kv_broadcast_shape, kv_unsqueezed, new_kv_shape;
             if (is_static) {
                 auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, Shape{}, {1});
                 kv_unsqueezed = std::make_shared<ov::op::v0::Unsqueeze>(kv, unsqueeze_axes);
 
-                auto kv_last_two_dims = get_dimensions(kv.get_node_shared_ptr(), {1, 2});
-                kv_broadcast_shape = std::make_shared<ov::op::v0::Concat>(
-                    ov::OutputVector{kv_batch_node, factor_node, kv_last_two_dims}, 0);
+                kv_broadcast_shape =
+                    ov::op::v0::Constant::create(ov::element::i64, {4}, {num_heads_kv, factor, (int64_t) 1, head_size});
                 new_kv_shape =
-                    std::make_shared<ov::op::v0::Concat>(ov::OutputVector{q_batch_node, kv_last_two_dims}, 0);
+                    ov::op::v0::Constant::create(ov::element::i64, {3}, {num_heads, (int64_t) -1, head_size});
             } else {
-                auto one_1d = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
                 auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, Shape{}, {2});
                 kv_unsqueezed = std::make_shared<ov::op::v0::Unsqueeze>(kv, unsqueeze_axes);
 
-                auto kv_last_two_dims = get_dimensions(kv.get_node_shared_ptr(), {2, 3});
-                kv_broadcast_shape = std::make_shared<ov::op::v0::Concat>(
-                    ov::OutputVector{one_1d, kv_batch_node, factor_node, kv_last_two_dims}, 0);
-                new_kv_shape =
-                    std::make_shared<ov::op::v0::Concat>(ov::OutputVector{one_1d, q_batch_node, kv_last_two_dims}, 0);
+                kv_broadcast_shape = ov::op::v0::Constant::create(
+                    ov::element::i64, {5}, {(int64_t) 1, num_heads_kv, factor, (int64_t) 1, head_size});
+                new_kv_shape = ov::op::v0::Constant::create(ov::element::i64, {4},
+                                                            {(int64_t) 1, num_heads, (int64_t) -1, head_size});
             }
 
-            kv = std::make_shared<ov::op::v3::Broadcast>(kv_unsqueezed, kv_broadcast_shape);
+            kv = std::make_shared<ov::op::v3::Broadcast>(kv_unsqueezed, kv_broadcast_shape,
+                                                         ov::op::BroadcastType::BIDIRECTIONAL);
             kv = std::make_shared<ov::op::v1::Reshape>(kv, new_kv_shape, false);
         }
         return kv;
     };
 
     auto q_shape = context.get_input_shape(0).to_shape();
     auto k_shape = context.get_input_shape(1).to_shape();
-    k = tile_kv(q_shape[0], k_shape[0], k, context.is_static());
-    v = tile_kv(q_shape[0], k_shape[0], v, context.is_static());
+    k = tile_kv(q_shape[0], k_shape[0], q_shape[2], k, context.is_static());
+    v = tile_kv(q_shape[0], k_shape[0], q_shape[2], v, context.is_static());
 
     auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(q, k, v, mask_sliced, scale_node, false);
     auto sdpa_f32 = std::make_shared<ov::op::v0::Convert>(sdpa, ov::element::f32);