From 18e5b70b5aca9b5d834b22ebdc35f2c70876da3a Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Fri, 23 Oct 2026 09:34:31 -0700 Subject: [PATCH 1/2] [Common] Keep cuda_driver::call's symbol cache alive through exit rtc::KernelManager's static destructor unloads its CUDA modules through cuda_driver::call. Each instantiation of call caches driver entry points in function-local statics, and those caches are usually first constructed after the KernelManager, so they are destroyed before it. ~Kernel then reads a destroyed unordered_map: with GCC 16's libstdc++ the process segfaults at exit; with GCC 13's it writes into freed memory instead. Allocate the cache and its mutex with new and never destroy them, and add a test that the process exits cleanly after running NVRTC kernels. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Connor Baker --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_process_exit.py | 36 ++++++++++++++++++++ transformer_engine/common/util/cuda_driver.h | 6 ++-- 3 files changed, 41 insertions(+), 2 deletions(-) create mode 100644 tests/pytorch/test_process_exit.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f45e9369d10..e8d4d012b00 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -31,6 +31,7 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_optional_flash_attn_import.xml $TE_PATH/tests/pytorch/test_optional_flash_attn_import.py || test_fail "test_optional_flash_attn_import.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_process_exit.xml $TE_PATH/tests/pytorch/test_process_exit.py || test_fail "test_process_exit.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_deferred_init.xml $TE_PATH/tests/pytorch/test_deferred_init.py || test_fail "test_deferred_init.py" diff --git a/tests/pytorch/test_process_exit.py b/tests/pytorch/test_process_exit.py new file mode 100644 index 00000000000..f8ac57ae465 --- /dev/null +++ b/tests/pytorch/test_process_exit.py @@ -0,0 +1,36 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""The process exits cleanly after NVRTC kernels have been loaded.""" + +import os +import subprocess +import sys + +import pytest +import torch + +# LayerNorm runs NVRTC-compiled kernels, so rtc::KernelManager holds loaded CUDA modules that its +# static destructor unloads at exit. +SCRIPT = """ +import torch +import transformer_engine.pytorch as te + +layer = te.LayerNorm(256).cuda() +x = torch.randn(64, 256, device="cuda", requires_grad=True) +layer(x).sum().backward() +torch.cuda.synchronize() +""" + + +@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}" diff --git a/transformer_engine/common/util/cuda_driver.h b/transformer_engine/common/util/cuda_driver.h index 16242347f12..71ca4024ad8 100644 --- a/transformer_engine/common/util/cuda_driver.h +++ b/transformer_engine/common/util/cuda_driver.h @@ -40,8 +40,10 @@ template inline CUresult call(const char *symbol, ArgTs... args) { using FuncT = CUresult(ArgTs...); - static std::unordered_map symbol_cache; - static std::mutex cache_mutex; + // Never destroyed. Static destructors may call into the driver during exit (rtc::KernelManager + // unloads its CUDA modules), and statics constructed after them are destroyed before them. + static auto &symbol_cache = *new std::unordered_map(); + static auto &cache_mutex = *new std::mutex(); FuncT *func; { From ec38bcc8daf532120bde266b2b703a2ef7957554 Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Fri, 23 Oct 2026 23:57:30 -0700 Subject: [PATCH 2/2] [Common] Never destroy rtc::KernelManager or the fused-attention graph caches Both were function-local statics whose destructors call into CUDA after main returns: ~Kernel unloads modules through the driver, and destroying a cached cuDNN plan that a CUDA graph retains releases a CUDA user object. The CUDA Programming Guide makes any use of CUDA during program termination undefined behavior, and both crashed at exit depending on static destruction order. Allocate both with new and never delete them, as TVMFFIConfigCache::create() and the fused-attention debug registry already do. Both caches only ever insert, so nothing changes while the process runs. Replace tests/pytorch/test_process_exit.py, whose LayerNorm only runs NVRTC kernels while the NVTE_NORM_FWD_USE_CUDNN and NVTE_NORM_BWD_USE_CUDNN switches are off, with MALLOC_PERTURB_=165 for the C++ util tests. UtilTest.NVRTC drives rtc::KernelManager directly, and ctest checks each test's exit status. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Connor Baker --- qa/L0_pytorch_unittest/test.sh | 1 - tests/cpp/util/CMakeLists.txt | 4 ++- tests/pytorch/test_process_exit.py | 36 ------------------- .../common/fused_attn/graph_cache.h | 4 ++- transformer_engine/common/util/cuda_driver.h | 4 +-- transformer_engine/common/util/rtc.cpp | 4 ++- 6 files changed, 11 insertions(+), 42 deletions(-) delete mode 100644 tests/pytorch/test_process_exit.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index e8d4d012b00..f45e9369d10 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -31,7 +31,6 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_optional_flash_attn_import.xml $TE_PATH/tests/pytorch/test_optional_flash_attn_import.py || test_fail "test_optional_flash_attn_import.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" -python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_process_exit.xml $TE_PATH/tests/pytorch/test_process_exit.py || test_fail "test_process_exit.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_deferred_init.xml $TE_PATH/tests/pytorch/test_deferred_init.py || test_fail "test_deferred_init.py" diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index 1dfd2fed4e4..b80a588fe35 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -21,4 +21,6 @@ target_link_libraries(test_util PUBLIC target_compile_options(test_util PRIVATE -O2 -fopenmp) include(GoogleTest) -gtest_discover_tests(test_util DISCOVERY_TIMEOUT 600) +# MALLOC_PERTURB_ makes glibc overwrite freed memory, so a destructor that reads a destroyed object +# at exit crashes the test process, which ctest reports as a failure, instead of passing silently. +gtest_discover_tests(test_util DISCOVERY_TIMEOUT 600 PROPERTIES ENVIRONMENT "MALLOC_PERTURB_=165") diff --git a/tests/pytorch/test_process_exit.py b/tests/pytorch/test_process_exit.py deleted file mode 100644 index f8ac57ae465..00000000000 --- a/tests/pytorch/test_process_exit.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""The process exits cleanly after NVRTC kernels have been loaded.""" - -import os -import subprocess -import sys - -import pytest -import torch - -# LayerNorm runs NVRTC-compiled kernels, so rtc::KernelManager holds loaded CUDA modules that its -# static destructor unloads at exit. -SCRIPT = """ -import torch -import transformer_engine.pytorch as te - -layer = te.LayerNorm(256).cuda() -x = torch.randn(64, 256, device="cuda", requires_grad=True) -layer(x).sum().backward() -torch.cuda.synchronize() -""" - - -@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}" diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h index b3020f25acf..d52a0daf70b 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -130,7 +130,9 @@ std::shared_ptr> cache_graph(GraphCache auto get_graph(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - static GraphCache cache; + // Never destroyed: destroying cached cuDNN plans calls into CUDA, which must not be used during + // program termination. + static auto &cache = *new GraphCache; cfg.check_derived(); return detail::cache_graph(cache, cfg.make_cache_key(kPass), kBackend, kPass, handle, [&] { return kCreateGraphFn(cfg); }); diff --git a/transformer_engine/common/util/cuda_driver.h b/transformer_engine/common/util/cuda_driver.h index 71ca4024ad8..3b85bd7cce4 100644 --- a/transformer_engine/common/util/cuda_driver.h +++ b/transformer_engine/common/util/cuda_driver.h @@ -40,8 +40,8 @@ template inline CUresult call(const char *symbol, ArgTs... args) { using FuncT = CUresult(ArgTs...); - // Never destroyed. Static destructors may call into the driver during exit (rtc::KernelManager - // unloads its CUDA modules), and statics constructed after them are destroyed before them. + // Never destroyed, so a lookup never reads a destroyed map, whatever order statics are + // constructed and destroyed in. static auto &symbol_cache = *new std::unordered_map(); static auto &cache_mutex = *new std::mutex(); FuncT *func; diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index 2ffbbf72334..f591620124c 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -153,7 +153,9 @@ int Kernel::occupancy_max_active_blocks_per_sm(int device_id, int block_size, KernelManager& KernelManager::instance() { NVTE_CHECK(is_enabled(), "NVRTC support is not enabled"); - static KernelManager instance_; + // Never destroyed: ~Kernel unloads modules through the CUDA driver, and CUDA must not be used + // during program termination. + static KernelManager& instance_ = *new KernelManager; return instance_; }