Add Torch 2.11 support for CUDA/CPU and remove Torch 2.9

.github/workflows/build_kernel_macos.yaml

@@ -28,7 +28,7 @@ jobs:
       # For now we only test that there are no regressions in building macOS
       # kernels. Also run tests once we have a macOS runner.
       - name: Build relu kernel
-        run: ( cd examples/kernels/relu && nix build .\#redistributable.torch29-metal-aarch64-darwin -L )
+        run: ( cd examples/kernels/relu && nix build .\#redistributable.torch210-metal-aarch64-darwin -L )
 
       - name: Build relu metal cpp kernel
-        run: ( cd examples/kernels/relu-metal-cpp && nix build .\#redistributable.torch29-metal-aarch64-darwin -L )
+        run: ( cd examples/kernels/relu-metal-cpp && nix build .\#redistributable.torch210-metal-aarch64-darwin -L )

docs/source/builder/build-variants.md

@@ -8,42 +8,42 @@ available. This list will be updated as new PyTorch versions are released.
 ## CPU aarch64-darwin
 
 - `torch210-cpu-aarch64-darwin`
-- `torch29-cpu-aarch64-darwin`
+- `torch211-cpu-aarch64-darwin`
 
 ## Metal aarch64-darwin
 
 - `torch210-metal-aarch64-darwin`
-- `torch29-metal-aarch64-darwin`
+- `torch211-metal-aarch64-darwin`
 
 ## CPU aarch64-linux
 
 - `torch210-cxx11-cpu-aarch64-linux`
-- `torch29-cxx11-cpu-aarch64-linux`
+- `torch211-cxx11-cpu-aarch64-linux`
 
 ## CUDA aarch64-linux
 
 - `torch210-cxx11-cu126-aarch64-linux`
 - `torch210-cxx11-cu128-aarch64-linux`
 - `torch210-cxx11-cu130-aarch64-linux`
-- `torch29-cxx11-cu126-aarch64-linux`
-- `torch29-cxx11-cu128-aarch64-linux`
+- `torch211-cxx11-cu126-aarch64-linux`
+- `torch211-cxx11-cu128-aarch64-linux`
+- `torch211-cxx11-cu130-aarch64-linux`
 - `torch29-cxx11-cu129-aarch64-linux`
-- `torch29-cxx11-cu130-aarch64-linux`
 
 ## CPU x86_64-linux
 
 - `torch210-cxx11-cpu-x86_64-linux`
-- `torch29-cxx11-cpu-x86_64-linux`
+- `torch211-cxx11-cpu-x86_64-linux`
 
 ## CUDA x86_64-linux
 
 - `torch210-cxx11-cu126-x86_64-linux`
 - `torch210-cxx11-cu128-x86_64-linux`
 - `torch210-cxx11-cu130-x86_64-linux`
-- `torch29-cxx11-cu126-x86_64-linux`
-- `torch29-cxx11-cu128-x86_64-linux`
+- `torch211-cxx11-cu126-x86_64-linux`
+- `torch211-cxx11-cu128-x86_64-linux`
+- `torch211-cxx11-cu130-x86_64-linux`
 - `torch29-cxx11-cu129-x86_64-linux`
-- `torch29-cxx11-cu130-x86_64-linux`
 
 ## ROCm x86_64-linux
 

examples/kernels/flake.nix

@@ -15,7 +15,7 @@
       inherit (kernel-builder.inputs.nixpkgs) lib;
 
       cudaVersion = "cu126";
-      torchVersion = "29";
+      torchVersion = "210";
       tvmFfiVersion = "01";
 
       # All example kernels to build in CI.
@@ -79,7 +79,7 @@
           drv = sys: out: out.packages.${sys}.default;
           torchVersions = _defaultVersions: [
             {
-              torchVersion = "2.9";
+              torchVersion = "2.10";
               cudaVersion = "12.8";
               systems = [
                 "x86_64-linux"

examples/kernels/relu-specific-torch/flake.nix

@@ -15,7 +15,7 @@
       path = ./.;
       torchVersions = defaultVersions: [
         {
-          torchVersion = "2.9";
+          torchVersion = "2.10";
           cudaVersion = "12.8";
           systems = [
             "x86_64-linux"

examples/kernels/relu-torch-bounds/build.toml

@@ -6,8 +6,8 @@ backends = [
 ]
 
 [torch]
-minver = "2.9"
-maxver = "2.9"
+minver = "2.10"
+maxver = "2.10"
 src = [
     "torch-ext/torch_binding.cpp",
     "torch-ext/torch_binding.h",

nix-builder/build-variants.json

@@ -2,41 +2,41 @@
   "aarch64-darwin": {
     "cpu": [
       "torch210-cpu-aarch64-darwin",
-      "torch29-cpu-aarch64-darwin"
+      "torch211-cpu-aarch64-darwin"
     ],
     "metal": [
       "torch210-metal-aarch64-darwin",
-      "torch29-metal-aarch64-darwin"
+      "torch211-metal-aarch64-darwin"
     ]
   },
   "aarch64-linux": {
     "cpu": [
       "torch210-cxx11-cpu-aarch64-linux",
-      "torch29-cxx11-cpu-aarch64-linux"
+      "torch211-cxx11-cpu-aarch64-linux"
     ],
     "cuda": [
       "torch210-cxx11-cu126-aarch64-linux",
       "torch210-cxx11-cu128-aarch64-linux",
       "torch210-cxx11-cu130-aarch64-linux",
-      "torch29-cxx11-cu126-aarch64-linux",
-      "torch29-cxx11-cu128-aarch64-linux",
-      "torch29-cxx11-cu129-aarch64-linux",
-      "torch29-cxx11-cu130-aarch64-linux"
+      "torch211-cxx11-cu126-aarch64-linux",
+      "torch211-cxx11-cu128-aarch64-linux",
+      "torch211-cxx11-cu130-aarch64-linux",
+      "torch29-cxx11-cu129-aarch64-linux"
     ]
   },
   "x86_64-linux": {
     "cpu": [
       "torch210-cxx11-cpu-x86_64-linux",
-      "torch29-cxx11-cpu-x86_64-linux"
+      "torch211-cxx11-cpu-x86_64-linux"
     ],
     "cuda": [
       "torch210-cxx11-cu126-x86_64-linux",
       "torch210-cxx11-cu128-x86_64-linux",
       "torch210-cxx11-cu130-x86_64-linux",
-      "torch29-cxx11-cu126-x86_64-linux",
-      "torch29-cxx11-cu128-x86_64-linux",
-      "torch29-cxx11-cu129-x86_64-linux",
-      "torch29-cxx11-cu130-x86_64-linux"
+      "torch211-cxx11-cu126-x86_64-linux",
+      "torch211-cxx11-cu128-x86_64-linux",
+      "torch211-cxx11-cu130-x86_64-linux",
+      "torch29-cxx11-cu129-x86_64-linux"
     ],
     "rocm": [
       "torch210-cxx11-rocm70-x86_64-linux",

nix-builder/lib/checks.nix

@@ -12,9 +12,10 @@ let
     path = ../../examples/kernels/relu-torch-bounds;
   };
 in
-assert lib.assertMsg (builtins.all (buildSet: buildSet.torch.version == "2.9.0") kernelBuildSets) ''
-  Torch minver/maxver filtering does not work.
-'';
+assert lib.assertMsg (builtins.all (buildSet: buildSet.torch.version == "2.10.0") kernelBuildSets)
+  ''
+    Torch minver/maxver filtering does not work.
+  '';
 runCommand "builder-nix-checks" { } ''
   touch $out
 ''

nix-builder/overlay.nix

@@ -114,6 +114,11 @@ in
           xpuPackages = final.xpuPackages_2025_3;
         };
 
+        torch-bin_2_11 = mkTorch {
+          version = "2.11";
+          xpuPackages = final.xpuPackages_2025_3;
+        };
+
         triton-xpu_2_9 = callPackage ./pkgs/python-modules/triton-xpu {
           torchVersion = "2.9";
           xpuPackages = final.xpuPackages_2025_2;

nix-builder/pkgs/python-modules/torch/archs.nix

@@ -176,4 +176,46 @@
     ];
   };
 
+  "2.11" = {
+    # https://github.com/pytorch/pytorch/blob/release/2.11/.ci/manywheel/build_cuda.sh
+    capsPerCudaVersion = {
+      "13.0" = [
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+        "10.0"
+        "11.0"
+        "12.0"
+      ];
+      # NOTE: 12.9 does not seem to be in RC builds, check if needed for final release.
+      #       https://download.pytorch.org/whl/test/torch/
+      "12.9" = [
+        "7.0"
+        "7.5"
+        "8.0"
+        "9.0"
+        "10.0"
+        "12.0"
+      ];
+      "12.8" = [
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+        "10.0"
+        "12.0"
+      ];
+      "12.6" = [
+        "5.0"
+        "6.0"
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+      ];
+    };
+  };
 }

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

@@ -230,6 +230,7 @@ buildPythonPackage {
   # the Nix store paths of the framework into the Torch libraries..
   pythonRemoveWheelDeps =
     lib.optionals cudaSupport [
+      "cuda-toolkit"
       "nvidia-cuda-runtime"
       "nvidia-cuda-nvrtc"
       "nvidia-cuda-cupti"

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

@@ -1,26 +1,11 @@
 {
   "2.9": {
     "x86_64-linux": {
-      "cu126": {
-        "url": "https://download.pytorch.org/whl/cu126/torch-2.9.0%2Bcu126-cp313-cp313-manylinux_2_28_x86_64.whl",
-        "hash": "sha256-DEW2Nu1sAszuYk7FPYRMMrHbm+sj1hnzfEEVJGeQTNU=",
-        "version": "2.9.0"
-      },
-      "cu128": {
-        "url": "https://download.pytorch.org/whl/cu128/torch-2.9.0%2Bcu128-cp313-cp313-manylinux_2_28_x86_64.whl",
-        "hash": "sha256-l97wCH+O8XG5AC6lALr/3UQMe91VnCPDi7+HgbZ+k2Q=",
-        "version": "2.9.0"
-      },
       "cu129": {
         "url": "https://download.pytorch.org/whl/cu129/torch-2.9.0%2Bcu129-cp313-cp313-manylinux_2_28_x86_64.whl",
         "hash": "sha256-a6V77xFJM5fBUddVM0CSKQQSkE5GyF7YYnegO7JP0To=",
         "version": "2.9.0"
       },
-      "cu130": {
-        "url": "https://download.pytorch.org/whl/cu130/torch-2.9.0%2Bcu130-cp313-cp313-manylinux_2_28_x86_64.whl",
-        "hash": "sha256-kqktts3jjQUiDB996Rzqz/AgVGOGxbegomjcquF7XBg=",
-        "version": "2.9.0"
-      },
       "rocm63": {
         "url": "https://download.pytorch.org/whl/rocm6.3/torch-2.9.0%2Brocm6.3-cp313-cp313-manylinux_2_28_x86_64.whl",
         "hash": "sha256-gw3ZjX2JlS7MV8aO6nCyTMAsDO1CSa7z9UzHsW/Wbic=",
@@ -31,49 +16,17 @@
         "hash": "sha256-yKiHbnxb2P3MOUKJB20rgyla3a1PXKYOo4ZRI9Xcn3Y=",
         "version": "2.9.0"
       },
-      "cpu": {
-        "url": "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp313-cp313-manylinux_2_28_x86_64.whl",
-        "hash": "sha256-bJshdYRACWPVtNrds3Eex6N3jqshHhhlT7oHbM47hoI=",
-        "version": "2.9.0"
-      },
       "xpu": {
         "url": "https://download.pytorch.org/whl/xpu/torch-2.9.0%2Bxpu-cp313-cp313-linux_x86_64.whl",
         "hash": "sha256-mNagbdf7GFh0NnsYvWCfBfFv3OQUKlmAypRGGUmWXNI=",
         "version": "2.9.0"
       }
     },
     "aarch64-linux": {
-      "cu126": {
-        "url": "https://download.pytorch.org/whl/cu126/torch-2.9.0%2Bcu126-cp313-cp313-manylinux_2_28_aarch64.whl",
-        "hash": "sha256-qn+dy/BTDsEKIRRxyK+1Mo1bvKtvFP0beLXm8vFO1hc=",
-        "version": "2.9.0"
-      },
-      "cu128": {
-        "url": "https://download.pytorch.org/whl/cu128/torch-2.9.0%2Bcu128-cp313-cp313-manylinux_2_28_aarch64.whl",
-        "hash": "sha256-TXb3E0WvR/Aix/pV7dDBgQ0Br4ncue3P36/j0qD3prg=",
-        "version": "2.9.0"
-      },
       "cu129": {
         "url": "https://download.pytorch.org/whl/cu129/torch-2.9.0%2Bcu129-cp313-cp313-manylinux_2_28_aarch64.whl",
         "hash": "sha256-wyc11mL4ogcYOM6S72NnjLT4x2Yc3809GwZQO3sAFiY=",
         "version": "2.9.0"
-      },
-      "cu130": {
-        "url": "https://download.pytorch.org/whl/cu130/torch-2.9.0%2Bcu130-cp313-cp313-manylinux_2_28_aarch64.whl",
-        "hash": "sha256-7PPSS9TA5uQlvXeKbembUieeACGmDX6xGrDC1mnz+bA=",
-        "version": "2.9.0"
-      },
-      "cpu": {
-        "url": "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp313-cp313-manylinux_2_28_aarch64.whl",
-        "hash": "sha256-vkQ42NrX8NWl5U8P7viok0RolOyH8QK7HYLcxFGFQuQ=",
-        "version": "2.9.0"
-      }
-    },
-    "aarch64-darwin": {
-      "cpu": {
-        "url": "https://download.pytorch.org/whl/cpu/torch-2.9.0-cp313-none-macosx_11_0_arm64.whl",
-        "hash": "sha256-4kg22Wi1TvTfsFWUABphlYcRrJIkAmKR5OP5L4Om/X8=",
-        "version": "2.9.0"
       }
     }
   },
@@ -144,5 +97,58 @@
         "version": "2.10.0"
       }
     }
+  },
+  "2.11": {
+    "x86_64-linux": {
+      "cu126": {
+        "url": "https://download.pytorch.org/whl/cu126/torch-2.11.0%2Bcu126-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-X24qVwmHbdJaX+DU8Kse8zszu2BDH8v+doq2Ro3P6Io=",
+        "version": "2.11.0"
+      },
+      "cu128": {
+        "url": "https://download.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-25ZLM8VQNacqs+IWIoevjxzCdgOcZdAVdAzIjCbc7fc=",
+        "version": "2.11.0"
+      },
+      "cu130": {
+        "url": "https://download.pytorch.org/whl/cu130/torch-2.11.0%2Bcu130-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-VURht28hIRknx3YFa8sLAPtClyNkeUtobXaOuwtYY2Y=",
+        "version": "2.11.0"
+      },
+      "cpu": {
+        "url": "https://download.pytorch.org/whl/cpu/torch-2.11.0%2Bcpu-cp313-cp313-manylinux_2_28_x86_64.whl",
+        "hash": "sha256-RQJdd1LbxrTHhMA6+u6cXxlzDOCEsuQ/yaL+FnfZ/4Y=",
+        "version": "2.11.0"
+      }
+    },
+    "aarch64-linux": {
+      "cu126": {
+        "url": "https://download.pytorch.org/whl/cu126/torch-2.11.0%2Bcu126-cp313-cp313-manylinux_2_28_aarch64.whl",
+        "hash": "sha256-O7v8ftFth4rvm5Go54wDO/Ex/7g1fiP/gJESEGg6MbM=",
+        "version": "2.11.0"
+      },
+      "cu128": {
+        "url": "https://download.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp313-cp313-manylinux_2_28_aarch64.whl",
+        "hash": "sha256-fbNYAQa7oETaW4lQ8/uP5fMZmeqrP2o6oqxdICw2hNI=",
+        "version": "2.11.0"
+      },
+      "cu130": {
+        "url": "https://download.pytorch.org/whl/cu130/torch-2.11.0%2Bcu130-cp313-cp313-manylinux_2_28_aarch64.whl",
+        "hash": "sha256-w9YPeWZrkQHjkUouXewugerINOE8rgvPWelNwaRl91Y=",
+        "version": "2.11.0"
+      },
+      "cpu": {
+        "url": "https://download.pytorch.org/whl/cpu/torch-2.11.0%2Bcpu-cp313-cp313-manylinux_2_28_aarch64.whl",
+        "hash": "sha256-SLPiGjEURazdCyfxODDiHZOt73DUch4FHp8Fm665uPk=",
+        "version": "2.11.0"
+      }
+    },
+    "aarch64-darwin": {
+      "cpu": {
+        "url": "https://download.pytorch.org/whl/cpu/torch-2.11.0-cp313-cp313-macosx_11_0_arm64.whl",
+        "hash": "sha256-RC7J3HhZJWT9rWnPC+qp2i+Cq4EMy08TkDhpqQvz8V0=",
+        "version": "2.11.0"
+      }
+    }
   }
 }