From 39e276eaeb9daed06a180f6a8d187bbb8790e97b Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Mon, 31 Aug 2026 19:39:27 +0300 Subject: [PATCH] [Structured Output] Let terminal grammars stop under min_tokens (II) (#54218) Signed-off-by: Artem Perevedentsev --- tests/v1/core/test_scheduler.py | 119 ++++++++++++----- tests/v1/e2e/general/test_min_tokens.py | 24 ++++ .../v1/logits_processors/test_correctness.py | 115 ++++++++++++++++- tests/v1/worker/test_gpu_logit_bias.py | 92 +++++++++++++ vllm/v1/core/sched/utils.py | 12 +- vllm/v1/sample/logits_processor/builtin.py | 121 +++++++++++++++--- vllm/v1/worker/gpu/sample/logit_bias.py | 78 ++++++++++- 7 files changed, 501 insertions(+), 60 deletions(-) create mode 100644 tests/v1/worker/test_gpu_logit_bias.py diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 7b730523fa..c630abe4f3 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -840,11 +840,10 @@ def test_stop_via_update_from_output(): def test_check_stop_min_tokens(): - """Test that requests don't stop when min_tokens requirement isn't met.""" + """Test that sampled stop tokens still finish under min_tokens.""" from vllm.v1.core.sched.utils import check_stop - # Test case 1: num_output_tokens < min_tokens - # Should return False (don't stop) + # Test case 1: regular tokens below min_tokens should keep running. sampling_params = SamplingParams( ignore_eos=False, max_tokens=20, @@ -857,30 +856,41 @@ def test_check_stop_min_tokens(): sampling_params=sampling_params, pooling_params=None, ) - # Simulate having generated 3 output tokens (less than min_tokens=5) - request.append_output_token_ids([10, 11, EOS_TOKEN_ID]) # EOS token present + request.append_output_token_ids([10, 11, 12]) result = check_stop(request, max_model_len=100) assert result is False, "Should not stop when num_output_tokens= min_tokens - # Should follow normal stopping logic (stop on EOS) - request.append_output_token_ids( - [ - 10, - 11, - 12, - 13, - 14, - EOS_TOKEN_ID, - ] - ) # 6 tokens > min_tokens + # Test case 2: if EOS is sampled, it should stop even below min_tokens. + request_eos = Request( + request_id="eos", + prompt_token_ids=[0, 1, 2], + sampling_params=sampling_params, + pooling_params=None, + ) + request_eos.append_output_token_ids([10, 11, EOS_TOKEN_ID]) + result = check_stop(request_eos, max_model_len=100) + assert result is True, "Sampled EOS should stop even below min_tokens" + assert request_eos.status == RequestStatus.FINISHED_STOPPED - result = check_stop(request, max_model_len=100) - assert result is True, "Should stop on EOS when min_tokens met" - assert request.status == RequestStatus.FINISHED_STOPPED + # Test case 3: ignore_eos should still prevent EOS-based stopping. + sampling_params_ignore_eos = SamplingParams( + ignore_eos=True, + max_tokens=20, + min_tokens=5, + ) + sampling_params_ignore_eos.update_from_generation_config({}, EOS_TOKEN_ID) + request_ignore_eos = Request( + request_id="ignore-eos", + prompt_token_ids=[0, 1, 2], + sampling_params=sampling_params_ignore_eos, + pooling_params=None, + ) + request_ignore_eos.append_output_token_ids([EOS_TOKEN_ID]) + result = check_stop(request_ignore_eos, max_model_len=100) + assert result is False, "ignore_eos should still suppress EOS stopping" - # Test case 3: min_tokens = 0, should follow normal stopping logic + # Test case 4: min_tokens = 0 should follow normal stopping logic. sampling_params_no_min = SamplingParams( ignore_eos=False, max_tokens=20, @@ -899,7 +909,7 @@ def test_check_stop_min_tokens(): assert result is True, "Should stop on EOS when min_tokens=0" assert request_no_min.status == RequestStatus.FINISHED_STOPPED - # Test case 4: min_tokens > 0 with stop token (not EOS) + # Test case 5: if a stop token is sampled, it should stop immediately. sampling_params_stop = SamplingParams( ignore_eos=False, max_tokens=20, @@ -913,22 +923,67 @@ def test_check_stop_min_tokens(): sampling_params=sampling_params_stop, pooling_params=None, ) - # Only 3 output tokens, less than min_tokens=5, but has stop token request_stop.append_output_token_ids([10, 11, 42]) result = check_stop(request_stop, max_model_len=100) - assert result is False, "Should not stop when num_output_tokens= min_tokens=5 - - result = check_stop(request_stop, max_model_len=100) - assert result is True, "Should stop on stop token when min_tokens met" + assert result is True, "Sampled stop tokens should stop even below min_tokens" assert request_stop.status == RequestStatus.FINISHED_STOPPED assert request_stop.stop_reason == 42 +def test_check_stop_length_cap_below_min_tokens(): + """Length caps must finish a request even below min_tokens. + + Guards against requests hanging forever when min_tokens cannot be + satisfied: min_tokens only defers EOS/stop tokens, so a request that + reaches max_model_len (or a max_tokens clamped below min_tokens) must + still finish as FINISHED_LENGTH_CAPPED instead of spinning in the + scheduler with zero new tokens to schedule. + """ + from vllm.v1.core.sched.utils import check_stop + + # Reaching max_model_len below min_tokens must stop the request. + sampling_params = SamplingParams( + ignore_eos=True, + max_tokens=500, + min_tokens=500, + ) + sampling_params.update_from_generation_config({}, EOS_TOKEN_ID) + request = Request( + request_id="0", + prompt_token_ids=[0, 1, 2], + sampling_params=sampling_params, + pooling_params=None, + ) + # 7 output tokens: num_tokens == 10 == max_model_len, min_tokens unmet. + request.append_output_token_ids([10, 11, 12, 13, 14, 15, 16]) + + result = check_stop(request, max_model_len=10) + assert result is True, "Should stop at max_model_len even below min_tokens" + assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED + + # Reaching max_tokens below min_tokens must stop the request. This can + # happen when max_tokens is clamped to the context window after + # SamplingParams validation (e.g. max_tokens=None on a long prompt). + sampling_params_capped = SamplingParams( + ignore_eos=True, + max_tokens=500, + min_tokens=500, + ) + sampling_params_capped.max_tokens = 3 + sampling_params_capped.update_from_generation_config({}, EOS_TOKEN_ID) + request_capped = Request( + request_id="1", + prompt_token_ids=[0, 1, 2], + sampling_params=sampling_params_capped, + pooling_params=None, + ) + request_capped.append_output_token_ids([10, 11, 12]) + + result = check_stop(request_capped, max_model_len=100) + assert result is True, "Should stop at max_tokens even below min_tokens" + assert request_capped.status == RequestStatus.FINISHED_LENGTH_CAPPED + + @pytest.mark.parametrize( "enable_prefix_caching, prompt_logprobs", [ diff --git a/tests/v1/e2e/general/test_min_tokens.py b/tests/v1/e2e/general/test_min_tokens.py index c5b6341fb9..0100a38ba1 100644 --- a/tests/v1/e2e/general/test_min_tokens.py +++ b/tests/v1/e2e/general/test_min_tokens.py @@ -13,11 +13,14 @@ Covers: 5) Multiple stop conditions """ +import json + import pytest from vllm import LLM, SamplingParams from vllm.exceptions import VLLMValidationError from vllm.outputs import RequestOutput +from vllm.sampling_params import StructuredOutputsParams # Test configuration TEST_MODEL = "facebook/opt-125m" # Small model for fast CI execution @@ -168,6 +171,27 @@ def llm_v1(): return llm +def test_min_tokens_with_terminal_structured_output(llm_v1: LLM): + """A grammar whose only legal continuation is EOS must stop under + min_tokens instead of sampling outside the grammar.""" + params = SamplingParams( + temperature=GREEDY, + min_tokens=20, + max_tokens=40, + structured_outputs=StructuredOutputsParams( + json=json.dumps({"type": "boolean"}) + ), + ) + + output = llm_v1.generate("Answer:", params)[0].outputs[0] + eos_token_id = llm_v1.get_tokenizer().eos_token_id + + assert output.finish_reason == "stop" + assert isinstance(json.loads(output.text), bool) + assert output.token_ids[-1] == eos_token_id + assert len(output.token_ids) < params.min_tokens + + def get_token_count(output: RequestOutput) -> int: """Extract token count from LLM output""" if not output.outputs: diff --git a/tests/v1/logits_processors/test_correctness.py b/tests/v1/logits_processors/test_correctness.py index a51157557e..7e60eeb2e3 100644 --- a/tests/v1/logits_processors/test_correctness.py +++ b/tests/v1/logits_processors/test_correctness.py @@ -20,7 +20,11 @@ from tests.v1.sample.utils import ( ) from vllm.config import VllmConfig from vllm.platforms import current_platform -from vllm.sampling_params import SamplingParams, validate_thinking_token_budget +from vllm.sampling_params import ( + SamplingParams, + StructuredOutputsParams, + validate_thinking_token_budget, +) from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.sample.logits_processor import ( BatchUpdate, @@ -371,7 +375,7 @@ def _min_tokens_validate( # Validate min-token logits processor state if min_tok: - (_, out_tok, all_stop_token_ids) = min_tok + (_, out_tok, all_stop_token_ids, _) = min_tok num_out_tokens = len(out_tok) if num_out_tokens != ref_num_out_tokens: _raise_error_invalid( @@ -448,6 +452,113 @@ def _min_tokens_validate( ) +def _make_min_tokens_processor() -> MinTokensLogitsProcessor: + return MinTokensLogitsProcessor(VllmConfig(), torch.device("cpu"), False) + + +# Constructing the processor allocates pinned host memory, which creates a +# CUDA driver context even though the tensors stay on CPU. Run these in their +# own process so the pytest process can still fork for `test_logitsprocs`. +@create_new_process_for_each_test() +def test_min_tokens_keeps_all_masked_behavior_without_structured_output(): + processor = _make_min_tokens_processor() + processor.update_state( + BatchUpdate( + batch_size=1, + removed=(), + added=[ + ( + 0, + SamplingParams(min_tokens=2, stop_token_ids=[0]), + None, + [], + ) + ], + moved=(), + ) + ) + + logits = torch.full((1, 3), -float("inf")) + logits[0, 0] = 1.0 + + processor.apply(logits) + + assert torch.isneginf(logits).all() + + +@create_new_process_for_each_test() +def test_min_tokens_restores_all_masked_structured_output_stop_token(): + processor = _make_min_tokens_processor() + processor.update_state( + BatchUpdate( + batch_size=1, + removed=(), + added=[ + ( + 0, + SamplingParams( + min_tokens=2, + stop_token_ids=[0], + structured_outputs=StructuredOutputsParams( + json={"type": "boolean"} + ), + ), + None, + [], + ) + ], + moved=(), + ) + ) + + logits = torch.full((1, 3), -float("inf")) + logits[0, 0] = 1.0 + + processor.apply(logits) + + assert logits[0, 0] == 1.0 + assert torch.isneginf(logits[0, 1:]).all() + + +@create_new_process_for_each_test() +def test_min_tokens_restores_all_masked_structured_output_stop_token_spec_decode(): + processor = _make_min_tokens_processor() + processor.update_state( + BatchUpdate( + batch_size=1, + removed=(), + added=[ + ( + 0, + SamplingParams( + min_tokens=2, + stop_token_ids=[0, 1], + structured_outputs=StructuredOutputsParams( + json={"type": "boolean"} + ), + ), + None, + [], + ) + ], + moved=(), + ) + ) + + logits = torch.full((2, 4), -float("inf")) + logits[0, 0] = 1.0 + logits[0, 1] = 2.0 + logits[1, 0] = 3.0 + + processor.apply_with_spec_decode(logits, [2]) + + assert logits[0, 0] == 1.0 + assert logits[0, 1] == 2.0 + assert logits[1, 0] == 3.0 + assert torch.isneginf(logits[1, 1:]).all() + assert torch.isneginf(logits[:, 2:]).all() + + def _thinking_budget_params(kwargs: dict) -> None: """Set SamplingParams kwargs for thinking token budget tests""" kwargs["thinking_token_budget"] = THINKING_TOKEN_BUDGET diff --git a/tests/v1/worker/test_gpu_logit_bias.py b/tests/v1/worker/test_gpu_logit_bias.py new file mode 100644 index 0000000000..3bf62380bb --- /dev/null +++ b/tests/v1/worker/test_gpu_logit_bias.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip( + "CUDA required for Model Runner V2 logit bias tests", + allow_module_level=True, + ) + +from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState + +DEVICE = torch.device("cuda") +VOCAB_SIZE = 128 +STOP_TOKEN = 3 +OTHER_TOKEN = 4 +PROMPT_LEN = 2 +MIN_TOKENS = 10 +# min_len is PROMPT_LEN + MIN_TOKENS, so this position is below the minimum +# and the kernel masks stop tokens. +POS = 5 + + +def _params(structured: bool) -> SamplingParams: + return SamplingParams( + min_tokens=MIN_TOKENS, + stop_token_ids=[STOP_TOKEN], + structured_outputs=( + StructuredOutputsParams(json={"type": "boolean"}) if structured else None + ), + ) + + +def _only_stop_token_left(num_rows: int) -> torch.Tensor: + """Logits as a terminal grammar leaves them: every token masked but one.""" + logits = torch.full((num_rows, VOCAB_SIZE), -float("inf"), device=DEVICE) + logits[:, STOP_TOKEN] = 1.0 + return logits + + +def _apply(logits: torch.Tensor, structured: list[bool]) -> torch.Tensor: + """Run the fused bias kernel with one logits row per request.""" + state = LogitBiasState(max_num_reqs=4, device=DEVICE) + for req_idx, is_structured in enumerate(structured): + state.add_request(req_idx, PROMPT_LEN, _params(is_structured)) + state.apply_staged_writes() + + num_reqs = len(structured) + state.apply_logit_bias( + logits, + torch.arange(num_reqs, dtype=torch.int32, device=DEVICE), + np.arange(num_reqs, dtype=np.intp), + torch.full((num_reqs,), POS, dtype=torch.int32, device=DEVICE), + ) + return logits.cpu() + + +def test_v2_min_tokens_restores_stop_token_when_row_would_be_empty(): + """Structured output plus min_tokens must not leave an all -inf row.""" + out = _apply(_only_stop_token_left(1), structured=[True]) + assert out[0, STOP_TOKEN] == 1.0 + + +def test_v2_min_tokens_keeps_empty_row_without_structured_output(): + """Plain min_tokens requests keep the pre-existing masking behavior.""" + out = _apply(_only_stop_token_left(1), structured=[False]) + assert torch.isneginf(out).all() + + +def test_v2_min_tokens_keeps_masking_when_other_candidates_remain(): + """The fallback must not weaken min_tokens while a non-stop token is legal.""" + logits = _only_stop_token_left(1) + logits[0, OTHER_TOKEN] = 2.0 + + out = _apply(logits, structured=[True]) + + assert torch.isneginf(out[0, STOP_TOKEN]) + assert out[0, OTHER_TOKEN] == 2.0 + + +def test_v2_min_tokens_mixed_batch_gates_restore_per_request(): + """CHECK_ALL_MASKED_ROWS is a batch-wide constexpr, so a mixed batch must + still consult the per-request flag.""" + out = _apply(_only_stop_token_left(2), structured=[True, False]) + + assert out[0, STOP_TOKEN] == 1.0 + assert torch.isneginf(out[1]).all() diff --git a/vllm/v1/core/sched/utils.py b/vllm/v1/core/sched/utils.py index c7cb6b9436..31fa840e77 100644 --- a/vllm/v1/core/sched/utils.py +++ b/vllm/v1/core/sched/utils.py @@ -97,9 +97,6 @@ def check_stop(request: Request, max_model_len: int) -> bool: sampling_params = request.sampling_params assert sampling_params is not None - if request.num_output_tokens < sampling_params.min_tokens: - return False - last_token_id = request.output_token_ids[-1] if last_token_id == sampling_params.eos_token_id: request.status = RequestStatus.FINISHED_STOPPED @@ -109,6 +106,7 @@ def check_stop(request: Request, max_model_len: int) -> bool: request.status = RequestStatus.FINISHED_STOPPED request.stop_reason = last_token_id return True + if ( request.num_tokens >= max_model_len or request.num_output_tokens >= request.max_tokens @@ -116,6 +114,14 @@ def check_stop(request: Request, max_model_len: int) -> bool: request.status = RequestStatus.FINISHED_LENGTH_CAPPED return True + # Note(arpera): + # Order of checks is important for min_tokens + # If you decide to change the order, first look at these PRs please: + # 47489, 49521, and 51299 + # These PRs show some reasoning about the order of checks + if request.num_output_tokens < sampling_params.min_tokens: + return False + repetition_detection = sampling_params.repetition_detection if repetition_detection is not None and ( check_sequence_repetition( diff --git a/vllm/v1/sample/logits_processor/builtin.py b/vllm/v1/sample/logits_processor/builtin.py index d7c9444380..45b5c47711 100644 --- a/vllm/v1/sample/logits_processor/builtin.py +++ b/vllm/v1/sample/logits_processor/builtin.py @@ -166,16 +166,22 @@ class MinTokensLogitsProcessor(LogitsProcessor): def __init__( self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool ): - # index -> (min_toks, output_token_ids, stop_token_ids) + # index -> (min_toks, output_token_ids, stop_token_ids, uses_structured_output) self.device = device - self.min_toks: dict[int, tuple[int, Sequence[int], set[int]]] = {} + self.min_toks: dict[int, tuple[int, Sequence[int], set[int], bool]] = {} - # (req_idx_tensor,eos_tok_id_tensor) + # (req_idx_tensor, stop_token_id_tensor) self.logits_slice: tuple[torch.Tensor, torch.Tensor] = ( self._device_tensor([], torch.int32), self._device_tensor([], torch.int32), ) + # Structured-output stop-token logits that may need to be restored. + self.restore_logits_slice: tuple[torch.Tensor, torch.Tensor] = ( + self._device_tensor([], torch.int32), + self._device_tensor([], torch.int32), + ) + self.neg_inf_tensor = torch.tensor( -float("inf"), dtype=torch.float32, device=self.device ) @@ -188,11 +194,16 @@ class MinTokensLogitsProcessor(LogitsProcessor): @staticmethod def add_request( params: SamplingParams, _: list[int] | None, output_tok_ids: list[int] - ) -> tuple[int, Sequence[int], set[int]] | None: + ) -> tuple[int, Sequence[int], set[int], bool] | None: min_tokens = params.min_tokens if not min_tokens or len(output_tok_ids) >= min_tokens: return None - return min_tokens, output_tok_ids, params.all_stop_token_ids + return ( + min_tokens, + output_tok_ids, + params.all_stop_token_ids, + params.structured_outputs is not None, + ) def update_state(self, batch_update: BatchUpdate | None): needs_update = process_dict_updates( @@ -202,7 +213,7 @@ class MinTokensLogitsProcessor(LogitsProcessor): # Check for any requests that have attained their min tokens. to_remove = tuple( index - for index, (min_toks, out_tok_ids, _) in self.min_toks.items() + for index, (min_toks, out_tok_ids, _, _) in self.min_toks.items() if len(out_tok_ids) >= min_toks ) if to_remove: @@ -214,22 +225,66 @@ class MinTokensLogitsProcessor(LogitsProcessor): if needs_update: reqs: list[int] = [] tok_ids: list[int] = [] - for req, (_, _, stop_tok_ids) in self.min_toks.items(): + restore_reqs: list[int] = [] + restore_tok_ids: list[int] = [] + for req, ( + _, + _, + stop_tok_ids, + uses_structured_output, + ) in self.min_toks.items(): reqs.extend([req] * len(stop_tok_ids)) tok_ids.extend(stop_tok_ids) + if uses_structured_output: + restore_reqs.extend([req] * len(stop_tok_ids)) + restore_tok_ids.extend(stop_tok_ids) self.logits_slice = ( self._device_tensor(reqs, torch.int32), self._device_tensor(tok_ids, torch.int32), ) + self.restore_logits_slice = ( + self._device_tensor(restore_reqs, torch.int32), + self._device_tensor(restore_tok_ids, torch.int32), + ) def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor: return async_tensor_h2d(data, device=self.device, dtype=dtype) + def _mask_stop_token_logits( + self, + logits: torch.Tensor, + logits_slice: tuple[torch.Tensor, torch.Tensor], + restore_logits_slice: tuple[torch.Tensor, torch.Tensor], + ) -> None: + restore_rows, restore_toks = restore_logits_slice + # Advanced indexing already returns a copy, so this survives the + # in-place masking below. + stop_logits = logits[restore_logits_slice] + logits.index_put_(logits_slice, self.neg_inf_tensor) + + # Stop-token entries for the same row are emitted together, so this + # checks each row once without sorting the index tensor. + unique_rows, inverse_indices = torch.unique_consecutive( + restore_rows, return_inverse=True + ) + row_needs_restore = torch.isneginf(logits[unique_rows]).all(dim=-1) + restore_mask = row_needs_restore[inverse_indices] & torch.isfinite(stop_logits) + restore_slice = ( + restore_rows[restore_mask], + restore_toks[restore_mask], + ) + logits.index_put_(restore_slice, stop_logits[restore_mask]) + def apply(self, logits: torch.Tensor) -> torch.Tensor: if self.min_toks: - # Inhibit EOS token for requests which have not reached min length - logits.index_put_(self.logits_slice, self.neg_inf_tensor) + # Inhibit stop tokens for requests which have not reached min length. + if self.restore_logits_slice[0].numel() == 0: + logits.index_put_(self.logits_slice, self.neg_inf_tensor) + else: + self._mask_stop_token_logits( + logits, self.logits_slice, self.restore_logits_slice + ) return logits def apply_with_spec_decode( @@ -238,10 +293,13 @@ class MinTokensLogitsProcessor(LogitsProcessor): num_draft_tokens: list[int], ) -> torch.Tensor: """Spec-decode version of apply(). - Priority: ``min_tokens`` > ``stop_token_ids`` / EOS. + + Stop tokens stay masked unless masking them would leave every logit + in a row at -inf. + Example: ``num_draft_tokens = [2, 3, 1]`` - → ``logits`` shape ``[6, V]``, ``cumsum = [0, 2, 5, 6]`` - → request 0 owns rows 0‑1, request 1 rows 2‑4, request 2 row 5. + -> ``logits`` shape ``[6, V]``, ``cumsum = [0, 2, 5, 6]`` + -> request 0 owns rows 0-1, request 1 rows 2-4, request 2 row 5. """ if not self.min_toks: return logits @@ -250,8 +308,19 @@ class MinTokensLogitsProcessor(LogitsProcessor): cumsum = np.concatenate([[0], np.cumsum(num_draft_arr)]) entries = [ - (req_idx, min_tok, len(out_tok_ids), list(stop_tok_ids)) - for req_idx, (min_tok, out_tok_ids, stop_tok_ids) in self.min_toks.items() + ( + req_idx, + min_tok, + len(out_tok_ids), + list(stop_tok_ids), + uses_structured_output, + ) + for req_idx, ( + min_tok, + out_tok_ids, + stop_tok_ids, + uses_structured_output, + ) in self.min_toks.items() if stop_tok_ids ] @@ -260,8 +329,10 @@ class MinTokensLogitsProcessor(LogitsProcessor): all_rows: list[np.ndarray] = [] # row indices to mask all_toks: list[np.ndarray] = [] # stop-token ids at those rows + restore_rows: list[np.ndarray] = [] + restore_toks: list[np.ndarray] = [] - for req_idx, min_tok, current_len, stop_toks in entries: + for req_idx, min_tok, current_len, stop_toks, uses_structured_output in entries: remaining = min_tok - current_len # How many leading draft positions still need stop-token masking. n_mask = int(min(max(remaining, 0), num_draft_arr[req_idx])) @@ -270,8 +341,13 @@ class MinTokensLogitsProcessor(LogitsProcessor): offset = cumsum[req_idx] row_indices = np.arange(offset, offset + n_mask, dtype=np.int64) n_stop = len(stop_toks) - all_rows.append(np.repeat(row_indices, n_stop)) - all_toks.append(np.tile(stop_toks, n_mask)) + rows = np.repeat(row_indices, n_stop) + toks = np.tile(stop_toks, n_mask) + all_rows.append(rows) + all_toks.append(toks) + if uses_structured_output: + restore_rows.append(rows) + restore_toks.append(toks) if all_rows: rows_arr = np.concatenate(all_rows) @@ -281,7 +357,16 @@ class MinTokensLogitsProcessor(LogitsProcessor): async_tensor_h2d(rows_arr, device=self.device), async_tensor_h2d(toks_arr, device=self.device), ) - logits.index_put_(logits_slice, self.neg_inf_tensor) + if restore_rows: + restore_rows_arr = np.concatenate(restore_rows) + restore_toks_arr = np.concatenate(restore_toks) + restore_logits_slice = ( + async_tensor_h2d(restore_rows_arr, device=self.device), + async_tensor_h2d(restore_toks_arr, device=self.device), + ) + self._mask_stop_token_logits(logits, logits_slice, restore_logits_slice) + else: + logits.index_put_(logits_slice, self.neg_inf_tensor) return logits diff --git a/vllm/v1/worker/gpu/sample/logit_bias.py b/vllm/v1/worker/gpu/sample/logit_bias.py index 6c95ed7aac..846edbe1e7 100644 --- a/vllm/v1/worker/gpu/sample/logit_bias.py +++ b/vllm/v1/worker/gpu/sample/logit_bias.py @@ -40,6 +40,9 @@ class LogitBiasState: # Min tokens. self.min_lens = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) self.num_stop_token_ids = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) + self.restore_when_all_masked = UvaBackedTensor( + self.max_num_reqs, dtype=torch.int32 + ) self.stop_token_ids = StagedWriteTensor( (self.max_num_reqs, MAX_NUM_STOP_TOKEN_IDS), dtype=torch.int32, @@ -100,9 +103,13 @@ class LogitBiasState: ) self.num_stop_token_ids.np[req_idx] = num_stop_token_ids self.stop_token_ids.stage_write(req_idx, 0, stop_token_ids) + self.restore_when_all_masked.np[req_idx] = int( + sampling_params.structured_outputs is not None + ) use_logit_bias = True else: self.num_stop_token_ids.np[req_idx] = 0 + self.restore_when_all_masked.np[req_idx] = 0 self.use_logit_bias[req_idx] = use_logit_bias @@ -116,6 +123,7 @@ class LogitBiasState: self.min_lens.copy_to_uva() self.num_stop_token_ids.copy_to_uva() + self.restore_when_all_masked.copy_to_uva() self.stop_token_ids.apply_write() def apply_logit_bias( @@ -129,6 +137,10 @@ class LogitBiasState: # No request uses logit bias. Skip the kernel launch. return + enable_stop_token_restore = bool( + np.any(self.restore_when_all_masked.np[idx_mapping_np]) + ) + apply_logit_bias( logits, expanded_idx_mapping, @@ -140,7 +152,9 @@ class LogitBiasState: self.logit_bias.gpu, self.min_lens.gpu, self.num_stop_token_ids.gpu, + self.restore_when_all_masked.gpu, self.stop_token_ids.gpu, + enable_stop_token_restore, ) @@ -164,10 +178,12 @@ def _bias_kernel( pos_ptr, min_lens_ptr, num_stop_token_ids_ptr, + restore_when_all_masked_ptr, stop_token_ids_ptr, stop_token_ids_stride, BLOCK_SIZE: tl.constexpr, LOGITS_BLOCK_SIZE: tl.constexpr, + CHECK_ALL_MASKED_ROWS: tl.constexpr, ): token_idx = tl.program_id(0).to(tl.int64) req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) @@ -232,11 +248,59 @@ def _bias_kernel( stop_token_ids_ptr + req_state_idx * stop_token_ids_stride + block, mask=mask, ) - tl.store( - logits_ptr + token_idx * logits_stride + stop_token_ids, - -float("inf"), - mask=mask, - ) + if CHECK_ALL_MASKED_ROWS: + should_restore_stop_logits = tl.load( + restore_when_all_masked_ptr + req_state_idx + ) + if should_restore_stop_logits: + stop_logits = tl.load( + logits_ptr + token_idx * logits_stride + stop_token_ids, + mask=mask, + other=-float("inf"), + ) + + # Save must read original logits before the -inf overwrite. + tl.debug_barrier() + + tl.store( + logits_ptr + token_idx * logits_stride + stop_token_ids, + -float("inf"), + mask=mask, + ) + + # The -inf overwrite must be visible to the row scan below. + tl.debug_barrier() + + row_max = tl.full((), -float("inf"), tl.float32) + for i in range(0, vocab_size, LOGITS_BLOCK_SIZE): + offset = i + tl.arange(0, LOGITS_BLOCK_SIZE) + logits = tl.load( + logits_ptr + token_idx * logits_stride + offset, + mask=offset < vocab_size, + other=-float("inf"), + ) + row_max = tl.maximum(row_max, tl.max(logits, axis=0)) + + if row_max == -float("inf"): + tl.store( + logits_ptr + token_idx * logits_stride + stop_token_ids, + stop_logits, + mask=mask + & (stop_logits > -float("inf")) + & (stop_logits < float("inf")), + ) + else: + tl.store( + logits_ptr + token_idx * logits_stride + stop_token_ids, + -float("inf"), + mask=mask, + ) + else: + tl.store( + logits_ptr + token_idx * logits_stride + stop_token_ids, + -float("inf"), + mask=mask, + ) def apply_logit_bias( @@ -250,7 +314,9 @@ def apply_logit_bias( logit_bias: torch.Tensor, min_lens: torch.Tensor, num_stop_token_ids: torch.Tensor, + restore_when_all_masked: torch.Tensor, stop_token_ids: torch.Tensor, + check_all_masked_rows: bool = False, ) -> None: num_tokens, vocab_size = logits.shape BLOCK_SIZE = triton.next_power_of_2( @@ -277,8 +343,10 @@ def apply_logit_bias( pos, min_lens, num_stop_token_ids, + restore_when_all_masked, stop_token_ids, stop_token_ids.stride(0), BLOCK_SIZE=BLOCK_SIZE, LOGITS_BLOCK_SIZE=LOGITS_BLOCK_SIZE, + CHECK_ALL_MASKED_ROWS=check_all_masked_rows, )