DeepSeek开源周，第三弹再次来袭，DeepGEMM

---

    在大型模型推理中，矩阵乘法（GEMM）是计算的核心瓶颈。DeepGEMM 应运而生——一款专为 FP8精度矩阵乘法 设计的轻量级CUDA库，由深度求索（DeepSeek）团队开源。它凭借极简代码（核心仅300行）、 无编译安装 的特性，以及针对Hopper架构的极致优化，性能超越主流库，为大模型推理提速高达 2.7倍！

---

---

核心亮点

1. 性能王者

   • 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倍。

2. 极简设计

   • 轻量核心

     仅需1个核心函数，代码量约300行，学习FP8矩阵乘法优化的绝佳范本。

   • 零编译依赖

     基于JIT（即时编译）实现，安装即用，无需复杂环境配置。

3. Hopper架构深度优化

   • TMA硬件加速

     利用Hopper的Tensor Memory Accelerator实现异步数据搬运，减少等待开销。

   • 创新指令调度

     通过FFMA指令交织、非对齐分块等黑科技，最大化SM利用率，解决小矩阵计算效率瓶颈。

---

---

适用场景

• 大模型推理

  支持预填充（Prefilling）和解码（Decoding）阶段，适配Transformer和MoE结构。

• MoE模型优化

  提供 连续布局（训练/预填充）和 掩码布局（解码）两种分组GEMM，灵活应对动态Token分配。

---

---

技术优势

1. 持久化线程束 specialization

   异步重叠数据搬运、张量核心计算与CUDA核心提升，实现计算效率最大化。

2. JIT极致优化

   矩阵形状、分块大小等参数编译期固定，寄存器占用更少，性能逼近手写汇编。

3. 自适应调度器

   智能选择分块策略、流水线阶段和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