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 6, 03:15 UTC
[Feature] limit thinking tokens (hard limit) (#20859)
Signed-off-by: Sungjae Lee <33976427+llsj14@users.noreply.github.com> Signed-off-by: Sungjae Lee <sung-jae.lee@navercorp.com> Signed-off-by: Chauncey <chaunceyjiang@gmail.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Sungjae Lee committed
Mar 25, 2026 at 01:53 UTC
4731884796d43bd8493769dbfb0e199a2a542092
13 files changed
+702
-12
docs/features/reasoning_outputs.md
+75
@@ -240,6 +240,81 @@ response = client.chat.completions.create(
240
)
241
```
242
243
+## Thinking Budget Control
244
+
245
+Some models, such as [Qwen3](https://qwen.readthedocs.io/en/latest/getting_started/quickstart.html#thinking-budget), [DeepSeek](https://www.alibabacloud.com/help/en/model-studio/deep-thinking), and [Nemotron3](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16), support a thinking budget that limits the maximum number of tokens used for reasoning.
246
+
247
+Token counting starts from `think_start_str`. Once the reasoning token count reaches the configured `thinking_token_budget`, vLLM forces the model to produce `think_end_str`, effectively terminating the reasoning block.
248
+
249
+To use this feature:
250
+
251
+- `--reasoning-parser` enables reasoning extraction.
252
+- `--reasoning-config` defines the reasoning boundary tokens (e.g., `think_start_str`, `think_end_str`).
253
+- `thinking_token_budget` (a sampling parameter) sets the per-request reasoning token limit.
254
+
255
+If `thinking_token_budget` is not specified, no explicit reasoning limit is applied beyond normal generation constraints such as `max_tokens`.
256
+
257
+`--reasoning-config` accepts a JSON object corresponding to
258
+[ReasoningConfig][vllm.config.ReasoningConfig] with the following fields:
259
+
260
+| Field | Type | Description |
261
+|-------------------|----------------|--------------------------------------------------|
262
+| `think_start_str` | `str \| null` | String that marks the start of reasoning content |
263
+| `think_end_str` | `str \| null` | String that marks the end of reasoning content |
264
+
265
+!!! note
266
+ `think_end_str` can include a transition phrase before the think end token. For example, setting `think_end_str` to `"I have to give the solution based on the thinking directly now.</think>"` instructs the model to emit that phrase when the budget is exhausted, making the reasoning termination more natural.
267
+
268
+### Online Serving
269
+
270
+```bash
271
+vllm serve Qwen/Qwen3-0.6B \
272
+ --reasoning-parser qwen3 \
273
+ --reasoning-config '{"think_start_str": "<think>", "think_end_str": "I have to give the solution based on the thinking directly now.</think>"}'
274
+```
275
+
276
+Then make a request with `thinking_token_budget` to limit the reasoning tokens:
277
+
278
+```bash
279
+curl http://localhost:8000/v1/chat/completions \
280
+ -H "Content-Type: application/json" \
281
+ -d '{
282
+ "model": "Qwen/Qwen3-0.6B",
283
+ "messages": [
284
+ { "role": "user", "content": "9.11 and 9.8, which is greater?" }
285
+ ],
286
+ "extra_body": {
287
+ "thinking_token_budget": 10
288
+ }
289
+ }'
290
+```
291
+
292
+### Offline Inference
293
+
294
+```python
295
+from vllm import LLM, SamplingParams
296
+from vllm.config import ReasoningConfig
297
+
298
+llm = LLM(
299
+ model="Qwen/Qwen3-0.6B",
300
+ reasoning_config=ReasoningConfig(
301
+ think_start_str="<think>",
302
+ think_end_str="I have to give the solution based on the thinking directly now.</think>",
303
+ ),
304
+)
305
+
306
+sampling_params = SamplingParams(thinking_token_budget=10)
307
+
308
+messages = [
309
+ {"role": "user", "content": "9.11 and 9.8, which is greater?"},
310
+]
311
+
312
+outputs = llm.chat(messages, sampling_params=sampling_params)
313
+
314
+for output in outputs:
315
+ print("text:", output.outputs[0].text)
316
+```
317
+
318
## Limitations
319
320
- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`).
tests/v1/entrypoints/openai/test_thinking_token_budget.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""E2E tests for thinking_token_budget with reasoning models."""
5
+
6
+import openai
7
+import pytest
8
+import pytest_asyncio
9
+
10
+from tests.utils import RemoteOpenAIServer
11
+
12
+MODEL_NAME = "Qwen/Qwen3-0.6B"
13
+MESSAGES = [{"role": "user", "content": "What is 1+1? Be concise."}]
14
+THINK_BUDGET = 5
15
+
16
+
17
+@pytest.fixture(scope="module")
18
+def server():
19
+ args = [
20
+ "--reasoning-parser",
21
+ "qwen3",
22
+ "--reasoning-config",
23
+ '{"think_start_str": "<think>", "think_end_str": "</think>"}',
24
+ "--max-model-len",
25
+ "2048",
26
+ "--enforce-eager",
27
+ "--no-async-scheduling",
28
+ ]
29
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
30
+ yield remote_server
31
+
32
+
33
+@pytest_asyncio.fixture
34
+async def client(server):
35
+ async with server.get_async_client() as async_client:
36
+ yield async_client
37
+
38
+
39
+@pytest.mark.asyncio
40
+async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI):
41
+ """Test that mixed requests (some with thinking_token_budget, some without)
42
+ complete successfully without errors."""
43
+
44
+ response_with_budget = await client.chat.completions.create(
45
+ model=MODEL_NAME,
46
+ messages=MESSAGES,
47
+ max_tokens=100,
48
+ extra_body={"thinking_token_budget": THINK_BUDGET},
49
+ )
50
+ response_without_budget = await client.chat.completions.create(
51
+ model=MODEL_NAME,
52
+ messages=MESSAGES,
53
+ max_tokens=100,
54
+ )
55
+
56
+ msg_with = response_with_budget.choices[0].message
57
+ msg_without = response_without_budget.choices[0].message
58
+
59
+ assert msg_with.content or getattr(msg_with, "reasoning", None)
60
+ assert msg_without.content or getattr(msg_without, "reasoning", None)
61
+
62
+
63
+@pytest.mark.asyncio
64
+async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI):
65
+ """Test that thinking_token_budget limits the number of reasoning tokens.
66
+
67
+ In streaming mode each reasoning delta corresponds to one token, so
68
+ counting non-empty reasoning_content chunks gives the exact token count.
69
+ """
70
+
71
+ reasoning_token_count = 0
72
+ stream = await client.chat.completions.create(
73
+ model=MODEL_NAME,
74
+ messages=MESSAGES,
75
+ max_tokens=100,
76
+ stream=True,
77
+ extra_body={"thinking_token_budget": THINK_BUDGET},
78
+ )
79
+ async for chunk in stream:
80
+ delta = chunk.choices[0].delta
81
+ if getattr(delta, "reasoning", None):
82
+ reasoning_token_count += 1
83
+
84
+ assert reasoning_token_count == THINK_BUDGET, (
85
+ f"reasoning tokens ({reasoning_token_count}) != "
86
+ f"thinking_token_budget ({THINK_BUDGET})"
87
+ )
tests/v1/logits_processors/test_correctness.py
+173
-10
@@ -30,6 +30,7 @@ from vllm.v1.sample.logits_processor import (
30
MinPLogitsProcessor,
31
MinTokensLogitsProcessor,
32
MoveDirectionality,
33
+ ThinkingTokenBudgetLogitsProcessor,
34
build_logitsprocs,
35
)
36
from vllm.v1.sample.metadata import SamplingMetadata
@@ -47,6 +48,11 @@ MIN_TOKENS_LEN_THRESHOLD = 5
48
REQS_PER_LOGITPROC = 50
49
STR_NO_LOGITPROC = "none"
50
51
+# ThinkingTokenBudgetLogitsProcessor testing constants
52
+THINKING_TOKEN_BUDGET = 5
53
+THINK_START_TOKEN_ID = 999
54
+THINK_END_TOKEN_ID = 998
55
+
56
# LogitsProcessor subclass or "none"
57
LogitprocType: TypeAlias = type[LogitsProcessor] | str
58
@@ -67,9 +73,24 @@ class LogitsProcsRequestParams:
73
self.workload_index = workload_index
74
self.logitproc_type = logitproc_type
75
# Number of output tokens is randomly 0 or twice the min-tokens
70
- # threshold which will be used in testing. Output token values
71
- # don't matter *for these tests* so use 0 as a dummy value
72
- self.out_tokens = [0] * (MIN_TOKENS_LEN_THRESHOLD * random.randint(0, 2))
76
+ # threshold which will be used in testing.
77
+ # Generate diverse random tokens for all processors (more realistic)
78
+ num_tokens = MIN_TOKENS_LEN_THRESHOLD * random.randint(0, 2)
79
+ if num_tokens > 0:
80
+ # Use diverse random tokens
81
+ self.out_tokens = [random.randint(1, 950) for _ in range(num_tokens)]
82
+ # Set first token for ThinkingTokenBudget testing
83
+ is_thinking_processor = (
84
+ logitproc_type is ThinkingTokenBudgetLogitsProcessor
85
+ or (
86
+ hasattr(logitproc_type, "__name__")
87
+ and logitproc_type.__name__ == "ThinkingTokenBudgetLogitsProcessor"
88
+ )
89
+ )
90
+ if is_thinking_processor:
91
+ self.out_tokens[0] = THINK_START_TOKEN_ID
92
+ else:
93
+ self.out_tokens = []
94
self.prompt_tokens = []
95
self.params = _sampling_params_from_logitproc(logitproc_type)
96
@@ -79,6 +100,13 @@ class LogitsProcsRequestParams:
100
return f"MyClass({summ})"
101
102
103
+class MockReasoningConfig:
104
+ """Mock reasoning config for testing ThinkingTokenBudgetLogitsProcessor."""
105
+
106
+ think_start_token_ids = [THINK_START_TOKEN_ID]
107
+ think_end_token_ids = [THINK_END_TOKEN_ID]
108
+
109
+
110
def _generate_fake_sampling_metadata(
111
num_output_tokens: int,
112
batch_size: int,
@@ -97,8 +125,12 @@ def _generate_fake_sampling_metadata(
125
0, vocab_size, size=np.random.randint(1, MAX_NUM_PROMPT_TOKENS)
126
).tolist()
127
)
128
+
129
+ vllm_config = VllmConfig()
130
+ vllm_config.reasoning_config = MockReasoningConfig()
131
+
132
logitsprocs = build_logitsprocs(
101
- vllm_config=VllmConfig(),
133
+ vllm_config=vllm_config,
134
device=device,
135
is_pin_memory=PIN_MEMORY_AVAILABLE,
136
is_pooling_model=False,
@@ -403,6 +435,127 @@ def _min_tokens_validate(
435
)
436
437
438
+def _thinking_budget_params(kwargs: dict) -> None:
439
+ """Set SamplingParams kwargs for thinking token budget tests"""
440
+ kwargs["thinking_token_budget"] = THINKING_TOKEN_BUDGET
441
+
442
+
443
+def _thinking_budget_validate(
444
+ test_fakes: LogitsprocsTestFakes,
445
+ persistent_batch: list[LogitsProcsRequestParams],
446
+ logits_new: torch.Tensor,
447
+ batch_index: int,
448
+ request_params: LogitsProcsRequestParams,
449
+ step_idx: int,
450
+) -> None:
451
+ """Validate thinking token budget processor behavior"""
452
+ # Get the ThinkingTokenBudgetLogitsProcessor instance
453
+ tb_processor: ThinkingTokenBudgetLogitsProcessor = next(
454
+ test_fakes.get_logitsprocs_by_cls(ThinkingTokenBudgetLogitsProcessor)
455
+ )
456
+
457
+ # Get current request state
458
+ state = tb_processor._state.get(batch_index)
459
+ params = request_params.params
460
+
461
+ # Validate thinking token budget configuration
462
+ if hasattr(params, "thinking_token_budget") and params.thinking_token_budget:
463
+ # State should exist for requests with thinking_token_budget
464
+ if state is None:
465
+ _raise_error_invalid(
466
+ msg_suffix=(
467
+ f"Expected state for batch {batch_index} "
468
+ f"with thinking_token_budget={params.thinking_token_budget}"
469
+ ),
470
+ batch_index=batch_index,
471
+ request_params=request_params,
472
+ step_idx=step_idx,
473
+ )
474
+
475
+ # Validate budget matches what was set
476
+ expected_budget = params.thinking_token_budget
477
+ actual_budget = state["thinking_token_budget"]
478
+
479
+ if actual_budget != expected_budget:
480
+ _raise_error_invalid(
481
+ msg_suffix=(
482
+ f"Budget mismatch: expected {expected_budget}, got {actual_budget}"
483
+ ),
484
+ batch_index=batch_index,
485
+ request_params=request_params,
486
+ step_idx=step_idx,
487
+ )
488
+
489
+ # Check if we're in thinking mode and validate token counting
490
+ output_tokens = request_params.out_tokens
491
+
492
+ # Find if thinking has started in output tokens
493
+ thinking_started = False
494
+ start_tokens = tb_processor.think_start_token_ids
495
+
496
+ if len(start_tokens) > 0:
497
+ for i in range(len(output_tokens) - len(start_tokens) + 1):
498
+ if output_tokens[i : i + len(start_tokens)] == start_tokens:
499
+ thinking_started = True
500
+ break
501
+
502
+ if thinking_started:
503
+ # If budget is exceeded, validate end token forcing
504
+ think_count = state["think_count"]
505
+ budget = state["thinking_token_budget"]
506
+
507
+ if think_count >= budget:
508
+ if not state["in_end"]:
509
+ _raise_error_invalid(
510
+ msg_suffix=(
511
+ f"Budget exceeded ({think_count} >= "
512
+ f"{budget}) but not "
513
+ "forcing end tokens"
514
+ ),
515
+ batch_index=batch_index,
516
+ request_params=request_params,
517
+ step_idx=step_idx,
518
+ )
519
+
520
+ # Validate that only end tokens are allowed
521
+ end_tokens = tb_processor.think_end_token_ids
522
+ if len(end_tokens) > 0:
523
+ expected_end_token_id = end_tokens[
524
+ min(state["end_count"], len(end_tokens) - 1)
525
+ ]
526
+
527
+ # Check logits masking
528
+ batch_logits = logits_new[batch_index]
529
+ for token_id in range(len(batch_logits)):
530
+ logit_value = batch_logits[token_id]
531
+
532
+ if token_id == expected_end_token_id:
533
+ # End token should not be masked
534
+ if logit_value == -float("inf"):
535
+ _raise_error_invalid(
536
+ msg_suffix=(
537
+ f"End token {token_id} should not be "
538
+ "masked but is"
539
+ ),
540
+ batch_index=batch_index,
541
+ request_params=request_params,
542
+ step_idx=step_idx,
543
+ )
544
+ else:
545
+ # All other tokens should be masked when forcing end
546
+ if logit_value != -float("inf"):
547
+ _raise_error_invalid(
548
+ msg_suffix=(
549
+ f"Token {token_id} should be masked "
550
+ f"when forcing end tokens, but "
551
+ f"logit={logit_value}"
552
+ ),
553
+ batch_index=batch_index,
554
+ request_params=request_params,
555
+ step_idx=step_idx,
556
+ )
557
+
558
+
559
def _none_validate(
560
test_fakes: LogitsprocsTestFakes,
561
persistent_batch: list[LogitsProcsRequestParams],
@@ -449,20 +602,30 @@ logitsprocs_test_mapping = {
602
MinTokensLogitsProcessor: LogitsprocTestHelpers(
603
gen_request_fxn=_min_tokens_params, eval_fxn=_min_tokens_validate
604
),
605
+ ThinkingTokenBudgetLogitsProcessor: LogitsprocTestHelpers(
606
+ gen_request_fxn=_thinking_budget_params, eval_fxn=_thinking_budget_validate
607
+ ),
608
}
609
610
611
def _get_test_cases() -> list[list[str]]:
612
"""Each test case is a set of logitsprocs"""
613
logitsprocs_types = list(logitsprocs_test_mapping.keys())
614
+
615
+ # Isolate ThinkingTokenBudgetLogitsProcessor from all other processors
616
+ # to avoid unexpected modification of logits interference
617
+ thinking_processor = ThinkingTokenBudgetLogitsProcessor
618
+ other_processors = [
619
+ p
620
+ for p in logitsprocs_types
621
+ if p != STR_NO_LOGITPROC and p != thinking_processor
622
+ ]
623
+
624
return (
625
[[STR_NO_LOGITPROC]]
460
- + [
461
- [logitproc_type, STR_NO_LOGITPROC]
462
- for logitproc_type in logitsprocs_types
463
- if logitproc_type != STR_NO_LOGITPROC
464
- ]
465
- + [logitsprocs_types]
626
+ + [[logitproc_type, STR_NO_LOGITPROC] for logitproc_type in other_processors]
627
+ + [other_processors]
628
+ + [[thinking_processor]]
629
)
630
631
vllm/config/__init__.py
+3
@@ -33,6 +33,7 @@ from vllm.config.offload import (
33
from vllm.config.parallel import EPLBConfig, ParallelConfig
34
from vllm.config.pooler import PoolerConfig
35
from vllm.config.profiler import ProfilerConfig
36
+from vllm.config.reasoning import ReasoningConfig
37
from vllm.config.scheduler import SchedulerConfig
38
from vllm.config.speculative import SpeculativeConfig
39
from vllm.config.speech_to_text import SpeechToTextConfig
@@ -101,6 +102,8 @@ __all__ = [
102
"ParallelConfig",
103
# From vllm.config.pooler
104
"PoolerConfig",
105
+ # From vllm.config.reasoning
106
+ "ReasoningConfig",
107
# From vllm.config.scheduler
108
"SchedulerConfig",
109
# From vllm.config.speculative
vllm/config/reasoning.py
new
+72
@@ -0,0 +1,72 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import field
5
+
6
+from vllm.config.model import ModelConfig
7
+from vllm.config.utils import config
8
+from vllm.tokenizers import cached_tokenizer_from_config
9
+
10
+
11
+@config
12
+class ReasoningConfig:
13
+ """Configuration for reasoning models.
14
+
15
+ Set `think_start_str` and `think_end_str` to the strings that delimit
16
+ the reasoning block (e.g. `"<think>"` and `"</think>"`). The
17
+ corresponding token IDs are derived automatically via
18
+ `initialize_token_ids` and are not intended to be set directly.
19
+ """
20
+
21
+ # NOTE: These parameters are temporary, the intent is to derive them
22
+ # automatically from the reasoning parser in a future version.
23
+ think_start_str: str = "<think>"
24
+ """String that indicates the start of reasoning."""
25
+ think_end_str: str = "</think>"
26
+ """String that indicates the end of reasoning content."""
27
+
28
+ _think_start_token_ids: list[int] | None = field(
29
+ default=None, init=False, repr=False
30
+ )
31
+ """Private backing field for `think_start_token_ids`. Set by
32
+ `initialize_token_ids`. Not intended to be configured directly."""
33
+ _think_end_token_ids: list[int] | None = field(default=None, init=False, repr=False)
34
+ """Private backing field for `think_end_token_ids`. Set by
35
+ `initialize_token_ids`. Not intended to be configured directly."""
36
+
37
+ @property
38
+ def think_start_token_ids(self) -> list[int] | None:
39
+ """Token IDs derived from `think_start_str`. Set automatically by
40
+ `initialize_token_ids`. Not intended to be configured directly."""
41
+ return self._think_start_token_ids
42
+
43
+ @property
44
+ def think_end_token_ids(self) -> list[int] | None:
45
+ """Token IDs derived from `think_end_str`. Set automatically by
46
+ `initialize_token_ids`. Not intended to be configured directly."""
47
+ return self._think_end_token_ids
48
+
49
+ def initialize_token_ids(self, model_config: ModelConfig) -> None:
50
+ """Initialize reasoning token IDs from strings using the tokenizer."""
51
+ if (
52
+ self._think_start_token_ids is not None
53
+ and self._think_end_token_ids is not None
54
+ ):
55
+ return
56
+
57
+ tokenizer = cached_tokenizer_from_config(model_config=model_config)
58
+
59
+ self._think_start_token_ids = tokenizer.encode(
60
+ self.think_start_str, add_special_tokens=False
61
+ )
62
+ self._think_end_token_ids = tokenizer.encode(
63
+ self.think_end_str, add_special_tokens=False
64
+ )
65
+
66
+ if not self._think_start_token_ids or not self._think_end_token_ids:
67
+ raise ValueError(
68
+ f"ReasoningConfig: failed to tokenize reasoning strings: "
69
+ f"think_start_str='{self.think_start_str}', "
70
+ f"think_end_str='{self.think_end_str}'. "
71
+ "Ensure the strings are valid tokens in the model's vocabulary."
72
+ )
vllm/config/vllm.py
+6
@@ -40,6 +40,7 @@ from .observability import ObservabilityConfig
40
from .offload import OffloadConfig
41
from .parallel import ParallelConfig
42
from .profiler import ProfilerConfig
43
+from .reasoning import ReasoningConfig
44
from .scheduler import SchedulerConfig
45
from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig
46
from .structured_outputs import StructuredOutputsConfig
@@ -302,6 +303,8 @@ class VllmConfig: # type: ignore[misc]
303
"""The configurations for event publishing."""
304
ec_transfer_config: ECTransferConfig | None = None
305
"""The configurations for distributed EC cache transfer."""
306
+ reasoning_config: ReasoningConfig | None = None
307
+ """The configurations for reasoning model."""
308
# some opaque config, only used to provide additional information
309
# for the hash computation, mainly used for testing, debugging or out of
310
# tree config registration.
@@ -1143,6 +1146,9 @@ class VllmConfig: # type: ignore[misc]
1146
if not self.instance_id:
1147
self.instance_id = random_uuid()[:5]
1148
1149
+ if self.reasoning_config is not None and self.model_config is not None:
1150
+ self.reasoning_config.initialize_token_ids(self.model_config)
1151
+
1152
# Hybrid KV cache manager (HMA) runtime rules:
1153
# - Explicit enable (--no-disable-kv-cache-manager): error if runtime
1154
# disables it
vllm/engine/arg_utils.py
+4
@@ -53,6 +53,7 @@ from vllm.config import (
53
PoolerConfig,
54
PrefetchOffloadConfig,
55
ProfilerConfig,
56
+ ReasoningConfig,
57
SchedulerConfig,
58
SpeculativeConfig,
59
StructuredOutputsConfig,
@@ -581,6 +582,7 @@ class EngineArgs:
582
kv_events_config: KVEventsConfig | None = None
583
584
ec_transfer_config: ECTransferConfig | None = None
585
+ reasoning_config: ReasoningConfig = get_field(VllmConfig, "reasoning_config")
586
587
generation_config: str = ModelConfig.generation_config
588
enable_sleep_mode: bool = ModelConfig.enable_sleep_mode
@@ -1297,6 +1299,7 @@ class EngineArgs:
1299
vllm_group.add_argument(
1300
"--attention-config", "-ac", **vllm_kwargs["attention_config"]
1301
)
1302
+ vllm_group.add_argument("--reasoning-config", **vllm_kwargs["reasoning_config"])
1303
vllm_group.add_argument("--kernel-config", **vllm_kwargs["kernel_config"])
1304
vllm_group.add_argument(
1305
"--additional-config", **vllm_kwargs["additional_config"]
@@ -1958,6 +1961,7 @@ class EngineArgs:
1961
kv_transfer_config=self.kv_transfer_config,
1962
kv_events_config=self.kv_events_config,
1963
ec_transfer_config=self.ec_transfer_config,
1964
+ reasoning_config=self.reasoning_config,
1965
profiler_config=self.profiler_config,
1966
additional_config=self.additional_config,
1967
optimization_level=self.optimization_level,
vllm/entrypoints/openai/chat_completion/protocol.py
+2
@@ -180,6 +180,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
180
| None
181
) = "none"
182
reasoning_effort: Literal["none", "low", "medium", "high"] | None = None
183
+ thinking_token_budget: int | None = None
184
include_reasoning: bool = True
185
parallel_tool_calls: bool | None = True
186
@@ -515,6 +516,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
516
structured_outputs=self.structured_outputs,
517
logit_bias=self.logit_bias,
518
bad_words=self.bad_words,
519
+ thinking_token_budget=self.thinking_token_budget,
520
allowed_token_ids=self.allowed_token_ids,
521
extra_args=extra_args or None,
522
skip_clone=True, # Created fresh per request, safe to skip clone
vllm/sampling_params.py
+5
@@ -281,6 +281,8 @@ class SamplingParams(
281
_bad_words_token_ids: list[list[int]] | None = None
282
283
skip_reading_prefix_cache: bool | None = None
284
+ thinking_token_budget: int | None = None
285
+ """Maximum number of tokens allowed for thinking operations."""
286
287
repetition_detection: RepetitionDetectionParams | None = None
288
"""Parameters for detecting repetitive N-gram patterns in output tokens.
@@ -304,6 +306,7 @@ class SamplingParams(
306
stop: str | list[str] | None = None,
307
stop_token_ids: list[int] | None = None,
308
bad_words: list[str] | None = None,
309
+ thinking_token_budget: int | None = None,
310
include_stop_str_in_output: bool = False,
311
ignore_eos: bool = False,
312
max_tokens: int | None = 16,
@@ -344,6 +347,7 @@ class SamplingParams(
347
stop=stop,
348
stop_token_ids=stop_token_ids,
349
bad_words=bad_words,
350
+ thinking_token_budget=thinking_token_budget,
351
include_stop_str_in_output=include_stop_str_in_output,
352
ignore_eos=ignore_eos,
353
max_tokens=max_tokens,
@@ -858,6 +862,7 @@ class SamplingParams(
862
f"stop={self.stop}, "
863
f"stop_token_ids={self.stop_token_ids}, "
864
f"bad_words={self.bad_words}, "
865
+ f"thinking_token_budget={self.thinking_token_budget}, "
866
f"include_stop_str_in_output={self.include_stop_str_in_output}, "
867
f"ignore_eos={self.ignore_eos}, "
868
f"max_tokens={self.max_tokens}, "
vllm/v1/engine/input_processor.py
+10
@@ -99,6 +99,16 @@ class InputProcessor:
99
self.structured_outputs_config,
100
self.tokenizer,
101
)
102
+
103
+ if (
104
+ params.thinking_token_budget is not None
105
+ and self.vllm_config.reasoning_config is None
106
+ ):
107
+ raise ValueError(
108
+ "thinking_token_budget is set but reasoning_config is "
109
+ "not configured. Please set --reasoning-config to use "
110
+ "thinking_token_budget."
111
+ )
112
elif isinstance(params, PoolingParams):
113
supported_pooling_tasks = [
114
task for task in supported_tasks if task in POOLING_TASKS
vllm/v1/sample/logits_processor/__init__.py
+3
@@ -18,6 +18,7 @@ from vllm.v1.sample.logits_processor.builtin import (
18
LogitBiasLogitsProcessor,
19
MinPLogitsProcessor,
20
MinTokensLogitsProcessor,
21
+ ThinkingTokenBudgetLogitsProcessor,
22
process_dict_updates,
23
)
24
from vllm.v1.sample.logits_processor.interface import (
@@ -50,6 +51,7 @@ BUILTIN_LOGITS_PROCESSORS: list[type[LogitsProcessor]] = [
51
MinTokensLogitsProcessor,
52
LogitBiasLogitsProcessor,
53
MinPLogitsProcessor,
54
+ ThinkingTokenBudgetLogitsProcessor,
55
]
56
57
@@ -354,4 +356,5 @@ __all__ = [
356
"STR_POOLING_REJECTS_LOGITSPROCS",
357
"LOGITSPROCS_GROUP",
358
"AdapterLogitsProcessor",
359
+ "ThinkingTokenBudgetLogitsProcessor",
360
]
vllm/v1/sample/logits_processor/builtin.py
+258
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Callable, Sequence
4
-from typing import TYPE_CHECKING, TypeVar
4
+from typing import TYPE_CHECKING, Any, TypeVar
5
6
import numpy as np
7
import torch
@@ -291,6 +291,263 @@ class MinTokensLogitsProcessor(LogitsProcessor):
291
return logits
292
293
294
+class ThinkingTokenBudgetLogitsProcessor(LogitsProcessor):
295
+ """Limits the number of tokens allowed inside a 'thinking' section."""
296
+
297
+ def __init__(
298
+ self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool
299
+ ):
300
+ reasoning_config = vllm_config.reasoning_config
301
+ max_num_reqs = vllm_config.scheduler_config.max_num_seqs
302
+
303
+ # Check if thinking is enabled
304
+ self.is_enabled = reasoning_config is not None
305
+
306
+ self.think_start_token_ids = getattr(
307
+ reasoning_config, "think_start_token_ids", []
308
+ )
309
+ self.think_end_token_ids = getattr(reasoning_config, "think_end_token_ids", [])
310
+
311
+ self.pin_memory = is_pin_memory
312
+ self.device = device
313
+ # Per-request state tracking for thinking token management
314
+ # Key: request_index, Value: state dict containing:
315
+ # "in_think": bool - currently in thinking mode
316
+ # "in_end": bool - currently forcing end tokens output
317
+ # "check_count_down": int - steps remaining until next think
318
+ # start/end token parsing
319
+ # "think_count": int - number of thinking tokens generated
320
+ # "end_count": int - number of end tokens forced so far
321
+ # "thinking_token_budget": int - max allowed thinking tokens
322
+ # "output_tok_ids": list[int] - generated output tokens
323
+ # "prev_output_length": int - previous output length for
324
+ # incremental processing
325
+ self._state: dict[int, dict[str, Any]] = {}
326
+
327
+ # Preallocate reusable tensors
328
+ self.mask = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
329
+ self.force_token_ids = torch.full(
330
+ (max_num_reqs,), -1, dtype=torch.long, device=device
331
+ )
332
+
333
+ @staticmethod
334
+ def _find_last_sequence_index(target_list: list[int], token_ids: list[int]) -> int:
335
+ """
336
+ Returns the index of the last occurrence of token_ids in target_list.
337
+
338
+ Args:
339
+ target_list (list[int]): The list of token IDs.
340
+ token_ids (list[int]): The sequence of token IDs to find.
341
+ """
342
+ if not token_ids:
343
+ return -1
344
+ for i in range(len(target_list) - len(token_ids), -1, -1):
345
+ if target_list[i : i + len(token_ids)] == token_ids:
346
+ return i
347
+ return -1
348
+
349
+ def _init_state_entry(
350
+ self, prompt_tok_ids: list[int] | None, thinking_token_budget: int
351
+ ) -> dict[str, Any]:
352
+ """Initializes the tracking state for a given sequence index."""
353
+ if prompt_tok_ids is None:
354
+ last_start = -1
355
+ last_end = -1
356
+ in_think = False
357
+ think_count = 0
358
+ else:
359
+ last_start = self._find_last_sequence_index(
360
+ prompt_tok_ids, self.think_start_token_ids
361
+ )
362
+ last_end = self._find_last_sequence_index(
363
+ prompt_tok_ids, self.think_end_token_ids
364
+ )
365
+ in_think = last_start > last_end
366
+ if in_think:
367
+ think_count = len(prompt_tok_ids) - (
368
+ last_start + len(self.think_start_token_ids)
369
+ )
370
+ else:
371
+ think_count = 0
372
+
373
+ return {
374
+ "in_think": in_think, # Currently in thinking mode
375
+ "in_end": in_think and thinking_token_budget == 0,
376
+ "check_count_down": thinking_token_budget,
377
+ "think_count": think_count, # Number of tokens in thinking section
378
+ "end_count": 0, # Number of end tokens forced so far
379
+ "prompt_tok_ids": prompt_tok_ids,
380
+ "output_tok_ids": [],
381
+ "thinking_token_budget": thinking_token_budget,
382
+ "prev_output_length": 0,
383
+ # Track previous output length for incremental updates
384
+ }
385
+
386
+ def _update_think_state(self, state: dict[str, Any]):
387
+ """Updates the state based on newly generated output tokens."""
388
+ if not state.get("in_end", False) and state.get("check_count_down", 0) > 0:
389
+ state["check_count_down"] -= 1
390
+ return
391
+
392
+ output = state.get("output_tok_ids", [])
393
+ if not output:
394
+ return
395
+
396
+ # Track previous output length for incremental processing
397
+ prev_length = state.get("prev_output_length", 0)
398
+ current_length = len(output)
399
+
400
+ if current_length <= prev_length:
401
+ return
402
+
403
+ # Process only newly added tokens
404
+ new_tokens = output[prev_length:]
405
+ state["prev_output_length"] = current_length
406
+
407
+ # Check if new tokens contain think start or end sequences
408
+ start_len = len(self.think_start_token_ids)
409
+ end_len = len(self.think_end_token_ids)
410
+
411
+ # Look for think sequences in recent tokens (including boundary)
412
+ # Check overlapping regions where sequences might span boundaries
413
+ check_start_idx = max(0, prev_length - max(start_len, end_len) + 1)
414
+ recent_tokens = output[check_start_idx:]
415
+
416
+ # Find any think start/end sequences in recent tokens
417
+ recent_start_pos = self._find_last_sequence_index(
418
+ recent_tokens, self.think_start_token_ids
419
+ )
420
+ recent_end_pos = self._find_last_sequence_index(
421
+ recent_tokens, self.think_end_token_ids
422
+ )
423
+
424
+ # Update state based on recent sequences
425
+ if not state["in_end"]:
426
+ if recent_start_pos >= 0 and recent_end_pos >= 0:
427
+ if recent_start_pos > recent_end_pos:
428
+ # Case: ...<end>...<start>... - entering think mode
429
+ absolute_start_pos = check_start_idx + recent_start_pos
430
+ new_think_count = current_length - (absolute_start_pos + start_len)
431
+ state["in_think"] = True
432
+ state["think_count"] = new_think_count
433
+ else:
434
+ # Case: ...<start>...<end>... - exiting think mode
435
+ state["in_think"] = False
436
+ state["think_count"] = 0
437
+ elif recent_start_pos >= 0:
438
+ # Found think start - entering think mode
439
+ absolute_start_pos = check_start_idx + recent_start_pos
440
+ new_think_count = current_length - (absolute_start_pos + start_len)
441
+ state["in_think"] = True
442
+ state["think_count"] = new_think_count
443
+ elif recent_end_pos >= 0:
444
+ # Found think end - exiting think mode
445
+ state["in_think"] = False
446
+ state["think_count"] = 0
447
+ elif state["in_think"]:
448
+ # Continue thinking mode, increment count by new tokens
449
+ state["think_count"] += len(new_tokens)
450
+
451
+ # Set countdown based on current state
452
+ if state["in_think"]:
453
+ remaining_budget = max(
454
+ 0, state["thinking_token_budget"] - state["think_count"]
455
+ )
456
+ state["check_count_down"] = max(0, remaining_budget - 1)
457
+ else:
458
+ state["check_count_down"] = state["thinking_token_budget"]
459
+
460
+ # Check if need to transition to end mode
461
+ if (
462
+ state["in_think"]
463
+ and state["think_count"] >= state["thinking_token_budget"]
464
+ ):
465
+ state["in_think"] = False
466
+ state["in_end"] = True
467
+ state["end_count"] = 0
468
+ state["check_count_down"] = state["thinking_token_budget"]
469
+ else:
470
+ # In end mode
471
+ state["end_count"] += 1
472
+ if state["end_count"] >= len(self.think_end_token_ids):
473
+ state.update(
474
+ {
475
+ "in_end": False,
476
+ "end_count": 0,
477
+ "check_count_down": state["thinking_token_budget"],
478
+ }
479
+ )
480
+
481
+ def is_argmax_invariant(self) -> bool:
482
+ """This logits processor can change the outcome of
483
+ greedy sampling by forcing that the thinking section
484
+ ends after a certain number of tokens."""
485
+ return False
486
+
487
+ def update_state(self, batch_update: BatchUpdate | None):
488
+ if not self.is_enabled:
489
+ return
490
+ if batch_update:
491
+ for index, params, prompt_tok_ids, output_tok_ids in batch_update.added:
492
+ thinking_token_budget = params.thinking_token_budget
493
+
494
+ if thinking_token_budget is not None:
495
+ self._state[index] = self._init_state_entry(
496
+ prompt_tok_ids, thinking_token_budget
497
+ )
498
+ self._state[index]["output_tok_ids"] = output_tok_ids
499
+ else:
500
+ # Remove state if no thinking budget
501
+ self._state.pop(index, None)
502
+
503
+ for index in batch_update.removed:
504
+ self._state.pop(index, {})
505
+
506
+ for i1, i2, direction in batch_update.moved:
507
+ if direction == MoveDirectionality.SWAP:
508
+ state1 = self._state.pop(i1, None)
509
+ state2 = self._state.pop(i2, None)
510
+ if state1 is not None:
511
+ self._state[i2] = state1
512
+ if state2 is not None:
513
+ self._state[i1] = state2
514
+ else:
515
+ state = self._state.pop(i1, None)
516
+ if state is not None:
517
+ self._state[i2] = state
518
+
519
+ for state in self._state.values():
520
+ self._update_think_state(state)
521
+
522
+ def apply(self, logits: torch.Tensor) -> torch.Tensor:
523
+ if not self.is_enabled or not self._state:
524
+ return logits
525
+
526
+ batch_size = logits.size(0)
527
+ self.mask[:batch_size] = False
528
+
529
+ for i in range(batch_size):
530
+ state = self._state.get(i)
531
+ if state and state["in_end"]:
532
+ self.mask[i] = True
533
+ self.force_token_ids[i] = self.think_end_token_ids[state["end_count"]]
534
+
535
+ # Check in CPU first not to sync with GPU
536
+ has_active_thinking = any(
537
+ state.get("in_end", False) for state in self._state.values()
538
+ )
539
+
540
+ if has_active_thinking:
541
+ current_mask = self.mask[:batch_size]
542
+ active_indices = current_mask.nonzero(as_tuple=False).view(-1)
543
+ if len(active_indices) > 0:
544
+ force_tokens = self.force_token_ids[active_indices]
545
+ # Apply a large value for the end thinking token id index
546
+ logits[active_indices, force_tokens] = 1e9
547
+
548
+ return logits
549
+
550
+
551
def process_dict_updates(
552
req_entries: dict[int, T],
553
batch_update: BatchUpdate | None,
vllm/v1/worker/gpu_model_runner.py
+4
-1
@@ -629,7 +629,10 @@ class GPUModelRunner(
629
),
630
# We currently don't know whether a particular custom logits processor
631
# uses output token ids so we set this conservatively.
632
- logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
632
+ # ThinkingTokenBudgetLogitsProcessor also needs output token ids to
633
+ # correctly track think start/end token sequences in async scheduling.
634
+ logitsprocs_need_output_token_ids=bool(custom_logitsprocs)
635
+ or self.vllm_config.reasoning_config is not None,
636
is_pooling_model=self.is_pooling_model,
637
cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
638
)