feat: CCCL CachingDeviceAllocator LD_PRELOAD — bypass CoreX expandable_segments ASSERT
从 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
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user