Expert Parallelism을 활용한 AMD Instinct MI300X GPU에서의 초당 21K 출력 토큰 DeepSeek 추론

개요

대규모 언어 모델(LLM)의 주요 혁신 중 하나는 Mixture-of-Experts(MoE) 아키텍처의 도입입니다. 단일 대형 feed-forward network(FFN) 레이어 대신, 모델은 각각 expert라고 불리는 여러 개의 소규모 FFN 레이어를 포함합니다. 각 토큰에 대해 경량 gating 네트워크가 각 단계에서 처리할 소수의 expert만을 동적으로 선택합니다. 이러한 조건부 연산을 통해 MoE 모델은 효율적인 연산 활용을 유지하면서도 매우 큰 파라미터 수로 확장할 수 있습니다. 예를 들어, 2025년 1월 출시 당시 큰 반향을 일으킨 DeepSeek-R1 671B 모델과 2025년 8월에 출시된 OpenAI의 인기 오픈소스 GPT-OSS 120B 모델 모두 MoE 아키텍처를 사용합니다.

그러나 대규모이며 희소한 설계로 인해, DeepSeek-R1과 같은 모델은 대규모로 효율적으로 서빙하기 위해 고급 최적화 기법이 필요합니다. 특히, Expert Parallelism(EP)을 적용하면 높은 처리량(즉, 초당 총 출력 토큰 수)을 달성할 수 있습니다. 이는 (1) GPU 메모리에 expert의 일부만 저장되어 더 큰 배치 크기가 가능하고, (2) 개별 expert의 파라미터가 로드된 후 더 많이 재사용되어 FLOPS per byte가 향상되기 때문입니다. 그러나 활성화된 expert를 담당하는 GPU로 활성화 값을 전송하는 dispatch/combine이라는 세밀한 통신 패턴을 효율적으로 구현하고, expert 간 불균형을 해소하는 것이 과제로 남습니다.

AMD의 소프트웨어 파트너인 Moreh는 최근 ROCm 소프트웨어 스택에서 EP를 구현하여 DeepSeek-R1 추론을 높은 처리량으로 실행할 수 있음을 보여주었습니다. LLM 추론에서 이러한 높은 성능을 실제로 달성하려면 스택의 여러 계층에 걸쳐 세심하게 최적화된 소프트웨어가 필요합니다. Moreh 팀은 라이브러리 및 모델 구현 수준 모두에서 다양한 최적화를 적용합니다. 그 결과, 8x AMD Instinct MI300X GPU가 장착된 서버에서 >21,000 tokens/sec의 디코딩 처리량을 달성했습니다.

병렬화 방법

DeepSeek R1 모델에서 각 decoder 블록은 256개의 expert를 포함하며, 각 토큰에 대해 그중 8개가 활성화됩니다. 이 외에도 라우팅과 관계없이 항상 실행되는 shared expert가 있으며, DeepSeek이 도입한 새로운 어텐션 메커니즘인 Multi-Head Latent Attention(MLA) 레이어가 있습니다. EP는 일반적으로 data parallelism(DP)과 결합하여 구현됩니다. 각 GPU는 서로 다른 expert 부분 집합을 보유하고 실행합니다. 한편, shared expert와 MLA처럼 항상 실행되는 구성 요소의 경우, 각 GPU가 별도의 요청 배치를 병렬로 처리합니다. 이에 따라, 노드당 DP=8 및 EP=8의 병렬화 구성을 적용했습니다. 즉, 각 GPU는 32개의 expert를 담당합니다.

또한, prefill-decode disaggregation — 즉, prefill과 decode 단계를 별도의 노드에서 실행하는 방식 — 을 적용하면 클러스터 시스템의 전체 처리량을 더욱 높일 수 있습니다. prefill 단계는 전체 입력 시퀀스를 한 번에 처리하며 연산 집약적(compute-bound)인 반면, decode 단계는 출력 토큰을 자기 회귀 방식으로 하나씩 생성하며 메모리 집약적(memory-bound)입니다. 실행을 분리하는 것뿐만 아니라 prefill과 decode의 서로 다른 연산 특성에 맞춰 병렬화 및 최적화 구성을 조정함으로써 GPU 효율을 더욱 향상시킬 수 있습니다. 본 글에서는 일반적인 애플리케이션에서 대부분의 GPU가 디코딩에 할당되므로, 주로 decode 단계의 성능 평가에 초점을 맞춥니다.

성능 최적화

새로운 병렬화 방법을 구현하는 것만으로는 최상의 성능이 자동으로 달성되지 않습니다. 하드웨어의 잠재력을 완전히 끌어내려면 라이브러리에서 모델 구현에 이르는 전체 소프트웨어 스택을 새로운 연산 및 통신 패턴에 맞게 최적화해야 합니다. Moreh 팀은 vLLM을 기반으로 아래와 같이 여러 최적화를 구현했습니다.

GPU 간 로드 밸런싱

EP 구현의 주요 과제 중 하나는 워크로드가 GPU 간에 균등하게 분배되지 않으며 라우팅 결과에 따라 달라진다는 것입니다. MoE 모델은 이러한 불균형을 최소화하도록 학습되지만, 여전히 본질적인 한계가 있습니다. DeepSeek도 이 문제를 인지하고 이를 해결하기 위해 EPLB라는 자체 로드 밸런싱 방법을 사용했다고 보고했습니다. 명시적인 로드 밸런싱 전략을 적용하지 않고 expert를 단순히 8개의 연속 청크로 나누었을 때 — 즉, expert 0-31을 첫 번째 GPU에, 32-63을 두 번째 GPU에, …, 224-255를 마지막 GPU에 할당했을 때 — GPU 간 최대 2배의 워크로드 불균형을 관찰했습니다. 우리는 각 expert의 활성화 빈도를 측정하고, 256개의 expert를 각 세트의 총 활성화 빈도가 최대한 균형을 이루도록 32개씩 8개 세트로 그룹화하는 알고리즘을 설계했습니다. (이는 균형 수 분할 문제의 변형으로 볼 수 있습니다.) 이 알고리즘은 각 세트가 연속적으로 배치되도록 expert를 재정렬하고 그 위에 EP를 적용합니다. expert 순열이 변경되었으므로 gating 네트워크도 새로운 expert 순열을 반영한 라우팅 결과를 생성해야 합니다. 그 결과, GPU 간 워크로드 불균형을 5% 이내로 줄이는 데 성공했습니다. 우리가 아는 한, 이는 AMD GPU에서의 최초의 EP 로드 밸런싱 구현입니다. 실제 서빙 시나리오에서는 활성화 빈도를 주기적으로 재측정하고 그에 따라 새로운 expert 순열을 생성함으로써 지속적인 로드 밸런싱을 달성할 수 있습니다.

연산-통신 오버래핑

EP 적용의 또 다른 과제는 dispatch/combine all-to-all 통신으로 인한 오버헤드를 최소화하는 것입니다. 배치를 두 개의 마이크로배치로 나누어 동시에 실행하면, 한 마이크로배치의 통신이 다른 마이크로배치의 연산과 오버랩될 수 있습니다. 배치 크기가 충분히 클 때 — 즉, 높은 처리량 시나리오에서 — 이 접근 방식은 연산 성능을 희생하지 않으면서 통신 지연 시간을 효과적으로 숨길 수 있습니다. 이 접근 방식에는 여러 구현 변형이 있으며, 그중 우리는 vLLM의 DBO(Dual Batch Overlap) 시스템을 채택했습니다. DBO 시스템을 MORI-EP를 사용하도록 수정했으며, 특히 MORI-EP 라이브러리의 통신 연산이 DBO의 두 워커 스레드에서 호출될 수 있고 다른 연산과 적절히 동기화되도록 개선했습니다.

라이브러리 최적화

Moreh 팀은 GPU 효율을 극대화하고 처리량(tokens/sec)과 지연 시간(첫 번째 토큰까지의 시간 및 토큰 간 지연 시간)을 모두 개선하기 위해 여러 라이브러리 최적화도 적용했습니다. 다음은 최적화의 주요 사례입니다:

• 최적 GEMM 및 Attention Kernel 선택: 온라인 프로파일링이나 수동 튜닝 없이 최적의 GEMM 및 Attention 커널을 동적으로 선택합니다.
• Fused MoE Kernel 최적화: 특히 작은 배치 크기에서 AMD의 AITER 라이브러리보다 더 나은 성능을 제공하는 고도로 최적화된 fused MoE 커널을 구현합니다.
• FP8 KV Cache 지원: KV cache를 FP8 형식으로 저장 및 로드할 수 있는 Multi-head Latent Attention(MLA) 커널을 구현하여, 특히 긴 컨텍스트 시나리오에서 성능을 개선합니다.
• 수직 및 수평 Kernel Fusion: 수직 fusion(예: fused RoPE 커널)과 수평 fusion(예: shared expert의 여러 GEMM 병합)을 모두 활용하여 커널 오버헤드를 줄이고 연산 효율을 향상시킵니다.

이러한 라이브러리 수준 최적화와 EP와 독립적인 성능 향상에 대한 자세한 내용은 Moreh의 기술 보고서를 참조하시기 바랍니다.

HIP Graphs 활용

HIP Graphs는 여러 GPU 연산을 정적 그래프로 캡처하여 한 번에 실행함으로써 CPU 런타임 오버헤드를 줄이는 기법입니다. 이는 개별 연산이 상대적으로 짧아 GPU를 완전히 활용하는 것이 중요한 decode 단계에서 특히 필수적입니다. 그러나 정적 그래프를 구성하려면 연산 간 전달되는 텐서 크기가 고정되어야 합니다. EP에서는 각 expert의 입력 크기가 본질적으로 동적이며 라우팅 결과에 의해 결정되므로, 이를 정적 그래프로 표현하기가 매우 어렵습니다. 우리는 vLLM의 DeepSeek 모델 구현을 수정하여 EP를 지원하면서도 텐서 크기를 최대한 정적으로 만들었습니다. 이를 통해 HIP Graph를 최대한 활용하여 CPU 오버헤드를 최소화할 수 있었습니다. MoRI 라이브러리 연산을 그래프에 통합하는 과정에서 여러 기술적 과제가 있었지만, 최종적으로 MORI-EP의 dispatch/combine 연산을 그래프 실행의 일부로 포함시키는 데 성공했습니다.

실험 방법

실험은 다음 사양의 MI300X 서버에서 수행되었습니다.

• CPU: 2x AMD EPYC 9474F 48-core 3.6 GHz
• Main memory: 2,304 GB
• GPU: 8x AMD Instinct MI300X OAM GPU 192 GB
• Server: Gigabyte G593-ZX1-AAX1
• Operating system: Ubuntu 22.04.4 LTS
• ROCm version: 6.4.1

평가 방법론은 LMSys 블로그 게시물을 참고하여 설계했습니다. 해당 게시물에서 EP 성능 평가의 핵심 지표(prefill/decode disaggregation을 가정)는 decode 노드당 총 처리량(tokens/sec)과 토큰 간 지연 시간입니다. LMSys는 12개 노드에 걸쳐 prefill/decode disaggregation을 구현했지만, 우리는 단일 노드 환경에서도 decode 성능을 측정하고자 했습니다. 이를 위해 vLLM을 다음과 같이 여러 가지로 수정했습니다:

• 서버 측에서 prefill 시간을 제외하고 decode 시간과 토큰 수만 누적하여 decode 처리량을 별도로 측정했습니다. 토큰 간 지연 시간은 표준 벤치마킹 접근 방식을 사용하여 클라이언트 측에서 여전히 측정할 수 있습니다.
• vLLM V1 엔진의 스케줄러는 원래 prefill과 decode 요청을 단일 배치로 실행하는데(mixed prefill이라 함), 이 경우 prefill 시간을 제외할 수 없습니다. prefill과 decode 요청이 별도의 큐에 발행되어 독립적으로 실행되도록 수정했습니다.

다음 명령을 사용하여 (수정된) vLLM 서버에 요청을 보내고 성능을 측정할 수 있습니다.

vllm bench serve \
  --backend vllm \
  --model "deepseek-ai/DeepSeek-R1" \
  --metric-percentiles "10,25,50,75,90" \
  --percentile-metrics "itl,tps,ttft,e2el" \
  --port 8001 \
  --num-prompts 2048 \
  --max-concurrency 2048 \
  --ignore-eos \
  --dataset-name sharegpt \
  --dataset-path /app/dataset/ShareGPT_V3_unfiltered_cleaned_split.json \
  --sharegpt-input-len 2000 \
  --sharegpt-output-len 100

실험 결과

지연 시간과 처리량 간의 트레이드오프를 관찰하기 위해, 4가지 동시성 수준(1024, 2048, 3072, 4096)에서 처리량(노드당)과 토큰 간 지연 시간을 측정했습니다. 결과는 다음과 같습니다.

최고 동시성 수준에서 21,224.6 tokens/sec의 처리량을 달성했습니다. 이는 SGLang 팀이 8x NVIDIA H100 GPU 서버에서 보고한 22,282 tokens/sec와 거의 동등한 수치로, 우리의 구현이 최첨단 EP 성능을 제공함을 확인합니다. 한편, 최소 지연 시간에 최적화된 구성(중앙값 토큰 간 지연 시간 91.48 ms)에서도 처리량 감소가 15% 미만이어서, 시스템이 목표 SLO(Service Level Objectives)에 따라 최대 동시성을 유연하게 조정할 수 있음을 보여줍니다.

결론

Moreh는 MoE 모델을 위한 다양한 기법을 적용하고 최적화하여 AMD Instinct MI300 시리즈 GPU에서 EP를 통한 고처리량 추론을 성공적으로 구현했습니다. 또한 본 글에서 설명한 효율적인 EP 구현을 포함하여, 다양한 분산 추론 기법을 AMD Instinct GPU 클러스터에 자동으로 적용하는 MoAI Inference Framework을 제공합니다. 자세한 내용은 Moreh 웹사이트를 방문하시기 바랍니다.