DeepGEMM：全球首个开源的 FP8 矩阵加速器技术解析

随着人工智能和深度学习的快速发展，计算效率和资源优化成为研究和工业应用中的关键问题。矩阵乘法（General Matrix Multiplication,GEMM）作为深度学习中最为基础且频繁使用的计算操作，其性能直接影响到模型训练和推理的速度。近年来，随着硬件技术的进步，尤其是 NVIDIA Hopper 架构的推出，低精度计算格式（如 FP8）逐渐成为研究热点。DeepSeek 开源的 DeepGEMM 矩阵加速器正是在这一背景下应运而生，它不仅填补了全球范围内 FP8 矩阵加速器的空白，还通过一系列创新技术实现了显著的性能提升。本文将详细解析 DeepGEMM 的技术原理、设计架构以及其背后的技术挑战和优化策略。

一、背景与需求

在深度学习中，矩阵乘法是卷积神经网络（CNN）、循环神经网络（RNN）以及 Transformer 架构等核心模型的基本运算单元。随着模型规模的不断扩大，计算资源的需求呈指数级增长。传统计算中，FP32（32 位浮点）格式被广泛使用，但其计算效率较低，难以满足大规模模型的实时训练和推理需求。近年来，随着硬件技术的发展，低精度计算格式（如 FP16 和 FP8）逐渐受到关注。这些低精度格式通过减少表示位数来提高计算速度和内存带宽利用率，同时显著降低功耗。

FP8 是一种新兴的 8 位浮点格式，相比 FP32 和 FP16，它在计算速度上有显著提升，但精度较低。为了在保持计算效率的同时减少精度损失，DeepGEMM 应运而生。它通过一系列优化策略，充分利用 NVIDIA Hopper 架构的硬件特性，实现了高效的 FP8 矩阵乘法加速。

二、DeepGEMM 的技术原理

（一）FP8 精度优化

FP8 是一种 8 位浮点格式，其位宽较短，因此计算速度更快，但精度较低。为了在低精度计算中减少误差，DeepGEMM 采用了细粒度的缩放策略（fine-grained scaling）。该策略通过动态调整计算过程中的缩放因子，确保在每一步计算中都能最大限度地保留有效信息。

此外，DeepGEMM 还引入了两级累加策略（two-level accumulation）。在传统的矩阵乘法中，累加操作可能会导致精度损失，尤其是在低精度计算中。两级累加策略通过在 CUDA 核心中实现中间精度提升，减少了累积误差，同时保持了计算速度。

（二）Just-In-Time(JIT)编译

DeepGEMM 的另一个核心技术是 Just-In-Time(JIT)编译。与传统的预编译方法不同，JIT 编译在运行时动态生成和优化内核代码。这种方法避免了安装时的复杂编译过程，能够根据不同的硬件环境和矩阵形状实时调整优化参数。

JIT 编译将矩阵乘法的形状、块大小和流水线阶段视为编译时常量，自动选择最优参数并全面展开 MMA（矩阵乘法累加）流水线。这种动态优化策略使得 DeepGEMM 能够在不同的硬件配置下表现出色，同时减少了开发和部署的复杂性。

（三）支持非对齐块大小

在矩阵乘法中，矩阵通常被划分为小块进行计算。传统方法中，块大小通常是固定的（如 128×128），但 DeepGEMM 支持非对齐块大小（如 112×112）。这种灵活性使得 DeepGEMM 能够更好地利用硬件资源，尤其是在处理非标准矩阵形状时。

非对齐块大小的支持不仅提高了计算效率，还减少了内存访问的开销。通过优化内存访问模式，DeepGEMM 能够进一步减少内存带宽瓶颈，提升整体性能。

（四）Hopper 架构优化

DeepGEMM 专为 NVIDIA Hopper 架构设计，充分利用了其 Tensor Memory Accelerator(TMA)特性。TMA 是 Hopper 架构中的一个关键特性，支持异步数据移动、多播和描述符预取。

通过 TMA，DeepGEMM 能够优化数据传输和存储，减少内存带宽瓶颈。此外，DeepGEMM 还针对 Hopper 架构的 Tensor Cores 进行了深度优化，确保在低精度计算中能够充分利用硬件资源。

（五）支持多种矩阵操作

DeepGEMM 不仅支持标准的密集矩阵乘法（dense GEMM），还支持混合专家模型（Mix-of-Experts,MoE）的分组矩阵乘法。MoE 模型是一种新兴的架构，广泛应用于大规模语言模型和多模态模型中。这种架构需要更灵活的计算策略，DeepGEMM 提供了连续布局（contiguous）和掩码布局（masked）两种 MoE 操作。

通过支持多种矩阵操作，DeepGEMM 能够满足不同应用场景的需求，进一步提升了其通用性和适用性。

（六）轻量级设计

DeepGEMM 的核心代码仅约 300 行，设计简洁，易于学习和二次开发。它避免了对现有库（如 CUTLASS 和 CuTe）的过度依赖，同时借鉴了它们的部分概念。这种轻量级设计不仅降低了开发和维护成本，还使得 DeepGEMM 更容易被集成到现有的深度学习框架中。

三、技术挑战与优化策略

（一）FP8 精度与性能的平衡

FP8 格式虽然计算效率高，但精度较低，容易导致累积误差。DeepGEMM 通过细粒度的缩放策略和两级累加机制来减少精度损失，同时保持高性能。然而，这种平衡并不容易实现。开发者需要在精度和速度之间进行权衡，确保在不同的应用场景中都能达到最优的性能。

（二）JIT 编译的复杂性

JIT 编译需要在运行时动态生成和优化内核代码。这不仅增加了开发难度，还需要确保编译器能够在短时间内生成高效的代码。为了实现这一点，DeepGEMM 的开发者需要对编译器技术有深入的理解，并且需要不断优化 JIT 编译器的性能。

（三）硬件适配与优化

DeepGEMM 针对 NVIDIA Hopper 架构进行了深度优化，充分利用了 TMA 和 Tensor Cores。然而，这种深度适配需要对硬件架构有深入的理解，并且需要不断调整以适应不同型号的 GPU。此外，随着硬件技术的不断进步，DeepGEMM 需要持续更新以支持新的硬件特性。

（四）支持多种矩阵操作

DeepGEMM 需要同时支持密集矩阵乘法和 MoE 模型的分组矩阵乘法。这要求库能够灵活处理不同类型的矩阵布局和计算策略，增加了代码复杂度。为了实现这一点，DeepGEMM 的开发者需要在代码设计上进行优化，确保库的灵活性和通用性。

（五）性能调优与通用性

DeepGEMM 在某些矩阵形状上表现优异，但在其他形状上可能表现不佳。因此，需要持续优化以提高其在更多场景下的通用性和性能。为了实现这一点，DeepGEMM 的开发者需要不断收集用户反馈，优化算法，并进行大量的实验验证。

四、DeepGEMM 的性能表现

根据 DeepSeek 的测试结果，DeepGEMM 在 NVIDIA Hopper 架构上实现了显著的性能提升。在标准的密集矩阵乘法中，DeepGEMM 的速度比传统的 FP32 实现快 2.7 倍。在 MoE 模型的分组矩阵乘法中，DeepGEMM 也表现出色，速度提升显著。

此外，DeepGEMM 的轻量级设计使得它能够轻松集成到现有的深度学习框架中。开发者可以通过简单的 API 调用，快速实现 FP8 矩阵乘法的加速。这种易用性进一步提升了 DeepGEMM 的应用价值。