| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable, Sequence |
| 4 | from typing import TYPE_CHECKING, TypeVar |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | |
| 9 | from vllm import SamplingParams |
| 10 | from vllm.utils.torch_utils import async_tensor_h2d |
| 11 | from vllm.v1.sample.logits_processor.interface import ( |
| 12 | BatchUpdate, |
| 13 | LogitsProcessor, |
| 14 | MoveDirectionality, |
| 15 | ) |
| 16 | |
| 17 | if TYPE_CHECKING: |
| 18 | from vllm.config import VllmConfig |
| 19 | |
| 20 | T = TypeVar("T") |
| 21 | |
| 22 | |
| 23 | class MinPLogitsProcessor(LogitsProcessor): |
| 24 | def __init__( |
| 25 | self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool |
| 26 | ): |
| 27 | max_num_reqs = vllm_config.scheduler_config.max_num_seqs |
| 28 | self.min_p_count: int = 0 |
| 29 | |
| 30 | self.min_p_cpu_tensor = torch.zeros( |
| 31 | (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=is_pin_memory |
| 32 | ) |
| 33 | self.min_p_cpu = self.min_p_cpu_tensor.numpy() |
| 34 | |
| 35 | self.use_double_tensor = torch.device(device).type != "cpu" |
| 36 | |
| 37 | if self.use_double_tensor: |
| 38 | # Pre-allocated device tensor |
| 39 | self.min_p_device: torch.Tensor = torch.empty( |
| 40 | (max_num_reqs,), dtype=torch.float32, device=device |
| 41 | ) |
| 42 | else: |
| 43 | self.min_p_device = self.min_p_cpu_tensor |
| 44 | # Current slice of the device tensor |
| 45 | self.min_p: torch.Tensor = self.min_p_device[:0] |
| 46 | |
| 47 | def is_argmax_invariant(self) -> bool: |
| 48 | """Min-p never impacts greedy sampling.""" |
| 49 | return True |
| 50 | |
| 51 | def get_min_p_by_index(self, index: int) -> float: |
| 52 | return float(self.min_p_cpu[index]) |
| 53 | |
| 54 | def update_state(self, batch_update: BatchUpdate | None): |
| 55 | if not batch_update: |
| 56 | return |
| 57 | |
| 58 | needs_update = False |
| 59 | # Process added requests. |
| 60 | for index, params, _, _ in batch_update.added: |
| 61 | min_p = params.min_p |
| 62 | min_p_before = self.min_p_cpu[index] |
| 63 | if min_p_before != min_p: |
| 64 | needs_update = True |
| 65 | self.min_p_cpu[index] = min_p |
| 66 | if min_p and not min_p_before: |
| 67 | self.min_p_count += 1 |
| 68 | elif not min_p and min_p_before: |
| 69 | self.min_p_count -= 1 |
| 70 | |
| 71 | if self.min_p_count: |
| 72 | # Process removed requests. |
| 73 | if batch_update.removed: |
| 74 | needs_update = True |
| 75 | for index in batch_update.removed: |
| 76 | if self.min_p_cpu[index]: |
| 77 | self.min_p_cpu[index] = 0 |
| 78 | self.min_p_count -= 1 |
| 79 | |
| 80 | # Process moved requests, unidirectional (a->b) and swap (a<->b). |
| 81 | for adx, bdx, direct in batch_update.moved: |
| 82 | min_p_a, min_p_b = self.min_p_cpu[adx], self.min_p_cpu[bdx] |
| 83 | if min_p_a != min_p_b: |
| 84 | needs_update = True |
| 85 | self.min_p_cpu[bdx] = min_p_a |
| 86 | if direct == MoveDirectionality.SWAP: |
| 87 | self.min_p_cpu[adx] = min_p_b |
| 88 | if direct == MoveDirectionality.UNIDIRECTIONAL: |
| 89 | if min_p_a: |
| 90 | self.min_p_cpu[adx] = 0 |
| 91 | if min_p_b: |
| 92 | self.min_p_count -= 1 |
| 93 | |
| 94 | # Update tensors if needed. |
| 95 | size = batch_update.batch_size |
| 96 | if self.min_p_count and (needs_update or self.min_p.shape[0] != size): |
| 97 | self.min_p = self.min_p_device[:size] |
| 98 | if self.use_double_tensor: |
| 99 | self.min_p.copy_(self.min_p_cpu_tensor[:size], non_blocking=True) |
| 100 | self.min_p.unsqueeze_(1) |
| 101 | |
| 102 | def apply(self, logits: torch.Tensor) -> torch.Tensor: |
| 103 | if not self.min_p_count: |
| 104 | return logits |
| 105 | |
| 106 | # Convert logits to probability distribution |
| 107 | probability_values = torch.nn.functional.softmax(logits, dim=-1) |
| 108 | # Calculate maximum probabilities per sequence |
| 109 | max_probabilities = torch.amax(probability_values, dim=-1, keepdim=True) |
| 110 | # Adjust min_p |
| 111 | adjusted_min_p = max_probabilities.mul_(self.min_p) |
| 112 | # Identify valid tokens using threshold comparison |
| 113 | invalid_token_mask = probability_values < adjusted_min_p |
| 114 | # Apply mask using boolean indexing |
| 115 | logits.masked_fill_(invalid_token_mask, -float("inf")) |
| 116 | return logits |
| 117 | |
| 118 | |
| 119 | class LogitBiasLogitsProcessor(LogitsProcessor): |
| 120 | def __init__(self, _, device: torch.device, is_pin_memory: bool): |
| 121 | self.device = device |
| 122 | self.biases: dict[int, dict[int, float]] = {} |
| 123 | |
| 124 | self.bias_tensor: torch.Tensor = torch.tensor(()) |
| 125 | self.logits_slice = ( |
| 126 | self._device_tensor([], torch.int32), |
| 127 | self._device_tensor([], torch.int32), |
| 128 | ) |
| 129 | |
| 130 | def is_argmax_invariant(self) -> bool: |
| 131 | """Logit bias can rebalance token probabilities and change the |
| 132 | outcome of argmax in greedy sampling.""" |
| 133 | return False |
| 134 | |
| 135 | def update_state(self, batch_update: BatchUpdate | None): |
| 136 | needs_update = process_dict_updates( |
| 137 | self.biases, batch_update, lambda params, _, __: params.logit_bias or None |
| 138 | ) |
| 139 | |
| 140 | # Update tensors if needed. |
| 141 | if needs_update: |
| 142 | reqs: list[int] = [] |
| 143 | tok_ids: list[int] = [] |
| 144 | biases: list[float] = [] |
| 145 | for req, lb in self.biases.items(): |
| 146 | reqs.extend([req] * len(lb)) |
| 147 | tok_ids.extend(lb.keys()) |
| 148 | biases.extend(lb.values()) |
| 149 | |
| 150 | self.bias_tensor = self._device_tensor(biases, torch.float32) |
| 151 | self.logits_slice = ( |
| 152 | self._device_tensor(reqs, torch.int32), |
| 153 | self._device_tensor(tok_ids, torch.int32), |
| 154 | ) |
| 155 | |
| 156 | def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor: |
| 157 | return async_tensor_h2d(data, device=self.device, dtype=dtype) |
| 158 | |
| 159 | def apply(self, logits: torch.Tensor) -> torch.Tensor: |
| 160 | if self.biases: |
| 161 | logits[self.logits_slice] += self.bias_tensor |
| 162 | return logits |
| 163 | |
| 164 | |
| 165 | class MinTokensLogitsProcessor(LogitsProcessor): |
| 166 | def __init__( |
| 167 | self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool |
| 168 | ): |
| 169 | # index -> (min_toks, output_token_ids, stop_token_ids, uses_structured_output) |
| 170 | self.device = device |
| 171 | self.min_toks: dict[int, tuple[int, Sequence[int], set[int], bool]] = {} |
| 172 | |
| 173 | # (req_idx_tensor, stop_token_id_tensor) |
| 174 | self.logits_slice: tuple[torch.Tensor, torch.Tensor] = ( |
| 175 | self._device_tensor([], torch.int32), |
| 176 | self._device_tensor([], torch.int32), |
| 177 | ) |
| 178 | |
| 179 | # Structured-output stop-token logits that may need to be restored. |
| 180 | self.restore_logits_slice: tuple[torch.Tensor, torch.Tensor] = ( |
| 181 | self._device_tensor([], torch.int32), |
| 182 | self._device_tensor([], torch.int32), |
| 183 | ) |
| 184 | |
| 185 | self.neg_inf_tensor = torch.tensor( |
| 186 | -float("inf"), dtype=torch.float32, device=self.device |
| 187 | ) |
| 188 | |
| 189 | def is_argmax_invariant(self) -> bool: |
| 190 | """By censoring stop tokens, min-tokens can change the outcome |
| 191 | of the argmax operation in greedy sampling.""" |
| 192 | return False |
| 193 | |
| 194 | @staticmethod |
| 195 | def add_request( |
| 196 | params: SamplingParams, _: list[int] | None, output_tok_ids: list[int] |
| 197 | ) -> tuple[int, Sequence[int], set[int], bool] | None: |
| 198 | min_tokens = params.min_tokens |
| 199 | if not min_tokens or len(output_tok_ids) >= min_tokens: |
| 200 | return None |
| 201 | return ( |
| 202 | min_tokens, |
| 203 | output_tok_ids, |
| 204 | params.all_stop_token_ids, |
| 205 | params.structured_outputs is not None, |
| 206 | ) |
| 207 | |
| 208 | def update_state(self, batch_update: BatchUpdate | None): |
| 209 | needs_update = process_dict_updates( |
| 210 | self.min_toks, batch_update, self.add_request |
| 211 | ) |
| 212 | if self.min_toks: |
| 213 | # Check for any requests that have attained their min tokens. |
| 214 | to_remove = tuple( |
| 215 | index |
| 216 | for index, (min_toks, out_tok_ids, _, _) in self.min_toks.items() |
| 217 | if len(out_tok_ids) >= min_toks |
| 218 | ) |
| 219 | if to_remove: |
| 220 | needs_update = True |
| 221 | for index in to_remove: |
| 222 | del self.min_toks[index] |
| 223 | |
| 224 | # Update tensors if needed. |
| 225 | if needs_update: |
| 226 | reqs: list[int] = [] |
| 227 | tok_ids: list[int] = [] |
| 228 | restore_reqs: list[int] = [] |
| 229 | restore_tok_ids: list[int] = [] |
| 230 | for req, ( |
| 231 | _, |
| 232 | _, |
| 233 | stop_tok_ids, |
| 234 | uses_structured_output, |
| 235 | ) in self.min_toks.items(): |
| 236 | reqs.extend([req] * len(stop_tok_ids)) |
| 237 | tok_ids.extend(stop_tok_ids) |
| 238 | if uses_structured_output: |
| 239 | restore_reqs.extend([req] * len(stop_tok_ids)) |
| 240 | restore_tok_ids.extend(stop_tok_ids) |
| 241 | |
| 242 | self.logits_slice = ( |
| 243 | self._device_tensor(reqs, torch.int32), |
| 244 | self._device_tensor(tok_ids, torch.int32), |
| 245 | ) |
| 246 | self.restore_logits_slice = ( |
| 247 | self._device_tensor(restore_reqs, torch.int32), |
| 248 | self._device_tensor(restore_tok_ids, torch.int32), |
| 249 | ) |
| 250 | |
| 251 | def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor: |
| 252 | return async_tensor_h2d(data, device=self.device, dtype=dtype) |
| 253 | |
| 254 | def _mask_stop_token_logits( |
| 255 | self, |
| 256 | logits: torch.Tensor, |
| 257 | logits_slice: tuple[torch.Tensor, torch.Tensor], |
| 258 | restore_logits_slice: tuple[torch.Tensor, torch.Tensor], |
| 259 | ) -> None: |
| 260 | restore_rows, restore_toks = restore_logits_slice |
| 261 | # Advanced indexing already returns a copy, so this survives the |
| 262 | # in-place masking below. |
| 263 | stop_logits = logits[restore_logits_slice] |
| 264 | logits.index_put_(logits_slice, self.neg_inf_tensor) |
| 265 | |
| 266 | # Stop-token entries for the same row are emitted together, so this |
| 267 | # checks each row once without sorting the index tensor. |
| 268 | unique_rows, inverse_indices = torch.unique_consecutive( |
| 269 | restore_rows, return_inverse=True |
| 270 | ) |
| 271 | row_needs_restore = torch.isneginf(logits[unique_rows]).all(dim=-1) |
| 272 | restore_mask = row_needs_restore[inverse_indices] & torch.isfinite(stop_logits) |
| 273 | restore_slice = ( |
| 274 | restore_rows[restore_mask], |
| 275 | restore_toks[restore_mask], |
| 276 | ) |
| 277 | logits.index_put_(restore_slice, stop_logits[restore_mask]) |
| 278 | |
| 279 | def apply(self, logits: torch.Tensor) -> torch.Tensor: |
| 280 | if self.min_toks: |
| 281 | # Inhibit stop tokens for requests which have not reached min length. |
| 282 | if self.restore_logits_slice[0].numel() == 0: |
| 283 | logits.index_put_(self.logits_slice, self.neg_inf_tensor) |
| 284 | else: |
| 285 | self._mask_stop_token_logits( |
| 286 | logits, self.logits_slice, self.restore_logits_slice |
| 287 | ) |
| 288 | return logits |
| 289 | |
| 290 | def apply_with_spec_decode( |
| 291 | self, |
| 292 | logits: torch.Tensor, |
| 293 | num_draft_tokens: list[int], |
| 294 | ) -> torch.Tensor: |
| 295 | """Spec-decode version of apply(). |
| 296 | |
| 297 | Stop tokens stay masked unless masking them would leave every logit |
| 298 | in a row at -inf. |
| 299 | |
| 300 | Example: ``num_draft_tokens = [2, 3, 1]`` |
| 301 | -> ``logits`` shape ``[6, V]``, ``cumsum = [0, 2, 5, 6]`` |
| 302 | -> request 0 owns rows 0-1, request 1 rows 2-4, request 2 row 5. |
| 303 | """ |
| 304 | if not self.min_toks: |
| 305 | return logits |
| 306 | |
| 307 | num_draft_arr = np.array(num_draft_tokens, dtype=np.int64) |
| 308 | cumsum = np.concatenate([[0], np.cumsum(num_draft_arr)]) |
| 309 | |
| 310 | entries = [ |
| 311 | ( |
| 312 | req_idx, |
| 313 | min_tok, |
| 314 | len(out_tok_ids), |
| 315 | list(stop_tok_ids), |
| 316 | uses_structured_output, |
| 317 | ) |
| 318 | for req_idx, ( |
| 319 | min_tok, |
| 320 | out_tok_ids, |
| 321 | stop_tok_ids, |
| 322 | uses_structured_output, |
| 323 | ) in self.min_toks.items() |
| 324 | if stop_tok_ids |
| 325 | ] |
| 326 | |
| 327 | if not entries: |
| 328 | return logits |
| 329 | |
| 330 | all_rows: list[np.ndarray] = [] # row indices to mask |
| 331 | all_toks: list[np.ndarray] = [] # stop-token ids at those rows |
| 332 | restore_rows: list[np.ndarray] = [] |
| 333 | restore_toks: list[np.ndarray] = [] |
| 334 | |
| 335 | for req_idx, min_tok, current_len, stop_toks, uses_structured_output in entries: |
| 336 | remaining = min_tok - current_len |
| 337 | # How many leading draft positions still need stop-token masking. |
| 338 | n_mask = int(min(max(remaining, 0), num_draft_arr[req_idx])) |
| 339 | |
| 340 | if n_mask > 0: |
| 341 | offset = cumsum[req_idx] |
| 342 | row_indices = np.arange(offset, offset + n_mask, dtype=np.int64) |
| 343 | n_stop = len(stop_toks) |
| 344 | rows = np.repeat(row_indices, n_stop) |
| 345 | toks = np.tile(stop_toks, n_mask) |
| 346 | all_rows.append(rows) |
| 347 | all_toks.append(toks) |
| 348 | if uses_structured_output: |
| 349 | restore_rows.append(rows) |
| 350 | restore_toks.append(toks) |
| 351 | |
| 352 | if all_rows: |
| 353 | rows_arr = np.concatenate(all_rows) |
| 354 | toks_arr = np.concatenate(all_toks) |
| 355 | # (row_indices, token_indices) for index_put_ to set -inf. |
| 356 | logits_slice = ( |
| 357 | async_tensor_h2d(rows_arr, device=self.device), |
| 358 | async_tensor_h2d(toks_arr, device=self.device), |
| 359 | ) |
| 360 | if restore_rows: |
| 361 | restore_rows_arr = np.concatenate(restore_rows) |
| 362 | restore_toks_arr = np.concatenate(restore_toks) |
| 363 | restore_logits_slice = ( |
| 364 | async_tensor_h2d(restore_rows_arr, device=self.device), |
| 365 | async_tensor_h2d(restore_toks_arr, device=self.device), |
| 366 | ) |
| 367 | self._mask_stop_token_logits(logits, logits_slice, restore_logits_slice) |
| 368 | else: |
| 369 | logits.index_put_(logits_slice, self.neg_inf_tensor) |
| 370 | |
| 371 | return logits |
| 372 | |
| 373 | |
| 374 | def process_dict_updates( |
| 375 | req_entries: dict[int, T], |
| 376 | batch_update: BatchUpdate | None, |
| 377 | new_state: Callable[[SamplingParams, list[int] | None, list[int]], T | None], |
| 378 | ) -> bool: |
| 379 | """Utility function to update dict state for sparse LogitsProcessors.""" |
| 380 | if not batch_update: |
| 381 | # Nothing to do. |
| 382 | return False |
| 383 | |
| 384 | updated = False |
| 385 | for index, params, prompt_tok_ids, output_tok_ids in batch_update.added: |
| 386 | if (state := new_state(params, prompt_tok_ids, output_tok_ids)) is not None: |
| 387 | req_entries[index] = state |
| 388 | updated = True |
| 389 | elif req_entries.pop(index, None) is not None: |
| 390 | updated = True |
| 391 | |
| 392 | if req_entries: |
| 393 | # Process removed requests. |
| 394 | for index in batch_update.removed: |
| 395 | if req_entries.pop(index, None): |
| 396 | updated = True |
| 397 | |
| 398 | # Process moved requests, unidirectional (a->b) and |
| 399 | # swapped (a<->b) |
| 400 | for a_index, b_index, direct in batch_update.moved: |
| 401 | a_entry = req_entries.pop(a_index, None) |
| 402 | b_entry = req_entries.pop(b_index, None) |
| 403 | if a_entry is not None: |
| 404 | req_entries[b_index] = a_entry |
| 405 | updated = True |
| 406 | if b_entry is not None: |
| 407 | updated = True |
| 408 | if direct == MoveDirectionality.SWAP: |
| 409 | req_entries[a_index] = b_entry |
| 410 | |
| 411 | return updated |