[Structured Output] Let terminal grammars stop under min_tokens (II) (#54218)

Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>

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"
 
-    # Test case 2: 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"
-
-    # Test case 5: min_tokens met, should stop on stop token
-    request_stop.append_output_token_ids(
-        [10, 11, 12, 13, 14, 42]
-    )  # 6 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",
     [

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:

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

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()

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(

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
 

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,
     )