[Bugfix][Spec Decode] Select earliest-completing stop string in check_stop_strings (#49391)
Signed-off-by: Junpu Yu <davidyu@nvidia.com>
Junpu Yu committed
Jul 23, 2026 at 03:14 UTC
1ad84fea866bc478942efa8550036ffa52a51283
2 files changed
+105
-11
tests/detokenizer/test_check_stop_strings.py
new
+76
@@ -0,0 +1,76 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for check_stop_strings.
4
+
5
+These are pure-function tests (no model / GPU). They pin down which stop
6
+string is selected when several stop strings match within the text that was
7
+appended in a single step -- which happens under speculative decoding, where
8
+multiple tokens (and therefore multiple stop strings) can be appended at once.
9
+"""
10
+
11
+import pytest
12
+
13
+from vllm.v1.engine.detokenizer import check_stop_strings
14
+
15
+
16
+@pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]])
17
+def test_earliest_completing_stop_wins_regardless_of_list_order(stop):
18
+ # " The user is a": " is a" (5 chars) was appended in one step. Both "is"
19
+ # (index 10) and " a" (index 13) land in the same window. "is" completes
20
+ # earlier in the text, so it must win over list order.
21
+ text = " The user is a"
22
+ new_char_count = len(" is a")
23
+
24
+ assert check_stop_strings(text, new_char_count, stop, include_in_output=False) == (
25
+ "is",
26
+ 10,
27
+ )
28
+
29
+
30
+@pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]])
31
+def test_earliest_completing_stop_include_in_output(stop):
32
+ text = " The user is a"
33
+ new_char_count = len(" is a")
34
+
35
+ # Truncate to the end of "is" (index 12) -> " The user is".
36
+ assert check_stop_strings(text, new_char_count, stop, include_in_output=True) == (
37
+ "is",
38
+ 12,
39
+ )
40
+
41
+
42
+def test_completion_position_not_start_position():
43
+ # "b" starts later than "abc" but completes earlier, so it must win.
44
+ text = "abc"
45
+ assert check_stop_strings(
46
+ text, len(text), ["abc", "b"], include_in_output=False
47
+ ) == ("b", 1)
48
+
49
+
50
+@pytest.mark.parametrize(
51
+ "stop,expected",
52
+ [
53
+ (["ab", "b"], ("ab", 0)),
54
+ (["b", "ab"], ("b", 1)),
55
+ ],
56
+)
57
+def test_ties_broken_by_list_order(stop, expected):
58
+ # "ab" and "b" both complete at index 2; list order decides the winner.
59
+ text = "ab"
60
+ assert (
61
+ check_stop_strings(text, len(text), stop, include_in_output=False) == expected
62
+ )
63
+
64
+
65
+def test_single_stop_in_window_unchanged():
66
+ # The common case (one stop in the window) is unaffected by the change.
67
+ text = "hello world."
68
+ assert check_stop_strings(text, 1, ["."], include_in_output=False) == (".", 11)
69
+ # Stop completes at the very end -> no truncation needed (-1).
70
+ assert check_stop_strings(text, 1, ["."], include_in_output=True) == (".", -1)
71
+
72
+
73
+def test_no_match_and_empty_inputs_return_none():
74
+ assert check_stop_strings("hello", 5, ["zzz"], include_in_output=False) is None
75
+ assert check_stop_strings("hello", 0, ["h"], include_in_output=False) is None
76
+ assert check_stop_strings("hello", 5, [], include_in_output=False) is None
vllm/v1/engine/detokenizer.py
+29
-11
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import sys
4
from abc import ABC, abstractmethod
5
6
import tokenizers
@@ -320,10 +321,19 @@ def check_stop_strings(
321
Where stop_string is the matched stop string and offset is the
322
length to which output_text should be truncated, or -1 for no
323
truncation.
324
+
325
+ When several stop strings match within the newly generated text (for
326
+ example when speculative decoding appends multiple tokens in a single
327
+ step), the stop string that completes earliest in the text is selected,
328
+ so the result matches appending one token at a time. Ties are broken by
329
+ stop-list order.
330
"""
331
if not new_char_count or not stop:
332
return None
333
334
+ best_stop_str: str | None = None
335
+ best_stop_index = 0
336
+ best_end = sys.maxsize
337
for stop_str in stop:
338
stop_string_len = len(stop_str)
339
# Avoid searching already-searched text.
@@ -331,14 +341,22 @@ def check_stop_strings(
341
if stop_index == -1:
342
continue
343
334
- if include_in_output:
335
- # Truncate to end of stop string.
336
- stop_index += stop_string_len
337
- if stop_index >= len(output_text):
338
- # No truncation required.
339
- return stop_str, -1
340
-
341
- # Truncate the output text to either the beginning
342
- # or end of the stop string.
343
- return stop_str, stop_index
344
- return None
344
+ # Prefer the stop string that completes earliest in the text.
345
+ end = stop_index + stop_string_len
346
+ if end < best_end:
347
+ best_stop_str = stop_str
348
+ best_stop_index = stop_index
349
+ best_end = end
350
+
351
+ if best_stop_str is None:
352
+ return None
353
+
354
+ if include_in_output:
355
+ # Truncate to end of stop string.
356
+ if best_end >= len(output_text):
357
+ # No truncation required.
358
+ return best_stop_str, -1
359
+ return best_stop_str, best_end
360
+
361
+ # Truncate the output text to the beginning of the stop string.
362
+ return best_stop_str, best_stop_index