biondizzle
aadde3ddf9
Key changes: - managed_alloc.cu: Add cudaMemPrefetchAsync to migrate pages to GPU immediately (prevents OOM from system RAM pinning on EGM systems where only ~102 GiB RAM remains). Add cudaMemAdviseSetAccessedBy for CPU so reads go over C2C NVLink without page migration. - vllm_managed_mem.py: Rewrite with idempotent patches, proper MemorySnapshot.measure() override, and torch.cuda tracking stubs for CUDAPluggableAllocator compatibility. - sitecustomize.py: Auto-loaded by Python in ALL subprocesses (including vLLM EngineCore). Applies allocator swap, torch patches, MemorySnapshot override, and request_memory override before any CUDA operations in spawned processes. - Dockerfile: Install sitecustomize.py into Python dist-packages. - README.md: Full rewrite with EGM problem statement, memory layout, architecture diagram, and build pipeline documentation.
95 lines
3.9 KiB
Plaintext
95 lines
3.9 KiB
Plaintext
// managed_alloc.cu - cudaMallocManaged allocator for PyTorch
|
|
// Compile: nvcc -shared -o libmanaged_alloc.so managed_alloc.cu -Xcompiler -fPIC
|
|
// Compatible with CUDA 13+ (uses cudaMemLocation API)
|
|
//
|
|
// Key design decisions for GH200 EGM:
|
|
// 1. cudaMallocManaged → allocations can page-fault across HBM + EGM
|
|
// 2. cudaMemAdviseSetPreferredLocation(GPU) → driver prefers keeping pages on GPU
|
|
// 3. cudaMemAdviseSetAccessedBy(CPU) → CPU can access over C2C NVLink without
|
|
// triggering page migration back to system RAM (critical: prevents OOM)
|
|
// 4. cudaMemPrefetchAsync(GPU) → actively migrates pages to GPU immediately,
|
|
// so subsequent writes go to HBM/EGM, not system RAM (prevents OOM on
|
|
// systems where EGM carved out most of system memory)
|
|
#include <cuda_runtime.h>
|
|
#include <stdio.h>
|
|
|
|
extern "C" {
|
|
|
|
// PyTorch pluggable allocator signature: void*(size_t, int, cudaStream_t)
|
|
void* managed_malloc(size_t size, int device, cudaStream_t stream) {
|
|
void* ptr = nullptr;
|
|
|
|
// Set the device before allocating
|
|
cudaError_t err = cudaSetDevice(device);
|
|
if (err != cudaSuccess) {
|
|
fprintf(stderr, "[managed_alloc] cudaSetDevice(%d) failed: %s\n",
|
|
device, cudaGetErrorString(err));
|
|
return nullptr;
|
|
}
|
|
|
|
// Use cudaMallocManaged - this is the key: allocations can page-fault
|
|
// across HBM and LPDDR on GH200 with EGM enabled
|
|
err = cudaMallocManaged(&ptr, size, cudaMemAttachGlobal);
|
|
if (err != cudaSuccess) {
|
|
fprintf(stderr, "[managed_alloc] cudaMallocManaged failed: %s "
|
|
"(size=%zu bytes / %.2f GiB)\n",
|
|
cudaGetErrorString(err), size, (double)size / (1024.0*1024.0*1024.0));
|
|
return nullptr;
|
|
}
|
|
|
|
// CUDA 13+ uses cudaMemLocation struct instead of int for device
|
|
cudaMemLocation gpu_loc;
|
|
gpu_loc.type = cudaMemLocationTypeDevice;
|
|
gpu_loc.id = device;
|
|
|
|
// Advise: prefer GPU placement. On GH200 with EGM, the hardware will
|
|
// migrate pages as needed, but the driver tries to keep them on GPU.
|
|
cudaMemAdvise(ptr, size, cudaMemAdviseSetPreferredLocation, gpu_loc);
|
|
|
|
// Advise: CPU will access this memory too. On GH200, this sets up
|
|
// remote mapping over C2C NVLink so CPU can read/write without
|
|
// triggering page migration back to system RAM. This is CRITICAL
|
|
// to prevent OOM on EGM systems where most system RAM was carved
|
|
// out for the GPU.
|
|
cudaMemLocation cpu_loc;
|
|
cpu_loc.type = cudaMemLocationTypeHost;
|
|
cpu_loc.id = cudaCpuDeviceId;
|
|
cudaMemAdvise(ptr, size, cudaMemAdviseSetAccessedBy, cpu_loc);
|
|
|
|
// Prefetch to GPU immediately. This actively migrates the virtual
|
|
// pages to the GPU side so that subsequent writes (e.g., model weight
|
|
// loading) go directly to HBM/EGM instead of pinning system RAM.
|
|
// Without this, the first write to each page faults into system RAM,
|
|
// which causes OOM when the OS only has ~102 GiB after EGM carveout.
|
|
//
|
|
// The prefetch is asynchronous on the given stream, so it won't block
|
|
// the calling thread. Subsequent operations on the same stream will
|
|
// wait for the prefetch to complete.
|
|
if (size > 0) {
|
|
err = cudaMemPrefetchAsync(ptr, size, gpu_loc, stream);
|
|
if (err != cudaSuccess) {
|
|
// Non-fatal: prefetch failure shouldn't prevent allocation.
|
|
// Pages will still be migrated on demand.
|
|
fprintf(stderr, "[managed_alloc] cudaMemPrefetchAsync warning: %s "
|
|
"(size=%.2f GiB, will use on-demand migration)\n",
|
|
cudaGetErrorString(err), (double)size / (1024.0*1024.0*1024.0));
|
|
}
|
|
}
|
|
|
|
return ptr;
|
|
}
|
|
|
|
// PyTorch pluggable allocator signature: void(void*, size_t, int, cudaStream_t)
|
|
void managed_free(void* ptr, size_t size, int device, cudaStream_t stream) {
|
|
if (ptr != nullptr) {
|
|
// Sync the stream before freeing to avoid use-after-free with
|
|
// managed memory (in-flight page faults can race with deallocation).
|
|
if (stream != nullptr) {
|
|
cudaStreamSynchronize(stream);
|
|
}
|
|
cudaFree(ptr);
|
|
}
|
|
}
|
|
|
|
} // extern "C"
|