Buckets:

hf-doc-build/doc / diffusers /v0.43.0 /ru /optimization /attention_backends.md
mishig's picture
|
download
raw
14.9 kB

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

Диспетчер внимания — это экспериментальная функция. Пожалуйста, создайте 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.

Проверки

Диспетчер внимания включает проверки отладки, которые выявляют типичные ошибки до того, как они вызовут проблемы.

  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

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

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

Развернуть

Имя бэкенда Семейство Описание
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.