FoundationVision/LlamaGen
64
1"""A layer that samples the next tokens from the model's outputs."""2import itertools3from typing import Dict, List, Optional, Tuple4 5import torch6import torch.nn as nn7 8from vllm.model_executor.layers.ops.sample import sample as sample_triton9from vllm.model_executor.sampling_metadata import (SamplingMetadata,10 SamplingTensors)11from vllm.sampling_params import SamplingParams, SamplingType12from vllm.sequence import (Logprob, PromptLogprobs, SampleLogprobs,13 SamplerOutput, SequenceData, SequenceGroupOutput,14 SequenceOutput)15 16 17class Sampler(nn.Module):18 """Samples the next tokens from the model's outputs.19 20 This layer does the following:21 1. Discard the hidden states that are not used for sampling (i.e., all22 tokens except the final one in each prompt).23 2. Compute the logits for the next tokens.24 3. Apply presence, frequency and repetition penalties.25 4. Apply temperature scaling.26 5. Apply top-p and top-k truncation.27 6. Sample the next tokens.28 Here, each sequence group within the batch can have different sampling29 parameters (e.g., sampling method, temperature, top-p, top-k, etc.).30 31 The structure of the logits tensor is coupled with the seq_groups in32 sampling_metadata. Typically, each sequence in each seq_group has one row in33 logits for the next token to be sampled; however, for a seq_group with a34 prompt request with the prompt_logprobs sampling parameter, there are rows35 in logits for each token in the input prompt.36 """37 38 def __init__(self, cfg_scale=1.0):39 super().__init__()40 self.cfg_scale = cfg_scale41 # Whether or not the SamplerOutput should have on-device tensors42 # containing the sampled token ids and probabilities. This is used by43 # speculative decoding.44 self.include_gpu_probs_tensor = False45 46 def forward(47 self,48 logits: torch.Tensor,49 sampling_metadata: SamplingMetadata,50 ) -> Optional[SamplerOutput]:51 assert logits is not None52 _, vocab_size = logits.shape53 54 if self.cfg_scale > 1.0:55 logits_combined = logits56 cond_logits, uncond_logits = torch.split(logits_combined, len(logits_combined) // 2, dim=0)57 logits = uncond_logits + (cond_logits - uncond_logits) * self.cfg_scale58 logits = torch.cat([logits, logits], dim=0)59 60 # Apply min_tokens penalty which sets stop tokens to -inf if min_tokens61 # have not been generated yet62 logits = _apply_min_tokens_penalty(logits, sampling_metadata)63 64 # Prepare sampling tensors with pinned memory to avoid blocking.65 (sampling_tensors, do_penalties, do_top_p_top_k,66 do_min_p) = SamplingTensors.from_sampling_metadata(67 sampling_metadata, vocab_size, logits.device, logits.dtype)68 69 # Apply presence and frequency penalties.70 if do_penalties:71 logits = _apply_penalties(logits, sampling_tensors.prompt_tokens,72 sampling_tensors.output_tokens,73 sampling_tensors.presence_penalties,74 sampling_tensors.frequency_penalties,75 sampling_tensors.repetition_penalties)76 77 # Apply temperature scaling.78 # Use in-place division to avoid creating a new tensor.79 logits.div_(sampling_tensors.temperatures.unsqueeze_(dim=1))80 81 if do_top_p_top_k:82 logits = _apply_top_k_top_p(logits, sampling_tensors.top_ps,83 sampling_tensors.top_ks)84 85 if do_min_p:86 logits = _apply_min_p(logits, sampling_tensors.min_ps)87 88 # We use float32 for probabilities and log probabilities.89 # Compute the probabilities.90 probs = torch.softmax(logits, dim=-1, dtype=torch.float)91 # Compute the log probabilities.92 # Use log_softmax to ensure numerical stability.93 logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)94 95 # Sample the next tokens.96 sample_results, maybe_sampled_tokens_tensor = _sample(97 probs,98 logprobs,99 sampling_metadata,100 sampling_tensors,101 include_gpu_probs_tensor=self.include_gpu_probs_tensor,102 modify_greedy_probs=self._should_modify_greedy_probs_inplace,103 )104 105 106 if self.cfg_scale > 1.0:107 cond_result = sample_results[:len(sample_results) // 2]108 sample_results = cond_result + cond_result109 110 111 if self.include_gpu_probs_tensor:112 assert maybe_sampled_tokens_tensor is not None113 sampled_tokens_tensor = maybe_sampled_tokens_tensor114 on_device_tensors = (probs, sampled_tokens_tensor)115 else:116 on_device_tensors = None117 118 # Get the logprobs query results.119 prompt_logprobs, sample_logprobs = _get_logprobs(120 logprobs, sampling_metadata, sample_results)121 return _build_sampler_output(sample_results,122 sampling_metadata,123 prompt_logprobs,124 sample_logprobs,125 on_device_tensors=on_device_tensors)126 127 @property128 def _should_modify_greedy_probs_inplace(self) -> bool:129 """Whether or not the sampler should modify the probability distribution130 of greedily-sampled tokens such that multinomial sampling would sample131 the greedily-sampled token.132 133 In other words, if True then we set the probability of the greedily-134 sampled token to 1.135 136 This is used by speculative decoding, which requires that the sampling137 method be encoded into the probability distribution.138 """139 # Modify greedy probs if include_gpu_probs_tensor is set.140 return self.include_gpu_probs_tensor141 142 143def _get_bin_counts_and_mask(144 tokens: torch.Tensor,145 vocab_size: int,146 num_seqs: int,147) -> Tuple[torch.Tensor, torch.Tensor]:148 # Compute the bin counts for the tokens.149 # vocab_size + 1 for padding.150 bin_counts = torch.zeros((num_seqs, vocab_size + 1),151 dtype=torch.long,152 device=tokens.device)153 bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens))154 bin_counts = bin_counts[:, :vocab_size]155 mask = bin_counts > 0156 157 return bin_counts, mask158 159 160def _apply_min_tokens_penalty(161 logits: torch.Tensor,162 sampling_metadata: SamplingMetadata,163) -> torch.Tensor:164 # list of indices in logits that will be set to -inf165 logits_to_penalize = []166 start_idx = 0167 for i, seq_group in enumerate(sampling_metadata.seq_groups):168 seq_ids, sampling_params = seq_group169 170 # handle prompt_logprobs by skipping rows in logits added for the prompt171 # tokens (prompt logprobs are not penalized)172 if (i < sampling_metadata.num_prompts173 and sampling_params.prompt_logprobs is not None):174 assert len(seq_ids) == 1175 start_idx += sampling_metadata.prompt_lens[i] - 1176 177 min_tokens = sampling_params.min_tokens178 if min_tokens > 0:179 seqs_to_penalize = []180 for i, seq_id in enumerate(seq_ids):181 seq_data = sampling_metadata.seq_data[seq_id]182 if len(seq_data.output_token_ids) < min_tokens:183 seqs_to_penalize.append(i)184 185 if seqs_to_penalize:186 # convert to the index into logits187 seqs_to_penalize = [start_idx + i for i in seqs_to_penalize]188 # use set() to remove any duplicates189 token_ids_to_penalize = set(sampling_params.stop_token_ids +190 [sampling_params.eos_token_id])191 # itertools.product pairs each seq index with every token id192 logits_to_penalize.extend(193 itertools.product(seqs_to_penalize, token_ids_to_penalize))194 195 start_idx += len(seq_ids)196 197 if logits_to_penalize:198 # use zip and * to group indices along each dimension199 # eg. [ (1,2), (1,3), (5,6) ] -> ( (1,1,5), (2,3,6) )200 logits[tuple(zip(*logits_to_penalize))] = -float("inf")201 202 # verifies that no rows in logits were missed unexpectedly203 assert start_idx == logits.shape[0]204 return logits205 206 207def _apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,208 output_tokens_tensor: torch.Tensor,209 presence_penalties: torch.Tensor,210 frequency_penalties: torch.Tensor,211 repetition_penalties: torch.Tensor) -> torch.Tensor:212 num_seqs, vocab_size = logits.shape213 _, prompt_mask = _get_bin_counts_and_mask(prompt_tokens_tensor, vocab_size,214 num_seqs)215 output_bin_counts, output_mask = _get_bin_counts_and_mask(216 output_tokens_tensor, vocab_size, num_seqs)217 218 repetition_penalties = repetition_penalties[:, None].repeat(1, vocab_size)219 repetition_penalties[~(prompt_mask | output_mask)] = 1.0220 logits = torch.where(logits > 0, logits / repetition_penalties,221 logits * repetition_penalties)222 223 # We follow the definition in OpenAI API.224 # Refer to https://platform.openai.com/docs/api-reference/parameter-details225 logits -= frequency_penalties.unsqueeze_(dim=1) * output_bin_counts226 logits -= presence_penalties.unsqueeze_(dim=1) * output_mask227 return logits228 229 230def _apply_top_k_top_p(231 logits: torch.Tensor,232 p: torch.Tensor,233 k: torch.Tensor,234) -> torch.Tensor:235 logits_sort, logits_idx = logits.sort(dim=-1, descending=False)236 237 # Apply top-k.238 top_k_mask = logits_sort.size(1) - k.to(torch.long)239 # Get all the top_k values.240 top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1))241 top_k_mask = logits_sort < top_k_mask242 logits_sort.masked_fill_(top_k_mask, -float("inf"))243 244 # Apply top-p.245 probs_sort = logits_sort.softmax(dim=-1)246 probs_sum = probs_sort.cumsum(dim=-1)247 top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)248 # at least one249 top_p_mask[:, -1] = False250 logits_sort.masked_fill_(top_p_mask, -float("inf"))251 252 # Re-sort the probabilities.253 src = torch.arange(logits_idx.shape[-1],254 device=logits_idx.device).expand_as(logits_idx)255 logits_idx_inv = torch.empty_like(logits_idx).scatter_(dim=-1,256 index=logits_idx,257 src=src)258 logits = torch.gather(logits_sort, dim=-1, index=logits_idx_inv)259 return logits260 261 262def _apply_min_p(263 logits: torch.Tensor,264 min_p: torch.Tensor,265) -> torch.Tensor:266 """267 Adapted from268 https://github.com/oobabooga/text-generation-webui/blob/3146124ec01f02c8fb1650a6517cf1b60b537aaf/modules/sampler_hijack.py#L16C17-L16C17269 """270 probs = torch.softmax(logits, dim=-1)271 top_probs, _ = probs.max(dim=-1, keepdim=True)272 scaled_min_p = min_p.unsqueeze_(dim=1) * top_probs273 tokens_to_remove = probs < scaled_min_p274 logits = logits.masked_fill_(tokens_to_remove, -float("inf"))275 276 return logits277 278 279def _greedy_sample(280 selected_seq_groups: List[Tuple[List[int], SamplingParams]],281 samples: torch.Tensor,282) -> List[Tuple[List[int], List[int]]]:283 samples = samples.tolist()284 sample_idx = 0285 results = []286 for seq_group in selected_seq_groups:287 seq_ids, _ = seq_group288 num_parent_seqs = len(seq_ids)289 assert num_parent_seqs == 1, (290 "Greedy sampling should have only one seq.")291 parent_ids = list(range(num_parent_seqs))292 next_token_ids = [samples[sample_idx]]293 results.append((next_token_ids, parent_ids))294 sample_idx += num_parent_seqs295 return results296 297 298def _random_sample(299 selected_seq_groups: List[Tuple[List[int], SamplingParams]],300 is_prompts: List[bool],301 random_samples: torch.Tensor,302) -> List[Tuple[List[int], List[int]]]:303 # Find the maximum best_of value of the prompt phase requests.304 random_samples = random_samples.cpu()305 sample_idx = 0306 results = []307 for seq_group, is_prompt in zip(selected_seq_groups, is_prompts):308 seq_ids, sampling_params = seq_group309 num_parent_seqs = len(seq_ids)310 if is_prompt:311 # Prompt phase.312 parent_ids = [0] * sampling_params.best_of313 next_token_ids = random_samples[314 sample_idx, :sampling_params.best_of].tolist()315 else:316 # Generation phase.317 parent_ids = list(range(num_parent_seqs))318 next_token_ids = random_samples[sample_idx:sample_idx +319 num_parent_seqs, 0].tolist()320 results.append((next_token_ids, parent_ids))321 sample_idx += num_parent_seqs322 return results323 324 325def _beam_search_sample(326 selected_seq_groups: List[Tuple[List[int], SamplingParams]],327 is_prompts: List[bool],328 seq_data: Dict[int, SequenceData],329 logprobs: torch.Tensor,330) -> List[Tuple[List[int], List[int]]]:331 # We sample 2 * beam_width candidates to make sure that with high332 # probability we can get `beam_width` candidates in addition to333 # the finished sequences for the next iteration. See334 # https://github.com/tensorflow/tensor2tensor/blob/bafdc1b67730430d38d6ab802cbd51f9d053ba2e/tensor2tensor/utils/beam_search.py#L557-L563335 # for details. See also HF reference:336 # https://github.com/huggingface/transformers/blob/a4dd53d88e4852f023332d284ff07a01afcd5681/src/transformers/generation/utils.py#L3063-L3065337 #338 # NOTE: Beam search is not vectorized, so its speed can be slower than339 # other sampling methods.340 sample_idx = 0341 results = []342 for seq_group, is_prompt in zip(selected_seq_groups, is_prompts):343 seq_ids, sampling_params = seq_group344 num_parent_seqs = len(seq_ids)345 beam_width = sampling_params.best_of346 seq_group_logprobs = logprobs[sample_idx:sample_idx + num_parent_seqs]347 if is_prompt:348 # Prompt phase.349 assert num_parent_seqs == 1, (350 "Prompt input should have only one seq.")351 parent_ids = [0] * (2 * beam_width)352 _, next_token_ids = torch.topk(seq_group_logprobs[0],353 2 * beam_width)354 next_token_ids = next_token_ids.tolist()355 else:356 # Generation phase.357 cumulative_logprobs = [358 seq_data[seq_id].cumulative_logprob for seq_id in seq_ids359 ]360 cumulative_logprobs = torch.tensor(361 cumulative_logprobs,362 dtype=torch.float,363 device=seq_group_logprobs.device)364 seq_group_logprobs = (seq_group_logprobs +365 cumulative_logprobs.unsqueeze(dim=1))366 _, topk_ids = torch.topk(seq_group_logprobs.flatten(),367 2 * beam_width)368 topk_ids = topk_ids.tolist()369 vocab_size = seq_group_logprobs.size(-1)370 parent_ids = [i // vocab_size for i in topk_ids]371 next_token_ids = [i % vocab_size for i in topk_ids]372 results.append((next_token_ids, parent_ids))373 sample_idx += num_parent_seqs374 assert sample_idx == logprobs.size(0)375 return results376 377 378# torch.multinomial forces a GPU<->CPU sync.379# Therefore, we use an optimized implementation instead.380# Note that we always sample with replacement.381# probs will be modified in place, but this is fine, as we pass382# in a copy already.383def _multinomial(384 probs: torch.Tensor,385 num_samples: int,386 seq_groups: Optional[List[Tuple[List[int], SamplingParams]]] = None,387 generators: Optional[List[torch.Generator]] = None,388) -> torch.Tensor:389 if num_samples > 1:390 # This is equivalent to torch.repeat_interleaved (which also391 # forces a GPU<->CPU sync).392 # This allows us to do sampling with replacement by creating393 # num_samples copies of each row in the tensor, and then394 # batch sampling the resulting tensor.395 probs = probs[:, None, :].expand(probs.shape[0], num_samples,396 probs.shape[1]).contiguous().view(397 -1, probs.shape[1])398 q = torch.empty_like(probs)399 if seq_groups is None:400 q.exponential_()401 else:402 sample_idx = 0403 for (seq_ids, _), generator in zip(seq_groups, generators):404 next_sample_idx = sample_idx + len(seq_ids) * num_samples405 q[sample_idx:next_sample_idx].exponential_(generator=generator)406 sample_idx = next_sample_idx407 return probs.div_(q).argmax(dim=1).view(-1, num_samples)408 409 410def _sample_with_torch(411 probs: torch.Tensor,412 logprobs: torch.Tensor,413 sampling_metadata: SamplingMetadata,414 include_gpu_probs_tensor: bool,415 modify_greedy_probs: bool,416) -> Tuple[List[Tuple[List[int], List[int]]], Optional[torch.Tensor]]:417 categorized_seq_group_ids = {t: [] for t in SamplingType}418 categorized_sample_indices = sampling_metadata.categorized_sample_indices419 for i, seq_group in enumerate(sampling_metadata.seq_groups):420 _, sampling_params = seq_group421 sampling_type = sampling_params.sampling_type422 categorized_seq_group_ids[sampling_type].append(i)423 424 sample_results_dict: Dict[int, Tuple[List[int], List[int]]] = {}425 sample_metadata = {}426 multinomial_samples = {}427 428 # Create output tensor for sampled token ids.429 if include_gpu_probs_tensor:430 sampled_token_ids_tensor = torch.empty(logprobs.shape[0],431 1,432 dtype=torch.long,433 device=logprobs.device)434 else:435 sampled_token_ids_tensor = None436 437 # Counterintiutively, having two loops here is actually faster.438 # The first loop can run without waiting on GPU<->CPU sync.439 for sampling_type in SamplingType:440 sample_indices = categorized_sample_indices[sampling_type][:, 0]441 num_tokens = len(sample_indices)442 if num_tokens == 0:443 continue444 seq_group_ids = categorized_seq_group_ids[sampling_type]445 seq_groups = [sampling_metadata.seq_groups[i] for i in seq_group_ids]446 is_prompts = [i < sampling_metadata.num_prompts for i in seq_group_ids]447 sample_metadata[sampling_type] = (seq_group_ids, seq_groups,448 is_prompts, sample_indices)449 long_sample_indices = sample_indices.long()450 451 if sampling_type == SamplingType.GREEDY:452 greedy_samples = torch.argmax(logprobs[long_sample_indices],453 dim=-1)454 455 if include_gpu_probs_tensor:456 # Store sampled tokens in output tensor.457 sampled_token_ids_tensor[458 long_sample_indices] = greedy_samples.unsqueeze(-1)459 460 if modify_greedy_probs:461 # If required, modify the probabilities such that sampling from462 # the modified distribution would always sample the argmax463 # token id.464 _modify_greedy_probs_inplace(logprobs, probs,465 long_sample_indices,466 greedy_samples)467 468 elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED):469 max_best_of_in_batch = 1470 for seq_group, is_prompt in zip(seq_groups, is_prompts):471 if is_prompt:472 _, sampling_params = seq_group473 max_best_of_in_batch = max(max_best_of_in_batch,474 sampling_params.best_of)475 seeded_args = {} if sampling_type == SamplingType.RANDOM else {476 "seq_groups": seq_groups,477 "generators": sampling_metadata.generators,478 }479 480 multinomial_samples[sampling_type] = _multinomial(481 probs[long_sample_indices], max_best_of_in_batch,482 **seeded_args)483 484 if include_gpu_probs_tensor:485 # Store sampled tokens in output tensor.486 sampled_token_ids_tensor[487 long_sample_indices] = multinomial_samples[sampling_type]488 489 elif sampling_type == SamplingType.BEAM:490 beam_search_logprobs = logprobs[sample_indices]491 else:492 raise ValueError(f"Unsupported sampling type: {sampling_type}")493 494 # GPU<->CPU sync happens in the loop below.495 # This also converts the sample output to Python objects.496 497 for sampling_type in SamplingType:498 if sampling_type not in sample_metadata:499 continue500 seq_group_ids, seq_groups, is_prompts, sample_indices = sample_metadata[501 sampling_type]502 if sampling_type == SamplingType.GREEDY:503 sample_results = _greedy_sample(seq_groups, greedy_samples)504 elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED):505 sample_results = _random_sample(seq_groups, is_prompts,506 multinomial_samples[sampling_type])507 elif sampling_type == SamplingType.BEAM:508 sample_results = _beam_search_sample(seq_groups, is_prompts,509 sampling_metadata.seq_data,510 beam_search_logprobs)511 sample_results_dict.update(zip(seq_group_ids, sample_results))512 513 sample_results = [514 sample_results_dict[i]515 for i in range(len(sampling_metadata.seq_groups))516 ]517 return sample_results, sampled_token_ids_tensor518 519 520def _sample_with_triton_kernel(521 probs: torch.Tensor,522 logprobs: torch.Tensor,523 sampling_metadata: SamplingMetadata,524 sampling_tensors: SamplingTensors,525) -> List[Tuple[List[int], List[int]]]:526 categorized_seq_group_ids = {t: [] for t in SamplingType}527 categorized_sample_indices = sampling_metadata.categorized_sample_indices528 for i, seq_group in enumerate(sampling_metadata.seq_groups):529 _, sampling_params = seq_group530 sampling_type = sampling_params.sampling_type531 categorized_seq_group_ids[sampling_type].append(i)532 533 sample_results_dict: Dict[int, Tuple[List[int], List[int]]] = {}534 sample_metadata = {}535 max_best_of_in_batch = 1536 537 # Counterintiutively, having two loops here is actually faster.538 # The first loop can run without waiting on GPU<->CPU sync.539 for sampling_type in SamplingType:540 sample_indices = categorized_sample_indices[sampling_type][:, 0]541 sampled_token_indices = categorized_sample_indices[sampling_type][:, 1]542 num_tokens = len(sample_indices)543 if num_tokens == 0:544 continue545 seq_group_ids = categorized_seq_group_ids[sampling_type]546 seq_groups = [sampling_metadata.seq_groups[i] for i in seq_group_ids]547 is_prompts = [i < sampling_metadata.num_prompts for i in seq_group_ids]548 sample_metadata[sampling_type] = (seq_group_ids, seq_groups,549 is_prompts, sample_indices,550 sampled_token_indices)551 if sampling_type in (SamplingType.GREEDY, SamplingType.RANDOM,552 SamplingType.RANDOM_SEED):553 for seq_group, is_prompt in zip(seq_groups, is_prompts):554 if is_prompt:555 _, sampling_params = seq_group556 max_best_of_in_batch = max(max_best_of_in_batch,557 sampling_params.best_of)558 elif sampling_type == SamplingType.BEAM:559 beam_search_logprobs = logprobs[sample_indices]560 else:561 raise ValueError(f"Unsupported sampling type: {sampling_type}")562 563 sampled_tokens, _, _ = sample_triton(564 probs=probs,565 seeds=sampling_tensors.sampling_seeds,566 max_best_of=max_best_of_in_batch,567 sample_indices=sampling_tensors.sample_indices,568 logprobs=logprobs,569 # don't save logprobs because we have logic for that below570 # TODO: use this instead of the CPU-based logic below571 save_logprobs=False,572 )573 574 # GPU<->CPU sync happens in the loop below.575 576 for sampling_type in SamplingType:577 if sampling_type not in sample_metadata:578 continue579 (seq_group_ids, seq_groups, is_prompts, sample_indices,580 sampled_token_indices) = sample_metadata[sampling_type]581 if sampling_type == SamplingType.GREEDY:582 sample_results = _greedy_sample(583 seq_groups, sampled_tokens[sampled_token_indices][:, 0])584 elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED):585 sample_results = _random_sample(586 seq_groups, is_prompts, sampled_tokens[sampled_token_indices])587 elif sampling_type == SamplingType.BEAM:588 sample_results = _beam_search_sample(seq_groups, is_prompts,589 sampling_metadata.seq_data,590 beam_search_logprobs)591 sample_results_dict.update(zip(seq_group_ids, sample_results))592 593 sample_results = [594 sample_results_dict[i]595 for i in range(len(sampling_metadata.seq_groups))596 ]597 return sample_results598 599 600def _sample(601 probs: torch.Tensor, logprobs: torch.Tensor,602 sampling_metadata: SamplingMetadata, sampling_tensors: SamplingTensors,603 include_gpu_probs_tensor: bool, modify_greedy_probs: bool604) -> Tuple[List[Tuple[List[int], List[int]]], Optional[torch.Tensor]]:605 return _sample_with_torch(606 probs,607 logprobs,608 sampling_metadata,609 include_gpu_probs_tensor=include_gpu_probs_tensor,610 modify_greedy_probs=modify_greedy_probs,611 )612 613 # TODO: Enable once Triton kernel & associated code is faster.614 # return _sample_with_triton_kernel(probs, logprobs, sampling_metadata,615 # sampling_tensors)616 617 618def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:619 """620 This function calculates the ranks of the chosen tokens in a logprob tensor.621 622 Args:623 x (torch.Tensor): 2D logprob tensor of shape (N, M)624 where N is the no. of tokens and M is the vocab dim.625 indices (torch.Tensor): List of chosen token indices.626 627 Returns:628 torch.Tensor: 1D tensor of shape (N,) where N is the no. of tokens.629 Each element in the returned tensor represents the rank 630 of the chosen token in the input logprob tensor.631 """632 vals = x[torch.arange(0, len(x), device=x.device, dtype=indices.dtype),633 indices]634 return (x > vals[:, None]).long().sum(1).add_(1)635 636 637def _get_logprobs(638 logprobs: torch.Tensor,639 sampling_metadata: SamplingMetadata,640 sample_results: List[Tuple[List[int], List[int]]],641) -> Tuple[List[Optional[List[Optional[Dict[int, float]]]]], List[List[Dict[642 int, float]]]]:643 # Prepare query indices644 batched_logprobs_query_seq_indices: List[int] = []645 batched_logprobs_query_token_indices: List[int] = []646 # at least get one logprob for each token647 largest_num_logprobs = 1648 sample_idx = 0649 for i, (seq_group, sample_result) in enumerate(650 zip(sampling_metadata.seq_groups, sample_results)):651 seq_ids, sampling_params = seq_group652 next_token_ids, parent_ids = sample_result653 num_parent_seqs = len(seq_ids)654 if (i < sampling_metadata.num_prompts655 and sampling_params.prompt_logprobs is not None):656 largest_num_logprobs = max(largest_num_logprobs,657 sampling_params.prompt_logprobs)658 prompt_len = sampling_metadata.prompt_lens[i]659 prompt_tokens = sampling_metadata.seq_data[660 seq_ids[0]].prompt_token_ids661 batched_logprobs_query_seq_indices.extend(662 sample_idx + j for j in range(prompt_len - 1))663 batched_logprobs_query_token_indices.extend(664 token_id for token_id in prompt_tokens[1:])665 sample_idx += prompt_len - 1666 batched_logprobs_query_seq_indices.extend(667 [sample_idx + parent_id for parent_id in parent_ids])668 batched_logprobs_query_token_indices.extend(next_token_ids)669 if sampling_params.logprobs is not None:670 largest_num_logprobs = max(largest_num_logprobs,671 sampling_params.logprobs)672 sample_idx += num_parent_seqs673 assert sample_idx == logprobs.size(0)674 675 batched_logprobs_query_seq_indices_gpu = torch.tensor(676 batched_logprobs_query_seq_indices, device=logprobs.device)677 batched_logprobs_query_token_indices_gpu = torch.tensor(678 batched_logprobs_query_token_indices, device=logprobs.device)679 680 # Batched query for logprobs of selected token681 batched_logprobs_query_result = logprobs[[682 batched_logprobs_query_seq_indices_gpu,683 batched_logprobs_query_token_indices_gpu684 ]]685 686 batched_ranks_query_result = _get_ranks(687 logprobs[batched_logprobs_query_seq_indices_gpu],688 batched_logprobs_query_token_indices_gpu)689 690 # Batched query for logprobs of topk tokens691 if largest_num_logprobs > 0:692 top_logprobs, top_token_ids = torch.topk(logprobs,693 largest_num_logprobs,694 dim=-1)695 top_logprobs = top_logprobs.cpu()696 top_token_ids = top_token_ids.cpu()697 else:698 top_logprobs, top_token_ids = None, None699 700 batched_logprobs_query_result = batched_logprobs_query_result.cpu()701 batched_ranks_query_result = batched_ranks_query_result.cpu()702 703 # Gather results704 result_prompt_logprobs: List[Optional[PromptLogprobs]] = []705 result_sample_logprobs: List[SampleLogprobs] = []706 sample_idx = 0707 query_result_idx = 0708 for i, (seq_group, sample_result) in enumerate(709 zip(sampling_metadata.seq_groups, sample_results)):710 seq_ids, sampling_params = seq_group711 next_token_ids, parent_ids = sample_result712 713 # Prompt logprobs714 if (i < sampling_metadata.num_prompts715 and sampling_params.prompt_logprobs is not None):716 num_logprobs = sampling_params.prompt_logprobs717 prompt_tokens = sampling_metadata.seq_data[718 seq_ids[0]].prompt_token_ids719 group_prompt_logprobs: PromptLogprobs = [None]720 for token_id in prompt_tokens[1:]:721 prompt_logprobs_dict = {722 token_id:723 (batched_logprobs_query_result[query_result_idx].item(),724 batched_ranks_query_result[query_result_idx].item())725 }726 if num_logprobs > 0:727 prompt_logprobs_dict.update(728 zip(729 top_token_ids[sample_idx, :num_logprobs].tolist(),730 zip(731 top_logprobs[732 sample_idx, :num_logprobs].tolist(),733 range(1, num_logprobs + 1))))734 group_prompt_logprobs.append({735 token_id: Logprob(*logprob_rank)736 for token_id, logprob_rank in prompt_logprobs_dict.items()737 })738 sample_idx += 1739 query_result_idx += 1740 result_prompt_logprobs.append(group_prompt_logprobs)741 else:742 result_prompt_logprobs.append(None)743 744 # Sample logprobs745 num_logprobs = sampling_params.logprobs746 if num_logprobs is None:747 num_logprobs = 0748 group_sample_logprobs: SampleLogprobs = []749 for next_token_id, parent_id in zip(next_token_ids, parent_ids):750 sample_logprobs_dict = {751 next_token_id:752 (batched_logprobs_query_result[query_result_idx].item(),753 batched_ranks_query_result[query_result_idx].item())754 }755 query_result_idx += 1756 if num_logprobs >= 0:757 sample_logprobs_dict.update(758 zip(759 top_token_ids[sample_idx +760 parent_id, :num_logprobs].tolist(),761 zip(762 top_logprobs[sample_idx +763 parent_id, :num_logprobs].tolist(),764 range(1, num_logprobs + 1))))765 group_sample_logprobs.append({766 token_id: Logprob(*logprob_rank)767 for token_id, logprob_rank in sample_logprobs_dict.items()768 })769 result_sample_logprobs.append(group_sample_logprobs)770 sample_idx += len(seq_ids)771 772 return result_prompt_logprobs, result_sample_logprobs773 774 775def _modify_greedy_probs_inplace(logprobs: torch.Tensor, probs: torch.Tensor,776 sample_indices: torch.Tensor,777 greedy_samples: torch.Tensor) -> None:778 """Modify the probability distributions of the greedily-sampled tokens such779 that each sampled token has a "probability" of 1.0. This is required by780 speculative decoding, which depends on the sampling method being encoded781 within the probability distribution for correctness.782 783 # Why do we only need to do this for greedy sampling?784 785 vLLM's sampler performs the following steps for greedy or multinomial786 (random) sampling:787 1. Get logits from model.788 2. Modify logits according to per-sequence sampling parameters.789 - Multiply by temperature, top-k and top-p masking, penalize tokens790 according to their frequency, etc.791 3. Sample a token.792 - Random sampling simply samples from the modified probability793 distribution.794 - Greedy sampling performs `argmax` to obtain the token with the795 highest likelihood.796 797 Ignoring greedy sampling for a moment, we find that the computed probability798 distribution has the following property: we can sample from it independently799 and find that the token sampled by the Sampler has a frequency corresponding800 to how often we see it in our sampling. In other words, for tokens sampled801 with vLLM's random SamplingType, the computed probability distribution802 encodes the sampling methodology completely.803 804 Greedy sampling does not normally have this property. vLLM modifies logits805 according to sampling params, then performs `argmax`, then returns the806 sampled token and the computed probability distribution. If we sample from807 the distribution, we'll find the likelihood of the greedily-sampled token808 is not always 1.0.809 810 Since lossless speculative decoding requires that the sampling methodology811 be encoded within the probability distribution, we are motivated to modify812 the probability distribution such that the sampled token has probability 1813 when speculative decoding is used.814 815 NOTE: Alternatively, we could use an extremely low temperature to achieve816 greedy sampling using multinomial computation and unite the codepaths. This817 has implications on the overall design of the sampler, e.g. how to record818 accurate logprobs for the user, so this improvement is deferred to later.819 """820 logprobs[sample_indices, :] = -float('inf')821 logprobs[sample_indices, greedy_samples] = 0.0822 probs[sample_indices, :] = 0823 probs[sample_indices, greedy_samples] = 1.0824 825 826def _build_sampler_output(827 sample_results: List[Tuple[List[int], List[int]]],828 sampling_metadata: SamplingMetadata,829 prompt_logprobs: List[Optional[PromptLogprobs]],830 sample_logprobs: List[SampleLogprobs],831 on_device_tensors: Optional[Tuple[torch.Tensor, torch.Tensor]],832) -> SamplerOutput:833 """Construct Python objects with the output of sampling.834 835 Args:836 on_device_tensors: Tuple containing on-device tensors with the837 probabilities used in sampling and the sampled token ids. This838 allows post-processing without copies to CPU/serialization, e.g. in839 speculative decoding rejection sampling.840 """841 842 sampler_output = []843 for (seq_group, sample_result, group_prompt_logprobs,844 group_sample_logprobs) in zip(sampling_metadata.seq_groups,845 sample_results, prompt_logprobs,846 sample_logprobs):847 seq_ids, _ = seq_group848 next_token_ids, parent_ids = sample_result849 seq_outputs = []850 for parent_id, next_token_id, logprobs in zip(parent_ids,851 next_token_ids,852 group_sample_logprobs):853 seq_outputs.append(854 SequenceOutput(seq_ids[parent_id], next_token_id, logprobs))855 sampler_output.append(856 SequenceGroupOutput(seq_outputs, group_prompt_logprobs))857 858 # If not specified, store None values in SamplerOutput.859 if on_device_tensors is not None:860 sampled_token_probs, sampled_token_ids = on_device_tensors861 else:862 sampled_token_probs, sampled_token_ids = (None, None)863 864 return SamplerOutput(865 outputs=sampler_output,866 sampled_token_probs=sampled_token_probs,867 sampled_token_ids=sampled_token_ids,868 )869