[Model Runner V2] Support custom logits processors (#56497)
Signed-off-by: David Cheung <d7cheung@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
David Cheung committed
Sep 19, 2026 at 10:46 UTC
5302d1fe40e6969b854e0ab288d34825653767c6
32 files changed
+1210
-505
docs/features/custom_logitsprocs.md
+72
-271
@@ -8,84 +8,40 @@ A "custom" logits processor is written by a user of vLLM and is loaded into vLLM
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
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.
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
19
-Custom logits processors must subclass `vllm.v1.sample.logits_processor.LogitsProcessor` and define (at minimum) the following methods:
25
+Custom logits processors must subclass `vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor` and define (at minimum) the following methods:
26
21
-* `validate_params(cls, sampling_params: SamplingParams)`:
22
- * Raise `ValueError` if `SamplingParams` has invalid arguments (especially custom arguments) used by logits processor.
23
- * When request is sent to entrypoint, `validate_params()` will validate `SamplingParams` and refuse request with invalid arguments.
24
- * **Note:** it's important to implement `validate_params()` to prevent invalid parameters for custom logits processor. Otherwise requests with invalid parameters can cause unexpected behaviour in custom logits processor.
25
-
26
-* `__init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool)`
27
+* `__init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState)`:
28
* `vllm_config`: engine configuration data structure
28
- * `device`: hardware accelerator device info
29
- * `is_pin_memory`: flag indicating whether pin memory is available to support logits processor implementation
30
-
31
-* `apply(self, logits: torch.Tensor) -> torch.Tensor`:
32
- * Consume a `(num_requests) x (vocab_size)` logits tensor (`logits`)
33
- * Apply logits processor transformation at batch granularity
34
- * Return a transformed `(num_requests) x (vocab_size)` logits tensor
35
- * You can modify the input logits processors in-place or out-of-place; in-place is more memory-efficient
36
-
37
-* `is_argmax_invariant(self) -> bool`:
38
- * 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
39
- * `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
40
-
41
-* `update_state(self, batch_update: Optional["BatchUpdate"]) -> None`:
42
- * Consume a `BatchUpdate` data structure representing persistent batch state changes at the beginning of the current engine step
43
- * Use the `BatchUpdate` members to update logits processor internal state
44
- * **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.
45
-
46
-### How the vLLM engine builds the `BatchUpdate` data structure
47
-
48
-!!! important
49
- Some logits processors design changes are still in progress. We expect
50
- that in the future you will not need to account for batch state changes
51
- when implementing a logits processor, and the information in this section
52
- will become irrelevant.
53
-
54
-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):
55
-
56
-1. Identify indices of requests which finished in the current engine step
57
-
58
-2. Identify new requests introduced in the current step
59
-
60
-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
61
-
62
-4. Based on the relative number of new and finished requests:
63
-
64
- 1. If the numbers of new and finished requests are the same, proceed to next step
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
66
- 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`
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
68
- 3. *If there are fewer new requests than finished requests:*
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
70
- * 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
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
72
- * **"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
73
-
74
- * **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
75
-
76
-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
77
-
78
-Notes:
79
-
80
-* A logits processor `update_state()` method must process batch update operations in the following order: removes, adds, moves
81
-
82
-* The index argument for Add operations refers to the index *at the time the Add occurred*, i.e. before any Move operations
83
- * 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
84
- * In other words Move operations can be assumed to be applied after Adds and Removes
85
-
86
-* Move operations can be assumed to be applied in the order in which they appear in `BatchUpdate.moved`
87
-
88
-* If there are no new/finished requests and there is no batch reordering, then the batch update for the logits processors will be `None`
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
@@ -93,7 +49,7 @@ Unlike built-in logits processors, custom logits processors may require configur
49
50
### Example Custom Logits Processor Implementation
51
96
-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:
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
@@ -101,193 +57,55 @@ The contrived example below implements a custom logits processor which consumes
57
import torch
58
from vllm.config import VllmConfig
59
from vllm.sampling_params import SamplingParams
104
- from vllm.v1.sample.logits_processor import (BatchUpdate,
105
- LogitsProcessor,
106
- MoveDirectionality)
60
+ from vllm.v1.worker.gpu.sample.logits_processor import (
61
+ LogitsProcRequestState,
62
+ LogitsContext,
63
+ LogitsProcessor,
64
+ )
65
108
- class DummyLogitsProcessor(LogitsProcessor):
109
- """Fake logit processor to support unit testing and examples"""
66
111
- @classmethod
112
- def validate_params(cls, params: SamplingParams):
113
- target_token: int | None = params.extra_args and params.extra_args.get(
114
- "target_token"
115
- )
116
- if target_token is not None and not isinstance(target_token, int):
117
- raise ValueError(f"target_token value {target_token} is not int")
118
-
119
- def __init__(self, vllm_config: "VllmConfig", device: torch.device,
120
- is_pin_memory: bool):
121
- self.req_info: dict[int, int] = {}
122
-
123
- def is_argmax_invariant(self) -> bool:
124
- """Never impacts greedy sampling"""
125
- return False
126
-
127
- def update_state(self, batch_update: BatchUpdate | None):
128
- if not batch_update:
129
- return
130
-
131
- # Process added requests.
132
- for index, params, _, _ in batch_update.added:
133
- assert params is not None
134
- self.validate_params(params)
135
- if params.extra_args and (target_token :=
136
- params.extra_args.get("target_token")):
137
- self.req_info[index] = target_token
138
- else:
139
- self.req_info.pop(index, None)
140
-
141
- if self.req_info:
142
- # Process removed requests.
143
- for index in batch_update.removed:
144
- self.req_info.pop(index, None)
145
-
146
- # Process moved requests, unidirectional move (a->b) and swap
147
- # (a<->b)
148
- for adx, bdx, direct in batch_update.moved:
149
- a_val = self.req_info.pop(adx, None)
150
- b_val = self.req_info.pop(bdx, None)
151
- if a_val is not None:
152
- self.req_info[bdx] = a_val
153
- if direct == MoveDirectionality.SWAP and b_val is not None:
154
- self.req_info[adx] = b_val
155
-
156
- def apply(self, logits: torch.Tensor) -> torch.Tensor:
157
- if not self.req_info:
158
- return logits
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
160
- # Save target values before modification
161
- cols = torch.tensor(
162
- list(self.req_info.values()), dtype=torch.long, device=logits.device
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
)
164
- rows = torch.tensor(
165
- list(self.req_info.keys()), dtype=torch.long, device=logits.device
77
+ self.target_token_dev = torch.full(
78
+ (req_states.max_num_reqs,), -1, dtype=torch.int64, device=req_states.device
79
)
167
- values_to_keep = logits[rows, cols].clone()
168
-
169
- # Mask all but target tokens
170
- logits[rows] = float('-inf')
171
- logits[rows, cols] = values_to_keep
172
-
173
- return logits
174
-
175
- ```
80
177
-In the rest of this document, we will use `DummyLogitsProcessor` as an example of a custom logits processor.
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
179
-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.
86
+ def apply_staged_writes(self) -> None:
87
+ self.target_token_dev.copy_(self.target_token, non_blocking=True)
88
181
-### Wrapping an Existing Request-Level Logits Processor
182
-
183
-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][vllm.logits_process]) conforming to the following type annotation:
184
-
185
-``` python
186
-RequestLogitsProcessor = Union[
187
-
188
- # (output token ids, logits tensor) -> logits tensor
189
- Callable[[list[int], Tensor], Tensor],
190
-
191
- # (prompt token ids, output token ids, logits tensor) -> logits tensor
192
- Callable[[list[int], list[int], Tensor], Tensor],
193
-]
194
-```
195
-
196
-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.
197
-
198
-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.):
199
-
200
-* Override `AdapterLogitsProcessor.validate_params(cls,params)` to validate request's sampling parameters.
201
-
202
-* Override `AdapterLogitsProcessor.is_argmax_invariant(self)` to accurately reflect whether your request-level logits processor may impact which token has the highest-value logit.
203
-
204
-* Override `AdapterLogitsProcessor.new_req_logits_processor(self,params)` to create a new request-level logits processor instance from a `SamplingParams` instance:
205
-
206
-??? code "Example of Wrapping a Request-Level Logits Processor"
207
-
208
- ``` python
209
- ...
210
-
211
- from vllm.v1.sample.logits_processor import (
212
- AdapterLogitsProcessor, # Wrapper base-class
213
- RequestLogitsProcessor, # Request-level logitsproc type annotation
214
- )
215
-
216
- ...
217
-
218
- # Stand-in for your request-level logits processor:
219
- class DummyPerReqLogitsProcessor:
220
- """The request-level logits processor masks out all logits except the
221
- token id identified by `target_token`"""
222
-
223
- def __init__(self, target_token: int) -> None:
224
- """Specify `target_token`"""
225
- self.target_token = target_token
226
-
227
- def __call__(
228
- self,
229
- output_ids: list[int],
230
- logits: torch.Tensor,
231
- ) -> torch.Tensor:
232
- val_to_keep = logits[self.target_token].item()
233
- logits[:] = float("-inf")
234
- logits[self.target_token] = val_to_keep
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
236
-
237
- ...
238
-
239
- # Example of wrapping the request-level logits processor:
240
- class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
241
- """Example of wrapping a fake request-level logit processor to create a
242
- batch-level logits processor"""
243
-
244
- @classmethod
245
- def validate_params(cls, params: SamplingParams):
246
- target_token: Any | None = params.extra_args and params.extra_args.get(
247
- "target_token"
248
- )
249
- if target_token is not None and not isinstance(target_token, int):
250
- raise ValueError(
251
- f"target_token value {target_token} is not int"
252
- )
253
-
254
- def is_argmax_invariant(self) -> bool:
255
- return False
256
-
257
- def new_req_logits_processor(
258
- self,
259
- params: SamplingParams,
260
- ) -> Optional[RequestLogitsProcessor]:
261
- """This method returns a new request-level logits processor, customized
262
- to the `target_token` value associated with a particular request.
263
-
264
- Returns None if the logits processor should not be applied to the
265
- particular request. To use the logits processor the request must have
266
- a "target_token" custom argument with an integer value.
267
-
268
- Args:
269
- params: per-request sampling params
270
-
271
- Returns:
272
- `Callable` request logits processor, or None
273
- """
274
- target_token: Any | None = params.extra_args and params.extra_args.get(
275
- "target_token"
276
- )
277
- if target_token is None:
278
- return None
279
- return DummyPerReqLogitsProcessor(target_token)
99
```
100
282
-!!! note
283
- 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.
284
-
285
-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.
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
@@ -314,7 +132,7 @@ See examples below:
132
# Pass in FQCN
133
llm = LLM(
134
model="facebook/opt-125m",
317
- logits_processors=["your.module.path:DummyLogitsProcessor"],
135
+ logits_processors=["your.module.path:TargetTokenLogitsProcessor"],
136
)
137
```
138
@@ -323,14 +141,14 @@ See examples below:
141
``` python
142
# Pass in FQCN
143
engine_args = AsyncEngineArgs(model="facebook/opt-125m",
326
- logits_processors=["your.module.path:DummyLogitsProcessor"])
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
333
- vllm serve facebook/opt-125m --logits_processors your.module.path:DummyLogitsProcessor
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
@@ -345,7 +163,7 @@ Suppose that you have developed a Python package that holds your custom logits p
163
164
``` toml
165
[project.entry-points."vllm.logits_processors"]
348
- dummy_logits_processor = "your.module.path:DummyLogitsProcessor"
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.
@@ -361,26 +179,26 @@ You can pass one or more custom logits processor class objects to the `LLM` and
179
180
``` python
181
# Import custom logits processor
364
- from some.module import DummyLogitsProcessor
182
+ from some.module import TargetTokenLogitsProcessor
183
184
# ...or...
185
186
# Define custom logits processor locally
369
- from vllm.v1.sample.logits_processor import LogitsProcessor
187
+ from vllm.v1.worker.gpu.sample.logits_processor import LogitsProcessor
188
371
- class DummyLogitsProcessor(LogitsProcessor):
372
- # See DummyLogitsProcessor implementation above
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",
378
- logits_processors=[DummyLogitsProcessor],
196
+ logits_processors=[TargetTokenLogitsProcessor],
197
)
198
199
# Pass class object to AsyncLLM constructor
200
engine_args = AsyncEngineArgs(model="facebook/opt-125m",
383
- logits_processors=[DummyLogitsProcessor])
201
+ logits_processors=[TargetTokenLogitsProcessor])
202
async_llm = AsyncLLM.from_engine_args(engine_args)
203
```
204
@@ -388,7 +206,7 @@ You can pass one or more custom logits processor class objects to the `LLM` and
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
391
-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.
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
@@ -419,7 +237,7 @@ The examples below show how a user would pass a custom argument (`target_token`)
237
??? code "Offline: configure custom logits processor for an `LLM` request"
238
239
``` python
422
- outputs_logitproc = llm.generate("your prompt",
240
+ outputs_logitproc = llm.generate("your prompt",
241
SamplingParams(...,
242
extra_args={"target_token": 67}))
243
```
@@ -438,31 +256,14 @@ The examples below show how a user would pass a custom argument (`target_token`)
256
257
## Best Practices for Writing Custom Logits Processors
258
441
-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.
442
-
443
-* Write efficient `apply()` and `update_state()` implementations in light of the fact that logits processors operate at batch granularity
444
- * For example, you may be able to use efficient vectorized operations to implement `apply()` or update internal state vectors in `update_state()`
445
- * 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
446
- * **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
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
450
- 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
451
-
452
- * **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.
453
-
454
- 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
455
-
456
- * **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
457
-
458
- 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
459
-
460
- * Additionally, an easy way to save compute in `update_state()` is to exit early when the `batch_update` is `None`
461
-
462
- * **Note:** for wrapped per-request logits processors, the `AdapterLogitsProcessor` base-class implements the above optimizations by default
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
464
-* 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)
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
466
- * **Note:** for wrapped per-request logits processors, the `AdapterLogitsProcessor` base-class handles this by default
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
468
-* `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
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
tests/test_config.py
+9
@@ -652,6 +652,15 @@ def test_v2_model_runner_supports_extract_hidden_states():
652
assert config._get_v2_model_runner_unsupported_features() == []
653
654
655
+def test_v2_model_runner_supports_custom_logits_processors():
656
+ config = VllmConfig()
657
+ config.model_config = cast(
658
+ ModelConfig, SimpleNamespace(logits_processors=["a.b:C"])
659
+ )
660
+
661
+ assert config._get_v2_model_runner_unsupported_features() == []
662
+
663
+
664
def test_dflash2_draft_forces_v2_model_runner():
665
"""A DFlash2 draft must reach the V2 speculator, the only one that runs its
666
candidate selector; on V1 it would draft as DFlash1 without raising."""
tests/test_sampling_params.py
+8
@@ -51,6 +51,14 @@ def test_non_diffusion_models_unaffected():
51
params.verify(MockModelConfig(), None, None, None)
52
53
54
+def test_verify_leaves_logits_processors_to_admission():
55
+ """verify() is runner-agnostic; LP validation lives in the admission
56
+ layer, so an unimportable FQCN must not fail verify()."""
57
+ SamplingParams().verify(
58
+ MockModelConfig(logits_processors=["no.such:Cls"]), None, None, None
59
+ )
60
+
61
+
62
@pytest.mark.parametrize("value", [-(2**63) - 1, 2**64])
63
def test_extra_args_rejects_nested_integer_overflow(value):
64
"""Reject extension values before they reach the engine transport."""
tests/v1/engine/test_input_processor_trace_replay.py
+1
@@ -85,6 +85,7 @@ def _validate(enable_trace_replay: bool) -> None:
85
speculative_config=None,
86
structured_outputs_config=None,
87
tokenizer=None,
88
+ validate_logits_processors_params=lambda params: None,
89
)
90
params = SamplingParams(trace_decode_token_ids=[1, 2, 3])
91
with patch.object(SamplingParams, "verify"):
tests/v1/logits_processors/test_custom_offline.py
+9
@@ -162,6 +162,9 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource
162
"""
163
# Test that logitproc info is passed to workers
164
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
165
+ # These tests exercise the V1-interface logits processor; Model Runner V2
166
+ # rejects V1-interface processors at load time by design.
167
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
168
set_random_seed(40)
169
170
# Choose LLM args based on logitproc source
@@ -217,6 +220,9 @@ def test_custom_logitsprocs_req(monkeypatch):
220
"""
221
# Test that logitproc info is passed to workers
222
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
223
+ # These tests exercise the V1-interface logits processor; Model Runner V2
224
+ # rejects V1-interface processors at load time by design.
225
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
226
set_random_seed(40)
227
_run_test(
228
{"logits_processors": [WrappedPerReqLogitsProcessor]}, logitproc_loaded=True
@@ -261,6 +267,9 @@ def test_rejects_custom_logitsprocs(
267
268
"""
269
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
270
+ # These tests exercise V1-interface rejection and V1 runner internals;
271
+ # Model Runner V2 rejects V1-interface processors at load time by design.
272
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
273
set_random_seed(40)
274
275
test_params: dict[str, dict[str, Any]] = {
tests/v1/logits_processors/test_custom_online.py
+3
@@ -79,6 +79,9 @@ def server(default_server_args, request, monkeypatch):
79
"""
80
# Test that logitproc info is passed to workers
81
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
82
+ # These tests exercise the V1-interface logits processor; Model Runner V2
83
+ # rejects V1-interface processors at load time by design.
84
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
85
86
if request.param:
87
# Launch server, append FQCN argument, inject dummy logitproc module
tests/v1/worker/test_gpu_bad_words.py
+49
-6
@@ -14,6 +14,10 @@ if not torch.cuda.is_available():
14
15
from vllm.sampling_params import SamplingParams
16
from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
17
+from vllm.v1.worker.gpu.sample.logits_processor import (
18
+ LogitsContext,
19
+ LogitsProcRequestState,
20
+)
21
from vllm.v1.worker.gpu.states import RequestState
22
23
DEVICE = torch.device("cuda")
@@ -47,7 +51,7 @@ def _make_state(bad_words_token_ids: list[list[int]]) -> tuple[BadWordsState, in
51
req_states.apply_staged_writes()
52
53
req_idx = req_states.req_id_to_index["req"]
50
- state = BadWordsState(req_states)
54
+ state = BadWordsState(None, LogitsProcRequestState.from_request_state(req_states))
55
state.add_request(req_idx, SamplingParams(_bad_words_token_ids=bad_words_token_ids))
56
state.apply_staged_writes()
57
return state, req_idx
@@ -61,16 +65,55 @@ def _apply(bad_words_token_ids: list[list[int]]) -> torch.Tensor:
65
expanded_idx_mapping = torch.tensor(
66
[req_idx] * num_logits, dtype=torch.int32, device=DEVICE
67
)
64
- state.apply_bad_words(
68
+ state.apply(
69
logits,
66
- expanded_idx_mapping,
67
- idx_mapping_np,
68
- torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE),
69
- torch.tensor(LOCAL_POS, dtype=torch.int32, device=DEVICE),
70
+ LogitsContext(
71
+ expanded_idx_mapping=expanded_idx_mapping,
72
+ idx_mapping=torch.tensor([req_idx], dtype=torch.int32, device=DEVICE),
73
+ idx_mapping_np=idx_mapping_np,
74
+ expanded_local_pos=torch.tensor(
75
+ LOCAL_POS, dtype=torch.int32, device=DEVICE
76
+ ),
77
+ input_ids=torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE),
78
+ pos=torch.zeros(num_logits, dtype=torch.int32, device=DEVICE),
79
+ seq_lens_upper_bound_np=np.full(1, num_logits, dtype=np.int64),
80
+ ),
81
)
82
return logits.cpu()
83
84
85
+def test_v2_bad_words_slot_reuse_clears_state():
86
+ """A slot recycled from a bad-words request must not keep masking.
87
+
88
+ There is no removal hook by design; safety comes from add_request()
89
+ overwriting the slot's state unconditionally.
90
+ """
91
+ state, req_idx = _make_state([[10, 11, 60]])
92
+ state.add_request(req_idx, SamplingParams())
93
+ state.apply_staged_writes()
94
+ assert int(state.num_bad_words.np[req_idx]) == 0
95
+
96
+ num_logits = len(INPUT_IDS)
97
+ idx_mapping = torch.tensor([req_idx], dtype=torch.int32, device=DEVICE)
98
+ out = state.apply(
99
+ torch.zeros((num_logits, VOCAB_SIZE), device=DEVICE),
100
+ LogitsContext(
101
+ expanded_idx_mapping=torch.tensor(
102
+ [req_idx] * num_logits, dtype=torch.int32, device=DEVICE
103
+ ),
104
+ idx_mapping=idx_mapping,
105
+ idx_mapping_np=np.array([req_idx], dtype=np.intp),
106
+ expanded_local_pos=torch.tensor(
107
+ LOCAL_POS, dtype=torch.int32, device=DEVICE
108
+ ),
109
+ input_ids=torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE),
110
+ pos=torch.zeros(num_logits, dtype=torch.int32, device=DEVICE),
111
+ seq_lens_upper_bound_np=np.full(1, num_logits, dtype=np.int64),
112
+ ),
113
+ )
114
+ assert not torch.isinf(out).any()
115
+
116
+
117
def test_v2_bad_words_prefix_inside_draft_tokens():
118
"""A prefix matching entirely within the draft tokens must mask the bad
119
word's last token at the draft position that completes the prefix."""
tests/v1/worker/test_gpu_batch_shard.py
+6
@@ -183,6 +183,12 @@ def test_local_batch_partition(tp_size: int, seed: int):
183
).all()
184
assert (_np(local.logits_indices) == _np(batch.logits_indices)[rows]).all()
185
assert (_np(local.seq_lens) == _np(batch.seq_lens)[owned]).all()
186
+ # Indexed by batch position, like seq_lens: a consumer holding local row
187
+ # indices reads another request's length if this is not localised too.
188
+ assert (
189
+ _np(local.seq_lens_cpu_upper_bound)
190
+ == _np(batch.seq_lens_cpu_upper_bound)[owned]
191
+ ).all()
192
assert batch.num_draft_tokens_per_req is not None
193
assert local.num_draft_tokens_per_req is not None
194
assert (
tests/v1/worker/test_gpu_logit_bias.py
+29
-6
@@ -14,6 +14,11 @@ if not torch.cuda.is_available():
14
15
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
16
from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
17
+from vllm.v1.worker.gpu.sample.logits_processor import (
18
+ LogitsContext,
19
+ LogitsProcRequestState,
20
+)
21
+from vllm.v1.worker.gpu.states import RequestState
22
23
DEVICE = torch.device("cuda")
24
VOCAB_SIZE = 128
@@ -45,17 +50,35 @@ def _only_stop_token_left(num_rows: int) -> torch.Tensor:
50
51
def _apply(logits: torch.Tensor, structured: list[bool]) -> torch.Tensor:
52
"""Run the fused bias kernel with one logits row per request."""
48
- state = LogitBiasState(max_num_reqs=4, device=DEVICE)
53
+ req_states = RequestState(
54
+ max_num_reqs=4,
55
+ max_model_len=64,
56
+ max_num_batched_tokens=16,
57
+ num_speculative_steps=1,
58
+ vocab_size=VOCAB_SIZE,
59
+ device=DEVICE,
60
+ )
61
+ # add_request() reads prompt_len per slot; every slot here uses the same one.
62
+ req_states.prompt_len.np[:] = PROMPT_LEN
63
+
64
+ state = LogitBiasState(None, LogitsProcRequestState.from_request_state(req_states))
65
for req_idx, is_structured in enumerate(structured):
50
- state.add_request(req_idx, PROMPT_LEN, _params(is_structured))
66
+ state.add_request(req_idx, _params(is_structured))
67
state.apply_staged_writes()
68
69
num_reqs = len(structured)
54
- state.apply_logit_bias(
70
+ idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
71
+ state.apply(
72
logits,
56
- torch.arange(num_reqs, dtype=torch.int32, device=DEVICE),
57
- np.arange(num_reqs, dtype=np.intp),
58
- torch.full((num_reqs,), POS, dtype=torch.int32, device=DEVICE),
73
+ LogitsContext(
74
+ expanded_idx_mapping=idx_mapping,
75
+ idx_mapping=idx_mapping,
76
+ idx_mapping_np=np.arange(num_reqs, dtype=np.intp),
77
+ expanded_local_pos=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
78
+ input_ids=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
79
+ pos=torch.full((num_reqs,), POS, dtype=torch.int32, device=DEVICE),
80
+ seq_lens_upper_bound_np=np.full(num_reqs, POS + 1, dtype=np.int64),
81
+ ),
82
)
83
return logits.cpu()
84
tests/v1/worker/test_gpu_logits_processors.py
new
+114
@@ -0,0 +1,114 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+from typing import Any
6
+
7
+import numpy as np
8
+import pytest
9
+import torch
10
+
11
+pytest.importorskip("triton")
12
+if not torch.cuda.is_available():
13
+ pytest.skip(
14
+ "CUDA required for Model Runner V2 logits processor tests",
15
+ allow_module_level=True,
16
+ )
17
+
18
+from vllm.sampling_params import SamplingParams
19
+from vllm.v1.worker.gpu.sample.logits_processor import (
20
+ LogitsContext,
21
+ LogitsProcessor,
22
+)
23
+from vllm.v1.worker.gpu.sample.sampler import Sampler
24
+from vllm.v1.worker.gpu.states import RequestState
25
+
26
+DEVICE = torch.device("cuda")
27
+VOCAB_SIZE = 128
28
+TARGET_TOKEN = 7
29
+BIAS = 5.0
30
+BOOST = 100.0
31
+TEMPERATURE = 2.0
32
+
33
+
34
+class MockReasoningConfig:
35
+ reasoning_start_token_ids = [90]
36
+ reasoning_end_token_ids = [91]
37
+ natural_reasoning_end_token_ids = [91]
38
+
39
+
40
+class RecordingProcessor(LogitsProcessor):
41
+ """Adds BOOST to TARGET_TOKEN and records every lifecycle callback."""
42
+
43
+ def __init__(self, vllm_config: Any, req_state: Any):
44
+ self.events: list[str] = []
45
+ self.seen_logits: torch.Tensor | None = None
46
+
47
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
48
+ self.events.append(f"add_request:{req_idx}")
49
+ return True
50
+
51
+ def apply_staged_writes(self) -> None:
52
+ self.events.append("apply_staged_writes")
53
+
54
+ def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
55
+ self.events.append("apply")
56
+ self.seen_logits = logits.clone()
57
+ logits[:, TARGET_TOKEN] += BOOST
58
+ return logits
59
+
60
+
61
+def _make_sampler(processor: RecordingProcessor) -> Sampler:
62
+ req_states = RequestState(
63
+ max_num_reqs=4,
64
+ max_model_len=64,
65
+ max_num_batched_tokens=16,
66
+ num_speculative_steps=1,
67
+ vocab_size=VOCAB_SIZE,
68
+ device=DEVICE,
69
+ )
70
+ return Sampler(
71
+ vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()),
72
+ max_num_reqs=4,
73
+ vocab_size=VOCAB_SIZE,
74
+ device=DEVICE,
75
+ req_states=req_states,
76
+ custom_logits_processors=[processor],
77
+ )
78
+
79
+
80
+def test_custom_processor_lifecycle_and_pipeline_position():
81
+ """A custom processor's hooks fire in order, it runs after the built-in
82
+ bias stage, and before temperature scaling."""
83
+ proc = RecordingProcessor(None, None)
84
+ sampler = _make_sampler(proc)
85
+ sampler.add_request(
86
+ 0,
87
+ SamplingParams(logit_bias={TARGET_TOKEN: BIAS}, temperature=TEMPERATURE),
88
+ )
89
+ sampler.apply_staged_writes()
90
+
91
+ num_reqs = 1
92
+ idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
93
+ raw = torch.zeros(num_reqs, VOCAB_SIZE, device=DEVICE)
94
+ out = sampler.apply_sampling_params(
95
+ raw,
96
+ expanded_idx_mapping=idx_mapping,
97
+ idx_mapping=idx_mapping,
98
+ idx_mapping_np=np.arange(num_reqs, dtype=np.intp),
99
+ pos=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
100
+ input_ids=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
101
+ expanded_local_pos=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
102
+ seq_lens_upper_bound_np=np.ones(num_reqs, dtype=np.int64),
103
+ )
104
+
105
+ assert proc.events == [
106
+ "add_request:0",
107
+ "apply_staged_writes",
108
+ "apply",
109
+ ]
110
+ # The processor sees the built-in bias applied but no temperature scaling.
111
+ assert proc.seen_logits is not None
112
+ assert proc.seen_logits[0, TARGET_TOKEN].item() == BIAS
113
+ # Its own modification is then temperature-scaled by the sampler.
114
+ assert out[0, TARGET_TOKEN].item() == (BIAS + BOOST) / TEMPERATURE
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+10
-1
@@ -56,7 +56,16 @@ class FakeEplbState:
56
def _make_runner(**overrides: Any) -> Any:
57
runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner)
58
runner.device = torch.device("cpu")
59
- runner.model_config = SimpleNamespace(model="test-model")
59
+ runner.model_config = SimpleNamespace(model="test-model", logits_processors=None)
60
+ runner.req_states = SimpleNamespace(
61
+ device=runner.device,
62
+ max_num_reqs=8,
63
+ vocab_size=0,
64
+ all_token_ids=None,
65
+ prompt_len=None,
66
+ prefill_len=None,
67
+ total_len=None,
68
+ )
69
runner.load_config = SimpleNamespace(load_format="hf")
70
runner.parallel_config = SimpleNamespace(
71
enable_eplb=True,
tests/v1/worker/test_gpu_rejection_sampler_chunking.py
+3
@@ -47,6 +47,9 @@ def test_chunked_scores_match_full_batch(logprobs_mode: str):
47
[np.arange(count, dtype=np.int32) for count in num_logits_per_req]
48
)
49
).to(device),
50
+ seq_lens_cpu_upper_bound=torch.from_numpy(
51
+ np.array([10, 20, 30, 40], dtype=np.int32)
52
+ ),
53
)
54
rejection_sampler = object.__new__(RejectionSampler)
55
rejection_sampler.sampler = SimpleNamespace(logprobs_mode=logprobs_mode)
tests/v1/worker/test_gpu_sampler_flags.py
+38
-7
@@ -1,6 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections.abc import Sequence
5
+from types import SimpleNamespace
6
+
7
import numpy as np
8
import pytest
9
import torch
@@ -10,6 +13,10 @@ if not torch.cuda.is_available():
13
pytest.skip("CUDA required for sampler flag tests", allow_module_level=True)
14
15
from vllm.sampling_params import SamplingParams
16
+from vllm.v1.worker.gpu.sample.logits_processor import (
17
+ LogitsContext,
18
+ LogitsProcessor,
19
+)
20
from vllm.v1.worker.gpu.sample.sampler import Sampler
21
from vllm.v1.worker.gpu.states import RequestState
22
@@ -23,7 +30,7 @@ class MockReasoningConfig:
30
natural_reasoning_end_token_ids = [91]
31
32
26
-def _make_sampler() -> Sampler:
33
+def _make_sampler(custom_logits_processors: Sequence[LogitsProcessor] = ()) -> Sampler:
34
req_states = RequestState(
35
max_num_reqs=4,
36
max_model_len=64,
@@ -33,11 +40,12 @@ def _make_sampler() -> Sampler:
40
device=DEVICE,
41
)
42
return Sampler(
43
+ vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()),
44
max_num_reqs=4,
45
vocab_size=VOCAB_SIZE,
46
device=DEVICE,
47
req_states=req_states,
40
- reasoning_config=MockReasoningConfig(),
48
+ custom_logits_processors=custom_logits_processors,
49
)
50
51
@@ -65,23 +73,46 @@ def test_logits_processing_cache_matches_request_features(
73
sampling_params: SamplingParams, expected: bool
74
):
75
sampler = _make_sampler()
68
- sampler.add_request(3, prompt_len=1, sampling_params=sampling_params)
76
+ sampler.add_request(3, sampling_params=sampling_params)
77
78
assert sampler.needs_logits_processing[3] == expected
79
80
81
def test_logits_processing_cache_is_overwritten_when_slot_is_reused():
82
sampler = _make_sampler()
75
- sampler.add_request(3, 1, SamplingParams.for_sampler_warmup())
76
- sampler.add_request(3, 1, SamplingParams())
83
+ sampler.add_request(3, SamplingParams.for_sampler_warmup())
84
+ sampler.add_request(3, SamplingParams())
85
86
assert not sampler.needs_logits_processing[3]
87
88
89
+class _GateProcessor(LogitsProcessor):
90
+ """Reports a fixed admission decision from add_request()."""
91
+
92
+ def __init__(self, admitted: bool):
93
+ self.admitted = admitted
94
+
95
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
96
+ return self.admitted
97
+
98
+ def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
99
+ raise AssertionError("the pipeline never runs in this test")
100
+
101
+
102
+@pytest.mark.parametrize("admitted", [True, False])
103
+def test_custom_processor_add_request_gates_pipeline_flag(admitted: bool):
104
+ """A custom processor's add_request() return value is OR-ed into the
105
+ per-request needs_logits_processing flag."""
106
+ sampler = _make_sampler(custom_logits_processors=[_GateProcessor(admitted)])
107
+ sampler.add_request(3, SamplingParams())
108
+
109
+ assert sampler.needs_logits_processing[3] == admitted
110
+
111
+
112
def test_logits_processing_cache_only_checks_active_requests():
113
sampler = _make_sampler()
83
- sampler.add_request(0, 1, SamplingParams(temperature=0.0))
84
- sampler.add_request(2, 1, SamplingParams.for_sampler_warmup())
114
+ sampler.add_request(0, SamplingParams(temperature=0.0))
115
+ sampler.add_request(2, SamplingParams.for_sampler_warmup())
116
117
sampling_only = np.array([0], dtype=np.int32)
118
with_processing = np.array([0, 2], dtype=np.int32)
tests/v1/worker/test_gpu_thinking_budget.py
+17
-7
@@ -1,6 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from types import SimpleNamespace
5
+
6
+import numpy as np
7
import pytest
8
import torch
9
@@ -12,6 +15,7 @@ if not torch.cuda.is_available():
15
)
16
17
from vllm.sampling_params import SamplingParams
18
+from vllm.v1.worker.gpu.sample.logits_processor import LogitsContext
19
from vllm.v1.worker.gpu.sample.sampler import Sampler
20
from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState
21
from vllm.v1.worker.gpu.states import RequestState
@@ -75,11 +79,17 @@ def _apply(
79
idx_mapping_np = idx_mapping.cpu().numpy()
80
state.apply(
81
logits,
78
- expanded_idx_mapping,
79
- idx_mapping,
80
- idx_mapping_np,
81
- torch.tensor(input_ids, dtype=torch.int32, device=DEVICE),
82
- torch.tensor(local_pos, dtype=torch.int32, device=DEVICE),
82
+ LogitsContext(
83
+ expanded_idx_mapping=expanded_idx_mapping,
84
+ idx_mapping=idx_mapping,
85
+ idx_mapping_np=idx_mapping_np,
86
+ expanded_local_pos=torch.tensor(
87
+ local_pos, dtype=torch.int32, device=DEVICE
88
+ ),
89
+ input_ids=torch.tensor(input_ids, dtype=torch.int32, device=DEVICE),
90
+ pos=torch.zeros(len(input_ids), dtype=torch.int32, device=DEVICE),
91
+ seq_lens_upper_bound_np=np.full(1, len(input_ids), dtype=np.int64),
92
+ ),
93
)
94
return logits.cpu()
95
@@ -190,15 +200,14 @@ def test_v2_greedy_sampling_applies_thinking_budget():
200
"""Greedy-only requests must not bypass thinking-budget processing."""
201
req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
202
sampler = Sampler(
203
+ vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()),
204
max_num_reqs=4,
205
vocab_size=VOCAB_SIZE,
206
device=DEVICE,
207
req_states=req_states,
197
- reasoning_config=MockReasoningConfig(),
208
)
209
sampler.add_request(
210
req_idx=3,
201
- prompt_len=1,
211
sampling_params=SamplingParams(
212
temperature=0.0,
213
thinking_token_budget=3,
@@ -219,6 +228,7 @@ def test_v2_greedy_sampling_applies_thinking_budget():
228
torch.tensor([4], dtype=torch.int32, device=DEVICE),
229
input_ids,
230
torch.tensor([0], dtype=torch.int32, device=DEVICE),
231
+ np.full(1, 4, dtype=np.int64),
232
)
233
234
assert out[0, END].item() == pytest.approx(1.0e9)
tests/v1/worker/test_logits_processor_loader.py
new
+228
@@ -0,0 +1,228 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU-only tests for the V2 custom logits processor loader.
4
+
5
+Load-time validation is pure Python (importlib + issubclass), so it runs
6
+without CUDA. Pipeline behavior is covered by test_gpu_logits_processors.py.
7
+"""
8
+
9
+import subprocess
10
+import sys
11
+from enum import Enum, auto
12
+from types import SimpleNamespace
13
+from typing import Any
14
+
15
+import pytest
16
+import torch
17
+
18
+import vllm.v1.worker.gpu.sample.logits_processor.loader as loader
19
+from vllm.exceptions import VLLMValidationError
20
+from vllm.sampling_params import SamplingParams
21
+from vllm.v1.sample.logits_processor import (
22
+ LogitsProcessor as V1LogitsProcessor,
23
+)
24
+from vllm.v1.worker.gpu.sample.logits_processor import (
25
+ LogitsContext,
26
+ LogitsProcessor,
27
+)
28
+
29
+
30
+class DummyV2Processor(LogitsProcessor):
31
+ """Records its constructor args so tests can assert how it was built."""
32
+
33
+ def __init__(self, vllm_config: Any, req_state: Any):
34
+ self.ctor_args = (vllm_config, req_state)
35
+
36
+ def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
37
+ return logits
38
+
39
+
40
+class AltV2Processor(DummyV2Processor):
41
+ """Second V2 processor, used to assert plugin/FQCN load ordering."""
42
+
43
+
44
+class DummyV1Processor(V1LogitsProcessor):
45
+ """A V1-interface processor. Never instantiated: the loader must reject
46
+ it at the issubclass check, before any constructor runs."""
47
+
48
+
49
+def not_a_processor() -> None:
50
+ """FQCN target that resolves but is not a class."""
51
+
52
+
53
+def _fake_entrypoint(cls: type) -> SimpleNamespace:
54
+ return SimpleNamespace(
55
+ name=cls.__name__.lower(),
56
+ value=f"fake:{cls.__name__}",
57
+ load=lambda: cls,
58
+ )
59
+
60
+
61
+@pytest.fixture(autouse=True)
62
+def no_installed_plugins(monkeypatch: pytest.MonkeyPatch):
63
+ """Isolate tests from entry points registered in the environment."""
64
+ monkeypatch.setattr(loader, "entry_points", lambda group: [])
65
+
66
+
67
+class ProcSource(Enum):
68
+ CLASS = auto()
69
+ FQCN = auto()
70
+ ENTRYPOINT = auto()
71
+
72
+
73
+@pytest.mark.parametrize("source", list(ProcSource))
74
+def test_loads_v2_processors(monkeypatch: pytest.MonkeyPatch, source: ProcSource):
75
+ """A V2 subclass loads via class object, FQCN string, or entrypoint."""
76
+ fake_req_states = SimpleNamespace(
77
+ device="cpu",
78
+ max_num_reqs=4,
79
+ vocab_size=128,
80
+ all_token_ids=None,
81
+ prompt_len=None,
82
+ prefill_len=None,
83
+ total_len=None,
84
+ )
85
+ custom: list[Any]
86
+ if source is ProcSource.CLASS:
87
+ custom = [DummyV2Processor]
88
+ elif source is ProcSource.FQCN:
89
+ custom = [f"{__name__}:DummyV2Processor"]
90
+ else:
91
+ monkeypatch.setattr(
92
+ loader, "entry_points", lambda group: [_fake_entrypoint(DummyV2Processor)]
93
+ )
94
+ # Plugins load before user-specified processors.
95
+ custom = [f"{__name__}:AltV2Processor"]
96
+
97
+ procs = loader.build_custom_logits_processors(None, fake_req_states, False, custom)
98
+
99
+ if source is ProcSource.ENTRYPOINT:
100
+ assert [type(p) for p in procs] == [DummyV2Processor, AltV2Processor]
101
+ else:
102
+ assert [type(p) for p in procs] == [DummyV2Processor]
103
+ vllm_config, req_state = procs[0].ctor_args
104
+ assert vllm_config is None
105
+ assert isinstance(req_state, loader.LogitsProcRequestState)
106
+ assert req_state.max_num_reqs == 4
107
+
108
+
109
+@pytest.mark.parametrize(
110
+ ("custom", "plugin_cls", "exc_type", "msg_fragment"),
111
+ [
112
+ pytest.param(
113
+ [DummyV1Processor],
114
+ None,
115
+ ValueError,
116
+ "not a subclass",
117
+ id="v1-class-object",
118
+ ),
119
+ pytest.param(
120
+ [f"{__name__}:DummyV1Processor"],
121
+ None,
122
+ ValueError,
123
+ "not a subclass",
124
+ id="v1-fqcn",
125
+ ),
126
+ pytest.param(
127
+ [],
128
+ DummyV1Processor,
129
+ ValueError,
130
+ "V1-interface plugins are not supported",
131
+ id="v1-entrypoint",
132
+ ),
133
+ pytest.param(
134
+ [f"{__name__}:not_a_processor"],
135
+ None,
136
+ ValueError,
137
+ "must be a type",
138
+ id="non-type-fqcn",
139
+ ),
140
+ pytest.param(
141
+ ["no.such.module:Nope"],
142
+ None,
143
+ RuntimeError,
144
+ "no.such.module:Nope",
145
+ id="missing-module",
146
+ ),
147
+ pytest.param(
148
+ ["no_colon_here"],
149
+ None,
150
+ ValueError,
151
+ "Expected format",
152
+ id="fqcn-missing-colon",
153
+ ),
154
+ pytest.param(
155
+ ["too:many:colons"],
156
+ None,
157
+ ValueError,
158
+ "Expected format",
159
+ id="fqcn-extra-colons",
160
+ ),
161
+ ],
162
+)
163
+def test_rejects_invalid(
164
+ monkeypatch: pytest.MonkeyPatch,
165
+ custom: list,
166
+ plugin_cls: type | None,
167
+ exc_type: type[Exception],
168
+ msg_fragment: str,
169
+):
170
+ """V1 processors and malformed references fail at load time, not at
171
+ sampling time, with a message naming the offending reference."""
172
+ if plugin_cls is not None:
173
+ monkeypatch.setattr(
174
+ loader, "entry_points", lambda group: [_fake_entrypoint(plugin_cls)]
175
+ )
176
+
177
+ with pytest.raises(exc_type, match=msg_fragment):
178
+ loader.build_custom_logits_processors(None, None, False, custom)
179
+
180
+
181
+def test_pooling_model_rejects_custom_logitsprocs():
182
+ """Pooling models reject custom processors instead of ignoring them."""
183
+ with pytest.raises(ValueError, match="Pooling models do not support"):
184
+ loader.build_custom_logits_processors(None, None, True, [DummyV2Processor])
185
+ assert loader.build_custom_logits_processors(None, None, True, []) == []
186
+
187
+
188
+class ValidatingProcessor(DummyV2Processor):
189
+ """Rejects target_token=-1 via the optional validate_params hook."""
190
+
191
+ @classmethod
192
+ def validate_params(cls, sampling_params: SamplingParams) -> None:
193
+ if (sampling_params.extra_args or {}).get("target_token") == -1:
194
+ raise ValueError("target_token must not be -1")
195
+
196
+
197
+def test_validate_params_runs_at_admission():
198
+ """The factory loads classes up front; the returned validator runs
199
+ validate_params per request and wraps rejections as VLLMValidationError."""
200
+ loader._cached_load_v2_logitsprocs.cache_clear()
201
+ validate = loader.build_custom_logits_processors_params_validator(
202
+ [ValidatingProcessor]
203
+ )
204
+ validate(SamplingParams(extra_args={"target_token": 1}))
205
+ with pytest.raises(VLLMValidationError, match="target_token"):
206
+ validate(SamplingParams(extra_args={"target_token": -1}))
207
+
208
+
209
+def test_validator_factory_loads_classes_eagerly():
210
+ """A bad reference fails when the validator is built (at startup), not on
211
+ the first request."""
212
+ loader._cached_load_v2_logitsprocs.cache_clear()
213
+ with pytest.raises(RuntimeError, match="no.such.module:Nope"):
214
+ loader.build_custom_logits_processors_params_validator(["no.such.module:Nope"])
215
+
216
+
217
+def test_loader_import_stays_frontend_safe():
218
+ """The frontend process imports the loader to validate params; importing
219
+ it must not pull in model-runner side modules."""
220
+ code = (
221
+ "import sys\n"
222
+ "import vllm.v1.worker.gpu.sample.logits_processor.loader\n"
223
+ "assert 'vllm.sampling_params' not in sys.modules\n"
224
+ "assert 'vllm.v1.sample.logits_processor' not in sys.modules\n"
225
+ "assert 'vllm.v1.worker.gpu.states' not in sys.modules\n"
226
+ "assert 'vllm.v1.worker.gpu.buffer_utils' not in sys.modules\n"
227
+ )
228
+ subprocess.run([sys.executable, "-c", code], check=True)
vllm/config/vllm.py
-12
@@ -2880,7 +2880,6 @@ class VllmConfig:
2880
def _get_v2_model_runner_unsupported_features(self) -> list[str]:
2881
"""Collect features not yet supported by the V2 model runner."""
2882
unsupported: list[str] = []
2883
- model_config = self.model_config
2883
speculative_config = self.speculative_config
2884
2885
if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
@@ -2929,17 +2928,6 @@ class VllmConfig:
2928
if self.parallel_config.enable_elastic_ep:
2929
unsupported.append("elastic expert parallelism")
2930
2932
- has_logitsproc_plugins = False
2933
- if model_config is not None:
2934
- from importlib.metadata import entry_points
2935
-
2936
- has_logitsproc_plugins = bool(entry_points(group="vllm.logits_processors"))
2937
-
2938
- if model_config is not None and (
2939
- model_config.logits_processors or has_logitsproc_plugins
2940
- ):
2941
- unsupported.append("custom logits processors")
2942
-
2931
if self.cache_config.mamba_cache_mode == "all":
2932
unsupported.append("mamba cache mode 'all'")
2933
vllm/model_executor/models/diffusion_gemma.py
+1
-1
@@ -1014,7 +1014,7 @@ class DiffusionSampler:
1014
# on the subsequent commit step when num_sampled=CANVAS_LEN.
1015
self._pending_logprobs: dict[int, LogprobsTensors] = {}
1016
1017
- def add_request(self, req_idx: int, prompt_len: int, sampling_params: Any) -> None:
1017
+ def add_request(self, req_idx: int, sampling_params: Any) -> None:
1018
if use_penalty(sampling_params):
1019
logger.warning_once(
1020
"DiffusionGemma does not support repetition/frequency/presence "
vllm/sampling_params.py
-8
@@ -823,7 +823,6 @@ class SamplingParams(
823
self._validate_logit_bias(model_config)
824
self._validate_trace_replay(model_config, speculative_config)
825
self._validate_stop_token_ids(model_config)
826
- self._validate_logits_processors(model_config)
826
self._validate_allowed_token_ids(model_config)
827
self._validate_spec_decode(speculative_config)
828
self._validate_diffusion(model_config)
@@ -993,13 +992,6 @@ class SamplingParams(
992
value=invalid_token_ids,
993
)
994
996
- def _validate_logits_processors(self, model_config: ModelConfig) -> None:
997
- from vllm.v1.sample.logits_processor import (
998
- validate_logits_processors_parameters,
999
- )
1000
-
1001
- validate_logits_processors_parameters(model_config.logits_processors, self)
1002
-
995
def _validate_allowed_token_ids(self, model_config: ModelConfig) -> None:
996
allowed_token_ids = self.allowed_token_ids
997
if allowed_token_ids is None:
vllm/v1/engine/input_processor.py
+28
-1
@@ -2,7 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import time
5
-from collections.abc import Mapping
5
+from collections.abc import Callable, Mapping
6
+from functools import partial
7
from typing import Any, Literal
8
9
import vllm.envs as envs
@@ -51,6 +52,11 @@ class InputProcessor:
52
self.speculative_config = vllm_config.speculative_config
53
self.structured_outputs_config = vllm_config.structured_outputs_config
54
self.observability_config = vllm_config.observability_config
55
+ # Load the custom logits processor classes once; the returned callable
56
+ # runs their validate_params hooks per request at admission.
57
+ self.validate_logits_processors_params = (
58
+ self._build_logits_processors_params_validator()
59
+ )
60
61
self.generation_config_fields = model_config.try_get_generation_config()
62
@@ -81,6 +87,25 @@ class InputProcessor:
87
def get_tokenizer(self) -> TokenizerLike:
88
return self.renderer.get_tokenizer()
89
90
+ def _build_logits_processors_params_validator(
91
+ self,
92
+ ) -> Callable[[SamplingParams], None]:
93
+ """Load the custom logits processor classes and return the per-request
94
+ params validator for the active model runner."""
95
+ custom_logitsprocs = self.model_config.logits_processors
96
+ if self.vllm_config.use_v2_model_runner:
97
+ from vllm.v1.worker.gpu.sample.logits_processor import (
98
+ build_custom_logits_processors_params_validator,
99
+ )
100
+
101
+ return build_custom_logits_processors_params_validator(custom_logitsprocs)
102
+
103
+ from vllm.v1.sample.logits_processor import (
104
+ validate_logits_processors_parameters,
105
+ )
106
+
107
+ return partial(validate_logits_processors_parameters, custom_logitsprocs)
108
+
109
def _validate_params(
110
self,
111
params: SamplingParams | PoolingParams,
@@ -101,6 +126,8 @@ class InputProcessor:
126
self.tokenizer,
127
)
128
129
+ self.validate_logits_processors_params(params)
130
+
131
if self.model_config.return_sampling_mask:
132
if params.temperature <= 0:
133
raise ValueError(
vllm/v1/watermarking/gpu_sampler.py
+2
-4
@@ -36,10 +36,8 @@ class GPUWatermarkSampler(Sampler):
36
self.watermarking.np.fill(True)
37
self.watermarking.copy_to_uva()
38
39
- def add_request(
40
- self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
41
- ) -> None:
42
- super().add_request(req_idx, prompt_len, sampling_params)
39
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
40
+ super().add_request(req_idx, sampling_params)
41
self.watermarking.np[req_idx] = sampling_params.watermarking
42
if sampling_params.watermarking and sampling_params.temperature == 0:
43
logger.warning_once(
vllm/v1/worker/gpu/model_runner.py
+8
-4
@@ -143,6 +143,7 @@ from vllm.v1.worker.gpu.sample.batch_shard import (
143
all_to_all_logits,
144
gather_sampler_output,
145
)
146
+from vllm.v1.worker.gpu.sample.logits_processor import build_custom_logits_processors
147
from vllm.v1.worker.gpu.sample.output import SamplerOutput
148
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
149
from vllm.v1.worker.gpu.sample.sampler import Sampler
@@ -447,8 +448,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
448
)
449
450
# Initialize samplers. Model states may override via custom_sampler().
451
+ config_processors = self.model_config.logits_processors or ()
452
+ custom_logits_processors = build_custom_logits_processors(
453
+ self.vllm_config, self.req_states, self.is_pooling_model, config_processors
454
+ )
455
if self.is_last_pp_rank and not self.is_pooling_model:
456
sampler_kwargs: dict[str, Any] = {
457
+ "vllm_config": self.vllm_config,
458
"max_num_reqs": self.max_num_reqs,
459
"vocab_size": self.vocab_size,
460
"device": self.device,
@@ -457,8 +463,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
463
"num_speculative_tokens": self.decode_query_len,
464
"use_fp64_gumbel": self.model_config.use_fp64_gumbel,
465
"enable_trace_replay": self.model_config.enable_trace_replay,
460
- "reasoning_config": self.vllm_config.reasoning_config,
466
"return_sampling_mask": self.model_config.return_sampling_mask,
467
+ "custom_logits_processors": custom_logits_processors,
468
}
469
if self.vllm_config.watermark_config is None:
470
self.sampler = Sampler(**sampler_kwargs)
@@ -1147,9 +1153,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1153
1154
if self.is_last_pp_rank and new_req_data.sampling_params is not None:
1155
assert self.sampler is not None
1150
- self.sampler.add_request(
1151
- req_index, prompt_len, new_req_data.sampling_params
1152
- )
1156
+ self.sampler.add_request(req_index, new_req_data.sampling_params)
1157
assert self.prompt_logprobs_worker is not None
1158
self.prompt_logprobs_worker.add_request(
1159
req_id, req_index, new_req_data.sampling_params
vllm/v1/worker/gpu/sample/bad_words.py
+27
-28
@@ -1,43 +1,47 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import numpy as np
3
+from typing import TYPE_CHECKING
4
+
5
import torch
6
7
import vllm.envs as envs
8
from vllm.sampling_params import SamplingParams
9
from vllm.triton_utils import tl, triton
10
from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
10
-from vllm.v1.worker.gpu.states import RequestState
11
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
12
+ LogitsContext,
13
+ LogitsProcessor,
14
+ LogitsProcRequestState,
15
+)
16
+
17
+if TYPE_CHECKING:
18
+ from vllm.config import VllmConfig
19
20
13
-class BadWordsState:
14
- def __init__(self, req_states: RequestState):
21
+class BadWordsState(LogitsProcessor):
22
+ def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState):
23
self.req_states = req_states
16
- self.max_num_reqs = req_states.max_num_reqs
17
- self.device = req_states.device
24
+ max_num_reqs = req_states.max_num_reqs
25
+ device = req_states.device
26
27
max_total_tokens = envs.VLLM_MAX_BAD_WORDS_TOTAL_TOKENS
28
max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS
29
# flattened bad word tokens: [max_num_reqs, VLLM_MAX_BAD_WORDS_TOTAL_TOKENS]
30
self.bad_word_token_ids = StagedWriteTensor(
23
- (self.max_num_reqs, max_total_tokens),
24
- dtype=torch.int32,
25
- device=self.device,
31
+ (max_num_reqs, max_total_tokens), dtype=torch.int32, device=device
32
)
33
# cumulative offsets of bad words: [max_num_reqs, VLLM_MAX_NUM_BAD_WORDS + 1]
34
self.bad_word_offsets = StagedWriteTensor(
29
- (self.max_num_reqs, max_num_bad_words + 1),
30
- dtype=torch.int32,
31
- device=self.device,
35
+ (max_num_reqs, max_num_bad_words + 1), dtype=torch.int32, device=device
36
)
37
# number of bad words per request
34
- self.num_bad_words = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
38
+ self.num_bad_words = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
39
36
- def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
40
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
41
bad_words_token_ids = sampling_params.bad_words_token_ids
42
if not bad_words_token_ids:
43
self.num_bad_words.np[req_idx] = 0
40
- return
44
+ return False
45
46
num_bad_words = len(bad_words_token_ids)
47
max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS
@@ -65,38 +69,33 @@ class BadWordsState:
69
self.bad_word_token_ids.stage_write(req_idx, 0, flattened_tokens)
70
self.bad_word_offsets.stage_write(req_idx, 0, offsets)
71
self.num_bad_words.np[req_idx] = num_bad_words
72
+ return True
73
74
def apply_staged_writes(self) -> None:
75
self.num_bad_words.copy_to_uva()
76
self.bad_word_token_ids.apply_write()
77
self.bad_word_offsets.apply_write()
78
74
- def apply_bad_words(
75
- self,
76
- logits: torch.Tensor,
77
- expanded_idx_mapping: torch.Tensor,
78
- idx_mapping_np: np.ndarray,
79
- input_ids: torch.Tensor,
80
- expanded_local_pos: torch.Tensor,
81
- ) -> None:
82
- max_num_bad_words = int(self.num_bad_words.np[idx_mapping_np].max())
79
+ def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
80
+ max_num_bad_words = int(self.num_bad_words.np[ctx.idx_mapping_np].max())
81
if max_num_bad_words == 0:
82
# No request uses bad words. Skip the kernel launch.
85
- return
83
+ return logits
84
85
apply_bad_words(
86
logits,
89
- expanded_idx_mapping,
87
+ ctx.expanded_idx_mapping,
88
self.bad_word_token_ids.gpu,
89
self.bad_word_offsets.gpu,
90
self.num_bad_words.gpu,
91
self.req_states.all_token_ids.gpu,
92
self.req_states.prompt_len.gpu,
93
self.req_states.total_len.gpu,
96
- input_ids,
97
- expanded_local_pos,
94
+ ctx.input_ids,
95
+ ctx.expanded_local_pos,
96
max_num_bad_words,
97
)
98
+ return logits
99
100
101
@triton.jit
vllm/v1/worker/gpu/sample/batch_shard.py
+4
@@ -200,6 +200,9 @@ class BatchSharder:
200
num_logits_per_req_np[local_req_indices_np], out=local_cu_num_logits_np[1:]
201
)
202
max_num_logits_per_req = int(num_logits_per_req_np.max()) if num_reqs else 1
203
+ local_seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[
204
+ torch.from_numpy(local_req_indices_np)
205
+ ]
206
207
# Shard the input batch GPU tensors.
208
sorted_logits_indices = torch.empty(
@@ -269,6 +272,7 @@ class BatchSharder:
272
expanded_idx_mapping=local_expanded_idx_mapping,
273
expanded_local_pos=local_expanded_local_pos,
274
seq_lens=local_seq_lens,
275
+ seq_lens_cpu_upper_bound=local_seq_lens_cpu_upper_bound,
276
logits_indices=local_logits_indices,
277
cu_num_logits=local_cu_num_logits,
278
cu_num_logits_np=local_cu_num_logits_np,
vllm/v1/worker/gpu/sample/logit_bias.py
+34
-37
@@ -1,60 +1,61 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import TYPE_CHECKING
4
+
5
import numpy as np
6
import torch
7
8
from vllm.sampling_params import SamplingParams
9
from vllm.triton_utils import tl, triton
10
from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
11
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
12
+ LogitsContext,
13
+ LogitsProcessor,
14
+ LogitsProcRequestState,
15
+)
16
+
17
+if TYPE_CHECKING:
18
+ from vllm.config import VllmConfig
19
20
MAX_NUM_ALLOWED_TOKEN_IDS = 1024
21
MAX_NUM_LOGIT_BIAS_TOKENS = 1024
22
MAX_NUM_STOP_TOKEN_IDS = 128
23
24
15
-class LogitBiasState:
16
- def __init__(self, max_num_reqs: int, device: torch.device):
17
- self.max_num_reqs = max_num_reqs
25
+class LogitBiasState(LogitsProcessor):
26
+ def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState):
27
+ self.prompt_len_np = req_states.prompt_len.np
28
+ max_num_reqs = req_states.max_num_reqs
29
+ device = req_states.device
30
31
# Allowed token IDs.
20
- self.num_allowed_token_ids = UvaBackedTensor(
21
- self.max_num_reqs, dtype=torch.int32
22
- )
32
+ self.num_allowed_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
33
self.allowed_token_ids = StagedWriteTensor(
24
- (self.max_num_reqs, MAX_NUM_ALLOWED_TOKEN_IDS),
25
- dtype=torch.int32,
26
- device=device,
34
+ (max_num_reqs, MAX_NUM_ALLOWED_TOKEN_IDS), dtype=torch.int32, device=device
35
)
36
# Logit bias.
29
- self.num_logit_bias = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
37
+ self.num_logit_bias = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
38
self.logit_bias_token_ids = StagedWriteTensor(
31
- (self.max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
32
- dtype=torch.int32,
33
- device=device,
39
+ (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS), dtype=torch.int32, device=device
40
)
41
self.logit_bias = StagedWriteTensor(
36
- (self.max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
42
+ (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
43
dtype=torch.float32,
44
device=device,
45
)
46
# Min tokens.
41
- self.min_lens = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
42
- self.num_stop_token_ids = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
43
- self.restore_when_all_masked = UvaBackedTensor(
44
- self.max_num_reqs, dtype=torch.int32
45
- )
47
+ self.min_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
48
+ self.num_stop_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
49
+ self.restore_when_all_masked = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
50
self.stop_token_ids = StagedWriteTensor(
47
- (self.max_num_reqs, MAX_NUM_STOP_TOKEN_IDS),
48
- dtype=torch.int32,
49
- device=device,
51
+ (max_num_reqs, MAX_NUM_STOP_TOKEN_IDS), dtype=torch.int32, device=device
52
)
53
54
# Using any of the above.
55
self.use_logit_bias = np.zeros(max_num_reqs, dtype=bool)
56
55
- def add_request(
56
- self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
57
- ) -> None:
57
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
58
+ prompt_len = int(self.prompt_len_np[req_idx])
59
# Using any logit bias.
60
use_logit_bias = False
61
@@ -112,6 +113,7 @@ class LogitBiasState:
113
self.restore_when_all_masked.np[req_idx] = 0
114
115
self.use_logit_bias[req_idx] = use_logit_bias
116
+ return use_logit_bias
117
118
def apply_staged_writes(self) -> None:
119
self.num_allowed_token_ids.copy_to_uva()
@@ -126,25 +128,19 @@ class LogitBiasState:
128
self.restore_when_all_masked.copy_to_uva()
129
self.stop_token_ids.apply_write()
130
129
- def apply_logit_bias(
130
- self,
131
- logits: torch.Tensor,
132
- expanded_idx_mapping: torch.Tensor,
133
- idx_mapping_np: np.ndarray,
134
- pos: torch.Tensor,
135
- ) -> None:
136
- if not np.any(self.use_logit_bias[idx_mapping_np]):
131
+ def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
132
+ if not np.any(self.use_logit_bias[ctx.idx_mapping_np]):
133
# No request uses logit bias. Skip the kernel launch.
138
- return
134
+ return logits
135
136
enable_stop_token_restore = bool(
141
- np.any(self.restore_when_all_masked.np[idx_mapping_np])
137
+ np.any(self.restore_when_all_masked.np[ctx.idx_mapping_np])
138
)
139
140
apply_logit_bias(
141
logits,
146
- expanded_idx_mapping,
147
- pos,
142
+ ctx.expanded_idx_mapping,
143
+ ctx.pos,
144
self.num_allowed_token_ids.gpu,
145
self.allowed_token_ids.gpu,
146
self.num_logit_bias.gpu,
@@ -156,6 +152,7 @@ class LogitBiasState:
152
self.stop_token_ids.gpu,
153
enable_stop_token_restore,
154
)
155
+ return logits
156
157
158
@triton.jit
vllm/v1/worker/gpu/sample/logits_processor/__init__.py
new
+26
@@ -0,0 +1,26 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Custom logits processors for the V2 model runner.
4
+
5
+Kept import-light: the frontend process imports this package to validate
6
+per-request params, so nothing here may pull in model-runner side modules
7
+(torch, triton, worker state) at import time.
8
+"""
9
+
10
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
11
+ LogitsContext,
12
+ LogitsProcessor,
13
+ LogitsProcRequestState,
14
+)
15
+from vllm.v1.worker.gpu.sample.logits_processor.loader import (
16
+ build_custom_logits_processors,
17
+ build_custom_logits_processors_params_validator,
18
+)
19
+
20
+__all__ = [
21
+ "LogitsProcRequestState",
22
+ "LogitsContext",
23
+ "LogitsProcessor",
24
+ "build_custom_logits_processors",
25
+ "build_custom_logits_processors_params_validator",
26
+]
vllm/v1/worker/gpu/sample/logits_processor/interface.py
new
+158
@@ -0,0 +1,158 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""V2 logits processor interface.
4
+
5
+Import-light by design: custom processor classes are also loaded in the
6
+frontend process (to validate per-request params), so this module must not
7
+pull in model-runner side modules (torch, triton, worker state) at import
8
+time. Heavy imports live under TYPE_CHECKING; annotations are deferred.
9
+"""
10
+
11
+from __future__ import annotations
12
+
13
+from abc import ABC, abstractmethod
14
+from dataclasses import dataclass
15
+from typing import TYPE_CHECKING
16
+
17
+if TYPE_CHECKING:
18
+ import numpy as np
19
+ import torch
20
+
21
+ from vllm.config import VllmConfig
22
+ from vllm.sampling_params import SamplingParams
23
+ from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
24
+ from vllm.v1.worker.gpu.states import RequestState
25
+
26
+
27
+@dataclass(frozen=True)
28
+class LogitsProcRequestState:
29
+ """State associated with active requests, shared with logits processors.
30
+
31
+ Wraps the model runner's per-slot buffers, which are mutated in place,
32
+ so reads always see current values. Processors must treat every field
33
+ as read-only.
34
+ """
35
+
36
+ device: torch.device
37
+ max_num_reqs: int
38
+ vocab_size: int
39
+
40
+ # [max_num_reqs, max_model_len] committed token ids per request slot.
41
+ all_token_ids: StagedWriteTensor
42
+ # [max_num_reqs] tokens in the user-provided prompt.
43
+ prompt_len: UvaBackedTensor
44
+ # [max_num_reqs] tokens fed at the latest (re)fill: the prompt plus any
45
+ # partial output on resumption after preemption.
46
+ prefill_len: UvaBackedTensor
47
+ # [max_num_reqs] prompt_len + output_len; grows as the request progresses.
48
+ total_len: StagedWriteTensor
49
+
50
+ @classmethod
51
+ def from_request_state(cls, req_states: RequestState) -> LogitsProcRequestState:
52
+ return cls(
53
+ device=req_states.device,
54
+ max_num_reqs=req_states.max_num_reqs,
55
+ vocab_size=req_states.vocab_size,
56
+ all_token_ids=req_states.all_token_ids,
57
+ prompt_len=req_states.prompt_len,
58
+ prefill_len=req_states.prefill_len,
59
+ total_len=req_states.total_len,
60
+ )
61
+
62
+
63
+@dataclass(frozen=True)
64
+class LogitsContext:
65
+ """The current step's batch layout, passed to every ``apply()`` call.
66
+
67
+ A row is a logits row, not a request: rows are reordered every step, and
68
+ under speculative decoding a request owns one row per draft token.
69
+ Committed tokens live in ``req_states.all_token_ids`` (valid up to
70
+ ``total_len``); this step's draft tokens are only in ``input_ids``.
71
+ """
72
+
73
+ # [num_logits_rows] row -> persistent request slot.
74
+ expanded_idx_mapping: torch.Tensor
75
+ # [num_reqs] batch position -> persistent request slot.
76
+ idx_mapping: torch.Tensor
77
+ # [num_reqs] batch position -> persistent request slot, on the host, for
78
+ # skipping work without a device sync.
79
+ idx_mapping_np: np.ndarray
80
+ # [num_logits_rows] row -> its offset among the rows of its own request.
81
+ expanded_local_pos: torch.Tensor
82
+ # [num_logits_rows] token fed to the model at each row's input position.
83
+ input_ids: torch.Tensor
84
+ # [num_logits_rows] position of each row within its sequence.
85
+ pos: torch.Tensor
86
+ # [num_reqs] batch position -> upper bound of tokens visible to the token
87
+ # being sampled this step, on the host. Exact when spec decoding isn't in use.
88
+ # Exact per-row lengths on device are `pos + 1`.
89
+ seq_lens_upper_bound_np: np.ndarray
90
+
91
+
92
+class LogitsProcessor(ABC):
93
+ """Custom logits processor for Model Runner V2.
94
+
95
+ Per-request state is keyed by the request slot index; slots are recycled
96
+ through a free list, so per-slot state must be fully (re)initialized in
97
+ ``add_request()``.
98
+
99
+ ``apply()`` runs after the built-in bias, penalty, bad-words and grammar
100
+ stages and before temperature, min_p and top-k/top-p, so it sees unscaled
101
+ logits and must not re-inflate grammar-masked tokens. Thinking-budget
102
+ forcing runs after ``apply()`` and wins over its edits for requests whose
103
+ budget is exhausted.
104
+
105
+ State that is constant for a request belongs in ``__init__()`` or
106
+ ``add_request()``; ``apply()`` receives only what changes per step.
107
+ """
108
+
109
+ def __init__( # noqa: B027
110
+ self, vllm_config: VllmConfig, req_states: LogitsProcRequestState
111
+ ):
112
+ """Capture what stays constant for the processor's lifetime.
113
+
114
+ ``req_states`` exposes the on-device token history and batch constants a
115
+ processor may read. Treat it as read-only.
116
+ """
117
+
118
+ @classmethod # noqa: B027
119
+ def validate_params(cls, sampling_params: SamplingParams) -> None:
120
+ """Raise ``ValueError`` for invalid per-request arguments.
121
+
122
+ Runs at request admission, so invalid arguments fail the request
123
+ with an error instead of reaching the sampler.
124
+ """
125
+
126
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
127
+ """Initialize per-slot state for a request entering the batch.
128
+
129
+ The slot may hold a previous occupant's state; overwrite or
130
+ neutralize all of it here.
131
+
132
+ Returns whether this processor modifies logits for the request.
133
+ """
134
+ return True
135
+
136
+ def apply_staged_writes(self) -> None: # noqa: B027
137
+ """Flush any host-side writes staged by ``add_request()`` to the device.
138
+
139
+ Called once per step before the forward pass, after the model runner
140
+ has flushed ``req_states``, so a processor that stages writes here can
141
+ read the request's tokens back on device.
142
+ """
143
+
144
+ @abstractmethod
145
+ def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
146
+ """Modify logits in place or return a new tensor. In-place modification
147
+ is preferred for efficiency.
148
+
149
+ ``apply()`` is called once for the whole batch, including rows of
150
+ requests this processor declined in ``add_request()``, so filter rows
151
+ via ``ctx.expanded_idx_mapping``.
152
+
153
+ Args:
154
+ logits: [num_logits_rows, vocab_size] float32 tensor.
155
+ ctx: this step's batch layout.
156
+
157
+ """
158
+ raise NotImplementedError
vllm/v1/worker/gpu/sample/logits_processor/loader.py
new
+211
@@ -0,0 +1,211 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Loading of custom logits processor classes for the V2 model runner.
4
+
5
+This module is deliberately import-light: the frontend process uses it to
6
+validate per-request params, so importing it must not pull in model-runner
7
+side modules (torch, triton, worker state). Imports needed to instantiate
8
+processors are deferred to ``build_custom_logits_processors``.
9
+"""
10
+
11
+import importlib
12
+from collections.abc import Callable, Sequence
13
+from functools import lru_cache
14
+from importlib.metadata import entry_points
15
+from typing import TYPE_CHECKING
16
+
17
+from vllm.exceptions import VLLMValidationError
18
+from vllm.logger import init_logger
19
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
20
+ LogitsProcessor,
21
+ LogitsProcRequestState,
22
+)
23
+
24
+if TYPE_CHECKING:
25
+ from vllm.config import VllmConfig
26
+ from vllm.sampling_params import SamplingParams
27
+ from vllm.v1.worker.gpu.states import RequestState
28
+
29
+logger = init_logger(__name__)
30
+
31
+# Same entry-points group as
32
+# vllm.v1.sample.logits_processor.LOGITSPROCS_GROUP; duplicated here so this
33
+# module stays importable without the V1 sampler machinery.
34
+LOGITSPROCS_GROUP = "vllm.logits_processors"
35
+
36
+
37
+def _load_v2_logitsprocs_plugins() -> list[type[LogitsProcessor]]:
38
+ """Load installed logits processor plugins.
39
+
40
+ Plugins registered under ``LOGITSPROCS_GROUP`` ship V1-interface
41
+ processors today; those crash under the V2 sampler, so reject them at
42
+ load time.
43
+ """
44
+ from vllm.utils.torch_utils import guard_cuda_initialization
45
+
46
+ installed_logitsprocs_plugins = entry_points(group=LOGITSPROCS_GROUP)
47
+ if len(installed_logitsprocs_plugins) == 0:
48
+ logger.debug("No logitsprocs plugins installed (group %s).", LOGITSPROCS_GROUP)
49
+ return []
50
+
51
+ logger.debug("Loading installed logitsprocs plugins (group %s):", LOGITSPROCS_GROUP)
52
+ classes: list[type[LogitsProcessor]] = []
53
+ for entrypoint in installed_logitsprocs_plugins:
54
+ try:
55
+ logger.debug(
56
+ "- Loading logitproc plugin entrypoint=%s target=%s",
57
+ entrypoint.name,
58
+ entrypoint.value,
59
+ )
60
+ with guard_cuda_initialization():
61
+ cls = entrypoint.load()
62
+ except Exception as e:
63
+ logger.error("Failed to load LogitsProcessor plugin %s: %s", entrypoint, e)
64
+ raise RuntimeError(
65
+ f"Failed to load LogitsProcessor plugin {entrypoint}"
66
+ ) from e
67
+ if not (isinstance(cls, type) and issubclass(cls, LogitsProcessor)):
68
+ raise ValueError(
69
+ f"LogitsProcessor plugin {entrypoint.name} is not a subclass of "
70
+ "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor; "
71
+ "V1-interface plugins are not supported by the V2 model runner."
72
+ )
73
+ classes.append(cls)
74
+ return classes
75
+
76
+
77
+def _load_v2_logitsprocs_by_fqcns(
78
+ logits_processors: Sequence[str | type],
79
+) -> list[type[LogitsProcessor]]:
80
+ """Resolve a mixed list of processor types and FQCN strings into types.
81
+
82
+ FQCN syntax is <module>:<type> i.e. x.y.z:CustomLogitProc.
83
+ """
84
+ from vllm.utils.torch_utils import guard_cuda_initialization
85
+
86
+ classes: list[type[LogitsProcessor]] = []
87
+ for ldx, logitproc in enumerate(logits_processors):
88
+ if isinstance(logitproc, type):
89
+ logger.debug(" - Already-loaded logit processor: %s", logitproc.__name__)
90
+ if not issubclass(logitproc, LogitsProcessor):
91
+ raise ValueError(
92
+ f"{logitproc.__name__} is not a subclass of "
93
+ "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor"
94
+ )
95
+ classes.append(logitproc)
96
+ continue
97
+
98
+ logger.debug("- Loading logits processor %s", logitproc)
99
+ parts = logitproc.split(":")
100
+ if len(parts) != 2:
101
+ raise ValueError(
102
+ f"Invalid logits processor FQCN {logitproc!r}. "
103
+ "Expected format: '<module>:<type>'"
104
+ )
105
+ module_path, qualname = parts
106
+
107
+ try:
108
+ with guard_cuda_initialization():
109
+ module = importlib.import_module(module_path)
110
+ except Exception as e:
111
+ logger.error(
112
+ "Failed to load %sth LogitsProcessor plugin %s: %s",
113
+ ldx,
114
+ logitproc,
115
+ e,
116
+ )
117
+ raise RuntimeError(
118
+ f"Failed to load {ldx}th LogitsProcessor plugin {logitproc}"
119
+ ) from e
120
+
121
+ obj = module
122
+ for attr in qualname.split("."):
123
+ obj = getattr(obj, attr)
124
+ if not isinstance(obj, type):
125
+ raise ValueError("Loaded logit processor must be a type.")
126
+ if not issubclass(obj, LogitsProcessor):
127
+ raise ValueError(
128
+ f"{obj.__name__} is not a subclass of "
129
+ "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor"
130
+ )
131
+ classes.append(obj)
132
+
133
+ return classes
134
+
135
+
136
+def _load_v2_logitsprocs(
137
+ logits_processors: Sequence[str | type] | None,
138
+) -> list[type[LogitsProcessor]]:
139
+ """Load all custom logits processors for the V2 model runner.
140
+
141
+ Combines installed plugins (validated as V2 processors) with the
142
+ user-specified list of processor types and FQCN strings.
143
+ """
144
+ return _load_v2_logitsprocs_plugins() + _load_v2_logitsprocs_by_fqcns(
145
+ logits_processors or ()
146
+ )
147
+
148
+
149
+@lru_cache
150
+def _cached_load_v2_logitsprocs(
151
+ custom_logitsprocs: tuple[str | type, ...],
152
+) -> list[type[LogitsProcessor]]:
153
+ return _load_v2_logitsprocs(list(custom_logitsprocs))
154
+
155
+
156
+def build_custom_logits_processors(
157
+ vllm_config: "VllmConfig",
158
+ req_states: "RequestState",
159
+ is_pooling_model: bool,
160
+ custom_logitsprocs: Sequence[str | type] = (),
161
+) -> list[LogitsProcessor]:
162
+ """Load and instantiate custom logits processors, entrypoint plugins first.
163
+
164
+ Raises:
165
+ ValueError: if a pooling model specifies custom processors, or a
166
+ loaded class does not implement the V2 interface.
167
+ RuntimeError: if an FQCN fails to import.
168
+
169
+ """
170
+ from vllm.v1.sample.logits_processor import STR_POOLING_REJECTS_LOGITSPROCS
171
+
172
+ if is_pooling_model:
173
+ if custom_logitsprocs:
174
+ raise ValueError(STR_POOLING_REJECTS_LOGITSPROCS)
175
+ logger.debug(
176
+ "Skipping logits processor loading because pooling models"
177
+ " do not support logits processors."
178
+ )
179
+ return []
180
+
181
+ custom_logitsprocs_classes = _load_v2_logitsprocs(custom_logitsprocs)
182
+ lp_req_state = LogitsProcRequestState.from_request_state(req_states)
183
+ return [ctor(vllm_config, lp_req_state) for ctor in custom_logitsprocs_classes]
184
+
185
+
186
+def build_custom_logits_processors_params_validator(
187
+ custom_logitsprocs: Sequence[str | type] | None,
188
+) -> "Callable[[SamplingParams], None]":
189
+ """Load custom processor classes once and return a params validator.
190
+
191
+ Called from the frontend at startup. The returned callable runs each
192
+ processor's ``validate_params`` at request admission.
193
+
194
+ Raises:
195
+ ValueError: if a loaded class does not implement the V2 interface.
196
+ RuntimeError: if an FQCN fails to import.
197
+
198
+ """
199
+ if not custom_logitsprocs:
200
+ return lambda _: None
201
+
202
+ classes = _cached_load_v2_logitsprocs(tuple(custom_logitsprocs))
203
+
204
+ def validate_params(sampling_params: "SamplingParams") -> None:
205
+ for cls in classes:
206
+ try:
207
+ cls.validate_params(sampling_params)
208
+ except ValueError as e:
209
+ raise VLLMValidationError(str(e)) from e
210
+
211
+ return validate_params
vllm/v1/worker/gpu/sample/penalties.py
+26
-28
@@ -1,5 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import TYPE_CHECKING
4
+
5
import numpy as np
6
import torch
7
@@ -8,16 +10,22 @@ from vllm.triton_utils import tl, triton
10
from vllm.utils.math_utils import cdiv
11
from vllm.utils.torch_utils import async_tensor_h2d
12
from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
11
-from vllm.v1.worker.gpu.states import RequestState
13
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
14
+ LogitsContext,
15
+ LogitsProcessor,
16
+ LogitsProcRequestState,
17
+)
18
19
+if TYPE_CHECKING:
20
+ from vllm.config import VllmConfig
21
14
-class PenaltiesState:
15
- def __init__(self, req_states: RequestState):
16
- self.req_states = req_states
22
18
- max_num_reqs = req_states.max_num_reqs
19
- self.vocab_size = req_states.vocab_size
23
+class PenaltiesState(LogitsProcessor):
24
+ def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState):
25
+ self.req_states = req_states
26
self.device = req_states.device
27
+ max_num_reqs = req_states.max_num_reqs
28
+ vocab_size = req_states.vocab_size
29
30
self.repetition_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
31
self.frequency_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
@@ -30,20 +38,17 @@ class PenaltiesState:
38
39
# Statistics for penalties.
40
self.prompt_bin_mask = torch.zeros(
33
- max_num_reqs,
34
- cdiv(self.vocab_size, 32),
35
- dtype=torch.int32,
36
- device=self.device,
41
+ max_num_reqs, cdiv(vocab_size, 32), dtype=torch.int32, device=self.device
42
)
43
# TODO(woosuk): This tensor is rarely used but can be very large, taking up
44
# GBs of GPU memory. Optimize the memory usage.
45
self.output_bin_counts = torch.zeros(
41
- max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
46
+ max_num_reqs, vocab_size, dtype=torch.int32, device=self.device
47
)
48
49
self._new_penalties_reqs: list[int] = []
50
46
- def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
51
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
52
self.repetition_penalty.np[req_idx] = sampling_params.repetition_penalty
53
self.frequency_penalty.np[req_idx] = sampling_params.frequency_penalty
54
self.presence_penalty.np[req_idx] = sampling_params.presence_penalty
@@ -52,13 +57,12 @@ class PenaltiesState:
57
self.use_penalty[req_idx] = do_penalty
58
if do_penalty:
59
self._new_penalties_reqs.append(req_idx)
60
+ return do_penalty
61
62
def apply_staged_writes(self) -> None:
63
if self._new_penalties_reqs:
64
idx_mapping = async_tensor_h2d(
59
- self._new_penalties_reqs,
60
- dtype=torch.int32,
61
- device=self.device,
65
+ self._new_penalties_reqs, dtype=torch.int32, device=self.device
66
)
67
68
prefill_lens = self.req_states.prefill_len.np[self._new_penalties_reqs]
@@ -78,29 +82,23 @@ class PenaltiesState:
82
self.frequency_penalty.copy_to_uva()
83
self.presence_penalty.copy_to_uva()
84
81
- def apply_penalties(
82
- self,
83
- logits: torch.Tensor,
84
- expanded_idx_mapping: torch.Tensor,
85
- idx_mapping_np: np.ndarray,
86
- input_ids: torch.Tensor,
87
- expanded_local_pos: torch.Tensor,
88
- ) -> None:
89
- if not np.any(self.use_penalty[idx_mapping_np]):
85
+ def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
86
+ if not np.any(self.use_penalty[ctx.idx_mapping_np]):
87
# No request uses penalties. Skip the kernel launch.
91
- return
88
+ return logits
89
90
apply_penalties(
91
logits,
95
- expanded_idx_mapping,
96
- input_ids,
97
- expanded_local_pos,
92
+ ctx.expanded_idx_mapping,
93
+ ctx.input_ids,
94
+ ctx.expanded_local_pos,
95
self.repetition_penalty.gpu,
96
self.frequency_penalty.gpu,
97
self.presence_penalty.gpu,
98
self.prompt_bin_mask,
99
self.output_bin_counts,
100
)
101
+ return logits
102
103
104
@triton.jit
vllm/v1/worker/gpu/sample/sampler.py
+59
-64
@@ -1,12 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections.abc import Sequence
5
+
6
import numpy as np
7
import torch
8
9
import vllm.envs as envs
10
+from vllm.config import VllmConfig
11
from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
9
-from vllm.config.reasoning import ReasoningConfig
12
from vllm.sampling_params import SamplingParams
13
from vllm.v1.sample.ops.topk_topp_sampler import (
14
apply_top_k_top_p,
@@ -18,6 +20,11 @@ from vllm.v1.worker.gpu.metrics.logits import get_num_nans
20
from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
21
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
22
from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
23
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
24
+ LogitsContext,
25
+ LogitsProcessor,
26
+ LogitsProcRequestState,
27
+)
28
from vllm.v1.worker.gpu.sample.logprob import (
29
LogprobTokenIdsState,
30
compute_topk_scores,
@@ -33,6 +40,7 @@ from vllm.v1.worker.gpu.states import RequestState
40
class Sampler:
41
def __init__(
42
self,
43
+ vllm_config: VllmConfig,
44
max_num_reqs: int,
45
vocab_size: int,
46
device: torch.device,
@@ -41,8 +49,8 @@ class Sampler:
49
num_speculative_tokens: int = 1,
50
use_fp64_gumbel: bool = False,
51
enable_trace_replay: bool = False,
44
- reasoning_config: ReasoningConfig | None = None,
52
return_sampling_mask: bool = False,
53
+ custom_logits_processors: Sequence[LogitsProcessor] = (),
54
):
55
self.logprobs_mode = logprobs_mode
56
self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.
@@ -50,11 +58,25 @@ class Sampler:
58
59
self.req_states = req_states
60
self.sampling_states = SamplingStates(max_num_reqs, vocab_size)
53
- self.penalties_state = PenaltiesState(req_states)
54
- self.logit_bias_state = LogitBiasState(max_num_reqs, device)
55
- self.bad_words_state = BadWordsState(req_states)
61
+
62
+ lp_req_state = LogitsProcRequestState.from_request_state(req_states)
63
+ self.penalties_state = PenaltiesState(vllm_config, lp_req_state)
64
+ logit_bias_state = LogitBiasState(vllm_config, lp_req_state)
65
+ bad_words_state = BadWordsState(vllm_config, lp_req_state)
66
+
67
+ # List order is pipeline order: bias adds, penalties scale, so the
68
+ # two do not commute.
69
+ self.logits_processors: list[LogitsProcessor] = [
70
+ logit_bias_state,
71
+ self.penalties_state,
72
+ bad_words_state,
73
+ *custom_logits_processors,
74
+ ]
75
+
76
self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
57
- self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config)
77
+ self.thinking_budget_state = ThinkingBudgetState(
78
+ req_states, vllm_config.reasoning_config
79
+ )
80
self.trace_replay_state = (
81
TraceReplayState(req_states) if enable_trace_replay else None
82
)
@@ -65,41 +87,25 @@ class Sampler:
87
not return_sampling_mask and flashinfer_sampler_supported()
88
)
89
68
- def add_request(
69
- self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
70
- ) -> None:
71
- self.sampling_states.add_request(req_idx, sampling_params)
72
- self.penalties_state.add_request(req_idx, sampling_params)
73
- self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params)
74
- self.bad_words_state.add_request(req_idx, sampling_params)
90
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
91
+ needs_processing = self.sampling_states.add_request(req_idx, sampling_params)
92
+ needs_processing |= self.thinking_budget_state.add_request(
93
+ req_idx, sampling_params
94
+ )
95
+ for processor in self.logits_processors:
96
+ needs_processing |= processor.add_request(req_idx, sampling_params)
97
+ self.needs_logits_processing[req_idx] = needs_processing
98
+
99
self.logprob_token_ids_state.add_request(req_idx, sampling_params)
76
- self.thinking_budget_state.add_request(req_idx, sampling_params)
100
if self.trace_replay_state is not None:
101
self.trace_replay_state.add_request(req_idx, sampling_params)
102
80
- states = self.sampling_states
81
- temperature = states.temperature.np[req_idx]
82
- self.needs_logits_processing[req_idx] = (
83
- self.logit_bias_state.use_logit_bias[req_idx]
84
- or self.penalties_state.use_penalty[req_idx]
85
- or self.bad_words_state.num_bad_words.np[req_idx] > 0
86
- or (
87
- self.thinking_budget_state.enabled
88
- and self.thinking_budget_state.use_thinking_budget[req_idx]
89
- )
90
- or (temperature != 0.0 and temperature != 1.0)
91
- or states.min_p.np[req_idx] != 0.0
92
- or states.top_k.np[req_idx] != states.vocab_size
93
- or states.top_p.np[req_idx] != 1.0
94
- )
95
-
103
def apply_staged_writes(self) -> None:
104
self.sampling_states.apply_staged_writes()
98
- self.penalties_state.apply_staged_writes()
99
- self.logit_bias_state.apply_staged_writes()
100
- self.bad_words_state.apply_staged_writes()
101
- self.logprob_token_ids_state.apply_staged_writes()
105
+ for processor in self.logits_processors:
106
+ processor.apply_staged_writes()
107
self.thinking_budget_state.apply_staged_writes()
108
+ self.logprob_token_ids_state.apply_staged_writes()
109
if self.trace_replay_state is not None:
110
self.trace_replay_state.apply_staged_writes()
111
@@ -131,6 +137,7 @@ class Sampler:
137
expanded_local_pos = input_batch.expanded_local_pos
138
pos = input_batch.positions[input_batch.logits_indices]
139
input_ids = input_batch.input_ids[input_batch.logits_indices]
140
+ seq_lens_upper_bound_np = input_batch.seq_lens_cpu_upper_bound.numpy()
141
142
# NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
143
# that num_nans is computed before applying penalties and temperature.
@@ -146,6 +153,7 @@ class Sampler:
153
pos,
154
input_ids,
155
expanded_local_pos,
156
+ seq_lens_upper_bound_np,
157
return_logprobs=logprobs_dims is not None,
158
)
159
@@ -215,6 +223,7 @@ class Sampler:
223
pos: torch.Tensor,
224
input_ids: torch.Tensor,
225
expanded_local_pos: torch.Tensor,
226
+ seq_lens_upper_bound_np: np.ndarray,
227
skip_top_k_top_p: bool = False,
228
) -> torch.Tensor:
229
if not np.any(self.needs_logits_processing[idx_mapping_np]):
@@ -223,39 +232,23 @@ class Sampler:
232
# Copy logits to a new FP32 tensor.
233
logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
234
226
- # Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
227
- self.logit_bias_state.apply_logit_bias(
228
- logits, expanded_idx_mapping, idx_mapping_np, pos
235
+ ctx = LogitsContext(
236
+ expanded_idx_mapping=expanded_idx_mapping,
237
+ idx_mapping=idx_mapping,
238
+ idx_mapping_np=idx_mapping_np,
239
+ expanded_local_pos=expanded_local_pos,
240
+ input_ids=input_ids,
241
+ pos=pos,
242
+ seq_lens_upper_bound_np=seq_lens_upper_bound_np,
243
)
244
231
- # Apply penalties in place.
232
- self.penalties_state.apply_penalties(
233
- logits,
234
- expanded_idx_mapping,
235
- idx_mapping_np,
236
- input_ids,
237
- expanded_local_pos,
238
- )
239
-
240
- # Apply bad words masking in place.
241
- self.bad_words_state.apply_bad_words(
242
- logits,
243
- expanded_idx_mapping,
244
- idx_mapping_np,
245
- input_ids,
246
- expanded_local_pos,
247
- )
245
+ # Apply logits processors (native + any custom).
246
+ for processor in self.logits_processors:
247
+ logits = processor.apply(logits, ctx)
248
249
- # Force the reasoning end marker once a request's thinking budget is
250
- # reached; applied before temperature so the forced token is always kept.
251
- self.thinking_budget_state.apply(
252
- logits,
253
- expanded_idx_mapping,
254
- idx_mapping,
255
- idx_mapping_np,
256
- input_ids,
257
- expanded_local_pos,
258
- )
249
+ # Forcing runs last so no stage can overwrite the forced end marker
250
+ # or weaken it by scaling.
251
+ self.thinking_budget_state.apply(logits, ctx)
252
253
# Apply temperature in place.
254
self.sampling_states.apply_temperature(
@@ -282,6 +275,7 @@ class Sampler:
275
pos: torch.Tensor,
276
input_ids: torch.Tensor,
277
expanded_local_pos: torch.Tensor,
278
+ seq_lens_upper_bound_np: np.ndarray,
279
return_logprobs: bool = False,
280
) -> tuple[torch.Tensor, torch.Tensor]:
281
processed_logits = self.apply_sampling_params(
@@ -292,6 +286,7 @@ class Sampler:
286
pos,
287
input_ids,
288
expanded_local_pos,
289
+ seq_lens_upper_bound_np,
290
skip_top_k_top_p=True,
291
)
292
top_k, top_p = self.sampling_states.get_top_k_top_p(
vllm/v1/worker/gpu/sample/states.py
+15
-4
@@ -38,14 +38,18 @@ class SamplingStates:
38
# -1 means no logprobs are requested.
39
self.num_logprobs.fill(NO_LOGPROBS)
40
41
- def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
42
- self.temperature.np[req_idx] = sampling_params.temperature
43
- self.top_p.np[req_idx] = sampling_params.top_p
41
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
42
+ temperature = sampling_params.temperature
43
top_k = sampling_params.top_k
44
+ top_p = sampling_params.top_p
45
+ min_p = sampling_params.min_p
46
+
47
+ self.temperature.np[req_idx] = temperature
48
+ self.top_p.np[req_idx] = top_p
49
if top_k <= 0 or top_k > self.vocab_size:
50
top_k = self.vocab_size
51
self.top_k.np[req_idx] = top_k
48
- self.min_p.np[req_idx] = sampling_params.min_p
52
+ self.min_p.np[req_idx] = min_p
53
54
seed = sampling_params.seed
55
self.seeds_set[req_idx] = seed is not None
@@ -60,6 +64,13 @@ class SamplingStates:
64
num_logprobs = self.vocab_size
65
self.num_logprobs[req_idx] = num_logprobs
66
67
+ return temperature != 0.0 and (
68
+ temperature != 1.0
69
+ or min_p != 0.0
70
+ or top_k != self.vocab_size
71
+ or top_p != 1.0
72
+ )
73
+
74
def apply_staged_writes(self) -> None:
75
self.temperature.copy_to_uva()
76
self.top_p.copy_to_uva()
vllm/v1/worker/gpu/sample/thinking_budget.py
+12
-16
@@ -9,6 +9,7 @@ from vllm.sampling_params import SamplingParams
9
from vllm.triton_utils import tl, triton
10
from vllm.utils.torch_utils import async_tensor_h2d
11
from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
12
+from vllm.v1.worker.gpu.sample.logits_processor import LogitsContext
13
from vllm.v1.worker.gpu.states import RequestState
14
15
if TYPE_CHECKING:
@@ -78,11 +79,12 @@ class ThinkingBudgetState:
79
end_ids, dtype=torch.int32, device=self.device
80
)
81
81
- def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
82
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
83
if not self.enabled:
83
- return
84
+ return False
85
budget = sampling_params.thinking_token_budget
85
- self.use_thinking_budget[req_idx] = budget is not None
86
+ use_thinking_budget = budget is not None
87
+ self.use_thinking_budget[req_idx] = use_thinking_budget
88
if budget is None:
89
budget = -1
90
else:
@@ -91,6 +93,7 @@ class ThinkingBudgetState:
93
if self.thinking_token_budget.np[req_idx] != budget:
94
self.thinking_token_budget.np[req_idx] = budget
95
self._budget_dirty = True
96
+ return use_thinking_budget
97
98
def apply_staged_writes(self) -> None:
99
if not self.enabled:
@@ -107,27 +110,20 @@ class ThinkingBudgetState:
110
self.thinking_token_budget.copy_to_uva()
111
self._budget_dirty = False
112
110
- def apply(
111
- self,
112
- logits: torch.Tensor,
113
- expanded_idx_mapping: torch.Tensor,
114
- idx_mapping: torch.Tensor,
115
- idx_mapping_np: np.ndarray,
116
- input_ids: torch.Tensor,
117
- expanded_local_pos: torch.Tensor,
118
- ) -> None:
113
+ def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> None:
114
+ idx_mapping_np = ctx.idx_mapping_np
115
if not self.enabled or not np.any(self.use_thinking_budget[idx_mapping_np]):
116
return
117
118
apply_thinking_budget(
119
logits,
124
- idx_mapping,
125
- expanded_idx_mapping,
120
+ ctx.idx_mapping,
121
+ ctx.expanded_idx_mapping,
122
self.thinking_token_budget.gpu,
123
self.req_states.all_token_ids.gpu,
124
self.req_states.total_len.gpu,
129
- input_ids,
130
- expanded_local_pos,
125
+ ctx.input_ids,
126
+ ctx.expanded_local_pos,
127
self.cached_last_start,
128
self.cached_last_end,
129
self.cached_scan_pos,
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+3
@@ -176,6 +176,7 @@ class RejectionSampler:
176
idx_mapping_np: np.ndarray,
177
expanded_idx_mapping: torch.Tensor,
178
expanded_local_pos: torch.Tensor,
179
+ seq_lens_upper_bound_np: np.ndarray,
180
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
181
processed_logits = self.sampler.apply_sampling_params(
182
logits,
@@ -185,6 +186,7 @@ class RejectionSampler:
186
pos,
187
draft_sampled,
188
expanded_local_pos,
189
+ seq_lens_upper_bound_np,
190
)
191
sampled, num_sampled = rejection_sample(
192
processed_logits,
@@ -251,6 +253,7 @@ class RejectionSampler:
253
input_batch.idx_mapping_np[start:end],
254
input_batch.expanded_idx_mapping[lo:hi],
255
input_batch.expanded_local_pos[lo:hi],
256
+ input_batch.seq_lens_cpu_upper_bound.numpy()[start:end],
257
)
258
chunk_logprobs = self._get_logprobs_tensors(
259
sampled,