| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn as nn |
| 6 | |
| 7 | from vllm import envs |
| 8 | from vllm._aiter_ops import rocm_aiter_ops |
| 9 | from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.platforms import CpuArchEnum, current_platform |
| 12 | from vllm.platforms.interface import DeviceCapability |
| 13 | from vllm.triton_utils import HAS_TRITON |
| 14 | |
| 15 | if HAS_TRITON: |
| 16 | from vllm.v1.sample.ops.topk_topp_triton import ( |
| 17 | _topk_topp, |
| 18 | _topp_split_mask, |
| 19 | _topp_split_stats, |
| 20 | _topp_split_step, |
| 21 | apply_top_k_top_p_triton, |
| 22 | ) |
| 23 | |
| 24 | logger = init_logger(__name__) |
| 25 | |
| 26 | |
| 27 | def register_top_k_top_p_warmups() -> None: |
| 28 | """Register every native accelerator sampling kernel used at runtime.""" |
| 29 | if HAS_TRITON and not current_platform.is_cpu(): |
| 30 | _topk_topp.register_warmup() |
| 31 | if current_platform.is_cuda_alike(): |
| 32 | _topp_split_stats.register_warmup() |
| 33 | _topp_split_step.register_warmup() |
| 34 | _topp_split_mask.register_warmup() |
| 35 | |
| 36 | |
| 37 | def _skip_aiter_sampler_on_gfx1250() -> bool: |
| 38 | # Lazy ROCm-only import; keeps arch detection out of import time on CUDA/CPU. |
| 39 | from vllm.platforms.rocm import on_gfx1250 |
| 40 | |
| 41 | return on_gfx1250() |
| 42 | |
| 43 | |
| 44 | def _flashinfer_jit_unsupported_reason(capability: DeviceCapability) -> str | None: |
| 45 | """Return why FlashInfer JIT codegen cannot target the current GPU, or |
| 46 | None if it can. |
| 47 | |
| 48 | FlashInfer swallows arch-detection errors when building its compilation |
| 49 | context (e.g. SM 12.x with a CUDA toolkit older than 12.9), leaving an |
| 50 | empty target-arch set that makes every JIT spec fail with a misleading |
| 51 | "requires sm75 or higher" error at first use — killing the engine during |
| 52 | startup profiling (https://github.com/vllm-project/vllm/issues/42393). |
| 53 | """ |
| 54 | try: |
| 55 | from flashinfer.jit.core import check_cuda_arch |
| 56 | except ImportError: |
| 57 | return None |
| 58 | try: |
| 59 | check_cuda_arch() |
| 60 | return None |
| 61 | except RuntimeError as e: |
| 62 | reason = str(e) |
| 63 | try: |
| 64 | # Re-derive the real error that FlashInfer swallowed during arch |
| 65 | # detection, e.g. "SM 12.x requires CUDA >= 12.9". |
| 66 | from flashinfer.compilation_context import CompilationContext |
| 67 | |
| 68 | CompilationContext._normalize_cuda_arch(capability.major, capability.minor) |
| 69 | except RuntimeError as e: |
| 70 | reason = str(e) |
| 71 | except Exception: |
| 72 | # FlashInfer internals changed; keep the original error message. |
| 73 | pass |
| 74 | return reason |
| 75 | |
| 76 | |
| 77 | def flashinfer_sampler_supported() -> bool: |
| 78 | """Decide whether FlashInfer's top-p/top-k sampler can be used. |
| 79 | |
| 80 | Returns False (with appropriate logging) when ``VLLM_USE_FLASHINFER_SAMPLER`` |
| 81 | is 0, when the platform isn't CUDA, when the GPU's compute capability is |
| 82 | unsupported, when the GPU has 16 or fewer SMs, or when FlashInfer cannot |
| 83 | JIT-compile for the current GPU/CUDA toolchain. Raises ``RuntimeError`` if |
| 84 | the user explicitly opted in via the env var but FlashInfer is unavailable. |
| 85 | |
| 86 | Assumes flashinfer is installed, as guaranteed by ``requirements/cuda.txt``; |
| 87 | otherwise importing the FlashInfer backend below raises ``ImportError``. |
| 88 | |
| 89 | Note: callers must additionally ensure ``logprobs_mode`` doesn't require |
| 90 | post-top-k/top-p logits/logprobs for any request whose logprobs will be |
| 91 | returned in this step, since FlashInfer doesn't expose those. |
| 92 | """ |
| 93 | if not current_platform.is_cuda(): |
| 94 | return False |
| 95 | if not envs.VLLM_USE_FLASHINFER_SAMPLER: |
| 96 | logger.info_once( |
| 97 | "FlashInfer top-p/top-k sampling disabled via " |
| 98 | "VLLM_USE_FLASHINFER_SAMPLER=0." |
| 99 | ) |
| 100 | return False |
| 101 | from vllm.v1.attention.backends.flashinfer import FlashInferBackend |
| 102 | |
| 103 | capability = current_platform.get_device_capability() |
| 104 | assert capability is not None |
| 105 | unsupported_reason: str | None = None |
| 106 | if not FlashInferBackend.supports_compute_capability(capability): |
| 107 | unsupported_reason = ( |
| 108 | f"unsupported compute capability {capability.as_version_str()}" |
| 109 | ) |
| 110 | elif ( |
| 111 | num_sms := current_platform.num_compute_units( |
| 112 | torch.accelerator.current_device_index() |
| 113 | ) |
| 114 | ) <= 16: |
| 115 | # FlashInfer 0.7+ rejects multi-CTA top-k masking on <=16 SMs because |
| 116 | # its cross-CTA software barrier cannot guarantee forward progress. |
| 117 | unsupported_reason = ( |
| 118 | f"top-k masking requires more than 16 SMs; device has {num_sms}" |
| 119 | ) |
| 120 | else: |
| 121 | unsupported_reason = _flashinfer_jit_unsupported_reason(capability) |
| 122 | |
| 123 | if unsupported_reason is None: |
| 124 | logger.info_once("Using FlashInfer for top-p & top-k sampling.", scope="global") |
| 125 | return True |
| 126 | if envs.is_set("VLLM_USE_FLASHINFER_SAMPLER"): |
| 127 | raise RuntimeError( |
| 128 | f"FlashInfer top-p/top-k sampling unavailable: {unsupported_reason}. " |
| 129 | "Unset VLLM_USE_FLASHINFER_SAMPLER=1." |
| 130 | ) |
| 131 | logger.warning_once( |
| 132 | "FlashInfer top-p/top-k sampling unavailable: %s; falling back. " |
| 133 | "Set VLLM_USE_FLASHINFER_SAMPLER=0 to silence.", |
| 134 | unsupported_reason, |
| 135 | ) |
| 136 | return False |
| 137 | |
| 138 | |
| 139 | def xpu_sampler_supported() -> bool: |
| 140 | """Decide whether the fused XPU top-k/top-p sampler kernel can be used. |
| 141 | |
| 142 | Returns False (with appropriate logging) when the platform isn't XPU, when |
| 143 | ``VLLM_XPU_USE_SAMPLER_KERNEL`` is 0. |
| 144 | |
| 145 | Note: callers must additionally ensure no request needs a per-request seed |
| 146 | or greedy sampling, since the kernel draws from the device's default |
| 147 | generator and always samples randomly. |
| 148 | """ |
| 149 | if not current_platform.is_xpu(): |
| 150 | return False |
| 151 | if not envs.VLLM_XPU_USE_SAMPLER_KERNEL: |
| 152 | logger.info_once( |
| 153 | "Fused XPU top-p/top-k sampling disabled via VLLM_XPU_USE_SAMPLER_KERNEL=0." |
| 154 | ) |
| 155 | return False |
| 156 | |
| 157 | logger.info_once("Using the fused XPU kernel for top-p & top-k sampling.") |
| 158 | return True |
| 159 | |
| 160 | |
| 161 | def xpu_sample( |
| 162 | logits: torch.Tensor, |
| 163 | k: torch.Tensor | None, |
| 164 | p: torch.Tensor | None, |
| 165 | logprobs_mode: LogprobsMode = "raw_logprobs", |
| 166 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 167 | """Sample from the logits using the fused XPU top-k/top-p kernel. |
| 168 | |
| 169 | Statistically equivalent to `random_sample`, but avoids sorting the vocab |
| 170 | and never materializes the probability tensor. The noise comes from the |
| 171 | device's default generator, so per-request generators aren't supported. |
| 172 | |
| 173 | Returns the sampled token ids and, for processed logprobs modes, the |
| 174 | post-top-k/top-p logits (or logprobs); otherwise None. |
| 175 | """ |
| 176 | logits = logits.to(dtype=torch.float32).contiguous() |
| 177 | sampled = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device) |
| 178 | logits_to_return = ( |
| 179 | torch.empty_like(logits) if logprobs_mode in PROCESSED_LOGPROBS_MODES else None |
| 180 | ) |
| 181 | |
| 182 | generator = torch.xpu.default_generators[logits.device.index] |
| 183 | state = generator.get_state() |
| 184 | seed, offset = state.view(torch.int64).tolist() |
| 185 | seeds = torch.tensor([seed, offset], dtype=torch.int64, device="cpu") |
| 186 | # The XPU kernel expects k as int64 (Long), but the input batch |
| 187 | # stores top_k as int32. Cast here to avoid dtype mismatch. |
| 188 | if k is not None: |
| 189 | k = k.to(torch.int64) |
| 190 | torch.ops.vllm.xpu_topk_topp_sampler( |
| 191 | sampled, logits_to_return, logits, k, p, logprobs_mode, seeds |
| 192 | ) |
| 193 | # The custom XPU sampler kernel consumes RNG values internally, so advance |
| 194 | # the default generator's offset to keep future draws deterministic. |
| 195 | # pytorch: offset must be multiple of 4 |
| 196 | offset = (offset + logits.numel() + 3) // 4 * 4 |
| 197 | state.view(torch.int64)[1] = offset |
| 198 | generator.set_state(state) |
| 199 | return sampled, logits_to_return |
| 200 | |
| 201 | |
| 202 | class TopKTopPSampler(nn.Module): |
| 203 | """Module that performs optional top-k and top-p filtering followed by |
| 204 | weighted random sampling of logits. |
| 205 | |
| 206 | Implementations may update the logits tensor in-place. |
| 207 | """ |
| 208 | |
| 209 | def __init__( |
| 210 | self, |
| 211 | logprobs_mode: LogprobsMode = "raw_logprobs", |
| 212 | use_fp64_gumbel: bool = False, |
| 213 | ) -> None: |
| 214 | super().__init__() |
| 215 | self.logprobs_mode = logprobs_mode |
| 216 | self.use_fp64_gumbel = use_fp64_gumbel |
| 217 | if current_platform.is_cuda(): |
| 218 | # FlashInfer doesn't expose post-top-k/top-p logits/logprobs, |
| 219 | # so it can't be used when the configured mode requires them. |
| 220 | can_use_flashinfer = ( |
| 221 | logprobs_mode not in PROCESSED_LOGPROBS_MODES |
| 222 | and flashinfer_sampler_supported() |
| 223 | ) |
| 224 | self.forward = ( |
| 225 | self.forward_cuda if can_use_flashinfer else self.forward_native |
| 226 | ) |
| 227 | elif current_platform.is_cpu(): |
| 228 | arch = current_platform.get_cpu_architecture() |
| 229 | # Fall back to native implementation for POWERPC and RISCV. |
| 230 | # On PowerPC argmax produces incorrect output with torch.compile. |
| 231 | # PR: https://github.com/vllm-project/vllm/pull/26987 |
| 232 | if arch in (CpuArchEnum.RISCV, CpuArchEnum.POWERPC): |
| 233 | self.forward = self.forward_native |
| 234 | else: |
| 235 | self.forward = self.forward_cpu |
| 236 | elif current_platform.is_xpu(): |
| 237 | if xpu_sampler_supported() and not envs.VLLM_BATCH_INVARIANT: |
| 238 | self.forward = self.forward_xpu |
| 239 | else: |
| 240 | if envs.VLLM_BATCH_INVARIANT: |
| 241 | logger.info_once( |
| 242 | "VLLM_BATCH_INVARIANT is enabled. Using the " |
| 243 | "PyTorch-native sampler on XPU." |
| 244 | ) |
| 245 | self.forward = self.forward_native |
| 246 | elif ( |
| 247 | logprobs_mode not in PROCESSED_LOGPROBS_MODES |
| 248 | and rocm_aiter_ops.is_enabled() |
| 249 | and not _skip_aiter_sampler_on_gfx1250() # TODO (JPVILLAM): Enable |
| 250 | ): |
| 251 | self.aiter_ops = None |
| 252 | self._aiter_ops_import_failed = False |
| 253 | logger.info_once( |
| 254 | "Using aiter sampler on ROCm (lazy import, sampling-only)." |
| 255 | ) |
| 256 | self.forward = self.forward_hip |
| 257 | else: |
| 258 | self.forward = self.forward_native |
| 259 | |
| 260 | # Every accelerator backend can fall back to native sampling at runtime. |
| 261 | register_top_k_top_p_warmups() |
| 262 | |
| 263 | def forward_native( |
| 264 | self, |
| 265 | logits: torch.Tensor, |
| 266 | generators: dict[int, torch.Generator], |
| 267 | k: torch.Tensor | None, |
| 268 | p: torch.Tensor | None, |
| 269 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 270 | """PyTorch-native implementation of top-k and top-p sampling. |
| 271 | |
| 272 | The logits tensor may be updated in-place. |
| 273 | """ |
| 274 | logits = apply_top_k_top_p(logits, k, p) |
| 275 | logits_to_return = None |
| 276 | if self.logprobs_mode == "processed_logits": |
| 277 | logits_to_return = logits |
| 278 | elif self.logprobs_mode == "processed_logprobs": |
| 279 | logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32) |
| 280 | probs = logits.softmax(dim=-1, dtype=torch.float32) |
| 281 | return ( |
| 282 | random_sample(probs, generators, self.use_fp64_gumbel), |
| 283 | logits_to_return, |
| 284 | ) |
| 285 | |
| 286 | def forward_cuda( |
| 287 | self, |
| 288 | logits: torch.Tensor, |
| 289 | generators: dict[int, torch.Generator], |
| 290 | k: torch.Tensor | None, |
| 291 | p: torch.Tensor | None, |
| 292 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 293 | """More optimized implementation for top-k and top-p sampling.""" |
| 294 | # Fall back to the PyTorch-native path when FlashInfer has nothing |
| 295 | # to do (no top-k / top-p filter) or when per-request generators |
| 296 | # are present (unsupported by FlashInfer 0.2.3+). |
| 297 | if (k is None and p is None) or generators: |
| 298 | if generators: |
| 299 | logger.debug_once( |
| 300 | "FlashInfer 0.2.3+ does not support " |
| 301 | "per-request generators. Falling back to " |
| 302 | "PyTorch-native implementation." |
| 303 | ) |
| 304 | return self.forward_native(logits, generators, k, p) |
| 305 | if self.use_fp64_gumbel: |
| 306 | return self.forward_native(logits, generators, k, p) |
| 307 | assert self.logprobs_mode not in PROCESSED_LOGPROBS_MODES, ( |
| 308 | "FlashInfer does not support returning logits/logprobs" |
| 309 | ) |
| 310 | # flashinfer sampling functions expect contiguous logits. |
| 311 | # In flex_attn/triton_attn fp32 inference, logits can be non-contiguous |
| 312 | # because of slicing operation in logits_processor. |
| 313 | return flashinfer_sample(logits.contiguous(), k, p, generators), None |
| 314 | |
| 315 | def forward_cpu( |
| 316 | self, |
| 317 | logits: torch.Tensor, |
| 318 | generators: dict[int, torch.Generator], |
| 319 | k: torch.Tensor | None, |
| 320 | p: torch.Tensor | None, |
| 321 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 322 | """Fused Gumbel-max sampling for CPU. |
| 323 | |
| 324 | Uses a precomputed Gumbel table + single-pass argmax over logits, |
| 325 | skipping softmax and intermediate allocations entirely. |
| 326 | Falls back to the native path when fp64 Gumbel noise is requested. |
| 327 | """ |
| 328 | logits = apply_top_k_top_p(logits, k, p) |
| 329 | logits_to_return = None |
| 330 | if self.logprobs_mode == "processed_logits": |
| 331 | logits_to_return = logits |
| 332 | elif self.logprobs_mode == "processed_logprobs": |
| 333 | logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32) |
| 334 | |
| 335 | if self.use_fp64_gumbel: |
| 336 | probs = logits.softmax(dim=-1, dtype=torch.float32) |
| 337 | q = empty_exponential_noise_like(probs, self.use_fp64_gumbel) |
| 338 | q.exponential_() |
| 339 | for i, generator in generators.items(): |
| 340 | q[i].exponential_(generator=generator) |
| 341 | return sample_with_exponential_noise(probs, q), logits_to_return |
| 342 | |
| 343 | batch_size = logits.shape[0] |
| 344 | seeds = torch.randint(0, 2**31, (batch_size,), dtype=torch.long) |
| 345 | for i, gen in generators.items(): |
| 346 | seeds[i] = torch.randint( |
| 347 | 0, 2**31, (1,), generator=gen, dtype=torch.long |
| 348 | ).item() |
| 349 | logits_f32 = logits.to(dtype=torch.float32) |
| 350 | return ( |
| 351 | torch.ops._C.fused_gumbel_argmax(logits_f32, seeds), |
| 352 | logits_to_return, |
| 353 | ) |
| 354 | |
| 355 | def _init_aiter_ops(self) -> bool: |
| 356 | if self._aiter_ops_import_failed: |
| 357 | return False |
| 358 | try: |
| 359 | import aiter.ops.sampling # noqa: F401 |
| 360 | except ImportError: |
| 361 | self._aiter_ops_import_failed = True |
| 362 | self.forward = self.forward_native |
| 363 | logger.warning_once( |
| 364 | "aiter.ops.sampling is not available on ROCm. " |
| 365 | "Falling back to PyTorch-native implementation." |
| 366 | ) |
| 367 | return False |
| 368 | self.aiter_ops = torch.ops.aiter |
| 369 | return True |
| 370 | |
| 371 | def forward_hip( |
| 372 | self, |
| 373 | logits: torch.Tensor, |
| 374 | generators: dict[int, torch.Generator], |
| 375 | k: torch.Tensor | None, |
| 376 | p: torch.Tensor | None, |
| 377 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 378 | """Optimized ROCm/aiter path (same structure as forward_cuda).""" |
| 379 | if (k is None and p is None) or generators: |
| 380 | if generators: |
| 381 | logger.warning_once( |
| 382 | "aiter sampler does not support per-request generators; " |
| 383 | "falling back to PyTorch-native." |
| 384 | ) |
| 385 | return self.forward_native(logits, generators, k, p) |
| 386 | if self.use_fp64_gumbel: |
| 387 | return self.forward_native(logits, generators, k, p) |
| 388 | assert self.logprobs_mode not in PROCESSED_LOGPROBS_MODES, ( |
| 389 | "aiter sampler does not support returning logits/logprobs." |
| 390 | ) |
| 391 | if self.aiter_ops is None and not self._init_aiter_ops(): |
| 392 | return self.forward_native(logits, generators, k, p) |
| 393 | return self.aiter_sample(logits, k, p, generators), None |
| 394 | |
| 395 | def aiter_sample( |
| 396 | self, |
| 397 | logits: torch.Tensor, |
| 398 | k: torch.Tensor | None, |
| 399 | p: torch.Tensor | None, |
| 400 | generators: dict[int, torch.Generator], |
| 401 | ) -> torch.Tensor: |
| 402 | """Sample from logits using aiter ops.""" |
| 403 | assert self.aiter_ops is not None |
| 404 | use_top_k = k is not None |
| 405 | use_top_p = p is not None |
| 406 | # Joint k+p path |
| 407 | if use_top_p and use_top_k: |
| 408 | probs = logits.softmax(dim=-1, dtype=torch.float32).contiguous() |
| 409 | next_token_ids = self.aiter_ops.top_k_top_p_sampling_from_probs( |
| 410 | probs, |
| 411 | None, |
| 412 | *_to_tensor_scalar_tuple(k), |
| 413 | *_to_tensor_scalar_tuple(p), |
| 414 | deterministic=True, |
| 415 | ) |
| 416 | return next_token_ids.view(-1) |
| 417 | # Top-p only path |
| 418 | elif use_top_p: |
| 419 | probs = logits.softmax(dim=-1, dtype=torch.float32).contiguous() |
| 420 | next_token_ids = self.aiter_ops.top_p_sampling_from_probs( |
| 421 | probs, None, *_to_tensor_scalar_tuple(p), deterministic=True |
| 422 | ) |
| 423 | return next_token_ids.view(-1) |
| 424 | # Top-k only path |
| 425 | elif use_top_k: |
| 426 | probs = logits.softmax(dim=-1, dtype=torch.float32).contiguous() |
| 427 | renorm_probs = self.aiter_ops.top_k_renorm_probs( |
| 428 | probs, *_to_tensor_scalar_tuple(k) |
| 429 | ) |
| 430 | return torch.multinomial(renorm_probs, num_samples=1).view(-1) |
| 431 | raise RuntimeError("aiter_sample was called with no active top-k or top-p.") |
| 432 | |
| 433 | def forward_xpu( |
| 434 | self, |
| 435 | logits: torch.Tensor, |
| 436 | generators: dict[int, torch.Generator], |
| 437 | k: torch.Tensor | None, |
| 438 | p: torch.Tensor | None, |
| 439 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 440 | if generators: |
| 441 | logger.warning_once( |
| 442 | "xpu kernel topk_topp_sampler does not support " |
| 443 | "per-request generators. Falling back to " |
| 444 | "PyTorch-native implementation." |
| 445 | ) |
| 446 | return self.forward_native(logits, generators, k, p) |
| 447 | return xpu_sample(logits, k, p, self.logprobs_mode) |
| 448 | |
| 449 | |
| 450 | # Note: this is a workaround for |
| 451 | # https://github.com/pytorch/pytorch/pull/151218 |
| 452 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 453 | def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor: |
| 454 | probs = logits.softmax(dim=-1, dtype=torch.float32) |
| 455 | q = torch.empty_like(probs) |
| 456 | q.exponential_() |
| 457 | return probs.div(q).argmax(dim=-1).view(-1) |
| 458 | |
| 459 | |
| 460 | def apply_top_k_top_p( |
| 461 | logits: torch.Tensor, k: torch.Tensor | None, p: torch.Tensor | None |
| 462 | ) -> torch.Tensor: |
| 463 | if p is None and k is None: |
| 464 | return logits |
| 465 | |
| 466 | if HAS_TRITON: |
| 467 | return apply_top_k_top_p_triton(logits, k, p) |
| 468 | |
| 469 | is_cpu = current_platform.is_cpu() |
| 470 | return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=is_cpu) |
| 471 | |
| 472 | |
| 473 | def apply_top_k_top_p_pytorch( |
| 474 | logits: torch.Tensor, |
| 475 | k: torch.Tensor | None, |
| 476 | p: torch.Tensor | None, |
| 477 | allow_cpu_sync: bool = False, |
| 478 | ) -> torch.Tensor: |
| 479 | """Apply top-k and top-p masks to the logits. |
| 480 | |
| 481 | If a top-p is used, this function will sort the logits tensor, |
| 482 | which can be slow for large batches. |
| 483 | |
| 484 | The logits tensor may be updated in-place. |
| 485 | """ |
| 486 | if p is None: |
| 487 | if k is None: |
| 488 | return logits |
| 489 | |
| 490 | if allow_cpu_sync: |
| 491 | # Avoid sorting vocab for top-k only case. |
| 492 | return apply_top_k_only(logits, k) |
| 493 | |
| 494 | logits_sort, logits_idx = logits.sort(dim=-1, descending=False) |
| 495 | |
| 496 | if k is not None: |
| 497 | # Apply top-k. |
| 498 | top_k_mask = logits_sort.size(1) - k.to(torch.long) # shape: B |
| 499 | # Get all the top_k values. |
| 500 | top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1)) |
| 501 | top_k_mask = logits_sort < top_k_mask |
| 502 | logits_sort.masked_fill_(top_k_mask, -float("inf")) |
| 503 | |
| 504 | if p is not None: |
| 505 | # Apply top-p. |
| 506 | probs_sort = logits_sort.softmax(dim=-1) |
| 507 | probs_sum = torch.cumsum(probs_sort, dim=-1, out=probs_sort) |
| 508 | top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1) |
| 509 | # at least one |
| 510 | top_p_mask[:, -1] = False |
| 511 | logits_sort.masked_fill_(top_p_mask, -float("inf")) |
| 512 | |
| 513 | # Re-sort the probabilities. |
| 514 | return logits.scatter_(dim=-1, index=logits_idx, src=logits_sort) |
| 515 | |
| 516 | |
| 517 | def apply_top_k_only(logits: torch.Tensor, k: torch.Tensor) -> torch.Tensor: |
| 518 | """Apply top-k mask to the logits. |
| 519 | |
| 520 | This implementation doesn't involve sorting the entire vocab. |
| 521 | Note however that it involves a GPU->CPU sync which can be detrimental for |
| 522 | async scheduling performance. |
| 523 | |
| 524 | The logits tensor may be updated in-place. |
| 525 | """ |
| 526 | no_top_k_mask = k == logits.shape[1] |
| 527 | # Set non-top-k rows to 1 so that we can gather. |
| 528 | k = k.masked_fill(no_top_k_mask, 1) |
| 529 | max_top_k = k.max() |
| 530 | # topk.values tensor has shape [batch_size, max_top_k]. |
| 531 | # Convert top k to 0-based index in range [0, max_top_k). |
| 532 | k_index = k.sub_(1).unsqueeze(1) |
| 533 | top_k_mask = logits.topk(max_top_k, dim=1).values.gather(1, k_index.long()) |
| 534 | # Handle non-topk rows. |
| 535 | top_k_mask.masked_fill_(no_top_k_mask.unsqueeze(1), -float("inf")) |
| 536 | return logits.masked_fill_(logits < top_k_mask, -float("inf")) |
| 537 | |
| 538 | |
| 539 | def empty_exponential_noise_like( |
| 540 | probs: torch.Tensor, use_fp64_gumbel: bool |
| 541 | ) -> torch.Tensor: |
| 542 | dtype = torch.float64 if use_fp64_gumbel else probs.dtype |
| 543 | return torch.empty(probs.shape, dtype=dtype, device=probs.device) |
| 544 | |
| 545 | |
| 546 | def sample_with_exponential_noise(probs: torch.Tensor, q: torch.Tensor) -> torch.Tensor: |
| 547 | if q.dtype == probs.dtype: |
| 548 | scores = probs.div_(q) |
| 549 | else: |
| 550 | scores = q.reciprocal_() |
| 551 | scores.mul_(probs) |
| 552 | return scores.argmax(dim=-1).view(-1) |
| 553 | |
| 554 | |
| 555 | def random_sample( |
| 556 | probs: torch.Tensor, |
| 557 | generators: dict[int, torch.Generator], |
| 558 | use_fp64_gumbel: bool = False, |
| 559 | ) -> torch.Tensor: |
| 560 | """Randomly sample from the probabilities. |
| 561 | |
| 562 | We use this function instead of torch.multinomial because torch.multinomial |
| 563 | causes CPU-GPU synchronization. |
| 564 | """ |
| 565 | q = empty_exponential_noise_like(probs, use_fp64_gumbel) |
| 566 | # NOTE(woosuk): To batch-process the requests without their own seeds, |
| 567 | # which is the common case, we first assume that every request does |
| 568 | # not have its own seed. Then, we overwrite the values for the requests |
| 569 | # that have their own seeds. |
| 570 | if len(generators) != probs.shape[0]: |
| 571 | q.exponential_() |
| 572 | if generators: |
| 573 | # TODO(woosuk): This can be slow because we handle each request |
| 574 | # one by one. Optimize this. |
| 575 | for i, generator in generators.items(): |
| 576 | q[i].exponential_(generator=generator) |
| 577 | return sample_with_exponential_noise(probs, q) |
| 578 | |
| 579 | |
| 580 | def flashinfer_sample( |
| 581 | logits: torch.Tensor, |
| 582 | k: torch.Tensor | None, |
| 583 | p: torch.Tensor | None, |
| 584 | generators: dict[int, torch.Generator] = {}, # noqa |
| 585 | ) -> torch.Tensor: |
| 586 | """Sample from the logits using FlashInfer. |
| 587 | |
| 588 | Statistically, this function is equivalent to the `random_sample` function. |
| 589 | However, this function is faster because it avoids sorting the logits tensor |
| 590 | via rejection sampling. |
| 591 | |
| 592 | NOTE: The outputs of this function do not necessarily match the outputs of |
| 593 | the `random_sample` function. It only guarantees that the outputs are |
| 594 | statistically equivalent. |
| 595 | """ |
| 596 | import flashinfer |
| 597 | |
| 598 | assert not (k is None and p is None) |
| 599 | if k is None: |
| 600 | # Top-p only. |
| 601 | probs = logits.softmax(dim=-1, dtype=torch.float32) |
| 602 | next_token_ids = flashinfer.sampling.top_p_sampling_from_probs( |
| 603 | probs, p, deterministic=True |
| 604 | ) |
| 605 | elif p is None: |
| 606 | # Top-k only. |
| 607 | probs = logits.softmax(dim=-1, dtype=torch.float32) |
| 608 | next_token_ids = flashinfer.sampling.top_k_sampling_from_probs( |
| 609 | probs, k, deterministic=True |
| 610 | ) |
| 611 | else: |
| 612 | # Both top-k and top-p. |
| 613 | next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits( |
| 614 | logits, k, p, deterministic=True |
| 615 | ) |
| 616 | |
| 617 | return next_token_ids.view(-1) |
| 618 | |
| 619 | |
| 620 | def _to_tensor_scalar_tuple(x): |
| 621 | return (x, 0) if isinstance(x, torch.Tensor) else (None, x) |