diffusers/main/ru/optimization/attention_backends.md

Бэкенды внимания

Диспетчер внимания — это экспериментальная функция. Пожалуйста, создайте issue, если у вас есть отзывы или вы столкнулись с какими-либо проблемами.

Большинство трансформерных моделей Diffusers направляют внимание через диспетчер внимания, что позволяет переключать оптимизированные бэкенды за единым API. Диспетчер управляет зарегистрированными реализациями и предоставляет для них единый путь вызова. Некоторые модели, например многие автоэнкодеры, не используют диспетчер из-за своей внутренней структуры, и переключение бэкенда не влияет на их слои внимания.

См. таблицу ниже для обзора доступных семейств внимания и раздел Доступные бэкенды для более полного списка. Самый быстрый бэкенд зависит от модели, GPU, dtype и формы входных данных.

семейство attention	основная функция
FlashAttention	минимизирует чтение/запись памяти с помощью тайлинга и пересчета
AI Tensor Engine для ROCm	Реализация FlashAttention, оптимизированная для ускорителей AMD ROCm
SageAttention	квантует внимание до int8
FlexAttention	PyTorch FlexAttention
Нативный PyTorch	встроенная реализация PyTorch с использованием scaled_dot_product_attention
xFormers	memory-efficient attention с поддержкой различных attention-ядер

Бэкенды Hub (имена *_hub) требуют библиотеку Kernels и загружают ядро при первом использовании. Бэкенды Hub для SageAttention также требуют DIFFUSERS_TRUST_REMOTE_KERNELS=true (см. Доверие удалённым ядрам). Остальные бэкенды требуют свои пакеты, например flash-attn или sageattention. В таблице Доступные бэкенды перечислены требования, которые Diffusers проверяет при включении бэкенда.

Установите бэкенд для модели

Метод set_attention_backend() проходит по слоям внимания модели и применяет выбранный бэкенд к каждому из них. Он также устанавливает активный бэкенд диспетчера для всего процесса в то же значение.

reset_attention_backend() очищает бэкенд только на слоях внимания. Он не очищает активный бэкенд для всего процесса. Для временного переключения, которое восстанавливает предыдущий активный бэкенд при выходе, используйте контекстный менеджер attention_backend.

Пример ниже включает _flash_3_hub (FlashAttention-3 из Hub) с device_map="cuda".

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"
)
pipeline.transformer.set_attention_backend("_flash_3_hub")

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""
pipeline(prompt).images[0]

Бэкенды FlashAttention-3, не использующие Hub (_flash_3, _flash_varlen_3), требуют сборки FlashAttention-3 из исходного кода и скоро будут сняты с поддержки. Вместо них используйте _flash_3_hub или _flash_3_varlen_hub.

Попробуйте бэкенд временно

Контекстный менеджер attention_backend() устанавливает активный бэкенд для всего процесса на время выполнения блока и восстанавливает предыдущий бэкенд при выходе из блока. Используйте его, чтобы протестировать бэкенд для одного вызова, не оставляя постоянный бэкенд, установленный через set_attention_backend().

import torch
from diffusers import QwenImagePipeline, attention_backend

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"
)
prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

with attention_backend("_flash_3_hub"):
    image = pipeline(prompt).images[0]

Большинство бэкендов внимания работают с torch.compile. Ускорит ли это ваш пайплайн, зависит от модели и бэкенда. См. Точность и компиляция.

Доверие к удалённым ядрам

Бэкендам Hub сначала требуется библиотека Kernels, после чего Diffusers загружает ядро Hub при первом использовании.

Бэкенды внимания Hub загружают вычислительные ядра с помощью библиотеки Kernels и запускают их локально. Большинство имён внимания Hub (_flash_3_hub, flash_hub и другие бэкенды *_hub) указывают на организацию kernels-community. Эта организация является доверенным издателем в Kernels, поэтому Diffusers загружает эти ядра внимания без установки переменной DIFFUSERS_TRUST_REMOTE_KERNELS.

Бэкенды SageAttention Hub (sage_hub и sage_blackwell_hub) загружаются из организации SageAttention. Для их использования установите DIFFUSERS_TRUST_REMOTE_KERNELS=true.

Другие функции, использующие ядра, такие как GGUF и Nunchaku Lite, могут загружать ядра от издателей, не входящих в kernels-community. Эти пути остаются заблокированными, если вы не разрешите их явно с помощью DIFFUSERS_TRUST_REMOTE_KERNELS. При установке этой переменной Diffusers передаёт trust_remote_code=True в Kernels, что позволяет загружать ядра и от ненадёжных издателей.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Включайте этот параметр только после проверки репозитория ядра. Без него загрузка ядра от ненадёжного издателя приводит к ошибке. Diffusers выполняет эту проверку самостоятельно, поэтому она также распространяется на kernels<0.14.0, которые предшествуют аргументу trust_remote_code. Установка DIFFUSERS_DISABLE_REMOTE_CODE=true глобально отключает удалённый код и имеет приоритет над DIFFUSERS_TRUST_REMOTE_KERNELS.

Проверки

Диспетчер внимания может выполнять отладочные проверки перед каждым вызовом диспетчеризованного внимания. Какие проверки выполняются, зависит от ограничений, зарегистрированных для активного бэкенда.

1. Проверки устройства подтверждают, что тензоры запроса, ключа и значения находятся на одном и том же устройстве.
2. Проверки типов данных, если они зарегистрированы, подтверждают совпадение dtypes и часто требуют bfloat16 или float16.
3. Проверки формы валидируют размеры тензоров и предотвращают смешивание масок внимания с каузальными флагами.

Включите проверки с помощью переменной окружения DIFFUSERS_ATTN_CHECKS. Проверки добавляют накладные расходы, поэтому по умолчанию они отключены.

export DIFFUSERS_ATTN_CHECKS=yes

При включённых проверках Diffusers выполняет эти ограничения перед каждым вызовом диспетчеризованного внимания. В приведённом ниже низкоуровневом примере dispatch_attention_fn вызывается напрямую. Для инференса пайплайна этот импорт не требуется. Достаточно установить бэкенд с помощью set_attention_backend() или attention_backend().

import torch
from diffusers.models.attention_dispatch import attention_backend, dispatch_attention_fn

query = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")
key = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")
value = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")

try:
    with attention_backend("flash"):
        output = dispatch_attention_fn(query, key, value)
        print("✓ Flash Attention works with checks enabled")
except Exception as e:
    print(f"✗ Flash Attention failed: {e}")

Доступные бэкенды

Список доступных бэкендов внимания и их вариантов приведён в таблице ниже. При включении бэкенда Diffusers проверяет наличие пакета и закреплённые версии.

Имя бэкенда	Семейство	Описание	Предварительные требования
native	нативный PyTorch	Бэкенд по умолчанию, использующий scaled_dot_product_attention из PyTorch	None
flex	FlexAttention	PyTorch FlexAttention	torch>=2.5.0
_native_cudnn	нативный PyTorch	Внимание с оптимизацией CuDNN	CUDA + CuDNN
_native_efficient	нативный PyTorch	Внимание с эффективной работой с памятью	Ничего, кроме PyTorch
_native_flash	нативный PyTorch	FlashAttention в PyTorch	CUDA
_native_math	нативный PyTorch	Внимание на основе математики (резервный вариант)	None
_native_npu	нативный PyTorch	Внимание, оптимизированное для NPU	torch_npu
_native_xla	нативный PyTorch	Внимание, оптимизированное XLA	torch_xla>=2.2
flash	FlashAttention	FlashAttention-2	flash-attn>=2.6.3
flash_hub	FlashAttention	FlashAttention-2 из Hub kernels	kernels>=0.12
flash_varlen	FlashAttention	FlashAttention с переменной длиной	flash-attn>=2.6.3
flash_varlen_hub	FlashAttention	FlashAttention переменной длины из ядра Hub	kernels>=0.12
aiter_fa2_hub	AI Tensor Engine для ROCm	FlashAttention-2 для AMD ROCm из Hub kernels (bfloat16)	kernels>=0.12, ROCm
flash_4_hub	FlashAttention	FlashAttention-4 из ядра Hub	kernels>=0.12.3
_flash_3	FlashAttention	FlashAttention-3 (локальный; скоро будет снят с поддержки)	Сборка FA3 из исходного кода
_flash_varlen_3	FlashAttention	FlashAttention-3 переменной длины (локальный; скоро будет снят с поддержки)	Сборка FA3 из исходного кода
_flash_3_hub	FlashAttention	FlashAttention-3 из Hub kernels	kernels>=0.12
_flash_3_varlen_hub	FlashAttention	FlashAttention-3 переменной длины из ядра Hub	kernels>=0.12
sage	SageAttention	Квантованное внимание (INT8 QK)	sageattention>=2.1.1
sage_hub	SageAttention	Квантованное внимание (INT8 QK) из ядра Hub	kernels>=0.12, DIFFUSERS_TRUST_REMOTE_KERNELS=true
sage_blackwell_hub	SageAttention	Внимание SageAttention3 FP4 для GPU SM120 Blackwell из ядра Hub	kernels>=0.12, DIFFUSERS_TRUST_REMOTE_KERNELS=true
sage_varlen	SageAttention	SageAttention переменной длины	sageattention>=2.1.1
_sage_qk_int8_pv_fp8_cuda	SageAttention	INT8 QK + FP8 PV (CUDA)	sageattention>=2.1.1
_sage_qk_int8_pv_fp8_cuda_sm90	SageAttention	INT8 QK + FP8 PV (SM90)	sageattention>=2.1.1; SM90
_sage_qk_int8_pv_fp16_cuda	SageAttention	INT8 QK + FP16 PV (CUDA)	sageattention>=2.1.1
_sage_qk_int8_pv_fp16_triton	SageAttention	INT8 QK + FP16 PV (Triton)	sageattention>=2.1.1
xformers	xFormers	Внимание с эффективной работой с памятью	xformers>=0.0.29