DeepGEMM：一款支持FP8通用矩阵乘法(GEMM)的开源库

DeepGEMM是什么

DeepGEMM是DeepSeek公司在其“开源周”第3天发布的一款开源库，专为高效且清晰的FP8通用矩阵乘法（GEMMs）设计。该库不仅支持标准的GEMM运算，还针对混合专家（MoE）模型进行了特别优化，为DeepSeek的V3和R1模型的训练和推理提供了强大的支持。DeepGEMM的发布标志着DeepSeek在推动AI技术透明化和社区协作方面迈出了重要一步。

功能特色

高性能

DeepGEMM在NVIDIA Hopper GPU上能够实现高达1350+ FP8 TFLOPS的性能。这一性能表现得益于其针对Hopper架构和FP8精度的深度优化，使其在处理大规模矩阵乘法运算时能够显著提升计算效率。

简洁高效的设计

DeepGEMM的核心代码仅约300行，却能在大多数矩阵尺寸上超越专家调优的内核。这一简洁高效的设计不仅降低了学习和使用的门槛，也使得库本身的维护和更新变得更加容易。同时，DeepGEMM采用了即时编译（Just-In-Time, JIT）技术，无需在安装时进行编译，而是在运行时动态编译所有内核，进一步提高了灵活性和效率。

支持多种数据排列方式

DeepGEMM支持标准矩阵排列和两种专为混合专家模型设计的特殊排列方式：连续排列和掩码排列。这使得DeepGEMM能够适应不同的应用场景和模型架构，为开发者提供了更大的灵活性。

细粒度缩放功能

DeepGEMM具备细粒度缩放功能，这一设计源于DeepSeek V3。通过细粒度缩放，DeepGEMM能够更精细地控制计算过程中的数值范围，从而提高计算的稳定性和准确性。

技术细节

CUDA编写与JIT编译

DeepGEMM采用CUDA编写，充分利用了NVIDIA GPU的并行处理能力。同时，DeepGEMM采用了JIT编译技术，所有计算内核都在实际运行时进行编译。这种即时编译方法可以根据当前的矩阵形状和硬件特性动态生成优化的代码，从而进一步提高性能。

两级累加技术

为了应对FP8张量核心累加不精确的问题，DeepGEMM采用了CUDA核心的两级累加（提升）方法。这一技术通过两次累加操作来减少累积误差，提高了计算的精度和稳定性。

非标准块大小支持

传统上，GPU计算通常使用标准大小的数据块（如128×128）。然而，DeepGEMM支持非标准块大小（如112×128），这能更好地适应特定矩阵形状，提高硬件资源利用率。例如，对于M=256,N=7168的矩阵，标准块大小只能利用112个计算单元，而使用非标准块大小可以利用128个，效率提升明显。

线程专业化技术

DeepGEMM采用了线程专业化技术，这是一种高效的任务分工方法。在这种设计中，不同的计算线程被分配专门负责特定任务：一些负责数据移动，一些负责核心计算，一些负责结果处理。这种分工使得数据移动、计算和后处理能够同时进行，形成高效的流水线，大大提高整体性能。

统一的计算任务调度系统

DeepGEMM设计了一套统一的计算任务调度系统，采用特殊的排布策略，增强缓存重用效率，减少内存访问，提高整体性能。

应用场景

深度学习模型训练与推理

DeepGEMM主要用于加速深度学习中的矩阵运算，特别是在大规模模型训练和推理中。它能够显著提升计算效率，减少训练时间和推理延迟，从而加快模型的开发和部署进程。

混合专家模型（MoE）

DeepGEMM针对混合专家模型进行了特别优化，提供了连续排列和掩码排列两种特殊数据排列方式。这使得DeepGEMM能够更高效地处理MoE模型的训练和推理任务，为构建更复杂、更高效的AI模型提供了有力支持。

高性能计算任务

除了深度学习领域外，DeepGEMM还可以用于其他需要高效矩阵乘法运算的高性能计算任务中。例如，在科学计算、金融分析等领域中，DeepGEMM都能够提供强大的计算支持。

相关官方链接

• GitHub链接：DeepGEMM

总结

DeepGEMM是一款专为高效且清晰的FP8通用矩阵乘法（GEMMs）设计的开源库。它支持标准矩阵计算和混合专家模型（MoE）计算，为DeepSeek的V3和R1模型的训练和推理提供了强大的支持。DeepGEMM采用CUDA编写，采用JIT编译技术，支持多种数据排列方式和细粒度缩放功能。其高性能、简洁高效的设计、灵活性和广泛的应用场景使其成为AI领域的一款重要工具。无论是深度学习模型训练与推理还是其他高性能计算任务中，DeepGEMM都能够提供强大的计算支持。通过GitHub和官方文档等资源，你可以方便地获取DeepGEMM的源代码、文档和示例代码，并与其他开发者一起交流和分享使用经验。