Paged attention transformation

src/bindings/python/src/openvino/_offline_transformations/__init__.py

@@ -17,3 +17,5 @@
 from openvino._pyopenvino._offline_transformations import compress_model_transformation
 from openvino._pyopenvino._offline_transformations import compress_quantize_weights_transformation
 from openvino._pyopenvino._offline_transformations import convert_sequence_to_tensor_iterator_transformation
+from openvino._pyopenvino._offline_transformations import paged_attention_transformation
+from openvino._pyopenvino._offline_transformations import print_model

src/bindings/python/src/pyopenvino/core/offline_transformations.cpp

@@ -18,6 +18,8 @@
 #include <transformations/flush_fp32_subnormals_to_zero.hpp>
 #include <transformations/op_conversions/convert_sequences_to_tensor_iterator.hpp>
 #include <transformations/smart_reshape/smart_reshape.hpp>
+#include <openvino/pass/sdpa_to_paged_attention.hpp>
+#include <transformations/print_model.hpp>
 
 #include "openvino/pass/low_latency.hpp"
 #include "openvino/pass/manager.hpp"
@@ -127,4 +129,26 @@ void regmodule_offline_transformations(py::module m) {
             manager.run_passes(model);
         },
         py::arg("model"));
+
+
+    m_offline_transformations.def(
+        "paged_attention_transformation",
+        [](std::shared_ptr<ov::Model> model) {
+            std::cout << "___CALLING SDPAToPagedAttention___" << std::endl;
+            ov::pass::Manager manager;
+            manager.register_pass<ov::pass::SDPAToPagedAttention>();
+            manager.run_passes(model);
+            std::cout << "___AFTER CALLING SDPAToPagedAttention___" << std::endl;
+        },
+        py::arg("model"));
+
+    m_offline_transformations.def(
+        "print_model",
+        [](std::shared_ptr<ov::Model> model, const std::string& file_name) {
+            ov::pass::Manager manager;
+            manager.register_pass<ov::pass::PrintModel>(file_name);
+            manager.run_passes(model);
+        },
+        py::arg("model"),
+        py::arg("file_name"));
 }

src/bindings/python/src/pyopenvino/graph/ops/paged_attention_extension.cpp

@@ -5,162 +5,13 @@
 #include "pyopenvino/graph/ops/paged_attention_extension.hpp"
 
 #include "openvino/op/op.hpp"
+#include "openvino/op/paged_attention.hpp"
 #include "pyopenvino/core/common.hpp"
 
 namespace py = pybind11;
 
-namespace {
-
-// This is an experimental operation that is implemented in the plugins.
-// Do not use in user applications, backward compatibility is not guaranteed in future releases.
-class PagedAttentionExtension : public ov::op::Op {
-public:
-    OPENVINO_OP("PagedAttentionExtension");
-
-    PagedAttentionExtension(const ov::OutputVector& args) : ov::op::Op(args) {
-        constructor_validate_and_infer_types();
-    }
-
-    void validate_and_infer_types() override {
-        auto value_cache_shape = get_input_partial_shape(4);
-        // m_num_kv_heads = value_cache_shape[1];
-        // m_head_size = value_cache_shape[2];
-        // m_block_size = value_cache_shape[3];
-        NODE_VALIDATION_CHECK(this, value_cache_shape.size() == 4, "Value cache shape must be 4 dims");
-
-        // key_cache: shape [num_blocks, num_kv_heads, head_size/x, block_size, x]
-        auto key_cache_shape = get_input_partial_shape(3);
-        NODE_VALIDATION_CHECK(this,
-                              value_cache_shape.size() == 4,
-                              // value_cache_shape[0] == key_cache_shape[0] && // num_blocks
-                              // key_cache_shape[1] == m_num_kv_heads &&
-                              // key_cache_shape[2] * key_cache_shape[4] == m_head_size &&
-                              // m_block_size == key_cache_shape[3], // block_size,
-                              "Key cache shape must be 4 dims");
-
-        // query: shape [batch_size, seq_len, num_heads * head_size]
-        auto query_type = get_input_element_type(0);
-        auto query_shape = get_input_partial_shape(0);
-        NODE_VALIDATION_CHECK(
-            this,
-            // query_type.is_real() &&
-            query_shape.size() == 3,
-            // query_shape[2] == m_num_heads * m_head_size,
-            "Query type must be real, shape must be like [batch_size, seq_len, num_heads * head_size]. ",
-            "Got element type ",
-            query_type,
-            ", shape ",
-            query_shape);
-
-        // key: shape [batch_size, seq_len, num_kv_heads * head_size]
-        auto key_type = get_input_element_type(1);
-        auto key_shape = get_input_partial_shape(1);
-        NODE_VALIDATION_CHECK(this,
-                              // query_type == key_type &&
-                              key_shape.size() == 3,
-                              "Key type must be the same as query, shape must be the same as query. "
-                              "Got element type ",
-                              key_type,
-                              ", shape ",
-                              key_shape);
-
-        // value: shape [batch_size, seq_len, num_kv_heads * head_size]
-        // auto value_type = get_input_element_type(2);
-        auto value_shape = get_input_partial_shape(2);
-
-        // is_prompt: boolean scalar
-        NODE_VALIDATION_CHECK(this,
-                              // get_input_element_type(5) == ov::element::boolean &&
-                              get_input_shape(5) == ov::Shape({}),
-                              "is_prompt validation failed. ",
-                              "Got element type ",
-                              get_input_element_type(5),
-                              ", shape ",
-                              get_input_shape(5));
-
-        // slot_mapping: shape [batch_size, max_context_len]
-        auto slot_mapping_shape = get_input_partial_shape(6);
-        NODE_VALIDATION_CHECK(this,
-                              // get_input_element_type(6) == ov::element::i64 &&
-                              slot_mapping_shape.size() == 2,
-                              "slot_mapping validation failed. ",
-                              "Got element type ",
-                              get_input_element_type(6),
-                              ", shape ",
-                              slot_mapping_shape);
-
-        // max_context_len: integer scalar
-        NODE_VALIDATION_CHECK(this,
-                              // get_input_element_type(7) == ov::element::i32 &&
-                              get_input_shape(7) == ov::Shape({}),
-                              "max_context_len validation failed. ",
-                              "Got element type ",
-                              get_input_element_type(7),
-                              ", shape ",
-                              get_input_shape(7));
-
-        // context_lens: shape [batch_size]
-        auto context_lens_shape = get_input_partial_shape(8);
-        NODE_VALIDATION_CHECK(this,
-                              // get_input_element_type(8) == ov::element::i32 &&
-                              context_lens_shape.size() == 1,
-                              "context_lens validation failed. ",
-                              "Got element type ",
-                              get_input_element_type(8),
-                              ", shape ",
-                              context_lens_shape);
-
-        // block_tables: shape [batch_size, max_block_per_request]
-        NODE_VALIDATION_CHECK(this,
-                              // get_input_element_type(9) == ov::element::i32 &&
-                              get_input_partial_shape(9).size() == 2,
-                              "block_tables validation failed. ",
-                              "Got element type ",
-                              get_input_element_type(9),
-                              ", shape ",
-                              get_input_partial_shape(9));
-
-        // scale: float scalar
-        NODE_VALIDATION_CHECK(this,
-                              // get_input_element_type(10) == ov::element::f32 &&
-                              get_input_shape(10) == ov::Shape({}),
-                              "block_tables validation failed. ",
-                              "Got element type ",
-                              get_input_element_type(10),
-                              ", shape ",
-                              get_input_shape(10));
-
-        // alibi_slopes: 1D float tensor
-        NODE_VALIDATION_CHECK(this,
-                              // get_input_element_type(11) == ov::element::f32 &&
-                              get_input_partial_shape(11).rank().get_length() == 1,
-                              "alibi_slopes should be a 1D float tensor. ",
-                              "Got element type ",
-                              get_input_element_type(11),
-                              ", shape ",
-                              get_input_partial_shape(11));
-
-        // sliding_window: int scalar
-        NODE_VALIDATION_CHECK(this,
-                              // get_input_element_type(12) == ov::element::i32 &&
-                              get_input_partial_shape(12).rank().get_length() == 0,
-                              "sliding_window argument should be an i32 scalar. ",
-                              "Got element type ",
-                              get_input_element_type(12),
-                              ", shape ",
-                              get_input_partial_shape(12));
-
-        set_output_type(0, query_type, query_shape);
-    }
-
-    std::shared_ptr<ov::Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override {
-        return std::make_shared<PagedAttentionExtension>(new_args);
-    }
-};
-
-}  // namespace
-
 void regclass_graph_op_PagedAttentionExtension(py::module m) {
+    using ov::op::PagedAttentionExtension;
     py::class_<PagedAttentionExtension, std::shared_ptr<PagedAttentionExtension>, ov::Node> cls(
         m,
         "_PagedAttentionExtension");

src/common/transformations/include/transformations/common_optimizations/PositionIDsReplacer.hpp

@@ -0,0 +1,56 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/op/parameter.hpp"
+#include "openvino/op/add.hpp"
+#include "openvino/pass/graph_rewrite.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+#include "openvino/cc/pass/itt.hpp"
+#include "transformations_visibility.hpp"
+#include "transformations/utils/utils.hpp"
+
+namespace ov {
+namespace pass {
+
+class TRANSFORMATIONS_API PositionIDsReplacer;
+
+}  // namespace pass
+}  // namespace ov
+
+using namespace ov::op;
+
+class ov::pass::PositionIDsReplacer : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("PositionIDsReplacer", "0");
+
+    // TODO: Instead of using the following transformation that matches quite a specific place in a model graph in case when position_ids parameter is missing,
+    //        consider replacing always existing attention_mask parameter with a sub-graph using a new slot_mapping parameter.
+    PositionIDsReplacer(const std::shared_ptr<Output<Node>>& position_ids) {
+        MATCHER_SCOPE(PositionIDsReplacer);
+
+        auto input_ids = pattern::any_input();
+        auto input_embed = pattern::wrap_type<v8::Gather>({pattern::any_input(), input_ids, pattern::any_input()});
+
+        auto position_ids_pattern = pattern::any_input();
+        auto offset = pattern::wrap_type<v0::Constant>();
+        auto add_offset = pattern::wrap_type<v1::Add>({position_ids_pattern, offset});
+        auto convert = pattern::wrap_type<v0::Convert>({add_offset});
+        auto position_embed = pattern::wrap_type<v8::Gather>({pattern::any_input(), convert, pattern::any_input()});
+
+        auto add = pattern::wrap_type<v1::Add>({input_embed, position_embed});
+
+        ov::matcher_pass_callback callback = [OV_CAPTURE_CPY_AND_THIS](ov::pass::pattern::Matcher& m) {
+            std::cout << "____" << matcher_name << "___Matched___" << std::endl;
+            const auto& pattern_map = m.get_pattern_value_map();
+            replace_node(pattern_map.at(position_ids_pattern).get_node_shared_ptr(), position_ids->get_node_shared_ptr());
+            std::cout << "APPLIED position_ids PARAMETER INSTEAD OF attention_mask-BASED SUB-GRAPH" << std::endl;
+            return true;
+        };
+
+        auto m = std::make_shared<ov::pass::pattern::Matcher>(add, matcher_name);
+        register_matcher(m, callback);
+    }
+};

src/common/transformations/include/transformations/common_optimizations/PrevSequenceLengthPattern.hpp

@@ -0,0 +1,56 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/pass/graph_rewrite.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+#include "openvino/cc/pass/itt.hpp"
+#include "transformations/utils/utils.hpp"
+#include "transformations_visibility.hpp"
+#include "openvino/op/subtract.hpp"
+#include "openvino/op/shape_of.hpp"
+
+namespace ov {
+namespace pass {
+
+class TRANSFORMATIONS_API PrevSequenceLengthPattern;
+
+}  // namespace pass
+}  // namespace ov
+
+using namespace ov::op;
+
+class ov::pass::PrevSequenceLengthPattern : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("PrevSequenceLengthPattern", "0");
+    PrevSequenceLengthPattern(const std::shared_ptr<ov::op::v1::Subtract>& prev_max_seq_len) {
+        MATCHER_SCOPE(PrevSequenceLengthPattern);
+
+        auto kv_past = pattern::wrap_type<v6::ReadValue>({pattern::any_input()});
+        auto kv_gather = pattern::wrap_type<v8::Gather>({kv_past, pattern::any_input(), pattern::any_input()});
+        auto kv_shape = pattern::wrap_type<v3::ShapeOf>({kv_gather});
+        auto seq = pattern::wrap_type<v8::Gather>({kv_past, pattern::any_input(), pattern::any_input()});
+
+        ov::matcher_pass_callback callback = [OV_CAPTURE_CPY_AND_THIS](ov::pass::pattern::Matcher& m) {
+        std::cout << "____" << matcher_name << "___Matched___" << std::endl;
+            // TODO: Check that seq has axis that really takes sequence len but not any other dimension -- use symbolics or look at the constant input
+            auto gather = m.get_match_root();
+            auto target_type = gather->get_output_element_type(0);
+            std::shared_ptr<Node> replacement;
+            if (prev_max_seq_len->get_output_element_type(0) != target_type) {
+                std::cout << "Converting " << prev_max_seq_len->get_output_element_type(0) << " of max_context_len to " << target_type << std::endl;
+                replacement = std::make_shared<v0::Convert>(prev_max_seq_len, target_type);
+            } else {
+                replacement = prev_max_seq_len;
+            }
+            replace_node(gather, replacement);
+            std::cout << "DETECTED PATTERN PrevSequenceLengthPattern, CONNECTED TO A DEDICATED PARAMETER" << std::endl;
+            return true;
+        };
+
+        auto m = std::make_shared<ov::pass::pattern::Matcher>(seq, matcher_name);
+        register_matcher(m, callback);
+    }
+};

src/common/transformations/include/transformations/common_optimizations/StateManagementPattern.hpp

@@ -0,0 +1,27 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/pass/graph_rewrite.hpp"
+#include "transformations_visibility.hpp"
+
+namespace ov {
+namespace pass {
+
+class TRANSFORMATIONS_API StateManagementPattern;
+
+}  // namespace pass
+}  // namespace ov
+
+class ov::pass::StateManagementPattern : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("StateManagementPattern", "0");
+    StateManagementPattern(ParameterVector& kv_parameters,
+                           const ParameterVector& model_remaining_params,
+                           const std::shared_ptr<ov::op::v0::Constant>& sliding_window,
+                           ParameterVector& parameters_to_remove,
+                           NodeVector& assignes_to_remove,
+                           int& layer_index);
+};