@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 25e02647c2
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 10:15 UTC

[Core] Add optional flags to check for repetitive token patterns in engine output (#35451)

Signed-off-by: aykoppol <aykoppol+git@gmail.com>

aykoppol committed Mar 2, 2026 at 20:23 UTC 25e02647c272eeb0ec076884d046816bdb86159e
7 files changed +433 -2
tests/v1/core/test_repetition_detection.py new
+290
@@ -0,0 +1,290 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +
5 +from vllm.sampling_params import RepetitionDetectionParams, SamplingParams
6 +from vllm.v1.core.sched.utils import check_sequence_repetition, check_stop
7 +from vllm.v1.request import Request, RequestStatus
8 +
9 +pytestmark = pytest.mark.cpu_test
10 +
11 +# ============================================================================
12 +# UNIT TESTS - check_sequence_repetition function
13 +# ============================================================================
14 +
15 +
16 +class TestCheckSequenceRepetition:
17 + """Unit tests for the check_sequence_repetition function"""
18 +
19 + def test_simple_repetition_detected(self):
20 + """Test detection of simple repetitive patterns"""
21 + token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 3]
22 + params = RepetitionDetectionParams(
23 + max_pattern_size=3,
24 + min_pattern_size=2,
25 + min_count=3,
26 + )
27 + assert check_sequence_repetition(token_ids, params)
28 +
29 + def test_repetition_below_min_count(self):
30 + """Test that pattern below min_count is not detected"""
31 + token_ids = [1, 2, 3, 1, 2, 3]
32 + params = RepetitionDetectionParams(
33 + max_pattern_size=3,
34 + min_pattern_size=2,
35 + min_count=3,
36 + )
37 + assert not check_sequence_repetition(token_ids, params)
38 +
39 + def test_two_token_pattern(self):
40 + """Test detection of 2-token patterns"""
41 + token_ids = [1, 2, 1, 2, 1, 2, 1, 2]
42 + params = RepetitionDetectionParams(
43 + max_pattern_size=5,
44 + min_pattern_size=2,
45 + min_count=4,
46 + )
47 + assert check_sequence_repetition(token_ids, params)
48 +
49 + def test_no_repetition_varied_sequence(self):
50 + """Test that non-repetitive sequences are not flagged"""
51 + token_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9]
52 + params = RepetitionDetectionParams(
53 + max_pattern_size=5,
54 + min_pattern_size=2,
55 + min_count=2,
56 + )
57 + assert not check_sequence_repetition(token_ids, params)
58 +
59 + def test_partial_repetition_not_detected(self):
60 + """Test that incomplete repetitions are not detected"""
61 + token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 4]
62 + params = RepetitionDetectionParams(
63 + max_pattern_size=3,
64 + min_pattern_size=2,
65 + min_count=3,
66 + )
67 + assert not check_sequence_repetition(token_ids, params)
68 +
69 + def test_empty_token_list(self):
70 + """Test with empty token list"""
71 + params = RepetitionDetectionParams(
72 + max_pattern_size=3,
73 + min_pattern_size=2,
74 + min_count=2,
75 + )
76 + assert not check_sequence_repetition([], params)
77 +
78 + def test_detection_disabled_max_size_zero(self):
79 + """Test that zero max_pattern_size disables detection"""
80 + token_ids = [1, 2, 1, 2, 1, 2]
81 + params = RepetitionDetectionParams()
82 + assert not check_sequence_repetition(token_ids, params)
83 +
84 + def test_invalid_min_count(self):
85 + """Test that min_count < 2 returns False"""
86 + token_ids = [1, 2, 1, 2]
87 + params = RepetitionDetectionParams()
88 + assert not check_sequence_repetition(token_ids, params)
89 +
90 + def test_repetition_at_end_of_sequence(self):
91 + """Test detection when repetition occurs at the end"""
92 + token_ids = [1, 2, 3, 4, 5, 6, 5, 6, 5, 6]
93 + params = RepetitionDetectionParams(
94 + max_pattern_size=3,
95 + min_pattern_size=2,
96 + min_count=3,
97 + )
98 + assert check_sequence_repetition(token_ids, params)
99 +
100 + def test_large_pattern_many_repetitions(self):
101 + """Test large pattern repeated many times"""
102 + token_ids = [1, 2, 3, 4, 5, 6, 7, 8] * 5
103 + params = RepetitionDetectionParams(
104 + max_pattern_size=10,
105 + min_pattern_size=2,
106 + min_count=3,
107 + )
108 + assert check_sequence_repetition(token_ids, params)
109 +
110 +
111 +# ============================================================================
112 +# INTEGRATION TESTS - check_stop with repetition detection
113 +# ============================================================================
114 +
115 +
116 +class TestRepetitionDetectionIntegration:
117 + """Integration tests for repetition detection in check_stop"""
118 +
119 + def test_basic_repetition_stops_generation(self):
120 + """Test that repetition is detected and stops generation"""
121 + params = SamplingParams(
122 + max_tokens=100,
123 + repetition_detection=RepetitionDetectionParams(
124 + max_pattern_size=5,
125 + min_pattern_size=2,
126 + min_count=3,
127 + ),
128 + )
129 + request = Request(
130 + request_id="test",
131 + prompt_token_ids=[1, 2, 3],
132 + sampling_params=params,
133 + pooling_params=None,
134 + )
135 + request.append_output_token_ids([10, 20, 10, 20, 10, 20])
136 + assert check_stop(request, max_model_len=1024)
137 + assert request.status == RequestStatus.FINISHED_REPETITION
138 + assert request.stop_reason == "repetition_detected"
139 +
140 + def test_detection_disabled_no_stop(self):
141 + """Test that disabled detection doesn't stop generation"""
142 + params = SamplingParams(
143 + max_tokens=100,
144 + )
145 + request = Request(
146 + request_id="test",
147 + prompt_token_ids=[1, 2, 3],
148 + sampling_params=params,
149 + pooling_params=None,
150 + )
151 + request.append_output_token_ids([10, 20, 10, 20, 10, 20])
152 + assert not check_stop(request, max_model_len=1024)
153 +
154 + def test_repetition_respects_min_tokens(self):
155 + """Test that repetition detection respects min_tokens"""
156 + params = SamplingParams(
157 + min_tokens=10,
158 + max_tokens=100,
159 + repetition_detection=RepetitionDetectionParams(
160 + max_pattern_size=5,
161 + min_pattern_size=2,
162 + min_count=3,
163 + ),
164 + )
165 + request = Request(
166 + request_id="test",
167 + prompt_token_ids=[1, 2, 3],
168 + sampling_params=params,
169 + pooling_params=None,
170 + )
171 + request.append_output_token_ids([10, 20, 10, 20, 10, 20])
172 + assert not check_stop(request, max_model_len=1024)
173 +
174 + def test_no_repetition_continues_generation(self):
175 + """Test that non-repetitive tokens don't stop generation"""
176 + params = SamplingParams(
177 + max_tokens=100,
178 + repetition_detection=RepetitionDetectionParams(
179 + max_pattern_size=5,
180 + min_pattern_size=2,
181 + min_count=3,
182 + ),
183 + )
184 + request = Request(
185 + request_id="test",
186 + prompt_token_ids=[1, 2, 3],
187 + sampling_params=params,
188 + pooling_params=None,
189 + )
190 + request.append_output_token_ids([10, 20, 30, 40, 50, 60])
191 + assert not check_stop(request, max_model_len=1024)
192 +
193 + def test_pattern_at_size_boundary(self):
194 + """Test detection at exact pattern size boundary"""
195 + params = SamplingParams(
196 + max_tokens=100,
197 + repetition_detection=RepetitionDetectionParams(
198 + max_pattern_size=3,
199 + min_pattern_size=3,
200 + min_count=2,
201 + ),
202 + )
203 + request = Request(
204 + request_id="test",
205 + prompt_token_ids=[1, 2],
206 + sampling_params=params,
207 + pooling_params=None,
208 + )
209 + request.append_output_token_ids([10, 20, 30, 10, 20, 30])
210 + assert check_stop(request, max_model_len=1024)
211 + assert request.status == RequestStatus.FINISHED_REPETITION
212 +
213 + def test_multiple_pattern_sizes_checked(self):
214 + """Test that function checks pattern sizes in range"""
215 + params = SamplingParams(
216 + max_tokens=100,
217 + repetition_detection=RepetitionDetectionParams(
218 + max_pattern_size=5,
219 + min_pattern_size=2,
220 + min_count=3,
221 + ),
222 + )
223 + request = Request(
224 + request_id="test",
225 + prompt_token_ids=[1],
226 + sampling_params=params,
227 + pooling_params=None,
228 + )
229 + request.append_output_token_ids([7, 8, 9, 10, 7, 8, 9, 10, 7, 8, 9, 10])
230 + assert check_stop(request, max_model_len=1024)
231 + assert request.status == RequestStatus.FINISHED_REPETITION
232 +
233 + def test_eos_takes_precedence_over_repetition(self):
234 + """Test that EOS token stops before repetition check"""
235 + params = SamplingParams(
236 + max_tokens=100,
237 + stop_token_ids=[999],
238 + repetition_detection=RepetitionDetectionParams(
239 + max_pattern_size=5,
240 + min_pattern_size=2,
241 + min_count=3,
242 + ),
243 + )
244 + request = Request(
245 + request_id="test",
246 + prompt_token_ids=[1, 2, 3],
247 + sampling_params=params,
248 + pooling_params=None,
249 + )
250 + request.append_output_token_ids([10, 20, 10, 20, 999])
251 + assert check_stop(request, max_model_len=1024)
252 + assert request.status == RequestStatus.FINISHED_STOPPED
253 +
254 + def test_min_pattern_size_filters_small_patterns(self):
255 + """Test that min_pattern_size filters out smaller patterns"""
256 + params = SamplingParams(
257 + max_tokens=100,
258 + repetition_detection=RepetitionDetectionParams(
259 + max_pattern_size=5,
260 + min_pattern_size=3,
261 + min_count=3,
262 + ),
263 + )
264 + request = Request(
265 + request_id="test",
266 + prompt_token_ids=[1],
267 + sampling_params=params,
268 + pooling_params=None,
269 + )
270 + request.append_output_token_ids([10, 20, 10, 20, 10, 20])
271 + assert not check_stop(request, max_model_len=1024)
272 +
273 + def test_high_repetition_threshold(self):
274 + """Test that high min_count requires many repetitions"""
275 + params = SamplingParams(
276 + max_tokens=100,
277 + repetition_detection=RepetitionDetectionParams(
278 + max_pattern_size=5,
279 + min_pattern_size=2,
280 + min_count=5,
281 + ),
282 + )
283 + request = Request(
284 + request_id="test",
285 + prompt_token_ids=[1],
286 + sampling_params=params,
287 + pooling_params=None,
288 + )
289 + request.append_output_token_ids([10, 20, 10, 20, 10, 20])
290 + assert not check_stop(request, max_model_len=1024)
vllm/entrypoints/openai/chat_completion/protocol.py
+12
@@ -38,6 +38,7 @@ from vllm.logprobs import Logprob
38 from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs
39 from vllm.sampling_params import (
40 BeamSearchParams,
41 + RepetitionDetectionParams,
42 RequestOutputKind,
43 SamplingParams,
44 StructuredOutputsParams,
@@ -336,6 +337,16 @@ class ChatCompletionRequest(OpenAIBaseModel):
337 ),
338 )
339
340 + repetition_detection: RepetitionDetectionParams | None = Field(
341 + default=None,
342 + description="Parameters for detecting repetitive N-gram patterns "
343 + "in output tokens. If such repetition is detected, generation will "
344 + "be ended early. LLMs can sometimes generate repetitive, unhelpful "
345 + "token patterns, stopping only when they hit the maximum output length "
346 + "(e.g. 'abcdabcdabcd...' or '\emoji \emoji \emoji ...'). This feature "
347 + "can detect such behavior and terminate early, saving time and tokens.",
348 + )
349 +
350 # --8<-- [end:chat-completion-extra-params]
351
352 def build_chat_params(
@@ -499,6 +510,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
510 allowed_token_ids=self.allowed_token_ids,
511 extra_args=extra_args or None,
512 skip_clone=True, # Created fresh per request, safe to skip clone
513 + repetition_detection=self.repetition_detection,
514 )
515
516 @model_validator(mode="before")
vllm/entrypoints/openai/completion/protocol.py
+12
@@ -26,6 +26,7 @@ from vllm.logprobs import Logprob
26 from vllm.renderers import TokenizeParams
27 from vllm.sampling_params import (
28 BeamSearchParams,
29 + RepetitionDetectionParams,
30 RequestOutputKind,
31 SamplingParams,
32 StructuredOutputsParams,
@@ -166,6 +167,16 @@ class CompletionRequest(OpenAIBaseModel):
167 ),
168 )
169
170 + repetition_detection: RepetitionDetectionParams | None = Field(
171 + default=None,
172 + description="Parameters for detecting repetitive N-gram patterns "
173 + "in output tokens. If such repetition is detected, generation will "
174 + "be ended early. LLMs can sometimes generate repetitive, unhelpful "
175 + "token patterns, stopping only when they hit the maximum output length "
176 + "(e.g. 'abcdabcdabcd...' or '\emoji \emoji \emoji ...'). This feature "
177 + "can detect such behavior and terminate early, saving time and tokens.",
178 + )
179 +
180 # --8<-- [end:completion-extra-params]
181
182 def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
@@ -310,6 +321,7 @@ class CompletionRequest(OpenAIBaseModel):
321 allowed_token_ids=self.allowed_token_ids,
322 extra_args=extra_args or None,
323 skip_clone=True, # Created fresh per request, safe to skip clone
324 + repetition_detection=self.repetition_detection,
325 )
326
327 @model_validator(mode="before")
vllm/sampling_params.py
+47
@@ -107,6 +107,43 @@ class StructuredOutputsParams:
107 )
108
109
110 +@dataclass
111 +class RepetitionDetectionParams:
112 + """Parameters for detecting repetitive N-gram patterns in output tokens."""
113 +
114 + max_pattern_size: int = 0
115 + """Maximum size of N-gram pattern to detect for sequence repetition.
116 + Set to 0 to disable. Must be used together with min_count."""
117 +
118 + min_pattern_size: int = 0
119 + """Minimum N-gram pattern size to check for sequence repetition.
120 + If set to 0, it defaults to 1.
121 + Must be <= max_pattern_size."""
122 +
123 + min_count: int = 0
124 + """Minimum number of times an N-gram pattern must repeat to trigger
125 + detection. Must be >= 2. Example: 3 for detecting a phrase repeated
126 + 3 times. Must be used together with max_pattern_size."""
127 +
128 + def __post_init__(self):
129 + if (
130 + self.max_pattern_size < 0
131 + or self.min_pattern_size < 0
132 + or self.min_pattern_size > self.max_pattern_size
133 + ):
134 + raise ValueError(
135 + "max_pattern_size, min_pattern_size must be >=0, "
136 + "with min_pattern_size <= max_pattern_size. "
137 + "Set both to 0 to disable repetitive pattern detection."
138 + )
139 + if self.max_pattern_size > 0 and self.min_count < 2:
140 + raise ValueError(
141 + "min_count must be >= 2 to detect repetitive patterns "
142 + "in engine output. If you do not wish to detect repetitive "
143 + "patterns, set max_pattern_size to 0."
144 + )
145 +
146 +
147 class RequestOutputKind(Enum):
148 # Return entire output so far in every RequestOutput
149 CUMULATIVE = 0
@@ -246,6 +283,14 @@ class SamplingParams(
283
284 skip_reading_prefix_cache: bool | None = None
285
286 + repetition_detection: RepetitionDetectionParams | None = None
287 + """Parameters for detecting repetitive N-gram patterns in output tokens.
288 + If such repetition is detected, generation will be ended early. LLMs can
289 + sometimes generate repetitive, unhelpful token patterns, stopping only
290 + when they hit the maximum output length (e.g. 'abcdabcdabcd...' or
291 + '\\emoji \\emoji \\emoji ...'). This feature can detect such behavior
292 + and terminate early, saving time and tokens."""
293 +
294 @staticmethod
295 def from_optional(
296 n: int | None = 1,
@@ -275,6 +320,7 @@ class SamplingParams(
320 allowed_token_ids: list[int] | None = None,
321 extra_args: dict[str, Any] | None = None,
322 skip_clone: bool = False,
323 + repetition_detection: RepetitionDetectionParams | None = None,
324 ) -> "SamplingParams":
325 if logit_bias is not None:
326 # Convert token_id to integer
@@ -314,6 +360,7 @@ class SamplingParams(
360 allowed_token_ids=allowed_token_ids,
361 extra_args=extra_args,
362 skip_clone=skip_clone,
363 + repetition_detection=repetition_detection,
364 )
365
366 def __post_init__(self) -> None:
vllm/v1/core/sched/utils.py
+66
@@ -1,10 +1,64 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import contextlib
4 +from collections.abc import Sequence
5
6 +from vllm.sampling_params import RepetitionDetectionParams
7 from vllm.v1.request import Request, RequestStatus
8
9
10 +def _has_repeating_pattern(
11 + token_ids: Sequence[int],
12 + pattern_len: int,
13 + repetition_min_count: int,
14 +) -> bool:
15 + """Check if the tail of token_ids contains a repeating pattern.
16 +
17 + Compares the last pattern_len tokens against the preceding
18 + (repetition_min_count - 1) repetitions of the same length.
19 + """
20 + for n in range(1, pattern_len + 1):
21 + target_token = token_ids[-n]
22 + for m in range(1, repetition_min_count):
23 + if token_ids[-(pattern_len * m + n)] != target_token:
24 + return False
25 + return True
26 +
27 +
28 +def check_sequence_repetition(
29 + token_ids: Sequence[int],
30 + params: RepetitionDetectionParams,
31 +) -> bool:
32 + """Check if a sequence of token IDs has a repetition pattern.
33 + Args:
34 + token_ids: List of token IDs
35 + params: Repetition detection parameters.
36 + Returns:
37 + True if a repetition pattern is found, False otherwise.
38 + """
39 + max_pattern_size = params.max_pattern_size
40 + min_pattern_size = params.min_pattern_size
41 + min_count = params.min_count
42 +
43 + if min_pattern_size <= 0:
44 + min_pattern_size = 1
45 +
46 + if max_pattern_size <= 0 or min_count < 2 or min_pattern_size > max_pattern_size:
47 + return False
48 +
49 + for pattern_len in range(
50 + min_pattern_size,
51 + max_pattern_size + 1,
52 + ):
53 + if pattern_len * min_count > len(token_ids):
54 + return False
55 +
56 + if _has_repeating_pattern(token_ids, pattern_len, min_count):
57 + return True
58 +
59 + return False
60 +
61 +
62 def remove_all(lst: list, items_to_remove: set) -> list:
63 """Remove all items from a list that are in the items_to_remove set.
64
@@ -61,4 +115,16 @@ def check_stop(request: Request, max_model_len: int) -> bool:
115 ):
116 request.status = RequestStatus.FINISHED_LENGTH_CAPPED
117 return True
118 +
119 + repetition_detection = sampling_params.repetition_detection
120 + if repetition_detection is not None and (
121 + check_sequence_repetition(
122 + request.output_token_ids,
123 + repetition_detection,
124 + )
125 + ):
126 + request.status = RequestStatus.FINISHED_REPETITION
127 + request.stop_reason = "repetition_detected"
128 + return True
129 +
130 return False
vllm/v1/engine/__init__.py
+4 -2
@@ -27,7 +27,7 @@ PauseMode = Literal["abort", "wait", "keep"]
27
28 # These are possible values of RequestOutput.finish_reason,
29 # so form part of the external API.
30 -FINISH_REASON_STRINGS = ("stop", "length", "abort", "error")
30 +FINISH_REASON_STRINGS = ("stop", "length", "abort", "error", "repetition")
31
32 EEP_NOTIFICATION_CALL_ID = -1
33
@@ -41,7 +41,7 @@ class EEPNotificationType(enum.Enum):
41
42 class FinishReason(enum.IntEnum):
43 """
44 - Reason a request finished - stop, length, abort, or error.
44 + Reason a request finished - stop, length, abort, error, or repetition.
45
46 Int rather than Str for more compact serialization.
47
@@ -50,6 +50,7 @@ class FinishReason(enum.IntEnum):
50 abort - aborted by client
51 error - retryable request-level internal error (e.g., KV load failure).
52 Invariant: always converted to 500 Internal Server Error.
53 + repetition - repetitive token pattern detected (hallucination)
54
55 """
56
@@ -57,6 +58,7 @@ class FinishReason(enum.IntEnum):
58 LENGTH = 1
59 ABORT = 2
60 ERROR = 3
61 + REPETITION = 4
62
63 def __str__(self):
64 return FINISH_REASON_STRINGS[self.value]
vllm/v1/request.py
+2
@@ -320,6 +320,7 @@ class RequestStatus(enum.IntEnum):
320 FINISHED_ABORTED = enum.auto()
321 FINISHED_IGNORED = enum.auto()
322 FINISHED_ERROR = enum.auto()
323 + FINISHED_REPETITION = enum.auto()
324
325 def __str__(self) -> str:
326 return self.name
@@ -344,4 +345,5 @@ _FINISHED_REASON_MAP = {
345 RequestStatus.FINISHED_IGNORED: FinishReason.LENGTH,
346 RequestStatus.FINISHED_ERROR: FinishReason.ERROR,
347 RequestStatus.WAITING_FOR_STREAMING_REQ: FinishReason.STOP,
348 + RequestStatus.FINISHED_REPETITION: FinishReason.REPETITION,
349 }