vllm.v1.watermarking

Modules:

• gpu_sampler –
• gumbel –

  Gumbel-max watermark generation and detection primitives.

• prfs –
• spec_decode –
• watermarker –

Classes:

• DualKeyGumbelWatermarkDetector –

  Detect a dual-key watermark with weighted early fusion.

• PhiloxPRF –

  Accelerator-friendly Philox4x32-10 watermark PRF.

• SupportsSpeculativeDecoding –

  A watermarker compatible with standard speculative rejection sampling.

• Watermarker –

DualKeyGumbelWatermarkDetector

Bases: GumbelWatermarkDetector

Detect a dual-key watermark with weighted early fusion.

alpha is the weight assigned to key B; key A receives 1 - alpha.

Source code in vllm/v1/watermarking/gumbel.py

class DualKeyGumbelWatermarkDetector(GumbelWatermarkDetector):
    """Detect a dual-key watermark with weighted early fusion.

    ``alpha`` is the weight assigned to key B; key A receives ``1 - alpha``.
    """

    def __init__(
        self,
        key: int,
        context_width: int = 4,
        p_value_threshold: float = 0.01,
        prf: WatermarkPRFName = "philox",
        deduplicate_contexts: bool = True,
        alpha: float = 0.1,
    ) -> None:
        if not 0 <= alpha <= 1:
            raise ValueError("alpha must be between 0 and 1")
        super().__init__(
            derive_watermark_key(key, b"key_a"),
            context_width,
            p_value_threshold,
            prf,
            deduplicate_contexts,
        )
        self.key_b_prf = create_prf(prf, derive_watermark_key(key, b"key_b"))
        self.alpha = alpha

    def _score_tokens(
        self, contexts: torch.Tensor, targets: torch.Tensor
    ) -> torch.Tensor:
        scores = torch.stack(
            [
                -torch.log1p(
                    -prf.uniform(contexts, targets.unsqueeze(-1))
                    .squeeze(-1)
                    .to(torch.float64)
                )
                for prf in (self.prf, self.key_b_prf)
            ],
            dim=-1,
        )
        weights = scores.new_tensor([1 - self.alpha, self.alpha])
        return scores @ weights

    def _get_p_value(self, score: float, num_scored_tokens: int) -> float:
        variance = (1 - self.alpha) ** 2 + self.alpha**2
        if self.alpha in (0, 0.5, 1):
            return _gamma_survival_integer_shape(
                score / variance, round(num_scored_tokens / variance)
            )
        shape = torch.tensor(num_scored_tokens / variance, dtype=torch.float64)
        scaled_score = torch.tensor(score / variance, dtype=torch.float64)
        return torch.special.gammaincc(shape, scaled_score).item()

PhiloxPRF

Bases: WatermarkPRF

Accelerator-friendly Philox4x32-10 watermark PRF.

Source code in vllm/v1/watermarking/prfs/philox.py

class PhiloxPRF(WatermarkPRF):
    """Accelerator-friendly Philox4x32-10 watermark PRF."""

    version = "philox4x32-10-v1"

    def __init__(self, key: int) -> None:
        if not 0 <= key <= 2**64 - 1:
            raise ValueError("Philox keys must fit in 64 bits")
        self.key = key

    def uniform(self, contexts: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
        key_words = self.key & _UINT32_MASK, self.key >> 32
        if contexts.device.type == "cuda":
            return _compiled_philox_uniform(*key_words, contexts, token_ids)
        return _philox_uniform(*key_words, contexts, token_ids)

SupportsSpeculativeDecoding

Bases: Protocol

A watermarker compatible with standard speculative rejection sampling.

Source code in vllm/v1/watermarking/watermarker.py

@runtime_checkable
class SupportsSpeculativeDecoding(Protocol):
    """A watermarker compatible with standard speculative rejection sampling."""

    draft_watermarker: Watermarker
    target_watermarker: Watermarker

Watermarker

Bases: ABC

Source code in vllm/v1/watermarking/watermarker.py

class Watermarker(ABC):
    @property
    @abstractmethod
    def context_width(self) -> int:
        raise NotImplementedError

    def sample(
        self,
        logits: torch.Tensor,
        contexts: torch.Tensor,
        random_sampler: RandomSampler | None = None,
        skip_mask: torch.Tensor | None = None,
    ) -> WatermarkSample:
        if skip_mask is not None:
            if random_sampler is None:
                raise ValueError("skip_mask requires a random sampler")
            mixed = self._try_sample_mixed(logits, contexts, skip_mask, random_sampler)
            if mixed is not None:
                return mixed

        watermarked = self._sample_watermarked(logits, contexts)
        if skip_mask is None:
            return watermarked

        assert random_sampler is not None
        token_ids = torch.where(
            skip_mask, random_sampler(logits), watermarked.token_ids
        )
        output_logits = watermarked.logits
        if output_logits is not logits:
            output_logits = torch.where(skip_mask.unsqueeze(-1), logits, output_logits)
        return WatermarkSample(token_ids, output_logits)

    @abstractmethod
    def _sample_watermarked(
        self,
        logits: torch.Tensor,
        contexts: torch.Tensor,
    ) -> WatermarkSample:
        raise NotImplementedError

    def _try_sample_mixed(
        self,
        logits: torch.Tensor,
        contexts: torch.Tensor,
        skip_mask: torch.Tensor,
        random_sampler: RandomSampler,
    ) -> WatermarkSample | None:
        return None

    def _try_sample_block(
        self,
        logits: torch.Tensor,
        num_steps: int,
        random_sampler: RandomSampler,
        state: DraftContextState,
    ) -> torch.Tensor | None:
        """Sample a (request, step) draft block in one fused pass.

        Returns None when there is no fused implementation, in which case the
        caller samples the block one step at a time.
        """
        return None

_try_sample_block(logits, num_steps, random_sampler, state)

Sample a (request, step) draft block in one fused pass.

Returns None when there is no fused implementation, in which case the caller samples the block one step at a time.

Source code in vllm/v1/watermarking/watermarker.py

def _try_sample_block(
    self,
    logits: torch.Tensor,
    num_steps: int,
    random_sampler: RandomSampler,
    state: DraftContextState,
) -> torch.Tensor | None:
    """Sample a (request, step) draft block in one fused pass.

    Returns None when there is no fused implementation, in which case the
    caller samples the block one step at a time.
    """
    return None