AMD PD/D PR ci (#17183)

Co-authored-by: YC Tseng <yctseng@amd.com>
Co-authored-by: Bingxu Chen <bingxche@amd.com>
Co-authored-by: bingxche <Bingxu.Chen@amd.com>

.github/workflows/pr-test-amd.yml

@@ -734,6 +734,118 @@ jobs:
         run: |
           bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 1 --timeout-per-file 3600
 
+  # =============================================== Disaggregation ====================================================
+  stage-b-test-large-8-gpu-35x-disaggregation-amd:
+    needs: [check-changes, stage-a-test-1-amd]
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'stage-b-test-large-8-gpu-disaggregation-amd') ||
+        (
+          !inputs.target_stage &&
+          (!failure() && !cancelled()) &&
+          ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
+    strategy:
+      fail-fast: false
+      matrix:
+        runner: [linux-mi35x-gpu-8.fabric]
+
+    runs-on: ${{matrix.runner}}
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
+
+      - name: Ensure VRAM is clear
+        run: bash scripts/ensure_vram_clear.sh rocm
+
+      - name: Check Host RDMA Environment
+        id: rdma_detect
+        run: |
+          set +e
+          echo "=== Checking Host RDMA Environment ==="
+
+          echo ""
+          echo "=== 1. Ionic driver library check ==="
+          ls -l /usr/lib/x86_64-linux-gnu/libibverbs/libionic* 2>/dev/null || echo "libionic not found in standard path"
+
+          echo ""
+          echo "=== 2. Infiniband devices ==="
+          ls -la /dev/infiniband/ 2>/dev/null || echo "/dev/infiniband not found"
+          ls -la /sys/class/infiniband/ 2>/dev/null || echo "/sys/class/infiniband not found"
+
+          echo ""
+          echo "=== 3. ibv_devinfo ==="
+          which ibv_devinfo 2>/dev/null && ibv_devinfo 2>&1 || echo "ibv_devinfo not available"
+
+          echo ""
+          echo "=== 4. Kernel modules ==="
+          lsmod 2>/dev/null | grep -E "ib_|rdma|ionic" || echo "No RDMA kernel modules loaded"
+
+          echo ""
+          echo "=== 5. Detect RDMA Devices for test environment ==="
+          if [ -d "/sys/class/infiniband" ]; then
+            RDMA_DEVS=$(ls /sys/class/infiniband | paste -sd "," -)
+            echo "Detected RDMA Devices: $RDMA_DEVS"
+            echo "SGLANG_TEST_RDMA_DEVICE=$RDMA_DEVS" >> $GITHUB_ENV
+          else
+            echo "No RDMA devices found in /sys/class/infiniband"
+            echo "SGLANG_TEST_RDMA_DEVICE=" >> $GITHUB_ENV
+          fi
+
+          echo ""
+          echo "=== Host RDMA Check Complete ==="
+
+      - name: Start Special Container
+        run: bash scripts/ci/amd/amd_ci_start_container_disagg.sh
+        env:
+          GITHUB_WORKSPACE: ${{ github.workspace }}
+
+      - name: Install dependencies
+        run: bash scripts/ci/amd/amd_ci_install_dependency.sh
+
+      - name: Verify RDMA in Container
+        run: |
+          docker exec -u root ci_sglang bash -c '
+            echo "=== Container RDMA Verification ==="
+            echo "Device nodes:"
+            ls -la /dev/infiniband/
+            echo ""
+            echo "Provider libraries:"
+            ls /usr/lib/x86_64-linux-gnu/libibverbs/ | grep -E "ionic|mlx" || echo "No Ionic/Mellanox providers"
+            echo ""
+            echo "HCA devices:"
+            HCA_COUNT=$(ibv_devinfo -list 2>&1 | grep -oE "^[0-9]+ HCAs? found" | grep -oE "^[0-9]+" || echo "0")
+            ibv_devinfo -list
+            if [ "$HCA_COUNT" -gt 0 ]; then
+              echo ""
+              echo "=== SUCCESS: RDMA setup complete. Found $HCA_COUNT HCA(s) ==="
+            else
+              echo ""
+              echo "=== WARNING: No HCAs detected. RDMA tests may fail ==="
+            fi
+          '
+
+      - name: Run Aiter Op Test (RMSNorm)
+        timeout-minutes: 10
+        run: |
+          echo "Running pre-check: test_rmsnorm2d.py"
+          docker exec \
+            -e MAX_JOBS=192 \
+            ci_sglang \
+            python /sgl-workspace/aiter/op_tests/test_rmsnorm2d.py
+
+      - name: Run test_disaggregation
+        timeout-minutes: 60
+        run: |
+          bash scripts/ci/amd/amd_ci_exec.sh \
+            -e SGLANG_TEST_RDMA_DEVICE="${{ env.SGLANG_TEST_RDMA_DEVICE }}" \
+            -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-large-8-gpu-35x-disaggregation-amd --timeout-per-file 1800
+
   pr-test-amd-finish:
     needs:
       [
@@ -750,6 +862,7 @@ jobs:
         stage-b-test-small-1-gpu-amd-mi35x,
         stage-b-test-large-1-gpu-amd,
         stage-b-test-large-2-gpu-amd,
+        stage-b-test-large-8-gpu-35x-disaggregation-amd,
         stage-c-test-large-8-gpu-amd,
         stage-c-test-large-8-gpu-amd-mi35x,
       ]

scripts/ci/amd/amd_ci_start_container_disagg.sh

@@ -0,0 +1,244 @@
+#!/bin/bash
+set -euo pipefail
+
+# Get version from git tags
+SGLANG_VERSION="v0.5.5"   # Default version, will be overridden if git tags are found
+
+# Fetch tags from origin to ensure we have the latest
+if git fetch --tags origin; then
+  # Get the latest version tag sorted by version number (e.g., v0.5.7)
+  VERSION_FROM_TAG=$(git tag -l 'v[0-9]*' --sort=-v:refname | head -1)
+  if [ -n "$VERSION_FROM_TAG" ]; then
+    SGLANG_VERSION="$VERSION_FROM_TAG"
+    echo "Using SGLang version from git tags: $SGLANG_VERSION"
+  else
+    echo "Warning: No version tags found; using default $SGLANG_VERSION" >&2
+  fi
+else
+  echo "Warning: Failed to fetch tags from origin; using default $SGLANG_VERSION" >&2
+fi
+
+
+# Default base tags (can be overridden by command line arguments)
+ROCM_VERSION="rocm700"
+DEFAULT_MI30X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi30x"
+DEFAULT_MI35X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi35x"
+
+# Parse command line arguments
+MI30X_BASE_TAG="${DEFAULT_MI30X_BASE_TAG}"
+MI35X_BASE_TAG="${DEFAULT_MI35X_BASE_TAG}"
+
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --mi30x-base-tag) MI30X_BASE_TAG="$2"; shift 2;;
+    --mi35x-base-tag) MI35X_BASE_TAG="$2"; shift 2;;
+    -h|--help)
+      echo "Usage: $0 [--mi30x-base-tag TAG] [--mi35x-base-tag TAG]"
+      exit 0
+      ;;
+    *) echo "Unknown option $1"; exit 1;;
+  esac
+done
+
+
+
+# Detect GPU architecture from the Kubernetes runner hostname
+HOSTNAME_VALUE=$(hostname)
+GPU_ARCH="mi30x"   # default
+
+# Host names look like: linux-mi35x-gpu-1-xxxxx-runner-zzzzz
+if [[ "${HOSTNAME_VALUE}" =~ ^linux-(mi[0-9]+[a-z]*)-gpu-[0-9]+ ]]; then
+  GPU_ARCH="${BASH_REMATCH[1]}"
+  echo "Detected GPU architecture from hostname: ${GPU_ARCH}"
+else
+  echo "Warning: could not parse GPU architecture from '${HOSTNAME_VALUE}', defaulting to ${GPU_ARCH}"
+fi
+
+# Normalise / collapse architectures we don’t yet build specifically for
+case "${GPU_ARCH}" in
+  mi35x)
+    echo "Runner uses ${GPU_ARCH}; will fetch mi35x image."
+    ;;
+  mi30x|mi300|mi325)
+    echo "Runner uses ${GPU_ARCH}; will fetch mi30x image."
+    GPU_ARCH="mi30x"
+    ;;
+  *)
+    echo "Runner architecture '${GPU_ARCH}' unrecognised; defaulting to mi30x image." >&2
+    GPU_ARCH="mi30x"
+    ;;
+esac
+
+
+# Set up DEVICE_FLAG based on Kubernetes pod info
+if [[ -f /etc/podinfo/gha-render-devices ]]; then
+  DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices)
+else
+  DEVICE_FLAG="--device /dev/dri"
+fi
+
+
+# Find the latest image
+find_latest_image() {
+  local gpu_arch=$1
+  local base_tag days_back image_tag
+
+  case "${gpu_arch}" in
+      mi30x) base_tag="${MI30X_BASE_TAG}" ;;
+      mi35x) base_tag="${MI35X_BASE_TAG}" ;;
+      *)     echo "Error: unsupported GPU architecture '${gpu_arch}'" >&2; return 1 ;;
+  esac
+
+  # First, check local cache
+  for days_back in {0..6}; do
+    image_tag="${base_tag}-$(date -d "${days_back} days ago" +%Y%m%d)"
+    local local_image="rocm/sgl-dev:${image_tag}"
+    image_id=$(docker images -q "${local_image}")
+    if [[ -n "$image_id" ]]; then
+        echo "Found cached image locally: ${local_image}" >&2
+        echo "${local_image}"
+        return 0
+    fi
+  done
+
+  # If not found locally, fall back to pulling from public registry
+  for days_back in {0..6}; do
+    image_tag="${base_tag}-$(date -d "${days_back} days ago" +%Y%m%d)"
+    echo "Checking for image: rocm/sgl-dev:${image_tag}" >&2
+    if docker manifest inspect "rocm/sgl-dev:${image_tag}" >/dev/null 2>&1; then
+      echo "Found available image: rocm/sgl-dev:${image_tag}" >&2
+      echo "rocm/sgl-dev:${image_tag}"
+      return 0
+    fi
+  done
+
+  # If still not found, try finding any image matching ROCm+arch from remote registry
+  echo "Exact version not found. Searching remote registry for any ${ROCM_VERSION}-${gpu_arch} image…" >&2
+  for days_back in {0..6}; do
+    local target_date=$(date -d "${days_back} days ago" +%Y%m%d)
+    local remote_tags=$(curl -s "https://registry.hub.docker.com/v2/repositories/rocm/sgl-dev/tags?page_size=100&name=${ROCM_VERSION}-${gpu_arch}-${target_date}" 2>/dev/null | grep -o '"name":"[^"]*"' | cut -d'"' -f4 | head -n 1)
+    if [[ -n "$remote_tags" ]]; then
+      echo "Found available image: rocm/sgl-dev:${remote_tags}" >&2
+      echo "rocm/sgl-dev:${remote_tags}"
+      return 0
+    fi
+  done
+
+  echo "No recent images found. Searching any cached local images matching ROCm+arch…" >&2
+  local any_local
+  any_local=$(docker images --format '{{.Repository}}:{{.Tag}}' --filter "reference=rocm/sgl-dev:*${ROCM_VERSION}*${gpu_arch}*" | sort -r | head -n 1)
+  if [[ -n "$any_local" ]]; then
+      echo "Using cached fallback image: ${any_local}" >&2
+      echo "${any_local}"
+      return 0
+  fi
+
+  echo "Error: no ${gpu_arch} image found in the last 7 days for base ${base_tag}" >&2
+  echo "Using hard-coded fallback…" >&2
+  if [[ "${gpu_arch}" == "mi35x" ]]; then
+    echo "rocm/sgl-dev:v0.5.5-rocm700-mi35x-20251110"
+  else
+    echo "rocm/sgl-dev:v0.5.5-rocm700-mi30x-20251110"
+  fi
+}
+
+# Pull and run the latest image
+IMAGE=$(find_latest_image "${GPU_ARCH}")
+echo "Pulling Docker image: ${IMAGE}"
+docker pull "${IMAGE}"
+
+CACHE_HOST=/home/runner/sgl-data
+if [[ -d "$CACHE_HOST" ]]; then
+    CACHE_VOLUME="-v $CACHE_HOST:/sgl-data"
+else
+    CACHE_VOLUME=""
+fi
+
+# Detect libionic library for RDMA support
+LIBIONIC_MOUNT=""
+IONIC_SYMLINK="/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so"
+if [[ -L "$IONIC_SYMLINK" ]]; then
+    LIBIONIC_LIB=$(readlink -f "$IONIC_SYMLINK" 2>/dev/null)
+    if [[ -f "$LIBIONIC_LIB" ]]; then
+        echo "Found libionic library: $LIBIONIC_LIB (resolved from symlink)"
+        LIBIONIC_MOUNT="-v ${LIBIONIC_LIB}:${LIBIONIC_LIB}:ro"
+    else
+        echo "Warning: libionic symlink exists but target does not: $LIBIONIC_LIB"
+    fi
+else
+    # Fallback: try to find directly
+    LIBIONIC_FOUND=$(find /usr/lib/x86_64-linux-gnu -maxdepth 1 -name "libionic.so.*" 2>/dev/null | head -1)
+    if [[ -n "$LIBIONIC_FOUND" ]]; then
+        LIBIONIC_LIB=$(readlink -f "$LIBIONIC_FOUND" 2>/dev/null)
+        if [[ -f "$LIBIONIC_LIB" ]]; then
+            echo "Found libionic library: $LIBIONIC_LIB"
+            LIBIONIC_MOUNT="-v ${LIBIONIC_LIB}:${LIBIONIC_LIB}:ro"
+        else
+            echo "Warning: libionic found but cannot resolve real path: $LIBIONIC_FOUND"
+        fi
+    else
+        echo "Warning: libionic library not found on host, RDMA may not work"
+    fi
+fi
+
+MOUNT_ARGS=""
+
+add_mount_if_exists() {
+    local name=$1
+    local search_pattern=$2
+    local path=$(find /lib/x86_64-linux-gnu /usr/lib/x86_64-linux-gnu /lib64 /usr/lib64 -name "$search_pattern" -print -quit 2>/dev/null)
+
+    if [ -n "$path" ]; then
+        echo "Found $name at: $path"
+        MOUNT_ARGS="$MOUNT_ARGS -v $path:$path:ro"
+    else
+        echo "WARNING: Could not find $name on host! (Pattern: $search_pattern)"
+    fi
+}
+
+IONIC_LINK="/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so"
+if [ -L "$IONIC_LINK" ]; then
+    IONIC_REAL=$(readlink -f "$IONIC_LINK")
+    if [ -f "$IONIC_REAL" ]; then
+        echo "Ionic Driver: $IONIC_REAL"
+        MOUNT_ARGS="$MOUNT_ARGS -v $IONIC_REAL:$IONIC_REAL:ro"
+    fi
+fi
+
+add_mount_if_exists "libnl-3" "libnl-3.so*"
+add_mount_if_exists "libmnl" "libmnl.so*"
+
+echo "Mount args: $MOUNT_ARGS"
+
+echo "Launching container: ci_sglang"
+docker run -dt --user root \
+  --device=/dev/kfd \
+  --device=/dev/dri \
+  ${DEVICE_FLAG} \
+  -v "${GITHUB_WORKSPACE:-$PWD}:/sglang-checkout" \
+  -v /sys/class/infiniband:/sys/class/infiniband:ro \
+  -v /sys/class/infiniband_verbs:/sys/class/infiniband_verbs:ro \
+  -v /sys/class/net:/sys/class/net:ro \
+  -v /etc/libibverbs.d:/etc/libibverbs.d:ro \
+  -v /usr/lib/x86_64-linux-gnu/libibverbs:/usr/lib/x86_64-linux-gnu/libibverbs:ro \
+  $MOUNT_ARGS \
+  $CACHE_VOLUME \
+  --privileged \
+  --network=host \
+  --ipc=host \
+  --ulimit memlock=-1 \
+  --cap-add=IPC_LOCK \
+  --cap-add=SYS_PTRACE \
+  --security-opt seccomp=unconfined \
+  --group-add video \
+  --group-add rdma \
+  --shm-size 32g \
+  -e HF_TOKEN="${HF_TOKEN:-}" \
+  -e HF_HOME=/sgl-data/hf-cache \
+  -e HF_HUB_ETAG_TIMEOUT=300 \
+  -e HF_HUB_DOWNLOAD_TIMEOUT=300 \
+  -e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \
+  -e MIOPEN_CUSTOM_CACHE_DIR=/sgl-data/miopen-cache \
+  -w /sglang-checkout \
+  --name ci_sglang \
+  "${IMAGE}"

test/registered/amd/disaggregation/test_disaggregation_basic.py

@@ -0,0 +1,422 @@
+import json
+import os
+import unittest
+from types import SimpleNamespace
+
+import openai
+import requests
+from transformers import AutoTokenizer
+
+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_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_pd_server,
+)
+
+register_amd_ci(est_time=600, suite="stage-b-test-large-8-gpu-35x-disaggregation-amd")
+
+
+class TestDisaggregationAccuracy(PDDisaggregationServerBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        # Configure ROCm RDMA environment
+        os.environ["SGLANG_USE_AITER"] = "1"
+        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.model = DEFAULT_MODEL_NAME_FOR_TEST
+        # DEFAULT_MODEL_NAME_FOR_TEST
+
+        # Non blocking start servers
+        cls.start_prefill()
+        cls.start_decode()
+
+        # Block until both
+        cls.wait_server_ready(cls.prefill_url + "/health")
+        cls.wait_server_ready(cls.decode_url + "/health")
+
+        cls.launch_lb()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp",
+            "1",
+            "--attention-backend",
+            "aiter",
+            "--log-level",
+            "debug",
+        ]
+        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",
+            "--tp",
+            "1",
+            "--base-gpu-id",
+            "1",
+            "--attention-backend",
+            "aiter",
+            "--mem-fraction-static",
+            "0.8",
+            "--log-level",
+            "debug",
+        ]
+        decode_args += cls.transfer_backend + cls.rdma_devices
+        print("Debug")
+        print(decode_args)
+        cls.process_decode = popen_launch_pd_server(
+            cls.model,
+            cls.decode_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=decode_args,
+        )
+
+    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)
+
+    def test_logprob(self):
+        prompt = "The capital of france is "
+        response = requests.post(
+            self.lb_url + "/generate",
+            json={
+                "text": prompt,
+                "sampling_params": {"temperature": 0},
+                "return_logprob": True,
+                "return_input_logprob": True,
+                "logprob_start_len": 0,
+            },
+        )
+
+        j = response.json()
+        completion_tokens = j["meta_info"]["completion_tokens"]
+        input_logprobs = j["meta_info"]["input_token_logprobs"]
+        output_logprobs = j["meta_info"]["output_token_logprobs"]
+
+        assert (
+            len(output_logprobs) == completion_tokens
+        ), f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}"
+        assert (
+            len(input_logprobs) > 0
+        ), f"input_logprobs should have at least one token, but got {len(input_logprobs)}"
+
+    def test_structured_output(self):
+        json_schema = json.dumps(
+            {
+                "type": "object",
+                "properties": {
+                    "name": {"type": "string", "pattern": "^[\\w]+$"},
+                    "population": {"type": "integer"},
+                },
+                "required": ["name", "population"],
+            }
+        )
+
+        # JSON
+        response = requests.post(
+            f"{self.lb_url}/generate",
+            json={
+                "text": "Here is the information of the capital of France in the JSON format.\n",
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 64,
+                    "json_schema": json_schema,
+                },
+            },
+        )
+        output = response.json()["text"]
+        # ensure the output is a valid JSON
+        json.loads(output)
+
+    def test_first_token_finish(self):
+        client = openai.Client(api_key="empty", base_url=f"{self.lb_url}/v1")
+        tokenizer = AutoTokenizer.from_pretrained(self.model)
+        eos_token = tokenizer.eos_token_id
+        prompt = "The best programming language for AI is"
+
+        # First token EOS
+        res = client.completions.create(
+            model="dummy", prompt=prompt, logit_bias={eos_token: 42}
+        ).model_dump()
+        print(f"{res=}")
+
+        assert res["usage"]["completion_tokens"] == 1, (
+            "Expected completion_tokens to be 1 when first token is EOS, "
+            f"but got {res['usage']['completion_tokens']}"
+        )
+
+        # First token EOS with ignore_eos
+        res = client.completions.create(
+            model="dummy",
+            prompt=prompt,
+            logit_bias={eos_token: 42},
+            extra_body={"ignore_eos": True},
+        ).model_dump()
+        print(f"{res=}")
+
+        assert res["usage"]["completion_tokens"] > 1, (
+            "Expected completion_tokens to be greater than 1 when ignore_eos is True, "
+            f"but got {res['usage']['completion_tokens']}"
+        )
+
+        # First token with specified stop token
+        stop_token_id = tokenizer.encode(" hello", add_special_tokens=False)[0]
+        res = client.completions.create(
+            model="dummy",
+            prompt=prompt,
+            logit_bias={stop_token_id: 42},
+            stop=[" hello"],
+        ).model_dump()
+        print(f"{res=}")
+
+        assert res["usage"]["completion_tokens"] == 1, (
+            "Expected completion_tokens to be 1 when first token is stop token, "
+            f"but got {res['usage']['completion_tokens']}"
+        )
+
+
+# register_amd_ci(est_time=300, suite="stage-b-test-large-2-gpu-amd")
+class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        # Configure ROCm RDMA environment
+        os.environ["SGLANG_USE_AITER"] = "1"
+        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 = []
+
+        # set DISAGGREGATION_TEST_FAILURE_PROB to simulate failure
+        os.environ["DISAGGREGATION_TEST_FAILURE_PROB"] = "0.05"
+
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+
+        # Non blocking start servers
+        cls.start_prefill()
+        cls.start_decode()
+
+        # Block until both
+        cls.wait_server_ready(cls.prefill_url + "/health")
+        cls.wait_server_ready(cls.decode_url + "/health")
+
+        cls.launch_lb()
+
+    @classmethod
+    def tearDownClass(cls):
+        os.environ.pop("DISAGGREGATION_TEST_FAILURE_PROB")
+        super().tearDownClass()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp",
+            "1",
+            "--attention-backend",
+            "aiter",
+            "--log-level",
+            "debug",
+        ]
+        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",
+            "--tp",
+            "1",
+            "--base-gpu-id",
+            "1",
+            "--attention-backend",
+            "aiter",
+            "--mem-fraction-static",
+            "0.8",
+            "--log-level",
+            "debug",
+        ]
+        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 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),
+        )
+
+        # Expect lots of failure but the server cannot crash
+        try:
+            metrics = run_eval_few_shot_gsm8k(args)
+            print(f"Evaluation metrics: {metrics}")
+        except Exception as e:
+            print(f"Test encountered expected errors: {e}")
+            # Check if servers are still healthy
+            try:
+                response = requests.get(self.prefill_url + "/health_generate")
+                assert response.status_code == 200
+                response = requests.get(self.decode_url + "/health_generate")
+                assert response.status_code == 200
+            except Exception as health_check_error:
+                # If health check fails, re-raise the original exception
+                raise e from health_check_error
+
+
+# register_amd_ci(est_time=300, suite="stage-b-test-large-2-gpu-amd")
+class TestDisaggregationSimulatedRetract(PDDisaggregationServerBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        # Configure ROCm RDMA environment
+        os.environ["SGLANG_USE_AITER"] = "1"
+        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 = []
+
+        os.environ["SGLANG_TEST_RETRACT"] = "true"
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+
+        # Non blocking start servers
+        cls.start_prefill()
+        cls.start_decode()
+
+        # Block until both
+        cls.wait_server_ready(cls.prefill_url + "/health")
+        cls.wait_server_ready(cls.decode_url + "/health")
+
+        cls.launch_lb()
+
+    @classmethod
+    def tearDownClass(cls):
+        os.environ.pop("SGLANG_TEST_RETRACT")
+        super().tearDownClass()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp",
+            "1",
+            "--attention-backend",
+            "aiter",
+            "--log-level",
+            "debug",
+        ]
+        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",
+            "--tp",
+            "1",
+            "--base-gpu-id",
+            "1",
+            "--attention-backend",
+            "aiter",
+            "--mem-fraction-static",
+            "0.8",
+            "--log-level",
+            "debug",
+        ]
+        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 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()

test/registered/amd/disaggregation/test_disaggregation_pp.py

@@ -0,0 +1,291 @@
+import os
+import time
+import unittest
+from types import SimpleNamespace
+
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.few_shot_gsm8k import run_eval
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_pd_server,
+    try_cached_model,
+)
+
+register_amd_ci(est_time=600, suite="stage-b-test-large-8-gpu-35x-disaggregation-amd")
+
+
+class TestDisaggregationPrefillPPAccuracy(PDDisaggregationServerBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        # set up ROCm env
+        os.environ["SGLANG_USE_AITER"] = "1"
+        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.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
+
+        # Non blocking start servers
+        cls.start_prefill()
+        cls.start_decode()
+
+        # Block until both
+        cls.wait_server_ready(cls.prefill_url + "/health")
+        cls.wait_server_ready(cls.decode_url + "/health")
+
+        cls.launch_lb()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp-size",
+            "2",
+            "--pp-size",
+            "2",
+            "--disable-overlap-schedule",
+            "--attention-backend",
+            "aiter",
+        ]
+        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",
+            "--tp-size",
+            "2",
+            "--base-gpu-id",
+            "4",
+            "--attention-backend",
+            "aiter",
+        ]
+        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 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(args)
+        print(f"{metrics=}")
+
+        self.assertGreater(metrics["accuracy"], 0.70)
+        # Wait a little bit so that the memory check happens.
+        time.sleep(5)
+
+
+# register_amd_ci(est_time=200, suite="stage-c-test-large-8-gpu-amd")
+class TestDisaggregationPrefillPPDynamicChunkAccuracy(PDDisaggregationServerBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        # set up ROCm env
+        os.environ["SGLANG_USE_AITER"] = "1"
+        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.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
+
+        # Non blocking start servers
+        cls.start_prefill()
+        cls.start_decode()
+
+        # Block until both
+        cls.wait_server_ready(cls.prefill_url + "/health")
+        cls.wait_server_ready(cls.decode_url + "/health")
+
+        cls.launch_lb()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp-size",
+            "2",
+            "--pp-size",
+            "2",
+            "--disable-overlap-schedule",
+            "--enable-dynamic-chunking",
+            "--attention-backend",
+            "aiter",
+        ]
+        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",
+            "--tp-size",
+            "2",
+            "--base-gpu-id",
+            "4",
+            "--attention-backend",
+            "aiter",
+        ]
+        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 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(args)
+        print(f"{metrics=}")
+
+        self.assertGreater(metrics["accuracy"], 0.70)
+        # Wait a little bit so that the memory check happens.
+        time.sleep(5)
+
+
+# register_amd_ci(est_time=200, suite="stage-c-test-large-8-gpu-amd")
+class TestDisaggregationDecodePPAccuracy(PDDisaggregationServerBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        # set up ROCm env
+        os.environ["SGLANG_USE_AITER"] = "1"
+        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.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
+
+        # Non blocking start servers
+        cls.start_prefill()
+        cls.start_decode()
+
+        # Block until both
+        cls.wait_server_ready(cls.prefill_url + "/health")
+        cls.wait_server_ready(cls.decode_url + "/health")
+
+        cls.launch_lb()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp-size",
+            "2",
+            "--pp-size",
+            "2",
+            "--disable-overlap-schedule",
+            "--attention-backend",
+            "aiter",
+        ]
+        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",
+            "--tp-size",
+            "2",
+            "--pp-size",
+            "2",
+            "--base-gpu-id",
+            "4",
+            "--attention-backend",
+            "aiter",
+        ]
+        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 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(args)
+        print(f"{metrics=}")
+
+        self.assertGreater(metrics["accuracy"], 0.70)
+        # Wait a little bit so that the memory check happens.
+        time.sleep(5)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/run_suite.py

@@ -22,6 +22,7 @@ PER_COMMIT_SUITES = {
         "stage-a-test-1-amd",
         "stage-b-test-small-1-gpu-amd",
         "stage-b-test-small-1-gpu-amd-mi35x",
+        "stage-b-test-large-8-gpu-35x-disaggregation-amd",
         "stage-b-test-large-1-gpu-amd",
         "stage-b-test-large-2-gpu-amd",
         "stage-c-test-large-8-gpu-amd-mi35x",