# facebookresearch/ads_model_kernel_library architecture

> Architecture diagram of the GitHub repository facebookresearch/ads_model_kernel_library, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/facebookresearch/ads_model_kernel_library
- Repository: https://github.com/facebookresearch/ads_model_kernel_library

## Overview

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: [`gdpa/src/interface.py`](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa/src/interface.py) (GDPA API), [`lp_fa4/src/lp_fa4/cute/interface.py`](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/lp_fa4/src/lp_fa4/cute/interface.py) (FlashAttention API), [`tlx_jfa/src/tlx_jagged_flash_attention.py`](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/tlx_jfa/src/tlx_jagged_flash_attention.py) (Jagged attention), [`block_attention/triton/tlx_block_attention.py`](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/block_attention/triton/tlx_block_attention.py) (Block attention)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_gdpa["GDPA attention"]
  node_gdpa_api["GDPA API<br/>[interface.py]"]
  node_gdpa_fwd["GDPA forward<br/>[flash_fwd_sm100.py]"]
  node_gdpa_bwd["GDPA backward<br/>[flash_bwd.py]"]
  node_gdpa_triton["Triton GDPA"]
end

subgraph group_block["Block attention"]
  node_block_api["Block attention"]
end

subgraph group_norm["Fused normalization"]
  node_norm_fwd["Matmul norm"]
  node_norm_layer_bwd["LayerNorm backward"]
  node_norm_layer_fwd["LayerNorm forward"]
  node_norm_rms_bwd["RMSNorm backward"]
  node_norm_rms_matmul["RMSNorm matmul"]
end

subgraph group_mega["GDPA megakernel"]
  node_mega_api["GDPA megakernel"]
end

subgraph group_jfa["Jagged attention"]
  node_jfa_api["Jagged attention"]
  node_jfa_bwd["1-CTA backward<br/>[bwd_1cta.py]"]
end

subgraph group_fa4["FlashAttention 4"]
  node_fa4_api["FlashAttention API<br/>[interface.py]"]
  node_fa4_fwd["FlashAttention forward<br/>[flash_fwd.py]"]
  node_fa4_bwd["FlashAttention backward<br/>[flash_bwd.py]"]
  node_fa4_paged["Paged KV support<br/>[paged_kv.py]"]
end

node_caller(("Kernel caller"))
node_result["Computed tensors"]

node_caller -->|"invokes"| node_gdpa_api
node_gdpa_api -->|"dispatches"| node_gdpa_fwd
node_gdpa_api -->|"dispatches"| node_gdpa_bwd
node_gdpa_api -->|"dispatches"| node_gdpa_triton
node_gdpa_api -->|"returns"| node_result
node_caller -->|"invokes"| node_block_api
node_block_api -->|"returns"| node_result
node_caller -->|"invokes"| node_norm_fwd
node_caller -->|"invokes"| node_norm_layer_fwd
node_caller -->|"invokes"| node_norm_layer_bwd
node_caller -->|"invokes"| node_norm_rms_bwd
node_caller -->|"invokes"| node_norm_rms_matmul
node_norm_fwd -->|"returns"| node_result
node_norm_layer_fwd -->|"returns"| node_result
node_norm_layer_bwd -->|"returns"| node_result
node_norm_rms_bwd -->|"returns"| node_result
node_norm_rms_matmul -->|"returns"| node_result
node_caller -->|"invokes"| node_mega_api
node_mega_api -->|"returns"| node_result
node_caller -->|"invokes"| node_jfa_api
node_jfa_api -->|"dispatches"| node_jfa_bwd
node_jfa_api -->|"returns"| node_result
node_caller -->|"invokes"| node_fa4_api
node_fa4_api -->|"dispatches"| node_fa4_fwd
node_fa4_api -->|"dispatches"| node_fa4_bwd
node_fa4_api -.->|"uses"| node_fa4_paged
node_fa4_api -->|"returns"| node_result

click node_gdpa_api "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa/src/interface.py"
click node_gdpa_fwd "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa/src/flash_fwd_sm100.py"
click node_gdpa_bwd "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa/src/flash_bwd.py"
click node_gdpa_triton "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa/triton/triton_generalized_dot_product_attention.py"
click node_block_api "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/block_attention/triton/tlx_block_attention.py"
click node_norm_fwd "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_rmsnorm_matmul.py"
click node_norm_layer_bwd "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_matmul_layernorm_bwd.py"
click node_norm_layer_fwd "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_matmul_layernorm.py"
click node_norm_rms_bwd "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_matmul_rmsnorm_bwd.py"
click node_norm_rms_matmul "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_rmsnorm_matmul.py"
click node_mega_api "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa_megakernel/src/tlx_gdpa_megakernel.py"
click node_jfa_api "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/tlx_jfa/src/tlx_jagged_flash_attention.py"
click node_jfa_bwd "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/tlx_jfa/src/bwd_1cta.py"
click node_fa4_api "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/lp_fa4/src/lp_fa4/cute/interface.py"
click node_fa4_fwd "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/lp_fa4/src/lp_fa4/cute/flash_fwd.py"
click node_fa4_bwd "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/lp_fa4/src/lp_fa4/cute/flash_bwd.py"
click node_fa4_paged "https://github.com/facebookresearch/ads_model_kernel_library/blob/main/lp_fa4/src/lp_fa4/cute/paged_kv.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_gdpa_api,node_gdpa_fwd,node_gdpa_bwd,node_gdpa_triton toneBlue
class node_block_api toneAmber
class node_norm_fwd,node_norm_layer_bwd,node_norm_layer_fwd,node_norm_rms_bwd,node_norm_rms_matmul toneMint
class node_mega_api toneRose
class node_jfa_api,node_jfa_bwd,node_caller,node_result toneIndigo
class node_fa4_api,node_fa4_fwd,node_fa4_bwd,node_fa4_paged toneTeal
```

## Components

### GDPA attention

- [GDPA API](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa/src/interface.py): `gdpa/src/interface.py`
- [GDPA forward](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa/src/flash_fwd_sm100.py): `gdpa/src/flash_fwd_sm100.py`
- [GDPA backward](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa/src/flash_bwd.py): `gdpa/src/flash_bwd.py`
- [Triton GDPA](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa/triton/triton_generalized_dot_product_attention.py): `gdpa/triton/triton_generalized_dot_product_attention.py`

### Block attention

- [Block attention](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/block_attention/triton/tlx_block_attention.py): `block_attention/triton/tlx_block_attention.py`

### Fused normalization

- [Matmul norm](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_rmsnorm_matmul.py): `multi_cta_norm_fusion/src/tlx_rmsnorm_matmul.py`
- [LayerNorm backward](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_matmul_layernorm_bwd.py): `multi_cta_norm_fusion/src/tlx_matmul_layernorm_bwd.py`
- [LayerNorm forward](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_matmul_layernorm.py): `multi_cta_norm_fusion/src/tlx_matmul_layernorm.py`
- [RMSNorm backward](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_matmul_rmsnorm_bwd.py): `multi_cta_norm_fusion/src/tlx_matmul_rmsnorm_bwd.py`
- [RMSNorm matmul](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/multi_cta_norm_fusion/src/tlx_rmsnorm_matmul.py): `multi_cta_norm_fusion/src/tlx_rmsnorm_matmul.py`

### GDPA megakernel

- [GDPA megakernel](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/gdpa_megakernel/src/tlx_gdpa_megakernel.py): `gdpa_megakernel/src/tlx_gdpa_megakernel.py`

### Jagged attention

- [Jagged attention](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/tlx_jfa/src/tlx_jagged_flash_attention.py): `tlx_jfa/src/tlx_jagged_flash_attention.py`
- [1-CTA backward](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/tlx_jfa/src/bwd_1cta.py): `tlx_jfa/src/bwd_1cta.py`

### FlashAttention 4

- [FlashAttention API](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/lp_fa4/src/lp_fa4/cute/interface.py): `lp_fa4/src/lp_fa4/cute/interface.py`
- [FlashAttention forward](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/lp_fa4/src/lp_fa4/cute/flash_fwd.py): `lp_fa4/src/lp_fa4/cute/flash_fwd.py`
- [FlashAttention backward](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/lp_fa4/src/lp_fa4/cute/flash_bwd.py): `lp_fa4/src/lp_fa4/cute/flash_bwd.py`
- [Paged KV support](https://github.com/facebookresearch/ads_model_kernel_library/blob/main/lp_fa4/src/lp_fa4/cute/paged_kv.py): `lp_fa4/src/lp_fa4/cute/paged_kv.py`

### Other components

- **Kernel caller**
- **Computed tensors**

## Connections

- 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