main
md 571 lines 25.3 KB
Rendered Raw
1 # 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 This document describes how the vLLM engine interacts with logits processors, and the programming model which vLLM supports for implementing logits processors.
8
9 ## Logits Processors Background
10
11 A logits processor adjusts the next-token probability distribution, usually with the intention of steering the model towards a desired type of behavior.
12
13 In vLLM, logits processors operate at batch granularity. During a given engine step, the logits processor consumes a `(num_requests) x (vocab_size)` tensor of raw logits output by the model. For all requests which enable the logits processor, the logits processor applies a transformation to the corresponding row of the logits tensor, while leaving other rows unmodified. The transformed logits tensor is then passed to softmax.
14
15 ## Logits Processors in the vLLM engine
16
17 The vLLM engine's persistent batch data structure maintains a list of loaded logits processors.
18
19 In order to operate on the entire batch at once, each logits processor may maintain metadata about the requests in the batch (i.e. each request's logits-processor-specific configuration settings). Therefore, logits processors are stateful.
20
21 In each engine step, the vLLM engine will (1) update each logits processor's internal state and (2) apply logits processors to the model output logits.
22
23 ### Updating Logits Processor Internal State
24
25 At the beginning of each engine step, the persistent batch may add, discard and/or reorder requests in response to the scheduler output. After the persistent batch has reorganized, the vLLM engine invokes each logits processor's `update_state()` method. This is necessary to ensure that logits processors' internal states are reorganized to match the new persistent batch state at the beginning of the engine step.
26
27 The pseudocode below shows the process by which the vLLM persistent batch notifies each logits processor of changes in batch state:
28
29 ??? code "Model Runner Updates Logits Processor States"
30
31 ``` python
32 # gpu_model_runner.py
33
34 class GPUModelRunner(...):
35
36 ...
37
38 def execute_model(self, scheduler_output, ...):
39 self._update_states(scheduler_output)
40
41 ...
42
43 def _update_states(...):
44
45 ...
46
47 # ...update persistent batch to reflect new/finished requests & reordering
48 # of requests within batch...
49
50 ...
51
52 self.input_batch.refresh_metadata()
53
54
55 # gpu_input_batch.py
56
57 class InputBatch:
58
59 ...
60
61 def refresh_metadata(self):
62
63 ...
64
65 # Update each logits processor's state to reflect persistent batch state
66 batch_update = self.batch_update_builder.get_and_reset(self.num_reqs)
67 for logit_proc in self.logitsprocs.all:
68 logit_proc.update_state(batch_update)
69
70 ...
71
72
73 # vllm/v1/sample/logits_processor/interface.py
74
75 @dataclass(frozen=True)
76 class BatchUpdate:
77 # Batch state-change data structure which is passed to logits processors'
78 # update_state() methods
79
80 batch_size: int
81
82 removed: Sequence[RemovedRequest]
83 added: Sequence[AddedRequest]
84 moved: Sequence[MovedRequest]
85
86 ```
87
88 ### Applying Logits Processors to the Model Output Logits
89
90 After updating persistent batch state, the vLLM model runner performs model inference to obtain logits. Then, the model runner invokes the sampler against the logits. In turn, part of the sampler's operation is to invoke the logits processors' `apply()` methods against the model output logit processors, yielding transformed logits (the `apply()` methods may modify the logits in-place or out-of-place, although in-place is more memory-efficient). This process is shown in the pseudocode below.
91
92 Note that the sampler will access the logits processors via `SamplingMetadata.logitsprocs`. When the vLLM engine constructs `SamplingMetadata` (not shown in the code below), the reference to the list of logits processors is passed from the persistent batch data structure to `SamplingMetadata`.
93
94 ??? code "Apply logits processors to model output logits"
95
96 ``` python
97 # gpu_model_runner.py
98
99 class GPUModelRunner(...):
100
101 ...
102
103 def execute_model(self, scheduler_output, ...):
104 # (discussed in previous section)
105 self._update_states(scheduler_output)
106
107 ...
108
109 # ...run model inference to obtain logits...
110
111 ...
112
113 # Invoke sampler, which applies logits processors
114 sampler_output = self.sampler(logits=logits,
115 sampling_metadata=sampling_metadata)
116
117 ...
118
119
120 # sampler.py
121
122 class Sampler(nn.Module):
123
124 ...
125
126 def forward(self, logits, sampling_metadata):
127
128 ...
129
130 # Apply non-argmax-invariant logits processors to model output logits
131 for processor in (sampling_metadata.logitsprocs.non_argmax_invariant):
132 logits = processor.apply(logits)
133
134 sampled = self.sample(logits, sampling_metadata)
135
136 ...
137
138 # ...return sampler output data structure...
139
140
141 def sample(self, logits, sampling_metadata)
142
143 ...
144
145 # ...exit early if all requests are greedy-sampling...
146
147 ...
148
149 # Apply argmax-invariant logits processors
150 for processor in sampling_metadata.logitsprocs.argmax_invariant:
151 logits = processor.apply(logits)
152
153 ...
154
155 # ...perform sampling and return sampling result...
156 ```
157
158 At sampling time, the sampler checks whether all requests in the persistent batch employ greedy sampling. If that is the case, the sampler saves compute by skipping "argmax-invariant" logits processors. Here, "argmax" is shorthand for the token ID with the highest logit value in a given row of the logits tensor (i.e. the token which the model weighted the highest for a given request).
159
160 * An **argmax-invariant logits processor** is a logits processor (such as Min-P) which does not modify the argmax. For example, a logits processor which masks out the lowest-probability tokens will not change which token ID has the max logit. Greedy sampling always picks the highest-logit-value token ID, and so conceptually an argmax-invariant logits processor can be skipped for greedy sampling requests.
161
162 * A **non-argmax-invariant logits processor** is a logits processor which may modify the argmax. For example, a logits processor which masks all tokens except for EOS after a certain number of steps in order to force decoding to terminate might end up masking the max-logit-value token and therefore change the argmax. Conceptually, these logits processors cannot be skipped for greedy sampling requests.
163
164 The vLLM logits processor abstraction requires the engine to apply logits processors at batch granularity; therefore in practice the argmax-invariant logits processors can only be skipped when the entire batch uses greedy sampling.
165
166 ## Logits Processor Programming Model
167
168 The previous sections alluded to the interfaces which vLLM logits processors must support. This section introduces in full the programming model for implementing logits processors that are compatible with the vLLM engine, including the `LogitsProcessor` base class and its interface methods as well as the `BatchUpdate` data structure for representing persistent batch state changes, both of which are shown in the code below:
169
170 ??? code "`LogitsProcessor` base class and `BatchUpdate` data structure"
171
172 ``` python
173 from abc import ABC, abstractmethod
174 from collections.abc import Sequence
175 from dataclasses import dataclass
176 from enum import Enum, auto
177 from typing import TYPE_CHECKING
178
179 import torch
180
181 from vllm import SamplingParams
182
183 if TYPE_CHECKING:
184 from vllm.config import VllmConfig
185
186
187 class MoveDirectionality(Enum):
188 # One-way i1->i2 req move within batch
189 UNIDIRECTIONAL = auto()
190 # Two-way i1<->i2 req swap within batch
191 SWAP = auto()
192
193
194 # (index, params, prompt_tok_ids, output_tok_ids) tuples for new
195 # requests added to the batch.
196 AddedRequest = tuple[int, SamplingParams, list[int], list[int]]
197
198 # (index 1, index 2, directionality) tuples representing
199 # one-way moves or two-way swaps of requests in batch
200 MovedRequest = tuple[int, int, MoveDirectionality]
201
202 # Batch indices of any removed requests.
203 RemovedRequest = int
204
205
206 @dataclass(frozen=True)
207 class BatchUpdate:
208 """Persistent batch state change info for logitsprocs"""
209 batch_size: int # Current num reqs in batch
210
211 # Metadata for requests added to, removed from, and moved
212 # within the persistent batch.
213 #
214 # Key assumption: the `output_tok_ids` list (which is an element of each
215 # tuple in `added`) is a reference to the request's running output tokens
216 # list; via this reference, the logits processors always see the latest
217 # list of generated output tokens
218 removed: Sequence[RemovedRequest]
219 moved: Sequence[MovedRequest]
220 added: Sequence[AddedRequest]
221
222
223 class LogitsProcessor(ABC):
224
225 @abstractmethod
226 def __init__(self, vllm_config: "VllmConfig", device: torch.device,
227 is_pin_memory: bool) -> None:
228 raise NotImplementedError
229
230 @abstractmethod
231 def apply(self, logits: torch.Tensor) -> torch.Tensor:
232 raise NotImplementedError
233
234 @abstractmethod
235 def is_argmax_invariant(self) -> bool:
236 """True if logits processor has no impact on the
237 argmax computation in greedy sampling.
238 NOTE: may or may not have the same value for all
239 instances of a given LogitsProcessor subclass,
240 depending on subclass implementation.
241 """
242 raise NotImplementedError
243
244 @abstractmethod
245 def update_state(
246 self,
247 batch_update: "BatchUpdate" | None,
248 ) -> None:
249 """Called when there are new output tokens, prior
250 to each forward pass.
251
252 Args:
253 batch_update is non-None iff there have been
254 changes to the batch makeup.
255 """
256 raise NotImplementedError
257
258 @classmethod
259 def validate_params(cls, sampling_params: SamplingParams):
260 """Validate sampling params for this logits processor.
261
262 Raise ValueError for invalid ones.
263 """
264 return None
265
266 ```
267
268 A vLLM logits processor must subclass `LogitsProcessor` and define (at minimum) the following methods:
269
270 * `__init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool)`
271 * `vllm_config`: engine configuration data structure
272 * `device`: hardware accelerator device info
273 * `is_pin_memory`: flag indicating whether pin memory is available to support logits processor implementation
274
275 * `apply(self, logits: torch.Tensor) -> torch.Tensor`:
276 * Consume a `(num_requests) x (vocab_size)` logits tensor (`logits`)
277 * Apply logits processor transformation at batch granularity
278 * Return a transformed `(num_requests) x (vocab_size)` logits tensor
279 * You can modify the input logits processors in-place or out-of-place; in-place is more memory-efficient
280
281 * `is_argmax_invariant(self) -> bool`:
282 * Return `True` if the logits processor is argmax invariant (never changes what is the highest-logit-value token ID for a given request), `False` if the logits processor may modify argmax
283 * `is_argmax_invariant()` is evaluated once at startup; if `True`, vLLM will skip applying this logits processor in a given step when all requests use greedy sampling
284
285 * `update_state(self, batch_update: "BatchUpdate" | None) -> None`:
286 * Consume a `BatchUpdate` data structure representing persistent batch state changes at the beginning of the current engine step
287 * Use the `BatchUpdate` members to update logits processor internal state
288 * **Note:** batch update data structure may be `None`, signaling no change to the batch constituents. In this case, the LogitsProcessor might still want to update its state based on the updated `output_token_ids` lists that it could have retained when they were added.
289
290 * `validate_params(cls, sampling_params: SamplingParams)`:
291 * Raise `ValueError` if `SamplingParams` has invalid arguments (especially custom arguments) used by logits processor.
292 * When request is sent to entrypoint, `validate_params()` will validate `SamplingParams` and refuse request with invalid arguments.
293
294 ### `BatchUpdate` data structure
295
296 The `BatchUpdate` abstraction models the persistent batch as a list of requests, supporting the following operations to change batch state (note that the order in which the operations are mentioned below reflects the order in which they should be processed in `update_state()`):
297
298 * **Remove:** remove (without replacement) request at index `i`
299
300 * A Remove is represented in `Batchupdate.removed` by an `int` (representing `i`)
301
302 * Effect of remove-at-index on batch:
303
304 ``` text
305 Batch: [A,B,C]
306 Remove @ i: 1
307
308 =>
309
310 New Batch: [A,x,C] # Discard B and leave an empty slot
311 ```
312
313 * **Add:** add (or replace existing request with) a new request at index `i`. If a request is replaced, its associated state should be discarded.
314
315 * An Add is represented in `Batchupdate.added` as a tuple of
316
317 ``` text
318 (index, new request SamplingParams, prompt token ids, output token ids)
319 ```
320
321 * `prompt token ids` and `output token ids` are references to the request's prompt token ids and output token ids lists, respectively. Note that the output token ids list grows with each engine step, and this growth is visible to the logits processor because output token ids are passed by reference. **This is important for LogitsProcessors that take into account the tokens generated so far**.
322
323 * The implementation of the particular logits processor subclass determines whether or how the fields in the added request tuple are digested into an internal representation. For example, a logits processor that does not utilize prompt or output token ids may only need to utilize `index` and `SamplingParams` and discard the other tuple fields
324
325 * If index `i` currently holds a request, a replacement occurs:
326
327 ``` text
328 Batch: [A,B,C]
329 New request to be added @ i: D @ 1
330
331 =>
332
333 New Batch: [A,D,C] # Add D, discard B
334 ```
335
336 * If index `i` does not currently hold a request (because `i` is out of bounds of the current batch size):
337
338 ``` text
339 Batch: [A,B,C]
340 New request to be added @ i: D @ 3
341
342 =>
343
344 New Batch: [A,B,C,D] # Add D, extending batch
345 ```
346
347 * **Move:** move request at index `s` to index `d` OR swap requests at indices `s` and `d`
348
349 * A Move is represented in `Batchupdate.moved` as a tuple of
350
351 ``` text
352 (s, d, UNIDIRECTIONAL or SWAP)
353 ```
354
355 * If the Move specifies `UNIDIRECTIONAL`:
356
357 * The request at index `s` is moved to index `d`; index `s` becomes an empty slot
358
359 ``` text
360 Batch: [A,x,C,D]
361 Unidirectionally Move s -> d: 3 -> 1
362
363 =>
364
365 New Batch: [A,D,C,x] # Move D to 1, leaving empty slot at 3
366 ```
367
368 * If another request already resided at index `d`, it is replaced and discarded
369
370 ``` text
371 Batch: [A,B,C,D]
372 Unidirectionally Move s -> d: 3 -> 1
373
374 =>
375
376 New Batch: [A,D,C,x] # Move D to 1, discarding B and leaving empty slot at 3
377 ```
378
379 * If the Move specifies `SWAP`, the requests at `s` and `d` exchange indices
380
381 ``` text
382 Batch: [A,B,C,D]
383 Swap Move s <-> d: 3 <-> 1
384
385 =>
386
387 New Batch: [A,D,C,B] # Swap B and D
388 ```
389
390 Additionally, the `BatchUpdate` data structure includes a representation (`batch_size`) of the size of the persistent batch at the beginning of the engine step.
391
392 ### How the vLLM engine builds the `BatchUpdate` data structure
393
394 Logits processor `update_state()` implementations should assume the following model for how the model runner updates persistent batch state (expressed here in terms of the `BatchUpdate` abstraction):
395
396 1. Identify indices of requests which finished in the current engine step
397
398 2. Identify new requests introduced in the current step
399
400 3. Use Add operations to replace as many finished requests with new requests, in order of increasing index of the replaced request starting with the lowest index
401
402 4. Based on the relative number of new and finished requests:
403
404 1. If the numbers of new and finished requests are the same, proceed to next step
405
406 2. *If there are more new requests than finished requests:* apply Add operations to extend the batch with the remaining new requests which did not replace finished requests. Assign consecutive indices to these new requests, starting with `current_max_batch_index + 1`
407
408 3. *If there are fewer new requests than finished requests:*
409
410 * Apply Remove operations to finished requests which were not replaced with new requests. These removed request indices will necessarily be greater than the greatest index of the finished requests which were replaced in the previous step. The Removes may leave the batch in a non-contiguous state
411
412 * **"Condense" the batch to be contiguous:** starting with the lowest-index empty slot (which was caused by a Remove), apply a Unidirectional Move from the current highest non-empty slot in the batch to fill the empty slot. Proceed with additional Unidirectional Move operations in order of increasing empty slot destination index and decreasing non-empty slot source index until the batch is contiguous
413
414 * **Shrink the batch:** a side effect of condensing the batch is that empty slots resulting from Remove operations are grouped in a contiguous block at the end of the batch array. Thus, after condensing, update `BatchUpdate.batch_size` to reflect the number of non-empty slots
415
416 5. Reorder the batch for improved efficiency. Depending on the attention backend implementation and the current characteristics of the batch, zero or more Swap Move operations may be applied to reorder the batch
417
418 Notes:
419
420 * A logits processor `update_state()` method must process batch update operations in the following order: removes, adds, moves
421
422 * The index argument for Add operations refers to the index *at the time the Add occurred*, i.e. before any Move operations
423 * Example: if a request is Added at index 5 and then swapped with index 3, the Add operation in `BatchUpdate.added` will be associated with index 5 not 3
424 * In other words Move operations can be assumed to be applied after Adds and Removes
425
426 * Move operations can be assumed to be applied in the order in which they appear in `BatchUpdate.moved`
427
428 * If there are no new/finished requests and there is no batch reordering, then the batch update for the logits processors will be `None`
429
430 #### Example: Batch Update with Fewer New Requests Than Finished Requests
431
432 The following example models an engine step where 1 new request is introduced and 2 finished requests are eliminated, additionally the attention backend performs a swap to optimize the batch ordering.
433
434 ``` text
435 Batch state (beginning of engine step): [A,B,C,D]
436 Batch size: 4
437
438 New requests: E
439
440 Finished requests: A, C
441
442 Processing steps (using BatchUpdate abstraction):
443
444 1. Add E at index 0
445
446 [E,B,C,D] # Discard A
447 Batch size: 4
448
449 2. Remove at index 2
450
451 [E,B,x,D] # Discard C, empty slot at index 2
452 Batch size: 4
453
454 3. Condense batch with a Unidirectional Move 3 -> 2 operation and shrink batch
455
456 [E,B,D] x # Empty slot is now outside batch
457 Batch size: 3
458
459 4. Attention backend optimization: reorder batch with Swap 0 <-> 1
460
461 [B,E,D]
462 Batch size: 3
463
464 ```
465
466 The resulting `BatchUpdate` data structure will look like
467
468 ``` text
469 BatchUpdate instance
470 * added: [(0,E's SamplingParams,E's prompt tokens ref,E's output tokens ref)]
471 * removed: [2] # request C was removed without replacement
472 * moved: [(3,2,UNIDIRECTIONAL),(0,1,SWAP)]
473 ```
474
475 #### Example: Batch Update with More New Requests Than Finished Requests
476
477 The following example models an engine step where 2 new requests are introduced and 1 finished request is eliminated, additionally the attention backend performs a swap to optimize the batch ordering.
478
479 ``` text
480 Batch state (beginning of engine step): [A,B,C,D]
481 Batch size: 4
482
483 New requests: E,F
484
485 Finished requests: C
486
487 Processing steps (using BatchUpdate abstraction):
488
489 1. Add E at index 2
490
491 [A,B,E,D] # Discard C
492 Batch size: 4
493
494 2. Add F at index 4 (current max batch index + 1)
495
496 [A,B,E,D,F] # Extend batch by 1
497 Batch size: 5
498
499 4. Attention backend optimization: reorder batch with Swap 0 <-> 1
500
501 [B,A,E,D,F]
502 Batch size: 5
503
504 ```
505
506 Note that batch condensation is skipped because there are no empty slots left behind by Remove operations.
507
508 The resulting `BatchUpdate` data structure will look like
509
510 ``` text
511 BatchUpdate instance
512 * added: [(2,E's SamplingParams,E's prompt tokens ref,E's output tokens ref),(4,F's SamplingParams,F's prompt tokens ref,F's output tokens ref)]
513 * removed: [] # no requests were removed without replacement
514 * moved: [(0,1,SWAP)]
515 ```
516
517 ## How to Introduce a New Logits Processor to vLLM
518
519 ### Best Practices for Writing Built-In Logits Processors
520
521 * Write efficient `apply()` and `update_state()` implementations in light of the fact that logits processors operate at batch granularity
522 * For example, you may be able to use efficient vectorized operations to implement `apply()` or update internal state vectors in `update_state()`
523 * However, if you think that a logits processor may be used infrequently, it may be appropriate to use a "sparse" representation of request state i.e. the class can represent request configuration using a dictionary which only stores metadata about requests that enable the logits processor
524
525 * It is up to the logits processor author to determine:
526
527 1. **The per-request attributes which configure the logits processor's behavior against that request.** For example, if you are writing a new built-in logits processor for vLLM, you may or may not need to add additional fields to `SamplingParams` and the vLLM REST API
528
529 2. **The conditions under which the logits processor is or is not enabled on a per-request basis.** Unless your intention is for the built-in 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`. Try to save compute and memory for requests which disable the logits processor
530
531 3. **The conditions under which the logits processor is short-circuited at the batch level.** Even if you have defined a way to disable the built-in logits processor at the request level, it may be difficult to translate this into compute savings i.e. if your `update_state()` and `apply()` implementations use efficient vectorized implementations that operate on the whole persistent batch in a single command. For example, you cannot skip an entire vectorized operation in `apply()` just because one request disabled the logits processor. To save compute in the edge-case where no running requests utilize the built-in logits processor, we recommend designing `apply()` to return the unmodified input tensor if all requests have the logits processor disabled. Similarly, consider whether steps can be skipped in `update_state()` if no requests enable the logits processor
532
533 * Additionally, an easy way to save compute in `update_state()` is to exit early when the batch_update is `None`
534
535 * Ensure that the logits processor `update_state` method discards information about finished requests (i.e. requests which are replaced by an Add or which are subject to a Remove)
536
537 * `is_argmax_invariant()` can be hard-coded to `True` or `False` if the logits processor has consistent behavior. However the argmax invariance may also be determined programmatically (i.e. if your logits processor is user-customizable in some way that impacts whether the logits processor is argmax invariant). For this reason, `is_argmax_invariant()` is not a class method
538
539 ### Built-In Logits Processors
540
541 Built-in logits processors are always loaded when the vLLM engine starts. See the existing vLLM built-in logits processors in `vllm/v1/sample/logits_processor/builtin.py` for examples of how to write a new built-in vLLM logits processor. It makes sense to write a PR to introduce a new logits processor as a built-in if it is likely to be useful to a wide audience. vLLM currently employs the following built-in logits processors based on the programming model described above:
542
543 * Min-P
544
545 * Logit bias
546
547 * Min-tokens
548
549 Review these logits processor implementations for guidance on writing built-in logits processors.
550
551 Additionally, the following logits-processor-like functionalities are hard-coded into the sampler and do not yet utilize the programming model described above. Most of them will be refactored to use the aforementioned logits processor programming model.
552
553 * Allowed token IDs
554
555 * Bad words
556
557 * Repetition penalty
558
559 * Frequency penalty
560
561 * Presence penalty
562
563 * Temperature
564
565 * Top-K
566
567 * Top-P
568
569 ### Custom Logits Processors
570
571 vLLM can be augmented with [user-provided custom logits processors](../features/custom_logitsprocs.md).