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 5, 19:45 UTC
[Performance] Add is_reasoning_end_streaming() override to GptOssReasoningParser (#35745)
Signed-off-by: Fergus <fergus.barratt00@gmail.com> Signed-off-by: fergus barratt <fergus.barratt00@gmail.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
Fergus committed
Apr 21, 2026 at 19:31 UTC
5544f8c18b9e9ae20e41e019d23d58260940f225
2 files changed
+89
-1
tests/reasoning/test_gptoss_reasoning_parser.py
+69
@@ -280,3 +280,72 @@ class TestGptOssStructuralTags:
280
assert tag["content"]["type"] == "any_text"
281
assert tag["end"] == "<|end|>"
282
assert tag["begin"].startswith("<|channel|>")
283
+
284
+
285
+@pytest.mark.parametrize(
286
+ "output, is_reasoning_end",
287
+ [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES],
288
+)
289
+def test_gptoss_is_reasoning_end_streaming(
290
+ output,
291
+ is_reasoning_end,
292
+ gpt_oss_tokenizer,
293
+):
294
+ """Streaming override must agree with is_reasoning_end for all cases."""
295
+ tokens = gpt_oss_tokenizer.tokenize(output)
296
+ parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer)
297
+ output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(tokens)
298
+ delta_ids = output_ids[-1:] if output_ids else []
299
+ actual = parser.is_reasoning_end_streaming(output_ids, delta_ids)
300
+ assert is_reasoning_end == actual
301
+
302
+
303
+@pytest.mark.parametrize(
304
+ "output, is_reasoning_end",
305
+ [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES],
306
+)
307
+def test_gptoss_is_reasoning_end_streaming_long_prefix(
308
+ output,
309
+ is_reasoning_end,
310
+ gpt_oss_tokenizer,
311
+):
312
+ """Windowing must produce correct results even with a long prefix."""
313
+ tokens = gpt_oss_tokenizer.tokenize(output)
314
+ parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer)
315
+ output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(tokens)
316
+ # Prepend 10k dummy reasoning tokens to simulate a long generation
317
+ long_prefix = [1] * 10_000
318
+ padded_ids = long_prefix + list(output_ids)
319
+ delta_ids = output_ids[-1:] if output_ids else []
320
+ actual = parser.is_reasoning_end_streaming(padded_ids, delta_ids)
321
+ assert is_reasoning_end == actual
322
+
323
+
324
+@pytest.mark.parametrize(
325
+ "output, is_reasoning_end",
326
+ [(t["output"], t["is_reasoning_end"]) for t in TEST_CASES],
327
+)
328
+def test_gptoss_is_reasoning_end_streaming_large_delta(
329
+ output,
330
+ is_reasoning_end,
331
+ gpt_oss_tokenizer,
332
+):
333
+ """Simulate speculative decoding where the entire test sequence arrives
334
+ as a single large delta appended after a long prefix. The window must
335
+ expand to cover delta_ids so the end pattern is never missed."""
336
+ tokens = gpt_oss_tokenizer.tokenize(output)
337
+ parser: ReasoningParser = GptOssReasoningParser(gpt_oss_tokenizer)
338
+ output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(tokens)
339
+ long_prefix = [1] * 10_000
340
+ padded_ids = long_prefix + list(output_ids)
341
+ # delta_ids = the entire test sequence (as if accepted in one spec step)
342
+ delta_ids = list(output_ids)
343
+ actual = parser.is_reasoning_end_streaming(padded_ids, delta_ids)
344
+ assert is_reasoning_end == actual
345
+
346
+
347
+def test_gptoss_is_reasoning_end_streaming_signature(gpt_oss_tokenizer):
348
+ """Verify the method is callable with the expected signature."""
349
+ parser = GptOssReasoningParser(gpt_oss_tokenizer)
350
+ result = parser.is_reasoning_end_streaming([], [])
351
+ assert result is False
vllm/reasoning/gptoss_reasoning_parser.py
+20
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import json
4
-from collections.abc import Sequence
4
+from collections.abc import Iterable, Sequence
5
from typing import TYPE_CHECKING
6
7
from transformers import PreTrainedTokenizerBase
@@ -112,6 +112,25 @@ class GptOssReasoningParser(ReasoningParser):
112
return True
113
return False
114
115
+ def is_reasoning_end_streaming(
116
+ self, input_ids: Sequence[int], delta_ids: Iterable[int]
117
+ ) -> bool:
118
+ # The pattern window covers the end-of-reasoning marker itself.
119
+ # We add len(delta_ids) so that under speculative decoding (where
120
+ # a single step can accept many tokens) the entire accepted chunk
121
+ # is always inside the scan region.
122
+ delta_ids = tuple(delta_ids)
123
+ pattern_len = (
124
+ len(self.reasoning_end_token_ids_prefix)
125
+ + self.reasoning_max_num_between_tokens
126
+ + len(self.reasoning_end_token_ids_suffix)
127
+ )
128
+ window = pattern_len + len(delta_ids)
129
+ n = len(input_ids)
130
+ if n <= window:
131
+ return self.is_reasoning_end(input_ids)
132
+ return self.is_reasoning_end(input_ids[n - window :])
133
+
134
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
135
_, content, _ = parse_chat_output(input_ids)
136
if content is None: