[sglang] HiSparse 도입: Sparse Attention 모델을 위한 효율적인 KV 캐시 관리

PR 링크: sgl-project/sglang#20343 상태: Merged | 변경: +1692 / -59

들어가며

최근 DeepSeek-V3.2나 GLM-5와 같은 모델들은 NSA(Non-uniform Sparse Attention) 메커니즘을 채택하여 효율적인 추론을 가능하게 합니다. 하지만 이러한 모델들은 방대한 KV 캐시를 필요로 하며, GPU 메모리 제한으로 인해 배치 사이즈를 늘리는 데 한계가 있습니다. SGLang의 이번 PR은 HiSparse라는 새로운 기법을 도입하여, 유휴(idle) KV 캐시를 CPU 메모리로 오프로딩(offloading)함으로써 GPU 메모리 효율을 극대화하고 처리량을 개선합니다.

코드 분석

1. GPU 커널 최적화: python/sglang/jit_kernel/csrc/hisparse.cuh

HiSparse의 핵심은 GPU와 CPU 메모리 간의 데이터를 효율적으로 전송하고, 캐시 적중(hit) 여부를 빠르게 판단하는 커널입니다. load_cache_to_device_buffer_kernel은 각 요청에 대해 토큰의 위치를 파악하고, 필요한 데이터를 디바이스 버퍼로 로드합니다.

Before/After 개념: 기존에는 모든 KV 캐시가 GPU 메모리에 상주해야 했으나, HiSparse는 다음과 같이 LRU(Least Recently Used) 기반의 캐시 교체 로직을 GPU 커널 수준에서 수행합니다.

// After: LRU 기반의 효율적인 캐시 관리 및 토큰 위치 매핑
if (is_hit) {
  int hit_offset = s_chunk_offset[chunk_idx] + local_hit_offset;
  s_lru_slots_out[hit_offset] = buf_slot;
}
if (is_evictable) {
  int evict_offset = s_evict_chunk_offset[chunk_idx] + local_evict_offset;
  s_lru_slots_out[HOT_BUFFER_SIZE - 1 - evict_offset] = buf_slot;
}

이 커널은 atomicCAS를 사용한 해시 테이블을 통해 토큰 ID를 빠르게 조회하며, __ballot_sync와 warp_inclusive_scan을 사용하여 워프(warp) 단위의 병렬 처리를 극대화했습니다.

2. 동기화 문제 해결: python/sglang/srt/managers/hisparse_coordinator.py

리뷰 과정에서 cudaErrorIllegalAddress 문제가 보고되었습니다. 이는 비동기 스트림 간의 레이스 컨디션 때문이었습니다. 이를 해결하기 위해 wait_stream을 명시적으로 추가하여 데이터 정합성을 보장했습니다.

# After: 스트림 동기화 추가로 레이스 컨디션 방지
def admit_request_into_staging(self, req: Req) -> None:
    if self.decode_producer_stream is not None:
        device_module.current_stream().wait_stream(self.decode_producer_stream)

왜 이게 좋은가

1. 메모리 효율성: GPU 메모리 부족으로 인해 포기해야 했던 대규모 배치 처리가 가능해집니다. CPU 메모리를 활용하여 KV 캐시를 확장함으로써 더 많은 요청을 동시에 처리할 수 있습니다.
2. 성능 최적화: 커널 수준에서 LRU 캐시 교체 로직을 구현하여, CPU-GPU 간 데이터 이동 오버헤드를 최소화했습니다.
3. 교훈: 복잡한 비동기 시스템(특히 event_loop_overlap 환경)에서는 스트림 간의 명시적인 동기화(wait_stream)가 필수적임을 보여줍니다. 또한, 커널 내에서 shared memory와 warp-level primitives를 적절히 활용하면 성능을 크게 향상시킬 수 있습니다.

리뷰어 피드백 반영

리뷰어들은 staging이라는 용어의 모호함과 테스트 커버리지 부족을 지적했습니다. 또한, 현재 HiSparse는 --disable-radix-cache 옵션과 함께 사용해야 한다는 제약 사항이 명확히 공유되었습니다. 향후 더 정교한 테스트와 기존 코드베이스와의 통합이 과제로 남아 있습니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#warp-level-primitives
• https://docs.sglang.ai/

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