From 00bd07b82bb1903a611b79ce35bc5fef53b5e2ea Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Thu, 12 Feb 2026 21:33:06 +0800 Subject: [PATCH] [XPU] Update Bagel's flash_attn_varlen_func to fa utils (#1295) Signed-off-by: zhenwei-intel --- vllm_omni/diffusion/attention/backends/utils/fa.py | 5 +++++ vllm_omni/diffusion/models/bagel/bagel_transformer.py | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index 3344c9963..eb8f72bad 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -30,6 +30,11 @@ if current_omni_platform.is_rocm(): from aiter import flash_attn_func, flash_attn_varlen_func # noqa: F401 except (ImportError, ModuleNotFoundError): pass +elif current_omni_platform.is_xpu(): + try: + from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func # noqa: F401 + except (ImportError, ModuleNotFoundError): + pass else: # CUDA: try FA3 -> FA2 fallback chain # Try FA3 from fa3-fwd PyPI package diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 7a476ee73..d2ed26d97 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -30,8 +30,8 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.transformers_utils.configs.bagel import BagelConfig -from vllm.vllm_flash_attn import flash_attn_varlen_func +from vllm_omni.diffusion.attention.backends.utils.fa import flash_attn_varlen_func from vllm_omni.diffusion.layers.rope import RotaryEmbedding