AMD Instinct MI308X에서의 Step3 추론 최적화: NVIDIA H20 대비 1.30배 높은 Decode 처리량

배경

StepFun의 Step3는 토큰당 38B 파라미터가 활성화되는 321B 파라미터 규모의 Mixture-of-Experts (MoE) 멀티모달 모델입니다. 61개 레이어로 구성되어 있으며, 그 중 56개가 3-in-48 expert selection 방식의 MoE 레이어이고, 토큰당 attention 비용을 DeepSeek V3 대비 약 22%로 줄이는 Multi-Matrix Factorization Attention (MFA)을 도입했습니다.

StepFun은 NVIDIA H20 GPU에서 Step3를 서빙하고 있었으며, AMD Instinct MI308X를 대안으로 평가하고자 했습니다. Moreh는 Step3가 오픈소스로 공개되기 전에, 동일한 아키텍처를 가진 비공개 모델의 MI308X 추론 최적화를 요청받았습니다. 이는 Moreh가 독자적인 모델 아키텍처에 맞춰 Moreh vLLM을 적용하는 커스텀 모델 최적화 서비스의 사례입니다.

MI308X가 Decode에 적합한 이유

AMD Instinct MI308X는 중국 시장에서 제공되는 MI300X의 변형 모델입니다. MI300X 대비 컴퓨트 코어가 1/4로 줄었지만 동일한 HBM3e 메모리 용량과 대역폭을 유지합니다. 이러한 특성 덕분에 MI308X는 LLM 추론의 decode 단계에 특히 적합합니다. Decode 단계는 컴퓨트 바운드가 아닌 메모리 대역폭 바운드이기 때문입니다. 토큰은 autoregressive 방식으로 한 번에 하나씩 생성되며, 병목은 행렬 곱셈이 아니라 메모리에서 모델 가중치와 KV-cache를 로딩하는 데 있습니다.

최적화 기법

• 커스텀 HIP attention 커널: 기본 vLLM Triton attention 커널이 GPU 시간의 약 50%를 차지하는 가장 큰 병목이었습니다. Step3의 MFA 구성(64 query heads, 1 KV head, head dimension 256)에 최적화된 커스텀 HIP attention 커널을 data parallelism과 함께 개발했습니다. 이 커널은 decode 배치에서 attention 지연을 72%, 혼합 prefill/decode 배치에서 37% 감소시켰습니다.
• CUDA graph: GPU 커널 지연이 크게 줄어든 후, CPU 측 오버헤드가 decode 단계의 다음 병목이 되었습니다. DP8-EP8 병렬처리를 사용하는 Step3 모델에 대해 전체 CUDA graph 캡처를 활성화하여, decode 처리량을 약 2,900에서 4,100 tok/s로 향상시켰습니다.
• 혼합 BF16–FP8 blockscale 양자화: BF16과 FP8 blockscale 연산 모두에 대해 철저한 GEMM 튜닝을 수행하여 최적의 정밀도–효율성 트레이드오프를 달성했습니다.
• 최적화된 MoE one-stage 커널: Step3의 MoE 레이어에 대한 커스텀 커널 최적화로, inter_dim 파라미터에 초점을 맞추었습니다.
• Shared-expert MLP 융합: MoE 레이어 내에 shared-expert MLP를 통합하여 중복 연산을 줄이고 추론 지연을 개선했습니다.
• MoRI EP 통합: AMD GPU에서 효율적인 expert-parallel all-to-all 통신을 위해 MoRI 라이브러리를 통합했습니다.

성능 결과

동일한 테스트 구성(ISL=4096, OSL=256, Concurrency=256, DP8-EP8 병렬처리(attention에 8-way data parallelism, MoE에 8-way expert parallelism))을 사용하여, 8× MI308X에서의 Moreh vLLM과 StepFun이 보고한 8× NVIDIA H20 수치를 비교 벤치마크했습니다.

결과는 두 단계 사이에 명확한 차이를 보여줍니다:

• Decode: MI308X에서 Moreh vLLM은 4,082 tok/s를 달성하여, StepFun의 H20 기준 대비 1.30× 높은 처리량과 1.30× 낮은 지연(63 ms vs. 82 ms)을 기록했습니다.
• Prefill: 컴퓨트 바운드인 prefill 단계에서는 H20이 우위를 유지합니다(13,780 vs. 9,601 tok/s). 이는 H20의 더 강력한 온칩 캐시 서브시스템을 고려하면 예상된 결과입니다.

Prefill–decode 분리(disaggregation) 방식의 프로덕션 서빙에서는 대부분의 GPU가 decode 단계에 할당됩니다. MI308X의 강력한 decode 성능은 대규모 서빙에서 직접적인 비용 효율성으로 이어집니다.

요약

이번 프로젝트는 Moreh vLLM의 모델별 최적화와 결합된 AMD Instinct MI308X가 대규모 MoE 모델에서 NVIDIA H20보다 높은 decode 처리량을 제공할 수 있음을 보여줍니다. MI308X는 컴퓨트 용량 대비 높은 메모리 대역폭을 갖추고 있어, 프로덕션 LLM 서빙 배포에서 GPU 할당의 대부분을 차지하는 decode 단계에 비용 효율적인 선택입니다.

Moreh는 독자적인 모델 및 파인튜닝된 모델을 위한 커스텀 vLLM 최적화를 제공합니다. AMD GPU에서 모델 운영을 평가하고 계신다면, 문의하기를 통해 저희가 어떻게 도움을 드릴 수 있는지 상담해 보세요.