DeepSeek 开源记-第3天-DeepGEMM

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 开源项目介绍

1. 项目概述

DeepGEMM 是一个高性能的 GEMM (通用矩阵乘法) CUDA 实现库，专门针对 NVIDIA Hopper 架构进行了优化。主要特点:

• 支持 FP8 (E4M3) 数据类型的矩阵乘法运算
• 利用 Hopper 架构的 TMA (Tensor Memory Accelerator) 特性
• 支持多种 GEMM 计算模式:普通模式、分组连续模式和分组掩码模式

2. 核心功能

1. 基础 GEMM 实现

• 实现 FP8 输入和 BF16 输出的矩阵乘法

2. 分组 GEMM 实现

• 连续分组模式: 适用于规则分组场景
• 掩码分组模式: 适用于不规则分组场景

3. JIT 编译系统

• 动态生成 CUDA 代码
• 运行时编译优化
• 缓存编译结果提高性能

3. 项目结构

DeepGEMM/
├── deep_gemm/              # 主要源代码
│   ├── include/           # C++ 头文件
│   ├── jit/              # JIT 编译系统
│   ├── jit_kernels/      # CUDA 核心实现
│   └── utils.py          # 工具函数
├── tests/                 # 测试用例
└── third-party/          # 第三方依赖
    └── cutlass/          # NVIDIA CUTLASS 库

4. 关键技术特性

1. 硬件要求

• NVIDIA Hopper 架构 GPU (sm_90a)
• CUDA 12.3 或更高版本(推荐 12.8+)
• Python 3.8+
• PyTorch 2.1+

2. 优化技术

• 使用 TMA 进行高效的内存访问
• WGMMA (Warp Group Matrix Multiply-Accumulate) 指令优化
• 支持多种矩阵块大小配置
• 自动调优系统

5. 使用方法

1. 安装

python setup.py install

2. 开发环境设置

git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
python setup.py develop

3. 测试

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