[Liger-Kernel] Liger-Kernel: SM90 MoE 통신 최적화 및 안정성 개선

PR 링크: linkedin/Liger-Kernel#1463 상태: Merged | 변경: +3453 / -576

들어가며

최근 대규모 언어 모델(LLM)의 Mixture-of-Experts(MoE) 구조가 보편화되면서, GPU 간 통신 효율이 전체 학습 속도를 결정짓는 핵심 요소가 되었습니다. 특히 NVIDIA H200(SM90)과 같은 최신 아키텍처에서 MoE 커널은 IBGDA(InfiniBand GDA)와 같은 고성능 통신 라이브러리를 사용하는데, 로컬 노드 내 통신에서도 불필요한 오버헤드가 발생하고 있었습니다. 이번 PR은 SM90 MoE 커널의 통신 경로를 로컬과 인터호스트로 분리하여 최적화하고, 백워드 패스에서 발생하던 동기화 행(hang) 문제를 해결하여 성능과 안정성을 동시에 확보했습니다.

코드 분석

1. SM90 통신 경로의 전문화 (Specialization)

기존에는 로컬 통신과 인터호스트 통신이 동일한 커널 경로를 공유하여, 로컬 환경에서도 불필요한 IBGDA 코드와 레지스터 압박이 존재했습니다. 이를 해결하기 위해 mlp_comms_sm90.cuh 등에서 통신 경로를 분리했습니다.

// Before: 공통 경로에서 IBGDA와 로컬 포인터를 혼용
// After: 로컬 전용 커널은 IBGDA 코드를 제거하고 TMA(Tensor Memory Accelerator)에 집중
if (is_local_only) {
    // Local specialization: Remove IBGDA overhead
    run_local_tma_kernel();
} else {
    // Inter-host specialization: Retain IBGDA for network
    run_ibgda_kernel();
}

2. 백워드 동기화 행(Hang) 수정

백워드 패스에서 MLP3 루프의 경계 조건이 명확하지 않아 발생하는 동기화 문제를 수정했습니다. 생산자/소비자(producer/consumer) 랑데부 지점을 명확히 하고, 빈 전문가 범위(empty expert range)나 K 슬라이스에서 커서가 멈추지 않도록 로직을 보강했습니다.

// Fix: Producer/Consumer rendezvous before blocking loop
publish_mlp3_loop_bounds();
perform_rendezvous();
for (int i = 0; i < num_experts; ++i) {
    if (is_empty_range) continue; // Advance cursors even on empty range
    advance_shared_cursors();
}

왜 이게 좋은가

이번 최적화는 H200 및 B300 환경에서 괄목할 만한 성능 향상을 보여줍니다. 특히 16개 GPU 환경에서 Llama 4 Scout 모델 기준, 이전 대비 훨씬 높은 TFLOP/s를 달성했습니다.

• 성능 수치: H200 16-GPU 환경에서 기존 대비 유의미한 처리량 향상을 기록했으며, 특히 B300에서는 1.4배 이상의 기하평균 속도 향상을 보였습니다.
• 교훈:
  1. 경로 분리(Path Specialization): 범용 커널보다 특정 하드웨어 토폴로지(로컬 vs 네트워크)에 최적화된 커널을 분리하는 것이 레지스터 압박을 줄이고 IPC(Instructions Per Cycle)를 높이는 데 효과적입니다.
  2. 동기화 명확화: 비동기 커널(TMA/NVSHMEM) 작성 시, 모든 분기(branch)에서 커서와 세마포어가 올바르게 갱신되도록 보장하는 것이 행(hang) 방지의 핵심입니다.

리뷰어 피드백 반영

리뷰어 yueyiming2009는 토폴로지 변경 시 모듈 재사용 문제를 지적했습니다. 이에 따라 ModuleState가 토폴로지를 검증하고, 변경 시 명시적으로 거부하도록 하여 안정성을 높였습니다. 또한, kolehma8은 빌드 시점의 설정(all-configs vs target-mixtral)이 핑거프린트에 포함되도록 하여, 잘못된 커널이 로드되는 것을 방지했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#tma-tensor-memory-accelerator

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