90 lines
3.5 KiB
Python
90 lines
3.5 KiB
Python
"""Custom attention processors that avoid SDPA on ROCm."""
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from diffusers.models.attention_processor import Attention, AttnProcessor
|
|
|
|
|
|
def simple_attention_forward(
|
|
attn: Attention,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | None = None,
|
|
attention_mask: torch.Tensor | None = None,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
"""Simple QKV attention using raw matmul, bypassing SDPA dispatch."""
|
|
batch_size, seq_len, _ = hidden_states.shape
|
|
inner_dim = attn.inner_dim if hasattr(attn, 'inner_dim') else hidden_states.shape[-1]
|
|
|
|
query = attn.to_q(hidden_states)
|
|
key = attn.to_k(hidden_states if encoder_hidden_states is None else encoder_hidden_states)
|
|
value = attn.to_v(hidden_states if encoder_hidden_states is None else encoder_hidden_states)
|
|
|
|
head_dim = inner_dim // attn.heads
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
scale = head_dim ** -0.5
|
|
attn_weights = query @ key.transpose(-2, -1) * scale
|
|
attn_weights = F.softmax(attn_weights, dim=-1).to(query.dtype)
|
|
hidden_states = attn_weights @ value
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
|
|
hidden_states = attn.to_out[0](hidden_states)
|
|
if len(attn.to_out) > 1:
|
|
hidden_states = attn.to_out[1](hidden_states)
|
|
return hidden_states
|
|
|
|
|
|
def _matmul_attention(query, key, value, attn_mask, dropout_p, is_causal, backend, parallel_config):
|
|
"""Simple matmul attention replacing dispatch_attention_fn. Returns 4D tensor."""
|
|
batch_size, seq_len, num_heads, head_dim = query.shape
|
|
scale = head_dim ** -0.5
|
|
q = query.transpose(1, 2)
|
|
k = key.transpose(1, 2)
|
|
v = value.transpose(1, 2)
|
|
w = q @ k.transpose(-2, -1) * scale
|
|
if attn_mask is not None:
|
|
w = w + attn_mask
|
|
w = F.softmax(w, dim=-1).to(q.dtype)
|
|
out = w @ v
|
|
return out.transpose(1, 2)
|
|
|
|
|
|
def patch_unet_attention(unet):
|
|
"""Replace all attention processors with simple matmul-based attention."""
|
|
class SimpleAttnProcessor(AttnProcessor):
|
|
def __call__(
|
|
self,
|
|
attn: Attention,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | None = None,
|
|
attention_mask: torch.Tensor | None = None,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
return simple_attention_forward(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
|
|
|
|
unet.set_attn_processor(SimpleAttnProcessor())
|
|
|
|
|
|
def patch_qwen_transformer(transformer):
|
|
"""Patch Qwen transformer attention.
|
|
|
|
If the ROCm experimental flash attention env var is set
|
|
(TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1), the default SDPA dispatch
|
|
will use flash attention — no patch needed.
|
|
|
|
Otherwise, fall back to a simple matmul-based attention that avoids
|
|
the unstable SDPA math backend on AMD GPUs.
|
|
"""
|
|
import os
|
|
if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL") == "1":
|
|
return # flash attention available, no patch needed
|
|
|
|
# Monkey-patch the dispatch function with matmul fallback
|
|
import diffusers.models.attention_dispatch as ad
|
|
ad.dispatch_attention_fn = _matmul_attention
|
|
import diffusers.models.transformers.transformer_qwenimage as tq
|
|
tq.dispatch_attention_fn = _matmul_attention
|