[flashinfer] FlashInfer: GB200/GB300 NVL72 환경에서의 LatentMoE 추론 최적화

PR 링크: flashinfer-ai/flashinfer#5624 상태: Merged | 변경: +6725 / -338

들어가며

최신 대규모 언어 모델(LLM)의 추론 성능을 극대화하기 위해서는 연산뿐만 아니라 통신 오버헤드를 줄이는 것이 핵심입니다. 특히 GB200/GB300 NVL72와 같은 고성능 환경에서 LatentMoE 모델의 통신 병목을 해결하기 위해 FlashInfer의 CAKE 팀은 새로운 커널 최적화(Round 4)를 도입했습니다. 이 PR은 토큰 수에 따라 커널을 동적으로 디스패치하여, 작은 배치에서는 weight-streaming을, 큰 배치에서는 persistent 파이프라인을 사용하여 SOTA 성능을 달성합니다.

코드 분석

1. cake_backend.py: 동적 커널 디스패치 및 파이프라인 제어

이 파일에서는 토큰 수(M)에 따라 최적의 커널을 선택하는 로직이 추가되었습니다. 특히 K3_PERSIST_MIN_TOKENS를 기준으로 커널 형태를 분기합니다.

Before (기존 방식):

# 기존에는 고정된 형태의 커널을 사용
def k3_kernel_key(num_tokens: int) -> str:
    return f"k3:{poll_schedule_for(num_tokens)}"

After (최적화된 방식):

# M >= 256일 때 persistent 파이프라인 사용
def k3_form_for(num_tokens: int) -> str:
    return "persist" if num_tokens >= K3_PERSIST_MIN_TOKENS else "lamport"

def k3_kernel_key(num_tokens: int) -> str:
    # ... 생략 ...
    prefix = "k3_persist" if k3_form_for(num_tokens) == "persist" else "k3"
    return f"{prefix}:{poll_schedule_for(num_tokens)}"

2. K2-stream 도입 (M <= 4)

작은 배치 사이즈(M <= 4)에서는 기존의 torch.mm (cuBLAS) 대신 SIMT 기반의 weight-streaming 커널을 사용하여 HBM 접근을 최적화했습니다.

# K2-stream 커널 디스패치
def k2_form_for(num_tokens: int) -> str:
    return "stream" if num_tokens <= K2_STREAM_MAX_TOKENS else "cublas"

왜 이게 좋은가

이번 최적화의 핵심은 통신과 연산의 중첩(Overlap)입니다.

1. Persistent K3 Pipeline: M >= 256인 경우, min(M, SM count) * 2개의 CTA를 사용하여 3단계 소프트웨어 파이프라인(scatter, reduce+add, gather)을 구성합니다. 이를 통해 한 토큰의 통신이 완료되기를 기다리지 않고 다음 토큰의 통신을 겹쳐서 처리함으로써, GB300 기준 4096 토큰에서 기존 대비 약 13%의 성능 향상을 보였습니다.
2. K2-stream: 작은 배치에서 cuBLAS의 오버헤드를 제거하고, 가중치를 HBM에서 직접 스트리밍하여 연산 효율을 높였습니다. 결과적으로 M=1에서 5배 이상의 속도 향상을 기록했습니다.

일반적인 교훈은 다음과 같습니다: 고성능 GPU 환경에서는 커널의 범용성보다 워크로드의 특성(토큰 수, 메모리 접근 패턴)에 맞춘 커스텀 커널 디스패치가 필수적이라는 점입니다. 특히 통신이 병목인 경우, 파이프라인의 단계를 세분화하여 fabric hop을 숨기는 전략이 매우 효과적입니다.

참고 자료

• FlashInfer GitHub
• CUDA Programming Guide

참고 자료

• https://github.com/flashinfer-ai/flashinfer

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