From 76f47d4a26b4ebed44d343e975a3f033883e6d07 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Mon, 3 Aug 2026 21:35:21 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .ci/scripts/wheel/cuda_arch_list.sh | 60 +++++++++++++++++++ .ci/scripts/wheel/envvar_linux.sh | 10 ++++ .../workflows/build-wheels-aarch64-linux.yml | 7 ++- .github/workflows/build-wheels-linux.yml | 7 ++- setup.py | 36 ++++++++++- 5 files changed, 117 insertions(+), 3 deletions(-) create mode 100644 .ci/scripts/wheel/cuda_arch_list.sh diff --git a/.ci/scripts/wheel/cuda_arch_list.sh b/.ci/scripts/wheel/cuda_arch_list.sh new file mode 100644 index 00000000000..d69f2535d7c --- /dev/null +++ b/.ci/scripts/wheel/cuda_arch_list.sh @@ -0,0 +1,60 @@ +#!/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, CMake 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 GPU +# generation. Detection is the right default for a local build and the wrong one for a +# published artifact. +# +# CUDA_ARCH_LIST_ names the architectures for that CUDA train. DESIRED_CUDA is +# supplied by the wheel build as cu126, cu130, and so on. + +# Data center and desktop parts on the CUDA 13 trains: Ampere, Hopper, Blackwell data center, +# and Blackwell desktop. +_cuda_arch_x86_64_cu130="80-real;90-real;100-real;120-real" +_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="90-real;100-real;110-real" +_cuda_arch_aarch64_cu132="${_cuda_arch_aarch64_cu130}" + +# The older CUDA train, where Orin is the target. +_cuda_arch_aarch64_cu126="87-real" +_cuda_arch_x86_64_cu126="80-real;90-real" + +executorch_cuda_arch_list() { + local machine + machine="$(uname -m)" + local train="${DESIRED_CUDA:-}" + if [ -z "${train}" ]; then + return 0 + fi + # DESIRED_CUDA 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}" ;; + 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}" ;; + esac + ;; + esac +} diff --git a/.ci/scripts/wheel/envvar_linux.sh b/.ci/scripts/wheel/envvar_linux.sh index 3b24b3f7188..1d4c972f4d6 100755 --- a/.ci/scripts/wheel/envvar_linux.sh +++ b/.ci/scripts/wheel/envvar_linux.sh @@ -9,3 +9,13 @@ # 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" +_executorch_cuda_arch="$(executorch_cuda_arch_list)" +if [ -n "${_executorch_cuda_arch}" ]; then + export CMAKE_ARGS="${CMAKE_ARGS} -DCMAKE_CUDA_ARCHITECTURES=${_executorch_cuda_arch}" + echo "CUDA architectures for this row: ${_executorch_cuda_arch}" +fi diff --git a/.github/workflows/build-wheels-aarch64-linux.yml b/.github/workflows/build-wheels-aarch64-linux.yml index 8adf4268228..fc01fa0294b 100644 --- a/.github/workflows/build-wheels-aarch64-linux.yml +++ b/.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"]' diff --git a/.github/workflows/build-wheels-linux.yml b/.github/workflows/build-wheels-linux.yml index 7428b68a773..49a80148ec6 100644 --- a/.github/workflows/build-wheels-linux.yml +++ b/.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"]' diff --git a/setup.py b/setup.py index 58258c1491d..ea978af1592 100644 --- a/setup.py +++ b/setup.py @@ -108,6 +108,40 @@ 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. + """ + train = os.environ.get("DESIRED_CUDA", "").strip().lower() + if not train: + return "" + train = train.removeprefix("cu").replace(".", "") + return train if train.isdigit() else "" + + +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, so the CPU rows + are unaffected. The major decides the package suffix, matching how these are published. + """ + train = _cuda_train() + if not train or not _is_env_flag_enabled("EXECUTORCH_BUILD_CUDA"): + return [] + major = train[:2] if train.startswith("13") else train[:2] + suffix = f"cu{major}" + # 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. + return [ + f"nvidia-cuda-runtime-{suffix}; platform_system == 'Linux'", + f"nvidia-curand-{suffix}; platform_system == 'Linux'", + f"nvidia-cublas-{suffix}; platform_system == 'Linux'", + ] + + def _minimal_cmake_flags() -> List[str]: return [ "-DEXECUTORCH_BUILD_COREML=OFF", @@ -1205,7 +1239,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(