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 2, 20:45 UTC
[Reasoning] [Bugfix] Reject invalid thinking_token_budget values (#43402)
Signed-off-by: linzm1007 <linzm1007@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
linzm1007 committed
May 26, 2026 at 18:37 UTC
739af5c7e1603b6c9a4d376b52dc7f4b19e1403d
5 files changed
+152
-5
tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py
new
+76
@@ -0,0 +1,76 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+from pydantic import ValidationError
6
+
7
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
8
+from vllm.entrypoints.openai.completion.protocol import CompletionRequest
9
+
10
+
11
+@pytest.mark.parametrize("raw_value", [-2, 0.6, 10.5])
12
+def test_chat_completion_request_rejects_invalid_thinking_token_budget(raw_value):
13
+ with pytest.raises(ValidationError, match="thinking_token_budget"):
14
+ ChatCompletionRequest.model_validate(
15
+ {
16
+ "model": "qwen",
17
+ "messages": [{"role": "user", "content": "hello"}],
18
+ "thinking_token_budget": raw_value,
19
+ }
20
+ )
21
+
22
+
23
+def test_chat_completion_request_accepts_valid_thinking_token_budget():
24
+ request = ChatCompletionRequest.model_validate(
25
+ {
26
+ "model": "qwen",
27
+ "messages": [{"role": "user", "content": "hello"}],
28
+ "thinking_token_budget": 10,
29
+ }
30
+ )
31
+ assert request.thinking_token_budget == 10
32
+
33
+
34
+def test_chat_completion_request_accepts_minus_one_as_unlimited():
35
+ request = ChatCompletionRequest.model_validate(
36
+ {
37
+ "model": "qwen",
38
+ "messages": [{"role": "user", "content": "hello"}],
39
+ "thinking_token_budget": -1,
40
+ }
41
+ )
42
+ assert request.thinking_token_budget is None
43
+
44
+
45
+@pytest.mark.parametrize("raw_value", [0.6, 3.14, -2])
46
+def test_completion_request_rejects_invalid_thinking_token_budget(raw_value):
47
+ with pytest.raises(ValidationError, match="thinking_token_budget"):
48
+ CompletionRequest.model_validate(
49
+ {
50
+ "model": "qwen",
51
+ "prompt": "hello",
52
+ "thinking_token_budget": raw_value,
53
+ }
54
+ )
55
+
56
+
57
+def test_completion_request_accepts_valid_thinking_token_budget():
58
+ request = CompletionRequest.model_validate(
59
+ {
60
+ "model": "qwen",
61
+ "prompt": "hello",
62
+ "thinking_token_budget": 5,
63
+ }
64
+ )
65
+ assert request.thinking_token_budget == 5
66
+
67
+
68
+def test_completion_request_accepts_minus_one_as_unlimited():
69
+ request = CompletionRequest.model_validate(
70
+ {
71
+ "model": "qwen",
72
+ "prompt": "hello",
73
+ "thinking_token_budget": -1,
74
+ }
75
+ )
76
+ assert request.thinking_token_budget is None
tests/v1/logits_processors/test_correctness.py
+35
-1
@@ -20,7 +20,7 @@ from tests.v1.sample.utils import (
20
)
21
from vllm.config import VllmConfig
22
from vllm.platforms import current_platform
23
-from vllm.sampling_params import SamplingParams
23
+from vllm.sampling_params import SamplingParams, validate_thinking_token_budget
24
from vllm.utils.platform_utils import is_pin_memory_available
25
from vllm.v1.sample.logits_processor import (
26
BatchUpdate,
@@ -1194,3 +1194,37 @@ def test_thinking_budget_enforced_without_penalties():
1194
"Budget exceeded: in_end should be True so that apply_to_logits "
1195
"forces the end token"
1196
)
1197
+
1198
+
1199
+@pytest.mark.parametrize(
1200
+ ("raw_value", "expected"),
1201
+ [
1202
+ (None, None),
1203
+ (-1, None),
1204
+ (10, 10),
1205
+ (0, 0),
1206
+ ],
1207
+)
1208
+def test_validate_thinking_token_budget(raw_value, expected):
1209
+ assert validate_thinking_token_budget(raw_value) == expected
1210
+
1211
+
1212
+def test_sampling_params_minus_one_normalizes_to_none():
1213
+ params = SamplingParams(thinking_token_budget=-1)
1214
+ assert params.thinking_token_budget is None
1215
+
1216
+
1217
+@pytest.mark.parametrize("invalid_budget", [-2, 0.6, 10.5, True])
1218
+def test_validate_thinking_token_budget_rejects_invalid(invalid_budget):
1219
+ from vllm.exceptions import VLLMValidationError
1220
+
1221
+ with pytest.raises(VLLMValidationError, match="thinking_token_budget"):
1222
+ validate_thinking_token_budget(invalid_budget)
1223
+
1224
+
1225
+@pytest.mark.parametrize("invalid_budget", [-2, 0.6, 10.5])
1226
+def test_thinking_budget_invalid_budget_rejected(invalid_budget):
1227
+ from vllm.exceptions import VLLMValidationError
1228
+
1229
+ with pytest.raises(VLLMValidationError, match="thinking_token_budget"):
1230
+ SamplingParams(thinking_token_budget=invalid_budget)
vllm/entrypoints/openai/chat_completion/protocol.py
+2
-1
@@ -41,6 +41,7 @@ from vllm.sampling_params import (
41
RequestOutputKind,
42
SamplingParams,
43
StructuredOutputsParams,
44
+ ThinkingTokenBudget,
45
)
46
from vllm.utils import random_uuid
47
@@ -225,7 +226,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
226
"part of the standard OpenAI API specification."
227
),
228
)
228
- thinking_token_budget: int | None = None
229
+ thinking_token_budget: ThinkingTokenBudget = None
230
include_reasoning: bool = True
231
parallel_tool_calls: bool | None = True
232
vllm/entrypoints/openai/completion/protocol.py
+4
-2
@@ -29,6 +29,7 @@ from vllm.sampling_params import (
29
RequestOutputKind,
30
SamplingParams,
31
StructuredOutputsParams,
32
+ ThinkingTokenBudget,
33
)
34
from vllm.utils import random_uuid
35
@@ -185,11 +186,12 @@ class CompletionRequest(OpenAIBaseModel):
186
"can detect such behavior and terminate early, saving time and tokens.",
187
)
188
188
- thinking_token_budget: int | None = Field(
189
+ thinking_token_budget: ThinkingTokenBudget = Field(
190
default=None,
191
description=(
192
"Maximum number of tokens allowed for thinking operations "
192
- "(reasoning models). -1 = unlimited."
193
+ "(reasoning models). Non-negative integer sets the limit; "
194
+ "-1 means unlimited (treated as unset)."
195
),
196
)
197
vllm/sampling_params.py
+35
-1
@@ -7,9 +7,10 @@ import json as json_mod
7
from dataclasses import field
8
from enum import Enum, IntEnum
9
from functools import cached_property
10
-from typing import Any
10
+from typing import Annotated, Any
11
12
import msgspec
13
+from pydantic import BeforeValidator
14
from pydantic.dataclasses import dataclass
15
16
import vllm.envs as envs
@@ -30,6 +31,35 @@ MAX_LOGPROB_TOKEN_IDS = 128
31
the per-request row width allocated by the sampler's `LogprobTokenIdsState`."""
32
33
34
+def validate_thinking_token_budget(value: int | float | bool | None) -> int | None:
35
+ """Validate ``thinking_token_budget``; return ``None`` if unset."""
36
+ if value is None:
37
+ return None
38
+ if isinstance(value, (bool, float)) or not isinstance(value, int):
39
+ raise VLLMValidationError(
40
+ "`thinking_token_budget` must be a non-negative integer "
41
+ "or -1 for unlimited.",
42
+ parameter="thinking_token_budget",
43
+ value=value,
44
+ )
45
+ if value == -1:
46
+ return None
47
+ if value < 0:
48
+ raise VLLMValidationError(
49
+ "`thinking_token_budget` must be a non-negative integer "
50
+ "or -1 for unlimited.",
51
+ parameter="thinking_token_budget",
52
+ value=value,
53
+ )
54
+ return value
55
+
56
+
57
+ThinkingTokenBudget = Annotated[
58
+ int | None,
59
+ BeforeValidator(validate_thinking_token_budget),
60
+]
61
+
62
+
63
class SamplingType(IntEnum):
64
GREEDY = 0
65
RANDOM = 1
@@ -409,6 +439,10 @@ class SamplingParams(
439
if self.seed == -1:
440
self.seed = None
441
442
+ self.thinking_token_budget = validate_thinking_token_budget(
443
+ self.thinking_token_budget
444
+ )
445
+
446
if self.stop is None:
447
self.stop = []
448
elif isinstance(self.stop, str):