[XPU] Add sage_attn backend (#3785)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

requirements/xpu.txt

@@ -1,2 +1,3 @@
 -r common.txt
 onnxruntime>=1.23.2
+auto-round-lib

vllm_omni/diffusion/attention/backends/sage_attn.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
+
 import torch
 from vllm.logger import init_logger
 
@@ -9,12 +11,38 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionImpl,
     AttentionMetadata,
 )
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
-try:
+_SAGE_ATTN_TENSOR_LAYOUT = os.environ.get("SAGE_ATTN_TENSOR_LAYOUT", "NHD").upper()
+assert _SAGE_ATTN_TENSOR_LAYOUT in ("NHD", "HND"), (
+    f"SAGE_ATTN_TENSOR_LAYOUT must be 'NHD' or 'HND', got '{_SAGE_ATTN_TENSOR_LAYOUT}'"
+)
+
+if current_omni_platform.is_xpu():
+    try:
+        import inspect
+
+        from auto_round_kernel import ARK
+
+        _ark = ARK()
+        xpu_sageattn = _ark.sagev1
+        _sagev1_params = inspect.signature(xpu_sageattn).parameters
+        _sagev1_has_tensor_layout = "tensor_layout" in _sagev1_params
+        _sagev1_scale_param = "sm_scale" if "sm_scale" in _sagev1_params else "scale"
+    except ImportError:
+        logger.warning(
+            "XPU SageAttention (auto_round_kernel.ARK.sagev1) is not available. "
+            "Install auto-round-lib for XPU sage attention support."
+        )
+        xpu_sageattn = None
+        _sagev1_has_tensor_layout = False
+        _sagev1_scale_param = "scale"
+else:
+    try:
         from sageattention import sageattn
-except ImportError:
+    except ImportError:
         logger.warning(
             "SageAttentionBackend is not available. You may install sage-attention"
             " by pip install git+https://github.com/thu-ml/SageAttention.git"
@@ -73,3 +101,50 @@ class SageAttentionImpl(AttentionImpl):
             sm_scale=self.softmax_scale,
         )
         return output
+
+    def forward_xpu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        if xpu_sageattn is None:
+            raise ImportError("XPU SageAttention requires auto-round-lib. Install with: pip install auto-round-lib")
+        orig_dtype = query.dtype
+        q = query.to(torch.float16) if orig_dtype != torch.float16 else query
+        k = key.to(torch.float16) if orig_dtype != torch.float16 else key
+        v = value.to(torch.float16) if orig_dtype != torch.float16 else value
+
+        if _sagev1_has_tensor_layout:
+            if _SAGE_ATTN_TENSOR_LAYOUT == "HND":
+                q = q.transpose(1, 2).contiguous()
+                k = k.transpose(1, 2).contiguous()
+                v = v.transpose(1, 2).contiguous()
+            output = xpu_sageattn(
+                q,
+                k,
+                v,
+                tensor_layout=_SAGE_ATTN_TENSOR_LAYOUT,
+                is_causal=self.causal,
+                **{_sagev1_scale_param: self.softmax_scale},
+            )
+            if _SAGE_ATTN_TENSOR_LAYOUT == "HND":
+                output = output.transpose(1, 2).contiguous()
+        else:
+            # No tensor_layout support: kernel expects HND [B, H, S, D]
+            q = q.transpose(1, 2).contiguous()
+            k = k.transpose(1, 2).contiguous()
+            v = v.transpose(1, 2).contiguous()
+            output = xpu_sageattn(
+                q,
+                k,
+                v,
+                is_causal=self.causal,
+                **{_sagev1_scale_param: self.softmax_scale},
+            )
+            output = output.transpose(1, 2).contiguous()
+
+        if orig_dtype != torch.float16:
+            output = output.to(orig_dtype)
+        return output