[Bugfix][Structured Output][Spec Decode] Advance grammar across reasoning boundary (#44993)
Signed-off-by: Allen.Yu <yuyue0225sc@163.com>
yue.yu committed
Jul 24, 2026 at 00:14 UTC
0416dab275d51327b331a1c6baaec754a68d7764
3 files changed
+146
-54
tests/v1/structured_output/test_reasoning_structured_output.py
+111
-19
@@ -73,6 +73,7 @@ class TestReasoningStructuredOutput:
73
request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8]
74
request.num_computed_tokens = 5
75
request.num_output_placeholders = 0
76
+ request.request_id = "mock_req"
77
return request
78
79
@pytest.fixture
@@ -208,53 +209,144 @@ class TestReasoningStructuredOutput:
209
mock_request_with_structured_output
210
)
211
211
- # Should set reasoning_ended to True but return False for this step
212
+ # The scheduler trims the reasoning prefix before advancing the grammar.
213
assert (
214
mock_request_with_structured_output.structured_output_request.reasoning_ended
215
is True
216
)
216
- assert result is False
217
+ assert result is True
218
218
- def test_should_advance_reasoning_just_ended_with_spec_decode_structural_tag(
219
+ def test_should_advance_reasoning_already_ended(
220
self,
221
manager_with_reasoner,
222
mock_request_with_structured_output,
223
):
223
- """When reasoning ends this step, advance immediately for structural
224
- tags with speculative decoding."""
224
+ """Test should_advance when reasoning has already ended."""
225
+ # Set reasoning as already ended
226
+ (
227
+ mock_request_with_structured_output.structured_output_request
228
+ ).reasoning_ended = True
229
+
230
+ result = manager_with_reasoner.should_advance(
231
+ mock_request_with_structured_output
232
+ )
233
+
234
+ # Should return True since reasoning has ended
235
+ assert result is True
236
+
237
+ def test_should_advance_uses_new_token_ids_when_provided(
238
+ self,
239
+ manager_with_reasoner,
240
+ mock_request_with_structured_output,
241
+ ):
242
+ """Regression for #43388: when caller passes new_token_ids, the
243
+ reasoner sees the exact multi-token delta rather than the
244
+ placeholder-derived window.
245
+ """
246
structured_req = mock_request_with_structured_output.structured_output_request
247
structured_req.reasoning_ended = False
227
- structured_req.structured_output_key = (
228
- StructuredOutputOptions.STRUCTURAL_TAG,
229
- "{}",
230
- )
248
+
249
+ end_token_id = 248069
250
+
251
reasoner = MockReasoner(tokenizer=Mock())
232
- reasoner.is_reasoning_end_streaming.return_value = True
252
+ # Detection mirrors the real Qwen3 parser: end token in the delta.
253
+ reasoner.is_reasoning_end_streaming = Mock(
254
+ side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids)
255
+ )
256
structured_req.reasoner = reasoner
257
235
- manager_with_reasoner.vllm_config.speculative_config = Mock()
258
+ # Scenario from #43388: async + spec decode K=4, 4 tokens accepted
259
+ # but only 1 placeholder remains (some drafts were rejected).
260
+ # The placeholder math would yield delta=[271] and miss </think>.
261
+ # Passing new_token_ids must override that.
262
+ new_token_ids = [9, 198, end_token_id, 271]
263
+ mock_request_with_structured_output.all_token_ids = [
264
+ 1,
265
+ 2,
266
+ 3,
267
+ 4,
268
+ 5,
269
+ ] + new_token_ids
270
+ mock_request_with_structured_output.num_computed_tokens = 9
271
+ mock_request_with_structured_output.num_output_placeholders = 1
272
273
result = manager_with_reasoner.should_advance(
238
- mock_request_with_structured_output
274
+ mock_request_with_structured_output,
275
+ new_token_ids=new_token_ids,
276
)
277
278
+ # First call to is_reasoning_end_streaming was with the full
279
+ # new_token_ids (not the truncated placeholder window).
280
+ first_call = reasoner.is_reasoning_end_streaming.call_args_list[0]
281
+ _, called_delta = first_call.args
282
+ assert list(called_delta) == new_token_ids
283
+
284
assert structured_req.reasoning_ended is True
285
assert result is True
286
244
- def test_should_advance_reasoning_already_ended(
287
+ def test_should_advance_without_new_token_ids_falls_back(
288
self,
289
manager_with_reasoner,
290
mock_request_with_structured_output,
291
):
249
- """Test should_advance when reasoning has already ended."""
250
- # Set reasoning as already ended
251
- (
252
- mock_request_with_structured_output.structured_output_request
253
- ).reasoning_ended = True
292
+ """Backward compat: callers that don't pass new_token_ids keep
293
+ the original placeholder-derived delta window.
294
+ """
295
+ structured_req = mock_request_with_structured_output.structured_output_request
296
+ structured_req.reasoning_ended = False
297
+ reasoner = MockReasoner(tokenizer=Mock())
298
+ reasoner.is_reasoning_end_streaming.return_value = False
299
+ structured_req.reasoner = reasoner
300
+
301
+ mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5]
302
+ mock_request_with_structured_output.num_computed_tokens = 5
303
+ mock_request_with_structured_output.num_output_placeholders = 2
304
305
result = manager_with_reasoner.should_advance(
306
mock_request_with_structured_output
307
)
308
259
- # Should return True since reasoning has ended
309
+ # placeholder window: start = 5 - 2 = 3, delta = [4, 5]
310
+ _, called_delta = reasoner.is_reasoning_end_streaming.call_args[0]
311
+ assert list(called_delta) == [4, 5]
312
+ assert result is False
313
+
314
+ def test_should_advance_trims_reasoning_prefix_for_json(
315
+ self,
316
+ manager_with_reasoner,
317
+ mock_request_with_structured_output,
318
+ ):
319
+ """JSON uses the common trim-then-advance path at the boundary."""
320
+ structured_req = mock_request_with_structured_output.structured_output_request
321
+ structured_req.reasoning_ended = False
322
+ structured_req.structured_output_key = (
323
+ StructuredOutputOptions.JSON_OBJECT,
324
+ "{}",
325
+ )
326
+
327
+ marker = 248069
328
+
329
+ class MarkerReasoner:
330
+ def __init__(self, *_, **__):
331
+ pass
332
+
333
+ def is_reasoning_end_streaming(self, input_ids, delta_ids):
334
+ return marker in list(delta_ids)
335
+
336
+ structured_req.reasoner = MarkerReasoner()
337
+
338
+ new_token_ids = [9, 198, marker, 271, 5005]
339
+ mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
340
+
341
+ result = manager_with_reasoner.should_advance(
342
+ mock_request_with_structured_output,
343
+ new_token_ids=new_token_ids,
344
+ )
345
+
346
+ structured_req.grammar.accept_tokens.assert_not_called()
347
+ assert structured_req.reasoning_ended is True
348
assert result is True
349
+ assert structured_req.reasoning_end_token_index == 5
350
+ assert manager_with_reasoner.trim_reasoning_for_advance(
351
+ mock_request_with_structured_output, new_token_ids
352
+ ) == [271, 5005]
vllm/v1/core/sched/scheduler.py
+3
-1
@@ -1716,7 +1716,9 @@ class Scheduler(SchedulerInterface):
1716
request.status = RequestStatus.FINISHED_STOPPED
1717
stopped = True
1718
1719
- if new_token_ids and self.structured_output_manager.should_advance(request):
1719
+ if new_token_ids and self.structured_output_manager.should_advance(
1720
+ request, new_token_ids=new_token_ids
1721
+ ):
1722
struct_output_request = request.structured_output_request
1723
assert struct_output_request is not None
1724
grammar = struct_output_request.grammar
vllm/v1/structured_output/__init__.py
+32
-34
@@ -7,7 +7,6 @@ from concurrent.futures import Future, ThreadPoolExecutor
7
from typing import TYPE_CHECKING
8
9
from vllm.config import VllmConfig
10
-from vllm.logger import init_logger
10
from vllm.reasoning import ReasoningParserManager
11
from vllm.tokenizers import cached_tokenizer_from_config
12
from vllm.utils.import_utils import LazyLoader
@@ -15,7 +14,6 @@ from vllm.v1.structured_output.backend_guidance import GuidanceBackend
14
from vllm.v1.structured_output.backend_types import (
15
StructuredOutputBackend,
16
StructuredOutputGrammar,
18
- StructuredOutputOptions,
17
)
18
from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
19
@@ -30,9 +28,6 @@ else:
28
torch = LazyLoader("torch", globals(), "torch")
29
30
33
-logger = init_logger(__name__)
34
-
35
-
31
class StructuredOutputManager:
32
"""Engine-level manager for structured output requests."""
33
@@ -379,7 +374,11 @@ class StructuredOutputManager:
374
return request.structured_output_request.reasoning_ended
375
return True
376
382
- def should_advance(self, request: "Request") -> bool:
377
+ def should_advance(
378
+ self,
379
+ request: "Request",
380
+ new_token_ids: list[int] | None = None,
381
+ ) -> bool:
382
if not request.use_structured_output:
383
return False
384
@@ -402,37 +401,36 @@ class StructuredOutputManager:
401
if structured_req.reasoning_ended:
402
return True
403
405
- # Check if reasoning ends in *this* step
406
- delta_from = request.num_computed_tokens - request.num_output_placeholders
404
+ # Check if reasoning ends in *this* step.
405
+ # When the caller passes new_token_ids (the tokens that were just
406
+ # appended this step), use it directly as the delta window. The
407
+ # placeholder-derived fallback assumes num_output_placeholders ==
408
+ # len(new_token_ids), which breaks under async scheduling + spec
409
+ # decode when some drafts are rejected (#43388): the placeholder
410
+ # count remains > 0 after the step and the computed delta window
411
+ # starts past the reasoning-end marker.
412
all_token_ids = request.all_token_ids
408
- start = (
409
- delta_from if delta_from >= 0 else max(len(all_token_ids) + delta_from, 0)
410
- )
411
- if reasoner.is_reasoning_end_streaming(
412
- all_token_ids, itertools.islice(all_token_ids, start, None)
413
- ):
413
+ if new_token_ids:
414
+ # The tokens were already appended this step, so the step window
415
+ # starts exactly len(new_token_ids) from the end.
416
+ start = len(all_token_ids) - len(new_token_ids)
417
+ delta_ids: Iterable[int] = new_token_ids
418
+ else:
419
+ delta_from = request.num_computed_tokens - request.num_output_placeholders
420
+ start = (
421
+ delta_from
422
+ if delta_from >= 0
423
+ else max(len(all_token_ids) + delta_from, 0)
424
+ )
425
+ delta_ids = itertools.islice(all_token_ids, start, None)
426
+ if reasoner.is_reasoning_end_streaming(all_token_ids, delta_ids):
427
structured_req.reasoning_ended = True
428
416
- # Reasoning just ended this step. Defer FSM advance until the next
417
- # pass (see reasoning_ended check above) for JSON/regex/choice/grammar:
418
- # advancing on the closing boundary token can accept tokens that still
419
- # belong to the reasoning stream. Structural tags are the only safe
420
- # same-step exception: they model phased output (e.g. thinking tag ->
421
- # answer tag), and speculative decoding must run grammar.validate_tokens
422
- # on draft tokens produced immediately after that transition.
423
- if (
424
- self.vllm_config.speculative_config is not None
425
- and structured_req.structured_output_key[0]
426
- == StructuredOutputOptions.STRUCTURAL_TAG
427
- ):
428
- # The scheduler will advance the grammar with this step's
429
- # tokens right away, but the step still contains reasoning
430
- # content up to and including the end marker. Record where
431
- # it ends so trim_reasoning_for_advance() can drop it.
432
- structured_req.reasoning_end_token_index = (
433
- self._find_reasoning_end_index(reasoner, all_token_ids, start)
434
- )
435
- return True
429
+ # Record the boundary so the scheduler can exclude reasoning tokens.
430
+ end_index = self._find_reasoning_end_index(reasoner, all_token_ids, start)
431
+
432
+ structured_req.reasoning_end_token_index = end_index
433
+ return True
434
435
return False
436