[Intel GPU][Encoder] Add xpu_attn backend for encoder vision attention (#26460)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

python/sglang/srt/layers/attention/vision.py

@@ -28,6 +28,7 @@ from sglang.srt.utils import (
     is_npu,
     is_xpu,
     print_info_once,
+    use_intel_xpu_backend,
 )
 from sglang.srt.utils.multi_stream_utils import (
     maybe_execute_in_parallel,
@@ -57,6 +58,8 @@ if _is_musa:
 
 if _is_npu:
     import torch_npu
+if _is_xpu:
+    from sgl_kernel.flash_attn import flash_attn_varlen_func
 
 from sglang.srt.distributed import (
     split_tensor_along_last_dim,
@@ -105,9 +108,11 @@ FLASHINFER_MAX_SEQLEN_BUCKETS = [
 @dataclasses.dataclass
 class SingletonCache:
     data: Any = None
+    _max_seqlen: Optional[int] = None
 
     def set_data(self, value: Any) -> None:
         self.data = value
+        self._max_seqlen = None
 
     def get_data(self) -> Optional[Any]:
         return self.data
@@ -154,6 +159,21 @@ def resolve_seqlens(
     return resolved_seqlens
 
 
+def resolve_max_seqlen(source, cu_seqlens: torch.Tensor) -> int:
+    """Return max segment length, caching it on a stable carrier so the
+    device->host sync (.item()) happens once per forward instead of once per layer.
+    """
+    if isinstance(source, SingletonCache) or isinstance(source, torch.Tensor):
+        cached = getattr(source, "_max_seqlen", None)
+        if cached is None:
+            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            cached = int(seq_lens.max().item())
+            source._max_seqlen = cached
+        return cached
+    seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+    return int(seq_lens.max().item())
+
+
 class VisionSdpaAttention(nn.Module):
     r"""
     Scaled Dot Product Attention inner product
@@ -791,6 +811,55 @@ class VisionAMXAttention(nn.Module):
         return output
 
 
+class VisionIntelXPUAttention(nn.Module):
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        if not (_is_xpu):
+            raise Exception("VisionIntelXPUAttention is only available for Intel XPU")
+        super().__init__()
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens: torch.Tensor | SingletonCache | None,
+        bsz: int,
+        seq_len: int,
+        softmax_scale: Optional[float] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        r"""
+        Args:
+            cu_seqlens: [b]
+        Returns:
+             [b * s, h, head_size]
+        """
+        window_size = kwargs.get("window_size", (-1, -1))
+        s_aux = kwargs.get("s_aux", None)
+
+        cu_seqlens_source = cu_seqlens
+        cu_seqlens = resolve_seqlens(cu_seqlens_source, bsz, seq_len, device=q.device)
+        cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
+        max_seqlen = resolve_max_seqlen(cu_seqlens_source, cu_seqlens)
+
+        fa_kwargs = dict(
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            softmax_scale=softmax_scale,
+            window_size=window_size,
+        )
+        if s_aux is not None:
+            fa_kwargs["sinks"] = s_aux
+        output = flash_attn_varlen_func(q, k, v, **fa_kwargs)
+
+        return output
+
+
 QKV_BACKEND_IMPL = {
     "triton_attn": VisionTritonAttention,
     "sdpa": VisionSdpaAttention,
@@ -800,6 +869,7 @@ QKV_BACKEND_IMPL = {
     "ascend_attn": VisionAscendAttention,
     "aiter_attn": VisionAiterAttention,
     "amx_attn": VisionAMXAttention,
+    "xpu_attn": VisionIntelXPUAttention,
 }
 
 
@@ -1030,7 +1100,7 @@ class VisionAttention(nn.Module):
         elif _is_cpu and _is_cpu_amx_available:
             backend = "amx_attn"
         elif _is_xpu:
-            backend = "triton_attn"
+            backend = "triton_attn" if not use_intel_xpu_backend() else "xpu_attn"
         else:
             backend = "sdpa"
         if backend == "fa3" and is_blackwell_supported():

python/sglang/srt/server_args.py

@@ -5710,6 +5710,7 @@ class ServerArgs:
                 "aiter_attn",
                 "flashinfer_cudnn",
                 "amx_attn",
+                "xpu_attn",
             ],
             default=ServerArgs.mm_attention_backend,
             help="Set multimodal attention backend.",