torch.nn.attention.sdpa_kernel

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

用於選擇要用於縮放點積注意力 (scaled dot product attention) 的後端的內容管理器。

警告

此功能為 Beta 版，可能會變更。

參數

• backends (Union[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(...)

這個上下文管理器可以用來選擇用於縮放點積注意力機制的後端。 退出上下文管理器後，標誌的先前狀態將被恢復，從而啟用所有後端。