python3Packages.triton: fix pytorch.compile

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 ];