@@ -454,8 +454,9 @@ EOF
     fi
     RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt_file \"$PROMPT_FILE\" --max_new_tokens 512"
     if [ "$MODE" = "solo-text-offgraph" ]; then
-      # The off-graph KV cache does not support CUDA graph yet.
-      RUNNER_ARGS="$RUNNER_ARGS --cuda_graph=false"
+      # Exercise the same artifact on both decode paths: eager, then the
+      # captured CUDA graph, which is recaptured whenever the cache grows.
+      RUNNER_VARIANTS="--cuda_graph=false --cuda_graph"
     else
       RUNNER_ARGS="$RUNNER_ARGS --cuda_graph"
     fi
@@ -481,50 +482,69 @@ EOF
     ;;
 esac
 
-OUTPUT=$(eval $RUNNER_BIN $RUNNER_ARGS 2>&1)
-EXIT_CODE=$?
-set -e
+# A mode that needs more than one runtime configuration sets RUNNER_VARIANTS to
+# a space-separated list of extra args; each is run and validated on its own.
+# Everything else runs once with no extra args.
+run_and_validate() {
+  VARIANT_ARGS="$1"
+  if [ -n "$VARIANT_ARGS" ]; then
+    echo "--- Runner pass: $VARIANT_ARGS ---"
+  fi
 
-echo "Runner output:"
-echo "$OUTPUT"
+  set +e
+  OUTPUT=$(eval $RUNNER_BIN $RUNNER_ARGS $VARIANT_ARGS 2>&1)
+  EXIT_CODE=$?
+  set -e
 
-if [ $EXIT_CODE -ne 0 ]; then
-  echo "Unexpected exit code: $EXIT_CODE"
-  exit $EXIT_CODE
-fi
+  echo "Runner output:"
+  echo "$OUTPUT"
 
-# Validate output for models that have expected output
-if [ -n "$EXPECTED_OUTPUT" ]; then
-  if ! echo "$OUTPUT" | grep -iq "$EXPECTED_OUTPUT"; then
-    echo "Expected output '$EXPECTED_OUTPUT' not found in output"
-    exit 1
+  if [ $EXIT_CODE -ne 0 ]; then
+    echo "Unexpected exit code: $EXIT_CODE"
+    exit $EXIT_CODE
+  fi
+
+  # Validate output for models that have expected output
+  if [ -n "$EXPECTED_OUTPUT" ]; then
+    if ! echo "$OUTPUT" | grep -iq "$EXPECTED_OUTPUT"; then
+      echo "Expected output '$EXPECTED_OUTPUT' not found in output"
+      exit 1
+    else
+      echo "Success: '$EXPECTED_OUTPUT' found in output"
+    fi
   else
-    echo "Success: '$EXPECTED_OUTPUT' found in output"
+    echo "SUCCESS: Runner completed successfully"
   fi
-else
-  echo "SUCCESS: Runner completed successfully"
-fi
 
-# Validate GPU peak memory usage for models with known memory budgets.
-# The runner prints "GPU peak memory usage: XXXX.X MiB" at the end.
-case "$MODEL_NAME" in
-  qwen3_5_moe)
-    MAX_MEMORY_MIB=20480  # 20 GB — must fit on a single GPU (e.g. 4090)
-    PEAK_MEM=$(echo "$OUTPUT" | grep -oP 'GPU peak memory usage: \K[0-9.]+' || true)
-    if [ -n "$PEAK_MEM" ]; then
-      # Compare as integers (truncate decimals)
-      PEAK_MEM_INT=${PEAK_MEM%%.*}
-      if [ "$PEAK_MEM_INT" -gt "$MAX_MEMORY_MIB" ]; then
-        echo "FAIL: GPU peak memory ${PEAK_MEM} MiB exceeds budget ${MAX_MEMORY_MIB} MiB"
-        exit 1
+  # Validate GPU peak memory usage for models with known memory budgets.
+  # The runner prints "GPU peak memory usage: XXXX.X MiB" at the end.
+  case "$MODEL_NAME" in
+    qwen3_5_moe)
+      MAX_MEMORY_MIB=20480  # 20 GB — must fit on a single GPU (e.g. 4090)
+      PEAK_MEM=$(echo "$OUTPUT" | grep -oP 'GPU peak memory usage: \K[0-9.]+' || true)
+      if [ -n "$PEAK_MEM" ]; then
+        # Compare as integers (truncate decimals)
+        PEAK_MEM_INT=${PEAK_MEM%%.*}
+        if [ "$PEAK_MEM_INT" -gt "$MAX_MEMORY_MIB" ]; then
+          echo "FAIL: GPU peak memory ${PEAK_MEM} MiB exceeds budget ${MAX_MEMORY_MIB} MiB"
+          exit 1
+        else
+          echo "Success: GPU peak memory ${PEAK_MEM} MiB within budget (max ${MAX_MEMORY_MIB} MiB)"
+        fi
       else
-        echo "Success: GPU peak memory ${PEAK_MEM} MiB within budget (max ${MAX_MEMORY_MIB} MiB)"
+        echo "WARNING: GPU peak memory usage not found in output"
       fi
-    else
-      echo "WARNING: GPU peak memory usage not found in output"
-    fi
-    ;;
-esac
+      ;;
+  esac
+}
+
+if [ -z "${RUNNER_VARIANTS:-}" ]; then
+  run_and_validate ""
+else
+  for VARIANT in $RUNNER_VARIANTS; do
+    run_and_validate "$VARIANT"
+  done
+fi
 echo "::endgroup::"
 
 if [ "$DEVICE" = "cuda" ] && [ "$MODEL_NAME" = "qwen3_5_moe" ]; then
 
@@ -206,6 +206,13 @@ class ET_EXPERIMENTAL CudaBackend final
         res.ok() ? reinterpret_cast<name##Func>(res.get()) : nullptr; \
   } while (0)
 
+    {
+      auto res =
+          get_function(so_handle, "AOTInductorModelContainerRunSingleThreaded");
+      handle->run_single_threaded = res.ok()
+          ? reinterpret_cast<AOTInductorModelContainerRunFunc>(res.get())
+          : nullptr;
+    }
     LOAD_OPTIONAL_SYMBOL(
         get_num_constants, AOTInductorModelContainerGetNumConstants);
     LOAD_OPTIONAL_SYMBOL(
@@ -924,14 +931,19 @@ class ET_EXPERIMENTAL CudaBackend final
       end_capture_guard.arm(cuda_stream);
     }
 
-    AOTIRuntimeError error = handle->run(
-        handle->container_handle,
-        reinterpret_cast<Tensor**>(slim_inputs.data()),
-        n_inputs,
-        reinterpret_cast<Tensor**>(slim_outputs.data()),
-        n_outputs,
-        static_cast<void*>(cuda_stream),
-        nullptr);
+    const auto run =
+        handle->cuda_graph_state.phase != CudaGraphPhase::Disabled &&
+            handle->run_single_threaded != nullptr
+        ? handle->run_single_threaded
+        : handle->run;
+    AOTIRuntimeError error =
+        run(handle->container_handle,
+            reinterpret_cast<Tensor**>(slim_inputs.data()),
+            n_inputs,
+            reinterpret_cast<Tensor**>(slim_outputs.data()),
+            n_outputs,
+            static_cast<void*>(cuda_stream),
+            nullptr);
     run_called = true;
 
     // Delete orphaned pre-created outputs that run() replaced.
@@ -971,6 +983,7 @@ class ET_EXPERIMENTAL CudaBackend final
           Internal,
           "cudaGraphInstantiate failed: %s",
           cudaGetErrorString(gerr));
+      handle->cuda_graph_state.note_graph_allocations();
 
       // Record static output pointers (stable under graph replay). Releasing
       // them from slim_outputs here, before the copies below, keeps the cleanup
 
@@ -12,6 +12,7 @@
 #include <executorch/backends/aoti/slim/core/slim_tensor.h>
 #include <executorch/extension/cuda/device_guard.h>
 #include <executorch/extension/cuda/runtime_api.h>
+#include <algorithm>
 #include <cstdint>
 #include <cstdlib>
 #include <memory>
@@ -125,21 +126,94 @@ struct CudaGraphState {
   std::vector<size_t> static_input_nbytes;
   std::vector<size_t> static_output_nbytes;
 
+  // Allocations the graph makes and never frees -- the outputs AOTI allocates
+  // while being captured. AutoFreeOnLaunch frees one launch's before the next
+  // launch of the same graph, so the last launch's outlive a graph that is
+  // dropped unless released with it. Outputs that alias other memory are not
+  // graph allocations and are not listed.
+  std::vector<void*> graph_allocations;
+
   CudaGraphState() = default;
 
   ~CudaGraphState() {
+    release();
+  }
+
+  // Records the graph's outstanding allocations: every memory-allocation node
+  // without a matching free node. Call once the graph is captured.
+  void note_graph_allocations() {
+    graph_allocations.clear();
+#if !defined(EXECUTORCH_USE_HIP)
+    size_t count = 0;
+    if (graph == nullptr ||
+        cudaGraphGetNodes(graph, nullptr, &count) != cudaSuccess) {
+      return;
+    }
+    std::vector<cudaGraphNode_t> nodes(count);
+    if (cudaGraphGetNodes(graph, nodes.data(), &count) != cudaSuccess) {
+      return;
+    }
+    std::vector<void*> allocated;
+    std::vector<void*> freed;
+    for (cudaGraphNode_t node : nodes) {
+      cudaGraphNodeType type;
+      if (cudaGraphNodeGetType(node, &type) != cudaSuccess) {
+        continue;
+      }
+      if (type == cudaGraphNodeTypeMemAlloc) {
+        cudaMemAllocNodeParams params{};
+        if (cudaGraphMemAllocNodeGetParams(node, &params) == cudaSuccess) {
+          allocated.push_back(params.dptr);
+        }
+      } else if (type == cudaGraphNodeTypeMemFree) {
+        void* ptr = nullptr;
+        if (cudaGraphMemFreeNodeGetParams(node, &ptr) == cudaSuccess) {
+          freed.push_back(ptr);
+        }
+      }
+    }
+    for (void* ptr : allocated) {
+      if (std::find(freed.begin(), freed.end(), ptr) == freed.end()) {
+        graph_allocations.push_back(ptr);
+      }
+    }
+#endif
+  }
+
+  // Frees the captured graph, the static inputs pinned for it, and the
+  // allocations its last launch left outstanding (see graph_allocations).
+  void release() {
     if (graph_exec) {
       (void)cudaGraphExecDestroy(graph_exec);
+      graph_exec = nullptr;
     }
     if (graph) {
       (void)cudaGraphDestroy(graph);
+      graph = nullptr;
     }
-    // Only free input buffers — output buffers are owned by the AOTI runtime
-    // (allocated during graph capture via the caching allocator).
     for (auto* ptr : static_input_ptrs) {
       if (ptr)
         (void)cudaFree(ptr);
     }
+    // cudaFree synchronises, so the last replay's reads of these are done.
+    for (auto* ptr : graph_allocations) {
+      (void)cudaFree(ptr);
+    }
+    graph_allocations.clear();
+    static_input_ptrs.clear();
+    static_output_ptrs.clear();
+    static_input_nbytes.clear();
+    static_output_nbytes.clear();
+  }
+
+  // Drops the captured graph so a new one is captured, for when memory the
+  // graph baked in has moved. One eager step first: rebinding constants resets
+  // AOTI's constant-fold state, and the fold the next run performs cannot run
+  // inside a stream capture. The kernels are already loaded, so one is enough.
+  void recapture() {
+    release();
+    phase = CudaGraphPhase::Warmup;
+    warmup_remaining = 1;
   }
 
   // Non-copyable: prevent double-free of CUDA resources
@@ -155,22 +229,15 @@ struct CudaGraphState {
         static_input_ptrs(std::move(other.static_input_ptrs)),
         static_output_ptrs(std::move(other.static_output_ptrs)),
         static_input_nbytes(std::move(other.static_input_nbytes)),
-        static_output_nbytes(std::move(other.static_output_nbytes)) {
+        static_output_nbytes(std::move(other.static_output_nbytes)),
+        graph_allocations(std::move(other.graph_allocations)) {
     other.graph = nullptr;
     other.graph_exec = nullptr;
   }
 
   CudaGraphState& operator=(CudaGraphState&& other) noexcept {
     if (this != &other) {
-      // Clean up existing resources
-      if (graph_exec)
-        (void)cudaGraphExecDestroy(graph_exec);
-      if (graph)
-        (void)cudaGraphDestroy(graph);
-      for (auto* ptr : static_input_ptrs) {
-        if (ptr)
-          (void)cudaFree(ptr);
-      }
+      release();
 
       phase = other.phase;
       warmup_remaining = other.warmup_remaining;
@@ -180,6 +247,7 @@ struct CudaGraphState {
       static_output_ptrs = std::move(other.static_output_ptrs);
       static_input_nbytes = std::move(other.static_input_nbytes);
       static_output_nbytes = std::move(other.static_output_nbytes);
+      graph_allocations = std::move(other.graph_allocations);
 
       other.graph = nullptr;
       other.graph_exec = nullptr;
@@ -190,6 +258,14 @@ struct CudaGraphState {
 
 // CUDA-specific delegate handle that extends AOTIDelegateHandle.
 struct CudaDelegateHandle : public aoti::AOTIDelegateHandle {
+  // AOTI's run() records a completion event on the stream and queries it at
+  // the start of the next run. Captured into a CUDA graph, that record belongs
+  // to the graph, so the first run() after the graph is dropped (to recapture
+  // against moved KV storage) fails the query. The single-threaded entry point
+  // skips the event, and a CUDA-graph method runs single-threaded by design.
+  // Null when the .so does not export it; run() is used then.
+  aoti::AOTInductorModelContainerRunFunc run_single_threaded{nullptr};
+
   // 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
 
@@ -503,7 +503,25 @@ class CudaSequenceKVCache final : public cache::SequenceCache,
     }
     metrics_.flat_capacity = new_rows;
     metrics_.growth_count++;
+    // Every program sharing this cache now points at freed storage: drop the
+    // bindings so each rebinds before its next run, and any captured CUDA
+    // graph so it is captured again against the new storage. prefill usually
+    // grows the cache while decode's graph sits idle, so this reaches every
+    // handle, not only the one stepping now.
+    //
+    // Rebinding also resets AOTI's constant-fold state, which must be run
+    // eagerly, so every graph-enabled handle gets at least one eager step
+    // before it captures -- including one that was about to capture for the
+    // first time, and without shortening a longer warmup still outstanding.
     bound_.clear();
+    for (auto& entry : descriptors_) {
+      CudaGraphState& graph = entry.first->cuda_graph_state;
+      if (graph.phase == CudaGraphPhase::Replay) {
+        graph.recapture();
+      } else if (graph.phase == CudaGraphPhase::Warmup) {
+        graph.warmup_remaining = std::max(graph.warmup_remaining, 1);
+      }
+    }
     ET_LOG(
         Info,
         "offgraph_kv: grew flat_capacity=%lld->%lld allocated_bytes=%lld "