From 676a4580a6fcf980556f3f0134648c1c9b07ce09 Mon Sep 17 00:00:00 2001
From: Mario Sieg <mario.sieg.64@gmail.com>
Date: Thu, 27 Aug 2026 19:48:26 +0000
Subject: [PATCH 1/9] Add full wheel compilation pipeline

---
 .github/workflows/build_kernels.yaml | 273 ++++++++++++++++++++++-----
 scripts/build_torchao.sh             |  58 ++++++
 scripts/install_deep_gemm.sh         |  62 ++++++
 scripts/install_ep_kernels.sh        | 137 ++++++++++++++
 scripts/smoke_test.py                |  91 +++++++++
 5 files changed, 572 insertions(+), 49 deletions(-)
 create mode 100755 scripts/build_torchao.sh
 create mode 100755 scripts/install_deep_gemm.sh
 create mode 100755 scripts/install_ep_kernels.sh
 create mode 100755 scripts/smoke_test.py

diff --git a/.github/workflows/build_kernels.yaml b/.github/workflows/build_kernels.yaml
index 83cb605..8dfa7b2 100644
--- a/.github/workflows/build_kernels.yaml
+++ b/.github/workflows/build_kernels.yaml
@@ -1,96 +1,181 @@
-name: Build Kernels
+name: Build and Release Kernels
 
-# Builds the `prime-kernels` wheel for x86_64 and aarch64. Runs on every push/PR that touches
-# the kernel sources to catch build breakage, and with `release_tag` to attach the wheels to
-# a GitHub release — which is how consumers (e.g. prime-rl's `kernels` extra) install them.
+# Builds every wheel prime-rl (and other consumers) pin via `[tool.uv.sources]`:
+# prime-kernels (this repo's own kernels.toml-declared extensions), deep-ep and
+# deep-gemm (from their pinned deepseek-ai revs), and our torchao (from the pytorch/ao
+# rev the MXFP8 training path requires). This repo is the single build+release site for
+# all four — prime-rl only consumes the resulting release, it does not build anything
+# itself (see scripts/install_ep_kernels.sh, scripts/install_deep_gemm.sh,
+# scripts/build_torchao.sh for the pinned revs).
 #
-# No GPU is needed: nvcc cross compiles for the architectures kernels.toml declares.
+# Matrix: {kernel} x {cuda: cu129, cu130} x {arch: x86_64, aarch64} — the two CUDA
+# variants are the ones torch 2.13 actually publishes that also meet DeepGEMM/DeepEP's
+# own "CUDA 12.8+" floor for Blackwell (sm_100a) compilation; cu126 (also published for
+# torch 2.13) is excluded because it's below that floor. TORCH_CUDA_ARCH_LIST is set
+# explicitly per kernel so every wheel covers Hopper + Blackwell + Blackwell Ultra — the
+# install scripts fall back to auto-detecting a single arch from `nvidia-smi` when this
+# is unset, which silently scopes the wheel to whatever GPU happened to build it (that
+# bug produced a Blackwell-only deep-ep wheel once already).
+#
+# No GPU is needed to build: nvcc cross compiles for every architecture we ship. GPU
+# smoke tests run separately, against the built wheels, on an actual GPU runner.
+#
+# Runs on every push to main and on PRs that touch the matrix inputs; `workflow_dispatch`
+# reruns it by hand (e.g. after a torch/CUDA bump, or to rebuild one rev).
 
 on:
-  pull_request:
+  push:
+    branches: [main]
     paths:
       - "prime_kernels/**"
-      - "setup.py"
+      - "scripts/install_ep_kernels.sh"
+      - "scripts/install_deep_gemm.sh"
+      - "scripts/build_torchao.sh"
+      - "scripts/smoke_test.py"
       - "pyproject.toml"
       - ".github/workflows/build_kernels.yaml"
-  push:
-    branches: [main]
+  pull_request:
     paths:
       - "prime_kernels/**"
-      - "setup.py"
+      - "scripts/install_ep_kernels.sh"
+      - "scripts/install_deep_gemm.sh"
+      - "scripts/build_torchao.sh"
+      - "scripts/smoke_test.py"
       - "pyproject.toml"
       - ".github/workflows/build_kernels.yaml"
   workflow_dispatch:
     inputs:
       release_tag:
-        description: "Existing release tag to attach the wheels to (e.g. v0.2.0). Empty = artifacts only."
-        required: false
-        type: string
-      ref:
-        description: "Ref to build the kernels from. Empty = the ref the workflow was started on."
+        description: "Tag for the release (e.g. v0.2.0). Empty = auto-generate from pyproject.toml version + short SHA."
         required: false
         type: string
 
-env:
-  # The CUDA toolkit wheels are built against. Bump alongside the minimum torch/CUDA support
-  # this repo targets.
-  CUDA_TAG: cu128
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 jobs:
   build:
-    name: Build wheel (${{ matrix.arch }})
-    # Uploading the wheels to a release needs write.
-    permissions:
-      contents: write
+    name: Build ${{ matrix.kernel }} (${{ matrix.cuda }}, ${{ matrix.arch }})
     strategy:
       fail-fast: false
       matrix:
+        kernel: [prime-kernels, deep-ep, deep-gemm, torchao]
+        cuda: [cu129, cu130]
+        arch: [x86_64, aarch64]
         include:
-          - runner: image-builder
-            arch: x86_64
-          - runner: image-builder-arm-2204
+          - arch: x86_64
+            runner: image-builder
+          - arch: aarch64
+            runner: image-builder-arm-2204
+          - cuda: cu129
+            container_image: nvidia/cuda:12.9.1-devel-ubuntu24.04
+            torch_version: 2.13.0+cu129
+          - cuda: cu130
+            container_image: nvidia/cuda:13.0.3-devel-ubuntu24.04
+            torch_version: 2.13.0+cu130
+        exclude:
+          # torchao is consumed on x86_64 only (see its marker in prime-rl's
+          # pyproject.toml); aarch64 resolves the PyPI release there instead.
+          - kernel: torchao
             arch: aarch64
     runs-on: ${{ matrix.runner }}
-    container: nvidia/cuda:12.8.1-devel-ubuntu24.04
+    container: ${{ matrix.container_image }}
     steps:
+      # libibverbs-dev/librdmacm-dev: NVSHMEM's IBGDA transport headers (infiniband/mlx5dv.h),
+      # compile-time only — deep-ep dlopens the real thing at runtime.
       - name: Install build tooling
         run: |
           apt-get update
-          apt-get install -y --no-install-recommends git curl ca-certificates build-essential gh
+          apt-get install -y --no-install-recommends git curl ca-certificates build-essential gh xz-utils libibverbs-dev librdmacm-dev
 
       - name: Checkout
         uses: actions/checkout@v5
-        with:
-          ref: ${{ inputs.ref }}
 
       - name: Install uv
         run: |
           curl -LsSf https://astral.sh/uv/install.sh | env INSTALLER_NO_MODIFY_PATH=1 UV_INSTALL_DIR=/usr/local/bin sh
           uv --version
 
+      # Every kernel here is a torch C++ extension, bound to the exact torch (and thus
+      # CUDA) it was compiled against.
       - name: Create build environment
         env:
-          UV_TORCH_BACKEND: ${{ env.CUDA_TAG }}
+          UV_TORCH_BACKEND: ${{ matrix.cuda }}
         run: |
+          echo "Building against torch==${{ matrix.torch_version }}"
           uv venv --python 3.12 /tmp/build-env
-          VIRTUAL_ENV=/tmp/build-env uv pip install "torch>=2.9.0" setuptools wheel ninja
+          VIRTUAL_ENV=/tmp/build-env uv pip install "torch==${{ matrix.torch_version }}" setuptools wheel ninja cmake
 
-      # torch C++ extensions are bound to the torch and CUDA they were compiled against, so
-      # the wheel carries both in its local version.
+      # prime-kernels' own version carries the ABI it was built against (deep-ep,
+      # deep-gemm and torchao keep their upstream `+<rev>` naming; the release tag they
+      # are attached to disambiguates the ABI for those three).
       - name: Stamp ABI into the version
+        if: matrix.kernel == 'prime-kernels'
         run: |
           TORCH_VERSION=$(/tmp/build-env/bin/python -c 'import torch; print(torch.__version__.split("+")[0])')
-          LOCAL="${CUDA_TAG}torch${TORCH_VERSION}"
+          LOCAL="${{ matrix.cuda }}torch${TORCH_VERSION}"
           sed -i -E "s/^version = \"([^\"]+)\"$/version = \"\1+${LOCAL}\"/" pyproject.toml
           grep '^version' pyproject.toml
 
-      - name: Build wheel
+      - name: Build prime-kernels wheel
+        if: matrix.kernel == 'prime-kernels'
         env:
           VIRTUAL_ENV: /tmp/build-env
           CUDA_HOME: /usr/local/cuda
           PRIME_KERNELS_REQUIRE: "1"
           MAX_JOBS: "8"
-        run: uv build --wheel --no-build-isolation --out-dir dist
+        run: uv build --wheel --no-build-isolation --out-dir dist .
+
+      - name: Build deep-ep wheel
+        if: matrix.kernel == 'deep-ep'
+        env:
+          VIRTUAL_ENV: /tmp/build-env
+          TORCH_CUDA_ARCH_LIST: "9.0;10.0;10.3"
+          MAX_JOBS: "8"
+        run: |
+          export PATH="/tmp/build-env/bin:$PATH"
+          bash scripts/install_ep_kernels.sh --workspace /tmp/ep_kernels_workspace --wheel-dir "$PWD/dist"
+
+      - name: Build deep-gemm wheel
+        if: matrix.kernel == 'deep-gemm'
+        env:
+          VIRTUAL_ENV: /tmp/build-env
+          CUDA_HOME: /usr/local/cuda
+          MAX_JOBS: "8"
+        run: |
+          export PATH="/tmp/build-env/bin:$PATH"
+          bash scripts/install_deep_gemm.sh --wheel-dir dist
+
+      - name: Build torchao wheel
+        if: matrix.kernel == 'torchao'
+        env:
+          VIRTUAL_ENV: /tmp/build-env
+          CUDA_HOME: /usr/local/cuda
+          TORCH_CUDA_ARCH_LIST: "9.0a;10.0a"
+          MAX_JOBS: "8"
+        run: |
+          export PATH="/tmp/build-env/bin:$PATH"
+          bash scripts/build_torchao.sh --wheel-dir dist
+
+      # deep-ep/deep-gemm/torchao wheel names carry only their upstream `+<rev>` version —
+      # nothing distinguishes a cu129 build from a cu130 build of the same rev. Building
+      # both into the same release would otherwise clobber one with the other on upload
+      # (identical filename). prime-kernels doesn't need this: its own version already
+      # got the CUDA tag stamped in above.
+      - name: Disambiguate wheel filename by CUDA toolkit
+        if: matrix.kernel != 'prime-kernels'
+        run: |
+          /tmp/build-env/bin/python - <<'PY'
+          import pathlib
+          wheel = next(pathlib.Path("dist").glob("*.whl"))
+          name, version, rest = wheel.name.split("-", 2)
+          cuda_tag = "${{ matrix.cuda }}"
+          version += f".{cuda_tag}" if "+" in version else f"+{cuda_tag}"
+          new_name = f"{name}-{version}-{rest}"
+          wheel.rename(wheel.with_name(new_name))
+          print(f"{wheel.name} -> {new_name}")
+          PY
 
       - name: Inspect wheel
         run: |
@@ -106,21 +191,111 @@ jobs:
       - name: Upload artifact
         uses: actions/upload-artifact@v4
         with:
-          name: prime-kernels-${{ matrix.arch }}
+          name: ${{ matrix.kernel }}-${{ matrix.cuda }}-${{ matrix.arch }}
           path: dist/*.whl
           if-no-files-found: error
 
-      - name: Attach to release
-        if: inputs.release_tag != ''
+  smoke-test:
+    name: Smoke test ${{ matrix.kernel }} (${{ matrix.cuda }})
+    needs: build
+    # x86_64 only: GPU runners registered for this org today are x86_64 (see the `vm`
+    # label prime-rl's gpu_tests.yaml uses). If an aarch64 GPU runner (e.g. GB200)
+    # becomes available, add an `arch` axis here the same way `build` has one.
+    strategy:
+      fail-fast: false
+      matrix:
+        kernel: [prime-kernels, deep-ep, deep-gemm, torchao]
+        cuda: [cu129, cu130]
+        include:
+          - cuda: cu129
+            container_image: nvidia/cuda:12.9.1-devel-ubuntu24.04
+            torch_version: 2.13.0+cu129
+          - cuda: cu130
+            container_image: nvidia/cuda:13.0.3-devel-ubuntu24.04
+            torch_version: 2.13.0+cu130
+    runs-on: vm
+    container:
+      image: ${{ matrix.container_image }}
+      options: --gpus all
+    steps:
+      - name: Install runtime tooling
+        run: |
+          apt-get update
+          apt-get install -y --no-install-recommends curl ca-certificates
+
+      - name: Checkout
+        uses: actions/checkout@v5
+
+      - name: Install uv
+        run: |
+          curl -LsSf https://astral.sh/uv/install.sh | env INSTALLER_NO_MODIFY_PATH=1 UV_INSTALL_DIR=/usr/local/bin sh
+
+      - name: Download wheel
+        uses: actions/download-artifact@v4
+        with:
+          name: ${{ matrix.kernel }}-${{ matrix.cuda }}-x86_64
+          path: dist
+
+      - name: Install wheel + matching torch
+        env:
+          UV_TORCH_BACKEND: ${{ matrix.cuda }}
+        run: |
+          uv venv --python 3.12 /tmp/test-env
+          VIRTUAL_ENV=/tmp/test-env uv pip install "torch==${{ matrix.torch_version }}"
+          VIRTUAL_ENV=/tmp/test-env uv pip install dist/*.whl
+
+      - name: Run smoke test
+        run: /tmp/test-env/bin/python scripts/smoke_test.py ${{ matrix.kernel }}
+
+  release:
+    name: Publish release
+    needs: [build, smoke-test]
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v5
+
+      - name: Determine release tag
+        id: tag
+        run: |
+          if [ -n "${{ inputs.release_tag }}" ]; then
+            TAG="${{ inputs.release_tag }}"
+          else
+            VERSION=$(grep -m1 '^version = ' pyproject.toml | sed -E 's/version = "(.*)"/\1/')
+            TAG="v${VERSION}-$(echo "${{ github.sha }}" | cut -c1-7)"
+          fi
+          echo "tag=$TAG" >> "$GITHUB_OUTPUT"
+
+      - name: Download all wheels
+        uses: actions/download-artifact@v4
+        with:
+          path: dist
+          pattern: "*"
+          merge-multiple: true
+
+      - name: Inspect release contents
+        run: ls -lh dist/
+
+      # Draft first, upload everything, publish last: if the upload step fails
+      # partway, the release stays a draft forever and no consumer ever sees a
+      # partial set of wheels — "atomic" from the outside.
+      - name: Create draft release
         env:
           GH_TOKEN: ${{ github.token }}
-          # gh must not resolve the repo from the checkout: inside the container the
-          # workspace is owned by another uid, which git rejects as dubious ownership.
-          GH_REPO: ${{ github.repository }}
-          TAG: ${{ inputs.release_tag }}
         run: |
-          gh release upload "$TAG" dist/*.whl --clobber
-          # Consumers' `[tool.uv.sources]` pin has to name the wheel exactly, ABI suffix and all.
-          WHEEL=$(basename dist/*.whl)
-          URL="https://github.com/${GITHUB_REPOSITORY}/releases/download/${TAG}/${WHEEL}"
-          echo "\`{ url = \"$URL\", marker = \"platform_machine == '${{ matrix.arch }}'\" },\`" >> "$GITHUB_STEP_SUMMARY"
+          gh release create "${{ steps.tag.outputs.tag }}" \
+            --draft \
+            --title "${{ steps.tag.outputs.tag }}" \
+            --notes "Kernel wheels built from $(git rev-parse HEAD)."
+
+      - name: Upload wheels
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: gh release upload "${{ steps.tag.outputs.tag }}" dist/*.whl --clobber
+
+      - name: Publish release
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: gh release edit "${{ steps.tag.outputs.tag }}" --draft=false
diff --git a/scripts/build_torchao.sh b/scripts/build_torchao.sh
new file mode 100755
index 0000000..021526a
--- /dev/null
+++ b/scripts/build_torchao.sh
@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+# Build our pinned torchao (pytorch/ao) as a wheel.
+#
+# Pinned to v0.18.0 (5f2baf9d) — the first tag that compiles against torch 2.13 /
+# CUDA 13: 0.17.0's `_C_cutlass_90a` extension hits a `STABLE_TORCH_LIBRARY_IMPL`
+# compile error there (default arguments on a function parameter). Bonus: 0.18.0
+# builds as `cp310-abi3` (torch stable ABI), so it should survive the next torch
+# bump without a rebuild.
+#
+# Usage:
+#   bash scripts/build_torchao.sh --wheel-dir dist
+#
+# Options:
+#   --ref REF        pytorch/ao commit hash (default: 5f2baf9d575cf732362594c998c399902942531f)
+#   --wheel-dir DIR  Output wheel to DIR (default: ./dist)
+#
+# Set TORCH_CUDA_ARCH_LIST to control the CUTLASS/MXFP8 arch list (default: "9.0a;10.0a" —
+# sm_90a enables the CUTLASS kernels, sm_100a the MXFP8 extension).
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(dirname "$SCRIPT_DIR")"
+
+TORCHAO_GIT_REF="5f2baf9d575cf732362594c998c399902942531f"
+WHEEL_DIR="$REPO_ROOT/dist"
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --ref)       TORCHAO_GIT_REF="$2"; shift 2 ;;
+        --wheel-dir) WHEEL_DIR="$2";       shift 2 ;;
+        *) echo "Unknown argument: $1" >&2; exit 1 ;;
+    esac
+done
+
+export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-9.0a;10.0a}"
+
+echo "================================================================"
+echo " Building torchao (${TORCHAO_GIT_REF:0:7})"
+echo " Arch list: ${TORCH_CUDA_ARCH_LIST}"
+echo "================================================================"
+
+TMPDIR=$(mktemp -d)
+cleanup() { rm -rf "$TMPDIR"; }
+trap cleanup EXIT
+
+git clone https://github.com/pytorch/ao "$TMPDIR/ao"
+git -C "$TMPDIR/ao" checkout "$TORCHAO_GIT_REF"
+git -C "$TMPDIR/ao" submodule update --init --recursive
+
+mkdir -p "$WHEEL_DIR"
+uv build --no-build-isolation --wheel --out-dir "$WHEEL_DIR" "$TMPDIR/ao"
+echo ""
+echo "Wheel built:"
+ls -lh "$WHEEL_DIR"/torchao*.whl
+echo "================================================================"
+echo " torchao build complete"
+echo "================================================================"
diff --git a/scripts/install_deep_gemm.sh b/scripts/install_deep_gemm.sh
new file mode 100755
index 0000000..7e48e05
--- /dev/null
+++ b/scripts/install_deep_gemm.sh
@@ -0,0 +1,62 @@
+#!/usr/bin/env bash
+# Build DeepGEMM (FP8 blockwise / grouped GEMM kernels) as a wheel.
+#
+# Ported from prime-rl's scripts/install_deep_gemm.sh: prime-kernels' own CI is the
+# canonical build site for this wheel now (see .github/workflows/build_and_release.yaml).
+#
+# Requires CUDA 12.8+ and a Hopper/Blackwell GPU (cross-compiles fine without one).
+#
+# Usage:
+#   bash scripts/install_deep_gemm.sh --wheel-dir dist
+#
+# Options:
+#   --ref REF        DeepGEMM commit hash (default: 891d57b)
+#   --wheel-dir DIR  Output wheel to DIR (default: ./dist)
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(dirname "$SCRIPT_DIR")"
+
+DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
+DEEPGEMM_GIT_REF="891d57b4db1071624b5c8fa0d1e51cb317fa709f"
+WHEEL_DIR="$REPO_ROOT/dist"
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --ref)       DEEPGEMM_GIT_REF="$2"; shift 2 ;;
+        --wheel-dir) WHEEL_DIR="$2";        shift 2 ;;
+        *) echo "Unknown argument: $1" >&2; exit 1 ;;
+    esac
+done
+
+SHORT_REF="${DEEPGEMM_GIT_REF:0:7}"
+
+CUDA_VERSION=$(nvcc --version 2>/dev/null | grep -oP 'release \K[\d.]+' || echo "")
+if [ -z "$CUDA_VERSION" ]; then
+    echo "ERROR: nvcc not found. CUDA toolkit required." >&2
+    exit 1
+fi
+
+echo "================================================================"
+echo " Building DeepGEMM (${SHORT_REF})"
+echo " CUDA: ${CUDA_VERSION}"
+echo "================================================================"
+
+TMPDIR=$(mktemp -d)
+cleanup() { rm -rf "$TMPDIR"; }
+trap cleanup EXIT
+
+git clone --recurse-submodules "$DEEPGEMM_GIT_REPO" "$TMPDIR/DeepGEMM"
+cd "$TMPDIR/DeepGEMM"
+git checkout "$DEEPGEMM_GIT_REF"
+git submodule update --init --recursive
+
+mkdir -p "$WHEEL_DIR"
+uv build --no-build-isolation --wheel --out-dir "$WHEEL_DIR" "$TMPDIR/DeepGEMM"
+echo ""
+echo "Wheel built:"
+ls -lh "$WHEEL_DIR"/deep_gemm*.whl
+echo "================================================================"
+echo " DeepGEMM build complete"
+echo "================================================================"
diff --git a/scripts/install_ep_kernels.sh b/scripts/install_ep_kernels.sh
new file mode 100755
index 0000000..fb21a50
--- /dev/null
+++ b/scripts/install_ep_kernels.sh
@@ -0,0 +1,137 @@
+#!/usr/bin/env bash
+# Build DeepEP (NVSHMEM-backed expert-parallel all-to-all kernels) as a wheel.
+#
+# Ported from prime-rl's scripts/install_ep_kernels.sh: prime-kernels' own CI is the
+# canonical build site for this wheel now (see .github/workflows/build_and_release.yaml),
+# so this script is self-contained rather than reaching into a sibling repo.
+#
+# Usage:
+#   bash scripts/install_ep_kernels.sh --wheel-dir dist
+#
+# Options:
+#   --workspace DIR       Build directory (default: ./ep_kernels_workspace)
+#   --wheel-dir DIR       Wheel output directory (default: ./dist)
+#   --deepep-ref REF      DeepEP commit hash (default: 29d31c0)
+#   --nvshmem-ver VER     NVSHMEM version (default: 3.3.24)
+#
+# Set TORCH_CUDA_ARCH_LIST to cross compile without a GPU (e.g. "9.0;10.0;10.3" in CI);
+# unset, the arch is detected from the GPU nvidia-smi reports.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(dirname "$SCRIPT_DIR")"
+
+DEEPEP_COMMIT_HASH="29d31c0"
+NVSHMEM_VER="3.3.24"
+WORKSPACE="$REPO_ROOT/ep_kernels_workspace"
+WHEEL_DIR="$REPO_ROOT/dist"
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --workspace)  WORKSPACE="$2";          shift 2 ;;
+        --wheel-dir)  WHEEL_DIR="$2";          shift 2 ;;
+        --deepep-ref) DEEPEP_COMMIT_HASH="$2"; shift 2 ;;
+        --nvshmem-ver) NVSHMEM_VER="$2";       shift 2 ;;
+        *) echo "Unknown argument: $1" >&2; exit 1 ;;
+    esac
+done
+
+# ── Auto-detect CUDA toolkit matching torch ───────────────────────────────────
+TORCH_CUDA_VER=$(python -c "import torch; print(torch.version.cuda)")
+CUDA_MAJOR_MINOR=$(echo "$TORCH_CUDA_VER" | grep -oP '^\d+\.\d+')
+CUDA_MAJOR=$(echo "$CUDA_MAJOR_MINOR" | cut -d. -f1)
+
+CUDA_HOME="/usr/local/cuda-${CUDA_MAJOR_MINOR}"
+if [ ! -x "$CUDA_HOME/bin/nvcc" ]; then
+    echo "ERROR: Could not find CUDA toolkit matching torch (cuda ${TORCH_CUDA_VER}) at ${CUDA_HOME}" >&2
+    exit 1
+fi
+export CUDA_HOME
+# CUDA 13 moved the CCCL headers (cuda/std/...) out of the main include dir.
+if [ -d "$CUDA_HOME/include/cccl" ]; then
+    export CPATH="$CUDA_HOME/include/cccl${CPATH:+:$CPATH}"
+fi
+
+NVCC_VER=$("$CUDA_HOME/bin/nvcc" --version | grep -oP 'release \K[\d.]+')
+echo "Torch CUDA: ${TORCH_CUDA_VER}, nvcc: ${NVCC_VER} (${CUDA_HOME})"
+
+# ── Auto-detect GPU architecture (honor a preset TORCH_CUDA_ARCH_LIST) ────────
+if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then
+    GPU_NAME=$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader,nounits -i 0 2>/dev/null)
+    case "$GPU_NAME" in
+        *H100*|*H200*|*H800*)   TORCH_CUDA_ARCH_LIST="9.0" ;;
+        *B100*|*B200*|*GB200*)  TORCH_CUDA_ARCH_LIST="10.0" ;;
+        *)
+            echo "Could not auto-detect GPU arch from '$GPU_NAME'. Set TORCH_CUDA_ARCH_LIST manually." >&2
+            exit 1
+            ;;
+    esac
+fi
+export TORCH_CUDA_ARCH_LIST
+
+echo "================================================================"
+echo " Building DeepEP kernels"
+echo " CUDA:       ${CUDA_HOME} (${NVCC_VER})"
+echo " GPU:        ${GPU_NAME:-cross-compile} (arch ${TORCH_CUDA_ARCH_LIST})"
+echo " DeepEP:     ${DEEPEP_COMMIT_HASH}"
+echo " NVSHMEM:    ${NVSHMEM_VER}"
+echo " Workspace:  ${WORKSPACE}"
+echo "================================================================"
+
+mkdir -p "$WORKSPACE"
+
+echo ""
+echo "--- Installing build dependencies ---"
+uv pip install cmake ninja
+
+echo ""
+echo "--- Setting up NVSHMEM ${NVSHMEM_VER} ---"
+
+ARCH=$(uname -m)
+case "${ARCH,,}" in
+    x86_64|amd64)  NVSHMEM_SUBDIR="linux-x86_64" ;;
+    aarch64|arm64) NVSHMEM_SUBDIR="linux-sbsa" ;;
+    *) echo "Unsupported architecture: ${ARCH}" >&2; exit 1 ;;
+esac
+
+NVSHMEM_DIR="$WORKSPACE/nvshmem"
+if [ ! -d "$NVSHMEM_DIR/lib" ]; then
+    NVSHMEM_FILE="libnvshmem-${NVSHMEM_SUBDIR}-${NVSHMEM_VER}_cuda${CUDA_MAJOR}-archive.tar.xz"
+    NVSHMEM_URL="https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/${NVSHMEM_SUBDIR}/${NVSHMEM_FILE}"
+
+    echo "Downloading ${NVSHMEM_URL}"
+    curl -fSL "${NVSHMEM_URL}" -o "$WORKSPACE/${NVSHMEM_FILE}"
+    tar -xf "$WORKSPACE/${NVSHMEM_FILE}" -C "$WORKSPACE"
+    mv "$WORKSPACE/${NVSHMEM_FILE%.tar.xz}" "$NVSHMEM_DIR"
+    rm -f "$WORKSPACE/${NVSHMEM_FILE}"
+    rm -rf "$NVSHMEM_DIR/lib/bin" "$NVSHMEM_DIR/lib/share"
+    echo "NVSHMEM extracted to ${NVSHMEM_DIR}"
+else
+    echo "NVSHMEM already present at ${NVSHMEM_DIR}, skipping download"
+fi
+
+export CMAKE_PREFIX_PATH="${NVSHMEM_DIR}/lib/cmake:${CMAKE_PREFIX_PATH:-}"
+export NVSHMEM_DIR
+
+echo ""
+echo "--- Building DeepEP (${DEEPEP_COMMIT_HASH}) ---"
+
+DEEPEP_DIR="$WORKSPACE/DeepEP"
+if [ ! -d "$DEEPEP_DIR/.git" ]; then
+    git clone https://github.com/deepseek-ai/DeepEP "$DEEPEP_DIR"
+fi
+
+cd "$DEEPEP_DIR"
+git fetch origin
+git checkout "$DEEPEP_COMMIT_HASH"
+
+mkdir -p "$WHEEL_DIR"
+python setup.py bdist_wheel --dist-dir "$WHEEL_DIR"
+
+WHEEL=$(ls "$WHEEL_DIR"/deep_ep*.whl | head -1)
+echo ""
+echo "--- DeepEP wheel built at: $WHEEL ---"
+echo "================================================================"
+echo " DeepEP build complete"
+echo "================================================================"
diff --git a/scripts/smoke_test.py b/scripts/smoke_test.py
new file mode 100755
index 0000000..52d06fc
--- /dev/null
+++ b/scripts/smoke_test.py
@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+"""GPU smoke tests run against a just-built kernel wheel in CI.
+
+Each check does real device work (not just an import) so a wheel built for the wrong
+arch, or one that links against the wrong torch ABI, fails here instead of silently
+shipping. Not a correctness suite — that lives in each kernel's own repo (or
+prime-rl's tests/unit/train/models/test_fused_moe.py for flash_moe); this only
+guards the thing cross-compilation can't verify on its own: does the compiled
+extension actually load and run on this GPU.
+
+Usage: python smoke_test.py <prime-kernels|deep-ep|deep-gemm|torchao>
+"""
+
+import sys
+
+import torch
+
+assert torch.cuda.is_available(), "smoke tests require a GPU"
+
+CAP = torch.cuda.get_device_capability()
+CAP_STR = f"sm_{CAP[0]}{CAP[1]}"
+print(f"Device: {torch.cuda.get_device_name()} ({CAP_STR}), torch {torch.__version__}")
+
+
+def smoke_prime_kernels():
+    import prime_kernels
+
+    status = prime_kernels.status()
+    print("prime_kernels.status():", status)
+    if CAP < (10, 0):
+        print(f"flash_moe is Blackwell-only (sm_100a); skipping on {CAP_STR}")
+        return
+    assert status.get("flash_moe") == "available", status
+    flash_moe = prime_kernels.load("flash_moe")
+    print("flash_moe loaded:", flash_moe)
+
+
+def smoke_deep_gemm():
+    import deep_gemm
+
+    print("deep_gemm", deep_gemm.__version__)
+    m, k, n = 128, 512, 512
+    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    out = torch.empty(m, n, device="cuda", dtype=torch.bfloat16)
+    deep_gemm.bf16_gemm_nt(a, b, out)
+    err = (out.float() - (a @ b.T).float()).abs().max().item()
+    print("bf16 gemm max err:", err)
+    assert err < 1.0, f"deep_gemm output diverges: max err {err}"
+
+
+def smoke_deep_ep():
+    import deep_ep
+
+    print("deep_ep", deep_ep.__file__)
+    # Real dispatch/combine needs a multi-GPU NVSHMEM group; out of scope for a
+    # single-process smoke test. Import + extension load is what cross-compilation
+    # can get wrong (wrong arch, wrong torch ABI), so that's what this checks.
+
+
+def smoke_torchao():
+    import torchao  # noqa: F401
+    from torchao.prototype.mx_formats.mx_tensor import MXTensor
+
+    print("torchao", torchao.__version__)
+    if CAP < (10, 0):
+        print(f"MXFP8 is Blackwell-only (sm_100a); skipping on {CAP_STR}")
+        return
+    x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16)
+    mx = MXTensor.to_mx(x, elem_dtype=torch.float8_e4m3fn, block_size=32)
+    print("MXTensor.to_mx OK:", mx.qdata.shape, mx.qdata.dtype)
+
+
+CHECKS = {
+    "prime-kernels": smoke_prime_kernels,
+    "deep-gemm": smoke_deep_gemm,
+    "deep-ep": smoke_deep_ep,
+    "torchao": smoke_torchao,
+}
+
+
+def main():
+    if len(sys.argv) != 2 or sys.argv[1] not in CHECKS:
+        print(f"Usage: {sys.argv[0]} <{'|'.join(CHECKS)}>", file=sys.stderr)
+        sys.exit(1)
+    CHECKS[sys.argv[1]]()
+    print("OK")
+
+
+if __name__ == "__main__":
+    main()

From bedb6768ad1a4e78e92bb643972c8e8be32c3f3e Mon Sep 17 00:00:00 2001
From: Mario Sieg <mario.sieg.64@gmail.com>
Date: Thu, 27 Aug 2026 19:55:19 +0000
Subject: [PATCH 2/9] Remove cuda 12.9

---
 .github/workflows/build_kernels.yaml | 80 ++++++++++------------------
 1 file changed, 28 insertions(+), 52 deletions(-)

diff --git a/.github/workflows/build_kernels.yaml b/.github/workflows/build_kernels.yaml
index 8dfa7b2..3e17b29 100644
--- a/.github/workflows/build_kernels.yaml
+++ b/.github/workflows/build_kernels.yaml
@@ -8,10 +8,8 @@ name: Build and Release Kernels
 # itself (see scripts/install_ep_kernels.sh, scripts/install_deep_gemm.sh,
 # scripts/build_torchao.sh for the pinned revs).
 #
-# Matrix: {kernel} x {cuda: cu129, cu130} x {arch: x86_64, aarch64} — the two CUDA
-# variants are the ones torch 2.13 actually publishes that also meet DeepGEMM/DeepEP's
-# own "CUDA 12.8+" floor for Blackwell (sm_100a) compilation; cu126 (also published for
-# torch 2.13) is excluded because it's below that floor. TORCH_CUDA_ARCH_LIST is set
+# Matrix: {kernel} x {arch: x86_64, aarch64}, all built against CUDA 13.0 (matches the
+# driver on our Blackwell fleet; see CUDA_TAG below). TORCH_CUDA_ARCH_LIST is set
 # explicitly per kernel so every wheel covers Hopper + Blackwell + Blackwell Ultra — the
 # install scripts fall back to auto-detecting a single arch from `nvidia-smi` when this
 # is unset, which silently scopes the wheel to whatever GPU happened to build it (that
@@ -20,8 +18,10 @@ name: Build and Release Kernels
 # No GPU is needed to build: nvcc cross compiles for every architecture we ship. GPU
 # smoke tests run separately, against the built wheels, on an actual GPU runner.
 #
-# Runs on every push to main and on PRs that touch the matrix inputs; `workflow_dispatch`
-# reruns it by hand (e.g. after a torch/CUDA bump, or to rebuild one rev).
+# Runs on every push to main and on PRs that touch the matrix inputs (build + smoke-test
+# only there — release is push-to-main / manual-dispatch only, see the `release` job's
+# `if`); `workflow_dispatch` reruns it by hand (e.g. after a torch/CUDA bump, or to
+# rebuild one rev).
 
 on:
   push:
@@ -54,33 +54,32 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
   cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
+env:
+  # Keep in sync with the CUDA toolkit the container images below use.
+  CUDA_TAG: cu130
+  CONTAINER_IMAGE: nvidia/cuda:13.0.3-devel-ubuntu24.04
+  TORCH_VERSION: 2.13.0+cu130
+
 jobs:
   build:
-    name: Build ${{ matrix.kernel }} (${{ matrix.cuda }}, ${{ matrix.arch }})
+    name: Build ${{ matrix.kernel }} (${{ matrix.arch }})
     strategy:
       fail-fast: false
       matrix:
         kernel: [prime-kernels, deep-ep, deep-gemm, torchao]
-        cuda: [cu129, cu130]
         arch: [x86_64, aarch64]
         include:
           - arch: x86_64
             runner: image-builder
           - arch: aarch64
             runner: image-builder-arm-2204
-          - cuda: cu129
-            container_image: nvidia/cuda:12.9.1-devel-ubuntu24.04
-            torch_version: 2.13.0+cu129
-          - cuda: cu130
-            container_image: nvidia/cuda:13.0.3-devel-ubuntu24.04
-            torch_version: 2.13.0+cu130
         exclude:
           # torchao is consumed on x86_64 only (see its marker in prime-rl's
           # pyproject.toml); aarch64 resolves the PyPI release there instead.
           - kernel: torchao
             arch: aarch64
     runs-on: ${{ matrix.runner }}
-    container: ${{ matrix.container_image }}
+    container: ${{ env.CONTAINER_IMAGE }}
     steps:
       # libibverbs-dev/librdmacm-dev: NVSHMEM's IBGDA transport headers (infiniband/mlx5dv.h),
       # compile-time only — deep-ep dlopens the real thing at runtime.
@@ -101,11 +100,11 @@ jobs:
       # CUDA) it was compiled against.
       - name: Create build environment
         env:
-          UV_TORCH_BACKEND: ${{ matrix.cuda }}
+          UV_TORCH_BACKEND: ${{ env.CUDA_TAG }}
         run: |
-          echo "Building against torch==${{ matrix.torch_version }}"
+          echo "Building against torch==${{ env.TORCH_VERSION }}"
           uv venv --python 3.12 /tmp/build-env
-          VIRTUAL_ENV=/tmp/build-env uv pip install "torch==${{ matrix.torch_version }}" setuptools wheel ninja cmake
+          VIRTUAL_ENV=/tmp/build-env uv pip install "torch==${{ env.TORCH_VERSION }}" setuptools wheel ninja cmake
 
       # prime-kernels' own version carries the ABI it was built against (deep-ep,
       # deep-gemm and torchao keep their upstream `+<rev>` naming; the release tag they
@@ -114,7 +113,7 @@ jobs:
         if: matrix.kernel == 'prime-kernels'
         run: |
           TORCH_VERSION=$(/tmp/build-env/bin/python -c 'import torch; print(torch.__version__.split("+")[0])')
-          LOCAL="${{ matrix.cuda }}torch${TORCH_VERSION}"
+          LOCAL="${{ env.CUDA_TAG }}torch${TORCH_VERSION}"
           sed -i -E "s/^version = \"([^\"]+)\"$/version = \"\1+${LOCAL}\"/" pyproject.toml
           grep '^version' pyproject.toml
 
@@ -158,25 +157,6 @@ jobs:
           export PATH="/tmp/build-env/bin:$PATH"
           bash scripts/build_torchao.sh --wheel-dir dist
 
-      # deep-ep/deep-gemm/torchao wheel names carry only their upstream `+<rev>` version —
-      # nothing distinguishes a cu129 build from a cu130 build of the same rev. Building
-      # both into the same release would otherwise clobber one with the other on upload
-      # (identical filename). prime-kernels doesn't need this: its own version already
-      # got the CUDA tag stamped in above.
-      - name: Disambiguate wheel filename by CUDA toolkit
-        if: matrix.kernel != 'prime-kernels'
-        run: |
-          /tmp/build-env/bin/python - <<'PY'
-          import pathlib
-          wheel = next(pathlib.Path("dist").glob("*.whl"))
-          name, version, rest = wheel.name.split("-", 2)
-          cuda_tag = "${{ matrix.cuda }}"
-          version += f".{cuda_tag}" if "+" in version else f"+{cuda_tag}"
-          new_name = f"{name}-{version}-{rest}"
-          wheel.rename(wheel.with_name(new_name))
-          print(f"{wheel.name} -> {new_name}")
-          PY
-
       - name: Inspect wheel
         run: |
           ls -lh dist/
@@ -191,12 +171,12 @@ jobs:
       - name: Upload artifact
         uses: actions/upload-artifact@v4
         with:
-          name: ${{ matrix.kernel }}-${{ matrix.cuda }}-${{ matrix.arch }}
+          name: ${{ matrix.kernel }}-${{ matrix.arch }}
           path: dist/*.whl
           if-no-files-found: error
 
   smoke-test:
-    name: Smoke test ${{ matrix.kernel }} (${{ matrix.cuda }})
+    name: Smoke test ${{ matrix.kernel }}
     needs: build
     # x86_64 only: GPU runners registered for this org today are x86_64 (see the `vm`
     # label prime-rl's gpu_tests.yaml uses). If an aarch64 GPU runner (e.g. GB200)
@@ -205,17 +185,9 @@ jobs:
       fail-fast: false
       matrix:
         kernel: [prime-kernels, deep-ep, deep-gemm, torchao]
-        cuda: [cu129, cu130]
-        include:
-          - cuda: cu129
-            container_image: nvidia/cuda:12.9.1-devel-ubuntu24.04
-            torch_version: 2.13.0+cu129
-          - cuda: cu130
-            container_image: nvidia/cuda:13.0.3-devel-ubuntu24.04
-            torch_version: 2.13.0+cu130
     runs-on: vm
     container:
-      image: ${{ matrix.container_image }}
+      image: ${{ env.CONTAINER_IMAGE }}
       options: --gpus all
     steps:
       - name: Install runtime tooling
@@ -233,15 +205,15 @@ jobs:
       - name: Download wheel
         uses: actions/download-artifact@v4
         with:
-          name: ${{ matrix.kernel }}-${{ matrix.cuda }}-x86_64
+          name: ${{ matrix.kernel }}-x86_64
           path: dist
 
       - name: Install wheel + matching torch
         env:
-          UV_TORCH_BACKEND: ${{ matrix.cuda }}
+          UV_TORCH_BACKEND: ${{ env.CUDA_TAG }}
         run: |
           uv venv --python 3.12 /tmp/test-env
-          VIRTUAL_ENV=/tmp/test-env uv pip install "torch==${{ matrix.torch_version }}"
+          VIRTUAL_ENV=/tmp/test-env uv pip install "torch==${{ env.TORCH_VERSION }}"
           VIRTUAL_ENV=/tmp/test-env uv pip install dist/*.whl
 
       - name: Run smoke test
@@ -250,6 +222,10 @@ jobs:
   release:
     name: Publish release
     needs: [build, smoke-test]
+    # Only cut a real release on push-to-main or an explicit manual dispatch — a PR
+    # should build and smoke-test, not publish. Without this, every push to a PR branch
+    # would create and publish a real release the moment build+smoke-test go green.
+    if: github.event_name != 'pull_request'
     runs-on: ubuntu-latest
     permissions:
       contents: write

From 120299461f1ac12af9a0d13a7c719c8aefe1b419 Mon Sep 17 00:00:00 2001
From: Mario Sieg <mario.sieg.64@gmail.com>
Date: Thu, 27 Aug 2026 19:58:50 +0000
Subject: [PATCH 3/9] Fix stale dir

---
 scripts/install_deep_gemm.sh | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/scripts/install_deep_gemm.sh b/scripts/install_deep_gemm.sh
index 7e48e05..3783b91 100755
--- a/scripts/install_deep_gemm.sh
+++ b/scripts/install_deep_gemm.sh
@@ -52,6 +52,10 @@ cd "$TMPDIR/DeepGEMM"
 git checkout "$DEEPGEMM_GIT_REF"
 git submodule update --init --recursive
 
+# Back to the invocation directory: a relative --wheel-dir must resolve there, not
+# inside $TMPDIR, which the EXIT trap deletes before the caller ever sees the wheel.
+cd "$REPO_ROOT"
+
 mkdir -p "$WHEEL_DIR"
 uv build --no-build-isolation --wheel --out-dir "$WHEEL_DIR" "$TMPDIR/DeepGEMM"
 echo ""

From 5061f3a6342c62e6b8d318f727fc7a47da3bb80d Mon Sep 17 00:00:00 2001
From: Mario Sieg <mario.sieg.64@gmail.com>
Date: Thu, 27 Aug 2026 20:01:20 +0000
Subject: [PATCH 4/9] Fix env

---
 .github/workflows/build_kernels.yaml | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/build_kernels.yaml b/.github/workflows/build_kernels.yaml
index 3e17b29..620b0ad 100644
--- a/.github/workflows/build_kernels.yaml
+++ b/.github/workflows/build_kernels.yaml
@@ -79,7 +79,9 @@ jobs:
           - kernel: torchao
             arch: aarch64
     runs-on: ${{ matrix.runner }}
-    container: ${{ env.CONTAINER_IMAGE }}
+    # `env` isn't available in `container:` (GitHub Actions only exposes it inside
+    # steps), so this has to be a literal — keep it in sync with CONTAINER_IMAGE above.
+    container: nvidia/cuda:13.0.3-devel-ubuntu24.04
     steps:
       # libibverbs-dev/librdmacm-dev: NVSHMEM's IBGDA transport headers (infiniband/mlx5dv.h),
       # compile-time only — deep-ep dlopens the real thing at runtime.
@@ -186,8 +188,10 @@ jobs:
       matrix:
         kernel: [prime-kernels, deep-ep, deep-gemm, torchao]
     runs-on: vm
+    # `env` isn't available in `container:` (GitHub Actions only exposes it inside
+    # steps), so this has to be a literal — keep it in sync with CONTAINER_IMAGE above.
     container:
-      image: ${{ env.CONTAINER_IMAGE }}
+      image: nvidia/cuda:13.0.3-devel-ubuntu24.04
       options: --gpus all
     steps:
       - name: Install runtime tooling

From e87f5f59e688ae84a5ef5ff515934fd4959f3003 Mon Sep 17 00:00:00 2001
From: Mario Sieg <mario.sieg.64@gmail.com>
Date: Thu, 27 Aug 2026 20:06:47 +0000
Subject: [PATCH 5/9] torchao aarch64

---
 .github/workflows/build_kernels.yaml | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/.github/workflows/build_kernels.yaml b/.github/workflows/build_kernels.yaml
index 620b0ad..c4d2b79 100644
--- a/.github/workflows/build_kernels.yaml
+++ b/.github/workflows/build_kernels.yaml
@@ -73,11 +73,6 @@ jobs:
             runner: image-builder
           - arch: aarch64
             runner: image-builder-arm-2204
-        exclude:
-          # torchao is consumed on x86_64 only (see its marker in prime-rl's
-          # pyproject.toml); aarch64 resolves the PyPI release there instead.
-          - kernel: torchao
-            arch: aarch64
     runs-on: ${{ matrix.runner }}
     # `env` isn't available in `container:` (GitHub Actions only exposes it inside
     # steps), so this has to be a literal — keep it in sync with CONTAINER_IMAGE above.

From 388e1bffe2aa431ae9329dfd4c304b1da470e471 Mon Sep 17 00:00:00 2001
From: Mario Sieg <mario.sieg.64@gmail.com>
Date: Thu, 27 Aug 2026 20:51:41 +0000
Subject: [PATCH 6/9] Skip smoke tests for now

---
 .github/workflows/build_kernels.yaml | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/build_kernels.yaml b/.github/workflows/build_kernels.yaml
index c4d2b79..d914bee 100644
--- a/.github/workflows/build_kernels.yaml
+++ b/.github/workflows/build_kernels.yaml
@@ -175,6 +175,11 @@ jobs:
   smoke-test:
     name: Smoke test ${{ matrix.kernel }}
     needs: build
+    # Flip to `true` once the `vm` GPU runner is registered for this repo (it's
+    # currently repo-scoped on prime-rl, not org-shared — see the runner-access
+    # discussion). `env`/`vars` aren't reliably available in a job-level `if`, so this
+    # is a literal, not a variable — that's deliberate, not an oversight.
+    if: false
     # x86_64 only: GPU runners registered for this org today are x86_64 (see the `vm`
     # label prime-rl's gpu_tests.yaml uses). If an aarch64 GPU runner (e.g. GB200)
     # becomes available, add an `arch` axis here the same way `build` has one.
@@ -224,7 +229,15 @@ jobs:
     # Only cut a real release on push-to-main or an explicit manual dispatch — a PR
     # should build and smoke-test, not publish. Without this, every push to a PR branch
     # would create and publish a real release the moment build+smoke-test go green.
-    if: github.event_name != 'pull_request'
+    #
+    # needs.smoke-test.result tolerates 'skipped' as well as 'success': smoke-test is
+    # currently disabled (see its `if: false`) until the GPU runner is registered, and
+    # a skipped dependency fails the default needs.*-implied success() check — without
+    # this, disabling smoke-test would also silently block every release.
+    if: |
+      github.event_name != 'pull_request' &&
+      needs.build.result == 'success' &&
+      (needs.smoke-test.result == 'success' || needs.smoke-test.result == 'skipped')
     runs-on: ubuntu-latest
     permissions:
       contents: write

From c293dc6cb697b4013db7b035a441fb83c9b6dbeb Mon Sep 17 00:00:00 2001
From: Mario Sieg <mario.sieg.64@gmail.com>
Date: Thu, 27 Aug 2026 21:09:29 +0000
Subject: [PATCH 7/9] glibc fix

---
 .github/workflows/build_kernels.yaml | 27 ++++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/build_kernels.yaml b/.github/workflows/build_kernels.yaml
index d914bee..84ac7d3 100644
--- a/.github/workflows/build_kernels.yaml
+++ b/.github/workflows/build_kernels.yaml
@@ -15,6 +15,14 @@ name: Build and Release Kernels
 # is unset, which silently scopes the wheel to whatever GPU happened to build it (that
 # bug produced a Blackwell-only deep-ep wheel once already).
 #
+# Building inside `pytorch/manylinux2_28-builder` (glibc 2.28, AlmaLinux 8), not a
+# plain Ubuntu 24.04 CUDA image: a wheel compiled against Ubuntu 24.04's glibc 2.39
+# fails to import (`GLIBC_2.38 not found`) on anything older — e.g. Ubuntu 22.04, still
+# a common dev-box OS here. glibc symbol versioning is backwards-compatible, so a
+# glibc-2.28 build imports fine everywhere newer, including Ubuntu 24.04 itself. This is
+# the same builder image PyTorch itself uses for its manylinux_2_28 wheels, matching the
+# platform tag vLLM's own prebuilt wheel already carries in `[tool.uv.sources]`.
+#
 # No GPU is needed to build: nvcc cross compiles for every architecture we ship. GPU
 # smoke tests run separately, against the built wheels, on an actual GPU runner.
 #
@@ -57,6 +65,10 @@ concurrency:
 env:
   # Keep in sync with the CUDA toolkit the container images below use.
   CUDA_TAG: cu130
+  # `build`'s container (glibc 2.28, for portability — see the comment above) and
+  # `smoke-test`'s (a stand-in for the real deployment target, which is Ubuntu 24.04 —
+  # see Dockerfile.cuda in prime-rl) are deliberately different images.
+  BUILD_CONTAINER_IMAGE: pytorch/manylinux2_28-builder:cuda13.0
   CONTAINER_IMAGE: nvidia/cuda:13.0.3-devel-ubuntu24.04
   TORCH_VERSION: 2.13.0+cu130
 
@@ -75,15 +87,16 @@ jobs:
             runner: image-builder-arm-2204
     runs-on: ${{ matrix.runner }}
     # `env` isn't available in `container:` (GitHub Actions only exposes it inside
-    # steps), so this has to be a literal — keep it in sync with CONTAINER_IMAGE above.
-    container: nvidia/cuda:13.0.3-devel-ubuntu24.04
+    # steps), so this has to be a literal — keep it in sync with BUILD_CONTAINER_IMAGE
+    # above. AlmaLinux 8 (glibc 2.28), not Ubuntu — see the portability comment up top.
+    container: pytorch/manylinux2_28-builder:cuda13.0
     steps:
-      # libibverbs-dev/librdmacm-dev: NVSHMEM's IBGDA transport headers (infiniband/mlx5dv.h),
-      # compile-time only — deep-ep dlopens the real thing at runtime.
+      # git, curl, gcc/g++/make and cmake/ninja already ship in this image (via
+      # /opt/conda and gcc-toolset-13). libibverbs-devel/librdmacm-devel: NVSHMEM's
+      # IBGDA transport headers (infiniband/mlx5dv.h), compile-time only — deep-ep
+      # dlopens the real thing at runtime. `xz` unpacks the NVSHMEM archive.
       - name: Install build tooling
-        run: |
-          apt-get update
-          apt-get install -y --no-install-recommends git curl ca-certificates build-essential gh xz-utils libibverbs-dev librdmacm-dev
+        run: dnf install -y --setopt=install_weak_deps=False ca-certificates xz libibverbs-devel librdmacm-devel
 
       - name: Checkout
         uses: actions/checkout@v5

From 14c699e27bd11e7cc35b13414700c69dc1053041 Mon Sep 17 00:00:00 2001
From: Mario Sieg <mario.sieg.64@gmail.com>
Date: Thu, 27 Aug 2026 21:23:51 +0000
Subject: [PATCH 8/9] glibc fix

---
 .github/workflows/build_kernels.yaml | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/build_kernels.yaml b/.github/workflows/build_kernels.yaml
index 84ac7d3..915cd9d 100644
--- a/.github/workflows/build_kernels.yaml
+++ b/.github/workflows/build_kernels.yaml
@@ -15,8 +15,10 @@ name: Build and Release Kernels
 # is unset, which silently scopes the wheel to whatever GPU happened to build it (that
 # bug produced a Blackwell-only deep-ep wheel once already).
 #
-# Building inside `pytorch/manylinux2_28-builder` (glibc 2.28, AlmaLinux 8), not a
-# plain Ubuntu 24.04 CUDA image: a wheel compiled against Ubuntu 24.04's glibc 2.39
+# Building inside PyTorch's manylinux builder images (glibc 2.28, AlmaLinux 8 — see
+# `build_image` per arch below; amd64 and arm64 are separate image repos, not one
+# multi-arch tag), not a plain Ubuntu 24.04 CUDA image: a wheel compiled against
+# Ubuntu 24.04's glibc 2.39
 # fails to import (`GLIBC_2.38 not found`) on anything older — e.g. Ubuntu 22.04, still
 # a common dev-box OS here. glibc symbol versioning is backwards-compatible, so a
 # glibc-2.28 build imports fine everywhere newer, including Ubuntu 24.04 itself. This is
@@ -65,10 +67,9 @@ concurrency:
 env:
   # Keep in sync with the CUDA toolkit the container images below use.
   CUDA_TAG: cu130
-  # `build`'s container (glibc 2.28, for portability — see the comment above) and
-  # `smoke-test`'s (a stand-in for the real deployment target, which is Ubuntu 24.04 —
-  # see Dockerfile.cuda in prime-rl) are deliberately different images.
-  BUILD_CONTAINER_IMAGE: pytorch/manylinux2_28-builder:cuda13.0
+  # `smoke-test`'s container is a stand-in for the real deployment target (Ubuntu 24.04
+  # — see Dockerfile.cuda in prime-rl), deliberately different from `build`'s (see the
+  # portability comment up top and the per-arch images in the `build` matrix below).
   CONTAINER_IMAGE: nvidia/cuda:13.0.3-devel-ubuntu24.04
   TORCH_VERSION: 2.13.0+cu130
 
@@ -81,15 +82,19 @@ jobs:
         kernel: [prime-kernels, deep-ep, deep-gemm, torchao]
         arch: [x86_64, aarch64]
         include:
+          # PyTorch publishes separate manylinux builder images per arch, not a
+          # multi-arch manifest under one tag — `manylinux2_28-builder` is amd64-only
+          # (confirmed: CI's aarch64 job failed to pull it with "no matching manifest
+          # for linux/arm64/v8"); `manylinuxaarch64-builder` is its aarch64 counterpart.
           - arch: x86_64
             runner: image-builder
+            build_image: pytorch/manylinux2_28-builder:cuda13.0
           - arch: aarch64
             runner: image-builder-arm-2204
+            build_image: pytorch/manylinuxaarch64-builder:cuda13.0
     runs-on: ${{ matrix.runner }}
-    # `env` isn't available in `container:` (GitHub Actions only exposes it inside
-    # steps), so this has to be a literal — keep it in sync with BUILD_CONTAINER_IMAGE
-    # above. AlmaLinux 8 (glibc 2.28), not Ubuntu — see the portability comment up top.
-    container: pytorch/manylinux2_28-builder:cuda13.0
+    # `matrix` (unlike `env`) is available in `container:`, so this can vary by arch.
+    container: ${{ matrix.build_image }}
     steps:
       # git, curl, gcc/g++/make and cmake/ninja already ship in this image (via
       # /opt/conda and gcc-toolset-13). libibverbs-devel/librdmacm-devel: NVSHMEM's

From dd539e18fffd63f58f36957ac815bfcbd41781a0 Mon Sep 17 00:00:00 2001
From: Mario Sieg <mario.sieg.64@gmail.com>
Date: Thu, 27 Aug 2026 21:34:15 +0000
Subject: [PATCH 9/9] Fix cuda driver runtime link

---
 scripts/install_ep_kernels.sh | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/scripts/install_ep_kernels.sh b/scripts/install_ep_kernels.sh
index fb21a50..749280b 100755
--- a/scripts/install_ep_kernels.sh
+++ b/scripts/install_ep_kernels.sh
@@ -53,6 +53,16 @@ if [ -d "$CUDA_HOME/include/cccl" ]; then
     export CPATH="$CUDA_HOME/include/cccl${CPATH:+:$CPATH}"
 fi
 
+# DeepEP links -lcuda (the driver stub, not the runtime). NVIDIA's own `devel` Ubuntu
+# images conveniently symlink it onto $CUDA_HOME/lib64/stubs, but PyTorch's manylinux
+# builder images don't — the stub only exists under
+# $CUDA_HOME/targets/<arch>-linux/lib/stubs (arch-linux name varies: x86_64-linux,
+# sbsa-linux, ...). Find it rather than hardcode the arch directory name.
+STUBS_DIR=$(find "$CUDA_HOME" -maxdepth 4 -type d -name stubs 2>/dev/null | head -1)
+if [ -n "$STUBS_DIR" ]; then
+    export LIBRARY_PATH="${STUBS_DIR}${LIBRARY_PATH:+:$LIBRARY_PATH}"
+fi
+
 NVCC_VER=$("$CUDA_HOME/bin/nvcc" --version | grep -oP 'release \K[\d.]+')
 echo "Torch CUDA: ${TORCH_CUDA_VER}, nvcc: ${NVCC_VER} (${CUDA_HOME})"