[sglang] SGLang의 Session-Aware Unified Radix Cache를 통한 추론 성능 최적화

2026년 8월 2일수정: 2026년 8월 2일

PR 링크: sgl-project/sglang#29173 상태: Merged | 변경: +1056 / -373

들어가며

LLM을 활용한 에이전트 워크로드나 멀티턴 대화형 서비스에서는 동일한 프롬프트 접두사(Prefix)를 반복적으로 사용하는 경우가 많습니다. 기존의 RadixCache는 메모리 압박이 발생하면 단순히 LRU(Least Recently Used) 정책에 따라 캐시를 비웠는데, 이 과정에서 현재 활성화된 세션이 참조 중인 중요한 KV 캐시까지 삭제되는 문제가 발생했습니다. 이로 인해 불필요한 재계산이 빈번해지고 전체적인 처리량(Throughput)이 저하되는 병목이 존재했습니다. 이번 SGLang의 UnifiedRadixCache 업데이트는 세션 인지(Session-aware) 기능을 도입하여 이 문제를 해결합니다.

코드 분석

1. 세션 라이프사이클 관리 (`UnifiedRadixCache`)

핵심은 요청별 session_id를 기반으로 KV 캐시의 생명주기를 관리하는 것입니다. UnifiedRadixCache는 이제 각 컴포넌트(Full Attention, SWA, Mamba)별로 세션 참조를 추적합니다.

Before (기존 방식):

# 단순히 캐시 정책에 따라 노드 삭제
self.evictable_size_ -= len(node.key)
if node in self.evictable_leaves:
    self.evictable_leaves.remove(node)

After (개선된 방식):

# 세션 참조를 고려하여 우선순위 기반 삭제
# 참조되지 않은 노드를 먼저 삭제하고, 참조된 노드는 세션 참조 카운트를 고려함
EvictionOrder = (is_referenced, session_ref_count, base_eviction_priority)

2. 세션 등록 및 해제 (`scheduler.py`)

요청이 완료되면 자동으로 캐시를 세션에 등록하고, 세션 종료 시 /close_session을 통해 참조를 해제합니다.

# scheduler.py 내 세션 종료 처리
def close_session(self, recv_req: CloseSessionReqInput):
    if self.enable_session_radix_cache:
        self.tree_cache.release_radix_session(recv_req.session_id)

왜 이게 좋은가

이 최적화의 핵심은 "참조 중인 KV 캐시를 소프트 보호(Soft-protection)한다"는 점입니다. 강제 고정(Pinning)이 아니므로 메모리가 정말 부족할 때는 참조된 데이터도 삭제될 수 있지만, 일반적인 상황에서는 불필요한 Eviction을 방지합니다.

성능 향상: 실제 SWE-agent 워크로드 테스트 결과, 캐시 적중률(Hit Rate)이 대폭 상승했습니다. 특히 배치 사이즈가 클 때 기존 방식보다 훨씬 안정적인 성능을 보여줍니다.
교훈: 캐시 정책을 설계할 때 단순히 시간적 지역성(Temporal Locality)만 고려하는 것이 아니라, 워크로드의 의미론적 단위(세션)를 고려하는 것이 대규모 추론 시스템의 효율을 결정짓는 핵심 요소임을 보여줍니다.

리뷰어 피드백 반영

리뷰 과정에서 ishandhanani와 krishs0404 등은 세션 생성(Generation) 관리와 우선순위 혼재 문제를 지적했습니다. 특히 세션 재시작 시 이전 요청이 잘못된 세션에 등록되는 문제를 방지하기 위해 session_generation을 도입하여 정합성을 확보했습니다. 또한, 성능을 위해 holder_set 연산을 최적화하여 O(1)에 가까운 성능을 유지하도록 개선되었습니다.