[sglang] SGLang HiCache에 AMD UMBP 기반 L3 스토리지 백엔드 도입

PR 링크: sgl-project/sglang#25377 상태: Merged | 변경: +2106 / -1

들어가며

최신 LLM 서빙 환경, 특히 대규모 컨텍스트를 처리하는 멀티-턴 워크로드에서 HBM(High Bandwidth Memory) 용량은 병목 현상의 주범입니다. SGLang의 HiCache는 HBM 기반의 L2 캐시를 제공하지만, 복잡한 에이전트 워크로드에서는 이조차 부족하여 빈번한 재계산(recomputation)이 발생합니다. 본 PR은 AMD의 mori 라이브러리 및 umbp(Unified Memory Buffer Pool) 서브시스템을 활용하여, 호스트 DRAM과 SSD를 활용하는 L3 계층형 스토리지 백엔드를 추가함으로써 이 문제를 해결합니다.

코드 분석

1. UMBP 호스트 메모리 할당자 (umbp_host_allocator.py)

새로운 UMBPHostTensorAllocator는 기존의 HostTensorAllocator를 확장하여, Hugepage를 지원하는 고성능 호스트 메모리 할당을 제공합니다.

# Before: 기본 torch.empty 기반 할당
# After: UMBP를 통한 Hugepage 할당
requested_backing = (
    self._mod.UMBPHostBufferBacking.AnonymousHugetlb
    if self._use_hugepage
    else self._mod.UMBPHostBufferBacking.Anonymous
)
handle = self._allocator.alloc(nbytes, requested_backing, ...)

이 할당자는 mori.umbp를 통해 커널 수준에서 Hugepage를 직접 관리하며, NUMA 노드 최적화 및 Prefaulting을 지원하여 호스트 메모리 접근 지연 시간을 최소화합니다.

2. UMBP 스토리지 백엔드 (umbp_store.py)

UMBPStore는 HiCache의 Zero-copy v1 인터페이스를 구현하여, KV 캐시 페이지를 호스트의 L3 계층으로 오프로드합니다.

# UMBPStore의 데이터 전송 로직
class UMBPStore(HiCacheStorage):
    def batch_get_v1(self, keys: List[int], ...):
        # UMBP 클라이언트를 통해 호스트 L3에서 데이터 인출
        return self._client.get(keys)

이 백엔드는 write_back 정책을 통해 비동기적으로 데이터를 저장하며, page_first 메모리 레이아웃을 강제하여 데이터 전송 효율을 극대화합니다.

왜 이게 좋은가

이번 최적화는 성능과 확장성 측면에서 큰 이점을 제공합니다.

1. TTFT(Time To First Token) 개선: 벤치마크 결과, HBM 전용 환경 대비 캐시 적중률이 약 47.6%p 상승했으며, 평균 TTFT가 약 3.5배 단축되었습니다.
2. 메모리 계층화: HBM(L2)이 가득 찼을 때, 데이터를 삭제하는 대신 호스트 DRAM(L3)으로 밀어냄으로써 재계산 비용을 획기적으로 줄였습니다.
3. Hugepage 활용: UMBPHostTensorAllocator를 통해 2MB Hugepage를 사용함으로써 TLB 미스를 줄이고, 대규모 KV 캐시 접근 시의 메모리 대역폭을 확보했습니다.

교훈: LLM 서빙에서 메모리 계층화는 단순히 용량을 늘리는 것이 아니라, 데이터 접근 패턴(LRU 등)과 하드웨어 특성(Hugepage, NUMA)을 결합할 때 가장 큰 성능 향상을 얻을 수 있습니다.

참고 자료

• mori library — AMD의 EP/IO 백엔드 및 UMBP 서브시스템

참고 자료

• https://github.com/ROCm/mori
• https://pytorch.org/docs/stable/generated/torch.frombuffer.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.