torch.nn.attention.sdpa_kernel

torch.nn.attention.sdpa_kernel(backends)

上下文管理器，用于选择用于缩放点积注意力的后端。

警告

此功能为测试版，可能会发生变化。

Parameters

后端 (联合[列表[SDPBackend], SDPBackend]) – 用于缩放点积注意力的后端或后端列表。

示例：

from torch.nn.functional import scaled_dot_product_attention
from torch.nn.attention import SDPBackend, sdpa_kernel
# 仅启用flash attention后端
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    scaled_dot_product_attention(...)

# 启用Math或Efficient attention后端
with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):
    scaled_dot_product_attention(...)

此上下文管理器可用于选择用于缩放点积注意力的后端。退出上下文管理器时，标志的先前状态将被恢复，启用所有后端。