chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

enable ut test for xpu devices (#11712)

Browse Source 
Co-authored-by: jundu <jun.du@intel.com>
Co-authored-by: Gao, Pengfei <pengfei.gao@intel.com>
This commit is contained in:
DiweiSun
2026-02-04 03:15:14 +08:00
committed by GitHub
parent 0a6925639b
commit 495290aefd
20 changed files with 238 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/moe/topk.py
View File

@@ -57,6 +57,7 @@ from sglang.srt.utils import (
is_cuda,
is_hip,
is_npu,
is_xpu,
)
from sglang.srt.utils.patch_torch import register_fake_if_exists
@@ -69,6 +70,7 @@ _is_cuda = is_cuda()
_is_hip = is_hip()
_is_cpu = is_cpu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_xpu = is_xpu()
_is_npu = is_npu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -85,7 +87,7 @@ if _is_cuda:
except ImportError as e:
pass
if _is_cuda or _is_hip:
if _is_cuda or _is_hip or _is_xpu:
from sgl_kernel import topk_softmax
try:

30
python/sglang/test/runners.py
View File

@@ -32,7 +32,7 @@ from transformers import (
from sglang.srt.entrypoints.engine import Engine
from sglang.srt.model_loader.ci_weight_validation import ci_validate_and_clean_hf_cache
from sglang.srt.utils import is_npu, load_image
from sglang.srt.utils import get_device, is_npu, load_image
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, calculate_rouge_l
@@ -122,7 +122,7 @@ def _get_sentence_transformer_embedding_model(
modules=[word_embedding_model, pooling_model], truncate_dim=matryoshka_dim
)
return model.cuda()
return model.to(get_device())
@dataclass
@@ -271,7 +271,7 @@ class HFRunner:
torch_dtype=torch_dtype,
trust_remote_code=self.trust_remote_code,
low_cpu_mem_usage=True,
).cuda()
).to(get_device())
elif self.model_type == "embedding":
if "gme-qwen2-vl" in model_path.lower():
self.model = AutoModelForVision2Seq.from_pretrained(
@@ -279,10 +279,10 @@ class HFRunner:
torch_dtype=torch_dtype,
trust_remote_code=False,
low_cpu_mem_usage=True,
).cuda()
).to(get_device())
self.processor = AutoProcessor.from_pretrained(model_path)
elif "clip" in model_path.lower():
self.model = AutoModel.from_pretrained(model_path).cuda()
self.model = AutoModel.from_pretrained(model_path).to(get_device())
self.processor = AutoProcessor.from_pretrained(model_path)
else:
self.model = _get_sentence_transformer_embedding_model(
@@ -295,7 +295,7 @@ class HFRunner:
model_path,
torch_dtype=torch_dtype,
trust_remote_code=self.needs_trust_remote_code(model_path),
).cuda()
).to(get_device())
else:
raise Exception(f"Unrecognized model type {self.model_type}")
self.tokenizer = get_tokenizer(
@@ -338,15 +338,19 @@ class HFRunner:
images=image[0], return_tensors="pt"
)
logits = self.model.get_image_features(
pixel_values=inputs.data["pixel_values"].cuda(),
pixel_values=inputs.data["pixel_values"].to(
get_device()
),
).tolist()
else:
inputs = self.tokenizer(
prompts, padding=True, return_tensors="pt"
)
logits = self.model.get_text_features(
input_ids=inputs.data["input_ids"].cuda(),
attention_mask=inputs.data["attention_mask"].cuda(),
input_ids=inputs.data["input_ids"].to(get_device()),
attention_mask=inputs.data["attention_mask"].to(
get_device()
),
).tolist()
else:
logits = self.model.encode(prompts).tolist()
@@ -354,7 +358,7 @@ class HFRunner:
elif self.model_type == "cross_encoder":
inputs = self.tokenizer(
prompts, padding=True, return_tensors="pt"
).to("cuda")
).to(get_device())
scores = self.model(**inputs).logits
scores = scores.squeeze().tolist()
if not isinstance(scores, list):
@@ -369,7 +373,7 @@ class HFRunner:
)
conv_tokenized = self.tokenizer(
conv_formatted, return_tensors="pt"
).to("cuda")
).to(get_device())
scores.append(
float(self.model(**conv_tokenized).logits[0][0].item())
)
@@ -426,9 +430,9 @@ class HFRunner:
for i, p in enumerate(prompts):
if isinstance(p, str):
input_ids = tokenizer.encode(p, return_tensors="pt").cuda()
input_ids = tokenizer.encode(p, return_tensors="pt").to(get_device())
else:
input_ids = torch.tensor([p], device="cuda")
input_ids = torch.tensor([p], device=get_device())
if lora_paths is not None and lora_paths[i] is not None:
from peft import PeftModel

40
python/sglang/test/test_utils.py
View File

@@ -37,7 +37,9 @@ from sglang.srt.environ import envs
from sglang.srt.utils import (
get_bool_env_var,
get_device,
is_cuda,
is_port_available,
is_xpu,
kill_process_tree,
retry,
)
@@ -2243,6 +2245,44 @@ def intel_amx_benchmark(extra_args=None, min_throughput=None):
return decorator
def get_gpu_count():
if get_device() == "cpu":
gpu_count = 0
else:
gpu_count = torch.accelerator.device_count()
return gpu_count
def empty_gpu_cache():
"""
Unified empty_cache for PyTorch 2.8 (no torch.accelerator)
and PyTorch 2.9+ (where torch.accelerator.empty_cache() exists).
"""
if hasattr(torch, "accelerator") and hasattr(torch.accelerator, "empty_cache"):
return torch.accelerator.empty_cache()
# CUDA
if hasattr(torch, "cuda") and torch.cuda.is_available():
torch.cuda.empty_cache()
return
# XPU (Intel)
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
return
return
def get_gpu_memory_gb():
if is_cuda():
return torch.cuda.device_memory_used() / 1024**3
elif is_xpu():
return torch.xpu.memory_allocated() / 1024**3
else:
return 0
def run_doctests(obj: Callable[..., Any] | ModuleType):
mod = inspect.getmodule(obj)
globals = dict(mod.__dict__)