[executorch][cuda] Add dynamic off-graph KV runtime manager (pytorch#23293)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#22816 by
@Gasoonjia
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/186/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/186/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/185/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/186/orig
Differential Revision:
[D119571986](https://our.internmc.facebook.com/intern/diff/D119571986/)
@diff-train-skip-merge

---------

Co-authored-by: gasoonjia <gasoonjia@icloud.com>

Author: pytorchbot

backends/cuda/CMakeLists.txt

@@ -349,6 +349,10 @@ set(_aoti_cuda_backend_sources
     runtime/cuda_backend.cpp runtime/cuda_mutable_state.cpp
     runtime/cuda_weight_cache.cpp
 )
+# The off-graph KV cache builds on the LLM extension's neutral cache.
+if(EXECUTORCH_BUILD_EXTENSION_LLM)
+  list(APPEND _aoti_cuda_backend_sources runtime/cuda_kv_cache.cpp)
+endif()
 if(_cuda_is_msvc_toolchain)
   # MSVC links aoti_cuda_backend into portable_lib without relying on C++
   # symbols exported from aoti_cuda_shims.dll.
@@ -401,6 +405,9 @@ target_link_libraries(
   PUBLIC cuda_platform extension_tensor ${_cuda_runtime_libraries}
          extension_cuda ${CMAKE_DL_LIBS}
 )
+if(EXECUTORCH_BUILD_EXTENSION_LLM)
+  target_link_libraries(aoti_cuda_backend PUBLIC extension_llm_cache)
+endif()
 
 if(_cuda_is_msvc_toolchain)
   # cuda_backend.cpp uses SlimTensor CUDA utilities (e.g. getCurrentCUDAStream)

backends/cuda/runtime/backend_options.h

@@ -0,0 +1,20 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+namespace executorch::backends::cuda {
+
+// Shared identifiers for the CUDA backend. Included by the backend itself (to
+// register) and by callers/runners (to route backend options and to ask
+// CacheFactory for a builder), keeping the string literals in one place.
+
+// Backend id under which the CUDA backend registers (see cuda_backend.cpp).
+inline constexpr char kCudaBackendId[] = "CudaBackend";
+
+} // namespace executorch::backends::cuda

backends/cuda/runtime/cuda_backend.cpp

@@ -202,6 +202,8 @@ class ET_EXPERIMENTAL CudaBackend final
         AOTInductorModelContainerGetConstantOriginalFQN);
     LOAD_OPTIONAL_SYMBOL(
         get_constant_dtype, AOTInductorModelContainerGetConstantDtype);
+    LOAD_OPTIONAL_SYMBOL(
+        get_constant_data_size, AOTInductorModelContainerGetConstantDataSize);
     LOAD_OPTIONAL_SYMBOL(
         extract_constants_map, AOTInductorModelContainerExtractConstantsMap);
     LOAD_OPTIONAL_SYMBOL(

backends/cuda/runtime/cuda_delegate_handle.h

@@ -26,6 +26,11 @@ using AOTInductorModelContainerGetConstantDtypeFunc =
         aoti::AOTInductorModelContainerHandle container_handle,
         size_t idx,
         int32_t* dtype);
+using AOTInductorModelContainerGetConstantDataSizeFunc =
+    aoti::AOTIRuntimeError (*)(
+        aoti::AOTInductorModelContainerHandle container_handle,
+        size_t idx,
+        size_t* data_size);
 struct CudaWeightStorage {
   void* data{nullptr};
   size_t nbytes{0};
@@ -170,6 +175,10 @@ struct CudaGraphState {
 struct CudaDelegateHandle : public aoti::AOTIDelegateHandle {
   // Extra AOTI metadata used to validate per-FQN weights before binding.
   AOTInductorModelContainerGetConstantDtypeFunc get_constant_dtype{nullptr};
+  // Bytes a constant's compiled shape spans; the off-graph KV cache checks its
+  // storage against it before binding.
+  AOTInductorModelContainerGetConstantDataSizeFunc get_constant_data_size{
+      nullptr};
 
   // The per-thread stream. Nothing owns it: the value is a fixed sentinel the
   // driver resolves to a different stream on each host thread, so releasing the