diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py index 861e919c102a..01630490a152 100644 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ b/tests/v1/structured_output/test_reasoning_structured_output.py @@ -258,3 +258,128 @@ def test_should_advance_reasoning_already_ended( # Should return True since reasoning has ended assert result is True + + def test_should_advance_with_new_token_ids_detects_end( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """new_token_ids parameter is used as delta instead of index arithmetic. + + Regression test for MTP speculative decoding: when a spec token is + accepted, num_computed_tokens is pre-incremented past the main token, + making the index-based delta skip . Passing new_token_ids + directly bypasses that off-by-one. + """ + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + + # Configure the reasoner to return True only when the delta contains + # the synthetic token id (999). + THINK_END_ID = 999 + + class EndTokenReasoner: + def __init__(self, tokenizer): + pass + + def is_reasoning_end(self, input_ids): + return False + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + return THINK_END_ID in list(delta_ids) + + manager_with_reasoner.reasoner_cls = EndTokenReasoner + structured_req.reasoner = None # force lazy rebuild + + # Simulate MTP: num_computed_tokens has been pre-incremented past the + # main token so the index-based delta would be [spec_token], not + # [think_end, spec_token]. new_token_ids carries both tokens correctly. + all_tokens = list(range(10)) + [THINK_END_ID, 42] # think_end at -2 + mock_request_with_structured_output.all_token_ids = all_tokens + # Index-based delta would start here (past think_end): + mock_request_with_structured_output.num_computed_tokens = len(all_tokens) - 1 + mock_request_with_structured_output.num_output_placeholders = 0 + + new_token_ids = [THINK_END_ID, 42] # main=think_end, spec=42 + + result = manager_with_reasoner.should_advance( + mock_request_with_structured_output, new_token_ids + ) + + # reasoning_ended must be set and advance deferred until next step + assert structured_req.reasoning_ended is True + assert result is False + + def test_should_advance_new_token_ids_no_end_token( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """When new_token_ids does not contain , reasoning stays open.""" + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + + THINK_END_ID = 999 + + class EndTokenReasoner: + def __init__(self, tokenizer): + pass + + def is_reasoning_end(self, input_ids): + return False + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + return THINK_END_ID in list(delta_ids) + + manager_with_reasoner.reasoner_cls = EndTokenReasoner + structured_req.reasoner = None + + mock_request_with_structured_output.all_token_ids = list(range(10)) + mock_request_with_structured_output.num_computed_tokens = 10 + mock_request_with_structured_output.num_output_placeholders = 0 + + result = manager_with_reasoner.should_advance( + mock_request_with_structured_output, new_token_ids=[7, 8] + ) + + assert structured_req.reasoning_ended is False + assert result is False + + def test_should_advance_fallback_delta_without_new_token_ids( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """Without new_token_ids, the index-based delta is used (existing path).""" + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + + THINK_END_ID = 999 + + class EndTokenReasoner: + def __init__(self, tokenizer): + pass + + def is_reasoning_end(self, input_ids): + return False + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + return THINK_END_ID in list(delta_ids) + + manager_with_reasoner.reasoner_cls = EndTokenReasoner + structured_req.reasoner = None + + # Place think_end exactly where the index-based delta would start. + all_tokens = list(range(5)) + [THINK_END_ID] + mock_request_with_structured_output.all_token_ids = all_tokens + # delta_from = num_computed_tokens - num_output_placeholders = 5 - 0 = 5 + # all_token_ids[5:] = [THINK_END_ID] + mock_request_with_structured_output.num_computed_tokens = 5 + mock_request_with_structured_output.num_output_placeholders = 0 + + result = manager_with_reasoner.should_advance( + mock_request_with_structured_output # no new_token_ids + ) + + assert structured_req.reasoning_ended is True + assert result is False diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 25ccf79bc3a7..51fca931b434 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1595,7 +1595,9 @@ def update_from_output( request.status = RequestStatus.FINISHED_STOPPED stopped = True - if new_token_ids and self.structured_output_manager.should_advance(request): + if new_token_ids and self.structured_output_manager.should_advance( + request, new_token_ids + ): struct_output_request = request.structured_output_request assert struct_output_request is not None assert struct_output_request.grammar is not None diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 30921f3d74a4..10bfcc1a8f10 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -322,7 +322,11 @@ def should_fill_bitmask(self, request: "Request") -> bool: return request.structured_output_request.reasoning_ended return True - def should_advance(self, request: "Request") -> bool: + def should_advance( + self, + request: "Request", + new_token_ids: list[int] | None = None, + ) -> bool: if not request.use_structured_output: return False @@ -345,15 +349,25 @@ def should_advance(self, request: "Request") -> bool: if structured_req.reasoning_ended: return True - # Check if reasoning ends in *this* step - delta_from = request.num_computed_tokens - request.num_output_placeholders + # Check if reasoning ends in *this* step. + # When new_token_ids is provided (call site: update_from_output), use it + # directly as the delta. This avoids the off-by-one that occurs with MTP + # speculative decoding: _update_after_schedule pre-increments + # num_computed_tokens by (1 main + N spec tokens) BEFORE execution, so + # when a spec token is accepted the index-based delta skips the main token + # and is silently missed. all_token_ids = request.all_token_ids - start = ( - delta_from if delta_from >= 0 else max(len(all_token_ids) + delta_from, 0) - ) - if reasoner.is_reasoning_end_streaming( - all_token_ids, itertools.islice(all_token_ids, start, None) - ): + if new_token_ids is not None: + delta: Iterable[int] = new_token_ids + else: + delta_from = request.num_computed_tokens - request.num_output_placeholders + start = ( + delta_from + if delta_from >= 0 + else max(len(all_token_ids) + delta_from, 0) + ) + delta = itertools.islice(all_token_ids, start, None) + if reasoner.is_reasoning_end_streaming(all_token_ids, delta): structured_req.reasoning_ended = True # Reasoning just ended this step. Defer FSM advance until the next