[flashinfer] FlashInfer에 NVIDIA Rubin(SM107) 및 Blackwell FP8 GEMM 지원 추가

PR 링크: flashinfer-ai/flashinfer#4526 상태: Merged | 변경: +20235 / -684

들어가며

최신 GPU 아키텍처인 NVIDIA Rubin(SM107)과 Blackwell(SM100/SM103)은 고성능 연산을 위해 FP8 데이터 포맷과 새로운 메모리 계층 구조를 도입했습니다. 이번 PR은 FlashInfer에 CuTe DSL 기반의 커널을 추가하여, Rubin 아키텍처에 대한 네이티브 지원과 Blackwell 아키텍처에서의 FP8 Batched GEMM 성능을 극대화하는 것을 목표로 합니다. 특히, 기존 C++ 기반의 bmm_fp8 경로를 대체하거나 보완하는 Python CuTe DSL 구현을 통해 유연한 오토튜닝과 커널 최적화를 가능하게 했습니다.

코드 분석

1. 아키텍처 호환성 및 런타임 감지 (flashinfer/cute_dsl/utils.py)

Rubin 커널은 CuTe DSL 4.8 이상에서 제공되는 rubin_helpers에 의존합니다. 이를 위해 하위 호환성을 유지하면서도 최신 기능을 활용할 수 있도록 지연 로딩(Lazy Loading) 방식을 도입했습니다.

@functools.cache
def is_rubin_cute_dsl_available() -> bool:
    return (
        is_cute_dsl_available()
        and importlib.util.find_spec("cutlass.utils.rubin_helpers") is not None
    )

2. SM107 아키텍처 지원 (flashinfer/cute_dsl/attention/compat.py)

Attention 커널의 TMEM(Tensor Memory) 할당 크기를 SM107에 맞게 조정하여 하드웨어 리소스를 최적으로 활용하도록 했습니다.

"sm_107": 576,

3. FP8 Batched GEMM 디스패치 (flashinfer/gemm/kernels/bmm_fp8_wrapper.py)

Blackwell와 Rubin 아키텍처 간의 GEMM 커널을 통합 관리하고, 아키텍처별로 최적화된 경로를 자동으로 선택하도록 설계되었습니다.

왜 이게 좋은가

이번 최적화의 핵심은 '유연한 아키텍처 대응'과 'DSL 기반의 생산성 향상'입니다.

1. 성능 최적화: 기존 C++ 커널 대비 CuTe DSL을 사용하면 Python 환경에서도 하드웨어 가속기(TMA, UMMA 등)를 직접 제어할 수 있어, 오토튜닝을 통한 커널 최적화가 훨씬 용이합니다.
2. 유지보수성: fused_moe 및 gemm 관련 유틸리티를 common/kernel_utils.py로 통합하여 코드 중복을 제거했습니다.
3. 하위 호환성: is_rubin_cute_dsl_available 함수를 통해, 구형 DSL 환경에서도 기존 기능이 깨지지 않도록 안전장치를 마련했습니다.

리뷰어 피드백 반영

리뷰 과정에서 코드 내 불필요한 주석 제거, 라이선스 헤더 정리, 그리고 cutlass 백엔드 노출 범위에 대한 논의가 있었습니다. 특히 gemm_base.py에서 cutlass 백엔드를 점진적으로 deprecate하고 cute-dsl로 전환하려는 방향성이 확인되었습니다.

결론

이번 PR은 FlashInfer가 차세대 NVIDIA GPU 아키텍처를 빠르게 수용할 수 있는 기반을 마련했습니다. 특히 Rubin 아키텍처의 특수성을 고려한 커널 설계와 Blackwell의 FP8 성능 최적화는 LLM 추론 가속기 분야에서 중요한 이정표가 될 것입니다.

참고 자료

• CuTe DSL Documentation
• FlashInfer GitHub Repository

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