@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bf7f470b22
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[V1] Logits processors extensibility (#19912)

Signed-off-by: Andrew Feldman <afeldman@redhat.com> Signed-off-by: Andrew Feldman <afeld2012@gmail.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: Andrew Feldman <afeld2012@gmail.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

afeldman-nm committed Aug 16, 2025 at 15:59 UTC bf7f470b22e8bf26e1edb30b3bf465ab7dd69f0c
22 files changed +1313 -335
.buildkite/test-pipeline.yaml
+1
@@ -253,6 +253,7 @@ steps:
253 - pytest -v -s v1/engine
254 - pytest -v -s v1/entrypoints
255 - pytest -v -s v1/sample
256 + - pytest -v -s v1/logits_processors
257 - pytest -v -s v1/worker
258 - pytest -v -s v1/structured_output
259 - pytest -v -s v1/spec_decode
examples/offline_inference/logits_processor.py new
+147
@@ -0,0 +1,147 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""This example demonstrates instantiating vLLM with a custom logits processor
5 +class object.
6 +
7 +For a basic example of implementing a custom logits processor, see
8 +the `DummyLogitsProcessor` implementation in `vllm/test_utils.py`.
9 +
10 +For testing purposes, a dummy logits processor is employed which, if
11 +`target_token` is passed as a keyword argument to `SamplingParams.extra_args`,
12 +will mask out all tokens except `target_token`.
13 +
14 +A batch is constructed with `temperature=0.0` and 50% of requests specifying
15 +`target_token`, and for these requests - and *only* these requests - we
16 +expect the `target_token` to be decoded in each step, yielding an output
17 +similar to that shown below:
18 +
19 +Generated Outputs:
20 +------------------------------------------------------------
21 +Prompt: 'Hello, my name is'
22 +Output: " ' ' ' ' ' ' ' ' ' ' ' ' ' ' ' '"
23 +------------------------------------------------------------
24 +Prompt: 'The president of the United States is'
25 +Output: " not a racist. He is a racist.\nHe's a racist because he"
26 +------------------------------------------------------------
27 +Prompt: 'The capital of France is'
28 +Output: ' also also also also also also also also also also also also also
29 + also also also'
30 +------------------------------------------------------------
31 +Prompt: 'The future of AI is'
32 +Output: ' in the hands of the people.\n\nThe future of AI is in the'
33 +------------------------------------------------------------
34 +"""
35 +
36 +from typing import Optional
37 +
38 +import torch
39 +
40 +from vllm import LLM, SamplingParams
41 +from vllm.config import VllmConfig
42 +from vllm.v1.sample.logits_processor import (
43 + BatchUpdate,
44 + LogitsProcessor,
45 + MoveDirectionality,
46 +)
47 +
48 +
49 +# Hypothetical custom logits processor
50 +class DummyLogitsProcessor(LogitsProcessor):
51 + """Fake logit processor to support unit testing and examples"""
52 +
53 + def __init__(
54 + self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool
55 + ):
56 + self.req_info: dict[int, SamplingParams] = {}
57 +
58 + def is_argmax_invariant(self) -> bool:
59 + """Never impacts greedy sampling"""
60 + return False
61 +
62 + def update_state(self, batch_update: Optional[BatchUpdate]):
63 + if not batch_update:
64 + return
65 +
66 + # Process added requests.
67 + for index, params, _, _ in batch_update.added:
68 + assert params is not None
69 + if params.extra_args and (
70 + target_token := params.extra_args.get("target_token")
71 + ):
72 + self.req_info[index] = target_token
73 +
74 + if self.req_info:
75 + # Process removed requests.
76 + for index in batch_update.removed:
77 + self.req_info.pop(index, None)
78 +
79 + # Process moved requests, unidirectional move (a->b) and swap
80 + # (a<->b)
81 + for adx, bdx, direct in batch_update.moved:
82 + a_val = self.req_info.pop(adx, None)
83 + b_val = self.req_info.pop(bdx, None)
84 + if a_val is not None:
85 + self.req_info[bdx] = a_val
86 + if direct == MoveDirectionality.SWAP and b_val is not None:
87 + self.req_info[adx] = b_val
88 +
89 + def apply(self, logits: torch.Tensor) -> torch.Tensor:
90 + if not self.req_info:
91 + return logits
92 +
93 + # Save target values before modification
94 + rows_list = list(self.req_info.keys())
95 + cols = torch.tensor(
96 + [self.req_info[i] for i in rows_list],
97 + dtype=torch.long,
98 + device=logits.device,
99 + )
100 + rows = torch.tensor(rows_list, dtype=torch.long, device=logits.device)
101 + values_to_keep = logits[rows, cols].clone()
102 +
103 + # Mask all but target tokens
104 + logits[rows] = float("-inf")
105 + logits[rows, cols] = values_to_keep
106 +
107 + return logits
108 +
109 +
110 +# Sample prompts.
111 +prompts = [
112 + "Hello, my name is",
113 + "The president of the United States is",
114 + "The capital of France is",
115 + "The future of AI is",
116 +]
117 +# Create a mixture of requests which do and don't utilize the dummy logitproc
118 +sampling_params_list = [
119 + SamplingParams(temperature=0.0, extra_args={"target_token": 128}),
120 + SamplingParams(temperature=0.0),
121 + SamplingParams(temperature=0.0, extra_args={"target_token": 67}),
122 + SamplingParams(temperature=0.0),
123 +]
124 +
125 +
126 +def main():
127 + # Create an LLM.
128 + llm = LLM(
129 + model="facebook/opt-125m",
130 + logits_processors=[DummyLogitsProcessor],
131 + )
132 + # Generate texts from the prompts.
133 + # The output is a list of RequestOutput objects
134 + # that contain the prompt, generated text, and other information.
135 + outputs = llm.generate(prompts, sampling_params_list)
136 + # Print the outputs.
137 + print("\nGenerated Outputs:\n" + "-" * 60)
138 + for output in outputs:
139 + prompt = output.prompt
140 + generated_text = output.outputs[0].text
141 + print(f"Prompt: {prompt!r}")
142 + print(f"Output: {generated_text!r}")
143 + print("-" * 60)
144 +
145 +
146 +if __name__ == "__main__":
147 + main()
tests/utils.py
+66 -13
@@ -13,6 +13,7 @@ import tempfile
13 import time
14 import warnings
15 from contextlib import contextmanager, suppress
16 +from multiprocessing import Process
17 from pathlib import Path
18 from typing import Any, Callable, Literal, Optional, Union
19
@@ -76,6 +77,23 @@ VLLM_PATH = Path(__file__).parent.parent
77 class RemoteOpenAIServer:
78 DUMMY_API_KEY = "token-abc123" # vLLM's OpenAI server does not need API key
79
80 + def _start_server(self, model: str, vllm_serve_args: list[str],
81 + env_dict: Optional[dict[str, str]]) -> None:
82 + """Subclasses override this method to customize server process launch
83 + """
84 + env = os.environ.copy()
85 + # the current process might initialize cuda,
86 + # to be safe, we should use spawn method
87 + env['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'
88 + if env_dict is not None:
89 + env.update(env_dict)
90 + self.proc: subprocess.Popen = subprocess.Popen(
91 + ["vllm", "serve", model, *vllm_serve_args],
92 + env=env,
93 + stdout=sys.stdout,
94 + stderr=sys.stderr,
95 + )
96 +
97 def __init__(self,
98 model: str,
99 vllm_serve_args: list[str],
@@ -128,18 +146,7 @@ class RemoteOpenAIServer:
146 model_loader = get_model_loader(load_config)
147 model_loader.download_model(model_config)
148
131 - env = os.environ.copy()
132 - # the current process might initialize cuda,
133 - # to be safe, we should use spawn method
134 - env['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'
135 - if env_dict is not None:
136 - env.update(env_dict)
137 - self.proc = subprocess.Popen(
138 - ["vllm", "serve", model, *vllm_serve_args],
139 - env=env,
140 - stdout=sys.stdout,
141 - stderr=sys.stderr,
142 - )
149 + self._start_server(model, vllm_serve_args, env_dict)
150 max_wait_seconds = max_wait_seconds or 240
151 self._wait_for_server(url=self.url_for("health"),
152 timeout=max_wait_seconds)
@@ -155,6 +162,10 @@ class RemoteOpenAIServer:
162 # force kill if needed
163 self.proc.kill()
164
165 + def _poll(self) -> Optional[int]:
166 + """Subclasses override this method to customize process polling"""
167 + return self.proc.poll()
168 +
169 def _wait_for_server(self, *, url: str, timeout: float):
170 # run health check
171 start = time.time()
@@ -169,7 +180,7 @@ class RemoteOpenAIServer:
180 # which means the server is not ready yet.
181 # the stack trace is not useful, so we suppress it
182 # by using `raise from None`.
172 - result = self.proc.poll()
183 + result = self._poll()
184 if result is not None and result != 0:
185 raise RuntimeError("Server exited unexpectedly.") from None
186
@@ -205,6 +216,48 @@ class RemoteOpenAIServer:
216 **kwargs)
217
218
219 +class RemoteOpenAIServerCustom(RemoteOpenAIServer):
220 + """Launch test server with custom child process"""
221 +
222 + def _start_server(self, model: str, vllm_serve_args: list[str],
223 + env_dict: Optional[dict[str, str]]) -> None:
224 + self.proc: Process = Process(
225 + target=self.child_process_fxn,
226 + args=(env_dict, model,
227 + vllm_serve_args)) # type: ignore[assignment]
228 + self.proc.start()
229 +
230 + def __init__(self,
231 + model: str,
232 + vllm_serve_args: list[str],
233 + child_process_fxn: Callable[
234 + [Optional[dict[str, str]], str, list[str]], None],
235 + *,
236 + env_dict: Optional[dict[str, str]] = None,
237 + seed: Optional[int] = 0,
238 + auto_port: bool = True,
239 + max_wait_seconds: Optional[float] = None) -> None:
240 + """Store custom child process function then invoke superclass
241 + constructor which will indirectly launch it."""
242 + self.child_process_fxn = child_process_fxn
243 + super().__init__(model=model,
244 + vllm_serve_args=vllm_serve_args,
245 + env_dict=env_dict,
246 + seed=seed,
247 + auto_port=auto_port,
248 + max_wait_seconds=max_wait_seconds)
249 +
250 + def _poll(self) -> Optional[int]:
251 + return self.proc.exitcode
252 +
253 + def __exit__(self, exc_type, exc_value, traceback):
254 + self.proc.terminate()
255 + self.proc.join(8)
256 + if self.proc.is_alive():
257 + # force kill if needed
258 + self.proc.kill()
259 +
260 +
261 def _test_completion(
262 client: openai.OpenAI,
263 model: str,
tests/v1/logits_processors/__init__.py
tests/v1/logits_processors/test_correctness.py renamed
+16 -8
@@ -9,11 +9,13 @@ import numpy as np
9 import pytest
10 import torch
11
12 +from tests.utils import create_new_process_for_each_test
13 from tests.v1.sample.utils import (LogitsprocsTestFakes, create_fake_logits,
14 create_penalty_tensor,
15 create_prompt_tokens_tensor,
16 fake_apply_logitsprocs,
17 fake_update_logitsprocs_state)
18 +from vllm.config import VllmConfig
19 from vllm.platforms import current_platform
20 from vllm.sampling_params import SamplingParams
21 from vllm.utils import is_pin_memory_available
@@ -24,7 +26,7 @@ from vllm.v1.sample.logits_processor import (BatchUpdate, BatchUpdateBuilder,
26 MinPLogitsProcessor,
27 MinTokensLogitsProcessor,
28 MoveDirectionality,
27 - init_builtin_logitsprocs)
29 + build_logitsprocs)
30 # yapf: enable
31 from vllm.v1.sample.metadata import SamplingMetadata
32
@@ -53,6 +55,7 @@ class LogitsProcsRequestParams:
55 workload_index: int
56 logitproc_type: LogitprocType # Logitproc enabled, specified by str id
57 out_tokens: list[int] # Output tokens required for min tokens test
58 + prompt_tokens: list[int] # Dummy prompt tokens placeholder
59 params: SamplingParams # Settings customized for logitproc
60
61 def __init__(self, workload_index: int, logitproc_type: LogitprocType):
@@ -63,6 +66,7 @@ class LogitsProcsRequestParams:
66 # don't matter *for these tests* so use 0 as a dummy value
67 self.out_tokens = ([0] *
68 (MIN_TOKENS_LEN_THRESHOLD * random.randint(0, 2)))
69 + self.prompt_tokens = []
70 self.params = _sampling_params_from_logitproc(logitproc_type)
71
72 def __str__(self):
@@ -88,11 +92,12 @@ def _generate_fake_sampling_metadata(
92 vocab_size,
93 size=np.random.randint(
94 1, MAX_NUM_PROMPT_TOKENS)).tolist())
91 - logitsprocs = init_builtin_logitsprocs(
92 - pin_memory_available=PIN_MEMORY_AVAILABLE,
93 - max_num_reqs=MAX_NUM_REQS + 1,
94 - device=device)
95 -
95 + logitsprocs = build_logitsprocs(
96 + vllm_config=VllmConfig(),
97 + device=device,
98 + is_pin_memory=PIN_MEMORY_AVAILABLE,
99 + is_pooling_model=False,
100 + )
101 fake_sampling_metadata = SamplingMetadata(
102 temperature=torch.full((batch_size, ), 0.0),
103 all_greedy=True,
@@ -462,7 +467,8 @@ def _generate_fake_step_update(
467 # Replace as many removed requests as possible with added requests
468 add_remove_idx = batch_update_builder.pop_removed()
469 batch_update_builder.added.append(
465 - (add_remove_idx, add_req_params.params, add_req_params.out_tokens))
470 + (add_remove_idx, add_req_params.params,
471 + add_req_params.prompt_tokens, add_req_params.out_tokens))
472 persistent_batch[add_remove_idx] = add_req_params
473
474 # Append remaining added requests to end of batch
@@ -470,7 +476,8 @@ def _generate_fake_step_update(
476 num_step_add_replace):(wdx +
477 num_step_add)]
478 batch_update_builder.added.extend([
473 - (adx + batch_size, add_req_params.params, add_req_params.out_tokens)
479 + (adx + batch_size, add_req_params.params, add_req_params.prompt_tokens,
480 + add_req_params.out_tokens)
481 for adx, add_req_params in enumerate(add_reqs_append)
482 ])
483 persistent_batch.extend(add_reqs_append)
@@ -561,6 +568,7 @@ def _assert_valid(
568 step_idx=step_idx)
569
570
571 +@create_new_process_for_each_test()
572 @pytest.mark.parametrize("device", CUDA_DEVICES)
573 @pytest.mark.parametrize("reqs_per_logitproc", [REQS_PER_LOGITPROC])
574 @pytest.mark.parametrize("logitsprocs_under_test", _get_test_cases())
tests/v1/logits_processors/test_custom_offline.py new
+237
@@ -0,0 +1,237 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import random
4 +import sys
5 +from typing import Union
6 +
7 +import pytest
8 +
9 +from tests.utils import create_new_process_for_each_test
10 +# yapf: disable
11 +from tests.v1.logits_processors.utils import (DUMMY_LOGITPROC_ARG,
12 + DUMMY_LOGITPROC_FQCN,
13 + DUMMY_LOGITPROC_MODULE,
14 + MAX_TOKENS, MODEL_NAME,
15 + POOLING_MODEL_NAME, TEMP_GREEDY,
16 + CustomLogitprocSource,
17 + DummyLogitsProcessor,
18 + dummy_module)
19 +from tests.v1.logits_processors.utils import entry_points as fake_entry_points
20 +from tests.v1.logits_processors.utils import prompts
21 +# yapf: enable
22 +from vllm import LLM, SamplingParams
23 +from vllm.v1.sample.logits_processor import (STR_POOLING_REJECTS_LOGITSPROCS,
24 + LogitsProcessor)
25 +
26 +# Create a mixture of requests which do and don't utilize the dummy logitproc
27 +sampling_params_list = [
28 + SamplingParams(temperature=TEMP_GREEDY,
29 + max_tokens=MAX_TOKENS,
30 + extra_args={DUMMY_LOGITPROC_ARG: 128}),
31 + SamplingParams(temperature=TEMP_GREEDY, max_tokens=MAX_TOKENS),
32 + SamplingParams(temperature=TEMP_GREEDY,
33 + max_tokens=MAX_TOKENS,
34 + extra_args={DUMMY_LOGITPROC_ARG: 67}),
35 + SamplingParams(temperature=TEMP_GREEDY, max_tokens=MAX_TOKENS),
36 +]
37 +
38 +
39 +def _run_test(kwargs: dict, logitproc_loaded: bool) -> None:
40 + """Compare `LLM` instance initialized with specified `kwargs` against
41 + reference `LLM` instance.
42 +
43 + Two scenarios:
44 + 1. Server has loaded dummy logitproc; test that requests which specify
45 + dummy logitproc arg value behave as if logitproc is operating (output
46 + token value should repeat), while requests that don't specify dummy
47 + logitproc arg value should match reference `LLM` output.
48 + 2. Server has *not* loaded dummy logitproc; test that all requests
49 + behave as if logitproc is *not* operating (output matches reference
50 + `LLM` output.)
51 +
52 + Args:
53 + kwargs: `LLM` constructor kwargs
54 + logitproc_loaded: server has loaded dummy logitproc if True
55 + """
56 +
57 + # Create a vLLM instance and load custom logitproc
58 + llm_logitproc = LLM(
59 + model=MODEL_NAME,
60 + gpu_memory_utilization=0.1,
61 + **kwargs,
62 + )
63 +
64 + # Create a reference vLLM instance without custom logitproc
65 + llm_ref = LLM(model=MODEL_NAME, gpu_memory_utilization=0.1)
66 +
67 + # Run inference with logitproc loaded
68 + outputs_logitproc = llm_logitproc.generate(prompts, sampling_params_list)
69 +
70 + # Reference run
71 + outputs_ref = llm_ref.generate(prompts, sampling_params_list)
72 +
73 + # Validate outputs
74 + for bdx, (out_lp, out_ref, params) in enumerate(
75 + zip(outputs_logitproc, outputs_ref, sampling_params_list)):
76 + lp_toks = out_lp.outputs[0].token_ids
77 + if logitproc_loaded and params.extra_args:
78 + # This request exercises custom logitproc; validate that logitproc
79 + # forces `target_token` to be decoded in each step
80 + target_token = params.extra_args[DUMMY_LOGITPROC_ARG]
81 + if not all(x == target_token for x in lp_toks):
82 + raise AssertionError(
83 + f"Request {bdx} generated {lp_toks}, shoud all be "
84 + f"{target_token}")
85 + else:
86 + # This request does not exercise custom logitproc (or custom
87 + # logitproc is not enabled on this server); validate against
88 + # reference result
89 + ref_toks = out_ref.outputs[0].token_ids
90 + if lp_toks != ref_toks:
91 + raise AssertionError(
92 + f"Request {bdx} generated {lp_toks}, should match "
93 + f"{ref_toks}")
94 +
95 +
96 +@create_new_process_for_each_test()
97 +@pytest.mark.parametrize("logitproc_source", list(CustomLogitprocSource))
98 +def test_custom_logitsprocs(monkeypatch,
99 + logitproc_source: CustomLogitprocSource):
100 + """Test offline Python interface for passing custom logitsprocs
101 +
102 + Construct an `LLM` instance which loads a custom logitproc that has a
103 + well-defined behavior (mask out all tokens except one `target_token`)
104 +
105 + Construct a reference `LLM` instance with no custom logitproc
106 +
107 + Pass in a batch of requests, 50% of which pass a `target_token` value
108 + in through `SamplingParams.extra_args`, 50% of which do not.
109 +
110 + Validate that
111 + * Requests which do not activate the custom logitproc, yield the same
112 + results for both `LLM` instances
113 + * Requests which activate the custom logitproc, only output `target_token`
114 +
115 + Test four scenarios, corresponding to `logitproc_source` value
116 + * No logitsprocs loaded - test that generated tokens match reference `LLM`
117 + instance output
118 + * Logitproc passed in via {entrypoint, class object, fully-qualified class
119 + name (FQCN)} - test that dummy logitproc is utilized correctly when
120 + provided via any of these three possible sources
121 +
122 + Args:
123 + monkeypatch: for setting env vars
124 + logitproc_source: what source (entrypoint, fully-qualified class name
125 + (FQCN), class object, or None) the user pulls the
126 + logitproc from
127 + """
128 +
129 + # Test that logitproc info is passed to workers
130 + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
131 + random.seed(40)
132 +
133 + # Choose LLM args based on logitproc source
134 + if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_NONE:
135 + # Scenario: the server does not load any custom logitproc
136 + # Every other scenario is a different way of loading a custom logitproc
137 + _run_test({}, logitproc_loaded=False)
138 + return
139 +
140 + if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_ENTRYPOINT:
141 + # Scenario: vLLM loads a logitproc from a preconfigured entrypoint
142 + # To that end, mock a dummy logitproc entrypoint
143 + import importlib.metadata
144 + importlib.metadata.entry_points = fake_entry_points # type: ignore
145 +
146 + # fork is required for workers to see entrypoint patch
147 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")
148 + _run_test({}, logitproc_loaded=True)
149 + return
150 +
151 + kwargs: dict[str, list[Union[str, type[LogitsProcessor]]]] = {}
152 + if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_FQCN:
153 + # Scenario: load logitproc based on fully-qualified class name (FQCN)
154 + # Inject dummy module which defines logitproc
155 + sys.modules[DUMMY_LOGITPROC_MODULE] = dummy_module
156 + kwargs["logits_processors"] = [DUMMY_LOGITPROC_FQCN]
157 + elif logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_CLASS:
158 + # Scenario: load logitproc from provided class object
159 + kwargs["logits_processors"] = [DummyLogitsProcessor]
160 +
161 + _run_test(kwargs, logitproc_loaded=True)
162 +
163 +
164 +@create_new_process_for_each_test()
165 +@pytest.mark.parametrize("logitproc_source", [
166 + CustomLogitprocSource.LOGITPROC_SOURCE_ENTRYPOINT,
167 + CustomLogitprocSource.LOGITPROC_SOURCE_FQCN,
168 + CustomLogitprocSource.LOGITPROC_SOURCE_CLASS,
169 +])
170 +def test_pooling_rejects_custom_logitsprocs(
171 + monkeypatch, logitproc_source: CustomLogitprocSource):
172 + """Validate that vLLM engine initialization properly rejects custom
173 + logitsprocs when the model is a pooling model.
174 +
175 + Use `LLM` entrypoint. We expect `LLM` initialization to fail before the
176 + logitproc is actually loaded.
177 +
178 + Scenario 1:
179 + * Mock a logitproc entrypoint
180 + * Validate that `LLM` does not load the logitproc
181 +
182 + Scenario 2:
183 + * Pass custom logitproc to `LLM` constructor
184 + * Scenario 2a: via FQCN
185 + * Scenario 2b: via class object
186 + * Validate that initialization fails with appropriate exception
187 +
188 + Args:
189 + monkeypatch: used to set environment variables
190 + logitproc_source: what source (entrypoint, fully-qualified class name
191 + (FQCN), or class object) the user pulls the
192 + logitproc from
193 + """
194 + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
195 + random.seed(40)
196 +
197 + if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_ENTRYPOINT:
198 + # Scenario: vLLM loads a pooling model and ignores a logitproc that is
199 + # available at a preconfigured entrypoint
200 +
201 + # Patch in dummy logitproc entrypoint
202 + import importlib.metadata
203 + importlib.metadata.entry_points = fake_entry_points # type: ignore
204 +
205 + # fork is required for entrypoint patch to be visible to workers,
206 + # although they should ignore the entrypoint patch anyway
207 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")
208 +
209 + llm = LLM(
210 + runner="pooling",
211 + model=POOLING_MODEL_NAME,
212 + gpu_memory_utilization=0.1,
213 + )
214 + # Require that no logitsprocs have been loaded
215 + assert sum([
216 + 1 for _ in llm.llm_engine.model_executor.driver_worker.worker.
217 + model_runner.input_batch.logitsprocs.all
218 + ]) == 0
219 + return
220 +
221 + kwargs: dict[str, list[Union[str, type[LogitsProcessor]]]] = {}
222 + if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_FQCN:
223 + # Scenario: load logitproc based on fully-qualified class name (FQCN)
224 + kwargs["logits_processors"] = [DUMMY_LOGITPROC_FQCN]
225 + elif logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_CLASS:
226 + # Scenario: load logitproc from provided class object
227 + kwargs["logits_processors"] = [DummyLogitsProcessor]
228 +
229 + with pytest.raises(ValueError, match=STR_POOLING_REJECTS_LOGITSPROCS):
230 + # Require that loading a pooling model alongside the logitproc raises
231 + # the appropriate exception.
232 + LLM(
233 + runner="pooling",
234 + model=POOLING_MODEL_NAME,
235 + gpu_memory_utilization=0.1,
236 + **kwargs,
237 + )
tests/v1/logits_processors/test_custom_online.py new
+180
@@ -0,0 +1,180 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import os
5 +import random
6 +import sys
7 +from typing import Any, Optional
8 +
9 +import openai
10 +import pytest
11 +import pytest_asyncio
12 +
13 +from tests.utils import (RemoteOpenAIServerCustom,
14 + create_new_process_for_each_test)
15 +# yapf: disable
16 +from tests.v1.logits_processors.utils import (DUMMY_LOGITPROC_ARG,
17 + DUMMY_LOGITPROC_FQCN,
18 + DUMMY_LOGITPROC_MODULE,
19 + MAX_TOKENS, MODEL_NAME,
20 + TEMP_GREEDY, dummy_module)
21 +from tests.v1.logits_processors.utils import entry_points as fake_entry_points
22 +from tests.v1.logits_processors.utils import prompts
23 +
24 +# yapf: enable
25 +
26 +
27 +def _server_with_logitproc_entrypoint(
28 + env_dict: Optional[dict[str, str]],
29 + model: str,
30 + vllm_serve_args: list[str],
31 +) -> None:
32 + """Start vLLM server, inject dummy logitproc entrypoint"""
33 +
34 + # Patch `entry_points` to inject logitproc entrypoint
35 + import importlib.metadata
36 + importlib.metadata.entry_points = fake_entry_points # type: ignore
37 + from vllm.entrypoints.cli import main
38 +
39 + # fork is required for workers to see entrypoint patch
40 + os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = "fork"
41 + if env_dict is not None:
42 + os.environ.update(env_dict)
43 +
44 + # Emulate `vllm serve <model> <CLI args>`
45 + sys.argv = ["vllm", "serve", model] + vllm_serve_args
46 + main.main()
47 +
48 +
49 +def _server_with_logitproc_module(
50 + env_dict: Optional[dict[str, str]],
51 + model: str,
52 + vllm_serve_args: list[str],
53 +) -> None:
54 + """Start vLLM server, inject module with dummy logitproc"""
55 +
56 + # Patch `modules` to inject dummy logitproc module
57 + from vllm.entrypoints.cli import main
58 + sys.modules[DUMMY_LOGITPROC_MODULE] = dummy_module
59 +
60 + # fork is required for workers to see entrypoint patch
61 + os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = "fork"
62 + if env_dict is not None:
63 + os.environ.update(env_dict)
64 +
65 + # Emulate `vllm serve <model> <CLI args>`
66 + sys.argv = ["vllm", "serve", model] + vllm_serve_args
67 + main.main()
68 +
69 +
70 +@pytest.fixture(scope="module")
71 +def default_server_args():
72 + return [
73 + # use half precision for speed and memory savings in CI environment
74 + "--dtype",
75 + "bfloat16",
76 + "--max-model-len",
77 + "2048",
78 + "--max-num-seqs",
79 + "128",
80 + ]
81 +
82 +
83 +@pytest.fixture(scope="function",
84 + params=[[], ["--logits-processors", DUMMY_LOGITPROC_FQCN]])
85 +def server(default_server_args, request, monkeypatch):
86 + """Consider two server configurations:
87 + (1) --logits-processors cli arg specifies dummy logits processor via fully-
88 + qualified class name (FQCN); patch in a dummy logits processor module
89 + (2) No --logits-processors cli arg; patch in a dummy logits processor
90 + entrypoint
91 + """
92 +
93 + # Test that logitproc info is passed to workers
94 + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
95 +
96 + if request.param:
97 + # Launch server, append FQCN argument, inject dummy logitproc module
98 + args = default_server_args + request.param
99 + _server_fxn = _server_with_logitproc_module
100 + else:
101 + # Launch server, inject dummy logitproc entrypoint
102 + args = default_server_args
103 + _server_fxn = _server_with_logitproc_entrypoint
104 +
105 + with RemoteOpenAIServerCustom(MODEL_NAME, args,
106 + _server_fxn) as remote_server:
107 + yield remote_server
108 +
109 +
110 +@pytest_asyncio.fixture
111 +async def client(server):
112 + async with server.get_async_client() as async_client:
113 + yield async_client
114 +
115 +
116 +# General request argument values for these tests
117 +api_keyword_args = {
118 + # Greedy sampling ensures that requests which receive the `target_token`
119 + # arg will decode it in every step
120 + "temperature": TEMP_GREEDY,
121 + # Since EOS will never be decoded (unless `target_token` is EOS)
122 + "max_tokens": MAX_TOKENS,
123 + # Return decoded token logprobs (as a way of getting token id)
124 + "logprobs": 0,
125 +}
126 +
127 +
128 +@create_new_process_for_each_test()
129 +@pytest.mark.asyncio
130 +@pytest.mark.parametrize(
131 + "model_name",
132 + [MODEL_NAME],
133 +)
134 +async def test_custom_logitsprocs(client: openai.AsyncOpenAI, model_name: str):
135 + """Test custom logitsprocs when starting OpenAI server from CLI
136 +
137 + Launch vLLM OpenAI-compatible server, configured to load a custom logitproc
138 + that has a well-defined behavior (mask out all tokens except one
139 + `target_token`).
140 +
141 + Pass in requests, 50% of which pass a `target_token` value
142 + in through `extra_body["vllm_xargs"]`, 50% of which do not.
143 +
144 + Validate that requests which activate the custom logitproc, repeat the same
145 + token
146 + """
147 +
148 + use_dummy_logitproc = True
149 + for prompt in prompts:
150 + # Build request arguments
151 + request_keyword_args: dict[str, Any] = {
152 + **api_keyword_args,
153 + }
154 + if use_dummy_logitproc:
155 + # 50% of requests pass target_token custom arg
156 + target_token = random.choice([128, 67])
157 + # For requests which activate the dummy logitproc, choose one of
158 + # two `target_token` values which are known not to be EOS tokens
159 + request_keyword_args["extra_body"] = {
160 + "vllm_xargs": {
161 + DUMMY_LOGITPROC_ARG: target_token
162 + }
163 + }
164 + batch = await client.completions.create(
165 + model=model_name,
166 + prompt=prompt,
167 + **request_keyword_args,
168 + )
169 +
170 + if use_dummy_logitproc:
171 + # Only for requests which activate dummy logitproc - validate that
172 + # output token is repeated
173 + choices: openai.types.CompletionChoice = batch.choices
174 + toks = choices[0].logprobs.tokens
175 + if not all([x == toks[0] for x in toks]):
176 + raise AssertionError(
177 + f"Generated {toks} should all be {toks[0]}")
178 +
179 + # Alternate whether to activate dummy logitproc for each request
180 + use_dummy_logitproc = not use_dummy_logitproc
tests/v1/logits_processors/utils.py new
+127
@@ -0,0 +1,127 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import types
5 +from enum import Enum, auto
6 +from typing import Optional
7 +
8 +import torch
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.sampling_params import SamplingParams
12 +from vllm.v1.sample.logits_processor import (LOGITSPROCS_GROUP, BatchUpdate,
13 + LogitsProcessor,
14 + MoveDirectionality)
15 +
16 +MODEL_NAME = "facebook/opt-125m"
17 +POOLING_MODEL_NAME = "BAAI/bge-base-en-v1.5"
18 +DUMMY_LOGITPROC_ARG = "target_token"
19 +TEMP_GREEDY = 0.0
20 +MAX_TOKENS = 20
21 +DUMMY_LOGITPROC_ENTRYPOINT = "dummy_logitproc"
22 +DUMMY_LOGITPROC_MODULE = "DummyModule"
23 +DUMMY_LOGITPROC_FQCN = f"{DUMMY_LOGITPROC_MODULE}:DummyLogitsProcessor"
24 +
25 +
26 +class CustomLogitprocSource(Enum):
27 + """How to source a logitproc for testing purposes"""
28 + LOGITPROC_SOURCE_NONE = auto() # No custom logitproc
29 + LOGITPROC_SOURCE_ENTRYPOINT = auto() # Via entrypoint
30 + LOGITPROC_SOURCE_FQCN = auto() # Via fully-qualified class name (FQCN)
31 + LOGITPROC_SOURCE_CLASS = auto() # Via provided class object
32 +
33 +
34 +# Sample prompts.
35 +prompts = [
36 + "Hello, my name is",
37 + "The president of the United States is",
38 + "The capital of France is",
39 + "The future of AI is",
40 +]
41 +
42 +
43 +class DummyLogitsProcessor(LogitsProcessor):
44 + """Fake logit processor to support unit testing and examples"""
45 +
46 + def __init__(self, vllm_config: "VllmConfig", device: torch.device,
47 + is_pin_memory: bool):
48 + self.req_info: dict[int, SamplingParams] = {}
49 +
50 + def is_argmax_invariant(self) -> bool:
51 + """Never impacts greedy sampling"""
52 + return False
53 +
54 + def update_state(self, batch_update: Optional[BatchUpdate]):
55 + if not batch_update:
56 + return
57 +
58 + # Process added requests.
59 + for index, params, _, _ in batch_update.added:
60 + assert params is not None
61 + if params.extra_args and (target_token :=
62 + params.extra_args.get("target_token")):
63 + self.req_info[index] = target_token
64 +
65 + if self.req_info:
66 + # Process removed requests.
67 + for index in batch_update.removed:
68 + self.req_info.pop(index, None)
69 +
70 + # Process moved requests, unidirectional move (a->b) and swap
71 + # (a<->b)
72 + for adx, bdx, direct in batch_update.moved:
73 + a_val = self.req_info.pop(adx, None)
74 + b_val = self.req_info.pop(bdx, None)
75 + if a_val is not None:
76 + self.req_info[bdx] = a_val
77 + if direct == MoveDirectionality.SWAP and b_val is not None:
78 + self.req_info[adx] = b_val
79 +
80 + def apply(self, logits: torch.Tensor) -> torch.Tensor:
81 + if not self.req_info:
82 + return logits
83 +
84 + # Save target values before modification
85 + rows_list = list(self.req_info.keys())
86 + cols = torch.tensor([self.req_info[i] for i in rows_list],
87 + dtype=torch.long,
88 + device=logits.device)
89 + rows = torch.tensor(rows_list, dtype=torch.long, device=logits.device)
90 + values_to_keep = logits[rows, cols].clone()
91 +
92 + # Mask all but target tokens
93 + logits[rows] = float('-inf')
94 + logits[rows, cols] = values_to_keep
95 +
96 + return logits
97 +
98 +
99 +"""Dummy module with dummy logitproc class"""
100 +dummy_module = types.ModuleType(DUMMY_LOGITPROC_MODULE)
101 +dummy_module.DummyLogitsProcessor = DummyLogitsProcessor # type: ignore
102 +
103 +
104 +class EntryPoint:
105 + """Dummy entrypoint class for logitsprocs testing"""
106 +
107 + def __init__(self):
108 + self.name = DUMMY_LOGITPROC_ENTRYPOINT
109 + self.value = DUMMY_LOGITPROC_FQCN
110 +
111 + def load(self):
112 + return DummyLogitsProcessor
113 +
114 +
115 +class EntryPoints(list):
116 + """Dummy EntryPoints class for logitsprocs testing"""
117 +
118 + def __init__(self, group: str):
119 + # Emulate list-like functionality
120 + eps = [EntryPoint()] if group == LOGITSPROCS_GROUP else []
121 + super().__init__(eps)
122 + # Extra attributes
123 + self.names = [ep.name for ep in eps]
124 +
125 +
126 +"""Fake version of importlib.metadata.entry_points"""
127 +entry_points = lambda group: EntryPoints(group)
tests/v1/sample/test_rejection_sampler.py
+2 -2
@@ -7,7 +7,7 @@ import torch
7 import torch.nn.functional as F
8
9 from vllm.platforms import current_platform
10 -from vllm.v1.sample.logits_processor import LogitsProcessorManager
10 +from vllm.v1.sample.logits_processor import LogitsProcessors
11 from vllm.v1.sample.metadata import SamplingMetadata
12 from vllm.v1.sample.rejection_sampler import (PLACEHOLDER_TOKEN_ID,
13 RejectionSampler)
@@ -69,7 +69,7 @@ def create_sampling_metadata(
69 output_token_ids=[],
70 allowed_token_ids_mask=None,
71 bad_words_token_ids={},
72 - logitsprocs=LogitsProcessorManager(),
72 + logitsprocs=LogitsProcessors(),
73 )
74
75
tests/v1/sample/test_sampler.py
+2 -2
@@ -9,7 +9,7 @@ import torch
9
10 from vllm.platforms import current_platform
11 from vllm.utils import is_pin_memory_available, make_tensor_with_pad
12 -from vllm.v1.sample.logits_processor import LogitsProcessorManager
12 +from vllm.v1.sample.logits_processor import LogitsProcessors
13 from vllm.v1.sample.metadata import SamplingMetadata
14 from vllm.v1.sample.sampler import Sampler
15
@@ -173,7 +173,7 @@ def _create_default_sampling_metadata(
173 no_penalties=True,
174 allowed_token_ids_mask=None,
175 bad_words_token_ids={},
176 - logitsprocs=LogitsProcessorManager(),
176 + logitsprocs=LogitsProcessors(),
177 )
178 return fake_sampling_metadata
179
tests/v1/worker/test_gpu_input_batch.py
+2 -2
@@ -13,7 +13,7 @@ from vllm.platforms import current_platform
13 from vllm.sampling_params import SamplingParams
14 from vllm.utils import is_pin_memory_available, make_tensor_with_pad
15 from vllm.v1.pool.metadata import PoolingMetadata
16 -from vllm.v1.sample.logits_processor import LogitsProcessorManager
16 +from vllm.v1.sample.logits_processor import LogitsProcessors
17 from vllm.v1.sample.metadata import SamplingMetadata
18 from vllm.v1.worker.block_table import BlockTable, MultiGroupBlockTable
19 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -169,7 +169,7 @@ def _construct_expected_sampling_metadata(
169 and all(x == 1 for x in repetition_penalties)),
170 allowed_token_ids_mask=allowed_token_ids_mask,
171 bad_words_token_ids=bad_words_token_ids,
172 - logitsprocs=LogitsProcessorManager(),
172 + logitsprocs=LogitsProcessors(),
173 )
174
175
vllm/config/__init__.py
+5
@@ -62,6 +62,7 @@ if TYPE_CHECKING:
62 QuantizationConfig)
63 from vllm.model_executor.model_loader import LoadFormats
64 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
65 + from vllm.v1.sample.logits_processor import LogitsProcessor
66
67 HfOverrides = Union[dict, Callable[[type], type]]
68 else:
@@ -72,6 +73,7 @@ else:
73 BaseModelLoader = Any
74 LoadFormats = Any
75 TensorizerConfig = Any
76 + LogitsProcessor = Any
77 HfOverrides = Union[dict[str, Any], Callable[[type], type]]
78
79 me_quant = LazyLoader("model_executor", globals(),
@@ -465,6 +467,9 @@ class ModelConfig:
467 - "transformers" will use the Transformers model implementation."""
468 override_attention_dtype: Optional[str] = None
469 """Override dtype for attention"""
470 + logits_processors: Optional[list[Union[str, type[LogitsProcessor]]]] = None
471 + """One or more logits processors' fully-qualified class names or class
472 + definitions"""
473
474 def compute_hash(self) -> str:
475 """
vllm/engine/arg_utils.py
+8
@@ -43,6 +43,7 @@ from vllm.transformers_utils.config import is_interleaved
43 from vllm.transformers_utils.utils import check_gguf_file
44 from vllm.utils import (STR_DUAL_CHUNK_FLASH_ATTN_VAL, FlexibleArgumentParser,
45 GiB_bytes, get_ip, is_in_ray_actor)
46 +from vllm.v1.sample.logits_processor import LogitsProcessor
47
48 # yapf: enable
49
@@ -435,6 +436,10 @@ class EngineArgs:
436 enable_multimodal_encoder_data_parallel: bool = \
437 ParallelConfig.enable_multimodal_encoder_data_parallel
438
439 + logits_processors: Optional[list[Union[
440 + str, type[LogitsProcessor]]]] = ModelConfig.logits_processors
441 + """Custom logitproc types"""
442 +
443 async_scheduling: bool = SchedulerConfig.async_scheduling
444 # DEPRECATED
445 enable_prompt_adapter: bool = False
@@ -549,6 +554,8 @@ class EngineArgs:
554 **model_kwargs["model_impl"])
555 model_group.add_argument("--override-attention-dtype",
556 **model_kwargs["override_attention_dtype"])
557 + model_group.add_argument("--logits-processors",
558 + **model_kwargs["logits_processors"])
559
560 # Model loading arguments
561 load_kwargs = get_kwargs(LoadConfig)
@@ -940,6 +947,7 @@ class EngineArgs:
947 enable_sleep_mode=self.enable_sleep_mode,
948 model_impl=self.model_impl,
949 override_attention_dtype=self.override_attention_dtype,
950 + logits_processors=self.logits_processors,
951 )
952
953 def validate_tensorizer_args(self):
vllm/entrypoints/llm.py
+4
@@ -55,6 +55,7 @@ from vllm.transformers_utils.tokenizer import (AnyTokenizer, MistralTokenizer,
55 get_cached_tokenizer)
56 from vllm.usage.usage_lib import UsageContext
57 from vllm.utils import Counter, Device, deprecate_kwargs, is_list_of
58 +from vllm.v1.sample.logits_processor import LogitsProcessor
59
60 if TYPE_CHECKING:
61 from vllm.v1.metrics.reader import Metric
@@ -198,6 +199,8 @@ class LLM:
199 override_pooler_config: Optional[PoolerConfig] = None,
200 compilation_config: Optional[Union[int, dict[str, Any],
201 CompilationConfig]] = None,
202 + logits_processors: Optional[list[Union[str,
203 + type[LogitsProcessor]]]] = None,
204 **kwargs,
205 ) -> None:
206 """LLM constructor."""
@@ -272,6 +275,7 @@ class LLM:
275 mm_processor_kwargs=mm_processor_kwargs,
276 override_pooler_config=override_pooler_config,
277 compilation_config=compilation_config_instance,
278 + logits_processors=logits_processors,
279 **kwargs,
280 )
281
vllm/utils/__init__.py
+1 -1
@@ -2562,7 +2562,7 @@ def direct_register_custom_op(
2562
2563 def resolve_obj_by_qualname(qualname: str) -> Any:
2564 """
2565 - Resolve an object by its fully qualified name.
2565 + Resolve an object by its fully-qualified class name.
2566 """
2567 module_name, obj_name = qualname.rsplit(".", 1)
2568 module = importlib.import_module(module_name)
vllm/v1/sample/logits_processor/__init__.py new
+185
@@ -0,0 +1,185 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import importlib
4 +import itertools
5 +from collections.abc import Sequence
6 +from typing import TYPE_CHECKING, Optional, Union
7 +
8 +import torch
9 +
10 +from vllm.logger import init_logger
11 +from vllm.v1.sample.logits_processor.builtin import (LogitBiasLogitsProcessor,
12 + MinPLogitsProcessor,
13 + MinTokensLogitsProcessor)
14 +from vllm.v1.sample.logits_processor.interface import (BatchUpdate,
15 + LogitsProcessor,
16 + MoveDirectionality)
17 +from vllm.v1.sample.logits_processor.state import (BatchUpdateBuilder,
18 + LogitsProcessors)
19 +
20 +if TYPE_CHECKING:
21 + from vllm.config import VllmConfig
22 +
23 +logger = init_logger(__name__)
24 +
25 +# Error message when the user tries to initialize vLLM with a pooling model
26 +# and custom logitsproces
27 +STR_POOLING_REJECTS_LOGITSPROCS = ("Pooling models do not support custom"
28 + " logits processors.")
29 +
30 +LOGITSPROCS_GROUP = 'vllm.logits_processors'
31 +
32 +BUILTIN_LOGITS_PROCESSORS: list[type[LogitsProcessor]] = [
33 + MinTokensLogitsProcessor,
34 + LogitBiasLogitsProcessor,
35 + MinPLogitsProcessor,
36 +]
37 +
38 +
39 +def _load_logitsprocs_plugins() -> list[type[LogitsProcessor]]:
40 + """Load all installed logit processor plugins"""
41 +
42 + import sys
43 +
44 + if sys.version_info < (3, 10):
45 + from importlib_metadata import entry_points
46 + else:
47 + from importlib.metadata import entry_points
48 +
49 + installed_logitsprocs_plugins = entry_points(group=LOGITSPROCS_GROUP)
50 + if len(installed_logitsprocs_plugins) == 0:
51 + logger.debug("No logitsprocs plugins installed (group %s).",
52 + LOGITSPROCS_GROUP)
53 + return []
54 +
55 + # Load logitsprocs plugins
56 + logger.debug("Loading installed logitsprocs plugins (group %s):",
57 + LOGITSPROCS_GROUP)
58 + classes: list[type[LogitsProcessor]] = []
59 + for entrypoint in installed_logitsprocs_plugins:
60 + try:
61 + logger.debug("- Loading logitproc plugin entrypoint=%s target=%s",
62 + entrypoint.name, entrypoint.value)
63 + classes.append(entrypoint.load())
64 + except Exception as e:
65 + raise RuntimeError(
66 + f"Failed to load LogitsProcessor plugin {entrypoint}") from e
67 + return classes
68 +
69 +
70 +def _load_logitsprocs_by_fqcns(
71 + logits_processors: Optional[Sequence[Union[str, type[LogitsProcessor]]]]
72 +) -> list[type[LogitsProcessor]]:
73 + """Load logit processor types, identifying them by fully-qualified class
74 + names (FQCNs).
75 +
76 + Effectively, a mixed list of logitproc types and FQCN strings is converted
77 + into a list of entirely logitproc types, by loading from the FQCNs.
78 +
79 + FQCN syntax is <module>:<type> i.e. x.y.z:CustomLogitProc
80 +
81 + Already-loaded logitproc types must be subclasses of LogitsProcessor
82 +
83 + Args:
84 + logits_processors: Potentially mixed list of logitsprocs types and FQCN
85 + strings for logitproc types
86 +
87 + Returns:
88 + List of logitproc types
89 +
90 + """
91 + if not logits_processors:
92 + return []
93 +
94 + logger.debug(
95 + "%s additional custom logits processors specified, checking whether "
96 + "they need to be loaded.", len(logits_processors))
97 +
98 + classes: list[type[LogitsProcessor]] = []
99 + for ldx, logitproc in enumerate(logits_processors):
100 + if isinstance(logitproc, type):
101 + logger.debug(" - Already-loaded logit processor: %s",
102 + logitproc.__name__)
103 + if not issubclass(logitproc, LogitsProcessor):
104 + raise ValueError(
105 + f"{logitproc.__name__} is not a subclass of LogitsProcessor"
106 + )
107 + classes.append(logitproc)
108 + continue
109 +
110 + logger.debug("- Loading logits processor %s", logitproc)
111 + module_path, qualname = logitproc.split(":")
112 +
113 + try:
114 + # Load module
115 + module = importlib.import_module(module_path)
116 + except Exception as e:
117 + raise RuntimeError(
118 + f"Failed to load {ldx}th LogitsProcessor plugin {logitproc}"
119 + ) from e
120 +
121 + # Walk down dotted name to get logitproc class
122 + obj = module
123 + for attr in qualname.split("."):
124 + obj = getattr(obj, attr)
125 + if not isinstance(obj, type):
126 + raise ValueError("Loaded logit processor must be a type.")
127 + if not issubclass(obj, LogitsProcessor):
128 + raise ValueError(
129 + f"{obj.__name__} must be a subclass of LogitsProcessor")
130 + classes.append(obj)
131 +
132 + return classes
133 +
134 +
135 +def _load_custom_logitsprocs(
136 + logits_processors: Optional[Sequence[Union[str, type[LogitsProcessor]]]],
137 +) -> list[type[LogitsProcessor]]:
138 + """Load all custom logits processors.
139 +
140 + * First load all installed logitproc plugins
141 + * Second load custom logitsprocs pass by the user at initialization time
142 +
143 + Args:
144 + logits_processors: potentially mixed list of logitproc types and
145 + logitproc type fully-qualified names (FQCNs)
146 + which need to be loaded
147 +
148 + Returns:
149 + A list of all loaded logitproc types
150 + """
151 + from vllm.platforms import current_platform
152 + if current_platform.is_tpu():
153 + # No logitsprocs specified by caller
154 + # TODO(andy) - vLLM V1 on TPU does not support custom logitsprocs
155 + return []
156 +
157 + return (_load_logitsprocs_plugins() +
158 + _load_logitsprocs_by_fqcns(logits_processors))
159 +
160 +
161 +def build_logitsprocs(
162 + vllm_config: "VllmConfig",
163 + device: torch.device,
164 + is_pin_memory: bool,
165 + is_pooling_model: bool,
166 + custom_logitsprocs: Sequence[Union[str, type[LogitsProcessor]]] = (),
167 +) -> LogitsProcessors:
168 + if is_pooling_model:
169 + if custom_logitsprocs:
170 + raise ValueError(STR_POOLING_REJECTS_LOGITSPROCS)
171 + logger.debug("Skipping logits processor loading because pooling models"
172 + " do not support logits processors.")
173 + return LogitsProcessors()
174 + custom_logitsprocs_classes = _load_custom_logitsprocs(custom_logitsprocs)
175 + return LogitsProcessors(
176 + ctor(vllm_config, device, is_pin_memory) for ctor in itertools.chain(
177 + BUILTIN_LOGITS_PROCESSORS, custom_logitsprocs_classes))
178 +
179 +
180 +__all__ = [
181 + "LogitsProcessor", "LogitBiasLogitsProcessor", "MinPLogitsProcessor",
182 + "MinTokensLogitsProcessor", "BatchUpdate", "BatchUpdateBuilder",
183 + "MoveDirectionality", "LogitsProcessors", "build_logitsprocs",
184 + "STR_POOLING_REJECTS_LOGITSPROCS", "LOGITSPROCS_GROUP"
185 +]
vllm/v1/sample/logits_processor/builtin.py renamed
+26 -270
@@ -1,241 +1,32 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import dataclasses
4 -from abc import ABC, abstractmethod
5 -from collections.abc import Iterator, Sequence
6 -from dataclasses import dataclass, field
7 -from enum import Enum
8 -from itertools import chain
9 -from typing import Optional, Union
3 +from collections.abc import Sequence
4 +from typing import TYPE_CHECKING, Optional
5
6 import torch
12 -from torch._prims_common import DeviceLikeType
13 -
14 -from vllm import PoolingParams, SamplingParams
15 -from vllm.logger import init_logger
16 -
17 -logger = init_logger(__name__)
18 -
19 -
20 -class MoveDirectionality(Enum):
21 - # One-way i1->i2 req move within batch
22 - UNIDIRECTIONAL = 0
23 - # Two-way i1<->i2 req swap within batch
24 - SWAP = 1
25 -
26 -
27 -# (index, params, output_tok_ids) tuples for new
28 -# requests added to the batch.
29 -AddedRequest = tuple[int, Union[SamplingParams, PoolingParams], list[int]]
30 -# (index 1, index 2, directionality) tuples representing
31 -# one-way moves or two-way swaps of requests in batch
32 -MovedRequest = tuple[int, int, MoveDirectionality]
33 -# Batch indices of any removed requests.
34 -RemovedRequest = int
35 -
36 -
37 -@dataclasses.dataclass(frozen=True)
38 -class BatchUpdate:
39 - """Persistent batch state change info for logitsprocs"""
40 - batch_size: int # Current num reqs in batch
41 -
42 - # Metadata for requests added to, removed from, and moved
43 - # within the persistent batch.
44 - #
45 - # Note: each added request is represented as
46 - # (index, params, output_tok_ids)
47 - # Key assumption: output_tok_ids is a reference to the
48 - # request's running output tokens list; in this way
49 - # the logits processors always see the latest list of
50 - # generated tokens
51 - removed: Sequence[RemovedRequest]
52 - moved: Sequence[MovedRequest]
53 - added: Sequence[AddedRequest]
54 -
55 -
56 -class BatchUpdateBuilder:
57 - """Helps track persistent batch state changes and build
58 - a batch update data structure for logitsprocs
59 -
60 - Assumptions:
61 - * All information about requests removed from persistent batch
62 - during a step is aggregated in self._removed through calls to
63 - self.removed_append() at the beginning of a step. This must happen
64 - before the first time that self.removed, self.pop_removed()
65 - or self.peek_removed() are invoked in a given step
66 - * After the first time that self.removed, self.pop_removed()
67 - or self.peek_removed() are read in a step, no new removals
68 - are registered using self.removed_append()
69 - * Elements of self._removed are never directly modified, added or
70 - removed (i.e. modification is only via self.removed_append() and
71 - self.pop_removed())
72 -
73 - Guarantees under above assumptions:
74 - * self.removed is always sorted in descending order
75 - * self.pop_removed() and self.peek_removed() both return
76 - the lowest removed request index in the current step
77 - """
78 -
79 - _removed: list[RemovedRequest]
80 - _is_removed_sorted: bool
81 - moved: list[MovedRequest]
82 - added: list[AddedRequest]
83 -
84 - def __init__(
85 - self,
86 - removed: Optional[list[RemovedRequest]] = None,
87 - moved: Optional[list[MovedRequest]] = None,
88 - added: Optional[list[AddedRequest]] = None,
89 - ) -> None:
90 - self._removed = removed or []
91 - self.moved = moved or []
92 - self.added = added or []
93 - self._is_removed_sorted = False
94 -
95 - def _ensure_removed_sorted(self) -> None:
96 - """Sort removed request indices in
97 - descending order.
98 -
99 - Idempotent after first call in a
100 - given step, until reset.
101 - """
102 - if not self._is_removed_sorted:
103 - self._removed.sort(reverse=True)
104 - self._is_removed_sorted = True
105 -
106 - @property
107 - def removed(self) -> list[RemovedRequest]:
108 - """Removed request indices sorted in
109 - descending order"""
110 - self._ensure_removed_sorted()
111 - return self._removed
112 -
113 - def removed_append(self, index: int) -> None:
114 - """Register the removal of a request from
115 - the persistent batch.
116 -
117 - Must not be called after the first time
118 - self.removed, self.pop_removed() or
119 - self.peek_removed() are invoked.
120 -
121 - Args:
122 - index: request index
123 - """
124 - if self._is_removed_sorted:
125 - raise RuntimeError("Cannot register new removed request after"
126 - " self.removed has been read.")
127 - self._removed.append(index)
128 -
129 - def has_removed(self) -> bool:
130 - return bool(self._removed)
131 -
132 - def peek_removed(self) -> Optional[int]:
133 - """Return lowest removed request index"""
134 - if self.has_removed():
135 - self._ensure_removed_sorted()
136 - return self._removed[-1]
137 - return None
138 -
139 - def pop_removed(self) -> Optional[int]:
140 - """Pop lowest removed request index"""
141 - if self.has_removed():
142 - self._ensure_removed_sorted()
143 - return self._removed.pop()
144 - return None
145 -
146 - def get_and_reset(self, batch_size: int) -> Optional[BatchUpdate]:
147 - """Generate a logitsprocs batch update data structure
148 - and reset internal batch update builder state.
149 -
150 - Args:
151 - batch_size: current persistent batch size
152 -
153 - Returns:
154 - Frozen logitsprocs batch update instance; `None` if no updates
155 - """
156 - # Reset removal-sorting logic
157 - self._is_removed_sorted = False
158 - if not any((self._removed, self.moved, self.added)):
159 - # No update; short-circuit
160 - return None
161 - # Build batch state update
162 - batch_update = BatchUpdate(
163 - batch_size=batch_size,
164 - removed=self._removed,
165 - moved=self.moved,
166 - added=self.added,
167 - )
168 - # Reset removed/moved/added update lists
169 - self._removed = []
170 - self.moved = []
171 - self.added = []
172 - return batch_update
173 -
174 -
175 -class LogitsProcessor(ABC):
176 -
177 - @abstractmethod
178 - def apply(self, logits: torch.Tensor) -> torch.Tensor:
179 - raise NotImplementedError
7
181 - @abstractmethod
182 - def is_argmax_invariant(self) -> bool:
183 - """True if logits processor has no impact on the
184 - argmax computation in greedy sampling.
185 - NOTE: may or may not have the same value for all
186 - instances of a given LogitsProcessor subclass,
187 - depending on subclass implementation.
188 - TODO(andy): won't be utilized until logits
189 - processors are user-extensible
190 - """
191 - raise NotImplementedError
192 -
193 - @abstractmethod
194 - def update_state(
195 - self,
196 - batch_update: Optional[BatchUpdate],
197 - ) -> None:
198 - """Called when there are new output tokens, prior
199 - to each forward pass.
200 -
201 - Args:
202 - batch_update is non-None iff there have been
203 - changes to the batch makeup.
204 - """
205 - raise NotImplementedError
206 -
207 -
208 -@dataclass
209 -class LogitsProcessorManager:
210 - """Encapsulates initialized logitsproc objects."""
211 - argmax_invariant: list[LogitsProcessor] = field(
212 - default_factory=list) # argmax-invariant logitsprocs
213 - non_argmax_invariant: list[LogitsProcessor] = field(
214 - default_factory=list) # non-argmax-invariant logitsprocs
215 -
216 - @property
217 - def all(self) -> Iterator[LogitsProcessor]:
218 - """Iterator over all logits processors."""
219 - return chain(self.argmax_invariant, self.non_argmax_invariant)
220 -
221 -
222 -###### ----- Built-in LogitsProcessor impls below here
8 +from vllm.v1.sample.logits_processor.interface import (BatchUpdate,
9 + LogitsProcessor,
10 + MoveDirectionality)
11 +
12 +if TYPE_CHECKING:
13 + from vllm.config import VllmConfig
14
15
16 class MinPLogitsProcessor(LogitsProcessor):
17
227 - def __init__(self, max_num_reqs: int, pin_memory: bool,
228 - device: DeviceLikeType):
229 - super().__init__()
18 + def __init__(self, vllm_config: "VllmConfig", device: torch.device,
19 + is_pin_memory: bool):
20 + max_num_reqs = vllm_config.scheduler_config.max_num_seqs
21 self.min_p_count: int = 0
22
23 self.min_p_cpu_tensor = torch.zeros((max_num_reqs, ),
24 dtype=torch.float32,
25 device="cpu",
235 - pin_memory=pin_memory)
26 + pin_memory=is_pin_memory)
27 self.min_p_cpu = self.min_p_cpu_tensor.numpy()
28
238 - self.use_double_tensor = torch.device("cpu") != torch.device(device)
29 + self.use_double_tensor = torch.device(device).type != "cpu"
30
31 if self.use_double_tensor:
32 # Pre-allocated device tensor
@@ -260,8 +51,8 @@ class MinPLogitsProcessor(LogitsProcessor):
51
52 needs_update = False
53 # Process added requests.
263 - for index, params, _ in batch_update.added:
264 - min_p = params.min_p if isinstance(params, SamplingParams) else 0.0
54 + for index, params, _, _ in batch_update.added:
55 + min_p = params.min_p
56 if self.min_p_cpu[index] != min_p:
57 needs_update = True
58 self.min_p_cpu[index] = min_p
@@ -316,11 +107,10 @@ class MinPLogitsProcessor(LogitsProcessor):
107
108 class LogitBiasLogitsProcessor(LogitsProcessor):
109
319 - def __init__(self, pin_memory: bool, device: torch.device):
320 - super().__init__()
321 - self.biases: dict[int, dict[int, float]] = {}
110 + def __init__(self, _, device: torch.device, is_pin_memory: bool):
111 self.device = device
323 - self.pin_memory = pin_memory
112 + self.pin_memory = is_pin_memory
113 + self.biases: dict[int, dict[int, float]] = {}
114
115 self.bias_tensor: torch.Tensor = torch.tensor(())
116 self.logits_slice = (self._device_tensor([], torch.int32),
@@ -337,9 +127,8 @@ class LogitBiasLogitsProcessor(LogitsProcessor):
127
128 needs_update: bool = False
129 # Process added requests.
340 - for index, params, _ in batch_update.added:
341 - if isinstance(params, SamplingParams) and (lb :=
342 - params.logit_bias):
130 + for index, params, _, _ in batch_update.added:
131 + if lb := params.logit_bias:
132 self.biases[index] = lb
133 needs_update = True
134 else:
@@ -400,12 +189,12 @@ class LogitBiasLogitsProcessor(LogitsProcessor):
189
190 class MinTokensLogitsProcessor(LogitsProcessor):
191
403 - def __init__(self, pin_memory: bool, device: torch.device):
192 + def __init__(self, vllm_config: "VllmConfig", device: torch.device,
193 + is_pin_memory: bool):
194 # index -> (min_toks, output_token_ids, stop_token_ids)
405 - super().__init__()
406 - self.min_toks: dict[int, tuple[int, Sequence[int], set[int]]] = {}
195 self.device = device
408 - self.pin_memory = pin_memory
196 + self.pin_memory = is_pin_memory
197 + self.min_toks: dict[int, tuple[int, Sequence[int], set[int]]] = {}
198
199 # (req_idx_tensor,eos_tok_id_tensor)
200 self.logits_slice: tuple[torch.Tensor,
@@ -424,9 +213,8 @@ class MinTokensLogitsProcessor(LogitsProcessor):
213
214 if batch_update:
215 # Process added requests.
427 - for index, params, output_tok_ids in batch_update.added:
428 - if (isinstance(params, SamplingParams)
429 - and (min_tokens := params.min_tokens)
216 + for index, params, _, output_tok_ids in batch_update.added:
217 + if ((min_tokens := params.min_tokens)
218 and len(output_tok_ids) < min_tokens):
219 # Replace request metadata at batch index
220 self.min_toks[index] = (min_tokens, output_tok_ids,
@@ -499,35 +287,3 @@ class MinTokensLogitsProcessor(LogitsProcessor):
287 # Inhibit EOS token for requests which have not reached min length
288 logits[self.logits_slice] = -float("inf")
289 return logits
502 -
503 -
504 -def init_builtin_logitsprocs(pin_memory_available: bool, max_num_reqs: int,
505 - device: torch.device) -> LogitsProcessorManager:
506 - """Construct 'builtin' vLLM logitsprocs which the engine
507 - loads by default.
508 -
509 - Args:
510 - pin_memory_available: pinned memory is available for use
511 - for use by logitsproc
512 - max_num_reqs: ceiling on request count in persistent batch
513 - device: inference device
514 -
515 - Returns:
516 - Data structure encapsulating loaded logitsprocs
517 - """
518 - min_tokens_logitproc = MinTokensLogitsProcessor(
519 - pin_memory=pin_memory_available, device=device)
520 - logit_bias_logitproc = LogitBiasLogitsProcessor(
521 - pin_memory=pin_memory_available, device=device)
522 - min_p_logitproc = MinPLogitsProcessor(
523 - pin_memory=pin_memory_available,
524 - device=device,
525 - # +1 for temporary swap space
526 - max_num_reqs=max_num_reqs + 1)
527 - return LogitsProcessorManager(
528 - non_argmax_invariant=[
529 - min_tokens_logitproc,
530 - logit_bias_logitproc,
531 - ],
532 - argmax_invariant=[min_p_logitproc],
533 - )
vllm/v1/sample/logits_processor/interface.py new
+86
@@ -0,0 +1,86 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from collections.abc import Sequence
5 +from dataclasses import dataclass
6 +from enum import Enum, auto
7 +from typing import TYPE_CHECKING, Optional
8 +
9 +import torch
10 +
11 +from vllm import SamplingParams
12 +
13 +if TYPE_CHECKING:
14 + from vllm.config import VllmConfig
15 +
16 +
17 +class MoveDirectionality(Enum):
18 + # One-way i1->i2 req move within batch
19 + UNIDIRECTIONAL = auto()
20 + # Two-way i1<->i2 req swap within batch
21 + SWAP = auto()
22 +
23 +
24 +# (index, params, prompt_tok_ids, output_tok_ids) tuples for new
25 +# requests added to the batch.
26 +AddedRequest = tuple[int, SamplingParams, list[int], list[int]]
27 +
28 +# (index 1, index 2, directionality) tuples representing
29 +# one-way moves or two-way swaps of requests in batch
30 +MovedRequest = tuple[int, int, MoveDirectionality]
31 +
32 +# Batch indices of any removed requests.
33 +RemovedRequest = int
34 +
35 +
36 +@dataclass(frozen=True)
37 +class BatchUpdate:
38 + """Persistent batch state change info for logitsprocs"""
39 + batch_size: int # Current num reqs in batch
40 +
41 + # Metadata for requests added to, removed from, and moved
42 + # within the persistent batch.
43 + #
44 + # Key assumption: the `output_tok_ids` list (which is an element of each
45 + # tuple in `added`) is a reference to the request's running output tokens
46 + # list; via this reference, the logits processors always see the latest
47 + # list of generated output tokens
48 + removed: Sequence[RemovedRequest]
49 + moved: Sequence[MovedRequest]
50 + added: Sequence[AddedRequest]
51 +
52 +
53 +class LogitsProcessor(ABC):
54 +
55 + @abstractmethod
56 + def __init__(self, vllm_config: "VllmConfig", device: torch.device,
57 + is_pin_memory: bool) -> None:
58 + raise NotImplementedError
59 +
60 + @abstractmethod
61 + def apply(self, logits: torch.Tensor) -> torch.Tensor:
62 + raise NotImplementedError
63 +
64 + @abstractmethod
65 + def is_argmax_invariant(self) -> bool:
66 + """True if logits processor has no impact on the
67 + argmax computation in greedy sampling.
68 + NOTE: may or may not have the same value for all
69 + instances of a given LogitsProcessor subclass,
70 + depending on subclass implementation.
71 + """
72 + raise NotImplementedError
73 +
74 + @abstractmethod
75 + def update_state(
76 + self,
77 + batch_update: Optional["BatchUpdate"],
78 + ) -> None:
79 + """Called when there are new output tokens, prior
80 + to each forward pass.
81 +
82 + Args:
83 + batch_update is non-None iff there have been
84 + changes to the batch makeup.
85 + """
86 + raise NotImplementedError
vllm/v1/sample/logits_processor/state.py new
+149
@@ -0,0 +1,149 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Iterator
4 +from itertools import chain
5 +from typing import TYPE_CHECKING, Optional
6 +
7 +from vllm.v1.sample.logits_processor.interface import (AddedRequest,
8 + BatchUpdate,
9 + MovedRequest,
10 + RemovedRequest)
11 +
12 +if TYPE_CHECKING:
13 + from vllm.v1.sample.logits_processor.interface import LogitsProcessor
14 +
15 +
16 +class BatchUpdateBuilder:
17 + """Helps track persistent batch state changes and build
18 + a batch update data structure for logitsprocs
19 + Assumptions:
20 + * All information about requests removed from persistent batch
21 + during a step is aggregated in self._removed through calls to
22 + self.removed_append() at the beginning of a step. This must happen
23 + before the first time that self.removed, self.pop_removed()
24 + or self.peek_removed() are invoked in a given step
25 + * After the first time that self.removed, self.pop_removed()
26 + or self.peek_removed() are read in a step, no new removals
27 + are registered using self.removed_append()
28 + * Elements of self._removed are never directly modified, added or
29 + removed (i.e. modification is only via self.removed_append() and
30 + self.pop_removed())
31 + Guarantees under above assumptions:
32 + * self.removed is always sorted in descending order
33 + * self.pop_removed() and self.peek_removed() both return
34 + the lowest removed request index in the current step
35 + """
36 +
37 + _removed: list[RemovedRequest]
38 + _is_removed_sorted: bool
39 + moved: list[MovedRequest]
40 + added: list[AddedRequest]
41 +
42 + def __init__(
43 + self,
44 + removed: Optional[list[RemovedRequest]] = None,
45 + moved: Optional[list[MovedRequest]] = None,
46 + added: Optional[list[AddedRequest]] = None,
47 + ) -> None:
48 + self._removed = removed or []
49 + self.moved = moved or []
50 + self.added = added or []
51 + self._is_removed_sorted = False
52 +
53 + def _ensure_removed_sorted(self) -> None:
54 + """Sort removed request indices in
55 + descending order.
56 + Idempotent after first call in a
57 + given step, until reset.
58 + """
59 + if not self._is_removed_sorted:
60 + self._removed.sort(reverse=True)
61 + self._is_removed_sorted = True
62 +
63 + @property
64 + def removed(self) -> list[RemovedRequest]:
65 + """Removed request indices sorted in
66 + descending order"""
67 + self._ensure_removed_sorted()
68 + return self._removed
69 +
70 + def removed_append(self, index: int) -> None:
71 + """Register the removal of a request from the persistent batch.
72 +
73 + Must not be called after the first time self.removed,
74 + self.pop_removed() or self.peek_removed() are invoked.
75 +
76 + Args:
77 + index: request index
78 + """
79 + if self._is_removed_sorted:
80 + raise RuntimeError("Cannot register new removed request after"
81 + " self.removed has been read.")
82 + self._removed.append(index)
83 +
84 + def has_removed(self) -> bool:
85 + return bool(self._removed)
86 +
87 + def peek_removed(self) -> Optional[int]:
88 + """Return lowest removed request index"""
89 + if self.has_removed():
90 + self._ensure_removed_sorted()
91 + return self._removed[-1]
92 + return None
93 +
94 + def pop_removed(self) -> Optional[int]:
95 + """Pop lowest removed request index"""
96 + if self.has_removed():
97 + self._ensure_removed_sorted()
98 + return self._removed.pop()
99 + return None
100 +
101 + def _is_update(self) -> bool:
102 + """True if there is a batch state change"""
103 + return any((self._removed, self.moved, self.added))
104 +
105 + def get_and_reset(self, batch_size: int) -> Optional[BatchUpdate]:
106 + """Generate a logitsprocs batch update data structure and reset
107 + internal batch update builder state.
108 +
109 + Args:
110 + batch_size: current persistent batch size
111 +
112 + Returns:
113 + Frozen logitsprocs batch update instance; `None` if no updates
114 + """
115 + # Reset removal-sorting logic
116 + self._is_removed_sorted = False
117 + if not self._is_update():
118 + # No update; short-circuit
119 + return None
120 + # Build batch state update
121 + batch_update = BatchUpdate(
122 + batch_size=batch_size,
123 + removed=self._removed,
124 + moved=self.moved,
125 + added=self.added,
126 + )
127 + self._removed = []
128 + self.moved = []
129 + self.added = []
130 + return batch_update
131 +
132 +
133 +class LogitsProcessors:
134 + """Encapsulates initialized logitsproc objects."""
135 +
136 + def __init__(
137 + self,
138 + logitsprocs: Optional[Iterator["LogitsProcessor"]] = None) -> None:
139 + self.argmax_invariant: list[LogitsProcessor] = []
140 + self.non_argmax_invariant: list[LogitsProcessor] = []
141 + if logitsprocs:
142 + for logitproc in logitsprocs:
143 + (self.argmax_invariant if logitproc.is_argmax_invariant() else
144 + self.non_argmax_invariant).append(logitproc)
145 +
146 + @property
147 + def all(self) -> Iterator["LogitsProcessor"]:
148 + """Iterator over all logits processors."""
149 + return chain(self.argmax_invariant, self.non_argmax_invariant)
vllm/v1/sample/metadata.py
+2 -2
@@ -6,7 +6,7 @@ from typing import Optional
6
7 import torch
8
9 -from vllm.v1.sample.logits_processor import LogitsProcessorManager
9 +from vllm.v1.sample.logits_processor import LogitsProcessors
10
11
12 @dataclass
@@ -40,4 +40,4 @@ class SamplingMetadata:
40 bad_words_token_ids: dict[int, list[list[int]]]
41
42 # Loaded logits processors
43 - logitsprocs: LogitsProcessorManager
43 + logitsprocs: LogitsProcessors
vllm/v1/worker/gpu_input_batch.py
+58 -33
@@ -18,8 +18,8 @@ from vllm.utils import swap_dict_values
18 from vllm.v1.outputs import LogprobsTensors
19 from vllm.v1.pool.metadata import PoolingMetadata
20 from vllm.v1.sample.logits_processor import (BatchUpdateBuilder,
21 - MoveDirectionality,
22 - init_builtin_logitsprocs)
21 + LogitsProcessors,
22 + MoveDirectionality)
23 from vllm.v1.sample.metadata import SamplingMetadata
24 from vllm.v1.spec_decode.utils import is_spec_decode_unsupported
25 from vllm.v1.utils import copy_slice
@@ -78,8 +78,11 @@ class InputBatch:
78 pin_memory: bool,
79 vocab_size: int,
80 block_sizes: list[int], # The block_size of each kv cache group
81 + logitsprocs: Optional[LogitsProcessors] = None,
82 is_spec_decode: bool = False,
83 + is_pooling_model: bool = False,
84 ):
85 + self.is_pooling_model = is_pooling_model
86 self.is_spec_decode = is_spec_decode
87 self.max_num_reqs = max_num_reqs
88 self.max_model_len = max_model_len
@@ -221,14 +224,6 @@ class InputBatch:
224 # updates. Should reset each step.
225 self.batch_update_builder = BatchUpdateBuilder()
226
224 - # Define logits processors.
225 - # TODO(andy): logits processor list should be extensible via engine
226 - # constructor argument; for now the list is fixed.
227 - self.logitsprocs = init_builtin_logitsprocs(
228 - pin_memory_available=pin_memory,
229 - max_num_reqs=max_num_reqs + 1,
230 - device=device)
231 -
227 # TODO convert this to LogitsProcessor
228 self.has_allowed_token_ids: set[str] = set()
229 # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
@@ -244,6 +239,10 @@ class InputBatch:
239
240 self.req_output_token_ids: list[Optional[list[int]]] = []
241
242 + # Store provided logitsprocs. If none are provided, initialize empty
243 + # data structure
244 + self.logitsprocs = logitsprocs or LogitsProcessors()
245 +
246 # This is updated each time the batch constituents change.
247 self.sampling_metadata = self._make_sampling_metadata()
248
@@ -255,28 +254,35 @@ class InputBatch:
254 # while performing state updates to the batch.
255 return cast(list[str], self._req_ids)
256
258 - def _get_next_add_index(self) -> int:
259 - if (req_index := self.batch_update_builder.pop_removed()) is not None:
260 - # Fill the empty index.
261 - return req_index
262 - # Append to end
263 - return self.num_reqs
264 -
257 def _register_add_request(self, request: "CachedRequestState") -> int:
266 - """Track add-request operations"""
267 - req_index = self._get_next_add_index()
268 - assert req_index < self.max_num_reqs
269 - params = (request.sampling_params
270 - if request.sampling_params else request.pooling_params)
258 + """Track add-request operations for logits processors.
259 + Not applicable to pooling models.
260 + """
261 +
262 + # Detailed added request metadata is only required for non-pooling
263 + # models, to support logitsprocs
264 + assert request.sampling_params
265 +
266 + # Fill the next empty index if there is one.
267 + if (new_req_index := self.batch_update_builder.pop_removed()) is None:
268 + # Append to end otherwise.
269 + new_req_index = self.num_reqs
270 +
271 + assert new_req_index < self.max_num_reqs
272 self.batch_update_builder.added.append(
272 - (req_index, params, request.output_token_ids))
273 - return req_index
273 + (new_req_index, request.sampling_params, request.prompt_token_ids,
274 + request.output_token_ids))
275 + return new_req_index
276
277 def add_request(
278 self,
279 request: "CachedRequestState",
280 ) -> int:
279 - req_index = self._register_add_request(request)
281 + if not self.is_pooling_model:
282 + # New request index bookkeeping for autoregressive models.
283 + req_index = self._register_add_request(request)
284 + else:
285 + req_index = self.num_reqs
286
287 req_id = request.req_id
288 if req_index == len(self._req_ids):
@@ -411,7 +417,10 @@ class InputBatch:
417 req_index = self.req_id_to_index.pop(req_id, None)
418 if req_index is None:
419 return None
414 - self.batch_update_builder.removed_append(req_index)
420 + if not self.is_pooling_model:
421 + # Autoregressive models require bookkeeping of removed requests to
422 + # support logitsprocs.
423 + self.batch_update_builder.removed_append(req_index)
424 self._req_ids[req_index] = None
425 self.req_output_token_ids[req_index] = None
426
@@ -446,6 +455,8 @@ class InputBatch:
455 return req_index
456
457 def swap_states(self, i1: int, i2: int) -> None:
458 + # For autoregressive models, track detailed request reordering info
459 + # to support logitsprocs
460 self.batch_update_builder.moved.append(
461 (i1, i2, MoveDirectionality.SWAP))
462 old_id_i1 = self._req_ids[i1]
@@ -513,11 +524,18 @@ class InputBatch:
524 swaps: list of (from,to) swap tuples for moved requests
525 empty_req_indices: indices not filled by condensation
526 """
527 + num_reqs = self.num_reqs
528 +
529 + if self.is_pooling_model:
530 + # Will be contiguous in pooling case, just trim the lists.
531 + del self._req_ids[num_reqs:]
532 + del self.req_output_token_ids[num_reqs:]
533 + return
534 +
535 if not (empty_req_indices := self.batch_update_builder.removed):
536 # All removed requests were replaced by added requests, or else no
537 # requests were removed at all. No condense() needed
538 return
520 - num_reqs = self.num_reqs
539 if num_reqs == 0:
540 # The batched states are empty.
541 self._req_ids.clear()
@@ -541,6 +559,8 @@ class InputBatch:
559 # Move active request down into empty request
560 # index.
561 self.batch_update_builder.pop_removed()
562 + # Autoregressive models require detailed tracking of condense
563 + # operations to support logitsprocs
564 self.batch_update_builder.moved.append(
565 (last_req_index, empty_index,
566 MoveDirectionality.UNIDIRECTIONAL))
@@ -596,15 +616,20 @@ class InputBatch:
616 last_req_index -= 1
617
618 # Trim lists to the batch size.
599 - del self._req_ids[self.num_reqs:]
600 - del self.req_output_token_ids[self.num_reqs:]
619 + del self._req_ids[num_reqs:]
620 + del self.req_output_token_ids[num_reqs:]
621
622 def refresh_metadata(self):
603 - """Apply batch updates, reset input batch at end of step
623 + """Apply any batch updates to sampling metadata."""
624
605 - * Apply batch add/remove/permute to logits procs' states
606 - * If batch state is modified, update sampling metadata
607 - """
625 + if self.is_pooling_model:
626 + # Batch changes every step for pooling models.
627 + self.sampling_metadata = self._make_sampling_metadata()
628 + return
629 +
630 + # For non-pooling models - generate and apply logitsprocs update;
631 + # reset batch update tracking.
632 + # Update sampling metadata if batch state is changed.
633 batch_update = self.batch_update_builder.get_and_reset(self.num_reqs)
634 for logit_proc in self.logitsprocs.all:
635 logit_proc.update_state(batch_update)
vllm/v1/worker/gpu_model_runner.py
+9 -2
@@ -68,6 +68,7 @@ from vllm.v1.kv_cache_interface import (AttentionSpec,
68 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
69 ModelRunnerOutput)
70 from vllm.v1.pool.metadata import PoolingMetadata
71 +from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs
72 from vllm.v1.sample.metadata import SamplingMetadata
73 from vllm.v1.sample.rejection_sampler import RejectionSampler
74 from vllm.v1.sample.sampler import Sampler
@@ -80,7 +81,6 @@ from vllm.v1.worker.kv_connector_model_runner_mixin import (
81 KVConnectorModelRunnerMixin, KVConnectorOutput)
82 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
83
83 -from ..sample.logits_processor import LogitsProcessorManager
84 from .utils import (AttentionGroup, MultiModalBudget, bind_kv_cache,
85 gather_mm_placeholders, initialize_kv_cache_for_kv_sharing,
86 sanity_check_mm_encoder_outputs, scatter_mm_placeholders)
@@ -221,6 +221,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
221 vocab_size=self.model_config.get_vocab_size(),
222 block_sizes=[self.cache_config.block_size],
223 is_spec_decode=bool(self.vllm_config.speculative_config),
224 + logitsprocs=build_logitsprocs(
225 + self.vllm_config, self.device, self.pin_memory,
226 + self.is_pooling_model,
227 + self.vllm_config.model_config.logits_processors),
228 + is_pooling_model=self.is_pooling_model,
229 )
230
231 # TODO(woosuk): Provide an option to tune the max cudagraph batch size.
@@ -2447,7 +2452,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2452 output_token_ids=[[] for _ in range(num_reqs)],
2453 allowed_token_ids_mask=None,
2454 bad_words_token_ids={},
2450 - logitsprocs=LogitsProcessorManager(),
2455 + logitsprocs=LogitsProcessors(),
2456 )
2457 try:
2458 sampler_output = self.sampler(logits=logits,
@@ -2968,6 +2973,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2973 vocab_size=self.model_config.get_vocab_size(),
2974 block_sizes=block_sizes,
2975 is_spec_decode=bool(self.vllm_config.speculative_config),
2976 + logitsprocs=self.input_batch.logitsprocs,
2977 + is_pooling_model=self.is_pooling_model,
2978 )
2979
2980 def _allocate_kv_cache_tensors(