Ship the CUDA runtime in the Windows CUDA wheel

.ci/scripts/install_nvidia_driver_windows.ps1

@@ -0,0 +1,28 @@
+# 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.
+
+# Installs the NVIDIA display driver on a Windows GPU runner. The runner images carry the
+# GPU but no driver, so without this nvcuda.dll is missing and every CUDA call fails with
+# cudaErrorInsufficientDriver. The same driver, from the same place, that PyTorch's Windows
+# CUDA smoke tests install first (.ci/pytorch/windows/internal/driver_update.bat in
+# pytorch/pytorch), so both projects test against one driver.
+$ErrorActionPreference = "Stop"
+$version = "580.88"
+$installer = Join-Path $env:RUNNER_TEMP "$version-data-center-tesla-desktop-win10-win11-64bit-dch-international.exe"
+$url = "https://ossci-windows.s3.amazonaws.com/$(Split-Path $installer -Leaf)"
+
+$nvcuda = Join-Path $env:SystemRoot "System32\nvcuda.dll"
+if (Test-Path $nvcuda) {
+    Write-Host "NVIDIA driver already present: $((Get-Item $nvcuda).VersionInfo.FileVersion)"
+    exit 0
+}
+curl.exe --retry 3 -fsSL $url --output $installer
+if ($LASTEXITCODE -ne 0) { throw "downloading the NVIDIA driver from $url failed" }
+$process = Start-Process -FilePath $installer -ArgumentList "-s", "-noreboot" -Wait -PassThru
+Remove-Item $installer -ErrorAction SilentlyContinue
+if ($process.ExitCode -ne 0) { throw "the NVIDIA driver installer exited with $($process.ExitCode)" }
+if (-not (Test-Path $nvcuda)) { throw "the NVIDIA driver installed but $nvcuda is missing" }
+Write-Host "NVIDIA driver $version installed: $((Get-Item $nvcuda).VersionInfo.FileVersion)"

.ci/scripts/test_cuda_runtime_windows.ps1

@@ -0,0 +1,48 @@
+# 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.
+
+# Builds and runs the CUDA backend's runtime C++ tests on a Windows GPU runner, the
+# counterpart of the unittest-cuda-runtime job in cuda.yml. The Windows runners carry
+# data-center GPUs in TCC mode, which have no memory pools, so this is also what covers the
+# allocator's synchronous fallback; the Linux job covers the pool path.
+$ErrorActionPreference = "Stop"
+
+# Windows PowerShell 5.1 does not stop on a failing native command.
+function Invoke-Native {
+    param([Parameter(Mandatory = $true)][scriptblock]$Command)
+    & $Command
+    if ($LASTEXITCODE -ne 0) {
+        throw "exit code ${LASTEXITCODE}: $Command"
+    }
+}
+
+& "$PSScriptRoot\install_nvidia_driver_windows.ps1"
+
+$cudaNvcc = Join-Path $env:CUDA_HOME "bin\nvcc.exe"
+if (-not (Test-Path $cudaNvcc)) { throw "CUDA compiler not found at '$cudaNvcc'" }
+$env:CUDACXX = $cudaNvcc
+$env:PATH = "$env:CUDA_HOME\bin\x64;$env:CUDA_HOME\bin;$env:PATH"
+# The runner's GPU has no memory pools, which is what makes this job cover the
+# allocator's fallback; fail rather than skip those tests if it ever has them.
+$env:EXECUTORCH_CUDA_TEST_REQUIRE_NO_MEMORY_POOLS = "1"
+
+$tests = @(
+    "test_cuda_allocator",
+    "test_cuda_mutable_state",
+    "test_cuda_weight_cache",
+    "test_cuda_sort_rand",
+    "test_cuda_guard",
+    "test_cuda_stream_guard"
+)
+$numCores = [Math]::Max([Environment]::ProcessorCount - 1, 1)
+Invoke-Native {
+    cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON -DCMAKE_CXX_STANDARD=20 `
+        -T "cuda=$env:CUDA_HOME" "-DCMAKE_CUDA_COMPILER=$cudaNvcc" "-DCUDAToolkit_ROOT=$env:CUDA_HOME"
+}
+Invoke-Native { cmake --build cmake-out --config Release -j $numCores --target $tests }
+foreach ($test in $tests) {
+    Invoke-Native { ctest --test-dir cmake-out -C Release -R "^$test`$" --output-on-failure -V }
+}

.ci/scripts/wheel/install_cuda_redist.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python
+# 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.
+
+"""Assemble part of a CUDA toolkit from NVIDIA's redistributable component archives.
+
+The Windows CUDA wheel end-to-end jobs cover several CUDA trains while the CI images carry
+one toolkit, and NVIDIA publishes no full Windows installer for every release. The redist
+archives exist for every release on both platforms, need no installer or admin rights, and
+are checksummed in a JSON index. Merging the components a job needs into one directory
+gives the layout nvcc, CMake's FindCUDAToolkit and torch's cross-compile step expect.
+
+    install_cuda_redist.py --train 13.4 --platform linux-x86_64 --dest /opt/cuda-13.4
+    install_cuda_redist.py --train 13.4 --platform windows-x86_64 --dest DIR --components cuda_cudart
+"""
+
+import argparse
+import hashlib
+import json
+import shutil
+import sys
+import tarfile
+import tempfile
+import urllib.request
+import zipfile
+from pathlib import Path
+
+_REDIST = "https://developer.download.nvidia.com/compute/cuda/redist"
+
+# The patch release each train resolves to: one entry per train in filter_cuda_matrix.py's
+# SUPPORTED_CUDA_VERSIONS, which the end-to-end jobs take their matrix from, and 13.0, which
+# is no longer published but still builds from source and is what local tests of it use. A
+# train missing here fails the job at argument parsing rather than testing a different
+# release.
+_RELEASES = {"13.0": "13.0.2", "13.2": "13.2.2", "13.4": "13.4.1"}
+
+# What a build of the CUDA backend needs. "a|b" names one component a release publishes
+# under either name (CCCL was renamed between releases).
+_BUILD_COMPONENTS = [
+    "cuda_nvcc",
+    "cuda_cudart",
+    "cuda_crt",
+    "cccl|cuda_cccl",
+    "libnvvm",
+    "cuda_cuobjdump",
+    "cuda_nvrtc",
+    "cuda_profiler_api",
+    "cuda_nvtx",
+    # rand.cu includes curand_kernel.h.
+    "libcurand",
+]
+
+
+def _fetch(url: str, destination: Path) -> None:
+    with urllib.request.urlopen(url, timeout=300) as response, open(
+        destination, "wb"
+    ) as handle:
+        shutil.copyfileobj(response, handle)
+
+
+def _extract(archive: Path, into: Path) -> Path:
+    if archive.suffix == ".zip":
+        with zipfile.ZipFile(archive) as bundle:
+            bundle.extractall(into)
+    else:
+        with tarfile.open(archive) as bundle:
+            if hasattr(tarfile, "fully_trusted_filter"):
+                bundle.extractall(into, filter="fully_trusted")
+            else:
+                bundle.extractall(into)
+    # Each archive holds one top-level directory named after itself.
+    (top,) = [entry for entry in into.iterdir() if entry.is_dir()]
+    return top
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    parser.add_argument("--train", required=True, choices=sorted(_RELEASES))
+    parser.add_argument(
+        "--platform", required=True, choices=["linux-x86_64", "windows-x86_64"]
+    )
+    parser.add_argument("--dest", required=True, type=Path)
+    parser.add_argument("--components", nargs="+", default=_BUILD_COMPONENTS)
+    arguments = parser.parse_args()
+
+    release = _RELEASES[arguments.train]
+    with urllib.request.urlopen(
+        f"{_REDIST}/redistrib_{release}.json", timeout=120
+    ) as response:
+        index = json.load(response)
+    arguments.dest.mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory() as work:
+        work = Path(work)
+        for choice in arguments.components:
+            names = [
+                name
+                for name in choice.split("|")
+                if arguments.platform in index.get(name, {})
+            ]
+            if not names:
+                sys.exit(
+                    f"CUDA {release} publishes none of {choice} for {arguments.platform}"
+                )
+            entry = index[names[0]][arguments.platform]
+            archive = work / Path(entry["relative_path"]).name
+            _fetch(f"{_REDIST}/{entry['relative_path']}", archive)
+            digest = hashlib.sha256(archive.read_bytes()).hexdigest()
+            if digest != entry["sha256"]:
+                sys.exit(f"checksum mismatch for {archive.name}")
+            staging = work / "staging"
+            top = _extract(archive, staging)
+            shutil.copytree(top, arguments.dest, symlinks=True, dirs_exist_ok=True)
+            shutil.rmtree(staging)
+            archive.unlink()
+            print(
+                f"merged {names[0]} {index[names[0]]['version']} into {arguments.dest}"
+            )
+    # The Linux archives keep libraries in lib/, while nvcc's link step and CMake's
+    # FindCUDAToolkit look in lib64/, which is where an installed toolkit has them.
+    lib, lib64 = arguments.dest / "lib", arguments.dest / "lib64"
+    if arguments.platform == "linux-x86_64" and lib.is_dir() and not lib64.exists():
+        lib64.symlink_to("lib", target_is_directory=True)
+
+
+if __name__ == "__main__":
+    main()

.ci/scripts/wheel/pre_build_script.sh

@@ -116,6 +116,23 @@ else
     export CMAKE_ARGS="${CMAKE_ARGS:-} -DEXECUTORCH_BUILD_CUDA=ON"
     echo "CMAKE_ARGS=${CMAKE_ARGS}" >> "${GITHUB_ENV}"
     echo "row '${CU_VERSION:-${DESIRED_CUDA:-}}' is a CUDA row, requiring the CUDA build"
+    # The Linux rows take the row's GPU list from envvar_cuda_linux.sh. Windows has no env-var
+    # script slot, so the same list is resolved here; without it the build compiles device
+    # code only for whatever GPU the builder has. The reusable workflow writes its own list
+    # into BUILD_ENV_FILE, which the build step sources after this hook and which would
+    # otherwise win, so the row's list is appended there as well as to GITHUB_ENV.
+    # Also needed for the ninja the Windows CUDA build generates with.
+    if [[ $UNAME_S == *"MINGW"* || $UNAME_S == *"MSYS"* ]]; then
+        source "${GITHUB_WORKSPACE}/${REPOSITORY}/.ci/scripts/wheel/cuda_arch_list.sh"
+        TORCH_CUDA_ARCH_LIST="$(EXECUTORCH_BUILD_CUDA=1 executorch_cuda_arch_list)"
+        export TORCH_CUDA_ARCH_LIST
+        echo "TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}" >> "${GITHUB_ENV}"
+        if [[ -n "${BUILD_ENV_FILE:-}" ]]; then
+            echo "export TORCH_CUDA_ARCH_LIST='${TORCH_CUDA_ARCH_LIST}'" >> "${BUILD_ENV_FILE}"
+        fi
+        echo "building device code for: ${TORCH_CUDA_ARCH_LIST}"
+        pip install --quiet ninja
+    fi
 fi
 
 # On Windows, enable symlinks and re-checkout the current revision to create
@@ -125,14 +142,6 @@ if [[ $UNAME_S == *"MINGW"* || $UNAME_S == *"MSYS"* ]]; then
     git config core.symlinks true
     git checkout -f HEAD
 
-    # Windows wheels are CPU-only (build-wheels-windows.yml sets
-    # with-cuda: disabled), but the Windows CI image ships a CUDA toolkit on
-    # PATH, which makes setup.py auto-enable EXECUTORCH_BUILD_CUDA. That bakes a
-    # CUDA _C into the CPU wheel, which then fails its DLL load in the
-    # smoke test ("DLL load failed while importing _C"). Force a
-    # CPU-only build.
-    export CMAKE_ARGS="${CMAKE_ARGS:-} -DEXECUTORCH_BUILD_CUDA=OFF"
-    echo "CMAKE_ARGS=${CMAKE_ARGS}" >> "${GITHUB_ENV}"
 fi
 
 # Manually install build requirements because `python setup.py bdist_wheel` does
@@ -202,7 +211,8 @@ if [[ "${EXECUTORCH_BUILD_VULKAN:-0}" != "0" \
     echo "glslc installed: $(command -v glslc)"
   fi
 else
-  # This is the Vulkan equivalent of the Windows CUDA force-off above (#20527).
+  # Off unless asked for: a builder with the Vulkan SDK would otherwise turn it on
+  # for a wheel that does not ship it, as a CUDA toolkit once did for Windows (#20527).
   export CMAKE_ARGS="${CMAKE_ARGS:-} -DEXECUTORCH_BUILD_VULKAN=OFF"
   echo "CMAKE_ARGS=${CMAKE_ARGS}" >> "${GITHUB_ENV}"
 fi

.ci/scripts/wheel/test_cpp_sdk.py

@@ -1741,11 +1741,13 @@ def test_pre_3_28_route_builds_a_consumer_through_variables(work_dir: Path) -> N
         "target_link_libraries(consumer PRIVATE ${EXECUTORCH_LIBRARIES})\n"
         "set_target_properties(consumer PROPERTIES CXX_STANDARD ${EXECUTORCH_CXX_STANDARD})\n"
         # No imported targets on this route, so no TARGET_RUNTIME_DLLS either. A Windows
-        # consumer copies the DLLs from the directory the package reports.
+        # consumer copies the DLLs from the directory the package reports, plus any the
+        # package lists separately because they ship elsewhere.
         "if(WIN32)\n"
         '  file(GLOB _dlls "${EXECUTORCH_RUNTIME_LIBRARY_DIR}/*.dll")\n'
         "  add_custom_command(TARGET consumer POST_BUILD COMMAND ${CMAKE_COMMAND} -E "
-        "copy_if_different ${_dlls} $<TARGET_FILE_DIR:consumer>)\n"
+        "copy_if_different ${_dlls} ${EXECUTORCH_RUNTIME_DLLS_EXTRA} "
+        "$<TARGET_FILE_DIR:consumer>)\n"
         "endif()\n"
     )
     build_dir = work_dir / "pre-328-build"