|
|
@@ -0,0 +1,451 @@
|
|
|
+#!/usr/bin/env python3
|
|
|
+"""Patch vLLM structured-output reasoning-boundary bugs.
|
|
|
+
|
|
|
+Ports upstream vLLM PRs #44297 and #44993 to the b12x fork's vLLM source.
|
|
|
+
|
|
|
+Background
|
|
|
+----------
|
|
|
+When a reasoning parser (e.g. glm45 -> Glm47MoeParser) emits the reasoning-end
|
|
|
+marker (</think>) under speculative decoding (MTP, num_speculative_tokens=4),
|
|
|
+the grammar FSM (used by response_format / guided decoding) fails to advance
|
|
|
+across the reasoning->content boundary. This is vLLM bug #34650, fixed
|
|
|
+upstream by:
|
|
|
+
|
|
|
+ #44297 - Constrain bitmask and trim grammar advance at the reasoning boundary
|
|
|
+ #44993 - Advance grammar across reasoning boundary (supersedes #44993's
|
|
|
+ earlier conservative version; final merged form)
|
|
|
+
|
|
|
+Without this patch, models emit prose or markdown-fence prefixes before JSON
|
|
|
+in the content channel (e.g. ```{..." or "Here{..."), causing structured
|
|
|
+output validation failures. The tool-eval-bench Structured Output category
|
|
|
+scores 50% (6/12) instead of 100%.
|
|
|
+
|
|
|
+The fix
|
|
|
+-------
|
|
|
+1. should_advance() accepts an optional new_token_ids argument and uses it as
|
|
|
+ the authoritative step delta (instead of reconstructing the window from
|
|
|
+ num_computed_tokens - num_output_placeholders, which skips the marker
|
|
|
+ under async scheduling + spec-decode draft rejection).
|
|
|
+
|
|
|
+2. At the boundary, should_advance() records reasoning_end_token_index and
|
|
|
+ returns True for ALL constraint types (not just STRUCTURAL_TAG under spec
|
|
|
+ decode). The scheduler then trims reasoning content before accept_tokens.
|
|
|
+
|
|
|
+3. grammar_bitmask() detects a mid-window reasoning-end marker and constrains
|
|
|
+ subsequent draft positions via the bitmask, skipping grammar advance
|
|
|
+ through the marker itself.
|
|
|
+
|
|
|
+4. The scheduler's update_from_output() passes new_token_ids to should_advance
|
|
|
+ and calls trim_reasoning_for_advance() before accept_tokens.
|
|
|
+
|
|
|
+Fixes: #34650, #43388, #48228.
|
|
|
+"""
|
|
|
+import sys
|
|
|
+from pathlib import Path
|
|
|
+
|
|
|
+VLLM_DIR = Path("/opt/venv/lib/python3.12/site-packages/vllm")
|
|
|
+
|
|
|
+
|
|
|
+def patch_request():
|
|
|
+ """Add reasoning_end_token_index field to StructuredOutputRequest."""
|
|
|
+ path = VLLM_DIR / "v1" / "structured_output" / "request.py"
|
|
|
+ content = path.read_text()
|
|
|
+ if "reasoning_end_token_index" in content:
|
|
|
+ print(f" {path.name}: already patched, skipping")
|
|
|
+ return
|
|
|
+ old = " reasoning_ended: bool | None = None\n"
|
|
|
+ new = (
|
|
|
+ " reasoning_ended: bool | None = None\n"
|
|
|
+ " # Absolute index into the request's all_token_ids of the last\n"
|
|
|
+ " # reasoning token (the reasoning-end marker). Tokens at or before\n"
|
|
|
+ " # this index are reasoning content and must never be fed to the\n"
|
|
|
+ " # grammar. Only set when reasoning ends in a step whose tokens\n"
|
|
|
+ " # the scheduler advances immediately (structural tags +\n"
|
|
|
+ " # speculative decoding, see #42452).\n"
|
|
|
+ " reasoning_end_token_index: int | None = None\n"
|
|
|
+ )
|
|
|
+ if old not in content:
|
|
|
+ print(f" {path.name}: ERROR - could not find reasoning_ended field",
|
|
|
+ file=sys.stderr)
|
|
|
+ sys.exit(1)
|
|
|
+ content = content.replace(old, new, 1)
|
|
|
+ path.write_text(content)
|
|
|
+ print(f" {path.name}: patched")
|
|
|
+
|
|
|
+
|
|
|
+def patch_init():
|
|
|
+ """Patch __init__.py: imports, grammar_bitmask, should_advance, helpers."""
|
|
|
+ path = VLLM_DIR / "v1" / "structured_output" / "__init__.py"
|
|
|
+ content = path.read_text()
|
|
|
+ if "trim_reasoning_for_advance" in content:
|
|
|
+ print(f" {path.name}: already patched, skipping")
|
|
|
+ return
|
|
|
+
|
|
|
+ # 1. Add Sequence to collections.abc import
|
|
|
+ old_import = "from collections.abc import Iterable\n"
|
|
|
+ new_import = "from collections.abc import Iterable, Sequence\n"
|
|
|
+ if old_import not in content:
|
|
|
+ print(f" {path.name}: ERROR - could not find Iterable import",
|
|
|
+ file=sys.stderr)
|
|
|
+ sys.exit(1)
|
|
|
+ content = content.replace(old_import, new_import, 1)
|
|
|
+
|
|
|
+ # 2. Remove StructuredOutputOptions from backend_types import
|
|
|
+ old_bo = (
|
|
|
+ "from vllm.v1.structured_output.backend_types import (\n"
|
|
|
+ " StructuredOutputBackend,\n"
|
|
|
+ " StructuredOutputGrammar,\n"
|
|
|
+ " StructuredOutputOptions,\n"
|
|
|
+ ")"
|
|
|
+ )
|
|
|
+ new_bo = (
|
|
|
+ "from vllm.v1.structured_output.backend_types import (\n"
|
|
|
+ " StructuredOutputBackend,\n"
|
|
|
+ " StructuredOutputGrammar,\n"
|
|
|
+ ")"
|
|
|
+ )
|
|
|
+ if old_bo not in content:
|
|
|
+ print(f" {path.name}: ERROR - could not find backend_types import",
|
|
|
+ file=sys.stderr)
|
|
|
+ sys.exit(1)
|
|
|
+ content = content.replace(old_bo, new_bo, 1)
|
|
|
+
|
|
|
+ # 3. Remove logger import and initialization (unused after should_advance
|
|
|
+ # no longer references StructuredOutputOptions)
|
|
|
+ old_logger_import = "from vllm.logger import init_logger\n"
|
|
|
+ if old_logger_import in content:
|
|
|
+ content = content.replace(old_logger_import, "", 1)
|
|
|
+ old_logger_init = "logger = init_logger(__name__)\n\n\n"
|
|
|
+ if old_logger_init in content:
|
|
|
+ content = content.replace(old_logger_init, "", 1)
|
|
|
+
|
|
|
+ # 4. Replace grammar_bitmask serial-filling loop
|
|
|
+ old_bitmask = ''' state_advancements = 0
|
|
|
+ req_tokens = scheduled_spec_decode_tokens.get(req_id, ())
|
|
|
+ if self.vllm_config.model_config.is_diffusion and req_tokens:
|
|
|
+ # Diffusion LLMs don't sample a bonus token after the
|
|
|
+ # scheduled positions, so don't append the -1 placeholder.
|
|
|
+ token_iter: Iterable[int] = req_tokens
|
|
|
+ else:
|
|
|
+ token_iter = itertools.chain(req_tokens, (-1,))
|
|
|
+ for token in token_iter:
|
|
|
+ self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
|
|
|
+ if token == -1:
|
|
|
+ # Stop advancing the grammar once we hit a padding token.
|
|
|
+ apply_bitmask = False
|
|
|
+ if apply_bitmask and not grammar.is_terminated():
|
|
|
+ accepted = grammar.accept_tokens(req_id, [token])
|
|
|
+ assert accepted, (token, req_id, scheduled_spec_decode_tokens)
|
|
|
+ state_advancements += 1
|
|
|
+ cumulative_index += 1
|
|
|
+ if state_advancements > 0:
|
|
|
+ grammar.rollback(state_advancements)'''
|
|
|
+
|
|
|
+ new_bitmask = ''' reasoner = self._get_reasoner(request)
|
|
|
+ detect_reasoning_end = (
|
|
|
+ not apply_bitmask
|
|
|
+ and reasoner is not None
|
|
|
+ and not self.enable_in_reasoning
|
|
|
+ )
|
|
|
+ simulated_buf: list[int] | None = None
|
|
|
+ history_len = 0
|
|
|
+
|
|
|
+ state_advancements = 0
|
|
|
+ post_reasoning_end_in_window = False
|
|
|
+ req_tokens = scheduled_spec_decode_tokens.get(req_id, ())
|
|
|
+ for i, token in enumerate(req_tokens):
|
|
|
+ self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
|
|
|
+ advance_grammar = apply_bitmask
|
|
|
+ if token == -1:
|
|
|
+ apply_bitmask = False
|
|
|
+ advance_grammar = False
|
|
|
+ elif (
|
|
|
+ detect_reasoning_end
|
|
|
+ and reasoner is not None
|
|
|
+ and not apply_bitmask
|
|
|
+ ):
|
|
|
+ if simulated_buf is None:
|
|
|
+ history = list(request.all_token_ids)
|
|
|
+ history_len = len(history)
|
|
|
+ simulated_buf = history + list(req_tokens)
|
|
|
+ simulated = simulated_buf[: history_len + i + 1]
|
|
|
+ if reasoner.is_reasoning_end_streaming(simulated, [token]):
|
|
|
+ # Reasoning ended mid-window. Constrain the rest
|
|
|
+ # of the window via bitmask. Skip grammar advance
|
|
|
+ # through the marker (it is reasoning content);
|
|
|
+ # try to advance through subsequent drafts so the
|
|
|
+ # next bitmask row reflects the post-advance state,
|
|
|
+ # but tolerate rejection since those drafts predate
|
|
|
+ # the bitmask and are not guaranteed valid.
|
|
|
+ apply_bitmask = True
|
|
|
+ advance_grammar = False
|
|
|
+ post_reasoning_end_in_window = True
|
|
|
+ if advance_grammar and not grammar.is_terminated():
|
|
|
+ accepted = grammar.accept_tokens(req_id, [token])
|
|
|
+ if accepted:
|
|
|
+ state_advancements += 1
|
|
|
+ elif not post_reasoning_end_in_window:
|
|
|
+ raise AssertionError(
|
|
|
+ (token, req_id, scheduled_spec_decode_tokens)
|
|
|
+ )
|
|
|
+ cumulative_index += 1
|
|
|
+ # Diffusion LLMs don't sample a bonus token after the
|
|
|
+ # scheduled positions, so skip its bitmask in that case.
|
|
|
+ if not (self.vllm_config.model_config.is_diffusion and req_tokens):
|
|
|
+ # bonus_apply must be True when the bonus-row position
|
|
|
+ # should be grammar-constrained. Two triggers:
|
|
|
+ # - should_fill_bitmask(request): reasoning was already
|
|
|
+ # over at step start (or no reasoner /
|
|
|
+ # enable_in_reasoning).
|
|
|
+ # - apply_bitmask: reasoning ended mid-window in this
|
|
|
+ # call and was flipped True after the marker;
|
|
|
+ # should_fill_bitmask still returns False here because
|
|
|
+ # reasoning_ended is only persisted later by
|
|
|
+ # should_advance.
|
|
|
+ bonus_apply = self.should_fill_bitmask(request) or apply_bitmask
|
|
|
+ self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),))
|
|
|
+ cumulative_index += 1
|
|
|
+ if state_advancements > 0:
|
|
|
+ grammar.rollback(state_advancements)'''
|
|
|
+
|
|
|
+ if old_bitmask not in content:
|
|
|
+ print(f" {path.name}: ERROR - could not find grammar_bitmask serial loop",
|
|
|
+ file=sys.stderr)
|
|
|
+ sys.exit(1)
|
|
|
+ content = content.replace(old_bitmask, new_bitmask, 1)
|
|
|
+
|
|
|
+ # 5. Replace should_advance and add helper methods
|
|
|
+ old_should_advance = ''' def should_advance(self, request: "Request") -> bool:
|
|
|
+ if not request.use_structured_output:
|
|
|
+ return False
|
|
|
+
|
|
|
+ # To determine whether we can advance the FSM.
|
|
|
+ # Supports thinking usage where we skip the reasoning components.
|
|
|
+ if TYPE_CHECKING:
|
|
|
+ assert request.structured_output_request is not None
|
|
|
+ assert request.structured_output_request.grammar is not None
|
|
|
+ # 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
|
|
|
+
|
|
|
+ # Check if reasoning ends in *this* step
|
|
|
+ delta_from = request.num_computed_tokens - request.num_output_placeholders
|
|
|
+ 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)
|
|
|
+ ):
|
|
|
+ structured_req.reasoning_ended = True
|
|
|
+
|
|
|
+ # Reasoning just ended this step. Defer FSM advance until the next
|
|
|
+ # pass (see reasoning_ended check above) for JSON/regex/choice/grammar:
|
|
|
+ # advancing on the closing boundary token can accept tokens that still
|
|
|
+ # belong to the reasoning stream. Structural tags are the only safe
|
|
|
+ # same-step exception: they model phased output (e.g. thinking tag ->
|
|
|
+ # answer tag), and speculative decoding must run grammar.validate_tokens
|
|
|
+ # on draft tokens produced immediately after that transition.
|
|
|
+ if (
|
|
|
+ self.vllm_config.speculative_config is not None
|
|
|
+ and structured_req.structured_output_key[0]
|
|
|
+ == StructuredOutputOptions.STRUCTURAL_TAG
|
|
|
+ ):
|
|
|
+ return True
|
|
|
+
|
|
|
+ return False'''
|
|
|
+
|
|
|
+ new_should_advance = ''' def should_advance(
|
|
|
+ self,
|
|
|
+ request: "Request",
|
|
|
+ new_token_ids: list[int] | None = None,
|
|
|
+ ) -> bool:
|
|
|
+ if not request.use_structured_output:
|
|
|
+ return False
|
|
|
+
|
|
|
+ # To determine whether we can advance the FSM.
|
|
|
+ # Supports thinking usage where we skip the reasoning components.
|
|
|
+ if TYPE_CHECKING:
|
|
|
+ assert request.structured_output_request is not None
|
|
|
+ assert request.structured_output_request.grammar is not None
|
|
|
+ # 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
|
|
|
+
|
|
|
+ # Check if reasoning ends in *this* step.
|
|
|
+ # When the caller passes new_token_ids (the tokens that were just
|
|
|
+ # appended this step), use it directly as the delta window. The
|
|
|
+ # placeholder-derived fallback assumes num_output_placeholders ==
|
|
|
+ # len(new_token_ids), which breaks under async scheduling + spec
|
|
|
+ # decode when some drafts are rejected (#43388): the placeholder
|
|
|
+ # count remains > 0 after the step and the computed delta window
|
|
|
+ # starts past the reasoning-end marker.
|
|
|
+ all_token_ids = request.all_token_ids
|
|
|
+ if new_token_ids:
|
|
|
+ # The tokens were already appended this step, so the step window
|
|
|
+ # starts exactly len(new_token_ids) from the end.
|
|
|
+ start = len(all_token_ids) - len(new_token_ids)
|
|
|
+ delta_ids: 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_ids = itertools.islice(all_token_ids, start, None)
|
|
|
+ if reasoner.is_reasoning_end_streaming(all_token_ids, delta_ids):
|
|
|
+ structured_req.reasoning_ended = True
|
|
|
+
|
|
|
+ # Record the boundary so the scheduler can exclude reasoning tokens.
|
|
|
+ end_index = self._find_reasoning_end_index(reasoner, all_token_ids, start)
|
|
|
+
|
|
|
+ structured_req.reasoning_end_token_index = end_index
|
|
|
+ return True
|
|
|
+
|
|
|
+ return False
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _find_reasoning_end_index(
|
|
|
+ reasoner: "ReasoningParser", all_token_ids: Sequence[int], start: int
|
|
|
+ ) -> int:
|
|
|
+ """Locates the last reasoning token within all_token_ids[start:].
|
|
|
+
|
|
|
+ Returns the absolute index of the token at which
|
|
|
+ is_reasoning_end_streaming first fires. Falls back to the final
|
|
|
+ index when no single token triggers the detection (e.g. a
|
|
|
+ multi-token marker only recognized on the full delta), which
|
|
|
+ conservatively treats the whole step as reasoning content.
|
|
|
+ """
|
|
|
+ prefix = list(itertools.islice(all_token_ids, start))
|
|
|
+ for idx in range(start, len(all_token_ids)):
|
|
|
+ token = all_token_ids[idx]
|
|
|
+ prefix.append(token)
|
|
|
+ if reasoner.is_reasoning_end_streaming(prefix, [token]):
|
|
|
+ return idx
|
|
|
+ return len(all_token_ids) - 1
|
|
|
+
|
|
|
+ def trim_reasoning_for_advance(
|
|
|
+ self, request: "Request", new_token_ids: list[int]
|
|
|
+ ) -> list[int]:
|
|
|
+ """Drops reasoning content from tokens about to advance the grammar.
|
|
|
+
|
|
|
+ When reasoning ends mid-step (see should_advance), the step's output
|
|
|
+ still contains reasoning tokens up to and including the end marker.
|
|
|
+ Those are not grammar content: feeding them to accept_tokens makes
|
|
|
+ the grammar reject the marker and kills the request (#44006).
|
|
|
+
|
|
|
+ Returns the suffix of new_token_ids that follows the reasoning-end
|
|
|
+ marker. Steps fully after the boundary are returned unchanged.
|
|
|
+ """
|
|
|
+ structured_req = request.structured_output_request
|
|
|
+ if structured_req is None:
|
|
|
+ return new_token_ids
|
|
|
+ end_idx = structured_req.reasoning_end_token_index
|
|
|
+ if end_idx is None:
|
|
|
+ return new_token_ids
|
|
|
+ first_idx = len(request.all_token_ids) - len(new_token_ids)
|
|
|
+ num_reasoning = end_idx + 1 - first_idx
|
|
|
+ if num_reasoning <= 0:
|
|
|
+ return new_token_ids
|
|
|
+ return new_token_ids[num_reasoning:]'''
|
|
|
+
|
|
|
+ if old_should_advance not in content:
|
|
|
+ print(f" {path.name}: ERROR - could not find should_advance method",
|
|
|
+ file=sys.stderr)
|
|
|
+ sys.exit(1)
|
|
|
+ content = content.replace(old_should_advance, new_should_advance, 1)
|
|
|
+
|
|
|
+ path.write_text(content)
|
|
|
+ print(f" {path.name}: patched")
|
|
|
+
|
|
|
+
|
|
|
+def patch_scheduler():
|
|
|
+ """Patch update_from_output to pass new_token_ids and trim before accept."""
|
|
|
+ path = VLLM_DIR / "v1" / "core" / "sched" / "scheduler.py"
|
|
|
+ content = path.read_text()
|
|
|
+ if "trim_reasoning_for_advance" in content:
|
|
|
+ print(f" {path.name}: already patched, skipping")
|
|
|
+ return
|
|
|
+
|
|
|
+ old = ''' if new_token_ids and self.structured_output_manager.should_advance(request):
|
|
|
+ struct_output_request = request.structured_output_request
|
|
|
+ assert struct_output_request is not None
|
|
|
+ assert struct_output_request.grammar is not None
|
|
|
+ if not struct_output_request.grammar.accept_tokens( # type: ignore[union-attr]
|
|
|
+ req_id, new_token_ids
|
|
|
+ ):
|
|
|
+ logger.error(
|
|
|
+ "Unexpected: grammar rejected tokens %s for request %s. "
|
|
|
+ "Terminating request.",
|
|
|
+ new_token_ids,
|
|
|
+ req_id,
|
|
|
+ )
|
|
|
+ request.status = RequestStatus.FINISHED_ERROR
|
|
|
+ request.resumable = False
|
|
|
+ stopped = True'''
|
|
|
+
|
|
|
+ new = ''' if new_token_ids and self.structured_output_manager.should_advance(
|
|
|
+ request, new_token_ids=new_token_ids
|
|
|
+ ):
|
|
|
+ struct_output_request = request.structured_output_request
|
|
|
+ assert struct_output_request is not None
|
|
|
+ grammar = struct_output_request.grammar
|
|
|
+ assert grammar is not None
|
|
|
+ # new_token_ids can be a mixed block of reasoning content, then
|
|
|
+ # the reasoning end marker, then the start of the grammar content.
|
|
|
+ # Trim the reasoning content so the grammar only sees grammar content.
|
|
|
+ advance_token_ids = (
|
|
|
+ self.structured_output_manager.trim_reasoning_for_advance(
|
|
|
+ request, new_token_ids
|
|
|
+ )
|
|
|
+ )
|
|
|
+ if advance_token_ids and not grammar.accept_tokens(
|
|
|
+ req_id, advance_token_ids
|
|
|
+ ):
|
|
|
+ logger.error(
|
|
|
+ "Unexpected: grammar rejected tokens %s for request %s. "
|
|
|
+ "Terminating request.",
|
|
|
+ advance_token_ids,
|
|
|
+ req_id,
|
|
|
+ )
|
|
|
+ request.status = RequestStatus.FINISHED_ERROR
|
|
|
+ request.resumable = False
|
|
|
+ stopped = True'''
|
|
|
+
|
|
|
+ if old not in content:
|
|
|
+ print(f" {path.name}: ERROR - could not find should_advance block",
|
|
|
+ file=sys.stderr)
|
|
|
+ sys.exit(1)
|
|
|
+ content = content.replace(old, new, 1)
|
|
|
+ path.write_text(content)
|
|
|
+ print(f" {path.name}: patched")
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ print("Applying structured output reasoning-boundary patch...")
|
|
|
+ print(f" vLLM dir: {VLLM_DIR}")
|
|
|
+ patch_request()
|
|
|
+ patch_init()
|
|
|
+ patch_scheduler()
|
|
|
+ print("Done.")
|
