-
FP8 행렬 곱셈(GEMM) 을 효율적으로 수행하는 라이브러리로, DeepSeek-V3에서 제안된 미세 조정 스케일링(fine-grained scaling) 방식을 지원함
- 일반 GEMM과 Mix-of-Experts(MoE) 그룹화 GEMM을 모두 지원
- CUDA 기반으로 구현되었으며, 설치 시 별도 컴파일 없이 경량 Just-In-Time(JIT) 모듈을 사용하여 런타임에서 커널을 컴파일함
- 현재 NVIDIA Hopper 텐서 코어 전용으로 지원
- FP8 텐서 코어의 부정확한 누적 연산을 보완하기 위해 CUDA 코어 기반 이중 누적(promotion) 사용
-
CUTLASS 및 CuTe의 일부 개념을 활용하지만, 복잡한 템플릿 의존성을 줄여 약 300줄의 커널 코드만 포함하는 단순한 설계
-
Hopper FP8 행렬 연산 및 최적화 기법을 학습하기에 적합
- 경량 설계에도 불구하고 다양한 행렬 크기에서 전문가 수준으로 튜닝된 라이브러리와 유사하거나 더 나은 성능을 보임
성능 평가
일반 GEMM 성능 (Dense 모델)
- 다양한 행렬 크기에 대해 DeepSeek-V3/R1 추론 환경에서 성능 평가 진행
- NVIDIA H800 GPU(NVCC 12.8) 환경에서 테스트
- 속도 향상(Speedup) 지표는 CUTLASS 3.6을 기반으로 한 내부 최적화 버전과 비교하여 계산
- 주요 성능 요약
- 작은 행렬 크기에서는 최대 2.7배 속도 향상
- 일부 큰 행렬 크기에서는 1.0~1.2배 수준의 성능 유지
- 메모리 대역폭 및 연산 성능을 최적화하여 Hopper 아키텍처에 최적화된 성능 발휘
MoE 모델을 위한 그룹화 GEMM 성능
-
연속된(contiguous) 레이아웃의 그룹화 GEMM
-
M축을 기준으로 그룹화하며, N과 K는 동일하게 유지
-
최대 1.2배 속도 향상을 보이며, 특정 설정에서 1.1배 수준 유지
-
마스킹(masked) 레이아웃의 그룹화 GEMM
-
CUDA 그래프 사용 시 CPU가 각 전문가의 토큰 수를 알 수 없는 경우를 위한 설계
- 마스크 텐서를 입력으로 제공하여 필요한 연산만 수행
- 성능 개선폭은 1.1배~1.2배 수준