Preserve suffix errors outside optional prefix selection

Author: mergennachin

extension/llm/batching/prefix_cache.h

@@ -348,11 +348,8 @@ class ET_EXPERIMENTAL PrefixCache {
     return lookup_identity(copy_tokens(request));
   }
 
-  // Skip candidates that cannot improve on an already available prefix.
-  std::optional<PrefixMatch> lookup_identity(
-      const PrefixIdentity& request,
-      std::size_t longer_than = 0) {
-    auto match = find_identity(request, longer_than);
+  std::optional<PrefixMatch> lookup_identity(const PrefixIdentity& request) {
+    auto match = find_identity(request);
     if (match) {
       touch(*match);
     }
@@ -361,6 +358,11 @@ class ET_EXPERIMENTAL PrefixCache {
 
   // Validate candidates before cloning, without refreshing recency. Call
   // touch() after adopting the match. The selector must not mutate this cache.
+  // Candidates arrive in descending length order, with repeated lengths
+  // allowed. Use immediately attempts a clone; failure advances to the next
+  // candidate. A successful clone returns before another selection. Stop and a
+  // miss both return nullopt; callers must retain any separate reason for
+  // stopping.
   std::optional<PrefixMatch> find_identity(
       const PrefixIdentity& request,
       std::size_t longer_than = 0,
@@ -404,22 +406,14 @@ class ET_EXPERIMENTAL PrefixCache {
             continue;
           }
         }
-#if ET_HAS_EXCEPTIONS
-        try {
-#endif
-          auto session =
-              candidate.entry->snapshot
-                  .clone_async(static_cast<Position>(candidate.matched))
-                  .get();
-          if (session) {
-            return PrefixMatch{
-                std::move(*session), candidate.matched, candidate.entry->id};
-          }
-#if ET_HAS_EXCEPTIONS
-        } catch (...) {
-          // A refused snapshot does not rule out the remaining candidates.
+        auto session =
+            candidate.entry->snapshot
+                .clone_async(static_cast<Position>(candidate.matched))
+                .get();
+        if (session) {
+          return PrefixMatch{
+              std::move(*session), candidate.matched, candidate.entry->id};
         }
-#endif
       }
 #if ET_HAS_EXCEPTIONS
     } catch (const std::bad_alloc&) {

extension/llm/batching/test/prefix_cache_test.cpp

@@ -898,22 +898,22 @@ TEST_F(PrefixCacheTest, EqualLengthMatchesPreferTheMostRecentlyUsedSnapshot) {
   EXPECT_EQ(executor.clones().back().source, first_id);
 }
 
-TEST_F(PrefixCacheTest, LookupThresholdSkipsClonesAndDoesNotRefreshLru) {
+TEST_F(PrefixCacheTest, FindThresholdSkipsClonesAndDoesNotRefreshLru) {
   PrefixCache cache(2);
   ASSERT_TRUE(cache.insert({1, 2, 3}, snapshot({1, 2, 3})));
   ASSERT_TRUE(cache.insert({8, 9}, snapshot({8, 9})));
   const auto request =
       token_identity(std::make_shared<const Tokens>(Tokens{1, 2, 3, 4}));
   const auto before = executor.clones().size();
-  EXPECT_FALSE(cache.lookup_identity(request, 3));
+  EXPECT_FALSE(cache.find_identity(request, 3));
   EXPECT_FALSE(
-      cache.lookup_identity(request, std::numeric_limits<size_t>::max()));
+      cache.find_identity(request, std::numeric_limits<size_t>::max()));
   EXPECT_EQ(executor.clones().size(), before);
   ASSERT_TRUE(cache.insert({6, 7}, snapshot({6, 7})));
   EXPECT_FALSE(cache.lookup({1, 2, 3, 4}));
   const auto other =
       token_identity(std::make_shared<const Tokens>(Tokens{8, 9, 10}));
-  auto match = cache.lookup_identity(other, 1);
+  auto match = cache.find_identity(other, 1);
   ASSERT_TRUE(match);
   EXPECT_EQ(match->matched_tokens, 2u);
   expect_cold_match(*match, {8, 9, 10});

extension/llm/serving/serving_runtime.cpp

@@ -783,6 +783,7 @@ struct ServingRuntime::Impl {
     }
     std::optional<batching::Session> cloned;
     std::optional<batching::PrefixMatch> match;
+    auto fallback_limit = upto ? upto : std::numeric_limits<std::size_t>::max();
     if (upto) {
       auto selection = lookup_prefix(request, match, upto);
       const auto* outcome = std::get_if<Selection>(&selection);
@@ -811,17 +812,15 @@ struct ServingRuntime::Impl {
           // Try a retained snapshot if the live source cannot be cloned.
         }
 #endif
+      } else {
+        fallback_limit = upto - 1;
       }
     }
     if (!cloned) {
       if (request->cancelled.load()) {
         return Selection::Cancelled;
       }
-      auto selection = lookup_prefix(
-          request,
-          match,
-          0,
-          upto ? upto : std::numeric_limits<std::size_t>::max());
+      auto selection = lookup_prefix(request, match, 0, fallback_limit);
       if (auto* outcome = std::get_if<Selection>(&selection);
           !outcome || *outcome != Selection::Selected) {
         return selection;
@@ -1023,36 +1022,48 @@ struct ServingRuntime::Impl {
     }
     auto original_delta = std::move(request->request.delta);
     SelectionResult result = Selection::Unavailable;
+    std::optional<std::size_t> previous_start;
+    SelectionResult selection = Selection::Unavailable;
     // A lookup clone uses the new request's reserved working-session slot.
 #if ET_HAS_EXCEPTIONS
     try {
 #endif
       match = prefix_cache_.find_identity(
           *request->text->prompt, longer_than, at_most, [&](std::size_t start) {
             using Decision = batching::PrefixCache::MatchDecision;
+            if (request->cancelled.load()) {
+              result = Selection::Cancelled;
+              return Decision::Stop;
+            }
+            if (previous_start != start) {
 #if ET_HAS_EXCEPTIONS
-            try {
+              try {
 #endif
-              auto selection = select_input(request, start);
-              if (const auto* outcome = std::get_if<Selection>(&selection)) {
-                if (*outcome == Selection::Selected) {
-                  return Decision::Use;
-                }
-                if (*outcome == Selection::Unavailable) {
-                  return Decision::Skip;
-                }
-              } else if (
-                  std::get<ServingError>(selection).code !=
-                  ErrorCode::NotReady) {
+                selection = select_input(request, start);
+#if ET_HAS_EXCEPTIONS
+              } catch (...) {
+                selection = ServingError{
+                    ErrorCode::Internal, "prepared suffix selection failed"};
+              }
+#endif
+              previous_start = start;
+            }
+            if (const auto* outcome = std::get_if<Selection>(&selection)) {
+              if (*outcome == Selection::Selected) {
+                return Decision::Use;
+              }
+              if (*outcome == Selection::Unavailable) {
                 return Decision::Skip;
               }
-              result = std::move(selection);
-              return Decision::Stop;
-#if ET_HAS_EXCEPTIONS
-            } catch (...) {
+            } else if (
+                longer_than &&
+                std::get<ServingError>(selection).code != ErrorCode::NotReady) {
+              // Only the preferred lookup may ignore a new boundary's error;
+              // ordinary lookup retains the existing input-validation contract.
               return Decision::Skip;
             }
-#endif
+            result = selection;
+            return Decision::Stop;
           });
 #if ET_HAS_EXCEPTIONS
     } catch (...) {