From 6bf15aa2b90608f89165a26db9d858177461643e Mon Sep 17 00:00:00 2001 From: Zhaoyi Li <36555117+Lzy17@users.noreply.github.com> Date: Mon, 29 Jun 2026 23:23:09 -0500 Subject: [PATCH] [AMD]: Enable NIXL PD disaggregation for ROCm(1/n) (#28348) Co-authored-by: bingxche --- .github/workflows/pr-test-amd-rocm720.yml | 17 ++ .github/workflows/pr-test-amd.yml | 17 ++ .../workflows/release-docker-amd-nightly.yml | 2 +- .../release-docker-amd-rocm720-nightly.yml | 2 +- docker/rocm.Dockerfile | 49 ++++ .../test_nixl_transfer_engine_e2e.py | 210 ++++++++++++++++++ 6 files changed, 295 insertions(+), 2 deletions(-) create mode 100644 test/registered/amd/disaggregation/test_nixl_transfer_engine_e2e.py diff --git a/.github/workflows/pr-test-amd-rocm720.yml b/.github/workflows/pr-test-amd-rocm720.yml index 82b8df2cd..164f68abf 100644 --- a/.github/workflows/pr-test-amd-rocm720.yml +++ b/.github/workflows/pr-test-amd-rocm720.yml @@ -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 ' diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 333ee7f52..31d233525 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -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 ' diff --git a/.github/workflows/release-docker-amd-nightly.yml b/.github/workflows/release-docker-amd-nightly.yml index 5cd04909e..db3d6ba2b 100644 --- a/.github/workflows/release-docker-amd-nightly.yml +++ b/.github/workflows/release-docker-amd-nightly.yml @@ -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 diff --git a/.github/workflows/release-docker-amd-rocm720-nightly.yml b/.github/workflows/release-docker-amd-rocm720-nightly.yml index cbad676b9..67eabdefd 100644 --- a/.github/workflows/release-docker-amd-rocm720-nightly.yml +++ b/.github/workflows/release-docker-amd-rocm720-nightly.yml @@ -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 diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 70d77d542..795fc8285 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -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. diff --git a/test/registered/amd/disaggregation/test_nixl_transfer_engine_e2e.py b/test/registered/amd/disaggregation/test_nixl_transfer_engine_e2e.py new file mode 100644 index 000000000..b8e36d852 --- /dev/null +++ b/test/registered/amd/disaggregation/test_nixl_transfer_engine_e2e.py @@ -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()