DeepGEMM
A CUDA kernel library supports LLM training and inference on NVIDIA SM90 and SM100 GPUs, with low-precision matrix operations and fused MoE computation.
DeepGEMM is an open-source CUDA kernel library for developers building LLM training and inference systems on their own NVIDIA GPUs. It supplies matrix operations and specialized computations for mixture-of-experts (MoE) models.
- Low-precision matrices: Matrix multiplication supports FP8, FP4 and BF16 for model computation.
- Grouped expert computation: Grouped FP8 kernels handle experts receiving different token counts, support masked decoding with CUDA graphs, and compute MoE weight gradients.
- Fused MoE: Mega MoE combines expert dispatch, two linear operations, SwiGLU activation and output combination in one kernel, overlapping NVLink communication with GPU computation.
- Attention and HyperConnection: Specialized kernels calculate MQA scores for the lightning indexer during prompt processing and decoding, alongside HyperConnection operations.
DeepGEMM uses the MIT license and requires NVIDIA SM90 or SM100 hardware, PyTorch and the CUDA toolkit. Mega MoE requires multiple processes with symmetric memory. DeepJIT compiles kernels at runtime.
Developers must handle input transposition and FP8 casting separately. Matrix layout support depends on the GPU: SM90 supports the NT layout, while SM100 supports NT, TN, NN and TT.
