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

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

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

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

Бэкенды 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 проверяет наличие пакета и закреплённые версии.