효율적이고 휴대 가능한 전문가 혼합 통신

당사는 표준 All-to-All 통신 기본 요소 대비 10배 더 빠른 성능을 달성하는 고성능 휴대용 오픈 소스 전문가 혼합(MoE) 통신 라이브러리를 소개합니다.

당사의 구현은 다음과 같은 MoE 통신 효율성을 제공하는 몇 가지 핵심 기술 혁신을 특징으로 합니다:

GPU 개시 통신(IBGDA): 직접적인 GPU-to-NIC 통신을 지원하여 CPU 관여를 우회함으로써 지연 시간을 크게 줄입니다.

통신 및 연산 중첩: 송신 및 수신 단계가 분리된 커널 아키텍처를 통해 네트워크 전송과 연산이 동시에 진행될 수 있습니다.

최고의 단일 노드 성능: 단일 노드 구성에서 기존 최고 구현 대비 2.5배 낮은 지연 시간 달성

효율적이고 휴대 가능한 다중 노드 성능: 당사의 구현은 표준 All-to-All 통신보다 최대 10배 더 빠른 속도를 달성합니다. 고도로 전문화된 구현보다는 약 2배 느리지만, 당사의 접근 방식은 NVSHMEM 버전 및 네트워크 환경(NVLink, CX-7 및 EFA) 전반에서 더 나은 휴대성을 제공합니다.

이 라이브러리는 완전한 오픈 소스이며 https://github.com/ppl-ai/pplx-kernels에서 이용 가능합니다.

본 문서에서는 대규모 MoE 모델에서의 전문가 병렬화 문제를 탐구하고 분산 GPU 환경 전반에서 효율적인 토큰 발송 및 결합을 위한 당사의 접근 방식을 설명합니다.

소개

DeepSeek R1 및 Mixtral 8x7B와 같은 전문가 혼합(MoE) 모델은 각 토큰에 대해 활성화되는 가중치 수를 제한함으로써 밀집형 모델을 개선합니다. MoE 모델은 레이어 수를 늘리거나 각 디코더 레이어의 다층 퍼셉트론(MLP) 내 선형 레이어 크기를 늘려 매개변수 수를 높이는 대신, 전통적인 MLP를 여러 전문가와 라우터로 대체합니다. 예를 들어, DeepSeek R1의 671B 매개변수 중 추론 중에 주어진 토큰과 곱해지는 매개변수는 37B에 불과합니다. 이는 밀집형 모델 대비 토큰을 디코딩하는 데 필요한 연산량을 크게 줄이고 지연 시간을 개선합니다.

MoE 모델은 밀집형 모델에 비해 추론과 관련하여 추가적인 문제를 제시합니다. 전문가 자체는 소규모 MLP 레이어이지만 각 디코더 레이어에는 토큰이 발송될 전문가를 결정하는 라우터가 포함되어 있으며, 각 토큰은 여러 전문가에게 발송됩니다. 라우터는 일반적으로 확률 분포를 생성하는 작은 선형 레이어입니다. 통상적으로 상위 K 점수를 받은 전문가가 선택되며, 최종 활성화는 라우터의 확률 분포에서 파생된 가중치를 곱한 전문가 출력을 가중 평균하여 계산됩니다.

분산 MoE 시스템의 지연 시간을 최소화하기 위해 여러 장치 전반에서 병렬화를 활용할 수 있으나, 이는 통신 문제를 야기합니다. Mixtral 8x7B 또는 Llama-70B와 같은 모델은 단일 노드의 8개 장치 내에 적합하며 빠르고 지연 시간이 짧은(최대 900Gbps) NVLink 상호 연결의 이점을 누립니다. 그러나 더 큰 모델은 InfiniBand(최대 400Gbps)를 사용하여 여러 노드에 전문가를 샤딩(sharding)해야 하며, 이는 추가적인 지연 시간 문제를 발생시킵니다.

확대

본 문서에서는 전문가 병렬화 문제를 탐구하고, 할당된 전문가가 있는 장치로 토큰을 효율적으로 발송하고 전문가 출력의 가중 합계를 계산하여 다시 수집하는 GPU 커널 구현을 설명합니다. 특정 네트워크 하드웨어를 대상으로 하는 적극적으로 최적화된 구현이 존재하지만, 당사는 최소한의 NVSHMEM 기본 요소에 의존하면서도 표준 All-to-All 통신보다 10배 더 빠른 성능을 제공하는 매우 휴대 가능한 구현을 제시합니다. 당사의 구현은 다양한 네트워크 환경 전반에서 우수한 휴대성을 유지하면서 단일 노드 구성에서 최첨단 성능을 달성합니다.

샤딩 및 병렬화

효율적인 추론을 위해서는 모델의 가중치를 장치 메모리에 유지하는 동시에, 활성화, KV 캐시 및 모델의 순방향 패스에 필요한 기타 버퍼를 위한 충분한 스페이스를 남겨두어야 합니다. 가장 성능이 뛰어난 모델도 가장 성능이 좋은 GPU의 용량을 초과하므로, 추론은 공동으로 모델의 가중치를 저장할 수 있는 여러 장치에 걸쳐 분산되어야 합니다. 가중치의 샤딩 방식에 따라 연산을 오프로드하고 장치를 동기화하기 위해 서로 다른 통신 및 연산 방식을 사용해야 합니다.

확대

첫 번째 그림에 예시된 전문가 병렬화(EP)는 전문가 연산만 병렬화합니다. 서로 다른 전문가가 가중치를 보유한 서로 다른 장치에 할당됩니다. 라우팅 후 토큰은 해당 장치로 전송되고 결과는 이후에 수집 및 누적됩니다. 라우팅의 복잡성은 모델의 다른 부분에서의 병렬화 정도에 따라 달라집니다. 다른 레이어를 복제하면 각 순위가 로컬로 복제된 라우팅 테이블에서 토큰을 선택할 수 있으므로 라우팅의 필요성이 완전히 제거될 수 있습니다. 그러나 하나의 순위만 라우팅을 실행하는 경우 토큰, 인덱스 및 가중치를 각 전문가에게 발송하기 위해 브로드캐스트가 필요합니다. 마지막으로, All-gather 또는 All-to-all 브로드캐스트는 모델의 실행을 계속하는 순위와 출력 토큰을 동기화합니다. 이러한 구현은 torch가 이미 필요한 기본 요소를 노출하므로 비교적 간단하지만, 전문가로부터 토큰을 수집하는 gather 연산과 누적을 융합함으로써 어느 정도 이점을 얻을 수 있습니다.

확대

전문가 전용 병렬화는 모델이 어텐션, 노름(norm) 및 샘플링과 같은 비전문가 레이어를 실행하는 동안 클러스터의 노드가 유휴 상태일 수 있으므로 이상적으로 확장되지 않습니다. 그러나 주로 어텐션인 이러한 레이어의 연산 또한 텐서 병렬화(TP)의 이점을 누릴 수 있습니다. 대부분의 모델은 멀티 헤드 어텐션에 의존하므로, 어텐션 헤드와 이에 해당하는 Q, K 및 V 프로젝션도 여러 레이어 간에 슬라이스를 복제하거나 수집하여 장치 전반에 샤딩할 수 있습니다. 어텐션이 모든 장치에 분산된 경우, All-gather가 활성화를 동기화하여 라우팅이 복제될 수 있게 하므로 효율적인 구현을 위해 전문가 병렬화 사례와 유사한 동기화 기본 요소가 필요합니다. 그러나 어텐션 헤드 수를 특정 임계값 미만으로 줄이면 수익이 감소하므로 이 수준의 병렬화에는 한계가 있습니다.

확대

모든 장치를 최대한 활용하고 최대 128개 또는 256개의 GPU에 이르는 매우 높은 수준의 전문가 병렬화를 지원하려면 데이터 병렬화(DP)가 필요합니다. 이 체계에서 모델이 분할된 장치는 요청을 동시에 처리하도록 그룹화되며, 로컬 그룹 전반에서 샤딩된 KV 캐시를 유지하고 어텐션을 계산합니다. 이러한 병렬 그룹의 여러 인스턴스가 전문가 평가에 협력하며 각 인스턴스는 전문가 가중치의 서로 다른 하위 집합을 호스팅합니다. 어텐션 헤드 수에 따라 그룹은 일반적으로 노드 내 통신이 더 빠르므로 전체 노드 크기까지 확장될 수 있습니다. 예를 들어 위 그림에서 하나의 DP 순위는 A와 B라는 두 요청을 독립적으로 서비스하며 요청에 대한 어텐션, 노름 및 기타 모든 부기(bookkeeping)를 처리합니다. 다른 DP 순위는 별도의 C 요청을 처리합니다. 그러나 첫 번째 노드는 전문가의 절반을 호스팅하고 다른 노드는 나머지 절반을 호스팅하므로, 토큰 라우팅 후 A와 B는 두 번째 노드로 전송될 수 있으며 그 반대도 마찬가지입니다. 이는 희소 통신 문제를 야기합니다. 즉, 각 장치가 다른 대상 순위로 서로 다른 수의 토큰을 전송할 수 있습니다. torch의 기존 기본 요소, 주로 all_to_all은 패딩이나 GPU-to-CPU 동기화 및 메타데이터 브로드캐스트가 필요할 수 있으므로 특별히 적합하지 않습니다. 통신을 효과적으로 구현하려면 라우팅 정보를 처리하고 CPU 개입 없이 장치에서 통신을 시작하는 사용자 정의 커널이 필요합니다. 라우팅 후 발송 커널은 토큰을 라우팅된 순위로 전송해야 하며, 결합 측에서는 현재 DP 그룹의 요청에 속하는 활성화를 수집해야 합니다. 또한 브로드캐스트를 최소화하고 후속 레이어를 위해 모든 순위에서 토큰을 올바르게 복제하기 위해 여러 장치가 있을 수 있는 DP 그룹 내에서 작업을 균형 있게 조정해야 합니다.

NVSHMEM

NVSHMEM은 NVIDIA 전용 OpenSHMEM 구현으로, 기본 하드웨어의 복잡성을 추상화하는 휴대용 장치 간 통신 기능을 제공합니다. API는 개별 전송 레이어의 기본 요소에 매핑되는 장치 간 읽기 및 쓰기를 표현할 수 있습니다. NVLink와 RDMA가 모두 지원되어 일정 수준의 휴대성을 제공하며 NVSHMEM 커널이 노드 내에서 또는 여러 노드의 클러스터 전반에서 작동할 수 있도록 합니다. 당사의 커널에서는 데이터를 전송하고 참여 장치를 동기화하기 위해 원격 메모리 쓰기 및 신호 작업을 사용합니다.

NVSHMEM 작업은 대칭 메모리 개념을 중심으로 구축됩니다. 즉, 장치 간 통신에 참여하는 모든 장치에 할당된 버퍼에서 작동합니다. 로컬 장치는 자체 로컬 버퍼 내에서 파생된 주소를 대상 순위의 인덱스와 함께 사용하여 작업의 대상을 지정할 수 있습니다. 아래 그림은 이 개념을 설명합니다. 두 GPU 모두 동일한 크기의 대칭 버퍼를 할당하고 src 포인터를 유지합니다. 첫 번째 순위는 3개의 정수를 버퍼의 시작 부분에 배치하여 두 번째 순위로 보내려고 합니다. nvshmem_int_put_nbi는 로컬 버퍼에서 시작 주소를 파생하고 대상 장치를 지정합니다. 두 번째 순위는 자체 버퍼에서 오프셋을 파생하여 첫 번째 장치로 하나의 요소를 전송하고 하나를 오프셋합니다. 대상 주소는 항상 nvshmem_alloc을 사용하여 할당된 대칭 버퍼여야 하지만, 소스 버퍼는 NVSHMEM에 사전 등록된 경우 장치 메모리의 임의 영역일 수 있습니다.

NVSHMEM은 광범위한 기본 요소를 제공하지만, 당사의 커널은 오직 3개의 함수에만 의존하며 그 위에 모든 동기화 및 펜싱을 구축합니다.

확대

nvshmemx_putmem_signal_nbi_warp: 데이터 블록을 한 장치에서 다른 장치로 전송하는 동시에 원격 장치에 플래그를 설정합니다. 이 작업은 64비트 위치를 설정(NVSHMEM_SIGNAL_SET)하거나 증가(NVSHMEM_SIGNAL_ADD)합니다. 전체 메모리 블록이 전송된 후 플래그가 업데이트됩니다. 원격 장치가 플래그의 변경 사항을 관찰하면 자체 메모리의 버퍼에 안전하게 액세스할 수 있습니다. 이 함수는 데이터 전송과 동기화를 통합하는 데 유용합니다.

nvshmemx_signal_op은 일반적으로 64비트 플래그인 단일 메모리 위치에서 작동하며 원자적으로 설정하거나 증가시킵니다. 메타데이터를 전송하고 장치를 동기화하는 데 유용합니다.

nvshmem_uint64_wait_until은 신호의 수신 측에서 사용되며, 원격 장치가 업데이트할 때까지 플래그를 폴링합니다.

nvshmem_quiet 또는 nvshmem_fence를 통해 작업을 정렬하는 것은 비용이 많이 들기 때문에, 당사의 커널 구현은 다른 NVSHMEM 기본 요소 사용을 피하고 개별 송수신 쌍 외의 종속성을 최소화합니다. 데이터를 전송할 때는 항상 함수의 비차단 버전을 사용하며 데이터가 로컬 순위에서 전송되기를 기다리지 않습니다. 커널 그룹 전반의 다른 암시적 동기화는 NIC가 버퍼에 액세스해야 할 때 이러한 버퍼가 파괴적으로 덮어쓰이지 않도록 보장합니다.

NVSHMEM은 GPUDirect RDMA를 사용하여 NIC가 CPU를 관여시키지 않고 데이터 교환을 위해 GPU 메모리를 직접 읽고 쓸 수 있도록 합니다. 또한, ConnectX NIC에서 NVSHMEM은 GPU 개시 통신(Infiniband GPUDirect Async 또는 IBGDA라고도 함)을 지원하여 GPU가 RDMA 작업을 NIC에 직접 게시하고 NIC의 완료 대기열을 직접 폴링할 수 있도록 합니다. GPU 개시 통신을 지원하지 않는 플랫폼에서는 NVSHMEM이 호스트 CPU에서 "프록시" 스레드를 생성하여 GPU를 대신하여 통신을 시작하고 완료를 폴링합니다. NVSHMEM 프로그램은 GPU 개시 통신 지원 여부와 관계없이 휴대 가능합니다. 그러나 GPU 개시 통신은 CPU로의 우회를 완전히 생략하므로 지연 시간을 크게 단축합니다.

휴대용 커널

당사는 발송 및 결합 커널 쌍을 통해 MoE 통신을 구현합니다. 발송 커널은 각 순위에서 토큰과 라우팅 정보를 읽고 적절한 전문가에게 발송하는 역할을 합니다. 결합 커널은 전문가가 생성한 활성화를 수집하여 소스 순위로 다시 전송하는 동시에 라우터가 계산한 가중치를 기반으로 선택된 전문가로부터 가중 평균을 계산합니다. 커널은 데이터 전송이 연산과 중첩될 수 있도록 송신 및 수신 구성 요소로 추가 분할됩니다. 송신 커널은 모두 비차단 및 비동기적이며 단순히 원격 장치로의 모든 쓰기를 발송합니다. 반대편에서 결합 커널은 메모리에서만 읽으며 모든 필수 데이터가 전송될 때까지 기다립니다. NIC에 작업을 발송한 후 데이터가 비동기적으로 네트워크를 통해 전송되는 동안 GPU는 공유 전문가를 적용하거나 어텐션을 계산하는 등 다른 유용한 로컬 작업을 수행할 수 있습니다.

각 커널 쌍은 자체 대칭 메모리 버퍼를 가지며, 각 순위는 다른 모든 순위로부터 데이터를 수신하기 위해 개인 저장소를 할당합니다. DP 그룹에서 발송될 수 있는 최대 토큰 수는 순위 전반에서 고정되어 있으며, 이는 각 순위가 각 DP 순위로부터 각 로컬 전문가에 대해 수신할 수 있는 상한선을 설정합니다. 이를 통해 송신 순위는 순위 간 동기화 없이도 대상 순위의 고유 주소를 파생하여 쓸 수 있습니다. 데이터가 수신된 후 커널은 토큰을 섞어 수신 커널에 필요한 최적의 형식으로 메모리에 배치합니다. 버퍼는 상당한 치수를 가지지만 모델의 모든 순차 레이어 전반에서 재사용됩니다.

확대

위 그림에 예시된 유일한 형태의 전역 동기화는 결합-수신 커널에서 구현됩니다. 송신 측에서 데이터가 발송되면, 순위는 다른 순위로부터 필요한 모든 토큰을 수신하면 실행을 계속하고, 처리되는 즉시 결합-송신 커널로 진입합니다. 결합-수신에서의 배리어는 어떤 순위도 앞서 나가서 다른 순위가 여전히 데이터를 수신하기를 기다리는 동안 발송-송신을 시작하여 잠재적으로 파괴적인 중첩을 일으키지 않도록 보장합니다. 단순성을 위해 결합 커널에서 동기화가 수행되며, 최적화로서 동일한 배리어가 발송-수신과 결합-송신 커널 사이에 구현될 수 있으며, 이는 GPU가 다른 커널을 실행하는 동안 NIC가 전송을 안정화할 시간을 더 많이 제공하여 지연 시간을 추가로 줄일 수 있습니다.

두 커널 모두 장치의 모든 가용 SM 전반에 분할됩니다. 발송-송신 및 결합-수신 커널은 순위당 최대 토큰 수(max_m) 전반에 병렬화되어야 하지만, 발송-수신 및 결합-송신 커널은 각 순위 및 로컬 전문가로부터 최대 max_m * (num_experts // EP) * (EP // TP)개의 토큰을 처리해야 합니다.

발송

발송 커널은 토큰과 블록 스케일링 요소를 토큰이 라우팅된 전문가를 호스팅하는 피어에게 발송하는 역할을 합니다. 송신 측은 워프 특수화(warp specialization)에 의존하여 두 가지 작업을 병렬화합니다. 각 순위에서 전송되는 토큰 수를 결정하기 위한 라우팅 정보 집계와 원격으로 발송될 메시지에 토큰 패킹이 그것입니다. 수신 측은 먼저 모든 토큰 수가 수신되기를 기다린 다음 공유 메모리에서 페이로드를 읽고 다운스트림 그룹화된 GEMM 커널로 전달될 텐서로 언패킹합니다. 또한 각 토큰을 어디로 다시 보낼지 나타내기 위해 결합 커널과 공유되는 버퍼에 정보를 저장합니다. 수신 커널이 비결정적 순서로 연속 버퍼에서 토큰을 섞기 때문에 이 정보가 필요합니다.

확대

송신부에서는 각 블록의 32개 스레드로 구성된 단일 워프가 라우팅 정보를 읽고 각 대상 전문가가 이 순위로부터 수신할 것으로 예상되는 토큰 수를 계산하여 nvshmem_signal_op을 사용하여 수 플러스 1을 전송합니다. 원격 측에서 0에서 0이 아닌 값으로의 전환은 소스 순위로부터의 수신을 의미하므로 카운트는 1만큼 증가합니다. 병렬로 나머지 워프는 협력하여 모든 블록에서 토큰을 대칭 메모리로 복사하고 활성화, 스케일링 요소 및 로컬 순위에서의 인덱스를 연속적인 청크로 패킹합니다. 이 인덱스는 결합 송신자가 토큰이 기록될 주소를 결정하는 데 필요합니다. 다음으로, 워프 그룹 전반의 배리어를 통해 모든 데이터가 복사되었는지 확인한 후, 워프는 다시 독립적으로 작동하며 각기 다른 전문가에게 동일한 버퍼를 병렬로 전송합니다. 토큰은 nvshmemx_putmem_signal_nbi_warp를 사용하여 전송되며, 이 함수는 원격 장치에서 로컬 순위로부터 전송된 토큰 수를 원자적으로 증가시킵니다. DP 그룹 내에서는 각 순위가 전송될 토큰의 복제본을 소유하므로 각 장치가 토큰의 하위 집합을 전송하여 발송이 균등하게 분산됩니다.

수신 측에서는 모든 커널이 먼저 총 토큰 수가 전송되기를 기다려 송신 순위와 동기화합니다. 그 후, putmem 호출의 의미론 덕분에 원자적으로 증가된 전송 토큰 수가 소스 순위로부터의 모든 페이로드도 전송되었음을 나타내는 총 수에 도달할 때까지 기다립니다. 커널은 모든 블록과 스레드 전반에 걸쳐 작업을 병렬화하면서 nvshmem_uint64_wait_until를 사용하여 카운트를 폴링합니다. 이후, 블록 간 배리어는 모든 데이터가 올바르게 수신되지 않는 한 어떤 블록도 버퍼를 읽지 않도록 보장합니다. 블록 전반에 분산되고 원자적으로 증가하는 카운터를 통해 동기화된 토큰은 대칭 메모리에서 전문가를 구현하는 다운스트림 커널의 입력이 되는 일반적인 연속 버퍼로 복사됩니다. 소스 순위, 전문가 인덱스 및 토큰 인덱스는 결합 커널이 활성화를 전송해야 하는 위치를 정확히 지정하여 별도로 저장됩니다. DP 그룹 내의 토큰은 서로 다른 장치에서 전송되지만, 모두 그룹화되어 해당 전문가에게 전달됩니다.

결합

결합 커널은 다시 송신 및 수신 절반으로 분할됩니다. 송신자는 비양자화된 16비트 활성화를 원격 장치로 복사하고, 수신 커널은 데이터가 전송되기를 기다려 모든 전문가 기여분의 가중 평균을 로컬에서 계산합니다. 또한, 각 순위는 송신 커널 진입 시 각 피어에 플래그를 설정하고 수신 커널은 플래그가 설정된 것을 확인하지 않으면 반환할 수 없도록 하여 발송-결합 시퀀스를 동기화하는 배리어 역할도 합니다. 동기화 지연 시간은 실제 통신 및 연산과 중첩되므로 최소화됩니다.

확대

송신 측에서 커널은 모든 로컬 전문가에게 할당된 토큰 목록을 병렬로 탐색하여 대상 순위의 버퍼에 기록합니다. 대상 순위, 전문가 인덱스 및 토큰 인덱스는 토큰 수신 시 scater 커널에 의해 채워진 토큰별 버퍼에서 읽습니다. 각 송신자는 위 그림에 표시된 대로 전문가당 고유한 개인 메모리 영역을 기록하여 동기화할 필요가 없습니다. 발송과 유사하게, 결합은 대상 순위에서 토큰별 카운터를 원자적으로 증가시켜 데이터 수신을 나타냅니다. 카운터가 토큰이 발송된 전문가 수와 일치하면 토큰 콘텐츠에 액세스할 수 있습니다.

수신 커널에서는 신호 작업에 의해 설정된 플래그를 폴링하여 콘텐츠가 도착하기를 기다리며 여러 블록 전반에 걸쳐 토큰 목록을 병렬로 탐색합니다. 도착 시 라우팅 테이블이 어떤 버퍼를 읽을지, 각 전문가에게 어떤 가중치를 할당할지 나타내어 개인 버퍼에서 페이로드를 읽습니다. 그런 다음 결과는 외부에서 할당된 텐서에 기록되며, 모든 장치가 배리어를 통과하면 커널 실행이 종료됩니다. DP 그룹 전반에서 모든 순위는 각 전문가 활성화의 복사본을 수신하여 자체 복제본을 계산합니다.

벤치마크

당사는 InfiniBand 및 CX-7 NIC를 통해 연결된 16개 노드 전반의 128개 GPU 클러스터에서 커널을 평가합니다. GPUDirect Async(IBGDA)와 CPU 프록시를 사용한 신뢰할 수 있는 연결(RC)을 모두 시도합니다. 이를 고도로 최적화된 DeepSeek 구현 및 PyTorch(NCCL을 통해) 또는 NVSHMEM에서 제공하는 밀집형 기본 요소와 비교합니다.

확대

밀집형 NVSHMEM All-to-All 작업에서 IBRC와 IBGDA의 성능은 비슷하지만(6378 µs 대 6180 µs), 희소 커널에서는 IBGDA가 훨씬 빠릅니다. 발송과 결합을 합치면 IBGDA는 902 µs를 사용하는 반면 IBRC는 3223 µs가 소요되어 지연 시간이 3.6배 개선됩니다. All-to-All은 대역폭 제한을 받지만 희소 커널은 훨씬 적은 데이터를 브로드캐스트하며 지연 시간 제한을 받습니다. CPU 프록시가 GPU와 NIC를 조정할 필요 없이 GPU에서 직접 네트워크 전송을 트리거함으로써 종단 간 지연 시간이 크게 단축됩니다.

휴대용 커널은 고도로 최적화된 DeepSeek 커널보다 약 2배 느리지만, 밀집형 커널 대비 지연 시간을 10배 개선합니다(902 µs 대 9944 µs).

또한 라이브러리 기본 요소와 달리 송신자-수신자 구성 요소로의 분할을 통해 지연 시간의 일부를 숨길 수 있습니다.

확대

단일 노드(EP8)에서 NVSHMEM은 NVLINK를 전송에 사용하여 노드 간 네트워킹보다 낮은 지연 시간과 높은 처리량을 제공합니다. 당사의 휴대용 커널은 단일 노드에서 DeepEP보다 약 2.5배 빠릅니다(186µs 대 481 µs).

추가 개선 사항

여기에 설명된 커널은 ML 프레임워크의 내장 기본 요소를 능가하며 NVLink 또는 InfiniBand와 같은 특정 장치 간 전송을 과도하게 전문화하지 않고도 괜찮은 성능을 제공합니다. 이미 언급된 기회 외에도 통신 기본 요소를 더 전문화된 버전으로 대체함으로써 추가적인 성능 향상을 달성할 수 있습니다. 예를 들어 NVLink 전반에서 대칭 메모리 사용을 장치 간 공유 버퍼로 대체하여 일부 복사를 제거하고 토큰 배치 대신 개별 토큰 전반에서 더 세분화된 동기화 체계를 허용할 수 있습니다. InfiniBand 전반에서 구현은 예상 토큰 수와 전송된 토큰 수를 모두 전송하고 토큰이 도착할 때 대기열에서 출력 버퍼로 토큰을 플러시할 필요를 없애고 기본 대기열 쌍에 직접 액세스할 수 있습니다. 그러나 이러한 구현은 휴대성을 희생해야 하며, 당사의 구현을 통해 특정 하드웨어를 최적화하기 위한 개발 노력을 기울이기 전에 특정 구성의 클러스터를 평가할 수 있습니다.

결론

당사는 다양한 하드웨어 구성 전반에서 호환성을 유지하면서 표준 All-to-All 통신 대비 10배 더 빠른 성능을 달성하는 MoE 통신을 위한 고성능 휴대용 라이브러리를 제시했습니다. NVLink를 사용하는 단일 노드 배포에서 당사의 솔루션은 이전 구현 대비 2.5배 낮은 지연 시간을 보여줍니다.

당사의 접근 방식은 GPU 개시 통신 지원, 연산-통신 중첩을 가능하게 하는 분할 커널 아키텍처, 최소한의 NVSHMEM 기본 요소를 사용한 효율적인 토큰 발송을 포함한 핵심 혁신을 통해 성능과 휴대성의 균형을 맞춥니다. 다중 노드 설정에서 고도로 전문화된 구현보다는 약 2배 느리지만, 당사의 라이브러리는 다양한 네트워크 환경(NVLink, CX-7 및 EFA) 전반에서 우수한 유연성을 제공합니다.

MoE 모델이 계속 확장됨에 따라 당사와 같은 효율적인 통신 전략은 실제 배포에 점점 더 중요해질 것입니다. 당사의 완전한 오픈 소스 구현은 https://github.com/ppl-ai/pplx-kernels에서 이용 가능합니다.

참고 자료

DeepEP

DeepSeek-V3 기술 보고서

NVIDIA NVSHMEM

인용

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}