[sglang] SGLang UnifiedRadixTree에 HiCache 프레임워크 도입: 하이브리드 모델 성능 최적화

2026년 5월 3일수정: 2026년 5월 3일

PR 링크: sgl-project/sglang#23316 상태: Merged | 변경: +None / -None

들어가며

최신 대규모 언어 모델(LLM) 환경에서 KV 캐시 관리는 추론 성능을 결정짓는 핵심 요소입니다. 특히 Hybrid Linear 모델이나 DeepSeek DSA와 같은 복합 아키텍처는 서로 다른 캐시 요구사항을 가집니다. 이번 PR은 SGLang의 UnifiedRadixTree에 HiCache 프레임워크를 통합하여, Full, Mamba, SWA 등 다양한 컴포넌트 간의 캐시 생명주기를 통일하고, 호스트 측 LRU 관리 및 계층적 캐시 로직을 최적화하는 것을 목표로 합니다.

코드 분석

1. Scheduler 및 캐시 초기화 (`python/sglang/srt/managers/scheduler.py`)

기존에는 캐시 타입별로 분기 처리가 파편화되어 있었으나, UnifiedRadixCache를 중심으로 통합되었습니다. 특히 HiCache가 활성화된 경우, tp_worker에 레이어 전송 카운터를 등록하여 데이터 전송 상태를 추적합니다.

# Before
elif self.enable_hierarchical_cache:
    # ... 기존 개별 캐시 로직 ...

# After
elif envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get():
    # ... UnifiedRadixCache 초기화 ...
    if self.enable_hierarchical_cache:
        self.tree_cache.init_hicache(server_args, params)
        self.tp_worker.register_hicache_layer_transfer_counter(
            self.tree_cache.cache_controller.layer_done_counter
        )

2. 하이브리드 풀 어셈블러 (`python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py`)

NSATokenToKVPool을 지원하기 위해 attach_hybrid_pool_to_unified_cache 함수가 확장되었습니다. 이제 NSA(Non-shared Attention) 스택에 대해서도 PoolName.INDEXER를 통해 공유 앵커 인덱스를 관리하며, HiCache 백업/로드 시 KV와 함께 처리되도록 최적화되었습니다.

# NSA 스택 지원 추가
nsa_stack = isinstance(kvcache, NSATokenToKVPool)
if nsa_stack:
    # ... NSAIndexerPoolHost 생성 및 등록 ...
    cache.register_hicache_anchor_kv_shared_indices_pool(
        PoolName.INDEXER,
        hit_policy=PoolHitPolicy.ALL_PAGES,
    )

3. FullComponent 매칭 로직 (`python/sglang/srt/mem_cache/unified_cache_components/full_component.py`)

HiCache가 활성화되면 evicted(퇴출된) 노드도 매칭 경계로 인정받아야 합니다. create_match_validator를 수정하여 백업된 노드도 유효한 매칭 대상으로 포함시켰습니다.

# After: evicted + backuped 노드도 유효한 매칭 경계로 처리
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
    return lambda node: (
        node.component_data[self.component_type].value is not None or node.backuped
    )

왜 이게 좋은가

이번 최적화는 단순히 기능을 추가하는 것을 넘어, 하이브리드 모델에서 발생하는 캐시 미스 문제를 해결합니다.

성능 향상: 공유 프리픽스 워크로드 테스트에서 Round 2/3 수행 시 Round 1 대비 TTFT(Time To First Token)가 약 27% 감소하는 성과를 보였습니다.
안정성: 16개 TP 랭크를 사용하는 대규모 환경에서 2048개의 프롬프트를 처리하는 동안 OOM이나 행(hang) 없이 안정적으로 동작했습니다.
데이터 무결성: SHA-256 핑거프린트 검증을 통해 D→H(Device to Host) 백업 및 H→D 로드 과정에서 바이트 단위의 정확성을 100% 보장함을 확인했습니다.

일반적인 교훈으로, 복잡한 캐시 계층 구조를 가질수록 컴포넌트별 생명주기를 통합 관리하는 UnifiedRadixTree와 같은 추상화 계층이 시스템의 유지보수성과 확장성을 크게 높여준다는 점을 시사합니다.