고성능 EP 커널 구조 분석

출처: Hacker News

대규모 언어 모델은 거대합니다. 크기가 크기 때문에 실행하려면 많은 GPU가 필요합니다. LLM 추론이 극히 병렬화가 쉬워서 각 GPU에서 독립적인 작업을 언제든지 수행할 수 있다면 좋겠지만, 현실은 그렇지 않습니다. LLM 추론에 많은 GPU를 사용하려면 GPU들 간에 서로 통신하도록 해야 합니다.

GPU들을 함께 작동시키는 방법은 다양합니다: 텐서 병렬화(Tensor Parallelism), 파이프라인 병렬화(Pipeline Parallelism), 컨텍스트 병렬화(Context Parallelism), 전문가 병렬화(Expert Parallelism) 등. 각각의 방법은 쓰임새가 있습니다. 하지만 MoE 모델에서는 MoE 레이어를 대규모로 서비스할 때 ‘wide Expert Parallelism’(wideEP) 가 가장 중요합니다. 
vLLM의 원본 DeepSeek 대규모 서비스 포스트에서 실제 운영 규모의 시연을 확인할 수 있습니다. 예를 들어 DeepSeek은 H200 클러스터에서 GPU당 2.2k 토큰/초를 달성했으며, wideEP와 데이터 병렬 어텐션을 사용해 서비스되었습니다.

다른 종류의 병렬화는 모두 GPU 간 통신이 필요하지만, 그 통신 패턴은 아키텍처에 의해 고정됩니다: 누가 보내고, 누가 받고, 얼마나 보낼지가 순전파가 시작되기 전에 모두 알려져 있으며, 매 단계마다 동일합니다. 이러한 통신은 표준 집합 연산(collective)으로 실행될 수 있습니다.

전문가 병렬화는 다릅니다. 어떤 토큰이 어떤 GPU에 도달해야 하는지는 라우터가 런타임에 데이터로부터 결정합니다. 그리고 토큰은 ‘데이터 병렬 어텐션’ 구성을 가정합니다. 즉, 각 토큰은 정확히 하나의 랭크(클러스터 내의 하나의 GPU)에 존재합니다. 전문가들은 같은 랭크들에 걸쳐 퍼져 있기 때문에, 토큰과 라우팅된 전문가가 일반적으로 같은 위치에 있지는 않습니다. 아래 예시는 8개의 GPU가 2개의 노드에 걸쳐 배치되고, GPU당 전문가 2명, 랭크당 토큰 1개, 토큰당 라우팅된 전문가 2명인 경우를 보여줍니다.

Hover a rank chip for its token’s round trip, or an expert for everything routed to it. Four of the sixteen experts drew no tokens at all this step: the routing is lumpy.

DISPATCH  EXPERTS  COMBINE
NODE 0 · NVLINK   NODE 1 · NVLINK
crossing = RDMA · within a node = NVLink
GPU 0 GPU 1 GPU 2 GPU 3 GPU 4 GPU 5 GPU 6 GPU 7
Expert 0 Expert 1 Expert 2 Expert 3 Expert 4 Expert 5 Expert 6 Expert 7
Expert 8 Expert 9 Expert 10 Expert 11 Expert 12 Expert 13 Expert 14 Expert 15
r0 r0 r1 r1 r2 r2 r3 r3 r4 r4 r5 r5 r6 r6 r7 r7

MoE 레이어를 실행할 때, 토큰은 네트워크 패브릭 어디에 있든 자신이 라우팅된 전문가와 만나야 합니다. 이를 실현하는 것이 EP(Expert Parallelism) 통신 커널의 역할입니다.

이러한 커널의 현대적 형태는 DeepSeek의 DeepEP 라이브러리에 의해 정립되었습니다. 이번 포스트에서는 DeepEP 스타일의 dispatch와 combine 커널을 단계별로 살펴보겠습니다: 먼저 고처리량(high‑throughput) 형태를, 그 다음 저지연(latency‑optimized) 형태를 다룹니다.

우리가 해야 할 일§

구성을 구체화해 보겠습니다. 우리는 2개의 노드에 걸쳐 8개의 GPU가 배치되어 있고, RDMA로 연결되어 있습니다. 각 데이터 병렬 랭크는 하나의 GPU를 소유합니다. 어텐션은 각 GPU에서 배치 크기 B_i(GPU마다 다를 수 있음)의 토큰에 대해 수행됩니다. 우리는 E = 16개의 전문가(전문가당 2명, GPU당 2명)를 사용하며, 각 토큰당 K = 2개의 전문가가 라우팅됩니다.

각 랭크 r_i의 EP 레이어 입구에서는 (B_i, H) 형태의 텐서가 있습니다(H는 은닉 차원). 라우팅 레이어는 로컬에서 실행되어 각 토큰에 대한 전문가 할당을 제공합니다. 우리는 16개 중 2개를 라우팅합니다: 각 토큰에 대해 라우터는 길이 16인 로짓(logits) 벡터를 출력하고((B_i, 16) 형태), 여기서 상위 2개의 인덱스를 선택해 (B_i, 2) 형태의 텐서를 얻습니다. 예를 들어 토큰 k가 전문가 3과 13에 라우팅되면, 행 k는 [3, 13]이 됩니다.

따라서 EP 레이어 입구에서 각 랭크는 두 가지 정보를 보유합니다: 활성화 행과 로컬 라우팅 후 각 행에 대한 top‑2 전문가 할당.

activations (Bᵢ, H)      router
expert logits (Bᵢ, E=16) assignment (Bᵢ, K=2)
0123456789 10 11 12 13 14 15
E: Wᵣ x H → E
t0 −1.2 0.3 −0.4 2.1 0.1 −0.8 0.5 −0.2 0.9 −0.6 0.2 −1.0 0.4 1.7 −0.3 0.6 3 13
t1 2.4 0.2 −0.5 9 −1.1 0.3 1.9 −0.2 0.7 −0.9 0.2 0.8 −0.3 0.5 2.0 −0.4 0.8
t2 −0.6 3.2 0.0 1 −0.2 0.7 1.0 −0.5 0.2 −0.3 0.6 −0.7 1.2 1.3
...

모든 전문가가 로컬에 있는 것은 아닙니다. 일부는 인접한 NVLink 피어에, 일부는 RDMA를 통해서만 접근 가능한 다른 노드에 있습니다. 전문가 병렬화 커널의 목표는 활성화를 필요한 위치로 이동시키고, 해당 위치에서 전문가 GEMM을 실행한 뒤, 결과를 다시 원래 위치로 되돌리는 것입니다.

통신을 다룰 때는 **처리량(throughput)**과 지연시간(latency) 중 무엇에 중점을 둘지에 따라 전략이 달라집니다. 이 구분은 추론의 두 단계와도 맞닿아 있습니다:

• Prefill 단계에서는 큰 배치를 처리하면서 다른 연산을 병렬로 수행해 통신 시간을 가릴 수 있습니다.
• Decode 단계에서는 거의 다른 작업이 없기 때문에, 전송 자체가 병목이 됩니다.

먼저 처리량 최적화된 표준 EP 형태를 살펴보고, 이후 지연시간 최적화 형태를 논의하겠습니다.

고처리량: 먼저 요청하고, 그 다음 전송§

Dispatch§

Dispatch의 목적은 그룹화된 GEMM에 입력을 공급하는 것입니다. 그룹화된 GEMM은 전문가마다 별도의 행 집합을 대상으로 한 매트릭스 곱을 하나의 커널 호출로 수행합니다. 라우팅 단계에서 생성되는 불규칙성은 바로 여기서 나타납니다: 각 전문가가 받은 토큰 수가 서로 다르기 때문입니다. 라우팅이 끝나면, 각 전문가는 자신에게 할당된 토큰들에 대해 매트릭스 곱을 수행해야 합니다. 라우팅 이전에는 토큰들이 클러스터 전역에 흩어져 있습니다. 따라서 각 랭크에서는 로컬 전문가용 토큰을 하나의 밀집 버퍼에 모아야 하며, 이 버퍼가 그룹화된 GEMM에 한 번에