[Bugfix][Spec Decode] Select earliest-completing stop string in check_stop_strings (#49391)

Signed-off-by: Junpu Yu <davidyu@nvidia.com>

tests/detokenizer/test_check_stop_strings.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for check_stop_strings.
+
+These are pure-function tests (no model / GPU). They pin down which stop
+string is selected when several stop strings match within the text that was
+appended in a single step -- which happens under speculative decoding, where
+multiple tokens (and therefore multiple stop strings) can be appended at once.
+"""
+
+import pytest
+
+from vllm.v1.engine.detokenizer import check_stop_strings
+
+
+@pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]])
+def test_earliest_completing_stop_wins_regardless_of_list_order(stop):
+    # " The user is a": " is a" (5 chars) was appended in one step. Both "is"
+    # (index 10) and " a" (index 13) land in the same window. "is" completes
+    # earlier in the text, so it must win over list order.
+    text = " The user is a"
+    new_char_count = len(" is a")
+
+    assert check_stop_strings(text, new_char_count, stop, include_in_output=False) == (
+        "is",
+        10,
+    )
+
+
+@pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]])
+def test_earliest_completing_stop_include_in_output(stop):
+    text = " The user is a"
+    new_char_count = len(" is a")
+
+    # Truncate to the end of "is" (index 12) -> " The user is".
+    assert check_stop_strings(text, new_char_count, stop, include_in_output=True) == (
+        "is",
+        12,
+    )
+
+
+def test_completion_position_not_start_position():
+    # "b" starts later than "abc" but completes earlier, so it must win.
+    text = "abc"
+    assert check_stop_strings(
+        text, len(text), ["abc", "b"], include_in_output=False
+    ) == ("b", 1)
+
+
+@pytest.mark.parametrize(
+    "stop,expected",
+    [
+        (["ab", "b"], ("ab", 0)),
+        (["b", "ab"], ("b", 1)),
+    ],
+)
+def test_ties_broken_by_list_order(stop, expected):
+    # "ab" and "b" both complete at index 2; list order decides the winner.
+    text = "ab"
+    assert (
+        check_stop_strings(text, len(text), stop, include_in_output=False) == expected
+    )
+
+
+def test_single_stop_in_window_unchanged():
+    # The common case (one stop in the window) is unaffected by the change.
+    text = "hello world."
+    assert check_stop_strings(text, 1, ["."], include_in_output=False) == (".", 11)
+    # Stop completes at the very end -> no truncation needed (-1).
+    assert check_stop_strings(text, 1, ["."], include_in_output=True) == (".", -1)
+
+
+def test_no_match_and_empty_inputs_return_none():
+    assert check_stop_strings("hello", 5, ["zzz"], include_in_output=False) is None
+    assert check_stop_strings("hello", 0, ["h"], include_in_output=False) is None
+    assert check_stop_strings("hello", 5, [], include_in_output=False) is None

vllm/v1/engine/detokenizer.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import sys
 from abc import ABC, abstractmethod
 
 import tokenizers
@@ -320,10 +321,19 @@ def check_stop_strings(
     Where stop_string is the matched stop string and offset is the
     length to which output_text should be truncated, or -1 for no
     truncation.
+
+    When several stop strings match within the newly generated text (for
+    example when speculative decoding appends multiple tokens in a single
+    step), the stop string that completes earliest in the text is selected,
+    so the result matches appending one token at a time. Ties are broken by
+    stop-list order.
     """
     if not new_char_count or not stop:
         return None
 
+    best_stop_str: str | None = None
+    best_stop_index = 0
+    best_end = sys.maxsize
     for stop_str in stop:
         stop_string_len = len(stop_str)
         # Avoid searching already-searched text.
@@ -331,14 +341,22 @@ def check_stop_strings(
         if stop_index == -1:
             continue
 
-        if include_in_output:
-            # Truncate to end of stop string.
-            stop_index += stop_string_len
-            if stop_index >= len(output_text):
-                # No truncation required.
-                return stop_str, -1
+        # Prefer the stop string that completes earliest in the text.
+        end = stop_index + stop_string_len
+        if end < best_end:
+            best_stop_str = stop_str
+            best_stop_index = stop_index
+            best_end = end
 
-        # Truncate the output text to either the beginning
-        # or end of the stop string.
-        return stop_str, stop_index
-    return None
+    if best_stop_str is None:
+        return None
+
+    if include_in_output:
+        # Truncate to end of stop string.
+        if best_end >= len(output_text):
+            # No truncation required.
+            return best_stop_str, -1
+        return best_stop_str, best_end
+
+    # Truncate the output text to the beginning of the stop string.
+    return best_stop_str, best_stop_index