[ConnorBaker]

This description was written entirely with AI (Claude Code, Claude Opus 5.5).

Description

Processes that ran an NVRTC-compiled kernel (norms and softmax since #3156, and others) or captured cuDNN fused attention into a CUDA graph can segfault at exit, after all their work is done. In pytest this turns a passing test file into exit status 139.

Both crashes come from static destructors that call CUDA after main returns. The CUDA Programming Guide, §2.1.6 says that "Using any of these interfaces (implicitly or explicitly) during program initiation or termination after main will result in undefined behavior", and it's Mandatory rule CUDA 1.1 (p. 8) in NVIDIA's CUDA C++ Guidelines v3.0.1.

• rtc::KernelManager: its destructor runs ~Kernel, which unloads modules through cuda_driver::call (rtc.cpp#L73-L92). cuda_driver::call caches driver entry points in function-local statics, which are usually first constructed after the KernelManager, when its first kernel is loaded. Statics are destroyed in reverse order of construction, so those caches are already destroyed when ~Kernel reads them.
• Fused-attention graph caches: destroying a cuDNN plan that a CUDA graph retains releases a CUDA user object (cuda_graph_retention.h#L159-L166), resolved through get_cuda_symbol()'s function-local map (cudnn_frontend_shim.h#L152-L171), which can likewise already be destroyed. That's [Bug]: CUDA graph plan retention can SIGSEGV during Python process teardown cudnn-frontend#1486. [Auto-FE] shim: keep the CUDA library handle map alive through process teardown (#1486) cudnn-frontend#1565 keeps that map alive on develop, but no release has it yet.

Whether the KernelManager case crashes depends on the libstdc++ TE was built with. With GCC 16's, the lookup walks freed nodes and faults. With GCC 13's (the NGC pytorch:26.09-py3 container), the destroyed map looks empty, so cuda_driver::call inserts into it, writing to freed memory, and the process exits 0. With MALLOC_PERTURB_=165, both crash.

Fix: KernelManager and the graph caches are allocated with new and never destroyed (rtc.cpp#L156-L158, graph_cache.h#L133-L135), so their destructors never run. TE already avoids exit-time teardown this way: TVMFFIConfigCache::create() (#3137) and the fused-attention debug registry (#2964) are never destroyed, and HandleManager never releases its handles (#1389). cuda_driver::call's symbol cache and mutex are also never destroyed (cuda_driver.h#L43-L46), so a lookup never reads a destroyed map.

Both caches only ever insert (rtc.cpp#L294, graph_cache.h#L120), so nothing changes before exit. The one behavioural difference is that a program that dlcloses libtransformer_engine.so and keeps running no longer unloads the NVRTC modules or destroys the cached plans.

Fixes # (no existing issue found)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• transformer_engine/common/util/rtc.cpp: KernelManager::instance() is never destroyed.
• transformer_engine/common/fused_attn/graph_cache.h: the fused-attention graph caches are never destroyed.
• transformer_engine/common/util/cuda_driver.h: cuda_driver::call's symbol cache and its mutex are never destroyed.
• tests/cpp/util/CMakeLists.txt: the util tests run with MALLOC_PERTURB_=165, so ctest fails a test whose process crashes at exit. UtilTest.NVRTC drives rtc::KernelManager directly (test_nvrtc.cpp#L66-L79).

Testing

On an RTX 4090 (driver 615.71.09) with CUDA 13.4, cuDNN 9.27, cudnn-frontend 1.31.0 (without NVIDIA/cudnn-frontend#1565), PyTorch 2.14.1 and GCC 16, building main (c9c9225) with and without this PR. The Python scripts ran 3 times plain and 3 times with MALLOC_PERTURB_=165.

	main	this PR
ctest on the util tests, with this PR's MALLOC_PERTURB_	UtilTest.NVRTC and UtilTest.NVRTCConcurrentCompile pass, then segfault at exit (3 of 3 runs)	all 5 pass (3 of 3 runs)
te.LayerNorm forward and backward	exit 139 in 5 of 6	exit 0 in 6 of 6
DotProductAttention (cuDNN) captured in a CUDA graph	exit 139 in 6 of 6	exit 0 in 6 of 6

I also tested both commits backported to TE 2.20.2 (NixOS/nixpkgs#577950) on the same 4090, with cudnn-frontend's fix applied to both builds. Without this PR, test_torch_compile.py exited 139 in 4 of 4 runs, and 12 of 20 L0 PyTorch test files exited 139 after their tests finished. With it, none did. The test counts and failing tests were the same: test_numerics.py tolerance failures, and out-of-memory errors in test_fused_rope.py on the 24 GiB card.

I haven't rerun the GB10 since adding the second commit.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code

[timmoon10]

Hm, this is an ugly bug and I don't have great clarity on the nicest solution. Leaking the CUDA Driver symbol cache does seem like it will work.

[timmoon10]

LayerNorm makes no guarantee to call the NVRTC kernel. It does by default, but it can switch to cuDNN by setting NVTE_NORM_FWD_USE_CUDNN=1/NVTE_NORM_BWD_USE_CUDNN=1. We ignore this envvar in certain cases (e.g. MXFP8 LayerNorm) and we may change the defaults in the future. It may be safer to make this test name more descriptive (test_exit_after_layernorm).

Suggested change

	@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU")
	@pytest.mark.skipif(bool(int(os.getenv("NVTE_DISABLE_NVRTC", "0"))), reason="Requires NVRTC")
	def test_exit_after_nvrtc_kernels():
	# MALLOC_PERTURB_ makes glibc overwrite freed memory, so an exit-time destructor that reads an
	# already-destroyed object crashes every time rather than only sometimes.
	env = dict(os.environ, MALLOC_PERTURB_="165")
	result = subprocess.run(
	[sys.executable, "-c", SCRIPT], env=env, capture_output=True, text=True, timeout=600
	)
	assert result.returncode == 0, f"exit status {result.returncode}\n{result.stderr}"
	def test_exit_after_layernorm():
	# MALLOC_PERTURB_ makes glibc overwrite freed memory, so an exit-time destructor that reads an
	# already-destroyed object crashes every time rather than only sometimes.
	env = dict(
	os.environ,
	# Make glibc overwrite freed memory, so an exit-time destructor that reads an
	# already-destroyed object crashes every time rather than only sometimes.
	MALLOC_PERTURB_="165",
	# Disable cuDNN LayerNorm kernel and use NVRTC kernel instead
	NVTE_NORM_FWD_USE_CUDNN="0",
	NVTE_NORM_BWD_USE_CUDNN="0",
	)
	result = subprocess.run(
	[sys.executable, "-c", SCRIPT], env=env, capture_output=True, text=True, timeout=600
	)
	assert result.returncode == 0, f"exit status {result.returncode}\n{result.stderr}"

[ConnorBaker]

This comment was written entirely with AI (Claude Code, Claude Opus 5.5).

Thanks for the review. I pushed ec38bcc, which changes the fix and the test.

Fix. The underlying problem is calling CUDA from static destructors at all. The CUDA Programming Guide, §2.1.6 says that "Using any of these interfaces (implicitly or explicitly) during program initiation or termination after main will result in undefined behavior". The legacy guide says the same in §6.2.1, and it's Mandatory rule CUDA 1.1 (p. 8) in NVIDIA's CUDA C++ Guidelines v3.0.1. Keeping the symbol cache alive only stopped ~Kernel's driver calls (rtc.cpp#L73-L92) from crashing; they still ran after main.

The fused-attention graph cache has the same problem. Destroying a cuDNN plan that a CUDA graph retains releases a CUDA user object (cuda_graph_retention.h#L159-L166), resolved through get_cuda_symbol()'s function-local map (cudnn_frontend_shim.h#L152-L171). That's NVIDIA/cudnn-frontend#1486.

KernelManager and the graph caches are now allocated with new and never deleted (rtc.cpp#L156-L158, graph_cache.h#L133-L135). TE already does this elsewhere:

• TVMFFIConfigCache::create(), from [Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI #3137 (tvm_ffi_bridge.h#L209-L212, #L383).
• The fused-attention debug registry, from [All] Refactor fused attention APIs with cuDNN-frontend support checks and opaque config/params handles #2964 (graph_cache_debug.h#L211-L223): "deliberately never freed … Leaking removes the ordering question".
• HandleManager, which never releases its handles (handle_manager.h#L17, Better cuBLAS handle management #1389).

Both caches only ever insert (rtc.cpp#L294, graph_cache.h#L120), so nothing changes before exit. The one behavioural difference is that a program that dlcloses libtransformer_engine.so and keeps running no longer unloads the NVRTC modules or destroys the cached plans. I kept the symbol-cache change as defence in depth.

Test. You're right that the LayerNorm test depended on the norm backend. NVTE_NORM_{FWD,BWD}_USE_CUDNN change production behaviour:

• Where: in common (common.cpp#L556-L566), PyTorch (normalization.cpp#L109-L130) and JAX (normalization.py#L66-L73).
• Undocumented: Document environment variables #2552 documented TE's environment variables while both switches existed (envvars.rst#L302-L318, common.cpp#L522-L527), but they aren't in docs/envvars.rst.
• The default may change: you suggested making cuDNN the default.

So test_process_exit.py is gone. Instead, UtilTest.NVRTC, which drives rtc::KernelManager directly (test_nvrtc.cpp#L66-L79), runs with MALLOC_PERTURB_=165 (CMakeLists.txt#L24-L26). ctest fails it if the process crashes at exit.

Results on an RTX 4090 (driver 615.71.09) with CUDA 13.4, cuDNN 9.27, cudnn-frontend 1.31.0 (without NVIDIA/cudnn-frontend#1565), PyTorch 2.14.1 and GCC 16, building main (c9c9225) with and without this PR. The Python scripts ran 3 times plain and 3 times with MALLOC_PERTURB_=165.

	main	this PR
ctest on the util tests, with this PR's MALLOC_PERTURB_	UtilTest.NVRTC and UtilTest.NVRTCConcurrentCompile pass, then segfault at exit (3 of 3 runs)	all 5 pass (3 of 3 runs)
te.LayerNorm forward and backward	exit 139 in 5 of 6	exit 0 in 6 of 6
DotProductAttention (cuDNN) captured in a CUDA graph	exit 139 in 6 of 6	exit 0 in 6 of 6

I haven't rerun the GB10.

The CUDA-graph attention case has no test here, because the C++ tests have no attention tests. I can add a PyTorch one that pins NVTE_FUSED_ATTN=1 and NVTE_FLASH_ATTN=0, if you want it in this PR.