@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ea40bb9e90

[feature] Watermarked generation and detection (Gumbel-max algorithm) (#54053)

Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Raphaël Rialland committed Sep 10, 2026 at 06:00 UTC ea40bb9e905f8d552281dd1ec074f91865a0a242
32 files changed +1693 -15
docs/features/watermarking.md new
+147
@@ -0,0 +1,147 @@
1 +# Text watermarking
2 +
3 +Text watermarking embeds a statistical signal in generated token choices. A
4 +detector that knows the generation parameters can test for this signal without
5 +access to the model weights.
6 +
7 +The design and scope are discussed in
8 +[RFC #53916](https://github.com/vllm-project/vllm/issues/53916).
9 +
10 +## Configuration
11 +
12 +Enable watermarking by configuring an algorithm and secret key at engine
13 +startup:
14 +
15 +```bash
16 +vllm serve MODEL \
17 + --watermark-config '{"algorithm":"gumbel","key":42}'
18 +```
19 +
20 +Watermarking is disabled when `--watermark-config` is omitted. Gumbel-max (see
21 +`gumbel`) is the default algorithm within an enabled `WatermarkConfig`. When
22 +watermarking is configured, it is enabled for requests by default.
23 +
24 +Requests can opt out without changing the engine-level algorithm or key:
25 +
26 +```python
27 +from vllm import SamplingParams
28 +
29 +sampling_params = SamplingParams(watermarking=False)
30 +```
31 +
32 +The OpenAI-compatible APIs accept the same `watermarking: false` request field.
33 +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
39 +edits because an insertion, deletion, or substitution changes more subsequent
40 +contexts. Values above 16 are allowed but emit a warning.
41 +
42 +## Architecture
43 +
44 +`WatermarkConfig` selects an algorithm and PRF. Model Runner V2 constructs the
45 +corresponding `Watermarker`, and `GPUWatermarkSampler` invokes it for the final
46 +stochastic token selection after temperature, min-p, top-k, and top-p are
47 +applied. A watermarker can either select a token directly or transform logits
48 +and delegate to vLLM's random sampler.
49 +
50 +Detection is separate from generation. vLLM provides detector primitives for
51 +the reference algorithms. `WatermarkDetector` consumes token IDs, so callers
52 +remain responsible for using the tokenizer and watermark profile that match
53 +generation.
54 +
55 +## Algorithms
56 +
57 +### Gumbel-max
58 +
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
62 +[Aaronson's original presentation](https://simons.berkeley.edu/sites/default/files/2024-10/LLM24-2%20Slides%20-%20Scott%20Aaronson.pdf).
63 +
64 +Gumbel-max requires stochastic sampling. Greedy requests (`temperature=0`)
65 +bypass watermarking and emit a warning once per worker.
66 +
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.
70 +
71 +### SynthID-Text
72 +
73 +[SynthID-Text](https://www.nature.com/articles/s41586-024-08025-4) is planned but
74 +not currently implemented.
75 +
76 +## Pseudorandom functions
77 +
78 +A watermark PRF turns the secret key, token context, and candidate token into
79 +reproducible random values. Generation and detection must produce identical
80 +values across devices and releases. The values should also be uniform and
81 +independent enough for the sampling algorithm and detector statistics. PRF
82 +selection is an advanced compatibility and performance setting; most users
83 +should keep the default.
84 +
85 +Watermarked generation currently supports the `philox` PRF:
86 +
87 +- `philox` is based on the counter-based Philox4x32-10
88 + generator from the [Random123 paper](https://doi.org/10.1145/2063384.2063405).
89 + It is parallel, vectorizes on accelerators, and avoids CPU transfers, but is
90 + not a cryptographic PRF and does not provide key-recovery or forgery
91 + resistance. vLLM versions its input mapping and provides compatibility
92 + vectors so generation and detection remain interoperable.
93 +
94 +## Detection
95 +
96 +The detector primitives operate on token IDs and do not require model weights:
97 +
98 +```python
99 +from transformers import AutoTokenizer
100 +
101 +from vllm.v1.watermarking import GumbelWatermarkDetector
102 +
103 +tokenizer = AutoTokenizer.from_pretrained(MODEL)
104 +token_ids = tokenizer.encode(text, add_special_tokens=False)
105 +result = GumbelWatermarkDetector(key=42, prf="philox").detect(token_ids)
106 +print(result.p_value, result.is_watermarked)
107 +```
108 +
109 +The tokenizer, algorithm, PRF, key, and context width must match generation.
110 +Gumbel-max detection scores repeated contexts once by default so identical PRF
111 +random vectors are not treated as independent evidence. Keep
112 +`deduplicate_contexts=True` unless the detector's calibration has been adjusted
113 +for correlated scores.
114 +
115 +The reported p-value is calibrated under the assumption that scored PRF inputs
116 +are independent. A deployment uses one fixed key, so repeated structures across
117 +documents reuse the same PRF values and can make the realized false-positive
118 +rate key-dependent even when contexts are deduplicated within each document.
119 +Measure the false-positive rate on representative unwatermarked traffic with the
120 +deployed key before relying on `is_watermarked` for decisions.
121 +
122 +A minimal HTTP detector is available in
123 +`examples/basic/online_serving/watermark_detection_server.py`:
124 +
125 +```bash
126 +python examples/basic/online_serving/watermark_detection_server.py \
127 + --tokenizer MODEL --key 42 --prf philox
128 +```
129 +
130 +```bash
131 +curl http://localhost:8000/detect \
132 + -H 'Content-Type: application/json' \
133 + -d '{"text":"Text to inspect"}'
134 +```
135 +
136 +Scores and p-values expose information about the per-token watermark signal.
137 +Repeated queries can use this information to construct text that imitates
138 +watermarked output or to modify watermarked text so it is no longer detected.
139 +
140 +## Limitations
141 +
142 +- Watermarking is currently available only with Model Runner V2.
143 +- Gumbel-max cannot be configured with speculative decoding.
144 +- Beam search expands candidates from model log probabilities and does not apply
145 + Gumbel-max watermarking.
146 +- Models that replace the vLLM sampler with a custom sampler cannot use
147 + configured watermarking.
examples/basic/online_serving/watermark_detection_server.py new
+70
@@ -0,0 +1,70 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Minimal reference server for watermark detection."""
5 +
6 +import argparse
7 +
8 +import uvicorn
9 +from fastapi import FastAPI
10 +from pydantic import BaseModel
11 +
12 +from vllm.tokenizers import TokenizerLike, cached_get_tokenizer
13 +from vllm.v1.watermarking import GumbelWatermarkDetector
14 +
15 +app = FastAPI()
16 +tokenizer: TokenizerLike | None = None
17 +detector: GumbelWatermarkDetector | None = None
18 +
19 +
20 +class DetectionRequest(BaseModel):
21 + text: str
22 +
23 +
24 +class DetectionResponse(BaseModel):
25 + score: float
26 + p_value: float
27 + num_scored_tokens: int
28 + is_watermarked: bool
29 +
30 +
31 +@app.post("/detect")
32 +def detect(request: DetectionRequest) -> DetectionResponse:
33 + assert tokenizer is not None
34 + assert detector is not None
35 + token_ids = tokenizer.encode(request.text, add_special_tokens=False)
36 + result = detector.detect(token_ids)
37 + return DetectionResponse(
38 + score=result.score,
39 + p_value=result.p_value,
40 + num_scored_tokens=result.num_scored_tokens,
41 + is_watermarked=result.is_watermarked,
42 + )
43 +
44 +
45 +def parse_args() -> argparse.Namespace:
46 + parser = argparse.ArgumentParser(description=__doc__)
47 + parser.add_argument("--tokenizer", required=True)
48 + parser.add_argument("--key", required=True, type=int)
49 + parser.add_argument("--prf", choices=("philox",), default="philox")
50 + parser.add_argument("--context-width", type=int, default=4)
51 + parser.add_argument("--p-value-threshold", type=float, default=0.01)
52 + parser.add_argument("--host", default="127.0.0.1")
53 + parser.add_argument("--port", type=int, default=8000)
54 + return parser.parse_args()
55 +
56 +
57 +def main(args: argparse.Namespace) -> None:
58 + global tokenizer, detector
59 + tokenizer = cached_get_tokenizer(args.tokenizer)
60 + detector = GumbelWatermarkDetector(
61 + key=args.key,
62 + context_width=args.context_width,
63 + p_value_threshold=args.p_value_threshold,
64 + prf=args.prf,
65 + )
66 + uvicorn.run(app, host=args.host, port=args.port)
67 +
68 +
69 +if __name__ == "__main__":
70 + main(parse_args())
tests/engine/test_arg_utils.py
+21
@@ -46,6 +46,27 @@ def test_optional_type():
46 assert optional_type_func("42") == 42
47
48
49 +def test_watermark_config_cli():
50 + parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
51 + args = parser.parse_args(
52 + [
53 + "--model",
54 + "dummy",
55 + "--watermark-config",
56 + '{"algorithm":"gumbel","key":42,"prf":"philox"}',
57 + ]
58 + )
59 +
60 + config = EngineArgs.from_cli_args(args).create_watermark_config()
61 +
62 + assert config is not None
63 + assert config.algorithm == "gumbel"
64 + assert config.key == 42
65 + assert config.context_width == 4
66 + assert config.prf == "philox"
67 + assert not config.supports_speculative_decoding
68 +
69 +
70 @pytest.mark.parametrize(
71 "options",
72 [
tests/entrypoints/openai/test_watermarking.py new
+32
@@ -0,0 +1,32 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
5 +from vllm.entrypoints.openai.completion.protocol import CompletionRequest
6 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
7 +
8 +
9 +def test_chat_request_disables_watermarking():
10 + request = ChatCompletionRequest(
11 + messages=[{"role": "user", "content": "hello"}], watermarking=False
12 + )
13 +
14 + params = request.to_sampling_params(max_tokens=1, default_sampling_params={})
15 +
16 + assert not params.watermarking
17 +
18 +
19 +def test_completion_request_disables_watermarking():
20 + request = CompletionRequest(prompt="hello", watermarking=False)
21 +
22 + params = request.to_sampling_params(max_tokens=1)
23 +
24 + assert not params.watermarking
25 +
26 +
27 +def test_responses_request_disables_watermarking():
28 + request = ResponsesRequest(input="hello", watermarking=False)
29 +
30 + params = request.to_sampling_params(default_max_tokens=1)
31 +
32 + assert not params.watermarking
tests/entrypoints/serve/utils/test_api_utils.py
+8 -1
@@ -120,17 +120,24 @@ class TestRedactSensitiveArgs:
120 API_KEY = "sk-test-secret-12345"
121
122 def test_redact_replaces_sensitive_values_only(self):
123 - args = {"api_key": self.API_KEY, "hf_token": "hf_secret", "other": "visible"}
123 + args = {
124 + "api_key": self.API_KEY,
125 + "hf_token": "hf_secret",
126 + "watermark_config": {"algorithm": "gumbel", "key": 42},
127 + "other": "visible",
128 + }
129 redacted = redact_sensitive_args(args)
130 assert redacted == {
131 "api_key": "***",
132 "hf_token": "***",
133 + "watermark_config": "***",
134 "other": "visible",
135 }
136 # original dict must not be mutated
137 assert args == {
138 "api_key": self.API_KEY,
139 "hf_token": "hf_secret",
140 + "watermark_config": {"algorithm": "gumbel", "key": 42},
141 "other": "visible",
142 }
143
tests/test_config.py
+61
@@ -32,6 +32,7 @@ from vllm.config import (
32 SchedulerConfig,
33 SpeculativeConfig,
34 VllmConfig,
35 + WatermarkConfig,
36 update_config,
37 )
38 from vllm.config.compilation import CompilationMode, CUDAGraphMode
@@ -2574,6 +2575,66 @@ def test_draft_sample_method_gumbel_is_rejected():
2575 )
2576
2577
2578 +def _watermarked_vllm_config() -> VllmConfig:
2579 + config = object.__new__(VllmConfig)
2580 + config.watermark_config = WatermarkConfig(key=42)
2581 + config.speculative_config = None
2582 + return config
2583 +
2584 +
2585 +def test_gumbel_watermark_rejects_speculative_decoding():
2586 + config = _watermarked_vllm_config()
2587 + config.speculative_config = SpeculativeConfig(
2588 + method="ngram",
2589 + num_speculative_tokens=1,
2590 + )
2591 +
2592 + with pytest.raises(ValueError, match="does not support speculative decoding"):
2593 + config._check_watermarking_unsupported()
2594 +
2595 +
2596 +def test_gumbel_watermark_rejects_beam_search():
2597 + with pytest.raises(ValueError, match="Beam search is not supported"):
2598 + _watermarked_vllm_config()._check_watermarking_unsupported(beam_search=True)
2599 +
2600 +
2601 +def test_gumbel_watermark_rejects_custom_sampler():
2602 + with pytest.raises(ValueError, match="custom samplers are not supported"):
2603 + _watermarked_vllm_config()._check_watermarking_unsupported(custom_sampler=True)
2604 +
2605 +
2606 +def test_watermark_key_must_fit_in_64_bits():
2607 + with pytest.raises(ValueError, match="64 bits"):
2608 + WatermarkConfig(key=2**64)
2609 +
2610 +
2611 +def test_unknown_watermark_prf_is_rejected():
2612 + with pytest.raises(ValidationError):
2613 + pydantic.TypeAdapter(WatermarkConfig).validate_python(
2614 + {"key": 42, "prf": "unsupported"}
2615 + )
2616 +
2617 +
2618 +def test_watermark_key_is_excluded_from_serialization():
2619 + serialized = pydantic.TypeAdapter(WatermarkConfig).dump_python(
2620 + WatermarkConfig(key=42), mode="json"
2621 + )
2622 +
2623 + assert "key" not in serialized
2624 +
2625 +
2626 +def test_watermarking_forces_model_runner_v2(monkeypatch):
2627 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
2628 +
2629 + with patch("vllm.config.vllm.logger.info_once") as info_once:
2630 + assert _watermarked_vllm_config().use_v2_model_runner
2631 +
2632 + info_once.assert_called_once_with(
2633 + "Watermarking requires Model Runner V2 and overrides "
2634 + "VLLM_USE_V2_MODEL_RUNNER=0."
2635 + )
2636 +
2637 +
2638 @patch("vllm.config.speculative.ModelConfig")
2639 def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls):
2640 """Regression test: MTP + runai_streamer should use model_weights (original
tests/watermarking/test_detection.py new
+91
@@ -0,0 +1,91 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.v1.watermarking import GumbelWatermarkDetector
5 +
6 +_NON_WATERMARKED_TOKEN_IDS = list(range(64))
7 +_WATERMARKED_TOKEN_IDS = [
8 + 85,
9 + 56,
10 + 123,
11 + 77,
12 + 105,
13 + 62,
14 + 9,
15 + 104,
16 + 16,
17 + 98,
18 + 123,
19 + 22,
20 + 16,
21 + 35,
22 + 47,
23 + 127,
24 + 77,
25 + 67,
26 + 19,
27 + 2,
28 + 10,
29 + 8,
30 + 75,
31 + 22,
32 + 127,
33 + 30,
34 + 95,
35 + 9,
36 + 101,
37 + 14,
38 + 113,
39 + 122,
40 + 85,
41 + 66,
42 + 122,
43 + 50,
44 + 38,
45 + 119,
46 + 15,
47 + 58,
48 + 68,
49 + 62,
50 + 24,
51 + 126,
52 + 80,
53 + 126,
54 + 68,
55 + 83,
56 + 81,
57 + 126,
58 + 62,
59 + 31,
60 + 18,
61 + 81,
62 + 64,
63 + 47,
64 + 22,
65 + 70,
66 + 47,
67 + 109,
68 + 54,
69 + 21,
70 + 65,
71 + 38,
72 +]
73 +
74 +
75 +def test_detection_on_known_sequence():
76 + detector = GumbelWatermarkDetector(key=42, context_width=4, prf="philox")
77 +
78 + assert detector.detect(_WATERMARKED_TOKEN_IDS).is_watermarked
79 + assert not detector.detect(_NON_WATERMARKED_TOKEN_IDS).is_watermarked
80 +
81 +
82 +def test_detector_handles_empty_input():
83 + detector = GumbelWatermarkDetector(key=42, context_width=4, prf="philox")
84 + contexts, targets = detector._prepare_inputs([])
85 + detection = detector.detect([])
86 +
87 + assert contexts.shape == (0, 4)
88 + assert targets.shape == (0,)
89 + assert detection.score == 0
90 + assert detection.p_value == 1
91 + assert not detection.is_watermarked
tests/watermarking/test_gumbel.py new
+84
@@ -0,0 +1,84 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.platforms import current_platform
8 +from vllm.v1.watermarking import GumbelWatermarkDetector, GumbelWatermarker
9 +from vllm.v1.watermarking.gumbel import _gamma_survival_integer_shape
10 +
11 +
12 +def test_gamma_survival_integer_shape():
13 + assert _gamma_survival_integer_shape(0.0, 4) == 1.0
14 + assert _gamma_survival_integer_shape(2.0, 1) == pytest.approx(0.1353352832)
15 +
16 +
17 +def test_detector_deduplicates_repeated_prf_inputs_by_default():
18 + detection = GumbelWatermarkDetector(key=42, context_width=1).detect([1, 1, 1, 1, 1])
19 +
20 + assert detection.num_scored_tokens == 2
21 +
22 +
23 +def test_detector_can_score_repeated_prf_inputs():
24 + detection = GumbelWatermarkDetector(
25 + key=42, context_width=1, deduplicate_contexts=False
26 + ).detect([1, 1, 1, 1, 1])
27 +
28 + assert detection.num_scored_tokens == 5
29 +
30 +
31 +def test_detector_deduplicates_context_even_when_target_differs():
32 + detection = GumbelWatermarkDetector(key=42, context_width=1).detect([1, 2, 1, 3])
33 +
34 + assert detection.num_scored_tokens == 3
35 +
36 +
37 +@pytest.mark.skipif(
38 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
39 +)
40 +@pytest.mark.parametrize("key", [1, 2**32, 42, 15726070495360670683])
41 +@pytest.mark.parametrize("context_width", [1, 4, 16])
42 +def test_fused_watermarker_matches_cpu(key: int, context_width: int):
43 + torch.manual_seed(0)
44 + contexts = torch.randint(0, 248320, (32, context_width), dtype=torch.int64)
45 + logits = torch.randn(32, 8193)
46 + watermarker = GumbelWatermarker(key, context_width)
47 +
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()
52 +
53 + torch.testing.assert_close(actual, expected, rtol=0, atol=0)
54 +
55 +
56 +@pytest.mark.skipif(
57 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
58 +)
59 +def test_fused_watermarker_handles_nan_logits():
60 + contexts = torch.zeros((2, 4), dtype=torch.int64, device="cuda")
61 + logits = torch.full((2, 1025), float("nan"), device="cuda")
62 + watermarker = GumbelWatermarker(key=42, context_width=4)
63 +
64 + token_ids = watermarker.sample(logits, contexts, lambda values: None).token_ids
65 +
66 + assert torch.all((token_ids >= 0) & (token_ids < logits.shape[-1]))
67 +
68 +
69 +@pytest.mark.skipif(
70 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
71 +)
72 +def test_fused_watermarker_handles_noncontiguous_inputs():
73 + contexts = torch.randint(0, 248320, (8, 8), dtype=torch.int64, device="cuda")[
74 + :, ::2
75 + ]
76 + logits = torch.randn(8, 2050, device="cuda")[:, ::2]
77 + watermarker = GumbelWatermarker(key=42, context_width=4)
78 +
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
83 +
84 + torch.testing.assert_close(actual, expected, rtol=0, atol=0)
tests/watermarking/test_prf.py new
+44
@@ -0,0 +1,44 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.platforms import current_platform
8 +from vllm.v1.watermarking.prfs import PhiloxPRF
9 +
10 +
11 +def test_philox_compatibility_vector():
12 + contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]])
13 + token_ids = torch.tensor([7, 8])
14 + expected_mantissas = [[1661833, 10571167], [9742331, 13375724]]
15 + expected = (
16 + (torch.tensor(expected_mantissas, dtype=torch.float64) + 1) / (2**24 + 1)
17 + ).to(torch.float32)
18 +
19 + assert torch.equal(PhiloxPRF(42).uniform(contexts, token_ids), expected)
20 +
21 +
22 +def test_prf_pairs_contexts_with_target_tokens():
23 + contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]])
24 + token_ids = torch.tensor([[7], [8]])
25 +
26 + assert PhiloxPRF(42).uniform(contexts, token_ids).shape == (2, 1)
27 +
28 +
29 +@pytest.mark.skipif(
30 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
31 +)
32 +@pytest.mark.parametrize("key", [42, 15726070495360670683])
33 +@pytest.mark.parametrize("context_width", [1, 4, 16])
34 +def test_philox_accelerator_matches_cpu(key: int, context_width: int):
35 + contexts = torch.arange(2 * context_width, dtype=torch.int64).reshape(
36 + 2, context_width
37 + )
38 + token_ids = torch.arange(1024, dtype=torch.int64)
39 + prf = PhiloxPRF(key)
40 +
41 + expected = prf.uniform(contexts, token_ids)
42 + actual = prf.uniform(contexts.cuda(), token_ids.cuda()).cpu()
43 +
44 + torch.testing.assert_close(actual, expected, rtol=0, atol=0)
tests/watermarking/test_watermarking.py new
+157
@@ -0,0 +1,157 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +
6 +import numpy as np
7 +import pytest
8 +import torch
9 +
10 +from vllm import SamplingParams
11 +from vllm.config.watermarking import WatermarkConfig
12 +from vllm.v1.watermarking import create_watermarker
13 +from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
14 +from vllm.v1.watermarking.watermarker import WatermarkSample
15 +from vllm.v1.worker.gpu.sample.sampler import Sampler
16 +
17 +
18 +@pytest.mark.parametrize("algorithm", ["gumbel"])
19 +def test_watermarker_contract(algorithm: str):
20 + watermarker = create_watermarker(
21 + WatermarkConfig(algorithm=algorithm, key=42, context_width=4)
22 + )
23 + logits = torch.zeros(2, 128)
24 + contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]])
25 + random_sample = lambda sample_logits: sample_logits.argmax(dim=-1)
26 +
27 + first = watermarker.sample(logits, contexts, random_sample)
28 + second = watermarker.sample(logits, contexts, random_sample)
29 +
30 + assert first.token_ids.shape == (2,)
31 + assert first.logits.shape == logits.shape
32 + assert torch.equal(first.token_ids, second.token_ids)
33 +
34 +
35 +def test_gumbel_config_warns_about_degenerate_generations(monkeypatch):
36 + messages: list[str] = []
37 + monkeypatch.setattr(
38 + "vllm.config.watermarking.logger.warning_once",
39 + lambda message, *, scope: messages.append(message),
40 + )
41 +
42 + WatermarkConfig(key=42)
43 +
44 + assert messages == [
45 + "Single-key Gumbel-max watermarking may increase the frequency of "
46 + "degenerate generations, including repetition loops."
47 + ]
48 +
49 +
50 +def test_large_context_width_warns_but_is_allowed():
51 + config = WatermarkConfig(key=42, context_width=17)
52 +
53 + with pytest.warns(UserWarning, match="reduce robustness to edits"):
54 + watermarker = create_watermarker(config)
55 +
56 + assert watermarker.context_width == 17
57 +
58 +
59 +def test_sampling_params_can_disable_watermarking():
60 + assert SamplingParams().watermarking
61 + assert not SamplingParams.from_optional(watermarking=False).watermarking
62 +
63 +
64 +def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch):
65 + sampler = object.__new__(GPUWatermarkSampler)
66 + sampler.watermarking = SimpleNamespace(np=np.ones(1, dtype=bool))
67 + messages: list[str] = []
68 + monkeypatch.setattr(Sampler, "add_request", lambda *args, **kwargs: None)
69 + monkeypatch.setattr(
70 + "vllm.v1.watermarking.gpu_sampler.logger.warning_once", messages.append
71 + )
72 +
73 + sampler.add_request(0, 1, SamplingParams(temperature=0))
74 + sampler.add_request(0, 1, SamplingParams(temperature=1))
75 + sampler.add_request(0, 1, SamplingParams(temperature=0, watermarking=False))
76 +
77 + assert messages == [
78 + (
79 + "Watermarking is enabled, but greedy decoding (temperature=0) cannot be "
80 + "watermarked. This request will use ordinary greedy sampling."
81 + )
82 + ]
83 +
84 +
85 +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 +
92 + sampler = object.__new__(GPUWatermarkSampler)
93 + sampler.watermarker = StubWatermarker()
94 + sampler.watermarking = SimpleNamespace(
95 + np=np.array([True, False]), gpu=torch.tensor([True, False])
96 + )
97 + sampler.sampling_states = SimpleNamespace(
98 + temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)),
99 + seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)),
100 + )
101 + sampler.use_fp64_gumbel = False
102 + sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
103 + 2, 1, dtype=torch.int64
104 + )
105 + monkeypatch.setattr(
106 + "vllm.v1.watermarking.gpu_sampler.gumbel_sample",
107 + lambda *args, **kwargs: torch.tensor([3, 4]),
108 + )
109 + logits = torch.zeros(2, 8)
110 +
111 + sampled, output_logits = sampler._sample_random(
112 + logits,
113 + torch.tensor([0, 1]),
114 + np.array([0, 1]),
115 + torch.zeros(2, dtype=torch.int64),
116 + None,
117 + None,
118 + False,
119 + )
120 +
121 + assert torch.equal(sampled, torch.tensor([7, 4]))
122 + assert torch.equal(output_logits[0], torch.full((8,), 10.0))
123 + assert torch.equal(output_logits[1], logits[1])
124 +
125 +
126 +def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch):
127 + class StubWatermarker:
128 + context_width = 1
129 +
130 + def sample(self, logits, contexts, random_sample):
131 + raise AssertionError("watermarker should not run for greedy requests")
132 +
133 + sampler = object.__new__(GPUWatermarkSampler)
134 + sampler.watermarker = StubWatermarker()
135 + sampler.watermarking = SimpleNamespace(
136 + np=np.array([True, True]), gpu=torch.tensor([True, True])
137 + )
138 + sampler.sampling_states = SimpleNamespace(
139 + temperature=SimpleNamespace(np=np.zeros(2), gpu=torch.zeros(2)),
140 + )
141 + expected = (torch.tensor([3, 4]), torch.zeros(2, 8))
142 + monkeypatch.setattr(
143 + "vllm.v1.watermarking.gpu_sampler.Sampler._sample_random",
144 + lambda *args, **kwargs: expected,
145 + )
146 +
147 + actual = sampler._sample_random(
148 + torch.zeros(2, 8),
149 + torch.tensor([0, 1]),
150 + np.array([0, 1]),
151 + torch.zeros(2, dtype=torch.int64),
152 + None,
153 + None,
154 + False,
155 + )
156 +
157 + assert actual is expected
vllm/config/__init__.py
+2
@@ -60,6 +60,7 @@ from vllm.config.vllm import (
60 get_layers_from_vllm_config,
61 set_current_vllm_config,
62 )
63 +from vllm.config.watermarking import WatermarkConfig
64 from vllm.config.weight_transfer import WeightTransferConfig
65
66 # __all__ should only contain classes and functions.
@@ -146,4 +147,5 @@ __all__ = [
147 "set_current_vllm_config",
148 "get_layers_from_vllm_config",
149 "WeightTransferConfig",
150 + "WatermarkConfig",
151 ]
vllm/config/vllm.py
+36
@@ -51,6 +51,7 @@ from .scheduler import SchedulerConfig
51 from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig
52 from .structured_outputs import StructuredOutputsConfig
53 from .utils import SupportsHash, config, replace
54 +from .watermarking import WatermarkConfig
55 from .weight_transfer import WeightTransferConfig
56
57 if TYPE_CHECKING:
@@ -381,6 +382,8 @@ class VllmConfig:
382 """LoRA configuration."""
383 speculative_config: SpeculativeConfig | None = None
384 """Speculative decoding configuration."""
385 + watermark_config: WatermarkConfig | None = None
386 + """Text watermarking configuration."""
387 diffusion_config: DiffusionConfig | None = None
388 """Diffusion LLM (dLLM) configuration."""
389
@@ -662,6 +665,14 @@ class VllmConfig:
665
666 @property
667 def use_v2_model_runner(self) -> bool:
668 + if getattr(self, "watermark_config", None) is not None:
669 + if envs.VLLM_USE_V2_MODEL_RUNNER is False:
670 + logger.info_once(
671 + "Watermarking requires Model Runner V2 and overrides "
672 + "VLLM_USE_V2_MODEL_RUNNER=0."
673 + )
674 + return True
675 +
676 use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
677 if use_v2_model_runner is not None:
678 return use_v2_model_runner
@@ -1091,6 +1102,30 @@ class VllmConfig:
1102 if not self.use_v2_model_runner:
1103 raise ValueError("trace replay requires Model Runner V2")
1104
1105 + def _check_watermarking_unsupported(
1106 + self,
1107 + *,
1108 + beam_search: bool = False,
1109 + custom_sampler: bool = False,
1110 + ) -> None:
1111 + watermark_config = getattr(self, "watermark_config", None)
1112 + if watermark_config is None:
1113 + return
1114 + if (
1115 + self.speculative_config is not None
1116 + and not watermark_config.supports_speculative_decoding
1117 + ):
1118 + raise ValueError(
1119 + f"The {watermark_config.algorithm} watermarking algorithm "
1120 + "does not support speculative decoding."
1121 + )
1122 + if beam_search:
1123 + raise ValueError("Beam search is not supported with watermarking.")
1124 + if custom_sampler:
1125 + raise ValueError(
1126 + "Model-specific custom samplers are not supported with watermarking."
1127 + )
1128 +
1129 def _resolve_and_verify_engram_config(self) -> None:
1130 """Resolve legacy offload settings and validate model and parallel configs."""
1131 if self.engram_config is None:
@@ -1124,6 +1159,7 @@ class VllmConfig:
1159 self.try_verify_and_update_config()
1160 self._resolve_and_verify_engram_config()
1161
1162 + self._check_watermarking_unsupported()
1163 # Models may have supplied their own DCP defaults above; anything still
1164 # unset falls back to the stock ones.
1165 self.parallel_config.set_dcp_defaults()
vllm/config/watermarking.py new
+49
@@ -0,0 +1,49 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Literal
5 +
6 +from pydantic import Field, model_validator
7 +from typing_extensions import Self
8 +
9 +from vllm.config.utils import config
10 +from vllm.logger import init_logger
11 +
12 +logger = init_logger(__name__)
13 +
14 +WatermarkingAlgorithm = Literal["gumbel"]
15 +WatermarkPRFName = Literal["philox"]
16 +
17 +_SPECULATIVE_DECODING_SUPPORT: dict[WatermarkingAlgorithm, bool] = {
18 + "gumbel": False,
19 +}
20 +
21 +
22 +@config
23 +class WatermarkConfig:
24 + """Configuration for text watermark generation."""
25 +
26 + key: int = Field(ge=0, repr=False, exclude=True)
27 + """Secret key used to watermark generated text."""
28 + algorithm: WatermarkingAlgorithm = "gumbel"
29 + """Algorithm used to watermark generated text."""
30 + context_width: int = Field(default=4, ge=1)
31 + """Number of prior output tokens used by the watermark PRF."""
32 + prf: WatermarkPRFName = "philox"
33 + """Pseudorandom function used by the watermarking algorithm."""
34 +
35 + @model_validator(mode="after")
36 + def validate_key(self) -> Self:
37 + if self.key > 2**64 - 1:
38 + raise ValueError("philox keys must fit in 64 bits")
39 + if self.algorithm == "gumbel":
40 + logger.warning_once(
41 + "Single-key Gumbel-max watermarking may increase the frequency of "
42 + "degenerate generations, including repetition loops.",
43 + scope="global",
44 + )
45 + return self
46 +
47 + @property
48 + def supports_speculative_decoding(self) -> bool:
49 + return _SPECULATIVE_DECODING_SUPPORT[self.algorithm]
vllm/engine/arg_utils.py
+14
@@ -106,6 +106,7 @@ from vllm.config.parallel import (
106 from vllm.config.scheduler import SchedulerPolicy
107 from vllm.config.utils import get_field
108 from vllm.config.vllm import OptimizationLevel, PerformanceMode
109 +from vllm.config.watermarking import WatermarkConfig
110 from vllm.logger import init_logger, suppress_logging
111 from vllm.platforms import CpuArchEnum, current_platform
112 from vllm.plugins import load_general_plugins
@@ -656,6 +657,7 @@ class EngineArgs:
657 spec_method: str | None = None
658 spec_model: str | None = None
659 spec_tokens: int | None = None
660 + watermark_config: dict[str, Any] | None = None
661 diffusion_config: dict[str, Any] | None = None
662
663 show_hidden_metrics_for_version: str | None = (
@@ -1680,6 +1682,8 @@ class EngineArgs:
1682 vllm_group.add_argument(
1683 "--spec-tokens", **speculative_kwargs["num_speculative_tokens"]
1684 )
1685 + vllm_kwargs["watermark_config"]["type"] = optional_type(json.loads)
1686 + vllm_group.add_argument("--watermark-config", **vllm_kwargs["watermark_config"])
1687 vllm_kwargs["diffusion_config"]["type"] = optional_type(json.loads)
1688 vllm_group.add_argument(
1689 "--diffusion-config", "-dc", **vllm_kwargs["diffusion_config"]
@@ -1986,6 +1990,14 @@ class EngineArgs:
1990 cfg = json.loads(cfg)
1991 return DiffusionConfig(**cfg)
1992
1993 + def create_watermark_config(self) -> WatermarkConfig | None:
1994 + if self.watermark_config is None:
1995 + return None
1996 + cfg = self.watermark_config
1997 + if isinstance(cfg, str):
1998 + cfg = json.loads(cfg)
1999 + return WatermarkConfig(**cfg)
2000 +
2001 def create_observability_config(self) -> ObservabilityConfig:
2002 return ObservabilityConfig(
2003 show_hidden_metrics_for_version=self.show_hidden_metrics_for_version,
@@ -2366,6 +2378,7 @@ class EngineArgs:
2378 target_parallel_config=parallel_config,
2379 )
2380 diffusion_config = self.create_diffusion_config()
2381 + watermark_config = self.create_watermark_config()
2382
2383 self._set_default_max_num_seqs_and_batched_tokens_args(
2384 usage_context,
@@ -2602,6 +2615,7 @@ class EngineArgs:
2615 kernel_config=kernel_config,
2616 lora_config=lora_config,
2617 speculative_config=speculative_config,
2618 + watermark_config=watermark_config,
2619 diffusion_config=diffusion_config,
2620 structured_outputs_config=self.structured_outputs_config,
2621 observability_config=observability_config,
vllm/entrypoints/generate/beam_search/offline.py
+1
@@ -82,6 +82,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
82 temperature = params.temperature
83 ignore_eos = params.ignore_eos
84 length_penalty = params.length_penalty
85 + self.llm_engine.vllm_config._check_watermarking_unsupported(beam_search=True)
86
87 tokenizer = self.renderer.get_tokenizer()
88 eos_token_id = tokenizer.eos_token_id
vllm/entrypoints/openai/chat_completion/protocol.py
+2
@@ -267,6 +267,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
267 top_k: int | None = None
268 min_p: float | None = None
269 repetition_penalty: float | None = None
270 + watermarking: bool = True
271 length_penalty: float = 1.0
272 stop_token_ids: list[int] | None = []
273 include_stop_str_in_output: bool = False
@@ -724,6 +725,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
725 frequency_penalty=self.frequency_penalty,
726 repetition_penalty=repetition_penalty,
727 temperature=temperature,
728 + watermarking=self.watermarking,
729 top_p=top_p,
730 top_k=top_k,
731 min_p=min_p,
vllm/entrypoints/openai/completion/protocol.py
+2
@@ -75,6 +75,7 @@ class CompletionRequest(OpenAIBaseModel):
75 top_k: int | None = None
76 min_p: float | None = None
77 repetition_penalty: float | None = None
78 + watermarking: bool = True
79 length_penalty: float = 1.0
80 stop_token_ids: list[int] | None = []
81 include_stop_str_in_output: bool = False
@@ -373,6 +374,7 @@ class CompletionRequest(OpenAIBaseModel):
374 frequency_penalty=self.frequency_penalty,
375 repetition_penalty=repetition_penalty,
376 temperature=temperature,
377 + watermarking=self.watermarking,
378 top_p=top_p,
379 top_k=top_k,
380 min_p=min_p,
vllm/entrypoints/openai/responses/protocol.py
+2
@@ -213,6 +213,7 @@ class ResponsesRequest(OpenAIBaseModel):
213 )
214
215 # --8<-- [start:responses-extra-params]
216 + watermarking: bool = True
217 request_id: str = Field(
218 default_factory=lambda: f"resp_{random_uuid()}",
219 description=(
@@ -436,6 +437,7 @@ class ResponsesRequest(OpenAIBaseModel):
437
438 return SamplingParams.from_optional(
439 temperature=temperature,
440 + watermarking=self.watermarking,
441 top_p=top_p,
442 top_k=top_k,
443 max_tokens=max_tokens,
vllm/entrypoints/serve/utils/api_utils.py
+1 -1
@@ -268,7 +268,7 @@ def jsonify_non_default_args(
268
269
270 # Fields whose values must never be logged verbatim.
271 -_SENSITIVE_ARG_FIELDS = frozenset({"api_key", "hf_token"})
271 +_SENSITIVE_ARG_FIELDS = frozenset({"api_key", "hf_token", "watermark_config"})
272
273
274 def redact_sensitive_args(args: dict[str, Any]) -> dict[str, Any]:
vllm/sampling_params.py
+5
@@ -253,6 +253,8 @@ class SamplingParams(
253 """Controls the randomness of the sampling. Lower values make the model
254 more deterministic, while higher values make the model more random. Zero
255 means greedy sampling."""
256 + watermarking: bool = True
257 + """Whether to apply the engine's configured watermark to this request."""
258 top_p: float = 1.0
259 """Controls the cumulative probability of the top tokens to consider. Must
260 be in (0, 1]. Set to 1 to consider all tokens."""
@@ -385,6 +387,7 @@ class SamplingParams(
387 frequency_penalty: float | None = 0.0,
388 repetition_penalty: float | None = 1.0,
389 temperature: float | None = 1.0,
390 + watermarking: bool = True,
391 top_p: float | None = 1.0,
392 top_k: int = 0,
393 min_p: float = 0.0,
@@ -450,6 +453,7 @@ class SamplingParams(
453 if repetition_penalty is None
454 else repetition_penalty,
455 temperature=1.0 if temperature is None else temperature,
456 + watermarking=watermarking,
457 top_p=1.0 if top_p is None else top_p,
458 top_k=top_k,
459 min_p=min_p,
@@ -1253,6 +1257,7 @@ class SamplingParams(
1257 f"frequency_penalty={self.frequency_penalty}, "
1258 f"repetition_penalty={self.repetition_penalty}, "
1259 f"temperature={self.temperature}, "
1260 + f"watermarking={self.watermarking}, "
1261 f"top_p={self.top_p}, "
1262 f"top_k={self.top_k}, "
1263 f"min_p={self.min_p}, "
vllm/v1/watermarking/__init__.py new
+28
@@ -0,0 +1,28 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.v1.watermarking.detector import WatermarkDetection, WatermarkDetector
5 +from vllm.v1.watermarking.factory import create_watermarker
6 +from vllm.v1.watermarking.gumbel import (
7 + GumbelWatermarkDetector,
8 + GumbelWatermarker,
9 +)
10 +from vllm.v1.watermarking.prfs import (
11 + PhiloxPRF,
12 + WatermarkPRF,
13 + create_prf,
14 +)
15 +from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample
16 +
17 +__all__ = [
18 + "GumbelWatermarkDetector",
19 + "GumbelWatermarker",
20 + "PhiloxPRF",
21 + "WatermarkDetection",
22 + "WatermarkDetector",
23 + "WatermarkPRF",
24 + "Watermarker",
25 + "WatermarkSample",
26 + "create_watermarker",
27 + "create_prf",
28 +]
vllm/v1/watermarking/detector.py new
+100
@@ -0,0 +1,100 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from abc import ABC, abstractmethod
5 +from dataclasses import dataclass
6 +from typing import TypeAlias
7 +
8 +import torch
9 +
10 +Context: TypeAlias = tuple[int, ...]
11 +
12 +
13 +@dataclass(frozen=True)
14 +class WatermarkDetection:
15 + score: float
16 + p_value: float
17 + num_scored_tokens: int
18 + is_watermarked: bool
19 +
20 +
21 +class WatermarkDetector(ABC):
22 + def __init__(
23 + self,
24 + context_width: int,
25 + p_value_threshold: float,
26 + deduplicate_contexts: bool = False,
27 + ) -> None:
28 + if context_width < 1:
29 + raise ValueError("context_width must be positive")
30 + if not 0 < p_value_threshold < 1:
31 + raise ValueError("p_value_threshold must be between 0 and 1")
32 + self.context_width = context_width
33 + self.p_value_threshold = p_value_threshold
34 + self.deduplicate_contexts = deduplicate_contexts
35 +
36 + def detect(self, token_ids: list[int]) -> WatermarkDetection:
37 + contexts, targets = self._prepare_inputs(token_ids)
38 + num_scored_tokens = len(targets)
39 + if num_scored_tokens == 0:
40 + return WatermarkDetection(0.0, 1.0, 0, False)
41 +
42 + token_scores = self._score_tokens(contexts, targets)
43 + score = self._aggregate_scores(token_scores)
44 + p_value = self._get_p_value(score, num_scored_tokens)
45 + return WatermarkDetection(
46 + score=score,
47 + p_value=p_value,
48 + num_scored_tokens=num_scored_tokens,
49 + is_watermarked=p_value <= self.p_value_threshold,
50 + )
51 +
52 + def _prepare_inputs(
53 + self, token_ids: list[int]
54 + ) -> tuple[torch.Tensor, torch.Tensor]:
55 + contexts: list[Context] = []
56 + targets = []
57 + seen_contexts: set[Context] = set()
58 + prefix = [-1] * self.context_width
59 + for token_id in token_ids:
60 + context = tuple(prefix[-self.context_width :])
61 + # Reusing a context reuses one keyed random vector, so the resulting
62 + # token scores are not independent evidence.
63 + if not self.deduplicate_contexts or not _is_repeated_context(
64 + context, seen_contexts
65 + ):
66 + contexts.append(context)
67 + targets.append(token_id)
68 + prefix.append(token_id)
69 +
70 + if not contexts:
71 + return (
72 + torch.empty((0, self.context_width), dtype=torch.int64),
73 + torch.empty(0, dtype=torch.int64),
74 + )
75 +
76 + return (
77 + torch.tensor(contexts, dtype=torch.int64).reshape(-1, self.context_width),
78 + torch.tensor(targets, dtype=torch.int64),
79 + )
80 +
81 + @abstractmethod
82 + def _score_tokens(
83 + self, contexts: torch.Tensor, targets: torch.Tensor
84 + ) -> torch.Tensor:
85 + raise NotImplementedError
86 +
87 + @abstractmethod
88 + def _aggregate_scores(self, token_scores: torch.Tensor) -> float:
89 + raise NotImplementedError
90 +
91 + @abstractmethod
92 + def _get_p_value(self, score: float, num_scored_tokens: int) -> float:
93 + raise NotImplementedError
94 +
95 +
96 +def _is_repeated_context(context: Context, seen_contexts: set[Context]) -> bool:
97 + if context in seen_contexts:
98 + return True
99 + seen_contexts.add(context)
100 + return False
vllm/v1/watermarking/factory.py new
+24
@@ -0,0 +1,24 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Callable
5 +
6 +from vllm.config.watermarking import WatermarkConfig
7 +from vllm.v1.watermarking.gumbel import GumbelWatermarker
8 +from vllm.v1.watermarking.watermarker import Watermarker
9 +
10 +
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 +
20 +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)
vllm/v1/watermarking/gpu_sampler.py new
+131
@@ -0,0 +1,131 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import numpy as np
5 +import torch
6 +
7 +from vllm.logger import init_logger
8 +from vllm.sampling_params import SamplingParams
9 +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
10 +from vllm.v1.watermarking.watermarker import Watermarker
11 +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
12 +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
13 +from vllm.v1.worker.gpu.sample.sampler import Sampler
14 +
15 +logger = init_logger(__name__)
16 +
17 +
18 +class GPUWatermarkSampler(Sampler):
19 + def __init__(
20 + self,
21 + watermarker: Watermarker,
22 + *args,
23 + **kwargs,
24 + ) -> None:
25 + super().__init__(*args, **kwargs)
26 + self.watermarker = watermarker
27 + self.watermarking = UvaBackedTensor(
28 + self.sampling_states.max_num_reqs, dtype=torch.bool
29 + )
30 + self.watermarking.np.fill(True)
31 + self.watermarking.copy_to_uva()
32 +
33 + def add_request(
34 + self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
35 + ) -> None:
36 + super().add_request(req_idx, prompt_len, sampling_params)
37 + self.watermarking.np[req_idx] = sampling_params.watermarking
38 + if sampling_params.watermarking and sampling_params.temperature == 0:
39 + logger.warning_once(
40 + "Watermarking is enabled, but greedy decoding "
41 + "(temperature=0) cannot be watermarked. This request will use "
42 + "ordinary greedy sampling."
43 + )
44 +
45 + def apply_staged_writes(self) -> None:
46 + super().apply_staged_writes()
47 + self.watermarking.copy_to_uva()
48 +
49 + def _sample_random(
50 + self,
51 + processed_logits: torch.Tensor,
52 + expanded_idx_mapping: torch.Tensor,
53 + idx_mapping_np: np.ndarray,
54 + pos: torch.Tensor,
55 + top_k: torch.Tensor | None,
56 + top_p: torch.Tensor | None,
57 + use_flashinfer: bool,
58 + ) -> tuple[torch.Tensor, torch.Tensor]:
59 + enabled = self.watermarking.np[idx_mapping_np] & (
60 + self.sampling_states.temperature.np[idx_mapping_np] != 0
61 + )
62 + if not np.any(enabled):
63 + return super()._sample_random(
64 + processed_logits,
65 + expanded_idx_mapping,
66 + idx_mapping_np,
67 + pos,
68 + top_k,
69 + top_p,
70 + use_flashinfer,
71 + )
72 +
73 + processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
74 + contexts = self._get_contexts(expanded_idx_mapping)
75 +
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,
86 + )
87 +
88 + output = self.watermarker.sample(
89 + processed_logits,
90 + contexts,
91 + random_sample,
92 + )
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
106 + sampled = torch.where(
107 + temperatures == 0,
108 + processed_logits.argmax(dim=-1),
109 + sampled,
110 + )
111 + return sampled, output_logits
112 +
113 + def _get_contexts(self, expanded_idx_mapping: torch.Tensor) -> torch.Tensor:
114 + context_width = self.watermarker.context_width
115 + req_indices = expanded_idx_mapping.to(torch.int64)
116 + valid_reqs = req_indices >= 0
117 + safe_req_indices = req_indices.clamp_min(0)
118 + total_lens = self.req_states.total_len.gpu[safe_req_indices].to(torch.int64)
119 + prompt_lens = self.req_states.prompt_len.gpu[safe_req_indices].to(torch.int64)
120 + offsets = torch.arange(
121 + -context_width, 0, dtype=torch.int64, device=req_indices.device
122 + )
123 + positions = total_lens.unsqueeze(-1) + offsets
124 + valid_positions = valid_reqs.unsqueeze(-1) & (
125 + positions >= prompt_lens.unsqueeze(-1)
126 + )
127 + positions = positions.clamp_min(0)
128 + contexts = self.req_states.all_token_ids.gpu[
129 + safe_req_indices.unsqueeze(-1), positions
130 + ]
131 + return torch.where(valid_positions, contexts, -1)
vllm/v1/watermarking/gumbel.py new
+111
@@ -0,0 +1,111 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Gumbel-max watermark generation and detection primitives."""
5 +
6 +import math
7 +import warnings
8 +
9 +import torch
10 +
11 +from vllm.config.watermarking import WatermarkPRFName
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 + Watermarker,
19 + WatermarkSample,
20 +)
21 +
22 +
23 +def _gamma_survival_integer_shape(score: float, shape: int) -> float:
24 + """Return the survival probability of a Gamma(shape, 1) variable."""
25 + if shape <= 0 or score <= 0:
26 + return 1.0
27 + log_score = math.log(score)
28 + log_term = 0.0
29 + max_log_term = 0.0
30 + log_terms = [log_term]
31 + for index in range(1, shape):
32 + log_term += log_score - math.log(index)
33 + log_terms.append(log_term)
34 + max_log_term = max(max_log_term, log_term)
35 + log_sum = max_log_term + math.log(
36 + sum(math.exp(term - max_log_term) for term in log_terms)
37 + )
38 + return min(1.0, math.exp(min(0.0, -score + log_sum)))
39 +
40 +
41 +class GumbelWatermarker(Watermarker):
42 + def __init__(
43 + self,
44 + key: int,
45 + context_width: int = 4,
46 + prf: WatermarkPRF | WatermarkPRFName = "philox",
47 + ) -> None:
48 + prf = create_prf(prf, key) if isinstance(prf, str) else prf
49 + _validate_context_width(context_width)
50 + self.prf = prf
51 + self._context_width = context_width
52 +
53 + @property
54 + def context_width(self) -> int:
55 + return self._context_width
56 +
57 + def sample(
58 + self,
59 + logits: torch.Tensor,
60 + 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
65 +
66 + return WatermarkSample(
67 + philox_gumbel_sample(logits, contexts, self.prf.key), logits
68 + )
69 + vocabulary = torch.arange(logits.shape[-1], device=logits.device)
70 + uniforms = self.prf.uniform(contexts, vocabulary)
71 + uniforms = uniforms.clamp_min(torch.finfo(torch.float32).tiny)
72 + noise = -torch.log(-torch.log(uniforms))
73 + return WatermarkSample(torch.argmax(logits + noise, dim=-1), logits)
74 +
75 +
76 +class GumbelWatermarkDetector(WatermarkDetector):
77 + def __init__(
78 + self,
79 + key: int,
80 + context_width: int = 4,
81 + p_value_threshold: float = 0.01,
82 + prf: WatermarkPRF | WatermarkPRFName = "philox",
83 + deduplicate_contexts: bool = True,
84 + ) -> None:
85 + prf = create_prf(prf, key) if isinstance(prf, str) else prf
86 + _validate_context_width(context_width)
87 + super().__init__(context_width, p_value_threshold, deduplicate_contexts)
88 + self.prf = prf
89 +
90 + def _score_tokens(
91 + self, contexts: torch.Tensor, targets: torch.Tensor
92 + ) -> torch.Tensor:
93 + uniforms = self.prf.uniform(contexts, targets.unsqueeze(-1)).squeeze(-1)
94 + return -torch.log1p(-uniforms.to(torch.float64))
95 +
96 + def _get_p_value(self, score: float, num_scored_tokens: int) -> float:
97 + return _gamma_survival_integer_shape(score, num_scored_tokens)
98 +
99 + def _aggregate_scores(self, token_scores: torch.Tensor) -> float:
100 + return token_scores.sum().item()
101 +
102 +
103 +def _validate_context_width(context_width: int) -> None:
104 + if context_width < 1:
105 + raise ValueError("context_width must be positive")
106 + if context_width > 16:
107 + warnings.warn(
108 + "context_width values greater than 16 reduce robustness to edits because "
109 + "each changed token affects more subsequent watermark contexts",
110 + stacklevel=3,
111 + )
vllm/v1/watermarking/prfs/__init__.py new
+19
@@ -0,0 +1,19 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.config.watermarking import WatermarkPRFName
5 +from vllm.v1.watermarking.prfs.base import WatermarkPRF
6 +from vllm.v1.watermarking.prfs.philox import PhiloxPRF
7 +
8 +
9 +def create_prf(name: WatermarkPRFName, key: int) -> WatermarkPRF:
10 + if name == "philox":
11 + return PhiloxPRF(key)
12 + raise ValueError(f"Unknown watermark PRF: {name}")
13 +
14 +
15 +__all__ = [
16 + "PhiloxPRF",
17 + "WatermarkPRF",
18 + "create_prf",
19 +]
vllm/v1/watermarking/prfs/base.py new
+19
@@ -0,0 +1,19 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from abc import ABC, abstractmethod
5 +
6 +import torch
7 +
8 +_UINT32_MASK = 2**32 - 1
9 +
10 +
11 +class WatermarkPRF(ABC):
12 + @abstractmethod
13 + def uniform(self, contexts: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
14 + raise NotImplementedError
15 +
16 +
17 +def uint32_to_uniform(values: torch.Tensor) -> torch.Tensor:
18 + mantissas = (values & _UINT32_MASK) >> 8
19 + return ((mantissas.to(torch.float64) + 1.0) / (2**24 + 1)).to(torch.float32)
vllm/v1/watermarking/prfs/philox.py new
+145
@@ -0,0 +1,145 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Philox4x32-10 as defined by https://doi.org/10.1145/2063384.2063405."""
5 +
6 +import torch
7 +
8 +from vllm.v1.watermarking.prfs.base import WatermarkPRF, uint32_to_uniform
9 +
10 +_UINT32_MASK = 2**32 - 1
11 +_PHILOX_M0 = 0xD2511F53
12 +_PHILOX_M1 = 0xCD9E8D57
13 +_PHILOX_W0 = 0x9E3779B9
14 +_PHILOX_W1 = 0xBB67AE85
15 +_CONTEXT_DOMAIN = 0x574D4358
16 +_TOKEN_DOMAIN = 0x574D544B
17 +
18 +
19 +def _mulhilo32(
20 + multiplier: int, values: torch.Tensor
21 +) -> tuple[torch.Tensor, torch.Tensor]:
22 + value_low = values & 0xFFFF
23 + value_high = values >> 16
24 + multiplier_low = multiplier & 0xFFFF
25 + multiplier_high = multiplier >> 16
26 +
27 + low_product = value_low * multiplier_low
28 + cross_product = value_high * multiplier_low + value_low * multiplier_high
29 + low_with_carry = low_product + ((cross_product & 0xFFFF) << 16)
30 + low = low_with_carry & _UINT32_MASK
31 + high = (
32 + value_high * multiplier_high + (cross_product >> 16) + (low_with_carry >> 32)
33 + ) & _UINT32_MASK
34 + return high, low
35 +
36 +
37 +def _philox4x32_10(
38 + counter: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
39 + key: tuple[torch.Tensor, torch.Tensor],
40 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
41 + counter_0, counter_1, counter_2, counter_3 = counter
42 + key_0, key_1 = key
43 + for round_index in range(10):
44 + high_0, low_0 = _mulhilo32(_PHILOX_M0, counter_0)
45 + high_1, low_1 = _mulhilo32(_PHILOX_M1, counter_2)
46 + counter_0, counter_1, counter_2, counter_3 = (
47 + (high_1 ^ counter_1 ^ key_0) & _UINT32_MASK,
48 + low_1,
49 + (high_0 ^ counter_3 ^ key_1) & _UINT32_MASK,
50 + low_0,
51 + )
52 + if round_index != 9:
53 + key_0 = (key_0 + _PHILOX_W0) & _UINT32_MASK
54 + key_1 = (key_1 + _PHILOX_W1) & _UINT32_MASK
55 + return counter_0, counter_1, counter_2, counter_3
56 +
57 +
58 +def _philox_uniform(
59 + key_0_value: int,
60 + key_1_value: int,
61 + contexts: torch.Tensor,
62 + token_ids: torch.Tensor,
63 +) -> torch.Tensor:
64 + contexts = contexts.to(torch.int64) & _UINT32_MASK
65 + prefix_shape = contexts.shape[:-1]
66 + device = contexts.device
67 + key_0 = torch.full(prefix_shape, key_0_value, dtype=torch.int64, device=device)
68 + key_1 = torch.full(prefix_shape, key_1_value, dtype=torch.int64, device=device)
69 + state = (
70 + torch.full(prefix_shape, _CONTEXT_DOMAIN, dtype=torch.int64, device=device),
71 + torch.full(
72 + prefix_shape,
73 + contexts.shape[-1],
74 + dtype=torch.int64,
75 + device=device,
76 + ),
77 + torch.zeros(prefix_shape, dtype=torch.int64, device=device),
78 + torch.zeros(prefix_shape, dtype=torch.int64, device=device),
79 + )
80 +
81 + for offset in range(0, contexts.shape[-1], 4):
82 + block = []
83 + for index in range(4):
84 + position = offset + index
85 + if position < contexts.shape[-1]:
86 + block.append(contexts[..., position])
87 + else:
88 + block.append(
89 + torch.full(
90 + prefix_shape,
91 + _UINT32_MASK - index,
92 + dtype=torch.int64,
93 + device=device,
94 + )
95 + )
96 + counter = tuple(
97 + (state[index] ^ block[index]) & _UINT32_MASK for index in range(4)
98 + )
99 + state = _philox4x32_10(counter, ((key_0 ^ offset) & _UINT32_MASK, key_1))
100 +
101 + output_shape = torch.broadcast_shapes(prefix_shape + (1,), token_ids.shape)
102 + tokens = torch.broadcast_to(token_ids.to(device=device), output_shape)
103 + state = tuple(
104 + torch.broadcast_to(word.unsqueeze(-1), output_shape) for word in state
105 + )
106 + candidate_key_0 = torch.broadcast_to(key_0.unsqueeze(-1), output_shape)
107 + candidate_key_1 = torch.broadcast_to(key_1.unsqueeze(-1), output_shape)
108 + token_words = tokens.to(torch.int64) & _UINT32_MASK
109 + outputs = _philox4x32_10(
110 + (
111 + token_words >> 2,
112 + state[0],
113 + state[1],
114 + state[2],
115 + ),
116 + (
117 + (candidate_key_0 ^ state[3]) & _UINT32_MASK,
118 + (candidate_key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK,
119 + ),
120 + )
121 + word_index = token_words & 3
122 + output = outputs[0]
123 + for index in range(1, 4):
124 + output = torch.where(word_index == index, outputs[index], output)
125 + return uint32_to_uniform(output)
126 +
127 +
128 +_compiled_philox_uniform = torch.compile(_philox_uniform, fullgraph=True, dynamic=True)
129 +
130 +
131 +class PhiloxPRF(WatermarkPRF):
132 + """Accelerator-friendly Philox4x32-10 watermark PRF."""
133 +
134 + version = "philox4x32-10-v1"
135 +
136 + def __init__(self, key: int) -> None:
137 + if not 0 <= key <= 2**64 - 1:
138 + raise ValueError("Philox keys must fit in 64 bits")
139 + self.key = key
140 +
141 + def uniform(self, contexts: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
142 + key_words = self.key & _UINT32_MASK, self.key >> 32
143 + if contexts.device.type == "cuda":
144 + return _compiled_philox_uniform(*key_words, contexts, token_ids)
145 + return _philox_uniform(*key_words, contexts, token_ids)
vllm/v1/watermarking/watermarker.py new
+33
@@ -0,0 +1,33 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from abc import ABC, abstractmethod
5 +from collections.abc import Callable
6 +from dataclasses import dataclass
7 +from typing import TypeAlias
8 +
9 +import torch
10 +
11 +RandomSampler: TypeAlias = Callable[[torch.Tensor], torch.Tensor]
12 +
13 +
14 +@dataclass(frozen=True)
15 +class WatermarkSample:
16 + token_ids: torch.Tensor
17 + logits: torch.Tensor
18 +
19 +
20 +class Watermarker(ABC):
21 + @property
22 + @abstractmethod
23 + def context_width(self) -> int:
24 + raise NotImplementedError
25 +
26 + @abstractmethod
27 + def sample(
28 + self,
29 + logits: torch.Tensor,
30 + contexts: torch.Tensor,
31 + random_sample: RandomSampler,
32 + ) -> WatermarkSample:
33 + raise NotImplementedError
vllm/v1/worker/gpu/model_runner.py
+20 -12
@@ -78,6 +78,8 @@ from vllm.v1.outputs import (
78 ModelRunnerOutput,
79 RoutedExpertsTensors,
80 )
81 +from vllm.v1.watermarking import create_watermarker
82 +from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
83 from vllm.v1.worker.block_table import get_block_table_width
84 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
85 from vllm.v1.worker.gpu import pcp_manager as pcp
@@ -452,21 +454,27 @@ class GPUModelRunner(LoRAModelRunnerMixin):
454
455 # Initialize samplers. Model states may override via custom_sampler().
456 if self.is_last_pp_rank and not self.is_pooling_model:
455 - self.sampler = Sampler(
456 - max_num_reqs=self.max_num_reqs,
457 - vocab_size=self.vocab_size,
458 - device=self.device,
459 - req_states=self.req_states,
460 - logprobs_mode=self.model_config.logprobs_mode,
461 - num_speculative_tokens=self.decode_query_len,
462 - use_fp64_gumbel=self.model_config.use_fp64_gumbel,
463 - enable_trace_replay=self.model_config.enable_trace_replay,
464 - reasoning_config=self.vllm_config.reasoning_config,
465 - return_sampling_mask=self.model_config.return_sampling_mask,
466 - )
457 + sampler_kwargs: dict[str, Any] = {
458 + "max_num_reqs": self.max_num_reqs,
459 + "vocab_size": self.vocab_size,
460 + "device": self.device,
461 + "req_states": self.req_states,
462 + "logprobs_mode": self.model_config.logprobs_mode,
463 + "num_speculative_tokens": self.decode_query_len,
464 + "use_fp64_gumbel": self.model_config.use_fp64_gumbel,
465 + "enable_trace_replay": self.model_config.enable_trace_replay,
466 + "reasoning_config": self.vllm_config.reasoning_config,
467 + "return_sampling_mask": self.model_config.return_sampling_mask,
468 + }
469 + if self.vllm_config.watermark_config is None:
470 + self.sampler = Sampler(**sampler_kwargs)
471 + else:
472 + watermarker = create_watermarker(self.vllm_config.watermark_config)
473 + self.sampler = GPUWatermarkSampler(watermarker, **sampler_kwargs)
474 custom = self.model_state.custom_sampler(self.sampler)
475
476 if custom:
477 + self.vllm_config._check_watermarking_unsupported(custom_sampler=True)
478 self.sampler, self.rejection_sampler = custom
479 elif self.speculative_config is not None:
480 self.rejection_sampler = RejectionSampler(
vllm/v1/worker/gpu/sample/sampler.py
+20 -1
@@ -307,7 +307,26 @@ class Sampler:
307 or self.sampling_states.any_explicit_seed(idx_mapping_np)
308 )
309
310 - # Sample the next token.
310 + return self._sample_random(
311 + processed_logits,
312 + expanded_idx_mapping,
313 + idx_mapping_np,
314 + pos,
315 + top_k,
316 + top_p,
317 + use_flashinfer,
318 + )
319 +
320 + def _sample_random(
321 + self,
322 + processed_logits: torch.Tensor,
323 + expanded_idx_mapping: torch.Tensor,
324 + idx_mapping_np: np.ndarray,
325 + pos: torch.Tensor,
326 + top_k: torch.Tensor | None,
327 + top_p: torch.Tensor | None,
328 + use_flashinfer: bool,
329 + ) -> tuple[torch.Tensor, torch.Tensor]:
330 if use_flashinfer:
331 sampled = flashinfer_sample(processed_logits, top_k, top_p).to(torch.int64)
332 else:
vllm/v1/worker/gpu/sample/watermark.py new
+214
@@ -0,0 +1,214 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.triton_utils import HAS_TRITON, tl, triton
7 +from vllm.v1.watermarking.prfs.philox import (
8 + _CONTEXT_DOMAIN as _CONTEXT_DOMAIN_VALUE,
9 +)
10 +from vllm.v1.watermarking.prfs.philox import (
11 + _PHILOX_M0 as _PHILOX_M0_VALUE,
12 +)
13 +from vllm.v1.watermarking.prfs.philox import (
14 + _PHILOX_M1 as _PHILOX_M1_VALUE,
15 +)
16 +from vllm.v1.watermarking.prfs.philox import (
17 + _PHILOX_W0 as _PHILOX_W0_VALUE,
18 +)
19 +from vllm.v1.watermarking.prfs.philox import (
20 + _PHILOX_W1 as _PHILOX_W1_VALUE,
21 +)
22 +from vllm.v1.watermarking.prfs.philox import (
23 + _TOKEN_DOMAIN as _TOKEN_DOMAIN_VALUE,
24 +)
25 +from vllm.v1.watermarking.prfs.philox import (
26 + _UINT32_MASK as _UINT32_MASK_VALUE,
27 +)
28 +
29 +if HAS_TRITON:
30 + from triton.language import math as tl_math
31 +else:
32 + tl_math = tl
33 +
34 +_UINT32_MASK = tl.constexpr(_UINT32_MASK_VALUE) if HAS_TRITON else _UINT32_MASK_VALUE
35 +_PHILOX_M0 = tl.constexpr(_PHILOX_M0_VALUE) if HAS_TRITON else _PHILOX_M0_VALUE
36 +_PHILOX_M1 = tl.constexpr(_PHILOX_M1_VALUE) if HAS_TRITON else _PHILOX_M1_VALUE
37 +_PHILOX_W0 = tl.constexpr(_PHILOX_W0_VALUE) if HAS_TRITON else _PHILOX_W0_VALUE
38 +_PHILOX_W1 = tl.constexpr(_PHILOX_W1_VALUE) if HAS_TRITON else _PHILOX_W1_VALUE
39 +_CONTEXT_DOMAIN = (
40 + tl.constexpr(_CONTEXT_DOMAIN_VALUE) if HAS_TRITON else _CONTEXT_DOMAIN_VALUE
41 +)
42 +_TOKEN_DOMAIN = tl.constexpr(_TOKEN_DOMAIN_VALUE) if HAS_TRITON else _TOKEN_DOMAIN_VALUE
43 +
44 +
45 +@triton.jit
46 +def _mulhilo32(multiplier: tl.constexpr, value):
47 + high = tl_math.umulhi(multiplier, value)
48 + low = tl.mul(multiplier, value, sanitize_overflow=False)
49 + return high, low
50 +
51 +
52 +@triton.jit
53 +def _philox4x32_10(counter_0, counter_1, counter_2, counter_3, key_0, key_1):
54 + for round_index in range(10):
55 + high_0, low_0 = _mulhilo32(_PHILOX_M0, counter_0)
56 + high_1, low_1 = _mulhilo32(_PHILOX_M1, counter_2)
57 + counter_0, counter_1, counter_2, counter_3 = (
58 + (high_1 ^ counter_1 ^ key_0) & _UINT32_MASK,
59 + low_1,
60 + (high_0 ^ counter_3 ^ key_1) & _UINT32_MASK,
61 + low_0,
62 + )
63 + if round_index != 9:
64 + key_0 = (key_0 + _PHILOX_W0) & _UINT32_MASK
65 + key_1 = (key_1 + _PHILOX_W1) & _UINT32_MASK
66 + return counter_0, counter_1, counter_2, counter_3
67 +
68 +
69 +@triton.jit
70 +def _uint32_to_uniform(value):
71 + mantissa = (value & _UINT32_MASK) >> 8
72 + scaled = (mantissa + 1).to(tl.float32) * 2**-24
73 + return (scaled.to(tl.uint32, bitcast=True) - 1).to(tl.float32, bitcast=True)
74 +
75 +
76 +@triton.jit
77 +def _gumbel_value(logits_ptr, output, mask):
78 + logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32)
79 + logits = tl.where(logits != logits, float("-inf"), logits)
80 + uniform = _uint32_to_uniform(output)
81 + return logits - tl.log(-tl.log(uniform))
82 +
83 +
84 +@triton.jit
85 +def _select_max(best_value, best_token, candidate_value, candidate):
86 + take_candidate = candidate_value > best_value
87 + return (
88 + tl.where(take_candidate, candidate_value, best_value),
89 + tl.where(take_candidate, candidate, best_token),
90 + )
91 +
92 +
93 +@triton.jit(do_not_specialize=["key_0_value", "key_1_value"])
94 +def _philox_gumbel_kernel(
95 + local_argmax_ptr,
96 + local_argmax_stride,
97 + local_max_ptr,
98 + local_max_stride,
99 + logits_ptr,
100 + logits_stride,
101 + contexts_ptr,
102 + context_stride,
103 + key_0_value,
104 + key_1_value,
105 + vocab_size,
106 + CONTEXT_WIDTH: tl.constexpr,
107 + BLOCK_SIZE: tl.constexpr,
108 +):
109 + row = tl.program_id(0).to(tl.int64)
110 + block_index = tl.program_id(1)
111 + groups = block_index * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4)
112 +
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
124 + )
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
130 + )
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
136 + )
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 + )
147 +
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,
157 + )
158 + candidate_0 = groups * 4
159 + candidate_1 = candidate_0 + 1
160 + candidate_2 = candidate_0 + 2
161 + candidate_3 = candidate_0 + 3
162 + logits_row = logits_ptr + row * logits_stride
163 + value_0 = _gumbel_value(
164 + logits_row + candidate_0, output_0, candidate_0 < vocab_size
165 + )
166 + value_1 = _gumbel_value(
167 + logits_row + candidate_1, output_1, candidate_1 < vocab_size
168 + )
169 + value_2 = _gumbel_value(
170 + logits_row + candidate_2, output_2, candidate_2 < vocab_size
171 + )
172 + value_3 = _gumbel_value(
173 + logits_row + candidate_3, output_3, candidate_3 < vocab_size
174 + )
175 + best_value = value_0
176 + best_token = candidate_0
177 + best_value, best_token = _select_max(best_value, best_token, value_1, candidate_1)
178 + best_value, best_token = _select_max(best_value, best_token, value_2, candidate_2)
179 + best_value, best_token = _select_max(best_value, best_token, value_3, candidate_3)
180 + value = tl.max(best_value, axis=0)
181 + token_id = tl.min(tl.where(best_value == value, best_token, vocab_size), axis=0)
182 + tl.store(local_argmax_ptr + row * local_argmax_stride + block_index, token_id)
183 + tl.store(local_max_ptr + row * local_max_stride + block_index, value)
184 +
185 +
186 +def philox_gumbel_sample(
187 + logits: torch.Tensor, contexts: torch.Tensor, key: int
188 +) -> torch.Tensor:
189 + if logits.stride(-1) != 1:
190 + logits = logits.contiguous()
191 + if contexts.stride(-1) != 1:
192 + contexts = contexts.contiguous()
193 + num_tokens, vocab_size = logits.shape
194 + block_size = 1024
195 + num_blocks = triton.cdiv(vocab_size, block_size)
196 + 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)
198 + _philox_gumbel_kernel[(num_tokens, num_blocks)](
199 + local_argmax,
200 + local_argmax.stride(0),
201 + local_max,
202 + local_max.stride(0),
203 + logits,
204 + logits.stride(0),
205 + contexts,
206 + contexts.stride(0),
207 + key & _UINT32_MASK_VALUE,
208 + key >> 32,
209 + vocab_size,
210 + CONTEXT_WIDTH=contexts.shape[-1],
211 + BLOCK_SIZE=block_size,
212 + )
213 + max_block_index = local_max.argmax(dim=-1, keepdim=True)
214 + return local_argmax.gather(dim=-1, index=max_block_index).view(-1)