Publish CUDA wheels so a GPU user does not build from source

.ci/scripts/wheel/cuda_arch_list.sh

@@ -0,0 +1,132 @@
+#!/usr/bin/env bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# GPU architectures to compile device code for, chosen per release row rather than detected
+# from the build machine.
+#
+# Without this, the build compiles for whichever GPU the builder happens to have. The wheel
+# then installs on every machine the row claims and fails when a model runs on a different
+# generation. Detection is the right default for a local build and the wrong one for a
+# published artifact.
+#
+# The value is published as TORCH_CUDA_ARCH_LIST rather than CMAKE_CUDA_ARCHITECTURES.
+# PyTorch's own CMake explicitly rejects the latter and overrides it with OFF, so setting it
+# alone silently reduces the build to a single detected architecture.
+
+# Data center and desktop parts on the CUDA 13 trains: Ampere, Hopper, Blackwell data center,
+# and Blackwell desktop.
+_cuda_arch_x86_64_cu130="8.0 9.0 10.0 12.0"
+_cuda_arch_x86_64_cu132="${_cuda_arch_x86_64_cu130}"
+
+# Server-class ARM plus the Jetson modules whose CUDA train matches: Hopper for Grace-Hopper,
+# Blackwell data center for GB200, and Thor.
+_cuda_arch_aarch64_cu130="9.0 10.0 11.0"
+_cuda_arch_aarch64_cu132="${_cuda_arch_aarch64_cu130}"
+
+# The older CUDA train on generic ARM. These are the server parts that train supports, matching
+# what the x86_64 row of the same train claims. Jetson is deliberately not here: a Jetson-only
+# architecture such as Orin's 8.7 in a generic manylinux wheel would advertise a device the row
+# cannot otherwise serve, since a Jetson also needs the pinned CUDA, TensorRT and PyTorch from its
+# own software release rather than the ones a generic wheel resolves.
+_cuda_arch_aarch64_cu126="8.0 9.0"
+_cuda_arch_x86_64_cu126="8.0 9.0"
+
+# A CUDA train with no architecture list would otherwise leave the build detecting the
+# builder's GPU, which is the failure this file exists to prevent. Adding a train to the release
+# matrix without adding its architectures should fail loudly.
+_executorch_unknown_train() {
+  echo "cuda_arch_list.sh: no GPU architecture list for CUDA train '$1' on $(uname -m)." >&2
+  echo "Add one before building this row, or the wheel ships device code for one GPU only." >&2
+  return 64
+}
+
+# The architectures for the current row, space separated in the dotted form PyTorch expects.
+# Empty when the row is unknown, which leaves the build detecting as before.
+executorch_cuda_arch_list() {
+  local machine
+  machine="$(uname -m)"
+  # The wheel build exports the row's CUDA train as CU_VERSION. DESIRED_CUDA is the name of
+  # the matrix field, not of the variable, so reading only that one leaves every row falling
+  # back to detecting the builder's GPU.
+  local train="${CU_VERSION:-${DESIRED_CUDA:-}}"
+  # A CPU row names no CUDA train and needs no architectures, so it is not an error.
+  case "${train}" in
+    "" | cpu | CPU | none | NONE) return 0 ;;
+  esac
+  # The value arrives as cu130; some callers pass 13.0 instead.
+  train="${train#cu}"
+  train="${train//./}"
+
+  case "${machine}" in
+    aarch64 | arm64)
+      case "${train}" in
+        126) printf '%s' "${_cuda_arch_aarch64_cu126}" ;;
+        130) printf '%s' "${_cuda_arch_aarch64_cu130}" ;;
+        132) printf '%s' "${_cuda_arch_aarch64_cu132}" ;;
+        *) _executorch_unknown_train "${train}" ;;
+      esac
+      ;;
+    x86_64)
+      case "${train}" in
+        126) printf '%s' "${_cuda_arch_x86_64_cu126}" ;;
+        130) printf '%s' "${_cuda_arch_x86_64_cu130}" ;;
+        132) printf '%s' "${_cuda_arch_x86_64_cu132}" ;;
+        *) _executorch_unknown_train "${train}" ;;
+      esac
+      ;;
+    *) _executorch_unknown_train "${train}" ;;
+  esac
+}
+
+# The highest architecture in a row, which is the one that gets a portable form as well.
+executorch_cuda_top_arch() {
+  local dotted
+  # Propagate a failed lookup rather than reporting an empty top architecture, since the caller
+  # cannot tell those apart and an unknown row would otherwise pass silently.
+  dotted="$(executorch_cuda_arch_list)" || return $?
+  [ -n "${dotted}" ] || return 0
+  printf '%s\n' ${dotted} | sort -g | tail -1
+}
+
+# The row's architectures plus a portable form of the highest one, which is what PyTorch's own
+# variable wants. Without it a GPU newer than every listed architecture has no kernel image at all
+# and the model cannot run. Measured to add no size, since the embedded text fits in space the
+# object already reserves.
+executorch_cuda_arch_list_with_ptx() {
+  local dotted top
+  # Capture the status explicitly. An unknown row makes the inner helper fail, and testing the
+  # captured string alone would discard that status and let the build ship a wheel with no device
+  # code at all.
+  dotted="$(executorch_cuda_arch_list)" || return $?
+  [ -n "${dotted}" ] || return 0
+  top="$(executorch_cuda_top_arch)" || return $?
+  printf '%s %s+PTX' "${dotted}" "${top}"
+}
+
+# The same architectures in CMake's own form, for targets outside PyTorch's CMake.
+executorch_cuda_cmake_arch_list() {
+  local dotted
+  # Propagate a failed lookup. Returning an empty list on an unknown row would let the build
+  # continue and produce an accelerator wheel carrying no device code.
+  dotted="$(executorch_cuda_arch_list)" || return $?
+  if [ -z "${dotted}" ]; then
+    return 0
+  fi
+  local out="" entry
+  for entry in ${dotted}; do
+    entry="${entry//./}"
+    out="${out:+${out};}${entry}-real"
+  done
+  # Also emit the highest architecture in its portable form, so a device newer than every listed
+  # one can still compile at load time rather than finding no kernel image.
+  local top
+  top="$(executorch_cuda_top_arch)"
+  if [ -n "${top}" ]; then
+    out="${out:+${out};}${top//./}-virtual"
+  fi
+  printf '%s' "${out}"
+}

.ci/scripts/wheel/envvar_linux.sh

@@ -9,3 +9,29 @@
 # any variables so that subprocesses will see them.
 
 source "${GITHUB_WORKSPACE}/${REPOSITORY}/.ci/scripts/wheel/envvar_base.sh"
+
+# Compile device code for the GPU architectures this release row claims, rather than for
+# whichever GPU the builder happens to have. A wheel built with detection alone installs on
+# every machine the row covers and then fails when a model runs on a different generation.
+source "${GITHUB_WORKSPACE}/${REPOSITORY}/.ci/scripts/wheel/cuda_arch_list.sh"
+# Check the status rather than only the output. An unrecognised row makes the resolver fail, but
+# this script is sourced and does not run under a failing-command shell, so ignoring the status
+# would leave both variables unset and let the build fall back to detecting the builder's own GPU.
+# That is exactly the outcome this file exists to prevent, and it would ship quietly.
+if ! _executorch_cuda_arch="$(executorch_cuda_arch_list_with_ptx)"; then
+  echo "could not resolve GPU architectures for CU_VERSION=${CU_VERSION:-unset}" >&2
+  exit 1
+fi
+if [ -n "${_executorch_cuda_arch}" ]; then
+  if ! _executorch_cuda_cmake_arch="$(executorch_cuda_cmake_arch_list)"; then
+    echo "could not resolve CMake GPU architectures for CU_VERSION=${CU_VERSION:-unset}" >&2
+    exit 1
+  fi
+  # PyTorch's CMake rejects CMAKE_CUDA_ARCHITECTURES and overrides it with OFF, which leaves
+  # the build compiling for one detected architecture, so the list has to go through the
+  # variable PyTorch reads. Both are set: targets that go through PyTorch's CMake honour the
+  # first, and any that do not honour the second.
+  export TORCH_CUDA_ARCH_LIST="${_executorch_cuda_arch}"
+  export CMAKE_ARGS="${CMAKE_ARGS} -DCMAKE_CUDA_ARCHITECTURES=${_executorch_cuda_cmake_arch}"
+  echo "CUDA architectures for this row: ${_executorch_cuda_arch}"
+fi

.github/workflows/build-wheels-aarch64-linux.yml

@@ -30,7 +30,12 @@ jobs:
       os: linux-aarch64
       test-infra-repository: pytorch/test-infra
       test-infra-ref: main
-      with-cuda: disabled
+      # CUDA enabled so the accelerator rows are built and published. The generator
+      # emits one cell per supported CUDA train, and each cell carries its own local
+      # version label so the CPU and accelerator artifacts stay distinguishable. A plain
+      # `pip install executorch` continues to resolve the CPU wheel from the default
+      # index; an accelerator wheel requires pointing at the matching index.
+      with-cuda: enable
       with-rocm: disabled
       python-versions: '["3.10", "3.11", "3.12", "3.13", "3.14"]'
 

.github/workflows/build-wheels-linux.yml

@@ -30,7 +30,12 @@ jobs:
       os: linux
       test-infra-repository: pytorch/test-infra
       test-infra-ref: main
-      with-cuda: disabled
+      # CUDA enabled so the accelerator rows are built and published. The generator
+      # emits one cell per supported CUDA train, and each cell carries its own local
+      # version label so the CPU and accelerator artifacts stay distinguishable. A plain
+      # `pip install executorch` continues to resolve the CPU wheel from the default
+      # index; an accelerator wheel requires pointing at the matching index.
+      with-cuda: enable
       with-rocm: disabled
       python-versions: '["3.10", "3.11", "3.12", "3.13", "3.14"]'
 

install_requirements.py

@@ -45,7 +45,11 @@ def install_requirements(use_pytorch_nightly):
 
     # Determine the appropriate PyTorch URL based on CUDA delegate status
     torch_url = determine_torch_url(TORCH_URL_BASE)
-    torchao_url = determine_torch_url(TORCHAO_URL_BASE)
+    # torchao only publishes its pinned nightly on the plain index. The CUDA-suffixed indexes
+    # carry it for x86_64 only, so suffixing this URL breaks every aarch64 accelerator row the
+    # moment CUDA is detected. The build needs torchao's Python API rather than its GPU kernels,
+    # so the plain index is the right source on every platform.
+    torchao_url = f"{TORCHAO_URL_BASE}/cpu"
 
     # pip packages needed by exir.
     TORCH_PACKAGE = [

setup.py

@@ -108,6 +108,67 @@ def _is_minimal_build() -> bool:
     return _is_env_flag_enabled("EXECUTORCH_BUILD_MINIMAL")
 
 
+def _cuda_train() -> str:
+    """The CUDA train this wheel is being built for, as a bare number like "130".
+
+    Read from the wheel build environment rather than detected from an installed compiler.
+    A CPU wheel built on a machine that happens to have a CUDA toolkit must not declare CUDA
+    dependencies, and detection cannot tell the two cases apart.
+
+    The wheel build exports this as CU_VERSION; DESIRED_CUDA is the matrix field name and is
+    accepted only so local invocations keep working.
+    """
+    train = (
+        (os.environ.get("CU_VERSION") or os.environ.get("DESIRED_CUDA") or "")
+        .strip()
+        .lower()
+    )
+    if not train:
+        return ""
+    train = train.removeprefix("cu").replace(".", "")
+    return train if train.isdigit() else ""
+
+
+# The published project names for the CUDA runtime components this wheel links but does not
+# bundle, keyed by CUDA major. These are not derivable from a suffix rule: the CUDA 12 wheels
+# carry a "-cu12" suffix while the CUDA 13 wheels are published under unsuffixed names, and the
+# suffixed CUDA 13 projects are placeholders that ship no binaries. A train with no entry gets
+# no declared dependencies, which is safer than requesting a name that may not exist.
+_CUDA_RUNTIME_PACKAGES = {
+    "12": ("nvidia-cuda-runtime-cu12", "nvidia-curand-cu12", "nvidia-cublas-cu12"),
+    "13": ("nvidia-cuda-runtime", "nvidia-curand", "nvidia-cublas"),
+}
+
+
+def _cuda_dependencies() -> List[str]:
+    """Runtime libraries a CUDA wheel needs but does not bundle.
+
+    Empty for a CPU wheel, and empty for a build whose CUDA train is unknown or unmapped, so
+    the CPU rows are unaffected.
+    """
+    train = _cuda_train()
+    if not train:
+        return []
+    # The CUDA train alone decides this. An earlier version also required
+    # EXECUTORCH_BUILD_CUDA, but that reaches the build as a CMake argument rather than an
+    # environment variable, so the condition was never true and a CUDA wheel shipped with no
+    # declared CUDA dependencies at all.
+    packages = _CUDA_RUNTIME_PACKAGES.get(train[:2])
+    if not packages:
+        return []
+    # Only what the delegate and its shim actually link. A shorter list keeps a CUDA install
+    # from pulling in libraries nothing in this wheel references.
+    #
+    # Bounded to the major this wheel was built against, which is the same key that selected the
+    # names. Without an upper bound, the day a newer CUDA major publishes, a fresh install of an
+    # unchanged wheel resolves to it, and a different major carries a different library version
+    # than the shipped code links, so the wheel installs and then fails to load.
+    major = int(train[:2])
+    return [
+        f"{name}>={major},<{major + 1}; platform_system == 'Linux'" for name in packages
+    ]
+
+
 def _minimal_cmake_flags() -> List[str]:
     return [
         "-DEXECUTORCH_BUILD_COREML=OFF",
@@ -1247,7 +1308,7 @@ def run(self):  # noqa C901
     setup_kwargs["packages"] = _minimal_packages()
     setup_kwargs["install_requires"] = _minimal_dependencies()
 else:
-    setup_kwargs["install_requires"] = _base_dependencies()
+    setup_kwargs["install_requires"] = _base_dependencies() + _cuda_dependencies()
 
 
 setup(