FlexAttention：兼具PyTorch灵活性和FlashAttention性能 (pytorch.org)

FlexAttention是一个新的PyTorch API，它在不损失灵活性的情况下，实现了与FlashAttention相当的性能。它允许用户通过自定义函数修改注意力分数，从而实现各种注意力机制，例如相对位置编码、ALiBi偏差、滑动窗口注意力、前缀语言模型和文档掩码。FlexAttention利用torch.compile将用户定义的函数降低到融合的FlashAttention内核中，并自动生成反向传播。它还支持利用注意力掩码中的稀疏性，从而显著提高性能。