一、项目概述

DeepGEMM 是一个专为高效 FP8 通用矩阵乘法（GEMMs）设计的库，支持普通和专家混合（MoE）分组 GEMMs，在 NVIDIA Hopper 张量核心上运行，采用 CUDA 编写，并通过轻量级即时编译（JIT）模块在运行时编译内核，无需安装时编译，其核心内核函数约 300 行代码，简洁易懂，有助于学习 Hopper FP8 矩阵乘法与优化技术。

二、性能表现

• 普通 GEMMs（密集模型）：在 H800 上使用 NVCC 12.8 测试 DeepSeek-V3/R1 推理中可能用到的矩阵形状，与基于 CUTLASS 3.6 的内部优化实现对比，如在 (64, 2112, 7168) 形状下计算性能达 206 TFLOPS，内存带宽 1688 GB/s，速度提升 2.7 倍；不同形状下性能提升在 1.1 – 2.7 倍之间，不过部分形状表现有待优化。
• 分组 GEMMs（MoE 模型）
  • 连续布局：如 (4, 8192, 4096, 7168) 配置下计算性能 1297 TFLOPS，内存带宽 418 GB/s，速度提升 1.2 倍左右。
  • 掩码布局：像 (1, 1024, 4096, 7168) 组合计算性能 1233 TFLOPS，内存带宽 924 GB/s，速度提升约 1.2 倍。

三、使用方法

• 环境要求：需 Hopper 架构 GPU（支持 sm_90a）、Python 3.8 及以上、CUDA 12.3 及以上（推荐 12.8 及以上）、PyTorch 2.1 及以上、CUTLASS 3.6 及以上。
• 安装步骤：先克隆仓库（git clone –recursive [email protected]:deepseek-ai/DeepGEMM.git），创建第三方库包含目录的符号链接（python setup.py develop），可测试 JIT 编译（python tests/test_jit.py）和所有 GEMM 实现（python tests/test_core.py），最后安装（python setup.py install），在 Python 项目中导入 deep_gemm 即可使用。

四、接口说明

• 普通密集 GEMMs：调用 deep_gemm.gemm_fp8_fp8_bf16_nt 函数执行非分组 FP8 GEMM，注意左乘因子需 TMA 对齐和转置，仅支持 NT 格式。
• 分组 GEMMs（连续布局）：按 MoE 模型专家共享形状需求仅在 M 轴分组，N 和 K 固定，训练或推理预填充时需按规则处理数据，参考 m_grouped_gemm_fp8_fp8_bf16_nt_contiguous 函数文档。
• 分组 GEMMs（掩码布局）：推理解码且启用 CUDA 图时，用 m_grouped_gemm_fp8_fp8_bf16_nt_masked 函数结合掩码张量计算有效部分，可参考 DeepEP 低延迟内核输出作为输入示例。

五、优化措施

• 持续线程专业化：借鉴 CUTLASS，实现数据移动、张量核心 MMA 指令和 CUDA 核心提升的重叠。
• Hopper TMA 特性利用：用于加载 / 存储矩阵及相关因子、LHS 多播和描述符预取。
• 通用细节优化：包括使用 stmatrix 指令、控制寄存器数量、最大化指令重叠、统一块调度、提高缓存重用等。
• 完全 JIT 设计：运行时编译，编译器可基于常量信息优化，自动选择参数提升小形状性能。
• 非对齐块大小支持：针对特定形状避免 SM 利用率低问题，提升性能但需精细优化。
• FFMA SASS 交织优化：修改 FFMA 指令提升细粒度缩放 FP8 GEMMs 性能。