vLLM에서 메모리 누수 디버깅

Mistral AI는 분리형 서빙(Disaggregated Serving)의 사전 생산 테스트 중 발생한 vLLM 시스템 메모리 누수를 식별하고 해결했습니다. 이 누수는 분리형 서빙 설정에서 InfiniBand 최적화를 위해 사용되는 UCX(통합 커뮤니케이션 X) 라이브러리의 메모리 훅 메커니즘에서 비롯되었으며, 시스템 메모리가 분당 약 400MB 증가하는 현상을 일으켰습니다.

누수: 조건과 증상

이 메모리 누수는 특정 조건에서만 나타났습니다: vLLM과 Mistral Medium 3.1 모델을 사용하고, 그래프 컴파일이 활성화되어 있으며, NIXL을 사용한 Prefill/Decode (P/D) 분리형 서빙 설정인 경우입니다. 누수는 NIXL과 UCX를 통해 KVCache 전송이 시작되는 디코딩 측에서만 발생했습니다.

P/D 분리형 설정에서는 프로세스가 두 단계로 나뉩니다:

Prefill 단계: 라우터가 vLLM 인스턴스에 Prefill 요청을 보내 KVCache를 계산합니다.
Decode 단계: 라우터가 KVCache 메타데이터와 디코딩 요청을 디코딩용 vLLM 인스턴스로 전송하여 전달된 KVCache를 사용해 토큰을 생성합니다.

진단 과정: 파이썬에서 커널 추적까지

Mistral AI 엔지니어링 팀은 소프트웨어 스택의 여러 계층을 거쳐 누수를 고립시키는 체계적인 접근 방식을 사용했습니다.

고수준 프로파일링

초기에는 파이썬 메모리 프로파일링 도구인 Memray와 Guppy 3를 사용했지만, 누수는 발견되지 않았습니다. GDB는 프로세스를 충돌시켰고, Valgrind는 vLLM의 무거운 설정에 비해 너무 느렸습니다. 이로 인해 팀은 vLLM 저장소에 GitHub 이슈를 열어 다른 사용자들도 동일한 문제를 겪고 있는지 확인했습니다.

Heaptrack을 통한 힙 분석

Heaptrack을 사용해 malloc과 free 작업을 모니터링했습니다. 힙 메모리는 안정적이었지만, 정점 주거 집합 크기(RSS) 에서 불일치를 관찰했습니다. 이는 누수가 힙 외부, 즉 glibc의 malloc이 아닌 mmap을 통해 할당된 익명 메모리 매핑 영역에서 발생하고 있음을 시사했습니다.

시스템 수준의 메모리 매핑

/proc/<pid>/maps를 읽기 위해 pmap 명령어를 사용했을 때, 일부 익명 메모리 영역이 증가하고 시작 주소가 변경되는 것을 확인했습니다. 이 행동은 mremap 또는 mmap과 munmap의 반복적인 사이클이 적절히 해제되지 않은 경우를 시사했습니다.

BPFtrace를 통한 커널 추적

할당의 원인을 확인하기 위해 BPFtrace를 사용해 모든 mmap, munmap, mremap 시스템 호출을 로깅했습니다. 누수되는 주소가 glibc의 원시 시스템 호출 래퍼(syscall+29)를 통해 발생하는 mmap 호출에서 비롯된다는 것을 발견했습니다. 이는 표준 glibc 래퍼와 LD_PRELOAD 훅을 우회하는 방식이었습니다.

타겟 GDB 자동화

BPFtrace는 일부 종속성에서 프레임 포인터가 비활성화되어 있어 전체 사용자 공간 스택 트레이스를 제공하지 못했습니다. 이에 팀은 GDB를 사용해 syscall 주소에 조건부 브레이크포인트를 설정했습니다. SYS_mmap에만 트리거되도록 설정하고 전체 스택 트레이스를 출력함으로써, 파이썬이 UCX를 통해 mmap을 호출하고 있음을 발견했습니다.

근본 원인: UCX mmap 훅 메커니즘

조사 결과, UCX는 InfiniBand 메모리 등록(등록 캐시 또는 RCache)을 최적화하기 위해 mmap 훅 메커니즘을 사용하고 있었습니다. 이 메커니즘은 기본적으로 mmap과 munmap의 글로벌 오프셋 테이블(GOT) 엔트리를 동적으로 수정하여 모든 호출을 가로챕니다.

이 가로채기로 인해 두 가지 주요 문제가 발생했습니다:

훅 우회: 표준 디버깅 도구와 LD_PRELOAD 훅이 누수를 추적하지 못하게 되었습니다.
메모리 누적: UCX는 munmap 호출 시 즉시 메모리를 해제하지 않고, 해당 영역을 무효화 대기열로 이동합니다. 이 특정한 에지 케이스에서는 이 대기열을 관리하는 메모리 풀이 동적으로 확장되어, munmap 작업 중에도 mmap 호출이 발생하게 되어 RSS가 선형적으로 증가했습니다.

해결 및 수정

이 누수는 두 가지 가능한 구성으로 해결되었습니다:

훅 비활성화: 환경 변수 UCX_MEM_MMAP_HOOK_MODE=none을 설정하면 훅 메커니즘을 완전히 비활성화할 수 있습니다. vLLM은 단 한 번만 큰 연속된 메모리 영역(KVCache 매니저 메모리)을 등록하면 되므로, 이 설정은 vLLM 성능에 부정적인 영향을 미치지 않았습니다.
캐시 제한: UCX_RCACHE_MAX_UNRELEASED=1024 (기본값 inf 대신)로 설정하면, 미해제된 메모리 영역의 임계값에 도달했을 때 UCX가 정리 작업을 시작하도록 강제할 수 있습니다.

Mistral AI는 vLLM 저장소에 수정 사항을 병합했습니다 (PR #32181)하고, NIXL 및 UCX 팀과 협력하여 향후 NIXL 릴리스에서 UCX_RCACHE_MAX_UNRELEASED의 기본 동작을 변경했습니다.