[XPU] Update Bagel's flash_attn_varlen_func to fa utils (vllm-project#1295)

Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>

Author: zhenwei-intel

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -30,6 +30,11 @@
             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

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -30,8 +30,8 @@
 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