tjh.dev
Clone of https://github.com/NixOS/nixpkgs.git (to stress-test knotserver)
1{
2 lib,
3 stdenv,
4 cudaPackages,
5 buildPythonPackage,
6 fetchurl,
7 python,
8 autoPatchelfHook,
9 zlib,
10}:
11
12buildPythonPackage (finalAttrs: {
13 pname = "triton";
14 version = "3.6.0";
15 format = "wheel";
16
17 src =
18 let
19 pyVerNoDot = lib.replaceStrings [ "." ] [ "" ] python.pythonVersion;
20 unsupported = throw "Unsupported system";
21 srcs =
22 (import ./binary-hashes.nix finalAttrs.version)."${stdenv.system}-${pyVerNoDot}" or unsupported;
23 in
24 fetchurl srcs;
25
26 pythonRemoveDeps = [
27 "cmake"
28 # torch and triton refer to each other so this hook is included to mitigate that.
29 "torch"
30 ];
31
32 buildInputs = [ zlib ];
33
34 nativeBuildInputs = [
35 autoPatchelfHook
36 ];
37
38 dontStrip = true;
39
40 # If this breaks, consider replacing with "${cuda_nvcc}/bin/ptxas"
41 postFixup = ''
42 mkdir -p $out/${python.sitePackages}/triton/third_party/cuda/bin/
43 ln -s ${cudaPackages.cuda_nvcc}/bin/ptxas $out/${python.sitePackages}/triton/third_party/cuda/bin/
44 '';
45
46 meta = {
47 description = "Language and compiler for custom Deep Learning operations";
48 homepage = "https://github.com/triton-lang/triton/";
49 changelog = "https://github.com/triton-lang/triton/releases/tag/v${finalAttrs.version}";
50 # Includes NVIDIA's ptxas, but redistributions of the binary are not limited.
51 # https://docs.nvidia.com/cuda/eula/index.html
52 # triton's license is MIT.
53 # triton-bin includes ptxas binary, therefore unfreeRedistributable is set.
54 license = with lib.licenses; [
55 unfreeRedistributable
56 mit
57 ];
58 sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
59 maintainers = with lib.maintainers; [
60 GaetanLepage
61 junjihashimoto
62 ];
63 };
64})