From 93e498197aa313d6197cac1a15290c6887e7953d Mon Sep 17 00:00:00 2001 From: dylanyunlon Date: Thu, 13 Aug 2026 09:42:26 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20thread=5Flocal=20reentrant=20guard=20?= =?UTF-8?q?=E2=80=94=20prevent=20cudaMalloc=20infinite=20recursion?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CUB CachingDeviceAllocator 内部在 cache miss 时调 cudaMalloc， 被我们的 LD_PRELOAD 再次拦截 → DeviceAllocate → cudaMalloc → 无限递归 → segfault。 加 thread_local bool inside_cub 标志: 外部调用 → CUB allocator (带缓存) CUB 内部调用 → 直接走 dlsym(RTLD_NEXT) 的真实 cudaMalloc --- .../cccl_preload/cccl_allocator_preload.cu | 58 +++++++++++-------- 1 file changed, 34 insertions(+), 24 deletions(-) diff --git a/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu b/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu index 8d3de4ec..9e304eac 100644 --- a/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu +++ b/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu @@ -8,6 +8,10 @@ * caching allocator. Strips expandable_segments from * PYTORCH_CUDA_ALLOC_CONF before libtorch reads it. * + * CRITICAL: CUB's allocator internally calls cudaMalloc/cudaFree for + * cache misses. We use a thread-local reentrant guard so internal calls + * go straight to the real CUDA runtime, avoiding infinite recursion. + * * Source: CCCL cub/cub/util_allocator.cuh (BSD-3, NVIDIA) * Build: bash build_cccl_preload.sh */ @@ -24,32 +28,21 @@ /* ======================================================================== * Configuration for BI-V100 (32GB × 4 cards) - * - * CUB CachingDeviceAllocator parameters: - * bin_growth = 2 (power-of-2 bins: 256B, 512B, 1KB, ... 4GB) - * min_bin = 8 (2^8 = 256B minimum allocation) - * max_bin = 32 (2^32 = 4GB maximum cached bin) - * max_cached = 8GB per device - * - * More granular bins (growth=2) than CUB default (growth=8) because - * PyTorch tensor sizes vary widely in inference. * ======================================================================== */ static constexpr unsigned int ALLOC_BIN_GROWTH = 2; -static constexpr unsigned int ALLOC_MIN_BIN = 8; /* 256 bytes */ -static constexpr unsigned int ALLOC_MAX_BIN = 32; /* 4 GB */ -static constexpr size_t ALLOC_MAX_CACHED = (size_t)8 * 1024 * 1024 * 1024; /* 8GB */ +static constexpr unsigned int ALLOC_MIN_BIN = 8; /* 2^8 = 256 bytes */ +static constexpr unsigned int ALLOC_MAX_BIN = 32; /* 2^32 = 4 GB */ +static constexpr size_t ALLOC_MAX_CACHED = (size_t)8 * 1024 * 1024 * 1024; /* ---- Global allocator singleton ---- */ using CubAllocator = CUB_NS_QUALIFIER::CachingDeviceAllocator; static CubAllocator& get_allocator() { static CubAllocator instance( - ALLOC_BIN_GROWTH, - ALLOC_MIN_BIN, - ALLOC_MAX_BIN, + ALLOC_BIN_GROWTH, ALLOC_MIN_BIN, ALLOC_MAX_BIN, ALLOC_MAX_CACHED, - true /* skip_cleanup: CoreX may tear down CUDA before our dtor */ + true /* skip_cleanup: CoreX may tear down CUDA before dtor */ ); return instance; } @@ -57,6 +50,17 @@ static CubAllocator& get_allocator() { static bool g_preload_active = false; static bool g_debug = false; +/* + * Reentrant guard: CUB's DeviceAllocate/DeviceFree internally call + * cudaMalloc/cudaFree (for cache misses and evictions). Without this + * guard, our intercept would call DeviceAllocate again → infinite + * recursion → stack overflow → segfault. + * + * When inside_cub == true, cudaMalloc/cudaFree go straight to the + * real CUDA runtime via dlsym(RTLD_NEXT). + */ +static thread_local bool inside_cub = false; + /* ---- Real cudaMalloc/cudaFree via dlsym(RTLD_NEXT) ---- */ using RealMalloc_t = cudaError_t (*)(void**, size_t); using RealFree_t = cudaError_t (*)(void*); @@ -109,11 +113,8 @@ static void cccl_preload_init() { alloc_conf, clean.empty() ? "(unset)" : clean.c_str()); } - /* Initialize allocator */ auto& alloc = get_allocator(); - if (g_debug) { - alloc.debug = true; - } + if (g_debug) alloc.debug = true; g_preload_active = true; fprintf(stderr, @@ -125,20 +126,29 @@ static void cccl_preload_init() { /* ======================================================================== * cudaMalloc / cudaFree intercepts + * + * outside CUB → route to CUB CachingDeviceAllocator (bin + cache) + * inside CUB → pass through to real cudaMalloc/cudaFree (no recursion) * ======================================================================== */ extern "C" cudaError_t cudaMalloc(void** devPtr, size_t size) { - if (!g_preload_active) { + if (!g_preload_active || inside_cub) { return get_real_malloc()(devPtr, size); } - return get_allocator().DeviceAllocate(devPtr, size); + inside_cub = true; + cudaError_t err = get_allocator().DeviceAllocate(devPtr, size); + inside_cub = false; + return err; } extern "C" cudaError_t cudaFree(void* devPtr) { - if (!g_preload_active || devPtr == nullptr) { + if (!g_preload_active || inside_cub || devPtr == nullptr) { return get_real_free()(devPtr); } - return get_allocator().DeviceFree(devPtr); + inside_cub = true; + cudaError_t err = get_allocator().DeviceFree(devPtr); + inside_cub = false; + return err; }