Attention 백엔드

Attention dispatcher는 실험적 기능입니다. 피드백이 있거나 문제가 발생하면 이슈를 열어 주세요.

Diffusers는 attention dispatcher를 통해 메모리와 계산 효율성이 더 높은 여러 최적화된 attention 알고리즘을 제공합니다. dispatcher는 다양한 attention 구현을 관리하고 전환하는 라우터 역할을 하며, 이러한 구현과 상호작용하기 위한 통합된 인터페이스를 제공합니다.

사용 가능한 어텐션 패밀리 개요는 아래 표를, 더 완전한 목록은 사용 가능한 백엔드 섹션을 참조하세요.

이 가이드에서는 다양한 attention backends를 설정하고 사용하는 방법을 보여줍니다.

set_attention_backend

set_attention_backend() 메서드는 모델의 모든 모듈을 순회하며 사용할 적절한 attention backend를 설정합니다. attention backend 설정은 reset_attention_backend()이 호출될 때까지 유지됩니다.

아래 예시는 kernels 라이브러리에서 FlashAttention-3의 _flash_3_hub 구현을 활성화하는 방법을 보여줍니다. 이를 통해 설정 과정 없이 Hub의 최적화된 compute kernel을 즉시 사용할 수 있습니다.

FlashAttention-3는 최소 Ampere GPU가 필요합니다.

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]

기본 attention 백엔드를 복원하려면 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]

대부분의 attention 백엔드는 그래프 브레이크 없이 torch.compile을 지원하며, 이를 통해 추론 속도를 추가로 높일 수 있습니다.

원격 커널에 대한 신뢰

Hub 백엔드 및 기타 커널 기반 기능(예: GGUF, Nunchaku Lite)은 kernels를 통해 Hub에서 컴퓨트 커널을 다운로드하고 해당 코드를 로컬에서 실행합니다.

기본적으로 kernels는 게시자가 Hub의 신뢰할 수 있는 커널 게시자인 경우에만 커널을 로드합니다. kernels-community 조직 아래 게시된 커널은 신뢰할 수 있으므로 Diffusers는 추가 구성 없이 이를 로드합니다. _flash_3_hub, flash_hub, sage_hub 및 기타 Hub 어텐션 백엔드는 모두 kernels-community 저장소로 해석됩니다.

다른 발행자의 Kernels은 검증되지 않습니다. 하나를 로드하면 Diffusers가 보증할 수 없는 코드가 다운로드되고 실행되므로, Diffusers는 DIFFUSERS_TRUST_REMOTE_KERNELS 환경 변수로 명시적으로 선택하지 않는 한 비활성화 상태로 유지합니다. 이 변수가 설정되면 Diffusers는 kernels에 trust_remote_code=True를 전달하여 신뢰할 수 없는 발행자의 Kernels도 로드합니다.

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

커널 저장소를 검사한 후에만 이 옵션을 활성화하세요. 이 옵션은 다운로드된 코드가 사용자의 머신에서 실행될 수 있는 권한을 부여하기 때문입니다. 이 옵션이 없으면 신뢰할 수 없는 게시자로부터 커널을 로드하면 오류가 발생합니다. Diffusers는 이 검사를 자체적으로 수행하므로, trust_remote_code 인자가 도입되기 이전인 kernels<0.14.0에도 적용됩니다. DIFFUSERS_DISABLE_REMOTE_CODE=true를 설정하면 원격 코드가 전역적으로 비활성화되며, DIFFUSERS_TRUST_REMOTE_KERNELS보다 우선합니다.

확인

Attention dispatcher에는 문제가 발생하기 전에 일반적인 오류를 포착하는 디버깅 체크가 포함되어 있습니다.

1. Device checks는 query, key, value 텐서가 동일한 device에 존재하는지 확인합니다.
2. Data type checks는 텐서가 동일한 dtype을 가지고 있는지 확인하며, bfloat16 또는 float16 중 하나를 사용해야 합니다.
3. Shape checks는 텐서 차원을 검증하고 attention mask와 causal flag의 혼합을 방지합니다.

DIFFUSERS_ATTN_CHECKS 환경 변수를 설정하여 이러한 검사를 활성화할 수 있습니다. 검사 기능은 모든 attention 연산에 오버헤드를 추가하므로 기본적으로 비활성화되어 있습니다.

export DIFFUSERS_ATTN_CHECKS=yes

이제 모든 attention 연산 전에 체크가 실행됩니다.

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

사용 가능한 백엔드

사용 가능한 어텐션 백엔드와 그 변형의 전체 목록은 아래 표를 참조하세요.