[flashinfer] [FlashInfer] Hopper SM90을 위한 MegaMoE 최적화: Combine Wire 및 GEMM 에필로그 개선

PR 링크: flashinfer-ai/flashinfer#6058 상태: Merged | 변경: +2450 / -1175

들어가며

대규모 언어 모델(LLM)의 효율성을 극대화하기 위해 Mixture of Experts(MoE) 구조가 널리 사용되고 있습니다. 하지만 MoE는 전문가(Expert) 간의 데이터를 주고받는 과정에서 발생하는 통신 오버헤드와 복잡한 커널 실행 구조로 인해 성능 병목이 발생하기 쉽습니다.

이번 글에서는 flashinfer-ai/flashinfer 레포지토리에 반영된 Hopper(SM90) 아키텍처 전용 BF16 MegaMoE 백엔드 최적화 PR을 분석합니다. 이 PR은 CAKE 팀에 의해 개발되었으며, 기존 SOTA(State-of-the-Art) 베이스라인을 뛰어넘는 성능을 보여줍니다. 핵심은 Combine Wire 방식의 혁신과 Grouped-GEMM 에필로그 최적화, 그리고 고정 비용(Fixed-cost) 레버 개선에 있습니다.

---

1. Combine Wire 최적화: 통신량의 획기적 절감

MoE의 'Combine' 단계는 각 전문가가 계산한 결과를 토큰의 원래 소유자에게 다시 보내 합치는 과정입니다. 기존 방식(per_route)은 각 경로(route)마다 별도의 행을 전송했으나, 이번 PR에서는 이를 사전 축소(pre-reduce)하는 방식을 도입했습니다.

코드 변경 사항: cake_config.py 및 cake_backend.py

먼저 설정 클래스에 combine_wire 옵션이 추가되었습니다.

Before (Round 1): 기존에는 단일한 통신 프로토콜만 사용되었습니다.

After (Round 2):

# flashinfer/moe_ep/backends/mega/kernel/sm90/bf16_bf16_bf16_push_cake/cake_config.py

class Sm90_Bf16_Bf16_Bf16_PushCake_MegaMoeConfig:
    # ... 기존 필드들 ...
    init_timeout_s: float = 600.0
    combine_wire: str | None = None # 신규 추가

이 combine_wire 옵션은 다음과 같은 세 가지 모드를 지원합니다:

1. prereduced (Default): 전문가 랭크에서 동일한 토큰에 대한 여러 경로의 결과를 FP32로 미리 합친 뒤(fmaf), BF16으로 한 번만 반올림하여 전송합니다. 이로 인해 통신량이 T * top_k에서 T * world_size로 줄어듭니다.
2. prereduced_hilo: 정밀도 손실을 최소화하기 위해 두 개의 BF16 행으로 나누어 전송하는 변형입니다.
3. per_route: 기존의 Round 1 방식입니다.

왜 좋은 최적화인가?

EP8(Expert Parallelism 8), Top-8 설정에서 기존 방식은 토큰당 7.0개의 행을 전송해야 했으나, prereduced 방식을 사용하면 4.6개로 줄어듭니다. 이는 네트워크 대역폭이 병목인 분산 환경에서 엄청난 이득을 제공합니다. 또한, FP32로 중간 합산을 수행하므로 수치적 정밀도(Numerics) 면에서도 기존 방식보다 더 나은 결과를 보여줍니다.

---

2. Protocol Fixed-cost Levers: 커널 런치 오버헤드 제거

추론(Inference) 시 토큰 수가 적은 Decode 단계에서는 커널 실행 자체의 고정 비용이 성능에 큰 영향을 미칩니다. 이번 PR은 불필요한 memset과 커널 런치를 제거했습니다.

코드 분석: cake_backend.py의 검증 로직 강화

# flashinfer/moe_ep/backends/mega/kernel/sm90/bf16_bf16_bf16_push_cake/cake_backend.py

def validate_init(
    # ...
):
    # ...
    from ......kernel_src.sm90.cake_bf16_megamoe import COMBINE_WIRES

    if kcfg.combine_wire is not None and kcfg.combine_wire not in COMBINE_WIRES:
        raise MoEEpConfigError(
            f"{_NAME} combine_wire must be one of {COMBINE_WIRES} or None, got "
            f"{kcfg.combine_wire!r}"
        )

이 외에도 다음과 같은 저수준 최적화가 적용되었습니다:

• Self-resets: 매 라운드마다 4번씩 호출되던 cudaMemset 대신, 스크래치 메모리를 스스로 리셋하는 방식을 도입했습니다.
• Fused Worklist Building: 사전 축소 작업 목록을 compact 패스 내부에서 직접 구축하여 커널 런치를 하나 줄였습니다.
• Static Grid Stride: 원자적 작업 큐(Atomic work queue) 대신 정적 그리드 스트라이드를 사용하여 동기화 비용을 낮췄습니다.

---

3. Grouped-GEMM 에필로그 최적화

Hopper 아키텍처의 성능을 끝까지 쥐어짜기 위해 GEMM(General Matrix Multiply)의 에필로그 부분도 개선되었습니다.

주요 변경점

• BF16x2 Packing: 데이터를 2개씩 묶어 처리하여 처리량을 높였습니다.
• Quad Transpose & 16-byte Stores: 메모리 쓰기 시 16바이트 단위로 저장하여 대역폭 효율을 극대화했습니다.

이 최적화 덕분에 T >= 1024인 상황에서 FC1+FC2 단계의 성능이 약 4.5% 향상되었습니다. 특히 레지스터 사용량(168 regs)을 유지하면서도 성능을 올린 점이 인상적입니다.

---

4. 성능 수치 및 결과

8개의 H100 노드에서 테스트한 결과, 기존 Round 1 대비 지오메트릭 평균(Geomean) 1.031배의 속도 향상을 기록했습니다. 특정 케이스(T=8192)에서는 최대 1.103배까지 빨라졌습니다.

EP	Geometry	T=1024	T=4096
8	A (7168 hidden)	1.076x	1.092x
8	C (4096 hidden)	1.060x	1.077x

수치적 정밀도 또한 rel-L2 오차가 0.00365에서 0.00341로 개선되어, 성능과 정확도라는 두 마리 토끼를 모두 잡았습니다.

---

마치며: 시니어 엔지니어의 관점

이번 PR은 단순히 "코드를 더 빨리 짜는 것"이 아니라, 아키텍처의 특성(SM90)과 알고리즘의 특성(MoE Combine)을 깊게 이해했을 때 어떤 최적화가 가능한지를 잘 보여줍니다.

1. 통신과 계산의 트레이드오프: 네트워크 대역폭을 아끼기 위해 GPU 내부에서 FP32 연산을 추가하는 결정은 현대 분산 학습/추론 시스템에서 매우 유효한 전략입니다.
2. 결정론적 연산(Determinism): 원자적 연산(Atomic)을 피하고 고정된 순서로 합산을 수행하여 결과의 재현성을 보장한 점은 상용 서비스 운영에서 매우 중요한 가치입니다.
3. 고정 비용의 무서움: 대규모 연산뿐만 아니라 작은 배치에서의 커널 런치 오버헤드를 줄이는 것이 실제 서비스 지연 시간(Latency) 개선에 핵심적임을 다시 한번 확인시켜 줍니다.

FlashInfer의 이러한 행보는 Hopper와 같은 최신 하드웨어를 사용하는 엔지니어들에게 훌륭한 레퍼런스가 될 것입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#mathematical-functions-fmaf
• https://github.com/flashinfer-ai/flashinfer

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.