Deepseek开源周第三天:DeepGEMM发布

前言

上周deepseek宣布,将在本周陆续发布五个开源项目,这些库已经在生产环境中经过了记录、部署和实战测试。

今天是deepseek开源周的第三天,deepseek发布了一个名为 DeepGEMM 的项目,截止发稿前,已有 3.2K 的 star

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

在人工智能和高性能计算领域,矩阵乘法(General Matrix Multiplication,简称 GEMM)是许多算法的核心运算。

DeepGEMM 的项目,旨在为 FP8(8 位浮点数)矩阵乘法提供高效、简洁的实现。

什么是 FP8 和矩阵乘法?

矩阵乘法是线性代数中的基本运算,用于计算两个矩阵的乘积。

在深度学习中,矩阵乘法被广泛应用于神经网络的前向传播和反向传播过程中。

然而,传统的 32 位浮点数(FP32)或 16 位浮点数(FP16)矩阵乘法在计算效率和内存占用方面存在瓶颈。

为了进一步提升计算效率和降低能耗,研究人员开始探索更低精度的浮点数格式。

FP8 是一种 8 位浮点数格式,它通过减少精度来换取更快的计算速度和更低的内存带宽需求。

尽管精度降低,但在许多深度学习任务中,FP8 仍然能够提供足够的数值精度,同时显著提升性能。

DeepGEMM 是什么?

DeepGEMM 是一个专注于 FP8 矩阵乘法的开源库,由 DeepSeek 团队开发。

它旨在为深度学习中的矩阵乘法提供高效、简洁的实现,同时支持普通矩阵乘法和混合专家模型(MoE)中的分组矩阵乘法。

DeepGEMM 的核心优势在于其简洁的设计和高效的实现。

它仅包含一个约 300 行代码的核心内核函数,这使得它易于理解和学习。

此外,DeepGEMM 采用了运行时编译(Just-In-Time,JIT)技术,无需在安装时进行编译,所有内核都在运行时动态编译,从而更好地适应不同的硬件和计算需求。

DeepGEMM 的技术亮点

  1. 高效利用 NVIDIA Hopper 架构

DeepGEMM 专为 NVIDIA Hopper 架构设计,充分利用了 Hopper 引入的 Tensor Memory Accelerator(TMA)特性。

TMA 是一种用于加速数据传输的硬件特性,能够显著提升矩阵乘法的性能。

  1. 细粒度缩放技术

在 FP8 计算中,由于精度较低,可能会导致数值溢出或下溢。

DeepGEMM 通过细粒度缩放技术解决了这一问题,确保计算的稳定性。

  1. 分组矩阵乘法支持

DeepGEMM 支持分组矩阵乘法(Grouped GEMM),这对于混合专家模型(MoE)尤为重要。

MoE 是一种大规模并行模型架构,通过将输入数据分配给不同的“专家”进行处理,从而提高模型的效率和可扩展性。

DeepGEMM 提供了针对 MoE 的优化实现,支持连续布局(contiguous layout)和掩码布局(masked layout)两种场景。

  1. 运行时编译(JIT)设计

DeepGEMM 采用 JIT 编译技术,所有内核在运行时动态编译。

这种设计不仅减少了安装时的复杂性,还允许在运行时根据矩阵形状和硬件特性进行优化,进一步提升性能。

  1. 性能优化

DeepGEMM 在多个矩阵形状上达到了与专家优化库相当甚至更高的性能。

例如,在某些测试中,DeepGEMM 的速度比基于 CUTLASS 的实现快 1.4 到 2.7 倍。

DeepGEMM 的应用场景

DeepGEMM 的高效实现使其在多个领域具有广泛的应用潜力:

  1. 深度学习推理

在深度学习模型的推理阶段,DeepGEMM 可以显著提升计算效率,降低延迟,尤其适用于大规模模型和实时应用。

  1. 混合专家模型(MoE)

MoE 是一种高效的并行模型架构,DeepGEMM 提供了针对 MoE 的优化支持,能够显著提升其性能。

  1. 高性能计算

在高性能计算领域,DeepGEMM 的低精度矩阵乘法可以用于加速科学计算和大规模数据分析。

总结

DeepGEMM 的开源为深度学习和高性能计算领域带来了新的可能性。

它不仅提供了高效、简洁的 FP8 矩阵乘法实现,还通过运行时编译和细粒度缩放技术,解决了低精度计算中的关键问题。

DeepGEMM 的开源也为研究人员和开发者提供了一个学习和优化 Hopper 架构的绝佳资源。

我们将持续关注接下来几天的项目发布,为读者带来第一手的技术资讯。


我码玄黄
1 声望1 粉丝