forked from Karylab-cklius/vllm
[ROCm] Upgrade NIXL and UCX (#49251)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
This commit is contained in:
Andreas Karatzas
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
|
||||
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
|
||||
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"
|
||||
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
|
||||
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"
|
||||
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"
|
||||
DEFAULT_CI_BASE_METADATA_VERSION="1"
|
||||
IMAGE_EXISTED_BEFORE_BUILD=0
|
||||
|
||||
@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
|
||||
metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
|
||||
metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
|
||||
metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
|
||||
metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
|
||||
metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
|
||||
metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
|
||||
metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
|
||||
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
|
||||
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
|
||||
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
|
||||
metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
|
||||
|
||||
@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
|
||||
return 0
|
||||
fi
|
||||
|
||||
for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
|
||||
for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
|
||||
if [[ -n "${!var:-}" ]]; then
|
||||
echo "Using provided ${var}: ${!var}"
|
||||
continue
|
||||
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
|
||||
compute_dependency_cache_keys() {
|
||||
local bake_dir=""
|
||||
local dockerfile_rocm=""
|
||||
local rixl_branch=""
|
||||
local nixl_branch=""
|
||||
local ucx_branch=""
|
||||
local rocshmem_branch=""
|
||||
local deepep_branch=""
|
||||
local rixl_material=""
|
||||
local nixl_material=""
|
||||
local rocshmem_material=""
|
||||
local deepep_material=""
|
||||
|
||||
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
|
||||
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
|
||||
rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
|
||||
nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
|
||||
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
|
||||
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
|
||||
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
|
||||
|
||||
if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
|
||||
rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
|
||||
RIXL_CACHE_KEY=$(
|
||||
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
|
||||
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
|
||||
NIXL_CACHE_KEY=$(
|
||||
compose_dependency_cache_key \
|
||||
"${rixl_branch}-ucx-${ucx_branch}" \
|
||||
"${rixl_material}"
|
||||
"${nixl_branch}-ucx-${ucx_branch}" \
|
||||
"${nixl_material}"
|
||||
)
|
||||
export RIXL_CACHE_KEY
|
||||
echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
|
||||
export NIXL_CACHE_KEY
|
||||
echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
|
||||
fi
|
||||
|
||||
if [[ -n "${rocshmem_branch}" ]]; then
|
||||
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
|
||||
local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
|
||||
|
||||
case "${target}" in
|
||||
rixl-rocm-ci)
|
||||
if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
|
||||
elif [[ -n "${RIXL_BRANCH:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
|
||||
nixl-rocm-ci)
|
||||
if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
|
||||
elif [[ -n "${NIXL_BRANCH:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
|
||||
fi
|
||||
;;
|
||||
rocshmem-rocm-ci)
|
||||
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {
|
||||
|
||||
resolve_ci_base_dependency_targets() {
|
||||
local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
|
||||
local rixl_ref=""
|
||||
local nixl_ref=""
|
||||
local rocshmem_ref=""
|
||||
local deepep_ref=""
|
||||
|
||||
@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
|
||||
case "${mode}" in
|
||||
always)
|
||||
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
|
||||
for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
|
||||
for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
|
||||
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
|
||||
add_dependency_cache_target "${target}"
|
||||
fi
|
||||
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
|
||||
;;
|
||||
esac
|
||||
|
||||
if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
|
||||
rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
|
||||
if dependency_cache_ref_exists "${rixl_ref}"; then
|
||||
echo "RIXL dependency cache exists: ${rixl_ref}"
|
||||
if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
|
||||
nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
|
||||
if dependency_cache_ref_exists "${nixl_ref}"; then
|
||||
echo "NIXL dependency cache exists: ${nixl_ref}"
|
||||
else
|
||||
echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
|
||||
add_dependency_cache_target "rixl-rocm-ci"
|
||||
echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
|
||||
add_dependency_cache_target "nixl-rocm-ci"
|
||||
fi
|
||||
fi
|
||||
|
||||
|
||||
+53
-31
@@ -339,18 +339,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# RIXL/UCX build stages
|
||||
FROM base AS build_rixl
|
||||
ARG RIXL_BRANCH="39be1de8"
|
||||
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
|
||||
ARG UCX_BRANCH="bfb51733"
|
||||
# NIXL/UCX build stages
|
||||
FROM base AS build_nixl
|
||||
ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
|
||||
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
|
||||
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
|
||||
ARG UCX_REPO="https://github.com/openucx/ucx.git"
|
||||
ENV ROCM_PATH=/opt/rocm
|
||||
ENV UCX_HOME=/usr/local/ucx
|
||||
ENV RIXL_HOME=/usr/local/rixl
|
||||
ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
|
||||
ENV NIXL_HOME=/usr/local/nixl
|
||||
|
||||
# RIXL build system dependences and RDMA support
|
||||
# NIXL build system dependencies and RDMA support
|
||||
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
|
||||
libgrpc-dev \
|
||||
libgrpc++-dev \
|
||||
@@ -368,7 +367,8 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system meson auditwheel patchelf tomlkit
|
||||
uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
|
||||
auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
cd /usr/local/src && \
|
||||
@@ -396,30 +396,50 @@ ENV PATH=/usr/local/ucx/bin:$PATH
|
||||
ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
git clone ${RIXL_REPO} /opt/rixl && \
|
||||
cd /opt/rixl && \
|
||||
git checkout ${RIXL_BRANCH} && \
|
||||
git clone ${NIXL_REPO} /opt/nixl && \
|
||||
cd /opt/nixl && \
|
||||
git checkout ${NIXL_BRANCH} && \
|
||||
CC="ccache gcc" CXX="ccache g++" \
|
||||
meson setup build --prefix=${RIXL_HOME} \
|
||||
meson setup build --prefix=${NIXL_HOME} \
|
||||
-Ducx_path=${UCX_HOME} \
|
||||
-Drocm_path=${ROCM_PATH} && \
|
||||
-Dwheel_variant=rocm \
|
||||
-Dbuild_tests=false \
|
||||
-Dbuild_examples=false && \
|
||||
cd build && \
|
||||
ninja -j$(nproc) && \
|
||||
ninja install
|
||||
ninja install && \
|
||||
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
|
||||
> /etc/ld.so.conf.d/nixl.conf && \
|
||||
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
|
||||
>> /etc/ld.so.conf.d/nixl.conf && \
|
||||
ldconfig
|
||||
|
||||
# Generate RIXL wheel
|
||||
# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
|
||||
# so configure the ROCm wheel variant directly through Meson.
|
||||
# Exclude libcore and libpull from auditwheel: transitive dependencies
|
||||
# that are not shipped in the wheel and vary across base images.
|
||||
RUN cd /opt/rixl && \
|
||||
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
|
||||
contrib/build-wheel.sh && \
|
||||
mkdir -p /app/install && \
|
||||
_ucx_install_dir=${UCX_HOME} \
|
||||
./contrib/build-wheel.sh \
|
||||
--output-dir /app/install \
|
||||
--rocm-dir ${ROCM_PATH} \
|
||||
RUN cd /opt/nixl && \
|
||||
./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
|
||||
CC="ccache gcc" CXX="ccache g++" \
|
||||
uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
|
||||
--python ${PYTHON_VERSION} \
|
||||
-Csetup-args=-Ducx_path=${UCX_HOME} \
|
||||
-Csetup-args=-Dwheel_variant=rocm \
|
||||
-Csetup-args=-Dbuild_tests=false \
|
||||
-Csetup-args=-Dbuild_examples=false && \
|
||||
mkdir -p /tmp/nixl_wheels/repaired /app/install && \
|
||||
auditwheel repair \
|
||||
--exclude 'libamdhip64*' \
|
||||
--exclude 'libcore*' \
|
||||
--exclude 'libpull*' \
|
||||
/tmp/nixl_wheels/nixl_rocm*.whl \
|
||||
--plat manylinux_2_34_$(uname -m) \
|
||||
--wheel-dir /tmp/nixl_wheels/repaired && \
|
||||
./contrib/wheel_add_ucx_plugins.py \
|
||||
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
|
||||
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
|
||||
--nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
|
||||
/tmp/nixl_wheels/repaired/*.whl && \
|
||||
cp /tmp/nixl_wheels/repaired/*.whl /app/install
|
||||
|
||||
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
|
||||
# invalidate the slow ROCShmem build.
|
||||
@@ -660,10 +680,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
|
||||
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
|
||||
fi
|
||||
|
||||
# Install RIXL + DeepEP wheels.
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
# Install NIXL + DeepEP wheels.
|
||||
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
|
||||
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
|
||||
uv pip install --system /rixl_install/*.whl /deep_install/*.whl
|
||||
uv pip install --system /nixl_install/*.whl /deep_install/*.whl
|
||||
|
||||
# Copy ROCShmem runtime libraries.
|
||||
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
|
||||
@@ -724,6 +744,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
|
||||
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
|
||||
# ROCm profiler limits workaround.
|
||||
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
|
||||
@@ -796,9 +817,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
|
||||
&& pip uninstall -y vllm \
|
||||
&& uv pip install --system *.whl
|
||||
|
||||
# Install RIXL wheel
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
uv pip install --system /rixl_install/*.whl
|
||||
# Install NIXL ROCm wheel
|
||||
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
|
||||
uv pip install --system /nixl_install/*.whl
|
||||
|
||||
ARG COMMON_WORKDIR
|
||||
ARG BASE_IMAGE
|
||||
@@ -813,6 +834,7 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
|
||||
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
|
||||
ENV TOKENIZERS_PARALLELISM=false
|
||||
|
||||
|
||||
+13
-13
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
|
||||
}
|
||||
|
||||
# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
|
||||
# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
|
||||
# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
|
||||
# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
|
||||
# stage and is irrelevant when building --target ci_base itself.
|
||||
variable "CI_BASE_IMAGE" {
|
||||
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
|
||||
# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
|
||||
# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
|
||||
# functions produce no entries when the variable is empty.
|
||||
variable "RIXL_BRANCH" {
|
||||
variable "NIXL_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
variable "RIXL_CACHE_KEY" {
|
||||
variable "NIXL_CACHE_KEY" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
|
||||
])
|
||||
}
|
||||
|
||||
# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
|
||||
# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
|
||||
# These stages are pinned to specific upstream commit hashes, so cache keys use
|
||||
# those hashes rather than the Buildkite commit. This means the cache persists
|
||||
# across all vLLM commits as long as the upstream dependency pins don't change.
|
||||
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
|
||||
function "get_cache_from_rocm_deps" {
|
||||
params = []
|
||||
result = compact([
|
||||
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}" : ""),
|
||||
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}" : ""),
|
||||
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}" : ""),
|
||||
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}" : ""),
|
||||
])
|
||||
}
|
||||
|
||||
function "get_cache_to_rocm_rixl" {
|
||||
function "get_cache_to_rocm_nixl" {
|
||||
params = []
|
||||
result = compact([
|
||||
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" : ""),
|
||||
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" : ""),
|
||||
])
|
||||
}
|
||||
|
||||
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
|
||||
# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
|
||||
# (e.g., requirements change), these stages are cache hits if their upstream
|
||||
# pins haven't changed -- saving ~35min of compilation.
|
||||
target "rixl-rocm-ci" {
|
||||
target "nixl-rocm-ci" {
|
||||
inherits = ["_common-rocm", "_ci-rocm"]
|
||||
target = "build_rixl"
|
||||
target = "build_nixl"
|
||||
cache-from = get_cache_from_rocm_deps()
|
||||
cache-to = get_cache_to_rocm_rixl()
|
||||
cache-to = get_cache_to_rocm_nixl()
|
||||
output = ["type=cacheonly"]
|
||||
}
|
||||
|
||||
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
|
||||
output = ["type=cacheonly"]
|
||||
}
|
||||
|
||||
# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
|
||||
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
|
||||
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
|
||||
# files drifts from the remote image label. Per-PR builds then pull the result
|
||||
# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
|
||||
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
|
||||
CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
|
||||
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
|
||||
]),
|
||||
# Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
|
||||
# Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
|
||||
# are cache hits even when ci_base itself needs rebuilding.
|
||||
get_cache_from_rocm_deps(),
|
||||
)
|
||||
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
|
||||
# Group for ci_base builds -- exports dependency stage caches alongside the
|
||||
# ci_base image so future rebuilds can reuse them independently.
|
||||
group "ci-base-rocm-ci-with-deps" {
|
||||
targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
|
||||
targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
|
||||
}
|
||||
|
||||
@@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" {
|
||||
# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
|
||||
# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
|
||||
# overrides) and writes a small HCL override before invoking CI targets.
|
||||
variable "RIXL_BRANCH" {
|
||||
variable "NIXL_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -106,7 +106,7 @@ target "test-rocm" {
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
|
||||
# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
|
||||
# torchcodec, requirements, etc.). Used by the weekly scheduled build and
|
||||
# the auto-rebuild trigger when requirements change in a PR.
|
||||
target "ci-base-rocm" {
|
||||
|
||||
@@ -13,11 +13,7 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat
|
||||
- Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
|
||||
- The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files
|
||||
|
||||
For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already.
|
||||
|
||||
- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
|
||||
- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
|
||||
- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages
|
||||
For ROCm, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source.
|
||||
|
||||
For non-cuda platform, please install nixl with ucx build from source, instructed as below.
|
||||
|
||||
|
||||
@@ -204,13 +204,12 @@ class FakeNixlWrapper:
|
||||
def _make_fake_nixl_pkg():
|
||||
"""Context manager that creates a temporary package making
|
||||
`from nixl._api import nixl_agent` resolve to our FakeNixlWrapper.
|
||||
Also creates rixl package for ROCm compatibility.
|
||||
Also creates the ROCm NIXL packages.
|
||||
|
||||
Automatically cleans up the temporary directory when done.
|
||||
"""
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
# Create both nixl and rixl packages for cross-platform compatibility
|
||||
for pkg_name in ["nixl", "rixl"]:
|
||||
for pkg_name in ["nixl", "nixl_rocm"]:
|
||||
pkg_root = os.path.join(td, pkg_name, "_api")
|
||||
os.makedirs(pkg_root, exist_ok=True)
|
||||
|
||||
@@ -3130,10 +3129,12 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
|
||||
)
|
||||
|
||||
|
||||
@patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
|
||||
FakeNixlWrapper,
|
||||
)
|
||||
def test_kv_both_deprecation_warning(default_vllm_config, dist_init):
|
||||
"""kv_role='kv_both' should emit a deprecation log warning."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.logger import _print_warning_once
|
||||
|
||||
_print_warning_once.cache_clear()
|
||||
@@ -3156,10 +3157,12 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init):
|
||||
assert "deprecated" in msg
|
||||
|
||||
|
||||
@patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
|
||||
FakeNixlWrapper,
|
||||
)
|
||||
def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init):
|
||||
"""kv_role='kv_consumer' or 'kv_producer' should NOT emit a warning."""
|
||||
from unittest.mock import patch
|
||||
|
||||
for role in ("kv_consumer", "kv_producer"):
|
||||
vllm_config = create_vllm_config(kv_role=role)
|
||||
with patch(
|
||||
|
||||
+7
-7
@@ -1,6 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Verify that GPU memory is fully released after RixlConnector shutdown on ROCm.
|
||||
"""Verify that GPU memory is released after NixlConnector shutdown on ROCm.
|
||||
|
||||
Regression test for ROCm/ucx#33: UCX rocm_ipc transport permanently pinned
|
||||
GPU memory via hsa_amd_ipc_memory_create during ucp_mem_map, causing
|
||||
@@ -62,7 +62,7 @@ def _full_gpu_cleanup():
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
|
||||
def test_gpu_memory_rixl_hma(model_name, sw_size):
|
||||
def test_gpu_memory_nixl_hma(model_name, sw_size):
|
||||
"""Track GPU memory through NixlConnector create/infer/shutdown cycle."""
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import KVTransferConfig
|
||||
@@ -84,7 +84,7 @@ def test_gpu_memory_rixl_hma(model_name, sw_size):
|
||||
}
|
||||
|
||||
print("\n" + "=" * 90)
|
||||
print("GPU MEMORY -- RIXL NixlConnector HMA (ROCm)")
|
||||
print("GPU MEMORY -- NIXL NixlConnector HMA (ROCm)")
|
||||
print("=" * 90)
|
||||
gc.collect()
|
||||
torch.accelerator.empty_cache()
|
||||
@@ -169,14 +169,14 @@ def test_gpu_memory_rixl_hma(model_name, sw_size):
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["google/gemma-3-1b-it"])
|
||||
def test_gpu_memory_no_rixl_baseline(model_name):
|
||||
def test_gpu_memory_no_nixl_baseline(model_name):
|
||||
"""Same workload without NixlConnector. Comparing driver-level memory
|
||||
between this and test_gpu_memory_rixl_hma isolates UCX/RIXL impact."""
|
||||
between this and test_gpu_memory_nixl_hma isolates UCX/NIXL impact."""
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
|
||||
|
||||
print("\n" + "=" * 90)
|
||||
print("CONTROL -- same model, no RIXL connector")
|
||||
print("CONTROL -- same model, no NIXL connector")
|
||||
print("=" * 90)
|
||||
gc.collect()
|
||||
torch.accelerator.empty_cache()
|
||||
@@ -209,7 +209,7 @@ def test_gpu_memory_no_rixl_baseline(model_name):
|
||||
drv_base = snap0["drv_used_mb"]
|
||||
drv_leaked = snap_final["drv_used_mb"] - drv_base
|
||||
drv_peak = snap_peak["drv_used_mb"] - drv_base
|
||||
print(f"\n Driver leaked (no rixl): {drv_leaked:.0f} MB")
|
||||
print(f"\n Driver leaked (no NIXL): {drv_leaked:.0f} MB")
|
||||
print("=" * 90)
|
||||
|
||||
leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0
|
||||
@@ -92,7 +92,7 @@ class EPLBConfig:
|
||||
Backend for EPLB expert weight communication:
|
||||
- "torch_nccl": Use torch.distributed on the device process group
|
||||
- "torch_gloo": Use torch.distributed gloo with CPU staging
|
||||
- "nixl": Use NIXL/ RIXL with staged send/recv buffers
|
||||
- "nixl": Use NIXL with staged send/recv buffers
|
||||
- "pynccl": Use PyNccl send/recv
|
||||
- None: Auto-select backend (prefers "nixl", falls back to "torch_gloo")
|
||||
"""
|
||||
|
||||
@@ -38,7 +38,7 @@ logger = init_logger(__name__)
|
||||
|
||||
|
||||
def has_nixl() -> bool:
|
||||
"""Whether the optional NIXL / RIXL package is available."""
|
||||
"""Whether the optional NIXL package is available."""
|
||||
return nixl_utils.NixlWrapper is not None
|
||||
|
||||
|
||||
@@ -266,7 +266,7 @@ class NixlEplbCommunicator(EplbCommunicator):
|
||||
assert expert_buffer, "NixlEplbCommunicator requires non-empty expert_buffer."
|
||||
nixl_wrapper_cls = nixl_utils.NixlWrapper
|
||||
if nixl_wrapper_cls is None:
|
||||
raise RuntimeError("NIXL/ RIXL is unavailable.")
|
||||
raise RuntimeError("NIXL is unavailable.")
|
||||
|
||||
self._cpu_group = cpu_group
|
||||
self._world_size = cpu_group.size()
|
||||
|
||||
@@ -21,7 +21,7 @@ def _maybe_set_ucx_rcache_limit() -> None:
|
||||
if "UCX_RCACHE_MAX_UNRELEASED" in os.environ:
|
||||
return
|
||||
|
||||
if "nixl" in sys.modules or "rixl" in sys.modules:
|
||||
if "nixl" in sys.modules or "nixl_rocm" in sys.modules:
|
||||
logger.warning_once(
|
||||
"NIXL was already imported, we can't reset "
|
||||
"UCX_RCACHE_MAX_UNRELEASED. "
|
||||
@@ -36,8 +36,12 @@ def _maybe_set_ucx_rcache_limit() -> None:
|
||||
os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024"
|
||||
|
||||
|
||||
def _get_nixl_package_name() -> str:
|
||||
return "nixl_rocm" if current_platform.is_rocm() else "nixl"
|
||||
|
||||
|
||||
def _get_nixl_module_name(name: str) -> str:
|
||||
package_name = "rixl" if current_platform.is_rocm() else "nixl"
|
||||
package_name = _get_nixl_package_name()
|
||||
if name == "nixlXferTelemetry":
|
||||
return f"{package_name}._bindings"
|
||||
return f"{package_name}._api"
|
||||
@@ -80,11 +84,11 @@ def __getattr__(name: str) -> Any:
|
||||
|
||||
|
||||
def is_nixl_available() -> bool:
|
||||
"""Lightweight check for nixl/rixl package without importing it."""
|
||||
"""Lightweight check for the platform's NIXL package without importing it."""
|
||||
import importlib.util
|
||||
|
||||
pkg = "rixl" if current_platform.is_rocm() else "nixl"
|
||||
return importlib.util.find_spec(pkg) is not None
|
||||
pkg = _get_nixl_package_name()
|
||||
return pkg in sys.modules or importlib.util.find_spec(pkg) is not None
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
||||
Reference in New Issue
Block a user