| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Philox4x32-10 as defined by https://doi.org/10.1145/2063384.2063405.""" |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.v1.watermarking.prfs.base import WatermarkPRF, uint32_to_uniform |
| 9 | |
| 10 | _UINT32_MASK = 2**32 - 1 |
| 11 | _PHILOX_M0 = 0xD2511F53 |
| 12 | _PHILOX_M1 = 0xCD9E8D57 |
| 13 | _PHILOX_W0 = 0x9E3779B9 |
| 14 | _PHILOX_W1 = 0xBB67AE85 |
| 15 | _CONTEXT_DOMAIN = 0x574D4358 |
| 16 | _TOKEN_DOMAIN = 0x574D544B |
| 17 | |
| 18 | |
| 19 | def _mulhilo32( |
| 20 | multiplier: int, values: torch.Tensor |
| 21 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 22 | value_low = values & 0xFFFF |
| 23 | value_high = values >> 16 |
| 24 | multiplier_low = multiplier & 0xFFFF |
| 25 | multiplier_high = multiplier >> 16 |
| 26 | |
| 27 | low_product = value_low * multiplier_low |
| 28 | cross_product = value_high * multiplier_low + value_low * multiplier_high |
| 29 | low_with_carry = low_product + ((cross_product & 0xFFFF) << 16) |
| 30 | low = low_with_carry & _UINT32_MASK |
| 31 | high = ( |
| 32 | value_high * multiplier_high + (cross_product >> 16) + (low_with_carry >> 32) |
| 33 | ) & _UINT32_MASK |
| 34 | return high, low |
| 35 | |
| 36 | |
| 37 | def _philox4x32_10( |
| 38 | counter: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], |
| 39 | key: tuple[torch.Tensor, torch.Tensor], |
| 40 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 41 | counter_0, counter_1, counter_2, counter_3 = counter |
| 42 | key_0, key_1 = key |
| 43 | for round_index in range(10): |
| 44 | high_0, low_0 = _mulhilo32(_PHILOX_M0, counter_0) |
| 45 | high_1, low_1 = _mulhilo32(_PHILOX_M1, counter_2) |
| 46 | counter_0, counter_1, counter_2, counter_3 = ( |
| 47 | (high_1 ^ counter_1 ^ key_0) & _UINT32_MASK, |
| 48 | low_1, |
| 49 | (high_0 ^ counter_3 ^ key_1) & _UINT32_MASK, |
| 50 | low_0, |
| 51 | ) |
| 52 | if round_index != 9: |
| 53 | key_0 = (key_0 + _PHILOX_W0) & _UINT32_MASK |
| 54 | key_1 = (key_1 + _PHILOX_W1) & _UINT32_MASK |
| 55 | return counter_0, counter_1, counter_2, counter_3 |
| 56 | |
| 57 | |
| 58 | def _philox_uniform( |
| 59 | key_0_value: int, |
| 60 | key_1_value: int, |
| 61 | contexts: torch.Tensor, |
| 62 | token_ids: torch.Tensor, |
| 63 | ) -> torch.Tensor: |
| 64 | contexts = contexts.to(torch.int64) & _UINT32_MASK |
| 65 | prefix_shape = contexts.shape[:-1] |
| 66 | device = contexts.device |
| 67 | key_0 = torch.full(prefix_shape, key_0_value, dtype=torch.int64, device=device) |
| 68 | key_1 = torch.full(prefix_shape, key_1_value, dtype=torch.int64, device=device) |
| 69 | state = ( |
| 70 | torch.full(prefix_shape, _CONTEXT_DOMAIN, dtype=torch.int64, device=device), |
| 71 | torch.full( |
| 72 | prefix_shape, |
| 73 | contexts.shape[-1], |
| 74 | dtype=torch.int64, |
| 75 | device=device, |
| 76 | ), |
| 77 | torch.zeros(prefix_shape, dtype=torch.int64, device=device), |
| 78 | torch.zeros(prefix_shape, dtype=torch.int64, device=device), |
| 79 | ) |
| 80 | |
| 81 | for offset in range(0, contexts.shape[-1], 4): |
| 82 | block = [] |
| 83 | for index in range(4): |
| 84 | position = offset + index |
| 85 | if position < contexts.shape[-1]: |
| 86 | block.append(contexts[..., position]) |
| 87 | else: |
| 88 | block.append( |
| 89 | torch.full( |
| 90 | prefix_shape, |
| 91 | _UINT32_MASK - index, |
| 92 | dtype=torch.int64, |
| 93 | device=device, |
| 94 | ) |
| 95 | ) |
| 96 | counter = tuple( |
| 97 | (state[index] ^ block[index]) & _UINT32_MASK for index in range(4) |
| 98 | ) |
| 99 | state = _philox4x32_10(counter, ((key_0 ^ offset) & _UINT32_MASK, key_1)) |
| 100 | |
| 101 | output_shape = torch.broadcast_shapes(prefix_shape + (1,), token_ids.shape) |
| 102 | tokens = torch.broadcast_to(token_ids.to(device=device), output_shape) |
| 103 | state = tuple( |
| 104 | torch.broadcast_to(word.unsqueeze(-1), output_shape) for word in state |
| 105 | ) |
| 106 | candidate_key_0 = torch.broadcast_to(key_0.unsqueeze(-1), output_shape) |
| 107 | candidate_key_1 = torch.broadcast_to(key_1.unsqueeze(-1), output_shape) |
| 108 | token_words = tokens.to(torch.int64) & _UINT32_MASK |
| 109 | outputs = _philox4x32_10( |
| 110 | ( |
| 111 | token_words >> 2, |
| 112 | state[0], |
| 113 | state[1], |
| 114 | state[2], |
| 115 | ), |
| 116 | ( |
| 117 | (candidate_key_0 ^ state[3]) & _UINT32_MASK, |
| 118 | (candidate_key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK, |
| 119 | ), |
| 120 | ) |
| 121 | word_index = token_words & 3 |
| 122 | output = outputs[0] |
| 123 | for index in range(1, 4): |
| 124 | output = torch.where(word_index == index, outputs[index], output) |
| 125 | return uint32_to_uniform(output) |
| 126 | |
| 127 | |
| 128 | _compiled_philox_uniform = torch.compile(_philox_uniform, fullgraph=True, dynamic=True) |
| 129 | |
| 130 | |
| 131 | class PhiloxPRF(WatermarkPRF): |
| 132 | """Accelerator-friendly Philox4x32-10 watermark PRF.""" |
| 133 | |
| 134 | version = "philox4x32-10-v1" |
| 135 | |
| 136 | def __init__(self, key: int) -> None: |
| 137 | if not 0 <= key <= 2**64 - 1: |
| 138 | raise ValueError("Philox keys must fit in 64 bits") |
| 139 | self.key = key |
| 140 | |
| 141 | def uniform(self, contexts: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: |
| 142 | key_words = self.key & _UINT32_MASK, self.key >> 32 |
| 143 | if contexts.device.type == "cuda": |
| 144 | return _compiled_philox_uniform(*key_words, contexts, token_ids) |
| 145 | return _philox_uniform(*key_words, contexts, token_ids) |