DeepSeek의 GPU GEMM 커널 라이브러리 오픈소스 공개

핵심 요약

• DeepGEMM은 NVIDIA GPU에 최적화된 통합 텐서 코어 커널 라이브러리로, 대언어모델(LLM)의 핵심 연산인 GEMM(FP8, FP4, BF16), Mega MoE, MQA 스코링 등을 단일 CUDA 코드베이스에 통합했다.
• CUTLASS와 CuTe의 개념을 참고하되 무거운 템플릿 의존성을 피하고, 런타임 JIT 컴파일(DeepJIT)로 설치 시 CUDA 컴파일을 요구하지 않는다.
• 각 행렬 크기에 따라 전문가 튜닝 라이브러리 수준 이상의 성능을 제공하며, H800에서 최대 1550 TFLOPS를 달성한다.
• SM90/SM100 아키텍처 GPU를 지원하고, 정방향·역방향 모두 MoE 가중치 그래디언트 커널을 제공한다.
• Mega MoE는 EP 디스패치, 선형 연산, SwiGLU, EP 병합을 단일 메가커널로 융합하여 NVLink 통신과 텐서 코어 계산을 겹쳐 실행한다.
• 그룹화된 GEMM을 지원하는데, M축만 그룹화하고 N·K는 고정하는 설계로 MoE 전문가들이 동일 형태를 공유하는 시나리오에 맞춰져 있다.
• 마스크된 그룹화 GEMM으로 추론 디코딩 단계에서 CUDA 그래프 활용 시 각 전문가가 받는 토큰 수를 동적으로 처리한다.
• DeepSeek v3.2 인덱서용 MQA 커널은 비페이징(프리필) 및 페이징(디코딩) 버전을 제공하며, FP8/FP4 데이터와 압축된 로짓 출력을 지원한다.
• 다양한 환경 변수와 유틸리티 함수로 SM 개수 제한, 텐서 코어 사용률 조정, JIT 캐시 관리, 컴파일 최적화 설정이 가능하다.
• Python 3.8+, PyTorch 2.3+, CUDA Toolkit 12.9+, CUTLASS 4.0+ 요구사항이 있고, 런타임 JIT 컴파일로 낮은 CPU 오버헤드를 유지한다.
• MIT 라이선스로 공개되었으며, CUTLASS 프로젝트에서 영감을 받아 개발됐다.