diff --git a/pkgs/development/python-modules/triton/default.nix b/pkgs/development/python-modules/triton/default.nix index 9b904f5b1728..6e088a1921f5 100644 --- a/pkgs/development/python-modules/triton/default.nix +++ b/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= + "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 ]; }; -} +})