@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d0df145c2a
Fork 0

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 14, 05:30 UTC

Add Olmo 3 reasoning parser (#26054)

Signed-off-by: Luca Soldaini <luca@soldaini.net>

Luca Soldaini committed Oct 4, 2025 at 02:48 UTC d0df145c2a760f915df0c4fd19a5e3517b64f598
3 files changed +453
tests/reasoning/test_olmo3_reasoning_parser.py new
+157
@@ -0,0 +1,157 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +from transformers import AutoTokenizer
6 +
7 +from tests.reasoning.utils import run_reasoning_extraction
8 +from vllm.reasoning import ReasoningParser, ReasoningParserManager
9 +
10 +parser_name = "olmo3"
11 +START_REASONING = "<think>"
12 +END_REASONING = "</think>"
13 +
14 +NO_REASONING = {
15 + "output": f"{START_REASONING}{END_REASONING}No thoughts, head empty!",
16 + "reasoning_content": None,
17 + "content": "No thoughts, head empty!",
18 +}
19 +
20 +NO_REASONING_WITH_NEWLINE = {
21 + "output":
22 + f"{START_REASONING}\n{END_REASONING}\n\nNo thoughts, head empty!",
23 + "reasoning_content": "\n",
24 + "content": "\n\nNo thoughts, head empty!",
25 +}
26 +
27 +SIMPLE_REASONING = {
28 + "output":
29 + f"{START_REASONING}This is a reasoning section{END_REASONING}This is the rest", # noqa: E501
30 + "reasoning_content": "This is a reasoning section",
31 + "content": "This is the rest",
32 +}
33 +
34 +SIMPLE_REASONING_WITH_NEWLINE = {
35 + "output":
36 + f"{START_REASONING} Look!\n\nI'm thinking...{END_REASONING}\nThis is the rest", # noqa: E501
37 + "reasoning_content": " Look!\n\nI'm thinking...",
38 + "content": "\nThis is the rest",
39 +}
40 +
41 +SIMPLE_REASONING_WITH_MULTIPLE_NEWLINES = {
42 + "output":
43 + f"{START_REASONING}\nLook!\nI'm thinking...\n\n{END_REASONING}\n\n\nThis is the rest", # noqa: E501
44 + "reasoning_content": "\nLook!\nI'm thinking...\n\n",
45 + "content": "\n\n\nThis is the rest",
46 +}
47 +
48 +NO_REASONING_ONLY_END_THINK = {
49 + "output": f"{END_REASONING}\n\nNo thoughts, head empty!",
50 + "reasoning_content": None,
51 + "content": "\n\nNo thoughts, head empty!",
52 +}
53 +
54 +REASONING_ONLY_END_THINK = {
55 + "output":
56 + f"The user is asking me not to think.{END_REASONING}No thoughts!",
57 + "reasoning_content": "The user is asking me not to think.",
58 + "content": "No thoughts!",
59 +}
60 +
61 +TEST_CASES = [
62 + pytest.param(
63 + False, # not streaming
64 + NO_REASONING,
65 + id="no_reasoning",
66 + ),
67 + pytest.param(
68 + False, # not streaming
69 + NO_REASONING_WITH_NEWLINE,
70 + id="no_reasoning_with_newline",
71 + ),
72 + pytest.param(
73 + False, # not streaming
74 + SIMPLE_REASONING,
75 + id="simple_reasoning",
76 + ),
77 + pytest.param(
78 + False, # not streaming
79 + SIMPLE_REASONING_WITH_NEWLINE,
80 + id="simple_reasoning_with_newline",
81 + ),
82 + pytest.param(
83 + True, # enable streaming
84 + SIMPLE_REASONING_WITH_MULTIPLE_NEWLINES,
85 + id="simple_reasoning_with_multiple_newlines",
86 + ),
87 + pytest.param(
88 + False, # not streaming
89 + NO_REASONING_ONLY_END_THINK,
90 + id="no_reasoning_only_end_think",
91 + ),
92 + pytest.param(
93 + False, # not streaming
94 + REASONING_ONLY_END_THINK,
95 + id="yes_reasoning_only_end_think",
96 + ),
97 + pytest.param(
98 + True, # enable streaming
99 + NO_REASONING,
100 + id="no_reasoning_streaming",
101 + ),
102 + pytest.param(
103 + True, # enable streaming
104 + NO_REASONING_WITH_NEWLINE,
105 + id="no_reasoning_with_newline_streaming",
106 + ),
107 + pytest.param(
108 + True, # enable streaming
109 + SIMPLE_REASONING,
110 + id="simple_reasoning_streaming",
111 + ),
112 + pytest.param(
113 + True, # enable streaming
114 + SIMPLE_REASONING_WITH_NEWLINE,
115 + id="simple_reasoning_with_newline_streaming",
116 + ),
117 + pytest.param(
118 + True, # enable streaming
119 + SIMPLE_REASONING_WITH_MULTIPLE_NEWLINES,
120 + id="simple_reasoning_with_multiple_newlines_streaming",
121 + ),
122 + pytest.param(
123 + True, # enable streaming
124 + NO_REASONING_ONLY_END_THINK,
125 + id="no_reasoning_only_end_think_streaming",
126 + ),
127 + pytest.param(
128 + True, # enable streaming
129 + REASONING_ONLY_END_THINK,
130 + id="yes_reasoning_only_end_think_streaming",
131 + ),
132 +]
133 +
134 +# Global tokenizer initialization to avoid repeated loading
135 +tokenizer = AutoTokenizer.from_pretrained("allenai/dolma2-tokenizer")
136 +
137 +
138 +@pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
139 +def test_reasoning(
140 + streaming: bool,
141 + param_dict: dict[str, str],
142 +):
143 + output = tokenizer.tokenize(param_dict["output"])
144 +
145 + # decode everything to tokens
146 + model_output: list[str] = [
147 + tokenizer.convert_tokens_to_string([token]) for token in output
148 + ]
149 + parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
150 + parser: ReasoningParser = parser_cls(tokenizer)
151 +
152 + reasoning, content = run_reasoning_extraction(reasoning_parser=parser,
153 + model_output=model_output,
154 + streaming=streaming)
155 +
156 + assert reasoning == param_dict["reasoning_content"]
157 + assert content == param_dict["content"]
vllm/reasoning/__init__.py
+2
@@ -9,6 +9,7 @@ from .gptoss_reasoning_parser import GptOssReasoningParser
9 from .granite_reasoning_parser import GraniteReasoningParser
10 from .hunyuan_a13b_reasoning_parser import HunyuanA13BReasoningParser
11 from .mistral_reasoning_parser import MistralReasoningParser
12 +from .olmo3_reasoning_parser import Olmo3ReasoningParser
13 from .qwen3_reasoning_parser import Qwen3ReasoningParser
14 from .seedoss_reasoning_parser import SeedOSSReasoningParser
15 from .step3_reasoning_parser import Step3ReasoningParser
@@ -23,6 +24,7 @@ __all__ = [
24 "Qwen3ReasoningParser",
25 "Glm4MoeModelReasoningParser",
26 "MistralReasoningParser",
27 + "Olmo3ReasoningParser",
28 "Step3ReasoningParser",
29 "GptOssReasoningParser",
30 "SeedOSSReasoningParser",
vllm/reasoning/olmo3_reasoning_parser.py new
+294
@@ -0,0 +1,294 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import dataclasses as dt
5 +import enum
6 +from collections.abc import Sequence
7 +from typing import TYPE_CHECKING, Optional, Union
8 +
9 +import regex as re
10 +
11 +if TYPE_CHECKING:
12 + from vllm.transformers_utils.tokenizer import AnyTokenizer
13 +
14 +from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
15 + DeltaMessage, ResponsesRequest)
16 +from vllm.logger import init_logger
17 +from vllm.reasoning import ReasoningParser, ReasoningParserManager
18 +
19 +logger = init_logger(__name__)
20 +
21 +
22 +class Olmo3ReasoningState(enum.Enum):
23 + REASONING = 1
24 + CONTENT = 2
25 +
26 +
27 +@dt.dataclass(frozen=True)
28 +class Indices:
29 + start: int
30 + end: int
31 +
32 + def __len__(self):
33 + return self.end - self.start
34 +
35 +
36 +def string_overlap(a: str,
37 + b: str) -> tuple[Optional[Indices], Optional[Indices]]:
38 + """
39 + Find the longest overlap where the end of string a matches the start
40 + of string b.
41 +
42 + Args:
43 + a: First string
44 + b: Second string
45 +
46 + Returns:
47 + Tuple of IndicesTuples representing the overlapping portions in each
48 + string, or a tuple of None if no overlap exists
49 + """
50 +
51 + # swap so a is always the shorter string
52 + a, b, swap = (a, b, False) if len(a) < len(b) else (b, a, True)
53 +
54 + # first check: is a fully contained in b?
55 + if a in b:
56 + ind_a = Indices(0, len(a))
57 + ind_b = Indices(b.index(a), b.index(a) + len(a))
58 + return (ind_b, ind_a) if swap else (ind_a, ind_b)
59 +
60 + # second check: does the end of a overlap with the
61 + # beginning of b?
62 + for i in range(len(a) - 1, 0, -1):
63 + if a[-i:] == b[:i]:
64 + ind_a = Indices(len(a) - i, len(a))
65 + ind_b = Indices(0, i)
66 + return (ind_b, ind_a) if swap else (ind_a, ind_b)
67 +
68 + # third check: does the beginning of a overlap with
69 + # the end of b?
70 + for i in range(len(a) - 1, 0, -1):
71 + if b[-i:] == a[:i]:
72 + ind_a = Indices(0, i)
73 + ind_b = Indices(len(b) - i, len(b))
74 + return (ind_b, ind_a) if swap else (ind_a, ind_b)
75 +
76 + return None, None
77 +
78 +
79 +@dt.dataclass
80 +class Olmo3ReasoningBuffer:
81 + think_start: str = "<think>"
82 + think_end: str = "</think>"
83 + buffer: str = ""
84 +
85 + # we start in reasoning state to support cases where we hardcode
86 + # <think> as the start of the reasoning block.
87 + # In those cases, the only token we will see is </think>, which
88 + # is when we switch to content state.
89 + state: Olmo3ReasoningState = Olmo3ReasoningState.REASONING
90 +
91 + def process_buffer(self) -> Optional[DeltaMessage]:
92 + start_think_idx = self.buffer.find(self.think_start)
93 +
94 + if start_think_idx >= 0:
95 + self.state = Olmo3ReasoningState.REASONING
96 + pretext, self.buffer = (
97 + self.buffer[:start_think_idx],
98 + self.buffer[start_think_idx + len(self.think_start):],
99 + )
100 + if start_think_idx > 0:
101 + # this covers the case there's content before
102 + # the start of the reasoning block
103 + return DeltaMessage(content=pretext)
104 +
105 + end_think_idx = self.buffer.rfind(self.think_end)
106 +
107 + if end_think_idx >= 0:
108 + self.state = Olmo3ReasoningState.CONTENT
109 + pretext, self.buffer = (
110 + self.buffer[:end_think_idx],
111 + self.buffer[end_think_idx + len(self.think_end):],
112 + )
113 + if end_think_idx > 0:
114 + # this covers the case there's content before
115 + # the end of the reasoning block
116 + return DeltaMessage(reasoning_content=pretext)
117 +
118 + if self.state == Olmo3ReasoningState.REASONING:
119 + # we are inside reasoning block, return and empty
120 + # the text buffer
121 + (
122 + text_buffer,
123 + self.buffer,
124 + ) = self.buffer, ""
125 + return DeltaMessage(reasoning_content=text_buffer)
126 +
127 + if self.state == Olmo3ReasoningState.CONTENT:
128 + # we are outside reasoning block, return and empty
129 + # the text buffer
130 + (
131 + text_buffer,
132 + self.buffer,
133 + ) = self.buffer, ""
134 + return DeltaMessage(content=text_buffer)
135 +
136 + # nothing to return unless we are in reasoning or content state
137 + return None
138 +
139 + def __len__(self):
140 + # is the length of the text buffer
141 + return len(self.buffer)
142 +
143 + def add_text(self, delta_text: str) -> Optional[DeltaMessage]:
144 + # we start by adding the delta text to the buffer
145 + self.buffer += delta_text
146 +
147 + # setting this to empty before starting
148 + delta_message: Optional[DeltaMessage] = None
149 +
150 + # we start by computing the overlap between the delta_text
151 + # and start/end of think tokens.
152 + _, overlap_think_start = string_overlap(delta_text, self.think_start)
153 + _, overlap_think_end = string_overlap(delta_text, self.think_end)
154 +
155 + partial_overlap_start = overlap_think_start is not None and len(
156 + overlap_think_start) < len(self.think_start)
157 + partial_overlap_end = overlap_think_end is not None and len(
158 + overlap_think_end) < len(self.think_end)
159 +
160 + if (partial_overlap_start and self.think_start in self.buffer
161 + and not partial_overlap_end):
162 + # we can only process the buffer if partial overlap
163 + # is the last part of think token (thus causing
164 + # text_buffer to contain the start of think token)
165 + # and there are no partial overlaps with end think
166 + delta_message = self.process_buffer()
167 +
168 + elif partial_overlap_end and self.think_end in self.buffer:
169 + # same as before (partial overlap only allowed)
170 + # if the buffer contains the end think token,
171 + # but we don't have to check for partial overlap
172 + # with start think token because they are handled
173 + # by the previous condition
174 + delta_message = self.process_buffer()
175 +
176 + elif partial_overlap_start or partial_overlap_end:
177 + # in general, if there are overlaps, we don't
178 + # process the buffer because we want to wait until
179 + # the think token is fully completed.
180 + return None
181 + else:
182 + # we process the buffer as normal
183 + delta_message = self.process_buffer()
184 +
185 + return delta_message
186 +
187 +
188 +@ReasoningParserManager.register_module("olmo3")
189 +class Olmo3ReasoningParser(ReasoningParser):
190 + """
191 + Reasoning parser for Olmo 3 model
192 +
193 + Olmo3ReasoningParser
194 +
195 + This class implements a reasoning parser specifically designed for the
196 + Olmo 3 family of models. Olmo 3 models do not use special tokens to
197 + indicate reasoning; rather, reasoning trace is wrapped in `<think>` and
198 + `</think>`, which are tokenized using standard vocabulary entries.
199 + Because of this, the parser operates in string space, accumulating the
200 + characters in a buffer until it sees `<think>` or `</think>`. tokens
201 + to switch modes.
202 +
203 + Key Features:
204 + - For non-stream output, Recognizes and extracts reasoning (text
205 + bracketed by `<think>` and `</think>`) and content (everything
206 + after the first `</think>`).
207 + - For stream process, it uses a buffer to accumulate delta text,
208 + and output progressive delta messages as soon as thinking starts
209 + or ends.
210 + - For reliability, some Olmo 3 models may hardcode the first
211 + `<think>` token is the input text (similar to Deepseek R1,
212 + or reasoning-only Qwen models). To support such variants, the
213 + parser can optionally work in cases where the first `<think>`
214 + token is missing from generation.
215 + """
216 +
217 + def __init__(self, tokenizer: "AnyTokenizer", *args, **kwargs):
218 + super().__init__(tokenizer, *args, **kwargs)
219 +
220 + self.think_start = r"<think>"
221 + self.think_end = r"</think>"
222 +
223 + # notice that the first think is optional; this allows template to
224 + # work in cases when we hardcode a <think> at the beginning of the
225 + # reasoning template.
226 + reasoning_expr = (rf"^(?:{self.think_start})?(?P<reasoning>.*?)" +
227 + rf"{self.think_end}(?P<content>.*)$")
228 + self.reasoning_regex = re.compile(reasoning_expr, re.DOTALL)
229 +
230 + self.buffer = Olmo3ReasoningBuffer(think_start=self.think_start,
231 + think_end=self.think_end)
232 +
233 + def is_reasoning_end(self, input_ids: list[int]) -> bool:
234 + text = self.model_tokenizer.decode(input_ids)
235 + return self.think_end in text
236 +
237 + def extract_content_ids(self, input_ids: list[int]) -> list[int]:
238 + # for Olmo 3 streaming reason parsing, the stream parse
239 + # will call first, and the same token will be called in
240 + # is_reasoning_end and extract_content_ids
241 + # this id is not part of content, so just return [] here.
242 + return []
243 +
244 + def extract_reasoning_content(
245 + self,
246 + model_output: str,
247 + request: Union[ChatCompletionRequest, ResponsesRequest],
248 + ) -> tuple[Optional[str], Optional[str]]:
249 + """Extract the reasoning content & content sections, respectively.
250 + If the sequence doesn't match what we expect, i.e., the model generates
251 + something else, all content is considered non-reasoning content.
252 +
253 + Args:
254 + model_output (str): Output of the model to be parsed.
255 + request (ChatCompletionRequest | ResponsesRequest): Request being
256 + processed.
257 +
258 + Returns:
259 + tuple[Optional[str], Optional[str]]: Tuple pair containing the
260 + reasoning content and non-reasoning content.
261 + """
262 +
263 + re_match = self.reasoning_regex.match(model_output)
264 + if re_match:
265 + reasoning_content = re_match.group("reasoning") or None
266 + content = re_match.group("content") or None
267 + return reasoning_content, content
268 +
269 + # no reasoning content
270 + return None, model_output
271 +
272 + def extract_reasoning_content_streaming(
273 + self,
274 + previous_text: str,
275 + current_text: str,
276 + delta_text: str,
277 + previous_token_ids: Sequence[int],
278 + current_token_ids: Sequence[int],
279 + delta_token_ids: Sequence[int],
280 + ) -> Union[DeltaMessage, None]:
281 + """Extract content using token ID sequence state machine"""
282 +
283 + delta_message = self.buffer.add_text(delta_text)
284 + if (delta_message is None
285 + and self.buffer.think_end in self.buffer.buffer):
286 + # this is a bit hacky, but, because of how the buffer is
287 + # constructed, if the last delta_text contains characters that
288 + # marks the end of thinking tokens, then messages in the buffer
289 + # would never be processed because we get no other turn. To get
290 + # around that, we check if the text buffer contains the end of
291 + # thinking tokens, and, if so, we reprocess the buffer again.
292 + delta_message = self.buffer.process_buffer()
293 +
294 + return delta_message