Add a CUDA 12.9 build variant for Torch 2.7 (#136)

This change adds a CUDA 12.9 build variant. Since this is not a variant
supported/distributed by upstream, we add an `upstreamVariant` attribute
to the Torch variants. A variant is only a candidate for a kernel's
`bundle` package when `upstreamVariant = true`.

Todo: also use `upstreamVariant` in generating variant JSON.

Author: danieldk

.github/workflows/check_variants.yaml

@@ -17,22 +17,22 @@ jobs:
         with:
           nix_path: nixpkgs=channel:nixos-unstable
       - name: Generate variants JSON
-        run: nix eval --raw .#lib.allBuildVariantsJSON | nix run nixpkgs#jq > build-variants.json
+        run: nix eval --raw .#lib.allBuildVariantsJSON | nix run nixpkgs#jq 'walk(if type == "array" then sort else . end)' > build-variants.json
       - name: Check if variants JSON is up-to-date
         run: |
           if git diff --exit-code build-variants.json; then
             echo "✅ variants.json is up-to-date"
           else
-            echo "🛑 regenerate variants.json: nix eval --raw .#lib.allBuildVariantsJSON | nix run nixpkgs#jq > build-variants.json"
+            echo "🛑 regenerate variants.json: nix eval --raw .#lib.allBuildVariantsJSON | nix run nixpkgs#jq 'walk(if type == "array" then sort else . end)' > build-variants.json"
             exit 1
           fi
       - name: Generate variants Markdown
-        run: scripts/gen_variants_markdown.py
+        run: nix run nixpkgs#python3 scripts/gen_variants_markdown.py
       - name: Check if variants Markdown is up-to-date
         run: |
           if git diff --exit-code docs/build-variants.md; then
             echo "✅ docs/build-variants.md is up-to-date"
           else
-            echo "🛑 regenerate docs/build-variants: scripts/gen_variants_markdown.py"
+            echo "🛑 regenerate docs/build-variants: nix run nixpkgs#python3 scripts/gen_variants_markdown.py"
             exit 1
           fi

build-variants.json

@@ -10,10 +10,10 @@
   "x86_64-linux": {
     "cuda": [
       "torch26-cxx11-cu118-x86_64-linux",
-      "torch26-cxx98-cu118-x86_64-linux",
       "torch26-cxx11-cu124-x86_64-linux",
-      "torch26-cxx98-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",

docs/build-variants.md

@@ -15,10 +15,10 @@ available. This list will be updated as new PyTorch versions are released.
 ## CUDA x86_64-linux
 
 - `torch26-cxx11-cu118-x86_64-linux`
-- `torch26-cxx98-cu118-x86_64-linux`
 - `torch26-cxx11-cu124-x86_64-linux`
-- `torch26-cxx98-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`

flake.lock

@@ -3,7 +3,8 @@
 
   inputs = {
     flake-utils.url = "github:numtide/flake-utils";
-    nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable-small";
+    #nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable-small";
+    nixpkgs.url = "github:danieldk/nixpkgs/kernel-builder-cuda-12.9.0";
     flake-compat.url = "github:edolstra/flake-compat";
     rocm-nix = {
       url = "github:huggingface/rocm-nix";
@@ -85,6 +86,7 @@
                 };
                 redistributable = build.buildDistTorchExtensions {
                   inherit path;
+                  buildSets = buildSetPerSystem.${system};
                   rev = revUnderscored;
                 };
                 buildTree =

flake.nix

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

lib/build-version.nix

@@ -43,6 +43,7 @@ rec {
     in
     lib.foldl (langs: kernel: langs // { ${kernelLang kernel} = true; }) init kernels;
 
+  # Filter buildsets that are applicable to a given kernel build config.
   applicableBuildSets =
     buildConfig: buildSets:
     let
@@ -61,6 +62,7 @@ rec {
       gpu,
       pkgs,
       torch,
+      upstreamVariant,
     }:
     {
       path,
@@ -126,7 +128,11 @@ rec {
 
   # Build multiple Torch extensions.
   buildDistTorchExtensions =
-    { path, rev }:
+    {
+      buildSets,
+      path,
+      rev,
+    }:
     let
       extensionForTorch =
         { path, rev }:
@@ -147,7 +153,11 @@ rec {
     let
       # We just need to get any nixpkgs for use by the path join.
       pkgs = (builtins.head buildSets).pkgs;
-      extensions = buildDistTorchExtensions { inherit path rev; };
+      upstreamBuildSets = builtins.filter (buildSet: buildSet.upstreamVariant) buildSets;
+      extensions = buildDistTorchExtensions {
+        inherit path rev;
+        buildSets = upstreamBuildSets;
+      };
       buildConfig = readBuildConfig path;
       namePaths =
         if buildConfig.torch.universal or false then

lib/build.nix

@@ -32,6 +32,7 @@ let
       rocmVersion ? "",
       torchVersion,
       cxx11Abi,
+      upstreamVariant ? false,
     }:
     let
       pkgs = if gpu == "cuda" then pkgsByCudaVer.${cudaVersion} else pkgsByRocmVer.${rocmVersion};
@@ -40,7 +41,12 @@ let
       };
     in
     {
-      inherit gpu pkgs torch;
+      inherit
+        gpu
+        pkgs
+        torch
+        upstreamVariant
+        ;
     };
 
   pkgsForRocm = import nixpkgs {

lib/buildsets.nix

@@ -126,7 +126,7 @@ let
   supportedTorchCudaCapabilities =
     let
       # https://github.com/pytorch/pytorch/blob/release/2.7/.ci/manywheel/build_cuda.sh
-      capsPerCudaVersion = {
+      capsPerCudaVersion = rec {
         "12.8" = [
           "7.5"
           "8.0"
@@ -163,6 +163,16 @@ let
           "8.6"
           "9.0"
         ];
+
+        # Not in upstream yet, so use same capabilities as 12.8.
+        "12.9" = [
+          "7.5"
+          "8.0"
+          "8.6"
+          "9.0"
+          "10.0"
+          "12.0"
+        ];
       };
       real = capsPerCudaVersion."${lib.versions.majorMinor cudaPackages.cudaVersion}";
       ptx = lists.map (x: "${x}+PTX") real;