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, 20:15 UTC
[CI]Test Group 'NixlConnector PD accuracy tests' is fixed (#31460)
Signed-off-by: qli88 <qiang.li2@amd.com>
qli88 committed
Dec 29, 2025 at 17:48 UTC
0f35429a0c4b2fd1b6b64e4af042f51afa2ba4a9
5 files changed
+157
-7
.buildkite/test-amd.yaml
+2
-2
@@ -1453,8 +1453,8 @@ steps:
1453
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
1454
- tests/v1/kv_connector/nixl_integration/
1455
commands:
1456
- - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
1457
- - bash v1/kv_connector/nixl_integration/tp_config_sweep_accuracy_test.sh
1456
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
1457
+ - VLLM_ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/tp_config_sweep_accuracy_test.sh
1458
1459
##### multi gpus test #####
1460
##### A100 test #####
docker/Dockerfile.rocm_base
+136
-1
@@ -12,6 +12,17 @@ ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
12
ARG AITER_BRANCH="6af8b687"
13
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14
15
+#TODO: When patch has been upstreamed, switch to the main repo/branch
16
+# ARG RIXL_BRANCH="<TODO>"
17
+# ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
18
+ARG RIXL_BRANCH="50d63d94"
19
+ARG RIXL_REPO="https://github.com/vcave/RIXL.git"
20
+# Needed by RIXL
21
+ARG ETCD_BRANCH="7c6e714f"
22
+ARG ETCD_REPO="https://github.com/etcd-cpp-apiv3/etcd-cpp-apiv3.git"
23
+ARG UCX_BRANCH="da3fac2a"
24
+ARG UCX_REPO="https://github.com/ROCm/ucx.git"
25
+
26
FROM ${BASE_IMAGE} AS base
27
28
ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
@@ -50,6 +61,10 @@ RUN apt-get update -y \
61
RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython
62
RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/*
63
64
+
65
+###
66
+### Triton Build
67
+###
68
FROM base AS build_triton
69
ARG TRITON_BRANCH
70
ARG TRITON_REPO
@@ -62,11 +77,19 @@ RUN cd triton \
77
RUN if [ -d triton/python/triton_kernels ]; then pip install build && cd triton/python/triton_kernels \
78
&& python3 -m build --wheel && cp dist/*.whl /app/install; fi
79
80
+
81
+###
82
+### AMD SMI Build
83
+###
84
FROM base AS build_amdsmi
85
RUN cd /opt/rocm/share/amd_smi \
86
&& pip wheel . --wheel-dir=dist
87
RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install
88
89
+
90
+###
91
+### Pytorch build
92
+###
93
FROM base AS build_pytorch
94
ARG PYTORCH_BRANCH
95
ARG PYTORCH_VISION_BRANCH
@@ -95,6 +118,96 @@ RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
118
&& cp /app/vision/dist/*.whl /app/install \
119
&& cp /app/audio/dist/*.whl /app/install
120
121
+
122
+###
123
+### RIXL Build
124
+###
125
+FROM build_pytorch AS build_rixl
126
+ARG RIXL_BRANCH
127
+ARG RIXL_REPO
128
+ARG ETCD_BRANCH
129
+ARG ETCD_REPO
130
+ARG UCX_BRANCH
131
+ARG UCX_REPO
132
+
133
+ENV ROCM_PATH=/opt/rocm
134
+ENV UCX_HOME=/usr/local/ucx
135
+ENV RIXL_HOME=/usr/local/rixl
136
+ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
137
+
138
+# RIXL build system dependences and RDMA support
139
+RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
140
+ libgrpc-dev \
141
+ libgrpc++-dev \
142
+ libprotobuf-dev \
143
+ protobuf-compiler-grpc \
144
+ libcpprest-dev \
145
+ libaio-dev \
146
+ librdmacm1 \
147
+ librdmacm-dev \
148
+ libibverbs1 \
149
+ libibverbs-dev \
150
+ ibverbs-utils \
151
+ rdmacm-utils \
152
+ ibverbs-providers
153
+
154
+RUN pip install meson auditwheel patchelf tomlkit
155
+
156
+WORKDIR /workspace
157
+
158
+RUN git clone ${ETCD_REPO} && \
159
+ cd etcd-cpp-apiv3 && \
160
+ git checkout ${ETCD_BRANCH} && \
161
+ mkdir build && cd build && \
162
+ cmake .. -DCMAKE_POLICY_VERSION_MINIMUM=3.5 && \
163
+ make -j$(nproc) && \
164
+ make install
165
+
166
+RUN cd /usr/local/src && \
167
+ git clone ${UCX_REPO} && \
168
+ cd ucx && \
169
+ git checkout ${UCX_BRANCH} && \
170
+ ./autogen.sh && \
171
+ mkdir build && cd build && \
172
+ ../configure \
173
+ --prefix=/usr/local/ucx \
174
+ --enable-shared \
175
+ --disable-static \
176
+ --disable-doxygen-doc \
177
+ --enable-optimizations \
178
+ --enable-devel-headers \
179
+ --with-rocm=/opt/rocm \
180
+ --with-verbs \
181
+ --with-dm \
182
+ --enable-mt && \
183
+ make -j && \
184
+ make -j install
185
+
186
+ENV PATH=/usr/local/ucx/bin:$PATH
187
+ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
188
+
189
+RUN git clone ${RIXL_REPO} /opt/rixl && \
190
+ cd /opt/rixl && \
191
+ git checkout ${RIXL_BRANCH} && \
192
+ meson setup build --prefix=${RIXL_HOME} \
193
+ -Ducx_path=${UCX_HOME} \
194
+ -Drocm_path=${ROCM_PATH} && \
195
+ cd build && \
196
+ ninja && \
197
+ ninja install
198
+
199
+# Generate RIXL wheel
200
+RUN cd /opt/rixl && mkdir -p /app/install && \
201
+ ./contrib/build-wheel.sh \
202
+ --output-dir /app/install \
203
+ --rocm-dir ${ROCM_PATH} \
204
+ --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
205
+ --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
206
+
207
+
208
+###
209
+### FlashAttention Build
210
+###
211
FROM base AS build_fa
212
ARG FA_BRANCH
213
ARG FA_REPO
@@ -107,6 +220,10 @@ RUN cd flash-attention \
220
&& GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist
221
RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
222
223
+
224
+###
225
+### AITER Build
226
+###
227
FROM base AS build_aiter
228
ARG AITER_BRANCH
229
ARG AITER_REPO
@@ -120,6 +237,10 @@ RUN cd aiter \
237
RUN pip install pyyaml && cd aiter && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist && ls /app/aiter/dist/*.whl
238
RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
239
240
+
241
+###
242
+### Final Build
243
+###
244
FROM base AS debs
245
RUN mkdir /app/debs
246
RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
@@ -132,6 +253,8 @@ RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
253
cp /install/*.whl /app/debs
254
RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
255
cp /install/*.whl /app/debs
256
+RUN --mount=type=bind,from=build_rixl,src=/app/install/,target=/install \
257
+ cp /install/*.whl /app/debs
258
259
FROM base AS final
260
RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
@@ -150,6 +273,12 @@ ARG FA_BRANCH
273
ARG FA_REPO
274
ARG AITER_BRANCH
275
ARG AITER_REPO
276
+ARG RIXL_BRANCH
277
+ARG RIXL_REPO
278
+ARG ETCD_BRANCH
279
+ARG ETCD_REPO
280
+ARG UCX_BRANCH
281
+ARG UCX_REPO
282
RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
283
&& echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
284
&& echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
@@ -162,4 +291,10 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
291
&& echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
292
&& echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
293
&& echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
165
- && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt
294
+ && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \
295
+ && echo "RIXL_BRANCH: ${RIXL_BRANCH}" >> /app/versions.txt \
296
+ && echo "RIXL_REPO: ${RIXL_REPO}" >> /app/versions.txt \
297
+ && echo "ETCD_BRANCH: ${ETCD_BRANCH}" >> /app/versions.txt \
298
+ && echo "ETCD_REPO: ${ETCD_REPO}" >> /app/versions.txt \
299
+ && echo "UCX_BRANCH: ${UCX_BRANCH}" >> /app/versions.txt \
300
+ && echo "UCX_REPO: ${UCX_REPO}" >> /app/versions.txt
docs/features/nixl_connector_usage.md
+7
-1
@@ -6,11 +6,17 @@ NixlConnector is a high-performance KV cache transfer connector for vLLM's disag
6
7
### Installation
8
9
-Install the NIXL library: `uv pip install nixl`, as a quick start.
9
+Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia platform.
10
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](../../requirements/kv_connectors.txt) and other relevant config files
13
14
+For ROCm platform, the [base ROCm docker file](../../docker/Dockerfile.rocm_base) includes RIXL and ucx already.
15
+
16
+- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
17
+- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
18
+- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages
19
+
20
For non-cuda platform, please install nixl with ucx build from source, instructed as below.
21
22
```bash
requirements/kv_connectors_rocm.txt
new
+2
@@ -0,0 +1,2 @@
1
+tblib
2
+lm_eval[api]
\ No newline at end of file
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+10
-3
@@ -84,8 +84,12 @@ logger = init_logger(__name__)
84
85
# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used
86
try:
87
- from nixl._api import nixl_agent as NixlWrapper
88
- from nixl._bindings import nixlXferTelemetry
87
+ if not current_platform.is_rocm():
88
+ from nixl._api import nixl_agent as NixlWrapper
89
+ from nixl._bindings import nixlXferTelemetry
90
+ else:
91
+ from rixl._api import nixl_agent as NixlWrapper
92
+ from rixl._bindings import nixlXferTelemetry
93
94
logger.info("NIXL is available")
95
except ImportError:
@@ -95,7 +99,10 @@ except ImportError:
99
100
101
try:
98
- from nixl._api import nixl_agent_config
102
+ if not current_platform.is_rocm():
103
+ from nixl._api import nixl_agent_config
104
+ else:
105
+ from rixl._api import nixl_agent_config
106
except ImportError:
107
nixl_agent_config = None
108
logger.warning("NIXL agent config is not available")