[sglang] HiCache 최적화: TMA를 활용한 Host-Device KV 캐시 전송 성능 2배 향상

PR 링크: sgl-project/sglang#40278 상태: Merged | 변경: +1060 / -32

들어가며

LLM 추론 엔진인 SGLang의 HiCache는 캐시된 프리픽스(prefix)를 호스트의 pinned memory에서 GPU 디바이스 메모리로 효율적으로 이동시키는 핵심 컴포넌트입니다. 기존의 레지스터 기반 JIT 커널은 바이트 전송 시 레지스터를 거쳐야 하는 구조적 한계로 인해, NVLink-C2C와 같은 고속 인터커넥트의 대역폭을 충분히 활용하지 못했습니다. 본 PR은 TMA(Tensor Memory Accelerator)를 활용한 새로운 커널 hicache_tma.cuh를 도입하여, H2D/D2H 전송 성능을 약 2배 향상시켰습니다.

코드 분석

1. TMA 기반의 비동기 전송 (hicache_tma.cuh)

기존 커널은 각 스레드가 레지스터에 데이터를 보관하며 전송을 수행했으나, 새로운 커널은 cp.async.bulk를 사용하여 레지스터를 거치지 않고 공유 메모리(Shared Memory)로 직접 데이터를 이동시킵니다.

Before (Register-based):

// 각 스레드가 레지스터에 64B를 들고 대기
// 데이터 전송이 레지스터 점유율에 의존하여 병목 발생

After (TMA-based):

// cp.async.bulk를 사용하여 레지스터를 통하지 않고 공유 메모리로 직접 전송
SGL_DEVICE void bulk_g2s(void* dst_smem, const void* src_gmem, uint32_t bytes, uint64_t* bar) {
  asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" 
               :: "r"(to_shared(dst_smem)), "l"(src_gmem), "r"(bytes), "r"(to_shared(bar)) : "memory");
}

2. 인덱스 프리페치 및 2D 텐서 매핑

단순 연속 데이터뿐만 아니라, 페이지 단위로 분산된 데이터도 cuTensorMapEncodeTiled를 통해 2D 박스 형태로 한 번에 전송하여 TMA 엔진의 효율을 극대화했습니다.

// 2D tiled tensor-map box를 활용한 효율적인 데이터 로드
SGL_DEVICE void bulk_tensor_2d_g2s(void* dst_smem, const CUtensorMap* map, int32_t x, int32_t y, uint64_t* bar) {
  asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" 
               :: "r"(to_shared(dst_smem)), "l"(map), "r"(x), "r"(y), "r"(to_shared(bar)) : "memory");
}

왜 이게 좋은가

1. 대역폭 극대화: GB300 환경에서 H2D 전송 속도가 97 GB/s에서 192 GB/s로, D2H는 93 GB/s에서 183 GB/s로 약 2배 향상되었습니다. 이는 호스트 링크의 물리적 한계치에 근접한 수치입니다.
2. SM 점유율 최적화: 레지스터 기반 커널은 6개 이상의 CTA가 필요했던 반면, TMA 커널은 단 4개의 CTA만으로 동일한 대역폭을 달성하여, forward pass 연산에 더 많은 SM 자원을 할당할 수 있게 되었습니다.
3. 교훈: 데이터 전송 시 레지스터를 경유하는 것은 고대역폭 환경에서 치명적인 병목이 됩니다. 하드웨어 가속기(TMA)를 활용해 데이터 경로를 공유 메모리로 직접 연결하고, mbarrier를 통해 비동기적으로 동기화하는 패턴은 현대 GPU 아키텍처(Hopper 이상)에서 필수적인 최적화 기법입니다.

리뷰 피드백

리뷰어들은 hicache 그룹에 대한 반복적인 테스트 실행을 통해 커널의 안정성을 검증했습니다. 특히 SGLANG_HICACHE_TMA_TRANSFER 환경 변수를 통해 기존 커널과 새 커널 간의 전환을 유연하게 처리한 점이 아키텍처 호환성 측면에서 높게 평가되었습니다.

참고 자료

• https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-copy-instructions-cp-async-bulk
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/#tma-operations

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