[executorch][cuda] Let the CUDA delegate drive the off-graph KV step

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #22820
* #22819
* #22818
* #22817
* __->__ #22989
* #22816
* #22815

The runner currently calls prepare_step/commit_step on CudaKVCache, which puts
a backend face in the runner's hands. The delegate is the natural owner: it
already runs on the host immediately before and after the compiled program, and
that is the only point where storage can grow and be rebound before the kernels
read it.

The one thing the delegate cannot work out for itself is how many tokens a step
writes. Lowering compiles the cache op away, so at runtime there is only
compiled code over pre-bound memory and nothing calls back out. The export side
knows the width from its own signature, so it declares where to find it as an
"input_index:dim" compile spec, and the delegate reads that input's extent --
shape metadata only, never contents, which are on device.

Committing from here also pairs the two halves on every exit: a run that fails
never advances the cache past tokens the kernels did not write.

Differential Revision: D121131873

Pull Request resolved: #22989

Author: Gasoonjia

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