Our latest blog presents work on Jagged Flash Attention (JFA) — the attention kernel behind
@Meta's Generative Ads Model (GEM) — on
@NVIDIA Blackwell (B200), built with TLX (Triton Low-level Extensions), which add explicit, hardware-aware control on top of Triton's high-level, tile-based programming model.
The TLX attention kernel is about 3.2K lines of concise Triton-level code — roughly 3× less than the ~10K-line CuteDSL kernels of the state-of-the-art FlashAttention-4 (FA4), and it outperforms FA4 on the jagged shapes that matter for GEM — by ~13% on the forward pass and ~50% on the backward pass.
Read it here:
✍️ Han Xu, Jacky Zhou, Jackie (Jiaqi) Xu, Hongtao Yu, Peng Chen (Dev Infra), Darren Liu, Dev (Devashish) Shankar (
@dev17392), Max Leung, Nick Riasanovsky, Hao Yan, Manman Ren, Yuanwei (Kevin) Fang