| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """V2 logits processor interface. |
| 4 | |
| 5 | Import-light by design: custom processor classes are also loaded in the |
| 6 | frontend process (to validate per-request params), so this module must not |
| 7 | pull in model-runner side modules (torch, triton, worker state) at import |
| 8 | time. Heavy imports live under TYPE_CHECKING; annotations are deferred. |
| 9 | """ |
| 10 | |
| 11 | from __future__ import annotations |
| 12 | |
| 13 | from abc import ABC, abstractmethod |
| 14 | from dataclasses import dataclass |
| 15 | from typing import TYPE_CHECKING |
| 16 | |
| 17 | if TYPE_CHECKING: |
| 18 | import numpy as np |
| 19 | import torch |
| 20 | |
| 21 | from vllm.config import VllmConfig |
| 22 | from vllm.sampling_params import SamplingParams |
| 23 | from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor |
| 24 | from vllm.v1.worker.gpu.states import RequestState |
| 25 | |
| 26 | |
| 27 | @dataclass(frozen=True) |
| 28 | class LogitsProcRequestState: |
| 29 | """State associated with active requests, shared with logits processors. |
| 30 | |
| 31 | Wraps the model runner's per-slot buffers, which are mutated in place, |
| 32 | so reads always see current values. Processors must treat every field |
| 33 | as read-only. |
| 34 | """ |
| 35 | |
| 36 | device: torch.device |
| 37 | max_num_reqs: int |
| 38 | vocab_size: int |
| 39 | |
| 40 | # [max_num_reqs, max_model_len] committed token ids per request slot. |
| 41 | all_token_ids: StagedWriteTensor |
| 42 | # [max_num_reqs] tokens in the user-provided prompt. |
| 43 | prompt_len: UvaBackedTensor |
| 44 | # [max_num_reqs] tokens fed at the latest (re)fill: the prompt plus any |
| 45 | # partial output on resumption after preemption. |
| 46 | prefill_len: UvaBackedTensor |
| 47 | # [max_num_reqs] prompt_len + output_len; grows as the request progresses. |
| 48 | total_len: StagedWriteTensor |
| 49 | |
| 50 | @classmethod |
| 51 | def from_request_state(cls, req_states: RequestState) -> LogitsProcRequestState: |
| 52 | return cls( |
| 53 | device=req_states.device, |
| 54 | max_num_reqs=req_states.max_num_reqs, |
| 55 | vocab_size=req_states.vocab_size, |
| 56 | all_token_ids=req_states.all_token_ids, |
| 57 | prompt_len=req_states.prompt_len, |
| 58 | prefill_len=req_states.prefill_len, |
| 59 | total_len=req_states.total_len, |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | @dataclass(frozen=True) |
| 64 | class LogitsContext: |
| 65 | """The current step's batch layout, passed to every ``apply()`` call. |
| 66 | |
| 67 | A row is a logits row, not a request: rows are reordered every step, and |
| 68 | under speculative decoding a request owns one row per draft token. |
| 69 | Committed tokens live in ``req_states.all_token_ids`` (valid up to |
| 70 | ``total_len``); this step's draft tokens are only in ``input_ids``. |
| 71 | """ |
| 72 | |
| 73 | # [num_logits_rows] row -> persistent request slot. |
| 74 | expanded_idx_mapping: torch.Tensor |
| 75 | # [num_reqs] batch position -> persistent request slot. |
| 76 | idx_mapping: torch.Tensor |
| 77 | # [num_reqs] batch position -> persistent request slot, on the host, for |
| 78 | # skipping work without a device sync. |
| 79 | idx_mapping_np: np.ndarray |
| 80 | # [num_logits_rows] row -> its offset among the rows of its own request. |
| 81 | expanded_local_pos: torch.Tensor |
| 82 | # [num_logits_rows] token fed to the model at each row's input position. |
| 83 | input_ids: torch.Tensor |
| 84 | # [num_logits_rows] position of each row within its sequence. |
| 85 | pos: torch.Tensor |
| 86 | # [num_reqs] batch position -> upper bound of tokens visible to the token |
| 87 | # being sampled this step, on the host. Exact when spec decoding isn't in use. |
| 88 | # Exact per-row lengths on device are `pos + 1`. |
| 89 | seq_lens_upper_bound_np: np.ndarray |
| 90 | |
| 91 | |
| 92 | class LogitsProcessor(ABC): |
| 93 | """Custom logits processor for Model Runner V2. |
| 94 | |
| 95 | Per-request state is keyed by the request slot index; slots are recycled |
| 96 | through a free list, so per-slot state must be fully (re)initialized in |
| 97 | ``add_request()``. |
| 98 | |
| 99 | ``apply()`` runs after the built-in bias, penalty, bad-words and grammar |
| 100 | stages and before temperature, min_p and top-k/top-p, so it sees unscaled |
| 101 | logits and must not re-inflate grammar-masked tokens. Thinking-budget |
| 102 | forcing runs after ``apply()`` and wins over its edits for requests whose |
| 103 | budget is exhausted. |
| 104 | |
| 105 | State that is constant for a request belongs in ``__init__()`` or |
| 106 | ``add_request()``; ``apply()`` receives only what changes per step. |
| 107 | """ |
| 108 | |
| 109 | def __init__( # noqa: B027 |
| 110 | self, vllm_config: VllmConfig, req_states: LogitsProcRequestState |
| 111 | ): |
| 112 | """Capture what stays constant for the processor's lifetime. |
| 113 | |
| 114 | ``req_states`` exposes the on-device token history and batch constants a |
| 115 | processor may read. Treat it as read-only. |
| 116 | """ |
| 117 | |
| 118 | @classmethod # noqa: B027 |
| 119 | def validate_params(cls, sampling_params: SamplingParams) -> None: |
| 120 | """Raise ``ValueError`` for invalid per-request arguments. |
| 121 | |
| 122 | Runs at request admission, so invalid arguments fail the request |
| 123 | with an error instead of reaching the sampler. |
| 124 | """ |
| 125 | |
| 126 | def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: |
| 127 | """Initialize per-slot state for a request entering the batch. |
| 128 | |
| 129 | The slot may hold a previous occupant's state; overwrite or |
| 130 | neutralize all of it here. |
| 131 | |
| 132 | Returns whether this processor modifies logits for the request. |
| 133 | """ |
| 134 | return True |
| 135 | |
| 136 | def apply_staged_writes(self) -> None: # noqa: B027 |
| 137 | """Flush any host-side writes staged by ``add_request()`` to the device. |
| 138 | |
| 139 | Called once per step before the forward pass, after the model runner |
| 140 | has flushed ``req_states``, so a processor that stages writes here can |
| 141 | read the request's tokens back on device. |
| 142 | """ |
| 143 | |
| 144 | @abstractmethod |
| 145 | def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: |
| 146 | """Modify logits in place or return a new tensor. In-place modification |
| 147 | is preferred for efficiency. |
| 148 | |
| 149 | ``apply()`` is called once for the whole batch, including rows of |
| 150 | requests this processor declined in ``add_request()``, so filter rows |
| 151 | via ``ctx.expanded_idx_mapping``. |
| 152 | |
| 153 | Args: |
| 154 | logits: [num_logits_rows, vocab_size] float32 tensor. |
| 155 | ctx: this step's batch layout. |
| 156 | |
| 157 | """ |
| 158 | raise NotImplementedError |