minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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-03 11:15:14 -08:00
committed by GitHub
co-authored by jundu Gao, Pengfei
parent 0a6925639b
commit 495290aefd
20 changed files with 237 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/rl/test_fp32_lm_head.py
+7 -4
View File
@@ -12,6 +12,7 @@ from sglang.srt.server_args import (
get_global_server_args,
set_global_server_args_for_scheduler,
)
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=9, suite="stage-b-test-small-1-gpu")
@@ -19,7 +20,7 @@ register_amd_ci(est_time=15, suite="stage-b-test-small-1-gpu-amd")
class LMHeadStub(nn.Module):
def __init__(self, vocab, hidden, dtype, device="cuda"):
def __init__(self, vocab, hidden, dtype, device=get_device()):
super().__init__()
self.weight = nn.Parameter(
torch.randn(vocab, hidden, dtype=dtype, device=device)
@@ -36,8 +37,10 @@ class DummyMeta:
class TestLMHeadFP32(unittest.TestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("needs CUDA GPU")
if not torch.cuda.is_available() and not (
hasattr(torch, "xpu") and torch.xpu.is_available()
):
raise unittest.SkipTest("needs CUDA GPU or XPU")
def _make_logprocessor(self, vocab_size, enable_fp32):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
@@ -54,7 +57,7 @@ class TestLMHeadFP32(unittest.TestCase):
expected_a_dtype,
expected_b_dtype,
):
device = "cuda"
device = get_device()
BATCH_SIZE, HIDDEN_SIZE, VOCAB_SIZE = 2, 64, 128
hidden_state = torch.randn(
BATCH_SIZE, HIDDEN_SIZE, dtype=hidden_state_dtype, device=device
test/registered/rl/test_release_memory_occupation.py
+13 -16
View File
@@ -31,7 +31,6 @@ import os
import time
import unittest
import torch
from transformers import AutoModelForCausalLM
import sglang as sgl
@@ -40,6 +39,7 @@ from sglang.srt.constants import (
GPU_MEMORY_TYPE_KV_CACHE,
GPU_MEMORY_TYPE_WEIGHTS,
)
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST,
@@ -48,6 +48,9 @@ from sglang.test.test_utils import (
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
CustomTestCase,
empty_gpu_cache,
get_gpu_count,
get_gpu_memory_gb,
)
register_cuda_ci(
@@ -60,10 +63,6 @@ register_cuda_ci(
_DEBUG_EXTRA = False
def get_gpu_memory_gb():
return torch.cuda.device_memory_used() / 1024**3
class TestReleaseMemoryOccupation(CustomTestCase):
def _setup_engine(
self,
@@ -120,9 +119,7 @@ class TestReleaseMemoryOccupation(CustomTestCase):
def test_release_and_resume_occupation(self):
# Without multi-stage release and resume, we need to carefully control the memory fraction to avoid OOM
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
assert (
torch.cuda.device_count() >= 2
), "Need at least 2 GPUs for tensor parallel tests"
assert get_gpu_count() >= 2, "Need at least 2 GPUs for tensor parallel tests"
for tp_size in [1, 2]:
@@ -165,13 +162,13 @@ class TestReleaseMemoryOccupation(CustomTestCase):
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map="cuda",
device_map=get_device(),
)
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
torch.cuda.empty_cache()
empty_gpu_cache()
print("generate (#2)")
outputs = engine.generate(params["prompt"], params["sampling_params"])[
@@ -232,7 +229,7 @@ class TestReleaseMemoryOccupation(CustomTestCase):
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
for tp_size in [1, 2]:
if tp_size == 2 and torch.cuda.device_count() < 2:
if tp_size == 2 and get_gpu_count() < 2:
continue
print(f"Testing tp_size={tp_size} for test_multi_stage_release_and_resume")
@@ -320,14 +317,14 @@ class TestReleaseMemoryOccupation(CustomTestCase):
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map="cuda",
device_map=get_device(),
)
gpu_memory_usage_after_loaded_hf_model = get_gpu_memory_gb()
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
torch.cuda.empty_cache()
empty_gpu_cache()
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE])
gpu_memory_usage_after_resume_kv_cache = get_gpu_memory_gb()
@@ -399,13 +396,13 @@ class TestReleaseMemoryOccupation(CustomTestCase):
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map="cuda",
device_map=get_device(),
)
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
torch.cuda.empty_cache()
empty_gpu_cache()
print("generate (#2)")
outputs = engine.generate(params["prompt_moe"], params["sampling_params_moe"])[
@@ -463,7 +460,7 @@ class TestReleaseMemoryOccupation(CustomTestCase):
engine.update_weights_from_disk(model_name)
# destroy the hf model
torch.cuda.empty_cache()
empty_gpu_cache()
print("generate (#2)")
outputs = engine.generate(