| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable |
| 5 | from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar |
| 6 | |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | |
| 10 | import vllm.ir |
| 11 | from vllm.config import VllmConfig, set_current_vllm_config |
| 12 | from vllm.distributed.kv_transfer.kv_connector.utils import get_current_attn_backends |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.lora.request import LoRARequest |
| 15 | from vllm.multimodal.cache import worker_receiver_cache_from_config |
| 16 | from vllm.tracing import instrument |
| 17 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 18 | from vllm.utils.system_utils import update_environment_variables |
| 19 | from vllm.v1.attention.backends.utils import ( |
| 20 | get_supported_kv_cache_layouts, |
| 21 | record_kv_cache_layout, |
| 22 | ) |
| 23 | from vllm.v1.kv_cache_interface import KVCacheSpec |
| 24 | |
| 25 | if TYPE_CHECKING: |
| 26 | from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput |
| 27 | from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput |
| 28 | else: |
| 29 | SchedulerOutput = object |
| 30 | GrammarOutput = object |
| 31 | AsyncModelRunnerOutput = object |
| 32 | ModelRunnerOutput = object |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | _R = TypeVar("_R") |
| 37 | |
| 38 | |
| 39 | class CompilationTimes(NamedTuple): |
| 40 | language_model: float |
| 41 | encoder: float |
| 42 | |
| 43 | |
| 44 | class WorkerBase: |
| 45 | """Worker interface that allows vLLM to cleanly separate implementations for |
| 46 | different hardware. Also abstracts control plane communication, e.g., to |
| 47 | communicate request metadata to other workers. |
| 48 | """ |
| 49 | |
| 50 | def __init__( |
| 51 | self, |
| 52 | vllm_config: VllmConfig, |
| 53 | local_rank: int, |
| 54 | rank: int, |
| 55 | distributed_init_method: str, |
| 56 | is_driver_worker: bool = False, |
| 57 | ) -> None: |
| 58 | """Initialize common worker components. |
| 59 | |
| 60 | Args: |
| 61 | vllm_config: Complete vLLM configuration |
| 62 | local_rank: Local device index |
| 63 | rank: Global rank in distributed setup |
| 64 | distributed_init_method: Distributed initialization method |
| 65 | is_driver_worker: Whether this worker handles driver |
| 66 | responsibilities |
| 67 | |
| 68 | """ |
| 69 | self.vllm_config = vllm_config |
| 70 | self.model_config = vllm_config.model_config |
| 71 | self.cache_config = vllm_config.cache_config |
| 72 | self.lora_config = vllm_config.lora_config |
| 73 | self.load_config = vllm_config.load_config |
| 74 | self.parallel_config = vllm_config.parallel_config |
| 75 | self.scheduler_config = vllm_config.scheduler_config |
| 76 | self.device_config = vllm_config.device_config |
| 77 | self.speculative_config = vllm_config.speculative_config |
| 78 | self.observability_config = vllm_config.observability_config |
| 79 | self.kv_transfer_config = vllm_config.kv_transfer_config |
| 80 | self.compilation_config = vllm_config.compilation_config |
| 81 | |
| 82 | from vllm.platforms import current_platform |
| 83 | |
| 84 | self.current_platform = current_platform |
| 85 | |
| 86 | self.parallel_config.rank = rank |
| 87 | self.local_rank = local_rank |
| 88 | self.rank = rank |
| 89 | self.distributed_init_method = distributed_init_method |
| 90 | self.is_driver_worker = is_driver_worker |
| 91 | |
| 92 | # Device and model state |
| 93 | self.device: torch.device | None = None |
| 94 | self.model_runner: nn.Module | None = None |
| 95 | |
| 96 | # IR op priority and torch-wrap state are constant for the worker's |
| 97 | # lifetime. |
| 98 | vllm_config.kernel_config.ir_op_priority.set_default() |
| 99 | vllm.ir.set_default_torch_wrap( |
| 100 | vllm_config.compilation_config.ir_enable_torch_wrap |
| 101 | ) |
| 102 | |
| 103 | def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: |
| 104 | """Get specifications for KV cache implementation.""" |
| 105 | raise NotImplementedError |
| 106 | |
| 107 | def get_supported_kv_cache_layouts(self) -> list[str]: |
| 108 | """Layout names every attention backend supports, most preferred first.""" |
| 109 | backends = get_current_attn_backends(self.vllm_config) |
| 110 | return [layout.name for layout in get_supported_kv_cache_layouts(backends)] |
| 111 | |
| 112 | def set_kv_cache_layout(self, kv_cache_layout: str) -> None: |
| 113 | """Adopt the KV cache layout resolved by the engine core.""" |
| 114 | record_kv_cache_layout(self.vllm_config.cache_config, kv_cache_layout) |
| 115 | |
| 116 | def compile_or_warm_up_model(self) -> CompilationTimes: |
| 117 | """Prepare model for execution through compilation/warmup. |
| 118 | |
| 119 | Returns: |
| 120 | Compilation times (language_model, encoder) in seconds. |
| 121 | |
| 122 | """ |
| 123 | raise NotImplementedError |
| 124 | |
| 125 | def check_health(self) -> None: |
| 126 | """Basic health check (override for device-specific checks).""" |
| 127 | return |
| 128 | |
| 129 | def synchronize_device(self) -> None: |
| 130 | """Block until in-flight device work completes; backends outside |
| 131 | ``torch.accelerator`` must override with their own wait.""" |
| 132 | if torch.accelerator.is_available(): |
| 133 | torch.accelerator.synchronize() |
| 134 | |
| 135 | def init_device(self) -> None: |
| 136 | """Initialize device state, such as loading the model or other on-device |
| 137 | memory allocations. |
| 138 | """ |
| 139 | raise NotImplementedError |
| 140 | |
| 141 | def reset_mm_cache(self) -> None: |
| 142 | reset_fn = getattr(self.model_runner, "reset_mm_cache", None) |
| 143 | if callable(reset_fn): |
| 144 | reset_fn() |
| 145 | |
| 146 | def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: |
| 147 | release_fn = getattr( |
| 148 | self.model_runner, "release_late_interaction_query_cache", None |
| 149 | ) |
| 150 | if callable(release_fn): |
| 151 | release_fn(query_keys) |
| 152 | |
| 153 | def get_model(self) -> nn.Module: |
| 154 | raise NotImplementedError |
| 155 | |
| 156 | def supports_draft_weight_updates(self) -> bool: |
| 157 | """Whether this worker can update its configured speculative model.""" |
| 158 | return False |
| 159 | |
| 160 | def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R: |
| 161 | """Apply a function on the model inside this worker.""" |
| 162 | return fn(self.get_model()) |
| 163 | |
| 164 | def get_model_inspection(self) -> str: |
| 165 | """Return a transformers-style hierarchical view of the model.""" |
| 166 | from vllm.model_inspection import format_model_inspection |
| 167 | |
| 168 | return format_model_inspection(self.get_model()) |
| 169 | |
| 170 | def load_model(self, *, load_dummy_weights: bool = False) -> None: |
| 171 | """Load model onto target device.""" |
| 172 | raise NotImplementedError |
| 173 | |
| 174 | def execute_model( |
| 175 | self, scheduler_output: SchedulerOutput |
| 176 | ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None: |
| 177 | """If this method returns None, sample_tokens should be called immediately after |
| 178 | to obtain the ModelRunnerOutput. |
| 179 | |
| 180 | Note that this design may be changed in future if/when structured outputs |
| 181 | parallelism is re-architected. |
| 182 | """ |
| 183 | raise NotImplementedError |
| 184 | |
| 185 | def sample_tokens( |
| 186 | self, grammar_output: GrammarOutput |
| 187 | ) -> ModelRunnerOutput | AsyncModelRunnerOutput: |
| 188 | """Should be called immediately after execute_model iff it returned None.""" |
| 189 | raise NotImplementedError |
| 190 | |
| 191 | def get_cache_block_size_bytes(self) -> int: |
| 192 | """Return the size of a single cache block, in bytes. Used in |
| 193 | speculative decoding. |
| 194 | """ |
| 195 | raise NotImplementedError |
| 196 | |
| 197 | def add_lora(self, lora_request: LoRARequest) -> bool: |
| 198 | raise NotImplementedError |
| 199 | |
| 200 | def remove_lora(self, lora_id: int) -> bool: |
| 201 | raise NotImplementedError |
| 202 | |
| 203 | def pin_lora(self, lora_id: int) -> bool: |
| 204 | raise NotImplementedError |
| 205 | |
| 206 | def list_loras(self) -> set[int]: |
| 207 | raise NotImplementedError |
| 208 | |
| 209 | @property |
| 210 | def vocab_size(self) -> int: |
| 211 | """Get vocabulary size from model configuration.""" |
| 212 | return self.model_config.get_vocab_size() |
| 213 | |
| 214 | def shutdown(self) -> None: |
| 215 | """Clean up resources held by the worker.""" |
| 216 | return |
| 217 | |
| 218 | |
| 219 | class WorkerWrapperBase: |
| 220 | """This class represents one process in an executor/engine. It is responsible |
| 221 | for lazily initializing the worker and handling the worker's lifecycle. |
| 222 | We first instantiate the WorkerWrapper, which remembers the worker module |
| 223 | and class name. Then, when we call `update_environment_variables`, and the |
| 224 | real initialization happens in `init_worker`. |
| 225 | """ |
| 226 | |
| 227 | def __init__( |
| 228 | self, |
| 229 | rpc_rank: int = 0, |
| 230 | global_rank: int | None = None, |
| 231 | ) -> None: |
| 232 | """Initialize the worker wrapper with the given vllm_config and rpc_rank. |
| 233 | Note: rpc_rank is the rank of the worker in the executor. In most cases, |
| 234 | it is also the rank of the worker in the distributed group. However, |
| 235 | when multiple executors work together, they can be different. |
| 236 | e.g. in the case of SPMD-style offline inference with TP=2, |
| 237 | users can launch 2 engines/executors, each with only 1 worker. |
| 238 | All workers have rpc_rank=0, but they have different ranks in the TP |
| 239 | group. |
| 240 | """ |
| 241 | self.rpc_rank: int = rpc_rank |
| 242 | self.global_rank: int = self.rpc_rank if global_rank is None else global_rank |
| 243 | |
| 244 | # Initialized after init_worker is called |
| 245 | self.worker: WorkerBase |
| 246 | self.vllm_config: VllmConfig |
| 247 | |
| 248 | def shutdown(self) -> None: |
| 249 | if self.worker is not None: |
| 250 | self.worker.shutdown() |
| 251 | |
| 252 | def update_environment_variables( |
| 253 | self, |
| 254 | envs_list: list[dict[str, str]], |
| 255 | ) -> None: |
| 256 | envs = envs_list[self.rpc_rank] |
| 257 | update_environment_variables(envs) |
| 258 | |
| 259 | @instrument(span_name="Worker init") |
| 260 | def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: |
| 261 | """Here we inject some common logic before initializing the worker. |
| 262 | Arguments are passed to the worker class constructor. |
| 263 | """ |
| 264 | kwargs = all_kwargs[self.rpc_rank] |
| 265 | |
| 266 | vllm_config: VllmConfig | None = kwargs.get("vllm_config") |
| 267 | assert vllm_config is not None, ( |
| 268 | "vllm_config is required to initialize the worker" |
| 269 | ) |
| 270 | self.vllm_config = vllm_config |
| 271 | |
| 272 | vllm_config.enable_trace_function_call_for_thread() |
| 273 | |
| 274 | from vllm.plugins import load_general_plugins |
| 275 | |
| 276 | load_general_plugins() |
| 277 | |
| 278 | # Let the platform replace core Triton kernels (e.g. CPU fallback |
| 279 | # implementations) before any worker or model code can launch them. |
| 280 | from vllm.platforms import current_platform |
| 281 | |
| 282 | current_platform.register_triton_kernel_overrides() |
| 283 | |
| 284 | parallel_config = vllm_config.parallel_config |
| 285 | if isinstance(parallel_config.worker_cls, str): |
| 286 | worker_class: type[WorkerBase] = resolve_obj_by_qualname( |
| 287 | parallel_config.worker_cls |
| 288 | ) |
| 289 | else: |
| 290 | raise ValueError( |
| 291 | "passing worker_cls is no longer supported. " |
| 292 | "Please pass keep the class in a separate module " |
| 293 | "and pass the qualified name of the class as a string." |
| 294 | ) |
| 295 | |
| 296 | if parallel_config.worker_extension_cls: |
| 297 | worker_extension_cls = resolve_obj_by_qualname( |
| 298 | parallel_config.worker_extension_cls |
| 299 | ) |
| 300 | extended_calls = [] |
| 301 | if worker_extension_cls not in worker_class.__bases__: |
| 302 | # check any conflicts between worker and worker_extension_cls |
| 303 | for attr in dir(worker_extension_cls): |
| 304 | if attr.startswith("__"): |
| 305 | continue |
| 306 | assert not hasattr(worker_class, attr), ( |
| 307 | f"Worker class {worker_class} already has an attribute" |
| 308 | f" {attr}, which conflicts with the worker" |
| 309 | f" extension class {worker_extension_cls}." |
| 310 | ) |
| 311 | if callable(getattr(worker_extension_cls, attr)): |
| 312 | extended_calls.append(attr) |
| 313 | # dynamically inherit the worker extension class |
| 314 | worker_class.__bases__ = worker_class.__bases__ + ( |
| 315 | worker_extension_cls, |
| 316 | ) |
| 317 | logger.info( |
| 318 | "Injected %s into %s for extended collective_rpc calls %s", |
| 319 | worker_extension_cls, |
| 320 | worker_class, |
| 321 | extended_calls, |
| 322 | ) |
| 323 | |
| 324 | assigned_physical_gpu_ids = kwargs.pop("assigned_physical_gpu_ids", None) |
| 325 | if assigned_physical_gpu_ids is not None: |
| 326 | vllm_config.parallel_config.assigned_physical_gpu_ids = ( |
| 327 | assigned_physical_gpu_ids |
| 328 | ) |
| 329 | |
| 330 | shared_worker_lock = kwargs.pop("shared_worker_lock", None) |
| 331 | self.mm_receiver_cache = worker_receiver_cache_from_config( |
| 332 | vllm_config, |
| 333 | shared_worker_lock, |
| 334 | ) |
| 335 | |
| 336 | with set_current_vllm_config(self.vllm_config): |
| 337 | # To make vLLM config available during worker initialization |
| 338 | self.worker = worker_class(**kwargs) |
| 339 | |
| 340 | def initialize_from_config(self, kv_cache_configs: list[Any]) -> None: |
| 341 | kv_cache_config = kv_cache_configs[self.global_rank] |
| 342 | assert self.vllm_config is not None |
| 343 | with set_current_vllm_config(self.vllm_config): |
| 344 | self.worker.initialize_from_config(kv_cache_config) # type: ignore |
| 345 | |
| 346 | def init_device(self): |
| 347 | assert self.vllm_config is not None |
| 348 | with set_current_vllm_config(self.vllm_config): |
| 349 | # To make vLLM config available during device initialization |
| 350 | self.worker.init_device() # type: ignore |
| 351 | |
| 352 | def __getattr__(self, attr: str): |
| 353 | return getattr(self.worker, attr) |
| 354 | |
| 355 | def _apply_mm_cache(self, scheduler_output: SchedulerOutput) -> None: |
| 356 | mm_cache = self.mm_receiver_cache |
| 357 | if mm_cache is None: |
| 358 | return |
| 359 | |
| 360 | for req_data in scheduler_output.scheduled_new_reqs: |
| 361 | req_data.mm_features = mm_cache.get_and_update_features( |
| 362 | req_data.mm_features |
| 363 | ) |
| 364 | |
| 365 | def execute_model( |
| 366 | self, scheduler_output: SchedulerOutput |
| 367 | ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None: |
| 368 | self._apply_mm_cache(scheduler_output) |
| 369 | |
| 370 | return self.worker.execute_model(scheduler_output) |
| 371 | |
| 372 | def reset_mm_cache(self) -> None: |
| 373 | mm_receiver_cache = self.mm_receiver_cache |
| 374 | if mm_receiver_cache is not None: |
| 375 | mm_receiver_cache.clear_cache() |
| 376 | |
| 377 | self.worker.reset_mm_cache() |