add reorder KV cache API

Author: Kotomi-Du

onnxruntime/core/providers/openvino/backend_manager.cc

@@ -892,5 +892,11 @@ void BackendManager::RewindKVCache(size_t index) {
   }
 }
 
+void BackendManager::ReorderKVCache(const std::vector<size_t>& src_indices, const std::vector<size_t>& dst_indices) {
+  if (concrete_backend_) {
+    concrete_backend_->ReorderKVCache(src_indices, dst_indices);
+  }
+}
+
 }  // namespace openvino_ep
 }  // namespace onnxruntime

onnxruntime/core/providers/openvino/backend_manager.h

@@ -31,6 +31,7 @@ class BackendManager {
   void TryExportCompiledBlobAsEPCtxNode(const onnxruntime::GraphViewer& subgraph, bool include_embed_data);
   ov::CompiledModel GetOVCompiledModel();
   void RewindKVCache(size_t index);
+  void ReorderKVCache(const std::vector<size_t>& src_indices, const std::vector<size_t>& dst_indices);
 
  private:
   std::unique_ptr<ONNX_NAMESPACE::ModelProto> GetModelProtoFromFusedNode(

onnxruntime/core/providers/openvino/backends/basic_backend.cc

@@ -334,6 +334,12 @@ void BasicBackend::RewindKVCache(size_t index) {
   });
 }
 
+void BasicBackend::ReorderKVCache(const std::vector<size_t>& src_indices, const std::vector<size_t>& dst_indices) {
+  infer_req_pool_->forEachIdleRequest([&](OVInferRequestPtr& infer_request) {
+    infer_request->ReorderKVCache(src_indices, dst_indices);
+  });
+}
+
 void BasicBackend::Infer(OrtKernelContext* ctx) const {
   Ort::KernelContext context(ctx);
 

onnxruntime/core/providers/openvino/backends/basic_backend.h

@@ -138,6 +138,7 @@ class BasicBackend : public IBackend {
     return exe_network_.Get();
   }
   void RewindKVCache(size_t index) override;
+  void ReorderKVCache(const std::vector<size_t>& src_indices, const std::vector<size_t>& dst_indices) override;
 
  private:
   bool ValidateSubgraph(std::map<std::string, std::shared_ptr<ov::Node>>& const_outputs_map);

onnxruntime/core/providers/openvino/ibackend.h

@@ -18,6 +18,7 @@ class IBackend {
   virtual ov::CompiledModel GetOVCompiledModel() = 0;
   virtual ~IBackend() = default;
   virtual void RewindKVCache(size_t index) {}
+  virtual void ReorderKVCache(const std::vector<size_t>& src_indices, const std::vector<size_t>& dst_indices) {}
 };
 using ptr_stream_t = std::unique_ptr<ModelBlobWrapper>;
 class BackendFactory {

onnxruntime/core/providers/openvino/openvino_execution_provider.cc

@@ -288,6 +288,68 @@ common::Status OpenVINOExecutionProvider::SetEpDynamicOptions(gsl::span<const ch
           LOGS_DEFAULT(WARNING) << "kvcache_rewind index is < 0:\t" << index;
         }
       }
+    } else if (key == "kvcache_reorder") {
+    // Convert kvcache_reorder value format "1,2,3;4,5,6" into two vectors
+      // src_indices = [1,2,3], dst_indices = [4,5,6]
+      size_t delimiter_pos = value.find(';');
+      if (delimiter_pos == std::string::npos) {
+        LOGS_DEFAULT(WARNING) << "kvcache_reorder value format is incorrect, expected format is 'x1,x2,x3;y1,y2,y3' where x and y are comma-separated int64_t lists";
+        return Status::OK();
+      }
+
+      std::string src_string = value.substr(0, delimiter_pos);
+      std::string dst_string = value.substr(delimiter_pos + 1);
+
+      std::vector<size_t> src_indices;
+      std::vector<size_t> dst_indices;
+
+      try {
+        // Parse source indices from comma-separated string
+        std::stringstream src_stream(src_string);
+        std::string src_token;
+        while (std::getline(src_stream, src_token, ',')) {
+          // Trim whitespace
+          src_token.erase(0, src_token.find_first_not_of(" \t"));
+          src_token.erase(src_token.find_last_not_of(" \t") + 1);
+
+          if (!src_token.empty()) {
+            int64_t index = std::stoll(src_token);
+            if (index >= 0) {
+              src_indices.push_back(static_cast<size_t>(index));
+            } else {
+              LOGS_DEFAULT(WARNING) << "kvcache_reorder src_index is < 0: " << index;
+            }
+          }
+        }
+
+        // Parse destination indices from comma-separated string
+        std::stringstream dst_stream(dst_string);
+        std::string dst_token;
+        while (std::getline(dst_stream, dst_token, ',')) {
+          // Trim whitespace
+          dst_token.erase(0, dst_token.find_first_not_of(" \t"));
+          dst_token.erase(dst_token.find_last_not_of(" \t") + 1);
+
+          if (!dst_token.empty()) {
+            int64_t index = std::stoll(dst_token);
+            if (index >= 0) {
+              dst_indices.push_back(static_cast<size_t>(index));
+            } else {
+              LOGS_DEFAULT(WARNING) << "kvcache_reorder dst_index is < 0: " << index;
+            }
+          }
+        }
+
+      } catch (const std::exception& e) {
+        LOGS_DEFAULT(WARNING) << "Conversion for kvcache_reorder string value to int64_t indices failed. "
+                              << "Exception: " << e.what();
+        return Status::OK();
+      }
+
+      // Trigger KVCache Reorder for target Backend with vector arguments
+      for (auto& backend : backend_managers_) {
+        backend.ReorderKVCache(src_indices, dst_indices);
+      }
     } else {
       // Handle unknown options
       LOGS_DEFAULT(WARNING) << "Unknown key/value pair - ignoring " << key << "/" << value;

onnxruntime/core/providers/openvino/ov_interface.h

@@ -132,6 +132,7 @@ class OVInferRequest {
     return ovInfReq;
   }
   virtual void RewindKVCache([[maybe_unused]] size_t index) {}
+  virtual void ReorderKVCache([[maybe_unused]] const std::vector<size_t>& src_indices,  [[maybe_unused]] const std::vector<size_t>& dst_indices) {}
 };
 
 class StatefulOVInferRequest : public OVInferRequest {
@@ -140,6 +141,7 @@ class StatefulOVInferRequest : public OVInferRequest {
 
   void Infer() override;
   void RewindKVCache(size_t index) override;
+  void ReorderKVCache(const std::vector<size_t>& src_indices, const std::vector<size_t>& dst_indices) override;
   void FillTensor(const std::string& tensor_name, const ov::element::Type& type,
                   const std::vector<size_t>& shape, int32_t fill_value);
   void CacheTensor(const std::string& tensor_name, std::vector<int64_t>& cache);