[vllm] ROCm Sleep Mode - AMD GPU 전력 절약 모드 이식

PR 링크: vllm-project/vllm#12695 상태: Merged | 변경: +582/-31

들어가며

vLLM의 sleep mode는 GPU가 유휴 상태일 때 가중치 메모리를 해제하여 다른 프로세스가 GPU를 사용할 수 있게 하는 기능이다. 이 PR은 CUDA 전용이었던 이 기능을 AMD ROCm 플랫폼으로 이식한다.

핵심 코드 분석

ROCm용 청크 기반 메모리 관리

// ROCm의 cuMemCreate 제한 대응: 청크 단위 할당
static const unsigned long long DEFAULT_MEMCREATE_CHUNK_SIZE =
    (256ULL * 1024ULL * 1024ULL);  // 256MB 기본 청크

static unsigned long long get_memcreate_chunk_size() {
    const char* env = getenv("VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE");
    // 환경변수로 청크 크기 조정 가능
    ...
}

ROCm의 cuMemCreate는 한 번에 할당할 수 있는 크기에 제한이 있다. 이를 우회하기 위해 256MB 청크 단위로 분할 할당하고, 환경변수로 조정 가능하게 했다.

청크 단위 map/unmap

// CUDA: 단일 할당
cuMemCreate(p_memHandle, size, &prop, 0);
cuMemMap(d_mem, size, 0, *p_memHandle, 0);

// ROCm: 청크 분할
for (auto i = 0; i < num_chunks; ++i) {
    cuMemCreate(p_memHandle[i], chunk_sizes[i], &prop, 0);
}
for (auto i = 0; i < num_chunks; ++i) {
    void* map_addr = (void*)((uintptr_t)d_mem + allocated_size);
    cuMemMap(map_addr, chunk_sizes[i], 0, *(p_memHandle[i]), 0);
    allocated_size += chunk_sizes[i];
}

CMake 빌드 지원

# HIP도 cumem allocator 빌드에 포함
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    message(STATUS "Enabling cumem allocator extension.")
    # ROCm: libamdhip64 링크
    find_library(AMDHIP64_LIB NAMES amdhip64 ...)

왜 이게 좋은가

1. AMD GPU 지원 확대: ROCm 사용자도 sleep mode의 전력 절약 혜택 활용
2. RL Training 시너지: 학습-추론 교대 시 유휴 GPU 메모리를 효율적으로 재활용
3. 에러 처리: 청크 단위 할당/해제에서의 부분 실패를 안전하게 처리
4. 환경변수 커스터마이징: VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE로 워크로드에 맞게 조정

정리

플랫폼 이식 PR이지만, ROCm의 메모리 API 제한을 청크 분할로 우회한 엔지니어링이 돋보인다. AMD GPU 사용자에게 CUDA와 동일한 기능을 제공하여 vLLM의 하드웨어 중립성을 강화한다.

참고 자료

• PR #12695

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.