[sglang] SGLang의 NIXL 통신 최적화: Prep+Make API 도입을 통한 KV 캐시 전송 성능 향상

PR 링크: sgl-project/sglang#26406 상태: Merged | 변경: +314 / -135

들어가며

LLM 추론 시스템에서 분산 환경의 KV 캐시 전송은 전체 시스템의 TTFT(Time To First Token)를 결정짓는 핵심 병목 구간 중 하나입니다. 특히 SGLang의 NIXL(Network Inter-process Xfer Library)을 사용하는 환경에서는 매 청크(chunk) 전송마다 전송 디스크립터(transfer descriptor)를 매번 재구성하는 비용이 발생하고 있었습니다. 본 PR은 NIXL의 2단계 API인 prep과 make를 분리하여, 디스크립터를 미리 준비(prepare)하고 재사용함으로써 이 오버헤드를 제거하는 최적화를 수행했습니다.

코드 분석

1. 전송 디스크립터의 사전 준비 (conn.py)

기존에는 전송 시마다 디스크립터를 생성했으나, 이제는 _init_equal_tp_prep_handle 및 _init_hetero_tp_prep_handle 메서드를 통해 피어(peer)별로 디스크립터 리스트를 미리 빌드합니다.

Before (개념적):

# 매 전송마다 디스크립터 생성
self.agent.xfer_dlist(peer_name, build_dlist_on_the_fly(), "VRAM")

After:

# 디스크립터를 미리 준비하여 핸들 저장
self.prep_handles[peer_name] = self.agent.prep_xfer_dlist(
    peer_name, np.vstack(arrays), "VRAM"
)
# 이후 전송 시에는 make_prepped_xfer 호출

2. 슬라이스 전송을 위한 인덱스 매핑 최적화

TP(Tensor Parallelism) 크기가 다른 이기종(heterogeneous) 환경에서도 효율적인 전송을 위해 expand_page_indices_for_slice와 같은 유틸리티 함수를 도입했습니다. 이는 메모리 레이아웃을 미리 계산하여, 전송 시점에 복잡한 연산 없이 데이터를 바로 복사할 수 있도록 합니다.

def expand_page_indices_for_slice(...):
    # Dlist 레이아웃을 미리 계산하여 flat indices로 변환
    pair_stride = num_slots * page_size * num_groups
    # ... (연산 생략) ...
    return (pair_offsets[:, None] + within_pair[None, :]).ravel().astype(np.int32)

왜 이게 좋은가

성능 개선 수치

이번 최적화는 특히 대규모 프롬프트 처리 시 탁월한 성능 향상을 보여줍니다. p2d4 환경에서 64개 프롬프트 처리 시 기존 대비 약 27.89%의 TTFT 개선을 달성했습니다. 이는 디스크립터 생성이라는 반복적인 CPU 연산 비용이 제거되었기 때문입니다.

기술적 교훈

1. API 설계의 중요성: prep과 make를 분리하는 것은 반복적인 작업이 많은 시스템 프로그래밍에서 매우 강력한 패턴입니다. 설정(Setup)과 실행(Execution)을 분리함으로써 핫 패스(Hot path)에서의 오버헤드를 최소화할 수 있습니다.
2. 메모리 레이아웃 최적화: 데이터 전송 시 인덱스 계산을 런타임이 아닌 준비 단계로 옮김으로써, GPU 메모리 접근 패턴을 최적화하고 CPU-GPU 간의 동기화 지연을 줄였습니다.

리뷰 피드백 반영

리뷰어들은 TP 비율에 따른 로직 분리(_init_equal_tp_prep_handle vs _init_hetero_tp_prep_handle)가 코드 가독성과 유지보수 측면에서 더 낫다는 점에 동의했습니다. 내부 로직이 완전히 다르기 때문에 하나의 함수로 합치기보다 분리하는 것이 더 나은 설계라는 점이 강조되었습니다.

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