[Structured Output] Let terminal grammars stop under min_tokens (II) (#54218)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
Artem Perevedentsev committed
Aug 31, 2026 at 19:39 UTC
39e276eaeb9daed06a180f6a8d187bbb8790e97b
7 files changed
+501
-60
tests/v1/core/test_scheduler.py
+87
-32
@@ -840,11 +840,10 @@ def test_stop_via_update_from_output():
840
841
842
def test_check_stop_min_tokens():
843
- """Test that requests don't stop when min_tokens requirement isn't met."""
843
+ """Test that sampled stop tokens still finish under min_tokens."""
844
from vllm.v1.core.sched.utils import check_stop
845
846
- # Test case 1: num_output_tokens < min_tokens
847
- # Should return False (don't stop)
846
+ # Test case 1: regular tokens below min_tokens should keep running.
847
sampling_params = SamplingParams(
848
ignore_eos=False,
849
max_tokens=20,
@@ -857,30 +856,41 @@ def test_check_stop_min_tokens():
856
sampling_params=sampling_params,
857
pooling_params=None,
858
)
860
- # Simulate having generated 3 output tokens (less than min_tokens=5)
861
- request.append_output_token_ids([10, 11, EOS_TOKEN_ID]) # EOS token present
859
+ request.append_output_token_ids([10, 11, 12])
860
861
result = check_stop(request, max_model_len=100)
862
assert result is False, "Should not stop when num_output_tokens<min_tokens"
863
866
- # Test case 2: num_output_tokens >= min_tokens
867
- # Should follow normal stopping logic (stop on EOS)
868
- request.append_output_token_ids(
869
- [
870
- 10,
871
- 11,
872
- 12,
873
- 13,
874
- 14,
875
- EOS_TOKEN_ID,
876
- ]
877
- ) # 6 tokens > min_tokens
864
+ # Test case 2: if EOS is sampled, it should stop even below min_tokens.
865
+ request_eos = Request(
866
+ request_id="eos",
867
+ prompt_token_ids=[0, 1, 2],
868
+ sampling_params=sampling_params,
869
+ pooling_params=None,
870
+ )
871
+ request_eos.append_output_token_ids([10, 11, EOS_TOKEN_ID])
872
+ result = check_stop(request_eos, max_model_len=100)
873
+ assert result is True, "Sampled EOS should stop even below min_tokens"
874
+ assert request_eos.status == RequestStatus.FINISHED_STOPPED
875
879
- result = check_stop(request, max_model_len=100)
880
- assert result is True, "Should stop on EOS when min_tokens met"
881
- assert request.status == RequestStatus.FINISHED_STOPPED
876
+ # Test case 3: ignore_eos should still prevent EOS-based stopping.
877
+ sampling_params_ignore_eos = SamplingParams(
878
+ ignore_eos=True,
879
+ max_tokens=20,
880
+ min_tokens=5,
881
+ )
882
+ sampling_params_ignore_eos.update_from_generation_config({}, EOS_TOKEN_ID)
883
+ request_ignore_eos = Request(
884
+ request_id="ignore-eos",
885
+ prompt_token_ids=[0, 1, 2],
886
+ sampling_params=sampling_params_ignore_eos,
887
+ pooling_params=None,
888
+ )
889
+ request_ignore_eos.append_output_token_ids([EOS_TOKEN_ID])
890
+ result = check_stop(request_ignore_eos, max_model_len=100)
891
+ assert result is False, "ignore_eos should still suppress EOS stopping"
892
883
- # Test case 3: min_tokens = 0, should follow normal stopping logic
893
+ # Test case 4: min_tokens = 0 should follow normal stopping logic.
894
sampling_params_no_min = SamplingParams(
895
ignore_eos=False,
896
max_tokens=20,
@@ -899,7 +909,7 @@ def test_check_stop_min_tokens():
909
assert result is True, "Should stop on EOS when min_tokens=0"
910
assert request_no_min.status == RequestStatus.FINISHED_STOPPED
911
902
- # Test case 4: min_tokens > 0 with stop token (not EOS)
912
+ # Test case 5: if a stop token is sampled, it should stop immediately.
913
sampling_params_stop = SamplingParams(
914
ignore_eos=False,
915
max_tokens=20,
@@ -913,22 +923,67 @@ def test_check_stop_min_tokens():
923
sampling_params=sampling_params_stop,
924
pooling_params=None,
925
)
916
- # Only 3 output tokens, less than min_tokens=5, but has stop token
926
request_stop.append_output_token_ids([10, 11, 42])
927
result = check_stop(request_stop, max_model_len=100)
919
- assert result is False, "Should not stop when num_output_tokens<min_tokens"
920
-
921
- # Test case 5: min_tokens met, should stop on stop token
922
- request_stop.append_output_token_ids(
923
- [10, 11, 12, 13, 14, 42]
924
- ) # 6 tokens >= min_tokens=5
925
-
926
- result = check_stop(request_stop, max_model_len=100)
927
- assert result is True, "Should stop on stop token when min_tokens met"
928
+ assert result is True, "Sampled stop tokens should stop even below min_tokens"
929
assert request_stop.status == RequestStatus.FINISHED_STOPPED
930
assert request_stop.stop_reason == 42
931
932
933
+def test_check_stop_length_cap_below_min_tokens():
934
+ """Length caps must finish a request even below min_tokens.
935
+
936
+ Guards against requests hanging forever when min_tokens cannot be
937
+ satisfied: min_tokens only defers EOS/stop tokens, so a request that
938
+ reaches max_model_len (or a max_tokens clamped below min_tokens) must
939
+ still finish as FINISHED_LENGTH_CAPPED instead of spinning in the
940
+ scheduler with zero new tokens to schedule.
941
+ """
942
+ from vllm.v1.core.sched.utils import check_stop
943
+
944
+ # Reaching max_model_len below min_tokens must stop the request.
945
+ sampling_params = SamplingParams(
946
+ ignore_eos=True,
947
+ max_tokens=500,
948
+ min_tokens=500,
949
+ )
950
+ sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
951
+ request = Request(
952
+ request_id="0",
953
+ prompt_token_ids=[0, 1, 2],
954
+ sampling_params=sampling_params,
955
+ pooling_params=None,
956
+ )
957
+ # 7 output tokens: num_tokens == 10 == max_model_len, min_tokens unmet.
958
+ request.append_output_token_ids([10, 11, 12, 13, 14, 15, 16])
959
+
960
+ result = check_stop(request, max_model_len=10)
961
+ assert result is True, "Should stop at max_model_len even below min_tokens"
962
+ assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED
963
+
964
+ # Reaching max_tokens below min_tokens must stop the request. This can
965
+ # happen when max_tokens is clamped to the context window after
966
+ # SamplingParams validation (e.g. max_tokens=None on a long prompt).
967
+ sampling_params_capped = SamplingParams(
968
+ ignore_eos=True,
969
+ max_tokens=500,
970
+ min_tokens=500,
971
+ )
972
+ sampling_params_capped.max_tokens = 3
973
+ sampling_params_capped.update_from_generation_config({}, EOS_TOKEN_ID)
974
+ request_capped = Request(
975
+ request_id="1",
976
+ prompt_token_ids=[0, 1, 2],
977
+ sampling_params=sampling_params_capped,
978
+ pooling_params=None,
979
+ )
980
+ request_capped.append_output_token_ids([10, 11, 12])
981
+
982
+ result = check_stop(request_capped, max_model_len=100)
983
+ assert result is True, "Should stop at max_tokens even below min_tokens"
984
+ assert request_capped.status == RequestStatus.FINISHED_LENGTH_CAPPED
985
+
986
+
987
@pytest.mark.parametrize(
988
"enable_prefix_caching, prompt_logprobs",
989
[
tests/v1/e2e/general/test_min_tokens.py
+24
@@ -13,11 +13,14 @@ Covers:
13
5) Multiple stop conditions
14
"""
15
16
+import json
17
+
18
import pytest
19
20
from vllm import LLM, SamplingParams
21
from vllm.exceptions import VLLMValidationError
22
from vllm.outputs import RequestOutput
23
+from vllm.sampling_params import StructuredOutputsParams
24
25
# Test configuration
26
TEST_MODEL = "facebook/opt-125m" # Small model for fast CI execution
@@ -168,6 +171,27 @@ def llm_v1():
171
return llm
172
173
174
+def test_min_tokens_with_terminal_structured_output(llm_v1: LLM):
175
+ """A grammar whose only legal continuation is EOS must stop under
176
+ min_tokens instead of sampling outside the grammar."""
177
+ params = SamplingParams(
178
+ temperature=GREEDY,
179
+ min_tokens=20,
180
+ max_tokens=40,
181
+ structured_outputs=StructuredOutputsParams(
182
+ json=json.dumps({"type": "boolean"})
183
+ ),
184
+ )
185
+
186
+ output = llm_v1.generate("Answer:", params)[0].outputs[0]
187
+ eos_token_id = llm_v1.get_tokenizer().eos_token_id
188
+
189
+ assert output.finish_reason == "stop"
190
+ assert isinstance(json.loads(output.text), bool)
191
+ assert output.token_ids[-1] == eos_token_id
192
+ assert len(output.token_ids) < params.min_tokens
193
+
194
+
195
def get_token_count(output: RequestOutput) -> int:
196
"""Extract token count from LLM output"""
197
if not output.outputs:
tests/v1/logits_processors/test_correctness.py
+113
-2
@@ -20,7 +20,11 @@ from tests.v1.sample.utils import (
20
)
21
from vllm.config import VllmConfig
22
from vllm.platforms import current_platform
23
-from vllm.sampling_params import SamplingParams, validate_thinking_token_budget
23
+from vllm.sampling_params import (
24
+ SamplingParams,
25
+ StructuredOutputsParams,
26
+ validate_thinking_token_budget,
27
+)
28
from vllm.utils.platform_utils import is_pin_memory_available
29
from vllm.v1.sample.logits_processor import (
30
BatchUpdate,
@@ -371,7 +375,7 @@ def _min_tokens_validate(
375
376
# Validate min-token logits processor state
377
if min_tok:
374
- (_, out_tok, all_stop_token_ids) = min_tok
378
+ (_, out_tok, all_stop_token_ids, _) = min_tok
379
num_out_tokens = len(out_tok)
380
if num_out_tokens != ref_num_out_tokens:
381
_raise_error_invalid(
@@ -448,6 +452,113 @@ def _min_tokens_validate(
452
)
453
454
455
+def _make_min_tokens_processor() -> MinTokensLogitsProcessor:
456
+ return MinTokensLogitsProcessor(VllmConfig(), torch.device("cpu"), False)
457
+
458
+
459
+# Constructing the processor allocates pinned host memory, which creates a
460
+# CUDA driver context even though the tensors stay on CPU. Run these in their
461
+# own process so the pytest process can still fork for `test_logitsprocs`.
462
+@create_new_process_for_each_test()
463
+def test_min_tokens_keeps_all_masked_behavior_without_structured_output():
464
+ processor = _make_min_tokens_processor()
465
+ processor.update_state(
466
+ BatchUpdate(
467
+ batch_size=1,
468
+ removed=(),
469
+ added=[
470
+ (
471
+ 0,
472
+ SamplingParams(min_tokens=2, stop_token_ids=[0]),
473
+ None,
474
+ [],
475
+ )
476
+ ],
477
+ moved=(),
478
+ )
479
+ )
480
+
481
+ logits = torch.full((1, 3), -float("inf"))
482
+ logits[0, 0] = 1.0
483
+
484
+ processor.apply(logits)
485
+
486
+ assert torch.isneginf(logits).all()
487
+
488
+
489
+@create_new_process_for_each_test()
490
+def test_min_tokens_restores_all_masked_structured_output_stop_token():
491
+ processor = _make_min_tokens_processor()
492
+ processor.update_state(
493
+ BatchUpdate(
494
+ batch_size=1,
495
+ removed=(),
496
+ added=[
497
+ (
498
+ 0,
499
+ SamplingParams(
500
+ min_tokens=2,
501
+ stop_token_ids=[0],
502
+ structured_outputs=StructuredOutputsParams(
503
+ json={"type": "boolean"}
504
+ ),
505
+ ),
506
+ None,
507
+ [],
508
+ )
509
+ ],
510
+ moved=(),
511
+ )
512
+ )
513
+
514
+ logits = torch.full((1, 3), -float("inf"))
515
+ logits[0, 0] = 1.0
516
+
517
+ processor.apply(logits)
518
+
519
+ assert logits[0, 0] == 1.0
520
+ assert torch.isneginf(logits[0, 1:]).all()
521
+
522
+
523
+@create_new_process_for_each_test()
524
+def test_min_tokens_restores_all_masked_structured_output_stop_token_spec_decode():
525
+ processor = _make_min_tokens_processor()
526
+ processor.update_state(
527
+ BatchUpdate(
528
+ batch_size=1,
529
+ removed=(),
530
+ added=[
531
+ (
532
+ 0,
533
+ SamplingParams(
534
+ min_tokens=2,
535
+ stop_token_ids=[0, 1],
536
+ structured_outputs=StructuredOutputsParams(
537
+ json={"type": "boolean"}
538
+ ),
539
+ ),
540
+ None,
541
+ [],
542
+ )
543
+ ],
544
+ moved=(),
545
+ )
546
+ )
547
+
548
+ logits = torch.full((2, 4), -float("inf"))
549
+ logits[0, 0] = 1.0
550
+ logits[0, 1] = 2.0
551
+ logits[1, 0] = 3.0
552
+
553
+ processor.apply_with_spec_decode(logits, [2])
554
+
555
+ assert logits[0, 0] == 1.0
556
+ assert logits[0, 1] == 2.0
557
+ assert logits[1, 0] == 3.0
558
+ assert torch.isneginf(logits[1, 1:]).all()
559
+ assert torch.isneginf(logits[:, 2:]).all()
560
+
561
+
562
def _thinking_budget_params(kwargs: dict) -> None:
563
"""Set SamplingParams kwargs for thinking token budget tests"""
564
kwargs["thinking_token_budget"] = THINKING_TOKEN_BUDGET
tests/v1/worker/test_gpu_logit_bias.py
new
+92
@@ -0,0 +1,92 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import numpy as np
5
+import pytest
6
+import torch
7
+
8
+pytest.importorskip("triton")
9
+if not torch.cuda.is_available():
10
+ pytest.skip(
11
+ "CUDA required for Model Runner V2 logit bias tests",
12
+ allow_module_level=True,
13
+ )
14
+
15
+from vllm.sampling_params import SamplingParams, StructuredOutputsParams
16
+from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
17
+
18
+DEVICE = torch.device("cuda")
19
+VOCAB_SIZE = 128
20
+STOP_TOKEN = 3
21
+OTHER_TOKEN = 4
22
+PROMPT_LEN = 2
23
+MIN_TOKENS = 10
24
+# min_len is PROMPT_LEN + MIN_TOKENS, so this position is below the minimum
25
+# and the kernel masks stop tokens.
26
+POS = 5
27
+
28
+
29
+def _params(structured: bool) -> SamplingParams:
30
+ return SamplingParams(
31
+ min_tokens=MIN_TOKENS,
32
+ stop_token_ids=[STOP_TOKEN],
33
+ structured_outputs=(
34
+ StructuredOutputsParams(json={"type": "boolean"}) if structured else None
35
+ ),
36
+ )
37
+
38
+
39
+def _only_stop_token_left(num_rows: int) -> torch.Tensor:
40
+ """Logits as a terminal grammar leaves them: every token masked but one."""
41
+ logits = torch.full((num_rows, VOCAB_SIZE), -float("inf"), device=DEVICE)
42
+ logits[:, STOP_TOKEN] = 1.0
43
+ return logits
44
+
45
+
46
+def _apply(logits: torch.Tensor, structured: list[bool]) -> torch.Tensor:
47
+ """Run the fused bias kernel with one logits row per request."""
48
+ state = LogitBiasState(max_num_reqs=4, device=DEVICE)
49
+ for req_idx, is_structured in enumerate(structured):
50
+ state.add_request(req_idx, PROMPT_LEN, _params(is_structured))
51
+ state.apply_staged_writes()
52
+
53
+ num_reqs = len(structured)
54
+ state.apply_logit_bias(
55
+ logits,
56
+ torch.arange(num_reqs, dtype=torch.int32, device=DEVICE),
57
+ np.arange(num_reqs, dtype=np.intp),
58
+ torch.full((num_reqs,), POS, dtype=torch.int32, device=DEVICE),
59
+ )
60
+ return logits.cpu()
61
+
62
+
63
+def test_v2_min_tokens_restores_stop_token_when_row_would_be_empty():
64
+ """Structured output plus min_tokens must not leave an all -inf row."""
65
+ out = _apply(_only_stop_token_left(1), structured=[True])
66
+ assert out[0, STOP_TOKEN] == 1.0
67
+
68
+
69
+def test_v2_min_tokens_keeps_empty_row_without_structured_output():
70
+ """Plain min_tokens requests keep the pre-existing masking behavior."""
71
+ out = _apply(_only_stop_token_left(1), structured=[False])
72
+ assert torch.isneginf(out).all()
73
+
74
+
75
+def test_v2_min_tokens_keeps_masking_when_other_candidates_remain():
76
+ """The fallback must not weaken min_tokens while a non-stop token is legal."""
77
+ logits = _only_stop_token_left(1)
78
+ logits[0, OTHER_TOKEN] = 2.0
79
+
80
+ out = _apply(logits, structured=[True])
81
+
82
+ assert torch.isneginf(out[0, STOP_TOKEN])
83
+ assert out[0, OTHER_TOKEN] == 2.0
84
+
85
+
86
+def test_v2_min_tokens_mixed_batch_gates_restore_per_request():
87
+ """CHECK_ALL_MASKED_ROWS is a batch-wide constexpr, so a mixed batch must
88
+ still consult the per-request flag."""
89
+ out = _apply(_only_stop_token_left(2), structured=[True, False])
90
+
91
+ assert out[0, STOP_TOKEN] == 1.0
92
+ assert torch.isneginf(out[1]).all()
vllm/v1/core/sched/utils.py
+9
-3
@@ -97,9 +97,6 @@ def check_stop(request: Request, max_model_len: int) -> bool:
97
sampling_params = request.sampling_params
98
assert sampling_params is not None
99
100
- if request.num_output_tokens < sampling_params.min_tokens:
101
- return False
102
-
100
last_token_id = request.output_token_ids[-1]
101
if last_token_id == sampling_params.eos_token_id:
102
request.status = RequestStatus.FINISHED_STOPPED
@@ -109,6 +106,7 @@ def check_stop(request: Request, max_model_len: int) -> bool:
106
request.status = RequestStatus.FINISHED_STOPPED
107
request.stop_reason = last_token_id
108
return True
109
+
110
if (
111
request.num_tokens >= max_model_len
112
or request.num_output_tokens >= request.max_tokens
@@ -116,6 +114,14 @@ def check_stop(request: Request, max_model_len: int) -> bool:
114
request.status = RequestStatus.FINISHED_LENGTH_CAPPED
115
return True
116
117
+ # Note(arpera):
118
+ # Order of checks is important for min_tokens
119
+ # If you decide to change the order, first look at these PRs please:
120
+ # 47489, 49521, and 51299
121
+ # These PRs show some reasoning about the order of checks
122
+ if request.num_output_tokens < sampling_params.min_tokens:
123
+ return False
124
+
125
repetition_detection = sampling_params.repetition_detection
126
if repetition_detection is not None and (
127
check_sequence_repetition(
vllm/v1/sample/logits_processor/builtin.py
+103
-18
@@ -166,16 +166,22 @@ class MinTokensLogitsProcessor(LogitsProcessor):
166
def __init__(
167
self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool
168
):
169
- # index -> (min_toks, output_token_ids, stop_token_ids)
169
+ # index -> (min_toks, output_token_ids, stop_token_ids, uses_structured_output)
170
self.device = device
171
- self.min_toks: dict[int, tuple[int, Sequence[int], set[int]]] = {}
171
+ self.min_toks: dict[int, tuple[int, Sequence[int], set[int], bool]] = {}
172
173
- # (req_idx_tensor,eos_tok_id_tensor)
173
+ # (req_idx_tensor, stop_token_id_tensor)
174
self.logits_slice: tuple[torch.Tensor, torch.Tensor] = (
175
self._device_tensor([], torch.int32),
176
self._device_tensor([], torch.int32),
177
)
178
179
+ # Structured-output stop-token logits that may need to be restored.
180
+ self.restore_logits_slice: tuple[torch.Tensor, torch.Tensor] = (
181
+ self._device_tensor([], torch.int32),
182
+ self._device_tensor([], torch.int32),
183
+ )
184
+
185
self.neg_inf_tensor = torch.tensor(
186
-float("inf"), dtype=torch.float32, device=self.device
187
)
@@ -188,11 +194,16 @@ class MinTokensLogitsProcessor(LogitsProcessor):
194
@staticmethod
195
def add_request(
196
params: SamplingParams, _: list[int] | None, output_tok_ids: list[int]
191
- ) -> tuple[int, Sequence[int], set[int]] | None:
197
+ ) -> tuple[int, Sequence[int], set[int], bool] | None:
198
min_tokens = params.min_tokens
199
if not min_tokens or len(output_tok_ids) >= min_tokens:
200
return None
195
- return min_tokens, output_tok_ids, params.all_stop_token_ids
201
+ return (
202
+ min_tokens,
203
+ output_tok_ids,
204
+ params.all_stop_token_ids,
205
+ params.structured_outputs is not None,
206
+ )
207
208
def update_state(self, batch_update: BatchUpdate | None):
209
needs_update = process_dict_updates(
@@ -202,7 +213,7 @@ class MinTokensLogitsProcessor(LogitsProcessor):
213
# Check for any requests that have attained their min tokens.
214
to_remove = tuple(
215
index
205
- for index, (min_toks, out_tok_ids, _) in self.min_toks.items()
216
+ for index, (min_toks, out_tok_ids, _, _) in self.min_toks.items()
217
if len(out_tok_ids) >= min_toks
218
)
219
if to_remove:
@@ -214,22 +225,66 @@ class MinTokensLogitsProcessor(LogitsProcessor):
225
if needs_update:
226
reqs: list[int] = []
227
tok_ids: list[int] = []
217
- for req, (_, _, stop_tok_ids) in self.min_toks.items():
228
+ restore_reqs: list[int] = []
229
+ restore_tok_ids: list[int] = []
230
+ for req, (
231
+ _,
232
+ _,
233
+ stop_tok_ids,
234
+ uses_structured_output,
235
+ ) in self.min_toks.items():
236
reqs.extend([req] * len(stop_tok_ids))
237
tok_ids.extend(stop_tok_ids)
238
+ if uses_structured_output:
239
+ restore_reqs.extend([req] * len(stop_tok_ids))
240
+ restore_tok_ids.extend(stop_tok_ids)
241
242
self.logits_slice = (
243
self._device_tensor(reqs, torch.int32),
244
self._device_tensor(tok_ids, torch.int32),
245
)
246
+ self.restore_logits_slice = (
247
+ self._device_tensor(restore_reqs, torch.int32),
248
+ self._device_tensor(restore_tok_ids, torch.int32),
249
+ )
250
251
def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor:
252
return async_tensor_h2d(data, device=self.device, dtype=dtype)
253
254
+ def _mask_stop_token_logits(
255
+ self,
256
+ logits: torch.Tensor,
257
+ logits_slice: tuple[torch.Tensor, torch.Tensor],
258
+ restore_logits_slice: tuple[torch.Tensor, torch.Tensor],
259
+ ) -> None:
260
+ restore_rows, restore_toks = restore_logits_slice
261
+ # Advanced indexing already returns a copy, so this survives the
262
+ # in-place masking below.
263
+ stop_logits = logits[restore_logits_slice]
264
+ logits.index_put_(logits_slice, self.neg_inf_tensor)
265
+
266
+ # Stop-token entries for the same row are emitted together, so this
267
+ # checks each row once without sorting the index tensor.
268
+ unique_rows, inverse_indices = torch.unique_consecutive(
269
+ restore_rows, return_inverse=True
270
+ )
271
+ row_needs_restore = torch.isneginf(logits[unique_rows]).all(dim=-1)
272
+ restore_mask = row_needs_restore[inverse_indices] & torch.isfinite(stop_logits)
273
+ restore_slice = (
274
+ restore_rows[restore_mask],
275
+ restore_toks[restore_mask],
276
+ )
277
+ logits.index_put_(restore_slice, stop_logits[restore_mask])
278
+
279
def apply(self, logits: torch.Tensor) -> torch.Tensor:
280
if self.min_toks:
231
- # Inhibit EOS token for requests which have not reached min length
232
- logits.index_put_(self.logits_slice, self.neg_inf_tensor)
281
+ # Inhibit stop tokens for requests which have not reached min length.
282
+ if self.restore_logits_slice[0].numel() == 0:
283
+ logits.index_put_(self.logits_slice, self.neg_inf_tensor)
284
+ else:
285
+ self._mask_stop_token_logits(
286
+ logits, self.logits_slice, self.restore_logits_slice
287
+ )
288
return logits
289
290
def apply_with_spec_decode(
@@ -238,10 +293,13 @@ class MinTokensLogitsProcessor(LogitsProcessor):
293
num_draft_tokens: list[int],
294
) -> torch.Tensor:
295
"""Spec-decode version of apply().
241
- Priority: ``min_tokens`` > ``stop_token_ids`` / EOS.
296
+
297
+ Stop tokens stay masked unless masking them would leave every logit
298
+ in a row at -inf.
299
+
300
Example: ``num_draft_tokens = [2, 3, 1]``
243
- → ``logits`` shape ``[6, V]``, ``cumsum = [0, 2, 5, 6]``
244
- → request 0 owns rows 0‑1, request 1 rows 2‑4, request 2 row 5.
301
+ -> ``logits`` shape ``[6, V]``, ``cumsum = [0, 2, 5, 6]``
302
+ -> request 0 owns rows 0-1, request 1 rows 2-4, request 2 row 5.
303
"""
304
if not self.min_toks:
305
return logits
@@ -250,8 +308,19 @@ class MinTokensLogitsProcessor(LogitsProcessor):
308
cumsum = np.concatenate([[0], np.cumsum(num_draft_arr)])
309
310
entries = [
253
- (req_idx, min_tok, len(out_tok_ids), list(stop_tok_ids))
254
- for req_idx, (min_tok, out_tok_ids, stop_tok_ids) in self.min_toks.items()
311
+ (
312
+ req_idx,
313
+ min_tok,
314
+ len(out_tok_ids),
315
+ list(stop_tok_ids),
316
+ uses_structured_output,
317
+ )
318
+ for req_idx, (
319
+ min_tok,
320
+ out_tok_ids,
321
+ stop_tok_ids,
322
+ uses_structured_output,
323
+ ) in self.min_toks.items()
324
if stop_tok_ids
325
]
326
@@ -260,8 +329,10 @@ class MinTokensLogitsProcessor(LogitsProcessor):
329
330
all_rows: list[np.ndarray] = [] # row indices to mask
331
all_toks: list[np.ndarray] = [] # stop-token ids at those rows
332
+ restore_rows: list[np.ndarray] = []
333
+ restore_toks: list[np.ndarray] = []
334
264
- for req_idx, min_tok, current_len, stop_toks in entries:
335
+ for req_idx, min_tok, current_len, stop_toks, uses_structured_output in entries:
336
remaining = min_tok - current_len
337
# How many leading draft positions still need stop-token masking.
338
n_mask = int(min(max(remaining, 0), num_draft_arr[req_idx]))
@@ -270,8 +341,13 @@ class MinTokensLogitsProcessor(LogitsProcessor):
341
offset = cumsum[req_idx]
342
row_indices = np.arange(offset, offset + n_mask, dtype=np.int64)
343
n_stop = len(stop_toks)
273
- all_rows.append(np.repeat(row_indices, n_stop))
274
- all_toks.append(np.tile(stop_toks, n_mask))
344
+ rows = np.repeat(row_indices, n_stop)
345
+ toks = np.tile(stop_toks, n_mask)
346
+ all_rows.append(rows)
347
+ all_toks.append(toks)
348
+ if uses_structured_output:
349
+ restore_rows.append(rows)
350
+ restore_toks.append(toks)
351
352
if all_rows:
353
rows_arr = np.concatenate(all_rows)
@@ -281,7 +357,16 @@ class MinTokensLogitsProcessor(LogitsProcessor):
357
async_tensor_h2d(rows_arr, device=self.device),
358
async_tensor_h2d(toks_arr, device=self.device),
359
)
284
- logits.index_put_(logits_slice, self.neg_inf_tensor)
360
+ if restore_rows:
361
+ restore_rows_arr = np.concatenate(restore_rows)
362
+ restore_toks_arr = np.concatenate(restore_toks)
363
+ restore_logits_slice = (
364
+ async_tensor_h2d(restore_rows_arr, device=self.device),
365
+ async_tensor_h2d(restore_toks_arr, device=self.device),
366
+ )
367
+ self._mask_stop_token_logits(logits, logits_slice, restore_logits_slice)
368
+ else:
369
+ logits.index_put_(logits_slice, self.neg_inf_tensor)
370
371
return logits
372
vllm/v1/worker/gpu/sample/logit_bias.py
+73
-5
@@ -40,6 +40,9 @@ class LogitBiasState:
40
# Min tokens.
41
self.min_lens = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
42
self.num_stop_token_ids = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
43
+ self.restore_when_all_masked = UvaBackedTensor(
44
+ self.max_num_reqs, dtype=torch.int32
45
+ )
46
self.stop_token_ids = StagedWriteTensor(
47
(self.max_num_reqs, MAX_NUM_STOP_TOKEN_IDS),
48
dtype=torch.int32,
@@ -100,9 +103,13 @@ class LogitBiasState:
103
)
104
self.num_stop_token_ids.np[req_idx] = num_stop_token_ids
105
self.stop_token_ids.stage_write(req_idx, 0, stop_token_ids)
106
+ self.restore_when_all_masked.np[req_idx] = int(
107
+ sampling_params.structured_outputs is not None
108
+ )
109
use_logit_bias = True
110
else:
111
self.num_stop_token_ids.np[req_idx] = 0
112
+ self.restore_when_all_masked.np[req_idx] = 0
113
114
self.use_logit_bias[req_idx] = use_logit_bias
115
@@ -116,6 +123,7 @@ class LogitBiasState:
123
124
self.min_lens.copy_to_uva()
125
self.num_stop_token_ids.copy_to_uva()
126
+ self.restore_when_all_masked.copy_to_uva()
127
self.stop_token_ids.apply_write()
128
129
def apply_logit_bias(
@@ -129,6 +137,10 @@ class LogitBiasState:
137
# No request uses logit bias. Skip the kernel launch.
138
return
139
140
+ enable_stop_token_restore = bool(
141
+ np.any(self.restore_when_all_masked.np[idx_mapping_np])
142
+ )
143
+
144
apply_logit_bias(
145
logits,
146
expanded_idx_mapping,
@@ -140,7 +152,9 @@ class LogitBiasState:
152
self.logit_bias.gpu,
153
self.min_lens.gpu,
154
self.num_stop_token_ids.gpu,
155
+ self.restore_when_all_masked.gpu,
156
self.stop_token_ids.gpu,
157
+ enable_stop_token_restore,
158
)
159
160
@@ -164,10 +178,12 @@ def _bias_kernel(
178
pos_ptr,
179
min_lens_ptr,
180
num_stop_token_ids_ptr,
181
+ restore_when_all_masked_ptr,
182
stop_token_ids_ptr,
183
stop_token_ids_stride,
184
BLOCK_SIZE: tl.constexpr,
185
LOGITS_BLOCK_SIZE: tl.constexpr,
186
+ CHECK_ALL_MASKED_ROWS: tl.constexpr,
187
):
188
token_idx = tl.program_id(0).to(tl.int64)
189
req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
@@ -232,11 +248,59 @@ def _bias_kernel(
248
stop_token_ids_ptr + req_state_idx * stop_token_ids_stride + block,
249
mask=mask,
250
)
235
- tl.store(
236
- logits_ptr + token_idx * logits_stride + stop_token_ids,
237
- -float("inf"),
238
- mask=mask,
239
- )
251
+ if CHECK_ALL_MASKED_ROWS:
252
+ should_restore_stop_logits = tl.load(
253
+ restore_when_all_masked_ptr + req_state_idx
254
+ )
255
+ if should_restore_stop_logits:
256
+ stop_logits = tl.load(
257
+ logits_ptr + token_idx * logits_stride + stop_token_ids,
258
+ mask=mask,
259
+ other=-float("inf"),
260
+ )
261
+
262
+ # Save must read original logits before the -inf overwrite.
263
+ tl.debug_barrier()
264
+
265
+ tl.store(
266
+ logits_ptr + token_idx * logits_stride + stop_token_ids,
267
+ -float("inf"),
268
+ mask=mask,
269
+ )
270
+
271
+ # The -inf overwrite must be visible to the row scan below.
272
+ tl.debug_barrier()
273
+
274
+ row_max = tl.full((), -float("inf"), tl.float32)
275
+ for i in range(0, vocab_size, LOGITS_BLOCK_SIZE):
276
+ offset = i + tl.arange(0, LOGITS_BLOCK_SIZE)
277
+ logits = tl.load(
278
+ logits_ptr + token_idx * logits_stride + offset,
279
+ mask=offset < vocab_size,
280
+ other=-float("inf"),
281
+ )
282
+ row_max = tl.maximum(row_max, tl.max(logits, axis=0))
283
+
284
+ if row_max == -float("inf"):
285
+ tl.store(
286
+ logits_ptr + token_idx * logits_stride + stop_token_ids,
287
+ stop_logits,
288
+ mask=mask
289
+ & (stop_logits > -float("inf"))
290
+ & (stop_logits < float("inf")),
291
+ )
292
+ else:
293
+ tl.store(
294
+ logits_ptr + token_idx * logits_stride + stop_token_ids,
295
+ -float("inf"),
296
+ mask=mask,
297
+ )
298
+ else:
299
+ tl.store(
300
+ logits_ptr + token_idx * logits_stride + stop_token_ids,
301
+ -float("inf"),
302
+ mask=mask,
303
+ )
304
305
306
def apply_logit_bias(
@@ -250,7 +314,9 @@ def apply_logit_bias(
314
logit_bias: torch.Tensor,
315
min_lens: torch.Tensor,
316
num_stop_token_ids: torch.Tensor,
317
+ restore_when_all_masked: torch.Tensor,
318
stop_token_ids: torch.Tensor,
319
+ check_all_masked_rows: bool = False,
320
) -> None:
321
num_tokens, vocab_size = logits.shape
322
BLOCK_SIZE = triton.next_power_of_2(
@@ -277,8 +343,10 @@ def apply_logit_bias(
343
pos,
344
min_lens,
345
num_stop_token_ids,
346
+ restore_when_all_masked,
347
stop_token_ids,
348
stop_token_ids.stride(0),
349
BLOCK_SIZE=BLOCK_SIZE,
350
LOGITS_BLOCK_SIZE=LOGITS_BLOCK_SIZE,
351
+ CHECK_ALL_MASKED_ROWS=check_all_masked_rows,
352
)