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ì.

Hầu hết các mô hình transformer của Diffusers định tuyến chú ý (attention) thông qua một attention dispatcher để bạn có thể chuyển đổi các backend tối ưu hóa phía sau một API duy nhất. Dispatcher quản lý các cài đặt đã đăng ký và cung cấp một đường gọi thống nhất cho chúng. Một số mô hình, chẳng hạn như nhiều autoencoder, không sử dụng dispatcher do cấu trúc bên trong của chúng, và việc chuyển đổi backend sẽ không có tác dụng đến các lớp chú ý của 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à phần Các backend khả dụng để có danh sách đầy đủ hơn. Backend nhanh nhất phụ thuộc vào model, GPU, dtype và hình dạng đầu vào.

Các backend Hub (các tên *_hub) cần thư viện Kernels và sẽ tải kernel về khi sử dụng lần đầu. Các backend Hub của SageAttention cũng cần DIFFUSERS_TRUST_REMOTE_KERNELS=true (xem Trusting remote kernels). Các backend khác cần gói riêng của chúng, chẳng hạn như flash-attn hoặc sageattention. Bảng Available backends liệt kê các yêu cầu mà Diffusers kiểm tra khi bạn bật một backend.

Đặt backend trên mô hình

Phương thức set_attention_backend() sẽ duyệt qua các lớp attention của mô hình và áp dụng backend đã chọn cho từng lớp. Nó cũng đặt backend hoạt động trên toàn bộ tiến trình của dispatcher thành cùng một giá trị đó.

reset_attention_backend() chỉ xóa backend trên các lớp attention. Nó không xóa backend đang hoạt động trên toàn bộ quy trình. Để chuyển đổi tạm thời và khôi phục lại backend đang hoạt động trước đó khi thoát, hãy sử dụng context manager attention_backend.

Ví dụ bên dưới bật _flash_3_hub (FlashAttention-3 từ Hub) với device_map="cuda".

import torch
from diffusers import QwenImagePipeline

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

Các backend FlashAttention-3 không thuộc Hub (_flash_3, _flash_varlen_3) yêu cầu phải build FlashAttention-3 từ mã nguồn và sẽ sớm bị ngừng hỗ trợ. Hãy dùng _flash_3_hub hoặc _flash_3_varlen_hub thay thế.

Thử một backend tạm thời

Trình quản lý ngữ cảnh attention_backend() đặt backend hoạt động cho toàn bộ tiến trình trong suốt thời gian khối thực thi và khôi phục backend trước đó khi khối kết thúc. Hãy sử dụng nó để thử một backend cho một lần gọi mà không để lại backend được áp dụng vĩnh viễn từ set_attention_backend().

import torch
from diffusers import QwenImagePipeline, attention_backend

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"
)
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 đều hoạt động với torch.compile. Việc điều này có làm tăng tốc pipeline của bạn hay không phụ thuộc vào model và backend. Xem Precision and compilation.

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

Các backend Hub cần thư viện Kernels trước, sau đó Diffusers sẽ tải kernel Hub từ Hub khi lần đầu sử dụng.

Các backend attention trên Hub tải các kernel tính toán bằng thư viện Kernels và chạy chúng cục bộ. Hầu hết các tên attention trên Hub (_flash_3_hub, flash_hub và các backend *_hub khác) được ánh xạ đến tổ chức kernels-community. Tổ chức này là một nhà xuất bản đáng tin cậy trong Kernels, do đó Diffusers tải các kernel attention đó mà không cần đặt DIFFUSERS_TRUST_REMOTE_KERNELS.

Các backend của SageAttention Hub (sage_hub và sage_blackwell_hub) được tải từ tổ chức SageAttention. Đặt DIFFUSERS_TRUST_REMOTE_KERNELS=true để sử dụng chúng.

Các tính năng khác dựa trên kernel như GGUF và Nunchaku Lite có thể lấy kernel từ các nhà phát hành bên ngoài kernels-community. Các đường dẫn này vẫn bị chặn trừ khi bạn chọn tham gia bằng DIFFUSERS_TRUST_REMOTE_KERNELS. Khi được đặt, Diffusers sẽ chuyển trust_remote_code=True sang Kernels để các nhà phát hành không đáng tin cậy cũng có thể tải lên.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

Chỉ bật tính năng này sau khi đã kiểm tra kho kernel. Nếu không bật, việc tải 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 attention có thể chạy các kiểm tra gỡ lỗi trước mỗi lần gọi attention được điều phối. Các kiểm tra được thực thi phụ thuộc vào các ràng buộc đã đăng ký cho backend đang hoạt động.

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. Các kiểm tra kiểu dữ liệu, khi được đăng ký, xác nhận sự khớp nhau của các dtype và thường yêu cầu 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 bằng biến môi trường DIFFUSERS_ATTN_CHECKS. Các kiểm tra làm tăng chi phí xử lý, do đó chúng bị tắt theo mặc định.

export DIFFUSERS_ATTN_CHECKS=yes

Khi bật các kiểm tra, Diffusers sẽ chạy các ràng buộc đó trước mỗi lần gọi attention được dispatch. Ví dụ cấp thấp bên dưới gọi trực tiếp dispatch_attention_fn. Suy luận pipeline không cần import đó. Nó chỉ cần đặt backend thông qua set_attention_backend() hoặc attention_backend().

import torch
from diffusers.models.attention_dispatch import attention_backend, dispatch_attention_fn

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

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. Diffusers sẽ kiểm tra tính khả dụng của gói và các phiên bản được cố định khi bạn bật một backend.