diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 1c2e3bb54f..f33c582554 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.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 diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index 7c601e13aa..f9ceda0547 100644 --- a/docs/features/watermarking.md +++ b/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 diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 71b9dbbab8..0262b42d11 100644 --- a/tests/engine/test_arg_utils.py +++ b/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( diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml new file mode 100644 index 0000000000..c1298bebfa --- /dev/null +++ b/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 diff --git a/tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt b/tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt new file mode 100644 index 0000000000..127b0b4750 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt @@ -0,0 +1 @@ +Qwen3-4B-FP8-DSpark-watermark-TP2.yaml diff --git a/tests/model_executor/test_spec_decode_rejection_warmup.py b/tests/model_executor/test_spec_decode_rejection_warmup.py new file mode 100644 index 0000000000..6abfb1a4d1 --- /dev/null +++ b/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) diff --git a/tests/model_executor/test_watermark_sample_warmup.py b/tests/model_executor/test_watermark_sample_warmup.py new file mode 100644 index 0000000000..73ecb95ee3 --- /dev/null +++ b/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 diff --git a/tests/test_config.py b/tests/test_config.py index 282f235122..e6d040caf5 100644 --- a/tests/test_config.py +++ b/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( diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 86787b53a1..a64b0eb6e9 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/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, + ) diff --git a/tests/watermarking/test_gumbel.py b/tests/watermarking/test_gumbel.py index e3b4125172..dd32bda627 100644 --- a/tests/watermarking/test_gumbel.py +++ b/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) diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index 0452957ffc..65ea9e93ad 100644 --- a/tests/watermarking/test_watermarking.py +++ b/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])) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 3c197df0da..3bdfe1f3ed 100644 --- a/vllm/config/vllm.py +++ b/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: diff --git a/vllm/config/watermarking.py b/vllm/config/watermarking.py index 2247dd8248..8701c7b72e 100644 --- a/vllm/config/watermarking.py +++ b/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 diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 2991dc1027..bbb0a66d53 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/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): diff --git a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py b/vllm/model_executor/warmup/spec_decode_rejection_warmup.py index 8cedbfc4db..6f6f08a2d5 100644 --- a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py +++ b/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( diff --git a/vllm/model_executor/warmup/watermark_sample_warmup.py b/vllm/model_executor/warmup/watermark_sample_warmup.py new file mode 100644 index 0000000000..d315f10310 --- /dev/null +++ b/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) diff --git a/vllm/v1/watermarking/__init__.py b/vllm/v1/watermarking/__init__.py index bb53981461..23e9fd20fd 100644 --- a/vllm/v1/watermarking/__init__.py +++ b/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", ] diff --git a/vllm/v1/watermarking/factory.py b/vllm/v1/watermarking/factory.py index d9111e7a9e..2bc1919aef 100644 --- a/vllm/v1/watermarking/factory.py +++ b/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}") diff --git a/vllm/v1/watermarking/gpu_sampler.py b/vllm/v1/watermarking/gpu_sampler.py index 7e943cc040..afa95ee878 100644 --- a/vllm/v1/watermarking/gpu_sampler.py +++ b/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) diff --git a/vllm/v1/watermarking/gumbel.py b/vllm/v1/watermarking/gumbel.py index 856d697aea..cacb620b5e 100644 --- a/vllm/v1/watermarking/gumbel.py +++ b/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") diff --git a/vllm/v1/watermarking/spec_decode.py b/vllm/v1/watermarking/spec_decode.py new file mode 100644 index 0000000000..c3675a3ab9 --- /dev/null +++ b/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), + ) diff --git a/vllm/v1/watermarking/watermarker.py b/vllm/v1/watermarking/watermarker.py index f00659393e..12570f4ae5 100644 --- a/vllm/v1/watermarking/watermarker.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 14e3823c04..0ce0288f31 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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, diff --git a/vllm/v1/worker/gpu/sample/watermark.py b/vllm/v1/worker/gpu/sample/watermark.py index a0105cc1b7..0a9c64cbf7 100644 --- a/vllm/v1/worker/gpu/sample/watermark.py +++ b/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) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index e9eca550ca..036c3014f3 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/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: diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index a9027e0550..609924e87f 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 29b20afea6..9d64c1ecbc 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/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. diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 4f20637d70..f670441219 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/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,