deepseek-ai/DeepGEMM

解决的问题

DeepGEMM 是一个高性能 Tensor Core 内核库，旨在优化现代大语言模型 (LLM) 的核心计算原语。它提供高度优化的通用矩阵乘法 (GEMM) 内核，能够媲美甚至超越专家调优的库，专门针对 NVIDIA SM90 和 SM100 架构进行了优化。

工作原理

该库使用轻量级的即时编译 (JIT) 模块在运行时编译内核，无需在安装期间进行 CUDA 编译。它实现了多个专用内核：

• Dense GEMMs：支持包括 FP8、FP4 和 BF16 在内的多种精度格式。
• Grouped GEMMs：针对混合专家 (MoE) 模型进行了优化，这些模型的专家共享形状，但处理不同的 Token 数量（连续和掩码布局）。
• Mega MoE：一种融合的“mega-kernel”，将 NVLink 通信 (EP dispatch/combine) 与 Tensor Core 计算 (Linear 1, SwiGLU, Linear 2) 进行重叠。
• MQA Kernels：为 DeepSeek v3.2 中的 lightning indexer 提供的专用评分内核。

适用对象

面向从事 LLM 推理和训练的开发人员和研究人员，特别是那些使用 NVIDIA H100/H800 (SM90) 或更新版本 (SM100) GPU 的用户，以及正在寻找简洁、易用的 GPU 内核优化学习资源的用户。

亮点

• 极致性能：在 H800 GPU 上实现高达 1550 TFLOPS 的性能。
• JIT 编译：内核在运行时编译，简化了安装和部署。
• 先进的 MoE 支持：包含将通信与计算重叠的融合 Mega MoE 内核。
• 广泛的精度支持：针对 FP8 和 FP4 精度进行了优化，以实现吞吐量最大化。
• 基于 C++20：使用现代 CUDA 代码库构建，避免了对复杂模板的沉重依赖。