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(...)
此上下文管理器可用于选择要用于缩放点积注意力的后端。退出上下文管理器后,将恢复标志的先前状态,从而启用所有后端。