deepseek-ai/DeepGEMM

DeepGEMM: clean and efficient BLAS kernel library on GPU

What it solves

DeepGEMM is a high-performance tensor core kernel library designed to optimize the core computation primitives of modern large language models (LLMs). It provides highly optimized General Matrix Multiplication (GEMM) kernels that match or exceed expert-tuned libraries, specifically targeting NVIDIA SM90 and SM100 architectures.

How it works

The library uses a lightweight Just-In-Time (JIT) module to compile kernels at runtime, removing the need for CUDA compilation during installation. It implements several specialized kernels:

  • Dense GEMMs: Supports various precision formats including FP8, FP4, and BF16.
  • Grouped GEMMs: Optimized for Mixture-of-Experts (MoE) models where experts share shapes but process different token counts (contiguous and masked layouts).
  • Mega MoE: A fused "mega-kernel" that overlaps NVLink communication (EP dispatch/combine) with tensor core computation (Linear 1, SwiGLU, Linear 2).
  • MQA Kernels: Specialized scoring kernels for the lightning indexer in DeepSeek v3.2.

Who it’s for

It is intended for developers and researchers working on LLM inference and training, particularly those utilizing NVIDIA H100/H800 (SM90) or newer (SM100) GPUs, and those looking for a clean, accessible resource to learn GPU kernel optimization.

Highlights

  • Extreme Performance: Achieves up to 1550 TFLOPS on H800 GPUs.
  • JIT Compilation: Kernels are compiled at runtime, simplifying installation and deployment.
  • Advanced MoE Support: Includes fused Mega MoE kernels that overlap communication and computation.
  • Broad Precision Support: Optimized for FP8 and FP4 precision to maximize throughput.
  • C++20 Based: Built with a modern CUDA codebase that avoids heavy reliance on complex templates.

Related

  • Project
    flashinfer-ai/flashinfer

    A high-performance GPU kernel library and generator for LLM inference that optimizes attention, GEMM, and MoE operations across diverse NVIDIA GPU architectures.

  • Project
    deepseek-ai/DeepEP

    A high-performance communication library for MoE models that provides optimized GPU kernels for expert parallelism to maximize throughput and minimize SM resource usage.

  • Project
    pytorch/FBGEMM

    A repository of highly-optimized kernels for deep learning, providing low-precision matrix multiplication and convolution libraries for server-side inference and generative AI.

  • Project
    NVIDIA/TensorRT-LLM

    An open-source library from NVIDIA that optimizes inference for LLMs and visual generation models using specialized kernels and an efficient runtime to maximize throughput and latency on NVIDIA GPUs.

  • Project
    Tencent/hpc-ops

    A production-grade operator library for LLM inference that provides highly optimized CUDA kernels for attention, MoE, and sampling to reduce serving latency.