[GPU] KVCache custom op and fusion pass (#21689)

src/plugins/intel_gpu/include/intel_gpu/op/kv_cache.hpp

@@ -0,0 +1,64 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/op/op.hpp"
+#include "openvino/op/util/variable.hpp"
+#include "openvino/op/util/variable_extension.hpp"
+
+namespace ov {
+namespace intel_gpu {
+namespace op {
+
+/// \brief Operator that implements Key-Values cache subgraph for large language models.
+/// This operation updates data of the corresponding Variable
+class KVCache : public ov::op::Op, public ov::op::util::VariableExtension {
+public:
+    OPENVINO_OP("KVCache", "gpu_opset");
+
+    KVCache() = default;
+
+    KVCache(const Output<Node>& past,
+            const Output<Node>& new_token_data,
+            const Output<Node>& beam_idx,
+            const std::shared_ptr<ov::op::util::Variable>& past_values,
+            int64_t concat_axis,
+            int64_t gather_axis,
+            const ov::element::Type output_type = ov::element::undefined);
+
+    KVCache(const Output<Node>& past,
+            const Output<Node>& new_token_data,
+            const std::shared_ptr<ov::op::util::Variable>& past_values,
+            int64_t concat_axis,
+            const ov::element::Type output_type = ov::element::undefined);
+
+    bool visit_attributes(ov::AttributeVisitor& visitor) override;
+
+    void validate_and_infer_types() override;
+
+    std::shared_ptr<Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override;
+
+    std::string get_variable_id() const override {
+        OPENVINO_ASSERT(m_variable, "Variable is not initialized. Variable_id is unavailable");
+        return m_variable->get_info().variable_id;
+    }
+
+    int64_t get_concat_axis() const { return m_concat_axis; }
+    int64_t get_gather_axis() const { return m_gather_axis; }
+
+    void set_concat_axis(int64_t axis) { m_concat_axis = axis; }
+    void set_gather_axis(int64_t axis) { m_gather_axis = axis; }
+
+private:
+    int64_t m_concat_axis;
+    int64_t m_gather_axis;
+    ov::element::Type m_output_type;
+};
+
+std::vector<ov::PartialShape> shape_infer(const KVCache* op, std::vector<ov::PartialShape> input_shapes);
+
+}   // namespace op
+}   // namespace intel_gpu
+}   // namespace ov

src/plugins/intel_gpu/include/intel_gpu/op/read_value.hpp

@@ -0,0 +1,40 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/op/op.hpp"
+#include "openvino/op/util/variable.hpp"
+#include "openvino/op/util/variable_extension.hpp"
+
+namespace ov {
+namespace intel_gpu {
+namespace op {
+
+/// \brief Similar to common v6::ReadValue, but it's not derived from ReadValueBase class to avoid ReadValue-Assign pairing check
+/// This is needed to have ReadValue-KVCache pair instead of ReadValue-Assign
+class ReadValue : public ov::op::Op, public ov::op::util::VariableExtension {
+public:
+    OPENVINO_OP("ReadValue", "gpu_opset");
+
+    ReadValue() = default;
+
+    ReadValue(const std::shared_ptr<ov::op::util::Variable>& variable);
+    ReadValue(const Output<Node>& variable_initializer, const std::shared_ptr<ov::op::util::Variable>& variable);
+
+    bool visit_attributes(ov::AttributeVisitor& visitor) override;
+
+    void validate_and_infer_types() override;
+
+    std::shared_ptr<Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override;
+
+    std::string get_variable_id() const override {
+        OPENVINO_ASSERT(m_variable, "Variable is not initialized. Variable_id is unavailable");
+        return m_variable->get_info().variable_id;
+    }
+};
+
+}   // namespace op
+}   // namespace intel_gpu
+}   // namespace ov

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

@@ -273,3 +273,5 @@ REGISTER_FACTORY(internal, FullyConnected);
 REGISTER_FACTORY(internal, FullyConnectedCompressed);
 REGISTER_FACTORY(internal, RMS);
 REGISTER_FACTORY(internal, GatherCompressed);
+REGISTER_FACTORY(internal, KVCache);
+REGISTER_FACTORY(internal, ReadValue);

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

@@ -0,0 +1,72 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+#include "openvino/core/partial_shape.hpp"
+#include "openvino/core/type/element_type.hpp"
+#include "openvino/op/util/variable.hpp"
+#include "primitive.hpp"
+#include <vector>
+
+namespace cldnn {
+
+struct kv_cache : public primitive_base<kv_cache> {
+    CLDNN_DECLARE_PRIMITIVE(kv_cache)
+
+    kv_cache() : primitive_base("", {}) {}
+
+    kv_cache(const primitive_id& id,
+             const std::vector<input_info>& inputs,
+             const ov::op::util::VariableInfo& variable_info,
+             const int64_t concat_axis,
+             const int64_t gather_axis,
+             const padding& output_padding = padding())
+        : primitive_base(id, inputs, {output_padding}), variable_info(variable_info), concat_axis(concat_axis), gather_axis(gather_axis) {}
+
+    ov::op::util::VariableInfo variable_info;
+    int64_t concat_axis = 0;
+    int64_t gather_axis = 0;
+
+    size_t hash() const override {
+        size_t seed = primitive::hash();
+        seed = hash_combine(seed, concat_axis);
+        seed = hash_combine(seed, gather_axis);
+        return seed;
+    }
+
+    bool operator==(const primitive& rhs) const override {
+        if (!compare_common_params(rhs))
+            return false;
+
+        auto rhs_casted = downcast<const kv_cache>(rhs);
+
+        return variable_info == rhs_casted.variable_info &&
+               concat_axis == rhs_casted.concat_axis &&
+               gather_axis == rhs_casted.gather_axis;
+    }
+
+    void save(BinaryOutputBuffer& ob) const override {
+        primitive_base<kv_cache>::save(ob);
+        ov::element::Type_t data_type = variable_info.data_type;
+        ob << variable_info.variable_id;
+        ob << variable_info.data_shape;
+        ob << make_data(&data_type, sizeof(ov::element::Type_t));
+        ob << concat_axis;
+        ob << gather_axis;
+    }
+
+    void load(BinaryInputBuffer& ib) override {
+        primitive_base<kv_cache>::load(ib);
+        ov::PartialShape data_shape;
+        ov::element::Type_t data_type;
+        std::string variable_id;
+        ib >> variable_id;
+        ib >> data_shape;
+        ib >> make_data(&data_type, sizeof(ov::element::Type_t));
+        variable_info = { data_shape, data_type, variable_id };
+        ib >> concat_axis;
+        ib >> gather_axis;
+    }
+};
+}  // namespace cldnn

src/plugins/intel_gpu/src/graph/graph_optimizer/prepare_buffer_fusing.cpp

@@ -3,6 +3,8 @@
 //
 #include "prepare_buffer_fusing.h"
 #include "pooling_inst.h"
+#include "kv_cache_inst.h"
+#include "gather_inst.h"
 #include "primitive_inst.h"
 #include "activation_inst.h"
 #include "concatenation_inst.h"
@@ -447,7 +449,7 @@ void prepare_buffer_fusing::run(program& p) {
         bool is_dynamic = node->is_dynamic();
         bool is_planar = format::is_default_format(node->get_output_layout().format);
         bool no_pad = !node->get_output_layout().data_padding && !node->get_input_layouts().empty() && !node->get_input_layout(0).data_padding;
-        if (node->is_type<read_value>())
+        if (node->is_type<read_value>() || node->is_type<kv_cache>())
             return true;
 
         if (node->is_type<reshape>() && is_dynamic && is_planar && no_pad && !node->is_output() && !node->has_fused_primitives()) {
@@ -620,6 +622,63 @@ void prepare_buffer_fusing::run(program& p) {
 
             node.can_be_optimized(can_reshape_be_optimized(node));
         });
+        program_helpers::do_for_types<kv_cache>(*node, [](kv_cache_node& node) {
+            auto kv_out_layout = node.get_output_layout();
+
+            program_node* rv_prim = nullptr;
+            program_node* gather_prim = nullptr;
+            if (node.get_dependency(0).is_type<read_value>()) {
+                rv_prim = &node.get_dependency(0);
+            } else {
+                if (node.get_dependency(0).is_type<gather>()) {
+                    gather_prim = &node.get_dependency(0);
+                } else {
+                    return;
+                }
+
+                if (gather_prim->get_dependency(0).is_type<read_value>()) {
+                    rv_prim = &gather_prim->get_dependency(0);
+                }
+            }
+
+            if (!rv_prim)
+                return;
+
+            if (kv_out_layout.data_type != rv_prim->get_output_layout().data_type)
+                return;
+
+            auto concat_axis_legacy = node.get_primitive()->concat_axis;
+            if (concat_axis_legacy >= 2) {
+                auto spatial_axis = concat_axis_legacy - 2;
+                // Default and minimum number of dimensions is 4
+                auto spatial_size = std::max<size_t>(kv_out_layout.get_partial_shape().size(), 4) - 2;
+                concat_axis_legacy = spatial_size - spatial_axis - 1 + 2;
+            }
+
+            if (kv_out_layout.is_dynamic()) {
+                // set dynamic pad dims for shape agnostic kernel
+                auto info_dynamic_pad = tensor(0).sizes();
+                info_dynamic_pad[concat_axis_legacy] = 1;
+                auto dynamic_pad_mask = tensor(info_dynamic_pad);
+                kv_out_layout.data_padding.set_dynamic_pad(dynamic_pad_mask);
+                node.set_output_layout(kv_out_layout);
+                node.can_share_buffer(false);
+
+                auto update_dep = [&dynamic_pad_mask](program_node* dep) {
+                    auto prev_layout = dep->get_output_layout();
+                    prev_layout.data_padding.set_dynamic_pad(dynamic_pad_mask);
+                    dep->set_output_layout(prev_layout);
+                    dep->can_share_buffer(false);
+                };
+
+                if (rv_prim) {
+                    update_dep(rv_prim);
+                }
+                if (gather_prim) {
+                    update_dep(gather_prim);
+                }
+            }
+        });
         program_helpers::do_for_types<read_value>(*node, [](read_value_node& node) {
             // Current implementation allows to avoid copy on read_value primitive
             // only in cases when it has single user

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

@@ -0,0 +1,141 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "intel_gpu/runtime/debug_configuration.hpp"
+#include "primitive_base.hpp"
+
+#include "kv_cache_inst.h"
+#include "concatenation/concatenation_kernel_selector.h"
+#include "concatenation/concatenation_kernel_base.h"
+
+namespace cldnn {
+namespace ocl {
+
+namespace {
+kernel_selector::concat_axis convert_axis(int64_t axis, size_t rank) {
+    auto cldnn_axis = axis >= 0 ? axis : axis + static_cast<int64_t>(rank);
+    if (cldnn_axis >= static_cast<int64_t>(rank))
+        OPENVINO_THROW("kv_cache axis exceeds number of dimensions");
+
+    // Difference in dimension ordering between IE and GPU plugin,
+    // reverse spatial dimensions after batch and feature.
+    if (cldnn_axis >= 2) {
+        auto spatial_axis = cldnn_axis - 2;
+        // Default and minimum number of dimensions is 4
+        auto spatial_size = std::max<size_t>(rank, 4) - 2;
+        cldnn_axis = spatial_size - spatial_axis - 1 + 2;
+    }
+
+    switch (cldnn_axis) {
+        case 0: return kernel_selector::concat_axis::BATCH;
+        case 1: return kernel_selector::concat_axis::FEATURE;
+        case 2: return kernel_selector::concat_axis::X;
+        case 3: return kernel_selector::concat_axis::Y;
+        case 4: return kernel_selector::concat_axis::Z;
+        case 5: return kernel_selector::concat_axis::W;
+        default: OPENVINO_THROW("Unsupported kv_cache axis: ", axis);
+    }
+
+    return kernel_selector::concat_axis::FEATURE;  // shouldn't get here
+}
+
+}  // namespace
+
+struct kv_cache_impl : typed_primitive_impl_ocl<kv_cache> {
+    using parent = typed_primitive_impl_ocl<kv_cache>;
+    using parent::parent;
+    using kernel_selector_t = kernel_selector::concatenation_kernel_selector;
+    using kernel_params_t = std::pair<kernel_selector::concatenation_params, kernel_selector::concatenation_optional_params>;
+
+    DECLARE_OBJECT_TYPE_SERIALIZATION(cldnn::ocl::kv_cache_impl)
+
+    std::unique_ptr<primitive_impl> clone() const override {
+        return make_unique<kv_cache_impl>(*this);
+    }
+
+    void load(BinaryInputBuffer& ib) override {
+        parent::load(ib);
+        if (is_dynamic()) {
+            auto& kernel_selector = kernel_selector_t::Instance();
+            auto kernel_impl = kernel_selector.GetImplementation(_kernel_data.kernelName);
+            kernel_impl->GetUpdateDispatchDataFunc(_kernel_data);
+        }
+    }
+
+    kernel_arguments_data get_arguments(const kv_cache_inst& instance) const override {
+        kernel_arguments_data args = parent::get_arguments(instance);
+
+        args.inputs = { instance.input_memory_ptr(0), instance.input_memory_ptr(1) };
+
+        return args;
+    }
+
+    event::ptr execute_impl(const std::vector<event::ptr>& events, kv_cache_inst& instance) override {
+        const bool can_be_optimized = instance.get_impl_params()->_can_be_optimized;
+        const auto& desc = instance.get_typed_desc<kv_cache>();
+        auto& variable = instance.get_network().get_variable(desc->variable_info.variable_id);
+        auto res_event = parent::execute_impl(events, instance);
+        variable.set();
+
+        if (can_be_optimized) {
+            // When primitive is optimized, concat kernel writes directly to variable memory
+            return res_event;
+        } else {
+            // Othwerise, we need to copy result from out buffer to state memory
+            auto& stream = instance.get_network().get_stream();
+
+            stream.enqueue_barrier();
+            return variable.get_memory()->copy_from(stream, instance.output_memory(0), false);
+        }
+    }
+
+    static kernel_params_t get_kernel_params(const kernel_impl_params& impl_param, bool is_shape_agnostic = false) {
+        const auto& primitive = impl_param.typed_desc<kv_cache>();
+        auto params = get_default_params<kernel_selector::concatenation_params>(impl_param, is_shape_agnostic);
+        auto optional_params = get_default_optional_params<kernel_selector::concatenation_optional_params>(impl_param.get_program());
+        auto axis = primitive->concat_axis;
+
+        const auto inputs_count = primitive->input_size();
+        params.inputs.resize(inputs_count);
+        for (size_t i = 0; i < inputs_count; ++i) {
+            params.inputs[i] = convert_data_tensor(impl_param.input_layouts[i]);
+        }
+
+        params.axis = convert_axis(axis, impl_param.get_output_layout().get_rank());
+        optional_params.kernelPerInput = true;
+
+        return {params, optional_params};
+    }
+
+    void update_dispatch_data(const kernel_impl_params& impl_param) override {
+        auto kernel_params = get_kernel_params(impl_param, true);
+        (_kernel_data.update_dispatch_data_func)(kernel_params.first, _kernel_data);
+        _kernel_data.kernels[0].skip_execution = impl_param._can_be_optimized || impl_param.get_input_layout(0).count() == 0;
+    }
+};
+
+namespace detail {
+
+attach_kv_cache_impl::attach_kv_cache_impl() {
+    auto types = { data_types::f16, data_types::f32 };
+    auto formats = { format::bfyx };
+    implementation_map<kv_cache>::add(impl_types::ocl,
+                                           shape_types::dynamic_shape,
+                                           typed_primitive_impl_ocl<kv_cache>::create<kv_cache_impl>,
+                                           types,
+                                           formats);
+
+    implementation_map<kv_cache>::add(impl_types::ocl,
+                                           shape_types::static_shape,
+                                           typed_primitive_impl_ocl<kv_cache>::create<kv_cache_impl>,
+                                           types,
+                                           formats);
+}
+
+}  // namespace detail
+}  // namespace ocl
+}  // namespace cldnn
+
+BIND_BINARY_BUFFER_WITH_TYPE(cldnn::ocl::kv_cache_impl)
+BIND_BINARY_BUFFER_WITH_TYPE(cldnn::kv_cache)

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

@@ -40,6 +40,7 @@ void register_implementations() {
     REGISTER_OCL(generate_proposals);
     REGISTER_OCL(grid_sample);
     REGISTER_OCL(group_normalization);
+    REGISTER_OCL(kv_cache);
     REGISTER_OCL(lrn);
     REGISTER_OCL(lstm_gemm);
     REGISTER_OCL(lstm_elt);

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

@@ -77,6 +77,7 @@
 #include "intel_gpu/primitives/non_zero.hpp"
 #include "intel_gpu/primitives/eye.hpp"
 #include "intel_gpu/primitives/unique.hpp"
+#include "intel_gpu/primitives/kv_cache.hpp"
 
 namespace cldnn {
 namespace ocl {
@@ -122,6 +123,7 @@ REGISTER_OCL(gemm);
 REGISTER_OCL(generate_proposals);
 REGISTER_OCL(grid_sample);
 REGISTER_OCL(group_normalization);
+REGISTER_OCL(kv_cache);
 REGISTER_OCL(lrn);
 REGISTER_OCL(lstm_gemm);
 REGISTER_OCL(lstm_elt);

src/plugins/intel_gpu/src/graph/include/kv_cache_inst.h

@@ -0,0 +1,45 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "intel_gpu/primitives/kv_cache.hpp"
+#include "primitive_inst.h"
+#include "variable.hpp"
+
+namespace cldnn {
+
+template <>
+struct typed_program_node<kv_cache> : public typed_program_node_base<kv_cache> {
+private:
+    using parent = typed_program_node_base<kv_cache>;
+
+public:
+    using parent::parent;
+
+    program_node& input() const { return get_dependency(0); }
+
+    std::vector<size_t> get_shape_infer_dependencies() const override { return {}; }
+};
+
+using kv_cache_node = typed_program_node<kv_cache>;
+
+template<>
+class typed_primitive_inst<kv_cache> : public typed_primitive_inst_base<kv_cache>, public memory_state::variable {
+    using parent = typed_primitive_inst_base<kv_cache>;
+
+public:
+    template<typename ShapeType>
+    static std::vector<layout> calc_output_layouts(kv_cache_node const& /*node*/, const kernel_impl_params& impl_param);
+    static layout calc_output_layout(const kv_cache_node& node, kernel_impl_params const& impl_param);
+
+    static std::string to_string(const kv_cache_node& node);
+
+    typed_primitive_inst(network& network, const kv_cache_node& desc);
+    typed_primitive_inst(network& network) : parent(network), memory_state::variable("") {}
+};
+
+using kv_cache_inst = typed_primitive_inst<kv_cache>;
+
+} // namespace cldnn

src/plugins/intel_gpu/src/graph/include/primitive_inst.h

@@ -232,6 +232,7 @@ public:
     void do_runtime_skip_reorder();
     void do_runtime_skip_gather();
     void do_runtime_in_place_concat();
+    void do_runtime_in_place_kv_cache();
     void configure_shape_of_dependencies();
 
     memory::ptr fused_memory(size_t dep_id) const {

src/plugins/intel_gpu/src/graph/kv_cache.cpp

@@ -0,0 +1,58 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "intel_gpu/op/kv_cache.hpp"
+#include "intel_gpu/plugin/common_utils.hpp"
+#include "intel_gpu/runtime/optionals.hpp"
+#include "kv_cache_inst.h"
+#include "primitive_type_base.h"
+#include <sstream>
+#include <json_object.h>
+
+namespace cldnn {
+GPU_DEFINE_PRIMITIVE_TYPE_ID(kv_cache)
+
+kv_cache_inst::typed_primitive_inst(network& network, const kv_cache_node& node) :
+    parent{network, node, false},
+    memory_state::variable{node.get_primitive()->variable_info.variable_id} {
+}
+
+layout kv_cache_inst::calc_output_layout(const kv_cache_node& node, kernel_impl_params const& impl_param) {
+    return impl_param.input_layouts[0];
+}
+
+template<typename ShapeType>
+std::vector<layout> kv_cache_inst::calc_output_layouts(kv_cache_node const& /*node*/, kernel_impl_params const& impl_param) {
+    auto desc = impl_param.typed_desc<kv_cache>();
+    auto output_data_type = desc->output_data_types[0].value_or(impl_param.get_input_layout().data_type);
+
+    ov::intel_gpu::op::KVCache op;
+    op.set_concat_axis(desc->concat_axis);
+    op.set_gather_axis(desc->gather_axis);
+
+    std::vector<ShapeType> input_shapes = {impl_param.get_input_layout(0).get<ShapeType>(), impl_param.get_input_layout(1).get<ShapeType>()};
+
+    std::vector<ShapeType> output_shapes = shape_infer(&op, input_shapes);
+
+    return {layout({output_shapes[0], output_data_type, impl_param.get_output_layout().format})};
+}
+
+template std::vector<layout> kv_cache_inst::calc_output_layouts<ov::PartialShape>(kv_cache_node const& node, const kernel_impl_params& impl_param);
+
+std::string kv_cache_inst::to_string(const kv_cache_node& node) {
+    auto node_info = node.desc_to_json();
+    json_composite kv_cache_info;
+    kv_cache_info.add("input id", node.input().id());
+    kv_cache_info.add("variable id", node.get_primitive()->variable_info.variable_id);
+    kv_cache_info.add("variable shape", node.get_primitive()->variable_info.data_shape);
+    kv_cache_info.add("variable type", node.get_primitive()->variable_info.data_type);
+    kv_cache_info.add("concat axis", node.get_primitive()->concat_axis);
+    kv_cache_info.add("gather axis", node.get_primitive()->gather_axis);
+    node_info->add("kv_cache info", kv_cache_info);
+    std::stringstream primitive_description;
+    node_info->dump(primitive_description);
+    return primitive_description.str();
+}
+
+} // namespace cldnn

src/plugins/intel_gpu/src/graph/primitive_inst.cpp

@@ -25,6 +25,7 @@
 #include "gemm_inst.h"
 #include "assign_inst.h"
 #include "read_value_inst.h"
+#include "kv_cache_inst.h"
 #include "condition_inst.h"
 #include "gather_inst.h"
 #include "experimental_detectron_roi_feature_extractor_inst.hpp"
@@ -425,7 +426,14 @@ void primitive_inst::update_shape() {
     if (get_node().is_type<assign>()) {
         auto desc = get_node().as<assign>().get_primitive();
         get_network().get_variable(desc->variable_id).set_layout(_impl_params->get_output_layout());
-        _impl_params->state_layout = _impl_params->get_output_layout();
+    }
+
+    if (get_node().is_type<read_value>()) {
+        auto desc = get_node().as<read_value>().get_primitive();
+        auto& variable = get_network().get_variable(desc->variable_id);
+        auto variable_layout = variable.get_layout();
+        // Custom output layout update as update_output_layout handles paddings incorrectly for optimized out read_value + kv_cache pattern
+        _impl_params->output_layouts[0] = variable_layout;
     }
 }
 
@@ -460,8 +468,20 @@ event::ptr primitive_inst::realloc_if_needed() {
     if (auto stateful_prim = dynamic_cast<memory_state::variable*>(this)) {
         std::string variable_id = stateful_prim->variable_id();
         auto& variable = get_network().get_variable(variable_id);
+        GPU_DEBUG_TRACE_DETAIL << "realloc_if_needed: variable " << id() << " set layout" << _impl_params->get_output_layout().to_string() << std::endl;
+        if (_node->is_type<kv_cache>()) {
+            // Reuse state memory as output for kv cache if possible
+            // otherwise clear _outputs for the cases when mem was reused previously
+            if (_impl_params->can_be_optimized()) {
+                _outputs[0] = variable.get_memory();
+                return ev;
+            } else if (_outputs[0] && variable.get_memory() && get_network().get_engine().is_the_same_buffer(*_outputs[0], *variable.get_memory())) {
+                _outputs[0] = nullptr;
+                max_output_layout_size = 0;
+            }
+        }
         variable.set_layout(actual_layout);
-        GPU_DEBUG_TRACE_DETAIL << id() << ": use variable memory " << variable.get_memory()
+        GPU_DEBUG_TRACE_DETAIL << id() << ": use variable memory " << variable.get_memory()->buffer_ptr()
                                << " (size=" << variable.get_memory()->size() << ")" << std::endl;
         // For nodes that can be optimized, variable memory is used as output memory
         // so there is no need for output memory reallocation
@@ -488,6 +508,19 @@ event::ptr primitive_inst::realloc_if_needed() {
         }
     }
 
+    // Clear out memory if if was previously reused, but now primitive can't be optimized
+    if (_node->is_type<gather>() && !can_be_optimized() && _outputs[0] && _network.get_engine().is_the_same_buffer(dep_memory(0), output_memory(0))) {
+        _outputs[0] = nullptr;
+        max_output_layout_size = 0;
+    }
+
+    // update layout to ensure that it repsects paddings for correct allocation size
+    if (_node->is_type<kv_cache>() && !_impl_params->can_be_optimized()) {
+        const auto current_buf_size = updated_layout.get_buffer_size().sizes();
+        ov::Shape current_shape(current_buf_size.begin(), current_buf_size.end());
+        updated_layout.set_partial_shape(current_shape);
+    }
+
     bool can_reuse_buffer = _outputs[0] && updated_layout.count() <= max_output_layout_size;
 
     // Handle runtime dynamic concat optimization
@@ -790,6 +823,95 @@ void primitive_inst::do_runtime_skip_reorder() {
     }
 }
 
+void primitive_inst::do_runtime_in_place_kv_cache() {
+    OV_ITT_SCOPED_TASK(ov::intel_gpu::itt::domains::intel_gpu_plugin, openvino::itt::handle("do_runtime_in_place_kv_cache: " + id()));
+    auto reset_pad = [](kernel_impl_params& params, const program_node* node) {
+        params.output_layouts[0].data_padding = node->get_output_layout(0).data_padding;
+    };
+    if (_node->is_type<read_value>()) {
+        auto& variable = get_network().get_variable(_node->as<read_value>().get_primitive()->variable_id);
+        // Reset paddings for read_value and users with dynamic pad when variable is reset
+        // to avoid wrong pad used for some nodes due to pad propagation logic (which uses previous iter pad values)
+        if (!variable.is_set()) {
+            primitive_inst* inst = this;
+            while (inst) {
+                reset_pad(*inst->_impl_params, inst->_node);
+                auto& users = inst->_node->get_users();
+                if (users.size() == 1 && users.front()->get_output_layout(0).data_padding.get_dynamic_pad_dims() != tensor(0)) {
+                    inst = inst->get_user_insts().front().get();
+                } else {
+                    inst = nullptr;
+                }
+            }
+        }
+        return;
+    }
+
+    if (_node->is_type<gather>() && _impl_params->output_layouts[0].data_padding.get_dynamic_pad_dims() != tensor(0)) {
+        if (can_be_optimized())
+            _impl_params->output_layouts[0] = _impl_params->input_layouts[0];
+        else
+            reset_pad(*_impl_params, _node);
+        return;
+    }
+
+    if (!_node->is_type<kv_cache>())
+        return;
+
+    _impl_params->_can_be_optimized = false;
+    if (_impl_params->get_input_layout(0).count() == 0) {
+        return;
+    }
+
+    auto desc = _node->as<kv_cache>().get_primitive();
+    auto& past_layout = _impl_params->input_layouts[0];
+    auto& present_layout = _impl_params->output_layouts[0];
+    const auto& sequence_axis = desc->concat_axis;
+
+    auto sequence_axis_legacy = sequence_axis;
+    if (sequence_axis_legacy >= 2) {
+        auto spatial_axis = sequence_axis_legacy - 2;
+        // Default and minimum number of dimensions is 4
+        auto spatial_size = std::max<size_t>(past_layout.get_partial_shape().size(), 4) - 2;
+        sequence_axis_legacy = spatial_size - spatial_axis - 1 + 2;
+    }
+
+    if (present_layout.data_padding.get_dynamic_pad_dims().sizes()[sequence_axis_legacy] != 1)
+        return;
+
+    const size_t total_elements = past_layout.count();
+    const int64_t concat_axis_size = past_layout.get_partial_shape()[sequence_axis].get_length();
+    const int64_t sequence_element_size = total_elements / concat_axis_size;
+
+    const int64_t max_sequence_elements = _deps[0].first->max_output_layout_size / sequence_element_size;
+    const int64_t max_pad = std::max<int64_t>(max_sequence_elements - concat_axis_size, 0);
+
+    if (max_pad > 0) {
+        auto update_pad = [&](layout& l, int64_t pad) {
+            const auto& dyn_pad_dims = l.data_padding.get_dynamic_pad_dims();
+            const auto& lower_padd = l.data_padding.lower_size().sizes();
+            auto upper_padd = l.data_padding.upper_size().sizes();
+            upper_padd[sequence_axis_legacy] = pad;
+            l.data_padding = padding(lower_padd, upper_padd, 0.f, dyn_pad_dims);
+        };
+
+        update_pad(present_layout, max_pad - 1);
+        auto& variable = get_network().get_variable(desc->variable_info.variable_id);
+        GPU_DEBUG_TRACE_DETAIL << "do_runtime_in_place_kv_cache set_layout: " << present_layout.to_string() << " is_set  = " << variable.is_set() << std::endl;
+        variable.set_layout(present_layout);
+        if (past_layout.data_padding.upper_size().sizes()[sequence_axis_legacy] > 0 && variable.is_set()) {
+            update_pad(past_layout, max_pad);
+            _impl_params->_can_be_optimized = true;
+        }
+        GPU_DEBUG_TRACE_DETAIL << "[do runtime kv_cache opt] concat_axis_size = " << concat_axis_size << std::endl;
+        GPU_DEBUG_TRACE_DETAIL << "[do runtime kv_cache opt] sequence_element_size = " << sequence_element_size << std::endl;
+        GPU_DEBUG_TRACE_DETAIL << "[do runtime kv_cache opt] max_sequence_elements = " << max_sequence_elements << std::endl;
+        GPU_DEBUG_TRACE_DETAIL << "[do runtime kv_cache opt] max_pad = " << max_pad << std::endl;
+        GPU_DEBUG_TRACE_DETAIL << "[do runtime kv_cache opt] can be optimized: " << _impl_params->_can_be_optimized << std::endl;
+        GPU_DEBUG_TRACE_DETAIL << "[do runtime kv_cache opt] " << present_layout.to_string() << std::endl;
+    }
+}
+
 void primitive_inst::do_runtime_skip_gather() {
     // Check pattern
     if (!get_node().is_type<gather>()
@@ -932,6 +1054,7 @@ event::ptr primitive_inst::execute(const std::vector<event::ptr>& events) {
         // if the user is can_be_optimized and output node then current nodes' output should be allocated to host.
         do_runtime_skip_reorder();
         do_runtime_skip_gather();
+        do_runtime_in_place_kv_cache();
         if (_impl_params->output_layouts[0].count() == 0) {
             GPU_DEBUG_TRACE_DETAIL << id() << " : Skipping because output data is empty " << std::endl;
             auto ev = get_network().get_stream().create_user_event(true);

src/plugins/intel_gpu/src/plugin/ops/kv_cache.cpp

@@ -0,0 +1,40 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "intel_gpu/op/kv_cache.hpp"
+#include "intel_gpu/plugin/program_builder.hpp"
+#include "intel_gpu/plugin/common_utils.hpp"
+#include "intel_gpu/primitives/kv_cache.hpp"
+
+namespace ov {
+namespace op {
+namespace internal {
+using KVCache = ov::intel_gpu::op::KVCache;
+}  // namespace internal
+}  // namespace op
+}  // namespace ov
+
+namespace ov {
+namespace intel_gpu {
+
+namespace {
+
+void CreateKVCacheOp(ProgramBuilder& p, const std::shared_ptr<ov::op::internal::KVCache>& op) {
+    validate_inputs_count(op, {2, 3});
+    auto inputs = p.GetInputInfo(op);
+    const auto prim = cldnn::kv_cache(layer_type_name_ID(op),
+                                      inputs,
+                                      op->get_variable()->get_info(),
+                                      op->get_concat_axis(),
+                                      op->get_gather_axis());
+
+    p.add_primitive(*op, prim);
+}
+
+} // namespace
+
+REGISTER_FACTORY_IMPL(internal, KVCache);
+
+}  // namespace intel_gpu
+}  // namespace ov

src/plugins/intel_gpu/src/plugin/ops/variable.cpp

@@ -8,9 +8,18 @@
 #include "openvino/op/assign.hpp"
 #include "openvino/op/read_value.hpp"
 #include "transformations/rt_info/original_precision_attribute.hpp"
+#include "intel_gpu/op/read_value.hpp"
 #include "intel_gpu/primitives/assign.hpp"
 #include "intel_gpu/primitives/read_value.hpp"
 
+namespace ov {
+namespace op {
+namespace internal {
+using ReadValue = ov::intel_gpu::op::ReadValue;
+}  // namespace internal
+}  // namespace op
+}  // namespace ov
+
 
 namespace ov {
 namespace intel_gpu {
@@ -41,6 +50,11 @@ void CreateReadValueOp(ProgramBuilder& p, const std::shared_ptr<ov::op::v3::Read
     CreateVariableAccessPrimitive<cldnn::read_value>(p, op, op->get_variable_id());
 }
 
+void CreateReadValueOp(ProgramBuilder& p, const std::shared_ptr<ov::intel_gpu::op::ReadValue>& op) {
+    validate_inputs_count(op, {0, 1});
+    CreateVariableAccessPrimitive<cldnn::read_value>(p, op, op->get_variable_id());
+}
+
 void CreateAssignOp(ProgramBuilder& p, const std::shared_ptr<ov::op::v3::Assign>& op) {
     validate_inputs_count(op, {1});
     CreateVariableAccessPrimitive<cldnn::assign>(p, op, op->get_variable_id());
@@ -62,6 +76,7 @@ REGISTER_FACTORY_IMPL(v3, Assign);
 REGISTER_FACTORY_IMPL(v6, Assign);
 REGISTER_FACTORY_IMPL(v3, ReadValue);
 REGISTER_FACTORY_IMPL(v6, ReadValue);
+REGISTER_FACTORY_IMPL(internal, ReadValue);
 
 }  // namespace intel_gpu
 }  // namespace ov

src/plugins/intel_gpu/src/plugin/transformations/kv_cache_fusion.cpp

@@ -0,0 +1,144 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "kv_cache_fusion.hpp"
+#include <memory>
+
+#include "intel_gpu/op/kv_cache.hpp"
+
+#include "intel_gpu/op/read_value.hpp"
+#include "intel_gpu/plugin/common_utils.hpp"
+#include "openvino/core/node_vector.hpp"
+#include "openvino/core/rt_info.hpp"
+#include "openvino/op/constant.hpp"
+#include "openvino/op/convert.hpp"
+#include "openvino/op/parameter.hpp"
+#include "openvino/op/read_value.hpp"
+#include "openvino/op/sink.hpp"
+#include "openvino/pass/graph_rewrite.hpp"
+#include "openvino/pass/pattern/op/label.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+#include "openvino/pass/pattern/op/or.hpp"
+#include "openvino/pass/visualize_tree.hpp"
+#include "transformations/utils/utils.hpp"
+
+namespace ov {
+namespace intel_gpu {
+
+class KVCacheFusionMatcher : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("KVCacheFusionMatcher", "0");
+    KVCacheFusionMatcher();
+};
+
+KVCacheFusionMatcher::KVCacheFusionMatcher() {
+    using namespace ov::pass::pattern;
+
+    auto past = wrap_type<ov::op::v6::ReadValue>();
+    auto convert_past = wrap_type<ov::op::v0::Convert>({past});
+    auto gather_input = std::make_shared<ov::pass::pattern::op::Or>(OutputVector{past, convert_past});
+    auto beam_idx = wrap_type<ov::op::v0::Parameter>();
+    auto gather_past = wrap_type<ov::op::v8::Gather>({gather_input, beam_idx, wrap_type<ov::op::v0::Constant>()});
+    auto concat_past_input = std::make_shared<ov::pass::pattern::op::Or>(OutputVector{past, convert_past, gather_past});
+    auto concat = wrap_type<ov::op::v0::Concat>({concat_past_input, any_input()});
+    auto convert_present = wrap_type<ov::op::v0::Convert>({concat});
+    auto present_input = std::make_shared<ov::pass::pattern::op::Or>(OutputVector{concat, convert_present});
+    auto present = wrap_type<ov::op::v6::Assign>({present_input});
+
+    ov::matcher_pass_callback callback = [=](ov::pass::pattern::Matcher& m) {
+        if (transformation_callback(m.get_match_root())) {
+            return false;
+        }
+
+        const auto& pattern_map = m.get_pattern_value_map();
+        auto concat_node = std::dynamic_pointer_cast<ov::op::v0::Concat>(pattern_map.at(concat).get_node_shared_ptr());
+
+        auto past_node = std::dynamic_pointer_cast<ov::op::v6::ReadValue>(pattern_map.at(past).get_node_shared_ptr());
+        auto present_node = std::dynamic_pointer_cast<ov::op::v6::Assign>(pattern_map.at(present).get_node_shared_ptr());
+
+        if (past_node->get_variable_id() != present_node->get_variable_id())
+            return false;
+
+        // TODO: Support conversion internally
+        if (concat_node->get_output_element_type(0) != past_node->get_output_element_type(0))
+            return false;
+
+        auto variable = past_node->get_variable();
+        auto concat_axis = concat_node->get_axis();
+
+        std::shared_ptr<ov::Node> variable_initializer = nullptr;
+        std::shared_ptr<ov::Node> kv_cache_node = nullptr;
+        if (past_node->get_input_size() == 1) {
+            variable_initializer = past_node->get_input_node_shared_ptr(0);
+        }
+
+        // Replace common ReadValue op with a custom one as common one expects paired Assign operation which is removed by this transform
+        auto new_read_value_node = variable_initializer ? std::make_shared<ov::intel_gpu::op::ReadValue>(variable_initializer, variable)
+                                                        : std::make_shared<ov::intel_gpu::op::ReadValue>(variable);
+        new_read_value_node->set_friendly_name(past_node->get_friendly_name());
+        ov::copy_runtime_info(past_node, new_read_value_node);
+        ov::replace_node(past_node, new_read_value_node);
+
+        if (pattern_map.count(gather_past) > 0) {
+            // TODO: Enable code below once KVCache custom op supports rearrange internally
+            // For now Gather is kept as standalone op
+            #if 0
+            auto gather_node = pattern_map.at(gather_past).get_node_shared_ptr();
+            auto gather_axis = std::dynamic_pointer_cast<ov::op::v0::Constant>(gather_node->get_input_node_shared_ptr(2))->cast_vector<int64_t>()[0];
+            kv_cache_node = std::make_shared<op::KVCache>(new_read_value_node,
+                                                          concat_node->get_input_node_shared_ptr(1),
+                                                          gather_node->get_input_node_shared_ptr(1),
+                                                          variable,
+                                                          concat_axis,
+                                                          gather_axis,
+                                                          concat_node->get_output_element_type(0));
+            #else
+            kv_cache_node = std::make_shared<op::KVCache>(pattern_map.at(gather_past).get_node_shared_ptr(),
+                                                          concat_node->get_input_node_shared_ptr(1),
+                                                          variable,
+                                                          concat_axis,
+                                                          new_read_value_node->get_output_element_type(0));
+            #endif
+        } else {
+            kv_cache_node = std::make_shared<op::KVCache>(new_read_value_node,
+                                                          concat_node->get_input_node_shared_ptr(1),
+                                                          variable,
+                                                          concat_axis,
+                                                          new_read_value_node->get_output_element_type(0));
+        }
+        kv_cache_node->set_friendly_name(concat_node->get_friendly_name());
+        ov::copy_runtime_info(m.get_matched_nodes(), kv_cache_node);
+        ov::replace_node(concat_node, kv_cache_node);
+
+        if (pattern_map.count(convert_present) > 0) {
+            present_node->set_argument(0, kv_cache_node->output(0));
+        }
+
+        return true;
+    };
+
+    auto m = std::make_shared<ov::pass::pattern::Matcher>(present, "KVCacheFusionMatcher");
+    this->register_matcher(m, callback);
+}
+
+bool KVCacheFusion::run_on_model(const std::shared_ptr<ov::Model>& m) {
+    bool res = pass::GraphRewrite::run_on_model(m);
+    if (res) {
+        ov::SinkVector sinks = m->get_sinks();
+        for (auto& sink : sinks) {
+            if (sink && sink->get_input_node_ptr(0)->get_type_info() == op::KVCache::get_type_info_static()) {
+                m->remove_sink(sink);
+            }
+        }
+    }
+
+    return res;
+}
+
+KVCacheFusion::KVCacheFusion() {
+    add_matcher<ov::intel_gpu::KVCacheFusionMatcher>();
+}
+
+}  // namespace intel_gpu
+}  // namespace ov

src/plugins/intel_gpu/src/plugin/transformations/kv_cache_fusion.hpp

@@ -0,0 +1,124 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/pass/graph_rewrite.hpp"
+
+namespace ov {
+namespace intel_gpu {
+
+/// 1. Trivial case (greedy search, no state initializer)
+///     ┌───────────┐      ┌───────────┐                                      ┌───────────┐
+///     │ ReadValue │      │  SomeOp   │                                      │  SomeOp   │
+///     | (past_kv) |      | new_token |                                      | new_token |
+///     └─────┬─────┘      └─────┬─────┘                                      └─────┬─────┘
+///           │                  │                                                  |
+///           │                  |                                                  |
+///           │   ┌────────┐     │                                             ┌────┴────┐       ┌──────────┐
+///           └───┤ Concat ├─────┘                    =>                       | KVCache |.......| Variable |
+///               └───┬────┘                                                   └────┬────┘       └──────────┘
+///                   │                                                             |
+///        ┌──────────┴────────────┐                                                |
+///   ┌────┴────────┐         ┌────┴──────┐                                    ┌────┴────┐
+///   │  Assign     │         │  SomeOp   │                                    | SomeOp  |
+///   | (present_kv |         |  (SDPA)   |                                    | (SDPA)  |
+///   └─────────────┘         └───────────┘                                    └─────────┘
+
+/// 2. With gather for beam search (or model which supports both greedy and beam search)
+///     ┌───────────┐      ┌────────────┐    ┌───────────┐                    ┌───────────┐     ┌────────────┐
+///     │ ReadValue │      │ Parameter  │    │  SomeOp   │                    │  SomeOp   │     │ Parameter  │
+///     | (past_kv) |      |  beam_idx  |    | new_token |                    | new_token |     |  beam_idx  |
+///     └─────┬─────┘      └─────┬──────┘    └─────┬─────┘                    └─────┬─────┘     └─────┬──────┘
+///           │                  │                 │                                |                 |
+///     ┌─────┴──────┐           |                 │                                |                 |
+///     |   Gather   |───────────┘                 │                                |                 |
+///     └─────┬──────┘                             │                                |                 |
+///           |                                    │                                |                 |
+///           |                                    │                                |                 |
+///           │                                    │                                |                 |
+///           │   ┌────────┐                       │                                |            ┌────┴────┐      ┌──────────┐
+///           └───┤ Concat ├───────────────────────┘        =>                      └────────────┤ KVCache |......| Variable |
+///               └───┬────┘                                                                     └────┬────┘      └──────────┘
+///                   │                                                                               |
+///        ┌──────────┴────────────┐                                                                  |
+///   ┌────┴────────┐         ┌────┴──────┐                                                      ┌────┴────┐
+///   │  Assign     │         │  SomeOp   │                                                      | SomeOp  |
+///   | (present_kv |         |  (SDPA)   |                                                      | (SDPA)  |
+///   └─────────────┘         └───────────┘                                                      └─────────┘
+
+/// 3. Similar to case 2, but with variable initializer
+///     ┌─────────────────────┐                                                                                  ┌─────────────────────┐
+///     │       SomeOp        │                                                                                  │       SomeOp        │
+///     | (state initializer) |                                                                                  | (state initializer) |
+///     └─────┬───────────────┘                                                                                  └─────┬───────────────┘
+///           |                                                                                                        |
+///     ┌─────┴─────┐      ┌────────────┐    ┌───────────┐                    ┌───────────┐     ┌────────────┐         |
+///     │ ReadValue │      │ Parameter  │    │  SomeOp   │                    │  SomeOp   │     │ Parameter  │         |
+///     | (past_kv) |      |  beam_idx  |    | new_token |                    | new_token |     |  beam_idx  |         |
+///     └─────┬─────┘      └─────┬──────┘    └─────┬─────┘                    └─────┬─────┘     └─────┬──────┘         |
+///           │                  │                 │                                |                 |      ┌─────────┘
+///     ┌─────┴──────┐           |                 │                                |                 |      |
+///     |   Gather   |───────────┘                 │                                |                 |      |
+///     └─────┬──────┘                             │                                |                 |      |
+///           |                                    │                                |                 |      |
+///           |                                    │                                |                 |      |
+///           │                                    │                                |                 |      |
+///           │   ┌────────┐                       │                                |            ┌────┴──────┴───────┐                       ┌──────────┐
+///           └───┤ Concat ├───────────────────────┘        =>                      └────────────┤      KVCache      |.......................| Variable |
+///               └───┬────┘                                                                     └────┬──────────────┘                       └──────────┘
+///                   │                                                                               |
+///        ┌──────────┴────────────┐                                                                  |
+///   ┌────┴────────┐         ┌────┴──────┐                                                      ┌────┴────┐
+///   │  Assign     │         │  SomeOp   │                                                      | SomeOp  |
+///   | (present_kv |         |  (SDPA)   |                                                      | (SDPA)  |
+///   └─────────────┘         └───────────┘                                                      └─────────┘
+
+/// 4. Pattern may also include Convert for state variable
+///     ┌─────────────────────┐                                                                                  ┌─────────────────────┐
+///     │       SomeOp        │                                                                                  │       SomeOp        │
+///     | (state initializer) |                                                                                  | (state initializer) |
+///     └─────┬───────────────┘                                                                                  └─────┬───────────────┘
+///           |                                                                                                        |
+///     ┌─────┴─────┐      ┌────────────┐    ┌───────────┐                    ┌───────────┐     ┌────────────┐         |
+///     │ ReadValue │      │ Parameter  │    │  SomeOp   │                    │  SomeOp   │     │ Parameter  │         |
+///     | (past_kv) |      |  beam_idx  |    | new_token |                    | new_token |     |  beam_idx  |         |
+///     └─────┬─────┘      └─────┬──────┘    └─────┬─────┘                    └─────┬─────┘     └─────┬──────┘         |
+///           │                  │                 │                                |                 |      ┌─────────┘
+///     ┌─────┴────┐             │                 │                                |                 |      |
+///     │  Convert │             │                 │                                |                 |      |
+///     └─────┬────┘             │                 │                                |                 |      |
+///           │                  │                 │                                |                 |      |
+///           │                  │                 │                                |                 |      |
+///     ┌─────┴──────┐           |                 │                                |                 |      |
+///     |   Gather   |───────────┘                 │                                |                 |      |
+///     └─────┬──────┘                             │                                |                 |      |
+///           |                                    │                                |                 |      |
+///           |                                    │                                |                 |      |
+///           │                                    │                                |                 |      |
+///           │   ┌────────┐                       │                                |            ┌────┴──────┴───────┐                       ┌──────────┐
+///           └───┤ Concat ├───────────────────────┘        =>                      └────────────┤      KVCache      |.......................| Variable |
+///               └───┬────┘                                                                     └────┬──────────────┘                       └──────────┘
+///                   │                                                                               |
+///        ┌──────────┴────────────┐                                                                  |
+///        |                       |                                                                  |
+///  ┌─────┴────┐                  |                                                                  |
+///  │  Convert │                  |                                                                  |
+///  └─────┬────┘                  |                                                                  |
+///        |                       |                                                                  |
+///   ┌────┴────────┐         ┌────┴──────┐                                                      ┌────┴────┐
+///   │  Assign     │         │  SomeOp   │                                                      | SomeOp  |
+///   | (present_kv |         |  (SDPA)   |                                                      | (SDPA)  |
+///   └─────────────┘         └───────────┘                                                      └─────────┘
+class KVCacheFusion : public ov::pass::GraphRewrite {
+public:
+    OPENVINO_RTTI("KVCacheFusion", "0");
+    KVCacheFusion();
+
+    bool run_on_model(const std::shared_ptr<ov::Model>& m) override;
+};
+
+
+}   // namespace intel_gpu
+}   // namespace ov

src/plugins/intel_gpu/src/plugin/transformations/op/kv_cache.cpp

@@ -0,0 +1,78 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "intel_gpu/op/kv_cache.hpp"
+#include "concat_shape_inference.hpp"
+#include "openvino/core/partial_shape.hpp"
+#include "openvino/op/concat.hpp"
+
+namespace ov {
+namespace intel_gpu {
+namespace op {
+
+KVCache::KVCache(const Output<Node>& past,
+                 const Output<Node>& new_token_data,
+                 const Output<Node>& beam_idx,
+                 const std::shared_ptr<ov::op::util::Variable>& past_variable,
+                 int64_t concat_axis,
+                 int64_t gather_axis,
+                 const ov::element::Type output_type)
+    : Op({past, new_token_data, beam_idx}), m_concat_axis(concat_axis), m_gather_axis(gather_axis), m_output_type(output_type) {
+    m_variable = past_variable;
+    validate_and_infer_types();
+}
+
+KVCache::KVCache(const Output<Node>& past,
+                 const Output<Node>& new_token_data,
+                 const std::shared_ptr<ov::op::util::Variable>& past_variable,
+                 int64_t concat_axis,
+                 const ov::element::Type output_type)
+    : Op({past, new_token_data}), m_concat_axis(concat_axis), m_gather_axis(0), m_output_type(output_type) {
+    m_variable = past_variable;
+    validate_and_infer_types();
+}
+
+bool KVCache::visit_attributes(ov::AttributeVisitor& visitor) {
+    visitor.on_attribute("concat_axis", m_concat_axis);
+    visitor.on_attribute("gather_axis", m_gather_axis);
+    visitor.on_attribute("output_type", m_output_type);
+    return true;
+}
+
+void KVCache::validate_and_infer_types() {
+    auto output_type = m_output_type == ov::element::undefined ? get_input_element_type(0) : m_output_type;
+    const std::vector<ov::PartialShape> input_shapes = {get_input_partial_shape(0), get_input_partial_shape(1)};
+
+    set_output_type(0, output_type, shape_infer(this, input_shapes)[0]);
+}
+
+std::shared_ptr<Node> KVCache::clone_with_new_inputs(const ov::OutputVector& new_args) const {
+    check_new_args_count(this, new_args);
+    if (new_args.size() == 2) {
+        return std::make_shared<KVCache>(new_args.at(0),
+                                         new_args.at(1),
+                                         m_variable,
+                                         m_concat_axis,
+                                         m_output_type);
+
+    } else {
+        return std::make_shared<KVCache>(new_args.at(0),
+                                         new_args.at(1),
+                                         new_args.at(2),
+                                         m_variable,
+                                         m_concat_axis,
+                                         m_gather_axis,
+                                         m_output_type);
+    }
+}
+
+std::vector<ov::PartialShape> shape_infer(const KVCache* op, std::vector<ov::PartialShape> input_shapes) {
+    ov::op::v0::Concat concat;
+    concat.set_axis(op->get_concat_axis());
+    return ov::op::v0::shape_infer(&concat, input_shapes);
+}
+
+}  // namespace op
+}  // namespace intel_gpu
+}  // namespace ov

src/plugins/intel_gpu/src/plugin/transformations/op/read_value.cpp

@@ -0,0 +1,94 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "intel_gpu/op/read_value.hpp"
+#include "openvino/core/partial_shape.hpp"
+
+namespace ov {
+namespace intel_gpu {
+namespace op {
+
+ReadValue::ReadValue(const std::shared_ptr<ov::op::util::Variable>& past_values) : Op() {
+    m_variable = past_values;
+    validate_and_infer_types();
+}
+ReadValue::ReadValue(const Output<Node>& variable_initializer, const std::shared_ptr<ov::op::util::Variable>& past_values) : Op({variable_initializer}) {
+    m_variable = past_values;
+    validate_and_infer_types();
+}
+
+bool ReadValue::visit_attributes(ov::AttributeVisitor& visitor) {
+    visitor.on_attribute("variable_id", m_variable);
+
+    auto variable_info = m_variable->get_info();
+    visitor.on_attribute("variable_type", variable_info.data_type);
+    visitor.on_attribute("variable_shape", variable_info.data_shape);
+    m_variable->update(variable_info);
+    return true;
+}
+
+void ReadValue::validate_and_infer_types() {
+    OPENVINO_ASSERT(m_variable, "Variable is not initialized.");
+    const auto& variable_info = m_variable->get_info();
+    const auto& variable_type = variable_info.data_type;
+    const auto& variable_shape = variable_info.data_shape;
+
+    // If no inputs provided, it means this ReadValue doesn't have initial subgraph. This is valid.
+    if (get_input_size() > 0) {
+        const auto& initial_type = get_input_element_type(0);
+        const auto& initial_shape = get_input_partial_shape(0);
+
+        // Variable's shape/type determine a permissible range of values for shape/type inferred from initial_subgraph.
+        // If initial_subgraph is set, then we need to check that shape/type inferred from initial_subgraph
+        // is within the permissible range.
+
+        const auto compatible_type = variable_type.is_dynamic() || initial_type == variable_type;
+        const auto compatible_shape = variable_shape.relaxes(initial_shape);
+
+        OPENVINO_ASSERT(compatible_shape,
+                        "The shape specified in the Variable has to extend (relax) the shape "
+                        "inferred from the initializing subgraph.",
+                        " Variable shape: ",
+                        variable_shape,
+                        " Initialization shape: ",
+                        initial_shape);
+        OPENVINO_ASSERT(compatible_type,
+                        "The type specified in the Variable has to extend (relax) the type "
+                        "inferred from the initializing subgraph.",
+                        " Variable type: ",
+                        variable_type,
+                        " Initialization type: ",
+                        initial_type);
+        // workaround:
+        // dynamic rank/type can be derived from the IRs generated via the prev versions of OV,
+        // but dynamic rank/type are not supported in plugins,
+        // so we are trying to fix them here using the rank/type of ReadValue 1st input, if it exists
+        if (get_input_size() > 0 && variable_info.data_shape.rank().is_dynamic() &&
+            variable_info.data_type.is_dynamic()) {
+            set_output_type(0, initial_type, initial_shape);
+            return;
+        }
+    }
+
+    set_output_type(0, variable_type, variable_shape);
+}
+
+std::shared_ptr<Node> ReadValue::clone_with_new_inputs(const ov::OutputVector& new_args) const {
+    check_new_args_count(this, new_args);
+        switch (new_args.size()) {
+    case 0:
+        return std::make_shared<ReadValue>(m_variable);
+    case 1:
+        return std::make_shared<ReadValue>(new_args[0], m_variable);
+    default:
+        OPENVINO_THROW("Unable to clone ReadValue ",
+                       this->get_friendly_name(),
+                       " Incorrect number of inputs. Expected: 0 or 1. Actual: ",
+                       new_args.size());
+    }
+}
+
+}  // namespace op
+}  // namespace intel_gpu
+}  // namespace ov

src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp

@@ -121,6 +121,7 @@
 #include "plugin/transformations/rms_fusion.hpp"
 #include "plugin/transformations/binary_conv_to_conv.hpp"
 #include "plugin/transformations/move_convert_after_gather.hpp"
+#include "plugin/transformations/kv_cache_fusion.hpp"
 
 #include "transformations/low_precision/mark_dequantization_subgraph.hpp"
 #include "low_precision/pull_reshape_through_dequantization.hpp"
@@ -703,6 +704,7 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
         manager.register_pass<ov::intel_gpu::ConvertFullyConnectedToFullyConnectedCompressed>();
         manager.register_pass<ov::intel_gpu::ConvertGatherToGatherCompressed>();
         manager.register_pass<ov::intel_gpu::RMSFusion>();
+        manager.register_pass<ov::intel_gpu::KVCacheFusion>();
 
         // This is supposed to be the last pass to ensure that we don't have name collisions until
         // GPU plugin stops using friendly names for program creation

src/plugins/intel_gpu/src/plugin/variable_state.cpp

@@ -74,7 +74,8 @@ void VariableState::update_device_buffer() {
 
     if (actual_size < m_layout.bytes_count()) {
         const auto alloc_type = m_context->get_engine().use_unified_shared_memory() ? cldnn::allocation_type::usm_device : cldnn::allocation_type::cl_mem;
-        const auto current_shape = get_tensor_shape(m_layout.get_partial_shape());
+        const auto current_buf_size = m_layout.get_buffer_size().sizes();
+        ov::Shape current_shape(current_buf_size.begin(), current_buf_size.end());
         const auto alloc_shape = predict_shape(m_name, current_shape, m_layout.data_type, *m_shape_predictor);
         const auto alloc_layout = cldnn::layout(alloc_shape, m_layout.data_type, m_layout.format);
         m_memory = m_context->get_engine().allocate_memory(alloc_layout, alloc_type, false);

src/plugins/intel_gpu/tests/unit/dynamic_execution/stateful_model.cpp

@@ -9,6 +9,9 @@
 #include <intel_gpu/primitives/data.hpp>
 #include <intel_gpu/primitives/gather.hpp>
 #include <intel_gpu/primitives/concatenation.hpp>
+#include <intel_gpu/primitives/kv_cache.hpp>
+#include <intel_gpu/primitives/read_value.hpp>
+#include <intel_gpu/runtime/internal_properties.hpp>
 
 #include "program_wrapper.h"
 
@@ -126,4 +129,47 @@ TEST(stateful_model, not_skip_gather_at_runtime) {
     network.execute();
     ASSERT_EQ(gather_inst->can_be_optimized(), false);
 }
+
+TEST(stateful_model, check_dynamic_pad_for_kv_cache) {
+    auto& engine = get_test_engine();
+
+    auto input_beam_idx_lay = layout{ov::PartialShape{-1}, data_types::i32, format::bfyx};
+    auto input_present_lay = layout{ov::PartialShape{-1, 32, -1, 128}, data_types::f32, format::bfyx};
+
+    ov::op::util::VariableInfo info{ov::PartialShape{-1, 32, -1, 128}, data_types::f32, "v0"};
+    auto input_kv_lay = layout{info.data_shape, info.data_type, format::bfyx};
+    topology topology(input_layout("beam_idx", input_beam_idx_lay),
+                      input_layout("present", input_present_lay),
+                      read_value("kv_cache", std::vector<input_info>{}, info.variable_id, input_kv_lay),
+                      gather("gather",
+                             input_info("kv_cache"),
+                             input_info("beam_idx"),
+                             0,                                       // axis
+                             input_kv_lay.get_partial_shape().size(), // input rank
+                             ov::Shape{},                             // output shape
+                             0,                                       // batch_dim
+                             true),                                   // support_neg_ind
+                      kv_cache("concat", {input_info("gather"), input_info("present")}, info, 0, 0),
+                      reorder("reorder", input_info("concat"), format::bfyx, data_types::f32)); /*output padding*/
+
+    ExecutionConfig config = get_test_default_config(engine);
+    config.set_property(ov::intel_gpu::allow_new_shape_infer(true));
+    config.set_property(ov::intel_gpu::optimize_data(true));
+
+    network network(engine, topology, config);
+    auto gather_inst = network.get_primitive("gather");
+    auto read_value_inst = network.get_primitive("kv_cache");
+    auto kv_cache_inst = network.get_primitive("concat");
+
+    ASSERT_EQ(gather_inst->get_node().can_be_optimized(), true);
+    ASSERT_EQ(gather_inst->can_be_optimized(), true);
+
+    auto pad = tensor(0);
+    pad.batch[0] = 1;
+
+    ASSERT_EQ(read_value_inst->get_output_layout(0).data_padding.get_dynamic_pad_dims(), pad);
+    ASSERT_EQ(gather_inst->get_output_layout(0).data_padding.get_dynamic_pad_dims(), pad);
+    ASSERT_EQ(kv_cache_inst->get_output_layout(0).data_padding.get_dynamic_pad_dims(), pad);
+}
+
 }  // stateful_model_tests

src/plugins/intel_gpu/tests/unit/shape_infer/kv_cache_si_test.cpp

@@ -0,0 +1,85 @@
+// Copyright (C) 2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "test_utils.h"
+
+#include <intel_gpu/primitives/input_layout.hpp>
+#include <intel_gpu/primitives/kv_cache.hpp>
+#include <intel_gpu/primitives/data.hpp>
+
+#include "kv_cache_inst.h"
+
+#include "program_wrapper.h"
+
+#include <cmath>
+#include <algorithm>
+
+using namespace cldnn;
+using namespace ::tests;
+
+namespace shape_infer_tests {
+
+struct kv_cache_test_params {
+    std::vector<layout> input_layouts;
+    int64_t concat_axis;
+    int64_t gather_axis;
+    layout expected_layout;
+};
+
+class kv_cache_test : public testing::TestWithParam<kv_cache_test_params> { };
+
+TEST_P(kv_cache_test, shape_infer) {
+    auto p = GetParam();
+
+    auto& engine = get_test_engine();
+
+    cldnn::program prog(engine);
+    std::vector<std::shared_ptr<input_layout>> input_prims;
+    std::vector<input_info> input_prims_ids;
+    for (size_t i = 0; i < p.input_layouts.size(); i++) {
+        auto prim_id = "data" + std::to_string(i);
+        auto data_layout_prim = std::make_shared<input_layout>(prim_id, p.input_layouts[i]);
+        input_prims.push_back(data_layout_prim);
+        input_prims_ids.push_back(input_info(prim_id));
+    }
+
+    ov::op::util::VariableInfo info{p.input_layouts[0].get_partial_shape(), p.input_layouts[0].data_type, "v0"};
+
+    auto kv_cache_prim = std::make_shared<kv_cache>("output", input_prims_ids, info, p.concat_axis, p.gather_axis);
+    auto& kv_cache_node = prog.get_or_create(kv_cache_prim);
+    for (size_t i = 0; i < p.input_layouts.size(); i++) {
+        auto& input_node = prog.get_or_create(input_prims[i]);
+        program_wrapper::add_connection(prog, input_node, kv_cache_node);
+    }
+
+    auto params = kv_cache_node.get_kernel_impl_params();
+    auto res = kv_cache_inst::calc_output_layouts<ov::PartialShape>(kv_cache_node, *params);
+
+    ASSERT_EQ(res.size(), 1);
+    ASSERT_EQ(res[0], p.expected_layout);
+}
+
+INSTANTIATE_TEST_SUITE_P(smoke, kv_cache_test,
+    testing::ValuesIn(std::vector<kv_cache_test_params>{
+        {
+            {
+                layout{ov::PartialShape{-1, 2, -1, 4}, data_types::f32, format::bfyx},
+                layout{ov::PartialShape{-1, 2, -1, 4}, data_types::f32, format::bfyx},
+            },
+            2,
+            0,
+            layout{ov::PartialShape{-1, 2, -1, 4}, data_types::f32, format::bfyx}
+        },
+        {
+            {
+                layout{ov::PartialShape{1, 2, 0, 4}, data_types::f16, format::bfyx},
+                layout{ov::PartialShape{1, 2, 10, 4}, data_types::f16, format::bfyx},
+            },
+            2,
+            0,
+            layout{ov::PartialShape{1, 2, 10, 4}, data_types::f16, format::bfyx}
+        },
+    }));
+
+}  // shape_infer_tests

src/plugins/intel_gpu/tests/unit/transformations/kv_cache_fusion_test.cpp

@@ -0,0 +1,78 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <gtest/gtest.h>
+
+#include <string>
+#include <memory>
+
+#include <openvino/core/model.hpp>
+#include <openvino/pass/manager.hpp>
+#include "common_test_utils/ov_test_utils.hpp"
+#include <transformations/utils/utils.hpp>
+
+#include <plugin/transformations/kv_cache_fusion.hpp>
+
+#include "intel_gpu/op/read_value.hpp"
+#include "openvino/core/node_vector.hpp"
+#include "openvino/core/partial_shape.hpp"
+#include "openvino/op/parameter.hpp"
+#include "openvino/op/result.hpp"
+#include "openvino/op/sink.hpp"
+#include "intel_gpu/op/kv_cache.hpp"
+
+using namespace testing;
+using namespace ov::intel_gpu;
+
+TEST_F(TransformationTestsF, KVCacheFusionTest1) {
+    {
+        auto variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{1, 32, -1, 80}, ov::element::f16, "v0"});
+        auto past = std::make_shared<ov::op::v6::ReadValue>(variable);
+        auto parameter = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::PartialShape{1, 32, -1, 80});
+        auto concat = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{past, parameter}, 2);
+        auto present = std::make_shared<ov::op::v6::Assign>(concat, variable);
+        auto result = std::make_shared<ov::op::v0::Result>(concat);
+
+        model = std::make_shared<ov::Model>(ov::ResultVector{result}, ov::SinkVector{present}, ov::ParameterVector{parameter});
+        manager.register_pass<KVCacheFusion>();
+    }
+    {
+        auto variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{1, 32, -1, 80}, ov::element::f16, "v0"});
+        auto parameter = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::PartialShape{1, 32, -1, 80});
+        auto past = std::make_shared<ov::intel_gpu::op::ReadValue>(variable);
+        auto kv_cache = std::make_shared<ov::intel_gpu::op::KVCache>(past, parameter, variable, 2, ov::element::f16);
+        auto result = std::make_shared<ov::op::v0::Result>(kv_cache);
+
+        model_ref = std::make_shared<ov::Model>(ov::ResultVector{result}, ov::ParameterVector{parameter});
+    }
+}
+
+TEST_F(TransformationTestsF, KVCacheFusionTest2) {
+    {
+        auto variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{1, 32, -1, 80}, ov::element::f32, "v0"});
+        auto past = std::make_shared<ov::op::v6::ReadValue>(variable);
+        auto parameter = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1, 32, -1, 80});
+        auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});
+        auto axis = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 0);
+        auto gather_past = std::make_shared<ov::op::v8::Gather>(past, beam_idx, axis);
+        auto concat = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{gather_past, parameter}, 2);
+        auto present = std::make_shared<ov::op::v6::Assign>(concat, variable);
+        auto result = std::make_shared<ov::op::v0::Result>(concat);
+
+        model = std::make_shared<ov::Model>(ov::ResultVector{result}, ov::SinkVector{present}, ov::ParameterVector{parameter, beam_idx});
+        manager.register_pass<KVCacheFusion>();
+    }
+    {
+        auto variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{1, 32, -1, 80}, ov::element::f32, "v0"});
+        auto parameter = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1, 32, -1, 80});
+        auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});
+        auto past = std::make_shared<ov::intel_gpu::op::ReadValue>(variable);
+        auto axis = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 0);
+        auto gather_past = std::make_shared<ov::op::v8::Gather>(past, beam_idx, axis);
+        auto kv_cache = std::make_shared<ov::intel_gpu::op::KVCache>(gather_past, parameter, variable, 2, ov::element::f32);
+        auto result = std::make_shared<ov::op::v0::Result>(kv_cache);
+
+        model_ref = std::make_shared<ov::Model>(ov::ResultVector{result}, ov::ParameterVector{parameter, beam_idx});
+    }
+}