biondizzle
bcc872c2c3
sitecustomize.py: No longer swaps CUDAPluggableAllocator globally. Sets VLLM_KV_CACHE_USE_MANAGED_MEMORY=1 instead. vllm_managed_mem.py: No global allocator swap, no torch.cuda patches.
272 lines
10 KiB
Python
272 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
vllm_managed_mem.py - Launch vLLM with cudaMallocManaged allocator
|
||
|
||
This MUST be the very first thing that runs before any torch.cuda calls.
|
||
It swaps PyTorch's CUDA allocator to use cudaMallocManaged, enabling
|
||
transparent page-fault access to EGM memory on GH200.
|
||
|
||
Usage:
|
||
python vllm_managed_mem.py [all normal vllm serve arguments]
|
||
|
||
Example:
|
||
python vllm_managed_mem.py --model google/gemma-4-31B-it \
|
||
--host 0.0.0.0 --port 80 --gpu-memory-utilization 0.90 \
|
||
--enforce-eager --max-model-len 32768
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import ctypes
|
||
|
||
|
||
def get_total_managed_memory_gb():
|
||
"""
|
||
Calculate total memory available via managed allocations on GH200.
|
||
First checks MANAGED_MEMORY_TOTAL_GB env var (explicit override).
|
||
Then tries /proc/iomem (requires --privileged container).
|
||
Falls back to HBM only if neither works.
|
||
"""
|
||
# Explicit override takes priority
|
||
env_val = os.environ.get('MANAGED_MEMORY_TOTAL_GB')
|
||
if env_val:
|
||
return float(env_val)
|
||
|
||
# Try /proc/iomem for EGM regions (requires --privileged or host mount)
|
||
egm_bytes = 0
|
||
try:
|
||
with open('/proc/iomem', 'r') as f:
|
||
for line in f:
|
||
if 'System RAM (NVIDIA)' in line:
|
||
parts = line.strip().split(':')[0].strip()
|
||
start_s, end_s = parts.split('-')
|
||
start = int(start_s, 16)
|
||
end = int(end_s, 16)
|
||
egm_bytes += (end - start + 1)
|
||
except (PermissionError, FileNotFoundError, ValueError):
|
||
pass
|
||
|
||
egm_gb = egm_bytes / (1024**3)
|
||
if egm_gb > 0:
|
||
# EGM detected via iomem, add HBM
|
||
import torch
|
||
_, hbm_total = torch.cuda.mem_get_info(0)
|
||
hbm_gb = hbm_total / (1024**3)
|
||
return egm_gb + hbm_gb
|
||
|
||
# No EGM detected - return 0 (caller should handle)
|
||
return 0
|
||
|
||
|
||
def swap_allocator():
|
||
"""
|
||
Replace PyTorch's default CUDA allocator with our managed memory allocator.
|
||
This MUST happen before any CUDA tensors are created.
|
||
If sitecustomize.py already swapped it, this is a no-op.
|
||
"""
|
||
lib_path = os.environ.get(
|
||
'MANAGED_ALLOC_LIB',
|
||
'/usr/local/lib/libmanaged_alloc.so'
|
||
)
|
||
|
||
if not os.path.exists(lib_path):
|
||
print(f"[managed_mem] ERROR: {lib_path} not found!", file=sys.stderr)
|
||
print(f"[managed_mem] Build it with: nvcc -shared -o {lib_path} "
|
||
f"managed_alloc.cu -Xcompiler -fPIC", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
# Verify the library loads
|
||
try:
|
||
lib = ctypes.CDLL(lib_path)
|
||
assert hasattr(lib, 'managed_malloc'), "managed_malloc symbol not found"
|
||
assert hasattr(lib, 'managed_free'), "managed_free symbol not found"
|
||
except Exception as e:
|
||
print(f"[managed_mem] ERROR loading {lib_path}: {e}", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
import torch
|
||
try:
|
||
alloc = torch.cuda.memory.CUDAPluggableAllocator(
|
||
lib_path, 'managed_malloc', 'managed_free'
|
||
)
|
||
torch.cuda.memory.change_current_allocator(alloc)
|
||
print(f"[managed_mem] Allocator swapped to cudaMallocManaged", file=sys.stderr)
|
||
except RuntimeError as e:
|
||
if "already initialized" in str(e):
|
||
print(f"[managed_mem] Allocator already initialized (sitecustomize)",
|
||
file=sys.stderr)
|
||
else:
|
||
raise
|
||
|
||
|
||
def patch_memory_snapshot():
|
||
"""
|
||
Patch MemorySnapshot.measure() to report the full managed memory
|
||
instead of just HBM.
|
||
|
||
This is the core fix: cudaMemGetInfo only reports HBM (~96 GiB),
|
||
but with cudaMallocManaged we can address HBM + EGM (~474 GiB).
|
||
We override the snapshot so all downstream code (request_memory,
|
||
KV cache sizing, etc.) sees the full managed memory capacity.
|
||
"""
|
||
from vllm.utils.mem_utils import MemorySnapshot
|
||
|
||
_original_measure = MemorySnapshot.measure
|
||
|
||
def patched_measure(self):
|
||
_original_measure(self)
|
||
# Override total_memory with managed memory capacity
|
||
managed_total_gb = os.environ.get('MANAGED_MEMORY_TOTAL_GB')
|
||
if managed_total_gb:
|
||
managed_total = int(float(managed_total_gb) * (1024**3))
|
||
self.total_memory = managed_total
|
||
# free_memory = total - what CUDA is using
|
||
self.free_memory = managed_total - self.cuda_memory
|
||
print(f"[managed_mem] MemorySnapshot patched: "
|
||
f"total={self.total_memory / (1024**3):.0f} GiB, "
|
||
f"free={self.free_memory / (1024**3):.0f} GiB",
|
||
file=sys.stderr)
|
||
|
||
MemorySnapshot.measure = patched_measure
|
||
print(f"[managed_mem] Patched MemorySnapshot.measure()", file=sys.stderr)
|
||
|
||
|
||
def patch_torch_memory_tracking():
|
||
"""
|
||
Patch PyTorch memory tracking functions that CUDAPluggableAllocator
|
||
doesn't support. If sitecustomize already patched them, this is a no-op.
|
||
"""
|
||
import torch
|
||
|
||
# Check if already patched by sitecustomize
|
||
if torch.cuda.reset_peak_memory_stats.__name__ == '_patched_reset_peak_memory_stats':
|
||
print(f"[managed_mem] torch.cuda already patched (sitecustomize)", file=sys.stderr)
|
||
return
|
||
|
||
_original_reset_peak = torch.cuda.reset_peak_memory_stats
|
||
_original_memory_stats = torch.cuda.memory_stats
|
||
_original_max_memory_allocated = torch.cuda.max_memory_allocated
|
||
|
||
def _patched_reset_peak_memory_stats(device=None):
|
||
"""No-op: CUDAPluggableAllocator doesn't support resetPeakStats."""
|
||
pass
|
||
|
||
def _patched_memory_stats(device=None):
|
||
"""Return minimal stats dict to avoid crashes."""
|
||
try:
|
||
return _original_memory_stats(device)
|
||
except RuntimeError:
|
||
return {
|
||
"allocated_bytes.all.current": 0,
|
||
"allocated_bytes.all.peak": 0,
|
||
"reserved_bytes.all.current": 0,
|
||
"reserved_bytes.all.peak": 0,
|
||
"num_alloc_retries": 0,
|
||
"num_ooms": 0,
|
||
}
|
||
|
||
def _patched_max_memory_allocated(device=None):
|
||
"""Return 0 since we can't track peak with pluggable allocator."""
|
||
try:
|
||
return _original_max_memory_allocated(device)
|
||
except RuntimeError:
|
||
return 0
|
||
|
||
torch.cuda.reset_peak_memory_stats = _patched_reset_peak_memory_stats
|
||
torch.cuda.memory_stats = _patched_memory_stats
|
||
torch.cuda.max_memory_allocated = _patched_max_memory_allocated
|
||
|
||
if hasattr(torch.accelerator, 'reset_peak_memory_stats'):
|
||
torch.accelerator.reset_peak_memory_stats = _patched_reset_peak_memory_stats
|
||
if hasattr(torch.accelerator, 'memory_stats'):
|
||
torch.accelerator.memory_stats = _patched_memory_stats
|
||
if hasattr(torch.accelerator, 'max_memory_allocated'):
|
||
torch.accelerator.max_memory_allocated = _patched_max_memory_allocated
|
||
|
||
print(f"[managed_mem] Patched torch.cuda memory tracking (no-op stubs)",
|
||
file=sys.stderr)
|
||
|
||
|
||
def patch_vllm_memory_check():
|
||
"""
|
||
Monkey-patch vLLM's memory validation to understand managed memory.
|
||
|
||
With managed memory, cudaMemGetInfo only reports HBM, so the free-memory
|
||
check would fail. When MANAGED_MEMORY_TOTAL_GB is set, we skip that check.
|
||
"""
|
||
import vllm.v1.worker.utils as worker_utils
|
||
|
||
_original_request_memory = worker_utils.request_memory
|
||
|
||
def patched_request_memory(init_snapshot, cache_config):
|
||
managed_total_gb = os.environ.get('MANAGED_MEMORY_TOTAL_GB')
|
||
if managed_total_gb:
|
||
total_bytes = float(managed_total_gb) * (1024**3)
|
||
gpu_util = cache_config.gpu_memory_utilization
|
||
requested = int(total_bytes * gpu_util)
|
||
print(f"[managed_mem] Overriding memory request: "
|
||
f"{float(managed_total_gb):.0f} GiB × {gpu_util} = "
|
||
f"{requested / (1024**3):.1f} GiB", file=sys.stderr)
|
||
return requested
|
||
else:
|
||
return _original_request_memory(init_snapshot, cache_config)
|
||
|
||
worker_utils.request_memory = patched_request_memory
|
||
print(f"[managed_mem] Patched vLLM request_memory", file=sys.stderr)
|
||
|
||
|
||
def main():
|
||
# Step 1: NO global allocator swap — model weights stay in HBM.
|
||
# KV cache uses cudaMallocManaged directly via
|
||
# VLLM_KV_CACHE_USE_MANAGED_MEMORY env var (set by sitecustomize.py).
|
||
# The global allocator swap broke cuBLAS GEMM operations because
|
||
# intermediate compute tensors ended up in managed memory.
|
||
print(f"[managed_mem] Using targeted KV cache managed allocation "
|
||
f"(no global allocator swap)", file=sys.stderr)
|
||
|
||
# Step 2: Calculate total managed memory and export it
|
||
total_managed_gb = get_total_managed_memory_gb()
|
||
if total_managed_gb <= 0:
|
||
print(f"[managed_mem] WARNING: No managed memory detected! "
|
||
f"Set MANAGED_MEMORY_TOTAL_GB env var explicitly.",
|
||
file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
os.environ['MANAGED_MEMORY_TOTAL_GB'] = str(total_managed_gb)
|
||
print(f"[managed_mem] MANAGED_MEMORY_TOTAL_GB={total_managed_gb:.0f}",
|
||
file=sys.stderr)
|
||
|
||
# Step 3: No torch.cuda memory tracking patches needed —
|
||
# we're not using CUDAPluggableAllocator anymore.
|
||
|
||
# Step 4: Patch MemorySnapshot.measure() to report full managed memory
|
||
# This is critical - without it, all downstream code only sees HBM
|
||
patch_memory_snapshot()
|
||
|
||
# Step 5: Patch request_memory as a safety net
|
||
patch_vllm_memory_check()
|
||
|
||
# Step 6: Launch vLLM's API server with remaining args
|
||
sys.argv = ['vllm.entrypoints.openai.api_server'] + sys.argv[1:]
|
||
print(f"[managed_mem] Launching vLLM with args: {sys.argv[1:]}",
|
||
file=sys.stderr)
|
||
|
||
# Import and run
|
||
from vllm.entrypoints.openai.api_server import run_server, FlexibleArgumentParser
|
||
import uvloop
|
||
|
||
parser = FlexibleArgumentParser(
|
||
description="vLLM OpenAI-compatible API server (managed memory)"
|
||
)
|
||
|
||
# Use vLLM's own argument parser
|
||
from vllm.entrypoints.openai.api_server import make_arg_parser
|
||
parser = make_arg_parser(parser)
|
||
args = parser.parse_args(sys.argv[1:])
|
||
|
||
uvloop.run(run_server(args))
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|