A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Chore] Separate out `vllm.utils.mem_utils` (#27143)
Signed-off-by: iAmir97 <Amir.balwel@embeddedllm.com> Signed-off-by: iAmir97 <71513472+iAmir97@users.noreply.github.com> Co-authored-by: iAmir97 <Amir.balwel@embeddedllm.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
iAmir97 committed
Oct 18, 2025 at 17:06 UTC
1d165d6d859d3c50720f0c07209db2363c4fd33b
19 files changed
+268
-252
tests/basic_correctness/test_cumem.py
+1
-1
@@ -6,7 +6,7 @@ import torch
6
7
from vllm import LLM, SamplingParams
8
from vllm.device_allocator.cumem import CuMemAllocator
9
-from vllm.utils import GiB_bytes
9
+from vllm.utils.mem_constants import GiB_bytes
10
11
from ..utils import create_new_process_for_each_test
12
tests/kernels/attention/test_attention.py
+1
-1
@@ -11,7 +11,7 @@ from tests.kernels.utils import opcheck
11
from vllm import _custom_ops as ops
12
from vllm.attention.layer import Attention, MultiHeadAttention
13
from vllm.platforms import current_platform
14
-from vllm.utils import get_max_shared_memory_bytes
14
+from vllm.utils.mem_utils import get_max_shared_memory_bytes
15
16
if not current_platform.is_rocm():
17
from xformers import ops as xops
tests/models/test_initialization.py
+1
-1
@@ -7,7 +7,7 @@ from unittest.mock import patch
7
import pytest
8
9
from vllm import LLM
10
-from vllm.utils import GiB_bytes
10
+from vllm.utils.mem_constants import GiB_bytes
11
from vllm.v1.core.kv_cache_utils import (
12
generate_scheduler_kv_cache_config,
13
get_kv_cache_configs,
tests/utils.py
+1
-1
@@ -46,10 +46,10 @@ from vllm.platforms import current_platform
46
from vllm.transformers_utils.tokenizer import get_tokenizer
47
from vllm.utils import (
48
FlexibleArgumentParser,
49
- GB_bytes,
49
cuda_device_count_stateless,
50
get_open_port,
51
)
52
+from vllm.utils.mem_constants import GB_bytes
53
54
if current_platform.is_rocm():
55
from amdsmi import (
tests/utils_/test_utils.py
+1
-2
@@ -23,7 +23,6 @@ from vllm.transformers_utils.detokenizer_utils import convert_ids_list_to_tokens
23
24
from vllm.utils import (
25
FlexibleArgumentParser,
26
- MemorySnapshot,
26
bind_kv_cache,
27
common_broadcastable_dtype,
28
current_stream,
@@ -33,13 +32,13 @@ from vllm.utils import (
32
join_host_port,
33
make_zmq_path,
34
make_zmq_socket,
36
- memory_profiling,
35
sha256,
36
split_host_port,
37
split_zmq_path,
38
unique_filepath,
39
)
40
41
+from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
42
from ..utils import create_new_process_for_each_test, flat_product
43
44
tests/v1/core/test_kv_cache_utils.py
+2
-1
@@ -14,7 +14,8 @@ from vllm.multimodal.inputs import (
14
PlaceholderRange,
15
)
16
from vllm.sampling_params import SamplingParams
17
-from vllm.utils import GiB_bytes, sha256, sha256_cbor
17
+from vllm.utils import sha256, sha256_cbor
18
+from vllm.utils.mem_constants import GiB_bytes
19
from vllm.v1.core.kv_cache_manager import KVCacheManager
20
from vllm.v1.core.kv_cache_utils import (
21
BlockHash,
tests/v1/tpu/worker/test_tpu_model_runner.py
+1
-1
@@ -13,7 +13,7 @@ from vllm.config import (
13
)
14
from vllm.pooling_params import PoolingParams
15
from vllm.sampling_params import SamplingParams
16
-from vllm.utils import GiB_bytes
16
+from vllm.utils.mem_constants import GiB_bytes
17
from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
18
from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
19
from vllm.v1.worker.tpu_model_runner import (
tests/v1/worker/test_gpu_model_runner.py
+2
-1
@@ -21,7 +21,8 @@ from vllm.distributed.parallel_state import (
21
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
22
from vllm.platforms import current_platform
23
from vllm.sampling_params import SamplingParams
24
-from vllm.utils import GiB_bytes, update_environment_variables
24
+from vllm.utils import update_environment_variables
25
+from vllm.utils.mem_constants import GiB_bytes
26
from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
27
from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
28
from vllm.v1.kv_cache_interface import (
tests/v1/worker/test_worker_memory_snapshot.py
+1
-1
@@ -11,7 +11,7 @@ import pytest
11
import torch
12
13
from vllm.engine.arg_utils import EngineArgs
14
-from vllm.utils import MemorySnapshot
14
+from vllm.utils.mem_utils import MemorySnapshot
15
from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
16
17
# Global queue to track operation order across processes
vllm/config/cache.py
+2
-1
@@ -10,7 +10,8 @@ from pydantic.dataclasses import dataclass
10
11
from vllm.config.utils import config
12
from vllm.logger import init_logger
13
-from vllm.utils import GiB_bytes, get_cpu_memory
13
+from vllm.utils.mem_constants import GiB_bytes
14
+from vllm.utils.mem_utils import get_cpu_memory
15
16
if TYPE_CHECKING:
17
from vllm.config.parallel import ParallelConfig
vllm/engine/arg_utils.py
+2
-1
@@ -81,7 +81,8 @@ from vllm.transformers_utils.config import (
81
maybe_override_with_speculators,
82
)
83
from vllm.transformers_utils.utils import check_gguf_file
84
-from vllm.utils import FlexibleArgumentParser, GiB_bytes, get_ip, is_in_ray_actor
84
+from vllm.utils import FlexibleArgumentParser, get_ip, is_in_ray_actor
85
+from vllm.utils.mem_constants import GiB_bytes
86
from vllm.v1.sample.logits_processor import LogitsProcessor
87
88
if TYPE_CHECKING:
vllm/multimodal/cache.py
+1
-1
@@ -17,9 +17,9 @@ from vllm.distributed.device_communicators.shm_object_storage import (
17
SingleWriterShmRingBuffer,
18
)
19
from vllm.logger import init_logger
20
-from vllm.utils import GiB_bytes, MiB_bytes
20
from vllm.utils.cache import CacheInfo, LRUCache
21
from vllm.utils.jsontree import json_count_leaves, json_map_leaves, json_reduce_leaves
22
+from vllm.utils.mem_constants import GiB_bytes, MiB_bytes
23
24
from .inputs import (
25
MultiModalBatchedField,
vllm/platforms/cpu.py
+1
-1
@@ -151,7 +151,7 @@ class CpuPlatform(Platform):
151
@classmethod
152
def get_device_total_memory(cls, device_id: int = 0) -> int:
153
import vllm.envs as envs
154
- from vllm.utils import GiB_bytes
154
+ from vllm.utils.mem_constants import GiB_bytes
155
156
kv_cache_space = envs.VLLM_CPU_KVCACHE_SPACE
157
if kv_cache_space is None:
vllm/utils/__init__.py
-234
@@ -4,7 +4,6 @@
4
import contextlib
5
import datetime
6
import enum
7
-import gc
7
import getpass
8
import hashlib
9
import importlib
@@ -21,7 +20,6 @@ import sys
20
import tempfile
21
import textwrap
22
import threading
24
-import time
23
import traceback
24
import uuid
25
import warnings
@@ -38,12 +36,10 @@ from collections import defaultdict
36
from collections.abc import (
37
Callable,
38
Collection,
41
- Generator,
39
Iterator,
40
Sequence,
41
)
42
from concurrent.futures.process import ProcessPoolExecutor
46
-from dataclasses import dataclass, field
43
from functools import cache, lru_cache, partial, wraps
44
from pathlib import Path
45
from typing import TYPE_CHECKING, Any, TextIO, TypeVar
@@ -58,7 +54,6 @@ import psutil
54
import regex as re
55
import setproctitle
56
import torch
61
-import torch.types
57
import yaml
58
import zmq
59
import zmq.asyncio
@@ -105,17 +100,6 @@ STR_XFORMERS_ATTN_VAL: str = "XFORMERS"
100
STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
101
STR_INVALID_VAL: str = "INVALID"
102
108
-MB_bytes = 1_000_000
109
-"""The number of bytes in one megabyte (MB)."""
110
-
111
-MiB_bytes = 1 << 20
112
-"""The number of bytes in one mebibyte (MiB)."""
113
-
114
-GB_bytes = 1_000_000_000
115
-"""The number of bytes in one gigabyte (GB)."""
116
-
117
-GiB_bytes = 1 << 30
118
-"""The number of bytes in one gibibyte (GiB)."""
103
104
# ANSI color codes
105
CYAN = "\033[1;36m"
@@ -180,23 +164,6 @@ class Counter:
164
self.counter = 0
165
166
183
-@cache
184
-def get_max_shared_memory_bytes(gpu: int = 0) -> int:
185
- """Returns the maximum shared memory per thread block in bytes."""
186
- from vllm import _custom_ops as ops
187
-
188
- max_shared_mem = ops.get_max_shared_memory_per_block_device_attribute(gpu)
189
- # value 0 will cause MAX_SEQ_LEN become negative and test_attention.py
190
- # will fail
191
- assert max_shared_mem > 0, "max_shared_mem can not be zero"
192
- return int(max_shared_mem)
193
-
194
-
195
-def get_cpu_memory() -> int:
196
- """Returns the total CPU memory of the node in bytes."""
197
- return psutil.virtual_memory().total
198
-
199
-
167
def random_uuid() -> str:
168
return str(uuid.uuid4().hex)
169
@@ -581,30 +548,6 @@ def is_uva_available() -> bool:
548
return is_pin_memory_available()
549
550
584
-class DeviceMemoryProfiler:
585
- def __init__(self, device: torch.types.Device | None = None):
586
- self.device = device
587
-
588
- def current_memory_usage(self) -> float:
589
- # Return the memory usage in bytes.
590
- from vllm.platforms import current_platform
591
-
592
- gc.collect()
593
- return current_platform.get_current_memory_usage(self.device)
594
-
595
- def __enter__(self):
596
- self.initial_memory = self.current_memory_usage()
597
- # This allows us to call methods of the context manager if needed
598
- return self
599
-
600
- def __exit__(self, exc_type, exc_val, exc_tb):
601
- self.final_memory = self.current_memory_usage()
602
- self.consumed_memory = self.final_memory - self.initial_memory
603
-
604
- # Force garbage collection
605
- gc.collect()
606
-
607
-
551
def make_ndarray_with_pad(
552
x: list[list[T]],
553
pad: T,
@@ -1642,183 +1585,6 @@ def kill_process_tree(pid: int):
1585
os.kill(pid, signal.SIGKILL)
1586
1587
1645
-@dataclass
1646
-class MemorySnapshot:
1647
- """Memory snapshot."""
1648
-
1649
- torch_peak: int = 0
1650
- free_memory: int = 0
1651
- total_memory: int = 0
1652
- cuda_memory: int = 0
1653
- torch_memory: int = 0
1654
- non_torch_memory: int = 0
1655
- timestamp: float = 0.0
1656
- auto_measure: bool = True
1657
-
1658
- def __post_init__(self):
1659
- if self.auto_measure:
1660
- self.measure()
1661
-
1662
- def measure(self):
1663
- from vllm.platforms import current_platform
1664
-
1665
- # we measure the torch peak memory usage via allocated_bytes,
1666
- # rather than `torch.cuda.memory_reserved()` .
1667
- # After `torch.cuda.reset_peak_memory_stats()`,
1668
- # `torch.cuda.memory_reserved()` will keep growing, and only shrink
1669
- # when we call `torch.cuda.empty_cache()` or OOM happens.
1670
- self.torch_peak = torch.cuda.memory_stats().get("allocated_bytes.all.peak", 0)
1671
-
1672
- self.free_memory, self.total_memory = torch.cuda.mem_get_info()
1673
- shared_sysmem_device_mem_sms = ((8, 7), (11, 0), (12, 1)) # Orin, Thor, Spark
1674
- if (
1675
- current_platform.is_cuda()
1676
- and current_platform.get_device_capability() in shared_sysmem_device_mem_sms
1677
- ):
1678
- # On UMA (Orin, Thor and Spark) platform,
1679
- # where both CPU and GPU rely on system memory,
1680
- # the cudaMemGetInfo function shows the amount of free system memory
1681
- # rather than what’s actually available.
1682
- # In the case,
1683
- # torch.cuda.mem_get_info() only reports "free" memory,
1684
- # which can be lower than what is actually
1685
- # available due to not including cache memory.
1686
- # There’s also a comprehensive reference page
1687
- # that explains how you can compute the proper value yourself.
1688
- # https://docs.nvidia.com/cuda/cuda-for-tegra-appnote/#estimating-total-allocatable-device-memory-on-an-integrated-gpu-device
1689
- self.free_memory = psutil.virtual_memory().available
1690
-
1691
- self.cuda_memory = self.total_memory - self.free_memory
1692
-
1693
- # torch.cuda.memory_reserved() is how many bytes
1694
- # PyTorch gets from cuda (by calling cudaMalloc, etc.)
1695
- # this is used to measure the non-torch memory usage
1696
- self.torch_memory = torch.cuda.memory_reserved()
1697
-
1698
- self.non_torch_memory = self.cuda_memory - self.torch_memory
1699
- self.timestamp = time.time()
1700
-
1701
- def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot":
1702
- return MemorySnapshot(
1703
- torch_peak=self.torch_peak - other.torch_peak,
1704
- free_memory=self.free_memory - other.free_memory,
1705
- total_memory=self.total_memory - other.total_memory,
1706
- cuda_memory=self.cuda_memory - other.cuda_memory,
1707
- torch_memory=self.torch_memory - other.torch_memory,
1708
- non_torch_memory=self.non_torch_memory - other.non_torch_memory,
1709
- timestamp=self.timestamp - other.timestamp,
1710
- auto_measure=False,
1711
- )
1712
-
1713
-
1714
-@dataclass
1715
-class MemoryProfilingResult:
1716
- """Memory profiling result. All numbers are in bytes."""
1717
-
1718
- non_kv_cache_memory: int = 0
1719
- torch_peak_increase: int = 0
1720
- non_torch_increase: int = 0
1721
- weights_memory: float = 0
1722
- before_create: MemorySnapshot = field(default_factory=MemorySnapshot)
1723
- before_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
1724
- after_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
1725
- profile_time: float = 0.0
1726
-
1727
- def __repr__(self) -> str:
1728
- return (
1729
- f"Memory profiling takes {self.profile_time:.2f} seconds. "
1730
- f"Total non KV cache memory: "
1731
- f"{(self.non_kv_cache_memory / GiB_bytes):.2f}GiB; "
1732
- f"torch peak memory increase: "
1733
- f"{(self.torch_peak_increase / GiB_bytes):.2f}GiB; "
1734
- f"non-torch forward increase memory: "
1735
- f"{(self.non_torch_increase / GiB_bytes):.2f}GiB; "
1736
- f"weights memory: {(self.weights_memory / GiB_bytes):.2f}GiB."
1737
- )
1738
-
1739
-
1740
-@contextlib.contextmanager
1741
-def memory_profiling(
1742
- baseline_snapshot: MemorySnapshot, weights_memory: int
1743
-) -> Generator[MemoryProfilingResult, None, None]:
1744
- """Memory profiling context manager.
1745
- baseline_snapshot: the memory snapshot before the current vLLM instance.
1746
- weights_memory: memory used by PyTorch when loading the model weights.
1747
- Note that, before loading the model weights, we also initialize the device
1748
- and distributed environment, which may consume some memory. This part is not
1749
- included in the weights_memory because PyTorch does not control it.
1750
-
1751
- The memory in one GPU can be classified into 3 categories:
1752
- 1. memory used by anything other than the current vLLM instance.
1753
- 2. memory used by torch in the current vLLM instance.
1754
- 3. memory used in the current vLLM instance, but not by torch.
1755
-
1756
- A quantitive example:
1757
-
1758
- Before creating the current vLLM instance:
1759
- category 1: 1 GiB
1760
- category 2: 0 GiB
1761
- category 3: 0 GiB
1762
-
1763
- After creating the current vLLM instance and loading the model,
1764
- (i.e. before profiling):
1765
- category 1: 1 GiB
1766
- category 2: 2 GiB (model weights take 2 GiB)
1767
- category 3: 0.5 GiB (memory used by NCCL)
1768
-
1769
- During profiling (peak):
1770
- category 1: 1 GiB
1771
- category 2: 4 GiB (peak activation tensors take 2 GiB)
1772
- category 3: 1 GiB (memory used by NCCL + buffers for some attention backends)
1773
-
1774
- After profiling:
1775
- category 1: 1 GiB
1776
- category 2: 3 GiB (after garbage-collecting activation tensors)
1777
- category 3: 1 GiB (memory used by NCCL + buffers for some attention backends)
1778
-
1779
- In this case, non-kv cache takes 5 GiB in total, including:
1780
- a. 2 GiB used by the model weights (category 2)
1781
- b. 2 GiB reserved for the peak activation tensors (category 2)
1782
- c. 1 GiB used by non-torch components (category 3)
1783
-
1784
- The memory used for loading weights (a.) is directly given from the argument `weights_memory`.
1785
-
1786
- The increase of `torch.cuda.memory_stats()["allocated_bytes.all.peak"]` during profiling gives (b.).
1787
-
1788
- The increase of `non_torch_memory` from creating the current vLLM instance until after profiling to get (c.).
1789
- """ # noqa
1790
- gc.collect()
1791
- torch.cuda.empty_cache()
1792
- torch.cuda.reset_peak_memory_stats()
1793
-
1794
- result = MemoryProfilingResult()
1795
-
1796
- result.before_create = baseline_snapshot
1797
- # the part of memory used for holding the model weights
1798
- result.weights_memory = weights_memory
1799
-
1800
- result.before_profile.measure()
1801
-
1802
- yield result
1803
-
1804
- gc.collect()
1805
- torch.cuda.empty_cache()
1806
-
1807
- result.after_profile.measure()
1808
-
1809
- diff_profile = result.after_profile - result.before_profile
1810
- diff_from_create = result.after_profile - result.before_create
1811
- result.torch_peak_increase = diff_profile.torch_peak
1812
- result.non_torch_increase = diff_from_create.non_torch_memory
1813
- result.profile_time = diff_profile.timestamp
1814
-
1815
- non_torch_memory = result.non_torch_increase
1816
- peak_activation_memory = result.torch_peak_increase
1817
- result.non_kv_cache_memory = (
1818
- non_torch_memory + peak_activation_memory + result.weights_memory
1819
- ) # noqa
1820
-
1821
-
1588
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 # noqa: E501
1589
def set_ulimit(target_soft_limit=65535):
1590
if sys.platform.startswith("win"):
vllm/utils/mem_constants.py
new
+13
@@ -0,0 +1,13 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+MB_bytes = 1_000_000
4
+"""The number of bytes in one megabyte (MB)."""
5
+
6
+MiB_bytes = 1 << 20
7
+"""The number of bytes in one mebibyte (MiB)."""
8
+
9
+GB_bytes = 1_000_000_000
10
+"""The number of bytes in one gigabyte (GB)."""
11
+
12
+GiB_bytes = 1 << 30
13
+"""The number of bytes in one gibibyte (GiB)."""
vllm/utils/mem_utils.py
new
+232
@@ -0,0 +1,232 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import contextlib
4
+import gc
5
+import time
6
+from collections.abc import Generator
7
+from dataclasses import dataclass, field
8
+from functools import cache
9
+
10
+import psutil
11
+import torch
12
+import torch.types
13
+
14
+from .mem_constants import GiB_bytes
15
+
16
+
17
+@cache
18
+def get_max_shared_memory_bytes(gpu: int = 0) -> int:
19
+ """Returns the maximum shared memory per thread block in bytes."""
20
+ from vllm import _custom_ops as ops
21
+
22
+ max_shared_mem = ops.get_max_shared_memory_per_block_device_attribute(gpu)
23
+ # value 0 will cause MAX_SEQ_LEN become negative and test_attention.py
24
+ # will fail
25
+ assert max_shared_mem > 0, "max_shared_mem can not be zero"
26
+ return int(max_shared_mem)
27
+
28
+
29
+def get_cpu_memory() -> int:
30
+ """Returns the total CPU memory of the node in bytes."""
31
+ return psutil.virtual_memory().total
32
+
33
+
34
+class DeviceMemoryProfiler:
35
+ def __init__(self, device: torch.types.Device | None = None):
36
+ self.device = device
37
+
38
+ def current_memory_usage(self) -> float:
39
+ # Return the memory usage in bytes.
40
+ from vllm.platforms import current_platform
41
+
42
+ gc.collect()
43
+ return current_platform.get_current_memory_usage(self.device)
44
+
45
+ def __enter__(self):
46
+ self.initial_memory = self.current_memory_usage()
47
+ # This allows us to call methods of the context manager if needed
48
+ return self
49
+
50
+ def __exit__(self, exc_type, exc_val, exc_tb):
51
+ self.final_memory = self.current_memory_usage()
52
+ self.consumed_memory = self.final_memory - self.initial_memory
53
+
54
+ # Force garbage collection
55
+ gc.collect()
56
+
57
+
58
+@dataclass
59
+class MemorySnapshot:
60
+ """Memory snapshot."""
61
+
62
+ torch_peak: int = 0
63
+ free_memory: int = 0
64
+ total_memory: int = 0
65
+ cuda_memory: int = 0
66
+ torch_memory: int = 0
67
+ non_torch_memory: int = 0
68
+ timestamp: float = 0.0
69
+ auto_measure: bool = True
70
+
71
+ def __post_init__(self):
72
+ if self.auto_measure:
73
+ self.measure()
74
+
75
+ def measure(self):
76
+ from vllm.platforms import current_platform
77
+
78
+ # we measure the torch peak memory usage via allocated_bytes,
79
+ # rather than `torch.cuda.memory_reserved()` .
80
+ # After `torch.cuda.reset_peak_memory_stats()`,
81
+ # `torch.cuda.memory_reserved()` will keep growing, and only shrink
82
+ # when we call `torch.cuda.empty_cache()` or OOM happens.
83
+ self.torch_peak = torch.cuda.memory_stats().get("allocated_bytes.all.peak", 0)
84
+
85
+ self.free_memory, self.total_memory = torch.cuda.mem_get_info()
86
+ shared_sysmem_device_mem_sms = ((8, 7), (11, 0), (12, 1)) # Orin, Thor, Spark
87
+ if (
88
+ current_platform.is_cuda()
89
+ and current_platform.get_device_capability() in shared_sysmem_device_mem_sms
90
+ ):
91
+ # On UMA (Orin, Thor and Spark) platform,
92
+ # where both CPU and GPU rely on system memory,
93
+ # the cudaMemGetInfo function shows the amount of free system memory
94
+ # rather than what’s actually available.
95
+ # In the case,
96
+ # torch.cuda.mem_get_info() only reports "free" memory,
97
+ # which can be lower than what is actually
98
+ # available due to not including cache memory.
99
+ # There’s also a comprehensive reference page
100
+ # that explains how you can compute the proper value yourself.
101
+ # https://docs.nvidia.com/cuda/cuda-for-tegra-appnote/#estimating-total-allocatable-device-memory-on-an-integrated-gpu-device
102
+ self.free_memory = psutil.virtual_memory().available
103
+
104
+ self.cuda_memory = self.total_memory - self.free_memory
105
+
106
+ # torch.cuda.memory_reserved() is how many bytes
107
+ # PyTorch gets from cuda (by calling cudaMalloc, etc.)
108
+ # this is used to measure the non-torch memory usage
109
+ self.torch_memory = torch.cuda.memory_reserved()
110
+
111
+ self.non_torch_memory = self.cuda_memory - self.torch_memory
112
+ self.timestamp = time.time()
113
+
114
+ def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot":
115
+ return MemorySnapshot(
116
+ torch_peak=self.torch_peak - other.torch_peak,
117
+ free_memory=self.free_memory - other.free_memory,
118
+ total_memory=self.total_memory - other.total_memory,
119
+ cuda_memory=self.cuda_memory - other.cuda_memory,
120
+ torch_memory=self.torch_memory - other.torch_memory,
121
+ non_torch_memory=self.non_torch_memory - other.non_torch_memory,
122
+ timestamp=self.timestamp - other.timestamp,
123
+ auto_measure=False,
124
+ )
125
+
126
+
127
+@dataclass
128
+class MemoryProfilingResult:
129
+ """Memory profiling result. All numbers are in bytes."""
130
+
131
+ non_kv_cache_memory: int = 0
132
+ torch_peak_increase: int = 0
133
+ non_torch_increase: int = 0
134
+ weights_memory: float = 0
135
+ before_create: MemorySnapshot = field(default_factory=MemorySnapshot)
136
+ before_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
137
+ after_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
138
+ profile_time: float = 0.0
139
+
140
+ def __repr__(self) -> str:
141
+ return (
142
+ f"Memory profiling takes {self.profile_time:.2f} seconds. "
143
+ f"Total non KV cache memory: "
144
+ f"{(self.non_kv_cache_memory / GiB_bytes):.2f}GiB; "
145
+ f"torch peak memory increase: "
146
+ f"{(self.torch_peak_increase / GiB_bytes):.2f}GiB; "
147
+ f"non-torch forward increase memory: "
148
+ f"{(self.non_torch_increase / GiB_bytes):.2f}GiB; "
149
+ f"weights memory: {(self.weights_memory / GiB_bytes):.2f}GiB."
150
+ )
151
+
152
+
153
+@contextlib.contextmanager
154
+def memory_profiling(
155
+ baseline_snapshot: MemorySnapshot, weights_memory: int
156
+) -> Generator[MemoryProfilingResult, None, None]:
157
+ """Memory profiling context manager.
158
+ baseline_snapshot: the memory snapshot before the current vLLM instance.
159
+ weights_memory: memory used by PyTorch when loading the model weights.
160
+ Note that, before loading the model weights, we also initialize the device
161
+ and distributed environment, which may consume some memory. This part is not
162
+ included in the weights_memory because PyTorch does not control it.
163
+
164
+ The memory in one GPU can be classified into 3 categories:
165
+ 1. memory used by anything other than the current vLLM instance.
166
+ 2. memory used by torch in the current vLLM instance.
167
+ 3. memory used in the current vLLM instance, but not by torch.
168
+
169
+ A quantitive example:
170
+
171
+ Before creating the current vLLM instance:
172
+ category 1: 1 GiB
173
+ category 2: 0 GiB
174
+ category 3: 0 GiB
175
+
176
+ After creating the current vLLM instance and loading the model,
177
+ (i.e. before profiling):
178
+ category 1: 1 GiB
179
+ category 2: 2 GiB (model weights take 2 GiB)
180
+ category 3: 0.5 GiB (memory used by NCCL)
181
+
182
+ During profiling (peak):
183
+ category 1: 1 GiB
184
+ category 2: 4 GiB (peak activation tensors take 2 GiB)
185
+ category 3: 1 GiB (memory used by NCCL + buffers for some attention backends)
186
+
187
+ After profiling:
188
+ category 1: 1 GiB
189
+ category 2: 3 GiB (after garbage-collecting activation tensors)
190
+ category 3: 1 GiB (memory used by NCCL + buffers for some attention backends)
191
+
192
+ In this case, non-kv cache takes 5 GiB in total, including:
193
+ a. 2 GiB used by the model weights (category 2)
194
+ b. 2 GiB reserved for the peak activation tensors (category 2)
195
+ c. 1 GiB used by non-torch components (category 3)
196
+
197
+ The memory used for loading weights (a.) is directly given from the argument `weights_memory`.
198
+
199
+ The increase of `torch.cuda.memory_stats()["allocated_bytes.all.peak"]` during profiling gives (b.).
200
+
201
+ The increase of `non_torch_memory` from creating the current vLLM instance until after profiling to get (c.).
202
+ """ # noqa
203
+ gc.collect()
204
+ torch.cuda.empty_cache()
205
+ torch.cuda.reset_peak_memory_stats()
206
+
207
+ result = MemoryProfilingResult()
208
+
209
+ result.before_create = baseline_snapshot
210
+ # the part of memory used for holding the model weights
211
+ result.weights_memory = weights_memory
212
+
213
+ result.before_profile.measure()
214
+
215
+ yield result
216
+
217
+ gc.collect()
218
+ torch.cuda.empty_cache()
219
+
220
+ result.after_profile.measure()
221
+
222
+ diff_profile = result.after_profile - result.before_profile
223
+ diff_from_create = result.after_profile - result.before_create
224
+ result.torch_peak_increase = diff_profile.torch_peak
225
+ result.non_torch_increase = diff_from_create.non_torch_memory
226
+ result.profile_time = diff_profile.timestamp
227
+
228
+ non_torch_memory = result.non_torch_increase
229
+ peak_activation_memory = result.torch_peak_increase
230
+ result.non_kv_cache_memory = (
231
+ non_torch_memory + peak_activation_memory + result.weights_memory
232
+ ) # noqa
vllm/v1/core/kv_cache_utils.py
+2
-1
@@ -12,7 +12,8 @@ from typing import Any, NewType, TypeAlias
12
from vllm import envs
13
from vllm.config import VllmConfig
14
from vllm.logger import init_logger
15
-from vllm.utils import GiB_bytes, cdiv, sha256_cbor
15
+from vllm.utils import cdiv, sha256_cbor
16
+from vllm.utils.mem_constants import GiB_bytes
17
from vllm.v1.kv_cache_interface import (
18
ChunkedLocalAttentionSpec,
19
FullAttentionSpec,
vllm/v1/worker/gpu_model_runner.py
+2
-2
@@ -74,8 +74,6 @@ from vllm.sequence import IntermediateTensors
74
from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
75
from vllm.utils import (
76
STR_DTYPE_TO_TORCH_DTYPE,
77
- DeviceMemoryProfiler,
78
- GiB_bytes,
77
cdiv,
78
check_use_alibi,
79
get_dtype_size,
@@ -85,6 +83,8 @@ from vllm.utils import (
83
supports_dynamo,
84
)
85
from vllm.utils.jsontree import json_map_leaves
86
+from vllm.utils.mem_constants import GiB_bytes
87
+from vllm.utils.mem_utils import DeviceMemoryProfiler
88
from vllm.v1.attention.backends.flash_attn import AttentionMetadata
89
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
90
from vllm.v1.attention.backends.utils import (
vllm/v1/worker/gpu_worker.py
+2
-1
@@ -28,7 +28,8 @@ from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
28
from vllm.platforms import current_platform
29
from vllm.sequence import IntermediateTensors
30
from vllm.tasks import SupportedTask
31
-from vllm.utils import GiB_bytes, MemorySnapshot, memory_profiling
31
+from vllm.utils.mem_constants import GiB_bytes
32
+from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
33
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
34
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
35
from vllm.v1.outputs import (