| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import ABC, abstractmethod |
| 4 | from collections.abc import Sequence |
| 5 | from dataclasses import dataclass |
| 6 | from enum import Enum, auto |
| 7 | from typing import TYPE_CHECKING |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm import SamplingParams |
| 12 | |
| 13 | if TYPE_CHECKING: |
| 14 | from vllm.config import VllmConfig |
| 15 | |
| 16 | |
| 17 | class MoveDirectionality(Enum): |
| 18 | # One-way i1->i2 req move within batch |
| 19 | UNIDIRECTIONAL = auto() |
| 20 | # Two-way i1<->i2 req swap within batch |
| 21 | SWAP = auto() |
| 22 | |
| 23 | |
| 24 | # Batch indices of any removed requests. |
| 25 | RemovedRequest = int |
| 26 | |
| 27 | # (index, params, prompt_tok_ids, output_tok_ids) tuples for new |
| 28 | # requests added to the batch. |
| 29 | AddedRequest = tuple[int, SamplingParams, list[int] | None, list[int]] |
| 30 | |
| 31 | # (index 1, index 2, directionality) tuples representing |
| 32 | # one-way moves or two-way swaps of requests in batch |
| 33 | MovedRequest = tuple[int, int, MoveDirectionality] |
| 34 | |
| 35 | |
| 36 | @dataclass(frozen=True) |
| 37 | class BatchUpdate: |
| 38 | """Persistent batch state change info for logitsprocs.""" |
| 39 | |
| 40 | batch_size: int # Current num reqs in batch |
| 41 | |
| 42 | # Metadata for requests added to, removed from, and moved |
| 43 | # within the persistent batch. |
| 44 | # |
| 45 | # Key assumption: the `output_tok_ids` list (which is an element of each |
| 46 | # tuple in `added`) is a reference to the request's running output tokens |
| 47 | # list; via this reference, the logits processors always see the latest |
| 48 | # list of generated output tokens. |
| 49 | # |
| 50 | # NOTE: |
| 51 | # * Added or moved requests may replace existing requests with the same |
| 52 | # index. |
| 53 | # * Operations should be processed in the following order: |
| 54 | # - removed, added, moved |
| 55 | removed: Sequence[RemovedRequest] |
| 56 | added: Sequence[AddedRequest] |
| 57 | moved: Sequence[MovedRequest] |
| 58 | |
| 59 | |
| 60 | class LogitsProcessor(ABC): |
| 61 | @classmethod |
| 62 | def validate_params(cls, sampling_params: SamplingParams): |
| 63 | """Validate sampling params for this logits processor. |
| 64 | |
| 65 | Raise ``VLLMValidationError`` (preferred) / ``ValueError`` (backward compatible) |
| 66 | for invalid params. Bare ``ValueError`` is converted to ``VLLMValidationError`` |
| 67 | at the engine boundary so online serving returns HTTP 400. |
| 68 | """ |
| 69 | return None |
| 70 | |
| 71 | @abstractmethod |
| 72 | def __init__( |
| 73 | self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool |
| 74 | ) -> None: |
| 75 | raise NotImplementedError |
| 76 | |
| 77 | @abstractmethod |
| 78 | def apply(self, logits: torch.Tensor) -> torch.Tensor: |
| 79 | """Apply LogitsProcessor to batch logits tensor. |
| 80 | |
| 81 | The updated tensor must be returned but may be |
| 82 | modified in-place. |
| 83 | """ |
| 84 | raise NotImplementedError |
| 85 | |
| 86 | @abstractmethod |
| 87 | def is_argmax_invariant(self) -> bool: |
| 88 | """True if logits processor has no impact on the |
| 89 | argmax computation in greedy sampling. |
| 90 | NOTE: may or may not have the same value for all |
| 91 | instances of a given LogitsProcessor subclass, |
| 92 | depending on subclass implementation. |
| 93 | """ |
| 94 | raise NotImplementedError |
| 95 | |
| 96 | @abstractmethod |
| 97 | def update_state( |
| 98 | self, |
| 99 | batch_update: "BatchUpdate | None", |
| 100 | ) -> None: |
| 101 | """Called when there are new output tokens, prior |
| 102 | to each forward pass. |
| 103 | |
| 104 | Args: |
| 105 | batch_update: Non-None iff there have been changes |
| 106 | to the batch makeup. |
| 107 | |
| 108 | """ |
| 109 | raise NotImplementedError |