Spaces:
Runtime error
Runtime error
File size: 9,055 Bytes
52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d 4504378 52ec15d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 | """
Paired with a good language model. Thanks!
FA3 is currently broken on Blackwell (sm_100) GPUs; this module detects that
at import time and falls back to PyTorch scaled-dot-product attention (SDPA)
automatically. The public class name / call signature are unchanged.
"""
import torch
import torch.nn.functional as F
from typing import Optional, Tuple
from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
# ---------------------------------------------------------------------------
# FA3 availability check
# ---------------------------------------------------------------------------
def _is_blackwell() -> bool:
"""Return True when the current default CUDA device is an sm_100 (Blackwell) GPU."""
if not torch.cuda.is_available():
return False
cap = torch.cuda.get_device_capability()
# Blackwell → compute capability 10.x (sm_100)
return cap[0] >= 10
_fa3_available: bool = False
_fa3_unavailable_reason: str = ""
_flash_attn_func = None
if _is_blackwell():
_fa3_unavailable_reason = (
"FlashAttention-3 is not yet supported on Blackwell (sm_100) GPUs. "
"Falling back to scaled-dot-product attention (SDPA)."
)
else:
try:
from kernels import get_kernel
_k = get_kernel("kernels-community/vllm-flash-attn3")
_flash_attn_func = _k.flash_attn_func
_fa3_available = True
except Exception as e:
_fa3_unavailable_reason = (
"FlashAttention-3 via Hugging Face `kernels` is unavailable. "
f"Tried `get_kernel('kernels-community/vllm-flash-attn3')` and failed with:\n{e}\n"
"Falling back to scaled-dot-product attention (SDPA)."
)
# ---------------------------------------------------------------------------
# FA3 custom op (registered only when the kernel loaded successfully)
# ---------------------------------------------------------------------------
if _fa3_available:
@torch.library.custom_op("flash::flash_attn_func", mutates_args=())
def flash_attn_func(
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False
) -> torch.Tensor:
# _flash_attn_func returns (output, softmax_lse); we only need output.
output, _lse = _flash_attn_func(q, k, v, causal=causal)
return output
@flash_attn_func.register_fake
def _flash_attn_func_fake(q, k, v, causal=False):
# output shape mirrors q: (batch, seq_len, num_heads, head_dim)
return torch.empty_like(q).contiguous()
else:
# Provide a stub so call-sites that import the symbol don't break at
# module load; the processor will route around it at runtime.
def flash_attn_func(
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False
) -> torch.Tensor:
raise RuntimeError(_fa3_unavailable_reason)
# ---------------------------------------------------------------------------
# SDPA fallback helper
# ---------------------------------------------------------------------------
def _sdpa_attention(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
causal: bool = False,
) -> torch.Tensor:
"""
Scaled dot-product attention using torch.nn.functional.scaled_dot_product_attention.
Input / output layout: (B, S, H, D_h) — same as the FA3 kernel.
"""
# SDPA expects (B, H, S, D_h)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
out = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
# Back to (B, S, H, D_h)
return out.transpose(1, 2)
# ---------------------------------------------------------------------------
# Attention processor
# ---------------------------------------------------------------------------
class QwenDoubleStreamAttnProcessorFA3:
"""
Attention processor for the Qwen double-stream architecture.
Preferred backend: vLLM FlashAttention-3 via Hugging Face ``kernels``.
Automatic fallback: PyTorch ``scaled_dot_product_attention`` (SDPA) when
FA3 is unavailable — e.g. on Blackwell (sm_100) GPUs where FA3 is not yet
supported, or when the ``kernels`` package is absent.
Notes / limitations
-------------------
- Arbitrary attention masks are not supported on the FA3 path. Pass
``attention_mask=None`` (the default) to stay on the fast path.
- On the SDPA path, ``attention_mask`` is likewise ignored; add explicit
support here if you need it.
- ``encoder_hidden_states`` (text stream) is required.
"""
_attention_backend: str # set in __init__ after capability detection
def __init__(self):
if _fa3_available:
self._attention_backend = "fa3"
else:
import warnings
warnings.warn(
f"QwenDoubleStreamAttnProcessorFA3: {_fa3_unavailable_reason}",
stacklevel=2,
)
self._attention_backend = "sdpa"
def _attend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
causal: bool = False,
) -> torch.Tensor:
"""Dispatch to FA3 or SDPA depending on what is available."""
if self._attention_backend == "fa3":
return flash_attn_func(q, k, v, causal=causal)
return _sdpa_attention(q, k, v, causal=causal)
@torch.no_grad()
def __call__(
self,
attn,
hidden_states: torch.FloatTensor, # (B, S_img, D_model)
encoder_hidden_states: torch.FloatTensor = None, # (B, S_txt, D_model)
encoder_hidden_states_mask: torch.FloatTensor = None, # unused
attention_mask: Optional[torch.FloatTensor] = None, # unsupported on FA3 path
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
if encoder_hidden_states is None:
raise ValueError(
"QwenDoubleStreamAttnProcessorFA3 requires encoder_hidden_states (text stream)."
)
if attention_mask is not None and self._attention_backend == "fa3":
raise NotImplementedError(
"attention_mask is not supported on the FA3 path. "
"Either drop the mask or let the processor fall back to SDPA."
)
B, S_img, _ = hidden_states.shape
S_txt = encoder_hidden_states.shape[1]
# ---- QKV projections ----
img_q = attn.to_q(hidden_states)
img_k = attn.to_k(hidden_states)
img_v = attn.to_v(hidden_states)
txt_q = attn.add_q_proj(encoder_hidden_states)
txt_k = attn.add_k_proj(encoder_hidden_states)
txt_v = attn.add_v_proj(encoder_hidden_states)
# ---- Reshape to (B, S, H, D_h) ----
H = attn.heads
img_q = img_q.unflatten(-1, (H, -1))
img_k = img_k.unflatten(-1, (H, -1))
img_v = img_v.unflatten(-1, (H, -1))
txt_q = txt_q.unflatten(-1, (H, -1))
txt_k = txt_k.unflatten(-1, (H, -1))
txt_v = txt_v.unflatten(-1, (H, -1))
# ---- Q/K normalization ----
if getattr(attn, "norm_q", None) is not None:
img_q = attn.norm_q(img_q)
if getattr(attn, "norm_k", None) is not None:
img_k = attn.norm_k(img_k)
if getattr(attn, "norm_added_q", None) is not None:
txt_q = attn.norm_added_q(txt_q)
if getattr(attn, "norm_added_k", None) is not None:
txt_k = attn.norm_added_k(txt_k)
# ---- RoPE (Qwen variant) ----
if image_rotary_emb is not None:
img_freqs, txt_freqs = image_rotary_emb
img_q = apply_rotary_emb_qwen(img_q, img_freqs, use_real=False)
img_k = apply_rotary_emb_qwen(img_k, img_freqs, use_real=False)
txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs, use_real=False)
txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs, use_real=False)
# ---- Joint attention over [text, image] along sequence axis ----
q = torch.cat([txt_q, img_q], dim=1) # (B, S_txt + S_img, H, D_h)
k = torch.cat([txt_k, img_k], dim=1)
v = torch.cat([txt_v, img_v], dim=1)
out = self._attend(q, k, v, causal=False) # (B, S_total, H, D_h)
# ---- Back to (B, S, D_model) ----
out = out.flatten(2, 3).to(q.dtype)
# ---- Split text / image segments ----
txt_attn_out = out[:, :S_txt, :]
img_attn_out = out[:, S_txt:, :]
# ---- Output projections ----
img_attn_out = attn.to_out[0](img_attn_out)
if len(attn.to_out) > 1:
img_attn_out = attn.to_out[1](img_attn_out) # dropout if present
txt_attn_out = attn.to_add_out(txt_attn_out)
return img_attn_out, txt_attn_out |