| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import hashlib |
| 5 | from typing import Literal |
| 6 | |
| 7 | from pydantic import Field, model_validator |
| 8 | from typing_extensions import Self |
| 9 | |
| 10 | from vllm.config.utils import config |
| 11 | from vllm.logger import init_logger |
| 12 | |
| 13 | logger = init_logger(__name__) |
| 14 | |
| 15 | WatermarkingAlgorithm = Literal["gumbel", "dual_key_gumbel"] |
| 16 | WatermarkPRFName = Literal["philox"] |
| 17 | WatermarkContextScope = Literal["none", "single_turn", "all"] |
| 18 | |
| 19 | _MIN_RECOMMENDED_DEDUP_HISTORY = 1024 |
| 20 | |
| 21 | |
| 22 | def derive_watermark_key(key: int, domain: bytes) -> int: |
| 23 | digest = hashlib.sha256(domain + key.to_bytes(8, "big")).digest() |
| 24 | return int.from_bytes(digest[:8], "big") |
| 25 | |
| 26 | |
| 27 | @config |
| 28 | class WatermarkConfig: |
| 29 | """Configuration for text watermark generation.""" |
| 30 | |
| 31 | key: int = Field(ge=0, repr=False, exclude=True) |
| 32 | """Secret key used to watermark generated text.""" |
| 33 | algorithm: WatermarkingAlgorithm = "gumbel" |
| 34 | """Algorithm used to watermark generated text.""" |
| 35 | alpha: float = Field(default=0.1, ge=0, le=1) |
| 36 | """Probability of selecting key B for dual-key watermarking.""" |
| 37 | context_width: int = Field(default=4, ge=1) |
| 38 | """Number of prior tokens used by the watermark PRF.""" |
| 39 | deduplicate_contexts: WatermarkContextScope = "single_turn" |
| 40 | """Which history is searched for a repeated context before a token is |
| 41 | sampled; a repeated context is sampled without the watermark. `none` |
| 42 | disables the search. `single_turn` (default) searches this request's |
| 43 | generated tokens. `all` also searches the prompt and samples the first |
| 44 | `context_width` generated tokens without watermarking.""" |
| 45 | deduplicate_contexts_max_history: int | None = Field(default=8192, ge=1) |
| 46 | """Number of most recent history positions searched (default 8192), or |
| 47 | `None` for the whole scope. Each position is compared over the |
| 48 | `context_width` tokens before it. Ignored when `deduplicate_contexts` is |
| 49 | `none`.""" |
| 50 | prf: WatermarkPRFName = "philox" |
| 51 | """Pseudorandom function used by the watermarking algorithm.""" |
| 52 | allow_target_only_watermarking: bool = False |
| 53 | """Allow speculative decoding without watermarking draft tokens.""" |
| 54 | |
| 55 | @property |
| 56 | def supports_speculative_decoding(self) -> bool: |
| 57 | return self.algorithm == "dual_key_gumbel" |
| 58 | |
| 59 | @model_validator(mode="after") |
| 60 | def validate_watermark_settings(self) -> Self: |
| 61 | if self.key > 2**64 - 1: |
| 62 | raise ValueError("philox keys must fit in 64 bits") |
| 63 | history_is_too_short = ( |
| 64 | self.deduplicate_contexts_max_history is not None |
| 65 | and self.deduplicate_contexts_max_history < _MIN_RECOMMENDED_DEDUP_HISTORY |
| 66 | ) |
| 67 | if self.algorithm in ("gumbel", "dual_key_gumbel") and ( |
| 68 | self.deduplicate_contexts == "none" or history_is_too_short |
| 69 | ): |
| 70 | logger.warning_once( |
| 71 | "Gumbel-max watermarking with context deduplication " |
| 72 | "disabled or limited to fewer than " |
| 73 | f"{_MIN_RECOMMENDED_DEDUP_HISTORY} positions may increase the " |
| 74 | "frequency of degenerate generations, including repetition loops. " |
| 75 | "Use deduplicate_contexts='single_turn' or 'all' with " |
| 76 | "deduplicate_contexts_max_history at least " |
| 77 | f"{_MIN_RECOMMENDED_DEDUP_HISTORY} or null to mitigate this.", |
| 78 | scope="global", |
| 79 | ) |
| 80 | return self |