vipshop/ffpa-attn

它解决了什么问题

FFPA（快速且内存高效的精确注意力）解决了在处理大头部维度（D > 256，最高 1024）时标准注意力机制的性能瓶颈。标准实现在这些规模下常常面临内存和寄存器压力，导致显著的速度下降。FFPA 提供了高度优化的内核，相比 PyTorch 原生的缩放点积注意力（SDPA）实现了大幅加速。

它是如何工作的

FFPA 使用两种主要技术来优化注意力：

• Split-D：这种分块策略将头部维度（D）轴拆分，以保持 SRAM 使用量恒定，有效解决内存瓶颈。
• TiledMMA：这种方法通过拆分矩阵乘法的 N 维度来优化寄存器使用，防止寄存器溢出，使内核能够处理比传统实现更大的头部维度。

该库支持多种后端，包括 CUDA、Triton 和 CuTe-DSL，并且可以通过一行代码集成到现有的 PyTorch 工作流中。

适用人群

本项目面向从事大规模模型训练和推理的 AI 研究人员和工程师，他们需要优化注意力层，特别是对于具有大头部维度或需要高吞吐量 FP8/FP4 精度的模型。

亮点

• 高性能：相比标准 PyTorch SDPA 实现 1.5 倍到 15 倍的加速。
• 大头部维度支持：将头部维度支持扩展到 1024，且无精度损失。
• 量化支持：包含实验性的 FP8 和 FP4 注意力内核，用于更快的推理。
• 灵活集成：提供简单的 API，可以猴子补丁 torch.nn.functional.scaled_dot_product_attention 以获得即时性能提升。
• 广泛兼容性：支持多种注意力类型，包括自注意力、GQA/MQA、交叉注意力以及因果/掩码注意力。