在大型模型推理中,矩阵乘法(GEMM)是计算的核心瓶颈。DeepGEMM 应运而生——一款专为 FP8精度矩阵乘法 设计的轻量级CUDA库,由深度求索(DeepSeek)团队开源。它凭借极简代码(核心仅300行)、 无编译安装 的特性,以及针对Hopper架构的极致优化,性能超越主流库,为大模型推理提速高达 2.7倍!
核心亮点
-
性能王者
- FP8精度加速
专为NVIDIA Hopper张量核心设计,支持普通GEMM和MoE分组GEMM,计算效率突破 1358 TFLOPS,带宽利用率高达 2668 GB/s。
- 速度碾压
相比CUTLASS 3.6优化版本,在典型推理场景(如M=64, N=2112, K=7168)下速度提升 2.7倍,MoE场景下提升 1.2倍。
- FP8精度加速
-
极简设计
- 轻量核心
仅需1个核心函数,代码量约300行,学习FP8矩阵乘法优化的绝佳范本。
- 零编译依赖
基于JIT(即时编译)实现,安装即用,无需复杂环境配置。
- 轻量核心
-
Hopper架构深度优化
- TMA硬件加速
利用Hopper的Tensor Memory Accelerator实现异步数据搬运,减少等待开销。
- 创新指令调度
通过FFMA指令交织、非对齐分块等黑科技,最大化SM利用率,解决小矩阵计算效率瓶颈。
- TMA硬件加速
适用场景
- 大模型推理
支持预填充(Prefilling)和解码(Decoding)阶段,适配Transformer和MoE结构。
- MoE模型优化
提供 连续布局(训练/预填充)和 掩码布局(解码)两种分组GEMM,灵活应对动态Token分配。
技术优势
- 持久化线程束 specialization
异步重叠数据搬运、张量核心计算与CUDA核心提升,实现计算效率最大化。
- JIT极致优化
矩阵形状、分块大小等参数编译期固定,寄存器占用更少,性能逼近手写汇编。
- 自适应调度器
智能选择分块策略、流水线阶段和TMA集群规模,无需手动调参。
快速上手
# 克隆仓库(需递归拉取子模块)
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
# 安装依赖
python setup.py develop
# 验证安装
python tests/test_jit.py
环境要求:
-
GPU:NVIDIA Hopper架构(如H100)
-
CUDA 12.3+(推荐12.8)
-
Python 3.8+、PyTorch 2.1+
使用示例
import deep_gemm
# 普通FP8 GEMM
deep_gemm.gemm_fp8_fp8_bf16_nt(lhs, rhs, out, scale_lhs, scale_rhs)
# MoE连续布局分组GEMM
deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(...)
🔗 项目地址:
https://github.com/deepseek-ai/DeepGEMM