DeepGEMM：让大模型矩阵乘法飞起来

痛点切入：当千亿参数卡在计算瓶颈

想象一下：你的DeepSeek-V3大模型正在进行推理，每生成一个token，都需要执行无数次矩阵乘法——而这些计算在普通GPU内核上跑得慢如牛车。H100显卡明明有上 TFLOPS 的算力，却被一个简单的GEMM（矩阵乘法）拖累到只能用三成。这是2024年大多数LLM团队的真实困境：硬件很贵，软件却拖了后腿。

DeepSeek团队也遇到了同样的问题。但他们没有选择忍——而是直接开源了自研的内核优化库 DeepGEMM，并在发布后迅速登上GitHub Trending榜首，在Hacker News上引发近400条讨论。

背景故事：DeepSeek的"造轮子"哲学

DeepSeek（深度求索）是一家来自中国的大模型公司，以"极低训练成本训练出世界顶级模型"闻名。2024年5月，他们在发布DeepSeek-V3时顺带开源了DeepGEMM——一个专注于FP8等低精度格式的GPU矩阵乘法内核库。

这个库的设计哲学非常有意思：不是为了炫技，而是为了可学习。DeepSeek在README中直接写到，库的设计借鉴了CUTLASS和CuTe的思想，但刻意减少了模板复杂度，让更多开发者能够读懂并学习GPU内核优化的精髓。这种"授人以渔"的态度，在大厂中是相当罕见的。

核心功能：不止是快，更是要简单

DeepGEMM的核心能力可以用一句话概括：在NVIDIA GPU上，以FP8/FP4等低精度格式，高效执行各类矩阵乘法。但如果展开来说，它支持的功能远不止于此：

第一，多精度格式支持。DeepGEMM同时支持FP8（8位浮点）、FP4（4位浮点）和BF16（16位浮点）三种格式的GEMM操作。这意味着它既可以用于训练时的前向/反向计算，也可以在推理阶段最大化吞吐量。FP8格式在H100/H800等新一代GPU上有专门张量核加速，相比FP16/BF16可以在保持模型精度的同时，将计算速度提升2-3倍。

第二，分组GEMM与MoE原生支持。这是DeepGEMM区别于其他GEMM库的关键特性。在Mixture-of-Experts（MoE）架构中，不同token会路由到不同的专家网络，导致需要执行大量小矩阵乘法——这些矩阵的M维度（batch维度）各不相同。DeepGEMM提供了m_grouped_gemm和k_grouped_gemm系列API，原生支持这种非规则形状的矩阵乘法，并且支持"掩码分组GEMM"，可以在CUDA Graph环境下优雅处理变长输出。

第三，Mega MoE：通信与计算融合。DeepGEMM的Mega MoE模块实现了DeepSeek-V3论文中描述的"重叠通信"机制。在分布式MoE训练中，专家需要接收来自其他节点的token——传统做法是等通信完成再计算，造成空闲等待。DeepGEMM通过SymmBuffer和融合内核，将All-to-All通信与GEMM计算流水线化，最大化GPU利用率。实测在H800上，Mega MoE可达到1350+ TFLOPS的吞吐。

第四，精细粒度缩放（Fine-grained Scaling）。FP8格式的挑战在于如何缩放——传统cuBLAS使用"tensorwise"缩放（整个张量用一个因子），DeepGEMM则实现了"blockwise"缩放（每个小矩阵块独立缩放），精度损失更小，性能也更高。这一设计在DeepSeek-V3论文中有详细理论支撑。

技术细节：从JIT编译到SASS优化

DeepGEMM的架构可以分为三层理解：

最上层是Python API层，通过pybind11将C++ CUDA内核暴露为Python可调用函数。用户直接import deep_gemm即可调用所有内核，无需编写CUDA代码。API设计遵循cuBLASLt的习惯（例如fp8_gemm_nt表示转置N×非转置T的GEMM），降低了学习成本。

中间层是C++ API层（csrc/apis/），定义了gemm.hpp、attention.hpp、einsum.hpp等头文件式接口。这些API是库的核心抽象，规定了不同操作的调用规范，也是学习内核优化的最佳入口——代码量不大（单个文件几百行），但涵盖了主流的GEMM变体。

最底层是CUDA JIT运行时（csrc/jit/）。这是DeepGEMM最具技术含量的部分。传统CUDA库需要在安装时编译PTX中间码（需要nvcc），DeepGEMM采用了完全不同的思路：运行时JIT编译。内核代码以字符串形式存储在C++源文件中，JIT模块在首次调用时动态编译到SASS（机器码），并缓存下来供后续使用。这意味着安装DeepGEMM不需要CUDA编译环境——pip install即可，甚至可以通过环境变量DG_JIT_USE_NVRTC=1启用NVRTC加速编译（编译速度提升10倍，略有性能损失）。

值得注意的是，DeepSeek团队在2025年7月进行了一次重大重构，彻底移除了NVRTC和后编译SASS优化，改为"无NVRTC、无后优化"的纯JIT方案。原因很实在：NVCC 12.9及以上版本会自动做FFMA SASS交织优化，手动优化反而可能帮倒忙。这个决定体现了DeepSeek团队"让代码更简单"的一贯理念。

上手体验：安装简单，门槛不低

安装方式有两种：通过pip直接安装预编译wheel，或者从源码编译。

pip安装只需一行，但要求CUDA版本、PyTorch版本、Python版本与wheel兼容。DeepGEMM通过setup.py从GitHub Releases下载对应版本的预编译包，文件命名包含CUDA版本、torch版本、cxx11abi版本等关键信息，确保二进制兼容。

从源码安装需要先安装依赖（CUDA Toolkit、PyTorch、CMake），然后运行install.sh。编译过程由setuptools调用torch的CUDAExtension自动处理，不需要手动nvcc。

使用门槛方面，DeepGEMM面向的是有深度学习系统背景的开发者。你需要理解FP8/FP4量化概念、GEMM的N/T布局差异，以及MoE架构的基本原理。如果是纯应用层用户，可能更适合使用基于DeepGEMM的上层框架（如vLLM），而不是直接调用内核API。

性能数据：H800上1550 TFLOPS

DeepSeek官方发布的benchmark数据相当亮眼：

在H800（受出口限制的H100降速版）上，DeepGEMM FP8 GEMM在多种矩阵形状下达到1350-1550 TFLOPS，全面超越cuBLASLt基线。在Mega MoE场景下，DeepGEMM的通信计算融合方案比传统"通信等计算"方案提升约40%吞吐量。

这些数据在GitHub README中有完整的测试条件和测试脚本可查，但需要注意：DeepGEMM的测试使用了最优的矩阵形状和CUDA配置。实际业务场景中，由于矩阵形状不规则或GPU资源竞争，性能可能低于理论峰值。

局限与争议

DeepGEMM并非没有短板：

第一，硬件依赖严重。DeepGEMM仅支持NVIDIA SM90（Hopper架构）和SM100（Blackwell架构）GPU，不支持更早的Ampere（A100）或Turing（A30）架构。这意味着如果你的GPU集群中有大量A100，将无法使用。

第二，平台锁定。目前仅支持Linux，不支持Windows或macOS。

第三，精度风险。FP4格式虽然能大幅提升吞吐，但在某些模型上可能导致精度下降。DeepGEMM在README中坦承"FP4可能不适合所有场景"，建议用户在实际部署前做充分验证。

第四，Scaling策略争议。2025年9月有用户在GitHub Issue中指出，DeepGEMM benchmark中的cuBLAS对比测试使用了tensorwise scaling而非blockwise scaling，导致cuBLAS性能被低估。这一争议提示我们：benchmark数据需谨慎解读，实际性能应以业务场景为准。

行业意义：开源基础设施的新标杆

DeepGEMM的开源具有多重意义：

对学术研究，DeepGEMM是目前学习现代GPU内核优化的最佳资源之一。代码量适中（核心JIT运行时约3000行），架构清晰，没有CUTLASS那样的模板地狱，非常适合作为内核优化课程的参考资料。

对工程实践，DeepGEMM为LLM推理框架（如vLLM、TensorRT-LLM）提供了底层加速能力。随着FP8在推理场景的普及，DeepGEMM很可能成为下一代LLM推理引擎的标准组件。

对开源生态，DeepSeek在半年内连续开源了FlashMLA、DeepEP、DeepGEMM、DeepSeek-V3论文等重磅项目，展现了"将内部技术积累回馈社区"的诚意。这种开放态度正在改变中国AI公司"封闭自研"的传统形象。

快速开始

import torch
import deep_gemm

# FP8矩阵乘法
A = torch.cuda.FloatTensor(...)
B = torch.cuda.FloatTensor(...)
D = deep_gemm.fp8_gemm_nt(A, B)  # N×T布局

# 分组GEMM（MoE场景）
D = deep_gemm.m_grouped_fp8_gemm_nt_contiguous(...)

图1：DeepGEMM支持的GEMM类型
（由于仓库图片在代理环境下无法访问，此处省略，可访问 https://github.com/deepseek-ai/DeepGEMM 查看官方benchmark图表）

DeepGEMM用简洁的代码实现了顶级性能，是LLM系统优化的重要基础设施。无论你是想学习GPU内核优化，还是在为LLM推理寻找加速方案，DeepGEMM都值得深入研究。