facebookresearch/ads_model_kernel_library
GitHub: facebookresearch/ads_model_kernel_library
Meta Ads AI 开源的高性能 GPU kernel 库，为广告与推荐系统模型训练提供经过真实业务负载验证的优化算子。
Stars: 32 | Forks: 8
# Ads Model Kernel 库
由 Meta Ads AI 开发的高性能 GPU kernel，专为 Meta Ads 推荐系统设计。本库提供了已发布的优化版 GPU kernel 实现。
## 项目
| 项目 | 描述 | 架构 | 路径 |
|---------|-------------|--------------|------|
| [GDPA](gdpa/) | 通用点积注意力 kernel | Blackwell (SM100) | `gdpa/` |
| [TLX Block Attention](block_attention/) | Triton TLX block attention kernel | Blackwell (SM100) | `block_attention/` |
| [TLX Multi-CTA Norm Fusion](multi_cta_norm_fusion/) | Triton TLX 融合 matmul 与 RMSNorm 及 LayerNorm 的 kernel | Blackwell (SM100) | `multi_cta_norm_fusion/` |
| [TLX GDPA Megakernel](gdpa_megakernel/) | Triton TLX 通用点积注意力 megakernel | Blackwell (SM100) | `gdpa_megakernel/` |
## 环境要求
- Python >= 3.10
- PyTorch >= 2.0
- 搭载 Hopper (SM90) 或 Blackwell (SM100) 架构的 NVIDIA GPU
- CUDA >= 12.0
- [nvidia-cutlass-dsl](https://github.com/NVIDIA/cutlass) >= 4.1.0
- 针对 TLX 包的 `fbtriton==3.6.1`
## 安装说明
对于 CuteDSL GDPA：
```
pip install nvidia-cutlass-dsl>=4.1.0 torch einops
```
对于 TLX 包，请参阅各个项目特定的 `environment.yml` 文件，或安装匹配的 Triton 构建版本：
```
pip install fbtriton==3.6.1
```
## 快速开始
有关详细用法，请参阅各个项目的 README：
- [GDPA 快速开始](gdpa/README.md#quick-start)
- [TLX Block Attention 快速开始](block_attention/README.md#quick-start)
- [TLX Multi-CTA Norm Fusion](multi_cta_norm_fusion/README.md)
- [TLX GDPA Megakernel](gdpa_megakernel/README.md)
## 贡献者
**Meta Ads AI：** Jiaqi Xu, Chao Chen, Hongtao Yu, Dev Shanker, Junqing (Jacky) Zhou, Han Xu, Jake Siso, Xiaoyi Liu, Huayu Li, Markus Hoehnerbach, Manman Ren
随着我们发布更多 kernel，将添加更多贡献者。
## 许可证
Apache License 2.0。详情请参阅 [LICENSE](LICENSE)。
标签：Vectored Exception Handling, 凭据扫描, 推荐系统, 模型训练, 注意力机制, 算子优化, 逆向工具, 高性能计算