DeepGEMM README

The project's README from github.com/deepseek-ai/DeepGEMM. As one markdown file: README.md.

DeepGEMM

DeepGEMM is a unified, high-performance tensor core kernel library that brings together the key computation primitives of modern large language models — GEMMs (FP8, FP4, BF16), fused MoE with overlapped communication (Mega MoE), MQA scoring for the lightning indexer, HyperConnection (HC), and more — into a single, cohesive CUDA codebase. All kernels are compiled at runtime through DeepJIT, requiring no CUDA compilation during installation.

DeepGEMM leverages some concepts from CUTLASS and CuTe, but avoids heavy reliance on their templates or algebras. The library is designed for simplicity, with only a limited number of core kernel functions, making it a clean and accessible resource for learning NVIDIA GPU kernel optimization techniques.

Despite its lightweight design, DeepGEMM's performance matches or exceeds expert-tuned libraries across various matrix shapes.

News

Quick start

Requirements

Development

# Submodule must be cloned
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git
cd DeepGEMM

# Link some essential includes and build the C++ extension
cat develop.sh
./develop.sh

Installation

cat install.sh
./install.sh

Then, import deep_gemm in your Python project, and enjoy!

Interfaces

Notices

This library provides optimized GEMM kernels for NVIDIA GPUs with a naming convention: D = C + A @ B. The input shape layout is NT (non-transposed A, transposed B). While the SM90 implementation supports only the NT memory layout (row-major, col-major), the SM100 implementation supports all memory layouts (NT, TN, NN, TT). For example, fp8_gemm_nt will do a D = C + A @ B.T

For both architectures, the LHS scaling factor is required to have a TMA-aligned and transposed layout. And the data format for the scaling factor of SM90 and SM100 is different:

Please note that operations like input transposition or FP8 casting must be handled separately by the user, please implement or fuse them into prior kernels independently. While the library provides some simple PyTorch utility functions, these may result in slower performance, but our primary focus is on optimizing the GEMM kernels themselves.

Normal dense GEMMs (non-grouped)

To perform a basic non-grouped FP8 GEMM, call the fp8_gemm_{nt, nn, tn, tt} function. For more details, please refer to the function documentation.

Grouped GEMMs (contiguous layout)

Unlike traditional grouped GEMMs in CUTLASS, DeepGEMM groups only the M-axis, while N and K must remain fixed. This design is tailored for scenarios where experts in an MoE model share the same shape. For training forward passes or inference prefilling, where each expert may process a varying number of tokens, we concatenate these tokens into a single tensor, referred to as the "contiguous" layout. Note that each expert segment must be aligned to the GEMM M block size (get_mk_alignment_for_contiguous_layout()). For more information, please refer to the m_grouped_fp8_gemm_{nt, nn}_contiguous function documentation.

We also provide a K-axis-grouped API for MoE weight backward (with M and N must remain fixed), please refer to k_grouped_fp8_gemm_tn_contiguous for more information.

Grouped GEMMs (masked layout)

During the inference decoding phase, when CUDA graph is enabled and the CPU is unaware of the number of tokens each expert receives, we support masked grouped GEMMs. By providing a mask tensor, the kernel computes only the valid portions.

Use m_grouped_fp8_gemm_nt_masked for this purpose and consult the relevant documentation. An example usage is to use the output of low-latency kernels from DeepEP as input.

V3.2 MQA kernels for the indexer

The kernel family has two versions, non-paged (for prefilling) and paged (for decoding). Take the non-paged version fp8_fp4_mqa_logits as an example. Its main inputs are:

The output is compressed to [seq_len, max_seqlen_k]; row i stores its valid KV span starting at column zero. For each token i in q, it will iterate all tokens j from [cu_seq_len_k_start[i], cu_seq_len_k_end[i]), and calculate the corresponding compressed logit as:

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]
out_ij = out_ij.sum()  # Scalar

For more details and the paged version fp8_fp4_paged_mqa_logits, please refer to tests/test_attention.py.

Mega MoE

Mega MoE fuses and overlaps EP dispatch, linear 1 and linear 2 (FP8xFP4 or FP8xFP8), SwiGLU, and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation. It requires multi-process launch with symmetric memory. Usage:

# Allocate symmetric memory buffer
# NOTES: requires PyTorch >= 2.9
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
    group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden,
    mma_type='fp8xfp4',  # Use 'fp8xfp8' for FP8 routed-expert weights
)

# Transform weights (FP4 or FP8 with UE8M0 SF) into the required layout
transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)

# (Optional) Localize weights into locality domains
transformed_l1 = (deep_gemm.localize(transformed_l1[0]), transformed_l1[1])
transformed_l2 = (deep_gemm.localize(transformed_l2[0]), transformed_l2[1])
deep_gemm.destroy_localizer()

# Copy inputs into the buffer before each call
# You may fuse these into previous kernels
buffer.x[:num_tokens].copy_(x_fp8)
buffer.x_sf[:num_tokens].copy_(x_sf)
buffer.topk_idx[:num_tokens].copy_(topk_idx)
buffer.topk_weights[:num_tokens].copy_(topk_weights)

# Run the fused mega MoE kernel
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)

For the full example with multi-process setup and benchmarking, please refer to tests/test_mega_moe.py.

Utilities

The library provides some utility functions besides the above kernels:

The library also provides some environment variables, which may be useful:

Each DG_JIT_* variable falls back to the corresponding global DJ_JIT_* variable when unset.

For additional examples and details, please refer to the test code or review the corresponding Python documentation.

Acknowledgement

DeepGEMM is inspired by the CUTLASS project. Thanks and respect to the developers!

License

This code repository is released under the MIT License.

Citation

@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},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/DeepGEMM}},
}

DeepGEMM (deepseek-ai/DeepGEMM)

Searching Database…

9K ·GitHub
DeepGEMM
What is this?DeepGEMM is a unified CUDA kernel library that provides high-performance tensor core GEMM operations (FP8, FP4, BF16), fused MoE kernels, and attention-indexer kernels used in large language models such as DeepSeek.
Is this a framework or a Python library I just import?It is a Python library backed by CUDA kernels — you import 'deepgemm' directly into your project rather than running it as a standalone application.
How do I add it to my project?Clone the repo (CUTLASS is pulled in via a Git submodule), build the CUDA extension, then simply import deepgemm in Python — individual kernels are compiled just-in-time on first use through DeepJIT, so no CUDA compilation step is needed during installation itself.
What GPU and software stack do I need?An NVIDIA SM90 (Hopper) or SM100 (Blackwell) GPU, Python 3.8+, a C++20-capable compiler, CUDA Toolkit 12.9+, PyTorch 2.3+, and CUTLASS 4.0+ (fetched via Git submodule).
Does it handle input preprocessing like FP8 casting for me?No — operations such as input transposition or FP8 casting must be implemented or fused by the user into prior kernels; DeepGEMM only optimizes the GEMM kernels themselves.
What is Mega MoE?Mega MoE fuses EP dispatch, two linear layers (FP8xFP4 or FP8xFP8), SwiGLU, and EP combine into a single mega-kernel that overlaps NVLink communication with tensor core computation, requiring a multi-process launch with symmetric memory.
What is DeepJIT?DeepJIT is DeepGEMM's just-in-time compilation system: kernels are compiled the first time they are used rather than during package installation, with environment variables available to control caching, debugging, and compiler behavior.
How fast is DeepGEMM compared to hand-tuned libraries?The project claims performance matching or exceeding expert-tuned libraries across various matrix shapes, and reports achieving up to 1550 TFLOPS on an H800 GPU.
Which GPU architectures are actually supported?DeepGEMM targets NVIDIA SM90 (Hopper) and SM100 (Blackwell) architectures; SM90 only supports the NT memory layout while SM100 supports NT, TN, NN, and TT layouts.
What are similar projects?tilelang, TensorRT-LLM, FlashMLA, DeepEP, DeepGEMM-Ascend, cupy, flash-attention, LeetCUDA
Which topics is it in?tensor, cuda-kernels, jit, gpu-programming, mixture-of-experts, attention-mechanism, high-performance, gpu-acceleration, quantization, transformer, large-language-models, artificial-intelligence
Creator
deepseek-ai
Size
2.4 MB
Stars
9K
Updated
CommitOct 14Releasev2.1.1.post3Oct 15
Age
1y
Commits · 90 Days