[Build] Download CUDA 12.9 nvcc/ptxas/cicc for better codegen in Docker build

Use CUDA 12.9.86 compiler toolchain (nvcc, ptxas, cicc) for better GPU
code generation while keeping CUDA 12.8 headers/libs for runtime
compatibility with PyTorch cu128 wheels.

nvcc is replaced (not just ptxas/cicc) so that setup.py's
get_nvcc_cuda_version() reports 12.9, which is required for FlashMLA
registration. PyTorch's CMake nvcc-vs-header version check is patched
to a warning in base (for extensions-build), csrc-build, and build
stages (re-applied after torch pip reinstall overwrites cmake files).

Controlled by ARG MIN_NVCC_TOOLCHAIN_VERSION (default 12.9.86):
- Skipped if system nvcc is already >= the specified version
- Set to empty string to disable entirely
- Archive integrity verified via SHA256 checksums from NVIDIA's
  redistrib manifest

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

.pre-commit-config.yaml

@@ -38,7 +38,7 @@ repos:
   rev: 0.9.1
   hooks:
     - id: pip-compile
-      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu129, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
+      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu128, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
       files: ^requirements/test\.(in|txt)$
 - repo: local
   hooks:

docker/Dockerfile

@@ -22,7 +22,7 @@
 #   docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
 # =============================================================================
 
-ARG CUDA_VERSION=12.9.1
+ARG CUDA_VERSION=12.8.1
 ARG PYTHON_VERSION=3.12
 
 # By parameterizing the base images, we allow third-party to use their own
@@ -183,6 +183,50 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # From versions.json: .torch.cuda_arch_list
 ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
+
+# Download newer CUDA compiler toolchain for better GPU code generation
+# while keeping CUDA 12.8 headers/libs for runtime compatibility with PyTorch cu128.
+# We replace nvcc, ptxas, and cicc so that setup.py's get_nvcc_cuda_version() reports
+# the newer version (required for FlashMLA registration which gates on >= 12.9).
+# PyTorch's CMake checks nvcc version vs CUDA header version; that check is patched
+# to a warning in the csrc-build and build stages (after torch is pip-installed).
+# Skipped if system nvcc is already >= this version (e.g. when building with CUDA 13+).
+# Set MIN_NVCC_TOOLCHAIN_VERSION="" to disable.
+ARG MIN_NVCC_TOOLCHAIN_VERSION=12.9.86
+# SHA256 checksums from https://developer.download.nvidia.com/compute/cuda/redist/redistrib_<cuda_version>.json
+ARG NVCC_TOOLCHAIN_SHA256_X86_64=7a1a5b652e5ef85c82b721d10672fc9a2dbaab44e9bd3c65a69517bf53998c35
+ARG NVCC_TOOLCHAIN_SHA256_AARCH64=2432ef8a7c12d0a9ce3332a8af42b123c07f256390b3390802b1b2c6254c6c74
+RUN if [ -n "${MIN_NVCC_TOOLCHAIN_VERSION}" ]; then \
+        CURRENT_VERSION=$(nvcc --version | sed -n 's/.*V\([0-9.]*\).*/\1/p') && \
+        SMALLEST=$(printf '%s\n' "${MIN_NVCC_TOOLCHAIN_VERSION}" "${CURRENT_VERSION}" | sort -V | head -n1) && \
+        if [ "${SMALLEST}" = "${MIN_NVCC_TOOLCHAIN_VERSION}" ]; then \
+            echo "System nvcc ${CURRENT_VERSION} >= ${MIN_NVCC_TOOLCHAIN_VERSION}, skipping toolchain download"; \
+        else \
+            ARCH=$(uname -m) && \
+            case "${ARCH}" in \
+                x86_64)  EXPECTED_SHA256="${NVCC_TOOLCHAIN_SHA256_X86_64}" ;; \
+                aarch64) EXPECTED_SHA256="${NVCC_TOOLCHAIN_SHA256_AARCH64}" ;; \
+                *) echo "Unsupported architecture: ${ARCH}" >&2; exit 1 ;; \
+            esac && \
+            curl -sL -o cuda_nvcc.tar.xz \
+                "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/linux-${ARCH}/cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive.tar.xz" && \
+            echo "${EXPECTED_SHA256}  cuda_nvcc.tar.xz" | sha256sum -c && \
+            tar xJf cuda_nvcc.tar.xz && \
+            cp -f cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive/bin/nvcc /usr/local/cuda/bin/nvcc && \
+            cp -f cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive/bin/ptxas /usr/local/cuda/bin/ptxas && \
+            cp -f cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive/nvvm/bin/cicc /usr/local/cuda/nvvm/bin/cicc && \
+            rm -rf cuda_nvcc-linux-${ARCH}-${MIN_NVCC_TOOLCHAIN_VERSION}-archive cuda_nvcc.tar.xz && \
+            echo "Upgraded nvcc/ptxas/cicc from ${CURRENT_VERSION} to ${MIN_NVCC_TOOLCHAIN_VERSION}"; \
+        fi; \
+    fi
+
+# Patch PyTorch cmake to allow nvcc/header version mismatch (for newer nvcc toolchain).
+# Stages that reinstall torch (csrc-build, build) must re-apply this patch.
+RUN python3 -c "\
+import pathlib; \
+files = list(pathlib.Path('/opt/venv').rglob('Caffe2/public/cuda.cmake')); \
+assert files, 'ERROR: Caffe2/public/cuda.cmake not found under /opt/venv'; \
+[print(f'Patching {f}') or f.write_text(f.read_text().replace('FATAL_ERROR \"FindCUDA says', 'WARNING \"FindCUDA says')) for f in files]"
 #################### BUILD BASE IMAGE ####################
 
 #################### CSRC BUILD IMAGE ####################
@@ -222,6 +266,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
     fi
 
+# Patch PyTorch cmake to allow nvcc/header version mismatch (for newer nvcc toolchain).
+# This must run after pip install since torch reinstall overwrites the cmake files.
+RUN python3 -c "\
+import pathlib; \
+files = list(pathlib.Path('/opt/venv').rglob('Caffe2/public/cuda.cmake')); \
+assert files, 'ERROR: Caffe2/public/cuda.cmake not found under /opt/venv'; \
+[print(f'Patching {f}') or f.write_text(f.read_text().replace('FATAL_ERROR \"FindCUDA says', 'WARNING \"FindCUDA says')) for f in files]"
+
 WORKDIR /workspace
 
 COPY pyproject.toml setup.py CMakeLists.txt ./
@@ -388,6 +440,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
     fi
 
+# Patch PyTorch cmake to allow nvcc/header version mismatch (for newer nvcc toolchain).
+# This must run after pip install since torch reinstall overwrites the cmake files.
+RUN python3 -c "\
+import pathlib; \
+files = list(pathlib.Path('/opt/venv').rglob('Caffe2/public/cuda.cmake')); \
+assert files, 'ERROR: Caffe2/public/cuda.cmake not found under /opt/venv'; \
+[print(f'Patching {f}') or f.write_text(f.read_text().replace('FATAL_ERROR \"FindCUDA says', 'WARNING \"FindCUDA says')) for f in files]"
+
 WORKDIR /workspace
 
 # Copy pre-built csrc wheel directly

docker/versions.json

@@ -2,16 +2,16 @@
   "_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
   "variable": {
     "CUDA_VERSION": {
-      "default": "12.9.1"
+      "default": "12.8.1"
     },
     "PYTHON_VERSION": {
       "default": "3.12"
     },
     "BUILD_BASE_IMAGE": {
-      "default": "nvidia/cuda:12.9.1-devel-ubuntu20.04"
+      "default": "nvidia/cuda:12.8.1-devel-ubuntu20.04"
     },
     "FINAL_BASE_IMAGE": {
-      "default": "nvidia/cuda:12.9.1-base-ubuntu22.04"
+      "default": "nvidia/cuda:12.8.1-base-ubuntu22.04"
     },
     "GET_PIP_URL": {
       "default": "https://bootstrap.pypa.io/get-pip.py"
@@ -31,6 +31,15 @@
     "TORCH_CUDA_ARCH_LIST": {
       "default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0"
     },
+    "MIN_NVCC_TOOLCHAIN_VERSION": {
+      "default": "12.9.86"
+    },
+    "NVCC_TOOLCHAIN_SHA256_X86_64": {
+      "default": "7a1a5b652e5ef85c82b721d10672fc9a2dbaab44e9bd3c65a69517bf53998c35"
+    },
+    "NVCC_TOOLCHAIN_SHA256_AARCH64": {
+      "default": "2432ef8a7c12d0a9ce3332a8af42b123c07f256390b3390802b1b2c6254c6c74"
+    },
     "MAX_JOBS": {
       "default": "2"
     },

requirements/test.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu129 --python-platform x86_64-manylinux_2_28 --python-version 3.12
+#    uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu128 --python-platform x86_64-manylinux_2_28 --python-version 3.12
 absl-py==2.1.0
     # via
     #   rouge-score
@@ -574,28 +574,28 @@ numpy==2.2.6
     #   tritonclient
     #   vocos
     #   xarray
-nvidia-cublas-cu12==12.9.1.4
+nvidia-cublas-cu12==12.8.4.1
     # via
     #   nvidia-cudnn-cu12
     #   nvidia-cusolver-cu12
     #   torch
-nvidia-cuda-cupti-cu12==12.9.79
+nvidia-cuda-cupti-cu12==12.8.90
     # via torch
-nvidia-cuda-nvrtc-cu12==12.9.86
+nvidia-cuda-nvrtc-cu12==12.8.93
     # via torch
-nvidia-cuda-runtime-cu12==12.9.79
+nvidia-cuda-runtime-cu12==12.8.90
     # via torch
 nvidia-cudnn-cu12==9.10.2.21
     # via torch
-nvidia-cufft-cu12==11.4.1.4
+nvidia-cufft-cu12==11.3.3.83
     # via torch
-nvidia-cufile-cu12==1.14.1.1
+nvidia-cufile-cu12==1.13.1.3
     # via torch
-nvidia-curand-cu12==10.3.10.19
+nvidia-curand-cu12==10.3.9.90
     # via torch
-nvidia-cusolver-cu12==11.7.5.82
+nvidia-cusolver-cu12==11.7.3.90
     # via torch
-nvidia-cusparse-cu12==12.5.10.65
+nvidia-cusparse-cu12==12.5.8.93
     # via
     #   nvidia-cusolver-cu12
     #   torch
@@ -603,7 +603,7 @@ nvidia-cusparselt-cu12==0.7.1
     # via torch
 nvidia-nccl-cu12==2.27.5
     # via torch
-nvidia-nvjitlink-cu12==12.9.86
+nvidia-nvjitlink-cu12==12.8.93
     # via
     #   nvidia-cufft-cu12
     #   nvidia-cusolver-cu12
@@ -611,7 +611,7 @@ nvidia-nvjitlink-cu12==12.9.86
     #   torch
 nvidia-nvshmem-cu12==3.4.5
     # via torch
-nvidia-nvtx-cu12==12.9.79
+nvidia-nvtx-cu12==12.8.90
     # via torch
 omegaconf==2.3.0
     # via
@@ -1154,7 +1154,7 @@ tomli==2.2.1
     # via schemathesis
 tomli-w==1.2.0
     # via schemathesis
-torch==2.10.0+cu129
+torch==2.10.0+cu128
     # via
     #   -r requirements/test.in
     #   accelerate
@@ -1179,7 +1179,7 @@ torch==2.10.0+cu129
     #   torchvision
     #   vector-quantize-pytorch
     #   vocos
-torchaudio==2.10.0+cu129
+torchaudio==2.10.0+cu128
     # via
     #   -r requirements/test.in
     #   encodec
@@ -1192,7 +1192,7 @@ torchmetrics==1.7.4
     #   pytorch-lightning
     #   terratorch
     #   torchgeo
-torchvision==0.25.0+cu129
+torchvision==0.25.0+cu128
     # via
     #   -r requirements/test.in
     #   lightly