DeepGEMM: 고성능 텐서 코어 커널 라이브러리

DeepGEMM은 현대 대규모 언어 모델(LLMs)의 핵심 계산 프리미티브인 GEMMs (FP8, FP4, BF16), 오버랩된 통신을 갖춘 융합 MoE (Mega MoE), 초고속 인덱서용 MQA 스코어링, HyperConnection (HC) 등을 단일하고 응집력 있는 CUDA 코드베이스로 통합한 고성능 텐서 코어 커널 라이브러리입니다. 모든 커널은 DeepJIT을 통해 런타임에 컴파일되므로 설치 시 CUDA 컴파일이 필요 없습니다.

DeepGEMM은 CUTLASS와 CuTe의 일부 개념을 활용하지만, 이들의 템플릿이나 대수학(algebras)에 크게 의존하지 않습니다. 이 라이브러리는 핵심 커널 함수가 제한적이어서 설계가 단순하며, NVIDIA GPU 커널 최적화 기술을 학습하기 위한 깨끗하고 접근성이 높은 리소스로 만들어졌습니다.

가벼운 디자인에도 불구하고, DeepGEMM의 성능은 다양한 행렬 크기에서 전문가가 튜닝한 라이브러리와 동등하거나 그 이상의 성능을 보여줍니다.

• 2026.09.30

• DeepGEMM Ascend가 이용 가능합니다! 자세한 내용은 DeepGEMM-Ascend를 확인하세요.

• 지역성 도메인(locality domain) 기능 등 추가 최적화는 #462를 확인하세요.

• 2026.09.10: Sparse Indexer, Mega Gate, Mega mHC, DeepJIT, MoE 및 Indexer 최적화 등이 이루어졌습니다.

• 자세한 내용은 #432를 참고하십시오.

• 2026.04.16: Mega MoE, FP8xFP4 GEMM, FP4 Indexer, PDL, 더 빠른 JIT 컴파일 등 개선 사항이 포함되었습니다.

• 2025.09.28: DeepGEMM은 DeepSeek v3.2의 초고속 인덱서용 스코어링 커널(가중치 ReLU MQA 로짓)을 지원하게 되었습니다.

• 자세한 내용은 #200을 참고하십시오.

• 2025.07.20: DeepGEMM은 SM90/SM100 모두를 지원하며, 낮은 CPU 오버헤드를 가진 JIT CPP 모듈로 전체 리팩토링되었습니다.

• NVCC 12.9가 FFMA 인터리빙을 자동으로 처리하므로, 모든 후처리 최적화는 더 이상 지원되지 않습니다.

• 자세한 내용은 #112를 참고하십시오.

• 2025.05.14: DeepGEMM은 밀집(dense) 및 MoE 역전파에 대한 가중치 기울기 커널을 제공합니다! 자세한 내용은 #95를 확인하세요.

• 2025.04.18: DeepGEMM이 H800에서 최대 1550 TFLOPS에 도달했습니다! 자세한 내용은 #74, #78, #81, #86 및 340d988을 확인하세요.

• NVIDIA SM90 또는 SM100 아키텍처 GPU

• Python 3.8 이상

• C++20을 지원하는 컴파일러 및 표준 라이브러리
  <format>

support - CUDA Toolkit 12.9 이상

• PyTorch 2.3 이상
• CUTLASS 4.0 이상 (Git submodule으로 클론 가능)

# Submodule은 반드시 클론해야 함
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git
cd DeepGEMM
...

cat install.sh
./install.sh

그런 다음, Python 프로젝트에서 deep_gemm을 임포트하고 사용하세요!

이 라이브러리는 명명 규칙(D = C + A @ B)을 가진 NVIDIA GPU용 최적화된 GEMM 커널을 제공합니다. 입력 텐서의 레이아웃은 NT(non-transposed A, transposed B)입니다. SM90 구현은 NT 메모리 레이아웃만 지원하는 반면 (row-major, col-major), SM100 구현은 모든 메모리 레이아웃(NT, TN, NN, TT)을 지원합니다. 예를 들어, fp8_gemm_nt는 D = C + A @ B.T를 수행합니다.

두 아키텍처 모두에서 LHS 스케일링 팩터는 TMA-aligned 및 전치된 레이아웃을 가져야 합니다. 그리고 SM90과 SM100의 스케일링 팩터 데이터 형식은 다릅니다:

• SM90은 FP32 형식의 스케일링 팩터를 요구합니다.
• SM100은 4개의 UE8M0을 단일 torch.int에 패킹하는 packed UE8M0 형식의 스케일링 팩터를 요구합니다.

사용자는 입력 전치(input transposition)나 FP8 캐스팅과 같은 연산을 별도로 처리해야 하며, 이를 구현하거나 이전 커널에 통합해야 합니다. 라이브러리가 일부 간단한 PyTorch 유틸리티 함수를 제공하지만, 이는 성능 저하를 초래할 수 있습니다. 우리의 주요 초점은 GEMM 커널 자체의 최적화입니다.

기본 비그룹(non-grouped) FP8 GEMM을 수행하려면 fp8_gemm_{nt, nn, tn, tt} 함수를 호출하세요. 더 자세한 내용은 함수 문서를 참조하십시오.

CUTLASS의 전통적인 그룹화된 GEMM과 달리 DeepGEMM은 M축만 그룹화하며, N과 K는 고정되어야 합니다. 이 설계는 MoE(Mixture-of-Experts) 모델에서 전문가들이 동일한 형태를 공유하는 시나리오에 맞춰져 있습니다. 각 전문가가 가변적인 수의 토큰을 처리할 수 있는 학습 순전파(training forward passes) 또는 추론 사전 채우기(inference prefilling)의 경우, 이러한 토큰들을 단일 텐서로 연결하며, 이를 '연속적(contiguous)' 레이아웃이라고 부릅니다. 각 전문가 세그먼트는 GEMM M 블록 크기에 맞춰 정렬되어야 함에 유의하십시오 (get_mk_alignment_for_contiguous_layout() 참조). 더 자세한 내용은 m_grouped_fp8_gemm_{nt, nn}_contiguous 함수 문서를 참고하십시오.

또한 MoE 가중치 역전파(weight backward)를 위한 K축 그룹화 API도 제공합니다 (이 경우 M과 N은 고정되어야 함). 더 자세한 내용은 k_grouped_fp8_gemm_tn_contiguous를 참조하십시오.

추론 디코딩 단계에서 CUDA 그래프가 활성화되고 CPU가 각 전문가가 받는 토큰 수를 알지 못하는 경우, 마스크 그룹화된 GEMM을 지원합니다. 마스크 텐서를 제공함으로써 커널은 유효한 부분만 계산합니다.

이 목적을 위해 m_grouped_fp8_gemm_nt_masked를 사용하고 관련 문서를 참고하십시오. 사용 예로는 DeepEP의 저지연(low-latency) 커널 출력을 입력으로 사용하는 것이 있습니다.

해당 커널 패밀리는 비페이지(non-paged, 사전 채우기용) 버전과 페이지(paged, 디코딩용) 버전 두 가지가 있습니다. 예시로 비페이지 버전인 fp8_fp4_mqa_logits를 사용합니다. 이의 주요 입력은 다음과 같습니다:

• q: (q_data, q_sf) 튜플 (q 데이터)
• kv: 모양이 [seq_len_kv, head_dim]인 (kv_data, kv_sf) 튜플. SM100은 패킹된 UE8M0 스케일의 MXFP4/MXFP8 데이터를 허용합니다.
• weights: 모양이 [seq_len, num_heads]인 텐서 (SM100에서는 BF16 사용).
• cu_seq_len_k_start 및 cu_seq_len_k_end: 모양이 [seq_len]인 정수 텐서.
• max_seqlen_k: 모든 쿼리 행의 최대 유효 KV 범위.

출력은 [seq_len, max_seqlen_k]로 압축됩니다. i번째 행은 컬럼 제로에서 시작하는 자신의 유효 KV 범위를 저장합니다. q의 토큰 i에 대해, 이는 모든 토큰 j를 반복하게 됩니다.

from [cu_seq_len_k_start[i], cu_seq_len_k_end[i])

, 그리고 해당 압축된 로짓(logit)을 계산합니다:

kv_j = kv[0][j, :] * kv[1][j].unsqueeze(1) # [head_dim]
out_ij = q[i, :, :] @ kv_j # [num_heads]
out_ij = out_ij.relu() * weights[i, :] # [num_heads]
...

더 자세한 내용과 paged 버전인 fp8_fp4_paged_mqa_logits에 대해서는 tests/test_attention.py를 참조하십시오.

Mega MoE는 EP 디스패치(dispatch), linear 1 및 linear 2 (FP8xFP4 또는 FP8xFP8), SwiGLU, 그리고 EP 조합을 단일 메가-커널(mega-kernel)로 통합하며, NVLink 통신과 텐서 코어 연산을 오버랩합니다. 이는 대칭 메모리(symmetric memory)를 사용하는 다중 프로세스 실행을 필요로 합니다. 사용법:

# 대칭 메모리 버퍼 할당
# 참고: PyTorch >= 2.9 필요
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
...

다중 프로세스 설정 및 벤치마킹을 포함한 전체 예제는 tests/test_mega_moe.py를 참조하십시오.

이 라이브러리는 위의 커널 외에도 몇 가지 유틸리티 함수를 제공합니다:

deep_gemm.set_num_sms / get_num_sms: 사용할 최대 SM(Streaming Multiprocessor) 개수를 설정/조회<br>deep_gemm.set_tc_util / get_tc_util: 근사화된 텐서 코어 활용률을 설정/조회<br>deep_gemm.set_pdl / get_pdl: 프로그램 의존적 실행(Programmatic Dependent Launch, PDL) 활성화/비활성화<br>deep_gemm.use_deterministic_algorithms: 결정론적 알고리즘 활성화/비활성화<br>deep_gemm.set_mk_alignment_for_contiguous_layout / get_mk_alignment_for_contiguous_layout: 연속 레이아웃에 대한 그룹 레벨 M/K 정렬을 설정/조회<br>deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout: 이론적인 최소 M/K 정렬 값을 조회<br>deep_gemm.set_ignore_compile_dims: JIT 컴파일 중 무시할 차원을 구성<br>deep_gemm.set_block_size_multiple_of / get_block_size_multiple_of: 블록 크기를 주어진 값의 배수로 제한<br>deep_gemm.transform_sf_into_required_layout: 스케일링 팩터(scaling factors)를 필요한 레이아웃으로 변환<br>deep_gemm.get_tma_aligned_size: 필요한 TMA 정렬 크기를 조회<br>deep_gemm.get_mn_major_tma_aligned_tensor

: MN-major TMA-aligned 텐서를 가져옴 deep_gemm.get_mn_major_tma_aligned_tensor
: MN-major TMA-aligned 텐서 (UE8M0으로 패킹된 FP32 포함)를 가져옴 deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
: K-그룹화 GEMM 패킹 커널

이 라이브러리는 또한 유용한 몇 가지 환경 변수를 제공합니다:

각 DG_JIT_* 변수는 설정되지 않은 경우 해당 전역 DJ_JIT_* 변수로 폴백(fallback)됩니다.

• 일반
  DG_JIT_DEBUG
  : '0' 또는 '1', 컴파일러 명령어 및 PTXAS 출력, 로드 시간 보고, 라인 정보, 그리고 PTX/SASS 덤프를 포함한 JIT 디버깅 기능을 활성화합니다. 기본값은 '0'
  DG_PRINT_CONFIGS
  : '0' 또는 '1', 각 모양에 대해 선택된 구성을 출력합니다. 기본값은 '0'
• JIT 캐시
  DG_JIT_CACHE_DIR
  : 문자열, 컴파일된 커널의 캐시 디렉토리 (또는 :로 구분된 디렉토리 목록); 조회는 모든 경로를 앞에서부터 뒤로 검색하며(첫 번째 일치 항목이 승리), 캐시 미스 발생 시 첫 번째 경로에 컴파일됩니다. 기본값은 $HOME/.dj
• 컴파일러 선택
  DG_JIT_NVCC_COMPILER
  : 문자열, NVCC 컴파일러 경로; 그렇지 않으면 CUDA_HOME, CUDA_PATH, which nvcc, 그리고 /usr/local/cuda를 통해 CUDA를 찾습니다.
  DG_JIT_CPP_STANDARD
  : 정수, C++ 표준 버전. 기본값은 '20'
• 컴파일러 출력
  DG_JIT_PRINT_COMPILER_COMMAND
  : '0' 또는 '1', 컴파일 명령어를 출력합니다. 기본값은 '0'
  DG_JIT_PTXAS_VERBOSE
  : '0' 또는 '1', 상세한 PTXAS 출력을 보여줍니다. 기본값은 '0'
  DG_JIT_CHECK_NO_SPILLS
  : '0' 또는 '1', 컴파일된 커널에 레지스터 스필(register spills)이 없음을 단언합니다. 기본값은 '0'
  DG_JIT_CHECK_NO_LOCAL_MEMORY
  : '0' 또는 '1', 컴파일된 커널에 로컬 메모리 사용이 없음을 단언합니다. 기본값은 '0'
  DG_JIT_PRINT_LOAD_TIME
  : '0' 또는 '1', 커널 로드 시간을 출력합니다. 기본값은 '0'
• 디버그 및 프로파일링
  DG_JIT_WITH_LINEINFO
  : '0' 또는 '1', 프로파일링 도구를 위해 소스 라인 정보를 임베드합니다. 기본값은 '0'
  DG_JIT_DUMP_ASM
  : '0' 또는 '1', PTX와 SASS를 모두 덤프합니다. 기본값은 '0'
  DG_JIT_DUMP_PTX
  : '0' 또는 '1', PTX 출력을 덤프합니다. 기본값은 '0'
  DG_JIT_DUMP_SASS
  : '0' 또는 '1', SASS 출력을 덤프합니다. 기본값은 '0'

기본적으로 DG_COMM_KERNEL_DEBUG는 '0' 또는 '1'이며, 디버깅을 위해 각 Mega MoE 호출 전에 제로 대칭 버퍼를 만듭니다. 기본값은 '0'입니다. 또한 DG_USE_NVIDIA_TOOLS는 '0' 또는 '1'이며, 외부 NVIDIA 도구 하에서 실행할 때 내부 프로파일링을 건너뜁니다. 기본값은 '0'입니다.

• 빌드 옵션
  DG_SKIP_CUDA_BUILD: '0' 또는 '1'이며, 설치 중 CUDA 확장 빌드를 건너뜁니다. 기본값은 '0'입니다. DG_FORCE_BUILD: '0' 또는 '1'이며, 사전 구축된 휠(pre-built wheels)을 다운로드하는 대신 로컬 빌드를 강제합니다. 기본값은 '0'입니다.

추가 예시 및 세부 정보는 테스트 코드를 참조하거나 해당 Python 문서를 검토해 주십시오.

DeepGEMM은 CUTLASS 프로젝트에서 영감을 받았습니다. 개발자분들께 감사와 존경을 표합니다!

이 코드 저장소는 MIT 라이선스 하에 공개되었습니다.

@misc{deepgemm2025,
title={DeepGEMM: clean and efficient BLAS kernel library on GPU},
author={Chenggang Zhao and Zhean Xu and Liang Zhao and Jiashi Li and Chenhao Xu and Anyi Xu and Shengyu Liu and Kexing Zhou and Kuai Yu},
...