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 4, 15:00 UTC
[Bugfix] Remove tokenizer encode/decode calls from Olmo3 reasoning parser (#40855)
Signed-off-by: Yifan <yzong@redhat.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>
yzong-rh committed
Apr 27, 2026 at 22:36 UTC
0d4f71420822a6f9eb386fe1f3f690ecbf31153b
2 files changed
+42
-16
tests/reasoning/test_olmo3_reasoning_parser.py
+12
-1
@@ -41,6 +41,12 @@ SIMPLE_REASONING_WITH_MULTIPLE_NEWLINES = {
41
"content": "\n\n\nThis is the rest",
42
}
43
44
+SIMPLE_REASONING_WITH_TRAILING_SPACE = {
45
+ "output": f"{START_REASONING}\nLook!\nI'm thinking... {END_REASONING}\nThis is the rest", # noqa: E501
46
+ "reasoning": "\nLook!\nI'm thinking... ",
47
+ "content": "\nThis is the rest",
48
+}
49
+
50
NO_REASONING_ONLY_END_THINK = {
51
"output": f"{END_REASONING}\n\nNo thoughts, head empty!",
52
"reasoning": None,
@@ -114,6 +120,11 @@ TEST_CASES = [
120
SIMPLE_REASONING_WITH_MULTIPLE_NEWLINES,
121
id="simple_reasoning_with_multiple_newlines_streaming",
122
),
123
+ pytest.param(
124
+ True, # enable streaming
125
+ SIMPLE_REASONING_WITH_TRAILING_SPACE,
126
+ id="simple_reasoning_with_trailing_space_streaming",
127
+ ),
128
pytest.param(
129
True, # enable streaming
130
NO_REASONING_ONLY_END_THINK,
@@ -127,7 +138,7 @@ TEST_CASES = [
138
]
139
140
# Global tokenizer initialization to avoid repeated loading
130
-tokenizer = AutoTokenizer.from_pretrained("allenai/dolma2-tokenizer")
141
+tokenizer = AutoTokenizer.from_pretrained("allenai/Olmo-3-7B-Think")
142
143
144
@pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
vllm/reasoning/olmo3_reasoning_parser.py
+30
-15
@@ -218,24 +218,32 @@ class Olmo3ReasoningParser(ReasoningParser):
218
token is missing from generation.
219
"""
220
221
+ think_start: str = r"<think>"
222
+ think_end: str = r"</think>"
223
+ # </think> is split in 3 by the pre-tokenizer, first split can be tokenized
224
+ # with an optional leading space, so there are 2 possible tokenizations
225
+ think_end_first_split: list[str] = [r"Ġ</", r"</"]
226
+ think_end_rest_split: list[str] = [r"think", r">"]
227
+ # notice that the first think is optional; this allows template to
228
+ # work in cases when we hardcode a <think> at the beginning of the
229
+ # reasoning template.
230
+ reasoning_regex: re.Pattern = re.compile(
231
+ rf"^(?:{think_start})?(?P<reasoning>.*?)"
232
+ rf"{think_end}(?P<content>.*)$",
233
+ re.DOTALL,
234
+ )
235
+
236
def __init__(self, tokenizer: "TokenizerLike", *args, **kwargs):
237
super().__init__(tokenizer, *args, **kwargs)
223
-
224
- self.think_start = r"<think>"
225
- self.think_end = r"</think>"
226
-
227
- # notice that the first think is optional; this allows template to
228
- # work in cases when we hardcode a <think> at the beginning of the
229
- # reasoning template.
230
- reasoning_expr = (
231
- rf"^(?:{self.think_start})?(?P<reasoning>.*?)"
232
- rf"{self.think_end}(?P<content>.*)$"
233
- )
234
- self.reasoning_regex = re.compile(reasoning_expr, re.DOTALL)
235
-
238
self.buffer = Olmo3ReasoningBuffer(
239
think_start=self.think_start, think_end=self.think_end
240
)
241
+ self.think_end_first_token_ids: list[int] = [
242
+ self.vocab[token] for token in self.think_end_first_split
243
+ ]
244
+ self.think_end_rest_token_ids: list[int] = [
245
+ self.vocab[token] for token in self.think_end_rest_split
246
+ ]
247
248
@property
249
def reasoning_start_str(self) -> str:
@@ -246,8 +254,15 @@ class Olmo3ReasoningParser(ReasoningParser):
254
return self.think_end
255
256
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
249
- text = self.model_tokenizer.decode(input_ids)
250
- return self.think_end in text
257
+ rest_ids = self.think_end_rest_token_ids
258
+ rest_len = len(rest_ids)
259
+ for i in range(len(input_ids) - rest_len, -1, -1):
260
+ if (
261
+ list(input_ids[i + 1 : i + 1 + rest_len]) == rest_ids
262
+ and input_ids[i] in self.think_end_first_token_ids
263
+ ):
264
+ return True
265
+ return False
266
267
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
268
# for Olmo 3 streaming reason parsing, the stream parse