7,104
// 项目简介
DeepGEMM 是一个统一的 CUDA 库,专门为现代大语言模型提供高性能 Tensor Core 内核。它具备轻量级的 Just-In-Time 编译模块,无需在安装过程中进行 CUDA 编译。该库为各种矩阵运算(包括 FP8、FP4 和 BF16 GEMM)以及融合的 MoE 和 MQA 评分提供了专家级优化性能。
// 技术分析
DeepGEMM 是一个高性能、统一的 CUDA 内核库,旨在为现代大语言模型提供必要的计算原语,包括 FP8/FP4 GEMM、融合 MoE 和 MQA 评分。通过利用轻量级的即时 (JIT) 编译模块,该库消除了安装过程中复杂的 CUDA 编译需求,同时保持了可与专家调优库相媲美的性能。该项目通过避免对复杂模板元编程的过度依赖,优先考虑简洁性和易用性,为开发者提供了一个清晰的代码库,以便学习和实现先进的 NVIDIA GPU 内核优化。
// 核心亮点
01
提供支持多种数据格式(包括 FP8、FP4 和 BF16)的高性能 GEMM 内核。
02
实现了 Mega MoE,这是一种将 NVLink 通信与 Tensor Core 计算重叠以最大化吞吐量的融合内核。
03
具备轻量级 JIT 模块,可在运行时编译内核,消除了手动 CUDA 编译的负担。
04
包含专为 lightning indexer 设计的专用 MQA 评分内核,支持非分页和分页内存布局。
05
支持 SM90 和 SM100 架构等先进硬件特性,并优化了 TMA 对齐的内存操作。
06
提供灵活的成组 GEMM API,适用于连续和掩码布局,专为高效的 MoE 专家处理而定制。
// 典型使用场景
01
面向 LLM 的高性能 FP8、FP4 和 BF16 GEMM 运算
02
融合了通信与 Tensor Core 计算的 Mega MoE 内核
03
用于大规模推理中闪电索引器的 MQA 评分内核
// 快速开始
首先,使用 'git clone --recursive' 递归克隆存储库,以确保包含所有子模块。运行提供的 'develop.sh' 脚本以链接必要的头文件并构建 C++ JIT 模块,然后运行 'install.sh' 完成安装。安装完成后,您可以直接在 Python 环境中导入 'deep_gemm' 模块以访问优化后的内核。