python3Packages.triton: fix libcublas dlopen, python3Packages.triton.gpuCheck: skip failing tests (#488887)

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

@@ -4,6 +4,7 @@
   config,
   buildPythonPackage,
   fetchFromGitHub,
+  pythonAtLeast,
 
   # patches
   replaceVars,
@@ -41,7 +42,7 @@
   cudaSupport ? config.cudaSupport,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "triton";
   version = "3.5.1";
   pyproject = true;
@@ -50,7 +51,7 @@ buildPythonPackage rec {
   src = fetchFromGitHub {
     owner = "triton-lang";
     repo = "triton";
-    tag = "v${version}";
+    tag = "v${finalAttrs.version}";
     hash = "sha256-dyNRtS1qtU8C/iAf0Udt/1VgtKGSvng1+r2BtvT9RB4=";
   };
 
@@ -99,6 +100,14 @@ buildPythonPackage rec {
       substituteInPlace cmake/AddTritonUnitTest.cmake \
         --replace-fail "include(\''${PROJECT_SOURCE_DIR}/unittest/googletest.cmake)" ""\
         --replace-fail "include(GoogleTest)" "find_package(GTest REQUIRED)"
+    ''
+
+    # triton will try dlopening libcublas.so at runtime
+    + lib.optionalString cudaSupport ''
+      substituteInPlace third_party/nvidia/include/cublas_instance.h \
+        --replace-fail \
+          '"libcublas.so"' \
+          '"${lib.getLib cudaPackages.libcublas}/lib/libcublas.so"'
     '';
 
   build-system = [ setuptools ];
@@ -217,9 +226,63 @@ buildPythonPackage rec {
 
       doCheck = true;
 
-      preCheck = ''
-        cd python/test/unit
-      '';
+      disabledTests = [
+        # triton.runtime.errors.OutOfResources: out of resource: shared memory,
+        # Required: 131072, Hardware limit: 101376. Reducing block sizes or `num_stages` may help.
+        "test_gather"
+        "test_gather_warp_shuffle"
+        "test_tensor_descriptor_batched_gemm_2d_tma"
+        "test_tensor_descriptor_batched_gemm_3d_tma"
+
+        # AssertionError: assert all(delta == 0 for delta in diff.values())
+        # ----------------------------- Captured stdout call -----------------------------
+        # Expected line "pid (0, 0, 0) idx ( 0,   0) x: 1" 1 time(s), but saw 0 time(s)
+        # ...
+        "test_print"
+
+        # This test ensures that the ptxas binary is available under .../site-packages/triton/backends/nvidia/bin/ptxas
+        # Usually, this is where the install script downloads and copies ptxas to.
+        # However, this is not the case here, as triton is built with TRITON_OFFLINE_BUILD=1
+        # and TRITON_PTXAS_PATH=<path_to_nix_store_ptxas>
+        "test_nvidia_tool"
+      ]
+      ++ lib.optionals (pythonAtLeast "3.14") [
+        # triton.compiler.errors.CompilationError
+        # AttributeError("module 'ast' has no attribute 'Num'"
+        "test_aggregate_modification_in_for_loop"
+        "test_call"
+        "test_call_in_loop"
+        "test_compile_only_k_loop"
+        "test_dot_mulbroadcasted"
+        "test_globaltimer"
+        "test_host_tensor_descriptor_matmul"
+        "test_make_tensor_descriptor_matmul"
+        "test_nested_while"
+        "test_preshuffle_scale_mxfp_cdna4"
+        "test_temp_var_in_loop"
+        "test_tensor_descriptor_rank_reducing_matmul"
+        "test_tensor_descriptor_reshape_matmul"
+      ];
+
+      disabledTestPaths = [
+        # torch.AcceleratorError: CUDA error: device-side assert triggered
+        "python/test/unit/test_debug.py"
+
+        # ptxas fatal   : Unexpected non-ASCII character encountered on line 1
+        # ptxas fatal   : Ptx assembly aborted due to errors
+        "python/test/unit/language/test_line_info.py"
+
+        # Triton Error [CUDA]: \n
+        "python/test/unit/tools/test_aot.py"
+
+        # ptxas fatal   : Unknown option 'sass'
+        "python/test/unit/tools/test_disasm.py"
+      ];
+
+      enabledTestPaths = [
+        "python/test/unit"
+      ];
+
       checkPhase = "pytestCheckPhase";
 
       installPhase = "touch $out";
@@ -314,11 +377,13 @@ buildPythonPackage rec {
   meta = {
     description = "Language and compiler for writing highly efficient custom Deep-Learning primitives";
     homepage = "https://github.com/triton-lang/triton";
+    changelog = "https://github.com/triton-lang/triton/releases/tag/${finalAttrs.src.tag}";
     platforms = lib.platforms.linux;
     license = lib.licenses.mit;
     maintainers = with lib.maintainers; [
+      GaetanLepage
       SomeoneSerge
       derdennisop
     ];
   };
-}
+})