综合介绍

DeepGEMM 是一个由 DeepSeek 团队开发的开源 FP8 GEMM（通用矩阵乘法）库，专注于提供高效的矩阵运算支持。它特别针对 NVIDIA Hopper 架构的 Tensor Core 设计，支持普通矩阵运算和混合专家模型（MoE）的分组 GEMM 操作。该库采用 CUDA 编写，通过轻量级的即时编译（JIT）技术实现运行时内核编译，无需安装时预编译，极大简化了部署流程。DeepGEMM 在保持简洁代码的同时，性能表现出色，在 Hopper GPU 上可达到超过 1350 TFLOPS 的 FP8 计算能力。它不仅适用于机器学习模型的训练和推理加速，还因其开源性质和易读性，成为学习 FP8 矩阵优化的绝佳资源。

功能列表

-支持 FP8 矩阵运算：提供高效的 FP8 通用矩阵乘法（GEMM），适用于高性能计算场景。
-MoE 模型优化：支持混合专家模型的分组 GEMM，仅对 M 轴分组，适配专家共享相同形状的场景。
-即时编译（JIT）：通过运行时编译内核，无需预编译即可适配不同硬件环境。
-高性能计算：在 NVIDIA Hopper GPU 上实现超过 1350 TFLOPS 的 FP8 计算吞吐量。
-简洁代码设计：核心代码约 300 行，易于学习和二次开发。
-兼容性强：支持普通 GEMM 和带掩码的分组 GEMM，适配多种推理场景。
-开源免费：基于 MIT 协议发布，适用于研究和商业用途。

使用帮助

DeepGEMM 是一个专为开发者设计的开源矩阵运算库，主要面向具备一定 CUDA 编程基础和机器学习背景的用户。以下是详细的使用指南，帮助你快速上手并将其集成到项目中。

安装流程

DeepGEMM 不需要复杂的预编译过程，只需几步即可完成安装和运行环境配置：
1.环境准备：

• 系统要求：支持 NVIDIA Hopper 架构的 GPU（如 H100）。
• 软件依赖：安装 CUDA Toolkit（建议版本 11.8 或更高）和 Python（3.8+）。
• 硬件支持：确保你的设备配备至少 40GB 显存的 NVIDIA GPU。
  2.克隆仓库：
  在终端运行以下命令，将 DeepGEMM 仓库下载到本地：

git clone https://github.com/deepseek-ai/DeepGEMM.git**
cd DeepGEMM**

3. 安装依赖：
   使用 Python 的包管理工具安装所需依赖：

pip install torch numpy

DeepGEMM 本身无需额外编译，因为它依赖即时编译技术，所有内核会在运行时自动生成。
4. 验证安装：
运行提供的测试脚本，确保环境配置正确：

python test/deep_gemm_test.py

如果输出显示正常的矩阵运算结果，说明安装成功。

主要功能操作

1. 执行基本的 FP8 GEMM 运算

DeepGEMM 提供了一个简单易用的接口，用于执行非分组的 FP8 矩阵乘法：

• 操作步骤：
  • 导入库和函数：

import torch
from deep_gemm import gemm_fp8_fp8_bf16_nt

• 准备输入数据（矩阵 A 和 B，必须是 FP8 格式）：

A = torch.randn(1024, 512, dtype=torch.float8_e4m3fn).cuda()
B = torch.randn(512, 1024, dtype=torch.float8_e4m3fn).cuda()

• 调用函数进行矩阵乘法：

C = gemm_fp8_fp8_bf16_nt(A, B)
print(C)

• 注意事项：
  • 输入矩阵需位于 GPU 上，且格式需为 FP8（E4M3 或 E5M2）。
  • 输出结果为 BF16 格式，适合后续计算或存储。

2. 支持 MoE 模型的分组 GEMM

对于需要处理 MoE 模型的用户，DeepGEMM 提供了分组 GEMM 支持：

• 操作步骤：
  • 导入分组 GEMM 函数：

from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_contiguous

• 准备连续布局的输入数据：

A = torch.randn(4096, 512, dtype=torch.float8_e4m3fn).cuda()  # 多个专家的输入拼接
B = torch.randn(512, 1024, dtype=torch.float8_e4m3fn).cuda()
group_sizes = [1024, 1024, 1024, 1024]  # 每个专家的 token 数

• 执行分组 GEMM：

C = m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(A, B, group_sizes)
print(C)

• 注意事项：
  • 输入矩阵 A 的 M 轴需按专家分组拼接，且每个分组的大小需对齐 GEMM M 块大小（可用 get_m_alignment_for_contiguous_layout() 获取）。
  • B 矩阵的 N 和 K 轴需保持固定。

3. 推理阶段的掩码分组 GEMM

在推理解码阶段，DeepGEMM 支持使用掩码的分组 GEMM，适用于动态 token 分配：

• 操作步骤：
  • 导入掩码分组函数：

from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_masked

• 准备输入数据和掩码：

A = torch.randn(4096, 512, dtype=torch.float8_e4m3fn).cuda()
B = torch.randn(512, 1024, dtype=torch.float8_e4m3fn).cuda()
mask = torch.ones(4096, dtype=torch.bool).cuda()  # 掩码指示有效 token

• 执行掩码分组 GEMM：

C = m_grouped_gemm_fp8_fp8_bf16_nt_masked(A, B, mask)
print(C)

• 注意事项：
  • 掩码用于指定哪些 token 需要计算，适合 CUDA 图启用时的动态推理。

特色功能操作流程

高性能优化与调试

DeepGEMM 的核心优势在于其高效性和简洁性，开发者可以通过以下步骤进一步优化和调试：

• 查看性能数据：
  在运行测试脚本时，添加日志输出以监控 TFLOPS：

  import logging
  logging.basicConfig(level=logging.INFO)
  C = gemm_fp8_fp8_bf16_nt(A, B)
  

调整参数*：
根据具体硬件调整块大小（TMA 参数）以优化数据移动和计算重叠，参考文档中的 test/ 文件夹中的示例。
学习与扩展*：
核心代码位于 deep_gemm/gemm_kernel.cu，约 300 行，开发者可直接阅读并修改以适配自定义需求。
使用建议
硬件要求*：目前仅支持 NVIDIA Hopper 架构 GPU，其他架构暂未适配。
文档参考*：详细函数说明和示例代码可在 GitHub 仓库的 README.md 和 test/ 文件夹中找到。
社区支持*：如遇问题，可在 GitHub Issues 页面提交反馈，DeepSeek 团队会积极响应。
通过以上步骤，你可以轻松将 DeepGEMM 集成到机器学习项目中，享受其高效的 FP8 矩阵运算能力。