PyTorch 发布 FlexAttention 注意力 API

PyTorch 团队发布新 API FlexAttention，让研究者通过 score_mod 和 mask_mod 两个用户函数，在几行 PyTorch 代码里实现相对位置编码等注意力变体，并称其保持 FlashAttention 级性能。 该 API 目前仅在 PyTorch nightly 版本中可用，计划作为原型功能随 2.5.0 发布。使用时要求所有序列长度为 128 的倍数，团队称将很快解决；GQA 支持尚未加入，当前需复制 kv heads 替代。团队还表示正改进性能，以在 H100 GPU 上追平 FlashAttention3。

AI 根据报道生成 · 56 分钟前更新

还没有足够的连续观测数据，暂不绘制趋势。