| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """A layer that compute logits from hidden_stats.""" |
| 4 | |
| 5 | from collections.abc import Callable |
| 6 | from functools import cache |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn.functional as F |
| 10 | |
| 11 | from vllm.config import get_current_vllm_config |
| 12 | from vllm.distributed import ( |
| 13 | tensor_model_parallel_all_gather, |
| 14 | tensor_model_parallel_gather, |
| 15 | ) |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.model_executor.custom_op import PluggableLayer |
| 18 | from vllm.model_executor.layers.linear import UnquantizedLinearMethod |
| 19 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 20 | UnquantizedEmbeddingMethod, |
| 21 | VocabParallelEmbedding, |
| 22 | ) |
| 23 | from vllm.platforms import current_platform |
| 24 | from vllm.utils.flashinfer import has_flashinfer |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | |
| 29 | @cache |
| 30 | def _flashinfer_topk() -> Callable[..., tuple[torch.Tensor, torch.Tensor]] | None: |
| 31 | """FlashInfer's radix top-k, or None for torch.topk. |
| 32 | |
| 33 | The top-k spans the vocabulary, where the radix kernel is about twice |
| 34 | torch.topk. |
| 35 | """ |
| 36 | if not current_platform.is_cuda(): |
| 37 | return None |
| 38 | if not has_flashinfer(): |
| 39 | logger.info_once( |
| 40 | "flashinfer is unavailable; vocab-parallel top-k uses torch.topk, " |
| 41 | "at roughly half the speed." |
| 42 | ) |
| 43 | return None |
| 44 | from flashinfer import top_k |
| 45 | |
| 46 | return top_k |
| 47 | |
| 48 | |
| 49 | def _topk(scores: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]: |
| 50 | impl = _flashinfer_topk() |
| 51 | if impl is None or not scores.is_cuda: |
| 52 | return torch.topk(scores, k, dim=-1) |
| 53 | return impl(scores, k, sorted=True, deterministic=True) |
| 54 | |
| 55 | |
| 56 | # --8<-- [start:logits_processor] |
| 57 | @PluggableLayer.register("logits_processor") |
| 58 | class LogitsProcessor(PluggableLayer): |
| 59 | """Process logits and apply logits processors from sampling metadata. |
| 60 | |
| 61 | This layer does the following: |
| 62 | 1. Gather logits from model hidden_states. |
| 63 | 2. Scale logits if needed. |
| 64 | 3. Apply logits processors (if any). |
| 65 | """ |
| 66 | |
| 67 | # --8<-- [end:logits_processor] |
| 68 | |
| 69 | def __init__( |
| 70 | self, |
| 71 | vocab_size: int, |
| 72 | org_vocab_size: int | None = None, |
| 73 | scale: float = 1.0, |
| 74 | logits_as_input: bool = False, |
| 75 | soft_cap: float | None = None, |
| 76 | ) -> None: |
| 77 | """Args: |
| 78 | scale: A scaling factor to apply to the logits. |
| 79 | |
| 80 | """ |
| 81 | super().__init__() |
| 82 | self.scale = scale |
| 83 | self.vocab_size = vocab_size |
| 84 | # Whether the input is logits (default is hidden states). |
| 85 | self.logits_as_input = logits_as_input |
| 86 | # original vocabulary size (without LoRA). |
| 87 | self.org_vocab_size = org_vocab_size or vocab_size |
| 88 | # Soft cap the logits. Used in Gemma 2. |
| 89 | self.soft_cap = soft_cap |
| 90 | # Whether to use gather or all-gather to gather the logits. |
| 91 | self.use_all_gather = current_platform.use_all_gather() |
| 92 | # Dtype of the lm_head projection. Defaults to the model dtype; an |
| 93 | # fp32 head (via `--hf-overrides '{"head_dtype": "float32"}'`) is |
| 94 | # required for RL training-inference consistency. |
| 95 | model_config = get_current_vllm_config().model_config |
| 96 | self.head_dtype = model_config.head_dtype if model_config is not None else None |
| 97 | |
| 98 | def forward( |
| 99 | self, |
| 100 | lm_head: VocabParallelEmbedding, |
| 101 | hidden_states: torch.Tensor, |
| 102 | embedding_bias: torch.Tensor | None = None, |
| 103 | skip_gather: bool = False, |
| 104 | ) -> torch.Tensor | None: |
| 105 | if self.logits_as_input: |
| 106 | logits = hidden_states |
| 107 | else: |
| 108 | # Get the logits for the next tokens. |
| 109 | logits = self._get_logits( |
| 110 | hidden_states, lm_head, embedding_bias, skip_gather |
| 111 | ) |
| 112 | if logits is not None: |
| 113 | if self.soft_cap is not None: |
| 114 | logits = logits / self.soft_cap |
| 115 | logits = torch.tanh(logits) |
| 116 | logits = logits * self.soft_cap |
| 117 | |
| 118 | if self.scale != 1.0: |
| 119 | logits *= self.scale |
| 120 | return logits |
| 121 | |
| 122 | def _gather_logits(self, logits: torch.Tensor) -> torch.Tensor: |
| 123 | """gather/all-gather the logits tensor across model parallel group.""" |
| 124 | if self.use_all_gather: |
| 125 | # Gather is not supported for some devices such as TPUs. |
| 126 | # Use all-gather instead. |
| 127 | # NOTE(woosuk): Here, the outputs of every device should not be None |
| 128 | # because XLA requires strict SPMD among all devices. Every device |
| 129 | # should execute the same operations after gathering the logits. |
| 130 | logits = tensor_model_parallel_all_gather(logits) |
| 131 | else: |
| 132 | # None may be returned for rank > 0 |
| 133 | logits = tensor_model_parallel_gather(logits) |
| 134 | return logits |
| 135 | |
| 136 | def _apply_head( |
| 137 | self, |
| 138 | lm_head: VocabParallelEmbedding, |
| 139 | hidden_states: torch.Tensor, |
| 140 | embedding_bias: torch.Tensor | None, |
| 141 | ) -> torch.Tensor: |
| 142 | """Project hidden states through the lm_head, honoring head_dtype.""" |
| 143 | if self.head_dtype is None or self.head_dtype == hidden_states.dtype: |
| 144 | return lm_head.quant_method.apply( |
| 145 | lm_head, hidden_states, bias=embedding_bias |
| 146 | ) |
| 147 | |
| 148 | # A quant config that excludes lm_head hands out UnquantizedLinearMethod |
| 149 | # rather than UnquantizedEmbeddingMethod, so accept both: either way the |
| 150 | # weight is plain and `lm_head.weight` can be cast directly. |
| 151 | if not isinstance( |
| 152 | lm_head.quant_method, (UnquantizedEmbeddingMethod, UnquantizedLinearMethod) |
| 153 | ): |
| 154 | raise ValueError( |
| 155 | "A head_dtype different from the model dtype is only " |
| 156 | "supported for an unquantized lm_head." |
| 157 | ) |
| 158 | if ( |
| 159 | self.head_dtype == torch.float32 |
| 160 | and (current_platform.is_cuda() or current_platform.is_rocm()) |
| 161 | and hidden_states.is_cuda |
| 162 | ): |
| 163 | # Accumulate the projection directly into fp32. This avoids |
| 164 | # materializing an fp32 copy of the lm_head weight on every step, |
| 165 | # unlike casting both operands. `torch.mm(out_dtype=...)` only |
| 166 | # supports fp32 output for fp16/bf16 inputs, and is only |
| 167 | # implemented for CUDA and ROCm (the latter via the non-Lt GEMM |
| 168 | # path); other platforms fall back to the cast path below. |
| 169 | flat = hidden_states.reshape(-1, hidden_states.shape[-1]) |
| 170 | logits = torch.mm(flat, lm_head.weight.t(), out_dtype=self.head_dtype) |
| 171 | if embedding_bias is not None: |
| 172 | logits = logits + embedding_bias.to(self.head_dtype) |
| 173 | return logits.reshape(*hidden_states.shape[:-1], -1) |
| 174 | return F.linear( |
| 175 | hidden_states.to(self.head_dtype), |
| 176 | lm_head.weight.to(self.head_dtype), |
| 177 | embedding_bias.to(self.head_dtype) if embedding_bias is not None else None, |
| 178 | ) |
| 179 | |
| 180 | def _get_logits( |
| 181 | self, |
| 182 | hidden_states: torch.Tensor, |
| 183 | lm_head: VocabParallelEmbedding, |
| 184 | embedding_bias: torch.Tensor | None, |
| 185 | skip_gather: bool = False, |
| 186 | ) -> torch.Tensor | None: |
| 187 | # Get the logits for the next tokens. |
| 188 | logits = self._apply_head(lm_head, hidden_states, embedding_bias) |
| 189 | if skip_gather: |
| 190 | return logits |
| 191 | |
| 192 | # Gather logits for TP |
| 193 | if lm_head.tp_size > 1: |
| 194 | logits = self._gather_logits(logits) |
| 195 | |
| 196 | # Remove paddings in vocab (if any). |
| 197 | if logits is not None: |
| 198 | logits = logits[..., : self.org_vocab_size] |
| 199 | return logits |
| 200 | |
| 201 | def get_top_tokens( |
| 202 | self, |
| 203 | lm_head: VocabParallelEmbedding, |
| 204 | hidden_states: torch.Tensor, |
| 205 | embedding_bias: torch.Tensor | None = None, |
| 206 | ) -> torch.Tensor: |
| 207 | """Vocab-parallel argmax without all-gathering full logits. |
| 208 | |
| 209 | Each TP rank computes local argmax, then only the (value, index) pairs |
| 210 | are gathered and reduced. Communication: O(batch * 2 * tp_size) vs |
| 211 | O(batch * vocab_size). |
| 212 | """ |
| 213 | if self.scale <= 0.0 and self.scale != 1.0: |
| 214 | raise ValueError( |
| 215 | "The local argmax reduction optimization is not supported for " |
| 216 | "non-positive logit scaling factors." |
| 217 | ) |
| 218 | tp_size = lm_head.tp_size |
| 219 | |
| 220 | logits = self._apply_head(lm_head, hidden_states, embedding_bias) |
| 221 | if self.soft_cap is not None: |
| 222 | logits = torch.tanh(logits / self.soft_cap) * self.soft_cap |
| 223 | if self.scale != 1.0: |
| 224 | logits = logits * self.scale |
| 225 | |
| 226 | # Mask out padding entries beyond org_vocab_size on this shard. |
| 227 | num_pad = lm_head.shard_indices.num_org_vocab_padding |
| 228 | if num_pad > 0: |
| 229 | logits[..., -num_pad:] = -float("inf") |
| 230 | |
| 231 | local_max_vals, local_max_indices = logits.max(dim=-1) |
| 232 | |
| 233 | # Convert shard-local indices to global vocab indices. |
| 234 | vocab_start = lm_head.shard_indices.org_vocab_start_index |
| 235 | global_indices = local_max_indices + vocab_start |
| 236 | |
| 237 | if tp_size == 1: |
| 238 | return global_indices |
| 239 | |
| 240 | # All-gather (value, index) pairs, then reduce to global argmax. |
| 241 | # Use float32 to avoid bf16 precision loss on large vocab indices. |
| 242 | local_pair = torch.stack( |
| 243 | [local_max_vals.float(), global_indices.float()], dim=-1 |
| 244 | ) |
| 245 | # [batch, 2] -> [batch, 2 * tp_size] |
| 246 | gathered = tensor_model_parallel_all_gather(local_pair, dim=-1) |
| 247 | # [batch, tp_size, 2] where [:, :, 0]=values, [:, :, 1]=indices |
| 248 | gathered = gathered.view(hidden_states.shape[0], tp_size, 2) |
| 249 | max_rank_idx = gathered[:, :, 0].argmax(dim=-1, keepdim=True) |
| 250 | top_tokens = gathered[:, :, 1].gather(dim=-1, index=max_rank_idx) |
| 251 | return top_tokens.squeeze(-1).to(torch.int64) |
| 252 | |
| 253 | def get_top_k_tokens( |
| 254 | self, |
| 255 | lm_head: VocabParallelEmbedding, |
| 256 | hidden_states: torch.Tensor, |
| 257 | k: int, |
| 258 | embedding_bias: torch.Tensor | None = None, |
| 259 | *, |
| 260 | return_log_probs: bool = False, |
| 261 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 262 | """Vocab-parallel top-k without all-gathering full logits. |
| 263 | |
| 264 | The `get_top_tokens` reduction widened from one token to k, returning |
| 265 | the values as well as the global ids. Communication is |
| 266 | O(batch * 2k * tp_size) rather than O(batch * vocab_size). |
| 267 | |
| 268 | With return_log_probs, values use the full-vocabulary log-partition. |
| 269 | |
| 270 | Otherwise scale and soft cap are applied to the k selected values rather than |
| 271 | the whole vocabulary; both are monotonic, so the selection is the same |
| 272 | and only k entries are touched. |
| 273 | """ |
| 274 | if self.scale <= 0.0 and self.scale != 1.0: |
| 275 | raise ValueError( |
| 276 | "The local top-k reduction optimization is not supported for " |
| 277 | "non-positive logit scaling factors." |
| 278 | ) |
| 279 | |
| 280 | logits = self._apply_head(lm_head, hidden_states, embedding_bias) |
| 281 | if return_log_probs: |
| 282 | if self.soft_cap is not None: |
| 283 | logits = torch.tanh(logits / self.soft_cap) * self.soft_cap |
| 284 | if self.scale != 1.0: |
| 285 | logits = logits * self.scale |
| 286 | |
| 287 | # Mask out padding entries beyond org_vocab_size on this shard. |
| 288 | num_pad = lm_head.shard_indices.num_org_vocab_padding |
| 289 | if num_pad > 0: |
| 290 | logits[..., -num_pad:] = -float("inf") |
| 291 | |
| 292 | if return_log_probs: |
| 293 | log_partition = torch.logsumexp(logits.float(), dim=-1, keepdim=True) |
| 294 | if lm_head.tp_size > 1: |
| 295 | log_partition = torch.logsumexp( |
| 296 | tensor_model_parallel_all_gather(log_partition, dim=-1), |
| 297 | dim=-1, |
| 298 | keepdim=True, |
| 299 | ) |
| 300 | |
| 301 | values, ids = _topk(logits, k) |
| 302 | # Convert shard-local indices to global vocab indices. |
| 303 | ids = ids.to(torch.int64) + lm_head.shard_indices.org_vocab_start_index |
| 304 | |
| 305 | if lm_head.tp_size > 1: |
| 306 | values = tensor_model_parallel_all_gather(values, dim=-1) |
| 307 | ids = tensor_model_parallel_all_gather(ids, dim=-1) |
| 308 | values, selected = _topk(values, k) |
| 309 | ids = ids.gather(-1, selected) |
| 310 | |
| 311 | values = values.float() |
| 312 | if return_log_probs: |
| 313 | return ids, values - log_partition |
| 314 | if self.scale != 1.0: |
| 315 | values = values * self.scale |
| 316 | if self.soft_cap is not None: |
| 317 | values = torch.tanh(values / self.soft_cap) * self.soft_cap |
| 318 | return ids, values |
| 319 | |
| 320 | def extra_repr(self) -> str: |
| 321 | s = f"vocab_size={self.vocab_size}" |
| 322 | s += f", org_vocab_size={self.org_vocab_size}" |
| 323 | s += f", scale={self.scale}, logits_as_input={self.logits_as_input}" |
| 324 | return s |