Merge pull request #326228 from SomeoneSerge/fix/openai-triton

python312Packages.openai-triton: unbreak cuda

pkgs/development/python-modules/openai-triton/default.nix

@@ -27,7 +27,7 @@
 }:
 
 let
-  ptxas = "${cudaPackages.cuda_nvcc}/bin/ptxas"; # Make sure cudaPackages is the right version each update (See python/setup.py)
+  ptxas = lib.getExe' cudaPackages.cuda_nvcc "ptxas"; # Make sure cudaPackages is the right version each update (See python/setup.py)
 in
 buildPythonPackage rec {
   pname = "triton";
@@ -48,9 +48,12 @@ buildPythonPackage rec {
         url = "https://github.com/openai/triton/commit/52c146f66b79b6079bcd28c55312fc6ea1852519.patch";
         hash = "sha256-098/TCQrzvrBAbQiaVGCMaF3o5Yc3yWDxzwSkzIuAtY=";
       })
+
+      # Upstream startded pinning CUDA version and falling back to downloading from Conda
+      # in https://github.com/triton-lang/triton/pull/1574/files#diff-eb8b42d9346d0a5d371facf21a8bfa2d16fb49e213ae7c21f03863accebe0fcfR120-R123
+      ./0000-dont-download-ptxas.patch
     ]
     ++ lib.optionals (!cudaSupport) [
-      ./0000-dont-download-ptxas.patch
       # openai-triton wants to get ptxas version even if ptxas is not
       # used, resulting in ptxas not found error.
       ./0001-ptxas-disable-version-key-for-non-cuda-targets.patch
@@ -130,6 +133,12 @@ buildPythonPackage rec {
     setuptools
   ];
 
+  NIX_CFLAGS_COMPILE = lib.optionals cudaSupport [
+    # Pybind11 started generating strange errors since python 3.12. Observed only in the CUDA branch.
+    # https://gist.github.com/SomeoneSerge/7d390b2b1313957c378e99ed57168219#file-gistfile0-txt-L1042
+    "-Wno-stringop-overread"
+  ];
+
   # Avoid GLIBCXX mismatch with other cuda-enabled python packages
   preConfigure =
     ''

pkgs/top-level/python-packages.nix

@@ -9290,7 +9290,6 @@ self: super: with self; {
 
   openai-triton = callPackage ../development/python-modules/openai-triton {
     llvm = pkgs.openai-triton-llvm;
-    cudaPackages = pkgs.cudaPackages_12_1;
   };
 
   openai-triton-cuda = self.openai-triton.override {