327c2c9044
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
100 lines
3.5 KiB
Bash
Executable File
100 lines
3.5 KiB
Bash
Executable File
#!/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"
|