| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.config.watermarking import WatermarkConfig, WatermarkContextScope |
| 7 | from vllm.v1.watermarking.factory import create_watermarker |
| 8 | from vllm.v1.watermarking.gumbel import GumbelWatermarker |
| 9 | from vllm.v1.watermarking.prfs import PhiloxPRF |
| 10 | from vllm.v1.watermarking.watermarker import ( |
| 11 | RandomSampler, |
| 12 | SupportsSpeculativeDecoding, |
| 13 | Watermarker, |
| 14 | ) |
| 15 | from vllm.v1.worker.gpu.sample.watermark import draft_watermarking_mask |
| 16 | from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample |
| 17 | |
| 18 | |
| 19 | class DraftWatermarker: |
| 20 | def __init__( |
| 21 | self, |
| 22 | watermarker: Watermarker, |
| 23 | max_num_reqs: int, |
| 24 | device: torch.device, |
| 25 | num_speculative_steps: int, |
| 26 | deduplicate_contexts: WatermarkContextScope, |
| 27 | deduplicate_contexts_max_history: int | None, |
| 28 | ) -> None: |
| 29 | self.watermarker = watermarker |
| 30 | self.num_speculative_steps = num_speculative_steps |
| 31 | self.deduplicate_contexts = deduplicate_contexts |
| 32 | self.deduplicate_contexts_max_history = deduplicate_contexts_max_history |
| 33 | self.contexts = torch.zeros( |
| 34 | max_num_reqs, |
| 35 | watermarker.context_width, |
| 36 | dtype=torch.int64, |
| 37 | device=device, |
| 38 | ) |
| 39 | self.enabled = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) |
| 40 | self.prior_contexts = torch.zeros( |
| 41 | max_num_reqs, |
| 42 | num_speculative_steps, |
| 43 | watermarker.context_width, |
| 44 | dtype=torch.int64, |
| 45 | device=device, |
| 46 | ) |
| 47 | self.all_token_ids: torch.Tensor | None = None |
| 48 | self.prompt_lens: torch.Tensor | None = None |
| 49 | self.total_lens: torch.Tensor | None = None |
| 50 | |
| 51 | def prepare( |
| 52 | self, |
| 53 | contexts: torch.Tensor, |
| 54 | enabled: torch.Tensor, |
| 55 | all_token_ids: torch.Tensor, |
| 56 | prompt_lens: torch.Tensor, |
| 57 | total_lens: torch.Tensor, |
| 58 | ) -> None: |
| 59 | num_reqs = contexts.shape[0] |
| 60 | self.contexts[:num_reqs].copy_(contexts) |
| 61 | self.enabled[:num_reqs].copy_(enabled) |
| 62 | self.all_token_ids = all_token_ids |
| 63 | self.prompt_lens = prompt_lens |
| 64 | self.total_lens = total_lens |
| 65 | |
| 66 | def _sampling_state( |
| 67 | self, |
| 68 | logits: torch.Tensor, |
| 69 | idx_mapping: torch.Tensor, |
| 70 | request_temperatures: torch.Tensor, |
| 71 | draft_step: int | torch.Tensor, |
| 72 | contexts: torch.Tensor, |
| 73 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 74 | num_rows = logits.shape[0] |
| 75 | steps = torch.as_tensor(draft_step, device=logits.device, dtype=torch.int64) |
| 76 | if steps.ndim == 0: |
| 77 | steps = steps.expand(num_rows) |
| 78 | enabled = self.enabled[:num_rows] & (request_temperatures != 0) |
| 79 | if self.deduplicate_contexts == "none": |
| 80 | return steps, enabled |
| 81 | assert self.all_token_ids is not None |
| 82 | assert self.prompt_lens is not None |
| 83 | assert self.total_lens is not None |
| 84 | enabled = draft_watermarking_mask( |
| 85 | self.all_token_ids, |
| 86 | idx_mapping, |
| 87 | self.prompt_lens, |
| 88 | self.total_lens, |
| 89 | self.prior_contexts, |
| 90 | contexts, |
| 91 | steps, |
| 92 | enabled, |
| 93 | self.deduplicate_contexts_max_history, |
| 94 | include_prompt=self.deduplicate_contexts == "all", |
| 95 | ) |
| 96 | return steps, enabled |
| 97 | |
| 98 | def sample( |
| 99 | self, |
| 100 | logits: torch.Tensor, |
| 101 | idx_mapping: torch.Tensor, |
| 102 | temperature: torch.Tensor, |
| 103 | seed: torch.Tensor, |
| 104 | pos: torch.Tensor, |
| 105 | apply_temperature: bool, |
| 106 | is_drafting: bool, |
| 107 | logits_cache: torch.Tensor | None = None, |
| 108 | logits_cache_col: torch.Tensor | None = None, |
| 109 | use_fp64: bool = False, |
| 110 | logits_cache_source: torch.Tensor | None = None, |
| 111 | ) -> torch.Tensor: |
| 112 | assert apply_temperature |
| 113 | assert is_drafting |
| 114 | assert logits_cache is not None |
| 115 | assert logits_cache_col is not None |
| 116 | num_rows = logits.shape[0] |
| 117 | request_temperatures = temperature[idx_mapping] |
| 118 | contexts = self.contexts[:num_rows] |
| 119 | steps, enabled = self._sampling_state( |
| 120 | logits, idx_mapping, request_temperatures, logits_cache_col, contexts |
| 121 | ) |
| 122 | |
| 123 | processed_logits = logits / torch.where( |
| 124 | request_temperatures == 0, 1, request_temperatures |
| 125 | ).unsqueeze(-1) |
| 126 | random_sampler = RandomSampler( |
| 127 | expanded_idx_mapping=idx_mapping, |
| 128 | temperatures=temperature, |
| 129 | seeds=seed, |
| 130 | positions=pos, |
| 131 | use_fp64=use_fp64, |
| 132 | is_drafting=True, |
| 133 | logits_cache=logits_cache, |
| 134 | logits_cache_col=steps, |
| 135 | logits_cache_source=( |
| 136 | logits if logits_cache_source is None else logits_cache_source |
| 137 | ), |
| 138 | ) |
| 139 | sampled = self.watermarker.sample( |
| 140 | processed_logits, |
| 141 | contexts, |
| 142 | random_sampler=random_sampler, |
| 143 | skip_mask=~enabled, |
| 144 | ).token_ids |
| 145 | |
| 146 | contexts.copy_(torch.cat((contexts[:, 1:], sampled.unsqueeze(-1)), dim=-1)) |
| 147 | return sampled |
| 148 | |
| 149 | |
| 150 | def create_speculative_target_watermarker(watermarker: Watermarker) -> Watermarker: |
| 151 | if isinstance(watermarker, SupportsSpeculativeDecoding): |
| 152 | return watermarker.target_watermarker |
| 153 | return watermarker |
| 154 | |
| 155 | |
| 156 | def create_speculative_draft_watermarker( |
| 157 | watermarker: Watermarker, |
| 158 | max_num_reqs: int, |
| 159 | device: torch.device, |
| 160 | allow_target_only: bool, |
| 161 | num_speculative_steps: int = 1, |
| 162 | deduplicate_contexts: WatermarkContextScope = "none", |
| 163 | deduplicate_contexts_max_history: int | None = 8192, |
| 164 | ) -> DraftWatermarker | None: |
| 165 | if isinstance(watermarker, SupportsSpeculativeDecoding): |
| 166 | return DraftWatermarker( |
| 167 | watermarker.draft_watermarker, |
| 168 | max_num_reqs, |
| 169 | device, |
| 170 | num_speculative_steps, |
| 171 | deduplicate_contexts, |
| 172 | deduplicate_contexts_max_history, |
| 173 | ) |
| 174 | if allow_target_only: |
| 175 | return None |
| 176 | raise ValueError( |
| 177 | f"The {type(watermarker).__name__} watermarking algorithm does not support " |
| 178 | "speculative decoding. Set allow_target_only_watermarking=true to " |
| 179 | "leave draft tokens unwatermarked." |
| 180 | ) |
| 181 | |
| 182 | |
| 183 | def _philox_key(watermarker: Watermarker) -> int: |
| 184 | if not ( |
| 185 | isinstance(watermarker, GumbelWatermarker) |
| 186 | and type(watermarker.prf) is PhiloxPRF |
| 187 | ): |
| 188 | raise NotImplementedError( |
| 189 | "In-kernel watermarked recovery supports only the Philox PRF " |
| 190 | "GumbelWatermarker" |
| 191 | ) |
| 192 | return watermarker.prf.key |
| 193 | |
| 194 | |
| 195 | def _resolve_watermark_key(watermarker: Watermarker) -> int: |
| 196 | """Return the Philox recovery key.""" |
| 197 | if isinstance(watermarker, SupportsSpeculativeDecoding): |
| 198 | raise ValueError( |
| 199 | f"{type(watermarker).__name__} keys the target role separately from " |
| 200 | "the draft. Pass create_speculative_target_watermarker(watermarker) " |
| 201 | "so the recovery draw carries the target's key." |
| 202 | ) |
| 203 | return _philox_key(watermarker) |
| 204 | |
| 205 | |
| 206 | def speculative_target_watermark_key( |
| 207 | watermark_config: WatermarkConfig | None, |
| 208 | ) -> int | None: |
| 209 | """Resolve the resample kernel's Philox key from configuration.""" |
| 210 | if watermark_config is None: |
| 211 | return None |
| 212 | |
| 213 | return _resolve_watermark_key( |
| 214 | create_speculative_target_watermarker(create_watermarker(watermark_config)) |
| 215 | ) |
| 216 | |
| 217 | |
| 218 | def watermarked_rejection_sample( |
| 219 | target_logits: torch.Tensor, |
| 220 | draft_logits: torch.Tensor | None, |
| 221 | draft_sampled: torch.Tensor, |
| 222 | cu_num_logits: torch.Tensor, |
| 223 | pos: torch.Tensor, |
| 224 | idx_mapping: torch.Tensor, |
| 225 | expanded_idx_mapping: torch.Tensor, |
| 226 | expanded_local_pos: torch.Tensor, |
| 227 | temperature: torch.Tensor, |
| 228 | seed: torch.Tensor, |
| 229 | num_speculative_steps: int, |
| 230 | contexts: torch.Tensor, |
| 231 | watermarking: torch.Tensor, |
| 232 | watermarker: Watermarker, |
| 233 | watermarking_skip_mask: torch.Tensor | None = None, |
| 234 | use_fp64: bool = False, |
| 235 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 236 | """Use the target key except at contexts selected for ordinary sampling.""" |
| 237 | assert contexts.shape == (target_logits.shape[0], watermarker.context_width) |
| 238 | return rejection_sample( |
| 239 | target_logits, |
| 240 | draft_logits, |
| 241 | draft_sampled, |
| 242 | cu_num_logits, |
| 243 | pos, |
| 244 | idx_mapping, |
| 245 | expanded_idx_mapping, |
| 246 | expanded_local_pos, |
| 247 | temperature, |
| 248 | seed, |
| 249 | num_speculative_steps, |
| 250 | use_fp64=use_fp64, |
| 251 | contexts=contexts, |
| 252 | watermarking=watermarking, |
| 253 | watermarking_skip_mask=watermarking_skip_mask, |
| 254 | watermark_key=_resolve_watermark_key(watermarker), |
| 255 | ) |