"vscode:/vscode.git/clone" did not exist on "ea5ff3a1f60e1b9f01af17260608009c184e7ff0"
__init__.py 83.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 inspect
12
import ipaddress
13
import json
14
import multiprocessing
15
import os
16
import pickle
17
import signal
18
import socket
19
import subprocess
20
import sys
21
import tempfile
22
import textwrap
23
import threading
24
import time
25
import traceback
Zhuohan Li's avatar
Zhuohan Li committed
26
import uuid
27
import warnings
28
import weakref
29
30
31
32
33
34
35
36
from argparse import (
    Action,
    ArgumentDefaultsHelpFormatter,
    ArgumentParser,
    ArgumentTypeError,
    RawDescriptionHelpFormatter,
    _ArgumentGroup,
)
37
from collections import defaultdict
38
from collections.abc import (
39
    Callable,
40
41
42
43
44
    Collection,
    Generator,
    Iterator,
    Sequence,
)
45
from concurrent.futures.process import ProcessPoolExecutor
46
from dataclasses import dataclass, field
47
from functools import cache, lru_cache, partial, wraps
48
from pathlib import Path
49
from typing import TYPE_CHECKING, Any, TextIO, TypeVar
50
from urllib.parse import urlparse
51
from uuid import uuid4
Zhuohan Li's avatar
Zhuohan Li committed
52

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

69
import vllm.envs as envs
70
from vllm.logger import enable_trace_function_call, init_logger
71
from vllm.ray.lazy_utils import is_in_ray_actor
72

73
if TYPE_CHECKING:
74
75
    from argparse import Namespace

76
    from vllm.config import ModelConfig, VllmConfig
77
    from vllm.sequence import IntermediateTensors
78
79
80
81
82
83
else:
    Namespace = object

    ModelConfig = object
    VllmConfig = object
    IntermediateTensors = object
84

85
86
logger = init_logger(__name__)

87
88
89
90
91
92
# 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

93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
# 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"

108
109
110
111
112
113
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)."""

114
115
116
GB_bytes = 1_000_000_000
"""The number of bytes in one gigabyte (GB)."""

117
118
119
GiB_bytes = 1 << 30
"""The number of bytes in one gibibyte (GiB)."""

120
# ANSI color codes
121
122
CYAN = "\033[1;36m"
RESET = "\033[0;0m"
123

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

137
138
139
140
141
142
143
144
145
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,
}

146
147
148
149
150
151
152
153
154
155

@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)


156
T = TypeVar("T")
157
U = TypeVar("U")
158

159

Woosuk Kwon's avatar
Woosuk Kwon committed
160
161
162
163
164
class Device(enum.Enum):
    GPU = enum.auto()
    CPU = enum.auto()


165
166
167
168
169
class LayerBlockType(enum.Enum):
    attention = "attention"
    mamba = "mamba"


Woosuk Kwon's avatar
Woosuk Kwon committed
170
171
172
173
class Counter:
    def __init__(self, start: int = 0) -> None:
        self.counter = start

Woosuk Kwon's avatar
Woosuk Kwon committed
174
    def __next__(self) -> int:
175
        i = self.counter
Woosuk Kwon's avatar
Woosuk Kwon committed
176
        self.counter += 1
177
        return i
Woosuk Kwon's avatar
Woosuk Kwon committed
178
179
180

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

182

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

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


195
def get_cpu_memory() -> int:
196
    """Returns the total CPU memory of the node in bytes."""
197
    return psutil.virtual_memory().total
Zhuohan Li's avatar
Zhuohan Li committed
198
199
200
201


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

203

204
def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]):
205
206
207
208
209
    for sock in sockets:
        if sock is not None:
            sock.close(linger=0)


210
def get_ip() -> str:
211
    host_ip = envs.VLLM_HOST_IP
212
213
214
215
    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"
216
            " interact with the container's network stack. Please "
217
            "use VLLM_HOST_IP instead to set the IP address for vLLM processes"
218
219
            " to communicate with each other."
        )
220
221
222
223
224
    if host_ip:
        return host_ip

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

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

    # try ipv6
    try:
235
        s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
236
237
238
        # 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
239
        return s.getsockname()[0]
240
241
242
243
244
    except Exception:
        pass

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


252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
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. "
276
277
            "Set the VLLM_LOOPBACK_IP environment variable explicitly."
        )
278
279


280
281
282
283
284
285
286
287
def is_valid_ipv6_address(address: str) -> bool:
    try:
        ipaddress.IPv6Address(address)
        return True
    except ValueError:
        return False


288
def split_host_port(host_port: str) -> tuple[str, int]:
289
    # ipv6
290
291
    if host_port.startswith("["):
        host, port = host_port.rsplit("]", 1)
292
        host = host[1:]
293
        port = port.split(":")[1]
294
295
        return host, int(port)
    else:
296
        host, port = host_port.split(":")
297
298
299
300
301
302
303
304
305
306
        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}"


307
def get_distributed_init_method(ip: str, port: int) -> str:
308
309
310
311
    return get_tcp_uri(ip, port)


def get_tcp_uri(ip: str, port: int) -> str:
312
313
314
315
    if is_valid_ipv6_address(ip):
        return f"tcp://[{ip}]:{port}"
    else:
        return f"tcp://{ip}:{port}"
316
317


318
319
320
321
322
def get_open_zmq_ipc_path() -> str:
    base_rpc_path = envs.VLLM_RPC_BASE_PATH
    return f"ipc://{base_rpc_path}/{uuid4()}"


323
324
325
326
def get_open_zmq_inproc_path() -> str:
    return f"inproc://{uuid4()}"


327
def get_open_port() -> int:
328
329
330
331
332
333
334
335
336
    """
    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:
337
338
        dp_master_port = envs.VLLM_DP_MASTER_PORT
        reserved_port_range = range(dp_master_port, dp_master_port + 10)
339
        while True:
340
341
342
            candidate_port = _get_open_port()
            if candidate_port not in reserved_port_range:
                return candidate_port
343
344
    return _get_open_port()

youkaichao's avatar
youkaichao committed
345

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


354
def _get_open_port() -> int:
355
    port = envs.VLLM_PORT
356
    if port is not None:
357
358
359
360
361
362
363
        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
364
                logger.info("Port %d is already in use, trying port %d", port - 1, port)
365
366
367
368
369
370
371
372
373
374
    # 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]
375
376


377
def find_process_using_port(port: int) -> psutil.Process | None:
378
379
380
381
382
383
384
    # 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

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


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


407
408
409
410
411
def cdiv(a: int, b: int) -> int:
    """Ceiling division."""
    return -(a // -b)


412
413
414
415
416
417
418
def next_power_of_2(n) -> int:
    """The next power of 2 (inclusive)"""
    if n < 1:
        return 1
    return 1 << (n - 1).bit_length()


419
420
421
422
423
424
425
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)


426
427
428
429
def round_up(x: int, y: int) -> int:
    return ((x + y - 1) // y) * y


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


434
def _generate_random_fp8(
435
    tensor: torch.Tensor,
436
437
438
439
440
441
    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.
442
    # For example, s.11111.00 in fp8e5m2 format represents Inf.
443
    #     | E4M3        | E5M2
444
    # -----|-------------|-------------------
445
446
    # Inf | N/A         | s.11111.00
    # NaN | s.1111.111  | s.11111.{01,10,11}
447
    from vllm import _custom_ops as ops
448

449
450
    tensor_tmp = torch.empty_like(tensor, dtype=torch.float16)
    tensor_tmp.uniform_(low, high)
451
    ops.convert_fp8(tensor, tensor_tmp)
452
453
454
    del tensor_tmp


455
def get_kv_cache_torch_dtype(
456
457
    cache_dtype: str | torch.dtype | None,
    model_dtype: str | torch.dtype | None = None,
458
) -> torch.dtype:
459
460
    if isinstance(cache_dtype, str):
        if cache_dtype == "auto":
461
            if isinstance(model_dtype, str) and model_dtype in STR_DTYPE_TO_TORCH_DTYPE:
462
463
464
465
466
                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}")
467
        elif cache_dtype in STR_DTYPE_TO_TORCH_DTYPE:
468
469
470
471
472
473
474
            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}")
475
476
477
478
479
480
481
482
483
    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,
484
485
    cache_dtype: str | torch.dtype | None,
    model_dtype: str | torch.dtype | None = None,
486
487
488
    seed: int | None = None,
    device: str | None = "cuda",
    cache_layout: str | None = "NHD",
489
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
490
    from vllm.platforms import current_platform
491

492
    current_platform.seed_everything(seed)
493

494
    dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
495
496
    generic_kv_cache_shape = (num_blocks, 2, block_size, num_heads, head_size)
    assert cache_layout in ("NHD", "HND")
497
    stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3, 2, 4)
498

499
    kv_cache_allocation_shape = tuple(generic_kv_cache_shape[i] for i in stride_order)
500
    scale = head_size**-0.5
501

502
503
    key_caches: list[torch.Tensor] = []
    value_caches: list[torch.Tensor] = []
504

505
    for _ in range(num_layers):
506
        key_value_cache = torch.empty(
507
            size=kv_cache_allocation_shape, dtype=dtype, device=device
508
        ).permute(*stride_order)
509
510
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
            key_value_cache.uniform_(-scale, scale)
511
        elif cache_dtype == "fp8":
512
513
            _generate_random_fp8(key_value_cache, -scale, scale)
        else:
514
            raise ValueError(f"Does not support key cache of type {cache_dtype}")
515
516
517
518
519
520
521
522
523
524
525
        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,
526
527
    cache_dtype: str | torch.dtype | None,
    model_dtype: str | torch.dtype | None = None,
528
529
    seed: int | None = None,
    device: str | None = "cuda",
530
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
Joe's avatar
Joe committed
531
532
533
534
    if cache_dtype == "fp8" and head_size % 16:
        raise ValueError(
            f"Does not support key cache of type fp8 with head_size {head_size}"
        )
535
    from vllm.platforms import current_platform
536

537
    current_platform.seed_everything(seed)
538

539
    dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
540
541

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

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


569
@cache
570
def is_pin_memory_available() -> bool:
571
    from vllm.platforms import current_platform
572

573
    return current_platform.is_pin_memory_available()
574
575


576
577
578
579
580
581
582
583
@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()


584
class DeviceMemoryProfiler:
585
    def __init__(self, device: torch.types.Device | None = None):
586
587
588
589
        self.device = device

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

592
        gc.collect()
593
        return current_platform.get_current_memory_usage(self.device)
594
595
596
597
598
599
600
601
602
603
604
605

    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()
606
607


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

    The padding is applied to the end of each inner list until it reaches
    `max_len`.
    """
621
622
623
624
625
    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)
626
627
    for ind, blocktb in enumerate(x):
        assert len(blocktb) <= max_len
628
        padded_x[ind, : len(blocktb)] = blocktb
629
630
631
632
633

    return padded_x


def make_tensor_with_pad(
634
    x: list[list[T]],
635
636
637
    pad: T,
    dtype: torch.dtype,
    *,
638
    max_len: int | None = None,
639
    device: str | torch.device | None = None,
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
    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
656
657
658
659
660


def async_tensor_h2d(
    data: list,
    dtype: torch.dtype,
661
    target_device: str | torch.device,
662
663
664
665
666
667
668
    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)


669
670
671
672
673
def get_dtype_size(dtype: torch.dtype) -> int:
    """Get the size of the data type in bytes."""
    return torch.tensor([], dtype=dtype).element_size()


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


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)
705
706
707
708
709
    return (
        src_info.min >= tgt_info.min
        and src_info.max <= tgt_info.max
        and src_info.resolution >= tgt_info.resolution
    )
710
711
712
713
714
715
716
717
718
719
720
721
722


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),
    )


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

731
    init_hf_modules()
732
733


734
@cache
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
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
750
    env_ld_library_path = envs.LD_LIBRARY_PATH
751
752
753
754
755
756
757
758
759
760
761
    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]


762
def find_nccl_library() -> str:
763
764
765
766
767
768
    """
    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.
    """
769
    so_file = envs.VLLM_NCCL_SO_PATH
770
771
772
773

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


787
def find_nccl_include_paths() -> list[str] | None:
788
789
    """
    We either use the nccl.h specified by the `VLLM_NCCL_INCLUDE_PATH`
790
791
    environment variable, or we find the library file brought by
    nvidia-nccl-cuXX. load_inline by default uses
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
    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:
        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
818
819
prev_set_stream = torch.cuda.set_stream

820
_current_stream_tls = threading.local()
youkaichao's avatar
youkaichao committed
821
822
823


def _patched_set_stream(stream: torch.cuda.Stream) -> None:
824
    _current_stream_tls.value = stream
youkaichao's avatar
youkaichao committed
825
826
827
828
829
830
    prev_set_stream(stream)


torch.cuda.set_stream = _patched_set_stream


831
832
833
834
835
class _StreamPlaceholder:
    def __init__(self):
        self.synchronize = lambda: None


youkaichao's avatar
youkaichao committed
836
837
838
839
840
841
842
843
844
845
846
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.
    """
847
    from vllm.platforms import current_platform
848
849

    if not hasattr(_current_stream_tls, "value") or _current_stream_tls.value is None:
youkaichao's avatar
youkaichao committed
850
851
        # when this function is called before any stream is set,
        # we return the default stream.
852
853
854
        # On ROCm using the default 0 stream in combination with RCCL
        # is hurting performance. Therefore creating a dedicated stream
        # per process
855
        if current_platform.is_rocm():
856
857
            # torch.cuda.set_stream here is the alias of _pathed_set_stream
            torch.cuda.set_stream(torch.cuda.Stream())
858
859
860
861
862
863
864
865
866
        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 "
867
868
                    "may not support current_stream with torch API"
                )
869
    return _current_stream_tls.value
youkaichao's avatar
youkaichao committed
870
871


872
def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
873
874
875
876
    """Set up function tracing for the current thread,
    if enabled via the VLLM_TRACE_FUNCTION environment variable
    """

877
    if envs.VLLM_TRACE_FUNCTION:
878
        tmp_dir = tempfile.gettempdir()
879
880
        # add username to tmp_dir to avoid permission issues
        tmp_dir = os.path.join(tmp_dir, getpass.getuser())
881
882
883
884
885
886
887
888
        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
        )
889
890
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        enable_trace_function_call(log_path)
891
892


893
@lru_cache(maxsize=8)
894
def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int:
895
896
897
898
899
900
901
902
903
904
    # 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

905
    from vllm.platforms import current_platform
906

907
908
    if not torch.cuda._is_compiled():
        return 0
909
    if current_platform.is_rocm():
910
911
        # ROCm uses amdsmi instead of nvml for stateless device count
        # This requires a sufficiently modern version of Torch 2.4.0
912
913
914
915
916
        raw_count = (
            torch.cuda._device_count_amdsmi()
            if (hasattr(torch.cuda, "_device_count_amdsmi"))
            else -1
        )
917
918
919
    else:
        raw_count = torch.cuda._device_count_nvml()
    r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count
920
921
922
923
924
925
    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.
926

927
928
929
930
931
932
933
    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)
934
935


936
937
938
939
940
941
942
def cuda_is_initialized() -> bool:
    """Check if CUDA is initialized."""
    if not torch.cuda._is_compiled():
        return False
    return torch.cuda.is_initialized()


943
944
945
946
947
948
949
def xpu_is_initialized() -> bool:
    """Check if XPU is initialized."""
    if not torch.xpu._is_compiled():
        return False
    return torch.xpu.is_initialized()


950
951
952
def cuda_get_device_properties(
    device, names: Sequence[str], init_cuda=False
) -> tuple[Any, ...]:
953
954
955
956
957
958
959
960
961
    """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:
962
        return executor.submit(cuda_get_device_properties, device, names, True).result()
963
964


965
966
967
def weak_bind(
    bound_method: Callable[..., Any],
) -> Callable[..., None]:
968
969
970
971
972
973
974
975
976
977
978
979
980
    """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


981
class StoreBoolean(Action):
982
983
984
985
986
987
    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:
988
989
990
            raise ValueError(
                f"Invalid boolean value: {values}. Expected 'true' or 'false'."
            )
991
992


993
class SortedHelpFormatter(ArgumentDefaultsHelpFormatter, RawDescriptionHelpFormatter):
994
995
    """SortedHelpFormatter that sorts arguments by their option strings."""

996
997
998
999
1000
1001
1002
    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
1003
1004
        single_newline = re.compile(r"(?<!\n)\n(?!\n)\s*")
        multiple_newlines = re.compile(r"\n{2,}\s*")
1005
        text = single_newline.sub(" ", text)
1006
1007
1008
        lines = re.split(multiple_newlines, text)
        return sum([textwrap.wrap(line, width) for line in lines], [])

1009
1010
    def add_arguments(self, actions):
        actions = sorted(actions, key=lambda x: x.option_strings)
1011
        super().add_arguments(actions)
1012
1013


1014
class FlexibleArgumentParser(ArgumentParser):
1015
1016
    """ArgumentParser that allows both underscore and dash in names."""

1017
    _deprecated: set[Action] = set()
1018
1019
1020
1021
1022
1023
1024
    _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'
1025
1026
        "   --json-arg.key4+ value3 --json-arg.key4+='value4,value5'\n\n"
    )
1027
    _search_keyword: str | None = None
1028

1029
    def __init__(self, *args, **kwargs):
1030
1031
1032
        # Set the default "formatter_class" to SortedHelpFormatter
        if "formatter_class" not in kwargs:
            kwargs["formatter_class"] = SortedHelpFormatter
1033
1034
        # Pop kwarg "add_json_tip" to control whether to add the JSON tip
        self.add_json_tip = kwargs.pop("add_json_tip", True)
1035
1036
        super().__init__(*args, **kwargs)

1037
    if sys.version_info < (3, 13):
1038
        # Enable the deprecated kwarg for Python 3.12 and below
1039

1040
        def parse_known_args(self, args=None, namespace=None):
1041
1042
1043
1044
            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(
1045
1046
                    "argument '--disable-log-requests' is deprecated and "
                    "replaced with '--enable-log-requests'. This will be "
1047
1048
                    "removed in v0.12.0."
                )
1049
1050
            namespace, args = super().parse_known_args(args, namespace)
            for action in FlexibleArgumentParser._deprecated:
1051
1052
1053
1054
                if (
                    hasattr(namespace, dest := action.dest)
                    and getattr(namespace, dest) != action.default
                ):
1055
                    logger.warning_once("argument '%s' is deprecated", dest)
1056
1057
            return namespace, args

1058
1059
        def add_argument(self, *args, **kwargs):
            deprecated = kwargs.pop("deprecated", False)
1060
            action = super().add_argument(*args, **kwargs)
1061
1062
            if deprecated:
                FlexibleArgumentParser._deprecated.add(action)
1063
1064
            return action

1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
        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
1077

1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
    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'
1090
            if search_keyword == "all":
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
                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
1109
1110
1111
                    if any(
                        search_keyword in opt.lower() for opt in action.option_strings
                    ):
1112
1113
                        matched_actions.append(action)
            if matched_actions:
1114
                formatter.start_section(f"Arguments matching '{search_keyword}'")
1115
1116
1117
1118
1119
1120
1121
1122
1123
                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 "
1124
1125
                "'--help=all' to see all available parameters."
            )
1126
1127
1128
            return formatter.format_help()

        # usage
1129
        formatter.add_usage(self.usage, self._actions, self._mutually_exclusive_groups)
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150

        # 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()
1151

1152
1153
1154
1155
1156
    def parse_args(  # type: ignore[override]
        self,
        args: list[str] | None = None,
        namespace: Namespace | None = None,
    ):
1157
1158
1159
        if args is None:
            args = sys.argv[1:]

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

                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
1196

1197
        if "--config" in args:
1198
            args = self._pull_args_from_config(args)
1199

1200
1201
1202
1203
1204
1205
1206
        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"(?<=--)[^\.]*")

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

1237
        def create_nested_dict(keys: list[str], value: str) -> dict[str, Any]:
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
            """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

1248
1249
1250
        def recursive_dict_update(
            original: dict[str, Any],
            update: dict[str, Any],
1251
1252
1253
1254
1255
        ) -> set[str]:
            """Recursively updates a dictionary with another dictionary.
            Returns a set of duplicate keys that were overwritten.
            """
            duplicates = set[str]()
1256
1257
            for k, v in update.items():
                if isinstance(v, dict) and isinstance(original.get(k), dict):
1258
1259
1260
1261
                    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
1262
                else:
1263
1264
                    if k in original:
                        duplicates.add(k)
1265
                    original[k] = v
1266
            return duplicates
1267

1268
1269
        delete = set[int]()
        dict_args = defaultdict[str, dict[str, Any]](dict)
1270
        duplicates = set[str]()
1271
        for i, processed_arg in enumerate(processed_args):
1272
1273
1274
1275
            if i in delete:  # skip if value from previous arg
                continue

            if processed_arg.startswith("-") and "." in processed_arg:
1276
                if "=" in processed_arg:
1277
                    processed_arg, value_str = processed_arg.split("=", 1)
1278
                    if "." not in processed_arg:
1279
                        # False positive, '.' was only in the value
1280
1281
                        continue
                else:
1282
                    value_str = processed_args[i + 1]
1283
                    delete.add(i + 1)
1284

1285
1286
1287
1288
                if processed_arg.endswith("+"):
                    processed_arg = processed_arg[:-1]
                    value_str = json.dumps(list(value_str.split(",")))

1289
                key, *keys = processed_arg.split(".")
1290
1291
1292
1293
1294
                try:
                    value = json.loads(value_str)
                except json.decoder.JSONDecodeError:
                    value = value_str

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

1305
1306
1307
1308
1309
        # 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))

1310
        return super().parse_args(processed_args, namespace)
1311

1312
1313
1314
1315
    def check_port(self, value):
        try:
            value = int(value)
        except ValueError:
1316
            msg = "Port must be an integer"
1317
            raise ArgumentTypeError(msg) from None
1318
1319

        if not (1024 <= value <= 65535):
1320
            raise ArgumentTypeError("Port must be between 1024 and 65535")
1321
1322
1323

        return value

1324
    def _pull_args_from_config(self, args: list[str]) -> list[str]:
1325
1326
        """Method to pull arguments specified in the config file
        into the command-line args variable.
1327
1328

        The arguments in config file will be inserted between
1329
        the argument list.
1330

1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
        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",
1341
1342
            "facebook/opt-12B",
            '--config', 'config.yaml',
1343
1344
1345
1346
            '-tp', '2'
        ]
        $: args = [
            "serve,chat,complete",
1347
1348
1349
            "facebook/opt-12B",
            '--port', '12323',
            '--tensor-parallel-size', '4',
1350
1351
1352
1353
1354
            '-tp', '2'
            ]
        ```

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

1360
        index = args.index("--config")
1361
        if index == len(args) - 1:
1362
1363
1364
1365
            raise ValueError(
                "No config file specified! \
                             Please check your command-line arguments."
            )
1366
1367
1368

        file_path = args[index + 1]

1369
        config_args = self.load_config_file(file_path)
1370

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

            if not model_in_cli and not model_in_config:
1385
                raise ValueError(
1386
                    "No model specified! Please specify model either "
1387
1388
                    "as a positional argument or in a config file."
                )
1389
1390
1391

            if model_in_cli:
                # Model specified as positional arg, keep CLI version
1392
1393
1394
1395
1396
1397
1398
                args = (
                    [args[0]]
                    + [args[1]]
                    + config_args
                    + args[2:index]
                    + args[index + 2 :]
                )
1399
1400
            else:
                # No model in CLI, use config if available
1401
                args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
1402
        else:
1403
            args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
1404
1405
1406

        return args

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

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

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

1443
        store_boolean_arguments = [
1444
            action.dest for action in self._actions if isinstance(action, StoreBoolean)
1445
1446
        ]

1447
        for key, value in config.items():
1448
1449
            if isinstance(value, bool) and key not in store_boolean_arguments:
                if value:
1450
                    processed_args.append("--" + key)
1451
1452
            elif isinstance(value, list):
                if value:
1453
                    processed_args.append("--" + key)
1454
1455
                    for item in value:
                        processed_args.append(str(item))
1456
            else:
1457
                processed_args.append("--" + key)
1458
                processed_args.append(str(value))
1459
1460
1461

        return processed_args

1462

1463
1464
1465
1466
1467
1468
# 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")
1469
1470


1471
# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
1472
def supports_xccl() -> bool:
1473
1474
1475
    return (
        is_torch_equal_or_newer("2.8.0.dev") and torch.distributed.is_xccl_available()
    )
1476
1477


1478
1479
1480
1481
1482
1483
# 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")


1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
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
1507
1508


1509
def weak_ref_tensor(tensor: Any) -> Any:
1510
1511
1512
1513
1514
    """
    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.
    """
1515
1516
1517
1518
    if isinstance(tensor, torch.Tensor):
        return torch.ops._C.weak_ref_tensor(tensor)
    else:
        return tensor
1519
1520
1521


def weak_ref_tensors(
1522
1523
1524
1525
1526
    tensors: torch.Tensor
    | list[torch.Tensor]
    | tuple[torch.Tensor]
    | IntermediateTensors,
) -> torch.Tensor | list[Any] | tuple[Any] | Any:
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
    """
    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)
1537
1538
1539

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

1541
    if isinstance(tensors, IntermediateTensors):
1542
1543
1544
        ret = IntermediateTensors(
            {key: weak_ref_tensor(val) for key, val in tensors.tensors.items()}
        )
1545
        return ret
1546
    raise ValueError("Invalid type for tensors")
1547
1548


1549
1550
1551
1552
1553
1554
1555
1556
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)


1557
1558
1559
1560
1561
# create a library to hold the custom op
vllm_lib = Library("vllm", "FRAGMENT")  # noqa


def direct_register_custom_op(
1562
1563
    op_name: str,
    op_func: Callable,
1564
1565
1566
1567
    mutates_args: list[str] | None = None,
    fake_impl: Callable | None = None,
    target_lib: Library | None = None,
    dispatch_key: str | None = None,
1568
    tags: tuple[torch.Tag, ...] = (),
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
):
    """
    `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.
    """
1585
    if not supports_custom_op():
1586
        from vllm.platforms import current_platform
1587

1588
1589
1590
1591
1592
        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 "
1593
1594
            "the required dependencies."
        )
1595
1596
        return

1597
1598
1599
1600
1601
    if mutates_args is None:
        mutates_args = []

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

1603
1604
        dispatch_key = current_platform.dispatch_key

1605
    import torch.library
1606

1607
    if hasattr(torch.library, "infer_schema"):
1608
        schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args)
1609
1610
1611
    else:
        # for pytorch 2.4
        import torch._custom_op.impl
1612

1613
        schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
1614
    my_lib = target_lib or vllm_lib
1615
    my_lib.define(op_name + schema_str, tags=tags)
1616
    my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
1617
1618
    if fake_impl is not None:
        my_lib._register_fake(op_name, fake_impl)
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
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)
1644
1645
1646
1647
1648


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

1650
    torch_peak: int = 0
1651
1652
    free_memory: int = 0
    total_memory: int = 0
1653
1654
1655
    cuda_memory: int = 0
    torch_memory: int = 0
    non_torch_memory: int = 0
1656
    timestamp: float = 0.0
1657
1658
1659
1660
1661
    auto_measure: bool = True

    def __post_init__(self):
        if self.auto_measure:
            self.measure()
1662
1663

    def measure(self):
1664
1665
        from vllm.platforms import current_platform

1666
1667
1668
1669
1670
        # 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.
1671
        self.torch_peak = torch.cuda.memory_stats().get("allocated_bytes.all.peak", 0)
1672

1673
        self.free_memory, self.total_memory = torch.cuda.mem_get_info()
1674
1675
1676
1677
1678
        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
        ):
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
            # 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

1692
        self.cuda_memory = self.total_memory - self.free_memory
1693

1694
1695
        # torch.cuda.memory_reserved() is how many bytes
        # PyTorch gets from cuda (by calling cudaMalloc, etc.)
1696
1697
1698
1699
        # 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
1700
1701
        self.timestamp = time.time()

1702
    def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot":
1703
        return MemorySnapshot(
1704
            torch_peak=self.torch_peak - other.torch_peak,
1705
1706
            free_memory=self.free_memory - other.free_memory,
            total_memory=self.total_memory - other.total_memory,
1707
1708
1709
1710
1711
1712
            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,
        )
1713
1714
1715
1716


@dataclass
class MemoryProfilingResult:
1717
1718
    """Memory profiling result. All numbers are in bytes."""

1719
1720
1721
1722
1723
    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)
1724
1725
1726
1727
    before_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
    after_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
    profile_time: float = 0.0

1728
    def __repr__(self) -> str:
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
        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."
        )
1739

1740
1741
1742

@contextlib.contextmanager
def memory_profiling(
1743
1744
    baseline_snapshot: MemorySnapshot, weights_memory: int
) -> Generator[MemoryProfilingResult, None, None]:
1745
    """Memory profiling context manager.
1746
1747
    baseline_snapshot: the memory snapshot before the current vLLM instance.
    weights_memory: memory used by PyTorch when loading the model weights.
1748
1749
        Note that, before loading the model weights, we also initialize the device
        and distributed environment, which may consume some memory. This part is not
1750
        included in the weights_memory because PyTorch does not control it.
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784

    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)

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

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

1789
    The increase of `non_torch_memory` from creating the current vLLM instance until after profiling to get (c.).
1790
    """  # noqa
1791
1792
    gc.collect()
    torch.cuda.empty_cache()
1793
1794
1795
1796
    torch.cuda.reset_peak_memory_stats()

    result = MemoryProfilingResult()

1797
    result.before_create = baseline_snapshot
1798
    # the part of memory used for holding the model weights
1799
    result.weights_memory = weights_memory
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809

    result.before_profile.measure()

    yield result

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

    result.after_profile.measure()

1810
1811
1812
1813
1814
    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
1815
1816
1817

    non_torch_memory = result.non_torch_increase
    peak_activation_memory = result.torch_peak_increase
1818
1819
1820
    result.non_kv_cache_memory = (
        non_torch_memory + peak_activation_memory + result.weights_memory
    )  # noqa
1821
1822


1823
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 # noqa: E501
1824
def set_ulimit(target_soft_limit=65535):
1825
    if sys.platform.startswith("win"):
1826
1827
1828
1829
        logger.info("Windows detected, skipping ulimit adjustment.")
        return

    import resource
1830

1831
1832
1833
1834
1835
    resource_type = resource.RLIMIT_NOFILE
    current_soft, current_hard = resource.getrlimit(resource_type)

    if current_soft < target_soft_limit:
        try:
1836
            resource.setrlimit(resource_type, (target_soft_limit, current_hard))
1837
1838
        except ValueError as e:
            logger.warning(
1839
1840
                "Found ulimit of %s and failed to automatically increase "
                "with error %s. This can cause fd limit errors like "
1841
                "`OSError: [Errno 24] Too many open files`. Consider "
1842
1843
1844
1845
                "increasing with ulimit -n",
                current_soft,
                e,
            )
1846
1847
1848
1849
1850
1851
1852
1853
1854


# 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


1855
def split_zmq_path(path: str) -> tuple[str, str, str]:
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
    """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


1876
def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str:
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
    """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.
    """
1887
    if port is None:
1888
1889
1890
1891
1892
1893
        return f"{scheme}://{host}"
    if is_valid_ipv6_address(host):
        return f"{scheme}://[{host}]:{port}"
    return f"{scheme}://{host}:{port}"


1894
1895
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501
def make_zmq_socket(
1896
    ctx: zmq.asyncio.Context | zmq.Context,  # type: ignore[name-defined]
1897
    path: str,
1898
    socket_type: Any,
1899
1900
1901
    bind: bool | None = None,
    identity: bytes | None = None,
    linger: int | None = None,
1902
) -> zmq.Socket | zmq.asyncio.Socket:  # type: ignore[name-defined]
1903
1904
1905
    """Make a ZMQ socket with the proper bind/connect semantics."""

    mem = psutil.virtual_memory()
1906
    socket = ctx.socket(socket_type)
1907
1908
1909
1910
1911
1912
1913
1914

    # 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
1915
    buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1
1916

1917
    if bind is None:
1918
        bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB)
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930

    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)

1931
1932
1933
    if linger is not None:
        socket.setsockopt(zmq.LINGER, linger)

1934
1935
1936
    if socket_type == zmq.XPUB:
        socket.setsockopt(zmq.XPUB_VERBOSE, True)

1937
1938
1939
1940
1941
1942
    # 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)

1943
    if bind:
1944
        socket.bind(path)
1945
    else:
1946
        socket.connect(path)
1947
1948
1949
1950
1951

    return socket


@contextlib.contextmanager
1952
1953
1954
def zmq_socket_ctx(
    path: str,
    socket_type: Any,
1955
    bind: bool | None = None,
1956
    linger: int = 0,
1957
    identity: bytes | None = None,
1958
) -> Iterator[zmq.Socket]:
1959
1960
    """Context manager for a ZMQ socket"""

1961
    ctx = zmq.Context()  # type: ignore[attr-defined]
1962
    try:
1963
        yield make_zmq_socket(ctx, path, socket_type, bind=bind, identity=identity)
1964
1965
1966
1967
    except KeyboardInterrupt:
        logger.debug("Got Keyboard Interrupt.")

    finally:
1968
        ctx.destroy(linger=linger)
1969
1970


1971
1972
1973
1974
1975
1976
1977
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

1978
1979
    reasons = []
    if is_in_ray_actor():
1980
1981
1982
1983
        # 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
1984

1985
        os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
1986
1987
1988
1989
1990
1991
        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")
1992

1993
    if reasons:
1994
1995
1996
        logger.warning(
            "We must use the `spawn` multiprocessing start method. "
            "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "
1997
            "See https://docs.vllm.ai/en/latest/usage/"
1998
            "troubleshooting.html#python-multiprocessing "
1999
2000
2001
            "for more information. Reasons: %s",
            "; ".join(reasons),
        )
2002
2003
2004
2005
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"


def get_mp_context():
2006
2007
2008
2009
2010
2011
2012
    """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()
2013
2014
    mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
    return multiprocessing.get_context(mp_method)
2015
2016
2017


def bind_kv_cache(
2018
2019
    ctx: dict[str, Any],
    kv_cache: list[list[torch.Tensor]],  # [virtual_engine][layer_index]
2020
    shared_kv_cache_layers: dict[str, str] | None = None,
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
) -> 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
2032
2033
2034
2035
    # 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 = {}
2036
2037
    from vllm.attention import AttentionType
    from vllm.model_executor.models.utils import extract_layer_index
2038

2039
    layer_need_kv_cache = [
2040
2041
2042
2043
2044
2045
2046
2047
        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
2048
2049
    ]
    layer_index_sorted = sorted(
2050
2051
        set(extract_layer_index(layer_name) for layer_name in layer_need_kv_cache)
    )
2052
    for layer_name in layer_need_kv_cache:
2053
        kv_cache_idx = layer_index_sorted.index(extract_layer_index(layer_name))
2054
2055
2056
2057
        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]
2058
2059
    if shared_kv_cache_layers is not None:
        for layer_name, target_layer_name in shared_kv_cache_layers.items():
2060
2061
2062
            assert extract_layer_index(target_layer_name) < extract_layer_index(
                layer_name
            ), "v0 doesn't support interleaving kv sharing"
2063
            ctx[layer_name].kv_cache = ctx[target_layer_name].kv_cache
2064
2065


2066
2067
def run_method(
    obj: Any,
2068
    method: str | bytes | Callable,
2069
2070
2071
    args: tuple[Any],
    kwargs: dict[str, Any],
) -> Any:
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
    """
    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:
2085
2086
2087
            raise NotImplementedError(
                f"Method {method!r} is not implemented."
            ) from None
2088
2089
2090
    else:
        func = partial(method, obj)  # type: ignore
    return func(*args, **kwargs)
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111


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.
2112
2113
2114
2115
2116
2117
2118
    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.
2119
    """
2120
    import vllm.third_party.pynvml as pynvml
2121

2122
    return pynvml
2123
2124


2125
def warn_for_unimplemented_methods(cls: type[T]) -> type[T]:
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
    """
    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
2141
            if attr_name.startswith("_"):
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
                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:
2155
2156
            method_names = ",".join(unimplemented_methods)
            msg = f"Methods {method_names} not implemented in {self}"
2157
            logger.debug(msg)
2158
2159
2160
2161
2162
2163

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

2164
    type.__setattr__(cls, "__init__", wrapped_init)
2165
    return cls
2166
2167


2168
@contextlib.contextmanager
2169
def cprofile_context(save_file: str | None = None):
2170
2171
2172
2173
    """Run a cprofile

    Args:
        save_file: path to save the profile result. "1" or
2174
            None will result in printing to stdout.
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
    """
    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")


2191
def cprofile(save_file: str | None = None, enabled: bool = True):
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
    """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
2213
2214


2215
2216
# Only relevant for models using ALiBi (e.g, MPT)
def check_use_alibi(model_config: ModelConfig) -> bool:
2217
    cfg = model_config.hf_text_config
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
    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)
                )
            )
        )
    )
2238
2239


2240
def sha256(input: Any) -> bytes:
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
    """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:
2251
        Bytes representing the SHA-256 hash of the serialized input.
2252
2253
    """
    input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
2254
    return hashlib.sha256(input_bytes).digest()
2255
2256


2257
def sha256_cbor(input: Any) -> bytes:
2258
    """
2259
    Hash objects using CBOR serialization and SHA-256.
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269

    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:
2270
        Bytes representing the SHA-256 hash of the CBOR serialized input.
2271
2272
    """
    input_bytes = cbor2.dumps(input, canonical=True)
2273
    return hashlib.sha256(input_bytes).digest()
2274
2275


2276
def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], bytes]:
2277
2278
2279
2280
2281
2282
2283
2284
2285
    """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
2286
2287
    if hash_fn_name == "sha256_cbor":
        return sha256_cbor
2288
2289
2290
2291

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


2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
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:
2302
        return _is_torch_equal_or_newer(str(torch.__version__), target)
2303
2304
    except Exception:
        # Fallback to PKG-INFO to load the package info, needed by the doc gen.
2305
        return Version(importlib.metadata.version("torch")) >= Version(target)
2306
2307
2308
2309
2310
2311


# 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)
2312
2313


2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
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)


2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
@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."""

2366
    return _has_module("deep_gemm")
2367
2368


2369
2370
2371
2372
2373
2374
def has_triton_kernels() -> bool:
    """Whether the optional `triton_kernels` package is available."""

    return _has_module("triton_kernels")


2375
2376
2377
2378
2379
2380
def has_tilelang() -> bool:
    """Whether the optional `tilelang` package is available."""

    return _has_module("tilelang")


2381
2382
2383
def set_process_title(
    name: str, suffix: str = "", prefix: str = envs.VLLM_PROCESS_NAME_PREFIX
) -> None:
2384
2385
2386
    """
    Set the current process title to a specific name with an
    optional suffix.
2387
2388

    Args:
2389
        name: The title to assign to the current process.
2390
        suffix: An optional suffix to append to the base name.
2391
        prefix: A prefix to prepend to the front separated by `::`.
2392
2393
2394
    """
    if suffix:
        name = f"{name}_{suffix}"
2395
    setproctitle.setproctitle(f"{prefix}::{name}")
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409


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
2410
        while (next_idx := s.find("\n", idx)) != -1:
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
            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]


2425
def decorate_logs(process_name: str | None = None) -> None:
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
    """
    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)
2445
2446
2447


def length_from_prompt_token_ids_or_embeds(
2448
2449
    prompt_token_ids: list[int] | None,
    prompt_embeds: torch.Tensor | None,
2450
) -> int:
2451
    """Calculate the request length (in number of tokens) give either
2452
2453
    prompt_token_ids or prompt_embeds.
    """
2454
2455
    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)
2456
2457
2458

    if prompt_token_len is None:
        if prompt_embeds_len is None:
2459
            raise ValueError("Neither prompt_token_ids nor prompt_embeds were defined.")
2460
2461
        return prompt_embeds_len
    else:
2462
        if prompt_embeds_len is not None and prompt_embeds_len != prompt_token_len:
2463
2464
2465
            raise ValueError(
                "Prompt token ids and prompt embeds had different lengths"
                f" prompt_token_ids={prompt_token_len}"
2466
2467
                f" prompt_embeds={prompt_embeds_len}"
            )
2468
        return prompt_token_len
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481


@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
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501


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