[sglang] SGLang HiCache: Mamba 브랜칭을 위한 증분 백업 최적화

PR 링크: sgl-project/sglang#33639 상태: Merged | 변경: +88 / -25

들어가며

SGLang의 UnifiedRadixCache는 효율적인 KV 캐시 관리를 위해 HiCache(Hierarchical Cache)를 사용합니다. 기존에는 새로운 Mamba 상태가 생성될 때 전체 KV를 다시 복사해야 하는 비효율이 존재했습니다. 이번 PR은 Mamba 브랜칭 시 전체 KV를 재복사하지 않고, 변경된 컴포넌트 데이터만 호스트로 백업하는 '증분 백업(Incremental Backup)' 메커니즘을 도입하여 캐시 적중률을 크게 향상시켰습니다.

코드 분석

1. unified_tree_core.py: 백업 로직 개선

핵심 변경 사항은 _should_backup_after_insert 메서드의 도입입니다. 기존에는 무조건 전체 KV를 백업하려 했으나, 이제는 컴포넌트별로 증분 백업이 필요한지 확인합니다.

# Before
if state.is_new_leaf and self._inc_hit_count_and_check(...):
    state.pending_actions.append(self._build_backup_kv_action(state.target_node))

# After
if self._should_backup_after_insert(state):
    state.pending_actions.append(self._build_backup_kv_action(state.target_node))

또한 _build_backup_spec에서 이미 백업된 노드에 대해 device_value를 빈 값([:0])으로 처리하여, 중복 전송을 방지하도록 수정되었습니다.

2. mamba_component.py: 증분 백업 감지

MambaComponent에 needs_incremental_backup 메서드를 추가하여, 호스트에 데이터가 없는 경우에만 백업이 트리거되도록 했습니다.

def needs_incremental_backup(self, node: UnifiedTreeNode) -> bool:
    data = node.component_data[self.component_type]
    return data.value is not None and data.host_value is None

3. memory_pool_host.py: 조건부 백업 실행

backup_from_device_all_layer 함수에서 host_indices가 존재할 때만 앵커 백업을 수행하도록 하여, 컴포넌트 전용 백업 시 불필요한 연산을 제거했습니다.

왜 이게 좋은가

이번 최적화는 특히 Mamba 모델과 같이 상태(State)가 복잡한 모델에서 큰 성능 향상을 가져옵니다. 벤치마크 결과에 따르면, 기존에는 Mamba 상태 복구 시 적중률이 15.74%에 불과했으나, 최적화 후에는 64.00%로 상승했습니다.

• 성능 수치: 6.4K 토큰 적중률이 1/10에서 10/10으로 개선됨.
• 교훈: 캐시 계층 구조에서 '전체 복사(Full Copy)'는 비용이 매우 큽니다. 데이터의 변경 지점(Delta)만을 추적하여 백업하는 증분 전략은 메모리 대역폭을 절약하고 캐시 효율을 극대화하는 핵심적인 설계 패턴입니다.

리뷰 피드백

리뷰 과정에서는 CI 테스트의 안정성을 확보하기 위해 여러 차례 재실행이 이루어졌으며, 특히 write-through 경로로 백업을 제한하여 동시성 문제를 방지하는 설계가 논의되었습니다. 이는 복잡한 캐시 동기화 환경에서 데이터 일관성을 유지하기 위한 필수적인 조치입니다.

참고 자료

• https://docs.sglang.ai/developer_guide/contribution_guide.html
• https://github.com/sgl-project/sglang/pull/31181

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