From 2586c3adf7562aaf5b54f3c2dc84994306fa7a19 Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Mon, 30 Dec 2024 21:18:46 +0000 Subject: [PATCH] python3Packages.triton: fix pytorch.compile --- pkgs/development/python-modules/triton/default.nix | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/pkgs/development/python-modules/triton/default.nix b/pkgs/development/python-modules/triton/default.nix index 214bc9de1e6a..69ddb4ad716c 100644 --- a/pkgs/development/python-modules/triton/default.nix +++ b/pkgs/development/python-modules/triton/default.nix @@ -53,7 +53,7 @@ buildPythonPackage { } // lib.optionalAttrs rocmSupport { libhipDir = "${lib.getLib rocmPackages.clr}/lib"; } // lib.optionalAttrs cudaSupport { - libcudaStubsDir = "${lib.getLib cudaPackages.cuda_cudart}/lib/stubs"; + libcudaStubsDir = "${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs"; ccCmdExtraFlags = "-Wl,-rpath,${addDriverRunpath.driverLink}/lib"; } )) @@ -82,6 +82,15 @@ buildPythonPackage { substituteInPlace unittest/CMakeLists.txt \ --replace-fail "include (\''${CMAKE_CURRENT_SOURCE_DIR}/googletest.cmake)" ""\ --replace-fail "include(GoogleTest)" "find_package(GTest REQUIRED)" + + # Patch the source code to make sure it doesn't specify a non-existent PTXAS version. + # CUDA 12.6 (the current default/max) tops out at PTXAS version 8.5. + # NOTE: This is fixed in `master`: + # https://github.com/triton-lang/triton/commit/f48dbc1b106c93144c198fbf3c4f30b2aab9d242 + substituteInPlace "$NIX_BUILD_TOP/$sourceRoot/third_party/nvidia/backend/compiler.py" \ + --replace-fail \ + 'return 80 + minor' \ + 'return 80 + min(minor, 5)' ''; build-system = [ setuptools ];