watermarking - vLLM
Skip to content

vllm.v1.watermarking

Modules:

  • gumbel –

    Gumbel-max watermark generation and detection primitives.

  • prfs –

Classes:

  • PhiloxPRF –

    Accelerator-friendly Philox4x32-10 watermark PRF.

PhiloxPRF

Bases: WatermarkPRF

Accelerator-friendly Philox4x32-10 watermark PRF.

Source code in vllm/v1/watermarking/prfs/philox.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
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)