Reorder KV cache using the new gather_by_axis API

Do a ScatterElementsUpdate-based reorder during execution

Get variable update lengths from incoming indices

Make changes to support new KVCache fusion

Add proper include

Author: mdvoretc-intel

onnxruntime/core/providers/openvino/ov_interface.cc

@@ -463,6 +463,26 @@ void StatefulOVInferRequest::PreProcessInferRequest() {
   // TODO(ankit): Address this issue and implement the fix at the appropriate layer.
   FillTensor("beam_idx", ov::element::i32, {1}, 0);
 
+  if (src_idx_val.size() > 0) {
+    ov::Tensor src_idx_tensor = ov::Tensor(ov::element::i32, {src_idx_val.size()});
+    for (int i = 0; i < src_idx_val.size(); ++i) {
+      src_idx_tensor.data<int32_t>()[i] = int32_t(src_idx_val[i]);
+    }
+    ovInfReq.set_tensor("src_idx", src_idx_tensor);
+    ov::Tensor dst_idx_tensor = ov::Tensor(ov::element::i32, {1, 32, dst_idx_val.size(), 96});
+    for (int i = 0; i < dst_idx_val.size(); ++i) {
+      for (int j = 0; j < 32; ++j) {
+        for (int k = 0; k < 96; ++k) {
+          dst_idx_tensor.data<int32_t>()[(j * dst_idx_val.size() + i) * 96 + k] = int32_t(dst_idx_val[i]);
+        }
+      }
+    }
+    ovInfReq.set_tensor("dst_idx", dst_idx_tensor);
+  } else {
+    FillTensor("src_idx", ov::element::i32, {0}, 0);
+    FillTensor("dst_idx", ov::element::i32, {1, 32, 0, 96}, 0);
+  }
+
   // If 'prefill use full chat history' mode is enabled, we need to cache input_ids and position_ids.
   if (prefill_use_full_chat_history) {
     auto input_ids_tensor = ovInfReq.get_tensor("input_ids");
@@ -501,6 +521,38 @@ void StatefulOVInferRequest::Infer() {
   OVInferRequest::Infer();
 }
 
+void StatefulOVInferRequest::ReorderKVCache(const std::vector<size_t>& src_indices, const std::vector<size_t>& dst_indices) {
+  // Validate input parameters
+  if (src_indices.size() != dst_indices.size()) {
+    ORT_THROW(log_tag + "ReorderKVCache: src_indices and dst_indices must have the same size. "
+              "Got src_indices.size()=" + std::to_string(src_indices.size()) +
+              ", dst_indices.size()=" + std::to_string(dst_indices.size()));
+  }
+
+  LOGS_DEFAULT(INFO) << log_tag << "ReorderKVCache: Reordering OpenVINO-internal KVCache state with "
+                     << src_indices.size() << " index pairs";
+
+  // set beam_idx and dst_idx based on provided values
+  src_idx_val.clear();
+  dst_idx_val.clear();
+  for (int i = 0; i < src_indices.size(); ++i) {
+    src_idx_val.emplace_back(src_indices[i]);
+    dst_idx_val.emplace_back(dst_indices[i]);
+  }
+  /*
+  // Retrieve KVCache states and reorder them based on the provided indices
+  auto states = ovInfReq.query_state();
+
+  for (auto& state : states) {
+    auto start_time = std::chrono::high_resolution_clock::now();
+    state.gather_by_axis(src_indices, dst_indices);
+    auto end_time = std::chrono::high_resolution_clock::now();
+    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time).count();
+    LOGS_DEFAULT(INFO) << log_tag << "gather_by_axis: " << duration << " microseconds";
+  }
+  */
+}
+
 void StatefulOVInferRequest::RewindKVCache(size_t index) {
   LOGS_DEFAULT(INFO) << log_tag << "RewindKVCache: Rewinding OpenVINO-internal KVCache state to index=" << index;
 

onnxruntime/core/providers/openvino/ov_interface.h

@@ -161,6 +161,8 @@ class StatefulOVInferRequest : public OVInferRequest {
 
   bool IsNPULogitsSliceRequired();
   bool _npu_logits_slice_required = false;
+  std::vector<int64_t> src_idx_val;
+  std::vector<int64_t> dst_idx_val;
 };
 
 }  // namespace openvino_ep

onnxruntime/core/providers/openvino/ov_stateful_patch_utils.cc

@@ -91,13 +91,26 @@ void FuseCacheReorder(std::shared_ptr<ov::Model> ov_model,
   std::string main_input_name = GetInputOutputName(ov_model, input_name_candidates);
 
   auto input_batch = ov_model->input(main_input_name).get_partial_shape()[0];
+  auto update_shape = ov_model->input(key_value_input_names[0]).get_partial_shape();
 
   auto beam_idx = std::make_shared<ov::opset13::Parameter>(ov::element::i32, ov::PartialShape({std::move(input_batch)}));
   beam_idx->set_friendly_name("beam_idx");
   beam_idx->output(0).get_tensor().add_names({"beam_idx"});
   ov_model->add_parameters({beam_idx});
   not_kv_inputs.push_back(beam_idx->get_friendly_name());
 
+  auto src_idx = std::make_shared<ov::opset13::Parameter>(ov::element::i32, ov::PartialShape({update_shape[2]}));
+  src_idx->set_friendly_name("src_idx");
+  src_idx->output(0).get_tensor().add_names({"src_idx"});
+  ov_model->add_parameters({src_idx});
+  not_kv_inputs.push_back(src_idx->get_friendly_name());
+
+  auto dst_idx = std::make_shared<ov::opset13::Parameter>(ov::element::i32, update_shape);
+  dst_idx->set_friendly_name("dst_idx");
+  dst_idx->output(0).get_tensor().add_names({"dst_idx"});
+  ov_model->add_parameters({dst_idx});
+  not_kv_inputs.push_back(dst_idx->get_friendly_name());
+
   // Go over all cache parameters and fuse _reorder_cache with indices provided by the new parameter beam_idx
   for (const auto& input_name : key_value_input_names) {
     auto parameter_output_port = ov_model->input(input_name);
@@ -108,9 +121,17 @@ void FuseCacheReorder(std::shared_ptr<ov::Model> ov_model,
                                               beam_idx,
                                               ov::opset13::Constant::create(ov::element::i64, {}, {gather_dim}));
 
+    auto update_gather_op =
+        std::make_shared<ov::opset13::Gather>(gather_op,
+                                              src_idx,
+                                              ov::opset13::Constant::create(ov::element::i64, {}, {2}));
+
+    auto update_op = std::make_shared<ov::opset12::ScatterElementsUpdate>(gather_op,
+        dst_idx, update_gather_op, ov::opset13::Constant::create(ov::element::i64, {}, {2}));
+
     // Replace the source output for all consumers of the input tensor
     for (auto& consumer : consumers) {
-      consumer.replace_source_output(gather_op->output(0));
+      consumer.replace_source_output(update_op->output(0));
     }
   }
 

onnxruntime/core/providers/openvino/ov_stateful_patch_utils.h

@@ -13,6 +13,8 @@
 
 #include "openvino/pass/manager.hpp"
 #include "openvino/pass/make_stateful.hpp"
+#include "openvino/opsets/opset3.hpp"
+#include "openvino/opsets/opset12.hpp"
 #include "openvino/opsets/opset13.hpp"
 
 namespace onnxruntime {