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

4
5
from __future__ import annotations

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

54
import cachetools
55
import cloudpickle
56
import numpy as np
57
import numpy.typing as npt
58
import psutil
59
import regex as re
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
from transformers.tokenization_utils_base import BatchEncoding
69
from typing_extensions import Never, ParamSpec, TypeIs, assert_never
70

71
import vllm.envs as envs
72
from vllm.logger import enable_trace_function_call, init_logger
73

74
if TYPE_CHECKING:
75
76
    from argparse import Namespace

77
    from vllm.config import ModelConfig, VllmConfig
78

79
80
logger = init_logger(__name__)

81
82
83
84
85
86
# 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

87
88
# Exception strings for non-implemented encoder/decoder scenarios

89
# Reminder: Please update docs/features/compatibility_matrix.md
90
91
# If the feature combo become valid

92
93
STR_NOT_IMPL_ENC_DEC_SWA = \
    "Sliding window attention for encoder/decoder models " + \
94
    "is not currently supported."
95
96
97

STR_NOT_IMPL_ENC_DEC_PREFIX_CACHE = \
    "Prefix caching for encoder/decoder models " + \
98
    "is not currently supported."
99
100
101

STR_NOT_IMPL_ENC_DEC_CHUNKED_PREFILL = \
    "Chunked prefill for encoder/decoder models " + \
102
    "is not currently supported."
103
104
105
106
107
108
109

STR_NOT_IMPL_ENC_DEC_LOGIT_SOFTCAP = (
    "Models with logits_soft_cap "
    "require FlashInfer backend, which is "
    "currently not supported for encoder/decoder "
    "models.")

Always-Naive's avatar
Always-Naive committed
110
STR_NOT_IMPL_ENC_DEC_LORA = ("LoRA is not currently "
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
                             "supported with encoder/decoder "
                             "models.")

STR_NOT_IMPL_ENC_DEC_PP = ("Pipeline parallelism is not "
                           "currently supported with "
                           "encoder/decoder models.")

STR_NOT_IMPL_ENC_DEC_MM = ("Multimodal is not currently "
                           "supported with encoder/decoder "
                           "models.")

STR_NOT_IMPL_ENC_DEC_SPEC_DEC = ("Speculative decoding is not "
                                 "currently supported with encoder/"
                                 "decoder models.")

126
127
STR_NOT_IMPL_ENC_DEC_BACKEND = ("XFormers and Flash-Attention are the only "
                                "backends currently supported with encoder/"
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
                                "decoder models.")

STR_NOT_IMPL_ENC_DEC_PROMPT_ADAPTER = ("Prompt adapters are not "
                                       "currently supported with encoder/"
                                       "decoder models.")

# Efficiently import all enc/dec error strings
# rather than having to import all of the above
STR_NOT_IMPL_ENC_DEC_ERR_STRS = {
    "STR_NOT_IMPL_ENC_DEC_SWA": STR_NOT_IMPL_ENC_DEC_SWA,
    "STR_NOT_IMPL_ENC_DEC_PREFIX_CACHE": STR_NOT_IMPL_ENC_DEC_PREFIX_CACHE,
    "STR_NOT_IMPL_ENC_DEC_CHUNKED_PREFILL":
    STR_NOT_IMPL_ENC_DEC_CHUNKED_PREFILL,
    "STR_NOT_IMPL_ENC_DEC_LOGIT_SOFTCAP": STR_NOT_IMPL_ENC_DEC_LOGIT_SOFTCAP,
    "STR_NOT_IMPL_ENC_DEC_LORA": STR_NOT_IMPL_ENC_DEC_LORA,
    "STR_NOT_IMPL_ENC_DEC_PP": STR_NOT_IMPL_ENC_DEC_PP,
    "STR_NOT_IMPL_ENC_DEC_MM": STR_NOT_IMPL_ENC_DEC_MM,
    "STR_NOT_IMPL_ENC_DEC_SPEC_DEC": STR_NOT_IMPL_ENC_DEC_SPEC_DEC,
    "STR_NOT_IMPL_ENC_DEC_BACKEND": STR_NOT_IMPL_ENC_DEC_BACKEND,
    "STR_NOT_IMPL_ENC_DEC_PROMPT_ADAPTER": STR_NOT_IMPL_ENC_DEC_PROMPT_ADAPTER,
}

# 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_ROCM_FLASH_ATTN_VAL: str = "ROCM_FLASH"
STR_XFORMERS_ATTN_VAL: str = "XFORMERS"
STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
164
STR_DUAL_CHUNK_FLASH_ATTN_VAL: str = "DUAL_CHUNK_FLASH_ATTN"
165
166
STR_INVALID_VAL: str = "INVALID"

167
168
169
GB_bytes = 1_000_000_000
"""The number of bytes in one gigabyte (GB)."""

170
171
172
GiB_bytes = 1 << 30
"""The number of bytes in one gibibyte (GiB)."""

173
174
175
176
STR_DTYPE_TO_TORCH_DTYPE = {
    "half": torch.half,
    "bfloat16": torch.bfloat16,
    "float": torch.float,
177
    "fp8": torch.uint8,
178
179
    "fp8_e4m3": torch.uint8,
    "fp8_e5m2": torch.uint8,
180
    "int8": torch.int8,
181
}
Zhuohan Li's avatar
Zhuohan Li committed
182

183
184
185
186
187
188
189
190
191
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,
}

192
193
194
195
196
197
198
199
200
201

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


202
203
P = ParamSpec('P')
T = TypeVar("T")
204
U = TypeVar("U")
205

206
207
_K = TypeVar("_K", bound=Hashable)
_V = TypeVar("_V")
208
_T = TypeVar("_T")
209

Woosuk Kwon's avatar
Woosuk Kwon committed
210

211
212
213
214
215
216
217
class _Sentinel:
    ...


ALL_PINNED_SENTINEL = _Sentinel()


Woosuk Kwon's avatar
Woosuk Kwon committed
218
219
220
221
222
class Device(enum.Enum):
    GPU = enum.auto()
    CPU = enum.auto()


223
224
225
226
227
class LayerBlockType(enum.Enum):
    attention = "attention"
    mamba = "mamba"


Woosuk Kwon's avatar
Woosuk Kwon committed
228
229
230
231
232
class Counter:

    def __init__(self, start: int = 0) -> None:
        self.counter = start

Woosuk Kwon's avatar
Woosuk Kwon committed
233
    def __next__(self) -> int:
234
        i = self.counter
Woosuk Kwon's avatar
Woosuk Kwon committed
235
        self.counter += 1
236
        return i
Woosuk Kwon's avatar
Woosuk Kwon committed
237
238
239

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

241

242
243
244
245
246
247
248
249
250
251
252
253
254
class _MappingOrderCacheView(UserDict[_K, _V]):

    def __init__(self, data: Mapping[_K, _V], ordered_keys: Mapping[_K, None]):
        super().__init__(data)
        self.ordered_keys = ordered_keys

    def __iter__(self) -> Iterator[_K]:
        return iter(self.ordered_keys)

    def keys(self) -> KeysView[_K]:
        return KeysView(self.ordered_keys)


255
256
257
258
259
260
261
262
263
264
265
class CacheInfo(NamedTuple):
    hits: int
    total: int

    @property
    def hit_ratio(self) -> float:
        if self.total == 0:
            return 0

        return self.hits / self.total

266
267
268
269
270
271
    def __sub__(self, other: CacheInfo):
        return CacheInfo(
            hits=self.hits - other.hits,
            total=self.total - other.total,
        )

272

273
class LRUCache(cachetools.LRUCache[_K, _V], Generic[_K, _V]):
274

275
276
277
278
    def __init__(self,
                 capacity: float,
                 getsizeof: Optional[Callable[[_V], float]] = None):
        super().__init__(capacity, getsizeof)
279

280
        self.pinned_items = set[_K]()
281

282
283
        self._hits = 0
        self._total = 0
284
285
286
287
288
289
290
291
292
293
        self._last_info = CacheInfo(hits=0, total=0)

    def __getitem__(self, key: _K, *, update_info: bool = True) -> _V:
        value = super().__getitem__(key)

        if update_info:
            self._hits += 1
            self._total += 1

        return value
294

295
296
    def __delitem__(self, key: _K) -> None:
        run_on_remove = key in self
297
298
        value = self.__getitem__(key,
                                 update_info=False)  # type: ignore[call-arg]
299
300
301
302
303
304
        super().__delitem__(key)
        if key in self.pinned_items:
            # Todo: add warning to inform that del pinned item
            self._unpin(key)
        if run_on_remove:
            self._on_remove(key, value)
305

306
307
308
309
310
311
    @property
    def cache(self) -> Mapping[_K, _V]:
        """Return the internal cache dictionary in order (read-only)."""
        return _MappingOrderCacheView(
            self._Cache__data,  # type: ignore
            self.order)
312

313
314
315
316
    @property
    def order(self) -> Mapping[_K, None]:
        """Return the internal order dictionary (read-only)."""
        return MappingProxyType(self._LRUCache__order)  # type: ignore
317

318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
    @property
    def capacity(self) -> float:
        return self.maxsize

    @property
    def usage(self) -> float:
        if self.maxsize == 0:
            return 0

        return self.currsize / self.maxsize

    def stat(self, *, delta: bool = False) -> CacheInfo:
        """
        Gets the cumulative number of hits and queries against this cache.

333
334
        If `delta=True`, instead gets these statistics
        since the last call that also passed `delta=True`.
335
336
337
338
339
340
341
342
343
        """
        info = CacheInfo(hits=self._hits, total=self._total)

        if delta:
            info_delta = info - self._last_info
            self._last_info = info
            info = info_delta

        return info
344

345
    def touch(self, key: _K) -> None:
346
347
348
349
        try:
            self._LRUCache__order.move_to_end(key)  # type: ignore
        except KeyError:
            self._LRUCache__order[key] = None  # type: ignore
350
351
352
353
354
355
356
357

    @overload
    def get(self, key: _K, /) -> Optional[_V]:
        ...

    @overload
    def get(self, key: _K, /, default: Union[_V, _T]) -> Union[_V, _T]:
        ...
358

359
360
361
362
363
364
365
    def get(self,
            key: _K,
            /,
            default: Optional[Union[_V,
                                    _T]] = None) -> Optional[Union[_V, _T]]:
        value: Optional[Union[_V, _T]]
        if key in self:
366
367
            value = self.__getitem__(
                key, update_info=False)  # type: ignore[call-arg]
368
369

            self._hits += 1
370
        else:
371
            value = default
372
373

        self._total += 1
374
375
        return value

376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
    @overload
    def pop(self, key: _K) -> _V:
        ...

    @overload
    def pop(self, key: _K, default: Union[_V, _T]) -> Union[_V, _T]:
        ...

    def pop(self,
            key: _K,
            default: Optional[Union[_V,
                                    _T]] = None) -> Optional[Union[_V, _T]]:
        value: Optional[Union[_V, _T]]
        if key not in self:
            return default

392
393
394
        value = self.__getitem__(key,
                                 update_info=False)  # type: ignore[call-arg]
        self.__delitem__(key)
395
396
        return value

397
    def put(self, key: _K, value: _V) -> None:
398
        self.__setitem__(key, value)
399

400
    def pin(self, key: _K) -> None:
401
402
403
404
        """
        Pins a key in the cache preventing it from being
        evicted in the LRU order.
        """
405
        if key not in self:
406
407
408
            raise ValueError(f"Cannot pin key: {key} not in cache.")
        self.pinned_items.add(key)

409
    def _unpin(self, key: _K) -> None:
410
411
412
413
        """
        Unpins a key in the cache allowing it to be
        evicted in the LRU order.
        """
414
415
        self.pinned_items.remove(key)

416
    def _on_remove(self, key: _K, value: Optional[_V]) -> None:
417
418
        pass

419
    def remove_oldest(self, *, remove_pinned: bool = False) -> None:
420
        if len(self) == 0:
421
            return
422

423
424
425
426
427
428
429
430
        self.popitem(remove_pinned=remove_pinned)

    def _remove_old_if_needed(self) -> None:
        while self.currsize > self.capacity:
            self.remove_oldest()

    def popitem(self, remove_pinned: bool = False):
        """Remove and return the `(key, value)` pair least recently used."""
431
432
433
        if not remove_pinned:
            # pop the oldest item in the cache that is not pinned
            lru_key = next(
434
                (key for key in self.order if key not in self.pinned_items),
435
436
437
438
439
                ALL_PINNED_SENTINEL)
            if lru_key is ALL_PINNED_SENTINEL:
                raise RuntimeError("All items are pinned, "
                                   "cannot remove oldest from the cache.")
        else:
440
441
442
            lru_key = next(iter(self.order))
        value = self.pop(cast(_K, lru_key))
        return (lru_key, value)
443

444
445
446
447
448
449
450
451
    def clear(self) -> None:
        while len(self) > 0:
            self.remove_oldest(remove_pinned=True)

        self._hits = 0
        self._total = 0
        self._last_info = CacheInfo(hits=0, total=0)

452

453
class PyObjectCache:
454
    """Used to cache python objects to avoid object allocations
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
    across scheduler iterations.
    """

    def __init__(self, obj_builder):
        self._obj_builder = obj_builder
        self._index = 0

        self._obj_cache = []
        for _ in range(128):
            self._obj_cache.append(self._obj_builder())

    def _grow_cache(self):
        # Double the size of the cache
        num_objs = len(self._obj_cache)
        for _ in range(num_objs):
            self._obj_cache.append(self._obj_builder())

    def get_object(self):
473
        """Returns a pre-allocated cached object. If there is not enough
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
        objects, then the cache size will double.
        """
        if self._index >= len(self._obj_cache):
            self._grow_cache()
            assert self._index < len(self._obj_cache)

        obj = self._obj_cache[self._index]
        self._index += 1

        return obj

    def reset(self):
        """Makes all cached-objects available for the next scheduler iteration.
        """
        self._index = 0


491
@cache
492
493
def get_max_shared_memory_bytes(gpu: int = 0) -> int:
    """Returns the maximum shared memory per thread block in bytes."""
494
    from vllm import _custom_ops as ops
495
    max_shared_mem = (
496
        ops.get_max_shared_memory_per_block_device_attribute(gpu))
497
498
    # value 0 will cause MAX_SEQ_LEN become negative and test_attention.py
    # will fail
499
    assert max_shared_mem > 0, "max_shared_mem can not be zero"
500
501
502
    return int(max_shared_mem)


503
def get_cpu_memory() -> int:
504
    """Returns the total CPU memory of the node in bytes."""
505
    return psutil.virtual_memory().total
Zhuohan Li's avatar
Zhuohan Li committed
506
507
508
509


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

511

512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
class AsyncMicrobatchTokenizer:
    """Asynchronous tokenizer with micro-batching.

    Pulls pending encode/decode requests from a queue and batches them 
    up to reduce overhead. A single-thread ThreadPoolExecutor is used 
    so the event loop stays responsive.
    """

    def __init__(
        self,
        tokenizer,
        max_batch_size: int = 32,
        batch_wait_timeout_s: float = 0.002,
    ) -> None:
        self.tokenizer = tokenizer
        self.max_batch_size = max_batch_size
        self.batch_wait_timeout_s = batch_wait_timeout_s

        self._loop = asyncio.get_running_loop()
        self._queues: dict[tuple,
                           asyncio.Queue[Union[tuple[str, dict,
                                                     asyncio.Future],
                                               tuple[list[int],
                                                     asyncio.Future]]]] = {}
        self._batcher_tasks: list[asyncio.Task] = []

        # Single-thread executor for blocking tokenizer calls.
        self._executor = ThreadPoolExecutor(max_workers=1)

    # === Public async API ===
    async def __call__(self, prompt, **kwargs):
        result_future: asyncio.Future = self._loop.create_future()
        key = self._queue_key("encode", kwargs)
        queue = self._get_queue(self._loop, key)
        await queue.put((prompt, kwargs, result_future))
        return await result_future

    async def decode(self, token_ids, **kwargs):
        result_future: asyncio.Future = self._loop.create_future()
        key = self._queue_key("decode", kwargs)
        queue = self._get_queue(self._loop, key)
        await queue.put((token_ids, result_future))
        return await result_future

    # === Internal helpers ===
    def _get_queue(
        self, loop: asyncio.AbstractEventLoop, key: tuple
    ) -> asyncio.Queue[Union[tuple[str, dict, asyncio.Future], tuple[
            list[int], asyncio.Future]]]:
        """Get the request queue for the given operation key, creating a new
        queue and batcher task if needed."""
        queue = self._queues.get(key)
        if queue is None:
            self._queues[key] = queue = asyncio.Queue()
            if key[0] == "encode":
                can_batch = key[1] != "other"
                coro = self._batch_encode_loop(queue, can_batch)
            else:
                assert key[0] == "decode", \
                    f"Unknown operation type: {key[0]}."
                coro = self._batch_decode_loop(queue)
            self._batcher_tasks.append(loop.create_task(coro))
        return queue

    async def _batch_encode_loop(self, queue: asyncio.Queue, can_batch: bool):
        """Batch incoming encode requests for efficiency."""
        while True:
            prompt, kwargs, result_future = await queue.get()
            prompts = [prompt]
            kwargs_list = [kwargs]
            result_futures = [result_future]
            deadline = self._loop.time() + self.batch_wait_timeout_s

            while len(prompts) < self.max_batch_size:
                timeout = deadline - self._loop.time()
                if timeout <= 0:
                    break
                try:
                    prompt, kwargs, result_future = await asyncio.wait_for(
                        queue.get(), timeout)
                    prompts.append(prompt)
                    result_futures.append(result_future)
                    if not can_batch:
                        kwargs_list.append(kwargs)
                except asyncio.TimeoutError:
                    break

            try:
                # If every request uses identical kwargs we can run a single
                # batched tokenizer call for a big speed-up.
                if can_batch and len(prompts) > 1:
                    encode_fn = partial(self.tokenizer, prompts, **kwargs)
                    results = await self._loop.run_in_executor(
                        self._executor, encode_fn)

                    for i, fut in enumerate(result_futures):
                        if not fut.done():
                            data = {k: v[i] for k, v in results.items()}
                            fut.set_result(BatchEncoding(data))
                else:
                    encode_fn = lambda prompts=prompts, kwargs=kwargs_list: [
                        self.tokenizer(p, **kw)
                        for p, kw in zip(prompts, kwargs)
                    ]
                    results = await self._loop.run_in_executor(
                        self._executor, encode_fn)

                    for fut, res in zip(result_futures, results):
                        if not fut.done():
                            fut.set_result(res)
            except Exception as e:
                for fut in result_futures:
                    if not fut.done():
                        fut.set_exception(e)

    async def _batch_decode_loop(self, queue: asyncio.Queue):
        """Batch incoming decode requests for efficiency."""
        while True:
            token_ids, result_future = await queue.get()
            token_ids_list = [token_ids]
            result_futures = [result_future]
            deadline = self._loop.time() + self.batch_wait_timeout_s

            while len(token_ids_list) < self.max_batch_size:
                timeout = deadline - self._loop.time()
                if timeout <= 0:
                    break
                try:
                    token_ids, result_future = await asyncio.wait_for(
                        queue.get(), timeout)
                    token_ids_list.append(token_ids)
                    result_futures.append(result_future)
                except asyncio.TimeoutError:
                    break

            try:
                # Perform a single batched decode call for all requests
                results = await self._loop.run_in_executor(
                    self._executor, self.tokenizer.batch_decode,
                    token_ids_list)
                for fut, res in zip(result_futures, results):
                    if not fut.done():
                        fut.set_result(res)
            except Exception as e:
                for fut in result_futures:
                    if not fut.done():
                        fut.set_exception(e)

    def _queue_key(self, op: str, kwargs: dict) -> tuple:
        """
        Return a normalized key describing operation + kwargs.
        
        - `add_special_tokens`: {True/False}
        - `truncation`: {True/False}
          - If `truncation` is False (`max_length` is None), 
            returns a key for a can_batch queue.
          - If `truncation` is True and `max_length` is None or equals
            `tokenizer.model_max_length`, returns a key for a can_batch queue.
          - Otherwise, returns a key for a cannot_batch queue.
        
        Examples:
          - Decode: ("decode",)
          - Encode typical: 
            ("encode", add_special_tokens, bool_truncation, max_length_label)
          - Fallback: ("encode", "other")
        """

        if op == "decode":
            return ("decode", )

        add_special_tokens = kwargs.get("add_special_tokens", True)
        truncation = kwargs.get("truncation", False)
        max_length = kwargs.get("max_length")

        if not truncation:
            return ("encode", add_special_tokens, False, None)

        model_max = getattr(self.tokenizer, "model_max_length", None)
        if max_length is None or (model_max is not None
                                  and max_length == model_max):
            return ("encode", add_special_tokens, True, "model_max")

        return ("encode", "other")

    def __del__(self):
        for task in self._batcher_tasks:
            if not task.done():
                task.cancel()


702
703
704
705
def make_async(
    func: Callable[P, T],
    executor: Optional[concurrent.futures.Executor] = None
) -> Callable[P, Awaitable[T]]:
706
707
708
709
710
711
712
    """Take a blocking function, and run it on in an executor thread.

    This function prevents the blocking function from blocking the
    asyncio event loop.
    The code in this function needs to be thread safe.
    """

713
    def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> asyncio.Future:
714
715
        loop = asyncio.get_event_loop()
        p_func = partial(func, *args, **kwargs)
716
        return loop.run_in_executor(executor=executor, func=p_func)
717
718
719
720

    return _async_wrapper


721
722
723
724
725
726
def _next_task(iterator: AsyncGenerator[T, None],
               loop: AbstractEventLoop) -> Task:
    # Can use anext() in python >= 3.10
    return loop.create_task(iterator.__anext__())  # type: ignore[arg-type]


727
async def merge_async_iterators(
728
    *iterators: AsyncGenerator[T,
729
                               None], ) -> AsyncGenerator[tuple[int, T], None]:
730
731
732
733
734
735
    """Merge multiple asynchronous iterators into a single iterator.

    This method handle the case where some iterators finish before others.
    When it yields, it yields a tuple (i, item) where i is the index of the
    iterator that yields the item.
    """
736
737
738
739
740
    if len(iterators) == 1:
        # Fast-path single iterator case.
        async for item in iterators[0]:
            yield 0, item
        return
741

742
743
744
    loop = asyncio.get_running_loop()

    awaits = {_next_task(pair[1], loop): pair for pair in enumerate(iterators)}
745
746
    try:
        while awaits:
747
748
            done, _ = await asyncio.wait(awaits.keys(),
                                         return_when=FIRST_COMPLETED)
749
750
751
752
753
            for d in done:
                pair = awaits.pop(d)
                try:
                    item = await d
                    i, it = pair
754
                    awaits[_next_task(it, loop)] = pair
755
756
757
758
759
760
761
762
763
                    yield i, item
                except StopAsyncIteration:
                    pass
    finally:
        # Cancel any remaining iterators
        for f, (_, it) in awaits.items():
            with contextlib.suppress(BaseException):
                f.cancel()
                await it.aclose()
764
765


766
async def collect_from_async_generator(
767
        iterator: AsyncGenerator[T, None]) -> list[T]:
768
769
770
771
772
773
774
    """Collect all items from an async generator into a list."""
    items = []
    async for item in iterator:
        items.append(item)
    return items


775
def get_ip() -> str:
776
    host_ip = envs.VLLM_HOST_IP
777
778
779
780
    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"
781
            " interact with the container's network stack. Please "
782
783
            "use VLLM_HOST_IP instead to set the IP address for vLLM processes"
            " to communicate with each other.")
784
785
786
787
788
    if host_ip:
        return host_ip

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

789
    # try ipv4
790
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
791
    try:
792
        s.connect(("8.8.8.8", 80))  # Doesn't need to be reachable
793
        return s.getsockname()[0]
794
795
796
797
798
    except Exception:
        pass

    # try ipv6
    try:
799
        s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
800
801
802
        # 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
803
        return s.getsockname()[0]
804
805
806
807
808
    except Exception:
        pass

    warnings.warn(
        "Failed to get the IP address, using 0.0.0.0 by default."
809
810
        "The value can be set by the environment variable"
        " VLLM_HOST_IP or HOST_IP.",
811
812
        stacklevel=2)
    return "0.0.0.0"
813
814


815
816
817
818
819
820
821
822
def is_valid_ipv6_address(address: str) -> bool:
    try:
        ipaddress.IPv6Address(address)
        return True
    except ValueError:
        return False


823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
def split_host_port(host_port: str) -> Tuple[str, int]:
    # ipv6
    if host_port.startswith('['):
        host, port = host_port.rsplit(']', 1)
        host = host[1:]
        port = port.split(':')[1]
        return host, int(port)
    else:
        host, port = host_port.split(':')
        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}"


842
def get_distributed_init_method(ip: str, port: int) -> str:
843
844
845
846
    return get_tcp_uri(ip, port)


def get_tcp_uri(ip: str, port: int) -> str:
847
848
849
850
    if is_valid_ipv6_address(ip):
        return f"tcp://[{ip}]:{port}"
    else:
        return f"tcp://{ip}:{port}"
851
852


853
854
855
856
857
def get_open_zmq_ipc_path() -> str:
    base_rpc_path = envs.VLLM_RPC_BASE_PATH
    return f"ipc://{base_rpc_path}/{uuid4()}"


858
859
860
861
def get_open_zmq_inproc_path() -> str:
    return f"inproc://{uuid4()}"


862
def get_open_port() -> int:
863
864
865
866
867
868
869
870
871
    """
    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:
872
873
        dp_master_port = envs.VLLM_DP_MASTER_PORT
        reserved_port_range = range(dp_master_port, dp_master_port + 10)
874
        while True:
875
876
877
            candidate_port = _get_open_port()
            if candidate_port not in reserved_port_range:
                return candidate_port
878
879
    return _get_open_port()

youkaichao's avatar
youkaichao committed
880

881
def _get_open_port() -> int:
882
    port = envs.VLLM_PORT
883
    if port is not None:
884
885
886
887
888
889
890
891
892
        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
                logger.info("Port %d is already in use, trying port %d",
                            port - 1, port)
893
894
895
896
897
898
899
900
901
902
    # 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]
903
904


905
def find_process_using_port(port: int) -> Optional[psutil.Process]:
906
907
908
909
910
911
912
    # 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

913
914
915
916
917
918
919
920
921
    for conn in psutil.net_connections():
        if conn.laddr.port == port:
            try:
                return psutil.Process(conn.pid)
            except psutil.NoSuchProcess:
                return None
    return None


922
def update_environment_variables(envs: dict[str, str]):
923
    for k, v in envs.items():
924
        if k in os.environ and os.environ[k] != v:
925
926
927
            logger.warning(
                "Overwriting environment variable %s "
                "from '%s' to '%s'", k, os.environ[k], v)
928
        os.environ[k] = v
929
930


931
def chunk_list(lst: list[T], chunk_size: int):
932
    """Yield successive chunk_size chunks from lst."""
933
934
    for i in range(0, len(lst), chunk_size):
        yield lst[i:i + chunk_size]
935
936
937
938
939
940
941


def cdiv(a: int, b: int) -> int:
    """Ceiling division."""
    return -(a // -b)


942
943
944
945
946
947
948
def next_power_of_2(n) -> int:
    """The next power of 2 (inclusive)"""
    if n < 1:
        return 1
    return 1 << (n - 1).bit_length()


949
950
951
952
def round_up(x: int, y: int) -> int:
    return ((x + y - 1) // y) * y


953
954
955
956
def round_down(x: int, y: int) -> int:
    return (x // y) * y


957
def _generate_random_fp8(
958
    tensor: torch.Tensor,
959
960
961
962
963
964
    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.
965
    # For example, s.11111.00 in fp8e5m2 format represents Inf.
966
    #     | E4M3        | E5M2
967
    # -----|-------------|-------------------
968
969
    # Inf | N/A         | s.11111.00
    # NaN | s.1111.111  | s.11111.{01,10,11}
970
    from vllm import _custom_ops as ops
971
972
    tensor_tmp = torch.empty_like(tensor, dtype=torch.float16)
    tensor_tmp.uniform_(low, high)
973
    ops.convert_fp8(tensor, tensor_tmp)
974
975
976
    del tensor_tmp


977
978
979
def get_kv_cache_torch_dtype(
        cache_dtype: Optional[Union[str, torch.dtype]],
        model_dtype: Optional[Union[str, torch.dtype]] = None) -> torch.dtype:
980
981
    if isinstance(cache_dtype, str):
        if cache_dtype == "auto":
982
983
            if isinstance(model_dtype,
                          str) and model_dtype in STR_DTYPE_TO_TORCH_DTYPE:
984
985
986
987
988
                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}")
989
        elif cache_dtype in STR_DTYPE_TO_TORCH_DTYPE:
990
991
992
993
994
995
996
            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}")
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
    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,
    cache_dtype: Optional[Union[str, torch.dtype]],
    model_dtype: Optional[Union[str, torch.dtype]] = None,
1008
    seed: Optional[int] = None,
1009
    device: Optional[str] = "cuda",
1010
    cache_layout: Optional[str] = "NHD",
1011
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
1012
    from vllm.platforms import current_platform
1013
    current_platform.seed_everything(seed)
1014
1015

    torch_dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
1016
1017
1018
1019
1020
1021
1022
    generic_kv_cache_shape = (num_blocks, 2, block_size, num_heads, head_size)
    assert cache_layout in ("NHD", "HND")
    stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3, 2,
                                                                  4)

    kv_cache_allocation_shape = tuple(generic_kv_cache_shape[i]
                                      for i in stride_order)
1023
    scale = head_size**-0.5
1024

1025
1026
    key_caches: list[torch.Tensor] = []
    value_caches: list[torch.Tensor] = []
1027

1028
    for _ in range(num_layers):
1029
        key_value_cache = torch.empty(size=kv_cache_allocation_shape,
1030
                                      dtype=torch_dtype,
1031
                                      device=device).permute(*stride_order)
1032
1033
1034
1035
1036
1037
1038
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
            key_value_cache.uniform_(-scale, scale)
        elif cache_dtype == 'fp8':
            _generate_random_fp8(key_value_cache, -scale, scale)
        else:
            raise ValueError(
                f"Does not support key cache of type {cache_dtype}")
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
        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,
    cache_dtype: Optional[Union[str, torch.dtype]],
    model_dtype: Optional[Union[str, torch.dtype]] = None,
1052
    seed: Optional[int] = None,
1053
    device: Optional[str] = "cuda",
1054
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
Joe's avatar
Joe committed
1055
1056
1057
1058
    if cache_dtype == "fp8" and head_size % 16:
        raise ValueError(
            f"Does not support key cache of type fp8 with head_size {head_size}"
        )
1059
    from vllm.platforms import current_platform
1060
    current_platform.seed_everything(seed)
1061
1062

    torch_dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
1063
1064
1065
1066

    scale = head_size**-0.5
    x = 16 // torch.tensor([], dtype=torch_dtype).element_size()
    key_cache_shape = (num_blocks, num_heads, head_size // x, block_size, x)
1067
    key_caches: list[torch.Tensor] = []
1068
1069
1070
1071
    for _ in range(num_layers):
        key_cache = torch.empty(size=key_cache_shape,
                                dtype=torch_dtype,
                                device=device)
1072
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
1073
            key_cache.uniform_(-scale, scale)
1074
1075
        elif cache_dtype == 'fp8':
            _generate_random_fp8(key_cache, -scale, scale)
1076
1077
1078
        else:
            raise ValueError(
                f"Does not support key cache of type {cache_dtype}")
1079
1080
1081
        key_caches.append(key_cache)

    value_cache_shape = (num_blocks, num_heads, head_size, block_size)
1082
    value_caches: list[torch.Tensor] = []
1083
1084
1085
1086
    for _ in range(num_layers):
        value_cache = torch.empty(size=value_cache_shape,
                                  dtype=torch_dtype,
                                  device=device)
1087
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
1088
            value_cache.uniform_(-scale, scale)
1089
1090
        elif cache_dtype == 'fp8':
            _generate_random_fp8(value_cache, -scale, scale)
1091
1092
1093
        else:
            raise ValueError(
                f"Does not support value cache of type {cache_dtype}")
1094
1095
        value_caches.append(value_cache)
    return key_caches, value_caches
1096
1097


1098
@cache
1099
def is_pin_memory_available() -> bool:
1100
    from vllm.platforms import current_platform
1101
    return current_platform.is_pin_memory_available()
1102
1103


1104
1105
1106
1107
1108
1109
1110
1111
@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()


1112
class DeviceMemoryProfiler:
1113

1114
    def __init__(self, device: Optional[torch.types.Device] = None):
1115
1116
1117
1118
        self.device = device

    def current_memory_usage(self) -> float:
        # Return the memory usage in bytes.
1119
        from vllm.platforms import current_platform
1120
        gc.collect()
1121
        return current_platform.get_current_memory_usage(self.device)
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133

    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()
1134
1135


1136
def make_ndarray_with_pad(
1137
    x: list[list[T]],
1138
1139
1140
1141
1142
1143
1144
    pad: T,
    dtype: npt.DTypeLike,
    *,
    max_len: Optional[int] = None,
) -> npt.NDArray:
    """
    Make a padded array from 2D inputs.
1145
1146
1147
1148

    The padding is applied to the end of each inner list until it reaches
    `max_len`.
    """
1149
1150
1151
1152
1153
    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)
1154
1155
1156
    for ind, blocktb in enumerate(x):
        assert len(blocktb) <= max_len
        padded_x[ind, :len(blocktb)] = blocktb
1157
1158
1159
1160
1161

    return padded_x


def make_tensor_with_pad(
1162
    x: list[list[T]],
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
    pad: T,
    dtype: torch.dtype,
    *,
    max_len: Optional[int] = None,
    device: Optional[Union[str, torch.device]] = None,
    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
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196


def async_tensor_h2d(
    data: list,
    dtype: torch.dtype,
    target_device: Union[str, torch.device],
    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)


1197
1198
1199
1200
1201
def get_dtype_size(dtype: torch.dtype) -> int:
    """Get the size of the data type in bytes."""
    return torch.tensor([], dtype=dtype).element_size()


1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
# bool = 0, int = 1, float = 2, complex = 3
def _get_precision_level(dtype: torch.dtype) -> int:
    # NOTE: Complex dtypes return `is_floating_point=False`
    return ((dtype != torch.bool) + dtype.is_floating_point +
            dtype.is_complex * 2)


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)
    return (src_info.min >= tgt_info.min and src_info.max <= tgt_info.max
            and src_info.resolution >= tgt_info.resolution)


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


1249
1250
1251
# `collections` helpers
def is_list_of(
    value: object,
1252
    typ: Union[type[T], tuple[type[T], ...]],
1253
1254
    *,
    check: Literal["first", "all"] = "first",
1255
) -> TypeIs[list[T]]:
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
    if not isinstance(value, list):
        return False

    if check == "first":
        return len(value) == 0 or isinstance(value[0], typ)
    elif check == "all":
        return all(isinstance(v, typ) for v in value)

    assert_never(check)


1267
def flatten_2d_lists(lists: Iterable[Iterable[T]]) -> list[T]:
1268
1269
1270
1271
    """Flatten a list of lists to a single list."""
    return [item for sublist in lists for item in sublist]


1272
1273
def full_groupby(values: Iterable[_V], *, key: Callable[[_V], _K]):
    """
1274
    Unlike [`itertools.groupby`][], groups are not broken by
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
    non-contiguous data.
    """
    groups = defaultdict[_K, list[_V]](list)

    for value in values:
        groups[key(value)].append(value)

    return groups.items()


1285
1286
# TODO: This function can be removed if transformer_modules classes are
# serialized by value when communicating between processes
1287
def init_cached_hf_modules() -> None:
1288
1289
1290
1291
1292
    """
    Lazy initialization of the Hugging Face modules.
    """
    from transformers.dynamic_module_utils import init_hf_modules
    init_hf_modules()
1293
1294


1295
@cache
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
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
1311
    env_ld_library_path = envs.LD_LIBRARY_PATH
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
    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]


1323
def find_nccl_library() -> str:
1324
1325
1326
1327
1328
1329
    """
    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.
    """
1330
    so_file = envs.VLLM_NCCL_SO_PATH
1331
1332
1333
1334

    # manually load the nccl library
    if so_file:
        logger.info(
1335
1336
            "Found nccl from environment variable VLLM_NCCL_SO_PATH=%s",
            so_file)
1337
1338
    else:
        if torch.version.cuda is not None:
1339
            so_file = "libnccl.so.2"
1340
        elif torch.version.hip is not None:
1341
            so_file = "librccl.so.1"
1342
1343
        else:
            raise ValueError("NCCL only supports CUDA and ROCm backends.")
1344
        logger.info("Found nccl from library %s", so_file)
1345
    return so_file
1346
1347


youkaichao's avatar
youkaichao committed
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
prev_set_stream = torch.cuda.set_stream

_current_stream = None


def _patched_set_stream(stream: torch.cuda.Stream) -> None:
    global _current_stream
    _current_stream = stream
    prev_set_stream(stream)


torch.cuda.set_stream = _patched_set_stream


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.
    """
1373
    from vllm.platforms import current_platform
youkaichao's avatar
youkaichao committed
1374
1375
1376
1377
    global _current_stream
    if _current_stream is None:
        # when this function is called before any stream is set,
        # we return the default stream.
1378
1379
1380
1381
1382
        # On ROCm using the default 0 stream in combination with RCCL
        # is hurting performance. Therefore creating a dedicated stream
        # per process
        _current_stream = torch.cuda.Stream() if current_platform.is_rocm(
        ) else torch.cuda.current_stream()
youkaichao's avatar
youkaichao committed
1383
1384
1385
    return _current_stream


1386
def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
1387
1388
1389
1390
    """Set up function tracing for the current thread,
    if enabled via the VLLM_TRACE_FUNCTION environment variable
    """

1391
    if envs.VLLM_TRACE_FUNCTION:
1392
        tmp_dir = tempfile.gettempdir()
1393
1394
        # add username to tmp_dir to avoid permission issues
        tmp_dir = os.path.join(tmp_dir, getpass.getuser())
1395
1396
1397
        filename = (f"VLLM_TRACE_FUNCTION_for_process_{os.getpid()}"
                    f"_thread_{threading.get_ident()}_"
                    f"at_{datetime.datetime.now()}.log").replace(" ", "_")
1398
1399
        log_path = os.path.join(tmp_dir, "vllm",
                                f"vllm-instance-{vllm_config.instance_id}",
1400
1401
1402
                                filename)
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        enable_trace_function_call(log_path)
1403
1404


1405
# `functools` helpers
1406
1407
def identity(value: T, **kwargs) -> T:
    """Returns the first provided value."""
1408
1409
1410
1411
1412
1413
    return value


F = TypeVar('F', bound=Callable[..., Any])


1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
def deprecate_args(
    start_index: int,
    is_deprecated: Union[bool, Callable[[], bool]] = True,
    additional_message: Optional[str] = None,
) -> Callable[[F], F]:
    if not callable(is_deprecated):
        is_deprecated = partial(identity, is_deprecated)

    def wrapper(fn: F) -> F:

        params = inspect.signature(fn).parameters
        pos_types = (
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
        )
        pos_kws = [
            kw for kw, param in params.items() if param.kind in pos_types
        ]

        @wraps(fn)
        def inner(*args, **kwargs):
            if is_deprecated():
                deprecated_args = pos_kws[start_index:len(args)]
                if deprecated_args:
                    msg = (
                        f"The positional arguments {deprecated_args} are "
                        "deprecated and will be removed in a future update.")
                    if additional_message is not None:
                        msg += f" {additional_message}"

                    warnings.warn(
                        DeprecationWarning(msg),
                        stacklevel=3,  # The inner function takes up one level
                    )

            return fn(*args, **kwargs)

        return inner  # type: ignore

    return wrapper


1456
def deprecate_kwargs(
1457
1458
1459
1460
    *kws: str,
    is_deprecated: Union[bool, Callable[[], bool]] = True,
    additional_message: Optional[str] = None,
) -> Callable[[F], F]:
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
    deprecated_kws = set(kws)

    if not callable(is_deprecated):
        is_deprecated = partial(identity, is_deprecated)

    def wrapper(fn: F) -> F:

        @wraps(fn)
        def inner(*args, **kwargs):
            if is_deprecated():
                deprecated_kwargs = kwargs.keys() & deprecated_kws
                if deprecated_kwargs:
                    msg = (
                        f"The keyword arguments {deprecated_kwargs} are "
                        "deprecated and will be removed in a future update.")
                    if additional_message is not None:
                        msg += f" {additional_message}"

                    warnings.warn(
                        DeprecationWarning(msg),
                        stacklevel=3,  # The inner function takes up one level
                    )

            return fn(*args, **kwargs)

        return inner  # type: ignore

    return wrapper
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503


@lru_cache(maxsize=8)
def _cuda_device_count_stateless(
        cuda_visible_devices: Optional[str] = None) -> int:
    # 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

1504
    from vllm.platforms import current_platform
1505
1506
    if not torch.cuda._is_compiled():
        return 0
1507
    if current_platform.is_rocm():
1508
1509
1510
1511
1512
1513
1514
        # ROCm uses amdsmi instead of nvml for stateless device count
        # This requires a sufficiently modern version of Torch 2.4.0
        raw_count = torch.cuda._device_count_amdsmi() if (hasattr(
            torch.cuda, "_device_count_amdsmi")) else -1
    else:
        raw_count = torch.cuda._device_count_nvml()
    r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count
1515
1516
1517
1518
1519
1520
    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.
1521

1522
1523
1524
1525
1526
1527
1528
    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)
1529
1530


1531
1532
1533
1534
1535
1536
1537
def cuda_is_initialized() -> bool:
    """Check if CUDA is initialized."""
    if not torch.cuda._is_compiled():
        return False
    return torch.cuda.is_initialized()


1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
def cuda_get_device_properties(device,
                               names: Sequence[str],
                               init_cuda=False) -> tuple[Any, ...]:
    """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:
        return executor.submit(cuda_get_device_properties, device, names,
                               True).result()


1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
def weak_bind(bound_method: Callable[..., Any], ) -> Callable[..., None]:
    """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


1568
# From: https://stackoverflow.com/a/4104188/2749989
1569
def run_once(f: Callable[P, None]) -> Callable[P, None]:
1570

1571
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> None:
1572
1573
1574
1575
1576
1577
        if not wrapper.has_run:  # type: ignore[attr-defined]
            wrapper.has_run = True  # type: ignore[attr-defined]
            return f(*args, **kwargs)

    wrapper.has_run = False  # type: ignore[attr-defined]
    return wrapper
1578
1579


1580
class StoreBoolean(Action):
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591

    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:
            raise ValueError(f"Invalid boolean value: {values}. "
                             "Expected 'true' or 'false'.")


1592
1593
class SortedHelpFormatter(ArgumentDefaultsHelpFormatter,
                          RawDescriptionHelpFormatter):
1594
1595
    """SortedHelpFormatter that sorts arguments by their option strings."""

1596
1597
1598
1599
1600
1601
1602
    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
1603
1604
        single_newline = re.compile(r"(?<!\n)\n(?!\n)\s*")
        multiple_newlines = re.compile(r"\n{2,}\s*")
1605
1606
1607
1608
        text = single_newline.sub(' ', text)
        lines = re.split(multiple_newlines, text)
        return sum([textwrap.wrap(line, width) for line in lines], [])

1609
1610
    def add_arguments(self, actions):
        actions = sorted(actions, key=lambda x: x.option_strings)
1611
        super().add_arguments(actions)
1612
1613


1614
class FlexibleArgumentParser(ArgumentParser):
1615
1616
    """ArgumentParser that allows both underscore and dash in names."""

1617
1618
    _deprecated: set[Action] = set()

1619
1620
1621
1622
1623
1624
    def __init__(self, *args, **kwargs):
        # Set the default 'formatter_class' to SortedHelpFormatter
        if 'formatter_class' not in kwargs:
            kwargs['formatter_class'] = SortedHelpFormatter
        super().__init__(*args, **kwargs)

1625
    if sys.version_info < (3, 13):
1626
        # Enable the deprecated kwarg for Python 3.12 and below
1627

1628
        def parse_known_args(self, args=None, namespace=None):
1629
1630
            namespace, args = super().parse_known_args(args, namespace)
            for action in FlexibleArgumentParser._deprecated:
1631
1632
1633
                if (hasattr(namespace, dest := action.dest)
                        and getattr(namespace, dest) != action.default):
                    logger.warning_once("argument '%s' is deprecated", dest)
1634
1635
            return namespace, args

1636
1637
        def add_argument(self, *args, **kwargs):
            deprecated = kwargs.pop("deprecated", False)
1638
            action = super().add_argument(*args, **kwargs)
1639
1640
            if deprecated:
                FlexibleArgumentParser._deprecated.add(action)
1641
1642
            return action

1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
        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
1656
1657
1658
1659
1660
1661

    def parse_args(  # type: ignore[override]
        self,
        args: list[str] | None = None,
        namespace: Namespace | None = None,
    ):
1662
1663
1664
        if args is None:
            args = sys.argv[1:]

1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
        # Check for --model in command line arguments first
        if args and args[0] == "serve":
            model_in_cli_args = any(arg == '--model' for arg in args)

            if model_in_cli_args:
                raise ValueError(
                    "With `vllm serve`, you should provide the model as a "
                    "positional argument or in a config file instead of via "
                    "the `--model` option.")

1675
        if '--config' in args:
1676
            args = self._pull_args_from_config(args)
1677

1678
1679
1680
1681
1682
1683
1684
        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"(?<=--)[^\.]*")

1685
        # Convert underscores to dashes and vice versa in argument names
1686
        processed_args = list[str]()
1687
        for i, arg in enumerate(args):
1688
            if arg.startswith('--'):
1689
1690
                if '=' in arg:
                    key, value = arg.split('=', 1)
1691
                    key = pattern.sub(repl, key, count=1)
1692
1693
                    processed_args.append(f'{key}={value}')
                else:
1694
1695
                    key = pattern.sub(repl, arg, count=1)
                    processed_args.append(key)
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
            elif arg.startswith('-O') and arg != '-O' and arg[2] != '.':
                # allow -O flag to be used without space, e.g. -O3 or -Odecode
                # -O.<...> handled later
                # also handle -O=<level> here
                level = arg[3:] if arg[2] == '=' else arg[2:]
                processed_args.append(f'-O.level={level}')
            elif arg == '-O' and i + 1 < len(args) and args[i + 1] in {
                    "0", "1", "2", "3"
            }:
                # Convert -O <n> to -O.level <n>
                processed_args.append('-O.level')
1707
1708
1709
            else:
                processed_args.append(arg)

1710
        def create_nested_dict(keys: list[str], value: str) -> dict[str, Any]:
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
            """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

1721
1722
1723
        def recursive_dict_update(
            original: dict[str, Any],
            update: dict[str, Any],
1724
1725
1726
1727
1728
        ) -> set[str]:
            """Recursively updates a dictionary with another dictionary.
            Returns a set of duplicate keys that were overwritten.
            """
            duplicates = set[str]()
1729
1730
            for k, v in update.items():
                if isinstance(v, dict) and isinstance(original.get(k), dict):
1731
1732
1733
1734
                    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
1735
                else:
1736
1737
                    if k in original:
                        duplicates.add(k)
1738
                    original[k] = v
1739
            return duplicates
1740

1741
1742
        delete = set[int]()
        dict_args = defaultdict[str, dict[str, Any]](dict)
1743
        duplicates = set[str]()
1744
        for i, processed_arg in enumerate(processed_args):
1745
1746
1747
1748
            if i in delete:  # skip if value from previous arg
                continue

            if processed_arg.startswith("-") and "." in processed_arg:
1749
                if "=" in processed_arg:
1750
                    processed_arg, value_str = processed_arg.split("=", 1)
1751
                    if "." not in processed_arg:
1752
                        # False positive, '.' was only in the value
1753
1754
                        continue
                else:
1755
                    value_str = processed_args[i + 1]
1756
                    delete.add(i + 1)
1757

1758
1759
1760
1761
                if processed_arg.endswith("+"):
                    processed_arg = processed_arg[:-1]
                    value_str = json.dumps(list(value_str.split(",")))

1762
                key, *keys = processed_arg.split(".")
1763
1764
1765
1766
1767
                try:
                    value = json.loads(value_str)
                except json.decoder.JSONDecodeError:
                    value = value_str

1768
1769
                # Merge all values with the same key into a single dict
                arg_dict = create_nested_dict(keys, value)
1770
1771
1772
                arg_duplicates = recursive_dict_update(dict_args[key],
                                                       arg_dict)
                duplicates |= {f'{key}.{d}' for d in arg_duplicates}
1773
1774
1775
1776
1777
                delete.add(i)
        # Filter out the dict args we set to None
        processed_args = [
            a for i, a in enumerate(processed_args) if i not in delete
        ]
1778
1779
1780
        if duplicates:
            logger.warning("Found duplicate keys %s", ", ".join(duplicates))

1781
1782
1783
1784
1785
        # 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))

1786
        return super().parse_args(processed_args, namespace)
1787

1788
1789
1790
1791
    def check_port(self, value):
        try:
            value = int(value)
        except ValueError:
1792
            msg = "Port must be an integer"
1793
            raise ArgumentTypeError(msg) from None
1794
1795

        if not (1024 <= value <= 65535):
1796
            raise ArgumentTypeError("Port must be between 1024 and 65535")
1797
1798
1799

        return value

1800
    def _pull_args_from_config(self, args: list[str]) -> list[str]:
1801
1802
        """Method to pull arguments specified in the config file
        into the command-line args variable.
1803
1804

        The arguments in config file will be inserted between
1805
        the argument list.
1806

1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
        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",
1817
1818
            "facebook/opt-12B",
            '--config', 'config.yaml',
1819
1820
1821
1822
            '-tp', '2'
        ]
        $: args = [
            "serve,chat,complete",
1823
1824
1825
            "facebook/opt-12B",
            '--port', '12323',
            '--tensor-parallel-size', '4',
1826
1827
1828
1829
1830
            '-tp', '2'
            ]
        ```

        Please note how the config args are inserted after the sub command.
1831
        this way the order of priorities is maintained when these are args
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
        parsed by super().
        """
        assert args.count(
            '--config') <= 1, "More than one config file specified!"

        index = args.index('--config')
        if index == len(args) - 1:
            raise ValueError("No config file specified! \
                             Please check your command-line arguments.")

        file_path = args[index + 1]

1844
        config_args = self._load_config_file(file_path)
1845
1846

        # 0th index is for {serve,chat,complete}
1847
        # optionally followed by model_tag (only for serve)
1848
1849
1850
1851
        # followed by config args
        # followed by rest of cli args.
        # maintaining this order will enforce the precedence
        # of cli > config > defaults
1852
        if args[0] == "serve":
1853
1854
1855
1856
            model_in_cli = len(args) > 1 and not args[1].startswith('-')
            model_in_config = any(arg == '--model' for arg in config_args)

            if not model_in_cli and not model_in_config:
1857
                raise ValueError(
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
                    "No model specified! Please specify model either "
                    "as a positional argument or in a config file.")

            if model_in_cli:
                # Model specified as positional arg, keep CLI version
                args = [args[0]] + [
                    args[1]
                ] + config_args + args[2:index] + args[index + 2:]
            else:
                # No model in CLI, use config if available
                args = [args[0]
                        ] + config_args + args[1:index] + args[index + 2:]
1870
1871
        else:
            args = [args[0]] + config_args + args[1:index] + args[index + 2:]
1872
1873
1874

        return args

1875
    def _load_config_file(self, file_path: str) -> list[str]:
1876
        """Loads a yaml file and returns the key value pairs as a
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
        flattened list with argparse like pattern
        ```yaml
            port: 12323
            tensor-parallel-size: 4
        ```
        returns:
            processed_args: list[str] = [
                '--port': '12323',
                '--tensor-parallel-size': '4'
            ]
        """
        extension: str = file_path.split('.')[-1]
        if extension not in ('yaml', 'yml'):
            raise ValueError(
                "Config file must be of a yaml/yml type.\
                              %s supplied", extension)

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

1897
        config: dict[str, Union[int, str]] = {}
1898
        try:
1899
            with open(file_path) as config_file:
1900
1901
1902
1903
1904
1905
1906
                config = yaml.safe_load(config_file)
        except Exception as ex:
            logger.error(
                "Unable to read the config file at %s. \
                Make sure path is correct", file_path)
            raise ex

1907
1908
1909
1910
1911
        store_boolean_arguments = [
            action.dest for action in self._actions
            if isinstance(action, StoreBoolean)
        ]

1912
        for key, value in config.items():
1913
1914
1915
1916
1917
1918
            if isinstance(value, bool) and key not in store_boolean_arguments:
                if value:
                    processed_args.append('--' + key)
            else:
                processed_args.append('--' + key)
                processed_args.append(str(value))
1919
1920
1921

        return processed_args

1922
1923
1924
1925
1926
1927

async def _run_task_with_lock(task: Callable, lock: asyncio.Lock, *args,
                              **kwargs):
    """Utility function to run async task in a lock"""
    async with lock:
        return await task(*args, **kwargs)
1928
1929


1930
1931
1932
def supports_kw(
    callable: Callable[..., object],
    kw_name: str,
1933
    *,
1934
1935
1936
1937
1938
1939
    requires_kw_only: bool = False,
    allow_var_kwargs: bool = True,
) -> bool:
    """Check if a keyword is a valid kwarg for a callable; if requires_kw_only
    disallows kwargs names that can also be positional arguments.
    """
1940
    params = inspect.signature(callable).parameters
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
    if not params:
        return False

    param_val = params.get(kw_name)

    # Types where the it may be valid, i.e., explicitly defined & nonvariadic
    passable_kw_types = set((inspect.Parameter.POSITIONAL_ONLY,
                             inspect.Parameter.POSITIONAL_OR_KEYWORD,
                             inspect.Parameter.KEYWORD_ONLY))

    if param_val:
        is_sig_param = param_val.kind in passable_kw_types
        # We want kwargs only, but this is passable as a positional arg
        if (requires_kw_only and is_sig_param
                and param_val.kind != inspect.Parameter.KEYWORD_ONLY):
            return False
        if ((requires_kw_only
             and param_val.kind == inspect.Parameter.KEYWORD_ONLY)
                or (not requires_kw_only and is_sig_param)):
            return True

    # If we're okay with var-kwargs, it's supported as long as
    # the kw_name isn't something like *args, **kwargs
    if allow_var_kwargs:
        # Get the last param; type is ignored here because params is a proxy
        # mapping, but it wraps an ordered dict, and they appear in order.
        # Ref: https://docs.python.org/3/library/inspect.html#inspect.Signature.parameters
        last_param = params[next(reversed(params))]  # type: ignore
        return (last_param.kind == inspect.Parameter.VAR_KEYWORD
                and last_param.name != kw_name)
1971

1972
1973
1974
1975
    return False


def resolve_mm_processor_kwargs(
1976
1977
    init_kwargs: Optional[Mapping[str, object]],
    inference_kwargs: Optional[Mapping[str, object]],
1978
    callable: Callable[..., object],
1979
1980
    *,
    requires_kw_only: bool = True,
1981
    allow_var_kwargs: bool = False,
1982
) -> dict[str, Any]:
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
    """Applies filtering to eliminate invalid mm_processor_kwargs, i.e.,
    those who are not explicit keywords to the given callable (of one is
    given; otherwise no filtering is done), then merges the kwarg dicts,
    giving priority to inference_kwargs if there are any collisions.

    In the case that no kwarg overrides are provided, returns an empty
    dict so that it can still be kwarg expanded into the callable later on.

    If allow_var_kwargs=True, allows for things that can be expanded into
    kwargs as long as they aren't naming collision for var_kwargs or potential
    positional arguments.
    """
    # Filter inference time multimodal processor kwargs provided
    runtime_mm_kwargs = get_allowed_kwarg_only_overrides(
        callable,
        overrides=inference_kwargs,
1999
2000
2001
        requires_kw_only=requires_kw_only,
        allow_var_kwargs=allow_var_kwargs,
    )
2002
2003
2004

    # Filter init time multimodal processor kwargs provided
    init_mm_kwargs = get_allowed_kwarg_only_overrides(
2005
2006
2007
2008
2009
        callable,
        overrides=init_kwargs,
        requires_kw_only=requires_kw_only,
        allow_var_kwargs=allow_var_kwargs,
    )
2010

2011
2012
2013
    # Merge the final processor kwargs, prioritizing inference
    # time values over the initialization time values.
    mm_processor_kwargs = {**init_mm_kwargs, **runtime_mm_kwargs}
2014

2015
    return mm_processor_kwargs
2016
2017


2018
2019
def get_allowed_kwarg_only_overrides(
    callable: Callable[..., object],
2020
    overrides: Optional[Mapping[str, object]],
2021
2022
    *,
    requires_kw_only: bool = True,
2023
    allow_var_kwargs: bool = False,
2024
) -> dict[str, Any]:
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
    """
    Given a callable which has one or more keyword only params and a dict
    mapping param names to values, drop values that can be not be kwarg
    expanded to overwrite one or more keyword-only args. This is used in a
    few places to handle custom processor overrides for multimodal models,
    e.g., for profiling when processor options provided by the user
    may affect the number of mm tokens per instance.

    Args:
        callable: Callable which takes 0 or more keyword only arguments.
2035
                  If None is provided, all overrides names are allowed.
2036
        overrides: Potential overrides to be used when invoking the callable.
2037
        allow_var_kwargs: Allows overrides that are expandable for var kwargs.
2038
2039
2040
2041
2042
2043
2044
2045
2046

    Returns:
        Dictionary containing the kwargs to be leveraged which may be used
        to overwrite one or more keyword only arguments when invoking the
        callable.
    """
    if not overrides:
        return {}

2047
2048
    # Drop any mm_processor_kwargs provided by the user that
    # are not kwargs, unless it can fit it var_kwargs param
2049
2050
2051
    filtered_overrides = {
        kwarg_name: val
        for kwarg_name, val in overrides.items()
2052
2053
        if supports_kw(callable,
                       kwarg_name,
2054
                       requires_kw_only=requires_kw_only,
2055
                       allow_var_kwargs=allow_var_kwargs)
2056
2057
2058
2059
2060
    }

    # If anything is dropped, log a warning
    dropped_keys = overrides.keys() - filtered_overrides.keys()
    if dropped_keys:
2061
2062
2063
        if requires_kw_only:
            logger.warning(
                "The following intended overrides are not keyword-only args "
2064
                "and will be dropped: %s", dropped_keys)
2065
2066
2067
        else:
            logger.warning(
                "The following intended overrides are not keyword args "
2068
                "and will be dropped: %s", dropped_keys)
2069
2070
2071
2072

    return filtered_overrides


2073
2074
2075
2076
2077
2078
# 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")
2079
2080


2081
2082
2083
2084
2085
2086
# Supports xccl with PyTorch versions >= 2.8.0 for XPU platform
def supports_xccl() -> bool:
    return is_torch_equal_or_newer(
        "2.8.0") and torch.distributed.is_xccl_available()


2087
2088
2089
2090
2091
2092
# 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")


2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
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
2116
2117
2118


# Adapted from: https://stackoverflow.com/a/47212782/5082708
2119
class LazyDict(Mapping[str, T], Generic[T]):
2120

2121
    def __init__(self, factory: dict[str, Callable[[], T]]):
2122
        self._factory = factory
2123
        self._dict: dict[str, T] = {}
2124

2125
    def __getitem__(self, key: str) -> T:
2126
2127
2128
2129
2130
2131
        if key not in self._dict:
            if key not in self._factory:
                raise KeyError(key)
            self._dict[key] = self._factory[key]()
        return self._dict[key]

2132
2133
2134
    def __setitem__(self, key: str, value: Callable[[], T]):
        self._factory[key] = value

2135
2136
2137
2138
2139
    def __iter__(self):
        return iter(self._factory)

    def __len__(self):
        return len(self._factory)
2140
2141


2142
class ClassRegistry(UserDict[type[T], _V]):
2143

2144
    def __getitem__(self, key: type[T]) -> _V:
2145
2146
2147
2148
2149
2150
2151
        for cls in key.mro():
            if cls in self.data:
                return self.data[cls]

        raise KeyError(key)

    def __contains__(self, key: object) -> bool:
2152
2153
2154
        return self.contains(key)

    def contains(self, key: object, *, strict: bool = False) -> bool:
2155
2156
2157
        if not isinstance(key, type):
            return False

2158
2159
2160
        if strict:
            return key in self.data

2161
2162
2163
        return any(cls in self.data for cls in key.mro())


2164
def weak_ref_tensor(tensor: Any) -> Any:
2165
2166
2167
2168
2169
    """
    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.
    """
2170
2171
2172
2173
    if isinstance(tensor, torch.Tensor):
        return torch.ops._C.weak_ref_tensor(tensor)
    else:
        return tensor
2174
2175
2176


def weak_ref_tensors(
2177
    tensors: Union[torch.Tensor, list[torch.Tensor], tuple[torch.Tensor]]
2178
) -> Union[torch.Tensor, list[Any], tuple[Any], Any]:
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
    """
    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)
    raise ValueError("Invalid type for tensors")
2190
2191


2192
2193
2194
2195
2196
2197
2198
2199
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)


2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
def import_from_path(module_name: str, file_path: Union[str, os.PathLike]):
    """
    Import a Python file according to its file path.

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

    assert spec.loader is not None

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


2219
@cache
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
def get_vllm_optional_dependencies():
    metadata = importlib.metadata.metadata("vllm")
    requirements = metadata.get_all("Requires-Dist", [])
    extras = metadata.get_all("Provides-Extra", [])

    return {
        extra: [
            re.split(r";|>=|<=|==", req)[0] for req in requirements
            if req.endswith(f'extra == "{extra}"')
        ]
        for extra in extras
    }


2234
2235
2236
2237
2238
class _PlaceholderBase:
    """
    Disallows downstream usage of placeholder modules.

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

2242
2243
    Info:
        [Special method lookup](https://docs.python.org/3/reference/datamodel.html#special-lookup)
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
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
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
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
    """

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

    # [Basic customization]

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

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

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

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

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

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

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

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

    # [Callable objects]

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

    # [Container types]

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

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

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

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

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

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

    # [Numeric Types]

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    # [Context managers]

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

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


class PlaceholderModule(_PlaceholderBase):
2391
2392
2393
2394
2395
2396
    """
    A placeholder object to use when a module does not exist.

    This enables more informative errors when trying to access attributes
    of a module that does not exists.
    """
2397
2398
2399
2400
2401
2402

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

        # Apply name mangling to avoid conflicting with module attributes
        self.__name = name
2403
2404
2405
2406
2407

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

    def __getattr__(self, key: str):
2408
        name = self.__name
2409
2410

        try:
2411
            importlib.import_module(name)
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
        except ImportError as exc:
            for extra, names in get_vllm_optional_dependencies().items():
                if name in names:
                    msg = f"Please install vllm[{extra}] for {extra} support"
                    raise ImportError(msg) from exc

            raise exc

        raise AssertionError("PlaceholderModule should not be used "
                             "when the original module can be imported")


2424
2425
2426
2427
2428
2429
2430
2431
class _PlaceholderModuleAttr(_PlaceholderBase):

    def __init__(self, module: PlaceholderModule, attr_path: str) -> None:
        super().__init__()

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

    def placeholder_attr(self, attr_path: str):
2434
2435
        return _PlaceholderModuleAttr(self.__module,
                                      f"{self.__attr_path}.{attr_path}")
2436
2437

    def __getattr__(self, key: str):
2438
        getattr(self.__module, f"{self.__attr_path}.{key}")
2439
2440
2441
2442
2443

        raise AssertionError("PlaceholderModule should not be used "
                             "when the original module can be imported")


2444
2445
2446
2447
2448
# create a library to hold the custom op
vllm_lib = Library("vllm", "FRAGMENT")  # noqa


def direct_register_custom_op(
2449
2450
2451
2452
2453
2454
        op_name: str,
        op_func: Callable,
        mutates_args: list[str],
        fake_impl: Optional[Callable] = None,
        target_lib: Optional[Library] = None,
        dispatch_key: str = "CUDA",
2455
        tags: tuple[torch.Tag, ...] = (),
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
):
    """
    `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.
    """
2472
    if not supports_custom_op():
2473
        from vllm.platforms import current_platform
2474
2475
2476
2477
2478
2479
2480
2481
        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 "
            "the required dependencies.")
        return

2482
2483
2484
2485
2486
2487
2488
2489
    import torch.library
    if hasattr(torch.library, "infer_schema"):
        schema_str = torch.library.infer_schema(op_func,
                                                mutates_args=mutates_args)
    else:
        # for pytorch 2.4
        import torch._custom_op.impl
        schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
2490
    my_lib = target_lib or vllm_lib
2491
    my_lib.define(op_name + schema_str, tags=tags)
2492
    my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
2493
2494
    if fake_impl is not None:
        my_lib._register_fake(op_name, fake_impl)
2495
2496
2497
2498
2499
2500
2501
2502
2503


def resolve_obj_by_qualname(qualname: str) -> Any:
    """
    Resolve an object by its fully qualified name.
    """
    module_name, obj_name = qualname.rsplit(".", 1)
    module = importlib.import_module(module_name)
    return getattr(module, obj_name)
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528


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)
2529
2530
2531
2532
2533


@dataclass
class MemorySnapshot:
    """Memory snapshot."""
2534
    torch_peak: int = 0
2535
2536
    free_memory: int = 0
    total_memory: int = 0
2537
2538
2539
    cuda_memory: int = 0
    torch_memory: int = 0
    non_torch_memory: int = 0
2540
    timestamp: float = 0.0
2541
2542
2543
2544
2545
    auto_measure: bool = True

    def __post_init__(self):
        if self.auto_measure:
            self.measure()
2546
2547

    def measure(self):
2548
2549
2550
2551
2552
2553
2554
2555
        # 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.
        self.torch_peak = torch.cuda.memory_stats().get(
            "allocated_bytes.all.peak", 0)

2556
2557
        self.free_memory, self.total_memory = torch.cuda.mem_get_info()
        self.cuda_memory = self.total_memory - self.free_memory
2558

2559
2560
        # torch.cuda.memory_reserved() is how many bytes
        # PyTorch gets from cuda (by calling cudaMalloc, etc.)
2561
2562
2563
2564
        # 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
2565
2566
        self.timestamp = time.time()

2567
    def __sub__(self, other: MemorySnapshot) -> MemorySnapshot:
2568
        return MemorySnapshot(
2569
            torch_peak=self.torch_peak - other.torch_peak,
2570
2571
            free_memory=self.free_memory - other.free_memory,
            total_memory=self.total_memory - other.total_memory,
2572
2573
2574
2575
2576
2577
            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,
        )
2578
2579
2580
2581


@dataclass
class MemoryProfilingResult:
2582
2583
2584
2585
2586
2587
2588
    """Memory profiling result. All numbers are in bytes.
    """
    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)
2589
2590
2591
2592
    before_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
    after_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
    profile_time: float = 0.0

2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
    def __repr__(self) -> str:
        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.")

2603
2604
2605

@contextlib.contextmanager
def memory_profiling(
2606
2607
        baseline_snapshot: MemorySnapshot,
        weights_memory: int) -> Generator[MemoryProfilingResult, None, None]:
2608
    """Memory profiling context manager.
2609
2610
    baseline_snapshot: the memory snapshot before the current vLLM instance.
    weights_memory: memory used by PyTorch when loading the model weights.
2611
2612
        Note that, before loading the model weights, we also initialize the device
        and distributed environment, which may consume some memory. This part is not
2613
        included in the weights_memory because PyTorch does not control it.
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647

    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)

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

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

2652
    The increase of `non_torch_memory` from creating the current vLLM instance until after profiling to get (c.).
2653
    """  # noqa
2654
2655
    gc.collect()
    torch.cuda.empty_cache()
2656
2657
2658
2659
    torch.cuda.reset_peak_memory_stats()

    result = MemoryProfilingResult()

2660
    result.before_create = baseline_snapshot
2661
    # the part of memory used for holding the model weights
2662
    result.weights_memory = weights_memory
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672

    result.before_profile.measure()

    yield result

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

    result.after_profile.measure()

2673
2674
2675
2676
2677
2678
    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
    result.non_kv_cache_memory = result.non_torch_increase + result.torch_peak_increase + result.weights_memory  # noqa
2679
2680


2681
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 # noqa: E501
2682
def set_ulimit(target_soft_limit=65535):
2683
2684
2685
2686
2687
    if sys.platform.startswith('win'):
        logger.info("Windows detected, skipping ulimit adjustment.")
        return

    import resource
2688
2689
2690
2691
2692
2693
2694
2695
2696
    resource_type = resource.RLIMIT_NOFILE
    current_soft, current_hard = resource.getrlimit(resource_type)

    if current_soft < target_soft_limit:
        try:
            resource.setrlimit(resource_type,
                               (target_soft_limit, current_hard))
        except ValueError as e:
            logger.warning(
2697
2698
                "Found ulimit of %s and failed to automatically increase "
                "with error %s. This can cause fd limit errors like "
2699
2700
                "`OSError: [Errno 24] Too many open files`. Consider "
                "increasing with ulimit -n", current_soft, e)
2701
2702
2703
2704
2705
2706
2707
2708
2709


# 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


2710
def split_zmq_path(path: str) -> tuple[str, str, str]:
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
    """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


2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
def make_zmq_path(scheme: str, host: str, port: Optional[int] = None) -> str:
    """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.
    """
2742
    if port is None:
2743
2744
2745
2746
2747
2748
        return f"{scheme}://{host}"
    if is_valid_ipv6_address(host):
        return f"{scheme}://[{host}]:{port}"
    return f"{scheme}://{host}:{port}"


2749
2750
2751
2752
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501
def make_zmq_socket(
    ctx: Union[zmq.asyncio.Context, zmq.Context],  # type: ignore[name-defined]
    path: str,
2753
    socket_type: Any,
2754
2755
    bind: Optional[bool] = None,
    identity: Optional[bytes] = None,
2756
    linger: Optional[int] = None,
2757
2758
2759
2760
) -> Union[zmq.Socket, zmq.asyncio.Socket]:  # type: ignore[name-defined]
    """Make a ZMQ socket with the proper bind/connect semantics."""

    mem = psutil.virtual_memory()
2761
    socket = ctx.socket(socket_type)
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774

    # 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
    if total_mem > 32 and available_mem > 16:
        buf_size = int(0.5 * 1024**3)  # 0.5GB in bytes
    else:
        buf_size = -1  # Use system default buffer size

2775
    if bind is None:
2776
        bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB)
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788

    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)

2789
2790
2791
    if linger is not None:
        socket.setsockopt(zmq.LINGER, linger)

2792
2793
2794
2795
2796
2797
    # 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)

2798
    if bind:
2799
        socket.bind(path)
2800
    else:
2801
        socket.connect(path)
2802
2803
2804
2805
2806

    return socket


@contextlib.contextmanager
2807
2808
2809
def zmq_socket_ctx(
    path: str,
    socket_type: Any,
2810
    bind: Optional[bool] = None,
2811
    linger: int = 0,
2812
    identity: Optional[bytes] = None,
2813
) -> Iterator[zmq.Socket]:
2814
2815
    """Context manager for a ZMQ socket"""

2816
    ctx = zmq.Context()  # type: ignore[attr-defined]
2817
    try:
2818
2819
2820
2821
2822
        yield make_zmq_socket(ctx,
                              path,
                              socket_type,
                              bind=bind,
                              identity=identity)
2823
2824
2825
2826
    except KeyboardInterrupt:
        logger.debug("Got Keyboard Interrupt.")

    finally:
2827
        ctx.destroy(linger=linger)
2828
2829


2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
def is_in_ray_actor():
    """Check if we are in a Ray actor."""

    try:
        import ray
        return (ray.is_initialized()
                and ray.get_runtime_context().get_actor_id() is not None)
    except ImportError:
        return False


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

    reason = None
    if cuda_is_initialized():
        reason = "CUDA is initialized"
    elif is_in_ray_actor():
2852
2853
2854
2855
2856
        # 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
        os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
2857
2858
2859
2860
2861
2862
        reason = "In a Ray actor and can only be spawned"

    if reason is not None:
        logger.warning(
            "We must use the `spawn` multiprocessing start method. "
            "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "
2863
            "See https://docs.vllm.ai/en/latest/usage/"
2864
2865
            "troubleshooting.html#python-multiprocessing "
            "for more information. Reason: %s", reason)
2866
2867
2868
2869
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"


def get_mp_context():
2870
2871
2872
2873
2874
2875
2876
    """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()
2877
2878
    mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
    return multiprocessing.get_context(mp_method)
2879
2880
2881


def bind_kv_cache(
2882
2883
        ctx: dict[str, Any],
        kv_cache: list[list[torch.Tensor]],  # [virtual_engine][layer_index]
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
) -> 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
    from vllm.attention import AttentionType
    from vllm.model_executor.models.utils import extract_layer_index
    layer_need_kv_cache = [
        layer_name for layer_name in ctx
2899
2900
        if (hasattr(ctx[layer_name], 'attn_type') and ctx[layer_name].attn_type
            in (AttentionType.DECODER, AttentionType.ENCODER_DECODER))
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
    ]
    layer_index_sorted = sorted(
        set(
            extract_layer_index(layer_name)
            for layer_name in layer_need_kv_cache))
    for layer_name in layer_need_kv_cache:
        kv_cache_idx = layer_index_sorted.index(
            extract_layer_index(layer_name))
        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]
2913
2914


2915
2916
def run_method(obj: Any, method: Union[str, bytes, Callable], args: tuple[Any],
               kwargs: dict[str, Any]) -> Any:
2917
2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
    """
    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:
            raise NotImplementedError(f"Method {method!r} is not"
                                      " implemented.") from None
    else:
        func = partial(method, obj)  # type: ignore
    return func(*args, **kwargs)
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955


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.
2956
2957
2958
2959
2960
2961
2962
    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.
2963
    """
2964
2965
    import vllm.third_party.pynvml as pynvml
    return pynvml
2966
2967


2968
def warn_for_unimplemented_methods(cls: type[T]) -> type[T]:
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999
    """
    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
            if attr_name.startswith('_'):
                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:
            method_names = ','.join(unimplemented_methods)
            msg = (f"Methods {method_names} not implemented in {self}")
3000
            logger.debug(msg)
3001
3002
3003
3004
3005
3006
3007
3008

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

    type.__setattr__(cls, '__init__', wrapped_init)
    return cls
3009
3010
3011
3012
3013
3014
3015
3016
3017
3018
3019
3020
3021
3022
3023
3024
3025
3026
3027
3028
3029
3030
3031
3032
3033
3034
3035
3036
3037
3038
3039
3040
3041
3042
3043
3044
3045
3046
3047
3048
3049
3050
3051
3052
3053
3054
3055
3056
3057
3058
3059


class LazyLoader(types.ModuleType):
    """
    LazyLoader module borrowed from Tensorflow
    https://github.com/tensorflow/tensorflow/blob/main/tensorflow/python/util/lazy_loader.py
    with a addition of "module caching".

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

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

        super().__init__(str(name))

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

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

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

    def __dir__(self) -> list[str]:
        if self._module is None:
            self._module = self._load()
        return dir(self._module)
3060
3061
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
3073
3074
3075


def swap_dict_values(obj: dict[_K, _V], key1: _K, key2: _K) -> None:
    """
    Helper function to swap values for two keys
    """
    v1 = obj.get(key1)
    v2 = obj.get(key2)
    if v1 is not None:
        obj[key2] = v1
    else:
        obj.pop(key2, None)
    if v2 is not None:
        obj[key1] = v2
    else:
        obj.pop(key1, None)
3076
3077
3078
3079
3080
3081
3082
3083
3084
3085
3086
3087
3088
3089
3090
3091
3092
3093
3094
3095
3096
3097
3098
3099
3100
3101
3102
3103
3104
3105
3106
3107
3108
3109
3110
3111
3112
3113
3114
3115
3116
3117
3118
3119
3120
3121
3122
3123


@contextlib.contextmanager
def cprofile_context(save_file: Optional[str] = None):
    """Run a cprofile

    Args:
        save_file: path to save the profile result. "1" or
          None will result in printing to stdout.
    """
    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")


def cprofile(save_file: Optional[str] = None, enabled: bool = True):
    """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
3124
3125


3126
3127
# Only relevant for models using ALiBi (e.g, MPT)
def check_use_alibi(model_config: ModelConfig) -> bool:
3128
3129
    cfg = model_config.hf_text_config
    return (getattr(cfg, "alibi", False)  # Falcon
3130
3131
            or ("BloomForCausalLM" in getattr(model_config.hf_config,
                                              "architectures", []))  # Bloom
3132
3133
3134
3135
3136
3137
3138
            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)))))
3139
3140


3141
3142
3143
3144
3145
3146
3147
3148
3149
3150
3151
3152
3153
3154
3155
3156
def sha256(input) -> int:
    """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:
        An integer representing the SHA-256 hash of the serialized input.
    """
    input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
    return int.from_bytes(hashlib.sha256(input_bytes).digest(),
                          byteorder="big")
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
3167
3168


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:
3169
        return _is_torch_equal_or_newer(str(torch.__version__), target)
3170
3171
3172
    except Exception:
        # Fallback to PKG-INFO to load the package info, needed by the doc gen.
        return Version(importlib.metadata.version('torch')) >= Version(target)
3173
3174
3175
3176
3177
3178


# 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)
3179
3180
3181
3182
3183
3184
3185
3186
3187
3188
3189
3190
3191
3192
3193
3194
3195
3196
3197
3198
3199
3200
3201
3202
3203
3204
3205


@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."""

3206
    return _has_module("deep_gemm")