| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import numpy as np |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm import SamplingParams |
| 11 | from vllm.config.watermarking import WatermarkConfig |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.v1.watermarking import ( |
| 14 | DualKeyGumbelWatermarker, |
| 15 | SupportsSpeculativeDecoding, |
| 16 | create_watermarker, |
| 17 | derive_watermark_key, |
| 18 | ) |
| 19 | from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler |
| 20 | from vllm.v1.watermarking.gumbel import GumbelWatermarker |
| 21 | from vllm.v1.watermarking.spec_decode import ( |
| 22 | DraftWatermarker, |
| 23 | _resolve_watermark_key, |
| 24 | create_speculative_draft_watermarker, |
| 25 | create_speculative_target_watermarker, |
| 26 | speculative_target_watermark_key, |
| 27 | ) |
| 28 | from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample |
| 29 | from vllm.v1.worker.gpu.sample.sampler import Sampler |
| 30 | from vllm.v1.worker.gpu.sample.watermark import ( |
| 31 | draft_watermarking_mask, |
| 32 | philox_gumbel_sample, |
| 33 | repeated_context_mask, |
| 34 | ) |
| 35 | from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator |
| 36 | from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator |
| 37 | |
| 38 | |
| 39 | class StubWatermarker(Watermarker): |
| 40 | context_width = 1 |
| 41 | |
| 42 | def _sample_watermarked(self, logits, contexts): |
| 43 | return WatermarkSample(torch.tensor([7, 7]), logits + 10) |
| 44 | |
| 45 | |
| 46 | def _argmax_sampler(logits: torch.Tensor) -> torch.Tensor: |
| 47 | return logits.argmax(dim=-1) |
| 48 | |
| 49 | |
| 50 | @pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"]) |
| 51 | def test_watermarker_contract(algorithm: str): |
| 52 | watermarker = create_watermarker( |
| 53 | WatermarkConfig(algorithm=algorithm, key=42, context_width=4) |
| 54 | ) |
| 55 | logits = torch.zeros(2, 128) |
| 56 | contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]]) |
| 57 | |
| 58 | first = watermarker.sample(logits, contexts, _argmax_sampler) |
| 59 | second = watermarker.sample(logits, contexts, _argmax_sampler) |
| 60 | |
| 61 | assert first.token_ids.shape == (2,) |
| 62 | assert first.logits.shape == logits.shape |
| 63 | assert torch.equal(first.token_ids, second.token_ids) |
| 64 | |
| 65 | |
| 66 | @pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"]) |
| 67 | @pytest.mark.parametrize( |
| 68 | "config_overrides", |
| 69 | [ |
| 70 | {"deduplicate_contexts": "none"}, |
| 71 | {"deduplicate_contexts_max_history": 1023}, |
| 72 | ], |
| 73 | ) |
| 74 | def test_gumbel_config_warns_when_context_deduplication_is_weak( |
| 75 | monkeypatch, algorithm, config_overrides |
| 76 | ): |
| 77 | messages: list[str] = [] |
| 78 | monkeypatch.setattr( |
| 79 | "vllm.config.watermarking.logger.warning_once", |
| 80 | lambda message, *, scope: messages.append(message), |
| 81 | ) |
| 82 | |
| 83 | WatermarkConfig(key=42, algorithm=algorithm) |
| 84 | WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=1024) |
| 85 | WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=None) |
| 86 | WatermarkConfig(key=42, algorithm=algorithm, **config_overrides) |
| 87 | |
| 88 | assert messages == [ |
| 89 | ( |
| 90 | "Gumbel-max watermarking with context deduplication disabled " |
| 91 | "or limited to fewer than 1024 positions may increase the frequency of " |
| 92 | "degenerate generations, including repetition loops. Use " |
| 93 | "deduplicate_contexts='single_turn' or 'all' with " |
| 94 | "deduplicate_contexts_max_history at least 1024 or null to mitigate this." |
| 95 | ) |
| 96 | ] |
| 97 | |
| 98 | |
| 99 | def test_context_deduplication_history_can_be_unbounded(): |
| 100 | config = WatermarkConfig(key=42, deduplicate_contexts_max_history=None) |
| 101 | |
| 102 | assert config.deduplicate_contexts_max_history is None |
| 103 | |
| 104 | |
| 105 | def test_large_context_width_warns_but_is_allowed(): |
| 106 | config = WatermarkConfig(key=42, context_width=17) |
| 107 | |
| 108 | with pytest.warns(UserWarning, match="reduce robustness to edits"): |
| 109 | watermarker = create_watermarker(config) |
| 110 | |
| 111 | assert watermarker.context_width == 17 |
| 112 | |
| 113 | |
| 114 | def test_dual_key_watermarker_uses_domain_separated_keys(): |
| 115 | config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) |
| 116 | |
| 117 | watermarker = create_watermarker(config) |
| 118 | assert isinstance(watermarker, SupportsSpeculativeDecoding) |
| 119 | draft = watermarker.draft_watermarker |
| 120 | target = watermarker.target_watermarker |
| 121 | |
| 122 | assert isinstance(watermarker, DualKeyGumbelWatermarker) |
| 123 | assert watermarker.alpha == 0.1 |
| 124 | assert draft.prf.key == derive_watermark_key(42, b"key_a") |
| 125 | assert target.prf.key == derive_watermark_key(42, b"key_b") |
| 126 | assert target.prf.key != draft.prf.key |
| 127 | |
| 128 | |
| 129 | def test_dual_key_watermarker_routes_tokens_with_alpha(): |
| 130 | watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25) |
| 131 | logits = torch.zeros(2, 16) |
| 132 | contexts = torch.tensor([[1, 2], [3, 4]]) |
| 133 | key_a = watermarker.draft_watermarker.sample(logits, contexts, lambda _: None) |
| 134 | key_b = watermarker.target_watermarker.sample(logits, contexts, lambda _: None) |
| 135 | |
| 136 | def route(routing_logits): |
| 137 | torch.testing.assert_close( |
| 138 | routing_logits.softmax(dim=-1), |
| 139 | torch.tensor([[0.75, 0.25], [0.75, 0.25]]), |
| 140 | ) |
| 141 | return torch.tensor([0, 1]) |
| 142 | |
| 143 | sampled = watermarker.sample(logits, contexts, route) |
| 144 | |
| 145 | assert torch.equal( |
| 146 | sampled.token_ids, torch.stack([key_a.token_ids[0], key_b.token_ids[1]]) |
| 147 | ) |
| 148 | |
| 149 | |
| 150 | def test_dual_key_routing_logits_are_cached(): |
| 151 | alpha = 0.25 |
| 152 | watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha) |
| 153 | logits = torch.zeros(2, 16) |
| 154 | contexts = torch.tensor([[1, 2], [3, 4]]) |
| 155 | recorded: list[torch.Tensor] = [] |
| 156 | |
| 157 | def route(routing_logits): |
| 158 | recorded.append(routing_logits) |
| 159 | return torch.tensor([0, 1]) |
| 160 | |
| 161 | watermarker.sample(logits, contexts, route) |
| 162 | watermarker.sample(logits, contexts, route) |
| 163 | |
| 164 | expected = torch.tensor([1 - alpha, alpha], dtype=torch.float32).log() |
| 165 | expected = expected.expand(2, -1) |
| 166 | assert len(recorded) == 2 |
| 167 | assert all(torch.equal(seen, expected) for seen in recorded) |
| 168 | assert len(watermarker._routing_logits_cache) == 1 |
| 169 | |
| 170 | |
| 171 | @pytest.mark.parametrize("alpha", [0.0, 0.25, 1.0]) |
| 172 | def test_dual_key_watermarker_skips_masked_rows(alpha: float): |
| 173 | watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha) |
| 174 | torch.manual_seed(0) |
| 175 | logits = torch.randn(4, 32) |
| 176 | contexts = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]]) |
| 177 | skip_mask = torch.tensor([True, False, True, False]) |
| 178 | |
| 179 | unmasked = watermarker.sample(logits, contexts, _argmax_sampler) |
| 180 | masked = watermarker.sample(logits, contexts, _argmax_sampler, skip_mask) |
| 181 | |
| 182 | ordinary = _argmax_sampler(logits) |
| 183 | assert torch.equal(masked.token_ids[skip_mask], ordinary[skip_mask]) |
| 184 | assert torch.equal(masked.token_ids[~skip_mask], unmasked.token_ids[~skip_mask]) |
| 185 | assert not torch.equal(unmasked.token_ids[skip_mask], ordinary[skip_mask]) |
| 186 | |
| 187 | |
| 188 | def test_dual_key_watermarker_routing_requires_a_random_sampler(): |
| 189 | watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25) |
| 190 | |
| 191 | with pytest.raises(ValueError, match="random sampler"): |
| 192 | watermarker.sample(torch.zeros(2, 16), torch.tensor([[1, 2], [3, 4]])) |
| 193 | |
| 194 | |
| 195 | def test_speculative_decoding_uses_fixed_dual_key_roles(): |
| 196 | watermarker = create_watermarker( |
| 197 | WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25) |
| 198 | ) |
| 199 | |
| 200 | target = create_speculative_target_watermarker(watermarker) |
| 201 | draft = create_speculative_draft_watermarker( |
| 202 | watermarker, |
| 203 | max_num_reqs=1, |
| 204 | device=torch.device("cpu"), |
| 205 | allow_target_only=False, |
| 206 | num_speculative_steps=3, |
| 207 | deduplicate_contexts="all", |
| 208 | deduplicate_contexts_max_history=None, |
| 209 | ) |
| 210 | |
| 211 | assert target.prf.key == derive_watermark_key(42, b"key_b") |
| 212 | assert draft is not None |
| 213 | assert draft.watermarker.prf.key == derive_watermark_key(42, b"key_a") |
| 214 | assert draft.num_speculative_steps == 3 |
| 215 | assert draft.deduplicate_contexts == "all" |
| 216 | assert draft.deduplicate_contexts_max_history is None |
| 217 | |
| 218 | |
| 219 | def test_recovery_key_rejects_the_unsplit_dual_key_watermarker(): |
| 220 | watermarker = create_watermarker( |
| 221 | WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25) |
| 222 | ) |
| 223 | |
| 224 | with pytest.raises(ValueError, match="keys the target role separately"): |
| 225 | _resolve_watermark_key(watermarker) |
| 226 | |
| 227 | target = create_speculative_target_watermarker(watermarker) |
| 228 | assert _resolve_watermark_key(target) == derive_watermark_key(42, b"key_b") |
| 229 | |
| 230 | |
| 231 | @pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"]) |
| 232 | def test_config_key_resolution_matches_the_model_runner(algorithm: str): |
| 233 | config = WatermarkConfig(algorithm=algorithm, key=42, alpha=0.25) |
| 234 | runtime_key = _resolve_watermark_key( |
| 235 | create_speculative_target_watermarker(create_watermarker(config)) |
| 236 | ) |
| 237 | |
| 238 | assert speculative_target_watermark_key(config) == runtime_key |
| 239 | |
| 240 | |
| 241 | def test_config_key_resolution_returns_none_when_watermarking_is_disabled(): |
| 242 | assert speculative_target_watermark_key(None) is None |
| 243 | |
| 244 | |
| 245 | def test_target_only_speculative_watermarking_skips_draft_watermarker(): |
| 246 | watermarker = create_watermarker(WatermarkConfig(algorithm="gumbel", key=42)) |
| 247 | |
| 248 | assert not isinstance(watermarker, SupportsSpeculativeDecoding) |
| 249 | with pytest.raises(ValueError, match="does not support speculative decoding"): |
| 250 | create_speculative_draft_watermarker( |
| 251 | watermarker, |
| 252 | max_num_reqs=1, |
| 253 | device=torch.device("cpu"), |
| 254 | allow_target_only=False, |
| 255 | ) |
| 256 | assert ( |
| 257 | create_speculative_draft_watermarker( |
| 258 | watermarker, |
| 259 | max_num_reqs=1, |
| 260 | device=torch.device("cpu"), |
| 261 | allow_target_only=True, |
| 262 | ) |
| 263 | is None |
| 264 | ) |
| 265 | |
| 266 | |
| 267 | def test_dual_key_derivation_is_stable(): |
| 268 | assert derive_watermark_key(32, b"key_a") == 16368605726115524094 |
| 269 | assert derive_watermark_key(32, b"key_b") == 4799302812959726346 |
| 270 | |
| 271 | |
| 272 | def test_sampling_params_can_disable_watermarking(): |
| 273 | assert SamplingParams().watermarking |
| 274 | assert not SamplingParams.from_optional(watermarking=False).watermarking |
| 275 | |
| 276 | |
| 277 | def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch): |
| 278 | sampler = object.__new__(GPUWatermarkSampler) |
| 279 | sampler.watermarking = SimpleNamespace(np=np.ones(1, dtype=bool)) |
| 280 | messages: list[str] = [] |
| 281 | monkeypatch.setattr(Sampler, "add_request", lambda *args, **kwargs: None) |
| 282 | monkeypatch.setattr( |
| 283 | "vllm.v1.watermarking.gpu_sampler.logger.warning_once", messages.append |
| 284 | ) |
| 285 | |
| 286 | sampler.add_request(0, SamplingParams(temperature=0)) |
| 287 | sampler.add_request(0, SamplingParams(temperature=1)) |
| 288 | sampler.add_request(0, SamplingParams(temperature=0, watermarking=False)) |
| 289 | |
| 290 | assert messages == [ |
| 291 | ( |
| 292 | "Watermarking is enabled, but greedy decoding (temperature=0) cannot be " |
| 293 | "watermarked. This request will use ordinary greedy sampling." |
| 294 | ) |
| 295 | ] |
| 296 | |
| 297 | |
| 298 | def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch): |
| 299 | sampler = object.__new__(GPUWatermarkSampler) |
| 300 | sampler.watermarker = StubWatermarker() |
| 301 | sampler.deduplicate_contexts = "single_turn" |
| 302 | sampler.watermarking = SimpleNamespace( |
| 303 | np=np.array([True, False]), gpu=torch.tensor([True, False]) |
| 304 | ) |
| 305 | sampler.sampling_states = SimpleNamespace( |
| 306 | temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)), |
| 307 | seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)), |
| 308 | ) |
| 309 | sampler.use_fp64_gumbel = False |
| 310 | sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros( |
| 311 | 2, 1, dtype=torch.int64 |
| 312 | ) |
| 313 | sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: torch.zeros( |
| 314 | 2, dtype=torch.bool |
| 315 | ) |
| 316 | monkeypatch.setattr( |
| 317 | "vllm.v1.watermarking.watermarker.gumbel_sample", |
| 318 | lambda *args, **kwargs: torch.tensor([3, 4]), |
| 319 | ) |
| 320 | logits = torch.zeros(2, 8) |
| 321 | |
| 322 | sampled, output_logits = sampler._sample_random( |
| 323 | logits, |
| 324 | torch.tensor([0, 1]), |
| 325 | np.array([0, 1]), |
| 326 | torch.zeros(2, dtype=torch.int64), |
| 327 | None, |
| 328 | None, |
| 329 | False, |
| 330 | ) |
| 331 | |
| 332 | assert torch.equal(sampled, torch.tensor([7, 4])) |
| 333 | assert torch.equal(output_logits[0], torch.full((8,), 10.0)) |
| 334 | assert torch.equal(output_logits[1], logits[1]) |
| 335 | |
| 336 | |
| 337 | def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch): |
| 338 | sampler = object.__new__(GPUWatermarkSampler) |
| 339 | sampler.watermarker = StubWatermarker() |
| 340 | sampler.deduplicate_contexts = "single_turn" |
| 341 | sampler.watermarking = SimpleNamespace( |
| 342 | np=np.array([True, True]), gpu=torch.tensor([True, True]) |
| 343 | ) |
| 344 | sampler.sampling_states = SimpleNamespace( |
| 345 | temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)), |
| 346 | seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)), |
| 347 | ) |
| 348 | sampler.use_fp64_gumbel = False |
| 349 | sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros( |
| 350 | 2, 1, dtype=torch.int64 |
| 351 | ) |
| 352 | sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: ( |
| 353 | torch.tensor([True, False]) |
| 354 | ) |
| 355 | monkeypatch.setattr( |
| 356 | "vllm.v1.watermarking.watermarker.gumbel_sample", |
| 357 | lambda *args, **kwargs: torch.tensor([3, 4]), |
| 358 | ) |
| 359 | logits = torch.zeros(2, 8) |
| 360 | |
| 361 | sampled, output_logits = sampler._sample_random( |
| 362 | logits, |
| 363 | torch.tensor([0, 1]), |
| 364 | np.array([0, 1]), |
| 365 | torch.zeros(2, dtype=torch.int64), |
| 366 | None, |
| 367 | None, |
| 368 | False, |
| 369 | ) |
| 370 | |
| 371 | assert torch.equal(sampled, torch.tensor([3, 7])) |
| 372 | assert torch.equal(output_logits[0], logits[0]) |
| 373 | assert torch.equal(output_logits[1], torch.full((8,), 10.0)) |
| 374 | |
| 375 | |
| 376 | def test_gpu_sampler_can_disable_context_deduplication(monkeypatch): |
| 377 | sampler = object.__new__(GPUWatermarkSampler) |
| 378 | sampler.watermarker = StubWatermarker() |
| 379 | sampler.deduplicate_contexts = "none" |
| 380 | sampler.watermarking = SimpleNamespace( |
| 381 | np=np.array([True, True]), gpu=torch.tensor([True, True]) |
| 382 | ) |
| 383 | sampler.sampling_states = SimpleNamespace( |
| 384 | temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)), |
| 385 | seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)), |
| 386 | ) |
| 387 | sampler.use_fp64_gumbel = False |
| 388 | sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros( |
| 389 | 2, 1, dtype=torch.int64 |
| 390 | ) |
| 391 | sampler._get_repeated_contexts = lambda *args: pytest.fail( |
| 392 | "context deduplication should not run" |
| 393 | ) |
| 394 | logits = torch.zeros(2, 8) |
| 395 | |
| 396 | sampled, output_logits = sampler._sample_random( |
| 397 | logits, |
| 398 | torch.tensor([0, 1]), |
| 399 | np.array([0, 1]), |
| 400 | torch.zeros(2, dtype=torch.int64), |
| 401 | None, |
| 402 | None, |
| 403 | False, |
| 404 | ) |
| 405 | |
| 406 | assert torch.equal(sampled, torch.tensor([7, 7])) |
| 407 | assert torch.equal(output_logits, torch.full((2, 8), 10.0)) |
| 408 | |
| 409 | |
| 410 | def test_repeated_context_mask_ignores_prompt_tokens(): |
| 411 | all_token_ids = torch.tensor( |
| 412 | [ |
| 413 | [8, 9, 1, 2, 1, 2], |
| 414 | [3, 4, 1, 2, 3, 4], |
| 415 | [0, 0, 0, 0, 0, 0], |
| 416 | ], |
| 417 | dtype=torch.int32, |
| 418 | ) |
| 419 | req_indices = torch.tensor([0, 1, -1]) |
| 420 | prompt_lens = torch.tensor([2, 2, 0]) |
| 421 | total_lens = torch.tensor([6, 6, 0]) |
| 422 | contexts = torch.tensor([[1, 2], [3, 4], [-1, -1]]) |
| 423 | |
| 424 | repeated = repeated_context_mask( |
| 425 | all_token_ids, |
| 426 | req_indices, |
| 427 | prompt_lens, |
| 428 | total_lens, |
| 429 | contexts, |
| 430 | ) |
| 431 | |
| 432 | assert torch.equal(repeated, torch.tensor([True, False, False])) |
| 433 | |
| 434 | |
| 435 | def test_repeated_context_mask_can_include_prompt_tokens(): |
| 436 | all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]) |
| 437 | req_indices = torch.tensor([0]) |
| 438 | prompt_lens = torch.tensor([4]) |
| 439 | total_lens = torch.tensor([10]) |
| 440 | contexts = torch.tensor([[1, 2, 3, 4]]) |
| 441 | |
| 442 | single_turn = repeated_context_mask( |
| 443 | all_token_ids, req_indices, prompt_lens, total_lens, contexts |
| 444 | ) |
| 445 | all_history = repeated_context_mask( |
| 446 | all_token_ids, |
| 447 | req_indices, |
| 448 | prompt_lens, |
| 449 | total_lens, |
| 450 | contexts, |
| 451 | include_prompt=True, |
| 452 | ) |
| 453 | |
| 454 | assert not single_turn.item() |
| 455 | assert all_history.item() |
| 456 | |
| 457 | |
| 458 | def test_repeated_context_mask_can_skip_partial_contexts(): |
| 459 | all_token_ids = torch.tensor([[10, 11, 12, 13, 1, 2, 3, 4]], dtype=torch.int32) |
| 460 | req_indices = torch.tensor([0]) |
| 461 | prompt_lens = torch.tensor([4]) |
| 462 | contexts = torch.tensor([[-1, -1, -1, 1]]) |
| 463 | |
| 464 | partial = repeated_context_mask( |
| 465 | all_token_ids, |
| 466 | req_indices, |
| 467 | prompt_lens, |
| 468 | torch.tensor([5]), |
| 469 | contexts, |
| 470 | include_prompt=True, |
| 471 | skip_partial_context=True, |
| 472 | ) |
| 473 | complete = repeated_context_mask( |
| 474 | all_token_ids, |
| 475 | req_indices, |
| 476 | prompt_lens, |
| 477 | torch.tensor([8]), |
| 478 | torch.tensor([[1, 2, 3, 4]]), |
| 479 | include_prompt=True, |
| 480 | skip_partial_context=True, |
| 481 | ) |
| 482 | |
| 483 | assert partial.item() |
| 484 | assert not complete.item() |
| 485 | |
| 486 | |
| 487 | def test_gpu_sampler_contexts_are_completion_local_for_all_scopes(): |
| 488 | sampler = object.__new__(GPUWatermarkSampler) |
| 489 | sampler.watermarker = SimpleNamespace(context_width=4) |
| 490 | sampler.req_states = SimpleNamespace( |
| 491 | all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])), |
| 492 | prompt_len=SimpleNamespace(gpu=torch.tensor([4])), |
| 493 | total_len=SimpleNamespace(gpu=torch.tensor([6])), |
| 494 | ) |
| 495 | request_indices = torch.tensor([0]) |
| 496 | |
| 497 | sampler.deduplicate_contexts = "single_turn" |
| 498 | single_turn = sampler._get_contexts(request_indices) |
| 499 | sampler.deduplicate_contexts = "all" |
| 500 | all_history = sampler._get_contexts(request_indices) |
| 501 | |
| 502 | assert torch.equal(single_turn, torch.tensor([[-1, -1, 1, 2]])) |
| 503 | assert torch.equal(all_history, single_turn) |
| 504 | |
| 505 | |
| 506 | def test_gpu_sampler_all_history_skips_partial_context(): |
| 507 | sampler = object.__new__(GPUWatermarkSampler) |
| 508 | sampler.watermarker = SimpleNamespace(context_width=4) |
| 509 | sampler.req_states = SimpleNamespace( |
| 510 | all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])), |
| 511 | prompt_len=SimpleNamespace(gpu=torch.tensor([4])), |
| 512 | total_len=SimpleNamespace(gpu=torch.tensor([6])), |
| 513 | ) |
| 514 | sampler.deduplicate_contexts_max_history = None |
| 515 | request_indices = torch.tensor([0]) |
| 516 | contexts = sampler._get_contexts(request_indices) |
| 517 | |
| 518 | sampler.deduplicate_contexts = "single_turn" |
| 519 | single_turn = sampler._get_repeated_contexts(request_indices, contexts) |
| 520 | sampler.deduplicate_contexts = "all" |
| 521 | all_history = sampler._get_repeated_contexts(request_indices, contexts) |
| 522 | |
| 523 | assert not single_turn.item() |
| 524 | assert all_history.item() |
| 525 | |
| 526 | |
| 527 | def test_repeated_context_mask_respects_max_history(): |
| 528 | all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]) |
| 529 | req_indices = torch.tensor([0]) |
| 530 | prompt_lens = torch.tensor([0]) |
| 531 | total_lens = torch.tensor([10]) |
| 532 | contexts = torch.tensor([[1, 2, 3, 4]]) |
| 533 | |
| 534 | full_history = repeated_context_mask( |
| 535 | all_token_ids, req_indices, prompt_lens, total_lens, contexts |
| 536 | ) |
| 537 | last_six = repeated_context_mask( |
| 538 | all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=6 |
| 539 | ) |
| 540 | last_five = repeated_context_mask( |
| 541 | all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=5 |
| 542 | ) |
| 543 | |
| 544 | assert full_history.item() |
| 545 | assert last_six.item() |
| 546 | assert not last_five.item() |
| 547 | |
| 548 | |
| 549 | @pytest.mark.parametrize("max_history", [0, -1]) |
| 550 | def test_repeated_context_mask_rejects_non_positive_max_history(max_history: int): |
| 551 | with pytest.raises(ValueError, match="max_history must be positive or None"): |
| 552 | repeated_context_mask( |
| 553 | torch.tensor([[1, 2, 3, 4]], dtype=torch.int32), |
| 554 | torch.tensor([0], dtype=torch.int32), |
| 555 | torch.tensor([0], dtype=torch.int32), |
| 556 | torch.tensor([4], dtype=torch.int32), |
| 557 | torch.tensor([[1, 2, 3, 4]], dtype=torch.int32), |
| 558 | max_history=max_history, |
| 559 | ) |
| 560 | |
| 561 | |
| 562 | def test_gpu_sampler_wires_request_history_scope(): |
| 563 | sampler = object.__new__(GPUWatermarkSampler) |
| 564 | sampler.req_states = SimpleNamespace( |
| 565 | all_token_ids=SimpleNamespace( |
| 566 | gpu=torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]) |
| 567 | ), |
| 568 | prompt_len=SimpleNamespace(gpu=torch.tensor([4])), |
| 569 | total_len=SimpleNamespace(gpu=torch.tensor([10])), |
| 570 | ) |
| 571 | sampler.deduplicate_contexts_max_history = None |
| 572 | request_indices = torch.tensor([0]) |
| 573 | contexts = torch.tensor([[1, 2, 3, 4]]) |
| 574 | |
| 575 | sampler.deduplicate_contexts = "single_turn" |
| 576 | single_turn = sampler._get_repeated_contexts(request_indices, contexts) |
| 577 | sampler.deduplicate_contexts = "all" |
| 578 | all_history = sampler._get_repeated_contexts(request_indices, contexts) |
| 579 | |
| 580 | assert not single_turn.item() |
| 581 | assert all_history.item() |
| 582 | |
| 583 | |
| 584 | @pytest.mark.skipif( |
| 585 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 586 | ) |
| 587 | @pytest.mark.parametrize("max_history", [5, 6]) |
| 588 | @pytest.mark.parametrize("include_prompt", [False, True]) |
| 589 | def test_repeated_context_mask_max_history_accelerator_parity( |
| 590 | max_history: int, include_prompt: bool |
| 591 | ): |
| 592 | inputs = ( |
| 593 | torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]), |
| 594 | torch.tensor([0]), |
| 595 | torch.tensor([0]), |
| 596 | torch.tensor([10]), |
| 597 | torch.tensor([[1, 2, 3, 4]]), |
| 598 | ) |
| 599 | |
| 600 | expected = repeated_context_mask( |
| 601 | *inputs, max_history=max_history, include_prompt=include_prompt |
| 602 | ) |
| 603 | actual = repeated_context_mask( |
| 604 | *(value.cuda() for value in inputs), |
| 605 | max_history=max_history, |
| 606 | include_prompt=include_prompt, |
| 607 | ).cpu() |
| 608 | |
| 609 | assert torch.equal(actual, expected) |
| 610 | |
| 611 | |
| 612 | @pytest.mark.skipif( |
| 613 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 614 | ) |
| 615 | def test_repeated_context_mask_partial_context_accelerator_parity(): |
| 616 | inputs = ( |
| 617 | torch.tensor([[10, 11, 12, 13, 1]], dtype=torch.int32), |
| 618 | torch.tensor([0], dtype=torch.int32), |
| 619 | torch.tensor([4], dtype=torch.int32), |
| 620 | torch.tensor([5], dtype=torch.int32), |
| 621 | torch.tensor([[-1, -1, -1, 1]], dtype=torch.int64), |
| 622 | ) |
| 623 | |
| 624 | expected = repeated_context_mask( |
| 625 | *inputs, include_prompt=True, skip_partial_context=True |
| 626 | ) |
| 627 | actual = repeated_context_mask( |
| 628 | *(value.cuda() for value in inputs), |
| 629 | include_prompt=True, |
| 630 | skip_partial_context=True, |
| 631 | ).cpu() |
| 632 | |
| 633 | assert torch.equal(actual, expected) |
| 634 | |
| 635 | |
| 636 | def _unaligned_max_history_inputs() -> tuple[torch.Tensor, ...]: |
| 637 | """One row whose only repeated context ends at position 680. |
| 638 | |
| 639 | 680 is not a multiple of the kernel's 512-position block, so the scan window |
| 640 | selected by ``max_history`` starts in the middle of a block. |
| 641 | """ |
| 642 | history = list(range(1_000, 2_200)) |
| 643 | history[676:680] = [1, 2, 3, 4] |
| 644 | return ( |
| 645 | torch.tensor([history], dtype=torch.int32), |
| 646 | torch.tensor([0], dtype=torch.int32), |
| 647 | torch.tensor([0], dtype=torch.int32), |
| 648 | torch.tensor([len(history)], dtype=torch.int32), |
| 649 | torch.tensor([[1, 2, 3, 4]], dtype=torch.int32), |
| 650 | ) |
| 651 | |
| 652 | |
| 653 | @pytest.mark.skipif( |
| 654 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 655 | ) |
| 656 | @pytest.mark.parametrize( |
| 657 | "max_history,expected", [(1200 - 680, True), (1200 - 680 - 1, False)] |
| 658 | ) |
| 659 | def test_repeated_context_mask_unaligned_max_history_accelerator_parity( |
| 660 | max_history: int, expected: bool |
| 661 | ): |
| 662 | inputs = _unaligned_max_history_inputs() |
| 663 | |
| 664 | reference = repeated_context_mask(*inputs, max_history=max_history) |
| 665 | actual = repeated_context_mask( |
| 666 | *(value.cuda() for value in inputs), max_history=max_history |
| 667 | ).cpu() |
| 668 | |
| 669 | assert torch.equal(reference, torch.tensor([expected])) |
| 670 | assert torch.equal(actual, reference) |
| 671 | |
| 672 | |
| 673 | def _repeated_context_inputs( |
| 674 | context_width: int, device: str = "cpu" |
| 675 | ) -> tuple[torch.Tensor, ...]: |
| 676 | """Build repeated, unique, and padded request rows for mask tests.""" |
| 677 | prompt = [101, 102, 103, 104] |
| 678 | repeated_output = [*range(1, context_width + 1)] * 2 |
| 679 | unique_output = [*range(1, context_width + 2)] |
| 680 | max_len = len(prompt) + len(repeated_output) |
| 681 | rows = [ |
| 682 | prompt + repeated_output, |
| 683 | prompt + unique_output, |
| 684 | [], |
| 685 | ] |
| 686 | all_token_ids = torch.zeros((3, max_len), dtype=torch.int32, device=device) |
| 687 | for row, token_ids in enumerate(rows): |
| 688 | all_token_ids[row, : len(token_ids)] = torch.tensor( |
| 689 | token_ids, dtype=torch.int32, device=device |
| 690 | ) |
| 691 | return ( |
| 692 | all_token_ids, |
| 693 | torch.tensor([0, 1, -1], dtype=torch.int32, device=device), |
| 694 | torch.tensor([len(prompt), len(prompt), 0], dtype=torch.int32, device=device), |
| 695 | torch.tensor([len(rows[0]), len(rows[1]), 0], dtype=torch.int32, device=device), |
| 696 | torch.tensor( |
| 697 | [ |
| 698 | repeated_output[-context_width:], |
| 699 | unique_output[-context_width:], |
| 700 | [-1] * context_width, |
| 701 | ], |
| 702 | dtype=torch.int64, |
| 703 | device=device, |
| 704 | ), |
| 705 | ) |
| 706 | |
| 707 | |
| 708 | @pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17]) |
| 709 | def test_repeated_context_mask(context_width: int): |
| 710 | repeated = repeated_context_mask(*_repeated_context_inputs(context_width)) |
| 711 | |
| 712 | assert torch.equal(repeated, torch.tensor([True, False, False])) |
| 713 | |
| 714 | |
| 715 | @pytest.mark.skipif( |
| 716 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 717 | ) |
| 718 | @pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17]) |
| 719 | def test_repeated_context_mask_accelerator_parity(context_width: int): |
| 720 | cpu_inputs = _repeated_context_inputs(context_width) |
| 721 | accelerator_inputs = list(_repeated_context_inputs(context_width, "cuda")) |
| 722 | contexts = accelerator_inputs[-1] |
| 723 | storage = torch.empty( |
| 724 | contexts.shape[0], contexts.shape[1] * 2, dtype=contexts.dtype, device="cuda" |
| 725 | ) |
| 726 | storage[:, ::2] = contexts |
| 727 | accelerator_inputs[-1] = storage[:, ::2] |
| 728 | |
| 729 | expected = repeated_context_mask(*cpu_inputs) |
| 730 | actual = repeated_context_mask(*accelerator_inputs).cpu() |
| 731 | |
| 732 | assert torch.equal(actual, expected) |
| 733 | |
| 734 | |
| 735 | @pytest.mark.skipif( |
| 736 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 737 | ) |
| 738 | def test_repeated_context_mask_scans_multiple_blocks(): |
| 739 | context = [1, 2, 3, 4] |
| 740 | repeated_output = [*range(10_000, 11_030), *context, 99, *context] |
| 741 | unique_output = list(range(20_000, 21_039)) |
| 742 | prompt = [101, 102, 103, 104] |
| 743 | rows = [prompt + repeated_output, prompt + unique_output] |
| 744 | all_token_ids = torch.zeros((2, len(rows[0])), dtype=torch.int32) |
| 745 | for row, token_ids in enumerate(rows): |
| 746 | all_token_ids[row, : len(token_ids)] = torch.tensor(token_ids) |
| 747 | inputs = ( |
| 748 | all_token_ids, |
| 749 | torch.tensor([0, 1], dtype=torch.int32), |
| 750 | torch.tensor([len(prompt), len(prompt)], dtype=torch.int32), |
| 751 | torch.tensor([len(rows[0]), len(rows[1])], dtype=torch.int32), |
| 752 | torch.tensor([context, unique_output[-4:]], dtype=torch.int64), |
| 753 | ) |
| 754 | |
| 755 | expected = repeated_context_mask(*inputs) |
| 756 | actual = repeated_context_mask(*(value.cuda() for value in inputs)).cpu() |
| 757 | |
| 758 | assert torch.equal(expected, torch.tensor([True, False])) |
| 759 | assert torch.equal(actual, expected) |
| 760 | |
| 761 | |
| 762 | @pytest.mark.skipif( |
| 763 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 764 | ) |
| 765 | def test_repeated_context_mask_last_request_at_capacity(): |
| 766 | max_model_len = 1024 |
| 767 | all_token_ids = torch.arange( |
| 768 | 2 * max_model_len, dtype=torch.int32, device="cuda" |
| 769 | ).reshape(2, max_model_len) |
| 770 | inputs = ( |
| 771 | all_token_ids, |
| 772 | torch.tensor([1], dtype=torch.int32, device="cuda"), |
| 773 | torch.tensor([0, 0], dtype=torch.int32, device="cuda"), |
| 774 | torch.tensor([max_model_len, max_model_len], dtype=torch.int32, device="cuda"), |
| 775 | all_token_ids[1, -4:].to(torch.int64).unsqueeze(0), |
| 776 | ) |
| 777 | |
| 778 | repeated = repeated_context_mask(*inputs) |
| 779 | |
| 780 | assert not repeated.item() |
| 781 | |
| 782 | |
| 783 | @pytest.mark.skipif( |
| 784 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 785 | ) |
| 786 | def test_gpu_sampler_uses_fused_gumbel_for_repeated_contexts(): |
| 787 | device = torch.device("cuda") |
| 788 | logits = torch.tensor([[0.0, 1.0, 2.0, 3.0], [3.0, 2.0, 1.0, 0.0]], device=device) |
| 789 | all_token_ids = torch.tensor( |
| 790 | [[101, 102, 1, 2, 1, 2], [101, 102, 3, 4, 5, 6]], |
| 791 | dtype=torch.int32, |
| 792 | device=device, |
| 793 | ) |
| 794 | request_indices = torch.tensor([0, 1], dtype=torch.int64, device=device) |
| 795 | request_indices_np = np.array([0, 1]) |
| 796 | temperatures = torch.ones(2, device=device) |
| 797 | seeds = torch.tensor([11, 22], dtype=torch.int64, device=device) |
| 798 | positions = torch.tensor([6, 6], dtype=torch.int64, device=device) |
| 799 | |
| 800 | sampler = object.__new__(GPUWatermarkSampler) |
| 801 | sampler.watermarker = GumbelWatermarker(key=42, context_width=1) |
| 802 | sampler.deduplicate_contexts = "single_turn" |
| 803 | sampler.deduplicate_contexts_max_history = None |
| 804 | sampler.watermarking = SimpleNamespace( |
| 805 | np=np.ones(2, dtype=bool), gpu=torch.ones(2, dtype=torch.bool, device=device) |
| 806 | ) |
| 807 | sampler.sampling_states = SimpleNamespace( |
| 808 | temperature=SimpleNamespace(np=np.ones(2), gpu=temperatures), |
| 809 | seeds=SimpleNamespace(gpu=seeds), |
| 810 | ) |
| 811 | sampler.req_states = SimpleNamespace( |
| 812 | all_token_ids=SimpleNamespace(gpu=all_token_ids), |
| 813 | prompt_len=SimpleNamespace( |
| 814 | gpu=torch.tensor([2, 2], dtype=torch.int32, device=device) |
| 815 | ), |
| 816 | total_len=SimpleNamespace( |
| 817 | gpu=torch.tensor([6, 6], dtype=torch.int32, device=device) |
| 818 | ), |
| 819 | ) |
| 820 | sampler.use_fp64_gumbel = False |
| 821 | |
| 822 | contexts = sampler._get_contexts(request_indices) |
| 823 | repeated = sampler._get_repeated_contexts(request_indices, contexts) |
| 824 | expected = philox_gumbel_sample( |
| 825 | logits, |
| 826 | contexts, |
| 827 | 42, |
| 828 | skip_mask=repeated, |
| 829 | expanded_idx_mapping=request_indices, |
| 830 | temperatures=temperatures, |
| 831 | seeds=seeds, |
| 832 | positions=positions, |
| 833 | ) |
| 834 | |
| 835 | actual, output_logits = sampler._sample_random( |
| 836 | logits, |
| 837 | request_indices, |
| 838 | request_indices_np, |
| 839 | positions, |
| 840 | None, |
| 841 | None, |
| 842 | False, |
| 843 | ) |
| 844 | |
| 845 | assert torch.equal(repeated, torch.tensor([True, False], device=device)) |
| 846 | assert torch.equal(actual, expected) |
| 847 | assert output_logits is logits |
| 848 | |
| 849 | |
| 850 | def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch): |
| 851 | class StubWatermarker: |
| 852 | context_width = 1 |
| 853 | |
| 854 | def sample(self, logits, contexts, random_sampler=None, skip_mask=None): |
| 855 | raise AssertionError("watermarker should not run for greedy requests") |
| 856 | |
| 857 | sampler = object.__new__(GPUWatermarkSampler) |
| 858 | sampler.watermarker = StubWatermarker() |
| 859 | sampler.watermarking = SimpleNamespace( |
| 860 | np=np.array([True, True]), gpu=torch.tensor([True, True]) |
| 861 | ) |
| 862 | sampler.sampling_states = SimpleNamespace( |
| 863 | temperature=SimpleNamespace(np=np.zeros(2), gpu=torch.zeros(2)), |
| 864 | ) |
| 865 | expected = (torch.tensor([3, 4]), torch.zeros(2, 8)) |
| 866 | monkeypatch.setattr( |
| 867 | "vllm.v1.watermarking.gpu_sampler.Sampler._sample_random", |
| 868 | lambda *args, **kwargs: expected, |
| 869 | ) |
| 870 | |
| 871 | actual = sampler._sample_random( |
| 872 | torch.zeros(2, 8), |
| 873 | torch.tensor([0, 1]), |
| 874 | np.array([0, 1]), |
| 875 | torch.zeros(2, dtype=torch.int64), |
| 876 | None, |
| 877 | None, |
| 878 | False, |
| 879 | ) |
| 880 | |
| 881 | assert actual is expected |
| 882 | |
| 883 | |
| 884 | def test_gpu_sampler_builds_speculative_contexts_from_drafts(): |
| 885 | sampler = object.__new__(GPUWatermarkSampler) |
| 886 | sampler.watermarker = SimpleNamespace(context_width=2) |
| 887 | sampler.req_states = SimpleNamespace( |
| 888 | all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 101, 10, 11, 0, 0]])), |
| 889 | prompt_len=SimpleNamespace(gpu=torch.tensor([2])), |
| 890 | total_len=SimpleNamespace(gpu=torch.tensor([4])), |
| 891 | ) |
| 892 | |
| 893 | contexts = sampler._get_contexts( |
| 894 | torch.tensor([0, 0, 0]), |
| 895 | torch.tensor([0, 1, 2]), |
| 896 | torch.tensor([11, 20, 21]), |
| 897 | ) |
| 898 | |
| 899 | assert torch.equal(contexts, torch.tensor([[10, 11], [11, 20], [20, 21]])) |
| 900 | |
| 901 | |
| 902 | def test_gpu_sampler_deduplicates_contexts_within_speculative_block(): |
| 903 | sampler = object.__new__(GPUWatermarkSampler) |
| 904 | sampler.watermarker = SimpleNamespace(context_width=2) |
| 905 | sampler.deduplicate_contexts = "single_turn" |
| 906 | sampler.deduplicate_contexts_max_history = None |
| 907 | sampler.num_speculative_tokens = 4 |
| 908 | sampler.req_states = SimpleNamespace( |
| 909 | all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 1, 2, 0, 0, 0]])), |
| 910 | prompt_len=SimpleNamespace(gpu=torch.tensor([1])), |
| 911 | total_len=SimpleNamespace(gpu=torch.tensor([3])), |
| 912 | ) |
| 913 | request_indices = torch.tensor([0, 0, 0, 0, 0]) |
| 914 | local_positions = torch.tensor([0, 1, 2, 3, 4]) |
| 915 | contexts = sampler._get_contexts( |
| 916 | request_indices, |
| 917 | local_positions, |
| 918 | torch.tensor([0, 3, 4, 1, 2]), |
| 919 | ) |
| 920 | |
| 921 | repeated = sampler._get_repeated_contexts( |
| 922 | request_indices, contexts, local_positions |
| 923 | ) |
| 924 | |
| 925 | assert torch.equal( |
| 926 | contexts, |
| 927 | torch.tensor([[1, 2], [2, 3], [3, 4], [4, 1], [1, 2]]), |
| 928 | ) |
| 929 | assert torch.equal(repeated, torch.tensor([False, False, False, False, True])) |
| 930 | |
| 931 | sampler.deduplicate_contexts_max_history = 3 |
| 932 | repeated = sampler._get_repeated_contexts( |
| 933 | request_indices, contexts, local_positions |
| 934 | ) |
| 935 | assert not repeated.any() |
| 936 | |
| 937 | |
| 938 | def test_gpu_sampler_all_scope_skips_only_partial_speculative_contexts(): |
| 939 | sampler = object.__new__(GPUWatermarkSampler) |
| 940 | sampler.watermarker = SimpleNamespace(context_width=2) |
| 941 | sampler.deduplicate_contexts = "all" |
| 942 | sampler.deduplicate_contexts_max_history = None |
| 943 | sampler.num_speculative_tokens = 3 |
| 944 | sampler.req_states = SimpleNamespace( |
| 945 | all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 101, 10, 0, 0]])), |
| 946 | prompt_len=SimpleNamespace(gpu=torch.tensor([2])), |
| 947 | total_len=SimpleNamespace(gpu=torch.tensor([3])), |
| 948 | ) |
| 949 | request_indices = torch.tensor([0, 0, 0]) |
| 950 | local_positions = torch.tensor([0, 1, 2]) |
| 951 | contexts = sampler._get_contexts( |
| 952 | request_indices, |
| 953 | local_positions, |
| 954 | torch.tensor([10, 20, 30]), |
| 955 | ) |
| 956 | |
| 957 | repeated = sampler._get_repeated_contexts( |
| 958 | request_indices, contexts, local_positions |
| 959 | ) |
| 960 | |
| 961 | assert torch.equal(contexts, torch.tensor([[-1, 10], [10, 20], [20, 30]])) |
| 962 | assert torch.equal(repeated, torch.tensor([True, False, False])) |
| 963 | |
| 964 | |
| 965 | @pytest.mark.skipif( |
| 966 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 967 | ) |
| 968 | @pytest.mark.parametrize("max_history,expected_last", [(None, True), (3, False)]) |
| 969 | def test_speculative_context_repetition_accelerator_parity( |
| 970 | max_history: int | None, expected_last: bool |
| 971 | ): |
| 972 | all_token_ids = torch.tensor([[10, 11]]) |
| 973 | req_indices = torch.tensor([0, 0, 0, 0, 0]) |
| 974 | prompt_lens = torch.tensor([0]) |
| 975 | total_lens = torch.tensor([2]) |
| 976 | local_positions = torch.tensor([0, 1, 2, 3, 4]) |
| 977 | contexts = torch.tensor([[1, 2], [2, 3], [3, 4], [4, 1], [1, 2]]) |
| 978 | |
| 979 | expected = repeated_context_mask( |
| 980 | all_token_ids, |
| 981 | req_indices, |
| 982 | prompt_lens, |
| 983 | total_lens, |
| 984 | contexts, |
| 985 | history_offsets=local_positions, |
| 986 | local_positions=local_positions, |
| 987 | num_speculative_steps=4, |
| 988 | max_history=max_history, |
| 989 | ) |
| 990 | actual = repeated_context_mask( |
| 991 | all_token_ids.cuda(), |
| 992 | req_indices.cuda(), |
| 993 | prompt_lens.cuda(), |
| 994 | total_lens.cuda(), |
| 995 | contexts.cuda(), |
| 996 | history_offsets=local_positions.cuda(), |
| 997 | local_positions=local_positions.cuda(), |
| 998 | num_speculative_steps=4, |
| 999 | max_history=max_history, |
| 1000 | ).cpu() |
| 1001 | |
| 1002 | assert torch.equal( |
| 1003 | expected, torch.tensor([False, False, False, False, expected_last]) |
| 1004 | ) |
| 1005 | assert torch.equal(actual, expected) |
| 1006 | |
| 1007 | |
| 1008 | @pytest.mark.skipif( |
| 1009 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 1010 | ) |
| 1011 | @pytest.mark.parametrize("steps_layout", ["per_row", "broadcast"]) |
| 1012 | @pytest.mark.parametrize("req_indices_layout", ["contiguous", "column_view"]) |
| 1013 | def test_draft_context_deduplication_accelerator_parity( |
| 1014 | steps_layout, req_indices_layout |
| 1015 | ): |
| 1016 | """Parity with the CPU reference for the layouts the speculators pass. |
| 1017 | |
| 1018 | A draft step is one 0-d tensor broadcast over the batch, and DSpark samples |
| 1019 | a step at a time, so its request indices are an `idx_map[:, i]` column. Both |
| 1020 | are read per row, so a row that ignores the layout consults another row's |
| 1021 | step or another request's history. |
| 1022 | """ |
| 1023 | num_rows, num_steps, step = 4, 3, 1 |
| 1024 | all_token_ids = torch.tensor([[10, 11, 20, 11], [30, 31, 11, 20]]) |
| 1025 | prompt_lens = torch.tensor([0, 0]) |
| 1026 | total_lens = torch.tensor([4, 4]) |
| 1027 | contexts = torch.tensor([[40, 41], [31, 11], [10, 11], [50, 51]]) |
| 1028 | prior_contexts = torch.zeros(num_rows, num_steps, 2, dtype=torch.int64) |
| 1029 | prior_contexts[0, 0] = torch.tensor([40, 41]) |
| 1030 | enabled = torch.ones(num_rows, dtype=torch.bool) |
| 1031 | req_indices = torch.tensor([0, 1, 0, 1]) |
| 1032 | |
| 1033 | def run(all_token_ids, req_indices, steps, prior): |
| 1034 | return draft_watermarking_mask( |
| 1035 | all_token_ids, |
| 1036 | req_indices, |
| 1037 | prompt_lens.to(all_token_ids.device), |
| 1038 | total_lens.to(all_token_ids.device), |
| 1039 | prior, |
| 1040 | contexts.to(all_token_ids.device), |
| 1041 | steps, |
| 1042 | enabled.to(all_token_ids.device), |
| 1043 | max_history=3, |
| 1044 | include_prompt=False, |
| 1045 | ) |
| 1046 | |
| 1047 | reference_prior = prior_contexts.clone() |
| 1048 | reference = run( |
| 1049 | all_token_ids, |
| 1050 | req_indices, |
| 1051 | torch.full((num_rows,), step), |
| 1052 | reference_prior, |
| 1053 | ) |
| 1054 | |
| 1055 | # Row 0 repeats its own step-0 context, rows 1 and 2 repeat a committed |
| 1056 | # context of their own request, row 3 repeats nothing. |
| 1057 | assert torch.equal(reference, torch.tensor([False, False, False, True])) |
| 1058 | expected_prior = prior_contexts.clone() |
| 1059 | expected_prior[:, step] = contexts |
| 1060 | assert torch.equal(reference_prior, expected_prior) |
| 1061 | |
| 1062 | if steps_layout == "per_row": |
| 1063 | steps = torch.full((num_rows,), step).cuda() |
| 1064 | else: |
| 1065 | steps = torch.tensor(step).cuda().expand(num_rows) |
| 1066 | assert steps.stride(0) == 0 |
| 1067 | |
| 1068 | if req_indices_layout == "contiguous": |
| 1069 | gpu_req_indices = req_indices.cuda() |
| 1070 | else: |
| 1071 | gpu_req_indices = ( |
| 1072 | req_indices.repeat_interleave(num_steps).view(num_rows, num_steps).cuda() |
| 1073 | )[:, 1] |
| 1074 | assert gpu_req_indices.stride(0) == num_steps |
| 1075 | |
| 1076 | actual_prior = prior_contexts.cuda() |
| 1077 | actual = run(all_token_ids.cuda(), gpu_req_indices, steps, actual_prior).cpu() |
| 1078 | |
| 1079 | assert torch.equal(actual, reference) |
| 1080 | assert torch.equal(actual_prior.cpu(), expected_prior) |
| 1081 | |
| 1082 | |
| 1083 | def test_gpu_sampler_builds_chunked_multi_request_speculative_contexts(): |
| 1084 | sampler = object.__new__(GPUWatermarkSampler) |
| 1085 | sampler.watermarker = SimpleNamespace(context_width=2) |
| 1086 | sampler.req_states = SimpleNamespace( |
| 1087 | all_token_ids=SimpleNamespace( |
| 1088 | gpu=torch.tensor( |
| 1089 | [ |
| 1090 | [100, 101, 10, 11, 0, 0], |
| 1091 | [200, 201, 30, 31, 0, 0], |
| 1092 | [300, 301, 50, 51, 0, 0], |
| 1093 | ] |
| 1094 | ) |
| 1095 | ), |
| 1096 | prompt_len=SimpleNamespace(gpu=torch.tensor([2, 2, 2])), |
| 1097 | total_len=SimpleNamespace(gpu=torch.tensor([4, 4, 4])), |
| 1098 | ) |
| 1099 | |
| 1100 | contexts = sampler._get_contexts( |
| 1101 | torch.tensor([2, 2, 1, 1]), |
| 1102 | torch.tensor([0, 1, 0, 1]), |
| 1103 | torch.tensor([51, 60, 31, 40]), |
| 1104 | ) |
| 1105 | |
| 1106 | assert torch.equal( |
| 1107 | contexts, |
| 1108 | torch.tensor([[50, 51], [51, 60], [30, 31], [31, 40]]), |
| 1109 | ) |
| 1110 | |
| 1111 | |
| 1112 | @pytest.mark.skipif( |
| 1113 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 1114 | ) |
| 1115 | def test_draft_context_deduplication_broadcasts_the_draft_step(): |
| 1116 | """The draft step is one 0-d tensor broadcast over the batch. |
| 1117 | |
| 1118 | The context occurs four positions back, which `max_history - step` puts |
| 1119 | just outside the scanned window, so a row that mistakes another row's step |
| 1120 | for its own reports the opposite verdict. The neighbours leg is the |
| 1121 | deterministic control: a row that reads a step of 0 instead of 2 widens its |
| 1122 | window by two positions and finds the context. The production leg asserts |
| 1123 | the layout the speculators pass, where the stride-less read lands outside |
| 1124 | the allocation and can return the right value by chance. |
| 1125 | """ |
| 1126 | num_rows = 4 |
| 1127 | all_token_ids = torch.tensor([[3, 4, 9, 9, 9, 9], [3, 4, 9, 9, 9, 9]]) |
| 1128 | req_indices = torch.tensor([0, 0, 1, 1]) |
| 1129 | prompt_lens = torch.tensor([0, 0]) |
| 1130 | total_lens = torch.tensor([6, 6]) |
| 1131 | prior_contexts = torch.zeros(num_rows, 3, 2, dtype=torch.int64) |
| 1132 | contexts = torch.tensor([[3, 4]]).expand(num_rows, 2).contiguous() |
| 1133 | enabled = torch.ones(num_rows, dtype=torch.bool) |
| 1134 | |
| 1135 | def run(steps, device): |
| 1136 | prior = prior_contexts.clone().to(device) |
| 1137 | active = draft_watermarking_mask( |
| 1138 | all_token_ids.to(device), |
| 1139 | req_indices.to(device), |
| 1140 | prompt_lens.to(device), |
| 1141 | total_lens.to(device), |
| 1142 | prior, |
| 1143 | contexts.to(device), |
| 1144 | steps, |
| 1145 | enabled.to(device), |
| 1146 | max_history=4, |
| 1147 | include_prompt=False, |
| 1148 | ) |
| 1149 | return active.cpu(), prior.cpu() |
| 1150 | |
| 1151 | expected, expected_prior = run(torch.full((num_rows,), 2), torch.device("cpu")) |
| 1152 | assert torch.equal(expected, torch.ones(num_rows, dtype=torch.bool)) |
| 1153 | |
| 1154 | materialized, materialized_prior = run( |
| 1155 | torch.full((num_rows,), 2).cuda(), torch.device("cuda") |
| 1156 | ) |
| 1157 | assert torch.equal(materialized, expected) |
| 1158 | assert torch.equal(materialized_prior, expected_prior) |
| 1159 | |
| 1160 | # Neighbouring steps are the value a row reads when the broadcast stride is |
| 1161 | # ignored; the production tensor reads uninitialized memory instead. |
| 1162 | neighbours = torch.tensor([2, 0, 0, 0]).cuda() |
| 1163 | broadcast, broadcast_prior = run( |
| 1164 | neighbours[0].expand(num_rows), torch.device("cuda") |
| 1165 | ) |
| 1166 | assert torch.equal(broadcast, expected) |
| 1167 | assert torch.equal(broadcast_prior, expected_prior) |
| 1168 | |
| 1169 | step = torch.tensor(2).cuda() |
| 1170 | production, production_prior = run(step.expand(num_rows), torch.device("cuda")) |
| 1171 | assert torch.equal(production, expected) |
| 1172 | assert torch.equal(production_prior, expected_prior) |
| 1173 | |
| 1174 | |
| 1175 | def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch): |
| 1176 | class StubSpeculator(DraftModelSpeculator): |
| 1177 | def capture(self): ... |
| 1178 | |
| 1179 | def init_cudagraph_manager(self, cudagraph_mode): ... |
| 1180 | |
| 1181 | def load_draft_model(self, target_model, target_attn_layer_names): ... |
| 1182 | |
| 1183 | def propose(self, *args, **kwargs): ... |
| 1184 | |
| 1185 | class StubModel: |
| 1186 | @staticmethod |
| 1187 | def compute_logits(hidden_states): |
| 1188 | return torch.zeros(hidden_states.shape[0], 8) |
| 1189 | |
| 1190 | class StubWatermarker: |
| 1191 | context_width = 2 |
| 1192 | |
| 1193 | @staticmethod |
| 1194 | def sample(logits, contexts, random_sampler=None, skip_mask=None): |
| 1195 | assert random_sampler is not None |
| 1196 | return WatermarkSample( |
| 1197 | torch.where(skip_mask, random_sampler(logits), 7), logits |
| 1198 | ) |
| 1199 | |
| 1200 | speculator = object.__new__(StubSpeculator) |
| 1201 | speculator.model = StubModel() |
| 1202 | speculator.use_fp64_gumbel = False |
| 1203 | speculator.acceptance_estimator = None |
| 1204 | draft_watermarker = DraftWatermarker( |
| 1205 | StubWatermarker(), |
| 1206 | max_num_reqs=2, |
| 1207 | device=torch.device("cpu"), |
| 1208 | num_speculative_steps=1, |
| 1209 | deduplicate_contexts="none", |
| 1210 | deduplicate_contexts_max_history=None, |
| 1211 | ) |
| 1212 | draft_watermarker.contexts.copy_(torch.tensor([[1, 2], [3, 4]])) |
| 1213 | draft_watermarker.enabled = torch.tensor([True, False]) |
| 1214 | speculator.draft_watermarker = draft_watermarker |
| 1215 | monkeypatch.setattr( |
| 1216 | "vllm.v1.watermarking.watermarker.gumbel_sample", |
| 1217 | lambda *args, **kwargs: torch.tensor([3, 4]), |
| 1218 | ) |
| 1219 | |
| 1220 | sampled = speculator.sample_draft( |
| 1221 | hidden_states=torch.zeros(2, 4), |
| 1222 | sample_src_positions=torch.zeros(2, dtype=torch.int64), |
| 1223 | idx_mapping=torch.tensor([0, 1]), |
| 1224 | temperature=torch.ones(2), |
| 1225 | seeds=torch.zeros(2, dtype=torch.int64), |
| 1226 | draft_step=torch.tensor(0), |
| 1227 | draft_logits=torch.zeros(2, 1, 8), |
| 1228 | ) |
| 1229 | |
| 1230 | assert torch.equal(sampled, torch.tensor([7, 4])) |
| 1231 | assert torch.equal(draft_watermarker.contexts, torch.tensor([[2, 7], [4, 4]])) |
| 1232 | |
| 1233 | |
| 1234 | def test_draft_sampler_uses_ordinary_samples_for_repeated_contexts(monkeypatch): |
| 1235 | class StubWatermarker: |
| 1236 | context_width = 1 |
| 1237 | |
| 1238 | def __init__(self): |
| 1239 | self.samples = iter((2, 1, 8)) |
| 1240 | |
| 1241 | def sample(self, logits, contexts, random_sampler=None, skip_mask=None): |
| 1242 | watermarked = torch.tensor([next(self.samples)]) |
| 1243 | return WatermarkSample( |
| 1244 | torch.where(skip_mask, random_sampler(logits), watermarked), logits |
| 1245 | ) |
| 1246 | |
| 1247 | draft_watermarker = DraftWatermarker( |
| 1248 | StubWatermarker(), |
| 1249 | max_num_reqs=1, |
| 1250 | device=torch.device("cpu"), |
| 1251 | num_speculative_steps=3, |
| 1252 | deduplicate_contexts="single_turn", |
| 1253 | deduplicate_contexts_max_history=None, |
| 1254 | ) |
| 1255 | draft_watermarker.prepare( |
| 1256 | contexts=torch.tensor([[1]]), |
| 1257 | enabled=torch.tensor([True]), |
| 1258 | all_token_ids=torch.tensor([[100, 1, 0, 0]]), |
| 1259 | prompt_lens=torch.tensor([1]), |
| 1260 | total_lens=torch.tensor([2]), |
| 1261 | ) |
| 1262 | idx_mapping = torch.tensor([0]) |
| 1263 | temperature = torch.tensor([1.0]) |
| 1264 | ordinary_samples = iter((torch.tensor([5]), torch.tensor([6]), torch.tensor([4]))) |
| 1265 | monkeypatch.setattr( |
| 1266 | "vllm.v1.watermarking.watermarker.gumbel_sample", |
| 1267 | lambda *args, **kwargs: next(ordinary_samples), |
| 1268 | ) |
| 1269 | seeds = torch.tensor([0]) |
| 1270 | draft_logits = torch.zeros(1, 3, 8) |
| 1271 | |
| 1272 | first = draft_watermarker.sample( |
| 1273 | torch.zeros(1, 8), |
| 1274 | idx_mapping, |
| 1275 | temperature=temperature, |
| 1276 | seed=seeds, |
| 1277 | pos=torch.tensor([0]), |
| 1278 | apply_temperature=True, |
| 1279 | is_drafting=True, |
| 1280 | logits_cache_col=torch.tensor(0), |
| 1281 | logits_cache=draft_logits, |
| 1282 | use_fp64=False, |
| 1283 | ) |
| 1284 | second = draft_watermarker.sample( |
| 1285 | torch.zeros(1, 8), |
| 1286 | idx_mapping, |
| 1287 | temperature=temperature, |
| 1288 | seed=seeds, |
| 1289 | pos=torch.tensor([1]), |
| 1290 | apply_temperature=True, |
| 1291 | is_drafting=True, |
| 1292 | logits_cache_col=torch.tensor(1), |
| 1293 | logits_cache=draft_logits, |
| 1294 | use_fp64=False, |
| 1295 | ) |
| 1296 | repeated = draft_watermarker.sample( |
| 1297 | torch.zeros(1, 8), |
| 1298 | idx_mapping, |
| 1299 | temperature=temperature, |
| 1300 | seed=seeds, |
| 1301 | pos=torch.tensor([2]), |
| 1302 | apply_temperature=True, |
| 1303 | is_drafting=True, |
| 1304 | logits_cache_col=torch.tensor(2), |
| 1305 | logits_cache=draft_logits, |
| 1306 | use_fp64=False, |
| 1307 | ) |
| 1308 | |
| 1309 | assert first.item() == 2 |
| 1310 | assert second.item() == 1 |
| 1311 | assert repeated.item() == 4 |
| 1312 | |
| 1313 | |
| 1314 | @pytest.mark.skipif( |
| 1315 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 1316 | ) |
| 1317 | def test_draft_context_deduplication_reads_strided_request_indices(): |
| 1318 | """DSpark samples a step at a time, so req_indices is an `idx_map[:, i]` column.""" |
| 1319 | num_reqs = 4 |
| 1320 | num_speculative_steps = 3 |
| 1321 | all_token_ids = torch.tensor( |
| 1322 | [[3, 4, 3, 4], [9, 9, 9, 9], [3, 4, 3, 4], [9, 9, 9, 9]] |
| 1323 | ).cuda() |
| 1324 | prompt_lens = torch.zeros(num_reqs, dtype=torch.int64).cuda() |
| 1325 | total_lens = torch.full((num_reqs,), 4).cuda() |
| 1326 | contexts = torch.tensor([[3, 4]]).expand(num_reqs, 2).contiguous().cuda() |
| 1327 | steps = torch.zeros(num_reqs, dtype=torch.int64).cuda() |
| 1328 | enabled = torch.ones(num_reqs, dtype=torch.bool).cuda() |
| 1329 | |
| 1330 | idx_map = ( |
| 1331 | torch.arange(num_reqs) |
| 1332 | .repeat_interleave(num_speculative_steps) |
| 1333 | .view(num_reqs, num_speculative_steps) |
| 1334 | .cuda() |
| 1335 | ) |
| 1336 | column = idx_map[:, 1] |
| 1337 | assert column.stride(0) == num_speculative_steps |
| 1338 | |
| 1339 | def run(req_indices): |
| 1340 | prior = torch.zeros( |
| 1341 | num_reqs, num_speculative_steps, 2, dtype=torch.int64, device="cuda" |
| 1342 | ) |
| 1343 | return draft_watermarking_mask( |
| 1344 | all_token_ids, |
| 1345 | req_indices, |
| 1346 | prompt_lens, |
| 1347 | total_lens, |
| 1348 | prior, |
| 1349 | contexts, |
| 1350 | steps, |
| 1351 | enabled, |
| 1352 | max_history=None, |
| 1353 | include_prompt=False, |
| 1354 | ).cpu() |
| 1355 | |
| 1356 | expected = torch.tensor([False, True, False, True]) |
| 1357 | assert torch.equal(run(column.contiguous()), expected) |
| 1358 | assert torch.equal(run(column), expected) |
| 1359 | |
| 1360 | |
| 1361 | def test_draft_sampler_deduplicates_against_committed_history(monkeypatch): |
| 1362 | class StubWatermarker: |
| 1363 | context_width = 1 |
| 1364 | |
| 1365 | @staticmethod |
| 1366 | def sample(logits, contexts, random_sampler=None, skip_mask=None): |
| 1367 | return WatermarkSample( |
| 1368 | torch.where(skip_mask, random_sampler(logits), 7), logits |
| 1369 | ) |
| 1370 | |
| 1371 | draft_watermarker = DraftWatermarker( |
| 1372 | StubWatermarker(), |
| 1373 | max_num_reqs=1, |
| 1374 | device=torch.device("cpu"), |
| 1375 | num_speculative_steps=1, |
| 1376 | deduplicate_contexts="single_turn", |
| 1377 | deduplicate_contexts_max_history=None, |
| 1378 | ) |
| 1379 | draft_watermarker.prepare( |
| 1380 | contexts=torch.tensor([[1]]), |
| 1381 | enabled=torch.tensor([True]), |
| 1382 | all_token_ids=torch.tensor([[100, 1, 2, 1]]), |
| 1383 | prompt_lens=torch.tensor([1]), |
| 1384 | total_lens=torch.tensor([4]), |
| 1385 | ) |
| 1386 | |
| 1387 | monkeypatch.setattr( |
| 1388 | "vllm.v1.watermarking.watermarker.gumbel_sample", |
| 1389 | lambda *args, **kwargs: torch.tensor([4]), |
| 1390 | ) |
| 1391 | sampled = draft_watermarker.sample( |
| 1392 | torch.zeros(1, 8), |
| 1393 | torch.tensor([0]), |
| 1394 | temperature=torch.tensor([1.0]), |
| 1395 | seed=torch.tensor([0]), |
| 1396 | pos=torch.tensor([0]), |
| 1397 | apply_temperature=True, |
| 1398 | is_drafting=True, |
| 1399 | logits_cache_col=torch.tensor(0), |
| 1400 | logits_cache=torch.zeros(1, 1, 8), |
| 1401 | use_fp64=False, |
| 1402 | ) |
| 1403 | |
| 1404 | assert sampled.item() == 4 |
| 1405 | |
| 1406 | |
| 1407 | @pytest.mark.skipif( |
| 1408 | not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" |
| 1409 | ) |
| 1410 | @pytest.mark.parametrize("idx_mapping_layout", ["contiguous", "column_view"]) |
| 1411 | def test_draft_watermarker_deduplicates_across_a_block_on_cuda( |
| 1412 | monkeypatch, idx_mapping_layout |
| 1413 | ): |
| 1414 | """Drive the whole draft sampler over a block with the production shapes. |
| 1415 | |
| 1416 | `draft_step` is a 0-d tensor, and a speculator that samples one step at a |
| 1417 | time hands over a column of its index map. Each step is checked against the |
| 1418 | committed history the window still reaches and against the earlier steps of |
| 1419 | the block. |
| 1420 | """ |
| 1421 | num_reqs, num_steps, vocab_size = 2, 3, 16 |
| 1422 | watermarker = GumbelWatermarker(key=42, context_width=2) |
| 1423 | draft_watermarker = DraftWatermarker( |
| 1424 | watermarker, |
| 1425 | max_num_reqs=num_reqs, |
| 1426 | device=torch.device("cuda"), |
| 1427 | num_speculative_steps=num_steps, |
| 1428 | deduplicate_contexts="single_turn", |
| 1429 | deduplicate_contexts_max_history=4, |
| 1430 | ) |
| 1431 | draft_watermarker.prepare( |
| 1432 | contexts=torch.tensor([[5, 6], [6, 3]]).cuda(), |
| 1433 | enabled=torch.ones(num_reqs, dtype=torch.bool).cuda(), |
| 1434 | all_token_ids=torch.tensor([[5, 6, 7, 5, 6], [8, 9, 6, 6, 3]]).cuda(), |
| 1435 | prompt_lens=torch.zeros(num_reqs, dtype=torch.int64).cuda(), |
| 1436 | total_lens=torch.full((num_reqs,), 5).cuda(), |
| 1437 | ) |
| 1438 | |
| 1439 | if idx_mapping_layout == "contiguous": |
| 1440 | idx_mapping = torch.arange(num_reqs, dtype=torch.int32).cuda() |
| 1441 | else: |
| 1442 | idx_mapping = ( |
| 1443 | torch.arange(num_reqs, dtype=torch.int32) |
| 1444 | .repeat_interleave(num_steps) |
| 1445 | .view(num_reqs, num_steps) |
| 1446 | .cuda() |
| 1447 | )[:, 1] |
| 1448 | assert idx_mapping.stride(0) == num_steps |
| 1449 | |
| 1450 | skip_masks: list[torch.Tensor] = [] |
| 1451 | original = philox_gumbel_sample |
| 1452 | |
| 1453 | def spy(*args, **kwargs): |
| 1454 | skip_masks.append(kwargs["skip_mask"].cpu()) |
| 1455 | return original(*args, **kwargs) |
| 1456 | |
| 1457 | monkeypatch.setattr("vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample", spy) |
| 1458 | |
| 1459 | draft_logits = torch.zeros(num_reqs, num_steps, vocab_size + 1).cuda() |
| 1460 | step_logits = [] |
| 1461 | for step in range(num_steps): |
| 1462 | # Token 6 wins by a margin no Gumbel draw can close, on either path. |
| 1463 | logits = torch.zeros(num_reqs, vocab_size).cuda() |
| 1464 | logits[:, 6] = 100.0 |
| 1465 | logits[:, 0] = float(step) + 1.0 |
| 1466 | step_logits.append(logits) |
| 1467 | sampled = draft_watermarker.sample( |
| 1468 | logits, |
| 1469 | idx_mapping, |
| 1470 | temperature=torch.ones(num_reqs).cuda(), |
| 1471 | seed=torch.zeros(num_reqs, dtype=torch.int64).cuda(), |
| 1472 | pos=torch.full((num_reqs,), 5 + step, dtype=torch.int64).cuda(), |
| 1473 | apply_temperature=True, |
| 1474 | is_drafting=True, |
| 1475 | logits_cache_col=torch.tensor(step).cuda(), |
| 1476 | logits_cache=draft_logits, |
| 1477 | use_fp64=False, |
| 1478 | ) |
| 1479 | assert torch.equal(sampled.cpu(), torch.tensor([6, 6])) |
| 1480 | |
| 1481 | # Step 0: request 0's [5, 6] is committed at position 2, request 1's [6, 3] |
| 1482 | # is new. Step 1: [6, 6] and [3, 6] are both new. Step 2: request 0's |
| 1483 | # [6, 6] repeats its own step-1 context and request 1's repeats position 4. |
| 1484 | assert [mask.tolist() for mask in skip_masks] == [ |
| 1485 | [True, False], |
| 1486 | [False, False], |
| 1487 | [True, True], |
| 1488 | ] |
| 1489 | assert torch.equal( |
| 1490 | draft_watermarker.prior_contexts[:num_reqs].cpu(), |
| 1491 | torch.tensor([[[5, 6], [6, 6], [6, 6]], [[6, 3], [3, 6], [6, 6]]]), |
| 1492 | ) |
| 1493 | assert torch.equal( |
| 1494 | draft_watermarker.contexts[:num_reqs].cpu(), torch.tensor([[6, 6], [6, 6]]) |
| 1495 | ) |
| 1496 | for step, logits in enumerate(step_logits): |
| 1497 | # The 0-d draft step selects the logits cache column. |
| 1498 | assert torch.equal(draft_logits[:, step, :vocab_size], logits) |
| 1499 | assert not draft_logits[:, :, -1].any() |
| 1500 | |
| 1501 | |
| 1502 | def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch): |
| 1503 | speculator = object.__new__(DSparkSpeculator) |
| 1504 | speculator.draft_logits = torch.zeros(2, 1, 8) |
| 1505 | speculator._d2t_scatter_index = torch.tensor([1, 5]) |
| 1506 | speculator._draft_scatter_buf = torch.full((2, 8), float("-inf")) |
| 1507 | speculator.temperature = torch.ones(2) |
| 1508 | speculator.seeds = torch.zeros(2, dtype=torch.int64) |
| 1509 | speculator._step_cols = torch.tensor([0]) |
| 1510 | speculator.use_fp64_gumbel = False |
| 1511 | speculator.acceptance_estimator = None |
| 1512 | watermark_logits: list[torch.Tensor] = [] |
| 1513 | |
| 1514 | def sample( |
| 1515 | logits, |
| 1516 | idx_mapping, |
| 1517 | temperature, |
| 1518 | seed, |
| 1519 | pos, |
| 1520 | apply_temperature, |
| 1521 | is_drafting, |
| 1522 | logits_cache=None, |
| 1523 | logits_cache_col=None, |
| 1524 | use_fp64=False, |
| 1525 | logits_cache_source=None, |
| 1526 | ): |
| 1527 | watermark_logits.append(logits.clone()) |
| 1528 | return torch.tensor([4, 5]) |
| 1529 | |
| 1530 | speculator.draft_watermarker = SimpleNamespace(sample=sample) |
| 1531 | |
| 1532 | sampled = speculator._sample_logits( |
| 1533 | torch.tensor([[10.0, 20.0], [30.0, 40.0]]), |
| 1534 | torch.tensor([0, 1]), |
| 1535 | torch.tensor([1, 1]), |
| 1536 | 0, |
| 1537 | ) |
| 1538 | |
| 1539 | assert torch.equal(sampled, torch.tensor([4, 5])) |
| 1540 | assert torch.equal( |
| 1541 | watermark_logits[0][:, [1, 5]], torch.tensor([[10, 20], [30, 40]]) |
| 1542 | ) |
| 1543 | assert torch.isneginf(watermark_logits[0][:, [0, 2, 3, 4, 6, 7]]).all() |
| 1544 | |
| 1545 | |
| 1546 | def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch): |
| 1547 | speculator = object.__new__(DSparkSpeculator) |
| 1548 | speculator.draft_logits = torch.zeros(2, 1, 8) |
| 1549 | speculator._d2t_scatter_index = None |
| 1550 | speculator.temperature = torch.ones(2) |
| 1551 | speculator.seeds = torch.zeros(2, dtype=torch.int64) |
| 1552 | speculator._step_cols = torch.tensor([0]) |
| 1553 | speculator.use_fp64_gumbel = False |
| 1554 | speculator.acceptance_estimator = None |
| 1555 | speculator.draft_watermarker = None |
| 1556 | monkeypatch.setattr( |
| 1557 | "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample", |
| 1558 | lambda *args, **kwargs: torch.tensor([3, 4]), |
| 1559 | ) |
| 1560 | |
| 1561 | sampled = speculator._sample_logits( |
| 1562 | torch.zeros(2, 8), |
| 1563 | torch.tensor([0, 1]), |
| 1564 | torch.tensor([1, 1]), |
| 1565 | 0, |
| 1566 | ) |
| 1567 | |
| 1568 | assert torch.equal(sampled, torch.tensor([3, 4])) |