[CI/Build] Add self-hosted MI350X ROCm external PD tier

.github/workflows/ci.yml

@@ -399,6 +399,20 @@ jobs:
     - name: Test CI service and smoke scripts
       run: python -m unittest discover -s scripts/ci/tests -v
 
+  tone-script-tests:
+    name: Test external CI lifecycle helpers
+    runs-on: ubuntu-22.04
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
+
+    - name: Run focused shell regression tests
+      run: |
+        bash scripts/tone_tests/tests/test_common_cleanup.sh
+        python3 scripts/tone_tests/tests/test_hicache_model_selection.py
+      shell: bash
+
   clang-format:
     name: Check code format
     runs-on: ubuntu-22.04
@@ -750,6 +764,7 @@ jobs:
       - docs-check
       - reshard-type-check
       - ci-script-tests
+      - tone-script-tests
       - build-wheel
       - unit-tests
       - build-flags
@@ -771,7 +786,8 @@ jobs:
                 needs.clang-format.result == 'success' &&
                 needs.python-lint.result == 'success' &&
                 needs.reshard-type-check.result == 'success' &&
-                needs.ci-script-tests.result == 'success' }}
+                needs.ci-script-tests.result == 'success' &&
+                needs.tone-script-tests.result == 'success' }}
         run: |
           echo "$NEEDS_JSON"
           test "$PASSED" = true

.github/workflows/ci_rocm.yml

@@ -3,11 +3,21 @@ name: 'Build Wheel (ROCm)'
 # ROCm/HIP CI parity with the standard CUDA wheel CI: build the AMD ROCm wheel
 # on PRs so packaging regressions are caught. Runs inside the ROCm dev image so
 # hipcc / HIP headers / hsa-runtime are available; no GPU is needed to compile.
-# On-device transfer tests (test_transfer_on_hip.py) require AMD hardware and
-# run outside GitHub-hosted runners.
+# On-device external-PD tests consume these artifacts later on the trusted
+# self-hosted ROCm runner in integration-test-rocm.yml.
 
 on:
-  workflow_call: {}
+  workflow_call:
+    inputs:
+      checkout_ref:
+        description: 'Git revision to build (used by trusted E2E dispatches)'
+        required: false
+        type: string
+      python_versions:
+        description: 'JSON array of Python versions to build'
+        required: false
+        default: '["3.10", "3.12"]'
+        type: string
 
 jobs:
   build-wheel-rocm:
@@ -17,9 +27,10 @@ jobs:
     # so it also loads on the ROCm 7.0 image variants.
     container: rocm/dev-ubuntu-22.04:7.2.3-complete
     strategy:
+      fail-fast: false
       matrix:
         # 3.10 covers the SGLang ROCm image; 3.12 covers vLLM ROCm.
-        python-version: ['3.10', '3.12']
+        python-version: ${{ fromJSON(inputs.python_versions) }}
     env:
       HIP_BUILD: "1"
       SCCACHE_GHA_ENABLED: "true"
@@ -37,15 +48,35 @@ jobs:
         apt-get update -y
         apt-get install -y --no-install-recommends git ca-certificates
 
-    - uses: actions/checkout@v4
+    - name: Check out source
+      uses: actions/checkout@v4
       with:
+        ref: ${{ inputs.checkout_ref || github.sha }}
         persist-credentials: false
-        submodules: recursive
+        submodules: false
 
     - name: Mark repository as safe
       run: git config --global --add safe.directory "$GITHUB_WORKSPACE" || true
       shell: bash
 
+    - name: Check out submodules
+      shell: bash
+      run: |
+        set -eo pipefail
+        git submodule sync --recursive
+        for attempt in 1 2 3; do
+          if git -c http.version=HTTP/1.1 submodule update \
+            --init --recursive --depth=1 --jobs=2; then
+            exit 0
+          fi
+          if [ "$attempt" -eq 3 ]; then
+            echo "Submodule checkout failed after $attempt attempts" >&2
+            exit 1
+          fi
+          echo "Submodule checkout attempt $attempt failed; retrying..." >&2
+          sleep $((attempt * 10))
+        done
+
     - name: Install toolchain and Python ${{ matrix.python-version }}
       shell: bash
       run: |
@@ -96,6 +127,7 @@ jobs:
         rm -rf build && mkdir build && cd build
         cmake -G Ninja .. \
           -DUSE_HIP=ON -DUSE_CUDA=OFF -DWITH_EP=OFF \
+          -DENABLE_MULTI_PROTOCOL=ON \
           -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON \
           -DBUILD_UNIT_TESTS=OFF -DENABLE_SCCACHE=ON \
           -DCMAKE_BUILD_TYPE=Release \
@@ -147,7 +179,8 @@ jobs:
           mooncake-wheel/dist-rocm-py${{ steps.generate_tag.outputs.python_version_tag }}/*.whl
         # ROCm runtime is excluded from the wheel and bound at load time.
         export LD_LIBRARY_PATH="/opt/rocm/lib:/usr/local/lib:${LD_LIBRARY_PATH:-}"
-        site="$("$smoke_venv/bin/python" -c 'import mooncake,os;print(os.path.dirname(mooncake.__file__))')"
+        "$smoke_venv/bin/python" scripts/tone_tests/python/verify_rocm_wheel.py
+        site="$("$smoke_venv/bin/python" -c 'import mooncake, os; print(os.path.dirname(mooncake.__file__))')"
         "$site/mooncake_master" --version
 
     - name: Upload Python wheel artifact

.github/workflows/e2e-ci.yml

@@ -37,6 +37,14 @@ jobs:
       checkout_ref: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}
     secrets: inherit
 
+  build-wheel-rocm:
+    if: >
+      github.event_name == 'workflow_dispatch' ||
+      github.event.label.name == 'run-e2e-ci'
+    uses: ./.github/workflows/ci_rocm.yml
+    with:
+      checkout_ref: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}
+
   integration-test:
     if: >
       github.event_name == 'workflow_dispatch' ||
@@ -47,6 +55,17 @@ jobs:
       pr_number: ${{ inputs.pr_number || github.event.pull_request.number }}
     secrets: inherit
 
+  rocm-integration-test:
+    if: >
+      github.event_name == 'workflow_dispatch' ||
+      github.event.label.name == 'run-e2e-ci'
+    needs:
+      - build-wheel-rocm
+    uses: ./.github/workflows/integration-test-rocm.yml
+    with:
+      checkout_ref: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}
+    secrets: inherit
+
   e2e-gate:
     name: E2E Gate
     if: >
@@ -55,7 +74,9 @@ jobs:
        github.event.label.name == 'run-e2e-ci')
     needs:
       - ascend-test
+      - build-wheel-rocm
       - integration-test
+      - rocm-integration-test
     runs-on: ubuntu-latest
     steps:
       - name: Check E2E results

.github/workflows/integration-test-rocm.yml

@@ -0,0 +1,167 @@
+name: 'Integration test (ROCm)'
+
+on:
+  workflow_call:
+    inputs:
+      checkout_ref:
+        description: 'Git revision whose ROCm wheel and integration tests are exercised'
+        required: false
+        type: string
+      run_sglang:
+        description: 'Run SGLang cases in addition to vLLM'
+        required: false
+        default: true
+        type: boolean
+
+jobs:
+  test-rocm-integration:
+    name: ROCm MI350X external PD
+    runs-on: [self-hosted, Linux, X64, amd, rocm, gfx950, mooncake-pd]
+    timeout-minutes: 390
+    concurrency:
+      group: mooncake-rocm-mi350x
+      cancel-in-progress: false
+    env:
+      REGISTRY_ADDR_SGLANG: 'lmsysorg/sglang:v0.5.18-rocm720-mi35x@sha256:6d68cd19206716cb3f1e31e2ad89cd0852d7ae614a792773c30a4277f8955c72'
+      REGISTRY_ADDR_VLLM: 'vllm/vllm-openai-rocm:v0.21.0@sha256:98a77b20df03adeb1cfc0ced009b4df6dd52b0a994ab99a32421f30876a9ae0c'
+    steps:
+      # Pytest previously ran as root in a bind-mounted container and left
+      # caches that the non-root runner could not remove on the next checkout.
+      - name: Remove root-owned Python test caches
+        shell: bash
+        run: |
+          set -euo pipefail
+          docker run --rm --network=none \
+            -v "${GITHUB_WORKSPACE}:/workspace" \
+            --entrypoint /bin/rm \
+            "${REGISTRY_ADDR_VLLM}" \
+            -rf \
+            /workspace/scripts/tone_tests/python/.pytest_cache \
+            /workspace/scripts/tone_tests/python/__pycache__
+
+      - name: Check out tested revision
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.checkout_ref || github.sha }}
+          persist-credentials: false
+          clean: true
+
+      - name: Download Python 3.10 ROCm wheel
+        if: inputs.run_sglang
+        uses: actions/download-artifact@v4
+        with:
+          name: mooncake-wheel-rocm-ubuntu-py310
+          path: ${{ runner.temp }}/mooncake-wheel-rocm-py310
+
+      - name: Download Python 3.12 ROCm wheel
+        uses: actions/download-artifact@v4
+        with:
+          name: mooncake-wheel-rocm-ubuntu-py312
+          path: ${{ runner.temp }}/mooncake-wheel-rocm-py312
+
+      - name: Run two-node ROCm external PD suite
+        shell: bash
+        env:
+          WHEEL_DIR_SGLANG: ${{ runner.temp }}/mooncake-wheel-rocm-py310
+          WHEEL_DIR_VLLM: ${{ runner.temp }}/mooncake-wheel-rocm-py312
+        run: |
+          set -euo pipefail
+          [ "$(id -u)" -ne 0 ] || {
+            echo "The self-hosted runner must not run as root" >&2
+            exit 1
+          }
+          [ -r /etc/mooncake-ci/runner.env ] || {
+            echo "Missing controller profile: /etc/mooncake-ci/runner.env" >&2
+            exit 1
+          }
+          set -a
+          # shellcheck disable=SC1091
+          source /etc/mooncake-ci/runner.env
+          set +a
+          export REGISTRY_ADDR_SGLANG REGISTRY_ADDR_VLLM
+          export WHEEL_DIR_SGLANG WHEEL_DIR_VLLM
+
+          required_vars=(
+            CI_ACCELERATOR LOCAL_IP REMOTE_IP REMOTE_SSH_TARGET
+            MODEL_CACHE MOONCAKE_RUNTIME_CACHE MOONCAKE_CI_TIER
+            MOONCAKE_GPU_INDICES
+            MOONCAKE_RENDER_DEVICES MOONCAKE_CPUSET_CPUS
+            MOONCAKE_CPUSET_MEMS MOONCAKE_SSH_CONFIG
+            MOONCAKE_LOCK_FILE REMOTE_TEST_DIR MOONCAKE_RDMA_DEVICES
+            MOONCAKE_RDMA_NETDEVS MOONCAKE_TRANSFER_DEVICE
+            MOONCAKE_GID_INDEX MOONCAKE_SGLANG_BASE_GPU_ID
+            MOONCAKE_EPD_ENCODER_GPU_ID MOONCAKE_EPD_PREFILL_GPU_ID
+            MOONCAKE_EPD_DECODE_GPU_ID MOONCAKE_VLLM_VISIBLE_DEVICES
+            MOONCAKE_SGLANG_MEM_FRACTION_STATIC AINIC_VERSION
+          )
+          for variable in "${required_vars[@]}"; do
+            [ -n "${!variable:-}" ] || {
+              echo "Missing controller setting: $variable" >&2
+              exit 1
+            }
+          done
+          [ "$CI_ACCELERATOR" = rocm ]
+          export SSH_CMD="ssh -F ${MOONCAKE_SSH_CONFIG}"
+
+          for command in docker flock python3 realpath rocm-smi rsync ssh; do
+            command -v "$command" >/dev/null || {
+              echo "Missing required command: $command" >&2
+              exit 1
+            }
+          done
+          test -c /dev/kfd
+          test -d /sys/class/kfd/kfd/topology/nodes
+          test -d /sys/class/kfd/kfd/proc
+          ssh -F "${MOONCAKE_SSH_CONFIG}" "${REMOTE_SSH_TARGET}" '
+            set -eu
+            for command in docker python3 rocm-smi; do
+              command -v "$command" >/dev/null || {
+                echo "Missing required remote command: $command" >&2
+                exit 1
+              }
+            done
+            test -c /dev/kfd
+            test -d /sys/class/kfd/kfd/topology/nodes
+            test -d /sys/class/kfd/kfd/proc
+          '
+          [ -w "$(dirname "$MOONCAKE_LOCK_FILE")" ] || {
+            echo "Runner cannot write the Mooncake lock directory" >&2
+            exit 1
+          }
+
+          exec 9>"${MOONCAKE_LOCK_FILE}"
+          flock -n 9 || {
+            echo "The MI350X Mooncake allocation is already in use" >&2
+            exit 1
+          }
+
+          cleanup_best_effort() {
+            local exit_code=$?
+            trap - EXIT
+            if [ -f scripts/tone_tests/run/.shrc ]; then
+              set +e
+              # shellcheck disable=SC1091
+              source scripts/tone_tests/run/.shrc
+              # shellcheck disable=SC1091
+              source scripts/tone_tests/scripts/common.sh
+              cleanup_test_env double || true
+            fi
+            exit "$exit_code"
+          }
+          trap cleanup_best_effort EXIT
+          trap 'exit 130' INT
+          trap 'exit 143' TERM
+          if ${{ inputs.run_sglang }}; then
+            bash scripts/tone_tests/scripts/run_test.sh run-all
+          else
+            bash scripts/tone_tests/scripts/run_test.sh run-all VLLM
+          fi
+
+      - name: Upload ROCm external PD logs
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: rocm-external-pd-logs-${{ github.run_id }}-${{ github.run_attempt }}
+          path: scripts/tone_tests/run/logs
+          if-no-files-found: warn
+          retention-days: 14

.github/workflows/nightly.yml

@@ -93,6 +93,17 @@ jobs:
       artifact_name: nightly-cuda13-x86-py312
     secrets: inherit
 
+  build-wheel-rocm:
+    needs: version-stamp
+    uses: ./.github/workflows/ci_rocm.yml
+
+  rocm-integration-test:
+    needs: build-wheel-rocm
+    uses: ./.github/workflows/integration-test-rocm.yml
+    with:
+      checkout_ref: ${{ github.sha }}
+    secrets: inherit
+
   build-musa:
     runs-on: ubuntu-22.04
     container: registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64
@@ -575,6 +586,8 @@ jobs:
       - version-stamp
       - build-wheels
       - tone-sglang-integration
+      - build-wheel-rocm
+      - rocm-integration-test
       - build-musa
       - build-docker
       - ascend-test