| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | import tempfile |
| 6 | from enum import Enum, auto |
| 7 | from pathlib import Path |
| 8 | from typing import Any |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from tests.utils import requires_spawn_multiprocessing |
| 13 | from vllm.config import VllmConfig |
| 14 | from vllm.exceptions import VLLMValidationError |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.sampling_params import SamplingParams |
| 17 | from vllm.utils.torch_utils import async_tensor_h2d |
| 18 | from vllm.v1.sample.logits_processor import ( |
| 19 | LOGITSPROCS_GROUP, |
| 20 | AdapterLogitsProcessor, |
| 21 | BatchUpdate, |
| 22 | LogitsProcessor, |
| 23 | RequestLogitsProcessor, |
| 24 | ) |
| 25 | from vllm.v1.sample.logits_processor.builtin import process_dict_updates |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | MODEL_NAME = "facebook/opt-125m" |
| 30 | POOLING_MODEL_NAME = "BAAI/bge-base-en-v1.5" |
| 31 | DUMMY_LOGITPROC_ARG = "target_token" |
| 32 | TEMP_GREEDY = 0.0 |
| 33 | MAX_TOKENS = 20 |
| 34 | DUMMY_LOGITPROC_ENTRYPOINT = "dummy_logitproc" |
| 35 | DUMMY_LOGITPROC_MODULE = "tests.v1.logits_processors.utils" |
| 36 | DUMMY_LOGITPROC_FQCN = f"{DUMMY_LOGITPROC_MODULE}:DummyLogitsProcessor" |
| 37 | |
| 38 | |
| 39 | class CustomLogitprocSource(Enum): |
| 40 | """How to source a logitproc for testing purposes.""" |
| 41 | |
| 42 | LOGITPROC_SOURCE_NONE = auto() # No custom logitproc |
| 43 | LOGITPROC_SOURCE_ENTRYPOINT = auto() # Via entrypoint |
| 44 | LOGITPROC_SOURCE_FQCN = auto() # Via fully-qualified class name (FQCN) |
| 45 | LOGITPROC_SOURCE_CLASS = auto() # Via provided class object |
| 46 | |
| 47 | |
| 48 | # Sample prompts. |
| 49 | prompts = [ |
| 50 | "Hello, my name is", |
| 51 | "The president of the United States is", |
| 52 | "The capital of France is", |
| 53 | "The future of AI is", |
| 54 | ] |
| 55 | |
| 56 | |
| 57 | class DummyLogitsProcessor(LogitsProcessor): |
| 58 | """Fake logit processor to support unit testing and examples.""" |
| 59 | |
| 60 | @staticmethod |
| 61 | def _get_target_token(params: SamplingParams) -> int | None: |
| 62 | target_token = ( |
| 63 | params.extra_args.get(DUMMY_LOGITPROC_ARG) if params.extra_args else None |
| 64 | ) |
| 65 | if target_token is not None and not isinstance(target_token, int): |
| 66 | raise VLLMValidationError( |
| 67 | f"target_token value {target_token} {type(target_token)} is not int" |
| 68 | ) |
| 69 | return target_token |
| 70 | |
| 71 | @classmethod |
| 72 | def validate_params(cls, params: SamplingParams): |
| 73 | cls._get_target_token(params) |
| 74 | |
| 75 | def __init__( |
| 76 | self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool |
| 77 | ): |
| 78 | self.req_info: dict[int, int] = {} |
| 79 | |
| 80 | def is_argmax_invariant(self) -> bool: |
| 81 | """Never impacts greedy sampling.""" |
| 82 | return False |
| 83 | |
| 84 | def update_state(self, batch_update: BatchUpdate | None): |
| 85 | def extract_extra_arg(params: SamplingParams) -> int | None: |
| 86 | return self._get_target_token(params) |
| 87 | |
| 88 | process_dict_updates( |
| 89 | self.req_info, |
| 90 | batch_update, |
| 91 | lambda params, _, __: extract_extra_arg(params), |
| 92 | ) |
| 93 | |
| 94 | def apply(self, logits: torch.Tensor) -> torch.Tensor: |
| 95 | if not self.req_info: |
| 96 | return logits |
| 97 | |
| 98 | # Save target values before modification. |
| 99 | cols = async_tensor_h2d( |
| 100 | list(self.req_info.values()), dtype=torch.long, device=logits.device |
| 101 | ) |
| 102 | rows = async_tensor_h2d( |
| 103 | list(self.req_info.keys()), dtype=torch.long, device=logits.device |
| 104 | ) |
| 105 | values_to_keep = logits[rows, cols].clone() |
| 106 | |
| 107 | # Mask all but target tokens. Use an on-device fill tensor so the |
| 108 | # scatter doesn't force a synchronizing scalar H2D. |
| 109 | fill = torch.full( |
| 110 | (rows.numel(), logits.size(-1)), |
| 111 | float("-inf"), |
| 112 | dtype=logits.dtype, |
| 113 | device=logits.device, |
| 114 | ) |
| 115 | logits[rows] = fill |
| 116 | logits[rows, cols] = values_to_keep |
| 117 | |
| 118 | return logits |
| 119 | |
| 120 | |
| 121 | class EntryPoint: |
| 122 | """Dummy entrypoint class for logitsprocs testing.""" |
| 123 | |
| 124 | def __init__(self): |
| 125 | self.name = DUMMY_LOGITPROC_ENTRYPOINT |
| 126 | self.value = DUMMY_LOGITPROC_FQCN |
| 127 | |
| 128 | def load(self): |
| 129 | return DummyLogitsProcessor |
| 130 | |
| 131 | |
| 132 | class EntryPoints(list): |
| 133 | """Dummy EntryPoints class for logitsprocs testing.""" |
| 134 | |
| 135 | def __init__(self, group: str): |
| 136 | # Emulate list-like functionality |
| 137 | eps = [EntryPoint()] if group == LOGITSPROCS_GROUP else [] |
| 138 | super().__init__(eps) |
| 139 | # Extra attributes |
| 140 | self.names = [ep.name for ep in eps] |
| 141 | |
| 142 | |
| 143 | class DummyPerReqLogitsProcessor: |
| 144 | """The request-level logits processor masks out all logits except the |
| 145 | token id identified by `target_token`""" |
| 146 | |
| 147 | def __init__(self, target_token: int) -> None: |
| 148 | """Specify `target_token`.""" |
| 149 | self.target_token = target_token |
| 150 | |
| 151 | def __call__( |
| 152 | self, |
| 153 | output_ids: list[int], |
| 154 | logits: torch.Tensor, |
| 155 | ) -> torch.Tensor: |
| 156 | val_to_keep = logits[self.target_token].clone() |
| 157 | logits[:] = float("-inf") |
| 158 | logits[self.target_token] = val_to_keep |
| 159 | return logits |
| 160 | |
| 161 | |
| 162 | class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor): |
| 163 | """Example of wrapping a fake request-level logit processor to create a |
| 164 | batch-level logits processor""" |
| 165 | |
| 166 | def is_argmax_invariant(self) -> bool: |
| 167 | return False |
| 168 | |
| 169 | def new_req_logits_processor( |
| 170 | self, |
| 171 | params: SamplingParams, |
| 172 | ) -> RequestLogitsProcessor | None: |
| 173 | """This method returns a new request-level logits processor, customized |
| 174 | to the `target_token` value associated with a particular request. |
| 175 | |
| 176 | Returns None if the logits processor should not be applied to the |
| 177 | particular request. To use the logits processor the request must have |
| 178 | a "target_token" custom argument with an integer value. |
| 179 | |
| 180 | Args: |
| 181 | params: per-request sampling params |
| 182 | |
| 183 | Returns: |
| 184 | `Callable` request logits processor, or None |
| 185 | |
| 186 | """ |
| 187 | target_token: Any | None = params.extra_args and params.extra_args.get( |
| 188 | "target_token" |
| 189 | ) |
| 190 | if target_token is None: |
| 191 | return None |
| 192 | if not isinstance(target_token, int): |
| 193 | logger.warning( |
| 194 | "target_token value %s is not int; not applying logits" |
| 195 | " processor to request.", |
| 196 | target_token, |
| 197 | ) |
| 198 | return None |
| 199 | return DummyPerReqLogitsProcessor(target_token) |
| 200 | |
| 201 | |
| 202 | def register_fake_entrypoint(monkeypatch) -> str: |
| 203 | """Register the dummy logitsproc entrypoint in a way that is visible |
| 204 | to spawned subprocesses by creating a real dist-info directory on disk. |
| 205 | |
| 206 | Unlike monkey-patching importlib.metadata.entry_points (which only works |
| 207 | with fork), this approach writes a real dist-info package that |
| 208 | importlib.metadata can discover in any subprocess via PYTHONPATH. |
| 209 | |
| 210 | Returns the temp directory path. |
| 211 | """ |
| 212 | tmpdir = Path(tempfile.mkdtemp(prefix="dummy-logitproc-")) |
| 213 | dist_info = tmpdir / "dummy_logitproc-0.1.dist-info" |
| 214 | dist_info.mkdir() |
| 215 | |
| 216 | # Write METADATA file (required by importlib.metadata) |
| 217 | (dist_info / "METADATA").write_text( |
| 218 | "Metadata-Version: 2.1\nName: dummy-logitproc\nVersion: 0.1\n", |
| 219 | encoding="utf-8", |
| 220 | ) |
| 221 | |
| 222 | # Write entry_points.txt |
| 223 | (dist_info / "entry_points.txt").write_text( |
| 224 | f"[{LOGITSPROCS_GROUP}]\n" |
| 225 | f"{DUMMY_LOGITPROC_ENTRYPOINT} = {DUMMY_LOGITPROC_FQCN}\n", |
| 226 | encoding="utf-8", |
| 227 | ) |
| 228 | |
| 229 | # Add to PYTHONPATH so spawned subprocesses can discover it |
| 230 | existing = os.environ.get("PYTHONPATH", "") |
| 231 | monkeypatch.setenv( |
| 232 | "PYTHONPATH", str(tmpdir) + (os.pathsep + existing if existing else "") |
| 233 | ) |
| 234 | |
| 235 | # Also update sys.path for the current process so the driver can |
| 236 | # discover the entrypoint. |
| 237 | monkeypatch.syspath_prepend(str(tmpdir)) |
| 238 | |
| 239 | return str(tmpdir) |
| 240 | |
| 241 | |
| 242 | def setup_fake_entrypoint(monkeypatch) -> None: |
| 243 | """Expose the dummy logitproc entrypoint for the current platform.""" |
| 244 | if requires_spawn_multiprocessing(): |
| 245 | register_fake_entrypoint(monkeypatch) |
| 246 | monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
| 247 | return |
| 248 | |
| 249 | import importlib.metadata |
| 250 | |
| 251 | original_entry_points = importlib.metadata.entry_points |
| 252 | |
| 253 | def fake_entry_points(**kwargs): |
| 254 | group = kwargs.get("group") |
| 255 | if group == LOGITSPROCS_GROUP: |
| 256 | return EntryPoints(group) |
| 257 | return original_entry_points(**kwargs) |
| 258 | |
| 259 | monkeypatch.setattr(importlib.metadata, "entry_points", fake_entry_points) |
| 260 | monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork") |