Add Torch 2.8, remove Torch 2.6

.github/workflows/build_kernel.yaml

@@ -24,17 +24,17 @@ jobs:
         env:
           USER: github_runner
       - name: Build activation kernel
-        run: ( cd examples/activation && nix build .\#redistributable.torch26-cxx98-cu124-x86_64-linux )
+        run: ( cd examples/activation && nix build .\#redistributable.torch27-cxx11-cu126-x86_64-linux )
       - name: Copy activation kernel
         run: cp -rL examples/activation/result activation-kernel
 
       - name: Build cutlass GEMM kernel
-        run: ( cd examples/cutlass-gemm && nix build .\#redistributable.torch26-cxx98-cu124-x86_64-linux )
+        run: ( cd examples/cutlass-gemm && nix build .\#redistributable.torch27-cxx11-cu126-x86_64-linux )
       - name: Copy cutlass GEMM kernel
         run: cp -rL examples/cutlass-gemm/result cutlass-gemm-kernel
 
       - name: Build relu kernel
-        run: ( cd examples/relu && nix build .\#redistributable.torch26-cxx98-cu124-x86_64-linux )
+        run: ( cd examples/relu && nix build .\#redistributable.torch27-cxx11-cu126-x86_64-linux )
       - name: Copy relu kernel
         run: cp -rL examples/relu/result relu-kernel
 
@@ -43,7 +43,7 @@ jobs:
         run: ( cd examples/relu-specific-torch && nix build . )
 
       - name: Build silu-and-mul-universal kernel
-        run: ( cd examples/silu-and-mul-universal && nix build .\#redistributable.torch26-cxx98-cu124-x86_64-linux )
+        run: ( cd examples/silu-and-mul-universal && nix build .\#redistributable.torch27-cxx11-cu126-x86_64-linux )
       - name: Copy silu-and-mul-universal kernel
         run: cp -rL examples/silu-and-mul-universal/result silu-and-mul-universal-kernel
 

.github/workflows/build_kernel_rocm.yaml

@@ -26,4 +26,4 @@ jobs:
       # For now we only test that there are no regressions in building ROCm
       # kernels. Also run tests once we have a ROCm runner.
       - name: Build relu kernel
-        run: ( cd examples/relu && nix build .\#redistributable.torch26-cxx11-rocm62-x86_64-linux -L )
+        run: ( cd examples/relu && nix build .\#redistributable.torch27-cxx11-rocm63-x86_64-linux -L )

build-variants.json

@@ -1,32 +1,29 @@
 {
   "aarch64-darwin": {
     "metal": [
-      "torch27-metal-aarch64-darwin"
+      "torch27-metal-aarch64-darwin",
+      "torch28-metal-aarch64-darwin"
     ]
   },
   "aarch64-linux": {
     "cuda": [
-      "torch26-cxx11-cu126-aarch64-linux",
-      "torch26-cxx98-cu126-aarch64-linux",
-      "torch27-cxx11-cu126-aarch64-linux",
-      "torch27-cxx11-cu128-aarch64-linux"
+      "torch27-cxx11-cu128-aarch64-linux",
+      "torch28-cxx11-cu129-aarch64-linux"
     ]
   },
   "x86_64-linux": {
     "cuda": [
-      "torch26-cxx11-cu118-x86_64-linux",
-      "torch26-cxx11-cu124-x86_64-linux",
-      "torch26-cxx11-cu126-x86_64-linux",
-      "torch26-cxx98-cu118-x86_64-linux",
-      "torch26-cxx98-cu124-x86_64-linux",
-      "torch26-cxx98-cu126-x86_64-linux",
       "torch27-cxx11-cu118-x86_64-linux",
       "torch27-cxx11-cu126-x86_64-linux",
-      "torch27-cxx11-cu128-x86_64-linux"
+      "torch27-cxx11-cu128-x86_64-linux",
+      "torch28-cxx11-cu126-x86_64-linux",
+      "torch28-cxx11-cu128-x86_64-linux",
+      "torch28-cxx11-cu129-x86_64-linux"
     ],
     "rocm": [
-      "torch26-cxx11-rocm62-x86_64-linux",
-      "torch27-cxx11-rocm63-x86_64-linux"
+      "torch27-cxx11-rocm63-x86_64-linux",
+      "torch28-cxx11-rocm63-x86_64-linux",
+      "torch28-cxx11-rocm64-x86_64-linux"
     ]
   }
 }

docs/build-variants.md

@@ -8,30 +8,27 @@ available. This list will be updated as new PyTorch versions are released.
 ## Metal aarch64-darwin
 
 - `torch27-metal-aarch64-darwin`
+- `torch28-metal-aarch64-darwin`
 
 ## CUDA aarch64-linux
 
-- `torch26-cxx11-cu126-aarch64-linux`
-- `torch26-cxx98-cu126-aarch64-linux`
-- `torch27-cxx11-cu126-aarch64-linux`
 - `torch27-cxx11-cu128-aarch64-linux`
+- `torch28-cxx11-cu129-aarch64-linux`
 
 ## CUDA x86_64-linux
 
-- `torch26-cxx11-cu118-x86_64-linux`
-- `torch26-cxx11-cu124-x86_64-linux`
-- `torch26-cxx11-cu126-x86_64-linux`
-- `torch26-cxx98-cu118-x86_64-linux`
-- `torch26-cxx98-cu124-x86_64-linux`
-- `torch26-cxx98-cu126-x86_64-linux`
 - `torch27-cxx11-cu118-x86_64-linux`
 - `torch27-cxx11-cu126-x86_64-linux`
 - `torch27-cxx11-cu128-x86_64-linux`
+- `torch28-cxx11-cu126-x86_64-linux`
+- `torch28-cxx11-cu128-x86_64-linux`
+- `torch28-cxx11-cu129-x86_64-linux`
 
 ## ROCm x86_64-linux
 
-- `torch26-cxx11-rocm62-x86_64-linux`
 - `torch27-cxx11-rocm63-x86_64-linux`
+- `torch28-cxx11-rocm63-x86_64-linux`
+- `torch28-cxx11-rocm64-x86_64-linux`
 
 ## Universal
 

examples/relu-specific-torch/flake.nix

@@ -22,7 +22,7 @@
             "x86_64-linux"
             "aarch64-linux"
           ];
-          upstreamVariant = true;
+          bundleBuild = true;
         }
       ];
     };

flake.nix

@@ -97,7 +97,7 @@
               revUnderscored = builtins.replaceStrings [ "-" ] [ "_" ] rev;
               pkgs = nixpkgs.legacyPackages.${system};
               shellTorch =
-                if system == "aarch64-darwin" then "torch27-metal-${system}" else "torch27-cxx11-cu126-${system}";
+                if system == "aarch64-darwin" then "torch28-metal-${system}" else "torch28-cxx11-cu126-${system}";
             in
             {
               devShells = rec {

lib/build-sets.nix

@@ -31,13 +31,13 @@ let
     let
       withCuda = builtins.filter (torchVersion: torchVersion ? cudaVersion) torchVersions;
     in
-    builtins.map (torchVersion: torchVersion.cudaVersion) withCuda;
+    lib.unique (builtins.map (torchVersion: torchVersion.cudaVersion) withCuda);
 
   rocmVersions =
     let
       withRocm = builtins.filter (torchVersion: torchVersion ? rocmVersion) torchVersions;
     in
-    builtins.map (torchVersion: torchVersion.rocmVersion) withRocm;
+    lib.unique (builtins.map (torchVersion: torchVersion.rocmVersion) withRocm);
 
   flattenVersion = version: lib.replaceStrings [ "." ] [ "_" ] (lib.versions.pad 2 version);
 
@@ -59,7 +59,7 @@ let
       torchVersion,
       cxx11Abi,
       system,
-      upstreamVariant ? false,
+      bundleBuild ? false,
     }:
     let
       pkgs =
@@ -80,7 +80,7 @@ let
         buildConfig
         pkgs
         torch
-        upstreamVariant
+        bundleBuild
         ;
     };
 
@@ -92,18 +92,6 @@ let
     ];
   };
 
-  pkgsForRocm = import nixpkgs {
-    inherit system;
-    config = {
-      allowUnfree = true;
-      rocmSupport = true;
-    };
-    overlays = [
-      hf-nix
-      overlay
-    ];
-  };
-
   # Instantiate nixpkgs for the given CUDA versions. Returns
   # an attribute set like `{ "12.4" = <nixpkgs with 12.4>; ... }`.
   pkgsForCudaVersions =

lib/build-variants.nix

@@ -19,7 +19,7 @@ rec {
     else
       throw "Could not find compute framework: no CUDA or ROCm version specified and Metal is not enabled";
 
-  # Upstream build variants.
+  # Build variants included in bundle builds.
   buildVariants =
     let
       inherit (import ./version-utils.nix { inherit lib; }) abiString flattenVersion;
@@ -39,7 +39,7 @@ rec {
           "torch${flattenVersion version.torchVersion}-${computeString version}-${version.system}"
         else
           "torch${flattenVersion version.torchVersion}-${abiString version.cxx11Abi}-${computeString version}-${version.system}";
-      upstreamVersions = lib.filter (version: version.upstreamVariant or false);
+      bundleBuildVersions = lib.filter (version: version.bundleBuild or false);
     in
     lib.foldl' (
       acc: version:
@@ -51,5 +51,5 @@ rec {
         pathVersions = lib.attrByPath path [ ] acc ++ [ (buildName version) ];
       in
       lib.recursiveUpdate acc (lib.setAttrByPath path pathVersions)
-    ) { } (flattenSystems (upstreamVersions torchVersions));
+    ) { } (flattenSystems (bundleBuildVersions torchVersions));
 }

lib/build-version.nix

@@ -2,7 +2,7 @@
   buildConfig,
   pkgs,
   torch,
-  upstreamVariant,
+  bundleBuild,
 }:
 let
   inherit (pkgs) lib;

lib/build.nix

@@ -90,7 +90,7 @@ rec {
       buildConfig,
       pkgs,
       torch,
-      upstreamVariant,
+      bundleBuild,
     }:
     {
       path,
@@ -199,10 +199,10 @@ rec {
     let
       # We just need to get any nixpkgs for use by the path join.
       pkgs = (builtins.head buildSets).pkgs;
-      upstreamBuildSets = builtins.filter (buildSet: buildSet.upstreamVariant) buildSets;
+      bundleBuildSets = builtins.filter (buildSet: buildSet.bundleBuild) buildSets;
       extensions = buildDistTorchExtensions {
         inherit path rev doGetKernelCheck;
-        buildSets = upstreamBuildSets;
+        buildSets = bundleBuildSets;
       };
       buildConfig = readBuildConfig path;
       namePaths =

lib/torch-extension/default.nix

@@ -121,7 +121,7 @@ stdenv.mkDerivation (prevAttrs: {
         libcusparse
       ]
     )
-    #++ lib.optionals rocmSupport (with rocmPackages; [ clr rocm-core ])
+    ++ lib.optionals rocmSupport (with rocmPackages; [ hipsparselt ])
     ++ lib.optionals stdenv.hostPlatform.isDarwin [
       apple-sdk_15
     ]

pkgs/stdenv-glibc-2_27/default.nix

@@ -6,6 +6,7 @@
   system,
   wrapBintoolsWith,
   wrapCCWith,
+  writeClosure,
   gcc12Stdenv,
   stdenv,
   bintools-unwrapped,
@@ -71,5 +72,11 @@ let
     in
     overrideCC stdenv compilerWrapped;
 
+  # Workaround for: https://github.com/NixOS/nixpkgs/issues/428546
+  ccWithoutHook = ((if cudaSupport then cudaPackages.backendStdenv else gcc12Stdenv).cc.cc).override {
+    # Just some null derivation to disable the hook.
+    sanitiseHeaderPathsHook = writeClosure [ ];
+  };
+
 in
-stdenvWith glibc_2_27 (if cudaSupport then cudaPackages.backendStdenv else gcc12Stdenv).cc.cc stdenv
+stdenvWith glibc_2_27 ccWithoutHook stdenv