Internal changes only.

runtime/components/stop_token_detector.cc

@@ -69,22 +69,16 @@ absl::Status StopTokenDetector::AddStopTokenSequence(
   }
 
   stop_sequences_storage_.push_back(stop_sequence);
-
-  // Add a progress tracker for the new stop sequence for each batch item.
-  for (auto& progress_vector_for_item : batch_item_match_progress_) {
-    progress_vector_for_item.push_back(0);
-  }
+  max_stop_sequence_length_ =
+      std::max(max_stop_sequence_length_, stop_sequence.size());
   return absl::OkStatus();
 }
 
 void StopTokenDetector::ResetBatch(size_t batch_size) {
   int new_batch_size = batch_size == 0 ? stop_token_found_.size() : batch_size;
   stop_token_found_.assign(new_batch_size, false);
   max_batch_item_match_progress_.assign(new_batch_size, 0);
-  // Initialize progress for each batch item for all currently defined stop
-  // sequences.
-  batch_item_match_progress_.assign(
-      new_batch_size, std::vector<int>(stop_sequences_storage_.size(), 0));
+  batch_item_token_history_.assign(new_batch_size, std::vector<int>());
   matched_stop_sequence_length_.assign(new_batch_size, 0);
 }
 
@@ -108,33 +102,51 @@ absl::Status StopTokenDetector::ProcessTokens(
       matched_stop_sequence_length_[i]++;
       continue;
     }
-    max_batch_item_match_progress_[i] = 0;
+
     int current_token_id = latest_tokens[i];
-    for (size_t k = 0; k < stop_sequences_storage_.size(); ++k) {
-      const auto& stop_seq_k =
-          stop_sequences_storage_[k];  // Guaranteed non-empty
-      int& current_match_len_for_k = batch_item_match_progress_[i][k];
-
-      if (current_match_len_for_k < stop_seq_k.size() &&
-          stop_seq_k[current_match_len_for_k] == current_token_id) {
-        current_match_len_for_k++;
-      } else {
-        // Mismatch or sequence completed; reset progress for this stop_seq_k.
-        // Check if current token starts stop_seq_k anew.
-        if (stop_seq_k[0] == current_token_id) {
-          current_match_len_for_k = 1;
-        } else {
-          current_match_len_for_k = 0;
+    auto& history = batch_item_token_history_[i];
+    history.push_back(current_token_id);
+    if (history.size() > max_stop_sequence_length_) {
+      history.erase(history.begin());
+    }
+
+    max_batch_item_match_progress_[i] = 0;
+    for (const auto& stop_seq : stop_sequences_storage_) {
+      // Check for full match.
+      if (history.size() >= stop_seq.size()) {
+        bool full_match = true;
+        for (size_t idx = 0; idx < stop_seq.size(); ++idx) {
+          if (history[history.size() - stop_seq.size() + idx] !=
+              stop_seq[idx]) {
+            full_match = false;
+            break;
+          }
+        }
+        if (full_match) {
+          stop_token_found_[i] = true;
+          matched_stop_sequence_length_[i] = stop_seq.size();
+          max_batch_item_match_progress_[i] = stop_seq.size();
+          break;  // Stop token found, no need to check other sequences.
         }
       }
 
-      if (current_match_len_for_k > 0 &&
-          current_match_len_for_k == stop_seq_k.size()) {
-        stop_token_found_[i] = true;
-        matched_stop_sequence_length_[i] = stop_seq_k.size();
+      // Check for partial match (only if stop token is not found yet).
+      if (!stop_token_found_[i]) {
+        for (int len = std::min(history.size(), stop_seq.size() - 1);
+             len > max_batch_item_match_progress_[i]; --len) {
+          bool partial_match = true;
+          for (int idx = 0; idx < len; ++idx) {
+            if (history[history.size() - len + idx] != stop_seq[idx]) {
+              partial_match = false;
+              break;
+            }
+          }
+          if (partial_match) {
+            max_batch_item_match_progress_[i] = len;
+            break;
+          }
+        }
       }
-      max_batch_item_match_progress_[i] =
-          std::max(max_batch_item_match_progress_[i], current_match_len_for_k);
     }
   }
   return absl::OkStatus();

runtime/components/stop_token_detector.h

@@ -99,12 +99,16 @@ class StopTokenDetector {
   // Stores all added stop sequences.
   std::vector<std::vector<int>> stop_sequences_storage_;
 
-  // batch_item_match_progress_[i][k]: current match length for batch item 'i'
-  // against stop_sequences_storage_[k].
-  std::vector<std::vector<int>> batch_item_match_progress_;
+  // Maximum length of any registered stop sequence.
+  size_t max_stop_sequence_length_ = 0;
 
-  // max_batch_item_match_progress_[i]: maximum match length for batch item 'i'
-  // against all stop_sequences_storage_[k].
+  // The token history buffer for each batch item.
+  // batch_item_token_history_[i] stores the last max_stop_sequence_length_
+  // tokens for batch item 'i'.
+  std::vector<std::vector<int>> batch_item_token_history_;
+
+  // max_batch_item_match_progress_[i]: maximum partial match length for batch
+  // item 'i' against all stop_sequences_storage_[k].
   std::vector<int> max_batch_item_match_progress_;
 
   // stop_token_found_[i]: true if batch item 'i' has matched a stop sequence.

runtime/components/stop_token_detector_test.cc

@@ -134,5 +134,27 @@ TEST(StopTokenDetectorTest, ResetBatch) {
   EXPECT_EQ(0, detector.GetStepsBeforeStopTokens()[0]);
 }
 
+TEST(StopTokenDetectorTest, ProcessTokensSelfOverlap) {
+  StopTokenDetector detector(1);
+  EXPECT_OK(detector.AddStopTokenSequence({1, 1, 2}));
+
+  std::vector<int> tokens = {1, 1, 1, 2};
+
+  size_t i;
+  for (i = 0; i < tokens.size(); ++i) {
+    std::vector<int> current_batch_tokens = {tokens[i]};
+    EXPECT_OK(detector.ProcessTokens(absl::MakeSpan(current_batch_tokens)));
+    auto all_done = detector.AllDone();
+    ASSERT_OK(all_done.status());
+    if (*all_done) {
+      break;
+    }
+  }
+  EXPECT_EQ(i, 3);
+  auto final_all_done = detector.AllDone();
+  ASSERT_OK(final_all_done.status());
+  EXPECT_TRUE(*final_all_done);
+}
+
 }  // namespace
 }  // namespace litert::lm