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

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

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

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

В этом руководстве показано, как настраивать и использовать различные бэкенды внимания.

set_attention_backend

Метод set_attention_backend() перебирает все модули в модели и устанавливает соответствующий бэкенд внимания для использования. Настройка бэкенда внимания сохраняется до вызова reset_attention_backend().

Пример ниже демонстрирует, как включить реализацию _flash_3_hub для FlashAttention-3 из библиотеки kernels, что позволяет мгновенно использовать оптимизированные вычислительные ядра из Hub без какой-либо настройки.

FlashAttention-3 требует как минимум GPU архитектуры Ampere.

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)
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]

Чтобы восстановить стандартный бэкенд внимания, вызовите reset_attention_backend().

pipeline.transformer.reset_attention_backend()

контекстный менеджер attention_backend

Контекстный менеджер attention_backend временно устанавливает бэкенд внимания для модели в пределах контекста. Вне контекста используется внимание по умолчанию (нативное масштабированное внимание на основе скалярного произведения в PyTorch). Это полезно, если вы хотите использовать разные бэкенды для разных частей конвейера или если вы хотите протестировать различные бэкенды.

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)
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 и другие функции, использующие ядра вычислений (например, GGUF и Nunchaku Lite), загружают вычислительные ядра с Hub с помощью kernels и выполняют их код локально.

По умолчанию kernels загружает ядро только в том случае, если его издатель является доверенным издателем ядер на Hub. Ядра, опубликованные в организации kernels-community, считаются доверенными, поэтому Diffusers загружает их без дополнительной настройки. Бэкенды внимания Hub _flash_3_hub, flash_hub, sage_hub и другие все разрешаются в репозитории kernels-community.

Ядра от любых других издателей не проходят проверку. Загрузка такого ядра приводит к скачиванию и запуску кода, за который Diffusers не может поручиться, поэтому Diffusers держит эту функцию отключённой, пока вы явно не включите её с помощью переменной окружения 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

Проверки теперь выполняются перед каждой операцией внимания.

import torch

query = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")  # or "mps", "xpu", "cpu"
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}")

Вы также можете настроить реестр напрямую.

from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True

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

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