minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Anupa Sajikumar
2026-09-08 12:46:41 +08:00
committed by GitHub
parent 80e9a4ec74
commit 91a45ea37e
19 changed files with 1262 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test-xpu.yml
+2 -2
View File
@@ -110,7 +110,7 @@ jobs:
timeout-minutes: 60
run: |
docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install --upgrade pip
docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install pytest expecttest ray huggingface_hub tabulate "lmcache>=0.3.9"
docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install pytest expecttest ray huggingface_hub tabulate "lmcache>=0.3.9" accelerate
docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip uninstall -y flashinfer-python sgl-kernel sglang
docker exec ci_sglang_xpu cp /sglang-checkout/python/pyproject_xpu.toml /sglang-checkout/python/pyproject.toml
# Fetch tags so setuptools_scm resolves a real version instead of
@@ -202,7 +202,7 @@ jobs:
timeout-minutes: 60
run: |
docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install --upgrade pip
docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install pytest expecttest ray huggingface_hub tabulate "lmcache>=0.3.9"
docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install pytest expecttest ray huggingface_hub tabulate "lmcache>=0.3.9" accelerate
docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip uninstall -y flashinfer-python sgl-kernel sglang
docker exec ci_sglang_xpu cp /sglang-checkout/python/pyproject_xpu.toml /sglang-checkout/python/pyproject.toml
# Fetch tags so setuptools_scm resolves a real version instead of
python/sglang/srt/arg_groups/overrides.py
+2
View File
@@ -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
+14 -2
View File
@@ -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
+9 -2
View File
@@ -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
+6 -4
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+12 -9
View File
@@ -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
+42 -10
View File
@@ -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)",
)
test/manual/lora/test_lora_xpu_basics.py
+1038
View File
File diff suppressed because it is too large Load Diff
test/manual/lora/test_torch_backend.py
+9 -5
View File
@@ -103,18 +103,20 @@ class TestTorchNativeLoRABackend(CustomTestCase):
self.lora_ranks, dtype=torch.int32, device=self.device
)
output_offset = torch.tensor(
[0, weight_out_dim], dtype=torch.int32, device="cpu"
)
expect_output = reference_sgmv_expand(
x,
weights,
weight_indices_tensor,
seg_len_tensor,
lora_ranks_tensor,
slice_offsets=torch.tensor(
[0, weight_out_dim], dtype=torch.int32, device="cpu"
),
slice_offsets=output_offset,
)
actual_output = self.backend.run_lora_b_sgemm(x, weights)
actual_output = self.backend.run_lora_b_sgemm(x, weights, output_offset)
self.assertTrue(torch.allclose(actual_output, expect_output))
@@ -237,7 +239,9 @@ class TestTorchNativeLoRABackend(CustomTestCase):
slice_offsets=output_offset,
)
actual_output = self.backend.run_gate_up_lora(x, gate_up_lora_a, gate_up_lora_b)
actual_output = self.backend.run_gate_up_lora(
x, gate_up_lora_a, gate_up_lora_b, output_offset
)
self.assertTrue(torch.allclose(actual_output, expect_output))
test/registered/lora/test_fused_moe_lora_kernel.py
+37 -19
View File
@@ -11,7 +11,8 @@ from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora
# IMPORT PREBUILT KERNEL
# ==============================================================================
from sglang.kernels.ops.moe.moe_lora_align import moe_lora_align_block_size
from sglang.srt.utils import set_random_seed
from sglang.srt.lora.lora_moe_runners import _naive_moe_lora_align_block_size
from sglang.srt.utils import get_device, is_xpu, set_random_seed
from sglang.test.ci.ci_register import register_cuda_ci
# ==============================================================================
@@ -160,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,
@@ -264,7 +282,7 @@ def use_torch(
DTYPES = [torch.float32, torch.float16, torch.bfloat16]
DEVICES = [f"cuda:{0}"]
DEVICES = [get_device(0)]
SEED = [42]
test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py
+25 -2
View File
@@ -284,7 +284,29 @@ REFERENCE_STATS = {
class TestMoELoraRegression(unittest.TestCase):
def test_sglang_moe_parity_strict(self):
def test_sglang_moe_parity_flashinfer(self):
# flashinfer is CUDA-only.
from sglang.srt.utils import is_cuda
if not is_cuda():
self.skipTest("flashinfer backend requires CUDA")
self._run_parity_strict(attention_backend="flashinfer")
def test_sglang_moe_parity_intel_xpu(self):
from sglang.srt.utils import is_xpu
if not is_xpu():
self.skipTest("intel_xpu backend requires XPU")
self._run_parity_strict(attention_backend="intel_xpu")
def test_sglang_moe_parity_triton(self):
from sglang.srt.utils import is_cuda, is_xpu
if is_cuda() or is_xpu():
self.skipTest("triton backend is the fallback for non-CUDA/non-XPU")
self._run_parity_strict(attention_backend="triton")
def _run_parity_strict(self, *, attention_backend, **runner_kwargs):
with SRTRunner(
model_path=MOE_BASE_MODEL_PATH,
@@ -295,8 +317,9 @@ class TestMoELoraRegression(unittest.TestCase):
tp_size=1,
trust_remote_code=True,
disable_radix_cache=True,
attention_backend="flashinfer",
attention_backend=attention_backend,
mem_fraction_static=0.80,
**runner_kwargs,
) as srt_runner:
srt_outputs = srt_runner.forward(
MOE_LORA_TEST_PROMPTS,
test/registered/lora/test_lora_overlap_loading.py
+7 -2
View File
@@ -23,6 +23,7 @@ from torch.cuda import Stream as CudaStream
from sglang.srt.lora.lora_manager import LoRAManager
from sglang.srt.lora.lora_overlap_loader import LoRAOverlapLoader, LoRAOverlapLoadStatus
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.lora_utils import (
CI_MULTI_LORA_MODELS,
@@ -33,6 +34,8 @@ from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=237, stage="base-b", runner_config="1-gpu-large")
register_amd_ci(est_time=75, suite="stage-b-test-1-gpu-small-amd")
DEVICE = get_device(0)
class TestLoRAOverlapLoading(CustomTestCase):
def test_ci_lora_models_batch_splitting(self):
@@ -56,6 +59,8 @@ class TestLoRAOverlapLoaderUnitTests(CustomTestCase):
self.mock_stream = MagicMock(spec=CudaStream)
self.mock_stream_context = MagicMock()
self.mock_event = MagicMock(spec=CudaEvent)
# Default to an in-flight copy; cases needing completion flip query().
self.mock_event.query.return_value = False
self.mock_device_module.Stream.return_value = self.mock_stream
self.mock_device_module.stream.return_value = self.mock_stream_context
@@ -64,7 +69,7 @@ class TestLoRAOverlapLoaderUnitTests(CustomTestCase):
self.mock_torch.cuda.current_stream.return_value = MagicMock(spec=CudaStream)
self.mock_lora_manager = MagicMock(spec=LoRAManager)
self.mock_lora_manager.device = "cuda:0"
self.mock_lora_manager.device = DEVICE
self.mock_lora_manager.memory_pool = MagicMock()
self.mock_lora_manager.memory_pool.uid_to_buffer_id = {}
self.mock_lora_manager.validate_lora_batch.return_value = True
@@ -166,7 +171,7 @@ class TestLoRAOverlapLoaderUnitTests(CustomTestCase):
def test_pending_load_is_synchronized_before_unload(self):
manager = LoRAManager.__new__(LoRAManager)
manager.device = torch.device("cuda:0")
manager.device = torch.device(DEVICE)
manager.pending_lora_load_events = {}
manager.memory_pool = MagicMock()
manager.configs = {"lora_A": object()}
test/registered/lora/test_lora_qwen3_8b_logprob_diff.py
+20 -3
View File
@@ -50,6 +50,25 @@ DECODE_ATTENTION_BACKEND = "fa4"
KL_THRESHOLD = 5e-3
def attention_backend_kwargs():
"""Engine attention-backend kwargs for the current platform.
fa4 and flashinfer are CUDA-only: fa4 dispatches into the CUTLASS CUTE DSL
kernel, which cannot import off CUDA. On XPU the equivalent fused path is
the intel_xpu backend, so select it there instead of forcing a backend the
device has no kernels for.
"""
from sglang.srt.utils import is_xpu
if is_xpu():
return {"attention_backend": "intel_xpu"}
return {
"attention_backend": "flashinfer",
"prefill_attention_backend": PREFILL_ATTENTION_BACKEND,
"decode_attention_backend": DECODE_ATTENTION_BACKEND,
}
def kl_v2(a, b):
a = torch.tensor(a) if not torch.is_tensor(a) else a
b = torch.tensor(b) if not torch.is_tensor(b) else b
@@ -137,9 +156,7 @@ class TestLoRAQwen3_8BLogprobDiff(CustomTestCase):
max_lora_rank=MAX_LORA_RANK,
lora_paths={"my_lora": adapter_path},
lora_backend=LORA_BACKEND,
attention_backend="flashinfer",
prefill_attention_backend=PREFILL_ATTENTION_BACKEND,
decode_attention_backend=DECODE_ATTENTION_BACKEND,
**attention_backend_kwargs(),
)
try:
test/registered/lora/test_moe_lora_info.py
+12 -4
View File
@@ -4,10 +4,18 @@ import pytest
import torch
from sglang.srt.lora.backend.base_backend import _compute_moe_lora_info
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cuda_ci,
register_xpu_ci,
)
register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=5, stage="stage-b", runner_config="1-gpu-small-amd")
register_xpu_ci(est_time=20, suite="stage-a-test-1-gpu-xpu")
DEVICE = get_device()
def _expected_adapter_enabled(
@@ -25,7 +33,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)
@@ -54,7 +62,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)
@@ -67,7 +75,7 @@ def test_compute_moe_lora_info_expands_segments(use_preallocated_buffers: bool):
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)
test/registered/lora/test_virtual_experts_kernels.py
+23 -10
View File
@@ -15,9 +15,10 @@ Covers two regression bugs that surface only with `--lora-use-virtual-experts`
wrap, or past-end) and don't get assigned to a real expert in the
consumer-block table.
Both kernels run on CUDA. The fallback is gated on `virtual_num_experts >= 1024`
in production, but we exercise it directly here at smaller sizes for cheaper
iteration; one test sticks to the >1024 regime to mirror the production trigger.
Both kernels run on CUDA or XPU; only the CUDA-JIT align variant is CUDA-only.
The fallback is gated on `virtual_num_experts >= 1024` in production, but we
exercise it directly here at smaller sizes for cheaper iteration; one test
sticks to the >1024 regime to mirror the production trigger.
Usage:
python -m pytest test/registered/lora/test_virtual_experts_kernels.py -v
@@ -27,10 +28,19 @@ 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=14, stage="base-b", runner_config="1-gpu-small")
register_xpu_ci(est_time=20, suite="stage-a-test-1-gpu-xpu")
def _require_accelerator_device():
if not (is_cuda() or is_xpu()):
raise unittest.SkipTest("CUDA or XPU required")
return get_device(0)
from sglang.kernels.ops.moe.virtual_experts import (
_align_block_size_jit,
@@ -45,9 +55,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 +147,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 +270,13 @@ class TestAlignBlockSizeTorchSentinelBucket(_AlignBlockSizeSentinelBucketBase):
return _align_block_size_torch(topk_ids, block_size, num_experts)
@unittest.skipIf(
is_xpu(),
"_align_block_size_jit builds a CUDA JIT kernel via tvm_ffi.load_inline "
"(requires a CUDA/nvcc install); it cannot run on XPU. The torch variant "
"(TestAlignBlockSizeTorchSentinelBucket) covers the same alignment logic "
"on that platform.",
)
class TestAlignBlockSizeJitSentinelBucket(_AlignBlockSizeSentinelBucketBase):
"""Test the CUDA JIT kernel path (with fused_sanitize_expert_ids, as in
production)."""