torch.nn.attention.activate_flash_attention_impl

torch.nn.attention.activate_flash_attention_impl(impl) [исходный код]

Активирует в диспетчере ранее зарегистрированную реализацию Flash Attention.

Примечание

Поставщикам бэкендов НЕ следует автоматически активировать свою реализацию при импорте. Пользователям следует явно включить её, вызвав эту функцию или задав переменные среды, чтобы обеспечить совместную работу нескольких библиотек-поставщиков.

Параметры:

impl (str | Literal['FA3', 'FA4']) – Идентификатор реализации для активации. Список доступных реализаций см. в list_flash_attention_impls(). Если вызываемый объект register_flash_attention_impl() бэкенда возвращает FlashAttentionHandle, реестр хранит эту ссылку до конца процесса (пока не появится явная поддержка удаления).

Пример

>>> activate_flash_attention_impl("FA4")