| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | from collections.abc import Sequence |
| 7 | from dataclasses import replace |
| 8 | from typing import TYPE_CHECKING, Any |
| 9 | |
| 10 | import torch |
| 11 | import torch.nn as nn |
| 12 | |
| 13 | from vllm.config.model import PROCESSED_LOGPROBS_MODES |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.model_executor.warmup.jit_warmup import WarmupChoices |
| 16 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 17 | DispatchSpec, |
| 18 | TritonWarmupTensor, |
| 19 | triton_kernel_dispatcher_with_warmup, |
| 20 | ) |
| 21 | from vllm.triton_utils import tl, triton |
| 22 | from vllm.utils.torch_utils import async_tensor_h2d |
| 23 | from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput |
| 24 | from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor |
| 25 | from vllm.v1.sample.metadata import SamplingMetadata |
| 26 | from vllm.v1.sample.ops.bad_words import apply_bad_words_with_drafts |
| 27 | from vllm.v1.sample.ops.penalties import apply_all_penalties |
| 28 | from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p |
| 29 | from vllm.v1.sample.sampler import Sampler |
| 30 | from vllm.v1.spec_decode.metadata import SpecDecodeMetadata |
| 31 | from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates |
| 32 | |
| 33 | if TYPE_CHECKING: |
| 34 | from vllm.config.speculative import SpeculativeConfig |
| 35 | |
| 36 | logger = init_logger(__name__) |
| 37 | |
| 38 | PLACEHOLDER_TOKEN_ID: tl.constexpr = -1 |
| 39 | GREEDY_TEMPERATURE: tl.constexpr = 0 |
| 40 | # Maximum number of speculative draft tokens allowed per request in a single |
| 41 | # step. This value is chosen to be large enough to handle typical use cases. |
| 42 | MAX_SPEC_LEN = 128 |
| 43 | |
| 44 | |
| 45 | class RejectionSampler(nn.Module): |
| 46 | """The implementation strictly follows the algorithm described in |
| 47 | https://arxiv.org/abs/2211.17192. |
| 48 | However, we want to clarify the terminology used in the implementation: |
| 49 | accepted tokens: tokens that are accepted based on the relationship |
| 50 | between the "raw" draft and target probabilities. |
| 51 | recovered tokens: tokens that are sampled based on the adjusted probability |
| 52 | distribution, which is derived from both the draft and target |
| 53 | probabilities. |
| 54 | bonus tokens: |
| 55 | If all proposed tokens are accepted, the bonus token is added to the |
| 56 | end of the sequence. The bonus token is only sampled from the target |
| 57 | probabilities. We pass in the bonus tokens instead of sampling them |
| 58 | in the rejection sampler to allow for more flexibility in the |
| 59 | sampling process. For example, we can use top_p, top_k sampling for |
| 60 | bonus tokens, while spec decode does not support these sampling |
| 61 | strategies. |
| 62 | output tokens: |
| 63 | Tokens are finally generated with the rejection sampler. |
| 64 | output tokens = accepted tokens + recovered tokens + bonus tokens |
| 65 | """ |
| 66 | |
| 67 | def __init__( |
| 68 | self, |
| 69 | sampler: Sampler, |
| 70 | spec_config: SpeculativeConfig | None = None, |
| 71 | device: torch.device | None = None, |
| 72 | ): |
| 73 | super().__init__() |
| 74 | self.sampler = sampler |
| 75 | self.use_fp64_gumbel = getattr(sampler, "use_fp64_gumbel", False) |
| 76 | logprobs_mode = self.sampler.logprobs_mode |
| 77 | self.is_processed_logprobs_mode = logprobs_mode in PROCESSED_LOGPROBS_MODES |
| 78 | self.is_logits_logprobs_mode = logprobs_mode in ( |
| 79 | "raw_logits", |
| 80 | "processed_logits", |
| 81 | ) |
| 82 | |
| 83 | self.synthetic_conditional_rates: torch.Tensor | None = None |
| 84 | if ( |
| 85 | spec_config is not None |
| 86 | and spec_config.rejection_sample_method == "synthetic" |
| 87 | ): |
| 88 | assert spec_config.synthetic_acceptance_rates is not None |
| 89 | self.synthetic_conditional_rates = torch.tensor( |
| 90 | unconditional_to_conditional_rates( |
| 91 | spec_config.synthetic_acceptance_rates |
| 92 | ), |
| 93 | dtype=torch.float32, |
| 94 | device=device, |
| 95 | ) |
| 96 | self.synthetic_mode = self.synthetic_conditional_rates is not None |
| 97 | _rejection_greedy_sample.register_warmup(synthetic_mode=self.synthetic_mode) |
| 98 | _rejection_random_sample.register_warmup( |
| 99 | synthetic_mode=self.synthetic_mode, |
| 100 | ) |
| 101 | _expand.register_warmup() |
| 102 | _sample_recovered_tokens.register_warmup( |
| 103 | use_fp64_gumbel=self.use_fp64_gumbel, |
| 104 | ) |
| 105 | |
| 106 | def forward( |
| 107 | self, |
| 108 | metadata: SpecDecodeMetadata, |
| 109 | # [num_tokens, vocab_size] |
| 110 | draft_probs: torch.Tensor | None, |
| 111 | # [num_tokens + batch_size, vocab_size] |
| 112 | logits: torch.Tensor, |
| 113 | sampling_metadata: SamplingMetadata, |
| 114 | ) -> SamplerOutput: |
| 115 | """Args: |
| 116 | metadata: |
| 117 | Metadata for spec decoding. |
| 118 | draft_probs (Optional[torch.Tensor]): |
| 119 | Probability distribution for the draft tokens. Shape is |
| 120 | [num_tokens, vocab_size]. Can be None if probabilities are |
| 121 | not provided, which is the case for ngram spec decode. |
| 122 | logits (torch.Tensor): |
| 123 | Target model's logits probability distribution. |
| 124 | Shape is [num_tokens + batch_size, vocab_size]. Here, |
| 125 | probabilities from different requests are flattened into a |
| 126 | single tensor because this is the shape of the output logits. |
| 127 | NOTE: `logits` can be updated in place to save memory. |
| 128 | sampling_metadata (vllm.v1.sample.metadata.SamplingMetadata): |
| 129 | Additional metadata needed for sampling, such as temperature, |
| 130 | top-k/top-p parameters, or other relevant information. |
| 131 | |
| 132 | Returns: |
| 133 | SamplerOutput: |
| 134 | Contains the final output token IDs and their logprobs if |
| 135 | requested. |
| 136 | |
| 137 | """ |
| 138 | assert metadata.max_spec_len <= MAX_SPEC_LEN |
| 139 | |
| 140 | bonus_logits_indices = metadata.bonus_logits_indices |
| 141 | target_logits_indices = metadata.target_logits_indices |
| 142 | |
| 143 | # When indexing with a tensor (bonus_logits_indices), PyTorch |
| 144 | # creates a new tensor with separate storage from the original |
| 145 | # logits tensor. This means any in-place operations on bonus_logits |
| 146 | # won't affect the original logits tensor. |
| 147 | assert logits is not None |
| 148 | bonus_logits = logits[bonus_logits_indices] |
| 149 | bonus_sampler_output = self.sampler( |
| 150 | logits=bonus_logits, |
| 151 | sampling_metadata=replace( |
| 152 | sampling_metadata, |
| 153 | max_num_logprobs=-1, |
| 154 | ), |
| 155 | predict_bonus_token=True, |
| 156 | # Override the logprobs mode to return logits because they are |
| 157 | # needed later to compute the accepted token logprobs. |
| 158 | logprobs_mode_override="processed_logits" |
| 159 | if self.is_processed_logprobs_mode |
| 160 | else "raw_logits", |
| 161 | ) |
| 162 | bonus_token_ids = bonus_sampler_output.sampled_token_ids |
| 163 | |
| 164 | # Just like `bonus_logits`, `target_logits` is a new tensor with |
| 165 | # separate storage from the original `logits` tensor. Therefore, |
| 166 | # it is safe to update `target_logits` in place. |
| 167 | raw_target_logits = logits[target_logits_indices] |
| 168 | # Use float32 for the target_logits. |
| 169 | raw_target_logits = raw_target_logits.to(torch.float32) |
| 170 | target_logits = raw_target_logits |
| 171 | if not self.is_processed_logprobs_mode: |
| 172 | # Clone raw_target_logits before applying processors to preserve |
| 173 | # the original raw logits for logprobs computation, since |
| 174 | # apply_logits_processors modifies the tensor in-place. |
| 175 | target_logits = target_logits.clone() |
| 176 | target_logits = self.apply_logits_processors( |
| 177 | target_logits, sampling_metadata, metadata |
| 178 | ) |
| 179 | # [num_tokens, vocab_size] |
| 180 | # NOTE(woosuk): `target_logits` can be updated in place inside the |
| 181 | # `apply_sampling_constraints` function. |
| 182 | target_logits = apply_sampling_constraints( |
| 183 | target_logits, |
| 184 | metadata.cu_num_draft_tokens, |
| 185 | sampling_metadata, |
| 186 | ) |
| 187 | |
| 188 | output_token_ids = rejection_sample( |
| 189 | metadata.draft_token_ids, |
| 190 | metadata.num_draft_tokens, |
| 191 | metadata.max_spec_len, |
| 192 | metadata.cu_num_draft_tokens, |
| 193 | draft_probs, |
| 194 | target_logits, |
| 195 | bonus_token_ids, |
| 196 | sampling_metadata, |
| 197 | synthetic_mode=self.synthetic_mode, |
| 198 | synthetic_conditional_rates=self.synthetic_conditional_rates, |
| 199 | use_fp64_gumbel=self.use_fp64_gumbel, |
| 200 | ) |
| 201 | |
| 202 | logprobs_tensors = None |
| 203 | if sampling_metadata.max_num_logprobs is not None: |
| 204 | logprobs_tensors = self._get_logprobs_tensors( |
| 205 | sampling_metadata.max_num_logprobs, |
| 206 | metadata, |
| 207 | logits, |
| 208 | target_logits if self.is_processed_logprobs_mode else raw_target_logits, |
| 209 | bonus_sampler_output.logprobs_tensors.logprobs, |
| 210 | output_token_ids, |
| 211 | ) |
| 212 | |
| 213 | return SamplerOutput( |
| 214 | sampled_token_ids=output_token_ids, |
| 215 | logprobs_tensors=logprobs_tensors, |
| 216 | ) |
| 217 | |
| 218 | def _get_logprobs_tensors( |
| 219 | self, |
| 220 | max_num_logprobs: int, |
| 221 | metadata: SpecDecodeMetadata, |
| 222 | logits: torch.Tensor, |
| 223 | target_logits: torch.Tensor, |
| 224 | bonus_logits: torch.Tensor, |
| 225 | sampled_token_ids: torch.Tensor, |
| 226 | ) -> LogprobsTensors: |
| 227 | cu_num_sampled_tokens = torch.zeros_like(metadata.cu_num_sampled_tokens) |
| 228 | cu_num_sampled_tokens[1:] = metadata.cu_num_sampled_tokens[:-1] |
| 229 | |
| 230 | # Collect target and bonus logits. |
| 231 | bonus_logits_indices = metadata.bonus_logits_indices |
| 232 | target_logits_indices = metadata.target_logits_indices |
| 233 | final_logits = torch.zeros_like(logits, dtype=torch.float32) |
| 234 | final_logits[target_logits_indices] = target_logits.to(torch.float32) |
| 235 | final_logits[bonus_logits_indices] = bonus_logits.to(torch.float32) |
| 236 | |
| 237 | # NOTE: To avoid cpu-gpu synchronization, we now simply compute indices for |
| 238 | # all draft tokens, including the rejected ones. The rejected tokens will |
| 239 | # be filtered out in the `parse_output`. |
| 240 | logit_start_indices = cu_num_sampled_tokens |
| 241 | offsets = torch.arange( |
| 242 | sampled_token_ids.shape[-1], |
| 243 | device=logit_start_indices.device, |
| 244 | dtype=logit_start_indices.dtype, |
| 245 | ) |
| 246 | accepted_logit_indices = ( |
| 247 | logit_start_indices.unsqueeze(1) + offsets.unsqueeze(0) |
| 248 | ).flatten() |
| 249 | accepted_logit_indices.clamp_(max=final_logits.shape[0] - 1) |
| 250 | accepted_tokens = sampled_token_ids.clone().flatten() |
| 251 | # we replace rejected token ids with 0 to avoid gather_logprobs error |
| 252 | accepted_tokens[accepted_tokens == PLACEHOLDER_TOKEN_ID] = 0 |
| 253 | |
| 254 | # Compute logprobs for accepted tokens. |
| 255 | accepted_logits = final_logits[accepted_logit_indices] |
| 256 | accepted_logprobs = ( |
| 257 | accepted_logits |
| 258 | if self.is_logits_logprobs_mode |
| 259 | else self.sampler.compute_logprobs(accepted_logits) |
| 260 | ) |
| 261 | return self.sampler.gather_logprobs( |
| 262 | accepted_logprobs, |
| 263 | max_num_logprobs, |
| 264 | accepted_tokens.to(torch.int64), |
| 265 | ) |
| 266 | |
| 267 | @staticmethod |
| 268 | def parse_output( |
| 269 | output_token_ids: torch.Tensor, |
| 270 | vocab_size: int, |
| 271 | discard_req_indices: Sequence[int] = (), |
| 272 | logprobs_tensors: LogprobsTensors | None = None, |
| 273 | ) -> tuple[list[list[int]], LogprobsLists | None]: |
| 274 | """Parse the output of the rejection sampler. |
| 275 | |
| 276 | Args: |
| 277 | output_token_ids: The sampled token IDs in shape |
| 278 | [batch_size, max_spec_len + 1]. The rejected tokens are |
| 279 | replaced with `PLACEHOLDER_TOKEN_ID` by the rejection sampler |
| 280 | and will be filtered out in this function. |
| 281 | vocab_size: The size of the vocabulary. |
| 282 | discard_req_indices: Optional row indices to discard tokens in. |
| 283 | logprobs_tensors: Optional logprobs tensors to filter. |
| 284 | |
| 285 | Returns: |
| 286 | A list of lists of token IDs. |
| 287 | |
| 288 | """ |
| 289 | output_token_ids_np = output_token_ids.cpu().numpy() |
| 290 | # Create mask for valid tokens. |
| 291 | valid_mask = (output_token_ids_np != PLACEHOLDER_TOKEN_ID) & ( |
| 292 | output_token_ids_np < vocab_size |
| 293 | ) |
| 294 | output_logprobs = None |
| 295 | if logprobs_tensors is not None: |
| 296 | cu_num_tokens = [0] + valid_mask.sum(axis=1).cumsum().tolist() |
| 297 | filtered_tensors = logprobs_tensors.filter(valid_mask.flatten()) |
| 298 | output_logprobs = filtered_tensors.tolists(cu_num_tokens) |
| 299 | |
| 300 | if len(discard_req_indices) > 0: |
| 301 | valid_mask[discard_req_indices] = False |
| 302 | outputs = [ |
| 303 | row[valid_mask[i]].tolist() for i, row in enumerate(output_token_ids_np) |
| 304 | ] |
| 305 | return outputs, output_logprobs |
| 306 | |
| 307 | def apply_logits_processors( |
| 308 | self, |
| 309 | logits: torch.Tensor, |
| 310 | sampling_metadata: SamplingMetadata, |
| 311 | metadata: SpecDecodeMetadata, |
| 312 | ) -> torch.Tensor: |
| 313 | has_penalties = not sampling_metadata.no_penalties |
| 314 | any_penalties_or_bad_words = ( |
| 315 | sampling_metadata.bad_words_token_ids or has_penalties |
| 316 | ) |
| 317 | output_token_ids = sampling_metadata.output_token_ids |
| 318 | if any_penalties_or_bad_words: |
| 319 | output_token_ids = self._combine_outputs_with_spec_tokens( |
| 320 | output_token_ids, |
| 321 | sampling_metadata.spec_token_ids, |
| 322 | ) |
| 323 | |
| 324 | # Calculate indices of target logits. |
| 325 | repeat_indices: torch.Tensor | None = None |
| 326 | need_repeat_indices = ( |
| 327 | sampling_metadata.allowed_token_ids_mask is not None or has_penalties |
| 328 | ) |
| 329 | if need_repeat_indices: |
| 330 | num_requests = len(metadata.num_draft_tokens) |
| 331 | num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu") |
| 332 | original_indices = torch.arange(num_requests, device="cpu") |
| 333 | repeat_indices_cpu = original_indices.repeat_interleave(num_draft_tokens) |
| 334 | repeat_indices = async_tensor_h2d(repeat_indices_cpu, logits.device) |
| 335 | logits = self.apply_penalties( |
| 336 | logits, sampling_metadata, metadata, repeat_indices, output_token_ids |
| 337 | ) |
| 338 | |
| 339 | # Apply allowed token ids. |
| 340 | if sampling_metadata.allowed_token_ids_mask is not None: |
| 341 | token_mask = sampling_metadata.allowed_token_ids_mask[repeat_indices] |
| 342 | logits.masked_fill_(token_mask, float("-inf")) |
| 343 | |
| 344 | # Apply bad words exclusion. |
| 345 | if bad_words_token_ids := sampling_metadata.bad_words_token_ids: |
| 346 | apply_bad_words_with_drafts( |
| 347 | logits, bad_words_token_ids, output_token_ids, metadata.num_draft_tokens |
| 348 | ) |
| 349 | |
| 350 | for processor in sampling_metadata.logitsprocs.non_argmax_invariant: |
| 351 | if isinstance(processor, MinTokensLogitsProcessor): |
| 352 | logits = processor.apply_with_spec_decode( |
| 353 | logits, metadata.num_draft_tokens |
| 354 | ) |
| 355 | holder = sampling_metadata.thinking_budget_state_holder |
| 356 | if holder is not None and holder.has_tracked_requests(): |
| 357 | logits = holder.apply_to_logits( |
| 358 | logits, |
| 359 | predict_bonus_token=False, |
| 360 | spec_token_ids=sampling_metadata.spec_token_ids, |
| 361 | ) |
| 362 | return logits |
| 363 | |
| 364 | @staticmethod |
| 365 | def apply_penalties( |
| 366 | logits: torch.Tensor, |
| 367 | sampling_metadata: SamplingMetadata, |
| 368 | metadata: SpecDecodeMetadata, |
| 369 | repeat_indices: torch.Tensor, |
| 370 | output_token_ids: list[list[int]], |
| 371 | ) -> torch.Tensor: |
| 372 | if sampling_metadata.no_penalties: |
| 373 | return logits |
| 374 | |
| 375 | assert sampling_metadata.prompt_token_ids is not None |
| 376 | |
| 377 | prompt_token_ids = sampling_metadata.prompt_token_ids[repeat_indices] |
| 378 | presence_penalties = sampling_metadata.presence_penalties[repeat_indices] |
| 379 | frequency_penalties = sampling_metadata.frequency_penalties[repeat_indices] |
| 380 | repetition_penalties = sampling_metadata.repetition_penalties[repeat_indices] |
| 381 | |
| 382 | logits = apply_all_penalties( |
| 383 | logits, |
| 384 | prompt_token_ids, |
| 385 | presence_penalties, |
| 386 | frequency_penalties, |
| 387 | repetition_penalties, |
| 388 | output_token_ids, |
| 389 | ) |
| 390 | return logits |
| 391 | |
| 392 | @staticmethod |
| 393 | def _combine_outputs_with_spec_tokens( |
| 394 | output_token_ids: list[list[int]], |
| 395 | spec_token_ids: list[list[int]] | None = None, |
| 396 | ) -> list[list[int]]: |
| 397 | if spec_token_ids is None: |
| 398 | return output_token_ids |
| 399 | |
| 400 | result = [] |
| 401 | for out, spec in zip(output_token_ids, spec_token_ids): |
| 402 | if len(spec) == 0: |
| 403 | continue |
| 404 | result.append(out) |
| 405 | for i in range(len(spec) - 1): |
| 406 | result.append([*result[-1], spec[i]]) |
| 407 | return result |
| 408 | |
| 409 | |
| 410 | def rejection_sample( |
| 411 | # [num_tokens] |
| 412 | draft_token_ids: torch.Tensor, |
| 413 | # [batch_size] |
| 414 | num_draft_tokens: list[int], |
| 415 | max_spec_len: int, |
| 416 | # [batch_size] |
| 417 | cu_num_draft_tokens: torch.Tensor, |
| 418 | # [num_tokens, vocab_size] |
| 419 | draft_probs: torch.Tensor | None, |
| 420 | # [num_tokens, vocab_size] |
| 421 | target_logits: torch.Tensor, |
| 422 | # [batch_size, 1] |
| 423 | bonus_token_ids: torch.Tensor, |
| 424 | sampling_metadata: SamplingMetadata, |
| 425 | synthetic_mode: bool = False, |
| 426 | synthetic_conditional_rates: torch.Tensor | None = None, |
| 427 | use_fp64_gumbel: bool = False, |
| 428 | ) -> torch.Tensor: |
| 429 | assert draft_token_ids.ndim == 1 |
| 430 | assert draft_probs is None or draft_probs.ndim == 2 |
| 431 | assert cu_num_draft_tokens.ndim == 1 |
| 432 | assert target_logits.ndim == 2 |
| 433 | |
| 434 | batch_size = len(num_draft_tokens) |
| 435 | num_tokens = draft_token_ids.shape[0] |
| 436 | vocab_size = target_logits.shape[-1] |
| 437 | device = target_logits.device |
| 438 | assert draft_token_ids.is_contiguous() |
| 439 | assert draft_probs is None or draft_probs.is_contiguous() |
| 440 | assert bonus_token_ids.is_contiguous() |
| 441 | assert target_logits.shape == (num_tokens, vocab_size) |
| 442 | |
| 443 | # Create output buffer. |
| 444 | output_token_ids = torch.full( |
| 445 | (batch_size, max_spec_len + 1), |
| 446 | PLACEHOLDER_TOKEN_ID, |
| 447 | dtype=torch.int32, # Consistent with SamplerOutput.sampled_token_ids. |
| 448 | device=device, |
| 449 | ) |
| 450 | |
| 451 | if sampling_metadata.all_greedy: |
| 452 | is_greedy = None |
| 453 | else: |
| 454 | is_greedy = sampling_metadata.temperature == GREEDY_TEMPERATURE |
| 455 | |
| 456 | # Generate uniform probabilities before either kernel because synthetic |
| 457 | # mode needs them in the greedy kernel too. Skip only when all requests |
| 458 | # are greedy *and* synthetic mode is off (the standard fast-path). |
| 459 | # [num_tokens] |
| 460 | uniform_probs: torch.Tensor | None = None |
| 461 | if synthetic_mode or not sampling_metadata.all_greedy: |
| 462 | uniform_probs = generate_uniform_probs( |
| 463 | num_tokens, |
| 464 | num_draft_tokens, |
| 465 | sampling_metadata.generators, |
| 466 | device, |
| 467 | ) |
| 468 | |
| 469 | if not sampling_metadata.all_random: |
| 470 | # Rejection sampling for greedy sampling requests. |
| 471 | target_argmax = target_logits.argmax(dim=-1) |
| 472 | _rejection_greedy_sample( |
| 473 | output_token_ids, |
| 474 | cu_num_draft_tokens, |
| 475 | draft_token_ids, |
| 476 | target_argmax, |
| 477 | bonus_token_ids, |
| 478 | is_greedy, |
| 479 | max_spec_len, |
| 480 | uniform_probs, |
| 481 | synthetic_conditional_rates, |
| 482 | ) |
| 483 | if sampling_metadata.all_greedy: |
| 484 | return output_token_ids |
| 485 | |
| 486 | # Compute probability distribution from target logits. |
| 487 | target_probs = target_logits.softmax(dim=-1, dtype=torch.float32) |
| 488 | assert target_probs.is_contiguous() |
| 489 | |
| 490 | # Sample recovered tokens for each position. |
| 491 | # [num_tokens] |
| 492 | recovered_token_ids = sample_recovered_tokens( |
| 493 | max_spec_len, |
| 494 | num_draft_tokens, |
| 495 | cu_num_draft_tokens, |
| 496 | draft_token_ids, |
| 497 | draft_probs, |
| 498 | target_probs, |
| 499 | sampling_metadata, |
| 500 | device, |
| 501 | use_fp64_gumbel, |
| 502 | ) |
| 503 | |
| 504 | # Rejection sampling for random sampling requests. |
| 505 | assert uniform_probs is not None |
| 506 | _rejection_random_sample( |
| 507 | output_token_ids, |
| 508 | cu_num_draft_tokens, |
| 509 | draft_token_ids, |
| 510 | draft_probs, |
| 511 | target_probs, |
| 512 | bonus_token_ids, |
| 513 | recovered_token_ids, |
| 514 | uniform_probs, |
| 515 | is_greedy, |
| 516 | max_spec_len, |
| 517 | synthetic_conditional_rates, |
| 518 | ) |
| 519 | return output_token_ids |
| 520 | |
| 521 | |
| 522 | def apply_sampling_constraints( |
| 523 | logits: torch.Tensor, # [num_tokens, vocab_size] |
| 524 | cu_num_draft_tokens: torch.Tensor, # [batch_size] |
| 525 | sampling_metadata: SamplingMetadata, |
| 526 | ) -> torch.Tensor: |
| 527 | """Process logits based on sampling metadata. |
| 528 | |
| 529 | This function applies temperature scaling to the logits, |
| 530 | as well as top-k and top-p. For greedy decoding, it returns |
| 531 | the original logits. |
| 532 | |
| 533 | Args: |
| 534 | logits: Input logits tensor to be processed. |
| 535 | cu_num_draft_tokens: Cumulative number of draft tokens. |
| 536 | sampling_metadata: Metadata containing sampling parameters such as |
| 537 | temperature and whether greedy sampling is used. |
| 538 | |
| 539 | Returns: |
| 540 | torch.Tensor: Processed logits if non-greedy sampling is used, |
| 541 | otherwise returns the original logits. |
| 542 | |
| 543 | """ |
| 544 | assert logits.ndim == 2 |
| 545 | assert cu_num_draft_tokens.ndim == 1 |
| 546 | if sampling_metadata.all_greedy: |
| 547 | return logits |
| 548 | |
| 549 | num_tokens = logits.shape[0] |
| 550 | temperature = expand_batch_to_tokens( |
| 551 | sampling_metadata.temperature, |
| 552 | cu_num_draft_tokens, |
| 553 | num_tokens, |
| 554 | replace_from=GREEDY_TEMPERATURE, |
| 555 | replace_to=1, |
| 556 | ) |
| 557 | # NOTE(woosuk): Update `logits` in place to avoid allocating a new tensor. |
| 558 | logits.div_(temperature.unsqueeze(-1)) |
| 559 | |
| 560 | # Get expanded top_k and top_p tensors. |
| 561 | top_k = None |
| 562 | if sampling_metadata.top_k is not None: |
| 563 | top_k = expand_batch_to_tokens( |
| 564 | sampling_metadata.top_k, |
| 565 | cu_num_draft_tokens, |
| 566 | num_tokens, |
| 567 | ) |
| 568 | top_p = None |
| 569 | if sampling_metadata.top_p is not None: |
| 570 | top_p = expand_batch_to_tokens( |
| 571 | sampling_metadata.top_p, |
| 572 | cu_num_draft_tokens, |
| 573 | num_tokens, |
| 574 | ) |
| 575 | |
| 576 | # NOTE(woosuk): `apply_top_k_top_p` uses sorting to calculate the mask, |
| 577 | # which is slow for large vocab sizes. This may cause performance issues. |
| 578 | return apply_top_k_top_p(logits, top_k, top_p) |
| 579 | |
| 580 | |
| 581 | def expand_batch_to_tokens( |
| 582 | x: torch.Tensor, # [batch_size] |
| 583 | cu_num_tokens: torch.Tensor, # [batch_size] |
| 584 | num_tokens: int, |
| 585 | replace_from: int = 0, |
| 586 | replace_to: int = 0, |
| 587 | ) -> torch.Tensor: |
| 588 | """Expand [batch_size] tensor to [num_tokens] tensor based on the number of |
| 589 | tokens per batch in cu_num_tokens. |
| 590 | |
| 591 | For example, if x = [a, b, c] and cu_num_tokens = [2, 5, 6], then |
| 592 | num_tokens = 6, and expanded_x = [a, a, b, b, b, c]. |
| 593 | |
| 594 | Args: |
| 595 | x: [batch_size] tensor to expand. |
| 596 | cu_num_tokens: [batch_size] tensor containing the cumulative number of |
| 597 | tokens per batch. Each element represents the total number of |
| 598 | tokens up to and including that batch. |
| 599 | num_tokens: Total number of tokens. |
| 600 | replace_from: int = 0 |
| 601 | Value to be replaced if it is found in x. |
| 602 | replace_to: int = 0 |
| 603 | Value to replace with when replace_from is found. |
| 604 | |
| 605 | Returns: |
| 606 | expanded_x: [num_tokens] tensor. |
| 607 | |
| 608 | """ |
| 609 | batch_size = x.shape[0] |
| 610 | assert cu_num_tokens.shape[0] == batch_size |
| 611 | expanded_x = x.new_empty(num_tokens) |
| 612 | _expand( |
| 613 | expanded_x, |
| 614 | x, |
| 615 | cu_num_tokens, |
| 616 | replace_from, |
| 617 | replace_to, |
| 618 | ) |
| 619 | return expanded_x |
| 620 | |
| 621 | |
| 622 | def generate_uniform_probs( |
| 623 | num_tokens: int, |
| 624 | num_draft_tokens: list[int], |
| 625 | generators: dict[int, torch.Generator], |
| 626 | device: torch.device, |
| 627 | ) -> torch.Tensor: |
| 628 | """Generates a batch of uniform random samples, with optional seeding |
| 629 | if available. |
| 630 | |
| 631 | This method creates a tensor of shape `(num_tokens, )` filled |
| 632 | with uniform random values in the range [0, 1). If `generators` is provided, |
| 633 | the requests with their own seeds will use the provided `torch.Generator` |
| 634 | for reproducibility. The samples for the other requests will be generated |
| 635 | without a seed. |
| 636 | |
| 637 | Args: |
| 638 | num_tokens: int |
| 639 | Total number of tokens. |
| 640 | num_draft_tokens: List[List[int]] |
| 641 | Number of draft tokens per request. |
| 642 | generators: Optional[Dict[int, torch.Generator]] |
| 643 | A dictionary mapping indices in the batch to |
| 644 | `torch.Generator` objects. |
| 645 | device: torch.device |
| 646 | The device on which to allocate the tensor. |
| 647 | |
| 648 | Returns: |
| 649 | uniform_rand: torch.Tensor |
| 650 | A tensor of shape `(num_tokens, )` containing uniform |
| 651 | random values in the range [0, 1). |
| 652 | |
| 653 | """ |
| 654 | # NOTE(woosuk): We deliberately use float64 instead of float32 here |
| 655 | # because when using float32, there's a non-negligible chance that |
| 656 | # uniform_prob is sampled to be exact 0.0 as reported in |
| 657 | # https://github.com/pytorch/pytorch/issues/16706. Using float64 |
| 658 | # mitigates the issue. |
| 659 | uniform_probs = torch.rand( |
| 660 | (num_tokens,), |
| 661 | dtype=torch.float64, |
| 662 | device=device, |
| 663 | ) |
| 664 | start_idx = 0 |
| 665 | for req_idx, n in enumerate(num_draft_tokens): |
| 666 | # Do not generate random numbers for requests with no draft tokens. |
| 667 | # This can be important for reproducibility. |
| 668 | if n == 0: |
| 669 | continue |
| 670 | end_idx = start_idx + n |
| 671 | generator = generators.get(req_idx) |
| 672 | if generator is not None: |
| 673 | uniform_probs[start_idx:end_idx].uniform_(generator=generator) |
| 674 | start_idx = end_idx |
| 675 | return uniform_probs |
| 676 | |
| 677 | |
| 678 | def sample_recovered_tokens( |
| 679 | max_spec_len: int, |
| 680 | num_draft_tokens: list[int], |
| 681 | # [batch_size] |
| 682 | cu_num_draft_tokens: torch.Tensor, |
| 683 | # [num_tokens] |
| 684 | draft_token_ids: torch.Tensor, |
| 685 | # [num_tokens, vocab_size] |
| 686 | draft_probs: torch.Tensor | None, |
| 687 | # [num_tokens, vocab_size] |
| 688 | target_probs: torch.Tensor, |
| 689 | sampling_metadata: SamplingMetadata, |
| 690 | device: torch.device, |
| 691 | use_fp64_gumbel: bool = False, |
| 692 | ) -> torch.Tensor: |
| 693 | # NOTE(woosuk): Create only one distribution for each request. |
| 694 | batch_size = len(num_draft_tokens) |
| 695 | vocab_size = target_probs.shape[-1] |
| 696 | q_dtype = torch.float64 if use_fp64_gumbel else torch.float32 |
| 697 | q = torch.empty( |
| 698 | (batch_size, vocab_size), |
| 699 | dtype=q_dtype, |
| 700 | device=device, |
| 701 | ) |
| 702 | q.exponential_() |
| 703 | for i, generator in sampling_metadata.generators.items(): |
| 704 | # Do not generate random numbers for requests with no draft tokens. |
| 705 | # This can be important for reproducibility. |
| 706 | if num_draft_tokens[i] > 0: |
| 707 | q[i].exponential_(generator=generator) |
| 708 | |
| 709 | inv_q = q.reciprocal() |
| 710 | |
| 711 | recovered_token_ids = torch.empty_like(draft_token_ids) |
| 712 | _sample_recovered_tokens( |
| 713 | recovered_token_ids, |
| 714 | cu_num_draft_tokens, |
| 715 | draft_token_ids, |
| 716 | draft_probs, |
| 717 | target_probs, |
| 718 | inv_q, |
| 719 | max_spec_len=max_spec_len, |
| 720 | ) |
| 721 | return recovered_token_ids |
| 722 | |
| 723 | |
| 724 | # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation. |
| 725 | @triton.jit(do_not_specialize=["max_spec_len"]) |
| 726 | def rejection_greedy_sample_kernel( |
| 727 | output_token_ids_ptr, # [batch_size, max_spec_len + 1] |
| 728 | cu_num_draft_tokens_ptr, # [batch_size] |
| 729 | draft_token_ids_ptr, # [num_tokens] |
| 730 | target_argmax_ptr, # [num_tokens] |
| 731 | bonus_token_ids_ptr, # [batch_size] |
| 732 | is_greedy_ptr, # [batch_size] or None |
| 733 | max_spec_len, |
| 734 | uniform_probs_ptr, # [num_tokens] or None (synthetic mode only) |
| 735 | synthetic_conditional_rates_ptr, # [num_speculative_tokens] or None |
| 736 | SYNTHETIC_MODE: tl.constexpr, |
| 737 | ): |
| 738 | req_idx = tl.program_id(0) |
| 739 | # FIXME(woosuk): Because is_greedy_ptr is not None at profiling run, |
| 740 | # re-compilation may happen during runtime when is_greedy_ptr is None. |
| 741 | is_greedy = True if is_greedy_ptr is None else tl.load(is_greedy_ptr + req_idx) |
| 742 | if not is_greedy: |
| 743 | # Early exit for non-greedy sampling requests. |
| 744 | return |
| 745 | |
| 746 | start_idx = ( |
| 747 | tl.zeros([], dtype=cu_num_draft_tokens_ptr.dtype.element_ty) |
| 748 | if req_idx == 0 |
| 749 | else tl.load(cu_num_draft_tokens_ptr + req_idx - 1) |
| 750 | ) |
| 751 | end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx) |
| 752 | num_draft_tokens = end_idx - start_idx |
| 753 | |
| 754 | rejected = False |
| 755 | for pos in range(num_draft_tokens): |
| 756 | if not rejected: |
| 757 | draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos) |
| 758 | target_argmax_id = tl.load(target_argmax_ptr + start_idx + pos).to(tl.int32) |
| 759 | if SYNTHETIC_MODE: |
| 760 | uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos) |
| 761 | rate = tl.load(synthetic_conditional_rates_ptr + pos) |
| 762 | # -1 is used for padded draft token ids that should be rejected. |
| 763 | accepted = (uniform_prob < rate) and draft_token_id >= 0 |
| 764 | token_id = draft_token_id if accepted else target_argmax_id |
| 765 | rejected = not accepted |
| 766 | else: |
| 767 | token_id = target_argmax_id |
| 768 | rejected = draft_token_id != target_argmax_id |
| 769 | tl.store( |
| 770 | output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, |
| 771 | token_id, |
| 772 | ) |
| 773 | |
| 774 | if not rejected: |
| 775 | # If all tokens are accepted, append the bonus token. |
| 776 | bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx) |
| 777 | tl.store( |
| 778 | output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens, |
| 779 | bonus_token_id, |
| 780 | ) |
| 781 | |
| 782 | |
| 783 | def _rejection_greedy_sample_warmup_inputs( |
| 784 | *, synthetic_mode: bool |
| 785 | ) -> dict[str, object]: |
| 786 | all_greedy = WarmupChoices(False, True) |
| 787 | int32 = TritonWarmupTensor(torch.int32) |
| 788 | return dict( |
| 789 | output_token_ids=int32, |
| 790 | cu_num_draft_tokens=int32, |
| 791 | draft_token_ids=int32, |
| 792 | target_argmax=TritonWarmupTensor(torch.int64), |
| 793 | bonus_token_ids=TritonWarmupTensor(torch.int64), |
| 794 | is_greedy=None if all_greedy else TritonWarmupTensor(torch.bool), |
| 795 | max_spec_len=5, |
| 796 | uniform_probs=( |
| 797 | None |
| 798 | if all_greedy and not synthetic_mode |
| 799 | else TritonWarmupTensor(torch.float64) |
| 800 | ), |
| 801 | synthetic_conditional_rates=( |
| 802 | TritonWarmupTensor(torch.float32) if synthetic_mode else None |
| 803 | ), |
| 804 | ) |
| 805 | |
| 806 | |
| 807 | @triton_kernel_dispatcher_with_warmup( |
| 808 | kernel=rejection_greedy_sample_kernel, |
| 809 | warmup_inputs=_rejection_greedy_sample_warmup_inputs, |
| 810 | ) |
| 811 | def _rejection_greedy_sample( |
| 812 | output_token_ids: torch.Tensor, |
| 813 | cu_num_draft_tokens: torch.Tensor, |
| 814 | draft_token_ids: torch.Tensor, |
| 815 | target_argmax: torch.Tensor, |
| 816 | bonus_token_ids: torch.Tensor, |
| 817 | is_greedy: torch.Tensor | None, |
| 818 | max_spec_len: int, |
| 819 | uniform_probs: torch.Tensor | None, |
| 820 | synthetic_conditional_rates: torch.Tensor | None, |
| 821 | ) -> DispatchSpec: |
| 822 | return (output_token_ids.shape[0],), dict( |
| 823 | SYNTHETIC_MODE=synthetic_conditional_rates is not None, |
| 824 | ) |
| 825 | |
| 826 | |
| 827 | # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation. |
| 828 | @triton.jit(do_not_specialize=["max_spec_len"]) |
| 829 | def rejection_random_sample_kernel( |
| 830 | output_token_ids_ptr, # [batch_size, max_spec_len + 1] |
| 831 | cu_num_draft_tokens_ptr, # [batch_size] |
| 832 | draft_token_ids_ptr, # [num_tokens] |
| 833 | draft_probs_ptr, # [num_tokens, vocab_size] or None |
| 834 | target_probs_ptr, # [num_tokens, vocab_size] |
| 835 | bonus_token_ids_ptr, # [batch_size] |
| 836 | recovered_token_ids_ptr, # [num_tokens] |
| 837 | uniform_probs_ptr, # [num_tokens] |
| 838 | is_greedy_ptr, # [batch_size] |
| 839 | max_spec_len, |
| 840 | vocab_size, |
| 841 | synthetic_conditional_rates_ptr, # [num_speculative_tokens] or None |
| 842 | NO_DRAFT_PROBS: tl.constexpr, |
| 843 | SYNTHETIC_MODE: tl.constexpr, |
| 844 | ): |
| 845 | req_idx = tl.program_id(0) |
| 846 | is_greedy = tl.load(is_greedy_ptr + req_idx) |
| 847 | if is_greedy: |
| 848 | # Early exit for greedy sampling requests. |
| 849 | return |
| 850 | |
| 851 | start_idx = ( |
| 852 | tl.zeros([], dtype=cu_num_draft_tokens_ptr.dtype.element_ty) |
| 853 | if req_idx == 0 |
| 854 | else tl.load(cu_num_draft_tokens_ptr + req_idx - 1) |
| 855 | ) |
| 856 | end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx) |
| 857 | num_draft_tokens = end_idx - start_idx |
| 858 | |
| 859 | rejected = False |
| 860 | for pos in range(num_draft_tokens): |
| 861 | if not rejected: |
| 862 | draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos) |
| 863 | uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos) |
| 864 | if draft_token_id < 0: |
| 865 | # -1 is used for padded draft token ids that should be rejected. |
| 866 | accepted = False |
| 867 | elif SYNTHETIC_MODE: |
| 868 | rate = tl.load(synthetic_conditional_rates_ptr + pos) |
| 869 | accepted = uniform_prob < rate |
| 870 | else: |
| 871 | if NO_DRAFT_PROBS: |
| 872 | draft_prob = 1 |
| 873 | else: |
| 874 | draft_prob = tl.load( |
| 875 | draft_probs_ptr |
| 876 | + (start_idx + pos) * vocab_size |
| 877 | + draft_token_id |
| 878 | ) |
| 879 | target_prob = tl.load( |
| 880 | target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id |
| 881 | ) |
| 882 | # NOTE(woosuk): While the draft probability should never be 0, |
| 883 | # we check it to avoid NaNs. If it happens to be 0, we reject. |
| 884 | accepted = draft_prob > 0 and target_prob / draft_prob >= uniform_prob |
| 885 | if accepted: |
| 886 | token_id = draft_token_id |
| 887 | else: |
| 888 | rejected = True |
| 889 | token_id = tl.load(recovered_token_ids_ptr + start_idx + pos) |
| 890 | tl.store( |
| 891 | output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id |
| 892 | ) |
| 893 | |
| 894 | if not rejected: |
| 895 | # If all tokens are accepted, append the bonus token. |
| 896 | bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx) |
| 897 | tl.store( |
| 898 | output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens, |
| 899 | bonus_token_id, |
| 900 | ) |
| 901 | |
| 902 | |
| 903 | def _rejection_random_sample_warmup_inputs( |
| 904 | *, synthetic_mode: bool |
| 905 | ) -> dict[str, object]: |
| 906 | vocab_size: Any = WarmupChoices(2, 16) |
| 907 | no_draft_probs = WarmupChoices(False, True) |
| 908 | int32 = TritonWarmupTensor(torch.int32) |
| 909 | probabilities = TritonWarmupTensor( |
| 910 | torch.float32, |
| 911 | shape=(1, vocab_size), |
| 912 | ) |
| 913 | return dict( |
| 914 | output_token_ids=int32, |
| 915 | cu_num_draft_tokens=int32, |
| 916 | draft_token_ids=int32, |
| 917 | draft_probs=None if no_draft_probs else probabilities, |
| 918 | target_probs=probabilities, |
| 919 | bonus_token_ids=TritonWarmupTensor(torch.int64), |
| 920 | recovered_token_ids=int32, |
| 921 | uniform_probs=TritonWarmupTensor(torch.float64), |
| 922 | is_greedy=TritonWarmupTensor(torch.bool), |
| 923 | max_spec_len=5, |
| 924 | synthetic_conditional_rates=( |
| 925 | TritonWarmupTensor(torch.float32) if synthetic_mode else None |
| 926 | ), |
| 927 | ) |
| 928 | |
| 929 | |
| 930 | @triton_kernel_dispatcher_with_warmup( |
| 931 | kernel=rejection_random_sample_kernel, |
| 932 | warmup_inputs=_rejection_random_sample_warmup_inputs, |
| 933 | ) |
| 934 | def _rejection_random_sample( |
| 935 | output_token_ids: torch.Tensor, |
| 936 | cu_num_draft_tokens: torch.Tensor, |
| 937 | draft_token_ids: torch.Tensor, |
| 938 | draft_probs: torch.Tensor | None, |
| 939 | target_probs: torch.Tensor, |
| 940 | bonus_token_ids: torch.Tensor, |
| 941 | recovered_token_ids: torch.Tensor, |
| 942 | uniform_probs: torch.Tensor, |
| 943 | is_greedy: torch.Tensor, |
| 944 | max_spec_len: int, |
| 945 | synthetic_conditional_rates: torch.Tensor | None, |
| 946 | ) -> DispatchSpec: |
| 947 | return (output_token_ids.shape[0],), dict( |
| 948 | vocab_size=target_probs.shape[1], |
| 949 | NO_DRAFT_PROBS=draft_probs is None, |
| 950 | SYNTHETIC_MODE=synthetic_conditional_rates is not None, |
| 951 | ) |
| 952 | |
| 953 | |
| 954 | # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation. |
| 955 | @triton.jit(do_not_specialize=["replace_from", "replace_to"]) |
| 956 | def expand_kernel( |
| 957 | output_ptr, # [num_tokens] |
| 958 | input_ptr, # [batch_size] |
| 959 | cu_num_tokens_ptr, # [batch_size] |
| 960 | replace_from, |
| 961 | replace_to, |
| 962 | MAX_NUM_TOKENS: tl.constexpr, |
| 963 | ): |
| 964 | req_idx = tl.program_id(0) |
| 965 | if req_idx == 0: |
| 966 | start_idx = tl.zeros([], dtype=cu_num_tokens_ptr.dtype.element_ty) |
| 967 | else: |
| 968 | start_idx = tl.load(cu_num_tokens_ptr + req_idx - 1) |
| 969 | end_idx = tl.load(cu_num_tokens_ptr + req_idx) |
| 970 | num_tokens = end_idx - start_idx |
| 971 | |
| 972 | src_val = tl.load(input_ptr + req_idx) |
| 973 | src_val = tl.where(src_val == replace_from, replace_to, src_val) |
| 974 | offset = tl.arange(0, MAX_NUM_TOKENS) |
| 975 | tl.store(output_ptr + start_idx + offset, src_val, mask=offset < num_tokens) |
| 976 | |
| 977 | |
| 978 | def _expand_warmup_inputs() -> dict[str, object]: |
| 979 | dtype = WarmupChoices(torch.float32, torch.int32, torch.int64) |
| 980 | data = TritonWarmupTensor(dtype) |
| 981 | return dict( |
| 982 | output=data, |
| 983 | input=data, |
| 984 | cu_num_tokens=TritonWarmupTensor(torch.int32), |
| 985 | replace_from=0, |
| 986 | replace_to=0, |
| 987 | ) |
| 988 | |
| 989 | |
| 990 | @triton_kernel_dispatcher_with_warmup( |
| 991 | kernel=expand_kernel, warmup_inputs=_expand_warmup_inputs |
| 992 | ) |
| 993 | def _expand( |
| 994 | output: torch.Tensor, |
| 995 | input: torch.Tensor, |
| 996 | cu_num_tokens: torch.Tensor, |
| 997 | replace_from: int, |
| 998 | replace_to: int, |
| 999 | ) -> DispatchSpec: |
| 1000 | return (input.shape[0],), dict(MAX_NUM_TOKENS=MAX_SPEC_LEN) |
| 1001 | |
| 1002 | |
| 1003 | @triton.jit |
| 1004 | def sample_recovered_tokens_kernel( |
| 1005 | output_token_ids_ptr, # [num_tokens] |
| 1006 | cu_num_draft_tokens_ptr, # [batch_size] |
| 1007 | draft_token_ids_ptr, # [num_tokens] |
| 1008 | draft_probs_ptr, # [num_tokens, vocab_size] or None |
| 1009 | target_probs_ptr, # [num_tokens, vocab_size] |
| 1010 | inv_q_ptr, # [batch_size, vocab_size] |
| 1011 | vocab_size, |
| 1012 | BLOCK_SIZE: tl.constexpr, |
| 1013 | NO_DRAFT_PROBS: tl.constexpr, |
| 1014 | USE_FP64_GUMBEL: tl.constexpr, |
| 1015 | ): |
| 1016 | req_idx = tl.program_id(0) |
| 1017 | start_idx = ( |
| 1018 | tl.zeros([], dtype=cu_num_draft_tokens_ptr.dtype.element_ty) |
| 1019 | if req_idx == 0 |
| 1020 | else tl.load(cu_num_draft_tokens_ptr + req_idx - 1) |
| 1021 | ) |
| 1022 | end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx) |
| 1023 | num_draft_tokens = end_idx - start_idx |
| 1024 | |
| 1025 | # Early exit for out-of-range positions. |
| 1026 | pos = tl.program_id(1) |
| 1027 | if pos >= num_draft_tokens: |
| 1028 | return |
| 1029 | |
| 1030 | token_idx = start_idx + pos |
| 1031 | |
| 1032 | if NO_DRAFT_PROBS: |
| 1033 | draft_token_id = tl.load(draft_token_ids_ptr + token_idx) |
| 1034 | |
| 1035 | if USE_FP64_GUMBEL: |
| 1036 | max_val = tl.full((), float("-inf"), tl.float64) |
| 1037 | else: |
| 1038 | max_val = tl.full((), float("-inf"), tl.float32) |
| 1039 | recovered_id = 0 |
| 1040 | for v in range(0, vocab_size, BLOCK_SIZE): |
| 1041 | vocab_offset = v + tl.arange(0, BLOCK_SIZE) |
| 1042 | vocab_mask = vocab_offset < vocab_size |
| 1043 | |
| 1044 | if NO_DRAFT_PROBS: |
| 1045 | prob = tl.load( |
| 1046 | target_probs_ptr + token_idx * vocab_size + vocab_offset, |
| 1047 | mask=(vocab_mask & (vocab_offset != draft_token_id)), |
| 1048 | other=0.0, |
| 1049 | ) |
| 1050 | else: |
| 1051 | draft_prob = tl.load( |
| 1052 | draft_probs_ptr + token_idx * vocab_size + vocab_offset, |
| 1053 | mask=vocab_mask, |
| 1054 | other=0.0, |
| 1055 | ) |
| 1056 | target_prob = tl.load( |
| 1057 | target_probs_ptr + token_idx * vocab_size + vocab_offset, |
| 1058 | mask=vocab_mask, |
| 1059 | other=0.0, |
| 1060 | ) |
| 1061 | prob = tl.maximum(target_prob - draft_prob, 0.0) |
| 1062 | # NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because |
| 1063 | # `tl.argmax` will select the maximum value. |
| 1064 | |
| 1065 | inv_q = tl.load( |
| 1066 | inv_q_ptr + req_idx * vocab_size + vocab_offset, |
| 1067 | mask=vocab_mask, |
| 1068 | other=0.0, |
| 1069 | ) |
| 1070 | |
| 1071 | # Local tile reduction. |
| 1072 | # Mask out-of-vocabulary entries to -inf so they can never win |
| 1073 | # the argmax — prevents producing recovered_id >= vocab_size |
| 1074 | # when all valid entries in the last tile have zero probability. |
| 1075 | score = prob * inv_q |
| 1076 | score = tl.where(vocab_mask, score, float("-inf")) |
| 1077 | local_max, local_id = tl.max(score, axis=0, return_indices=True) |
| 1078 | |
| 1079 | if local_max > max_val: |
| 1080 | max_val = local_max |
| 1081 | recovered_id = v + local_id |
| 1082 | |
| 1083 | recovered_id = tl.minimum(recovered_id, vocab_size - 1) |
| 1084 | tl.store(output_token_ids_ptr + token_idx, recovered_id) |
| 1085 | |
| 1086 | |
| 1087 | def _sample_recovered_tokens_warmup_inputs( |
| 1088 | *, use_fp64_gumbel: bool |
| 1089 | ) -> dict[str, object]: |
| 1090 | vocab_size: Any = WarmupChoices(2, 16) |
| 1091 | no_draft_probs = WarmupChoices(False, True) |
| 1092 | int32 = TritonWarmupTensor(torch.int32) |
| 1093 | probabilities = TritonWarmupTensor( |
| 1094 | torch.float32, |
| 1095 | shape=(1, vocab_size), |
| 1096 | ) |
| 1097 | return dict( |
| 1098 | output_token_ids=int32, |
| 1099 | cu_num_draft_tokens=int32, |
| 1100 | draft_token_ids=int32, |
| 1101 | draft_probs=None if no_draft_probs else probabilities, |
| 1102 | target_probs=probabilities, |
| 1103 | inv_q=TritonWarmupTensor( |
| 1104 | torch.float64 if use_fp64_gumbel else torch.float32, |
| 1105 | shape=(1, vocab_size), |
| 1106 | ), |
| 1107 | max_spec_len=5, |
| 1108 | ) |
| 1109 | |
| 1110 | |
| 1111 | @triton_kernel_dispatcher_with_warmup( |
| 1112 | kernel=sample_recovered_tokens_kernel, |
| 1113 | warmup_inputs=_sample_recovered_tokens_warmup_inputs, |
| 1114 | ) |
| 1115 | def _sample_recovered_tokens( |
| 1116 | output_token_ids: torch.Tensor, |
| 1117 | cu_num_draft_tokens: torch.Tensor, |
| 1118 | draft_token_ids: torch.Tensor, |
| 1119 | draft_probs: torch.Tensor | None, |
| 1120 | target_probs: torch.Tensor, |
| 1121 | inv_q: torch.Tensor, |
| 1122 | *, |
| 1123 | max_spec_len: int, |
| 1124 | ) -> DispatchSpec: |
| 1125 | return (cu_num_draft_tokens.shape[0], max_spec_len), dict( |
| 1126 | vocab_size=target_probs.shape[1], |
| 1127 | BLOCK_SIZE=8192, |
| 1128 | NO_DRAFT_PROBS=draft_probs is None, |
| 1129 | USE_FP64_GUMBEL=inv_q.dtype == torch.float64, |
| 1130 | ) |