diff --git a/computility-run.yaml b/computility-run.yaml index 947b6a6b..38cf2789 100644 --- a/computility-run.yaml +++ b/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' diff --git a/qwen3_6_scripts/build_cccl_preload_allocator.sh b/qwen3_6_scripts/build_cccl_preload_allocator.sh new file mode 100755 index 00000000..e57580ea --- /dev/null +++ b/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" diff --git a/qwen3_6_scripts/cccl_preload_allocator.cu b/qwen3_6_scripts/cccl_preload_allocator.cu new file mode 100644 index 00000000..fc52498c --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include + +// 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; +using BusyBlocks = std::multiset; + +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 cached_bytes; + + // Stats + std::atomic stat_hits{0}; + std::atomic stat_misses{0}; + std::atomic stat_frees{0}; + std::atomic 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 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 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" diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 7604c881..a678bd1c 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/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}" || \ diff --git a/qwen3_6_scripts/test_cccl_preload.sh b/qwen3_6_scripts/test_cccl_preload.sh new file mode 100755 index 00000000..5b26dfcd --- /dev/null +++ b/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"