From c32495d58ff580b518cf4c5e5d6b69adc4d38351 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Rialland?= <36076211+TQCB@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:37:13 +0200 Subject: [PATCH] [watermark] compatibility validation (#56801) Signed-off-by: Raphael Rialland Signed-off-by: Simon Veitner Co-authored-by: Simon Veitner --- .buildkite/test_areas/lm_eval.yaml | 4 + docs/features/watermarking.md | 53 ++- .../test_anthropic_messages_conversion.py | 27 ++ .../cohere/test_serving_conversion.py | 21 ++ .../test_generative_scoring.py | 5 + .../test_completion_with_function_calling.py | 37 +- tests/entrypoints/openai/test_watermarking.py | 231 +++++++++++- .../realtime/test_realtime_validation.py | 42 +++ .../test_speech_to_text_cancellation.py | 2 + tests/samplers/test_beam_search_online.py | 175 ++++++++- tests/test_config.py | 54 ++- tests/test_sampling_params.py | 1 + .../test_input_processor_trace_replay.py | 7 +- tests/v1/engine/test_kv_transfer_rejection.py | 27 ++ tests/v1/sample/test_logprobs.py | 1 + tests/v1/worker/test_mixed_warmup_gate.py | 34 ++ tests/watermarking/test_compatibility.py | 353 ++++++++++++++++++ .../watermarking/test_engine_compatibility.py | 204 ++++++++++ tests/watermarking/test_gumbel.py | 25 +- .../test_parallel_compatibility.py | 106 ++++++ tests/watermarking/test_watermarking.py | 180 +++++++-- vllm/benchmarks/latency.py | 1 + vllm/benchmarks/mm_processor.py | 4 +- vllm/benchmarks/throughput.py | 1 + vllm/config/vllm.py | 54 ++- vllm/entrypoints/anthropic/protocol.py | 6 + vllm/entrypoints/anthropic/serving.py | 1 + vllm/entrypoints/cohere/protocol.py | 6 + vllm/entrypoints/cohere/serving.py | 1 + .../generate/beam_search/offline.py | 4 +- .../generate/beam_search/online.py | 2 + .../generate/generative_scoring/serving.py | 1 + .../openai/chat_completion/protocol.py | 4 +- .../entrypoints/openai/completion/protocol.py | 3 +- vllm/entrypoints/openai/responses/protocol.py | 2 +- .../speech_to_text/base/serving.py | 1 + .../speech_to_text/transcription/protocol.py | 5 + .../speech_to_text/translation/protocol.py | 5 + vllm/sampling_params.py | 10 +- vllm/v1/engine/async_llm.py | 8 + vllm/v1/engine/input_processor.py | 46 ++- vllm/v1/engine/llm_engine.py | 7 + vllm/v1/watermarking/gumbel.py | 2 +- vllm/v1/worker/gpu/model_runner.py | 2 +- vllm/v1/worker/gpu/warmup.py | 6 +- 45 files changed, 1670 insertions(+), 101 deletions(-) create mode 100644 tests/v1/engine/test_kv_transfer_rejection.py create mode 100644 tests/watermarking/test_compatibility.py create mode 100644 tests/watermarking/test_engine_compatibility.py create mode 100644 tests/watermarking/test_parallel_compatibility.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 04337d03d2..bcb1110848 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -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" diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index 7ab02d3363..66b6465eb8 100644 --- a/docs/features/watermarking.md +++ b/docs/features/watermarking.md @@ -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. diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index fd0d0dc2a9..1bccee4601 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -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 diff --git a/tests/entrypoints/cohere/test_serving_conversion.py b/tests/entrypoints/cohere/test_serving_conversion.py index 5cbe13df02..58fa7fe793 100644 --- a/tests/entrypoints/cohere/test_serving_conversion.py +++ b/tests/entrypoints/cohere/test_serving_conversion.py @@ -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 diff --git a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py index f413869c90..69be10b19d 100644 --- a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py +++ b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py @@ -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 diff --git a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py index 61fd2f1576..72442670a8 100644 --- a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py +++ b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py @@ -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( diff --git a/tests/entrypoints/openai/test_watermarking.py b/tests/entrypoints/openai/test_watermarking.py index 0425b86a02..e6a750eeb7 100644 --- a/tests/entrypoints/openai/test_watermarking.py +++ b/tests/entrypoints/openai/test_watermarking.py @@ -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 diff --git a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py index c74246b2f8..a8c4f35c04 100644 --- a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py +++ b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py @@ -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 diff --git a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py index 0ff8adb6b5..8ed31e4cdb 100644 --- a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py +++ b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py @@ -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) diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index cbed58363a..22afce2d78 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -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 diff --git a/tests/test_config.py b/tests/test_config.py index b1348480cc..9d4567874b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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(): diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index a6e057ad7a..f32a3d4c36 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -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",)) diff --git a/tests/v1/engine/test_input_processor_trace_replay.py b/tests/v1/engine/test_input_processor_trace_replay.py index 957f70b718..043de170af 100644 --- a/tests/v1/engine/test_input_processor_trace_replay.py +++ b/tests/v1/engine/test_input_processor_trace_replay.py @@ -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, diff --git a/tests/v1/engine/test_kv_transfer_rejection.py b/tests/v1/engine/test_kv_transfer_rejection.py new file mode 100644 index 0000000000..a21351900c --- /dev/null +++ b/tests/v1/engine/test_kv_transfer_rejection.py @@ -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 diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 4152662fbd..813bea0bf3 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -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"): diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index 04f456a7c9..84ddf692a1 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -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.""" diff --git a/tests/watermarking/test_compatibility.py b/tests/watermarking/test_compatibility.py new file mode 100644 index 0000000000..748223d94b --- /dev/null +++ b/tests/watermarking/test_compatibility.py @@ -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 diff --git a/tests/watermarking/test_engine_compatibility.py b/tests/watermarking/test_engine_compatibility.py new file mode 100644 index 0000000000..1320306cb3 --- /dev/null +++ b/tests/watermarking/test_engine_compatibility.py @@ -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 diff --git a/tests/watermarking/test_gumbel.py b/tests/watermarking/test_gumbel.py index d789692a0e..aa0ffa2799 100644 --- a/tests/watermarking/test_gumbel.py +++ b/tests/watermarking/test_gumbel.py @@ -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( diff --git a/tests/watermarking/test_parallel_compatibility.py b/tests/watermarking/test_parallel_compatibility.py new file mode 100644 index 0000000000..6e47a65b6f --- /dev/null +++ b/tests/watermarking/test_parallel_compatibility.py @@ -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() diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index 4de6022e34..384cc30d57 100644 --- a/tests/watermarking/test_watermarking.py +++ b/tests/watermarking/test_watermarking.py @@ -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): diff --git a/vllm/benchmarks/latency.py b/vllm/benchmarks/latency.py index acb3efef15..0c75c04563 100644 --- a/vllm/benchmarks/latency.py +++ b/vllm/benchmarks/latency.py @@ -119,6 +119,7 @@ def main(args: argparse.Namespace): beam_width=args.n, max_tokens=args.output_len, ignore_eos=True, + watermarking=False, ), ) diff --git a/vllm/benchmarks/mm_processor.py b/vllm/benchmarks/mm_processor.py index 40b3494151..7a52faf29d 100644 --- a/vllm/benchmarks/mm_processor.py +++ b/vllm/benchmarks/mm_processor.py @@ -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, diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index 620ce2b08f..d7c3c01454 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -178,6 +178,7 @@ def _run_vllm_requests( beam_width=n, max_tokens=output_len, ignore_eos=True, + watermarking=False, ), ) if do_profile: diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 53d4102697..dec1f27fe1 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -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() diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 5abfaffed3..2a14c62933 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -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.", diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index cf5fd9ecc1..05c1e24f2a 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -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, diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py index 1432e6eb11..45e0d310dc 100644 --- a/vllm/entrypoints/cohere/protocol.py +++ b/vllm/entrypoints/cohere/protocol.py @@ -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.", diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 63d00f0c57..bdaa098659 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -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, diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index b7095131af..1ad451ac13 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -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 diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index ac7f0f884f..5a95f40d31 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -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 = [ diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index cb0c4995bf..4d5a15c2f5 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -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 diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 228472ab6b..6a70624be4 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -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 diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index a0974cb91c..78722b45ec 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -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, diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 9c44a04b3a..8fa96ecdb0 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -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=( diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index 34b7e9c249..871f69666f 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -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( diff --git a/vllm/entrypoints/speech_to_text/transcription/protocol.py b/vllm/entrypoints/speech_to_text/transcription/protocol.py index e9570fb926..7e66c512c7 100644 --- a/vllm/entrypoints/speech_to_text/transcription/protocol.py +++ b/vllm/entrypoints/speech_to_text/transcription/protocol.py @@ -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, diff --git a/vllm/entrypoints/speech_to_text/translation/protocol.py b/vllm/entrypoints/speech_to_text/translation/protocol.py index ac8c69fdc2..91ec56e2fb 100644 --- a/vllm/entrypoints/speech_to_text/translation/protocol.py +++ b/vllm/entrypoints/speech_to_text/translation/protocol.py @@ -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, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 246e2ce90c..ab9f271114 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -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") diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 4a7a91cb8e..37025a8632 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -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(), diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 8b5a69bcf4..17e8af5cae 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -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 ) diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index c3fb3efe7e..509d46703d 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -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, diff --git a/vllm/v1/watermarking/gumbel.py b/vllm/v1/watermarking/gumbel.py index 655267d1e4..2f953dd389 100644 --- a/vllm/v1/watermarking/gumbel.py +++ b/vllm/v1/watermarking/gumbel.py @@ -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") diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1b9d0608e2..b9880f43f2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -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( diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 2af97207ce..c29d9958fd 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -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 = [