From f6cb12e27ac3c96c93e3717b79f085c7b52ba2d6 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sun, 7 Jun 2026 01:00:44 +0200
Subject: [PATCH 1/2] python3Packages.triton-bin: enable __structuredAttrs

---
 pkgs/development/python-modules/triton/bin.nix | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pkgs/development/python-modules/triton/bin.nix b/pkgs/development/python-modules/triton/bin.nix
index b78e4fd38707..3e0af134cc1e 100644
--- a/pkgs/development/python-modules/triton/bin.nix
+++ b/pkgs/development/python-modules/triton/bin.nix
@@ -13,6 +13,7 @@ buildPythonPackage (finalAttrs: {
   pname = "triton";
   version = "3.6.0";
   format = "wheel";
+  __structuredAttrs = true;
 
   src =
     let

From f012a712b4576d41bccf8b4a791f581f7fd4c347 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sun, 7 Jun 2026 01:01:10 +0200
Subject: [PATCH 2/2] python3Packages.triton-bin: revisit cuda & rocm patching

---
 .../development/python-modules/triton/bin.nix | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/pkgs/development/python-modules/triton/bin.nix b/pkgs/development/python-modules/triton/bin.nix
index 3e0af134cc1e..f6589b72645a 100644
--- a/pkgs/development/python-modules/triton/bin.nix
+++ b/pkgs/development/python-modules/triton/bin.nix
@@ -1,12 +1,17 @@
 {
   lib,
   stdenv,
+  config,
   cudaPackages,
   buildPythonPackage,
   fetchurl,
   python,
   autoPatchelfHook,
   zlib,
+  rocmPackages,
+
+  rocmSupport ? config.rocmSupport,
+  addDriverRunpath,
 }:
 
 buildPythonPackage (finalAttrs: {
@@ -38,10 +43,25 @@ buildPythonPackage (finalAttrs: {
 
   dontStrip = true;
 
-  # If this breaks, consider replacing with "${cuda_nvcc}/bin/ptxas"
   postFixup = ''
     mkdir -p $out/${python.sitePackages}/triton/third_party/cuda/bin/
     ln -s ${cudaPackages.cuda_nvcc}/bin/ptxas $out/${python.sitePackages}/triton/third_party/cuda/bin/
+  ''
+  # Ugly patch to circumvent the heuristic-based search logic of libcuda.so
+  + ''
+    substituteInPlace "${placeholder "out"}/${python.sitePackages}/triton/backends/nvidia/driver.py" \
+      --replace-fail \
+        "if env_libcuda_path := knobs.nvidia.libcuda_path:" \
+        "return '${addDriverRunpath.driverLink}/lib/libcuda.so'" \
+      --replace-fail \
+        "return [env_libcuda_path]" \
+        ""
+  ''
+  + lib.optionalString rocmSupport ''
+    substituteInPlace "${placeholder "out"}/${python.sitePackages}/triton/backends/amd/driver.py" \
+      --replace-fail \
+        "lib_name = "libamdhip64.so" \
+        "return '${rocmPackages.clr}/lib/libamdhip64.so'"
   '';
 
   meta = {