Backend perhatian

Dispatcher perhatian adalah fitur eksperimental. Silakan buka masalah jika Anda memiliki umpan balik atau mengalami masalah.

Diffusers menyediakan beberapa algoritma perhatian yang dioptimalkan, yang lebih hemat memori dan komputasi, melalui attention dispatcher-nya. Dispatcher ini berfungsi sebagai router untuk mengelola dan beralih antara berbagai implementasi perhatian, serta menyediakan antarmuka terpadu untuk berinteraksi dengan implementasi tersebut.

Lihat tabel di bawah ini untuk ringkasan keluarga perhatian yang tersedia dan bagian Backend yang tersedia untuk daftar yang lebih lengkap.

Panduan ini akan menunjukkan cara mengatur dan menggunakan berbagai backend perhatian.

set_attention_backend

Metode set_attention_backend() mengiterasi semua modul dalam model dan mengatur backend perhatian yang sesuai untuk digunakan. Pengaturan backend perhatian ini tetap berlaku hingga reset_attention_backend() dipanggil.

Contoh di bawah ini menunjukkan cara mengaktifkan implementasi _flash_3_hub untuk FlashAttention-3 dari pustaka kernels, yang memungkinkan Anda menggunakan kernel komputasi yang dioptimalkan dari Hub secara instan tanpa memerlukan konfigurasi apa pun.

FlashAttention-3 membutuhkan GPU Ampere sebagai minimum.

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]

Untuk memulihkan backend perhatian default, panggil reset_attention_backend().

pipeline.transformer.reset_attention_backend()

context manager attention_backend

Context manager attention_backend secara sementara mengatur backend attention untuk sebuah model di dalam konteks tersebut. Di luar konteks, attention default (scaled dot product attention native dari PyTorch) akan digunakan. Hal ini berguna jika Anda ingin menggunakan backend yang berbeda untuk bagian-bagian berbeda dari sebuah pipeline atau jika Anda ingin menguji berbagai backend.

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]

Sebagian besar backend perhatian mendukung torch.compile tanpa pemutusan graf dan dapat digunakan untuk mempercepat inferensi lebih lanjut.

Memercayai kernel jarak jauh

Backend Hub dan fitur lain yang didukung kernel (seperti GGUF dan Nunchaku Lite) mengunduh kernel komputasi dari Hub dengan kernels dan mengeksekusi kodenya secara lokal.

Secara default, kernels hanya memuat kernel ketika penerbitnya adalah penerbit kernel yang tepercaya di Hub. Kernel yang diterbitkan di bawah organisasi kernels-community tepercaya, sehingga Diffusers memuatnya tanpa konfigurasi tambahan. _flash_3_hub, flash_hub, sage_hub, dan backend perhatian Hub lainnya semuanya merujuk ke repositori kernels-community.

Kernel dari penerbit lain tidak melalui proses verifikasi. Memuat salah satunya akan mengunduh dan menjalankan kode yang tidak dapat dijamin keamanannya oleh Diffusers, sehingga Diffusers tetap menonaktifkannya kecuali Anda secara eksplisit mengaktifkannya dengan variabel lingkungan DIFFUSERS_TRUST_REMOTE_KERNELS. Jika disetel, Diffusers meneruskan trust_remote_code=True ke kernels sehingga kernel dari penerbit yang tidak tepercaya juga dimuat.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Hanya aktifkan ini setelah memeriksa repositori kernel, karena hal ini memberikan kemampuan pada kode yang diunduh untuk dijalankan di mesin Anda. Tanpa ini, memuat kernel dari penerbit yang tidak tepercaya akan memicu error. Diffusers melakukan pemeriksaan ini sendiri, sehingga juga berlaku untuk kernels<0.14.0, yang mendahului argumen trust_remote_code. Mengatur DIFFUSERS_DISABLE_REMOTE_CODE=true akan menonaktifkan kode jarak jauh secara global dan memiliki prioritas lebih tinggi daripada DIFFUSERS_TRUST_REMOTE_KERNELS.

Pemeriksaan

Dispatcher perhatian mencakup pemeriksaan debug yang menangkap kesalahan umum sebelum menyebabkan masalah.

1. Periksa perangkat memverifikasi bahwa tensor query, key, dan value berada di perangkat yang sama.
2. Pemeriksaan tipe data memastikan tensor memiliki dtype yang cocok dan menggunakan bfloat16 atau float16.
3. Pemeriksaan bentuk memvalidasi dimensi tensor dan mencegah pencampuran masker perhatian dengan bendera kausal.

Aktifkan pemeriksaan ini dengan mengatur variabel lingkungan DIFFUSERS_ATTN_CHECKS. Pemeriksaan menambah overhead pada setiap operasi perhatian, sehingga dinonaktifkan secara default.

export DIFFUSERS_ATTN_CHECKS=yes

Pemeriksaan sekarang dijalankan sebelum setiap operasi perhatian.

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}")

Anda juga dapat mengatur registry secara langsung.

from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True

Backend yang tersedia

Lihat tabel di bawah ini untuk daftar lengkap backend perhatian yang tersedia beserta variasinya.