From d81042b8598f30b59befaa15fce9ed28aab3e2db Mon Sep 17 00:00:00 2001
From: gasoonjia <gasoonjia@icloud.com>
Date: Mon, 19 Oct 2026 13:07:56 -0700
Subject: [PATCH] [executorch][cuda] Test prefix-cache reuse on CudaExecutor

Pull Request resolved: https://github.com/pytorch/executorch/pull/23732

The serving runtime's opt-in prefix cache (batching::PrefixCache, enabled by
`--prefix_cache_entries`) reuses KV through Session::clone_async, which on CUDA
reaches CudaExecutor::clone -> SessionTable::clone -> CellCache::seq_clone. The
cell layout clones by reference: the snapshot's id joins the owner bits of the
source's cells below the boundary, so reuse shares the prompt's KV instead of
copying it. Nothing needed to change to enable it; this adds the coverage that
shows it is correct on the GPU.

PrefixCacheResumesFromACapturedPrompt, on each toy artifact (dense, sparse,
device_sampling): every prompt's cold generation leaves a snapshot captured at
its prompt boundary. A lookup for the same prompt reuses all but its last token
and one for a prompt extending it reuses the whole snapshot; both resume on the
eager greedy continuation. So the clone carried the prompt's exact KV, and the
source's later decode writes, padding rows and the clones' own writes did not
disturb one another.
ghstack-source-id: 451855087
@exported-using-ghexport

Differential Revision: [D124871960](https://our.internmc.facebook.com/intern/diff/D124871960/)
---
 .../cuda/batching/test/test_cuda_executor.cpp | 78 +++++++++++++++++++
 1 file changed, 78 insertions(+)

diff --git a/backends/cuda/batching/test/test_cuda_executor.cpp b/backends/cuda/batching/test/test_cuda_executor.cpp
index 6c66cd1de15..34c54930319 100644
--- a/backends/cuda/batching/test/test_cuda_executor.cpp
+++ b/backends/cuda/batching/test/test_cuda_executor.cpp
@@ -14,6 +14,7 @@
 
 #include <executorch/backends/cuda/batching/cuda_executor.h>
 #include <executorch/extension/llm/batching/decode_first_scheduler.h>
+#include <executorch/extension/llm/batching/prefix_cache.h>
 #include <executorch/extension/llm/batching/runner.h>
 
 #include <cuda_runtime.h>
@@ -169,6 +170,34 @@ class CudaExecutorTest : public ::testing::TestWithParam<const char*> {
     return out;
   }
 
+  // Generates `new_tokens` greedily from `delta` on `session`, which may hold
+  // a reused prefix. A capture, if given, wraps the callback.
+  static std::vector<batching::Token> continue_on(
+      const batching::Session& session,
+      std::vector<batching::Token> delta,
+      int new_tokens,
+      batching::PrefixCache::PromptCapture* capture = nullptr) {
+    std::vector<batching::Token> tokens;
+    std::mutex mutex;
+    batching::GenConfig config;
+    config.max_new_tokens = new_tokens;
+    config.seed = 0;
+    batching::GenerationCallback on_update =
+        [&](const batching::GenerationUpdate& update) {
+          std::lock_guard<std::mutex> guard(mutex);
+          tokens.insert(
+              tokens.end(), update.tokens.begin(), update.tokens.end());
+        };
+    if (capture != nullptr) {
+      on_update = capture->wrap(std::move(on_update));
+    }
+    auto handle =
+        session.generate_async(std::move(delta), config, std::move(on_update));
+    handle.wait();
+    EXPECT_EQ(handle.finish_reason(), batching::FinishReason::NewTokenLimit);
+    return tokens;
+  }
+
   static std::unique_ptr<batching::Scheduler> scheduler() {
     // Up to four decodes per forward and prefill chunks of eight, so decodes
     // and prefills share forwards and a long prompt runs in several chunks.
@@ -407,6 +436,55 @@ TEST_P(
   }
 }
 
+TEST_P(CudaExecutorTest, PrefixCacheResumesFromACapturedPrompt) {
+  // Each prompt's cold generation leaves a snapshot cloned at its prompt
+  // boundary. That snapshot then serves the same prompt again and a prompt
+  // extending it; both resume on the eager greedy path, so the clone carried
+  // the prompt's exact KV and later writes did not disturb it.
+  auto exec = executor();
+  ASSERT_NE(exec, nullptr);
+  batching::Runner runner(*exec, scheduler());
+  batching::PrefixCache prefixes(2);
+  constexpr int kKnown = 2;
+  for (const Case& c : cases_) {
+    if (c.prompt.size() < 2) {
+      continue;
+    }
+    {
+      auto cold = runner.open_session_async().get();
+      ASSERT_TRUE(cold.has_value());
+      auto capture = prefixes.capture_prompt(*cold, c.prompt);
+      EXPECT_EQ(continue_on(*cold, c.prompt, kNewTokens, &capture), c.expected);
+      EXPECT_TRUE(capture.collect());
+    }
+
+    auto same = prefixes.lookup(c.prompt);
+    ASSERT_TRUE(same.has_value());
+    EXPECT_EQ(same->matched_tokens, c.prompt.size() - 1);
+    EXPECT_EQ(
+        same->session.position(),
+        static_cast<batching::Position>(c.prompt.size() - 1));
+    EXPECT_EQ(
+        continue_on(same->session, {c.prompt.back()}, kNewTokens), c.expected);
+
+    std::vector<batching::Token> extended = c.prompt;
+    extended.insert(
+        extended.end(), c.expected.begin(), c.expected.begin() + kKnown);
+    auto longer = prefixes.lookup(extended);
+    ASSERT_TRUE(longer.has_value());
+    EXPECT_EQ(longer->matched_tokens, c.prompt.size());
+    const std::vector<batching::Token> suffix(
+        extended.begin() + longer->matched_tokens, extended.end());
+    EXPECT_EQ(
+        continue_on(longer->session, suffix, kNewTokens - kKnown),
+        std::vector<batching::Token>(
+            c.expected.begin() + kKnown, c.expected.end()));
+  }
+  prefixes.clear();
+  runner.shutdown();
+  EXPECT_EQ(runner.metrics().steps_failed, 0u);
+}
+
 TEST_P(CudaExecutorTest, RefusesLimitsThePoolCannotHold) {
   // 8 sessions of the full 64-token context need 512 cells; the program has
   // 256.