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

Co-authored-by: bingxche <bingxche@amd.com>
2026-06-29 21:23:09 -07:00
parent 2068ae7eec
commit 6bf15aa2b9
6 changed files with 295 additions and 2 deletions
@@ -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 '
@@ -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 '
@@ -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
@@ -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
@@ -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.
@@ -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()