[功能] [RFC] 支持可互换的注意力后端

作者: fffffgggg54创建于 2025年11月8日更新于 2025年11月9日
标签enhancement

**您的功能请求是否与问题有关?请描述。** 目前,许多模型依赖于标准的多头自注意力运算器。Timm 目前允许用户在两个版本之间选择,即 PyTorch 的迫切实现和 PyTorch 提供的融合实现(`torch.nn.functional.scaled_dot_product_attention`),以及通过 PyTorch 提供的 3 种实现(FA2、内存高效的注意力、迫切)。这可能会产生限制(其他地方提供的更好实现、阻碍 PT SDPA 正常工作的上游问题)或导致性能下降(FA3 和其他较新的实现)。为 timm 添加更多支持的后端,供用户选择(并最终允许用户注册自己的实现),将缓解这一限制。总的来说,目前对迫切与 SDPA 的处理方式也有些“黑科技”。

内容来源: huggingface/pytorch-image-models