[SGLang] CUDA IPC Pool Handle 캐싱으로 멀티모달 전송 최적화

PR 링크: sgl-project/sglang#21418 상태: Merged | 변경: +195 / -50

들어가며

SGLang에서 멀티모달 모델 추론 시, 이미지/비디오 feature를 토크나이저 프로세스에서 스케줄러 프로세스로 전송할 때 CUDA IPC(Inter-Process Communication)를 사용한다. 기존 구현은 매 텐서마다 cudaIpcOpenMemHandle을 호출하여 IPC 핸들을 열었는데, 이 호출은 CUDA 드라이버 수준에서 비용이 크다.

이 PR은 메모리 풀 수준에서 IPC 핸들을 한 번만 열고 캐싱하여, 이후 전송에서는 캐시된 storage에서 byte offset으로 슬라이싱하는 방식으로 전환한다.

핵심 코드 분석

Pool 수준 IPC 핸들 캐시 도입

Before:

# 매 텐서마다 개별 IPC 핸들을 직렬화/역직렬화
self.proxy_state = self.get_proxy_state(data, info_data)
# get_proxy_state 내부에서 data.untyped_storage()._share_cuda_() 호출

After:

# 풀 수준 핸들을 사전 계산하여 캐싱
storage = self.memory_pool.untyped_storage()
self._pool_ipc_handle = storage._share_cuda_()
self._pool_device_index = self.memory_pool.device.index

# 텐서 전송 시 풀 핸들 + byte offset만 전달
if pool_ipc_handle is not None:
    self.proxy_state = {
        "ipc_extra": {
            "pool_handle": pool_ipc_handle,
            "pool_byte_offset": pool_byte_offset,
            "pool_device_index": pool_device_index,
            "shape": data.shape,
            "dtype": data.dtype,
            ...
        },
        "tensor_data": None,
    }

소비자 측 Double-checked Locking 캐시

수신 측에서도 풀 핸들을 한 번만 열고 재사용한다.

_pool_storage_cache: dict = {}
_pool_cache_lock = threading.Lock()

def _pool_handle_cache_get_or_open(cache_key, pool_handle):
    storage = _pool_storage_cache.get(cache_key)
    if storage is None:
        with _pool_cache_lock:
            storage = _pool_storage_cache.get(cache_key)
            if storage is None:
                storage = _open_pooled_storage_uncached(pool_handle)
                _pool_storage_cache[cache_key] = storage
    return storage

복원 시 캐시 실패 대비 Fallback

try:
    slice_tensor, _, cache_key, _ = self._reconstruct_from_ipc_extra(
        ipc_extra, use_cache=True
    )
except Exception as e:
    _pool_handle_cache_invalidate(cache_key)
    slice_tensor, _, _, _ = self._reconstruct_from_ipc_extra(
        ipc_extra, use_cache=False
    )

왜 이게 좋은가

1. IPC 핸들 오픈 비용 제거: cudaIpcOpenMemHandle은 CUDA 드라이버 내부에서 mutex 획득과 메모리 매핑을 수행하는 무거운 연산이다. 풀 단위 캐싱으로 이를 최초 1회로 줄인다.
2. Thread-safe 설계: Double-checked locking 패턴과 cache invalidation/clear API로 멀티스레드 환경에서도 안전하다.
3. 점진적 도입: SGLANG_USE_IPC_POOL_HANDLE_CACHE 환경변수로 opt-in 방식이라 기존 동작에 영향 없이 테스트 가능하다.

정리

멀티모달 추론에서 프로세스 간 텐서 전송은 숨겨진 병목이 될 수 있다. 이 PR은 풀 수준 IPC 핸들 캐싱이라는 단순한 아이디어로 반복적인 드라이버 호출을 제거하며, fallback 로직으로 안정성도 확보했다.

참고 자료

• CUDA IPC Documentation
• PyTorch UntypedStorage.share_cuda

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석과 해석에서 오류가 있을 수 있으니, 정확한 내용은 원본 PR을 참고해주세요.