Paged attention

Author: sshlyapn

src/plugins/intel_gpu/include/intel_gpu/plugin/program_builder.hpp

@@ -172,6 +172,7 @@ class ProgramBuilder final {
     void CreateSingleLayerPrimitive(cldnn::topology& topology, const std::shared_ptr<ov::Node>& op);
 };
 
+void CreatePagedAttention(ProgramBuilder& p, const std::shared_ptr<ov::Node>& op);
 void CreateCustomOp(ProgramBuilder& p, const std::shared_ptr<ov::Node>& node, CustomLayerPtr customLayer);
 void CreateUnaryEltwiseOp(ProgramBuilder& p, const std::shared_ptr<ov::Node>& node,
                           cldnn::activation_func func, cldnn::activation_additional_params params);

src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp

@@ -0,0 +1,53 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+#include "primitive.hpp"
+#include "intel_gpu/graph/program.hpp"
+
+#include <vector>
+
+namespace cldnn {
+
+struct paged_attention : public primitive_base<paged_attention> {
+    CLDNN_DECLARE_PRIMITIVE(paged_attention)
+
+    paged_attention() : primitive_base("", {}) {}
+
+    paged_attention(const primitive_id& id,
+                    const std::vector<input_info>& inputs,
+                    const padding& output_padding = padding())
+        : primitive_base(id, inputs, {output_padding}) {
+            OPENVINO_ASSERT(inputs.size() == 13, "[GPU] Unexpected inputs number for PagedAttention primitive: ", inputs.size());
+    }
+
+    bool operator==(const primitive& rhs) const override {
+        return compare_common_params(rhs);
+    }
+
+    void save(BinaryOutputBuffer& ob) const override {
+        primitive_base<paged_attention>::save(ob);
+        ob << head_size;
+        ob << heads_num;
+        ob << kv_heads_num;
+        ob << block_size;
+        ob << x_block_size;
+    }
+
+    void load(BinaryInputBuffer& ib) override {
+        primitive_base<paged_attention>::load(ib);
+        ib >> head_size;
+        ib >> heads_num;
+        ib >> kv_heads_num;
+        ib >> block_size;
+        ib >> x_block_size;
+    }
+
+    size_t head_size;
+    size_t heads_num;
+    size_t kv_heads_num;
+    size_t block_size;
+    size_t x_block_size;
+};
+}  // namespace cldnn

src/plugins/intel_gpu/src/graph/impls/ocl/gemm.cpp

@@ -33,6 +33,12 @@ struct gemm_impl : multi_stage_primitive<gemm> {
         return make_unique<gemm_impl>(*this);
     }
 
+    gemm_impl() = default;
+
+    gemm_impl(const std::vector<kernel_selector::kernel_data>& kd) : parent(kd) {
+        this->can_reuse_memory = true;
+    }
+
     void load(BinaryInputBuffer& ib) override {
         parent::load(ib);
         if (is_dynamic()) {

src/plugins/intel_gpu/src/graph/impls/ocl/multi_stage_primitive.hpp

@@ -53,6 +53,7 @@ struct multi_stage_primitive : public typed_primitive_impl<PType> {
         }
         this->can_reuse_memory = false;
         this->_kernel_name = other._kernel_name;
+        this->can_reuse_memory = other.can_reuse_memory;
         this->_is_dynamic = other._is_dynamic;
     }