CI: Migrate MoRI to ROCm 10.1

.claude/skills/deploy-mori/SKILL.md

@@ -52,7 +52,7 @@ done
 ## Step 1: Start the Docker container
 
 Use the container name from the user's args if provided; ask if not.
-Default image: `rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0`
+Default image: `rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0`
 (use unless the user specifies another) — same base the CI workflows build on.
 
 ```bash
@@ -427,35 +427,44 @@ mst status -v
 
 ## Step 4: Install MORI
 
-### 4.0 — ROCm 7.14 dev toolchain (rocm7.14 base images only)
+### 4.0 — ROCm dev toolchain (rocm10.x / rocm7.14 base images only)
 
 **Skip on pre-7.14 images** — they already ship a full `/opt/rocm` and must not
 get a second ROCm stacked on top.
 
-The rocm7.14 pytorch image has **no `/opt/rocm` at all**: ROCm comes from the pip
-`rocm-sdk-core` wheel, which ships runtime `.so` files and headers but no CMake
-package configs. Without this step the build dies at configure time with
+The rocm10.x and rocm7.14 pytorch images have **no `/opt/rocm` at all**: ROCm
+comes from the pip `rocm-sdk-core` wheel, which ships runtime `.so` files and
+headers but no CMake package configs. Without this step the build dies at
+configure time with
 `does not contain the HIP runtime CMake package ... hip-lang-config.cmake`.
 
+Install the release in the image's tag — mixing them makes every `import mori`
+warn that it was built against another ROCm. For ROCm 10.1 (an Ubuntu 22.04
+image takes `ubuntu2204` in the repo URL):
+
 ```bash
 sudo docker exec $CONTAINER_NAME bash -c '
 set -e
 [ -d /opt/rocm/bin ] && { echo "/opt/rocm already present — skipping"; exit 0; }
 apt-get install -y --no-install-recommends wget gnupg ca-certificates
 mkdir -p --mode=0755 /etc/apt/keyrings
-wget -qO - https://repo.amd.com/rocm/packages-multi-arch/gpg/rocm.gpg \
+wget -qO - https://stable.repo.amd.com/rocm/gpg/packages.gpg \
   | gpg --dearmor > /etc/apt/keyrings/amdrocm.gpg
-echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdrocm.gpg] https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404 stable main" \
+echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdrocm.gpg] https://stable.repo.amd.com/rocm/core/packages/ubuntu2404 stable main" \
   > /etc/apt/sources.list.d/amdrocm.list
 apt-get update
-apt-get install -y --no-install-recommends amdrocm-core-dev7.14-gfx942 amdrocm-core-dev7.14-gfx950
-hipconfig --version
+apt-get install -y --no-install-recommends amdrocm-core-dev10.1-gfx942 amdrocm-core-dev10.1-gfx950
+cat /opt/rocm/core/.info/version
 '
 ```
 
+A rocm7.14 image takes `amdrocm-core-dev7.14-<arch>` from the
+`https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404` repo instead, signed
+by `https://repo.amd.com/rocm/packages-multi-arch/gpg/rocm.gpg`.
+
 No environment setup is needed for the rest of the session. CMake resolves ROCm
-as `-DROCM_PATH` → `$ROCM_PATH` → `$HIP_PATH` → `/opt/rocm` (where both the
-7.2.4 and 7.14 stacks live) and pins the compiler to an absolute
+as `-DROCM_PATH` → `$ROCM_PATH` → `$HIP_PATH` → `/opt/rocm` (where the 7.2.4,
+7.14 and 10.x stacks all live) and pins the compiler to an absolute
 `<rocm>/llvm/bin/amdclang++`; mori's libraries carry
 `RUNPATH "$ORIGIN:<rocm>/lib"`, and the JIT builds absolute `<rocm>/bin/...`
 paths off the same default. Set `ROCM_PATH` only for ROCm somewhere else (e.g. a

.github/workflows/ci.yml

@@ -31,7 +31,7 @@ concurrency:
 
 env:
   IMAGE: rocm/mori:ci
-  BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+  BASE_IMAGE: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
   CONTAINER: mori_ci_${{ github.run_id }}
   CT: docker
 
@@ -83,6 +83,12 @@ jobs:
           $CT exec $CONTAINER \
             git config --global --add safe.directory $GITHUB_WORKSPACE
 
+      - name: Report ROCm version
+        run: |
+          $CT exec $CONTAINER bash -c "
+            ls -d /opt/rocm/core-* && /opt/rocm/bin/hipconfig --version && echo
+          "
+
       - name: Build and install mori
         run: $CT exec $CONTAINER pip install .
 
@@ -454,19 +460,17 @@ jobs:
                 chown -R $(id -u):$(id -g) $GITHUB_WORKSPACE 2>/dev/null || true
           fi
 
-  # Gate for "does mori still build and run on the next ROCm major". Kept to the
-  # single EP intranode test on one runner: the rest of the intranode suite is
-  # ROCm-version agnostic, and the ROCm 10 base image is a separate ~30GB pull.
-  intranode-test-rocm10:
-    name: mori intranode test (ROCm 10, MI355X_AINIC)
+  # Gate for "does mori still build and run on ROCm 7.14"; every other job runs
+  # on 10.1. Kept to the EP intranode test on one runner: the rest of the
+  # intranode suite is ROCm-version agnostic.
+  intranode-test-rocm714:
+    name: mori intranode test (ROCm 7.14, MI355X_AINIC)
     needs: build
     runs-on: [self-hosted, MI355X-AINIC-PIT-G17]
     env:
-      # py3.11 rather than the 7.14 lane's py3.12 — this is the tag the ROCm 10
-      # build was validated on.
-      BASE_IMAGE: rocm/pytorch:rocm10.0_ubuntu24.04_py3.11_pytorch_release_2.13.0
-      IMAGE: rocm/mori:ci-rocm10
-      CONTAINER: mori_ci_rocm10_${{ github.run_id }}
+      BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+      IMAGE: rocm/mori:ci-rocm714
+      CONTAINER: mori_ci_rocm714_${{ github.run_id }}
       MORI_RDMA_DEVICES: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
       MORI_RDMA_SL: 3
       MORI_RDMA_TC: 104
@@ -491,14 +495,14 @@ jobs:
           # main; otherwise it would reuse the merge commit of the first attempt.
           ref: ${{ github.event_name == 'pull_request' && github.run_attempt > 1 && format('refs/pull/{0}/merge', github.event.pull_request.number) || '' }}
 
-      # ROCm 10 ships in a different apt repo than 7.x, hence the ROCM_APT_* args.
+      # ROCm 7.14 ships in a different apt repo than 10.x, hence the ROCM_APT_* args.
       - name: Build CI image
         run: |
           $CT build --network=host \
             --build-arg BASE_IMAGE=$BASE_IMAGE \
-            --build-arg ROCM_APT_VERSION=10.0 \
-            --build-arg ROCM_APT_GPG=https://stable.repo.amd.com/rocm/gpg/packages.gpg \
-            --build-arg ROCM_APT_LIST="https://stable.repo.amd.com/rocm/core/packages/ubuntu2404 stable main" \
+            --build-arg ROCM_APT_VERSION=7.14 \
+            --build-arg ROCM_APT_GPG=https://repo.amd.com/rocm/packages-multi-arch/gpg/rocm.gpg \
+            --build-arg ROCM_APT_LIST="https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404 stable main" \
             -t $IMAGE -f docker/Dockerfile.dev .
 
       - name: Start container
@@ -522,10 +526,12 @@ jobs:
             ls -d /opt/rocm/core-* && /opt/rocm/bin/hipconfig --version && echo
           "
 
+      # The same build as intranode-test's, so a ROCm 10-only API anywhere in it
+      # fails to compile here.
       - name: Install mori
         run: |
           $CT exec $CONTAINER bash -c \
-            "cd $GITHUB_WORKSPACE && pip install ."
+            "cd $GITHUB_WORKSPACE && BUILD_BENCHMARK=ON BUILD_CCO_SDMA=ON BUILD_EXAMPLES=ON pip install ."
 
       - name: Verify installation
         run: $CT exec $CONTAINER bash -c "cd $GITHUB_WORKSPACE && python tools/verify_install.py"
@@ -556,20 +562,19 @@ jobs:
           - platform: MI355X_AINIC
             runner: [self-hosted, MI355X-AINIC-PIT-G17]
             rdma_devices: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
-            gpu_arch: gfx950
             socket_ifname: eno0
           # Runner 325_bnxt_108 is decommissioned.
           # - platform: MI325X_BNXT
           #   runner: [self-hosted, 325_bnxt_108]
           #   rdma_devices: bnxt_re0,bnxt_re1,bnxt_re2,bnxt_re3,bnxt_re4,bnxt_re5,bnxt_re7,bnxt_re8
-          #   gpu_arch: gfx942
           #   socket_ifname: enp49s0f1np1
     env:
       # jax.distributed rendezvouses over localhost, but mori.jax.shmem_init_attr
       # goes through MORI's own bootstrap, which needs the same eno0 pin as
       # intranode-test.
       MORI_SOCKET_IFNAME: ${{ matrix.socket_ifname }}
       NCCL_SOCKET_IFNAME: ${{ matrix.socket_ifname }}
+      JAX_BASE_IMAGE: rocm/jax:rocm10.1.0_ubuntu24.04_py3.12_jax_release_0.10.2
     timeout-minutes: 90
     steps:
       - name: Fix workspace permissions
@@ -588,15 +593,18 @@ jobs:
           # main; otherwise it would reuse the merge commit of the first attempt.
           ref: ${{ github.event_name == 'pull_request' && github.run_attempt > 1 && format('refs/pull/{0}/merge', github.event.pull_request.number) || '' }}
 
+      # The JAX image has no ROCm dev stack either, so it takes the same
+      # Dockerfile as the torch images.
       - name: Build JAX CI image
         run: |
-          BUILD_ARGS="--network=host --build-arg GPU_TARGET=${{ matrix.gpu_arch }}"
+          BUILD_ARGS="--network=host"
           if [ "$CT" = "podman" ]; then
             BUILD_ARGS="$BUILD_ARGS --isolation=chroot"
           fi
           $CT build $BUILD_ARGS \
+            --build-arg BASE_IMAGE=$JAX_BASE_IMAGE \
             -t rocm/mori:jax-ci \
-            -f docker/Dockerfile.jax091-rocm711 .
+            -f docker/Dockerfile.dev .
 
       - name: Start container
         run: |
@@ -616,10 +624,11 @@ jobs:
           # Resolve the XLA FFI headers from the system jax *outside* pip's
           # isolated build env (build isolation hides system site-packages, so
           # CMake can't `import jax` during configure). Pass it via env var.
+          # The JAX image ships no pytest.
           $CT exec ${CONTAINER}_jax bash -c \
             "cd $GITHUB_WORKSPACE && \
              export XLA_FFI_INCLUDE_DIR=\$(python3 -c 'from jax import ffi; print(ffi.include_dir())') && \
-             BUILD_UMBP=OFF BUILD_XLA_FFI_OPS=ON pip install --break-system-packages ."
+             BUILD_UMBP=OFF BUILD_XLA_FFI_OPS=ON pip install . && pip install pytest"
 
       - name: MORI-EP JAX test
         run: |
@@ -760,6 +769,14 @@ jobs:
             rdma_tc: 104
             ct: podman
             timeout_minutes: 45
+            # These hosts' amdgpu-dkms lacks 93aa919 ("drm/amdkfd: fix vram
+            # allocation failure for a special case"): hipMemCreate (VRAM with no VA
+            # yet) fails with EADDRINUSE, reported as out of memory, while any SVM
+            # range sits below its size, and libhsakmt registers every pageable
+            # buffer HIP pins for a copy as one. HSA_USE_SVM=0 registers them as
+            # userptr instead. Not in the image: with SVM off every pageable copy is
+            # pinned again. Drop it once the hosts' amdgpu-dkms carries the fix.
+            hsa_use_svm: 0
           # 8x ConnectX-7 NDR on InfiniBand, so no rdma_sl/rdma_tc: mlx5 applies
           # both only on the RoCE path. mlx5_8 is deliberately not in
           # rdma_devices -- it is the 100GbE port behind ens14np0 (the TCP
@@ -806,6 +823,7 @@ jobs:
       MORI_RDMA_SL: ${{ matrix.rdma_sl }}
       MORI_RDMA_TC: ${{ matrix.rdma_tc }}
       MORI_SOCKET_IFNAME: ${{ matrix.node1_ifname }}
+      HSA_USE_SVM: ${{ matrix.hsa_use_svm }}
     timeout-minutes: ${{ matrix.timeout_minutes }}
     steps:
       - name: Fix workspace permissions
@@ -850,6 +868,7 @@ jobs:
             -e MORI_RDMA_DEVICES=$MORI_RDMA_DEVICES \
             ${MORI_RDMA_SL:+-e MORI_RDMA_SL=$MORI_RDMA_SL} \
             ${MORI_RDMA_TC:+-e MORI_RDMA_TC=$MORI_RDMA_TC} \
+            ${HSA_USE_SVM:+-e HSA_USE_SVM=$HSA_USE_SVM} \
             -e MORI_SOCKET_IFNAME=$MORI_SOCKET_IFNAME \
             -e GLOO_SOCKET_IFNAME=$NODE1_IFNAME \
             -e NCCL_SOCKET_IFNAME=$NODE1_IFNAME \
@@ -892,6 +911,7 @@ jobs:
               -e MORI_RDMA_DEVICES=$MORI_RDMA_DEVICES \
               ${MORI_RDMA_SL:+-e MORI_RDMA_SL=$MORI_RDMA_SL} \
               ${MORI_RDMA_TC:+-e MORI_RDMA_TC=$MORI_RDMA_TC} \
+              ${HSA_USE_SVM:+-e HSA_USE_SVM=$HSA_USE_SVM} \
               -e MORI_SOCKET_IFNAME=$MORI_SOCKET_IFNAME \
               -e GLOO_SOCKET_IFNAME=$NODE2_IFNAME \
               -e NCCL_SOCKET_IFNAME=$NODE2_IFNAME \

.github/workflows/ci_cco.yml

@@ -12,7 +12,7 @@ concurrency:
 
 env:
   IMAGE: rocm/mori:ci-cco
-  BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+  BASE_IMAGE: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
   CONTAINER: mori_cco_ci_${{ github.run_id }}
   CT: podman  # MI355X-AINIC runner is podman rootless (no docker daemon)
 
@@ -29,10 +29,19 @@ jobs:
             rdma_devices: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
             rdma_sl: 3
             rdma_tc: 104
+            # These hosts' amdgpu-dkms lacks 93aa919 ("drm/amdkfd: fix vram
+            # allocation failure for a special case"): hipMemCreate (VRAM with no VA
+            # yet) fails with EADDRINUSE, reported as out of memory, while any SVM
+            # range sits below its size, and libhsakmt registers every pageable
+            # buffer HIP pins for a copy as one. HSA_USE_SVM=0 registers them as
+            # userptr instead. Not in the image: with SVM off every pageable copy is
+            # pinned again. Drop it once the hosts' amdgpu-dkms carries the fix.
+            hsa_use_svm: 0
     env:
       MORI_RDMA_DEVICES: ${{ matrix.rdma_devices }}
       MORI_RDMA_SL: ${{ matrix.rdma_sl }}
       MORI_RDMA_TC: ${{ matrix.rdma_tc }}
+      HSA_USE_SVM: ${{ matrix.hsa_use_svm }}
     timeout-minutes: 60
     steps:
       - name: Checkout
@@ -56,6 +65,7 @@ jobs:
             -e MORI_RDMA_DEVICES=$MORI_RDMA_DEVICES \
             -e MORI_RDMA_SL=$MORI_RDMA_SL \
             -e MORI_RDMA_TC=$MORI_RDMA_TC \
+            -e HSA_USE_SVM=$HSA_USE_SVM \
             -v $GITHUB_WORKSPACE:$GITHUB_WORKSPACE \
             -w $GITHUB_WORKSPACE \
             $IMAGE sleep infinity
@@ -433,6 +443,7 @@ jobs:
             rdma_devices: rocep105s0,rocep121s0,rocep137s0,rocep153s0,rocep233s0,rocep249s0,rocep25s0,rocep9s0
             rdma_sl: 3
             rdma_tc: 104
+            hsa_use_svm: 0  # the amdgpu-dkms workaround, see cco-unit-test
     env:
       # Overrides the workflow-level name, which is keyed by run_id alone: sharing
       # it would have this job's cleanup delete the other job's container.
@@ -447,6 +458,7 @@ jobs:
       MORI_RDMA_SL: ${{ matrix.rdma_sl }}
       MORI_RDMA_TC: ${{ matrix.rdma_tc }}
       MORI_SOCKET_IFNAME: ${{ matrix.node1_ifname }}
+      HSA_USE_SVM: ${{ matrix.hsa_use_svm }}
       ENTRY: tests/python/ops/dispatch_combine_v2/test_dispatch_combine_v2_internode.py
     timeout-minutes: 60
     steps:
@@ -474,6 +486,7 @@ jobs:
             -e MORI_SOCKET_IFNAME=$MORI_SOCKET_IFNAME \
             -e GLOO_SOCKET_IFNAME=$NODE1_IFNAME \
             -e NCCL_SOCKET_IFNAME=$NODE1_IFNAME \
+            -e HSA_USE_SVM=$HSA_USE_SVM \
             -v $GITHUB_WORKSPACE:$GITHUB_WORKSPACE \
             -w $GITHUB_WORKSPACE \
             $IMAGE sleep infinity
@@ -502,6 +515,7 @@ jobs:
               -e MORI_SOCKET_IFNAME=$MORI_SOCKET_IFNAME \
               -e GLOO_SOCKET_IFNAME=$NODE2_IFNAME \
               -e NCCL_SOCKET_IFNAME=$NODE2_IFNAME \
+              -e HSA_USE_SVM=$HSA_USE_SVM \
               -v $GITHUB_WORKSPACE:$GITHUB_WORKSPACE \
               -w $GITHUB_WORKSPACE \
               $IMAGE sleep infinity &&

.github/workflows/codeql.yml

@@ -32,7 +32,7 @@ concurrency:
 
 env:
   IMAGE: rocm/mori:ci
-  BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+  BASE_IMAGE: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
   CONTAINER: mori_codeql_${{ github.run_id }}
   CT: docker
   # falls back to cpp,python on push events (where inputs.languages is empty)