vllm.v1.watermarking.gumbel

Gumbel-max watermark generation and detection primitives.

Classes:

• DualKeyGumbelWatermarkDetector –

  Detect a dual-key watermark with weighted early fusion.

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()

_gamma_survival_integer_shape(score, shape)

Return the survival probability of a Gamma(shape, 1) variable.

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

def _gamma_survival_integer_shape(score: float, shape: int) -> float:
    """Return the survival probability of a Gamma(shape, 1) variable."""
    if shape <= 0 or score <= 0:
        return 1.0
    log_score = math.log(score)
    log_term = 0.0
    max_log_term = 0.0
    log_terms = [log_term]
    for index in range(1, shape):
        log_term += log_score - math.log(index)
        log_terms.append(log_term)
        max_log_term = max(max_log_term, log_term)
    log_sum = max_log_term + math.log(
        sum(math.exp(term - max_log_term) for term in log_terms)
    )
    return min(1.0, math.exp(min(0.0, -score + log_sum)))