DeepGEMM: 세밀한 스케일링을 통한 깨끗하고 효율적인 FP8 GEMM 커널

DeepGEMM

DeepGEMM은 FP8 일반 행렬 곱셈(GEMM)을 위한 라이브러리로, DeepSeek-V3에서 제안된 세밀한 스케일링을 지원함. 이 라이브러리는 일반 및 Mix-of-Experts(MoE) 그룹화된 GEMM을 지원하며, CUDA로 작성되어 설치 시 컴파일이 필요하지 않음. NVIDIA Hopper 텐서 코어를 지원하며, FP8 텐서 코어 누적의 부정확성을 해결하기 위해 CUDA 코어 2단계 누적을 사용함. CUTLASS와 CuTe의 개념을 일부 활용하지만, 템플릿이나 대수에 대한 의존성을 최소화하여 단순함을 유지함. 약 300줄의 코드로 구성된 하나의 핵심 커널 함수로 Hopper FP8 행렬 곱셈 및 최적화 기술을 학습하기에 적합한 자원임. 경량 설계에도 불구하고 다양한 행렬 형태에서 전문가가 조정한 라이브러리와 성능이 같거나 뛰어남.

성능

DeepSeek-V3/R1 추론에서 사용될 수 있는 모든 형태를 H800 SXM5에서 NVCC 12.8로 테스트함. 모든 속도 향상 지표는 CUTLASS 3.6을 기반으로 내부적으로 최적화된 구현과 비교하여 계산됨. 일부 형태에서는 성능이 좋지 않을 수 있으며, 최적화 PR을 환영함.

일반 GEMM (밀집 모델)

• 다양한 행렬 크기에서 DeepGEMM의 성능을 측정한 결과, 특정 크기에서 최대 2.7배의 속도 향상을 보임.

MoE 모델을 위한 그룹화된 GEMM (연속 레이아웃)

• 그룹 수와 각 그룹의 행렬 크기에 따라 최대 1.2배의 속도 향상을 보임.

MoE 모델을 위한 그룹화된 GEMM (마스크 레이아웃)

• 마스크 레이아웃을 사용하여 최대 1.2배의 속도 향상을 보임.

빠른 시작

요구 사항

• Hopper 아키텍처 GPU, sm_90a 지원 필요
• Python 3.8 이상
• CUDA 12.3 이상 (최상의 성능을 위해 12.8 이상 권장)
• PyTorch 2.1 이상
• CUTLASS 3.6 이상

개발

• 서브모듈 클론, 심볼릭 링크 생성, JIT 컴파일 및 모든 GEMM 구현 테스트를 포함한 개발 과정 설명.

설치

• deep_gemm을 Python 프로젝트에 가져와 사용 가능.

인터페이스

주의사항

• 이 라이브러리는 GEMM 커널만 포함하며, NT 형식만 지원함. 전치나 다른 FP8 캐스팅 작업은 독립적으로 구현해야 함.

일반 밀집 GEMM (비그룹화)

• 기본 비그룹화 FP8 GEMM 수행을 위한 함수 제공.

그룹화된 GEMM (연속 레이아웃)

• MoE 모델에서 전문가가 동일한 형태를 공유하는 시나리오에 맞춰 설계됨.

그룹화된 GEMM (마스크 레이아웃)

• 추론 디코딩 단계에서 마스크 텐서를 제공하여 유효한 부분만 계산.

유틸리티

• 다양한 유틸리티 함수 및 환경 변수를 제공하여 성능 최적화에 도움을 줌.

최적화

지속적인 워프 특화

• CUTLASS 디자인을 따르며, 데이터 이동, 텐서 코어 MMA 명령어, CUDA 코어 승격을 겹치게 함.

Hopper TMA 기능

• TMA를 활용하여 데이터 이동을 가속화함.

공통 세부 최적화

• 다양한 최적화 기법을 통해 성능을 향상시킴.

통합 및 최적화된 블록 스케줄러

• 모든 비그룹화 및 그룹화 커널에 대한 스케줄러 제공.

완전한 JIT 설계

• 설치 시 컴파일이 필요 없는 JIT 설계를 통해 성능을 향상시킴.

정렬되지 않은 블록 크기

• 특정 형태에서 SM 활용을 극대화하기 위해 정렬되지 않은 블록 크기를 지원함.

FFMA SASS 인터리빙

• 성능 향상을 위해 FFMA 명령어를 수정하여 워프 수준 병렬성을 향상시킴.

감사

• DeepGEMM은 CUTLASS 프로젝트에서 영감을 받았으며, 개발자들에게 감사와 존경을 표함.

라이선스

• MIT 라이선스로 공개됨.