| 1 | # Custom Logits Processors |
| 2 | |
| 3 | !!! important |
| 4 | Some logits processors design changes are still in progress and the API may |
| 5 | change in the near future. We hope to stabilize this part of the API soon |
| 6 | |
| 7 | A "custom" logits processor is written by a user of vLLM and is loaded into vLLM at initialization without needing to modify or recompile the vLLM source code. It is the opposite of a built-in logits processor. |
| 8 | |
| 9 | This document shows how to write, load and use a custom logits processor. |
| 10 | |
| 11 | !!! note |
| 12 | This document covers the Model Runner V2 (MRV2) interface. The legacy V1 |
| 13 | model runner (`VLLM_USE_V2_MODEL_RUNNER=0`) uses a different interface, |
| 14 | `vllm.v1.sample.logits_processor.LogitsProcessor`, which is not covered |
| 15 | here. |
| 16 | |
| 17 | ## Logits Processors Background |
| 18 | |
| 19 | A logits processor adjusts the next-token probability distribution, usually with the intention of steering the model towards a desired type of behavior. |
| 20 | |
| 21 | In vLLM, logits processors operate at batch granularity. During a given engine step, the logits processor consumes a `(num_logits_rows) x (vocab_size)` tensor of raw logits output by the model. Note that a logits row is not a request: rows are reordered every step, and under speculative decoding a request owns one row per draft token. The logits processor applies a transformation to the rows of the logits tensor, while leaving other rows unmodified. The transformed logits tensor is then used for sampling. |
| 22 | |
| 23 | ## Creating a Custom Logits Processor |
| 24 | |
| 25 | Custom logits processors must subclass `vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor` and define (at minimum) the following methods: |
| 26 | |
| 27 | * `__init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState)`: |
| 28 | * `vllm_config`: engine configuration data structure |
| 29 | * `req_states`: a narrow, read-only view of the persistent batch, exposing the on-device token history (`all_token_ids`, `prompt_len`, `prefill_len`, `total_len`) plus `device`, `max_num_reqs` and `vocab_size` |
| 30 | |
| 31 | * `add_request(self, req_idx, sampling_params) -> bool`: |
| 32 | * Initialize per-slot state for a request entering the batch. Slots are recycled through a free list, so per-slot state must be fully (re)initialized here; there is no removal hook, since freed slots are never read |
| 33 | * Return whether this processor modifies logits for the request. The sampler skips the logits-processing pipeline for batches in which no request needs it |
| 34 | |
| 35 | * `apply_staged_writes(self) -> None` (optional): |
| 36 | * Flush host-side writes staged by `add_request()` to the device; called once per step before the forward pass |
| 37 | |
| 38 | * `validate_params(cls, sampling_params) -> None` (optional classmethod): |
| 39 | * Raise `ValueError` for invalid per-request arguments (especially custom arguments); runs at request admission, so invalid arguments fail the request instead of reaching the sampler |
| 40 | |
| 41 | * `apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor`: |
| 42 | * Consume a `(num_logits_rows) x (vocab_size)` logits tensor and the step's batch layout (`ctx`: row-to-slot mappings, `input_ids`, positions) |
| 43 | * Apply logits processor transformation at batch granularity |
| 44 | * Return a transformed logits tensor. You can modify the input logits in-place or out-of-place; in-place is more memory-efficient |
| 45 | |
| 46 | ### Passing Custom Argument to a Custom Logits Processor |
| 47 | |
| 48 | Unlike built-in logits processors, custom logits processors may require configuration arguments that are not hard-coded into `SamplingParams` or the vLLM server REST API. To solve this problem, custom logits processors may leverage vLLM [custom arguments](./custom_arguments.md) support to receive configuration settings from the user (although you are also free to design a custom logits processor which utilizes the pre-existing fields in `SamplingParams`.) |
| 49 | |
| 50 | ### Example Custom Logits Processor Implementation |
| 51 | |
| 52 | The contrived example below implements a custom logits processor which masks out all tokens except for one (`target_token`) for the requests that enable it. The processor is disabled for any request that does not specify `target_token`; `add_request()` reports this via its return value, so the sampler can skip the logits-processing pipeline entirely when no request enables the processor. To determine whether the logits processor is enabled and which token to leave unmasked, the processor checks `SamplingParams.extra_args` for a `target_token` custom argument associated with each request: |
| 53 | |
| 54 | ??? code "Example custom logits processor definition" |
| 55 | |
| 56 | ``` python |
| 57 | import torch |
| 58 | from vllm.config import VllmConfig |
| 59 | from vllm.sampling_params import SamplingParams |
| 60 | from vllm.v1.worker.gpu.sample.logits_processor import ( |
| 61 | LogitsProcRequestState, |
| 62 | LogitsContext, |
| 63 | LogitsProcessor, |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | class TargetTokenLogitsProcessor(LogitsProcessor): |
| 68 | """Masks out all tokens except `target_token` (a per-request custom |
| 69 | argument); requests without it are left alone.""" |
| 70 | |
| 71 | def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState): |
| 72 | # Per-slot target; -1 means disabled. Staged on the host and |
| 73 | # flushed to the device once per step in apply_staged_writes(). |
| 74 | self.target_token = torch.full( |
| 75 | (req_states.max_num_reqs,), -1, dtype=torch.int64 |
| 76 | ) |
| 77 | self.target_token_dev = torch.full( |
| 78 | (req_states.max_num_reqs,), -1, dtype=torch.int64, device=req_states.device |
| 79 | ) |
| 80 | |
| 81 | def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: |
| 82 | target = (sampling_params.extra_args or {}).get("target_token") |
| 83 | self.target_token[req_idx] = target if target is not None else -1 |
| 84 | return target is not None |
| 85 | |
| 86 | def apply_staged_writes(self) -> None: |
| 87 | self.target_token_dev.copy_(self.target_token, non_blocking=True) |
| 88 | |
| 89 | def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: |
| 90 | # ctx.expanded_idx_mapping maps each logits row to its request slot. |
| 91 | cols = self.target_token_dev[ctx.expanded_idx_mapping.long()] |
| 92 | rows = torch.nonzero(cols >= 0).squeeze(1) |
| 93 | if rows.numel() == 0: |
| 94 | return logits |
| 95 | kept = logits[rows, cols[rows]].clone() |
| 96 | logits[rows] = float("-inf") |
| 97 | logits[rows, cols[rows]] = kept |
| 98 | return logits |
| 99 | ``` |
| 100 | |
| 101 | Per-request state is keyed by the request slot index, and slots are recycled through a free list. The example keeps a per-slot state where `-1` means disabled; note that `add_request()` overwrites the slot's entry unconditionally, which is what keeps recycled slots from leaking state between requests. |
| 102 | |
| 103 | ## Ways to Load Your Custom Logits Processor in vLLM |
| 104 | |
| 105 | Logits processors are loaded at initialization. Critically, the set of loaded logits processors cannot be modified after the vLLM engine finishes loading, and new logits processors cannot be loaded on-demand for individual requests. |
| 106 | |
| 107 | Loaded classes are validated against the Model Runner V2 interface; a class that does not subclass `vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor` is rejected with a clear error. |
| 108 | |
| 109 | This section details different ways of making your logits processor visible to vLLM and triggering vLLM to load your logits processor. |
| 110 | |
| 111 | ### Method 1: Pass the Custom Logits Processor Fully-Qualified Class Name (FQCN) to vLLM at Initialization Time |
| 112 | |
| 113 | This method is supported in both offline and online vLLM usage scenarios. The custom logits processor's FQCN (in the form of `dotted.path.to.module:ClassName`) can be passed as an argument to the `LLM` and `AsyncLLM` Python constructors, or as a CLI argument to `vllm serve` with the following syntax |
| 114 | |
| 115 | ``` bash |
| 116 | vllm serve ... --logits_processors <logits processor 1> <logits processor 2> ... |
| 117 | ``` |
| 118 | |
| 119 | The only requirements on the FQCN are |
| 120 | |
| 121 | 1. Python's `importlib.import_module()` must be able to resolve the dotted path portion of the FQCN and load it as a module |
| 122 | |
| 123 | 2. The class-name portion of the FQCN must be possible to import from the loaded module |
| 124 | |
| 125 | 3. The object pointed to by the FQCN must be a subclass of `LogitsProcessor` |
| 126 | |
| 127 | See examples below: |
| 128 | |
| 129 | ??? code "Passing custom logits processor FQCN to `LLM` in Python" |
| 130 | |
| 131 | ``` python |
| 132 | # Pass in FQCN |
| 133 | llm = LLM( |
| 134 | model="facebook/opt-125m", |
| 135 | logits_processors=["your.module.path:TargetTokenLogitsProcessor"], |
| 136 | ) |
| 137 | ``` |
| 138 | |
| 139 | ??? code "Passing custom logits processor FQCN to `AsyncLLM` in Python" |
| 140 | |
| 141 | ``` python |
| 142 | # Pass in FQCN |
| 143 | engine_args = AsyncEngineArgs(model="facebook/opt-125m", |
| 144 | logits_processors=["your.module.path:TargetTokenLogitsProcessor"]) |
| 145 | async_llm = AsyncLLM.from_engine_args(engine_args) |
| 146 | ``` |
| 147 | |
| 148 | ??? code "Passing custom logits processor FQCN to vLLM server via CLI" |
| 149 | |
| 150 | ```bash |
| 151 | vllm serve facebook/opt-125m --logits_processors your.module.path:TargetTokenLogitsProcessor |
| 152 | ``` |
| 153 | |
| 154 | ### Method 2: Automatically Detect Custom Logits Processors Installed in Your Python Environment As Entry Points |
| 155 | |
| 156 | [`setuptools`](https://setuptools.pypa.io/en/latest/userguide/entry_point.html) can enable installed packages to make themselves available as plugins to other Python programs, via pieces of metadata known as "entry points". |
| 157 | |
| 158 | During initialization, vLLM automatically scans the `vllm.logits_processors` entry point group and loads any installed logits processors which it finds. |
| 159 | |
| 160 | Suppose that you have developed a Python package that holds your custom logits processors. You can expose each logits processor to vLLM by adding a unique entrypoint for each logits processor to your logits processor Python package. The example below shows how to add an entrypoint to your project's `pyproject.toml` file: |
| 161 | |
| 162 | ??? code "Exposing a custom logits processor as a Python entrypoint" |
| 163 | |
| 164 | ``` toml |
| 165 | [project.entry-points."vllm.logits_processors"] |
| 166 | target_token_logits_processor = "your.module.path:TargetTokenLogitsProcessor" |
| 167 | ``` |
| 168 | |
| 169 | Once your package is installed, your custom logits processor will be loaded automatically whenever vLLM is initialized. You do *not* need to pass the custom logits processor to the `LLM` or `AsyncLLM` constructors or to the vLLM server explicitly at initialization time if your logits processor is exposed as an entry point. |
| 170 | |
| 171 | !!! note |
| 172 | vLLM will *always* load *all* logits processors which are exposed via entrypoints under the `vllm.logits_processors` grouping. |
| 173 | |
| 174 | ### Method 3 (Offline-only): Pass a Python Class Object to the vLLM Constructor |
| 175 | |
| 176 | You can pass one or more custom logits processor class objects to the `LLM` and `AsyncLLM` constructors. This option is very flexible, as the logits processor classes may either be (1) defined locally within the same Python source file where `LLM` or `AsyncLLM` is instantiated, or (2) imported from a Python package. |
| 177 | |
| 178 | ??? code "Passing custom logits processor class object to `LLM` or `AsyncLLM` in Python" |
| 179 | |
| 180 | ``` python |
| 181 | # Import custom logits processor |
| 182 | from some.module import TargetTokenLogitsProcessor |
| 183 | |
| 184 | # ...or... |
| 185 | |
| 186 | # Define custom logits processor locally |
| 187 | from vllm.v1.worker.gpu.sample.logits_processor import LogitsProcessor |
| 188 | |
| 189 | class TargetTokenLogitsProcessor(LogitsProcessor): |
| 190 | # See TargetTokenLogitsProcessor implementation above |
| 191 | ... |
| 192 | |
| 193 | # Pass class object to LLM constructor |
| 194 | llm = LLM( |
| 195 | model="facebook/opt-125m", |
| 196 | logits_processors=[TargetTokenLogitsProcessor], |
| 197 | ) |
| 198 | |
| 199 | # Pass class object to AsyncLLM constructor |
| 200 | engine_args = AsyncEngineArgs(model="facebook/opt-125m", |
| 201 | logits_processors=[TargetTokenLogitsProcessor]) |
| 202 | async_llm = AsyncLLM.from_engine_args(engine_args) |
| 203 | ``` |
| 204 | |
| 205 | ## Invoking a Custom Logits Processor Against a Request |
| 206 | |
| 207 | The design of the custom logits processor determines whether the logits processor must be enabled/disabled for a given request, and what arguments must be provided to configure the logits processor. |
| 208 | |
| 209 | The examples below show how a user would pass a custom argument (`target_token`) to `TargetTokenLogitsProcessor` in order to (1) enable the logits processor for that particular request and (2) control the logits processor's behavior. |
| 210 | |
| 211 | ??? code "vLLM REST API: configure custom logits processor for a request" |
| 212 | |
| 213 | ``` bash |
| 214 | curl http://localhost:8000/v1/completions \ |
| 215 | -H "Content-Type: application/json" \ |
| 216 | -d '{ |
| 217 | "model": "Qwen/Qwen2.5-1.5B-Instruct", |
| 218 | ... |
| 219 | "vllm_xargs": {"target_token": 67} |
| 220 | }' |
| 221 | ``` |
| 222 | |
| 223 | ??? code "OpenAI SDK: configure custom logits processor for a request" |
| 224 | |
| 225 | ``` python |
| 226 | batch = await client.completions.create( |
| 227 | model="Qwen/Qwen2.5-1.5B-Instruct", |
| 228 | ..., |
| 229 | extra_body={ |
| 230 | "vllm_xargs": { |
| 231 | "target_token": 67 |
| 232 | } |
| 233 | } |
| 234 | ) |
| 235 | ``` |
| 236 | |
| 237 | ??? code "Offline: configure custom logits processor for an `LLM` request" |
| 238 | |
| 239 | ``` python |
| 240 | outputs_logitproc = llm.generate("your prompt", |
| 241 | SamplingParams(..., |
| 242 | extra_args={"target_token": 67})) |
| 243 | ``` |
| 244 | |
| 245 | ??? code "Offline: configure custom logits processor for an `AsyncLLM` request" |
| 246 | |
| 247 | ``` python |
| 248 | async for out in engine.generate(request_id="your request id", |
| 249 | prompt="your prompt", |
| 250 | sampling_params=SamplingParams(..., |
| 251 | extra_args={"target_token": 67})): |
| 252 | |
| 253 | # Process async request outputs |
| 254 | ... |
| 255 | ``` |
| 256 | |
| 257 | ## Best Practices for Writing Custom Logits Processors |
| 258 | |
| 259 | * Write an efficient `apply()` implementation in light of the fact that logits processors operate at batch granularity. For example, you may be able to use efficient vectorized operations to implement `apply()`, and to stage per-request state into contiguous tensors in `add_request()` instead of per-row Python loops |
| 260 | |
| 261 | * It is up to the logits processor author to determine: |
| 262 | |
| 263 | 1. **The per-request attributes which configure the logits processor's behavior against that request.** Your custom logits processor's `add_request()` override determines how `SamplingParams` fields are mapped into logits processor state |
| 264 | |
| 265 | 2. **The conditions under which the logits processor is or is not enabled on a per-request basis.** Unless your intention is for the custom logits processor to act on all requests all the time, you should write your logits processor in such a way that it is possible to disable the logits processor for a given request, i.e. by defaulting an argument to `None` or by passing in a specific do-nothing argument value i.e. `0.0`. Return `False` from `add_request()` for such requests so that the sampler can skip the logits-processing pipeline for batches in which no request needs it |
| 266 | |
| 267 | * Stage host-side writes in `add_request()` and flush them once per step in `apply_staged_writes()`, instead of writing device tensors directly; this is the discipline every built-in sampler state follows |
| 268 | |
| 269 | * Since there is no removal hook, `add_request()` must fully (re)initialize the slot's state; do not rely on state left behind by the slot's previous occupant |