[Bugfix][Outlines] Fix EOS termination and unconstrained masks after rejected drafts (#58612)
Nick Hill committed
Sep 24, 2026 at 14:55 UTC
36f94d5fe8a3953f1ad340de3c5aaec909765243
3 files changed
+172
-34
tests/v1/structured_output/test_structured_output_manager.py
+98
-1
@@ -19,6 +19,8 @@ THINK_END = "\n" # reasoning-end marker (single GPT-2 token)
19
EOS = "<|eos|>" # resolved to tokenizer.eos_token_id
20
JSON_SCHEMA = '{"type": "object"}'
21
BACKENDS = ("xgrammar", "guidance")
22
+REGEX_BACKENDS = (*BACKENDS, "outlines")
23
+DIGITS_REGEX = "[0-9]+"
24
MAX_WAIT_SECONDS = 5
25
NUM_SPEC_TOKENS = 8
26
@@ -88,6 +90,7 @@ def _build_harness(
90
reasoning_ended: bool | None = None,
91
enable_in_reasoning: bool = False,
92
reasoning_parser_kwargs: dict | None = None,
93
+ structured_outputs: StructuredOutputsParams | None = None,
94
) -> tuple[StructuredOutputManager, Request]:
95
vllm_config = VllmConfig(
96
model_config=ModelConfig(tokenizer=TOKENIZER),
@@ -106,7 +109,8 @@ def _build_harness(
109
manager.reasoner_cls = MockReasoner
110
111
sampling_params = SamplingParams(
109
- structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA)
112
+ structured_outputs=structured_outputs
113
+ or StructuredOutputsParams(json=JSON_SCHEMA)
114
)
115
sampling_params.structured_outputs._backend = backend
116
sampling_params.update_from_generation_config({}, tokenizer.eos_token_id)
@@ -412,3 +416,96 @@ def test_initial_constraint_activation(
416
structured_req = request.structured_output_request
417
assert structured_req is not None
418
assert structured_req.reasoner is None
419
+
420
+
421
+@pytest.mark.parametrize("backend", REGEX_BACKENDS)
422
+@pytest.mark.parametrize(
423
+ ("raw_drafts", "expected_validated", "expected_row_pattern", "terminated"),
424
+ [
425
+ pytest.param(("1", "2"), None, "CCC", False, id="all_valid"),
426
+ # EOS is only valid once the regex has matched; rows after EOS are
427
+ # unconstrained since the request stops there.
428
+ pytest.param(("1", EOS, "2"), ("1", EOS), "CCUU", True, id="terminates"),
429
+ pytest.param((EOS, "1"), (), "CUU", False, id="early_eos"),
430
+ ],
431
+)
432
+def test_regex_flow(
433
+ tokenizer,
434
+ backend: str,
435
+ raw_drafts: tuple[str, ...],
436
+ expected_validated: tuple[str, ...] | None,
437
+ expected_row_pattern: str,
438
+ terminated: bool,
439
+):
440
+ manager, request = _build_harness(
441
+ tokenizer,
442
+ backend,
443
+ use_reasoner=False,
444
+ structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX),
445
+ )
446
+ validated = _to_token_ids(
447
+ tokenizer, raw_drafts if expected_validated is None else expected_validated
448
+ )
449
+ _run_real_flow(
450
+ manager,
451
+ request,
452
+ raw_drafts=_to_token_ids(tokenizer, raw_drafts),
453
+ expected_validated=validated,
454
+ expected_row_pattern=expected_row_pattern,
455
+ # accept_tokens() latches reasoning_ended once it accepts any token.
456
+ expected_reasoning=True if validated else None,
457
+ expect_terminated=terminated,
458
+ )
459
+
460
+
461
+@pytest.mark.parametrize("backend", REGEX_BACKENDS)
462
+def test_rejected_draft_keeps_later_rows_constrained(tokenizer, backend: str):
463
+ """A draft the grammar rejects during bitmask fill must not unconstrain
464
+ the remaining rows or the bonus row."""
465
+ manager, request = _build_harness(
466
+ tokenizer,
467
+ backend,
468
+ use_reasoner=False,
469
+ structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX),
470
+ )
471
+ drafts = _to_token_ids(tokenizer, ("1", "z", "2"))
472
+ bitmask = manager.grammar_bitmask(
473
+ requests={request.request_id: request},
474
+ structured_output_request_ids=[request.request_id],
475
+ scheduled_spec_decode_tokens={request.request_id: drafts},
476
+ )
477
+ assert bitmask is not None
478
+ assert _row_pattern(bitmask) == "CCCC"
479
+ assert (bitmask[2] == bitmask[1]).all()
480
+ assert (bitmask[3] == bitmask[1]).all()
481
+
482
+
483
+def test_outlines_termination(tokenizer):
484
+ """outlines_core never advances on EOS; termination must still be
485
+ tracked exactly once per accepted EOS and undone by rollback."""
486
+ manager, request = _build_harness(
487
+ tokenizer,
488
+ "outlines",
489
+ use_reasoner=False,
490
+ structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX),
491
+ )
492
+ structured_req = request.structured_output_request
493
+ assert structured_req is not None
494
+ grammar = structured_req.grammar
495
+ assert grammar is not None
496
+ one, eos = _to_token_ids(tokenizer, ("1", EOS))
497
+
498
+ assert not grammar.accept_tokens(request.request_id, [eos])
499
+ assert grammar.accept_tokens(request.request_id, [one])
500
+ # Polling must not change the result.
501
+ assert not any(grammar.is_terminated() for _ in range(3))
502
+ assert grammar.accept_tokens(request.request_id, [eos])
503
+ assert all(grammar.is_terminated() for _ in range(3))
504
+ assert grammar.num_processed_tokens == 2
505
+
506
+ grammar.rollback(2)
507
+ assert not grammar.is_terminated()
508
+ assert grammar.num_processed_tokens == 0
509
+ assert grammar.validate_tokens([eos]) == []
510
+ assert grammar.accept_tokens(request.request_id, [one, eos, one])
511
+ assert grammar.is_terminated()
vllm/v1/structured_output/__init__.py
+17
-11
@@ -389,11 +389,16 @@ class StructuredOutputManager:
389
)
390
state_advancements = 0
391
seen_padding = False
392
- failed = False
392
+ # Row filled from the last valid grammar state before a draft
393
+ # was rejected; later rows reuse it rather than unconstraining.
394
+ failed_index: int | None = None
395
+ bitmask = self._grammar_bitmask
396
for i, token in enumerate(req_tokens):
394
- apply_bitmask = (
395
- not failed and not seen_padding and i >= constraint_start
396
- )
397
+ if failed_index is not None:
398
+ bitmask[cumulative_index].copy_(bitmask[failed_index])
399
+ cumulative_index += 1
400
+ continue
401
+ apply_bitmask = not seen_padding and i >= constraint_start
402
self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
403
if token == -1:
404
seen_padding = True
@@ -403,7 +408,7 @@ class StructuredOutputManager:
408
):
409
state_advancements += 1
410
else:
406
- failed = True
411
+ failed_index = cumulative_index
412
logger.error(
413
"Unexpected: grammar terminated or rejected draft "
414
"token %s for request %s during bitmask fill.",
@@ -415,12 +420,13 @@ class StructuredOutputManager:
420
# Diffusion LLMs don't sample a bonus token after the
421
# scheduled positions, so skip its bitmask in that case.
422
if not (self.vllm_config.model_config.is_diffusion and req_tokens):
418
- bonus_apply = (
419
- not failed
420
- and not seen_padding
421
- and constraint_start <= len(req_tokens)
422
- )
423
- self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),))
423
+ if failed_index is not None:
424
+ bitmask[cumulative_index].copy_(bitmask[failed_index])
425
+ else:
426
+ bonus_apply = not seen_padding and constraint_start <= len(
427
+ req_tokens
428
+ )
429
+ self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),))
430
cumulative_index += 1
431
432
if state_advancements > 0:
vllm/v1/structured_output/backend_outlines.py
+57
-22
@@ -97,6 +97,7 @@ class OutlinesBackend(StructuredOutputBackend):
97
)
98
return OutlinesGrammar(
99
vocab_size=self.vocab_size,
100
+ eos_token_id=self.vocabulary.inner.get_eos_token_id(),
101
guide=oc.Guide(index, max_rollback=max_rollback_tokens),
102
)
103
@@ -115,60 +116,94 @@ class OutlinesBackend(StructuredOutputBackend):
116
@dataclass
117
class OutlinesGrammar(StructuredOutputGrammar):
118
vocab_size: int
119
+ eos_token_id: int
120
guide: oc.Guide = field(hash=False)
121
num_processed_tokens: int = field(
122
default_factory=lambda: 0, repr=False, hash=False, init=False
123
)
124
123
- # outlines_core signals done on DFA accept; vLLM expects done after EOS.
124
- # We delay the finished flag by one step so EOS can still be emitted.
125
- _prev_finished: bool = field(default=False, init=False, repr=False, hash=False)
125
+ # outlines_core signals done on DFA accept and never advances on EOS
126
+ # (its mask allows only EOS in a final state, but advance(EOS) fails),
127
+ # so accepting EOS in a final state is tracked here instead.
128
+ _is_terminated: bool = field(default=False, init=False, repr=False, hash=False)
129
130
def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
131
"""Accepts a list of tokens and advances the FSM.
132
130
- Returns True if the FSM was advanced successfully.
131
- Returns False if the FSM failed to advance.
133
+ Returns True if all grammar-constrained tokens were accepted.
134
+ Tokens after termination (EOS) are ignored. Returns False if the FSM
135
+ failed to advance.
136
"""
133
- if self.guide.accepts_tokens(tokens):
134
- # Advance can fail when the next state reached after advancing with
135
- # the current tokens is a dead state. This is because Guide.accepts_tokens()
136
- # only checks whether the current tokens can be accepted,
137
- # whereas guide.advance() additionally checks the next state
138
- # after all tokens are accepted.
139
- # We need to be aware that the FSM must be prepared without dead states.
140
- for t in tokens:
141
- self.guide.advance(t)
142
- self.num_processed_tokens += 1
137
+ if self._is_terminated:
138
return True
144
- return False
139
+ eos_index = next(
140
+ (i for i, t in enumerate(tokens) if t == self.eos_token_id), None
141
+ )
142
+ if eos_index is not None:
143
+ tokens = tokens[:eos_index]
144
+ # Advance can fail when the next state reached after advancing with
145
+ # the current tokens is a dead state. This is because Guide.accepts_tokens()
146
+ # only checks whether the current tokens can be accepted,
147
+ # whereas guide.advance() additionally checks the next state
148
+ # after all tokens are accepted.
149
+ # We need to be aware that the FSM must be prepared without dead states.
150
+ if tokens and not self.guide.accepts_tokens(tokens):
151
+ return False
152
+ for t in tokens:
153
+ self.guide.advance(t)
154
+ if eos_index is not None:
155
+ if not self.guide.is_finished():
156
+ if tokens:
157
+ self.guide.rollback_state(len(tokens))
158
+ return False
159
+ self._is_terminated = True
160
+ self.num_processed_tokens += len(tokens) + self._is_terminated
161
+ return True
162
163
def rollback(self, num_tokens: int) -> None:
147
- self.guide.rollback_state(num_tokens)
164
+ if num_tokens <= 0:
165
+ return
166
self.num_processed_tokens -= num_tokens
167
+ if self._is_terminated:
168
+ self._is_terminated = False
169
+ num_tokens -= 1
170
+ if num_tokens:
171
+ self.guide.rollback_state(num_tokens)
172
173
def validate_tokens(self, tokens: list[int]) -> list[int]:
174
+ if self._is_terminated:
175
+ return []
176
accepted: list[int] = []
177
for tok in tokens:
178
+ if tok == self.eos_token_id:
179
+ if self._is_finished_after(accepted):
180
+ accepted.append(tok)
181
+ break
182
accepted.append(tok)
183
if not self.guide.accepts_tokens(accepted):
184
accepted.pop()
185
break
186
return accepted
187
188
+ def _is_finished_after(self, tokens: list[int]) -> bool:
189
+ if not tokens:
190
+ return self.guide.is_finished()
191
+ for t in tokens:
192
+ self.guide.advance(t)
193
+ finished = self.guide.is_finished()
194
+ self.guide.rollback_state(len(tokens))
195
+ return finished
196
+
197
def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> None:
198
mask = bitmask[idx]
199
self.guide.write_mask_into(mask.data_ptr(), mask.numel(), mask.element_size())
200
201
def is_terminated(self) -> bool:
164
- curr = self.guide.is_finished()
165
- prev = self._prev_finished
166
- self._prev_finished = curr
167
- return prev
202
+ return self._is_terminated
203
204
def reset(self):
205
self.num_processed_tokens = 0
171
- self._prev_finished = False
206
+ self._is_terminated = False
207
self.guide.reset()
208
209