[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.",

test/registered/xpu/test_encoder_attention_backend.py

@@ -0,0 +1,143 @@
+"""
+python3 -m unittest test_encoder_attention_backend.py
+"""
+
+import json
+import os
+import unittest
+from pathlib import Path
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils.hf_transformers import get_tokenizer
+from sglang.test.ci.ci_register import register_xpu_ci
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_xpu_ci(est_time=360, suite="stage-b-test-1-gpu-xpu")
+
+
+class TestEncoderAttention(CustomTestCase):
+    # Test "xpu_attn" attention backend
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen3-VL-2B-Thinking"
+        cls.tokenizer = get_tokenizer(cls.model)
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.image_path = str(
+            (Path(__file__).resolve().parents[3] / "examples/assets/example_image.png")
+        )
+        if not os.path.exists(cls.image_path):
+            raise FileNotFoundError(f"Image not found: {cls.image_path}")
+        cls.common_args = [
+            "--device",
+            "xpu",
+            "--mm-attention-backend",
+            "xpu_attn",
+        ]
+        os.environ["SGLANG_USE_SGL_XPU"] = "1"
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                *cls.common_args,
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        """Fixture that is run once after all tests in the class."""
+        if hasattr(cls, "process") and cls.process:
+            cls.process.terminate()
+            try:
+                cls.process.wait(timeout=30)
+            except Exception:
+                # Force kill if it didn't exit cleanly in time
+                kill_process_tree(cls.process.pid)
+
+    def get_request_json(self, max_new_tokens=32, n=1):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": "<image>\n Tell me what you can see from the image.",
+                "image_data": self.image_path,
+                "sampling_params": {
+                    "temperature": 0 if n == 1 else 1.0,
+                    "max_new_tokens": max_new_tokens,
+                },
+            },
+        )
+        return response.json()
+
+    def run_decode(
+        self,
+        max_new_tokens=128,
+        n=1,
+    ):
+
+        ret = self.get_request_json(max_new_tokens=max_new_tokens, n=n)
+        print(json.dumps(ret, indent=2))
+
+        def assert_one_item(item):
+            if item["meta_info"]["finish_reason"]["type"] == "stop":
+                self.assertEqual(
+                    item["meta_info"]["finish_reason"]["matched"],
+                    self.tokenizer.eos_token_id,
+                )
+            elif item["meta_info"]["finish_reason"]["type"] == "length":
+                self.assertEqual(
+                    len(item["output_ids"]), item["meta_info"]["completion_tokens"]
+                )
+                self.assertEqual(len(item["output_ids"]), max_new_tokens)
+
+        # Determine whether to assert a single item or multiple items based on n
+        if n == 1:
+            assert_one_item(ret)
+        else:
+            self.assertEqual(len(ret), n)
+            for i in range(n):
+                assert_one_item(ret[i])
+
+        print("=" * 100)
+
+    def test_run(self):
+        self.run_decode()
+
+
+class TestEncoderAttention_Triton(TestEncoderAttention):
+    # Test "triton_attn" attention backend
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen3-VL-2B-Thinking"
+        cls.tokenizer = get_tokenizer(cls.model)
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.image_path = str(
+            (Path(__file__).resolve().parents[3] / "examples/assets/example_image.png")
+        )
+        if not os.path.exists(cls.image_path):
+            raise FileNotFoundError(f"Image not found: {cls.image_path}")
+        cls.common_args = [
+            "--device",
+            "xpu",
+            "--mm-attention-backend",
+            "triton_attn",
+        ]
+        os.environ["SGLANG_USE_SGL_XPU"] = "0"
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                *cls.common_args,
+            ],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()