| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.triton_utils import HAS_TRITON, tl, triton |
| 7 | from vllm.v1.watermarking.prfs.philox import ( |
| 8 | _CONTEXT_DOMAIN as _CONTEXT_DOMAIN_VALUE, |
| 9 | ) |
| 10 | from vllm.v1.watermarking.prfs.philox import ( |
| 11 | _PHILOX_M0 as _PHILOX_M0_VALUE, |
| 12 | ) |
| 13 | from vllm.v1.watermarking.prfs.philox import ( |
| 14 | _PHILOX_M1 as _PHILOX_M1_VALUE, |
| 15 | ) |
| 16 | from vllm.v1.watermarking.prfs.philox import ( |
| 17 | _PHILOX_W0 as _PHILOX_W0_VALUE, |
| 18 | ) |
| 19 | from vllm.v1.watermarking.prfs.philox import ( |
| 20 | _PHILOX_W1 as _PHILOX_W1_VALUE, |
| 21 | ) |
| 22 | from vllm.v1.watermarking.prfs.philox import ( |
| 23 | _TOKEN_DOMAIN as _TOKEN_DOMAIN_VALUE, |
| 24 | ) |
| 25 | from vllm.v1.watermarking.prfs.philox import ( |
| 26 | _UINT32_MASK as _UINT32_MASK_VALUE, |
| 27 | ) |
| 28 | from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax |
| 29 | |
| 30 | if HAS_TRITON: |
| 31 | from triton.language import math as tl_math |
| 32 | else: |
| 33 | tl_math = tl |
| 34 | |
| 35 | _UINT32_MASK = tl.constexpr(_UINT32_MASK_VALUE) if HAS_TRITON else _UINT32_MASK_VALUE |
| 36 | _PHILOX_M0 = tl.constexpr(_PHILOX_M0_VALUE) if HAS_TRITON else _PHILOX_M0_VALUE |
| 37 | _PHILOX_M1 = tl.constexpr(_PHILOX_M1_VALUE) if HAS_TRITON else _PHILOX_M1_VALUE |
| 38 | _PHILOX_W0 = tl.constexpr(_PHILOX_W0_VALUE) if HAS_TRITON else _PHILOX_W0_VALUE |
| 39 | _PHILOX_W1 = tl.constexpr(_PHILOX_W1_VALUE) if HAS_TRITON else _PHILOX_W1_VALUE |
| 40 | _CONTEXT_DOMAIN = ( |
| 41 | tl.constexpr(_CONTEXT_DOMAIN_VALUE) if HAS_TRITON else _CONTEXT_DOMAIN_VALUE |
| 42 | ) |
| 43 | _TOKEN_DOMAIN = tl.constexpr(_TOKEN_DOMAIN_VALUE) if HAS_TRITON else _TOKEN_DOMAIN_VALUE |
| 44 | |
| 45 | |
| 46 | @triton.jit |
| 47 | def _repeated_context_mask_kernel( |
| 48 | output_ptr, |
| 49 | all_token_ids_ptr, |
| 50 | all_token_ids_stride, |
| 51 | req_indices_ptr, |
| 52 | prompt_lens_ptr, |
| 53 | total_lens_ptr, |
| 54 | history_offsets_ptr, |
| 55 | history_offsets_stride, |
| 56 | contexts_ptr, |
| 57 | context_stride, |
| 58 | local_positions_ptr, |
| 59 | prior_contexts_ptr, |
| 60 | prior_contexts_stride_0, |
| 61 | prior_contexts_stride_1, |
| 62 | steps_ptr, |
| 63 | steps_stride, |
| 64 | enabled_ptr, |
| 65 | CONTEXT_WIDTH: tl.constexpr, |
| 66 | NUM_SPECULATIVE_STEPS: tl.constexpr, |
| 67 | MAX_HISTORY: tl.constexpr, |
| 68 | INCLUDE_PROMPT: tl.constexpr, |
| 69 | SKIP_PARTIAL_CONTEXT: tl.constexpr, |
| 70 | HAS_HISTORY_OFFSETS: tl.constexpr, |
| 71 | SPECULATIVE_CONTEXTS: tl.constexpr, |
| 72 | DRAFT_CONTEXTS: tl.constexpr, |
| 73 | BLOCK: tl.constexpr, |
| 74 | ): |
| 75 | row = tl.program_id(0).to(tl.int64) |
| 76 | req_idx = tl.load(req_indices_ptr + row).to(tl.int64) |
| 77 | if req_idx < 0: |
| 78 | tl.store(output_ptr + row, 0) |
| 79 | return |
| 80 | prompt_len = tl.load(prompt_lens_ptr + req_idx) |
| 81 | total_len = tl.load(total_lens_ptr + req_idx) |
| 82 | if SKIP_PARTIAL_CONTEXT and total_len - prompt_len < CONTEXT_WIDTH: |
| 83 | tl.store(output_ptr + row, 1) |
| 84 | return |
| 85 | sequence_start = 0 if INCLUDE_PROMPT else prompt_len |
| 86 | history_len = total_len - sequence_start |
| 87 | history_offset = 0 |
| 88 | if HAS_HISTORY_OFFSETS: |
| 89 | history_offset = tl.load(history_offsets_ptr + row * history_offsets_stride) |
| 90 | |
| 91 | offsets = tl.arange(0, BLOCK) |
| 92 | repeated = tl.full((), 0, tl.int32) |
| 93 | scan_start = 0 |
| 94 | if MAX_HISTORY > 0: |
| 95 | remaining_history = tl.maximum(MAX_HISTORY - history_offset, 0) |
| 96 | scan_start = tl.maximum(history_len - remaining_history, 0) |
| 97 | # Aligned loop bounds let the per-lane offsets fold into immediates. |
| 98 | aligned_start = (scan_start // BLOCK) * BLOCK |
| 99 | for block_start in tl.range(aligned_start, history_len, BLOCK): |
| 100 | previous_pos = block_start + offsets |
| 101 | in_window = (previous_pos >= scan_start) & (previous_pos < history_len) |
| 102 | matches = in_window |
| 103 | for offset in range(CONTEXT_WIDTH): |
| 104 | context_token = tl.load(contexts_ptr + row * context_stride + offset) |
| 105 | historical_pos = previous_pos + offset - CONTEXT_WIDTH |
| 106 | historical_token = tl.load( |
| 107 | all_token_ids_ptr |
| 108 | + req_idx * all_token_ids_stride |
| 109 | + sequence_start |
| 110 | + tl.maximum(historical_pos, 0), |
| 111 | mask=in_window & (historical_pos >= 0), |
| 112 | other=-1, |
| 113 | ) |
| 114 | matches &= historical_token == context_token |
| 115 | repeated |= tl.max(matches.to(tl.int32), axis=0) |
| 116 | |
| 117 | if SPECULATIVE_CONTEXTS: |
| 118 | local_pos = tl.load(local_positions_ptr + row) |
| 119 | if INCLUDE_PROMPT: |
| 120 | partial_context = tl.full((), 0, tl.int32) |
| 121 | for context_offset in range(CONTEXT_WIDTH): |
| 122 | context_token = tl.load( |
| 123 | contexts_ptr + row * context_stride + context_offset |
| 124 | ) |
| 125 | partial_context |= context_token < 0 |
| 126 | repeated |= partial_context |
| 127 | for offset in range(1, NUM_SPECULATIVE_STEPS + 1): |
| 128 | prior_row = row - offset |
| 129 | in_history = (prior_row >= 0) & (local_pos >= offset) |
| 130 | if MAX_HISTORY > 0: |
| 131 | in_history &= offset <= MAX_HISTORY |
| 132 | prior_req_idx = tl.load( |
| 133 | req_indices_ptr + tl.maximum(prior_row, 0), |
| 134 | mask=in_history, |
| 135 | other=-1, |
| 136 | ) |
| 137 | prior_local_pos = tl.load( |
| 138 | local_positions_ptr + tl.maximum(prior_row, 0), |
| 139 | mask=in_history, |
| 140 | other=-1, |
| 141 | ) |
| 142 | matches = ( |
| 143 | in_history |
| 144 | & (req_idx == prior_req_idx) |
| 145 | & (local_pos == prior_local_pos + offset) |
| 146 | ) |
| 147 | for context_offset in range(CONTEXT_WIDTH): |
| 148 | context_token = tl.load( |
| 149 | contexts_ptr + row * context_stride + context_offset |
| 150 | ) |
| 151 | prior_token = tl.load( |
| 152 | contexts_ptr |
| 153 | + tl.maximum(prior_row, 0) * context_stride |
| 154 | + context_offset, |
| 155 | mask=in_history, |
| 156 | other=-2, |
| 157 | ) |
| 158 | matches &= context_token == prior_token |
| 159 | repeated |= matches |
| 160 | |
| 161 | if DRAFT_CONTEXTS: |
| 162 | step = tl.load(steps_ptr + row * steps_stride).to(tl.int64) |
| 163 | partial_context = tl.full((), 0, tl.int32) |
| 164 | for context_offset in range(CONTEXT_WIDTH): |
| 165 | context_token = tl.load( |
| 166 | contexts_ptr + row * context_stride + context_offset |
| 167 | ) |
| 168 | partial_context |= context_token < 0 |
| 169 | if INCLUDE_PROMPT: |
| 170 | repeated |= partial_context |
| 171 | |
| 172 | for prior_step in range(NUM_SPECULATIVE_STEPS): |
| 173 | in_history = prior_step < step |
| 174 | if MAX_HISTORY > 0: |
| 175 | in_history &= prior_step >= step - MAX_HISTORY |
| 176 | matches = in_history |
| 177 | for context_offset in range(CONTEXT_WIDTH): |
| 178 | context_token = tl.load( |
| 179 | contexts_ptr + row * context_stride + context_offset |
| 180 | ) |
| 181 | prior_token = tl.load( |
| 182 | prior_contexts_ptr |
| 183 | + row * prior_contexts_stride_0 |
| 184 | + prior_step * prior_contexts_stride_1 |
| 185 | + context_offset |
| 186 | ) |
| 187 | matches &= context_token == prior_token |
| 188 | repeated |= matches |
| 189 | |
| 190 | for context_offset in range(CONTEXT_WIDTH): |
| 191 | context_token = tl.load( |
| 192 | contexts_ptr + row * context_stride + context_offset |
| 193 | ) |
| 194 | tl.store( |
| 195 | prior_contexts_ptr |
| 196 | + row * prior_contexts_stride_0 |
| 197 | + step * prior_contexts_stride_1 |
| 198 | + context_offset, |
| 199 | context_token, |
| 200 | ) |
| 201 | enabled = tl.load(enabled_ptr + row) |
| 202 | repeated = enabled & (repeated == 0) |
| 203 | tl.store(output_ptr + row, repeated) |
| 204 | |
| 205 | |
| 206 | def _repeated_context_mask_cpu( |
| 207 | all_token_ids: torch.Tensor, |
| 208 | req_indices: torch.Tensor, |
| 209 | prompt_lens: torch.Tensor, |
| 210 | total_lens: torch.Tensor, |
| 211 | contexts: torch.Tensor, |
| 212 | max_history: int | None = None, |
| 213 | include_prompt: bool = False, |
| 214 | skip_partial_context: bool = False, |
| 215 | history_offsets: torch.Tensor | None = None, |
| 216 | ) -> torch.Tensor: |
| 217 | """Reference implementation; the parity tests check the Triton kernel against it.""" |
| 218 | repeated = torch.zeros(len(req_indices), dtype=torch.bool) |
| 219 | for row, req_idx_tensor in enumerate(req_indices): |
| 220 | req_idx = int(req_idx_tensor) |
| 221 | if req_idx < 0: |
| 222 | continue |
| 223 | prompt_len = int(prompt_lens[req_idx]) |
| 224 | total_len = int(total_lens[req_idx]) |
| 225 | if skip_partial_context and total_len - prompt_len < contexts.shape[-1]: |
| 226 | repeated[row] = True |
| 227 | continue |
| 228 | sequence_start = 0 if include_prompt else prompt_len |
| 229 | history_tokens = all_token_ids[req_idx, sequence_start:total_len].tolist() |
| 230 | prefix = [-1] * contexts.shape[-1] |
| 231 | current_context = tuple(contexts[row].tolist()) |
| 232 | history_offset = 0 if history_offsets is None else int(history_offsets[row]) |
| 233 | history_start = ( |
| 234 | max(0, len(history_tokens) - max(0, max_history - history_offset)) |
| 235 | if max_history is not None |
| 236 | else 0 |
| 237 | ) |
| 238 | for history_pos, token_id in enumerate(history_tokens): |
| 239 | if history_pos >= history_start and ( |
| 240 | tuple(prefix[-contexts.shape[-1] :]) == current_context |
| 241 | ): |
| 242 | repeated[row] = True |
| 243 | break |
| 244 | prefix.append(token_id) |
| 245 | return repeated |
| 246 | |
| 247 | |
| 248 | def repeated_context_mask( |
| 249 | all_token_ids: torch.Tensor, |
| 250 | req_indices: torch.Tensor, |
| 251 | prompt_lens: torch.Tensor, |
| 252 | total_lens: torch.Tensor, |
| 253 | contexts: torch.Tensor, |
| 254 | max_history: int | None = None, |
| 255 | include_prompt: bool = False, |
| 256 | skip_partial_context: bool = False, |
| 257 | history_offsets: torch.Tensor | None = None, |
| 258 | local_positions: torch.Tensor | None = None, |
| 259 | num_speculative_steps: int = 0, |
| 260 | ) -> torch.Tensor: |
| 261 | """Return, per row, whether the row's context already occurred in its history. |
| 262 | |
| 263 | Args: |
| 264 | all_token_ids: `[max_num_reqs, max_model_len]` token ids of every request. |
| 265 | req_indices: Request slot per sampled row; -1 marks a padding row, which |
| 266 | is reported as not repeated. |
| 267 | prompt_lens: Prompt length per request slot. |
| 268 | total_lens: Prompt plus generated length per request slot. |
| 269 | contexts: `[num_rows, context_width]` context per row, padded with -1 |
| 270 | before the start of the scanned history. |
| 271 | max_history: Number of most recent history positions searched, or |
| 272 | ``None`` for all of them. The window compared at each position |
| 273 | reaches `context_width` tokens further back. |
| 274 | include_prompt: Search the prompt as well as the generated tokens. |
| 275 | skip_partial_context: Mark contexts containing start padding so they use |
| 276 | ordinary sampling. |
| 277 | history_offsets: Number of newer, non-committed positions preceding each |
| 278 | context. These positions count toward `max_history`. Must be 1-D |
| 279 | over rows; a stride-0 broadcast row is fine. |
| 280 | local_positions: Position within each request's speculative block. When |
| 281 | provided, compare each context with earlier contexts in that block. |
| 282 | num_speculative_steps: Maximum number of draft tokens in the block. |
| 283 | |
| 284 | """ |
| 285 | if max_history is not None and max_history < 1: |
| 286 | raise ValueError("max_history must be positive or None") |
| 287 | if (local_positions is None) != (num_speculative_steps == 0): |
| 288 | raise ValueError( |
| 289 | "local_positions and positive num_speculative_steps must be " |
| 290 | "provided together" |
| 291 | ) |
| 292 | if all_token_ids.device.type == "cpu": |
| 293 | repeated = _repeated_context_mask_cpu( |
| 294 | all_token_ids, |
| 295 | req_indices, |
| 296 | prompt_lens, |
| 297 | total_lens, |
| 298 | contexts, |
| 299 | max_history, |
| 300 | include_prompt, |
| 301 | skip_partial_context, |
| 302 | history_offsets, |
| 303 | ) |
| 304 | if local_positions is not None: |
| 305 | max_offset = num_speculative_steps + 1 |
| 306 | if max_history is not None: |
| 307 | max_offset = min(max_offset, max_history + 1) |
| 308 | for offset in range(1, max_offset): |
| 309 | prior_contexts = torch.cat( |
| 310 | (contexts[:offset], contexts[:-offset]), dim=0 |
| 311 | ) |
| 312 | prior_requests = torch.cat( |
| 313 | (req_indices[:offset], req_indices[:-offset]), dim=0 |
| 314 | ) |
| 315 | prior_local_pos = torch.cat( |
| 316 | (local_positions[:offset], local_positions[:-offset]), dim=0 |
| 317 | ) |
| 318 | repeated |= ( |
| 319 | (req_indices >= 0) |
| 320 | & (local_positions >= offset) |
| 321 | & (req_indices == prior_requests) |
| 322 | & (local_positions == prior_local_pos + offset) |
| 323 | & (contexts == prior_contexts).all(dim=-1) |
| 324 | ) |
| 325 | if include_prompt: |
| 326 | repeated |= (contexts < 0).any(dim=-1) |
| 327 | return repeated |
| 328 | |
| 329 | if contexts.stride(-1) != 1: |
| 330 | contexts = contexts.contiguous() |
| 331 | # Only history_offsets passes its row stride to the kernel. |
| 332 | req_indices = req_indices.contiguous() |
| 333 | if local_positions is not None: |
| 334 | local_positions = local_positions.contiguous() |
| 335 | # `prompt_lens` and `total_lens` are indexed flat by request slot, and |
| 336 | # `all_token_ids` takes a row stride but is read flat within the row. Those |
| 337 | # three are contiguous by contract; every caller owns them outright. |
| 338 | repeated = torch.empty(len(req_indices), dtype=torch.bool, device=contexts.device) |
| 339 | _repeated_context_mask_kernel[(len(req_indices),)]( |
| 340 | repeated, |
| 341 | all_token_ids, |
| 342 | all_token_ids.stride(0), |
| 343 | req_indices, |
| 344 | prompt_lens, |
| 345 | total_lens, |
| 346 | history_offsets, |
| 347 | 0 if history_offsets is None else history_offsets.stride(0), |
| 348 | contexts, |
| 349 | contexts.stride(0), |
| 350 | local_positions, |
| 351 | None, |
| 352 | 0, |
| 353 | 0, |
| 354 | None, |
| 355 | 0, |
| 356 | None, |
| 357 | CONTEXT_WIDTH=contexts.shape[-1], |
| 358 | NUM_SPECULATIVE_STEPS=num_speculative_steps, |
| 359 | MAX_HISTORY=0 if max_history is None else max_history, |
| 360 | INCLUDE_PROMPT=include_prompt, |
| 361 | SKIP_PARTIAL_CONTEXT=skip_partial_context, |
| 362 | HAS_HISTORY_OFFSETS=history_offsets is not None, |
| 363 | SPECULATIVE_CONTEXTS=local_positions is not None, |
| 364 | DRAFT_CONTEXTS=False, |
| 365 | BLOCK=512, |
| 366 | ) |
| 367 | return repeated |
| 368 | |
| 369 | |
| 370 | def draft_watermarking_mask( |
| 371 | all_token_ids: torch.Tensor, |
| 372 | req_indices: torch.Tensor, |
| 373 | prompt_lens: torch.Tensor, |
| 374 | total_lens: torch.Tensor, |
| 375 | prior_contexts: torch.Tensor, |
| 376 | contexts: torch.Tensor, |
| 377 | steps: torch.Tensor, |
| 378 | enabled: torch.Tensor, |
| 379 | max_history: int | None, |
| 380 | include_prompt: bool, |
| 381 | ) -> torch.Tensor: |
| 382 | if contexts.device.type == "cpu": |
| 383 | repeated = repeated_context_mask( |
| 384 | all_token_ids, |
| 385 | req_indices, |
| 386 | prompt_lens, |
| 387 | total_lens, |
| 388 | contexts, |
| 389 | max_history, |
| 390 | include_prompt=include_prompt, |
| 391 | history_offsets=steps, |
| 392 | ) |
| 393 | prior_steps = torch.arange(prior_contexts.shape[1]).unsqueeze(0) |
| 394 | in_history = prior_steps < steps.unsqueeze(1) |
| 395 | if max_history is not None: |
| 396 | in_history &= prior_steps >= steps.unsqueeze(1) - max_history |
| 397 | repeated |= ( |
| 398 | (prior_contexts[: len(contexts)] == contexts.unsqueeze(1)).all(dim=-1) |
| 399 | & in_history |
| 400 | ).any(dim=-1) |
| 401 | if include_prompt: |
| 402 | repeated |= (contexts < 0).any(dim=-1) |
| 403 | row_indices = torch.arange(len(contexts)) |
| 404 | prior_contexts.index_put_((row_indices, steps), contexts) |
| 405 | return enabled & ~repeated |
| 406 | |
| 407 | if contexts.stride(-1) != 1: |
| 408 | contexts = contexts.contiguous() |
| 409 | # Only `steps` passes its row stride to the kernel. |
| 410 | req_indices = req_indices.contiguous() |
| 411 | enabled = enabled.contiguous() |
| 412 | active = torch.empty(len(contexts), dtype=torch.bool, device=contexts.device) |
| 413 | _repeated_context_mask_kernel[(len(contexts),)]( |
| 414 | active, |
| 415 | all_token_ids, |
| 416 | all_token_ids.stride(0), |
| 417 | req_indices, |
| 418 | prompt_lens, |
| 419 | total_lens, |
| 420 | steps, |
| 421 | steps.stride(0), |
| 422 | contexts, |
| 423 | contexts.stride(0), |
| 424 | None, |
| 425 | prior_contexts, |
| 426 | prior_contexts.stride(0), |
| 427 | prior_contexts.stride(1), |
| 428 | steps, |
| 429 | steps.stride(0), |
| 430 | enabled, |
| 431 | NUM_SPECULATIVE_STEPS=prior_contexts.shape[1], |
| 432 | CONTEXT_WIDTH=contexts.shape[1], |
| 433 | MAX_HISTORY=0 if max_history is None else max_history, |
| 434 | INCLUDE_PROMPT=include_prompt, |
| 435 | SKIP_PARTIAL_CONTEXT=False, |
| 436 | HAS_HISTORY_OFFSETS=True, |
| 437 | SPECULATIVE_CONTEXTS=False, |
| 438 | DRAFT_CONTEXTS=True, |
| 439 | BLOCK=512, |
| 440 | ) |
| 441 | return active |
| 442 | |
| 443 | |
| 444 | @triton.jit |
| 445 | def _mulhilo32(multiplier: tl.constexpr, value): |
| 446 | high = tl_math.umulhi(multiplier, value) |
| 447 | low = tl.mul(multiplier, value, sanitize_overflow=False) |
| 448 | return high, low |
| 449 | |
| 450 | |
| 451 | @triton.jit |
| 452 | def _philox4x32_10(counter_0, counter_1, counter_2, counter_3, key_0, key_1): |
| 453 | for round_index in range(10): |
| 454 | high_0, low_0 = _mulhilo32(_PHILOX_M0, counter_0) |
| 455 | high_1, low_1 = _mulhilo32(_PHILOX_M1, counter_2) |
| 456 | counter_0, counter_1, counter_2, counter_3 = ( |
| 457 | (high_1 ^ counter_1 ^ key_0) & _UINT32_MASK, |
| 458 | low_1, |
| 459 | (high_0 ^ counter_3 ^ key_1) & _UINT32_MASK, |
| 460 | low_0, |
| 461 | ) |
| 462 | if round_index != 9: |
| 463 | key_0 = (key_0 + _PHILOX_W0) & _UINT32_MASK |
| 464 | key_1 = (key_1 + _PHILOX_W1) & _UINT32_MASK |
| 465 | return counter_0, counter_1, counter_2, counter_3 |
| 466 | |
| 467 | |
| 468 | @triton.jit |
| 469 | def _uint32_to_uniform(value): |
| 470 | mantissa = (value & _UINT32_MASK) >> 8 |
| 471 | scaled = (mantissa + 1).to(tl.float32) * 2**-24 |
| 472 | return (scaled.to(tl.uint32, bitcast=True) - 1).to(tl.float32, bitcast=True) |
| 473 | |
| 474 | |
| 475 | @triton.jit |
| 476 | def _philox_context_state( |
| 477 | contexts_row_ptr, |
| 478 | key_0, |
| 479 | key_1, |
| 480 | CONTEXT_WIDTH: tl.constexpr, |
| 481 | ): |
| 482 | state_0 = tl.full((), _CONTEXT_DOMAIN, tl.uint32) |
| 483 | state_1 = tl.full((), CONTEXT_WIDTH, tl.uint32) |
| 484 | state_2 = tl.full((), 0, tl.uint32) |
| 485 | state_3 = tl.full((), 0, tl.uint32) |
| 486 | for offset in range(0, CONTEXT_WIDTH, 4): |
| 487 | context_0 = tl.load(contexts_row_ptr + offset).to(tl.uint32) |
| 488 | if offset + 1 < CONTEXT_WIDTH: |
| 489 | context_1 = tl.load(contexts_row_ptr + offset + 1).to(tl.uint32) |
| 490 | else: |
| 491 | context_1 = tl.full((), _UINT32_MASK - 1, tl.uint32) |
| 492 | if offset + 2 < CONTEXT_WIDTH: |
| 493 | context_2 = tl.load(contexts_row_ptr + offset + 2).to(tl.uint32) |
| 494 | else: |
| 495 | context_2 = tl.full((), _UINT32_MASK - 2, tl.uint32) |
| 496 | if offset + 3 < CONTEXT_WIDTH: |
| 497 | context_3 = tl.load(contexts_row_ptr + offset + 3).to(tl.uint32) |
| 498 | else: |
| 499 | context_3 = tl.full((), _UINT32_MASK - 3, tl.uint32) |
| 500 | state_0, state_1, state_2, state_3 = _philox4x32_10( |
| 501 | (state_0 ^ context_0) & _UINT32_MASK, |
| 502 | (state_1 ^ context_1) & _UINT32_MASK, |
| 503 | (state_2 ^ context_2) & _UINT32_MASK, |
| 504 | (state_3 ^ context_3) & _UINT32_MASK, |
| 505 | (key_0 ^ offset) & _UINT32_MASK, |
| 506 | key_1, |
| 507 | ) |
| 508 | return state_0, state_1, state_2, state_3 |
| 509 | |
| 510 | |
| 511 | @triton.jit |
| 512 | def _philox_candidate_words(groups, state_0, state_1, state_2, state_3, key_0, key_1): |
| 513 | candidate_words = groups.to(tl.uint32) |
| 514 | vector_zero = candidate_words * 0 |
| 515 | return _philox4x32_10( |
| 516 | candidate_words & _UINT32_MASK, |
| 517 | state_0 + vector_zero, |
| 518 | state_1 + vector_zero, |
| 519 | state_2 + vector_zero, |
| 520 | ((key_0 ^ state_3) & _UINT32_MASK) + vector_zero, |
| 521 | ((key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK) + vector_zero, |
| 522 | ) |
| 523 | |
| 524 | |
| 525 | @triton.jit |
| 526 | def _philox_gumbel_from_logits(logits, word): |
| 527 | logits = tl.where(logits != logits, float("-inf"), logits) |
| 528 | uniform = _uint32_to_uniform(word) |
| 529 | return logits - tl.log(-tl.log(uniform)) |
| 530 | |
| 531 | |
| 532 | @triton.jit |
| 533 | def _gumbel_value(logits_ptr, output, mask): |
| 534 | logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32) |
| 535 | return _philox_gumbel_from_logits(logits, output) |
| 536 | |
| 537 | |
| 538 | @triton.jit |
| 539 | def philox_gumbel_block_argmax( |
| 540 | logits, |
| 541 | mask, |
| 542 | block_idx, |
| 543 | contexts_row_ptr, |
| 544 | key_0, |
| 545 | key_1, |
| 546 | CONTEXT_WIDTH: tl.constexpr, |
| 547 | BLOCK_SIZE: tl.constexpr, |
| 548 | ): |
| 549 | """Apply the detector-compatible PRF to one vocabulary block.""" |
| 550 | tl.static_assert(BLOCK_SIZE % 4 == 0) |
| 551 | state_0, state_1, state_2, state_3 = _philox_context_state( |
| 552 | contexts_row_ptr, key_0, key_1, CONTEXT_WIDTH |
| 553 | ) |
| 554 | groups = block_idx * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4) |
| 555 | output_0, output_1, output_2, output_3 = _philox_candidate_words( |
| 556 | groups, state_0, state_1, state_2, state_3, key_0, key_1 |
| 557 | ) |
| 558 | # Lay the per-group words out in token order: 4g+0, 4g+1, 4g+2, 4g+3. |
| 559 | words = tl.interleave( |
| 560 | tl.interleave(output_0, output_2), |
| 561 | tl.interleave(output_1, output_3), |
| 562 | ) |
| 563 | values = tl.where( |
| 564 | mask, |
| 565 | _philox_gumbel_from_logits(logits.to(tl.float32), words), |
| 566 | float("-inf"), |
| 567 | ) |
| 568 | return tl.max(values, axis=0, return_indices=True) |
| 569 | |
| 570 | |
| 571 | @triton.jit |
| 572 | def _select_max(best_value, best_token, candidate_value, candidate): |
| 573 | take_candidate = candidate_value > best_value |
| 574 | return ( |
| 575 | tl.where(take_candidate, candidate_value, best_value), |
| 576 | tl.where(take_candidate, candidate, best_token), |
| 577 | ) |
| 578 | |
| 579 | |
| 580 | @triton.jit(do_not_specialize=["key_0_value", "key_1_value"]) |
| 581 | def _philox_gumbel_kernel( |
| 582 | local_argmax_ptr, |
| 583 | local_argmax_stride, |
| 584 | local_max_ptr, |
| 585 | local_max_stride, |
| 586 | logits_ptr, |
| 587 | logits_stride, |
| 588 | contexts_ptr, |
| 589 | context_stride, |
| 590 | skip_mask_ptr, |
| 591 | expanded_idx_mapping_ptr, |
| 592 | seeds_ptr, |
| 593 | pos_ptr, |
| 594 | temp_ptr, |
| 595 | logits_cache_ptr, |
| 596 | logits_cache_stride_0, |
| 597 | logits_cache_stride_1, |
| 598 | logits_cache_col_ptr, |
| 599 | logits_cache_source_ptr, |
| 600 | logits_cache_source_stride, |
| 601 | key_0_value, |
| 602 | key_1_value, |
| 603 | vocab_size, |
| 604 | CONTEXT_WIDTH: tl.constexpr, |
| 605 | BLOCK_SIZE: tl.constexpr, |
| 606 | USE_FP64: tl.constexpr, |
| 607 | IS_DRAFTING: tl.constexpr, |
| 608 | PER_TOKEN_COL: tl.constexpr, |
| 609 | ): |
| 610 | row = tl.program_id(0).to(tl.int64) |
| 611 | block_index = tl.program_id(1) |
| 612 | groups = block_index * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4) |
| 613 | candidate = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 614 | mask = candidate < vocab_size |
| 615 | |
| 616 | req_state_idx = row |
| 617 | valid_req = True |
| 618 | if expanded_idx_mapping_ptr is not None: |
| 619 | req_state_idx = tl.load(expanded_idx_mapping_ptr + row).to(tl.int64) |
| 620 | valid_req = req_state_idx >= 0 |
| 621 | |
| 622 | if logits_cache_ptr is not None: |
| 623 | if PER_TOKEN_COL: |
| 624 | col = tl.load(logits_cache_col_ptr + row) |
| 625 | else: |
| 626 | col = tl.load(logits_cache_col_ptr) |
| 627 | cached_logits = tl.load( |
| 628 | logits_cache_source_ptr + row * logits_cache_source_stride + candidate, |
| 629 | mask=mask, |
| 630 | ) |
| 631 | tl.store( |
| 632 | logits_cache_ptr |
| 633 | + req_state_idx * logits_cache_stride_0 |
| 634 | + col * logits_cache_stride_1 |
| 635 | + candidate, |
| 636 | cached_logits, |
| 637 | mask=mask & valid_req, |
| 638 | ) |
| 639 | |
| 640 | if skip_mask_ptr is not None: # noqa: SIM102 |
| 641 | if tl.load(skip_mask_ptr + row): |
| 642 | temp = tl.load(temp_ptr + req_state_idx, mask=valid_req, other=0.0).to( |
| 643 | tl.float32 |
| 644 | ) |
| 645 | seed = tl.load(seeds_ptr + req_state_idx, mask=valid_req, other=0) |
| 646 | pos = tl.load(pos_ptr + row) |
| 647 | logits_row = logits_ptr + row * logits_stride |
| 648 | logits = tl.load(logits_row + candidate, mask=mask, other=float("-inf")).to( |
| 649 | tl.float32 |
| 650 | ) |
| 651 | value, index = gumbel_noised_argmax( |
| 652 | logits, |
| 653 | candidate, |
| 654 | mask, |
| 655 | seed, |
| 656 | pos, |
| 657 | temp, |
| 658 | IS_DRAFTING=IS_DRAFTING, |
| 659 | USE_FP64=USE_FP64, |
| 660 | APPLY_TEMPERATURE=False, |
| 661 | ) |
| 662 | token_id = block_index * BLOCK_SIZE + index |
| 663 | tl.store( |
| 664 | local_argmax_ptr + row * local_argmax_stride + block_index, |
| 665 | token_id, |
| 666 | ) |
| 667 | tl.store(local_max_ptr + row * local_max_stride + block_index, value) |
| 668 | return |
| 669 | |
| 670 | key_0 = key_0_value.to(tl.uint32) |
| 671 | key_1 = key_1_value.to(tl.uint32) |
| 672 | state_0, state_1, state_2, state_3 = _philox_context_state( |
| 673 | contexts_ptr + row * context_stride, key_0, key_1, CONTEXT_WIDTH |
| 674 | ) |
| 675 | output_0, output_1, output_2, output_3 = _philox_candidate_words( |
| 676 | groups, state_0, state_1, state_2, state_3, key_0, key_1 |
| 677 | ) |
| 678 | candidate_0 = groups * 4 |
| 679 | candidate_1 = candidate_0 + 1 |
| 680 | candidate_2 = candidate_0 + 2 |
| 681 | candidate_3 = candidate_0 + 3 |
| 682 | logits_row = logits_ptr + row * logits_stride |
| 683 | value_0 = _gumbel_value( |
| 684 | logits_row + candidate_0, output_0, candidate_0 < vocab_size |
| 685 | ) |
| 686 | value_1 = _gumbel_value( |
| 687 | logits_row + candidate_1, output_1, candidate_1 < vocab_size |
| 688 | ) |
| 689 | value_2 = _gumbel_value( |
| 690 | logits_row + candidate_2, output_2, candidate_2 < vocab_size |
| 691 | ) |
| 692 | value_3 = _gumbel_value( |
| 693 | logits_row + candidate_3, output_3, candidate_3 < vocab_size |
| 694 | ) |
| 695 | best_value = value_0 |
| 696 | best_token = candidate_0 |
| 697 | best_value, best_token = _select_max(best_value, best_token, value_1, candidate_1) |
| 698 | best_value, best_token = _select_max(best_value, best_token, value_2, candidate_2) |
| 699 | best_value, best_token = _select_max(best_value, best_token, value_3, candidate_3) |
| 700 | value = tl.max(best_value, axis=0) |
| 701 | token_id = tl.min(tl.where(best_value == value, best_token, vocab_size), axis=0) |
| 702 | tl.store(local_argmax_ptr + row * local_argmax_stride + block_index, token_id) |
| 703 | tl.store(local_max_ptr + row * local_max_stride + block_index, value) |
| 704 | |
| 705 | |
| 706 | def philox_gumbel_sample( |
| 707 | logits: torch.Tensor, |
| 708 | contexts: torch.Tensor, |
| 709 | key: int, |
| 710 | *, |
| 711 | skip_mask: torch.Tensor | None = None, |
| 712 | expanded_idx_mapping: torch.Tensor | None = None, |
| 713 | temperatures: torch.Tensor | None = None, |
| 714 | seeds: torch.Tensor | None = None, |
| 715 | positions: torch.Tensor | None = None, |
| 716 | use_fp64: bool = False, |
| 717 | is_drafting: bool = False, |
| 718 | logits_cache: torch.Tensor | None = None, |
| 719 | logits_cache_col: torch.Tensor | None = None, |
| 720 | logits_cache_source: torch.Tensor | None = None, |
| 721 | ) -> torch.Tensor: |
| 722 | sampling_state = (expanded_idx_mapping, temperatures, seeds, positions) |
| 723 | if skip_mask is None and any(value is not None for value in sampling_state): |
| 724 | raise ValueError("sampling state requires skip_mask") |
| 725 | if skip_mask is not None and any(value is None for value in sampling_state): |
| 726 | raise ValueError("skip_mask requires complete sampling state") |
| 727 | if logits_cache is not None: |
| 728 | if logits_cache_col is None or logits_cache_source is None: |
| 729 | raise ValueError("logits cache requires its column and source logits") |
| 730 | if expanded_idx_mapping is None: |
| 731 | raise ValueError("logits cache requires expanded_idx_mapping") |
| 732 | assert logits_cache_source.shape == logits.shape, ( |
| 733 | "logits cache source must match sampled logits shape" |
| 734 | ) |
| 735 | assert logits_cache_source.device == logits.device, ( |
| 736 | "logits cache source must be on the sampled logits device" |
| 737 | ) |
| 738 | assert logits_cache_source.dtype == logits_cache.dtype, ( |
| 739 | "logits cache source and destination must have the same dtype" |
| 740 | ) |
| 741 | assert logits_cache.size(-1) >= logits.shape[-1], ( |
| 742 | f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower " |
| 743 | f"than the sampled logits ({logits.shape[-1]}). Cached logits would be " |
| 744 | "truncated." |
| 745 | ) |
| 746 | elif logits_cache_col is not None or logits_cache_source is not None: |
| 747 | raise ValueError("logits cache column and source require logits_cache") |
| 748 | |
| 749 | if logits.stride(-1) != 1: |
| 750 | logits = logits.contiguous() |
| 751 | if contexts.stride(-1) != 1: |
| 752 | contexts = contexts.contiguous() |
| 753 | if skip_mask is not None: |
| 754 | skip_mask = skip_mask.contiguous() |
| 755 | if expanded_idx_mapping is not None: |
| 756 | expanded_idx_mapping = expanded_idx_mapping.contiguous() |
| 757 | if positions is not None: |
| 758 | positions = positions.contiguous() |
| 759 | if logits_cache_col is not None: |
| 760 | logits_cache_col = logits_cache_col.contiguous() |
| 761 | if logits_cache_source is not None and logits_cache_source.stride(-1) != 1: |
| 762 | logits_cache_source = logits_cache_source.contiguous() |
| 763 | num_tokens, vocab_size = logits.shape |
| 764 | block_size = 1024 |
| 765 | num_blocks = triton.cdiv(vocab_size, block_size) |
| 766 | local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64) |
| 767 | local_max = logits.new_empty( |
| 768 | num_tokens, |
| 769 | num_blocks, |
| 770 | dtype=torch.float64 if use_fp64 else torch.float32, |
| 771 | ) |
| 772 | _philox_gumbel_kernel[(num_tokens, num_blocks)]( |
| 773 | local_argmax, |
| 774 | local_argmax.stride(0), |
| 775 | local_max, |
| 776 | local_max.stride(0), |
| 777 | logits, |
| 778 | logits.stride(0), |
| 779 | contexts, |
| 780 | contexts.stride(0), |
| 781 | skip_mask, |
| 782 | expanded_idx_mapping, |
| 783 | seeds, |
| 784 | positions, |
| 785 | temperatures, |
| 786 | logits_cache, |
| 787 | logits_cache.stride(0) if logits_cache is not None else 0, |
| 788 | logits_cache.stride(1) if logits_cache is not None else 0, |
| 789 | logits_cache_col, |
| 790 | logits_cache_source, |
| 791 | logits_cache_source.stride(0) if logits_cache_source is not None else 0, |
| 792 | key & _UINT32_MASK_VALUE, |
| 793 | key >> 32, |
| 794 | vocab_size, |
| 795 | CONTEXT_WIDTH=contexts.shape[-1], |
| 796 | BLOCK_SIZE=block_size, |
| 797 | USE_FP64=use_fp64, |
| 798 | IS_DRAFTING=is_drafting, |
| 799 | PER_TOKEN_COL=logits_cache_col is not None and logits_cache_col.dim() > 0, |
| 800 | ) |
| 801 | max_block_index = local_max.argmax(dim=-1, keepdim=True) |
| 802 | return local_argmax.gather(dim=-1, index=max_block_index).view(-1) |