Skip to content

Commit 6879edf

Browse files
Han Xumeta-codesync[bot]
Han Xu
authored and
meta-codesync[bot]
committed
Add OSS TLX jagged flash attention package (#14)
Summary: Pull Request resolved: #14 - Copy the TLX jagged flash attention kernel into `ads_mkl/ops/oss/tlx_jfa` for OSS publication as a separate package from the GDPA megakernel. - Keep the kernel and local support helpers under `src/`, and add standalone correctness tests under `tests/`. - Replace internal `ads_mkl.ops.*` imports with local helper modules for packed-f32x2 math, custom registration, and index-width selection. - Drop the IKBO forward branch and the `ADS_MKL_AUTOTUNE_CONFIG_SET` / `ADS_MKL_DISABLE_AUTOTUNE` machinery; only the default autotune configs remain. - Split the general 1-CTA backward into `bwd_1cta.py` and the shared JIT helpers into `kernel_common.py`, and enable the 2-CTA collaborative-MMA backward for the PMA case; every other shape falls back to 1-CTA. - Add conda/pip setup around `fbtriton==3.6.1` and document BUCK-free unittest execution. - Add the `tlx_jfa` package to the top-level OSS project table and quick-start section. Reviewed By: jackiexu77 Differential Revision: D115074159 fbshipit-source-id: 703515301f718f305a63a4842e3086a6c8755517
1 parent e9f07d5 commit 6879edf

11 files changed

Lines changed: 5731 additions & 0 deletions

‎README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ High-performance GPU kernels for Meta Ads Recommendation Systems, developed by M
1010
| [TLX Block Attention](block_attention/) | Triton TLX block attention kernels | Blackwell (SM100) | `block_attention/` | [PyTorch Blog](https://pytorch.org/blog/tlx-block-attention-a-warp-specialized-blackwell-kernel-for-fixed-block-sparse-self-attention/) |
1111
| [TLX Multi-CTA Norm Fusion](multi_cta_norm_fusion/) | Triton TLX fused matmul with RMSNorm and LayerNorm kernels | Blackwell (SM100) | `multi_cta_norm_fusion/` | [PyTorch Blog](https://pytorch.org/blog/towards-free-normalization-fusing-normalization-into-gemm-and-attention-kernels/) |
1212
| [TLX GDPA Megakernel](gdpa_megakernel/) | Triton TLX generalized dot product attention megakernel | Blackwell (SM100) | `gdpa_megakernel/` | [PyTorch Blog](https://pytorch.org/blog/towards-free-normalization-fusing-normalization-into-gemm-and-attention-kernels/) |
13+
| [TLX Jagged Flash Attention](tlx_jfa/) | Triton TLX jagged flash attention kernel for variable-length sequences | Blackwell (SM100) | `tlx_jfa/` | Coming soon |
1314

1415
## Requirements
1516

@@ -42,6 +43,7 @@ See individual project READMEs for detailed usage:
4243
- [TLX Block Attention Quick Start](block_attention/README.md#quick-start)
4344
- [TLX Multi-CTA Norm Fusion](multi_cta_norm_fusion/README.md)
4445
- [TLX GDPA Megakernel](gdpa_megakernel/README.md)
46+
- [TLX Jagged Flash Attention](tlx_jfa/README.md)
4547

4648
## Contributors
4749

‎tlx_jfa/README.md‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
# TLX Jagged Flash Attention
2+
3+
Triton TLX jagged flash attention kernel for Blackwell, with warp-specialized
4+
forward and backward passes over variable-length (jagged) sequences.
5+
6+
## Supported variants
7+
8+
- **Jagged self- and cross-attention.** Query and key/value sequence lengths are
9+
independent, each given as `[batch_size + 1]` prefix-sum offsets.
10+
- **PMA (pooling by multi-head attention).** A single shared query sequence
11+
attends to per-batch jagged key/value sequences, via `broadcast_q=True`. See
12+
the Kunlun paper (<https://arxiv.org/abs/2602.10016>) for the definition. This
13+
is the case routed through the 2-CTA collaborative-MMA backward.
14+
- **Sliding window.** `window_size=W` restricts attention to a symmetric band of
15+
`+/- W` positions.
16+
- **Grouped query attention** in the forward, where the query head count is a
17+
multiple of the key/value head count. The backward supports a single query
18+
group only.
19+
20+
Forward and backward are both supported; backward runs through autograd.
21+
22+
## Layout
23+
24+
- `src/tlx_jagged_flash_attention.py` - public API, forward kernel, and the 2-CTA
25+
(cluster) collaborative-MMA backward
26+
- `src/bwd_1cta.py` - general 1-CTA backward, used for every configuration the
27+
2-CTA path does not cover
28+
- `src/kernel_common.py` - JIT helpers shared by the forward and both backwards
29+
- `src/tlx_math.py`, `src/register_helpers.py`, `src/utils.py` - local support helpers
30+
- `tests/` - GPU correctness tests
31+
32+
The 2-CTA backward is scoped to the PMA case: `broadcast_q` with
33+
`head_dim == 128`, a single query group, no sliding window, and load balancing
34+
enabled. The launcher routes every other shape to the 1-CTA backward
35+
automatically, so all supported variants work regardless of which kernel runs.
36+
37+
## Run
38+
39+
```bash
40+
conda env create -f environment.yml
41+
conda activate tlx-jfa
42+
TRITON_ALLOW_NON_CONSTEXPR_GLOBALS=1 python -m unittest discover -s tests -p "test_*.py"
43+
```
44+
45+
If conda channel access is restricted, create the environment with any available Python 3.12 conda channel and install the Python packages with pip:
46+
47+
```bash
48+
conda create -n tlx-jfa python=3.12 pip
49+
conda activate tlx-jfa
50+
pip install --upgrade pip setuptools wheel
51+
pip install --extra-index-url https://download.pytorch.org/whl/cu128 torch
52+
pip install fbtriton==3.6.1
53+
TRITON_ALLOW_NON_CONSTEXPR_GLOBALS=1 python -m unittest discover -s tests -p "test_*.py"
54+
```
55+
56+
To verify TLX is importable:
57+
58+
```bash
59+
python -c 'import triton.language.extra.tlx as tlx; print(tlx)'
60+
```
61+
62+
For interactive use outside the tests, add the kernel sources to `PYTHONPATH`:
63+
64+
```bash
65+
export PYTHONPATH="$PWD/src:$PYTHONPATH"
66+
```
67+
68+
## Usage
69+
70+
`query`, `key`, and `value` are jagged tensors of shape `[total_seq_len, num_heads, head_dim]`,
71+
where `total_seq_len` is the sum of the per-batch sequence lengths. `query_offset` and
72+
`key_offset` are `[batch_size + 1]` int32 prefix-sum offsets delimiting each sequence.
73+
74+
```python
75+
from tlx_jagged_flash_attention import jagged_flash_attention
76+
77+
out = jagged_flash_attention(
78+
query=q, # [total_q, H, D]
79+
key=k, # [total_kv, H, D]
80+
value=v, # [total_kv, H, D]
81+
query_offset=q_offsets, # [B + 1], int32
82+
key_offset=kv_offsets, # [B + 1], int32
83+
max_seq_len_q=max_seq_len_q,
84+
max_seq_len_kv=max_seq_len_kv,
85+
sm_scale=head_dim**-0.5,
86+
)
87+
```
88+
89+
Optional arguments:
90+
91+
- `window_size` - restricts attention to a symmetric band of `+/- window_size` positions.
92+
- `broadcast_q` - PMA: shares a single query sequence across all batches. `query`
93+
holds that one sequence, `query_offset` is `[0, q_len]`, and `output_offset`
94+
gives the per-batch output slots. The gradient of the shared query accumulates
95+
the contributions of every batch.
96+
- `sm_scale` - softmax scale; defaults to `1 / sqrt(head_dim)`.
97+
98+
Backward is supported through autograd; call `.backward()` on the output.
99+
100+
## Limitations
101+
102+
- Requires a Blackwell (SM100+) GPU.
103+
- `head_dim` must be `<= 128`. Larger head dims need the D-tiled IKBO forward, which
104+
is not part of this package.
105+
- The backward requires a single query group (query and key/value head counts equal).
106+
107+
## Reference
108+
109+
- Kunlun: <https://arxiv.org/abs/2602.10016> (PMA definition)
110+
- A PyTorch blog post covering this kernel is in preparation.

‎tlx_jfa/environment.yml‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
name: tlx-jfa
2+
channels:
3+
- conda-forge
4+
dependencies:
5+
- python=3.12
6+
- pip
7+
- pip:
8+
- --extra-index-url https://download.pytorch.org/whl/cu128
9+
- torch
10+
- fbtriton==3.6.1

0 commit comments

Comments
 (0)