nix-builder: add support for Torch 2.15.0 RC3

docs/source/builder/build-variants.md

@@ -9,6 +9,7 @@ available. This list will be updated as new PyTorch versions are released.
 
 - `torch213-cpu-aarch64-darwin`
 - `torch214-cpu-aarch64-darwin`
+- `torch215-cpu-aarch64-darwin`
 
 ## Metal aarch64-darwin
 
@@ -17,11 +18,13 @@ available. This list will be updated as new PyTorch versions are released.
 - `torch212-metal-aarch64-darwin`
 - `torch213-metal-aarch64-darwin`
 - `torch214-metal-aarch64-darwin`
+- `torch215-metal-aarch64-darwin`
 
 ## CPU aarch64-linux
 
 - `torch213-cpu-aarch64-linux`
 - `torch214-cpu-aarch64-linux`
+- `torch215-cpu-aarch64-linux`
 
 ## CUDA aarch64-linux
 
@@ -31,11 +34,15 @@ available. This list will be updated as new PyTorch versions are released.
 - `torch214-cu126-aarch64-linux`
 - `torch214-cu130-aarch64-linux`
 - `torch214-cu132-aarch64-linux`
+- `torch215-cu130-aarch64-linux`
+- `torch215-cu132-aarch64-linux`
+- `torch215-cu134-aarch64-linux`
 
 ## CPU x86_64-linux
 
 - `torch213-cpu-x86_64-linux`
 - `torch214-cpu-x86_64-linux`
+- `torch215-cpu-x86_64-linux`
 
 ## CUDA x86_64-linux
 
@@ -45,13 +52,18 @@ available. This list will be updated as new PyTorch versions are released.
 - `torch214-cu126-x86_64-linux`
 - `torch214-cu130-x86_64-linux`
 - `torch214-cu132-x86_64-linux`
+- `torch215-cu130-x86_64-linux`
+- `torch215-cu132-x86_64-linux`
+- `torch215-cu134-x86_64-linux`
 
 ## ROCm x86_64-linux
 
 - `torch213-rocm71-x86_64-linux`
 - `torch213-rocm72-x86_64-linux`
 - `torch214-rocm714-x86_64-linux`
 - `torch214-rocm72-x86_64-linux`
+- `torch215-rocm100-x86_64-linux`
+- `torch215-rocm714-x86_64-linux`
 
 ## XPU x86_64-linux
 
@@ -60,6 +72,7 @@ available. This list will be updated as new PyTorch versions are released.
 - `torch212-xpu20253-x86_64-linux`
 - `torch213-xpu20260-x86_64-linux`
 - `torch214-xpu20261-x86_64-linux`
+- `torch215-xpu20261-x86_64-linux`
 
 ## Python-only kernels
 

nix-builder/build-variants.json

@@ -2,55 +2,68 @@
   "aarch64-darwin": {
     "cpu": [
       "torch213-cpu-aarch64-darwin",
-      "torch214-cpu-aarch64-darwin"
+      "torch214-cpu-aarch64-darwin",
+      "torch215-cpu-aarch64-darwin"
     ],
     "metal": [
       "torch210-metal-aarch64-darwin",
       "torch211-metal-aarch64-darwin",
       "torch212-metal-aarch64-darwin",
       "torch213-metal-aarch64-darwin",
-      "torch214-metal-aarch64-darwin"
+      "torch214-metal-aarch64-darwin",
+      "torch215-metal-aarch64-darwin"
     ]
   },
   "aarch64-linux": {
     "cpu": [
       "torch213-cpu-aarch64-linux",
-      "torch214-cpu-aarch64-linux"
+      "torch214-cpu-aarch64-linux",
+      "torch215-cpu-aarch64-linux"
     ],
     "cuda": [
       "torch213-cu126-aarch64-linux",
       "torch213-cu130-aarch64-linux",
       "torch213-cu132-aarch64-linux",
       "torch214-cu126-aarch64-linux",
       "torch214-cu130-aarch64-linux",
-      "torch214-cu132-aarch64-linux"
+      "torch214-cu132-aarch64-linux",
+      "torch215-cu130-aarch64-linux",
+      "torch215-cu132-aarch64-linux",
+      "torch215-cu134-aarch64-linux"
     ]
   },
   "x86_64-linux": {
     "cpu": [
       "torch213-cpu-x86_64-linux",
-      "torch214-cpu-x86_64-linux"
+      "torch214-cpu-x86_64-linux",
+      "torch215-cpu-x86_64-linux"
     ],
     "cuda": [
       "torch213-cu126-x86_64-linux",
       "torch213-cu130-x86_64-linux",
       "torch213-cu132-x86_64-linux",
       "torch214-cu126-x86_64-linux",
       "torch214-cu130-x86_64-linux",
-      "torch214-cu132-x86_64-linux"
+      "torch214-cu132-x86_64-linux",
+      "torch215-cu130-x86_64-linux",
+      "torch215-cu132-x86_64-linux",
+      "torch215-cu134-x86_64-linux"
     ],
     "rocm": [
       "torch213-rocm71-x86_64-linux",
       "torch213-rocm72-x86_64-linux",
       "torch214-rocm714-x86_64-linux",
-      "torch214-rocm72-x86_64-linux"
+      "torch214-rocm72-x86_64-linux",
+      "torch215-rocm100-x86_64-linux",
+      "torch215-rocm714-x86_64-linux"
     ],
     "xpu": [
       "torch210-xpu20253-x86_64-linux",
       "torch211-xpu20253-x86_64-linux",
       "torch212-xpu20253-x86_64-linux",
       "torch213-xpu20260-x86_64-linux",
-      "torch214-xpu20261-x86_64-linux"
+      "torch214-xpu20261-x86_64-linux",
+      "torch215-xpu20261-x86_64-linux"
     ]
   }
 }

nix-builder/overlay.nix

@@ -108,21 +108,29 @@ final: prev:
           triton_3_7_0
           triton_3_7_1
           triton_3_8_0
-          triton_3_8_0-nightly
+          triton_3_9_0
+          ;
+        inherit (triton-rocm)
+          triton-rocm_3_7_0
+          triton-rocm_3_7_1
+          triton-rocm_3_8_0
+          triton-rocm_3_9_0
           ;
-        inherit (triton-rocm) triton-rocm_3_7_0;
         inherit (triton-xpu)
           triton-xpu_3_6_0
           triton-xpu_3_7_0
           triton-xpu_3_7_1
           triton-xpu_3_7_2
           triton-xpu_3_8_0
+          triton-xpu_3_9_0
           ;
 
         cuda-bindings = python-self.callPackage ./pkgs/python-modules/cuda-bindings { };
 
         cuda-pathfinder = python-self.callPackage ./pkgs/python-modules/cuda-pathfinder { };
 
+        cupti-python = python-self.callPackage ./pkgs/python-modules/cupti-python { };
+
         # Starting with the CUDA 12.8 version, cuda-python is a metapackage
         # that pulls in relevant dependencies. For CUDA 12.6 it is just
         # cuda-bindings.
@@ -267,26 +275,25 @@ final: prev:
         torch-bin_2_13 = mkTorch {
           version = "2.13";
           triton-cuda = triton_3_7_1;
-          triton-rocm = triton-rocm_3_7_0;
+          triton-rocm = triton-rocm_3_7_1;
           triton-xpu = triton-xpu_3_7_2;
           xpuPackages = final.xpuPackages_2026_0_0;
         };
 
         torch-bin_2_14 = mkTorch {
           version = "2.14";
           triton-cuda = triton_3_8_0;
-          triton-rocm = triton-rocm_3_7_0;
+          triton-rocm = triton-rocm_3_8_0;
           triton-xpu = triton-xpu_3_8_0;
           xpuPackages = final.xpuPackages_2026_1_0;
         };
 
-        # Nightly for CUDA 13.4 development work.
         torch-bin_2_15 = mkTorch {
           version = "2.15";
-          triton-cuda = triton_3_8_0-nightly;
-          triton-rocm = null;
-          triton-xpu = null;
-          xpuPackages = null;
+          triton-cuda = triton_3_9_0;
+          triton-rocm = triton-rocm_3_9_0;
+          triton-xpu = triton-xpu_3_9_0;
+          xpuPackages = final.xpuPackages_2026_1_2;
         };
 
         transformers = python-super.transformers.overridePythonAttrs (prevAttrs: rec {

nix-builder/pkgs/aotriton/default.nix

@@ -1,11 +1,12 @@
 {
   callPackage,
   fetchurl,
+  rocmPackages,
   stdenvNoCC,
 }:
 
 let
-  generic = callPackage ./generic.nix { };
+  generic = callPackage ./generic.nix { inherit rocmPackages; };
   mkImages =
     version: srcs:
     stdenvNoCC.mkDerivation {
@@ -130,4 +131,49 @@ in
     ];
   };
 
+  aotriton_0_14_2 = generic rec {
+    version = "0.14.2b";
+
+    # Upstream names the ROCm 10.0 package after its HIP version (7.15).
+    hipVersions = {
+      "10.0" = "7.15";
+    };
+
+    hashes = {
+      "7.14" = "sha256-3HVoJLXGOmiRefdlxR8qp0fT1e6EKxe92SOK3PI3+Rs=";
+      "10.0" = "sha256-bthA43swN/1+Znfn0p1KuP0VeHV4IM4AtAF2u7UasFM=";
+    };
+
+    images = mkImages version [
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.14.2b/aotriton-0.14.2b-images-amd-gfx90a.tar.gz";
+        hash = "sha256-fMKTgDqnO9IjVU1Rp3+rUP03sL251eh09+KZf0OQucU=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.14.2b/aotriton-0.14.2b-images-amd-gfx942.tar.gz";
+        hash = "sha256-6uXTbJdNl08IXPvez3d2DbcTJPqYFO1Q2KP8Qr++PkE=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.14.2b/aotriton-0.14.2b-images-amd-gfx950.tar.gz";
+        hash = "sha256-I3Fsv4Vd+d8cp9SAqQ4PspsDVrtCQBXEJvwOWpc/aCs=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.14.2b/aotriton-0.14.2b-images-amd-gfx110x.tar.gz";
+        hash = "sha256-pU/s+bThZCVoF4XrpHWRLrmudh41pN8+NXm9vXxv5Ns=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.14.2b/aotriton-0.14.2b-images-amd-gfx115x.tar.gz";
+        hash = "sha256-5L2nb9aFenM817Sz65zLIcD0e1Zc4wR3emG7sT0SrW4=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.14.2b/aotriton-0.14.2b-images-amd-gfx120x.tar.gz";
+        hash = "sha256-8XNnHMURKJfhNPHSCb1vVy7NUH9NaNyEy28Htn4T6FE=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.14.2b/aotriton-0.14.2b-images-amd-gfx1250.tar.gz";
+        hash = "sha256-ajVL+EoY12ia/675U+pKLR9pQ1adpKTKJv1ea162eg8=";
+      })
+    ];
+  };
+
 }

nix-builder/pkgs/aotriton/generic.nix

@@ -11,10 +11,15 @@
   version,
   images,
   hashes,
+
+  # Maps a ROCm version to the HIP version that upstream uses in the
+  # binary package name, for ROCm versions where they differ.
+  hipVersions ? { },
 }:
 
 let
   rocmVersion = lib.versions.majorMinor rocmPackages.version;
+  hipVersion = hipVersions.${rocmVersion} or rocmVersion;
   hash =
     hashes.${rocmVersion}
       or (throw "aotriton ${version} binary package is not specified for ROCm ${rocmVersion}");
@@ -24,7 +29,7 @@ stdenv.mkDerivation {
   inherit version;
 
   src = fetchurl {
-    url = "https://github.com/ROCm/aotriton/releases/download/${version}/aotriton-${version}-manylinux_2_28_x86_64-rocm${rocmVersion}-shared.tar.gz";
+    url = "https://github.com/ROCm/aotriton/releases/download/${version}/aotriton-${version}-manylinux_2_28_x86_64-rocm${hipVersion}-shared.tar.gz";
     inherit hash;
   };
 

nix-builder/pkgs/python-modules/cuda-bindings/default.nix

@@ -52,6 +52,24 @@ let
   format = "wheel";
   pyShortVersion = "cp" + builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
 
+  cudaMajor = lib.versions.major cudaPackages.cudaMajorMinorVersion;
+
+  # The extension modules load CUDA libraries at runtime through
+  # cuda-pathfinder, add dependencies so that they can be loaded
+  # and lacated through rpath.
+  linkedLibraries = {
+    "_bindings/cynvrtc" = "libnvrtc.so.${cudaMajor}";
+    "_internal/cufile" = "libcufile.so.0";
+    "_internal/nvjitlink" = "libnvJitLink.so.${cudaMajor}";
+    "_internal/nvvm" = "libnvvm.so.4";
+  }
+  // lib.optionalAttrs (cudaMajor == "12") {
+    "_internal/_nvml" = "libnvidia-ml.so.1";
+  }
+  // lib.optionalAttrs (cudaMajor != "12") {
+    "_internal/nvfatbin" = "libnvfatbin.so.${cudaMajor}";
+    "_internal/nvml" = "libnvidia-ml.so.1";
+  };
 in
 buildPythonPackage {
   inherit format;
@@ -68,11 +86,40 @@ buildPythonPackage {
   };
 
   nativeBuildInputs = [
-    autoAddDriverRunpath
+    # autoAddDriverRunpath must come after autoPatchelfHook, because
+    # autoPatchelf rewrites the runpath and would drop the driver path.
     autoPatchelfHook
+    autoAddDriverRunpath
     pythonWheelDepsCheckHook
   ];
 
+  buildInputs = map lib.getLib (
+    with cudaPackages;
+    [
+      cuda_nvrtc
+      libcufile
+      libnvjitlink
+    ]
+    ++ lib.optionals (cudaMajor != "12") [ libnvfatbin ]
+  );
+
+  # libnvidia-ml is part of the driver, found through the driver runpath.
+  autoPatchelfIgnoreMissingDeps = [ "libnvidia-ml.so.1" ];
+
+  postInstall = ''
+    ${lib.concatStrings (
+      lib.mapAttrsToList (module: soname: ''
+        patchelf --add-needed ${soname} \
+          $out/${python.sitePackages}/cuda/bindings/${module}.*.so
+      '') linkedLibraries
+    )}
+  '';
+
+  # libnvvm is in a subdirectory of cuda_nvcc.
+  preFixup = ''
+    addAutoPatchelfSearchPath --no-recurse ${lib.getLib cudaPackages.cuda_nvcc}/nvvm/lib
+  '';
+
   dependencies = [ cuda-pathfinder ];
 
   pythonImportsCheck = [ "cuda.bindings" ];

nix-builder/pkgs/python-modules/cuda-pathfinder/default.nix

@@ -2,23 +2,30 @@
   buildPythonPackage,
   fetchFromGitHub,
   setuptools,
+  setuptools-scm,
 }:
 
 buildPythonPackage rec {
   pname = "cuda-pathfinder";
-  version = "1.3.3";
+  version = "1.4.1";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "NVIDIA";
     repo = "cuda-python";
     tag = "${pname}-v${version}";
-    hash = "sha256-aF06WwThJmGEqUcVP4dooym1uqYjUM45arqZaxjlTuA=";
+    hash = "sha256-rdDJGkA+20NB5dvwJHOSUGBzvx3lfLuuUqMhvz1dmXY=";
   };
 
   sourceRoot = "source/cuda_pathfinder";
 
-  build-system = [ setuptools ];
+  build-system = [
+    setuptools
+    setuptools-scm
+  ];
+
+  # The source tarball has no git metadata for setuptools-scm.
+  env.SETUPTOOLS_SCM_PRETEND_VERSION = version;
 
   pythonImportsCheck = [ "cuda.pathfinder" ];
 }