torch.nn.attention.sdpa_kernel

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

Менеджер контекста для выбора бэкенда, используемого для внимания на основе масштабированного скалярного произведения.

Предупреждение

Эта функция находится в статусе бета-версии и может измениться.

Параметры:

• backends (Union[List[SDPBackend], SDPBackend]) – Бэкенд или список бэкендов для внимания на основе масштабированного скалярного произведения.
• set_priority (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(...)

# Enable the cuDNN or flash attention backends, and in that order
with sdpa_kernel(
    [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION], set_priority=True
):
    scaled_dot_product_attention(...)

Этот менеджер контекста можно использовать для выбора бэкенда, применяемого для внимания на основе масштабированного скалярного произведения. При выходе из менеджера контекста предыдущее состояние флагов будет восстановлено, что позволит использовать все бэкенды.