minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 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/attention/test_create_kvindices.py
+9 -8
View File
@@ -4,6 +4,7 @@ import numpy as np
import torch
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -15,30 +16,28 @@ register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")
class TestCreateKvIndices(CustomTestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
torch.set_default_device(get_device())
def _run_test(self, batch, max_batch, max_context_len):
req_to_token = torch.arange(
max_batch * max_context_len, dtype=torch.int32, device="cuda"
max_batch * max_context_len, dtype=torch.int32, device=get_device()
).reshape((max_batch, max_context_len))
req_pool_indices = torch.tensor(
torch.from_numpy(
np.random.choice(range(max_batch), size=batch, replace=False)
),
dtype=torch.int32,
device="cuda",
device=get_device(),
)
paged_kernel_lens = torch.tensor(
torch.from_numpy(
np.random.choice(range(max_context_len), size=batch, replace=False)
),
dtype=torch.int32,
device="cuda",
device=get_device(),
)
kv_indptr = torch.zeros((batch + 1,), dtype=torch.int32, device="cuda")
kv_indptr = torch.zeros((batch + 1,), dtype=torch.int32, device=get_device())
kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
# ref
@@ -53,7 +52,9 @@ class TestCreateKvIndices(CustomTestCase):
).contiguous()
# triton
kv_indices_triton = torch.empty(kv_indptr[-1], dtype=torch.int32, device="cuda")
kv_indices_triton = torch.empty(
kv_indptr[-1], dtype=torch.int32, device=get_device()
)
create_flashinfer_kv_indices_triton[(batch,)](
req_to_token,
req_pool_indices,
test/registered/attention/test_wave_attention_kernels.py
+64 -39
View File
@@ -21,6 +21,7 @@ from sglang.srt.layers.attention.wave_ops.extend_attention import extend_attenti
from sglang.srt.layers.attention.wave_ops.prefill_attention import (
prefill_attention_wave,
)
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci
# Wave attention kernel unit tests (AMD only - requires wave_lang)
@@ -47,24 +48,24 @@ class TestWaveAttention(unittest.TestCase):
extend_seq_len = 1024
b_seq_len_prefix = torch.full(
(B,), N_CTX // B, dtype=torch.int32, device="cuda"
(B,), N_CTX // B, dtype=torch.int32, device=get_device()
)
b_seq_len_extend = torch.full(
(B,), extend_seq_len, dtype=torch.int32, device="cuda"
(B,), extend_seq_len, dtype=torch.int32, device=get_device()
)
b_seq_len = b_seq_len_prefix + b_seq_len_extend
max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
b_req_idx = torch.arange(B, dtype=torch.int32, device="cuda")
b_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda")
b_req_idx = torch.arange(B, dtype=torch.int32, device=get_device())
b_start_loc = torch.zeros((B,), dtype=torch.int32, device=get_device())
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device="cuda")
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=get_device())
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda")
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=get_device())
kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
kv_indices = torch.zeros(
(b_seq_len_prefix.sum().item(),), dtype=torch.int32, device="cuda"
(b_seq_len_prefix.sum().item(),), dtype=torch.int32, device=get_device()
)
for i in range(B):
@@ -75,15 +76,21 @@ class TestWaveAttention(unittest.TestCase):
total_token_num = torch.sum(b_seq_len).item()
extend_token_num = torch.sum(b_seq_len_extend).item()
k_buffer = torch.empty(
(total_token_num, H_KV, D), dtype=dtype, device="cuda"
(total_token_num, H_KV, D), dtype=dtype, device=get_device()
).normal_(mean=0.1, std=0.2)
v_buffer = torch.empty(
(total_token_num, H_KV, D), dtype=dtype, device="cuda"
(total_token_num, H_KV, D), dtype=dtype, device=get_device()
).normal_(mean=0.1, std=0.2)
k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda")
v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda")
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
k_extend = torch.empty(
(extend_token_num, H_KV, D), dtype=dtype, device=get_device()
)
v_extend = torch.empty(
(extend_token_num, H_KV, D), dtype=dtype, device=get_device()
)
q_extend = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
@@ -96,20 +103,22 @@ class TestWaveAttention(unittest.TestCase):
extend_start_in_buffer:extend_end_in_buffer
]
q_extend[extend_start:extend_end] = torch.empty(
(b_seq_len_extend[i], H_Q, D), dtype=dtype, device="cuda"
(b_seq_len_extend[i], H_Q, D), dtype=dtype, device=get_device()
).normal_(mean=0.1, std=0.2)
o_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
o_extend = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
o_extend_mask = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device="cuda"
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
o_redundant = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device="cuda"
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
b_seq_len_extend = b_seq_len - b_seq_len_prefix
max_len_extend = torch.max(b_seq_len_extend, 0)[0].item()
qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda")
qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=get_device())
qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)
custom_mask = None
@@ -129,7 +138,9 @@ class TestWaveAttention(unittest.TestCase):
is_causal = True
o_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
o_extend = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
extend_attention_fwd(
q_extend,
k_extend,
@@ -146,7 +157,9 @@ class TestWaveAttention(unittest.TestCase):
max_len_extend,
)
o_wave = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
o_wave = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
extend_attention_wave(
q_extend,
k_extend,
@@ -181,33 +194,37 @@ class TestWaveAttention(unittest.TestCase):
total_tokens = B * seq_len
sm_scale = 1.0 / (D**0.5)
max_kv_splits = 8
num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device="cuda")
num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device=get_device())
# q represents the new token being generated, one per batch
q = torch.randn(B, H_Q, D, dtype=dtype, device="cuda")
q = torch.randn(B, H_Q, D, dtype=dtype, device=get_device())
# k_buffer and v_buffer represent all previous tokens
k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device="cuda")
k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=get_device())
v_buffer = torch.randn(
total_tokens, H_KV, D_V, dtype=dtype, device=get_device()
)
# o will have the same shape as q
o_triton = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
o_triton = torch.zeros(B, H_Q, D_V, dtype=dtype, device=get_device())
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=get_device())
req_to_token = torch.arange(total_tokens, device="cuda", dtype=torch.int32)
b_req_idx = torch.zeros(B + 1, device="cuda", dtype=torch.int32)
b_seq_len = torch.full((B,), seq_len, device="cuda", dtype=torch.int32)
req_to_token = torch.arange(
total_tokens, device=get_device(), dtype=torch.int32
)
b_req_idx = torch.zeros(B + 1, device=get_device(), dtype=torch.int32)
b_seq_len = torch.full((B,), seq_len, device=get_device(), dtype=torch.int32)
b_req_idx[1 : B + 1] = torch.cumsum(b_seq_len, dim=0)
attn_logits = torch.empty(
(B, H_Q, max_kv_splits, D_V + 1),
dtype=torch.float32,
device="cuda",
device=get_device(),
)
attn_lse = torch.empty(
(B, H_Q, max_kv_splits),
dtype=torch.float32,
device="cuda",
device=get_device(),
)
logit_cap = 0.0
@@ -233,13 +250,13 @@ class TestWaveAttention(unittest.TestCase):
attn_logits = torch.empty(
attn_logits_shape,
dtype=torch.float32,
device="cuda",
device=get_device(),
)
attn_logits_max = torch.empty(
attn_logits_max_shape,
dtype=torch.float32,
device="cuda",
device=get_device(),
)
decode_attention_wave(
@@ -288,17 +305,25 @@ class TestWaveAttention(unittest.TestCase):
max_seq_len = max(seq_lens)
# Create random input tensors
q = torch.randn(sum(seq_lens), num_heads, head_dim, dtype=dtype, device="cuda")
k = torch.randn(sum(seq_lens), kv_heads, head_dim, dtype=dtype, device="cuda")
v = torch.randn(sum(seq_lens), kv_heads, head_dim, dtype=dtype, device="cuda")
q = torch.randn(
sum(seq_lens), num_heads, head_dim, dtype=dtype, device=get_device()
)
k = torch.randn(
sum(seq_lens), kv_heads, head_dim, dtype=dtype, device=get_device()
)
v = torch.randn(
sum(seq_lens), kv_heads, head_dim, dtype=dtype, device=get_device()
)
o_triton = torch.zeros(
sum(seq_lens), num_heads, head_dim, dtype=dtype, device="cuda"
sum(seq_lens), num_heads, head_dim, dtype=dtype, device=get_device()
)
o = torch.zeros(
sum(seq_lens), num_heads, head_dim, dtype=dtype, device=get_device()
)
o = torch.zeros(sum(seq_lens), num_heads, head_dim, dtype=dtype, device="cuda")
# Create b_start_loc and b_seq_len tensors
b_start_loc = torch.tensor([0, seq_lens[0]], device="cuda")
b_seq_len = torch.tensor(seq_lens, device="cuda")
b_start_loc = torch.tensor([0, seq_lens[0]], device=get_device())
b_seq_len = torch.tensor(seq_lens, device=get_device())
context_attention_fwd(
q, k, v, o_triton, b_start_loc, b_seq_len, max_seq_len, is_causal=is_causal
test/registered/core/test_hidden_states.py
+3 -3
View File
@@ -4,7 +4,7 @@ import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import sglang as sgl
from sglang.srt.utils import is_hip
from sglang.srt.utils import get_device, is_hip
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
@@ -57,7 +57,7 @@ class TestHiddenState(CustomTestCase):
)
model = AutoModelForCausalLM.from_pretrained(
model_path, torch_dtype=torch.bfloat16, device_map="cuda"
model_path, torch_dtype=torch.bfloat16, device_map=get_device()
)
for input_id, output in zip(input_ids, outputs):
@@ -75,7 +75,7 @@ class TestHiddenState(CustomTestCase):
i.unsqueeze(0) if len(i.shape) == 1 else i
for i in output["meta_info"]["hidden_states"]
]
).to("cuda")
).to(get_device())
print("=== SRT Hiddens ===")
print(sg_hidden_states)
test/registered/layers/mamba/test_causal_conv1d.py
+7 -11
View File
@@ -18,6 +18,8 @@ from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
causal_conv1d_fn,
causal_conv1d_update,
)
from sglang.srt.utils import get_device
from sglang.test.test_utils import empty_gpu_cache
def causal_conv1d_ref(
@@ -154,10 +156,8 @@ def causal_conv1d_opcheck_fn(
@pytest.mark.parametrize("width", [4])
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, itype):
if not torch.cuda.is_available():
pytest.skip("CUDA device not available")
device = "cuda"
device = get_device()
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
if itype == torch.bfloat16:
rtol, atol = 1e-2, 5e-2
@@ -193,10 +193,8 @@ def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, ity
def test_causal_conv1d_update_with_batch_gather(
batch_size, with_padding, dim, width, seqlen, has_bias, silu_activation, itype
):
if not torch.cuda.is_available():
pytest.skip("CUDA device not available")
device = "cuda"
device = get_device()
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
if itype == torch.bfloat16:
rtol, atol = 1e-2, 5e-2
@@ -273,11 +271,9 @@ def test_causal_conv1d_update_with_batch_gather(
def test_causal_conv1d_varlen(
batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
):
if not torch.cuda.is_available():
pytest.skip("CUDA device not available")
device = "cuda"
torch.cuda.empty_cache()
device = get_device()
empty_gpu_cache()
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
if itype == torch.bfloat16:
rtol, atol = 1e-2, 5e-2
@@ -336,7 +332,7 @@ def test_causal_conv1d_varlen(
weight,
bias=bias,
conv_states=final_states,
query_start_loc=cumsum.cuda(),
query_start_loc=cumsum.to(get_device()),
seq_lens_cpu=torch.tensor(seqlens[0]),
cache_indices=padded_state_indices,
has_initial_state=has_initial_states,
test/registered/model_loading/test_modelopt_loader.py
+2 -1
View File
@@ -15,6 +15,7 @@ from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES
from sglang.srt.model_loader.loader import ModelOptModelLoader
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -62,7 +63,7 @@ class TestModelOptModelLoader(CustomTestCase):
self.model_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
self.load_config = LoadConfig()
self.device_config = DeviceConfig(device="cuda")
self.device_config = DeviceConfig(device=get_device())
# Create a basic model config with unified quantization flag
self.model_config = ModelConfig(
test/registered/moe/test_fused_moe.py
+20 -20
View File
@@ -9,9 +9,9 @@ from sglang.srt.layers.moe.topk import TopKConfig, select_experts
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.quantization.fp8_utils import normalize_e4m3fn_to_e4m3fnuz
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.srt.utils import is_hip
from sglang.srt.utils import get_device, get_device_capability, is_hip
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
from sglang.test.test_utils import CustomTestCase, empty_gpu_cache
register_cuda_ci(est_time=80, suite="stage-b-test-large-1-gpu")
register_amd_ci(est_time=30, suite="stage-b-test-small-1-gpu-amd")
@@ -25,8 +25,8 @@ class TestFusedMOE(CustomTestCase):
TOP_KS = [2, 6]
@staticmethod
def create_random_cuda_tensor(shape, dtype, mean=0, std=0.01):
"""Create a random CUDA tensor
def create_random_gpu_tensor(shape, dtype, mean=0, std=0.01):
"""Create a random Torch(device) tensor
Args:
shape: Tensor shape
@@ -35,9 +35,9 @@ class TestFusedMOE(CustomTestCase):
std: Standard deviation
Returns:
torch.Tensor: Randomly initialized CUDA tensor
torch.Tensor: Randomly initialized Torch(device) tensor
"""
return torch.empty(shape, dtype=dtype, device="cuda").normal_(mean, std)
return torch.empty(shape, dtype=dtype, device=get_device()).normal_(mean, std)
def get_tolerance(self, dtype):
"""Get tolerance values for different data types
@@ -109,20 +109,20 @@ class TestFusedMOE(CustomTestCase):
if use_fp8_w8a8:
# AssertionError: fp8e4nv data type is not supported on CUDA arch < 89
capability = torch.cuda.get_device_capability()
capability = get_device_capability()
if not _is_hip and not (capability[0] >= 9 or capability == (8, 9)):
return
a = self.create_random_cuda_tensor((m, k), dtype)
w1 = self.create_random_cuda_tensor((e, 2 * n, k), dtype)
w2 = self.create_random_cuda_tensor((e, k, n), dtype)
a = self.create_random_gpu_tensor((m, k), dtype)
w1 = self.create_random_gpu_tensor((e, 2 * n, k), dtype)
w2 = self.create_random_gpu_tensor((e, k, n), dtype)
w1 = w1.to(torch.float8_e4m3fn)
w2 = w2.to(torch.float8_e4m3fn)
score = self.create_random_cuda_tensor((m, e), dtype)
w1_scale = self.create_random_cuda_tensor(e, torch.float32)
w2_scale = self.create_random_cuda_tensor(e, torch.float32)
a1_scale = self.create_random_cuda_tensor(1, torch.float32)
a2_scale = self.create_random_cuda_tensor(1, torch.float32)
score = self.create_random_gpu_tensor((m, e), dtype)
w1_scale = self.create_random_gpu_tensor(e, torch.float32)
w2_scale = self.create_random_gpu_tensor(e, torch.float32)
a1_scale = self.create_random_gpu_tensor(1, torch.float32)
a2_scale = self.create_random_gpu_tensor(1, torch.float32)
# Handle HIP case: normalize float8 weights so fused kernel doesn't break
# on ROCm.
@@ -172,10 +172,10 @@ class TestFusedMOE(CustomTestCase):
sglang_output, torch_output, rtol=rtol, atol=atol
)
else:
a = self.create_random_cuda_tensor((m, k), dtype)
w1 = self.create_random_cuda_tensor((e, 2 * n, k), dtype)
w2 = self.create_random_cuda_tensor((e, k, n), dtype)
score = self.create_random_cuda_tensor((m, e), dtype)
a = self.create_random_gpu_tensor((m, k), dtype)
w1 = self.create_random_gpu_tensor((e, 2 * n, k), dtype)
w2 = self.create_random_gpu_tensor((e, k, n), dtype)
score = self.create_random_gpu_tensor((m, e), dtype)
topk_output = select_experts(
hidden_states=a,
@@ -236,7 +236,7 @@ class TestFusedMOE(CustomTestCase):
dtype,
use_fp8_w8a8=use_fp8_w8a8,
)
torch.cuda.empty_cache()
empty_gpu_cache()
pbar.update(1)
test/registered/quant/test_awq_dequant.py
+2 -1
View File
@@ -16,12 +16,13 @@ from sglang.srt.layers.quantization.awq_triton import (
awq_dequantize_triton,
awq_gemm_triton,
)
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
register_amd_ci(est_time=2, suite="stage-a-test-1-amd")
device = "cuda"
device = get_device()
def reverse_awq_order(t: torch.Tensor) -> torch.Tensor:
test/registered/quant/test_gptqmodel_dynamic.py
+2 -2
View File
@@ -5,7 +5,7 @@ import requests
import torch
from sglang.srt.server_args import set_global_server_args_for_scheduler
from sglang.srt.utils import kill_process_tree
from sglang.srt.utils import get_device, kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -49,7 +49,7 @@ def check_quant_method(model_path: str, use_marlin_kernel: bool):
model_config = ModelConfig.from_server_args(server_args)
load_config = LoadConfig()
device_config = DeviceConfig("cuda")
device_config = DeviceConfig(get_device())
model = get_model(
model_config=model_config, load_config=load_config, device_config=device_config
)
test/registered/radix_cache/test_mamba_unittest.py
+4 -3
View File
@@ -17,6 +17,7 @@ from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToToke
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import ServerArgs, 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")
@@ -39,7 +40,7 @@ class TestMamba(unittest.TestCase):
num_layers = 48
global_interval = 4
dtype = torch.bfloat16
device = "cuda"
device = get_device()
full_attention_layer_ids = [
i for i in range(global_interval - 1, num_layers, global_interval)
]
@@ -67,7 +68,7 @@ class TestMamba(unittest.TestCase):
max_num_reqs = 10
mamba_cache_size = 20
max_context_len = 128
device = "cuda"
device = get_device()
global_interval = 4
num_layers = 48
full_attention_layer_ids = [
@@ -152,7 +153,7 @@ class TestMamba(unittest.TestCase):
max_num_reqs = 10
mamba_cache_size = 20
max_context_len = 128
device = "cuda"
device = get_device()
full_attention_layer_ids = [
i for i in range(global_interval - 1, num_layers, global_interval)
]
test/registered/radix_cache/test_swa_unittest.py
+4 -3
View File
@@ -13,6 +13,7 @@ from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
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=8, suite="stage-b-test-large-1-gpu")
@@ -37,7 +38,7 @@ class TestSWA(unittest.TestCase):
num_layers = 48
global_interval = 4
dtype = torch.bfloat16
device = "cuda"
device = get_device()
full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
full_attention_layer_ids_set = set(full_attention_layer_ids)
swa_attention_layer_ids = [
@@ -89,7 +90,7 @@ class TestSWA(unittest.TestCase):
num_layers = 48
global_interval = 4
dtype = torch.bfloat16
device = "cuda"
device = get_device()
full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
full_attention_layer_ids_set = set(full_attention_layer_ids)
swa_attention_layer_ids = [
@@ -243,7 +244,7 @@ class TestSWA(unittest.TestCase):
num_layers = 48
global_interval = 4
dtype = torch.bfloat16
device = "cuda"
device = get_device()
full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
full_attention_layer_ids_set = set(full_attention_layer_ids)
swa_attention_layer_ids = [
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(
test/registered/spec/utils/test_build_eagle_tree.py
+10 -9
View File
@@ -6,6 +6,7 @@ from sglang.srt.speculative.eagle_utils import (
build_tree_kernel_efficient,
organize_draft_results,
)
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=3, suite="stage-b-test-small-1-gpu")
@@ -17,7 +18,7 @@ class TestBuildEagleTree(unittest.TestCase):
def test_build_tree_kernel_efficient(self):
"""Test the build_tree_kernel_efficient function with known inputs and expected outputs."""
verified_id = torch.tensor([29974, 13], device="cuda", dtype=torch.int32)
verified_id = torch.tensor([29974, 13], device=get_device(), dtype=torch.int32)
score_list = [
torch.tensor(
[
@@ -25,7 +26,7 @@ class TestBuildEagleTree(unittest.TestCase):
[[9.7476e-01, 2.2219e-02, 6.5031e-04, 1.3212e-04]],
],
dtype=torch.float32,
device="cuda",
device=get_device(),
),
torch.tensor(
[
@@ -43,7 +44,7 @@ class TestBuildEagleTree(unittest.TestCase):
],
],
dtype=torch.float32,
device="cuda",
device=get_device(),
),
torch.tensor(
[
@@ -61,7 +62,7 @@ class TestBuildEagleTree(unittest.TestCase):
],
],
dtype=torch.float32,
device="cuda",
device=get_device(),
),
torch.tensor(
[
@@ -79,14 +80,14 @@ class TestBuildEagleTree(unittest.TestCase):
],
],
dtype=torch.float32,
device="cuda",
device=get_device(),
),
]
token_list = [
torch.tensor(
[[29896, 29906, 29900, 29945], [13, 2, 29871, 28956]],
dtype=torch.int64,
device="cuda",
device=get_device(),
),
torch.tensor(
[
@@ -127,7 +128,7 @@ class TestBuildEagleTree(unittest.TestCase):
259,
],
],
device="cuda",
device=get_device(),
),
torch.tensor(
[
@@ -168,7 +169,7 @@ class TestBuildEagleTree(unittest.TestCase):
2186,
],
],
device="cuda",
device=get_device(),
),
torch.tensor(
[
@@ -209,7 +210,7 @@ class TestBuildEagleTree(unittest.TestCase):
13,
],
],
device="cuda",
device=get_device(),
),
]
parents_list = [