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 6, 00:45 UTC
[CPU][UX][Perf] Enable tcmalloc by default (#37607)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Mar 25, 2026 at 12:39 UTC
34d317dcec3935e588d6c8ee8a7a57abb7a3e731
2 files changed
+94
-3
setup.py
+65
@@ -82,6 +82,66 @@ def is_freethreaded():
82
return bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
83
84
85
+def should_bundle_tcmalloc() -> bool:
86
+ import platform
87
+
88
+ return (
89
+ VLLM_TARGET_DEVICE == "cpu"
90
+ and sys.platform.startswith("linux")
91
+ and platform.machine() in ("aarch64", "x86_64")
92
+ )
93
+
94
+
95
+def find_tcmalloc() -> Path | None:
96
+ try:
97
+ # get all shared libs the dynamic loader knows about
98
+ output = subprocess.check_output(
99
+ ["ldconfig", "-p"],
100
+ text=True,
101
+ stderr=subprocess.DEVNULL,
102
+ )
103
+ except Exception:
104
+ return None
105
+
106
+ # search for libtcmalloc and libtcmalloc_minimal
107
+ for library_pattern in (
108
+ r"\blibtcmalloc_minimal\.so\.(\d+)\b",
109
+ r"\blibtcmalloc\.so\.(\d+)\b",
110
+ ):
111
+ candidates: list[tuple[int, Path]] = []
112
+ for line in output.splitlines():
113
+ match = re.search(library_pattern, line)
114
+ if match is None or "=>" not in line:
115
+ continue
116
+ candidate = Path(line.split("=>")[1].strip())
117
+ if candidate.exists():
118
+ candidates.append((int(match.group(1)), candidate))
119
+
120
+ if candidates:
121
+ # if multiple candidates are found, pick the one with the highest
122
+ # version number
123
+ return max(candidates, key=lambda item: item[0])[1]
124
+
125
+ return None
126
+
127
+
128
+def bundle_tcmalloc(build_lib: str) -> None:
129
+ tcmalloc_library = find_tcmalloc()
130
+ if tcmalloc_library is None:
131
+ logger.warning(
132
+ "Failed to locate tcmalloc. For best performance, "
133
+ "please install tcmalloc (e.g. `sudo apt-get "
134
+ "install -y --no-install-recommends libtcmalloc-minimal4`)"
135
+ )
136
+ return
137
+
138
+ bundle_dir = os.path.join(build_lib, "vllm", "libs")
139
+ os.makedirs(bundle_dir, exist_ok=True)
140
+ bundle_path = os.path.join(bundle_dir, tcmalloc_library.name)
141
+ shutil.copy2(tcmalloc_library, bundle_path)
142
+ logger.info("Bundled tcmalloc into wheel: %s", bundle_path)
143
+
144
+
145
class CMakeExtension(Extension):
146
def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
147
super().__init__(name, sources=[], py_limited_api=not is_freethreaded(), **kwa)
@@ -285,6 +345,10 @@ class cmake_build_ext(build_ext):
345
# First, run the standard build_ext command to compile the extensions
346
super().run()
347
348
+ # bundle tcmalloc into CPU wheels for best OOB perf
349
+ if should_bundle_tcmalloc():
350
+ bundle_tcmalloc(self.build_lib)
351
+
352
# copy vllm/vllm_flash_attn/**/*.py from self.build_lib to current
353
# directory so that they can be included in the editable build
354
import glob
@@ -944,6 +1008,7 @@ if _build_custom_ops():
1008
package_data = {
1009
"vllm": [
1010
"py.typed",
1011
+ "libs/*.so*",
1012
"model_executor/layers/fused_moe/configs/*.json",
1013
"model_executor/layers/quantization/utils/configs/*.json",
1014
"entrypoints/serve/instrumentator/static/*.js",
vllm/platforms/cpu.py
+29
-3
@@ -284,8 +284,9 @@ class CpuPlatform(Platform):
284
# Avoid inductor generates num_thread() and breaks the thread binding
285
os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1"
286
287
- # Intel OpenMP setting
287
ld_preload_str = os.getenv("LD_PRELOAD", "")
288
+
289
+ # Intel OpenMP setting
290
if "libiomp5.so" in ld_preload_str:
291
# The time(milliseconds) that a thread should wait after
292
# completing the execution of a parallel region, before sleeping.
@@ -297,10 +298,35 @@ class CpuPlatform(Platform):
298
os.environ["KMP_PLAIN_BARRIER_PATTERN"] = "dist,dist"
299
os.environ["KMP_REDUCTION_BARRIER_PATTERN"] = "dist,dist"
300
301
+ cpu_architecture = Platform.get_cpu_architecture()
302
+
303
+ # LD_PRELOAD libtcmalloc, bundled under vllm/libs to reduce
304
+ # memory allocation overhead
305
+ if (
306
+ platform.system() == "Linux"
307
+ and cpu_architecture in (CpuArchEnum.ARM, CpuArchEnum.X86)
308
+ and "libtcmalloc" not in ld_preload_str
309
+ ):
310
+ vllm_pkg = os.path.dirname(os.path.dirname(__file__))
311
+ tcmalloc_so = None
312
+ for pattern in ("libtcmalloc_minimal*.so*", "libtcmalloc.so*"):
313
+ tcmalloc_so_candidates = glob.glob(
314
+ os.path.join(vllm_pkg, "libs", pattern)
315
+ )
316
+ if tcmalloc_so_candidates:
317
+ tcmalloc_so = tcmalloc_so_candidates[0]
318
+ break
319
+
320
+ if tcmalloc_so is not None:
321
+ if ld_preload_str:
322
+ ld_preload_str = f"{tcmalloc_so}:{ld_preload_str}"
323
+ else:
324
+ ld_preload_str = tcmalloc_so
325
+ os.environ["LD_PRELOAD"] = ld_preload_str
326
+
327
if (
328
platform.system() == "Linux"
302
- and Platform.get_cpu_architecture()
303
- in (CpuArchEnum.ARM, CpuArchEnum.POWERPC)
329
+ and cpu_architecture in (CpuArchEnum.ARM, CpuArchEnum.POWERPC)
330
and not ("libomp" in ld_preload_str or "libgomp" in ld_preload_str)
331
):
332
# We need to LD_PRELOAD PyTorch's libgomp, otherwise only