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 4, 22:30 UTC
Enable building MoRI with AMD AINIC stack (#38371)
Signed-off-by: Theresa Shan <thshan@smci355-ccs-aus-n08-21.prov.aus.ccs.cpe.ice.amd.com> Signed-off-by: Theresa Shan <theresa.shan@amd.com> Co-authored-by: Theresa Shan <thshan@smci355-ccs-aus-n08-21.prov.aus.ccs.cpe.ice.amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Theresa Shan committed
Apr 21, 2026 at 02:17 UTC
2390caf157cbfb5f73c5994158bca654ddae2706
2 files changed
+54
-4
docker/Dockerfile.rocm
+53
-3
@@ -2,6 +2,11 @@
2
ARG REMOTE_VLLM="0"
3
ARG COMMON_WORKDIR=/app
4
ARG BASE_IMAGE=rocm/vllm-dev:base
5
+# AMD NIC backend
6
+ARG NIC_BACKEND=none
7
+# AMD AINIC apt repo settings
8
+ARG AINIC_VERSION=1.117.5
9
+ARG UBUNTU_CODENAME=jammy
10
11
# Sccache configuration (only used in release pipeline)
12
ARG USE_SCCACHE
@@ -216,6 +221,47 @@ RUN git clone ${DEEPEP_REPO} \
221
&& git checkout ${DEEPEP_BRANCH} \
222
&& python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
223
224
+# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
225
+# not force users to rebuild the long-lived Dockerfile.rocm_base image.
226
+FROM base AS mori_base
227
+ARG NIC_BACKEND
228
+ARG AINIC_VERSION
229
+ARG UBUNTU_CODENAME
230
+RUN /bin/bash -lc 'set -euo pipefail; \
231
+ echo "[MORI] Install MoRI proxy deps"; \
232
+ pip install --quiet --ignore-installed blinker && \
233
+ pip install --quiet quart msgpack aiohttp pyzmq; \
234
+ echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
235
+ \
236
+ # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override).
237
+ # Only vendor packages are installed here for dlopen (e.g. libionic.so); no compile-time flags needed.
238
+ case "${NIC_BACKEND}" in \
239
+ # default: mlx5
240
+ none) \
241
+ ;; \
242
+ # AMD NIC
243
+ ainic) \
244
+ apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \
245
+ rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \
246
+ curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
247
+ echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
248
+ > /etc/apt/sources.list.d/amdainic.list; \
249
+ apt-get update && apt-get install -y --no-install-recommends \
250
+ libionic-dev \
251
+ ionic-common \
252
+ ; \
253
+ rm -rf /var/lib/apt/lists/*; \
254
+ ;; \
255
+ # TODO: Add Broadcom bnxt packages/repos here later.
256
+ # bnxt) \
257
+ # echo "[MORI] Add Broadcom bnxt packages/repos here later."; \
258
+ # ;; \
259
+ *) \
260
+ echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \
261
+ exit 2; \
262
+ ;; \
263
+ esac;'
264
+
265
# -----------------------
266
# vLLM wheel release build stage (for building distributable wheels)
267
# This stage pins dependencies to custom ROCm wheel versions and handles version detection
@@ -294,7 +340,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
340
341
# -----------------------
342
# Test vLLM image
297
-FROM base AS test
343
+FROM mori_base AS test
344
345
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
346
@@ -385,7 +431,7 @@ RUN printf '%s\n' \
431
432
# -----------------------
433
# Final vLLM image
388
-FROM base AS final
434
+FROM mori_base AS final
435
436
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
437
@@ -422,6 +468,8 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
468
469
ARG COMMON_WORKDIR
470
ARG BASE_IMAGE
471
+ARG NIC_BACKEND
472
+ARG AINIC_VERSION
473
474
# Copy over the benchmark scripts as well
475
COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
@@ -439,7 +487,9 @@ ENV HIP_FORCE_DEV_KERNARG=1
487
# Workaround for ROCm profiler limits
488
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
489
ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
442
-RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt
490
+RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
491
+ && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
492
+ && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
493
494
CMD ["/bin/bash"]
495
docker/Dockerfile.rocm_base
+1
-1
@@ -11,7 +11,7 @@ ARG FA_BRANCH="0e60e394"
11
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
12
ARG AITER_BRANCH="v0.1.10.post3"
13
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14
-ARG MORI_BRANCH="2d02c6a9"
14
+ARG MORI_BRANCH="v1.1.0"
15
ARG MORI_REPO="https://github.com/ROCm/mori.git"
16
17
# Sccache configuration (only used in release pipeline)