AWS EFA에서 조 단위 매개변수 모델 구현

Perplexity에서는 제품, API 및 연구 팀을 위해 최상의 모델을 사용합니다. Kimi-K2와 같은 대규모 오픈 소스 Mixture-of-Experts 모델은 8x NVIDIA H200 GPU를 탑재한 가장 큰 추론 노드에서도 효율적으로 수용할 수 없어 다중 노드 배포가 필요하기 때문에 특별한 과제를 안겨줍니다. 저희는 ConnectX-7에서 최첨단 지연 시간을 달성하여 DeepEP의 성능을 능가하는 전문가 병렬 처리를 위한 일련의 커널을 선보입니다. 동일한 커널은 또한 AWS EFA(Elastic Fabric Adapter)에서 실행 가능한 지연 시간을 최초로 달성하여 조 단위 매개변수 모델 배포를 가능하게 합니다.

GitHub에서 저희 커널을 사용해 보시고 arXiv에서 전체 연구 논문을 읽어보십시오.

Mixture-of-Experts(MoE)는 수천억 또는 조 단위의 매개변수로 확장하는 동시에 추론 시 합리적인 지연 시간을 제공하는 표준 아키텍처로 부상했습니다. MoE 모델은 트랜스포머의 밀집 계층을 전문가 세트와 토큰에 곱할 전문가를 결정하는 작은 선형 라우팅 계층으로 대체합니다. 동일한 전문가와 관련된 토큰을 그룹화함으로써, 이제 여러 노드에 분산된 서로 다른 GPU에서 각기 다른 전문가를 호스팅할 수 있게 되어 이 문제는 희소해지고 사소하게 병렬화할 수 있게 됩니다.

torch.distributed 및 NCCL과 같은 집단 통신 라이브러리를 사용하여 쉽게 구현할 수 있는 텐서 병렬 처리(TP) 및 데이터 병렬 처리(DP)와 같은 다른 형태의 병렬 처리와 달리, MoE 라우팅은 희소한 P2P 통신을 포함합니다. 디스패치 커널은 들어오는 토큰 세트를 분할하여 전문가를 호스팅하는 랭크로 전달해야 하며, 이와 대응되는 결합 커널은 처리된 토큰을 수신하고 가중 평균을 계산해야 합니다. 이러한 전송은 집단 통신에 매핑되지 않으며 낮은 지연 시간을 달성하기 위해 사용자 지정 커널이 필요합니다.

Perplexity에서는 휴대용 노드 간 커널과 특수 노드 내 커널을 모두 개발해 왔으며, 이전 MoE 통신 관련 기사에서 저희 작업을 더 깊이 있게 소개했습니다. 또한 다른 블로그 게시물을 통해 다중 노드 배포가 대규모 MoE 모델에 대해 더 낮은 지연 시간과 더 높은 처리량을 가져올 수 있음을 입증했습니다.

이 기사에서는 NVSHMEM 기반의 이전 작업 성능을 크게 개선한 새로운 노드 간 커널 구현에 대해 자세히 살펴봅니다. ConnectX-7에서 최첨단 결과를 선보일 뿐만 아니라, 추론 런타임 내에서 커널을 사용하여 가장 큰 오픈 소스 모델을 효과적으로 서비스할 수 있음을 증명합니다.

텐서 병렬 처리와 전문가 병렬 처리는 NVLink가 마이크로초 단위의 지연 시간으로 높은 처리량(900GB/s)을 제공하므로 단일 노드 내에서 효과적으로 확장될 수 있습니다. NVIDIA GB200 및 GB300 시스템은 조 단위 매개변수 모델의 가중치와 KV 캐시를 보유할 충분한 메모리를 갖추고 있지만, 최대 8개의 H200 GPU를 탑재한 AWS p5en 인스턴스에서의 배포는 1120GB의 HBM으로 제한되며, 이는 모델 가중치와 KV 캐시 모두에 분할되어야 합니다. 결과적으로 특정 모델은 InfiniBand를 통한 노드 간 배포가 필요합니다.

InfiniBand는 여전히 매우 빠르지만 400Gbps(50GB/s)의 처리량만 제공하므로 MoE 라우팅에 수십에서 수백 마이크로초의 지연 시간이 추가됩니다. 다행히도 공유 전문가, 마이크로 배치 및 계산-통신 중첩을 사용하여 추가된 지연 시간을 숨김으로써 경쟁력 있는 비용으로 실행 가능한 처리량을 제공할 수 있습니다. DeepSeek는 H800 GPU 클러스터 및 NVIDIA ConnectX-7 InfiniBand 어댑터에 배포된 DeepSeek-R1을 통해 이러한 기술을 개척했습니다.

AWS 클러스터의 경우, 자체 Elastic Fabric Adapter(EFA)를 사용하기 때문에 상황이 다릅니다. 이들은 일반적으로 훈련 워크로드에 사용되는 집단 통신에서 400Gbps의 피크 처리량을 제공하지만, MoE 디스패치 및 결합 중에 교환되는 메시지 크기에서는 ConnectX-7에 약간 미치지 못합니다. 또한 EFA는 GPUDirect Async를 지원하지 않으므로 전송을 시작하기 위해 GPU와 NIC를 연결하는 CPU 프록시 스레드가 필요합니다. 추가적인 복잡성 외에도, 추가 PCIe 트랜잭션은 모든 전송에 수 마이크로초의 오버헤드를 더합니다. NVSHMEM과 같은 프레임워크를 사용하면 커널 작성자가 장치 독립적인 구현을 제공할 수 있지만, 일반 API는 비용 및 지연 시간 측면에서 배포를 비효율적으로 만드는 오버헤드를 발생시킵니다.

이전 커널은 NVSHMEM을 통해 기본 전송 구현을 추상화했습니다. 이들은 nvshmemx_putmem_signal_nbi_warp에 의존하여 원자 카운터를 사용해 피어와 동기화하며 토큰을 하나씩 전송했습니다. 이는 IBGDA가 있는 ConnectX-7 NIC에서는 합리적인 성능을 제공했지만, 프록시 기반의 ConnectX-7 IBRC 구현은 훨씬 느렸고, EFA 구현은 1밀리초 미만의 지연 시간을 제공하지 못했습니다.

저희는 이 커널을 기본 하드웨어에 대한 테스트베드로 사용하면서 모델에 노출할 가장 적합한 인터페이스를 결정했습니다. 저희는 커널과 NIC 또는 프록시 스레드 간의 빈번한 전송 및 상호 작용이 성능을 저해한다는 가설을 세웠습니다. 또한, 저희 구현은 신뢰성을 제외하고 하드웨어의 어떠한 보장에도 의존하지 않았으며, 동기화 기본 요소의 오버헤드를 줄여 성능을 높일 수 있음을 시사했습니다.

이러한 관찰에 따라 저희는 동일한 인터페이스를 노출하는 EFA 기반 커널을 고안하여, 최고의 성능을 위해 장치 커널과 호스트 프록시 간의 상호 작용을 최적화했습니다.

호스트 프록시를 사용한 RDMA 및 NVLink 하이브리드 통신으로, 시작 및 수신 랭크, 프록시 워커, NIC를 보여줍니다.

새로운 커널은 GPU 커널이 장치의 모델과 상호 작용하고 호스트 CPU에서 실행되는 프록시 스레드가 NIC와의 상호 작용을 처리하는 하이브리드 CPU-GPU 아키텍처를 따릅니다. 새로운 커널은 KV 캐시 전송을 위해 구축한 동일한 TransferEngine을 재사용합니다.

MoE 라우팅은 디스패치 커널에 의해 수행되며, 이 커널은 모델로부터 라우팅할 토큰과 토큰별 전문가를 수신하여 전문가를 호스팅하는 랭크로 네트워크를 통해 전송합니다. 그룹화된 GEMM 커널이 토큰과 전문가 가중치를 곱한 후, 결합 커널은 토큰을 소스 랭크로 전송하고 가중 평균을 계산한 후 모델로 반환합니다.

디스패치 및 결합 커널은 모두 중첩 및 마이크로 배칭을 용이하게 하기 위해 전송자 및 수신자 절반으로 분할됩니다. 전송자는 토큰을 전송 버퍼로 섞어 단일 RDMA 쓰기 작업을 통해 피어에 쓸 수 있도록 합니다. 전송자를 실행한 후, GPU는 백그라운드에서 네트워크 전송이 대기 중인 동안 다른 작업(예: 공유 전문가 곱셈)을 수행할 수 있습니다. 수신자 커널은 모든 전송이 수신될 때까지 차단됩니다. 백그라운드에서는 호스트 측 작업자 스레드가 전송을 처리하며, 전송 버퍼가 채워지면 GPU로부터 신호를 받아 RDMA 쓰기를 디스패치합니다. NIC가 완료 신호를 보내면, 수신자를 차단 해제하고 수신자는 토큰을 모델로 반환합니다.

커널은 호스트와 긴밀하게 결합되어 통합 메모리와 GDRCopy를 통해 서로 통신합니다. CUDA 통합 메모리를 사용하면 장치 메모리나 호스트 메모리를 둘 다의 주소 스페이스에 매핑하여 PCIe를 통한 통신을 투명하게 처리할 수 있습니다.

통합 메모리는 대량 전송에 적합하지만, 짧은 지연 시간의 폴링은 GDRCopy를 통해 수행됩니다. GDRCopy는 GPUDirect RDMA를 사용하여 마이크로초 지연 시간으로 GPU 메모리를 읽고 씁니다.

RDMA 외에도 노드 내에서 커널은 토큰 교환을 위해 NVLink를 활용합니다. 이는 트래픽의 상당 부분(EP32의 경우 최대 1/4, EP64의 경우 최대 1/8)을 오프로드하여 전반적으로 더 낮은 지연 시간을 달성할 수 있습니다.

디스패치 커널을 위한 전송 버퍼 설계는 호스트 프록시와 관련된 실행 가능한 지연 시간을 달성하는 데 중요합니다. N개의 랭크가 각각 T개의 토큰을 디스패치하고, 각 토큰이 E개의 전문가 중 R개의 전문가에게 라우팅되는 경우, 모든 토큰이 특정 랭크로 라우팅되는 최악의 경우 해당 랭크는 최대 N * T * max(R, E / N)개의 토큰을 수신할 수 있습니다. 또한, 전용 수신 버퍼는 너무 많은 메모리를 사용하므로, 이 최소 한계는 전송자가 쓰기를 연속적으로 배치할 때만 달성할 수 있습니다. 전송자가 충돌 없이 쓰기를 연속적으로 배치하려면 라우팅 정보에 대한 완전한 지식이 각 랭크에 배포되어야 합니다. 결과적으로 저희 디스패치 커널은 먼저 전문가별 토큰 수를 교환한 후, 각각 자신의 쓰기가 어디로 가는지, 각 랭크로부터 토큰을 수신하는 오프셋이 무엇인지 결정할 수 있습니다.

단순히 라우팅 정보를 교환하는 것은 각 랭크가 자신의 쓰기를 디스패치하기 전에 쓰기가 도착할 때까지 기다려야 하므로 불필요하게 높은 지연 시간을 초래할 것입니다. 대신, 저희는 각 수신자 랭크에 소량의 전용 전송자별 스페이스를 예약하고 라우팅 정보 쓰기와 함께 토큰의 일부를 디스패치합니다. 이는 디스패치 과정 전반에 걸쳐 완전한 대역폭 활용을 보장합니다. 경로가 교환된 후, 나머지 토큰은 원격 피어당 단일 쓰기를 통해 전송됩니다. 라우팅 정보의 교환과 고정 크기 버퍼, 그리고 이를 채우는 전송 작업은 아래 그림에 설명되어 있습니다.

MoE 디스패치용 버퍼 레이아웃으로, 토큰 데이터와 라우팅 메타데이터가 어떻게 배치되고 두 랭크 간에 교환되는지 보여줍니다.

결합 커널은 디스패치 단계에서 라우팅 정보를 이미 사용할 수 있으므로 훨씬 간단합니다. 이들은 피어당 단일 쓰기로 토큰을 복사할 수 있으며, 그 후 수신 버퍼에서 읽어 토큰의 가중 평균을 계산하여 출력 텐서로 만듭니다.

NVLink의 경우, 노드 내 전송 및 수신 버퍼가 피어에 매핑되어 인접 랭크로부터 토큰을 직접 푸시하거나 풀(pull)할 수 있습니다.

아래 그림은 GPU, CPU 및 NIC 간의 타이밍과 상호 작용을 자세히 보여주며, 모델 커널과 교차하는 디스패치 및 결합의 한 반복을 나타냅니다.

작업은 디스패치 전송에서 시작되며, 커널은 단일 블록 내의 공유 메모리에서 원자 누적을 통해 라우팅 정보를 집계하고 전문가별 토큰 수를 계산합니다. 토큰 수를 사용할 수 있게 되면 호스트 메모리로 복사되고 프록시는 즉시 다른 모든 피어 랭크로 라우팅 정보를 디스패치하도록 통보받습니다. 병렬로, 디스패치는 토큰 수의 누적 합계를 계속 계산하여 피어별 전송 버퍼에 토큰을 채우는 데 필요한 오프셋을 식별합니다. 오프셋을 사용할 수 있게 되면 EFA를 통해 전송될 모든 토큰이 전송 버퍼에 채워집니다. 프록시는 노드 간 피어의 전송 버퍼에 쓰기를 게시하도록 통보받습니다. 노드 내에서는 커널이 NVLink를 통해 동일 노드 내의 전용 버퍼로 즉시 토큰을 푸시합니다.

MoE 디스패치의 타이밍과 조정으로, 단계적으로 진행되는 GPU, CPU, NIC 작업과 그 동기화를 보여줍니다.

MoE 결합의 타이밍과 조정으로, GPU, CPU, NIC 작업과 그 동기화를 보여줍니다.

랭크가 모든 피어로부터 라우팅 정보를 수신하면 이를 중앙 집중화하고 디스패치 및 결합을 위해 토큰을 식별하는 데 필요한 다양한 오프셋을 계산합니다. 이 작업은 전송이 대기 중인 동안 CPU에서 수행되므로 지연 시간에 중요하지 않습니다. 이후 모든 오프셋을 알게 되면 전용 버퍼 크기를 초과하는 나머지 토큰에 대한 쓰기가 게시되고 디스패치의 전송 단계가 완료됩니다.

수신자 측에서 커널은 먼저 NVLink를 통해 전용 버퍼 용량을 초과하는 동일 노드의 토큰을 가져옵니다. 다음으로, 호스트 프록시가 모든 전송 완료 신호를 보낼 때까지 차단됩니다. 신호가 수신되면 토큰은 수신 버퍼에서 그룹 GEMM에 전달된 출력 텐서로 섞입니다. 버퍼는 디스패치와 결합 간에 재사용되므로, NVLink 및 RDMA 장벽은 다음 커널이 덮어쓰기 전에 수신 버퍼가 읽혔음을 보장합니다.

결합 커널은 토큰을 섞고 프록시에 신호를 보내 완료를 기다리는 유사한 패턴을 따릅니다. 그러나 쓰기를 게시하기 위한 메타데이터를 즉시 사용할 수 있고 모든 원격 오프셋을 알 수 있으므로, 노드 내에서 전송된 모든 토큰은 NVLink를 통해 수신자로 푸시되고 노드 간 토큰은 단일 쓰기를 통해 전송됩니다. 결과를 반환하기 전에 커널은 수신 버퍼에서 토큰을 누적하여 출력 텐서에 씁니다.

구현 프로토타입을 제작하기 위해, 저희는 처음에 RDMA 쓰기를 발행하고 추적하기 위해 분리된 프리필 및 디코딩을 위해 구축된 Transfer Engine을 재사용했습니다. 그러나 각 노드가 동일한 피어에 많은 수의 쓰기를 게시하는 분리된 디코딩과 달리, MoE 디스패치는 한 번의 작업으로 모든 피어에 작은 전송을 게시합니다. 즉시 사용 가능한 상태에서 성능은 최적이 아니었으며, 저희 구현은 ConnectX-7에서 DeepEP보다 약 20us 뒤처졌습니다. 결과적으로 저희는 많은 피어와의 상호 작용을 최적화하기 위해 scatter 및 barrier라는 두 가지 작업을 지원하도록 TransferEngine을 특화했습니다.

Scatter 및 barrier 작업은 모두 기본 전송 구현(작업 요청(WR) 템플릿)에서 사용하는 데이터 구조를 미리 채우기 위해 사전 등록할 수 있는 피어 그룹을 노출합니다. Scatter는 동일한 소스 버퍼의 서로 다른 슬라이스를 피어 그룹의 모든 멤버에 있는 서로 다른 오프셋으로 복사합니다. Barrier는 immediate를 통해 모든 피어에 신호를 보냅니다. 이러한 작업을 대량으로 대기열에 넣고 처리하면 임계 경로에서 귀중한 마이크로초를 절약할 수 있습니다. 400Gbps 대역폭을 제공하기 위해 두 개의 200Gbps NIC가 집계되는 EFA에서는 각 전송의 바이트를 샤딩하는 대신 그룹 피어를 두 NIC에 걸쳐 샤딩합니다.

처음에는 libfabric을 사용하여 EFA에서만 TransferEngine을 개발했지만, TransferEngine에 libibverbs를 사용한 ConnectX-7 지원을 추가한 후 MoE 커널은 코드를 단 한 줄도 변경하지 않고 ConnectX-7에서 작동합니다. Scalable Reliable Datagram(SRD) 기반 EFA와 비교할 때, ConnectX-7은 연결 설정 및 피어 관리를 위해 추가 작업이 필요합니다. 저희는 Reliable Connection(RC)을 부트스트랩하기 위해 Unreliable Datagram(UD) 큐 쌍을 사용합니다. 각 원격 피어에 대해 두 개의 큐 쌍을 생성하는데, 하나는 양방향 SEND/RECV용이고 다른 하나는 일방향 WRITE_IMM용입니다. 이는 RECV와 WRITE_IMM이 순서대로 RECV 작업을 소비하지만 두 작업이 다른 버퍼 크기를 필요로 하기 때문입니다. 또한, ConnectX-7과 GPU 메모리 간의 PCIe 트랜잭션 순서를 완화합니다. EFA와 유사하게 WR 템플릿 최적화를 구현합니다. 또한 동일한 피어에 여러 작은 메시지를 보낼 때 메시지 속도를 향상시키기 위해 WR 체이닝을 사용합니다.

GPU 측 작업은 두 어댑터 모두에서 동일하므로, 저희는 EFA와 ConnectX-7 MoE 라우팅 간의 격차가 기본 하드웨어의 기능에 의해 결정된다는 점을 강조하기 위해 기본 전송 구현을 한계까지 밀어붙여 DeepEP의 성능을 능가했습니다.

저희는 7168의 숨겨진 차원, 블록 스케일 fp8 디스패치 및 bf16 결합을 사용하는 표준 DeepSeek-V3 및 Kimi-K2 구성에서 ConnectX-7 및 EFA 어댑터가 장착된 H200 시스템의 디스패치 및 결합 지연 시간을 평가했습니다. 디스패치에는 128, 결합에는 4096의 배치 크기를 고려합니다. ConnectX-7에서는 GPUDirect Async(IBGDA)와 호스트 프록시(IBRC)를 모두 사용하여 DeepEP 및 이전 NVSHMEM 기반 pplx-커널과 비교합니다. EP16, EP32 및 EP64에서 각각 459us, 582us 및 692us의 결합 지연 시간을 달성합니다. 저희는 519us, 966us 및 1159us를 달성하는 UCCL-EP와 같은 일반 프록시에 의존하는 다른 EFA 기반 커널보다 뛰어난 성능을 보입니다.

EP64 MoE 디코드의 디스패치 및 결합 지연 시간으로, 표시된 구현과 통신 경로를 비교합니다.

디코딩 시, 저희 커널이 DeepEP보다 몇 마이크로초 느림에도 불구하고 결합 커널이 훨씬 빠르기 때문에 전반적으로 더 낮은 지연 시간을 달성합니다. 호스트 프록시를 사용하지만, 대량 전송으로 인해 모든 IBGDA 기반 프록시리스 구현의 성능을 능가합니다. 또한 저희의 고도로 특화된 프록시가 일반 NVSHMEM 구현보다 10배 빠르기 때문에 프록시 구현이 중요하다는 것을 보여줍니다. EFA가 MoE 라우팅에서 교환되는 일반적인 패킷 크기에 가까운 256KB 패킷을 전송할 때 ConnectX-7보다 두 배 느리다는 사실에도 불구하고, 두 어댑터 간의 격차는 전반적으로 덜 상당합니다.

EP64 MoE 프리필의 디스패치 및 결합 지연 시간으로, pplx와 DeepEP 통신 경로를 비교합니다.

저희의 디코딩 최적화 커널이 추가 조정 없이도 프리필 배치 크기로 확장될 수 있음에도 불구하고, 성능은 DeepEP에 뒤처집니다. 이는 주로 결합의 수신자 측에서만 토큰을 누적하는 반면 DeepEP는 전송되는 데이터 양을 줄이기 위해 전송자 측에서 누적하기 때문입니다. 그러나 이 최적화는 정확도에 영향을 미칠 수 있습니다. 또한 커널이 한 번의 쓰기로 데이터를 전송하므로 전송 및 수신 버퍼에 더 많은 메모리를 사용합니다.

저희 커널의 실행 가능성을 증명하기 위해, Elastic Fabric Adapter(EFA) 및 H200 GPU가 장착된 여러 p5en 인스턴스에서 DeepSeek-V3(671B 매개변수) 및 Kimi-K2(1T 매개변수)를 서비스하는 사내 추론 엔진의 엔드투엔드 성능을 벤치마킹합니다. 모델을 초안 검증 모드에서 실행하고 단일 MTP 레이어를 사용하여 80%의 수락률을 가정하여 지연 시간을 보고합니다.

DeepSeek-V3에서는 효율적인 NVLink 전용 디스패치 및 결합 커널을 실행하는 EP=8 노드 내 베이스라인과 비교합니다. 노드 간 배포는 전반적으로 단일 노드 배포의 성능과 일치하거나 능가하며, 특히 중간 크기 배치에서 노드 간 확장이 전반적인 처리량 향상을 가져옴을 보여줍니다.

Kimi-K2에서 커널은 서비스 제공을 가능하게 합니다. 크기로 인해 모델은 H200 노드에서 데이터 병렬 처리(DP)를 사용하여 동일한 노드 내에 맞을 수 없으므로 모델을 서비스하려면 노드 간 배포가 필요합니다. 모델이 더 적은 어텐션 헤드를 가지고 있기 때문에 DeepSeek-V3보다 더 나은 지연 시간을 달성하며 실행 가능한 지연 시간으로 EFA에서 모델을 서비스할 수 있습니다.

EP32, EP16, EP8 구성에서 DeepSeek-V3의 GPU당 처리량과 출력 속도를 비교합니다.

EP32 및 EP16 구성에서 Kimi-K2의 GPU당 처리량과 출력 속도를 비교합니다.

저희는 마이크로 배칭도 고려했지만, 이는 매우 높은 배치 크기나 두 개 이상의 초안 토큰을 사용할 때만 지연 시간을 개선했습니다. 그룹화된 GEMM 문제를 두 개로 분할하여 이점을 얻으려면 토큰 배치 크기가 충분히 조밀해야 합니다. 시간의 대부분은 여전히 MoE 디스패치 및 결합에 소비되며, 이는 커널에 대한 추가 개선이 결과적으로 지연 시간을 개선하고 마이크로 배칭을 촉진할 것임을 나타냅니다.

저희는 EFA 성능을 개선하기 위해 AWS 엔지니어들과 긴밀히 협력하고 있습니다. 저희 측에서는 데이터 평면 오버헤드를 줄일 libfabric 업데이트를 따를 것입니다. 또한, 사용자 공간 오버헤드를 더욱 줄이고 하드웨어와 직접 인터페이스하기 위해 efa-direct를 실험할 것입니다.

NVIDIA Magnum IO NVSHMEM 및 GPUDirect Async를 사용한 HPC 시스템의 네트워크 성능 향상

BibTeX

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}

소개

InfiniBand 기반 추론

NVSHMEM 기반 pplx-커널

EFA를 통한 MoE 올투올(all-to-all)

전송 및 수신 버퍼

디스패치 및 결합

MoE 라우팅을 위한 TransferEngine

평가

디스패치/결합 지연 시간

Kimi-K2 / DeepSeek-V3 처리량

향후 작업

참고 문헌

인용