Add multi-cuda wheel build (#289)

Co-authored-by: Ashwin Srinath <shwina@users.noreply.github.com>
Co-authored-by: Nader Al Awar <naderalawar@gmail.com>

.github/workflows/build-and-test-python-wheels.yml

@@ -9,9 +9,8 @@ defaults:
     shell: bash --noprofile --norc -euo pipefail {0}
 
 jobs:
-  # Build wheels for all CUDA/Python combinations
   build-wheels:
-    name: Build wheel (CUDA ${{ matrix.cuda }}, Python ${{ matrix.python }})
+    name: Build wheel (Python ${{ matrix.python }})
     runs-on: ubuntu-latest
     permissions:
       id-token: write
@@ -19,7 +18,6 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        cuda: ['12', '13']
         python: ['3.10', '3.11', '3.12', '3.13']
 
     steps:
@@ -32,19 +30,20 @@ jobs:
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
 
-      - name: Build wheel
+      - name: Build multi-CUDA wheel
         run: |
-          bash ci/build_pynvbench_wheel.sh -py-version ${{ matrix.python }} -cuda-version ${{ matrix.cuda }}
+          bash ci/build_multi_cuda_wheel.sh -py-version ${{ matrix.python }}
 
       - name: Upload wheel artifact
         uses: actions/upload-artifact@v4
         with:
-          name: wheel-pynvbench-cu${{ matrix.cuda }}-py${{ matrix.python }}
+          name: wheel-pynvbench-py${{ matrix.python }}
           path: wheelhouse/*.whl
           retention-days: 7
           if-no-files-found: error
 
-  # Test wheels for all CUDA/Python combinations
+  # Test wheels on all CUDA/Python combinations
+  # Each wheel contains extensions for both CUDA 12 & 13, runtime detection picks the right one
   test-wheels:
     name: Test wheel (CUDA ${{ matrix.cuda }}, Python ${{ matrix.python }})
     needs: build-wheels
@@ -68,7 +67,7 @@ jobs:
       - name: Download wheel artifact
         uses: actions/download-artifact@v4
         with:
-          name: wheel-pynvbench-cu${{ matrix.cuda }}-py${{ matrix.python }}
+          name: wheel-pynvbench-py${{ matrix.python }}
           path: wheelhouse
 
       - name: Test wheel
@@ -76,8 +75,10 @@ jobs:
           # Use the same rapidsai/ci-wheel Docker image as build
           if [[ "${{ matrix.cuda }}" == "12" ]]; then
             cuda_full_version="12.9.1"
+            cuda_extra="cu12"
           else
             cuda_full_version="13.0.1"
+            cuda_extra="cu13"
           fi
 
           docker run --rm \
@@ -86,6 +87,7 @@ jobs:
             --mount type=bind,source=$(pwd),target=/workspace/ \
             --env py_version=${{ matrix.python }} \
             --env cuda_version=${{ matrix.cuda }} \
+            --env cuda_extra="${cuda_extra}" \
             rapidsai/ci-wheel:25.12-cuda${cuda_full_version}-rockylinux8-py${{ matrix.python }} \
             /workspace/ci/test_pynvbench_inner.sh
 

.github/workflows/build-wheels.yml

@@ -1,58 +0,0 @@
-name: Build Python Wheels (Manual)
-
-on:
-  workflow_dispatch:
-  workflow_call:
-
-defaults:
-  run:
-    shell: bash --noprofile --norc -euo pipefail {0}
-
-jobs:
-  build-wheels:
-    name: Build wheel (CUDA ${{ matrix.cuda }}, Python ${{ matrix.python }})
-    runs-on: ubuntu-latest
-    permissions:
-      id-token: write
-      contents: read
-    strategy:
-      fail-fast: false
-      matrix:
-        cuda: ['12', '13']
-        python: ['3.10', '3.11', '3.12', '3.13']
-
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
-          persist-credentials: false
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-
-      - name: Build wheel
-        run: |
-          bash ci/build_pynvbench_wheel.sh -py-version ${{ matrix.python }} -cuda-version ${{ matrix.cuda }}
-
-      - name: Upload wheel artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: wheel-pynvbench-cu${{ matrix.cuda }}-py${{ matrix.python }}
-          path: wheelhouse/*.whl
-          retention-days: 7
-          if-no-files-found: error
-
-  verify-wheels:
-    name: Verify all wheels built successfully
-    if: ${{ always() }}
-    needs: build-wheels
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check build results
-        run: |
-          if [[ "${{ needs.build-wheels.result }}" != "success" ]]; then
-            echo "Wheel builds failed!"
-            exit 1
-          fi
-          echo "All wheels built successfully!"