[executorch][cuda] Let the CUDA delegate drive the off-graph KV step (#23294)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #22989 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/191/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/191/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/186/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/191/orig
Differential Revision:
[D121131873](https://our.internmc.facebook.com/intern/diff/D121131873/)
@diff-train-skip-merge

---------

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

Author: pytorchbot

backends/cuda/CMakeLists.txt

@@ -407,6 +407,9 @@ target_link_libraries(
 )
 if(EXECUTORCH_BUILD_EXTENSION_LLM)
   target_link_libraries(aoti_cuda_backend PUBLIC extension_llm_cache)
+  target_compile_definitions(
+    aoti_cuda_backend PRIVATE EXECUTORCH_BUILD_EXTENSION_LLM
+  )
 endif()
 
 if(_cuda_is_msvc_toolchain)

backends/cuda/cuda_backend.py

@@ -25,8 +25,10 @@
     trim_host_memory,
 )
 from executorch.backends.cuda.passes.lower_offgraph_kv import (
+    CheckOffGraphKVStepWidthPass,
     LowerOffGraphKVPass,
     OFFGRAPH_KV_COMPILE_SPEC,
+    OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC,
     parse_offgraph_kv_manifest,
 )
 from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import (
@@ -782,6 +784,9 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                     )
                 triton_kernel_mode = mode
         passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()]
+        for spec in compile_specs:
+            if spec.key == OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC:
+                passes.append(CheckOffGraphKVStepWidthPass(spec.value))
         for spec in compile_specs:
             if spec.key == OFFGRAPH_KV_COMPILE_SPEC:
                 passes.append(

backends/cuda/passes/lower_offgraph_kv.py

@@ -219,6 +219,84 @@ def fn(q, k, v, position, k_storage, v_storage):
     return fn
 
 
+# Ops that pass a tensor's values through unchanged, which the export of a
+# position input may place between the method input and the cache op.
+_VALUE_PRESERVING_OPS = (
+    "view",
+    "reshape",
+    "unsqueeze",
+    "squeeze",
+    "expand",
+    "clone",
+    "alias",
+    "_to_copy",
+)
+
+
+def _source_input(node):
+    while node.op == "call_function" and any(
+        name in str(node.target) for name in _VALUE_PRESERVING_OPS
+    ):
+        node = node.args[0]
+    return node
+
+
+class CheckOffGraphKVStepWidthPass:
+    """Check the declared step width against the delegated program.
+
+    The runtime reads each step's token count from input ``input`` at ``dim``
+    of the delegate. That index is the delegate's own input order, which
+    partitioning decides, so it is checked here, after partitioning: every
+    cache op's position must come from that input, whose ``dim`` must be the
+    position's length.
+    """
+
+    requires_exported_program = True
+
+    def __init__(self, value: bytes) -> None:
+        text = value.decode("utf-8")
+        fields = text.split(":")
+        if len(fields) != 2 or not all(f.isdigit() for f in fields):
+            raise ValueError(
+                f"{OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC} must be "
+                f'"input_index:dim", got {text!r}'
+            )
+        self._input, self._dim = (int(f) for f in fields)
+
+    def __call__(self, exported_program: ExportedProgram) -> ExportedProgram:
+        graph = exported_program.graph_module.graph
+        user_inputs = set(exported_program.graph_signature.user_inputs)
+        inputs = [
+            n for n in graph.nodes if n.op == "placeholder" and n.name in user_inputs
+        ]
+        target = exir_ops.edge.kvcache.update_and_attend.default
+        for node in graph.nodes:
+            if node.op != "call_function" or node.target != target:
+                continue
+            if self._input >= len(inputs):
+                raise ValueError(
+                    f"off-graph KV step width names input {self._input}, but "
+                    f"the delegate takes {len(inputs)}"
+                )
+            declared = inputs[self._input]
+            position = node.args[3]
+            if _source_input(position) is not declared:
+                raise ValueError(
+                    f"off-graph KV step width names input {self._input} "
+                    f"({declared.name}), but the cache op's position comes "
+                    f"from {_source_input(position).name}"
+                )
+            shape = declared.meta["val"].shape
+            width = position.meta["val"].shape[0]
+            if self._dim >= len(shape) or shape[self._dim] != width:
+                raise ValueError(
+                    f"off-graph KV step width names dim {self._dim} of "
+                    f"{declared.name} {tuple(shape)}, which is not the step's "
+                    f"{width} positions"
+                )
+        return exported_program
+
+
 class LowerOffGraphKVPass:
     requires_exported_program = True
 

backends/cuda/runtime/cuda_backend.cpp

@@ -45,11 +45,13 @@
 #include <executorch/backends/aoti/utils.h>
 #include <executorch/backends/cuda/runtime/cuda_allocator.h>
 #include <executorch/backends/cuda/runtime/cuda_delegate_handle.h>
+#include <executorch/backends/cuda/runtime/cuda_kv_cache.h>
 #include <executorch/backends/cuda/runtime/cuda_mutable_state.h>
 #include <executorch/backends/cuda/runtime/cuda_weight_cache.h>
 #include <executorch/backends/cuda/runtime/platform/platform.h>
 #include <executorch/backends/cuda/runtime/shims/memory.h>
 #include <executorch/backends/cuda/runtime/utils.h>
+#include <executorch/extension/llm/cache/cache_registry.h>
 
 namespace executorch::backends::cuda {
 
@@ -87,6 +89,17 @@ constexpr int kCudaGraphWarmupSteps = 3;
 constexpr char kWeightSharingAcrossMethods[] = "weight_sharing_across_methods";
 } // anonymous namespace
 
+// Advances the off-graph KV cache past a step's tokens. Every successful exit
+// of execute() ends here, and only those: a run that failed never wrote its
+// tokens, and committing them would desynchronise the cache's length from the
+// history the kernels actually hold.
+Error commit_offgraph_kv_step(CudaDelegateHandle* handle, int64_t width) {
+  if (handle->kv_cache == nullptr) {
+    return Error::Ok;
+  }
+  return handle->kv_cache->commit_step(width);
+}
+
 class ET_EXPERIMENTAL CudaBackend final
     : public ::executorch::runtime::BackendInterface {
  private:
@@ -306,7 +319,14 @@ class ET_EXPERIMENTAL CudaBackend final
       ArrayRef<CompileSpec> compile_specs // This will be my empty list
   ) const override {
     std::string method_name;
+    std::optional<OffGraphKVStepWidth> kv_step_width;
     for (const CompileSpec& spec : compile_specs) {
+      if (std::strcmp(spec.key, kOffGraphKVStepWidthSpec) == 0) {
+        auto parsed = parse_offgraph_kv_step_width(std::string_view(
+            static_cast<const char*>(spec.value.buffer), spec.value.nbytes));
+        ET_CHECK_OK_OR_RETURN_ERROR(parsed.error());
+        kv_step_width = parsed.get();
+      }
       if (std::strcmp(spec.key, "method_name") == 0) {
         method_name.assign(
             static_cast<const char*>(spec.value.buffer),
@@ -443,6 +463,47 @@ class ET_EXPERIMENTAL CudaBackend final
 
     handle->container_handle = container_handle;
 
+    // Runtime-owned off-graph buffers must capture their AOTI names before
+    // the serialized constants update installs the ordinary weight set.
+    //
+    // The cache is found by the key the runner published as a runtime spec.
+    // No key means an in-graph model, whose KV state is ordinary mutable
+    // buffers -- that is the whole of the off-graph/in-graph detection, and it
+    // needs nothing from the user.
+    //
+    // The key alone does not decide it, though: a program whose storage was
+    // lowered off-graph carries no bytes for it, so running one without a cache
+    // would hand its kernels unbound memory. That requirement is read from the
+    // program's own constants and checked whether or not a key was passed.
+    const char* cache_key = nullptr;
+    if (auto spec = context.get_runtime_spec<const char*>(
+            ::executorch::extension::llm::cache::kCacheKeyOption);
+        spec.ok() && spec.get() != nullptr && *spec.get() != '\0') {
+      cache_key = spec.get();
+    }
+    const auto needs_offgraph_kv = requires_offgraph_kv_storage(*handle);
+    ET_CHECK_OK_OR_RETURN_ERROR(needs_offgraph_kv.error());
+    if (cache_key != nullptr) {
+#if defined(EXECUTORCH_BUILD_EXTENSION_LLM)
+      ET_CHECK_OK_OR_RETURN_ERROR(
+          attach_offgraph_kv_cache(*handle, cache_key, kv_step_width));
+#else
+      ET_LOG(
+          Error,
+          "init: an off-graph KV cache was supplied, but this CUDA backend was "
+          "built without EXECUTORCH_BUILD_EXTENSION_LLM");
+      return Error::NotSupported;
+#endif
+    }
+    ET_CHECK_OR_RETURN_ERROR(
+        !needs_offgraph_kv.get() || handle->kv_cache != nullptr,
+        InvalidArgument,
+        "init: method '%s' keeps its KV cache off-graph but no cache was "
+        "installed under the '%s' backend option (a CUDA backend built without "
+        "EXECUTORCH_BUILD_EXTENSION_LLM cannot run it)",
+        method_name.c_str(),
+        ::executorch::extension::llm::cache::kCacheKeyOption);
+
     // Versioned artifacts load each (device, FQN) through the same process-wide
     // cross-method cache model used by the legacy path. The payload only adds
     // the tensor metadata needed to reconstruct independently named PTD blobs.
@@ -591,6 +652,23 @@ class ET_EXPERIMENTAL CudaBackend final
       }
     }
 
+    // The cache has to grow, and its storage be rebound, before the compiled
+    // program reads it -- and nothing inside that program can call back out, so
+    // the step is driven from here rather than by the runner.
+    int64_t kv_step_width = 0;
+    if (handle->kv_cache != nullptr) {
+      const auto step_width = read_offgraph_kv_step_width(
+          handle->kv_step_width, Span<EValue*>(args.data(), n_inputs));
+      ET_CHECK_OK_OR_RETURN_ERROR(step_width.error());
+      kv_step_width = step_width.get();
+      // Growth is ordered on the stream this step runs on, which was just
+      // installed as current above.
+      const Result<cudaStream_t> step_stream = getCurrentCUDAStream(0);
+      ET_CHECK_OK_OR_RETURN_ERROR(step_stream.error());
+      ET_CHECK_OK_OR_RETURN_ERROR(
+          handle->kv_cache->prepare_step(kv_step_width, step_stream.get()));
+      ET_CHECK_OK_OR_RETURN_ERROR(handle->kv_cache->rebind_for_execute(handle));
+    }
     ET_CHECK_OK_OR_RETURN_ERROR(mutable_state_rebind_for_execute(handle));
 
     // ---------------------------------------------------------------
@@ -641,7 +719,7 @@ class ET_EXPERIMENTAL CudaBackend final
       }
       ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamSynchronize(cs));
 
-      return Error::Ok;
+      return commit_offgraph_kv_step(handle, kv_step_width);
     }
 
     // ---------------------------------------------------------------
@@ -934,7 +1012,7 @@ class ET_EXPERIMENTAL CudaBackend final
       // Last failure point is behind us, so the captured state is now the state
       // the next call should replay from.
       capture_guard.disarm();
-      return Error::Ok;
+      return commit_offgraph_kv_step(handle, kv_step_width);
     }
 
     // ----- Normal / WARMUP execution continues here -----
@@ -980,7 +1058,7 @@ class ET_EXPERIMENTAL CudaBackend final
       slim_outputs[i] = nullptr;
     }
 
-    return Error::Ok;
+    return commit_offgraph_kv_step(handle, kv_step_width);
   }
 
   void destroy(DelegateHandle* handle_) const override {
@@ -990,6 +1068,9 @@ class ET_EXPERIMENTAL CudaBackend final
     cuda::CudaDelegateHandle* handle = (cuda::CudaDelegateHandle*)handle_;
 
     mutable_state_forget_handle(handle);
+    if (handle->kv_cache != nullptr) {
+      handle->kv_cache->forget_handle(handle);
+    }
 
     // NOTE: AOTInductorModelContainerDelete does not work correctly with
     // multiple .so files. Deleting one container frees shared resources,
@@ -1465,7 +1546,7 @@ class ET_EXPERIMENTAL CudaBackend final
 namespace executorch::backends {
 namespace {
 auto cls = cuda::CudaBackend();
-executorch::runtime::Backend backend{"CudaBackend", &cls};
+executorch::runtime::Backend backend{cuda::kCudaBackendId, &cls};
 static executorch::runtime::Error success_with_compiler =
     register_backend(backend);
 

backends/cuda/runtime/cuda_delegate_handle.h

@@ -18,9 +18,26 @@
 #include <vector>
 
 namespace executorch {
+
+// Forward declarations for the off-graph KV cache the delegate may be given.
+// The concrete types live in cuda_kv_cache.h, which includes this header.
+namespace extension::llm::cache {
+class Cache;
+} // namespace extension::llm::cache
+
 namespace backends {
 namespace cuda {
 
+class CudaKVCache;
+
+// Where a method's inputs carry the number of tokens a step writes: an index
+// into execute()'s inputs and a dimension of that tensor. Declared by the
+// export side as the "offgraph_kv_step_width" compile spec, "input_index:dim".
+struct OffGraphKVStepWidth {
+  int input{0};
+  int dim{0};
+};
+
 using AOTInductorModelContainerGetConstantDtypeFunc =
     aoti::AOTIRuntimeError (*)(
         aoti::AOTInductorModelContainerHandle container_handle,
@@ -199,6 +216,22 @@ struct CudaDelegateHandle : public aoti::AOTIDelegateHandle {
   // SlimTensor handles alive for as long as AOTI may reference their views.
   std::vector<std::shared_ptr<CudaWeightStorage>> fqn_weight_storages;
   std::vector<std::unique_ptr<aoti::slim::SlimTensor>> fqn_weight_tensors;
+
+  // The off-graph KV cache the runner installed for this model, resolved from
+  // the registry at init. Null for an in-graph model, which is the signal that
+  // this program owns its KV state as ordinary (mutable) buffers.
+  //
+  // The shared_ptr is the handle's own claim on the cache, so the cache
+  // outlives the delegate even if the runner drops its guard first; the raw
+  // pointer is the backend face of that same object. Forward-declared rather
+  // than included: cuda_kv_cache.h includes this header.
+  std::shared_ptr<::executorch::extension::llm::cache::Cache> kv_cache_shared;
+  CudaKVCache* kv_cache{nullptr};
+
+  // Where this method's inputs carry the step width. The compiled program
+  // cannot report it, because lowering replaced the cache op with kernels over
+  // pre-bound memory.
+  OffGraphKVStepWidth kv_step_width;
 };
 
 } // namespace cuda