diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index 39cb8f7e2386..87edd2d9a2c6 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/python/sglang/srt/layers/rotary_embedding/base.py @@ -69,7 +69,19 @@ ) 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: + # Older/CUDA sgl_kernel builds don't ship this XPU fused kernel. Don't + # let a missing symbol break module import (it would deregister every + # model that imports get_rope and silently fall back to the generic + # Transformers backbone). forward_xpu() handles the None case below. + 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(MultiPlatformOp): @@ -453,8 +465,12 @@ def forward_xpu( self._match_cos_sin_cache_dtype(query) - # Fused_qk_rope only supports aligned head_size - if self.head_size in [128, 256, 512]: + # Fused_qk_rope only supports aligned head_size, and requires the fused + # XPU kernel to be available in sgl_kernel. + 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) diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index 4ddeebeff894..9913a8955bd9 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -40,7 +40,14 @@ import torch_npu if _is_xpu: - from sgl_kernel import multimodal_rotary_embedding + try: + from sgl_kernel import multimodal_rotary_embedding + except ImportError: + # Older/CUDA sgl_kernel builds don't ship this XPU kernel. Avoid a hard + # import failure (it would deregister every model importing get_rope and + # silently fall back to the generic Transformers backbone). forward_xpu() + # falls back to forward_native when the kernel is None. + multimodal_rotary_embedding = None import triton import triton.language as tl @@ -386,7 +393,11 @@ def forward_xpu( fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: assert positions.ndim in (1, 2) - 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, diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 7a47d10cdd76..6b16870312b4 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/python/sglang/srt/lora/backend/chunked_backend.py @@ -224,7 +224,7 @@ def init_cuda_graph_batch_info( (num_tokens_per_bs + MIN_CHUNK_SIZE - 1) // MIN_CHUNK_SIZE ) * max_bs_in_cuda_graph max_num_tokens = max_bs_in_cuda_graph * num_tokens_per_bs - 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, diff --git a/python/sglang/srt/lora/backend/torch_backend.py b/python/sglang/srt/lora/backend/torch_backend.py index 0f05a3f12f7f..64e4ccedaaed 100644 --- a/python/sglang/srt/lora/backend/torch_backend.py +++ b/python/sglang/srt/lora/backend/torch_backend.py @@ -166,7 +166,7 @@ def init_cuda_graph_batch_info( max_bs_in_cuda_graph: int, num_tokens_per_bs: 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, diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index 47bca9c6d7ae..b1c04cd231b0 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -146,7 +146,7 @@ def init_cuda_graph_batch_info( ): max_tokens = max_bs_in_cuda_graph * num_tokens_per_bs 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, diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index f44bc486a02b..3f3ee05cbcca 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/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 diff --git a/python/sglang/srt/lora/lora_overlap_loader.py b/python/sglang/srt/lora/lora_overlap_loader.py index 6d5845ba084d..380e5db9e57b 100644 --- a/python/sglang/srt/lora/lora_overlap_loader.py +++ b/python/sglang/srt/lora/lora_overlap_loader.py @@ -72,7 +72,7 @@ def _drain_completed_overlap_loads(self) -> None: 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( diff --git a/python/sglang/test/lora_utils.py b/python/sglang/test/lora_utils.py index 0bfd7fa07a80..ad3d38c0f3c9 100644 --- a/python/sglang/test/lora_utils.py +++ b/python/sglang/test/lora_utils.py @@ -4,9 +4,34 @@ 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: @@ -625,17 +650,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_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 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_no_lora_outputs.output_strs[i], + hf_no_lora_outputs.output_strs[i], + rouge_tol, + f"base '{base_path}', backend '{backend}' (no-LoRA baseline)", ) diff --git a/test/registered/lora/test_chunked_sgmv_backend.py b/test/registered/lora/test_chunked_sgmv_backend.py index e075043a7012..fbb3c2ccd383 100644 --- a/test/registered/lora/test_chunked_sgmv_backend.py +++ b/test/registered/lora/test_chunked_sgmv_backend.py @@ -20,7 +20,8 @@ from sglang.srt.lora.triton_ops.chunked_sgmv_shrink import _chunked_lora_shrink_kernel from sglang.srt.lora.utils import LoRABatchInfo, get_lm_head_pruned_lens from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci from sglang.test.lora_utils import ( reference_embedding_lora_a_shrink, reference_sgmv_expand, @@ -30,11 +31,27 @@ CHUNK_SIZE = 16 register_cuda_ci(est_time=60, suite="nightly-1-gpu", nightly=True) +register_xpu_ci(est_time=60, suite="stage-a-test-1-gpu-xpu") + + +def _clear_jit_cache(kernel): + """Clear a triton JITFunction's compiled-kernel cache across triton versions. + + Older triton exposes ``_clear_cache()``; newer builds (e.g. Intel triton on + XPU) keep per-device caches in ``device_caches`` instead. + """ + clear = getattr(kernel, "_clear_cache", None) + if callable(clear): + clear() + return + device_caches = getattr(kernel, "device_caches", None) + if isinstance(device_caches, dict): + device_caches.clear() def reset_kernel_cache(): - _chunked_lora_shrink_kernel._clear_cache() - _chunked_lora_expand_kernel._clear_cache() + _clear_jit_cache(_chunked_lora_shrink_kernel) + _clear_jit_cache(_chunked_lora_expand_kernel) class BatchComposition(Enum): @@ -110,7 +127,7 @@ def setUp(self): torch.manual_seed(42) random.seed(42) - self.device = torch.device("cuda") + self.device = torch.device(get_device()) self.dtype = torch.float16 self.input_dim = 2560 # Hidden dimension self.max_seq_len = 1024 diff --git a/test/registered/lora/test_fused_moe_lora_kernel.py b/test/registered/lora/test_fused_moe_lora_kernel.py index 9b152464cf2c..20019373632e 100644 --- a/test/registered/lora/test_fused_moe_lora_kernel.py +++ b/test/registered/lora/test_fused_moe_lora_kernel.py @@ -9,13 +9,15 @@ # IMPORT PREBUILT KERNEL # ============================================================================== from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size +from sglang.srt.lora.lora_moe_runners import _naive_moe_lora_align_block_size from sglang.srt.lora.triton_ops import fused_moe_lora -from sglang.srt.utils import set_random_seed -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.utils import get_device, is_xpu, set_random_seed +from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci # ============================================================================== register_cuda_ci(est_time=28, stage="base-b", runner_config="1-gpu-large") +register_xpu_ci(est_time=28, suite="stage-a-test-1-gpu-xpu") def round_up(x, base): @@ -159,23 +161,40 @@ def use_fused_moe_lora_kernel( adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32, device=device) lora_ids = torch.arange(max_loras, dtype=torch.int32, device=device) - # call kernel - moe_lora_align_block_size( - topk_ids, - seg_indptr, - req_to_lora, - num_experts, - block_size, - max_loras, - max_num_tokens_padded, - max_num_m_blocks, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - adapter_enabled, - lora_ids, - None, # maybe_expert_map - ) + # call kernel — the fused CUDA align kernel exists only for CUDA; on XPU + # use the pure-torch native alignment. + if is_xpu(): + sorted_token_ids, expert_ids, num_tokens_post_padded = ( + _naive_moe_lora_align_block_size( + topk_ids, + seg_indptr, + req_to_lora, + num_experts, + block_size, + max_loras, + max_num_tokens_padded, + max_num_m_blocks, + adapter_enabled, + device, + ) + ) + else: + moe_lora_align_block_size( + topk_ids, + seg_indptr, + req_to_lora, + num_experts, + block_size, + max_loras, + max_num_tokens_padded, + max_num_m_blocks, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + adapter_enabled, + lora_ids, + None, # maybe_expert_map + ) config = { "BLOCK_SIZE_M": 16, @@ -263,7 +282,7 @@ def use_torch( DTYPES = [torch.float32, torch.float16, torch.bfloat16] -DEVICES = [f"cuda:{0}"] +DEVICES = [get_device(0)] SEED = [42] diff --git a/test/registered/lora/test_lora_hf_sgl_logprob_diff.py b/test/registered/lora/test_lora_hf_sgl_logprob_diff.py index efebc2e1b33c..a4ca453ec748 100644 --- a/test/registered/lora/test_lora_hf_sgl_logprob_diff.py +++ b/test/registered/lora/test_lora_hf_sgl_logprob_diff.py @@ -35,7 +35,11 @@ import numpy as np import torch -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.ci.ci_register import ( + register_amd_ci, + register_cuda_ci, + register_xpu_ci, +) from sglang.test.runners import HFRunner, SRTRunner from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, CustomTestCase @@ -48,6 +52,7 @@ est_time=250, suite="stage-b-test-1-gpu-small-amd", ) +register_xpu_ci(est_time=250, suite="stage-a-test-1-gpu-xpu") # Test configuration constants BASE_MODEL = "meta-llama/Llama-2-7b-hf" LORA_PATHS = ["yushengsu/sglang_lora_logprob_diff_without_tuning"] diff --git a/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py b/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py index a35e69dd9e54..c9a09d550c7c 100644 --- a/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py +++ b/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py @@ -16,7 +16,7 @@ import torch -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci from sglang.test.lora_utils import ( MOE_BASE_MODEL_PATH, MOE_LORA_PATH, @@ -29,6 +29,7 @@ stage="base-b", runner_config="1-gpu-large", ) +register_xpu_ci(est_time=60, suite="stage-a-test-1-gpu-xpu") # Format: [{"text": "result string", "lps": [0.1, 0.2, ...]}, ...] VLLM_CACHED_RESULTS = [ @@ -287,6 +288,17 @@ class TestMoELoraRegression(unittest.TestCase): def test_sglang_moe_parity_strict(self): + # flashinfer is CUDA-only; on other devices (e.g. XPU) use the + # platform-appropriate attention backend. + from sglang.srt.utils import is_cuda, is_xpu + + if is_cuda(): + attention_backend = "flashinfer" + elif is_xpu(): + attention_backend = "intel_xpu" + else: + attention_backend = "triton" + with SRTRunner( model_path=MOE_BASE_MODEL_PATH, torch_dtype=torch.bfloat16, @@ -296,7 +308,7 @@ def test_sglang_moe_parity_strict(self): tp_size=1, trust_remote_code=True, disable_radix_cache=True, - attention_backend="flashinfer", + attention_backend=attention_backend, mem_fraction_static=0.80, ) as srt_runner: diff --git a/test/registered/lora/test_lora_tied_lm_head.py b/test/registered/lora/test_lora_tied_lm_head.py index 4070a53217c5..ce95fea4b237 100644 --- a/test/registered/lora/test_lora_tied_lm_head.py +++ b/test/registered/lora/test_lora_tied_lm_head.py @@ -46,11 +46,20 @@ from transformers import AutoModelForCausalLM -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci from sglang.test.runners import HFRunner, SRTRunner from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, CustomTestCase register_cuda_ci(est_time=120, suite="nightly-1-gpu", nightly=True) +register_xpu_ci(est_time=120, suite="stage-a-test-1-gpu-xpu") + + +def _empty_device_cache(): + """Device-agnostic cache release (cuda/xpu/...).""" + module = torch.get_device_module(get_device()) + if hasattr(module, "empty_cache"): + module.empty_cache() # Use a small model with tie_word_embeddings=True BASE_MODEL = "Qwen/Qwen2.5-0.5B" @@ -123,7 +132,7 @@ def create_lora_adapter_with_lm_head(base_model_name: str, output_dir: str): # Clean up the model to free memory del peft_model, model - torch.cuda.empty_cache() + _empty_device_cache() class TestLoRATiedLMHead(CustomTestCase): @@ -174,7 +183,7 @@ def test_tied_lm_head_lora_hf_sgl_logprob_match(self): lora_paths=[self._adapter_dir] * len(prompts), ) - torch.cuda.empty_cache() + _empty_device_cache() # Run HuggingFace with LoRA (via PEFT) with HFRunner( diff --git a/test/registered/lora/test_lora_tp.py b/test/registered/lora/test_lora_tp.py index 3ed859d11e27..667264f7267f 100644 --- a/test/registered/lora/test_lora_tp.py +++ b/test/registered/lora/test_lora_tp.py @@ -17,6 +17,7 @@ import unittest from typing import List, Optional +from sglang.srt.utils import is_xpu from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.lora_utils import ( ALL_OTHER_LORA_MODELS, @@ -56,6 +57,12 @@ def _run_tp_on_model_cases( if not model_case.skip_long_prompt else [p for p in DEFAULT_PROMPTS if len(p) < 1000] ) + _backend = "csgmv" + _attention_backend = "fa3" + if is_xpu(): + _backend = "torch_native" + _attention_backend = "triton" # Using triton instead of intel_xpu to avoid NaN issue + for tp_size in tp_list: model_case.tp_size = tp_size for torch_dtype in TORCH_DTYPES: @@ -66,7 +73,8 @@ def _run_tp_on_model_cases( max_new_tokens=32, enable_lora_overlap_loading=enable_lora_overlap_loading, test_tag=f"tp={tp_size}, enable_lora_overlap_loading={enable_lora_overlap_loading}", - attention_backend="fa3", + backend=_backend, + attention_backend=_attention_backend, ) def test_ci_lora_models(self): diff --git a/test/registered/lora/test_moe_lora_info.py b/test/registered/lora/test_moe_lora_info.py index 44dd871c27f4..d94c2e3b5fd0 100644 --- a/test/registered/lora/test_moe_lora_info.py +++ b/test/registered/lora/test_moe_lora_info.py @@ -4,9 +4,13 @@ import torch from sglang.srt.lora.backend.base_backend import _compute_moe_lora_info -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small") +register_xpu_ci(est_time=5, suite="stage-a-test-1-gpu-xpu") + +DEVICE = get_device() def _expected_adapter_enabled( @@ -24,7 +28,7 @@ def _expected_adapter_enabled( @pytest.mark.parametrize("use_preallocated_buffers", [False, True]) def test_compute_moe_lora_info_expands_segments(use_preallocated_buffers: bool): - device = "cuda" + device = DEVICE seg_lens = torch.tensor([5, 1, 7, 3, 9, 2], dtype=torch.int32, device=device) seg_indptr = torch.zeros((seg_lens.numel() + 1,), dtype=torch.int32, device=device) seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) @@ -53,7 +57,7 @@ def test_compute_moe_lora_info_expands_segments(use_preallocated_buffers: bool): token_lora_mapping, max_len=int(seg_lens.max().item()), ) - torch.cuda.synchronize() + torch.get_device_module(device).synchronize() expected_mapping = torch.repeat_interleave(weight_indices, seg_lens) expected_enabled = _expected_adapter_enabled(lora_ranks, weight_indices) @@ -65,8 +69,13 @@ def test_compute_moe_lora_info_expands_segments(use_preallocated_buffers: bool): assert actual_mapping.data_ptr() == token_lora_mapping.data_ptr() +@pytest.mark.skipif( + not (hasattr(torch, "cuda") and torch.cuda.is_available()), + reason="The launch-coverage assertion guards the CUDA triton kernel path only; " + "non-CUDA devices (e.g. XPU) use the pure-torch fallback which has no such launch.", +) def test_compute_moe_lora_info_rejects_undercovered_launch(): - device = "cuda" + device = DEVICE seg_indptr = torch.tensor([0, 300], dtype=torch.int32, device=device) weight_indices = torch.tensor([0], dtype=torch.int32, device=device) lora_ranks = torch.tensor([16], dtype=torch.int32, device=device) diff --git a/test/registered/lora/test_multi_lora_backend.py b/test/registered/lora/test_multi_lora_backend.py index 5f3e372e8e9e..7f408fb0b5eb 100644 --- a/test/registered/lora/test_multi_lora_backend.py +++ b/test/registered/lora/test_multi_lora_backend.py @@ -16,7 +16,11 @@ import os import unittest -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.ci.ci_register import ( + register_amd_ci, + register_cuda_ci, + register_xpu_ci, +) from sglang.test.lora_utils import ( ALL_OTHER_MULTI_LORA_MODELS, CI_MULTI_LORA_MODELS, @@ -27,6 +31,7 @@ register_cuda_ci(est_time=99, stage="base-b", runner_config="1-gpu-large") register_amd_ci(est_time=100, suite="stage-b-test-1-gpu-small-amd") +register_xpu_ci(est_time=120, suite="stage-a-test-1-gpu-xpu") class TestMultiLoRABackend(CustomTestCase): diff --git a/test/registered/lora/test_virtual_experts_kernels.py b/test/registered/lora/test_virtual_experts_kernels.py index 91423abc88f9..c16364d06cbf 100644 --- a/test/registered/lora/test_virtual_experts_kernels.py +++ b/test/registered/lora/test_virtual_experts_kernels.py @@ -23,14 +23,28 @@ python -m pytest test/registered/lora/test_virtual_experts_kernels.py -v """ +import importlib.util import unittest import torch -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.utils import get_device, is_cuda, is_xpu +from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small") +register_xpu_ci(est_time=15, suite="stage-a-test-1-gpu-xpu") + + +def _require_accelerator_device(): + """Return the active GPU device string, or skip if no GPU is available. + + The virtual-experts kernels are triton kernels that run on CUDA or XPU + (Intel triton). They have no CPU implementation. + """ + if not (is_cuda() or is_xpu()): + raise unittest.SkipTest("CUDA or XPU required") + return get_device(0) from sglang.srt.lora.triton_ops.virtual_experts import ( _align_block_size_jit, @@ -45,9 +59,7 @@ class TestFusedVirtualTopkIdsPreservesSentinels(CustomTestCase): @classmethod def setUpClass(cls): - if not torch.cuda.is_available(): - raise unittest.SkipTest("CUDA required") - cls.device = "cuda:0" + cls.device = _require_accelerator_device() def test_negative_sentinels_preserved(self): # Mix of valid topk_ids in [0, num_experts), -1 sentinels (typical @@ -139,11 +151,9 @@ class _AlignBlockSizeSentinelBucketBase(CustomTestCase): @classmethod def setUpClass(cls): - if not torch.cuda.is_available(): - raise unittest.SkipTest("CUDA required") if cls is _AlignBlockSizeSentinelBucketBase: raise unittest.SkipTest("Base class") - cls.device = "cuda:0" + cls.device = _require_accelerator_device() def _align(self, topk_ids, block_size, num_experts): raise NotImplementedError @@ -264,6 +274,13 @@ def _align(self, topk_ids, block_size, num_experts): return _align_block_size_torch(topk_ids, block_size, num_experts) +@unittest.skipUnless( + is_cuda() and importlib.util.find_spec("tvm_ffi") is not None, + "_align_block_size_jit builds a CUDA JIT kernel via tvm_ffi.load_inline " + "(requires a CUDA/nvcc install); it cannot run on non-CUDA platforms such " + "as XPU. The torch variant (TestAlignBlockSizeTorchSentinelBucket) covers " + "the same alignment logic on those platforms.", +) class TestAlignBlockSizeJitSentinelBucket(_AlignBlockSizeSentinelBucketBase): """Test the CUDA JIT kernel path (with fused_sanitize_expert_ids, as in production).""" diff --git a/test/registered/unit/lora/test_mem_pool_ep_unit.py b/test/registered/unit/lora/test_mem_pool_ep_unit.py index 5856eecbaa38..703245f7661b 100644 --- a/test/registered/unit/lora/test_mem_pool_ep_unit.py +++ b/test/registered/unit/lora/test_mem_pool_ep_unit.py @@ -13,11 +13,16 @@ python -m pytest test/registered/unit/lora/test_mem_pool_ep_unit.py -v """ -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.ci.ci_register import ( + register_amd_ci, + register_cuda_ci, + register_xpu_ci, +) # CPU-only unit test; no CUDA/distributed dependencies. register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd") +register_xpu_ci(est_time=9, suite="stage-a-test-1-gpu-xpu") import types import unittest