Fix InternVL and vision attention for non-CUDA backends (e.g. XPU) (#19997)

Co-authored-by: Yang Wang <mr.yang.wang@outlook.com>

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

@@ -22,6 +22,7 @@ from sglang.srt.utils import (
     is_cuda,
     is_hip,
     is_npu,
+    is_xpu,
     print_info_once,
 )
 from sglang.srt.utils.multi_stream_utils import (
@@ -32,6 +33,7 @@ from sglang.srt.utils.multi_stream_utils import (
 _is_cuda = is_cuda()
 _is_npu = is_npu()
 _is_hip = is_hip()
+_is_xpu = is_xpu()
 
 if _is_cuda:
     from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
@@ -366,8 +368,8 @@ class VisionTritonAttention(nn.Module):
                 k,
                 v,
                 output,
-                cu_seqlens.cuda(),
-                seq_lens.cuda(),
+                cu_seqlens.to(q.device),
+                seq_lens.to(q.device),
                 max_seqlen,
                 is_causal=False,
             )
@@ -905,6 +907,8 @@ class VisionAttention(nn.Module):
                 backend = "aiter_attn"
             else:
                 backend = "triton_attn"
+        elif _is_xpu:
+            backend = "triton_attn"
         else:
             backend = "sdpa"
         if backend == "fa3" and is_blackwell_supported():

python/sglang/srt/managers/scheduler.py

@@ -1237,7 +1237,7 @@ class Scheduler(
         self.schedule_stream = self.device_module.Stream(priority=0)
         if self.device == "cpu":
             self.schedule_stream.synchronize = lambda: None  # No-op for CPU
-        with CudaStreamContext(self.schedule_stream):
+        with self.device_module.StreamContext(self.schedule_stream):
             dispatch_event_loop(self)
 
     @DynamicGradMode()

python/sglang/srt/multimodal/processors/internvl.py

@@ -19,6 +19,7 @@ from sglang.srt.multimodal.processors.base_processor import (
     BaseMultiModalProcessorOutput,
     MultimodalSpecialTokens,
 )
+from sglang.srt.utils import get_device
 from sglang.srt.utils.video_decoder import VideoDecoderWrapper
 
 logger = logging.getLogger(__name__)
@@ -434,7 +435,7 @@ class InternVLProcessor(BaseMultimodalProcessor):
             len(base_output.videos),
         )
 
-        mean, std = self._get_normalize_tensors(device="cuda")
+        mean, std = self._get_normalize_tensors(device=get_device())
 
         # ----- Images -> tiles -----
         num_patches_list: List[int] = []
@@ -444,10 +445,11 @@ class InternVLProcessor(BaseMultimodalProcessor):
             if isinstance(image, Image.Image):
                 img_np = np.array(image.convert("RGB"))
                 tensor = (
-                    torch.from_numpy(img_np).permute(2, 0, 1).cuda().float() / 255.0
+                    torch.from_numpy(img_np).permute(2, 0, 1).to(get_device()).float()
+                    / 255.0
                 )
             else:
-                tensor = image.cuda()
+                tensor = image.to(get_device())
 
             tensor = (tensor - mean) / std
             tiles = self.dynamic_preprocess(
@@ -496,7 +498,11 @@ class InternVLProcessor(BaseMultimodalProcessor):
                 for fi in frame_indices:
                     img_np = vr[int(fi)]
                     frame_t = (
-                        torch.from_numpy(img_np).permute(2, 0, 1).cuda().float() / 255.0
+                        torch.from_numpy(img_np)
+                        .permute(2, 0, 1)
+                        .to(get_device())
+                        .float()
+                        / 255.0
                     )
                     frame_t = (frame_t - mean) / std
 
@@ -568,14 +574,14 @@ class InternVLProcessor(BaseMultimodalProcessor):
         image_offsets = []
         if image_tensor is not None:
             image_offsets = self.get_mm_items_offset(
-                input_ids=input_ids_tensor.to("cuda"),
+                input_ids=input_ids_tensor.to(get_device()),
                 mm_token_id=self.img_context_token_id,
             )
 
         video_offsets = []
         if video_tensor is not None and self.video_token_id is not None:
             video_offsets = self.get_mm_items_offset(
-                input_ids=input_ids_tensor.to("cuda"),
+                input_ids=input_ids_tensor.to(get_device()),
                 mm_token_id=self.video_token_id,
             )
 
@@ -633,7 +639,7 @@ class InternVLProcessor(BaseMultimodalProcessor):
             discard_alpha_channel=True,
         )
 
-        mean, std = self._get_normalize_tensors(device="cuda")
+        mean, std = self._get_normalize_tensors(device=get_device())
 
         num_patches_list: List[int] = []
         pixel_values_list: List[torch.Tensor] = []
@@ -642,10 +648,11 @@ class InternVLProcessor(BaseMultimodalProcessor):
             if isinstance(image, Image.Image):
                 img_np = np.array(image.convert("RGB"))
                 tensor = (
-                    torch.from_numpy(img_np).permute(2, 0, 1).cuda().float() / 255.0
+                    torch.from_numpy(img_np).permute(2, 0, 1).to(get_device()).float()
+                    / 255.0
                 )
             else:
-                tensor = image.cuda()
+                tensor = image.to(get_device())
 
             tensor = (tensor - mean) / std
             tiles = self.dynamic_preprocess(
@@ -688,7 +695,7 @@ class InternVLProcessor(BaseMultimodalProcessor):
         image_offsets = []
         if pixel_values is not None:
             image_offsets = self.get_mm_items_offset(
-                input_ids=input_ids_tensor.to("cuda"),
+                input_ids=input_ids_tensor.to(get_device()),
                 mm_token_id=self.img_context_token_id,
             )
 

python/sglang/srt/server_args.py

@@ -751,6 +751,7 @@ class ServerArgs:
         self._handle_hpu_backends()
         self._handle_cpu_backends()
         self._handle_npu_backends()
+        self._handle_xpu_backends()
 
         # Handle piecewise CUDA graph.
         self._handle_piecewise_cuda_graph()
@@ -1027,6 +1028,15 @@ class ServerArgs:
                 )
                 self.piecewise_cuda_graph_compiler = "eager"
 
+    def _handle_xpu_backends(self):
+        if self.device == "xpu":
+            if not self.disable_piecewise_cuda_graph:
+                logger.warning(
+                    "XPU platform does not support piecewise CUDA graph, ignoring --disable-piecewise-cuda-graph"
+                    " flag and disabling piecewise CUDA graph."
+                )
+            self.disable_piecewise_cuda_graph = True
+
     def _handle_piecewise_cuda_graph(self):
         # Skip auto-disable when enforce flag is set (for testing)
         if self.enforce_piecewise_cuda_graph: