Flex Attention API 应用 Notebook 代码速览

文章预览

对FlexAttention的常见API的使用方法做一个解读，博客来源：https://github.com/pytorch-labs/attention-gym/blob/main/examples/flex_attn.ipynb ，在此基础上我对部分代码添加了一些解释，修复了几个代码中的bug并使用PyTorch的nightly版本运行了示例，得到了每个custom attention的输出，展示在了下面的每个示例代码后面。最后还补充了一下torch compile inductor后端中实现FlexAttention的入口的代码浏览。 FlexAttention API 使用 NoteBook 本笔记本演示了新的 FlexAttention API 的使用方法，该 API 允许用户指定对缩放点积注意力（SDPA）中计算的注意力分数进行修改。 介绍 FlexAttention API 允许用户在Fused Scaled Dot Product Attention Kernel中指定对注意力分数的自定义修改。这使得各种注意力模式和偏置能够高效地实现，并具有潜在的运行时和内存节省。API 还将根据用户定义的修改生成融合的反向kernel。 设 ………………………………