torch.nn.attention.sdpa_kernel

torch.nn.attention.sdpa_kernel(backends, set_priority=False)

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

警告

此函数处于测试阶段，可能会发生变化。

参数:

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

• set_priority_order (python:bool=False) – 是否将后端排序解释为优先级顺序。

示例：

from torch.nn.functional import scaled_dot_product_attention
from torch.nn.attention import SDPBackend, sdpa_kernel
# Only enable flash attention backend
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    scaled_dot_product_attention(...)

# Enable the Math or Efficient attention backends
with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):
    scaled_dot_product_attention(...)

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