[flashinfer] FlashInfer의 MXFP8 GEMM 성능 최적화: Warp-level Split-K 도입

PR 링크: flashinfer-ai/flashinfer#6167 상태: Merged | 변경: +1332 / -16

들어가며

최신 GPU 아키텍처(SM100/SM103)에서 MXFP8(Microscaling Formats)을 활용한 행렬 곱셈(GEMM)은 LLM 추론 성능의 핵심입니다. 하지만 기존의 tcgen05 기반 러너만으로는 특정 입력 크기(특히 작은 M)에서 최적의 성능을 내기 어려웠습니다. 이번 PR은 warp-level split-K 기법을 도입하여, 4개의 연산 워프(compute warp)가 K 차원을 분할하여 처리하고 공유 메모리 메일박스를 통해 결과를 합산함으로써, 작은 M과 K 차원에서 최대 1.8배 이상의 성능 향상을 달성했습니다.

코드 분석

1. flashinfer/gemm/gemm_base.py: 러너 선택 로직 개선

기존에는 단일 러너만 존재했으나, 이제는 _cute_dsl_mxfp8_runners를 통해 입력 크기에 따라 warp-level 러너와 tcgen05 러너를 동적으로 선택합니다.

# Before: 단일 러너 사용
runners: List[TunableRunner] = [
    backend_to_runner_factory[cur_backend]() for cur_backend in backends
]

# After: 입력 크기에 따른 러너 리스트 구성
if warp.supports_inputs(inputs) and (m <= 8 or k <= 2048):
    return [warp, tcgen05]
return [tcgen05, warp]

이 변경은 오토튜닝 과정에서 두 러너를 모두 프로파일링하고, 특정 조건(M <= 8 또는 K <= 2048)에서 warp-level 러너를 우선적으로 선택하게 합니다.

2. flashinfer/gemm/kernels/dense_mxfp8_gemm_warp_splitk.py: Warp-level Split-K 커널 구현

이 커널은 cp.async 파이프라이닝을 사용하여 데이터를 로드하고, 4개의 워프가 각각 K의 일부를 계산합니다. 특히 FP8 mma.sync가 없는 환경을 고려하여 cvt.rn.f16x2.e4m3x2를 통해 FP16으로 언팩 후 연산을 수행합니다.

# 커널 핵심 로직: FP8 E4M3를 FP16으로 변환하여 연산
# sm_100/sm_103 have no FP8 mma.sync: ptxas lowers m16n8k32.e4m3 to FP16 unpacks
_MMA_SHAPE = (16, 8, 32)

왜 이게 좋은가

이 최적화의 핵심은 연산 효율성(Compute Efficiency)과 데이터 재사용성입니다.

1. 성능 향상: 벤치마크 결과, M=1, N=2304, K=1536 환경에서 기존 대비 1.80배의 속도 향상을 보였습니다. 이는 작은 행렬 연산에서 발생하는 오버헤드를 워프 레벨의 병렬 처리를 통해 효과적으로 상쇄했기 때문입니다.
2. 유연한 폴백(Fallback): 오토튜닝 과정에서 특정 입력이 warp-level 러너에 적합하지 않을 경우, 기존 tcgen05 러너로 자연스럽게 전환되는 구조를 갖추어 범용성을 확보했습니다.
3. 교훈: GPU 커널 최적화 시, 단순히 큰 행렬에 집중하기보다 작은 M/K 차원에서의 병목(데이터 로드 및 워프 유휴 시간)을 파악하고, 이를 워프 레벨의 분할 연산으로 해결하는 것이 중요함을 보여줍니다.

리뷰어 피드백 반영

리뷰 과정에서 StudyingShao는 PDL 의존성 동기화 문제와 오토튜닝 캐시의 유효성 검증에 대해 중요한 지적을 했습니다. 특히 griddepcontrol_wait()를 통해 가중치 로드와 연산 간의 의존성을 명확히 함으로써, 비동기 연산 시 발생할 수 있는 결과 오류(결과값이 0으로 나오는 현상)를 방지했습니다.

참고 자료

• Cutlass CuTe Documentation — 본 PR의 기반이 된 CuTe 프로그래밍 모델 문서

참고 자료

• https://github.com/NVIDIA/cutlass/blob/main/media/docs/cute/00_introduction.md

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