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

[Test] Add XPU device support to unit tests (#22236)

Browse Source 
Co-authored-by: vshekhawat-hlab <vshekhawat@habana.ai>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
shubham singhal
2026-05-01 07:18:51 +08:00
committed by GitHub
co-authored by vshekhawat-hlab Ma Mingfei
parent 8b23d32ec1
commit e35ac95cdc
3 changed files with 45 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/test/attention/test_prefix_chunk_info.py
+6 -2
View File
@@ -4,6 +4,7 @@ import torch
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.utils.common import get_device
from sglang.test.test_utils import CustomTestCase
TEST_CASES = [
@@ -131,14 +132,17 @@ def check_kv_indices(forward_batch):
assert torch.allclose(computed_kv_indices, ref_kv_indices)
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
@unittest.skipIf(
not (torch.cuda.is_available() or torch.xpu.is_available()),
"Test requires CUDA or XPU",
)
class TestPrefixChunkInfo(CustomTestCase):
def setUp(self):
# Common test parameters
self.num_local_heads = 128
self.kv_lora_rank = 512
self.qk_rope_head_dim = 64
self.device = torch.device("cuda")
self.device = get_device()
self.dtype = torch.bfloat16
self.extend_len = 64
self.max_bs = 4
test/registered/attention/test_kda_kernels.py
+22 -11
View File
@@ -11,33 +11,40 @@ from sglang.srt.layers.attention.fla.kda import (
fused_recurrent_kda,
kda_gate_chunk_cumsum,
)
from sglang.srt.utils.common import get_device
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=12, suite="stage-b-test-1-gpu-large")
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
@unittest.skipIf(
not (torch.cuda.is_available() or torch.xpu.is_available()),
"Test requires CUDA or XPU",
)
class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase):
def setUp(self):
self.device = get_device()
self.token_num = 4
self.query_start_loc = torch.tensor([0, 1, 2, 3, 4], device="cuda")
self.cache_indices = torch.tensor([0, 2, 5, 8], device="cuda")
self.query_start_loc = torch.tensor([0, 1, 2, 3, 4], device=self.device)
self.cache_indices = torch.tensor([0, 2, 5, 8], device=self.device)
self.local_num_heads = 8
self.head_dim = 128
self.cache_len = 64
self.A_log = torch.randn(
1, 1, self.local_num_heads, 1, dtype=torch.float32, device="cuda"
1, 1, self.local_num_heads, 1, dtype=torch.float32, device=self.device
)
self.a = torch.randn(
1,
self.token_num,
self.local_num_heads * self.head_dim,
dtype=torch.bfloat16,
device="cuda",
device=self.device,
)
self.dt_bias = torch.randn(
self.local_num_heads * self.head_dim, dtype=torch.bfloat16, device="cuda"
self.local_num_heads * self.head_dim,
dtype=torch.bfloat16,
device=self.device,
)
self.softplus_beta = 1.0
self.softplus_threshold = 20.0
@@ -47,7 +54,7 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase):
self.local_num_heads,
self.head_dim,
dtype=torch.bfloat16,
device="cuda",
device=self.device,
)
self.k = torch.randn(
1,
@@ -55,7 +62,7 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase):
self.local_num_heads,
self.head_dim,
dtype=torch.bfloat16,
device="cuda",
device=self.device,
)
self.v = torch.randn(
1,
@@ -63,10 +70,14 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase):
self.local_num_heads,
self.head_dim,
dtype=torch.bfloat16,
device="cuda",
device=self.device,
)
self.beta = torch.randn(
1, self.token_num, self.local_num_heads, dtype=torch.bfloat16, device="cuda"
1,
self.token_num,
self.local_num_heads,
dtype=torch.bfloat16,
device=self.device,
)
self.ssm_states = torch.zeros(
@@ -75,7 +86,7 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase):
self.head_dim,
self.head_dim,
dtype=torch.float32,
device="cuda",
device=self.device,
)
def run_fused(self):
test/registered/quant/test_triton_scaled_mm.py
+17 -11
View File
@@ -5,6 +5,7 @@ import torch
import torch.testing
from sglang.srt.layers.quantization.fp8_kernel import triton_scaled_mm
from sglang.srt.utils.common import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -31,20 +32,25 @@ def torch_scaled_mm(
class TestScaledMM(CustomTestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("This test requires a CUDA device.")
torch.set_default_device("cuda")
if not (torch.cuda.is_available() or torch.xpu.is_available()):
raise unittest.SkipTest("No CUDA or XPU device available")
cls._device = get_device()
torch.set_default_device(cls._device)
def _make_inputs(self, M, K, N, in_dtype):
if in_dtype == torch.int8:
a = torch.randint(-8, 8, (M, K), dtype=in_dtype, device="cuda")
b = torch.randint(-8, 8, (K, N), dtype=in_dtype, device="cuda")
a = torch.randint(-8, 8, (M, K), dtype=in_dtype, device=self._device)
b = torch.randint(-8, 8, (K, N), dtype=in_dtype, device=self._device)
else: # fp8
a = torch.clamp(
0.1 * torch.randn((M, K), dtype=torch.float16, device="cuda"), -0.3, 0.3
0.1 * torch.randn((M, K), dtype=torch.float16, device=self._device),
-0.3,
0.3,
).to(in_dtype)
b = torch.clamp(
0.1 * torch.randn((K, N), dtype=torch.float16, device="cuda"), -0.3, 0.3
0.1 * torch.randn((K, N), dtype=torch.float16, device=self._device),
-0.3,
0.3,
).to(in_dtype)
return a, b
@@ -56,7 +62,7 @@ class TestScaledMM(CustomTestCase):
]
try:
torch.tensor([1.0], dtype=torch.float8_e4m3fn, device="cuda")
torch.tensor([1.0], dtype=torch.float8_e4m3fn, device=self._device)
test_configs.append((32, 32, 32, torch.float8_e4m3fn, torch.float16, False))
except:
print("FP8 not supported, skipping")
@@ -68,13 +74,13 @@ class TestScaledMM(CustomTestCase):
input, weight = self._make_inputs(M, K, N, in_dtype)
scale_a = 0.1 + 0.05 * torch.rand(
(M, 1), dtype=torch.float32, device="cuda"
(M, 1), dtype=torch.float32, device=self._device
)
scale_b = 0.1 + 0.05 * torch.rand(
(N, 1), dtype=torch.float32, device="cuda"
(N, 1), dtype=torch.float32, device=self._device
)
bias = (
0.01 * torch.randn((M, N), dtype=out_dtype, device="cuda")
0.01 * torch.randn((M, N), dtype=out_dtype, device=self._device)
if with_bias
else None
)