python3Packages.triton-bin: init at 2.0.0
This commit is contained in:
parent
608908ed7a
commit
8c168726c5
82
pkgs/development/python-modules/openai-triton/bin.nix
Normal file
82
pkgs/development/python-modules/openai-triton/bin.nix
Normal file
|
@ -0,0 +1,82 @@
|
|||
{ lib
|
||||
, stdenv
|
||||
, addOpenGLRunpath
|
||||
, cudaPackages
|
||||
, buildPythonPackage
|
||||
, fetchurl
|
||||
, isPy38
|
||||
, isPy39
|
||||
, isPy310
|
||||
, isPy311
|
||||
, python
|
||||
, autoPatchelfHook
|
||||
, filelock
|
||||
, lit
|
||||
, pythonRelaxDepsHook
|
||||
, zlib
|
||||
}:
|
||||
|
||||
buildPythonPackage rec {
|
||||
pname = "triton";
|
||||
version = "2.0.0";
|
||||
format = "wheel";
|
||||
|
||||
src =
|
||||
let pyVerNoDot = lib.replaceStrings [ "." ] [ "" ] python.pythonVersion;
|
||||
unsupported = throw "Unsupported system";
|
||||
srcs = (import ./binary-hashes.nix version)."${stdenv.system}-${pyVerNoDot}" or unsupported;
|
||||
in fetchurl srcs;
|
||||
|
||||
disabled = !(isPy38 || isPy39 || isPy310 || isPy311);
|
||||
|
||||
pythonRemoveDeps = [ "cmake" "torch" ];
|
||||
|
||||
buildInputs = [ zlib ];
|
||||
|
||||
nativeBuildInputs = [
|
||||
pythonRelaxDepsHook # torch and triton refer to each other so this hook is included to mitigate that.
|
||||
autoPatchelfHook
|
||||
];
|
||||
|
||||
propagatedBuildInputs = [
|
||||
filelock
|
||||
lit
|
||||
zlib
|
||||
];
|
||||
|
||||
dontStrip = true;
|
||||
|
||||
# If this breaks, consider replacing with "${cuda_nvcc}/bin/ptxas"
|
||||
postFixup = ''
|
||||
chmod +x "$out/${python.sitePackages}/triton/third_party/cuda/bin/ptxas"
|
||||
'' +
|
||||
(let
|
||||
# Bash was getting weird without linting,
|
||||
# but basically upstream contains [cc, ..., "-lcuda", ...]
|
||||
# and we replace it with [..., "-lcuda", "-L/run/opengl-driver/lib", "-L$stubs", ...]
|
||||
old = [ "-lcuda" ];
|
||||
new = [ "-lcuda" "-L${addOpenGLRunpath.driverLink}" "-L${cudaPackages.cuda_cudart}/lib/stubs/" ];
|
||||
|
||||
quote = x: ''"${x}"'';
|
||||
oldStr = lib.concatMapStringsSep ", " quote old;
|
||||
newStr = lib.concatMapStringsSep ", " quote new;
|
||||
in
|
||||
''
|
||||
substituteInPlace $out/${python.sitePackages}/triton/compiler.py \
|
||||
--replace '${oldStr}' '${newStr}'
|
||||
'');
|
||||
|
||||
meta = with lib; {
|
||||
description = "A language and compiler for custom Deep Learning operations";
|
||||
homepage = "https://github.com/openai/triton/";
|
||||
changelog = "https://github.com/openai/triton/releases/tag/v${version}";
|
||||
# Includes NVIDIA's ptxas, but redistributions of the binary are not limited.
|
||||
# https://docs.nvidia.com/cuda/eula/index.html
|
||||
# triton's license is MIT.
|
||||
# openai-triton-bin includes ptxas binary, therefore unfreeRedistributable is set.
|
||||
license = with licenses; [ unfreeRedistributable mit ];
|
||||
sourceProvenance = with sourceTypes; [ binaryNativeCode ];
|
||||
platforms = [ "x86_64-linux" ];
|
||||
maintainers = with maintainers; [ junjihashimoto ];
|
||||
};
|
||||
}
|
|
@ -0,0 +1,31 @@
|
|||
# Warning: Need to update at the same time as torch-bin
|
||||
#
|
||||
# Precompiled wheels can be found at:
|
||||
# https://download.pytorch.org/whl/torch_stable.html
|
||||
|
||||
# To add a new version, run "prefetch.sh 'new-version'" to paste the generated file as follows.
|
||||
|
||||
version : builtins.getAttr version {
|
||||
"2.0.0" = {
|
||||
x86_64-linux-38 = {
|
||||
name = "triton-2.0.0-1-cp38-cp38-linux_x86_64.whl";
|
||||
url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.whl";
|
||||
hash = "sha256-nUl4KYt0/PWadf5x5TXAkrAjCIkzsvHfkz7DJhXkvu8=";
|
||||
};
|
||||
x86_64-linux-39 = {
|
||||
name = "triton-2.0.0-1-cp39-cp39-linux_x86_64.whl";
|
||||
url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl";
|
||||
hash = "sha256-dPEYwStDf7LKJeGgR1kXO1F1gvz0x74RkTMWx2QhNlY=";
|
||||
};
|
||||
x86_64-linux-310 = {
|
||||
name = "triton-2.0.0-1-cp310-cp310-linux_x86_64.whl";
|
||||
url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl";
|
||||
hash = "sha256-OIBu6WY/Sw981keQ6WxXk3QInlj0mqxKZggSGqVeJQU=";
|
||||
};
|
||||
x86_64-linux-311 = {
|
||||
name = "triton-2.0.0-1-cp311-cp311-linux_x86_64.whl";
|
||||
url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl";
|
||||
hash = "sha256-ImlBx7hZUhnd71mh/bgh6MdEKJoTJBXd1YT6zt60dbE=";
|
||||
};
|
||||
};
|
||||
}
|
40
pkgs/development/python-modules/openai-triton/prefetch.sh
Executable file
40
pkgs/development/python-modules/openai-triton/prefetch.sh
Executable file
|
@ -0,0 +1,40 @@
|
|||
#!/usr/bin/env nix-shell
|
||||
#!nix-shell -i bash -p nix-prefetch-scripts
|
||||
|
||||
set -eou pipefail
|
||||
|
||||
version=$1
|
||||
|
||||
linux_bucket="https://download.pytorch.org/whl"
|
||||
|
||||
url_and_key_list=(
|
||||
"x86_64-linux-38 $linux_bucket/triton-${version}-1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp38-cp38-linux_x86_64.whl"
|
||||
"x86_64-linux-39 $linux_bucket/triton-${version}-1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp39-cp39-linux_x86_64.whl"
|
||||
"x86_64-linux-310 $linux_bucket/triton-${version}-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp310-cp310-linux_x86_64.whl"
|
||||
"x86_64-linux-311 $linux_bucket/triton-${version}-1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp311-cp311-linux_x86_64.whl"
|
||||
)
|
||||
|
||||
hashfile=binary-hashes-"$version".nix
|
||||
echo " \"$version\" = {" >> $hashfile
|
||||
|
||||
for url_and_key in "${url_and_key_list[@]}"; do
|
||||
key=$(echo "$url_and_key" | cut -d' ' -f1)
|
||||
url=$(echo "$url_and_key" | cut -d' ' -f2)
|
||||
name=$(echo "$url_and_key" | cut -d' ' -f3)
|
||||
|
||||
echo "prefetching ${url}..."
|
||||
hash=$(nix hash to-sri --type sha256 `nix-prefetch-url "$url" --name "$name"`)
|
||||
|
||||
cat << EOF >> $hashfile
|
||||
$key = {
|
||||
name = "$name";
|
||||
url = "$url";
|
||||
hash = "$hash";
|
||||
};
|
||||
EOF
|
||||
|
||||
echo
|
||||
done
|
||||
|
||||
echo " };" >> $hashfile
|
||||
echo "done."
|
|
@ -6856,6 +6856,8 @@ self: super: with self; {
|
|||
|
||||
openai-triton = callPackage ../development/python-modules/openai-triton { llvmPackages = pkgs.llvmPackages_rocm; };
|
||||
|
||||
openai-triton-bin = callPackage ../development/python-modules/openai-triton/bin.nix { };
|
||||
|
||||
openai-whisper = callPackage ../development/python-modules/openai-whisper {
|
||||
cudaSupport = pkgs.config.cudaSupport or false;
|
||||
};
|
||||
|
|
Loading…
Reference in a new issue