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:
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