[Intel][XPU][LoRA] Enable LoRA on Intel XPU (#30345)

python/sglang/srt/arg_groups/overrides.py

@@ -564,6 +564,8 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset(
 def supports_mamba_cache_extra_buffer(view: Any, model_arch: str) -> bool:
     """Whether ``model_arch`` supports the extra_buffer strategy on the
     configured linear-attention backend (pure read)."""
+    if get_platform().is_xpu:
+        return False
     if model_arch in _MAMBA_EXTRA_BUFFER_ARCHS:
         return view.linear_attn_backend == "triton"
     return False

python/sglang/srt/layers/rotary_embedding/base.py

@@ -72,7 +72,15 @@ if _is_hip:
     )
 
 if _is_xpu:
-    from sgl_kernel import fused_qk_rope_with_cos_sin_cache_inplace
+    try:
+        from sgl_kernel import fused_qk_rope_with_cos_sin_cache_inplace
+    except ImportError:
+        fused_qk_rope_with_cos_sin_cache_inplace = None
+        logger.warning(
+            "sgl_kernel.fused_qk_rope_with_cos_sin_cache_inplace is unavailable; "
+            "XPU rotary embedding will use the generic rotary_embedding kernel. "
+            "Upgrade sgl_kernel to enable the fused XPU kernel."
+        )
 
 
 class RotaryEmbedding(BaseFusedOp):
@@ -454,7 +462,11 @@ class RotaryEmbedding(BaseFusedOp):
         positions = torch.add(positions, offsets) if offsets is not None else positions
 
         # Fused_qk_rope only supports aligned head_size
-        if self.head_size in [128, 256, 512]:
+        if fused_qk_rope_with_cos_sin_cache_inplace is not None and self.head_size in [
+            128,
+            256,
+            512,
+        ]:
             num_tokens = positions.size(0)
             q_rope = query.view(num_tokens, -1, self.head_size)
             k_rope = key.view(num_tokens, -1, self.head_size)

python/sglang/srt/layers/rotary_embedding/mrope.py

@@ -39,7 +39,10 @@ if _is_npu:
     import torch_npu
 
 if _is_xpu:
-    from sgl_kernel import multimodal_rotary_embedding
+    try:
+        from sgl_kernel import multimodal_rotary_embedding
+    except ImportError:
+        multimodal_rotary_embedding = None
 
 from sglang.kernels.ops.attention.mrope import apply_interleaved_rope_triton
 
@@ -309,7 +312,11 @@ class MRotaryEmbedding(RotaryEmbedding):
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         assert positions.ndim in (1, 2)
         self._match_cos_sin_cache_dtype(query)
-        if positions.ndim == 2 and self.mrope_section:
+        if (
+            multimodal_rotary_embedding is not None
+            and positions.ndim == 2
+            and self.mrope_section
+        ):
             multimodal_rotary_embedding(
                 query,
                 key,

python/sglang/srt/lora/backend/base_backend.py

@@ -440,10 +440,8 @@ def _compute_moe_lora_info(
     adapter_enabled.zero_()
 
     has_segments = weight_indices.numel() != 0
-    use_cuda_kernel = (
-        num_tokens != 0 and has_segments and seg_indptr.device.type == "cuda"
-    )
-    if use_cuda_kernel:
+    needs_launch = num_tokens != 0 and has_segments
+    if needs_launch:
         block_size = 256
         tiles_per_segment = triton.cdiv(max_len, block_size)
         grid_size = tiles_per_segment * weight_indices.numel()
@@ -451,6 +449,10 @@ def _compute_moe_lora_info(
             f"MoE LoRA token-mapping launch under-covers tokens: "
             f"{grid_size=} {block_size=} {num_tokens=}"
         )
+
+    # Triton kernel on CUDA only; every other device (e.g. XPU) falls through to
+    # the native torch path below, which yields the same mapping.
+    if needs_launch and seg_indptr.device.type == "cuda":
         _compute_moe_lora_info_kernel[(grid_size,)](
             seg_indptr,
             lora_ranks,

python/sglang/srt/lora/backend/chunked_backend.py

@@ -223,7 +223,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
             (num_tokens_per_req + MIN_CHUNK_SIZE - 1) // MIN_CHUNK_SIZE
         ) * max_bs_in_cuda_graph
         max_num_tokens = max_bs_in_cuda_graph * num_tokens_per_req
-        with torch.device("cuda"):
+        with torch.device(self.device):
             self.cuda_graph_batch_info = LoRABatchInfo(
                 bs=max_bs_in_cuda_graph,
                 use_cuda_graph=True,

python/sglang/srt/lora/backend/torch_backend.py

@@ -166,7 +166,7 @@ class TorchNativeLoRABackend(BaseLoRABackend):
         max_bs_in_cuda_graph: int,
         num_tokens_per_req: int,
     ):
-        with torch.device("cuda"):
+        with torch.device(self.device):
             self.cuda_graph_batch_info = TorchNativeLoRABatchInfo(
                 use_cuda_graph=True,
                 bs=max_bs_in_cuda_graph,

python/sglang/srt/lora/backend/triton_backend.py

@@ -163,7 +163,7 @@ class TritonLoRABackend(BaseLoRABackend):
     ):
         max_tokens = max_bs_in_cuda_graph * num_tokens_per_req
         mlpb = self.max_loras_per_batch
-        with torch.device("cuda"):
+        with torch.device(self.device):
             self.cuda_graph_batch_info = LoRABatchInfo(
                 bs=max_bs_in_cuda_graph,
                 use_cuda_graph=True,

python/sglang/srt/lora/lora_moe_runners.py

@@ -227,7 +227,7 @@ def _compute_lora_alignment(
 
     device = topk_ids.device
 
-    use_naive = (
+    use_naive = _is_xpu or (
         cg is None
         and M * topk_ids.shape[1] * _SPARSITY_FACTOR
         <= lora_info.num_experts * max_loras

python/sglang/srt/lora/lora_overlap_loader.py

@@ -44,16 +44,19 @@ class LoRAOverlapLoader:
         lora_pipeline_load_status = self._check_overlap_load_status(lora_id)
         if lora_pipeline_load_status == LoRAOverlapLoadStatus.LOADING:
             return False
-        elif lora_pipeline_load_status == LoRAOverlapLoadStatus.NOT_LOADED:
-            res = self._try_start_overlap_load(lora_id, running_loras)
-            if res:
-                logger.debug(f"Loading LoRA adapter {lora_id} asynchronously")
-
-            return False
-        else:
-            assert lora_pipeline_load_status == LoRAOverlapLoadStatus.LOADED
+        elif lora_pipeline_load_status == LoRAOverlapLoadStatus.LOADED:
             return True
 
+        assert lora_pipeline_load_status == LoRAOverlapLoadStatus.NOT_LOADED
+        if not self._try_start_overlap_load(lora_id, running_loras):
+            return False
+        logger.debug(f"Loading LoRA adapter {lora_id} asynchronously")
+
+        # Report an already-finished copy as LOADED, or a sibling's load could
+        # evict it before it is ever scheduled. No-op while still in flight.
+        self._drain_completed_overlap_loads()
+        return self._check_overlap_load_status(lora_id) == LoRAOverlapLoadStatus.LOADED
+
     def _check_overlap_load_status(
         self, lora_id: Optional[str]
     ) -> LoRAOverlapLoadStatus:
@@ -74,7 +77,7 @@ class LoRAOverlapLoader:
             if event.query()
         ]
         for lora_id, event in completed_loads:
-            torch.cuda.current_stream().wait_event(event)
+            self.device_module.current_stream().wait_event(event)
             del self.lora_to_overlap_load_event[lora_id]
 
     def _try_start_overlap_load(

python/sglang/test/lora_utils.py

@@ -4,9 +4,36 @@ from typing import List, Optional
 
 import torch
 
+from sglang.srt.utils import is_xpu
 from sglang.test.runners import HFRunner, SRTRunner
 from sglang.test.test_utils import calculate_rouge_l
 
+_IS_XPU = is_xpu()
+
+
+def _assert_lora_output_match(
+    srt_str: str, hf_str: str, rouge_tol: float, context: str
+):
+    """Compare SRT vs HF greedy output strings.
+
+    Everywhere except XPU we keep the historical strict exact-match (SGLang and HF
+    kernels agree numerically enough for greedy argmax to pick identical tokens).
+    On XPU, small kernel-level fp differences can make greedy decoding diverge
+    after a shared prefix even when the LoRA math is correct, so we fall back to
+    the same ROUGE-L tolerance the per-adaptor comparison path uses.
+    """
+    srt_str = srt_str.strip(" ")
+    hf_str = hf_str.strip(" ")
+    if not _IS_XPU:
+        assert srt_str == hf_str, (srt_str, hf_str)
+        return
+    rouge_score = calculate_rouge_l([srt_str], [hf_str])[0]
+    if rouge_score < rouge_tol:
+        raise AssertionError(
+            f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} for {context}. "
+            f"SRT: {srt_str!r} HF: {hf_str!r}"
+        )
+
 
 @dataclasses.dataclass
 class LoRAAdaptor:
@@ -624,17 +651,22 @@ def run_lora_test_by_batch(
         print("HF output:", hf_output_str)
         print("SRT no lora output:", srt_no_lora_outputs.output_strs[i].strip())
         print("HF no lora output:", hf_no_lora_outputs.output_strs[i].strip())
-        assert srt_outputs.output_strs[i].strip(" ") == hf_outputs.output_strs[i].strip(
-            " "
-        ), (
-            srt_outputs.output_strs[i].strip(" "),
-            hf_outputs.output_strs[i].strip(" "),
+        rouge_tol = (
+            adaptors[i].rouge_l_tolerance
+            if adaptors[i].rouge_l_tolerance is not None
+            else model_case.rouge_l_tolerance
         )
-        assert srt_no_lora_outputs.output_strs[i].strip(
-            " "
-        ) == hf_no_lora_outputs.output_strs[i].strip(" "), (
-            srt_no_lora_outputs.output_strs[i].strip(" "),
-            hf_no_lora_outputs.output_strs[i].strip(" "),
+        _assert_lora_output_match(
+            srt_outputs.output_strs[i],
+            hf_outputs.output_strs[i],
+            rouge_tol,
+            f"base '{base_path}', adaptor '{adaptor_names[i]}', backend '{backend}' (LoRA)",
+        )
+        _assert_lora_output_match(
+            srt_no_lora_outputs.output_strs[i],
+            hf_no_lora_outputs.output_strs[i],
+            rouge_tol,
+            f"base '{base_path}', backend '{backend}' (no-LoRA baseline)",
         )