Fix FlashAttention-3 kernel version lookup

#3
by evalstate HF Staff - opened
Files changed (1) hide show
  1. qwenimage/qwen_fa3_processor.py +3 -3
qwenimage/qwen_fa3_processor.py CHANGED
@@ -8,7 +8,7 @@ from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb
8
 
9
  try:
10
  from kernels import get_kernel
11
- _k = get_kernel("kernels-community/vllm-flash-attn3")
12
  _flash_attn_func = _k.flash_attn_func
13
  except Exception as e:
14
  _flash_attn_func = None
@@ -19,7 +19,7 @@ def _ensure_fa3_available():
19
  if _flash_attn_func is None:
20
  raise ImportError(
21
  "FlashAttention-3 via Hugging Face `kernels` is required. "
22
- "Tried `get_kernel('kernels-community/vllm-flash-attn3')` and failed with:\n"
23
  f"{_kernels_err}"
24
  )
25
 
@@ -139,4 +139,4 @@ class QwenDoubleStreamAttnProcessorFA3:
139
 
140
  txt_attn_out = attn.to_add_out(txt_attn_out)
141
 
142
- return img_attn_out, txt_attn_out
 
8
 
9
  try:
10
  from kernels import get_kernel
11
+ _k = get_kernel("kernels-community/vllm-flash-attn3", version=1)
12
  _flash_attn_func = _k.flash_attn_func
13
  except Exception as e:
14
  _flash_attn_func = None
 
19
  if _flash_attn_func is None:
20
  raise ImportError(
21
  "FlashAttention-3 via Hugging Face `kernels` is required. "
22
+ "Tried `get_kernel('kernels-community/vllm-flash-attn3', version=1)` and failed with:\n"
23
  f"{_kernels_err}"
24
  )
25
 
 
139
 
140
  txt_attn_out = attn.to_add_out(txt_attn_out)
141
 
142
+ return img_attn_out, txt_attn_out