| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import ABC, abstractmethod |
| 5 | from dataclasses import dataclass |
| 6 | from typing import Protocol, runtime_checkable |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample |
| 11 | |
| 12 | |
| 13 | @dataclass(frozen=True) |
| 14 | class WatermarkSample: |
| 15 | token_ids: torch.Tensor |
| 16 | logits: torch.Tensor |
| 17 | |
| 18 | |
| 19 | @dataclass(frozen=True) |
| 20 | class RandomSampler: |
| 21 | expanded_idx_mapping: torch.Tensor |
| 22 | temperatures: torch.Tensor |
| 23 | seeds: torch.Tensor |
| 24 | positions: torch.Tensor |
| 25 | use_fp64: bool = False |
| 26 | is_drafting: bool = False |
| 27 | logits_cache: torch.Tensor | None = None |
| 28 | logits_cache_col: torch.Tensor | None = None |
| 29 | logits_cache_source: torch.Tensor | None = None |
| 30 | |
| 31 | def __call__(self, logits: torch.Tensor) -> torch.Tensor: |
| 32 | return gumbel_sample( |
| 33 | logits, |
| 34 | self.expanded_idx_mapping, |
| 35 | self.temperatures, |
| 36 | self.seeds, |
| 37 | self.positions, |
| 38 | apply_temperature=False, |
| 39 | is_drafting=self.is_drafting, |
| 40 | logits_cache=self.logits_cache, |
| 41 | logits_cache_col=self.logits_cache_col, |
| 42 | logits_cache_source=self.logits_cache_source, |
| 43 | use_fp64=self.use_fp64, |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | class Watermarker(ABC): |
| 48 | @property |
| 49 | @abstractmethod |
| 50 | def context_width(self) -> int: |
| 51 | raise NotImplementedError |
| 52 | |
| 53 | def sample( |
| 54 | self, |
| 55 | logits: torch.Tensor, |
| 56 | contexts: torch.Tensor, |
| 57 | random_sampler: RandomSampler | None = None, |
| 58 | skip_mask: torch.Tensor | None = None, |
| 59 | ) -> WatermarkSample: |
| 60 | if skip_mask is not None: |
| 61 | if random_sampler is None: |
| 62 | raise ValueError("skip_mask requires a random sampler") |
| 63 | mixed = self._try_sample_mixed(logits, contexts, skip_mask, random_sampler) |
| 64 | if mixed is not None: |
| 65 | return mixed |
| 66 | |
| 67 | watermarked = self._sample_watermarked(logits, contexts) |
| 68 | if skip_mask is None: |
| 69 | return watermarked |
| 70 | |
| 71 | assert random_sampler is not None |
| 72 | token_ids = torch.where( |
| 73 | skip_mask, random_sampler(logits), watermarked.token_ids |
| 74 | ) |
| 75 | output_logits = watermarked.logits |
| 76 | if output_logits is not logits: |
| 77 | output_logits = torch.where(skip_mask.unsqueeze(-1), logits, output_logits) |
| 78 | return WatermarkSample(token_ids, output_logits) |
| 79 | |
| 80 | @abstractmethod |
| 81 | def _sample_watermarked( |
| 82 | self, |
| 83 | logits: torch.Tensor, |
| 84 | contexts: torch.Tensor, |
| 85 | ) -> WatermarkSample: |
| 86 | raise NotImplementedError |
| 87 | |
| 88 | def _try_sample_mixed( |
| 89 | self, |
| 90 | logits: torch.Tensor, |
| 91 | contexts: torch.Tensor, |
| 92 | skip_mask: torch.Tensor, |
| 93 | random_sampler: RandomSampler, |
| 94 | ) -> WatermarkSample | None: |
| 95 | return None |
| 96 | |
| 97 | |
| 98 | @runtime_checkable |
| 99 | class SupportsSpeculativeDecoding(Protocol): |
| 100 | """A watermarker compatible with standard speculative rejection sampling.""" |
| 101 | |
| 102 | draft_watermarker: Watermarker |
| 103 | target_watermarker: Watermarker |