DeepGEMM

AICUDAGEMMDeep LearningGPU OptimizationPyTorch

近期趋势热门

总收藏数6,348

在 GitHub 上查看完整项目前往 GitHub →

项目简介

DeepGEMM 是一个轻量级的 CUDA 库，专为高效的通用矩阵乘法（GEMM）设计，支持 FP8 和 BF16 数据格式。它利用即时编译（JIT）模块消除了预安装内核编译的需求，同时保持了与专家调优库相当的性能。该库为密集型和 MoE 分组 GEMM 提供了专门的 API，是学习 GPU 内核优化的绝佳资源。

技术分析

DeepGEMM 是一个专为高性能通用矩阵乘法 (GEMMs) 设计的专业 CUDA 库，专门针对密集型和混合专家 (MoE) 架构中的 FP8 和 BF16 数据格式进行了优化。通过利用轻量级的即时 (JIT) 编译模块，它无需预安装内核编译，同时保持了媲美专家调优库的性能。该项目通过避免对模板的过度依赖来优先考虑简洁性和易用性，既可作为 DeepSeek 风格模型的生产级工具，也可作为 NVIDIA GPU 内核优化的教育资源。

核心亮点

支持针对密集型和 MoE 模型架构的高性能 FP8 和 BF16 GEMM 操作。

利用轻量级 JIT 编译模块在运行时编译内核，消除了对复杂预安装构建的需求。

提供针对连续和掩码布局的专用分组 GEMM API，针对 MoE 训练和推理场景进行了优化。

包含专为 DeepSeek v3.2 中使用的闪电索引器 (lightning indexer) 设计的专用 MQA (Multi-Query Attention) 评分内核。

在现代 NVIDIA 架构上实现高性能，在 H800 GPU 上可达到 1550 TFLOPS。

提供一套用于管理 TMA 对齐、Tensor Core 利用率和缩放因子转换的实用函数。

典型使用场景

为 NVIDIA GPU 提供高性能的 FP8 和 BF16 密集矩阵乘法

为连续和掩码布局提供高效的 MoE 分组 GEMM 操作

为高级模型索引和推理提供专门的 MQA logit 内核

快速开始

首先，递归克隆存储库以确保包含所有子模块。运行提供的 'develop.sh' 脚本以链接必要的包含文件并构建 C++ JIT 模块，然后执行 'tests/' 目录中的测试脚本以验证功能。最后，运行 'install.sh' 以完成设置，然后即可将 'deep_gemm' 导入到您的 Python 项目中。