[ConnorBaker]

cudnn-frontend 1.31.0, on master since 2026-10-09, keeps execution plans alive for the CUDA graphs recorded from them (NVIDIA/cudnn-frontend#1195). Those plans are released through get_cuda_symbol(), whose library handle map is a function-local static. A static constructed before that map is destroyed after it at exit, so when it releases such a plan, the lookup reads freed memory. TransformerEngine's fused-attention graph cache is one such static, and its processes can crash with SIGSEGV at exit (NVIDIA/cudnn-frontend#1486).

Upstream fixed this on develop in NVIDIA/cudnn-frontend#1565 (NVIDIA/cudnn-frontend@473a2c8), which never destroys the map or its mutex. No release has it yet; 1.31.0 is the latest. This applies that commit with fetchpatch.

python3Packages.nvidia-cudnn-frontend builds its extension from the same source and already reuses the package's src and postPatch, so it now takes patches as well. python3Packages.transformer-engine compiles against cudaPackages.cudnn-frontend's headers, so it gets the fix through them.

TransformerEngine has a second crash at exit, in its own code, which this PR also fixes. cuda_driver::call caches driver entry points in function-local statics. The static KernelManager that owns its NVRTC-compiled kernels (LayerNorm, RMSNorm, softmax and others) is usually constructed before those caches, so it's destroyed after them, and its destructor unloads the kernels' modules through the destroyed caches. The second commit applies both commits of NVIDIA/TransformerEngine#3733 (open) with fetchpatch. The first never destroys those caches. The second never destroys the KernelManager or the fused-attention graph caches, so neither calls CUDA at exit, and TransformerEngine no longer releases retained cuDNN plans at exit. Only the library changes are taken; the rest of those commits is tests and CI wiring.

This also adds cudaPackages.tests.cudnn-frontend.teardown, built with AddressSanitizer. It resolves a CUDA symbol through get_cuda_symbol() from the destructor of a static constructed before the map. Resolving a symbol only loads libcudart, so it needs no GPU.

Testing

• cudaPackages.cudnn-frontend (all outputs, including samples and tests) and python3Packages.nvidia-cudnn-frontend build on x86_64-linux and aarch64-linux, and the installed cudnn_frontend_shim.h has the fix.
• cudaPackages.tests.cudnn-frontend.teardown passes on x86_64-linux and aarch64-linux.
  • Against the same headers without the patch, it fails on both with AddressSanitizer: heap-use-after-free at cudnn_frontend_shim.h:163, the map lookup in get_cuda_symbol().
  • The read comes from the static's destructor during exit, on memory freed by the map's destructor.
• On x86_64-linux with cudaSupport = true and cudaCapabilities = [ "8.9" ], I built python3Packages.transformer-engine from this branch, with JAX and PyTorch. That includes Torch 2.14.1, the JAX CUDA plugin and PJRT, onnxruntime and torchvision. I also built its passthru.tests: withCusolvermp, withMpi and withNvshmem; withJax, withPytorch and withOutNcclEp are the default derivation.
• On an RTX 4090, all with the same Torch, I compared TransformerEngine built from three commits: the parent (3ec6363), the cudnn-frontend commit alone (2c6636c), and this branch.
  • A fused-attention script captured in a CUDA graph exited with 139 in 6 of 6 runs from the parent (3 plain, 3 with MALLOC_PERTURB_=165), and with 0 in 6 of 6 from both of the others. This crash is the cudnn-frontend one.
  • test_torch_compile.py exited with 139 in 4 of 4 runs (2 plain, 2 with MALLOC_PERTURB_=165) from both the parent and the cudnn-frontend commit alone, after all its tests passed. From this branch it exited with 0 in 4 of 4. All three ran 92 passed, 71 skipped, 1 xpassed. For the cudnn-frontend commit alone, gdb places the crash in transformer_engine::rtc::KernelManager::~KernelManager → rtc::Kernel::~Kernel, reading cuda_driver::call's symbol cache. That's the TransformerEngine one.
  • TransformerEngine's PyTorch L0 tests (20 files from 2.20.2's tests/pytorch): with the cudnn-frontend commit alone, 12 of 20 files exited with 139 after their tests finished. From this branch, none did. Both have the same failing tests: 72 test_numerics.py tolerance failures and CUDA out-of-memory errors in test_fused_rope.py on the 24 GiB card.
• Earlier, to isolate the cudnn-frontend fix, I built it in a downstream overlay: the same commit applied to cudnn-frontend 1.31.0, and the first commit of [Common] Don't call CUDA from static destructors at exit NVIDIA/TransformerEngine#3733 (then its only one, NVIDIA/TransformerEngine@18e5b70), which leaves the graph caches alone, applied to both arms.
  • RTX 4090, TransformerEngine main at NVIDIA/TransformerEngine@c9c92255: the fused-attention script exited with 139 in 6 of 6 runs without the fix and with 0 in 6 of 6 with it.
  • test_torch_compile.py on the same machine: without the fix, it exited with 139 in 1 of 2 plain runs and in 2 of 2 runs with MALLOC_PERTURB_=165. With the fix, it exited with 0 in every run. Both arms had the same test counts: 277 passed, 148 skipped, 1 xpassed.
  • GB10 (aarch64): reading the freed map did not fault there. With the fix, gdb showed the map intact when it was read at exit, and 12 of 12 runs exited with 0.
• treefmt, ci/nixpkgs-vet.sh and nix-build ci -A eval.baseline for x86_64-linux and aarch64-linux pass locally at this branch's head.

The changes, the test and this description were written with Claude Code (Claude Opus 5.5); the commits carry an Assisted-by: trailer.

Things done

• Fits CONTRIBUTING.md, pkgs/README.md, maintainers/README.md and other READMEs.
• Follows the automation/AI policy.
• Built on platform:
  • x86_64-linux
  • aarch64-linux
  • aarch64-darwin
• Tested, as applicable:
  • Basic functionality of all binary files, usually in ./result/bin/.
  • Ran nixpkgs-review on this PR. See nixpkgs-review usage.
  • Package tests at passthru.tests.
  • NixOS tests in nixos/tests.
  • Tests in lib/tests or pkgs/test for functions and "core" functionality.
• Nixpkgs Release Notes
  • Package update: when the change is major or breaking.
• NixOS Release Notes
  • Module addition: when adding a new NixOS module.
  • Module update: when the change is significant.

🤖 Generated with Claude Code