diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md new file mode 100644 index 0000000000..7c601e13aa --- /dev/null +++ b/docs/features/watermarking.md @@ -0,0 +1,147 @@ +# Text watermarking + +Text watermarking embeds a statistical signal in generated token choices. A +detector that knows the generation parameters can test for this signal without +access to the model weights. + +The design and scope are discussed in +[RFC #53916](https://github.com/vllm-project/vllm/issues/53916). + +## Configuration + +Enable watermarking by configuring an algorithm and secret key at engine +startup: + +```bash +vllm serve MODEL \ + --watermark-config '{"algorithm":"gumbel","key":42}' +``` + +Watermarking is disabled when `--watermark-config` is omitted. Gumbel-max (see +`gumbel`) is the default algorithm within an enabled `WatermarkConfig`. When +watermarking is configured, it is enabled for requests by default. + +Requests can opt out without changing the engine-level algorithm or key: + +```python +from vllm import SamplingParams + +sampling_params = SamplingParams(watermarking=False) +``` + +The OpenAI-compatible APIs accept the same `watermarking: false` request field. +Deployments that require watermarking must restrict this field to trusted +callers, or strip and validate it at the ingress boundary, so untrusted clients +cannot opt out. + +`context_width` controls how many prior output tokens seed each watermark +decision and defaults to 4. Larger values make the watermark less robust to +edits because an insertion, deletion, or substitution changes more subsequent +contexts. Values above 16 are allowed but emit a warning. + +## Architecture + +`WatermarkConfig` selects an algorithm and PRF. Model Runner V2 constructs the +corresponding `Watermarker`, and `GPUWatermarkSampler` invokes it for the final +stochastic token selection after temperature, min-p, top-k, and top-p are +applied. A watermarker can either select a token directly or transform logits +and delegate to vLLM's random sampler. + +Detection is separate from generation. vLLM provides detector primitives for +the reference algorithms. `WatermarkDetector` consumes token IDs, so callers +remain responsible for using the tokenizer and watermark profile that match +generation. + +## Algorithms + +### Gumbel-max + +Gumbel-max derives a deterministic pseudorandom value from the key, prior +generated-token context, and every candidate token, then uses the resulting +Gumbel noise for categorical sampling. See +[Aaronson's original presentation](https://simons.berkeley.edu/sites/default/files/2024-10/LLM24-2%20Slides%20-%20Scott%20Aaronson.pdf). + +Gumbel-max requires stochastic sampling. Greedy requests (`temperature=0`) +bypass watermarking and emit a warning once per worker. + +Single-key Gumbel-max watermarking may increase degenerate generations. In +particular, reusing the same keyed random vector when a context repeats can +reinforce repetition loops that continue until the maximum token limit. + +### SynthID-Text + +[SynthID-Text](https://www.nature.com/articles/s41586-024-08025-4) is planned but +not currently implemented. + +## Pseudorandom functions + +A watermark PRF turns the secret key, token context, and candidate token into +reproducible random values. Generation and detection must produce identical +values across devices and releases. The values should also be uniform and +independent enough for the sampling algorithm and detector statistics. PRF +selection is an advanced compatibility and performance setting; most users +should keep the default. + +Watermarked generation currently supports the `philox` PRF: + +- `philox` is based on the counter-based Philox4x32-10 + generator from the [Random123 paper](https://doi.org/10.1145/2063384.2063405). + It is parallel, vectorizes on accelerators, and avoids CPU transfers, but is + not a cryptographic PRF and does not provide key-recovery or forgery + resistance. vLLM versions its input mapping and provides compatibility + vectors so generation and detection remain interoperable. + +## Detection + +The detector primitives operate on token IDs and do not require model weights: + +```python +from transformers import AutoTokenizer + +from vllm.v1.watermarking import GumbelWatermarkDetector + +tokenizer = AutoTokenizer.from_pretrained(MODEL) +token_ids = tokenizer.encode(text, add_special_tokens=False) +result = GumbelWatermarkDetector(key=42, prf="philox").detect(token_ids) +print(result.p_value, result.is_watermarked) +``` + +The tokenizer, algorithm, PRF, key, and context width must match generation. +Gumbel-max detection scores repeated contexts once by default so identical PRF +random vectors are not treated as independent evidence. Keep +`deduplicate_contexts=True` unless the detector's calibration has been adjusted +for correlated scores. + +The reported p-value is calibrated under the assumption that scored PRF inputs +are independent. A deployment uses one fixed key, so repeated structures across +documents reuse the same PRF values and can make the realized false-positive +rate key-dependent even when contexts are deduplicated within each document. +Measure the false-positive rate on representative unwatermarked traffic with the +deployed key before relying on `is_watermarked` for decisions. + +A minimal HTTP detector is available in +`examples/basic/online_serving/watermark_detection_server.py`: + +```bash +python examples/basic/online_serving/watermark_detection_server.py \ + --tokenizer MODEL --key 42 --prf philox +``` + +```bash +curl http://localhost:8000/detect \ + -H 'Content-Type: application/json' \ + -d '{"text":"Text to inspect"}' +``` + +Scores and p-values expose information about the per-token watermark signal. +Repeated queries can use this information to construct text that imitates +watermarked output or to modify watermarked text so it is no longer detected. + +## Limitations + +- Watermarking is currently available only with Model Runner V2. +- Gumbel-max cannot be configured with speculative decoding. +- Beam search expands candidates from model log probabilities and does not apply + Gumbel-max watermarking. +- Models that replace the vLLM sampler with a custom sampler cannot use + configured watermarking. diff --git a/examples/basic/online_serving/watermark_detection_server.py b/examples/basic/online_serving/watermark_detection_server.py new file mode 100644 index 0000000000..124850d75e --- /dev/null +++ b/examples/basic/online_serving/watermark_detection_server.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Minimal reference server for watermark detection.""" + +import argparse + +import uvicorn +from fastapi import FastAPI +from pydantic import BaseModel + +from vllm.tokenizers import TokenizerLike, cached_get_tokenizer +from vllm.v1.watermarking import GumbelWatermarkDetector + +app = FastAPI() +tokenizer: TokenizerLike | None = None +detector: GumbelWatermarkDetector | None = None + + +class DetectionRequest(BaseModel): + text: str + + +class DetectionResponse(BaseModel): + score: float + p_value: float + num_scored_tokens: int + is_watermarked: bool + + +@app.post("/detect") +def detect(request: DetectionRequest) -> DetectionResponse: + assert tokenizer is not None + assert detector is not None + token_ids = tokenizer.encode(request.text, add_special_tokens=False) + result = detector.detect(token_ids) + return DetectionResponse( + score=result.score, + p_value=result.p_value, + num_scored_tokens=result.num_scored_tokens, + is_watermarked=result.is_watermarked, + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokenizer", required=True) + parser.add_argument("--key", required=True, type=int) + parser.add_argument("--prf", choices=("philox",), default="philox") + parser.add_argument("--context-width", type=int, default=4) + parser.add_argument("--p-value-threshold", type=float, default=0.01) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + return parser.parse_args() + + +def main(args: argparse.Namespace) -> None: + global tokenizer, detector + tokenizer = cached_get_tokenizer(args.tokenizer) + detector = GumbelWatermarkDetector( + key=args.key, + context_width=args.context_width, + p_value_threshold=args.p_value_threshold, + prf=args.prf, + ) + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main(parse_args()) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index eb96af2972..f894861022 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -46,6 +46,27 @@ def test_optional_type(): assert optional_type_func("42") == 42 +def test_watermark_config_cli(): + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + args = parser.parse_args( + [ + "--model", + "dummy", + "--watermark-config", + '{"algorithm":"gumbel","key":42,"prf":"philox"}', + ] + ) + + config = EngineArgs.from_cli_args(args).create_watermark_config() + + assert config is not None + assert config.algorithm == "gumbel" + assert config.key == 42 + assert config.context_width == 4 + assert config.prf == "philox" + assert not config.supports_speculative_decoding + + @pytest.mark.parametrize( "options", [ diff --git a/tests/entrypoints/openai/test_watermarking.py b/tests/entrypoints/openai/test_watermarking.py new file mode 100644 index 0000000000..0425b86a02 --- /dev/null +++ b/tests/entrypoints/openai/test_watermarking.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + + +def test_chat_request_disables_watermarking(): + request = ChatCompletionRequest( + messages=[{"role": "user", "content": "hello"}], watermarking=False + ) + + params = request.to_sampling_params(max_tokens=1, default_sampling_params={}) + + assert not params.watermarking + + +def test_completion_request_disables_watermarking(): + request = CompletionRequest(prompt="hello", watermarking=False) + + params = request.to_sampling_params(max_tokens=1) + + assert not params.watermarking + + +def test_responses_request_disables_watermarking(): + request = ResponsesRequest(input="hello", watermarking=False) + + params = request.to_sampling_params(default_max_tokens=1) + + assert not params.watermarking diff --git a/tests/entrypoints/serve/utils/test_api_utils.py b/tests/entrypoints/serve/utils/test_api_utils.py index 4f8f7864fe..a41f39187f 100644 --- a/tests/entrypoints/serve/utils/test_api_utils.py +++ b/tests/entrypoints/serve/utils/test_api_utils.py @@ -120,17 +120,24 @@ class TestRedactSensitiveArgs: API_KEY = "sk-test-secret-12345" def test_redact_replaces_sensitive_values_only(self): - args = {"api_key": self.API_KEY, "hf_token": "hf_secret", "other": "visible"} + args = { + "api_key": self.API_KEY, + "hf_token": "hf_secret", + "watermark_config": {"algorithm": "gumbel", "key": 42}, + "other": "visible", + } redacted = redact_sensitive_args(args) assert redacted == { "api_key": "***", "hf_token": "***", + "watermark_config": "***", "other": "visible", } # original dict must not be mutated assert args == { "api_key": self.API_KEY, "hf_token": "hf_secret", + "watermark_config": {"algorithm": "gumbel", "key": 42}, "other": "visible", } diff --git a/tests/test_config.py b/tests/test_config.py index eae797adf5..96ca8f36e4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -32,6 +32,7 @@ from vllm.config import ( SchedulerConfig, SpeculativeConfig, VllmConfig, + WatermarkConfig, update_config, ) from vllm.config.compilation import CompilationMode, CUDAGraphMode @@ -2574,6 +2575,66 @@ def test_draft_sample_method_gumbel_is_rejected(): ) +def _watermarked_vllm_config() -> VllmConfig: + config = object.__new__(VllmConfig) + config.watermark_config = WatermarkConfig(key=42) + config.speculative_config = None + return config + + +def test_gumbel_watermark_rejects_speculative_decoding(): + config = _watermarked_vllm_config() + config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + ) + + with pytest.raises(ValueError, match="does not support speculative decoding"): + config._check_watermarking_unsupported() + + +def test_gumbel_watermark_rejects_beam_search(): + with pytest.raises(ValueError, match="Beam search is not supported"): + _watermarked_vllm_config()._check_watermarking_unsupported(beam_search=True) + + +def test_gumbel_watermark_rejects_custom_sampler(): + with pytest.raises(ValueError, match="custom samplers are not supported"): + _watermarked_vllm_config()._check_watermarking_unsupported(custom_sampler=True) + + +def test_watermark_key_must_fit_in_64_bits(): + with pytest.raises(ValueError, match="64 bits"): + WatermarkConfig(key=2**64) + + +def test_unknown_watermark_prf_is_rejected(): + with pytest.raises(ValidationError): + pydantic.TypeAdapter(WatermarkConfig).validate_python( + {"key": 42, "prf": "unsupported"} + ) + + +def test_watermark_key_is_excluded_from_serialization(): + serialized = pydantic.TypeAdapter(WatermarkConfig).dump_python( + WatermarkConfig(key=42), mode="json" + ) + + assert "key" not in serialized + + +def test_watermarking_forces_model_runner_v2(monkeypatch): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") + + with patch("vllm.config.vllm.logger.info_once") as info_once: + assert _watermarked_vllm_config().use_v2_model_runner + + info_once.assert_called_once_with( + "Watermarking requires Model Runner V2 and overrides " + "VLLM_USE_V2_MODEL_RUNNER=0." + ) + + @patch("vllm.config.speculative.ModelConfig") def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls): """Regression test: MTP + runai_streamer should use model_weights (original diff --git a/tests/watermarking/test_detection.py b/tests/watermarking/test_detection.py new file mode 100644 index 0000000000..4d10051d4d --- /dev/null +++ b/tests/watermarking/test_detection.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.v1.watermarking import GumbelWatermarkDetector + +_NON_WATERMARKED_TOKEN_IDS = list(range(64)) +_WATERMARKED_TOKEN_IDS = [ + 85, + 56, + 123, + 77, + 105, + 62, + 9, + 104, + 16, + 98, + 123, + 22, + 16, + 35, + 47, + 127, + 77, + 67, + 19, + 2, + 10, + 8, + 75, + 22, + 127, + 30, + 95, + 9, + 101, + 14, + 113, + 122, + 85, + 66, + 122, + 50, + 38, + 119, + 15, + 58, + 68, + 62, + 24, + 126, + 80, + 126, + 68, + 83, + 81, + 126, + 62, + 31, + 18, + 81, + 64, + 47, + 22, + 70, + 47, + 109, + 54, + 21, + 65, + 38, +] + + +def test_detection_on_known_sequence(): + detector = GumbelWatermarkDetector(key=42, context_width=4, prf="philox") + + assert detector.detect(_WATERMARKED_TOKEN_IDS).is_watermarked + assert not detector.detect(_NON_WATERMARKED_TOKEN_IDS).is_watermarked + + +def test_detector_handles_empty_input(): + detector = GumbelWatermarkDetector(key=42, context_width=4, prf="philox") + contexts, targets = detector._prepare_inputs([]) + detection = detector.detect([]) + + assert contexts.shape == (0, 4) + assert targets.shape == (0,) + assert detection.score == 0 + assert detection.p_value == 1 + assert not detection.is_watermarked diff --git a/tests/watermarking/test_gumbel.py b/tests/watermarking/test_gumbel.py new file mode 100644 index 0000000000..e3b4125172 --- /dev/null +++ b/tests/watermarking/test_gumbel.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.watermarking import GumbelWatermarkDetector, GumbelWatermarker +from vllm.v1.watermarking.gumbel import _gamma_survival_integer_shape + + +def test_gamma_survival_integer_shape(): + assert _gamma_survival_integer_shape(0.0, 4) == 1.0 + assert _gamma_survival_integer_shape(2.0, 1) == pytest.approx(0.1353352832) + + +def test_detector_deduplicates_repeated_prf_inputs_by_default(): + detection = GumbelWatermarkDetector(key=42, context_width=1).detect([1, 1, 1, 1, 1]) + + assert detection.num_scored_tokens == 2 + + +def test_detector_can_score_repeated_prf_inputs(): + detection = GumbelWatermarkDetector( + key=42, context_width=1, deduplicate_contexts=False + ).detect([1, 1, 1, 1, 1]) + + assert detection.num_scored_tokens == 5 + + +def test_detector_deduplicates_context_even_when_target_differs(): + detection = GumbelWatermarkDetector(key=42, context_width=1).detect([1, 2, 1, 3]) + + assert detection.num_scored_tokens == 3 + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +@pytest.mark.parametrize("key", [1, 2**32, 42, 15726070495360670683]) +@pytest.mark.parametrize("context_width", [1, 4, 16]) +def test_fused_watermarker_matches_cpu(key: int, context_width: int): + torch.manual_seed(0) + contexts = torch.randint(0, 248320, (32, context_width), dtype=torch.int64) + logits = torch.randn(32, 8193) + watermarker = GumbelWatermarker(key, context_width) + + expected = watermarker.sample(logits, contexts, lambda values: None).token_ids + actual = watermarker.sample( + logits.cuda(), contexts.cuda(), lambda values: None + ).token_ids.cpu() + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_fused_watermarker_handles_nan_logits(): + contexts = torch.zeros((2, 4), dtype=torch.int64, device="cuda") + logits = torch.full((2, 1025), float("nan"), device="cuda") + watermarker = GumbelWatermarker(key=42, context_width=4) + + token_ids = watermarker.sample(logits, contexts, lambda values: None).token_ids + + assert torch.all((token_ids >= 0) & (token_ids < logits.shape[-1])) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_fused_watermarker_handles_noncontiguous_inputs(): + contexts = torch.randint(0, 248320, (8, 8), dtype=torch.int64, device="cuda")[ + :, ::2 + ] + logits = torch.randn(8, 2050, device="cuda")[:, ::2] + watermarker = GumbelWatermarker(key=42, context_width=4) + + expected = watermarker.sample( + logits.contiguous(), contexts.contiguous(), lambda values: None + ).token_ids + actual = watermarker.sample(logits, contexts, lambda values: None).token_ids + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/watermarking/test_prf.py b/tests/watermarking/test_prf.py new file mode 100644 index 0000000000..437a1da695 --- /dev/null +++ b/tests/watermarking/test_prf.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.watermarking.prfs import PhiloxPRF + + +def test_philox_compatibility_vector(): + contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]]) + token_ids = torch.tensor([7, 8]) + expected_mantissas = [[1661833, 10571167], [9742331, 13375724]] + expected = ( + (torch.tensor(expected_mantissas, dtype=torch.float64) + 1) / (2**24 + 1) + ).to(torch.float32) + + assert torch.equal(PhiloxPRF(42).uniform(contexts, token_ids), expected) + + +def test_prf_pairs_contexts_with_target_tokens(): + contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]]) + token_ids = torch.tensor([[7], [8]]) + + assert PhiloxPRF(42).uniform(contexts, token_ids).shape == (2, 1) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +@pytest.mark.parametrize("key", [42, 15726070495360670683]) +@pytest.mark.parametrize("context_width", [1, 4, 16]) +def test_philox_accelerator_matches_cpu(key: int, context_width: int): + contexts = torch.arange(2 * context_width, dtype=torch.int64).reshape( + 2, context_width + ) + token_ids = torch.arange(1024, dtype=torch.int64) + prf = PhiloxPRF(key) + + expected = prf.uniform(contexts, token_ids) + actual = prf.uniform(contexts.cuda(), token_ids.cuda()).cpu() + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py new file mode 100644 index 0000000000..0452957ffc --- /dev/null +++ b/tests/watermarking/test_watermarking.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm import SamplingParams +from vllm.config.watermarking import WatermarkConfig +from vllm.v1.watermarking import create_watermarker +from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler +from vllm.v1.watermarking.watermarker import WatermarkSample +from vllm.v1.worker.gpu.sample.sampler import Sampler + + +@pytest.mark.parametrize("algorithm", ["gumbel"]) +def test_watermarker_contract(algorithm: str): + watermarker = create_watermarker( + WatermarkConfig(algorithm=algorithm, key=42, context_width=4) + ) + logits = torch.zeros(2, 128) + contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]]) + random_sample = lambda sample_logits: sample_logits.argmax(dim=-1) + + first = watermarker.sample(logits, contexts, random_sample) + second = watermarker.sample(logits, contexts, random_sample) + + assert first.token_ids.shape == (2,) + assert first.logits.shape == logits.shape + assert torch.equal(first.token_ids, second.token_ids) + + +def test_gumbel_config_warns_about_degenerate_generations(monkeypatch): + messages: list[str] = [] + monkeypatch.setattr( + "vllm.config.watermarking.logger.warning_once", + lambda message, *, scope: messages.append(message), + ) + + WatermarkConfig(key=42) + + assert messages == [ + "Single-key Gumbel-max watermarking may increase the frequency of " + "degenerate generations, including repetition loops." + ] + + +def test_large_context_width_warns_but_is_allowed(): + config = WatermarkConfig(key=42, context_width=17) + + with pytest.warns(UserWarning, match="reduce robustness to edits"): + watermarker = create_watermarker(config) + + assert watermarker.context_width == 17 + + +def test_sampling_params_can_disable_watermarking(): + assert SamplingParams().watermarking + assert not SamplingParams.from_optional(watermarking=False).watermarking + + +def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch): + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarking = SimpleNamespace(np=np.ones(1, dtype=bool)) + messages: list[str] = [] + monkeypatch.setattr(Sampler, "add_request", lambda *args, **kwargs: None) + monkeypatch.setattr( + "vllm.v1.watermarking.gpu_sampler.logger.warning_once", messages.append + ) + + sampler.add_request(0, 1, SamplingParams(temperature=0)) + sampler.add_request(0, 1, SamplingParams(temperature=1)) + sampler.add_request(0, 1, SamplingParams(temperature=0, watermarking=False)) + + assert messages == [ + ( + "Watermarking is enabled, but greedy decoding (temperature=0) cannot be " + "watermarked. This request will use ordinary greedy sampling." + ) + ] + + +def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch): + class StubWatermarker: + context_width = 1 + + def sample(self, logits, contexts, random_sample): + return WatermarkSample(torch.tensor([7, 7]), logits + 10) + + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = StubWatermarker() + sampler.watermarking = SimpleNamespace( + np=np.array([True, False]), gpu=torch.tensor([True, False]) + ) + sampler.sampling_states = SimpleNamespace( + temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)), + seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)), + ) + sampler.use_fp64_gumbel = False + sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros( + 2, 1, dtype=torch.int64 + ) + monkeypatch.setattr( + "vllm.v1.watermarking.gpu_sampler.gumbel_sample", + lambda *args, **kwargs: torch.tensor([3, 4]), + ) + logits = torch.zeros(2, 8) + + sampled, output_logits = sampler._sample_random( + logits, + torch.tensor([0, 1]), + np.array([0, 1]), + torch.zeros(2, dtype=torch.int64), + None, + None, + False, + ) + + assert torch.equal(sampled, torch.tensor([7, 4])) + assert torch.equal(output_logits[0], torch.full((8,), 10.0)) + assert torch.equal(output_logits[1], logits[1]) + + +def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch): + class StubWatermarker: + context_width = 1 + + def sample(self, logits, contexts, random_sample): + raise AssertionError("watermarker should not run for greedy requests") + + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = StubWatermarker() + sampler.watermarking = SimpleNamespace( + np=np.array([True, True]), gpu=torch.tensor([True, True]) + ) + sampler.sampling_states = SimpleNamespace( + temperature=SimpleNamespace(np=np.zeros(2), gpu=torch.zeros(2)), + ) + expected = (torch.tensor([3, 4]), torch.zeros(2, 8)) + monkeypatch.setattr( + "vllm.v1.watermarking.gpu_sampler.Sampler._sample_random", + lambda *args, **kwargs: expected, + ) + + actual = sampler._sample_random( + torch.zeros(2, 8), + torch.tensor([0, 1]), + np.array([0, 1]), + torch.zeros(2, dtype=torch.int64), + None, + None, + False, + ) + + assert actual is expected diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 8036d29093..1cfa372d60 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -60,6 +60,7 @@ from vllm.config.vllm import ( get_layers_from_vllm_config, set_current_vllm_config, ) +from vllm.config.watermarking import WatermarkConfig from vllm.config.weight_transfer import WeightTransferConfig # __all__ should only contain classes and functions. @@ -146,4 +147,5 @@ __all__ = [ "set_current_vllm_config", "get_layers_from_vllm_config", "WeightTransferConfig", + "WatermarkConfig", ] diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 0564cd420e..294f80b143 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -51,6 +51,7 @@ from .scheduler import SchedulerConfig from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig from .structured_outputs import StructuredOutputsConfig from .utils import SupportsHash, config, replace +from .watermarking import WatermarkConfig from .weight_transfer import WeightTransferConfig if TYPE_CHECKING: @@ -381,6 +382,8 @@ class VllmConfig: """LoRA configuration.""" speculative_config: SpeculativeConfig | None = None """Speculative decoding configuration.""" + watermark_config: WatermarkConfig | None = None + """Text watermarking configuration.""" diffusion_config: DiffusionConfig | None = None """Diffusion LLM (dLLM) configuration.""" @@ -662,6 +665,14 @@ class VllmConfig: @property def use_v2_model_runner(self) -> bool: + if getattr(self, "watermark_config", None) is not None: + if envs.VLLM_USE_V2_MODEL_RUNNER is False: + logger.info_once( + "Watermarking requires Model Runner V2 and overrides " + "VLLM_USE_V2_MODEL_RUNNER=0." + ) + return True + use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER if use_v2_model_runner is not None: return use_v2_model_runner @@ -1091,6 +1102,30 @@ class VllmConfig: if not self.use_v2_model_runner: raise ValueError("trace replay requires Model Runner V2") + def _check_watermarking_unsupported( + self, + *, + beam_search: bool = False, + custom_sampler: bool = False, + ) -> None: + watermark_config = getattr(self, "watermark_config", None) + if watermark_config is None: + return + if ( + self.speculative_config is not None + and not watermark_config.supports_speculative_decoding + ): + raise ValueError( + f"The {watermark_config.algorithm} watermarking algorithm " + "does not support speculative decoding." + ) + if beam_search: + raise ValueError("Beam search is not supported with watermarking.") + if custom_sampler: + raise ValueError( + "Model-specific custom samplers are not supported with watermarking." + ) + def _resolve_and_verify_engram_config(self) -> None: """Resolve legacy offload settings and validate model and parallel configs.""" if self.engram_config is None: @@ -1124,6 +1159,7 @@ class VllmConfig: self.try_verify_and_update_config() self._resolve_and_verify_engram_config() + self._check_watermarking_unsupported() # Models may have supplied their own DCP defaults above; anything still # unset falls back to the stock ones. self.parallel_config.set_dcp_defaults() diff --git a/vllm/config/watermarking.py b/vllm/config/watermarking.py new file mode 100644 index 0000000000..2247dd8248 --- /dev/null +++ b/vllm/config/watermarking.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Literal + +from pydantic import Field, model_validator +from typing_extensions import Self + +from vllm.config.utils import config +from vllm.logger import init_logger + +logger = init_logger(__name__) + +WatermarkingAlgorithm = Literal["gumbel"] +WatermarkPRFName = Literal["philox"] + +_SPECULATIVE_DECODING_SUPPORT: dict[WatermarkingAlgorithm, bool] = { + "gumbel": False, +} + + +@config +class WatermarkConfig: + """Configuration for text watermark generation.""" + + key: int = Field(ge=0, repr=False, exclude=True) + """Secret key used to watermark generated text.""" + algorithm: WatermarkingAlgorithm = "gumbel" + """Algorithm used to watermark generated text.""" + context_width: int = Field(default=4, ge=1) + """Number of prior output tokens used by the watermark PRF.""" + prf: WatermarkPRFName = "philox" + """Pseudorandom function used by the watermarking algorithm.""" + + @model_validator(mode="after") + def validate_key(self) -> Self: + if self.key > 2**64 - 1: + raise ValueError("philox keys must fit in 64 bits") + if self.algorithm == "gumbel": + logger.warning_once( + "Single-key Gumbel-max watermarking may increase the frequency of " + "degenerate generations, including repetition loops.", + scope="global", + ) + return self + + @property + def supports_speculative_decoding(self) -> bool: + return _SPECULATIVE_DECODING_SUPPORT[self.algorithm] diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 2fed4f3b6a..3a36021b93 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -106,6 +106,7 @@ from vllm.config.parallel import ( from vllm.config.scheduler import SchedulerPolicy from vllm.config.utils import get_field from vllm.config.vllm import OptimizationLevel, PerformanceMode +from vllm.config.watermarking import WatermarkConfig from vllm.logger import init_logger, suppress_logging from vllm.platforms import CpuArchEnum, current_platform from vllm.plugins import load_general_plugins @@ -656,6 +657,7 @@ class EngineArgs: spec_method: str | None = None spec_model: str | None = None spec_tokens: int | None = None + watermark_config: dict[str, Any] | None = None diffusion_config: dict[str, Any] | None = None show_hidden_metrics_for_version: str | None = ( @@ -1680,6 +1682,8 @@ class EngineArgs: vllm_group.add_argument( "--spec-tokens", **speculative_kwargs["num_speculative_tokens"] ) + vllm_kwargs["watermark_config"]["type"] = optional_type(json.loads) + vllm_group.add_argument("--watermark-config", **vllm_kwargs["watermark_config"]) vllm_kwargs["diffusion_config"]["type"] = optional_type(json.loads) vllm_group.add_argument( "--diffusion-config", "-dc", **vllm_kwargs["diffusion_config"] @@ -1986,6 +1990,14 @@ class EngineArgs: cfg = json.loads(cfg) return DiffusionConfig(**cfg) + def create_watermark_config(self) -> WatermarkConfig | None: + if self.watermark_config is None: + return None + cfg = self.watermark_config + if isinstance(cfg, str): + cfg = json.loads(cfg) + return WatermarkConfig(**cfg) + def create_observability_config(self) -> ObservabilityConfig: return ObservabilityConfig( show_hidden_metrics_for_version=self.show_hidden_metrics_for_version, @@ -2366,6 +2378,7 @@ class EngineArgs: target_parallel_config=parallel_config, ) diffusion_config = self.create_diffusion_config() + watermark_config = self.create_watermark_config() self._set_default_max_num_seqs_and_batched_tokens_args( usage_context, @@ -2602,6 +2615,7 @@ class EngineArgs: kernel_config=kernel_config, lora_config=lora_config, speculative_config=speculative_config, + watermark_config=watermark_config, diffusion_config=diffusion_config, structured_outputs_config=self.structured_outputs_config, observability_config=observability_config, diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index 7814d4d44e..f412f5aebe 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -82,6 +82,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): temperature = params.temperature ignore_eos = params.ignore_eos length_penalty = params.length_penalty + self.llm_engine.vllm_config._check_watermarking_unsupported(beam_search=True) tokenizer = self.renderer.get_tokenizer() eos_token_id = tokenizer.eos_token_id diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 89c1579865..b7b292e835 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -267,6 +267,7 @@ class ChatCompletionRequest(OpenAIBaseModel): top_k: int | None = None min_p: float | None = None repetition_penalty: float | None = None + watermarking: bool = True length_penalty: float = 1.0 stop_token_ids: list[int] | None = [] include_stop_str_in_output: bool = False @@ -724,6 +725,7 @@ class ChatCompletionRequest(OpenAIBaseModel): frequency_penalty=self.frequency_penalty, repetition_penalty=repetition_penalty, temperature=temperature, + watermarking=self.watermarking, top_p=top_p, top_k=top_k, min_p=min_p, diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index b8d66925c0..6d7a49bacd 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -75,6 +75,7 @@ class CompletionRequest(OpenAIBaseModel): top_k: int | None = None min_p: float | None = None repetition_penalty: float | None = None + watermarking: bool = True length_penalty: float = 1.0 stop_token_ids: list[int] | None = [] include_stop_str_in_output: bool = False @@ -373,6 +374,7 @@ class CompletionRequest(OpenAIBaseModel): frequency_penalty=self.frequency_penalty, repetition_penalty=repetition_penalty, temperature=temperature, + watermarking=self.watermarking, top_p=top_p, top_k=top_k, min_p=min_p, diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 790024fad3..27c60967e5 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -213,6 +213,7 @@ class ResponsesRequest(OpenAIBaseModel): ) # --8<-- [start:responses-extra-params] + watermarking: bool = True request_id: str = Field( default_factory=lambda: f"resp_{random_uuid()}", description=( @@ -436,6 +437,7 @@ class ResponsesRequest(OpenAIBaseModel): return SamplingParams.from_optional( temperature=temperature, + watermarking=self.watermarking, top_p=top_p, top_k=top_k, max_tokens=max_tokens, diff --git a/vllm/entrypoints/serve/utils/api_utils.py b/vllm/entrypoints/serve/utils/api_utils.py index 51c94a8e1f..aa089cff44 100644 --- a/vllm/entrypoints/serve/utils/api_utils.py +++ b/vllm/entrypoints/serve/utils/api_utils.py @@ -268,7 +268,7 @@ def jsonify_non_default_args( # Fields whose values must never be logged verbatim. -_SENSITIVE_ARG_FIELDS = frozenset({"api_key", "hf_token"}) +_SENSITIVE_ARG_FIELDS = frozenset({"api_key", "hf_token", "watermark_config"}) def redact_sensitive_args(args: dict[str, Any]) -> dict[str, Any]: diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 7108803065..52a65aefbc 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -253,6 +253,8 @@ class SamplingParams( """Controls the randomness of the sampling. Lower values make the model more deterministic, while higher values make the model more random. Zero means greedy sampling.""" + watermarking: bool = True + """Whether to apply the engine's configured watermark to this request.""" top_p: float = 1.0 """Controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to 1 to consider all tokens.""" @@ -385,6 +387,7 @@ class SamplingParams( frequency_penalty: float | None = 0.0, repetition_penalty: float | None = 1.0, temperature: float | None = 1.0, + watermarking: bool = True, top_p: float | None = 1.0, top_k: int = 0, min_p: float = 0.0, @@ -450,6 +453,7 @@ class SamplingParams( if repetition_penalty is None else repetition_penalty, temperature=1.0 if temperature is None else temperature, + watermarking=watermarking, top_p=1.0 if top_p is None else top_p, top_k=top_k, min_p=min_p, @@ -1253,6 +1257,7 @@ class SamplingParams( f"frequency_penalty={self.frequency_penalty}, " f"repetition_penalty={self.repetition_penalty}, " f"temperature={self.temperature}, " + f"watermarking={self.watermarking}, " f"top_p={self.top_p}, " f"top_k={self.top_k}, " f"min_p={self.min_p}, " diff --git a/vllm/v1/watermarking/__init__.py b/vllm/v1/watermarking/__init__.py new file mode 100644 index 0000000000..bb53981461 --- /dev/null +++ b/vllm/v1/watermarking/__init__.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.v1.watermarking.detector import WatermarkDetection, WatermarkDetector +from vllm.v1.watermarking.factory import create_watermarker +from vllm.v1.watermarking.gumbel import ( + GumbelWatermarkDetector, + GumbelWatermarker, +) +from vllm.v1.watermarking.prfs import ( + PhiloxPRF, + WatermarkPRF, + create_prf, +) +from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample + +__all__ = [ + "GumbelWatermarkDetector", + "GumbelWatermarker", + "PhiloxPRF", + "WatermarkDetection", + "WatermarkDetector", + "WatermarkPRF", + "Watermarker", + "WatermarkSample", + "create_watermarker", + "create_prf", +] diff --git a/vllm/v1/watermarking/detector.py b/vllm/v1/watermarking/detector.py new file mode 100644 index 0000000000..42759b03c4 --- /dev/null +++ b/vllm/v1/watermarking/detector.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TypeAlias + +import torch + +Context: TypeAlias = tuple[int, ...] + + +@dataclass(frozen=True) +class WatermarkDetection: + score: float + p_value: float + num_scored_tokens: int + is_watermarked: bool + + +class WatermarkDetector(ABC): + def __init__( + self, + context_width: int, + p_value_threshold: float, + deduplicate_contexts: bool = False, + ) -> None: + if context_width < 1: + raise ValueError("context_width must be positive") + if not 0 < p_value_threshold < 1: + raise ValueError("p_value_threshold must be between 0 and 1") + self.context_width = context_width + self.p_value_threshold = p_value_threshold + self.deduplicate_contexts = deduplicate_contexts + + def detect(self, token_ids: list[int]) -> WatermarkDetection: + contexts, targets = self._prepare_inputs(token_ids) + num_scored_tokens = len(targets) + if num_scored_tokens == 0: + return WatermarkDetection(0.0, 1.0, 0, False) + + token_scores = self._score_tokens(contexts, targets) + score = self._aggregate_scores(token_scores) + p_value = self._get_p_value(score, num_scored_tokens) + return WatermarkDetection( + score=score, + p_value=p_value, + num_scored_tokens=num_scored_tokens, + is_watermarked=p_value <= self.p_value_threshold, + ) + + def _prepare_inputs( + self, token_ids: list[int] + ) -> tuple[torch.Tensor, torch.Tensor]: + contexts: list[Context] = [] + targets = [] + seen_contexts: set[Context] = set() + prefix = [-1] * self.context_width + for token_id in token_ids: + context = tuple(prefix[-self.context_width :]) + # Reusing a context reuses one keyed random vector, so the resulting + # token scores are not independent evidence. + if not self.deduplicate_contexts or not _is_repeated_context( + context, seen_contexts + ): + contexts.append(context) + targets.append(token_id) + prefix.append(token_id) + + if not contexts: + return ( + torch.empty((0, self.context_width), dtype=torch.int64), + torch.empty(0, dtype=torch.int64), + ) + + return ( + torch.tensor(contexts, dtype=torch.int64).reshape(-1, self.context_width), + torch.tensor(targets, dtype=torch.int64), + ) + + @abstractmethod + def _score_tokens( + self, contexts: torch.Tensor, targets: torch.Tensor + ) -> torch.Tensor: + raise NotImplementedError + + @abstractmethod + def _aggregate_scores(self, token_scores: torch.Tensor) -> float: + raise NotImplementedError + + @abstractmethod + def _get_p_value(self, score: float, num_scored_tokens: int) -> float: + raise NotImplementedError + + +def _is_repeated_context(context: Context, seen_contexts: set[Context]) -> bool: + if context in seen_contexts: + return True + seen_contexts.add(context) + return False diff --git a/vllm/v1/watermarking/factory.py b/vllm/v1/watermarking/factory.py new file mode 100644 index 0000000000..d9111e7a9e --- /dev/null +++ b/vllm/v1/watermarking/factory.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable + +from vllm.config.watermarking import WatermarkConfig +from vllm.v1.watermarking.gumbel import GumbelWatermarker +from vllm.v1.watermarking.watermarker import Watermarker + + +def _create_gumbel(config: WatermarkConfig) -> Watermarker: + return GumbelWatermarker(config.key, config.context_width, config.prf) + + +_WATERMARKERS: dict[str, Callable[[WatermarkConfig], Watermarker]] = { + "gumbel": _create_gumbel, +} + + +def create_watermarker(config: WatermarkConfig) -> Watermarker: + watermarker_factory = _WATERMARKERS.get(config.algorithm) + if watermarker_factory is None: + raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}") + return watermarker_factory(config) diff --git a/vllm/v1/watermarking/gpu_sampler.py b/vllm/v1/watermarking/gpu_sampler.py new file mode 100644 index 0000000000..7e943cc040 --- /dev/null +++ b/vllm/v1/watermarking/gpu_sampler.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import torch + +from vllm.logger import init_logger +from vllm.sampling_params import SamplingParams +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p +from vllm.v1.watermarking.watermarker import Watermarker +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.sample.sampler import Sampler + +logger = init_logger(__name__) + + +class GPUWatermarkSampler(Sampler): + def __init__( + self, + watermarker: Watermarker, + *args, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self.watermarker = watermarker + self.watermarking = UvaBackedTensor( + self.sampling_states.max_num_reqs, dtype=torch.bool + ) + self.watermarking.np.fill(True) + self.watermarking.copy_to_uva() + + def add_request( + self, req_idx: int, prompt_len: int, sampling_params: SamplingParams + ) -> None: + super().add_request(req_idx, prompt_len, sampling_params) + self.watermarking.np[req_idx] = sampling_params.watermarking + if sampling_params.watermarking and sampling_params.temperature == 0: + logger.warning_once( + "Watermarking is enabled, but greedy decoding " + "(temperature=0) cannot be watermarked. This request will use " + "ordinary greedy sampling." + ) + + def apply_staged_writes(self) -> None: + super().apply_staged_writes() + self.watermarking.copy_to_uva() + + def _sample_random( + self, + processed_logits: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + idx_mapping_np: np.ndarray, + pos: torch.Tensor, + top_k: torch.Tensor | None, + top_p: torch.Tensor | None, + use_flashinfer: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + enabled = self.watermarking.np[idx_mapping_np] & ( + self.sampling_states.temperature.np[idx_mapping_np] != 0 + ) + if not np.any(enabled): + return super()._sample_random( + processed_logits, + expanded_idx_mapping, + idx_mapping_np, + pos, + top_k, + top_p, + use_flashinfer, + ) + + processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p) + contexts = self._get_contexts(expanded_idx_mapping) + + def random_sample(sample_logits: torch.Tensor) -> torch.Tensor: + return gumbel_sample( + sample_logits, + expanded_idx_mapping, + self.sampling_states.temperature.gpu, + self.sampling_states.seeds.gpu, + pos, + apply_temperature=False, + is_drafting=False, + use_fp64=self.use_fp64_gumbel, + ) + + output = self.watermarker.sample( + processed_logits, + contexts, + random_sample, + ) + temperatures = self.sampling_states.temperature.gpu[expanded_idx_mapping] + watermarking = self.watermarking.gpu[expanded_idx_mapping] & (temperatures != 0) + if not np.all(enabled): + unwatermarked = random_sample(processed_logits) + sampled = torch.where(watermarking, output.token_ids, unwatermarked) + output_logits = output.logits + if output.logits is not processed_logits: + output_logits = torch.where( + watermarking.unsqueeze(-1), output.logits, processed_logits + ) + else: + sampled = output.token_ids + output_logits = output.logits + sampled = torch.where( + temperatures == 0, + processed_logits.argmax(dim=-1), + sampled, + ) + return sampled, output_logits + + def _get_contexts(self, expanded_idx_mapping: torch.Tensor) -> torch.Tensor: + context_width = self.watermarker.context_width + req_indices = expanded_idx_mapping.to(torch.int64) + valid_reqs = req_indices >= 0 + safe_req_indices = req_indices.clamp_min(0) + total_lens = self.req_states.total_len.gpu[safe_req_indices].to(torch.int64) + prompt_lens = self.req_states.prompt_len.gpu[safe_req_indices].to(torch.int64) + offsets = torch.arange( + -context_width, 0, dtype=torch.int64, device=req_indices.device + ) + positions = total_lens.unsqueeze(-1) + offsets + valid_positions = valid_reqs.unsqueeze(-1) & ( + positions >= prompt_lens.unsqueeze(-1) + ) + positions = positions.clamp_min(0) + contexts = self.req_states.all_token_ids.gpu[ + safe_req_indices.unsqueeze(-1), positions + ] + return torch.where(valid_positions, contexts, -1) diff --git a/vllm/v1/watermarking/gumbel.py b/vllm/v1/watermarking/gumbel.py new file mode 100644 index 0000000000..856d697aea --- /dev/null +++ b/vllm/v1/watermarking/gumbel.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Gumbel-max watermark generation and detection primitives.""" + +import math +import warnings + +import torch + +from vllm.config.watermarking import WatermarkPRFName +from vllm.v1.watermarking.detector import ( + WatermarkDetector, +) +from vllm.v1.watermarking.prfs import PhiloxPRF, WatermarkPRF, create_prf +from vllm.v1.watermarking.watermarker import ( + RandomSampler, + Watermarker, + WatermarkSample, +) + + +def _gamma_survival_integer_shape(score: float, shape: int) -> float: + """Return the survival probability of a Gamma(shape, 1) variable.""" + if shape <= 0 or score <= 0: + return 1.0 + log_score = math.log(score) + log_term = 0.0 + max_log_term = 0.0 + log_terms = [log_term] + for index in range(1, shape): + log_term += log_score - math.log(index) + log_terms.append(log_term) + max_log_term = max(max_log_term, log_term) + log_sum = max_log_term + math.log( + sum(math.exp(term - max_log_term) for term in log_terms) + ) + return min(1.0, math.exp(min(0.0, -score + log_sum))) + + +class GumbelWatermarker(Watermarker): + def __init__( + self, + key: int, + context_width: int = 4, + prf: WatermarkPRF | WatermarkPRFName = "philox", + ) -> None: + prf = create_prf(prf, key) if isinstance(prf, str) else prf + _validate_context_width(context_width) + self.prf = prf + self._context_width = context_width + + @property + def context_width(self) -> int: + return self._context_width + + def sample( + self, + logits: torch.Tensor, + contexts: torch.Tensor, + _random_sample: RandomSampler, + ) -> WatermarkSample: + if type(self.prf) is PhiloxPRF and logits.device.type == "cuda": + from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample + + return WatermarkSample( + philox_gumbel_sample(logits, contexts, self.prf.key), logits + ) + vocabulary = torch.arange(logits.shape[-1], device=logits.device) + uniforms = self.prf.uniform(contexts, vocabulary) + uniforms = uniforms.clamp_min(torch.finfo(torch.float32).tiny) + noise = -torch.log(-torch.log(uniforms)) + return WatermarkSample(torch.argmax(logits + noise, dim=-1), logits) + + +class GumbelWatermarkDetector(WatermarkDetector): + def __init__( + self, + key: int, + context_width: int = 4, + p_value_threshold: float = 0.01, + prf: WatermarkPRF | WatermarkPRFName = "philox", + deduplicate_contexts: bool = True, + ) -> None: + prf = create_prf(prf, key) if isinstance(prf, str) else prf + _validate_context_width(context_width) + super().__init__(context_width, p_value_threshold, deduplicate_contexts) + self.prf = prf + + def _score_tokens( + self, contexts: torch.Tensor, targets: torch.Tensor + ) -> torch.Tensor: + uniforms = self.prf.uniform(contexts, targets.unsqueeze(-1)).squeeze(-1) + return -torch.log1p(-uniforms.to(torch.float64)) + + def _get_p_value(self, score: float, num_scored_tokens: int) -> float: + return _gamma_survival_integer_shape(score, num_scored_tokens) + + def _aggregate_scores(self, token_scores: torch.Tensor) -> float: + return token_scores.sum().item() + + +def _validate_context_width(context_width: int) -> None: + if context_width < 1: + raise ValueError("context_width must be positive") + if context_width > 16: + warnings.warn( + "context_width values greater than 16 reduce robustness to edits because " + "each changed token affects more subsequent watermark contexts", + stacklevel=3, + ) diff --git a/vllm/v1/watermarking/prfs/__init__.py b/vllm/v1/watermarking/prfs/__init__.py new file mode 100644 index 0000000000..b88844f17f --- /dev/null +++ b/vllm/v1/watermarking/prfs/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.config.watermarking import WatermarkPRFName +from vllm.v1.watermarking.prfs.base import WatermarkPRF +from vllm.v1.watermarking.prfs.philox import PhiloxPRF + + +def create_prf(name: WatermarkPRFName, key: int) -> WatermarkPRF: + if name == "philox": + return PhiloxPRF(key) + raise ValueError(f"Unknown watermark PRF: {name}") + + +__all__ = [ + "PhiloxPRF", + "WatermarkPRF", + "create_prf", +] diff --git a/vllm/v1/watermarking/prfs/base.py b/vllm/v1/watermarking/prfs/base.py new file mode 100644 index 0000000000..02e4ea2e4e --- /dev/null +++ b/vllm/v1/watermarking/prfs/base.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod + +import torch + +_UINT32_MASK = 2**32 - 1 + + +class WatermarkPRF(ABC): + @abstractmethod + def uniform(self, contexts: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + +def uint32_to_uniform(values: torch.Tensor) -> torch.Tensor: + mantissas = (values & _UINT32_MASK) >> 8 + return ((mantissas.to(torch.float64) + 1.0) / (2**24 + 1)).to(torch.float32) diff --git a/vllm/v1/watermarking/prfs/philox.py b/vllm/v1/watermarking/prfs/philox.py new file mode 100644 index 0000000000..47555d1611 --- /dev/null +++ b/vllm/v1/watermarking/prfs/philox.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Philox4x32-10 as defined by https://doi.org/10.1145/2063384.2063405.""" + +import torch + +from vllm.v1.watermarking.prfs.base import WatermarkPRF, uint32_to_uniform + +_UINT32_MASK = 2**32 - 1 +_PHILOX_M0 = 0xD2511F53 +_PHILOX_M1 = 0xCD9E8D57 +_PHILOX_W0 = 0x9E3779B9 +_PHILOX_W1 = 0xBB67AE85 +_CONTEXT_DOMAIN = 0x574D4358 +_TOKEN_DOMAIN = 0x574D544B + + +def _mulhilo32( + multiplier: int, values: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + value_low = values & 0xFFFF + value_high = values >> 16 + multiplier_low = multiplier & 0xFFFF + multiplier_high = multiplier >> 16 + + low_product = value_low * multiplier_low + cross_product = value_high * multiplier_low + value_low * multiplier_high + low_with_carry = low_product + ((cross_product & 0xFFFF) << 16) + low = low_with_carry & _UINT32_MASK + high = ( + value_high * multiplier_high + (cross_product >> 16) + (low_with_carry >> 32) + ) & _UINT32_MASK + return high, low + + +def _philox4x32_10( + counter: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + key: tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + counter_0, counter_1, counter_2, counter_3 = counter + key_0, key_1 = key + for round_index in range(10): + high_0, low_0 = _mulhilo32(_PHILOX_M0, counter_0) + high_1, low_1 = _mulhilo32(_PHILOX_M1, counter_2) + counter_0, counter_1, counter_2, counter_3 = ( + (high_1 ^ counter_1 ^ key_0) & _UINT32_MASK, + low_1, + (high_0 ^ counter_3 ^ key_1) & _UINT32_MASK, + low_0, + ) + if round_index != 9: + key_0 = (key_0 + _PHILOX_W0) & _UINT32_MASK + key_1 = (key_1 + _PHILOX_W1) & _UINT32_MASK + return counter_0, counter_1, counter_2, counter_3 + + +def _philox_uniform( + key_0_value: int, + key_1_value: int, + contexts: torch.Tensor, + token_ids: torch.Tensor, +) -> torch.Tensor: + contexts = contexts.to(torch.int64) & _UINT32_MASK + prefix_shape = contexts.shape[:-1] + device = contexts.device + key_0 = torch.full(prefix_shape, key_0_value, dtype=torch.int64, device=device) + key_1 = torch.full(prefix_shape, key_1_value, dtype=torch.int64, device=device) + state = ( + torch.full(prefix_shape, _CONTEXT_DOMAIN, dtype=torch.int64, device=device), + torch.full( + prefix_shape, + contexts.shape[-1], + dtype=torch.int64, + device=device, + ), + torch.zeros(prefix_shape, dtype=torch.int64, device=device), + torch.zeros(prefix_shape, dtype=torch.int64, device=device), + ) + + for offset in range(0, contexts.shape[-1], 4): + block = [] + for index in range(4): + position = offset + index + if position < contexts.shape[-1]: + block.append(contexts[..., position]) + else: + block.append( + torch.full( + prefix_shape, + _UINT32_MASK - index, + dtype=torch.int64, + device=device, + ) + ) + counter = tuple( + (state[index] ^ block[index]) & _UINT32_MASK for index in range(4) + ) + state = _philox4x32_10(counter, ((key_0 ^ offset) & _UINT32_MASK, key_1)) + + output_shape = torch.broadcast_shapes(prefix_shape + (1,), token_ids.shape) + tokens = torch.broadcast_to(token_ids.to(device=device), output_shape) + state = tuple( + torch.broadcast_to(word.unsqueeze(-1), output_shape) for word in state + ) + candidate_key_0 = torch.broadcast_to(key_0.unsqueeze(-1), output_shape) + candidate_key_1 = torch.broadcast_to(key_1.unsqueeze(-1), output_shape) + token_words = tokens.to(torch.int64) & _UINT32_MASK + outputs = _philox4x32_10( + ( + token_words >> 2, + state[0], + state[1], + state[2], + ), + ( + (candidate_key_0 ^ state[3]) & _UINT32_MASK, + (candidate_key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK, + ), + ) + word_index = token_words & 3 + output = outputs[0] + for index in range(1, 4): + output = torch.where(word_index == index, outputs[index], output) + return uint32_to_uniform(output) + + +_compiled_philox_uniform = torch.compile(_philox_uniform, fullgraph=True, dynamic=True) + + +class PhiloxPRF(WatermarkPRF): + """Accelerator-friendly Philox4x32-10 watermark PRF.""" + + version = "philox4x32-10-v1" + + def __init__(self, key: int) -> None: + if not 0 <= key <= 2**64 - 1: + raise ValueError("Philox keys must fit in 64 bits") + self.key = key + + def uniform(self, contexts: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + key_words = self.key & _UINT32_MASK, self.key >> 32 + if contexts.device.type == "cuda": + return _compiled_philox_uniform(*key_words, contexts, token_ids) + return _philox_uniform(*key_words, contexts, token_ids) diff --git a/vllm/v1/watermarking/watermarker.py b/vllm/v1/watermarking/watermarker.py new file mode 100644 index 0000000000..f00659393e --- /dev/null +++ b/vllm/v1/watermarking/watermarker.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass +from typing import TypeAlias + +import torch + +RandomSampler: TypeAlias = Callable[[torch.Tensor], torch.Tensor] + + +@dataclass(frozen=True) +class WatermarkSample: + token_ids: torch.Tensor + logits: torch.Tensor + + +class Watermarker(ABC): + @property + @abstractmethod + def context_width(self) -> int: + raise NotImplementedError + + @abstractmethod + def sample( + self, + logits: torch.Tensor, + contexts: torch.Tensor, + random_sample: RandomSampler, + ) -> WatermarkSample: + raise NotImplementedError diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 7dec923c0a..1ee6a73072 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -78,6 +78,8 @@ from vllm.v1.outputs import ( ModelRunnerOutput, RoutedExpertsTensors, ) +from vllm.v1.watermarking import create_watermarker +from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler from vllm.v1.worker.block_table import get_block_table_width from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu import pcp_manager as pcp @@ -452,21 +454,27 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Initialize samplers. Model states may override via custom_sampler(). if self.is_last_pp_rank and not self.is_pooling_model: - self.sampler = Sampler( - max_num_reqs=self.max_num_reqs, - vocab_size=self.vocab_size, - device=self.device, - req_states=self.req_states, - logprobs_mode=self.model_config.logprobs_mode, - num_speculative_tokens=self.decode_query_len, - use_fp64_gumbel=self.model_config.use_fp64_gumbel, - enable_trace_replay=self.model_config.enable_trace_replay, - reasoning_config=self.vllm_config.reasoning_config, - return_sampling_mask=self.model_config.return_sampling_mask, - ) + sampler_kwargs: dict[str, Any] = { + "max_num_reqs": self.max_num_reqs, + "vocab_size": self.vocab_size, + "device": self.device, + "req_states": self.req_states, + "logprobs_mode": self.model_config.logprobs_mode, + "num_speculative_tokens": self.decode_query_len, + "use_fp64_gumbel": self.model_config.use_fp64_gumbel, + "enable_trace_replay": self.model_config.enable_trace_replay, + "reasoning_config": self.vllm_config.reasoning_config, + "return_sampling_mask": self.model_config.return_sampling_mask, + } + if self.vllm_config.watermark_config is None: + self.sampler = Sampler(**sampler_kwargs) + else: + watermarker = create_watermarker(self.vllm_config.watermark_config) + self.sampler = GPUWatermarkSampler(watermarker, **sampler_kwargs) custom = self.model_state.custom_sampler(self.sampler) if custom: + self.vllm_config._check_watermarking_unsupported(custom_sampler=True) self.sampler, self.rejection_sampler = custom elif self.speculative_config is not None: self.rejection_sampler = RejectionSampler( diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 614431d356..b069604c29 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -307,7 +307,26 @@ class Sampler: or self.sampling_states.any_explicit_seed(idx_mapping_np) ) - # Sample the next token. + return self._sample_random( + processed_logits, + expanded_idx_mapping, + idx_mapping_np, + pos, + top_k, + top_p, + use_flashinfer, + ) + + def _sample_random( + self, + processed_logits: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + idx_mapping_np: np.ndarray, + pos: torch.Tensor, + top_k: torch.Tensor | None, + top_p: torch.Tensor | None, + use_flashinfer: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: if use_flashinfer: sampled = flashinfer_sample(processed_logits, top_k, top_p).to(torch.int64) else: diff --git a/vllm/v1/worker/gpu/sample/watermark.py b/vllm/v1/worker/gpu/sample/watermark.py new file mode 100644 index 0000000000..a0105cc1b7 --- /dev/null +++ b/vllm/v1/worker/gpu/sample/watermark.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.v1.watermarking.prfs.philox import ( + _CONTEXT_DOMAIN as _CONTEXT_DOMAIN_VALUE, +) +from vllm.v1.watermarking.prfs.philox import ( + _PHILOX_M0 as _PHILOX_M0_VALUE, +) +from vllm.v1.watermarking.prfs.philox import ( + _PHILOX_M1 as _PHILOX_M1_VALUE, +) +from vllm.v1.watermarking.prfs.philox import ( + _PHILOX_W0 as _PHILOX_W0_VALUE, +) +from vllm.v1.watermarking.prfs.philox import ( + _PHILOX_W1 as _PHILOX_W1_VALUE, +) +from vllm.v1.watermarking.prfs.philox import ( + _TOKEN_DOMAIN as _TOKEN_DOMAIN_VALUE, +) +from vllm.v1.watermarking.prfs.philox import ( + _UINT32_MASK as _UINT32_MASK_VALUE, +) + +if HAS_TRITON: + from triton.language import math as tl_math +else: + tl_math = tl + +_UINT32_MASK = tl.constexpr(_UINT32_MASK_VALUE) if HAS_TRITON else _UINT32_MASK_VALUE +_PHILOX_M0 = tl.constexpr(_PHILOX_M0_VALUE) if HAS_TRITON else _PHILOX_M0_VALUE +_PHILOX_M1 = tl.constexpr(_PHILOX_M1_VALUE) if HAS_TRITON else _PHILOX_M1_VALUE +_PHILOX_W0 = tl.constexpr(_PHILOX_W0_VALUE) if HAS_TRITON else _PHILOX_W0_VALUE +_PHILOX_W1 = tl.constexpr(_PHILOX_W1_VALUE) if HAS_TRITON else _PHILOX_W1_VALUE +_CONTEXT_DOMAIN = ( + tl.constexpr(_CONTEXT_DOMAIN_VALUE) if HAS_TRITON else _CONTEXT_DOMAIN_VALUE +) +_TOKEN_DOMAIN = tl.constexpr(_TOKEN_DOMAIN_VALUE) if HAS_TRITON else _TOKEN_DOMAIN_VALUE + + +@triton.jit +def _mulhilo32(multiplier: tl.constexpr, value): + high = tl_math.umulhi(multiplier, value) + low = tl.mul(multiplier, value, sanitize_overflow=False) + return high, low + + +@triton.jit +def _philox4x32_10(counter_0, counter_1, counter_2, counter_3, key_0, key_1): + for round_index in range(10): + high_0, low_0 = _mulhilo32(_PHILOX_M0, counter_0) + high_1, low_1 = _mulhilo32(_PHILOX_M1, counter_2) + counter_0, counter_1, counter_2, counter_3 = ( + (high_1 ^ counter_1 ^ key_0) & _UINT32_MASK, + low_1, + (high_0 ^ counter_3 ^ key_1) & _UINT32_MASK, + low_0, + ) + if round_index != 9: + key_0 = (key_0 + _PHILOX_W0) & _UINT32_MASK + key_1 = (key_1 + _PHILOX_W1) & _UINT32_MASK + return counter_0, counter_1, counter_2, counter_3 + + +@triton.jit +def _uint32_to_uniform(value): + mantissa = (value & _UINT32_MASK) >> 8 + scaled = (mantissa + 1).to(tl.float32) * 2**-24 + return (scaled.to(tl.uint32, bitcast=True) - 1).to(tl.float32, bitcast=True) + + +@triton.jit +def _gumbel_value(logits_ptr, output, mask): + logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32) + logits = tl.where(logits != logits, float("-inf"), logits) + uniform = _uint32_to_uniform(output) + return logits - tl.log(-tl.log(uniform)) + + +@triton.jit +def _select_max(best_value, best_token, candidate_value, candidate): + take_candidate = candidate_value > best_value + return ( + tl.where(take_candidate, candidate_value, best_value), + tl.where(take_candidate, candidate, best_token), + ) + + +@triton.jit(do_not_specialize=["key_0_value", "key_1_value"]) +def _philox_gumbel_kernel( + local_argmax_ptr, + local_argmax_stride, + local_max_ptr, + local_max_stride, + logits_ptr, + logits_stride, + contexts_ptr, + context_stride, + key_0_value, + key_1_value, + vocab_size, + CONTEXT_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + block_index = tl.program_id(1) + groups = block_index * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4) + + key_0 = key_0_value.to(tl.uint32) + key_1 = key_1_value.to(tl.uint32) + state_0 = tl.full((), _CONTEXT_DOMAIN, tl.uint32) + state_1 = tl.full((), CONTEXT_WIDTH, tl.uint32) + state_2 = tl.full((), 0, tl.uint32) + state_3 = tl.full((), 0, tl.uint32) + for offset in range(0, CONTEXT_WIDTH, 4): + context_0 = tl.load(contexts_ptr + row * context_stride + offset).to(tl.uint32) + if offset + 1 < CONTEXT_WIDTH: + context_1 = tl.load(contexts_ptr + row * context_stride + offset + 1).to( + tl.uint32 + ) + else: + context_1 = tl.full((), _UINT32_MASK - 1, tl.uint32) + if offset + 2 < CONTEXT_WIDTH: + context_2 = tl.load(contexts_ptr + row * context_stride + offset + 2).to( + tl.uint32 + ) + else: + context_2 = tl.full((), _UINT32_MASK - 2, tl.uint32) + if offset + 3 < CONTEXT_WIDTH: + context_3 = tl.load(contexts_ptr + row * context_stride + offset + 3).to( + tl.uint32 + ) + else: + context_3 = tl.full((), _UINT32_MASK - 3, tl.uint32) + state_0, state_1, state_2, state_3 = _philox4x32_10( + (state_0 ^ context_0) & _UINT32_MASK, + (state_1 ^ context_1) & _UINT32_MASK, + (state_2 ^ context_2) & _UINT32_MASK, + (state_3 ^ context_3) & _UINT32_MASK, + (key_0 ^ offset) & _UINT32_MASK, + key_1, + ) + + candidate_words = groups.to(tl.uint32) + vector_zero = candidate_words * 0 + output_0, output_1, output_2, output_3 = _philox4x32_10( + candidate_words & _UINT32_MASK, + state_0 + vector_zero, + state_1 + vector_zero, + state_2 + vector_zero, + ((key_0 ^ state_3) & _UINT32_MASK) + vector_zero, + ((key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK) + vector_zero, + ) + candidate_0 = groups * 4 + candidate_1 = candidate_0 + 1 + candidate_2 = candidate_0 + 2 + candidate_3 = candidate_0 + 3 + logits_row = logits_ptr + row * logits_stride + value_0 = _gumbel_value( + logits_row + candidate_0, output_0, candidate_0 < vocab_size + ) + value_1 = _gumbel_value( + logits_row + candidate_1, output_1, candidate_1 < vocab_size + ) + value_2 = _gumbel_value( + logits_row + candidate_2, output_2, candidate_2 < vocab_size + ) + value_3 = _gumbel_value( + logits_row + candidate_3, output_3, candidate_3 < vocab_size + ) + best_value = value_0 + best_token = candidate_0 + best_value, best_token = _select_max(best_value, best_token, value_1, candidate_1) + best_value, best_token = _select_max(best_value, best_token, value_2, candidate_2) + best_value, best_token = _select_max(best_value, best_token, value_3, candidate_3) + value = tl.max(best_value, axis=0) + token_id = tl.min(tl.where(best_value == value, best_token, vocab_size), axis=0) + tl.store(local_argmax_ptr + row * local_argmax_stride + block_index, token_id) + tl.store(local_max_ptr + row * local_max_stride + block_index, value) + + +def philox_gumbel_sample( + logits: torch.Tensor, contexts: torch.Tensor, key: int +) -> torch.Tensor: + if logits.stride(-1) != 1: + logits = logits.contiguous() + if contexts.stride(-1) != 1: + contexts = contexts.contiguous() + num_tokens, vocab_size = logits.shape + block_size = 1024 + num_blocks = triton.cdiv(vocab_size, block_size) + local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64) + local_max = logits.new_empty(num_tokens, num_blocks, dtype=torch.float32) + _philox_gumbel_kernel[(num_tokens, num_blocks)]( + local_argmax, + local_argmax.stride(0), + local_max, + local_max.stride(0), + logits, + logits.stride(0), + contexts, + contexts.stride(0), + key & _UINT32_MASK_VALUE, + key >> 32, + vocab_size, + CONTEXT_WIDTH=contexts.shape[-1], + BLOCK_SIZE=block_size, + ) + max_block_index = local_max.argmax(dim=-1, keepdim=True) + return local_argmax.gather(dim=-1, index=max_block_index).view(-1)