cudaPackages.cudnn-frontend: 1.31.0 -> 1.32.0; python3Packages.transformer-engine: backport fixes for crashes at process exit (#577950)

pkgs/development/cuda-modules/packages/cudnn-frontend/package.nix

@@ -35,13 +35,13 @@ backendStdenv.mkDerivation (finalAttrs: {
   name = "${cudaNamePrefix}-${finalAttrs.pname}-${finalAttrs.version}";
 
   pname = "cudnn-frontend";
-  version = "1.31.0";
+  version = "1.32.0";
 
   src = fetchFromGitHub {
     owner = "NVIDIA";
     repo = "cudnn-frontend";
     tag = "v${finalAttrs.version}";
-    hash = "sha256-/KMprsFaxKdNRJ9APbb+NH7bL1izSayO6l9BVRxOCRE=";
+    hash = "sha256-xrLk0gD5cvbX6qxdIFcFmHsHl6JpGTBG+U1Eu+3pJ1Q=";
   };
 
   # nlohmann_json should be the only vendored dependency.

pkgs/development/python-modules/nvidia-cudnn-frontend/default.nix

@@ -5,8 +5,8 @@
 
   # build-system
   cmake,
+  nanobind_3,
   ninja,
-  pybind11,
   setuptools,
 
   # nativeBuildInputs
@@ -42,11 +42,17 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs
     cudaPackages.cudnn-frontend.postPatch
     + ''
       substituteInPlace pyproject.toml \
-        --replace-fail '"ninja==1.11.1.1"' '"ninja"' \
-        --replace-fail '"pybind11[global]>=2.13,<3"' '"pybind11"'
+        --replace-fail '"ninja==1.11.1.1;' '"ninja;'
 
       sed -i '/cmake_args =/a\\f"-DCUDNN_FRONTEND_USE_SYSTEM_DLPACK=ON",' setup.py
     ''
+    # The bindings include standard headers before Python.h, so pyconfig.h redefines the
+    # _POSIX_C_SOURCE and _XOPEN_SOURCE that glibc already defined, with different values.
+    # That warning is an error under upstream's -Werror.
+    + ''
+      substituteInPlace CMakeLists.txt \
+        --replace-fail '-Wpedantic -Werror ' '-Wpedantic '
+    ''
     + ''
       substituteInPlace python/cudnn/__init__.py \
         --replace-fail \
@@ -56,8 +62,8 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs
 
   build-system = [
     cmake
+    nanobind_3
     ninja
-    pybind11
     setuptools
   ];
 
@@ -101,4 +107,9 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs
     requiredSystemFeatures = [ "cuda" ];
     doCheck = true;
   };
+
+  # The version and source are cudaPackages.cudnn-frontend's, so updating it updates this package too.
+  passthru.updateScript = cudaPackages.cudnn-frontend.updateScript // {
+    attrPath = "cudaPackages.cudnn-frontend";
+  };
 })

pkgs/development/python-modules/transformer-engine/default.nix

@@ -3,6 +3,7 @@
   config,
   buildPythonPackage,
   fetchFromGitHub,
+  fetchpatch,
   replaceVars,
   python,
   cudaPackages,
@@ -107,15 +108,32 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
     hash = "sha256-YXo0LIq13HNVrYa19npDoLRrB2NY4IjZb3E+asblBnQ=";
   };
 
-  patches = optionals cudaSupport [
-    (replaceVars ./cuda-libs-paths.patch {
-      libcudnn_so = "${getLib cudaPackages.cudnn}/lib/libcudnn.so";
-      libnvrtc_so = "${getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so";
-      libcurand_so = "${getLib cudaPackages.libcurand}/lib/libcurand.so";
+  patches =
+    # Never destroy cuda_driver::call's symbol caches, the KernelManager (which unloads its NVRTC modules) or the
+    # fused-attention graph caches (which destroy cuDNN plans retained by CUDA graphs). Without these, processes that ran
+    # an NVRTC kernel or captured fused attention into a CUDA graph can segfault at exit.
+    # https://github.com/NVIDIA/TransformerEngine/pull/3733 (merged after 2.20.2); drop it once a release contains it.
+    [
+      (fetchpatch {
+        name = "never-destroy-exit-statics.patch";
+        url = "https://github.com/NVIDIA/TransformerEngine/commit/dd628d3a11b7b52908ea5b7bba7d38d739df6f55.patch";
+        includes = [
+          "transformer_engine/common/fused_attn/graph_cache.h"
+          "transformer_engine/common/util/cuda_driver.h"
+          "transformer_engine/common/util/rtc.cpp"
+        ];
+        hash = "sha256-96O66tXnfVMWw7L39if1ME4GTlTvSAroOyh+fRFIIC8=";
+      })
+    ]
+    ++ optionals cudaSupport [
+      (replaceVars ./cuda-libs-paths.patch {
+        libcudnn_so = "${getLib cudaPackages.cudnn}/lib/libcudnn.so";
+        libnvrtc_so = "${getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so";
+        libcurand_so = "${getLib cudaPackages.libcurand}/lib/libcurand.so";
 
-      cudart_include_dir = "${getInclude cudaPackages.cuda_cudart}/include";
-    })
-  ];
+        cudart_include_dir = "${getInclude cudaPackages.cuda_cudart}/include";
+      })
+    ];
 
   postPatch =
     # Patch build-system requirements: