fix(structured_output): pass new_token_ids to should_advance() to fix MTP spec-decode off-by-one

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 </think>.  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 </think> 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 </think>, 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

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

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 </think> 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