| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | from vllm import SamplingParams |
| 7 | from vllm.platforms import current_platform |
| 8 | from vllm.v1.watermarking import GumbelWatermarkDetector |
| 9 | |
| 10 | |
| 11 | @pytest.mark.skipif( |
| 12 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 13 | ) |
| 14 | def test_llm_generated_sequence_is_watermarked(vllm_runner): |
| 15 | """Engine-level watermarking must produce tokens the detector recognizes.""" |
| 16 | watermark_config = {"key": 42, "context_width": 4} |
| 17 | detector = GumbelWatermarkDetector(**watermark_config) |
| 18 | |
| 19 | runner = vllm_runner( |
| 20 | "facebook/opt-125m", |
| 21 | watermark_config=watermark_config, |
| 22 | seed=0, |
| 23 | enforce_eager=True, |
| 24 | max_model_len=512, |
| 25 | gpu_memory_utilization=0.2, |
| 26 | ) |
| 27 | request = "Tell me a story about an explorer who discovers a mysterious island." |
| 28 | params_use_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=True) |
| 29 | params_no_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=False) |
| 30 | |
| 31 | with runner: |
| 32 | output_use_wm = runner.llm.generate(request, params_use_wm) |
| 33 | output_no_wm = runner.llm.generate(request, params_no_wm) |
| 34 | |
| 35 | result_use_wm = detector.detect(list(output_use_wm[0].outputs[0].token_ids)) |
| 36 | result_no_wm = detector.detect(list(output_no_wm[0].outputs[0].token_ids)) |
| 37 | |
| 38 | assert result_use_wm.is_watermarked, result_use_wm |
| 39 | recorded_p_value = 9.83e-20 |
| 40 | assert result_use_wm.p_value < recorded_p_value * 10 # tolerance |
| 41 | |
| 42 | assert not result_no_wm.is_watermarked, result_no_wm |