__init__.py 92.7 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import contextlib
5
import datetime
Woosuk Kwon's avatar
Woosuk Kwon committed
6
import enum
7
import gc
8
import getpass
9
import hashlib
10
import importlib
11
import importlib.metadata
12
import importlib.util
13
import inspect
14
import ipaddress
15
import json
16
import multiprocessing
17
import os
18
import pickle
19
import signal
20
import socket
21
import subprocess
22
import sys
23
import tempfile
24
import textwrap
25
import threading
26
import time
27
import traceback
28
import types
Zhuohan Li's avatar
Zhuohan Li committed
29
import uuid
30
import warnings
31
import weakref
32
33
34
35
36
37
38
39
from argparse import (
    Action,
    ArgumentDefaultsHelpFormatter,
    ArgumentParser,
    ArgumentTypeError,
    RawDescriptionHelpFormatter,
    _ArgumentGroup,
)
40
from collections import defaultdict
41
from collections.abc import (
42
    Callable,
43
44
45
46
47
    Collection,
    Generator,
    Iterator,
    Sequence,
)
48
from concurrent.futures.process import ProcessPoolExecutor
49
from dataclasses import dataclass, field
50
from functools import cache, lru_cache, partial, wraps
51
from pathlib import Path
52
from typing import TYPE_CHECKING, Any, TextIO, TypeVar
53
from urllib.parse import urlparse
54
from uuid import uuid4
Zhuohan Li's avatar
Zhuohan Li committed
55

56
import cbor2
57
import cloudpickle
58
import numpy as np
59
import numpy.typing as npt
60
import psutil
61
import regex as re
62
import setproctitle
Zhuohan Li's avatar
Zhuohan Li committed
63
import torch
64
import torch.types
65
import yaml
66
67
import zmq
import zmq.asyncio
68
from packaging import version
69
from packaging.version import Version
70
from torch.library import Library
71
from typing_extensions import Never
72

73
import vllm.envs as envs
74
from vllm.logger import enable_trace_function_call, init_logger
75
from vllm.ray.lazy_utils import is_in_ray_actor
76

77
if TYPE_CHECKING:
78
79
    from argparse import Namespace

80
    from vllm.config import ModelConfig, VllmConfig
81
    from vllm.sequence import IntermediateTensors
82
83
84
85
86
87
else:
    Namespace = object

    ModelConfig = object
    VllmConfig = object
    IntermediateTensors = object
88

89
90
logger = init_logger(__name__)

91
92
93
94
95
96
# This value is chosen to have a balance between ITL and TTFT. Note it is
# not optimized for throughput.
DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048
POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768
MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120

97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
# Constants related to forcing the attention backend selection

# String name of register which may be set in order to
# force auto-selection of attention backend by Attention
# wrapper
STR_BACKEND_ENV_VAR: str = "VLLM_ATTENTION_BACKEND"

# Possible string values of STR_BACKEND_ENV_VAR
# register, corresponding to possible backends
STR_FLASHINFER_ATTN_VAL: str = "FLASHINFER"
STR_TORCH_SDPA_ATTN_VAL: str = "TORCH_SDPA"
STR_XFORMERS_ATTN_VAL: str = "XFORMERS"
STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
STR_INVALID_VAL: str = "INVALID"

112
113
114
115
116
117
MB_bytes = 1_000_000
"""The number of bytes in one megabyte (MB)."""

MiB_bytes = 1 << 20
"""The number of bytes in one mebibyte (MiB)."""

118
119
120
GB_bytes = 1_000_000_000
"""The number of bytes in one gigabyte (GB)."""

121
122
123
GiB_bytes = 1 << 30
"""The number of bytes in one gibibyte (GiB)."""

124
# ANSI color codes
125
126
CYAN = "\033[1;36m"
RESET = "\033[0;0m"
127

128
STR_DTYPE_TO_TORCH_DTYPE = {
129
    "float32": torch.float32,
130
131
132
    "half": torch.half,
    "bfloat16": torch.bfloat16,
    "float": torch.float,
133
    "fp8": torch.uint8,
134
135
    "fp8_e4m3": torch.uint8,
    "fp8_e5m2": torch.uint8,
136
    "int8": torch.int8,
137
    "fp8_inc": torch.float8_e4m3fn,
138
    "fp8_ds_mla": torch.uint8,
139
}
Zhuohan Li's avatar
Zhuohan Li committed
140

141
142
143
144
145
146
147
148
149
TORCH_DTYPE_TO_NUMPY_DTYPE = {
    torch.float16: np.float16,
    torch.float32: np.float32,
    torch.float64: np.float64,
    torch.uint8: np.uint8,
    torch.int32: np.int32,
    torch.int64: np.int64,
}

150
151
152
153
154
155
156
157
158
159

@contextlib.contextmanager
def set_default_torch_num_threads(num_threads: int):
    """Sets the default number of threads for PyTorch to the given value."""
    old_num_threads = torch.get_num_threads()
    torch.set_num_threads(num_threads)
    yield
    torch.set_num_threads(old_num_threads)


160
T = TypeVar("T")
161
U = TypeVar("U")
162

163

Woosuk Kwon's avatar
Woosuk Kwon committed
164
165
166
167
168
class Device(enum.Enum):
    GPU = enum.auto()
    CPU = enum.auto()


169
170
171
172
173
class LayerBlockType(enum.Enum):
    attention = "attention"
    mamba = "mamba"


Woosuk Kwon's avatar
Woosuk Kwon committed
174
175
176
177
class Counter:
    def __init__(self, start: int = 0) -> None:
        self.counter = start

Woosuk Kwon's avatar
Woosuk Kwon committed
178
    def __next__(self) -> int:
179
        i = self.counter
Woosuk Kwon's avatar
Woosuk Kwon committed
180
        self.counter += 1
181
        return i
Woosuk Kwon's avatar
Woosuk Kwon committed
182
183
184

    def reset(self) -> None:
        self.counter = 0
Zhuohan Li's avatar
Zhuohan Li committed
185

186

187
@cache
188
189
def get_max_shared_memory_bytes(gpu: int = 0) -> int:
    """Returns the maximum shared memory per thread block in bytes."""
190
    from vllm import _custom_ops as ops
191
192

    max_shared_mem = ops.get_max_shared_memory_per_block_device_attribute(gpu)
193
194
    # value 0 will cause MAX_SEQ_LEN become negative and test_attention.py
    # will fail
195
    assert max_shared_mem > 0, "max_shared_mem can not be zero"
196
197
198
    return int(max_shared_mem)


199
def get_cpu_memory() -> int:
200
    """Returns the total CPU memory of the node in bytes."""
201
    return psutil.virtual_memory().total
Zhuohan Li's avatar
Zhuohan Li committed
202
203
204
205


def random_uuid() -> str:
    return str(uuid.uuid4().hex)
206

207

208
def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]):
209
210
211
212
213
    for sock in sockets:
        if sock is not None:
            sock.close(linger=0)


214
def get_ip() -> str:
215
    host_ip = envs.VLLM_HOST_IP
216
217
218
219
    if "HOST_IP" in os.environ and "VLLM_HOST_IP" not in os.environ:
        logger.warning(
            "The environment variable HOST_IP is deprecated and ignored, as"
            " it is often used by Docker and other software to"
220
            " interact with the container's network stack. Please "
221
            "use VLLM_HOST_IP instead to set the IP address for vLLM processes"
222
223
            " to communicate with each other."
        )
224
225
226
227
228
    if host_ip:
        return host_ip

    # IP is not set, try to get it from the network interface

229
    # try ipv4
230
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
231
    try:
232
        s.connect(("8.8.8.8", 80))  # Doesn't need to be reachable
233
        return s.getsockname()[0]
234
235
236
237
238
    except Exception:
        pass

    # try ipv6
    try:
239
        s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
240
241
242
        # Google's public DNS server, see
        # https://developers.google.com/speed/public-dns/docs/using#addresses
        s.connect(("2001:4860:4860::8888", 80))  # Doesn't need to be reachable
243
        return s.getsockname()[0]
244
245
246
247
248
    except Exception:
        pass

    warnings.warn(
        "Failed to get the IP address, using 0.0.0.0 by default."
249
250
        "The value can be set by the environment variable"
        " VLLM_HOST_IP or HOST_IP.",
251
252
        stacklevel=2,
    )
253
    return "0.0.0.0"
254
255


256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
def test_loopback_bind(address, family):
    try:
        s = socket.socket(family, socket.SOCK_DGRAM)
        s.bind((address, 0))  # Port 0 = auto assign
        s.close()
        return True
    except OSError:
        return False


def get_loopback_ip() -> str:
    loopback_ip = envs.VLLM_LOOPBACK_IP
    if loopback_ip:
        return loopback_ip

    # VLLM_LOOPBACK_IP is not set, try to get it based on network interface

    if test_loopback_bind("127.0.0.1", socket.AF_INET):
        return "127.0.0.1"
    elif test_loopback_bind("::1", socket.AF_INET6):
        return "::1"
    else:
        raise RuntimeError(
            "Neither 127.0.0.1 nor ::1 are bound to a local interface. "
280
281
            "Set the VLLM_LOOPBACK_IP environment variable explicitly."
        )
282
283


284
285
286
287
288
289
290
291
def is_valid_ipv6_address(address: str) -> bool:
    try:
        ipaddress.IPv6Address(address)
        return True
    except ValueError:
        return False


292
def split_host_port(host_port: str) -> tuple[str, int]:
293
    # ipv6
294
295
    if host_port.startswith("["):
        host, port = host_port.rsplit("]", 1)
296
        host = host[1:]
297
        port = port.split(":")[1]
298
299
        return host, int(port)
    else:
300
        host, port = host_port.split(":")
301
302
303
304
305
306
307
308
309
310
        return host, int(port)


def join_host_port(host: str, port: int) -> str:
    if is_valid_ipv6_address(host):
        return f"[{host}]:{port}"
    else:
        return f"{host}:{port}"


311
def get_distributed_init_method(ip: str, port: int) -> str:
312
313
314
315
    return get_tcp_uri(ip, port)


def get_tcp_uri(ip: str, port: int) -> str:
316
317
318
319
    if is_valid_ipv6_address(ip):
        return f"tcp://[{ip}]:{port}"
    else:
        return f"tcp://{ip}:{port}"
320
321


322
323
324
325
326
def get_open_zmq_ipc_path() -> str:
    base_rpc_path = envs.VLLM_RPC_BASE_PATH
    return f"ipc://{base_rpc_path}/{uuid4()}"


327
328
329
330
def get_open_zmq_inproc_path() -> str:
    return f"inproc://{uuid4()}"


331
def get_open_port() -> int:
332
333
334
335
336
337
338
339
340
    """
    Get an open port for the vLLM process to listen on.
    An edge case to handle, is when we run data parallel,
    we need to avoid ports that are potentially used by
    the data parallel master process.
    Right now we reserve 10 ports for the data parallel master
    process. Currently it uses 2 ports.
    """
    if "VLLM_DP_MASTER_PORT" in os.environ:
341
342
        dp_master_port = envs.VLLM_DP_MASTER_PORT
        reserved_port_range = range(dp_master_port, dp_master_port + 10)
343
        while True:
344
345
346
            candidate_port = _get_open_port()
            if candidate_port not in reserved_port_range:
                return candidate_port
347
348
    return _get_open_port()

youkaichao's avatar
youkaichao committed
349

350
351
def get_open_ports_list(count: int = 5) -> list[int]:
    """Get a list of open ports."""
352
    ports = set[int]()
353
354
355
356
357
    while len(ports) < count:
        ports.add(get_open_port())
    return list(ports)


358
def _get_open_port() -> int:
359
    port = envs.VLLM_PORT
360
    if port is not None:
361
362
363
364
365
366
367
        while True:
            try:
                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    s.bind(("", port))
                    return port
            except OSError:
                port += 1  # Increment port number if already in use
368
                logger.info("Port %d is already in use, trying port %d", port - 1, port)
369
370
371
372
373
374
375
376
377
378
    # try ipv4
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("", 0))
            return s.getsockname()[1]
    except OSError:
        # try ipv6
        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
            s.bind(("", 0))
            return s.getsockname()[1]
379
380


381
def find_process_using_port(port: int) -> psutil.Process | None:
382
383
384
385
386
387
388
    # TODO: We can not check for running processes with network
    # port on macOS. Therefore, we can not have a full graceful shutdown
    # of vLLM. For now, let's not look for processes in this case.
    # Ref: https://www.florianreinhard.de/accessdenied-in-psutil/
    if sys.platform.startswith("darwin"):
        return None

389
    our_pid = os.getpid()
390
    for conn in psutil.net_connections():
391
        if conn.laddr.port == port and (conn.pid is not None and conn.pid != our_pid):
392
393
394
395
396
397
398
            try:
                return psutil.Process(conn.pid)
            except psutil.NoSuchProcess:
                return None
    return None


399
def update_environment_variables(envs: dict[str, str]):
400
    for k, v in envs.items():
401
        if k in os.environ and os.environ[k] != v:
402
            logger.warning(
403
404
405
406
407
                "Overwriting environment variable %s from '%s' to '%s'",
                k,
                os.environ[k],
                v,
            )
408
        os.environ[k] = v
409
410


411
412
413
414
415
def cdiv(a: int, b: int) -> int:
    """Ceiling division."""
    return -(a // -b)


416
417
418
419
420
421
422
def next_power_of_2(n) -> int:
    """The next power of 2 (inclusive)"""
    if n < 1:
        return 1
    return 1 << (n - 1).bit_length()


423
424
425
426
427
428
429
def prev_power_of_2(n: int) -> int:
    """The previous power of 2 (inclusive)"""
    if n <= 0:
        return 0
    return 1 << (n.bit_length() - 1)


430
431
432
433
def round_up(x: int, y: int) -> int:
    return ((x + y - 1) // y) * y


434
435
436
437
def round_down(x: int, y: int) -> int:
    return (x // y) * y


438
def _generate_random_fp8(
439
    tensor: torch.Tensor,
440
441
442
443
444
445
    low: float,
    high: float,
) -> None:
    # NOTE(zhaoyang): Due to NaN and Inf representation for fp8 data type,
    # it may occur Inf or NaN if we directly use torch.randint
    # to generate random data for fp8 data.
446
    # For example, s.11111.00 in fp8e5m2 format represents Inf.
447
    #     | E4M3        | E5M2
448
    # -----|-------------|-------------------
449
450
    # Inf | N/A         | s.11111.00
    # NaN | s.1111.111  | s.11111.{01,10,11}
451
    from vllm import _custom_ops as ops
452

453
454
    tensor_tmp = torch.empty_like(tensor, dtype=torch.float16)
    tensor_tmp.uniform_(low, high)
455
    ops.convert_fp8(tensor, tensor_tmp)
456
457
458
    del tensor_tmp


459
def get_kv_cache_torch_dtype(
460
461
    cache_dtype: str | torch.dtype | None,
    model_dtype: str | torch.dtype | None = None,
462
) -> torch.dtype:
463
464
    if isinstance(cache_dtype, str):
        if cache_dtype == "auto":
465
            if isinstance(model_dtype, str) and model_dtype in STR_DTYPE_TO_TORCH_DTYPE:
466
467
468
469
470
                torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
            elif isinstance(model_dtype, torch.dtype):
                torch_dtype = model_dtype
            else:
                raise ValueError(f"Invalid model dtype: {model_dtype}")
471
        elif cache_dtype in STR_DTYPE_TO_TORCH_DTYPE:
472
473
474
475
476
477
478
            torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
        else:
            raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
    elif isinstance(cache_dtype, torch.dtype):
        torch_dtype = cache_dtype
    else:
        raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
479
480
481
482
483
484
485
486
487
    return torch_dtype


def create_kv_caches_with_random_flash(
    num_blocks: int,
    block_size: int,
    num_layers: int,
    num_heads: int,
    head_size: int,
488
489
    cache_dtype: str | torch.dtype | None,
    model_dtype: str | torch.dtype | None = None,
490
491
492
    seed: int | None = None,
    device: str | None = "cuda",
    cache_layout: str | None = "NHD",
493
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
494
    from vllm.platforms import current_platform
495

496
    current_platform.seed_everything(seed)
497

498
    dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
499
500
    generic_kv_cache_shape = (num_blocks, 2, block_size, num_heads, head_size)
    assert cache_layout in ("NHD", "HND")
501
    stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3, 2, 4)
502

503
    kv_cache_allocation_shape = tuple(generic_kv_cache_shape[i] for i in stride_order)
504
    scale = head_size**-0.5
505

506
507
    key_caches: list[torch.Tensor] = []
    value_caches: list[torch.Tensor] = []
508

509
    for _ in range(num_layers):
510
        key_value_cache = torch.empty(
511
            size=kv_cache_allocation_shape, dtype=dtype, device=device
512
        ).permute(*stride_order)
513
514
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
            key_value_cache.uniform_(-scale, scale)
515
        elif cache_dtype == "fp8":
516
517
            _generate_random_fp8(key_value_cache, -scale, scale)
        else:
518
            raise ValueError(f"Does not support key cache of type {cache_dtype}")
519
520
521
522
523
524
525
526
527
528
529
        key_caches.append(key_value_cache[:, 0])
        value_caches.append(key_value_cache[:, 1])
    return key_caches, value_caches


def create_kv_caches_with_random(
    num_blocks: int,
    block_size: int,
    num_layers: int,
    num_heads: int,
    head_size: int,
530
531
    cache_dtype: str | torch.dtype | None,
    model_dtype: str | torch.dtype | None = None,
532
533
    seed: int | None = None,
    device: str | None = "cuda",
534
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
Joe's avatar
Joe committed
535
536
537
538
    if cache_dtype == "fp8" and head_size % 16:
        raise ValueError(
            f"Does not support key cache of type fp8 with head_size {head_size}"
        )
539
    from vllm.platforms import current_platform
540

541
    current_platform.seed_everything(seed)
542

543
    dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
544
545

    scale = head_size**-0.5
546
    x = 16 // torch.tensor([], dtype=dtype).element_size()
547
    key_cache_shape = (num_blocks, num_heads, head_size // x, block_size, x)
548
    key_caches: list[torch.Tensor] = []
549
    for _ in range(num_layers):
550
        key_cache = torch.empty(size=key_cache_shape, dtype=dtype, device=device)
551
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
552
            key_cache.uniform_(-scale, scale)
553
        elif cache_dtype == "fp8":
554
            _generate_random_fp8(key_cache, -scale, scale)
555
        else:
556
            raise ValueError(f"Does not support key cache of type {cache_dtype}")
557
558
559
        key_caches.append(key_cache)

    value_cache_shape = (num_blocks, num_heads, head_size, block_size)
560
    value_caches: list[torch.Tensor] = []
561
    for _ in range(num_layers):
562
        value_cache = torch.empty(size=value_cache_shape, dtype=dtype, device=device)
563
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
564
            value_cache.uniform_(-scale, scale)
565
        elif cache_dtype == "fp8":
566
            _generate_random_fp8(value_cache, -scale, scale)
567
        else:
568
            raise ValueError(f"Does not support value cache of type {cache_dtype}")
569
570
        value_caches.append(value_cache)
    return key_caches, value_caches
571
572


573
@cache
574
def is_pin_memory_available() -> bool:
575
    from vllm.platforms import current_platform
576

577
    return current_platform.is_pin_memory_available()
578
579


580
581
582
583
584
585
586
587
@cache
def is_uva_available() -> bool:
    """Check if Unified Virtual Addressing (UVA) is available."""
    # UVA requires pinned memory.
    # TODO: Add more requirements for UVA if needed.
    return is_pin_memory_available()


588
class DeviceMemoryProfiler:
589
    def __init__(self, device: torch.types.Device | None = None):
590
591
592
593
        self.device = device

    def current_memory_usage(self) -> float:
        # Return the memory usage in bytes.
594
        from vllm.platforms import current_platform
595

596
        gc.collect()
597
        return current_platform.get_current_memory_usage(self.device)
598
599
600
601
602
603
604
605
606
607
608
609

    def __enter__(self):
        self.initial_memory = self.current_memory_usage()
        # This allows us to call methods of the context manager if needed
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.final_memory = self.current_memory_usage()
        self.consumed_memory = self.final_memory - self.initial_memory

        # Force garbage collection
        gc.collect()
610
611


612
def make_ndarray_with_pad(
613
    x: list[list[T]],
614
615
616
    pad: T,
    dtype: npt.DTypeLike,
    *,
617
    max_len: int | None = None,
618
619
620
) -> npt.NDArray:
    """
    Make a padded array from 2D inputs.
621
622
623
624

    The padding is applied to the end of each inner list until it reaches
    `max_len`.
    """
625
626
627
628
629
    if max_len is None:
        # Unlike for most functions, map is faster than a genexpr over `len`
        max_len = max(map(len, x), default=0)

    padded_x = np.full((len(x), max_len), pad, dtype=dtype)
630
631
    for ind, blocktb in enumerate(x):
        assert len(blocktb) <= max_len
632
        padded_x[ind, : len(blocktb)] = blocktb
633
634
635
636
637

    return padded_x


def make_tensor_with_pad(
638
    x: list[list[T]],
639
640
641
    pad: T,
    dtype: torch.dtype,
    *,
642
    max_len: int | None = None,
643
    device: str | torch.device | None = None,
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
    pin_memory: bool = False,
) -> torch.Tensor:
    """
    Make a padded tensor from 2D inputs.

    The padding is applied to the end of each inner list until it reaches
    `max_len`.
    """
    np_dtype = TORCH_DTYPE_TO_NUMPY_DTYPE[dtype]
    padded_x = make_ndarray_with_pad(x, pad, np_dtype, max_len=max_len)

    tensor = torch.from_numpy(padded_x).to(device)
    if pin_memory:
        tensor = tensor.pin_memory()

    return tensor
660
661
662
663
664


def async_tensor_h2d(
    data: list,
    dtype: torch.dtype,
665
    target_device: str | torch.device,
666
667
668
669
670
671
672
    pin_memory: bool,
) -> torch.Tensor:
    """Asynchronously create a tensor and copy it from host to device."""
    t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu")
    return t.to(device=target_device, non_blocking=True)


673
674
675
676
677
def get_dtype_size(dtype: torch.dtype) -> int:
    """Get the size of the data type in bytes."""
    return torch.tensor([], dtype=dtype).element_size()


678
679
680
# bool = 0, int = 1, float = 2, complex = 3
def _get_precision_level(dtype: torch.dtype) -> int:
    # NOTE: Complex dtypes return `is_floating_point=False`
681
    return (dtype != torch.bool) + dtype.is_floating_point + dtype.is_complex * 2
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708


def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype):
    """
    Test whether it is lossless to cast a tensor from
    `src_dtype` to `tgt_dtype`.
    """
    if src_dtype == tgt_dtype:
        return True

    src_level = _get_precision_level(src_dtype)
    tgt_level = _get_precision_level(tgt_dtype)

    if src_level < tgt_level:
        return True
    if src_level > tgt_level:
        return False

    # Compare integral types
    if not src_dtype.is_floating_point and not src_dtype.is_complex:
        src_info = torch.iinfo(src_dtype)
        tgt_info = torch.iinfo(tgt_dtype)
        return src_info.min >= tgt_info.min and src_info.max <= tgt_info.max

    # Compare floating-point types
    src_info = torch.finfo(src_dtype)
    tgt_info = torch.finfo(tgt_dtype)
709
710
711
712
713
    return (
        src_info.min >= tgt_info.min
        and src_info.max <= tgt_info.max
        and src_info.resolution >= tgt_info.resolution
    )
714
715
716
717
718
719
720
721
722
723
724
725
726


def common_broadcastable_dtype(dtypes: Collection[torch.dtype]):
    """
    Get the common `dtype` where all of the other `dtypes` can be
    cast to it without losing any information.
    """
    return max(
        dtypes,
        key=lambda dtype: sum(is_lossless_cast(dt, dtype) for dt in dtypes),
    )


727
728
# TODO: This function can be removed if transformer_modules classes are
# serialized by value when communicating between processes
729
def init_cached_hf_modules() -> None:
730
731
732
733
    """
    Lazy initialization of the Hugging Face modules.
    """
    from transformers.dynamic_module_utils import init_hf_modules
734

735
    init_hf_modules()
736
737


738
@cache
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
def find_library(lib_name: str) -> str:
    """
    Find the library file in the system.
    `lib_name` is full filename, with both prefix and suffix.
    This function resolves `lib_name` to the full path of the library.
    """
    # Adapted from https://github.com/openai/triton/blob/main/third_party/nvidia/backend/driver.py#L19 # noqa
    # According to https://en.wikipedia.org/wiki/Filesystem_Hierarchy_Standard
    # `/sbin/ldconfig` should exist in all Linux systems.
    # `/sbin/ldconfig` searches the library in the system
    libs = subprocess.check_output(["/sbin/ldconfig", "-p"]).decode()
    # each line looks like the following:
    # libcuda.so.1 (libc6,x86-64) => /lib/x86_64-linux-gnu/libcuda.so.1
    locs = [line.split()[-1] for line in libs.splitlines() if lib_name in line]
    # `LD_LIBRARY_PATH` searches the library in the user-defined paths
754
    env_ld_library_path = envs.LD_LIBRARY_PATH
755
756
757
758
759
760
761
762
763
764
765
    if not locs and env_ld_library_path:
        locs = [
            os.path.join(dir, lib_name)
            for dir in env_ld_library_path.split(":")
            if os.path.exists(os.path.join(dir, lib_name))
        ]
    if not locs:
        raise ValueError(f"Cannot find {lib_name} in the system.")
    return locs[0]


766
def find_nccl_library() -> str:
767
768
769
770
771
772
    """
    We either use the library file specified by the `VLLM_NCCL_SO_PATH`
    environment variable, or we find the library file brought by PyTorch.
    After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be
    found by `ctypes` automatically.
    """
773
    so_file = envs.VLLM_NCCL_SO_PATH
774
775
776
777

    # manually load the nccl library
    if so_file:
        logger.info(
778
779
            "Found nccl from environment variable VLLM_NCCL_SO_PATH=%s", so_file
        )
780
781
    else:
        if torch.version.cuda is not None:
782
            so_file = "libnccl.so.2"
783
        elif torch.version.hip is not None:
784
            so_file = "librccl.so.1"
785
786
        else:
            raise ValueError("NCCL only supports CUDA and ROCm backends.")
787
        logger.debug_once("Found nccl from library %s", so_file)
788
    return so_file
789
790


791
def find_nccl_include_paths() -> list[str] | None:
792
793
    """
    We either use the nccl.h specified by the `VLLM_NCCL_INCLUDE_PATH`
794
795
    environment variable, or we find the library file brought by
    nvidia-nccl-cuXX. load_inline by default uses
796
797
798
799
800
801
802
803
804
    torch.utils.cpp_extension.include_paths
    """
    paths: list[str] = []
    inc = envs.VLLM_NCCL_INCLUDE_PATH
    if inc and os.path.isdir(inc):
        paths.append(inc)

    try:
        import importlib.util
805

806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
        spec = importlib.util.find_spec("nvidia.nccl")
        if spec and getattr(spec, "submodule_search_locations", None):
            for loc in spec.submodule_search_locations:
                inc_dir = os.path.join(loc, "include")
                if os.path.exists(os.path.join(inc_dir, "nccl.h")):
                    paths.append(inc_dir)
    except Exception:
        pass

    seen = set()
    out: list[str] = []
    for p in paths:
        if p and p not in seen:
            out.append(p)
            seen.add(p)
    return out or None


youkaichao's avatar
youkaichao committed
824
825
prev_set_stream = torch.cuda.set_stream

826
_current_stream_tls = threading.local()
youkaichao's avatar
youkaichao committed
827
828
829


def _patched_set_stream(stream: torch.cuda.Stream) -> None:
830
    _current_stream_tls.value = stream
youkaichao's avatar
youkaichao committed
831
832
833
834
835
836
    prev_set_stream(stream)


torch.cuda.set_stream = _patched_set_stream


837
838
839
840
841
class _StreamPlaceholder:
    def __init__(self):
        self.synchronize = lambda: None


youkaichao's avatar
youkaichao committed
842
843
844
845
846
847
848
849
850
851
852
def current_stream() -> torch.cuda.Stream:
    """
    replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`.
    it turns out that `torch.cuda.current_stream()` is quite expensive,
    as it will construct a new stream object at each call.
    here we patch `torch.cuda.set_stream` to keep track of the current stream
    directly, so that we can avoid calling `torch.cuda.current_stream()`.

    the underlying hypothesis is that we do not call `torch._C._cuda_setStream`
    from C/C++ code.
    """
853
    from vllm.platforms import current_platform
854
855

    if not hasattr(_current_stream_tls, "value") or _current_stream_tls.value is None:
youkaichao's avatar
youkaichao committed
856
857
        # when this function is called before any stream is set,
        # we return the default stream.
858
859
860
        # On ROCm using the default 0 stream in combination with RCCL
        # is hurting performance. Therefore creating a dedicated stream
        # per process
861
        if current_platform.is_rocm():
862
863
            # torch.cuda.set_stream here is the alias of _pathed_set_stream
            torch.cuda.set_stream(torch.cuda.Stream())
864
865
866
867
868
869
870
871
872
        elif current_platform.is_cpu():
            _current_stream_tls.value = _StreamPlaceholder()
        else:
            current_stream = current_platform.current_stream
            if current_stream is not None:
                _current_stream_tls.value = current_stream()
            else:
                raise ValueError(
                    "Fail to set current stream, current platform "
873
874
                    "may not support current_stream with torch API"
                )
875
    return _current_stream_tls.value
youkaichao's avatar
youkaichao committed
876
877


878
def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
879
880
881
882
    """Set up function tracing for the current thread,
    if enabled via the VLLM_TRACE_FUNCTION environment variable
    """

883
    if envs.VLLM_TRACE_FUNCTION:
884
        tmp_dir = tempfile.gettempdir()
885
886
        # add username to tmp_dir to avoid permission issues
        tmp_dir = os.path.join(tmp_dir, getpass.getuser())
887
888
889
890
891
892
893
894
        filename = (
            f"VLLM_TRACE_FUNCTION_for_process_{os.getpid()}"
            f"_thread_{threading.get_ident()}_"
            f"at_{datetime.datetime.now()}.log"
        ).replace(" ", "_")
        log_path = os.path.join(
            tmp_dir, "vllm", f"vllm-instance-{vllm_config.instance_id}", filename
        )
895
896
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        enable_trace_function_call(log_path)
897
898


899
@lru_cache(maxsize=8)
900
def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int:
901
902
903
904
905
906
907
908
909
910
    # Note: cuda_visible_devices is not used, but we keep it as an argument for
    # LRU Cache purposes.

    # Code below is based on
    # https://github.com/pytorch/pytorch/blob/
    # c1cd946818442aca8c7f812b16d187ce1586c3bc/
    # torch/cuda/__init__.py#L831C1-L831C17
    import torch.cuda
    import torch.version

911
    from vllm.platforms import current_platform
912

913
914
    if not torch.cuda._is_compiled():
        return 0
915
    if current_platform.is_rocm():
916
917
        # ROCm uses amdsmi instead of nvml for stateless device count
        # This requires a sufficiently modern version of Torch 2.4.0
918
919
920
921
922
        raw_count = (
            torch.cuda._device_count_amdsmi()
            if (hasattr(torch.cuda, "_device_count_amdsmi"))
            else -1
        )
923
924
925
    else:
        raw_count = torch.cuda._device_count_nvml()
    r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count
926
927
928
929
930
931
    return r


def cuda_device_count_stateless() -> int:
    """Get number of CUDA devices, caching based on the value of
    CUDA_VISIBLE_DEVICES at the time of call.
932

933
934
935
936
937
938
939
    This should be used instead of torch.cuda.device_count()
    unless CUDA_VISIBLE_DEVICES has already been set to the desired
    value."""

    # This can be removed and simply replaced with torch.cuda.get_device_count
    # after https://github.com/pytorch/pytorch/pull/122815 is released.
    return _cuda_device_count_stateless(envs.CUDA_VISIBLE_DEVICES)
940
941


942
943
944
945
946
947
948
def cuda_is_initialized() -> bool:
    """Check if CUDA is initialized."""
    if not torch.cuda._is_compiled():
        return False
    return torch.cuda.is_initialized()


949
950
951
952
953
954
955
def xpu_is_initialized() -> bool:
    """Check if XPU is initialized."""
    if not torch.xpu._is_compiled():
        return False
    return torch.xpu.is_initialized()


956
957
958
def cuda_get_device_properties(
    device, names: Sequence[str], init_cuda=False
) -> tuple[Any, ...]:
959
960
961
962
963
964
965
966
967
    """Get specified CUDA device property values without initializing CUDA in
    the current process."""
    if init_cuda or cuda_is_initialized():
        props = torch.cuda.get_device_properties(device)
        return tuple(getattr(props, name) for name in names)

    # Run in subprocess to avoid initializing CUDA as a side effect.
    mp_ctx = multiprocessing.get_context("fork")
    with ProcessPoolExecutor(max_workers=1, mp_context=mp_ctx) as executor:
968
        return executor.submit(cuda_get_device_properties, device, names, True).result()
969
970


971
972
973
def weak_bind(
    bound_method: Callable[..., Any],
) -> Callable[..., None]:
974
975
976
977
978
979
980
981
982
983
984
985
986
    """Make an instance method that weakly references
    its associated instance and no-ops once that
    instance is collected."""
    ref = weakref.ref(bound_method.__self__)  # type: ignore[attr-defined]
    unbound = bound_method.__func__  # type: ignore[attr-defined]

    def weak_bound(*args, **kwargs) -> None:
        if inst := ref():
            unbound(inst, *args, **kwargs)

    return weak_bound


987
class StoreBoolean(Action):
988
989
990
991
992
993
    def __call__(self, parser, namespace, values, option_string=None):
        if values.lower() == "true":
            setattr(namespace, self.dest, True)
        elif values.lower() == "false":
            setattr(namespace, self.dest, False)
        else:
994
995
996
            raise ValueError(
                f"Invalid boolean value: {values}. Expected 'true' or 'false'."
            )
997
998


999
class SortedHelpFormatter(ArgumentDefaultsHelpFormatter, RawDescriptionHelpFormatter):
1000
1001
    """SortedHelpFormatter that sorts arguments by their option strings."""

1002
1003
1004
1005
1006
1007
1008
    def _split_lines(self, text, width):
        """
        1. Sentences split across lines have their single newlines removed.
        2. Paragraphs and explicit newlines are split into separate lines.
        3. Each line is wrapped to the specified width (width of terminal).
        """
        # The patterns also include whitespace after the newline
1009
1010
        single_newline = re.compile(r"(?<!\n)\n(?!\n)\s*")
        multiple_newlines = re.compile(r"\n{2,}\s*")
1011
        text = single_newline.sub(" ", text)
1012
1013
1014
        lines = re.split(multiple_newlines, text)
        return sum([textwrap.wrap(line, width) for line in lines], [])

1015
1016
    def add_arguments(self, actions):
        actions = sorted(actions, key=lambda x: x.option_strings)
1017
        super().add_arguments(actions)
1018
1019


1020
class FlexibleArgumentParser(ArgumentParser):
1021
1022
    """ArgumentParser that allows both underscore and dash in names."""

1023
    _deprecated: set[Action] = set()
1024
1025
1026
1027
1028
1029
1030
    _json_tip: str = (
        "When passing JSON CLI arguments, the following sets of arguments "
        "are equivalent:\n"
        '   --json-arg \'{"key1": "value1", "key2": {"key3": "value2"}}\'\n'
        "   --json-arg.key1 value1 --json-arg.key2.key3 value2\n\n"
        "Additionally, list elements can be passed individually using +:\n"
        '   --json-arg \'{"key4": ["value3", "value4", "value5"]}\'\n'
1031
1032
        "   --json-arg.key4+ value3 --json-arg.key4+='value4,value5'\n\n"
    )
1033
    _search_keyword: str | None = None
1034

1035
    def __init__(self, *args, **kwargs):
1036
1037
1038
        # Set the default "formatter_class" to SortedHelpFormatter
        if "formatter_class" not in kwargs:
            kwargs["formatter_class"] = SortedHelpFormatter
1039
1040
        # Pop kwarg "add_json_tip" to control whether to add the JSON tip
        self.add_json_tip = kwargs.pop("add_json_tip", True)
1041
1042
        super().__init__(*args, **kwargs)

1043
    if sys.version_info < (3, 13):
1044
        # Enable the deprecated kwarg for Python 3.12 and below
1045

1046
        def parse_known_args(self, args=None, namespace=None):
1047
1048
1049
1050
            if args is not None and "--disable-log-requests" in args:
                # Special case warning because the warning below won't trigger
                # if –-disable-log-requests because its value is default.
                logger.warning_once(
1051
1052
                    "argument '--disable-log-requests' is deprecated and "
                    "replaced with '--enable-log-requests'. This will be "
1053
1054
                    "removed in v0.12.0."
                )
1055
1056
            namespace, args = super().parse_known_args(args, namespace)
            for action in FlexibleArgumentParser._deprecated:
1057
1058
1059
1060
                if (
                    hasattr(namespace, dest := action.dest)
                    and getattr(namespace, dest) != action.default
                ):
1061
                    logger.warning_once("argument '%s' is deprecated", dest)
1062
1063
            return namespace, args

1064
1065
        def add_argument(self, *args, **kwargs):
            deprecated = kwargs.pop("deprecated", False)
1066
            action = super().add_argument(*args, **kwargs)
1067
1068
            if deprecated:
                FlexibleArgumentParser._deprecated.add(action)
1069
1070
            return action

1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
        class _FlexibleArgumentGroup(_ArgumentGroup):
            def add_argument(self, *args, **kwargs):
                deprecated = kwargs.pop("deprecated", False)
                action = super().add_argument(*args, **kwargs)
                if deprecated:
                    FlexibleArgumentParser._deprecated.add(action)
                return action

        def add_argument_group(self, *args, **kwargs):
            group = self._FlexibleArgumentGroup(self, *args, **kwargs)
            self._action_groups.append(group)
            return group
1083

1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
    def format_help(self):
        # Only use custom help formatting for bottom level parsers
        if self._subparsers is not None:
            return super().format_help()

        formatter = self._get_formatter()

        # Handle keyword search of the args
        if (search_keyword := self._search_keyword) is not None:
            # Normalise the search keyword
            search_keyword = search_keyword.lower().replace("_", "-")
            # Return full help if searching for 'all'
1096
            if search_keyword == "all":
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
                self.epilog = self._json_tip
                return super().format_help()

            # Return group help if searching for a group title
            for group in self._action_groups:
                if group.title and group.title.lower() == search_keyword:
                    formatter.start_section(group.title)
                    formatter.add_text(group.description)
                    formatter.add_arguments(group._group_actions)
                    formatter.end_section()
                    formatter.add_text(self._json_tip)
                    return formatter.format_help()

            # Return matched args if searching for an arg name
            matched_actions = []
            for group in self._action_groups:
                for action in group._group_actions:
                    # search option name
1115
1116
1117
                    if any(
                        search_keyword in opt.lower() for opt in action.option_strings
                    ):
1118
1119
                        matched_actions.append(action)
            if matched_actions:
1120
                formatter.start_section(f"Arguments matching '{search_keyword}'")
1121
1122
1123
1124
1125
1126
1127
1128
1129
                formatter.add_arguments(matched_actions)
                formatter.end_section()
                formatter.add_text(self._json_tip)
                return formatter.format_help()

            # No match found
            formatter.add_text(
                f"No group or arguments matching '{search_keyword}'.\n"
                "Use '--help' to see available groups or "
1130
1131
                "'--help=all' to see all available parameters."
            )
1132
1133
1134
            return formatter.format_help()

        # usage
1135
        formatter.add_usage(self.usage, self._actions, self._mutually_exclusive_groups)
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156

        # description
        formatter.add_text(self.description)

        # positionals, optionals and user-defined groups
        formatter.start_section("Config Groups")
        config_groups = ""
        for group in self._action_groups:
            if not group._group_actions:
                continue
            title = group.title
            description = group.description or ""
            config_groups += f"{title: <24}{description}\n"
        formatter.add_text(config_groups)
        formatter.end_section()

        # epilog
        formatter.add_text(self.epilog)

        # determine help from format above
        return formatter.format_help()
1157

1158
1159
1160
1161
1162
    def parse_args(  # type: ignore[override]
        self,
        args: list[str] | None = None,
        namespace: Namespace | None = None,
    ):
1163
1164
1165
        if args is None:
            args = sys.argv[1:]

1166
1167
        # Check for --model in command line arguments first
        if args and args[0] == "serve":
1168
1169
            try:
                model_idx = next(
1170
1171
1172
1173
                    i
                    for i, arg in enumerate(args)
                    if arg == "--model" or arg.startswith("--model=")
                )
1174
                logger.warning(
1175
1176
                    "With `vllm serve`, you should provide the model as a "
                    "positional argument or in a config file instead of via "
1177
                    "the `--model` option. "
1178
1179
                    "The `--model` option will be removed in v0.13."
                )
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201

                if args[model_idx] == "--model":
                    model_tag = args[model_idx + 1]
                    rest_start_idx = model_idx + 2
                else:
                    model_tag = args[model_idx].removeprefix("--model=")
                    rest_start_idx = model_idx + 1

                # Move <model> to the front, e,g:
                # [Before]
                # vllm serve -tp 2 --model <model> --enforce-eager --port 8001
                # [After]
                # vllm serve <model> -tp 2 --enforce-eager --port 8001
                args = [
                    "serve",
                    model_tag,
                    *args[1:model_idx],
                    *args[rest_start_idx:],
                ]
                print("args", args)
            except StopIteration:
                pass
1202

1203
        if "--config" in args:
1204
            args = self._pull_args_from_config(args)
1205

1206
1207
1208
1209
1210
1211
1212
        def repl(match: re.Match) -> str:
            """Replaces underscores with dashes in the matched string."""
            return match.group(0).replace("_", "-")

        # Everything between the first -- and the first .
        pattern = re.compile(r"(?<=--)[^\.]*")

1213
        # Convert underscores to dashes and vice versa in argument names
1214
        processed_args = list[str]()
1215
        for i, arg in enumerate(args):
1216
            if arg.startswith("--help="):
1217
                FlexibleArgumentParser._search_keyword = arg.split("=", 1)[-1].lower()
1218
                processed_args.append("--help")
1219
1220
1221
            elif arg.startswith("--"):
                if "=" in arg:
                    key, value = arg.split("=", 1)
1222
                    key = pattern.sub(repl, key, count=1)
1223
                    processed_args.append(f"{key}={value}")
1224
                else:
1225
1226
                    key = pattern.sub(repl, arg, count=1)
                    processed_args.append(key)
1227
            elif arg.startswith("-O") and arg != "-O" and arg[2] != ".":
1228
1229
                # allow -O flag to be used without space, e.g. -O3 or -Odecode
                # -O.<...> handled later
1230
1231
1232
                # also handle -O=<mode> here
                mode = arg[3:] if arg[2] == "=" else arg[2:]
                processed_args.append(f"-O.mode={mode}")
1233
1234
1235
1236
1237
            elif (
                arg == "-O"
                and i + 1 < len(args)
                and args[i + 1] in {"0", "1", "2", "3"}
            ):
1238
1239
                # Convert -O <n> to -O.mode <n>
                processed_args.append("-O.mode")
1240
1241
1242
            else:
                processed_args.append(arg)

1243
        def create_nested_dict(keys: list[str], value: str) -> dict[str, Any]:
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
            """Creates a nested dictionary from a list of keys and a value.

            For example, `keys = ["a", "b", "c"]` and `value = 1` will create:
            `{"a": {"b": {"c": 1}}}`
            """
            nested_dict: Any = value
            for key in reversed(keys):
                nested_dict = {key: nested_dict}
            return nested_dict

1254
1255
1256
        def recursive_dict_update(
            original: dict[str, Any],
            update: dict[str, Any],
1257
1258
1259
1260
1261
        ) -> set[str]:
            """Recursively updates a dictionary with another dictionary.
            Returns a set of duplicate keys that were overwritten.
            """
            duplicates = set[str]()
1262
1263
            for k, v in update.items():
                if isinstance(v, dict) and isinstance(original.get(k), dict):
1264
1265
1266
1267
                    nested_duplicates = recursive_dict_update(original[k], v)
                    duplicates |= {f"{k}.{d}" for d in nested_duplicates}
                elif isinstance(v, list) and isinstance(original.get(k), list):
                    original[k] += v
1268
                else:
1269
1270
                    if k in original:
                        duplicates.add(k)
1271
                    original[k] = v
1272
            return duplicates
1273

1274
1275
        delete = set[int]()
        dict_args = defaultdict[str, dict[str, Any]](dict)
1276
        duplicates = set[str]()
1277
        for i, processed_arg in enumerate(processed_args):
1278
1279
1280
1281
            if i in delete:  # skip if value from previous arg
                continue

            if processed_arg.startswith("-") and "." in processed_arg:
1282
                if "=" in processed_arg:
1283
                    processed_arg, value_str = processed_arg.split("=", 1)
1284
                    if "." not in processed_arg:
1285
                        # False positive, '.' was only in the value
1286
1287
                        continue
                else:
1288
                    value_str = processed_args[i + 1]
1289
                    delete.add(i + 1)
1290

1291
1292
1293
1294
                if processed_arg.endswith("+"):
                    processed_arg = processed_arg[:-1]
                    value_str = json.dumps(list(value_str.split(",")))

1295
                key, *keys = processed_arg.split(".")
1296
1297
1298
1299
1300
                try:
                    value = json.loads(value_str)
                except json.decoder.JSONDecodeError:
                    value = value_str

1301
1302
                # Merge all values with the same key into a single dict
                arg_dict = create_nested_dict(keys, value)
1303
1304
                arg_duplicates = recursive_dict_update(dict_args[key], arg_dict)
                duplicates |= {f"{key}.{d}" for d in arg_duplicates}
1305
1306
                delete.add(i)
        # Filter out the dict args we set to None
1307
        processed_args = [a for i, a in enumerate(processed_args) if i not in delete]
1308
1309
1310
        if duplicates:
            logger.warning("Found duplicate keys %s", ", ".join(duplicates))

1311
1312
1313
1314
1315
        # Add the dict args back as if they were originally passed as JSON
        for dict_arg, dict_value in dict_args.items():
            processed_args.append(dict_arg)
            processed_args.append(json.dumps(dict_value))

1316
        return super().parse_args(processed_args, namespace)
1317

1318
1319
1320
1321
    def check_port(self, value):
        try:
            value = int(value)
        except ValueError:
1322
            msg = "Port must be an integer"
1323
            raise ArgumentTypeError(msg) from None
1324
1325

        if not (1024 <= value <= 65535):
1326
            raise ArgumentTypeError("Port must be between 1024 and 65535")
1327
1328
1329

        return value

1330
    def _pull_args_from_config(self, args: list[str]) -> list[str]:
1331
1332
        """Method to pull arguments specified in the config file
        into the command-line args variable.
1333
1334

        The arguments in config file will be inserted between
1335
        the argument list.
1336

1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
        example:
        ```yaml
            port: 12323
            tensor-parallel-size: 4
        ```
        ```python
        $: vllm {serve,chat,complete} "facebook/opt-12B" \
            --config config.yaml -tp 2
        $: args = [
            "serve,chat,complete",
1347
1348
            "facebook/opt-12B",
            '--config', 'config.yaml',
1349
1350
1351
1352
            '-tp', '2'
        ]
        $: args = [
            "serve,chat,complete",
1353
1354
1355
            "facebook/opt-12B",
            '--port', '12323',
            '--tensor-parallel-size', '4',
1356
1357
1358
1359
1360
            '-tp', '2'
            ]
        ```

        Please note how the config args are inserted after the sub command.
1361
        this way the order of priorities is maintained when these are args
1362
1363
        parsed by super().
        """
1364
        assert args.count("--config") <= 1, "More than one config file specified!"
1365

1366
        index = args.index("--config")
1367
        if index == len(args) - 1:
1368
1369
1370
1371
            raise ValueError(
                "No config file specified! \
                             Please check your command-line arguments."
            )
1372
1373
1374

        file_path = args[index + 1]

1375
        config_args = self.load_config_file(file_path)
1376

1377
        # 0th index might be the sub command {serve,chat,complete,...}
1378
        # optionally followed by model_tag (only for serve)
1379
1380
1381
1382
        # followed by config args
        # followed by rest of cli args.
        # maintaining this order will enforce the precedence
        # of cli > config > defaults
1383
        if args[0].startswith("-"):
1384
            # No sub command (e.g., api_server entry point)
1385
            args = config_args + args[0:index] + args[index + 2 :]
1386
        elif args[0] == "serve":
1387
1388
            model_in_cli = len(args) > 1 and not args[1].startswith("-")
            model_in_config = any(arg == "--model" for arg in config_args)
1389
1390

            if not model_in_cli and not model_in_config:
1391
                raise ValueError(
1392
                    "No model specified! Please specify model either "
1393
1394
                    "as a positional argument or in a config file."
                )
1395
1396
1397

            if model_in_cli:
                # Model specified as positional arg, keep CLI version
1398
1399
1400
1401
1402
1403
1404
                args = (
                    [args[0]]
                    + [args[1]]
                    + config_args
                    + args[2:index]
                    + args[index + 2 :]
                )
1405
1406
            else:
                # No model in CLI, use config if available
1407
                args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
1408
        else:
1409
            args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
1410
1411
1412

        return args

1413
    def load_config_file(self, file_path: str) -> list[str]:
1414
        """Loads a yaml file and returns the key value pairs as a
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
        flattened list with argparse like pattern
        ```yaml
            port: 12323
            tensor-parallel-size: 4
        ```
        returns:
            processed_args: list[str] = [
                '--port': '12323',
                '--tensor-parallel-size': '4'
            ]
        """
1426
1427
        extension: str = file_path.split(".")[-1]
        if extension not in ("yaml", "yml"):
1428
1429
            raise ValueError(
                "Config file must be of a yaml/yml type.\
1430
1431
1432
                              %s supplied",
                extension,
            )
1433
1434

        # only expecting a flat dictionary of atomic types
1435
        processed_args: list[str] = []
1436

1437
        config: dict[str, int | str] = {}
1438
        try:
1439
            with open(file_path) as config_file:
1440
1441
1442
1443
                config = yaml.safe_load(config_file)
        except Exception as ex:
            logger.error(
                "Unable to read the config file at %s. \
1444
1445
1446
                Make sure path is correct",
                file_path,
            )
1447
1448
            raise ex

1449
        store_boolean_arguments = [
1450
            action.dest for action in self._actions if isinstance(action, StoreBoolean)
1451
1452
        ]

1453
        for key, value in config.items():
1454
1455
            if isinstance(value, bool) and key not in store_boolean_arguments:
                if value:
1456
                    processed_args.append("--" + key)
1457
1458
            elif isinstance(value, list):
                if value:
1459
                    processed_args.append("--" + key)
1460
1461
                    for item in value:
                        processed_args.append(str(item))
1462
            else:
1463
                processed_args.append("--" + key)
1464
                processed_args.append(str(value))
1465
1466
1467

        return processed_args

1468

1469
1470
1471
1472
1473
1474
# Using dynamo with vLLM doesn't really work well with PyTorch versions < 2.4.0.
# In particular, the FakeScalarType is not supported for earlier versions of
# PyTorch which breaks dynamo for any ops registered using ScalarType.
def supports_dynamo() -> bool:
    base_torch_version = Version(Version(torch.__version__).base_version)
    return base_torch_version >= Version("2.4.0")
1475
1476


1477
# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
1478
def supports_xccl() -> bool:
1479
1480
1481
    return (
        is_torch_equal_or_newer("2.8.0.dev") and torch.distributed.is_xccl_available()
    )
1482
1483


1484
1485
1486
1487
1488
1489
# Some backends use pytorch version < 2.4.0 which doesn't
# support `torch.library.custom_op`.
def supports_custom_op() -> bool:
    return hasattr(torch.library, "custom_op")


1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
class AtomicCounter:
    """An atomic, thread-safe counter"""

    def __init__(self, initial=0):
        """Initialize a new atomic counter to given initial value"""
        self._value = initial
        self._lock = threading.Lock()

    def inc(self, num=1):
        """Atomically increment the counter by num and return the new value"""
        with self._lock:
            self._value += num
            return self._value

    def dec(self, num=1):
        """Atomically decrement the counter by num and return the new value"""
        with self._lock:
            self._value -= num
            return self._value

    @property
    def value(self):
        return self._value
1513
1514


1515
def weak_ref_tensor(tensor: Any) -> Any:
1516
1517
1518
1519
1520
    """
    Create a weak reference to a tensor.
    The new tensor will share the same data as the original tensor,
    but will not keep the original tensor alive.
    """
1521
1522
1523
1524
    if isinstance(tensor, torch.Tensor):
        return torch.ops._C.weak_ref_tensor(tensor)
    else:
        return tensor
1525
1526
1527


def weak_ref_tensors(
1528
1529
1530
1531
1532
    tensors: torch.Tensor
    | list[torch.Tensor]
    | tuple[torch.Tensor]
    | IntermediateTensors,
) -> torch.Tensor | list[Any] | tuple[Any] | Any:
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
    """
    Convenience function to create weak references to tensors,
    for single tensor, list of tensors or tuple of tensors.
    """
    if isinstance(tensors, torch.Tensor):
        return weak_ref_tensor(tensors)
    if isinstance(tensors, list):
        return [weak_ref_tensor(t) for t in tensors]
    if isinstance(tensors, tuple):
        return tuple(weak_ref_tensor(t) for t in tensors)
1543
1544
1545

    # For IntermediateTensors used in pipeline parallelism
    from vllm.sequence import IntermediateTensors
1546

1547
    if isinstance(tensors, IntermediateTensors):
1548
1549
1550
        ret = IntermediateTensors(
            {key: weak_ref_tensor(val) for key, val in tensors.tensors.items()}
        )
1551
        return ret
1552
    raise ValueError("Invalid type for tensors")
1553
1554


1555
1556
1557
1558
1559
1560
1561
1562
def get_cuda_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:
    """
    Get a CUDA view of a CPU tensor using Unified Virtual Addressing (UVA).
    """
    assert cpu_tensor.is_pinned(), "CPU tensor must be pinned"
    return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)


1563
def import_from_path(module_name: str, file_path: str | os.PathLike):
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
    """
    Import a Python file according to its file path.

    Based on the official recipe:
    https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
    """
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    if spec is None:
        raise ModuleNotFoundError(f"No module named '{module_name}'")

    assert spec.loader is not None

    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


1582
@cache
1583
1584
1585
1586
1587
1588
1589
def get_vllm_optional_dependencies():
    metadata = importlib.metadata.metadata("vllm")
    requirements = metadata.get_all("Requires-Dist", [])
    extras = metadata.get_all("Provides-Extra", [])

    return {
        extra: [
1590
1591
            re.split(r";|>=|<=|==", req)[0]
            for req in requirements
1592
1593
1594
1595
1596
1597
            if req.endswith(f'extra == "{extra}"')
        ]
        for extra in extras
    }


1598
1599
1600
1601
1602
class _PlaceholderBase:
    """
    Disallows downstream usage of placeholder modules.

    We need to explicitly override each dunder method because
1603
1604
    [`__getattr__`][vllm.utils._PlaceholderBase.__getattr__]
    is not called when they are accessed.
1605

1606
1607
    Info:
        [Special method lookup](https://docs.python.org/3/reference/datamodel.html#special-lookup)
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
    """

    def __getattr__(self, key: str) -> Never:
        """
        The main class should implement this to throw an error
        for attribute accesses representing downstream usage.
        """
        raise NotImplementedError

    # [Basic customization]

    def __lt__(self, other: object):
        return self.__getattr__("__lt__")

    def __le__(self, other: object):
        return self.__getattr__("__le__")

    def __eq__(self, other: object):
        return self.__getattr__("__eq__")

    def __ne__(self, other: object):
        return self.__getattr__("__ne__")

    def __gt__(self, other: object):
        return self.__getattr__("__gt__")

    def __ge__(self, other: object):
        return self.__getattr__("__ge__")

    def __hash__(self):
        return self.__getattr__("__hash__")

    def __bool__(self):
        return self.__getattr__("__bool__")

    # [Callable objects]

    def __call__(self, *args: object, **kwargs: object):
        return self.__getattr__("__call__")

    # [Container types]

    def __len__(self):
        return self.__getattr__("__len__")

    def __getitem__(self, key: object):
        return self.__getattr__("__getitem__")

    def __setitem__(self, key: object, value: object):
        return self.__getattr__("__setitem__")

    def __delitem__(self, key: object):
        return self.__getattr__("__delitem__")

    # __missing__ is optional according to __getitem__ specification,
    # so it is skipped

    # __iter__ and __reversed__ have a default implementation
    # based on __len__ and __getitem__, so they are skipped.

    # [Numeric Types]

    def __add__(self, other: object):
        return self.__getattr__("__add__")

    def __sub__(self, other: object):
        return self.__getattr__("__sub__")

    def __mul__(self, other: object):
        return self.__getattr__("__mul__")

    def __matmul__(self, other: object):
        return self.__getattr__("__matmul__")

    def __truediv__(self, other: object):
        return self.__getattr__("__truediv__")

    def __floordiv__(self, other: object):
        return self.__getattr__("__floordiv__")

    def __mod__(self, other: object):
        return self.__getattr__("__mod__")

    def __divmod__(self, other: object):
        return self.__getattr__("__divmod__")

    def __pow__(self, other: object, modulo: object = ...):
        return self.__getattr__("__pow__")

    def __lshift__(self, other: object):
        return self.__getattr__("__lshift__")

    def __rshift__(self, other: object):
        return self.__getattr__("__rshift__")

    def __and__(self, other: object):
        return self.__getattr__("__and__")

    def __xor__(self, other: object):
        return self.__getattr__("__xor__")

    def __or__(self, other: object):
        return self.__getattr__("__or__")

    # r* and i* methods have lower priority than
    # the methods for left operand so they are skipped

    def __neg__(self):
        return self.__getattr__("__neg__")

    def __pos__(self):
        return self.__getattr__("__pos__")

    def __abs__(self):
        return self.__getattr__("__abs__")

    def __invert__(self):
        return self.__getattr__("__invert__")

    # __complex__, __int__ and __float__ have a default implementation
    # based on __index__, so they are skipped.

    def __index__(self):
        return self.__getattr__("__index__")

    def __round__(self, ndigits: object = ...):
        return self.__getattr__("__round__")

    def __trunc__(self):
        return self.__getattr__("__trunc__")

    def __floor__(self):
        return self.__getattr__("__floor__")

    def __ceil__(self):
        return self.__getattr__("__ceil__")

    # [Context managers]

    def __enter__(self):
        return self.__getattr__("__enter__")

    def __exit__(self, *args: object, **kwargs: object):
        return self.__getattr__("__exit__")


class PlaceholderModule(_PlaceholderBase):
1755
1756
1757
1758
    """
    A placeholder object to use when a module does not exist.

    This enables more informative errors when trying to access attributes
1759
    of a module that does not exist.
1760
    """
1761
1762
1763
1764
1765
1766

    def __init__(self, name: str) -> None:
        super().__init__()

        # Apply name mangling to avoid conflicting with module attributes
        self.__name = name
1767
1768
1769
1770
1771

    def placeholder_attr(self, attr_path: str):
        return _PlaceholderModuleAttr(self, attr_path)

    def __getattr__(self, key: str):
1772
        name = self.__name
1773
1774

        try:
1775
            importlib.import_module(name)
1776
1777
1778
1779
1780
1781
1782
1783
        except ImportError as exc:
            for extra, names in get_vllm_optional_dependencies().items():
                if name in names:
                    msg = f"Please install vllm[{extra}] for {extra} support"
                    raise ImportError(msg) from exc

            raise exc

1784
1785
1786
1787
        raise AssertionError(
            "PlaceholderModule should not be used "
            "when the original module can be imported"
        )
1788
1789


1790
1791
1792
1793
1794
1795
1796
class _PlaceholderModuleAttr(_PlaceholderBase):
    def __init__(self, module: PlaceholderModule, attr_path: str) -> None:
        super().__init__()

        # Apply name mangling to avoid conflicting with module attributes
        self.__module = module
        self.__attr_path = attr_path
1797
1798

    def placeholder_attr(self, attr_path: str):
1799
        return _PlaceholderModuleAttr(self.__module, f"{self.__attr_path}.{attr_path}")
1800
1801

    def __getattr__(self, key: str):
1802
        getattr(self.__module, f"{self.__attr_path}.{key}")
1803

1804
1805
1806
1807
        raise AssertionError(
            "PlaceholderModule should not be used "
            "when the original module can be imported"
        )
1808
1809


1810
1811
1812
1813
1814
# create a library to hold the custom op
vllm_lib = Library("vllm", "FRAGMENT")  # noqa


def direct_register_custom_op(
1815
1816
    op_name: str,
    op_func: Callable,
1817
1818
1819
1820
    mutates_args: list[str] | None = None,
    fake_impl: Callable | None = None,
    target_lib: Library | None = None,
    dispatch_key: str | None = None,
1821
    tags: tuple[torch.Tag, ...] = (),
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
):
    """
    `torch.library.custom_op` can have significant overhead because it
    needs to consider complicated dispatching logic. This function
    directly registers a custom op and dispatches it to the CUDA backend.
    See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
    for more details.

    By default, the custom op is registered to the vLLM library. If you
    want to register it to a different library, you can pass the library
    object to the `target_lib` argument.

    IMPORTANT: the lifetime of the operator is tied to the lifetime of the
    library object. If you want to bind the operator to a different library,
    make sure the library object is alive when the operator is used.
    """
1838
    if not supports_custom_op():
1839
        from vllm.platforms import current_platform
1840

1841
1842
1843
1844
1845
        assert not current_platform.is_cuda_alike(), (
            "cuda platform needs torch>=2.4 to support custom op, "
            "chances are you are using an old version of pytorch "
            "or a custom build of pytorch. It is recommended to "
            "use vLLM in a fresh new environment and let it install "
1846
1847
            "the required dependencies."
        )
1848
1849
        return

1850
1851
1852
1853
1854
    if mutates_args is None:
        mutates_args = []

    if dispatch_key is None:
        from vllm.platforms import current_platform
1855

1856
1857
        dispatch_key = current_platform.dispatch_key

1858
    import torch.library
1859

1860
    if hasattr(torch.library, "infer_schema"):
1861
        schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args)
1862
1863
1864
    else:
        # for pytorch 2.4
        import torch._custom_op.impl
1865

1866
        schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
1867
    my_lib = target_lib or vllm_lib
1868
    my_lib.define(op_name + schema_str, tags=tags)
1869
    my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
1870
1871
    if fake_impl is not None:
        my_lib._register_fake(op_name, fake_impl)
1872
1873
1874
1875


def resolve_obj_by_qualname(qualname: str) -> Any:
    """
1876
    Resolve an object by its fully-qualified class name.
1877
1878
1879
1880
    """
    module_name, obj_name = qualname.rsplit(".", 1)
    module = importlib.import_module(module_name)
    return getattr(module, obj_name)
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905


def kill_process_tree(pid: int):
    """
    Kills all descendant processes of the given pid by sending SIGKILL.

    Args:
        pid (int): Process ID of the parent process
    """
    try:
        parent = psutil.Process(pid)
    except psutil.NoSuchProcess:
        return

    # Get all children recursively
    children = parent.children(recursive=True)

    # Send SIGKILL to all children first
    for child in children:
        with contextlib.suppress(ProcessLookupError):
            os.kill(child.pid, signal.SIGKILL)

    # Finally kill the parent
    with contextlib.suppress(ProcessLookupError):
        os.kill(pid, signal.SIGKILL)
1906
1907
1908
1909
1910


@dataclass
class MemorySnapshot:
    """Memory snapshot."""
1911

1912
    torch_peak: int = 0
1913
1914
    free_memory: int = 0
    total_memory: int = 0
1915
1916
1917
    cuda_memory: int = 0
    torch_memory: int = 0
    non_torch_memory: int = 0
1918
    timestamp: float = 0.0
1919
1920
1921
1922
1923
    auto_measure: bool = True

    def __post_init__(self):
        if self.auto_measure:
            self.measure()
1924
1925

    def measure(self):
1926
1927
        from vllm.platforms import current_platform

1928
1929
1930
1931
1932
        # we measure the torch peak memory usage via allocated_bytes,
        # rather than `torch.cuda.memory_reserved()` .
        # After `torch.cuda.reset_peak_memory_stats()`,
        # `torch.cuda.memory_reserved()` will keep growing, and only shrink
        # when we call `torch.cuda.empty_cache()` or OOM happens.
1933
        self.torch_peak = torch.cuda.memory_stats().get("allocated_bytes.all.peak", 0)
1934

1935
        self.free_memory, self.total_memory = torch.cuda.mem_get_info()
1936
1937
1938
1939
1940
        shared_sysmem_device_mem_sms = ((8, 7), (11, 0), (12, 1))  # Orin, Thor, Spark
        if (
            current_platform.is_cuda()
            and current_platform.get_device_capability() in shared_sysmem_device_mem_sms
        ):
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
            # On UMA (Orin, Thor and Spark) platform,
            # where both CPU and GPU rely on system memory,
            # the cudaMemGetInfo function shows the amount of free system memory
            # rather than what’s actually available.
            # In the case,
            # torch.cuda.mem_get_info() only reports "free" memory,
            # which can be lower than what is actually
            # available due to not including cache memory.
            # There’s also a comprehensive reference page
            # that explains how you can compute the proper value yourself.
            # https://docs.nvidia.com/cuda/cuda-for-tegra-appnote/#estimating-total-allocatable-device-memory-on-an-integrated-gpu-device
            self.free_memory = psutil.virtual_memory().available

1954
        self.cuda_memory = self.total_memory - self.free_memory
1955

1956
1957
        # torch.cuda.memory_reserved() is how many bytes
        # PyTorch gets from cuda (by calling cudaMalloc, etc.)
1958
1959
1960
1961
        # this is used to measure the non-torch memory usage
        self.torch_memory = torch.cuda.memory_reserved()

        self.non_torch_memory = self.cuda_memory - self.torch_memory
1962
1963
        self.timestamp = time.time()

1964
    def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot":
1965
        return MemorySnapshot(
1966
            torch_peak=self.torch_peak - other.torch_peak,
1967
1968
            free_memory=self.free_memory - other.free_memory,
            total_memory=self.total_memory - other.total_memory,
1969
1970
1971
1972
1973
1974
            cuda_memory=self.cuda_memory - other.cuda_memory,
            torch_memory=self.torch_memory - other.torch_memory,
            non_torch_memory=self.non_torch_memory - other.non_torch_memory,
            timestamp=self.timestamp - other.timestamp,
            auto_measure=False,
        )
1975
1976
1977
1978


@dataclass
class MemoryProfilingResult:
1979
1980
    """Memory profiling result. All numbers are in bytes."""

1981
1982
1983
1984
1985
    non_kv_cache_memory: int = 0
    torch_peak_increase: int = 0
    non_torch_increase: int = 0
    weights_memory: float = 0
    before_create: MemorySnapshot = field(default_factory=MemorySnapshot)
1986
1987
1988
1989
    before_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
    after_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
    profile_time: float = 0.0

1990
    def __repr__(self) -> str:
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
        return (
            f"Memory profiling takes {self.profile_time:.2f} seconds. "
            f"Total non KV cache memory: "
            f"{(self.non_kv_cache_memory / GiB_bytes):.2f}GiB; "
            f"torch peak memory increase: "
            f"{(self.torch_peak_increase / GiB_bytes):.2f}GiB; "
            f"non-torch forward increase memory: "
            f"{(self.non_torch_increase / GiB_bytes):.2f}GiB; "
            f"weights memory: {(self.weights_memory / GiB_bytes):.2f}GiB."
        )
2001

2002
2003
2004

@contextlib.contextmanager
def memory_profiling(
2005
2006
    baseline_snapshot: MemorySnapshot, weights_memory: int
) -> Generator[MemoryProfilingResult, None, None]:
2007
    """Memory profiling context manager.
2008
2009
    baseline_snapshot: the memory snapshot before the current vLLM instance.
    weights_memory: memory used by PyTorch when loading the model weights.
2010
2011
        Note that, before loading the model weights, we also initialize the device
        and distributed environment, which may consume some memory. This part is not
2012
        included in the weights_memory because PyTorch does not control it.
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046

    The memory in one GPU can be classified into 3 categories:
    1. memory used by anything other than the current vLLM instance.
    2. memory used by torch in the current vLLM instance.
    3. memory used in the current vLLM instance, but not by torch.

    A quantitive example:

    Before creating the current vLLM instance:
        category 1: 1 GiB
        category 2: 0 GiB
        category 3: 0 GiB

    After creating the current vLLM instance and loading the model,
    (i.e. before profiling):
        category 1: 1 GiB
        category 2: 2 GiB (model weights take 2 GiB)
        category 3: 0.5 GiB (memory used by NCCL)

    During profiling (peak):
        category 1: 1 GiB
        category 2: 4 GiB (peak activation tensors take 2 GiB)
        category 3: 1 GiB (memory used by NCCL + buffers for some attention backends)

    After profiling:
        category 1: 1 GiB
        category 2: 3 GiB (after garbage-collecting activation tensors)
        category 3: 1 GiB (memory used by NCCL + buffers for some attention backends)

    In this case, non-kv cache takes 5 GiB in total, including:
    a. 2 GiB used by the model weights (category 2)
    b. 2 GiB reserved for the peak activation tensors (category 2)
    c. 1 GiB used by non-torch components (category 3)

2047
    The memory used for loading weights (a.) is directly given from the argument `weights_memory`.
2048

2049
    The increase of `torch.cuda.memory_stats()["allocated_bytes.all.peak"]` during profiling gives (b.).
2050

2051
    The increase of `non_torch_memory` from creating the current vLLM instance until after profiling to get (c.).
2052
    """  # noqa
2053
2054
    gc.collect()
    torch.cuda.empty_cache()
2055
2056
2057
2058
    torch.cuda.reset_peak_memory_stats()

    result = MemoryProfilingResult()

2059
    result.before_create = baseline_snapshot
2060
    # the part of memory used for holding the model weights
2061
    result.weights_memory = weights_memory
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071

    result.before_profile.measure()

    yield result

    gc.collect()
    torch.cuda.empty_cache()

    result.after_profile.measure()

2072
2073
2074
2075
2076
    diff_profile = result.after_profile - result.before_profile
    diff_from_create = result.after_profile - result.before_create
    result.torch_peak_increase = diff_profile.torch_peak
    result.non_torch_increase = diff_from_create.non_torch_memory
    result.profile_time = diff_profile.timestamp
2077
2078
2079

    non_torch_memory = result.non_torch_increase
    peak_activation_memory = result.torch_peak_increase
2080
2081
2082
    result.non_kv_cache_memory = (
        non_torch_memory + peak_activation_memory + result.weights_memory
    )  # noqa
2083
2084


2085
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 # noqa: E501
2086
def set_ulimit(target_soft_limit=65535):
2087
    if sys.platform.startswith("win"):
2088
2089
2090
2091
        logger.info("Windows detected, skipping ulimit adjustment.")
        return

    import resource
2092

2093
2094
2095
2096
2097
    resource_type = resource.RLIMIT_NOFILE
    current_soft, current_hard = resource.getrlimit(resource_type)

    if current_soft < target_soft_limit:
        try:
2098
            resource.setrlimit(resource_type, (target_soft_limit, current_hard))
2099
2100
        except ValueError as e:
            logger.warning(
2101
2102
                "Found ulimit of %s and failed to automatically increase "
                "with error %s. This can cause fd limit errors like "
2103
                "`OSError: [Errno 24] Too many open files`. Consider "
2104
2105
2106
2107
                "increasing with ulimit -n",
                current_soft,
                e,
            )
2108
2109
2110
2111
2112
2113
2114
2115
2116


# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/utils.py#L28 # noqa: E501
def get_exception_traceback():
    etype, value, tb = sys.exc_info()
    err_str = "".join(traceback.format_exception(etype, value, tb))
    return err_str


2117
def split_zmq_path(path: str) -> tuple[str, str, str]:
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
    """Split a zmq path into its parts."""
    parsed = urlparse(path)
    if not parsed.scheme:
        raise ValueError(f"Invalid zmq path: {path}")

    scheme = parsed.scheme
    host = parsed.hostname or ""
    port = str(parsed.port or "")

    if scheme == "tcp" and not all((host, port)):
        # The host and port fields are required for tcp
        raise ValueError(f"Invalid zmq path: {path}")

    if scheme != "tcp" and port:
        # port only makes sense with tcp
        raise ValueError(f"Invalid zmq path: {path}")

    return scheme, host, port


2138
def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str:
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
    """Make a ZMQ path from its parts.

    Args:
        scheme: The ZMQ transport scheme (e.g. tcp, ipc, inproc).
        host: The host - can be an IPv4 address, IPv6 address, or hostname.
        port: Optional port number, only used for TCP sockets.

    Returns:
        A properly formatted ZMQ path string.
    """
2149
    if port is None:
2150
2151
2152
2153
2154
2155
        return f"{scheme}://{host}"
    if is_valid_ipv6_address(host):
        return f"{scheme}://[{host}]:{port}"
    return f"{scheme}://{host}:{port}"


2156
2157
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501
def make_zmq_socket(
2158
    ctx: zmq.asyncio.Context | zmq.Context,  # type: ignore[name-defined]
2159
    path: str,
2160
    socket_type: Any,
2161
2162
2163
    bind: bool | None = None,
    identity: bytes | None = None,
    linger: int | None = None,
2164
) -> zmq.Socket | zmq.asyncio.Socket:  # type: ignore[name-defined]
2165
2166
2167
    """Make a ZMQ socket with the proper bind/connect semantics."""

    mem = psutil.virtual_memory()
2168
    socket = ctx.socket(socket_type)
2169
2170
2171
2172
2173
2174
2175
2176

    # Calculate buffer size based on system memory
    total_mem = mem.total / 1024**3
    available_mem = mem.available / 1024**3
    # For systems with substantial memory (>32GB total, >16GB available):
    # - Set a large 0.5GB buffer to improve throughput
    # For systems with less memory:
    # - Use system default (-1) to avoid excessive memory consumption
2177
    buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1
2178

2179
    if bind is None:
2180
        bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB)
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192

    if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER):
        socket.setsockopt(zmq.RCVHWM, 0)
        socket.setsockopt(zmq.RCVBUF, buf_size)

    if socket_type in (zmq.PUSH, zmq.DEALER, zmq.ROUTER):
        socket.setsockopt(zmq.SNDHWM, 0)
        socket.setsockopt(zmq.SNDBUF, buf_size)

    if identity is not None:
        socket.setsockopt(zmq.IDENTITY, identity)

2193
2194
2195
    if linger is not None:
        socket.setsockopt(zmq.LINGER, linger)

2196
2197
2198
    if socket_type == zmq.XPUB:
        socket.setsockopt(zmq.XPUB_VERBOSE, True)

2199
2200
2201
2202
2203
2204
    # Determine if the path is a TCP socket with an IPv6 address.
    # Enable IPv6 on the zmq socket if so.
    scheme, host, _ = split_zmq_path(path)
    if scheme == "tcp" and is_valid_ipv6_address(host):
        socket.setsockopt(zmq.IPV6, 1)

2205
    if bind:
2206
        socket.bind(path)
2207
    else:
2208
        socket.connect(path)
2209
2210
2211
2212
2213

    return socket


@contextlib.contextmanager
2214
2215
2216
def zmq_socket_ctx(
    path: str,
    socket_type: Any,
2217
    bind: bool | None = None,
2218
    linger: int = 0,
2219
    identity: bytes | None = None,
2220
) -> Iterator[zmq.Socket]:
2221
2222
    """Context manager for a ZMQ socket"""

2223
    ctx = zmq.Context()  # type: ignore[attr-defined]
2224
    try:
2225
        yield make_zmq_socket(ctx, path, socket_type, bind=bind, identity=identity)
2226
2227
2228
2229
    except KeyboardInterrupt:
        logger.debug("Got Keyboard Interrupt.")

    finally:
2230
        ctx.destroy(linger=linger)
2231
2232


2233
2234
2235
2236
2237
2238
2239
def _maybe_force_spawn():
    """Check if we need to force the use of the `spawn` multiprocessing start
    method.
    """
    if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") == "spawn":
        return

2240
2241
    reasons = []
    if is_in_ray_actor():
2242
2243
2244
2245
        # even if we choose to spawn, we need to pass the ray address
        # to the subprocess so that it knows how to connect to the ray cluster.
        # env vars are inherited by subprocesses, even if we use spawn.
        import ray
2246

2247
        os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
2248
2249
2250
2251
2252
2253
        reasons.append("In a Ray actor and can only be spawned")

    if cuda_is_initialized():
        reasons.append("CUDA is initialized")
    elif xpu_is_initialized():
        reasons.append("XPU is initialized")
2254

2255
    if reasons:
2256
2257
2258
        logger.warning(
            "We must use the `spawn` multiprocessing start method. "
            "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "
2259
            "See https://docs.vllm.ai/en/latest/usage/"
2260
            "troubleshooting.html#python-multiprocessing "
2261
2262
2263
            "for more information. Reasons: %s",
            "; ".join(reasons),
        )
2264
2265
2266
2267
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"


def get_mp_context():
2268
2269
2270
2271
2272
2273
2274
    """Get a multiprocessing context with a particular method (spawn or fork).
    By default we follow the value of the VLLM_WORKER_MULTIPROC_METHOD to
    determine the multiprocessing method (default is fork). However, under
    certain conditions, we may enforce spawn and override the value of
    VLLM_WORKER_MULTIPROC_METHOD.
    """
    _maybe_force_spawn()
2275
2276
    mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
    return multiprocessing.get_context(mp_method)
2277
2278
2279


def bind_kv_cache(
2280
2281
    ctx: dict[str, Any],
    kv_cache: list[list[torch.Tensor]],  # [virtual_engine][layer_index]
2282
    shared_kv_cache_layers: dict[str, str] | None = None,
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
) -> None:
    # Bind the kv_cache tensor to Attention modules, similar to
    # ctx[layer_name].kv_cache[ve]=kv_cache[ve][extract_layer_index(layer_name)]
    # Special things handled here:
    # 1. Some models have non-attention layers, e.g., Jamba
    # 2. Pipeline parallelism, each rank only has a subset of layers
    # 3. Encoder attention has no kv cache
    # 4. Encoder-decoder models, encoder-decoder attention and decoder-only
    #    attention of the same layer (e.g., bart's decoder.layers.1.self_attn
    #    and decoder.layers.1.encoder_attn) is mapped to the same kv cache
    #    tensor
2294
2295
2296
2297
    # 5. Some models have attention layers that share kv cache with previous
    #    layers, this is specified through shared_kv_cache_layers
    if shared_kv_cache_layers is None:
        shared_kv_cache_layers = {}
2298
2299
    from vllm.attention import AttentionType
    from vllm.model_executor.models.utils import extract_layer_index
2300

2301
    layer_need_kv_cache = [
2302
2303
2304
2305
2306
2307
2308
2309
        layer_name
        for layer_name in ctx
        if (
            hasattr(ctx[layer_name], "attn_type")
            and ctx[layer_name].attn_type
            in (AttentionType.DECODER, AttentionType.ENCODER_DECODER)
        )
        and ctx[layer_name].kv_sharing_target_layer_name is None
2310
2311
    ]
    layer_index_sorted = sorted(
2312
2313
        set(extract_layer_index(layer_name) for layer_name in layer_need_kv_cache)
    )
2314
    for layer_name in layer_need_kv_cache:
2315
        kv_cache_idx = layer_index_sorted.index(extract_layer_index(layer_name))
2316
2317
2318
2319
        forward_ctx = ctx[layer_name]
        assert len(forward_ctx.kv_cache) == len(kv_cache)
        for ve, ve_kv_cache in enumerate(kv_cache):
            forward_ctx.kv_cache[ve] = ve_kv_cache[kv_cache_idx]
2320
2321
    if shared_kv_cache_layers is not None:
        for layer_name, target_layer_name in shared_kv_cache_layers.items():
2322
2323
2324
            assert extract_layer_index(target_layer_name) < extract_layer_index(
                layer_name
            ), "v0 doesn't support interleaving kv sharing"
2325
            ctx[layer_name].kv_cache = ctx[target_layer_name].kv_cache
2326
2327


2328
2329
def run_method(
    obj: Any,
2330
    method: str | bytes | Callable,
2331
2332
2333
    args: tuple[Any],
    kwargs: dict[str, Any],
) -> Any:
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
    """
    Run a method of an object with the given arguments and keyword arguments.
    If the method is string, it will be converted to a method using getattr.
    If the method is serialized bytes and will be deserialized using
    cloudpickle.
    If the method is a callable, it will be called directly.
    """
    if isinstance(method, bytes):
        func = partial(cloudpickle.loads(method), obj)
    elif isinstance(method, str):
        try:
            func = getattr(obj, method)
        except AttributeError:
2347
2348
2349
            raise NotImplementedError(
                f"Method {method!r} is not implemented."
            ) from None
2350
2351
2352
    else:
        func = partial(method, obj)  # type: ignore
    return func(*args, **kwargs)
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373


def import_pynvml():
    """
    Historical comments:

    libnvml.so is the library behind nvidia-smi, and
    pynvml is a Python wrapper around it. We use it to get GPU
    status without initializing CUDA context in the current process.
    Historically, there are two packages that provide pynvml:
    - `nvidia-ml-py` (https://pypi.org/project/nvidia-ml-py/): The official
        wrapper. It is a dependency of vLLM, and is installed when users
        install vLLM. It provides a Python module named `pynvml`.
    - `pynvml` (https://pypi.org/project/pynvml/): An unofficial wrapper.
        Prior to version 12.0, it also provides a Python module `pynvml`,
        and therefore conflicts with the official one. What's worse,
        the module is a Python package, and has higher priority than
        the official one which is a standalone Python file.
        This causes errors when both of them are installed.
        Starting from version 12.0, it migrates to a new module
        named `pynvml_utils` to avoid the conflict.
2374
2375
2376
2377
2378
2379
2380
    It is so confusing that many packages in the community use the
    unofficial one by mistake, and we have to handle this case.
    For example, `nvcr.io/nvidia/pytorch:24.12-py3` uses the unofficial
    one, and it will cause errors, see the issue
    https://github.com/vllm-project/vllm/issues/12847 for example.
    After all the troubles, we decide to copy the official `pynvml`
    module to our codebase, and use it directly.
2381
    """
2382
    import vllm.third_party.pynvml as pynvml
2383

2384
    return pynvml
2385
2386


2387
def warn_for_unimplemented_methods(cls: type[T]) -> type[T]:
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
    """
    A replacement for `abc.ABC`.
    When we use `abc.ABC`, subclasses will fail to instantiate
    if they do not implement all abstract methods.
    Here, we only require `raise NotImplementedError` in the
    base class, and log a warning if the method is not implemented
    in the subclass.
    """

    original_init = cls.__init__

    def find_unimplemented_methods(self: object):
        unimplemented_methods = []
        for attr_name in dir(self):
            # bypass inner method
2403
            if attr_name.startswith("_"):
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
                continue

            try:
                attr = getattr(self, attr_name)
                # get the func of callable method
                if callable(attr):
                    attr_func = attr.__func__
            except AttributeError:
                continue
            src = inspect.getsource(attr_func)
            if "NotImplementedError" in src:
                unimplemented_methods.append(attr_name)
        if unimplemented_methods:
2417
2418
            method_names = ",".join(unimplemented_methods)
            msg = f"Methods {method_names} not implemented in {self}"
2419
            logger.debug(msg)
2420
2421
2422
2423
2424
2425

    @wraps(original_init)
    def wrapped_init(self, *args, **kwargs) -> None:
        original_init(self, *args, **kwargs)
        find_unimplemented_methods(self)

2426
    type.__setattr__(cls, "__init__", wrapped_init)
2427
    return cls
2428
2429
2430
2431
2432
2433


class LazyLoader(types.ModuleType):
    """
    LazyLoader module borrowed from Tensorflow
    https://github.com/tensorflow/tensorflow/blob/main/tensorflow/python/util/lazy_loader.py
2434
    with an addition of "module caching".
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478

    Lazily import a module, mainly to avoid pulling in large dependencies.
    Modules such as `xgrammar` might do additional side effects, so we
    only want to use this when it is needed, delaying all eager effects
    """

    def __init__(
        self,
        local_name: str,
        parent_module_globals: dict[str, Any],
        name: str,
    ):
        self._local_name = local_name
        self._parent_module_globals = parent_module_globals
        self._module: types.ModuleType | None = None

        super().__init__(str(name))

    def _load(self) -> types.ModuleType:
        # Import the target module and insert it into the parent's namespace
        try:
            module = importlib.import_module(self.__name__)
            self._parent_module_globals[self._local_name] = module
            # The additional add to sys.modules
            # ensures library is actually loaded.
            sys.modules[self._local_name] = module
        except ModuleNotFoundError as err:
            raise err from None

        # Update this object's dict so that if someone keeps a
        # reference to the LazyLoader, lookups are efficient
        # (__getattr__ is only called on lookups that fail).
        self.__dict__.update(module.__dict__)
        return module

    def __getattr__(self, item: Any) -> Any:
        if self._module is None:
            self._module = self._load()
        return getattr(self._module, item)

    def __dir__(self) -> list[str]:
        if self._module is None:
            self._module = self._load()
        return dir(self._module)
2479
2480


2481
@contextlib.contextmanager
2482
def cprofile_context(save_file: str | None = None):
2483
2484
2485
2486
    """Run a cprofile

    Args:
        save_file: path to save the profile result. "1" or
2487
            None will result in printing to stdout.
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
    """
    import cProfile

    prof = cProfile.Profile()
    prof.enable()

    try:
        yield
    finally:
        prof.disable()
        if save_file and save_file != "1":
            prof.dump_stats(save_file)
        else:
            prof.print_stats(sort="cumtime")


2504
def cprofile(save_file: str | None = None, enabled: bool = True):
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
    """Decorator to profile a Python method using cProfile.

    Args:
        save_file: Path to save the profile result.
            If "1", None, or "", results will be printed to stdout.
        enabled: Set to false to turn this into a no-op
    """

    def decorator(func: Callable):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not enabled:
                # If profiling is disabled, just call the function directly.
                return func(*args, **kwargs)

            with cprofile_context(save_file):
                return func(*args, **kwargs)

        return wrapper

    return decorator
2526
2527


2528
2529
# Only relevant for models using ALiBi (e.g, MPT)
def check_use_alibi(model_config: ModelConfig) -> bool:
2530
    cfg = model_config.hf_text_config
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
    return (
        getattr(cfg, "alibi", False)  # Falcon
        or (
            "BloomForCausalLM" in getattr(model_config.hf_config, "architectures", [])
        )  # Bloom
        or getattr(cfg, "position_encoding_type", "") == "alibi"  # codellm_1b_alibi
        or (
            hasattr(cfg, "attn_config")  # MPT
            and (
                (
                    isinstance(cfg.attn_config, dict)
                    and cfg.attn_config.get("alibi", False)
                )
                or (
                    not isinstance(cfg.attn_config, dict)
                    and getattr(cfg.attn_config, "alibi", False)
                )
            )
        )
    )
2551
2552


2553
def sha256(input: Any) -> bytes:
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
    """Hash any picklable Python object using SHA-256.

    The input is serialized using pickle before hashing, which allows
    arbitrary Python objects to be used. Note that this function does
    not use a hash seed—if you need one, prepend it explicitly to the input.

    Args:
        input: Any picklable Python object.

    Returns:
2564
        Bytes representing the SHA-256 hash of the serialized input.
2565
2566
    """
    input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
2567
    return hashlib.sha256(input_bytes).digest()
2568
2569


2570
def sha256_cbor(input: Any) -> bytes:
2571
    """
2572
    Hash objects using CBOR serialization and SHA-256.
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582

    This option is useful for non-Python-dependent serialization and hashing.

    Args:
        input: Object to be serialized and hashed. Supported types include
            basic Python types and complex structures like lists, tuples, and
            dictionaries.
            Custom classes must implement CBOR serialization methods.

    Returns:
2583
        Bytes representing the SHA-256 hash of the CBOR serialized input.
2584
2585
    """
    input_bytes = cbor2.dumps(input, canonical=True)
2586
    return hashlib.sha256(input_bytes).digest()
2587
2588


2589
def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], bytes]:
2590
2591
2592
2593
2594
2595
2596
2597
2598
    """Get a hash function by name, or raise an error if
    the function is not found.
    Args:
        hash_fn_name: Name of the hash function.
    Returns:
        A hash function.
    """
    if hash_fn_name == "sha256":
        return sha256
2599
2600
    if hash_fn_name == "sha256_cbor":
        return sha256_cbor
2601
2602
2603
2604

    raise ValueError(f"Unsupported hash function: {hash_fn_name}")


2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
def is_torch_equal_or_newer(target: str) -> bool:
    """Check if the installed torch version is >= the target version.

    Args:
        target: a version string, like "2.6.0".

    Returns:
        Whether the condition meets.
    """
    try:
2615
        return _is_torch_equal_or_newer(str(torch.__version__), target)
2616
2617
    except Exception:
        # Fallback to PKG-INFO to load the package info, needed by the doc gen.
2618
        return Version(importlib.metadata.version("torch")) >= Version(target)
2619
2620
2621
2622
2623
2624


# Helper function used in testing.
def _is_torch_equal_or_newer(torch_version: str, target: str) -> bool:
    torch_version = version.parse(torch_version)
    return torch_version >= version.parse(target)
2625
2626


2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
def _is_torch_equal(target: str) -> bool:
    assert target.count(".") == 2
    torch_version = str(torch.__version__)
    torch_version = version.parse(torch_version)
    # torch version is like "2.6.0.dev20240101" or "2.6.0.dev20240101+cpu"
    # or "2.6.0+cu128" but never "2.6.0.1"
    return (
        torch_version >= version.parse(target)
        and version.parse(target + ".1") > torch_version
    )


def is_torch_equal(target: str) -> bool:
    """Check if the installed torch version is == the target version.

    Args:
        target: a version string, like "2.6.0".

    Returns:
        Whether the condition meets.
    """
    try:
        return _is_torch_equal(target)
    except Exception:
        return Version(importlib.metadata.version("torch")) == Version(target)


2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677
2678
@cache
def _has_module(module_name: str) -> bool:
    """Return True if *module_name* can be found in the current environment.

    The result is cached so that subsequent queries for the same module incur
    no additional overhead.
    """
    return importlib.util.find_spec(module_name) is not None


def has_pplx() -> bool:
    """Whether the optional `pplx_kernels` package is available."""

    return _has_module("pplx_kernels")


def has_deep_ep() -> bool:
    """Whether the optional `deep_ep` package is available."""

    return _has_module("deep_ep")


def has_deep_gemm() -> bool:
    """Whether the optional `deep_gemm` package is available."""

2679
    return _has_module("deep_gemm")
2680
2681


2682
2683
2684
2685
2686
2687
def has_triton_kernels() -> bool:
    """Whether the optional `triton_kernels` package is available."""

    return _has_module("triton_kernels")


2688
2689
2690
2691
2692
2693
def has_tilelang() -> bool:
    """Whether the optional `tilelang` package is available."""

    return _has_module("tilelang")


2694
2695
2696
def set_process_title(
    name: str, suffix: str = "", prefix: str = envs.VLLM_PROCESS_NAME_PREFIX
) -> None:
2697
2698
2699
    """
    Set the current process title to a specific name with an
    optional suffix.
2700
2701

    Args:
2702
        name: The title to assign to the current process.
2703
        suffix: An optional suffix to append to the base name.
2704
        prefix: A prefix to prepend to the front separated by `::`.
2705
2706
2707
    """
    if suffix:
        name = f"{name}_{suffix}"
2708
    setproctitle.setproctitle(f"{prefix}::{name}")
2709
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722


def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None:
    """Prepend each output line with process-specific prefix"""

    prefix = f"{CYAN}({worker_name} pid={pid}){RESET} "
    file_write = file.write

    def write_with_prefix(s: str):
        if not s:
            return
        if file.start_new_line:  # type: ignore[attr-defined]
            file_write(prefix)
        idx = 0
2723
        while (next_idx := s.find("\n", idx)) != -1:
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
            next_idx += 1
            file_write(s[idx:next_idx])
            if next_idx == len(s):
                file.start_new_line = True  # type: ignore[attr-defined]
                return
            file_write(prefix)
            idx = next_idx
        file_write(s[idx:])
        file.start_new_line = False  # type: ignore[attr-defined]

    file.start_new_line = True  # type: ignore[attr-defined]
    file.write = write_with_prefix  # type: ignore[method-assign]


2738
def decorate_logs(process_name: str | None = None) -> None:
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
    """
    Adds a process-specific prefix to each line of output written to stdout and
    stderr.

    This function is intended to be called before initializing the api_server,
    engine_core, or worker classes, so that all subsequent output from the
    process is prefixed with the process name and PID. This helps distinguish
    log output from different processes in multi-process environments.

    Args:
        process_name: Optional; the name of the process to use in the prefix.
            If not provided, the current process name from the multiprocessing
            context is used.
    """
    if process_name is None:
        process_name = get_mp_context().current_process().name
    pid = os.getpid()
    _add_prefix(sys.stdout, process_name, pid)
    _add_prefix(sys.stderr, process_name, pid)
2758
2759
2760


def length_from_prompt_token_ids_or_embeds(
2761
2762
    prompt_token_ids: list[int] | None,
    prompt_embeds: torch.Tensor | None,
2763
) -> int:
2764
    """Calculate the request length (in number of tokens) give either
2765
2766
    prompt_token_ids or prompt_embeds.
    """
2767
2768
    prompt_token_len = None if prompt_token_ids is None else len(prompt_token_ids)
    prompt_embeds_len = None if prompt_embeds is None else len(prompt_embeds)
2769
2770
2771

    if prompt_token_len is None:
        if prompt_embeds_len is None:
2772
            raise ValueError("Neither prompt_token_ids nor prompt_embeds were defined.")
2773
2774
        return prompt_embeds_len
    else:
2775
        if prompt_embeds_len is not None and prompt_embeds_len != prompt_token_len:
2776
2777
2778
            raise ValueError(
                "Prompt token ids and prompt embeds had different lengths"
                f" prompt_token_ids={prompt_token_len}"
2779
2780
                f" prompt_embeds={prompt_embeds_len}"
            )
2781
        return prompt_token_len
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794


@contextlib.contextmanager
def set_env_var(key, value):
    old = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:
        if old is None:
            del os.environ[key]
        else:
            os.environ[key] = old
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814


def unique_filepath(fn: Callable[[int], Path]) -> Path:
    """
    unique_filepath returns a unique path by trying
    to include an integer in increasing order.

    fn should be a callable that returns a path that
    includes the passed int at a fixed location.

    Note: This function has a TOCTOU race condition.
    Caller should use atomic operations (e.g., open with 'x' mode)
    when creating the file to ensure thread safety.
    """
    i = 0
    while True:
        p = fn(i)
        if not p.exists():
            return p
        i += 1