From 8dc6462a2b9382991c03f460d92f4514cdaf921e Mon Sep 17 00:00:00 2001 From: dylanyunlon Date: Thu, 13 Aug 2026 09:21:57 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20CCCL=20CachingDeviceAllocator=20LD=5FPR?= =?UTF-8?q?ELOAD=20=E2=80=94=20bypass=20CoreX=20expandable=5Fsegments=20AS?= =?UTF-8?q?SERT?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 从 CCCL upstream cub/cub/util_allocator.cuh 提取 CachingDeviceAllocator 核心算法，去掉所有 CUB/CCCL 宏依赖，编译为独立 .so。 用 LD_PRELOAD 拦截 cudaMalloc/cudaFree，路由到 CUB 的 geometric-bin 缓存分配器。同时在 constructor 中 strip PYTORCH_CUDA_ALLOC_CONF 里的 expandable_segments 配置，避免 CoreX CUDACachingAllocator.cpp:545 ASSERT。 BI-V100 调优参数: bin_growth=8, min_bin=3 (512B), max_bin=13 (~550MB) max_cached_bytes=4GB per device (32GB卡的合理上限) 真机测试步骤: 1. bash build_cccl_preload.sh 2. LD_PRELOAD=./libcccl_allocator.so CCCL_ALLOC_DEBUG=1 \ PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ python3 verify_preload.py --- .../cccl_preload/build_cccl_preload.sh | 116 +++-- .../cccl_preload/cccl_allocator_preload.cu | 465 +++++++++++++++--- .../cccl_preload/verify_preload.py | 2 +- 3 files changed, 456 insertions(+), 127 deletions(-) diff --git a/qwen3_6_scripts/cccl_preload/build_cccl_preload.sh b/qwen3_6_scripts/cccl_preload/build_cccl_preload.sh index b97fde7c..0d7c53c7 100755 --- a/qwen3_6_scripts/cccl_preload/build_cccl_preload.sh +++ b/qwen3_6_scripts/cccl_preload/build_cccl_preload.sh @@ -1,43 +1,41 @@ #!/usr/bin/env bash -# Build libcccl_allocator.so -# -# Full CCCL dependency chain (288 headers) in ./include/ -# Source: cccl_upstream/cub/cub/util_allocator.cuh + transitive deps +# Build libcccl_allocator.so — LD_PRELOAD .so for CUB CachingDeviceAllocator # # Usage: # bash build_cccl_preload.sh [output_dir] +# +# On BI-V100 with CoreX SDK: +# bash build_cccl_preload.sh /workspace/qwen3_6_scripts/cccl_preload +# +# The .so intercepts cudaMalloc/cudaFree and routes through CUB's +# caching allocator, bypassing CoreX's "expandable segment not supported" +# ASSERT in CUDACachingAllocator.cpp:545. set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" OUTPUT_DIR="${1:-${SCRIPT_DIR}}" SRC="${SCRIPT_DIR}/cccl_allocator_preload.cu" -INC="${SCRIPT_DIR}/include" OUT="${OUTPUT_DIR}/libcccl_allocator.so" -[[ -d "${INC}/cub" ]] || { echo "CCCL include tree missing: ${INC}/cub"; exit 2; } -[[ -d "${INC}/cuda" ]] || { echo "CCCL include tree missing: ${INC}/cuda"; exit 2; } +# Find CoreX clang++ (preferred) or system g++ +if [[ -x /usr/local/corex-3.2.3/bin/clang++ ]]; then + CXX=/usr/local/corex-3.2.3/bin/clang++ + echo "[build] Using CoreX clang++: ${CXX}" +elif [[ -x /usr/local/corex/bin/clang++ ]]; then + CXX=/usr/local/corex/bin/clang++ + echo "[build] Using CoreX clang++ (alt): ${CXX}" +else + CXX=g++ + echo "[build] CoreX clang++ not found, falling back to g++" +fi -# Find compiler -CXX="" -for candidate in \ - /usr/local/corex-3.2.3/bin/clang++ \ - /usr/local/corex/bin/clang++ \ - /usr/local/corex/lib64/clang/16/bin/clang++ \ - ; do - if [[ -x "${candidate}" ]]; then - CXX="${candidate}" - break - fi -done -[[ -n "${CXX}" ]] || { CXX=g++; echo "[build] no CoreX clang++, falling back to g++"; } -echo "[build] CXX=${CXX}" - -# Find CUDA headers (for cuda_runtime_api.h) +# Find CUDA include path CUDA_INC="" for candidate in \ /usr/local/corex/include \ /usr/local/cuda/include \ + /usr/local/corex/lib64/clang/16/include \ ; do if [[ -f "${candidate}/cuda_runtime_api.h" ]]; then CUDA_INC="${candidate}" @@ -45,7 +43,7 @@ for candidate in \ fi done -# Find CUDA libs +# Find CUDA lib path for linking CUDA_LIB="" for candidate in \ /usr/local/corex/lib64 \ @@ -57,43 +55,57 @@ for candidate in \ fi done +if [[ -z "${CUDA_INC}" ]]; then + echo "[WARN] cuda_runtime_api.h not found — trying compile anyway" +fi + echo "[build] CUDA include: ${CUDA_INC:-system}" -echo "[build] CUDA lib: ${CUDA_LIB:-system}" -echo "[build] CCCL include: ${INC} ($(find "${INC}" -type f | wc -l) files)" -echo "[build] Source: ${SRC}" -echo "[build] Output: ${OUT}" - -COMMON_FLAGS=( - -shared -fPIC -O2 -std=c++17 - -I"${INC}" - ${CUDA_INC:+-I"${CUDA_INC}"} - ${CUDA_LIB:+-L"${CUDA_LIB}"} - -lcudart -ldl - # Suppress CCCL warnings that don't affect correctness - -Wno-unused-function - -Wno-unknown-pragmas - # CUB needs these for non-NVCC compilers - -D_CCCL_COMPILER_GCC=1 - -D__CUDA_ARCH_LIST__=700 - -DCUB_DISABLE_NAMESPACE_MAGIC - -DCUB_WRAPPED_NAMESPACE=cccl_preload -) +echo "[build] CUDA lib: ${CUDA_LIB:-system}" +echo "[build] Source: ${SRC}" +echo "[build] Output: ${OUT}" +# Build as shared library +# -x cuda or -x c++ depending on compiler if [[ "${CXX}" == *clang++* ]]; then - "${CXX}" "${COMMON_FLAGS[@]}" -x c++ -o "${OUT}" "${SRC}" 2>&1 + # CoreX clang++ can compile .cu natively + ${CXX} \ + -shared -fPIC \ + -O2 \ + ${CUDA_INC:+-I"${CUDA_INC}"} \ + ${CUDA_LIB:+-L"${CUDA_LIB}"} \ + -lcudart \ + -ldl \ + -std=c++17 \ + -o "${OUT}" \ + "${SRC}" else - "${CXX}" "${COMMON_FLAGS[@]}" -x c++ -o "${OUT}" "${SRC}" 2>&1 + # g++ needs .cu renamed or treated as C++ + # cuda_runtime_api.h should still work with host compiler + ${CXX} \ + -shared -fPIC \ + -O2 \ + ${CUDA_INC:+-I"${CUDA_INC}"} \ + ${CUDA_LIB:+-L"${CUDA_LIB}"} \ + -lcudart \ + -ldl \ + -std=c++17 \ + -x c++ \ + -o "${OUT}" \ + "${SRC}" fi if [[ -f "${OUT}" ]]; then - SIZE=$(stat -c%s "${OUT}" 2>/dev/null || echo "?") - echo "" + SIZE=$(stat -c%s "${OUT}" 2>/dev/null || stat -f%z "${OUT}" 2>/dev/null || echo "?") echo "[build] SUCCESS: ${OUT} (${SIZE} bytes)" echo "" - echo "Test:" - echo " LD_PRELOAD=${OUT} CCCL_ALLOC_DEBUG=1 \\" - echo " PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \\" - echo " python3 verify_preload.py" + echo "Usage:" + echo " LD_PRELOAD=${OUT} CCCL_ALLOC_DEBUG=1 python3 -c 'import torch; t=torch.zeros(1024, device=\"cuda\")'" + echo "" + echo "In computility-run.yaml, add to env:" + echo " - name: LD_PRELOAD" + echo " value: /workspace/qwen3_6_scripts/cccl_preload/libcccl_allocator.so" + echo " - name: PYTORCH_CUDA_ALLOC_CONF" + echo " value: expandable_segments:True" else echo "[build] FAILED" exit 1 diff --git a/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu b/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu index 50061632..2812d2ae 100644 --- a/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu +++ b/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu @@ -1,90 +1,409 @@ /* * cccl_allocator_preload.cu * - * LD_PRELOAD .so — CUB CachingDeviceAllocator from CCCL upstream. - * Full dependency chain (288 files) extracted into include/. + * LD_PRELOAD .so that replaces PyTorch's CUDA memory allocator with + * CUB's CachingDeviceAllocator (extracted from CCCL upstream). * - * Intercepts cudaMalloc/cudaFree, routes through CUB's geometric-bin - * caching allocator. Strips expandable_segments from - * PYTORCH_CUDA_ALLOC_CONF before libtorch reads it. + * Purpose: CoreX's CUDACachingAllocator.cpp:545 asserts + * "expandable segment not supported". Instead of patching libtorch, + * we intercept cudaMalloc/cudaFree at the dynamic linker level and + * route them through CUB's battle-tested caching allocator. * - * Source: CCCL cub/cub/util_allocator.cuh (BSD-3, NVIDIA) - * Build: bash build_cccl_preload.sh + * Source: cccl_upstream/cub/cub/util_allocator.cuh + * License: BSD-3 (NVIDIA/CUB) + * + * Build (on BI-V100 with CoreX clang++): + * bash build_cccl_preload.sh + * + * Usage: + * LD_PRELOAD=/workspace/qwen3_6_scripts/cccl_preload/libcccl_allocator.so \ + * CCCL_ALLOC_DEBUG=0 \ + * PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512 \ + * python3 -m vllm.entrypoints.openai.api_server ... */ -/* ---- CCCL include chain (288 files from cccl_upstream) ---- */ -#include - -/* ---- System ---- */ +#include #include #include #include #include -#include +#include +#include +#include /* ======================================================================== - * Configuration for BI-V100 (32GB × 4 cards) + * CUB CachingDeviceAllocator — extracted from CCCL + * cccl_upstream/cub/cub/util_allocator.cuh * - * 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. + * All CUB/CCCL macro dependencies replaced with plain C++. * ======================================================================== */ -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 bool g_cccl_debug = false; -/* ---- Global allocator singleton ---- */ -static cub::CachingDeviceAllocator& get_allocator() { - static cub::CachingDeviceAllocator instance( - ALLOC_BIN_GROWTH, - ALLOC_MIN_BIN, - ALLOC_MAX_BIN, - ALLOC_MAX_CACHED, - true /* skip_cleanup: CoreX may tear down CUDA before our dtor */ - ); +#define CcclDebug(e) (e) +#define CcclLog(...) \ + do { \ + if (g_cccl_debug) { \ + fprintf(stderr, "[cccl_alloc] "); \ + fprintf(stderr, __VA_ARGS__); \ + } \ + } while (0) + +struct CachingDeviceAllocator +{ + static constexpr unsigned int INVALID_BIN = (unsigned int) -1; + static constexpr size_t INVALID_SIZE = (size_t) -1; + static constexpr int INVALID_DEVICE_ORDINAL = -1; + + struct BlockDescriptor + { + void* d_ptr; + size_t bytes; + unsigned int bin; + int device; + cudaStream_t associated_stream; + cudaEvent_t ready_event; + + BlockDescriptor(void* d_ptr_, int device_) + : d_ptr(d_ptr_), bytes(0), bin(INVALID_BIN), device(device_), + associated_stream(nullptr), ready_event(nullptr) {} + + BlockDescriptor(int device_) + : d_ptr(nullptr), bytes(0), bin(INVALID_BIN), device(device_), + associated_stream(nullptr), ready_event(nullptr) {} + + static bool PtrCompare(const BlockDescriptor& a, const BlockDescriptor& b) { + return (a.device == b.device) ? (a.d_ptr < b.d_ptr) : (a.device < b.device); + } + static bool SizeCompare(const BlockDescriptor& a, const BlockDescriptor& b) { + return (a.device == b.device) ? (a.bytes < b.bytes) : (a.device < b.device); + } + }; + + using Compare = bool (*)(const BlockDescriptor&, const BlockDescriptor&); + + struct TotalBytes { size_t free; size_t live; TotalBytes() : free(0), live(0) {} }; + + using CachedBlocks = std::multiset; + using BusyBlocks = std::multiset; + using GpuCachedBytes = std::map; + + static unsigned int IntPow(unsigned int base, unsigned int exp) { + unsigned int retval = 1; + while (exp > 0) { + if (exp & 1) retval *= base; + base *= base; + exp >>= 1; + } + return retval; + } + + void NearestPowerOf(unsigned int& power, size_t& rounded_bytes, + unsigned int base, size_t value) { + power = 0; + rounded_bytes = 1; + if (value * base < value) { + power = sizeof(size_t) * 8; + rounded_bytes = size_t(0) - 1; + return; + } + while (rounded_bytes < value) { + rounded_bytes *= base; + power++; + } + } + + std::mutex mutex; + unsigned int bin_growth; + unsigned int min_bin; + unsigned int max_bin; + size_t min_bin_bytes; + size_t max_bin_bytes; + size_t max_cached_bytes; + bool skip_cleanup; + GpuCachedBytes cached_bytes; + CachedBlocks cached_blocks; + BusyBlocks live_blocks; + + /* + * Constructor tuned for BI-V100 (32GB per card, 4 cards): + * bin_growth=8, min_bin=3 (512B), max_bin=13 (~550MB) + * max_cached_bytes = 4GB per device (reasonable for 32GB card) + * + * This replaces PyTorch's expandable_segments with a proven + * geometric-bin caching strategy from CUB/CCCL. + */ + CachingDeviceAllocator() + : bin_growth(8) + , min_bin(3) /* 8^3 = 512B minimum allocation */ + , max_bin(13) /* 8^13 = ~550MB maximum cached bin */ + , min_bin_bytes(IntPow(8, 3)) + , max_bin_bytes(IntPow(8, 13)) + , max_cached_bytes((size_t)4 * 1024 * 1024 * 1024) /* 4GB per device */ + , skip_cleanup(true) /* CoreX may tear down CUDA before our dtor */ + , cached_blocks(BlockDescriptor::SizeCompare) + , live_blocks(BlockDescriptor::PtrCompare) + { + CcclLog("CachingDeviceAllocator init: bin_growth=%u min_bin=%u " + "max_bin=%u max_cached=%.1fGB\n", + bin_growth, min_bin, max_bin, + (double)max_cached_bytes / (1024.0*1024.0*1024.0)); + } + + /* ---- Real cudaMalloc/cudaFree via dlsym(RTLD_NEXT) ---- */ + using RealMalloc_t = cudaError_t (*)(void**, size_t); + using RealFree_t = cudaError_t (*)(void*); + + static RealMalloc_t get_real_malloc() { + static RealMalloc_t fn = (RealMalloc_t)dlsym(RTLD_NEXT, "cudaMalloc"); + return fn; + } + static RealFree_t get_real_free() { + static RealFree_t fn = (RealFree_t)dlsym(RTLD_NEXT, "cudaFree"); + return fn; + } + + cudaError_t DeviceAllocate(int device, void** d_ptr, size_t bytes, + cudaStream_t active_stream = nullptr) + { + *d_ptr = nullptr; + int entrypoint_device = INVALID_DEVICE_ORDINAL; + cudaError_t error = cudaSuccess; + + if (device == INVALID_DEVICE_ORDINAL) { + error = cudaGetDevice(&entrypoint_device); + if (error != cudaSuccess) return error; + device = entrypoint_device; + } + + bool found = false; + BlockDescriptor search_key(device); + search_key.associated_stream = active_stream; + NearestPowerOf(search_key.bin, search_key.bytes, bin_growth, bytes); + + if (search_key.bin > max_bin) { + search_key.bin = INVALID_BIN; + search_key.bytes = bytes; + } else { + mutex.lock(); + if (search_key.bin < min_bin) { + search_key.bin = min_bin; + search_key.bytes = min_bin_bytes; + } + + CachedBlocks::iterator block_itr = cached_blocks.lower_bound(search_key); + while ((block_itr != cached_blocks.end()) && + (block_itr->device == device) && + (block_itr->bin == search_key.bin)) + { + bool is_reusable = false; + if (active_stream == block_itr->associated_stream) { + is_reusable = true; + } else { + cudaError_t event_status = cudaEventQuery(block_itr->ready_event); + if (event_status != cudaErrorNotReady) { + is_reusable = true; + } + } + + if (is_reusable) { + found = true; + search_key = *block_itr; + search_key.associated_stream = active_stream; + live_blocks.insert(search_key); + cached_bytes[device].free -= search_key.bytes; + cached_bytes[device].live += search_key.bytes; + + CcclLog("reuse %p (%zu bytes) dev=%d\n", + search_key.d_ptr, search_key.bytes, device); + cached_blocks.erase(block_itr); + break; + } + block_itr++; + } + mutex.unlock(); + } + + if (!found) { + if (device != entrypoint_device) { + if (entrypoint_device == INVALID_DEVICE_ORDINAL) + cudaGetDevice(&entrypoint_device); + cudaSetDevice(device); + } + + /* Use real cudaMalloc, not ourselves */ + error = get_real_malloc()(&search_key.d_ptr, search_key.bytes); + + if (error == cudaErrorMemoryAllocation) { + CcclLog("OOM for %zu bytes on dev=%d, freeing cache...\n", + search_key.bytes, device); + cudaGetLastError(); /* reset */ + + mutex.lock(); + BlockDescriptor free_key(device); + CachedBlocks::iterator block_itr = cached_blocks.lower_bound(free_key); + while ((block_itr != cached_blocks.end()) && + (block_itr->device == device)) + { + error = get_real_free()(block_itr->d_ptr); + if (error != cudaSuccess) break; + cudaEventDestroy(block_itr->ready_event); + cached_bytes[device].free -= block_itr->bytes; + block_itr = cached_blocks.erase(block_itr); + } + mutex.unlock(); + + if (error != cudaSuccess) return error; + error = get_real_malloc()(&search_key.d_ptr, search_key.bytes); + if (error != cudaSuccess) return error; + } else if (error != cudaSuccess) { + return error; + } + + cudaEventCreateWithFlags(&search_key.ready_event, cudaEventDisableTiming); + + mutex.lock(); + live_blocks.insert(search_key); + cached_bytes[device].live += search_key.bytes; + mutex.unlock(); + + CcclLog("alloc %p (%zu bytes, bin=%u) dev=%d\n", + search_key.d_ptr, search_key.bytes, search_key.bin, device); + + if ((entrypoint_device != INVALID_DEVICE_ORDINAL) && + (entrypoint_device != device)) + cudaSetDevice(entrypoint_device); + } + + *d_ptr = search_key.d_ptr; + return cudaSuccess; + } + + cudaError_t DeviceAllocate(void** d_ptr, size_t bytes, + cudaStream_t active_stream = nullptr) { + return DeviceAllocate(INVALID_DEVICE_ORDINAL, d_ptr, bytes, active_stream); + } + + cudaError_t DeviceFree(int device, void* d_ptr) + { + int entrypoint_device = INVALID_DEVICE_ORDINAL; + cudaError_t error = cudaSuccess; + + if (d_ptr == nullptr) return cudaSuccess; + + if (device == INVALID_DEVICE_ORDINAL) { + error = cudaGetDevice(&entrypoint_device); + if (error != cudaSuccess) return error; + device = entrypoint_device; + } + + mutex.lock(); + bool recached = false; + BlockDescriptor search_key(d_ptr, device); + BusyBlocks::iterator block_itr = live_blocks.find(search_key); + + if (block_itr != live_blocks.end()) { + search_key = *block_itr; + live_blocks.erase(block_itr); + cached_bytes[device].live -= search_key.bytes; + + if ((search_key.bin != INVALID_BIN) && + (cached_bytes[device].free + search_key.bytes <= max_cached_bytes)) + { + recached = true; + cached_blocks.insert(search_key); + cached_bytes[device].free += search_key.bytes; + CcclLog("cache %p (%zu bytes) dev=%d\n", + d_ptr, search_key.bytes, device); + } + } + mutex.unlock(); + + if (device != entrypoint_device) { + if (entrypoint_device == INVALID_DEVICE_ORDINAL) + cudaGetDevice(&entrypoint_device); + cudaSetDevice(device); + } + + if (recached) { + cudaEventRecord(search_key.ready_event, search_key.associated_stream); + } else { + /* Not tracked or cache full — real free */ + CcclLog("free %p dev=%d (not cached)\n", d_ptr, device); + error = get_real_free()(d_ptr); + if (block_itr != live_blocks.end()) + cudaEventDestroy(search_key.ready_event); + } + + if ((entrypoint_device != INVALID_DEVICE_ORDINAL) && + (entrypoint_device != device)) + cudaSetDevice(entrypoint_device); + + return error; + } + + cudaError_t DeviceFree(void* d_ptr) { + return DeviceFree(INVALID_DEVICE_ORDINAL, d_ptr); + } + + cudaError_t FreeAllCached() + { + cudaError_t error = cudaSuccess; + int entrypoint_device = INVALID_DEVICE_ORDINAL; + int current_device = INVALID_DEVICE_ORDINAL; + + mutex.lock(); + while (!cached_blocks.empty()) { + CachedBlocks::iterator begin = cached_blocks.begin(); + if (entrypoint_device == INVALID_DEVICE_ORDINAL) + cudaGetDevice(&entrypoint_device); + if (begin->device != current_device) { + cudaSetDevice(begin->device); + current_device = begin->device; + } + get_real_free()(begin->d_ptr); + cudaEventDestroy(begin->ready_event); + cached_bytes[current_device].free -= begin->bytes; + cached_blocks.erase(begin); + } + mutex.unlock(); + + if (entrypoint_device != INVALID_DEVICE_ORDINAL) + cudaSetDevice(entrypoint_device); + return error; + } +}; + +/* ======================================================================== + * Global singleton + LD_PRELOAD intercepts + * ======================================================================== */ + +static CachingDeviceAllocator& get_allocator() { + static CachingDeviceAllocator instance; return instance; } static bool g_preload_active = false; -static bool g_debug = false; -/* ---- Real cudaMalloc/cudaFree via dlsym(RTLD_NEXT) ---- */ -using RealMalloc_t = cudaError_t (*)(void**, size_t); -using RealFree_t = cudaError_t (*)(void*); - -static RealMalloc_t get_real_malloc() { - static RealMalloc_t fn = (RealMalloc_t)dlsym(RTLD_NEXT, "cudaMalloc"); - return fn; -} -static RealFree_t get_real_free() { - static RealFree_t fn = (RealFree_t)dlsym(RTLD_NEXT, "cudaFree"); - return fn; -} - -/* ======================================================================== - * Constructor: runs at LD_PRELOAD load time - * ======================================================================== */ +/* Called once at .so load time */ __attribute__((constructor)) static void cccl_preload_init() { const char* debug_env = getenv("CCCL_ALLOC_DEBUG"); - g_debug = (debug_env && atoi(debug_env) > 0); + g_cccl_debug = (debug_env && atoi(debug_env) > 0); const char* disable_env = getenv("CCCL_ALLOC_DISABLE"); if (disable_env && atoi(disable_env) > 0) { fprintf(stderr, "[cccl_alloc] DISABLED by CCCL_ALLOC_DISABLE=1\n"); + g_preload_active = false; return; } - /* Strip expandable_segments from PYTORCH_CUDA_ALLOC_CONF */ + /* Strip expandable_segments from PYTORCH_CUDA_ALLOC_CONF + * so CoreX's allocator doesn't hit the assert. + * We handle the caching ourselves. */ const char* alloc_conf = getenv("PYTORCH_CUDA_ALLOC_CONF"); if (alloc_conf) { + /* Build a new conf string without expandable_segments */ std::string conf(alloc_conf); std::string clean; size_t pos = 0; @@ -92,51 +411,49 @@ static void cccl_preload_init() { size_t comma = conf.find(',', pos); if (comma == std::string::npos) comma = conf.size(); std::string token = conf.substr(pos, comma - pos); + /* Skip expandable_segments:* */ if (token.find("expandable_segments") == std::string::npos) { if (!clean.empty()) clean += ","; clean += token; } pos = comma + 1; } - if (clean.empty()) + if (clean.empty()) { unsetenv("PYTORCH_CUDA_ALLOC_CONF"); - else + } else { setenv("PYTORCH_CUDA_ALLOC_CONF", clean.c_str(), 1); - - fprintf(stderr, "[cccl_alloc] PYTORCH_CUDA_ALLOC_CONF: \"%s\" -> \"%s\"\n", + } + fprintf(stderr, "[cccl_alloc] stripped expandable_segments from " + "PYTORCH_CUDA_ALLOC_CONF: \"%s\" -> \"%s\"\n", alloc_conf, clean.empty() ? "(unset)" : clean.c_str()); } - /* Initialize allocator */ - auto& alloc = get_allocator(); - if (g_debug) { - alloc.debug = true; - } - + /* Force-initialize the allocator singleton */ + (void)get_allocator(); g_preload_active = true; - fprintf(stderr, - "[cccl_alloc] LD_PRELOAD active — CUB CachingDeviceAllocator " - "(growth=%u, bins=[%u..%u], max_cached=%.1fGB)\n", - ALLOC_BIN_GROWTH, ALLOC_MIN_BIN, ALLOC_MAX_BIN, - (double)ALLOC_MAX_CACHED / (1024.0*1024.0*1024.0)); + fprintf(stderr, "[cccl_alloc] LD_PRELOAD active — CUB CachingDeviceAllocator " + "replacing cudaMalloc/cudaFree\n"); } -/* ======================================================================== - * cudaMalloc / cudaFree intercepts - * ======================================================================== */ - +/* ---- cudaMalloc intercept ---- */ extern "C" cudaError_t cudaMalloc(void** devPtr, size_t size) { if (!g_preload_active) { - return get_real_malloc()(devPtr, size); + /* Fallback to real cudaMalloc during init or if disabled */ + static auto real_fn = (CachingDeviceAllocator::RealMalloc_t) + dlsym(RTLD_NEXT, "cudaMalloc"); + return real_fn(devPtr, size); } return get_allocator().DeviceAllocate(devPtr, size); } +/* ---- cudaFree intercept ---- */ extern "C" cudaError_t cudaFree(void* devPtr) { if (!g_preload_active || devPtr == nullptr) { - return get_real_free()(devPtr); + static auto real_fn = (CachingDeviceAllocator::RealFree_t) + dlsym(RTLD_NEXT, "cudaFree"); + return real_fn(devPtr); } return get_allocator().DeviceFree(devPtr); } diff --git a/qwen3_6_scripts/cccl_preload/verify_preload.py b/qwen3_6_scripts/cccl_preload/verify_preload.py index 270447b1..752d10f4 100644 --- a/qwen3_6_scripts/cccl_preload/verify_preload.py +++ b/qwen3_6_scripts/cccl_preload/verify_preload.py @@ -28,7 +28,7 @@ if not torch.cuda.is_available(): device = torch.device("cuda:0") print(f"Device: {torch.cuda.get_device_name(0)}") -print(f"Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB") +print(f"Memory: {torch.cuda.get_device_properties(0).total_mem / 1024**3:.1f} GB") print() # Test 1: Basic allocation