feat: CCCL CachingDeviceAllocator preload — 完整依赖链 288 files

从 cccl_upstream 递归追踪 cub/util_allocator.cuh 的全部 include 依赖:
  cub/         9 files (config, util_*, version, detect_cuda_runtime)
  cuda/        libcudacxx type_traits, concepts, algorithm, iterator...
  nv/          target macros, preprocessor

总计 288 个头文件 (1.4MB)，打包到 include/ 目录，编译时 -I include
即可完全脱离 CCCL 原始目录结构。

.cu 文件直接 #include <cub/util_allocator.cuh>，
走原版 CUB CachingDeviceAllocator，零 mock。

BI-V100 参数: growth=2 bins=[8..32] max_cached=8GB/device

qwen3_6_scripts/cccl_preload/build_cccl_preload.sh

@@ -1,41 +1,43 @@
 #!/usr/bin/env bash
-# Build libcccl_allocator.so — LD_PRELOAD .so for CUB CachingDeviceAllocator
+# Build libcccl_allocator.so
+#
+# Full CCCL dependency chain (288 headers) in ./include/
+# Source: cccl_upstream/cub/cub/util_allocator.cuh + transitive deps
 #
 # 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"
 
-# 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
+[[ -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 CUDA include path
+# 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)
 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}"
@@ -43,7 +45,7 @@ for candidate in \
     fi
 done
 
-# Find CUDA lib path for linking
+# Find CUDA libs
 CUDA_LIB=""
 for candidate in \
     /usr/local/corex/lib64 \
@@ -55,57 +57,43 @@ 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] Source: ${SRC}"
-echo "[build] Output: ${OUT}"
+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
+)
 
-# Build as shared library
-# -x cuda or -x c++ depending on compiler
 if [[ "${CXX}" == *clang++* ]]; then
-    # 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}"
+    "${CXX}" "${COMMON_FLAGS[@]}" -x c++ -o "${OUT}" "${SRC}" 2>&1
 else
-    # 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}"
+    "${CXX}" "${COMMON_FLAGS[@]}" -x c++ -o "${OUT}" "${SRC}" 2>&1
 fi
 
 if [[ -f "${OUT}" ]]; then
-    SIZE=$(stat -c%s "${OUT}" 2>/dev/null || stat -f%z "${OUT}" 2>/dev/null || echo "?")
+    SIZE=$(stat -c%s "${OUT}" 2>/dev/null || echo "?")
+    echo ""
     echo "[build] SUCCESS: ${OUT} (${SIZE} bytes)"
     echo ""
-    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"
+    echo "Test:"
+    echo "  LD_PRELOAD=${OUT} CCCL_ALLOC_DEBUG=1 \\"
+    echo "  PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \\"
+    echo "  python3 verify_preload.py"
 else
     echo "[build] FAILED"
     exit 1

qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu

@@ -1,409 +1,90 @@
 /*
  * cccl_allocator_preload.cu
  *
- * LD_PRELOAD .so that replaces PyTorch's CUDA memory allocator with
- * CUB's CachingDeviceAllocator (extracted from CCCL upstream).
+ * LD_PRELOAD .so — CUB CachingDeviceAllocator from CCCL upstream.
+ * Full dependency chain (288 files) extracted into include/.
  *
- * 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.
+ * Intercepts cudaMalloc/cudaFree, routes through CUB's geometric-bin
+ * caching allocator. Strips expandable_segments from
+ * PYTORCH_CUDA_ALLOC_CONF before libtorch reads it.
  *
- * 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 ...
+ * Source: CCCL cub/cub/util_allocator.cuh (BSD-3, NVIDIA)
+ * Build:  bash build_cccl_preload.sh
  */
 
-#include <cuda_runtime_api.h>
+/* ---- CCCL include chain (288 files from cccl_upstream) ---- */
+#include <cub/util_allocator.cuh>
+
+/* ---- System ---- */
 #include <dlfcn.h>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
-#include <map>
-#include <mutex>
-#include <set>
+#include <string>
 
 /* ========================================================================
- * CUB CachingDeviceAllocator — extracted from CCCL
- * cccl_upstream/cub/cub/util_allocator.cuh
+ * Configuration for BI-V100 (32GB × 4 cards)
  *
- * All CUB/CCCL macro dependencies replaced with plain C++.
+ * 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 bool g_cccl_debug = false;
+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 */
 
-#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<BlockDescriptor, Compare>;
-    using BusyBlocks = std::multiset<BlockDescriptor, Compare>;
-    using GpuCachedBytes = std::map<int, TotalBytes>;
-
-    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;
+/* ---- 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 */
+    );
     return instance;
 }
 
 static bool g_preload_active = false;
+static bool g_debug = false;
 
-/* Called once at .so load time */
+/* ---- 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
+ * ======================================================================== */
 __attribute__((constructor))
 static void cccl_preload_init() {
     const char* debug_env = getenv("CCCL_ALLOC_DEBUG");
-    g_cccl_debug = (debug_env && atoi(debug_env) > 0);
+    g_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
-     * so CoreX's allocator doesn't hit the assert.
-     * We handle the caching ourselves. */
+    /* Strip expandable_segments from PYTORCH_CUDA_ALLOC_CONF */
     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;
@@ -411,49 +92,51 @@ 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] stripped expandable_segments from "
-                "PYTORCH_CUDA_ALLOC_CONF: \"%s\" -> \"%s\"\n",
+
+        fprintf(stderr, "[cccl_alloc] PYTORCH_CUDA_ALLOC_CONF: \"%s\" -> \"%s\"\n",
                 alloc_conf, clean.empty() ? "(unset)" : clean.c_str());
     }
 
-    /* Force-initialize the allocator singleton */
-    (void)get_allocator();
+    /* Initialize allocator */
+    auto& alloc = get_allocator();
+    if (g_debug) {
+        alloc.debug = true;
+    }
+
     g_preload_active = true;
-    fprintf(stderr, "[cccl_alloc] LD_PRELOAD active — CUB CachingDeviceAllocator "
-            "replacing cudaMalloc/cudaFree\n");
+    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));
 }
 
-/* ---- cudaMalloc intercept ---- */
+/* ========================================================================
+ * cudaMalloc / cudaFree intercepts
+ * ======================================================================== */
+
 extern "C" cudaError_t cudaMalloc(void** devPtr, size_t size)
 {
     if (!g_preload_active) {
-        /* 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_real_malloc()(devPtr, size);
     }
     return get_allocator().DeviceAllocate(devPtr, size);
 }
 
-/* ---- cudaFree intercept ---- */
 extern "C" cudaError_t cudaFree(void* devPtr)
 {
     if (!g_preload_active || devPtr == nullptr) {
-        static auto real_fn = (CachingDeviceAllocator::RealFree_t)
-            dlsym(RTLD_NEXT, "cudaFree");
-        return real_fn(devPtr);
+        return get_real_free()(devPtr);
     }
     return get_allocator().DeviceFree(devPtr);
 }