[ROCm] Upgrade NIXL and UCX (#49251)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Jul 21, 2026 at 22:30 UTC
16aca639b747301c398bc4eed1cf5f0af8510538
10 files changed
+127
-102
.buildkite/scripts/ci-bake-rocm.sh
+28
-28
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
17
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
18
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
19
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
20
-DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base"
20
+DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
21
DEFAULT_CI_BASE_METADATA_VERSION="1"
22
IMAGE_EXISTED_BEFORE_BUILD=0
23
@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
1159
metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
1160
metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
1161
metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
1162
- metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
1163
- metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
1162
+ metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
1163
+ metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
1164
metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
1165
metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
1166
metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
1169
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
1170
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
1171
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
1172
- metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
1172
+ metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
1173
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
1174
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
1175
@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
1686
return 0
1687
fi
1688
1689
- for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
1689
+ for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
1690
if [[ -n "${!var:-}" ]]; then
1691
echo "Using provided ${var}: ${!var}"
1692
continue
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
1706
compute_dependency_cache_keys() {
1707
local bake_dir=""
1708
local dockerfile_rocm=""
1709
- local rixl_branch=""
1709
+ local nixl_branch=""
1710
local ucx_branch=""
1711
local rocshmem_branch=""
1712
local deepep_branch=""
1713
- local rixl_material=""
1713
+ local nixl_material=""
1714
local rocshmem_material=""
1715
local deepep_material=""
1716
1717
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
1718
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
1719
- rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
1719
+ nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
1720
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
1721
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
1722
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
1723
1724
- if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
1725
- rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
1726
- RIXL_CACHE_KEY=$(
1724
+ if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
1725
+ nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
1726
+ NIXL_CACHE_KEY=$(
1727
compose_dependency_cache_key \
1728
- "${rixl_branch}-ucx-${ucx_branch}" \
1729
- "${rixl_material}"
1728
+ "${nixl_branch}-ucx-${ucx_branch}" \
1729
+ "${nixl_material}"
1730
)
1731
- export RIXL_CACHE_KEY
1732
- echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
1731
+ export NIXL_CACHE_KEY
1732
+ echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
1733
fi
1734
1735
if [[ -n "${rocshmem_branch}" ]]; then
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
1780
local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
1781
1782
case "${target}" in
1783
- rixl-rocm-ci)
1784
- if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
1785
- printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
1786
- elif [[ -n "${RIXL_BRANCH:-}" ]]; then
1787
- printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
1783
+ nixl-rocm-ci)
1784
+ if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
1785
+ printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
1786
+ elif [[ -n "${NIXL_BRANCH:-}" ]]; then
1787
+ printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
1788
fi
1789
;;
1790
rocshmem-rocm-ci)
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {
1815
1816
resolve_ci_base_dependency_targets() {
1817
local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
1818
- local rixl_ref=""
1818
+ local nixl_ref=""
1819
local rocshmem_ref=""
1820
local deepep_ref=""
1821
@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
1824
case "${mode}" in
1825
always)
1826
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
1827
- for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
1827
+ for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
1828
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
1829
add_dependency_cache_target "${target}"
1830
fi
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
1844
;;
1845
esac
1846
1847
- if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
1848
- rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
1849
- if dependency_cache_ref_exists "${rixl_ref}"; then
1850
- echo "RIXL dependency cache exists: ${rixl_ref}"
1847
+ if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
1848
+ nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
1849
+ if dependency_cache_ref_exists "${nixl_ref}"; then
1850
+ echo "NIXL dependency cache exists: ${nixl_ref}"
1851
else
1852
- echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
1853
- add_dependency_cache_target "rixl-rocm-ci"
1852
+ echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
1853
+ add_dependency_cache_target "nixl-rocm-ci"
1854
fi
1855
fi
1856
docker/Dockerfile.rocm
+54
-32
@@ -339,18 +339,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
339
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
340
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
341
342
-# RIXL/UCX build stages
343
-FROM base AS build_rixl
344
-ARG RIXL_BRANCH="39be1de8"
345
-ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
346
-ARG UCX_BRANCH="bfb51733"
342
+# NIXL/UCX build stages
343
+FROM base AS build_nixl
344
+ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
345
+ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
346
+ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
347
ARG UCX_REPO="https://github.com/openucx/ucx.git"
348
ENV ROCM_PATH=/opt/rocm
349
ENV UCX_HOME=/usr/local/ucx
350
-ENV RIXL_HOME=/usr/local/rixl
351
-ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
350
+ENV NIXL_HOME=/usr/local/nixl
351
353
-# RIXL build system dependences and RDMA support
352
+# NIXL build system dependencies and RDMA support
353
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
354
libgrpc-dev \
355
libgrpc++-dev \
@@ -368,7 +367,8 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
367
&& rm -rf /var/lib/apt/lists/*
368
369
RUN --mount=type=cache,target=/root/.cache/uv \
371
- uv pip install --system meson auditwheel patchelf tomlkit
370
+ uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
371
+ auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
372
373
RUN --mount=type=cache,target=/root/.cache/ccache \
374
cd /usr/local/src && \
@@ -396,30 +396,50 @@ ENV PATH=/usr/local/ucx/bin:$PATH
396
ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
397
398
RUN --mount=type=cache,target=/root/.cache/ccache \
399
- git clone ${RIXL_REPO} /opt/rixl && \
400
- cd /opt/rixl && \
401
- git checkout ${RIXL_BRANCH} && \
399
+ git clone ${NIXL_REPO} /opt/nixl && \
400
+ cd /opt/nixl && \
401
+ git checkout ${NIXL_BRANCH} && \
402
CC="ccache gcc" CXX="ccache g++" \
403
- meson setup build --prefix=${RIXL_HOME} \
403
+ meson setup build --prefix=${NIXL_HOME} \
404
-Ducx_path=${UCX_HOME} \
405
- -Drocm_path=${ROCM_PATH} && \
405
+ -Dwheel_variant=rocm \
406
+ -Dbuild_tests=false \
407
+ -Dbuild_examples=false && \
408
cd build && \
409
ninja -j$(nproc) && \
408
- ninja install
409
-
410
-# Generate RIXL wheel
410
+ ninja install && \
411
+ echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
412
+ > /etc/ld.so.conf.d/nixl.conf && \
413
+ echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
414
+ >> /etc/ld.so.conf.d/nixl.conf && \
415
+ ldconfig
416
+
417
+# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
418
+# so configure the ROCm wheel variant directly through Meson.
419
# Exclude libcore and libpull from auditwheel: transitive dependencies
420
# that are not shipped in the wheel and vary across base images.
413
-RUN cd /opt/rixl && \
414
- sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
415
- contrib/build-wheel.sh && \
416
- mkdir -p /app/install && \
417
- _ucx_install_dir=${UCX_HOME} \
418
- ./contrib/build-wheel.sh \
419
- --output-dir /app/install \
420
- --rocm-dir ${ROCM_PATH} \
421
+RUN cd /opt/nixl && \
422
+ ./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
423
+ CC="ccache gcc" CXX="ccache g++" \
424
+ uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
425
+ --python ${PYTHON_VERSION} \
426
+ -Csetup-args=-Ducx_path=${UCX_HOME} \
427
+ -Csetup-args=-Dwheel_variant=rocm \
428
+ -Csetup-args=-Dbuild_tests=false \
429
+ -Csetup-args=-Dbuild_examples=false && \
430
+ mkdir -p /tmp/nixl_wheels/repaired /app/install && \
431
+ auditwheel repair \
432
+ --exclude 'libamdhip64*' \
433
+ --exclude 'libcore*' \
434
+ --exclude 'libpull*' \
435
+ /tmp/nixl_wheels/nixl_rocm*.whl \
436
+ --plat manylinux_2_34_$(uname -m) \
437
+ --wheel-dir /tmp/nixl_wheels/repaired && \
438
+ ./contrib/wheel_add_ucx_plugins.py \
439
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
422
- --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
440
+ --nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
441
+ /tmp/nixl_wheels/repaired/*.whl && \
442
+ cp /tmp/nixl_wheels/repaired/*.whl /app/install
443
444
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
445
# invalidate the slow ROCShmem build.
@@ -660,10 +680,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
680
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
681
fi
682
663
-# Install RIXL + DeepEP wheels.
664
-RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
683
+# Install NIXL + DeepEP wheels.
684
+RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
685
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
666
- uv pip install --system /rixl_install/*.whl /deep_install/*.whl
686
+ uv pip install --system /nixl_install/*.whl /deep_install/*.whl
687
688
# Copy ROCShmem runtime libraries.
689
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
@@ -724,6 +744,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
744
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
745
# See: https://github.com/ROCm/rocm-libraries/issues/6266
746
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
747
+ENV UCX_RMA_PPLN_ENABLE=y
748
749
# ROCm profiler limits workaround.
750
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
@@ -796,9 +817,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
817
&& pip uninstall -y vllm \
818
&& uv pip install --system *.whl
819
799
-# Install RIXL wheel
800
-RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
801
- uv pip install --system /rixl_install/*.whl
820
+# Install NIXL ROCm wheel
821
+RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
822
+ uv pip install --system /nixl_install/*.whl
823
824
ARG COMMON_WORKDIR
825
ARG BASE_IMAGE
@@ -813,6 +834,7 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
834
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
835
# See: https://github.com/ROCm/rocm-libraries/issues/6266
836
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
837
+ENV UCX_RMA_PPLN_ENABLE=y
838
839
ENV TOKENIZERS_PARALLELISM=false
840
docker/ci-rocm.hcl
+13
-13
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
59
}
60
61
# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
62
-# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
62
+# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
63
# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
64
# stage and is irrelevant when building --target ci_base itself.
65
variable "CI_BASE_IMAGE" {
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
75
# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
76
# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
77
# functions produce no entries when the variable is empty.
78
-variable "RIXL_BRANCH" {
78
+variable "NIXL_BRANCH" {
79
default = ""
80
}
81
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
91
default = ""
92
}
93
94
-variable "RIXL_CACHE_KEY" {
94
+variable "NIXL_CACHE_KEY" {
95
default = ""
96
}
97
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
236
])
237
}
238
239
-# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
239
+# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
240
# These stages are pinned to specific upstream commit hashes, so cache keys use
241
# those hashes rather than the Buildkite commit. This means the cache persists
242
# across all vLLM commits as long as the upstream dependency pins don't change.
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
244
function "get_cache_from_rocm_deps" {
245
params = []
246
result = compact([
247
- RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
247
+ NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
248
ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
249
DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
250
])
251
}
252
253
-function "get_cache_to_rocm_rixl" {
253
+function "get_cache_to_rocm_nixl" {
254
params = []
255
result = compact([
256
- RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
256
+ NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
257
])
258
}
259
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
372
# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
373
# (e.g., requirements change), these stages are cache hits if their upstream
374
# pins haven't changed -- saving ~35min of compilation.
375
-target "rixl-rocm-ci" {
375
+target "nixl-rocm-ci" {
376
inherits = ["_common-rocm", "_ci-rocm"]
377
- target = "build_rixl"
377
+ target = "build_nixl"
378
cache-from = get_cache_from_rocm_deps()
379
- cache-to = get_cache_to_rocm_rixl()
379
+ cache-to = get_cache_to_rocm_nixl()
380
output = ["type=cacheonly"]
381
}
382
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
396
output = ["type=cacheonly"]
397
}
398
399
-# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
399
+# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
400
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
401
# files drifts from the remote image label. Per-PR builds then pull the result
402
# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
412
CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
413
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
414
]),
415
- # Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
415
+ # Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
416
# are cache hits even when ci_base itself needs rebuilding.
417
get_cache_from_rocm_deps(),
418
)
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
424
# Group for ci_base builds -- exports dependency stage caches alongside the
425
# ci_base image so future rebuilds can reuse them independently.
426
group "ci-base-rocm-ci-with-deps" {
427
- targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
427
+ targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
428
}
docker/docker-bake-rocm.hcl
+2
-2
@@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" {
53
# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
54
# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
55
# overrides) and writes a small HCL override before invoking CI targets.
56
-variable "RIXL_BRANCH" {
56
+variable "NIXL_BRANCH" {
57
default = ""
58
}
59
@@ -106,7 +106,7 @@ target "test-rocm" {
106
output = ["type=docker"]
107
}
108
109
-# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
109
+# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
110
# torchcodec, requirements, etc.). Used by the weekly scheduled build and
111
# the auto-rebuild trigger when requirements change in a PR.
112
target "ci-base-rocm" {
docs/features/nixl_connector_usage.md
+1
-5
@@ -13,11 +13,7 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat
13
- Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
14
- The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files
15
16
-For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already.
17
-
18
-- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
19
-- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
20
-- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages
16
+For ROCm, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source.
17
18
For non-cuda platform, please install nixl with ucx build from source, instructed as below.
19
tests/v1/kv_connector/unit/test_nixl_connector.py
+10
-7
@@ -204,13 +204,12 @@ class FakeNixlWrapper:
204
def _make_fake_nixl_pkg():
205
"""Context manager that creates a temporary package making
206
`from nixl._api import nixl_agent` resolve to our FakeNixlWrapper.
207
- Also creates rixl package for ROCm compatibility.
207
+ Also creates the ROCm NIXL packages.
208
209
Automatically cleans up the temporary directory when done.
210
"""
211
with tempfile.TemporaryDirectory() as td:
212
- # Create both nixl and rixl packages for cross-platform compatibility
213
- for pkg_name in ["nixl", "rixl"]:
212
+ for pkg_name in ["nixl", "nixl_rocm"]:
213
pkg_root = os.path.join(td, pkg_name, "_api")
214
os.makedirs(pkg_root, exist_ok=True)
215
@@ -3130,10 +3129,12 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
3129
)
3130
3131
3132
+@patch(
3133
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
3134
+ FakeNixlWrapper,
3135
+)
3136
def test_kv_both_deprecation_warning(default_vllm_config, dist_init):
3137
"""kv_role='kv_both' should emit a deprecation log warning."""
3135
- from unittest.mock import patch
3136
-
3138
from vllm.logger import _print_warning_once
3139
3140
_print_warning_once.cache_clear()
@@ -3156,10 +3157,12 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init):
3157
assert "deprecated" in msg
3158
3159
3160
+@patch(
3161
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
3162
+ FakeNixlWrapper,
3163
+)
3164
def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init):
3165
"""kv_role='kv_consumer' or 'kv_producer' should NOT emit a warning."""
3161
- from unittest.mock import patch
3162
-
3166
for role in ("kv_consumer", "kv_producer"):
3167
vllm_config = create_vllm_config(kv_role=role)
3168
with patch(
tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py
renamed
+7
-7
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Verify that GPU memory is fully released after RixlConnector shutdown on ROCm.
3
+"""Verify that GPU memory is released after NixlConnector shutdown on ROCm.
4
5
Regression test for ROCm/ucx#33: UCX rocm_ipc transport permanently pinned
6
GPU memory via hsa_amd_ipc_memory_create during ucp_mem_map, causing
@@ -62,7 +62,7 @@ def _full_gpu_cleanup():
62
63
64
@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
65
-def test_gpu_memory_rixl_hma(model_name, sw_size):
65
+def test_gpu_memory_nixl_hma(model_name, sw_size):
66
"""Track GPU memory through NixlConnector create/infer/shutdown cycle."""
67
from vllm import LLM, SamplingParams
68
from vllm.config import KVTransferConfig
@@ -84,7 +84,7 @@ def test_gpu_memory_rixl_hma(model_name, sw_size):
84
}
85
86
print("\n" + "=" * 90)
87
- print("GPU MEMORY -- RIXL NixlConnector HMA (ROCm)")
87
+ print("GPU MEMORY -- NIXL NixlConnector HMA (ROCm)")
88
print("=" * 90)
89
gc.collect()
90
torch.accelerator.empty_cache()
@@ -169,14 +169,14 @@ def test_gpu_memory_rixl_hma(model_name, sw_size):
169
170
171
@pytest.mark.parametrize("model_name", ["google/gemma-3-1b-it"])
172
-def test_gpu_memory_no_rixl_baseline(model_name):
172
+def test_gpu_memory_no_nixl_baseline(model_name):
173
"""Same workload without NixlConnector. Comparing driver-level memory
174
- between this and test_gpu_memory_rixl_hma isolates UCX/RIXL impact."""
174
+ between this and test_gpu_memory_nixl_hma isolates UCX/NIXL impact."""
175
from vllm import LLM, SamplingParams
176
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
177
178
print("\n" + "=" * 90)
179
- print("CONTROL -- same model, no RIXL connector")
179
+ print("CONTROL -- same model, no NIXL connector")
180
print("=" * 90)
181
gc.collect()
182
torch.accelerator.empty_cache()
@@ -209,7 +209,7 @@ def test_gpu_memory_no_rixl_baseline(model_name):
209
drv_base = snap0["drv_used_mb"]
210
drv_leaked = snap_final["drv_used_mb"] - drv_base
211
drv_peak = snap_peak["drv_used_mb"] - drv_base
212
- print(f"\n Driver leaked (no rixl): {drv_leaked:.0f} MB")
212
+ print(f"\n Driver leaked (no NIXL): {drv_leaked:.0f} MB")
213
print("=" * 90)
214
215
leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0
vllm/config/parallel.py
+1
-1
@@ -92,7 +92,7 @@ class EPLBConfig:
92
Backend for EPLB expert weight communication:
93
- "torch_nccl": Use torch.distributed on the device process group
94
- "torch_gloo": Use torch.distributed gloo with CPU staging
95
- - "nixl": Use NIXL/ RIXL with staged send/recv buffers
95
+ - "nixl": Use NIXL with staged send/recv buffers
96
- "pynccl": Use PyNccl send/recv
97
- None: Auto-select backend (prefers "nixl", falls back to "torch_gloo")
98
"""
vllm/distributed/eplb/eplb_communicator.py
+2
-2
@@ -38,7 +38,7 @@ logger = init_logger(__name__)
38
39
40
def has_nixl() -> bool:
41
- """Whether the optional NIXL / RIXL package is available."""
41
+ """Whether the optional NIXL package is available."""
42
return nixl_utils.NixlWrapper is not None
43
44
@@ -266,7 +266,7 @@ class NixlEplbCommunicator(EplbCommunicator):
266
assert expert_buffer, "NixlEplbCommunicator requires non-empty expert_buffer."
267
nixl_wrapper_cls = nixl_utils.NixlWrapper
268
if nixl_wrapper_cls is None:
269
- raise RuntimeError("NIXL/ RIXL is unavailable.")
269
+ raise RuntimeError("NIXL is unavailable.")
270
271
self._cpu_group = cpu_group
272
self._world_size = cpu_group.size()
vllm/distributed/nixl_utils.py
+9
-5
@@ -21,7 +21,7 @@ def _maybe_set_ucx_rcache_limit() -> None:
21
if "UCX_RCACHE_MAX_UNRELEASED" in os.environ:
22
return
23
24
- if "nixl" in sys.modules or "rixl" in sys.modules:
24
+ if "nixl" in sys.modules or "nixl_rocm" in sys.modules:
25
logger.warning_once(
26
"NIXL was already imported, we can't reset "
27
"UCX_RCACHE_MAX_UNRELEASED. "
@@ -36,8 +36,12 @@ def _maybe_set_ucx_rcache_limit() -> None:
36
os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024"
37
38
39
+def _get_nixl_package_name() -> str:
40
+ return "nixl_rocm" if current_platform.is_rocm() else "nixl"
41
+
42
+
43
def _get_nixl_module_name(name: str) -> str:
40
- package_name = "rixl" if current_platform.is_rocm() else "nixl"
44
+ package_name = _get_nixl_package_name()
45
if name == "nixlXferTelemetry":
46
return f"{package_name}._bindings"
47
return f"{package_name}._api"
@@ -80,11 +84,11 @@ def __getattr__(name: str) -> Any:
84
85
86
def is_nixl_available() -> bool:
83
- """Lightweight check for nixl/rixl package without importing it."""
87
+ """Lightweight check for the platform's NIXL package without importing it."""
88
import importlib.util
89
86
- pkg = "rixl" if current_platform.is_rocm() else "nixl"
87
- return importlib.util.find_spec(pkg) is not None
90
+ pkg = _get_nixl_package_name()
91
+ return pkg in sys.modules or importlib.util.find_spec(pkg) is not None
92
93
94
__all__ = [