[watermark] compatibility validation (#56801)

Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai>
Signed-off-by: Simon Veitner <sveitner@redhat.com>
Co-authored-by: Simon Veitner <sveitner@redhat.com>
This commit is contained in:
Raphaël Rialland
2026-10-05 14:37:13 -04:00
committed by GitHub
co-authored by Simon Veitner
parent 88f3c601b3
commit c32495d58f
45 changed files with 1670 additions and 101 deletions
+4
View File
@@ -270,6 +270,9 @@ steps:
source_file_dependencies:
- .buildkite/test_areas/lm_eval.yaml
- tests/evals/gsm8k/
- tests/watermarking/test_compatibility.py
- tests/watermarking/test_engine_compatibility.py
- tests/watermarking/test_parallel_compatibility.py
- vllm/config/speculative.py
- vllm/config/vllm.py
- vllm/config/watermarking.py
@@ -282,6 +285,7 @@ steps:
autorun_on_main: true
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
- pytest -s -v watermarking/test_compatibility.py watermarking/test_engine_compatibility.py watermarking/test_parallel_compatibility.py
mirror:
amd:
label: ":amd: (MI355) LM Eval Watermark Feature Combination"
+42 -11
View File
@@ -18,8 +18,16 @@ vllm serve MODEL \
```
Watermarking is disabled when `--watermark-config` is omitted. Gumbel-max (see
`gumbel`) is the default algorithm within an enabled `WatermarkConfig`. When
watermarking is configured, it is enabled for requests by default.
`gumbel`) is the default algorithm within an enabled `WatermarkConfig`.
The request-level `watermarking` option has two effective states:
- Omitted (`None` in Python) or `true` requests watermarking.
- `false` explicitly opts out.
When a request asks for watermarking but the engine or generation mode cannot
apply it, vLLM emits a server-side warning once and runs the request without a
watermark.
Requests can opt out without changing the engine-level algorithm or key:
@@ -30,14 +38,12 @@ sampling_params = SamplingParams(watermarking=False)
```
The Python OpenAI-compatible APIs and the Rust chat/completions APIs accept the
same `watermarking: false` request field. The Rust token API accepts it in
`sampling_params`, and the Rust gRPC `GenerateRequest` accepts
the optional `watermarking` field. Omitting the field keeps watermarking enabled
when the engine has a watermark configuration.
same optional `watermarking` request field. The Rust token API accepts it in
`sampling_params`, and the Rust gRPC `GenerateRequest` accepts it directly.
Omitting the field requests watermarking by default.
Deployments that require watermarking must restrict this field to trusted
callers, or strip and validate it at the ingress boundary, so untrusted clients
cannot opt out.
Deployments that require watermarking must restrict explicit opt-out to trusted
callers, or strip and validate the field at the ingress boundary.
`context_width` controls how many prior tokens seed each watermark decision
and defaults to 4. Larger values make the watermark less robust to
@@ -63,6 +69,31 @@ the reference algorithms. `WatermarkDetector` consumes token IDs, so callers
remain responsible for using the tokenizer and watermark profile that match
generation.
## Interactions with generation features
- Greedy decoding (`temperature=0`) cannot apply a Gumbel watermark. The first
such request emits a server-side warning; it and subsequent greedy requests
use ordinary greedy sampling without a watermark. This includes
transcription, translation, and realtime transcription, which default to
greedy decoding. Transcription and translation can use a nonzero temperature
to apply the configured watermark.
- Beam search cannot apply a watermark. The first watermarked beam request emits
a server-side warning and runs without watermarking.
- Trace replay replaces the sampled token and therefore cannot apply a
watermark. The first watermarked trace replay request emits a server-side
warning and runs without watermarking.
- Parallel sampling (`n > 1`) is supported, but plain Gumbel candidates can be
identical while their contexts remain identical.
- Structured outputs and tool grammars are applied before watermark sampling.
Restrictive constraints can reduce the statistical evidence available to a
detector.
- Automatic language detection and generative scoring use unwatermarked
auxiliary generations. Language-detection tokens are not part of the returned
transcription, and generative scoring returns scores rather than generated
text.
- Speculative decoding has additional requirements and partial-watermarking
modes described in [Speculative decoding](#speculative-decoding).
## Speculative decoding
Watermarking requires speculative decoding to use probabilistic draft sampling,
@@ -245,7 +276,7 @@ watermarked output or to modify watermarked text so it is no longer detected.
- Watermarking is currently available only with Model Runner V2.
- Not all watermarking algorithms have native speculative-decoding support.
- Beam search expands candidates from model log probabilities and does not apply
Gumbel-max watermarking.
- Models that replace the vLLM sampler with a custom sampler cannot use
configured watermarking.
- Global custom logits processors are unavailable because Model Runner V2 does
not support them.
@@ -2117,3 +2117,30 @@ class TestMidConversationToolChanges:
ValidationError, match='only allowed in messages with role "system"'
):
request_cls(model="test-model", max_tokens=128, **fields)
class TestWatermarking:
def test_defaults_to_unspecified(self):
request = _make_request([{"role": "user", "content": "hi"}])
assert _convert(request).watermarking is None
def test_forwards_explicit_enable(self):
request = _make_request([{"role": "user", "content": "hi"}], watermarking=True)
assert _convert(request).watermarking is True
def test_forwards_the_opt_out(self):
request = _make_request([{"role": "user", "content": "hi"}], watermarking=False)
assert _convert(request).watermarking is False
@pytest.mark.parametrize("watermarking", [None, True, False])
def test_reaches_sampling_params(self, watermarking):
request = _make_request(
[{"role": "user", "content": "hi"}], watermarking=watermarking
)
params = _convert(request).to_sampling_params(16, {})
assert params.watermarking is watermarking
@@ -2510,3 +2510,24 @@ class TestCreateErrorResponse:
assert err.error.message == "oops"
assert err.error.code == 400
assert err.error.type == "bad_request"
def test_watermarking_defaults_to_unspecified():
assert _convert(_make_request()).watermarking is None
def test_watermarking_explicit_enable_is_forwarded():
assert _convert(_make_request(watermarking=True)).watermarking is True
def test_watermarking_opt_out_is_forwarded():
assert _convert(_make_request(watermarking=False)).watermarking is False
@pytest.mark.parametrize("watermarking", [None, True, False])
def test_watermarking_reaches_sampling_params(watermarking):
request = _make_request(watermarking=watermarking)
params = _convert(request).to_sampling_params(16, {})
assert params.watermarking is watermarking
@@ -12,6 +12,7 @@ Tests cover:
import math
from dataclasses import dataclass, field
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock
@@ -302,8 +303,10 @@ class TestGeneration:
mock_logprobs = {1234: -0.5, 5678: -2.0, 100: -3.0}
mock_output = _create_mock_request_output(mock_logprobs)
state = SimpleNamespace(sampling_params=None)
async def mock_generate(*args, **kwargs):
state.sampling_params = args[1]
yield mock_output
mock_engine.generate = mock_generate
@@ -317,6 +320,8 @@ class TestGeneration:
result = await serving.create_generative_scoring(request, None)
assert isinstance(result, GenerativeScoringResponse)
assert state.sampling_params is not None
assert state.sampling_params.watermarking is False
assert len(result.data) == 2
for item_result in result.data:
assert 0.0 <= item_result.score <= 1.0
@@ -154,6 +154,40 @@ async def client(server):
yield async_client
@pytest.fixture(scope="module")
def watermarked_server():
args = [
"--dtype",
"half",
"--enable-auto-tool-choice",
"--structured-outputs-config.backend",
"xgrammar",
"--tool-call-parser",
"hermes",
"--reasoning-parser",
"qwen3",
"--gpu-memory-utilization",
"0.4",
"--enforce-eager",
"--watermark-config",
'{"algorithm":"gumbel","key":42}',
] + ROCM_EXTRA_ARGS
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@pytest.fixture(params=["server", "watermarked_server"], ids=["plain", "watermarked"])
def tool_use_server(request: pytest.FixtureRequest):
return request.getfixturevalue(request.param)
@pytest_asyncio.fixture
async def tool_use_client(tool_use_server):
async with tool_use_server.get_async_client() as async_client:
yield async_client
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("stream", [True, False])
@@ -167,12 +201,13 @@ async def client(server):
)
@pytest.mark.parametrize("enable_thinking", [True, False])
async def test_function_tool_use(
client: openai.AsyncOpenAI,
tool_use_client: openai.AsyncOpenAI,
model_name: str,
stream: bool,
tool_choice: str | dict,
enable_thinking: bool,
):
client = tool_use_client
if not stream:
# Non-streaming test
chat_completion = await client.chat.completions.create(
+224 -7
View File
@@ -1,32 +1,249 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from io import BytesIO
import pytest
from fastapi import UploadFile
from vllm.entrypoints.openai.chat_completion.protocol import (
BatchChatCompletionRequest,
ChatCompletionRequest,
)
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.entrypoints.speech_to_text.transcription.protocol import (
TranscriptionRequest,
)
from vllm.entrypoints.speech_to_text.translation.protocol import TranslationRequest
from vllm.sampling_params import StructuredOutputsParams
def test_chat_request_disables_watermarking():
def test_chat_request_preserves_watermarking_opt_out():
request = ChatCompletionRequest(
messages=[{"role": "user", "content": "hello"}], watermarking=False
)
params = request.to_sampling_params(max_tokens=1, default_sampling_params={})
assert not params.watermarking
assert params.watermarking is False
def test_completion_request_disables_watermarking():
def test_completion_request_preserves_watermarking_opt_out():
request = CompletionRequest(prompt="hello", watermarking=False)
params = request.to_sampling_params(max_tokens=1)
assert not params.watermarking
assert params.watermarking is False
def test_responses_request_disables_watermarking():
def test_responses_request_preserves_watermarking_opt_out():
request = ResponsesRequest(input="hello", watermarking=False)
params = request.to_sampling_params(default_max_tokens=1)
assert not params.watermarking
assert params.watermarking is False
def test_batch_chat_request_preserves_watermarking_opt_out():
request = BatchChatCompletionRequest(
messages=[[{"role": "user", "content": "hello"}]], watermarking=False
)
converted = request.to_chat_completion_request(request.messages[0])
params = converted.to_sampling_params(max_tokens=1, default_sampling_params={})
assert params.watermarking is False
@pytest.mark.parametrize(
"api_request,conversion,kwargs",
[
(
ChatCompletionRequest(
messages=[{"role": "user", "content": "hello"}],
use_beam_search=True,
watermarking=False,
),
"to_beam_search_params",
{"max_tokens": 1, "default_sampling_params": {}},
),
(
CompletionRequest(prompt="hello", use_beam_search=True, watermarking=False),
"to_beam_search_params",
{"max_tokens": 1, "default_sampling_params": {}},
),
(
TranscriptionRequest(
file=UploadFile(file=BytesIO(), filename="audio.wav"),
watermarking=False,
),
"to_sampling_params",
{"default_max_tokens": 1},
),
(
TranscriptionRequest(
file=UploadFile(file=BytesIO(), filename="audio.wav"),
watermarking=False,
),
"to_beam_search_params",
{"default_max_tokens": 1},
),
(
TranslationRequest(
file=UploadFile(file=BytesIO(), filename="audio.wav"),
watermarking=False,
),
"to_sampling_params",
{"default_max_tokens": 1},
),
(
TranslationRequest(
file=UploadFile(file=BytesIO(), filename="audio.wav"),
watermarking=False,
),
"to_beam_search_params",
{"default_max_tokens": 1},
),
],
)
def test_request_preserves_watermarking_opt_out(api_request, conversion, kwargs):
params = getattr(api_request, conversion)(**kwargs)
assert params.watermarking is False
@pytest.mark.parametrize("request_cls", [TranscriptionRequest, TranslationRequest])
def test_speech_requests_leave_watermarking_unspecified(request_cls):
request = request_cls(
file=UploadFile(file=BytesIO(), filename="audio.wav"),
temperature=0.8,
)
sampling_params = request.to_sampling_params(default_max_tokens=1)
beam_params = request.to_beam_search_params(default_max_tokens=1)
assert request.watermarking is None
assert sampling_params.watermarking is None
assert beam_params.watermarking is None
@pytest.mark.parametrize(
"api_request,kwargs",
[
(
ChatCompletionRequest(
messages=[{"role": "user", "content": "choose A or B"}],
temperature=0.8,
structured_outputs=StructuredOutputsParams(choice=["A", "B"]),
),
{"max_tokens": 8, "default_sampling_params": {}},
),
(
CompletionRequest(
prompt="choose A or B",
temperature=0.8,
structured_outputs=StructuredOutputsParams(choice=["A", "B"]),
),
{"max_tokens": 8},
),
(
ResponsesRequest(
input="choose A or B",
temperature=0.8,
structured_outputs=StructuredOutputsParams(choice=["A", "B"]),
),
{"default_max_tokens": 8},
),
],
)
def test_request_preserves_unspecified_watermarking_with_structured_outputs(
api_request, kwargs
):
params = api_request.to_sampling_params(**kwargs)
assert params.watermarking is None
assert params.structured_outputs is not None
assert params.structured_outputs.choice == ["A", "B"]
@pytest.mark.parametrize(
"api_request,conversion,kwargs",
[
(
ChatCompletionRequest(
messages=[{"role": "user", "content": "hello"}], watermarking=True
),
"to_sampling_params",
{"max_tokens": 1, "default_sampling_params": {}},
),
(
ChatCompletionRequest(
messages=[{"role": "user", "content": "hello"}],
use_beam_search=True,
watermarking=True,
),
"to_beam_search_params",
{"max_tokens": 1, "default_sampling_params": {}},
),
(
CompletionRequest(prompt="hello", watermarking=True),
"to_sampling_params",
{"max_tokens": 1},
),
(
CompletionRequest(prompt="hello", use_beam_search=True, watermarking=True),
"to_beam_search_params",
{"max_tokens": 1},
),
(
ResponsesRequest(input="hello", watermarking=True),
"to_sampling_params",
{"default_max_tokens": 1},
),
(
TranscriptionRequest(
file=UploadFile(file=BytesIO(), filename="audio.wav"),
watermarking=True,
),
"to_sampling_params",
{"default_max_tokens": 1},
),
(
TranslationRequest(
file=UploadFile(file=BytesIO(), filename="audio.wav"),
watermarking=True,
),
"to_sampling_params",
{"default_max_tokens": 1},
),
],
)
def test_request_preserves_explicit_watermarking(api_request, conversion, kwargs):
params = getattr(api_request, conversion)(**kwargs)
assert params.watermarking is True
def test_batch_chat_request_preserves_explicit_watermarking():
request = BatchChatCompletionRequest(
messages=[[{"role": "user", "content": "hello"}]], watermarking=True
)
converted = request.to_chat_completion_request(request.messages[0])
params = converted.to_sampling_params(max_tokens=1, default_sampling_params={})
assert converted.watermarking is True
assert params.watermarking is True
def test_batch_chat_request_inherits_watermarking():
request = BatchChatCompletionRequest(
messages=[[{"role": "user", "content": "hello"}]]
)
converted = request.to_chat_completion_request(request.messages[0])
params = converted.to_sampling_params(max_tokens=1, default_sampling_params={})
assert converted.watermarking is None
assert params.watermarking is None
@@ -4,6 +4,7 @@
import asyncio
import json
import warnings
from types import SimpleNamespace
import numpy as np
import pybase64 as base64
@@ -12,6 +13,7 @@ import websockets
from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
from vllm.assets.audio import AudioAsset
from vllm.entrypoints.speech_to_text.realtime.connection import RealtimeConnection
from vllm.multimodal.media.audio import load_audio
MISTRAL_FORMAT_ARGS = [
@@ -26,6 +28,46 @@ MISTRAL_FORMAT_ARGS = [
MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
def test_realtime_generation_preserves_watermarking():
transcription_outputs = []
class EngineClient:
def generate(self, **kwargs):
async def outputs():
output = SimpleNamespace(
outputs=[],
prompt_token_ids=[],
is_watermarked=kwargs["sampling_params"].watermarking,
)
transcription_outputs.append(output)
yield output
return outputs()
async def streaming_input():
chunk: np.ndarray
for chunk in ():
yield chunk
async def send(_event):
return None
connection = object.__new__(RealtimeConnection)
connection.connection_id = "test"
connection.serving = SimpleNamespace(
model_cls=SimpleNamespace(realtime_max_tokens=1),
engine_client=EngineClient(),
)
connection.audio_queue = asyncio.Queue()
connection._is_connected = True
connection.send = send
asyncio.run(connection._run_generation(streaming_input(), asyncio.Queue()))
assert len(transcription_outputs) == 1
assert transcription_outputs[0].is_watermarked is None
def _get_websocket_url(server: RemoteOpenAIServer) -> str:
"""Convert HTTP URL to WebSocket URL for realtime endpoint."""
http_url = server.url_root
@@ -198,4 +198,6 @@ async def test_language_detection_cancel_aborts_engine_request():
with pytest.raises(asyncio.CancelledError):
await task
sampling_params = engine_client.generate.call_args.args[1]
assert sampling_params.watermarking is False
engine_client.abort.assert_awaited_once_with(request_id)
+166 -9
View File
@@ -1,13 +1,33 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from collections.abc import Callable, Sequence
from typing import Any, cast
from unittest.mock import MagicMock
from vllm import CompletionOutput, RequestOutput
import pytest
import torch
from vllm import CompletionOutput, PromptType, RequestOutput, SamplingParams
from vllm import logger as vllm_logger
from vllm.config import VllmConfig
from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin
from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin
from vllm.inputs import TokensInput
from vllm.entrypoints.generate.beam_search.utils import (
BeamSearchInstance,
BeamSearchSequence,
)
from vllm.inputs import EngineInput, TokensInput
from vllm.logprobs import Logprob, SampleLogprobs
from vllm.sampling_params import BeamSearchParams
from vllm.v1.structured_output.backend_types import StructuredOutputBackend
@pytest.fixture
def reset_warning_once():
vllm_logger._print_warning_once.cache_clear()
yield
vllm_logger._print_warning_once.cache_clear()
class _Tokenizer:
@@ -27,8 +47,26 @@ class _Renderer:
return _Tokenizer()
class _VllmConfig:
watermark_config = object()
speculative_config = None
def _check_supports_watermarking(self, params: BeamSearchParams) -> bool:
return VllmConfig._check_supports_watermarking(cast(VllmConfig, self), params)
class _InputProcessor:
vllm_config = _VllmConfig()
def resolve_watermarking(self, params):
return self.vllm_config._check_supports_watermarking(params)
class _EngineClient:
input_processor = _InputProcessor()
async def generate(self, prompt, *args, **kwargs):
assert args[0].watermarking is False
yield RequestOutput(
request_id=kwargs.get("request_id", "test-request"),
prompt=prompt.get("prompt"),
@@ -56,11 +94,42 @@ class _EngineClient:
)
class _Serving(BeamSearchOnlineMixin):
class _AsyncServing(BeamSearchOnlineMixin):
renderer = _Renderer()
engine_client = _EngineClient()
class _OfflineServing(BeamSearchOfflineMixin):
renderer = _Renderer()
llm_engine = _EngineClient()
def _preprocess_cmpl(
self,
prompts: Sequence[PromptType],
tokenization_kwargs: dict[str, Any] | None = None,
mm_processor_kwargs: dict[str, Any] | None = None,
) -> Sequence[EngineInput]:
return cast(Sequence[EngineInput], prompts)
def _lora_request_to_seq(self, lora_request, num_requests):
return [None] * num_requests
def _beam_search_step(
self,
instances_batch: list[BeamSearchInstance],
base_sampling_params: SamplingParams,
eos_token_id: int | None,
ignore_eos: bool,
beam_width: int,
sort_beams_key: Callable[..., Any],
structured_output_backend: StructuredOutputBackend | None,
structured_output_key: tuple[Any, ...] | None,
structured_output_bitmask: torch.Tensor | None,
) -> bool:
assert base_sampling_params.watermarking is False
return True
@pytest.mark.asyncio
async def test_beam_search_handles_extra_logprob_candidates() -> None:
prompt: TokensInput = {
@@ -68,10 +137,11 @@ async def test_beam_search_handles_extra_logprob_candidates() -> None:
"prompt": "prompt",
"prompt_token_ids": [1],
}
params = BeamSearchParams(beam_width=2, max_tokens=1)
params = BeamSearchParams(beam_width=2, max_tokens=1, watermarking=False)
outputs = [
output async for output in _Serving().beam_search(prompt, "request", params)
output
async for output in _AsyncServing().beam_search(prompt, "request", params)
]
assert len(outputs) == 1
@@ -101,10 +171,12 @@ async def test_beam_search_respects_skip_special_tokens(
max_tokens=1,
ignore_eos=True,
skip_special_tokens=skip_special_tokens,
watermarking=False,
)
outputs = [
output async for output in _Serving().beam_search(prompt, "request", params)
output
async for output in _AsyncServing().beam_search(prompt, "request", params)
]
assert outputs[0].outputs[0].text == expected_text
@@ -139,7 +211,7 @@ async def test_beam_search_abort_returns_partial_outputs(
output.finish_reason = "abort"
yield result
serving = _Serving()
serving = _AsyncServing()
monkeypatch.setattr(serving.engine_client, "generate", generate)
prompt: TokensInput = {
"type": "token",
@@ -149,7 +221,9 @@ async def test_beam_search_abort_returns_partial_outputs(
outputs = [
output
async for output in serving.beam_search(
prompt, "request", BeamSearchParams(beam_width=2, max_tokens=4)
prompt,
"request",
BeamSearchParams(beam_width=2, max_tokens=4, watermarking=False),
)
]
@@ -165,3 +239,86 @@ async def test_beam_search_abort_returns_partial_outputs(
assert output.logprobs is not None
assert len(output.logprobs) == abort_after
assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1 * abort_after)
@pytest.mark.asyncio
async def test_beam_search_warns_and_disables_watermarking(
caplog_vllm, reset_warning_once
) -> None:
prompt = {
"type": "token",
"prompt": "prompt",
"prompt_token_ids": [1],
}
params = BeamSearchParams(beam_width=2, max_tokens=1)
with caplog_vllm.at_level("WARNING"):
output = await anext(_AsyncServing().beam_search(prompt, "request", params))
assert "beam search requests will run without watermarking" in caplog_vllm.text
# The caller's params object is never rewritten by admission.
assert params.watermarking is None
assert output.finished
def test_offline_beam_search_warns_and_disables_watermarking(
caplog_vllm, reset_warning_once
) -> None:
params = BeamSearchParams(beam_width=2, max_tokens=1)
with caplog_vllm.at_level("WARNING"):
outputs = _OfflineServing().beam_search(
cast(list[PromptType], [{"type": "token", "prompt_token_ids": [1]}]),
params,
)
assert "beam search requests will run without watermarking" in caplog_vllm.text
# The caller's params object is never rewritten by admission.
assert params.watermarking is None
assert len(outputs) == 1
def test_offline_beam_search_disables_internal_watermarking() -> None:
outputs = _OfflineServing().beam_search(
cast(list[PromptType], [{"type": "token", "prompt_token_ids": [1]}]),
BeamSearchParams(beam_width=2, max_tokens=1, watermarking=False),
)
assert len(outputs) == 1
def test_offline_structured_beam_search_disables_internal_watermarking() -> None:
grammar = MagicMock()
grammar.is_terminated.return_value = False
grammar.fill_bitmask.side_effect = lambda bitmask, index: bitmask[index].fill_(1)
backend = MagicMock()
backend.compile_grammar.return_value = grammar
serving = _OfflineServing()
serving.model_config = MagicMock()
serving.model_config.get_vocab_size.return_value = 32
beam = BeamSearchSequence(
orig_prompt={"type": "token", "prompt_token_ids": [1]},
tokens=[1],
logprobs=[],
)
base_params = SamplingParams(
logprobs=2,
watermarking=False,
skip_clone=True,
)
entries = serving._build_beam_sampling_params(
[beam],
base_params,
backend,
("regex", ".*"),
torch.zeros((1, 1), dtype=torch.int32),
)
assert entries[0] is not None
beam_params = entries[0][0]
assert beam_params is not base_params
assert beam_params.logprobs == base_params.logprobs
assert beam_params.watermarking is False
assert beam_params.allowed_token_ids == [0]
assert base_params.allowed_token_ids is None
+39 -15
View File
@@ -47,6 +47,7 @@ from vllm.config.speculative import _validate_qwen3_omni_dspark
from vllm.config.utils import get_field
from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel
from vllm.platforms import current_platform
from vllm.sampling_params import BeamSearchParams
from vllm.transformers_utils.config import (
_patch_hf_transformers_nested_rope_validation,
get_pooling_config,
@@ -3630,7 +3631,7 @@ def test_target_only_gumbel_allows_speculative_decoding(caplog_vllm, disable_log
)
with caplog_vllm.at_level(logging.WARNING):
config._check_watermarking_unsupported()
config._check_supports_watermarking()
assert "Target-only watermarking leaves accepted draft tokens" in caplog_vllm.text
@@ -3651,7 +3652,7 @@ def test_speculative_watermarking_without_context_dedup_does_not_warn(
caplog_vllm.clear()
with caplog_vllm.at_level(logging.WARNING):
config._check_watermarking_unsupported()
config._check_supports_watermarking()
assert "dedup" not in caplog_vllm.text.lower()
@@ -3672,7 +3673,7 @@ def test_speculative_context_dedup_logs_no_unsupported_warning(
caplog_vllm.clear()
with caplog_vllm.at_level(logging.WARNING):
config._check_watermarking_unsupported()
config._check_supports_watermarking()
assert "dedup" not in caplog_vllm.text.lower()
@@ -3690,7 +3691,7 @@ def test_gumbel_rejects_speculative_decoding_without_target_only():
)
with pytest.raises(ValueError, match="'gumbel'.*allow_target_only_watermarking"):
config._check_watermarking_unsupported()
config._check_supports_watermarking()
def test_dual_key_gumbel_warns_that_configured_alpha_is_unused(
@@ -3708,7 +3709,7 @@ def test_dual_key_gumbel_warns_that_configured_alpha_is_unused(
)
with caplog_vllm.at_level(logging.WARNING):
config._check_watermarking_unsupported()
config._check_supports_watermarking()
assert "The configured alpha=0.25 is not used" in caplog_vllm.text
@@ -3734,7 +3735,7 @@ def test_dual_key_gumbel_alpha_warning_is_not_emitted(
)
with caplog_vllm.at_level(logging.WARNING):
config._check_watermarking_unsupported()
config._check_supports_watermarking()
assert "is not used" not in caplog_vllm.text
@@ -3750,7 +3751,7 @@ def test_dual_key_gumbel_requires_probabilistic_drafting():
)
with pytest.raises(ValueError, match="draft_sample_method='probabilistic'"):
config._check_watermarking_unsupported()
config._check_supports_watermarking()
@pytest.mark.parametrize("method", ["eagle", "eagle3", "mtp"])
@@ -3764,7 +3765,7 @@ def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(method):
parallel_drafting=False,
)
config._check_watermarking_unsupported()
config._check_supports_watermarking()
def test_dual_key_gumbel_supports_dspark():
@@ -3777,7 +3778,7 @@ def test_dual_key_gumbel_supports_dspark():
parallel_drafting=True,
)
config._check_watermarking_unsupported()
config._check_supports_watermarking()
def test_dual_key_gumbel_rejects_non_autoregressive_speculation():
@@ -3791,7 +3792,7 @@ def test_dual_key_gumbel_rejects_non_autoregressive_speculation():
)
with pytest.raises(ValueError, match="autoregressive model-based"):
config._check_watermarking_unsupported()
config._check_supports_watermarking()
@pytest.mark.parametrize(
@@ -3815,17 +3816,40 @@ def test_dual_key_gumbel_rejects_incompatible_speculative_modes(overrides, match
config.speculative_config = SimpleNamespace(**values)
with pytest.raises(ValueError, match=match):
config._check_watermarking_unsupported()
config._check_supports_watermarking()
def test_gumbel_watermark_rejects_beam_search():
with pytest.raises(ValueError, match="Beam search is not supported"):
_watermarked_vllm_config()._check_watermarking_unsupported(beam_search=True)
def test_gumbel_watermark_disables_beam_search_watermarking(caplog_vllm):
enabled = _watermarked_vllm_config()._check_supports_watermarking(
BeamSearchParams(beam_width=2, max_tokens=1)
)
assert not enabled
assert "beam search requests will run without watermarking" in caplog_vllm.text
def test_gumbel_watermark_allows_unwatermarked_beam_search():
_watermarked_vllm_config()._check_supports_watermarking(
BeamSearchParams(beam_width=2, max_tokens=1, watermarking=False)
)
def test_unwatermarked_beam_does_not_bypass_engine_incompatibilities():
config = _watermarked_vllm_config()
config.speculative_config = SpeculativeConfig(
method="ngram",
num_speculative_tokens=1,
)
with pytest.raises(ValueError, match="Speculative decoding with watermarking"):
config._check_supports_watermarking(
BeamSearchParams(beam_width=2, max_tokens=1, watermarking=False)
)
def test_gumbel_watermark_rejects_custom_sampler():
with pytest.raises(ValueError, match="custom samplers are not supported"):
_watermarked_vllm_config()._check_watermarking_unsupported(custom_sampler=True)
_watermarked_vllm_config()._check_supports_watermarking(custom_sampler=True)
def test_watermark_key_must_fit_in_64_bits():
+1
View File
@@ -112,6 +112,7 @@ def test_diffusion_canvas_width_is_accepted_with_either_scheduler(
diffusion_config=DiffusionConfig(canvas_length=8),
tokenizer=None,
validate_logits_processors_params=lambda params: None,
resolve_watermarking=lambda params: False,
)
params = SamplingParams(extra_args=extra_args)
InputProcessor._validate_params(processor, params, ("generate",))
@@ -82,7 +82,12 @@ def _validate(enable_trace_replay: bool) -> None:
enable_trace_replay=enable_trace_replay,
is_diffusion=False,
),
vllm_config=SimpleNamespace(reasoning_config=None),
vllm_config=SimpleNamespace(
reasoning_config=None,
# Engine has no watermark configuration.
_check_supports_watermarking=lambda params: False,
),
resolve_watermarking=lambda params: False,
speculative_config=None,
structured_outputs_config=None,
diffusion_config=None,
@@ -0,0 +1,27 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from vllm.v1.engine.async_llm import AsyncLLM
@pytest.mark.asyncio
async def test_kv_transfer_rejection_stub_disables_watermarking():
engine = object.__new__(AsyncLLM)
add_request = AsyncMock()
engine.engine_core = SimpleNamespace(
add_request_async=add_request,
shutdown=lambda **kwargs: None,
)
await engine.notify_kv_transfer_request_rejected(
"request",
{"do_remote_prefill": True},
)
request = add_request.call_args.args[0]
assert request.sampling_params.watermarking is False
+1
View File
@@ -518,6 +518,7 @@ def test_prompt_logprob_token_ids_require_v2_model_runner():
structured_outputs_config=None,
tokenizer=None,
validate_logits_processors_params=lambda params: None,
resolve_watermarking=lambda params: False,
)
params = SamplingParams(prompt_logprob_token_ids=[[1, 2]])
with patch.object(SamplingParams, "verify"):
+34
View File
@@ -6,6 +6,7 @@ from types import SimpleNamespace
import pytest
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.worker.gpu import warmup
from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup
@@ -31,6 +32,39 @@ def test_mixed_warmup_skipped_for_single_seq(max_num_reqs):
)
def test_mixed_warmup_disables_watermarking():
outputs: list[SchedulerOutput] = []
runner = SimpleNamespace(
is_pooling_model=False,
max_num_reqs=2,
kv_cache_config=SimpleNamespace(
kv_cache_groups=[
SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16))
],
num_blocks=128,
),
vllm_config=SimpleNamespace(num_lookahead_tokens=0),
max_model_len=128,
model_state=SimpleNamespace(max_encoder_len=0),
kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
)
assert run_mixed_prefill_decode_warmup(
runner,
worker_execute_model=outputs.append,
worker_sample_tokens=lambda grammar_output: None,
num_tokens=32,
)
sampling_params = [
request.sampling_params
for output in outputs
for request in output.scheduled_new_reqs
]
assert len(sampling_params) == 2
assert all(params.watermarking is False for params in sampling_params)
@pytest.mark.parametrize("fail_warmup", [False, True])
def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_warmup):
"""Startup must warm fixed drafts before calibration and retain its manager."""
+353
View File
@@ -0,0 +1,353 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from vllm import SamplingParams
from vllm import logger as vllm_logger
from vllm.config import VllmConfig, WatermarkConfig
from vllm.renderers import BaseRenderer
from vllm.sampling_params import BeamSearchParams
from vllm.v1.engine import EngineCoreRequest
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.engine.input_processor import InputProcessor
from vllm.v1.engine.llm_engine import LLMEngine
def _input_processor(server_uses_watermarking: bool = True) -> InputProcessor:
config = VllmConfig()
config.model_config = SimpleNamespace(
try_get_generation_config=lambda: {},
logits_processors=None,
return_sampling_mask=False,
enable_trace_replay=True,
is_multimodal_model=False,
supports_multimodal_inputs=False,
runner_type="generate",
is_diffusion=False,
)
config.watermark_config = (
WatermarkConfig(key=42) if server_uses_watermarking else None
)
renderer = MagicMock(spec=BaseRenderer)
renderer.tokenizer = None
renderer._executor = None
return InputProcessor(config, renderer)
def _validate(
params: SamplingParams, server_uses_watermarking: bool = True
) -> bool | None:
processor = _input_processor(server_uses_watermarking)
with patch.object(SamplingParams, "verify"):
return processor._validate_params(params, ("generate",))
def _engine_core_request(params: SamplingParams) -> EngineCoreRequest:
return EngineCoreRequest(
request_id="request",
prompt_token_ids=[1],
mm_features=None,
sampling_params=params,
pooling_params=None,
arrival_time=0.0,
lora_request=None,
cache_salt=None,
data_parallel_rank=None,
)
@pytest.fixture
def reset_warning_once():
vllm_logger._print_warning_once.cache_clear()
yield
vllm_logger._print_warning_once.cache_clear()
def test_greedy_watermarked_requests_warn_once(caplog_vllm, reset_warning_once):
first = SamplingParams(temperature=0)
second = SamplingParams(temperature=0)
with caplog_vllm.at_level("WARNING"):
first_resolved = _validate(first)
second_resolved = _validate(second)
message = "subsequent greedy requests will use ordinary greedy sampling"
assert caplog_vllm.text.count(message) == 1
assert first_resolved is False
assert second_resolved is False
# The caller's params objects are left untouched.
assert first.watermarking is None
assert second.watermarking is None
def test_default_watermarking_resolves_when_configured():
enabled = SamplingParams()
assert _validate(enabled) is True
# The caller's params object keeps inheriting.
assert enabled.watermarking is None
@pytest.mark.parametrize("watermarking", [None, True])
def test_requested_watermarking_without_engine_config_warns_and_disables(
watermarking, caplog_vllm, reset_warning_once
):
params = SamplingParams(watermarking=watermarking)
with caplog_vllm.at_level("WARNING"):
resolved = _validate(params, server_uses_watermarking=False)
assert "engine has no watermark configuration" in caplog_vllm.text
assert resolved is False
assert params.watermarking is watermarking
@pytest.mark.parametrize("watermarking", [None, True])
def test_explicit_and_default_greedy_warn_and_disable(
watermarking, caplog_vllm, reset_warning_once
):
params = SamplingParams(temperature=0, watermarking=watermarking)
with caplog_vllm.at_level("WARNING"):
resolved = _validate(params)
assert "ordinary greedy sampling" in caplog_vllm.text
assert resolved is False
assert params.watermarking is watermarking
def test_trace_replay_with_watermarking_warns_and_disables(
caplog_vllm, reset_warning_once
):
params = SamplingParams(trace_decode_token_ids=[1])
with caplog_vllm.at_level("WARNING"):
resolved = _validate(params)
assert "trace replay requests will run without watermarking" in caplog_vllm.text
assert resolved is False
assert params.watermarking is None
@pytest.mark.parametrize(
"params",
[
SamplingParams(temperature=0, watermarking=False),
SamplingParams(trace_decode_token_ids=[1], watermarking=False),
],
)
def test_incompatible_modes_are_allowed_when_watermarking_is_disabled(params):
_validate(params)
@pytest.mark.parametrize(
"params",
[
SamplingParams(temperature=0),
SamplingParams(trace_decode_token_ids=[1]),
],
)
def test_requests_without_engine_config_are_unwatermarked(
params, caplog_vllm, reset_warning_once
):
with caplog_vllm.at_level("WARNING"):
resolved = _validate(params, server_uses_watermarking=False)
assert "engine has no watermark configuration" in caplog_vllm.text
assert resolved is False
assert params.watermarking is None
def test_trace_replay_takes_precedence_over_greedy_warning(
caplog_vllm, reset_warning_once
):
params = SamplingParams(temperature=0, trace_decode_token_ids=[1])
with caplog_vllm.at_level("WARNING"):
_validate(params)
assert "trace replay requests will run without watermarking" in caplog_vllm.text
assert "ordinary greedy sampling" not in caplog_vllm.text
@pytest.mark.parametrize(
"params",
[
SamplingParams(seed=42),
SamplingParams(n=2),
SamplingParams(n=2, seed=42),
],
)
def test_seeded_and_parallel_sampling_are_allowed(params):
_validate(params)
def test_direct_engine_trace_replay_warns_and_disables(
caplog_vllm, monkeypatch, reset_warning_once
):
engine = object.__new__(LLMEngine)
engine.input_processor = _input_processor()
engine.vllm_config = engine.input_processor.vllm_config
assign_request_id = MagicMock()
monkeypatch.setattr(engine.input_processor, "assign_request_id", assign_request_id)
engine.output_processor = MagicMock()
engine.engine_core = MagicMock()
request = _engine_core_request(SamplingParams(trace_decode_token_ids=[1]))
with caplog_vllm.at_level("WARNING"):
engine.add_request(
request.request_id,
request,
SamplingParams(watermarking=False),
)
assert "trace replay requests will run without watermarking" in caplog_vllm.text
assert request.sampling_params.watermarking is False
assert request.sampling_params._watermarking_skipped is True
assign_request_id.assert_called_once_with(request)
def test_direct_engine_request_uses_embedded_opt_out(monkeypatch):
engine = object.__new__(LLMEngine)
engine.input_processor = _input_processor()
engine.vllm_config = engine.input_processor.vllm_config
assign_request_id = MagicMock()
monkeypatch.setattr(engine.input_processor, "assign_request_id", assign_request_id)
engine.output_processor = MagicMock()
engine.engine_core = MagicMock()
request = _engine_core_request(SamplingParams(temperature=0, watermarking=False))
engine.add_request(
request.request_id,
request,
SamplingParams(temperature=0),
)
assert request.sampling_params._watermarking_skipped is False
assign_request_id.assert_called_once_with(request)
def test_direct_async_engine_trace_replay_warns_and_disables(
caplog_vllm, monkeypatch, reset_warning_once
):
engine = object.__new__(AsyncLLM)
engine.engine_core = SimpleNamespace(
resources=SimpleNamespace(engine_dead=False),
shutdown=lambda **kwargs: None,
add_request_async=AsyncMock(),
)
engine.output_handler = None
engine.input_processor = _input_processor()
engine.vllm_config = engine.input_processor.vllm_config
engine.output_processor = MagicMock()
engine.log_requests = False
monkeypatch.setattr(engine, "_run_output_handler", MagicMock())
request = _engine_core_request(SamplingParams(trace_decode_token_ids=[1]))
async def add_request() -> None:
await engine.add_request(
request.request_id,
request,
SamplingParams(watermarking=False),
)
with caplog_vllm.at_level("WARNING"):
asyncio.run(add_request())
assert "trace replay requests will run without watermarking" in caplog_vllm.text
assert request.sampling_params.watermarking is False
assert request.sampling_params._watermarking_skipped is True
def test_direct_async_engine_greedy_request_warns_and_disables_watermarking(
caplog_vllm, monkeypatch, reset_warning_once
):
engine = object.__new__(AsyncLLM)
engine.engine_core = SimpleNamespace(
resources=SimpleNamespace(engine_dead=False),
shutdown=lambda **kwargs: None,
add_request_async=AsyncMock(),
)
engine.output_handler = None
engine.input_processor = _input_processor()
engine.vllm_config = engine.input_processor.vllm_config
engine.output_processor = MagicMock()
engine.log_requests = False
monkeypatch.setattr(engine, "_run_output_handler", MagicMock())
request = _engine_core_request(SamplingParams(temperature=0))
async def add_request() -> None:
await engine.add_request(
request.request_id,
request,
SamplingParams(watermarking=False),
)
with caplog_vllm.at_level("WARNING"):
asyncio.run(add_request())
assert "subsequent greedy requests will use ordinary greedy sampling" in (
caplog_vllm.text
)
assert request.sampling_params.watermarking is False
assert request.sampling_params._watermarking_skipped is True
@pytest.mark.parametrize(
("kwargs", "server_uses_watermarking", "skipped"),
[
({}, True, False),
({"temperature": 0}, True, True),
({"temperature": 0, "watermarking": False}, True, False),
({"temperature": 0}, False, False),
],
)
def test_apply_watermarking_records_skipped_requests(
kwargs, server_uses_watermarking, skipped, reset_warning_once
):
params = SamplingParams(**kwargs)
processor = _input_processor(server_uses_watermarking)
processor.apply_watermarking(params, processor.resolve_watermarking(params))
assert params._watermarking_skipped is skipped
# Resolving an already resolved request keeps the flag.
processor.apply_watermarking(params, processor.resolve_watermarking(params))
assert params._watermarking_skipped is skipped
def test_shared_sampling_params_are_not_rewritten_by_admission(reset_warning_once):
"""A params object reused across prompts/engines must keep inheriting."""
params = SamplingParams(temperature=0.8)
assert _validate(params, server_uses_watermarking=False) is False
assert params.watermarking is None
# Same object submitted to a watermarked engine must still be watermarked.
assert _validate(params, server_uses_watermarking=True) is True
greedy = SamplingParams(temperature=0.0)
assert _validate(greedy) is False
greedy.temperature = 0.8
assert _validate(greedy) is True
def test_shared_beam_search_params_are_not_rewritten_by_admission(
caplog_vllm, reset_warning_once
):
"""Beam search must keep resolving per engine after reuse."""
params = BeamSearchParams(beam_width=2, max_tokens=4, temperature=0.8)
# A plain engine leaves the tri-state unresolved on the caller's object.
assert _input_processor(False).resolve_watermarking(params) is False
assert params.watermarking is None
# The same object on a watermarked engine is warned about and still
# resolves to False, without the first decision being written back.
with caplog_vllm.at_level("WARNING"):
assert _input_processor(True).resolve_watermarking(params) is False
assert "beam search requests will run without watermarking" in caplog_vllm.text
assert params.watermarking is None
@@ -0,0 +1,204 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Generator
import pytest
from vllm import LLM, SamplingParams
from vllm.platforms import current_platform
from vllm.sampling_params import StructuredOutputsParams
MODEL = "hmellor/tiny-random-LlamaForCausalLM"
MAX_NUM_BATCHED_TOKENS = 32
pytestmark = pytest.mark.skipif(
not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
)
@pytest.fixture(scope="module")
def watermarked_llm(vllm_runner) -> Generator[LLM, None, None]:
with vllm_runner(
MODEL,
dtype="half",
enforce_eager=True,
enable_prefix_caching=True,
enable_chunked_prefill=True,
max_num_batched_tokens=MAX_NUM_BATCHED_TOKENS,
max_model_len=128,
max_num_seqs=4,
disable_log_stats=False,
watermark_config={"algorithm": "gumbel", "key": 42},
) as runner:
yield runner.llm
def test_watermark_with_combined_sampling_controls(watermarked_llm: LLM):
params = SamplingParams(
temperature=0.8,
top_k=8,
top_p=0.9,
min_p=0.01,
presence_penalty=0.2,
frequency_penalty=0.2,
repetition_penalty=1.1,
logit_bias={10: 100.0},
allowed_token_ids=[10],
max_tokens=8,
ignore_eos=True,
)
output = watermarked_llm.generate("Sampling smoke test", params)[0]
assert list(output.outputs[0].token_ids) == [10] * 8
# The assertion above holds with or without watermarking, because
# allowed_token_ids leaves a single candidate. Repeat with a wider allowed
# set, where only watermarked sampling is reproducible without a seed.
wide = SamplingParams(
temperature=0.8,
top_k=8,
top_p=0.9,
min_p=0.01,
presence_penalty=0.2,
frequency_penalty=0.2,
repetition_penalty=1.1,
allowed_token_ids=list(range(10, 210)),
max_tokens=8,
ignore_eos=True,
)
first = watermarked_llm.generate("Sampling smoke test", wide)[0]
second = watermarked_llm.generate("Sampling smoke test", wide)[0]
assert set(first.outputs[0].token_ids) <= set(range(10, 210))
assert first.outputs[0].token_ids == second.outputs[0].token_ids
def test_watermark_with_penalties_and_nucleus_sampling(watermarked_llm: LLM):
params = SamplingParams(
temperature=0.8,
top_k=8,
top_p=0.9,
min_p=0.01,
presence_penalty=0.2,
frequency_penalty=0.2,
repetition_penalty=1.1,
max_tokens=8,
ignore_eos=True,
)
first = watermarked_llm.generate("Penalty smoke test", params)[0]
second = watermarked_llm.generate("Penalty smoke test", params)[0]
assert len(first.outputs[0].token_ids) == 8
assert first.outputs[0].token_ids == second.outputs[0].token_ids
def test_plain_gumbel_seeded_parallel_sampling_is_deterministic(
watermarked_llm: LLM,
):
output = watermarked_llm.generate(
"Seeded parallel sampling smoke test",
SamplingParams(
n=2,
seed=42,
temperature=0.8,
max_tokens=8,
ignore_eos=True,
),
)[0]
assert len(output.outputs) == 2
assert all(len(candidate.token_ids) == 8 for candidate in output.outputs)
assert output.outputs[0].token_ids == output.outputs[1].token_ids
# Plain Gumbel keys on the token context, so the request seed does not
# change a watermarked draw. Pin that down rather than implying the seed
# is what makes the candidates match.
unseeded = watermarked_llm.generate(
"Seeded parallel sampling smoke test",
SamplingParams(
n=2,
temperature=0.8,
max_tokens=8,
ignore_eos=True,
),
)[0]
assert unseeded.outputs[0].token_ids == output.outputs[0].token_ids
@pytest.mark.parametrize(
"structured_outputs",
[
StructuredOutputsParams(choice=["A", "B"]),
StructuredOutputsParams(grammar='root ::= "A" | "B"'),
],
ids=["choice", "grammar"],
)
def test_watermark_with_structured_output(
watermarked_llm: LLM, structured_outputs: StructuredOutputsParams
):
params = SamplingParams(
temperature=0.8,
max_tokens=8,
structured_outputs=structured_outputs,
)
output = watermarked_llm.generate("Choose A or B:", params)[0]
assert output.outputs[0].text.strip() in ("A", "B")
def test_watermark_with_wide_grammar_is_deterministic(watermarked_llm: LLM):
# A wider grammar makes deterministic watermarking observable.
params = SamplingParams(
temperature=0.8,
max_tokens=8,
structured_outputs=StructuredOutputsParams(
grammar="root ::= [A-Z] [A-Z] [A-Z] [A-Z] [A-Z] [A-Z]"
),
)
first = watermarked_llm.generate("Emit six letters:", params)[0]
second = watermarked_llm.generate("Emit six letters:", params)[0]
assert first.outputs[0].text.strip().isupper()
assert first.outputs[0].token_ids == second.outputs[0].token_ids
def test_watermark_with_prefix_cache_and_chunked_prefill(watermarked_llm: LLM):
prompt = {"prompt_token_ids": [101] * 48}
params = SamplingParams(temperature=0.8, max_tokens=8, ignore_eos=True)
assert len(prompt["prompt_token_ids"]) > MAX_NUM_BATCHED_TOKENS
cold = watermarked_llm.generate(prompt, params)[0]
cached = watermarked_llm.generate(prompt, params)[0]
assert cached.num_cached_tokens > 0
assert cold.outputs[0].token_ids == cached.outputs[0].token_ids
def test_mixed_watermarked_batch(watermarked_llm: LLM):
prompts = [
{"prompt_token_ids": [101] * 32 + [102]},
{"prompt_token_ids": [201] * 32 + [202]},
]
marked = SamplingParams(temperature=0.8, max_tokens=8, ignore_eos=True)
ordinary = SamplingParams(
temperature=0.8,
seed=123,
max_tokens=8,
ignore_eos=True,
watermarking=False,
)
batched = watermarked_llm.generate(prompts, [marked, ordinary])
marked_alone = watermarked_llm.generate(prompts[0], marked)[0]
ordinary_alone = watermarked_llm.generate(prompts[1], ordinary)[0]
assert batched[0].outputs[0].token_ids == marked_alone.outputs[0].token_ids
assert batched[1].outputs[0].token_ids == ordinary_alone.outputs[0].token_ids
+23 -2
View File
@@ -4,9 +4,11 @@
import pytest
import torch
from vllm.config.watermarking import WatermarkConfig
from vllm.platforms import current_platform
from vllm.v1.watermarking import (
DualKeyGumbelWatermarkDetector,
DualKeyGumbelWatermarker,
GumbelWatermarkDetector,
GumbelWatermarker,
derive_watermark_key,
@@ -42,6 +44,21 @@ def test_detector_deduplicates_context_even_when_target_differs():
assert detection.num_scored_tokens == 3
def test_watermarker_respects_filtered_token_support():
logits = torch.full((2, 8), -torch.inf)
logits[0, 3] = 0
logits[1, 6] = 0
contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]])
token_ids = (
GumbelWatermarker(key=42)
.sample(logits, contexts, lambda values: None)
.token_ids
)
assert torch.equal(token_ids, torch.tensor([3, 6]))
def test_dual_key_detector_scores_each_token_against_both_keys():
token_ids = [1, 2, 3, 4, 5]
key_a_detector = GumbelWatermarkDetector(
@@ -99,8 +116,12 @@ def test_dual_key_detector_rejects_invalid_alpha():
DualKeyGumbelWatermarkDetector(key=42, alpha=1.1)
def test_dual_key_detector_default_alpha():
assert DualKeyGumbelWatermarkDetector(key=42).alpha == 0.2
def test_dual_key_alpha_defaults_match():
config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
watermarker = DualKeyGumbelWatermarker(key=42)
detector = DualKeyGumbelWatermarkDetector(key=42)
assert detector.alpha == watermarker.alpha == config.alpha
@pytest.mark.skipif(
@@ -0,0 +1,106 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import pytest
from tests.utils import multi_gpu_test
from vllm import SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.v1.engine.async_llm import AsyncLLM
MODEL = "hmellor/tiny-random-LlamaForCausalLM"
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize(
"parallel_kwargs",
[
{"tensor_parallel_size": 2, "enable_batch_sharded_sampling": True},
{"pipeline_parallel_size": 2},
],
ids=["tp2-sharded-sampling", "pp2"],
)
def test_watermarked_generation_with_model_parallelism(vllm_runner, parallel_kwargs):
prompts = ["Parallel watermark smoke"] * 2
with vllm_runner(
MODEL,
dtype="half",
enforce_eager=True,
max_model_len=128,
max_num_seqs=4,
watermark_config={"algorithm": "gumbel", "key": 42},
**parallel_kwargs,
) as runner:
first = runner.llm.generate(
prompts,
SamplingParams(
temperature=0.8,
seed=1,
max_tokens=8,
ignore_eos=True,
),
)
second = runner.llm.generate(
prompts,
SamplingParams(
temperature=0.8,
seed=2,
max_tokens=8,
ignore_eos=True,
),
)
first_tokens = [output.outputs[0].token_ids for output in first]
second_tokens = [output.outputs[0].token_ids for output in second]
assert all(len(token_ids) == 8 for token_ids in first_tokens)
assert first_tokens[0] == first_tokens[1]
assert first_tokens == second_tokens
@multi_gpu_test(num_gpus=2)
def test_watermarked_generation_with_data_parallelism():
async def run() -> None:
engine = AsyncLLM.from_engine_args(
AsyncEngineArgs(
model=MODEL,
dtype="half",
data_parallel_size=2,
enforce_eager=True,
max_model_len=128,
max_num_seqs=4,
watermark_config={"algorithm": "gumbel", "key": 42},
)
)
async def generate(index: int) -> tuple[int, ...]:
final_output = None
async for output in engine.generate(
"Data-parallel watermark smoke",
SamplingParams(
temperature=0.8,
seed=index,
max_tokens=8,
ignore_eos=True,
),
request_id=f"watermark-dp-{index}",
data_parallel_rank=index,
):
final_output = output
assert final_output is not None
return tuple(final_output.outputs[0].token_ids)
try:
token_ids = await asyncio.gather(generate(0), generate(1))
assert len(token_ids[0]) == 8
assert token_ids[0] == token_ids[1]
finally:
engine.shutdown()
try:
asyncio.run(run())
finally:
cleanup_dist_env_and_memory()
+147 -33
View File
@@ -8,8 +8,11 @@ import pytest
import torch
from vllm import SamplingParams
from vllm import logger as vllm_logger
from vllm.config import VllmConfig
from vllm.config.watermarking import WatermarkConfig
from vllm.platforms import current_platform
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
from vllm.v1.watermarking import (
DualKeyGumbelWatermarker,
SupportsSpeculativeDecoding,
@@ -26,8 +29,8 @@ from vllm.v1.watermarking.spec_decode import (
speculative_target_watermark_key,
)
from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample
from vllm.v1.worker.gpu import buffer_utils
from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
from vllm.v1.worker.gpu.sample.sampler import Sampler
from vllm.v1.worker.gpu.sample.watermark import (
draft_watermarking_mask,
philox_gumbel_sample,
@@ -48,6 +51,44 @@ def _argmax_sampler(logits: torch.Tensor) -> torch.Tensor:
return logits.argmax(dim=-1)
@pytest.fixture
def make_gpu_watermark_sampler(monkeypatch):
class CPUUvaBuffer:
def __init__(self, size, dtype):
self.cpu = torch.zeros(size, dtype=dtype)
self.np = self.cpu.numpy()
def uva(self, n=None):
return self.cpu[:n] if n is not None else self.cpu
monkeypatch.setattr(buffer_utils, "UvaBuffer", CPUUvaBuffer)
def make(watermarker, *, max_num_reqs=2, vocab_size=8, **kwargs):
req_states = SimpleNamespace(
max_num_reqs=max_num_reqs,
vocab_size=vocab_size,
device=torch.device("cpu"),
all_token_ids=SimpleNamespace(gpu=torch.zeros(max_num_reqs, 1)),
prompt_len=SimpleNamespace(
np=np.zeros(max_num_reqs, dtype=np.int32),
gpu=torch.zeros(max_num_reqs, dtype=torch.int32),
),
prefill_len=SimpleNamespace(np=np.zeros(max_num_reqs, dtype=np.int32)),
total_len=SimpleNamespace(gpu=torch.zeros(max_num_reqs, dtype=torch.int32)),
)
return GPUWatermarkSampler(
watermarker,
vllm_config=VllmConfig(),
max_num_reqs=max_num_reqs,
vocab_size=vocab_size,
device=torch.device("cpu"),
req_states=req_states,
**kwargs,
)
return make
@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
def test_watermarker_contract(algorithm: str):
watermarker = create_watermarker(
@@ -271,49 +312,47 @@ def test_dual_key_derivation_is_stable():
def test_sampling_params_can_disable_watermarking():
assert SamplingParams().watermarking
assert not SamplingParams.from_optional(watermarking=False).watermarking
assert SamplingParams().watermarking is None
assert SamplingParams.from_optional(watermarking=False).watermarking is False
def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch):
sampler = object.__new__(GPUWatermarkSampler)
sampler.watermarking = SimpleNamespace(np=np.ones(1, dtype=bool))
def test_gpu_sampler_warns_about_unexpected_greedy_watermarking(
monkeypatch, make_gpu_watermark_sampler
):
vllm_logger._print_warning_once.cache_clear()
sampler = make_gpu_watermark_sampler(StubWatermarker(), max_num_reqs=1)
messages: list[str] = []
monkeypatch.setattr(Sampler, "add_request", lambda *args, **kwargs: None)
monkeypatch.setattr(
"vllm.v1.watermarking.gpu_sampler.logger.warning_once", messages.append
"vllm.v1.watermarking.gpu_sampler.logger.warning",
lambda message, *_args, **_kwargs: messages.append(message),
)
sampler.add_request(0, SamplingParams(temperature=0))
sampler.add_request(0, SamplingParams(temperature=1))
sampler.add_request(0, SamplingParams(temperature=0, watermarking=True))
assert sampler.watermarking.np[0]
sampler.add_request(0, SamplingParams(temperature=0, watermarking=True))
assert sampler.watermarking.np[0]
sampler.add_request(0, SamplingParams(temperature=0, watermarking=False))
assert not sampler.watermarking.np[0]
assert messages == [
(
"Watermarking is enabled, but greedy decoding (temperature=0) cannot be "
"watermarked. This request will use ordinary greedy sampling."
"Watermarking is enabled, but greedy decoding "
"(temperature=0) cannot be watermarked. This request will use "
"ordinary greedy sampling."
)
]
vllm_logger._print_warning_once.cache_clear()
def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch):
sampler = object.__new__(GPUWatermarkSampler)
sampler.watermarker = StubWatermarker()
sampler.deduplicate_contexts = "single_turn"
sampler.watermarking = SimpleNamespace(
np=np.array([True, False]), gpu=torch.tensor([True, False])
)
sampler.sampling_states = SimpleNamespace(
temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)),
seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)),
)
sampler.use_fp64_gumbel = False
def test_gpu_sampler_respects_mixed_request_watermarking(
monkeypatch, make_gpu_watermark_sampler
):
sampler = make_gpu_watermark_sampler(StubWatermarker(), deduplicate_contexts="none")
sampler.add_request(0, SamplingParams(watermarking=True))
sampler.add_request(1, SamplingParams(watermarking=False))
sampler.apply_staged_writes()
sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
2, 1, dtype=torch.int64
)
sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: torch.zeros(
2, dtype=torch.bool
)
monkeypatch.setattr(
"vllm.v1.watermarking.watermarker.gumbel_sample",
lambda *args, **kwargs: torch.tensor([3, 4]),
@@ -322,17 +361,92 @@ def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch):
sampled, output_logits = sampler._sample_random(
logits,
torch.tensor([0, 1]),
np.array([0, 1]),
torch.tensor([1, 0]),
np.array([1, 0]),
torch.zeros(2, dtype=torch.int64),
None,
None,
False,
)
assert torch.equal(sampled, torch.tensor([7, 4]))
assert torch.equal(output_logits[0], torch.full((8,), 10.0))
assert torch.equal(output_logits[1], logits[1])
assert torch.equal(sampled, torch.tensor([3, 7]))
assert torch.equal(output_logits[0], logits[0])
assert torch.equal(output_logits[1], torch.full((8,), 10.0))
def test_gpu_sampler_filters_top_k_top_p_before_watermarking(
monkeypatch, make_gpu_watermark_sampler
):
# ``apply_top_k_top_p`` dispatches to a Triton kernel whenever Triton is
# installed, which cannot read this test's CPU tensors. Force the torch
# reference path; the assertion is about the call order, not the kernel.
monkeypatch.setattr(
"vllm.v1.watermarking.gpu_sampler.apply_top_k_top_p",
lambda logits, k, p: apply_top_k_top_p_pytorch(
logits, k, p, allow_cpu_sync=True
),
)
class CapturingWatermarker:
context_width = 1
captured_logits = None
def sample(self, logits, contexts, random_sampler=None, skip_mask=None):
self.captured_logits = logits
return WatermarkSample(logits.argmax(dim=-1), logits)
watermarker = CapturingWatermarker()
sampler = make_gpu_watermark_sampler(
watermarker, max_num_reqs=1, vocab_size=4, deduplicate_contexts="none"
)
sampler.add_request(0, SamplingParams(top_k=2, top_p=0.8, watermarking=True))
sampler.apply_staged_writes()
sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
1, 1, dtype=torch.int64
)
logits = torch.tensor([[5.0, 4.0, 3.0, 2.0]])
sampled, _ = sampler._sample_random(
logits,
torch.tensor([0]),
np.array([0]),
torch.zeros(1, dtype=torch.int64),
torch.tensor([2]),
torch.tensor([0.8]),
False,
)
assert watermarker.captured_logits is not None
assert torch.isneginf(watermarker.captured_logits[0, 2:]).all()
assert sampled.item() in (0, 1)
def test_watermark_context_ignores_prefix_cache_bookkeeping():
sampler = object.__new__(GPUWatermarkSampler)
sampler.watermarker = SimpleNamespace(context_width=3)
sampler.req_states = SimpleNamespace(
total_len=SimpleNamespace(gpu=torch.tensor([4, 0, 5])),
prompt_len=SimpleNamespace(gpu=torch.tensor([3, 0, 2])),
all_token_ids=SimpleNamespace(
gpu=torch.tensor(
[
[10, 11, 12, 40, 0, 0],
[0, 0, 0, 0, 0, 0],
[20, 21, 50, 51, 52, 0],
]
)
),
num_computed_tokens=SimpleNamespace(gpu=torch.tensor([0, 0, 0])),
)
request_slots = torch.tensor([2, 0])
cold_contexts = sampler._get_contexts(request_slots)
sampler.req_states.num_computed_tokens.gpu[:] = torch.tensor([3, 0, 2])
cached_contexts = sampler._get_contexts(request_slots)
expected = torch.tensor([[50, 51, 52], [-1, -1, 40]])
assert torch.equal(cold_contexts, expected)
assert torch.equal(cached_contexts, expected)
def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch):
+1
View File
@@ -119,6 +119,7 @@ def main(args: argparse.Namespace):
beam_width=args.n,
max_tokens=args.output_len,
ignore_eos=True,
watermarking=False,
),
)
+3 -1
View File
@@ -241,6 +241,7 @@ def benchmark_multimodal_processor(
temperature=0.0,
max_tokens=output_len,
detokenize=True,
watermarking=False,
)
for output_len in expected_output_lens
]
@@ -262,7 +263,8 @@ def benchmark_multimodal_processor(
warmup_prompts = [req.prompt for req in warmup_requests]
warmup_output_lens = [req.expected_output_len for req in warmup_requests]
warmup_sampling_params = [
SamplingParams(max_tokens=output_len) for output_len in warmup_output_lens
SamplingParams(max_tokens=output_len, watermarking=False)
for output_len in warmup_output_lens
]
llm.chat(
warmup_prompts,
+1
View File
@@ -178,6 +178,7 @@ def _run_vllm_requests(
beam_width=n,
max_tokens=output_len,
ignore_eos=True,
watermarking=False,
),
)
if do_profile:
+47 -7
View File
@@ -60,6 +60,7 @@ if TYPE_CHECKING:
from transformers import PreTrainedConfig
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.sampling_params import BeamSearchParams, SamplingParams
from vllm.v1.kv_cache_interface import KVCacheConfig
else:
PreTrainedConfig = Any
@@ -1314,15 +1315,22 @@ class VllmConfig:
if not self.use_v2_model_runner:
raise ValueError("trace replay requires Model Runner V2")
def _check_watermarking_unsupported(
def _check_supports_watermarking(
self,
config: "SamplingParams | BeamSearchParams | None" = None,
*,
beam_search: bool = False,
custom_sampler: bool = False,
) -> None:
) -> bool:
watermark_config = getattr(self, "watermark_config", None)
if watermark_config is None:
return
if config is not None and config.watermarking is not False:
logger.warning_once(
"Watermarking is enabled for this request, but the engine has no "
"watermark configuration. This and subsequent requests will run "
"without watermarking.",
scope="global",
)
return False
if self.speculative_config is not None:
speculative_config = self.speculative_config
if speculative_config.draft_sample_method != "probabilistic":
@@ -1377,12 +1385,44 @@ class VllmConfig:
watermark_config.alpha,
scope="global",
)
if beam_search:
raise ValueError("Beam search is not supported with watermarking.")
if custom_sampler:
raise ValueError(
"Model-specific custom samplers are not supported with watermarking."
)
if config is None:
return True
if config.watermarking is False:
return False
from vllm.sampling_params import BeamSearchParams, SamplingParams
if isinstance(config, BeamSearchParams):
logger.warning_once(
"Watermarking is enabled, but beam search cannot be watermarked. "
"This and subsequent beam search requests will run without "
"watermarking.",
scope="global",
)
return False
if not isinstance(config, SamplingParams):
raise TypeError(f"Unsupported watermarking config: {type(config).__name__}")
if config.trace_decode_token_ids is not None:
logger.warning_once(
"Watermarking is enabled, but trace replay cannot be watermarked. "
"This and subsequent trace replay requests will run without "
"watermarking.",
scope="global",
)
return False
if config.temperature == 0:
logger.warning_once(
"Watermarking is enabled, but greedy decoding "
"(temperature=0) cannot be watermarked. This and subsequent "
"greedy requests will use ordinary greedy sampling.",
scope="global",
)
return False
return True
def _resolve_and_verify_engram_config(self) -> None:
"""Resolve defaults and validate n-gram embedding settings."""
@@ -1446,7 +1486,7 @@ class VllmConfig:
self.try_verify_and_update_config()
self._resolve_and_verify_engram_config()
self._check_watermarking_unsupported()
self._check_supports_watermarking()
# Models may have supplied their own DCP defaults above; anything still
# unset falls back to the stock ones.
self.parallel_config.set_dcp_defaults()
+6
View File
@@ -242,6 +242,12 @@ class AnthropicMessagesRequest(BaseModel):
"to 256 bit)."
),
)
watermarking: bool | None = Field(
default=None,
description=(
"Whether to apply the engine's configured watermark to this request."
),
)
kv_transfer_params: dict[str, Any] | None = Field(
default=None,
description="KVTransfer parameters used for disaggregated serving.",
+1
View File
@@ -613,6 +613,7 @@ class AnthropicServingMessages(OpenAIServingChat):
temperature=anthropic_request.temperature,
top_p=anthropic_request.top_p,
top_k=anthropic_request.top_k,
watermarking=anthropic_request.watermarking,
cache_salt=anthropic_request.cache_salt,
kv_transfer_params=anthropic_request.kv_transfer_params,
ec_transfer_params=anthropic_request.ec_transfer_params,
+6
View File
@@ -175,6 +175,12 @@ class CohereChatV2Request(BaseModel):
# vLLM-specific extensions (not in Cohere spec). These mirror what the
# Anthropic and OpenAI surfaces already expose so V2 callers can reach
# the same engine knobs when needed.
watermarking: bool | None = Field(
default=None,
description=(
"Whether to apply the engine's configured watermark to this request."
),
)
kv_transfer_params: dict[str, Any] | None = Field(
default=None,
description="KVTransfer parameters used for disaggregated serving.",
+1
View File
@@ -557,6 +557,7 @@ class CohereServingChatV2(OpenAIServingChat):
temperature=request.temperature,
top_p=request.p,
top_k=request.k,
watermarking=request.watermarking,
seed=request.seed,
frequency_penalty=request.frequency_penalty,
presence_penalty=request.presence_penalty,
@@ -82,7 +82,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
temperature = params.temperature
ignore_eos = params.ignore_eos
length_penalty = params.length_penalty
self.llm_engine.vllm_config._check_watermarking_unsupported(beam_search=True)
self.llm_engine.input_processor.resolve_watermarking(params)
tokenizer = self.renderer.get_tokenizer()
eos_token_id = tokenizer.eos_token_id
@@ -120,6 +120,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
logprobs=2 * beam_width,
max_tokens=1,
temperature=temperature,
watermarking=False,
detokenize=False,
skip_clone=True, # Internal beam search, safe to skip clone
)
@@ -457,6 +458,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
logprobs=base_params.logprobs,
max_tokens=1,
temperature=base_params.temperature,
watermarking=base_params.watermarking,
detokenize=False,
allowed_token_ids=(
allowed_ids
@@ -40,6 +40,7 @@ class BeamSearchOnlineMixin(ABC):
temperature = params.temperature
length_penalty = params.length_penalty
include_stop_str_in_output = params.include_stop_str_in_output
self.engine_client.input_processor.resolve_watermarking(params)
tokenizer = self.renderer.get_tokenizer()
eos_token_id = tokenizer.eos_token_id
@@ -62,6 +63,7 @@ class BeamSearchOnlineMixin(ABC):
logprobs=logprobs_num,
max_tokens=1,
temperature=temperature,
watermarking=False,
detokenize=False,
)
all_beams = [
@@ -261,6 +261,7 @@ class ServingGenerativeScoring(BaseServing):
logprobs=len(request.label_token_ids),
logprob_token_ids=request.label_token_ids,
n=1,
watermarking=False,
)
# Get trace headers
@@ -268,7 +268,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
top_k: int | None = None
min_p: float | None = None
repetition_penalty: float | None = None
watermarking: bool = True
watermarking: bool | None = None
length_penalty: float = 1.0
stop_token_ids: list[int] | None = []
include_stop_str_in_output: bool = False
@@ -647,6 +647,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
max_tokens=max_tokens,
ignore_eos=self.ignore_eos,
temperature=temperature,
watermarking=self.watermarking,
length_penalty=self.length_penalty,
include_stop_str_in_output=self.include_stop_str_in_output,
skip_special_tokens=self.skip_special_tokens,
@@ -1100,6 +1101,7 @@ class BatchChatCompletionRequest(OpenAIBaseModel):
top_k: int | None = None
min_p: float | None = None
repetition_penalty: float | None = None
watermarking: bool | None = None
length_penalty: float | None = 1.0
early_stopping: bool = False
structured_outputs: StructuredOutputsParams | None = None
@@ -75,7 +75,7 @@ class CompletionRequest(OpenAIBaseModel):
top_k: int | None = None
min_p: float | None = None
repetition_penalty: float | None = None
watermarking: bool = True
watermarking: bool | None = None
length_penalty: float = 1.0
stop_token_ids: list[int] | None = []
include_stop_str_in_output: bool = False
@@ -303,6 +303,7 @@ class CompletionRequest(OpenAIBaseModel):
max_tokens=max_tokens,
ignore_eos=self.ignore_eos,
temperature=temperature,
watermarking=self.watermarking,
length_penalty=self.length_penalty,
include_stop_str_in_output=self.include_stop_str_in_output,
skip_special_tokens=self.skip_special_tokens,
@@ -245,7 +245,7 @@ class ResponsesRequest(OpenAIBaseModel):
)
# --8<-- [start:responses-extra-params]
watermarking: bool = True
watermarking: bool | None = None
request_id: str = Field(
default_factory=lambda: f"resp_{random_uuid()}",
description=(
@@ -229,6 +229,7 @@ class SpeechToTextBaseServing(GenerateBaseServing):
max_tokens=1,
temperature=0.0,
allowed_token_ids=allowed_token_ids,
watermarking=False,
)
result_generator = self.engine_client.generate(
@@ -151,6 +151,9 @@ class TranscriptionRequest(OpenAIBaseModel):
to automatically increase the temperature until certain thresholds are hit.
"""
watermarking: bool | None = None
"""Whether to apply the engine's configured watermark to this request."""
top_p: float | None = None
"""Enables nucleus (top-p) sampling, where tokens are selected from the
smallest possible set whose cumulative probability exceeds `p`.
@@ -226,6 +229,7 @@ class TranscriptionRequest(OpenAIBaseModel):
beam_width=n,
max_tokens=max_tokens,
temperature=temperature,
watermarking=self.watermarking,
length_penalty=self.length_penalty,
include_stop_str_in_output=self.include_stop_str_in_output,
)
@@ -264,6 +268,7 @@ class TranscriptionRequest(OpenAIBaseModel):
return SamplingParams.from_optional(
temperature=temperature,
watermarking=self.watermarking,
max_tokens=max_tokens,
seed=self.seed,
top_p=top_p,
@@ -102,6 +102,9 @@ class TranslationRequest(OpenAIBaseModel):
to automatically increase the temperature until certain thresholds are hit.
"""
watermarking: bool | None = None
"""Whether to apply the engine's configured watermark to this request."""
top_p: float | None = None
"""Enables nucleus (top-p) sampling, where tokens are selected from the
smallest possible set whose cumulative probability exceeds `p`.
@@ -215,6 +218,7 @@ class TranslationRequest(OpenAIBaseModel):
beam_width=n,
max_tokens=max_tokens,
temperature=temperature,
watermarking=self.watermarking,
length_penalty=self.length_penalty,
include_stop_str_in_output=self.include_stop_str_in_output,
)
@@ -251,6 +255,7 @@ class TranslationRequest(OpenAIBaseModel):
return SamplingParams.from_optional(
temperature=temperature,
watermarking=self.watermarking,
max_tokens=max_tokens,
seed=self.seed,
top_p=top_p,
+7 -3
View File
@@ -251,8 +251,9 @@ class SamplingParams(
"""Controls the randomness of the sampling. Lower values make the model
more deterministic, while higher values make the model more random. Zero
means greedy sampling."""
watermarking: bool = True
"""Whether to apply the engine's configured watermark to this request."""
watermarking: bool | None = None
"""Whether to apply the engine's configured watermark to this request.
``None`` defaults to enabled."""
top_p: float = 1.0
"""Controls the cumulative probability of the top tokens to consider. Must
be in (0, 1]. Set to 1 to consider all tokens."""
@@ -342,6 +343,7 @@ class SamplingParams(
output_text_buffer_length: int = 0
_eos_token_id: int | None = None
_all_stop_token_ids: set[int] = msgspec.field(default_factory=set)
_watermarking_skipped: bool = False
# Fields used to construct logits processors
structured_outputs: StructuredOutputsParams | None = None
@@ -395,7 +397,7 @@ class SamplingParams(
frequency_penalty: float | None = 0.0,
repetition_penalty: float | None = 1.0,
temperature: float | None = 1.0,
watermarking: bool = True,
watermarking: bool | None = None,
top_p: float | None = 1.0,
top_k: int = 0,
min_p: float = 0.0,
@@ -1411,6 +1413,7 @@ class SamplingParams(
_bad_words_token_ids=[[0], [1, 2]],
logprobs=5,
prompt_logprobs=1,
watermarking=True,
)
@@ -1430,6 +1433,7 @@ class BeamSearchParams(
include_stop_str_in_output: bool = False
structured_outputs: StructuredOutputsParams | None = None
skip_special_tokens: bool = True
watermarking: bool | None = None
def __post_init__(self) -> None:
_verify_num_sequences(self.beam_width, "beam_width")
+8
View File
@@ -443,6 +443,13 @@ class AsyncLLM(EngineClient):
"does not match the EngineCoreRequest.request_id attribute. The "
"latter will be used, and the former will be ignored."
)
request_params = request.params
if isinstance(request_params, SamplingParams):
# This request object is owned by the engine from here on.
self.input_processor.apply_watermarking(
request_params,
self.input_processor.resolve_watermarking(request_params),
)
else:
if isinstance(prompt, dict) and "type" in prompt:
# Rendered EngineInput; no blocking preprocessing needed.
@@ -910,6 +917,7 @@ class AsyncLLM(EngineClient):
sampling_params=SamplingParams(
max_tokens=1,
extra_args={"kv_transfer_params": dict(kv_transfer_params)},
watermarking=False,
),
pooling_params=None,
arrival_time=time.time(),
+42 -4
View File
@@ -25,7 +25,7 @@ from vllm.platforms import current_platform
from vllm.pooling_params import PoolingParams
from vllm.renderers import BaseRenderer, renderer_from_config
from vllm.renderers.inputs.preprocess import parse_model_prompt
from vllm.sampling_params import SamplingParams
from vllm.sampling_params import BeamSearchParams, SamplingParams
from vllm.tasks import GENERATION_TASKS, POOLING_TASKS, SupportedTask
from vllm.tokenizers import TokenizerLike
from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid
@@ -112,12 +112,41 @@ class InputProcessor:
return partial(validate_logits_processors_parameters, custom_logitsprocs)
def resolve_watermarking(self, params: SamplingParams | BeamSearchParams) -> bool:
"""Return whether watermarking will be applied to this request.
``None`` and ``True`` request watermarking, while ``False`` opts out.
If watermarking is unavailable or incompatible with the request, emit
a warning and return ``False``.
"""
return self.vllm_config._check_supports_watermarking(params)
def apply_watermarking(self, params: SamplingParams, watermarking: bool) -> None:
"""Store the resolved `watermarking` on request-local params.
Also records whether watermarking was requested on a watermarking
engine but could not be applied, as opposed to an explicit opt-out.
"""
# Sticky, so resolving an already resolved request keeps the flag.
params._watermarking_skipped = params._watermarking_skipped or (
self.vllm_config.watermark_config is not None
and params.watermarking is not False
and not watermarking
)
params.watermarking = watermarking
def _validate_params(
self,
params: SamplingParams | PoolingParams,
supported_tasks: tuple[SupportedTask, ...],
) -> None:
"""Raise `ValueError` if SamplingParams or PoolingParams is not valid."""
) -> bool | None:
"""Raise `ValueError` if SamplingParams or PoolingParams is not valid.
Returns the resolved `watermarking` value for `SamplingParams` (`None`
for `PoolingParams`). The caller applies it to its private copy of the
params so that a caller-owned params object is never mutated.
"""
watermarking: bool | None = None
if isinstance(params, SamplingParams):
supported_generation_tasks = [
task for task in supported_tasks if task in GENERATION_TASKS
@@ -189,6 +218,9 @@ class InputProcessor:
"enabled. Start the engine with --enable-trace-replay "
"to use it."
)
# Resolved last so that feature gates above report the missing
# feature rather than its incompatibility with watermarking.
watermarking = self.resolve_watermarking(params)
elif isinstance(params, PoolingParams):
supported_pooling_tasks = [
task for task in supported_tasks if task in POOLING_TASKS
@@ -217,6 +249,8 @@ class InputProcessor:
f"but got {type(params).__name__}"
)
return watermarking
def _normalize_trace_replay_params(
self, sampling_params: SamplingParams, prompt_len: int
) -> None:
@@ -352,7 +386,7 @@ class InputProcessor:
session_id: str | None = None,
kv_hints: KvHintsEnvelope | None = None,
) -> EngineCoreRequest:
self._validate_params(params, supported_tasks)
watermarking = self._validate_params(params, supported_tasks)
self._validate_lora(lora_request)
parallel_config = self.vllm_config.parallel_config
@@ -413,6 +447,10 @@ class InputProcessor:
if isinstance(params, SamplingParams):
# TODO: can we avoid cloning here in multiproc case?
sampling_params = params.clone()
# Resolve on the request-local copy: `params` may be shared across
# prompts, requests and even engines by the caller.
assert watermarking is not None
self.apply_watermarking(sampling_params, watermarking)
prompt_len = length_from_prompt_token_ids_or_embeds(
prompt_token_ids, prompt_embeds
)
+7
View File
@@ -254,6 +254,13 @@ class LLMEngine:
"does not match the EngineCoreRequest.request_id attribute. The "
"latter will be used, and the former will be ignored."
)
request_params = request.params
if isinstance(request_params, SamplingParams):
# This request object is owned by the engine from here on.
self.input_processor.apply_watermarking(
request_params,
self.input_processor.resolve_watermarking(request_params),
)
else:
request = self.input_processor.process_inputs(
request_id,
+1 -1
View File
@@ -221,7 +221,7 @@ class DualKeyGumbelWatermarkDetector(GumbelWatermarkDetector):
p_value_threshold: float = 0.01,
prf: WatermarkPRFName = "philox",
deduplicate_contexts: bool = True,
alpha: float = 0.2,
alpha: float = 0.1,
) -> None:
if not 0 <= alpha <= 1:
raise ValueError("alpha must be between 0 and 1")
+1 -1
View File
@@ -482,7 +482,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
custom = self.model_state.custom_sampler(self.sampler)
if custom:
self.vllm_config._check_watermarking_unsupported(custom_sampler=True)
self.vllm_config._check_supports_watermarking(custom_sampler=True)
self.sampler, self.rejection_sampler = custom
elif self.speculative_config is not None:
self.rejection_sampler = RejectionSampler(
+5 -1
View File
@@ -139,7 +139,11 @@ def run_mixed_prefill_decode_warmup(
next_block_id += num_blocks
return block_ids
sampling_params = SamplingParams(max_tokens=2, temperature=0.0)
sampling_params = SamplingParams(
max_tokens=2,
temperature=0.0,
watermarking=False,
)
decode_prefill_output = SchedulerOutput.make_empty()
decode_prefill_output.scheduled_new_reqs = [