feat(ci): cross-platform Docker matrix and hardware CI for CUDA/ROCm (closes #173)

.github/actionlint.yaml

@@ -0,0 +1,4 @@
+self-hosted-runner:
+  labels:
+    - cuda
+    - rocm

.github/actions/setup-runpod/action.yml

@@ -0,0 +1,45 @@
+name: Setup RunPod CLI
+description: >-
+  Install runpodctl, authenticate it, and materialize the SSH private key used
+  to reach hosted GPU instances. Exports RUNPOD_SSH_KEY_PATH for later steps.
+
+inputs:
+  api-key:
+    description: RunPod API key
+    required: true
+  ssh-private-key:
+    description: Ed25519 private key authorized for hosted GPU instances
+    required: true
+
+runs:
+  using: composite
+  steps:
+    - name: Install runpodctl
+      shell: bash
+      env:
+        RUNPODCTL_VERSION: v2.6.1
+        RUNPODCTL_SHA256: eaecd8d7110022b0a8d870b7462628eab4bb47c67947121be7ef201ce914ae4f
+      run: |
+        curl -fsSL -o runpodctl "https://github.com/runpod/runpodctl/releases/download/${RUNPODCTL_VERSION}/runpodctl-linux-amd64"
+        echo "${RUNPODCTL_SHA256}  runpodctl" | sha256sum -c -
+        chmod +x runpodctl
+        sudo mv runpodctl /usr/local/bin/runpodctl
+        runpodctl version
+
+    - name: Configure runpodctl
+      shell: bash
+      run: runpodctl config --apiKey "${RUNPOD_API_KEY}"
+      env:
+        RUNPOD_API_KEY: ${{ inputs.api-key }}
+
+    - name: Setup SSH key
+      shell: bash
+      run: |
+        mkdir -p ~/.ssh && chmod 700 ~/.ssh
+        KEY_PATH="$RUNNER_TEMP/hosted-gpu-ci-key"
+        printf '%s\n' "$RUNPOD_SSH_PRIVATE_KEY" > "$KEY_PATH"
+        chmod 600 "$KEY_PATH"
+        ssh-keygen -y -f "$KEY_PATH" > /dev/null
+        echo "RUNPOD_SSH_KEY_PATH=$KEY_PATH" >> "$GITHUB_ENV"
+      env:
+        RUNPOD_SSH_PRIVATE_KEY: ${{ inputs.ssh-private-key }}

.github/workflows/build-ci-image.yml

@@ -1,29 +1,106 @@
-name: Build CI Image (NVIDIA/CUDA)
+name: Build CI Images
 
 on:
   push:
     branches: [ main ]
     paths:
-      - 'docker/Dockerfile.cuda'
+      - 'docker/**'
       - 'pyproject.toml'
       - 'setup.py'
       - 'requirements*.txt'
       - 'csrc/**'
+      - '.github/workflows/build-ci-image.yml'
+  pull_request:
+    branches: [ main ]
+    paths:
+      - 'docker/**'
+      - 'pyproject.toml'
+      - 'setup.py'
+      - 'requirements*.txt'
+      - 'csrc/**'
+      - '.github/workflows/build-ci-image.yml'
   workflow_dispatch:
 
 jobs:
+  build-pr:
+    if: github.event_name == 'pull_request'
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - backend: cuda
+            dockerfile: docker/Dockerfile.cuda
+            tag: cuda
+          - backend: rocm
+            dockerfile: docker/Dockerfile.rocm
+            tag: rocm
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      # The CUDA/ROCm devel base images are tens of GB; ubuntu-latest only has
+      # ~14 GB free by default. Drop preinstalled toolchains we never use.
+      - name: Free runner disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \
+            /opt/hostedtoolcache/CodeQL /usr/local/share/boost
+          sudo docker image prune -af
+          df -h /
+
+      - name: Set lower case image name
+        run: |
+          REPO="${GITHUB_REPOSITORY,,}"
+          echo "IMAGE=ghcr.io/$REPO/rl-kernel-ci" >> "$GITHUB_ENV"
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Build image
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: ${{ matrix.dockerfile }}
+          platforms: linux/amd64
+          push: false
+          tags: ${{ env.IMAGE }}:${{ matrix.tag }}
+          cache-from: type=gha,scope=ci-image-${{ matrix.backend }}
+
   build-and-push:
+    if: github.event_name != 'pull_request'
     runs-on: ubuntu-latest
     permissions:
       contents: read
       packages: write
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - backend: cuda
+            dockerfile: docker/Dockerfile.cuda
+            tag: cuda
+          - backend: rocm
+            dockerfile: docker/Dockerfile.rocm
+            tag: rocm
     steps:
       - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      - name: Free runner disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \
+            /opt/hostedtoolcache/CodeQL /usr/local/share/boost
+          sudo docker image prune -af
+          df -h /
 
       - name: Set lower case image name
         run: |
           REPO="${GITHUB_REPOSITORY,,}"
-          echo "IMAGE=ghcr.io/$REPO/rl-kernel-ci:cuda" >> "$GITHUB_ENV"
+          echo "IMAGE=ghcr.io/$REPO/rl-kernel-ci" >> "$GITHUB_ENV"
 
       - name: Log in to GitHub Container Registry
         uses: docker/login-action@v3
@@ -32,12 +109,16 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: Build and push
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Build and push image
         uses: docker/build-push-action@v6
         with:
           context: .
-          file: docker/Dockerfile.cuda
+          file: ${{ matrix.dockerfile }}
+          platforms: linux/amd64
           push: true
-          tags: ${{ env.IMAGE }}
-          cache-from: type=gha
-          cache-to: type=gha,mode=max
+          tags: ${{ env.IMAGE }}:${{ matrix.tag }}
+          cache-from: type=gha,scope=ci-image-${{ matrix.backend }}
+          cache-to: type=gha,mode=max,scope=ci-image-${{ matrix.backend }}

.github/workflows/ci.yml

@@ -61,13 +61,19 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install torch --index-url https://download.pytorch.org/whl/cpu
-          pip install -e ".[dev]"
+          pip install -e ".[test]"
 
       - name: Run Mocked Hardware Discovery Tests
         run: |
           python -m pytest rl_engine/tests/test_dispatch.py -v
+          python -m pytest tests/test_kernel_registry.py -q
           PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/test_attention_correctness.py -q -rs
 
+      - name: Run CPU-only Operator Tests
+        run: |
+          python -m pytest tests/test_linear_logp.py -q -rs -k "native"
+          python -m pytest tests/test_cpu_hal.py -q
+
       - name: Run Attention Ground-Truth Tests (CPU-safe)
         run: |
           python -m pytest tests/test_attention.py -v -k "not large and not gpu"

.github/workflows/gpu-ci.yml

@@ -1,4 +1,4 @@
-name: GPU CI
+name: Hardware GPU CI
 
 on:
   pull_request_target:
@@ -10,47 +10,150 @@ on:
       - 'pyproject.toml'
       - 'setup.py'
       - 'requirements*.txt'
-      - 'docker/Dockerfile.cuda'
+      - 'docker/**'
       - 'ci/run_gpu_ci.sh'
+      - 'ci/run_cuda_tests.sh'
+      - 'ci/run_rocm_ci.sh'
+      - 'ci/run_rocm_container.sh'
       - '.github/workflows/gpu-ci.yml'
-    types: [ opened, synchronize, reopened, labeled ]
+      - '.github/actions/setup-runpod/**'
+    types: [ labeled ]
 
 concurrency:
-  group: gpu-ci-serial
+  group: hardware-gpu-ci
   cancel-in-progress: false
 
+permissions:
+  contents: read
 
 jobs:
-  gpu-tests:
-    if: contains(github.event.pull_request.labels.*.name, 'needs-gpu-ci')
+  # ── Option A: RunPod CUDA runner ──────────────────────────────────────────
+  # Triggered by the 'needs-gpu-ci' label. Runs ci/run_cuda_tests.sh on an
+  # ephemeral RunPod pod and removes the pod afterwards.
+  cuda-cloud:
+    if: github.event.action == 'labeled' && github.event.label.name == 'needs-gpu-ci'
     runs-on: ubuntu-latest
     timeout-minutes: 60
+    strategy:
+      fail-fast: false
+      max-parallel: 1
+      matrix:
+        include:
+          - name: cuda-a4000-tp2
+            gpu_id: NVIDIA RTX A4000
+            gpu_count: 2
+            fallback_gpu_id: NVIDIA A40
+            fallback_gpu_count: 1
+            torch_cuda_arch_list: "8.6"
+          - name: cuda-a40-tp1
+            gpu_id: NVIDIA A40
+            gpu_count: 1
+            fallback_gpu_id: NVIDIA RTX A4000
+            fallback_gpu_count: 1
+            torch_cuda_arch_list: "8.6"
     steps:
+      # The checkout uses the base SHA, so both the orchestrator script and the
+      # local setup-runpod action come from trusted (already-merged) code.
       - name: Checkout secure orchestrator script from base branch
         uses: actions/checkout@v4
         with:
           ref: ${{ github.event.pull_request.base.sha }}
+          persist-credentials: false
 
-      - name: Install runpodctl
-        run: |
-          wget -qO runpodctl https://github.com/runpod/runpodctl/releases/latest/download/runpodctl-linux-amd64
-          chmod +x runpodctl
-          sudo mv runpodctl /usr/local/bin/runpodctl
-          runpodctl version
+      - name: Setup RunPod CLI and SSH key
+        uses: ./.github/actions/setup-runpod
+        with:
+          api-key: ${{ secrets.RUNPOD_API_KEY }}
+          ssh-private-key: ${{ secrets.RUNPOD_SSH_PRIVATE_KEY }}
+
+      - name: Run CUDA hardware tests
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          RUNPOD_GPU_ID: ${{ matrix.gpu_id }}
+          RUNPOD_GPU_COUNT: ${{ matrix.gpu_count }}
+          RUNPOD_FALLBACK_GPU_ID: ${{ matrix.fallback_gpu_id }}
+          RUNPOD_FALLBACK_GPU_COUNT: ${{ matrix.fallback_gpu_count }}
+          RUNPOD_PROFILE_NAME: ${{ matrix.name }}
+          TORCH_CUDA_ARCH_LIST: ${{ matrix.torch_cuda_arch_list }}
+          PR_REPO_URL: ${{ github.event.pull_request.head.repo.clone_url }}
+          PR_SHA: ${{ github.event.pull_request.head.sha }}
+        run: bash ci/run_gpu_ci.sh
+
+  # ── Option B: self-hosted CUDA runner (no cloud account needed) ───────────
+  # Triggered by the 'needs-gpu-ci-self-hosted' label. The runner must have
+  # CUDA drivers installed and be registered with the label 'cuda'.
+  cuda-self-hosted:
+    if: github.event.action == 'labeled' && github.event.label.name == 'needs-gpu-ci-self-hosted'
+    runs-on: [ self-hosted, linux, x64, cuda ]
+    timeout-minutes: 90
+    steps:
+      # Check out trusted CI scripts from the base branch. The test script still
+      # executes the PR commit on this runner, so only apply this label to PRs
+      # that are allowed to run on the configured CUDA hardware.
+      - name: Checkout secure CI scripts from base branch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.base.sha }}
+          persist-credentials: false
+
+      - name: Run CUDA tests
+        env:
+          PR_REPO_URL: ${{ github.event.pull_request.head.repo.clone_url }}
+          PR_SHA: ${{ github.event.pull_request.head.sha }}
+        run: bash ci/run_cuda_tests.sh
 
-      - name: Configure runpodctl
-        run: runpodctl config --apiKey "${{ secrets.RUNPOD_API_KEY }}"
+  # ── Option C: RunPod CUDA SM90 runner ─────────────────────────────────────
+  # Triggered separately because H100 capacity is expensive and not needed for
+  # every GPU label. This validates the Hopper/TMA build path with SM90 enabled.
+  cuda-sm90-cloud:
+    if: github.event.action == 'labeled' && github.event.label.name == 'needs-gpu-ci-sm90'
+    runs-on: ubuntu-latest
+    timeout-minutes: 90
+    steps:
+      - name: Checkout secure orchestrator script from base branch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.base.sha }}
+          persist-credentials: false
 
-      - name: Setup SSH key
-        run: |
-          mkdir -p ~/.ssh && chmod 700 ~/.ssh
-          printf '%s\n' "${{ secrets.RUNPOD_SSH_PRIVATE_KEY }}" > ~/.ssh/id_ed25519
-          chmod 600 ~/.ssh/id_ed25519
-          ssh-keygen -y -f ~/.ssh/id_ed25519 > /dev/null && echo "key OK" || echo "key BROKEN"
+      - name: Setup RunPod CLI and SSH key
+        uses: ./.github/actions/setup-runpod
+        with:
+          api-key: ${{ secrets.RUNPOD_API_KEY }}
+          ssh-private-key: ${{ secrets.RUNPOD_SSH_PRIVATE_KEY }}
 
-      - name: Run GPU tests on RunPod
+      - name: Run SM90 CUDA hardware tests
         env:
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          RUNPOD_GPU_ID: NVIDIA H100 80GB HBM3
+          RUNPOD_GPU_COUNT: 1
+          RUNPOD_FALLBACK_GPU_ID: NVIDIA H100 NVL
+          RUNPOD_FALLBACK_GPU_COUNT: 1
+          RUNPOD_PROFILE_NAME: cuda-h100-sm90
+          TORCH_CUDA_ARCH_LIST: auto
+          KERNEL_ALIGN_FORCE_SM90: "1"
           PR_REPO_URL: ${{ github.event.pull_request.head.repo.clone_url }}
           PR_SHA: ${{ github.event.pull_request.head.sha }}
         run: bash ci/run_gpu_ci.sh
+
+  # ── Option D: self-hosted ROCm runner ─────────────────────────────────────
+  rocm-self-hosted:
+    if: github.event.action == 'labeled' && github.event.label.name == 'needs-rocm-ci'
+    runs-on: [ self-hosted, linux, x64, rocm ]
+    timeout-minutes: 90
+    steps:
+      # Check out trusted CI scripts from the base branch. The test script still
+      # executes the PR commit on this runner, so only apply this label to PRs
+      # that are allowed to run on the configured ROCm hardware.
+      - name: Checkout secure CI scripts from base branch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.base.sha }}
+          persist-credentials: false
+
+      - name: Run ROCm hardware tests
+        env:
+          RL_KERNEL_ROCM_ATTN_BACKEND: sdpa
+          PR_REPO_URL: ${{ github.event.pull_request.head.repo.clone_url }}
+          PR_SHA: ${{ github.event.pull_request.head.sha }}
+        run: bash ci/run_rocm_ci.sh