| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """A layer that samples the next tokens from the model's outputs.""" |
| 4 | |
| 5 | import torch |
| 6 | import torch.nn as nn |
| 7 | |
| 8 | from vllm.config.model import LogprobsMode |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 11 | from vllm.utils.torch_utils import PIN_MEMORY |
| 12 | from vllm.v1.outputs import LogprobsTensors, SamplerOutput |
| 13 | from vllm.v1.sample.metadata import SamplingMetadata |
| 14 | from vllm.v1.sample.ops.bad_words import apply_bad_words |
| 15 | from vllm.v1.sample.ops.logprobs import batched_count_greater_than |
| 16 | from vllm.v1.sample.ops.penalties import apply_all_penalties |
| 17 | from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler |
| 18 | |
| 19 | _SAMPLING_EPS = 1e-5 |
| 20 | |
| 21 | |
| 22 | class Sampler(nn.Module): |
| 23 | """A layer that samples the next tokens from the model's outputs |
| 24 | with the following steps in order: |
| 25 | |
| 26 | 1. If logprobs are requested: |
| 27 | a) If `logprobs_mode` is `raw_logprobs`, compute logprobs |
| 28 | as the final logprobs to return. |
| 29 | b) If `logprobs_mode` is `raw_logits`, clone the logits |
| 30 | as the final logprobs to return. |
| 31 | 2. Convert logits to float32. |
| 32 | 3. Apply allowed token ids whitelist. |
| 33 | 4. Apply bad words exclusion. |
| 34 | 5. Apply logit processors which are not argmax-invariant, |
| 35 | i.e. that can impact greedy sampling. |
| 36 | a) Min tokens processor |
| 37 | b) Logit bias processor |
| 38 | 6. Apply penalties |
| 39 | a) Repetition penalty |
| 40 | b) Frequency penalty |
| 41 | c) Presence penalty |
| 42 | 7. Sample the next tokens. `sample` method performs the following steps: |
| 43 | a) If not `all_random`, perform greedy sampling. If `all_greedy`, |
| 44 | return the greedily sampled tokens and final logprobs if requested. |
| 45 | b) Apply temperature. |
| 46 | c) Apply logit processors which are argmax-invariant, by default |
| 47 | the min_p processor. |
| 48 | d) Apply top_k and/or top_p. |
| 49 | e) Sample the next tokens with the probability distribution. |
| 50 | f) If `all_random` or temperature >= epsilon (1e-5), return the |
| 51 | randomly sampled tokens and final logprobs if requested. Else, |
| 52 | return the greedily sampled tokens and logprobs if requested. |
| 53 | 8. Gather the logprobs of the top `max_num_logprobs` and sampled token |
| 54 | (if requested). Note that if the sampled token is within the top |
| 55 | `max_num_logprobs`, the logprob will be eventually merged in |
| 56 | `LogprobsProcessor` during output processing. Therefore, the |
| 57 | final output may contain either `max_num_logprobs + 1` or |
| 58 | `max_num_logprobs` logprobs. |
| 59 | 9. Return the final `SamplerOutput`. |
| 60 | """ |
| 61 | |
| 62 | def __init__( |
| 63 | self, |
| 64 | logprobs_mode: LogprobsMode = "raw_logprobs", |
| 65 | use_fp64_gumbel: bool = False, |
| 66 | ): |
| 67 | super().__init__() |
| 68 | self.topk_topp_sampler = TopKTopPSampler(logprobs_mode, use_fp64_gumbel) |
| 69 | self.pin_memory = PIN_MEMORY |
| 70 | self.logprobs_mode = logprobs_mode |
| 71 | self.use_fp64_gumbel = use_fp64_gumbel |
| 72 | |
| 73 | def forward( |
| 74 | self, |
| 75 | logits: torch.Tensor, |
| 76 | sampling_metadata: SamplingMetadata, |
| 77 | predict_bonus_token: bool = False, |
| 78 | logprobs_mode_override: LogprobsMode | None = None, |
| 79 | ) -> SamplerOutput: |
| 80 | logprobs_mode = logprobs_mode_override or self.logprobs_mode |
| 81 | # NOTE(woosuk): Use the original logits (before any penalties or |
| 82 | # temperature scaling) for the top-k logprobs. |
| 83 | # This is different from the V0 sampler, which uses the logits that |
| 84 | # is used for sampling (after penalties and temperature scaling). |
| 85 | num_logprobs = sampling_metadata.max_num_logprobs |
| 86 | raw_logprobs: torch.Tensor | None = None |
| 87 | if num_logprobs is not None or sampling_metadata.logprob_token_ids: |
| 88 | if logprobs_mode == "raw_logprobs": |
| 89 | raw_logprobs = self.compute_logprobs(logits) |
| 90 | elif logprobs_mode == "raw_logits": |
| 91 | if logits.dtype == torch.float32: |
| 92 | raw_logprobs = logits.clone() |
| 93 | else: |
| 94 | raw_logprobs = logits.to(torch.float32) |
| 95 | |
| 96 | # Use float32 for the logits. |
| 97 | logits = logits.to(torch.float32) |
| 98 | |
| 99 | logits = self.apply_logits_processors( |
| 100 | logits, sampling_metadata, predict_bonus_token |
| 101 | ) |
| 102 | # Sample the next token. |
| 103 | sampled, processed_logprobs = self.sample(logits, sampling_metadata) |
| 104 | if processed_logprobs is not None: |
| 105 | raw_logprobs = processed_logprobs |
| 106 | # Convert sampled token ids to int64 (long) type to ensure compatibility |
| 107 | # with subsequent operations that may use these values as indices. |
| 108 | # This conversion is necessary because FlashInfer sampling operations |
| 109 | # return int32 (while PyTorch argmax and topk return int64). |
| 110 | sampled = sampled.long() |
| 111 | |
| 112 | # Handle logprob_token_ids if specified (more efficient than full vocab) |
| 113 | # This is used by generative_scoring API to get logprobs for specific tokens |
| 114 | logprob_token_ids_tensors = None |
| 115 | if sampling_metadata.logprob_token_ids: |
| 116 | assert raw_logprobs is not None |
| 117 | logprob_token_ids_tensors = self.gather_specific_token_logprobs( |
| 118 | raw_logprobs, sampling_metadata.logprob_token_ids, sampled |
| 119 | ) |
| 120 | |
| 121 | if num_logprobs is None: |
| 122 | logprobs_tensors = logprob_token_ids_tensors |
| 123 | elif num_logprobs == -1: |
| 124 | # Return the full unsorted and unranked logprobs. |
| 125 | logprobs_tensors = LogprobsTensors( |
| 126 | torch.empty(0), raw_logprobs, torch.empty(0) |
| 127 | ) |
| 128 | else: |
| 129 | # Gather the logprobs and ranks of the topk and sampled token. |
| 130 | logprobs_tensors = self.gather_logprobs( |
| 131 | raw_logprobs, num_logprobs, token_ids=sampled |
| 132 | ) |
| 133 | |
| 134 | # If we have both num_logprobs and logprob_token_ids, prefer |
| 135 | # logprob_token_ids as it's more specific |
| 136 | if logprob_token_ids_tensors is not None and num_logprobs is not None: |
| 137 | logprobs_tensors = logprob_token_ids_tensors |
| 138 | |
| 139 | # Use int32 to reduce the tensor size. |
| 140 | sampled = sampled.to(torch.int32) |
| 141 | |
| 142 | # These are GPU tensors. |
| 143 | sampler_output = SamplerOutput( |
| 144 | # The sampled tokens are expanded to 2D tensor with shape |
| 145 | # [num_requests, 1], where each row represents one generated |
| 146 | # token per request. |
| 147 | sampled_token_ids=sampled.unsqueeze(-1), |
| 148 | logprobs_tensors=logprobs_tensors, |
| 149 | ) |
| 150 | return sampler_output |
| 151 | |
| 152 | def gather_specific_token_logprobs( |
| 153 | self, |
| 154 | logprobs: torch.Tensor, |
| 155 | logprob_token_ids: dict[int, list[int]], |
| 156 | sampled: torch.Tensor, |
| 157 | ) -> LogprobsTensors | None: |
| 158 | """Gather logprobs for specific token IDs requested per request. |
| 159 | |
| 160 | Used by the generative_scoring API to return logprobs for an explicit |
| 161 | set of token ids rather than the top-k. Handles heterogeneous token |
| 162 | id lists across requests by padding shorter lists to the max length. |
| 163 | |
| 164 | Args: |
| 165 | logprobs: [batch_size, vocab_size] tensor of (raw) logprobs to |
| 166 | gather from. |
| 167 | logprob_token_ids: dict mapping req_index -> list of token IDs |
| 168 | sampled: [batch_size] tensor of sampled token IDs |
| 169 | |
| 170 | Returns: |
| 171 | LogprobsTensors with logprobs for the specified tokens, or None |
| 172 | if no requests have logprob_token_ids. |
| 173 | |
| 174 | """ |
| 175 | if not logprob_token_ids: |
| 176 | return None |
| 177 | |
| 178 | batch_size = logprobs.shape[0] |
| 179 | device = logprobs.device |
| 180 | |
| 181 | # Find max number of tokens across all requests |
| 182 | max_num_tokens = max(len(tids) for tids in logprob_token_ids.values()) |
| 183 | pin = self.pin_memory |
| 184 | |
| 185 | # Build the padded token_ids and valid_mask matrices on pinned CPU, |
| 186 | # then upload non-blocking. |
| 187 | token_ids_cpu = torch.zeros( |
| 188 | batch_size, max_num_tokens + 1, dtype=torch.int64, pin_memory=pin |
| 189 | ) |
| 190 | # Create mask for valid positions (True = valid, False = padded) |
| 191 | valid_mask_cpu = torch.zeros( |
| 192 | batch_size, max_num_tokens + 1, dtype=torch.bool, pin_memory=pin |
| 193 | ) |
| 194 | valid_mask_cpu[:, 0] = True # Sampled token is always valid |
| 195 | for req_idx, token_ids in logprob_token_ids.items(): |
| 196 | num_tokens = len(token_ids) |
| 197 | token_ids_cpu[req_idx, 1 : num_tokens + 1] = torch.as_tensor( |
| 198 | token_ids, dtype=torch.int64 |
| 199 | ) |
| 200 | valid_mask_cpu[req_idx, 1 : num_tokens + 1] = True |
| 201 | |
| 202 | token_ids_tensor = token_ids_cpu.to(device, non_blocking=True) |
| 203 | valid_mask = valid_mask_cpu.to(device, non_blocking=True) |
| 204 | # Sampled token in column 0 — fill on-device from the sampled GPU |
| 205 | # tensor so we don't need to D2H + re-upload. |
| 206 | token_ids_tensor[:, 0] = sampled |
| 207 | |
| 208 | # Gather logprobs at the requested token ids. |
| 209 | gathered_logprobs = logprobs.gather(-1, token_ids_tensor) |
| 210 | |
| 211 | # Mask invalid (padded) positions with -inf |
| 212 | gathered_logprobs = gathered_logprobs.masked_fill(~valid_mask, float("-inf")) |
| 213 | |
| 214 | # Compute ranks for the sampled token. log_softmax is monotonic w.r.t. |
| 215 | # the original logits, so ranks computed from logprobs are equivalent. |
| 216 | sampled_logprobs = logprobs.gather(-1, sampled.unsqueeze(-1)) |
| 217 | # Avoid 0/1 specialization recompile on the batch dimension of the |
| 218 | # compiled batched_count_greater_than. See gather_logprobs for context. |
| 219 | torch._dynamo.decorators.mark_unbacked(logprobs, 0) |
| 220 | torch._dynamo.decorators.mark_unbacked(sampled_logprobs, 0) |
| 221 | token_ranks = batched_count_greater_than(logprobs, sampled_logprobs) |
| 222 | |
| 223 | return LogprobsTensors( |
| 224 | logprob_token_ids=token_ids_tensor.to(torch.int32), |
| 225 | logprobs=gathered_logprobs, |
| 226 | selected_token_ranks=token_ranks, |
| 227 | ) |
| 228 | |
| 229 | @staticmethod |
| 230 | def apply_temperature( |
| 231 | logits: torch.Tensor, |
| 232 | temp: torch.Tensor, |
| 233 | all_random: bool, |
| 234 | ) -> torch.Tensor: |
| 235 | # Use in-place division to avoid creating a new tensor. |
| 236 | # Avoid division by zero if there are greedy requests. |
| 237 | if not all_random: |
| 238 | temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp) |
| 239 | return logits.div_(temp.unsqueeze(dim=1)) |
| 240 | |
| 241 | @staticmethod |
| 242 | def greedy_sample(logits: torch.Tensor) -> torch.Tensor: |
| 243 | if current_platform.is_cpu(): |
| 244 | return torch.ops._C.greedy_argmax(logits) |
| 245 | return logits.argmax(dim=-1).view(-1) |
| 246 | |
| 247 | def sample( |
| 248 | self, |
| 249 | logits: torch.Tensor, |
| 250 | sampling_metadata: SamplingMetadata, |
| 251 | logprobs_mode_override: LogprobsMode | None = None, |
| 252 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 253 | """Sample logits based on sampling metadata. |
| 254 | |
| 255 | The various logits processing functions called in this method |
| 256 | may update the logits tensor in-place. |
| 257 | """ |
| 258 | logprobs_mode = logprobs_mode_override or self.logprobs_mode |
| 259 | assert not (sampling_metadata.all_greedy and sampling_metadata.all_random) |
| 260 | if sampling_metadata.all_random: |
| 261 | greedy_sampled = None |
| 262 | else: |
| 263 | greedy_sampled = self.greedy_sample(logits) |
| 264 | if sampling_metadata.all_greedy: |
| 265 | processed_logprobs = None |
| 266 | if ( |
| 267 | sampling_metadata.max_num_logprobs is not None |
| 268 | or sampling_metadata.logprob_token_ids |
| 269 | ): |
| 270 | if logprobs_mode == "processed_logits": |
| 271 | processed_logprobs = logits |
| 272 | elif logprobs_mode == "processed_logprobs": |
| 273 | processed_logprobs = self.compute_logprobs(logits) |
| 274 | return greedy_sampled, processed_logprobs |
| 275 | |
| 276 | assert sampling_metadata.temperature is not None |
| 277 | |
| 278 | # Apply temperature. |
| 279 | logits = self.apply_temperature( |
| 280 | logits, sampling_metadata.temperature, sampling_metadata.all_random |
| 281 | ) |
| 282 | |
| 283 | # Apply logits processors that only apply to random sampling |
| 284 | # (argmax invariant) |
| 285 | for processor in sampling_metadata.logitsprocs.argmax_invariant: |
| 286 | logits = processor.apply(logits) |
| 287 | |
| 288 | # Apply top_k and/or top_p. |
| 289 | random_sampled, processed_logprobs = self.topk_topp_sampler( |
| 290 | logits, |
| 291 | sampling_metadata.generators, |
| 292 | sampling_metadata.top_k, |
| 293 | sampling_metadata.top_p, |
| 294 | ) |
| 295 | |
| 296 | if greedy_sampled is None: |
| 297 | return random_sampled, processed_logprobs |
| 298 | |
| 299 | sampled = torch.where( |
| 300 | sampling_metadata.temperature < _SAMPLING_EPS, |
| 301 | greedy_sampled, |
| 302 | random_sampled, |
| 303 | out=greedy_sampled, # Reuse tensor |
| 304 | ) |
| 305 | return sampled, processed_logprobs |
| 306 | |
| 307 | @staticmethod |
| 308 | def compute_logprobs(logits: torch.Tensor) -> torch.Tensor: |
| 309 | return logits.log_softmax(dim=-1, dtype=torch.float32) |
| 310 | |
| 311 | @staticmethod |
| 312 | def gather_logprobs( |
| 313 | logprobs: torch.Tensor, |
| 314 | num_logprobs: int, |
| 315 | token_ids: torch.Tensor, |
| 316 | ) -> LogprobsTensors: |
| 317 | """Gather logprobs for topk and sampled/prompt token. |
| 318 | |
| 319 | Args: |
| 320 | logprobs: (num tokens) x (vocab) tensor |
| 321 | num_logprobs: maximum number of logprobs to |
| 322 | retain per token |
| 323 | token_ids: prompt tokens (if prompt logprobs) |
| 324 | or sampled tokens (if sampled |
| 325 | logprobs); 1D token ID tensor |
| 326 | with (num tokens) elements |
| 327 | Must be int64. |
| 328 | |
| 329 | Returns: |
| 330 | Top-k int indices tensor, (num tokens) x (num_logprobs + 1) |
| 331 | Top-k float logprobs tensor, (num tokens) x (num_logprobs + 1) |
| 332 | Sampled token rank tensor, (num tokens) |
| 333 | |
| 334 | """ |
| 335 | assert token_ids.dtype == torch.int64 |
| 336 | # Find the topK values. |
| 337 | topk_logprobs, topk_indices = torch.topk(logprobs, num_logprobs, dim=-1) |
| 338 | |
| 339 | # Get with the logprob of the prompt or sampled token. |
| 340 | token_ids = token_ids.unsqueeze(-1) |
| 341 | token_logprobs = logprobs.gather(-1, token_ids) |
| 342 | |
| 343 | # Compute the ranks of the actual token. |
| 344 | # Avoid 0/1 specialization recompile on the batch dimension |
| 345 | # of the compiled batched_count_greater_than. mark_unbacked makes |
| 346 | # the size fully symbolic so dynamo doesn't specialize when |
| 347 | # batch_size transitions from 1 to >=2. |
| 348 | with gpu_sync_allowed(first_only=True): |
| 349 | torch._dynamo.decorators.mark_unbacked(logprobs, 0) |
| 350 | torch._dynamo.decorators.mark_unbacked(token_logprobs, 0) |
| 351 | token_ranks = batched_count_greater_than(logprobs, token_logprobs) |
| 352 | |
| 353 | # Concatenate together with the topk. |
| 354 | indices = torch.cat((token_ids, topk_indices), dim=1) |
| 355 | logprobs = torch.cat((token_logprobs, topk_logprobs), dim=1) |
| 356 | |
| 357 | # Use int32 to reduce the tensor size. |
| 358 | indices = indices.to(torch.int32) |
| 359 | |
| 360 | return LogprobsTensors(indices, logprobs, token_ranks) |
| 361 | |
| 362 | @staticmethod |
| 363 | def _combine_outputs_with_spec_tokens( |
| 364 | output_token_ids: list[list[int]], |
| 365 | spec_token_ids: list[list[int]] | None = None, |
| 366 | ) -> list[list[int]]: |
| 367 | if spec_token_ids is None: |
| 368 | return output_token_ids |
| 369 | |
| 370 | return [ |
| 371 | [*out, *spec] if spec else out |
| 372 | for out, spec in zip(output_token_ids, spec_token_ids) |
| 373 | ] |
| 374 | |
| 375 | def apply_logits_processors( |
| 376 | self, |
| 377 | logits: torch.Tensor, |
| 378 | sampling_metadata: SamplingMetadata, |
| 379 | predict_bonus_token: bool, |
| 380 | ) -> torch.Tensor: |
| 381 | bad_words_token_ids = sampling_metadata.bad_words_token_ids |
| 382 | any_penalties_or_bad_words = ( |
| 383 | bool(bad_words_token_ids) or not sampling_metadata.no_penalties |
| 384 | ) |
| 385 | output_token_ids = sampling_metadata.output_token_ids |
| 386 | if predict_bonus_token and any_penalties_or_bad_words: |
| 387 | # Combine base outputs with spec tokens when speculative decoding |
| 388 | # is enabled. |
| 389 | output_token_ids = self._combine_outputs_with_spec_tokens( |
| 390 | output_token_ids, |
| 391 | sampling_metadata.spec_token_ids, |
| 392 | ) |
| 393 | |
| 394 | # Apply allowed token ids. |
| 395 | if sampling_metadata.allowed_token_ids_mask is not None: |
| 396 | logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf")) |
| 397 | |
| 398 | # Apply bad words exclusion. |
| 399 | if bad_words_token_ids: |
| 400 | apply_bad_words(logits, bad_words_token_ids, output_token_ids) |
| 401 | |
| 402 | # Apply logits processors which can impact greedy sampling. |
| 403 | for processor in sampling_metadata.logitsprocs.non_argmax_invariant: |
| 404 | logits = processor.apply(logits) |
| 405 | |
| 406 | # Apply penalties (e.g., freq_penalties). |
| 407 | logits = self.apply_penalties(logits, sampling_metadata, output_token_ids) |
| 408 | holder = sampling_metadata.thinking_budget_state_holder |
| 409 | if holder is not None and holder.has_tracked_requests(): |
| 410 | # Committed outputs only; spec drafts live in ``spec_token_ids``. |
| 411 | holder.update_state( |
| 412 | sampling_metadata.output_token_ids, |
| 413 | sampling_metadata.spec_token_ids, |
| 414 | repeat_indices=None, |
| 415 | ) |
| 416 | logits = holder.apply_to_logits( |
| 417 | logits, |
| 418 | predict_bonus_token, |
| 419 | sampling_metadata.spec_token_ids, |
| 420 | ) |
| 421 | return logits |
| 422 | |
| 423 | @staticmethod |
| 424 | def apply_penalties( |
| 425 | logits: torch.Tensor, |
| 426 | sampling_metadata: SamplingMetadata, |
| 427 | output_token_ids: list[list[int]], |
| 428 | ) -> torch.Tensor: |
| 429 | if sampling_metadata.no_penalties: |
| 430 | return logits |
| 431 | |
| 432 | assert sampling_metadata.prompt_token_ids is not None |
| 433 | return apply_all_penalties( |
| 434 | logits, |
| 435 | sampling_metadata.prompt_token_ids, |
| 436 | sampling_metadata.presence_penalties, |
| 437 | sampling_metadata.frequency_penalties, |
| 438 | sampling_metadata.repetition_penalties, |
| 439 | output_token_ids, |
| 440 | ) |