| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The warmup must cover every specialization the sampler will launch. |
| 4 | |
| 5 | ``_philox_gumbel_kernel`` is keyed on the Philox key halves, on the logits |
| 6 | dtype (model dtype when the sampler passes logits through, fp32 when logits |
| 7 | processing copies them) and on whether a skip mask is passed, so a missing |
| 8 | combination leaves a request paying the compile this module exists to remove. |
| 9 | """ |
| 10 | |
| 11 | import logging |
| 12 | from types import SimpleNamespace |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | |
| 17 | from vllm.config.watermarking import WatermarkConfig, derive_watermark_key |
| 18 | from vllm.model_executor.warmup.watermark_sample_warmup import watermark_sample_warmup |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.v1.watermarking.factory import create_watermarker |
| 21 | from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler |
| 22 | from vllm.v1.watermarking.spec_decode import create_speculative_target_watermarker |
| 23 | |
| 24 | MASTER_KEY = 20250910 |
| 25 | KEY_A = derive_watermark_key(MASTER_KEY, b"key_a") |
| 26 | KEY_B = derive_watermark_key(MASTER_KEY, b"key_b") |
| 27 | |
| 28 | |
| 29 | @pytest.fixture(autouse=True) |
| 30 | def cuda_alike_platform(monkeypatch: pytest.MonkeyPatch) -> None: |
| 31 | """The warmup returns early off CUDA, so pin the platform it reads. |
| 32 | |
| 33 | Everything here runs on the fake worker's cpu device; without this the |
| 34 | tests below would pass only on a CUDA-alike host and silently warm nothing |
| 35 | anywhere else. |
| 36 | """ |
| 37 | monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True) |
| 38 | |
| 39 | |
| 40 | def _sampler( |
| 41 | watermark_config: WatermarkConfig, |
| 42 | speculative: bool, |
| 43 | use_fp64_gumbel: bool, |
| 44 | ) -> GPUWatermarkSampler: |
| 45 | """The sampler ``GPUModelRunner`` builds for this config. |
| 46 | |
| 47 | Mirrors ``vllm/v1/worker/gpu/model_runner.py:473-486``. ``__new__`` skips |
| 48 | ``Sampler.__init__``, which allocates UVA-backed device buffers; the warmup |
| 49 | only reads the three attributes set here. |
| 50 | """ |
| 51 | watermarker = create_watermarker(watermark_config) |
| 52 | if speculative: |
| 53 | watermarker = create_speculative_target_watermarker(watermarker) |
| 54 | sampler = GPUWatermarkSampler.__new__(GPUWatermarkSampler) |
| 55 | sampler.watermarker = watermarker |
| 56 | sampler.deduplicate_contexts = watermark_config.deduplicate_contexts |
| 57 | sampler.use_fp64_gumbel = use_fp64_gumbel |
| 58 | return sampler |
| 59 | |
| 60 | |
| 61 | def _worker( |
| 62 | watermark_config: WatermarkConfig | None, |
| 63 | speculative: bool = False, |
| 64 | use_fp64_gumbel: bool = False, |
| 65 | with_sampler: bool = True, |
| 66 | ) -> SimpleNamespace: |
| 67 | sampler = ( |
| 68 | _sampler(watermark_config, speculative, use_fp64_gumbel) |
| 69 | if watermark_config is not None and with_sampler |
| 70 | else None |
| 71 | ) |
| 72 | return SimpleNamespace( |
| 73 | device=torch.device("cpu"), |
| 74 | model_runner=SimpleNamespace(sampler=sampler), |
| 75 | vllm_config=SimpleNamespace( |
| 76 | model_config=SimpleNamespace( |
| 77 | get_vocab_size=lambda: 128, |
| 78 | dtype=torch.bfloat16, |
| 79 | ), |
| 80 | watermark_config=watermark_config, |
| 81 | ), |
| 82 | ) |
| 83 | |
| 84 | |
| 85 | def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]: |
| 86 | calls: list[dict] = [] |
| 87 | |
| 88 | def fake_philox_gumbel_sample(logits, contexts, key, **kwargs): |
| 89 | calls.append({"logits": logits, "contexts": contexts, "key": key, **kwargs}) |
| 90 | return torch.zeros(logits.shape[0], dtype=torch.int64) |
| 91 | |
| 92 | monkeypatch.setattr( |
| 93 | "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample", |
| 94 | fake_philox_gumbel_sample, |
| 95 | ) |
| 96 | return calls |
| 97 | |
| 98 | |
| 99 | def _variants(calls: list[dict]) -> set[tuple[int, torch.dtype, bool]]: |
| 100 | return { |
| 101 | (call["key"], call["logits"].dtype, call.get("skip_mask") is not None) |
| 102 | for call in calls |
| 103 | } |
| 104 | |
| 105 | |
| 106 | def test_no_warmup_without_a_watermark_config( |
| 107 | monkeypatch: pytest.MonkeyPatch, |
| 108 | ) -> None: |
| 109 | calls = _capture_calls(monkeypatch) |
| 110 | watermark_sample_warmup(_worker(None)) |
| 111 | |
| 112 | assert calls == [] |
| 113 | |
| 114 | |
| 115 | def test_no_warmup_without_a_watermarked_sampler( |
| 116 | monkeypatch: pytest.MonkeyPatch, |
| 117 | ) -> None: |
| 118 | """A rank that does not sample launches nothing to warm. |
| 119 | |
| 120 | Only the final pipeline stage of a generative model holds a |
| 121 | ``GPUWatermarkSampler``; on any other rank the kernel is never launched. |
| 122 | """ |
| 123 | calls = _capture_calls(monkeypatch) |
| 124 | watermark_sample_warmup( |
| 125 | _worker(WatermarkConfig(key=MASTER_KEY), with_sampler=False) |
| 126 | ) |
| 127 | |
| 128 | assert calls == [] |
| 129 | |
| 130 | |
| 131 | def test_no_warmup_off_cuda(monkeypatch: pytest.MonkeyPatch) -> None: |
| 132 | """``GumbelWatermarker`` launches the kernel only on a CUDA-alike device. |
| 133 | |
| 134 | Elsewhere -- an XPU worker inherits ``kernel_warmup`` from the CUDA worker |
| 135 | -- the sampler takes the torch path, so warming would compile a kernel |
| 136 | nothing launches. |
| 137 | """ |
| 138 | calls = _capture_calls(monkeypatch) |
| 139 | monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: False) |
| 140 | watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY))) |
| 141 | |
| 142 | assert calls == [] |
| 143 | |
| 144 | |
| 145 | def test_single_key_warms_both_logits_dtypes( |
| 146 | monkeypatch: pytest.MonkeyPatch, |
| 147 | ) -> None: |
| 148 | calls = _capture_calls(monkeypatch) |
| 149 | watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY))) |
| 150 | |
| 151 | assert _variants(calls) == { |
| 152 | (MASTER_KEY, torch.bfloat16, True), |
| 153 | (MASTER_KEY, torch.float32, True), |
| 154 | } |
| 155 | assert len(calls) == len(_variants(calls)) |
| 156 | |
| 157 | |
| 158 | def test_dual_key_warms_both_derived_keys(monkeypatch: pytest.MonkeyPatch) -> None: |
| 159 | """``DualKeyGumbelWatermarker.sample`` launches the kernel once per key. |
| 160 | |
| 161 | The halves of a derived key exceed 2**31, so each key is typed i64/i32 by |
| 162 | magnitude and a wrong key compiles a specialization the engine never uses. |
| 163 | """ |
| 164 | calls = _capture_calls(monkeypatch) |
| 165 | watermark_sample_warmup( |
| 166 | _worker(WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel")) |
| 167 | ) |
| 168 | |
| 169 | assert _variants(calls) == { |
| 170 | (KEY_A, torch.bfloat16, True), |
| 171 | (KEY_A, torch.float32, True), |
| 172 | (KEY_B, torch.bfloat16, True), |
| 173 | (KEY_B, torch.float32, True), |
| 174 | } |
| 175 | |
| 176 | |
| 177 | @pytest.mark.parametrize("alpha,key", [(0.0, KEY_A), (1.0, KEY_B)]) |
| 178 | def test_dual_key_alpha_bounds_warm_one_key( |
| 179 | monkeypatch: pytest.MonkeyPatch, alpha: float, key: int |
| 180 | ) -> None: |
| 181 | calls = _capture_calls(monkeypatch) |
| 182 | watermark_sample_warmup( |
| 183 | _worker( |
| 184 | WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel", alpha=alpha) |
| 185 | ) |
| 186 | ) |
| 187 | |
| 188 | assert _variants(calls) == { |
| 189 | (key, torch.bfloat16, True), |
| 190 | (key, torch.float32, True), |
| 191 | } |
| 192 | |
| 193 | |
| 194 | def test_speculative_config_warms_the_target_role_only( |
| 195 | monkeypatch: pytest.MonkeyPatch, |
| 196 | ) -> None: |
| 197 | """Under spec decode the sampler holds ``create_target_watermarker()``. |
| 198 | |
| 199 | Rows without draft tokens sample through it, so the sampler path launches |
| 200 | key B alone; the draft key runs in the speculator's own sampler. |
| 201 | """ |
| 202 | calls = _capture_calls(monkeypatch) |
| 203 | watermark_sample_warmup( |
| 204 | _worker( |
| 205 | WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel"), |
| 206 | speculative=True, |
| 207 | ) |
| 208 | ) |
| 209 | |
| 210 | assert _variants(calls) == { |
| 211 | (KEY_B, torch.bfloat16, True), |
| 212 | (KEY_B, torch.float32, True), |
| 213 | } |
| 214 | |
| 215 | |
| 216 | def test_dedup_none_also_warms_the_mask_free_specialization( |
| 217 | monkeypatch: pytest.MonkeyPatch, |
| 218 | ) -> None: |
| 219 | """Without deduplication an all-watermarked batch passes no skip mask. |
| 220 | |
| 221 | A later batch mixing in an opted-out request passes one, so both |
| 222 | specializations are reachable and both must be warmed. |
| 223 | """ |
| 224 | calls = _capture_calls(monkeypatch) |
| 225 | watermark_sample_warmup( |
| 226 | _worker(WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none")) |
| 227 | ) |
| 228 | |
| 229 | assert _variants(calls) == { |
| 230 | (MASTER_KEY, torch.bfloat16, True), |
| 231 | (MASTER_KEY, torch.float32, True), |
| 232 | (MASTER_KEY, torch.bfloat16, False), |
| 233 | (MASTER_KEY, torch.float32, False), |
| 234 | } |
| 235 | for call in calls: |
| 236 | if call.get("skip_mask") is None: |
| 237 | assert "temperatures" not in call |
| 238 | |
| 239 | |
| 240 | def test_warmup_argument_dtypes_match_the_runtime_buffers( |
| 241 | monkeypatch: pytest.MonkeyPatch, |
| 242 | ) -> None: |
| 243 | """Contexts are int32, the sampling state int64 indices and fp32 temperatures. |
| 244 | |
| 245 | ``GPUWatermarkSampler._get_contexts`` gathers from the int32 request-state |
| 246 | token ids, ``expanded_idx_mapping``/``positions``/``seeds`` are int64 and |
| 247 | ``SamplingStates.temperature`` is fp32; warming other dtypes compiles |
| 248 | specializations no engine launches. |
| 249 | """ |
| 250 | calls = _capture_calls(monkeypatch) |
| 251 | watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY, context_width=3))) |
| 252 | |
| 253 | assert calls |
| 254 | for call in calls: |
| 255 | num_tokens = call["logits"].shape[0] |
| 256 | assert call["contexts"].dtype == torch.int32 |
| 257 | assert call["contexts"].shape == (num_tokens, 3) |
| 258 | assert call["skip_mask"].dtype == torch.bool |
| 259 | assert call["expanded_idx_mapping"].dtype == torch.int64 |
| 260 | assert call["positions"].dtype == torch.int64 |
| 261 | assert call["seeds"].dtype == torch.int64 |
| 262 | assert call["temperatures"].dtype == torch.float32 |
| 263 | |
| 264 | |
| 265 | def test_use_fp64_follows_the_path_that_forwards_it( |
| 266 | monkeypatch: pytest.MonkeyPatch, |
| 267 | ) -> None: |
| 268 | """Only the mixed path forwards the sampler's ``use_fp64_gumbel``. |
| 269 | |
| 270 | ``GumbelWatermarker._sample_watermarked`` launches the mask-free variant |
| 271 | with the ``use_fp64=False`` default, so warming it with fp64 would compile |
| 272 | a specialization the engine never launches and miss the one it does. |
| 273 | """ |
| 274 | calls = _capture_calls(monkeypatch) |
| 275 | watermark_sample_warmup( |
| 276 | _worker( |
| 277 | WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none"), |
| 278 | use_fp64_gumbel=True, |
| 279 | ) |
| 280 | ) |
| 281 | |
| 282 | assert calls |
| 283 | for call in calls: |
| 284 | assert call["use_fp64"] is (call.get("skip_mask") is not None) |
| 285 | |
| 286 | |
| 287 | def test_warmup_failure_does_not_abort_startup( |
| 288 | monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture |
| 289 | ) -> None: |
| 290 | def failing_philox_gumbel_sample(*args, **kwargs): |
| 291 | raise RuntimeError("no triton") |
| 292 | |
| 293 | monkeypatch.setattr( |
| 294 | "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample", |
| 295 | failing_philox_gumbel_sample, |
| 296 | ) |
| 297 | with caplog.at_level(logging.WARNING): |
| 298 | watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY))) |
| 299 | |
| 300 | assert "Skipping watermark sampler warmup." in caplog.text |