| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable, Sequence |
| 5 | from pathlib import Path |
| 6 | from typing import TYPE_CHECKING, Any |
| 7 | |
| 8 | import cloudpickle |
| 9 | import torch.nn as nn |
| 10 | from pydantic import ValidationError |
| 11 | from tqdm.auto import tqdm |
| 12 | from typing_extensions import overload |
| 13 | |
| 14 | from vllm.config import ( |
| 15 | AttentionConfig, |
| 16 | CompilationConfig, |
| 17 | PoolerConfig, |
| 18 | ProfilerConfig, |
| 19 | StructuredOutputsConfig, |
| 20 | is_init_field, |
| 21 | ) |
| 22 | from vllm.config.compilation import CompilationMode |
| 23 | from vllm.config.model import ( |
| 24 | ConvertOption, |
| 25 | HfOverrides, |
| 26 | ModelDType, |
| 27 | RunnerOption, |
| 28 | TokenizerMode, |
| 29 | ) |
| 30 | from vllm.config.quantization import QuantizationConfigArgs |
| 31 | from vllm.distributed.weight_transfer.base import ( |
| 32 | WeightTransferInitRequest, |
| 33 | WeightTransferUpdateRequest, |
| 34 | ) |
| 35 | from vllm.engine.arg_utils import EngineArgs |
| 36 | from vllm.entrypoints.chat_utils import ( |
| 37 | ChatCompletionMessageParam, |
| 38 | ChatTemplateContentFormatOption, |
| 39 | load_chat_template, |
| 40 | ) |
| 41 | from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin |
| 42 | from vllm.entrypoints.pooling.offline import PoolingOfflineMixin |
| 43 | from vllm.entrypoints.serve.utils.api_utils import log_non_default_args |
| 44 | from vllm.inputs import PromptType |
| 45 | from vllm.logger import configure_logging_if_needed, init_logger |
| 46 | from vllm.lora.request import LoRARequest |
| 47 | from vllm.model_executor.layers.quantization import QuantizationMethods |
| 48 | from vllm.outputs import PoolingRequestOutput, RequestOutput |
| 49 | from vllm.platforms import current_platform |
| 50 | from vllm.sampling_params import SamplingParams |
| 51 | from vllm.tokenizers import TokenizerLike |
| 52 | from vllm.usage.usage_lib import UsageContext |
| 53 | from vllm.utils.counter import Counter |
| 54 | from vllm.v1.engine import PauseMode |
| 55 | from vllm.v1.engine.llm_engine import LLMEngine |
| 56 | from vllm.v1.sample.logits_processor import LogitsProcessor |
| 57 | |
| 58 | from ..renderers import ChatParams |
| 59 | from .offline_utils import _O, _R, OfflineInferenceMixin |
| 60 | |
| 61 | if TYPE_CHECKING: |
| 62 | from vllm.v1.metrics.reader import Metric |
| 63 | |
| 64 | logger = init_logger(__name__) |
| 65 | |
| 66 | |
| 67 | class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): |
| 68 | """An LLM for generating texts from given prompts and sampling parameters. |
| 69 | |
| 70 | This class includes a tokenizer, a language model (possibly distributed |
| 71 | across multiple GPUs), and GPU memory space allocated for intermediate |
| 72 | states (aka KV cache). Given a batch of prompts and sampling parameters, |
| 73 | this class generates texts from the model, using an intelligent batching |
| 74 | mechanism and efficient memory management. |
| 75 | |
| 76 | Args: |
| 77 | model: The name or path of a HuggingFace Transformers model. |
| 78 | tokenizer: The name or path of a HuggingFace Transformers tokenizer. |
| 79 | tokenizer_mode: The tokenizer mode. See |
| 80 | [ModelConfig.tokenizer_mode][vllm.config.ModelConfig.tokenizer_mode]. |
| 81 | skip_tokenizer_init: If true, skip initialization of tokenizer and |
| 82 | detokenizer. Expect valid prompt_token_ids and None for prompt |
| 83 | from the input. |
| 84 | trust_remote_code: Trust remote code (e.g., from HuggingFace) when |
| 85 | downloading the model and tokenizer. |
| 86 | allowed_local_media_path: Allowing API requests to read local images |
| 87 | or videos from directories specified by the server file system. |
| 88 | This is a security risk. Should only be enabled in trusted |
| 89 | environments. |
| 90 | allowed_media_domains: If set, only media URLs that belong to this |
| 91 | domain can be used for multi-modal inputs. |
| 92 | tensor_parallel_size: The number of GPUs to use for distributed |
| 93 | execution with tensor parallelism. |
| 94 | dtype: The data type for the model weights and activations. Currently, |
| 95 | we support `float32`, `float16`, and `bfloat16`. If `auto`, we use |
| 96 | the `dtype` attribute of the Transformers model's config. However, |
| 97 | if the `dtype` in the config is `float32`, we will use `float16` instead. |
| 98 | quantization: The method used to quantize the model weights. Currently, |
| 99 | we support "awq", "gptq", and "fp8" (experimental). |
| 100 | If None, we first check the `quantization_config` attribute in the |
| 101 | model config file. If that is None, we assume the model weights are |
| 102 | not quantized and use `dtype` to determine the data type of |
| 103 | the weights. |
| 104 | revision: The specific model version to use. It can be a branch name, |
| 105 | a tag name, or a commit id. |
| 106 | tokenizer_revision: The specific tokenizer version to use. It can be a |
| 107 | branch name, a tag name, or a commit id. |
| 108 | chat_template: The chat template to apply. |
| 109 | seed: The seed to initialize the random number generator for sampling. |
| 110 | gpu_memory_utilization: The ratio (between 0 and 1) of GPU memory to |
| 111 | reserve for the model weights, activations, and KV cache. Higher |
| 112 | values will increase the KV cache size and thus improve the model's |
| 113 | throughput. However, if the value is too high, it may cause out-of- |
| 114 | memory (OOM) errors. |
| 115 | kv_cache_memory_bytes: Size of KV Cache per GPU in bytes. By default, |
| 116 | this is set to None and vllm can automatically infer the kv cache |
| 117 | size based on gpu_memory_utilization. However, users may want to |
| 118 | manually specify the kv cache memory size. kv_cache_memory_bytes |
| 119 | allows more fine-grain control of how much memory gets used when |
| 120 | compared with using gpu_memory_utilization. Note that |
| 121 | kv_cache_memory_bytes (when not-None) ignores |
| 122 | gpu_memory_utilization |
| 123 | cpu_offload_gb: The size (GiB) of CPU memory to use for offloading |
| 124 | the model weights. This virtually increases the GPU memory space |
| 125 | you can use to hold the model weights, at the cost of CPU-GPU data |
| 126 | transfer for every forward pass. |
| 127 | offload_group_size: Prefetch offloading: Group every N layers |
| 128 | together. Offload last `offload_num_in_group` layers of each group. |
| 129 | Default is 0 (disabled). |
| 130 | offload_num_in_group: Prefetch offloading: Number of layers to |
| 131 | offload per group. Default is 1. |
| 132 | offload_prefetch_step: Prefetch offloading: Number of layers to |
| 133 | prefetch ahead. Higher values hide more latency but use more GPU |
| 134 | memory. Default is 1. |
| 135 | offload_params: Prefetch offloading: Set of parameter name segments |
| 136 | to selectively offload. Only parameters whose names contain one of |
| 137 | these segments will be offloaded (e.g., {"gate_up_proj", "down_proj"} |
| 138 | for MLP weights, or {"w13_weight", "w2_weight"} for MoE expert |
| 139 | weights). If None or empty, all parameters are offloaded. |
| 140 | enforce_eager: Whether to enforce eager execution. If True, we will |
| 141 | disable CUDA graph and always execute the model in eager mode. |
| 142 | If False, we will use CUDA graph and eager execution in hybrid. |
| 143 | enable_return_routed_experts: Whether to return routed experts. |
| 144 | disable_custom_all_reduce: See |
| 145 | [ParallelConfig][vllm.config.ParallelConfig]. |
| 146 | hf_token: The token to use as HTTP bearer authorization for remote files |
| 147 | . If `True`, will use the token generated when running |
| 148 | `hf auth login` (stored in `~/.cache/huggingface/token`). |
| 149 | hf_overrides: If a dictionary, contains arguments to be forwarded to the |
| 150 | HuggingFace config. If a callable, it is called to update the |
| 151 | HuggingFace config. |
| 152 | mm_processor_kwargs: Arguments to be forwarded to the model's processor |
| 153 | for multi-modal data, e.g., image processor. Overrides for the |
| 154 | multi-modal processor obtained from `AutoProcessor.from_pretrained`. |
| 155 | The available overrides depend on the model that is being run. |
| 156 | For example, for Phi-3-Vision: `{"num_crops": 4}`. |
| 157 | pooler_config: Initialize non-default pooling config for the pooling model, |
| 158 | e.g., `PoolerConfig(seq_pooling_type="MEAN", use_activation=False)`. |
| 159 | compilation_config: Either an integer or a dictionary. If it is an |
| 160 | integer, it is used as the mode of compilation optimization. If it |
| 161 | is a dictionary, it can specify the full compilation configuration. |
| 162 | attention_config: Configuration for attention mechanisms. Can be a |
| 163 | dictionary or an AttentionConfig instance. If a dictionary, it will |
| 164 | be converted to an AttentionConfig. Allows specifying the attention |
| 165 | backend and other attention-related settings. |
| 166 | return_sampling_mask: Return each sampled token's post-processing |
| 167 | support set. Requires Model Runner V2 and processed log probabilities. |
| 168 | spec_method: Top-level alias for `speculative_config["method"]`. |
| 169 | spec_model: Top-level alias for `speculative_config["model"]`. |
| 170 | spec_tokens: Top-level alias for |
| 171 | `speculative_config["num_speculative_tokens"]`. |
| 172 | **kwargs: Arguments for [`EngineArgs`][vllm.EngineArgs]. |
| 173 | |
| 174 | Note: |
| 175 | This class is intended to be used for offline inference. For online |
| 176 | serving, use the [AsyncLLMEngine][vllm.AsyncLLMEngine] class instead. |
| 177 | |
| 178 | """ |
| 179 | |
| 180 | def __init__( |
| 181 | self, |
| 182 | model: str, |
| 183 | *, |
| 184 | runner: RunnerOption = "auto", |
| 185 | convert: ConvertOption = "auto", |
| 186 | tokenizer: str | None = None, |
| 187 | tokenizer_mode: TokenizerMode | str = "auto", |
| 188 | skip_tokenizer_init: bool = False, |
| 189 | trust_remote_code: bool = False, |
| 190 | allowed_local_media_path: str = "", |
| 191 | allowed_media_domains: list[str] | None = None, |
| 192 | tensor_parallel_size: int = 1, |
| 193 | dtype: ModelDType = "auto", |
| 194 | quantization: QuantizationMethods | None = None, |
| 195 | revision: str | None = None, |
| 196 | tokenizer_revision: str | None = None, |
| 197 | chat_template: Path | str | None = None, |
| 198 | seed: int = 0, |
| 199 | gpu_memory_utilization: float = 0.92, |
| 200 | cpu_offload_gb: float = 0, |
| 201 | offload_group_size: int = 0, |
| 202 | offload_num_in_group: int = 1, |
| 203 | offload_prefetch_step: int = 1, |
| 204 | offload_params: set[str] | None = None, |
| 205 | enforce_eager: bool = False, |
| 206 | enable_return_routed_experts: bool = False, |
| 207 | return_sampling_mask: bool = False, |
| 208 | disable_custom_all_reduce: bool = False, |
| 209 | hf_token: bool | str | None = None, |
| 210 | hf_overrides: HfOverrides | None = None, |
| 211 | mm_processor_kwargs: dict[str, Any] | None = None, |
| 212 | pooler_config: PoolerConfig | None = None, |
| 213 | structured_outputs_config: dict[str, Any] |
| 214 | | StructuredOutputsConfig |
| 215 | | None = None, |
| 216 | profiler_config: dict[str, Any] | ProfilerConfig | None = None, |
| 217 | attention_config: dict[str, Any] | AttentionConfig | None = None, |
| 218 | kv_cache_memory_bytes: int | None = None, |
| 219 | compilation_config: int | dict[str, Any] | CompilationConfig | None = None, |
| 220 | quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None, |
| 221 | logits_processors: list[str | type[LogitsProcessor]] | None = None, |
| 222 | spec_method: str | None = None, |
| 223 | spec_model: str | None = None, |
| 224 | spec_tokens: int | None = None, |
| 225 | **kwargs: Any, |
| 226 | ) -> None: |
| 227 | """LLM constructor.""" |
| 228 | if "disable_log_stats" not in kwargs: |
| 229 | kwargs["disable_log_stats"] = True |
| 230 | |
| 231 | if "worker_cls" in kwargs: |
| 232 | worker_cls = kwargs["worker_cls"] |
| 233 | # if the worker_cls is not qualified string name, |
| 234 | # we serialize it using cloudpickle to avoid pickling issues |
| 235 | if isinstance(worker_cls, type): |
| 236 | kwargs["worker_cls"] = cloudpickle.dumps(worker_cls) |
| 237 | |
| 238 | if "kv_transfer_config" in kwargs and isinstance( |
| 239 | kwargs["kv_transfer_config"], dict |
| 240 | ): |
| 241 | from vllm.config.kv_transfer import KVTransferConfig |
| 242 | |
| 243 | raw_config_dict = kwargs["kv_transfer_config"] |
| 244 | try: |
| 245 | kwargs["kv_transfer_config"] = KVTransferConfig(**raw_config_dict) |
| 246 | except ValidationError as e: |
| 247 | logger.error( |
| 248 | "Failed to convert 'kv_transfer_config' dict to " |
| 249 | "KVTransferConfig object. Dict: %s. Error: %s", |
| 250 | raw_config_dict, |
| 251 | e, |
| 252 | ) |
| 253 | # Consider re-raising a more specific vLLM error or ValueError |
| 254 | # to provide better context to the user. |
| 255 | raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e |
| 256 | |
| 257 | if hf_overrides is None: |
| 258 | hf_overrides = {} |
| 259 | |
| 260 | def _make_config(value: Any, cls: type[_R]) -> _R: |
| 261 | """Convert dict/None/instance to a config instance.""" |
| 262 | if value is None: |
| 263 | return cls() |
| 264 | if isinstance(value, dict): |
| 265 | return cls(**{k: v for k, v in value.items() if is_init_field(cls, k)}) # type: ignore[arg-type] |
| 266 | return value |
| 267 | |
| 268 | if isinstance(compilation_config, int): |
| 269 | compilation_config_instance = CompilationConfig( |
| 270 | mode=CompilationMode(compilation_config) |
| 271 | ) |
| 272 | else: |
| 273 | compilation_config_instance = _make_config( |
| 274 | compilation_config, CompilationConfig |
| 275 | ) |
| 276 | |
| 277 | structured_outputs_instance = _make_config( |
| 278 | structured_outputs_config, StructuredOutputsConfig |
| 279 | ) |
| 280 | profiler_config_instance = _make_config(profiler_config, ProfilerConfig) |
| 281 | attention_config_instance = _make_config(attention_config, AttentionConfig) |
| 282 | |
| 283 | # warn about single-process data parallel usage. |
| 284 | _dp_size = int(kwargs.get("data_parallel_size", 1)) |
| 285 | _distributed_executor_backend = kwargs.get("distributed_executor_backend") |
| 286 | if ( |
| 287 | _dp_size > 1 |
| 288 | and not _distributed_executor_backend == "external_launcher" |
| 289 | and not current_platform.is_tpu() |
| 290 | ): |
| 291 | raise ValueError( |
| 292 | f"LLM(data_parallel_size={_dp_size}) is not supported for single-" |
| 293 | "process usage and may hang. Please use " |
| 294 | "the explicit multi-process data-parallel example at " |
| 295 | "'examples/features/data_parallel/data_parallel_offline.py'." |
| 296 | ) |
| 297 | |
| 298 | engine_args = EngineArgs( |
| 299 | model=model, |
| 300 | runner=runner, |
| 301 | convert=convert, |
| 302 | tokenizer=tokenizer, |
| 303 | tokenizer_mode=tokenizer_mode, |
| 304 | skip_tokenizer_init=skip_tokenizer_init, |
| 305 | trust_remote_code=trust_remote_code, |
| 306 | allowed_local_media_path=allowed_local_media_path, |
| 307 | allowed_media_domains=allowed_media_domains, |
| 308 | tensor_parallel_size=tensor_parallel_size, |
| 309 | dtype=dtype, |
| 310 | quantization=quantization, |
| 311 | revision=revision, |
| 312 | tokenizer_revision=tokenizer_revision, |
| 313 | seed=seed, |
| 314 | gpu_memory_utilization=gpu_memory_utilization, |
| 315 | kv_cache_memory_bytes=kv_cache_memory_bytes, |
| 316 | cpu_offload_gb=cpu_offload_gb, |
| 317 | offload_group_size=offload_group_size, |
| 318 | offload_num_in_group=offload_num_in_group, |
| 319 | offload_prefetch_step=offload_prefetch_step, |
| 320 | offload_params=offload_params or set(), |
| 321 | enforce_eager=enforce_eager, |
| 322 | enable_return_routed_experts=enable_return_routed_experts, |
| 323 | return_sampling_mask=return_sampling_mask, |
| 324 | disable_custom_all_reduce=disable_custom_all_reduce, |
| 325 | hf_token=hf_token, |
| 326 | hf_overrides=hf_overrides, |
| 327 | mm_processor_kwargs=mm_processor_kwargs, |
| 328 | pooler_config=pooler_config, |
| 329 | structured_outputs_config=structured_outputs_instance, |
| 330 | profiler_config=profiler_config_instance, |
| 331 | attention_config=attention_config_instance, |
| 332 | compilation_config=compilation_config_instance, |
| 333 | quantization_config=quantization_config, |
| 334 | logits_processors=logits_processors, |
| 335 | spec_method=spec_method, |
| 336 | spec_model=spec_model, |
| 337 | spec_tokens=spec_tokens, |
| 338 | **kwargs, |
| 339 | ) |
| 340 | |
| 341 | configure_logging_if_needed(engine_args.create_logging_config()) |
| 342 | log_non_default_args(engine_args) |
| 343 | |
| 344 | self.llm_engine = LLMEngine.from_engine_args( |
| 345 | engine_args=engine_args, usage_context=UsageContext.LLM_CLASS |
| 346 | ) |
| 347 | self.model_config = self.llm_engine.model_config |
| 348 | self.engine_class = type(self.llm_engine) |
| 349 | |
| 350 | self.request_counter = Counter() |
| 351 | self.default_sampling_params: dict[str, Any] | None = None |
| 352 | |
| 353 | supported_tasks = self.llm_engine.get_supported_tasks() |
| 354 | self.supported_tasks = supported_tasks |
| 355 | |
| 356 | self.runner_type = self.model_config.runner_type |
| 357 | self.renderer = self.llm_engine.renderer |
| 358 | self.chat_template = load_chat_template(chat_template) |
| 359 | self.input_processor = self.llm_engine.input_processor |
| 360 | |
| 361 | self.renderer.warmup(ChatParams(chat_template=self.chat_template)) |
| 362 | |
| 363 | # The renderer thread pool is only consumed by the async renderer |
| 364 | # path; the synchronous `LLM` entrypoint runs multimodal |
| 365 | # preprocessing serially. Warn so the setting is not a silent |
| 366 | # no-op. See vllm-project/vllm#42901. |
| 367 | if self.model_config.renderer_num_workers > 1 and self.runner_type != "pooling": |
| 368 | logger.warning_once( |
| 369 | "`renderer_num_workers=%d` was set, but the offline `LLM` " |
| 370 | "entrypoint uses the synchronous renderer path and runs " |
| 371 | "multimodal preprocessing serially across prompts. The " |
| 372 | "renderer thread pool is only consumed by the async " |
| 373 | "renderer path used by `vllm serve` / `AsyncLLM`, so this " |
| 374 | "setting has no effect here.", |
| 375 | self.model_config.renderer_num_workers, |
| 376 | ) |
| 377 | |
| 378 | PoolingOfflineMixin.__init__(self) |
| 379 | |
| 380 | # Cache for __repr__ to avoid repeated collective_rpc calls |
| 381 | self._cached_repr: str | None = None |
| 382 | |
| 383 | @classmethod |
| 384 | def from_engine_args(cls, engine_args: EngineArgs) -> "LLM": |
| 385 | """Create an LLM instance from EngineArgs.""" |
| 386 | return cls(**vars(engine_args)) |
| 387 | |
| 388 | def get_tokenizer(self) -> TokenizerLike: |
| 389 | return self.llm_engine.get_tokenizer() |
| 390 | |
| 391 | def get_world_size(self, include_dp: bool = True) -> int: |
| 392 | """Get the world size from the parallel config. |
| 393 | |
| 394 | Args: |
| 395 | include_dp: If True (default), returns the world size including |
| 396 | data parallelism (TP * PP * DP). If False, returns the world |
| 397 | size without data parallelism (TP * PP). |
| 398 | |
| 399 | Returns: |
| 400 | The world size (tensor_parallel_size * pipeline_parallel_size), |
| 401 | optionally multiplied by data_parallel_size if include_dp is True. |
| 402 | |
| 403 | """ |
| 404 | parallel_config = self.llm_engine.vllm_config.parallel_config |
| 405 | if include_dp: |
| 406 | return parallel_config.world_size_across_dp |
| 407 | return parallel_config.world_size |
| 408 | |
| 409 | def reset_mm_cache(self) -> None: |
| 410 | self.renderer.clear_mm_cache() |
| 411 | self.llm_engine.reset_mm_cache() |
| 412 | |
| 413 | def get_default_sampling_params(self) -> SamplingParams: |
| 414 | if self.default_sampling_params is None: |
| 415 | self.default_sampling_params = self.model_config.get_diff_sampling_param() |
| 416 | if self.default_sampling_params: |
| 417 | return SamplingParams.from_optional(**self.default_sampling_params) |
| 418 | return SamplingParams() |
| 419 | |
| 420 | def generate( |
| 421 | self, |
| 422 | prompts: PromptType | Sequence[PromptType], |
| 423 | sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, |
| 424 | *, |
| 425 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 426 | lora_request: Sequence[LoRARequest] | LoRARequest | None = None, |
| 427 | priority: list[int] | None = None, |
| 428 | tokenization_kwargs: dict[str, Any] | None = None, |
| 429 | mm_processor_kwargs: dict[str, Any] | None = None, |
| 430 | ) -> list[RequestOutput]: |
| 431 | """Generates the completions for the input prompts. |
| 432 | |
| 433 | This class automatically batches the given prompts, considering |
| 434 | the memory constraint. For the best performance, put all of your prompts |
| 435 | into a single list and pass it to this method. |
| 436 | |
| 437 | Args: |
| 438 | prompts: The prompts to the LLM. You may pass a sequence of prompts |
| 439 | for batch inference. See [PromptType][vllm.inputs.PromptType] |
| 440 | for more details about the format of each prompt. |
| 441 | sampling_params: The sampling parameters for text generation. If |
| 442 | None, we use the default sampling parameters. |
| 443 | When it is a single value, it is applied to every prompt. |
| 444 | When it is a list, the list must have the same length as the |
| 445 | prompts and it is paired one by one with the prompt. |
| 446 | use_tqdm: If `True`, shows a tqdm progress bar. |
| 447 | If a callable (e.g., `functools.partial(tqdm, leave=False)`), |
| 448 | it is used to create the progress bar. |
| 449 | If `False`, no progress bar is created. |
| 450 | lora_request: LoRA request to use for generation, if any. |
| 451 | priority: The priority of the requests, if any. |
| 452 | Only applicable when priority scheduling policy is enabled. |
| 453 | If provided, must be a list of integers matching the length |
| 454 | of `prompts`, where each priority value corresponds to the prompt |
| 455 | at the same index. |
| 456 | tokenization_kwargs: Overrides for `tokenizer.encode`. |
| 457 | mm_processor_kwargs: Overrides for `processor.__call__`. |
| 458 | |
| 459 | Returns: |
| 460 | A list of `RequestOutput` objects containing the |
| 461 | generated completions in the same order as the input prompts. |
| 462 | |
| 463 | """ |
| 464 | runner_type = self.model_config.runner_type |
| 465 | if runner_type != "generate": |
| 466 | raise ValueError( |
| 467 | "LLM.generate() is only supported for generative models. " |
| 468 | "Try passing `--runner generate` to use the model as a " |
| 469 | "generative model." |
| 470 | ) |
| 471 | |
| 472 | if sampling_params is None: |
| 473 | sampling_params = self.get_default_sampling_params() |
| 474 | |
| 475 | return self._run_completion( |
| 476 | prompts=prompts, |
| 477 | params=sampling_params, |
| 478 | output_type=RequestOutput, |
| 479 | use_tqdm=use_tqdm, |
| 480 | lora_request=lora_request, |
| 481 | tokenization_kwargs=tokenization_kwargs, |
| 482 | priority=priority, |
| 483 | mm_processor_kwargs=mm_processor_kwargs, |
| 484 | ) |
| 485 | |
| 486 | def enqueue( |
| 487 | self, |
| 488 | prompts: PromptType | Sequence[PromptType], |
| 489 | sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, |
| 490 | lora_request: Sequence[LoRARequest] | LoRARequest | None = None, |
| 491 | priority: list[int] | None = None, |
| 492 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 493 | tokenization_kwargs: dict[str, Any] | None = None, |
| 494 | mm_processor_kwargs: dict[str, Any] | None = None, |
| 495 | ) -> list[str]: |
| 496 | """Enqueue prompts for generation without waiting for completion. |
| 497 | |
| 498 | This method adds requests to the engine queue but does not start |
| 499 | processing them. Use wait_for_completion() to process the queued |
| 500 | requests and get results. |
| 501 | |
| 502 | Args: |
| 503 | prompts: The prompts to the LLM. See generate() for details. |
| 504 | sampling_params: The sampling parameters for text generation. |
| 505 | lora_request: LoRA request to use for generation, if any. |
| 506 | priority: The priority of the requests, if any. |
| 507 | use_tqdm: If True, shows a tqdm progress bar while adding requests. |
| 508 | tokenization_kwargs: Overrides for `tokenizer.encode`. |
| 509 | mm_processor_kwargs: Overrides for `processor.__call__`. |
| 510 | |
| 511 | Returns: |
| 512 | A list of request IDs for the enqueued requests. |
| 513 | |
| 514 | """ |
| 515 | runner_type = self.model_config.runner_type |
| 516 | if runner_type != "generate": |
| 517 | raise ValueError("LLM.enqueue() is only supported for generative models.") |
| 518 | |
| 519 | if sampling_params is None: |
| 520 | sampling_params = self.get_default_sampling_params() |
| 521 | |
| 522 | return self._add_completion_requests( |
| 523 | prompts=prompts, |
| 524 | params=sampling_params, |
| 525 | use_tqdm=use_tqdm, |
| 526 | lora_request=lora_request, |
| 527 | priority=priority, |
| 528 | tokenization_kwargs=tokenization_kwargs, |
| 529 | mm_processor_kwargs=mm_processor_kwargs, |
| 530 | ) |
| 531 | |
| 532 | @overload |
| 533 | def wait_for_completion( |
| 534 | self, |
| 535 | *, |
| 536 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 537 | ) -> list[RequestOutput | PoolingRequestOutput]: ... |
| 538 | |
| 539 | @overload |
| 540 | def wait_for_completion( |
| 541 | self, |
| 542 | output_type: type[_O] | tuple[type[_O], ...], |
| 543 | *, |
| 544 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 545 | ) -> list[_O]: ... |
| 546 | |
| 547 | def wait_for_completion( |
| 548 | self, |
| 549 | output_type: type[Any] | tuple[type[Any], ...] | None = None, |
| 550 | *, |
| 551 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 552 | ) -> list[Any]: |
| 553 | """Wait for all enqueued requests to complete and return results. |
| 554 | |
| 555 | This method processes all requests currently in the engine queue |
| 556 | and returns their outputs. Use after enqueue() to get results. |
| 557 | |
| 558 | Args: |
| 559 | output_type: The expected output type(s). If not provided, accepts |
| 560 | both RequestOutput and PoolingRequestOutput. |
| 561 | use_tqdm: If True, shows a tqdm progress bar. |
| 562 | |
| 563 | Returns: |
| 564 | A list of output objects for all completed requests. |
| 565 | |
| 566 | """ |
| 567 | if output_type is None: |
| 568 | output_type = (RequestOutput, PoolingRequestOutput) |
| 569 | |
| 570 | return self._run_engine(output_type, use_tqdm=use_tqdm) |
| 571 | |
| 572 | def collective_rpc( |
| 573 | self, |
| 574 | method: str | Callable[..., _R], |
| 575 | timeout: float | None = None, |
| 576 | args: tuple = (), |
| 577 | kwargs: dict[str, Any] | None = None, |
| 578 | ) -> list[_R]: |
| 579 | """Execute an RPC call on all workers. |
| 580 | |
| 581 | Args: |
| 582 | method: Name of the worker method to execute, or a callable that |
| 583 | is serialized and sent to all workers to execute. |
| 584 | |
| 585 | If the method is a callable, it should accept an additional |
| 586 | `self` argument, in addition to the arguments passed in `args` |
| 587 | and `kwargs`. The `self` argument will be the worker object. |
| 588 | timeout: Maximum time in seconds to wait for execution. Raises a |
| 589 | [`TimeoutError`][] on timeout. `None` means wait indefinitely. |
| 590 | args: Positional arguments to pass to the worker method. |
| 591 | kwargs: Keyword arguments to pass to the worker method. |
| 592 | |
| 593 | Returns: |
| 594 | A list containing the results from each worker. |
| 595 | |
| 596 | Note: |
| 597 | It is recommended to use this API to only pass control messages, |
| 598 | and set up data-plane communication to pass data. |
| 599 | |
| 600 | """ |
| 601 | return self.llm_engine.collective_rpc(method, timeout, args, kwargs) |
| 602 | |
| 603 | def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]: |
| 604 | """Run a function directly on the model inside each worker, |
| 605 | returning the result for each of them. |
| 606 | |
| 607 | !!! warning |
| 608 | To reduce the overhead of data transfer, avoid returning large |
| 609 | arrays or tensors from this method. If you must return them, |
| 610 | make sure you move them to CPU first to avoid taking up additional |
| 611 | VRAM! |
| 612 | """ |
| 613 | return self.llm_engine.apply_model(func) |
| 614 | |
| 615 | def chat( |
| 616 | self, |
| 617 | messages: list[ChatCompletionMessageParam] |
| 618 | | Sequence[list[ChatCompletionMessageParam]], |
| 619 | sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, |
| 620 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 621 | lora_request: Sequence[LoRARequest] | LoRARequest | None = None, |
| 622 | chat_template: str | None = None, |
| 623 | chat_template_content_format: ChatTemplateContentFormatOption = "auto", |
| 624 | add_generation_prompt: bool = True, |
| 625 | continue_final_message: bool = False, |
| 626 | tools: list[dict[str, Any]] | None = None, |
| 627 | chat_template_kwargs: dict[str, Any] | None = None, |
| 628 | tokenization_kwargs: dict[str, Any] | None = None, |
| 629 | mm_processor_kwargs: dict[str, Any] | None = None, |
| 630 | ) -> list[RequestOutput]: |
| 631 | """Generate responses for a chat conversation. |
| 632 | |
| 633 | The chat conversation is converted into a text prompt using the |
| 634 | tokenizer and calls the [generate][vllm.LLM.generate] method to generate |
| 635 | the responses. |
| 636 | |
| 637 | Multi-modal inputs can be passed in the same way you would pass them |
| 638 | to the OpenAI API. |
| 639 | |
| 640 | Args: |
| 641 | messages: A sequence of conversations or a single conversation. |
| 642 | |
| 643 | - Each conversation is represented as a list of messages. |
| 644 | - Each message is a dictionary with 'role' and 'content' keys. |
| 645 | |
| 646 | sampling_params: The sampling parameters for text generation. |
| 647 | If None, we use the default sampling parameters. When it |
| 648 | is a single value, it is applied to every prompt. When it |
| 649 | is a list, the list must have the same length as the |
| 650 | prompts and it is paired one by one with the prompt. |
| 651 | use_tqdm: If `True`, shows a tqdm progress bar. |
| 652 | If a callable (e.g., `functools.partial(tqdm, leave=False)`), |
| 653 | it is used to create the progress bar. |
| 654 | If `False`, no progress bar is created. |
| 655 | lora_request: LoRA request to use for generation, if any. |
| 656 | chat_template: The template to use for structuring the chat. |
| 657 | If not provided, the model's default chat template will be used. |
| 658 | chat_template_content_format: The format to render message content. |
| 659 | |
| 660 | - "string" will render the content as a string. |
| 661 | Example: `"Who are you?"` |
| 662 | - "openai" will render the content as a list of dictionaries, |
| 663 | similar to OpenAI schema. |
| 664 | Example: `[{"type": "text", "text": "Who are you?"}]` |
| 665 | |
| 666 | add_generation_prompt: If True, adds a generation template |
| 667 | to each message. |
| 668 | continue_final_message: If True, continues the final message in |
| 669 | the conversation instead of starting a new one. Cannot be |
| 670 | `True` if `add_generation_prompt` is also `True`. |
| 671 | chat_template_kwargs: Additional kwargs to pass to the chat |
| 672 | template. |
| 673 | tokenization_kwargs: Overrides for `tokenizer.encode`. |
| 674 | mm_processor_kwargs: Overrides for `processor.__call__`. |
| 675 | tools: Tools to make available to the model, if any. |
| 676 | |
| 677 | Returns: |
| 678 | A list of `RequestOutput` objects containing the generated |
| 679 | responses in the same order as the input messages. |
| 680 | |
| 681 | """ |
| 682 | model_config = self.model_config |
| 683 | runner_type = model_config.runner_type |
| 684 | if runner_type != "generate": |
| 685 | raise ValueError( |
| 686 | "LLM.chat() is only supported for generative models. " |
| 687 | "Try passing `--runner generate` to use the model as a " |
| 688 | "generative model." |
| 689 | ) |
| 690 | |
| 691 | if sampling_params is None: |
| 692 | sampling_params = self.get_default_sampling_params() |
| 693 | |
| 694 | return self._run_chat( |
| 695 | messages=messages, |
| 696 | params=sampling_params, |
| 697 | output_type=RequestOutput, |
| 698 | use_tqdm=use_tqdm, |
| 699 | lora_request=lora_request, |
| 700 | chat_template=chat_template, |
| 701 | chat_template_content_format=chat_template_content_format, |
| 702 | chat_template_kwargs=chat_template_kwargs, |
| 703 | add_generation_prompt=add_generation_prompt, |
| 704 | continue_final_message=continue_final_message, |
| 705 | tools=tools, |
| 706 | tokenization_kwargs=tokenization_kwargs, |
| 707 | mm_processor_kwargs=mm_processor_kwargs, |
| 708 | ) |
| 709 | |
| 710 | def enqueue_chat( |
| 711 | self, |
| 712 | messages: list[ChatCompletionMessageParam] |
| 713 | | Sequence[list[ChatCompletionMessageParam]], |
| 714 | sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, |
| 715 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 716 | lora_request: Sequence[LoRARequest] | LoRARequest | None = None, |
| 717 | priority: list[int] | None = None, |
| 718 | chat_template: str | None = None, |
| 719 | chat_template_content_format: ChatTemplateContentFormatOption = "auto", |
| 720 | add_generation_prompt: bool = True, |
| 721 | continue_final_message: bool = False, |
| 722 | tools: list[dict[str, Any]] | None = None, |
| 723 | chat_template_kwargs: dict[str, Any] | None = None, |
| 724 | tokenization_kwargs: dict[str, Any] | None = None, |
| 725 | mm_processor_kwargs: dict[str, Any] | None = None, |
| 726 | ) -> list[str]: |
| 727 | """Enqueue chat conversations for generation without waiting. |
| 728 | |
| 729 | This method renders chat conversations and adds the resulting requests |
| 730 | to the engine queue. Use wait_for_completion() to get results. To |
| 731 | guarantee that all requests are queued before scheduling starts, pause |
| 732 | scheduling with sleep(level=0) before calling this method and resume it |
| 733 | with wake_up(tags=["scheduling"]) afterward. |
| 734 | |
| 735 | Args: |
| 736 | messages: A sequence of conversations or a single conversation. |
| 737 | Each conversation is represented as a list of messages. |
| 738 | sampling_params: The sampling parameters for text generation. |
| 739 | If None, we use the default sampling parameters. |
| 740 | use_tqdm: If `True`, shows a tqdm progress bar while rendering |
| 741 | conversations. |
| 742 | lora_request: LoRA request to use for generation, if any. |
| 743 | priority: The priority of the requests, if any. |
| 744 | chat_template: The template to use for structuring the chat. |
| 745 | chat_template_content_format: The format to render message content. |
| 746 | add_generation_prompt: If True, adds a generation template |
| 747 | to each message. |
| 748 | continue_final_message: If True, continues the final message in |
| 749 | the conversation instead of starting a new one. |
| 750 | tools: Tools to make available to the model, if any. |
| 751 | chat_template_kwargs: Additional kwargs to pass to the chat |
| 752 | template. |
| 753 | tokenization_kwargs: Overrides for `tokenizer.encode`. |
| 754 | mm_processor_kwargs: Overrides for `processor.__call__`. |
| 755 | |
| 756 | Returns: |
| 757 | A list of request IDs for the enqueued requests. |
| 758 | |
| 759 | """ |
| 760 | model_config = self.model_config |
| 761 | runner_type = model_config.runner_type |
| 762 | if runner_type != "generate": |
| 763 | raise ValueError( |
| 764 | "LLM.enqueue_chat() is only supported for generative models. " |
| 765 | "Try passing `--runner generate` to use the model as a " |
| 766 | "generative model." |
| 767 | ) |
| 768 | |
| 769 | if sampling_params is None: |
| 770 | sampling_params = self.get_default_sampling_params() |
| 771 | |
| 772 | return self._add_chat_requests( |
| 773 | messages=messages, |
| 774 | params=sampling_params, |
| 775 | use_tqdm=use_tqdm, |
| 776 | lora_request=lora_request, |
| 777 | priority=priority, |
| 778 | chat_template=chat_template, |
| 779 | chat_template_content_format=chat_template_content_format, |
| 780 | chat_template_kwargs=chat_template_kwargs, |
| 781 | add_generation_prompt=add_generation_prompt, |
| 782 | continue_final_message=continue_final_message, |
| 783 | tools=tools, |
| 784 | tokenization_kwargs=tokenization_kwargs, |
| 785 | mm_processor_kwargs=mm_processor_kwargs, |
| 786 | ) |
| 787 | |
| 788 | def start_profile(self, profile_prefix: str | None = None) -> None: |
| 789 | """Start profiling with optional custom trace prefix. |
| 790 | |
| 791 | Args: |
| 792 | profile_prefix: Optional prefix for the trace file names. If provided, |
| 793 | trace files will be named as "<prefix>_dp<X>_pp<Y>_tp<Z>". |
| 794 | If not provided, default naming will be used. |
| 795 | |
| 796 | """ |
| 797 | self.llm_engine.start_profile(profile_prefix) |
| 798 | |
| 799 | def stop_profile(self) -> None: |
| 800 | self.llm_engine.stop_profile() |
| 801 | |
| 802 | def reset_prefix_cache( |
| 803 | self, reset_running_requests: bool = False, reset_connector: bool = False |
| 804 | ) -> bool: |
| 805 | return self.llm_engine.reset_prefix_cache( |
| 806 | reset_running_requests, reset_connector |
| 807 | ) |
| 808 | |
| 809 | def sleep(self, level: int = 1, mode: PauseMode = "abort"): |
| 810 | """Put the engine to sleep. The engine should not process any requests. |
| 811 | The caller should guarantee that no requests are being processed |
| 812 | during the sleep period, before `wake_up` is called. |
| 813 | |
| 814 | Args: |
| 815 | level: The sleep level. |
| 816 | - Level 0: Pause scheduling but continue accepting requests. |
| 817 | Requests are queued but not processed. |
| 818 | - Level 1: Offload model weights to CPU, discard KV cache. |
| 819 | The content of kv cache is forgotten. Good for |
| 820 | sleeping and waking up the engine to run the same |
| 821 | model again. Please make sure there's enough CPU |
| 822 | memory to store the model weights. |
| 823 | - Level 2: Discard all GPU memory (weights + KV cache). |
| 824 | Good for sleeping and waking up the engine to run |
| 825 | a different model or update the model, where |
| 826 | previous model weights are not needed. It reduces |
| 827 | CPU memory pressure. |
| 828 | mode: How to handle any existing requests, can be "abort", "wait", |
| 829 | or "keep". |
| 830 | |
| 831 | """ |
| 832 | self.llm_engine.sleep(level=level, mode=mode) |
| 833 | |
| 834 | def release_kv_cache_memory(self) -> None: |
| 835 | """Release the GPU physical memory backing the KV cache. |
| 836 | |
| 837 | Requires a completed ``sleep(level=0)`` and resident executor memory. |
| 838 | Restore with ``wake_up(tags=["kv_cache"])``; kept requests are recomputed. |
| 839 | """ |
| 840 | self.llm_engine.release_kv_cache_memory() |
| 841 | |
| 842 | def wake_up(self, tags: list[str] | None = None) -> bool: |
| 843 | """Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep] |
| 844 | method for more details. |
| 845 | |
| 846 | Args: |
| 847 | tags: An optional list of tags to reallocate the engine memory |
| 848 | for specific memory allocations. Values must be in |
| 849 | `("weights", "kv_cache", "scheduling")`. If None, all memory |
| 850 | is reallocated. wake_up should be called with all tags |
| 851 | (or None) before the engine is used again. |
| 852 | Use tags=["scheduling"] to resume from level 0 sleep. |
| 853 | |
| 854 | Returns: |
| 855 | Whether the engine is fully awake. |
| 856 | |
| 857 | """ |
| 858 | return self.llm_engine.wake_up(tags) |
| 859 | |
| 860 | def get_metrics(self) -> list["Metric"]: |
| 861 | """Return a snapshot of aggregated metrics from Prometheus. |
| 862 | |
| 863 | Returns: |
| 864 | A `MetricSnapshot` instance capturing the current state |
| 865 | of all aggregated metrics from Prometheus. |
| 866 | |
| 867 | Note: |
| 868 | This method is only available with the V1 LLM engine. |
| 869 | |
| 870 | """ |
| 871 | return self.llm_engine.get_metrics() |
| 872 | |
| 873 | def init_weight_transfer_engine( |
| 874 | self, request: WeightTransferInitRequest | dict |
| 875 | ) -> None: |
| 876 | """Initialize weight transfer for RL training. |
| 877 | |
| 878 | Args: |
| 879 | request: Weight transfer initialization request with backend-specific info |
| 880 | |
| 881 | """ |
| 882 | init_info_dict = ( |
| 883 | request["init_info"] if isinstance(request, dict) else request.init_info |
| 884 | ) |
| 885 | |
| 886 | self.llm_engine.collective_rpc( |
| 887 | "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} |
| 888 | ) |
| 889 | |
| 890 | def start_weight_update(self) -> None: |
| 891 | """Start a new weight update.""" |
| 892 | self.llm_engine.collective_rpc("start_weight_update") |
| 893 | |
| 894 | def start_draft_weight_update(self) -> None: |
| 895 | """Start a new weight update targeting the speculative draft model.""" |
| 896 | self.llm_engine.collective_rpc("start_draft_weight_update") |
| 897 | |
| 898 | def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: |
| 899 | """Update the weights of the model. |
| 900 | |
| 901 | Args: |
| 902 | request: Weight update request with backend-specific update info |
| 903 | |
| 904 | """ |
| 905 | update_info_dict = ( |
| 906 | request["update_info"] if isinstance(request, dict) else request.update_info |
| 907 | ) |
| 908 | |
| 909 | self.llm_engine.collective_rpc( |
| 910 | "update_weights", kwargs={"update_info": update_info_dict} |
| 911 | ) |
| 912 | |
| 913 | def finish_weight_update(self, weight_version: str | None = None) -> None: |
| 914 | """Finish the weight update and set its version if provided.""" |
| 915 | self.llm_engine.collective_rpc("finish_weight_update") |
| 916 | if weight_version is not None: |
| 917 | self.llm_engine.set_weight_version(weight_version) |
| 918 | |
| 919 | def update_weight_version(self, new_version: str) -> None: |
| 920 | """Set the weight version without updating weights.""" |
| 921 | self.llm_engine.set_weight_version(new_version) |
| 922 | |
| 923 | def get_weight_version(self) -> str: |
| 924 | """Return the latest committed weight version.""" |
| 925 | return self.llm_engine.get_weight_version() |
| 926 | |
| 927 | def __repr__(self) -> str: |
| 928 | """Return a transformers-style hierarchical view of the model.""" |
| 929 | # Cache the result to avoid repeated collective_rpc calls |
| 930 | if self._cached_repr is None: |
| 931 | results = self.llm_engine.collective_rpc("get_model_inspection") |
| 932 | # In distributed settings, we get results from all workers |
| 933 | # Just return the first one (they should all be the same) |
| 934 | if results: |
| 935 | self._cached_repr = results[0] |
| 936 | else: |
| 937 | self._cached_repr = f"LLM(model={self.model_config.model!r})" |
| 938 | return self._cached_repr |