Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .github/workflows/test_e2e.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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.13 grep once it is released.
VARIANT=$(nix run $GITHUB_WORKSPACE#kernel-builder -- list-variants . | grep -v torch213 | tail -1)
echo "name=$VARIANT" >> $GITHUB_OUTPUT
echo "Building variant: $VARIANT"

Expand Down
13 changes: 13 additions & 0 deletions docs/source/builder/build-variants.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,19 @@ available. This list will be updated as new PyTorch versions are released.

- `torch211-cpu-aarch64-darwin`
- `torch212-cpu-aarch64-darwin`
- `torch213-cpu-aarch64-darwin`

## Metal aarch64-darwin

- `torch211-metal-aarch64-darwin`
- `torch212-metal-aarch64-darwin`
- `torch213-metal-aarch64-darwin`

## CPU aarch64-linux

- `torch211-cxx11-cpu-aarch64-linux`
- `torch212-cxx11-cpu-aarch64-linux`
- `torch213-cxx11-cpu-aarch64-linux`

## CUDA aarch64-linux

Expand All @@ -28,11 +31,15 @@ available. This list will be updated as new PyTorch versions are released.
- `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`

## CPU x86_64-linux

- `torch211-cxx11-cpu-x86_64-linux`
- `torch212-cxx11-cpu-x86_64-linux`
- `torch213-cxx11-cpu-x86_64-linux`

## CUDA x86_64-linux

Expand All @@ -42,18 +49,24 @@ available. This list will be updated as new PyTorch versions are released.
- `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`

## ROCm x86_64-linux

- `torch211-cxx11-rocm71-x86_64-linux`
- `torch211-cxx11-rocm72-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`

## XPU x86_64-linux

- `torch211-cxx11-xpu20253-x86_64-linux`
- `torch212-cxx11-xpu20253-x86_64-linux`
- `torch213-cxx11-xpu20253-x86_64-linux`

## Python-only kernels

Expand Down
29 changes: 21 additions & 8 deletions nix-builder/build-variants.json
Original file line number Diff line number Diff line change
Expand Up @@ -2,49 +2,62 @@
"aarch64-darwin": {
"cpu": [
"torch211-cpu-aarch64-darwin",
"torch212-cpu-aarch64-darwin"
"torch212-cpu-aarch64-darwin",
"torch213-cpu-aarch64-darwin"
],
"metal": [
"torch211-metal-aarch64-darwin",
"torch212-metal-aarch64-darwin"
"torch212-metal-aarch64-darwin",
"torch213-metal-aarch64-darwin"
]
},
"aarch64-linux": {
"cpu": [
"torch211-cxx11-cpu-aarch64-linux",
"torch212-cxx11-cpu-aarch64-linux"
"torch212-cxx11-cpu-aarch64-linux",
"torch213-cxx11-cpu-aarch64-linux"
],
"cuda": [
"torch211-cxx11-cu126-aarch64-linux",
"torch211-cxx11-cu128-aarch64-linux",
"torch211-cxx11-cu130-aarch64-linux",
"torch212-cxx11-cu126-aarch64-linux",
"torch212-cxx11-cu130-aarch64-linux",
"torch212-cxx11-cu132-aarch64-linux"
"torch212-cxx11-cu132-aarch64-linux",
"torch213-cxx11-cu126-aarch64-linux",
"torch213-cxx11-cu130-aarch64-linux",
"torch213-cxx11-cu132-aarch64-linux"
]
},
"x86_64-linux": {
"cpu": [
"torch211-cxx11-cpu-x86_64-linux",
"torch212-cxx11-cpu-x86_64-linux"
"torch212-cxx11-cpu-x86_64-linux",
"torch213-cxx11-cpu-x86_64-linux"
],
"cuda": [
"torch211-cxx11-cu126-x86_64-linux",
"torch211-cxx11-cu128-x86_64-linux",
"torch211-cxx11-cu130-x86_64-linux",
"torch212-cxx11-cu126-x86_64-linux",
"torch212-cxx11-cu130-x86_64-linux",
"torch212-cxx11-cu132-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"
],
"rocm": [
"torch211-cxx11-rocm71-x86_64-linux",
"torch211-cxx11-rocm72-x86_64-linux",
"torch212-cxx11-rocm71-x86_64-linux",
"torch212-cxx11-rocm72-x86_64-linux"
"torch212-cxx11-rocm72-x86_64-linux",
"torch213-cxx11-rocm71-x86_64-linux",
"torch213-cxx11-rocm72-x86_64-linux"
],
"xpu": [
"torch211-cxx11-xpu20253-x86_64-linux",
"torch212-cxx11-xpu20253-x86_64-linux"
"torch212-cxx11-xpu20253-x86_64-linux",
"torch213-cxx11-xpu20253-x86_64-linux"
]
}
}
25 changes: 23 additions & 2 deletions nix-builder/overlay.nix
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,19 @@ final: prev:
triton-xpu = callPackage ./pkgs/python-modules/triton-xpu { };
in
{
inherit (triton) triton_3_5_0 triton_3_6_0 triton_3_7_0;
inherit (triton)
triton_3_5_0
triton_3_6_0
triton_3_7_0
triton_3_7_1
;
inherit (triton-rocm) triton-rocm_3_6_0 triton-rocm_3_7_0;
inherit (triton-xpu) triton-xpu_3_6_0 triton-xpu_3_7_0 triton-xpu_3_7_1;
inherit (triton-xpu)
triton-xpu_3_6_0
triton-xpu_3_7_0
triton-xpu_3_7_1
triton-xpu_3_7_2
;

cuda-bindings = python-self.callPackage ./pkgs/python-modules/cuda-bindings { };

Expand Down Expand Up @@ -145,6 +155,8 @@ final: prev:

mkTorch = callPackage ./pkgs/python-modules/torch/binary { };

pyzes = callPackage ./pkgs/python-modules/pyzes { };

scipy = python-super.scipy.overrideAttrs (
_: prevAttrs: {
# Three tests have a slight deviance.
Expand Down Expand Up @@ -179,6 +191,14 @@ final: prev:
xpuPackages = final.xpuPackages_2025_3_2;
};

torch-bin_2_13 = mkTorch {
version = "2.13";
triton-cuda = triton_3_7_1;
triton-rocm = triton-rocm_3_7_0;
triton-xpu = triton-xpu_3_7_2;
xpuPackages = final.xpuPackages_2026_0_0;
};

transformers = python-super.transformers.overridePythonAttrs (prevAttrs: rec {
version = "5.3.0";
src = python-super.fetchPypi {
Expand Down Expand Up @@ -247,6 +267,7 @@ final: prev:
xpuVersions = [
"2025.3.1"
"2025.3.2"
"2026.0.0"
];
newXpuPackages = final.callPackage ./pkgs/xpu-packages { };
in
Expand Down
37 changes: 37 additions & 0 deletions nix-builder/pkgs/aotriton/default.nix
Original file line number Diff line number Diff line change
Expand Up @@ -88,4 +88,41 @@ in
];
};

aotriton_0_12 = generic rec {
version = "0.12b";

hashes = {
"7.0" = "sha256-9quh/lkxIATM0T39pNCp41RXUn/sGDSNLnWilKEFHvU=";
"7.1" = "sha256-odcxdFkpthWY0IjuqtMdioKicDKqUeOnyDHkWpnglcI=";
"7.2" = "sha256-W5fo0EGxYMhAhZYfPTvXuYkGQrFGussEyZGqmtao3Kg=";
};

images = mkImages version [
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.12b/aotriton-0.12b-images-amd-gfx90a.tar.gz";
hash = "sha256-u4vyI3t3/FA7wpZ+oNmdbKQZEmxHnpUepCtxJzcSgIY=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.12b/aotriton-0.12b-images-amd-gfx942.tar.gz";
hash = "sha256-8I7az4PJzPHEvctR8cqwUtFoCr6jHJ4DXz+frbLxO6Q=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.12b/aotriton-0.12b-images-amd-gfx950.tar.gz";
hash = "sha256-MHo31ynNo6ISBEmQnlGSzXHCutzL038CInhgmOacepE=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.12b/aotriton-0.12b-images-amd-gfx110x.tar.gz";
hash = "sha256-ycrHz28ncWjhZZrC8EcG+II1gLfH4+iV9aVQPta91V8=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.12b/aotriton-0.12b-images-amd-gfx115x.tar.gz";
hash = "sha256-MXc4ehXGeLMAV/RYTR/BuPjbVhY4kMtcmPJ0UCCfWns=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.12b/aotriton-0.12b-images-amd-gfx120x.tar.gz";
hash = "sha256-aFclEc5kh6g/kBS9JVvWnIlD+H0Mk71XstqsX7xsecE=";
})
];
};

}
39 changes: 39 additions & 0 deletions nix-builder/pkgs/python-modules/pyzes/default.nix
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
{

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where is this module needed?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's a new dependency of Torch on XPU.

lib,

buildPythonPackage,
fetchPypi,
setuptools,

python,

level-zero,
}:

buildPythonPackage rec {
pname = "pyzes";
version = "0.1.2";
pyproject = true;

src = fetchPypi {
inherit pname version;
hash = "sha256-IWZ/CiOxX8+BwSZ9vIb1qbDVEYKXyuZF5KsxxQ6efWs=";
};

build-system = [ setuptools ];

postInstall = ''
substituteInPlace $out/${python.sitePackages}/pyzes.py \
--replace-fail 'libName = "/usr/lib/x86_64-linux-gnu/lib" + libName + ".so.1"' \
'libName = "${level-zero}/lib/libze_loader.so.1"'
'';

pythonImportsCheck = [ "pyzes" ];

meta = with lib; {
description = "Python bindings to the Intel Level-Zero-Driver Library";
homepage = "https://github.com/oneapi-src/level-zero";
license = lib.licenses.mit;
sourceProvenance = with lib.sourceTypes; [ fromSource ];
};
}
24 changes: 20 additions & 4 deletions nix-builder/pkgs/python-modules/torch/binary/generic.nix
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

# Build inputs
cudaPackages,
numactl,
rocmPackages,
xpuPackages,
zlib,
Expand All @@ -35,6 +36,7 @@
networkx,
numpy,
pyyaml,
pyzes,
requests,
setuptools,
sympy,
Expand All @@ -51,6 +53,8 @@
effectiveStdenv ? if cudaSupport then cudaPackages.backendStdenv else stdenv,
}:
let
torchMajorMinor = lib.versions.majorMinor version;

effectiveTriton =
if cudaSupport then
triton-cuda
Expand All @@ -64,12 +68,10 @@ let
aotritonVersions = with rocmPackages; {
"2.11" = aotriton_0_11_2;
"2.12" = aotriton_0_11_2;
"2.13" = aotriton_0_12;
};

aotriton =
let
torchMajorMinor = lib.versions.majorMinor version;
in
aotritonVersions.${torchMajorMinor}
or (throw "aotriton version is not specified Torch ${torchMajorMinor}");

Expand Down Expand Up @@ -105,7 +107,7 @@ let
rocsparse
roctracer
]
++ lib.optionals (lib.versions.majorMinor version == "2.12") [
++ lib.optionals (lib.versionAtLeast version "2.12") [
rocprofiler-sdk
];

Expand Down Expand Up @@ -229,10 +231,14 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
]
++ lib.optionals (cudaSupport && lib.versionAtLeast version "2.10") [
cuda-bindings
]
++ lib.optionals (xpuSupport && lib.versionAtLeast version "2.13") [
pyzes
];

pythonRelaxWheelDeps = [
"cuda-bindings"
"pyzes"
"sympy"
];

Expand Down Expand Up @@ -350,6 +356,16 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
printWords "''${propagatedCxxBuildInputs[@]}" >> "$cxxdev/nix-support/propagated-build-inputs"
'';

# 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") ''
# 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.
patchelf --add-rpath ${numactl}/lib $out/${python.sitePackages}/torch/lib/libtorch_rocshmem.so
'';

dontStrip = true;

pythonImportsCheck = [ "torch" ];
Expand Down
Loading
Loading