[flashinfer] FlashInfer SM107(Rubin) 아키텍처를 위한 MXFP8 GEMM 최적화 분석

PR 링크: flashinfer-ai/flashinfer#6709 상태: Merged | 변경: +96 / -7

들어가며

최신 GPU 아키텍처인 SM107(Rubin) 환경에서 mm_mxfp8(backend="cute-dsl")을 사용할 때, 소규모 M(batch size) 행렬 곱셈 연산이 비효율적으로 동작하는 문제가 있었습니다. 기존 SM107 전용 커널은 최소 64개의 출력 열을 커버하는 큰 타일을 사용하는데, M이 작을 경우(예: M <= 16) GPU의 SM(Streaming Multiprocessor) 자원을 충분히 활용하지 못하고 HBM 대역폭의 일부만 사용하는 병목 현상이 발생했습니다. 본 PR은 SM107 환경에 최적화된 소규모 M 커널 러너들을 등록하고, 상황에 맞는 최적의 커널을 선택하도록 하여 성능을 개선했습니다.

코드 분석

flashinfer/gemm/gemm_base.py

핵심 변경 사항은 SM107을 위한 새로운 러너 등록 함수 _cute_dsl_mxfp8_runners_sm107의 도입입니다. 기존에는 SM107 전용 커널만 사용했으나, 이제는 상황에 따라 warp split-K 및 low-M 경로를 선택할 수 있게 되었습니다.

Before (기존 방식):

# SM107은 단일 커널만 사용
_cute_dsl_gemm_mxfp8_sm107_runner(out_dtype)

After (개선된 방식):

# M 크기에 따라 적절한 러너를 우선순위대로 반환
if m <= 16 and out.is_contiguous():
    return [low_m, sm107] + tail
return [sm107, low_m] + tail

또한, _sm107_low_m_split_k 함수를 추가하여 split-K 슬라이스 수를 동적으로 결정합니다. 이는 GPU SM 개수와 K 블록 크기를 고려하여 연산 효율을 극대화합니다.

tests/gemm/test_mm_mxfp8.py

테스트 코드에서는 M=1, M=24 등 다양한 케이스를 추가하여, 소규모 M에서는 low-M 러너가, 대규모 M에서는 SM107 전용 커널이 선택되는지 검증합니다.

왜 이게 좋은가

이번 최적화는 단순히 커널을 교체하는 것을 넘어, GPU의 연산 자원 활용도를 극대화했습니다.

1. SM 활용도 향상: 기존에는 48개의 CTA만 실행되던 작업이 split-K를 통해 더 많은 CTA로 분산되어 SM 유휴 시간을 줄였습니다.
2. 성능 수치: 측정 결과, M <= 16인 구간에서 기하평균(geomean) 기준 약 1.24배의 성능 향상을 보였으며, 특정 케이스에서는 최대 1.8배까지 성능이 개선되었습니다.
3. 일반적 교훈: 특정 아키텍처에서 고성능 커널이 항상 최선은 아닙니다. 특히 소규모 연산(Small-M)에서는 타일 크기가 작고 유연한 커널이 오버헤드를 줄이는 데 훨씬 유리하다는 점을 시사합니다.

리뷰 과정에서 M=24~32 구간의 회귀(regression) 문제가 발견되었으나, low_m_only 로직을 정교화하여 M <= 16으로 제한함으로써 성능 저하를 방지하고 안정성을 확보했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions

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