[watermarking] Dual-key gumbel-max watermarking for speculative decoding support (#56122)
Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai> Signed-off-by: Simon Veitner <sveitner@redhat.com> Signed-off-by: Tomas Ruiz <tomas.ruiz.te@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: Tomas Ruiz <tomas.ruiz.te@gmail.com>
Raphaël Rialland committed
Sep 12, 2026 at 20:42 UTC
7ee8a6dd013819838da8012ca549d724bee7c6c6
28 files changed
+3451
-236
.buildkite/test_areas/lm_eval.yaml
+22
@@ -122,6 +122,28 @@ steps:
122
commands:
123
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
124
125
+- label: ":nvidia: (H100) LM Eval Watermark Feature Combination"
126
+ key: lm-eval-dspark-watermark-2xh100
127
+ device: h100
128
+ num_devices: 2
129
+ # Local TP2 runs took 84–379 seconds, including startup.
130
+ timeout_in_minutes: 15
131
+ source_file_dependencies:
132
+ - .buildkite/test_areas/lm_eval.yaml
133
+ - tests/evals/gsm8k/
134
+ - vllm/config/speculative.py
135
+ - vllm/config/vllm.py
136
+ - vllm/config/watermarking.py
137
+ - vllm/sampling_params.py
138
+ - vllm/model_executor/models/qwen3_dspark.py
139
+ - vllm/model_executor/models/qwen3_dflash.py
140
+ - vllm/model_executor/warmup/
141
+ - vllm/v1/watermarking/
142
+ - vllm/v1/worker/gpu/
143
+ autorun_on_main: true
144
+ commands:
145
+ - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
146
+
147
- label: ":nvidia: (B200) LM Eval Spec Decode"
148
key: lm-eval-spec-decode-4xb200
149
timeout_in_minutes: 120
docs/features/watermarking.md
+110
-10
@@ -34,11 +34,17 @@ Deployments that require watermarking must restrict this field to trusted
34
callers, or strip and validate it at the ingress boundary, so untrusted clients
35
cannot opt out.
36
37
-`context_width` controls how many prior output tokens seed each watermark
38
-decision and defaults to 4. Larger values make the watermark less robust to
37
+`context_width` controls how many prior tokens seed each watermark decision
38
+and defaults to 4. Larger values make the watermark less robust to
39
edits because an insertion, deletion, or substitution changes more subsequent
40
contexts. Values above 16 are allowed but emit a warning.
41
42
+`allow_target_only_watermarking` defaults to false and only has an effect when
43
+speculative decoding is enabled. It permits speculative decoding with a
44
+watermarking algorithm that does not support it natively, at the cost of
45
+weaker detectability. See
46
+[Speculative decoding](#speculative-decoding).
47
+
48
## Architecture
49
50
`WatermarkConfig` selects an algorithm and PRF. Model Runner V2 constructs the
@@ -52,21 +58,108 @@ the reference algorithms. `WatermarkDetector` consumes token IDs, so callers
58
remain responsible for using the tokenizer and watermark profile that match
59
generation.
60
61
+## Speculative decoding
62
+
63
+Watermarking requires speculative decoding to use probabilistic draft sampling,
64
+standard rejection sampling, and an autoregressive model-based method (`dspark`,
65
+`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dspark`.
66
+
67
+A watermarking algorithm without native speculative-decoding support is
68
+rejected before model loading. Set
69
+`"allow_target_only_watermarking": true` to allow it: accepted draft tokens are
70
+not watermarked, while target-side rejection recovery and bonus sampling remain
71
+watermarked. The watermark signal is diluted in proportion to the share of
72
+output tokens supplied by accepted drafts; rejected drafts do not dilute it
73
+because their recovery tokens are watermarked.
74
+
75
+Speculative-decoding token paths do not currently support generation-side
76
+context deduplication. The configured `deduplicate_contexts` policy is not
77
+applied to accepted drafts, rejection-recovery tokens, or bonus tokens.
78
+
79
+For `dual_key_gumbel`, `alpha` has no effect under speculative decoding. The
80
+speculative protocol selects the key for each token instead.
81
+
82
## Algorithms
83
84
### Gumbel-max
85
59
-Gumbel-max derives a deterministic pseudorandom value from the key, prior
60
-generated-token context, and every candidate token, then uses the resulting
61
-Gumbel noise for categorical sampling. See
86
+Gumbel-max derives a deterministic pseudorandom value from the key, prior token
87
+context, and every candidate token, then uses the resulting Gumbel noise for
88
+categorical sampling. See
89
[Aaronson's original presentation](https://simons.berkeley.edu/sites/default/files/2024-10/LLM24-2%20Slides%20-%20Scott%20Aaronson.pdf).
90
91
Gumbel-max requires stochastic sampling. Greedy requests (`temperature=0`)
92
bypass watermarking and emit a warning once per worker.
93
67
-Single-key Gumbel-max watermarking may increase degenerate generations. In
68
-particular, reusing the same keyed random vector when a context repeats can
69
-reinforce repetition loops that continue until the maximum token limit.
94
+When a token context is repeated, generation can use ordinary sampling for that
95
+occurrence. Reusing the repeated context leads to a bias over the sequence, as
96
+certain token choices would be correlated. Ordinary sampling at these positions
97
+allows for single-sequence non-distortion (see section G.3 of the
98
+[SynthID-Text supplementary materials](https://media.springernature.com/original/springer-static/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf)).
99
+The detector independently deduplicates contexts so repeated keyed random
100
+vectors are not treated as independent evidence, meaning that context
101
+deduplication at generation time does not reduce the watermarking signal, unless
102
+a user tampers with the output and removes the first occurrence while leaving an
103
+unwatermarked repetition.
104
+
105
+`deduplicate_contexts` controls which history is searched for a repeated
106
+context:
107
+
108
+- `"none"` disables context deduplication.
109
+- `"single_turn"`, the default, searches the tokens generated for the current
110
+ request. This gives single-turn non-distortion.
111
+- `"all"` also searches the prompt, which extends non-distortion across the
112
+ turns of a conversation. The first `context_width` generated tokens use
113
+ ordinary sampling, and subsequent tokens whose context already occurs in the
114
+ prompt are not watermarked. Previous turns of a session, or a prompt that
115
+ already contains the structure of the answer such as a tool result the model
116
+ extends, may therefore leave little of the answer watermarked. Given the
117
+ impact of `"all"` on both performance (sampling steps must scan over a
118
+ potentially large amount of previous context) and watermarking detectability
119
+ (we may significantly reduce the amount of watermarked tokens), you should
120
+ only use `"all"` when non-distortion across multiple turns is required.
121
+
122
+`deduplicate_contexts_max_history` limits the search to the most recent
123
+positions and defaults to 8,192. Each position is compared over the
124
+`context_width` tokens before it, so the window of tokens read is that much
125
+longer. A smaller value reduces scanning cost but only provides the guarantee
126
+within that window. Set it to `null` to search back to the start of the
127
+generation for `"single_turn"`, or of the request for `"all"`; an unbounded
128
+search costs more as the sequence grows. The setting has no effect when
129
+`deduplicate_contexts` is `"none"`. Values below 1,024 emit a warning because a
130
+short window can miss repetition loops whose contexts recur farther apart.
131
+
132
+For example, this checks prompt and completion history within the default
133
+8,192-position window:
134
+
135
+```bash
136
+vllm serve MODEL \
137
+ --watermark-config \
138
+ '{"algorithm":"gumbel","key":42,"deduplicate_contexts":"all"}'
139
+```
140
+
141
+### Dual-key Gumbel-max
142
+
143
+Dual-key Gumbel-max derives independent keys A and B from one configured master
144
+key. During ordinary generation, each token uses key A with probability
145
+`1 - alpha` and key B with probability `alpha`. The algorithm-specific `alpha`
146
+parameter defaults to 0.1. Detection scores every token against both keys.
147
+
148
+The same two key streams support speculative decoding without changing its
149
+acceptance rate. In this mode, the speculative protocol selects the key instead
150
+of `alpha`: draft tokens use key A, while rejection recovery and bonus tokens
151
+use key B. The ordinary target-to-draft probability-ratio test remains unchanged,
152
+implementing [SynthID-Text Supplementary Algorithm
153
+6](https://media.springernature.com/original/springer-static/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf).
154
+
155
+Select `dual_key_gumbel` together with probabilistic drafting:
156
+
157
+```bash
158
+vllm serve MODEL \
159
+ --speculative-config \
160
+ '{"method":"mtp","num_speculative_tokens":3,"draft_sample_method":"probabilistic"}' \
161
+ --watermark-config '{"algorithm":"dual_key_gumbel","key":42,"alpha":0.1}'
162
+```
163
164
### SynthID-Text
165
@@ -106,7 +199,14 @@ result = GumbelWatermarkDetector(key=42, prf="philox").detect(token_ids)
199
print(result.p_value, result.is_watermarked)
200
```
201
109
-The tokenizer, algorithm, PRF, key, and context width must match generation.
202
+The detection configuration must match the generation configuration, including
203
+the tokenizer, PRF, watermarking algorithm, algorithm-specific watermarking
204
+configuration, and key. In practice, this information is often unavailable
205
+when checking a piece of text. Deployments should therefore retain the set of
206
+candidate configurations they have served, test the text against each
207
+candidate, and correct for multiple testing, for example with a Bonferroni
208
+correction to the resulting p-values.
209
+
210
Gumbel-max detection scores repeated contexts once by default so identical PRF
211
random vectors are not treated as independent evidence. Keep
212
`deduplicate_contexts=True` unless the detector's calibration has been adjusted
@@ -140,7 +240,7 @@ watermarked output or to modify watermarked text so it is no longer detected.
240
## Limitations
241
242
- Watermarking is currently available only with Model Runner V2.
143
-- Gumbel-max cannot be configured with speculative decoding.
243
+- Not all watermarking algorithms have native speculative-decoding support.
244
- Beam search expands candidates from model log probabilities and does not apply
245
Gumbel-max watermarking.
246
- Models that replace the vLLM sampler with a custom sampler cannot use
tests/engine/test_arg_utils.py
+22
-3
@@ -53,18 +53,37 @@ def test_watermark_config_cli():
53
"--model",
54
"dummy",
55
"--watermark-config",
56
- '{"algorithm":"gumbel","key":42,"prf":"philox"}',
56
+ '{"algorithm":"dual_key_gumbel","key":42,"prf":"philox","alpha":0.25,'
57
+ '"allow_target_only_watermarking":true}',
58
]
59
)
60
61
config = EngineArgs.from_cli_args(args).create_watermark_config()
62
63
assert config is not None
63
- assert config.algorithm == "gumbel"
64
+ assert config.algorithm == "dual_key_gumbel"
65
assert config.key == 42
66
+ assert config.alpha == 0.25
67
assert config.context_width == 4
68
+ assert config.deduplicate_contexts == "single_turn"
69
+ assert config.deduplicate_contexts_max_history == 8192
70
assert config.prf == "philox"
67
- assert not config.supports_speculative_decoding
71
+ assert config.allow_target_only_watermarking
72
+
73
+ args = parser.parse_args(
74
+ [
75
+ "--model",
76
+ "dummy",
77
+ "--watermark-config",
78
+ '{"key":42,"deduplicate_contexts":"none",'
79
+ '"deduplicate_contexts_max_history":32}',
80
+ ]
81
+ )
82
+ config = EngineArgs.from_cli_args(args).create_watermark_config()
83
+
84
+ assert config is not None
85
+ assert config.deduplicate_contexts == "none"
86
+ assert config.deduplicate_contexts_max_history == 32
87
88
89
@pytest.mark.parametrize(
tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml
new
+29
@@ -0,0 +1,29 @@
1
+model_name: "Qwen/Qwen3-4B-FP8"
2
+num_questions: 1319
3
+num_fewshot: 5
4
+max_tokens: 256
5
+use_chat_completions: false
6
+temperature: 1.0
7
+seed: 42
8
+
9
+# Reference accuracy from the non-watermarked DSpark regression test.
10
+accuracy_threshold: 0.801
11
+tolerance: 0.02
12
+# Measured TP1 acceptance length: 4.115; allow a 0.2 regression.
13
+min_acceptance_length: 3.915
14
+
15
+env:
16
+ VLLM_USE_V2_MODEL_RUNNER: "1"
17
+server_args: >-
18
+ --tensor-parallel-size 2
19
+ --max-model-len 4096
20
+ --gpu-memory-utilization 0.92
21
+ --no-enable-prefix-caching
22
+ --speculative-config.method dspark
23
+ --speculative-config.model deepseek-ai/dspark_qwen3_4b_block7
24
+ --speculative-config.num_speculative_tokens 7
25
+ --speculative-config.draft_sample_method probabilistic
26
+ --watermark-config.algorithm dual_key_gumbel
27
+ --watermark-config.key 42
28
+ --watermark-config.context_width 4
29
+ --watermark-config.deduplicate_contexts none
tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt
new
+1
@@ -0,0 +1 @@
1
+Qwen3-4B-FP8-DSpark-watermark-TP2.yaml
tests/model_executor/test_spec_decode_rejection_warmup.py
new
+121
@@ -0,0 +1,121 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from types import SimpleNamespace
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.config.watermarking import WatermarkConfig, derive_watermark_key
9
+from vllm.model_executor.warmup.spec_decode_rejection_warmup import (
10
+ spec_decode_rejection_warmup,
11
+)
12
+
13
+MASTER_KEY = 20250910
14
+
15
+
16
+def _worker(
17
+ watermark_config: WatermarkConfig | None,
18
+ draft_sample_method: str = "greedy",
19
+) -> SimpleNamespace:
20
+ return SimpleNamespace(
21
+ device=torch.device("cpu"),
22
+ vllm_config=SimpleNamespace(
23
+ speculative_config=SimpleNamespace(
24
+ num_speculative_tokens=2,
25
+ rejection_sample_method="standard",
26
+ draft_sample_method=draft_sample_method,
27
+ ),
28
+ model_config=SimpleNamespace(
29
+ get_vocab_size=lambda: 128, dtype=torch.bfloat16
30
+ ),
31
+ watermark_config=watermark_config,
32
+ ),
33
+ )
34
+
35
+
36
+def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]:
37
+ calls: list[dict] = []
38
+
39
+ def fake_rejection_sample(**kwargs):
40
+ calls.append(kwargs)
41
+ return None, None
42
+
43
+ monkeypatch.setattr(
44
+ "vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils.rejection_sample",
45
+ fake_rejection_sample,
46
+ )
47
+ return calls
48
+
49
+
50
+def test_warmup_token_id_dtypes_match_the_runtime_buffers(
51
+ monkeypatch: pytest.MonkeyPatch,
52
+) -> None:
53
+ calls = _capture_calls(monkeypatch)
54
+ spec_decode_rejection_warmup(_worker(None))
55
+
56
+ assert calls
57
+ for kwargs in calls:
58
+ assert kwargs["draft_sampled"].dtype == torch.int32
59
+ assert kwargs["pos"].dtype == torch.int64
60
+
61
+
62
+def test_warmup_index_mapping_dtypes_match_the_runtime_buffers(
63
+ monkeypatch: pytest.MonkeyPatch,
64
+) -> None:
65
+ calls = _capture_calls(monkeypatch)
66
+ spec_decode_rejection_warmup(_worker(None))
67
+
68
+ assert calls
69
+ for kwargs in calls:
70
+ assert kwargs["idx_mapping"].dtype == torch.int64
71
+ assert kwargs["expanded_idx_mapping"].dtype == torch.int64
72
+ assert kwargs["expanded_local_pos"].dtype == torch.int32
73
+
74
+
75
+def test_warmup_watermark_args_match_the_runtime_kernel_arguments(
76
+ monkeypatch: pytest.MonkeyPatch,
77
+) -> None:
78
+ watermark_config = WatermarkConfig(
79
+ key=MASTER_KEY, algorithm="dual_key_gumbel", context_width=3
80
+ )
81
+ calls = _capture_calls(monkeypatch)
82
+ spec_decode_rejection_warmup(
83
+ _worker(watermark_config, draft_sample_method="probabilistic")
84
+ )
85
+
86
+ assert calls
87
+ key_b = derive_watermark_key(MASTER_KEY, b"key_b")
88
+ for kwargs in calls:
89
+ contexts = kwargs["contexts"]
90
+ assert contexts.dtype == torch.int32
91
+ assert contexts.shape == (3, 3) and contexts.stride() == (3, 1)
92
+ assert kwargs["watermarking"].dtype == torch.bool
93
+ assert kwargs["watermark_key"] == key_b
94
+ assert key_b & 0xFFFFFFFF >= 2**31 and key_b >> 32 >= 2**31
95
+
96
+
97
+def test_warmup_is_skipped_when_the_recovery_key_cannot_be_derived(
98
+ monkeypatch: pytest.MonkeyPatch,
99
+) -> None:
100
+ calls = _capture_calls(monkeypatch)
101
+ monkeypatch.setattr(
102
+ "vllm.v1.watermarking.spec_decode.speculative_target_watermark_key",
103
+ lambda _: (_ for _ in ()).throw(NotImplementedError("no philox")),
104
+ )
105
+ spec_decode_rejection_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
106
+
107
+ assert calls == []
108
+
109
+
110
+def test_warmup_draft_logits_match_the_runtime_draft_sample_method(
111
+ monkeypatch: pytest.MonkeyPatch,
112
+) -> None:
113
+ calls = _capture_calls(monkeypatch)
114
+ spec_decode_rejection_warmup(_worker(None))
115
+ assert calls
116
+ assert all(kwargs["draft_logits"] is None for kwargs in calls)
117
+
118
+ calls = _capture_calls(monkeypatch)
119
+ spec_decode_rejection_warmup(_worker(None, draft_sample_method="probabilistic"))
120
+ assert calls
121
+ assert all(kwargs["draft_logits"] is not None for kwargs in calls)
tests/model_executor/test_watermark_sample_warmup.py
new
+300
@@ -0,0 +1,300 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""The warmup must cover every specialization the sampler will launch.
4
+
5
+``_philox_gumbel_kernel`` is keyed on the Philox key halves, on the logits
6
+dtype (model dtype when the sampler passes logits through, fp32 when logits
7
+processing copies them) and on whether a skip mask is passed, so a missing
8
+combination leaves a request paying the compile this module exists to remove.
9
+"""
10
+
11
+import logging
12
+from types import SimpleNamespace
13
+
14
+import pytest
15
+import torch
16
+
17
+from vllm.config.watermarking import WatermarkConfig, derive_watermark_key
18
+from vllm.model_executor.warmup.watermark_sample_warmup import watermark_sample_warmup
19
+from vllm.platforms import current_platform
20
+from vllm.v1.watermarking.factory import create_watermarker
21
+from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
22
+from vllm.v1.watermarking.spec_decode import create_speculative_target_watermarker
23
+
24
+MASTER_KEY = 20250910
25
+KEY_A = derive_watermark_key(MASTER_KEY, b"key_a")
26
+KEY_B = derive_watermark_key(MASTER_KEY, b"key_b")
27
+
28
+
29
+@pytest.fixture(autouse=True)
30
+def cuda_alike_platform(monkeypatch: pytest.MonkeyPatch) -> None:
31
+ """The warmup returns early off CUDA, so pin the platform it reads.
32
+
33
+ Everything here runs on the fake worker's cpu device; without this the
34
+ tests below would pass only on a CUDA-alike host and silently warm nothing
35
+ anywhere else.
36
+ """
37
+ monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
38
+
39
+
40
+def _sampler(
41
+ watermark_config: WatermarkConfig,
42
+ speculative: bool,
43
+ use_fp64_gumbel: bool,
44
+) -> GPUWatermarkSampler:
45
+ """The sampler ``GPUModelRunner`` builds for this config.
46
+
47
+ Mirrors ``vllm/v1/worker/gpu/model_runner.py:473-486``. ``__new__`` skips
48
+ ``Sampler.__init__``, which allocates UVA-backed device buffers; the warmup
49
+ only reads the three attributes set here.
50
+ """
51
+ watermarker = create_watermarker(watermark_config)
52
+ if speculative:
53
+ watermarker = create_speculative_target_watermarker(watermarker)
54
+ sampler = GPUWatermarkSampler.__new__(GPUWatermarkSampler)
55
+ sampler.watermarker = watermarker
56
+ sampler.deduplicate_contexts = watermark_config.deduplicate_contexts
57
+ sampler.use_fp64_gumbel = use_fp64_gumbel
58
+ return sampler
59
+
60
+
61
+def _worker(
62
+ watermark_config: WatermarkConfig | None,
63
+ speculative: bool = False,
64
+ use_fp64_gumbel: bool = False,
65
+ with_sampler: bool = True,
66
+) -> SimpleNamespace:
67
+ sampler = (
68
+ _sampler(watermark_config, speculative, use_fp64_gumbel)
69
+ if watermark_config is not None and with_sampler
70
+ else None
71
+ )
72
+ return SimpleNamespace(
73
+ device=torch.device("cpu"),
74
+ model_runner=SimpleNamespace(sampler=sampler),
75
+ vllm_config=SimpleNamespace(
76
+ model_config=SimpleNamespace(
77
+ get_vocab_size=lambda: 128,
78
+ dtype=torch.bfloat16,
79
+ ),
80
+ watermark_config=watermark_config,
81
+ ),
82
+ )
83
+
84
+
85
+def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]:
86
+ calls: list[dict] = []
87
+
88
+ def fake_philox_gumbel_sample(logits, contexts, key, **kwargs):
89
+ calls.append({"logits": logits, "contexts": contexts, "key": key, **kwargs})
90
+ return torch.zeros(logits.shape[0], dtype=torch.int64)
91
+
92
+ monkeypatch.setattr(
93
+ "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample",
94
+ fake_philox_gumbel_sample,
95
+ )
96
+ return calls
97
+
98
+
99
+def _variants(calls: list[dict]) -> set[tuple[int, torch.dtype, bool]]:
100
+ return {
101
+ (call["key"], call["logits"].dtype, call.get("skip_mask") is not None)
102
+ for call in calls
103
+ }
104
+
105
+
106
+def test_no_warmup_without_a_watermark_config(
107
+ monkeypatch: pytest.MonkeyPatch,
108
+) -> None:
109
+ calls = _capture_calls(monkeypatch)
110
+ watermark_sample_warmup(_worker(None))
111
+
112
+ assert calls == []
113
+
114
+
115
+def test_no_warmup_without_a_watermarked_sampler(
116
+ monkeypatch: pytest.MonkeyPatch,
117
+) -> None:
118
+ """A rank that does not sample launches nothing to warm.
119
+
120
+ Only the final pipeline stage of a generative model holds a
121
+ ``GPUWatermarkSampler``; on any other rank the kernel is never launched.
122
+ """
123
+ calls = _capture_calls(monkeypatch)
124
+ watermark_sample_warmup(
125
+ _worker(WatermarkConfig(key=MASTER_KEY), with_sampler=False)
126
+ )
127
+
128
+ assert calls == []
129
+
130
+
131
+def test_no_warmup_off_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
132
+ """``GumbelWatermarker`` launches the kernel only on a CUDA-alike device.
133
+
134
+ Elsewhere -- an XPU worker inherits ``kernel_warmup`` from the CUDA worker
135
+ -- the sampler takes the torch path, so warming would compile a kernel
136
+ nothing launches.
137
+ """
138
+ calls = _capture_calls(monkeypatch)
139
+ monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: False)
140
+ watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
141
+
142
+ assert calls == []
143
+
144
+
145
+def test_single_key_warms_both_logits_dtypes(
146
+ monkeypatch: pytest.MonkeyPatch,
147
+) -> None:
148
+ calls = _capture_calls(monkeypatch)
149
+ watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
150
+
151
+ assert _variants(calls) == {
152
+ (MASTER_KEY, torch.bfloat16, True),
153
+ (MASTER_KEY, torch.float32, True),
154
+ }
155
+ assert len(calls) == len(_variants(calls))
156
+
157
+
158
+def test_dual_key_warms_both_derived_keys(monkeypatch: pytest.MonkeyPatch) -> None:
159
+ """``DualKeyGumbelWatermarker.sample`` launches the kernel once per key.
160
+
161
+ The halves of a derived key exceed 2**31, so each key is typed i64/i32 by
162
+ magnitude and a wrong key compiles a specialization the engine never uses.
163
+ """
164
+ calls = _capture_calls(monkeypatch)
165
+ watermark_sample_warmup(
166
+ _worker(WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel"))
167
+ )
168
+
169
+ assert _variants(calls) == {
170
+ (KEY_A, torch.bfloat16, True),
171
+ (KEY_A, torch.float32, True),
172
+ (KEY_B, torch.bfloat16, True),
173
+ (KEY_B, torch.float32, True),
174
+ }
175
+
176
+
177
+@pytest.mark.parametrize("alpha,key", [(0.0, KEY_A), (1.0, KEY_B)])
178
+def test_dual_key_alpha_bounds_warm_one_key(
179
+ monkeypatch: pytest.MonkeyPatch, alpha: float, key: int
180
+) -> None:
181
+ calls = _capture_calls(monkeypatch)
182
+ watermark_sample_warmup(
183
+ _worker(
184
+ WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel", alpha=alpha)
185
+ )
186
+ )
187
+
188
+ assert _variants(calls) == {
189
+ (key, torch.bfloat16, True),
190
+ (key, torch.float32, True),
191
+ }
192
+
193
+
194
+def test_speculative_config_warms_the_target_role_only(
195
+ monkeypatch: pytest.MonkeyPatch,
196
+) -> None:
197
+ """Under spec decode the sampler holds ``create_target_watermarker()``.
198
+
199
+ Rows without draft tokens sample through it, so the sampler path launches
200
+ key B alone; the draft key runs in the speculator's own sampler.
201
+ """
202
+ calls = _capture_calls(monkeypatch)
203
+ watermark_sample_warmup(
204
+ _worker(
205
+ WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel"),
206
+ speculative=True,
207
+ )
208
+ )
209
+
210
+ assert _variants(calls) == {
211
+ (KEY_B, torch.bfloat16, True),
212
+ (KEY_B, torch.float32, True),
213
+ }
214
+
215
+
216
+def test_dedup_none_also_warms_the_mask_free_specialization(
217
+ monkeypatch: pytest.MonkeyPatch,
218
+) -> None:
219
+ """Without deduplication an all-watermarked batch passes no skip mask.
220
+
221
+ A later batch mixing in an opted-out request passes one, so both
222
+ specializations are reachable and both must be warmed.
223
+ """
224
+ calls = _capture_calls(monkeypatch)
225
+ watermark_sample_warmup(
226
+ _worker(WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none"))
227
+ )
228
+
229
+ assert _variants(calls) == {
230
+ (MASTER_KEY, torch.bfloat16, True),
231
+ (MASTER_KEY, torch.float32, True),
232
+ (MASTER_KEY, torch.bfloat16, False),
233
+ (MASTER_KEY, torch.float32, False),
234
+ }
235
+ for call in calls:
236
+ if call.get("skip_mask") is None:
237
+ assert "temperatures" not in call
238
+
239
+
240
+def test_warmup_argument_dtypes_match_the_runtime_buffers(
241
+ monkeypatch: pytest.MonkeyPatch,
242
+) -> None:
243
+ """Contexts are int32, the sampling state int64 indices and fp32 temperatures.
244
+
245
+ ``GPUWatermarkSampler._get_contexts`` gathers from the int32 request-state
246
+ token ids, ``expanded_idx_mapping``/``positions``/``seeds`` are int64 and
247
+ ``SamplingStates.temperature`` is fp32; warming other dtypes compiles
248
+ specializations no engine launches.
249
+ """
250
+ calls = _capture_calls(monkeypatch)
251
+ watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY, context_width=3)))
252
+
253
+ assert calls
254
+ for call in calls:
255
+ num_tokens = call["logits"].shape[0]
256
+ assert call["contexts"].dtype == torch.int32
257
+ assert call["contexts"].shape == (num_tokens, 3)
258
+ assert call["skip_mask"].dtype == torch.bool
259
+ assert call["expanded_idx_mapping"].dtype == torch.int64
260
+ assert call["positions"].dtype == torch.int64
261
+ assert call["seeds"].dtype == torch.int64
262
+ assert call["temperatures"].dtype == torch.float32
263
+
264
+
265
+def test_use_fp64_follows_the_path_that_forwards_it(
266
+ monkeypatch: pytest.MonkeyPatch,
267
+) -> None:
268
+ """Only the mixed path forwards the sampler's ``use_fp64_gumbel``.
269
+
270
+ ``GumbelWatermarker._sample_watermarked`` launches the mask-free variant
271
+ with the ``use_fp64=False`` default, so warming it with fp64 would compile
272
+ a specialization the engine never launches and miss the one it does.
273
+ """
274
+ calls = _capture_calls(monkeypatch)
275
+ watermark_sample_warmup(
276
+ _worker(
277
+ WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none"),
278
+ use_fp64_gumbel=True,
279
+ )
280
+ )
281
+
282
+ assert calls
283
+ for call in calls:
284
+ assert call["use_fp64"] is (call.get("skip_mask") is not None)
285
+
286
+
287
+def test_warmup_failure_does_not_abort_startup(
288
+ monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
289
+) -> None:
290
+ def failing_philox_gumbel_sample(*args, **kwargs):
291
+ raise RuntimeError("no triton")
292
+
293
+ monkeypatch.setattr(
294
+ "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample",
295
+ failing_philox_gumbel_sample,
296
+ )
297
+ with caplog.at_level(logging.WARNING):
298
+ watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
299
+
300
+ assert "Skipping watermark sampler warmup." in caplog.text
tests/test_config.py
+179
-4
@@ -2856,14 +2856,183 @@ def _watermarked_vllm_config() -> VllmConfig:
2856
return config
2857
2858
2859
-def test_gumbel_watermark_rejects_speculative_decoding():
2859
+def test_target_only_gumbel_allows_speculative_decoding(caplog_vllm, disable_log_dedup):
2860
config = _watermarked_vllm_config()
2861
- config.speculative_config = SpeculativeConfig(
2861
+ config.watermark_config = WatermarkConfig(
2862
+ key=42, allow_target_only_watermarking=True
2863
+ )
2864
+ config.speculative_config = SimpleNamespace(
2865
+ method="mtp",
2866
+ draft_sample_method="probabilistic",
2867
+ rejection_sample_method="standard",
2868
+ parallel_drafting=False,
2869
+ )
2870
+
2871
+ with caplog_vllm.at_level(logging.WARNING):
2872
+ config._check_watermarking_unsupported()
2873
+
2874
+ assert "Target-only watermarking leaves accepted draft tokens" in caplog_vllm.text
2875
+ assert "Context deduplication is not supported" in caplog_vllm.text
2876
+
2877
+
2878
+def test_speculative_watermarking_without_context_dedup_does_not_warn(
2879
+ caplog_vllm, disable_log_dedup
2880
+):
2881
+ config = _watermarked_vllm_config()
2882
+ config.watermark_config = WatermarkConfig(
2883
+ algorithm="dual_key_gumbel", key=42, deduplicate_contexts="none"
2884
+ )
2885
+ config.speculative_config = SimpleNamespace(
2886
+ method="mtp",
2887
+ draft_sample_method="probabilistic",
2888
+ rejection_sample_method="standard",
2889
+ parallel_drafting=False,
2890
+ )
2891
+
2892
+ with caplog_vllm.at_level(logging.WARNING):
2893
+ config._check_watermarking_unsupported()
2894
+
2895
+ assert "Context deduplication is not supported" not in caplog_vllm.text
2896
+
2897
+
2898
+def test_gumbel_rejects_speculative_decoding_without_target_only():
2899
+ config = _watermarked_vllm_config()
2900
+ config.watermark_config = WatermarkConfig(
2901
+ algorithm="gumbel", key=42, allow_target_only_watermarking=False
2902
+ )
2903
+ config.speculative_config = SimpleNamespace(
2904
+ method="mtp",
2905
+ draft_sample_method="probabilistic",
2906
+ rejection_sample_method="standard",
2907
+ parallel_drafting=False,
2908
+ )
2909
+
2910
+ with pytest.raises(ValueError, match="'gumbel'.*allow_target_only_watermarking"):
2911
+ config._check_watermarking_unsupported()
2912
+
2913
+
2914
+def test_dual_key_gumbel_warns_that_configured_alpha_is_unused(
2915
+ caplog_vllm, disable_log_dedup
2916
+):
2917
+ config = _watermarked_vllm_config()
2918
+ config.watermark_config = WatermarkConfig(
2919
+ algorithm="dual_key_gumbel", key=42, alpha=0.25
2920
+ )
2921
+ config.speculative_config = SimpleNamespace(
2922
+ method="mtp",
2923
+ draft_sample_method="probabilistic",
2924
+ rejection_sample_method="standard",
2925
+ parallel_drafting=False,
2926
+ )
2927
+
2928
+ with caplog_vllm.at_level(logging.WARNING):
2929
+ config._check_watermarking_unsupported()
2930
+
2931
+ assert "The configured alpha=0.25 is not used" in caplog_vllm.text
2932
+
2933
+
2934
+@pytest.mark.parametrize(
2935
+ ("alpha", "speculative"),
2936
+ [(0.1, True), (0.25, False)],
2937
+ ids=["default-alpha", "no-specdec"],
2938
+)
2939
+def test_dual_key_gumbel_alpha_warning_is_not_emitted(
2940
+ caplog_vllm, disable_log_dedup, alpha, speculative
2941
+):
2942
+ config = _watermarked_vllm_config()
2943
+ config.watermark_config = WatermarkConfig(
2944
+ algorithm="dual_key_gumbel", key=42, alpha=alpha
2945
+ )
2946
+ if speculative:
2947
+ config.speculative_config = SimpleNamespace(
2948
+ method="mtp",
2949
+ draft_sample_method="probabilistic",
2950
+ rejection_sample_method="standard",
2951
+ parallel_drafting=False,
2952
+ )
2953
+
2954
+ with caplog_vllm.at_level(logging.WARNING):
2955
+ config._check_watermarking_unsupported()
2956
+
2957
+ assert "is not used" not in caplog_vllm.text
2958
+
2959
+
2960
+def test_dual_key_gumbel_requires_probabilistic_drafting():
2961
+ config = _watermarked_vllm_config()
2962
+ config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
2963
+ config.speculative_config = SimpleNamespace(
2964
+ method="mtp",
2965
+ draft_sample_method="greedy",
2966
+ rejection_sample_method="standard",
2967
+ parallel_drafting=False,
2968
+ )
2969
+
2970
+ with pytest.raises(ValueError, match="draft_sample_method='probabilistic'"):
2971
+ config._check_watermarking_unsupported()
2972
+
2973
+
2974
+@pytest.mark.parametrize("method", ["eagle", "eagle3", "mtp"])
2975
+def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(method):
2976
+ config = _watermarked_vllm_config()
2977
+ config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
2978
+ config.speculative_config = SimpleNamespace(
2979
+ method=method,
2980
+ draft_sample_method="probabilistic",
2981
+ rejection_sample_method="standard",
2982
+ parallel_drafting=False,
2983
+ )
2984
+
2985
+ config._check_watermarking_unsupported()
2986
+
2987
+
2988
+def test_dual_key_gumbel_supports_dspark():
2989
+ config = _watermarked_vllm_config()
2990
+ config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
2991
+ config.speculative_config = SimpleNamespace(
2992
+ method="dspark",
2993
+ draft_sample_method="probabilistic",
2994
+ rejection_sample_method="standard",
2995
+ parallel_drafting=True,
2996
+ )
2997
+
2998
+ config._check_watermarking_unsupported()
2999
+
3000
+
3001
+def test_dual_key_gumbel_rejects_non_autoregressive_speculation():
3002
+ config = _watermarked_vllm_config()
3003
+ config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
3004
+ config.speculative_config = SimpleNamespace(
3005
method="ngram",
2863
- num_speculative_tokens=1,
3006
+ draft_sample_method="probabilistic",
3007
+ rejection_sample_method="standard",
3008
+ parallel_drafting=False,
3009
)
3010
2866
- with pytest.raises(ValueError, match="does not support speculative decoding"):
3011
+ with pytest.raises(ValueError, match="autoregressive model-based"):
3012
+ config._check_watermarking_unsupported()
3013
+
3014
+
3015
+@pytest.mark.parametrize(
3016
+ ("overrides", "match"),
3017
+ [
3018
+ ({"rejection_sample_method": "synthetic"}, "rejection_sample_method"),
3019
+ ({"rejection_sample_method": "block"}, "rejection_sample_method"),
3020
+ ({"parallel_drafting": True}, "Parallel speculative drafting"),
3021
+ ],
3022
+)
3023
+def test_dual_key_gumbel_rejects_incompatible_speculative_modes(overrides, match):
3024
+ config = _watermarked_vllm_config()
3025
+ config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
3026
+ values = {
3027
+ "method": "mtp",
3028
+ "draft_sample_method": "probabilistic",
3029
+ "rejection_sample_method": "standard",
3030
+ "parallel_drafting": False,
3031
+ **overrides,
3032
+ }
3033
+ config.speculative_config = SimpleNamespace(**values)
3034
+
3035
+ with pytest.raises(ValueError, match=match):
3036
config._check_watermarking_unsupported()
3037
3038
@@ -2882,6 +3051,12 @@ def test_watermark_key_must_fit_in_64_bits():
3051
WatermarkConfig(key=2**64)
3052
3053
3054
+@pytest.mark.parametrize("alpha", [-0.1, 1.1])
3055
+def test_watermark_alpha_must_be_a_probability(alpha):
3056
+ with pytest.raises(ValidationError):
3057
+ WatermarkConfig(key=42, algorithm="dual_key_gumbel", alpha=alpha)
3058
+
3059
+
3060
def test_unknown_watermark_prf_is_rejected():
3061
with pytest.raises(ValidationError):
3062
pydantic.TypeAdapter(WatermarkConfig).validate_python(
tests/v1/spec_decode/test_rejection_sampler_utils.py
+279
@@ -6,7 +6,10 @@ import math
6
import pytest
7
import torch
8
9
+from vllm.v1.watermarking import GumbelWatermarker
10
+from vllm.v1.watermarking.spec_decode import watermarked_rejection_sample
11
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
12
+from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
13
from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
14
rejection_sample,
15
)
@@ -139,6 +142,93 @@ def _assert_distribution_match(
142
)
143
144
145
+def test_watermarked_bonus_uses_target_key():
146
+ target_logits = torch.randn(16, device="cuda")
147
+ inputs = _build_rejection_sample_inputs(
148
+ target_logits,
149
+ target_logits,
150
+ num_speculative_steps=1,
151
+ temperature=1.0,
152
+ num_trials=1,
153
+ )
154
+ draft_token = inputs["draft_sampled"][1]
155
+ contexts = torch.tensor(
156
+ [[10, 11], [11, draft_token.item()]], dtype=torch.int64, device="cuda"
157
+ )
158
+
159
+ sampled, num_sampled = watermarked_rejection_sample(
160
+ **inputs,
161
+ num_speculative_steps=1,
162
+ contexts=contexts,
163
+ watermarking=torch.tensor([True], device="cuda"),
164
+ watermarker=GumbelWatermarker(key=42, context_width=2),
165
+ )
166
+ expected = GumbelWatermarker(key=42, context_width=2).sample(
167
+ target_logits.unsqueeze(0), contexts[1:], lambda _: None
168
+ )
169
+
170
+ assert num_sampled.item() == 2
171
+ assert sampled[0, 1] == expected.token_ids[0]
172
+
173
+
174
+def test_watermarked_recovery_uses_target_key():
175
+ target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0], device="cuda")
176
+ draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda")
177
+ inputs = _build_rejection_sample_inputs(
178
+ target_logits,
179
+ draft_logits,
180
+ num_speculative_steps=1,
181
+ temperature=1.0,
182
+ num_trials=1,
183
+ )
184
+ inputs["draft_sampled"][1] = 0
185
+ contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda")
186
+
187
+ sampled, num_sampled = watermarked_rejection_sample(
188
+ **inputs,
189
+ num_speculative_steps=1,
190
+ contexts=contexts,
191
+ watermarking=torch.tensor([True], device="cuda"),
192
+ watermarker=GumbelWatermarker(key=42, context_width=2),
193
+ )
194
+ expected = GumbelWatermarker(key=42, context_width=2).sample(
195
+ target_logits.unsqueeze(0), contexts[:1], lambda _: None
196
+ )
197
+
198
+ assert num_sampled.item() == 1
199
+ assert sampled[0, 0] == expected.token_ids[0]
200
+
201
+
202
+def test_watermarked_recovery_supports_smaller_draft_vocabulary():
203
+ target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0, 100.0], device="cuda")
204
+ draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda")
205
+ inputs = _build_rejection_sample_inputs(
206
+ target_logits[:4],
207
+ draft_logits,
208
+ num_speculative_steps=1,
209
+ temperature=1.0,
210
+ num_trials=1,
211
+ )
212
+ inputs["target_logits"] = target_logits.unsqueeze(0).expand(2, -1).contiguous()
213
+ inputs["draft_sampled"][1] = 0
214
+ contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda")
215
+ watermarker = GumbelWatermarker(key=42, context_width=2)
216
+
217
+ sampled, num_sampled = watermarked_rejection_sample(
218
+ **inputs,
219
+ num_speculative_steps=1,
220
+ contexts=contexts,
221
+ watermarking=torch.tensor([True], device="cuda"),
222
+ watermarker=watermarker,
223
+ )
224
+ expected = watermarker.sample(
225
+ target_logits[:4].unsqueeze(0), contexts[:1], lambda _: None
226
+ )
227
+
228
+ assert num_sampled.item() == 1
229
+ assert sampled[0, 0] == expected.token_ids[0]
230
+
231
+
232
@pytest.mark.parametrize(
233
"num_speculative_steps,temperature",
234
[
@@ -747,3 +837,192 @@ def test_chunked_requests_match_full_batch(has_draft_logits: bool):
837
steps = torch.arange(num_speculative_steps + 1, device=device)
838
valid = steps.unsqueeze(0) < num_sampled.unsqueeze(1)
839
assert torch.equal(chunked_sampled[valid], sampled[valid])
840
+
841
+
842
+def _residual_recovery_logits(
843
+ inputs: dict,
844
+ num_sampled: torch.Tensor,
845
+ vocab_size: int,
846
+) -> tuple[torch.Tensor, torch.Tensor]:
847
+ target_logits = inputs["target_logits"]
848
+ draft_logits = inputs["draft_logits"]
849
+ cu_num_logits = inputs["cu_num_logits"]
850
+ resample_rows = cu_num_logits[:-1].to(torch.int64) + num_sampled.to(torch.int64) - 1
851
+ req_state_indices = inputs["expanded_idx_mapping"][resample_rows].to(torch.int64)
852
+ temperatures = inputs["temperature"][req_state_indices]
853
+
854
+ is_bonus = resample_rows == cu_num_logits[1:] - 1
855
+ next_rows = (resample_rows + 1).clamp_max(inputs["draft_sampled"].shape[0] - 1)
856
+ needs_residual = ~is_bonus & (inputs["draft_sampled"][next_rows] >= 0)
857
+
858
+ target_rows = target_logits[resample_rows][:, :vocab_size].float()
859
+ target_rows = torch.where(target_rows.isnan(), float("-inf"), target_rows)
860
+ if draft_logits is None:
861
+ rejected = inputs["draft_sampled"][next_rows].clamp_min(0)
862
+ one_hot = torch.arange(vocab_size, device=target_rows.device).unsqueeze(
863
+ 0
864
+ ) == rejected.unsqueeze(-1)
865
+ residual_logits = torch.where(one_hot, float("-inf"), target_rows)
866
+ else:
867
+ draft_steps = (num_sampled.to(torch.int64) - 1).clamp_max(
868
+ draft_logits.shape[1] - 1
869
+ )
870
+ draft_rows = draft_logits[req_state_indices, draft_steps][
871
+ :, :vocab_size
872
+ ].float()
873
+ draft_rows = torch.where(draft_rows.isnan(), float("-inf"), draft_rows)
874
+ draft_rows = draft_rows / temperatures.unsqueeze(-1)
875
+ target_log_probs = torch.log_softmax(target_rows, dim=-1)
876
+ draft_log_probs = torch.log_softmax(draft_rows, dim=-1)
877
+ ratio = torch.exp(draft_log_probs - target_log_probs)
878
+ residual_logits = torch.where(
879
+ ratio < 1,
880
+ target_log_probs + torch.log1p(-ratio.clamp_max(1)),
881
+ float("-inf"),
882
+ )
883
+ recovery = torch.where(needs_residual.unsqueeze(-1), residual_logits, target_rows)
884
+ return recovery, resample_rows
885
+
886
+
887
+def _watermark_inputs(
888
+ num_trials: int,
889
+ num_speculative_steps: int,
890
+ context_width: int,
891
+ vocab_size: int,
892
+ one_hot_draft: bool,
893
+ temperature: float = 1.0,
894
+) -> tuple[dict, torch.Tensor]:
895
+ device = "cuda"
896
+ target_logits_1d = torch.randn(vocab_size, device=device) / temperature
897
+ draft_logits_1d = torch.randn(vocab_size, device=device)
898
+ inputs = _build_rejection_sample_inputs(
899
+ target_logits_1d,
900
+ draft_logits_1d,
901
+ num_speculative_steps,
902
+ temperature=temperature,
903
+ num_trials=num_trials,
904
+ )
905
+ if one_hot_draft:
906
+ inputs["draft_logits"] = None
907
+ inputs["draft_sampled"].view(num_trials, num_speculative_steps + 1)[::5, 1] = -1
908
+ num_logits = inputs["target_logits"].shape[0]
909
+ contexts = torch.randint(
910
+ 0, vocab_size, (num_logits, context_width), dtype=torch.int64, device=device
911
+ )
912
+ contexts[::7, 0] = -1
913
+ return inputs, contexts
914
+
915
+
916
+@pytest.mark.parametrize("one_hot_draft", [False, True])
917
+# Cover one full block and a partial multi-block vocabulary.
918
+@pytest.mark.parametrize("vocab_size", [1024, 3500])
919
+def test_watermarked_recovery_matches_philox_gumbel_sample(
920
+ one_hot_draft: bool, vocab_size: int
921
+):
922
+ torch.manual_seed(7)
923
+ key = 0xDA39A3EE5E6B4B0D
924
+ context_width = 4
925
+ inputs, contexts = _watermark_inputs(
926
+ num_trials=128,
927
+ num_speculative_steps=2,
928
+ context_width=context_width,
929
+ vocab_size=vocab_size,
930
+ one_hot_draft=one_hot_draft,
931
+ )
932
+ watermarker = GumbelWatermarker(key=key, context_width=context_width)
933
+
934
+ sampled, num_sampled = watermarked_rejection_sample(
935
+ **inputs,
936
+ num_speculative_steps=2,
937
+ contexts=contexts,
938
+ watermarking=torch.ones(128, dtype=torch.bool, device="cuda"),
939
+ watermarker=watermarker,
940
+ )
941
+
942
+ recovery, resample_rows = _residual_recovery_logits(inputs, num_sampled, vocab_size)
943
+ expected = philox_gumbel_sample(recovery, contexts[resample_rows], key)
944
+ recovered = sampled[torch.arange(128, device="cuda"), num_sampled.long() - 1]
945
+ torch.testing.assert_close(recovered, expected, rtol=0, atol=0)
946
+
947
+ is_bonus = resample_rows == inputs["cu_num_logits"][1:] - 1
948
+ next_rows = (resample_rows + 1).clamp_max(inputs["draft_sampled"].shape[0] - 1)
949
+ is_placeholder = ~is_bonus & (inputs["draft_sampled"][next_rows] < 0)
950
+ assert int((~is_bonus & ~is_placeholder).sum()) > 0
951
+ assert int(is_placeholder.sum()) > 0
952
+ if not one_hot_draft:
953
+ # One-hot drafts almost never reach the bonus row.
954
+ assert int(is_bonus.sum()) > 0
955
+
956
+
957
+def test_watermarked_recovery_leaves_disabled_rows_on_the_stock_draw():
958
+ torch.manual_seed(11)
959
+ num_trials = 96
960
+ vocab_size = 512
961
+ inputs, contexts = _watermark_inputs(
962
+ num_trials=num_trials,
963
+ num_speculative_steps=3,
964
+ context_width=2,
965
+ vocab_size=vocab_size,
966
+ one_hot_draft=False,
967
+ )
968
+ watermarking = torch.zeros(num_trials, dtype=torch.bool, device="cuda")
969
+ watermarking[::2] = True
970
+ inputs["temperature"][:8] = 0.0
971
+
972
+ stock_sampled, stock_num_sampled = rejection_sample(
973
+ **inputs, num_speculative_steps=3
974
+ )
975
+ sampled, num_sampled = watermarked_rejection_sample(
976
+ **inputs,
977
+ num_speculative_steps=3,
978
+ contexts=contexts,
979
+ watermarking=watermarking,
980
+ watermarker=GumbelWatermarker(key=99, context_width=2),
981
+ )
982
+
983
+ torch.testing.assert_close(num_sampled, stock_num_sampled, rtol=0, atol=0)
984
+ rows = torch.arange(num_trials, device="cuda")
985
+ positions = num_sampled.long() - 1
986
+ disabled = ~watermarking | (inputs["temperature"] == 0)
987
+ torch.testing.assert_close(
988
+ sampled[rows[disabled], positions[disabled]],
989
+ stock_sampled[rows[disabled], positions[disabled]],
990
+ rtol=0,
991
+ atol=0,
992
+ )
993
+ for req in range(num_trials):
994
+ n = int(num_sampled[req]) - 1
995
+ torch.testing.assert_close(
996
+ sampled[req, :n], stock_sampled[req, :n], rtol=0, atol=0
997
+ )
998
+ assert int(disabled.sum()) > 0 and int((~disabled).sum()) > 0
999
+
1000
+
1001
+def test_watermarked_recovery_is_unchanged_by_fp64_gumbel():
1002
+ torch.manual_seed(13)
1003
+ inputs, contexts = _watermark_inputs(
1004
+ num_trials=64,
1005
+ num_speculative_steps=2,
1006
+ context_width=4,
1007
+ vocab_size=2048,
1008
+ one_hot_draft=False,
1009
+ )
1010
+ watermarker = GumbelWatermarker(key=2**63 + 12345, context_width=4)
1011
+ kwargs = dict(
1012
+ num_speculative_steps=2,
1013
+ contexts=contexts,
1014
+ watermarking=torch.ones(64, dtype=torch.bool, device="cuda"),
1015
+ watermarker=watermarker,
1016
+ )
1017
+ sampled32, num_sampled32 = watermarked_rejection_sample(**inputs, **kwargs)
1018
+ sampled64, num_sampled64 = watermarked_rejection_sample(
1019
+ **inputs, **kwargs, use_fp64=True
1020
+ )
1021
+ rows = torch.arange(64, device="cuda")
1022
+ torch.testing.assert_close(num_sampled64, num_sampled32, rtol=0, atol=0)
1023
+ torch.testing.assert_close(
1024
+ sampled64[rows, num_sampled64.long() - 1],
1025
+ sampled32[rows, num_sampled32.long() - 1],
1026
+ rtol=0,
1027
+ atol=0,
1028
+ )
tests/watermarking/test_gumbel.py
+125
-10
@@ -5,8 +5,15 @@ import pytest
5
import torch
6
7
from vllm.platforms import current_platform
8
-from vllm.v1.watermarking import GumbelWatermarkDetector, GumbelWatermarker
8
+from vllm.v1.watermarking import (
9
+ DualKeyGumbelWatermarkDetector,
10
+ GumbelWatermarkDetector,
11
+ GumbelWatermarker,
12
+ derive_watermark_key,
13
+)
14
from vllm.v1.watermarking.gumbel import _gamma_survival_integer_shape
15
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
16
+from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
17
18
19
def test_gamma_survival_integer_shape():
@@ -34,6 +41,67 @@ def test_detector_deduplicates_context_even_when_target_differs():
41
assert detection.num_scored_tokens == 3
42
43
44
+def test_dual_key_detector_scores_each_token_against_both_keys():
45
+ token_ids = [1, 2, 3, 4, 5]
46
+ key_a_detector = GumbelWatermarkDetector(
47
+ key=derive_watermark_key(42, b"key_a"),
48
+ context_width=1,
49
+ deduplicate_contexts=False,
50
+ )
51
+ key_b_detector = GumbelWatermarkDetector(
52
+ key=derive_watermark_key(42, b"key_b"),
53
+ context_width=1,
54
+ deduplicate_contexts=False,
55
+ )
56
+ detector = DualKeyGumbelWatermarkDetector(
57
+ key=42,
58
+ context_width=1,
59
+ deduplicate_contexts=False,
60
+ alpha=0.5,
61
+ )
62
+
63
+ key_a = key_a_detector.detect(token_ids)
64
+ key_b = key_b_detector.detect(token_ids)
65
+ dual = detector.detect(token_ids)
66
+
67
+ assert dual.score == pytest.approx((key_a.score + key_b.score) / 2)
68
+ assert dual.p_value == pytest.approx(
69
+ _gamma_survival_integer_shape(dual.score * 2, dual.num_scored_tokens * 2)
70
+ )
71
+
72
+
73
+@pytest.mark.parametrize("alpha", [0.2, 0.7])
74
+def test_dual_key_detector_weights_and_recalibrates_scores(alpha: float):
75
+ token_ids = [1, 2, 3, 4, 5]
76
+ key_a = GumbelWatermarkDetector(
77
+ key=derive_watermark_key(42, b"key_a"), deduplicate_contexts=False
78
+ ).detect(token_ids)
79
+ key_b = GumbelWatermarkDetector(
80
+ key=derive_watermark_key(42, b"key_b"), deduplicate_contexts=False
81
+ ).detect(token_ids)
82
+ dual = DualKeyGumbelWatermarkDetector(
83
+ key=42, deduplicate_contexts=False, alpha=alpha
84
+ ).detect(token_ids)
85
+
86
+ variance = (1 - alpha) ** 2 + alpha**2
87
+ expected_p_value = torch.special.gammaincc(
88
+ torch.tensor(dual.num_scored_tokens / variance, dtype=torch.float64),
89
+ torch.tensor(dual.score / variance, dtype=torch.float64),
90
+ ).item()
91
+
92
+ assert dual.score == pytest.approx((1 - alpha) * key_a.score + alpha * key_b.score)
93
+ assert dual.p_value == pytest.approx(expected_p_value)
94
+
95
+
96
+def test_dual_key_detector_rejects_invalid_alpha():
97
+ with pytest.raises(ValueError, match="alpha must be between 0 and 1"):
98
+ DualKeyGumbelWatermarkDetector(key=42, alpha=1.1)
99
+
100
+
101
+def test_dual_key_detector_default_alpha():
102
+ assert DualKeyGumbelWatermarkDetector(key=42).alpha == 0.2
103
+
104
+
105
@pytest.mark.skipif(
106
not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
107
)
@@ -45,10 +113,8 @@ def test_fused_watermarker_matches_cpu(key: int, context_width: int):
113
logits = torch.randn(32, 8193)
114
watermarker = GumbelWatermarker(key, context_width)
115
48
- expected = watermarker.sample(logits, contexts, lambda values: None).token_ids
49
- actual = watermarker.sample(
50
- logits.cuda(), contexts.cuda(), lambda values: None
51
- ).token_ids.cpu()
116
+ expected = watermarker.sample(logits, contexts).token_ids
117
+ actual = watermarker.sample(logits.cuda(), contexts.cuda()).token_ids.cpu()
118
119
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
120
@@ -61,7 +127,7 @@ def test_fused_watermarker_handles_nan_logits():
127
logits = torch.full((2, 1025), float("nan"), device="cuda")
128
watermarker = GumbelWatermarker(key=42, context_width=4)
129
64
- token_ids = watermarker.sample(logits, contexts, lambda values: None).token_ids
130
+ token_ids = watermarker.sample(logits, contexts).token_ids
131
132
assert torch.all((token_ids >= 0) & (token_ids < logits.shape[-1]))
133
@@ -76,9 +142,58 @@ def test_fused_watermarker_handles_noncontiguous_inputs():
142
logits = torch.randn(8, 2050, device="cuda")[:, ::2]
143
watermarker = GumbelWatermarker(key=42, context_width=4)
144
79
- expected = watermarker.sample(
80
- logits.contiguous(), contexts.contiguous(), lambda values: None
81
- ).token_ids
82
- actual = watermarker.sample(logits, contexts, lambda values: None).token_ids
145
+ expected = watermarker.sample(logits.contiguous(), contexts.contiguous()).token_ids
146
+ actual = watermarker.sample(logits, contexts).token_ids
147
+
148
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
149
+
150
+
151
+@pytest.mark.skipif(
152
+ not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
153
+)
154
+@pytest.mark.parametrize("use_fp64", [False, True])
155
+def test_philox_gumbel_sample_skip_mask_matches_separate_samplers(use_fp64: bool):
156
+ torch.manual_seed(0)
157
+ logits = torch.randn(8, 8193, device="cuda")
158
+ context_storage = torch.randint(0, 248320, (8, 8), dtype=torch.int64, device="cuda")
159
+ contexts = context_storage[:, ::2]
160
+ repeated_mask = torch.tensor(
161
+ [False, False, False, True, False, False, True, False], device="cuda"
162
+ )
163
+ watermarking = torch.tensor(
164
+ [True, False, True, True, False, True, True, False], device="cuda"
165
+ )
166
+ req_indices = torch.tensor([3, 1, 7, 0, 4, 2, 6, 5], device="cuda")
167
+ temperatures = torch.ones(8, dtype=torch.float32, device="cuda")
168
+ temperatures[2] = 0
169
+ seeds = torch.arange(8, dtype=torch.int64, device="cuda") + 1000
170
+ positions = torch.arange(8, dtype=torch.int64, device="cuda") + 100
171
+ watermark_mask = (
172
+ watermarking[req_indices] & (temperatures[req_indices] != 0) & ~repeated_mask
173
+ )
174
+
175
+ watermarked = philox_gumbel_sample(logits, contexts, 42)
176
+ ordinary = gumbel_sample(
177
+ logits,
178
+ req_indices,
179
+ temperatures,
180
+ seeds,
181
+ positions,
182
+ apply_temperature=False,
183
+ is_drafting=False,
184
+ use_fp64=use_fp64,
185
+ )
186
+ expected = torch.where(watermark_mask, watermarked, ordinary)
187
+ actual = philox_gumbel_sample(
188
+ logits,
189
+ contexts,
190
+ 42,
191
+ skip_mask=~watermark_mask,
192
+ expanded_idx_mapping=req_indices,
193
+ temperatures=temperatures,
194
+ seeds=seeds,
195
+ positions=positions,
196
+ use_fp64=use_fp64,
197
+ )
198
199
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
tests/watermarking/test_watermarking.py
+889
-18
@@ -9,44 +9,98 @@ import torch
9
10
from vllm import SamplingParams
11
from vllm.config.watermarking import WatermarkConfig
12
-from vllm.v1.watermarking import create_watermarker
12
+from vllm.platforms import current_platform
13
+from vllm.v1.watermarking import (
14
+ DualKeyGumbelWatermarker,
15
+ SupportsSpeculativeDecoding,
16
+ create_watermarker,
17
+ derive_watermark_key,
18
+)
19
from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
14
-from vllm.v1.watermarking.watermarker import WatermarkSample
20
+from vllm.v1.watermarking.gumbel import GumbelWatermarker
21
+from vllm.v1.watermarking.spec_decode import (
22
+ DraftWatermarker,
23
+ _resolve_watermark_key,
24
+ create_speculative_draft_watermarker,
25
+ create_speculative_target_watermarker,
26
+ speculative_target_watermark_key,
27
+)
28
+from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample
29
from vllm.v1.worker.gpu.sample.sampler import Sampler
30
+from vllm.v1.worker.gpu.sample.watermark import (
31
+ philox_gumbel_sample,
32
+ repeated_context_mask,
33
+)
34
+from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
35
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
36
37
18
-@pytest.mark.parametrize("algorithm", ["gumbel"])
38
+class StubWatermarker(Watermarker):
39
+ context_width = 1
40
+
41
+ def _sample_watermarked(self, logits, contexts):
42
+ return WatermarkSample(torch.tensor([7, 7]), logits + 10)
43
+
44
+
45
+def _argmax_sampler(logits: torch.Tensor) -> torch.Tensor:
46
+ return logits.argmax(dim=-1)
47
+
48
+
49
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
50
def test_watermarker_contract(algorithm: str):
51
watermarker = create_watermarker(
52
WatermarkConfig(algorithm=algorithm, key=42, context_width=4)
53
)
54
logits = torch.zeros(2, 128)
55
contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]])
25
- random_sample = lambda sample_logits: sample_logits.argmax(dim=-1)
56
27
- first = watermarker.sample(logits, contexts, random_sample)
28
- second = watermarker.sample(logits, contexts, random_sample)
57
+ first = watermarker.sample(logits, contexts, _argmax_sampler)
58
+ second = watermarker.sample(logits, contexts, _argmax_sampler)
59
60
assert first.token_ids.shape == (2,)
61
assert first.logits.shape == logits.shape
62
assert torch.equal(first.token_ids, second.token_ids)
63
64
35
-def test_gumbel_config_warns_about_degenerate_generations(monkeypatch):
65
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
66
+@pytest.mark.parametrize(
67
+ "config_overrides",
68
+ [
69
+ {"deduplicate_contexts": "none"},
70
+ {"deduplicate_contexts_max_history": 1023},
71
+ ],
72
+)
73
+def test_gumbel_config_warns_when_context_deduplication_is_weak(
74
+ monkeypatch, algorithm, config_overrides
75
+):
76
messages: list[str] = []
77
monkeypatch.setattr(
78
"vllm.config.watermarking.logger.warning_once",
79
lambda message, *, scope: messages.append(message),
80
)
81
42
- WatermarkConfig(key=42)
82
+ WatermarkConfig(key=42, algorithm=algorithm)
83
+ WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=1024)
84
+ WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=None)
85
+ WatermarkConfig(key=42, algorithm=algorithm, **config_overrides)
86
87
assert messages == [
45
- "Single-key Gumbel-max watermarking may increase the frequency of "
46
- "degenerate generations, including repetition loops."
88
+ (
89
+ "Gumbel-max watermarking with context deduplication disabled "
90
+ "or limited to fewer than 1024 positions may increase the frequency of "
91
+ "degenerate generations, including repetition loops. Use "
92
+ "deduplicate_contexts='single_turn' or 'all' with "
93
+ "deduplicate_contexts_max_history at least 1024 or null to mitigate this."
94
+ )
95
]
96
97
98
+def test_context_deduplication_history_can_be_unbounded():
99
+ config = WatermarkConfig(key=42, deduplicate_contexts_max_history=None)
100
+
101
+ assert config.deduplicate_contexts_max_history is None
102
+
103
+
104
def test_large_context_width_warns_but_is_allowed():
105
config = WatermarkConfig(key=42, context_width=17)
106
@@ -56,6 +110,158 @@ def test_large_context_width_warns_but_is_allowed():
110
assert watermarker.context_width == 17
111
112
113
+def test_dual_key_watermarker_uses_domain_separated_keys():
114
+ config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
115
+
116
+ watermarker = create_watermarker(config)
117
+ assert isinstance(watermarker, SupportsSpeculativeDecoding)
118
+ draft = watermarker.draft_watermarker
119
+ target = watermarker.target_watermarker
120
+
121
+ assert isinstance(watermarker, DualKeyGumbelWatermarker)
122
+ assert watermarker.alpha == 0.1
123
+ assert draft.prf.key == derive_watermark_key(42, b"key_a")
124
+ assert target.prf.key == derive_watermark_key(42, b"key_b")
125
+ assert target.prf.key != draft.prf.key
126
+
127
+
128
+def test_dual_key_watermarker_routes_tokens_with_alpha():
129
+ watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25)
130
+ logits = torch.zeros(2, 16)
131
+ contexts = torch.tensor([[1, 2], [3, 4]])
132
+ key_a = watermarker.draft_watermarker.sample(logits, contexts, lambda _: None)
133
+ key_b = watermarker.target_watermarker.sample(logits, contexts, lambda _: None)
134
+
135
+ def route(routing_logits):
136
+ torch.testing.assert_close(
137
+ routing_logits.softmax(dim=-1),
138
+ torch.tensor([[0.75, 0.25], [0.75, 0.25]]),
139
+ )
140
+ return torch.tensor([0, 1])
141
+
142
+ sampled = watermarker.sample(logits, contexts, route)
143
+
144
+ assert torch.equal(
145
+ sampled.token_ids, torch.stack([key_a.token_ids[0], key_b.token_ids[1]])
146
+ )
147
+
148
+
149
+def test_dual_key_routing_logits_are_cached():
150
+ alpha = 0.25
151
+ watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha)
152
+ logits = torch.zeros(2, 16)
153
+ contexts = torch.tensor([[1, 2], [3, 4]])
154
+ recorded: list[torch.Tensor] = []
155
+
156
+ def route(routing_logits):
157
+ recorded.append(routing_logits)
158
+ return torch.tensor([0, 1])
159
+
160
+ watermarker.sample(logits, contexts, route)
161
+ watermarker.sample(logits, contexts, route)
162
+
163
+ expected = torch.tensor([1 - alpha, alpha], dtype=torch.float32).log()
164
+ expected = expected.expand(2, -1)
165
+ assert len(recorded) == 2
166
+ assert all(torch.equal(seen, expected) for seen in recorded)
167
+ assert len(watermarker._routing_logits_cache) == 1
168
+
169
+
170
+@pytest.mark.parametrize("alpha", [0.0, 0.25, 1.0])
171
+def test_dual_key_watermarker_skips_masked_rows(alpha: float):
172
+ watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha)
173
+ torch.manual_seed(0)
174
+ logits = torch.randn(4, 32)
175
+ contexts = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
176
+ skip_mask = torch.tensor([True, False, True, False])
177
+
178
+ unmasked = watermarker.sample(logits, contexts, _argmax_sampler)
179
+ masked = watermarker.sample(logits, contexts, _argmax_sampler, skip_mask)
180
+
181
+ ordinary = _argmax_sampler(logits)
182
+ assert torch.equal(masked.token_ids[skip_mask], ordinary[skip_mask])
183
+ assert torch.equal(masked.token_ids[~skip_mask], unmasked.token_ids[~skip_mask])
184
+ assert not torch.equal(unmasked.token_ids[skip_mask], ordinary[skip_mask])
185
+
186
+
187
+def test_dual_key_watermarker_routing_requires_a_random_sampler():
188
+ watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25)
189
+
190
+ with pytest.raises(ValueError, match="random sampler"):
191
+ watermarker.sample(torch.zeros(2, 16), torch.tensor([[1, 2], [3, 4]]))
192
+
193
+
194
+def test_speculative_decoding_uses_fixed_dual_key_roles():
195
+ watermarker = create_watermarker(
196
+ WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25)
197
+ )
198
+
199
+ target = create_speculative_target_watermarker(watermarker)
200
+ draft = create_speculative_draft_watermarker(
201
+ watermarker,
202
+ max_num_reqs=1,
203
+ device=torch.device("cpu"),
204
+ allow_target_only=False,
205
+ )
206
+
207
+ assert target.prf.key == derive_watermark_key(42, b"key_b")
208
+ assert draft is not None
209
+ assert draft.watermarker.prf.key == derive_watermark_key(42, b"key_a")
210
+
211
+
212
+def test_recovery_key_rejects_the_unsplit_dual_key_watermarker():
213
+ watermarker = create_watermarker(
214
+ WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25)
215
+ )
216
+
217
+ with pytest.raises(ValueError, match="keys the target role separately"):
218
+ _resolve_watermark_key(watermarker)
219
+
220
+ target = create_speculative_target_watermarker(watermarker)
221
+ assert _resolve_watermark_key(target) == derive_watermark_key(42, b"key_b")
222
+
223
+
224
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
225
+def test_config_key_resolution_matches_the_model_runner(algorithm: str):
226
+ config = WatermarkConfig(algorithm=algorithm, key=42, alpha=0.25)
227
+ runtime_key = _resolve_watermark_key(
228
+ create_speculative_target_watermarker(create_watermarker(config))
229
+ )
230
+
231
+ assert speculative_target_watermark_key(config) == runtime_key
232
+
233
+
234
+def test_config_key_resolution_returns_none_when_watermarking_is_disabled():
235
+ assert speculative_target_watermark_key(None) is None
236
+
237
+
238
+def test_target_only_speculative_watermarking_skips_draft_watermarker():
239
+ watermarker = create_watermarker(WatermarkConfig(algorithm="gumbel", key=42))
240
+
241
+ assert not isinstance(watermarker, SupportsSpeculativeDecoding)
242
+ with pytest.raises(ValueError, match="does not support speculative decoding"):
243
+ create_speculative_draft_watermarker(
244
+ watermarker,
245
+ max_num_reqs=1,
246
+ device=torch.device("cpu"),
247
+ allow_target_only=False,
248
+ )
249
+ assert (
250
+ create_speculative_draft_watermarker(
251
+ watermarker,
252
+ max_num_reqs=1,
253
+ device=torch.device("cpu"),
254
+ allow_target_only=True,
255
+ )
256
+ is None
257
+ )
258
+
259
+
260
+def test_dual_key_derivation_is_stable():
261
+ assert derive_watermark_key(32, b"key_a") == 16368605726115524094
262
+ assert derive_watermark_key(32, b"key_b") == 4799302812959726346
263
+
264
+
265
def test_sampling_params_can_disable_watermarking():
266
assert SamplingParams().watermarking
267
assert not SamplingParams.from_optional(watermarking=False).watermarking
@@ -83,14 +289,9 @@ def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch):
289
290
291
def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch):
86
- class StubWatermarker:
87
- context_width = 1
88
-
89
- def sample(self, logits, contexts, random_sample):
90
- return WatermarkSample(torch.tensor([7, 7]), logits + 10)
91
-
292
sampler = object.__new__(GPUWatermarkSampler)
293
sampler.watermarker = StubWatermarker()
294
+ sampler.deduplicate_contexts = "single_turn"
295
sampler.watermarking = SimpleNamespace(
296
np=np.array([True, False]), gpu=torch.tensor([True, False])
297
)
@@ -102,8 +303,11 @@ def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch):
303
sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
304
2, 1, dtype=torch.int64
305
)
306
+ sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: torch.zeros(
307
+ 2, dtype=torch.bool
308
+ )
309
monkeypatch.setattr(
106
- "vllm.v1.watermarking.gpu_sampler.gumbel_sample",
310
+ "vllm.v1.watermarking.watermarker.gumbel_sample",
311
lambda *args, **kwargs: torch.tensor([3, 4]),
312
)
313
logits = torch.zeros(2, 8)
@@ -123,11 +327,524 @@ def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch):
327
assert torch.equal(output_logits[1], logits[1])
328
329
330
+def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch):
331
+ sampler = object.__new__(GPUWatermarkSampler)
332
+ sampler.watermarker = StubWatermarker()
333
+ sampler.deduplicate_contexts = "single_turn"
334
+ sampler.watermarking = SimpleNamespace(
335
+ np=np.array([True, True]), gpu=torch.tensor([True, True])
336
+ )
337
+ sampler.sampling_states = SimpleNamespace(
338
+ temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)),
339
+ seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)),
340
+ )
341
+ sampler.use_fp64_gumbel = False
342
+ sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
343
+ 2, 1, dtype=torch.int64
344
+ )
345
+ sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: (
346
+ torch.tensor([True, False])
347
+ )
348
+ monkeypatch.setattr(
349
+ "vllm.v1.watermarking.watermarker.gumbel_sample",
350
+ lambda *args, **kwargs: torch.tensor([3, 4]),
351
+ )
352
+ logits = torch.zeros(2, 8)
353
+
354
+ sampled, output_logits = sampler._sample_random(
355
+ logits,
356
+ torch.tensor([0, 1]),
357
+ np.array([0, 1]),
358
+ torch.zeros(2, dtype=torch.int64),
359
+ None,
360
+ None,
361
+ False,
362
+ )
363
+
364
+ assert torch.equal(sampled, torch.tensor([3, 7]))
365
+ assert torch.equal(output_logits[0], logits[0])
366
+ assert torch.equal(output_logits[1], torch.full((8,), 10.0))
367
+
368
+
369
+def test_gpu_sampler_can_disable_context_deduplication(monkeypatch):
370
+ sampler = object.__new__(GPUWatermarkSampler)
371
+ sampler.watermarker = StubWatermarker()
372
+ sampler.deduplicate_contexts = "none"
373
+ sampler.watermarking = SimpleNamespace(
374
+ np=np.array([True, True]), gpu=torch.tensor([True, True])
375
+ )
376
+ sampler.sampling_states = SimpleNamespace(
377
+ temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)),
378
+ seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)),
379
+ )
380
+ sampler.use_fp64_gumbel = False
381
+ sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
382
+ 2, 1, dtype=torch.int64
383
+ )
384
+ sampler._get_repeated_contexts = lambda *args: pytest.fail(
385
+ "context deduplication should not run"
386
+ )
387
+ logits = torch.zeros(2, 8)
388
+
389
+ sampled, output_logits = sampler._sample_random(
390
+ logits,
391
+ torch.tensor([0, 1]),
392
+ np.array([0, 1]),
393
+ torch.zeros(2, dtype=torch.int64),
394
+ None,
395
+ None,
396
+ False,
397
+ )
398
+
399
+ assert torch.equal(sampled, torch.tensor([7, 7]))
400
+ assert torch.equal(output_logits, torch.full((2, 8), 10.0))
401
+
402
+
403
+def test_repeated_context_mask_ignores_prompt_tokens():
404
+ all_token_ids = torch.tensor(
405
+ [
406
+ [8, 9, 1, 2, 1, 2],
407
+ [3, 4, 1, 2, 3, 4],
408
+ [0, 0, 0, 0, 0, 0],
409
+ ],
410
+ dtype=torch.int32,
411
+ )
412
+ req_indices = torch.tensor([0, 1, -1])
413
+ prompt_lens = torch.tensor([2, 2, 0])
414
+ total_lens = torch.tensor([6, 6, 0])
415
+ contexts = torch.tensor([[1, 2], [3, 4], [-1, -1]])
416
+
417
+ repeated = repeated_context_mask(
418
+ all_token_ids,
419
+ req_indices,
420
+ prompt_lens,
421
+ total_lens,
422
+ contexts,
423
+ )
424
+
425
+ assert torch.equal(repeated, torch.tensor([True, False, False]))
426
+
427
+
428
+def test_repeated_context_mask_can_include_prompt_tokens():
429
+ all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
430
+ req_indices = torch.tensor([0])
431
+ prompt_lens = torch.tensor([4])
432
+ total_lens = torch.tensor([10])
433
+ contexts = torch.tensor([[1, 2, 3, 4]])
434
+
435
+ single_turn = repeated_context_mask(
436
+ all_token_ids, req_indices, prompt_lens, total_lens, contexts
437
+ )
438
+ all_history = repeated_context_mask(
439
+ all_token_ids,
440
+ req_indices,
441
+ prompt_lens,
442
+ total_lens,
443
+ contexts,
444
+ include_prompt=True,
445
+ )
446
+
447
+ assert not single_turn.item()
448
+ assert all_history.item()
449
+
450
+
451
+def test_repeated_context_mask_can_skip_partial_contexts():
452
+ all_token_ids = torch.tensor([[10, 11, 12, 13, 1, 2, 3, 4]], dtype=torch.int32)
453
+ req_indices = torch.tensor([0])
454
+ prompt_lens = torch.tensor([4])
455
+ contexts = torch.tensor([[-1, -1, -1, 1]])
456
+
457
+ partial = repeated_context_mask(
458
+ all_token_ids,
459
+ req_indices,
460
+ prompt_lens,
461
+ torch.tensor([5]),
462
+ contexts,
463
+ include_prompt=True,
464
+ skip_partial_context=True,
465
+ )
466
+ complete = repeated_context_mask(
467
+ all_token_ids,
468
+ req_indices,
469
+ prompt_lens,
470
+ torch.tensor([8]),
471
+ torch.tensor([[1, 2, 3, 4]]),
472
+ include_prompt=True,
473
+ skip_partial_context=True,
474
+ )
475
+
476
+ assert partial.item()
477
+ assert not complete.item()
478
+
479
+
480
+def test_gpu_sampler_contexts_are_completion_local_for_all_scopes():
481
+ sampler = object.__new__(GPUWatermarkSampler)
482
+ sampler.watermarker = SimpleNamespace(context_width=4)
483
+ sampler.req_states = SimpleNamespace(
484
+ all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])),
485
+ prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
486
+ total_len=SimpleNamespace(gpu=torch.tensor([6])),
487
+ )
488
+ request_indices = torch.tensor([0])
489
+
490
+ sampler.deduplicate_contexts = "single_turn"
491
+ single_turn = sampler._get_contexts(request_indices)
492
+ sampler.deduplicate_contexts = "all"
493
+ all_history = sampler._get_contexts(request_indices)
494
+
495
+ assert torch.equal(single_turn, torch.tensor([[-1, -1, 1, 2]]))
496
+ assert torch.equal(all_history, single_turn)
497
+
498
+
499
+def test_gpu_sampler_all_history_skips_partial_context():
500
+ sampler = object.__new__(GPUWatermarkSampler)
501
+ sampler.watermarker = SimpleNamespace(context_width=4)
502
+ sampler.req_states = SimpleNamespace(
503
+ all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])),
504
+ prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
505
+ total_len=SimpleNamespace(gpu=torch.tensor([6])),
506
+ )
507
+ sampler.deduplicate_contexts_max_history = None
508
+ request_indices = torch.tensor([0])
509
+ contexts = sampler._get_contexts(request_indices)
510
+
511
+ sampler.deduplicate_contexts = "single_turn"
512
+ single_turn = sampler._get_repeated_contexts(request_indices, contexts)
513
+ sampler.deduplicate_contexts = "all"
514
+ all_history = sampler._get_repeated_contexts(request_indices, contexts)
515
+
516
+ assert not single_turn.item()
517
+ assert all_history.item()
518
+
519
+
520
+def test_repeated_context_mask_respects_max_history():
521
+ all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
522
+ req_indices = torch.tensor([0])
523
+ prompt_lens = torch.tensor([0])
524
+ total_lens = torch.tensor([10])
525
+ contexts = torch.tensor([[1, 2, 3, 4]])
526
+
527
+ full_history = repeated_context_mask(
528
+ all_token_ids, req_indices, prompt_lens, total_lens, contexts
529
+ )
530
+ last_six = repeated_context_mask(
531
+ all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=6
532
+ )
533
+ last_five = repeated_context_mask(
534
+ all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=5
535
+ )
536
+
537
+ assert full_history.item()
538
+ assert last_six.item()
539
+ assert not last_five.item()
540
+
541
+
542
+@pytest.mark.parametrize("max_history", [0, -1])
543
+def test_repeated_context_mask_rejects_non_positive_max_history(max_history: int):
544
+ with pytest.raises(ValueError, match="max_history must be positive or None"):
545
+ repeated_context_mask(
546
+ torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
547
+ torch.tensor([0], dtype=torch.int32),
548
+ torch.tensor([0], dtype=torch.int32),
549
+ torch.tensor([4], dtype=torch.int32),
550
+ torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
551
+ max_history=max_history,
552
+ )
553
+
554
+
555
+def test_gpu_sampler_wires_request_history_scope():
556
+ sampler = object.__new__(GPUWatermarkSampler)
557
+ sampler.req_states = SimpleNamespace(
558
+ all_token_ids=SimpleNamespace(
559
+ gpu=torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
560
+ ),
561
+ prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
562
+ total_len=SimpleNamespace(gpu=torch.tensor([10])),
563
+ )
564
+ sampler.deduplicate_contexts_max_history = None
565
+ request_indices = torch.tensor([0])
566
+ contexts = torch.tensor([[1, 2, 3, 4]])
567
+
568
+ sampler.deduplicate_contexts = "single_turn"
569
+ single_turn = sampler._get_repeated_contexts(request_indices, contexts)
570
+ sampler.deduplicate_contexts = "all"
571
+ all_history = sampler._get_repeated_contexts(request_indices, contexts)
572
+
573
+ assert not single_turn.item()
574
+ assert all_history.item()
575
+
576
+
577
+@pytest.mark.skipif(
578
+ not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
579
+)
580
+@pytest.mark.parametrize("max_history", [5, 6])
581
+@pytest.mark.parametrize("include_prompt", [False, True])
582
+def test_repeated_context_mask_max_history_accelerator_parity(
583
+ max_history: int, include_prompt: bool
584
+):
585
+ inputs = (
586
+ torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]),
587
+ torch.tensor([0]),
588
+ torch.tensor([0]),
589
+ torch.tensor([10]),
590
+ torch.tensor([[1, 2, 3, 4]]),
591
+ )
592
+
593
+ expected = repeated_context_mask(
594
+ *inputs, max_history=max_history, include_prompt=include_prompt
595
+ )
596
+ actual = repeated_context_mask(
597
+ *(value.cuda() for value in inputs),
598
+ max_history=max_history,
599
+ include_prompt=include_prompt,
600
+ ).cpu()
601
+
602
+ assert torch.equal(actual, expected)
603
+
604
+
605
+@pytest.mark.skipif(
606
+ not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
607
+)
608
+def test_repeated_context_mask_partial_context_accelerator_parity():
609
+ inputs = (
610
+ torch.tensor([[10, 11, 12, 13, 1]], dtype=torch.int32),
611
+ torch.tensor([0], dtype=torch.int32),
612
+ torch.tensor([4], dtype=torch.int32),
613
+ torch.tensor([5], dtype=torch.int32),
614
+ torch.tensor([[-1, -1, -1, 1]], dtype=torch.int64),
615
+ )
616
+
617
+ expected = repeated_context_mask(
618
+ *inputs, include_prompt=True, skip_partial_context=True
619
+ )
620
+ actual = repeated_context_mask(
621
+ *(value.cuda() for value in inputs),
622
+ include_prompt=True,
623
+ skip_partial_context=True,
624
+ ).cpu()
625
+
626
+ assert torch.equal(actual, expected)
627
+
628
+
629
+def _unaligned_max_history_inputs() -> tuple[torch.Tensor, ...]:
630
+ """One row whose only repeated context ends at position 680.
631
+
632
+ 680 is not a multiple of the kernel's 512-position block, so the scan window
633
+ selected by ``max_history`` starts in the middle of a block.
634
+ """
635
+ history = list(range(1_000, 2_200))
636
+ history[676:680] = [1, 2, 3, 4]
637
+ return (
638
+ torch.tensor([history], dtype=torch.int32),
639
+ torch.tensor([0], dtype=torch.int32),
640
+ torch.tensor([0], dtype=torch.int32),
641
+ torch.tensor([len(history)], dtype=torch.int32),
642
+ torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
643
+ )
644
+
645
+
646
+@pytest.mark.skipif(
647
+ not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
648
+)
649
+@pytest.mark.parametrize(
650
+ "max_history,expected", [(1200 - 680, True), (1200 - 680 - 1, False)]
651
+)
652
+def test_repeated_context_mask_unaligned_max_history_accelerator_parity(
653
+ max_history: int, expected: bool
654
+):
655
+ inputs = _unaligned_max_history_inputs()
656
+
657
+ reference = repeated_context_mask(*inputs, max_history=max_history)
658
+ actual = repeated_context_mask(
659
+ *(value.cuda() for value in inputs), max_history=max_history
660
+ ).cpu()
661
+
662
+ assert torch.equal(reference, torch.tensor([expected]))
663
+ assert torch.equal(actual, reference)
664
+
665
+
666
+def _repeated_context_inputs(
667
+ context_width: int, device: str = "cpu"
668
+) -> tuple[torch.Tensor, ...]:
669
+ """Build repeated, unique, and padded request rows for mask tests."""
670
+ prompt = [101, 102, 103, 104]
671
+ repeated_output = [*range(1, context_width + 1)] * 2
672
+ unique_output = [*range(1, context_width + 2)]
673
+ max_len = len(prompt) + len(repeated_output)
674
+ rows = [
675
+ prompt + repeated_output,
676
+ prompt + unique_output,
677
+ [],
678
+ ]
679
+ all_token_ids = torch.zeros((3, max_len), dtype=torch.int32, device=device)
680
+ for row, token_ids in enumerate(rows):
681
+ all_token_ids[row, : len(token_ids)] = torch.tensor(
682
+ token_ids, dtype=torch.int32, device=device
683
+ )
684
+ return (
685
+ all_token_ids,
686
+ torch.tensor([0, 1, -1], dtype=torch.int32, device=device),
687
+ torch.tensor([len(prompt), len(prompt), 0], dtype=torch.int32, device=device),
688
+ torch.tensor([len(rows[0]), len(rows[1]), 0], dtype=torch.int32, device=device),
689
+ torch.tensor(
690
+ [
691
+ repeated_output[-context_width:],
692
+ unique_output[-context_width:],
693
+ [-1] * context_width,
694
+ ],
695
+ dtype=torch.int64,
696
+ device=device,
697
+ ),
698
+ )
699
+
700
+
701
+@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17])
702
+def test_repeated_context_mask(context_width: int):
703
+ repeated = repeated_context_mask(*_repeated_context_inputs(context_width))
704
+
705
+ assert torch.equal(repeated, torch.tensor([True, False, False]))
706
+
707
+
708
+@pytest.mark.skipif(
709
+ not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
710
+)
711
+@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17])
712
+def test_repeated_context_mask_accelerator_parity(context_width: int):
713
+ cpu_inputs = _repeated_context_inputs(context_width)
714
+ accelerator_inputs = list(_repeated_context_inputs(context_width, "cuda"))
715
+ contexts = accelerator_inputs[-1]
716
+ storage = torch.empty(
717
+ contexts.shape[0], contexts.shape[1] * 2, dtype=contexts.dtype, device="cuda"
718
+ )
719
+ storage[:, ::2] = contexts
720
+ accelerator_inputs[-1] = storage[:, ::2]
721
+
722
+ expected = repeated_context_mask(*cpu_inputs)
723
+ actual = repeated_context_mask(*accelerator_inputs).cpu()
724
+
725
+ assert torch.equal(actual, expected)
726
+
727
+
728
+@pytest.mark.skipif(
729
+ not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
730
+)
731
+def test_repeated_context_mask_scans_multiple_blocks():
732
+ context = [1, 2, 3, 4]
733
+ repeated_output = [*range(10_000, 11_030), *context, 99, *context]
734
+ unique_output = list(range(20_000, 21_039))
735
+ prompt = [101, 102, 103, 104]
736
+ rows = [prompt + repeated_output, prompt + unique_output]
737
+ all_token_ids = torch.zeros((2, len(rows[0])), dtype=torch.int32)
738
+ for row, token_ids in enumerate(rows):
739
+ all_token_ids[row, : len(token_ids)] = torch.tensor(token_ids)
740
+ inputs = (
741
+ all_token_ids,
742
+ torch.tensor([0, 1], dtype=torch.int32),
743
+ torch.tensor([len(prompt), len(prompt)], dtype=torch.int32),
744
+ torch.tensor([len(rows[0]), len(rows[1])], dtype=torch.int32),
745
+ torch.tensor([context, unique_output[-4:]], dtype=torch.int64),
746
+ )
747
+
748
+ expected = repeated_context_mask(*inputs)
749
+ actual = repeated_context_mask(*(value.cuda() for value in inputs)).cpu()
750
+
751
+ assert torch.equal(expected, torch.tensor([True, False]))
752
+ assert torch.equal(actual, expected)
753
+
754
+
755
+@pytest.mark.skipif(
756
+ not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
757
+)
758
+def test_repeated_context_mask_last_request_at_capacity():
759
+ max_model_len = 1024
760
+ all_token_ids = torch.arange(
761
+ 2 * max_model_len, dtype=torch.int32, device="cuda"
762
+ ).reshape(2, max_model_len)
763
+ inputs = (
764
+ all_token_ids,
765
+ torch.tensor([1], dtype=torch.int32, device="cuda"),
766
+ torch.tensor([0, 0], dtype=torch.int32, device="cuda"),
767
+ torch.tensor([max_model_len, max_model_len], dtype=torch.int32, device="cuda"),
768
+ all_token_ids[1, -4:].to(torch.int64).unsqueeze(0),
769
+ )
770
+
771
+ repeated = repeated_context_mask(*inputs)
772
+
773
+ assert not repeated.item()
774
+
775
+
776
+@pytest.mark.skipif(
777
+ not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
778
+)
779
+def test_gpu_sampler_uses_fused_gumbel_for_repeated_contexts():
780
+ device = torch.device("cuda")
781
+ logits = torch.tensor([[0.0, 1.0, 2.0, 3.0], [3.0, 2.0, 1.0, 0.0]], device=device)
782
+ all_token_ids = torch.tensor(
783
+ [[101, 102, 1, 2, 1, 2], [101, 102, 3, 4, 5, 6]],
784
+ dtype=torch.int32,
785
+ device=device,
786
+ )
787
+ request_indices = torch.tensor([0, 1], dtype=torch.int64, device=device)
788
+ request_indices_np = np.array([0, 1])
789
+ temperatures = torch.ones(2, device=device)
790
+ seeds = torch.tensor([11, 22], dtype=torch.int64, device=device)
791
+ positions = torch.tensor([6, 6], dtype=torch.int64, device=device)
792
+
793
+ sampler = object.__new__(GPUWatermarkSampler)
794
+ sampler.watermarker = GumbelWatermarker(key=42, context_width=1)
795
+ sampler.deduplicate_contexts = "single_turn"
796
+ sampler.deduplicate_contexts_max_history = None
797
+ sampler.watermarking = SimpleNamespace(
798
+ np=np.ones(2, dtype=bool), gpu=torch.ones(2, dtype=torch.bool, device=device)
799
+ )
800
+ sampler.sampling_states = SimpleNamespace(
801
+ temperature=SimpleNamespace(np=np.ones(2), gpu=temperatures),
802
+ seeds=SimpleNamespace(gpu=seeds),
803
+ )
804
+ sampler.req_states = SimpleNamespace(
805
+ all_token_ids=SimpleNamespace(gpu=all_token_ids),
806
+ prompt_len=SimpleNamespace(
807
+ gpu=torch.tensor([2, 2], dtype=torch.int32, device=device)
808
+ ),
809
+ total_len=SimpleNamespace(
810
+ gpu=torch.tensor([6, 6], dtype=torch.int32, device=device)
811
+ ),
812
+ )
813
+ sampler.use_fp64_gumbel = False
814
+
815
+ contexts = sampler._get_contexts(request_indices)
816
+ repeated = sampler._get_repeated_contexts(request_indices, contexts)
817
+ expected = philox_gumbel_sample(
818
+ logits,
819
+ contexts,
820
+ 42,
821
+ skip_mask=repeated,
822
+ expanded_idx_mapping=request_indices,
823
+ temperatures=temperatures,
824
+ seeds=seeds,
825
+ positions=positions,
826
+ )
827
+
828
+ actual, output_logits = sampler._sample_random(
829
+ logits,
830
+ request_indices,
831
+ request_indices_np,
832
+ positions,
833
+ None,
834
+ None,
835
+ False,
836
+ )
837
+
838
+ assert torch.equal(repeated, torch.tensor([True, False], device=device))
839
+ assert torch.equal(actual, expected)
840
+ assert output_logits is logits
841
+
842
+
843
def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch):
844
class StubWatermarker:
845
context_width = 1
846
130
- def sample(self, logits, contexts, random_sample):
847
+ def sample(self, logits, contexts, random_sampler=None, skip_mask=None):
848
raise AssertionError("watermarker should not run for greedy requests")
849
850
sampler = object.__new__(GPUWatermarkSampler)
@@ -155,3 +872,157 @@ def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch):
872
)
873
874
assert actual is expected
875
+
876
+
877
+def test_gpu_sampler_builds_speculative_contexts_from_drafts():
878
+ sampler = object.__new__(GPUWatermarkSampler)
879
+ sampler.watermarker = SimpleNamespace(context_width=2)
880
+ sampler.req_states = SimpleNamespace(
881
+ all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 101, 10, 11, 0, 0]])),
882
+ prompt_len=SimpleNamespace(gpu=torch.tensor([2])),
883
+ total_len=SimpleNamespace(gpu=torch.tensor([4])),
884
+ )
885
+
886
+ contexts = sampler._get_contexts(
887
+ torch.tensor([0, 0, 0]),
888
+ torch.tensor([0, 1, 2]),
889
+ torch.tensor([11, 20, 21]),
890
+ )
891
+
892
+ assert torch.equal(contexts, torch.tensor([[10, 11], [11, 20], [20, 21]]))
893
+
894
+
895
+def test_gpu_sampler_builds_chunked_multi_request_speculative_contexts():
896
+ sampler = object.__new__(GPUWatermarkSampler)
897
+ sampler.watermarker = SimpleNamespace(context_width=2)
898
+ sampler.req_states = SimpleNamespace(
899
+ all_token_ids=SimpleNamespace(
900
+ gpu=torch.tensor(
901
+ [
902
+ [100, 101, 10, 11, 0, 0],
903
+ [200, 201, 30, 31, 0, 0],
904
+ [300, 301, 50, 51, 0, 0],
905
+ ]
906
+ )
907
+ ),
908
+ prompt_len=SimpleNamespace(gpu=torch.tensor([2, 2, 2])),
909
+ total_len=SimpleNamespace(gpu=torch.tensor([4, 4, 4])),
910
+ )
911
+
912
+ contexts = sampler._get_contexts(
913
+ torch.tensor([2, 2, 1, 1]),
914
+ torch.tensor([0, 1, 0, 1]),
915
+ torch.tensor([51, 60, 31, 40]),
916
+ )
917
+
918
+ assert torch.equal(
919
+ contexts,
920
+ torch.tensor([[50, 51], [51, 60], [30, 31], [31, 40]]),
921
+ )
922
+
923
+
924
+def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch):
925
+ class StubSpeculator(DraftModelSpeculator):
926
+ def capture(self): ...
927
+
928
+ def init_cudagraph_manager(self, cudagraph_mode): ...
929
+
930
+ def load_draft_model(self, target_model, target_attn_layer_names): ...
931
+
932
+ def propose(self, *args, **kwargs): ...
933
+
934
+ class StubModel:
935
+ @staticmethod
936
+ def compute_logits(hidden_states):
937
+ return torch.zeros(hidden_states.shape[0], 8)
938
+
939
+ class StubWatermarker:
940
+ @staticmethod
941
+ def sample(logits, contexts, random_sampler=None):
942
+ assert random_sampler is None
943
+ return WatermarkSample(torch.tensor([7, 7]), logits)
944
+
945
+ speculator = object.__new__(StubSpeculator)
946
+ speculator.model = StubModel()
947
+ speculator.use_fp64_gumbel = False
948
+ draft_watermarker = object.__new__(DraftWatermarker)
949
+ draft_watermarker.watermarker = StubWatermarker()
950
+ draft_watermarker.contexts = torch.tensor([[1, 2], [3, 4]])
951
+ draft_watermarker.enabled = torch.tensor([True, False])
952
+ speculator.draft_watermarker = draft_watermarker
953
+ monkeypatch.setattr(
954
+ "vllm.v1.worker.gpu.spec_decode.speculator.gumbel_sample",
955
+ lambda *args, **kwargs: torch.tensor([3, 4]),
956
+ )
957
+
958
+ sampled = speculator.sample_draft(
959
+ hidden_states=torch.zeros(2, 4),
960
+ sample_src_positions=torch.zeros(2, dtype=torch.int64),
961
+ idx_mapping=torch.tensor([0, 1]),
962
+ temperature=torch.ones(2),
963
+ seeds=torch.zeros(2, dtype=torch.int64),
964
+ draft_step=torch.tensor(0),
965
+ draft_logits=torch.zeros(2, 1, 8),
966
+ )
967
+
968
+ assert torch.equal(sampled, torch.tensor([7, 4]))
969
+ assert torch.equal(draft_watermarker.contexts, torch.tensor([[2, 7], [4, 4]]))
970
+
971
+
972
+def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):
973
+ speculator = object.__new__(DSparkSpeculator)
974
+ speculator.draft_logits = torch.zeros(2, 1, 8)
975
+ speculator._d2t_scatter_index = torch.tensor([1, 5])
976
+ speculator._draft_scatter_buf = torch.full((2, 8), float("-inf"))
977
+ speculator.temperature = torch.ones(2)
978
+ speculator.seeds = torch.zeros(2, dtype=torch.int64)
979
+ speculator._step_cols = torch.tensor([0])
980
+ speculator.use_fp64_gumbel = False
981
+ watermark_logits: list[torch.Tensor] = []
982
+
983
+ def sample(logits, sampled, idx_map, temperature):
984
+ watermark_logits.append(logits.clone())
985
+ return sampled + 1
986
+
987
+ speculator.draft_watermarker = SimpleNamespace(sample=sample)
988
+ monkeypatch.setattr(
989
+ "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",
990
+ lambda *args, **kwargs: torch.tensor([3, 4]),
991
+ )
992
+
993
+ sampled = speculator._sample_logits(
994
+ torch.tensor([[10.0, 20.0], [30.0, 40.0]]),
995
+ torch.tensor([0, 1]),
996
+ torch.tensor([1, 1]),
997
+ 0,
998
+ )
999
+
1000
+ assert torch.equal(sampled, torch.tensor([4, 5]))
1001
+ assert torch.equal(
1002
+ watermark_logits[0][:, [1, 5]], torch.tensor([[10, 20], [30, 40]])
1003
+ )
1004
+ assert torch.isneginf(watermark_logits[0][:, [0, 2, 3, 4, 6, 7]]).all()
1005
+
1006
+
1007
+def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch):
1008
+ speculator = object.__new__(DSparkSpeculator)
1009
+ speculator.draft_logits = torch.zeros(2, 1, 8)
1010
+ speculator._d2t_scatter_index = None
1011
+ speculator.temperature = torch.ones(2)
1012
+ speculator.seeds = torch.zeros(2, dtype=torch.int64)
1013
+ speculator._step_cols = torch.tensor([0])
1014
+ speculator.use_fp64_gumbel = False
1015
+ speculator.draft_watermarker = None
1016
+ monkeypatch.setattr(
1017
+ "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",
1018
+ lambda *args, **kwargs: torch.tensor([3, 4]),
1019
+ )
1020
+
1021
+ sampled = speculator._sample_logits(
1022
+ torch.zeros(2, 8),
1023
+ torch.tensor([0, 1]),
1024
+ torch.tensor([1, 1]),
1025
+ 0,
1026
+ )
1027
+
1028
+ assert torch.equal(sampled, torch.tensor([3, 4]))
vllm/config/vllm.py
+62
-9
@@ -50,7 +50,7 @@ from .reasoning import ReasoningConfig
50
from .scheduler import SchedulerConfig
51
from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig
52
from .structured_outputs import StructuredOutputsConfig
53
-from .utils import SupportsHash, config, replace
53
+from .utils import SupportsHash, config, get_field, replace
54
from .watermarking import WatermarkConfig
55
from .weight_transfer import WeightTransferConfig
56
@@ -1142,14 +1142,67 @@ class VllmConfig:
1142
watermark_config = getattr(self, "watermark_config", None)
1143
if watermark_config is None:
1144
return
1145
- if (
1146
- self.speculative_config is not None
1147
- and not watermark_config.supports_speculative_decoding
1148
- ):
1149
- raise ValueError(
1150
- f"The {watermark_config.algorithm} watermarking algorithm "
1151
- "does not support speculative decoding."
1152
- )
1145
+ if self.speculative_config is not None:
1146
+ speculative_config = self.speculative_config
1147
+ if speculative_config.draft_sample_method != "probabilistic":
1148
+ raise ValueError(
1149
+ "Speculative decoding with watermarking requires "
1150
+ "draft_sample_method='probabilistic'."
1151
+ )
1152
+ if speculative_config.rejection_sample_method != "standard":
1153
+ raise ValueError(
1154
+ "Speculative decoding with watermarking requires "
1155
+ "rejection_sample_method='standard'."
1156
+ )
1157
+ if (
1158
+ speculative_config.parallel_drafting
1159
+ and speculative_config.method != "dspark"
1160
+ ):
1161
+ raise ValueError(
1162
+ "Parallel speculative drafting is not supported with watermarking."
1163
+ )
1164
+ if speculative_config.method not in ("dspark", "eagle", "eagle3", "mtp"):
1165
+ raise ValueError(
1166
+ "Watermarking supports only autoregressive model-based "
1167
+ "speculative decoding."
1168
+ )
1169
+ if (
1170
+ not watermark_config.allow_target_only_watermarking
1171
+ and not watermark_config.supports_speculative_decoding
1172
+ ):
1173
+ raise ValueError(
1174
+ f"The '{watermark_config.algorithm}' watermarking algorithm "
1175
+ "does not support speculative decoding. Set "
1176
+ "allow_target_only_watermarking=true to leave draft tokens "
1177
+ "unwatermarked."
1178
+ )
1179
+ if (
1180
+ watermark_config.allow_target_only_watermarking
1181
+ and not watermark_config.supports_speculative_decoding
1182
+ ):
1183
+ logger.warning_once(
1184
+ "Target-only watermarking leaves accepted draft tokens "
1185
+ "unwatermarked, weakening detectability in proportion to the "
1186
+ "share of output tokens supplied by accepted drafts.",
1187
+ scope="global",
1188
+ )
1189
+ if watermark_config.deduplicate_contexts != "none":
1190
+ logger.warning_once(
1191
+ "Context deduplication is not supported with speculative "
1192
+ "decoding and will not be applied to accepted drafts, "
1193
+ "rejection-recovery tokens, or bonus tokens.",
1194
+ scope="global",
1195
+ )
1196
+ if watermark_config.algorithm == "dual_key_gumbel" and (
1197
+ watermark_config.alpha != get_field(WatermarkConfig, "alpha").default
1198
+ ):
1199
+ logger.warning_once(
1200
+ "Speculative decoding selects the watermark key by token role: "
1201
+ "draft tokens use key A, recovery and bonus tokens use key B. "
1202
+ "The configured alpha=%s is not used.",
1203
+ watermark_config.alpha,
1204
+ scope="global",
1205
+ )
1206
if beam_search:
1207
raise ValueError("Beam search is not supported with watermarking.")
1208
if custom_sampler:
vllm/config/watermarking.py
+44
-13
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import hashlib
5
from typing import Literal
6
7
from pydantic import Field, model_validator
@@ -11,12 +12,16 @@ from vllm.logger import init_logger
12
13
logger = init_logger(__name__)
14
14
-WatermarkingAlgorithm = Literal["gumbel"]
15
+WatermarkingAlgorithm = Literal["gumbel", "dual_key_gumbel"]
16
WatermarkPRFName = Literal["philox"]
17
+WatermarkContextScope = Literal["none", "single_turn", "all"]
18
17
-_SPECULATIVE_DECODING_SUPPORT: dict[WatermarkingAlgorithm, bool] = {
18
- "gumbel": False,
19
-}
19
+_MIN_RECOMMENDED_DEDUP_HISTORY = 1024
20
+
21
+
22
+def derive_watermark_key(key: int, domain: bytes) -> int:
23
+ digest = hashlib.sha256(domain + key.to_bytes(8, "big")).digest()
24
+ return int.from_bytes(digest[:8], "big")
25
26
27
@config
@@ -27,23 +32,49 @@ class WatermarkConfig:
32
"""Secret key used to watermark generated text."""
33
algorithm: WatermarkingAlgorithm = "gumbel"
34
"""Algorithm used to watermark generated text."""
35
+ alpha: float = Field(default=0.1, ge=0, le=1)
36
+ """Probability of selecting key B for dual-key watermarking."""
37
context_width: int = Field(default=4, ge=1)
31
- """Number of prior output tokens used by the watermark PRF."""
38
+ """Number of prior tokens used by the watermark PRF."""
39
+ deduplicate_contexts: WatermarkContextScope = "single_turn"
40
+ """Which history is searched for a repeated context before a token is
41
+ sampled; a repeated context is sampled without the watermark. `none`
42
+ disables the search. `single_turn` (default) searches this request's
43
+ generated tokens. `all` also searches the prompt and samples the first
44
+ `context_width` generated tokens without watermarking."""
45
+ deduplicate_contexts_max_history: int | None = Field(default=8192, ge=1)
46
+ """Number of most recent history positions searched (default 8192), or
47
+ `None` for the whole scope. Each position is compared over the
48
+ `context_width` tokens before it. Ignored when `deduplicate_contexts` is
49
+ `none`."""
50
prf: WatermarkPRFName = "philox"
51
"""Pseudorandom function used by the watermarking algorithm."""
52
+ allow_target_only_watermarking: bool = False
53
+ """Allow speculative decoding without watermarking draft tokens."""
54
+
55
+ @property
56
+ def supports_speculative_decoding(self) -> bool:
57
+ return self.algorithm == "dual_key_gumbel"
58
59
@model_validator(mode="after")
36
- def validate_key(self) -> Self:
60
+ def validate_watermark_settings(self) -> Self:
61
if self.key > 2**64 - 1:
62
raise ValueError("philox keys must fit in 64 bits")
39
- if self.algorithm == "gumbel":
63
+ history_is_too_short = (
64
+ self.deduplicate_contexts_max_history is not None
65
+ and self.deduplicate_contexts_max_history < _MIN_RECOMMENDED_DEDUP_HISTORY
66
+ )
67
+ if self.algorithm in ("gumbel", "dual_key_gumbel") and (
68
+ self.deduplicate_contexts == "none" or history_is_too_short
69
+ ):
70
logger.warning_once(
41
- "Single-key Gumbel-max watermarking may increase the frequency of "
42
- "degenerate generations, including repetition loops.",
71
+ "Gumbel-max watermarking with context deduplication "
72
+ "disabled or limited to fewer than "
73
+ f"{_MIN_RECOMMENDED_DEDUP_HISTORY} positions may increase the "
74
+ "frequency of degenerate generations, including repetition loops. "
75
+ "Use deduplicate_contexts='single_turn' or 'all' with "
76
+ "deduplicate_contexts_max_history at least "
77
+ f"{_MIN_RECOMMENDED_DEDUP_HISTORY} or null to mitigate this.",
78
scope="global",
79
)
80
return self
46
-
47
- @property
48
- def supports_speculative_decoding(self) -> bool:
49
- return _SPECULATIVE_DECODING_SUPPORT[self.algorithm]
vllm/model_executor/warmup/kernel_warmup.py
+4
@@ -44,6 +44,9 @@ from vllm.model_executor.warmup.replayssm_warmup import (
44
from vllm.model_executor.warmup.spec_decode_rejection_warmup import (
45
spec_decode_rejection_warmup,
46
)
47
+from vllm.model_executor.warmup.watermark_sample_warmup import (
48
+ watermark_sample_warmup,
49
+)
50
from vllm.platforms import current_platform
51
from vllm.utils.deep_gemm import is_deep_gemm_supported
52
from vllm.utils.flashinfer import has_flashinfer
@@ -207,6 +210,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
210
if enable_jit_warmup:
211
kimi_k3_triton_warmup(worker)
212
spec_decode_rejection_warmup(worker)
213
+ watermark_sample_warmup(worker)
214
qwen4_exp_qsa_triton_warmup(worker)
215
216
if enable_jit_warmup and current_platform.is_device_capability_family(100):
vllm/model_executor/warmup/spec_decode_rejection_warmup.py
+54
-26
@@ -1,17 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Warm up spec-decode rejection-sampler Triton kernels.
4
-
5
-The rejection sampler kernels (``_compute_local_logits_stats_kernel``,
6
-``_rejection_kernel``, ``_resample_kernel``) are JIT-compiled by Triton on
7
-first use. Without warmup, the first spec-decode request pays a multi-second
8
-compilation cost. This pre-compiles them with dummy data matching the
9
-server's vocab size and speculative config.
10
-"""
3
+"""Warm up speculative-decoding rejection kernels."""
4
5
from __future__ import annotations
6
14
-from typing import TYPE_CHECKING
7
+from typing import TYPE_CHECKING, Any
8
9
import torch
10
@@ -39,42 +32,76 @@ def spec_decode_rejection_warmup(worker: Worker) -> None:
32
if num_spec <= 0 or vocab_size <= 0:
33
return
34
42
- # Mirror the constexpr-relevant flags the runtime uses.
35
rejection_method = getattr(spec_config, "rejection_sample_method", None)
36
use_block_verification = rejection_method == "block"
37
use_synthetic = rejection_method == "synthetic"
38
+ # HAS_DRAFT_LOGITS is a constexpr, so the warmup must match the runtime.
39
+ use_draft_logits = (
40
+ getattr(spec_config, "draft_sample_method", "greedy") == "probabilistic"
41
+ )
42
43
device = worker.device
44
num_reqs = 1
45
tokens_per_req = num_spec + 1
46
num_logits = num_reqs * tokens_per_req
47
52
- # Triton JIT-specializes on tensor dtypes. The target logits may be fp32
53
- # (apply_sampling_params copies to fp32 when processing is needed) or the
54
- # model dtype (pass-through otherwise), while draft logits are always the
55
- # model dtype. Warm every (target, draft) combination the runtime can hit.
48
+ watermark_config = getattr(worker.vllm_config, "watermark_config", None)
49
+ watermark_kwargs: dict[str, Any] = {}
50
+ if watermark_config is not None:
51
+ from vllm.v1.watermarking.spec_decode import speculative_target_watermark_key
52
+
53
+ try:
54
+ watermark_key = speculative_target_watermark_key(watermark_config)
55
+ except Exception:
56
+ logger.warning(
57
+ "Skipping spec-decode rejection sampler warmup: could not derive "
58
+ "the watermark recovery key.",
59
+ exc_info=True,
60
+ )
61
+ return
62
+ # Tensor dtypes and integer widths affect Triton specialization.
63
+ watermark_kwargs = {
64
+ "contexts": torch.zeros(
65
+ (num_logits, watermark_config.context_width),
66
+ dtype=torch.int32,
67
+ device=device,
68
+ ),
69
+ "watermarking": torch.zeros(num_reqs, dtype=torch.bool, device=device),
70
+ "watermark_key": watermark_key,
71
+ }
72
+
73
+ # Sampling-parameter processing may promote either distribution to fp32.
74
model_dtype = model_config.dtype
57
- warmup_dtype_pairs = {
58
- (model_dtype, model_dtype),
59
- (torch.float32, torch.float32),
60
- (torch.float32, model_dtype),
61
- (model_dtype, torch.float32),
62
- }
75
+ warmup_dtype_pairs: set[tuple[torch.dtype, torch.dtype | None]]
76
+ if use_draft_logits:
77
+ warmup_dtype_pairs = {
78
+ (model_dtype, model_dtype),
79
+ (torch.float32, torch.float32),
80
+ (torch.float32, model_dtype),
81
+ (model_dtype, torch.float32),
82
+ }
83
+ else:
84
+ warmup_dtype_pairs = {(model_dtype, None), (torch.float32, None)}
85
86
logger.info(
87
"Warming up spec-decode rejection sampler kernels "
66
- "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s).",
88
+ "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s, watermark=%s).",
89
vocab_size,
90
num_spec,
91
[(str(t), str(d)) for t, d in warmup_dtype_pairs],
92
use_block_verification,
93
+ watermark_config is not None,
94
)
95
for tgt_dtype, draft_dtype in warmup_dtype_pairs:
96
target_logits = torch.zeros(
97
(num_logits, vocab_size), dtype=tgt_dtype, device=device
98
)
76
- draft_logits = torch.zeros(
77
- (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device
99
+ draft_logits = (
100
+ torch.zeros(
101
+ (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device
102
+ )
103
+ if draft_dtype is not None
104
+ else None
105
)
106
synthetic_rates = (
107
torch.full((num_spec,), 0.5, dtype=torch.float32, device=device)
@@ -85,14 +112,14 @@ def spec_decode_rejection_warmup(worker: Worker) -> None:
112
rejection_sample(
113
target_logits=target_logits,
114
draft_logits=draft_logits,
88
- draft_sampled=torch.zeros(num_logits, dtype=torch.int64, device=device),
115
+ draft_sampled=torch.zeros(num_logits, dtype=torch.int32, device=device),
116
cu_num_logits=torch.tensor(
117
[0, num_logits], dtype=torch.int32, device=device
118
),
119
pos=torch.zeros(num_logits, dtype=torch.int64, device=device),
93
- idx_mapping=torch.zeros(num_reqs, dtype=torch.int32, device=device),
120
+ idx_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device),
121
expanded_idx_mapping=torch.zeros(
95
- num_logits, dtype=torch.int32, device=device
122
+ num_logits, dtype=torch.int64, device=device
123
),
124
expanded_local_pos=torch.arange(
125
num_logits, dtype=torch.int32, device=device
@@ -103,6 +130,7 @@ def spec_decode_rejection_warmup(worker: Worker) -> None:
130
synthetic_conditional_rates=synthetic_rates,
131
use_fp64=False,
132
use_block_verification=use_block_verification,
133
+ **watermark_kwargs,
134
)
135
except Exception:
136
logger.warning(
vllm/model_executor/warmup/watermark_sample_warmup.py
new
+155
@@ -0,0 +1,155 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Warm up the watermarked sampler's Triton kernel.
4
+
5
+``_philox_gumbel_kernel`` is JIT-compiled once per (key, logits dtype,
6
+skip-mask, ``USE_FP64``) specialization. The generic sampler warmup samples with
7
+``SamplingParams.for_sampler_warmup()``, whose feature-heavy logits processing
8
+forces the fp32 copy in ``apply_sampling_params``, so the first ordinary
9
+temperature-1.0 request is the first launch with model-dtype logits and pays
10
+the compilation inside inference. This pre-compiles every specialization the
11
+sampler path can launch for the configured watermark.
12
+"""
13
+
14
+from __future__ import annotations
15
+
16
+from typing import TYPE_CHECKING, cast
17
+
18
+import torch
19
+
20
+from vllm.logger import init_logger
21
+from vllm.platforms import current_platform
22
+
23
+if TYPE_CHECKING:
24
+ from vllm.v1.watermarking.watermarker import Watermarker
25
+ from vllm.v1.worker.gpu_worker import Worker
26
+
27
+logger = init_logger(__name__)
28
+
29
+_NUM_WARMUP_TOKENS = 1
30
+
31
+
32
+def _philox_key(watermarker: Watermarker) -> int | None:
33
+ """None when the watermarker takes the torch path and compiles no kernel."""
34
+ from vllm.v1.watermarking.gumbel import GumbelWatermarker
35
+ from vllm.v1.watermarking.prfs import PhiloxPRF
36
+
37
+ if not isinstance(watermarker, GumbelWatermarker):
38
+ return None
39
+ prf = watermarker.prf
40
+ return prf.key if type(prf) is PhiloxPRF else None
41
+
42
+
43
+def _philox_sampler_keys(watermarker: Watermarker) -> list[int]:
44
+ """Philox keys the sampler launches the kernel with."""
45
+ from vllm.v1.watermarking.gumbel import DualKeyGumbelWatermarker
46
+
47
+ watermarkers: list[Watermarker] = [watermarker]
48
+ if isinstance(watermarker, DualKeyGumbelWatermarker):
49
+ # ``DualKeyGumbelWatermarker.sample`` launches once per key, and
50
+ # collapses to a single key at the alpha bounds, where the other key's
51
+ # specialization is never launched.
52
+ if watermarker.alpha == 0:
53
+ watermarkers = [watermarker.draft_watermarker]
54
+ elif watermarker.alpha == 1:
55
+ watermarkers = [watermarker.target_watermarker]
56
+ else:
57
+ watermarkers = [
58
+ watermarker.draft_watermarker,
59
+ watermarker.target_watermarker,
60
+ ]
61
+ return [key for key in map(_philox_key, watermarkers) if key is not None]
62
+
63
+
64
+@torch.inference_mode()
65
+def watermark_sample_warmup(worker: Worker) -> None:
66
+ if getattr(worker.vllm_config, "watermark_config", None) is None:
67
+ return
68
+ # ``GumbelWatermarker`` launches the kernel only on CUDA
69
+ # (vllm/v1/watermarking/gumbel.py:63,:82); elsewhere the sampler takes the
70
+ # torch path and warming here would compile a kernel nothing launches.
71
+ if not current_platform.is_cuda_alike():
72
+ return
73
+
74
+ from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
75
+ from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
76
+
77
+ # Read the sampler the runner built instead of rebuilding it from config:
78
+ # its watermarker (the target role under a speculative config), its context
79
+ # deduplication setting and its fp64 flag are what the launches have to
80
+ # match, and cannot drift from the runner. A rank that does not sample -- a
81
+ # non-final pipeline stage, a pooling model -- holds no such sampler.
82
+ sampler = getattr(worker.model_runner, "sampler", None)
83
+ if not isinstance(sampler, GPUWatermarkSampler):
84
+ return
85
+
86
+ model_config = worker.vllm_config.model_config
87
+ device = worker.device
88
+ try:
89
+ keys = _philox_sampler_keys(sampler.watermarker)
90
+ if not keys:
91
+ return
92
+
93
+ # The sampler passes model-dtype logits through when no logits
94
+ # processing is needed and fp32 otherwise. ``ModelConfig.dtype`` is
95
+ # resolved to a torch dtype during validation.
96
+ dtypes = {cast("torch.dtype", model_config.dtype), torch.float32}
97
+ # ``GPUWatermarkSampler._sample_random`` builds a skip mask whenever
98
+ # context deduplication is on; with it off only a batch that mixes in an
99
+ # opted-out or greedy row does, and the mask-free specialization stays
100
+ # reachable.
101
+ with_skip_mask = (
102
+ (True,) if sampler.deduplicate_contexts != "none" else (True, False)
103
+ )
104
+
105
+ vocab_size = model_config.get_vocab_size()
106
+ contexts = torch.zeros(
107
+ (_NUM_WARMUP_TOKENS, sampler.watermarker.context_width),
108
+ dtype=torch.int32,
109
+ device=device,
110
+ )
111
+ # Dtypes mirror the runtime buffers the sampler passes and are part of
112
+ # the specialization key: int32 request-state token ids, int64 index and
113
+ # seed buffers, fp32 sampling temperatures.
114
+ sampling_state = {
115
+ "skip_mask": torch.zeros(
116
+ _NUM_WARMUP_TOKENS, dtype=torch.bool, device=device
117
+ ),
118
+ "expanded_idx_mapping": torch.zeros(
119
+ _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device
120
+ ),
121
+ "temperatures": torch.ones(
122
+ _NUM_WARMUP_TOKENS, dtype=torch.float32, device=device
123
+ ),
124
+ "seeds": torch.zeros(_NUM_WARMUP_TOKENS, dtype=torch.int64, device=device),
125
+ "positions": torch.zeros(
126
+ _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device
127
+ ),
128
+ }
129
+
130
+ logger.info(
131
+ "Warming up watermark sampler kernel (vocab=%d, keys=%d, dtypes=%s, "
132
+ "skip_mask=%s).",
133
+ vocab_size,
134
+ len(keys),
135
+ [str(dtype) for dtype in dtypes],
136
+ list(with_skip_mask),
137
+ )
138
+ for dtype in dtypes:
139
+ logits = torch.zeros(
140
+ (_NUM_WARMUP_TOKENS, vocab_size), dtype=dtype, device=device
141
+ )
142
+ for key in keys:
143
+ for use_skip_mask in with_skip_mask:
144
+ philox_gumbel_sample(
145
+ logits,
146
+ contexts,
147
+ key,
148
+ # Only the mixed path forwards the sampler's setting;
149
+ # ``GumbelWatermarker._sample_watermarked`` launches
150
+ # the mask-free variant with the fp32 default.
151
+ use_fp64=use_skip_mask and sampler.use_fp64_gumbel,
152
+ **(sampling_state if use_skip_mask else {}),
153
+ )
154
+ except Exception:
155
+ logger.warning("Skipping watermark sampler warmup.", exc_info=True)
vllm/v1/watermarking/__init__.py
+12
-1
@@ -1,9 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from vllm.config.watermarking import derive_watermark_key
5
from vllm.v1.watermarking.detector import WatermarkDetection, WatermarkDetector
6
from vllm.v1.watermarking.factory import create_watermarker
7
from vllm.v1.watermarking.gumbel import (
8
+ DualKeyGumbelWatermarkDetector,
9
+ DualKeyGumbelWatermarker,
10
GumbelWatermarkDetector,
11
GumbelWatermarker,
12
)
@@ -12,12 +15,19 @@ from vllm.v1.watermarking.prfs import (
15
WatermarkPRF,
16
create_prf,
17
)
15
-from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample
18
+from vllm.v1.watermarking.watermarker import (
19
+ SupportsSpeculativeDecoding,
20
+ Watermarker,
21
+ WatermarkSample,
22
+)
23
24
__all__ = [
25
+ "DualKeyGumbelWatermarkDetector",
26
+ "DualKeyGumbelWatermarker",
27
"GumbelWatermarkDetector",
28
"GumbelWatermarker",
29
"PhiloxPRF",
30
+ "SupportsSpeculativeDecoding",
31
"WatermarkDetection",
32
"WatermarkDetector",
33
"WatermarkPRF",
@@ -25,4 +35,5 @@ __all__ = [
35
"WatermarkSample",
36
"create_watermarker",
37
"create_prf",
38
+ "derive_watermark_key",
39
]
vllm/v1/watermarking/factory.py
+14
-16
@@ -1,24 +1,22 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
5
-
4
from vllm.config.watermarking import WatermarkConfig
7
-from vllm.v1.watermarking.gumbel import GumbelWatermarker
5
+from vllm.v1.watermarking.gumbel import (
6
+ DualKeyGumbelWatermarker,
7
+ GumbelWatermarker,
8
+)
9
from vllm.v1.watermarking.watermarker import Watermarker
10
11
11
-def _create_gumbel(config: WatermarkConfig) -> Watermarker:
12
- return GumbelWatermarker(config.key, config.context_width, config.prf)
13
-
14
-
15
-_WATERMARKERS: dict[str, Callable[[WatermarkConfig], Watermarker]] = {
16
- "gumbel": _create_gumbel,
17
-}
18
-
19
-
12
def create_watermarker(config: WatermarkConfig) -> Watermarker:
21
- watermarker_factory = _WATERMARKERS.get(config.algorithm)
22
- if watermarker_factory is None:
23
- raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}")
24
- return watermarker_factory(config)
13
+ if config.algorithm == "gumbel":
14
+ return GumbelWatermarker(config.key, config.context_width, config.prf)
15
+ if config.algorithm == "dual_key_gumbel":
16
+ return DualKeyGumbelWatermarker(
17
+ config.key,
18
+ config.context_width,
19
+ config.prf,
20
+ config.alpha,
21
+ )
22
+ raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}")
vllm/v1/watermarking/gpu_sampler.py
+97
-34
@@ -1,16 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from typing import Literal
5
+
6
import numpy as np
7
import torch
8
9
from vllm.logger import init_logger
10
from vllm.sampling_params import SamplingParams
11
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
10
-from vllm.v1.watermarking.watermarker import Watermarker
12
+from vllm.v1.watermarking.watermarker import RandomSampler, Watermarker
13
from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
12
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
14
from vllm.v1.worker.gpu.sample.sampler import Sampler
15
+from vllm.v1.worker.gpu.sample.watermark import repeated_context_mask
16
17
logger = init_logger(__name__)
18
@@ -20,10 +22,14 @@ class GPUWatermarkSampler(Sampler):
22
self,
23
watermarker: Watermarker,
24
*args,
25
+ deduplicate_contexts: Literal["none", "single_turn", "all"] = "single_turn",
26
+ deduplicate_contexts_max_history: int | None = 8192,
27
**kwargs,
28
) -> None:
29
super().__init__(*args, **kwargs)
30
self.watermarker = watermarker
31
+ self.deduplicate_contexts = deduplicate_contexts
32
+ self.deduplicate_contexts_max_history = deduplicate_contexts_max_history
33
self.watermarking = UvaBackedTensor(
34
self.sampling_states.max_num_reqs, dtype=torch.bool
35
)
@@ -72,37 +78,38 @@ class GPUWatermarkSampler(Sampler):
78
79
processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
80
contexts = self._get_contexts(expanded_idx_mapping)
81
+ repeated_contexts = None
82
+ if self.deduplicate_contexts != "none":
83
+ repeated_contexts = self._get_repeated_contexts(
84
+ expanded_idx_mapping, contexts
85
+ )
86
76
- def random_sample(sample_logits: torch.Tensor) -> torch.Tensor:
77
- return gumbel_sample(
78
- sample_logits,
79
- expanded_idx_mapping,
80
- self.sampling_states.temperature.gpu,
81
- self.sampling_states.seeds.gpu,
82
- pos,
83
- apply_temperature=False,
84
- is_drafting=False,
85
- use_fp64=self.use_fp64_gumbel,
87
+ temperatures = self.sampling_states.temperature.gpu[expanded_idx_mapping]
88
+ needs_mixed_sampling = repeated_contexts is not None or not np.all(enabled)
89
+ skip_mask = None
90
+ if needs_mixed_sampling:
91
+ watermarking = self.watermarking.gpu[expanded_idx_mapping] & (
92
+ temperatures != 0
93
)
94
+ if repeated_contexts is not None:
95
+ watermarking &= ~repeated_contexts
96
+ skip_mask = ~watermarking
97
98
+ random_sampler = RandomSampler(
99
+ expanded_idx_mapping=expanded_idx_mapping,
100
+ temperatures=self.sampling_states.temperature.gpu,
101
+ seeds=self.sampling_states.seeds.gpu,
102
+ positions=pos,
103
+ use_fp64=self.use_fp64_gumbel,
104
+ )
105
output = self.watermarker.sample(
106
processed_logits,
107
contexts,
91
- random_sample,
108
+ random_sampler,
109
+ skip_mask=skip_mask,
110
)
93
- temperatures = self.sampling_states.temperature.gpu[expanded_idx_mapping]
94
- watermarking = self.watermarking.gpu[expanded_idx_mapping] & (temperatures != 0)
95
- if not np.all(enabled):
96
- unwatermarked = random_sample(processed_logits)
97
- sampled = torch.where(watermarking, output.token_ids, unwatermarked)
98
- output_logits = output.logits
99
- if output.logits is not processed_logits:
100
- output_logits = torch.where(
101
- watermarking.unsqueeze(-1), output.logits, processed_logits
102
- )
103
- else:
104
- sampled = output.token_ids
105
- output_logits = output.logits
111
+ sampled = output.token_ids
112
+ output_logits = output.logits
113
sampled = torch.where(
114
temperatures == 0,
115
processed_logits.argmax(dim=-1),
@@ -110,22 +117,78 @@ class GPUWatermarkSampler(Sampler):
117
)
118
return sampled, output_logits
119
113
- def _get_contexts(self, expanded_idx_mapping: torch.Tensor) -> torch.Tensor:
120
+ def _get_repeated_contexts(
121
+ self,
122
+ expanded_idx_mapping: torch.Tensor,
123
+ contexts: torch.Tensor,
124
+ ) -> torch.Tensor:
125
+ return repeated_context_mask(
126
+ self.req_states.all_token_ids.gpu,
127
+ expanded_idx_mapping,
128
+ self.req_states.prompt_len.gpu,
129
+ self.req_states.total_len.gpu,
130
+ contexts,
131
+ self.deduplicate_contexts_max_history,
132
+ include_prompt=self.deduplicate_contexts == "all",
133
+ skip_partial_context=self.deduplicate_contexts == "all",
134
+ )
135
+
136
+ def _get_contexts(
137
+ self,
138
+ expanded_idx_mapping: torch.Tensor,
139
+ expanded_local_pos: torch.Tensor | None = None,
140
+ draft_sampled: torch.Tensor | None = None,
141
+ ) -> torch.Tensor:
142
+ """Build contexts for flattened verification rows.
143
+
144
+ With committed tokens ``[10, 11]``, drafts ``[20, 21]``, and width 2,
145
+ local positions ``[0, 1, 2]`` use contexts ``[10, 11]``, ``[11, 20]``,
146
+ and ``[20, 21]``. The first verification row precedes the first draft.
147
+ """
148
context_width = self.watermarker.context_width
149
req_indices = expanded_idx_mapping.to(torch.int64)
150
valid_reqs = req_indices >= 0
151
safe_req_indices = req_indices.clamp_min(0)
152
total_lens = self.req_states.total_len.gpu[safe_req_indices].to(torch.int64)
153
prompt_lens = self.req_states.prompt_len.gpu[safe_req_indices].to(torch.int64)
154
+ if expanded_local_pos is None:
155
+ expanded_local_pos = torch.zeros_like(req_indices)
156
+ else:
157
+ expanded_local_pos = expanded_local_pos.to(torch.int64)
158
offsets = torch.arange(
159
-context_width, 0, dtype=torch.int64, device=req_indices.device
160
)
123
- positions = total_lens.unsqueeze(-1) + offsets
124
- valid_positions = valid_reqs.unsqueeze(-1) & (
125
- positions >= prompt_lens.unsqueeze(-1)
161
+ positions = (
162
+ total_lens.unsqueeze(-1) + expanded_local_pos.unsqueeze(-1) + offsets
163
)
127
- positions = positions.clamp_min(0)
128
- contexts = self.req_states.all_token_ids.gpu[
129
- safe_req_indices.unsqueeze(-1), positions
164
+ committed = positions < total_lens.unsqueeze(-1)
165
+ valid_committed = (
166
+ valid_reqs.unsqueeze(-1)
167
+ & committed
168
+ & (positions >= prompt_lens.unsqueeze(-1))
169
+ )
170
+ committed_contexts = self.req_states.all_token_ids.gpu[
171
+ safe_req_indices.unsqueeze(-1),
172
+ positions.clamp(0, self.req_states.all_token_ids.gpu.shape[1] - 1),
173
]
131
- return torch.where(valid_positions, contexts, -1)
174
+ contexts = torch.where(valid_committed, committed_contexts, -1)
175
+ if draft_sampled is None:
176
+ return contexts
177
+
178
+ row_indices = torch.arange(
179
+ len(req_indices), dtype=torch.int64, device=req_indices.device
180
+ )
181
+ draft_offsets = positions - total_lens.unsqueeze(-1)
182
+ draft_indices = (
183
+ row_indices.unsqueeze(-1)
184
+ - expanded_local_pos.unsqueeze(-1)
185
+ + draft_offsets
186
+ + 1
187
+ )
188
+ valid_drafts = (
189
+ valid_reqs.unsqueeze(-1)
190
+ & ~committed
191
+ & (draft_offsets < expanded_local_pos.unsqueeze(-1))
192
+ )
193
+ draft_contexts = draft_sampled[draft_indices.clamp(0, len(draft_sampled) - 1)]
194
+ return torch.where(valid_drafts, draft_contexts, contexts)
vllm/v1/watermarking/gumbel.py
+162
-3
@@ -8,13 +8,14 @@ import warnings
8
9
import torch
10
11
-from vllm.config.watermarking import WatermarkPRFName
11
+from vllm.config.watermarking import WatermarkPRFName, derive_watermark_key
12
from vllm.v1.watermarking.detector import (
13
WatermarkDetector,
14
)
15
from vllm.v1.watermarking.prfs import PhiloxPRF, WatermarkPRF, create_prf
16
from vllm.v1.watermarking.watermarker import (
17
RandomSampler,
18
+ SupportsSpeculativeDecoding,
19
Watermarker,
20
WatermarkSample,
21
)
@@ -54,11 +55,10 @@ class GumbelWatermarker(Watermarker):
55
def context_width(self) -> int:
56
return self._context_width
57
57
- def sample(
58
+ def _sample_watermarked(
59
self,
60
logits: torch.Tensor,
61
contexts: torch.Tensor,
61
- _random_sample: RandomSampler,
62
) -> WatermarkSample:
63
if type(self.prf) is PhiloxPRF and logits.device.type == "cuda":
64
from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
@@ -72,6 +72,110 @@ class GumbelWatermarker(Watermarker):
72
noise = -torch.log(-torch.log(uniforms))
73
return WatermarkSample(torch.argmax(logits + noise, dim=-1), logits)
74
75
+ def _try_sample_mixed(
76
+ self,
77
+ logits: torch.Tensor,
78
+ contexts: torch.Tensor,
79
+ skip_mask: torch.Tensor,
80
+ random_sampler: RandomSampler,
81
+ ) -> WatermarkSample | None:
82
+ if type(self.prf) is not PhiloxPRF or logits.device.type != "cuda":
83
+ return None
84
+
85
+ from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
86
+
87
+ token_ids = philox_gumbel_sample(
88
+ logits,
89
+ contexts,
90
+ self.prf.key,
91
+ skip_mask=skip_mask,
92
+ expanded_idx_mapping=random_sampler.expanded_idx_mapping,
93
+ temperatures=random_sampler.temperatures,
94
+ seeds=random_sampler.seeds,
95
+ positions=random_sampler.positions,
96
+ use_fp64=random_sampler.use_fp64,
97
+ )
98
+ return WatermarkSample(token_ids, logits)
99
+
100
+
101
+class DualKeyGumbelWatermarker(Watermarker, SupportsSpeculativeDecoding):
102
+ def __init__(
103
+ self,
104
+ key: int,
105
+ context_width: int = 4,
106
+ prf: WatermarkPRFName = "philox",
107
+ alpha: float = 0.1,
108
+ ) -> None:
109
+ if not 0 <= alpha <= 1:
110
+ raise ValueError("alpha must be between 0 and 1")
111
+ self.alpha = alpha
112
+ self.draft_watermarker = GumbelWatermarker(
113
+ derive_watermark_key(key, b"key_a"), context_width, prf
114
+ )
115
+ self.target_watermarker = GumbelWatermarker(
116
+ derive_watermark_key(key, b"key_b"), context_width, prf
117
+ )
118
+ self._routing_logits_cache: dict[torch.device, torch.Tensor] = {}
119
+
120
+ def _routing_logits(self, device: torch.device) -> torch.Tensor:
121
+ # Building this tensor on every CUDA call stages a pageable H2D copy and
122
+ # synchronizes the stream.
123
+ cached = self._routing_logits_cache.get(device)
124
+ if cached is None:
125
+ cached = torch.tensor(
126
+ [1 - self.alpha, self.alpha],
127
+ dtype=torch.float32,
128
+ device=device,
129
+ ).log()
130
+ self._routing_logits_cache[device] = cached
131
+ return cached
132
+
133
+ @property
134
+ def context_width(self) -> int:
135
+ return self.draft_watermarker.context_width
136
+
137
+ def sample(
138
+ self,
139
+ logits: torch.Tensor,
140
+ contexts: torch.Tensor,
141
+ random_sampler: RandomSampler | None = None,
142
+ skip_mask: torch.Tensor | None = None,
143
+ ) -> WatermarkSample:
144
+ if self.alpha == 0:
145
+ return self.draft_watermarker.sample(
146
+ logits, contexts, random_sampler, skip_mask
147
+ )
148
+ if self.alpha == 1:
149
+ return self.target_watermarker.sample(
150
+ logits, contexts, random_sampler, skip_mask
151
+ )
152
+ if random_sampler is None:
153
+ raise ValueError("dual-key Gumbel routing requires a random sampler")
154
+
155
+ key_a_sample = self.draft_watermarker.sample(
156
+ logits, contexts, random_sampler, skip_mask
157
+ )
158
+ key_b_sample = self.target_watermarker.sample(
159
+ logits, contexts, random_sampler, skip_mask
160
+ )
161
+ routing_logits = self._routing_logits(logits.device).expand(logits.shape[0], -1)
162
+ use_key_a = random_sampler(routing_logits) == 0
163
+ return WatermarkSample(
164
+ torch.where(
165
+ use_key_a,
166
+ key_a_sample.token_ids,
167
+ key_b_sample.token_ids,
168
+ ),
169
+ logits,
170
+ )
171
+
172
+ def _sample_watermarked(
173
+ self,
174
+ logits: torch.Tensor,
175
+ contexts: torch.Tensor,
176
+ ) -> WatermarkSample:
177
+ raise ValueError("dual-key Gumbel routing requires a random sampler")
178
+
179
180
class GumbelWatermarkDetector(WatermarkDetector):
181
def __init__(
@@ -100,6 +204,61 @@ class GumbelWatermarkDetector(WatermarkDetector):
204
return token_scores.sum().item()
205
206
207
+class DualKeyGumbelWatermarkDetector(GumbelWatermarkDetector):
208
+ """Detect a dual-key watermark with weighted early fusion.
209
+
210
+ ``alpha`` is the weight assigned to key B; key A receives ``1 - alpha``.
211
+ """
212
+
213
+ def __init__(
214
+ self,
215
+ key: int,
216
+ context_width: int = 4,
217
+ p_value_threshold: float = 0.01,
218
+ prf: WatermarkPRFName = "philox",
219
+ deduplicate_contexts: bool = True,
220
+ alpha: float = 0.2,
221
+ ) -> None:
222
+ if not 0 <= alpha <= 1:
223
+ raise ValueError("alpha must be between 0 and 1")
224
+ super().__init__(
225
+ derive_watermark_key(key, b"key_a"),
226
+ context_width,
227
+ p_value_threshold,
228
+ prf,
229
+ deduplicate_contexts,
230
+ )
231
+ self.key_b_prf = create_prf(prf, derive_watermark_key(key, b"key_b"))
232
+ self.alpha = alpha
233
+
234
+ def _score_tokens(
235
+ self, contexts: torch.Tensor, targets: torch.Tensor
236
+ ) -> torch.Tensor:
237
+ scores = torch.stack(
238
+ [
239
+ -torch.log1p(
240
+ -prf.uniform(contexts, targets.unsqueeze(-1))
241
+ .squeeze(-1)
242
+ .to(torch.float64)
243
+ )
244
+ for prf in (self.prf, self.key_b_prf)
245
+ ],
246
+ dim=-1,
247
+ )
248
+ weights = scores.new_tensor([1 - self.alpha, self.alpha])
249
+ return scores @ weights
250
+
251
+ def _get_p_value(self, score: float, num_scored_tokens: int) -> float:
252
+ variance = (1 - self.alpha) ** 2 + self.alpha**2
253
+ if self.alpha in (0, 0.5, 1):
254
+ return _gamma_survival_integer_shape(
255
+ score / variance, round(num_scored_tokens / variance)
256
+ )
257
+ shape = torch.tensor(num_scored_tokens / variance, dtype=torch.float64)
258
+ scaled_score = torch.tensor(score / variance, dtype=torch.float64)
259
+ return torch.special.gammaincc(shape, scaled_score).item()
260
+
261
+
262
def _validate_context_width(context_width: int) -> None:
263
if context_width < 1:
264
raise ValueError("context_width must be positive")
vllm/v1/watermarking/spec_decode.py
new
+153
@@ -0,0 +1,153 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.config.watermarking import WatermarkConfig
7
+from vllm.v1.watermarking.factory import create_watermarker
8
+from vllm.v1.watermarking.gumbel import GumbelWatermarker
9
+from vllm.v1.watermarking.prfs import PhiloxPRF
10
+from vllm.v1.watermarking.watermarker import (
11
+ SupportsSpeculativeDecoding,
12
+ Watermarker,
13
+)
14
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample
15
+
16
+
17
+class DraftWatermarker:
18
+ def __init__(
19
+ self,
20
+ watermarker: Watermarker,
21
+ max_num_reqs: int,
22
+ device: torch.device,
23
+ ) -> None:
24
+ self.watermarker = watermarker
25
+ self.contexts = torch.zeros(
26
+ max_num_reqs,
27
+ watermarker.context_width,
28
+ dtype=torch.int64,
29
+ device=device,
30
+ )
31
+ self.enabled = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
32
+
33
+ def prepare(self, contexts: torch.Tensor, enabled: torch.Tensor) -> None:
34
+ num_reqs = contexts.shape[0]
35
+ self.contexts[:num_reqs].copy_(contexts)
36
+ self.enabled[:num_reqs].copy_(enabled)
37
+
38
+ def sample(
39
+ self,
40
+ logits: torch.Tensor,
41
+ ordinary_sampled: torch.Tensor,
42
+ idx_mapping: torch.Tensor,
43
+ temperature: torch.Tensor,
44
+ ) -> torch.Tensor:
45
+ request_temperatures = temperature[idx_mapping]
46
+ processed_logits = logits / torch.where(
47
+ request_temperatures == 0, 1, request_temperatures
48
+ ).unsqueeze(-1)
49
+ watermarked = self.watermarker.sample(
50
+ processed_logits,
51
+ self.contexts[: logits.shape[0]],
52
+ ).token_ids
53
+ enabled = self.enabled[: logits.shape[0]] & (request_temperatures != 0)
54
+ sampled = torch.where(enabled, watermarked, ordinary_sampled)
55
+ contexts = self.contexts[: logits.shape[0]]
56
+ contexts.copy_(torch.cat((contexts[:, 1:], sampled.unsqueeze(-1)), dim=-1))
57
+ return sampled
58
+
59
+
60
+def create_speculative_target_watermarker(watermarker: Watermarker) -> Watermarker:
61
+ if isinstance(watermarker, SupportsSpeculativeDecoding):
62
+ return watermarker.target_watermarker
63
+ return watermarker
64
+
65
+
66
+def create_speculative_draft_watermarker(
67
+ watermarker: Watermarker,
68
+ max_num_reqs: int,
69
+ device: torch.device,
70
+ allow_target_only: bool,
71
+) -> DraftWatermarker | None:
72
+ if isinstance(watermarker, SupportsSpeculativeDecoding):
73
+ return DraftWatermarker(watermarker.draft_watermarker, max_num_reqs, device)
74
+ if allow_target_only:
75
+ return None
76
+ raise ValueError(
77
+ f"The {type(watermarker).__name__} watermarking algorithm does not support "
78
+ "speculative decoding. Set allow_target_only_watermarking=true to "
79
+ "leave draft tokens unwatermarked."
80
+ )
81
+
82
+
83
+def _philox_key(watermarker: Watermarker) -> int:
84
+ if not (
85
+ isinstance(watermarker, GumbelWatermarker)
86
+ and type(watermarker.prf) is PhiloxPRF
87
+ ):
88
+ raise NotImplementedError(
89
+ "In-kernel watermarked recovery supports only the Philox PRF "
90
+ "GumbelWatermarker"
91
+ )
92
+ return watermarker.prf.key
93
+
94
+
95
+def _resolve_watermark_key(watermarker: Watermarker) -> int:
96
+ """Return the Philox recovery key."""
97
+ if isinstance(watermarker, SupportsSpeculativeDecoding):
98
+ raise ValueError(
99
+ f"{type(watermarker).__name__} keys the target role separately from "
100
+ "the draft. Pass create_speculative_target_watermarker(watermarker) "
101
+ "so the recovery draw carries the target's key."
102
+ )
103
+ return _philox_key(watermarker)
104
+
105
+
106
+def speculative_target_watermark_key(
107
+ watermark_config: WatermarkConfig | None,
108
+) -> int | None:
109
+ """Resolve the resample kernel's Philox key from configuration."""
110
+ if watermark_config is None:
111
+ return None
112
+
113
+ return _resolve_watermark_key(
114
+ create_speculative_target_watermarker(create_watermarker(watermark_config))
115
+ )
116
+
117
+
118
+def watermarked_rejection_sample(
119
+ target_logits: torch.Tensor,
120
+ draft_logits: torch.Tensor | None,
121
+ draft_sampled: torch.Tensor,
122
+ cu_num_logits: torch.Tensor,
123
+ pos: torch.Tensor,
124
+ idx_mapping: torch.Tensor,
125
+ expanded_idx_mapping: torch.Tensor,
126
+ expanded_local_pos: torch.Tensor,
127
+ temperature: torch.Tensor,
128
+ seed: torch.Tensor,
129
+ num_speculative_steps: int,
130
+ contexts: torch.Tensor,
131
+ watermarking: torch.Tensor,
132
+ watermarker: Watermarker,
133
+ use_fp64: bool = False,
134
+) -> tuple[torch.Tensor, torch.Tensor]:
135
+ """Use the target key for rejection recovery and bonus tokens."""
136
+ assert contexts.shape == (target_logits.shape[0], watermarker.context_width)
137
+ return rejection_sample(
138
+ target_logits,
139
+ draft_logits,
140
+ draft_sampled,
141
+ cu_num_logits,
142
+ pos,
143
+ idx_mapping,
144
+ expanded_idx_mapping,
145
+ expanded_local_pos,
146
+ temperature,
147
+ seed,
148
+ num_speculative_steps,
149
+ use_fp64=use_fp64,
150
+ contexts=contexts,
151
+ watermarking=watermarking,
152
+ watermark_key=_resolve_watermark_key(watermarker),
153
+ )
vllm/v1/watermarking/watermarker.py
+68
-5
@@ -2,13 +2,12 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from abc import ABC, abstractmethod
5
-from collections.abc import Callable
5
from dataclasses import dataclass
7
-from typing import TypeAlias
6
+from typing import Protocol, runtime_checkable
7
8
import torch
9
11
-RandomSampler: TypeAlias = Callable[[torch.Tensor], torch.Tensor]
10
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
11
12
13
@dataclass(frozen=True)
@@ -17,17 +16,81 @@ class WatermarkSample:
16
logits: torch.Tensor
17
18
19
+@dataclass(frozen=True)
20
+class RandomSampler:
21
+ expanded_idx_mapping: torch.Tensor
22
+ temperatures: torch.Tensor
23
+ seeds: torch.Tensor
24
+ positions: torch.Tensor
25
+ use_fp64: bool = False
26
+
27
+ def __call__(self, logits: torch.Tensor) -> torch.Tensor:
28
+ return gumbel_sample(
29
+ logits,
30
+ self.expanded_idx_mapping,
31
+ self.temperatures,
32
+ self.seeds,
33
+ self.positions,
34
+ apply_temperature=False,
35
+ is_drafting=False,
36
+ use_fp64=self.use_fp64,
37
+ )
38
+
39
+
40
class Watermarker(ABC):
41
@property
42
@abstractmethod
43
def context_width(self) -> int:
44
raise NotImplementedError
45
26
- @abstractmethod
46
def sample(
47
self,
48
logits: torch.Tensor,
49
contexts: torch.Tensor,
31
- random_sample: RandomSampler,
50
+ random_sampler: RandomSampler | None = None,
51
+ skip_mask: torch.Tensor | None = None,
52
+ ) -> WatermarkSample:
53
+ if skip_mask is not None:
54
+ if random_sampler is None:
55
+ raise ValueError("skip_mask requires a random sampler")
56
+ mixed = self._try_sample_mixed(logits, contexts, skip_mask, random_sampler)
57
+ if mixed is not None:
58
+ return mixed
59
+
60
+ watermarked = self._sample_watermarked(logits, contexts)
61
+ if skip_mask is None:
62
+ return watermarked
63
+
64
+ assert random_sampler is not None
65
+ token_ids = torch.where(
66
+ skip_mask, random_sampler(logits), watermarked.token_ids
67
+ )
68
+ output_logits = watermarked.logits
69
+ if output_logits is not logits:
70
+ output_logits = torch.where(skip_mask.unsqueeze(-1), logits, output_logits)
71
+ return WatermarkSample(token_ids, output_logits)
72
+
73
+ @abstractmethod
74
+ def _sample_watermarked(
75
+ self,
76
+ logits: torch.Tensor,
77
+ contexts: torch.Tensor,
78
) -> WatermarkSample:
79
raise NotImplementedError
80
+
81
+ def _try_sample_mixed(
82
+ self,
83
+ logits: torch.Tensor,
84
+ contexts: torch.Tensor,
85
+ skip_mask: torch.Tensor,
86
+ random_sampler: RandomSampler,
87
+ ) -> WatermarkSample | None:
88
+ return None
89
+
90
+
91
+@runtime_checkable
92
+class SupportsSpeculativeDecoding(Protocol):
93
+ """A watermarker compatible with standard speculative rejection sampling."""
94
+
95
+ draft_watermarker: Watermarker
96
+ target_watermarker: Watermarker
vllm/v1/worker/gpu/model_runner.py
+29
-2
@@ -79,6 +79,10 @@ from vllm.v1.outputs import (
79
)
80
from vllm.v1.watermarking import create_watermarker
81
from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
82
+from vllm.v1.watermarking.spec_decode import (
83
+ create_speculative_target_watermarker,
84
+ speculative_target_watermark_key,
85
+)
86
from vllm.v1.worker.block_table import get_block_table_width
87
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
88
from vllm.v1.worker.gpu import pcp_manager as pcp
@@ -472,8 +476,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
476
if self.vllm_config.watermark_config is None:
477
self.sampler = Sampler(**sampler_kwargs)
478
else:
475
- watermarker = create_watermarker(self.vllm_config.watermark_config)
476
- self.sampler = GPUWatermarkSampler(watermarker, **sampler_kwargs)
479
+ wm_config = self.vllm_config.watermark_config
480
+ watermarker = create_watermarker(wm_config)
481
+ if self.speculative_config is not None:
482
+ watermarker = create_speculative_target_watermarker(watermarker)
483
+ self.sampler = GPUWatermarkSampler(
484
+ watermarker,
485
+ deduplicate_contexts=wm_config.deduplicate_contexts,
486
+ deduplicate_contexts_max_history=(
487
+ wm_config.deduplicate_contexts_max_history
488
+ ),
489
+ **sampler_kwargs,
490
+ )
491
custom = self.model_state.custom_sampler(self.sampler)
492
493
if custom:
@@ -484,6 +498,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
498
self.sampler,
499
self.speculative_config,
500
self.device,
501
+ watermark_key=speculative_target_watermark_key(
502
+ self.vllm_config.watermark_config
503
+ ),
504
)
505
self.prompt_logprobs_worker = PromptLogprobsWorker(
506
self.max_num_reqs,
@@ -860,6 +877,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
877
if hasattr(self.model, "get_mtp_target_hidden_states"):
878
pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
879
spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]
880
+ if isinstance(self.sampler, GPUWatermarkSampler):
881
+ self.speculator.prepare_watermarking(
882
+ self.sampler._get_contexts(input_batch.idx_mapping),
883
+ self.sampler.watermarking.gpu[input_batch.idx_mapping],
884
+ )
885
with use_workspace_lane(self._draft_workspace_lane):
886
self.speculator.propose(
887
input_batch=input_batch,
@@ -2086,6 +2108,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2108
if hasattr(self.model, "get_mtp_target_hidden_states"):
2109
pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
2110
spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]
2111
+ if isinstance(self.sampler, GPUWatermarkSampler):
2112
+ self.speculator.prepare_watermarking(
2113
+ self.sampler._get_contexts(input_batch.idx_mapping),
2114
+ self.sampler.watermarking.gpu[input_batch.idx_mapping],
2115
+ )
2116
with use_workspace_lane(self._draft_workspace_lane):
2117
draft_tokens = self.speculator.propose(
2118
input_batch,
vllm/v1/worker/gpu/sample/watermark.py
+324
-45
@@ -25,6 +25,7 @@ from vllm.v1.watermarking.prfs.philox import (
25
from vllm.v1.watermarking.prfs.philox import (
26
_UINT32_MASK as _UINT32_MASK_VALUE,
27
)
28
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax
29
30
if HAS_TRITON:
31
from triton.language import math as tl_math
@@ -42,6 +43,162 @@ _CONTEXT_DOMAIN = (
43
_TOKEN_DOMAIN = tl.constexpr(_TOKEN_DOMAIN_VALUE) if HAS_TRITON else _TOKEN_DOMAIN_VALUE
44
45
46
+@triton.jit
47
+def _repeated_context_mask_kernel(
48
+ output_ptr,
49
+ all_token_ids_ptr,
50
+ all_token_ids_stride,
51
+ req_indices_ptr,
52
+ prompt_lens_ptr,
53
+ total_lens_ptr,
54
+ contexts_ptr,
55
+ context_stride,
56
+ CONTEXT_WIDTH: tl.constexpr,
57
+ MAX_HISTORY: tl.constexpr,
58
+ INCLUDE_PROMPT: tl.constexpr,
59
+ SKIP_PARTIAL_CONTEXT: tl.constexpr,
60
+ BLOCK: tl.constexpr,
61
+):
62
+ row = tl.program_id(0).to(tl.int64)
63
+ req_idx = tl.load(req_indices_ptr + row).to(tl.int64)
64
+ if req_idx < 0:
65
+ tl.store(output_ptr + row, 0)
66
+ return
67
+ prompt_len = tl.load(prompt_lens_ptr + req_idx)
68
+ total_len = tl.load(total_lens_ptr + req_idx)
69
+ if SKIP_PARTIAL_CONTEXT and total_len - prompt_len < CONTEXT_WIDTH:
70
+ tl.store(output_ptr + row, 1)
71
+ return
72
+ sequence_start = 0 if INCLUDE_PROMPT else prompt_len
73
+ history_len = total_len - sequence_start
74
+
75
+ offsets = tl.arange(0, BLOCK)
76
+ repeated = tl.full((), 0, tl.int32)
77
+ scan_start = 0
78
+ if MAX_HISTORY > 0:
79
+ scan_start = tl.maximum(history_len - MAX_HISTORY, 0)
80
+ # Aligned loop bounds let the per-lane offsets fold into immediates.
81
+ aligned_start = (scan_start // BLOCK) * BLOCK
82
+ for block_start in tl.range(aligned_start, history_len, BLOCK):
83
+ previous_pos = block_start + offsets
84
+ in_window = (previous_pos >= scan_start) & (previous_pos < history_len)
85
+ matches = in_window
86
+ for offset in range(CONTEXT_WIDTH):
87
+ context_token = tl.load(contexts_ptr + row * context_stride + offset)
88
+ historical_pos = previous_pos + offset - CONTEXT_WIDTH
89
+ historical_token = tl.load(
90
+ all_token_ids_ptr
91
+ + req_idx * all_token_ids_stride
92
+ + sequence_start
93
+ + tl.maximum(historical_pos, 0),
94
+ mask=in_window & (historical_pos >= 0),
95
+ other=-1,
96
+ )
97
+ matches &= historical_token == context_token
98
+ repeated |= tl.max(matches.to(tl.int32), axis=0)
99
+ tl.store(output_ptr + row, repeated)
100
+
101
+
102
+def _repeated_context_mask_cpu(
103
+ all_token_ids: torch.Tensor,
104
+ req_indices: torch.Tensor,
105
+ prompt_lens: torch.Tensor,
106
+ total_lens: torch.Tensor,
107
+ contexts: torch.Tensor,
108
+ max_history: int | None = None,
109
+ include_prompt: bool = False,
110
+ skip_partial_context: bool = False,
111
+) -> torch.Tensor:
112
+ """Reference implementation; the parity tests check the Triton kernel against it."""
113
+ repeated = torch.zeros(len(req_indices), dtype=torch.bool)
114
+ for row, req_idx_tensor in enumerate(req_indices):
115
+ req_idx = int(req_idx_tensor)
116
+ if req_idx < 0:
117
+ continue
118
+ prompt_len = int(prompt_lens[req_idx])
119
+ total_len = int(total_lens[req_idx])
120
+ if skip_partial_context and total_len - prompt_len < contexts.shape[-1]:
121
+ repeated[row] = True
122
+ continue
123
+ sequence_start = 0 if include_prompt else prompt_len
124
+ history_tokens = all_token_ids[req_idx, sequence_start:total_len].tolist()
125
+ prefix = [-1] * contexts.shape[-1]
126
+ current_context = tuple(contexts[row].tolist())
127
+ history_start = (
128
+ max(0, len(history_tokens) - max_history) if max_history is not None else 0
129
+ )
130
+ for history_pos, token_id in enumerate(history_tokens):
131
+ if history_pos >= history_start and (
132
+ tuple(prefix[-contexts.shape[-1] :]) == current_context
133
+ ):
134
+ repeated[row] = True
135
+ break
136
+ prefix.append(token_id)
137
+ return repeated
138
+
139
+
140
+def repeated_context_mask(
141
+ all_token_ids: torch.Tensor,
142
+ req_indices: torch.Tensor,
143
+ prompt_lens: torch.Tensor,
144
+ total_lens: torch.Tensor,
145
+ contexts: torch.Tensor,
146
+ max_history: int | None = None,
147
+ include_prompt: bool = False,
148
+ skip_partial_context: bool = False,
149
+) -> torch.Tensor:
150
+ """Return, per row, whether the row's context already occurred in its history.
151
+
152
+ Args:
153
+ all_token_ids: `[max_num_reqs, max_model_len]` token ids of every request.
154
+ req_indices: Request slot per sampled row; -1 marks a padding row, which
155
+ is reported as not repeated.
156
+ prompt_lens: Prompt length per request slot.
157
+ total_lens: Prompt plus generated length per request slot.
158
+ contexts: `[num_rows, context_width]` context per row, padded with -1
159
+ before the start of the scanned history.
160
+ max_history: Number of most recent history positions searched, or
161
+ ``None`` for all of them. The window compared at each position
162
+ reaches `context_width` tokens further back.
163
+ include_prompt: Search the prompt as well as the generated tokens.
164
+ skip_partial_context: Mark contexts containing start padding so they use
165
+ ordinary sampling.
166
+ """
167
+ if max_history is not None and max_history < 1:
168
+ raise ValueError("max_history must be positive or None")
169
+ if all_token_ids.device.type == "cpu":
170
+ return _repeated_context_mask_cpu(
171
+ all_token_ids,
172
+ req_indices,
173
+ prompt_lens,
174
+ total_lens,
175
+ contexts,
176
+ max_history,
177
+ include_prompt,
178
+ skip_partial_context,
179
+ )
180
+
181
+ if contexts.stride(-1) != 1:
182
+ contexts = contexts.contiguous()
183
+ repeated = torch.empty(len(req_indices), dtype=torch.bool, device=contexts.device)
184
+ _repeated_context_mask_kernel[(len(req_indices),)](
185
+ repeated,
186
+ all_token_ids,
187
+ all_token_ids.stride(0),
188
+ req_indices,
189
+ prompt_lens,
190
+ total_lens,
191
+ contexts,
192
+ contexts.stride(0),
193
+ CONTEXT_WIDTH=contexts.shape[-1],
194
+ MAX_HISTORY=0 if max_history is None else max_history,
195
+ INCLUDE_PROMPT=include_prompt,
196
+ SKIP_PARTIAL_CONTEXT=skip_partial_context,
197
+ BLOCK=512,
198
+ )
199
+ return repeated
200
+
201
+
202
@triton.jit
203
def _mulhilo32(multiplier: tl.constexpr, value):
204
high = tl_math.umulhi(multiplier, value)
@@ -74,13 +231,101 @@ def _uint32_to_uniform(value):
231
232
233
@triton.jit
77
-def _gumbel_value(logits_ptr, output, mask):
78
- logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32)
234
+def _philox_context_state(
235
+ contexts_row_ptr,
236
+ key_0,
237
+ key_1,
238
+ CONTEXT_WIDTH: tl.constexpr,
239
+):
240
+ state_0 = tl.full((), _CONTEXT_DOMAIN, tl.uint32)
241
+ state_1 = tl.full((), CONTEXT_WIDTH, tl.uint32)
242
+ state_2 = tl.full((), 0, tl.uint32)
243
+ state_3 = tl.full((), 0, tl.uint32)
244
+ for offset in range(0, CONTEXT_WIDTH, 4):
245
+ context_0 = tl.load(contexts_row_ptr + offset).to(tl.uint32)
246
+ if offset + 1 < CONTEXT_WIDTH:
247
+ context_1 = tl.load(contexts_row_ptr + offset + 1).to(tl.uint32)
248
+ else:
249
+ context_1 = tl.full((), _UINT32_MASK - 1, tl.uint32)
250
+ if offset + 2 < CONTEXT_WIDTH:
251
+ context_2 = tl.load(contexts_row_ptr + offset + 2).to(tl.uint32)
252
+ else:
253
+ context_2 = tl.full((), _UINT32_MASK - 2, tl.uint32)
254
+ if offset + 3 < CONTEXT_WIDTH:
255
+ context_3 = tl.load(contexts_row_ptr + offset + 3).to(tl.uint32)
256
+ else:
257
+ context_3 = tl.full((), _UINT32_MASK - 3, tl.uint32)
258
+ state_0, state_1, state_2, state_3 = _philox4x32_10(
259
+ (state_0 ^ context_0) & _UINT32_MASK,
260
+ (state_1 ^ context_1) & _UINT32_MASK,
261
+ (state_2 ^ context_2) & _UINT32_MASK,
262
+ (state_3 ^ context_3) & _UINT32_MASK,
263
+ (key_0 ^ offset) & _UINT32_MASK,
264
+ key_1,
265
+ )
266
+ return state_0, state_1, state_2, state_3
267
+
268
+
269
+@triton.jit
270
+def _philox_candidate_words(groups, state_0, state_1, state_2, state_3, key_0, key_1):
271
+ candidate_words = groups.to(tl.uint32)
272
+ vector_zero = candidate_words * 0
273
+ return _philox4x32_10(
274
+ candidate_words & _UINT32_MASK,
275
+ state_0 + vector_zero,
276
+ state_1 + vector_zero,
277
+ state_2 + vector_zero,
278
+ ((key_0 ^ state_3) & _UINT32_MASK) + vector_zero,
279
+ ((key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK) + vector_zero,
280
+ )
281
+
282
+
283
+@triton.jit
284
+def _philox_gumbel_from_logits(logits, word):
285
logits = tl.where(logits != logits, float("-inf"), logits)
80
- uniform = _uint32_to_uniform(output)
286
+ uniform = _uint32_to_uniform(word)
287
return logits - tl.log(-tl.log(uniform))
288
289
290
+@triton.jit
291
+def _gumbel_value(logits_ptr, output, mask):
292
+ logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32)
293
+ return _philox_gumbel_from_logits(logits, output)
294
+
295
+
296
+@triton.jit
297
+def philox_gumbel_block_argmax(
298
+ logits,
299
+ mask,
300
+ block_idx,
301
+ contexts_row_ptr,
302
+ key_0,
303
+ key_1,
304
+ CONTEXT_WIDTH: tl.constexpr,
305
+ BLOCK_SIZE: tl.constexpr,
306
+):
307
+ """Apply the detector-compatible PRF to one vocabulary block."""
308
+ tl.static_assert(BLOCK_SIZE % 4 == 0)
309
+ state_0, state_1, state_2, state_3 = _philox_context_state(
310
+ contexts_row_ptr, key_0, key_1, CONTEXT_WIDTH
311
+ )
312
+ groups = block_idx * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4)
313
+ output_0, output_1, output_2, output_3 = _philox_candidate_words(
314
+ groups, state_0, state_1, state_2, state_3, key_0, key_1
315
+ )
316
+ # Lay the per-group words out in token order: 4g+0, 4g+1, 4g+2, 4g+3.
317
+ words = tl.interleave(
318
+ tl.interleave(output_0, output_2),
319
+ tl.interleave(output_1, output_3),
320
+ )
321
+ values = tl.where(
322
+ mask,
323
+ _philox_gumbel_from_logits(logits.to(tl.float32), words),
324
+ float("-inf"),
325
+ )
326
+ return tl.max(values, axis=0, return_indices=True)
327
+
328
+
329
@triton.jit
330
def _select_max(best_value, best_token, candidate_value, candidate):
331
take_candidate = candidate_value > best_value
@@ -100,60 +345,63 @@ def _philox_gumbel_kernel(
345
logits_stride,
346
contexts_ptr,
347
context_stride,
348
+ skip_mask_ptr,
349
+ expanded_idx_mapping_ptr,
350
+ seeds_ptr,
351
+ pos_ptr,
352
+ temp_ptr,
353
key_0_value,
354
key_1_value,
355
vocab_size,
356
CONTEXT_WIDTH: tl.constexpr,
357
BLOCK_SIZE: tl.constexpr,
358
+ USE_FP64: tl.constexpr,
359
):
360
row = tl.program_id(0).to(tl.int64)
361
block_index = tl.program_id(1)
362
groups = block_index * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4)
363
113
- key_0 = key_0_value.to(tl.uint32)
114
- key_1 = key_1_value.to(tl.uint32)
115
- state_0 = tl.full((), _CONTEXT_DOMAIN, tl.uint32)
116
- state_1 = tl.full((), CONTEXT_WIDTH, tl.uint32)
117
- state_2 = tl.full((), 0, tl.uint32)
118
- state_3 = tl.full((), 0, tl.uint32)
119
- for offset in range(0, CONTEXT_WIDTH, 4):
120
- context_0 = tl.load(contexts_ptr + row * context_stride + offset).to(tl.uint32)
121
- if offset + 1 < CONTEXT_WIDTH:
122
- context_1 = tl.load(contexts_ptr + row * context_stride + offset + 1).to(
123
- tl.uint32
364
+ if skip_mask_ptr is not None:
365
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + row).to(tl.int64)
366
+ valid_req = req_state_idx >= 0
367
+ if tl.load(skip_mask_ptr + row):
368
+ temp = tl.load(temp_ptr + req_state_idx, mask=valid_req, other=0.0).to(
369
+ tl.float32
370
)
125
- else:
126
- context_1 = tl.full((), _UINT32_MASK - 1, tl.uint32)
127
- if offset + 2 < CONTEXT_WIDTH:
128
- context_2 = tl.load(contexts_ptr + row * context_stride + offset + 2).to(
129
- tl.uint32
371
+ seed = tl.load(seeds_ptr + req_state_idx, mask=valid_req, other=0)
372
+ pos = tl.load(pos_ptr + row)
373
+ candidate = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
374
+ mask = candidate < vocab_size
375
+ logits_row = logits_ptr + row * logits_stride
376
+ logits = tl.load(logits_row + candidate, mask=mask, other=float("-inf")).to(
377
+ tl.float32
378
)
131
- else:
132
- context_2 = tl.full((), _UINT32_MASK - 2, tl.uint32)
133
- if offset + 3 < CONTEXT_WIDTH:
134
- context_3 = tl.load(contexts_ptr + row * context_stride + offset + 3).to(
135
- tl.uint32
379
+ value, index = gumbel_noised_argmax(
380
+ logits,
381
+ candidate,
382
+ mask,
383
+ seed,
384
+ pos,
385
+ temp,
386
+ IS_DRAFTING=False,
387
+ USE_FP64=USE_FP64,
388
+ APPLY_TEMPERATURE=False,
389
)
137
- else:
138
- context_3 = tl.full((), _UINT32_MASK - 3, tl.uint32)
139
- state_0, state_1, state_2, state_3 = _philox4x32_10(
140
- (state_0 ^ context_0) & _UINT32_MASK,
141
- (state_1 ^ context_1) & _UINT32_MASK,
142
- (state_2 ^ context_2) & _UINT32_MASK,
143
- (state_3 ^ context_3) & _UINT32_MASK,
144
- (key_0 ^ offset) & _UINT32_MASK,
145
- key_1,
146
- )
390
+ token_id = block_index * BLOCK_SIZE + index
391
+ tl.store(
392
+ local_argmax_ptr + row * local_argmax_stride + block_index,
393
+ token_id,
394
+ )
395
+ tl.store(local_max_ptr + row * local_max_stride + block_index, value)
396
+ return
397
148
- candidate_words = groups.to(tl.uint32)
149
- vector_zero = candidate_words * 0
150
- output_0, output_1, output_2, output_3 = _philox4x32_10(
151
- candidate_words & _UINT32_MASK,
152
- state_0 + vector_zero,
153
- state_1 + vector_zero,
154
- state_2 + vector_zero,
155
- ((key_0 ^ state_3) & _UINT32_MASK) + vector_zero,
156
- ((key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK) + vector_zero,
398
+ key_0 = key_0_value.to(tl.uint32)
399
+ key_1 = key_1_value.to(tl.uint32)
400
+ state_0, state_1, state_2, state_3 = _philox_context_state(
401
+ contexts_ptr + row * context_stride, key_0, key_1, CONTEXT_WIDTH
402
+ )
403
+ output_0, output_1, output_2, output_3 = _philox_candidate_words(
404
+ groups, state_0, state_1, state_2, state_3, key_0, key_1
405
)
406
candidate_0 = groups * 4
407
candidate_1 = candidate_0 + 1
@@ -184,17 +432,42 @@ def _philox_gumbel_kernel(
432
433
434
def philox_gumbel_sample(
187
- logits: torch.Tensor, contexts: torch.Tensor, key: int
435
+ logits: torch.Tensor,
436
+ contexts: torch.Tensor,
437
+ key: int,
438
+ *,
439
+ skip_mask: torch.Tensor | None = None,
440
+ expanded_idx_mapping: torch.Tensor | None = None,
441
+ temperatures: torch.Tensor | None = None,
442
+ seeds: torch.Tensor | None = None,
443
+ positions: torch.Tensor | None = None,
444
+ use_fp64: bool = False,
445
) -> torch.Tensor:
446
+ sampling_state = (expanded_idx_mapping, temperatures, seeds, positions)
447
+ if skip_mask is None and any(value is not None for value in sampling_state):
448
+ raise ValueError("sampling state requires skip_mask")
449
+ if skip_mask is not None and any(value is None for value in sampling_state):
450
+ raise ValueError("skip_mask requires complete sampling state")
451
+
452
if logits.stride(-1) != 1:
453
logits = logits.contiguous()
454
if contexts.stride(-1) != 1:
455
contexts = contexts.contiguous()
456
+ if skip_mask is not None:
457
+ skip_mask = skip_mask.contiguous()
458
+ if expanded_idx_mapping is not None:
459
+ expanded_idx_mapping = expanded_idx_mapping.contiguous()
460
+ if positions is not None:
461
+ positions = positions.contiguous()
462
num_tokens, vocab_size = logits.shape
463
block_size = 1024
464
num_blocks = triton.cdiv(vocab_size, block_size)
465
local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)
197
- local_max = logits.new_empty(num_tokens, num_blocks, dtype=torch.float32)
466
+ local_max = logits.new_empty(
467
+ num_tokens,
468
+ num_blocks,
469
+ dtype=torch.float64 if use_fp64 else torch.float32,
470
+ )
471
_philox_gumbel_kernel[(num_tokens, num_blocks)](
472
local_argmax,
473
local_argmax.stride(0),
@@ -204,11 +477,17 @@ def philox_gumbel_sample(
477
logits.stride(0),
478
contexts,
479
contexts.stride(0),
480
+ skip_mask,
481
+ expanded_idx_mapping,
482
+ seeds,
483
+ positions,
484
+ temperatures,
485
key & _UINT32_MASK_VALUE,
486
key >> 32,
487
vocab_size,
488
CONTEXT_WIDTH=contexts.shape[-1],
489
BLOCK_SIZE=block_size,
490
+ USE_FP64=use_fp64,
491
)
492
max_block_index = local_max.argmax(dim=-1, keepdim=True)
493
return local_argmax.gather(dim=-1, index=max_block_index).view(-1)
vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+9
-1
@@ -137,7 +137,7 @@ class DSparkSpeculator(DFlashSpeculator):
137
138
# sample_pos is the predicted token's position P. Sampling keys a draw
139
# by the position before the sampled token, P-1.
140
- return gumbel_sample(
140
+ sampled = gumbel_sample(
141
logits,
142
idx_map,
143
self.temperature,
@@ -149,6 +149,14 @@ class DSparkSpeculator(DFlashSpeculator):
149
logits_cache_col=self._step_cols[step],
150
use_fp64=self.use_fp64_gumbel,
151
)
152
+ if self.draft_watermarker is not None:
153
+ sampled = self.draft_watermarker.sample(
154
+ logits,
155
+ sampled,
156
+ idx_map,
157
+ self.temperature,
158
+ )
159
+ return sampled
160
161
def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
162
if self._draft_topk is not None:
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+27
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable, Iterator
4
+from typing import Any
5
6
import numpy as np
7
import torch
@@ -10,6 +11,7 @@ from vllm.config.model import PROCESSED_LOGPROBS_MODES
11
from vllm.triton_utils import tl, triton
12
from vllm.v1.outputs import LogprobsTensors
13
from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
14
+from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
15
from vllm.v1.worker.gpu.input_batch import (
16
InputBatch,
17
get_num_sampled_and_rejected,
@@ -78,8 +80,13 @@ class RejectionSampler:
80
sampler: Sampler,
81
spec_config: SpeculativeConfig,
82
device: torch.device,
83
+ *,
84
+ watermark_key: int | None = None,
85
):
86
self.sampler = sampler
87
+ self.watermark_key = watermark_key
88
+ if watermark_key is not None:
89
+ assert isinstance(sampler, GPUWatermarkSampler)
90
self.num_speculative_steps = spec_config.num_speculative_tokens
91
self.enable_adaptive_verification = spec_config.enable_adaptive_verification
92
rejection_sample_method = spec_config.rejection_sample_method
@@ -141,6 +148,23 @@ class RejectionSampler:
148
in ("raw_logits", "processed_logits"),
149
)
150
151
+ def _watermarking_kwargs(
152
+ self,
153
+ draft_sampled: torch.Tensor,
154
+ expanded_idx_mapping: torch.Tensor,
155
+ expanded_local_pos: torch.Tensor,
156
+ ) -> dict[str, Any]:
157
+ if self.watermark_key is None:
158
+ return {}
159
+ assert isinstance(self.sampler, GPUWatermarkSampler)
160
+ return {
161
+ "contexts": self.sampler._get_contexts(
162
+ expanded_idx_mapping, expanded_local_pos, draft_sampled
163
+ ),
164
+ "watermarking": self.sampler.watermarking.gpu,
165
+ "watermark_key": self.watermark_key,
166
+ }
167
+
168
def _verify(
169
self,
170
logits: torch.Tensor,
@@ -177,6 +201,9 @@ class RejectionSampler:
201
self.synthetic_conditional_rates,
202
use_fp64=self.sampler.use_fp64_gumbel,
203
use_block_verification=self.use_block_verification,
204
+ **self._watermarking_kwargs(
205
+ draft_sampled, expanded_idx_mapping, expanded_local_pos
206
+ ),
207
)
208
return processed_logits, sampled, num_sampled
209
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+128
-35
@@ -4,6 +4,7 @@ import torch
4
5
from vllm.triton_utils import tl, tldevice, triton
6
from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32
7
+from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_block_argmax
8
9
10
@triton.jit
@@ -693,6 +694,39 @@ def _rejection_kernel(
694
695
696
@triton.jit
697
+def _seeded_resample_argmax(
698
+ residual_logits,
699
+ block,
700
+ mask,
701
+ resample_token_idx,
702
+ expanded_idx_mapping_ptr,
703
+ temp_ptr,
704
+ seed_ptr,
705
+ pos_ptr,
706
+ vocab_size,
707
+ USE_FP64: tl.constexpr,
708
+):
709
+ return gumbel_block_argmax(
710
+ residual_logits,
711
+ block,
712
+ mask,
713
+ resample_token_idx,
714
+ expanded_idx_mapping_ptr,
715
+ temp_ptr,
716
+ seed_ptr,
717
+ pos_ptr,
718
+ None, # logits_cache_ptr
719
+ 0, # logits_cache_stride_0
720
+ 0, # logits_cache_stride_1
721
+ None, # logits_cache_col_ptr
722
+ vocab_size,
723
+ IS_DRAFTING=False,
724
+ APPLY_TEMPERATURE=False,
725
+ USE_FP64=USE_FP64,
726
+ )
727
+
728
+
729
+@triton.jit(do_not_specialize=["watermark_key_0", "watermark_key_1"])
730
def _resample_kernel(
731
# [num_reqs, num_blocks]
732
resampled_local_argmax_ptr,
@@ -727,11 +761,20 @@ def _resample_kernel(
761
pos_ptr,
762
# [num_logits]
763
cumulative_log_p_ptr,
764
+ # [num_logits, CONTEXT_WIDTH]
765
+ contexts_ptr,
766
+ contexts_stride,
767
+ # [max_num_reqs], uint8 view of a bool tensor
768
+ watermarking_ptr,
769
+ watermark_key_0,
770
+ watermark_key_1,
771
vocab_size,
772
BLOCK_SIZE: tl.constexpr,
773
HAS_DRAFT_LOGITS: tl.constexpr,
774
USE_FP64: tl.constexpr,
775
USE_BLOCK_VERIFICATION: tl.constexpr,
776
+ CONTEXT_WIDTH: tl.constexpr,
777
+ WATERMARK: tl.constexpr,
778
):
779
req_idx = tl.program_id(0)
780
resample_idx = tl.load(rejected_step_ptr + req_idx)
@@ -763,10 +806,7 @@ def _resample_kernel(
806
other=float("-inf"),
807
).to(tl.float32)
808
766
- # Compute the residual logits to resample the rejected token from.
809
if is_bonus or not is_valid_rejected_draft:
768
- # Bonus token (no rejections) or -1 placeholder token. In either case,
769
- # directly use the target logits.
810
residual_logits = target_logits
811
elif HAS_DRAFT_LOGITS:
812
# draft_logits is stored pre-temperature, so apply scale first.
@@ -796,12 +836,7 @@ def _resample_kernel(
836
)
837
target_log_probs += log_p_tau
838
draft_log_probs = draft_logits - draft_lse
799
- # Compute the residual:
800
- # r(x) = max(p(x) - q(x), 0)
801
- # Gumbel sampling needs logits, so we compute it in log space:
802
- # log(r(x)) = log(max(exp(log_p(x)) - exp(log_q(x)), 0))
803
- # The more numerically stable form is:
804
- # log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0))
839
+ # Compute log(max(p - q, 0)) without subtracting probabilities.
840
ratio = tl.exp(draft_log_probs - target_log_probs)
841
residual_logits = tl.where(
842
ratio < 1.0,
@@ -809,38 +844,62 @@ def _resample_kernel(
844
float("-inf"),
845
).to(tl.float32)
846
else:
812
- # One-hot draft. The residual is just the target distribution with
813
- # the rejected draft token probability zeroed out.
814
- # NOTE: During block verification, the residual becomes:
815
- # 0 if x == rejected_draft_token
816
- # p_tau * M_b(x) / Z otherwise
817
- # Therefore p_tau is a constant that cancels under normalization,
818
- # and does not need to be applied.
847
+ # The block-verification factor is constant and cancels on normalization.
848
residual_logits = tl.where(
849
block != rejected_draft_token,
850
target_logits,
851
float("-inf"),
852
).to(tl.float32)
853
825
- # Resample the rejected/bonus token.
826
- value, idx = gumbel_block_argmax(
827
- residual_logits,
828
- block,
829
- mask,
830
- resample_token_idx,
831
- expanded_idx_mapping_ptr,
832
- temp_ptr,
833
- seed_ptr,
834
- pos_ptr,
835
- None, # logits_cache_ptr
836
- 0, # logits_cache_stride_0
837
- 0, # logits_cache_stride_1
838
- None, # logits_cache_col_ptr
839
- vocab_size,
840
- IS_DRAFTING=False,
841
- APPLY_TEMPERATURE=False,
842
- USE_FP64=USE_FP64,
843
- )
854
+ if WATERMARK:
855
+ # Padded and greedy rows retain the stock draw.
856
+ is_watermarked = (
857
+ tl.load(
858
+ watermarking_ptr + req_state_idx,
859
+ mask=req_state_idx >= 0,
860
+ other=0,
861
+ )
862
+ != 0
863
+ )
864
+ if is_watermarked & (temp != 0.0):
865
+ watermark_value, idx = philox_gumbel_block_argmax(
866
+ residual_logits,
867
+ mask,
868
+ block_idx,
869
+ contexts_ptr + resample_token_idx * contexts_stride,
870
+ watermark_key_0.to(tl.uint32),
871
+ watermark_key_1.to(tl.uint32),
872
+ CONTEXT_WIDTH,
873
+ BLOCK_SIZE,
874
+ )
875
+ # Detector compatibility fixes the keyed draw at fp32.
876
+ value = watermark_value.to(tl.float64) if USE_FP64 else watermark_value
877
+ else:
878
+ value, idx = _seeded_resample_argmax(
879
+ residual_logits,
880
+ block,
881
+ mask,
882
+ resample_token_idx,
883
+ expanded_idx_mapping_ptr,
884
+ temp_ptr,
885
+ seed_ptr,
886
+ pos_ptr,
887
+ vocab_size,
888
+ USE_FP64=USE_FP64,
889
+ )
890
+ else:
891
+ value, idx = _seeded_resample_argmax(
892
+ residual_logits,
893
+ block,
894
+ mask,
895
+ resample_token_idx,
896
+ expanded_idx_mapping_ptr,
897
+ temp_ptr,
898
+ seed_ptr,
899
+ pos_ptr,
900
+ vocab_size,
901
+ USE_FP64=USE_FP64,
902
+ )
903
token_id = block_idx * BLOCK_SIZE + idx
904
tl.store(
905
resampled_local_argmax_ptr
@@ -947,6 +1006,9 @@ def rejection_sample(
1006
synthetic_conditional_rates: torch.Tensor | None = None,
1007
use_fp64: bool = False,
1008
use_block_verification: bool = False,
1009
+ contexts: torch.Tensor | None = None,
1010
+ watermarking: torch.Tensor | None = None,
1011
+ watermark_key: int | None = None,
1012
) -> tuple[torch.Tensor, torch.Tensor]:
1013
assert target_logits.ndim == 2 and target_logits.stride(-1) == 1
1014
assert draft_logits is None or (
@@ -954,6 +1016,30 @@ def rejection_sample(
1016
)
1017
num_reqs = cu_num_logits.shape[0] - 1
1018
num_logits, vocab_size = target_logits.shape
1019
+
1020
+ watermark = contexts is not None
1021
+ assert watermark == (watermarking is not None) == (watermark_key is not None), (
1022
+ "contexts, watermarking and watermark_key must be set together."
1023
+ )
1024
+ contexts_stride = 0
1025
+ context_width = 1
1026
+ watermarking_bytes: torch.Tensor | None = None
1027
+ watermark_key_0 = 0
1028
+ watermark_key_1 = 0
1029
+ if contexts is not None:
1030
+ assert watermarking is not None and watermark_key is not None
1031
+ assert contexts.ndim == 2 and contexts.shape[0] == num_logits
1032
+ # Context words are hashed as uint32, so int32 (the request-state token
1033
+ # dtype) and int64 both land on the same PRF stream, -1 included.
1034
+ assert contexts.dtype in (torch.int32, torch.int64)
1035
+ if contexts.stride(-1) != 1:
1036
+ contexts = contexts.contiguous()
1037
+ assert watermarking.ndim == 1 and watermarking.dtype == torch.bool
1038
+ contexts_stride = contexts.stride(0)
1039
+ context_width = contexts.shape[-1]
1040
+ watermarking_bytes = watermarking.view(torch.uint8)
1041
+ watermark_key_0 = watermark_key & 0xFFFFFFFF
1042
+ watermark_key_1 = watermark_key >> 32
1043
draft_logits_stride_0 = 0
1044
draft_logits_stride_1 = 0
1045
if has_draft_logits := draft_logits is not None:
@@ -1165,11 +1251,18 @@ def rejection_sample(
1251
seed,
1252
pos,
1253
cumulative_log_p,
1254
+ contexts,
1255
+ contexts_stride,
1256
+ watermarking_bytes,
1257
+ watermark_key_0,
1258
+ watermark_key_1,
1259
vocab_size,
1260
BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
1261
HAS_DRAFT_LOGITS=has_draft_logits,
1262
USE_FP64=use_fp64,
1263
USE_BLOCK_VERIFICATION=use_block_verification,
1264
+ CONTEXT_WIDTH=context_width,
1265
+ WATERMARK=watermark,
1266
)
1267
1268
# Insert the resampled tokens into the output sampled.
vllm/v1/worker/gpu/spec_decode/speculator.py
+32
-1
@@ -17,6 +17,11 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
17
from vllm.model_executor.models import supports_multimodal_embeddings
18
from vllm.multimodal import MULTIMODAL_REGISTRY
19
from vllm.v1.kv_cache_interface import KVCacheConfig
20
+from vllm.v1.watermarking import create_watermarker
21
+from vllm.v1.watermarking.spec_decode import (
22
+ DraftWatermarker,
23
+ create_speculative_draft_watermarker,
24
+)
25
from vllm.v1.worker.gpu.attn_utils import (
26
build_attn_metadata,
27
init_attn_backend,
@@ -168,6 +173,17 @@ class DraftModelSpeculator(BaseSpeculator):
173
device=device,
174
)
175
176
+ self.draft_watermarker: DraftWatermarker | None = None
177
+ watermark_config = getattr(vllm_config, "watermark_config", None)
178
+ if watermark_config is not None:
179
+ watermarker = create_watermarker(watermark_config)
180
+ self.draft_watermarker = create_speculative_draft_watermarker(
181
+ watermarker,
182
+ self.max_num_reqs,
183
+ device,
184
+ watermark_config.allow_target_only_watermarking,
185
+ )
186
+
187
self.supports_mm_inputs = False
188
self.pcp_manager: PCPManager | None = None
189
@@ -373,7 +389,7 @@ class DraftModelSpeculator(BaseSpeculator):
389
) -> torch.Tensor:
390
if draft_logits is not None:
391
logits = self.model.compute_logits(hidden_states)
376
- return gumbel_sample(
392
+ sampled = gumbel_sample(
393
logits,
394
idx_mapping,
395
temperature,
@@ -385,8 +401,23 @@ class DraftModelSpeculator(BaseSpeculator):
401
logits_cache_col=draft_step,
402
use_fp64=self.use_fp64_gumbel,
403
)
404
+ if self.draft_watermarker is not None:
405
+ sampled = self.draft_watermarker.sample(
406
+ logits,
407
+ sampled,
408
+ idx_mapping,
409
+ temperature,
410
+ )
411
+ return sampled
412
return self._greedy_sample_draft(hidden_states)
413
414
+ def prepare_watermarking(
415
+ self, contexts: torch.Tensor, watermarking: torch.Tensor
416
+ ) -> None:
417
+ if self.draft_watermarker is None:
418
+ return
419
+ self.draft_watermarker.prepare(contexts, watermarking)
420
+
421
def _copy_request_inputs(
422
self,
423
num_reqs: int,