DeepGEMM FP8GEMM 库
DeepGEMM 是一个专为简洁高效的 FP8 通用矩阵乘法（GEMM）设计的库，具有细粒度缩放功能，如 DeepSeek-V3 中所述。它支持普通和混合专家（MoE）分组的 GEMM。该库采用 CUDA 编写，在安装过程中无需编译，通过使用轻量级的即时编译（JIT）模块在运行时编译所有内核。
目前，DeepGEMM 仅支持 NVIDIA Hopper 张量核心。为了解决不精确的 FP8 张量核心累积问题，它采用了 CUDA-core two-level accumulation (promotion)。虽然它利用了CUTLASS和CuTe的一些概念，但避免了对其模板或代数的过度依赖。相反，该库的设计非常简单，只有一个核心内核函数，包含大约300 行代码。这使其成为学习 Hopper FP8 矩阵乘法和优化技术的干净且易于访问的资源。
要求
- Hopper 架构 GPU，
sm_90a必须支持 - Python 3.8 或更高版本
- CUDA 12.3 或更高版本
- 但强烈建议使用 12.8 或更高版本以获得最佳性能
- PyTorch 2.1 或更高版本
- CUTLASS 3.6 或更高版本（可以通过 Git 子模块克隆）
Development
# Submodule must be cloned
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
# Make symbolic links for third-party (CUTLASS and CuTe) include directories
python setup.py develop
# Test JIT compilation
python tests/test_jit.py
# Test all GEMM implements (normal, contiguous-grouped and masked-grouped)
python tests/test_core.py
Installation
python setup.py install
评论