[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>

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.

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())

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",
     [

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

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",
         }
 

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

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

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)

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)

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

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",
 ]

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()

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]

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,

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

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,

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,

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,

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]:

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}, "

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",
+]

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

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)

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)

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,
+        )

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",
+]

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)

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)

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

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(

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:

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)