Unverified Commit 8f0af006 authored by Gaétan Lepage's avatar Gaétan Lepage Committed by GitHub
Browse files

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

parents a8fef1cf 4b374913

pkgs/development/python-modules/torch/source/default.nix

+7 −0
Original line number Diff line number Diff line
@@ -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

0 → 100644
+45 −0
Original line number Diff line number Diff line 
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