Skip to content

README.md

Latest commit

 

History

History
56 lines (38 loc) · 2.64 KB

README.md

File metadata and controls

56 lines (38 loc) · 2.64 KB
Raw

Ads Model Kernel Library

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.

Projects

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

Requirements

  • 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.1 for TLX packages

Installation

For CuteDSL GDPA:

pip install nvidia-cutlass-dsl>=4.1.0 torch einops

For TLX packages, see the project-specific environment.yml files or install the matching Triton build:

pip install fbtriton==3.6.1

Quick Start

See individual project READMEs for detailed usage:

Contributors

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.

License

Apache License 2.0. See LICENSE for details.