Add support for Torch 2.14 (2.14-rc6)

.github/workflows/test_e2e.yaml

@@ -84,7 +84,8 @@ jobs:
         id: variant
         run: |
           cd /tmp/${{ env.E2E_REPO_NAME }}
-          VARIANT=$(nix run $GITHUB_WORKSPACE#kernel-builder -- list-variants . | tail -1)
+          # NOTE: Remove Torch 2.14 grep once it is released.
+          VARIANT=$(nix run $GITHUB_WORKSPACE#kernel-builder -- list-variants . | grep -v torch214 | tail -1)
           echo "name=$VARIANT" >> $GITHUB_OUTPUT
           echo "Building variant: $VARIANT"
 

docs/source/builder/build-variants.md

@@ -9,16 +9,19 @@ available. This list will be updated as new PyTorch versions are released.
 
 - `torch212-cpu-aarch64-darwin`
 - `torch213-cpu-aarch64-darwin`
+- `torch214-cpu-aarch64-darwin`
 
 ## Metal aarch64-darwin
 
 - `torch212-metal-aarch64-darwin`
 - `torch213-metal-aarch64-darwin`
+- `torch214-metal-aarch64-darwin`
 
 ## CPU aarch64-linux
 
 - `torch212-cxx11-cpu-aarch64-linux`
 - `torch213-cxx11-cpu-aarch64-linux`
+- `torch214-cxx11-cpu-aarch64-linux`
 
 ## CUDA aarch64-linux
 
@@ -28,11 +31,15 @@ available. This list will be updated as new PyTorch versions are released.
 - `torch213-cxx11-cu126-aarch64-linux`
 - `torch213-cxx11-cu130-aarch64-linux`
 - `torch213-cxx11-cu132-aarch64-linux`
+- `torch214-cxx11-cu126-aarch64-linux`
+- `torch214-cxx11-cu130-aarch64-linux`
+- `torch214-cxx11-cu132-aarch64-linux`
 
 ## CPU x86_64-linux
 
 - `torch212-cxx11-cpu-x86_64-linux`
 - `torch213-cxx11-cpu-x86_64-linux`
+- `torch214-cxx11-cpu-x86_64-linux`
 
 ## CUDA x86_64-linux
 
@@ -42,18 +49,23 @@ available. This list will be updated as new PyTorch versions are released.
 - `torch213-cxx11-cu126-x86_64-linux`
 - `torch213-cxx11-cu130-x86_64-linux`
 - `torch213-cxx11-cu132-x86_64-linux`
+- `torch214-cxx11-cu126-x86_64-linux`
+- `torch214-cxx11-cu130-x86_64-linux`
+- `torch214-cxx11-cu132-x86_64-linux`
 
 ## ROCm x86_64-linux
 
 - `torch212-cxx11-rocm71-x86_64-linux`
 - `torch212-cxx11-rocm72-x86_64-linux`
 - `torch213-cxx11-rocm71-x86_64-linux`
 - `torch213-cxx11-rocm72-x86_64-linux`
+- `torch214-cxx11-rocm72-x86_64-linux`
 
 ## XPU x86_64-linux
 
 - `torch212-cxx11-xpu20253-x86_64-linux`
 - `torch213-cxx11-xpu20260-x86_64-linux`
+- `torch214-cxx11-xpu20261-x86_64-linux`
 
 ## Python-only kernels
 

nix-builder/build-variants.json

@@ -2,49 +2,61 @@
   "aarch64-darwin": {
     "cpu": [
       "torch212-cpu-aarch64-darwin",
-      "torch213-cpu-aarch64-darwin"
+      "torch213-cpu-aarch64-darwin",
+      "torch214-cpu-aarch64-darwin"
     ],
     "metal": [
       "torch212-metal-aarch64-darwin",
-      "torch213-metal-aarch64-darwin"
+      "torch213-metal-aarch64-darwin",
+      "torch214-metal-aarch64-darwin"
     ]
   },
   "aarch64-linux": {
     "cpu": [
       "torch212-cxx11-cpu-aarch64-linux",
-      "torch213-cxx11-cpu-aarch64-linux"
+      "torch213-cxx11-cpu-aarch64-linux",
+      "torch214-cxx11-cpu-aarch64-linux"
     ],
     "cuda": [
       "torch212-cxx11-cu126-aarch64-linux",
       "torch212-cxx11-cu130-aarch64-linux",
       "torch212-cxx11-cu132-aarch64-linux",
       "torch213-cxx11-cu126-aarch64-linux",
       "torch213-cxx11-cu130-aarch64-linux",
-      "torch213-cxx11-cu132-aarch64-linux"
+      "torch213-cxx11-cu132-aarch64-linux",
+      "torch214-cxx11-cu126-aarch64-linux",
+      "torch214-cxx11-cu130-aarch64-linux",
+      "torch214-cxx11-cu132-aarch64-linux"
     ]
   },
   "x86_64-linux": {
     "cpu": [
       "torch212-cxx11-cpu-x86_64-linux",
-      "torch213-cxx11-cpu-x86_64-linux"
+      "torch213-cxx11-cpu-x86_64-linux",
+      "torch214-cxx11-cpu-x86_64-linux"
     ],
     "cuda": [
       "torch212-cxx11-cu126-x86_64-linux",
       "torch212-cxx11-cu130-x86_64-linux",
       "torch212-cxx11-cu132-x86_64-linux",
       "torch213-cxx11-cu126-x86_64-linux",
       "torch213-cxx11-cu130-x86_64-linux",
-      "torch213-cxx11-cu132-x86_64-linux"
+      "torch213-cxx11-cu132-x86_64-linux",
+      "torch214-cxx11-cu126-x86_64-linux",
+      "torch214-cxx11-cu130-x86_64-linux",
+      "torch214-cxx11-cu132-x86_64-linux"
     ],
     "rocm": [
       "torch212-cxx11-rocm71-x86_64-linux",
       "torch212-cxx11-rocm72-x86_64-linux",
       "torch213-cxx11-rocm71-x86_64-linux",
-      "torch213-cxx11-rocm72-x86_64-linux"
+      "torch213-cxx11-rocm72-x86_64-linux",
+      "torch214-cxx11-rocm72-x86_64-linux"
     ],
     "xpu": [
       "torch212-cxx11-xpu20253-x86_64-linux",
-      "torch213-cxx11-xpu20260-x86_64-linux"
+      "torch213-cxx11-xpu20260-x86_64-linux",
+      "torch214-cxx11-xpu20261-x86_64-linux"
     ]
   }
 }

nix-builder/overlay.nix

@@ -78,11 +78,13 @@ final: prev:
         inherit (triton)
           triton_3_7_0
           triton_3_7_1
+          triton_3_8_0
           ;
         inherit (triton-rocm) triton-rocm_3_7_0;
         inherit (triton-xpu)
           triton-xpu_3_7_1
           triton-xpu_3_7_2
+          triton-xpu_3_8_0
           ;
 
         cuda-bindings = python-self.callPackage ./pkgs/python-modules/cuda-bindings { };
@@ -229,6 +231,14 @@ final: prev:
           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-xpu = triton-xpu_3_8_0;
+          xpuPackages = final.xpuPackages_2026_1_0;
+        };
+
         transformers = python-super.transformers.overridePythonAttrs (prevAttrs: rec {
           version = "5.3.0";
           src = python-super.fetchPypi {
@@ -296,6 +306,7 @@ final: prev:
     xpuVersions = [
       "2025.3.2"
       "2026.0.0"
+      "2026.1.0"
     ];
     newXpuPackages = final.callPackage ./pkgs/xpu-packages { };
   in

nix-builder/pkgs/aotriton/default.nix

@@ -90,4 +90,43 @@ in
     ];
   };
 
+  aotriton_0_13 = generic rec {
+    version = "0.13b";
+
+    hashes = {
+      "7.2" = "sha256-HN7rt+9hq2kfuh2B2pGbnbXYvvKCaciSowvROgSVt6A=";
+    };
+
+    images = mkImages version [
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.13b/aotriton-0.13b-images-amd-gfx90a.tar.gz";
+        hash = "sha256-o9GmhozikLqBGGGCBwk+eFJS7/ThimT0lXUstaA//tY=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.13b/aotriton-0.13b-images-amd-gfx942.tar.gz";
+        hash = "sha256-zNvH49loOb5Ile4ATyFTHMVdWQyQGJN7njFLujY7OSc=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.13b/aotriton-0.13b-images-amd-gfx950.tar.gz";
+        hash = "sha256-UY/QcusFlI/ApsJaIIMlkcZAbfhl47aRoq7/P9TFzh0=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.13b/aotriton-0.13b-images-amd-gfx110x.tar.gz";
+        hash = "sha256-7+dz56LIrcmV2Q7NDayi24AShURe4QQy3yspxn1bEdI=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.13b/aotriton-0.13b-images-amd-gfx115x.tar.gz";
+        hash = "sha256-G8UOiqi2vaNBDpKIbMyo/UXfPmCmy9qf/FiyxUHv1cI=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.13b/aotriton-0.13b-images-amd-gfx120x.tar.gz";
+        hash = "sha256-akZdvAMUi7qKLXjEwqPLgxVeygD09/dJ5nZALXZglow=";
+      })
+      (fetchurl {
+        url = "https://github.com/ROCm/aotriton/releases/download/0.13b/aotriton-0.13b-images-amd-gfx1250.tar.gz";
+        hash = "sha256-Sq9x1uUQVJ1ZN1fl+IWY3x5KKcvNkfcHUO6KdvZcAn8=";
+      })
+    ];
+  };
+
 }

nix-builder/pkgs/python-modules/torch/binary/generic.nix

@@ -62,6 +62,8 @@
   effectiveStdenv ? if cudaSupport then cudaPackages.backendStdenv else stdenv,
 }:
 let
+  inherit (lib) versionAtLeast;
+
   torchMajorMinor = lib.versions.majorMinor version;
 
   effectiveTriton =
@@ -77,6 +79,7 @@ let
   aotritonVersions = with rocmPackages; {
     "2.12" = aotriton_0_11_2;
     "2.13" = aotriton_0_12;
+    "2.14" = aotriton_0_13;
   };
 
   aotriton =
@@ -115,7 +118,7 @@ let
         rocsparse
         roctracer
       ]
-      ++ lib.optionals (lib.versionAtLeast version "2.12") [
+      ++ lib.optionals (versionAtLeast version "2.12") [
         rocprofiler-sdk
       ];
 
@@ -227,7 +230,7 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
         ))
       ]
     )
-    ++ lib.optionals (cudaSupport && lib.versionAtLeast version "2.9") [
+    ++ lib.optionals (cudaSupport && versionAtLeast version "2.9") [
       cudaPackages.libnvshmem
     ]
     ++ lib.optionals rocmSupport ([
@@ -273,10 +276,10 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
   ++ lib.optionals tritonSupport [
     effectiveTriton
   ]
-  ++ lib.optionals (cudaSupport && lib.versionAtLeast version "2.10") [
+  ++ lib.optionals (cudaSupport && versionAtLeast version "2.10") [
     cuda-bindings
   ]
-  ++ lib.optionals (xpuSupport && lib.versionAtLeast version "2.13") [
+  ++ lib.optionals (xpuSupport && versionAtLeast version "2.13") [
     pyzes
   ];
 
@@ -403,7 +406,7 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
   # We need to add this rpath after postFixup (otherwise it might be ordered
   # before autoPatchelfHook), but before the import check to ensure that the
   # rpath is used during the import check.
-  preInstallCheck = lib.optionalString (rocmSupport && torchMajorMinor == "2.13") ''
+  preInstallCheck = lib.optionalString (rocmSupport && versionAtLeast torchMajorMinor "2.13") ''
     # libtorch_shmem dynamically loads libnuma. For added fun, if libnuma
     # cannot be found, it exit()s inside a constructor, calling an atexit()
     # hook, which then proceeds to deadlock.

nix-builder/pkgs/python-modules/torch/binary/torch-versions-hash.json

@@ -143,5 +143,83 @@
         "version": "2.13.0"
       }
     }
+  },
+  "2.14": {
+    "x86_64-linux": {
+      "cu130": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcu130-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-Rek9qAsDlogrAyWNaX+obKDwZt0x0VdwRNHpTdhymxs=",
+        "version": "2.14.0"
+      },
+      "cu132": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcu132-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-wYSnxRk5txrxNkbv06j3Nki9N2xpIumcTGOE02dgXVg=",
+        "version": "2.14.0"
+      },
+      "cpu": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcpu-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-J4TuoRFcRBGEHDvMEbfgUm7R00fcz4GwFZGIXlcmDFM=",
+        "version": "2.14.0"
+      },
+      "cu134": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcu134-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-tF/80uEFZy+BeCg6yUICZlR4iLxGG2G7tQtSIc3dgps=",
+        "version": "2.14.0"
+      },
+      "cu126": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcu126-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-I8A6bMBJeOyv1onF/k/d0m53elSTRfw0trjxgxJnyPo=",
+        "version": "2.14.0"
+      },
+      "xpu": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-biy7nCq44+TNguilNL+vVh8vUMqxJm62NI9TKwMnYiU=",
+        "version": "2.14.0"
+      },
+      "rocm714": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Brocm7.14-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-AzDhFg7FWQPXQWTlAALNInmfH6zSGzQ85GOdBeNnr1Q=",
+        "version": "2.14.0"
+      },
+      "rocm72": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Brocm7.2-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-vKWLdiy6vxl7xWeAoXMI41SNINmoRo5J+RSzV/pFtWg=",
+        "version": "2.14.0"
+      }
+    },
+    "aarch64-linux": {
+      "cu126": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcu126-cp313-cp313-manylinux_2_28_aarch64.whl",
+        "hash": "sha256-OHT6clnwB+1o9Kts/uPllSUkXNQvYhnfHMx0LaLWOoQ=",
+        "version": "2.14.0"
+      },
+      "cu130": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcu130-cp313-cp313-manylinux_2_28_aarch64.whl",
+        "hash": "sha256-98TXTngVFaFh4gBBvjuHx/e1Nzl4c3xgMnFseLW83VU=",
+        "version": "2.14.0"
+      },
+      "cu134": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcu134-cp313-cp313-manylinux_2_28_aarch64.whl",
+        "hash": "sha256-9FLL2bhpyXxMdrOnmHA0pWoUmNTcMRzA0O/gppiOJnE=",
+        "version": "2.14.0"
+      },
+      "cpu": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcpu-cp313-cp313-manylinux_2_28_aarch64.whl",
+        "hash": "sha256-nRYFP0amJ/JSrxinD74EGQDDcZhZUvWeS66TrxddtAU=",
+        "version": "2.14.0"
+      },
+      "cu132": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0%2Bcu132-cp313-cp313-manylinux_2_28_aarch64.whl",
+        "hash": "sha256-0nJppAuxuA9qxjakegmyeL1sXQvMDxkKWpH21xz8XPM=",
+        "version": "2.14.0"
+      }
+    },
+    "aarch64-darwin": {
+      "cpu": {
+        "url": "https://huggingface.co/buckets/danieldk/pytorch-rc/resolve/2.14.0/rc6/torch-2.14.0-cp313-cp313-macosx_14_0_arm64.whl",
+        "hash": "sha256-KRX4fbuKEzAe6HIFDQ7N+S3R6f4kWnZsyQ8Zi23QAas=",
+        "version": "2.14.0"
+      }
+    }
   }
 }