Buckets:
Бэкенды внимания
Диспетчер внимания — это экспериментальная функция. Пожалуйста, создайте issue, если у вас есть отзывы или вы столкнулись с какими-либо проблемами.
Diffusers предоставляет несколько оптимизированных алгоритмов внимания, которые более эффективны по памяти и вычислениям, через свой attention dispatcher. Диспетчер выполняет роль маршрутизатора для управления и переключения между различными реализациями внимания и предоставляет единый интерфейс для взаимодействия с ними.
См. таблицу ниже для обзора доступных семейств внимания и раздел Доступные бэкенды для более полного списка.
| семейство attention | основная функция |
|---|---|
| FlashAttention | минимизирует чтение/запись памяти с помощью тайлинга и пересчета |
| AI Tensor Engine для ROCm | Реализация FlashAttention, оптимизированная для ускорителей AMD ROCm |
| SageAttention | квантует внимание до int8 |
| Нативный PyTorch | встроенная реализация PyTorch с использованием scaled_dot_product_attention |
| xFormers | memory-efficient attention с поддержкой различных attention-ядер |
В этом руководстве показано, как настраивать и использовать различные бэкенды внимания.
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.
Проверки
Диспетчер внимания включает проверки отладки, которые выявляют типичные ошибки до того, как они вызовут проблемы.
- Проверки устройства подтверждают, что тензоры запроса, ключа и значения находятся на одном и том же устройстве.
- Проверки типа данных подтверждают, что тензоры имеют совпадающие dtypes и используют либо bfloat16, либо float16.
- Проверки формы валидируют размеры тензоров и предотвращают смешивание масок внимания с каузальными флагами.
Включите эти проверки, задав переменную окружения 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
Доступные бэкенды
Список доступных бэкендов внимания и их вариантов приведён в таблице ниже.
Развернуть
| Имя бэкенда | Семейство | Описание |
|---|---|---|
native |
нативный PyTorch | Бэкенд по умолчанию, использующий scaled_dot_product_attention из PyTorch |
flex |
FlexAttention | Реализация PyTorch FlexAttention |
_native_cudnn |
нативный PyTorch | Внимание с оптимизацией CuDNN |
_native_efficient |
нативный PyTorch | Внимание с эффективной работой с памятью |
_native_flash |
нативный PyTorch | FlashAttention в PyTorch |
_native_math |
нативный PyTorch | Внимание на основе математики (резервный вариант) |
_native_npu |
нативный PyTorch | Внимание, оптимизированное для NPU |
_native_xla |
нативный PyTorch | Внимание, оптимизированное XLA |
flash |
FlashAttention | FlashAttention-2 |
flash_hub |
FlashAttention | FlashAttention-2 из ядер |
flash_varlen |
FlashAttention | FlashAttention с переменной длиной |
flash_varlen_hub |
FlashAttention | FlashAttention переменной длины из ядер |
aiter_fa2_hub |
AI Tensor Engine для ROCm | FlashAttention-2 для AMD ROCm из kernels |
aiter_sage_hub |
AI Tensor Engine для ROCm | Квантованное внимание (INT8 QK, FP8 PV) для AMD RDNA4 (gfx1201) из kernels |
flash_4_hub |
FlashAttention | FlashAttention-4 |
_flash_3 |
FlashAttention | FlashAttention-3 |
_flash_varlen_3 |
FlashAttention | FlashAttention-3 с переменной длиной |
_flash_3_hub |
FlashAttention | FlashAttention-3 из kernels |
_flash_3_varlen_hub |
FlashAttention | FlashAttention-3 переменной длины из kernels |
sage |
SageAttention | Квантованное внимание (INT8 QK) |
sage_hub |
SageAttention | Квантованное внимание (INT8 QK) из ядер |
sage_blackwell_hub |
SageAttention | SageAttention3 FP4-внимание для SM120 Blackwell GPU из kernels |
sage_varlen |
SageAttention | SageAttention переменной длины |
_sage_qk_int8_pv_fp8_cuda |
SageAttention | INT8 QK + FP8 PV (CUDA) |
_sage_qk_int8_pv_fp8_cuda_sm90 |
SageAttention | INT8 QK + FP8 PV (SM90) |
_sage_qk_int8_pv_fp16_cuda |
SageAttention | INT8 QK + FP16 PV (CUDA) |
_sage_qk_int8_pv_fp16_triton |
SageAttention | INT8 QK + FP16 PV (Triton) |
xformers |
xFormers | Внимание с эффективной работой с памятью |
Xet Storage Details
- Size:
- 14.9 kB
- Xet hash:
- 5116f87d6050959404365b0e8a982e6c97fd38bd4c1ab839b8c10ae171f34b8b
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.