DeepGEMM 架构设计
DeepGEMM 采用了精心设计的架构，专为 NVIDIA Hopper 张量核心优化，实现高效的 FP8 矩阵乘法。
完全 JIT 设计
DeepGEMM 采用完全即时编译(JIT)设计，安装时无需编译。所有内核在运行时使用轻量级 JIT 模块编译，将 GEMM 形状、块大小和流水线阶段数作为编译时常量处理，节省寄存器并允许编译器进行更多优化。
Hopper TMA 特性
充分利用 Hopper 架构的张量内存加速器(TMA)进行更快速和异步的数据移动，包括 TMA 加载、存储、多播和描述符预取，显著提升性能。
核心优化技术
DeepGEMM 采用了多种先进的优化技术，超越传统 GEMM 库，实现卓越性能。
持久化线程束专用化
遵循 CUTLASS 设计，使数据移动、张量核心 MMA 指令和 CUDA 核心提升重叠，提高计算效率。
非对齐块大小
支持非对齐块大小（如 112），以更好地利用 SM。例如，对于 M=256, N=7168，使用 BLOCK_M=128, BLOCK_N=112 可以使 128 个 SM 工作，而不是使用对齐块大小时的 112 个。
FFMA SASS 交错
修改编译后二进制文件中的 FFMA 指令，翻转 yield 位和 reuse 位，为 MMA 指令与提升 FFMA 指令的重叠创造更多机会，在某些情况下可提高 10%+ 的性能。
统一优化的块调度器
一个调度器适用于所有非分组和分组内核，采用光栅化以增强 L2 缓存重用，提高整体性能。
代码解析
DeepGEMM 的核心内核函数仅约 300 行代码，设计简洁，便于学习 Hopper FP8 矩阵乘法和优化技术。
// Simplified version of the core GEMM kernel
template <
int BLOCK_M,
int BLOCK_N,
int BLOCK_K,
int CLUSTER_M,
int CLUSTER_N,
int STAGES,
int WARPGROUP_SIZE,
bool ENABLE_INTERLEAVE
>
__global__ void gemm_fp8_fp8_bf16_nt_kernel(
const void* __restrict__ lhs,
const void* __restrict__ lhs_scale,
const void* __restrict__ rhs,
void* __restrict__ output,
const void* __restrict__ rhs_scale,
int m, int n, int k,
int lhs_stride, int rhs_stride, int output_stride) {
// Using TMA to load LHS, LHS scale factors, and RHS
// Using CUDA cores for two-level accumulation (promotion)
// Using TMA to store output
// Persistent thread specialization implementation
// Unaligned block size support
// FFMA SASS interleaving optimization
}代码结构特点
- 模板参数化设计，支持不同的块大小和优化配置
- 充分利用 Hopper 架构的 TMA 特性进行高效数据移动
- 实现两级累加解决 FP8 张量核心累加不精确的问题
- 持久化线程束专用化实现，优化计算和数据移动重叠
接口设计
- 普通密集 GEMM: 使用 deep_gemm.gemm_fp8_fp8_bf16_nt 函数
- 分组 GEMM（连续布局）: 使用 m_grouped_gemm_fp8_fp8_bf16_nt_contiguous 函数，适用于专家共享相同形状的 MoE 模型
- 分组 GEMM（掩码布局）: 使用 m_grouped_gemm_fp8_fp8_bf16_nt_masked 函数，用于启用 CUDA 图的推理解码阶段，CPU 不知道每个专家接收的令牌数量
性能表现
DeepGEMM 在各种矩阵形状上的性能与专家调优的库相当或更好，在某些情况下可达到 2.7 倍的加速。
性能亮点
- 普通 GEMM：在某些形状上可达到 2.7 倍的加速
- 连续布局的分组 GEMM：可达到 1.2 倍的加速
- 掩码布局的分组 GEMM：也可达到 1.2 倍的加速
- 在 H800 SXM5 上测试的计算性能最高可达 1358 TFLOPS
- 内存带宽利用率高，最高可达 2668 GB/s
DeepGEMM 的主要特性
卓越性能
DeepGEMM 在 NVIDIA Hopper 架构上相比专家调优的库可提供高达 2.7 倍的加速，计算性能超过 1350 TFLOPS
先进优化
DeepGEMM 采用持久化线程束专用化、非对齐块大小和 FFMA SASS 交错等技术，实现最大计算效率
灵活集成
DeepGEMM 通过清晰、文档完善的接口，与深度学习框架和科学计算库无缝集成
细粒度缩放
DeepGEMM 通过优化的块调度和光栅化技术，高效地从小型矩阵扩展到大规模计算
原生 FP8 支持
DeepGEMM 提供一流的 FP8 精度支持，专为 NVIDIA Hopper 张量核心优化
高内存带宽
DeepGEMM 实现高达 2668 GB/s 的内存带宽利用率，最大化硬件能力
清晰代码设计
DeepGEMM 的核心内核仅约 300 行代码，设计模式清晰，易于理解和扩展
优化的分组 GEMM
DeepGEMM 在连续和掩码布局的分组 GEMM 操作上表现出色，非常适合混合专家(MoE)模型
常见问题
- 什么是 DeepGEMM，它有什么特别之处？
- DeepGEMM 是一个为 NVIDIA Hopper 架构优化的高效 FP8 矩阵乘法库。它的特别之处在于简洁的设计（核心代码仅约 300 行）、卓越的性能（相比专家调优的库可提供高达 2.7 倍的加速）以及先进的优化技术，如持久化线程束专用化和 FFMA SASS 交错。
- DeepGEMM 针对哪些硬件进行了优化？
- DeepGEMM 专为 NVIDIA Hopper 架构优化，充分利用了 Hopper 的张量内存加速器(TMA)特性和张量核心进行 FP8 矩阵乘法。它在 H800 SXM5 GPU 上进行了测试，计算性能可达 1358 TFLOPS。
- DeepGEMM 如何实现性能提升？
- DeepGEMM 通过多种先进技术实现性能提升：完全 JIT 设计和编译时常量、非对齐块大小以更好地利用 SM、FFMA SASS 交错、两级累加解决 FP8 精度问题，以及采用光栅化增强 L2 缓存重用的统一优化块调度器。
- DeepGEMM 支持哪些类型的 GEMM 操作？
- DeepGEMM 支持三种主要类型的 GEMM 操作：普通密集 GEMM（使用 deep_gemm.gemm_fp8_fp8_bf16_nt 函数）、连续布局的分组 GEMM 和掩码布局的分组 GEMM。后两种特别适用于混合专家(MoE)模型。
- DeepGEMM 适合混合专家(MoE)模型吗？
- 是的，DeepGEMM 非常适合 MoE 模型。它提供了专门的连续和掩码布局分组 GEMM 操作，针对 MoE 架构中的特定计算模式进行了优化。这些操作相比其他库可提供高达 1.2 倍的加速。
- DeepGEMM 支持哪些精度？
- DeepGEMM 主要关注 FP8 精度，这对高效的 AI 模型推理越来越重要。它实现了两级累加方法来解决 FP8 张量核心可能出现的不精确累加问题，确保性能和精度兼顾。
- DeepGEMM 的内存带宽利用率与其他库相比如何？
- DeepGEMM 实现了卓越的内存带宽利用率，高达 2668 GB/s。这得益于其高效利用 Hopper 的 TMA 特性进行更快速和异步的数据移动，包括 TMA 加载、存储、多播和描述符预取。
- 我可以从 DeepGEMM 的实现中学习优化自己的 CUDA 代码吗？
- 当然可以！DeepGEMM 的核心内核仅约 300 行代码，设计清晰，是很好的学习资源。它展示了先进的 CUDA 优化技术、Hopper 架构特性的高效利用，以及可应用于其他高性能计算任务的巧妙矩阵乘法方法。