[watermarking] Dual-key gumbel-max watermarking for speculative decoding support (#56122)

Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai>
Signed-off-by: Simon Veitner <sveitner@redhat.com>
Signed-off-by: Tomas Ruiz <tomas.ruiz.te@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Simon Veitner <sveitner@redhat.com>
Co-authored-by: Tomas Ruiz <tomas.ruiz.te@gmail.com>

.buildkite/test_areas/lm_eval.yaml

@@ -122,6 +122,28 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
 
+- label: ":nvidia: (H100) LM Eval Watermark Feature Combination"
+  key: lm-eval-dspark-watermark-2xh100
+  device: h100
+  num_devices: 2
+  # Local TP2 runs took 84–379 seconds, including startup.
+  timeout_in_minutes: 15
+  source_file_dependencies:
+  - .buildkite/test_areas/lm_eval.yaml
+  - tests/evals/gsm8k/
+  - vllm/config/speculative.py
+  - vllm/config/vllm.py
+  - vllm/config/watermarking.py
+  - vllm/sampling_params.py
+  - vllm/model_executor/models/qwen3_dspark.py
+  - vllm/model_executor/models/qwen3_dflash.py
+  - vllm/model_executor/warmup/
+  - vllm/v1/watermarking/
+  - vllm/v1/worker/gpu/
+  autorun_on_main: true
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
+
 - label: ":nvidia: (B200) LM Eval Spec Decode"
   key: lm-eval-spec-decode-4xb200
   timeout_in_minutes: 120

docs/features/watermarking.md

@@ -34,11 +34,17 @@ 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
+`context_width` controls how many prior 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.
 
+`allow_target_only_watermarking` defaults to false and only has an effect when
+speculative decoding is enabled. It permits speculative decoding with a
+watermarking algorithm that does not support it natively, at the cost of
+weaker detectability. See
+[Speculative decoding](#speculative-decoding).
+
 ## Architecture
 
 `WatermarkConfig` selects an algorithm and PRF. Model Runner V2 constructs the
@@ -52,21 +58,108 @@ the reference algorithms. `WatermarkDetector` consumes token IDs, so callers
 remain responsible for using the tokenizer and watermark profile that match
 generation.
 
+## Speculative decoding
+
+Watermarking requires speculative decoding to use probabilistic draft sampling,
+standard rejection sampling, and an autoregressive model-based method (`dspark`,
+`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dspark`.
+
+A watermarking algorithm without native speculative-decoding support is
+rejected before model loading. Set
+`"allow_target_only_watermarking": true` to allow it: accepted draft tokens are
+not watermarked, while target-side rejection recovery and bonus sampling remain
+watermarked. The watermark signal is diluted in proportion to the share of
+output tokens supplied by accepted drafts; rejected drafts do not dilute it
+because their recovery tokens are watermarked.
+
+Speculative-decoding token paths do not currently support generation-side
+context deduplication. The configured `deduplicate_contexts` policy is not
+applied to accepted drafts, rejection-recovery tokens, or bonus tokens.
+
+For `dual_key_gumbel`, `alpha` has no effect under speculative decoding. The
+speculative protocol selects the key for each token instead.
+
 ## 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
+Gumbel-max derives a deterministic pseudorandom value from the key, prior 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.
+When a token context is repeated, generation can use ordinary sampling for that
+occurrence. Reusing the repeated context leads to a bias over the sequence, as
+certain token choices would be correlated. Ordinary sampling at these positions
+allows for single-sequence non-distortion (see section G.3 of the
+[SynthID-Text supplementary materials](https://media.springernature.com/original/springer-static/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf)).
+The detector independently deduplicates contexts so repeated keyed random
+vectors are not treated as independent evidence, meaning that context
+deduplication at generation time does not reduce the watermarking signal, unless
+a user tampers with the output and removes the first occurrence while leaving an
+unwatermarked repetition.
+
+`deduplicate_contexts` controls which history is searched for a repeated
+context:
+
+- `"none"` disables context deduplication.
+- `"single_turn"`, the default, searches the tokens generated for the current
+  request. This gives single-turn non-distortion.
+- `"all"` also searches the prompt, which extends non-distortion across the
+  turns of a conversation. The first `context_width` generated tokens use
+  ordinary sampling, and subsequent tokens whose context already occurs in the
+  prompt are not watermarked. Previous turns of a session, or a prompt that
+  already contains the structure of the answer such as a tool result the model
+  extends, may therefore leave little of the answer watermarked. Given the
+  impact of `"all"` on both performance (sampling steps must scan over a
+  potentially large amount of previous context) and watermarking detectability
+  (we may significantly reduce the amount of watermarked tokens), you should
+  only use `"all"` when non-distortion across multiple turns is required.
+
+`deduplicate_contexts_max_history` limits the search to the most recent
+positions and defaults to 8,192. Each position is compared over the
+`context_width` tokens before it, so the window of tokens read is that much
+longer. A smaller value reduces scanning cost but only provides the guarantee
+within that window. Set it to `null` to search back to the start of the
+generation for `"single_turn"`, or of the request for `"all"`; an unbounded
+search costs more as the sequence grows. The setting has no effect when
+`deduplicate_contexts` is `"none"`. Values below 1,024 emit a warning because a
+short window can miss repetition loops whose contexts recur farther apart.
+
+For example, this checks prompt and completion history within the default
+8,192-position window:
+
+```bash
+vllm serve MODEL \
+  --watermark-config \
+  '{"algorithm":"gumbel","key":42,"deduplicate_contexts":"all"}'
+```
+
+### Dual-key Gumbel-max
+
+Dual-key Gumbel-max derives independent keys A and B from one configured master
+key. During ordinary generation, each token uses key A with probability
+`1 - alpha` and key B with probability `alpha`. The algorithm-specific `alpha`
+parameter defaults to 0.1. Detection scores every token against both keys.
+
+The same two key streams support speculative decoding without changing its
+acceptance rate. In this mode, the speculative protocol selects the key instead
+of `alpha`: draft tokens use key A, while rejection recovery and bonus tokens
+use key B. The ordinary target-to-draft probability-ratio test remains unchanged,
+implementing [SynthID-Text Supplementary Algorithm
+6](https://media.springernature.com/original/springer-static/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf).
+
+Select `dual_key_gumbel` together with probabilistic drafting:
+
+```bash
+vllm serve MODEL \
+  --speculative-config \
+  '{"method":"mtp","num_speculative_tokens":3,"draft_sample_method":"probabilistic"}' \
+  --watermark-config '{"algorithm":"dual_key_gumbel","key":42,"alpha":0.1}'
+```
 
 ### SynthID-Text
 
@@ -106,7 +199,14 @@ 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.
+The detection configuration must match the generation configuration, including
+the tokenizer, PRF, watermarking algorithm, algorithm-specific watermarking
+configuration, and key. In practice, this information is often unavailable
+when checking a piece of text. Deployments should therefore retain the set of
+candidate configurations they have served, test the text against each
+candidate, and correct for multiple testing, for example with a Bonferroni
+correction to the resulting p-values.
+
 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
@@ -140,7 +240,7 @@ 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.
+- Not all watermarking algorithms have native speculative-decoding support.
 - 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

tests/engine/test_arg_utils.py

@@ -53,18 +53,37 @@ def test_watermark_config_cli():
             "--model",
             "dummy",
             "--watermark-config",
-            '{"algorithm":"gumbel","key":42,"prf":"philox"}',
+            '{"algorithm":"dual_key_gumbel","key":42,"prf":"philox","alpha":0.25,'
+            '"allow_target_only_watermarking":true}',
         ]
     )
 
     config = EngineArgs.from_cli_args(args).create_watermark_config()
 
     assert config is not None
-    assert config.algorithm == "gumbel"
+    assert config.algorithm == "dual_key_gumbel"
     assert config.key == 42
+    assert config.alpha == 0.25
     assert config.context_width == 4
+    assert config.deduplicate_contexts == "single_turn"
+    assert config.deduplicate_contexts_max_history == 8192
     assert config.prf == "philox"
-    assert not config.supports_speculative_decoding
+    assert config.allow_target_only_watermarking
+
+    args = parser.parse_args(
+        [
+            "--model",
+            "dummy",
+            "--watermark-config",
+            '{"key":42,"deduplicate_contexts":"none",'
+            '"deduplicate_contexts_max_history":32}',
+        ]
+    )
+    config = EngineArgs.from_cli_args(args).create_watermark_config()
+
+    assert config is not None
+    assert config.deduplicate_contexts == "none"
+    assert config.deduplicate_contexts_max_history == 32
 
 
 @pytest.mark.parametrize(

tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml

@@ -0,0 +1,29 @@
+model_name: "Qwen/Qwen3-4B-FP8"
+num_questions: 1319
+num_fewshot: 5
+max_tokens: 256
+use_chat_completions: false
+temperature: 1.0
+seed: 42
+
+# Reference accuracy from the non-watermarked DSpark regression test.
+accuracy_threshold: 0.801
+tolerance: 0.02
+# Measured TP1 acceptance length: 4.115; allow a 0.2 regression.
+min_acceptance_length: 3.915
+
+env:
+  VLLM_USE_V2_MODEL_RUNNER: "1"
+server_args: >-
+  --tensor-parallel-size 2
+  --max-model-len 4096
+  --gpu-memory-utilization 0.92
+  --no-enable-prefix-caching
+  --speculative-config.method dspark
+  --speculative-config.model deepseek-ai/dspark_qwen3_4b_block7
+  --speculative-config.num_speculative_tokens 7
+  --speculative-config.draft_sample_method probabilistic
+  --watermark-config.algorithm dual_key_gumbel
+  --watermark-config.key 42
+  --watermark-config.context_width 4
+  --watermark-config.deduplicate_contexts none

tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt

@@ -0,0 +1 @@
+Qwen3-4B-FP8-DSpark-watermark-TP2.yaml

tests/model_executor/test_spec_decode_rejection_warmup.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.config.watermarking import WatermarkConfig, derive_watermark_key
+from vllm.model_executor.warmup.spec_decode_rejection_warmup import (
+    spec_decode_rejection_warmup,
+)
+
+MASTER_KEY = 20250910
+
+
+def _worker(
+    watermark_config: WatermarkConfig | None,
+    draft_sample_method: str = "greedy",
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        device=torch.device("cpu"),
+        vllm_config=SimpleNamespace(
+            speculative_config=SimpleNamespace(
+                num_speculative_tokens=2,
+                rejection_sample_method="standard",
+                draft_sample_method=draft_sample_method,
+            ),
+            model_config=SimpleNamespace(
+                get_vocab_size=lambda: 128, dtype=torch.bfloat16
+            ),
+            watermark_config=watermark_config,
+        ),
+    )
+
+
+def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]:
+    calls: list[dict] = []
+
+    def fake_rejection_sample(**kwargs):
+        calls.append(kwargs)
+        return None, None
+
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils.rejection_sample",
+        fake_rejection_sample,
+    )
+    return calls
+
+
+def test_warmup_token_id_dtypes_match_the_runtime_buffers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    spec_decode_rejection_warmup(_worker(None))
+
+    assert calls
+    for kwargs in calls:
+        assert kwargs["draft_sampled"].dtype == torch.int32
+        assert kwargs["pos"].dtype == torch.int64
+
+
+def test_warmup_index_mapping_dtypes_match_the_runtime_buffers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    spec_decode_rejection_warmup(_worker(None))
+
+    assert calls
+    for kwargs in calls:
+        assert kwargs["idx_mapping"].dtype == torch.int64
+        assert kwargs["expanded_idx_mapping"].dtype == torch.int64
+        assert kwargs["expanded_local_pos"].dtype == torch.int32
+
+
+def test_warmup_watermark_args_match_the_runtime_kernel_arguments(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    watermark_config = WatermarkConfig(
+        key=MASTER_KEY, algorithm="dual_key_gumbel", context_width=3
+    )
+    calls = _capture_calls(monkeypatch)
+    spec_decode_rejection_warmup(
+        _worker(watermark_config, draft_sample_method="probabilistic")
+    )
+
+    assert calls
+    key_b = derive_watermark_key(MASTER_KEY, b"key_b")
+    for kwargs in calls:
+        contexts = kwargs["contexts"]
+        assert contexts.dtype == torch.int32
+        assert contexts.shape == (3, 3) and contexts.stride() == (3, 1)
+        assert kwargs["watermarking"].dtype == torch.bool
+        assert kwargs["watermark_key"] == key_b
+        assert key_b & 0xFFFFFFFF >= 2**31 and key_b >> 32 >= 2**31
+
+
+def test_warmup_is_skipped_when_the_recovery_key_cannot_be_derived(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    monkeypatch.setattr(
+        "vllm.v1.watermarking.spec_decode.speculative_target_watermark_key",
+        lambda _: (_ for _ in ()).throw(NotImplementedError("no philox")),
+    )
+    spec_decode_rejection_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
+
+    assert calls == []
+
+
+def test_warmup_draft_logits_match_the_runtime_draft_sample_method(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    spec_decode_rejection_warmup(_worker(None))
+    assert calls
+    assert all(kwargs["draft_logits"] is None for kwargs in calls)
+
+    calls = _capture_calls(monkeypatch)
+    spec_decode_rejection_warmup(_worker(None, draft_sample_method="probabilistic"))
+    assert calls
+    assert all(kwargs["draft_logits"] is not None for kwargs in calls)

tests/model_executor/test_watermark_sample_warmup.py

@@ -0,0 +1,300 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The warmup must cover every specialization the sampler will launch.
+
+``_philox_gumbel_kernel`` is keyed on the Philox key halves, on the logits
+dtype (model dtype when the sampler passes logits through, fp32 when logits
+processing copies them) and on whether a skip mask is passed, so a missing
+combination leaves a request paying the compile this module exists to remove.
+"""
+
+import logging
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.config.watermarking import WatermarkConfig, derive_watermark_key
+from vllm.model_executor.warmup.watermark_sample_warmup import watermark_sample_warmup
+from vllm.platforms import current_platform
+from vllm.v1.watermarking.factory import create_watermarker
+from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
+from vllm.v1.watermarking.spec_decode import create_speculative_target_watermarker
+
+MASTER_KEY = 20250910
+KEY_A = derive_watermark_key(MASTER_KEY, b"key_a")
+KEY_B = derive_watermark_key(MASTER_KEY, b"key_b")
+
+
+@pytest.fixture(autouse=True)
+def cuda_alike_platform(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The warmup returns early off CUDA, so pin the platform it reads.
+
+    Everything here runs on the fake worker's cpu device; without this the
+    tests below would pass only on a CUDA-alike host and silently warm nothing
+    anywhere else.
+    """
+    monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
+
+
+def _sampler(
+    watermark_config: WatermarkConfig,
+    speculative: bool,
+    use_fp64_gumbel: bool,
+) -> GPUWatermarkSampler:
+    """The sampler ``GPUModelRunner`` builds for this config.
+
+    Mirrors ``vllm/v1/worker/gpu/model_runner.py:473-486``. ``__new__`` skips
+    ``Sampler.__init__``, which allocates UVA-backed device buffers; the warmup
+    only reads the three attributes set here.
+    """
+    watermarker = create_watermarker(watermark_config)
+    if speculative:
+        watermarker = create_speculative_target_watermarker(watermarker)
+    sampler = GPUWatermarkSampler.__new__(GPUWatermarkSampler)
+    sampler.watermarker = watermarker
+    sampler.deduplicate_contexts = watermark_config.deduplicate_contexts
+    sampler.use_fp64_gumbel = use_fp64_gumbel
+    return sampler
+
+
+def _worker(
+    watermark_config: WatermarkConfig | None,
+    speculative: bool = False,
+    use_fp64_gumbel: bool = False,
+    with_sampler: bool = True,
+) -> SimpleNamespace:
+    sampler = (
+        _sampler(watermark_config, speculative, use_fp64_gumbel)
+        if watermark_config is not None and with_sampler
+        else None
+    )
+    return SimpleNamespace(
+        device=torch.device("cpu"),
+        model_runner=SimpleNamespace(sampler=sampler),
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                get_vocab_size=lambda: 128,
+                dtype=torch.bfloat16,
+            ),
+            watermark_config=watermark_config,
+        ),
+    )
+
+
+def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]:
+    calls: list[dict] = []
+
+    def fake_philox_gumbel_sample(logits, contexts, key, **kwargs):
+        calls.append({"logits": logits, "contexts": contexts, "key": key, **kwargs})
+        return torch.zeros(logits.shape[0], dtype=torch.int64)
+
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample",
+        fake_philox_gumbel_sample,
+    )
+    return calls
+
+
+def _variants(calls: list[dict]) -> set[tuple[int, torch.dtype, bool]]:
+    return {
+        (call["key"], call["logits"].dtype, call.get("skip_mask") is not None)
+        for call in calls
+    }
+
+
+def test_no_warmup_without_a_watermark_config(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(_worker(None))
+
+    assert calls == []
+
+
+def test_no_warmup_without_a_watermarked_sampler(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A rank that does not sample launches nothing to warm.
+
+    Only the final pipeline stage of a generative model holds a
+    ``GPUWatermarkSampler``; on any other rank the kernel is never launched.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(WatermarkConfig(key=MASTER_KEY), with_sampler=False)
+    )
+
+    assert calls == []
+
+
+def test_no_warmup_off_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    """``GumbelWatermarker`` launches the kernel only on a CUDA-alike device.
+
+    Elsewhere -- an XPU worker inherits ``kernel_warmup`` from the CUDA worker
+    -- the sampler takes the torch path, so warming would compile a kernel
+    nothing launches.
+    """
+    calls = _capture_calls(monkeypatch)
+    monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: False)
+    watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
+
+    assert calls == []
+
+
+def test_single_key_warms_both_logits_dtypes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
+
+    assert _variants(calls) == {
+        (MASTER_KEY, torch.bfloat16, True),
+        (MASTER_KEY, torch.float32, True),
+    }
+    assert len(calls) == len(_variants(calls))
+
+
+def test_dual_key_warms_both_derived_keys(monkeypatch: pytest.MonkeyPatch) -> None:
+    """``DualKeyGumbelWatermarker.sample`` launches the kernel once per key.
+
+    The halves of a derived key exceed 2**31, so each key is typed i64/i32 by
+    magnitude and a wrong key compiles a specialization the engine never uses.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel"))
+    )
+
+    assert _variants(calls) == {
+        (KEY_A, torch.bfloat16, True),
+        (KEY_A, torch.float32, True),
+        (KEY_B, torch.bfloat16, True),
+        (KEY_B, torch.float32, True),
+    }
+
+
+@pytest.mark.parametrize("alpha,key", [(0.0, KEY_A), (1.0, KEY_B)])
+def test_dual_key_alpha_bounds_warm_one_key(
+    monkeypatch: pytest.MonkeyPatch, alpha: float, key: int
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(
+            WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel", alpha=alpha)
+        )
+    )
+
+    assert _variants(calls) == {
+        (key, torch.bfloat16, True),
+        (key, torch.float32, True),
+    }
+
+
+def test_speculative_config_warms_the_target_role_only(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Under spec decode the sampler holds ``create_target_watermarker()``.
+
+    Rows without draft tokens sample through it, so the sampler path launches
+    key B alone; the draft key runs in the speculator's own sampler.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(
+            WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel"),
+            speculative=True,
+        )
+    )
+
+    assert _variants(calls) == {
+        (KEY_B, torch.bfloat16, True),
+        (KEY_B, torch.float32, True),
+    }
+
+
+def test_dedup_none_also_warms_the_mask_free_specialization(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Without deduplication an all-watermarked batch passes no skip mask.
+
+    A later batch mixing in an opted-out request passes one, so both
+    specializations are reachable and both must be warmed.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none"))
+    )
+
+    assert _variants(calls) == {
+        (MASTER_KEY, torch.bfloat16, True),
+        (MASTER_KEY, torch.float32, True),
+        (MASTER_KEY, torch.bfloat16, False),
+        (MASTER_KEY, torch.float32, False),
+    }
+    for call in calls:
+        if call.get("skip_mask") is None:
+            assert "temperatures" not in call
+
+
+def test_warmup_argument_dtypes_match_the_runtime_buffers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Contexts are int32, the sampling state int64 indices and fp32 temperatures.
+
+    ``GPUWatermarkSampler._get_contexts`` gathers from the int32 request-state
+    token ids, ``expanded_idx_mapping``/``positions``/``seeds`` are int64 and
+    ``SamplingStates.temperature`` is fp32; warming other dtypes compiles
+    specializations no engine launches.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY, context_width=3)))
+
+    assert calls
+    for call in calls:
+        num_tokens = call["logits"].shape[0]
+        assert call["contexts"].dtype == torch.int32
+        assert call["contexts"].shape == (num_tokens, 3)
+        assert call["skip_mask"].dtype == torch.bool
+        assert call["expanded_idx_mapping"].dtype == torch.int64
+        assert call["positions"].dtype == torch.int64
+        assert call["seeds"].dtype == torch.int64
+        assert call["temperatures"].dtype == torch.float32
+
+
+def test_use_fp64_follows_the_path_that_forwards_it(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Only the mixed path forwards the sampler's ``use_fp64_gumbel``.
+
+    ``GumbelWatermarker._sample_watermarked`` launches the mask-free variant
+    with the ``use_fp64=False`` default, so warming it with fp64 would compile
+    a specialization the engine never launches and miss the one it does.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(
+            WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none"),
+            use_fp64_gumbel=True,
+        )
+    )
+
+    assert calls
+    for call in calls:
+        assert call["use_fp64"] is (call.get("skip_mask") is not None)
+
+
+def test_warmup_failure_does_not_abort_startup(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    def failing_philox_gumbel_sample(*args, **kwargs):
+        raise RuntimeError("no triton")
+
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample",
+        failing_philox_gumbel_sample,
+    )
+    with caplog.at_level(logging.WARNING):
+        watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
+
+    assert "Skipping watermark sampler warmup." in caplog.text

tests/test_config.py

@@ -2856,14 +2856,183 @@ def _watermarked_vllm_config() -> VllmConfig:
     return config
 
 
-def test_gumbel_watermark_rejects_speculative_decoding():
+def test_target_only_gumbel_allows_speculative_decoding(caplog_vllm, disable_log_dedup):
     config = _watermarked_vllm_config()
-    config.speculative_config = SpeculativeConfig(
-        method="ngram",
-        num_speculative_tokens=1,
+    config.watermark_config = WatermarkConfig(
+        key=42, allow_target_only_watermarking=True
+    )
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
     )
 
-    with pytest.raises(ValueError, match="does not support speculative decoding"):
+    with caplog_vllm.at_level(logging.WARNING):
+        config._check_watermarking_unsupported()
+
+    assert "Target-only watermarking leaves accepted draft tokens" in caplog_vllm.text
+    assert "Context deduplication is not supported" in caplog_vllm.text
+
+
+def test_speculative_watermarking_without_context_dedup_does_not_warn(
+    caplog_vllm, disable_log_dedup
+):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(
+        algorithm="dual_key_gumbel", key=42, deduplicate_contexts="none"
+    )
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with caplog_vllm.at_level(logging.WARNING):
+        config._check_watermarking_unsupported()
+
+    assert "Context deduplication is not supported" not in caplog_vllm.text
+
+
+def test_gumbel_rejects_speculative_decoding_without_target_only():
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(
+        algorithm="gumbel", key=42, allow_target_only_watermarking=False
+    )
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with pytest.raises(ValueError, match="'gumbel'.*allow_target_only_watermarking"):
+        config._check_watermarking_unsupported()
+
+
+def test_dual_key_gumbel_warns_that_configured_alpha_is_unused(
+    caplog_vllm, disable_log_dedup
+):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(
+        algorithm="dual_key_gumbel", key=42, alpha=0.25
+    )
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with caplog_vllm.at_level(logging.WARNING):
+        config._check_watermarking_unsupported()
+
+    assert "The configured alpha=0.25 is not used" in caplog_vllm.text
+
+
+@pytest.mark.parametrize(
+    ("alpha", "speculative"),
+    [(0.1, True), (0.25, False)],
+    ids=["default-alpha", "no-specdec"],
+)
+def test_dual_key_gumbel_alpha_warning_is_not_emitted(
+    caplog_vllm, disable_log_dedup, alpha, speculative
+):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(
+        algorithm="dual_key_gumbel", key=42, alpha=alpha
+    )
+    if speculative:
+        config.speculative_config = SimpleNamespace(
+            method="mtp",
+            draft_sample_method="probabilistic",
+            rejection_sample_method="standard",
+            parallel_drafting=False,
+        )
+
+    with caplog_vllm.at_level(logging.WARNING):
+        config._check_watermarking_unsupported()
+
+    assert "is not used" not in caplog_vllm.text
+
+
+def test_dual_key_gumbel_requires_probabilistic_drafting():
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="greedy",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with pytest.raises(ValueError, match="draft_sample_method='probabilistic'"):
+        config._check_watermarking_unsupported()
+
+
+@pytest.mark.parametrize("method", ["eagle", "eagle3", "mtp"])
+def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(method):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    config.speculative_config = SimpleNamespace(
+        method=method,
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    config._check_watermarking_unsupported()
+
+
+def test_dual_key_gumbel_supports_dspark():
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    config.speculative_config = SimpleNamespace(
+        method="dspark",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=True,
+    )
+
+    config._check_watermarking_unsupported()
+
+
+def test_dual_key_gumbel_rejects_non_autoregressive_speculation():
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    config.speculative_config = SimpleNamespace(
+        method="ngram",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with pytest.raises(ValueError, match="autoregressive model-based"):
+        config._check_watermarking_unsupported()
+
+
+@pytest.mark.parametrize(
+    ("overrides", "match"),
+    [
+        ({"rejection_sample_method": "synthetic"}, "rejection_sample_method"),
+        ({"rejection_sample_method": "block"}, "rejection_sample_method"),
+        ({"parallel_drafting": True}, "Parallel speculative drafting"),
+    ],
+)
+def test_dual_key_gumbel_rejects_incompatible_speculative_modes(overrides, match):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    values = {
+        "method": "mtp",
+        "draft_sample_method": "probabilistic",
+        "rejection_sample_method": "standard",
+        "parallel_drafting": False,
+        **overrides,
+    }
+    config.speculative_config = SimpleNamespace(**values)
+
+    with pytest.raises(ValueError, match=match):
         config._check_watermarking_unsupported()
 
 
@@ -2882,6 +3051,12 @@ def test_watermark_key_must_fit_in_64_bits():
         WatermarkConfig(key=2**64)
 
 
+@pytest.mark.parametrize("alpha", [-0.1, 1.1])
+def test_watermark_alpha_must_be_a_probability(alpha):
+    with pytest.raises(ValidationError):
+        WatermarkConfig(key=42, algorithm="dual_key_gumbel", alpha=alpha)
+
+
 def test_unknown_watermark_prf_is_rejected():
     with pytest.raises(ValidationError):
         pydantic.TypeAdapter(WatermarkConfig).validate_python(

tests/v1/spec_decode/test_rejection_sampler_utils.py

@@ -6,7 +6,10 @@ import math
 import pytest
 import torch
 
+from vllm.v1.watermarking import GumbelWatermarker
+from vllm.v1.watermarking.spec_decode import watermarked_rejection_sample
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
     rejection_sample,
 )
@@ -139,6 +142,93 @@ def _assert_distribution_match(
     )
 
 
+def test_watermarked_bonus_uses_target_key():
+    target_logits = torch.randn(16, device="cuda")
+    inputs = _build_rejection_sample_inputs(
+        target_logits,
+        target_logits,
+        num_speculative_steps=1,
+        temperature=1.0,
+        num_trials=1,
+    )
+    draft_token = inputs["draft_sampled"][1]
+    contexts = torch.tensor(
+        [[10, 11], [11, draft_token.item()]], dtype=torch.int64, device="cuda"
+    )
+
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=1,
+        contexts=contexts,
+        watermarking=torch.tensor([True], device="cuda"),
+        watermarker=GumbelWatermarker(key=42, context_width=2),
+    )
+    expected = GumbelWatermarker(key=42, context_width=2).sample(
+        target_logits.unsqueeze(0), contexts[1:], lambda _: None
+    )
+
+    assert num_sampled.item() == 2
+    assert sampled[0, 1] == expected.token_ids[0]
+
+
+def test_watermarked_recovery_uses_target_key():
+    target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0], device="cuda")
+    draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda")
+    inputs = _build_rejection_sample_inputs(
+        target_logits,
+        draft_logits,
+        num_speculative_steps=1,
+        temperature=1.0,
+        num_trials=1,
+    )
+    inputs["draft_sampled"][1] = 0
+    contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda")
+
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=1,
+        contexts=contexts,
+        watermarking=torch.tensor([True], device="cuda"),
+        watermarker=GumbelWatermarker(key=42, context_width=2),
+    )
+    expected = GumbelWatermarker(key=42, context_width=2).sample(
+        target_logits.unsqueeze(0), contexts[:1], lambda _: None
+    )
+
+    assert num_sampled.item() == 1
+    assert sampled[0, 0] == expected.token_ids[0]
+
+
+def test_watermarked_recovery_supports_smaller_draft_vocabulary():
+    target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0, 100.0], device="cuda")
+    draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda")
+    inputs = _build_rejection_sample_inputs(
+        target_logits[:4],
+        draft_logits,
+        num_speculative_steps=1,
+        temperature=1.0,
+        num_trials=1,
+    )
+    inputs["target_logits"] = target_logits.unsqueeze(0).expand(2, -1).contiguous()
+    inputs["draft_sampled"][1] = 0
+    contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda")
+    watermarker = GumbelWatermarker(key=42, context_width=2)
+
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=1,
+        contexts=contexts,
+        watermarking=torch.tensor([True], device="cuda"),
+        watermarker=watermarker,
+    )
+    expected = watermarker.sample(
+        target_logits[:4].unsqueeze(0), contexts[:1], lambda _: None
+    )
+
+    assert num_sampled.item() == 1
+    assert sampled[0, 0] == expected.token_ids[0]
+
+
 @pytest.mark.parametrize(
     "num_speculative_steps,temperature",
     [
@@ -747,3 +837,192 @@ def test_chunked_requests_match_full_batch(has_draft_logits: bool):
     steps = torch.arange(num_speculative_steps + 1, device=device)
     valid = steps.unsqueeze(0) < num_sampled.unsqueeze(1)
     assert torch.equal(chunked_sampled[valid], sampled[valid])
+
+
+def _residual_recovery_logits(
+    inputs: dict,
+    num_sampled: torch.Tensor,
+    vocab_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    target_logits = inputs["target_logits"]
+    draft_logits = inputs["draft_logits"]
+    cu_num_logits = inputs["cu_num_logits"]
+    resample_rows = cu_num_logits[:-1].to(torch.int64) + num_sampled.to(torch.int64) - 1
+    req_state_indices = inputs["expanded_idx_mapping"][resample_rows].to(torch.int64)
+    temperatures = inputs["temperature"][req_state_indices]
+
+    is_bonus = resample_rows == cu_num_logits[1:] - 1
+    next_rows = (resample_rows + 1).clamp_max(inputs["draft_sampled"].shape[0] - 1)
+    needs_residual = ~is_bonus & (inputs["draft_sampled"][next_rows] >= 0)
+
+    target_rows = target_logits[resample_rows][:, :vocab_size].float()
+    target_rows = torch.where(target_rows.isnan(), float("-inf"), target_rows)
+    if draft_logits is None:
+        rejected = inputs["draft_sampled"][next_rows].clamp_min(0)
+        one_hot = torch.arange(vocab_size, device=target_rows.device).unsqueeze(
+            0
+        ) == rejected.unsqueeze(-1)
+        residual_logits = torch.where(one_hot, float("-inf"), target_rows)
+    else:
+        draft_steps = (num_sampled.to(torch.int64) - 1).clamp_max(
+            draft_logits.shape[1] - 1
+        )
+        draft_rows = draft_logits[req_state_indices, draft_steps][
+            :, :vocab_size
+        ].float()
+        draft_rows = torch.where(draft_rows.isnan(), float("-inf"), draft_rows)
+        draft_rows = draft_rows / temperatures.unsqueeze(-1)
+        target_log_probs = torch.log_softmax(target_rows, dim=-1)
+        draft_log_probs = torch.log_softmax(draft_rows, dim=-1)
+        ratio = torch.exp(draft_log_probs - target_log_probs)
+        residual_logits = torch.where(
+            ratio < 1,
+            target_log_probs + torch.log1p(-ratio.clamp_max(1)),
+            float("-inf"),
+        )
+    recovery = torch.where(needs_residual.unsqueeze(-1), residual_logits, target_rows)
+    return recovery, resample_rows
+
+
+def _watermark_inputs(
+    num_trials: int,
+    num_speculative_steps: int,
+    context_width: int,
+    vocab_size: int,
+    one_hot_draft: bool,
+    temperature: float = 1.0,
+) -> tuple[dict, torch.Tensor]:
+    device = "cuda"
+    target_logits_1d = torch.randn(vocab_size, device=device) / temperature
+    draft_logits_1d = torch.randn(vocab_size, device=device)
+    inputs = _build_rejection_sample_inputs(
+        target_logits_1d,
+        draft_logits_1d,
+        num_speculative_steps,
+        temperature=temperature,
+        num_trials=num_trials,
+    )
+    if one_hot_draft:
+        inputs["draft_logits"] = None
+    inputs["draft_sampled"].view(num_trials, num_speculative_steps + 1)[::5, 1] = -1
+    num_logits = inputs["target_logits"].shape[0]
+    contexts = torch.randint(
+        0, vocab_size, (num_logits, context_width), dtype=torch.int64, device=device
+    )
+    contexts[::7, 0] = -1
+    return inputs, contexts
+
+
+@pytest.mark.parametrize("one_hot_draft", [False, True])
+# Cover one full block and a partial multi-block vocabulary.
+@pytest.mark.parametrize("vocab_size", [1024, 3500])
+def test_watermarked_recovery_matches_philox_gumbel_sample(
+    one_hot_draft: bool, vocab_size: int
+):
+    torch.manual_seed(7)
+    key = 0xDA39A3EE5E6B4B0D
+    context_width = 4
+    inputs, contexts = _watermark_inputs(
+        num_trials=128,
+        num_speculative_steps=2,
+        context_width=context_width,
+        vocab_size=vocab_size,
+        one_hot_draft=one_hot_draft,
+    )
+    watermarker = GumbelWatermarker(key=key, context_width=context_width)
+
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=2,
+        contexts=contexts,
+        watermarking=torch.ones(128, dtype=torch.bool, device="cuda"),
+        watermarker=watermarker,
+    )
+
+    recovery, resample_rows = _residual_recovery_logits(inputs, num_sampled, vocab_size)
+    expected = philox_gumbel_sample(recovery, contexts[resample_rows], key)
+    recovered = sampled[torch.arange(128, device="cuda"), num_sampled.long() - 1]
+    torch.testing.assert_close(recovered, expected, rtol=0, atol=0)
+
+    is_bonus = resample_rows == inputs["cu_num_logits"][1:] - 1
+    next_rows = (resample_rows + 1).clamp_max(inputs["draft_sampled"].shape[0] - 1)
+    is_placeholder = ~is_bonus & (inputs["draft_sampled"][next_rows] < 0)
+    assert int((~is_bonus & ~is_placeholder).sum()) > 0
+    assert int(is_placeholder.sum()) > 0
+    if not one_hot_draft:
+        # One-hot drafts almost never reach the bonus row.
+        assert int(is_bonus.sum()) > 0
+
+
+def test_watermarked_recovery_leaves_disabled_rows_on_the_stock_draw():
+    torch.manual_seed(11)
+    num_trials = 96
+    vocab_size = 512
+    inputs, contexts = _watermark_inputs(
+        num_trials=num_trials,
+        num_speculative_steps=3,
+        context_width=2,
+        vocab_size=vocab_size,
+        one_hot_draft=False,
+    )
+    watermarking = torch.zeros(num_trials, dtype=torch.bool, device="cuda")
+    watermarking[::2] = True
+    inputs["temperature"][:8] = 0.0
+
+    stock_sampled, stock_num_sampled = rejection_sample(
+        **inputs, num_speculative_steps=3
+    )
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=3,
+        contexts=contexts,
+        watermarking=watermarking,
+        watermarker=GumbelWatermarker(key=99, context_width=2),
+    )
+
+    torch.testing.assert_close(num_sampled, stock_num_sampled, rtol=0, atol=0)
+    rows = torch.arange(num_trials, device="cuda")
+    positions = num_sampled.long() - 1
+    disabled = ~watermarking | (inputs["temperature"] == 0)
+    torch.testing.assert_close(
+        sampled[rows[disabled], positions[disabled]],
+        stock_sampled[rows[disabled], positions[disabled]],
+        rtol=0,
+        atol=0,
+    )
+    for req in range(num_trials):
+        n = int(num_sampled[req]) - 1
+        torch.testing.assert_close(
+            sampled[req, :n], stock_sampled[req, :n], rtol=0, atol=0
+        )
+    assert int(disabled.sum()) > 0 and int((~disabled).sum()) > 0
+
+
+def test_watermarked_recovery_is_unchanged_by_fp64_gumbel():
+    torch.manual_seed(13)
+    inputs, contexts = _watermark_inputs(
+        num_trials=64,
+        num_speculative_steps=2,
+        context_width=4,
+        vocab_size=2048,
+        one_hot_draft=False,
+    )
+    watermarker = GumbelWatermarker(key=2**63 + 12345, context_width=4)
+    kwargs = dict(
+        num_speculative_steps=2,
+        contexts=contexts,
+        watermarking=torch.ones(64, dtype=torch.bool, device="cuda"),
+        watermarker=watermarker,
+    )
+    sampled32, num_sampled32 = watermarked_rejection_sample(**inputs, **kwargs)
+    sampled64, num_sampled64 = watermarked_rejection_sample(
+        **inputs, **kwargs, use_fp64=True
+    )
+    rows = torch.arange(64, device="cuda")
+    torch.testing.assert_close(num_sampled64, num_sampled32, rtol=0, atol=0)
+    torch.testing.assert_close(
+        sampled64[rows, num_sampled64.long() - 1],
+        sampled32[rows, num_sampled32.long() - 1],
+        rtol=0,
+        atol=0,
+    )

tests/watermarking/test_gumbel.py

@@ -5,8 +5,15 @@ import pytest
 import torch
 
 from vllm.platforms import current_platform
-from vllm.v1.watermarking import GumbelWatermarkDetector, GumbelWatermarker
+from vllm.v1.watermarking import (
+    DualKeyGumbelWatermarkDetector,
+    GumbelWatermarkDetector,
+    GumbelWatermarker,
+    derive_watermark_key,
+)
 from vllm.v1.watermarking.gumbel import _gamma_survival_integer_shape
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
 
 
 def test_gamma_survival_integer_shape():
@@ -34,6 +41,67 @@ def test_detector_deduplicates_context_even_when_target_differs():
     assert detection.num_scored_tokens == 3
 
 
+def test_dual_key_detector_scores_each_token_against_both_keys():
+    token_ids = [1, 2, 3, 4, 5]
+    key_a_detector = GumbelWatermarkDetector(
+        key=derive_watermark_key(42, b"key_a"),
+        context_width=1,
+        deduplicate_contexts=False,
+    )
+    key_b_detector = GumbelWatermarkDetector(
+        key=derive_watermark_key(42, b"key_b"),
+        context_width=1,
+        deduplicate_contexts=False,
+    )
+    detector = DualKeyGumbelWatermarkDetector(
+        key=42,
+        context_width=1,
+        deduplicate_contexts=False,
+        alpha=0.5,
+    )
+
+    key_a = key_a_detector.detect(token_ids)
+    key_b = key_b_detector.detect(token_ids)
+    dual = detector.detect(token_ids)
+
+    assert dual.score == pytest.approx((key_a.score + key_b.score) / 2)
+    assert dual.p_value == pytest.approx(
+        _gamma_survival_integer_shape(dual.score * 2, dual.num_scored_tokens * 2)
+    )
+
+
+@pytest.mark.parametrize("alpha", [0.2, 0.7])
+def test_dual_key_detector_weights_and_recalibrates_scores(alpha: float):
+    token_ids = [1, 2, 3, 4, 5]
+    key_a = GumbelWatermarkDetector(
+        key=derive_watermark_key(42, b"key_a"), deduplicate_contexts=False
+    ).detect(token_ids)
+    key_b = GumbelWatermarkDetector(
+        key=derive_watermark_key(42, b"key_b"), deduplicate_contexts=False
+    ).detect(token_ids)
+    dual = DualKeyGumbelWatermarkDetector(
+        key=42, deduplicate_contexts=False, alpha=alpha
+    ).detect(token_ids)
+
+    variance = (1 - alpha) ** 2 + alpha**2
+    expected_p_value = torch.special.gammaincc(
+        torch.tensor(dual.num_scored_tokens / variance, dtype=torch.float64),
+        torch.tensor(dual.score / variance, dtype=torch.float64),
+    ).item()
+
+    assert dual.score == pytest.approx((1 - alpha) * key_a.score + alpha * key_b.score)
+    assert dual.p_value == pytest.approx(expected_p_value)
+
+
+def test_dual_key_detector_rejects_invalid_alpha():
+    with pytest.raises(ValueError, match="alpha must be between 0 and 1"):
+        DualKeyGumbelWatermarkDetector(key=42, alpha=1.1)
+
+
+def test_dual_key_detector_default_alpha():
+    assert DualKeyGumbelWatermarkDetector(key=42).alpha == 0.2
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
 )
@@ -45,10 +113,8 @@ def test_fused_watermarker_matches_cpu(key: int, context_width: int):
     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()
+    expected = watermarker.sample(logits, contexts).token_ids
+    actual = watermarker.sample(logits.cuda(), contexts.cuda()).token_ids.cpu()
 
     torch.testing.assert_close(actual, expected, rtol=0, atol=0)
 
@@ -61,7 +127,7 @@ def test_fused_watermarker_handles_nan_logits():
     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
+    token_ids = watermarker.sample(logits, contexts).token_ids
 
     assert torch.all((token_ids >= 0) & (token_ids < logits.shape[-1]))
 
@@ -76,9 +142,58 @@ def test_fused_watermarker_handles_noncontiguous_inputs():
     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
+    expected = watermarker.sample(logits.contiguous(), contexts.contiguous()).token_ids
+    actual = watermarker.sample(logits, contexts).token_ids
+
+    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"
+)
+@pytest.mark.parametrize("use_fp64", [False, True])
+def test_philox_gumbel_sample_skip_mask_matches_separate_samplers(use_fp64: bool):
+    torch.manual_seed(0)
+    logits = torch.randn(8, 8193, device="cuda")
+    context_storage = torch.randint(0, 248320, (8, 8), dtype=torch.int64, device="cuda")
+    contexts = context_storage[:, ::2]
+    repeated_mask = torch.tensor(
+        [False, False, False, True, False, False, True, False], device="cuda"
+    )
+    watermarking = torch.tensor(
+        [True, False, True, True, False, True, True, False], device="cuda"
+    )
+    req_indices = torch.tensor([3, 1, 7, 0, 4, 2, 6, 5], device="cuda")
+    temperatures = torch.ones(8, dtype=torch.float32, device="cuda")
+    temperatures[2] = 0
+    seeds = torch.arange(8, dtype=torch.int64, device="cuda") + 1000
+    positions = torch.arange(8, dtype=torch.int64, device="cuda") + 100
+    watermark_mask = (
+        watermarking[req_indices] & (temperatures[req_indices] != 0) & ~repeated_mask
+    )
+
+    watermarked = philox_gumbel_sample(logits, contexts, 42)
+    ordinary = gumbel_sample(
+        logits,
+        req_indices,
+        temperatures,
+        seeds,
+        positions,
+        apply_temperature=False,
+        is_drafting=False,
+        use_fp64=use_fp64,
+    )
+    expected = torch.where(watermark_mask, watermarked, ordinary)
+    actual = philox_gumbel_sample(
+        logits,
+        contexts,
+        42,
+        skip_mask=~watermark_mask,
+        expanded_idx_mapping=req_indices,
+        temperatures=temperatures,
+        seeds=seeds,
+        positions=positions,
+        use_fp64=use_fp64,
+    )
 
     torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/watermarking/test_watermarking.py

@@ -9,44 +9,98 @@ import torch
 
 from vllm import SamplingParams
 from vllm.config.watermarking import WatermarkConfig
-from vllm.v1.watermarking import create_watermarker
+from vllm.platforms import current_platform
+from vllm.v1.watermarking import (
+    DualKeyGumbelWatermarker,
+    SupportsSpeculativeDecoding,
+    create_watermarker,
+    derive_watermark_key,
+)
 from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
-from vllm.v1.watermarking.watermarker import WatermarkSample
+from vllm.v1.watermarking.gumbel import GumbelWatermarker
+from vllm.v1.watermarking.spec_decode import (
+    DraftWatermarker,
+    _resolve_watermark_key,
+    create_speculative_draft_watermarker,
+    create_speculative_target_watermarker,
+    speculative_target_watermark_key,
+)
+from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample
 from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.sample.watermark import (
+    philox_gumbel_sample,
+    repeated_context_mask,
+)
+from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
 
 
-@pytest.mark.parametrize("algorithm", ["gumbel"])
+class StubWatermarker(Watermarker):
+    context_width = 1
+
+    def _sample_watermarked(self, logits, contexts):
+        return WatermarkSample(torch.tensor([7, 7]), logits + 10)
+
+
+def _argmax_sampler(logits: torch.Tensor) -> torch.Tensor:
+    return logits.argmax(dim=-1)
+
+
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_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)
+    first = watermarker.sample(logits, contexts, _argmax_sampler)
+    second = watermarker.sample(logits, contexts, _argmax_sampler)
 
     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):
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
+@pytest.mark.parametrize(
+    "config_overrides",
+    [
+        {"deduplicate_contexts": "none"},
+        {"deduplicate_contexts_max_history": 1023},
+    ],
+)
+def test_gumbel_config_warns_when_context_deduplication_is_weak(
+    monkeypatch, algorithm, config_overrides
+):
     messages: list[str] = []
     monkeypatch.setattr(
         "vllm.config.watermarking.logger.warning_once",
         lambda message, *, scope: messages.append(message),
     )
 
-    WatermarkConfig(key=42)
+    WatermarkConfig(key=42, algorithm=algorithm)
+    WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=1024)
+    WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=None)
+    WatermarkConfig(key=42, algorithm=algorithm, **config_overrides)
 
     assert messages == [
-        "Single-key Gumbel-max watermarking may increase the frequency of "
-        "degenerate generations, including repetition loops."
+        (
+            "Gumbel-max watermarking with context deduplication disabled "
+            "or limited to fewer than 1024 positions may increase the frequency of "
+            "degenerate generations, including repetition loops. Use "
+            "deduplicate_contexts='single_turn' or 'all' with "
+            "deduplicate_contexts_max_history at least 1024 or null to mitigate this."
+        )
     ]
 
 
+def test_context_deduplication_history_can_be_unbounded():
+    config = WatermarkConfig(key=42, deduplicate_contexts_max_history=None)
+
+    assert config.deduplicate_contexts_max_history is None
+
+
 def test_large_context_width_warns_but_is_allowed():
     config = WatermarkConfig(key=42, context_width=17)
 
@@ -56,6 +110,158 @@ def test_large_context_width_warns_but_is_allowed():
     assert watermarker.context_width == 17
 
 
+def test_dual_key_watermarker_uses_domain_separated_keys():
+    config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+
+    watermarker = create_watermarker(config)
+    assert isinstance(watermarker, SupportsSpeculativeDecoding)
+    draft = watermarker.draft_watermarker
+    target = watermarker.target_watermarker
+
+    assert isinstance(watermarker, DualKeyGumbelWatermarker)
+    assert watermarker.alpha == 0.1
+    assert draft.prf.key == derive_watermark_key(42, b"key_a")
+    assert target.prf.key == derive_watermark_key(42, b"key_b")
+    assert target.prf.key != draft.prf.key
+
+
+def test_dual_key_watermarker_routes_tokens_with_alpha():
+    watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25)
+    logits = torch.zeros(2, 16)
+    contexts = torch.tensor([[1, 2], [3, 4]])
+    key_a = watermarker.draft_watermarker.sample(logits, contexts, lambda _: None)
+    key_b = watermarker.target_watermarker.sample(logits, contexts, lambda _: None)
+
+    def route(routing_logits):
+        torch.testing.assert_close(
+            routing_logits.softmax(dim=-1),
+            torch.tensor([[0.75, 0.25], [0.75, 0.25]]),
+        )
+        return torch.tensor([0, 1])
+
+    sampled = watermarker.sample(logits, contexts, route)
+
+    assert torch.equal(
+        sampled.token_ids, torch.stack([key_a.token_ids[0], key_b.token_ids[1]])
+    )
+
+
+def test_dual_key_routing_logits_are_cached():
+    alpha = 0.25
+    watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha)
+    logits = torch.zeros(2, 16)
+    contexts = torch.tensor([[1, 2], [3, 4]])
+    recorded: list[torch.Tensor] = []
+
+    def route(routing_logits):
+        recorded.append(routing_logits)
+        return torch.tensor([0, 1])
+
+    watermarker.sample(logits, contexts, route)
+    watermarker.sample(logits, contexts, route)
+
+    expected = torch.tensor([1 - alpha, alpha], dtype=torch.float32).log()
+    expected = expected.expand(2, -1)
+    assert len(recorded) == 2
+    assert all(torch.equal(seen, expected) for seen in recorded)
+    assert len(watermarker._routing_logits_cache) == 1
+
+
+@pytest.mark.parametrize("alpha", [0.0, 0.25, 1.0])
+def test_dual_key_watermarker_skips_masked_rows(alpha: float):
+    watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha)
+    torch.manual_seed(0)
+    logits = torch.randn(4, 32)
+    contexts = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
+    skip_mask = torch.tensor([True, False, True, False])
+
+    unmasked = watermarker.sample(logits, contexts, _argmax_sampler)
+    masked = watermarker.sample(logits, contexts, _argmax_sampler, skip_mask)
+
+    ordinary = _argmax_sampler(logits)
+    assert torch.equal(masked.token_ids[skip_mask], ordinary[skip_mask])
+    assert torch.equal(masked.token_ids[~skip_mask], unmasked.token_ids[~skip_mask])
+    assert not torch.equal(unmasked.token_ids[skip_mask], ordinary[skip_mask])
+
+
+def test_dual_key_watermarker_routing_requires_a_random_sampler():
+    watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25)
+
+    with pytest.raises(ValueError, match="random sampler"):
+        watermarker.sample(torch.zeros(2, 16), torch.tensor([[1, 2], [3, 4]]))
+
+
+def test_speculative_decoding_uses_fixed_dual_key_roles():
+    watermarker = create_watermarker(
+        WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25)
+    )
+
+    target = create_speculative_target_watermarker(watermarker)
+    draft = create_speculative_draft_watermarker(
+        watermarker,
+        max_num_reqs=1,
+        device=torch.device("cpu"),
+        allow_target_only=False,
+    )
+
+    assert target.prf.key == derive_watermark_key(42, b"key_b")
+    assert draft is not None
+    assert draft.watermarker.prf.key == derive_watermark_key(42, b"key_a")
+
+
+def test_recovery_key_rejects_the_unsplit_dual_key_watermarker():
+    watermarker = create_watermarker(
+        WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25)
+    )
+
+    with pytest.raises(ValueError, match="keys the target role separately"):
+        _resolve_watermark_key(watermarker)
+
+    target = create_speculative_target_watermarker(watermarker)
+    assert _resolve_watermark_key(target) == derive_watermark_key(42, b"key_b")
+
+
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
+def test_config_key_resolution_matches_the_model_runner(algorithm: str):
+    config = WatermarkConfig(algorithm=algorithm, key=42, alpha=0.25)
+    runtime_key = _resolve_watermark_key(
+        create_speculative_target_watermarker(create_watermarker(config))
+    )
+
+    assert speculative_target_watermark_key(config) == runtime_key
+
+
+def test_config_key_resolution_returns_none_when_watermarking_is_disabled():
+    assert speculative_target_watermark_key(None) is None
+
+
+def test_target_only_speculative_watermarking_skips_draft_watermarker():
+    watermarker = create_watermarker(WatermarkConfig(algorithm="gumbel", key=42))
+
+    assert not isinstance(watermarker, SupportsSpeculativeDecoding)
+    with pytest.raises(ValueError, match="does not support speculative decoding"):
+        create_speculative_draft_watermarker(
+            watermarker,
+            max_num_reqs=1,
+            device=torch.device("cpu"),
+            allow_target_only=False,
+        )
+    assert (
+        create_speculative_draft_watermarker(
+            watermarker,
+            max_num_reqs=1,
+            device=torch.device("cpu"),
+            allow_target_only=True,
+        )
+        is None
+    )
+
+
+def test_dual_key_derivation_is_stable():
+    assert derive_watermark_key(32, b"key_a") == 16368605726115524094
+    assert derive_watermark_key(32, b"key_b") == 4799302812959726346
+
+
 def test_sampling_params_can_disable_watermarking():
     assert SamplingParams().watermarking
     assert not SamplingParams.from_optional(watermarking=False).watermarking
@@ -83,14 +289,9 @@ def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch):
 
 
 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.deduplicate_contexts = "single_turn"
     sampler.watermarking = SimpleNamespace(
         np=np.array([True, False]), gpu=torch.tensor([True, False])
     )
@@ -102,8 +303,11 @@ def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch):
     sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
         2, 1, dtype=torch.int64
     )
+    sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: torch.zeros(
+        2, dtype=torch.bool
+    )
     monkeypatch.setattr(
-        "vllm.v1.watermarking.gpu_sampler.gumbel_sample",
+        "vllm.v1.watermarking.watermarker.gumbel_sample",
         lambda *args, **kwargs: torch.tensor([3, 4]),
     )
     logits = torch.zeros(2, 8)
@@ -123,11 +327,524 @@ def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch):
     assert torch.equal(output_logits[1], logits[1])
 
 
+def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch):
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = StubWatermarker()
+    sampler.deduplicate_contexts = "single_turn"
+    sampler.watermarking = SimpleNamespace(
+        np=np.array([True, True]), gpu=torch.tensor([True, True])
+    )
+    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
+    )
+    sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: (
+        torch.tensor([True, False])
+    )
+    monkeypatch.setattr(
+        "vllm.v1.watermarking.watermarker.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([3, 7]))
+    assert torch.equal(output_logits[0], logits[0])
+    assert torch.equal(output_logits[1], torch.full((8,), 10.0))
+
+
+def test_gpu_sampler_can_disable_context_deduplication(monkeypatch):
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = StubWatermarker()
+    sampler.deduplicate_contexts = "none"
+    sampler.watermarking = SimpleNamespace(
+        np=np.array([True, True]), gpu=torch.tensor([True, True])
+    )
+    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
+    )
+    sampler._get_repeated_contexts = lambda *args: pytest.fail(
+        "context deduplication should not run"
+    )
+    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, 7]))
+    assert torch.equal(output_logits, torch.full((2, 8), 10.0))
+
+
+def test_repeated_context_mask_ignores_prompt_tokens():
+    all_token_ids = torch.tensor(
+        [
+            [8, 9, 1, 2, 1, 2],
+            [3, 4, 1, 2, 3, 4],
+            [0, 0, 0, 0, 0, 0],
+        ],
+        dtype=torch.int32,
+    )
+    req_indices = torch.tensor([0, 1, -1])
+    prompt_lens = torch.tensor([2, 2, 0])
+    total_lens = torch.tensor([6, 6, 0])
+    contexts = torch.tensor([[1, 2], [3, 4], [-1, -1]])
+
+    repeated = repeated_context_mask(
+        all_token_ids,
+        req_indices,
+        prompt_lens,
+        total_lens,
+        contexts,
+    )
+
+    assert torch.equal(repeated, torch.tensor([True, False, False]))
+
+
+def test_repeated_context_mask_can_include_prompt_tokens():
+    all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
+    req_indices = torch.tensor([0])
+    prompt_lens = torch.tensor([4])
+    total_lens = torch.tensor([10])
+    contexts = torch.tensor([[1, 2, 3, 4]])
+
+    single_turn = repeated_context_mask(
+        all_token_ids, req_indices, prompt_lens, total_lens, contexts
+    )
+    all_history = repeated_context_mask(
+        all_token_ids,
+        req_indices,
+        prompt_lens,
+        total_lens,
+        contexts,
+        include_prompt=True,
+    )
+
+    assert not single_turn.item()
+    assert all_history.item()
+
+
+def test_repeated_context_mask_can_skip_partial_contexts():
+    all_token_ids = torch.tensor([[10, 11, 12, 13, 1, 2, 3, 4]], dtype=torch.int32)
+    req_indices = torch.tensor([0])
+    prompt_lens = torch.tensor([4])
+    contexts = torch.tensor([[-1, -1, -1, 1]])
+
+    partial = repeated_context_mask(
+        all_token_ids,
+        req_indices,
+        prompt_lens,
+        torch.tensor([5]),
+        contexts,
+        include_prompt=True,
+        skip_partial_context=True,
+    )
+    complete = repeated_context_mask(
+        all_token_ids,
+        req_indices,
+        prompt_lens,
+        torch.tensor([8]),
+        torch.tensor([[1, 2, 3, 4]]),
+        include_prompt=True,
+        skip_partial_context=True,
+    )
+
+    assert partial.item()
+    assert not complete.item()
+
+
+def test_gpu_sampler_contexts_are_completion_local_for_all_scopes():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = SimpleNamespace(context_width=4)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
+        total_len=SimpleNamespace(gpu=torch.tensor([6])),
+    )
+    request_indices = torch.tensor([0])
+
+    sampler.deduplicate_contexts = "single_turn"
+    single_turn = sampler._get_contexts(request_indices)
+    sampler.deduplicate_contexts = "all"
+    all_history = sampler._get_contexts(request_indices)
+
+    assert torch.equal(single_turn, torch.tensor([[-1, -1, 1, 2]]))
+    assert torch.equal(all_history, single_turn)
+
+
+def test_gpu_sampler_all_history_skips_partial_context():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = SimpleNamespace(context_width=4)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
+        total_len=SimpleNamespace(gpu=torch.tensor([6])),
+    )
+    sampler.deduplicate_contexts_max_history = None
+    request_indices = torch.tensor([0])
+    contexts = sampler._get_contexts(request_indices)
+
+    sampler.deduplicate_contexts = "single_turn"
+    single_turn = sampler._get_repeated_contexts(request_indices, contexts)
+    sampler.deduplicate_contexts = "all"
+    all_history = sampler._get_repeated_contexts(request_indices, contexts)
+
+    assert not single_turn.item()
+    assert all_history.item()
+
+
+def test_repeated_context_mask_respects_max_history():
+    all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
+    req_indices = torch.tensor([0])
+    prompt_lens = torch.tensor([0])
+    total_lens = torch.tensor([10])
+    contexts = torch.tensor([[1, 2, 3, 4]])
+
+    full_history = repeated_context_mask(
+        all_token_ids, req_indices, prompt_lens, total_lens, contexts
+    )
+    last_six = repeated_context_mask(
+        all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=6
+    )
+    last_five = repeated_context_mask(
+        all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=5
+    )
+
+    assert full_history.item()
+    assert last_six.item()
+    assert not last_five.item()
+
+
+@pytest.mark.parametrize("max_history", [0, -1])
+def test_repeated_context_mask_rejects_non_positive_max_history(max_history: int):
+    with pytest.raises(ValueError, match="max_history must be positive or None"):
+        repeated_context_mask(
+            torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
+            torch.tensor([0], dtype=torch.int32),
+            torch.tensor([0], dtype=torch.int32),
+            torch.tensor([4], dtype=torch.int32),
+            torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
+            max_history=max_history,
+        )
+
+
+def test_gpu_sampler_wires_request_history_scope():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(
+            gpu=torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
+        ),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
+        total_len=SimpleNamespace(gpu=torch.tensor([10])),
+    )
+    sampler.deduplicate_contexts_max_history = None
+    request_indices = torch.tensor([0])
+    contexts = torch.tensor([[1, 2, 3, 4]])
+
+    sampler.deduplicate_contexts = "single_turn"
+    single_turn = sampler._get_repeated_contexts(request_indices, contexts)
+    sampler.deduplicate_contexts = "all"
+    all_history = sampler._get_repeated_contexts(request_indices, contexts)
+
+    assert not single_turn.item()
+    assert all_history.item()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+@pytest.mark.parametrize("max_history", [5, 6])
+@pytest.mark.parametrize("include_prompt", [False, True])
+def test_repeated_context_mask_max_history_accelerator_parity(
+    max_history: int, include_prompt: bool
+):
+    inputs = (
+        torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]),
+        torch.tensor([0]),
+        torch.tensor([0]),
+        torch.tensor([10]),
+        torch.tensor([[1, 2, 3, 4]]),
+    )
+
+    expected = repeated_context_mask(
+        *inputs, max_history=max_history, include_prompt=include_prompt
+    )
+    actual = repeated_context_mask(
+        *(value.cuda() for value in inputs),
+        max_history=max_history,
+        include_prompt=include_prompt,
+    ).cpu()
+
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+def test_repeated_context_mask_partial_context_accelerator_parity():
+    inputs = (
+        torch.tensor([[10, 11, 12, 13, 1]], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        torch.tensor([4], dtype=torch.int32),
+        torch.tensor([5], dtype=torch.int32),
+        torch.tensor([[-1, -1, -1, 1]], dtype=torch.int64),
+    )
+
+    expected = repeated_context_mask(
+        *inputs, include_prompt=True, skip_partial_context=True
+    )
+    actual = repeated_context_mask(
+        *(value.cuda() for value in inputs),
+        include_prompt=True,
+        skip_partial_context=True,
+    ).cpu()
+
+    assert torch.equal(actual, expected)
+
+
+def _unaligned_max_history_inputs() -> tuple[torch.Tensor, ...]:
+    """One row whose only repeated context ends at position 680.
+
+    680 is not a multiple of the kernel's 512-position block, so the scan window
+    selected by ``max_history`` starts in the middle of a block.
+    """
+    history = list(range(1_000, 2_200))
+    history[676:680] = [1, 2, 3, 4]
+    return (
+        torch.tensor([history], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        torch.tensor([len(history)], dtype=torch.int32),
+        torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
+    )
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+@pytest.mark.parametrize(
+    "max_history,expected", [(1200 - 680, True), (1200 - 680 - 1, False)]
+)
+def test_repeated_context_mask_unaligned_max_history_accelerator_parity(
+    max_history: int, expected: bool
+):
+    inputs = _unaligned_max_history_inputs()
+
+    reference = repeated_context_mask(*inputs, max_history=max_history)
+    actual = repeated_context_mask(
+        *(value.cuda() for value in inputs), max_history=max_history
+    ).cpu()
+
+    assert torch.equal(reference, torch.tensor([expected]))
+    assert torch.equal(actual, reference)
+
+
+def _repeated_context_inputs(
+    context_width: int, device: str = "cpu"
+) -> tuple[torch.Tensor, ...]:
+    """Build repeated, unique, and padded request rows for mask tests."""
+    prompt = [101, 102, 103, 104]
+    repeated_output = [*range(1, context_width + 1)] * 2
+    unique_output = [*range(1, context_width + 2)]
+    max_len = len(prompt) + len(repeated_output)
+    rows = [
+        prompt + repeated_output,
+        prompt + unique_output,
+        [],
+    ]
+    all_token_ids = torch.zeros((3, max_len), dtype=torch.int32, device=device)
+    for row, token_ids in enumerate(rows):
+        all_token_ids[row, : len(token_ids)] = torch.tensor(
+            token_ids, dtype=torch.int32, device=device
+        )
+    return (
+        all_token_ids,
+        torch.tensor([0, 1, -1], dtype=torch.int32, device=device),
+        torch.tensor([len(prompt), len(prompt), 0], dtype=torch.int32, device=device),
+        torch.tensor([len(rows[0]), len(rows[1]), 0], dtype=torch.int32, device=device),
+        torch.tensor(
+            [
+                repeated_output[-context_width:],
+                unique_output[-context_width:],
+                [-1] * context_width,
+            ],
+            dtype=torch.int64,
+            device=device,
+        ),
+    )
+
+
+@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17])
+def test_repeated_context_mask(context_width: int):
+    repeated = repeated_context_mask(*_repeated_context_inputs(context_width))
+
+    assert torch.equal(repeated, torch.tensor([True, False, False]))
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17])
+def test_repeated_context_mask_accelerator_parity(context_width: int):
+    cpu_inputs = _repeated_context_inputs(context_width)
+    accelerator_inputs = list(_repeated_context_inputs(context_width, "cuda"))
+    contexts = accelerator_inputs[-1]
+    storage = torch.empty(
+        contexts.shape[0], contexts.shape[1] * 2, dtype=contexts.dtype, device="cuda"
+    )
+    storage[:, ::2] = contexts
+    accelerator_inputs[-1] = storage[:, ::2]
+
+    expected = repeated_context_mask(*cpu_inputs)
+    actual = repeated_context_mask(*accelerator_inputs).cpu()
+
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+def test_repeated_context_mask_scans_multiple_blocks():
+    context = [1, 2, 3, 4]
+    repeated_output = [*range(10_000, 11_030), *context, 99, *context]
+    unique_output = list(range(20_000, 21_039))
+    prompt = [101, 102, 103, 104]
+    rows = [prompt + repeated_output, prompt + unique_output]
+    all_token_ids = torch.zeros((2, len(rows[0])), dtype=torch.int32)
+    for row, token_ids in enumerate(rows):
+        all_token_ids[row, : len(token_ids)] = torch.tensor(token_ids)
+    inputs = (
+        all_token_ids,
+        torch.tensor([0, 1], dtype=torch.int32),
+        torch.tensor([len(prompt), len(prompt)], dtype=torch.int32),
+        torch.tensor([len(rows[0]), len(rows[1])], dtype=torch.int32),
+        torch.tensor([context, unique_output[-4:]], dtype=torch.int64),
+    )
+
+    expected = repeated_context_mask(*inputs)
+    actual = repeated_context_mask(*(value.cuda() for value in inputs)).cpu()
+
+    assert torch.equal(expected, torch.tensor([True, False]))
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+def test_repeated_context_mask_last_request_at_capacity():
+    max_model_len = 1024
+    all_token_ids = torch.arange(
+        2 * max_model_len, dtype=torch.int32, device="cuda"
+    ).reshape(2, max_model_len)
+    inputs = (
+        all_token_ids,
+        torch.tensor([1], dtype=torch.int32, device="cuda"),
+        torch.tensor([0, 0], dtype=torch.int32, device="cuda"),
+        torch.tensor([max_model_len, max_model_len], dtype=torch.int32, device="cuda"),
+        all_token_ids[1, -4:].to(torch.int64).unsqueeze(0),
+    )
+
+    repeated = repeated_context_mask(*inputs)
+
+    assert not repeated.item()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+def test_gpu_sampler_uses_fused_gumbel_for_repeated_contexts():
+    device = torch.device("cuda")
+    logits = torch.tensor([[0.0, 1.0, 2.0, 3.0], [3.0, 2.0, 1.0, 0.0]], device=device)
+    all_token_ids = torch.tensor(
+        [[101, 102, 1, 2, 1, 2], [101, 102, 3, 4, 5, 6]],
+        dtype=torch.int32,
+        device=device,
+    )
+    request_indices = torch.tensor([0, 1], dtype=torch.int64, device=device)
+    request_indices_np = np.array([0, 1])
+    temperatures = torch.ones(2, device=device)
+    seeds = torch.tensor([11, 22], dtype=torch.int64, device=device)
+    positions = torch.tensor([6, 6], dtype=torch.int64, device=device)
+
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = GumbelWatermarker(key=42, context_width=1)
+    sampler.deduplicate_contexts = "single_turn"
+    sampler.deduplicate_contexts_max_history = None
+    sampler.watermarking = SimpleNamespace(
+        np=np.ones(2, dtype=bool), gpu=torch.ones(2, dtype=torch.bool, device=device)
+    )
+    sampler.sampling_states = SimpleNamespace(
+        temperature=SimpleNamespace(np=np.ones(2), gpu=temperatures),
+        seeds=SimpleNamespace(gpu=seeds),
+    )
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=all_token_ids),
+        prompt_len=SimpleNamespace(
+            gpu=torch.tensor([2, 2], dtype=torch.int32, device=device)
+        ),
+        total_len=SimpleNamespace(
+            gpu=torch.tensor([6, 6], dtype=torch.int32, device=device)
+        ),
+    )
+    sampler.use_fp64_gumbel = False
+
+    contexts = sampler._get_contexts(request_indices)
+    repeated = sampler._get_repeated_contexts(request_indices, contexts)
+    expected = philox_gumbel_sample(
+        logits,
+        contexts,
+        42,
+        skip_mask=repeated,
+        expanded_idx_mapping=request_indices,
+        temperatures=temperatures,
+        seeds=seeds,
+        positions=positions,
+    )
+
+    actual, output_logits = sampler._sample_random(
+        logits,
+        request_indices,
+        request_indices_np,
+        positions,
+        None,
+        None,
+        False,
+    )
+
+    assert torch.equal(repeated, torch.tensor([True, False], device=device))
+    assert torch.equal(actual, expected)
+    assert output_logits is logits
+
+
 def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch):
     class StubWatermarker:
         context_width = 1
 
-        def sample(self, logits, contexts, random_sample):
+        def sample(self, logits, contexts, random_sampler=None, skip_mask=None):
             raise AssertionError("watermarker should not run for greedy requests")
 
     sampler = object.__new__(GPUWatermarkSampler)
@@ -155,3 +872,157 @@ def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch):
     )
 
     assert actual is expected
+
+
+def test_gpu_sampler_builds_speculative_contexts_from_drafts():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = SimpleNamespace(context_width=2)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 101, 10, 11, 0, 0]])),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([2])),
+        total_len=SimpleNamespace(gpu=torch.tensor([4])),
+    )
+
+    contexts = sampler._get_contexts(
+        torch.tensor([0, 0, 0]),
+        torch.tensor([0, 1, 2]),
+        torch.tensor([11, 20, 21]),
+    )
+
+    assert torch.equal(contexts, torch.tensor([[10, 11], [11, 20], [20, 21]]))
+
+
+def test_gpu_sampler_builds_chunked_multi_request_speculative_contexts():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = SimpleNamespace(context_width=2)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(
+            gpu=torch.tensor(
+                [
+                    [100, 101, 10, 11, 0, 0],
+                    [200, 201, 30, 31, 0, 0],
+                    [300, 301, 50, 51, 0, 0],
+                ]
+            )
+        ),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([2, 2, 2])),
+        total_len=SimpleNamespace(gpu=torch.tensor([4, 4, 4])),
+    )
+
+    contexts = sampler._get_contexts(
+        torch.tensor([2, 2, 1, 1]),
+        torch.tensor([0, 1, 0, 1]),
+        torch.tensor([51, 60, 31, 40]),
+    )
+
+    assert torch.equal(
+        contexts,
+        torch.tensor([[50, 51], [51, 60], [30, 31], [31, 40]]),
+    )
+
+
+def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch):
+    class StubSpeculator(DraftModelSpeculator):
+        def capture(self): ...
+
+        def init_cudagraph_manager(self, cudagraph_mode): ...
+
+        def load_draft_model(self, target_model, target_attn_layer_names): ...
+
+        def propose(self, *args, **kwargs): ...
+
+    class StubModel:
+        @staticmethod
+        def compute_logits(hidden_states):
+            return torch.zeros(hidden_states.shape[0], 8)
+
+    class StubWatermarker:
+        @staticmethod
+        def sample(logits, contexts, random_sampler=None):
+            assert random_sampler is None
+            return WatermarkSample(torch.tensor([7, 7]), logits)
+
+    speculator = object.__new__(StubSpeculator)
+    speculator.model = StubModel()
+    speculator.use_fp64_gumbel = False
+    draft_watermarker = object.__new__(DraftWatermarker)
+    draft_watermarker.watermarker = StubWatermarker()
+    draft_watermarker.contexts = torch.tensor([[1, 2], [3, 4]])
+    draft_watermarker.enabled = torch.tensor([True, False])
+    speculator.draft_watermarker = draft_watermarker
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.spec_decode.speculator.gumbel_sample",
+        lambda *args, **kwargs: torch.tensor([3, 4]),
+    )
+
+    sampled = speculator.sample_draft(
+        hidden_states=torch.zeros(2, 4),
+        sample_src_positions=torch.zeros(2, dtype=torch.int64),
+        idx_mapping=torch.tensor([0, 1]),
+        temperature=torch.ones(2),
+        seeds=torch.zeros(2, dtype=torch.int64),
+        draft_step=torch.tensor(0),
+        draft_logits=torch.zeros(2, 1, 8),
+    )
+
+    assert torch.equal(sampled, torch.tensor([7, 4]))
+    assert torch.equal(draft_watermarker.contexts, torch.tensor([[2, 7], [4, 4]]))
+
+
+def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):
+    speculator = object.__new__(DSparkSpeculator)
+    speculator.draft_logits = torch.zeros(2, 1, 8)
+    speculator._d2t_scatter_index = torch.tensor([1, 5])
+    speculator._draft_scatter_buf = torch.full((2, 8), float("-inf"))
+    speculator.temperature = torch.ones(2)
+    speculator.seeds = torch.zeros(2, dtype=torch.int64)
+    speculator._step_cols = torch.tensor([0])
+    speculator.use_fp64_gumbel = False
+    watermark_logits: list[torch.Tensor] = []
+
+    def sample(logits, sampled, idx_map, temperature):
+        watermark_logits.append(logits.clone())
+        return sampled + 1
+
+    speculator.draft_watermarker = SimpleNamespace(sample=sample)
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",
+        lambda *args, **kwargs: torch.tensor([3, 4]),
+    )
+
+    sampled = speculator._sample_logits(
+        torch.tensor([[10.0, 20.0], [30.0, 40.0]]),
+        torch.tensor([0, 1]),
+        torch.tensor([1, 1]),
+        0,
+    )
+
+    assert torch.equal(sampled, torch.tensor([4, 5]))
+    assert torch.equal(
+        watermark_logits[0][:, [1, 5]], torch.tensor([[10, 20], [30, 40]])
+    )
+    assert torch.isneginf(watermark_logits[0][:, [0, 2, 3, 4, 6, 7]]).all()
+
+
+def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch):
+    speculator = object.__new__(DSparkSpeculator)
+    speculator.draft_logits = torch.zeros(2, 1, 8)
+    speculator._d2t_scatter_index = None
+    speculator.temperature = torch.ones(2)
+    speculator.seeds = torch.zeros(2, dtype=torch.int64)
+    speculator._step_cols = torch.tensor([0])
+    speculator.use_fp64_gumbel = False
+    speculator.draft_watermarker = None
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",
+        lambda *args, **kwargs: torch.tensor([3, 4]),
+    )
+
+    sampled = speculator._sample_logits(
+        torch.zeros(2, 8),
+        torch.tensor([0, 1]),
+        torch.tensor([1, 1]),
+        0,
+    )
+
+    assert torch.equal(sampled, torch.tensor([3, 4]))

vllm/config/vllm.py

@@ -50,7 +50,7 @@ from .reasoning import ReasoningConfig
 from .scheduler import SchedulerConfig
 from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig
 from .structured_outputs import StructuredOutputsConfig
-from .utils import SupportsHash, config, replace
+from .utils import SupportsHash, config, get_field, replace
 from .watermarking import WatermarkConfig
 from .weight_transfer import WeightTransferConfig
 
@@ -1142,14 +1142,67 @@ class VllmConfig:
         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 self.speculative_config is not None:
+            speculative_config = self.speculative_config
+            if speculative_config.draft_sample_method != "probabilistic":
+                raise ValueError(
+                    "Speculative decoding with watermarking requires "
+                    "draft_sample_method='probabilistic'."
+                )
+            if speculative_config.rejection_sample_method != "standard":
+                raise ValueError(
+                    "Speculative decoding with watermarking requires "
+                    "rejection_sample_method='standard'."
+                )
+            if (
+                speculative_config.parallel_drafting
+                and speculative_config.method != "dspark"
+            ):
+                raise ValueError(
+                    "Parallel speculative drafting is not supported with watermarking."
+                )
+            if speculative_config.method not in ("dspark", "eagle", "eagle3", "mtp"):
+                raise ValueError(
+                    "Watermarking supports only autoregressive model-based "
+                    "speculative decoding."
+                )
+            if (
+                not watermark_config.allow_target_only_watermarking
+                and not watermark_config.supports_speculative_decoding
+            ):
+                raise ValueError(
+                    f"The '{watermark_config.algorithm}' watermarking algorithm "
+                    "does not support speculative decoding. Set "
+                    "allow_target_only_watermarking=true to leave draft tokens "
+                    "unwatermarked."
+                )
+            if (
+                watermark_config.allow_target_only_watermarking
+                and not watermark_config.supports_speculative_decoding
+            ):
+                logger.warning_once(
+                    "Target-only watermarking leaves accepted draft tokens "
+                    "unwatermarked, weakening detectability in proportion to the "
+                    "share of output tokens supplied by accepted drafts.",
+                    scope="global",
+                )
+            if watermark_config.deduplicate_contexts != "none":
+                logger.warning_once(
+                    "Context deduplication is not supported with speculative "
+                    "decoding and will not be applied to accepted drafts, "
+                    "rejection-recovery tokens, or bonus tokens.",
+                    scope="global",
+                )
+            if watermark_config.algorithm == "dual_key_gumbel" and (
+                watermark_config.alpha != get_field(WatermarkConfig, "alpha").default
+            ):
+                logger.warning_once(
+                    "Speculative decoding selects the watermark key by token role: "
+                    "draft tokens use key A, recovery and bonus tokens use key B. "
+                    "The configured alpha=%s is not used.",
+                    watermark_config.alpha,
+                    scope="global",
+                )
         if beam_search:
             raise ValueError("Beam search is not supported with watermarking.")
         if custom_sampler:

vllm/config/watermarking.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import hashlib
 from typing import Literal
 
 from pydantic import Field, model_validator
@@ -11,12 +12,16 @@ from vllm.logger import init_logger
 
 logger = init_logger(__name__)
 
-WatermarkingAlgorithm = Literal["gumbel"]
+WatermarkingAlgorithm = Literal["gumbel", "dual_key_gumbel"]
 WatermarkPRFName = Literal["philox"]
+WatermarkContextScope = Literal["none", "single_turn", "all"]
 
-_SPECULATIVE_DECODING_SUPPORT: dict[WatermarkingAlgorithm, bool] = {
-    "gumbel": False,
-}
+_MIN_RECOMMENDED_DEDUP_HISTORY = 1024
+
+
+def derive_watermark_key(key: int, domain: bytes) -> int:
+    digest = hashlib.sha256(domain + key.to_bytes(8, "big")).digest()
+    return int.from_bytes(digest[:8], "big")
 
 
 @config
@@ -27,23 +32,49 @@ class WatermarkConfig:
     """Secret key used to watermark generated text."""
     algorithm: WatermarkingAlgorithm = "gumbel"
     """Algorithm used to watermark generated text."""
+    alpha: float = Field(default=0.1, ge=0, le=1)
+    """Probability of selecting key B for dual-key watermarking."""
     context_width: int = Field(default=4, ge=1)
-    """Number of prior output tokens used by the watermark PRF."""
+    """Number of prior tokens used by the watermark PRF."""
+    deduplicate_contexts: WatermarkContextScope = "single_turn"
+    """Which history is searched for a repeated context before a token is
+    sampled; a repeated context is sampled without the watermark. `none`
+    disables the search. `single_turn` (default) searches this request's
+    generated tokens. `all` also searches the prompt and samples the first
+    `context_width` generated tokens without watermarking."""
+    deduplicate_contexts_max_history: int | None = Field(default=8192, ge=1)
+    """Number of most recent history positions searched (default 8192), or
+    `None` for the whole scope. Each position is compared over the
+    `context_width` tokens before it. Ignored when `deduplicate_contexts` is
+    `none`."""
     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
+    allow_target_only_watermarking: bool = False
+    """Allow speculative decoding without watermarking draft tokens."""
 
     @property
     def supports_speculative_decoding(self) -> bool:
-        return _SPECULATIVE_DECODING_SUPPORT[self.algorithm]
+        return self.algorithm == "dual_key_gumbel"
+
+    @model_validator(mode="after")
+    def validate_watermark_settings(self) -> Self:
+        if self.key > 2**64 - 1:
+            raise ValueError("philox keys must fit in 64 bits")
+        history_is_too_short = (
+            self.deduplicate_contexts_max_history is not None
+            and self.deduplicate_contexts_max_history < _MIN_RECOMMENDED_DEDUP_HISTORY
+        )
+        if self.algorithm in ("gumbel", "dual_key_gumbel") and (
+            self.deduplicate_contexts == "none" or history_is_too_short
+        ):
+            logger.warning_once(
+                "Gumbel-max watermarking with context deduplication "
+                "disabled or limited to fewer than "
+                f"{_MIN_RECOMMENDED_DEDUP_HISTORY} positions may increase the "
+                "frequency of degenerate generations, including repetition loops. "
+                "Use deduplicate_contexts='single_turn' or 'all' with "
+                "deduplicate_contexts_max_history at least "
+                f"{_MIN_RECOMMENDED_DEDUP_HISTORY} or null to mitigate this.",
+                scope="global",
+            )
+        return self

vllm/model_executor/warmup/kernel_warmup.py

@@ -44,6 +44,9 @@ from vllm.model_executor.warmup.replayssm_warmup import (
 from vllm.model_executor.warmup.spec_decode_rejection_warmup import (
     spec_decode_rejection_warmup,
 )
+from vllm.model_executor.warmup.watermark_sample_warmup import (
+    watermark_sample_warmup,
+)
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import is_deep_gemm_supported
 from vllm.utils.flashinfer import has_flashinfer
@@ -207,6 +210,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
     if enable_jit_warmup:
         kimi_k3_triton_warmup(worker)
         spec_decode_rejection_warmup(worker)
+        watermark_sample_warmup(worker)
         qwen4_exp_qsa_triton_warmup(worker)
 
     if enable_jit_warmup and current_platform.is_device_capability_family(100):

vllm/model_executor/warmup/spec_decode_rejection_warmup.py

@@ -1,17 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Warm up spec-decode rejection-sampler Triton kernels.
-
-The rejection sampler kernels (``_compute_local_logits_stats_kernel``,
-``_rejection_kernel``, ``_resample_kernel``) are JIT-compiled by Triton on
-first use. Without warmup, the first spec-decode request pays a multi-second
-compilation cost. This pre-compiles them with dummy data matching the
-server's vocab size and speculative config.
-"""
+"""Warm up speculative-decoding rejection kernels."""
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 
@@ -39,42 +32,76 @@ def spec_decode_rejection_warmup(worker: Worker) -> None:
     if num_spec <= 0 or vocab_size <= 0:
         return
 
-    # Mirror the constexpr-relevant flags the runtime uses.
     rejection_method = getattr(spec_config, "rejection_sample_method", None)
     use_block_verification = rejection_method == "block"
     use_synthetic = rejection_method == "synthetic"
+    # HAS_DRAFT_LOGITS is a constexpr, so the warmup must match the runtime.
+    use_draft_logits = (
+        getattr(spec_config, "draft_sample_method", "greedy") == "probabilistic"
+    )
 
     device = worker.device
     num_reqs = 1
     tokens_per_req = num_spec + 1
     num_logits = num_reqs * tokens_per_req
 
-    # Triton JIT-specializes on tensor dtypes. The target logits may be fp32
-    # (apply_sampling_params copies to fp32 when processing is needed) or the
-    # model dtype (pass-through otherwise), while draft logits are always the
-    # model dtype. Warm every (target, draft) combination the runtime can hit.
+    watermark_config = getattr(worker.vllm_config, "watermark_config", None)
+    watermark_kwargs: dict[str, Any] = {}
+    if watermark_config is not None:
+        from vllm.v1.watermarking.spec_decode import speculative_target_watermark_key
+
+        try:
+            watermark_key = speculative_target_watermark_key(watermark_config)
+        except Exception:
+            logger.warning(
+                "Skipping spec-decode rejection sampler warmup: could not derive "
+                "the watermark recovery key.",
+                exc_info=True,
+            )
+            return
+        # Tensor dtypes and integer widths affect Triton specialization.
+        watermark_kwargs = {
+            "contexts": torch.zeros(
+                (num_logits, watermark_config.context_width),
+                dtype=torch.int32,
+                device=device,
+            ),
+            "watermarking": torch.zeros(num_reqs, dtype=torch.bool, device=device),
+            "watermark_key": watermark_key,
+        }
+
+    # Sampling-parameter processing may promote either distribution to fp32.
     model_dtype = model_config.dtype
-    warmup_dtype_pairs = {
-        (model_dtype, model_dtype),
-        (torch.float32, torch.float32),
-        (torch.float32, model_dtype),
-        (model_dtype, torch.float32),
-    }
+    warmup_dtype_pairs: set[tuple[torch.dtype, torch.dtype | None]]
+    if use_draft_logits:
+        warmup_dtype_pairs = {
+            (model_dtype, model_dtype),
+            (torch.float32, torch.float32),
+            (torch.float32, model_dtype),
+            (model_dtype, torch.float32),
+        }
+    else:
+        warmup_dtype_pairs = {(model_dtype, None), (torch.float32, None)}
 
     logger.info(
         "Warming up spec-decode rejection sampler kernels "
-        "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s).",
+        "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s, watermark=%s).",
         vocab_size,
         num_spec,
         [(str(t), str(d)) for t, d in warmup_dtype_pairs],
         use_block_verification,
+        watermark_config is not None,
     )
     for tgt_dtype, draft_dtype in warmup_dtype_pairs:
         target_logits = torch.zeros(
             (num_logits, vocab_size), dtype=tgt_dtype, device=device
         )
-        draft_logits = torch.zeros(
-            (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device
+        draft_logits = (
+            torch.zeros(
+                (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device
+            )
+            if draft_dtype is not None
+            else None
         )
         synthetic_rates = (
             torch.full((num_spec,), 0.5, dtype=torch.float32, device=device)
@@ -85,14 +112,14 @@ def spec_decode_rejection_warmup(worker: Worker) -> None:
             rejection_sample(
                 target_logits=target_logits,
                 draft_logits=draft_logits,
-                draft_sampled=torch.zeros(num_logits, dtype=torch.int64, device=device),
+                draft_sampled=torch.zeros(num_logits, dtype=torch.int32, device=device),
                 cu_num_logits=torch.tensor(
                     [0, num_logits], dtype=torch.int32, device=device
                 ),
                 pos=torch.zeros(num_logits, dtype=torch.int64, device=device),
-                idx_mapping=torch.zeros(num_reqs, dtype=torch.int32, device=device),
+                idx_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device),
                 expanded_idx_mapping=torch.zeros(
-                    num_logits, dtype=torch.int32, device=device
+                    num_logits, dtype=torch.int64, device=device
                 ),
                 expanded_local_pos=torch.arange(
                     num_logits, dtype=torch.int32, device=device
@@ -103,6 +130,7 @@ def spec_decode_rejection_warmup(worker: Worker) -> None:
                 synthetic_conditional_rates=synthetic_rates,
                 use_fp64=False,
                 use_block_verification=use_block_verification,
+                **watermark_kwargs,
             )
         except Exception:
             logger.warning(

vllm/model_executor/warmup/watermark_sample_warmup.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Warm up the watermarked sampler's Triton kernel.
+
+``_philox_gumbel_kernel`` is JIT-compiled once per (key, logits dtype,
+skip-mask, ``USE_FP64``) specialization. The generic sampler warmup samples with
+``SamplingParams.for_sampler_warmup()``, whose feature-heavy logits processing
+forces the fp32 copy in ``apply_sampling_params``, so the first ordinary
+temperature-1.0 request is the first launch with model-dtype logits and pays
+the compilation inside inference. This pre-compiles every specialization the
+sampler path can launch for the configured watermark.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+if TYPE_CHECKING:
+    from vllm.v1.watermarking.watermarker import Watermarker
+    from vllm.v1.worker.gpu_worker import Worker
+
+logger = init_logger(__name__)
+
+_NUM_WARMUP_TOKENS = 1
+
+
+def _philox_key(watermarker: Watermarker) -> int | None:
+    """None when the watermarker takes the torch path and compiles no kernel."""
+    from vllm.v1.watermarking.gumbel import GumbelWatermarker
+    from vllm.v1.watermarking.prfs import PhiloxPRF
+
+    if not isinstance(watermarker, GumbelWatermarker):
+        return None
+    prf = watermarker.prf
+    return prf.key if type(prf) is PhiloxPRF else None
+
+
+def _philox_sampler_keys(watermarker: Watermarker) -> list[int]:
+    """Philox keys the sampler launches the kernel with."""
+    from vllm.v1.watermarking.gumbel import DualKeyGumbelWatermarker
+
+    watermarkers: list[Watermarker] = [watermarker]
+    if isinstance(watermarker, DualKeyGumbelWatermarker):
+        # ``DualKeyGumbelWatermarker.sample`` launches once per key, and
+        # collapses to a single key at the alpha bounds, where the other key's
+        # specialization is never launched.
+        if watermarker.alpha == 0:
+            watermarkers = [watermarker.draft_watermarker]
+        elif watermarker.alpha == 1:
+            watermarkers = [watermarker.target_watermarker]
+        else:
+            watermarkers = [
+                watermarker.draft_watermarker,
+                watermarker.target_watermarker,
+            ]
+    return [key for key in map(_philox_key, watermarkers) if key is not None]
+
+
+@torch.inference_mode()
+def watermark_sample_warmup(worker: Worker) -> None:
+    if getattr(worker.vllm_config, "watermark_config", None) is None:
+        return
+    # ``GumbelWatermarker`` launches the kernel only on CUDA
+    # (vllm/v1/watermarking/gumbel.py:63,:82); elsewhere the sampler takes the
+    # torch path and warming here would compile a kernel nothing launches.
+    if not current_platform.is_cuda_alike():
+        return
+
+    from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
+    from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
+
+    # Read the sampler the runner built instead of rebuilding it from config:
+    # its watermarker (the target role under a speculative config), its context
+    # deduplication setting and its fp64 flag are what the launches have to
+    # match, and cannot drift from the runner. A rank that does not sample -- a
+    # non-final pipeline stage, a pooling model -- holds no such sampler.
+    sampler = getattr(worker.model_runner, "sampler", None)
+    if not isinstance(sampler, GPUWatermarkSampler):
+        return
+
+    model_config = worker.vllm_config.model_config
+    device = worker.device
+    try:
+        keys = _philox_sampler_keys(sampler.watermarker)
+        if not keys:
+            return
+
+        # The sampler passes model-dtype logits through when no logits
+        # processing is needed and fp32 otherwise. ``ModelConfig.dtype`` is
+        # resolved to a torch dtype during validation.
+        dtypes = {cast("torch.dtype", model_config.dtype), torch.float32}
+        # ``GPUWatermarkSampler._sample_random`` builds a skip mask whenever
+        # context deduplication is on; with it off only a batch that mixes in an
+        # opted-out or greedy row does, and the mask-free specialization stays
+        # reachable.
+        with_skip_mask = (
+            (True,) if sampler.deduplicate_contexts != "none" else (True, False)
+        )
+
+        vocab_size = model_config.get_vocab_size()
+        contexts = torch.zeros(
+            (_NUM_WARMUP_TOKENS, sampler.watermarker.context_width),
+            dtype=torch.int32,
+            device=device,
+        )
+        # Dtypes mirror the runtime buffers the sampler passes and are part of
+        # the specialization key: int32 request-state token ids, int64 index and
+        # seed buffers, fp32 sampling temperatures.
+        sampling_state = {
+            "skip_mask": torch.zeros(
+                _NUM_WARMUP_TOKENS, dtype=torch.bool, device=device
+            ),
+            "expanded_idx_mapping": torch.zeros(
+                _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device
+            ),
+            "temperatures": torch.ones(
+                _NUM_WARMUP_TOKENS, dtype=torch.float32, device=device
+            ),
+            "seeds": torch.zeros(_NUM_WARMUP_TOKENS, dtype=torch.int64, device=device),
+            "positions": torch.zeros(
+                _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device
+            ),
+        }
+
+        logger.info(
+            "Warming up watermark sampler kernel (vocab=%d, keys=%d, dtypes=%s, "
+            "skip_mask=%s).",
+            vocab_size,
+            len(keys),
+            [str(dtype) for dtype in dtypes],
+            list(with_skip_mask),
+        )
+        for dtype in dtypes:
+            logits = torch.zeros(
+                (_NUM_WARMUP_TOKENS, vocab_size), dtype=dtype, device=device
+            )
+            for key in keys:
+                for use_skip_mask in with_skip_mask:
+                    philox_gumbel_sample(
+                        logits,
+                        contexts,
+                        key,
+                        # Only the mixed path forwards the sampler's setting;
+                        # ``GumbelWatermarker._sample_watermarked`` launches
+                        # the mask-free variant with the fp32 default.
+                        use_fp64=use_skip_mask and sampler.use_fp64_gumbel,
+                        **(sampling_state if use_skip_mask else {}),
+                    )
+    except Exception:
+        logger.warning("Skipping watermark sampler warmup.", exc_info=True)

vllm/v1/watermarking/__init__.py

@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from vllm.config.watermarking import derive_watermark_key
 from vllm.v1.watermarking.detector import WatermarkDetection, WatermarkDetector
 from vllm.v1.watermarking.factory import create_watermarker
 from vllm.v1.watermarking.gumbel import (
+    DualKeyGumbelWatermarkDetector,
+    DualKeyGumbelWatermarker,
     GumbelWatermarkDetector,
     GumbelWatermarker,
 )
@@ -12,12 +15,19 @@ from vllm.v1.watermarking.prfs import (
     WatermarkPRF,
     create_prf,
 )
-from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample
+from vllm.v1.watermarking.watermarker import (
+    SupportsSpeculativeDecoding,
+    Watermarker,
+    WatermarkSample,
+)
 
 __all__ = [
+    "DualKeyGumbelWatermarkDetector",
+    "DualKeyGumbelWatermarker",
     "GumbelWatermarkDetector",
     "GumbelWatermarker",
     "PhiloxPRF",
+    "SupportsSpeculativeDecoding",
     "WatermarkDetection",
     "WatermarkDetector",
     "WatermarkPRF",
@@ -25,4 +35,5 @@ __all__ = [
     "WatermarkSample",
     "create_watermarker",
     "create_prf",
+    "derive_watermark_key",
 ]

vllm/v1/watermarking/factory.py

@@ -1,24 +1,22 @@
 # 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.gumbel import (
+    DualKeyGumbelWatermarker,
+    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)
+    if config.algorithm == "gumbel":
+        return GumbelWatermarker(config.key, config.context_width, config.prf)
+    if config.algorithm == "dual_key_gumbel":
+        return DualKeyGumbelWatermarker(
+            config.key,
+            config.context_width,
+            config.prf,
+            config.alpha,
+        )
+    raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}")

vllm/v1/watermarking/gpu_sampler.py

@@ -1,16 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from typing import Literal
+
 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.watermarking.watermarker import RandomSampler, 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
+from vllm.v1.worker.gpu.sample.watermark import repeated_context_mask
 
 logger = init_logger(__name__)
 
@@ -20,10 +22,14 @@ class GPUWatermarkSampler(Sampler):
         self,
         watermarker: Watermarker,
         *args,
+        deduplicate_contexts: Literal["none", "single_turn", "all"] = "single_turn",
+        deduplicate_contexts_max_history: int | None = 8192,
         **kwargs,
     ) -> None:
         super().__init__(*args, **kwargs)
         self.watermarker = watermarker
+        self.deduplicate_contexts = deduplicate_contexts
+        self.deduplicate_contexts_max_history = deduplicate_contexts_max_history
         self.watermarking = UvaBackedTensor(
             self.sampling_states.max_num_reqs, dtype=torch.bool
         )
@@ -72,37 +78,38 @@ class GPUWatermarkSampler(Sampler):
 
         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,
+        repeated_contexts = None
+        if self.deduplicate_contexts != "none":
+            repeated_contexts = self._get_repeated_contexts(
+                expanded_idx_mapping, contexts
             )
 
+        temperatures = self.sampling_states.temperature.gpu[expanded_idx_mapping]
+        needs_mixed_sampling = repeated_contexts is not None or not np.all(enabled)
+        skip_mask = None
+        if needs_mixed_sampling:
+            watermarking = self.watermarking.gpu[expanded_idx_mapping] & (
+                temperatures != 0
+            )
+            if repeated_contexts is not None:
+                watermarking &= ~repeated_contexts
+            skip_mask = ~watermarking
+
+        random_sampler = RandomSampler(
+            expanded_idx_mapping=expanded_idx_mapping,
+            temperatures=self.sampling_states.temperature.gpu,
+            seeds=self.sampling_states.seeds.gpu,
+            positions=pos,
+            use_fp64=self.use_fp64_gumbel,
+        )
         output = self.watermarker.sample(
             processed_logits,
             contexts,
-            random_sample,
+            random_sampler,
+            skip_mask=skip_mask,
         )
-        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 = output.token_ids
+        output_logits = output.logits
         sampled = torch.where(
             temperatures == 0,
             processed_logits.argmax(dim=-1),
@@ -110,22 +117,78 @@ class GPUWatermarkSampler(Sampler):
         )
         return sampled, output_logits
 
-    def _get_contexts(self, expanded_idx_mapping: torch.Tensor) -> torch.Tensor:
+    def _get_repeated_contexts(
+        self,
+        expanded_idx_mapping: torch.Tensor,
+        contexts: torch.Tensor,
+    ) -> torch.Tensor:
+        return repeated_context_mask(
+            self.req_states.all_token_ids.gpu,
+            expanded_idx_mapping,
+            self.req_states.prompt_len.gpu,
+            self.req_states.total_len.gpu,
+            contexts,
+            self.deduplicate_contexts_max_history,
+            include_prompt=self.deduplicate_contexts == "all",
+            skip_partial_context=self.deduplicate_contexts == "all",
+        )
+
+    def _get_contexts(
+        self,
+        expanded_idx_mapping: torch.Tensor,
+        expanded_local_pos: torch.Tensor | None = None,
+        draft_sampled: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Build contexts for flattened verification rows.
+
+        With committed tokens ``[10, 11]``, drafts ``[20, 21]``, and width 2,
+        local positions ``[0, 1, 2]`` use contexts ``[10, 11]``, ``[11, 20]``,
+        and ``[20, 21]``. The first verification row precedes the first draft.
+        """
         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)
+        if expanded_local_pos is None:
+            expanded_local_pos = torch.zeros_like(req_indices)
+        else:
+            expanded_local_pos = expanded_local_pos.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 = (
+            total_lens.unsqueeze(-1) + expanded_local_pos.unsqueeze(-1) + offsets
         )
-        positions = positions.clamp_min(0)
-        contexts = self.req_states.all_token_ids.gpu[
-            safe_req_indices.unsqueeze(-1), positions
+        committed = positions < total_lens.unsqueeze(-1)
+        valid_committed = (
+            valid_reqs.unsqueeze(-1)
+            & committed
+            & (positions >= prompt_lens.unsqueeze(-1))
+        )
+        committed_contexts = self.req_states.all_token_ids.gpu[
+            safe_req_indices.unsqueeze(-1),
+            positions.clamp(0, self.req_states.all_token_ids.gpu.shape[1] - 1),
         ]
-        return torch.where(valid_positions, contexts, -1)
+        contexts = torch.where(valid_committed, committed_contexts, -1)
+        if draft_sampled is None:
+            return contexts
+
+        row_indices = torch.arange(
+            len(req_indices), dtype=torch.int64, device=req_indices.device
+        )
+        draft_offsets = positions - total_lens.unsqueeze(-1)
+        draft_indices = (
+            row_indices.unsqueeze(-1)
+            - expanded_local_pos.unsqueeze(-1)
+            + draft_offsets
+            + 1
+        )
+        valid_drafts = (
+            valid_reqs.unsqueeze(-1)
+            & ~committed
+            & (draft_offsets < expanded_local_pos.unsqueeze(-1))
+        )
+        draft_contexts = draft_sampled[draft_indices.clamp(0, len(draft_sampled) - 1)]
+        return torch.where(valid_drafts, draft_contexts, contexts)

vllm/v1/watermarking/gumbel.py

@@ -8,13 +8,14 @@ import warnings
 
 import torch
 
-from vllm.config.watermarking import WatermarkPRFName
+from vllm.config.watermarking import WatermarkPRFName, derive_watermark_key
 from vllm.v1.watermarking.detector import (
     WatermarkDetector,
 )
 from vllm.v1.watermarking.prfs import PhiloxPRF, WatermarkPRF, create_prf
 from vllm.v1.watermarking.watermarker import (
     RandomSampler,
+    SupportsSpeculativeDecoding,
     Watermarker,
     WatermarkSample,
 )
@@ -54,11 +55,10 @@ class GumbelWatermarker(Watermarker):
     def context_width(self) -> int:
         return self._context_width
 
-    def sample(
+    def _sample_watermarked(
         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
@@ -72,6 +72,110 @@ class GumbelWatermarker(Watermarker):
         noise = -torch.log(-torch.log(uniforms))
         return WatermarkSample(torch.argmax(logits + noise, dim=-1), logits)
 
+    def _try_sample_mixed(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
+        skip_mask: torch.Tensor,
+        random_sampler: RandomSampler,
+    ) -> WatermarkSample | None:
+        if type(self.prf) is not PhiloxPRF or logits.device.type != "cuda":
+            return None
+
+        from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
+
+        token_ids = philox_gumbel_sample(
+            logits,
+            contexts,
+            self.prf.key,
+            skip_mask=skip_mask,
+            expanded_idx_mapping=random_sampler.expanded_idx_mapping,
+            temperatures=random_sampler.temperatures,
+            seeds=random_sampler.seeds,
+            positions=random_sampler.positions,
+            use_fp64=random_sampler.use_fp64,
+        )
+        return WatermarkSample(token_ids, logits)
+
+
+class DualKeyGumbelWatermarker(Watermarker, SupportsSpeculativeDecoding):
+    def __init__(
+        self,
+        key: int,
+        context_width: int = 4,
+        prf: WatermarkPRFName = "philox",
+        alpha: float = 0.1,
+    ) -> None:
+        if not 0 <= alpha <= 1:
+            raise ValueError("alpha must be between 0 and 1")
+        self.alpha = alpha
+        self.draft_watermarker = GumbelWatermarker(
+            derive_watermark_key(key, b"key_a"), context_width, prf
+        )
+        self.target_watermarker = GumbelWatermarker(
+            derive_watermark_key(key, b"key_b"), context_width, prf
+        )
+        self._routing_logits_cache: dict[torch.device, torch.Tensor] = {}
+
+    def _routing_logits(self, device: torch.device) -> torch.Tensor:
+        # Building this tensor on every CUDA call stages a pageable H2D copy and
+        # synchronizes the stream.
+        cached = self._routing_logits_cache.get(device)
+        if cached is None:
+            cached = torch.tensor(
+                [1 - self.alpha, self.alpha],
+                dtype=torch.float32,
+                device=device,
+            ).log()
+            self._routing_logits_cache[device] = cached
+        return cached
+
+    @property
+    def context_width(self) -> int:
+        return self.draft_watermarker.context_width
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
+        random_sampler: RandomSampler | None = None,
+        skip_mask: torch.Tensor | None = None,
+    ) -> WatermarkSample:
+        if self.alpha == 0:
+            return self.draft_watermarker.sample(
+                logits, contexts, random_sampler, skip_mask
+            )
+        if self.alpha == 1:
+            return self.target_watermarker.sample(
+                logits, contexts, random_sampler, skip_mask
+            )
+        if random_sampler is None:
+            raise ValueError("dual-key Gumbel routing requires a random sampler")
+
+        key_a_sample = self.draft_watermarker.sample(
+            logits, contexts, random_sampler, skip_mask
+        )
+        key_b_sample = self.target_watermarker.sample(
+            logits, contexts, random_sampler, skip_mask
+        )
+        routing_logits = self._routing_logits(logits.device).expand(logits.shape[0], -1)
+        use_key_a = random_sampler(routing_logits) == 0
+        return WatermarkSample(
+            torch.where(
+                use_key_a,
+                key_a_sample.token_ids,
+                key_b_sample.token_ids,
+            ),
+            logits,
+        )
+
+    def _sample_watermarked(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
+    ) -> WatermarkSample:
+        raise ValueError("dual-key Gumbel routing requires a random sampler")
+
 
 class GumbelWatermarkDetector(WatermarkDetector):
     def __init__(
@@ -100,6 +204,61 @@ class GumbelWatermarkDetector(WatermarkDetector):
         return token_scores.sum().item()
 
 
+class DualKeyGumbelWatermarkDetector(GumbelWatermarkDetector):
+    """Detect a dual-key watermark with weighted early fusion.
+
+    ``alpha`` is the weight assigned to key B; key A receives ``1 - alpha``.
+    """
+
+    def __init__(
+        self,
+        key: int,
+        context_width: int = 4,
+        p_value_threshold: float = 0.01,
+        prf: WatermarkPRFName = "philox",
+        deduplicate_contexts: bool = True,
+        alpha: float = 0.2,
+    ) -> None:
+        if not 0 <= alpha <= 1:
+            raise ValueError("alpha must be between 0 and 1")
+        super().__init__(
+            derive_watermark_key(key, b"key_a"),
+            context_width,
+            p_value_threshold,
+            prf,
+            deduplicate_contexts,
+        )
+        self.key_b_prf = create_prf(prf, derive_watermark_key(key, b"key_b"))
+        self.alpha = alpha
+
+    def _score_tokens(
+        self, contexts: torch.Tensor, targets: torch.Tensor
+    ) -> torch.Tensor:
+        scores = torch.stack(
+            [
+                -torch.log1p(
+                    -prf.uniform(contexts, targets.unsqueeze(-1))
+                    .squeeze(-1)
+                    .to(torch.float64)
+                )
+                for prf in (self.prf, self.key_b_prf)
+            ],
+            dim=-1,
+        )
+        weights = scores.new_tensor([1 - self.alpha, self.alpha])
+        return scores @ weights
+
+    def _get_p_value(self, score: float, num_scored_tokens: int) -> float:
+        variance = (1 - self.alpha) ** 2 + self.alpha**2
+        if self.alpha in (0, 0.5, 1):
+            return _gamma_survival_integer_shape(
+                score / variance, round(num_scored_tokens / variance)
+            )
+        shape = torch.tensor(num_scored_tokens / variance, dtype=torch.float64)
+        scaled_score = torch.tensor(score / variance, dtype=torch.float64)
+        return torch.special.gammaincc(shape, scaled_score).item()
+
+
 def _validate_context_width(context_width: int) -> None:
     if context_width < 1:
         raise ValueError("context_width must be positive")

vllm/v1/watermarking/spec_decode.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.config.watermarking import WatermarkConfig
+from vllm.v1.watermarking.factory import create_watermarker
+from vllm.v1.watermarking.gumbel import GumbelWatermarker
+from vllm.v1.watermarking.prfs import PhiloxPRF
+from vllm.v1.watermarking.watermarker import (
+    SupportsSpeculativeDecoding,
+    Watermarker,
+)
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample
+
+
+class DraftWatermarker:
+    def __init__(
+        self,
+        watermarker: Watermarker,
+        max_num_reqs: int,
+        device: torch.device,
+    ) -> None:
+        self.watermarker = watermarker
+        self.contexts = torch.zeros(
+            max_num_reqs,
+            watermarker.context_width,
+            dtype=torch.int64,
+            device=device,
+        )
+        self.enabled = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
+
+    def prepare(self, contexts: torch.Tensor, enabled: torch.Tensor) -> None:
+        num_reqs = contexts.shape[0]
+        self.contexts[:num_reqs].copy_(contexts)
+        self.enabled[:num_reqs].copy_(enabled)
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        ordinary_sampled: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        temperature: torch.Tensor,
+    ) -> torch.Tensor:
+        request_temperatures = temperature[idx_mapping]
+        processed_logits = logits / torch.where(
+            request_temperatures == 0, 1, request_temperatures
+        ).unsqueeze(-1)
+        watermarked = self.watermarker.sample(
+            processed_logits,
+            self.contexts[: logits.shape[0]],
+        ).token_ids
+        enabled = self.enabled[: logits.shape[0]] & (request_temperatures != 0)
+        sampled = torch.where(enabled, watermarked, ordinary_sampled)
+        contexts = self.contexts[: logits.shape[0]]
+        contexts.copy_(torch.cat((contexts[:, 1:], sampled.unsqueeze(-1)), dim=-1))
+        return sampled
+
+
+def create_speculative_target_watermarker(watermarker: Watermarker) -> Watermarker:
+    if isinstance(watermarker, SupportsSpeculativeDecoding):
+        return watermarker.target_watermarker
+    return watermarker
+
+
+def create_speculative_draft_watermarker(
+    watermarker: Watermarker,
+    max_num_reqs: int,
+    device: torch.device,
+    allow_target_only: bool,
+) -> DraftWatermarker | None:
+    if isinstance(watermarker, SupportsSpeculativeDecoding):
+        return DraftWatermarker(watermarker.draft_watermarker, max_num_reqs, device)
+    if allow_target_only:
+        return None
+    raise ValueError(
+        f"The {type(watermarker).__name__} watermarking algorithm does not support "
+        "speculative decoding. Set allow_target_only_watermarking=true to "
+        "leave draft tokens unwatermarked."
+    )
+
+
+def _philox_key(watermarker: Watermarker) -> int:
+    if not (
+        isinstance(watermarker, GumbelWatermarker)
+        and type(watermarker.prf) is PhiloxPRF
+    ):
+        raise NotImplementedError(
+            "In-kernel watermarked recovery supports only the Philox PRF "
+            "GumbelWatermarker"
+        )
+    return watermarker.prf.key
+
+
+def _resolve_watermark_key(watermarker: Watermarker) -> int:
+    """Return the Philox recovery key."""
+    if isinstance(watermarker, SupportsSpeculativeDecoding):
+        raise ValueError(
+            f"{type(watermarker).__name__} keys the target role separately from "
+            "the draft. Pass create_speculative_target_watermarker(watermarker) "
+            "so the recovery draw carries the target's key."
+        )
+    return _philox_key(watermarker)
+
+
+def speculative_target_watermark_key(
+    watermark_config: WatermarkConfig | None,
+) -> int | None:
+    """Resolve the resample kernel's Philox key from configuration."""
+    if watermark_config is None:
+        return None
+
+    return _resolve_watermark_key(
+        create_speculative_target_watermarker(create_watermarker(watermark_config))
+    )
+
+
+def watermarked_rejection_sample(
+    target_logits: torch.Tensor,
+    draft_logits: torch.Tensor | None,
+    draft_sampled: torch.Tensor,
+    cu_num_logits: torch.Tensor,
+    pos: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    expanded_local_pos: torch.Tensor,
+    temperature: torch.Tensor,
+    seed: torch.Tensor,
+    num_speculative_steps: int,
+    contexts: torch.Tensor,
+    watermarking: torch.Tensor,
+    watermarker: Watermarker,
+    use_fp64: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Use the target key for rejection recovery and bonus tokens."""
+    assert contexts.shape == (target_logits.shape[0], watermarker.context_width)
+    return rejection_sample(
+        target_logits,
+        draft_logits,
+        draft_sampled,
+        cu_num_logits,
+        pos,
+        idx_mapping,
+        expanded_idx_mapping,
+        expanded_local_pos,
+        temperature,
+        seed,
+        num_speculative_steps,
+        use_fp64=use_fp64,
+        contexts=contexts,
+        watermarking=watermarking,
+        watermark_key=_resolve_watermark_key(watermarker),
+    )

vllm/v1/watermarking/watermarker.py

@@ -2,13 +2,12 @@
 # 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
+from typing import Protocol, runtime_checkable
 
 import torch
 
-RandomSampler: TypeAlias = Callable[[torch.Tensor], torch.Tensor]
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 
 
 @dataclass(frozen=True)
@@ -17,17 +16,81 @@ class WatermarkSample:
     logits: torch.Tensor
 
 
+@dataclass(frozen=True)
+class RandomSampler:
+    expanded_idx_mapping: torch.Tensor
+    temperatures: torch.Tensor
+    seeds: torch.Tensor
+    positions: torch.Tensor
+    use_fp64: bool = False
+
+    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
+        return gumbel_sample(
+            logits,
+            self.expanded_idx_mapping,
+            self.temperatures,
+            self.seeds,
+            self.positions,
+            apply_temperature=False,
+            is_drafting=False,
+            use_fp64=self.use_fp64,
+        )
+
+
 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,
+        random_sampler: RandomSampler | None = None,
+        skip_mask: torch.Tensor | None = None,
+    ) -> WatermarkSample:
+        if skip_mask is not None:
+            if random_sampler is None:
+                raise ValueError("skip_mask requires a random sampler")
+            mixed = self._try_sample_mixed(logits, contexts, skip_mask, random_sampler)
+            if mixed is not None:
+                return mixed
+
+        watermarked = self._sample_watermarked(logits, contexts)
+        if skip_mask is None:
+            return watermarked
+
+        assert random_sampler is not None
+        token_ids = torch.where(
+            skip_mask, random_sampler(logits), watermarked.token_ids
+        )
+        output_logits = watermarked.logits
+        if output_logits is not logits:
+            output_logits = torch.where(skip_mask.unsqueeze(-1), logits, output_logits)
+        return WatermarkSample(token_ids, output_logits)
+
+    @abstractmethod
+    def _sample_watermarked(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
     ) -> WatermarkSample:
         raise NotImplementedError
+
+    def _try_sample_mixed(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
+        skip_mask: torch.Tensor,
+        random_sampler: RandomSampler,
+    ) -> WatermarkSample | None:
+        return None
+
+
+@runtime_checkable
+class SupportsSpeculativeDecoding(Protocol):
+    """A watermarker compatible with standard speculative rejection sampling."""
+
+    draft_watermarker: Watermarker
+    target_watermarker: Watermarker

vllm/v1/worker/gpu/model_runner.py

@@ -79,6 +79,10 @@ from vllm.v1.outputs import (
 )
 from vllm.v1.watermarking import create_watermarker
 from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
+from vllm.v1.watermarking.spec_decode import (
+    create_speculative_target_watermarker,
+    speculative_target_watermark_key,
+)
 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
@@ -472,8 +476,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             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)
+                wm_config = self.vllm_config.watermark_config
+                watermarker = create_watermarker(wm_config)
+                if self.speculative_config is not None:
+                    watermarker = create_speculative_target_watermarker(watermarker)
+                self.sampler = GPUWatermarkSampler(
+                    watermarker,
+                    deduplicate_contexts=wm_config.deduplicate_contexts,
+                    deduplicate_contexts_max_history=(
+                        wm_config.deduplicate_contexts_max_history
+                    ),
+                    **sampler_kwargs,
+                )
             custom = self.model_state.custom_sampler(self.sampler)
 
             if custom:
@@ -484,6 +498,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     self.sampler,
                     self.speculative_config,
                     self.device,
+                    watermark_key=speculative_target_watermark_key(
+                        self.vllm_config.watermark_config
+                    ),
                 )
             self.prompt_logprobs_worker = PromptLogprobsWorker(
                 self.max_num_reqs,
@@ -860,6 +877,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             if hasattr(self.model, "get_mtp_target_hidden_states"):
                 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
                 spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]  # type: ignore[union-attr]
+            if isinstance(self.sampler, GPUWatermarkSampler):
+                self.speculator.prepare_watermarking(
+                    self.sampler._get_contexts(input_batch.idx_mapping),
+                    self.sampler.watermarking.gpu[input_batch.idx_mapping],
+                )
             with use_workspace_lane(self._draft_workspace_lane):
                 self.speculator.propose(
                     input_batch=input_batch,
@@ -2086,6 +2108,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             if hasattr(self.model, "get_mtp_target_hidden_states"):
                 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
                 spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]
+            if isinstance(self.sampler, GPUWatermarkSampler):
+                self.speculator.prepare_watermarking(
+                    self.sampler._get_contexts(input_batch.idx_mapping),
+                    self.sampler.watermarking.gpu[input_batch.idx_mapping],
+                )
             with use_workspace_lane(self._draft_workspace_lane):
                 draft_tokens = self.speculator.propose(
                     input_batch,

vllm/v1/worker/gpu/sample/watermark.py

@@ -25,6 +25,7 @@ from vllm.v1.watermarking.prfs.philox import (
 from vllm.v1.watermarking.prfs.philox import (
     _UINT32_MASK as _UINT32_MASK_VALUE,
 )
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax
 
 if HAS_TRITON:
     from triton.language import math as tl_math
@@ -42,6 +43,162 @@ _CONTEXT_DOMAIN = (
 _TOKEN_DOMAIN = tl.constexpr(_TOKEN_DOMAIN_VALUE) if HAS_TRITON else _TOKEN_DOMAIN_VALUE
 
 
+@triton.jit
+def _repeated_context_mask_kernel(
+    output_ptr,
+    all_token_ids_ptr,
+    all_token_ids_stride,
+    req_indices_ptr,
+    prompt_lens_ptr,
+    total_lens_ptr,
+    contexts_ptr,
+    context_stride,
+    CONTEXT_WIDTH: tl.constexpr,
+    MAX_HISTORY: tl.constexpr,
+    INCLUDE_PROMPT: tl.constexpr,
+    SKIP_PARTIAL_CONTEXT: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    req_idx = tl.load(req_indices_ptr + row).to(tl.int64)
+    if req_idx < 0:
+        tl.store(output_ptr + row, 0)
+        return
+    prompt_len = tl.load(prompt_lens_ptr + req_idx)
+    total_len = tl.load(total_lens_ptr + req_idx)
+    if SKIP_PARTIAL_CONTEXT and total_len - prompt_len < CONTEXT_WIDTH:
+        tl.store(output_ptr + row, 1)
+        return
+    sequence_start = 0 if INCLUDE_PROMPT else prompt_len
+    history_len = total_len - sequence_start
+
+    offsets = tl.arange(0, BLOCK)
+    repeated = tl.full((), 0, tl.int32)
+    scan_start = 0
+    if MAX_HISTORY > 0:
+        scan_start = tl.maximum(history_len - MAX_HISTORY, 0)
+    # Aligned loop bounds let the per-lane offsets fold into immediates.
+    aligned_start = (scan_start // BLOCK) * BLOCK
+    for block_start in tl.range(aligned_start, history_len, BLOCK):
+        previous_pos = block_start + offsets
+        in_window = (previous_pos >= scan_start) & (previous_pos < history_len)
+        matches = in_window
+        for offset in range(CONTEXT_WIDTH):
+            context_token = tl.load(contexts_ptr + row * context_stride + offset)
+            historical_pos = previous_pos + offset - CONTEXT_WIDTH
+            historical_token = tl.load(
+                all_token_ids_ptr
+                + req_idx * all_token_ids_stride
+                + sequence_start
+                + tl.maximum(historical_pos, 0),
+                mask=in_window & (historical_pos >= 0),
+                other=-1,
+            )
+            matches &= historical_token == context_token
+        repeated |= tl.max(matches.to(tl.int32), axis=0)
+    tl.store(output_ptr + row, repeated)
+
+
+def _repeated_context_mask_cpu(
+    all_token_ids: torch.Tensor,
+    req_indices: torch.Tensor,
+    prompt_lens: torch.Tensor,
+    total_lens: torch.Tensor,
+    contexts: torch.Tensor,
+    max_history: int | None = None,
+    include_prompt: bool = False,
+    skip_partial_context: bool = False,
+) -> torch.Tensor:
+    """Reference implementation; the parity tests check the Triton kernel against it."""
+    repeated = torch.zeros(len(req_indices), dtype=torch.bool)
+    for row, req_idx_tensor in enumerate(req_indices):
+        req_idx = int(req_idx_tensor)
+        if req_idx < 0:
+            continue
+        prompt_len = int(prompt_lens[req_idx])
+        total_len = int(total_lens[req_idx])
+        if skip_partial_context and total_len - prompt_len < contexts.shape[-1]:
+            repeated[row] = True
+            continue
+        sequence_start = 0 if include_prompt else prompt_len
+        history_tokens = all_token_ids[req_idx, sequence_start:total_len].tolist()
+        prefix = [-1] * contexts.shape[-1]
+        current_context = tuple(contexts[row].tolist())
+        history_start = (
+            max(0, len(history_tokens) - max_history) if max_history is not None else 0
+        )
+        for history_pos, token_id in enumerate(history_tokens):
+            if history_pos >= history_start and (
+                tuple(prefix[-contexts.shape[-1] :]) == current_context
+            ):
+                repeated[row] = True
+                break
+            prefix.append(token_id)
+    return repeated
+
+
+def repeated_context_mask(
+    all_token_ids: torch.Tensor,
+    req_indices: torch.Tensor,
+    prompt_lens: torch.Tensor,
+    total_lens: torch.Tensor,
+    contexts: torch.Tensor,
+    max_history: int | None = None,
+    include_prompt: bool = False,
+    skip_partial_context: bool = False,
+) -> torch.Tensor:
+    """Return, per row, whether the row's context already occurred in its history.
+
+    Args:
+        all_token_ids: `[max_num_reqs, max_model_len]` token ids of every request.
+        req_indices: Request slot per sampled row; -1 marks a padding row, which
+            is reported as not repeated.
+        prompt_lens: Prompt length per request slot.
+        total_lens: Prompt plus generated length per request slot.
+        contexts: `[num_rows, context_width]` context per row, padded with -1
+            before the start of the scanned history.
+        max_history: Number of most recent history positions searched, or
+            ``None`` for all of them. The window compared at each position
+            reaches `context_width` tokens further back.
+        include_prompt: Search the prompt as well as the generated tokens.
+        skip_partial_context: Mark contexts containing start padding so they use
+            ordinary sampling.
+    """
+    if max_history is not None and max_history < 1:
+        raise ValueError("max_history must be positive or None")
+    if all_token_ids.device.type == "cpu":
+        return _repeated_context_mask_cpu(
+            all_token_ids,
+            req_indices,
+            prompt_lens,
+            total_lens,
+            contexts,
+            max_history,
+            include_prompt,
+            skip_partial_context,
+        )
+
+    if contexts.stride(-1) != 1:
+        contexts = contexts.contiguous()
+    repeated = torch.empty(len(req_indices), dtype=torch.bool, device=contexts.device)
+    _repeated_context_mask_kernel[(len(req_indices),)](
+        repeated,
+        all_token_ids,
+        all_token_ids.stride(0),
+        req_indices,
+        prompt_lens,
+        total_lens,
+        contexts,
+        contexts.stride(0),
+        CONTEXT_WIDTH=contexts.shape[-1],
+        MAX_HISTORY=0 if max_history is None else max_history,
+        INCLUDE_PROMPT=include_prompt,
+        SKIP_PARTIAL_CONTEXT=skip_partial_context,
+        BLOCK=512,
+    )
+    return repeated
+
+
 @triton.jit
 def _mulhilo32(multiplier: tl.constexpr, value):
     high = tl_math.umulhi(multiplier, value)
@@ -73,12 +230,100 @@ def _uint32_to_uniform(value):
     return (scaled.to(tl.uint32, bitcast=True) - 1).to(tl.float32, bitcast=True)
 
 
+@triton.jit
+def _philox_context_state(
+    contexts_row_ptr,
+    key_0,
+    key_1,
+    CONTEXT_WIDTH: tl.constexpr,
+):
+    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_row_ptr + offset).to(tl.uint32)
+        if offset + 1 < CONTEXT_WIDTH:
+            context_1 = tl.load(contexts_row_ptr + 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_row_ptr + 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_row_ptr + 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,
+        )
+    return state_0, state_1, state_2, state_3
+
+
+@triton.jit
+def _philox_candidate_words(groups, state_0, state_1, state_2, state_3, key_0, key_1):
+    candidate_words = groups.to(tl.uint32)
+    vector_zero = candidate_words * 0
+    return _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,
+    )
+
+
+@triton.jit
+def _philox_gumbel_from_logits(logits, word):
+    logits = tl.where(logits != logits, float("-inf"), logits)
+    uniform = _uint32_to_uniform(word)
+    return logits - tl.log(-tl.log(uniform))
+
+
 @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))
+    return _philox_gumbel_from_logits(logits, output)
+
+
+@triton.jit
+def philox_gumbel_block_argmax(
+    logits,
+    mask,
+    block_idx,
+    contexts_row_ptr,
+    key_0,
+    key_1,
+    CONTEXT_WIDTH: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Apply the detector-compatible PRF to one vocabulary block."""
+    tl.static_assert(BLOCK_SIZE % 4 == 0)
+    state_0, state_1, state_2, state_3 = _philox_context_state(
+        contexts_row_ptr, key_0, key_1, CONTEXT_WIDTH
+    )
+    groups = block_idx * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4)
+    output_0, output_1, output_2, output_3 = _philox_candidate_words(
+        groups, state_0, state_1, state_2, state_3, key_0, key_1
+    )
+    # Lay the per-group words out in token order: 4g+0, 4g+1, 4g+2, 4g+3.
+    words = tl.interleave(
+        tl.interleave(output_0, output_2),
+        tl.interleave(output_1, output_3),
+    )
+    values = tl.where(
+        mask,
+        _philox_gumbel_from_logits(logits.to(tl.float32), words),
+        float("-inf"),
+    )
+    return tl.max(values, axis=0, return_indices=True)
 
 
 @triton.jit
@@ -100,60 +345,63 @@ def _philox_gumbel_kernel(
     logits_stride,
     contexts_ptr,
     context_stride,
+    skip_mask_ptr,
+    expanded_idx_mapping_ptr,
+    seeds_ptr,
+    pos_ptr,
+    temp_ptr,
     key_0_value,
     key_1_value,
     vocab_size,
     CONTEXT_WIDTH: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
+    USE_FP64: 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)
 
+    if skip_mask_ptr is not None:
+        req_state_idx = tl.load(expanded_idx_mapping_ptr + row).to(tl.int64)
+        valid_req = req_state_idx >= 0
+        if tl.load(skip_mask_ptr + row):
+            temp = tl.load(temp_ptr + req_state_idx, mask=valid_req, other=0.0).to(
+                tl.float32
+            )
+            seed = tl.load(seeds_ptr + req_state_idx, mask=valid_req, other=0)
+            pos = tl.load(pos_ptr + row)
+            candidate = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = candidate < vocab_size
+            logits_row = logits_ptr + row * logits_stride
+            logits = tl.load(logits_row + candidate, mask=mask, other=float("-inf")).to(
+                tl.float32
+            )
+            value, index = gumbel_noised_argmax(
+                logits,
+                candidate,
+                mask,
+                seed,
+                pos,
+                temp,
+                IS_DRAFTING=False,
+                USE_FP64=USE_FP64,
+                APPLY_TEMPERATURE=False,
+            )
+            token_id = block_index * BLOCK_SIZE + index
+            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)
+            return
+
     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,
+    state_0, state_1, state_2, state_3 = _philox_context_state(
+        contexts_ptr + row * context_stride, key_0, key_1, CONTEXT_WIDTH
+    )
+    output_0, output_1, output_2, output_3 = _philox_candidate_words(
+        groups, state_0, state_1, state_2, state_3, key_0, key_1
     )
     candidate_0 = groups * 4
     candidate_1 = candidate_0 + 1
@@ -184,17 +432,42 @@ def _philox_gumbel_kernel(
 
 
 def philox_gumbel_sample(
-    logits: torch.Tensor, contexts: torch.Tensor, key: int
+    logits: torch.Tensor,
+    contexts: torch.Tensor,
+    key: int,
+    *,
+    skip_mask: torch.Tensor | None = None,
+    expanded_idx_mapping: torch.Tensor | None = None,
+    temperatures: torch.Tensor | None = None,
+    seeds: torch.Tensor | None = None,
+    positions: torch.Tensor | None = None,
+    use_fp64: bool = False,
 ) -> torch.Tensor:
+    sampling_state = (expanded_idx_mapping, temperatures, seeds, positions)
+    if skip_mask is None and any(value is not None for value in sampling_state):
+        raise ValueError("sampling state requires skip_mask")
+    if skip_mask is not None and any(value is None for value in sampling_state):
+        raise ValueError("skip_mask requires complete sampling state")
+
     if logits.stride(-1) != 1:
         logits = logits.contiguous()
     if contexts.stride(-1) != 1:
         contexts = contexts.contiguous()
+    if skip_mask is not None:
+        skip_mask = skip_mask.contiguous()
+    if expanded_idx_mapping is not None:
+        expanded_idx_mapping = expanded_idx_mapping.contiguous()
+    if positions is not None:
+        positions = positions.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)
+    local_max = logits.new_empty(
+        num_tokens,
+        num_blocks,
+        dtype=torch.float64 if use_fp64 else torch.float32,
+    )
     _philox_gumbel_kernel[(num_tokens, num_blocks)](
         local_argmax,
         local_argmax.stride(0),
@@ -204,11 +477,17 @@ def philox_gumbel_sample(
         logits.stride(0),
         contexts,
         contexts.stride(0),
+        skip_mask,
+        expanded_idx_mapping,
+        seeds,
+        positions,
+        temperatures,
         key & _UINT32_MASK_VALUE,
         key >> 32,
         vocab_size,
         CONTEXT_WIDTH=contexts.shape[-1],
         BLOCK_SIZE=block_size,
+        USE_FP64=use_fp64,
     )
     max_block_index = local_max.argmax(dim=-1, keepdim=True)
     return local_argmax.gather(dim=-1, index=max_block_index).view(-1)

vllm/v1/worker/gpu/spec_decode/dspark/speculator.py

@@ -137,7 +137,7 @@ class DSparkSpeculator(DFlashSpeculator):
 
         # sample_pos is the predicted token's position P. Sampling keys a draw
         # by the position before the sampled token, P-1.
-        return gumbel_sample(
+        sampled = gumbel_sample(
             logits,
             idx_map,
             self.temperature,
@@ -149,6 +149,14 @@ class DSparkSpeculator(DFlashSpeculator):
             logits_cache_col=self._step_cols[step],
             use_fp64=self.use_fp64_gumbel,
         )
+        if self.draft_watermarker is not None:
+            sampled = self.draft_watermarker.sample(
+                logits,
+                sampled,
+                idx_map,
+                self.temperature,
+            )
+        return sampled
 
     def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
         if self._draft_topk is not None:

vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable, Iterator
+from typing import Any
 
 import numpy as np
 import torch
@@ -10,6 +11,7 @@ from vllm.config.model import PROCESSED_LOGPROBS_MODES
 from vllm.triton_utils import tl, triton
 from vllm.v1.outputs import LogprobsTensors
 from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
+from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
 from vllm.v1.worker.gpu.input_batch import (
     InputBatch,
     get_num_sampled_and_rejected,
@@ -78,8 +80,13 @@ class RejectionSampler:
         sampler: Sampler,
         spec_config: SpeculativeConfig,
         device: torch.device,
+        *,
+        watermark_key: int | None = None,
     ):
         self.sampler = sampler
+        self.watermark_key = watermark_key
+        if watermark_key is not None:
+            assert isinstance(sampler, GPUWatermarkSampler)
         self.num_speculative_steps = spec_config.num_speculative_tokens
         self.enable_adaptive_verification = spec_config.enable_adaptive_verification
         rejection_sample_method = spec_config.rejection_sample_method
@@ -141,6 +148,23 @@ class RejectionSampler:
             in ("raw_logits", "processed_logits"),
         )
 
+    def _watermarking_kwargs(
+        self,
+        draft_sampled: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        expanded_local_pos: torch.Tensor,
+    ) -> dict[str, Any]:
+        if self.watermark_key is None:
+            return {}
+        assert isinstance(self.sampler, GPUWatermarkSampler)
+        return {
+            "contexts": self.sampler._get_contexts(
+                expanded_idx_mapping, expanded_local_pos, draft_sampled
+            ),
+            "watermarking": self.sampler.watermarking.gpu,
+            "watermark_key": self.watermark_key,
+        }
+
     def _verify(
         self,
         logits: torch.Tensor,
@@ -177,6 +201,9 @@ class RejectionSampler:
             self.synthetic_conditional_rates,
             use_fp64=self.sampler.use_fp64_gumbel,
             use_block_verification=self.use_block_verification,
+            **self._watermarking_kwargs(
+                draft_sampled, expanded_idx_mapping, expanded_local_pos
+            ),
         )
         return processed_logits, sampled, num_sampled
 

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

@@ -4,6 +4,7 @@ import torch
 
 from vllm.triton_utils import tl, tldevice, triton
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32
+from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_block_argmax
 
 
 @triton.jit
@@ -693,6 +694,39 @@ def _rejection_kernel(
 
 
 @triton.jit
+def _seeded_resample_argmax(
+    residual_logits,
+    block,
+    mask,
+    resample_token_idx,
+    expanded_idx_mapping_ptr,
+    temp_ptr,
+    seed_ptr,
+    pos_ptr,
+    vocab_size,
+    USE_FP64: tl.constexpr,
+):
+    return gumbel_block_argmax(
+        residual_logits,
+        block,
+        mask,
+        resample_token_idx,
+        expanded_idx_mapping_ptr,
+        temp_ptr,
+        seed_ptr,
+        pos_ptr,
+        None,  # logits_cache_ptr
+        0,  # logits_cache_stride_0
+        0,  # logits_cache_stride_1
+        None,  # logits_cache_col_ptr
+        vocab_size,
+        IS_DRAFTING=False,
+        APPLY_TEMPERATURE=False,
+        USE_FP64=USE_FP64,
+    )
+
+
+@triton.jit(do_not_specialize=["watermark_key_0", "watermark_key_1"])
 def _resample_kernel(
     # [num_reqs, num_blocks]
     resampled_local_argmax_ptr,
@@ -727,11 +761,20 @@ def _resample_kernel(
     pos_ptr,
     # [num_logits]
     cumulative_log_p_ptr,
+    # [num_logits, CONTEXT_WIDTH]
+    contexts_ptr,
+    contexts_stride,
+    # [max_num_reqs], uint8 view of a bool tensor
+    watermarking_ptr,
+    watermark_key_0,
+    watermark_key_1,
     vocab_size,
     BLOCK_SIZE: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
     USE_FP64: tl.constexpr,
     USE_BLOCK_VERIFICATION: tl.constexpr,
+    CONTEXT_WIDTH: tl.constexpr,
+    WATERMARK: tl.constexpr,
 ):
     req_idx = tl.program_id(0)
     resample_idx = tl.load(rejected_step_ptr + req_idx)
@@ -763,10 +806,7 @@ def _resample_kernel(
         other=float("-inf"),
     ).to(tl.float32)
 
-    # Compute the residual logits to resample the rejected token from.
     if is_bonus or not is_valid_rejected_draft:
-        # Bonus token (no rejections) or -1 placeholder token. In either case,
-        # directly use the target logits.
         residual_logits = target_logits
     elif HAS_DRAFT_LOGITS:
         # draft_logits is stored pre-temperature, so apply scale first.
@@ -796,12 +836,7 @@ def _resample_kernel(
                 )
             target_log_probs += log_p_tau
         draft_log_probs = draft_logits - draft_lse
-        # Compute the residual:
-        #   r(x) = max(p(x) - q(x), 0)
-        # Gumbel sampling needs logits, so we compute it in log space:
-        #   log(r(x)) = log(max(exp(log_p(x)) - exp(log_q(x)), 0))
-        # The more numerically stable form is:
-        #   log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0))
+        # Compute log(max(p - q, 0)) without subtracting probabilities.
         ratio = tl.exp(draft_log_probs - target_log_probs)
         residual_logits = tl.where(
             ratio < 1.0,
@@ -809,38 +844,62 @@ def _resample_kernel(
             float("-inf"),
         ).to(tl.float32)
     else:
-        # One-hot draft. The residual is just the target distribution with
-        # the rejected draft token probability zeroed out.
-        # NOTE: During block verification, the residual becomes:
-        #   0                   if x == rejected_draft_token
-        #   p_tau * M_b(x) / Z  otherwise
-        # Therefore p_tau is a constant that cancels under normalization,
-        # and does not need to be applied.
+        # The block-verification factor is constant and cancels on normalization.
         residual_logits = tl.where(
             block != rejected_draft_token,
             target_logits,
             float("-inf"),
         ).to(tl.float32)
 
-    # Resample the rejected/bonus token.
-    value, idx = gumbel_block_argmax(
-        residual_logits,
-        block,
-        mask,
-        resample_token_idx,
-        expanded_idx_mapping_ptr,
-        temp_ptr,
-        seed_ptr,
-        pos_ptr,
-        None,  # logits_cache_ptr
-        0,  # logits_cache_stride_0
-        0,  # logits_cache_stride_1
-        None,  # logits_cache_col_ptr
-        vocab_size,
-        IS_DRAFTING=False,
-        APPLY_TEMPERATURE=False,
-        USE_FP64=USE_FP64,
-    )
+    if WATERMARK:
+        # Padded and greedy rows retain the stock draw.
+        is_watermarked = (
+            tl.load(
+                watermarking_ptr + req_state_idx,
+                mask=req_state_idx >= 0,
+                other=0,
+            )
+            != 0
+        )
+        if is_watermarked & (temp != 0.0):
+            watermark_value, idx = philox_gumbel_block_argmax(
+                residual_logits,
+                mask,
+                block_idx,
+                contexts_ptr + resample_token_idx * contexts_stride,
+                watermark_key_0.to(tl.uint32),
+                watermark_key_1.to(tl.uint32),
+                CONTEXT_WIDTH,
+                BLOCK_SIZE,
+            )
+            # Detector compatibility fixes the keyed draw at fp32.
+            value = watermark_value.to(tl.float64) if USE_FP64 else watermark_value
+        else:
+            value, idx = _seeded_resample_argmax(
+                residual_logits,
+                block,
+                mask,
+                resample_token_idx,
+                expanded_idx_mapping_ptr,
+                temp_ptr,
+                seed_ptr,
+                pos_ptr,
+                vocab_size,
+                USE_FP64=USE_FP64,
+            )
+    else:
+        value, idx = _seeded_resample_argmax(
+            residual_logits,
+            block,
+            mask,
+            resample_token_idx,
+            expanded_idx_mapping_ptr,
+            temp_ptr,
+            seed_ptr,
+            pos_ptr,
+            vocab_size,
+            USE_FP64=USE_FP64,
+        )
     token_id = block_idx * BLOCK_SIZE + idx
     tl.store(
         resampled_local_argmax_ptr
@@ -947,6 +1006,9 @@ def rejection_sample(
     synthetic_conditional_rates: torch.Tensor | None = None,
     use_fp64: bool = False,
     use_block_verification: bool = False,
+    contexts: torch.Tensor | None = None,
+    watermarking: torch.Tensor | None = None,
+    watermark_key: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert target_logits.ndim == 2 and target_logits.stride(-1) == 1
     assert draft_logits is None or (
@@ -954,6 +1016,30 @@ def rejection_sample(
     )
     num_reqs = cu_num_logits.shape[0] - 1
     num_logits, vocab_size = target_logits.shape
+
+    watermark = contexts is not None
+    assert watermark == (watermarking is not None) == (watermark_key is not None), (
+        "contexts, watermarking and watermark_key must be set together."
+    )
+    contexts_stride = 0
+    context_width = 1
+    watermarking_bytes: torch.Tensor | None = None
+    watermark_key_0 = 0
+    watermark_key_1 = 0
+    if contexts is not None:
+        assert watermarking is not None and watermark_key is not None
+        assert contexts.ndim == 2 and contexts.shape[0] == num_logits
+        # Context words are hashed as uint32, so int32 (the request-state token
+        # dtype) and int64 both land on the same PRF stream, -1 included.
+        assert contexts.dtype in (torch.int32, torch.int64)
+        if contexts.stride(-1) != 1:
+            contexts = contexts.contiguous()
+        assert watermarking.ndim == 1 and watermarking.dtype == torch.bool
+        contexts_stride = contexts.stride(0)
+        context_width = contexts.shape[-1]
+        watermarking_bytes = watermarking.view(torch.uint8)
+        watermark_key_0 = watermark_key & 0xFFFFFFFF
+        watermark_key_1 = watermark_key >> 32
     draft_logits_stride_0 = 0
     draft_logits_stride_1 = 0
     if has_draft_logits := draft_logits is not None:
@@ -1165,11 +1251,18 @@ def rejection_sample(
         seed,
         pos,
         cumulative_log_p,
+        contexts,
+        contexts_stride,
+        watermarking_bytes,
+        watermark_key_0,
+        watermark_key_1,
         vocab_size,
         BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
         HAS_DRAFT_LOGITS=has_draft_logits,
         USE_FP64=use_fp64,
         USE_BLOCK_VERIFICATION=use_block_verification,
+        CONTEXT_WIDTH=context_width,
+        WATERMARK=watermark,
     )
 
     # Insert the resampled tokens into the output sampled.

vllm/v1/worker/gpu/spec_decode/speculator.py

@@ -17,6 +17,11 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.models import supports_multimodal_embeddings
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.watermarking import create_watermarker
+from vllm.v1.watermarking.spec_decode import (
+    DraftWatermarker,
+    create_speculative_draft_watermarker,
+)
 from vllm.v1.worker.gpu.attn_utils import (
     build_attn_metadata,
     init_attn_backend,
@@ -168,6 +173,17 @@ class DraftModelSpeculator(BaseSpeculator):
                 device=device,
             )
 
+        self.draft_watermarker: DraftWatermarker | None = None
+        watermark_config = getattr(vllm_config, "watermark_config", None)
+        if watermark_config is not None:
+            watermarker = create_watermarker(watermark_config)
+            self.draft_watermarker = create_speculative_draft_watermarker(
+                watermarker,
+                self.max_num_reqs,
+                device,
+                watermark_config.allow_target_only_watermarking,
+            )
+
         self.supports_mm_inputs = False
         self.pcp_manager: PCPManager | None = None
 
@@ -373,7 +389,7 @@ class DraftModelSpeculator(BaseSpeculator):
     ) -> torch.Tensor:
         if draft_logits is not None:
             logits = self.model.compute_logits(hidden_states)
-            return gumbel_sample(
+            sampled = gumbel_sample(
                 logits,
                 idx_mapping,
                 temperature,
@@ -385,8 +401,23 @@ class DraftModelSpeculator(BaseSpeculator):
                 logits_cache_col=draft_step,
                 use_fp64=self.use_fp64_gumbel,
             )
+            if self.draft_watermarker is not None:
+                sampled = self.draft_watermarker.sample(
+                    logits,
+                    sampled,
+                    idx_mapping,
+                    temperature,
+                )
+            return sampled
         return self._greedy_sample_draft(hidden_states)
 
+    def prepare_watermarking(
+        self, contexts: torch.Tensor, watermarking: torch.Tensor
+    ) -> None:
+        if self.draft_watermarker is None:
+            return
+        self.draft_watermarker.prepare(contexts, watermarking)
+
     def _copy_request_inputs(
         self,
         num_reqs: int,