DeepSeek 开源工作第3天
🚀 Day 3 of #OpenSourceWeek: DeepGEMM
Introducing DeepGEMM - an FP8 GEMM library that supports both dense and MoE GEMMs, powering V3/R1 training and inference.
⚡ Up to 1350+ FP8 TFLOPS on Hopper GPUs
✅ No heavy dependency, as clean as a tutorial
✅ Fully Just-In-Time compiled
✅ Core logic at ~300 lines - yet outperforms expert-tuned kernels across most matrix sizes
✅ Supports dense layout and two MoE layouts
🔗 GitHub: https://github.com/deepseek-ai/DeepGEMM
🚀 第3天 #OpenSourceWeek:DeepGEMM
OpenSourceWeek第3天:DeepGEMM
介绍DeepGEMM-一个支持密集和MoE GEMM的FP8 GEMM库,支持V3/R1训练和推理。
⚡ Hopper GPU上高达1350+FP8 TFLOPS
✅ 没有重度依赖,像教程一样干净
✅ 完全即时编译
✅ 约300行的核心逻辑-但在大多数矩阵大小上优于专家调整的内核
✅ 支持密集布局和两种MoE布局
🔗 在GitHub上探索:https://github.com/deepseek-ai/DeepGEMM
DeepGEMM 开源项目介绍
我来详细解释下 DeepGEMM 这个项目:
1. 项目概述
DeepGEMM 是一个高性能的 GEMM (通用矩阵乘法) CUDA 实现库,专门针对 NVIDIA Hopper 架构进行了优化。主要特点:
- 支持 FP8 (E4M3) 数据类型的矩阵乘法运算
- 利用 Hopper 架构的 TMA (Tensor Memory Accelerator) 特性
- 支持多种 GEMM 计算模式:普通模式、分组连续模式和分组掩码模式
2. 核心功能
- 基础 GEMM 实现
: 实现 FP8 输入和 BF16 输出的矩阵乘法
- 分组 GEMM 实现
- 连续分组模式: 适用于规则分组场景
- 掩码分组模式: 适用于不规则分组场景
- JIT 编译系统
- 动态生成 CUDA 代码
- 运行时编译优化
- 缓存编译结果提高性能
3. 项目结构
DeepGEMM/
├── deep_gemm/ # 主要源代码
│ ├── include/ # C++ 头文件
│ ├── jit/ # JIT 编译系统
│ ├── jit_kernels/ # CUDA 核心实现
│ └── utils.py # 工具函数
├── tests/ # 测试用例
└── third-party/ # 第三方依赖└── cutlass/ # NVIDIA CUTLASS 库
4. 关键技术特性
- 硬件要求
- NVIDIA Hopper 架构 GPU (sm_90a)
- CUDA 12.3 或更高版本(推荐 12.8+)
- Python 3.8+
- PyTorch 2.1+
- 优化技术
- 使用 TMA 进行高效的内存访问
- WGMMA (Warp Group Matrix Multiply-Accumulate) 指令优化
- 支持多种矩阵块大小配置
- 自动调优系统
5. 使用方法
- 安装
python setup.py install
- 开发环境设置
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
python setup.py develop
- 测试
python tests/test_jit.py # 测试 JIT 编译
python tests/test_core.py # 测试核心功能
6. 许可证
项目采用 MIT 许可证,允许自由使用、修改和分发。
这个项目的主要目标是提供一个高性能的 GEMM 实现,特别适合在 AI 训练和推理场景中使用 FP8 数据类型的情况。通过充分利用 Hopper 架构的新特性,实现了优秀的计算性能。
DeepSeek 开源组织及X账号
- https://github.com/deepseek-ai
- https://x.com/deepseek_ai