A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[V1] Logits processor docs (#22919)
Signed-off-by: Andrew Feldman <afeldman@redhat.com> Signed-off-by: afeldman-nm <156691304+afeldman-nm@users.noreply.github.com> Co-authored-by: Joseph Marinier <Joseph.Marinier@gmail.com>
afeldman-nm committed
Sep 17, 2025 at 14:53 UTC
7ae9887542bf0fc45cc44de583290c41234a09c4
7 files changed
+1065
-16
docs/design/logits_processors.md
new
+559
@@ -0,0 +1,559 @@
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_metadta)
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, Optional
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: Optional["BatchUpdate"],
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
+ ```
259
+
260
+A vLLM logits processor must subclass `LogitsProcessor` and define (at minimum) the following methods:
261
+
262
+* `__init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool)`
263
+ * `vllm_config`: engine configuration data structure
264
+ * `device`: hardware accelerator device info
265
+ * `is_pin_memory`: flag indicating whether pin memory is available to support logits processor implementation
266
+
267
+* `apply(self, logits: torch.Tensor) -> torch.Tensor`:
268
+ * Consume a `(num_requests) x (vocab_size)` logits tensor (`logits`)
269
+ * Apply logits processor transformation at batch granularity
270
+ * Return a transformed `(num_requests) x (vocab_size)` logits tensor
271
+ * You can modify the input logits processors in-place or out-of-place; in-place is more memory-efficient
272
+
273
+* `is_argmax_invariant(self) -> bool`:
274
+ * 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
275
+ * `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
276
+
277
+* `update_state(self, batch_update: Optional["BatchUpdate"]) -> None`:
278
+ * Consume a `BatchUpdate` data structure representing persistent batch state changes at the beginning of the current engine step
279
+ * Use the `BatchUpdate` members to update logits processor internal state
280
+ * **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.
281
+
282
+### `BatchUpdate` data structure
283
+
284
+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()`):
285
+
286
+* **Remove:** remove (without replacement) request at index `i`
287
+
288
+ * A Remove is represented in `Batchupdate.removed` by an `int` (representing `i`)
289
+
290
+ * Effect of remove-at-index on batch:
291
+
292
+ ``` text
293
+ Batch: [A,B,C]
294
+ Remove @ i: 1
295
+
296
+ =>
297
+
298
+ New Batch: [A,x,C] # Discard B and leave an empty slot
299
+ ```
300
+
301
+* **Add:** add (or replace existing request with) a new request at index `i`. If a request is replaced, its associated state should be discarded.
302
+
303
+ * An Add is represented in `Batchupdate.added` as a tuple of
304
+
305
+ ``` text
306
+ (index, new request SamplingParams, prompt token ids, output token ids)
307
+ ```
308
+
309
+ * `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**.
310
+
311
+ * 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
312
+
313
+ * If index `i` currently holds a request, a replacement occurs:
314
+
315
+ ``` text
316
+ Batch: [A,B,C]
317
+ New request to be added @ i: D @ 1
318
+
319
+ =>
320
+
321
+ New Batch: [A,D,C] # Add D, discard B
322
+ ```
323
+
324
+ * If index `i` does not currently hold a request (because `i` is out of bounds of the current batch size):
325
+
326
+ ``` text
327
+ Batch: [A,B,C]
328
+ New request to be added @ i: D @ 3
329
+
330
+ =>
331
+
332
+ New Batch: [A,B,C,D] # Add D, extending batch
333
+ ```
334
+
335
+* **Move:** move request at index `s` to index `d` OR swap requests at indices `s` and `d`
336
+
337
+ * A Move is represented in `Batchupdate.moved` as a tuple of
338
+
339
+ ``` text
340
+ (s, d, UNIDIRECTIONAL or SWAP)
341
+ ```
342
+
343
+ * If the Move specifies `UNIDRECTIONAL`:
344
+
345
+ * The request at index `s` is moved to index `d`; index `s` becomes an empty slot
346
+
347
+ ``` text
348
+ Batch: [A,x,C,D]
349
+ Unidirectionally Move s -> d: 3 -> 1
350
+
351
+ =>
352
+
353
+ New Batch: [A,D,C,x] # Move D to 1, leaving empty slot at 3
354
+ ```
355
+
356
+ * If another request already resided at index `d`, it is replaced and discarded
357
+
358
+ ``` text
359
+ Batch: [A,B,C,D]
360
+ Unidirectionally Move s -> d: 3 -> 1
361
+
362
+ =>
363
+
364
+ New Batch: [A,D,C,x] # Move D to 1, discarding B and leaving empty slot at 3
365
+ ```
366
+
367
+ * If the Move specifies `SWAP`, the requests at `s` and `d` exchange indices
368
+
369
+ ``` text
370
+ Batch: [A,B,C,D]
371
+ Swap Move s <-> d: 3 <-> 1
372
+
373
+ =>
374
+
375
+ New Batch: [A,D,C,B] # Swap B and D
376
+ ```
377
+
378
+Additionally, the `BatchUpdate` data structure includes a representation (`batch_size`) of the size of the persistent batch at the beginning of the engine step.
379
+
380
+### How the vLLM engine builds the `BatchUpdate` data structure
381
+
382
+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):
383
+
384
+1. Identify indices of requests which finished in the current engine step
385
+
386
+2. Identify new requests introduced in the current step
387
+
388
+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
389
+
390
+4. Based on the relative number of new and finished requests:
391
+
392
+ 1. If the numbers of new and finished requests are the same, proceed to next step
393
+
394
+ 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`
395
+
396
+ 3. *If there are fewer new requests than finished requests:*
397
+
398
+ * 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
399
+
400
+ * **"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
401
+
402
+ * **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
403
+
404
+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
405
+
406
+Notes:
407
+
408
+* A logits processor `update_state()` method must process batch update operations in the following order: removes, adds, moves
409
+
410
+* The index argument for Add operations refers to the index *at the time the Add occurred*, i.e. before any Move operations
411
+ * 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
412
+ * In other words Move operations can be assumed to be applied after Adds and Removes
413
+
414
+* Move operations can be assumed to be applied in the order in which they appear in `BatchUpdate.moved`
415
+
416
+* If there are no new/finished requests and there is no batch reordering, then the batch update for the logits processors will be `None`
417
+
418
+#### Example: Batch Update with Fewer New Requests Than Finished Requests
419
+
420
+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.
421
+
422
+``` text
423
+Batch state (beginning of engine step): [A,B,C,D]
424
+Batch size: 4
425
+
426
+New requests: E
427
+
428
+Finished requests: A, C
429
+
430
+Processing steps (using BatchUpdate abstraction):
431
+
432
+1. Add E at index 0
433
+
434
+[E,B,C,D] # Discard A
435
+Batch size: 4
436
+
437
+2. Remove at index 2
438
+
439
+[E,B,x,D] # Discard C, empty slot at index 2
440
+Batch size: 4
441
+
442
+3. Condense batch with a Unidirectional Move 3 -> 2 operation and shrink batch
443
+
444
+[E,B,D] x # Empty slot is now outside batch
445
+Batch size: 3
446
+
447
+4. Attention backend optimization: reorder batch with Swap 0 <-> 1
448
+
449
+[B,E,D]
450
+Batch size: 3
451
+
452
+```
453
+
454
+The resulting `BatchUpdate` data structure will look like
455
+
456
+``` text
457
+BatchUpdate instance
458
+* added: [(0,E's SamplingParams,E's prompt tokens ref,E's output tokens ref)]
459
+* removed: [2] # request C was removed without replacement
460
+* moved: [(3,2,UNIDIRECTIONAL),(0,1,SWAP)]
461
+```
462
+
463
+#### Example: Batch Update with More New Requests Than Finished Requests
464
+
465
+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.
466
+
467
+``` text
468
+Batch state (beginning of engine step): [A,B,C,D]
469
+Batch size: 4
470
+
471
+New requests: E,F
472
+
473
+Finished requests: C
474
+
475
+Processing steps (using BatchUpdate abstraction):
476
+
477
+1. Add E at index 2
478
+
479
+[A,B,E,D] # Discard C
480
+Batch size: 4
481
+
482
+2. Add F at index 4 (current max batch index + 1)
483
+
484
+[A,B,E,D,F] # Extend batch by 1
485
+Batch size: 5
486
+
487
+4. Attention backend optimization: reorder batch with Swap 0 <-> 1
488
+
489
+[B,A,E,D,F]
490
+Batch size: 5
491
+
492
+```
493
+
494
+Note that batch condensation is skipped because there are no empty slots left behind by Remove operations.
495
+
496
+The resulting `BatchUpdate` data structure will look like
497
+
498
+``` text
499
+BatchUpdate instance
500
+* 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)]
501
+* removed: [] # no requests were removed without replacement
502
+* moved: [(0,1,SWAP)]
503
+```
504
+
505
+## How to Introduce a New Logits Processor to vLLM
506
+
507
+### Best Practices for Writing Built-In Logits Processors
508
+
509
+* Write efficient `apply()` and `update_state()` implementations in light of the fact that logits processors operate at batch granularity
510
+ * For example, you may be able to use efficient vectorized operations to implement `apply()` or update internal state vectors in `update_state()`
511
+ * 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
512
+
513
+* It is up to the logits processor author to determine:
514
+
515
+ 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
516
+
517
+ 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
518
+
519
+ 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
520
+
521
+ * Additionally, an easy way to save compute in `update_state()` is to exit early when the batch_update is `None`
522
+
523
+* 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)
524
+
525
+* `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
526
+
527
+### Built-In Logits Processors
528
+
529
+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:
530
+
531
+* Min-P
532
+
533
+* Logit bias
534
+
535
+* Min-tokens
536
+
537
+Review these logits processor implementations for guidance on writing built-in logits processors.
538
+
539
+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 aforemented logits processor programming model.
540
+
541
+* Allowed token IDs
542
+
543
+* Bad words
544
+
545
+* Repetition penalty
546
+
547
+* Frequency penalty
548
+
549
+* Presence penalty
550
+
551
+* Temperature
552
+
553
+* Top-K
554
+
555
+* Top-P
556
+
557
+### Custom Logits Processors
558
+
559
+vLLM can be augmented with [user-provided custom logits processors](../features/custom_logitsprocs.md).
docs/features/custom_arguments.md
new
+46
@@ -0,0 +1,46 @@
1
+# Custom Arguments
2
+
3
+You can use vLLM *custom arguments* to pass in arguments which are not part of the vLLM `SamplingParams` and REST API specifications. Adding or removing a vLLM custom argument does not require recompiling vLLM, since the custom arguments are passed in as a dictionary.
4
+
5
+Custom arguments can be useful if, for example, you want to use a [custom logits processor](./custom_logitsprocs.md) without modifying the vLLM source code.
6
+
7
+## Offline Custom Arguments
8
+
9
+Custom arguments passed to `SamplingParams.extra_args` as a `dict` will be visible to any code which has access to `SamplingParams`:
10
+
11
+``` python
12
+SamplingParams(extra_args={"your_custom_arg_name": 67})
13
+```
14
+
15
+This allows arguments which are not already part of `SamplingParams` to be passed into `LLM` as part of a request.
16
+
17
+## Online Custom Arguments
18
+
19
+The vLLM REST API allows custom arguments to be passed to the vLLM server via `vllm_xargs`. The example below integrates custom arguments into a vLLM REST API request:
20
+
21
+``` bash
22
+curl http://localhost:8000/v1/completions \
23
+ -H "Content-Type: application/json" \
24
+ -d '{
25
+ "model": "Qwen/Qwen2.5-1.5B-Instruct",
26
+ ...
27
+ "vllm_xargs": {"your_custom_arg": 67}
28
+ }'
29
+```
30
+
31
+Furthermore, OpenAI SDK users can access `vllm_xargs` via the `extra_body` argument:
32
+
33
+``` python
34
+batch = await client.completions.create(
35
+ model="Qwen/Qwen2.5-1.5B-Instruct",
36
+ ...,
37
+ extra_body={
38
+ "vllm_xargs": {
39
+ "your_custom_arg": 67
40
+ }
41
+ }
42
+)
43
+```
44
+
45
+!!! note
46
+ `vllm_xargs` is assigned to `SamplingParams.extra_args` under the hood, so code which uses `SamplingParams.extra_args` is compatible with both offline and online scenarios.
docs/features/custom_logitsprocs.md
new
+445
@@ -0,0 +1,445 @@
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
+## Logits Processors Background
12
+
13
+A logits processor adjusts the next-token probability distribution, usually with the intention of steering the model towards a desired type of behavior.
14
+
15
+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.
16
+
17
+## Creating a Custom Logits Processor
18
+
19
+Custom logits processors must subclass `vllm.v1.sample.logits_processor.LogitsProcessor` and define (at minimum) the following methods:
20
+
21
+* `__init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool)`
22
+ * `vllm_config`: engine configuration data structure
23
+ * `device`: hardware accelerator device info
24
+ * `is_pin_memory`: flag indicating whether pin memory is available to support logits processor implementation
25
+
26
+* `apply(self, logits: torch.Tensor) -> torch.Tensor`:
27
+ * Consume a `(num_requests) x (vocab_size)` logits tensor (`logits`)
28
+ * Apply logits processor transformation at batch granularity
29
+ * Return a transformed `(num_requests) x (vocab_size)` logits tensor
30
+ * You can modify the input logits processors in-place or out-of-place; in-place is more memory-efficient
31
+
32
+* `is_argmax_invariant(self) -> bool`:
33
+ * 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
34
+ * `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
35
+
36
+* `update_state(self, batch_update: Optional["BatchUpdate"]) -> None`:
37
+ * Consume a `BatchUpdate` data structure representing persistent batch state changes at the beginning of the current engine step
38
+ * Use the `BatchUpdate` members to update logits processor internal state
39
+ * **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.
40
+
41
+### How the vLLM engine builds the `BatchUpdate` data structure
42
+
43
+!!! important
44
+ Some logits processors design changes are still in progress. We expect
45
+ that in the future you will not need to account for batch state changes
46
+ when implementing a logits processor, and the information in this section
47
+ will become irrelevant.
48
+
49
+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):
50
+
51
+1. Identify indices of requests which finished in the current engine step
52
+
53
+2. Identify new requests introduced in the current step
54
+
55
+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
56
+
57
+4. Based on the relative number of new and finished requests:
58
+
59
+ 1. If the numbers of new and finished requests are the same, proceed to next step
60
+
61
+ 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`
62
+
63
+ 3. *If there are fewer new requests than finished requests:*
64
+
65
+ * 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
66
+
67
+ * **"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
68
+
69
+ * **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
70
+
71
+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
72
+
73
+Notes:
74
+
75
+* A logits processor `update_state()` method must process batch update operations in the following order: removes, adds, moves
76
+
77
+* The index argument for Add operations refers to the index *at the time the Add occurred*, i.e. before any Move operations
78
+ * 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
79
+ * In other words Move operations can be assumed to be applied after Adds and Removes
80
+
81
+* Move operations can be assumed to be applied in the order in which they appear in `BatchUpdate.moved`
82
+
83
+* If there are no new/finished requests and there is no batch reordering, then the batch update for the logits processors will be `None`
84
+
85
+### Passing Custom Argument to a Custom Logits Processor
86
+
87
+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`.)
88
+
89
+### Example Custom Logits Processor Implementation
90
+
91
+The contrived example below implements a custom logits processor which consumes a `(num\_requests) \times (vocab\_size)` logits tensor and masks out all tokens except for one (`target_token`) with `float(-inf)`. The logits processor is disabled for any request that does not specify `target_token`. To determine whether the logits processor is enabled and which token to leave unmasked, the logits processor checks `SamplingParams.extra_args` for a `target_token` custom argument associated with each request:
92
+
93
+??? code "Example custom logits processor definition"
94
+
95
+ ``` python
96
+ from typing import Optional
97
+ import torch
98
+ from vllm.config import VllmConfig
99
+ from vllm.sampling_params import SamplingParams
100
+ from vllm.v1.sample.logits_processor import (BatchUpdate,
101
+ LogitsProcessor,
102
+ MoveDirectionality)
103
+
104
+ class DummyLogitsProcessor(LogitsProcessor):
105
+ """Fake logit processor to support unit testing and examples"""
106
+
107
+ def __init__(self, vllm_config: "VllmConfig", device: torch.device,
108
+ is_pin_memory: bool):
109
+ self.req_info: dict[int, int] = {}
110
+
111
+ def is_argmax_invariant(self) -> bool:
112
+ """Never impacts greedy sampling"""
113
+ return False
114
+
115
+ def update_state(self, batch_update: Optional[BatchUpdate]):
116
+ if not batch_update:
117
+ return
118
+
119
+ # Process added requests.
120
+ for index, params, _, _ in batch_update.added:
121
+ assert params is not None
122
+ if params.extra_args and (target_token :=
123
+ params.extra_args.get("target_token")):
124
+ self.req_info[index] = target_token
125
+ else:
126
+ self.req_info.pop(index, None)
127
+
128
+ if self.req_info:
129
+ # Process removed requests.
130
+ for index in batch_update.removed:
131
+ self.req_info.pop(index, None)
132
+
133
+ # Process moved requests, unidirectional move (a->b) and swap
134
+ # (a<->b)
135
+ for adx, bdx, direct in batch_update.moved:
136
+ a_val = self.req_info.pop(adx, None)
137
+ b_val = self.req_info.pop(bdx, None)
138
+ if a_val is not None:
139
+ self.req_info[bdx] = a_val
140
+ if direct == MoveDirectionality.SWAP and b_val is not None:
141
+ self.req_info[adx] = b_val
142
+
143
+ def apply(self, logits: torch.Tensor) -> torch.Tensor:
144
+ if not self.req_info:
145
+ return logits
146
+
147
+ # Save target values before modification
148
+ cols = torch.tensor(
149
+ list(self.req_info.values()), dtype=torch.long, device=logits.device
150
+ )
151
+ rows = torch.tensor(
152
+ list(self.req_info.keys()), dtype=torch.long, device=logits.device
153
+ )
154
+ values_to_keep = logits[rows, cols].clone()
155
+
156
+ # Mask all but target tokens
157
+ logits[rows] = float('-inf')
158
+ logits[rows, cols] = values_to_keep
159
+
160
+ return logits
161
+ ```
162
+
163
+In the rest of this document, we will use `DummyLogitsProcessor` as an example of a custom logits processor.
164
+
165
+The `DummyLogitsProcessor.update_state()` implementation maintains a "sparse" representation of the batched requests in the `self.req_info` dictionary: only those requests which specify a `target_token` value have a key in the dictionary. `update_state()` adjusts the stored request indices and `target_token` values (keys and values respectively in `self.req_info`) in response to Add, Remove and Move operations against the persistent batch.
166
+
167
+### Wrapping an Existing Request-Level Logits Processor
168
+
169
+Although the vLLM engine applies logits processors at batch granularity, some users may want to use vLLM with a "request-level" logits processor implementation - an implementation which operates on individual requests. This will be especially true if your logits processor was developed for vLLM version 0, which required it to be a `Callable` (as described [here](https://docs.vllm.ai/en/v0.10.1.1/api/vllm/logits_process.html)) conforming to the following type annotation:
170
+
171
+``` python
172
+RequestLogitsProcessor = Union[
173
+
174
+ # (output token ids, logits tensor) -> logits tensor
175
+ Callable[[list[int], Tensor], Tensor],
176
+
177
+ # (prompt token ids, output token ids, logits tensor) -> logits tensor
178
+ Callable[[list[int], list[int], Tensor], Tensor],
179
+]
180
+```
181
+
182
+While request-level logits processors are explicitly *not* supported in the vLLM engine, vLLM *does* provide a convenient process to wrap an existing `Callable` request-level logits processor and create a batch-level logits processor that is compatible with vLLM. The `Callable` must conform to the type annotation above; if your request-level logits processor has a different interface, then in order to wrap it, you may need to modify it or implement an additional wrapper layer to comply with the interface specification above.
183
+
184
+You can wrap the request-level logits processor by subclassing `AdapterLogitsProcessor` as shown in the example below (in this example, `DummyPerReqLogitsProcessor` is a stand-in for your request-level logits processor which needs to be wrapped.) Override `AdapterLogitsProcessor.is_argmax_invariant(self)` to accurately reflect whether your request-level logits processor may impact which token has the highest-value logit. Override `AdapterLogitsProcessor.new_req_logits_processor(self,params)` to create a new request-level logits processor instance from a `SamplingParams` instance:
185
+
186
+??? code "Example of Wrapping a Request-Level Logits Processor"
187
+
188
+ ``` python
189
+ ...
190
+
191
+ from vllm.v1.sample.logits_processor import (
192
+ AdapterLogitsProcessor, # Wrapper base-class
193
+ RequestLogitsProcessor, # Request-level logitsproc type annotation
194
+ )
195
+
196
+ ...
197
+
198
+ # Stand-in for your request-level logits processor:
199
+ class DummyPerReqLogitsProcessor:
200
+ """The request-level logits processor masks out all logits except the
201
+ token id identified by `target_token`"""
202
+
203
+ def __init__(self, target_token: int) -> None:
204
+ """Specify `target_token`"""
205
+ self.target_token = target_token
206
+
207
+ def __call__(
208
+ self,
209
+ output_ids: list[int],
210
+ logits: torch.Tensor,
211
+ ) -> torch.Tensor:
212
+ val_to_keep = logits[self.target_token].item()
213
+ logits[:] = float("-inf")
214
+ logits[self.target_token] = val_to_keep
215
+ return logits
216
+
217
+ ...
218
+
219
+ # Example of wrapping the request-level logits processor:
220
+ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
221
+ """Example of wrapping a fake request-level logit processor to create a
222
+ batch-level logits processor"""
223
+
224
+ def is_argmax_invariant(self) -> bool:
225
+ return False
226
+
227
+ def new_req_logits_processor(
228
+ self,
229
+ params: SamplingParams,
230
+ ) -> Optional[RequestLogitsProcessor]:
231
+ """This method returns a new request-level logits processor, customized
232
+ to the `target_token` value associated with a particular request.
233
+
234
+ Returns None if the logits processor should not be applied to the
235
+ particular request. To use the logits processor the request must have
236
+ a "target_token" custom argument with an integer value.
237
+
238
+ Args:
239
+ params: per-request sampling params
240
+
241
+ Returns:
242
+ `Callable` request logits processor, or None
243
+ """
244
+ target_token: Optional[Any] = params.extra_args and params.extra_args.get(
245
+ "target_token"
246
+ )
247
+ if target_token is None:
248
+ return None
249
+ if not isinstance(target_token, int):
250
+ logger.warning(
251
+ "target_token value %s is not int; not applying logits"
252
+ " processor to request.",
253
+ target_token,
254
+ )
255
+ return None
256
+ return DummyPerReqLogitsProcessor(target_token)
257
+ ```
258
+
259
+!!! note
260
+ Your `new_req_logits_processor()` override can return `None` to signal that the wrapped logits processor should not be applied to the request in question.
261
+
262
+Once you have created a custom subclass (like `WrappedPerReqLogitsProcessor`) which wraps your request level logits processor, you can pass the custom subclass to vLLM via any of the methods described in the following section.
263
+
264
+## Ways to Load Your Custom Logits Processor in vLLM
265
+
266
+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 logits processors cannot be loaded on-demand for individual requests.
267
+
268
+This section details different ways of making your logits processor visible to vLLM and triggering vLLM to load your logits processor.
269
+
270
+### Method 1: Pass the Custom Logits Processor Fully-Qualified Class Name (FQCN) to vLLM at Initialization Time
271
+
272
+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
273
+
274
+``` bash
275
+vllm serve ... --logits_processors <logits processor 1> <logits processor 2> ...
276
+```
277
+
278
+The only requirements on the FQCN are
279
+
280
+1. Python's `importlib.import_module()` must be able to resolve the dotted path portion of the FQCN and load it as a module
281
+
282
+2. The class-name portion of the FQCN must be possible to import from the loaded module
283
+
284
+3. The object pointed to by the FQCN must be a subclass of `LogitsProcessor`
285
+
286
+See examples below:
287
+
288
+??? code "Passing custom logits processor FQCN to `LLM` in Python"
289
+
290
+ ``` python
291
+ # Pass in FQCN
292
+ llm = LLM(
293
+ model="facebook/opt-125m",
294
+ logits_processors=["your.module.path:DummyLogitsProcessor"],
295
+ )
296
+ ```
297
+
298
+??? code "Passing custom logits processor FQCN to `AsyncLLM` in Python"
299
+
300
+ ``` python
301
+ # Pass in FQCN
302
+ engine_args = AsyncEngineArgs(model="facebook/opt-125m",
303
+ logits_processors=["your.module.path:DummyLogitsProcessor"])
304
+ async_llm = AsyncLLM.from_engine_args(engine_args)
305
+ ```
306
+
307
+??? code "Passing custom logits processor FQCN to vLLM server via CLI"
308
+
309
+ ```bash
310
+ vllm serve facebook/opt-125m --logits_processors your.module.path:DummyLogitsProcessor
311
+ ```
312
+
313
+### Method 2: Automatically Detect Custom Logits Processors Installed in Your Python Environment As Entry Points
314
+
315
+[`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".
316
+
317
+During initialization, vLLM automatically scans the `vllm.logits_processors` entry point group and loads any installed logits processors which it finds.
318
+
319
+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:
320
+
321
+??? code "Exposing a custom logits processor as a Python entrypoint"
322
+
323
+ ``` toml
324
+ [project.entry-points."vllm.logits_processors"]
325
+ dummy_logits_processor = "your.module.path:DummyLogitsProcessor"
326
+ ```
327
+
328
+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.
329
+
330
+!!! note
331
+ vLLM will *always* load *all* logits processors which are exposed via entrypoints under the `vllm.logits_processors` grouping.
332
+
333
+### Method 3 (Offline-only): Pass a Python Class Object to the vLLM Constructor
334
+
335
+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.
336
+
337
+??? code "Passing custom logits processor class object to `LLM` or `AsyncLLM` in Python"
338
+
339
+ ``` python
340
+ # Import custom logits processor
341
+ from some.module import DummyLogitsProcessor
342
+
343
+ # ...or...
344
+
345
+ # Define custom logits processor locally
346
+ from vllm.v1.sample.logits_processor import LogitsProcessor
347
+
348
+ class DummyLogitsProcessor(LogitsProcessor):
349
+ # See DummyLogitsProcessor implementation above
350
+ ...
351
+
352
+ # Pass class object to LLM constructor
353
+ llm = LLM(
354
+ model="facebook/opt-125m",
355
+ logits_processors=[DummyLogitsProcessor],
356
+ )
357
+
358
+ # Pass class object to AsyncLLM constructor
359
+ engine_args = AsyncEngineArgs(model="facebook/opt-125m",
360
+ logits_processors=[DummyLogitsProcessor])
361
+ async_llm = AsyncLLM.from_engine_args(engine_args)
362
+ ```
363
+
364
+## Invoking a Custom Logits Processor Against a Request
365
+
366
+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.
367
+
368
+The examples below show how a user would pass a custom argument (`target_token`) to `DummyLogitsProcessor` in order to (1) enable the logits processor for that particular request and (2) control the logits processor's behavior.
369
+
370
+??? code "vLLM REST API: configure custom logits processor for a request"
371
+
372
+ ``` bash
373
+ curl http://localhost:8000/v1/completions \
374
+ -H "Content-Type: application/json" \
375
+ -d '{
376
+ "model": "Qwen/Qwen2.5-1.5B-Instruct",
377
+ ...
378
+ "vllm_xargs": {"target_token": 67}
379
+ }'
380
+ ```
381
+
382
+??? code "OpenAI SDK: configure custom logits processor for a request"
383
+
384
+ ``` python
385
+ batch = await client.completions.create(
386
+ model="Qwen/Qwen2.5-1.5B-Instruct",
387
+ ...,
388
+ extra_body={
389
+ "vllm_xargs": {
390
+ "target_token": 67
391
+ }
392
+ }
393
+ )
394
+ ```
395
+
396
+??? code "Offline: configure custom logits processor for an `LLM` request"
397
+
398
+ ``` python
399
+ outputs_logitproc = llm.generate("your prompt",
400
+ SamplingParams(...,
401
+ extra_args={"target_token": 67}))
402
+ ```
403
+
404
+??? code "Offline: configure custom logits processor for an `AsyncLLM` request"
405
+
406
+ ``` python
407
+ async for out in engine.generate(request_id="your request id",
408
+ prompt="your prompt",
409
+ sampling_params=SamplingParams(...,
410
+ extra_args={"target_token": 67})):
411
+
412
+ # Process async request outputs
413
+ ...
414
+ ```
415
+
416
+## Best Practices for Writing Custom Logits Processors
417
+
418
+Once vLLM loads a logits processor during initialization, then vLLM will invoke `update_state()` and `apply()` against that logits processor in every engine step. Both methods operate on all requests which currently reside in the vLLM persistent batch. Thus it is important to implement these methods efficiently.
419
+
420
+* Write efficient `apply()` and `update_state()` implementations in light of the fact that logits processors operate at batch granularity
421
+ * For example, you may be able to use efficient vectorized operations to implement `apply()` or update internal state vectors in `update_state()`
422
+ * 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
423
+ * **Note:** wrapped request-level logits processors do not need to implement `apply()` and `update_state()`; the default `AdapterLogitsProcessor.update_state()` implementation maintains a sparse representation of request state, wherein requests for which `new_req_logits_processor()` returns `None` are not represented in the base-class state dictionary. The default implementation of `AdapterLogitsProcessor.apply()` applies the request-level logits processor to each row of input logits sequentially and assembles the output logits tensor. If the performance of this `AdapterLogitsProcessor` default implementation is insufficient, then avoid wrapping your request-level logits processor and instead re-implement it as a `LogitsProcessor` subclass with optimized `apply()` and `update_state()` implementations that operate at batch granularity
424
+
425
+* It is up to the logits processor author to determine:
426
+
427
+ 1. **The per-request attributes which configure the logits processor's behavior against that request.** Your custom logits processor's `update_state()` override determines how `SamplingParams` fields are mapped into logits processor state
428
+
429
+ * **Note:** for wrapped request-level logits processors, `new_req_logits_processor()` determines how `SamplingParams` fields are used to initialize a request-level logits processor instance.
430
+
431
+ 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`. Try to save compute and memory for requests which disable the logits processor
432
+
433
+ * **Note:** for wrapped per-request logits processors, the default `AdapterLogitsProcessor.update_state()` implementation ensures that the request-level logits processor is disabled when `new_req_logits_processor()` returns `None` for that request
434
+
435
+ 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 custom 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 custom 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
436
+
437
+ * Additionally, an easy way to save compute in `update_state()` is to exit early when the `batch_update` is `None`
438
+
439
+ * **Note:** for wrapped per-request logits processors, the `AdapterLogitsProcessor` base-class implements the above optimizations by default
440
+
441
+* 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)
442
+
443
+ * **Note:** for wrapped per-request logits processors, the `AdapterLogitsProcessor` base-class handles this by default
444
+
445
+* `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
examples/offline_inference/logits_processor/custom.py
+4
-6
@@ -56,7 +56,6 @@ class DummyLogitsProcessor(LogitsProcessor):
56
self.req_info: dict[int, int] = {}
57
58
def is_argmax_invariant(self) -> bool:
59
- """Never impacts greedy sampling"""
59
return False
60
61
def update_state(self, batch_update: Optional[BatchUpdate]):
@@ -75,13 +74,12 @@ class DummyLogitsProcessor(LogitsProcessor):
74
return logits
75
76
# Save target values before modification
78
- rows_list = list(self.req_info.keys())
77
cols = torch.tensor(
80
- [self.req_info[i] for i in rows_list],
81
- dtype=torch.long,
82
- device=logits.device,
78
+ list(self.req_info.values()), dtype=torch.long, device=logits.device
79
+ )
80
+ rows = torch.tensor(
81
+ list(self.req_info.keys()), dtype=torch.long, device=logits.device
82
)
84
- rows = torch.tensor(rows_list, dtype=torch.long, device=logits.device)
83
values_to_keep = logits[rows, cols].clone()
84
85
# Mask all but target tokens
tests/v1/logits_processors/utils.py
+4
-3
@@ -69,11 +69,12 @@ class DummyLogitsProcessor(LogitsProcessor):
69
return logits
70
71
# Save target values before modification
72
- rows_list = list(self.req_info.keys())
73
- cols = torch.tensor([self.req_info[i] for i in rows_list],
72
+ cols = torch.tensor(list(self.req_info.values()),
73
+ dtype=torch.long,
74
+ device=logits.device)
75
+ rows = torch.tensor(list(self.req_info.keys()),
76
dtype=torch.long,
77
device=logits.device)
76
- rows = torch.tensor(rows_list, dtype=torch.long, device=logits.device)
78
values_to_keep = logits[rows, cols].clone()
79
80
# Mask all but target tokens
vllm/v1/sample/logits_processor/interface.py
+3
-3
@@ -21,6 +21,9 @@ class MoveDirectionality(Enum):
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], list[int]]
@@ -29,9 +32,6 @@ AddedRequest = tuple[int, SamplingParams, list[int], list[int]]
32
# one-way moves or two-way swaps of requests in batch
33
MovedRequest = tuple[int, int, MoveDirectionality]
34
32
-# Batch indices of any removed requests.
33
-RemovedRequest = int
34
-
35
36
@dataclass(frozen=True)
37
class BatchUpdate:
vllm/v1/sample/logits_processor/state.py
+4
-4
@@ -36,18 +36,18 @@ class BatchUpdateBuilder:
36
37
_removed: list[RemovedRequest]
38
_is_removed_sorted: bool
39
- moved: list[MovedRequest]
39
added: list[AddedRequest]
40
+ moved: list[MovedRequest]
41
42
def __init__(
43
self,
44
removed: Optional[list[RemovedRequest]] = None,
45
- moved: Optional[list[MovedRequest]] = None,
45
added: Optional[list[AddedRequest]] = None,
46
+ moved: Optional[list[MovedRequest]] = None,
47
) -> None:
48
self._removed = removed or []
49
- self.moved = moved or []
49
self.added = added or []
50
+ self.moved = moved or []
51
self._is_removed_sorted = False
52
53
# Used to track changes in the pooling case
@@ -107,8 +107,8 @@ class BatchUpdateBuilder:
107
"""Returns True if there were any changes to the batch."""
108
self._is_removed_sorted = False
109
self._removed.clear()
110
- self.moved.clear()
110
self.added.clear()
111
+ self.moved.clear()
112
batch_changed = self.batch_changed
113
self.batch_changed = False
114
return batch_changed