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.

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.2,
    ) -> 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