[watermarking hardening] Add e2e test and basic GSM8K quality tests (#56309)
Signed-off-by: Tomas Ruiz <tomas.ruiz.te@gmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Tomas Ruiz committed
Sep 14, 2026 at 17:29 UTC
b7e8dd8f372c1324c376ba5dc03d92264eae8f24
4 files changed
+81
.buildkite/test_areas/lm_eval.yaml
+16
@@ -30,6 +30,22 @@ steps:
30
- vllm/_aiter_ops.py
31
- vllm/platforms/rocm.py
32
33
+- label: ":nvidia: (H200 MIG 35GB) LM Eval Watermarking"
34
+ device: h200_35gb
35
+ key: lm-eval-watermarking
36
+ # Recorded runs took 80–140 seconds
37
+ timeout_in_minutes: 10
38
+ source_file_dependencies:
39
+ - .buildkite/test_areas/lm_eval.yaml
40
+ - tests/evals/gsm8k/
41
+ - vllm/config/watermarking.py
42
+ - vllm/sampling_params.py
43
+ - vllm/v1/watermarking/
44
+ - vllm/v1/worker/gpu/
45
+ autorun_on_main: true
46
+ commands:
47
+ - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt
48
+
49
# - label: LM Eval Large Models (4xA100)
50
# key: lm-eval-large-models-4xa100
51
# device: a100
tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml
new
+22
@@ -0,0 +1,22 @@
1
+model_name: "Qwen/Qwen3.5-4B"
2
+# full GSM8K set is 1319 questions
3
+num_questions: 1319
4
+num_fewshot: 5
5
+max_concurrency: 128
6
+max_tokens: 1024
7
+use_chat_completions: true
8
+temperature: 0.7
9
+seed: 42
10
+
11
+# Calibrated by running non-watermaked 13 times
12
+# the threshold is the average accuracy, and the tolerance is set to 3 sigma (standard deviations)
13
+accuracy_threshold: 0.885
14
+tolerance: 0.015
15
+
16
+server_args: >-
17
+ --language-model-only
18
+ --max-model-len 8192
19
+ --max-num-seqs 128
20
+ --gpu-memory-utilization 0.4
21
+ --default-chat-template-kwargs '{"enable_thinking": false}'
22
+ --watermark-config '{"key": 42, "context_width": 4}'
tests/evals/gsm8k/configs/models-watermark.txt
new
+1
@@ -0,0 +1 @@
1
+Qwen3.5-4B-watermark.yaml
tests/watermarking/test_watermarking_e2e.py
new
+42
@@ -0,0 +1,42 @@
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