diff --git a/tests/sampling/test_prompt_logprobs_security.py b/tests/sampling/test_prompt_logprobs_security.py new file mode 100644 index 000000000000..6efc0918ba67 --- /dev/null +++ b/tests/sampling/test_prompt_logprobs_security.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Security tests for prompt_logprobs resource-bound validation. + +Replaces the original sentinel-only rejection (-1) with a resource bound +(``VLLM_MAX_PROMPT_LOGPROBS``). The bound closes the bypass where an +operator-configured ``max_logprobs=-1`` makes the allowed maximum equal to +the vocabulary size, so ``prompt_logprobs=vocab_size`` (or any large value) +allocates the identical full-vocabulary tensor that ``-1`` would have. +""" + +import os + +from unittest.mock import MagicMock + +import pytest + +from vllm.sampling_params import SamplingParams + + +def _mock_model_config(vocab_size=154880, max_logprobs=-1): + """Create a mock ModelConfig with a large vocab and max_logprobs=-1.""" + config = MagicMock() + config.get_vocab_size.return_value = vocab_size + config.max_logprobs = max_logprobs + config.get_hidden_size.return_value = 4096 + config.is_encoder_decoder = False + config.is_diffusion = False + return config + + +# --------------------------------------------------------------------------- +# Original five tests (preserved, adapted to the cap-based approach) +# --------------------------------------------------------------------------- + + +def test_prompt_logprobs_minus_one_rejected(): + """B9: prompt_logprobs=-1 must be rejected regardless of max_logprobs. + + With the cap-based approach, -1 resolves to vocab_size (154,880 for + GLM-5.2) which exceeds the default cap of 20. + """ + from vllm.exceptions import VLLMValidationError + + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(prompt_logprobs=-1) + + with pytest.raises(VLLMValidationError, match="greater than max allowed"): + sp._validate_logprobs(model_config) + + +def test_prompt_logprobs_positive_still_works(): + """A concrete prompt_logprobs value within the cap should pass.""" + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(prompt_logprobs=20) + sp._validate_logprobs(model_config) + + +def test_prompt_logprobs_zero_still_works(): + """prompt_logprobs=0 (disabled) should not trigger the rejection.""" + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(prompt_logprobs=0) + sp._validate_logprobs(model_config) + + +def test_prompt_logprobs_none_still_works(): + """prompt_logprobs=None (unset) should not trigger the rejection.""" + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams() + sp._validate_logprobs(model_config) + + +def test_prompt_logprobs_exceeds_max_still_rejected(): + """prompt_logprobs > max_logprobs is still rejected by the existing + check (confirms the existing behavior is preserved).""" + from vllm.exceptions import VLLMValidationError + + model_config = _mock_model_config(vocab_size=154880, max_logprobs=20) + sp = SamplingParams(prompt_logprobs=100) + + with pytest.raises(VLLMValidationError, match="greater than max allowed"): + sp._validate_logprobs(model_config) + + +# --------------------------------------------------------------------------- +# New tests: resource-bound bypass coverage (B9) +# --------------------------------------------------------------------------- + + +def test_prompt_logprobs_vocab_size_rejected(): + """B9: prompt_logprobs=vocab_size must be rejected even when + max_logprobs=-1 (the bypass the reviewers found).""" + from vllm.exceptions import VLLMValidationError + + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(prompt_logprobs=154880) + + with pytest.raises(VLLMValidationError, match="greater than max allowed"): + sp._validate_logprobs(model_config) + + +def test_prompt_logprobs_vocab_size_minus_one_rejected(): + """B9: prompt_logprobs=vocab_size-1 must also be rejected; values just + below vocab_size are equally expensive.""" + from vllm.exceptions import VLLMValidationError + + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(prompt_logprobs=154879) + + with pytest.raises(VLLMValidationError, match="greater than max allowed"): + sp._validate_logprobs(model_config) + + +def test_prompt_logprobs_at_cap_accepted(): + """prompt_logprobs equal to the cap (VLLM_MAX_PROMPT_LOGPROBS, default 20) + must be accepted.""" + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(prompt_logprobs=20) + sp._validate_logprobs(model_config) + + +def test_prompt_logprobs_above_cap_rejected(): + """prompt_logprobs=cap+1 must be rejected.""" + from vllm.exceptions import VLLMValidationError + + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(prompt_logprobs=21) + + with pytest.raises(VLLMValidationError, match="greater than max allowed"): + sp._validate_logprobs(model_config) + + +def test_prompt_logprobs_bypass_with_max_logprobs_minus_one(): + """B9: the original bypass — max_logprobs=-1 makes the allowed maximum + vocab_size, so large positive values pass the > check. The cap must + still reject them.""" + from vllm.exceptions import VLLMValidationError + + # Simulate the exact bypass scenario: max_logprobs=-1 (unlimited) + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(prompt_logprobs=154880) + + with pytest.raises(VLLMValidationError, match="VLLM_MAX_PROMPT_LOGPROBS"): + sp._validate_logprobs(model_config) + + +# --------------------------------------------------------------------------- +# C5: same bound applied to sampling logprobs +# --------------------------------------------------------------------------- + + +def test_sample_logprobs_vocab_size_rejected(): + """C5: logprobs=vocab_size (sampling, not prompt) must also be capped.""" + from vllm.exceptions import VLLMValidationError + + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(logprobs=154880) + + with pytest.raises(VLLMValidationError, match="VLLM_MAX_LOGPROBS"): + sp._validate_logprobs(model_config) + + +def test_sample_logprobs_minus_one_rejected(): + """C5: logprobs=-1 resolves to vocab_size and must be capped.""" + from vllm.exceptions import VLLMValidationError + + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(logprobs=-1) + + with pytest.raises(VLLMValidationError, match="greater than max allowed"): + sp._validate_logprobs(model_config) + + +def test_sample_logprobs_at_cap_accepted(): + """C5: logprobs equal to the cap must be accepted.""" + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(logprobs=20) + sp._validate_logprobs(model_config) + + +# --------------------------------------------------------------------------- +# Custom cap via env var +# --------------------------------------------------------------------------- + + +def test_prompt_logprobs_custom_cap(monkeypatch): + """An operator who raises VLLM_MAX_PROMPT_LOGPROBS should be able to + request more than the default 20.""" + monkeypatch.setenv("VLLM_MAX_PROMPT_LOGPROBS", "100") + # Need to re-import to pick up the new env var value — envs.py caches + # via lambda, so the value is read fresh each time. + model_config = _mock_model_config(vocab_size=154880, max_logprobs=-1) + sp = SamplingParams(prompt_logprobs=100) + sp._validate_logprobs(model_config) + + # But 101 should still fail + from vllm.exceptions import VLLMValidationError + + sp2 = SamplingParams(prompt_logprobs=101) + with pytest.raises(VLLMValidationError, match="greater than max allowed"): + sp2._validate_logprobs(model_config) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index b70114446f06..4befcdbc046e 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -46,6 +46,7 @@ StructuredOutputsParams, ThinkingTokenBudget, ) +import vllm.envs as envs from vllm.utils import random_uuid logger = init_logger(__name__) @@ -794,14 +795,32 @@ def check_logprobs(cls, data): if prompt_logprobs < 0 and prompt_logprobs != -1: raise VLLMValidationError( - "`prompt_logprobs` must be a positive value or -1.", + "`prompt_logprobs` must be non-negative or -1 (all " + "logprobs, subject to VLLM_MAX_PROMPT_LOGPROBS cap).", + parameter="prompt_logprobs", + value=prompt_logprobs, + ) + if prompt_logprobs > 0 and prompt_logprobs > envs.VLLM_MAX_PROMPT_LOGPROBS: + raise VLLMValidationError( + f"`prompt_logprobs`={prompt_logprobs} exceeds the " + f"maximum allowed: {envs.VLLM_MAX_PROMPT_LOGPROBS} " + f"(VLLM_MAX_PROMPT_LOGPROBS).", parameter="prompt_logprobs", value=prompt_logprobs, ) if (top_logprobs := data.get("top_logprobs")) is not None: if top_logprobs < 0 and top_logprobs != -1: raise VLLMValidationError( - "`top_logprobs` must be a positive value or -1.", + "`top_logprobs` must be non-negative or -1 (all " + "logprobs, subject to VLLM_MAX_LOGPROBS cap).", + parameter="top_logprobs", + value=top_logprobs, + ) + if top_logprobs > 0 and top_logprobs > envs.VLLM_MAX_LOGPROBS: + raise VLLMValidationError( + f"`top_logprobs`={top_logprobs} exceeds the " + f"maximum allowed: {envs.VLLM_MAX_LOGPROBS} " + f"(VLLM_MAX_LOGPROBS).", parameter="top_logprobs", value=top_logprobs, ) diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 1784d0f5364c..b5f1b5cad663 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -505,16 +505,34 @@ def check_logprobs(cls, data): if prompt_logprobs < 0 and prompt_logprobs != -1: raise VLLMValidationError( - "`prompt_logprobs` must be a positive value or -1.", + "`prompt_logprobs` must be non-negative or -1 (all " + "logprobs, subject to VLLM_MAX_PROMPT_LOGPROBS cap).", parameter="prompt_logprobs", value=prompt_logprobs, ) - if (logprobs := data.get("logprobs")) is not None and logprobs < 0: - raise VLLMValidationError( - "`logprobs` must be a positive value.", - parameter="logprobs", - value=logprobs, - ) + if prompt_logprobs > 0 and prompt_logprobs > envs.VLLM_MAX_PROMPT_LOGPROBS: + raise VLLMValidationError( + f"`prompt_logprobs`={prompt_logprobs} exceeds the " + f"maximum allowed: {envs.VLLM_MAX_PROMPT_LOGPROBS} " + f"(VLLM_MAX_PROMPT_LOGPROBS).", + parameter="prompt_logprobs", + value=prompt_logprobs, + ) + if (logprobs := data.get("logprobs")) is not None: + if logprobs < 0 and logprobs != -1: + raise VLLMValidationError( + "`logprobs` must be non-negative or -1 (all logprobs, " + "subject to VLLM_MAX_LOGPROBS cap).", + parameter="logprobs", + value=logprobs, + ) + if logprobs > 0 and logprobs > envs.VLLM_MAX_LOGPROBS: + raise VLLMValidationError( + f"`logprobs`={logprobs} exceeds the maximum allowed: " + f"{envs.VLLM_MAX_LOGPROBS} (VLLM_MAX_LOGPROBS).", + parameter="logprobs", + value=logprobs, + ) return data diff --git a/vllm/envs.py b/vllm/envs.py index 7c05ab3d3a7c..a938e9048387 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -191,6 +191,8 @@ VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 VLLM_PROMPT_LOGPROBS_CHUNK_SIZE: int = 1024 + VLLM_MAX_PROMPT_LOGPROBS: int = 20 + VLLM_MAX_LOGPROBS: int = 20 VLLM_MLA_DISABLE: bool = False VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH: bool = False VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH_WINDOW: int = 8 @@ -1573,6 +1575,20 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_PROMPT_LOGPROBS_CHUNK_SIZE": lambda: int( os.getenv("VLLM_PROMPT_LOGPROBS_CHUNK_SIZE", "1024") ), + # Maximum number of prompt logprobs that may be requested per prompt + # token. Caps the [num_prompt_tokens, prompt_logprobs] allocation to + # prevent OOM (upstream vllm-project/vllm#14239). The sentinel -1 (all + # vocab logprobs) is resolved to vocab_size before the cap is applied, + # so it is effectively rejected for any real vocabulary. + "VLLM_MAX_PROMPT_LOGPROBS": lambda: int( + os.getenv("VLLM_MAX_PROMPT_LOGPROBS", "20") + ), + # Maximum number of sample (output) logprobs per token. Same OOM + # rationale as VLLM_MAX_PROMPT_LOGPROBS but for the sampling logprobs + # path (logprobs / top_logprobs). + "VLLM_MAX_LOGPROBS": lambda: int( + os.getenv("VLLM_MAX_LOGPROBS", "20") + ), # If set, vLLM will disable the MLA attention optimizations. "VLLM_MLA_DISABLE": lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))), # Physically shorten DSpark's next draft block from the historical diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 08580f6e8f67..fd677120cbc0 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -271,10 +271,12 @@ class SamplingParams( likely tokens, as well as the chosen tokens. Note that the implementation follows the OpenAI API: The API will always return the log probability of the sampled token, so there may be up to `logprobs+1` elements in the - response. When set to -1, return all `vocab_size` log probabilities.""" + response. When set to -1, returns all `vocab_size` log probabilities, + subject to the ``VLLM_MAX_LOGPROBS`` cap (default 20).""" prompt_logprobs: int | None = None """Number of log probabilities to return per prompt token. - When set to -1, return all `vocab_size` log probabilities.""" + When set to -1, returns all `vocab_size` log probabilities, subject to + the ``VLLM_MAX_PROMPT_LOGPROBS`` cap (default 20).""" logprob_token_ids: list[int] | None = None """Specific token IDs to return logprobs for. More efficient than logprobs=-1 when you only need logprobs for a small set of tokens. @@ -587,7 +589,8 @@ def _verify_args(self) -> None: ) if self.logprobs is not None and self.logprobs != -1 and self.logprobs < 0: raise VLLMValidationError( - f"logprobs must be non-negative or -1, got {self.logprobs}.", + f"logprobs must be non-negative or -1 (all logprobs, " + f"capped by VLLM_MAX_LOGPROBS), got {self.logprobs}.", parameter="logprobs", value=self.logprobs, ) @@ -597,7 +600,8 @@ def _verify_args(self) -> None: and self.prompt_logprobs < 0 ): raise VLLMValidationError( - f"prompt_logprobs must be non-negative or -1, got " + f"prompt_logprobs must be non-negative or -1 (all logprobs, " + f"capped by VLLM_MAX_PROMPT_LOGPROBS), got " f"{self.prompt_logprobs}.", parameter="prompt_logprobs", value=self.prompt_logprobs, @@ -761,10 +765,19 @@ def _validate_logprobs(self, model_config: ModelConfig) -> None: if num_logprobs := self.logprobs: if num_logprobs == -1: num_logprobs = model_config.get_vocab_size() - if num_logprobs > max_logprobs: + # Hard cap from env var, immune to max_logprobs=-1 configuration. + # Without this, an operator who sets max_logprobs=-1 (unlimited) + # makes the allowed maximum the full vocabulary size, so + # logprobs=vocab_size passes the > check and allocates the + # identical full-vocabulary tensor that -1 would have. + cap = envs.VLLM_MAX_LOGPROBS + effective_max = min(max_logprobs, cap) + if num_logprobs > effective_max: raise VLLMValidationError( f"Requested sample logprobs of {num_logprobs}, " - f"which is greater than max allowed: {max_logprobs}", + f"which is greater than max allowed: {effective_max} " + f"(model max_logprobs={model_config.max_logprobs}, " + f"VLLM_MAX_LOGPROBS={cap})", parameter="logprobs", value=num_logprobs, ) @@ -802,14 +815,23 @@ def _validate_logprobs(self, model_config: ModelConfig) -> None: value=n, ) - # Validate prompt logprobs. if num_prompt_logprobs := self.prompt_logprobs: if num_prompt_logprobs == -1: + # Resolve the sentinel to its effective value before capping. num_prompt_logprobs = model_config.get_vocab_size() - if num_prompt_logprobs > max_logprobs: + # Hard cap from env var, immune to max_logprobs=-1 configuration. + # Without this, an operator who sets max_logprobs=-1 (unlimited) + # makes the allowed maximum the full vocabulary size, so + # prompt_logprobs=vocab_size passes the > check and allocates the + # identical full-vocabulary tensor that -1 would have. + cap = envs.VLLM_MAX_PROMPT_LOGPROBS + effective_max = min(max_logprobs, cap) + if num_prompt_logprobs > effective_max: raise VLLMValidationError( f"Requested prompt logprobs of {num_prompt_logprobs}, " - f"which is greater than max allowed: {max_logprobs}", + f"which is greater than max allowed: {effective_max} " + f"(model max_logprobs={model_config.max_logprobs}, " + f"VLLM_MAX_PROMPT_LOGPROBS={cap})", parameter="prompt_logprobs", value=num_prompt_logprobs, )