PyTorch 博客· Han Xu, Jacky Zhou, Jackie (Jiaqi) Xu, Hongtao Yu, Peng Chen (Dev Infra), Darren Liu, Dev (Devashish) Shankar, Max Leung, Nick Riasanovsky, Hao Yan, Manman Ren, Yuanwei (Kevin) Fang·· 3 天前AI 评分44
TLX 优化 Jagged Flash Attention：迈向 Blackwell 上 SOTA FA4 之路
Optimizing Jagged Flash Attention with TLX: The Road Toward SOTA FA4 on Blackwell
AI 导读
Meta 团队使用 TLX 在 NVIDIA Blackwell 上优化了其生成式广告模型 GEM 的 Jagged Flash Attention 内核。该 TLX 内核代码量约为 3.2K 行，比当前 SOTA 的 FlashAttention-4 少约 3 倍，并在 GEM 关键的不规则形状上，前向传递性能提升约 13%，反向传递提升约 50%。
来源：PyTorch 博客 · pytorch.org