Support KVCache rewind for stateful LLMs via SetEpDynamicOptions

Author: RyanMetcalfeInt8

onnxruntime/core/providers/openvino/backend_manager.cc

@@ -550,5 +550,11 @@ void BackendManager::ShutdownBackendManager() {
   concrete_backend_.reset();
 }
 
+void BackendManager::RewindKVCache(size_t index) {
+  if (concrete_backend_) {
+    concrete_backend_->RewindKVCache(index);
+  }
+}
+
 }  // namespace openvino_ep
 }  // namespace onnxruntime

onnxruntime/core/providers/openvino/backend_manager.h

@@ -30,6 +30,7 @@ class BackendManager {
   SessionContext& GetSessionContext();
   Status ExportCompiledBlobAsEPCtxNode(const onnxruntime::GraphViewer& subgraph);
   ov::CompiledModel& GetOVCompiledModel();
+  void RewindKVCache(size_t index);
 
  private:
   std::unique_ptr<ONNX_NAMESPACE::ModelProto> GetModelProtoFromFusedNode(

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

@@ -358,6 +358,13 @@ void BasicBackend::SetNumThreads(ov::AnyMap& device_config) {
     device_config.emplace(ov::inference_num_threads(session_context_.num_of_threads));
 }
 
+void BasicBackend::RewindKVCache(size_t index) {
+  OVInferRequestPtr infer_request;
+  infer_request = inferRequestsQueue_->getIdleRequest();
+  infer_request->RewindKVCache(index);
+  inferRequestsQueue_->putIdleRequest(std::move(infer_request));
+}
+
 // Starts an asynchronous inference request for data in slice indexed by batch_slice_idx on
 // an Infer Request indexed by infer_req_idx
 void BasicBackend::StartAsyncInference(Ort::KernelContext& context, OVInferRequestPtr infer_request) {

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

@@ -41,6 +41,7 @@ class BasicBackend : public IBackend {
   ov::CompiledModel& GetOVCompiledModel() override {
     return exe_network_.Get();
   }
+  void RewindKVCache(size_t index) override;
 
  private:
   void PopulateCompiledDirectory(std::string, std::string&, std::string&, bool&);

onnxruntime/core/providers/openvino/ibackend.h

@@ -17,6 +17,7 @@ class IBackend {
   virtual void Infer(OrtKernelContext* context) = 0;
   virtual ov::CompiledModel& GetOVCompiledModel() = 0;
   virtual ~IBackend() = default;
+  virtual void RewindKVCache(size_t index) {};
 };
 using ptr_stream_t = std::unique_ptr<std::istream>;
 class BackendFactory {

onnxruntime/core/providers/openvino/openvino_execution_provider.cc

@@ -294,6 +294,24 @@ common::Status OpenVINOExecutionProvider::SetEpDynamicOptions(gsl::span<const ch
           ov_compiled_model.set_property(ov::workload_type(workload_type));
         }
       }
+    } else if (key == "kvcache_rewind") {
+      // convert kvcache_rewind value to int64_t
+      int64_t index;
+      try {
+        index = std::stoll(value);
+      } catch (const std::exception& e) {
+        LOGS_DEFAULT(WARNING) << "Could not convert kvcache_rewind value string to index. Exception: " + std::string(e.what());
+        return Status::OK();
+      }
+
+      // Trigger KVCache rewind for backed
+      for (auto& backend : backend_managers_) {
+        if (index >= 0) {
+          backend.RewindKVCache(static_cast<size_t>(index));
+        } else {
+          LOGS_DEFAULT(WARNING) << "kvcache_rewind index is < 0: " << index;
+        }
+      }
     } else {
       // Handle unknown options
       LOGS_DEFAULT(WARNING) << "Unknown key/value pair - ignoring " << key << "/" << value;

onnxruntime/core/providers/openvino/ov_interface.cc

@@ -105,7 +105,7 @@ OVExeNetwork OVCore::StatefulCompileModel(std::shared_ptr<OVNetwork>& model,
 
   if (hw_target.find("NPU") != std::string::npos) {
     KVDesc kv_desc;
-    kv_desc.max_prompt_len = PopIntAndCast(config, "MAX_PROMPT_LEN").value_or(3072u);
+    kv_desc.max_prompt_len = PopIntAndCast(config, "MAX_PROMPT_LEN").value_or(1024u);
     kv_desc.min_response_len = PopIntAndCast(config, "MIN_RESPONSE_LEN").value_or(128u);
 
     if (onnxruntime::openvino_ep::backend_utils::IsDebugEnabled()) {
@@ -488,5 +488,50 @@ void StatefulOVInferRequest::Infer() {
   OVInferRequest::Infer();
 }
 
+void StatefulOVInferRequest::RewindKVCache(size_t index) {
+  if (device == "NPU") {
+    std::cout << "RewindKVCache on NPU: Trimming cached input_ids / position_ids to length "
+              << index << std::endl;
+    if (cached_input_ids.size() > index) {
+      cached_input_ids.resize(index);
+    }
+
+    if (cached_position_ids.size() > index) {
+      cached_position_ids.resize(index);
+    }
+  } else {
+    std::cout << "OVInferRequest::RewindKVCache: Trimming internal states to length = "
+              << index << std::endl;
+    if (index == 0) {
+      // in this case, since we're trimming *all* of the KVCache, just reset the state.
+      ovInfReq.reset_state();
+    } else {
+      // retrieve kvcache states, and trim...
+      // Most of this code was grabbed from here:
+      // https://github.com/openvinotoolkit/openvino.genai/blob/releases/2025/1/src/cpp/src/utils.cpp#L329
+      auto states = ovInfReq.query_state();
+      for (auto& state : states) {
+        ov::Tensor old_tensor = state.get_state();
+        // [BATCH_SIZE, num_kv_heads, seq_len, head_size]
+        auto shape = old_tensor.get_shape();
+
+        if (shape[2] > index) {
+          shape[2] = index;
+
+          ov::Coordinate new_shape_begin{0, 0, 0, 0};
+          ov::Coordinate new_shape_end{shape};
+
+          auto trimmed_tensor = ov::Tensor(old_tensor, new_shape_begin, new_shape_end);
+
+          ov::Tensor new_tensor(old_tensor.get_element_type(), shape);
+          trimmed_tensor.copy_to(new_tensor);
+
+          state.set_state(new_tensor);
+        }
+      }
+    }
+  }
+}
+
 }  // namespace openvino_ep
 }  // namespace onnxruntime

onnxruntime/core/providers/openvino/ov_interface.h

@@ -133,6 +133,7 @@ class OVInferRequest {
   ov::InferRequest& GetNewObj() {
     return ovInfReq;
   }
+  virtual void RewindKVCache(size_t index) {};
 };
 
 class StatefulOVInferRequest : public OVInferRequest {
@@ -141,6 +142,7 @@ class StatefulOVInferRequest : public OVInferRequest {
 
   void StartAsync() override;
   void Infer() override;
+  void RewindKVCache(size_t index) override;
 
  private:
   void _pre_infer();