Skip to content

Add ROCm support for Torch 2.11 #401

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
danieldk merged 6 commits into from
Mar 26, 2026
Merged

Add ROCm support for Torch 2.11 #401

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e973349
rocmPackages_7_2: init at 7.2.1
danieldk Mar 26, 2026
f6d6233
aotriton_0_11_2: init at 0.11.2
danieldk Mar 26, 2026
847bdf0
python3Packages.torch_2_11: add ROCm builds
danieldk Mar 26, 2026
9c3c029
Add Torch 2.11 ROCm variants, remove 2.9
danieldk Mar 26, 2026
2204378
Remove unused aotriton/ROCm versions
danieldk Mar 26, 2026
a95eb96
CI: do not build Torch 2.9 ROCm kernels
danieldk Mar 26, 2026
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
4 changes: 2 additions & 2 deletions .github/workflows/build_kernel_rocm.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ 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/kernels/relu && nix build .\#redistributable.torch29-cxx11-rocm63-x86_64-linux -L )
run: ( cd examples/kernels/relu && nix build .\#redistributable.torch210-cxx11-rocm70-x86_64-linux -L )

- name: Build relu kernel (compiler flags)
run: ( cd examples/kernels/relu-compiler-flags && nix build .\#redistributable.torch29-cxx11-rocm63-x86_64-linux )
run: ( cd examples/kernels/relu-compiler-flags && nix build .\#redistributable.torch210-cxx11-rocm70-x86_64-linux )
4 changes: 2 additions & 2 deletions docs/source/builder/build-variants.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,8 @@ available. This list will be updated as new PyTorch versions are released.

- `torch210-cxx11-rocm70-x86_64-linux`
- `torch210-cxx11-rocm71-x86_64-linux`
- `torch29-cxx11-rocm63-x86_64-linux`
- `torch29-cxx11-rocm64-x86_64-linux`
- `torch211-cxx11-rocm71-x86_64-linux`
- `torch211-cxx11-rocm72-x86_64-linux`

## XPU x86_64-linux

Expand Down
4 changes: 2 additions & 2 deletions nix-builder/build-variants.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,8 @@
"rocm": [
"torch210-cxx11-rocm70-x86_64-linux",
"torch210-cxx11-rocm71-x86_64-linux",
"torch29-cxx11-rocm63-x86_64-linux",
"torch29-cxx11-rocm64-x86_64-linux"
"torch211-cxx11-rocm71-x86_64-linux",
"torch211-cxx11-rocm72-x86_64-linux"
],
"xpu": [
"torch210-cxx11-xpu20253-x86_64-linux",
Expand Down
3 changes: 1 addition & 2 deletions nix-builder/overlay.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,10 +191,9 @@ in
flattenVersion = prev.lib.strings.replaceStrings [ "." ] [ "_" ];
readPackageMetadata = path: (builtins.fromJSON (builtins.readFile path));
versions = [
"6.3.4"
"6.4.2"
"7.0.2"
"7.1.1"
"7.2.1"
];
newRocmPackages = final.callPackage ./pkgs/rocm-packages { };
in
Expand Down
93 changes: 23 additions & 70 deletions nix-builder/pkgs/aotriton/default.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,68 +32,21 @@ let
};
in
{
aotriton_0_10 = generic rec {
Copy link
Copy Markdown
Member

@sayakpaul sayakpaul Mar 26, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

Do we not need it anymore?

Copy link
Copy Markdown
Member Author

@danieldk danieldk Mar 26, 2026

Choose a reason for hiding this comment

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

No, Torch 2.10 and 2.11 use aotriton 0.11.1 and 0.11.2 respectively.

Copy link
Copy Markdown
Member

@sayakpaul sayakpaul Mar 26, 2026

Choose a reason for hiding this comment

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

Ah yeah. I see the modifications now.

version = "0.10b";

src = fetchFromGitHub {
owner = "ROCm";
repo = "aotriton";
rev = version;
hash = "sha256-stAHnsqChkNv69wjlhM/qUetrJpNwI1i7rGnPMwsNz0=";
leaveDotGit = true;
inherit postFetch;
};

patches = [
# A bunch of implicit type narrowing issues that are rejected by newer
# compilers.
./v0.10b-explicit-cast-for-narrowing.diff
# Fails with: ld.lld: error: unable to insert .comment after .comment
./v0.10b-no-ld-script.diff

# CMakeLists.txt: AOTRITON_INHERIT_SYSTEM_SITE_TRITON flag
(fetchpatch {
url = "https://github.com/ROCm/aotriton/commit/9734c3e999c412a07d2b35671998650942b26ed4.patch";
hash = "sha256-tBmjjhRJmLv3K6F2+4OcMuwf8dH7efPPECMQjh6QdUA=";
})
];

gpuTargets = [
# aotriton GPU support list:
# https://github.com/ROCm/aotriton/blob/main/v2python/gpu_targets.py
"gfx90a"
"gfx942"
"gfx950"
"gfx1100"
"gfx1101"
"gfx1201"
];

images = mkImages version [
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.10b/aotriton-0.10b-manylinux_2_28_x86_64-rocm6.3-shared.tar.gz";
hash = "sha256-hhzZ90ee7JQ5M8J8uGkgJH5bXdE5vHwTdsgYCKu31/4=";
})
];

extraPythonDepends = ps: [ ps.pandas ];
};

aotriton_0_11 = generic rec {
version = "0.11b";
aotriton_0_11_1 = generic rec {
version = "0.11.1b";

src = fetchFromGitHub {
owner = "ROCm";
repo = "aotriton";
rev = version;
hash = "sha256-QXkNB3vNmPg4/m23OMuBBX4cjZQ3zQPotaeimFMbclc=";
hash = "sha256-F7JjyS+6gMdCpOFLldTsNJdVzzVwd6lwW7+V8ZOZfig=";
leaveDotGit = true;
inherit postFetch;
};

patches = [
# Fails with: ld.lld: error: unable to insert .comment after .comment
./v0.11b-no-ld-script.diff
./v0.11.1b-no-ld-script.diff
];

gpuTargets = [
Expand All @@ -109,38 +62,38 @@ in

images = mkImages version [
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11b/aotriton-0.11b-images-amd-gfx90a.tar.gz";
hash = "sha256-wZpByUgFEKsy5vsF5u0KODLWsHY08FC4NrdgIAvvpzU=";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx90a.tar.gz";
hash = "sha256-/p8Etmv1KsJ80CXh2Jz9BJdN0/s64HYZL3g2QaTYD98=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11b/aotriton-0.11b-images-amd-gfx942.tar.gz";
hash = "sha256-OgapmXHd23cDowN48cXWtBRo2SbqUYIRVtG2hXuYW8Q=";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx942.tar.gz";
hash = "sha256-CnvO4Z07ttVIcyJIwyNPe5JzbCq3p6rmUpS4en/WTAY=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11b/aotriton-0.11b-images-amd-gfx950.tar.gz";
hash = "sha256-J/wh9nYdV5h6cAQ23ozynL3Z7u6RMY3+1ZbusUfSGa0=";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx950.tar.gz";
hash = "sha256-wbo7/oQhf9Z9890fi2fICn97M9CtTXS0HWVnA24DKs4=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11b/aotriton-0.11b-images-amd-gfx11xx.tar.gz";
hash = "sha256-7BNAMghzRBdmlVBdtlk4c3TRkWrf7hbw20fe442chgM=";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx11xx.tar.gz";
hash = "sha256-ZjIEDEBdgzvm/3ICkknHdoOLr18Do8E7pOjTeoe3p0A=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11b/aotriton-0.11b-images-amd-gfx120x.tar.gz";
hash = "sha256-/sBSBXR/9RZJseFRVFJn1aogN7qdAzjK0oaIKRW5QbA=";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx120x.tar.gz";
hash = "sha256-Ck/zJL/9rAwv3oeop/cFY9PISoCtTo8xNF8rQKE4TpU=";
})
];

extraPythonDepends = ps: [ ps.pandas ];
};

aotriton_0_11_1 = generic rec {
version = "0.11.1b";
aotriton_0_11_2 = generic rec {
version = "0.11.2b";

src = fetchFromGitHub {
owner = "ROCm";
repo = "aotriton";
rev = version;
hash = "sha256-F7JjyS+6gMdCpOFLldTsNJdVzzVwd6lwW7+V8ZOZfig=";
hash = "sha256-VIwwQR1fl40NLNOwO8KhQK/xOK6wb2l8qBugJ1cRjm4=";
leaveDotGit = true;
inherit postFetch;
};
Expand All @@ -163,23 +116,23 @@ in

images = mkImages version [
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx90a.tar.gz";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.2b/aotriton-0.11.2b-images-amd-gfx90a.tar.gz";
hash = "sha256-/p8Etmv1KsJ80CXh2Jz9BJdN0/s64HYZL3g2QaTYD98=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx942.tar.gz";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.2b/aotriton-0.11.2b-images-amd-gfx942.tar.gz";
hash = "sha256-CnvO4Z07ttVIcyJIwyNPe5JzbCq3p6rmUpS4en/WTAY=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx950.tar.gz";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.2b/aotriton-0.11.2b-images-amd-gfx950.tar.gz";
hash = "sha256-wbo7/oQhf9Z9890fi2fICn97M9CtTXS0HWVnA24DKs4=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx11xx.tar.gz";
hash = "sha256-ZjIEDEBdgzvm/3ICkknHdoOLr18Do8E7pOjTeoe3p0A=";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.2b/aotriton-0.11.2b-images-amd-gfx11xx.tar.gz";
hash = "sha256-g5KZY3/MsT++Pngj1X0bLc0OC+14q7y3AF6l9P2CuSg=";
})
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx120x.tar.gz";
url = "https://github.com/ROCm/aotriton/releases/download/0.11.2b/aotriton-0.11.2b-images-amd-gfx120x.tar.gz";
hash = "sha256-Ck/zJL/9rAwv3oeop/cFY9PISoCtTo8xNF8rQKE4TpU=";
})
];
Expand Down
3 changes: 1 addition & 2 deletions nix-builder/pkgs/python-modules/torch/binary/generic.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,9 +69,8 @@ let
inherit (archs) supportedTorchRocmArchs;

aotritonVersions = with rocmPackages; {
"2.8" = aotriton_0_10;
"2.9" = aotriton_0_11;
"2.10" = aotriton_0_11_1;
"2.11" = aotriton_0_11_2;
};

aotriton =
Expand Down
20 changes: 10 additions & 10 deletions nix-builder/pkgs/python-modules/torch/binary/torch-versions-hash.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,16 +6,6 @@
"hash": "sha256-a6V77xFJM5fBUddVM0CSKQQSkE5GyF7YYnegO7JP0To=",
"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=",
"version": "2.9.0"
},
"rocm64": {
"url": "https://download.pytorch.org/whl/rocm6.4/torch-2.9.0%2Brocm6.4-cp313-cp313-manylinux_2_28_x86_64.whl",
"hash": "sha256-yKiHbnxb2P3MOUKJB20rgyla3a1PXKYOo4ZRI9Xcn3Y=",
"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=",
Expand Down Expand Up @@ -119,6 +109,16 @@
"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"
},
"rocm71": {
"url": "https://download.pytorch.org/whl/rocm7.1/torch-2.11.0%2Brocm7.1-cp313-cp313-manylinux_2_28_x86_64.whl",
"hash": "sha256-Oy30B0GP3Vh1YbZUiy7VOy2UU1REkYYPRu8MKOxrfo4=",
"version": "2.11.0"
},
"rocm72": {
"url": "https://download.pytorch.org/whl/rocm7.2/torch-2.11.0%2Brocm7.2-cp313-cp313-manylinux_2_28_x86_64.whl",
"hash": "sha256-CupGzRf2SN6z/1kOFipiQ6SCV2W9AHJ0OWDvqdSlvB0=",
"version": "2.11.0"
}
},
"aarch64-linux": {
Expand Down
20 changes: 10 additions & 10 deletions nix-builder/pkgs/python-modules/torch/binary/torch-versions.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,6 @@
"cudaVersion": "12.9",
"systems": ["x86_64-linux", "aarch64-linux"]
},
{
"torchVersion": "2.9.0",
"rocmVersion": "6.3",
"systems": ["x86_64-linux"]
},
{
"torchVersion": "2.9.0",
"rocmVersion": "6.4",
"systems": ["x86_64-linux"]
},
{
"torchVersion": "2.9.0",
"xpuVersion": "2025.2.1",
Expand Down Expand Up @@ -85,5 +75,15 @@
"torchVersion": "2.11.0",
"metal": true,
"systems": ["aarch64-darwin"]
},
{
"torchVersion": "2.11.0",
"rocmVersion": "7.1",
"systems": ["x86_64-linux"]
},
{
"torchVersion": "2.11.0",
"rocmVersion": "7.2",
"systems": ["x86_64-linux"]
}
]
3 changes: 1 addition & 2 deletions nix-builder/pkgs/rocm-packages/default.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,8 @@ let
# Add aotriton
(final: prev: {
inherit (prev.callPackage ../aotriton { })
aotriton_0_10
aotriton_0_11
aotriton_0_11_1
aotriton_0_11_2
;
})
];
Expand Down
13 changes: 10 additions & 3 deletions nix-builder/pkgs/rocm-packages/overrides.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,7 @@ applyOverrides {
hipblas-common-devel ? null,
}:
prevAttrs: {
# Only available starting ROCm 6.3.
propagatedBuildInputs =
prevAttrs.buildInputs ++ lib.optionals (hipblas-common-devel != null) [ hipblas-common-devel ];
propagatedBuildInputs = prevAttrs.buildInputs ++ [ hipblas-common-devel ];
};

hipblaslt =
Expand Down Expand Up @@ -94,6 +92,15 @@ applyOverrides {
];
};

hipsparselt =
{ hip-runtime-amd, roctracer }:
prevAttrs: {
buildInputs = prevAttrs.buildInputs ++ [
hip-runtime-amd
roctracer
];
};

hsa-rocr =
{
elfutils,
Expand Down
Loading
Loading