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 3, 16:30 UTC
Update Dockerfile.rocm for AINIC & Thor NIC (#40453)
Signed-off-by: root <root@gbt350-odcdh5-wbb3.png-odc.dcgpu> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com> Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Co-authored-by: root <root@gbt350-odcdh5-wbb3.png-odc.dcgpu> Co-authored-by: Jhao-Ting Chen <jhaotingc@nvidia.com> Co-authored-by: simondanielsson <simon.danielsson99@hotmail.com>
haic0 committed
May 14, 2026 at 15:24 UTC
ce29c26b31d432b1b4bc028c46bb2c3b07a667d8
1 file changed
+53
-32
docker/Dockerfile.rocm
+53
-32
@@ -2,10 +2,18 @@
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
5
+# NIC backend for MoRI RDMA support.
6
+# By default (all), drivers and userspace libraries for all supported NIC types
7
+# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
8
+# To install drivers for a single NIC type only, set NIC_BACKEND explicitly:
9
+# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only
10
+# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only
11
+# --build-arg NIC_BACKEND=none # Install nothing.
12
+ARG NIC_BACKEND=all
13
# AMD AINIC apt repo settings
8
-ARG AINIC_VERSION=1.117.5
14
+# Users can specify a custom version compatible with their host drivers.
15
+# The default version has been tested with ioinic-dkms=25.11.1.001
16
+ARG AINIC_VERSION=1.117.3-hydra
17
ARG UBUNTU_CODENAME=jammy
18
19
# Sccache configuration (only used in release pipeline)
@@ -231,39 +239,52 @@ ARG NIC_BACKEND
239
ARG AINIC_VERSION
240
ARG UBUNTU_CODENAME
241
RUN /bin/bash -lc 'set -euo pipefail; \
242
+ \
243
+ install_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/*; \
246
+ mkdir -p /etc/apt/keyrings; \
247
+ curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
248
+ echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
249
+ > /etc/apt/sources.list.d/amdainic.list; \
250
+ apt-get update && apt-get install -y --no-install-recommends \
251
+ libionic-dev \
252
+ ionic-common \
253
+ ; \
254
+ rm -rf /var/lib/apt/lists/*; \
255
+ }; \
256
+ \
257
+ # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \
258
+ # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \
259
+ install_bnxt() { \
260
+ install -m 0755 -d /etc/apt/keyrings; \
261
+ curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \
262
+ -o /etc/apt/keyrings/broadcom-nic.asc; \
263
+ chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \
264
+ echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \
265
+ > /etc/apt/sources.list.d/broadcom-nic.list; \
266
+ apt-get update && apt-get install -y --no-install-recommends \
267
+ bnxt-rocelib=235.2.86.0 \
268
+ ; \
269
+ cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \
270
+ ldconfig; \
271
+ rm -rf /var/lib/apt/lists/*; \
272
+ }; \
273
+ \
274
echo "[MORI] Install MoRI proxy deps"; \
275
pip install --quiet --ignore-installed blinker && \
276
pip install --quiet quart msgpack aiohttp pyzmq; \
277
echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
238
- \
239
- # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override).
240
- # Only vendor packages are installed here for dlopen (e.g. libionic.so); no compile-time flags needed.
241
- case "${NIC_BACKEND}" in \
242
- # default: mlx5
243
- none) \
244
- ;; \
245
- # AMD NIC
246
- ainic) \
247
- apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \
248
- rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \
249
- curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
250
- echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
251
- > /etc/apt/sources.list.d/amdainic.list; \
252
- apt-get update && apt-get install -y --no-install-recommends \
253
- libionic-dev \
254
- ionic-common \
255
- ; \
256
- rm -rf /var/lib/apt/lists/*; \
257
- ;; \
258
- # TODO: Add Broadcom bnxt packages/repos here later.
259
- # bnxt) \
260
- # echo "[MORI] Add Broadcom bnxt packages/repos here later."; \
261
- # ;; \
262
- *) \
263
- echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \
264
- exit 2; \
265
- ;; \
266
- esac;'
278
+ \
279
+ # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \
280
+ # Only vendor packages are installed here for dlopen; no compile-time flags needed. \
281
+ case "${NIC_BACKEND}" in \
282
+ none) ;; \
283
+ all) install_ainic; install_bnxt ;; \
284
+ ainic) install_ainic ;; \
285
+ bnxt) install_bnxt ;; \
286
+ *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
287
+ esac'
288
289
# -----------------------
290
# vLLM wheel release build stage (for building distributable wheels)