minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD]: Enable NIXL PD disaggregation for ROCm(1/n) (#28348)

Browse Source 
Co-authored-by: bingxche <bingxche@amd.com>
This commit is contained in:
Zhaoyi Li
2026-06-29 21:23:09 -07:00
committed by GitHub
co-authored by bingxche
parent 2068ae7eec
commit 6bf15aa2b9
6 changed files with 295 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test-amd-rocm720.yml
+17
View File
@@ -1085,6 +1085,23 @@ jobs:
- name: Install dependencies
run: bash scripts/ci/amd/amd_ci_install_dependency.sh
- name: Verify NIXL in Container
run: |
docker exec ci_sglang bash -lc '
set -eux
export LD_LIBRARY_PATH=/opt/ucx/lib:/opt/rocm/lib:${LD_LIBRARY_PATH:-}
python -c "import nixl; print(\"nixl OK:\", nixl.__file__)"
test -d /opt/ucx/lib
UCX_INFO=$(command -v ucx_info || true)
if [ -z "$UCX_INFO" ]; then
UCX_INFO=/opt/ucx/bin/ucx_info
fi
test -x "$UCX_INFO"
"$UCX_INFO" -v
"$UCX_INFO" -d | grep -iq rocm
python -c "from nixl._api import nixl_agent, nixl_agent_config; nixl_agent(\"ci-smoke\", nixl_agent_config(backends=[\"UCX\"])); print(\"nixl UCX backend init OK\")"
'
- name: Verify RDMA in Container
run: |
docker exec -u root ci_sglang bash -c '
.github/workflows/pr-test-amd.yml
+17
View File
@@ -1134,6 +1134,23 @@ jobs:
- name: Install dependencies
run: bash scripts/ci/amd/amd_ci_install_dependency.sh
- name: Verify NIXL in Container
run: |
docker exec ci_sglang bash -lc '
set -eux
export LD_LIBRARY_PATH=/opt/ucx/lib:/opt/rocm/lib:${LD_LIBRARY_PATH:-}
python -c "import nixl; print(\"nixl OK:\", nixl.__file__)"
test -d /opt/ucx/lib
UCX_INFO=$(command -v ucx_info || true)
if [ -z "$UCX_INFO" ]; then
UCX_INFO=/opt/ucx/bin/ucx_info
fi
test -x "$UCX_INFO"
"$UCX_INFO" -v
"$UCX_INFO" -d | grep -iq rocm
python -c "from nixl._api import nixl_agent, nixl_agent_config; nixl_agent(\"ci-smoke\", nixl_agent_config(backends=[\"UCX\"])); print(\"nixl UCX backend init OK\")"
'
- name: Verify RDMA in Container
run: |
docker exec -u root ci_sglang bash -c '
.github/workflows/release-docker-amd-nightly.yml
+1 -1
View File
@@ -87,7 +87,7 @@ jobs:
# remove --build-arg NIC_BACKEND=ainic for auto detection nic support in mori
# UBUNTU_MIRROR forces apt over HTTPS to dodge port-80 reachability flakes
# to Canonical's archive.ubuntu.com mirror IPs from the amd-docker-scale runner.
docker build . -f docker/rocm.Dockerfile --build-arg SGL_BRANCH=${{ github.ref_name }} --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} --build-arg ENABLE_MORI=1 --build-arg SETUPTOOLS_SCM_PRETEND_VERSION=${pretend_version} --build-arg UBUNTU_MIRROR=https://archive.ubuntu.com -t rocm/sgl-dev:${tag}-${{ env.DATE }} --no-cache
docker build . -f docker/rocm.Dockerfile --build-arg SGL_BRANCH=${{ github.ref_name }} --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} --build-arg ENABLE_MORI=1 --build-arg ENABLE_NIXL=1 --build-arg SETUPTOOLS_SCM_PRETEND_VERSION=${pretend_version} --build-arg UBUNTU_MIRROR=https://archive.ubuntu.com -t rocm/sgl-dev:${tag}-${{ env.DATE }} --no-cache
docker push rocm/sgl-dev:${tag}-${{ env.DATE }}
# Persist the tag right after rocm/sgl-dev push succeeds so the local
.github/workflows/release-docker-amd-rocm720-nightly.yml
+1 -1
View File
@@ -96,7 +96,7 @@ jobs:
# remove --build-arg NIC_BACKEND=ainic for auto detection nic support in mori
# UBUNTU_MIRROR forces apt over HTTPS to dodge port-80 reachability flakes
# to Canonical's archive.ubuntu.com mirror IPs from the amd-docker-scale runner.
docker build . -f docker/rocm.Dockerfile --build-arg SGL_BRANCH=${{ github.ref_name }} --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} --build-arg ENABLE_MORI=1 --build-arg SETUPTOOLS_SCM_PRETEND_VERSION=${pretend_version} --build-arg UBUNTU_MIRROR=https://archive.ubuntu.com -t rocm/sgl-dev:${tag}-${{ env.DATE }} --no-cache
docker build . -f docker/rocm.Dockerfile --build-arg SGL_BRANCH=${{ github.ref_name }} --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} --build-arg ENABLE_MORI=1 --build-arg ENABLE_NIXL=1 --build-arg SETUPTOOLS_SCM_PRETEND_VERSION=${pretend_version} --build-arg UBUNTU_MIRROR=https://archive.ubuntu.com -t rocm/sgl-dev:${tag}-${{ env.DATE }} --no-cache
docker push rocm/sgl-dev:${tag}-${{ env.DATE }}
# Persist the tag right after rocm/sgl-dev push succeeds so the local
docker/rocm.Dockerfile
+49
View File
@@ -14,6 +14,12 @@
# docker build --build-arg SGL_BRANCH=v0.5.10.post1 --build-arg GPU_ARCH=gfx950 --build-arg ENABLE_MORI=1 -t v0.5.10.post1-rocm700-mi35x -f rocm.Dockerfile .
# docker build --build-arg SGL_BRANCH=v0.5.10.post1 --build-arg GPU_ARCH=gfx950-rocm720 --build-arg ENABLE_MORI=1 -t v0.5.10.post1-rocm720-mi35x -f rocm.Dockerfile .
# Usage (to build SGLang ROCm + NIXL docker image, for prefill/decode disaggregation):
# Builds UCX (--with-rocm) and upstream ai-dynamo/nixl from source by default.
# Set ENABLE_NIXL=0 to skip NIXL.
# At runtime use --disaggregation-transfer-backend nixl (env is wired via /etc/bash.bashrc).
# docker build --build-arg SGL_BRANCH=v0.5.10.post1 --build-arg GPU_ARCH=gfx950-rocm720 -t v0.5.10.post1-rocm720-mi35x -f rocm.Dockerfile .
# Default base images
ARG BASE_IMAGE_942="rocm/sgl-dev:rocm7-vllm-20250904"
ARG BASE_IMAGE_942_ROCM720="rocm/pytorch:rocm7.2_ubuntu22.04_py3.10_pytorch_release_2.9.1"
@@ -106,6 +112,15 @@ ARG NIC_BACKEND=none
ARG MORI_REPO="https://github.com/ROCm/mori.git"
ARG MORI_COMMIT="bf99bdf18fc69887a346913ca01c315c2aa9bd4c"
# NIXL (upstream ai-dynamo/nixl) — KV transfer backend for prefill/decode disaggregation.
# Built from source for ROCm; needs UCX built --with-rocm (built here from openucx).
# Enabled by default; disable with --build-arg ENABLE_NIXL=0.
ARG ENABLE_NIXL=1
ARG UCX_REPO="https://github.com/openucx/ucx.git"
ARG UCX_BRANCH="v1.19.x"
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
ARG NIXL_COMMIT="c28061f9782e099f975bcc79198b7b5a1a36cc40"
# AMD AINIC apt repo settings
ARG AINIC_VERSION=1.117.5-a-38
ARG UBUNTU_CODENAME=jammy
@@ -488,6 +503,40 @@ RUN /bin/bash -lc 'set -euo pipefail; \
echo "export PYTHONPATH=/sgl-workspace/mori:\${PYTHONPATH}" >> /etc/bash.bashrc; \
echo "[MORI] Done."'
# -----------------------
# NIXL — upstream ai-dynamo/nixl KV transfer backend for PD disaggregation on ROCm.
# Builds UCX (--with-rocm) + nixl from source by default; skip with ENABLE_NIXL=0.
# --no-build-isolation reuses the image's ROCm torch (nixl pins torch==2.11.* as a build dep,
# which would otherwise pull a multi-GB CUDA torch); --no-deps keeps CUDA runtime deps out.
# wheel_variant=rocm names the pkg nixl_rocm, so symlink `nixl` since SGLang imports plain nixl.
# taskflow (header-only) is provided via pkg-config so meson skips its broken upstream wrap
# download (GitHub regenerated the v3.10.0 tarball, breaking the pinned source_hash).
RUN /bin/bash -lc 'set -euo pipefail; \
[ "${ENABLE_NIXL}" = "1" ] || { echo "[NIXL] skip (ENABLE_NIXL=${ENABLE_NIXL})"; exit 0; }; \
apt-get update && apt-get install -y --no-install-recommends \
build-essential autoconf automake libtool pkg-config git \
libibverbs-dev librdmacm-dev rdma-core && rm -rf /var/lib/apt/lists/*; \
pip install --no-cache-dir meson ninja pybind11 meson-python patchelf pyyaml; \
git clone --depth=1 -b "${UCX_BRANCH}" "${UCX_REPO}" /sgl-workspace/ucx; \
cd /sgl-workspace/ucx && ./autogen.sh && mkdir build && cd build && \
../configure --prefix=/opt/ucx --enable-shared --disable-static --disable-doxygen-doc \
--enable-optimizations --enable-devel-headers \
--with-rocm=/opt/rocm --with-verbs --with-dm --enable-mt && \
make -j"$(nproc)" && make install; \
git clone --depth=1 -b v3.10.0 https://github.com/taskflow/taskflow.git /sgl-workspace/taskflow; \
cp -r /sgl-workspace/taskflow/taskflow /usr/local/include/; \
mkdir -p /usr/local/lib/pkgconfig; \
printf "Name: taskflow\nDescription: Taskflow\nVersion: 3.10.0\nCflags: -I/usr/local/include\n" > /usr/local/lib/pkgconfig/taskflow.pc; \
git clone "${NIXL_REPO}" /sgl-workspace/nixl && cd /sgl-workspace/nixl && git checkout -f "${NIXL_COMMIT}"; \
CXXFLAGS="-Wno-error" LD_LIBRARY_PATH="/opt/ucx/lib:/opt/rocm/lib" PKG_CONFIG_PATH="/usr/local/lib/pkgconfig" \
pip install . --no-deps --no-build-isolation \
--config-settings=setup-args="-Ducx_path=/opt/ucx" \
--config-settings=setup-args="-Dwheel_variant=rocm" \
--config-settings=setup-args="-Denable_plugins=UCX,POSIX"; \
SITE=$(python3 -c "import sysconfig; print(sysconfig.get_paths()[\"purelib\"])"); \
ln -sfn nixl_rocm "$SITE/nixl"; \
echo "export LD_LIBRARY_PATH=/opt/ucx/lib:\${LD_LIBRARY_PATH}" >> /etc/bash.bashrc'
# -----------------------
# Hot patch: torch-ROCm
# The artifact hardcoded the supported triton version to be 3.5.1.
test/registered/amd/disaggregation/test_nixl_transfer_engine_e2e.py
+210
View File
@@ -0,0 +1,210 @@
import os
import unittest
from types import SimpleNamespace
import requests
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
try_cached_model,
)
register_amd_ci(est_time=900, suite="stage-b-test-large-8-gpu-mi35x-disaggregation-amd")
class NixlTransferEngineBase(PDDisaggregationServerBase):
"""PD-over-NIXL e2e on ROCm. NIXL (upstream ai-dynamo/nixl + UCX --with-rocm)
is enabled by default in the ROCm image; when the image was built with
`--build-arg ENABLE_NIXL=0`, `import nixl` fails and the test skips rather
than failing the suite."""
port_delta = 0
prefill_tp = 1
decode_tp = 1
decode_base_gpu_id = 1
required_gpus = 2
model_default = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
model_env_var = "SGLANG_NIXL_E2E_TEST_MODEL"
extra_prefill_args: list = []
extra_decode_args: list = []
@classmethod
def setUpClass(cls):
try:
import torch
if not torch.cuda.is_available():
raise unittest.SkipTest("torch.cuda is not available.")
if torch.cuda.device_count() < cls.required_gpus:
raise unittest.SkipTest(
f"NIXL PD smoke test requires >= {cls.required_gpus} visible GPUs."
)
except unittest.SkipTest:
raise
except Exception as e:
raise unittest.SkipTest(f"torch is not available/usable: {e}")
try:
import nixl # noqa: F401
except Exception as e:
raise unittest.SkipTest(
"nixl not importable; image may have been built with "
"--build-arg ENABLE_NIXL=0 "
f"({e})."
)
super().setUpClass()
cls._old_use_aiter = os.environ.get("SGLANG_USE_AITER")
os.environ["SGLANG_USE_AITER"] = "1"
# The shared fixture defaults to Mooncake in CI; pin NIXL explicitly here.
cls.transfer_backend = ["--disaggregation-transfer-backend", "nixl"]
rdma_env = os.environ.get("SGLANG_TEST_RDMA_DEVICE")
if rdma_env:
cls.rdma_devices = ["--disaggregation-ib-device", rdma_env]
print(f"Found RDMA devices in env: {rdma_env}")
else:
print("SGLANG_TEST_RDMA_DEVICE is not set! Running without RDMA.")
cls.rdma_devices = []
cls._shift_ports()
cls.model = try_cached_model(
os.environ.get(cls.model_env_var, cls.model_default)
)
cls.start_prefill()
cls.start_decode()
cls.wait_server_ready(
cls.prefill_url + "/health",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
process=cls.process_prefill,
)
cls.wait_server_ready(
cls.decode_url + "/health",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
process=cls.process_decode,
)
cls.launch_lb()
@classmethod
def tearDownClass(cls):
if getattr(cls, "_old_use_aiter", None) is None:
os.environ.pop("SGLANG_USE_AITER", None)
else:
os.environ["SGLANG_USE_AITER"] = cls._old_use_aiter
super().tearDownClass()
@classmethod
def _shift_ports(cls):
if cls.port_delta == 0:
return
cls.lb_port = str(int(cls.lb_port) + cls.port_delta)
cls.prefill_port = str(int(cls.prefill_port) + cls.port_delta)
cls.decode_port = str(int(cls.decode_port) + cls.port_delta)
cls.bootstrap_port = str(int(cls.bootstrap_port) + cls.port_delta)
cls.prefill_url = f"http://{cls.base_host}:{cls.prefill_port}"
cls.decode_url = f"http://{cls.base_host}:{cls.decode_port}"
cls.lb_url = f"http://{cls.base_host}:{cls.lb_port}"
cls.base_url = cls.lb_url
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
str(cls.prefill_tp),
"--attention-backend",
"aiter",
] + list(cls.extra_prefill_args)
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
str(cls.decode_tp),
"--base-gpu-id",
str(cls.decode_base_gpu_id),
"--attention-backend",
"aiter",
"--mem-fraction-static",
"0.8",
] + list(cls.extra_decode_args)
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def _assert_generate_smoke(self):
resp = requests.post(
self.lb_url + "/generate",
json={
"text": "Hello",
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
},
timeout=120,
)
self.assertEqual(resp.status_code, 200, resp.text)
out = resp.json()
self.assertIn("text", out)
self.assertIsInstance(out["text"], str)
self.assertGreater(len(out["text"]), 0)
class TestNixlTransferEngineE2E(NixlTransferEngineBase):
def test_generate_smoke(self):
self._assert_generate_smoke()
class TestNixlTransferEngineAccuracy(NixlTransferEngineBase):
port_delta = 10
model_default = DEFAULT_MODEL_NAME_FOR_TEST
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.70)
if __name__ == "__main__":
unittest.main()