DeepGEMM
github.com/deepseek-ai/deepgemm- Category
- Developer Tools
- Rank
- No. 300Tools index
- Pricing
- Open Source
- Type
- TOOL
- Builder
- deepseek-ai
- GitHub
- 7.8k stars
- Latest release
- nv_dev_f8e8fb5
- Date
About
A high-performance CUDA library for FP8/FP4 tensor operations in large language models, featuring optimized GEMM kernels, MoE fusion, and runtime JIT compilation. Designed for NVIDIA GPUs with clean, accessible code for learning GPU optimization techniques.
What it does
DeepGEMM supplies PyTorch-accessible GPU primitives for the costly matrix and routing work inside large language models. It covers dense and grouped multiplication, low-precision einsum, attention scoring, HyperConnection, and a fused mixture-of-experts pipeline. A small C++ extension exposes the interface, while specialized kernels are compiled when needed and cached.
Why it's ranked here
The appeal is unusually broad coverage without hiding the hardware mechanics. Dense, grouped, masked, attention, and end-to-end mixture-of-experts paths share one package. Runtime compilation avoids building every CUDA kernel during installation. The value is concentrated, however, on recent NVIDIA architectures and workloads prepared around strict layouts, scaling formats, and alignment rules.
What's good
The grouped kernels directly model expert workloads, including variable token counts during training or prefill and masked execution during graph-based decoding. Mega MoE combines expert dispatch, two low-precision linear stages, SwiGLU, and result combining while overlapping NVLink traffic with tensor-core computation. Utilities cover quantization, scale packing, layout preparation, distributed setup, benchmarking, and numeric comparison.
Tradeoffs
Current requirements are demanding: recent NVIDIA GPU architectures, modern CUDA, PyTorch, CUTLASS, a C++ compiler, and Linux packaging. Older A100 support sits in a legacy area that may be deprecated. Users must separately handle input transposition and low-precision casting, and convenient PyTorch helpers may be slower. Scaling representation differs by GPU generation, while grouped operations constrain unchanged dimensions and require block alignment.
How to use it well
Use it when optimizing training or inference kernels for recent NVIDIA systems, especially mixture-of-experts models, low-precision matrix multiplication, or specialized attention scoring. Start from the supplied tests and utilities, validate numerical behavior, then benchmark exact production shapes because compilation and configuration are shape-sensitive. It does not replace upstream preprocessing, general model orchestration, portable accelerator support, or communication infrastructure outside its fused expert path.
Technical notes+
setup.py packages deep_gemm, builds the deep_gemm._C CUDA extension from csrc/python_api.cpp, bundles CUTLASS and CuTe headers, generates _C.pyi through scripts/generate_pyi.py, and can fetch a matching wheel before falling back to a source build. csrc/python_api.cpp registers attention, einsum, hyperconnection, GEMM, layout, Mega MoE, and runtime APIs through pybind11. deep_gemm/__init__.py initializes _C, exports kernel bindings, loads legacy A100 Triton kernels opportunistically, and resolves the CUDA installation. deep_gemm/mega/__init__.py manages symmetric-memory buffers and fused FP8/FP4 or BF16 MoE execution. deep_gemm/utils/math.py implements FP8 and packed FP4 conversion plus UE8M0 scaling helpers. deep_gemm/testing/bench.py provides CUDA event and Kineto benchmarking, while scripts/quick_plot_pm.py extracts and plots selected Nsight Compute metrics.
Observed
- Primary languages
- CUDA/C++ kernel extension with a Python package and utilities
- Interface
- Importable Python library backed by a pybind11 C++ extension
- Installation surface
- Shell installation script and setuptools wheel or source-build packaging
- Supported GPU platforms
- NVIDIA SM90 and SM100; legacy Triton kernels are included for A100
- Runtime requirements
- Linux, Python 3.8 or higher, PyTorch 2.1 or higher, CUDA Toolkit, CUTLASS, fmt, and a C++20-capable compiler
- Compilation model
- GPU kernels compile at runtime through a lightweight JIT module and use a configurable cache
- Testing structure
- The documented examples reference tests for core kernels, attention, and multi-process Mega MoE
Read from README.md, setup.py, csrc/python_api.cpp, deep_gemm/__init__.py, scripts/generate_pyi.py, scripts/quick_plot_pm.py, deep_gemm/utils/dist.py, deep_gemm/utils/math.py, deep_gemm/utils/layout.py, deep_gemm/mega/__init__.py, deep_gemm/testing/bench.py, deep_gemm/testing/utils.py, deep_gemm/utils/__init__.py, deep_gemm/legacy/__init__.py, deep_gemm/testing/numeric.py.
What it can do
Perform high-performance FP8/FP4 tensor matrix multiplication
FP8 or FP4 tensor matrices → Computed matrix multiplication results
Execute optimized GEMM operations on NVIDIA GPUs
General matrix multiplication parameters and data → Accelerated matrix computation results
Perform fused Mixture of Experts (MoE) computations
MoE model parameters and input tensors → Fused MoE computation results
Execute Multi-Query Attention (MQA) scoring operations
Attention query, key, and value tensors → Attention scores and weighted outputs
Compile CUDA kernels at runtime using JIT compilation
Kernel specifications and runtime parameters → Optimized compiled CUDA kernels
Optimize tensor operations for large language model inference
LLM computational workloads and tensors → Accelerated LLM computation results
Tags
Comments (0)
No comments yet
Editorially curated, with community endorsements as a secondary signal. Corrections welcome.