| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.platforms import current_platform |
| 8 | from vllm.v1.watermarking import ( |
| 9 | DualKeyGumbelWatermarkDetector, |
| 10 | GumbelWatermarkDetector, |
| 11 | GumbelWatermarker, |
| 12 | derive_watermark_key, |
| 13 | ) |
| 14 | from vllm.v1.watermarking.gumbel import _gamma_survival_integer_shape |
| 15 | from vllm.v1.watermarking.watermarker import RandomSampler |
| 16 | from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample |
| 17 | from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample |
| 18 | |
| 19 | |
| 20 | def test_gamma_survival_integer_shape(): |
| 21 | assert _gamma_survival_integer_shape(0.0, 4) == 1.0 |
| 22 | assert _gamma_survival_integer_shape(2.0, 1) == pytest.approx(0.1353352832) |
| 23 | |
| 24 | |
| 25 | def test_detector_deduplicates_repeated_prf_inputs_by_default(): |
| 26 | detection = GumbelWatermarkDetector(key=42, context_width=1).detect([1, 1, 1, 1, 1]) |
| 27 | |
| 28 | assert detection.num_scored_tokens == 2 |
| 29 | |
| 30 | |
| 31 | def test_detector_can_score_repeated_prf_inputs(): |
| 32 | detection = GumbelWatermarkDetector( |
| 33 | key=42, context_width=1, deduplicate_contexts=False |
| 34 | ).detect([1, 1, 1, 1, 1]) |
| 35 | |
| 36 | assert detection.num_scored_tokens == 5 |
| 37 | |
| 38 | |
| 39 | def test_detector_deduplicates_context_even_when_target_differs(): |
| 40 | detection = GumbelWatermarkDetector(key=42, context_width=1).detect([1, 2, 1, 3]) |
| 41 | |
| 42 | assert detection.num_scored_tokens == 3 |
| 43 | |
| 44 | |
| 45 | def test_dual_key_detector_scores_each_token_against_both_keys(): |
| 46 | token_ids = [1, 2, 3, 4, 5] |
| 47 | key_a_detector = GumbelWatermarkDetector( |
| 48 | key=derive_watermark_key(42, b"key_a"), |
| 49 | context_width=1, |
| 50 | deduplicate_contexts=False, |
| 51 | ) |
| 52 | key_b_detector = GumbelWatermarkDetector( |
| 53 | key=derive_watermark_key(42, b"key_b"), |
| 54 | context_width=1, |
| 55 | deduplicate_contexts=False, |
| 56 | ) |
| 57 | detector = DualKeyGumbelWatermarkDetector( |
| 58 | key=42, |
| 59 | context_width=1, |
| 60 | deduplicate_contexts=False, |
| 61 | alpha=0.5, |
| 62 | ) |
| 63 | |
| 64 | key_a = key_a_detector.detect(token_ids) |
| 65 | key_b = key_b_detector.detect(token_ids) |
| 66 | dual = detector.detect(token_ids) |
| 67 | |
| 68 | assert dual.score == pytest.approx((key_a.score + key_b.score) / 2) |
| 69 | assert dual.p_value == pytest.approx( |
| 70 | _gamma_survival_integer_shape(dual.score * 2, dual.num_scored_tokens * 2) |
| 71 | ) |
| 72 | |
| 73 | |
| 74 | @pytest.mark.parametrize("alpha", [0.2, 0.7]) |
| 75 | def test_dual_key_detector_weights_and_recalibrates_scores(alpha: float): |
| 76 | token_ids = [1, 2, 3, 4, 5] |
| 77 | key_a = GumbelWatermarkDetector( |
| 78 | key=derive_watermark_key(42, b"key_a"), deduplicate_contexts=False |
| 79 | ).detect(token_ids) |
| 80 | key_b = GumbelWatermarkDetector( |
| 81 | key=derive_watermark_key(42, b"key_b"), deduplicate_contexts=False |
| 82 | ).detect(token_ids) |
| 83 | dual = DualKeyGumbelWatermarkDetector( |
| 84 | key=42, deduplicate_contexts=False, alpha=alpha |
| 85 | ).detect(token_ids) |
| 86 | |
| 87 | variance = (1 - alpha) ** 2 + alpha**2 |
| 88 | expected_p_value = torch.special.gammaincc( |
| 89 | torch.tensor(dual.num_scored_tokens / variance, dtype=torch.float64), |
| 90 | torch.tensor(dual.score / variance, dtype=torch.float64), |
| 91 | ).item() |
| 92 | |
| 93 | assert dual.score == pytest.approx((1 - alpha) * key_a.score + alpha * key_b.score) |
| 94 | assert dual.p_value == pytest.approx(expected_p_value) |
| 95 | |
| 96 | |
| 97 | def test_dual_key_detector_rejects_invalid_alpha(): |
| 98 | with pytest.raises(ValueError, match="alpha must be between 0 and 1"): |
| 99 | DualKeyGumbelWatermarkDetector(key=42, alpha=1.1) |
| 100 | |
| 101 | |
| 102 | def test_dual_key_detector_default_alpha(): |
| 103 | assert DualKeyGumbelWatermarkDetector(key=42).alpha == 0.2 |
| 104 | |
| 105 | |
| 106 | @pytest.mark.skipif( |
| 107 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 108 | ) |
| 109 | @pytest.mark.parametrize("key", [1, 2**32, 42, 15726070495360670683]) |
| 110 | @pytest.mark.parametrize("context_width", [1, 4, 16]) |
| 111 | def test_fused_watermarker_matches_cpu(key: int, context_width: int): |
| 112 | torch.manual_seed(0) |
| 113 | contexts = torch.randint(0, 248320, (32, context_width), dtype=torch.int64) |
| 114 | logits = torch.randn(32, 8193) |
| 115 | watermarker = GumbelWatermarker(key, context_width) |
| 116 | |
| 117 | expected = watermarker.sample(logits, contexts).token_ids |
| 118 | actual = watermarker.sample(logits.cuda(), contexts.cuda()).token_ids.cpu() |
| 119 | |
| 120 | torch.testing.assert_close(actual, expected, rtol=0, atol=0) |
| 121 | |
| 122 | |
| 123 | @pytest.mark.skipif( |
| 124 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 125 | ) |
| 126 | def test_fused_watermarker_handles_nan_logits(): |
| 127 | contexts = torch.zeros((2, 4), dtype=torch.int64, device="cuda") |
| 128 | logits = torch.full((2, 1025), float("nan"), device="cuda") |
| 129 | watermarker = GumbelWatermarker(key=42, context_width=4) |
| 130 | |
| 131 | token_ids = watermarker.sample(logits, contexts).token_ids |
| 132 | |
| 133 | assert torch.all((token_ids >= 0) & (token_ids < logits.shape[-1])) |
| 134 | |
| 135 | |
| 136 | @pytest.mark.skipif( |
| 137 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 138 | ) |
| 139 | def test_fused_watermarker_handles_noncontiguous_inputs(): |
| 140 | contexts = torch.randint(0, 248320, (8, 8), dtype=torch.int64, device="cuda")[ |
| 141 | :, ::2 |
| 142 | ] |
| 143 | logits = torch.randn(8, 2050, device="cuda")[:, ::2] |
| 144 | watermarker = GumbelWatermarker(key=42, context_width=4) |
| 145 | |
| 146 | expected = watermarker.sample(logits.contiguous(), contexts.contiguous()).token_ids |
| 147 | actual = watermarker.sample(logits, contexts).token_ids |
| 148 | |
| 149 | torch.testing.assert_close(actual, expected, rtol=0, atol=0) |
| 150 | |
| 151 | |
| 152 | @pytest.mark.skipif( |
| 153 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 154 | ) |
| 155 | @pytest.mark.parametrize("use_fp64", [False, True]) |
| 156 | def test_philox_gumbel_sample_skip_mask_matches_separate_samplers(use_fp64: bool): |
| 157 | torch.manual_seed(0) |
| 158 | logits = torch.randn(8, 8193, device="cuda") |
| 159 | context_storage = torch.randint(0, 248320, (8, 8), dtype=torch.int64, device="cuda") |
| 160 | contexts = context_storage[:, ::2] |
| 161 | repeated_mask = torch.tensor( |
| 162 | [False, False, False, True, False, False, True, False], device="cuda" |
| 163 | ) |
| 164 | watermarking = torch.tensor( |
| 165 | [True, False, True, True, False, True, True, False], device="cuda" |
| 166 | ) |
| 167 | req_indices = torch.tensor([3, 1, 7, 0, 4, 2, 6, 5], device="cuda") |
| 168 | temperatures = torch.ones(8, dtype=torch.float32, device="cuda") |
| 169 | temperatures[2] = 0 |
| 170 | seeds = torch.arange(8, dtype=torch.int64, device="cuda") + 1000 |
| 171 | positions = torch.arange(8, dtype=torch.int64, device="cuda") + 100 |
| 172 | watermark_mask = ( |
| 173 | watermarking[req_indices] & (temperatures[req_indices] != 0) & ~repeated_mask |
| 174 | ) |
| 175 | |
| 176 | watermarked = philox_gumbel_sample(logits, contexts, 42) |
| 177 | ordinary = gumbel_sample( |
| 178 | logits, |
| 179 | req_indices, |
| 180 | temperatures, |
| 181 | seeds, |
| 182 | positions, |
| 183 | apply_temperature=False, |
| 184 | is_drafting=False, |
| 185 | use_fp64=use_fp64, |
| 186 | ) |
| 187 | expected = torch.where(watermark_mask, watermarked, ordinary) |
| 188 | actual = philox_gumbel_sample( |
| 189 | logits, |
| 190 | contexts, |
| 191 | 42, |
| 192 | skip_mask=~watermark_mask, |
| 193 | expanded_idx_mapping=req_indices, |
| 194 | temperatures=temperatures, |
| 195 | seeds=seeds, |
| 196 | positions=positions, |
| 197 | use_fp64=use_fp64, |
| 198 | ) |
| 199 | |
| 200 | torch.testing.assert_close(actual, expected, rtol=0, atol=0) |
| 201 | |
| 202 | |
| 203 | @pytest.mark.skipif( |
| 204 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 205 | ) |
| 206 | def test_philox_mixed_draft_sampling_caches_unprocessed_logits(): |
| 207 | torch.manual_seed(0) |
| 208 | num_rows, vocab_size, num_steps = 5, 4099, 3 |
| 209 | logits = torch.randn(num_rows, vocab_size, device="cuda", dtype=torch.float16) |
| 210 | temperatures = torch.tensor([0.0, 0.3, 0.5, 1.0, 2.0], device="cuda") |
| 211 | processed_logits = logits / torch.where( |
| 212 | temperatures == 0, 1, temperatures |
| 213 | ).unsqueeze(-1) |
| 214 | contexts = torch.randint(0, vocab_size, (num_rows, 4), device="cuda") |
| 215 | idx_mapping = torch.arange(num_rows, dtype=torch.int32, device="cuda") |
| 216 | seeds = torch.arange(num_rows, dtype=torch.int64, device="cuda") + 10 |
| 217 | positions = torch.arange(num_rows, dtype=torch.int64, device="cuda") + 20 |
| 218 | # The skipped rows must divide by a temperature other than 1, or the |
| 219 | # unprocessed logits the cache stores are the logits the sampler sees. |
| 220 | # Row 0 keeps a greedy row in the skipped set. |
| 221 | skip_mask = torch.tensor([True, True, False, False, True], device="cuda") |
| 222 | cols = torch.arange(num_rows, dtype=torch.int32, device="cuda") % num_steps |
| 223 | cache = torch.zeros( |
| 224 | num_rows, num_steps, vocab_size + 1, dtype=logits.dtype, device="cuda" |
| 225 | ) |
| 226 | |
| 227 | ordinary = gumbel_sample( |
| 228 | logits, |
| 229 | idx_mapping, |
| 230 | temperatures, |
| 231 | seeds, |
| 232 | positions, |
| 233 | apply_temperature=True, |
| 234 | is_drafting=True, |
| 235 | ) |
| 236 | watermarked = philox_gumbel_sample(processed_logits, contexts, 42) |
| 237 | expected = torch.where(skip_mask, ordinary, watermarked) |
| 238 | actual = ( |
| 239 | GumbelWatermarker(key=42, context_width=4) |
| 240 | .sample( |
| 241 | processed_logits, |
| 242 | contexts, |
| 243 | random_sampler=RandomSampler( |
| 244 | expanded_idx_mapping=idx_mapping, |
| 245 | temperatures=temperatures, |
| 246 | seeds=seeds, |
| 247 | positions=positions, |
| 248 | is_drafting=True, |
| 249 | logits_cache=cache, |
| 250 | logits_cache_col=cols, |
| 251 | logits_cache_source=logits, |
| 252 | ), |
| 253 | skip_mask=skip_mask, |
| 254 | ) |
| 255 | .token_ids |
| 256 | ) |
| 257 | |
| 258 | torch.testing.assert_close(actual, expected, rtol=0, atol=0) |
| 259 | stored = cache[torch.arange(num_rows, device="cuda"), cols.long(), :vocab_size] |
| 260 | assert torch.equal(stored.view(torch.int16), logits.view(torch.int16)) |
| 261 | assert not cache[:, :, -1].any() |