From db143e5212f0b5d166b8b1625d3cc98f6ee06f52 Mon Sep 17 00:00:00 2001 From: jianan-gu Date: Tue, 9 Jun 2026 09:47:44 +0800 Subject: [PATCH] [Intel GPU][Encoder] Add xpu_attn backend for encoder vision attention (#26460) Co-authored-by: Ma Mingfei --- python/sglang/srt/layers/attention/vision.py | 72 ++++++++- python/sglang/srt/server_args.py | 1 + .../xpu/test_encoder_attention_backend.py | 143 ++++++++++++++++++ 3 files changed, 215 insertions(+), 1 deletion(-) create mode 100644 test/registered/xpu/test_encoder_attention_backend.py diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 2718ccfe4..cc37c8a08 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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(): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9435f9254..5fdf50ab3 100644 --- a/python/sglang/srt/server_args.py +++ b/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.", diff --git a/test/registered/xpu/test_encoder_attention_backend.py b/test/registered/xpu/test_encoder_attention_backend.py new file mode 100644 index 000000000..1f64c15df --- /dev/null +++ b/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": "\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()