A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[V1] Wrapper which plumbs request-level logits processors into vLLM batch-level logits processing (#23656)
Signed-off-by: Andrew Feldman <afeldman@redhat.com>
afeldman-nm committed
Sep 2, 2025 at 22:52 UTC
136d853e65a91a21b08227217b51daaba2d5cc71
6 files changed
+524
-5
examples/offline_inference/logits_processor/custom.py
renamed
examples/offline_inference/logits_processor/custom_req.py
new
+151
@@ -0,0 +1,151 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""This example demonstrates wrapping a request-level logits processor to be
5
+compatible with vLLM's batch-level logits processing
6
+
7
+For demo purposes, a dummy logits processor is employed which, if
8
+`target_token` is passed as a keyword argument to `SamplingParams.extra_args`,
9
+will mask out all tokens except `target_token`. This logits processor can be
10
+applied to a vector of logits associated with a single decode step for a single
11
+request. The logits processor cannot be applied to a request which does not
12
+pass in a `target_token` custom argument.
13
+
14
+The request-level dummy logits processor is wrapped to create a batch-level
15
+logits processor, which can apply the logits processor to output logits from
16
+all requests in the persistent batch in a given decode step. For requests which
17
+do not provide a `target_token` argument, the corresponding row of `logits`
18
+will not be modified.
19
+
20
+A batch is constructed with `temperature=0.0` and 50% of requests specifying
21
+`target_token`, and for these requests - and *only* these requests - we
22
+expect the `target_token` to be decoded in each step, yielding an output
23
+similar to that shown below:
24
+
25
+Generated Outputs:
26
+------------------------------------------------------------
27
+Prompt: 'Hello, my name is'
28
+Output: " ' ' ' ' ' ' ' ' ' ' ' ' ' ' ' '"
29
+------------------------------------------------------------
30
+Prompt: 'The president of the United States is'
31
+Output: " not a racist. He is a racist.\nHe's a racist because he"
32
+------------------------------------------------------------
33
+Prompt: 'The capital of France is'
34
+Output: ' also also also also also also also also also also also also also
35
+ also also also'
36
+------------------------------------------------------------
37
+Prompt: 'The future of AI is'
38
+Output: ' in the hands of the people.\n\nThe future of AI is in the'
39
+------------------------------------------------------------
40
+"""
41
+
42
+from typing import Any, Optional
43
+
44
+import torch
45
+
46
+from vllm import LLM, SamplingParams
47
+from vllm.logger import init_logger
48
+from vllm.v1.sample.logits_processor import (
49
+ AdapterLogitsProcessor,
50
+ RequestLogitsProcessor,
51
+)
52
+
53
+logger = init_logger(__name__)
54
+
55
+
56
+class DummyPerReqLogitsProcessor:
57
+ """The request-level logits processor masks out all logits except the
58
+ token id identified by `target_token`"""
59
+
60
+ def __init__(self, target_token: int) -> None:
61
+ """Specify `target_token`"""
62
+ self.target_token = target_token
63
+
64
+ def __call__(
65
+ self,
66
+ output_ids: list[int],
67
+ logits: torch.Tensor,
68
+ ) -> torch.Tensor:
69
+ val_to_keep = logits[self.target_token].item()
70
+ logits[:] = float("-inf")
71
+ logits[self.target_token] = val_to_keep
72
+ return logits
73
+
74
+
75
+class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
76
+ """Example of wrapping a fake request-level logit processor to create a
77
+ batch-level logits processor"""
78
+
79
+ def is_argmax_invariant(self) -> bool:
80
+ return False
81
+
82
+ def new_req_logits_processor(
83
+ self,
84
+ params: SamplingParams,
85
+ ) -> Optional[RequestLogitsProcessor]:
86
+ """This method returns a new request-level logits processor, customized
87
+ to the `target_token` value associated with a particular request.
88
+
89
+ Returns None if the logits processor should not be applied to the
90
+ particular request. To use the logits processor the request must have
91
+ a "target_token" custom argument with an integer value.
92
+
93
+ Args:
94
+ params: per-request sampling params
95
+
96
+ Returns:
97
+ `Callable` request logits processor, or None
98
+ """
99
+ target_token: Optional[Any] = params.extra_args and params.extra_args.get(
100
+ "target_token"
101
+ )
102
+ if target_token is None:
103
+ return None
104
+ if not isinstance(target_token, int):
105
+ logger.warning(
106
+ "target_token value %s is not int; not applying logits"
107
+ " processor to request.",
108
+ target_token,
109
+ )
110
+ return None
111
+ return DummyPerReqLogitsProcessor(target_token)
112
+
113
+
114
+# Sample prompts.
115
+prompts = [
116
+ "Hello, my name is",
117
+ "The president of the United States is",
118
+ "The capital of France is",
119
+ "The future of AI is",
120
+]
121
+# Create a mixture of requests which do and don't utilize the dummy logitproc
122
+sampling_params_list = [
123
+ SamplingParams(temperature=0.0, extra_args={"target_token": 128}),
124
+ SamplingParams(temperature=0.0),
125
+ SamplingParams(temperature=0.0, extra_args={"target_token": 67}),
126
+ SamplingParams(temperature=0.0),
127
+]
128
+
129
+
130
+def main():
131
+ # Create an LLM.
132
+ llm = LLM(
133
+ model="facebook/opt-125m",
134
+ logits_processors=[WrappedPerReqLogitsProcessor],
135
+ )
136
+ # Generate texts from the prompts.
137
+ # The output is a list of RequestOutput objects
138
+ # that contain the prompt, generated text, and other information.
139
+ outputs = llm.generate(prompts, sampling_params_list)
140
+ # Print the outputs.
141
+ print("\nGenerated Outputs:\n" + "-" * 60)
142
+ for output in outputs:
143
+ prompt = output.prompt
144
+ generated_text = output.outputs[0].text
145
+ print(f"Prompt: {prompt!r}")
146
+ print(f"Output: {generated_text!r}")
147
+ print("-" * 60)
148
+
149
+
150
+if __name__ == "__main__":
151
+ main()
examples/offline_inference/logits_processor/custom_req_init.py
new
+165
@@ -0,0 +1,165 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""This example demonstrates a special case of wrapping a request-level logits
5
+processor, namely the case where it is necessary to utilize engine config or
6
+environment info passed to the constructor. The subclass must override the
7
+wrapper base class `__init__()` method to access the engine config, the device
8
+identifier, or the flag which indicates whether pinned memory is available.
9
+
10
+For demo purposes, a request-level dummy logits processor is employed which
11
+causes the same token (`target_token`) to be decoded in each step. The
12
+request-level dummy logits processor is wrapped to create a batch-level logits
13
+processor, which can apply the logits processor to output logits from all
14
+requests in the persistent batch in a given decode step.
15
+
16
+The wrapped dummy logits processor below models a scenario where we must
17
+disable the logits processor on non-"cuda" platforms. The wrapper base class
18
+`__init__()` is overridden in order to check this condition and set a flag.
19
+
20
+A batch is constructed with `temperature=0.0` and 50% of requests specifying
21
+`target_token`, and for these requests - and *only* these requests - we
22
+expect that on a "cuda" device the output will look something like:
23
+
24
+Generated Outputs:
25
+------------------------------------------------------------
26
+Prompt: 'Hello, my name is'
27
+Output: " ' ' ' ' ' ' ' ' ' ' ' ' ' ' ' '"
28
+------------------------------------------------------------
29
+Prompt: 'The president of the United States is'
30
+Output: " not a racist. He is a racist.\nHe's a racist because he"
31
+------------------------------------------------------------
32
+Prompt: 'The capital of France is'
33
+Output: ' also also also also also also also also also also also also also
34
+ also also also'
35
+------------------------------------------------------------
36
+Prompt: 'The future of AI is'
37
+Output: ' in the hands of the people.\n\nThe future of AI is in the'
38
+------------------------------------------------------------
39
+
40
+which indicates that the logits processor is running. However, on a non-"cuda"
41
+device, the first and third requests would not repeat the same token.
42
+"""
43
+
44
+from typing import Optional
45
+
46
+import torch
47
+
48
+from vllm import LLM, SamplingParams
49
+from vllm.config import VllmConfig
50
+from vllm.logger import init_logger
51
+from vllm.v1.sample.logits_processor import (
52
+ AdapterLogitsProcessor,
53
+ RequestLogitsProcessor,
54
+)
55
+
56
+logger = init_logger(__name__)
57
+
58
+
59
+class DummyPerReqLogitsProcessor:
60
+ """The request-level logits processor masks out all logits except the
61
+ token id identified by `target_token`"""
62
+
63
+ def __init__(self, target_token: int) -> None:
64
+ """Specify `target_token`"""
65
+ self.target_token = target_token
66
+
67
+ def __call__(
68
+ self,
69
+ output_ids: list[int],
70
+ logits: torch.Tensor,
71
+ ) -> torch.Tensor:
72
+ val_to_keep = logits[self.target_token].item()
73
+ logits[:] = float("-inf")
74
+ logits[self.target_token] = val_to_keep
75
+ return logits
76
+
77
+
78
+class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
79
+ """Example of overriding the wrapper class `__init__()` in order to utilize
80
+ info about the device type"""
81
+
82
+ def __init__(
83
+ self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool
84
+ ):
85
+ super().__init__(vllm_config, device, is_pin_memory)
86
+ self.is_cuda = device.type == "cuda"
87
+
88
+ def is_argmax_invariant(self) -> bool:
89
+ return False
90
+
91
+ def new_req_logits_processor(
92
+ self,
93
+ params: SamplingParams,
94
+ ) -> Optional[RequestLogitsProcessor]:
95
+ """This method returns a new request-level logits processor, customized
96
+ to the `target_token` value associated with a particular request.
97
+
98
+ Returns None if the logits processor should not be applied to the
99
+ particular request. To use the logits processor the request must have
100
+ a "target_token" custom argument with an integer value, and the device
101
+ must be "cuda"-type
102
+
103
+ Args:
104
+ params: per-request sampling params
105
+
106
+ Returns:
107
+ `Callable` request logits processor, or None
108
+ """
109
+ if (
110
+ not self.is_cuda
111
+ or (
112
+ target_token := params.extra_args
113
+ and params.extra_args.get("target_token")
114
+ )
115
+ is None
116
+ ):
117
+ return None
118
+ if not isinstance(target_token, int):
119
+ logger.warning(
120
+ "target_token value %s is not int; not applying logits"
121
+ " processor to request.",
122
+ target_token,
123
+ )
124
+ return None
125
+ return DummyPerReqLogitsProcessor(target_token)
126
+
127
+
128
+# Sample prompts.
129
+prompts = [
130
+ "Hello, my name is",
131
+ "The president of the United States is",
132
+ "The capital of France is",
133
+ "The future of AI is",
134
+]
135
+# Create a mixture of requests which do and don't utilize the dummy logitproc
136
+sampling_params_list = [
137
+ SamplingParams(temperature=0.0, extra_args={"target_token": 128}),
138
+ SamplingParams(temperature=0.0),
139
+ SamplingParams(temperature=0.0, extra_args={"target_token": 67}),
140
+ SamplingParams(temperature=0.0),
141
+]
142
+
143
+
144
+def main():
145
+ # Create an LLM.
146
+ llm = LLM(
147
+ model="facebook/opt-125m",
148
+ logits_processors=[WrappedPerReqLogitsProcessor],
149
+ )
150
+ # Generate texts from the prompts.
151
+ # The output is a list of RequestOutput objects
152
+ # that contain the prompt, generated text, and other information.
153
+ outputs = llm.generate(prompts, sampling_params_list)
154
+ # Print the outputs.
155
+ print("\nGenerated Outputs:\n" + "-" * 60)
156
+ for output in outputs:
157
+ prompt = output.prompt
158
+ generated_text = output.outputs[0].text
159
+ print(f"Prompt: {prompt!r}")
160
+ print(f"Output: {generated_text!r}")
161
+ print("-" * 60)
162
+
163
+
164
+if __name__ == "__main__":
165
+ main()
tests/v1/logits_processors/test_custom_offline.py
+33
@@ -15,6 +15,7 @@ from tests.v1.logits_processors.utils import (DUMMY_LOGITPROC_ARG,
15
POOLING_MODEL_NAME, TEMP_GREEDY,
16
CustomLogitprocSource,
17
DummyLogitsProcessor,
18
+ WrappedPerReqLogitsProcessor,
19
dummy_module)
20
from tests.v1.logits_processors.utils import entry_points as fake_entry_points
21
from tests.v1.logits_processors.utils import prompts
@@ -161,6 +162,38 @@ def test_custom_logitsprocs(monkeypatch,
162
_run_test(kwargs, logitproc_loaded=True)
163
164
165
+@create_new_process_for_each_test()
166
+def test_custom_logitsprocs_req(monkeypatch):
167
+ """Test passing request-level logits processor to offline Python interface
168
+
169
+ Wrap a request-level logits processor to create a batch level logits
170
+ processor that has a well-defined behavior (mask out all tokens except one
171
+ `target_token`)
172
+
173
+ Construct an `LLM` instance which loads the wrapped logits processor. Pass
174
+ the custom logitproc as a class object.
175
+
176
+ Construct a reference `LLM` instance with no custom logitproc
177
+
178
+ Pass in a batch of requests, 50% of which pass a `target_token` value
179
+ in through `SamplingParams.extra_args`, 50% of which do not.
180
+
181
+ Validate that
182
+ * Requests which do not activate the custom logitproc, yield the same
183
+ results for both `LLM` instances
184
+ * Requests which activate the custom logitproc, only output `target_token`
185
+
186
+ Args:
187
+ monkeypatch: for setting env vars
188
+ """
189
+
190
+ # Test that logitproc info is passed to workers
191
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
192
+ random.seed(40)
193
+ _run_test({"logits_processors": [WrappedPerReqLogitsProcessor]},
194
+ logitproc_loaded=True)
195
+
196
+
197
@create_new_process_for_each_test()
198
@pytest.mark.parametrize("logitproc_source", [
199
CustomLogitprocSource.LOGITPROC_SOURCE_ENTRYPOINT,
tests/v1/logits_processors/utils.py
+64
-3
@@ -3,15 +3,21 @@
3
4
import types
5
from enum import Enum, auto
6
-from typing import Optional
6
+from typing import Any, Optional
7
8
import torch
9
10
from vllm.config import VllmConfig
11
-from vllm.v1.sample.logits_processor import (LOGITSPROCS_GROUP, BatchUpdate,
12
- LogitsProcessor)
11
+from vllm.logger import init_logger
12
+from vllm.sampling_params import SamplingParams
13
+from vllm.v1.sample.logits_processor import (LOGITSPROCS_GROUP,
14
+ AdapterLogitsProcessor,
15
+ BatchUpdate, LogitsProcessor,
16
+ RequestLogitsProcessor)
17
from vllm.v1.sample.logits_processor.builtin import process_dict_updates
18
19
+logger = init_logger(__name__)
20
+
21
MODEL_NAME = "facebook/opt-125m"
22
POOLING_MODEL_NAME = "BAAI/bge-base-en-v1.5"
23
DUMMY_LOGITPROC_ARG = "target_token"
@@ -104,5 +110,60 @@ class EntryPoints(list):
110
self.names = [ep.name for ep in eps]
111
112
113
+class DummyPerReqLogitsProcessor:
114
+ """The request-level logits processor masks out all logits except the
115
+ token id identified by `target_token`"""
116
+
117
+ def __init__(self, target_token: int) -> None:
118
+ """Specify `target_token`"""
119
+ self.target_token = target_token
120
+
121
+ def __call__(
122
+ self,
123
+ output_ids: list[int],
124
+ logits: torch.Tensor,
125
+ ) -> torch.Tensor:
126
+ val_to_keep = logits[self.target_token].item()
127
+ logits[:] = float("-inf")
128
+ logits[self.target_token] = val_to_keep
129
+ return logits
130
+
131
+
132
+class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
133
+ """Example of wrapping a fake request-level logit processor to create a
134
+ batch-level logits processor"""
135
+
136
+ def is_argmax_invariant(self) -> bool:
137
+ return False
138
+
139
+ def new_req_logits_processor(
140
+ self,
141
+ params: SamplingParams,
142
+ ) -> Optional[RequestLogitsProcessor]:
143
+ """This method returns a new request-level logits processor, customized
144
+ to the `target_token` value associated with a particular request.
145
+
146
+ Returns None if the logits processor should not be applied to the
147
+ particular request. To use the logits processor the request must have
148
+ a "target_token" custom argument with an integer value.
149
+
150
+ Args:
151
+ params: per-request sampling params
152
+
153
+ Returns:
154
+ `Callable` request logits processor, or None
155
+ """
156
+ target_token: Optional[
157
+ Any] = params.extra_args and params.extra_args.get("target_token")
158
+ if target_token is None:
159
+ return None
160
+ if not isinstance(target_token, int):
161
+ logger.warning(
162
+ "target_token value %s is not int; not applying logits"
163
+ " processor to request.", target_token)
164
+ return None
165
+ return DummyPerReqLogitsProcessor(target_token)
166
+
167
+
168
"""Fake version of importlib.metadata.entry_points"""
169
entry_points = lambda group: EntryPoints(group)
vllm/v1/sample/logits_processor/__init__.py
+111
-2
@@ -1,16 +1,22 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import importlib
4
+import inspect
5
import itertools
6
+from abc import abstractmethod
7
from collections.abc import Sequence
8
+from functools import partial
9
from typing import TYPE_CHECKING, Optional, Union
10
11
import torch
12
13
from vllm.logger import init_logger
14
+from vllm.logits_process import LogitsProcessor as RequestLogitsProcessor
15
+from vllm.sampling_params import SamplingParams
16
from vllm.v1.sample.logits_processor.builtin import (LogitBiasLogitsProcessor,
17
MinPLogitsProcessor,
13
- MinTokensLogitsProcessor)
18
+ MinTokensLogitsProcessor,
19
+ process_dict_updates)
20
from vllm.v1.sample.logits_processor.interface import (BatchUpdate,
21
LogitsProcessor,
22
MoveDirectionality)
@@ -177,9 +183,112 @@ def build_logitsprocs(
183
BUILTIN_LOGITS_PROCESSORS, custom_logitsprocs_classes))
184
185
186
+class AdapterLogitsProcessor(LogitsProcessor):
187
+ """Wrapper for per-request logits processors
188
+
189
+ To wrap a specific per-request logits processor,
190
+ * Subclass `AdapterLogitsProcessor`
191
+ * Implement `self.is_argmax_invariant()` base-class method
192
+ * Implement `self.new_req_logits_processor(params)`
193
+
194
+ `self.__init__(vllm_config, device, is_pin_memory)` does not need to be
195
+ overridden in general. However, to implement custom constructor behavior -
196
+ especially any logic which operates on or stores `vllm_config`, `device`,
197
+ or `is_pin_memory` - `self.__init__(vllm_config, device, is_pin_memory)`
198
+ must be overriden and the override must call
199
+ `super().__init__(vllm_config, device, is_pin_memory)`
200
+ """
201
+
202
+ def __init__(self, vllm_config: "VllmConfig", device: torch.device,
203
+ is_pin_memory: bool):
204
+ """Subclass must invoke
205
+ `super().__init__(vllm_config, device, is_pin_memory)`.
206
+
207
+ Subclass constructor may find it useful to utilize the `vllm_config`,
208
+ `device` and `is_pin_memory` argument. However regardless of whether
209
+ these arguments are used, the vLLM logits processor interface requires
210
+ all three arguments to be present.
211
+ """
212
+
213
+ # Map req index -> logits processor state
214
+ #
215
+ # State representation is a partial[Tensor] comprising a request-level
216
+ # logits processor with the output token ids argument and (if required)
217
+ # the prompt token ids argument pre-populated
218
+ #
219
+ # Note that the partial carries a *reference* to output token ids, and
220
+ # will thus always operate on the list as it is currently, not as it
221
+ # was when the partial was created.
222
+ self.req_info: dict[int, partial[torch.Tensor]] = {}
223
+
224
+ @abstractmethod
225
+ def new_req_logits_processor(
226
+ self,
227
+ params: SamplingParams,
228
+ ) -> Optional[RequestLogitsProcessor]:
229
+ """Consume request info; return a per-request logits processor.
230
+
231
+ Return None if logits processor does not need to be applied to request
232
+
233
+ Args:
234
+ params: request sampling params
235
+
236
+ Returns:
237
+ None if logits processor should not be applied to request; otherwise
238
+ returns a `RequestLogitsProcessor` instance
239
+
240
+ """
241
+ raise NotImplementedError
242
+
243
+ def _new_state(
244
+ self,
245
+ params: SamplingParams,
246
+ prompt_ids: list[int],
247
+ output_ids: list[int],
248
+ ) -> Optional[partial[torch.Tensor]]:
249
+ """Return state representation for new request
250
+
251
+ Returns None if logits processor is not applicable to request
252
+
253
+ Args:
254
+ params: request sampling params
255
+ prompt_ids: request prompt token ids
256
+ output_ids: decoded tokens so far for this request
257
+
258
+ Returns:
259
+ logits processor partial[Tensor] or None
260
+
261
+ """
262
+ if req_lp := self.new_req_logits_processor(params):
263
+ args = [prompt_ids, output_ids] if (len(
264
+ inspect.signature(req_lp).parameters) == 3) else [output_ids]
265
+ return partial(req_lp, *args)
266
+ return None
267
+
268
+ def update_state(self, batch_update: Optional[BatchUpdate]):
269
+ process_dict_updates(
270
+ self.req_info,
271
+ batch_update,
272
+ self._new_state,
273
+ )
274
+
275
+ def apply(self, logits: torch.Tensor) -> torch.Tensor:
276
+ if self.req_info:
277
+ # Apply per-request logits processors to corresponding rows of
278
+ # logits tensor
279
+ for req_idx, req_lp in self.req_info.items():
280
+ req_logits = logits[req_idx]
281
+ new_logits = req_lp(req_logits)
282
+ if new_logits is not req_logits:
283
+ # Modify logits tensor row in-place if necessary
284
+ logits[req_idx] = new_logits
285
+ return logits
286
+
287
+
288
__all__ = [
289
"LogitsProcessor", "LogitBiasLogitsProcessor", "MinPLogitsProcessor",
290
"MinTokensLogitsProcessor", "BatchUpdate", "BatchUpdateBuilder",
291
"MoveDirectionality", "LogitsProcessors", "build_logitsprocs",
184
- "STR_POOLING_REJECTS_LOGITSPROCS", "LOGITSPROCS_GROUP"
292
+ "STR_POOLING_REJECTS_LOGITSPROCS", "LOGITSPROCS_GROUP",
293
+ "AdapterLogitsProcessor"
294
]