From 5ae2d38821ba381ff2091b52b463296784b56217 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 18 Aug 2026 09:30:41 +1000 Subject: [PATCH] [Perf][Structured Output] Skip unused request-local reasoners (#52573) Signed-off-by: Bugen Zhao Co-authored-by: OpenAI Codex Co-authored-by: Nick Hill --- .../test_reasoning_structured_output.py | 65 +++++++++++++++++-- vllm/v1/structured_output/__init__.py | 40 ++++++------ 2 files changed, 79 insertions(+), 26 deletions(-) diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py index ad5f1d5d79..00a6cd87bc 100644 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ b/tests/v1/structured_output/test_reasoning_structured_output.py @@ -84,17 +84,62 @@ class TestReasoningStructuredOutput: return manager def test_should_fill_bitmask_with_enable_in_reasoning( - self, mock_vllm_config, mock_request_with_structured_output + self, manager_with_reasoner, mock_request_with_structured_output ): """Test should_fill_bitmask when enable_in_reasoning is True.""" - # Enable enable_in_reasoning - mock_vllm_config.structured_outputs_config.enable_in_reasoning = True - - manager = StructuredOutputManager(mock_vllm_config) + manager_with_reasoner.enable_in_reasoning = True # Should always return True when enable_in_reasoning is enabled - result = manager.should_fill_bitmask(mock_request_with_structured_output) + result = manager_with_reasoner.should_fill_bitmask( + mock_request_with_structured_output + ) assert result is True + assert ( + mock_request_with_structured_output.structured_output_request.reasoner + is None + ) + + def test_should_fill_bitmask_reasoning_already_ended( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """An active grammar does not need a request-local reasoner.""" + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = True + + result = manager_with_reasoner.should_fill_bitmask( + mock_request_with_structured_output + ) + + assert result is True + assert structured_req.reasoner is None + + def test_grammar_bitmask_skips_reasoner_when_already_active( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + """Bitmask generation skips reasoning-boundary detection.""" + manager_with_reasoner.vllm_config.num_speculative_tokens = 1 + manager_with_reasoner._grammar_bitmask = Mock() + manager_with_reasoner._grammar_bitmask.shape = (1, 1) + expected_bitmask = Mock() + manager_with_reasoner._grammar_bitmask.numpy.return_value = expected_bitmask + manager_with_reasoner._fill_bitmasks = Mock() + + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = True + request_id = mock_request_with_structured_output.request_id + + bitmask = manager_with_reasoner.grammar_bitmask( + requests={request_id: mock_request_with_structured_output}, + structured_output_request_ids=[request_id], + scheduled_spec_decode_tokens={}, + ) + + assert bitmask is expected_bitmask + assert structured_req.reasoner is None def test_should_fill_bitmask_without_enable_in_reasoning( self, @@ -171,6 +216,10 @@ class TestReasoningStructuredOutput: mock_request_with_structured_output ) assert result is True + assert ( + mock_request_with_structured_output.structured_output_request.reasoner + is None + ) def test_should_advance_reasoning_not_ended( self, @@ -233,6 +282,10 @@ class TestReasoningStructuredOutput: # Should return True since reasoning has ended assert result is True + assert ( + mock_request_with_structured_output.structured_output_request.reasoner + is None + ) def test_should_advance_uses_new_token_ids_when_provided( self, diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index f2ba4ac368..6a2409c8a5 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -289,12 +289,8 @@ class StructuredOutputManager: assert isinstance(grammar, StructuredOutputGrammar) apply_bitmask = self.should_fill_bitmask(request) - reasoner = self._get_reasoner(request) - detect_reasoning_end = ( - not apply_bitmask - and reasoner is not None - and not self.enable_in_reasoning - ) + reasoner = None if apply_bitmask else self._get_reasoner(request) + detect_reasoning_end = reasoner is not None simulated_buf: list[int] | None = None history_len = 0 @@ -369,20 +365,24 @@ class StructuredOutputManager: # NOTE (Hanchen) if enable_in_reasoning is True, it means that # the model needs to be constrained in reasoning. So we should always # enable the bitmask filling. + structured_req = request.structured_output_request + if self.enable_in_reasoning or ( + structured_req is not None and structured_req.reasoning_ended is True + ): + return True + reasoner = self._get_reasoner(request) if reasoner is not None: - if self.enable_in_reasoning: - return True - assert request.structured_output_request is not None - if request.structured_output_request.reasoning_ended is None: + assert structured_req is not None + if structured_req.reasoning_ended is None: # This should be removed here, but since `openai_gptoss` # is an independent code path, it is kept for now. # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, # it can be removed. - request.structured_output_request.reasoning_ended = ( - reasoner.is_reasoning_end(request.prompt_token_ids or []) + structured_req.reasoning_ended = reasoner.is_reasoning_end( + request.prompt_token_ids or [] ) - return request.structured_output_request.reasoning_ended + return structured_req.reasoning_ended return True def should_advance( @@ -398,19 +398,19 @@ class StructuredOutputManager: if TYPE_CHECKING: assert request.structured_output_request is not None assert request.structured_output_request.grammar is not None + structured_req = request.structured_output_request + if self.enable_in_reasoning or ( + structured_req is not None and structured_req.reasoning_ended is True + ): + return True + # by default, we should always advance # for cases that don't use thinking mode. reasoner = self._get_reasoner(request) if reasoner is None: return True - # if the model needs structured in reasoning, we should advance - if self.enable_in_reasoning: - return True - - structured_req = request.structured_output_request - if structured_req.reasoning_ended: - return True + assert structured_req is not None # Check if reasoning ends in *this* step. # When the caller passes new_token_ids (the tokens that were just