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 10, 14:00 UTC
[refactor] refactor memory constants usage (#31865)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Ning Xie committed
Jan 8, 2026 at 02:37 UTC
c907d22158dbf31508e73e89cb2bb272bb59e37e
9 files changed
+33
-30
tests/basic_correctness/test_cumem.py
-1
@@ -247,7 +247,6 @@ def test_deep_sleep_async():
247
248
@requires_fp8
249
def test_deep_sleep_fp8_kvcache():
250
- GiB_bytes = 1 << 30
250
model = "Qwen/Qwen2-0.5B"
251
used_bytes_baseline = current_platform.get_current_memory_usage()
252
vllm/config/cache.py
+5
-4
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import math
5
from dataclasses import field
6
from typing import TYPE_CHECKING, Any, Literal
7
@@ -10,7 +11,7 @@ from pydantic.dataclasses import dataclass
11
from vllm.config.utils import config
12
from vllm.logger import init_logger
13
from vllm.utils.mem_constants import GiB_bytes
13
-from vllm.utils.mem_utils import get_cpu_memory
14
+from vllm.utils.mem_utils import format_gib, get_cpu_memory
15
16
if TYPE_CHECKING:
17
from vllm.config.parallel import ParallelConfig
@@ -214,7 +215,7 @@ class CacheConfig:
215
self,
216
parallel_config: ParallelConfig,
217
) -> None:
217
- swap_space_bytes = self.swap_space * GiB_bytes
218
+ swap_space_bytes = math.ceil(self.swap_space * GiB_bytes)
219
total_cpu_memory = get_cpu_memory()
220
# FIXME(woosuk): Here, it is assumed that the GPUs in a tensor parallel
221
# group are in the same node. However, the GPUs may span multiple nodes.
@@ -222,8 +223,8 @@ class CacheConfig:
223
cpu_memory_usage = swap_space_bytes * num_gpus_per_node
224
225
msg = (
225
- f"{cpu_memory_usage / GiB_bytes:.2f} GiB out of the "
226
- f"{total_cpu_memory / GiB_bytes:.2f} GiB total CPU memory "
226
+ f"{format_gib(cpu_memory_usage)} GiB out of the "
227
+ f"{format_gib(total_cpu_memory)} GiB total CPU memory "
228
"is allocated for the swap space."
229
)
230
if cpu_memory_usage > 0.7 * total_cpu_memory:
vllm/multimodal/cache.py
+3
-2
@@ -20,6 +20,7 @@ from vllm.logger import init_logger
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
+from vllm.utils.mem_utils import format_gib
24
25
from .inputs import (
26
MultiModalBatchedField,
@@ -130,9 +131,9 @@ class MultiModalCache:
131
if debug:
132
leaf_count = json_count_leaves(value)
133
logger.debug(
133
- "Calculated size of %s to be %.2f GiB (%d leaves)",
134
+ "Calculated size of %s to be %s GiB (%d leaves)",
135
type(value),
135
- size / GiB_bytes,
136
+ format_gib(size),
137
leaf_count,
138
)
139
vllm/platforms/cpu.py
+3
-3
@@ -140,6 +140,7 @@ class CpuPlatform(Platform):
140
@classmethod
141
def get_device_total_memory(cls, device_id: int = 0) -> int:
142
from vllm.utils.mem_constants import GiB_bytes
143
+ from vllm.utils.mem_utils import format_gib
144
145
kv_cache_space = envs.VLLM_CPU_KVCACHE_SPACE
146
node_dir = "/sys/devices/system/node"
@@ -153,10 +154,9 @@ class CpuPlatform(Platform):
154
free_cpu_memory = psutil.virtual_memory().total // num_numa_nodes
155
DEFAULT_CPU_MEM_UTILIZATION = 0.5
156
kv_cache_space = int(free_cpu_memory * DEFAULT_CPU_MEM_UTILIZATION)
156
- kv_cache_space_gib = kv_cache_space / GiB_bytes
157
logger.warning_once(
158
- "VLLM_CPU_KVCACHE_SPACE not set. Using "
159
- f"{kv_cache_space_gib:.2f} GiB for KV cache."
158
+ "VLLM_CPU_KVCACHE_SPACE not set. Using %s GiB for KV cache.",
159
+ format_gib(kv_cache_space),
160
)
161
else:
162
kv_cache_space *= GiB_bytes
vllm/utils/mem_utils.py
+7
-3
@@ -11,11 +11,15 @@ import psutil
11
import torch
12
import torch.types
13
14
-from .mem_constants import GiB_bytes
14
+from .mem_constants import GiB_bytes, MiB_bytes
15
16
17
-def format_gib(b: int) -> float:
18
- return round(b / GiB_bytes, 2)
17
+def format_mib(b: int) -> str:
18
+ return f"{round(b / MiB_bytes, 2)}"
19
+
20
+
21
+def format_gib(b: int) -> str:
22
+ return f"{round(b / GiB_bytes, 2)}"
23
24
25
@cache
vllm/v1/core/kv_cache_utils.py
+5
-5
@@ -14,7 +14,7 @@ from vllm.config import VllmConfig
14
from vllm.logger import init_logger
15
from vllm.utils.hashing import sha256_cbor, xxhash_cbor
16
from vllm.utils.math_utils import cdiv
17
-from vllm.utils.mem_constants import GiB_bytes
17
+from vllm.utils.mem_utils import format_gib
18
from vllm.v1.kv_cache_interface import (
19
ChunkedLocalAttentionSpec,
20
FullAttentionSpec,
@@ -633,9 +633,9 @@ def _check_enough_kv_cache_memory(
633
634
raise ValueError(
635
f"To serve at least one request with the models's max seq len "
636
- f"({max_model_len}), ({needed_memory / GiB_bytes:.2f} GiB KV "
636
+ f"({max_model_len}), ({format_gib(needed_memory)} GiB KV "
637
f"cache is needed, which is larger than the available KV cache "
638
- f"memory ({available_memory / GiB_bytes:.2f} GiB). {estimated_msg}"
638
+ f"memory ({format_gib(available_memory)} GiB). {estimated_msg}"
639
f"Try increasing `gpu_memory_utilization` or decreasing `max_model_len` "
640
f"when initializing the engine. "
641
f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ "
@@ -1441,10 +1441,10 @@ def _auto_fit_max_model_len(
1441
vllm_config.model_config.max_model_len = auto_fit_max
1442
logger.info_once(
1443
"Auto-fit max_model_len: reduced from %d to %d to fit in "
1444
- "available GPU memory (%.2f GiB available for KV cache)",
1444
+ "available GPU memory (%s GiB available for KV cache)",
1445
original_max,
1446
auto_fit_max,
1447
- min_available_memory / GiB_bytes,
1447
+ format_gib(min_available_memory),
1448
scope="local",
1449
)
1450
vllm/v1/worker/gpu/model_runner.py
+3
-4
@@ -14,8 +14,7 @@ from vllm.config.compilation import CUDAGraphMode
14
from vllm.forward_context import set_forward_context
15
from vllm.logger import init_logger
16
from vllm.model_executor.model_loader import get_model_loader
17
-from vllm.utils.mem_constants import GiB_bytes
18
-from vllm.utils.mem_utils import DeviceMemoryProfiler
17
+from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
18
from vllm.utils.platform_utils import is_pin_memory_available
19
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
20
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -165,8 +164,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
164
165
self.model_memory_usage = m.consumed_memory
166
logger.info(
168
- "Model loading took %.4f GiB and %.6f seconds",
169
- m.consumed_memory / GiB_bytes,
167
+ "Model loading took %s GiB and %.6f seconds",
168
+ format_gib(m.consumed_memory),
169
time_after_load - time_before_load,
170
)
171
vllm/v1/worker/gpu_model_runner.py
+3
-4
@@ -93,8 +93,7 @@ from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
93
from vllm.utils import length_from_prompt_token_ids_or_embeds
94
from vllm.utils.jsontree import json_map_leaves
95
from vllm.utils.math_utils import cdiv, round_up
96
-from vllm.utils.mem_constants import GiB_bytes
97
-from vllm.utils.mem_utils import DeviceMemoryProfiler
96
+from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
97
from vllm.utils.nvtx_pytorch_hooks import PytHooks
98
from vllm.utils.platform_utils import is_pin_memory_available
99
from vllm.utils.torch_utils import (
@@ -3899,8 +3898,8 @@ class GPUModelRunner(
3898
logger.error(combined_msg)
3899
raise e
3900
logger.info_once(
3902
- "Model loading took %.4f GiB memory and %.6f seconds",
3903
- self.model_memory_usage / GiB_bytes,
3901
+ "Model loading took %s GiB memory and %.6f seconds",
3902
+ format_gib(self.model_memory_usage),
3903
time_after_load - time_before_load,
3904
scope="local",
3905
)
vllm/v1/worker/gpu_worker.py
+4
-4
@@ -125,7 +125,7 @@ class Worker(WorkerBase):
125
used_bytes = total - free_bytes_after_sleep
126
assert freed_bytes >= 0, "Memory usage increased after sleeping."
127
logger.info(
128
- "Sleep mode freed %f GiB memory, %f GiB memory is still in use.",
128
+ "Sleep mode freed %s GiB memory, %s GiB memory is still in use.",
129
format_gib(freed_bytes),
130
format_gib(used_bytes),
131
)
@@ -342,19 +342,19 @@ class Worker(WorkerBase):
342
343
unrequested_memory = self.init_snapshot.free_memory - self.requested_memory
344
logger.debug(
345
- "Initial free memory: %f GiB; Requested memory: %f (util), %f GiB",
345
+ "Initial free memory: %s GiB; Requested memory: %f (util), %s GiB",
346
format_gib(self.init_snapshot.free_memory),
347
self.cache_config.gpu_memory_utilization,
348
format_gib(self.requested_memory),
349
)
350
logger.debug(
351
- "Free memory after profiling: %f GiB (total), %f GiB (within requested)",
351
+ "Free memory after profiling: %s GiB (total), %s GiB (within requested)",
352
format_gib(free_gpu_memory),
353
format_gib(free_gpu_memory - unrequested_memory),
354
)
355
logger.debug(profile_result)
356
logger.info_once(
357
- "Available KV cache memory: %f GiB",
357
+ "Available KV cache memory: %s GiB",
358
format_gib(self.available_kv_cache_memory_bytes),
359
scope="local",
360
)