병목을 따라가다: AMD MI355X에서 MiniMax M3를 최적화한 방법

들어가며

MoE 모델의 추론 최적화는 가장 빠른 커널 하나를 찾는 것으로 끝나지 않는다. vLLM이 AMD Instinct MI355X에서 MiniMax M3를 최적화한 과정은 병목이 이동할 때마다 관찰 지점을 바꾸는 방법을 보여준다. 각 rank에 실제로 어떤 로컬 형태가 전달되는지, 레이어나 토큰 사이에서 어떤 작업이 반복되는지, 어떤 바이트가 이동하는지, 의도한 고속 경로가 실제로 선택되는지, 그리고 커널이 더 이상 전체 시간을 지배하지 않을 때 어느 대기열이 길어지는지를 확인하는 방식이다.

누적된 최적화의 결과

고정 토폴로지 MXFP8 표준 서빙에서 동시성 32의 출력 처리량은 109.1에서 342.4 token/s/GPU로 높아져 3.14배가 됐다. 중앙값 TTFT는 1.46초에서 0.67초로, 평균 TPOT는 69.1밀리초에서 22.1밀리초로 줄었다. 동시성 128에서는 같은 TP4/EP1 4GPU 경로가 297.8에서 623.7 token/s/GPU로 증가했다. MXFP4는 TP4/EP1에서 212.1에서 716.8로 높아졌고, 이후 TP2/EP1 구성에서 943.5에 도달했다. 다만 마지막 비교에는 배치 밀도 변화도 포함되므로 고정 토폴로지의 순수한 속도 향상으로 해석해서는 안 된다. EAGLE3는 동시성 128의 TP4/EP1에서 682.4를 기록했다. P/D 분리와 토폴로지 재조정을 적용한 뒤에는 동시성 512에서 6,370.5 total token/s/GPU, 중앙값 TTFT 1.32초에 도달했다.

최적화 과정의 핵심 다섯 가지

• 모델 전체가 아니라 로컬 형태를 본다. 텐서 병렬화, KV 헤드 복제, padding, 토큰 라우팅이 각 rank의 실제 M, N, K를 결정한다. TP8에서는 Q 헤드는 분할되지만 4개의 KV 헤드와 4개의 인덱스 헤드는 복제되므로, fused QKV projection의 로컬 N은 전체 차원을 8로 나눈 값이 아니라 1536이다. 큰 M과 작은 M을 분리한 launcher는 TP8 8K/1K 처리량을 7.8%~9.4% 높였다.
• 필요한 계산과 분리된 실행 경로를 구분한다. 공유 전문가는 원래 별도의 dense MLP로 실행돼 추가 launch와 중간 데이터 이동을 만들었다. 이를 routed expert 테이블에 추가하고 grouped GEMM으로 함께 처리하면서 모델 계산은 유지하고 실행 비용을 줄였다. 향상 폭은 동시성 1에서 30.2%, 동시성 128에서 5.6%였다.
• 변하지 않는 작업은 서빙 루프 밖으로 옮긴다. MXFP8 가중치와 scale의 재배치는 모델 로드 시점에 수행하도록 바뀌었다. 투기적 디코딩에서는 MSA indexer를 투기 토큰별 workgroup에서 요청별 workgroup으로 변경해 여러 draft 위치를 묶고 key 로드를 재사용했다. 인덱스 커널은 최대 48.9%, 엔드투엔드 서빙은 약 3.3% 개선됐다.
• 제어면의 바이트도 계산한다. 희소 어텐션은 연산량을 줄이지만 top-k 선택, 물리 페이지 매핑, 메타데이터 전달을 추가한다. 인접한 희소 레이어가 비슷한 블록을 고르는 특성을 활용해 인덱스를 공유하자, 평균 TPOT는 동시성 1에서 약 10%, 더 높은 동시성에서 약 4% 감소했다.
• 커널이 평탄해지면 대기열을 추적한다. 그래프 실행은 launch 비용을 줄이지만, 최적 백엔드는 입력 길이와 동시성에 따라 달라진다. native MXFP8 linear, emulated linear, sparse paged attention은 서로 다른 구간에서 유리하다. P/D 분리 역시 먼저 KV 전달을 검증한 뒤 실제 요청이 기다리는 대기열에 용량을 추가해야 한다.

의미와 영향

이 사례는 희소 모델과 MoE 서빙에 재사용할 수 있는 성능 분석 틀을 제시한다. “prefill”, “decode”, “TP4” 같은 이름만으로는 실제 로컬 형태와 실행 백엔드를 충분히 설명할 수 없다. 커널의 개선률도 서비스 전체의 개선률과 같지 않다. 인덱스 커널이 최대 48.9% 빨라졌지만 엔드투엔드 개선은 약 3.3%였던 것은 다른 단계가 여전히 중요한 경로에 남아 있었기 때문이다. 가장 실용적인 방법은 로컬 형태를 측정하고, 반복 작업을 제거하고, 데이터 이동과 디스패치를 확인한 다음, 병목이 이동한 대기열을 계속 추적하는 순환 과정이다.

출처: vLLM Blog