[Feature]: Limit thinking tokens

[gohar94]

🚀 The feature, motivation and pitch

Is there plan to support an option (e.g., max_thinking_tokens) to limit the thinking tokens for reasoning models? Perhaps both a minimum and a maximum limit would be useful. This is an example of one approach to achieving this: https://simonwillison.net/2025/Jan/22/r1py/

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[csvance]

If I understand these models correctly, it may be possible to already to do this with two separate generation prompts / LLM calls, albeit with the limitation that the limit is characters not tokens. You will have to turn off vLLM reasoning model support and use your own parser on the application side.

First prompt starts with the assistant message <thinking> and uses regex constrained output like .{0,1024}</thinking>. Then the second chat template starts with assistant message <thinking>{reasoning_content}</thinking> for generation. Maybe you need to replace <thinking> and </thinking> with the appropriate tokens for your model, but if you do I'm not sure how that would work with the regular expression.

Of course it would be nice to be able to not have to work against the built in reasoning model support and have a token length limit instead of a character one. This feature would make reasoning models so much easier to use for many applications where we could have appropriate constraints for reasoning tokens depending on the task.

[yelaman-starrez]

I have come up with logits processor approach:

from typing import List, Optional
import torch

class ThinkLogitsProcessor:
    """A logits processor that limit the number of thinking tokens."""
    
    def __init__(self, think_start_token, think_end_token, num_think_tokens: int = 100):
        """
        Initialize the think logits processor.
        
        Args:
            tokenizer: The tokenizer used for the model
            num_think_tokens: Maximum number of tokens allowed in thinking section
        """
        self.num_think_tokens = num_think_tokens
        self.think_start_token = think_start_token
        self.think_end_token = think_end_token
        
    def __call__(
        self,
        input_ids: List[int],
        logits: torch.Tensor,
    ) -> torch.Tensor:
        """
        Process the logits to enforce </think> token when needed.
        
        Args:
            input_ids: List of input token IDs.
            logits: Tensor of logits for the next token.
            
        Returns:
            Processed logits tensor.
        """
        # Check if we're in a thinking section
        if self.think_start_token in input_ids and self.think_end_token not in input_ids:
            # Find the position of the last <think> token
            think_start_pos = len(input_ids) - 1 - input_ids[::-1].index(self.think_start_token)
            
            # Calculate number of tokens since <think>
            tokens_since_think = len(input_ids) - think_start_pos - 1
            
            # If we've reached the maximum thinking length, force </think>
            if tokens_since_think >= self.num_think_tokens:
                # Set all other logits to -inf except for </think>
                logits = torch.full_like(logits, float('-inf'))
                logits[self.think_end_token] = 1.0
                
        return logits 

And then use it in sampling params:

sampling_params = SamplingParams(
            temperature=temperature,
            max_tokens=max_tokens,
            logits_processors=[ThinkLogitsProcessor(think_start_token, think_end_token, num_think_tokens)]
        )

[Kenwwww]

I have come up with logits processor approach:

from typing import List, Optional
import torch

class ThinkLogitsProcessor:
    """A logits processor that limit the number of thinking tokens."""
    
    def __init__(self, think_start_token, think_end_token, num_think_tokens: int = 100):
        """
        Initialize the think logits processor.
        
        Args:
            tokenizer: The tokenizer used for the model
            num_think_tokens: Maximum number of tokens allowed in thinking section
        """
        self.num_think_tokens = num_think_tokens
        self.think_start_token = think_start_token
        self.think_end_token = think_end_token
        
    def __call__(
        self,
        input_ids: List[int],
        logits: torch.Tensor,
    ) -> torch.Tensor:
        """
        Process the logits to enforce </think> token when needed.
        
        Args:
            input_ids: List of input token IDs.
            logits: Tensor of logits for the next token.
            
        Returns:
            Processed logits tensor.
        """
        # Check if we're in a thinking section
        if self.think_start_token in input_ids and self.think_end_token not in input_ids:
            # Find the position of the last <think> token
            think_start_pos = len(input_ids) - 1 - input_ids[::-1].index(self.think_start_token)
            
            # Calculate number of tokens since <think>
            tokens_since_think = len(input_ids) - think_start_pos - 1
            
            # If we've reached the maximum thinking length, force </think>
            if tokens_since_think >= self.num_think_tokens:
                # Set all other logits to -inf except for </think>
                logits = torch.full_like(logits, float('-inf'))
                logits[self.think_end_token] = 1.0
                
        return logits 

And then use it in sampling params:

sampling_params = SamplingParams(
            temperature=temperature,
            max_tokens=max_tokens,
            logits_processors=[ThinkLogitsProcessor(think_start_token, think_end_token, num_think_tokens)]
        )

how to use it in openai server？

[vadimkantorov]

Also think this is a crucial feature, as sometimes a bad checkpoint might get into looping or other bad modes, so forcing it to finish with a forced </think> or </think> Now let's formulate the final answer:, and then letting it finish for also some limited number of tokens would be very convenient

[csvance]

Being able to limit thinking tokens would allow my application to use the same LLM for utility tasks as it does for generation. The amount of time a model might spend thinking can vary wildly especially when using constrained output. You might get paragraphs of tokens during thinking only to generate a simple r'(Yes|No)' . Not having control over this means I have to tie up more GPU memory for additional models.

[dhruvmullick]

I tried a similar logit processor approach, but this leads to slowdowns because vLLM v1 doesn't support logit processors yet.

[renjie0]

I tried a similar logit processor approach, but this leads to slowdowns because vLLM v1 doesn't support logit processors yet.

Why does it cause slowdown? Is it hard to make logit processor run fast in v1?

[samsja]

I also think this would be a nice feature to have following this paper https://arxiv.org/abs/2505.05315

[aarnphm]

Once we have #16728 merged, this should be easier to be implemented here. We can probably use this to mimic OpenAI's reasoning_effort feature 🤔

[llsj14]

I’m currently implementing this feature, and I’ve encountered some problems as follows.

First, extending the builtin LogitsProcessor has been challenging because its apply method doesn’t take sampling_metadata as an argument, so it’s not straightforward to handle input or output tokens within the current design.
I found that I can use a reference to the output token IDs, so I may need to add a reference to the input token IDs in the same way.

Passing reasoning token IDs from the reasoning parser and tokenizer to the LogitsProcessor through the InputBatch also feels misaligned with the current codebase’s philosophy and might require some refactoring.

Second, we could consider token-level limits for DeepSeek-R1 or Qwen reasoning models. However, after reviewing the Graphite reasoning parser, I found that its reasoning tokens are not represented as single tokens. So for now, I’m limiting the thinking tokens only for DeepSeek-R1 and Qwen models.

I hope to share my drafts soon and get your feedback.