Merge pull request #329310 from SomeoneSerge/refactor/move-triton

python312Packages.triton: rename from openai-triton

pkgs/by-name/tr/triton-llvm/package.nix

@@ -36,7 +36,7 @@ let
   llvmTargetsToBuild' = [ "AMDGPU" "NVPTX" ] ++ builtins.map inferNativeTarget llvmTargetsToBuild;
 
   # This LLVM version can't seem to find pygments/pyyaml,
-  # but a later update will likely fix this (openai-triton-2.1.0)
+  # but a later update will likely fix this (triton-2.1.0)
   python =
     if buildTests
     then python3Packages.python.withPackages (p: with p; [ psutil pygments pyyaml ])
@@ -44,7 +44,7 @@ let
 
   isNative = stdenv.hostPlatform == stdenv.buildPlatform;
 in stdenv.mkDerivation (finalAttrs: {
-  pname = "openai-triton-llvm";
+  pname = "triton-llvm";
   version = "17.0.0-c5dede880d17";
 
   outputs = [
@@ -55,7 +55,7 @@ in stdenv.mkDerivation (finalAttrs: {
     "man"
   ];
 
-  # See https://github.com/openai/triton/blob/main/python/setup.py
+  # See https://github.com/triton-lang/triton/blob/main/python/setup.py
   # and https://github.com/ptillet/triton-llvm-releases/releases
   src = fetchFromGitHub {
     owner = "llvm";

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

@@ -15,7 +15,7 @@
   more-itertools,
   numba,
   numpy,
-  openai-triton,
+  triton,
   tiktoken,
   torch,
   tqdm,
@@ -53,7 +53,7 @@ buildPythonPackage rec {
     tiktoken
     torch
     tqdm
-  ] ++ lib.optionals (lib.meta.availableOn stdenv.hostPlatform openai-triton) [ openai-triton ];
+  ] ++ lib.optionals (lib.meta.availableOn stdenv.hostPlatform triton) [ triton ];
 
   preCheck = ''
     export HOME=$TMPDIR

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

@@ -40,7 +40,7 @@
   tabulate,
   tiktoken,
   transformers,
-  openai-triton,
+  triton,
   xformers,
 }:
 
@@ -117,7 +117,7 @@ buildPythonPackage rec {
       # auto-gptq
     ]; # ++ autogptq.optional-dependencies.triton;
     grpc = [ bentoml ] ++ bentoml.optional-dependencies.grpc;
-    mpt = [ openai-triton ];
+    mpt = [ triton ];
     openai = [
       openai
       tiktoken

pkgs/development/python-modules/torch/bin.nix

@@ -22,7 +22,7 @@
   jinja2,
   networkx,
   filelock,
-  openai-triton,
+  triton,
 }:
 
 let
@@ -88,7 +88,7 @@ buildPythonPackage {
     jinja2
     networkx
     filelock
-  ] ++ lib.optionals (stdenv.isLinux && stdenv.isx86_64) [ openai-triton ];
+  ] ++ lib.optionals (stdenv.isLinux && stdenv.isx86_64) [ triton ];
 
   postInstall = ''
     # ONNX conversion

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

@@ -53,9 +53,9 @@
   cffi,
   click,
   typing-extensions,
-  # ROCm build and `torch.compile` requires `openai-triton`
+  # ROCm build and `torch.compile` requires `triton`
   tritonSupport ? (!stdenv.isDarwin),
-  openai-triton,
+  triton,
 
   # Unit tests
   hypothesis,
@@ -486,7 +486,7 @@ buildPythonPackage rec {
       CoreServices
       libobjc
     ]
-    ++ lib.optionals tritonSupport [ openai-triton ]
+    ++ lib.optionals tritonSupport [ triton ]
     ++ lib.optionals MPISupport [ mpi ]
     ++ lib.optionals rocmSupport [ rocmtoolkit_joined ];
 
@@ -514,7 +514,7 @@ buildPythonPackage rec {
 
     # torch/csrc requires `pybind11` at runtime
     pybind11
-  ] ++ lib.optionals tritonSupport [ openai-triton ];
+  ] ++ lib.optionals tritonSupport [ triton ];
 
   propagatedCxxBuildInputs =
     [ ] ++ lib.optionals MPISupport [ mpi ] ++ lib.optionals rocmSupport [ rocmtoolkit_joined ];

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

@@ -80,12 +80,12 @@ buildPythonPackage rec {
 
   meta = with lib; {
     description = "Language and compiler for custom Deep Learning operations";
-    homepage = "https://github.com/openai/triton/";
-    changelog = "https://github.com/openai/triton/releases/tag/v${version}";
+    homepage = "https://github.com/triton-lang/triton/";
+    changelog = "https://github.com/triton-lang/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.
+    # triton-bin includes ptxas binary, therefore unfreeRedistributable is set.
     license = with licenses; [
       unfreeRedistributable
       mit

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

@@ -53,7 +53,7 @@ buildPythonPackage rec {
       ./0000-dont-download-ptxas.patch
     ]
     ++ lib.optionals (!cudaSupport) [
-      # openai-triton wants to get ptxas version even if ptxas is not
+      # 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
     ];
@@ -127,7 +127,7 @@ buildPythonPackage rec {
 
   propagatedBuildInputs = [
     filelock
-    # openai-triton uses setuptools at runtime:
+    # triton uses setuptools at runtime:
     # https://github.com/NixOS/nixpkgs/pull/286763/#discussion_r1480392652
     setuptools
   ];
@@ -197,7 +197,7 @@ buildPythonPackage rec {
     # and pythonImportsCheck is commented back.
     import-triton =
       runCommand "import-triton"
-        { nativeBuildInputs = [ (python.withPackages (ps: [ ps.openai-triton ])) ]; }
+        { nativeBuildInputs = [ (python.withPackages (ps: [ ps.triton ])) ]; }
         ''
           python << \EOF
           import triton

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

@@ -18,7 +18,7 @@
   fairscale,
   scipy,
   cmake,
-  openai-triton,
+  triton,
   networkx,
   #, apex
   einops,
@@ -103,7 +103,7 @@ buildPythonPackage {
     scipy
     cmake
     networkx
-    openai-triton
+    triton
     # apex
     einops
     transformers

pkgs/development/rocm-modules/5/miopengemm/default.nix

@@ -67,7 +67,7 @@ in stdenv.mkDerivation (finalAttrs: {
     openblas
   ] ++ lib.optionals buildBenchmarks [
     clblast
-    python3Packages.openai-triton
+    python3Packages.triton
   ];
 
   cmakeFlags = [

pkgs/top-level/aliases.nix

@@ -1399,6 +1399,7 @@ mapAliases ({
   transfig = fig2dev; # Added 2022-02-15
   transifex-client = transifex-cli; # Added 2023-12-29
   trezor_agent = trezor-agent; # Added 2024-01-07
+  openai-triton-llvm = triton-llvm; # added 2024-07-18
   trustedGrub = throw "trustedGrub has been removed, because it is not maintained upstream anymore"; # Added 2023-05-10
   trustedGrub-for-HP = throw "trustedGrub-for-HP has been removed, because it is not maintained upstream anymore"; # Added 2023-05-10
   tumpa = throw "tumpa has been removed, as it is broken"; # Added 2024-07-15

pkgs/top-level/python-aliases.nix

@@ -339,6 +339,10 @@ mapAliases ({
   notifymuch = throw "notifymuch has been promoted to a top-level attribute name: `pkgs.notifymuch`"; # added 2022-10-02
   Nuitka = nuitka; # added 2023-02-19
   ntlm-auth = throw "ntlm-auth has been removed, because it relies on the md4 implementation provided by openssl. Use pyspnego instead.";
+  openai-triton = triton; # added 2024-07-18
+  openai-triton-bin = triton-bin; # added 2024-07-18
+  openai-triton-cuda = triton-cuda; # added 2024-07-18
+  openai-triton-no-cuda = triton-no-cuda; # added 2024-07-18
   openapi-schema-pydantic = throw "openapi-schema-pydantic has been removed, since it is no longer maintained"; # added 2023-10-30
   opencv3 = throw "opencv3 has been removed as it is obsolete"; # added 2023-10-12
   opsdroid_get_image_size = opsdroid-get-image-size; # added 2023-10-16

pkgs/top-level/python-packages.nix

@@ -4572,7 +4572,7 @@ self: super: with self; {
   oelint-parser = callPackage ../development/python-modules/oelint-parser { };
 
   openllm = callPackage ../development/python-modules/openllm {
-    openai-triton = self.openai-triton-cuda;
+    triton = self.triton-cuda;
   };
 
   openllm-client = callPackage ../development/python-modules/openllm-client { };
@@ -9367,20 +9367,6 @@ self: super: with self; {
 
   open-meteo = callPackage ../development/python-modules/open-meteo { };
 
-  openai-triton = callPackage ../development/python-modules/openai-triton {
-    llvm = pkgs.openai-triton-llvm;
-  };
-
-  openai-triton-cuda = self.openai-triton.override {
-    cudaSupport = true;
-  };
-
-  openai-triton-no-cuda = self.openai-triton.override {
-    cudaSupport = false;
-  };
-
-  openai-triton-bin = callPackage ../development/python-modules/openai-triton/bin.nix { };
-
   openai-whisper = callPackage ../development/python-modules/openai-whisper { };
 
   openant = callPackage ../development/python-modules/openant { };
@@ -15580,13 +15566,13 @@ self: super: with self; {
   torch-pitch-shift = callPackage ../development/python-modules/torch-pitch-shift { };
 
   torch-bin = callPackage ../development/python-modules/torch/bin.nix {
-    openai-triton = self.openai-triton-bin;
+    triton = self.triton-bin;
   };
 
   torchsnapshot = callPackage ../development/python-modules/torchsnapshot { };
 
   torchWithCuda = self.torch.override {
-    openai-triton = self.openai-triton-cuda;
+    triton = self.triton-cuda;
     cudaSupport = true;
     rocmSupport = false;
   };
@@ -15596,7 +15582,7 @@ self: super: with self; {
   };
 
   torchWithRocm = self.torch.override {
-    openai-triton = self.openai-triton-no-cuda;
+    triton = self.triton-no-cuda;
     rocmSupport = true;
     cudaSupport = false;
   };
@@ -15755,6 +15741,20 @@ self: super: with self; {
 
   trio-websocket = callPackage ../development/python-modules/trio-websocket { };
 
+  triton = callPackage ../development/python-modules/triton {
+    llvm = pkgs.triton-llvm;
+  };
+
+  triton-cuda = self.triton.override {
+    cudaSupport = true;
+  };
+
+  triton-no-cuda = self.triton.override {
+    cudaSupport = false;
+  };
+
+  triton-bin = callPackage ../development/python-modules/triton/bin.nix { };
+
   tritonclient = callPackage ../development/python-modules/tritonclient { };
 
   troposphere = callPackage ../development/python-modules/troposphere { };

pkgs/top-level/release-cuda.nix

@@ -132,7 +132,7 @@ let
         mxnet = linux;
         numpy = linux; # Only affected by MKL?
         onnx = linux;
-        openai-triton = linux;
+        triton = linux;
         openai-whisper = linux;
         opencv4 = linux;
         opensfm = linux;