From a50fdb0528768ad138c760fac575240b494ae39f Mon Sep 17 00:00:00 2001
From: nimlgen <138685161+nimlgen@users.noreply.github.com>
Date: Tue, 17 Mar 2026 17:25:33 +0800
Subject: [PATCH] nvcc macos (#15308)

* fix nvcc install macos

* um

* arm

* per

* tm
---
 extra/setup_nvcc_osx.sh | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/extra/setup_nvcc_osx.sh b/extra/setup_nvcc_osx.sh
index d95348a342..1244c2ea87 100755
--- a/extra/setup_nvcc_osx.sh
+++ b/extra/setup_nvcc_osx.sh
@@ -1,9 +1,9 @@
 #!/bin/sh
 install_loc="$HOME/.local/bin"
-docker build --platform=linux/amd64 -t cuda-nvcc:12.8 - <<'EOF'
+docker build --platform=linux/arm64 -t cuda-nvcc:12.8 - <<'EOF'
 FROM ubuntu:22.04
 RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates && \
-  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
+  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb && \
   dpkg -i cuda-keyring_1.1-1_all.deb && \
   apt-get update && apt-get install -y --no-install-recommends cuda-nvcc-12-8 cuda-nvdisasm-12-8 cuda-cuobjdump-12-8 && rm -rf /var/lib/apt/lists/*
 ENV PATH=/usr/local/cuda/bin:$PATH
@@ -13,10 +13,14 @@ mkdir -p "$install_loc"
 tee "$install_loc/nvccshim" >/dev/null <<'EOF'
 #!/bin/sh
 set -eu
-# assume the final arg is the input path
-# mount it so that container can read it
-dir=$(dirname "${@: -1}")
-exec docker run --rm --platform=linux/amd64 -v "$dir":"$dir" cuda-nvcc:12.8 "$(basename "$0")" "$@"
+cname="cuda-nvcc-persistent"
+if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q true; then
+  docker rm -f "$cname" 2>/dev/null || true
+  docker run -d --platform=linux/arm64 --name "$cname" \
+    -v /var/folders:/var/folders -v "$HOME":"$HOME" \
+    cuda-nvcc:12.8 sleep 300 >/dev/null
+fi
+exec docker exec "$cname" "$(basename "$0")" "$@"
 EOF
 chmod +x "$install_loc/nvccshim"
 for t in nvcc nvdisasm; do