Add max_thinking_tokens for reasoning models (issue #42111)
Author: AndresAlgaba
Pull request: yes

Add feature described in Issue #42111 

# What does this PR do?
A built-in way to cap how many tokens a reasoning model spends inside its ``<think> … </think>`` block. Today, we can only control the total response length via ``max_new_tokens``. No parameter limits the internal reasoning segment when ``enable_thinking=True``.

Motivation:
- Reasoning models (e.g., Qwen3 series) often produce very long thought blocks, which can blow past latency budgets before the final answer starts.
- Users need a simple, model-agnostic control to bound that “thinking” cost without disabling reasoning entirely.
- The Qwen docs (https://qwen.readthedocs.io/en/latest/getting_started/quickstart.html#thinking-budget) already describe a brute-force approach (two-step generation) to implement “thinking budgets”.

This PR:
- Extends ``GenerationConfig`` with:
``max_thinking_tokens``: integer budget for reasoning tokens.
``begin_thinking_token_id / end_thinking_token_id``: marker IDs so generation knows where the thinking span begins/ends.
- Add a ``MaxThinkingTokensLogitsProcessor`` that watches the active ``<think>`` block. Once the budget is reached, it forces end_thinking_token_id, ensuring the model exits reasoning and continues with the final response.
- Document the new parameter in reasoning-model guides (EXAONE, CWM, etc.) and show how to wire the thinking-token IDs until configs do it automatically.
- Provide unit coverage so ``_get_logits_processor`` injects the new processor whenever the config is fully specified.

Fixes #42111 


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ArthurZucker @Cyrilvallez @zucchini-nlp

Comments:
One thing that still bothers me is that:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3-4B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")

messages = [{"role": "user", "content": "Give me a concise proof that sqrt(2) is irrational."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)

think_start = tokenizer.convert_tokens_to_ids("<think>")
think_end = tokenizer.convert_tokens_to_ids("</think>")
model.generation_config.begin_thinking_token_id = think_start
model.generation_config.end_thinking_token_id = think_end

out = model.generate(
    **inputs,
    max_new_tokens=128,
    max_thinking_tokens=64,
    do_sample=False,
)

print(tokenizer.decode(out[0], skip_special_tokens=False))
```

will result in: 
```python

user
Give me a concise proof that sqrt(2) is irrational.


assistant
<think>
Okay, so I need to prove that the square root of 2 is irrational. Hmm, I remember that an irrational number is a number that can't be expressed as a fraction of two integers. So, I need to show that there are no integers a and b where a/b equals sqrt(2). Let</think>

To prove that √2 is irrational, we use a **proof by contradiction**. Here's a concise version:

1. **Assume** √2 is rational. Then it can be written as a fraction $ \frac{a}{b} $, where $ a $ and $ b
```
where the ```end_thinking_token``` is appended to the output, whereas by default ```\n<\think>``` is generated by the model, resulting in:
```python

user
Give me a concise proof that sqrt(2) is irrational.


assistant
<think>
Okay, so I need to prove that the square root of 2 is irrational. Hmm, I remember that an irrational number is a number that can't be expressed as a fraction of two integers. So, I need to show that there are no integers a and b where a/b equals sqrt(2). Let
</think>

To prove that √2 is irrational, we use a **proof by contradiction**. Here's a concise version:

1. **Assume** √2 is rational. Then it can be written as a fraction $ \frac{a}{b} $, where $ a $ and $ b
```

Additionally, the early stopping of the Qwen example (https://qwen.readthedocs.io/en/latest/getting_started/quickstart.html#thinking-budget) adds:
```\n\nConsidering the limited time by the user, I have to give the solution based on the thinking directly now.\n</think>\n\n```

Not sure whether such customization should be allowed? Kind of related to the comment above, but seems like a different logic and maybe for a future PR?

I might have overcomplicated the ```MaxThinkingTokensLogitsProcessor```, but it is to avoid token overcount when 1) people use ```<think>``` in their user prompt, or 2) people feed previous reasoning traces which include ```<think>...</think>``` (although this is not best practice, see number 4 in https://huggingface.co/Qwen/Qwen3-4B#best-practices).

Another issue is when ```enable_thinking=True``` pre-fills the prompt with a ```<think>``` token as this becomes part of the ```prompt_length``` and will be ignored. Any thoughts on this? A simple solution here would be ```prompt_length-1```, but I noticed that, for example, Qwen3-4B doesn't do the pre-filling.

While testing this, another problem is that the model might output a second ```<\think>``` token after ```\n\n``` (so not directly after the enforced ```<\think>``` token):

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3-4B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")

messages = [
    {
        "role": "user",
        "content": "Give me a concise proof that sqrt(2) is irrational.",
    },
]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)

think_start = tokenizer.convert_tokens_to_ids("<think>")
think_end = tokenizer.convert_tokens_to_ids("</think>")
model.generation_config.begin_thinking_token_id = think_start
model.generation_config.end_thinking_token_id = think_end

out = model.generate(
    **inputs,
    max_new_tokens=64,
    max_thinking_tokens=32,
    do_sample=False,
)

print(tokenizer.decode(out[0], skip_special_tokens=False))
```

```python

user
Give me a concise proof that sqrt(2) is irrational.


assistant
<think>
Okay, so I need to prove that the square root of 2 is irrational. Hmm, I remember that an irrational number is a number that can't</think>

</think>

Sure! Here's a concise proof that $\sqrt{2}$ is irrational:

---

**Proof that $\sqrt{2}$ is irrational
```



Update should satisfy following requirements:

- When a user uses past reasoning traces (i.e., previous assistant chunks contain `<think> ... </think>`), it should ignore this.
- When a user uses `<think>`, `</think>`, or `<think> ... </think>`, in their prompt, it should be ignored.
- It should count the number of tokens only in the current reasoning chain.
It should be robust against the model generating additional `<think>` tokens in the current reasoning chain (not restart the count).
- It should be robust against the model automatically appending the `<think>` token when `enable_thinking=True` (right now it becomes part of the prompt and gets ignored because of it).
- The model sometimes generates an additional `</think>` token after the enforced closure `</think>` token. Not always directly after, sometimes it first generated `\n\n` and then the `</think>` token. Right now, the `</think>` token can only be generated if a new `<think>` token has been generated.

@AndresAlgaba hey!
Great job, thanks a lot for the PR! I think it a super useful feature to have. I left a few general comment and question above, mostly to make sure I understand the flow and to simplify it when possible

I'm actually not sure about this PR! I think forcibly ending model thinking mid-stream will force them out of distribution and cause them to output very low-quality answers. We could add it as a feature for advanced users, but I worry that if we do that then people will assume that because it's a feature that it works well in practice, when I think it probably won't, right?

Note that in the [Qwen example](https://qwen.readthedocs.io/en/latest/getting_started/quickstart.html#thinking-budget), they don't cut off thinking directly, but instead insert an "early stopping prompt" that is likely fine-tuned for Qwen and will not work well on other thinking models.

@zucchini-nlp Thank you for the feedback, working on it!

@Rocketknight1 Thanks for the comment! Indeed, I agree it is more of an experimental feature. It is in line with Gemini's thinking_budget (https://ai.google.dev/gemini-api/docs/thinking#set-budget- and Claude's budget_tokens (https://docs.claude.com/en/docs/build-with-claude/extended-thinking#how-to-use-extended-thinking). I remember that DeepSeek was planning to implement a max_thinking_tokens parameter at some point, but I can't find it anymore. I implemented it for my own research and thought others might be interested in this feature, but with the warranty that it is experimental. So not sure whether it should then be added to the library like this?

hi @AndresAlgaba, yes, but note that that feature in Gemini or Claude does not actually cap the reasoning tokens or forcibly stop thinking - they specifically say the model can overflow the budget. They've probably just trained the model so that it (roughly) obeys thinking budgets, and then they probably tell it the budget in the prompt somewhere

True, same for the reasoning effort in OpenAI models. I still believe some people might be interested in the feature, see a similar feature request in vllm: https://github.com/vllm-project/vllm/issues/17887 (and PR: https://github.com/vllm-project/vllm/pull/20859, where they also propose a logit processor). Is there any way we can indicate it's experimental? Is it an option to keep the processor without adding the```max_thinking_tokens``` argument for now?

I believe that adding the ```\n\nConsidering the limited time by the user, I have to give the solution based on the thinking directly now.\n</think>\n\n``` part is interesting, but part of different PR?

Hmn, if we can mention in the docs that it's experimental and might reduce answer quality it's probably okay! cc @zucchini-nlp, are you happy with that?

Thanks again for the feedback! I have made some changes after further experimentation and have updated the docs (including a warning that the feature is experimental).

Seems like a good idea! I agree that hard truncation is often not the way to go. Thank you for offering help, I will look into it and keep you posted here!


Hi @zucchini-nlp, sorry for not following up on this but I decided to look further into the "early stopping" schemes and write a paper on it (now available on arxiv: https://arxiv.org/abs/2601.23163). I will experiment further and then start to add to the ```transformers-community``` org!

Amazing @AndresAlgaba ! If you have a paper/blog-post, feel free to link it directly in the hub page via `hub papers` :)