pyrox.dev
python312Packages.triton: fix torch.compile (#369495)
+10
-1
pkgs/development/python-modules/triton/default.nix
+10
-1
pkgs/development/python-modules/triton/default.nix
···
53
53
}
54
54
// lib.optionalAttrs rocmSupport { libhipDir = "${lib.getLib rocmPackages.clr}/lib"; }
55
55
// lib.optionalAttrs cudaSupport {
56
-
libcudaStubsDir = "${lib.getLib cudaPackages.cuda_cudart}/lib/stubs";
56
+
libcudaStubsDir = "${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs";
57
57
ccCmdExtraFlags = "-Wl,-rpath,${addDriverRunpath.driverLink}/lib";
58
58
}
59
59
))
···
82
82
substituteInPlace unittest/CMakeLists.txt \
83
83
--replace-fail "include (\''${CMAKE_CURRENT_SOURCE_DIR}/googletest.cmake)" ""\
84
84
--replace-fail "include(GoogleTest)" "find_package(GTest REQUIRED)"
85
+
86
+
# Patch the source code to make sure it doesn't specify a non-existent PTXAS version.
87
+
# CUDA 12.6 (the current default/max) tops out at PTXAS version 8.5.
88
+
# NOTE: This is fixed in `master`:
89
+
# https://github.com/triton-lang/triton/commit/f48dbc1b106c93144c198fbf3c4f30b2aab9d242
90
+
substituteInPlace "$NIX_BUILD_TOP/$sourceRoot/third_party/nvidia/backend/compiler.py" \
91
+
--replace-fail \
92
+
'return 80 + minor' \
93
+
'return 80 + min(minor, 5)'
85
94
'';
86
95
87
96
build-system = [ setuptools ];