torch.nn.attention#

创建于: 2024 年 1 月 24 日 | 最后更新于: 2025 年 11 月 12 日

此模块包含修改 torch.nn.functional.scaled_dot_product_attention 的行为的函数和类

Utils#

`sdpa_kernel`	用于选择缩放点积注意力后端的上下文管理器。
`SDPBackend`	包含缩放点积注意力不同后端的类似枚举的类。
`register_flash_attention_impl`	注册用于激活闪光注意力实现的调用器。
`activate_flash_attention_impl`	将先前注册的闪光注意力实现激活到调度器中。
`list_flash_attention_impls`	返回所有可用闪光注意力实现的名称。
`current_flash_attention_impl`	返回当前激活的闪光注意力实现的名称（如果存在）。

`flex_attention`	此模块实现了 PyTorch 中 flex_attention 的面向用户的 API。
`bias`	定义与 scaled_dot_product_attention 一起使用的偏置子类
`experimental`
`varlen`	使用 Flash Attention 实现的可变长度注意力。