观点 / 方法普通
用 TLX 优化 Jagged Flash Attention:迈向 Blackwell 上的 SOTA FA4
原始标题:Optimizing Jagged Flash Attention with TLX: The Road Toward SOTA FA4 on Blackwell
内容摘要
Meta 用 TLX(Triton Low-level Extensions)重写 GEM 广告模型背后的 Jagged Flash Attention 内核,在 NVIDIA B200 上以约 3.2K 行 Triton 代码实现,比 FA4 约 10K 行的 CuteDSL 内核少约 3 倍。