[ConnorBaker]

Description

Processes that ran an NVRTC-compiled kernel (LayerNorm, RMSNorm, softmax and others since #3156) can segfault at exit, after all their work is done. In pytest this turns a passing test file into exit status 139.

Cause: an exit-time destructor uses function-local statics that were already destroyed.

• cuda_driver::call caches driver entry points in function-local statics, symbol_cache and cache_mutex. Each instantiation of the template has its own, and call<int *, int> (for cuDeviceGet) is a different instantiation from call<CUcontext *, int> (for cuDevicePrimaryCtxRetain).
• Kernel::~Kernel unloads its modules through five driver calls: cuDeviceGet, cuDevicePrimaryCtxRetain, cuCtxSetCurrent, cuModuleUnload and cuDevicePrimaryCtxRelease.
• It runs inside KernelManager::~KernelManager, the destructor of the KernelManager::instance() static. Static destructors run in reverse order of construction (__cxa_atexit). So if an instantiation's cache was constructed after the KernelManager, it's already destroyed when ~Kernel uses it.

In the common case, that's every one of them:

1. The first KernelManager::compile constructs the KernelManager.
2. The first launch then calls Kernel::get_function, which makes the first cuDeviceGet, cuDevicePrimaryCtxRetain and similar calls in the process. That constructs those instantiations' caches after the KernelManager.

Which caches outlive the KernelManager depends on which driver calls ran before the first NVRTC compile. For example, cuda::supports_multicast() also calls cuDeviceGet.

Traced in gdb on an unpatched build (watchpoints on the guard variables, and breakpoints on __cxa_atexit and on operator delete of the cache's bucket array):

1. The KernelManager is constructed, and its destructor registered, from compile.
2. The call<int *, int> cache is constructed, and its destructor registered, from get_function.
3. At exit, __run_exit_handlers frees that cache's bucket array.
4. Only then does ~KernelManager start. ~Kernel faults in cuda_driver::call("cuDeviceGet", ...), reading that cache:

#0  std::_Hashtable<std::string, std::pair<const std::string, void*>, ...>::_M_locate_tr(...)
#1  transformer_engine::rtc::Kernel::~Kernel()
#2  transformer_engine::rtc::KernelManager::~KernelManager()
#3  __run_exit_handlers

Calling transformer_engine_torch.device_supports_multicast() before the first LayerNorm constructs the call<int *, int> cache first. Its destructor then runs after ~KernelManager, and the crash moves to the next call in ~Kernel: the call<CUcontext *, int> cache, for cuDevicePrimaryCtxRetain. That cache was still constructed by get_function. This held in all 3 crashes sampled with each order.

Whether it crashes depends on the C++ standard library the extension was built with. The use of a destroyed cache happens either way.

• With GCC 16's libstdc++: ~_Hashtable frees the nodes and the bucket array but leaves _M_element_count and _M_before_begin as they were. The lookup then walks a freed node whose first two words glibc has overwritten with free-list links, and faults.
  • A TE 2.20.2 build on a DGX Spark crashed in 15 of 15 runs of a LayerNorm-only script, and in 6 of 7 runs of a LayerNormMLP script. The one clean run had PYTHONFAULTHANDLER=1 set.
• With GCC 13's libstdc++ (the NGC pytorch:26.09-py3 container, TE 2.19.0+5e52befd): ~_Hashtable calls clear() first, so the destroyed cache looks empty. The lookup misses, and cuda_driver::call inserts the symbol back into the destroyed map. At exit, the map holds a new node while its bucket pointer is still the freed array, which is a write to freed memory, and the process exits 0.
• With MALLOC_PERTURB_=165 (glibc overwrites freed memory), both crashed in 3 of 3 runs, with the stack above.

The code involved is unchanged on main.

Fix: allocate symbol_cache and cache_mutex with new and never destroy them. Any static destructor can then call into the driver through cuda_driver::call, whatever the construction order. Today that's ~KernelManager, but it would also cover later ones. The cost is one map and one mutex per instantiation, left for the OS to reclaim at exit.

~Kernel itself is unchanged and still unloads its modules. With the fix, at exit on the GB10, all five of its driver calls returned CUDA_SUCCESS for each of the 3 loaded modules.

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/cuda_driver.h: cuda_driver::call's symbol cache and its mutex are never destroyed.
• tests/pytorch/test_process_exit.py (new): runs LayerNorm forward and backward in a subprocess with MALLOC_PERTURB_=165 and checks that the subprocess exits with status 0.
• qa/L0_pytorch_unittest/test.sh: runs the new test.

Testing

On each platform, two builds of TE against PyTorch 2.14.1 and CUDA 13.4, differing only by this change:

• DGX Spark (GB10, SM 12.1, aarch64): TE 2.20.2 and main (c9c9225).
• RTX 4090 (SM 8.9, x86_64): main.

	Without this change	With it
LayerNormMLP script (gelu, swiglu, geglu; forward and backward)	exit 139: GB10 2.20.2 in 6 of 7, GB10 main in 3 of 3, 4090 main in 5 of 5	exit 0 in every run (5, 3 and 5)
Same, MALLOC_PERTURB_=165	exit 139 in every run (3 per build)	exit 0 in every run (5, 3 and 3)
te.LayerNorm only, MALLOC_PERTURB_=165	exit 139 in every run (3 per build)	exit 0 in every run (3 per build)
te.LayerNorm after device_supports_multicast(), MALLOC_PERTURB_=165	exit 139 in every run (GB10 2.20.2: 3; GB10 main: 2; 4090 main: 3), in the cuDevicePrimaryCtxRetain cache	exit 0 (GB10 main: 2; 4090 main: 3)
te.Linear only (no NVRTC kernels), MALLOC_PERTURB_=165	exit 0 in every run	exit 0 in every run
tests/pytorch/test_process_exit.py	fails on all three (subprocess exit -11)	passes on all three

L0 files: the GB10 2.20.2 builds ran 18 files of qa/L0_pytorch_unittest/test.sh, and the 4090 main builds ran 7. With and without the change, every file had the same pass, fail, skip and xfail counts.

• Files whose tests all passed exited with 139 without the change and 0 with it: test_recipe.py, test_jit.py, test_cuda_graphs.py, test_quantized_tensor.py, test_torch_compile.py and test_fused_optimizer.py.
• Files with failing tests went from 139 to 1. Those failures are the same with and without the change, and the GB10 ones also occur in the NGC container: for example, fp32 gated activations exceed the GB10's shared-memory limit, and grouped GEMM's supported SM range excludes SM 12.1.

On the RTX 4090, test_torch_compile.py can still exit 139 with this change, from a different bug: NVIDIA/cudnn-frontend#1486.

• At exit, fused_attn's GraphCache static destroys cuDNN plans that were captured into a CUDA graph. Releasing their graph user objects goes through cudnn-frontend's get_cuda_symbol(), whose function-local handle map has already been destroyed. It's the same pattern as this bug.
• [Auto-FE] shim: keep the CUDA library handle map alive through process teardown (#1486) cudnn-frontend#1565 fixes it on develop the same way this PR does, but it isn't in a release yet: 1.31.0 predates it, and TE requires nvidia-cudnn-frontend>=1.29.0.
• On the 4090 with only this change, a fused-attention script that captures into a CUDA graph exited 139 in 6 of 6 runs, half of them with MALLOC_PERTURB_=165. test_torch_compile.py exited 139 in 1 of 2 plain runs and 2 of 2 runs with MALLOC_PERTURB_=165.
• With this change plus [JAX] Unbalanced Context Parallelism with THD format #1565 applied to cudnn-frontend 1.31.0, all of those exited 0, with the same test counts. On the GB10, gdb showed that cudnn-frontend's handle map is intact when it's read at exit.

The compiler warnings are the same with and without the change, for all three pairs of builds, apart from nvcc's per-build temporary file names and hashes.

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}"