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, 04:00 UTC
[CI] Move rixl/ucx from Dockerfile.rocm_base to Dockerfile.rocm (#32295)
Signed-off-by: Qiang Li <qiang.li2@amd.com>
qli88 committed
Jan 14, 2026 at 10:53 UTC
3a612322eb9351594bc0c35fe5b85426931be36e
2 files changed
+76
-111
docker/Dockerfile.rocm
+76
@@ -67,6 +67,78 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
67
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
68
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
69
70
+# -----------------------
71
+# RIXL/UCX build stages
72
+FROM base AS build_rixl
73
+ARG RIXL_BRANCH="f33a5599"
74
+ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
75
+ARG UCX_BRANCH="da3fac2a"
76
+ARG UCX_REPO="https://github.com/ROCm/ucx.git"
77
+ENV ROCM_PATH=/opt/rocm
78
+ENV UCX_HOME=/usr/local/ucx
79
+ENV RIXL_HOME=/usr/local/rixl
80
+ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
81
+
82
+# RIXL build system dependences and RDMA support
83
+RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
84
+ libgrpc-dev \
85
+ libgrpc++-dev \
86
+ libprotobuf-dev \
87
+ protobuf-compiler-grpc \
88
+ libcpprest-dev \
89
+ libaio-dev \
90
+ librdmacm1 \
91
+ librdmacm-dev \
92
+ libibverbs1 \
93
+ libibverbs-dev \
94
+ ibverbs-utils \
95
+ rdmacm-utils \
96
+ ibverbs-providers \
97
+ && rm -rf /var/lib/apt/lists/*
98
+
99
+RUN uv pip install --system meson auditwheel patchelf tomlkit
100
+
101
+RUN cd /usr/local/src && \
102
+ git clone ${UCX_REPO} && \
103
+ cd ucx && \
104
+ git checkout ${UCX_BRANCH} && \
105
+ ./autogen.sh && \
106
+ mkdir build && cd build && \
107
+ ../configure \
108
+ --prefix=/usr/local/ucx \
109
+ --enable-shared \
110
+ --disable-static \
111
+ --disable-doxygen-doc \
112
+ --enable-optimizations \
113
+ --enable-devel-headers \
114
+ --with-rocm=/opt/rocm \
115
+ --with-verbs \
116
+ --with-dm \
117
+ --enable-mt && \
118
+ make -j && \
119
+ make install
120
+
121
+ENV PATH=/usr/local/ucx/bin:$PATH
122
+ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
123
+
124
+RUN git clone ${RIXL_REPO} /opt/rixl && \
125
+ cd /opt/rixl && \
126
+ git checkout ${RIXL_BRANCH} && \
127
+ meson setup build --prefix=${RIXL_HOME} \
128
+ -Ducx_path=${UCX_HOME} \
129
+ -Drocm_path=${ROCM_PATH} && \
130
+ cd build && \
131
+ ninja && \
132
+ ninja install
133
+
134
+# Generate RIXL wheel
135
+RUN cd /opt/rixl && mkdir -p /app/install && \
136
+ ./contrib/build-wheel.sh \
137
+ --output-dir /app/install \
138
+ --rocm-dir ${ROCM_PATH} \
139
+ --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
140
+ --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
141
+
142
# -----------------------
143
# Test vLLM image
144
FROM base AS test
@@ -83,6 +155,10 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
155
&& pip uninstall -y vllm \
156
&& uv pip install --system *.whl
157
158
+# Install RIXL wheel
159
+RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
160
+ uv pip install --system /rixl_install/*.whl
161
+
162
WORKDIR /vllm-workspace
163
ARG COMMON_WORKDIR
164
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
docker/Dockerfile.rocm_base
-111
@@ -14,17 +14,6 @@ ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14
ARG MORI_BRANCH="2d02c6a9"
15
ARG MORI_REPO="https://github.com/ROCm/mori.git"
16
17
-#TODO: When patch has been upstreamed, switch to the main repo/branch
18
-# ARG RIXL_BRANCH="<TODO>"
19
-# ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
20
-ARG RIXL_BRANCH="50d63d94"
21
-ARG RIXL_REPO="https://github.com/vcave/RIXL.git"
22
-# Needed by RIXL
23
-ARG ETCD_BRANCH="7c6e714f"
24
-ARG ETCD_REPO="https://github.com/etcd-cpp-apiv3/etcd-cpp-apiv3.git"
25
-ARG UCX_BRANCH="da3fac2a"
26
-ARG UCX_REPO="https://github.com/ROCm/ucx.git"
27
-
17
FROM ${BASE_IMAGE} AS base
18
19
ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
@@ -138,92 +127,6 @@ RUN cd mori \
127
RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install
128
129
141
-###
142
-### RIXL Build
143
-###
144
-FROM build_pytorch AS build_rixl
145
-ARG RIXL_BRANCH
146
-ARG RIXL_REPO
147
-ARG ETCD_BRANCH
148
-ARG ETCD_REPO
149
-ARG UCX_BRANCH
150
-ARG UCX_REPO
151
-
152
-ENV ROCM_PATH=/opt/rocm
153
-ENV UCX_HOME=/usr/local/ucx
154
-ENV RIXL_HOME=/usr/local/rixl
155
-ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
156
-
157
-# RIXL build system dependences and RDMA support
158
-RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
159
- libgrpc-dev \
160
- libgrpc++-dev \
161
- libprotobuf-dev \
162
- protobuf-compiler-grpc \
163
- libcpprest-dev \
164
- libaio-dev \
165
- librdmacm1 \
166
- librdmacm-dev \
167
- libibverbs1 \
168
- libibverbs-dev \
169
- ibverbs-utils \
170
- rdmacm-utils \
171
- ibverbs-providers
172
-
173
-RUN pip install meson auditwheel patchelf tomlkit
174
-
175
-WORKDIR /workspace
176
-
177
-RUN git clone ${ETCD_REPO} && \
178
- cd etcd-cpp-apiv3 && \
179
- git checkout ${ETCD_BRANCH} && \
180
- mkdir build && cd build && \
181
- cmake .. -DCMAKE_POLICY_VERSION_MINIMUM=3.5 && \
182
- make -j$(nproc) && \
183
- make install
184
-
185
-RUN cd /usr/local/src && \
186
- git clone ${UCX_REPO} && \
187
- cd ucx && \
188
- git checkout ${UCX_BRANCH} && \
189
- ./autogen.sh && \
190
- mkdir build && cd build && \
191
- ../configure \
192
- --prefix=/usr/local/ucx \
193
- --enable-shared \
194
- --disable-static \
195
- --disable-doxygen-doc \
196
- --enable-optimizations \
197
- --enable-devel-headers \
198
- --with-rocm=/opt/rocm \
199
- --with-verbs \
200
- --with-dm \
201
- --enable-mt && \
202
- make -j && \
203
- make -j install
204
-
205
-ENV PATH=/usr/local/ucx/bin:$PATH
206
-ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
207
-
208
-RUN git clone ${RIXL_REPO} /opt/rixl && \
209
- cd /opt/rixl && \
210
- git checkout ${RIXL_BRANCH} && \
211
- meson setup build --prefix=${RIXL_HOME} \
212
- -Ducx_path=${UCX_HOME} \
213
- -Drocm_path=${ROCM_PATH} && \
214
- cd build && \
215
- ninja && \
216
- ninja install
217
-
218
-# Generate RIXL wheel
219
-RUN cd /opt/rixl && mkdir -p /app/install && \
220
- ./contrib/build-wheel.sh \
221
- --output-dir /app/install \
222
- --rocm-dir ${ROCM_PATH} \
223
- --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
224
- --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
225
-
226
-
130
###
131
### FlashAttention Build
132
###
@@ -274,8 +177,6 @@ RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
177
cp /install/*.whl /app/debs
178
RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
179
cp /install/*.whl /app/debs
277
-RUN --mount=type=bind,from=build_rixl,src=/app/install/,target=/install \
278
- cp /install/*.whl /app/debs
180
181
FROM base AS final
182
RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
@@ -294,12 +195,6 @@ ARG FA_BRANCH
195
ARG FA_REPO
196
ARG AITER_BRANCH
197
ARG AITER_REPO
297
-ARG RIXL_BRANCH
298
-ARG RIXL_REPO
299
-ARG ETCD_BRANCH
300
-ARG ETCD_REPO
301
-ARG UCX_BRANCH
302
-ARG UCX_REPO
198
ARG MORI_BRANCH
199
ARG MORI_REPO
200
RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
@@ -315,11 +210,5 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
210
&& echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
211
&& echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
212
&& echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \
318
- && echo "RIXL_BRANCH: ${RIXL_BRANCH}" >> /app/versions.txt \
319
- && echo "RIXL_REPO: ${RIXL_REPO}" >> /app/versions.txt \
320
- && echo "ETCD_BRANCH: ${ETCD_BRANCH}" >> /app/versions.txt \
321
- && echo "ETCD_REPO: ${ETCD_REPO}" >> /app/versions.txt \
322
- && echo "UCX_BRANCH: ${UCX_BRANCH}" >> /app/versions.txt \
323
- && echo "UCX_REPO: ${UCX_REPO}" >> /app/versions.txt \
213
&& echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \
214
&& echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt