在大型模型推理中，矩阵乘法（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
