feat(CCCL): LD_PRELOAD CachingDeviceAllocator — intercept cudaMalloc/cudaFree

Route C: replace PyTorch's cudaMalloc/cudaFree with CCCL CUB's
CachingDeviceAllocator via LD_PRELOAD. Eliminates driver-level allocation
overhead by reusing freed GPU memory from a bin-based cache.

Based on cccl_upstream/cub/cub/util_allocator.cuh (901 lines).
Self-contained .so with no CCCL header dependencies at compile time.

Files:
- cccl_preload_allocator.cu: the allocator (405 lines)
- build_cccl_preload_allocator.sh: build script (corex clang++ or g++ fallback)
- test_cccl_preload.sh: smoke test suite for BI-V100
- patch_ops.sh: build during docker build
- computility-run.yaml: LD_PRELOAD env var for runtime

Config via env:
  CCCL_ALLOC_BIN_GROWTH=8, MIN_BIN=3, MAX_BIN=13, MAX_CACHED_MB=4096

Test on real machine:
  cd qwen3_6_scripts && bash test_cccl_preload.sh

computility-run.yaml

@@ -49,3 +49,13 @@ env:
       value: '1'
     - name: PYTORCH_CUDA_ALLOC_CONF
       value: expandable_segments:True
+    - name: LD_PRELOAD
+      value: /workspace/qwen3_6_scripts/cccl_preload_allocator.so
+    - name: CCCL_ALLOC_BIN_GROWTH
+      value: '8'
+    - name: CCCL_ALLOC_MIN_BIN
+      value: '3'
+    - name: CCCL_ALLOC_MAX_BIN
+      value: '13'
+    - name: CCCL_ALLOC_MAX_CACHED_MB
+      value: '4096'

qwen3_6_scripts/build_cccl_preload_allocator.sh

@@ -0,0 +1,99 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Build the CCCL CachingDeviceAllocator LD_PRELOAD .so
+#
+# Usage: bash build_cccl_preload_allocator.sh [output_dir]
+# Default output: ./cccl_preload_allocator.so
+#
+# Test:  LD_PRELOAD=./cccl_preload_allocator.so python3 -c "import torch; x=torch.zeros(1024,device='cuda'); del x; print('OK')"
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT_DIR=${1:-${SCRIPT_DIR}}
+OUTPUT=${OUTPUT_DIR}/cccl_preload_allocator.so
+SOURCE=${SCRIPT_DIR}/cccl_preload_allocator.cu
+
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+
+# Check if we have the corex compiler
+if [[ -x "${COREX_ROOT}/bin/clang++" ]]; then
+    COMPILER="${COREX_ROOT}/bin/clang++"
+    echo "[build] Using corex clang++: ${COMPILER}"
+
+    "${COMPILER}" \
+        -std=c++17 -O3 -shared -fPIC \
+        --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+        --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+        -I"${COREX_ROOT}/include" \
+        "${SOURCE}" \
+        -L"${COREX_ROOT}/lib64" -lcudart -ldl \
+        -Wl,-rpath,"${COREX_ROOT}/lib64" \
+        -o "${OUTPUT}"
+else
+    # Fallback: try to compile as pure C++ (no CUDA kernels needed)
+    # The allocator is entirely host-side code
+    echo "[build] corex clang++ not found, trying system g++"
+    echo "[build] Note: this is HOST-only code, no GPU kernels involved"
+
+    # Find cuda include path
+    CUDA_INC=""
+    for p in /usr/local/corex-3.2.3/include /usr/local/cuda/include /usr/local/corex/include; do
+        if [[ -d "$p" ]]; then CUDA_INC="$p"; break; fi
+    done
+
+    CUDA_LIB=""
+    for p in /usr/local/corex-3.2.3/lib64 /usr/local/cuda/lib64 /usr/local/corex/lib64; do
+        if [[ -d "$p" ]]; then CUDA_LIB="$p"; break; fi
+    done
+
+    if [[ -z "$CUDA_INC" || -z "$CUDA_LIB" ]]; then
+        echo "[build] ERROR: Cannot find CUDA headers/libs" >&2
+        exit 2
+    fi
+
+    # Rename .cu → .cpp for g++ (it's all host code anyway)
+    TMP_CPP=$(mktemp /tmp/cccl_alloc_XXXXXX.cpp)
+    cp "${SOURCE}" "${TMP_CPP}"
+
+    g++ -std=c++17 -O3 -shared -fPIC \
+        -D_GLIBCXX_USE_CXX11_ABI=0 \
+        -I"${CUDA_INC}" \
+        "${TMP_CPP}" \
+        -L"${CUDA_LIB}" -lcudart -ldl \
+        -Wl,-rpath,"${CUDA_LIB}" \
+        -o "${OUTPUT}"
+
+    rm -f "${TMP_CPP}"
+fi
+
+# Verify
+if [[ ! -s "${OUTPUT}" ]]; then
+    echo "[build] ERROR: output is empty" >&2
+    exit 2
+fi
+
+# Check it's a proper shared library with our symbols
+if command -v nm &>/dev/null; then
+    HAS_MALLOC=$(nm -D "${OUTPUT}" 2>/dev/null | grep -c "T cudaMalloc" || true)
+    HAS_FREE=$(nm -D "${OUTPUT}" 2>/dev/null | grep -c "T cudaFree" || true)
+    if [[ "$HAS_MALLOC" -gt 0 && "$HAS_FREE" -gt 0 ]]; then
+        echo "[build] OK: cudaMalloc and cudaFree symbols exported"
+    else
+        echo "[build] WARNING: symbol check inconclusive (nm output may differ)"
+    fi
+fi
+
+echo "[build] Built: ${OUTPUT} ($(stat -c%s "${OUTPUT}" 2>/dev/null || stat -f%z "${OUTPUT}") bytes)"
+echo ""
+echo "Test command:"
+echo "  LD_PRELOAD=${OUTPUT} python3 -c \"import torch; x=torch.zeros(1024,device='cuda'); del x; print('OK')\""
+echo ""
+echo "Production usage (add to computility-run.yaml or Dockerfile CMD):"
+echo "  LD_PRELOAD=${OUTPUT} python3 -m vllm.entrypoints.openai.api_server ..."
+echo ""
+echo "Tuning env vars:"
+echo "  CCCL_ALLOC_BIN_GROWTH=8       # Geometric growth factor"
+echo "  CCCL_ALLOC_MIN_BIN=3          # Min bin (growth^3 = 512B)"
+echo "  CCCL_ALLOC_MAX_BIN=13         # Max bin (8^13 = ~550MB)"
+echo "  CCCL_ALLOC_MAX_CACHED_MB=4096 # Max 4GB cached per device"
+echo "  CCCL_ALLOC_DEBUG=1            # Print every alloc/free"

qwen3_6_scripts/cccl_preload_allocator.cu

@@ -0,0 +1,405 @@
+// cccl_preload_allocator.cu — LD_PRELOAD interception of cudaMalloc/cudaFree
+//
+// Replaces the default cudaMalloc/cudaFree with CCCL CUB CachingDeviceAllocator.
+// This eliminates the CUDA driver's allocation overhead (cudaMalloc is slow on
+// BI-V100: ~2-50ms per call) by reusing freed blocks from a bin-based cache.
+//
+// Build on BI-V100:
+//   /usr/local/corex-3.2.3/bin/clang++ -std=c++17 -O3 -shared -fPIC \
+//     --cuda-path=/usr/local/corex-3.2.3 --cuda-gpu-arch=ivcore10 \
+//     --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+//     -I/usr/local/corex-3.2.3/include \
+//     cccl_preload_allocator.cu \
+//     -L/usr/local/corex-3.2.3/lib64 -lcudart -ldl \
+//     -o cccl_preload_allocator.so
+//
+// Usage:
+//   LD_PRELOAD=/workspace/qwen3_6_scripts/cccl_preload_allocator.so python3 -m vllm.entrypoints.openai.api_server ...
+//
+// Tuning (env vars):
+//   CCCL_ALLOC_BIN_GROWTH=8      Geometric growth factor (default 8)
+//   CCCL_ALLOC_MIN_BIN=3         Min bin exponent (default 3 → 512B)
+//   CCCL_ALLOC_MAX_BIN=13        Max bin exponent (default 13 → 512MB for growth=8; was 7→2MB)
+//   CCCL_ALLOC_MAX_CACHED_MB=4096  Max cached bytes per device in MB (default 4096=4GB)
+//   CCCL_ALLOC_DEBUG=0           Print alloc/free events (default 0)
+//
+// Design notes:
+// - Only intercepts cudaMalloc and cudaFree (the synchronous variants).
+// - cudaMallocAsync/cudaFreeAsync are NOT intercepted (PyTorch on BI-V100
+//   doesn't use them; the corex runtime may not support them).
+// - Thread-safe via CUB's internal mutex.
+// - Stream association: all allocations use the default stream (nullptr).
+//   PyTorch's CUDACachingAllocator handles stream ordering itself, so we
+//   don't need to track streams here.
+// - Large allocations (> max_bin_bytes) pass through to real cudaMalloc.
+// - The allocator is process-global (static singleton).
+
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <dlfcn.h>
+#include <mutex>
+#include <map>
+#include <set>
+#include <atomic>
+
+// We inline the essential logic from CUB CachingDeviceAllocator rather than
+// #include it, because the corex toolchain may not have full CCCL headers
+// installed, and we need to link against corex's cudart, not NVIDIA's.
+
+// Forward declare the real CUDA functions we'll dlsym
+typedef int cudaError_t;
+static constexpr cudaError_t cudaSuccess = 0;
+
+typedef void* cudaStream_t;
+typedef void* cudaEvent_t;
+
+// Real function pointers (resolved via dlsym on first call)
+using cudaMalloc_fn = cudaError_t(*)(void**, size_t);
+using cudaFree_fn = cudaError_t(*)(void*);
+using cudaGetDevice_fn = cudaError_t(*)(int*);
+using cudaEventCreate_fn = cudaError_t(*)(cudaEvent_t*);
+using cudaEventRecord_fn = cudaError_t(*)(cudaEvent_t, cudaStream_t);
+using cudaEventQuery_fn = cudaError_t(*)(cudaEvent_t);
+using cudaEventDestroy_fn = cudaError_t(*)(cudaEvent_t);
+using cudaEventSynchronize_fn = cudaError_t(*)(cudaEvent_t);
+
+static cudaMalloc_fn  real_cudaMalloc  = nullptr;
+static cudaFree_fn    real_cudaFree    = nullptr;
+static cudaGetDevice_fn real_cudaGetDevice = nullptr;
+static cudaEventCreate_fn real_cudaEventCreate = nullptr;
+static cudaEventRecord_fn real_cudaEventRecord = nullptr;
+static cudaEventQuery_fn  real_cudaEventQuery  = nullptr;
+static cudaEventDestroy_fn real_cudaEventDestroy = nullptr;
+static cudaEventSynchronize_fn real_cudaEventSynchronize = nullptr;
+
+static std::once_flag resolve_flag;
+
+static void resolve_real_functions() {
+    real_cudaMalloc = (cudaMalloc_fn)dlsym(RTLD_NEXT, "cudaMalloc");
+    real_cudaFree = (cudaFree_fn)dlsym(RTLD_NEXT, "cudaFree");
+    real_cudaGetDevice = (cudaGetDevice_fn)dlsym(RTLD_NEXT, "cudaGetDevice");
+    real_cudaEventCreate = (cudaEventCreate_fn)dlsym(RTLD_NEXT, "cudaEventCreate");
+    real_cudaEventRecord = (cudaEventRecord_fn)dlsym(RTLD_NEXT, "cudaEventRecord");
+    real_cudaEventQuery = (cudaEventQuery_fn)dlsym(RTLD_NEXT, "cudaEventQuery");
+    real_cudaEventDestroy = (cudaEventDestroy_fn)dlsym(RTLD_NEXT, "cudaEventDestroy");
+    real_cudaEventSynchronize = (cudaEventSynchronize_fn)dlsym(RTLD_NEXT, "cudaEventSynchronize");
+
+    if (!real_cudaMalloc || !real_cudaFree) {
+        fprintf(stderr, "[CCCL_PRELOAD] FATAL: cannot resolve cudaMalloc/cudaFree via dlsym\n");
+        abort();
+    }
+}
+
+// ============================================================================
+// Simplified CachingDeviceAllocator (from CCCL cub/util_allocator.cuh)
+// Stripped to essentials: no debug logging macros, no CCCL config dependencies
+// ============================================================================
+
+struct BlockDescriptor {
+    void*         d_ptr;
+    size_t        bytes;
+    unsigned int  bin;
+    int           device;
+    cudaStream_t  associated_stream;
+    cudaEvent_t   ready_event;
+
+    BlockDescriptor(void* p, int dev)
+        : d_ptr(p), bytes(0), bin(~0u), device(dev),
+          associated_stream(nullptr), ready_event(nullptr) {}
+
+    BlockDescriptor(int dev)
+        : d_ptr(nullptr), bytes(0), bin(~0u), device(dev),
+          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&);
+using CachedBlocks = std::multiset<BlockDescriptor, Compare>;
+using BusyBlocks = std::multiset<BlockDescriptor, Compare>;
+
+struct DeviceBytes { size_t free = 0; size_t live = 0; };
+
+struct CachingAllocator {
+    std::mutex     mtx;
+    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           debug;
+    CachedBlocks   cached_blocks;
+    BusyBlocks     live_blocks;
+    std::map<int, DeviceBytes> cached_bytes;
+
+    // Stats
+    std::atomic<uint64_t> stat_hits{0};
+    std::atomic<uint64_t> stat_misses{0};
+    std::atomic<uint64_t> stat_frees{0};
+    std::atomic<uint64_t> stat_bypasses{0};
+
+    static unsigned int IntPow(unsigned int base, unsigned int exp) {
+        unsigned int r = 1;
+        while (exp > 0) {
+            if (exp & 1) r *= base;
+            base *= base;
+            exp >>= 1;
+        }
+        return r;
+    }
+
+    void NearestPowerOf(unsigned int& power, size_t& rounded,
+                        unsigned int base, size_t value) {
+        power = 0; rounded = 1;
+        if (value * base < value) {
+            power = sizeof(size_t) * 8;
+            rounded = size_t(-1);
+            return;
+        }
+        while (rounded < value) { rounded *= base; power++; }
+    }
+
+    CachingAllocator()
+        : cached_blocks(BlockDescriptor::SizeCompare),
+          live_blocks(BlockDescriptor::PtrCompare) {
+        // Read config from env
+        auto env_or = [](const char* name, int def) -> int {
+            const char* v = getenv(name);
+            return v ? atoi(v) : def;
+        };
+
+        bin_growth      = env_or("CCCL_ALLOC_BIN_GROWTH", 8);
+        min_bin         = env_or("CCCL_ALLOC_MIN_BIN", 3);
+        max_bin         = env_or("CCCL_ALLOC_MAX_BIN", 13);
+        int max_mb      = env_or("CCCL_ALLOC_MAX_CACHED_MB", 4096);
+        debug           = env_or("CCCL_ALLOC_DEBUG", 0) != 0;
+        min_bin_bytes   = IntPow(bin_growth, min_bin);
+        max_bin_bytes   = IntPow(bin_growth, max_bin);
+        max_cached_bytes = (size_t)max_mb * 1024ULL * 1024ULL;
+
+        fprintf(stderr, "[CCCL_PRELOAD] CachingDeviceAllocator: growth=%u "
+                "bins=[%u..%u] bin_bytes=[%zu..%zu] max_cached=%zuMB\n",
+                bin_growth, min_bin, max_bin,
+                min_bin_bytes, max_bin_bytes, max_cached_bytes / (1024*1024));
+    }
+
+    ~CachingAllocator() {
+        fprintf(stderr, "[CCCL_PRELOAD] Stats: hits=%lu misses=%lu frees=%lu bypasses=%lu\n",
+                stat_hits.load(), stat_misses.load(),
+                stat_frees.load(), stat_bypasses.load());
+        // Free all cached blocks
+        for (auto& b : cached_blocks) {
+            if (b.ready_event) real_cudaEventDestroy(b.ready_event);
+            real_cudaFree(b.d_ptr);
+        }
+    }
+
+    cudaError_t Allocate(void** d_ptr, size_t bytes) {
+        std::call_once(resolve_flag, resolve_real_functions);
+        *d_ptr = nullptr;
+
+        // Get current device
+        int device = 0;
+        if (real_cudaGetDevice) real_cudaGetDevice(&device);
+
+        // Bin classification
+        unsigned int bin;
+        size_t rounded_bytes;
+        bool oversized = false;
+
+        if (bytes > max_bin_bytes) {
+            // Too large for caching — pass through
+            bin = max_bin + 1;
+            rounded_bytes = bytes;
+            oversized = true;
+        } else {
+            NearestPowerOf(bin, rounded_bytes, bin_growth, bytes);
+            if (bin < min_bin) {
+                bin = min_bin;
+                rounded_bytes = min_bin_bytes;
+            }
+        }
+
+        BlockDescriptor search_key(device);
+        search_key.bytes = rounded_bytes;
+        search_key.bin = bin;
+
+        // Lock
+        std::lock_guard<std::mutex> lock(mtx);
+
+        if (!oversized) {
+            // Search cached blocks for a match
+            auto range = cached_blocks.equal_range(search_key);
+            for (auto it = range.first; it != range.second; ++it) {
+                if (it->device == device && it->bin == bin) {
+                    // Check if the stream work has completed
+                    bool ready = true;
+                    if (it->ready_event) {
+                        cudaError_t ev_status = real_cudaEventQuery(it->ready_event);
+                        if (ev_status != cudaSuccess) {
+                            // Event not ready — try to synchronize briefly
+                            // For BI-V100 with enforce_eager, events should be ready
+                            real_cudaEventSynchronize(it->ready_event);
+                        }
+                        real_cudaEventDestroy(it->ready_event);
+                    }
+
+                    // Reuse this block
+                    search_key.d_ptr = it->d_ptr;
+                    search_key.bytes = it->bytes;
+                    live_blocks.insert(search_key);
+                    cached_bytes[device].free -= it->bytes;
+                    cached_bytes[device].live += it->bytes;
+                    cached_blocks.erase(it);
+
+                    *d_ptr = search_key.d_ptr;
+                    stat_hits++;
+
+                    if (debug) {
+                        fprintf(stderr, "[CCCL_PRELOAD] HIT  dev=%d bin=%u "
+                                "req=%zu alloc=%zu ptr=%p\n",
+                                device, bin, bytes, search_key.bytes, *d_ptr);
+                    }
+                    return cudaSuccess;
+                }
+            }
+        }
+
+        // Cache miss — allocate new block
+        cudaError_t err = real_cudaMalloc(&search_key.d_ptr, rounded_bytes);
+
+        // If OOM, try evicting cached blocks and retry
+        if (err != cudaSuccess) {
+            // Free all cached blocks on this device
+            auto it = cached_blocks.begin();
+            while (it != cached_blocks.end()) {
+                if (it->device == device) {
+                    if (it->ready_event) {
+                        real_cudaEventSynchronize(it->ready_event);
+                        real_cudaEventDestroy(it->ready_event);
+                    }
+                    real_cudaFree(it->d_ptr);
+                    cached_bytes[device].free -= it->bytes;
+                    it = cached_blocks.erase(it);
+                } else {
+                    ++it;
+                }
+            }
+            // Retry
+            err = real_cudaMalloc(&search_key.d_ptr, rounded_bytes);
+        }
+
+        if (err != cudaSuccess) {
+            return err;
+        }
+
+        search_key.bytes = rounded_bytes;
+        live_blocks.insert(search_key);
+        cached_bytes[device].live += rounded_bytes;
+
+        *d_ptr = search_key.d_ptr;
+
+        if (oversized) {
+            stat_bypasses++;
+        } else {
+            stat_misses++;
+        }
+
+        if (debug) {
+            fprintf(stderr, "[CCCL_PRELOAD] %s dev=%d bin=%u "
+                    "req=%zu alloc=%zu ptr=%p\n",
+                    oversized ? "PASS" : "MISS",
+                    device, bin, bytes, rounded_bytes, *d_ptr);
+        }
+        return cudaSuccess;
+    }
+
+    cudaError_t Free(void* d_ptr) {
+        std::call_once(resolve_flag, resolve_real_functions);
+
+        if (d_ptr == nullptr) return cudaSuccess;
+
+        int device = 0;
+        if (real_cudaGetDevice) real_cudaGetDevice(&device);
+
+        BlockDescriptor search_key(d_ptr, device);
+
+        std::lock_guard<std::mutex> lock(mtx);
+
+        auto it = live_blocks.find(search_key);
+        if (it == live_blocks.end()) {
+            // Not tracked by us — pass through to real cudaFree
+            return real_cudaFree(d_ptr);
+        }
+
+        search_key.bytes = it->bytes;
+        search_key.bin = it->bin;
+        cached_bytes[device].live -= it->bytes;
+        live_blocks.erase(it);
+
+        stat_frees++;
+
+        // Check if this block is too large or would exceed cache limit
+        bool should_cache = (search_key.bin <= max_bin) &&
+                            (cached_bytes[device].free + search_key.bytes <= max_cached_bytes);
+
+        if (should_cache) {
+            // Record an event so we know when it's safe to reuse
+            if (real_cudaEventCreate) {
+                cudaEvent_t event = nullptr;
+                cudaError_t ev_err = real_cudaEventCreate(&event);
+                if (ev_err == cudaSuccess && real_cudaEventRecord) {
+                    real_cudaEventRecord(event, nullptr);  // default stream
+                    search_key.ready_event = event;
+                }
+            }
+
+            cached_blocks.insert(search_key);
+            cached_bytes[device].free += search_key.bytes;
+
+            if (debug) {
+                fprintf(stderr, "[CCCL_PRELOAD] CACHE dev=%d bin=%u "
+                        "bytes=%zu cached_free=%zu\n",
+                        device, search_key.bin, search_key.bytes,
+                        cached_bytes[device].free);
+            }
+            return cudaSuccess;
+        } else {
+            // Don't cache — actually free
+            if (debug) {
+                fprintf(stderr, "[CCCL_PRELOAD] FREE  dev=%d bin=%u bytes=%zu\n",
+                        device, search_key.bin, search_key.bytes);
+            }
+            return real_cudaFree(d_ptr);
+        }
+    }
+};
+
+// Global singleton
+static CachingAllocator& get_allocator() {
+    static CachingAllocator alloc;
+    return alloc;
+}
+
+// ============================================================================
+// LD_PRELOAD interception points
+// ============================================================================
+
+extern "C" {
+
+cudaError_t cudaMalloc(void** devPtr, size_t size) {
+    return get_allocator().Allocate(devPtr, size);
+}
+
+cudaError_t cudaFree(void* devPtr) {
+    return get_allocator().Free(devPtr);
+}
+
+}  // extern "C"

qwen3_6_scripts/patch_ops.sh

@@ -246,6 +246,10 @@ if source != installed:
     raise SystemExit("runtime api_server overlay identity mismatch")
 PY
 
+build_stage "compiling CCCL CachingDeviceAllocator LD_PRELOAD module"
+bash ./build_cccl_preload_allocator.sh /workspace/qwen3_6_scripts || \
+    echo "[WARN] CCCL preload allocator build failed — will use default allocator"
+
 build_stage "compiling CoreX CUDA extensions (moe_index_combine + gdn_chunk_recurrent)"
 if [[ -x /usr/local/corex-3.2.3/bin/clang++ ]]; then
     bash ./build_corex_moe_index_combine.sh "${VLLM_ROOT}" || \

qwen3_6_scripts/test_cccl_preload.sh

@@ -0,0 +1,121 @@
+#!/usr/bin/env bash
+# Quick test for CCCL preload allocator on BI-V100
+# Usage: cd /home/dylan/project_6/qwen3_6_scripts && bash test_cccl_preload.sh
+set -eo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+SO="${SCRIPT_DIR}/cccl_preload_allocator.so"
+
+echo "=== Step 1: Build ==="
+bash "${SCRIPT_DIR}/build_cccl_preload_allocator.sh" "${SCRIPT_DIR}"
+echo ""
+
+if [[ ! -f "$SO" ]]; then
+    echo "BUILD FAILED: $SO not found"
+    exit 1
+fi
+
+echo "=== Step 2: Basic smoke test (torch.zeros on GPU) ==="
+echo "Without preload:"
+python3 -c "
+import time, torch
+t0=time.time()
+for i in range(100):
+    x=torch.zeros(1024*1024, device='cuda')
+    del x
+    torch.cuda.synchronize()
+print(f'100 alloc+free cycles: {time.time()-t0:.3f}s')
+"
+
+echo ""
+echo "With CCCL preload:"
+CCCL_ALLOC_DEBUG=0 LD_PRELOAD="$SO" python3 -c "
+import time, torch
+t0=time.time()
+for i in range(100):
+    x=torch.zeros(1024*1024, device='cuda')
+    del x
+    torch.cuda.synchronize()
+print(f'100 alloc+free cycles: {time.time()-t0:.3f}s')
+" 2>&1
+
+echo ""
+echo "=== Step 3: Varied sizes (simulating model inference allocations) ==="
+CCCL_ALLOC_DEBUG=0 LD_PRELOAD="$SO" python3 -c "
+import time, torch
+
+# Simulate inference: repeated allocs of same sizes (should hit cache)
+sizes = [512, 4096, 32768, 262144, 1048576, 4194304, 16777216]  # 512B to 16MB
+tensors = []
+
+print('First pass (cold cache):')
+t0 = time.time()
+for s in sizes:
+    x = torch.empty(s // 2, dtype=torch.float16, device='cuda')  # s bytes
+    tensors.append(x)
+t1 = time.time()
+print(f'  {len(sizes)} allocs: {(t1-t0)*1000:.1f}ms')
+
+print('Free all:')
+del tensors
+torch.cuda.synchronize()
+t2 = time.time()
+print(f'  {len(sizes)} frees: {(t2-t1)*1000:.1f}ms')
+
+print('Second pass (warm cache - should be faster):')
+tensors2 = []
+for s in sizes:
+    x = torch.empty(s // 2, dtype=torch.float16, device='cuda')
+    tensors2.append(x)
+t3 = time.time()
+print(f'  {len(sizes)} allocs: {(t3-t2)*1000:.1f}ms')
+
+print('Third pass (reuse same sizes 100x):')
+for _ in range(100):
+    for s in sizes:
+        x = torch.empty(s // 2, dtype=torch.float16, device='cuda')
+        del x
+t4 = time.time()
+print(f'  700 alloc+free: {(t4-t3)*1000:.1f}ms ({(t4-t3)/700*1000000:.0f}μs/op)')
+" 2>&1
+
+echo ""
+echo "=== Step 4: Large allocation test (model weights sized) ==="
+CCCL_ALLOC_DEBUG=0 LD_PRELOAD="$SO" python3 -c "
+import torch
+# Simulate KV cache blocks (typical: 256KB-2MB each)
+blocks = []
+for i in range(100):
+    b = torch.empty(256*1024 // 2, dtype=torch.float16, device='cuda')
+    blocks.append(b)
+print(f'Allocated 100 x 256KB blocks = {100*256/1024:.0f}MB')
+del blocks
+torch.cuda.synchronize()
+print('Freed all blocks')
+# Reallocate (should hit cache)
+blocks2 = []
+for i in range(100):
+    b = torch.empty(256*1024 // 2, dtype=torch.float16, device='cuda')
+    blocks2.append(b)
+print('Re-allocated 100 blocks (from cache)')
+print('OK: large allocation test passed')
+" 2>&1
+
+echo ""
+echo "=== Step 5: Stats output ==="
+CCCL_ALLOC_DEBUG=0 LD_PRELOAD="$SO" python3 -c "
+import torch
+for _ in range(50):
+    x = torch.zeros(1024*1024, device='cuda')
+    del x
+# Stats print on process exit
+" 2>&1
+
+echo ""
+echo "=== DONE ==="
+echo "If all tests passed, add to your launch command:"
+echo "  LD_PRELOAD=$SO python3 -m vllm.entrypoints.openai.api_server ..."
+echo ""
+echo "Or set in computility-run.yaml env:"
+echo "  - name: LD_PRELOAD"
+echo "    value: /workspace/qwen3_6_scripts/cccl_preload_allocator.so"