DeepSeek 开源活动来到了第三天,新项目如约而至。
此次开源库名为 DeepGEMM,GitHub 地址在文末。
DeepSeek 表示,这是一个支持密集和混合专家(MoE,Mixture of Experts)通用矩阵乘法(GEMM,General Matrix Multiplication)的 FP8 GEMM 库,为 V3/R1 的训练和推理提供支持。
在 Hopper GPU 上最高可达 1350+FP8 TFLOPS。其他优点包括:
✅没有过多的依赖,像教程一样简洁
✅完全即时编译
✅核心逻辑约为 300 行,但在大多数矩阵大小上均优于专家调优的内核
✅支持密集布局和两种 MoE 布局
(来源:DeepSeek)
据 DeepSeek 介绍,DeepGEMM 是一个专门设计的、干净且高效的工具库,它的核心任务是进行一种叫做 GEMM 的数学运算,这是 AI 模型训练和运行中非常常见的一种计算。
DeepGEMM 的特别之处在于,它使用了一种叫做 FP8 的超高效、低精度计算方式,能让运算速度更快,同时占用更少的内存。这种方式在 DeepSeek-V3 中被提出,并且支持精细的缩放调整(fine-grained scaling),让计算更灵活。
除了普通的矩阵乘法,DeepGEMM 还能处理混合专家矩阵乘法。
(来源:DeepSeek)
目前,DeepGEMM 只支持英伟达 Hopper 架构的张量核心。张量核心是 GPU 里的一种特殊硬件,专门用来加速矩阵运算。不过,Hopper 的张量核心在做 FP8 计算时,会出现累加(accumulation,就是把结果一点点加起来的过程)不够精确的问题。
为了解决这个问题,DeepGEMM 采用了两级累加(two-level accumulation)的办法。它借助 CUDA 核心来做更精确的累加,确保结果不会因为硬件限制而出错。
DeepGEMM 借鉴了一些来自 CUTLASS 和 CuTe 的概念。这两个也是 NVIDIA 的高性能矩阵运算库。不过,DeepGEMM 没有完全依赖它们的复杂模板或数学系统。
相反,该库的设计非常简单,只有一个核心内核函数,包含大约 300 行代码。这使其成为学习 Hopper FP8 矩阵乘法和优化技术的、干净且易于访问的资源。
尽管 DeepGEMM 设计轻量,但它的性能足以媲美那些由专家精心调优的库,甚至在某些矩阵形状(matrix shapes,指矩阵的大小和结构)上表现得更好。
在性能方面,DeepSeek 在搭载 NVCC 12.8 的 H800 上测试了 DeepSeek-V3/R1 推理中可能使用的所有形状(包括预填充和解码,但没有张量并行性)。
从性能对比表格上可见,DeepGEMM 最高能加速 2.7 倍。
图 | 密集模型加速数据(来源:DeepSeek)
图 | MoE 模型加速数据(来源:DeepSeek)
所有加速指标都是与 DeepSeek 基于 CUTLASS 3.6 的内部精心优化的实现进行比较计算的。
鉴于 DeepGEMM 在某些矩阵形状上的表现并不是很好,DeepSeek 也邀请各路大神帮助优化这个开源库。
在优化方面,DeepSeek 实现了一些 CUTLASS 设计之外的优化。这些也是 DeepGEMM 最具创新的地方。
(来源:DeepSeek)
首先是完全 JIT 设计,安装时无需编译。所有内核均使用轻量级 JIT 实现在运行时进行编译。这种方法的优点包括:
- GEMM 形状、块大小和管道阶段数被视为编译时常量
- 自动选择块大小、warpgroups 数量、最佳管道阶段和 TMA 集群大小
- 全面展开 MMA 管道,为编译器提供更多优化机会
(来源:DeepSeek)
其次是支持非对齐的块大小。在做矩阵乘法时,矩阵会被分成一个个小块(block)来计算。块的大小通常是固定的,比如 128×128(因为 128 是 2 的幂,计算机喜欢这种数字)。但有时候,矩阵的尺寸跟这些块大小不完全匹配,就会浪费计算资源。
举个例子:
假设矩阵的行数(M)是 256,列数(N)是 7168。如果用常见的块大小 BLOCK_M=128(行)和 BLOCK_N=128(列),那么行方向可以分成 256÷128=2 个块,列方向可以分成 7168÷128=56 个块。总共用到的 SM(GPU 的计算单元)是 2×56=112 个。
GPU 通常有更多 SM(比如 128 个),但这里只用了 112 个,剩下的就闲着了,没充分利用。
DeepGEMM 的办法是支持非对齐的块大小,也就是不一定非要是 2 的幂(比如 128),可以根据矩阵尺寸灵活调整。
还是上面的例子,如果把 BLOCK_N 改成 112(而不是 128,不是 2 的幂),行方向还是 256÷128=2 个块,列方向变成了 7168÷112=64 个块。总共用到的 SM 变成 2×64=128 个。
这下正好把 128 个 SM 全用上,没有浪费。
(来源:DeepSeek)
最后一点是,DeepSeek 通过观察和修改底层代码(SASS 汇编指令)实现了性能优化。
简单来说,DeepSeek 发现新版编译器(NVCC 12.3)里有些底层代码变了,性能变得更好。深入研究后,他们觉得这是因为线程调度(yield)的方式变了,能让 GPU 同时干更多工作。
于是,他们模仿这个变化,写了个脚本修改自己的代码(调整 FFMA 指令的某些位),让 MMA 指令(矩阵乘累加)和提升指令更高效地“重叠在一起执行”。
结果是 DeepGEMM 在某些场景下快了 10% 以上,特别适合那种需要灵活调整的 FP8 矩阵乘法。
最后,DeepSeek 致谢了 CUTLASS 项目,称其是 DeepGEMM 的灵感来源。
参考资料:
https://x.com/deepseek_ai/status/1894553164235640933
https://github.com/deepseek-ai/DeepGEMM
排版:刘雅坤