High-performance GPU kernels for Meta Ads Recommendation Systems, developed by Meta Ads AI. This library provides optimized GPU kernel implementations that have been published.
| Project | Description | Architecture | Path | Blog |
|---|---|---|---|---|
| GDPA | Generalized Dot Product Attention kernels | Blackwell (SM100) | gdpa/ |
PyTorch Blog |
| TLX Block Attention | Triton TLX block attention kernels | Blackwell (SM100) | block_attention/ |
PyTorch Blog |
| TLX Multi-CTA Norm Fusion | Triton TLX fused matmul with RMSNorm and LayerNorm kernels | Blackwell (SM100) | multi_cta_norm_fusion/ |
PyTorch Blog |
| TLX GDPA Megakernel | Triton TLX generalized dot product attention megakernel | Blackwell (SM100) | gdpa_megakernel/ |
PyTorch Blog |
| TLX Jagged Flash Attention | Triton TLX jagged flash attention kernel for variable-length sequences | Blackwell (SM100) | tlx_jfa/ |
Coming soon |
- Python >= 3.10
- PyTorch >= 2.0
- NVIDIA GPU with Hopper (SM90) or Blackwell (SM100) architecture
- CUDA >= 12.0
- nvidia-cutlass-dsl >= 4.1.0
fbtriton==3.6.1for TLX packages
For CuteDSL GDPA:
pip install nvidia-cutlass-dsl>=4.1.0 torch einopsFor TLX packages, see the project-specific environment.yml files or install the matching Triton build:
pip install fbtriton==3.6.1See individual project READMEs for detailed usage:
- GDPA Quick Start
- TLX Block Attention Quick Start
- TLX Multi-CTA Norm Fusion
- TLX GDPA Megakernel
- TLX Jagged Flash Attention
Meta Ads AI: Jiaqi Xu, Hongtao Yu, Dev Shanker, Junqing (Jacky) Zhou, Han Xu, Jake Siso, Xiaoyi Liu, Huayu Li, Markus Hoehnerbach, Manman Ren, Chao Chen, Hao Yan, Weinan Song
More contributors will be added as we publish more kernels.
Apache License 2.0. See LICENSE for details.