gador/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

python3Packages.torch: don't copy AOTriton to output (#515088)

Browse Source
This commit is contained in:
Gaétan Lepage
2026-04-30 22:12:07 +00:00
committed by GitHub
parent a8fef1cf38 4b374913c8
commit 8f0af00611
2 changed files with 52 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkgs/development/python-modules/torch/source/default.nix
+7
View File
@@ -339,6 +339,11 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
# with the Nix store, which fails. Simply remove this step to get
# rpaths that point to the Nix store.
./disable-cmake-mkl-rpath.patch
]
++ lib.optionals rocmSupport [
# [ROCm] Make AOTriton bundling optional via BUILD_AOTRITON_INTO_WHEEL flag
# https://github.com/pytorch/pytorch/pull/182030
./no-bundle-aotriton.patch
];
postPatch = ''
@@ -506,6 +511,8 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
}
// lib.optionalAttrs rocmSupport {
AOTRITON_INSTALLED_PREFIX = "${rocmPackages.aotriton}";
# Don't copy AOTriton to output, load from AOTriton package
BUILD_AOTRITON_INTO_WHEEL = false;
# Broken HIP flag setup, fails to compile due to not finding rocthrust
# Only supports gfx942 so let's turn it off for now
USE_FBGEMM_GENAI = setBool false;
pkgs/development/python-modules/torch/source/no-bundle-aotriton.patch
+45
View File
@@ -0,0 +1,45 @@
From 6ed1a4dbd96b1e2efe0513b9e6ff283c5f74cfa1 Mon Sep 17 00:00:00 2001
From: Luna Nova <git@lunnova.dev>
Date: Thu, 30 Apr 2026 07:29:42 -0700
Subject: [PATCH] [ROCm] Make AOTriton bundling optional via
`BUILD_AOTRITON_INTO_WHEEL` flag
Previously AOTriton was always copied to Torch's output/wheel.
This is unnecessary and can waste a gigabyte of space easily for distros
which have AOTriton as an installable package.
---
cmake/External/aotriton.cmake | 16 ++++++++++++----
1 file changed, 12 insertions(+), 4 deletions(-)
diff --git a/cmake/External/aotriton.cmake b/cmake/External/aotriton.cmake
index d2f1e15ff11..1bf580f8059 100644
--- a/cmake/External/aotriton.cmake
+++ b/cmake/External/aotriton.cmake
@@ -239,12 +239,20 @@ if(NOT __AOTRITON_INCLUDED)
message(STATUS "Download AOTriton pre-compiled GPU images from ${__AOTRITON_URL}.")
endfunction()
+ # By default AOTriton's lib/ and include/ are distributed with torch
+ # adding hundreds of MiB and avoiding a separate AOTriton install
+ # For distros which have AOTriton as an installable package this can be disabled
+ # to avoid duplicating AOTriton into torch
+ option(BUILD_AOTRITON_INTO_WHEEL "Copy AOTriton lib/ and include/ into the torch install tree so they ship in the wheel" ON)
+
# Note it is INSTALL"ED"
if(DEFINED ENV{AOTRITON_INSTALLED_PREFIX})
- install(DIRECTORY
- $ENV{AOTRITON_INSTALLED_PREFIX}/lib
- $ENV{AOTRITON_INSTALLED_PREFIX}/include
- DESTINATION ${__AOTRITON_INSTALL_DIR})
+ if(BUILD_AOTRITON_INTO_WHEEL)
+ install(DIRECTORY
+ $ENV{AOTRITON_INSTALLED_PREFIX}/lib
+ $ENV{AOTRITON_INSTALLED_PREFIX}/include
+ DESTINATION ${__AOTRITON_INSTALL_DIR})
+ endif()
set(__AOTRITON_INSTALL_DIR "$ENV{AOTRITON_INSTALLED_PREFIX}")
message(STATUS "Using Preinstalled AOTriton at ${__AOTRITON_INSTALL_DIR}")
elseif(DEFINED ENV{AOTRITON_INSTALL_FROM_SOURCE})
--
2.53.0