Các backend chú ý

Bộ điều phối chú ý (attention dispatcher) là một tính năng thử nghiệm. Vui lòng mở một issue nếu bạn có bất kỳ phản hồi nào hoặc gặp phải vấn đề gì.

Diffusers cung cấp một số thuật toán attention được tối ưu hóa, hiệu quả hơn về bộ nhớ và tính toán thông qua attention dispatcher của nó. Dispatcher đóng vai trò như một router để quản lý và chuyển đổi giữa các triển khai attention khác nhau, đồng thời cung cấp một giao diện thống nhất để tương tác với chúng.

Tham khảo bảng bên dưới để có cái nhìn tổng quan về các họ attention khả dụng và xem phần Các backend khả dụng để có danh sách đầy đủ hơn.

Hướng dẫn này sẽ cho bạn thấy cách thiết lập và sử dụng các backend chú ý khác nhau.

set_attention_backend

Phương thức set_attention_backend() sẽ lặp qua tất cả các module trong mô hình và thiết lập backend attention phù hợp để sử dụng. Cài đặt backend attention sẽ được giữ nguyên cho đến khi reset_attention_backend() được gọi.

Ví dụ dưới đây minh họa cách bật triển khai _flash_3_hub cho FlashAttention-3 từ thư viện kernels, cho phép bạn sử dụng ngay các kernel tính toán tối ưu từ Hub mà không cần bất kỳ cấu hình nào.

FlashAttention-3 yêu cầu tối thiểu là 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]

Để khôi phục backend chú ý mặc định, hãy gọi reset_attention_backend().

pipeline.transformer.reset_attention_backend()

context manager attention_backend

Trình quản lý ngữ cảnh attention_backend tạm thời đặt một backend attention cho một mô hình trong phạm vi ngữ cảnh đó. Bên ngoài ngữ cảnh, attention mặc định (scaled dot product attention bản địa của PyTorch) sẽ được sử dụng. Điều này hữu ích nếu bạn muốn sử dụng các backend khác nhau cho các phần khác nhau của một pipeline hoặc nếu bạn muốn kiểm tra các backend khác nhau.

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]

Hầu hết các backend attention hỗ trợ torch.compile mà không bị ngắt đồ thị và có thể được sử dụng để tăng tốc suy luận thêm.

Tin tưởng các kernel từ xa

Các backend của Hub và các tính năng khác dựa trên kernel (chẳng hạn như GGUF và Nunchaku Lite) sẽ tải các kernel tính toán từ Hub bằng kernels và thực thi mã của chúng cục bộ.

Mặc định, kernels chỉ tải một kernel khi nhà phát hành của nó là một nhà phát hành kernel được tin cậy trên Hub. Các kernel được phát hành dưới tổ chức kernels-community được tin cậy, vì vậy Diffusers sẽ tải chúng mà không cần cấu hình bổ sung. Các backend attention Hub khác như _flash_3_hub, flash_hub, sage_hub đều trỏ đến các kho kernels-community.

Các kernel từ bất kỳ nhà phát hành nào khác đều không được kiểm duyệt. Việc tải một kernel sẽ tải xuống và chạy mã mà Diffusers không thể đảm bảo, do đó Diffusers giữ tính năng này ở trạng thái tắt trừ khi bạn chủ động bật bằng biến môi trường DIFFUSERS_TRUST_REMOTE_KERNELS. Khi được đặt, Diffusers sẽ truyền trust_remote_code=True đến kernels để nó tải cả các kernel từ các nhà phát hành không đáng tin cậy.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Chỉ bật tính năng này sau khi đã kiểm tra kho kernel, vì nó cấp cho mã đã tải xuống khả năng chạy trên máy của bạn. Nếu không có nó, việc tải một kernel từ nhà phát hành không đáng tin cậy sẽ gây ra lỗi. Diffusers tự thực hiện kiểm tra này, do đó nó cũng áp dụng cho kernels<0.14.0, phiên bản ra đời trước đối số trust_remote_code. Đặt DIFFUSERS_DISABLE_REMOTE_CODE=true sẽ vô hiệu hóa mã từ xa trên toàn cục và có ưu tiên cao hơn DIFFUSERS_TRUST_REMOTE_KERNELS.

Kiểm tra

Bộ điều phối chú ý bao gồm các kiểm tra gỡ lỗi để bắt các lỗi phổ biến trước khi chúng gây ra vấn đề.

1. Kiểm tra thiết bị xác minh rằng các tensor truy vấn, khóa và giá trị nằm trên cùng một thiết bị.
2. Kiểm tra kiểu dữ liệu xác nhận các tensor có cùng dtype và sử dụng bfloat16 hoặc float16.
3. Kiểm tra hình dạng xác thực các chiều của tensor và ngăn chặn việc trộn lẫn các mặt nạ chú ý với các cờ nhân quả.

Bật các kiểm tra này bằng cách đặt biến môi trường DIFFUSERS_ATTN_CHECKS. Các kiểm tra này làm tăng chi phí cho mỗi thao tác attention, do đó chúng bị tắt theo mặc định.

export DIFFUSERS_ATTN_CHECKS=yes

Các kiểm tra được thực hiện trước mỗi thao tác chú ý.

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

Bạn cũng có thể cấu hình trực tiếp registry.

from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True

Các backend khả dụng

Xem bảng bên dưới để có danh sách đầy đủ các backend attention khả dụng và các biến thể của chúng.