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
[NIXL][non-cuda] Add install script for nixl with non-cuda ucx (#25959)
Signed-off-by: Chendi Xue <Chendi.Xue@intel.com>
Chendi.Xue committed
Oct 8, 2025 at 09:19 UTC
9fc983c707a9371951135d2e5a62183da7e09369
3 files changed
+217
-1
docs/features/nixl_connector_usage.md
+6
@@ -11,6 +11,12 @@ Install the NIXL library: `uv pip install nixl`, as a quick start.
11
- Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
12
- The specified required NIXL version can be found in [requirements/kv_connectors.txt](gh-file:requirements/kv_connectors.txt) and other relevant config files
13
14
+For non-cuda platform, please install nixl with ucx build from source, instructed as below.
15
+
16
+```bash
17
+python tools/install_nixl_from_source_ubuntu.py
18
+```
19
+
20
### Transport Configuration
21
22
NixlConnector uses NIXL library for underlying communication, which supports multiple transport backends. UCX (Unified Communication X) is the primary default transport library used by NIXL. Configure transport environment variables:
docs/serving/expert_parallel_deployment.md
+1
-1
@@ -191,7 +191,7 @@ For production deployments requiring strict SLA guarantees for time-to-first-tok
191
192
### Setup Steps
193
194
-1. **Install gdrcopy/ucx/nixl**: For maximum performance, run the [install_gdrcopy.sh](gh-file:tools/install_gdrcopy.sh) script to install `gdrcopy` (e.g., `install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "12.8" "x64"`). You can find available OS versions [here](https://developer.download.nvidia.com/compute/redist/gdrcopy/CUDA%2012.8/). If `gdrcopy` is not installed, things will still work with a plain `pip install nixl`, just with lower performance. `nixl` and `ucx` are installed as dependencies via pip.
194
+1. **Install gdrcopy/ucx/nixl**: For maximum performance, run the [install_gdrcopy.sh](gh-file:tools/install_gdrcopy.sh) script to install `gdrcopy` (e.g., `install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "12.8" "x64"`). You can find available OS versions [here](https://developer.download.nvidia.com/compute/redist/gdrcopy/CUDA%2012.8/). If `gdrcopy` is not installed, things will still work with a plain `pip install nixl`, just with lower performance. `nixl` and `ucx` are installed as dependencies via pip. For non-cuda platform to install nixl with non-cuda UCX build, run the [install_nixl_from_source_ubuntu.py](gh-file:tools/install_nixl_from_source_ubuntu.py) script.
195
196
2. **Configure Both Instances**: Add this flag to both prefill and decode instances `--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}`. Noted, you may also specify one or multiple NIXL_Backend. Such as: `--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'`
197
tools/install_nixl_from_source_ubuntu.py
new
+210
@@ -0,0 +1,210 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# install_prerequisites.py
4
+import argparse
5
+import glob
6
+import os
7
+import subprocess
8
+import sys
9
+
10
+# --- Configuration ---
11
+WHEELS_CACHE_HOME = os.environ.get("WHEELS_CACHE_HOME", "/tmp/wheels_cache")
12
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
13
+UCX_DIR = os.path.join('/tmp', 'ucx_source')
14
+NIXL_DIR = os.path.join('/tmp', 'nixl_source')
15
+UCX_INSTALL_DIR = os.path.join('/tmp', 'ucx_install')
16
+UCX_REPO_URL = 'https://github.com/openucx/ucx.git'
17
+NIXL_REPO_URL = 'https://github.com/ai-dynamo/nixl.git'
18
+
19
+
20
+# --- Helper Functions ---
21
+def run_command(command, cwd='.', env=None):
22
+ """Helper function to run a shell command and check for errors."""
23
+ print(f"--> Running command: {' '.join(command)} in '{cwd}'", flush=True)
24
+ subprocess.check_call(command, cwd=cwd, env=env)
25
+
26
+
27
+def is_pip_package_installed(package_name):
28
+ """Checks if a package is installed via pip without raising an exception."""
29
+ result = subprocess.run(
30
+ [sys.executable, '-m', 'pip', 'show', package_name],
31
+ stdout=subprocess.DEVNULL,
32
+ stderr=subprocess.DEVNULL)
33
+ return result.returncode == 0
34
+
35
+
36
+def find_nixl_wheel_in_cache(cache_dir):
37
+ """Finds a nixl wheel file in the specified cache directory."""
38
+ # The repaired wheel will have a 'manylinux' tag, but this glob still works.
39
+ search_pattern = os.path.join(cache_dir, "nixl-*.whl")
40
+ wheels = glob.glob(search_pattern)
41
+ if wheels:
42
+ # Sort to get the most recent/highest version if multiple exist
43
+ wheels.sort()
44
+ return wheels[-1]
45
+ return None
46
+
47
+
48
+def install_system_dependencies():
49
+ """Installs required system packages using apt-get if run as root."""
50
+ if os.geteuid() != 0:
51
+ print("\n---", flush=True)
52
+ print("WARNING: Not running as root. \
53
+ Skipping system dependency installation.",
54
+ flush=True)
55
+ print(
56
+ "Please ensure the listed packages are installed on your system:",
57
+ flush=True)
58
+ print(" patchelf build-essential git cmake ninja-build \
59
+ autotools-dev automake meson libtool libtool-bin",
60
+ flush=True)
61
+ print("---\n", flush=True)
62
+ return
63
+
64
+ print("--- Running as root. Installing system dependencies... ---",
65
+ flush=True)
66
+ apt_packages = [
67
+ "patchelf", # <-- Add patchelf here
68
+ "build-essential",
69
+ "git",
70
+ "cmake",
71
+ "ninja-build",
72
+ "autotools-dev",
73
+ "automake",
74
+ "meson",
75
+ "libtool",
76
+ "libtool-bin"
77
+ ]
78
+ run_command(['apt-get', 'update'])
79
+ run_command(['apt-get', 'install', '-y'] + apt_packages)
80
+ print("--- System dependencies installed successfully. ---\n", flush=True)
81
+
82
+
83
+def build_and_install_prerequisites(args):
84
+ """Builds UCX and NIXL from source, creating a self-contained wheel."""
85
+
86
+ if not args.force_reinstall and is_pip_package_installed('nixl'):
87
+ print("--> NIXL is already installed. Nothing to do.", flush=True)
88
+ return
89
+
90
+ cached_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME)
91
+ if not args.force_reinstall and cached_wheel:
92
+ print(f"\n--> Found self-contained wheel: \
93
+ {os.path.basename(cached_wheel)}.",
94
+ flush=True)
95
+ print("--> Installing from cache, skipping all source builds.",
96
+ flush=True)
97
+ install_command = [
98
+ sys.executable, '-m', 'pip', 'install', cached_wheel
99
+ ]
100
+ run_command(install_command)
101
+ print("\n--- Installation from cache complete. ---", flush=True)
102
+ return
103
+
104
+ print("\n--> No installed package or cached wheel found. \
105
+ Starting full build process...",
106
+ flush=True)
107
+ print("\n--> Installing auditwheel...", flush=True)
108
+ run_command([sys.executable, '-m', 'pip', 'install', 'auditwheel'])
109
+ install_system_dependencies()
110
+ ucx_install_path = os.path.abspath(UCX_INSTALL_DIR)
111
+ print(f"--> Using wheel cache directory: {WHEELS_CACHE_HOME}", flush=True)
112
+ os.makedirs(WHEELS_CACHE_HOME, exist_ok=True)
113
+
114
+ # -- Step 1: Build UCX from source --
115
+ print("\n[1/3] Configuring and building UCX from source...", flush=True)
116
+ if not os.path.exists(UCX_DIR):
117
+ run_command(['git', 'clone', UCX_REPO_URL, UCX_DIR])
118
+ ucx_source_path = os.path.abspath(UCX_DIR)
119
+ run_command(['git', 'checkout', 'v1.19.x'], cwd=ucx_source_path)
120
+ run_command(['./autogen.sh'], cwd=ucx_source_path)
121
+ configure_command = [
122
+ './configure',
123
+ f'--prefix={ucx_install_path}',
124
+ '--enable-shared',
125
+ '--disable-static',
126
+ '--disable-doxygen-doc',
127
+ '--enable-optimizations',
128
+ '--enable-cma',
129
+ '--enable-devel-headers',
130
+ '--with-verbs',
131
+ '--enable-mt',
132
+ ]
133
+ run_command(configure_command, cwd=ucx_source_path)
134
+ run_command(['make', '-j', str(os.cpu_count() or 1)], cwd=ucx_source_path)
135
+ run_command(['make', 'install'], cwd=ucx_source_path)
136
+ print("--- UCX build and install complete ---", flush=True)
137
+
138
+ # -- Step 2: Build NIXL wheel from source --
139
+ print("\n[2/3] Building NIXL wheel from source...", flush=True)
140
+ if not os.path.exists(NIXL_DIR):
141
+ run_command(['git', 'clone', NIXL_REPO_URL, NIXL_DIR])
142
+
143
+ build_env = os.environ.copy()
144
+ build_env['PKG_CONFIG_PATH'] = os.path.join(ucx_install_path, 'lib',
145
+ 'pkgconfig')
146
+ ucx_lib_path = os.path.join(ucx_install_path, 'lib')
147
+ ucx_plugin_path = os.path.join(ucx_lib_path, 'ucx')
148
+ existing_ld_path = os.environ.get('LD_LIBRARY_PATH', '')
149
+ build_env['LD_LIBRARY_PATH'] = \
150
+ f"{ucx_lib_path}:{ucx_plugin_path}:{existing_ld_path}".strip(':')
151
+ print(f"--> Using LD_LIBRARY_PATH: {build_env['LD_LIBRARY_PATH']}",
152
+ flush=True)
153
+
154
+ temp_wheel_dir = os.path.join(ROOT_DIR, 'temp_wheelhouse')
155
+ run_command([
156
+ sys.executable, '-m', 'pip', 'wheel', '.', '--no-deps',
157
+ f'--wheel-dir={temp_wheel_dir}'
158
+ ],
159
+ cwd=os.path.abspath(NIXL_DIR),
160
+ env=build_env)
161
+
162
+ # -- Step 3: Repair the wheel by copying UCX libraries --
163
+ print("\n[3/3] Repairing NIXL wheel to include UCX libraries...",
164
+ flush=True)
165
+ unrepaired_wheel = find_nixl_wheel_in_cache(temp_wheel_dir)
166
+ if not unrepaired_wheel:
167
+ raise RuntimeError("Failed to find the NIXL wheel after building it.")
168
+
169
+ # We tell auditwheel to ignore the plugin that mesonpy already handled.
170
+ auditwheel_command = [
171
+ 'auditwheel',
172
+ 'repair',
173
+ '--exclude',
174
+ 'libplugin_UCX.so', # <-- Exclude because mesonpy already includes it
175
+ unrepaired_wheel,
176
+ f'--wheel-dir={WHEELS_CACHE_HOME}'
177
+ ]
178
+ run_command(auditwheel_command, env=build_env)
179
+
180
+ # --- CLEANUP ---
181
+ # No more temporary files to remove, just the temp wheelhouse
182
+ run_command(['rm', '-rf', temp_wheel_dir])
183
+ # --- END CLEANUP ---
184
+
185
+ newly_built_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME)
186
+ if not newly_built_wheel:
187
+ raise RuntimeError("Failed to find the repaired NIXL wheel.")
188
+
189
+ print(f"--> Successfully built self-contained wheel: \
190
+ {os.path.basename(newly_built_wheel)}. Now installing...",
191
+ flush=True)
192
+ install_command = [
193
+ sys.executable, '-m', 'pip', 'install', newly_built_wheel
194
+ ]
195
+ if args.force_reinstall:
196
+ install_command.insert(-1, '--force-reinstall')
197
+
198
+ run_command(install_command)
199
+ print("--- NIXL installation complete ---", flush=True)
200
+
201
+
202
+if __name__ == "__main__":
203
+ parser = argparse.ArgumentParser(
204
+ description="Build and install UCX and NIXL dependencies.")
205
+ parser.add_argument('--force-reinstall',
206
+ action='store_true',
207
+ help='Force rebuild and reinstall of UCX and NIXL \
208
+ even if they are already installed.')
209
+ args = parser.parse_args()
210
+ build_and_install_prerequisites(args)