How facebookresearch/ads_model_kernel_library works
This repository is a portfolio of GPU kernel projects for attention and fused normalization, rather than one unified application. The principal workflow is a caller invoking a project’s attention or normalization API and receiving computed tensors; attention projects include forward and, where implemented, backward computation. The README documents all six projects. Sampled source confirms selected kernel launches and autograd registration, but not every public entry point or cross-module call, so uncertain wiring is omitted.
Main entry points: src/interface.py GDPA API · cute/interface.py FlashAttention API · tlx_jagged_flash_attention.py Jagged attention · tlx_block_attention.py Block attention
Components 19
GDPA attention
- GDPA API
gdpa/src/interface.py - GDPA forward
gdpa/src/flash_fwd_sm100.py - GDPA backward
gdpa/src/flash_bwd.py - Triton GDPA
gdpa/triton/triton_generalized_dot_product_attention.py
Block attention
- Block attention
block_attention/triton/tlx_block_attention.py
Fused normalization
- Matmul norm
multi_cta_norm_fusion/src/tlx_rmsnorm_matmul.py - LayerNorm backward
multi_cta_norm_fusion/src/tlx_matmul_layernorm_bwd.py - LayerNorm forward
multi_cta_norm_fusion/src/tlx_matmul_layernorm.py - RMSNorm backward
multi_cta_norm_fusion/src/tlx_matmul_rmsnorm_bwd.py - RMSNorm matmul
multi_cta_norm_fusion/src/tlx_rmsnorm_matmul.py
GDPA megakernel
- GDPA megakernel
gdpa_megakernel/src/tlx_gdpa_megakernel.py
Jagged attention
- Jagged attention
tlx_jfa/src/tlx_jagged_flash_attention.py - 1-CTA backward
tlx_jfa/src/bwd_1cta.py
FlashAttention 4
- FlashAttention API
lp_fa4/src/lp_fa4/cute/interface.py - FlashAttention forward
lp_fa4/src/lp_fa4/cute/flash_fwd.py - FlashAttention backward
lp_fa4/src/lp_fa4/cute/flash_bwd.py - Paged KV support
lp_fa4/src/lp_fa4/cute/paged_kv.py
Other components
- Kernel caller
- Computed tensors
Connections 27
- Kernel caller → GDPA API: invokes
- GDPA API → GDPA forward: dispatches
- GDPA API → GDPA backward: dispatches
- GDPA API → Triton GDPA: dispatches
- GDPA API → Computed tensors: returns
- Kernel caller → Block attention: invokes
- Block attention → Computed tensors: returns
- Kernel caller → Matmul norm: invokes
- Kernel caller → LayerNorm forward: invokes
- Kernel caller → LayerNorm backward: invokes
- Kernel caller → RMSNorm backward: invokes
- Kernel caller → RMSNorm matmul: invokes
- Matmul norm → Computed tensors: returns
- LayerNorm forward → Computed tensors: returns
- LayerNorm backward → Computed tensors: returns
- RMSNorm backward → Computed tensors: returns
- RMSNorm matmul → Computed tensors: returns
- Kernel caller → GDPA megakernel: invokes
- GDPA megakernel → Computed tensors: returns
- Kernel caller → Jagged attention: invokes
- Jagged attention → 1-CTA backward: dispatches
- Jagged attention → Computed tensors: returns
- Kernel caller → FlashAttention API: invokes
- FlashAttention API → FlashAttention forward: dispatches
- FlashAttention API → FlashAttention backward: dispatches
- FlashAttention API → Paged KV support: uses
- FlashAttention API → Computed tensors: returns