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
This commit is contained in:
Claude
@@ -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'
|
||||
|
||||
99
qwen3_6_scripts/build_cccl_preload_allocator.sh
Executable file
99
qwen3_6_scripts/build_cccl_preload_allocator.sh
Executable file
@@ -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"
|
||||
405
qwen3_6_scripts/cccl_preload_allocator.cu
Normal file
405
qwen3_6_scripts/cccl_preload_allocator.cu
Normal file
@@ -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"
|
||||
@@ -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}" || \
|
||||
|
||||
121
qwen3_6_scripts/test_cccl_preload.sh
Executable file
121
qwen3_6_scripts/test_cccl_preload.sh
Executable file
@@ -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"
|
||||
Reference in New Issue
Block a user