[sglang] SGLang의 Unified Radix Cache를 위한 SWA HiCache 지원 최적화

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

들어가며

최신 대규모 언어 모델(LLM) 추론에서 KV Cache 관리는 성능의 핵심입니다. 특히 긴 문맥(Long Context)을 처리할 때 메모리 병목 현상이 빈번하게 발생합니다. SGLang은 이를 해결하기 위해 Unified Radix Cache와 HiCache를 도입해왔습니다. 이번 PR은 기존의 HiCache 구조에 SWA(Sliding Window Attention) 지원을 추가하여, SWA를 사용하는 모델에서도 계층적 캐싱(Hierarchical Caching)의 이점을 누릴 수 있도록 최적화했습니다.

코드 분석

1. hybrid_cache_controller.py: 동적 할당 로직 개선

기존에는 KV 캐시 할당이 고정적이었으나, SWA 지원을 위해 full_attn_allocator를 동적으로 참조하도록 변경되었습니다.

# Before
device_indices = self.mem_pool_device_allocator.alloc(len(host_indices))

# After
full_allocator = getattr(
    self.mem_pool_device_allocator,
    "full_attn_allocator",
    self.mem_pool_device_allocator,
)
device_indices = full_allocator.alloc(len(host_indices))

이 변경은 SWA와 같은 특수 어텐션 메커니즘이 별도의 할당기(Allocator)를 가질 수 있도록 유연성을 제공합니다.

2. hybrid_pool_assembler.py: SWA 스택 빌드 로직 추가

build_hybrid_swa_stack 함수가 새로 추가되어 SWA 전용 풀을 통합 캐시 구조에 등록합니다.

# SWA 전용 풀 엔트리 등록
build_pool_entry(
    name=PoolName.SWA,
    host_pool=swa_host_pool,
    device_pool=swa_kv_pool,
    device_alloc_fn=swa_attn_allocator.alloc,
    device_free_fn=swa_attn_allocator.free,
)

이를 통해 HybridCacheController가 SWA 메모리 풀의 할당과 해제를 직접 제어할 수 있게 되었습니다.

왜 이게 좋은가

이번 최적화의 핵심은 메모리 계층화(Memory Tiering)를 SWA 모델까지 확장했다는 점입니다.

1. 성능 향상: 벤치마크 결과, HiCache를 적용했을 때 긴 문맥 처리에서 처리량(Throughput)이 크게 개선되었습니다. 특히 KV Cache가 메모리 병목인 환경에서 SWA를 효율적으로 관리하여 더 많은 요청을 동시에 처리할 수 있습니다.
2. 안정성: 리뷰 과정에서 발견된 AssertionError와 같은 레이스 컨디션 이슈들을 atomic rollback 로직(할당 실패 시 이전 할당분을 모두 해제)을 통해 해결했습니다.

교훈: 복잡한 캐시 시스템을 설계할 때는 할당 실패 시의 원자적 롤백(Atomic Rollback)이 필수적입니다. 또한, getattr을 활용한 동적 할당기 참조는 다양한 어텐션 패턴을 지원하는 확장성 있는 아키텍처를 만드는 데 매우 효과적인 패턴입니다.

리뷰어 피드백 반영

실제 프로덕션 환경에서 사용 중인 사용자들의 피드백을 통해, unified_radix_cache.py 내의 캐시 보호 로직과 LRU 관리에서 발생하던 충돌을 수정했습니다. 이는 오픈소스 프로젝트에서 커뮤니티의 실사용 피드백이 얼마나 중요한지 보여주는 사례입니다.

참고 자료

• https://github.com/sgl-project/sglang
• https://pytorch.org/docs/stable/generated/torch.Tensor.html

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