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:
dylanyunlon
2026-08-13 09:21:57 +00:00
committed by root
parent 887e0981ad
commit 8dc6462a2b
3 changed files with 456 additions and 127 deletions

View File

@@ -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