feat: limit thinking tokens

Signed-off-by: Sungjae Lee <33976427+llsj14@users.noreply.github.com>

Author: llsj14

vllm/config.py

@@ -4404,6 +4404,17 @@ def set_splitting_ops_for_v1(self):
                 "vllm.unified_attention_with_output",
             ]
 
+class ReasoningConfig:
+    """Configuration for reasoning models."""
+
+    think_start_token_id: Optional[int] = None
+    """Token ID that indicates the start of reasoning."""
+    think_end_token_id: Optional[int] = None
+    """Token ID that indicates the end of reasoning."""
+
+    def __init__(self, think_start_token_id: Optional[int] = None, think_end_token_id: Optional[int] = None):
+        self.think_start_token_id = think_start_token_id
+        self.think_end_token_id = think_end_token_id
 
 @config
 @dataclass(config=ConfigDict(arbitrary_types_allowed=True))
@@ -4461,6 +4472,8 @@ class VllmConfig:
     # some opaque config, only used to provide additional information
     # for the hash computation, mainly used for testing, debugging or out of
     # tree config registration.
+    reasoning_config: Optional[ReasoningConfig] = None
+    """The configurations for reasoning model."""
     additional_config: Union[dict, SupportsHash] = field(default_factory=dict)
     """Additional config for specified platform. Different platforms may
     support different configs. Make sure the configs are valid for the platform

vllm/entrypoints/openai/protocol.py

@@ -404,6 +404,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
     prompt_logprobs: Optional[int] = None
     allowed_token_ids: Optional[list[int]] = None
     bad_words: list[str] = Field(default_factory=list)
+    max_think_tokens: Optional[int] = None
     # --8<-- [end:chat-completion-sampling-params]
 
     # --8<-- [start:chat-completion-extra-params]
@@ -670,6 +671,7 @@ def to_sampling_params(
             guided_decoding=guided_decoding,
             logit_bias=self.logit_bias,
             bad_words= self.bad_words,
+            max_think_tokens=self.max_think_tokens,
             allowed_token_ids=self.allowed_token_ids,
             extra_args=extra_args or None,
         )

vllm/reasoning/deepseek_r1_reasoning_parser.py

@@ -23,8 +23,8 @@ class DeepSeekR1ReasoningParser(ReasoningParser):
     text. This parser extracts the reasoning content from the model output.
     """
 
-    start_token_id: int
-    end_token_id: int
+    think_start_token_id: int
+    think_end_token_id: int
 
     start_token: str = "<think>"
     end_token: str = "</think>"
@@ -37,24 +37,24 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase):
                 "The model tokenizer must be passed to the ReasoningParser "
                 "constructor during construction.")
 
-        self.start_token_id = self.vocab.get(self.start_token)
-        self.end_token_id = self.vocab.get(self.end_token)
-        if self.start_token_id is None or self.end_token_id is None:
+        self.think_start_token_id = self.vocab.get(self.start_token)
+        self.think_end_token_id = self.vocab.get(self.end_token)
+        if self.think_start_token_id is None or self.think_end_token_id is None:
             raise RuntimeError(
                 "DeepSeek R1 reasoning parser could not locate think start/end "
                 "tokens in the tokenizer!")
 
     def is_reasoning_end(self, input_ids: list[int]) -> bool:
-        return self.end_token_id in input_ids
+        return self.think_end_token_id in input_ids
 
     def extract_content_ids(self, input_ids: list[int]) -> list[int]:
         """
         Extract the content after the end tokens
         """
-        if self.end_token_id not in input_ids[:-1]:
+        if self.think_end_token_id not in input_ids[:-1]:
             return []
         else:
-            return input_ids[input_ids.index(self.end_token_id) + 1:]
+            return input_ids[input_ids.index(self.think_end_token_id) + 1:]
 
     def extract_reasoning_content_streaming(
         self,
@@ -75,14 +75,14 @@ def extract_reasoning_content_streaming(
         """
         # Skip single special tokens
         if len(delta_token_ids) == 1 and (delta_token_ids[0] in [
-                self.start_token_id, self.end_token_id
+                self.think_start_token_id, self.think_end_token_id
         ]):
             return None
 
         # Check if <think> is present in previous or delta.
         # Keep compatibility with models that don't generate <think> tokens.
-        if self.start_token_id in previous_token_ids:
-            if self.end_token_id in delta_token_ids:
+        if self.think_start_token_id in previous_token_ids:
+            if self.think_end_token_id in delta_token_ids:
                 # <think> in previous, </think> in delta,
                 # extract reasoning content
                 end_index = delta_text.find(self.end_token)
@@ -92,16 +92,16 @@ def extract_reasoning_content_streaming(
                     reasoning_content=reasoning_content,
                     content=content if content else None,
                 )
-            elif self.end_token_id in previous_token_ids:
+            elif self.think_end_token_id in previous_token_ids:
                 # <think> in previous, </think> in previous,
                 # reasoning content continues
                 return DeltaMessage(content=delta_text)
             else:
                 # <think> in previous, no </think> in previous or delta,
                 # reasoning content continues
                 return DeltaMessage(reasoning_content=delta_text)
-        elif self.start_token_id in delta_token_ids:
-            if self.end_token_id in delta_token_ids:
+        elif self.think_start_token_id in delta_token_ids:
+            if self.think_end_token_id in delta_token_ids:
                 # <think> in delta, </think> in delta, extract reasoning content
                 start_index = delta_text.find(self.start_token)
                 end_index = delta_text.find(self.end_token)
@@ -120,7 +120,7 @@ def extract_reasoning_content_streaming(
             # No <think> in previous or delta, also need to check for </think>.
             # Because the model may have generated </think> without <think>
             # Ref https://huggingface.co/deepseek-ai/DeepSeek-R1/commit/8a58a132790c9935686eb97f042afa8013451c9f
-            if self.end_token_id in delta_token_ids:
+            if self.think_end_token_id in delta_token_ids:
                 # </think> in delta with more tokens,
                 # extract reasoning content and content
                 end_index = delta_text.find(self.end_token)
@@ -130,7 +130,7 @@ def extract_reasoning_content_streaming(
                     reasoning_content=reasoning_content,
                     content=content if content else None,
                 )
-            elif self.end_token_id in previous_token_ids:
+            elif self.think_end_token_id in previous_token_ids:
                 # </think> in previous, thinking content ends
                 return DeltaMessage(content=delta_text)
             else:

vllm/sampling_params.py

@@ -200,6 +200,7 @@ class SamplingParams(
         extra_args: Arbitrary additional args, that can be used by custom
             sampling implementations, plugins, etc. Not used by any in-tree
             sampling implementations.
+        max_think_tokens: Maximum number of tokens allowed for thinking
     """
 
     n: int = 1
@@ -248,6 +249,9 @@ class SamplingParams(
     bad_words: Optional[list[str]] = None
     _bad_words_token_ids: Optional[list[list[int]]] = None
 
+    # Maximum number of tokens allowed for thinking operations.
+    max_think_tokens: Optional[int] = None
+
     @staticmethod
     def from_optional(
         n: Optional[int] = 1,
@@ -263,6 +267,7 @@ def from_optional(
         stop: Optional[Union[str, list[str]]] = None,
         stop_token_ids: Optional[list[int]] = None,
         bad_words: Optional[list[str]] = None,
+        max_think_tokens: Optional[int] = None,
         include_stop_str_in_output: bool = False,
         ignore_eos: bool = False,
         max_tokens: Optional[int] = 16,
@@ -306,6 +311,7 @@ def from_optional(
             stop=stop,
             stop_token_ids=stop_token_ids,
             bad_words=bad_words,
+            max_think_tokens=max_think_tokens,
             include_stop_str_in_output=include_stop_str_in_output,
             ignore_eos=ignore_eos,
             max_tokens=max_tokens,
@@ -574,6 +580,7 @@ def __repr__(self) -> str:
             f"stop={self.stop}, "
             f"stop_token_ids={self.stop_token_ids}, "
             f"bad_words={self.bad_words}, "
+            f"max_think_tokens={self.max_think_tokens}, "
             f"include_stop_str_in_output={self.include_stop_str_in_output}, "
             f"ignore_eos={self.ignore_eos}, "
             f"max_tokens={self.max_tokens}, "

vllm/v1/engine/core.py

@@ -86,6 +86,7 @@ def __init__(self,
         self.collective_rpc("initialize_cache",
                             args=(num_gpu_blocks, num_cpu_blocks))
 
+        # EngineCore holds StructuredOutputManager to handle and it has vllm config as an arg.
         self.structured_output_manager = StructuredOutputManager(vllm_config)
 
         # Setup scheduler.

vllm/v1/sample/logits_processor.py

@@ -12,6 +12,7 @@
 from torch._prims_common import DeviceLikeType
 
 from vllm import PoolingParams, SamplingParams
+from vllm.config import ReasoningConfig
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
@@ -24,9 +25,9 @@ class MoveDirectionality(Enum):
     SWAP = 1
 
 
-# (index, params, output_tok_ids) tuples for new
+# (index, params, prompt_tok_ids, output_tok_ids) tuples for new
 # requests added to the batch.
-AddedRequest = tuple[int, Union[SamplingParams, PoolingParams], list[int]]
+AddedRequest = tuple[int, Union[SamplingParams, PoolingParams], list[int], list[int]]
 # (index 1, index 2, directionality) tuples representing
 # one-way moves or two-way swaps of requests in batch
 MovedRequest = tuple[int, int, MoveDirectionality]
@@ -43,9 +44,9 @@ class BatchUpdate:
     # within the persistent batch.
     #
     # Note: each added request is represented as
-    # (index, params, output_tok_ids)
-    # Key assumption: output_tok_ids is a reference to the
-    # request's running output tokens list; in this way
+    # (index, params, prompt_tok_ids, output_tok_ids)
+    # Key assumption: prompt_tok_ids, output_tok_ids is a reference to the
+    # request's prompt and running output tokens list; in this way
     # the logits processors always see the latest list of
     # generated tokens
     removed: Sequence[RemovedRequest]
@@ -260,7 +261,7 @@ def update_state(self, batch_update: Optional[BatchUpdate]):
 
         needs_update = False
         # Process added requests.
-        for index, params, _ in batch_update.added:
+        for index, params, _, _ in batch_update.added:
             min_p = params.min_p if isinstance(params, SamplingParams) else 0.0
             if self.min_p_cpu[index] != min_p:
                 needs_update = True
@@ -337,7 +338,7 @@ def update_state(self, batch_update: Optional[BatchUpdate]):
 
         # Process added requests.
         needs_update = bool(batch_update.added)
-        for index, params, _ in batch_update.added:
+        for index, params, _, _ in batch_update.added:
             if isinstance(params, SamplingParams) and (lb :=
                                                        params.logit_bias):
                 self.biases[index] = lb
@@ -420,7 +421,7 @@ def update_state(self, batch_update: Optional[BatchUpdate]):
         if batch_update:
             # Process added requests.
             needs_update |= bool(batch_update.added)
-            for index, params, output_tok_ids in batch_update.added:
+            for index, params, _, output_tok_ids in batch_update.added:
                 if (isinstance(params, SamplingParams)
                         and (min_tokens := params.min_tokens)
                         and len(output_tok_ids) < min_tokens):
@@ -493,8 +494,113 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor:
         return logits
 
 
+class MaxThinkTokensLogitsProcessor(LogitsProcessor):
+    """A logits processor that limits the maximum number of thinking tokens."""
+
+    def __init__(self, reasoning_config: ReasoningConfig, pin_memory: bool, device: torch.device):
+        """
+        Args:
+            think_start_token_id (int): Token ID for the start of thinking section.
+            think_end_token_id (int): Token ID for the end of thinking section.
+            pin_memory (bool): Whether to use pinned memory for tensors.
+            device (torch.device): Device to use for tensor operations.
+        """
+        super().__init__()
+        self.think_start_token_id = reasoning_config.think_start_token_id
+        self.think_end_token_id = reasoning_config.think_end_token_id
+        self.pin_memory = pin_memory
+        self.device = device
+        self._state = {}
+
+    def _find_last_token_index(self, tokens, token_id):
+        try:
+            return len(tokens) - tokens[::-1].index(token_id) - 1
+        except ValueError:
+            return -1
+
+    def is_argmax_invariant(self) -> bool:
+        """This logits processor can change the outcome of greedy sampling
+        by forcing that the thinking section ends after a certain number of tokens."""
+        return False
+
+    def update_state(self, batch_update: Optional[BatchUpdate]):
+        if batch_update is None:
+            return
+
+        for index, params, prompt_tok_ids, output_tok_ids in batch_update.added:
+            max_think_tokens = params.max_think_tokens if isinstance(params, SamplingParams) else None
+
+            if max_think_tokens is None:
+                continue
+
+            last_think_start_idx = self._find_last_token_index(prompt_tok_ids, self.think_start_token_id)
+            last_think_end_idx = self._find_last_token_index(prompt_tok_ids, self.think_end_token_id)
+
+            in_think = False
+            count = 0
+
+            if last_think_start_idx > last_think_end_idx:
+                in_think = True
+                count = len(prompt_tok_ids) - (last_think_start_idx + 1)
+
+            self._state[index] = {
+                "in_think": in_think,
+                "count": count,
+                "prompt_tok_ids": prompt_tok_ids,
+                "output_tok_ids": output_tok_ids,
+                "max_think_tokens": max_think_tokens,
+            }
+
+        for index in batch_update.removed:
+            self._state.pop(index, None)
+
+        for i1, i2, direction in batch_update.moved:
+            if direction == MoveDirectionality.SWAP:
+                self._state[i1], self._state[i2] = self._state[i2], self._state[i1]
+            else:
+                self._state[i2] = self._state.pop(i1, None)
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        batch_size = logits.size(0)
+        if batch_size == 0:
+            return logits
+
+        mask = torch.zeros(batch_size, dtype=torch.bool, device=logits.device)
+        end_token_id = self.think_end_token_id
+
+        for index in range(batch_size):
+            state = self._state.get(index, None)
+            if not state or not state.get("output_tok_ids"):
+                continue
+
+            last_tok = state["output_tok_ids"][-1]
+            in_think = state["in_think"]
+            count = state["count"]
+
+            if last_tok == self.think_start_token_id:
+                in_think = True
+                count = 0
+            elif last_tok == self.think_end_token_id:
+                in_think = False
+                count = 0
+            elif in_think:
+                count += 1
+
+            state["in_think"] = in_think
+            state["count"] = count
+
+            if state["in_think"] and state["count"] >= state["max_think_tokens"]:
+                mask[index] = True
+
+        if mask.any():
+            logits[mask] = -float("inf")
+            logits[mask, end_token_id] = 0.0
+
+        return logits
+
+
 def init_builtin_logitsprocs(pin_memory_available: bool, max_num_reqs: int,
-                             device: torch.device) -> LogitsProcessorManager:
+                             device: torch.device, reasoning_config: ReasoningConfig) -> LogitsProcessorManager:
     """Construct 'builtin' vLLM logitsprocs which the engine
     loads by default.
 
@@ -516,10 +622,16 @@ def init_builtin_logitsprocs(pin_memory_available: bool, max_num_reqs: int,
         device=device,
         # +1 for temporary swap space
         max_num_reqs=max_num_reqs + 1)
+    max_think_tokens_logitproc = MaxThinkTokensLogitsProcessor(
+        reasoning_config=reasoning_config,
+        pin_memory=pin_memory_available,
+        device=device,
+    )
     return LogitsProcessorManager(
         non_argmax_invariant=[
             min_tokens_logitproc,
             logit_bias_logitproc,
+            max_think_tokens_logitproc
         ],
         argmax_invariant=[min_p_logitproc],
     )