deepseek-ai/DeepGEMM

해결하는 문제

DeepGEMM은 현대적인 대규모 언어 모델(LLM)의 핵심 계산 프리미티브를 최적화하도록 설계된 고성능 Tensor Core 커널 라이브러리입니다. NVIDIA SM90 및 SM100 아키텍처를 특별히 타겟팅하여, 전문가가 튜닝한 라이브러리에 필적하거나 이를 능가하는 고도로 최적화된 일반 행렬 곱셈(GEMM) 커널을 제공합니다.

작동 방식

이 라이브러리는 경량 Just-In-Time(JIT) 모듈을 사용하여 런타임에 커널을 컴파일하므로 설치 중 CUDA 컴파일이 필요하지 않습니다. 다음과 같은 특화된 커널을 구현합니다:

• Dense GEMMs: FP8, FP4, BF16을 포함한 다양한 정밀도 형식을 지원합니다.
• Grouped GEMMs: 전문가가 동일한 형상을 공유하지만 서로 다른 토큰 수를 처리하는 Mixture-of-Experts(MoE) 모델에 최적화되었습니다(contiguous 및 masked 레이아웃).
• Mega MoE: NVLink 통신(EP dispatch/combine)과 Tensor Core 계산(Linear 1, SwiGLU, Linear 2)을 중첩시키는 융합된 "mega-kernel"입니다.
• MQA Kernels: DeepSeek v3.2의 lightning indexer를 위한 특화된 스코어링 커널입니다.

대상 사용자

LLM 추론 및 학습을 연구하는 개발자와 연구자, 특히 NVIDIA H100/H800 (SM90) 또는 최신 (SM100) GPU를 사용하는 분들과 GPU 커널 최적화를 배우기 위한 깔끔하고 접근하기 쉬운 리소스를 찾는 분들을 위한 것입니다.

주요 특징

• 극한의 성능: H800 GPU에서 최대 1550 TFLOPS를 달성합니다.
• JIT 컴파일: 커널이 런타임에 컴파일되어 설치 및 배포가 간편합니다.
• 고급 MoE 지원: 통신과 계산을 중첩시키는 융합된 Mega MoE 커널을 포함합니다.
• 폭넓은 정밀도 지원: 처리량을 극대화하기 위해 FP8 및 FP4 정밀도에 최적화되었습니다.
• C++20 기반: 복잡한 템플릿에 대한 과도한 의존을 피하는 현대적인 CUDA 코드베이스로 구축되었습니다.