Perplexity, MoE 통신 라이브러리 오픈소스 공개: 최대 10배 속도 향상

pplx-kernels는 Perplexity가 오픈소스로 공개한 MoE(Mixture-of-Experts) 전용 초고속 GPU 통신 라이브러리입니다. 대규모 MoE 모델을 여러 GPU에 분산 배치해 추론할 때 발생하는 통신 병목을 해결하며, 표준 All-to-All 통신 대비 최대 10배 빠른 전송 속도를 냅니다.

DeepSeek R1이나 Mixtral 같은 MoE 모델을 운영할 때 가장 큰 지연을 유발하는 요소는 연산량 자체보다 GPU 간 통신 대기 시간입니다. DeepSeek R1의 경우 전체 671B 파라미터 중 토큰당 37B만 활성화되지만, 활성화된 전문가 레이어가 여러 노드의 GPU에 흩어져 있어 장치 간 잦은 데이터 교환이 필수적입니다.

무엇을 해결하는가

기존 PyTorch 환경에서는 토큰 분배에 주로 기본 all_to_all 원어를 사용해 왔습니다. 하지만 이런 방식은 데이터 패딩이 발생하기 쉽고, CPU와 GPU 간 동기화 및 메타데이터 브로드캐스트로 인해 많은 시간을 낭비합니다.

pplx-kernels는 토큰을 해당 전문가가 위치한 장치로 보내는 디스패치와 결과를 가중합으로 취합하는 컴바인 과정을 커스텀 커널로 최적화했습니다.

단일 노드 환경에서는 기존 최고 성능 구현체보다 지연 시간을 2.5배 단축했습니다. 여러 노드로 구성된 분산 환경에서도 표준 All-to-All 대비 최대 10배 빠른 통신 성능을 달성합니다.

구현 원리

이 라이브러리의 핵심은 CPU 개입을 원천 배제하고 통신과 연산을 겹쳐서 실행하는 구조입니다.

ConnectX 네트워크 카드 환경에서 GPU 주도 통신(IBGDA)을 지원합니다. GPU가 호스트 CPU를 거치지 않고 NIC에 직접 RDMA 작업을 요청하고 완료 큐를 감시하므로 지연 시간이 대폭 줄어듭니다.

또한 전송과 수신 단계를 분리한 분할 커널 구조를 적용했습니다. 네트워크를 통해 데이터가 오가는 동안 GPU는 로컬 연산을 동시에 수행해 전송 대기 시간을 숨깁니다.

NVSHMEM 라이브러리 활용 역시 단순화했습니다. 비용이 큰 전역 동기화나 펜스 작업을 피하고, 단 세 개의 비차단 원어 함수만으로 데이터 전송과 신호 전달을 처리해 장치 간 종속성을 최소화했습니다.

호환성과 이용 방법

특정 전용 하드웨어에 극단적으로 맞춘 비공개 구현체보다는 약 2배 느리지만, pplx-kernels는 폭넓은 이식성을 핵심 무기로 삼았습니다.

단일 노드 내 900Gbps NVLink 통신은 물론, 노드 간 400Gbps InfiniBand, ConnectX-7, AWS EFA 네트워크 환경을 두루 지원합니다.

현재 전체 코드가 오픈소스로 공개되어 있으며, GitHub의 ppl-ai/pplx-kernels 저장소에서 직접 다운로드받아 빌드할 수 있습니다.

왜 중요한가

DeepSeek R1 등 대형 오픈소스 MoE 모델을 자체 인프라에서 서빙해야 하는 엔지니어링 팀에게 분산 통신 지연 시간 단축은 곧바로 서비스 비용 절감과 직결됩니다.

pplx-kernels는 특정 벤더의 전용 인프라에 갇히지 않으면서도 표준 네트워크 환경에서 즉각 적용할 수 있는 고성능 통신 해법을 제공합니다. 기존 프레임워크의 비효율적인 통신 루틴을 대체해 추론 속도를 끌어올릴 가치가 충분합니다.