Skip to content

Commit 461eb63

Browse files
Re-sync with internal repository (#10)
The internal and external repositories are out of sync. This Pull Request attempts to brings them back in sync by patching the GitHub repository. Please carefully review this patch. You must disable ShipIt for your project in order to merge this pull request. DO NOT IMPORT this pull request. Instead, merge it directly on GitHub using the MERGE BUTTON. Re-enable ShipIt after merging. fbshipit-source-id: 0a51b956d84f761124079edc03a7c7fc0c706c82 Co-authored-by: Facebook Community Bot <6422482+facebook-github-bot@users.noreply.github.com>
1 parent 0e4008d commit 461eb63

52 files changed

Lines changed: 31610 additions & 0 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

README.md

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# Ads Model Kernel Library
2+
3+
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.
4+
5+
## Projects
6+
7+
More coming soon!
8+
9+
| Project | Description | Architecture | Path |
10+
|---------|-------------|--------------|------|
11+
| [GDPA](gdpa/) | Generalized Dot Product Attention kernels | Blackwell (SM100) | `gdpa/` |
12+
| [TLX Block Attention](block_attention/) | Triton TLX block attention kernels | Blackwell (SM100) | `block_attention/` |
13+
14+
## Requirements
15+
16+
- Python >= 3.10
17+
- PyTorch >= 2.0
18+
- NVIDIA GPU with Hopper (SM90) or Blackwell (SM100) architecture
19+
- CUDA >= 12.0
20+
- [nvidia-cutlass-dsl](https://github.com/NVIDIA/cutlass) >= 4.1.0
21+
22+
## Installation
23+
24+
```bash
25+
pip install nvidia-cutlass-dsl>=4.1.0 torch einops
26+
```
27+
28+
## Quick Start
29+
30+
See individual project READMEs for detailed usage:
31+
32+
- [GDPA Quick Start](gdpa/README.md#quick-start)
33+
34+
## Contributors
35+
36+
**Meta Ads AI:** Jiaqi Xu, Chao Chen, Hongtao Yu, Dev Shanker, Jacky Zhou, Han Xu, Jake Siso, Xiaoyi Liu, Huayu Li, Markus Hoehnerbach, Manman Ren
37+
38+
More contributors will be added as we publish more kernels.
39+
40+
## License
41+
42+
Apache License 2.0. See [LICENSE](LICENSE) for details.

block_attention/README.md

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# TLX Block Attention
2+
3+
Triton TLX-based block attention kernels for variable-length attention workloads.
4+
5+
## Layout
6+
7+
- `triton/` - Triton TLX kernel implementation
8+
- `tests/` - project tests
9+
10+
## Quick Start
11+
12+
```python
13+
from ads_mkl.ops.oss.block_attention.triton.tlx_block_attention import (
14+
block_attention_api,
15+
)
16+
```
17+
18+
## Run
19+
20+
```bash
21+
buck2 test fbcode//ads_mkl/ops/oss/block_attention:test_tlx_block_attention
22+
```
23+
24+
This test target includes:
25+
- `fbcode/ads_mkl/ops/oss/block_attention/triton/tests/test_tlx_block_attention.py`
26+
- `fbcode/ads_mkl/ops/oss/block_attention/tests/block_attention_test.py`

0 commit comments

Comments
 (0)