DeepSeek开源第三弹:惊人的300行代码驱动R1和V3的训练与推理,超越各种专家内核

人工智能 新闻
这是由 DeepSeek-AI 团队精心打造的 FP8 通用矩阵乘法 (GEMM) 加速库,专为追求极致性能和代码简洁而生。

DeepSeek开源周第三弹!DeepSeek-AI 重磅发布高效FP8 GEMM库 DeepGEMM:极致性能,代码精简,助力V3/R1模型训练与推理!

简单来说这是由 DeepSeek-AI 团队精心打造的 FP8 通用矩阵乘法 (GEMM) 加速库,专为追求极致性能和代码简洁而生

图片

我们一起来扒一扒:

Hopper GPU 上狂飙 1350+ FP8 TFLOPS!

在算力为王的 AI 时代,矩阵乘法 (GEMM) 的效率至关重要。DeepGEMM 正是为了解决这一痛点!它充分利用 NVIDIA Hopper 架构 GPU的强大算力,在 FP8 精度下,性能 高达 1350+ TFLOPS!这意味着更快的模型训练速度,更流畅的推理体验,以及更低的计算成本!

DeepGEMM 不仅适用于传统的 稠密模型,更完美支持 混合专家模型 (MoE) 的 GEMM 计算,无论是 DeepSeek-V3 还是 R1 模型,都能得到强力加速!

代码极简!核心逻辑仅 300 行,堪比教程级!

你没听错!DeepGEMM 的核心 kernel 函数代码量仅有 惊人的 ~300 行! DeepSeek-AI 团队秉持着 “大道至简” 的设计理念,在保证极致性能的同时,力求代码的 可读性 和 可维护性。 即使是刚入门 CUDA 开发的同学,也能轻松理解 DeepGEMM 的实现原理,甚至可以作为学习 Hopper FP8 矩阵乘法和优化的绝佳教程!

✅ 无需编译!完全 Just-In-Time (JIT) 编译,即装即用! ✅

告别繁琐的编译过程!DeepGEMM 采用了 全 Just-In-Time (JIT) 编译 技术,所有 kernel 都在运行时动态编译,无需在安装时进行任何预编译。 这意味着你可以 即装即用 DeepGEMM,省去了大量的配置和编译时间,让你可以更专注于模型开发和实验。

DeepGEMM 的 JIT 设计还带来了额外的优势:它可以根据不同的 GEMM 形状、block size 等参数进行 动态优化,始终选择最佳的 kernel 配置,保证在各种场景下都能发挥出最佳性能。

💪 硬核技术解析:DeepGEMM 的性能秘诀

DeepGEMM 在代码简洁的同时,性能却能比肩甚至超越一些专家调优的库,这背后离不开一系列硬核技术的加持:

  • • 精细粒度 Scaling (Fine-grained Scaling): DeepGEMM 采用了 DeepSeek-V3 论文中提出的精细粒度 scaling 技术,更有效地利用 FP8 的动态范围,提升计算精度和性能
  • • CUDA-core 双层累加 (Two-level Accumulation): 为了解决 FP8 tensor core 累加精度不足的问题,DeepGEMM 巧妙地使用了 CUDA-core 双层累加技术,保证了计算结果的准确性
  • • Persistent Warp-specialization (持久 Warp 特化): 借鉴 CUTLASS 的设计思想,DeepGEMM 的 kernel 进行了 warp 特化,实现了数据移动、tensor-core MMA 指令和 CUDA-core promotion 的高效重叠,最大化利用硬件资源
  • • Tensor Memory Accelerator (TMA): DeepGEMM 充分利用 Hopper 架构引入的 TMA 特性,加速 LHS、RHS 矩阵和 scaling factor 的加载,以及输出矩阵的存储,实现更快的数据访问速度
  • • 统一优化 Block Scheduler 和 Rasterization (栅格化): DeepGEMM 采用统一的 block scheduler,并结合 Rasterization 技术,提升 L2 cache 的复用率,进一步优化性能
  • • FFMA SASS Interleaving: DeepGEMM 甚至深入到 SASS 汇编层面进行优化,通过调整 FFMA 指令的 interleaving 模式,提升 warp 级别的并行度,榨干硬件的每一丝潜力

DeepGEMM 虽然借鉴了 CUTLASS 和 CuTe 的一些概念,但它并没有过度依赖于复杂的模板或代数库,而是更加注重 简洁性 和 易用性。 这使得 DeepGEMM 不仅是一个高性能的计算库,更是一个学习 Hopper FP8 矩阵乘法和优化的优秀资源

📊 实测性能数据:实力说话! 📊

DeepGEMM 的性能究竟如何?我们用数据说话!在 DeepSeek-V3/R1 模型常用的各种 shape 上,DeepGEMM 都展现出了惊人的性能:

  • • Normal GEMMs for dense models (稠密模型 GEMM)

M

N

K

Computation

Memory bandwidth

Speedup

64

2112

7168

206 TFLOPS

1688 GB/s

2.7x

64

24576

1536

289 TFLOPS

2455 GB/s

1.7x

...

...

...

...

...

...

4096

7168

2048

1025 TFLOPS

697 GB/s

1.1x

  • • Grouped GEMMs for MoE models (contiguous layout) (MoE 模型 GEMM - 连续布局)

#Groups

M per group

N

K

Computation

Memory bandwidth

Speedup

4

8192

4096

7168

1297 TFLOPS

418 GB/s

1.2x

4

8192

7168

2048

1099 TFLOPS

681 GB/s

1.2x

...

...

...

...

...

...

...

8

4096

7168

2048

1093 TFLOPS

743 GB/s

1.1x

  • • Grouped GEMMs for MoE models (masked layout) (MoE 模型 GEMM - Masked 布局)

#Groups

M per group

N

K

Computation

Memory bandwidth

Speedup

1

1024

4096

7168

1233 TFLOPS

924 GB/s

1.2x

1

1024

7168

2048

925 TFLOPS

968 GB/s

1.2x

...

...

...

...

...

...

...

4

256

7168

2048

815 TFLOPS

2047 GB/s

1.2x

从数据中可以看出,DeepGEMM 在各种矩阵 shape 下都表现出色,速度提升明显! 💪

快速上手 DeepGEMM:只需几步!

想要体验 DeepGEMM 的强大性能? 上手非常简单!

环境要求:

  • • NVIDIA Hopper 架构 GPU (sm_90a)
  • • Python 3.8+
  • • CUDA 12.3+ (推荐 12.8+ 获得最佳性能)
  • • PyTorch 2.1+
  • • CUTLASS 3.6+ (可以通过 Git submodule 克隆)

安装步骤:

  1. 1. 克隆 DeepGEMM 代码库 (需要递归克隆 submodule):
git clone --recursive https://github.com/deepseek-ai/DeepGEMM.git
  1. 2. 创建 third-party 库的符号链接 (CUTLASS 和 CuTe):
python setup.py develop
  1. 3. 测试 JIT 编译:
python tests/test_jit.py
  1. 4. 测试所有 GEMM 实现 (normal, contiguous-grouped, masked-grouped):
python tests/test_core.py
  1. 5. 安装 DeepGEMM:
python setup.py install

安装完成后,只需在你的 Python 项目中 import deep_gemm 即可开始使用!

写在最后:

🔗 项目地址:

https://github.com/deepseek-ai/DeepGEMM

DeepGEMM 现已以MIT许可方式正式开源!

DeepGEMM 的灵感来自 CUTLASS 项目

责任编辑:张燕妮 来源: AI寒武纪
相关推荐

2025-02-26 11:16:18

2025-03-03 08:17:00

DeepSeek模型数据

2025-02-26 11:13:51

2025-02-26 10:24:51

2024-12-27 09:50:00

模型数据测试

2025-03-03 09:00:00

DeepSeekAI人工智能

2024-12-30 20:32:36

2025-01-27 12:30:07

2025-02-17 09:33:00

AI算法模型

2025-01-21 11:53:53

2025-02-20 15:32:28

2025-02-03 14:17:27

2025-02-10 06:50:00

AIDeepSeek推理模型

2025-02-27 00:00:05

2025-02-07 13:10:06

2025-02-03 00:00:55

DeepSeekRAG系统

2025-02-08 14:03:25

2025-02-20 08:45:41

V3GPU资源
点赞
收藏

51CTO技术栈公众号