DeepGEMM: GPU용 간결·고성능 BLAS 커널 라이브러리

AI 요약

DeepGEMM은 대규모 언어 모델에 필요한 GPU 연산 커널을 하나의 CUDA 코드베이스에 모은 라이브러리입니다. FP8·FP4·BF16 GEMM, MoE(Mixture of Experts)용 연산, Lightning Indexer의 MQA 점수 계산, HyperConnection(HC) 등을 제공합니다. CUTLASS와 CuTe의 개념을 참고하되 템플릿과 추상화에 크게 의존하지 않고, 제한된 수의 커널 함수로 구현했습니다. 커널은 DeepJIT로 실행 시점에 컴파일하므로 설치 단계에서 CUDA 코드를 미리 빌드하지 않습니다.

지원 연산과 GPU 아키텍처

일반 GEMM은 D = C + A @ B 형태로 계산합니다. SM90 구현은 A가 비전치, B가 전치된 NT 레이아웃을 지원하고, SM100 구현은 NT·TN·NN·TT 레이아웃을 지원합니다. 예를 들어 fp8_gemm_nt는 D = C + A @ B.T를 계산합니다. 입력 전치나 FP8 변환은 사용자가 별도로 처리해야 하며, 라이브러리는 주로 GEMM 커널 최적화에 초점을 둡니다.

스케일링 인자는 TMA에 맞는 전치 레이아웃이어야 합니다. SM90은 FP32 스케일을, SM100은 UE8M0 스케일 네 개를 torch.int 하나에 담은 패킹 형식을 사용합니다. 연속형(contiguous) MoE GEMM은 M축으로 그룹을 나누며 N과 K는 고정합니다. 전문가별 토큰 수가 다른 경우 토큰을 이어 붙이고, 각 전문가 구간을 GEMM의 M 블록 크기에 맞춰 정렬해야 합니다. MoE 가중치의 역전파에는 K축 그룹 API도 제공합니다. 디코딩 중 CUDA Graph를 사용해 CPU가 전문가별 토큰 수를 알지 못하는 상황에는 마스크 기반 grouped GEMM으로 유효한 구간만 계산합니다.

MoE와 인덱서 커널

Mega MoE는 전문가 병렬화(EP) 디스패치, 첫 번째 선형층, SwiGLU, 두 번째 선형층, EP 결합을 하나의 커널로 묶습니다. FP8xFP4 또는 FP8xFP8 가중치를 처리하며, NVLink 통신과 Tensor Core 계산을 겹쳐 실행합니다. 다중 프로세스 실행과 대칭 메모리가 필요하고, 입력을 대칭 버퍼에 복사한 뒤 커널을 호출합니다. 가중치는 요구되는 레이아웃으로 변환해야 하며, 선택적으로 locality domain에 맞춰 배치할 수 있습니다.

인덱서 점수 커널은 비페이징(prefill)과 페이징(decoding) 버전으로 나뉩니다. 쿼리와 키-값 벡터의 점수를 계산한 뒤 ReLU를 적용하고 가중치를 곱해 합산합니다. SM100에서는 패킹된 UE8M0 스케일을 사용하는 MXFP4·MXFP8 입력을 지원합니다.

설치와 개발 설정

요구 환경은 NVIDIA SM90 또는 SM100 GPU, Python 3.8 이상, C++20의 <format>을 지원하는 컴파일러와 표준 라이브러리, CUDA Toolkit 12.9 이상, PyTorch 2.3 이상, CUTLASS 4.0 이상입니다. 저장소를 서브모듈과 함께 복제한 뒤 개발용 또는 설치용 스크립트를 실행하고 Python에서 deep_gemm을 불러옵니다. JIT 설정에는 컴파일러 경로와 캐시 디렉터리 지정, 컴파일 명령 출력, PTX·SASS 덤프, 레지스터 스필 검사 등이 포함됩니다. 저장소는 2025년 4월 H800에서 최대 1550 TFLOPS를 기록했다고 안내합니다.

원문: GitHub / 번역·요약: Trawling