快捷方式

torch.nn.attention.sdpa_kernel

torch.nn.attention.sdpa_kernel(backends)[源代码]

上下文管理器,用于选择要用于缩放点积注意力(scaled dot product attention)的后端。

警告

此函数处于 Beta 阶段,可能会发生变化。

参数

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

示例

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(...)

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源