[flashinfer] FlashInfer: NVIDIA Blackwell(SM120)을 위한 고성능 FP8 MoE GEMM 최적화

PR 링크: flashinfer-ai/flashinfer#3891 상태: Merged | 변경: +4070 / -1035

들어가며

최신 LLM 추론에서 Mixture-of-Experts(MoE) 모델의 효율성은 매우 중요합니다. 특히 NVIDIA의 최신 Blackwell(SM120) 아키텍처에서 FP8 연산을 효과적으로 활용하는 것은 추론 속도 향상의 핵심입니다. 이번 FlashInfer 업데이트에서는 CuTe C++를 사용하여 SM120 아키텍처에 최적화된 FP8 groupwise MoE GEMM 커널을 도입했습니다. 기존의 cutlass 기반 구현이나 cudnn 대비, 토큰 패딩을 제거하고 하드웨어 특화된 타일링 전략을 통해 성능을 크게 개선했습니다.

코드 분석

1. csrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_op.cu

이 파일은 Python 인터페이스와 C++ 커널을 연결하는 핵심 엔트리 포인트입니다. scale_granularity_mnk를 (1, 128, 128)로 고정하여 groupwise scaling을 지원하며, 입력 텐서의 정렬과 패딩을 검증합니다.

// Before: N/A (New Feature)
// After: 입력 텐서의 정렬 및 패딩 검증 로직 추가
auto ceil_div = [](int64_t x, int64_t y) { return (x + y - 1) / y; };
auto compute_padded_offset = [](int64_t offset, int64_t problem_idx) {
  constexpr int64_t kAlignment = 4;
  return (offset + problem_idx * (kAlignment - 1)) / kAlignment * kAlignment;
};

2. csrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cu

이 파일은 하드웨어 리소스(SM) 활용을 극대화하기 위한 타일 크기 선택 로직을 포함합니다. select_fp8_flat_tile_m 함수는 연산량과 SM 개수를 고려하여 최적의 타일 크기(8, 32, 64, 128)를 동적으로 선택합니다.

// 타일 크기 선택 로직
if (shape_m <= 8 || underfilled_flat_work || bounded_small_m_work) {
  return 8;
}
// ... (중략) ...
return 32;

왜 이게 좋은가

이번 최적화의 핵심은 'Zero-padding'과 'Multi-tile selector'입니다.

1. 토큰 패딩 제거: 기존 cutlass 구현은 m_indptr의 4배수 정렬을 위해 불필요한 패딩이 발생했으나, 이번 구현은 토큰 패딩 없이 (cum_m, k) 레이아웃을 직접 사용하여 메모리 대역폭을 절약합니다.
2. 성능 향상: RTX PRO 6000 Blackwell(SM120) 환경에서 벤치마크 결과, 작은 M(token 수) 환경에서 기존 방식 대비 100%~200% 이상의 성능 향상을 보였습니다. 특히 m_pe=1인 경우 cutlass 대비 약 178%~194% 더 빠른 속도를 기록했습니다.

교훈: 하드웨어 아키텍처(SM120)에 최적화된 CuTe 레이아웃을 사용하고, 데이터의 정렬(Alignment)을 정밀하게 제어함으로써 범용 라이브러리보다 훨씬 높은 효율을 낼 수 있습니다.

리뷰어 피드백

리뷰 과정에서 sf_fp8_tma_load.cuh 내의 불필요한 코드(MXFP8 전용 로직)에 대한 논의가 있었으나, 이는 FP8 float-scale 커널의 특성을 고려할 때 false positive로 확인되었습니다. 또한, runner 헤더의 네이밍 규칙을 유지하여 기존 MXFP8 커널과의 대칭성을 보존했습니다.

참고 자료

• https://github.com/flashinfer-ai/flashinfer
• https://developer.nvidia.com/blog/nvidia-cute-a-new-way-to-program-tensor-cores/

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