From 4b374913c803bc2bd29ef32c5fd46cc870da1783 Mon Sep 17 00:00:00 2001
From: Luna Nova <git@lunnova.dev>
Date: Thu, 30 Apr 2026 07:48:04 -0700
Subject: [PATCH] python3Packages.torch: don't copy AOTriton to output

Avoids duplicating AOTriton in torch's store path.
---
 .../python-modules/torch/source/default.nix   |  7 +++
 .../torch/source/no-bundle-aotriton.patch     | 45 +++++++++++++++++++
 2 files changed, 52 insertions(+)
 create mode 100644 pkgs/development/python-modules/torch/source/no-bundle-aotriton.patch

diff --git a/pkgs/development/python-modules/torch/source/default.nix b/pkgs/development/python-modules/torch/source/default.nix
index 0f5f5d57137e..d9a475d7512e 100644
--- a/pkgs/development/python-modules/torch/source/default.nix
+++ b/pkgs/development/python-modules/torch/source/default.nix
@@ -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;
diff --git a/pkgs/development/python-modules/torch/source/no-bundle-aotriton.patch b/pkgs/development/python-modules/torch/source/no-bundle-aotriton.patch
new file mode 100644
index 000000000000..dfc204d4aca7
--- /dev/null
+++ b/pkgs/development/python-modules/torch/source/no-bundle-aotriton.patch
@@ -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
+